Skip to content

[WIP] Fewer guard cell exchanges when applying Darwin linear operator - #7177

Open
RemiLehe wants to merge 2 commits into
BLAST-WarpX:developmentfrom
RemiLehe:darwin-fused-operator
Open

[WIP] Fewer guard cell exchanges when applying Darwin linear operator#7177
RemiLehe wants to merge 2 commits into
BLAST-WarpX:developmentfrom
RemiLehe:darwin-fused-operator

Conversation

@RemiLehe

@RemiLehe RemiLehe commented Aug 19, 2026

Copy link
Copy Markdown
Member

Motivation

The linear operator that GMRES applies on every iteration of the Darwin solve is

A Z = bilaplacian(Z) + curl(chi curl(Z))

where chi is the deposited mass matrix scaled by 2 mu_0 / dt. Each stage of that expression is a short, local stencil, so the composition is itself a single banded stencil on Z, with a predictable width set by the curl stencils and the mass matrix stencil. DarwinLinearFieldOperator::apply() nonetheless evaluated it as three separate passes with a FillBoundaryAndSync round 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 Zscratch at 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 the FillBoundary that 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_fp twice — 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_fp and Efield_fp to hold its intermediates, so Efield_fp keeps 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 it
and sized once in define().

Fewer kernel launches, and one less full-array setVal(0) per intermediate.

Main changes

  • ComputeCurlA / ComputeCurlB take an optional ngrow argument (defaulting to
    zero, 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 at
    the 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 ComputeCurlA variants assert ngrow == 0 rather than pretending to support
    it.

  • ImplicitSolver::MassMatricesStencilHalfWidth() reports how far the mass matrix
    gather reaches — max over the nine blocks of ncomp_ab[dir]/2, which is the furthest
    cell read regardless of the relative staggering of the two components.

  • DarwinLinearFieldOperator owns the intermediates (m_dA_*, m_chidA_*) alongside
    m_Zscratch_* / m_lapZ_*, with the guard widths derived in define() by walking the
    stencil chain backwards from the output:

    field ghost cells why
    chidA 1 the final curl_A reads i..i+1
    dA 1 + half_width ApplyMassMatrices clamps its gather to the input's guard region
    Zscratch max(2, dA_ng + 1) curl_B reads i-1..i; the nabla^4 stencil reads i-2..i+2

    For the common case of shape 1 with direct deposition in 3D, half_width is 2, so
    Zscratch goes 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 ApplyMassMatrices clamps its gather to the input's allocated guard
    region rather than to a caller-supplied validity box. Tightening that would be a
    natural follow-up.

  • SemiImplicitDarwin::Define() calls InitializeMassMatrices() before defining the
    linear operator, since the operator now sizes its scratch from the mass matrix stencil.
    Its only other dependency is m_nlsolver, and only when m_use_mass_matrices_pc is
    true, which the Darwin solver sets to false.

  • SemiImplicitDarwin::ApplyScaledMassMatrices() no longer ends with a
    FillBoundaryAndSync on its output. ApplyMassMatrices() was never communicating
    those guard cells — it computes them, from the same wide-stencil read of dA it uses
    for 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 FillBoundaryAndSync on rhs_vec in apply() is gone too. Its comment
    claimed 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 from Zscratch, which was synced, and from mass
    matrices that SumBoundaryJ leaves equal at duplicated nodes, so the two cells are
    computed by identical operations on identical inputs. This was not free — rhs_vec has
    zero ghost cells, so the call was purely the OverrideSync half, i.e. real
    communication once per GMRES iteration.

ApplyMassMatrices() itself needed no change: it already fills min(out ng, S ng) guard
cells 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 SumBoundaryJ
passes dst_nghost = nGrowVect().

@RemiLehe RemiLehe changed the title [WIP] Darwin fused operator [WIP] Fewer guard cell exchanges when applying Darwin linear operator Aug 19, 2026
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
RemiLehe force-pushed the darwin-fused-operator branch from 1c31398 to 7208a92 Compare August 20, 2026 12:48
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
RemiLehe force-pushed the darwin-fused-operator branch from 7208a92 to cfe7f98 Compare August 20, 2026 15:48
@roelof-groenewald

Copy link
Copy Markdown
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.
Why do you have it marked as work-in-progress? Are there outstanding changes / testing you would still like to do before merging this?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants