Nested differentiation through a custom rule whose body calls a function returning a mixed pointer/isbits struct: CallingConventionMismatchError in classify_arguments (Julia <= 1.11)
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
Under nested (forward-over-forward) differentiation, when a custom rule's body calls a non-inlined function that returns an immutable struct with both pointer and isbits fields, Enzyme's module type-classification fails while compiling the OUTER differentiation:
```
ERROR: LoadError: CallingConventionMismatchError: Enzyme hit an internal error trying to parse the julia calling convention definition for:
makecache(::Int64, ::Int64)
@ Main ~/claude_tmp/claude-25022/-home-crackauc-sandbox-tmp-20260710-063942-23305/7bf73f89-ee49-42ce-86e5-b5335e721373/scratchpad/mwe4_callconv.jl:26
Hint: catch this exception as `err` and call `code_typed(err)` to inspect the errornous code.
expectLen != length(parameters(f))
void ({ {} addrspace(10)*, {} addrspace(10)*, i8 }*, [2 x {} addrspace(10)*]*, i64, i64)
expectLen=3
has_swiftself=false
has_sret=true
has_returnroots=true
source_sig=Tuple{typeof(makecache), Int64, Int64}
parmsRemoved=UInt64[0x0000000000000001]
Stacktrace:
[1] classify_arguments(source_sig::Type, codegen_ft::LLVM.FunctionType, has_sret::Bool, has_returnroots::Bool, has_swiftself::Bool, parmsRemoved::Vector{UInt64}, mi::Core.MethodInstance, world::UInt64)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/CCx0o/src/typeutils/jltypes.jl:339
[2] set_module_types!(interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, mod::LLVM.Module, primalf::LLVM.Function, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, edges::Vector{Any}, run_enzyme::Bool, mode::Enzyme.API.CDerivativeMode)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/CCx0o/src/compiler.jl:1163
[3] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/CCx0o/src/compiler.jl:5532
[4] #compile#204
@ ~/.julia/packages/GPUCompiler/BSi1T/src/driver.jl:67 [inlined]
[5] compile
@ ~/.julia/packages/GPUCompiler/BSi1T/src/driver.jl:55 [inlined]
[6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, postopt::Bool)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/CCx0o/src/compiler.jl:7011
[7] _thunk
@ ~/.julia/packages/Enzyme/CCx0o/src/compiler.jl:7009 [inlined]
[8] cached_compilation
@ ~/.julia/packages/Enzyme/CCx0o/src/compiler.jl:7079 [inlined]
[9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{<:Annotation}, A::Type{<:Annotation}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width:
```
First-order differentiation through the same rule works — the parse only happens when the inner rule body becomes part of the outer-differentiated module (`set_module_types!` -> `classify_arguments`, sret + returnroots, `expectLen=2`).
## Reproducer
```julia
# MWE: under nested differentiation, Enzyme's module type-classification
# fails to parse a function returning an immutable struct with mixed
# pointer/isbits fields when it is reachable from a custom rule body:
#
# CallingConventionMismatchError: ... expectLen != length(parameters(f))
# has_sret=true has_returnroots=true
# in classify_arguments at src/typeutils/jltypes.jl, called from
# set_module_types! during compilation of the OUTER differentiation.
#
# First-order differentiation through the same rule works; only the nested
# (forward-over-forward) case, where the inner rule body becomes part of the
# outer-differentiated module, triggers the parse. Marking the helper
# @inline avoids it (no standalone function left in the module).
#
# Observed with Enzyme v0.13.181 on Julia 1.11.9 and 1.10.11.
using Enzyme
using Enzyme.EnzymeCore: EnzymeRules, Annotation, Const, Duplicated
struct Cache
du::Matrix{Float64}
extra::Vector{Float64}
flag::Bool
end
@noinline makecache(n::Int, m::Int) = Cache(zeros(n, m), zeros(n * m), true)
getbuf(c::Cache, t) = c.du
function EnzymeRules.forward(
config::EnzymeRules.FwdConfig,
::Const{typeof(getbuf)}, ::Type{<:Annotation},
c::Annotation{Cache}, t::Annotation
)
du = getbuf(c.val, t.val)
if !EnzymeRules.needs_shadow(config)
return EnzymeRules.needs_primal(config) ? du : nothing
end
sc = makecache(size(du)...) # struct-returning call inside the rule body
return EnzymeRules.needs_primal(config) ? Duplicated(du, sc.du) : sc.du
end
function loss(t, c)
buf = getbuf(c, t)
buf .= t
return sum(abs2, buf)
end
c = Cache(zeros(5, 3), zeros(15), true)
dloss(t, c) = only(Enzyme.autodiff(Forward, loss, Duplicated(t, 1.0), Const(c)))
println("first order: ", dloss(0.3, c)) # works
d2 = only(
Enzyme.autodiff(
set_runtime_activity(Forward), Const(dloss), Duplicated(0.3, 1.0), Const(c)
)
)
println("second order: ", d2)
```
## Notes
- Enzyme v0.13.181 / EnzymeCore v0.8.21; fails on Julia 1.11.9 and 1.10.11. On Julia 1.12+ this particular parse does not trigger (there the nested case instead hits #3315).
- Workaround: marking the struct-returning helper `@inline` removes the standalone function from the module and everything compiles and produces correct second derivatives (this is what SciML/PreallocationTools.jl#183 now does for its `makeshadow` helpers).
- Found while removing the #3310 workaround after v0.13.181 (thanks for the quick fixes of #3310/#3311 — both confirmed against their MWEs).
Contributor guide
Assessment
This issue has not been assessed yet.