Portable MFIter Loops w/ Threading

docs/source/usage/compute.rst

@@ -72,6 +72,98 @@ For many small CPU and GPU examples on how to compute on fields, see the followi
         :caption: This files is in ``tests/test_multifab.py``.
 
 
+.. _usage-compute-portable:
+
+Portable Kernels, OpenMP and Threading
+--------------------------------------
+
+A frequent question is how OpenMP parallelizes a pyAMReX ``MFIter`` loop.
+It does not, and it cannot: in C++, OpenMP works by wrapping the loop in ``#pragma omp parallel``, so that ``MFIter`` hands each thread a subset of the tiles.
+There is no way to express that from Python, and the GIL serializes the loop body anyway.
+**Your Python** ``for mfi in mfab:`` **loop is always single-threaded, whether or not you built with** ``AMREX_OMP=ON``.
+
+OpenMP is still doing work for you, just not there.
+Every AMReX operation you call from Python runs its own OpenMP-parallel loop internally: :py:meth:`~amrex.space3d.MultiFab.saxpy`, ``lin_comb``, the norms and reductions, ``FillBoundary``, ``ParallelCopy``, ``average_down``, plotfile I/O, particle ``Redistribute``, and so on.
+So the first lever is to prefer AMReX's own operations over a hand-written loop wherever one exists.
+
+For everything else, this is how to write a custom kernel once and have it run CPU-serial, CPU-threaded and on GPU:
+
+.. literalinclude:: ../../../tests/test_multifab.py
+   :language: python3
+   :dedent: 4
+   :start-after: # Manual: Portable Kernel START
+   :end-before: # Manual: Portable Kernel END
+
+:py:func:`~amrex.space3d.for_each_tile` is the ``MFIter`` loop as a decorator, so the Python reads in the same order as the C++ it replaces.
+The kernel receives the tilebox followed by one ``Array4`` per field.
+Index those with the box: ``f(bx)`` is the tile, and ``px(bx, di=-1)`` is the analogue of C++ ``px(i-1,j,k)`` over that tile.
+Unlike :py:meth:`~amrex.space3d.Array4_double.to_xp`, which is a locally zero-based view of the whole fab, this indexing is in AMReX global index space -- which is what makes it correct under tiling, where a whole-array expression would otherwise be applied once per tile to the entire fab.
+
+``amr.xp`` is the array namespace matching your build: NumPy on CPU, CuPy for CUDA/HIP, dpnp for SYCL.
+
+On-node parallelism
+^^^^^^^^^^^^^^^^^^^
+
+In rough order of what to try:
+
+#. **More MPI ranks.** This is the primary parallelism in pyAMReX and usually the best answer.
+   Over-decompose with ``ba.max_size(...)`` and run ``mpirun -np <cores> python script.py``.
+#. **Let AMReX run the loop**, using its built-in ``MultiFab`` operations where one fits.
+#. **The** ``threads=`` **argument**, as above.
+   It runs the per-tile kernels on a thread pool, which parallelizes because NumPy and CuPy release the GIL for the array operations a kernel body is made of.
+   It only helps if the per-tile work is element-local -- no ghost exchange, no cross-tile dependencies.
+#. **The GPU**, by building with CUDA/HIP/SYCL. The same kernel source then runs on the device.
+
+Do not reach for ``multiprocessing``: the field memory lives in one process and AMReX+MPI is not fork-safe.
+
+.. note::
+
+   ``threads=`` is a *separate* thread pool from the one an ``AMREX_OMP=ON`` build uses for AMReX's own kernels.
+   Using both oversubscribes the node. Pick one, or set ``OMP_NUM_THREADS=1``.
+
+Tiling
+^^^^^^
+
+:py:func:`~amrex.space3d.TilingIfNotGPU` tiles on CPU and never on GPU, like its C++ namesake, but the tile size has **no default** and tiling is opt-in.
+Plain ``for mfi in mfab:`` does not tile, and there is no global switch that changes that -- the ``fabarray.mfiter_tile_size`` runtime parameter only sets the size used *when* tiling is requested.
+
+Tiling costs more in Python than in C++, where a tile is nearly free: here each tile is a loop iteration plus an array view per field.
+It is worth it in one situation, when you have fewer boxes per rank than threads and want to feed the thread pool -- which is the same job it does for OpenMP threads in C++.
+Size tiles so that you get roughly one to a few per thread.
+Do not reuse AMReX's C++ default of ``(1024000, 8, 8)``: in Python it produces thousands of tiny work units and runs *slower than serial*.
+
+.. note::
+
+   Two things do not carry over from C++, by design.
+   There is no ``ParallelFor``, because device lambdas cannot be written in Python; and there is no OpenMP region.
+   Portability in pyAMReX means *array-expression* portability across NumPy/CuPy/dpnp, not *scalar-kernel* portability.
+   A kernel with per-cell control flow has to become a masked array expression rather than translating line for line.
+
+GPU streams and synchronization
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You do not need to synchronize by hand.
+AMReX keeps a pool of GPU streams -- four by default, set with the ``amrex.max_gpu_streams`` runtime parameter -- and ``MFIter`` round-robins over them, advancing the stream index on every tile.
+``MFIter::Finalize()`` then synchronizes the streams it used, and pyAMReX runs it however the loop is left, including ``break``, ``return`` and exceptions, exactly as ``~MFIter()`` does in C++.
+
+Two consequences are specific to Python kernels.
+
+First, **you do not get the multi-stream overlap that C++ gets.**
+The round-robin only helps if the work is launched on AMReX's current stream, which is what ``ParallelFor`` does.
+A CuPy kernel launches on *CuPy's* current stream instead, so all tiles of a Python kernel serialize onto that one stream no matter which stream index ``MFIter`` has selected.
+
+Second, when you mix custom kernels with AMReX's own operations, correctness rests on those two stream sets being ordered.
+Today they are, because AMReX creates its pool with default flags -- making them *blocking* streams, which implicitly synchronize against the legacy default stream that CuPy uses by default.
+That is a coincidence of defaults, not a guarantee: it is lost if you open an explicit ``cupy.cuda.Stream()`` or set ``CUPY_CUDA_PER_THREAD_DEFAULT_STREAM=1``.
+In that case synchronize yourself between a Python kernel and the next AMReX call.
+
+.. note::
+
+   AMReX can also adopt a caller-supplied stream, via ``amrex::Gpu::setExternalGpuStream()`` and the RAII ``ExternalGpuStreamRegion``, which makes ``numGpuStreams()`` report 1 for as long as it is active.
+   Handing AMReX CuPy's stream that way would put both on one stream and remove the reliance on default-stream semantics.
+   pyAMReX does not bind this yet.
+
+
 .. _usage-compute-particles:
 
 Particles

