Add jl_nthfield in reverse mode

[ChrisRackauckas]

I was testing out main, things have really progressed well! Though Lux doesn't quite work yet, so I built an MWE:

using Enzyme

A = Any[2.0 3.0
     2.0 4.0]
function f(x::Array{Float64}, y::Array{Float64})
    y[1] = (A * x)[1]
    return nothing
end

x  = [2.0, 2.0]
bx = [0.0, 0.0]
y  = [0.0]
by = [1.0];

Enzyme.autodiff(Reverse, f, Duplicated(x, bx), Duplicated(y, by)); # Works fine!

using ComponentArrays, Lux, Random

rng = Random.default_rng()
dudt2 = Lux.Chain(x -> x.^3,
                  Lux.Dense(2, 50, tanh),
                  Lux.Dense(50, 2))
p, st = Lux.setup(rng, dudt2)

dudt2(x, p, st)[1]
y = zeros(2)

function f(x::Array{Float64}, y::Array{Float64})
    y .= dudt2(x, p, st)[1]
    return nothing
end

Enzyme.autodiff(Forward, f, Duplicated(x, bx), Duplicated(y, by)); # works

Enzyme.autodiff(Reverse, f, Duplicated(x, bx), Duplicated(y, by)); # fails

#=
warning: didn't implement memmove, using memcpy as fallback which can result in errors
ERROR: Enzyme: not yet implemented in reverse mode, jl_nthfield
=#

I don't know what causes jl_nthfield. When trying to make a simple MWE I found #644 because I thought it was the namedtuple indexing, but apparently that's a separate issue. I'll help make a simpler MWE if I can find out what actually causes jl_nthfield to be called.

[ChrisRackauckas]

The julia function is called jl_get_nth_field_checked. It comes from any ill-inferred getfield

@avik-pal do you know why the call is ill-inferred? Though getfield also doesn't work right now, but it seems like that's the last step if the inference is worked out.

[wsmoses]

We should at least give a partial backtrace so you know where in Lux the type instability/etc is coming from.

[avik-pal]

In this example, rewriting the function as

function f1(x::Array{Float64}, y::Array{Float64})
    y .= first(dudt2(x, p, st))
    return nothing
end

makes reverse mode work but the following way:

function f2(x::Array{Float64}, y::Array{Float64})
    z, st_ = dudt2(x, p, st)
    y .= z
    return nothing
end

Enzyme.autodiff(Reverse, f2, Duplicated(x, bx), Duplicated(y, by));

leads to

  %"'ip_phi7" = phi {} addrspace(10)* , !dbg !14


Stacktrace:
 [1] f2
   @ /mnt/software/Lux.jl/wip/enzyme.jl:49
 [2] f2
   @ /mnt/software/Lux.jl/wip/enzyme.jl:0

Stacktrace:
  [1] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing})
    @ Enzyme.Compiler /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:3735
  [2] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{Enzyme.API.CDIFFE_TYPE}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, width::Int64, additionalArg::Ptr{Nothing}, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{Bool}, augmented::Ptr{Nothing}, atomicAdd::Bool)
    @ Enzyme.API /mnt/julia/packages/Enzyme/DIkTv/src/api.jl:123
  [3] enzyme!(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(f2), Tuple{Vector{Float64}, Vector{Float64}}}}, mod::LLVM.Module, primalf::LLVM.Function, adjoint::GPUCompiler.FunctionSpec{typeof(f2), Tuple{Duplicated{Vector{Float64}}, Duplicated{Vector{Float64}}}}, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, dupClosure::Bool, wrap::Bool, modifiedBetween::Bool, returnPrimal::Bool, jlrules::Vector{String})
    @ Enzyme.Compiler /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:5187
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(f2), Tuple{Vector{Float64}, Vector{Float64}}}}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.ThreadSafeContext, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:6383
  [5] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(f2), Tuple{Vector{Float64}, Vector{Float64}}}}, ctx::Nothing)
    @ Enzyme.Compiler /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:6870
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(f2), Tuple{Vector{Float64}, Vector{Float64}}}})
    @ Enzyme.Compiler /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:6864
  [7] cached_compilation(job::GPUCompiler.CompilerJob, key::UInt64, specid::UInt64)
    @ Enzyme.Compiler /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:6908
  [8] #s906#169
    @ /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:6968 [inlined]
  [9] var"#s906#169"(F::Any, Fn::Any, DF::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, specid::Any, ReturnPrimal::Any, ShadowInit::Any, ::Any, #unused#::Type, f::Any, df::Any, #unused#::Type, tt::Any, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Any)
    @ Enzyme.Compiler ./none:0
 [10] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:602
 [11] thunk
    @ /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:7001 [inlined]
 [12] thunk (repeats 3 times)
    @ /mnt/julia/packages/Enzyme/DIkTv/src/compiler.jl:6994 [inlined]
 [13] autodiff
    @ /mnt/julia/packages/Enzyme/DIkTv/src/Enzyme.jl:211 [inlined]
 [14] autodiff(::EnzymeCore.ReverseMode, ::typeof(f2), ::Duplicated{Vector{Float64}}, ::Duplicated{Vector{Float64}})
    @ Enzyme /mnt/julia/packages/Enzyme/DIkTv/src/Enzyme.jl:248
 [15] top-level scope
    @ /mnt/software/Lux.jl/wip/enzyme.jl:56

The complete stacktrace is too large :(

[wsmoses]

Getting the full backtrace is unfortunately an issue on Julia itself (JuliaLang/julia#49056)

[wsmoses]

Also @avik-pal are you on the latest main, that is a different issue?

[avik-pal]

I updated to main but it leads to the same set of results - f1 works but f and f2 fail

For f we get:

ERROR: Enzyme execution failed.
Enzyme: not yet implemented in reverse mode, jl_nthfield
Stacktrace:
 [1] getindex
   @ ./tuple.jl:29
 [2] getindex
   @ ./tuple.jl:0

Stacktrace:
 [1] throwerr(cstr::Cstring)
   @ Enzyme.Compiler /mnt/julia/packages/Enzyme/wSCid/src/compiler.jl:2387

For f2, we get: https://gist.github.com/avik-pal/042080caa5be0b1417e83c2d0034ce1d

[wsmoses]

Fixed by #806