Custom rule not detected if defined after call to autodiff

[gaurav-arya]

MWE on main:

julia> using Enzyme

julia> import .EnzymeRules: forward

julia> f(x) = x^2
f (generic function with 1 method)

julia> autodiff(Forward, f, Duplicated(1.0, 1.0))
(2.0,)

julia> forward(::Const{typeof(f)}, ::Type{<:DuplicatedNoNeed}, x::Duplicated) = 10+2*x.val*x.dval
forward (generic function with 1 method)

julia> forward(func::Const{typeof(f)}, ::Type{<:Duplicated}, x::Duplicated) = Duplicated(func.val(x.val), 10+2*x.val*x.dval)
forward (generic function with 2 methods)

julia> autodiff(Forward, f, Duplicated(1.0, 1.0)) # answer unchanged
(2.0,)

julia> f(x) = x^2 # redefine f and we get the correct answer
f (generic function with 1 method)

julia> autodiff(Forward, f, Duplicated(1.0, 1.0))
(12.0,)

The reason for the behaviour might be this line:

Enzyme.jl/src/Enzyme.jl

Line 313 in 939f9b4

world = GPUCompiler.get_world(Core.Typeof(f.val), tt)

[vchuravy]

We have no invalidation edge from forward/reverse/augmented_forward to the result of our compilation.
So the code is not invalidated when a new rule is being defined.

[gaurav-arya]

Is this something that you think is somewhat easy to fix? Happy to try to make a PR if you think so. The workaround is to redefine the original function each time the custom rule is modified, but it's easy to forget.

[vchuravy]

Hm not trivial, we need to add a fictitious edge to a potentially not yet existing method.

Somewhere in https://github.com/EnzymeAD/Enzyme.jl/blob/939f9b4086d62b07eca8107db9523a2f8fe043d3/src/compiler/interpreter.jl or

Enzyme.jl/src/compiler.jl

Line 7510 in 939f9b4

has_custom_rule = EnzymeRules.has_frule_from_sig(specTypes; world)

Maybe @aviatesk has some ideas. I suspect we will need to intercept abstract_gf_call or something similar and add inference edge... But we don't necessarily know which version of the rule we will call due to activity analysis.

That we only know around

Enzyme.jl/src/compiler.jl

Line 3699 in 939f9b4

llvmf = nested_codegen!(mode, mod, EnzymeRules.reverse, rev_TT, world)

We should definitely try to add edges from there... But I am unsure how to add a speculative edge for an undefined method.

[gaurav-arya]

speculative edge for an undefined method

Definitely a bit out of my depth here, but: would backedges be easier if forward etc. had a default fallback of nothing? (So rule detection would change from method detection to checking if the method returns nothing). That's the situation for frule and friends: https://github.com/JuliaDiff/ChainRulesCore.jl/blob/79ba4ef03afdf5715b6fa0294e5accfe4e95c79b/src/rules.jl#L61

[gaurav-arya]

Some relevant code from Zygote:

https://github.com/FluxML/Zygote.jl/blob/dc16b2e35cdfa9dda62f636e95271490548cc574/src/compiler/chainrules.jl#L18
https://github.com/FluxML/Zygote.jl/blob/dc16b2e35cdfa9dda62f636e95271490548cc574/src/compiler/interface2.jl#L20

[aviatesk]

I believe I can fix this. How can we access the MethodInstance object representing the caller inside GPUCompiler.codegen? We need it to add appropriate backedges to it.

[wsmoses]

Enzyme.jl/src/compiler.jl

Line 7503 in 939f9b4

for (mi, k) in meta.compiled

this contains every compiled methodinstance [at least that we see]