From 33ccffc2d38e78888b13989cf0cef19aecdcd9fd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:47 +0200 Subject: [PATCH 1/2] GPU marker buffer: column-major storage; marker rows prepared inside the push kernels - Particles allocates the marker buffer column-major on the CuPy backend: kernels and the Python column operations access a few columns of many markers, so with a row-major buffer every access touched whole rows and the threads of a warp read addresses a row apart. The CUDA kernels address the buffer through its strides (Array2D in MarkerArgs); NumPy keeps C order for the pyccel kernels. - push_eta_stage and push_vxb_analytic prepare each marker row themselves at stage 0 (MarkerArgs.prepare_rows), so Pusher._push skips the separate prepare_push pass when no init or evaluation kernel runs before them. prepare_push is launched on a 2D grid. Co-Authored-By: Claude Opus 5.5 --- src/struphy/kernel_arguments/pusher_args.cuh | 1 + .../kernel_arguments/pusher_args_cuda.py | 18 +++++++-- src/struphy/pic/base.py | 7 +++- .../push_eta_stage/push_eta_stage_cuda.cu | 3 ++ .../push_vxb_analytic_cuda.cu | 4 ++ src/struphy/pic/pushing/prepare_push_cuda.cu | 26 ++++++------- src/struphy/pic/pushing/pusher.py | 37 +++++++++++++------ .../pic/pushing/pusher_utilities_kernels.cuh | 14 +++++++ src/struphy/pic/tests/test_kernel_backends.py | 2 +- 9 files changed, 80 insertions(+), 32 deletions(-) diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh index e3a547e1e..0ace87644 100644 --- a/src/struphy/kernel_arguments/pusher_args.cuh +++ b/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 { diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index feac42098..9b2a2cabc 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -12,15 +12,21 @@ from cunumpy.arguments import CudaStructArguments -def _device_array(name, arr, dtype, ndim=None): - """Check that ``arr`` is a C-contiguous CuPy array of ``dtype`` (and ``ndim``) on the current device (never copies).""" +def _device_array(name, arr, dtype, ndim=None, strided=False): + """Check that ``arr`` is a CuPy array of ``dtype`` (and ``ndim``) on the current device (never copies). + + It must be C-contiguous, unless ``strided``: arrays passed as strided views (``Array2D``, e.g. the + column-major marker buffer) may have any strides. + """ import cupy as cp if not isinstance(arr, cp.ndarray): raise TypeError(f"{name} must be a CuPy array, got {type(arr)}") if arr.device.id != cp.cuda.runtime.getDevice(): raise ValueError(f"{name} must be on the current CUDA device") - if arr.dtype != dtype or not arr.flags.c_contiguous: + if arr.dtype != dtype: + raise TypeError(f"{name} must be an array of dtype {np.dtype(dtype)}") + if not strided and not arr.flags.c_contiguous: raise TypeError(f"{name} must be a C-contiguous array of dtype {np.dtype(dtype)}") if ndim is not None and arr.ndim != ndim: raise TypeError(f"{name} must be a {ndim}-dimensional array, got {arr.ndim} dimensions") @@ -46,6 +52,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 row at stage 0 (Pusher._push), else 0 + ("prepare_rows", "int"), ) def __init__( @@ -64,7 +72,8 @@ def __init__( bc_type, pole: bool, ): - self.markers = _device_array("markers", markers, np.float64) + # a strided view in MarkerArgs: Particles allocates the buffer column-major on the GPU + self.markers = _device_array("markers", markers, np.float64, strided=True) if markers.ndim != 2: raise TypeError("markers must be a two-dimensional array.") self.valid_mks = _device_array("valid_mks", valid_mks, np.bool_) @@ -80,6 +89,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() diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index e3529b9c4..8f238267b 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2634,7 +2634,12 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float) + # column-major on the GPU: the kernels (and the column operations on the Python side) access a few columns of + # many markers, so with a row-major buffer every access touches the whole row (n_cols * 8 bytes) and the + # threads of a warp read addresses a row apart. The CUDA kernels address the buffer through its strides + # (Array2D in MarkerArgs). The pyccel kernels expect C order, so NumPy keeps it. + order = "F" if xp.get_backend() == "cupy" else "C" + self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float, order=order) # allocate auxiliary arrays self._holes = xp.zeros(self.n_rows, dtype=bool) diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 9a6c24d52..eeb1ac9ca 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/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 row 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.; diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index 7d2a270ec..d4822afcf 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/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 row 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 diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 71c02556d..d22518b55 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -5,24 +5,20 @@ * markers[:, first_shift_idx:residual_idx] = 0.0 * markers[:, residual_idx:-2] = 0.0 * - * One thread per entry of the contiguous column range [first_pusher_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). + * A 2D launch, n_threads = (n_rows, n_cols - 2 - first_pusher_idx): x is the row, y the column within the range + * [first_pusher_idx, n_cols - 2), so the whole range is written in a single pass. The buffer is addressed through + * a strided view: on the CuPy backend it is column-major (Particles allocates it in Fortran order), so neighbouring + * threads (rows) of one column write neighbouring addresses. * - * @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 markers Marker buffer (n_rows x n_cols, any strides), every row including holes. * @param first_pusher_idx First column of the saved initial phase space coordinates. * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ -extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_pusher_idx, - int n_init) { - long long width = n_cols - 2 - first_pusher_idx; - long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; - if (t >= n_rows * width) return; +#include "cunumpy/array_view.cuh" - long long ip = t / width; - int col = first_pusher_idx + (int)(t - ip * width); - double* row = markers + ip * n_cols; - row[col] = col < first_pusher_idx + n_init ? row[col - first_pusher_idx] : 0.0; +extern "C" __global__ void prepare_push(Array2D markers, int first_pusher_idx, int n_init) { + const long long ip = (long long)blockDim.x * blockIdx.x + threadIdx.x; + const long long k = (long long)blockDim.y * blockIdx.y + threadIdx.y; + if (ip >= markers.shape[0] || k >= markers.shape[1] - 2 - first_pusher_idx) return; + markers(ip, first_pusher_idx + k) = k < n_init ? markers(ip, k) : 0.0; } diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index c9271883f..3617791a5 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -18,6 +18,11 @@ logger = logging.getLogger("struphy") +# CUDA kernels that prepare each marker row 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` @@ -217,17 +222,22 @@ def _push(self, dt: float): logger.debug(f"{residual_idx =}") logger.debug(f"{self.particles.n_cols =}") - if self._cuda: + # On the GPU the kernels of _KERNELS_PREPARING_ROWS prepare each row themselves at stage 0, saving a pass over + # the marker buffer; only if no other kernel runs before them (init and evaluation kernels may read the + # prepared columns). + 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 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), - ) + _prepare_push_cuda()(markers, 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) @@ -279,15 +289,20 @@ 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, ) + if prepare_in_kernel and args_markers.prepare_rows: + args_markers.prepare_rows = 0 # markers have moved if self.pushes_eta: diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index 259091753..3e0f5e8ea 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -43,6 +43,20 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, * the host. The marker view carries shape and strides; j is a CUDA-only * column index replacing Pyccel array slices. */ +/** + * Prepare marker row ip for a push, as the kernel prepare_push (prepare_push_cuda.cu) does for every row: copy the + * phase space coordinates (columns 0 to 3 + vdim) to the columns from first_pusher_idx and zero the columns 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 row 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[1] - 2; + for (int col = first; col < end; ++col) + args_markers.markers(ip, col) = col < first + n_init ? args_markers.markers(ip, col - first) : 0.0; +} + __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_markers, const DomainArgs& args_domain, bool newton) { const long long* bc_type = args_markers.bc_type; diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 75e3178d3..add0614d3 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/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), ( From 074151ef16ff7956915225d4b6e3d1f26b518f37 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:28:24 +0200 Subject: [PATCH 2/2] Column-major marker buffer opt-in (STRUPHY_MARKERS_COLUMN_MAJOR=1); prepare_rows in the CUDA emulation Kernels without a CUDA version run on host copies of the device arrays, which keep the F order (cunumpy.to_numpy uses get(order='A')), and the pyccel kernels reject it: the sorting-box kernels failed on the fake CuPy (test_sorting_boxes_on_cupy_fake_cupy), and the same holds on a GPU for every marker kernel without a CUDA version. The buffer is therefore C-ordered unless the variable is set. The CUDA emulation of the tests reads the struct fields from the pyccel argument objects, which do not have the CUDA-only prepare_rows flag (every emulated marker kernel failed). Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 21 ++++++++++++++++----- src/struphy/pic/tests/cuda_emulation.py | 5 ++++- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 8f238267b..75947a4ea 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -75,6 +75,16 @@ """Orbit quantities of the position columns, which post-processing maps to physical coordinates.""" +def markers_column_major() -> bool: + """Whether the marker buffer is allocated column-major on the CuPy backend (``STRUPHY_MARKERS_COLUMN_MAJOR=1``). + + Off by default: kernels without a CUDA version run on host copies of the device arrays, which keep the F order + (``cunumpy.to_numpy``), and the pyccel kernels reject it (e.g. the sorting-box kernels). Switch it on for runs + whose marker kernels all have CUDA versions (e.g. the Vlasov push profiling case). + """ + return os.environ.get("STRUPHY_MARKERS_COLUMN_MAJOR", "0").lower() in ("1", "true", "yes") + + class Particles(metaclass=ABCMeta): r""" Base class for particle species. @@ -2634,11 +2644,12 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - # column-major on the GPU: the kernels (and the column operations on the Python side) access a few columns of - # many markers, so with a row-major buffer every access touches the whole row (n_cols * 8 bytes) and the - # threads of a warp read addresses a row apart. The CUDA kernels address the buffer through its strides - # (Array2D in MarkerArgs). The pyccel kernels expect C order, so NumPy keeps it. - order = "F" if xp.get_backend() == "cupy" else "C" + # column-major on the GPU (opt-in, see markers_column_major): the kernels (and the column operations on the + # Python side) access a few columns of many markers, so with a row-major buffer every access touches the whole + # row (n_cols * 8 bytes) and the threads of a warp read addresses a row apart. The CUDA kernels address the + # buffer through its strides (Array2D in MarkerArgs). The pyccel kernels expect C order, so NumPy + # keeps it. + order = "F" if xp.get_backend() == "cupy" and markers_column_major() else "C" self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float, order=order) # allocate auxiliary arrays diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index ca3c76d01..31071a699 100644 --- a/src/struphy/pic/tests/cuda_emulation.py +++ b/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 rows are prepared beforehand + return getattr(value, name, 0) return getattr(value, name)