Mooncake Upgrades (#645)

Author: willtebbutt

DifferentiationInterface/Project.toml

@@ -55,13 +55,13 @@ FiniteDifferences = "0.12.31"
 ForwardDiff = "0.10.36"
 JuliaFormatter = "1"
 LinearAlgebra = "<0.0.1,1"
-Mooncake = "0.4.11"
+Mooncake = "0.4.52"
 PolyesterForwardDiff = "0.1.2"
 ReverseDiff = "1.15.1"
 SparseArrays = "<0.0.1,1"
 SparseConnectivityTracer = "0.5.0,0.6"
-StaticArrays = "1.9.7"
 SparseMatrixColorings = "0.4.9"
+StaticArrays = "1.9.7"
 Symbolics = "5.27.1, 6"
 Tracker = "0.2.33"
 Zygote = "0.6.69"

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -6,22 +6,13 @@ using DifferentiationInterface: Context, PullbackPrep, unwrap
 using Mooncake:
     CoDual,
     Config,
-    NoRData,
-    NoTangent,
-    build_rrule,
-    fdata,
-    get_interpreter,
-    increment!!,
     primal,
-    rdata,
-    set_to_zero!!,
     tangent,
     tangent_type,
     value_and_pullback!!,
-    zero_codual,
-    zero_fcodual,
     zero_tangent,
-    __value_and_pullback!!
+    prepare_pullback_cache,
+    Mooncake
 
 DI.check_available(::AutoMooncake) = true
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -1,65 +1,49 @@
-struct MooncakeOneArgPullbackPrep{Y,R,DX,DY} <: PullbackPrep
-    y_prototype::Y
-    rrule::R
-    dx_righttype::DX
+struct MooncakeOneArgPullbackPrep{Tcache,DY} <: PullbackPrep
+    cache::Tcache
     dy_righttype::DY
 end
 
 function DI.prepare_pullback(
     f, backend::AutoMooncake, x, ty::NTuple, contexts::Vararg{Context,C}
 ) where {C}
-    y = f(x, map(unwrap, contexts)...)
     config = get_config(backend)
-    rrule = build_rrule(
-        get_interpreter(),
-        Tuple{typeof(f),typeof(x),typeof.(map(unwrap, contexts))...};
-        debug_mode=config.debug_mode,
-        silence_debug_messages=config.silence_debug_messages,
+    cache = prepare_pullback_cache(
+        f, x, map(unwrap, contexts)...; config.debug_mode, config.silence_debug_messages
     )
-    dx_righttype = zero_tangent(x)
+    y = f(x, map(unwrap, contexts)...)
     dy_righttype = zero_tangent(y)
-    prep = MooncakeOneArgPullbackPrep(y, rrule, dx_righttype, dy_righttype)
-    DI.value_and_pullback(f, prep, backend, x, ty, contexts...)  # warm up
+    prep = MooncakeOneArgPullbackPrep(cache, dy_righttype)
+    DI.value_and_pullback(f, prep, backend, x, ty, contexts...)
     return prep
 end
 
 function DI.value_and_pullback(
-    f,
+    f::F,
     prep::MooncakeOneArgPullbackPrep{Y},
     ::AutoMooncake,
     x,
     ty::NTuple{1},
     contexts::Vararg{Context,C},
-) where {Y,C}
+) where {F,Y,C}
     dy = only(ty)
     dy_righttype = dy isa tangent_type(Y) ? dy : copyto!!(prep.dy_righttype, dy)
-    new_y, (_, new_dx) = value_and_pullback!!(
-        prep.rrule, dy_righttype, f, x, map(unwrap, contexts)...
+    new_y, (_, new_dx) = Mooncake.value_and_pullback!!(
+        prep.cache, dy_righttype, f, x, map(unwrap, contexts)...
     )
-    return new_y, (new_dx,)
+    return new_y, (copy(new_dx),)
 end
 
 function DI.value_and_pullback!(
     f,
     tx::NTuple{1},
     prep::MooncakeOneArgPullbackPrep{Y},
-    ::AutoMooncake,
+    backend::AutoMooncake,
     x,
     ty::NTuple{1},
     contexts::Vararg{Context,C},
 ) where {Y,C}
-    dx, dy = only(tx), only(ty)
-    dy_righttype = dy isa tangent_type(Y) ? dy : copyto!!(prep.dy_righttype, dy)
-    dx_righttype = set_to_zero!!(prep.dx_righttype)
-    contexts_coduals = map(zero_codual ∘ unwrap, contexts)
-    y, (_, new_dx) = __value_and_pullback!!(
-        prep.rrule,
-        dy_righttype,
-        zero_codual(f),
-        CoDual(x, dx_righttype),
-        contexts_coduals...,
-    )
-    copyto!(dx, new_dx)
+    y, (new_dx,) = DI.value_and_pullback(f, prep, backend, x, ty, contexts...)
+    copyto!(only(tx), new_dx)
     return y, tx
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/twoarg.jl

@@ -1,32 +1,28 @@
-struct MooncakeTwoArgPullbackPrep{R,F,Y,DX,DY} <: PullbackPrep
-    rrule::R
-    df!::F
-    y_copy::Y
-    dx_righttype::DX
+struct MooncakeTwoArgPullbackPrep{Tcache,DY,F} <: PullbackPrep
+    cache::Tcache
     dy_righttype::DY
-    dy_righttype_after::DY
+    target_function::F
 end
 
 function DI.prepare_pullback(
     f!, y, backend::AutoMooncake, x, ty::NTuple, contexts::Vararg{Context,C}
 ) where {C}
+    target_function = function (f!, y, x, contexts...)
+        f!(y, x, contexts...)
+        return y
+    end
     config = get_config(backend)
-    rrule = build_rrule(
-        get_interpreter(),
-        Tuple{typeof(f!),typeof(y),typeof(x),typeof.(map(unwrap, contexts))...};
+    cache = prepare_pullback_cache(
+        target_function,
+        f!,
+        y,
+        x,
+        map(unwrap, contexts)...;
         debug_mode=config.debug_mode,
         silence_debug_messages=config.silence_debug_messages,
     )
-    df! = zero_tangent(f!)
-    y_copy = copy(y)
-    dx_righttype = zero_tangent(x)
-    dy_righttype = zero_tangent(y)
     dy_righttype_after = zero_tangent(y)
-    prep = MooncakeTwoArgPullbackPrep(
-        rrule, df!, y_copy, dx_righttype, dy_righttype, dy_righttype_after
-    )
-    DI.value_and_pullback(f!, y, prep, backend, x, ty, contexts...)  # warm up
-    return prep
+    return MooncakeTwoArgPullbackPrep(cache, dy_righttype_after, target_function)
 end
 
 function DI.value_and_pullback(
@@ -38,39 +34,17 @@ function DI.value_and_pullback(
     ty::NTuple{1},
     contexts::Vararg{Context,C},
 ) where {C}
-    dy = only(ty)
-
-    # Set all tangent storage to zero.
-    df! = set_to_zero!!(prep.df!)
-    dx_righttype = set_to_zero!!(prep.dx_righttype)
-    dy_righttype = set_to_zero!!(prep.dy_righttype)
-
     # Prepare cotangent to add after the forward pass.
-    dy_righttype_after = copyto!(prep.dy_righttype_after, dy)
-
-    contexts_coduals = map(zero_fcodual ∘ unwrap, contexts)
+    dy = only(ty)
+    dy_righttype_after = copyto!(prep.dy_righttype, dy)
 
-    # Run the forward pass
-    out, pb!! = prep.rrule(
-        CoDual(f!, fdata(df!)),
-        CoDual(prep.y_copy, fdata(dy_righttype)),
-        CoDual(x, fdata(dx_righttype)),
-        contexts_coduals...,
+    # Run the reverse-pass and return the results.
+    contexts = map(unwrap, contexts)
+    y_after, (_, _, _, dx) = Mooncake.value_and_pullback!!(
+        prep.cache, dy_righttype_after, prep.target_function, f!, y, x, contexts...
     )
-
-    # Verify that the output is non-differentiable.
-    @assert primal(out) === nothing
-
-    # Increment the desired cotangent dy.
-    dy_righttype = increment!!(dy_righttype, dy_righttype_after)
-
-    # Record the state of y before running the reverse pass.
-    y = copyto!(y, prep.y_copy)
-
-    # Run the reverse pass.
-    _, _, new_dx = pb!!(NoRData())
-
-    return y, (tangent(copy(fdata(dx_righttype)), new_dx),)  # TODO: remove this allocation in `value_and_pullback!`
+    copyto!(y, y_after)
+    return y, (copy(dx),) # TODO: remove this allocation in `value_and_pullback!`
 end
 
 function DI.value_and_pullback(