diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index 1c05580..1a852ce 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/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,8 +201,12 @@ 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}, @@ -181,9 +214,12 @@ function EnzymeRules.augmented_primal( 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 diff --git a/test/Enzyme/enzyme_tests.jl b/test/Enzyme/enzyme_tests.jl index 5a8f5f9..9696c0b 100644 --- a/test/Enzyme/enzyme_tests.jl +++ b/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