Skip to content

ECT: atomic test-and-subtract for the face-extension area accounting #7152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
RemiLehe merged 4 commits into from
Sep 1, 2026
Merged

ECT: atomic test-and-subtract for the face-extension area accounting #7152

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
934808c
ECT: atomic test-and-subtract for the face-extension area accounting
clarkse-he Aug 11, 2026
c5e8e5a
Merge branch 'development' into ect-atomic-extensions
RemiLehe Aug 24, 2026
2daeef5
Remove EOL whitespace from the face-extension predicate
clarkse-he Aug 24, 2026
0986772
ECT: make the face extension all-or-nothing, and address review
clarkse-he Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 143 additions & 15 deletions Source/EmbeddedBoundary/WarpXFaceExtensions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
#include <ablastr/warn_manager/WarnManager.H>
#include <ablastr/fields/MultiFabRegister.H>

#include <AMReX_Functional.H>
#include <AMReX_GpuAtomic.H>
#include <AMReX_Scan.H>
#include <AMReX_iMultiFab.H>
#include <AMReX_MultiFab.H>
Expand Down Expand Up @@ -357,6 +359,54 @@ namespace
amrex::Abort("SetNeigh: dim must be 0, 1 or 2");
}


/**
* \brief Get the address of the value of arr in the neighbor (i_n, j_n) on
* the plane with normal 'dim' (same indexing convention as GetNeigh), for
* atomic updates of the neighbor's value.
*
* \param[in] arr data to be accessed
* \param[in] i, j, k the indices of the "center" cell
* \param[in] i_n the offset of the neighbor in the first direction
* \param[in] j_n the offset of the neighbor in the second direction
* \param[in] dim normal direction to the plane in consideration (0 for x, 1 for y, 2 for z)
*/
template <class T>
AMREX_GPU_DEVICE AMREX_FORCE_INLINE
constexpr
T*
GetNeighPtr(const amrex::Array4<T>& arr,
const int i, const int j, const int k,
const int i_n, const int j_n, const int dim){

if(dim == 0){
return &arr(i, j + i_n, k + j_n);
}
#ifdef WARPX_DIM_XZ
else if(dim == 1 || (dim == 2)){
return &arr(i + i_n, j + j_n, k);
}
#elif defined(WARPX_DIM_3D)
else if(dim == 1){
return &arr(i + i_n, j, k + j_n);
}
else if(dim == 2){
return &arr(i + i_n, j + j_n, k);
}
#else
else if(dim == 1){
amrex::Abort("GetNeighPtr: Only implemented in 2D3V and 3D3V");
}
else if(dim == 2){
return &arr(i + i_n, j + j_n, k);
}
#endif

amrex::Abort("GetNeighPtr: dim must be 0, 1 or 2");

return nullptr;
}

#ifdef AMREX_USE_EB

