[WIP] Fewer guard cell exchanges when applying Darwin linear operator - #7177
Open
RemiLehe wants to merge 2 commits into
Open
[WIP] Fewer guard cell exchanges when applying Darwin linear operator#7177RemiLehe wants to merge 2 commits into
RemiLehe wants to merge 2 commits into
Conversation
The Darwin linear operator
A Z = bilaplacian(Z) + curl(chi curl(Z))
is built from three short, local stencils, so the composition is itself a
single banded stencil on Z. It was evaluated as three separate passes with a
FillBoundaryAndSync round between each (five rounds per GMRES iteration in
total, one of which was redundantly repeated), using dA_fp and Efield_fp as
scratch space.
Instead, compute each intermediate over enough of its own guard region to
feed the next stage, so the only communication left is the single exchange
on Zscratch. This is the same trade the hybrid-PIC solver makes in BLAST-WarpX#6571,
over a longer chain (curl -> mass matrices -> curl). Besides being cheaper,
it removes a class of bug: an intermediate exchange reconciles the
duplicated periodic-image cells of the nodal components independently at
each stage, which is what made the two-pass laplacian(laplacian(Z)) develop
a parasitic sign-alternating mode.
- ComputeCurlA/ComputeCurlB take an optional `ngrow` argument (default zero,
so existing call sites are unchanged) that fills that many guard cells of
the output. They assert that the input, and the EB update flags, hold
valid data over the region they are read on.
- ImplicitSolver::MassMatricesStencilHalfWidth() reports how far the mass
matrix gather reaches, which is what sets the guard widths.
- The operator now owns its scratch, so dA_fp and Efield_fp are no longer
clobbered during the solve.
- Both remaining exchanges of intermediates are dropped: the duplicated
nodal cells now agree because they are computed by identical operations on
identical inputs, rather than by being reconciled after the fact.
ApplyMassMatrices() needed no change: it already fills as many guard cells
of its output as it and the mass matrices have, so the ghost width of the
output scratch is the control knob. The mass matrices themselves are already
valid in their guard region (SumBoundaryJ passes dst_nghost = nGrowVect()).
Verified to be a pure restructuring: test_1d_darwin_solver_em_modes_picmi
gives checksums bit-for-bit identical to an unmodified build of BLAST-WarpX#6293.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
RemiLehe
force-pushed
the
darwin-fused-operator
branch
from
August 20, 2026 12:48
1c31398 to
7208a92
Compare
Fusing the operator reorders floating-point work: guard cells that used to be communicated are now computed, sums land in a different order, and OverrideSync selects among values that are equal in exact arithmetic but not bitwise. The result drifts by roundoff, which over 50 steps of GMRES reaches 1.7e-9 on Ez in test_2d_darwin_solver_em_modes_es_picmi - just past the 1e-9 tolerance. Every other key already agreed to ~10 digits, and the 1D test passes unchanged. Values taken from the CI run of this branch (Azure build 5977). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
RemiLehe
force-pushed
the
darwin-fused-operator
branch
from
August 20, 2026 15:48
7208a92 to
cfe7f98
Compare
Member
|
I just looked over these changes and they all look good to me. I think this will significantly help with the parallel scaling of this solver. |
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
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.
Motivation
The linear operator that GMRES applies on every iteration of the Darwin solve is
where
chiis the deposited mass matrix scaled by2 mu_0 / dt. Each stage of that expression is a short, local stencil, so the composition is itself a single banded stencil onZ, with a predictable width set by the curl stencils and the mass matrix stencil.DarwinLinearFieldOperator::apply()nonetheless evaluated it as three separate passes with aFillBoundaryAndSyncround between each.Instead, each intermediate is now computed over enough of its own guard region to feed the next stage, so the only communication left in the whole operator is the single exchange on
Zscratchat the top.This is the same trade the hybrid-PIC solver already makes in #6571 — allocate a few more guard cells, compute into them with
mfi.tilebox(type, ngrow), and delete theFillBoundarythat followed the curl. Here the chain is longer (curl → mass matrices → curl), so the saving is proportionally larger.Advantages
Fewer guard-cell exchanges. Per GMRES iteration, from 5 rounds of
FillBoundaryAndSync(Zscratch,dA_fp,Efield_fptwice — one of those was already redundant — and the result), i.e. 15 calls, down to a single round of 3. Since the operator is applied once per Krylov vector, this is the dominant communication cost of the field solve.No global scratch fields. The operator no longer borrows
dA_fpandEfield_fpto hold its intermediates, soEfield_fpkeeps holding-grad(phi)for the whole solve rather than being clobbered and restored. A Python callback or diagnostic firing during the solve no longer observes garbage. The scratch the operator does need is owned by itand sized once in
define().Fewer kernel launches, and one less full-array
setVal(0)per intermediate.Main changes
ComputeCurlA/ComputeCurlBtake an optionalngrowargument (defaulting tozero, so every existing call site is unchanged) that fills that many guard cells of the
output, via the
mfi.tilebox(type, ngrow)overload — which grows only tiles sitting atthe boundary of the valid box, so tiles never overlap. Both assert that the input, and
the EB update flags, hold valid data over the region they are read on. The RZ and
spherical
ComputeCurlAvariants assertngrow == 0rather than pretending to supportit.
ImplicitSolver::MassMatricesStencilHalfWidth()reports how far the mass matrixgather reaches —
maxover the nine blocks ofncomp_ab[dir]/2, which is the furthestcell read regardless of the relative staggering of the two components.
DarwinLinearFieldOperatorowns the intermediates (m_dA_*,m_chidA_*) alongsidem_Zscratch_*/m_lapZ_*, with the guard widths derived indefine()by walking thestencil chain backwards from the output:
chidAcurl_Areadsi..i+1dA1 + half_widthApplyMassMatricesclamps its gather to the input's guard regionZscratchmax(2, dA_ng + 1)curl_Breadsi-1..i; thenabla^4stencil readsi-2..i+2For the common case of shape 1 with direct deposition in 3D,
half_widthis 2, soZscratchgoes from 2 ghost cells to 4. One of those is slack rather than necessity:the true stencil half-width of the composed operator is 3, and the extra cell exists
only because
ApplyMassMatricesclamps its gather to the input's allocated guardregion rather than to a caller-supplied validity box. Tightening that would be a
natural follow-up.
SemiImplicitDarwin::Define()callsInitializeMassMatrices()before defining thelinear operator, since the operator now sizes its scratch from the mass matrix stencil.
Its only other dependency is
m_nlsolver, and only whenm_use_mass_matrices_pcistrue, which the Darwin solver sets to
false.SemiImplicitDarwin::ApplyScaledMassMatrices()no longer ends with aFillBoundaryAndSyncon its output.ApplyMassMatrices()was never communicatingthose guard cells — it computes them, from the same wide-stencil read of
dAit usesfor the valid region. All it needs is enough valid input, which the guard-cell budget
now guarantees. (The old code also exchanged the same field a second time immediately
afterwards in
apply(); one of the two was already dead.)The closing
FillBoundaryAndSynconrhs_vecinapply()is gone too. Its commentclaimed that "nothing guarantees the stencil evaluations above produced identical values
at the duplicate periodic-image cells" of the nodal components — which mattered because
GMRES builds every subsequent Krylov vector from this result with element-wise
arithmetic that has no notion of that duplication. Something does guarantee it now:
everything in
apply()is evaluated fromZscratch, which was synced, and from massmatrices that
SumBoundaryJleaves equal at duplicated nodes, so the two cells arecomputed by identical operations on identical inputs. This was not free —
rhs_vechaszero ghost cells, so the call was purely the
OverrideSynchalf, i.e. realcommunication once per GMRES iteration.
ApplyMassMatrices()itself needed no change: it already fillsmin(out ng, S ng)guardcells of its output, so the ghost width of the output scratch is the control knob. The
mass matrices are likewise already valid in their own guard region, since
SumBoundaryJpasses
dst_nghost = nGrowVect().