Split reverse-mode thunk for the IIP Enzyme rule (alternative to #63)

ext/FunctionWrappersWrappersEnzymeExt.jl

@@ -154,6 +154,26 @@ end
 # Reverse mode rules
 # =============================================================================
 
+# A `ReverseSplitNoPrimal` mode reflecting the rule's RevConfig runtime-activity
+# / strong-zero flags, so the split thunk we delegate to inherits the caller's
+# outer settings.
+@inline function _rev_split_mode(config::EnzymeRules.RevConfig)
+    mode = Enzyme.ReverseSplitNoPrimal
+    EnzymeRules.runtime_activity(config) && (mode = Enzyme.set_runtime_activity(mode))
+    EnzymeRules.strong_zero(config) && (mode = Enzyme.set_strong_zero(mode))
+    return mode
+end
+
+# Build one slot of a reverse rule's return tuple: `nothing` for every
+# non-Active arg (their gradients accumulate in-place), and the concrete
+# gradient for each `Active` arg.  Dispatching on the annotation *type* keeps
+# the resulting tuple exactly typed (e.g. `Tuple{Nothing, Nothing, Float64}`)
+# even though the raw `autodiff`/thunk return is `Any`-typed inside the rule —
+# Enzyme rejects a union-typed return.  `g` is the matching entry of that raw
+# per-argument gradient tuple.
+@inline _revslot(::EnzymeCore.Active{T}, g) where {T} = convert(T, g)::T
+@inline _revslot(::EnzymeCore.Annotation, @nospecialize(g)) = nothing
+
 function EnzymeRules.augmented_primal(
         config::EnzymeRules.RevConfig,
         func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
@@ -172,18 +192,28 @@ function EnzymeRules.augmented_primal(
 end
 
 # Const return (e.g. IIP functions returning Nothing, or any non-differentiated
-# return). Just run the primal for its side effects; no tape is needed because
-# the reverse pass has nothing to propagate back from the return.
+# return). Delegate to a split reverse-mode thunk on the *unwrapped* function so
+# Enzyme differentiates it directly, exactly as it would if the function were
+# never wrapped. We forward the rule's `overwritten` flags as the thunk's
+# `ModifiedBetween`, so Enzyme's own tape caches any argument the caller mutates
+# before the reverse pass (the ODE-integrator pattern: `wf(du, u, p, t)` then an
+# in-place step on `u`). The forward thunk's tape and the reverse thunk are
+# stashed for the reverse rule.
 function EnzymeRules.augmented_primal(
         config::EnzymeRules.RevConfig,
         func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
         RT::Type{<:EnzymeCore.Const},
         args::Vararg{EnzymeCore.Annotation, N}
     ) where {N}
     f_orig = unwrap(func.val)
-    pargs = ntuple(i -> args[i].val, Val(N))
-    f_orig(pargs...)
-    return EnzymeRules.AugmentedReturn(nothing, nothing, nothing)
+    mode = Enzyme.ReverseSplitModified(
+        _rev_split_mode(config), Val(EnzymeRules.overwritten(config))
+    )
+    fwd_thunk, rev_thunk = Enzyme.autodiff_thunk(
+        mode, EnzymeCore.Const{typeof(f_orig)}, EnzymeCore.Const, map(typeof, args)...
+    )
+    tape = fwd_thunk(EnzymeCore.Const(f_orig), args...)[1]
+    return EnzymeRules.AugmentedReturn(nothing, nothing, (tape, rev_thunk))
 end
 
 # Duplicated / BatchDuplicated return: record the primal so that reverse has
@@ -296,14 +326,15 @@ end
 # Const return — Enzyme passes the RT as a `Type{<:Const}` to `reverse`, not
 # as an instance.  Delegate the reverse pass to
 # `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)` so gradients
-# accumulate into any `Duplicated` arg shadow buffers (the SciML IIP
-# pattern).  Simply returning `nothing` left Duplicated shadows at zero.
+# accumulate into any `Duplicated` arg shadow buffers (the SciML IIP pattern).
 #
-# Per Enzyme's rule return-type protocol, `Active` args require a concrete
-# scalar gradient (not `nothing`).  Under a `Const` return there is no
-# gradient source, so Active arg gradients are zero.  `Duplicated` /
-# `BatchDuplicated` args return `nothing` because their gradients are
-# accumulated in-place by the `Enzyme.autodiff(Reverse, …)` call above.
+# The reverse thunk stashed by `augmented_primal` reads its cached tape (which
+# already captured any `ModifiedBetween` args) and accumulates gradients into
+# the `Duplicated` arg shadows in place; its `[1]` return is the per-argument
+# gradient tuple.  We rebuild that through `map(_revslot, …)` so it is exactly
+# typed (Enzyme rejects a union-typed `Tuple{Union{Nothing,Float64},…}`), which
+# also makes `Active` args (e.g. `t` in a time-dependent IIP rhs) correct rather
+# than zeroed.
 function EnzymeRules.reverse(
         config::EnzymeRules.RevConfig,
         func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
@@ -312,18 +343,11 @@ function EnzymeRules.reverse(
         args::Vararg{EnzymeCore.Annotation, N}
     ) where {N}
     f_orig = unwrap(func.val)
-    # Only worth invoking Enzyme.autodiff when at least one arg is
-    # Duplicated/BatchDuplicated — otherwise there's nothing to accumulate.
-    if any(a -> a isa EnzymeCore.Duplicated || a isa EnzymeCore.BatchDuplicated, args)
-        Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)
-    end
-    return ntuple(Val(N)) do i
-        if args[i] isa EnzymeCore.Active
-            zero(eltype(typeof(args[i])))
-        else
-            nothing
-        end
-    end
+    tape_data, rev_thunk = tape
+    raw = rev_thunk(EnzymeCore.Const(f_orig), args..., tape_data)[1]::NTuple{N, Any}
+    # `map` over tuples specialises per element, so dispatching `_revslot` on
+    # each arg's concrete annotation type yields an exactly-typed result.
+    return map(_revslot, args, raw)
 end
 
 end

test/Enzyme/enzyme_tests.jl

@@ -164,26 +164,29 @@ end
 
 @testset "Enzyme reverse mode, Const return — augmented_primal runs primal" begin
     # Mirrors the forward {false, false} case on the reverse side. Augmented
-    # primal runs the wrapped function for its side effects and returns
-    # AugmentedReturn(nothing, nothing, nothing).  Reverse returns `nothing`
-    # per arg since there is no return derivative to propagate.
+    # primal delegates to a split reverse-mode thunk on the unwrapped function:
+    # it runs the forward pass once (which executes the primal for its side
+    # effects) and stashes `(enzyme_tape, reverse_thunk)` so the reverse rule
+    # can run the matching reverse pass.  Reverse returns `nothing`/zero per arg
+    # since there is no return derivative to propagate.
     counter = Ref(0)
     g(x, y) = (counter[] += 1; x + y)  # returns Float64 (ignored via Const RT)
     fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,))
 
     # Construct a concrete RevConfig. Fields:
     # (NeedsPrimal, NeedsShadow, Width, Overwritten, RuntimeActivity, StrongZero)
-    rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}()
+    # Overwritten is indexed (func, args...) — here (func, x, y).
+    rconfig = EnzymeRules.RevConfig{false, false, 1, (false, true, false), false, false}()
 
     counter[] = 0
     aug = EnzymeRules.augmented_primal(
         rconfig, Const(fww), EnzymeCore.Const{Float64},
         Active(3.0), Active(4.0)
     )
