Merge pull request #45 from ChrisRackauckas-Claude/enzyme-forward-rule

Propagate set_runtime_activity through FWW Enzyme forward rules

Author: ChrisRackauckas

ext/FunctionWrappersWrappersEnzymeExt.jl

@@ -5,6 +5,32 @@ using Enzyme
 using EnzymeCore
 using EnzymeCore.EnzymeRules
 
+# =============================================================================
+# Helper: build a Forward mode from FwdConfig flags
+# =============================================================================
+# The outer caller may invoke `Enzyme.autodiff(set_runtime_activity(Forward), …)`
+# or `set_strong_zero(Forward)` or `ForwardWithPrimal`.  Those settings flow
+# into the `EnzymeRules.FwdConfig{NeedsPrimal, NeedsShadow, Width,
+# RuntimeActivity, StrongZero}` type parameters of the rule's first argument.
+# Before this fix the rules hard-coded plain `Forward` in their inner
+# `Enzyme.autodiff` delegation, which dropped both `RuntimeActivity` and
+# `StrongZero` — breaking users who need `set_runtime_activity(Forward)` to
+# avoid `EnzymeRuntimeActivityError` inside the wrapped function (the SciML
+# `Rosenbrock23(autodiff = AutoEnzyme(set_runtime_activity(Forward)))` path
+# on an in-place time-dependent RHS; see OrdinaryDiffEq.jl PR #3518).
+#
+# `_fwd_mode(needs_primal, RuntimeActivity, StrongZero)` returns the
+# `ForwardMode` matching the outer config so the delegated call inherits
+# those flags.
+@inline function _fwd_mode(
+    ::Val{NeedsPrimal}, ::Val{RuntimeActivity}, ::Val{StrongZero}
+) where {NeedsPrimal, RuntimeActivity, StrongZero}
+    mode = NeedsPrimal ? ForwardWithPrimal : Forward
+    RuntimeActivity && (mode = Enzyme.set_runtime_activity(mode))
+    StrongZero && (mode = Enzyme.set_strong_zero(mode))
+    return mode
+end
+
 # =============================================================================
 # Forward mode rules — generalized to arbitrary batch width W
 # =============================================================================
@@ -17,11 +43,12 @@ function EnzymeRules.forward(
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {T, W, N, RuntimeActivity, StrongZero}
     f_orig = unwrap(func.val)
+    mode = _fwd_mode(Val(false), Val(RuntimeActivity), Val(StrongZero))
     if W == 1
-        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{T}, args...)
+        shadow_result = Enzyme.autodiff(mode, Const(f_orig), Duplicated{T}, args...)
         return shadow_result[1]::T
     else
-        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), BatchDuplicated{T, W}, args...)
+        shadow_result = Enzyme.autodiff(mode, Const(f_orig), BatchDuplicated{T, W}, args...)
         return shadow_result[1]::NTuple{W, T}
     end
 end
@@ -36,12 +63,16 @@ function EnzymeRules.forward(
     f_orig = unwrap(func.val)
     pargs = ntuple(i -> args[i].val, Val(N))
     primal = f_orig(pargs...)::T
+    # Use plain Forward (not ForwardWithPrimal) here — we already have the
+    # primal above, and `Duplicated{T}` / `BatchDuplicated{T,W}` as the RT
+    # annotation asks only for the shadow.
+    mode = _fwd_mode(Val(false), Val(RuntimeActivity), Val(StrongZero))
     if W == 1
-        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{T}, args...)
+        shadow_result = Enzyme.autodiff(mode, Const(f_orig), Duplicated{T}, args...)
         shadow = shadow_result[1]::T
         return Duplicated(primal, shadow)
     else
-        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), BatchDuplicated{T, W}, args...)
+        shadow_result = Enzyme.autodiff(mode, Const(f_orig), BatchDuplicated{T, W}, args...)
         shadows = shadow_result[1]::NTuple{W, T}
         return BatchDuplicated(primal, shadows)
     end
