perf(cam): compile the CAM driver — jit + lax.scan (plan 027) - #80
Open
aliakherati wants to merge 3 commits into
Open
perf(cam): compile the CAM driver — jit + lax.scan (plan 027)#80aliakherati wants to merge 3 commits into
aliakherati wants to merge 3 commits into
Conversation
Phase-B (ADR-004) for the plan-025 driver, mirroring the M6 PR-J1/J2 pattern from the E3SM side: - cam_run_step / cam_run_timesteps become thin UNJITTED wrappers whose only jobs are resolving topology=None outside the trace (get_topology() deliberately raises inside a jit trace) and stabilising the carry pytree (qaerwat pre-populated with the pbuf's zeros default), over jitted inners. - The substep loop and the step loop are lax.scan. The per-substep reseed flag and the per-step first_step flag ride the scan xs as traced booleans, so the box-reference reseed mode and production CAM's lagged mode share one compiled body. - Statics per ADR-020's split: topology (frozen/hashable, the #65 staleness rule), strat, n_substeps (scan length), the do_* toggles, and the two compat flags. so2_to_h2so4_rate and the whole state are traced operands: a rate sweep reuses one compilation — locked by test_rate_is_traced_and_substeps_are_static, alongside the inverse assertion that a new n_substeps DOES retrace. - The scan trajectory is BIT-IDENTICAL to repeated cam_run_step calls (same jitted body; assert_array_equal, not allclose). No physics or API change; every existing test passes unchanged, including all four end-to-end G5 parity bars. Measured on the G5 workload (120 steps x 16 substeps, strat, cam_mam4): eager 81.6 s -> ~1.0 s cold (compile + run) / 0.02 s warm — ~4000x. The trajectory scan traces the step body once regardless of n_steps. Full suite green (272).
Six figures under docs/figures (cam_strat_rh_*) regenerated end-to-end by scripts/cam_rh_scenarios.py — exploratory scenario science, not validation (no Fortran reference attached): - 6-h strong-forcing sweep (SO2 = 0.1 ppmv, volcanic-plume scale): the classic nucleation-growth banana; number and mass are RH-insensitive because the regime is production-limited. - 12-h weak-forcing experiments (H2SO4 = 1e7 molec/cm3) isolating the RH -> water uptake -> wet diameter -> condensation-sink chain: burst (no SO2) shows RH-ordered gas-depletion floors spanning 3.5 decades; background (SO2 = 1e-10 vmr) shows a ~12% number spread with MORE particles at LOW RH (the weaker sink leaves gas standing for nucleation); condensation-only (nucleation off) fans the wet aitken diameter 31.5 -> 36.2 nm at identical dry size with 35% lower steady gas at RH 50%. Also the empirical answer to 'is 1e7/cm3 enough': at 232 K it nucleates completely within one 30 s step at any RH, so RH-sensitive NUMBER requires the sustained weak-production regime. Figure conventions: per-mode categorical colors in fixed order across all figures; ordered dimensions (time, RH) on sequential ramps; per-mode breakouts everywhere (no collapsed mode fields).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #74 (base =
feat/cam-driver) while merges tomainare held. Phase-B optimization per ADR-004 — plan 025 established correctness; this PR only compiles it. Plan:docs/plans/027-cam-driver-jit.md.What
cam_run_step/cam_run_timestepsbecome thin unjitted wrappers — resolvingtopology=Noneoutside the trace (get_topology()deliberately raises inside a jit trace) and stabilising the scan-carry pytree (qaerwatpre-populated with the pbuf's zeros default) — over jitted inners:lax.scan; the per-substep reseed flag and the per-stepfirst_stepflag ride the scanxsas traced booleans, so the box-reference reseed mode and production CAM's lagged mode share one compiled body;topology(frozen/hashable — the feat(topology): make the jit-staleness hazard impossible to hit silently #65 staleness rule),strat,n_substeps(scan length), thedo_*toggles, the two compat flags.so2_to_h2so4_rateand the state are traced: a rate sweep reuses one compilation.No physics or API change.
Verification
cam_run_stepcalls (assert_array_equal, notallclose) — same jitted body either way.n_substepsmust.Measured
G5 workload (120 steps × 16 substeps, strat,
cam_mam4):The trajectory scan traces the step body once regardless of
n_steps.Docs updated in-PR per rule 5: plan 027 (with results),
PROGRESS.md, and theFEATURES.mddriver row.