GPU stack 6/8: marker rows prepared inside the push kernels

src/struphy/kernel_arguments/pusher_args.cuh

@@ -19,6 +19,7 @@ struct MarkerArgs {
     int mu_idx;
     long long* bc_type;
     bool pole;
+    int prepare_rows;
 };
 
 struct DerhamArgs {

src/struphy/kernel_arguments/pusher_args_cuda.py

@@ -46,6 +46,8 @@ class CudaMarkerArguments(CudaStructArguments):
         ("mu_idx", "int"),
         ("bc_type", "long long*"),
         ("pole", "bool"),
+        # CUDA-only: 1 while a pusher has its kernel prepare each marker at stage 0 (Pusher._push), else 0
+        ("prepare_rows", "int"),
     )
 
     def __init__(
@@ -81,6 +83,7 @@ def __init__(
         self.mu_idx = mu_idx
         self.bc_type = _device_array("bc_type", bc_type, np.int64)
         self.pole = pole
+        self.prepare_rows = 0
         self.pack()
 
 

src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu

@@ -26,6 +26,9 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_
     int n_markers = args_markers.n_markers;
     if (ip >= n_markers) return;
 
+    // CUDA-only: prepare this marker here instead of in a separate pass of prepare_push (Pusher._push)
+    if (stage == 0 && args_markers.prepare_rows) pusher_utilities_kernels::prepare_marker_row(ip, args_markers);
+
     int first_pusher_idx = args_markers.first_pusher_idx;
     int first_free_idx = args_markers.first_free_idx;
     double last = stage == n_stages - 1 ? 1. : 0.;

src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu

@@ -2,6 +2,7 @@
 #include "struphy/bsplines/evaluation_kernels_3d.cuh"
 #include "struphy/geometry/evaluation_kernels.cuh"
 #include "struphy/linear_algebra/linalg_kernels.cuh"
+#include "struphy/pic/pushing/pusher_utilities_kernels.cuh"
 
 using namespace struphy_cuda;
 
@@ -30,6 +31,9 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar
     int n_markers = args_markers.n_markers;
     if (ip >= n_markers) return;
 
+    // CUDA-only: prepare this marker here instead of in a separate pass of prepare_push (Pusher._push)
+    if (stage == 0 && args_markers.prepare_rows) pusher_utilities_kernels::prepare_marker_row(ip, args_markers);
+
     int first_pusher_idx = args_markers.first_pusher_idx;
 
     // check if marker is a hole

src/struphy/pic/pushing/pusher.py

@@ -18,6 +18,11 @@
 logger = logging.getLogger("struphy")
 
 
+# CUDA kernels that prepare each marker themselves at stage 0 when MarkerArgs.prepare_rows is set (see
+# pusher_utilities_kernels.prepare_marker_row), so that Pusher._push can skip the separate pass of prepare_push
+_KERNELS_PREPARING_ROWS = frozenset({"push_eta_stage", "push_vxb_analytic"})
+
+
 @cache
 def _prepare_push_cuda() -> CudaKernel:
     """The CUDA kernel that fuses the column-slice assignments at the start of :meth:`Pusher._push`
@@ -220,7 +225,19 @@ def _push(self, dt: float):
         # only the markers (columns) that can hold markers; the columns after them are holes, which nothing reads
         used = markers[:, : self.particles.n_markers_used]
 
-        if self._cuda:
+        # On the GPU the kernels of _KERNELS_PREPARING_ROWS prepare each marker themselves at stage 0 (the same markers
+        # as below: one thread per used marker), saving a pass over the marker buffer; only if no other kernel runs
+        # before them (init and evaluation kernels may read the prepared rows).
+        prepare_in_kernel = (
+            self._cuda
+            and self.kernel.name in _KERNELS_PREPARING_ROWS
+            and not self.init_kernels
+            and not self.eval_kernels
+        )
+
+        if prepare_in_kernel:
+            pass
+        elif self._cuda:
             # the three assignments below in one launch
             n_cols, n_markers = used.shape
             _prepare_push_cuda()(
@@ -280,16 +297,21 @@ def _push(self, dt: float):
                 # sort according to alpha-weighted average
                 self._sort_for_alpha(self._alpha_in_kernel)
 
-                # push markers
+                # push markers (the first call prepares the rows if prepare_in_kernel)
+                args_markers = self.particles.args_markers
+                if prepare_in_kernel and stage == 0 and k == 1:
+                    args_markers.prepare_rows = 1
                 with ProfileManager.profile_region("kernel: " + self.kernel.name):
                     self.kernel(
                         dt,
                         stage,
-                        self.particles.args_markers,
+                        args_markers,
                         self._args_domain,
                         *self._args_kernel,
                         n_threads=self.particles.args_markers.n_markers,  # CUDA: one thread per marker
                     )
+                if prepare_in_kernel and args_markers.prepare_rows:
+                    args_markers.prepare_rows = 0
 
                 # markers have moved
                 if self.pushes_eta:

src/struphy/pic/pushing/pusher_utilities_kernels.cuh

@@ -28,6 +28,20 @@ __device__ inline void reflect_velocity(long long ip, Array2D<double> markers,
     for (int j = 0; j < 3; ++j) markers(3 + j, ip) = v[j];
 }
 
+/**
+ * Prepare marker ip for a push, as the kernel prepare_push (prepare_push_cuda.cu) does for every marker: copy the
+ * phase space coordinates (rows 0 to 3 + vdim of the buffer) to the rows from first_pusher_idx and zero the rows after
+ * them up to n_cols - 2. Called at stage 0 by the push kernels that take over the preparation from prepare_push
+ * (args_markers.prepare_rows, see Pusher._push), for every marker including holes, so that it is exactly the same.
+ */
+__device__ inline void prepare_marker_row(int ip, const MarkerArgs& args_markers) {
+    const int first = args_markers.first_pusher_idx;
+    const int n_init = 3 + args_markers.vdim;
+    const int end = (int)args_markers.markers.shape[0] - 2;
+    for (int j = first; j < end; ++j)
+        args_markers.markers(j, ip) = j < first + n_init ? args_markers.markers(j - first, ip) : 0.0;
+}
+
 /**
  * Apply kinetic boundary conditions, as in pusher_utilities_kernels.apply_kinetic_bc_marker.
  *

src/struphy/pic/tests/cuda_emulation.py

@@ -50,10 +50,13 @@ def _struct_field(value, name):
     """Field `name` of a struct, read from the pyccel argument object `value`.
 
     The CUDA structs have the attributes of their pyccel classes plus the knot lengths ``nt1``, ``nt2``,
-    ``nt3`` of ``DerhamArgs``, which pyccel gets from ``len(tn1)``.
+    ``nt3`` of ``DerhamArgs``, which pyccel gets from ``len(tn1)``, and the flag ``prepare_rows`` of ``MarkerArgs``.
     """
     if name in ("nt1", "nt2", "nt3"):
         return len(getattr(value, "tn" + name[-1]))
+    if name == "prepare_rows":
+        # CUDA-only flag of MarkerArgs (set by Pusher._push around a launch); 0: the markers are prepared beforehand
+        return getattr(value, name, 0)
     return getattr(value, name)
 
 

src/struphy/pic/tests/test_kernel_backends.py

@@ -37,7 +37,7 @@
 
 # (CUDA class, pyccel source, pyccel class, struct members that only the CUDA class has, its header)
 ARGUMENT_PAIRS = (
-    (CudaMarkerArguments, "pusher_args_kernels.py", "MarkerArguments", {"n_markers"}, HEADER),
+    (CudaMarkerArguments, "pusher_args_kernels.py", "MarkerArguments", {"n_markers", "prepare_rows"}, HEADER),
     (CudaDerhamArguments, "pusher_args_kernels.py", "DerhamArguments", {"nt1", "nt2", "nt3"}, HEADER),
     (CudaDomainArguments, "pusher_args_kernels.py", "DomainArguments", set(), HEADER),
     (