Module allocatables as run state, character values and dummies, stub-dropped reads, use only-lists: what ELM's PhotosynthesisMod forced

[chenyueqi]

Rules two ELM units forced, none about ELM; each with a test naming the case. With them, Photosynthesis (E3SM e508436e45, components/elm/src/biogeophys/PhotosynthesisMod.F90) translates with zero deferred sites, is bit-exact on a recording of the model's own run (61,440 points, 0 ULP, read/write sets matching on 93 blocks), and lowers -- root finders included, as fixed-count loops -- to lax kernels within 10 ULP (day 1) and 21 ULP dominant (noon) of the same recordings, every output compared, with jax.grad agreeing with forward mode to 8e-15 and forward mode with finite differences of the anchor to 3.5e-7 through the root finders (private extension a85tract/recast-elm, docs/photosynthesis.md).

Plan, adapters, recorder

1. A module allocatable is run state with a shape: its own module's ALLOCATE (allocate (vcmax_np1 (0:mxpft))) is the shape, the extent spelled by the module variables that size it, which become recorded scalars; the Fortran adapter allocates over the original bounds; the recorder takes such an extent from size(). A reader that use-imports the array subscripts it over the companion's lower bound.
2. A character parameter is a value (namep = 'pft'): a str node, rendered by repr.
3. A character scalar dummy (phase = 'sun') is flat: character(len=*) in the adapter, text in the recorder and replay, admitted by the bit-exact gate under the recorded convention.
4. The recorder's own argument is recast_phase; the allocate emitter no longer crashes on a conflicting-bounds record.
5. The recorder guards every allocated array (associated()/allocated()), records NaN of the planned shape otherwise, and flushes after every section.
6. A pointer assignment is an alias: a write through psn_z => photosyns_vars%psnsun_z_patch writes every target it may have, where the alias is assigned (not merely pointed). Left out, the routine's own outputs (leaf photosynthesis, stomatal resistance, ci) were planned read-only and never compared: the first bit-exact verdict stood on 20 outputs, the present one on 42 (129,024 points).

Translation and the read/write gate
7. A use ..., only: bounds what a companion exports (an associate alias c3psn resolved to pftvarcon.c3psn and wrote photosyns_vars%qe_patch into pftvarcon.qe).
8. A stubbed call's argument reads are named to the read/write gate (dropped_reads), excused only where the target does not read them; a block that is nothing but an abort is waived by name.
9. The array transforms (reshape, spread, pack, unpack, transpose, cshift, eoshift, maxloc, minloc, lbound) are intrinsics to the analysis; the vocabulary test that kept the divergence records the gate run that ended it.

The JAX port
10. A subroutine's early returns of its output names merge on the flag alone (hybrid's if (f0 == 0) return).
11. A call's positional actuals follow the emitted convention (non-OUT dummies, then OUT array buffers): hybrid(..., gs_mol, iter, atm2lnd_vars, photosyns_vars) omits the OUT scalars.
12. A pointer local pointed in a branch starts the kernel as its first target (fourteen sun/shade views a lax.cond could not carry as None).
13. An integer literal into a real local is a float (nscaler = 1).
14. The JAX runtime carries _f_sqrt and _f_cfold as the NumPy one does (mpmath exempted by the contract the same way).
15. A while True its counter caps through a break is a fixed count of passes under the done flag (and under any loop condition, the single-exit flag included): lax.fori_loop with static bounds, which reverse mode transposes, where lax.while_loop has no rule.
16. A write through a branch-pointed local reaches every target it may have pointed at: psn_z(p,iv) = ... rebinds a JAX local, and the flat output it aliased came back as it went in (d(psn)/d(psn) was exactly one; the ULP gate read 9e18 ULP once it compared those outputs). Before the return each target takes the local under the condition of its pointing.

