ITADN

`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

#3007Openxkykai 创建于 2026-06-25
X
xkykaicommented
## 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 条评论