perf(evolution): ⚡ stop retaining the P-sized exchange state per layer

[diagonal-hamiltonian]

🤖 AI text below 🤖

The defect

Inside the engine rank_count is mpi::size(comm). On a partitioned run the comm is
Kind::Hybrid, whose size() returns the flat world P = mpi_ranks x partitions
(HybridComm.h:81), not the MPI rank count. So every "per-rank" array in a graph layer is
P long, each MPI process holds one per partition, and the graph retains one per layer.
Per-slot state therefore costs O(P²) across the job.

Measured on the pauli kicked-Ising anchor, main @ 6abd839, with 91,273,861 terms and
an identical config on every rung — one problem measured four times, not four problems:

P = ranks x partitions	geometry	operator	graph	summed peak RSS
16	1x16, 1 node	6.93 GB	3.81 GB	12.10 GB
128	8x16, 1 node	6.97 GB	7.61 GB	18.72 GB
256	8x16, 2 nodes	6.61 GB	18.62 GB	35.79 GB
512	8x16, 4 nodes	6.52 GB	61.94 GB	97.99 GB

The operator is flat — it partitions perfectly. The graph is not. At P=512 the quadratic term
is 93% of the graph.

It tracks P, not the MPI rank count

This is the observation that could have refuted the diagnosis. Holding P fixed and moving the
rank count across its whole range:

P	geometries	graph total_bytes
16	1x16, 8x2	366,649,352 — identical
128	1x128, 8x16, 2x64, 4x32	2,769,389,680 — identical

Six geometries, two values. A single MPI rank at 1x128 pays 7.6x what the same single rank
pays at 1x16. The rank-count reading predicts that cell is cheap; it is the expensive one.

The layer count falls out of the instrument independently: 86,720 layer-cores over 16 partitions
is 5,420 layers = 20 x (127 X + 144 heavy-hex ZZ), the circuit's gate count.

What was retained per layer, and what is left

Per world slot, per layer, per partition:

structure	main	after	why
CrossRankPartnerRange	32 B	16 B	held the same offset/count pair twice
evolution counts + displs	8 B	0	a prefix sum of what the records already say
derivative layout (counts, displs, recv)	16 B*	0	the evolution layout at scale 2
evolution recv_cache	8 B	0	the recv layout is the send layout

_{* lazily allocated, so only after a gradient has run.}

Plus a per-layer cost, not per-slot: sizeof(LayerCore) falls 416 B -> 168 B (measured),
because two LayerExchangeLayout structs, an embedded RecvLayoutCache and the generation id that
guarded it all left the struct.

So the retained slot-proportional footprint goes 40 -> 16 B/slot counted — 48 -> 16
counting what the ledger does not — plus a second, linear-in-P term of 248 B per layer-core.

The measured saving understates the real one. The derivative layout was a diagnostic outside
total_bytes(), so its 16 B/slot never appeared in graph and its removal cannot appear there
either. Counting what the ledger does not, at P=512 (d_ fields, port arm, measured):

uncounted, per slot	main	after
recv_cache	8 B	0
derivative layout	16 B	0

At P=512 that is 31.76 GiB of uncounted state on main against 0 now — none of which the
tables below can show. d_recv_cache_bytes was measured at 11,366,563,840 B (exactly 8 B/slot)
after the second commit and 0 after the third, both on the port arm. The main figure is
computed from the struct layout, because main has no such binding, so treat that one as an
estimate rather than a measurement.

It is corroborated independently: gradient peak transient memory on the worst rank falls
0.88 -> 0.04 GiB at P=512, and a per-rank derivative layout at 16 B/slot works out to ~0.66 GiB.
Same order, from a metric that knows nothing about the ledger.

Nothing sized by P is retained per layer any more. What is left of the quadratic term is the
slot records themselves, at 16 B/slot.

The three commits

1d62773 — the record was storing the same range twice. CrossRankPartnerRange carried an
offset and a count for each of B and D. They are always equal: GraphSink::finalize resizes
both vectors from the same P + Q expression, so the counts match per slot and their prefix sums
match with them. That equality is now a checked precondition rather than a comment — unchecked
it would not throw, it would mis-derive Q and read a wrong-but-valid endpoint.

