Skip to content

EnzymeRules: bind free TypeVars in has_{f,r}rule_from_sig query type#3121

Open
Keno wants to merge 1 commit into
EnzymeAD:mainfrom
KenoAIStaging:fix/wrap-free-typevars-in-query-sig
Open

EnzymeRules: bind free TypeVars in has_{f,r}rule_from_sig query type#3121
Keno wants to merge 1 commit into
EnzymeAD:mainfrom
KenoAIStaging:fix/wrap-free-typevars-in-query-sig

Conversation

@Keno
Copy link
Copy Markdown

@Keno Keno commented May 24, 2026

_annotate_tt synthesizes a method-table query like Tuple{<:FwdConfig, <:Annotation{ft}, Type{<:Annotation}, tt...} from the caller's signature. The tt... entries came from _annotate, which builds fresh TypeVars but never wraps them in a UnionAll; ft (the first parameter of unwrap_unionall(TT0)) likewise could contain TypeVars bound by TT0's outer where clauses. Both sources left the constructed query with free TypeVars, which used to be tolerated by Julia's subtype/intersection but no longer is under JuliaLang/julia#61876 (free TypeVars now compare by identity, so e.g. typeintersect(Tuple{…,Const{Float64}}, Tuple{…,V<:Annotation{Float64}}) collapses to Union{}).

_annotate_tt now returns the new TypeVars alongside ft, tt, with ft rewrapped in TT0's UnionAll chain. Each has_*_from_sig caller wraps the constructed query in a UnionAll per fresh TypeVar via foldr(UnionAll, …). The internal callers in Enzyme.Compiler.tfunc are updated to the new return shape.

The behavior of free TypeVars in subtyping was previously undefined (in effect a bug); JuliaLang/julia#61876 makes them well-defined singleton identities, which is what surfaces this latent issue.

This change was AI-generated and should be reviewed carefully before merging.

`_annotate_tt` synthesizes a method-table query like
`Tuple{<:FwdConfig, <:Annotation{ft}, Type{<:Annotation}, tt...}` from the
caller's signature. The `tt...` entries came from `_annotate`, which builds
fresh `TypeVar`s but never wraps them in a `UnionAll`; `ft` (the first
parameter of `unwrap_unionall(TT0)`) likewise could contain `TypeVar`s bound by
`TT0`'s outer `where` clauses. Both sources left the constructed query with
free `TypeVar`s, which used to be tolerated by Julia's subtype/intersection
but no longer is under JuliaLang/julia#61876 (free `TypeVar`s now compare by
identity, so e.g. `typeintersect(Tuple{…,Const{Float64}},
Tuple{…,V<:Annotation{Float64}})` collapses to `Union{}`).

`_annotate_tt` now returns the new `TypeVar`s alongside `ft, tt`, with `ft`
rewrapped in `TT0`'s `UnionAll` chain. Each `has_*_from_sig` caller wraps the
constructed query in a `UnionAll` per fresh `TypeVar` via `foldr(UnionAll, …)`.
The internal callers in `Enzyme.Compiler.tfunc` are updated to the new return
shape.

The behavior of free TypeVars in subtyping was previously undefined (in effect
a bug); JuliaLang/julia#61876 makes them well-defined singleton identities,
which is what surfaces this latent issue.

This change was AI-generated and should be reviewed carefully before merging.

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
@github-actions
Copy link
Copy Markdown
Contributor

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.
diff --git a/lib/EnzymeCore/src/rules.jl b/lib/EnzymeCore/src/rules.jl
index d3a87f3a..440a6d25 100644
--- a/lib/EnzymeCore/src/rules.jl
+++ b/lib/EnzymeCore/src/rules.jl
@@ -355,7 +355,7 @@ function has_frule_from_sig(@nospecialize(TT);
                             method_table::Union{Nothing,Core.Compiler.MethodTableView}=nothing,
                             caller::Union{Nothing,Core.MethodInstance}=nothing)::Bool
     ft, tt, tvars = _annotate_tt(TT)
-    TT = foldr(UnionAll, tvars; init=Tuple{<:FwdConfig, <:Annotation{ft}, Type{<:Annotation}, tt...})
+    TT = foldr(UnionAll, tvars; init = Tuple{<:FwdConfig, <:Annotation{ft}, Type{<:Annotation}, tt...})
     return isapplicable(forward, TT; world, method_table, caller)
 end
 
@@ -364,7 +364,7 @@ function has_rrule_from_sig(@nospecialize(TT);
                             method_table::Union{Nothing,Core.Compiler.MethodTableView}=nothing,
                             caller::Union{Nothing,Core.MethodInstance}=nothing)::Bool
     ft, tt, tvars = _annotate_tt(TT)
-    TT = foldr(UnionAll, tvars; init=Tuple{<:RevConfig, <:Annotation{ft}, Type{<:Annotation}, tt...})
+    TT = foldr(UnionAll, tvars; init = Tuple{<:RevConfig, <:Annotation{ft}, Type{<:Annotation}, tt...})
     return isapplicable(augmented_primal, TT; world, method_table, caller)
 end
 
diff --git a/src/compiler/tfunc.jl b/src/compiler/tfunc.jl
index 34ba158d..982b198a 100644
--- a/src/compiler/tfunc.jl
+++ b/src/compiler/tfunc.jl
@@ -4,7 +4,7 @@ import EnzymeCore.EnzymeRules: FwdConfig, RevConfig, forward, augmented_primal,
 function has_frule_from_sig(@nospecialize(interp::Core.Compiler.AbstractInterpreter),
     @nospecialize(TT::Type), sv::Core.Compiler.AbsIntState, partialedge::Bool=true)::Bool
     ft, tt, tvars = _annotate_tt(TT)
-    TT = foldr(UnionAll, tvars; init=Tuple{<:FwdConfig,<:Annotation{ft},Type{<:Annotation},tt...})
+    TT = foldr(UnionAll, tvars; init = Tuple{<:FwdConfig, <:Annotation{ft}, Type{<:Annotation}, tt...})
     fwd_sig = Tuple{typeof(EnzymeRules.forward), <:EnzymeRules.FwdConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
     return isapplicable(interp, forward, TT, sv, fwd_sig)
 end
@@ -12,7 +12,7 @@ end
 function has_rrule_from_sig(@nospecialize(interp::Core.Compiler.AbstractInterpreter),
     @nospecialize(TT::Type), sv::Core.Compiler.AbsIntState, partialedge::Bool=true)::Bool
     ft, tt, tvars = _annotate_tt(TT)
-    TT = foldr(UnionAll, tvars; init=Tuple{<:RevConfig,<:Annotation{ft},Type{<:Annotation},tt...})
+    TT = foldr(UnionAll, tvars; init = Tuple{<:RevConfig, <:Annotation{ft}, Type{<:Annotation}, tt...})
     rev_sig = Tuple{typeof(EnzymeRules.augmented_primal), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
     return isapplicable(interp, augmented_primal, TT, sv, rev_sig)
 end

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant