`raise=true` fails to lower a KA kernel that takes a struct argument containing a traced **scalar** (`raise_triton_custom_call` operand mismatch) on Julia 1.12 but not Julia 1.11
## Summary
On **Julia 1.12**, `@compile raise=true` of a KernelAbstractions kernel launched
on a `ReactantState` grid fails when one of its arguments is a **struct that
contains a traced scalar** (`TracedRNumber`) — concretely, an Oceananigans
`Clock` whose `time::TracedRNumber{Float64}` is passed by value into the kernel.
The same code compiles and runs on **Julia 1.11**.
On Reactant `main` / 0.2.268 the failure is a cryptic raise-pass crash:
```
error: 'llvm.call' op operand type mismatch for operand 5:
'!llvm.ptr' != '!llvm.struct<(f64, array<1 x array<2 x f64>>)>'
caused by: "failed to run pass manager on module"
MLIR pass pipeline "raise_triton_custom_call" failed
```
With **[PR #2782](https://github.com/EnzymeAD/Reactant.jl/pull/2782)** (`pb/ll-inline-roots`, "Use LLVM type for inline roots
detection") the cryptic crash is **replaced by a clear, correct diagnostic** —
but it still does not compile:
```
GPU kernel argument of type
@NamedTuple{time::TracedRNumber{Float64}, last_Δt::TracedRNumber{Float64},
last_stage_Δt::TracedRNumber{Float64}, iteration::TracedRNumber{Int64}, stage::Int64}
contains a non-concrete traced value at field: time
```
So PR #2782 is a strict improvement (it pinpoints the offending argument/field)
but the underlying capability is still missing: **a traced scalar nested inside
a by-value struct kernel argument is not hoisted to a custom-call root.**
## Environment
- Julia **1.12.6** (fails) vs **1.11.9** (succeeds) — only the Julia version differs.
- Reactant: reproduced on **`main`** (HEAD `667be290…`, 2026-06-25) and **0.2.268**;
PR #2782 (`pb/ll-inline-roots`) turns the raise crash into the diagnostic above.
- CPU backend (`set_default_backend("cpu")`), macOS aarch64.
- No Enzyme / no autodiff — a pure forward raise.
## Root cause
The failing kernel `_apply_air_land_radiative_fluxes!` takes, among other args, a
`clock` (`@NamedTuple{time::TracedRNumber{Float64}, …}`) and an
`interface_radiative_flux` (a `NamedTuple` of **Fields**, i.e. traced arrays).
Key observation, from the *same* kernel's lowered `llvm.call` operands:
- the **NamedTuple-of-Fields** (traced **arrays**) lowers fine — its array data
pointers are hoisted to roots (`array<… x struct<(array<1 x ptr<1>>, …)>>`);
- the **clock** (traced **scalars**) is the argument that breaks.
So Reactant already hoists traced *arrays* nested in struct arguments, but not
traced *scalars*. Passing the bare `clock.time` (a top-level `TracedRNumber`,
which **is** handled as a root) instead of the whole `clock` makes the kernel
compile and run on both 1.11 and 1.12 — confirming the scalar-in-struct path is
the gap.
## Potential fix
1. Extend inline-roots handling so traced **scalars**
nested in a struct argument are hoisted to custom-call roots and the struct is
rematerialized in the kernel prologue — symmetric to the existing handling of
traced arrays. This is the area PR #2782 reworks; it currently *detects* the
nested scalar and errors rather than hoisting it.
## Reproducer
Deps: NumericalEarth, Oceananigans, Reactant, CUDA — no Enzyme. Compiles+runs on
Julia 1.11, fails on 1.12 with the error above. `apply_air_land_radiative_fluxes!`
launches the offending `_apply_air_land_radiative_fluxes!` kernel, which takes the
`clock` by value.
```julia
using NumericalEarth
using Oceananigans
using Oceananigans.Units
using Oceananigans.TimeSteppers: Clock, update_state!
using Reactant, CUDA
using Oceananigans.Architectures: ReactantState
Reactant.set_default_backend("cpu")
maximum_water_storage = 400.0
times = range(0, 4 * 3600, step = 1hour)
function bucket_model(grid)
slab_land = SlabLand(grid;
energy = SlabEnergy(eltype(grid); dry_heat_capacity = 0.1 * 1500 * 1480),
hydrology = BucketHydrology(eltype(grid); maximum_water_storage))
atmosphere = PrescribedAtmosphere(grid, times; surface_layer_height = 10, boundary_layer_height = 512)
set!(atmosphere; u = 3.0, q = 0.005, p = 101325)
for n in eachindex(times)
set!(atmosphere.tracers.T[n], 275.0)
set!(atmosphere.freshwater_flux.rain[n], 6e-3)
end
update_state!(atmosphere)
radiation = PrescribedRadiation(grid, times;
land_surface = SurfaceRadiationProperties(0.2, 0.97),
ocean_surface = nothing, sea_ice_surface = nothing)
for n in eachindex(times)
set!(radiation.downwelling_shortwave[n], 300.0)
set!(radiation.downwelling_longwave[n], 300.0)
end
update_state!(radiation)
al_interface = atmosphere_land_interface(grid, atmosphere, slab_land;
specific_humidity = FractionalHumidity(efficiency = CriticalSaturation(0.75)),
solver_stop_criteria = FixedIterations(8))
return AtmosphereLandModel(atmosphere, slab_land; radiation,
atmosphere_land_interface = al_interface,
clock = Clock(grid))
end
apply_rad!(model) = (NumericalEarth.EarthSystemModels.apply_air_land_radiative_fluxes!(model); nothing)
grid = RectilinearGrid(ReactantState(); size = (), topology = (Flat, Flat, Flat))
model = bucket_model(grid)
# Fails on Julia 1.12 (`raise_triton_custom_call` operand mismatch / with PR #2782,
# "non-concrete traced value at field: time"); compiles + runs on Julia 1.11.
compiled = Reactant.@compile raise=true raise_first=true sync=true apply_rad!(model)
compiled(model)
```
Simpler variants compile fine: passing the bare `clock.time` scalar, or a kernel
taking a struct of concrete (non-traced) values, both lower without error. Only a
traced scalar nested inside a struct kernel argument triggers the failure.
@dkytezab @Pangoraw @glwagner
1 条评论