5f33a71 — derive the layouts instead of retaining them. counts[r] is
(r == my_rank ? 0 : cross_rank.sin_send_size(r)) and displs is its running prefix, so both are
derived into the per-thread scratch that already owns the send and recv buffers.

The derivative round needs no collective of its own. Its counts are the evolution counts at a
hardcoded scale of 2 applied identically on every rank, and displacements are prefix sums of
counts, so scaling commutes with the transpose: the derivative recv layout is 2x the evolution
recv layout. main ran a separate resolve_recv per layer for the derivative round, each with
its own cache. One resolve_recv per layer per evaluation now serves both rounds.

The predicate this needed first

Sharing scratch across layers is only sound once resolve_recv can tell one send pattern from
another. Its predicate was

if (cache.comm_size == comm_size && cache.layout.counts.size() == n) return cache.layout;

— effectively "have we ever resolved anything for a communicator this size", true for every layer
after the first. Correct only while each cache belonged to the one layout that produced it.

The fix is not a checksum of the counts. A miss runs alltoall_counts, a collective, so two
ranks disagreeing about validity is a distributed hang, not a wrong answer — and any rank-local
key can collide on one rank and not on another. The cache now carries a generation assigned per
LayerCore at build time. Build order is identical on every rank, so every rank misses on a
layer's first resolve and hits afterwards: the decision is rank-uniform even though the id
values are not.

The third commit deletes all of this. With no cache there is no miss path, so there is no
collective for ranks to split on and nothing for a generation id to guard. The hazard is removed
rather than managed. The reasoning is kept here because it is what the second commit shipped, and
because it is the argument that had to be right for the third one to be reachable.

Hoisting the slot resolution

Resolving a world slot is an index into the P-sized ranges array, and the element accessors were
doing it per endpoint — three times per term on the recv side, four times per rotation pair in
the self-slot gradient loop. cross_rank_slot() resolves once and the accessors take that view.
No behaviour change, and it is the precondition for ever storing slots sparsely.

cff597e — the transpose was the send layout all along

The second commit kept a RecvLayoutCache per layer on the grounds that a transpose is the one
thing a rank cannot work out alone. That was wrong, and reading the layer-build sink says why:
slot r on rank m holds the queries r sent m, followed by the queries m sent r; rank
r's slot for m holds those two swapped. The counts are therefore equal, and displacements are
prefix sums of counts, so the recv layout is the send layout. MPI_Alltoallv reads
recvcounts/recvdispls rather than writing them, so the same two arrays now serve both sides.

Deleted with it: resolve_recv, RecvLayoutCache, RecvLayout.h, the per-layer
alltoall_counts on the miss path, and exchange_generation.

Symmetry is an invariant of the routing, not of this file, so it is checked where it can actually
break. MONOPROP_CHECK_EXCHANGE_SYMMETRY=1 re-adds the alltoall and throws naming the slot and
both counts. Unguarded, a routing change that broke the invariant would surface as a peer blocked
in MPI_Alltoallv against a size nobody sends — a hang with no line number. It is off by default
because on, it costs exactly the collective the commit removes.

Measured effect

Interleaved A/B, one allocation per cell, order flipped per (rep, cell), 6 reps,
--bench-rounds=1, ratio = median of per-rep paired ratios. main 0.8.1.dev37+g6abd839e6
(= origin/main) against the port build of this branch. Arms confirmed distinct by .so hash
(f73d17f2… vs d37b0f1e…), not by version string. Term counts identical across arms on every
cell; every cell pinned at distinct_pinned_cpus == PARTITIONS.

Memory — the graph ledger