src/amrex/_module_api.py

@@ -17,6 +17,7 @@
 from .extensions.PODVector import register_PODVector_extension
 from .extensions.SmallMatrix import register_SmallMatrix_extension
 from .extensions.StructOfArrays import register_SoA_extension
+from .extensions.Tiling import TilingIfNotGPU, for_each_tile
 
 
 def setup_module(ns, amr):
@@ -79,6 +80,26 @@ def read_particles_(
     read_particles_.__name__ = "read_particles"
     read_particles_.__qualname__ = "read_particles"
 
+    def TilingIfNotGPU_(tile=None):
+        """MFItInfo that tiles on CPU and never on GPU.
+
+        See :py:func:`amrex.extensions.Tiling.TilingIfNotGPU` for details.
+        """
+        return TilingIfNotGPU(amr, tile)
+
+    TilingIfNotGPU_.__name__ = "TilingIfNotGPU"
+    TilingIfNotGPU_.__qualname__ = "TilingIfNotGPU"
+
+    def for_each_tile_(mfab, *others, tile=None, threads=1):
+        """Run the decorated kernel over every box or tile of a field.
+
+        See :py:func:`amrex.extensions.Tiling.for_each_tile` for details.
+        """
+        return for_each_tile(amr, mfab, *others, tile=tile, threads=threads)
+
+    for_each_tile_.__name__ = "for_each_tile"
+    for_each_tile_.__qualname__ = "for_each_tile"
+
     def module_getattr(attr):
         """Resolve ``xp`` lazily (PEP 562).
 
@@ -123,7 +144,15 @@ def module_getattr(attr):
     ns["Print"] = Print
     ns["read_particles"] = read_particles_
     ns["list_particle_species"] = list_particle_species
+    ns["TilingIfNotGPU"] = TilingIfNotGPU_
+    ns["for_each_tile"] = for_each_tile_
     ns["__getattr__"] = module_getattr
 
-    for injected in (Print, read_particles_, module_getattr):
+    for injected in (
+        Print,
+        read_particles_,
+        TilingIfNotGPU_,
+        for_each_tile_,
+        module_getattr,
+    ):
         injected.__module__ = name

src/amrex/extensions/Array4.py

@@ -157,6 +157,73 @@ def array4_to_xp(self, copy=False, order="F"):
     return getattr(self, "to_" + xp_module_name(amr))(copy, order)
 
 
+def array4_call(self, bx, di=0, dj=0, dk=0, comp=0):
+    """Provide a view of this Array4 over a Box, in AMReX global indexing.
+
+    This is the array-expression analogue of indexing an ``Array4`` in C++:
+    ``a(bx, di=-1)`` corresponds to ``a(i-1, j, k)`` evaluated for every
+    ``(i,j,k)`` in ``bx``, so a stencil written in C++ as::
+
+        d(i,j,k) = inv2dr * (ex(i-1,j,k) - ex(i+1,j,k));
+
+    becomes::
+
+        d(bx)[...] = inv2dr * (ex(bx, di=-1) - ex(bx, di=+1))
+
+    Unlike :py:meth:`to_xp`, which is a view of the whole fab in local 0-based
+    indexing, ``bx`` here is in AMReX global index space. That is what makes it
+    usable under tiling: ``mfi.tilebox()`` names a sub-region of the fab, and
+    without restricting to it a whole-array expression would be applied once
+    per tile to the entire fab.
+
+    Reading offset cells (``di``/``dj``/``dk``) reaches into the guard cells,
+    so ``bx`` grown by the offsets must stay inside the fab; for a valid-region
+    ``bx`` that means the field needs enough ghost cells for the stencil.
+
+    Parameters
+    ----------
+    self : amrex.Array4_*
+        An Array4 class in pyAMReX.
+    bx : amrex.Box
+        Index-space region to view, in AMReX global indices.
+    di, dj, dk : int, optional
+        Shift the region by this many cells per direction (default 0).
+    comp : int, optional
+        Component to select (default 0).
+
+    Returns
+    -------
+    xp.array
+        A non-copying NumPy, CuPy or dpnp view of ``bx`` shifted by
+        ``(di, dj, dk)``, with ``AMREX_SPACEDIM`` dimensions.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+
+    arr = self.to_xp(copy=False, order="F")
+    lo = amr.lbound(self)
+    lo = (lo.x, lo.y, lo.z)
+    shift = (di, dj, dk)
+
+    # An Array4 is always 4D (i,j,k,n): unused directions carry extent 1 rather
+    # than being dropped. So always index all three spatial axes -- taking only
+    # AMREX_SPACEDIM of them would make `comp` land on k in 1D/2D. The unused
+    # ones are indexed with a scalar rather than sliced, which drops them, so
+    # the result has AMREX_SPACEDIM dimensions in every build.
+    dims = amr.Config.spacedim
+    slices = tuple(
+        slice(
+            bx.small_end[d] + shift[d] - lo[d],
+            bx.big_end[d] + shift[d] - lo[d] + 1,
+        )
+        if d < dims
+        else 0
+        for d in range(3)
+    )
+    return arr[slices + (comp,)]
+
+
 def register_Array4_extension(amr):
     """Array4 helper methods"""
     import inspect
@@ -175,3 +242,4 @@ def register_Array4_extension(amr):
         Array4_type.to_cupy = array4_to_cupy
         Array4_type.to_dpnp = array4_to_dpnp
         Array4_type.to_xp = array4_to_xp
+        Array4_type.__call__ = array4_call

src/amrex/extensions/Iterator.py

@@ -40,6 +40,34 @@ def next(self):
     return self
 
 
+def iterate(it):
+    """Drive an ``MFIter``/``ParIter`` as a generator.
+
+    This is what ``__iter__`` returns, so it is the object a ``for`` loop holds
+    on to. The ``finally`` clause runs on normal exhaustion *and* on ``break``,
+    ``return`` and exceptions -- the same coverage C++ gets from ``~MFIter()``
+    at scope exit.
+
+    That matters because the C++ destructor is not a reliable stand-in here:
+    ``__next__`` yields the iterator itself, so after ``break`` the loop
+    variable still references it and it is not destroyed. Its ``Finalize()``
+    would then be deferred, which (a) leaves ``MFIter::depth`` at 1 so the next
+    iterator construction trips ``AMREX_ALWAYS_ASSERT(depth == 1)`` and aborts,
+    and (b) skips the ``Gpu::streamSynchronize()`` that ``Finalize()`` performs.
+
+    ``MFIter::Finalize()`` is idempotent (guarded by its ``finalized`` flag), so
+    an explicit ``it.finalize()`` in user code stays safe.
+
+    it: the C++ iterator to drive; yielded unchanged on every step
+    """
+    try:
+        while it.is_valid:
+            yield it
+            it._incr()
+    finally:
+        it.finalize()
+
+
 def getitem(self, name):
     """Access (read/write) particle vectors."""
     if not self.is_soa_particle:

src/amrex/extensions/MultiFab.py

@@ -8,7 +8,8 @@
 
 import numpy as np
 
-from .Iterator import next
+from .Iterator import iterate, next
+from .Tiling import ix_type, tiles
 
 
 def mf_to_numpy(self, copy=False, order="F"):
@@ -696,11 +697,25 @@ def register_MultiFab_extension(amr):
     # register member functions for the MFIter type
     amr.MFIter.__next__ = next
 
+    # Iterating an MFIter yields a generator rather than the MFIter itself, so
+    # that leaving the loop early (break/return/exception) still finalizes it.
+    # See amrex.extensions.Iterator.iterate for why the C++ destructor alone is
+    # not enough here.
+    #
+    # Note: __iter__ must return the generator, not something that would itself
+    # need another iter() call -- Python invokes __next__ directly on whatever
+    # __iter__ returns and does not re-iter() it. Registering iterate() on
+    # MFIter alone would therefore not cover `for mfi in mfab:`.
+    amr.MFIter.__iter__ = iterate
+
     # FabArrayBase: iterate as data access in Box index space
-    amr.FabArrayBase.__iter__ = lambda fab: amr.MFIter(fab)
+    amr.FabArrayBase.__iter__ = lambda fab: iterate(amr.MFIter(fab))
 
     # register member functions for the MultiFab type
-    amr.MultiFab.__iter__ = lambda mfab: amr.MFIter(mfab)
+    amr.MultiFab.__iter__ = lambda mfab: iterate(amr.MFIter(mfab))
+    amr.MultiFab.tiles = lambda self, tile=None: tiles(amr, self, tile)
+    amr.MultiFab.tiles.__doc__ = tiles.__doc__
+    amr.MultiFab.ix_type = property(ix_type)
 
     amr.MultiFab.to_numpy = mf_to_numpy
     amr.MultiFab.to_cupy = mf_to_cupy
@@ -717,7 +732,10 @@ def register_MultiFab_extension(amr):
     amr.MultiFab.__setitem__ = __setitem__
 
     # iMultiFab
-    amr.iMultiFab.__iter__ = lambda imfab: amr.MFIter(imfab)
+    amr.iMultiFab.__iter__ = lambda imfab: iterate(amr.MFIter(imfab))
+    amr.iMultiFab.tiles = lambda self, tile=None: tiles(amr, self, tile)
+    amr.iMultiFab.tiles.__doc__ = tiles.__doc__
+    amr.iMultiFab.ix_type = property(ix_type)
 
     amr.iMultiFab.to_numpy = mf_to_numpy
     amr.iMultiFab.to_cupy = mf_to_cupy

src/amrex/extensions/ParticleContainer.py

@@ -9,7 +9,7 @@
 import os
 import warnings
 
-from .Iterator import getitem, next
+from .Iterator import getitem, iterate, next
 
 
 def iterator(self, *args, level=None):
@@ -351,7 +351,10 @@ def register_ParticleContainer_extension(amr):
         ),
     ):
         ParIter_type.__next__ = next
-        ParIter_type.__iter__ = lambda self: self
+        # a generator, so that leaving the loop early still finalizes the
+        # iterator -- ParIterBase derives from MFIter and inherits its
+        # Finalize(); see amrex.extensions.Iterator.iterate
+        ParIter_type.__iter__ = iterate
         ParIter_type.__getitem__ = getitem
 
     # register member functions for every ParticleContainer_* type