diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index 1c05580..b401cf8 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/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,8 +192,13 @@ 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}, @@ -181,9 +206,14 @@ function EnzymeRules.augmented_primal( 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 diff --git a/test/Enzyme/enzyme_tests.jl b/test/Enzyme/enzyme_tests.jl index 5a8f5f9..1a72020 100644 --- a/test/Enzyme/enzyme_tests.jl +++ b/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