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