#ifndef WARPX_DIM_RZ
Expand Down Expand Up @@ -676,17 +726,33 @@ WarpX::ComputeOneWayExtensions ()
// has given away already some area, so we use Sz_red rather than Sz.
// If no face is available we don't do anything and we will need to use the
// multi-face extensions.
const int flag_neigh =
::GetNeigh(flag_info_face, i, j, k, i_n, j_n, idim);
if (::GetNeigh(S_mod, i, j, k, i_n, j_n, idim) > S_ext
&& (flag_neigh == FaceInfo::available
|| flag_neigh == FaceInfo::intruded)
&& flag_ext_face(i, j, k)) {

::SetNeigh(S_mod,
::GetNeigh(S_mod, i, j, k, i_n, j_n, idim) - S_ext,
i, j, k, i_n, j_n, idim);

// The area is taken with an atomic test-and-subtract: on GPU
// several faces can try to borrow from the same intruded face
// concurrently, and a plain read-test-write lets the intruded

@RemiLehe RemiLehe Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add a comment that explains the syntax of amrex::Gpu::Atomic::If?
Right now, it is difficult to parse what this code is doing.

@RemiLehe RemiLehe Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The syntax is here: https://amrex-codes.github.io/amrex/doxygen/namespaceamrex_1_1Gpu_1_1Atomic.html#a54472d6cb90e23b6479dc667f7e02669

and is not necessarily very intuitive (the condition is given at the end)

@clarkse-he clarkse-he Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added — each argument is now documented inline at the call site, with a link to the AMReX docs. You're right that the trailing condition is the confusing part; I also noted that the flag_ext_face test inside the predicate is what stops the loop after the first successful borrow.

// face give the same area away more than once.
//
// *Atomically* decrement `S_mod` of the neighboring cell by
// `S_ext` under the condition that this is possible (i.e. that
// S_mod-S_ext is positive, and the cell is marked as available
// for borrowing). If this indeed updated `S_mod`, it returns
// `true` for `borrowed`. For the syntax, see
// https://amrex-codes.github.io/amrex/doxygen/namespaceamrex_1_1Gpu_1_1Atomic.html
const bool borrowed = amrex::Gpu::Atomic::If(
::GetNeighPtr(S_mod, i, j, k, i_n, j_n, idim), // address to atomically update
S_ext, // value to combine
amrex::Minus<amrex::Real>(), // operation to perform when combining the value
// condition: callable that gets called with rem=S_mod-S_ext.
// The `flag_ext_face` test also stops this loop after the
// first successful borrow, as it is cleared just below.
[=] (amrex::Real rem) {
const int flag_neigh = ::GetNeigh(flag_info_face, i, j, k, i_n, j_n, idim);
return rem > amrex::Real(0.)
&& (flag_neigh == FaceInfo::available
|| flag_neigh == FaceInfo::intruded)
&& flag_ext_face(i, j, k);
});
Comment thread
RemiLehe marked this conversation as resolved.

if (borrowed) {
// Insert the index of the face info
borrowing_inds[ps] = ps;
// Store the information about the intruded face in the dataset of the
Expand All @@ -705,6 +771,17 @@ WarpX::ComputeOneWayExtensions ()
}
}
}
// The counting pass reserved one slot for this face, but the atomic
// test-and-subtract above fails if a concurrently extended face drained
// the intruded face in the meantime. The face is then still flagged, and
// has to report that it borrowed nothing: the solver reads
// `borrowing_size` entries starting at `*borrowing_inds_pointer` (see
// EvolveBCartesianECT), which were never filled in. The face itself is
// left to the eight-ways extension.
if (flag_ext_face(i, j, k)) {
borrowing_size(i, j, k) = 0;
borrowing_inds_pointer(i, j, k) = nullptr;
}
}
}, amrex::Scan::Type::exclusive);
}
Expand Down Expand Up @@ -854,10 +931,36 @@ WarpX::ComputeEightWaysExtensions ()
if(denom >= S_ext){
S_mod(i, j, k) = S(i, j, k);
int count = 0;
// The extension is all-or-nothing: a face that got only some of its
// patches would not reach its stable area, and the area it did take
// would be lost to the faces that lent it, since the ECT update of an
// intruded face assumes that the area it lent is accounted for by the
// face that borrowed it (see EvolveBCartesianECT).
bool all_borrowed = true;

@RemiLehe RemiLehe Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add a comment explaining why this is needed

@clarkse-he clarkse-he Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — this came in from #7104 along with the rest of the partial-extension bookkeeping. There, a face's lent area is tracked in a dedicated lent_area MultiFab and reduced across boxes by sync_lent_areas(), which is what makes a partially satisfied borrow representable in the first place.

Without that ledger it was simply wrong here: a face that got only some of its patches left the intruded neighbours' S_mod short with nothing to compensate, which breaks the S_mod + (sum of lent areas) == S identity that the intruded-face branch of EvolveBCartesianECT relies on.

So I've made the extension all-or-nothing instead: if any patch cannot be taken, the successful ones are given back and the face is left flagged, so it gets stabilized by the BCK correction. all_borrowed now only drives that rollback, and says so in a comment. Partial and cross-box borrowing stay with #7104, which will want a rebase on top of whatever lands here.

for (int i_n = -1; i_n < 2; i_n++) {
for (int j_n = -1; j_n < 2; j_n++) {
if(local_avail(i_n + 1, j_n + 1) != 0_rt){
if (count == nborrow) {
// The borrowing pass found more available neighbors
// than the counting pass reserved slots for
all_borrowed = false;

@aarontran aarontran Aug 27, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi Eric, I'm working with Remi to study this & subsequent PRs.

Do you have an example situation where borrowing pass gets count > nborrow and thus triggers this check?

If the first lambda function (counting pass) in amrex::Scan::PrefixSum<int> is guaranteed to complete before any calls to the second lambda function (borrowing pass), then it seems count must be strictly <= nborrow because during the first pass, the neighbor faces still have their "maximal" available area in S_mod as no atomic updates have occurred yet. If no such guarantee exists, then I could see the check triggering...

So the question might be translated to: what guarantees, if any, do we have on the execution order of the lambdas in PrefixSum (and where can humans best find information about any such guarantees or lack thereof?) ( ... deferring the question of cross-box or cross-multifab contention til your PR 7104 ... )

@clarkse-he clarkse-he Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi Aaron — thanks for the careful read, and for the measurements above.

Short answer: no, there is no such guarantee on CUDA or HIP — and separately, this check isn't really about fin/fout ordering, it's about the fout calls racing with each other, so it would survive even if the ordering guarantee existed.

1. What amrex::Scan::PrefixSum actually guarantees

I couldn't find this documented anywhere, so I went through the implementation at the AMReX commit WarpX 26.08 pins (0579402, Src/Base/AMReX_Scan.H). It's backend-dependent:

backend structure all fin before any fout?
CUDA (L810) one kernel; CUB BlockLoad + BlockScan with cub::TilePrefixCallbackOp decoupled look-back No. Each block calls fin only over its own tile (L876), and block 0 goes straight to fout (L935) right after its BlockScan; block k only waits on look-back over tiles < k. Nothing ever waits on a later block's fin.
HIP (L642) same shape, rocprim::detail::lookback_scan_prefix_op No (fin L732, fout L786).
SYCL, nblocks > 1 (L196, PrefixSum_mp) three separate kernel launches: all fin (L247), scan the block sums (L299), all fout (L377) Yes, strictly separated — unless AMREX_SYCL_NO_MULTIPASS_SCAN.
SYCL, single block (L400) hand-rolled Merrill & Garland n/a
no GPU (L1117) plain serial loop, x = fin(i); …; fout(i, y) interleaved per cell, single-threaded

So on CUDA/HIP the two lambdas overlap in wall-clock time across blocks. As far as I can tell, reading AMReX_Scan.H is the only way to know this; a one-line note on PrefixSum saying fin and fout are not separated by a barrier on the look-back backends would be a worthwhile small AMReX PR if you agree it's contract-level information.

One correction while I'm here, since it bears on how you read my earlier summary: on host, amrex::Gpu::Atomic::If is a plain non-atomic read-test-write (AMReX_GpuAtomic.H L256, the AMREX_IF_ON_HOST branch). Combined with the serial non-GPU scan and the fact that neither ComputeOneWayExtensions nor ComputeEightWaysExtensions has an OpenMP pragma over its MFIter, a CPU build never exercises the atomic path at all within a box. I overstated it above when I said the OMP_NUM_THREADS-unset runs exercised the atomic path under real concurrency — they didn't; those were only a regression check that the restructuring still reproduces the committed benchmarks.

2. Why the check survives even given a strict fin/fout fence

count > nborrow requires the borrowing pass to find more live entries in local_avail than the counting pass did. Working through the inputs:

  • S is never written in these kernels.
  • flag_info_face is only ever written available → intruded (L969), and local_avail accepts both, so flag updates can't change membership.
  • That leaves S_mod. A neighbour is dropped in the while loop when S_mod(neigh) - patch <= 0, so the set can only grow if some neighbour's S_mod is larger at borrowing time than it was at counting time.

Nearly every write to S_mod decreases it, and those land in the !borrowed / denom < S_ext branches instead (count < nborrow). The exception is the rollback this PR introduced, at L984:

amrex::Gpu::Atomic::AddNoRet(
    ::GetNeighPtr(S_mod, i, j, k, vec(0), vec(1), idim),
    borrowing_area[ps + n]);

Face B failing all-or-nothing hands area back to a shared neighbour N, concurrently with face A's borrowing loop. A's counting pass had excluded N as drained; A's borrowing pass recomputes local_avail after the refund and includes it — one slot more than was reserved. A and B are both in fout, so a fence between the passes wouldn't help.

(The other writes that raise S_mod — S_mod(i,j,k) = S(i,j,k), += patch, and = S(i,j,k) + S_ext in the one-way pass — are all on the extending face's own cell, and are harmless here: a face awaiting extension keeps flag_info_face == FaceInfo::extended, which is neither available nor intruded, so it's never a member of anybody's local_avail.)

3. Do I have a concrete trace of it firing?

No — honestly it's a guard, not a fix for something I caught in the act. I added it because the failure mode without it isn't a wrong number: it's borrowing_inds[ps + count] running past this face's reserved slots into the next face's region of the scan-sized DeviceVector. I'd rather pay a compare than rely on the argument that two independent re-derivations of local_avail stay in lockstep across a race.

If you'd like to settle it empirically, your sphere case is a much better instrument than anything I have (5,552 unstable faces per dimension, and you're on GPU). A device counter in that branch —

if (count == nborrow) {
    amrex::Gpu::Atomic::AddNoRet(p_ncapped, 1);
    all_borrowed = false;
    continue;
}

— over a few hundred steps would say whether the path is live in practice. By the analysis in §2 it can only fire on a step where some face also rolled back, so I'd expect it to be rare if it fires at all, and on a CPU build it can't fire by construction. I'd genuinely like to know either way: if it turns out to be unreachable, I'd rather it become an AMREX_ASSERT than stay as silent flow control.

Happy to fold any of this into #7104, which has to rebase on this anyway.

@aarontran aarontran Sep 1, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, that clarifies & i think no followup needed for now,

  1. Agree no guarantees of 1st/2nd function order across threadblocks (based only on quick skim of your reply & AMReX source)
  2. Agree on count > nborrow possible due to rollback (after sketching out one possible execution order of two adjoining threadblocks that can trigger this code path on paper; omitted for brevity)
  3. Agree on rationale for pre-emptive catch

continue;
}
const amrex::Real patch = S_ext * ::GetNeigh(S, i, j, k, i_n, j_n, idim) / denom;
// Atomic test-and-subtract, for the same reason as in
// ComputeOneWayExtensions: an intruded face shared by
// concurrently extended faces must not give the same
// area away more than once
const bool borrowed = amrex::Gpu::Atomic::If(
::GetNeighPtr(S_mod, i, j, k, i_n, j_n, idim),
patch, amrex::Minus<amrex::Real>(),
[=] (amrex::Real rem) {
return rem > amrex::Real(0.);
});
if (!borrowed) {
all_borrowed = false;

@RemiLehe RemiLehe Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you explain why this is needed

@clarkse-he clarkse-he Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above — this is what triggers the rollback now, and it is commented.

While re-checking it I also found a hole in the count < nborrow cap: the borrowing pass can find more available neighbours than the counting pass reserved slots for, and silently skipping them could unflag a face that never actually reached S_stab. That path now sets all_borrowed = false as well, so the face rolls back and goes to BCK.

continue;
}
borrowing_inds[ps + count] = ps + count;
FaceInfoBox::addConnectedNeighbor(i_n, j_n, ps + count,
borrowing_neigh_faces);
Expand All @@ -868,14 +971,39 @@ WarpX::ComputeEightWaysExtensions ()
i, j, k, i_n, j_n, idim);

S_mod(i, j, k) += patch;
::SetNeigh(S_mod,
::GetNeigh(S_mod, i, j, k, i_n, j_n, idim) - patch,
i, j, k, i_n, j_n, idim);
count +=1;
}
}
}
flag_ext_face(i, j, k) = false;
if (!all_borrowed) {
// Give the area of the successful patches back, and leave the face
// flagged so that it is stabilized by the BCK correction instead
for (int n = 0; n < count; n++) {
auto const vec =
FaceInfoBox::uint8_to_inds(borrowing_neigh_faces[ps + n]);
amrex::Gpu::Atomic::AddNoRet(
::GetNeighPtr(S_mod, i, j, k, vec(0), vec(1), idim),
borrowing_area[ps + n]);
}
count = 0;
S_mod(i, j, k) = S(i, j, k);
}
// The recorded size has to match the entries actually written, since
// the solver reads `borrowing_size` entries starting at
// `*borrowing_inds_pointer` (see EvolveBCartesianECT)
borrowing_size(i, j, k) = count;

@RemiLehe RemiLehe Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you check whether this is still needed?

@clarkse-he clarkse-he Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, for the same reason as the one-way case: the solver trusts borrowing_size when it walks the borrowing list, and count could come out below the nborrow slots the counting pass reserved. With the all-or-nothing change it is now either 0 or the full count.

if (count == 0) {
borrowing_inds_pointer(i, j, k) = nullptr;
} else {
flag_ext_face(i, j, k) = false;
}
}
else {
// The area available shrank between the counting and the borrowing
// pass: the face cannot be extended after all, and has to report
// that it borrowed nothing
borrowing_size(i, j, k) = 0;
borrowing_inds_pointer(i, j, k) = nullptr;
}
}
}, amrex::Scan::Type::exclusive);
Expand Down
Loading