Forward-mode rules for jl_eqtable_get/jl_eqtable_put (#3135 follow-up)

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
 

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