Add rules for jl_eqtable_pop and jl_eqtable_nextind (#3042)#3043
Open
Add rules for jl_eqtable_pop and jl_eqtable_nextind (#3042)#3043
Conversation
Implement augfwd, fwd, and rev rules for jl_eqtable_pop mimicking eqtableput Add argument activity annotations for jl_eqtable_pop Mark jl_eqtable_nextind as pure read-only and inactive Add both functions to validation whitelist Add GPUArraysCore to test dependencies and include an @allowscalar test in iddict.jl Fixes #3042 Co-authored-by: Gemini 3.1 Pro (High) <gemini@google.com>
Contributor
|
Your PR requires formatting changes to meet the project's style guidelines. Click here to view the suggested changes.diff --git a/src/compiler/validation.jl b/src/compiler/validation.jl
index a4704e18..69394911 100644
--- a/src/compiler/validation.jl
+++ b/src/compiler/validation.jl
@@ -89,8 +89,8 @@ function __init__()
"ijl_eqtable_get",
"jl_eqtable_put",
"ijl_eqtable_put",
- "jl_eqtable_pop",
- "ijl_eqtable_pop",
+ "jl_eqtable_pop",
+ "ijl_eqtable_pop",
"memcmp",
"memchr",
"jl_get_nth_field_checked",
diff --git a/src/llvm/attributes.jl b/src/llvm/attributes.jl
index 03c0a0ac..4b1bd0c4 100644
--- a/src/llvm/attributes.jl
+++ b/src/llvm/attributes.jl
@@ -210,10 +210,10 @@ const nofreefns = Set{String}((
"cuCtxGetId",
"cuDeviceGetName",
"ijl_eqtable_get",
- "jl_eqtable_pop",
- "ijl_eqtable_pop",
- "jl_eqtable_nextind",
- "ijl_eqtable_nextind",
+ "jl_eqtable_pop",
+ "ijl_eqtable_pop",
+ "jl_eqtable_nextind",
+ "ijl_eqtable_nextind",
"cuCtxGetApiVersion",
"cuCtxSetCurrent",
))
@@ -344,8 +344,8 @@ const inactivefns = Set{String}((
"jl_array_to_string",
"ijl_array_to_string",
"pcre2_jit_compile_8",
- "jl_eqtable_nextind",
- "ijl_eqtable_nextind",
+ "jl_eqtable_nextind",
+ "ijl_eqtable_nextind",
# "jl_"
))
@@ -762,8 +762,8 @@ function annotate!(mod::LLVM.Module)
"jl_reshape_array",
"ijl_eqtable_get",
"jl_eqtable_get",
- "jl_eqtable_pop",
- "ijl_eqtable_pop",
+ "jl_eqtable_pop",
+ "ijl_eqtable_pop",
"jl_gc_run_pending_finalizers",
"ijl_try_substrtod",
"jl_try_substrtod",
@@ -1128,8 +1128,8 @@ function annotate!(mod::LLVM.Module)
"memory",
MemoryEffect(
(MRI_ModRef << getLocationPos(ArgMem)) |
- (MRI_NoModRef << getLocationPos(InaccessibleMem)) |
- (MRI_NoModRef << getLocationPos(Other)),
+ (MRI_NoModRef << getLocationPos(InaccessibleMem)) |
+ (MRI_NoModRef << getLocationPos(Other)),
).data,
),
)
@@ -1153,8 +1153,8 @@ function annotate!(mod::LLVM.Module)
"memory",
MemoryEffect(
(MRI_Ref << getLocationPos(ArgMem)) |
- (MRI_NoModRef << getLocationPos(InaccessibleMem)) |
- (MRI_NoModRef << getLocationPos(Other)),
+ (MRI_NoModRef << getLocationPos(InaccessibleMem)) |
+ (MRI_NoModRef << getLocationPos(Other)),
).data,
),
)
diff --git a/src/rules/llvmrules.jl b/src/rules/llvmrules.jl
index 11980cce..bcd451fa 100644
--- a/src/rules/llvmrules.jl
+++ b/src/rules/llvmrules.jl
@@ -1166,25 +1166,25 @@ end
B,
orig,
"Enzyme: Not yet implemented constant table in jl_eqtable_get " *
- string(origh) *
- " " *
- string(orig) *
- " result: " *
- string(absint(orig)) *
- " " *
- string(abs_typeof(orig, true)) *
- " dict: " *
- string(absint(origh)) *
- " " *
- string(abs_typeof(origh, true)) *
- " key " *
- string(absint(origkey)) *
- " " *
- string(abs_typeof(origkey, true)) *
- " dflt " *
- string(absint(origdflt)) *
- " " *
- string(abs_typeof(origdflt, true)),
+ string(origh) *
+ " " *
+ string(orig) *
+ " result: " *
+ string(absint(orig)) *
+ " " *
+ string(abs_typeof(orig, true)) *
+ " dict: " *
+ string(absint(origh)) *
+ " " *
+ string(abs_typeof(origh, true)) *
+ " key " *
+ string(absint(origkey)) *
+ " " *
+ string(abs_typeof(origkey, true)) *
+ " dflt " *
+ string(absint(origdflt)) *
+ " " *
+ string(abs_typeof(origdflt, true)),
)
return false
end
@@ -1196,9 +1196,9 @@ end
Base.unsafe_convert(
Cstring,
"Mixed activity for default of jl_eqtable_get " *
- string(orig) *
- " " *
- string(origdflt),
+ string(orig) *
+ " " *
+ string(origdflt),
),
orig.ref,
API.ET_MixedActivityError,
@@ -1215,7 +1215,7 @@ end
else
ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(nop)))
shadowm = LLVM.UndefValue(ST)
- for j = 1:width
+ for j in 1:width
shadowm = insert_value!(B, shadowm, nop, j - 1)
end
shadowm
@@ -1387,9 +1387,9 @@ end
Base.unsafe_convert(
Cstring,
"Mixed activity for val of jl_eqtable_put " *
- string(orig) *
- " " *
- string(origval),
+ string(orig) *
+ " " *
+ string(origval),
),
orig.ref,
API.ET_MixedActivityError,
@@ -1406,7 +1406,7 @@ end
else
ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(nop)))
shadowm = LLVM.UndefValue(ST)
- for j = 1:width
+ for j in 1:width
shadowm = insert_value!(B, shadowm, nop, j - 1)
end
shadowm
@@ -1419,10 +1419,10 @@ 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)),
+ shadowh,
+ new_from_original(gutils, origkey),
+ shadowval,
+ LLVM.null(value_type(originserted)),
]
shadowres = batch_call_same_with_inverted_arg_if_active!(B, gutils, orig, newops, newvals, false) #=lookup=#
@@ -1520,9 +1520,9 @@ end
Base.unsafe_convert(
Cstring,
"Mixed activity for default of jl_eqtable_pop " *
- string(orig) *
- " " *
- string(origdflt),
+ string(orig) *
+ " " *
+ string(origdflt),
),
orig.ref,
API.ET_MixedActivityError,
@@ -1539,7 +1539,7 @@ end
else
ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(nop)))
shadowm = LLVM.UndefValue(ST)
- for j = 1:width
+ for j in 1:width
shadowm = insert_value!(B, shadowm, nop, j - 1)
end
shadowm
@@ -1552,10 +1552,10 @@ 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),
- shadowdflt,
- LLVM.null(value_type(origfound)),
+ shadowh,
+ new_from_original(gutils, origkey),
+ shadowdflt,
+ LLVM.null(value_type(origfound)),
]
shadowres = batch_call_same_with_inverted_arg_if_active!(B, gutils, orig, newops, newvals, false) #=lookup=#
@@ -1584,9 +1584,9 @@ end
Base.unsafe_convert(
Cstring,
"Mixed activity for default of jl_eqtable_pop " *
- string(orig) *
- " " *
- string(origdflt),
+ string(orig) *
+ " " *
+ string(origdflt),
),
orig.ref,
API.ET_MixedActivityError,
@@ -1603,7 +1603,7 @@ end
else
ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(nop)))
shadowm = LLVM.UndefValue(ST)
- for j = 1:width
+ for j in 1:width
shadowm = insert_value!(B, shadowm, nop, j - 1)
end
shadowm
@@ -1616,10 +1616,10 @@ 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),
- shadowdflt,
- LLVM.null(value_type(origfound)),
+ shadowh,
+ new_from_original(gutils, origkey),
+ shadowdflt,
+ LLVM.null(value_type(origfound)),
]
shadowres = batch_call_same_with_inverted_arg_if_active!(B, gutils, orig, newops, newvals, false) #=lookup=# |
Contributor
Benchmark Results
Benchmark PlotsA plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/25425001320/artifacts/6826042797. |
Properly compute shadow values for jl_eqtable_get, jl_eqtable_put, and jl_eqtable_pop in forward mode, conditionally on shadowR. Co-authored-by: Gemini 3.1 Pro (High) <gemini@google.com>
Since IdDict stores boxed values, which are represented as pointers in Enzyme, we do not need custom reverse-mode tracking. Gradients naturally accumulate into the shadow boxes. Added tests for storing and reading active arrays and floats in IdDict with inactive keys. Co-authored-by: Gemini 3.1 Pro (High) <gemini@google.com>
Since IdDict maps seamlessly to its shadow representation, Duplicated(IdDict(), IdDict()) works naturally with our implementation. Co-authored-by: Gemini 3.1 Pro (High) <gemini@google.com>
Added tests for Enzyme.make_zero, Enzyme.make_zero!, and Enzyme.remake_zero! on IdDict. Also tested Enzyme.jacobian in both Forward and Reverse modes with an IdDict as input by passing explicit shadows. Co-authored-by: Gemini 3.1 Pro (High) <gemini@google.com>
Added a test that populates an IdDict with an active Float64, an active Vector{Float64}, and an inactive Int. This verifies that Enzyme correctly accumulates gradients when interacting with dynamically-typed boxes and respects mixed activity in runtime structures.
Co-authored-by: Gemini 3.1 Pro (High) <gemini@google.com>
Member
Author
|
The x86 error is JuliaLLVM/LLVM.jl#548 |
wsmoses
reviewed
Apr 30, 2026
| if shadowR != C_NULL | ||
| unsafe_store!(shadowR, shadowres.ref) | ||
| end | ||
| return false |
Member
There was a problem hiding this comment.
I'm not sure this is correct, if the popped value guaranteed to be the same jlvalue within the list?
Exercises cases left uncovered by the existing tests, where pop!'s return value is used in computation (existing tests only call pop! for side effects via @allowscalar/delete!): - pop! return value used in forward and reverse mode - pop! with non-trivial stored expression (shadow must match put!, not default) - pop! with key present but explicit default (default shadow must not leak) - split reverse mode: gradient flows into the object captured at forward time, not into whatever is currently in the shadow dict (mutation-between-passes test) - pop! on missing key returns active default with correct gradient - overwrite same key twice: pop! must carry second put's shadow, not first - double pop same key: second (missing) pop must not inherit first's shadow - two distinct keys: per-key shadows must not mix - pop then put back: shadow table reflects re-inserted transformed value Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implement augfwd, fwd, and rev rules for jl_eqtable_pop mimicking eqtableput
Add argument activity annotations for jl_eqtable_pop
Mark jl_eqtable_nextind as pure read-only and inactive
Add both functions to validation whitelist
Add GPUArraysCore to test dependencies and include an @allowscalar test in iddict.jl
Fixes #3042
Co-authored-by: Gemini 3.1 Pro (High) gemini@google.com