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 4443a6791..7a125123b 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/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() 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 caefe86d6..6c711014a 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 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.; 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 8eb12c785..c0eadfbe9 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 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 diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 423960096..ff1290f3d 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 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: diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index 91f2feaf5..59fac4686 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -28,6 +28,20 @@ __device__ inline void reflect_velocity(long long ip, Array2D 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. * diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index ca3c76d01..0fe1a9743 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 markers are prepared beforehand + return getattr(value, name, 0) return getattr(value, name) diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index b7e56cb47..bf71b8e70 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), (