-    @test counter[] == 1                       # primal ran exactly once
+    @test counter[] == 1                       # primal ran exactly once (forward pass)
     @test aug.primal === nothing               # NeedsPrimal == false
     @test aug.shadow === nothing
-    @test aug.tape === nothing
+    @test length(aug.tape) == 2                # (enzyme_tape, reverse_thunk)
 
     # Reverse step — dret is Const (passed as TYPE not instance in reverse
     # rules).  Enzyme's rule protocol requires concrete gradients for Active
@@ -330,6 +333,101 @@ end
     @test u_shadow[1] ≈ expected_u_grad
 end
 
+# =============================================================================
+# Regression for the wrong gradient when a wrapped IIP function's arguments are
+# MUTATED AFTER the call — the ODE-integrator pattern that the whole-solve
+# Enzyme adjoint exercises (and the root cause of the EnsembleProblem adjoint
+# failure, SciMLSensitivity.jl#1424).
+#
+# The Const-return reverse rule re-runs `Enzyme.autodiff(Reverse, …)` on the
+# unwrapped function during the reverse pass.  If it differentiates about the
+# arguments' *current* state rather than their *call-time* state, then any
+# caller that steps `u` after the RHS call gets a silently wrong gradient.
+# Before the snapshot/restore tape fix these end-to-end gradients were wrong.
+# =============================================================================
+
+@testset "Enzyme Reverse: IIP wrapper, args mutated after call (single step)" begin
+    f!(du, u, p, t) = (du[1] = p[1] * u[1]; du[2] = p[2] * u[2]^2; nothing)
+    ARGT = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64}
+
+    function loss(p)
+        u = [1.5, 2.0]
+        du = zero(u)
+        wf = FunctionWrappersWrapper(f!, (ARGT,), (Nothing,))
+        wf(du, u, p, 0.0)
+        @inbounds for k in 1:2
+            u[k] += 0.05 * du[k]          # mutate u AFTER the wrapped call
+        end
+        return du[1]^2 + du[2]^2          # loss depends on du only
+    end
+
+    p = [0.7, 0.4]
+    # du = [p1*1.5, p2*4];  loss = (1.5 p1)^2 + (4 p2)^2
+    # ∂loss/∂p = [2*1.5^2*p1, 2*4^2*p2] = [4.5 p1, 32 p2]; evaluated at CALL-TIME u
+    g = collect(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), loss, p)[1])
+    @test g ≈ [4.5 * p[1], 32 * p[2]]
+end
+
+@testset "Enzyme Reverse: IIP wrapper in a multi-step integrator" begin
+    f!(du, u, p, t) = (
+        du[1] = -p[1] * u[1] + p[2] * u[2];
+        du[2] = -p[3] * u[2] + p[4] * u[1]; nothing
+    )
+    ARGT = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64}
+
+    function loss(p)
+        u = [1.0, 2.0]
+        du = zero(u)
+        wf = FunctionWrappersWrapper(f!, (ARGT,), (Nothing,))
+        for _ in 1:8
+            wf(du, u, p, 0.0)
+            @inbounds for k in 1:2
+                u[k] += 0.05 * du[k]      # integrator step mutates u each call
+            end
+        end
+        return sum(abs2, u)
+    end
+
+    p = [1.0, 0.5, 2.0, 0.3]
+    g = collect(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), loss, p)[1])
+
+    # central finite-difference reference (no extra deps)
+    fd = map(eachindex(p)) do i
+        h = 1.0e-6
+        pp = copy(p); pp[i] += h
+        pm = copy(p); pm[i] -= h
+        (loss(pp) - loss(pm)) / (2h)
+    end
+    @test g ≈ fd rtol = 1.0e-4
+end
+
+@testset "Enzyme Reverse: IIP wrapper with a mix of Duplicated and Active args" begin
+    # A time-dependent in-place rhs, differentiated so the reverse rule sees
+    # (Duplicated du, Duplicated u, Duplicated p, Active t).  The rule must
+    # return the *real* gradient for the Active `t` with an exact-typed tuple
+    # (Nothing per Duplicated arg, Float64 for the Active).  Before the fix this
+    # returned a union-typed `(nothing, …, 0.0)` — Enzyme rejected it with a
+    # `ReverseRuleReturnError`, and the `t`-gradient was zeroed rather than
+    # computed.
+    f!(du, u, p, t) = (du[1] = p[1] * u[1] + t * u[2]; du[2] = p[2] * u[2]; nothing)
+    ARGT = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64}
+
+    function loss(x)              # x = [p1, p2, t]
+        u = [1.5, 2.0]
+        du = zero(u)
+        wf = FunctionWrappersWrapper(f!, (ARGT,), (Nothing,))
+        wf(du, u, [x[1], x[2]], x[3])   # t = x[3] flows in as an Active scalar
+        return du[1]^2 + du[2]^2
+    end
+
+    x = [0.7, 0.4, 0.9]
+    g = collect(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), loss, x)[1])
+    # du = [p1*u1 + t*u2, p2*u2]; loss = du1^2 + du2^2
+    du1 = x[1] * 1.5 + x[3] * 2.0
+    du2 = x[2] * 2.0
+    @test g ≈ [2 * du1 * 1.5, 2 * du2 * 2.0, 2 * du1 * 2.0]   # ∂/∂t = 2*du1*u2 ≠ 0
+end
+
 # =============================================================================
 # Runtime-activity propagation through the FWW forward rules.
 #
@@ -434,7 +532,8 @@ end
     du = [0.0];       du_shadow = [1.0]
     u = [3.0];       u_shadow = [0.0]
 
-    rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}()
+    # Overwritten indexed (func, du, u); none modified between fwd and rev here.
+    rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false, false), false, false}()
     aug = EnzymeRules.augmented_primal(
         rconfig,
         Duplicated(fww, fww),                # <-- Duplicated FWW