[dkytezab]

WIP

[dkytezab]

Tracker: compiling the initialization of a NumericalEarth nested model on ReactantState

In draft PRs:

• Breeze: Create separate initialize! method Breeze.jl#1056 (initialize!), Recompute the Exner reference state inside a compiled function Breeze.jl#1112 (Exner reset, fixes Recomputing the reference state inside a Reactant-compiled function fails (set!(model; compute_reference_state = true)) Breeze.jl#1111)
• Oceananigans: Move LatitudeLongitudeGrid vertical coordinates to ReactantState via Adapt CliMA/Oceananigans.jl#6007 (vertical coordinates to the device through Add ConcreteRArrayAdaptor, an Adapt adaptor to ConcreteRArray EnzymeAD/Reactant.jl#3405, merged), Keep which LatitudeLongitudeGrid boundaries are poles in the grid's type, so Fields allocate inside compiled functions CliMA/Oceananigans.jl#6334 (polar boundaries, fixes Allocating a Field on a ReactantState LatitudeLongitudeGrid inside a compiled function fails at the polar boundary check CliMA/Oceananigans.jl#6333), Reactant run!: peel the first step only for Adams-Bashforth CliMA/Oceananigans.jl#6330 (compile size, fixes Reactant run! traces the first time step twice for every timestepper; only Adams-Bashforth needs the peel CliMA/Oceananigans.jl#6329)
• NumericalEarth: Hold as many derived levels in the exchanger as the parent holds #803 (exchanger window, fixes Parent-first nested_atmosphere_model gives the exchanger a 3-level window even when the parent is fully resident #802), Fill the ERA5 level pressure through set! #807 (pressure_level_field, fixes pressure_level_field fails on ReactantState: broadcast into interior of a device field #805), Compute surface_elevation through set! #808 (surface_elevation, fixes surface_elevation fails on ReactantState: broadcast between interiors of device fields #806), Default atmosphere_model's and SlabLand's clock to Clock(grid) #811 (atmosphere_model's and SlabLand's clock, fixes atmosphere_model's default host clock freezes time in a compiled step on ReactantState #810), Reduce the mean sea-level pressure with the Field's own sum #815 (mean_sea_level_pressure, fixes mean_sea_level_pressure fails on ReactantState: generic sum over interior of a device field #814)
• Reactant: Clamp loaded values of CuTracedRNumber arguments in kernels EnzymeAD/Reactant.jl#3568 (clamp over converted traced numbers, fixes Raised kernel: clamp over numbers converted to the kernel's traced type returns zeros EnzymeAD/Reactant.jl#3567), Accept types as kernel arguments EnzymeAD/Reactant.jl#3571 (types as kernel arguments, fixes Kernel-argument check rejects a type passed to a kernel (DataType), which CUDA and the CPU backend accept EnzymeAD/Reactant.jl#3570)

Without a draft PR yet. init = this PR's compiled initialization (synthetic parent); ERA5 = the same with an ERA5 pressure-level parent; window = compiling the time step afterwards.

repo	item	needed for	state
NumericalEarth	#812 set_region_data!: Missing-typed NetCDF slices (EnzymeAD/Reactant.jl#3569), reshaped slices (EnzymeAD/Reactant.jl#3204), the float type as a kernel argument (EnzymeAD/Reactant.jl#3570)	ERA5	issue; NumericalEarth-side fixes on branch debug-nestedsim, mixed with other changes
NumericalEarth	#813 column_fractional_z_index returns zeros under raising; convert(FT, clamp(kk, k_sfc, grid.Nz)) avoids it	ERA5, window	issue; one-line fix
NumericalEarth	the nested step subtracts the parent's host clock time from the child's and branches on it (time_step!(::NestedModel)); the parent should tick on a device clock unconditionally	window	not filed
NumericalEarth	printing a PrescribedAtmosphere on the device takes extrema over device views	display	goes with CliMA/Oceananigans.jl#6325
Oceananigans	CliMA/Oceananigans.jl#6325 xnodes/ynodes, CliMA/Oceananigans.jl#6326 reduced-Field getindex, CliMA/Oceananigans.jl#6327 index_binary_search on ReactantState	init (#6325, #6326), ERA5 (#6327)	issues; fixes on one branch, not yet pushed
Reactant	EnzymeAD/Reactant.jl#3204 reshape(::CuTracedArray) calls an undefined _derived_array; EnzymeAD/Reactant.jl#3569 Ops.constant recurses on an unsupported element type	ERA5 (through #812)	issues; #812's NumericalEarth-side fixes avoid both
Breeze	acoustic RK3 cold start: stage 1 of step 1 builds the scalar tendencies from a zero transport velocity on ReactantState (first step only, small)	window	not filed; MWE ready
Enzyme-JAX	raising a traced Bool kernel argument; reverse mode through a loop nested in one with a dynamic trip count	window, gradients	not filed; MWEs ready

Dropped: the Interpolated boundary's clock_time conversion (does not reproduce), and raising a loop with an early return (refused with a clear error, and CliMA/Oceananigans.jl#6327's fixed-trip search avoids it).