EarthSciIntegrator: decisions and tradeoff analyses

Facts are in survey.md and the design is in architecture.md . This file records each decision and the analysis behind it.

Decisions

1. Adjoint strategy: tradeoffs

Goal: a full-simulation adjoint, as in reseact. That means the gradient of a scalar functional of the trajectory with respect to parameters, initial state and possibly gridded fields. It must be exact for the trajectory that was actually computed: reseact checks to 1e-11 against a reference, with dot-product and FD gates.

1.1 Parts every option needs

The bulk of reseact’s tools/adjoint_gradient.jl is this machinery, and it is needed whichever way the per-step derivative is produced:

• A step tape. Record the accepted (t, dt) sequence, the method per leaf, discrete decisions (event times, which maps fired, clamp masks) and forcing epochs. The backward sweep replays this tape with the controller off.
• Checkpointing. A per-macro-step state checkpoint costs about 680 kB per step at CONUS 4×5°. At 0.25° a 5-day run would need about 188 GB, so we need multi-level checkpointing (memory plus disk) and revolve-style schedules for long or high-resolution runs.
• Adjoint of composition. The adjoint of Lie(A, B) is B* then A*, and Strang, Sequence, Subcycle and Every compose the same way. Splitting is therefore not a problem for discrete adjoints.
• Classification and functionals. Parameters are classified as active or inactive, by where each value is consumed. Cost functionals and observation operators are defined over outputs.
• Validation. Dot-product identity per step, FD of the frozen-dt map, and replay-versus-checkpoint consistency.
• Controller derivative. The derivative of the controller (dt as a function of state) is ignored, which is standard. reseact’s forward-sensitivity lane does carry dt tangents.

1.2 The choice: what does automatic or symbolic differentiation see?

1.3 Where the model primitives come from (for A)

1. EarthSciAST emits them by source-to-source transformation (recommended): f, jvp, vjp_u, vjp_p, hvp and the structured Jacobian (EarthSciASTDiff bands) from the same expression tree.
   • This works the same way for Julia, Rust and later Python, and can be conformance-tested.
   • Recurrences (faq sweeps such as Thomas solves or RRTM) need reverse-mode code generation for loops. The compiler has the structure to do that.
2. Julia: Enzyme on the compiled RHS only. This is much smaller and more robust than Enzyme on a step.
3. Rust: std::autodiff on the RHS, if it stabilizes.

1.4 Recommendation

• A is the spec’d contract in both languages. IFT is used for every nonlinear and algebraic solve, and EarthSciAST emits the model primitives.
• Julia additionally offers B (Enzyme on a step) for leaves that have no hand adjoint. B also serves as the conformance oracle that checks A’s hand adjoints.
• C is not supported.
• Rosenbrock’s second-derivative requirement is the main extra primitive A needs. For per-cell chemistry it is a small block-local HVP. Where it is unavailable, running SDIRK in adjoint mode avoids it, since SDIRK needs only first derivatives plus transposed solves.

1.5 SciMLSensitivity.jl (formerly DiffEqSensitivity.jl), considered 2026-10-14

Its catalogue is exactly the B and C columns of §1.2, so it cannot supply A:

No method-level discrete adjoint exists in it or elsewhere in SciML. What is reused: ReverseDiffAdjoint through OrdinaryDiffEq as the per-family adjoint oracle (D31); SteadyStateAdjoint, the IFT adjoint of a nonlinear solve, as an optional implementation for solve nodes (D32); the ForwardDiffSensitivity(convert_tspan=true) pattern for forward sensitivities with dt tangents (D36). Its callback handling (freeze firing times, differentiate the affect map by AD) matches architecture §9’s rule but is not lifted, since our maps are compiled .esm kernels with band models (U6, U7). Its VJP wrappers are replaced by DifferentiationInterface (D38).

2. Partitioning and parallelism: tradeoffs

2.1 What “every stage is global” loses

If the RHS is one whole-domain function and parallelism lives only inside it, every stage of every leaf touches the whole state and ends in a barrier. The costs:

1. One dt per leaf for the whole domain. One stiff cell, such as the pre-dawn NO minimum, rejects the step for every cell. reseact measured its global chemistry dt at about 6× the per-cell ideal in cell-steps.
2. Memory and locality. An implicit leaf holds stage vectors and Jacobians for all cells at once. On EqAtmChem d05 that is 6.7M cells × 64² × 8 B ≈ 220 GB of CBM-Z block Jacobians (est.); run per shard, it is O(shard). Each explicit stage also streams the full state through memory, so tiles cannot stay in cache across stages.
3. Parallel efficiency is whatever the kernel achieves. In reseact, transport runs on about one core and XLA:CPU parallel efficiency is about 9%. There is a barrier per stage. Load imbalance cannot be rebalanced: chemistry costs 9.6 s/step pre-dawn against 35 s/step at midday.
4. No distributed memory. State, checkpoints and I/O share one address space. That rules out WRF d05 with chemistry, and the roughly 188 GB of 0.25° adjoint checkpoints.
5. No heterogeneous executors (processes, devices, ranks). reseact needed separate processes because XLA:CPU executions in one process do not overlap.
6. Compile cost grows with the domain on tracing backends. reseact’s lesson was that compile and memory must be per tile.

