Second order derivative error

[jgreener64]

I am on Enzyme main (b63f36f) and Julia 1.10.6. The following used to work.

using Enzyme
f_nest(x) = 2 * x^4
deriv(f, x) = first(first(autodiff(Reverse, f, Active(x))))
f′(x) = deriv(f_nest, x)
f′′(x) = deriv(f′, x)
f′(2.0) # Works, 64
f′′(2.0) # Errors, should be 96

Unreachable reached at 0x7f7d85ca034f

[181735] signal (4.2): Illegal instruction
in expression starting at REPL[7]:1
unknown function (ip: 0x7f7d85ca034f)
macro expansion at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:8369 [inlined]
enzyme_call at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:7932 [inlined]
PrimalErrorThunk at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:7689 [inlined]
autodiff at /home/jgreener/.julia/dev/Enzyme/src/Enzyme.jl:491 [inlined]
autodiff at /home/jgreener/.julia/dev/Enzyme/src/Enzyme.jl:537
unknown function (ip: 0x7f7d85ca1216)
_jl_invoke at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:3077
autodiff at /home/jgreener/.julia/dev/Enzyme/src/Enzyme.jl:504 [inlined]
deriv at ./REPL[3]:1
unknown function (ip: 0x7f7d85ca11ad)
f′′ at ./REPL[5]:1
unknown function (ip: 0x7f7d85ca0979)
_jl_invoke at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_apply at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_call at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/interpreter.c:126
eval_value at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/interpreter.c:223
eval_stmt_value at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/interpreter.c:174 [inlined]
eval_body at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/interpreter.c:617
eval_body at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/interpreter.c:544
jl_interpret_toplevel_thunk at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/interpreter.c:775
jl_toplevel_eval_flex at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/toplevel.c:934
jl_toplevel_eval_flex at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/toplevel.c:877
ijl_toplevel_eval_in at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/toplevel.c:985
eval at ./boot.jl:385 [inlined]
eval_user_input at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/usr/share/julia/stdlib/v1.10/REPL/src/REPL.jl:150
repl_backend_loop at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/usr/share/julia/stdlib/v1.10/REPL/src/REPL.jl:246
#start_repl_backend#46 at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/usr/share/julia/stdlib/v1.10/REPL/src/REPL.jl:231
start_repl_backend at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/usr/share/julia/stdlib/v1.10/REPL/src/REPL.jl:228
_jl_invoke at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:3077
#run_repl#59 at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/usr/share/julia/stdlib/v1.10/REPL/src/REPL.jl:389
run_repl at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/usr/share/julia/stdlib/v1.10/REPL/src/REPL.jl:375
jfptr_run_repl_91949.1 at /home/jgreener/soft/julia/julia-1.10.6/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:3077
#1013 at ./client.jl:437
jfptr_YY.1013_82918.1 at /home/jgreener/soft/julia/julia-1.10.6/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_apply at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
jl_f__call_latest at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/builtins.c:812
#invokelatest#2 at ./essentials.jl:892 [inlined]
invokelatest at ./essentials.jl:889 [inlined]
run_main_repl at ./client.jl:421
exec_options at ./client.jl:338
_start at ./client.jl:557
jfptr__start_82944.1 at /home/jgreener/soft/julia/julia-1.10.6/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_apply at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
true_main at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/jlapi.c:582
jl_repl_entrypoint at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/src/jlapi.c:731
main at /cache/build/builder-amdci5-5/julialang/julia-release-1-dot-10/cli/loader_exe.c:58
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x4010b8)
Allocations: 25179959 (Pool: 25131195; Big: 48764); GC: 37
Illegal instruction (core dumped)

[wsmoses]

using Enzyme

f_nest(x) = x*x*x*x

function df1(x)
	first(first(autodiff_deferred(Reverse, Const(f_nest), Active, Active(x))))
end

function df1_nd(x)
	first(first(autodiff(Reverse, Const(f_nest), Active(x))))
end


function df2(x)
	first(first(autodiff(Reverse, df1, Active(x))))
end

function df2_nd(x)
	first(first(autodiff(Reverse, df1_nd, Active(x))))
end

@code_typed df1(2.0)
@code_typed df2(2.0)

@code_typed df1_nd(2.0)
@code_typed df2_nd(2.0)

[wsmoses]

so I think it's an issue with the deferred auto upgrade, investigating

[wsmoses]

@aviatesk I think I need some absint help here: x/ref https://github.com/EnzymeAD/Enzyme.jl/pull/2115/files

So basically the problem at hand is that we call autodiff of autodiff. the abstract interpreter plugin we added should change autodiff -> autodiff_deferred here:

Enzyme.jl/src/compiler/interpreter.jl

Line 862 in 1292a53

if interp.deferred_lower && f === Enzyme.autodiff && length(argtypes) >= 4

the outer autodiff asks what the return type of its inner differentiated function is (e.g. here it asks the return type of e.g. df1_nd). However, it gets told the result is a Union{}, here:

Enzyme.jl/src/compiler.jl

Line 3129 in 1292a53

rt = Union{}

This seems to imply the overload doesn't happen, yet on all runs (e.g. see the debug logs I added to the autodiff interpreter overload linked above), it definiteyl changes autodiff->autodiff_deferred, incuding in the type computation.

The way we do this (reproduced below) has something that begs a question.

    if interp.deferred_lower && f === Enzyme.autodiff && length(argtypes) >= 4
        if widenconst(argtypes[2]) <: Enzyme.Mode &&
           widenconst(argtypes[3]) <: Enzyme.Annotation &&
           widenconst(argtypes[4]) <: Type{<:Enzyme.Annotation}
            arginfo2 = ArgInfo(
                fargs isa Nothing ? nothing :
                [:(Enzyme.autodiff_deferred), fargs[2:end]...],
                [Core.Const(Enzyme.autodiff_deferred), argtypes[2:end]...],
            )
            ccall(:jl_, Any, (Any,), (f, argtypes, arginfo2, si, typeof(sv)))
            # throw(AssertionError("wat"))
            return abstract_call_known(
                interp,
                Enzyme.autodiff_deferred,
                arginfo2,
                si,
                sv,
                max_methods,
            )
        end
    end

Specifically, we modify arginfo and f, but we left si/sv alone. Is there something we need to do to change these, and if so what. Similarly could this be the cause of the issue?