Assertion error with inlined function on recent Enzyme versions

[jgreener64]

I am on Julia 1.9.2 and Enzyme 628c9a4. The following used to work, e.g. on Enzyme 0.11.2.

It seems to be something to do with inlining, since adding @noinline before function g makes it work.

using Enzyme, StaticArrays, LinearAlgebra

struct HarmonicAngle
    k::Float64
    θ0::Float64
end

Base.zero(::HarmonicAngle) = HarmonicAngle(0.0, 0.0)

Base.:+(a1::HarmonicAngle, a2::HarmonicAngle) = HarmonicAngle(a1.k + a2.k, a1.θ0 + a2.θ0)

struct InteractionList3Atoms
    is::Vector{Int}
    js::Vector{Int}
    ks::Vector{Int}
    inters::Vector{HarmonicAngle}
    types::Vector{String}
end

function Base.:+(il1::InteractionList3Atoms, il2::InteractionList3Atoms)
    return InteractionList3Atoms(
        il1.is,
        il1.js,
        il1.ks,
        il1.inters .+ il2.inters,
        il1.types,
    )
end

function g(a, coords_i, coords_j, coords_k)
    return (a.k / 2) * (2.0 - a.θ0) ^ 2
end

function f!(s, inter_list, coords)
    si = 0.0
    for (i, j, k, b) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.inters)
        si += g(b, coords[i], coords[j], coords[k])
    end
    s[1] = si
    return nothing
end

inter_list = InteractionList3Atoms(
    [1, 3],
    [2, 2],
    [3, 1],
    [HarmonicAngle(1.0, 0.1), HarmonicAngle(1.0, 0.1)],
    ["", ""],
)
d_inter_list = InteractionList3Atoms(
    [0, 0],
    [0, 0],
    [0, 0],
    [HarmonicAngle(0.0, 0.0), HarmonicAngle(0.0, 0.0)],
    ["", ""],
)
coords   = [SVector(1.0, 2.0, 3.0), SVector(1.1, 2.1, 3.1), SVector(1.2, 2.2, 3.2)]
d_coords = [SVector(0.0, 0.0, 0.0), SVector(0.0, 0.0, 0.0), SVector(0.0, 0.0, 0.0)]

s, ds = [0.0], [1.0]
autodiff(
    Reverse,
    f!,
    Const,
    Duplicated(s, ds),
    Duplicated(inter_list, d_inter_list),
    Duplicated(coords, d_coords),
)

f = LLVM.Function("julia_f__868_inner.1")
(gty, inst, v) = (LLVM.IntegerType[LLVM.IntegerType(i64), LLVM.IntegerType(i64)], LLVM.PHIInst(%.sroa.675.0723.in = phi double addrspace(13)* [ %.sroa.4219.0..sroa_cast, %idxend.i.lr.ph ], [ %.sroa.4148.0..sroa_idx149, %L118.i ]), LLVM.BitCastInst(%.sroa.4219.0..sroa_cast = bitcast i8 addrspace(13)* %.sroa.4219.0..sroa_idx to double addrspace(13)*, !dbg !80))
f = LLVM.Function("julia_f__1550_inner.1")
(gty, inst, v) = (LLVM.IntegerType[LLVM.IntegerType(i64), LLVM.IntegerType(i64)], LLVM.PHIInst(%.sroa.675.0723.in = phi double addrspace(13)* [ %.sroa.4219.0..sroa_cast, %idxend.i.lr.ph ], [ %.sroa.4148.0..sroa_idx149, %L118.i ]), LLVM.BitCastInst(%.sroa.4219.0..sroa_cast = bitcast i8 addrspace(13)* %.sroa.4219.0..sroa_idx to double addrspace(13)*, !dbg !80))
ERROR: LoadError: AssertionError: false
Stacktrace:
  [1] nodecayed_phis!(mod::LLVM.Module)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler/optimize.jl:198
  [2] optimize!
    @ ~/.julia/dev/Enzyme/src/compiler/optimize.jl:1097 [inlined]
  [3] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:9145
  [4] codegen
    @ ~/.julia/dev/Enzyme/src/compiler.jl:8766 [inlined]
  [5] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:9713
  [6] _thunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:9713 [inlined]
  [7] cached_compilation
    @ ~/.julia/dev/Enzyme/src/compiler.jl:9747 [inlined]
  [8] (::Enzyme.Compiler.var"#475#476"{DataType, UnionAll, DataType, Enzyme.API.CDerivativeMode, NTuple{4, Bool}, Int64, Bool, Bool, UInt64, DataType})(ctx::LLVM.Context)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:9810
  [9] JuliaContext(f::Enzyme.Compiler.var"#475#476"{DataType, UnionAll, DataType, Enzyme.API.CDerivativeMode, NTuple{4, Bool}, Int64, Bool, Bool, UInt64, DataType})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:47
 [10] #s292#474
    @ ~/.julia/dev/Enzyme/src/compiler.jl:9765 [inlined]
 [11] var"#s292#474"(FA::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, ReturnPrimal::Any, ShadowInit::Any, World::Any, ABI::Any, ::Any, #unused#::Type, #unused#::Type, #unused#::Type, tt::Any, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Any)
    @ Enzyme.Compiler ./none:0
 [12] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:602
 [13] autodiff(::ReverseMode{false, FFIABI}, ::Const{typeof(f!)}, ::Type{Const}, ::Duplicated{Vector{Float64}}, ::Vararg{Any})
    @ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:207
 [14] autodiff(::ReverseMode{false, FFIABI}, ::typeof(f!), ::Type, ::Duplicated{Vector{Float64}}, ::Vararg{Any})
    @ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:222
 [15] top-level scope
    @ ~/dms/molly_dev/enzyme_err27.jl:66

[jgreener64]

Still erroring on main (c29e611), any chance of getting this one looked at? It causes the Molly tests to fail.

[wsmoses]

using Enzyme

struct HarmonicAngle
    k::Float64
    θ0::Float64
end

function g(a, coords_i)
    return (a.k) * a.θ0
end

function f!(inter_list, coords, inters)
    si = 0.0
    for (i, b) in zip(inter_list, inters)
        si += g(b, coords[i])
    end
    return si
end

inters = [HarmonicAngle(1.0, 0.1), HarmonicAngle(1.0, 0.1)]
inter_list = [1, 3]
dinters = [HarmonicAngle(0.0, 0.0), HarmonicAngle(0.0, 0.0)]
coords   = [(1.0, 2.0, 3.0), (1.1, 2.1, 3.1), (1.2, 2.2, 3.2)]
d_coords = [(0.0, 0.0, 0.0), (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)]

autodiff(
    Reverse,
    f!,
    Const,
    Const(inter_list),
    Duplicated(coords, d_coords),
    Duplicated(inters, dinters),
)