Gained by staying global:

• simplicity;
• boundary conditions and halos are handled once, by the compiler;
• global operators are natural: SSA, iceberg connected components, global norms;
• a single GPU likes large kernels.

2.2 The partition follows each operator’s locality

reseact-style stiffness sharding is the pointwise row. We adopt its executor contract from tools/shard_exec.jl:

• species-major layout at the boundary;
• per-shard partial error norms reduced in a fixed order, for determinism;
• per-epoch lane refresh;
• per-shard dJ/dp;
• executors read no driver state.

Executors are threads by default, then processes, devices and MPI ranks. reseact used processes only because of XLA:CPU.

With per-shard dt, the adjoint tape is per shard, so tape memory scales with the number of shards, not cells.

The hierarchy is rank → tile → shard. In shared memory, tiles are index ranges over one global array and the halo “exchange” is a barrier. MPI later replaces this with real exchange, with no change to the leaf algorithms.

What the compiler must provide:

• Pointwise and column partitions: nothing new. reseact’s capacity builds re-instantiate the operator on a lane grid with gathered inputs, at a build cost that depends only on the lane count.
• Exact stencil tiles: a range-restricted RHS that evaluates cells in a range and reads halo cells from a neighbour buffer. Today’s flat f cannot do this. A tile re-instantiated as a small model would apply boundary stencils at its edges, so it would not match the global run.

3. Compiler contract: inferring strategy from the AST, integrator side

“Option 2” means: EarthSciAST changes only by solver.external. The integrator takes the flattened AST, infers the structure itself, and compiles each piece with the stock compiler. That is how reseact already works:

1. flatten;
2. EarthSciASTSplitter classifies terms by locality;
3. each part is compiled separately;
4. EarthSciASTDiff’s band model gives the Jacobian;
5. capacity builds give the shards.

3.1 What option 2 gives

• No waiting on spec RFCs.
• Strategy inference becomes a feature, like DifferentialEquations’ default algorithm choice but structure-aware:
  • block-diagonal → batched Rosenbrock;
  • banded column → block tridiagonal;
  • stencil → SSP-RK with a Gershgorin bound on the band Jacobian as the step limit;
  • algebraic → Newton–Krylov.
• It fits the .esm philosophy: solver-relevant classes are derived, not declared.

3.2 What it gives up or makes harder

1. Authorial intent that structure cannot show.
   • Locality separates pointwise, column and stencil terms. It cannot tell radiation from PBL (both column), acoustic from advective (both stencil), microphysics-after-RK, the radt cadence, or PISM’s 12-step order.
   • The run config must name these. After flattening, a term’s component can only be recovered from the namespaced variables it references. That usually works, but it fails when a term references two components’ variables.
2. Derivatives come from separately compiled “derivative models”. This is the band-model trick, and it works (reseact).
   • It is fragile across versions. reseact’s band model went silently empty after an upstream node rename (aggregate → faq), which is why validate_plan (about 338 s) exists. Note what that check is: it compares the gather plan against EarthSciASTDiff’s own index maps at one point and refuses an empty Jacobian (rx_sym_block_jac.jl:288-337). It is not a band-versus-derivative value check; no reusable tool for that exists in any sibling repo (EarthSciASTDiff’s ForwardDiff oracle lives only in its test suite), so architecture §3.2 pass 8 is new integrator work.
   • EarthSciASTDiff fails loudly on filtered faq nodes and non-+ semirings. It has no ∂f/∂(forcing) (D5), and does not list recurrence (column-sweep) support.
   • Rust has no AD fallback, so the Rust adjoint depends entirely on EarthSciASTDiff coverage.
3. No exact stencil tiling and no MPI (§2.2). Pointwise and column sharding still work.
4. Events, algebraic equations and step maps are refused by every fast compiler. The integrator must strip them from the AST and compile the conditions, affects and residuals as separate observed models, in each language.
5. Duplication. AST analysis and surgery would be implemented in Julia and in Rust, and later Python, unless it is done once (§3.3).
6. Split parts are full-length RHSs over the shared u. The Splitter produces no DOF subsets, so work is wasted on zeros. This is fixable on the integrator side.

Unavoidable even under option 2:

• solver.external;
• a public, thread-safe RHS entry point in Rust: today debug_eval_rhs is #[doc(hidden)] and EsmProblem is built on Rc/RefCell;
• Float32 integration not being refused (WRF).

Cheap additions that recover most of the loss:

