GPU stack 7/8: chunked marker loading, row-based weight update, leaner MPI sorting; static field in PushVxB

src/struphy/pic/base.py

@@ -1236,7 +1236,10 @@ def draw_markers(
                 # counting integers
                 num_loaded_particles_loc = 0  # number of particles alreday loaded (local)
                 num_loaded_particles_glob = 0  # number of particles already loaded (each clone)
-                chunk_size = 10000  # TODO: number of particle chunk
+                # markers drawn per chunk: large chunks keep the number of iterations (each a few small kernels and
+                # a device sync on the GPU) low, e.g. 16 instead of 1600 for 16M markers; a chunk of temporaries
+                # takes ~(3 + vdim + 4) * 8 bytes per marker (~50 MB). The NumPy draws do not depend on it.
+                chunk_size = 2**20
 
                 # Total number of markers to draw (sum over all clones)
                 while num_loaded_particles_glob < int(self.Np):
@@ -1561,16 +1564,26 @@ def update_weights(self):
         """
         from struphy.pic.particles import ParticlesSPH
 
-        if self.n_mks_loc == 0:
+        # the indices of the valid markers, found once: each property (positions, weights0, ...) would select them
+        # again with the boolean mask (and n_mks_loc counts them with one more pass and a host read)
+        idx = xp.flatnonzero(self.valid_mks)
+        if idx.size == 0:
             return
+        markers = self._markers
+
+        def quantity(j):
+            """Row j of the marker buffer (one quantity) at the valid markers."""
+            return xp.take(markers[j], idx)
 
         if isinstance(self, ParticlesSPH):
             f0 = self.f0.n0(self.positions)
         else:
             # in case of CanonicalMaxwellian, evaluate constants_of_motion
             # if isinstance(self.f0, CanonicalMaxwellian):
             #     self.save_constants_of_motion()
-            f0 = self.f0(*self.f_coords)
+            coords = self.f_coords_index
+            rows = range(markers.shape[0])[coords] if isinstance(coords, slice) else np.atleast_1d(coords).tolist()
+            f0 = self.f0(*[quantity(j) for j in rows])
 
         # if f_init is vol-form, transform to 0-form
         if self.is_volume_form[0]:
@@ -1579,7 +1592,7 @@ def update_weights(self):
         if self.is_volume_form[1]:
             f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords)
 
-        self.weights = self.weights0 - f0 / self.sampling_density_values / self.Np
+        markers[self.index["weights"], idx] = quantity(self.index["w0"]) - f0 / quantity(self.index["s0"]) / self.Np
 
     @profile
     def binning(
@@ -1823,10 +1836,25 @@ def mpi_sort_markers(
         (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`).
         """
+        # Ghost particles (-2 in the last row) only exist with sorting boxes (SPH ghost boxes,
+        # _prepare_ghost_particles); without them the ghost mask stays all False and its
+        # bookkeeping (full passes over the buffer, a nonzero) can be skipped.
+        ghosts_possible = self._initialized_sorting
+
+        # On the GPU the markers go through cunumpy.mpi.exchange, which waits for every send and
+        # receive (_sendrecv_markers_device): no barriers are needed around the exchange. The host path
+        # does not wait for its sends, so it keeps them.
+        barriers = xp.get_array_backend(self._markers) != "cupy"
+
         if remove_ghost:
-            self._remove_ghost_particles()
+            if ghosts_possible:
+                self._remove_ghost_particles()
+            else:
+                # what _remove_ghost_particles leaves when there are no ghosts: refreshed holes
+                self.update_holes()
 
-        self._Barrier()
+        if barriers:
+            self._Barrier()
 
         # before sorting, apply kinetic bc
         if apply_bc:
@@ -1835,32 +1863,38 @@ def mpi_sort_markers(
         if isinstance(alpha, int) or isinstance(alpha, float):
             alpha = (alpha, alpha, alpha)
 
-        # create new markers_to_be_sent array and make corresponding holes in markers array
-        hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha)
+        if self._use_device_sort():
+            # destinations in one kernel pass, one gather per destination rank (see _mpi_sort_device)
+            n_emptied, recv_info = self._mpi_sort_device(alpha)
+        else:
+            # create new markers_to_be_sent array and make corresponding holes in markers array
+            hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha)
 
-        # determine where to send markers_to_be_sent
-        send_info = self._sendrecv_get_destinations(send_inds)
+            # determine where to send markers_to_be_sent
+            send_info = self._sendrecv_get_destinations(send_inds)
 
-        # set new holes in markers array to -1
-        self._markers[:, send_inds] = -1.0
+            # set new holes in markers array to -1
+            self._markers[:, send_inds] = -1.0
 
