Marker exchange between ranks on the device (CUDA-aware MPI)

.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

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 `<name>_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).

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",

src/struphy/pic/base.py

@@ -1757,6 +1757,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 +4369,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 +4481,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 +4549,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 +4577,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 +4624,80 @@ 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):
+        """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]
+            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: "
+                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()
+
+        recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=self.markers.dtype)
+        first_row = np.cumsum(recv_info) - recv_info
+
+        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
+
 
 class Tesselation:
     """