feat(distributed): compact diffusion leaf on partition slabs - #66
Open
kylebeggs wants to merge 1 commit into
Open
feat(distributed): compact diffusion leaf on partition slabs#66kylebeggs wants to merge 1 commit into
kylebeggs wants to merge 1 commit into
Conversation
|
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. |
Stage 2 of #56, closes #57. Diffusion joins the prepare_distributed whitelist on the same coefficient terms as ScalingOp. The face averaging reads κ one cell past every face, the partition cut included, so _slab_field's zero coefficient ghosts do not serve it. It still needs no communication: _slab_op runs on the host holding the global κ, whose ghosts diffusion already extended by an even mirror and a periodic wrap, and slab padded index p is global padded index first(local_range[d]) - 1 + p. One padded window — _slab_coeff_field — therefore lands every ghost on the value it should hold with no per-face logic: an Interface ghost onto a global interior plane (the neighbour's κ), a wall ghost onto the global mirror, a periodic cut onto the global wrap. Widening the slice is a setup-time indexing change, not a transport, so the per-apply exchange count is unchanged from the Laplacian baseline. _slab_field's invariant is narrowed, not broken — it still holds verbatim for ScalingOp, which reads its coefficient pointwise. The public diffusion(g, κ) keeps rejecting Interface faces; _slab_op builds through the inner constructor, which is the seam that docstring reserved. This is also the first production path to reach the leaf's mechanical transpose: a slab carries Interface faces, so apply_adjoint! takes the adjoint_gather! branch rather than the self-adjoint shortcut.
kylebeggs
force-pushed
the
feat/distributed-diffusion-leaf
branch
from
August 14, 2026 16:31
5f37161 to
437a878
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.
Stage 2 of #56. Closes #57. Stacked on #60 — review that first; this PR's diff is only the distributed layer.
Diffusionjoins theprepare_distributedwhitelist on the same coefficient terms asScalingOp: a real-eltypeFieldon an undistributedCartesianGrid, under either averaging policy.It needs no communication
The face averaging reads κ one cell past every face, the partition cut included, so
_slab_field's zero coefficient ghosts do not serve it. But_slab_opruns on the host holding the global κ, whose ghostsdiffusionalready extended at construction by an even mirror and a periodic wrap. Slab padded indexpis global padded indexfirst(local_range[d]) - 1 + p, so one padded window lands every ghost on the value it should hold, with no per-face logic:Interface(cut plane)The window is
first(zr):(last(zr) + 2h) ⊆ 1:(n + 2h), so it never needs clamping. Widening the slice is a setup-time indexing change, not a transport, and the per-apply exchange count is unchanged from theLaplacianbaseline — #56's "one halo exchange per application" invariant holds trivially. That is legitimate only because κ is constant through a solve; a coefficient that changed per iteration would owe an exchange per iteration._slab_field's invariant is narrowed, not broken. It still holds verbatim forScalingOp, which reads its coefficient pointwise;Diffusionopts into the widened_slab_coeff_fieldinstead.Deviation from the issue. #57 sketches
_slab_coeff_fieldas an interior slice, thenfill_coefficient_ghosts!, then an explicit per-face copy of neighbour interior planes with its own periodic-wrap branch. The padded window above produces the same values in six lines with no branches, so that is what landed.What did not change
diffusion(g, κ)still refusesInterfacefaces._slab_opbuilds through the inner constructor, which is the seam its docstring reserved. Its error message now points atprepare_distributedinstead of claiming slabs are unsupported.operator_diagonal(::Diffusion)keeps itsInterfacerejection — it serves multigrid, which is not distributed.This PR is also the first production path to reach the leaf's mechanical transpose: a slab carries
Interfacefaces, soapply_adjoint!takes theadjoint_gather!branch that scatters cotangents into ghosts for the slab reduction to fold, rather than the self-adjoint shortcut an all-physical grid takes. That code landed in #48 with a hand-built test; here it runs for real.Coverage
test/partitioning.jlgains 2514 assertions over 1-D/2-D/3-D grids (the cut always falls on dimension N, so grid rank is how "a cut in each dimension" is reached), both averaging policies, Dirichlet/Neumann and periodic cuts, and 2 and 3 partitions: padded-slice exactness cell by cell against an independently spelled oracle, bitwise forward parity, the dot-product adjoint identity, dense forward/transpose structure viadist_materialize,boundary_rhsparity, and an assertion that the exchange gating matches theLaplacian's node for node.Every positive claim is paired with a control. The load-bearing one zeroes only the cut-plane ghosts of each slab's localized κ — exactly what
_slab_fieldwould have left — and demands the answer move. Verified by sabotage: swapping_slab_coeff_fieldfor_slab_fieldin_slab_opturns 3372 passes into 3040 passes / 332 failures, concentrated in the dense-parity, ghost-poison, andboundary_rhstestsets. The κ used throughout varies across the cut and is asymmetric about it; a constant κ would hide a zeroed ghost entirely.test/mdla_gpu.jlgains the extension-path twins — accepted list, a padded-κ upload check, forward parity, adjoint identity and dense structure,boundary_rhsparity, and-1.0 * diffusion(g, κ)as an SPD system in the CG end-to-end. Those were not run locally:test/mdla.jlis gated onMFO_TEST_MDLA=truewith ≥2 CUDA devices, and this machine has neither. Their tree-shape and padded-window assertions were validated against the CPU harness, but the GPU matrix is the real check.Local: full
Pkg.test()green — 8163 pass, 3 pre-existing broken, 0 fail.🤖 Written by Claude, who slid one index window sideways and called it a distributed algorithm.