From 7de72859bedf05720d8c5e3ce3b8dcf30e3cacf2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 7 Oct 2026 19:24:11 +0200 Subject: [PATCH] Fuse the pusher preamble into one CUDA kernel --- src/struphy/pic/pushing/prepare_push_cuda.cu | 28 ++++++++++++++ src/struphy/pic/pushing/pusher.py | 39 +++++++++++++++----- 2 files changed, 58 insertions(+), 9 deletions(-) create mode 100644 src/struphy/pic/pushing/prepare_push_cuda.cu diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu new file mode 100644 index 000000000..5c1dbd28e --- /dev/null +++ b/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; +} diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 7f3d64961..762dd3bd9 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/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} ...")