Segfault in `lower-jit` (`rewriteKernelCallABI`) lowering a KernelAbstractions launch to `gpu.LaunchFuncOp`

This is a reduced version of https://github.com/NumericalEarth/NumericalEarth.jl/issues/407 and a temporary solution in Oceananigans is https://github.com/NumericalEarth/NumericalEarth.jl/pull/408. Compiling a KernelAbstractions `launch!` on the **CUDA** backend segfaults inside the `lower-jit` pass, in `rewriteKernelCallABI`, while lowering an `enzymexla.jit_call` to a `gpu.LaunchFuncOp` — but only when the launched kernel has **no writeable device-array output** (the resulting `gpu.LaunchFuncOp` produces no results). The same program compiles fine with `raise=true` (no `jit_call` is produced) and on the CPU backend. ## MWE (pure Oceananigans, ~10 lines) ```julia using Oceananigans, Reactant, CUDA using Oceananigans.Architectures: ReactantState, architecture using Oceananigans.Utils: launch! using KernelAbstractions: @kernel, @index Reactant.set_default_backend("gpu") @kernel function _empty!(grid) # receives the grid, writes nothing i, j = @index(Global, NTuple) end grid = RectilinearGrid(ReactantState(); size = (), topology = (Flat, Flat, Flat)) Reactant.@jit launch!(architecture(grid), grid, (1, 1), _empty!, grid) # ^ segfaults (signal 11) inside lower-jit / rewriteKernelCallABI ``` Run in a **fresh** session, un-raised (the crash is order/state dependent — see Notes). An Oceananigans coupled-model regridder kernel whose output fields happen to be `nothing` (so it writes nothing) is the first launch compiled when building the model, and it crashes here. ## Where it crashes `lower-jit` walks each `enzymexla.jit_call` and, lowering one to a `gpu.LaunchFuncOp`, segfaults. Backtrace (crash frame first, down to `@jit`): ``` signal 11 (Segmentation fault) [libReactantExtra.so] # crash: walk callback lowering a JITCallOp to gpu::LaunchFuncOp inside rewriteKernelCallABI llvm::function_ref<...>::callback_fn<mlir::detail::walk<..., rewriteKernelCallABI(...)::$_2, mlir::gpu::LaunchFuncOp, void>>(long, mlir::Operation*) mlir::detail::walk<mlir::ForwardIterator>(...) rewriteKernelCallABI(mlir::ModuleOp&, mlir::Location, ..., mlir::enzymexla::JITCallOp, ...) CompileCall(..., mlir::enzymexla::JITCallOp, ...) (anonymous namespace)::LowerJITPass::runOnOperation()::{lambda(mlir::enzymexla::JITCallOp)#1} mlir::detail::walk<mlir::ForwardIterator>(...) (anonymous namespace)::LowerJITPass::runOnOperation() mlir::PassManager::run(mlir::Operation*) mlirPassManagerRunOnOp run_pass_pipeline! Reactant/src/compiler/OptimizationPasses.jl:208 (pipeline key = "all") compile_mlir! Reactant/src/compiler/Compiler.jl:238 @jit ``` The pass, as invoked on an sm_75 (T4): ``` lower-jit{backend=cuda compileLaunch=true cubinChip=sm_75 cubinFeatures=+ptx90 cubinFormat=bin cubinTriple=nvptx64-nvidia-cuda indexBitWidth=32 jit=true run_init=true toolkitPath=… …} ``` ## What is / isn't required to trigger it (bisected) Each row is a separately verified run (fresh GPU session, un-raised, first compile). The grid is passed as a kernel argument in every row, so it is not the discriminator: | Kernel body | Launch worksize | Writes a device-array output? | Result | |-------------|-----------------|-------------------------------|--------| | `a[i,j,1] = b[i,j,1]` (two `Field`s, no grid arg) | dynamic `(1,1)` | yes | compiles, runs | | `out[i,j,1] = size(grid, 3)` (`Field` + grid) | static `KernelParameters(1:1,1:1)` | yes | compiles, runs | | empty body (grid only) | static `KernelParameters(1:1,1:1)` | **no** | **segfault** | | empty body (grid only) | dynamic `(1,1)` | **no** | **segfault** | So the discriminator is **whether the kernel writes a real output field**: a launch whose `gpu.LaunchFuncOp` has *no results* crashes `rewriteKernelCallABI`; a launch that writes an output compiles fine. The launch worksize (static `KernelParameters` vs dynamic tuple) does not matter — both output-less variants crash. A hand-written KA kernel with an `OffsetArray`-over-device-array argument (no Oceananigans) fails *earlier*, in GPU codegen (dynamic `getindex`), so it can't reach `lower-jit`; the Oceananigans `Field` / `launch!` machinery is what gets an output-less kernel to the crashing pass. ## Notes toward the cause - **Order/state dependent.** Reproduces only when the **un-raised** compile is the *first* touch of that function in the session. If a `@jit raise=true` compile of the *same* call runs first, the subsequent un-raised `lower-jit` lowering then **succeeds** (emits a `gpu.binary`, no crash). Strong evidence the trigger is in-session state, not the IR text alone. - **A textual round-trip does not reproduce it.** Dumping the module right before `lower-jit` (`Reactant.@code_hlo optimize=:before_jit raise=false …`) and re-running only `lower-jit` on the parsed textual IR does not segfault — corroborating the in-session-state dependence. - **`raise=true` is the current workaround** — it avoids emitting the `jit_call` entirely, so nothing reaches `rewriteKernelCallABI`. - **`DEBUG_KERNEL[] = true` does not make it catchable.** That injects `debug=true cuResultHandlerPtr=<ptr>` into `lower-jit`, but that handler is for *runtime* CUDA results; this fault is at *compile time* inside the rewrite, before any kernel runs, so a Julia `try/catch` never sees it. ## Environment Reactant 0.2.268, Enzyme 0.13.173, Oceananigans 0.110.6, CUDA.jl 6.2.0, Julia 1.12.6, NVIDIA Tesla T4 (sm_75), CUDA runtime 13.1, cuDNN 9.14, XLA/PJRT CUDA. @dkytezab @Pangoraw