# Profiling and benchmarking Two tools, one question: *where does a step's time go, and is a change faster without changing the answer?* ## `profile_graph` ```python from maddening import GraphManager from maddening.core.simulation.profiler import profile_graph from maddening.nodes import SpringDamperNode gm = GraphManager() gm.add_node(SpringDamperNode("spring", 0.01, stiffness=30.0, initial_position=0.5)) gm.compile() report = profile_graph(gm, n_steps=200, n_warmup=5, trace=True) print(report) ``` `python -m maddening.examples.advanced.profiling_demo` runs it on a coupled pair with every measurement on (`measure_coupling`, `counts`, `count_scan_steps`, `trace`) and reads each part of the report back: the {term}`coupling ` overhead and its cost per extra iteration, the bottleneck and recommendations, the compile counts (checked identical on a freshly built graph) and the trace summary, with the Perfetto JSON and the trace written to a temporary directory. `python -m maddening.examples.advanced.profile_lbm_step` is the short Perfetto-export example. What it measures, and how: | number | how it is obtained | |--------|--------------------| | `jit_compile_ms` | the first step (compile + run) | | `mean_step_ms`, `median`, `p95` | `n_steps` warmed steps, each blocked on the device | | `dispatch_floor_ms` | a jitted identity on the same `(state, ext, params)` {term}`pytree `: the cost of dispatching a step before any physics runs | | `node_times_ms` | each {term}`node `'s `update` jitted and timed in isolation (an estimate: one dispatch per node, no coupling) | | `coupling_iter_stats[group]` | over up to 50 further steps: iterations used (mean / min / max) against `max_iterations`, fraction of steps that exited at the cap, fraction whose residual met the threshold, from the always-on `_meta` diagnostics | | `coupling_overhead_ms` (`"measured"`) | the graph is recompiled with every group capped at one iteration and timed; the difference to the real step is what the extra iterations cost, and `coupling_per_iteration_ms` divides it by the mean number of extra passes: the real step's `total_iterations` less the one pass per waveform sweep the capped step still runs (before 0.4.0's fix it divided by `iterations - 1`, the largest sweep's count, which overstated the cost of a pass by nearly `waveform_iterations` times on a group that runs several sweeps). The graph, its state and compiled step are restored afterwards. `measure_coupling=False` falls back to the old estimate (step minus the sum of isolated node times). The two windows are timed at different points of the trajectory -- `mean_step_ms` just after a reset, the capped steps from wherever the run and the statistics pass left the state -- which is sound because a cap of one returns after a single staggered pass with no data-dependent control flow, so the capped step costs the same from any state (measured on the compute-bound `expensive-pair` fixture: pinning both windows moves the number by less than one standard error of the difference -- 0.03 SE over 10 repeats on a quiet box, 0.62 SE over 16 on a loaded one). **Reported signed**: on a graph whose step is dominated by dispatch the difference is smaller than its own scatter and falls either side of zero (measured: negative in 11 of 30 runs of a two-spring pair), so a negative value is the honest report of an overhead the run could not resolve. It used to be clamped at zero, which biased the mean upward by discarding only the negative half of the noise | | `coupling_overhead_se_ms` | the standard error of that difference, propagated from the two windows' own per-step scatter -- no chosen threshold. Below it the overhead was certainly not resolved and `str(report)` says `below resolution`; above it, not necessarily, because it captures the scatter *within* each window and not drift between two windows timed minutes apart. On a shared machine the spread over repeated runs is the better guide. `"measured"` only | | `trace` (`trace=True`) | a short `jax.profiler` trace; device kernel time per step, kernels per step, host dispatch of the step's own `PjitFunction`, and kernel time by scope | ### Reading the trace attribution Node updates run under `jax.named_scope("node:")`, the residual under `coupling:residual`, acceleration under `coupling:accelerate`, interface corrections under `coupling:interface_override`, mappings under `edge:mapping`. {term}`XLA` labels a *fused* kernel with the longest common prefix of the fused ops' names, so a kernel that fuses two nodes' work is reported under the shared prefix (`while/body` for anything inside the coupling loop) rather than under either node. Attribution is therefore exact for unfused kernels and coarse inside the loop body — which still separates "inside the {term}`fixed-point ` loop" from "outside it". `device_busy_fraction` is the number to look at first on a GPU: kernel time over wall step time. Well below one with hundreds of kernels per step means the step is **launch-bound** — fewer, larger nodes or fewer iterations help; optimising a single node's arithmetic does not. ### What the profiler does *not* tell you It cannot see a compile that happens inside a caller's timing loop. A weak-typed leaf in the seed state (`jnp.array(0.0)` without a dtype) retraces the step when the leaf comes back strongly typed; `compile()` normalises the seed state, but a driver that assigns its own arrays into `gm._state` can reintroduce it — use `gm.set_node_state()` and `gm.reset_state()` instead, and check that `gm._compiled_step._cache_size()` stays at 1 across a run. `gm.trace_count` covers the compiled step only. `run_scan`, `run_scan_with_history`, `run_sweep` and `run_adaptive_scan` build a separate `lax.scan` program around that step, counted by `gm.scan_trace_count`: one per `compile()` per (entry point, step count) in a healthy run. A driver that loops over `run_scan` and sees that number climb is paying a compile per call, which `trace_count` alone would not show. ## `benchmarks/bench_coupling.py` ```bash python benchmarks/bench_coupling.py --graph mime-ar4 \ --experiment ../MIME/experiments/ar4_helical_drive --trace --json out.json JAX_PLATFORMS=cpu python benchmarks/bench_coupling.py --graph heat-chain --n-cells 256 ``` Runs `profile_graph` on the graph with its {term}`coupling group `(s) and on the same graph with them disabled (staggered back-{term}`edges `), prints both reports and a summary (coupling cost vs baseline, iterations used vs cap, PERF-1 acceptance at `--acceptance-ms`, default 30), and writes a JSON record. Keep before/after records under `benchmarks/results/` when a change claims a speed-up; a claim without the two JSON files is a guess. ## `benchmarks/bench_coupling_sweep.py` ```bash JAX_PLATFORMS=cpu python benchmarks/bench_coupling_sweep.py \ --json benchmarks/results/coupling_sweep_cpu.json ``` Sweeps every `iteration_mode` x `acceleration` x `convergence_norm` over the graph shapes in `benchmarks/coupling_fixtures.py` — chain, star, ring, a stiffness sweep, two grid fixtures, a two-group graph and a slow-drift case — and records per configuration the step time, the iterations used against the cap, the fraction of steps converged, the final residual and the launch-bound / compute-bound verdict. The results and what they mean for a given graph shape are in [Choosing a coupling algorithm](coupling_algorithm_guide.md). `--steps` changes only how many timings are averaged. The iteration counts, convergence fractions and residuals come from a separate statistics pass whose length and starting point are the *fixture's*, not the run's (`--stat-steps` overrides it, `profile_graph`'s `n_stat_steps` is the underlying knob). That matters because those three are the numbers people quote across runs: while the pass was `min(n_steps, 50)` steps taken from wherever the timed run stopped, a shortened run moved the window as well as the sample size, and on a periodically driven graph the mean iteration count moved with it. ## Persistent compilation cache ```bash export MADDENING_COMPILATION_CACHE_DIR=~/.cache/maddening/xla ``` `GraphManager.compile()` picks the variable up and points {term}`JAX`'s persistent cache at it with thresholds that cache MADDENING's sub-second compiles (JAX's defaults skip them). `compile_cache.enable(dir)` does the same from code, and `compile_cache.warm_cache(gm_factory, n_steps=2, scan_steps=N)` compiles a graph's `step` and `run_scan` ahead of a run. The cache key covers the JAX/XLA version, backend and traced program, so a stale entry is a miss, never a wrong executable. **Which directory.** An explicit directory wins over the variable: `enable(dir)` and `warm_cache(..., cache_dir=dir)` use `dir` whatever `MADDENING_COMPILATION_CACHE_DIR` says, and `GraphManager.compile()` enables the variable's directory only while no directory is enabled yet. `enable()` and `warm_cache()` without a directory keep the one already enabled, then fall back to the variable, then to `~/.cache/maddening/xla`. ```python import os import tempfile from maddening import GraphManager from maddening.core.simulation import compile_cache from maddening.nodes import SpringDamperNode with tempfile.TemporaryDirectory() as mine, tempfile.TemporaryDirectory() as env_dir: os.environ[compile_cache.ENV_VAR] = env_dir # as if set in the shell assert compile_cache.enable(mine) == mine # explicit: wins over the variable gm = GraphManager() gm.add_node(SpringDamperNode("s", 0.01)) gm.compile() # leaves the enabled directory alone gm.step() assert compile_cache.enabled_dir() == mine assert os.listdir(mine) and not os.listdir(env_dir) ``` **A change of directory after the process has compiled.** JAX opens its persistent cache once, at the first compile that finds a directory set, and ignores a later change of the setting. So `enable(other)` after that also drops JAX's open cache (`jax.experimental.compilation_cache.compilation_cache.reset_cache`), and the next compile opens `other`. Programs the process has already compiled stay in memory and are not written to `other`. On a JAX without that reset, the move is refused with a `RuntimeError`: enable the directory before the first compile, or set the variable before the process starts. `compile_cache.enabled_dir()` reads JAX's own state and answers `None` rather than name a directory JAX is not using (after a direct `jax.config.update("jax_compilation_cache_dir", ...)`, say). `python -m maddening.examples.advanced.compile_cache_demo` runs a cold and a warm process against a temporary directory, with the variable pointed at another, and prints both compile times. ## Compilation counts, and the regression gate Compile time is what a user of MADDENING feels first — the pitch is that the whole graph compiles to one jitted step — and until v0.4.0 nothing stopped a silent 3× compile-time regression from shipping. The gate that guards it does not read the clock. This repository has already paid for a wall-clock gate: a test comparing per-step cost at two coupling iteration caps with a 3× bound, on quantities around 1e-4 s, failed CI at 3.27× on a pull request that changed 122 documentation files and no code at all. On an idle box the real ratio is 0.97–1.31. The number it read was the runner. So the gate reads integers. `profile_graph` reports a `CompileCounts` alongside the timings (`compile_counts(gm)` gets them on their own): | count | what it is | why it matters | | --- | --- | --- | | `retrace_count` | Python traces of the compiled step since `compile()`, i.e. XLA compilations | One is healthy. An unexpected retrace is the classic silent regression — a weak-typed leaf, a dtype that drifts, an argument that stopped being static — and costs a full compile on every run | | `jaxpr_primitive_count` | primitives in the step's jaxpr, recursing into `scan`/`while`/`cond` bodies | how much work the graph builder emits. A body counts once, not once per iteration, so the number describes the program and not the trip count | | `hlo_op_count` | operations in the lowered StableHLO module | what MADDENING hands to XLA. Deliberately *pre*-optimisation: the post-fusion count is the backend's decision and differs by platform | `scan_*` variants cover a `run_scan` program when `scan_steps` is passed. All of these reproduce exactly: the same graph on the same JAX version gives the same integers on any machine, under any load. Wall-clock stays in the report, worth trending, and never gates. ### Running the gate ```bash python scripts/compile_counts.py --check # what CI runs python scripts/compile_counts.py --show # print, write nothing python scripts/compile_counts.py # regenerate the baseline ``` Six workloads are measured — a single node, a coupled pair, a {term}`multi-rate ` graph, a heat chain with a `run_scan` program, a `ShardedStencilNode` over a four-device mesh, and two coupling groups with `diagnostics=True` — and compared to `benchmarks/compile_counts_baseline.json`. The script pins `JAX_PLATFORMS=cpu` and the device count before importing JAX, because both change the counts; JAX fixes its device count when its backend initialises, so `tests/core/test_compile_counts.py` runs the script as a subprocess rather than calling it. The pin **overrides** an inherited `--xla_force_host_platform_device_count`; it does not defer to it. That matters more than it sounds. `tests/cloud/multigpu/conftest.py` appends `--xla_force_host_platform_device_count=16` to `os.environ` when it is imported, so in a whole-suite run — which is what CI does — every subprocess spawned after collection inherits sixteen virtual devices, while running `tests/core/test_compile_counts.py` on its own inherits none. If the gate took the inherited value, its device count would be a function of which tests were collected beside it, and no contributor could regenerate a baseline CI would accept. Do not "fix" the pin to respect the caller. ### If the gate fails on your branch It prints every count that moved, both tables, and the regenerate command. Two cases: - **A retrace count moved.** Compared exactly, always, because it is counted in Python and does not depend on the JAX version at all. An increase is an extra XLA compile on every run; find the leaf whose dtype or weak type changes after the first step, or the argument that stopped being static, before touching the baseline. - **An op count moved.** Compared against a band: ±max(2 ops, 2%) when the running JAX matches the baseline's, ±max(10 ops, 25%) otherwise. If the change is intended, regenerate and commit the new baseline, saying in the commit message why the counts moved. - **The device count moved.** Never banded and never widened: the {term}`sharded ` workload meshes over every device, so counts taken on another topology are not comparable at any tolerance. A baseline recording a different `device_count` — or none at all — is rejected rather than compared. Regenerating fixes it, because the script pins the count itself; if you see this, the pin failed, which means something imported JAX before the script ran. Regenerate under the JAX version CI pins where you can — the baseline records the version it was taken on, and one matching CI's gets the tighter band. The counts were measured identical on JAX 0.10.2 and 0.11.0, so this costs strictness and nothing else. `CompileCounts` reads the counts after four warmup steps, not one. The retrace bugs this project has actually had did not appear on the first step: the weak-typed-seed bug traced again on step 2 when a weak leaf came back strongly typed, and a third time on step 3. Measured after a single step it looked healthy — which is how the gate's own mutation test found the flaw.