Portable MFIter Loops w/ Threading

[ax3l]

Lets you write one field kernel that runs CPU-serial, CPU-threaded and on GPU,
mirroring the C++ MFIter + ParallelFor idiom. Motivated by
#607, which asked how OpenMP
parallelizes a pyAMReX MFIter loop; the honest answer is that it does not and
cannot w/o making at least a subset of the AMReX public APIs threadsafe #614, so this provides what people actually need instead.

The kernel block

C++ today, e.g. WarpX ComputeDivE.cpp:

#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (MFIter mfi(divE, TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
    Box const& bx = mfi.tilebox(divE.ixType().toIntVect());
    Array4<Real> const& d  = divE.array(mfi);
    Array4<Real> const& ex = Ex.array(mfi);

    amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
    {
        d(i,j,k) = inv2dr[0] * (ex(i-1,j,k) - ex(i+1,j,k)) + ...;
    });
}

The same thing in Python:

@amr.for_each_tile(divE, Ex, Ey, Ez, tile=(64,) * 3, threads=8)
def _(bx, d, ex, ey, ez):
    d(bx)[...] = (
        inv2dr[0] * (ex(bx, di=-1) - ex(bx, di=+1))
        + inv2dr[1] * (ey(bx, dj=-1) - ey(bx, dj=+1))
        + inv2dr[2] * (ez(bx, dk=-1) - ez(bx, dk=+1))
    )

The decorator line stands in for the pragma and the MFIter line, its parameter
list for the tilebox and array(mfi) extractions, and its body for the
ParallelFor lambda. Same reading order, no nesting.

User-facing API

API	Purpose
Array4.__call__(bx, di=0, dj=0, dk=0, comp=0)	Box-shaped view in AMReX global index space. a(bx, di=-1) is the analogue of C++ a(i-1,j,k) over bx. Returns an AMREX_SPACEDIM-dimensional NumPy/CuPy/dpnp view.
MultiFab.tiles(tile=None)	Iterate local boxes, optionally split into tiles. iMultiFab too.
MultiFab.ix_type	The field's IntVect index type; mirrors C++ mf.ixType().toIntVect().
amr.TilingIfNotGPU(tile=None)	MFItInfo that tiles on CPU and never on GPU.
amr.for_each_tile(mfab, *others, tile=None, threads=1)	The MFIter loop as a decorator, with an optional thread pool.
for mfi in amr.MFIter(mf, info):	MFIter gained __iter__, so an explicitly constructed (e.g. tiled) iterator is finally usable.

Array4.__call__ is the load-bearing one. to_xp() is a locally 0-based view of
the whole fab; indexing by Box in global index space is what makes ghost-cell
stencils writable, and what makes them correct under tiling, where a
whole-array expression would otherwise be applied once per tile to the entire
fab.

Two bug fixes

Leaving an MFIter loop early aborted the process. __next__ returned
self, so the loop variable pinned the iterator and finalize() was only
reached on the StopIteration path. After break, the next MFIter
construction tripped AMREX_ALWAYS_ASSERT(depth == 1) and killed the process;
on GPU the Gpu::streamSynchronize() in Finalize() was skipped too.
Iteration now yields a generator whose finally finalizes on break, return and
exceptions alike -- the coverage C++ gets from ~MFIter() at scope exit.

ParIter had the identical bug (it derives from MFIter and inherits its
Finalize()) and is fixed the same way.

Behavior change

iter(mfab) is now a generator rather than the MFIter itself. Use
amr.MFIter(mfab) if you want the object; the two tests that asserted
iter(mfab).length were updated.

Notes on the design

TilingIfNotGPU() deliberately diverges from C++: the tile size has no
default, tiling is opt-in. AMReX's (1024000,8,8) is sized for OpenMP, where
a tile is nearly free; in Python each tile costs a loop iteration plus an array
view per field, and with that tile size a representative kernel measured 0.4x,
i.e. slower than serial. A large default would instead be a silent no-op
against typical 16-64^3 boxes. Tiling pays off in one case: fewer boxes than
threads, feeding the thread pool.

threads= uses a thread pool, which parallelizes because NumPy/CuPy release the
GIL for the array operations a kernel body is made of. A single serial MFIter
pass snapshots the per-tile arguments first; that is required, not an
optimization, because the MFIter mutates in place and yields itself, and AMReX
permits only one live MFIter at a time. It is a separate pool from an
AMREX_OMP=ON build's -- using both oversubscribes the node.

Measured on 256^3 / 64 tiles with a compute-bound kernel:
1.0 / 1.7 / 2.9 / 4.6 / 5.4x at 1/2/4/8/20 threads.

Explicit non-goals, now stated in the docs: there is no ParallelFor (device
lambdas cannot be written in Python) and no OpenMP region. Portability here means
array-expression portability across NumPy/CuPy/dpnp, not scalar-kernel
portability.

Docs

New "Portable Kernels, OpenMP and Threading" section in
docs/source/usage/compute.rst -- the page previously never mentioned OpenMP,
which is what prompted #607. Covers the recipe, on-node parallelism options,
tiling guidance, the GPU stream pool and synchronization, and the non-goals.

Testing

• CPU AMReX_SPACEDIM="1;2;3", AMReX_OMP=ON: ctest green; the new
  Array4.__call__ slicing verified per dimensionality (an Array4 is always 4D
  with extent-1 padding, so 1D/2D would otherwise index k with comp -- a bug
  the 3D suite cannot catch).
• CUDA (RTX A2000, sm_86): 311 passed, 32 skipped.
• SYCL (Intel Iris Xe, Level Zero, single precision -- the device has no native
  fp64): all new tests pass, amr.xp resolves to dpnp, views are real
  dpnp_arrays.
• 2 MPI ranks: new tests pass.
• Kernel correctness is checked against an analytic gradient by porting the 3D
  branch of ImpactX ForceFromSelfFields.cpp; with linear input fields the
  central differences are exact. Untiled, tiled and threaded runs agree.

Known pre-existing SYCL flake, unrelated to this PR: test_podvector.py errors
intermittently via the DLPack finalize guard. Measured with an equal-N A/B --
this branch touches no C++, so the .so is identical and only the Python files
are swapped: 12/40 failures on this branch vs 11/40 on development, i.e.
indistinguishable. test_imfab_numpy on SYCL is pre-existing by the same method.
Both are worth a separate issue.

🤖 Generated with Claude Code

[ax3l]

Similar to here / alternate to explore:
#614 (comment)