Fuse the pusher preamble into one CUDA kernel

src/struphy/pic/pushing/prepare_push_cuda.cu

@@ -0,0 +1,28 @@
+/**
+ * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push:
+ *
+ *     markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init]
+ *     markers[:, first_shift_idx:residual_idx] = 0.0
+ *     markers[:, residual_idx:-2] = 0.0
+ *
+ * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that
+ * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written
+ * in a single pass (instead of one strided pass over the buffer per slice).
+ *
+ * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes.
+ * @param n_rows Number of rows of the buffer.
+ * @param n_cols Number of columns of the buffer.
+ * @param first_init_idx First column of the saved initial phase space coordinates.
+ * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init).
+ */
+extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx,
+                                        int n_init) {
+    long long width = n_cols - 2 - first_init_idx;
+    long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x;
+    if (t >= n_rows * width) return;
+
+    long long ip = t / width;
+    int col = first_init_idx + (int)(t - ip * width);
+    double* row = markers + ip * n_cols;
+    row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0;
+}

src/struphy/pic/pushing/pusher.py

@@ -1,9 +1,11 @@
 "Accelerated particle pushing."
 
 import logging
+from functools import cache
+from pathlib import Path
 
 import cunumpy as xp
-from cunumpy.kernels import Kernel, PyccelKernel
+from cunumpy.kernels import CudaKernel, Kernel, PyccelKernel
 from line_profiler import profile
 from maybempi import MPI
 from scope_profiler import ProfileManager
@@ -16,6 +18,13 @@
 logger = logging.getLogger("struphy")
 
 
+@cache
+def _prepare_push_cuda() -> CudaKernel:
+    """The CUDA kernel that fuses the column-slice assignments at the start of :meth:`Pusher._push`
+    (compiled on first call)."""
+    return CudaKernel.from_file(Path(__file__).with_name("prepare_push_cuda.cu"), "prepare_push")
+
+
 class Pusher:
     r"""
     Class for solving particle ODEs
@@ -208,17 +217,29 @@ def _push(self, dt: float):
         logger.debug(f"{residual_idx =}")
         logger.debug(f"{self.particles.n_cols =}")
 
-        init_slice = slice(first_pusher_idx, first_shift_idx)
-        shift_slice = slice(first_shift_idx, residual_idx)
+        if self._cuda:
+            # the three assignments below in one pass over the buffer (each slice is a strided pass on its own)
+            n_rows, n_cols = markers.shape
+            _prepare_push_cuda()(
+                markers,
+                n_rows,
+                n_cols,
+                first_pusher_idx,
+                3 + vdim,
+                n_threads=n_rows * (n_cols - 2 - first_pusher_idx),
+            )
+        else:
+            init_slice = slice(first_pusher_idx, first_shift_idx)
+            shift_slice = slice(first_shift_idx, residual_idx)
 
-        # save initial phase space coordinates
-        markers[:, init_slice] = markers[:, : 3 + vdim]
+            # save initial phase space coordinates
+            markers[:, init_slice] = markers[:, : 3 + vdim]
 
-        # set boundary shifts to zero
-        markers[:, shift_slice] = 0.0
+            # set boundary shifts to zero
+            markers[:, shift_slice] = 0.0
 
-        # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box
-        markers[:, residual_idx:-2] = 0.0
+            # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box
+            markers[:, residual_idx:-2] = 0.0
 
         rank = self.particles.mpi_rank
         logger.debug(f"rank {rank}: starting {self.kernel} ...")