perf(amr): fuse each coarse–fine ghost fill into one gather pass - #96
Merged
Merged
Conversation
kylebeggs
force-pushed
the
perf/fused-gather-fills
branch
from
September 10, 2026 12:33
b9e9250 to
5155be4
Compare
kylebeggs
changed the title
Benchmark ResultsTime
Memory and allocations
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
kylebeggs
added this pull request to stack #99
September 10, 2026 20:44
kylebeggs
force-pushed
the
perf/fused-gather-fills
branch
from
September 10, 2026 20:49
5155be4 to
36e4109
Compare
kylebeggs
removed this pull request from stack #99
September 10, 2026 20:49
kylebeggs
added this pull request to stack #100
September 10, 2026 20:50
|
The host fill sweeps ran one broadcast per CSR term — 7 per fill in 2D, 19 in 3D — over a handful of dst cells each, so broadcast setup, not arithmetic, dominated the exchange on a refined forest. Evaluate the whole weighted sum per cell inside a single broadcast instead: the K term views are built once per fill by recursion over Val(K) down the row (closure-free, so the two-array walk keeps inlining) and ride an immutable Ref subtype as a broadcast scalar, which a mutable RefValue cannot do without escaping the un-inlined gather body and heap-allocating ~2 KB per 3D fill inside a rule seam that must not allocate. halo_update! on a refined 32³ forest: 853 → 205 µs packed and 751 → 229 µs per-block at 8³ blocks, 3.0 → 0.67 ms at 4³ (min-of-N, alternating processes, three rounds, ±4%). Term and accumulation order are unchanged, so the result is bit-identical to the retired loop whenever the field eltype is the schedule's weight type, and — unlike that loop, which rounded its partial sum into the field K−1 times — bit-identical to the device CSR kernel in mixed precision too. _close_fill now asserts the two row invariants the fusion rests on: every term window has the dst slab's shape, and every term window is disjoint from the dst slab (the views ride a non-AbstractArray scalar, so Base's aliasing machinery never sees them). The adjoint keeps its per-term scatter: a fill's term windows collide on shared source cells, so a dst-centric single pass would reorder the accumulation. The bit-identical fused route is a source-centric transposed CSR, scoped to #94. Claude-Session: https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP
…cally A bit-parity testset alone cannot see a revert: its reference is a re-written copy of the per-term loop the fusion replaces, and the zero-allocation testset passes either way. Redefining _run_fills! back to the loop left the suite green. So count what the sweep does. ProbeArray is block storage that tallies every element write and every broadcast materialized into a view of it; one phase run over it must issue exactly one broadcast per fill and write every dst cell exactly once. A per-term loop over a K-wide row issues K and writes K times, and now fails. Also added: bit-parity against the retired loop in Float32 and Float64, 2D and 3D, both storage layouts, scalar and SVector fields, forward and adjoint; a mixed-precision case (Float32 field on a Float64 forest) checked against the device CSR kernel on the CPU backend, where the fused gather is exact and the retired loop differed on 73 cells in 2D and 4420 in 3D; the row invariants _close_fill asserts, on hand-built rows and on every row the emitters produce; and zero allocation for both layouts, forward and adjoint, on refined 2D and 3D forests. Claude-Session: https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP
What the fusion buys, what it is bit-identical to (the retired loop at matched eltypes; the device kernel always, which the loop was not), the two row invariants that make it legal, the compile-latency price, and why the adjoint stays per-term until #94. Claude-Session: https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP
kylebeggs
force-pushed
the
perf/fused-gather-fills
branch
from
September 11, 2026 12:37
36e4109 to
ab2664f
Compare
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.
Closes #84.
Stacked on #42 (
perf/ghostfill-isbits) — merge that first; this branch is built on its CSR term buffers and does not revert any of it.Problem
_run_fills!wrote each coarse–fine ghost slab once per term:dst .= w₁·src₁then onedst .+= wₖ·srcₖper remaining term. An interpolation row is1 + 2·3^(N−1)terms — 7 in 2D, 19 in 3D — and there are 16 fills per coarse–fine face in 3D, so one 3D interface face cost roughly 300 separate broadcasts over slabs a few cells wide, each re-reading and re-writing the same handful ofdstcells. Broadcast setup, not arithmetic, was the cost: on a refined 3D forest the fill phases were about 90% of the whole forward exchange.Fix
One fused gather per fill:
dst[I] = Σₖ wₖ · srcₖ[I], evaluated in a single broadcast over the dst slab'sCartesianIndices. The phase's row width is a function of the dimension alone, so_run_fills!takes it as aVal(K)(_interp_width/_restrict_widthfoldVal(N)toVal(7)/Val(19)/Val(5)/Val(9)), and the K term views and weights are built by recursion down the CSR row rather than byntuple(Val(K)) do kover two arrays — the same non-inlining trap_diff_axesdocuments inoperators/diffusion.jl.Three things make that legal, and all three are now enforced rather than assumed:
GhostFilldocstring and asserted per term in_close_fillat schedule build. Disjointness became load-bearing with the fusion: the term views ride an immutableRefsubtype (_AsScalar) so that the un-inlined gather body does not heap-allocate them, and Base'sbroadcast_unaliasnever sees them as a result.@inboundswalk downterms, so a malformed row throws instead of reading out of bounds.One intentional behavioural change. In mixed precision — a
Float32field on aFloat64forest — the fused gather is not bit-identical to the retired loop: the loop rounded its partial sum into the field's eltype K−1 times, the fused body keeps the promoted accumulator and rounds once. That is the device CSR kernel's arithmetic, so the fusion removes a pre-existing host/device divergence rather than introducing one, and the claim in the code and in DESIGN.md is narrowed to say exactly that. Measured on cascaded 3-level refined forests: 155 of 6400 cells differ in 2D and 3645 of 158 976 in 3D, all within 4·eps(Float32) absolute (max |Δ| = 1.8e-7 in 2D, 4.8e-7 in 3D).Adjoint: deliberately out of scope, filed as #94
halo_update_adjoint!keeps its per-term scatter-add and is untouched here, which makes it the slower half of the exchange — 2.6× the fused forward in 2D and 4.4–4.5× on the refined 3D forests. That is a known, measured, accepted state, not an oversight.The reason it was not fused: a fill's term windows collide on shared source cells, so a dst-centric single pass would accumulate into a colliding cell in dst-cell order instead of term order and change the roundoff. Measured collisions: 40 960 of 77 824 scatter-adds collide for 3D 32³/8³ interpolation, and 0 for every restriction phase and for 3D 32³/4³ interpolation. Fusing only the collision-free rows was considered and rejected — it would make the exchange's cost depend on the blocksize in a way nothing in the API hints at. The bit-identical route is a source-centric transposed CSR built at schedule time: cell-granular, ~78 k contributions for one 3D interpolation phase, rebuilt on every regrid, and unable to precompute linear indices because
BlockFieldandPackedBlockFieldaddress the same cell differently. That is the descriptor bloat #42 just removed, so it needs its own design pass. Issue #94 carries the collision data, the lopsidedness numbers, and the acceptance criteria.Tests
test/exchange_schedule.jlgains aProbeArray— block storage that tallies every element write and every broadcast materialized into a view of it — and asserts that one phase run issues exactly one broadcast per fill and writes every dst cell exactly once. This is a structural guard, not a numeric one: I verified it bites by@evaling_run_fills!back to the per-term loop in the loaded module and re-running the file. It fails with 26 assertions — the per-fill broadcast count, the per-cell write count, the row-width throw, and the device-kernel parity below — where the base branch's suite passes unchanged under the same reversion.Also added: the mixed-precision case, comparing the fused host fills against
_run_fills_device!on the CPU backend for aFloat32field on aFloat64forest. The fused result matches the device kernel exactly; the retired loop does not, and the test asserts both directions. And_close_fillinvariant tests — hand-built rows that are the wrong width, the wrong shape, or that read their own dst all throw, and both invariants are re-checked on every row the real emitters produce in 2D and 3D.test/forest_parity.jl,test/exchange_kernels.jl,test/forest_amr.jl,test/forest_packed.jlandtest/blockforest.jlall pass unchanged. Independently of the suite, I byte-compared the exchanged padded storage (ghosts included) againstmainfor both layouts, forward and adjoint, on refined 2D 256²/32², refined 3D 32³/8³, refined 3D 32³/4³ and cascaded 3-level 2D/3D forests: identical everywhere, with the sole exception of the documented mixed-precision forward case above.Performance
Method: Apple M5 Pro, macOS (Darwin 25.6.0), Julia 1.13.0 aarch64,
-t1, machine otherwise idle (no other Julia processes at any point). Three project environments, one per revision, eachPkg.developing a separate worktree. One fresh process per (revision, round); rounds run main → #42 → this branch, five times. Within a process: 3 warmup calls, then the minimum over 60 reps of five calls,GC.gc(false)between reps. The table is the minimum of the five per-round minima. Round-to-round spread (max/min of the per-round minima) is 2.5% median, 7.2% at p90, 9.2% worst — the forward margins below are 2.7× to 7.2×, far outside it. Forests arerefine!(bf, p -> p[1] < 0.5), maxlevel 2. The#42column is this branch's base, so the middle ratio is what this PR contributes.BlockFieldhalo_update!BlockField_run_fills!BlockFieldhalo_update_adjoint!BlockField_run_fills_adjoint!halo_update!_run_fills!halo_update_adjoint!_run_fills_adjoint!BlockFieldhalo_update!BlockField_run_fills!BlockFieldhalo_update_adjoint!BlockField_run_fills_adjoint!halo_update!_run_fills!halo_update_adjoint!_run_fills_adjoint!BlockFieldhalo_update!BlockField_run_fills!BlockFieldhalo_update_adjoint!BlockField_run_fills_adjoint!halo_update!_run_fills!halo_update_adjoint!_run_fills_adjoint!So the forward fill sweep is 2.7× to 7.2× faster and the whole forward exchange 1.6× to 4.0× faster, with the largest wins on packed 3D — the layout and dimension the AMR work is aimed at. The adjoint legs move by at most 2%, which is inside the spread, and nothing regresses against
mainon any path.Allocations stay at 0 bytes per call for both sweeps, both directions, both layouts, on all three forests — asserted in the suite with
@inferredon both directions, and confirmed here out of band (@allocatedover a 200-call loop, divided). The_AsScalarimmutableRefis what buys this: a mutableRefValueescapes the un-inlined gather body and heap-allocates the whole tuple of views on every fill, inside the_exchange_storage!rule seam that must not allocate.Schedule build cost is unchanged despite the new per-term disjointness assertion: 0.56 / 17.8 / 73.8 ms on the three forests, against 0.58 / 17.9 / 75.5 on #42 and 0.59 / 18.6 / 77.2 on
main.Compile latency
This is the real cost, and it is not small. First
halo_update!on a refined 3D 32³/8³ forest in a fresh process, two runs each:BlockFieldmainRoughly a doubling, about +550 ms one-time per specialization, inherent to specializing the gather body over
Val(K). Against a forward saving of ~570 µs per call on that forest, it pays for itself after about a thousand exchanges — fine for a Krylov or multigrid solve, a net loss for a script that exchanges a few dozen times. The firsthalo_update_adjoint!after it also rises, 205 → 306 ms, since it now follows a heavier forward compile.Not measured
MFO_TEST_GPUunset. The one path that reaches this host broadcast on a GPU is a refinedBlockFieldof device arrays; it ships K deviceSubArrays as kernel parameters, about 2.4 kB for a 19-term 3D interpolation fill against CUDA's 4 kB parameter budget on pre-sm_90 hardware. It fits with under 2× headroom and none for a wider stencil or a 4th dimension, and that is documented beside_AsScalar. Packed fields are unaffected — they take the batched CSR kernels. The_AsScalarAdapt rule has only been exercised against a hand-rolled adaptor.test/enzyme_rules.jl46/46 in 3m47 against 3m36 on perf(amr): make GhostFill isbits with an NTuple of terms #42, andtest/autodiff.jl54/54 in 4m49 against 4m44. So the latency does not propagate there. But that is Rosetta, not native Linux, and per CLAUDE.md a local green run is not evidence about CI — read the matrix.ProbeArraybroadcast count runs theBlocksLayoutonly.PackedBlockFieldgoes through the same_run_fills!method so the count applies to it, but it is not separately counted.https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP