From 77a25b5613dd93b99af23651552b83fa59ca0dc7 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 1/2] Marker exchange between ranks on the device (CUDA-aware MPI) On the CuPy backend, Particles._sendrecv_markers hands the marker rows to Isend/Irecv as device buffers through cunumpy.mpi.mpi_buffer (directly after synchronize_for_mpi with CUDA-aware MPI, staged through host memory otherwise). Rows from each rank arrive in a contiguous slice of one device receive buffer, and one device scatter fills the holes. The per-rank counts are small NumPy arrays on every backend (they are host integers already), and a uniform alpha is a Python scalar, so the exchange makes no host/device copy and no implicit host read of a device array. The NumPy path is unchanged. Adds pic/tests/test_sorting_device.py (fake CuPy under mpirun, and a GPU variant) and a CI step running it on the fake CuPy. Solves #698, part of #650. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-PR-unit-mpi.yml | 6 + CUDA_STRATEGY.md | 20 +- src/struphy/pic/base.py | 129 ++++++++++++- src/struphy/pic/tests/test_sorting_device.py | 190 +++++++++++++++++++ 4 files changed, 338 insertions(+), 7 deletions(-) create mode 100644 src/struphy/pic/tests/test_sorting_device.py diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index e3a13dbc5..7bb53fbe2 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -76,3 +76,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} pytest -v --with-mpi -m mpi_pic pic/tests + + - name: Run the marker exchange on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 6d2c60e72..68035f68c 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -365,7 +365,7 @@ matrix data). - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side. -- **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. +- **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. - **Single-source alternatives.** Hand-written CUDA stays the default. Generating whole kernels from the Python source (`cupyx.jit`, numba-cuda, or a pyccel CUDA backend) is worth a look before the guiding-center kernels (the largest ones) are ported. Those tools take flat arguments, which `fields` also provides. - **Polar splines and MHD equilibria on the GPU** (left after PR 19): spline mappings run on the device, but `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the @@ -644,3 +644,21 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## Marker exchange implementation notes (#698) + +- `Particles._sendrecv_markers` on CuPy (`_sendrecv_markers_device`): the rows to send (gathered on the device per + destination) and one device receive buffer go to `Isend`/`Irecv` through `xp.mpi.mpi_buffer`, i.e. directly after + `synchronize_for_mpi` with CUDA-aware MPI, staged through (pinned) host memory otherwise (with a warning, once). + Rows from rank `i` arrive in a contiguous slice of the receive buffer; one device scatter fills the holes. + Zero-count messages are skipped (both sides know the counts). The NumPy path is unchanged. +- Whether MPI is CUDA-aware: the answer recorded by `xp.mpi.mpi_is_cuda_aware()`/`set_mpi_cuda_aware()`, else a + collective probe of `mpi_comm` at the first exchange (`mpi_sort_markers` is collective). +- Counts: `send_info`/`recv_info` are small NumPy arrays on every backend. The send counts are sizes of `nonzero` + results, which are host integers already (CuPy synchronizes to size them), so the `Alltoall` of the counts needs no + copy. Uniform `alpha` is a Python scalar (no host-to-device copy of a 3-vector). +- Tests: `pic/tests/test_sorting_device.py` sorts the same markers on NumPy and on CuPy over all ranks and compares + the marker arrays; with CUDA-aware MPI no transfer (`count_transfers`) and no implicit host read of a device array + during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake + CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI + workflow); the GPU variant has not run (no GPU here). diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 99db7dbfd..cae11c476 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,6 +2,7 @@ import os import warnings from abc import ABCMeta, abstractmethod +from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -1757,6 +1758,12 @@ def mpi_sort_markers( remove_ghost : bool Remove ghost particles before send. + + Notes + ----- + On the CuPy backend, the marker rows are exchanged as device buffers with CUDA-aware MPI + (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The + per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ if remove_ghost: self._remove_ghost_particles() @@ -4363,9 +4370,19 @@ def _sendrecv_determine_mtbs( """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): - alpha = xp.array(alpha, dtype=float) - assert alpha.size == 3 - assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + # check a tuple/list on the host: on CuPy, asserting on a device comparison reads it back + alpha_host = np.array(alpha, dtype=float) + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + if alpha_host[0] == alpha_host[1] == alpha_host[2]: + # the common case (alpha = 1 or alpha = 0.5 in every direction): a Python scalar, + # which needs no host-to-device copy on CuPy and gives the same values on NumPy + alpha = float(alpha_host[0]) + else: + alpha = xp.asarray(alpha_host) + else: + assert alpha.size == 3 + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx self._sorting_etas[:] = ( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) @@ -4465,8 +4482,11 @@ def _sendrecv_get_destinations(self, send_inds): Amount of particles sent to i-th process. """ - # One entry for each process - send_info = xp.zeros(self.mpi_size, dtype=int) + # One entry for each process. The counts are host integers on every backend: the + # size of a nonzero() result is known on the host (CuPy synchronizes to size it), so + # a small NumPy array holds them without any host/device copy, and the Alltoall of + # the counts (_sendrecv_all_to_all) passes host buffers to MPI. + send_info = np.zeros(self.mpi_size, dtype=int) # Gathered once and reused for every rank below, instead of re-gathering # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every @@ -4530,9 +4550,16 @@ def _sendrecv_all_to_all(self, send_info): ------- recv_info : array[int] Amount of marticles to be received from i-th process. + + Notes + ----- + ``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every + backend: the counts size the receive buffers and the holes to fill, which Python needs + on the host anyway. On CuPy, only the marker rows themselves are device buffers + (:meth:`_sendrecv_markers`). """ - recv_info = xp.zeros(self.mpi_size, dtype=int) + recv_info = np.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) @@ -4551,6 +4578,10 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): Indices of empty rows in markers after send. """ + if xp.get_array_backend(self._markers) == "cupy": + self._sendrecv_markers_device(recv_info, hole_inds_after_send) + return + # i-th entry holds the number (not the index) of the first hole to be filled by data from process i first_hole = xp.cumsum(recv_info) - recv_info @@ -4594,6 +4625,92 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): for req in send_reqs: req.Wait() + def _mpi_cuda_aware(self) -> bool: + """Whether MPI takes device (CuPy) buffers, for :meth:`_sendrecv_markers_device`. + + Uses the answer recorded by :func:`cunumpy.mpi.mpi_is_cuda_aware` or + :func:`cunumpy.mpi.set_mpi_cuda_aware` (e.g. by feectools or at startup). If nothing + was recorded, probes :attr:`mpi_comm` once; the probe is collective, which is safe + here because :meth:`mpi_sort_markers` is called by every rank of :attr:`mpi_comm`. + """ + cuda_aware = xp.mpi.get_mpi_cuda_aware() + if cuda_aware is None: + cuda_aware = xp.mpi.mpi_is_cuda_aware(self.mpi_comm) + if not cuda_aware and not Particles._warned_mpi_staging: + Particles._warned_mpi_staging = True + logger.warning( + "MPI is not CUDA-aware: marker rows exchanged between ranks are staged through host memory.", + ) + return cuda_aware + + _warned_mpi_staging = False + + def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): + """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. + + With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, + gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive + buffer are handed to ``Isend``/``Irecv`` directly, after + :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI + reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of + the receive buffer, and one device scatter puts all of them into the holes. Without + CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host + memory (counted by :func:`cunumpy.profiling.count_transfers`). + + The bookkeeping stays on the host without host/device copies: the counts are host + integers (``recv_info`` from the Alltoall, the send counts are the row counts of the + send buffers), and so is the number of holes. Ranks exchange no message when the + count is zero (both sides know the count, so they skip consistently). + + Parameters + ---------- + recv_info : array[int] + Amount of markers to be received from i-th process (NumPy array). + + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). + """ + recv_info = np.asarray(recv_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f'Strong load imbalance detected: \ +number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ +is smaller than number of incoming particles ({n_recv}). \ +Increasing the value of "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + cuda_aware = self._mpi_cuda_aware() + + # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + first_row = np.cumsum(recv_info) - recv_info + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) + + with ExitStack() as buffers: + if n_recv > 0: + recv_mpi = buffers.enter_context( + xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), + ) + reqs = [] + for i in range(self.mpi_size): + if i == self.mpi_rank: + continue + if recv_info[i] > 0: + rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] + reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) + data = self._send_list[i] + if data.shape[0] > 0: + send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) + reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) + # wait inside the block: the buffers must stay valid until MPI is done, and staged + # receive rows are copied to the device when the block ends + for req in reqs: + req.Wait() + + if n_recv > 0: + self._markers[hole_inds_after_send[:n_recv]] = recvbuf + class Tesselation: """ diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py new file mode 100644 index 000000000..55c5d09f5 --- /dev/null +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -0,0 +1,190 @@ +"""Marker exchange between ranks on the CuPy backend (:meth:`Particles.mpi_sort_markers`). + +The marker rows are exchanged as device buffers with CUDA-aware MPI, and staged through +host memory otherwise. Without a GPU, the tests run on cunumpy's fake CuPy, which must be +installed before Python imports cunumpy, so they are skipped unless the fake is active:: + + mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi src/struphy/pic/tests/test_sorting_device.py + +Fake CuPy arrays live in host memory, so a non-CUDA-aware MPI library takes them, and the +CUDA-aware path (the probe of :func:`cunumpy.mpi.mpi_is_cuda_aware` succeeds) is what runs by +default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. +""" + +import sys +from contextlib import contextmanager + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from maybempi import MPI + +from struphy import BoundaryParameters, LoadingParameters +from struphy.pic.particles import Particles6D + + +def _fake_cupy_active() -> bool: + return bool(getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False)) + + +requires_fake_cupy = pytest.mark.skipif( + not _fake_cupy_active(), + reason="needs the fake CuPy: run under `mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest ...`", +) + + +@contextmanager +def _count_host_reads(): + """Count reads of fake CuPy arrays by the host (``int()``, ``bool()``, iteration, ``.get()``, ...). + + :func:`cunumpy.profiling.count_transfers` sees the copies made through cunumpy; this also sees + the implicit ones (a device scalar in an ``if``, ``list(device_array)``), which synchronize + and copy on a GPU. + """ + import cupy + + reads = [] + names = ("__int__", "__float__", "__bool__", "__index__", "__iter__", "get") + originals = {name: getattr(cupy.ndarray, name) for name in names} + original_getattr = cupy.ndarray.__getattr__ + + def counting(name, method): + def wrapper(self, *args, **kwargs): + reads.append(name) + return method(self, *args, **kwargs) + + return wrapper + + def counting_getattr(self, name): + if name in ("item", "tolist"): + reads.append(name) + return original_getattr(self, name) + + for name, method in originals.items(): + setattr(cupy.ndarray, name, counting(name, method)) + cupy.ndarray.__getattr__ = counting_getattr + try: + yield reads + finally: + for name, method in originals.items(): + setattr(cupy.ndarray, name, method) + cupy.ndarray.__getattr__ = original_getattr + + +def _make_particles(comm, bc): + return Particles6D( + comm_world=comm, + loading_params=LoadingParameters(Np=2000, seed=1234), + boundary_params=BoundaryParameters(bc=(bc, bc, bc)), + ) + + +def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.ndarray: + """New positions in [0, 1) for the valid rows, far enough away that most markers change rank.""" + new = markers.copy() + shift = np.array([0.37, 0.21, 0.53]) * (step + 1) + new[valid, :3] = (markers[valid, :3] + shift) % 1.0 + return new + + +def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): + """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. + + ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` + force the device-buffer or the host-staged exchange. + """ + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + + ref = _make_particles(comm, bc) + ref.draw_markers(sort=False) + ref.mpi_sort_markers() + + # the reference is sorted on NumPy, the device markers inside `use_backend("cupy")` blocks + on_device = lambda: xp.use_backend("cupy") # noqa: E731 + + previous = xp.mpi.get_mpi_cuda_aware() + try: + with on_device(): + if cuda_aware is not None: + xp.mpi.set_mpi_cuda_aware(cuda_aware) + else: + xp.mpi.set_mpi_cuda_aware(None) + cuda_aware = xp.mpi.mpi_is_cuda_aware(comm) + + dev = _make_particles(comm, bc) + assert xp.get_array_backend(dev.markers) == "cupy" + assert dev.markers.shape == ref.markers.shape + + for step in range(2): + new = _shift_positions(ref.markers, ref.valid_mks, step) + + # number of rows this rank has to send, from the host reference + etas = new[ref.valid_mks, :3] + left, right = ref.domain_array[rank, 0::3], ref.domain_array[rank, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=1))) + n_send_total = comm.allreduce(n_send) + assert size == 1 or n_send_total > 0 + + ref.markers[:] = new + ref.update_holes() + n_before = ref.n_mks_loc + ref.mpi_sort_markers(do_test=True) + n_recv = ref.n_mks_loc - (n_before - n_send) + + with on_device(): + dev.markers[:] = xp.to_cupy(new) + dev.update_holes() + with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + dev.mpi_sort_markers() + + if cuda_aware: + assert transfers.total == 0, transfers.report() + else: + # one staged copy per non-empty send buffer, one copy back of the receive buffer + n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") + assert sent_bytes == n_send * ref.markers.shape[1] * 8 + + if count_reads and cuda_aware: + # (the staged exchange reads the send buffers with .get(), counted above) + assert reads == [], f"host reads of device arrays during the exchange: {reads}" + + with on_device(): + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + assert np.array_equal(xp.to_numpy(dev.valid_mks), ref.valid_mks) + n_dev = dev.n_mks_loc + assert n_dev == ref.n_mks_loc + + with on_device(): + # all markers are on the right rank now: nothing moves + dev.mpi_sort_markers(do_test=True) + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + finally: + xp.mpi.set_mpi_cuda_aware(previous) + + +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [None, True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): + """On the fake CuPy: the exchange gives the NumPy result, without host/device copies when MPI is CUDA-aware.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_gpu(bc): + """On a GPU (``mpirun -n 2``, one GPU per rank): the exchange gives the NumPy result; with + CUDA-aware MPI without host/device copies.""" + check_mpi_sort_markers_device(None, bc=bc, count_reads=False) + + +if __name__ == "__main__": + check_mpi_sort_markers_device(None) From 0ea9085afb0ee05ad4ecaa5e3da9b9222d8c7759 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 9 Oct 2026 16:21:37 +0200 Subject: [PATCH 2/2] Use xp.mpi.exchange --- pyproject.toml | 2 +- src/struphy/pic/base.py | 85 +++++++++----------- src/struphy/pic/tests/test_sorting_device.py | 7 +- 3 files changed, 42 insertions(+), 52 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c75808225..fdeda2015 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy >= 0.6.1, <= 0.6.1", + "cunumpy @ git+https://github.com/max-models/cunumpy.git@26a640b1d47d5677663785e49b1c44d343540df1", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index cae11c476..5e1d217b2 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,7 +2,6 @@ import os import warnings from abc import ABCMeta, abstractmethod -from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -4646,67 +4645,55 @@ def _mpi_cuda_aware(self) -> bool: _warned_mpi_staging = False def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): - """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. - - With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, - gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive - buffer are handed to ``Isend``/``Irecv`` directly, after - :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI - reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of - the receive buffer, and one device scatter puts all of them into the holes. Without - CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host - memory (counted by :func:`cunumpy.profiling.count_transfers`). - - The bookkeeping stays on the host without host/device copies: the counts are host - integers (``recv_info`` from the Alltoall, the send counts are the row counts of the - send buffers), and so is the number of holes. Ranks exchange no message when the - count is zero (both sides know the count, so they skip consistently). + """Exchange marker rows and scatter incoming rows into empty slots. + + :func:`cunumpy.mpi.exchange` handles GPU synchronization, optional host + staging, and completion of all MPI transfers. Parameters ---------- - recv_info : array[int] - Amount of markers to be received from i-th process (NumPy array). - - hole_inds_after_send : array[int] - Indices of empty rows in markers after send (device array). + recv_info : array[int] + Number of markers to receive from each process (NumPy array). + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). """ recv_info = np.asarray(recv_info) n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: warnings.warn( - f'Strong load imbalance detected: \ -number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ -is smaller than number of incoming particles ({n_recv}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.', + f"Strong load imbalance detected: " + f"number of holes ({hole_inds_after_send.size}) " + f"on rank {self.mpi_rank} is smaller than number of " + f"incoming particles ({n_recv}). " + 'Increase "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() - cuda_aware = self._mpi_cuda_aware() - - # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=self.markers.dtype) first_row = np.cumsum(recv_info) - recv_info - recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) - with ExitStack() as buffers: - if n_recv > 0: - recv_mpi = buffers.enter_context( - xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), - ) - reqs = [] - for i in range(self.mpi_size): - if i == self.mpi_rank: - continue - if recv_info[i] > 0: - rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] - reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) - data = self._send_list[i] - if data.shape[0] > 0: - send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) - reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) - # wait inside the block: the buffers must stay valid until MPI is done, and staged - # receive rows are copied to the device when the block ends - for req in reqs: - req.Wait() + sends = [] + receives = [] + for rank in range(self.mpi_size): + if rank == self.mpi_rank: + continue + + if recv_info[rank] > 0: + start = int(first_row[rank]) + stop = start + int(recv_info[rank]) + receives.append((recvbuf[start:stop], rank, rank)) + + data = self._send_list[rank] + if data.shape[0] > 0: + sends.append((data, rank, self.mpi_rank)) + + xp.mpi.exchange( + self.mpi_comm, + sends=sends, + receives=receives, + cuda_aware=self._mpi_cuda_aware(), + ) if n_recv > 0: self._markers[hole_inds_after_send[:n_recv]] = recvbuf diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index 55c5d09f5..dc8fa9462 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -143,12 +143,15 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", if cuda_aware: assert transfers.total == 0, transfers.report() else: - # one staged copy per non-empty send buffer, one copy back of the receive buffer + # exchange stages each non-empty send and receive buffer separately n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[0] > 0) assert transfers.to_host == n_dest, transfers.report() - assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") assert sent_bytes == n_send * ref.markers.shape[1] * 8 + received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") + assert received_bytes == n_recv * ref.markers.shape[1] * 8 if count_reads and cuda_aware: # (the staged exchange reads the send buffers with .get(), counted above)