cell	P	main	port	saved	ratio
c12 1x16	16	1.442 GiB	1.397 GiB	0.045 GiB	1.03x
c12 8x2	16	1.442 GiB	1.397 GiB	0.045 GiB	1.03x
c12 1x128	128	4.982 GiB	2.888 GiB	2.093 GiB	1.72x
c12 8x16	128	4.982 GiB	2.888 GiB	2.093 GiB	1.72x
c14 8x16, 1 node	128	7.091 GiB	4.998 GiB	2.093 GiB	1.42x
c14 8x16, 2 nodes	256	17.340 GiB	9.183 GiB	8.157 GiB	1.89x
c14 8x16, 4 nodes	512	57.686 GiB	25.494 GiB	32.192 GiB	2.26x

The saving is predicted with nothing fitted, and it is exact. Job-total slots = L·P² and
layer-cores = L·P, so

saved = 24 B x slots  +  168 B x layer_cores

cell	P	predicted	measured	ratio
c12	16	47,869,440 B	47,869,440 B	1.00000
c12	128	2,247,782,400 B	2,247,782,400 B	1.00000
c14	128	2,247,782,400 B	2,247,782,400 B	1.00000
c14	256	8,758,026,240 B	8,758,026,240 B	1.00000
c14	512	34,565,898,240 B	34,565,898,240 B	1.00000

To the byte, across a 4x range in P and a 3.1x range in term count, with zero free parameters.
The 24 B/slot is the two record fields plus counts/displs; the 168 B/core is
sizeof(LayerCore) falling 416 -> 248.

That second term is why the per-slot saving looks inconsistent (34.50 B/slot at P=16 against
25.31 B/slot at P=128): the linear term is a larger share of the total when P is small. It is not
noise, and the earlier one-term reading of it was wrong.

The operator is byte-identical between arms on every cell (port/main = 1.0000). The change is
confined to the graph.

Time

* marks 6/6, which clears a sign test at p=0.031. 5/6 is p=0.109 and does not clear, so it is
reported but not claimed.

cell	P	build_graph	propagate	energy	gradient
c12 1x16	16	1.00x (4/6)	1.03x faster (5/6)	1.04x slower (6/6*)	1.01x slower (6/6*)
c12 8x2	16	1.01x slower (5/6)	1.00x (4/6)	1.04x slower (6/6*)	1.02x faster (6/6*)
c12 1x128	128	1.23x faster (5/6)	1.11x faster (3/6)	1.01x faster (5/6)	2.10x faster (6/6*)
c12 8x16	128	1.02x slower (5/6)	1.01x faster (3/6)	1.00x (3/6)	1.15x faster (6/6*)
c14 N1	128	1.02x faster (6/6*)	1.00x (3/6)	1.01x slower (4/6)	1.03x faster (5/6)
c14 N2	256	1.02x faster (5/6)	1.00x (4/6)	1.06x faster (6/6*)	1.13x faster (6/6*)
c14 N4	512	1.01x faster (5/6)	1.01x faster (4/6)	1.26x faster (6/6*)	1.43x faster (6/6*)

gradient is the operation this moves, and the win grows with P — 1.15x at P=128, 1.13x at
P=256, 1.43x at P=512, and 2.10x at 1x128 where 128 partitions sit inside one rank (per-rep ratios
0.466–0.486, the tightest cell in the campaign). That is the derivative round no longer resolving
its own transpose: main paid one extra alltoall_counts per layer on the first gradient.

Transient memory on gradient falls by more than the ledger does, because the retained derivative
layout is never allocated at all. Peak dmem on the worst rank, port against main:

cell	P=16 (1x16 / 8x2)	c12 P=128 (8x16 / 1x128)	c14 P=128	P=256	P=512
gradient dmem	1.06x / 1.04x	3.95x / 4.58x	2.28x	6.83x	20.53x

It is not monotone in P alone — c14 at P=128 sits below c12 at the same P — because the cell's
own working set scales with the term count while the removed layout scales with P.

build_graph and propagate are flat everywhere. The 1.23x on build_graph at 1x128 is the
largest number in the table and the one I would most like to claim; it is 5/6 and not a result.

The honest negative

At P=16 the change is a small loss: energy is 1.04x slower on both geometries at 6/6. That
is the expected shape rather than a surprise — deriving counts is a fixed cost paid per exchange,
while the saving grows with P, so a small world pays the cost without earning the benefit.

