Julia 1.12: segfault from @cfunction ABI mismatch (gcstack_arg) when custom rules execute nested-compiled host code

[vchuravy]

Summary

When a custom rule executes host code that was compiled into the differentiated module by nested_codegen! (src/rules/customrules.jl:1704), any @cfunction in that code produces a trampoline whose ABI disagrees with what Julia 1.12's runtime cfunction converter hands back, leading to a segfault. Found while debugging the long-red CUDA CI "Reverse Kernel" testset; this is the next failure after the jl-inst-simplify fix (EnzymeAD/Enzyme#2919).

Reproducer

Julia 1.12.6, Enzyme main + EnzymeAD/Enzyme#2919 (via deps/build_local.jl), CUDA v6:

using CUDA, Enzyme

function square_kernel!(x)
    i = threadIdx().x
    x[i] *= x[i]
    return nothing
end

function square!(x)
    @cuda blocks = 1 threads = length(x) square_kernel!(x)
    return nothing
end

A = CuArray(collect(Float32, 1:64))
dA = CUDA.ones(64)
Enzyme.autodiff(Reverse, square!, Duplicated(A, dA))  # SIGSEGV

CUDACore's augmented_primal calls Enzyme.tape_type at runtime; the nested compile's optimize! invokes LLVM.jl custom passes, and the process crashes inside module_callback (LLVM.jl newpm.jl:140). The identical tape_type call works at top level — only the copy of run! compiled into the differentiated module misbehaves.

Root cause

1. nested_codegen! compiles LLVM.jl's run! into the module using GPUCompiler's CodegenParams(gcstack_arg = false) (GPUCompiler jlgen.jl:791), so specsig functions there do not take pgcstack as an argument.
2. Julia 1.12 lowers the @cfunction(module_callback, ...) inside run! to the abi-converter pattern (gen_cfun_wrapper): a jlcapi_module_callback_* trampoline, cache slots, and a jl_get_abi_converter guard. Under gcstack_arg=false the trampoline's indirect converter call passes only (ptr, ptr) — no pgcstack (verified in the post-AD module dump; cfuncdata flags have specsig=1).
3. At runtime jl_get_abi_converter (runtime_ccall.cpp:445-449) finds module_callback's code instance compiled by the stock JIT (gcstack_arg=true, pgcstack in swiftself/r13), sees an exact declrt/sigt match, and returns the raw specsig pointer.
4. The trampoline calls it without pgcstack. gdb at the fault: rdi = LLVM Module, rsi = thunk, and r13 = leftover thunk pointer — the callee treats r13 as pgcstack, pushes a GC frame into the CustomPassState object, and segfaults.

Not Enzyme-specific in principle: any GPUCompiler-compiled host code that executes an @cfunction on 1.12 has this mismatch. Enzyme is the main consumer that actually runs nested-compiled host code, because rules pull their whole call graph (here: the entire Enzyme.jl compiler plus LLVM.jl's pass machinery — the post-AD module grows from 647 to 3615 functions) into the differentiated module.

Ruled out

• GC rooting inside rule-executed code is sound (WeakRef + GC.@preserve probe inside a custom rule passes).
• The registered callback pointer in the emitted IR is the correct in-module trampoline; the bad pointer comes from the runtime converter at world-guard refresh.

Fix options

• (a) Julia: encode gcstack_arg in the cfuncdata flags so jl_get_abi_converter refuses the raw-specsig fast path when the caller's ABI differs from the native JIT's. The runtime currently assumes all callers share the stock ABI.
• (b) GPUCompiler gcstack_arg=true for host-executable targets: tested, not viable — Enzyme.jl fails its own precompilation with "AssertionError: Swiftself attribute coming from differentiable context is not supported" (lower_convention).
• (c) Enzyme.jl: stop compiling rule-reachable compiler code into the differentiated module and call rules through a dynamic-dispatch boundary. This would also remove the exposure that made the jl-inst-simplify miscompile (nondeterminstic segfault during compilation #2919) reachable: unsound-but-latent pass behavior currently gets applied to the whole Enzyme/GPUCompiler/LLVM.jl compiler stack on every rule compile.

Happy to file the corresponding JuliaLang/julia issue for (a) if that is the preferred direction.

🤖 Generated with Claude Code