What the hydraulic-stress unit forced (PhotosynthesisHydraulicStress, the same module under use_hydrstress = .true., recorded on a sibling case: 96 calls, 159 inputs, 50 outputs; bit-exact on 267,264 points in NumPy, and every one of them within 23 ULP as a lax kernel -- the plant-hydraulics Newton solve and both capped root finders as fixed counts of passes)
17. A module-state object is reached by its name; its type is named only where a dummy of it is declared. params_inst is public and protected while photo_params_type is private to the module: the recorder's use PhotosynthesisMod, only: photo_params_type did not compile.
18. A borrowed ALLOCATE bound is spelled in the companion's namespace. ColumnType allocates dz(begc:endc, -nlevsno+1:nlevgrnd); PhotosynthesisMod subscripts dz over that bound and imports nothing named nlevsno, so the subscript read a bare nlevsno (NameError at the first call). The name binds to the module that declares it, imported if the reader has not; module-level and derived-type component allocations alike.
19. SUM is a left fold in element order. np.sum is pairwise past eight elements; over the ten soil layers it put 925 of 267,264 recorded points up to 75 ULP off, and a sequential fold replayed on the same dumps put every one of them on the recording. _f_vsum in both runtimes, as _f_vdot already is; a DIM reduction keeps the library's order.
20. The ULP gate judges what the backend lowered; a routine forwarded to the host fails by name. The JAX module binds what it could not lower to the anchor (f = _host.f) and lists what it did lower in _JAX_KERNELS; the tolerance verifier fetched the name from the module and awarded the port a bit-exact verdict on the hydraulic-stress kernel it had never emitted. A name outside the lowered list is not compared: it fails with the backend's recorded reason unless declared ungated (issue #32, row 22).
21. A call's OUT arrays are buffers in their declared place. The anchor's positional convention is declaration order with OUT scalars omitted and an OUT array passed where it stands (spacf(p, c, x, f, qflx_sun, ...)); the rewrite appended the buffers after the other dummies, shifting every actual behind one, so calcstress handed spacF a flux for atm2lnd_inst ("object not passed" -- a misalignment, not a missing object).
22. A buffered call's OUT scalars bind to temporaries. In _out = calcstress(...); bsun = _out[1] an OUT scalar has no actual to bind to ("output bsun has no target"); it binds to a temporary the unpack reads, and an object's unpack (photosyns_inst = _out[7]) is dropped since its components bound at the call.
23. A nested function's return is its own. PhotosynthesisHydraulicStress carries its ft/fth/fth25 statement functions as nested defs; their returns counted as early returns of the kernel and the single-exit rewrite refused the routine ("an early return in a function with several outputs").
24. A nested def at the top of a body is not descended by the early-return scan (the statement functions again, as the first statements of the body: the scan entered them from the root).
25. A loop around nothing lowers to nothing. brent_PHS checks its bracket per phase and ends the run; with the raise dropped the loop carried nothing and the backend refused the kernel ("loop with no carried effects"); as a guard around nothing already did, it lowers to nothing.
26. A module's own state object is spelled by its components in a kernel. The anchor spells PhotosynthesisMod's own params_inst bare (params_inst.psi50 inside plc); the flat spelling knew dummies and companion state only, so every hydraulic-stress helper kernel took params_inst as an opaque parameter no calling kernel had -- latent while the SP recordings never ran them, NameError the first time the hydraulic-stress kernel did.
27. A kernel's state closure is what its body still names. With the object spelled flat, the record's module_state_read still said params_inst, and the closure taken from it gave plc's kernel an opaque trailing parameter that every calling kernel passed bare.
28. A while loop's done flag starts False at the top of the kernel. calcstress runs its Newton loop inside if (.not. night); the branch carried the loop's flag and the lax.cond operand tuple read _done_1 before anything had bound it (UnboundLocalError). Like the goto-region flags, every done flag is initialized at the top.
29. Consecutive guarded statements share one block. After an early return or a break each following statement was wrapped in a guard of its own, so a call and the unpack of its result became two lax.conds and the temporary the first bound was not in the second's carry (NameError: _t1 in brent_PHS); a statement that returns or breaks closes the block.
30. A Fortran-order copy is the array itself in a kernel. transpose(a) is np.asfortranarray(a.T) in the anchor, a layout the F-order reshape downstream relies on; jnp has no layouts and no asfortranarray (SimpleMathMod's kernel raised AttributeError under the hydraulic-stress kernel).
31. A name that sizes an array or bounds a range is a Python int in a kernel. brent_PHS's integer, parameter :: nphs = 2 sizes its bracket arrays; the backend strengthened the literal to jnp.int32(2) for a stable carry dtype and jnp.zeros((nphs,)) refused a traced shape.
32. An alias bound twice to one target is still an alias. The anchor binds a pointer alias twice, as the view it opens the body with and again where the Fortran's => stands; counted as two stores it was a plain local, rebound in the loops and never written back, and the hydraulic-stress kernel returned every output as it came in (the gate read 9e18 ULP). Every top-level store binding the same value is one alias.
33. Nested capped loops have flags of their own. The capped while lowering gave its counter back, so hybrid_PHS's two capped loops were both _done_1 and the inner loop's exit ended the outer after one pass (the kernel's iter1 stayed 1 where the anchor ran to 4).

What the nutrient-limited path forced (Photosynthesis under use_cn = .true., nu_com = 'ECA', recorded on a third sibling case held in accelerated decomposition; bit-exact on 129,024 points in NumPy)
34. A return inside a loop ends the pass and the loop: the statements after the flag in any block run under if not _ret, a while takes the flag in its test (its body stays a body, so the cap it breaks on stays in the capped lowering's sight and the loop remains a fixed count of passes), a for, which has no test, runs its whole body under the flag -- no break, since a for lowers to a scan whose body cannot break. brent returns from inside its bracketing loop ahead of the next ci_func evaluation; the flag alone left the rest of the pass and the following passes running, and the kernel's gs_mol was an evaluation the Fortran never made (this recording is the first on which hybrid reaches brent).
35. A store inside an inner loop, and a subscript store, is a store of the enclosing block: the carry of a while_loop is computed before lowering, and brent_PHS's c(phase) = ... inside its do phase loops was counted by neither walk (UnboundLocalError: c in the pass).

Full suite passes (tests, minus the C kernels and JAX-diff suites).

🤖 Generated with Claude Code

https://claude.ai/code/session_01BwmCpkkFKaGihSPvqAQnaK