• a per-term provenance tag carried through flatten (fixes 1);
• a range-restricted RHS (fixes 3).

3.3 Proposal: a planner that writes artifacts

Structure inference runs once, in a planner, which writes portable artifacts:

• a resolved run config: an explicit schedule, inspectable and editable, and the file solver.external can point to;
• derived .esm documents: split parts, Jacobian and VJP band models, residual models for algebraic constraints, and event-condition models.

Both runtimes consume these artifacts. Inference is implemented once, and the runtimes stay thin and conformance-testable.

Superseded by D8: the planner is implemented in both languages. The artifact boundary stays, as the conformance seam: both planners must produce equivalent plans for the same document.

4. GPU via XLA: tradeoffs

Feasibility. Writing the integrator’s kernels in XLA works, and has been done:

• reseact’s tools/reactant_handoff/rx_traced_integrator.jl already runs SSPRK43, Rosenbrock23 (with the batched 13×13 block solve), the PI controller and adaptive while loops as one XLA program per macro step.
• Both EarthSciAST bindings already emit the RHS as XLA. Julia emits StableHLO through Reactant. Rust uses XlaBuilder via the opt-in xla 0.4 crate (simulate_array/tape/xla_emit.rs, 2,121 lines), which links a 144 MB prebuilt xla_extension.

So a whole leaf step, RHS plus stages plus linear solves plus controller, can be one XLA program on CPU or GPU. No XLA↔CUDA buffer interop is needed.

What fits XLA’s static-shape, data-parallel model:

• explicit RK stages;
• Rosenbrock and SDIRK stages with per-cell LU, unrolled over the n×n entries and vectorized across lanes (reseact’s blocksolve);
• column Thomas sweeps (a loop over levels);
• error norms, PI control and fixed-tree reductions;
• Newton and Krylov while loops. These run forward only, since D4 never reverse-differentiates through them.

What is awkward:

• Lane pools need a fixed lane capacity and refill in masked chunks; finished lanes idle until the chunk ends. reseact’s masked device loop wasted 20–28%.
• Dynamic re-sharding works if the composition is a runtime gather-index input rather than a shape, so nothing recompiles.
• MPI halos and distributed Krylov. XLA has no MPI. The options are:
  • split programs per stage and exchange halos on the host — reseact warned these round trips “would DOMINATE on an A100/H100”;
  • XLA collectives (collective-permute, all-reduce) — fine across one node’s GPUs; multi-node needs PJRT’s distributed runtime.
• Determinism (D10) needs XLA’s deterministic-ops mode on GPU. reseact also hit a data race on XLA:CPU (one NaN in about 2,000 calls).

Measured costs (reseact):

• On CPU, XLA was 2.8× slower than native Julia, had about 9% parallel efficiency, and a roughly 3 ms per-call floor.
• Compiles took minutes per program. Pathological ones ran for hours or exhausted memory (HloCSE, LLVM section memory). Compiling per tile bounds this.
• Upstream bugs: miscompiles, and @compile donating input buffers.

Effort by language:

• Julia: trace generic kernel code with Reactant; reseact’s integrator is a working template. One source that is both fast native code and traceable is possible but may be suboptimal for one of the two.
• Rust: build programs op by op with XlaBuilder. This is more verbose, but EarthSciAST’s tape→XLA emitter is precedent.

Recommendation.

• Keep native kernels for CPU (D11).
• Give the integrator’s kernel-primitive layer an XLA lowering for GPU: each leaf step becomes one XLA program, and the host is touched only at macro-step and epoch boundaries and for MPI.
• Accept these rules on the XLA path: fixed lane capacity, composition as gather indices, per-tile programs, deterministic-ops mode, and XLA collectives within a node.

5. Open questions

• The GPU hardware target (D28 keeps both kinds viable).

6. Alternatives recorded but not adopted (2026-10-14 review)

• Sequence the Rust planner after the Rust runtime. The plan artifact is the seam, so the Rust runtime could consume Julia-produced plan.json and derived documents until a Rust planner lands (after M5), with D8 as the end state. Not adopted; D8 and M2 stand as written. Recorded in architecture §15 as the fallback if planner effort threatens M3.
• A separate, later Rust M4 gate. Rust’s Rosenbrock adjoint depends on U10 W0 → U7, U8, U9 (filtered faq) → U12, six upstream RFCs, while Julia’s M4 can use the Enzyme-on-RHS fallback. Not adopted as a milestone change; recorded in architecture §15 risk 3 so that a Rust slip does not hold the Julia gate.
• Wrapping OrdinaryDiffEq per leaf for the forward pass of stencil, column and global leaves on Julia’s native path. Rejected for the same reasons as D31 (the adjoint would have to read OrdinaryDiffEq cache internals), but it remains the fastest route to a forward-only M3 transport leaf if M1 slips.