From 3c3e020e949ae2a4f98ae52e5d3a2b47f5f2c2bb Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 23 Jun 2026 09:59:47 -0400 Subject: [PATCH] Implement forward-mode rules for jl_eqtable_get/jl_eqtable_put Forward mode had erroring rules for the eqtable builtins ("Not yet implemented forward for jl_eqtable_get/_put"), so forward-over-reverse nesting (HVPs/Hessians) failed whenever the inner reverse's make_zero recursed through its `seen` IdDict to dedup aliased mutable objects. This is the eqtable follow-up deferred in #3137 (#3135). Replace the erroring `eqtableget_fwd` / `eqtableput_fwd` bodies with real shadow-propagating rules mirroring the existing reverse augmented-forward rules: invert the table (and default/value) pointers and re-issue the call on the inverted arguments via batch_call_same_with_inverted_arg_if_active!. Unlike the reverse put rule, forward mode does not use the `eqtable_shadow_active` guard: in forward mode the tangent is itself the shadow value, so storing an active scalar value is correct (the reverse rule rejects it because active scalars are accumulated via the tape). This also enables plain forward-mode differentiation through Dict/IdDict build + lookup. The constant-result short-circuit is unchanged, so this can only enable previously-erroring forward cases. The struct-with-undef-able-field make_zero path (EnzymeAD/Enzyme.jl#3135 MWE, SciML/SciMLSensitivity.jl#1427) advances past jl_eqtable_get with this change but still needs forward-mode mixed-activity jl_new_struct support; that remains a separate follow-up. Co-Authored-By: Chris Rackauckas --- src/rules/llvmrules.jl | 140 ++++++++++++++++++++++++++++++++++------- test/make_zero.jl | 58 +++++++++++++++++ 2 files changed, 174 insertions(+), 24 deletions(-) diff --git a/src/rules/llvmrules.jl b/src/rules/llvmrules.jl index 1be5f8f976..84881e4fa7 100644 --- a/src/rules/llvmrules.jl +++ b/src/rules/llvmrules.jl @@ -1174,25 +1174,70 @@ end return true end - err = emit_error(B, orig, "Enzyme: Not yet implemented forward for jl_eqtable_get") + width = get_width(gutils) - newo = new_from_original(gutils, orig) - API.moveBefore(newo, err, B) - normal = - (unsafe_load(normalR) != C_NULL) ? LLVM.Instruction(unsafe_load(normalR)) : nothing - if shadowR != C_NULL && normal !== nothing - width = get_width(gutils) - shadowres = UndefValue(LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(orig)))) - for idx in 1:width + origh, origkey, origdflt = arg_operands_view(orig) + + if is_constant_value(gutils, origh) + emit_error( + B, + orig, + "Enzyme: Not yet implemented forward constant table in jl_eqtable_get " * + string(orig), + ) + return false + end + + shadowh = invert_pointer(gutils, origh, B) + + shadowdflt = if is_constant_value(gutils, origdflt) + shadowdflt2 = julia_error( + Base.unsafe_convert( + Cstring, + "Mixed activity for default of jl_eqtable_get " * + string(orig) * + " " * + string(origdflt), + ), + orig.ref, + API.ET_MixedActivityError, + gutils.ref, + origdflt.ref, + B.ref, + ) + if shadowdflt2 != C_NULL + LLVM.Value(shadowdflt2) + else + nop = new_from_original(gutils, origdflt) if width == 1 - shadowres = normal + nop else - shadowres = insert_value!(B, shadowres, normal, idx - 1) + ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(nop))) + shadowm = LLVM.UndefValue(ST) + for j in 1:width + shadowm = insert_value!(B, shadowm, nop, j - 1) + end + shadowm end end - unsafe_store!(shadowR, shadowres.ref) + else + invert_pointer(gutils, origdflt, B) end + newvals = API.CValueType[API.VT_Shadow, API.VT_Primal, API.VT_Shadow] + + newops = LLVM.Value[shadowh, new_from_original(gutils, origkey), shadowdflt] + + shadowres = batch_call_same_with_inverted_arg_if_active!( + B, + gutils, + orig, + newops, + newvals, + false, + ) + + unsafe_store!(shadowR, shadowres.ref) return false end @@ -1330,25 +1375,72 @@ end if is_constant_value(gutils, orig) && is_constant_inst(gutils, orig) return true end - err = emit_error(B, orig, "Enzyme: Not yet implemented forward for jl_eqtable_put") - newo = new_from_original(gutils, orig) - API.moveBefore(newo, err, B) - normal = - (unsafe_load(normalR) != C_NULL) ? LLVM.Instruction(unsafe_load(normalR)) : nothing - if shadowR != C_NULL && normal !== nothing - width = get_width(gutils) - shadowres = UndefValue(LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(orig)))) - for idx in 1:width + width = get_width(gutils) + + origh, origkey, origval, originserted = arg_operands_view(orig) + + @assert !is_constant_value(gutils, origh) + + shadowh = invert_pointer(gutils, origh, B) + + shadowval = if is_constant_value(gutils, origval) + shadowval2 = julia_error( + Base.unsafe_convert( + Cstring, + "Mixed activity for val of jl_eqtable_put " * + string(orig) * + " " * + string(origval), + ), + orig.ref, + API.ET_MixedActivityError, + gutils.ref, + origval.ref, + B.ref, + ) + if shadowval2 != C_NULL + LLVM.Value(shadowval2) + else + nop = new_from_original(gutils, origval) if width == 1 - shadowres = normal + nop else - shadowres = insert_value!(B, shadowres, normal, idx - 1) + ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(nop))) + shadowm = LLVM.UndefValue(ST) + for j in 1:width + shadowm = insert_value!(B, shadowm, nop, j - 1) + end + shadowm end end - unsafe_store!(shadowR, shadowres.ref) + else + invert_pointer(gutils, origval, B) end + newvals = API.CValueType[API.VT_Shadow, API.VT_Primal, API.VT_Shadow, API.VT_None] + + newops = LLVM.Value[ + shadowh, + new_from_original(gutils, origkey), + shadowval, + LLVM.null(value_type(originserted)), + ] + + # Unlike the reverse augmented-forward rule (which uses `eqtable_shadow_active` + # to reject storing an active value, since active scalars need tape + # accumulation), forward mode stores the tangent itself as the shadow value, + # so an active scalar value is handled directly with no preprocess guard. + shadowres = batch_call_same_with_inverted_arg_if_active!( + B, + gutils, + orig, + newops, + newvals, + false, + ) + + unsafe_store!(shadowR, shadowres.ref) return false end diff --git a/test/make_zero.jl b/test/make_zero.jl index 56f0499b0f..4a127d6b94 100644 --- a/test/make_zero.jl +++ b/test/make_zero.jl @@ -780,4 +780,62 @@ end @test Enzyme.autodiff(Forward, isdefined_field_walk, Duplicated(2.0, 1.0)) == (2.0,) end +# Forward-mode shadow propagation through make_zero's `seen` IdDict (the +# `jl_eqtable_get` / `jl_eqtable_put` bookkeeping). make_zero uses `seen` to +# dedup aliased mutable objects, so forward-over-reverse over an aliased value +# exercises both builtins. They previously had erroring forward rules +# ("Not yet implemented forward for jl_eqtable_get"); see EnzymeAD/Enzyme.jl#3135. +function aliased_sumsq(x) + return sum(abs2, x[1]) + sum(abs2, x[2]) +end + +@testset "forward-over-reverse through make_zero seen IdDict (jl_eqtable)" begin + RA_R = set_runtime_activity(Reverse) + RA_F = set_runtime_activity(Forward) + a = [1.0, 2.0, 3.0] + x = Any[a, a] # aliased -> make_zero recurses through the `seen` IdDict + # loss(a) = 2*sum(abs2, a); gradient = 4a; Hessian = 4I, so HVP in any + # direction d is 4d. Seed the same direction in both aliases. + seed(d) = Any[copy(d), copy(d)] + r1 = Enzyme.autodiff( + RA_F, + Const(y -> Enzyme.gradient(RA_R, aliased_sumsq, y)[1]), + Duplicated(x, seed([1.0, 0.0, 0.0])), + ) + @test r1[1][1] ≈ [4.0, 0.0, 0.0] + @test r1[1][2] ≈ [4.0, 0.0, 0.0] + + r2 = Enzyme.autodiff( + RA_F, + Const(y -> Enzyme.gradient(RA_R, aliased_sumsq, y)[1]), + Duplicated(x, seed([0.0, 1.0, 0.0])), + ) + @test r2[1][1] ≈ [0.0, 4.0, 0.0] + @test r2[1][2] ≈ [0.0, 4.0, 0.0] +end + +# The same `jl_eqtable_get` / `jl_eqtable_put` forward rules also enable plain +# forward-mode differentiation through Dict/IdDict build + lookup. Forward mode +# stores the tangent directly as the shadow value, so an active scalar value +# (which the reverse rule rejects via `eqtable_shadow_active`) is fine here. +function dict_build_lookup(x) + d = Dict{Int, Float64}() + d[1] = x + d[2] = 2x + return d[1] + d[2] # = 3x +end +function iddict_build_lookup(x) + d = IdDict{Any, Float64}() + k1 = [1] + k2 = [2] + d[k1] = x + d[k2] = 3x + return d[k1] + d[k2] # = 4x +end + +@testset "forward-mode through Dict/IdDict (jl_eqtable)" begin + @test Enzyme.autodiff(Forward, dict_build_lookup, Duplicated(5.0, 1.0)) == (3.0,) + @test Enzyme.autodiff(Forward, iddict_build_lookup, Duplicated(2.0, 1.0)) == (4.0,) +end + end # module MakeZeroTests