-        # transpose send_info
-        recv_info = self._sendrecv_all_to_all(send_info)
+            # transpose send_info
+            recv_info = self._sendrecv_all_to_all(send_info)
 
-        # send and receive markers
-        self._sendrecv_markers(recv_info, hole_inds_after_send)
+            # send and receive markers
+            self._sendrecv_markers(recv_info, hole_inds_after_send)
+            n_emptied = int(send_inds.size)
 
         # the sent markers are holes now, the received ones filled the lowest holes (host counts, no device read)
         n_received = int(np.sum(recv_info)) - int(recv_info[self.mpi_rank])
-        self._track_markers(n_emptied=int(send_inds.size), n_filled=n_received)
+        self._track_markers(n_emptied=n_emptied, n_filled=n_received)
 
         # new holes and new number of holes and markers on process
         self.update_holes()
 
         # refresh ghost mask: received markers may land in columns that previously held
         # ghost particles. update_holes alone recomputes valid_mks from a stale
         # _ghost_particles mask, which would wrongly exclude these incoming real markers.
-        self._update_ghost_particles()
+        if ghosts_possible:
+            self._update_ghost_particles()
 
         # check if all markers are on the right process after sorting
         if do_test:
@@ -1874,7 +1908,8 @@ def mpi_sort_markers(
             assert all_on_right_proc
             # assert self.phasespace_coords.size > 0, f'No particles on process {self.mpi_rank}, please rebalance, aborting ...'
 
-        self._Barrier()
+        if barriers:
+            self._Barrier()
 
     @profile
     @ProfileManager.profile("apply_kinetic_bc")
@@ -4576,10 +4611,19 @@ def _sendrecv_determine_mtbs(
         n = self.n_markers_used
         bi = self.first_pusher_idx
         sorting_etas = self._sorting_etas[:, :n]
-        sorting_etas[:] = (
-            alpha * (self.markers[:3, :n] + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n])
-            + (1.0 - alpha) * self.markers[bi : bi + 3, :n]
-        )
+        if isinstance(alpha, float) and alpha == 1.0:
+            # the common case: alpha * (pos + shift) + 0 * init is exactly pos + shift for finite values,
+            # computed with one kernel instead of five
+            xp.add(
+                self.markers[:3, :n],
+                self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n],
+                out=sorting_etas,
+            )
+        else:
+            sorting_etas[:] = (
+                alpha * (self.markers[:3, :n] + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n])
+                + (1.0 - alpha) * self.markers[bi : bi + 3, :n]
+            )
 
         # for non-periodic axes, eta = 1 is the right end of the domain and must not be wrapped to eta = 0
         non_periodic_axes = [axis for axis in range(3) if axis not in self._periodic_axes]
