within_autodiff weirdly sticky

[vchuravy]

julia> using Enzyme

julia> function error_not_within_autodiff()
           !Enzyme.within_autodiff() && error("Not within autodiff")
           return nothing
       end
error_not_within_autodiff (generic function with 1 method)

julia> fwd_within_autodiff(Mode, RT) = Enzyme.autodiff(Mode, error_not_within_autodiff, RT)
fwd_within_autodiff (generic function with 1 method)

julia> error_not_within_autodiff()
ERROR: Not within autodiff
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] error_not_within_autodiff()
   @ Main ./REPL[5]:2
 [3] top-level scope
   @ REPL[7]:1
 [4] top-level scope
   @ none:1

julia> fwd_within_autodiff(Forward, Const)
()

julia> error_not_within_autodiff()

julia> Enzyme.within_autodiff()
false

I am extremely confused why is the 2nd call not throw an error here. Only happens if I call fwd_within_autodiff in between. cc @wsmoses this is in isolation from Reactant

Originally posted by @avik-pal in EnzymeAD/Reactant.jl#490 (comment)

[vchuravy]

It get's weirder... Using --trace-precompile I see no additional compilation.

[vchuravy]

Oof that is a nasty bug

julia> m = methods(error_not_within_autodiff) |> only

julia> m.specializations.cache.owner |> typeof
Enzyme.Compiler.EnzymeCacheToken

julia> m.specializations.cache.next.owner |> typeof
Enzyme.Compiler.EnzymeCacheToken

julia> m.specializations.cache.next.next.owner |> typeof
ERROR: UndefRefError: access to undefined reference
Stacktrace:
 [1] getproperty(x::Core.CodeInstance, f::Symbol)
   @ Base ./Base.jl:49
 [2] top-level scope
   @ REPL[14]:1

So the CodeInstance cache has two entries... both are owned by Enzyme, and not Base.

Somehow Base instead of yelling or complaining is using the Enzyme compiled entry directly.

[vchuravy]

So:

julia> m = methods(error_not_within_autodiff) |> only
error_not_within_autodiff()
     @ Main ~/src/Enzyme/within_ad.jl:3

julia> m.specializations
svec()

julia> error_not_within_autodiff()
ERROR: Not within autodiff
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] error_not_within_autodiff()
   @ Main ~/src/Enzyme/within_ad.jl:4
 [3] top-level scope
   @ REPL[3]:1

julia> m.specializations
MethodInstance for error_not_within_autodiff()

julia> m.specializations.cache.owner |> typeof
Nothing

So far so good.

julia> fwd_within_autodiff(Forward, Const)
()

julia> m.specializations.cache.owner |> typeof
Enzyme.Compiler.EnzymeCacheToken

julia> m.specializations.cache.next.owner |> typeof
Enzyme.Compiler.EnzymeCacheToken

julia> m.specializations.cache.next.next.owner |> typeof
Nothing

Okay all the entries are there (although why are there two Enzyme entries that is weird).

But who ever chooses which method to call is not filtering upon owner.

[vchuravy]

Should be fixed by JuliaLang/julia#59558

[vchuravy]

Will be fixed by Julia 1.11.8