gradient at P=16 is unresolved: 1.01x slower on 1x16 and 1.02x faster on 8x2, each at 6/6, in
opposite directions. Two geometries at the same P disagreeing in direction while both reach 6/6
means the effect is per-geometry, not per-P, and neither figure should be read as the P=16
gradient result.

The third commit, measured on its own

main -> HEAD mixes all three commits, and two independently measured ratios cannot be differenced
into an increment. So a third arm was built from the second commit's binaries and run against HEAD
under the same protocol (arms 5ad94fee vs 97dc2277 by .so hash).

Memory, against origin/main, cells gws3-*:

	P=128 (c12)	P=512 (c14, N=4)
main	4.982 GiB	57.686 GiB
after 5f33a71	2.888 GiB	25.494 GiB
after cff597e	2.837 GiB	25.288 GiB
main -> HEAD	1.76x	2.28x (32.399 GiB saved)

The increment was predicted before the cells ran, and the collator asserts it rather than
printing it: 80 B per layer core, and nothing per slot. Measured 80.00 B/core at both cells,
ratio measured/predicted 1.00000. The per-slot half is the load-bearing one — the cache sat
outside total_bytes(), so any per-slot movement would have meant the model of that field was
wrong. cores / L = P and slots / cores = P came out exact at both cells for L = 5,420.

The larger half moves no shipped metric: d_recv_cache_bytes 0.662 -> 0 GiB at P=128 and
10.586 -> 0 GiB at P=512, measured at exactly 8.00 B/slot.

Time, port-vs-port, 6 reps, same protocol:

operation	P=128	agree	P=512	agree
build_graph	1.12x faster	6/6*	1.05x faster	6/6*
propagate	unresolved	5/6	flat	3/6
energy	flat	5/6	1.14x faster	6/6*
gradient	1.05x faster	6/6*	1.17x faster	6/6*
build_graph dmem	1.13x smaller	—	1.29x smaller	—

The mechanism is 5,420 alltoall_counts per first evaluation that no longer happen. build_graph
dmem falling is what locates the old cache's allocation in the build rather than the first eval.

A second honest negative, 6/6 and unexplained. gradient dmem at P=512 is 1.16x LARGER
after this commit, 0.04 -> 0.05 GiB. It is ~10 MiB against a 10.6 GiB resident saving in the same
cell, and gradient time is 1.17x faster, so it is not a blocker — but it agrees on all six reps,
which is not noise, and I have no mechanism for it.

ab_summary refused both port-vs-port cells, correctly by its own rule and wrongly in fact:
both arms are the same worktree, so they report the same monoprop_version. That is the blind spot
this PR's own instrumentation notes describe — a version is a git describe, not a fingerprint of
the extension. The arms were confirmed distinct by .so md5, and each venv confirmed to import its
own, before submission. The table is quoted on that basis and no other.

Instrumentation

graph_memory_breakdown() gives the graph a per-field split, and reports slot occupancy — the
fraction of P slots carrying any traffic — which nothing reported before.
exchange_layout_bytes and derivative_layout_bytes now report 0 rather than being removed,
so an A/B against a build that did retain them shows the drop instead of losing the row.

recv_cache_bytes stays a diagnostic outside total_bytes(), so graph_memory_bytes() means
the same thing before and after.

Correctness

Full gate on 5f33a71, MPI build, MAX_NUM_MODES=1024, built with ninja reporting no work to do (up to date against every declared dependency, not an mtime guess):

gate	result
C++ suite, whole binary	226 cases, no errors
ctest -L serial	214/214 passed
partition_* (in-process cross-rank gradient cover)	10 cases, no errors
graph_encoding_*	17 cases, no errors
Python suite, 1 rank	617 passed, 8 skipped
Python --with-mpi, 2x1 / 2x8 / 4x1 / 4x8 / 8x1 / 8x8	625 passed each