@@ -69,14 +100,20 @@ end
 # orders of magnitude in OrdinaryDiffEq.jl v7.  Delegate to `Enzyme.autodiff`
 # on the unwrapped function with a Const return annotation so the Duplicated
 # arg shadows are propagated correctly and no return is produced.
+#
+# IMPORTANT: forward the `RuntimeActivity` and `StrongZero` flags from the
+# outer config into the delegated `Enzyme.autodiff` call.  Prior to this
+# fix the rule hard-coded `Forward`, silently dropping
+# `set_runtime_activity(Forward)` on the way down into `f_orig`. 
 function EnzymeRules.forward(
     ::EnzymeRules.FwdConfig{false, false, W, RuntimeActivity, StrongZero},
     func::EnzymeCore.Const{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {W, N, RuntimeActivity, StrongZero}
     f_orig = unwrap(func.val)
-    Enzyme.autodiff(Forward, Const(f_orig), Const, args...)
+    mode = _fwd_mode(Val(false), Val(RuntimeActivity), Val(StrongZero))
+    Enzyme.autodiff(mode, Const(f_orig), Const, args...)
     return nothing
 end
 
@@ -151,11 +188,21 @@ function EnzymeRules.augmented_primal(
     end
 end
 
+# Helper: build a Forward mode reflecting a RevConfig's runtime_activity /
+# strong_zero flags so the internal forward-mode delegation inside reverse
+# rules inherits the user's outer config.
+@inline function _fwd_mode_from_rev(config::EnzymeRules.RevConfig)
+    mode = Forward
+    EnzymeRules.runtime_activity(config) && (mode = Enzyme.set_runtime_activity(mode))
+    EnzymeRules.strong_zero(config) && (mode = Enzyme.set_strong_zero(mode))
+    return mode
+end
+
 # Varargs reverse: compute each partial via forward-mode AD on the unwrapped
 # function, then scale by dret. This avoids type-inference issues that arise
 # from calling autodiff(Reverse, Const{Any}(...), ...).
 @generated function _fww_reverse_grads(
-    f_orig, dret_val::T, args::Vararg{EnzymeCore.Active, N}
+    mode, f_orig, dret_val::T, args::Vararg{EnzymeCore.Active, N}
 ) where {T, N}
     # Build forward-mode calls for each partial derivative
     exprs = []
@@ -164,7 +211,7 @@ end
         dups = [:(Duplicated(args[$j].val, $(seeds[j]))) for j in 1:N]
         Ti = :(eltype(typeof(args[$i])))
         push!(exprs, quote
-            fwd = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{$T}, $(dups...))
+            fwd = Enzyme.autodiff(mode, Const(f_orig), Duplicated{$T}, $(dups...))
             $Ti(fwd[1] * dret_val)::$Ti
         end)
     end
@@ -179,7 +226,7 @@ function EnzymeRules.reverse(
     args::Vararg{EnzymeCore.Active, N}
 ) where {T, N}
     f_orig = unwrap(func.val)
-    return _fww_reverse_grads(f_orig, dret.val, args...)
+    return _fww_reverse_grads(_fwd_mode_from_rev(config), f_orig, dret.val, args...)
 end
 
 # Handle mixed Active/Const args: return nothing for Const, gradient for Active
@@ -192,6 +239,7 @@ function EnzymeRules.reverse(
 ) where {N}
     f_orig = unwrap(func.val)
     dret_val = dret.val
+    mode = _fwd_mode_from_rev(config)
     return ntuple(Val(N)) do i
         if args[i] isa EnzymeCore.Const
             nothing
@@ -204,7 +252,7 @@ function EnzymeRules.reverse(
                     Duplicated(args[j].val, zero(eltype(typeof(args[j]))))
                 end
             end
-            fwd = Enzyme.autodiff(Forward, Const(f_orig), Duplicated, dup_args...)
+            fwd = Enzyme.autodiff(mode, Const(f_orig), Duplicated, dup_args...)
             fwd[1] * dret_val
         end
     end

test/enzyme_tests.jl

@@ -277,3 +277,83 @@ end
     @test u_shadow[1] ≈ expected_u_grad
 end
 
+# =============================================================================
+# Runtime-activity propagation through the FWW forward rules.
+#
+# Prior to this fix the rules hard-coded plain `Forward` when delegating to
+# `Enzyme.autodiff`, silently dropping the caller's
+# `set_runtime_activity(Forward)` flag.  Enzyme's static IR-level activity
+# analysis can't see through `FunctionWrappersWrapper`'s opaque cfunction
+# indirection, so the inner call raised `EnzymeRuntimeActivityError` inside
+# `@.` broadcast's `broadcast_unalias` → `mightalias` — despite
+# `set_runtime_activity` being set on the outer `autodiff` call.
+#
+# Upstream motivation: OrdinaryDiffEq.jl PR #3518 —
+#   Rosenbrock23(autodiff = AutoEnzyme(set_runtime_activity(Enzyme.Forward)))
+# on any time-dependent in-place RHS routed through DiffEqBase's
+# `wrapfun_iip`.  Here we reproduce the shape (`f!(du, u, p, t) = @. du = …`)
+# in a 4-arg `FunctionWrappersWrapper` matching DiffEqBase's
+# `wrapfun_iip` output, and assert both that (a) the call completes without
+# an `EnzymeRuntimeActivityError` and (b) the resulting tangent is
+# numerically correct.
+# =============================================================================
+
+@testset "Enzyme Forward: set_runtime_activity propagates through FWW (IIP, time-dependent)" begin
+    # DiffEqBase's `wrapfun_iip(ff, (u, u, p, t))` shape.
+    const_INPUTS = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64}
+
+    # 1) Time-independent RHS — ∂du/∂t = 0.
+    f!(du, u, p, t) = (@. du = p * u; nothing)
+    fww = FunctionWrappersWrapper(f!, (const_INPUTS,), (Nothing,))
+
+    u = [1.0, 2.0, 3.0]
+    p = [0.5, 0.5, 0.5]
+    t = 1.7
+    du = zero(u); ddu = zero(u); dt = 1.0
+
+    Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Forward),
+        Const(fww), Const,
+        Duplicated(du, ddu),
+        Const(u), Const(p),
+        Duplicated(t, dt),
+    )
+    @test du ≈ p .* u
+    @test all(iszero, ddu)
+
+    # 2) Non-trivial time dependence: g!(du,u,p,t) = @. sin(t)*u.
+    #    Expected ∂du/∂t = cos(t) .* u.
+    g!(du, u, p, t) = (@. du = sin(t) * u; nothing)
+    gww = FunctionWrappersWrapper(g!, (const_INPUTS,), (Nothing,))
+
+    du2 = zero(u); ddu2 = zero(u)
+    Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Forward),
+        Const(gww), Const,
+        Duplicated(du2, ddu2),
+        Const(u), Const(p),
+        Duplicated(t, 1.0),
+    )
+    @test du2 ≈ sin(t) .* u
+    @test ddu2 ≈ cos(t) .* u
+
+    # 3) Confirm the rule also propagates `set_strong_zero(Forward)` (the
+    #    other ForwardMode flag carried in FwdConfig) — another RHS that
+    #    doesn't need runtime activity but exercises a distinct flag.
+    h!(du, u, p, t) = (du[1] = u[1] * t; nothing)
+    hww = FunctionWrappersWrapper(
+        h!, (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},),
+        (Nothing,)
+    )
+    du_h = [0.0]; ddu_h = [0.0]
+    Enzyme.autodiff(
+        Enzyme.set_strong_zero(Forward),
+        Const(hww), Const,
+        Duplicated(du_h, ddu_h),
+        Const([2.0]), Const([0.0]),
+        Duplicated(3.5, 1.0),
+    )
+    @test du_h[1] ≈ 2.0 * 3.5     # primal: u[1] * t = 7.0
+    @test ddu_h[1] ≈ 2.0          # ∂(u[1]*t)/∂t = u[1]
+end
+