diff --git a/docs/source/usage/compute.rst b/docs/source/usage/compute.rst index 8e5fcc62..3e0bed80 100644 --- a/docs/source/usage/compute.rst +++ b/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 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 diff --git a/src/amrex/_module_api.py b/src/amrex/_module_api.py index a48973e7..674d6213 100644 --- a/src/amrex/_module_api.py +++ b/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 diff --git a/src/amrex/extensions/Array4.py b/src/amrex/extensions/Array4.py index b353bfa0..cc4d7446 100644 --- a/src/amrex/extensions/Array4.py +++ b/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 diff --git a/src/amrex/extensions/Iterator.py b/src/amrex/extensions/Iterator.py index c57d9fcb..3a199af8 100644 --- a/src/amrex/extensions/Iterator.py +++ b/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: diff --git a/src/amrex/extensions/MultiFab.py b/src/amrex/extensions/MultiFab.py index a4d52443..e2bef270 100644 --- a/src/amrex/extensions/MultiFab.py +++ b/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 diff --git a/src/amrex/extensions/ParticleContainer.py b/src/amrex/extensions/ParticleContainer.py index ebb367be..5960425b 100644 --- a/src/amrex/extensions/ParticleContainer.py +++ b/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 diff --git a/src/amrex/extensions/Tiling.py b/src/amrex/extensions/Tiling.py new file mode 100644 index 00000000..728e087b --- /dev/null +++ b/src/amrex/extensions/Tiling.py @@ -0,0 +1,213 @@ +""" +This file is part of pyAMReX + +Copyright 2026 AMReX community +Authors: Axel Huebl +License: BSD-3-Clause-LBNL +""" + +from concurrent.futures import ThreadPoolExecutor + +from .Iterator import iterate + + +def TilingIfNotGPU(amr, tile=None): + """Return an ``MFItInfo`` that tiles on CPU and never on GPU. + + This is the analogue of C++ ``amrex::TilingIfNotGPU()``, with one + deliberate difference: ``tile`` has **no default**, so tiling is opt-in. + + In C++ the default tile size is ``FabArrayBase::mfiter_tile_size`` + (``1024000,8,8`` in 3D), which is sized for OpenMP: it hands threads many + small work units, and the per-tile cost of an ``MFIter`` step is + negligible. In Python the per-tile cost is a loop iteration plus an array + view per field, which is not negligible -- with that tile size a + representative kernel measured *slower than serial*. A large default + instead would silently do nothing for the typical 16-64^3 box. + + Tiling in Python only pays off when there are fewer boxes than threads and + you are threading over them (see :py:func:`for_each_tile`), so the caller + has to say so, and say how big. + + Parameters + ---------- + amr : module + The dimensionality-specific pyAMReX module. + tile : sequence of int, optional + Tile size, ``AMREX_SPACEDIM`` entries. ``None`` (default) means do not + tile, i.e. iterate whole boxes. Ignored on GPU. + + Returns + ------- + amr.MFItInfo + Info object to hand to ``amr.MFIter``. + """ + info = amr.MFItInfo() + if tile is not None and not amr.Config.have_gpu: + info.enable_tiling(amr.IntVect(*tile)) + return info + + +def tiles(amr, mfab, tile=None): + """Iterate the local boxes of ``mfab``, optionally split into tiles. + + Equivalent to ``iter(mfab)`` when ``tile`` is ``None``. Like every + iteration path in pyAMReX this yields the ``MFIter`` itself, and finalizes + it however the loop is left. + + Parameters + ---------- + amr : module + The dimensionality-specific pyAMReX module. + mfab : amr.MultiFab or amr.iMultiFab + The field to iterate. + tile : sequence of int, optional + Tile size; see :py:func:`TilingIfNotGPU`. + + Yields + ------ + amr.MFIter + The iterator, positioned on the current box or tile. + """ + return iterate(amr.MFIter(mfab, TilingIfNotGPU(amr, tile))) + + +def _sync_device(amr, sample): + """Wait for the array library's stream/queue, on GPU builds. + + Kernels written against ``amr.xp`` launch on CuPy's stream or dpnp's SYCL + queue, not on the AMReX stream that ``MFIter::Finalize()`` synchronizes, so + that one has to be waited on separately. + + Parameters + ---------- + amr : module + The dimensionality-specific pyAMReX module. + sample : amr.Array4_* + An Array4 of a field the kernel wrote. Only SYCL needs it, to find the + queue the work went to, so the array view is built only in that case. + + Raises + ------ + ImportError + If the build's array library (CuPy or dpnp) is not installed. Those are + optional dependencies of pyAMReX; reaching here means the kernel just + ran on device arrays, so one of them is necessarily already imported. + """ + if amr.Config.gpu_backend == "SYCL": + # dpnp has no module-level synchronize(); the queue lives on the array + sample.to_xp(copy=False).sycl_queue.wait() + else: # CUDA, HIP + import cupy + + cupy.cuda.get_current_stream().synchronize() + + +def ix_type(self): + """The index type (staggering/centering) of this field, as an ``IntVect``. + + The analogue of C++ ``mf.ixType().toIntVect()``, for passing to + ``mfi.tilebox(...)`` so that the returned box carries the right centering. + """ + return self.box_array().ix_type().to_IntVect() + + +def for_each_tile(amr, mfab, *others, tile=None, threads=1): + """Run the decorated kernel over every box or tile of ``mfab``. + + This is the ``MFIter`` loop expressed as a decorator, so that a portable + kernel reads in the same order as the equivalent C++ block:: + + #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 const& d = divE.array(mfi); + Array4 const& ex = Ex.array(mfi); + + amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + d(i,j,k) = inv2dr * (ex(i-1,j,k) - ex(i+1,j,k)); + }); + } + + becomes:: + + @amr.for_each_tile(divE, Ex, tile=(64,) * 3, threads=8) + def _(bx, d, ex): + d(bx)[...] = inv2dr * (ex(bx, di=-1) - ex(bx, di=+1)) + + The decorator line stands in for the pragma and the ``MFIter`` line, the + parameter list for the ``tilebox`` and ``array(mfi)`` extractions, and the + body for the ``ParallelFor`` lambda. The kernel receives the tilebox + followed by one ``Array4`` per field passed here, in order; index those + with the box, as ``d(bx)`` or ``ex(bx, di=-1)``. + + The decorated function is called immediately and returned unchanged, so + naming it ``_`` is conventional but not required. + + On threading: a single serial ``MFIter`` pass collects the per-tile + arguments before any kernel runs, and this is required rather than an + optimization. The ``MFIter`` mutates in place on each step and yields + itself, so its state cannot be handed to a worker to use later; and AMReX + permits only one live ``MFIter`` at a time + (``AMREX_ALWAYS_ASSERT(depth == 1)``), so workers cannot drive their own. + Threading then works because numpy/cupy release the GIL for the array + operations the kernel body is made of. + + Note that ``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. + + Parameters + ---------- + amr : module + The dimensionality-specific pyAMReX module. + mfab : amr.MultiFab or amr.iMultiFab + Field defining the iteration space; also the first kernel argument. + *others : amr.MultiFab or amr.iMultiFab + Further fields, passed to the kernel after ``mfab``. They must share + ``mfab``'s BoxArray and DistributionMapping. + tile : sequence of int, optional + Tile size; see :py:func:`TilingIfNotGPU`. + threads : int, optional + Worker threads for the kernel (default 1, i.e. serial). Forced to 1 on + GPU, where the device provides the parallelism. + + Returns + ------- + callable + A decorator that runs the kernel and returns it unchanged. + """ + + def decorate(kernel): + nthreads = 1 if amr.Config.have_gpu else threads + + # Snapshot per-tile arguments in one serial pass; see docstring. + ixt = ix_type(mfab) + tasks = [ + (mfi.tilebox(ixt), mfab.array(mfi), *[o.array(mfi) for o in others]) + for mfi in tiles(amr, mfab, tile) + ] + + if nthreads == 1: + for task in tasks: + kernel(*task) + else: + with ThreadPoolExecutor(max_workers=nthreads) as pool: + # list() so that exceptions raised in a worker propagate here + list(pool.map(lambda task: kernel(*task), tasks)) + + # The MFIter is already finalized at this point -- the snapshot pass + # above ran it to completion -- so its Gpu::streamSynchronize() came + # *before* any of these kernels launched and cannot cover them. The + # kernels also ran on the array library's stream rather than AMReX's, + # so synchronize that one to make the work complete on return. + if amr.Config.have_gpu and tasks: + _sync_device(amr, tasks[0][1]) + + return kernel + + return decorate diff --git a/tests/test_imultifab.py b/tests/test_imultifab.py index 60ad036b..5623391c 100644 --- a/tests/test_imultifab.py +++ b/tests/test_imultifab.py @@ -242,14 +242,16 @@ def test_imfab_ops(boxarr, distmap, nghost): def test_imfab_mfiter(imfab): assert len(imfab) == 8 - assert iter(imfab).is_valid - assert iter(imfab).length == 8 + # iter(imfab) is a generator driving an MFIter, not the MFIter itself, so + # that leaving a loop early still finalizes it. Ask the MFIter directly. + assert amr.MFIter(imfab).is_valid + assert amr.MFIter(imfab).length == 8 cnt = 0 for _mfi in imfab: cnt += 1 - assert iter(imfab).length == cnt + assert amr.MFIter(imfab).length == cnt def test_imfab_mfiter_keeps_imfab_alive(imfab, assert_keeps_python_alive): diff --git a/tests/test_multifab.py b/tests/test_multifab.py index d02d3dc3..187f6efd 100644 --- a/tests/test_multifab.py +++ b/tests/test_multifab.py @@ -265,14 +265,16 @@ def test_mfab_ops(boxarr, distmap, nghost): def test_mfab_mfiter(mfab): assert len(mfab) == 8 - assert iter(mfab).is_valid - assert iter(mfab).length == 8 + # iter(mfab) is a generator driving an MFIter, not the MFIter itself, so + # that leaving a loop early still finalizes it. Ask the MFIter directly. + assert amr.MFIter(mfab).is_valid + assert amr.MFIter(mfab).length == 8 cnt = 0 for _mfi in mfab: cnt += 1 - assert iter(mfab).length == cnt + assert amr.MFIter(mfab).length == cnt def test_mfab_mfiter_keeps_mfab_alive(mfab, assert_keeps_python_alive): @@ -515,3 +517,253 @@ def test_mfab_copy(mfab): # check new mfab is the original data for i in range(new_mfab.num_comp): np.testing.assert_allclose(new_mfab.max(i), 42.0) + + +def test_mfab_mfiter_early_exit(mfab): + """Leaving an MFIter loop early must still finalize the iterator. + + Otherwise MFIter::depth stays at 1 and the *next* iterator construction + trips AMREX_ALWAYS_ASSERT(depth == 1), aborting the process. On GPU the + Gpu::streamSynchronize() in MFIter::Finalize() would also be skipped. + """ + + def count(): + return sum(1 for _mfi in mfab) + + n = count() + + for _mfi in mfab: + break + assert count() == n + + for _mfi in mfab.tiles(tile=(16, 16, 16)): + break + assert count() == n + + for _mfi in amr.MFIter(mfab, amr.TilingIfNotGPU((16, 16, 16))): + break + assert count() == n + + def raises(): + for _mfi in mfab: + raise ValueError("kernel error") + + with pytest.raises(ValueError): + raises() + assert count() == n + + # Finalize() is idempotent, so calling it by hand stays safe + for mfi in mfab: + mfi.finalize() + break + assert count() == n + + +def test_mfab_tiles(mfab): + nboxes = sum(1 for _mfi in mfab) + + # no tile size given -> whole boxes, same as plain iteration + assert sum(1 for _mfi in mfab.tiles()) == nboxes + + tiled = [mfi.tilebox() for mfi in mfab.tiles(tile=(16, 16, 16))] + if amr.Config.have_gpu: + # TilingIfNotGPU: never tile on GPU, where whole boxes are wanted + assert len(tiled) == nboxes + else: + # a tile smaller than the 32^3 boxes subdivides them + assert len(tiled) == 8 * nboxes + + # either way the tiles cover this rank's valid region exactly once + local_pts = sum(mfi.validbox().num_pts for mfi in mfab) + assert sum(bx.num_pts for bx in tiled) == local_pts + + # ix_type mirrors C++ mf.ixType().toIntVect() + assert mfab.ix_type == mfab.box_array().ix_type().to_IntVect() + + +def assert_xp_equal(actual, desired): + """Array equality that works on NumPy, CuPy and dpnp alike. + + numpy.testing on a device array raises "Implicit conversion to a NumPy + array is not allowed", so compare on the device and coerce only the + resulting scalar. + """ + assert actual.shape == desired.shape + assert bool((actual == desired).all()) + + +def test_mfab_array4_call(mfab): + """Array4.__call__ is a global-indexed, Box-shaped view.""" + ng = mfab.n_grow_vect + + for mfi in mfab: + bx = mfi.tilebox() + arr = mfab.array(mfi) + + view = arr(bx) + # exactly AMREX_SPACEDIM dimensions: an Array4 is always 4D (i,j,k,n) + # with extent-1 padding, and the unused axes must be dropped rather + # than sliced, or `comp` would land on k in a 1D/2D build + assert view.ndim == amr.Config.spacedim + assert view.shape == tuple(bx.size) + + # a view, not a copy: a write through it is visible in a fresh view. + # Checked behaviorally rather than via .base, which dpnp arrays do + # not have. + view[...] = 7.0 + assert bool((arr(bx) == 7.0).all()) + + # matches a hand-sliced reference against the fab's global lower bound + ref = arr.to_xp(copy=False, order="F") + lo = amr.lbound(arr) + lo = (lo.x, lo.y, lo.z) + expected = ref[ + tuple( + slice(bx.small_end[d] - lo[d], bx.big_end[d] - lo[d] + 1) + for d in range(amr.Config.spacedim) + ) + + (0,) + ] + assert_xp_equal(view, expected) + + # component selection + for comp in range(mfab.num_comp): + arr(bx, comp=comp)[...] = float(comp) + for comp in range(mfab.num_comp): + np.testing.assert_allclose(float(arr(bx, comp=comp).min()), float(comp)) + np.testing.assert_allclose(float(arr(bx, comp=comp).max()), float(comp)) + + # shifted access reaches the guard cells + if ng.max > 0: + shifted = arr(bx, di=-1) + assert shifted.shape == view.shape + assert_xp_equal( + shifted, + ref[ + tuple( + slice( + bx.small_end[d] + (-1 if d == 0 else 0) - lo[d], + bx.big_end[d] + (-1 if d == 0 else 0) - lo[d] + 1, + ) + for d in range(amr.Config.spacedim) + ) + + (0,) + ], + ) + del shifted + + # release the views: on a GPU build these are DLPack exports, and + # amrex.finalize() refuses to run while any are still alive + del view, ref, expected, arr + break + + +def test_mfab_for_each_tile(mfab): + """Serial, tiled and threaded runs must all agree.""" + + def fill(tile, threads): + mfab.set_val(0.0) + + @amr.for_each_tile(mfab, tile=tile, threads=threads) + def _(bx, a): + a(bx)[...] = 42.0 + + return [mfab.min(c) for c in range(mfab.num_comp)] + [ + mfab.max(c) for c in range(mfab.num_comp) + ] + + reference = fill(None, 1) + assert reference[0] == 42.0 + assert fill((16, 16, 16), 1) == reference + assert fill((16, 16, 16), 4) == reference + assert fill(None, 4) == reference + + # the decorator returns the kernel unchanged + @amr.for_each_tile(mfab) + def my_kernel(bx, a): + pass + + assert callable(my_kernel) + assert my_kernel.__name__ == "my_kernel" + + +def test_mfab_portable_kernel(boxarr, distmap): + """A ghost-cell stencil written once for CPU-serial, CPU-threaded and GPU. + + Ported from the 3D branch of ImpactX ForceFromSelfFields.cpp. The input + fields are linear in the index coordinates, so the central differences are + exact and can be checked against the analytic gradient. + """ + dr = (0.1, 0.2, 0.4) + coef = (1.0, 2.0, 3.0) + + phi = [amr.MultiFab(boxarr, distmap, 1, 1) for _ in range(3)] + for mfi in phi[0]: + gb = mfi.fabbox() + idx = amr.xp.meshgrid( + *[ + amr.xp.arange(gb.small_end[d], gb.big_end[d] + 1) + for d in range(amr.Config.spacedim) + ], + indexing="ij", + ) + for c in range(3): + phi[c].array(mfi)(gb)[...] = coef[c] * idx[c] + + field = amr.MultiFab(boxarr, distmap, 1, 0) + field.set_val(0.0) + + # Manual: Portable Kernel START + # C++ equivalent, e.g. ImpactX ForceFromSelfFields.cpp: + # + # #ifdef AMREX_USE_OMP + # #pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) + # #endif + # for (MFIter mfi(field, TilingIfNotGPU()); mfi.isValid(); ++mfi) + # { + # Box const& bx = mfi.tilebox(field.ixType().toIntVect()); + # Array4 const& f = field.array(mfi); + # Array4 const& px = phi_x.array(mfi); + # + # amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + # { + # f(i,j,k) = inv2dr[0] * (px(i-1,j,k) - px(i+1,j,k)) + ...; + # }); + # } + # + # In Python the decorator 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. Index the Array4s with the Box: f(bx) is the + # tile, px(bx, di=-1) is the same tile shifted by one cell in x. + # + # There is no OpenMP here: the Python loop is always serial, so on-node + # parallelism is either MPI ranks, `threads=` below, or the GPU. + inv2dr = [0.5 / d for d in dr] + + @amr.for_each_tile(field, *phi, tile=(16, 16, 16), threads=4) + def _(bx, f, px, py, pz): + f(bx)[...] = ( + inv2dr[0] * (px(bx, di=-1) - px(bx, di=+1)) + + inv2dr[1] * (py(bx, dj=-1) - py(bx, dj=+1)) + + inv2dr[2] * (pz(bx, dk=-1) - pz(bx, dk=+1)) + ) + + # Manual: Portable Kernel END + + expected = -sum(coef[c] / dr[c] for c in range(3)) + np.testing.assert_allclose(field.min(0), expected) + np.testing.assert_allclose(field.max(0), expected) + + # the untiled, unthreaded path must agree + field.set_val(0.0) + for mfi in field.tiles(): + bx = mfi.tilebox(field.ix_type) + f = field.array(mfi) + px, py, pz = (p.array(mfi) for p in phi) + f(bx)[...] = ( + inv2dr[0] * (px(bx, di=-1) - px(bx, di=+1)) + + inv2dr[1] * (py(bx, dj=-1) - py(bx, dj=+1)) + + inv2dr[2] * (pz(bx, dk=-1) - pz(bx, dk=+1)) + ) + np.testing.assert_allclose(field.min(0), expected) + np.testing.assert_allclose(field.max(0), expected)