Re-gated on cff597e (job 1826413, md5 97dc2277, the exact binary these numbers come from):
214 ctest -L serial, and 625 Python tests on each of four geometries twice — world 2, 32,
32 and 256, once on the production path and once with MONOPROP_CHECK_EXCHANGE_SYMMETRY=1
asserting the recv-equals-send equality on every layer. The second pass is the one that
carries weight; the first only says nobody noticed.

Independently of that, a temporary probe comparing derived counts against a real alltoall on every
resolve saw 0 mismatches in 550M slot comparisons at world 32 and 256, over the full MPI suite
and a pauli c12 energy+gradient run.

The claim of the derivation is equivalence, so it is tested against an independent
construction rather than literals: graph_encoding_derived_layout_matches_the_layout_it_replaces
asserts the derived layout equals build_layer_exchange_layout's output elementwise — counts,
displs and total_count — for every my_rank and both scales. build_layer_exchange_layout now
has no production caller and is kept deliberately as that oracle, which the header says so a
later reader does not delete it as dead code.

Also covered: a zero-traffic slot still gets a valid non-decreasing displacement (where an
off-by-one in a prefix sum would hide), the scratch is reused rather than reallocated, the
derivative overflow throws at build time, and a built layer retains no exchange layout at all.

Notes for review

• This overlaps perf/sparse-hot-path. That branch is refactoring MPICompat.h into leased
  alltoallv buffer sets (begin_alltoallv / wait_into). Both changes answer "who owns the
  buffers for an exchange": this one derives counts into a thread_local scratch, that one leases
  a set per call. Whoever merges second should move the derivation inside the lease. Note the
  third commit deletes resolve_recv outright, so if that branch grew a caller for it after I
  last looked, this is a real conflict rather than the clean one it was.
• The thread_local scratch is sound only if at most one exchange is in flight per thread,
  which MPI_Ialltoallv requires anyway (the send counts must stay valid until the wait, and
  HybridComm/ShmComm publish the raw pointer for peer partitions to read across barriers).
  That invariant is not new — send_buffer has always required it.
• Build-time derivation is retained purely as eager validation: an int overflow has to throw
  from build_graph, not from inside the exchange where peers are already committed to a transfer
  of that size.
• The symmetry invariant is the one thing to push on in review. It is a property of how layer
  build assigns endpoints to slots, not of the exchange code, so a future routing change could
  break it from a distance. That is why the audit exists and why I would run the multi-rank CI
  geometries with MONOPROP_CHECK_EXCHANGE_SYMMETRY=1 set.

Corrected from the earlier version of this PR

The earlier description said sparsifying ranges was the largest remaining lever, "worth about 3x"
the record shrink. That was wrong. Occupancy of ~25% bounds what sparsifying could reclaim as a
fraction of itself; it says nothing about its size beside the other levers. In bytes per slot,
sparsifying ranges saves ~12 B — less than the 16 B the record shrink already shipped — and it
is the one option that changes a structure three consumers index by global slot. The retained
layouts were the larger prize and needed no such change.

It also said recv_cache was "the one piece that genuinely cannot be derived — it takes a
collective". Also wrong, and shipped as the third commit. I asserted it from the general fact
that a transpose usually needs a collective, without checking whether this particular one carries
data both sides already hold.

Occupancy at the measured sizes, which does still pick the next fix: 31.5% (c12 P=16), 23.9%
(c12 P=128), 27.8% (c14 P=128), 17.5% (c14 P=512).

None of this removes the P² term, and it is worth being plain about that. What remains is
16 B/slot of slot records, of which the 8 B sin_send_offset is itself a prefix sum and derivable
by the same trick. The floor for a dense layout is 8 B/slot; sparse at ~25% occupancy is ~3.4 B.
Only dropping the L = 5,420 retained-layer factor — streaming or checkpointing the graph —
makes the quadratic irrelevant, and that trades it for forward-rebuild time on the gradient's
reverse pass.

---

Results are not reproducible from this diff: the benchmark harness, the campaign driver and the
symmetry probe all live on an unmerged branch, and the numbers come from 11 cells plus 3 gate jobs
on up to 4 Deucalion nodes.

[github-actions]

Docs preview: https://pr-237.monoprop-docs.pages.dev