Snapshot call-time args in the IIP Enzyme reverse rule (fixes #62)

ext/FunctionWrappersWrappersEnzymeExt.jl

@@ -154,6 +154,35 @@ end
 # Reverse mode rules
 # =============================================================================
 
+# The Const-return (IIP) reverse rule re-runs `Enzyme.autodiff(Reverse, …)` on
+# the unwrapped function during the *reverse* pass.  That recomputation reads
+# the arguments' current `.val`s — but a caller may mutate those arguments
+# between the forward call and the reverse pass (an ODE integrator steps `u`
+# after every RHS call).  Without snapshotting, the VJP is then taken about the
+# wrong state and the gradient is silently wrong.  `_snapshot`/`_restore!` let
+# the rule pin the call-time argument values in its tape.
+#
+# To keep this cheap, only arguments that Enzyme's `overwritten(config)`
+# analysis flags as possibly modified between the forward and reverse passes
+# are copied; everything else (and immutables like `t`) tapes `nothing` and is
+# never copied.  `overwritten` is a conservative over-approximation, so gating
+# on it can never *miss* an argument that actually needs snapshotting.  A
+# non-mutating call therefore allocates nothing here.
+@inline _snapshot(v::AbstractArray) = copy(v)
+@inline _snapshot(v) = v
+@inline _restore!(dst::AbstractArray, snap::AbstractArray) = (copyto!(dst, snap); nothing)
+@inline _restore!(@nospecialize(dst), @nospecialize(snap)) = nothing
+
+# 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 +201,25 @@ 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). Run the primal for its side effects, and snapshot the call-time
+# argument values into the tape so the reverse pass can take the VJP about the
+# state at the time of *this* call — even if the caller mutates the arguments
+# afterwards (the ODE-integrator pattern: `wf(du, u, p, t)` followed by an
+# in-place step on `u`). Snapshot BEFORE running the primal in case `f_orig`
+# mutates its inputs.
 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)
+    # `overwritten` is indexed (func, args...), so arg `i` is entry `i + 1`.
+    ow = EnzymeRules.overwritten(config)
+    tape = ntuple(i -> (ow[i + 1] ? _snapshot(args[i].val) : nothing), Val(N))
     pargs = ntuple(i -> args[i].val, Val(N))
     f_orig(pargs...)
-    return EnzymeRules.AugmentedReturn(nothing, nothing, nothing)
+    return EnzymeRules.AugmentedReturn(nothing, nothing, tape)
 end
 
 # Duplicated / BatchDuplicated return: record the primal so that reverse has
@@ -299,11 +335,14 @@ end
 # accumulate into any `Duplicated` arg shadow buffers (the SciML IIP
 # pattern).  Simply returning `nothing` left Duplicated shadows at zero.
 #
-# 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.
+# `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)[1]` already returns
+# the per-argument gradient tuple in exactly the form a reverse rule must hand
+# back: `nothing` for each `Const`/`Duplicated`/`BatchDuplicated` arg (their
+# gradients are accumulated in-place), and the concrete gradient for each
+# `Active` arg — with exact per-slot types (Enzyme rejects a union-typed
+# `Tuple{Union{Nothing,Float64},…}`).  So we forward that tuple directly, 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 +351,23 @@ 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
+    # Nothing to do if every argument is Const.
+    if any(a -> !(a isa EnzymeCore.Const), args)
+        # Temporarily restore the call-time argument values (snapshotted in the
+        # tape during augmented_primal) so the VJP is taken about the correct
+        # state, then put the live values back so the surrounding reverse pass
+        # is left undisturbed.  Without this, a caller that mutates an argument
+        # after the forward call (an ODE integrator stepping `u`) makes the
+        # recomputation differentiate `f_orig` about the wrong state.
+        live = ntuple(i -> (tape[i] === nothing ? nothing : _snapshot(args[i].val)), Val(N))
+        map((a, snap) -> _restore!(a.val, snap), args, tape)
+        raw = Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)[1]::NTuple{N, Any}
+        map((a, snap) -> _restore!(a.val, snap), args, live)
+        # `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
+    return ntuple(_ -> nothing, Val(N))
 end
 
 end

test/Enzyme/enzyme_tests.jl

@@ -164,16 +164,20 @@ 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 runs the wrapped function for its side effects and tapes a snapshot
+    # of the call-time argument values (so a later mutation by the caller can't
+    # make the reverse pass differentiate about the wrong state).  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).  Mark only `x`
+    # as overwritten so we can check the rule snapshots exactly that arg.
+    rconfig = EnzymeRules.RevConfig{false, false, 1, (false, true, false), false, false}()
 
     counter[] = 0
     aug = EnzymeRules.augmented_primal(
@@ -183,7 +187,7 @@ end
     @test counter[] == 1                       # primal ran exactly once
     @test aug.primal === nothing               # NeedsPrimal == false
     @test aug.shadow === nothing
-    @test aug.tape === nothing
+    @test aug.tape == (3.0, nothing)           # only the overwritten arg (x) is snapshotted
 
     # Reverse step — dret is Const (passed as TYPE not instance in reverse
     # rules).  Enzyme's rule protocol requires concrete gradients for Active
@@ -330,6 +334,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 +533,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