@@ -4610,6 +4654,151 @@ def _sendrecv_determine_mtbs(
 
         return hole_inds_after_send, send_inds
 
+    def _use_device_sort(self) -> bool:
+        """Whether :meth:`mpi_sort_markers` uses :meth:`_mpi_sort_device`: markers on a GPU (not on cunumpy's fake CuPy,
+        which cannot launch kernels). ``STRUPHY_DEVICE_SORT=1`` or ``0`` forces it on (also on the fake CuPy, e.g. in
+        tests with emulated launches) or off."""
+        forced = os.environ.get("STRUPHY_DEVICE_SORT", "").strip().lower()
+        if forced in ("0", "false", "no"):
+            return False
+        if xp.get_array_backend(self._markers) != "cupy":
+            return False
+        if forced in ("1", "true", "yes"):
+            return True
+        from cunumpy.kernel_testing import fake_cupy_active
+
+        return not fake_cupy_active()
+
+    def _mpi_sort_device(self, alpha) -> tuple[int, np.ndarray]:
+        """The exchange of :meth:`mpi_sort_markers` for markers on the GPU.
+
+        The same result as :meth:`_sendrecv_determine_mtbs`, :meth:`_sendrecv_get_destinations` and
+        :meth:`_sendrecv_markers` (the same markers in the same columns), with fewer passes over the marker buffer:
+
+        - one kernel finds the destination of every marker (:func:`~struphy.pic.mpi_sort_cuda.sort_destinations`)
+          and counts the markers per destination, instead of a pass per direction and per rank;
+        - the leaving markers are listed once, stably sorted by destination (so the markers to each rank keep their
+          order, as before) and gathered once per destination rank.
+
+        Markers that leave towards no rank's domain (e.g. a NaN position) become holes and are not sent, as before.
+
+        Parameters
+        ----------
+        alpha : tuple
+            Weights of the sorting position, see :meth:`mpi_sort_markers`.
+
+        Returns
+        -------
+        n_emptied : int
+            Number of markers that left (became holes), for :meth:`_track_markers`.
+        recv_info : numpy.ndarray
+            Number of markers received from each rank.
+        """
+        from struphy.pic import mpi_sort_cuda
+
+        alpha_host = np.array(alpha, dtype=float).ravel()
+        assert alpha_host.size == 3
+        assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0)
+
+        if not hasattr(self, "_dev_sort_bufs"):
+            ranks = np.array(self._neighbor_ranks + self._non_neighbor_ranks, dtype=np.int32)
+            periodic = np.array([1 if axis in self._periodic_axes else 0 for axis in range(3)], dtype=np.int32)
+            bounds = mpi_sort_cuda.domain_bounds(self._domain_array_host)
+            self._dev_sort_bufs = {
+                "ranks": xp.asarray(ranks),
+                "periodic": xp.asarray(periodic),
+                "bounds": xp.asarray(bounds),
+                "dest": xp.empty(self.n_markers_max, dtype=np.int32),
+                "counts": xp.zeros(self.mpi_size + 2, dtype=np.uint64),
+                "alpha": xp.zeros(3, dtype=float),
+            }
+        bufs = self._dev_sort_bufs
+        use_alpha = not np.all(alpha_host == 1.0)
+        if use_alpha:
+            bufs["alpha"][:] = xp.asarray(alpha_host)
+
+        n = self.n_markers_used
+        bi = self.first_pusher_idx
+        dest = bufs["dest"][:n]
+        counts_dev = bufs["counts"]
+        counts_dev.fill(0)
+        mpi_sort_cuda.sort_destinations(
+            self._markers,
+            self.holes,
+            self.ghost_particles,
+            bi + 3 + self.vdim,
+            bi,
+            bufs["alpha"],
+            use_alpha,
+            bufs["periodic"],
+            bufs["bounds"],
+            bufs["ranks"],
+            self.mpi_rank,
+            n,
+            bufs["dest"],
+            counts_dev,
+        )
+        # markers per destination (lost first): a read of 8 bytes per rank to size the messages, like the one the
+        # nonzero() of the host path makes to size its result (marker data stays on the device)
+        counts = xp.to_numpy(counts_dev).astype(int)
+        n_lost = int(counts[0])
+        send_info = counts[2:].copy()
+
+        # the leaving markers, sorted by destination (stable: ascending columns per destination); lost markers first
+        leaving = dest != -1
+        leave_inds = xp.flatnonzero(leaving)
+        n_leaving = n_lost + int(send_info.sum())
+        if n_leaving:
+            leave_sorted = leave_inds[xp.argsort(dest[leave_inds], kind="stable")]
+        else:
+            leave_sorted = leave_inds
+
+        # the columns that are holes after the send: the old holes (all columns from n_markers_used on are holes)
+        # and the leaving markers; the received markers fill the lowest of them
+        holes_after = self.holes.copy()
+        xp.logical_or(holes_after[:n], leaving, out=holes_after[:n])
+        hole_inds_after_send = xp.flatnonzero(holes_after)
+
+        # one contiguous buffer (n_cols, count) per destination rank, in the order of the host path
+        first_send = np.cumsum(send_info) - send_info
+        sends = []
+        for rank in range(self.mpi_size):
+            if rank == self.mpi_rank or send_info[rank] == 0:
+                self._send_list[rank] = self._markers[:, :0]
+                continue
+            start = n_lost + int(first_send[rank])
+            cols = leave_sorted[start : start + int(send_info[rank])]
+            self._send_list[rank] = xp.ascontiguousarray(self._markers[:, cols])
+            sends.append((self._send_list[rank], rank, self.mpi_rank))
+
+        # the leaving markers are holes now (also the lost ones)
+        if n_leaving:
+            self._markers[:, leave_inds] = -1.0
+
+        recv_info = self._sendrecv_all_to_all(send_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 "
+                f"{self.mpi_rank} is smaller than number of incoming particles ({n_recv}). "
+                'Increase "bufsize" in the markers parameters for the next run.',
+            )
+            self.mpi_comm.Abort()
+
+        first_recv = np.cumsum(recv_info) - recv_info
+        receives = []
+        for rank in range(self.mpi_size):
+            if rank != self.mpi_rank and recv_info[rank] > 0:
+                receives.append((xp.empty((self.n_cols, int(recv_info[rank])), dtype=self._markers.dtype), rank, rank))
+
+        xp.mpi.exchange(self.mpi_comm, sends=sends, receives=receives, cuda_aware=self._mpi_cuda_aware())
+
+        for recvbuf, rank, _ in receives:
+            start = int(first_recv[rank])
+            self._markers[:, hole_inds_after_send[start : start + int(recv_info[rank])]] = recvbuf
+
+        return n_leaving, recv_info
+
     def _compute_neighbor_ranks(self) -> tuple[list[int], list[int]]:
         """Split every other rank into geometric neighbours of this rank's sub-domain
         box and everyone else, for :meth:`_sendrecv_get_destinations`.

src/struphy/pic/mpi_sort_cuda.cu

@@ -0,0 +1,79 @@
+// Destinations of the markers in Particles.mpi_sort_markers on the GPU (struphy.pic.mpi_sort_cuda).
+#include "cunumpy/array_view.cuh"
+#include "cunumpy/atomic.cuh"
+#include "cunumpy/index.cuh"
+
+// The domain test of Particles._is_on_domain_of_rank for one rank: [left, right) in every direction.
+__device__ inline bool on_domain(const double* eta, const double* bounds, int r) {
+    for (int a = 0; a < 3; ++a) {
+        const double left = bounds[(r * 3 + a) * 2], right = bounds[(r * 3 + a) * 2 + 1];
+        if (!(eta[a] >= left && eta[a] < right)) return false;
+    }
+    return true;
+}
+
+/**
+ * The rank each marker has to go to, as Particles._sendrecv_determine_mtbs and _sendrecv_get_destinations decide it.
+ *
+ * Marker ip (a column of the buffer, ip < n_used) is sorted by eta = pos + shift (use_alpha == 0) or
+ * alpha * (pos + shift) + (1 - alpha) * init per direction (rows 0-2, shift_row.., init_row..), reduced modulo 1 as
+ * numpy.mod does; on a non-periodic axis eta = 1 stays 1. It stays (dest -1) if it is a hole or a ghost or lies in
+ * the domain of this rank, else it goes to the rank whose domain contains eta (bounds: [left, right) per rank and
+ * direction, the right end 1 shifted to the next float, see domain_bounds), or it is lost (dest -2) if no domain
+ * contains it (e.g. a NaN position). The domains do not overlap, so the rank is unique.
+ *
+ * One thread per marker (n_threads = n_used). counts[dest + 2] counts the leaving markers (dest -2 or a rank).
+ *
+ * @param markers Marker buffer (n_cols x n_markers_max, one column per marker).
+ * @param holes, ghosts Hole and ghost masks (n_markers_max).
+ * @param shift_row, init_row First rows of the boundary shifts and of the initial positions.
+ * @param alpha, use_alpha The three weights, used if use_alpha (else alpha = 1, eta = pos + shift).
+ * @param periodic Per direction, 1 if periodic.
+ * @param bounds (n_ranks, 3, 2): left and right end of each rank's domain per direction.
+ * @param ranks The other ranks in the order they are checked (neighbours first), n_others entries.
+ * @param n_others Number of entries of ranks.
+ * @param rank This rank.
+ * @param n_used Number of markers (columns) that can hold markers.
+ * @param dest Output (n_markers_max): -1 stay, -2 lost, else the destination rank (written for ip < n_used).
+ * @param counts Output (n_ranks + 2), added to: leaving markers per destination (dest + 2).
+ */
+extern "C" __global__ void sort_destinations(Array2D<double> markers, const bool* holes, const bool* ghosts,
+                                             int shift_row, int init_row, const double* alpha, int use_alpha,
+                                             const int* periodic,
+                                             const double* bounds, const int* ranks, int n_others, int rank,
+                                             long long n_used, int* dest, unsigned long long* counts) {
+    CUNUMPY_THREAD_1D(ip, n_used);
+    if (holes[ip] || ghosts[ip]) {
+        dest[ip] = -1;
+        return;
+    }
+
+    double eta[3];
+    for (int a = 0; a < 3; ++a) {
+        const double moved = markers(a, ip) + markers(shift_row + a, ip);
+        double e = use_alpha ? alpha[a] * moved + (1.0 - alpha[a]) * markers(init_row + a, ip) : moved;
+        const bool right_end = !periodic[a] && e == 1.0;
+        // numpy.mod(e, 1.0): the remainder with the sign of the divisor
+        double r = fmod(e, 1.0);
+        if (r != 0.0) {
+            if (r < 0.0) r += 1.0;
+        } else {
+            r = 0.0;
+        }
+        eta[a] = right_end ? 1.0 : r;
+    }
+
+    if (on_domain(eta, bounds, rank)) {
+        dest[ip] = -1;
+        return;
+    }
+    int d = -2;
+    for (int k = 0; k < n_others; ++k) {
+        if (on_domain(eta, bounds, ranks[k])) {
+            d = ranks[k];
+            break;
+        }
+    }
+    dest[ip] = d;
+    cunumpy_atomic_add(&counts[d + 2], 1ULL);
+}