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.

7. Data inputs: tradeoffs (D87, 2026-11-02)

The problem. Real models read gridded data. reseact.esm reads GEOS-FP meteorology (GEOSFP.PBLH, GEOSFP.CLDTOT, winds, temperature and more) as parameters with a data update (esm-spec §5.4, §8.5), mapped by param_to_var onto shaped parameters with a leading two-record axis (Transport3D.F_PBLH[gf_t, gf_lat, gf_lon]) that the model blends in time with a weight it computes from t. Until D87, M2 refused every such model (unsupported_parameter_update, D75), and spec/operator-interface.md §8.1 said derived documents carry no data providers. So pass 8’s checks and §5.5.2’s stiffness bound had no values for these parameters, and a runtime building an operator from derived documents had no way to give its right-hand side time-varying fields. No fixture or toy problem had a data input, so none of this was exercised, and it stops the reseact.esm planner smoke test.

What EarthSciAST already gives. rhs_program takes forcing as an argument: every parameter whose update has from (or a handler) is a forcing slot, a 0-d slot for a scalar, out of the scalar parameter vector; a forcing set is immutable and carries an epoch stamp; refresh is a swap (API_SPEC §5.8). Data providers are always the caller’s: no binding fetches from data_sources on its own (Julia providers, Rust ProblemOptions::providers, the EarthSciIO adapters opt-in). EarthSciASTDiff’s forcing axis takes the forcing set as a list of names (esm-jacobian-spec §4.3) and already uses the leading record axis convention. EarthSciIO’s PrefetchGroup gives the epoch boundaries. The planner half uses these as they are, and needed two upstream fixes to stop patching documents (spec/plan.md §10 items 16, 17, resolved by EarthSciAST#634 and EarthSciASTDiff#55). The runtime half rests on upstream work that arrived on 2026-11-02 and on what is still missing (§10 items 18–22, §7.3 below).

7.1 How a data-fed parameter appears in a derived document

7.2 What the planner’s checks evaluate data at

7.3 How the runtime supplies data

The forcing set becomes the last argument of every operator call at M2 (spec/operator-interface.md §8.4). The runtime owns the data: it builds providers from the model’s own data_sources (not from the derived documents), turns a sample into a buffer exactly as esm_problem does, and swaps one immutable set per forcing epoch; steps never cross an epoch, and time interpolation stays the model’s (D5’s per-record gradients rely on the model’s blend being explicit). Rejected: interpolating in time in the runtime (it would put a second, hidden blend beside the model’s own, and break D5’s accumulation per record); giving each operator its own provider (several decodes of one field per epoch); passing data through $p$ (only scalars, and it would change the parameter vector’s length with the data).

Locating the data. The model locates it: url_template resolves relative to the model’s file (esm-spec §8.2.1), so a plan runs wherever the model and its data sit together. Rejected: naming data files in the run config (a second declaration of what the model already declares, which can disagree with it) or in the plan (the plan would carry paths, which §2.1 forbids, and go stale when data moves). Cache directory, offline mode and URL overrides are run-time options of the runtime’s API: they change where bytes come from, never what is computed.

What the runtimes call (spec/operator-interface.md §8.4.2): EarthSciAST’s providers_from_document (Julia, from its EarthSciIO extension; Rust, feature esio), built on EarthSciIO’s constructor of a provider from a document’s data_sources entry; provider_sample / CadenceProvider::refresh at each epoch; EarthSciAST’s public sample-to-buffer conversion; EarthSciIO’s refresh times for the epochs. The proposal had the runtimes write a small data layer of their own over a subset of esm-spec §8; review (2026-11-02) chose the full feature set instead, through EarthSciAST and EarthSciIO, so that a planned run and EarthSciAST’s solve read the same numbers by construction and reseact runs end to end at M2. The cost is that M2’s runtime half waits for the upstream work below. Rejected as well: building the original model’s whole esm_problem with providers and copying its forcing_buffers at each epoch, which compiles the whole model’s right-hand side only to read its buffers, and whose refresh is a solver callback, not a call a runtime can make at an epoch boundary.

Upstream work (spec/plan.md §10 items 16–22, as of 2026-11-02):

• Arrived: EarthSciAST#636 (update.from.records_per_sample, overriding the source’s temporal.records_per_sample, reported in loader_fields; metadata.esio_format defined in esm-spec §8.1.1), #632 (forcing_buffer in Julia and Rust, which does not apply unit_conversion because providers already do), #634 (after param_to_var the source’s declaration survives, so defaults go on sources), #635 (providers_from_document on EarthSciIO’s document constructor in all three languages); EarthSciIO#27–#30 (esm templates with substitutions and mirrors, exact calendar cadences, document-built providers, a Rust CSV reader); EarthSciASTDiff#55 (the evaluation document keeps data rules and their sources).
• Still missing: #635 does not pass records_per_sample through yet; no json reader in any language; the file’s own time axis is used only in Python; determinism is not checked; providers have no time_origin shift, so the runtimes convert; and three Rust EarthSciAST bugs (rhs_program misses array-default data-fed parameters as forcing slots, may return defaults instead of data with qualified provider keys depending on map order, and samples discrete providers only when a solve starts).

7.4 Milestone

Both halves are M2 (decided 2026-11-02, amending D75; the proposal had the runtime half at M3), with the full data layer: the planner half for the reseact.esm smoke test, the runtime half so that a planned run reads data end to end, reseact included, gated by data_decay, data_monthly and data_points (spec/toy-problems.md §4.4–§4.6). M3 keeps reseact’s forward run at scale, shards and lanes, which take forcing slices per lane (architecture §7.3). M4 adds forcing derivatives (U7, U8); M5 rank windows, forcing halos and asynchronous prefetch (U11).

7.5 Questions settled in review (2026-11-02)

First round: the stub data source is accepted; Gershgorin stiffness on operators that read data is refused; a used data-fed parameter without a nonzero default is refused (forcing_without_scale), not sampled around 1; band-value files list forcing values; the two derived-document gaps (§10 items 16, 17) are fixed upstream, with the planner’s patches as the fallback; both halves are M2. Second round: the two-record request and esio_format become esm fields (the run-config forcing.records stand-in is dropped, not kept as a fallback, since the planner and the runtimes read the model’s field); the full data layer is M2’s, through EarthSciAST and EarthSciIO; a file variable’s dimensions map onto a parameter’s by position, as esm_problem does; data identity is recorded in the run’s report, never in plan.json.

Third round, a review of the pull request (2026-11-03), fixed inconsistencies without changing the design: (1) only used data-fed parameters reach the plan: the plan’s forcing lists them, an operator’s derived documents declare as forcing inputs only the parameters in its forcing (another used one, which only ic equations read there, is a plain parameter at its default; an unused one is not declared), so a program’s forcing slots are the data it reads, the runtime builds no provider for data nothing reads, and band-value files list only the canonical forcing order; (2) the set a call takes is chosen by the call’s $t$, so a stage at exactly an epoch’s end reads the next epoch (architecture §7.5), and $F$ is constant on each step’s half-open interval rather than over the closed step, which a two-record blend with a $\mathrm{frac}$ weight needs to be continuous; (3) a refresh time at $t_f$ starts a last epoch that the stages at $t_f$ read, so its records must exist; (4) epochs and forcing_refresh_not_hit follow the parameters operators read, and a parameter only an ic equation reads is sampled once, at $t_0$; (5) a run’s initial state evaluates such an ic equation with the data at $t_0$, from EarthSciAST’s esm_problem with the providers, as EarthSciAST’s own solve does, while the planner’s $u_0$ keeps defaults (spec/operator-interface.md §8.4.5); (6) the wrt prefixes params: and fields: exclude data-fed parameters, which match forcing: only; (7) every call’s forcing argument is written into spec/operator-interface.md §3 and §8.3; (8) forcing_records_invalid covers only a temporal block that is not time-varying, since EarthSciAST already refuses a missing block and a record axis not of size 2 at load, and its fixture was changed to match; (9) run-config’s rule-16 paragraph no longer lists the data refusals.