From 339921189f7496872bbd0ed2c41663c900c85f57 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 8 Oct 2026 23:00:41 +0200 Subject: [PATCH 1/4] Set up sorting boxes on the CuPy backend SortingBoxes could not be created on CuPy: n_cols used xp.sqrt (a 0-d device array that round() rejects) and the pyccel initialize_neighbours was called with a device array. - SortingBoxes._set_boxes: n_cols with math.sqrt; neighbours built on the host with initialize_neighbours and moved to the backend once. - Particles: n_boxes with math.prod; is_domain_boundary and _get_neighbouring_proc (SPH) read a host copy of domain_array. - put_particles_in_boxes / do_sort: the pyccel box-sorting kernels are wrapped in PyccelKernel, an explicit host fallback on CuPy (device -> host -> device per call) until #694 ports them. - Test: Particles6D and ParticlesSPH with boxes_per_dim on CuPy (real or fake CuPy in a serial subprocess on rank 0) match NumPy, including one do_sort. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 32 +++++---- src/struphy/pic/sorting.py | 15 +++-- src/struphy/pic/tests/test_sorting.py | 94 ++++++++++++++++++++++++++- 3 files changed, 122 insertions(+), 19 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1d08e9116..aeb99e5c9 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1,4 +1,5 @@ import logging +import math import os import warnings from abc import ABCMeta, abstractmethod @@ -7,6 +8,7 @@ import cunumpy as xp import h5py import numpy as np +from cunumpy.kernels import PyccelKernel from feectools.ddm.partition import factorint from line_profiler import profile from maybempi import MPI, SerialComm @@ -63,6 +65,14 @@ logger = logging.getLogger("struphy") +# Box sorting has no CUDA port yet (#694). On the CuPy backend these wrappers copy the +# device arrays to the host, run the pyccel kernel there and copy the declared outputs +# back: an explicit host/device round trip on every call. On NumPy they call the +# kernel directly. +_assign_box_to_each_particle = PyccelKernel(assign_box_to_each_particle, outputs=(0,)) +_assign_particles_to_boxes = PyccelKernel(assign_particles_to_boxes, outputs=(2, 3)) +_sort_boxed_particles = PyccelKernel(sort_boxed_particles, outputs=(0, 1, 2, 4, 5)) + def _to_numpy_for_kernel(value): """Convert CuPy arrays to NumPy for compiled kernel calls.""" @@ -297,7 +307,7 @@ def __init__( assert all([nboxes % nproc == 0 for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), ( f"Number of boxes {self.boxes_per_dim =} must be divisible by number of processes {self.nprocs =} in each direction." ) - n_boxes = xp.prod(xp.array(self.boxes_per_dim), dtype=int) * self.num_clones + n_boxes = math.prod(self.boxes_per_dim) * self.num_clones # total number of markers (Np) and particles per cell (ppc) Np = self.loading_params.Np @@ -1947,7 +1957,7 @@ def put_particles_in_boxes(self): neighbouring boxes of neighbouring processes are also communicated (as ghost particles).""" self._remove_ghost_particles() - assign_box_to_each_particle( + _assign_box_to_each_particle( self.markers, self.holes, self._sorting_boxes.nx, @@ -1992,7 +2002,7 @@ def do_sort(self, use_numpy_argsort=False): if use_numpy_argsort: self._sort_boxed_particles_numpy() else: - sort_boxed_particles( + _sort_boxed_particles( self._markers, self._sorting_boxes._swap_line_1, self._sorting_boxes._swap_line_2, @@ -2579,14 +2589,10 @@ def _initialize_sorting_boxes(self): # split boxes across MPI processes nboxes = [nboxes // nproc for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)] - # check whether this process touches the domain boundary + # check whether this process touches the domain boundary (host-side bookkeeping: + # Python bools, also on the CuPy backend) + x_l, x_r, _, y_l, y_r, _, z_l, z_r, _ = xp.to_numpy(self.domain_array[self.mpi_rank]).tolist() is_domain_boundary = {} - x_l = self.domain_array[self.mpi_rank, 0] - x_r = self.domain_array[self.mpi_rank, 1] - y_l = self.domain_array[self.mpi_rank, 3] - y_r = self.domain_array[self.mpi_rank, 4] - z_l = self.domain_array[self.mpi_rank, 6] - z_r = self.domain_array[self.mpi_rank, 7] is_domain_boundary["x_m"] = x_l == 0.0 is_domain_boundary["x_p"] = x_r == 1.0 is_domain_boundary["y_m"] = y_l == 0.0 @@ -3092,7 +3098,7 @@ def _check_and_assign_particles_to_boxes(self): warnings.warn(msg) self.mpi_comm.Abort() - assign_particles_to_boxes( + _assign_particles_to_boxes( self.markers, self.holes, self._sorting_boxes._boxes, @@ -3834,8 +3840,8 @@ def _get_neighbouring_proc(self): periodic2 = self.bc_sph[1] == "periodic" periodic3 = self.bc_sph[2] == "periodic" - # Determine which proc are on which side - dd = self.domain_array + # Determine which proc are on which side (setup on the host, also on the CuPy backend) + dd = xp.to_numpy(self.domain_array) rank = self.mpi_rank x_l = dd[rank][0] diff --git a/src/struphy/pic/sorting.py b/src/struphy/pic/sorting.py index 2d331d6c1..e4521d171 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -1,7 +1,9 @@ import logging +import math from typing import TYPE_CHECKING import cunumpy as xp +import numpy as np from struphy.pic.sorting_kernels import flatten_index, initialize_neighbours @@ -225,19 +227,22 @@ def _set_boxes(self): n_particles = self._markers_shape[0] n_mkr = int(n_particles / n_box_in) + 1 + # host-side bookkeeping: plain Python numbers, also on the CuPy backend n_cols = round( - n_mkr * (1 + 1 / xp.sqrt(n_mkr) + self._box_bufsize), + n_mkr * (1 + 1 / math.sqrt(n_mkr) + self._box_bufsize), ) # cartesian boxes (extra last row stores holes/outside particles) self._boxes = xp.full((self._n_boxes + 1, n_cols), -1, dtype=int) self._next_index = xp.zeros((self._n_boxes + 1), dtype=int) self._cumul_next_index = xp.zeros((self._n_boxes + 2), dtype=int) - self._neighbours = xp.zeros((self._n_boxes, 27), dtype=int) - # A particle on box i only sees particles in boxes that belong to neighbours[i] - initialize_neighbours(self._neighbours, self.nx, self.ny, self.nz) - # logger.info(f"{self._rank = }\n{self._neighbours = }") + # A particle on box i only sees particles in boxes that belong to neighbours[i]. + # Setup data: built once on the host with the pyccel kernel, then moved to the + # active backend (one host-to-device copy on CuPy). + neighbours = np.zeros((self._n_boxes, 27), dtype=int) + initialize_neighbours(neighbours, self.nx, self.ny, self.nz) + self._neighbours = xp.to_cunumpy(neighbours) self._swap_line_1 = xp.zeros(self._markers_shape[1]) self._swap_line_2 = xp.zeros(self._markers_shape[1]) diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index c3748a722..95ce62aff 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -1,4 +1,6 @@ import logging +import subprocess +import sys from time import time import cunumpy as xp @@ -7,8 +9,9 @@ from struphy import BoundaryParameters, LoadingParameters, SortingParameters, WeightsParameters, domains from struphy.feec.psydac_derham import Derham +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env from struphy.io.options import DerhamOptions -from struphy.pic.particles import Particles6D +from struphy.pic.particles import Particles6D, ParticlesSPH from struphy.topology.grids import TensorProductGrid logger = logging.getLogger("struphy") @@ -194,6 +197,95 @@ def test_mpi_sort_markers_on_rank_boundary(bc): assert n_found == 1, f"marker at {eta} found on {n_found} processes" +def check_sorting_boxes_on_cupy(): + """Create ``Particles6D`` with sorting boxes on the CuPy backend and compare the boxes with NumPy. + + Checks the setup (``SortingBoxes``: neighbours and box arrays, also for SPH particles) and one :meth:`do_sort` of the same markers on + both backends. Box sorting has no CUDA port yet (#694); on CuPy it runs the pyccel kernels on the host. Serial + (no MPI communicator). Runs with the real CuPy on a GPU and with cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``). + """ + import numpy as np + + def make_particles(): + return Particles6D( + loading_params=LoadingParameters(Np=3000, seed=1234), + sorting_params=SortingParameters(boxes_per_dim=(4, 3, 2)), + ) + + with xp.use_backend("numpy"): + ref = make_particles() + ref.draw_markers(sort=False) + ref_boxes = ref.sorting_boxes + ref_setup = { + name: np.array(getattr(ref_boxes, name), copy=True) + for name in ("_neighbours", "_boxes", "_next_index", "_cumul_next_index") + } + markers_in = np.array(ref.markers, copy=True) + ref.do_sort() + + with xp.use_backend("cupy"): + particles = make_particles() + boxes = particles.sorting_boxes + assert xp.get_array_backend(boxes.neighbours) == "cupy" + assert xp.get_array_backend(boxes.boxes) == "cupy" + for name, value in ref_setup.items(): + dev = getattr(boxes, name) + assert dev.shape == value.shape and dev.dtype == value.dtype, name + assert np.array_equal(xp.to_numpy(dev), value), name + + # the same markers as on NumPy, then sort on both backends + particles.markers[:] = xp.to_cunumpy(markers_in) + particles.update_holes() + particles.do_sort() + assert np.array_equal(xp.to_numpy(particles.markers), ref.markers) + for name in ("_boxes", "_next_index", "_cumul_next_index"): + assert np.array_equal(xp.to_numpy(getattr(boxes, name)), getattr(ref_boxes, name)), name + + # SPH particles also set up the boundary boxes and the neighbouring processes (on the host) + def make_sph(): + return ParticlesSPH( + loading_params=LoadingParameters(Np=3000, seed=1234), + sorting_params=SortingParameters(boxes_per_dim=(4, 3, 2)), + ) + + with xp.use_backend("numpy"): + ref_sph = make_sph() + with xp.use_backend("cupy"): + sph = make_sph() + assert sph.sorting_boxes.communicate + assert sph.sorting_boxes.is_domain_boundary == ref_sph.sorting_boxes.is_domain_boundary + assert all(type(v) is bool for v in sph.sorting_boxes.is_domain_boundary.values()) + assert sph.sorting_boxes._bnd_boxes_x_m == ref_sph.sorting_boxes._bnd_boxes_x_m + assert (sph._x_m_proc, sph._y_p_proc, sph._x_p_y_p_z_p_proc) == ( + ref_sph._x_m_proc, + ref_sph._y_p_proc, + ref_sph._x_p_y_p_z_p_proc, + ) + assert np.array_equal(xp.to_numpy(sph.sorting_boxes.neighbours), ref_sph.sorting_boxes.neighbours) + + +@pytest.mark.skipif(not xp.cupy_available(), reason="CuPy/GPU not available") +def test_sorting_boxes_on_cupy(): + """Sorting boxes on the CuPy backend match the NumPy ones (see :func:`check_sorting_boxes_on_cupy`).""" + check_sorting_boxes_on_cupy() + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +def test_sorting_boxes_on_cupy_fake_cupy(): + """Without a GPU: :func:`check_sorting_boxes_on_cupy` with cunumpy's fake CuPy, which rejects host/device mixing. + + Runs in a serial subprocess because the fake CuPy must be installed before cunumpy is imported; under ``mpirun`` + only rank 0 starts it. + """ + if MPI.COMM_WORLD.Get_rank() != 0: + pytest.skip("serial subprocess test, runs on rank 0 only") + code = "from struphy.pic.tests.test_sorting import check_sorting_boxes_on_cupy; check_sorting_boxes_on_cupy()" + result = subprocess.run( + [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr[-4000:] + + if __name__ == "__main__": test_flattening_roundtrip(8, 8, 8, "c_ordering") # test_sorting( From 2395b8e1f46271dc9797f5d52eb92284431fd4f3 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Oct 2026 17:30:44 +0200 Subject: [PATCH 2/4] Address review of #736: box kernels in kernel folders, one host round trip per sort - Box sorting goes through cunumpy Kernel dispatch: four entry kernels in pic/kernels// (box_particles, fill_boxes, sort_and_fill_boxes, box_and_sort_particles), declared with Kernel.from_folder and missing_cuda="fallback", so #694 only has to add _cuda.cu. - do_sort (no SPH ghosts) runs assign-box, the per-box count, the sort and the box refill in one kernel: on CuPy the markers cross once in each direction. The count is done on the host (no device bincount, no implicit sync). - The swap lines are scratch buffers, no longer copied back to the device. - boxes is rebuilt after the sort, so it points to the sorted rows. - The box column is taken from SortingBoxes.box_index everywhere. - One host copy of domain_array (_domain_array_host), used for n_cells, the per-rank cell counts, the boundary flags, the neighbouring procs, the neighbour ranks and the box kernels. - Drop the misleading host-bookkeeping comment in SortingBoxes._set_boxes. - Tests: run_fake_cupy_child, mpi_skip for the GPU test, check that boxes matches the markers after do_sort (both sort paths), count the marker transfers, and a multi-rank check on the fake CuPy in test_sorting_device.py. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 160 +++++++++--------- src/struphy/pic/kernels/__init__.py | 10 ++ .../box_and_sort_particles/__init__.py | 16 ++ .../box_and_sort_particles_kernels.py | 48 ++++++ .../pic/kernels/box_particles/__init__.py | 16 ++ .../box_particles/box_particles_kernels.py | 32 ++++ .../pic/kernels/fill_boxes/__init__.py | 16 ++ .../kernels/fill_boxes/fill_boxes_kernels.py | 26 +++ .../kernels/sort_and_fill_boxes/__init__.py | 16 ++ .../sort_and_fill_boxes_kernels.py | 37 ++++ src/struphy/pic/sorting.py | 1 - src/struphy/pic/sorting_kernels.py | 28 +++ src/struphy/pic/tests/test_cuda_parity.py | 1 + src/struphy/pic/tests/test_sorting.py | 98 +++++++---- src/struphy/pic/tests/test_sorting_device.py | 105 +++++++++++- 15 files changed, 494 insertions(+), 116 deletions(-) create mode 100644 src/struphy/pic/kernels/__init__.py create mode 100644 src/struphy/pic/kernels/box_and_sort_particles/__init__.py create mode 100644 src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py create mode 100644 src/struphy/pic/kernels/box_particles/__init__.py create mode 100644 src/struphy/pic/kernels/box_particles/box_particles_kernels.py create mode 100644 src/struphy/pic/kernels/fill_boxes/__init__.py create mode 100644 src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py create mode 100644 src/struphy/pic/kernels/sort_and_fill_boxes/__init__.py create mode 100644 src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 769e1dac2..0a31159c5 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -8,7 +8,6 @@ import cunumpy as xp import h5py import numpy as np -from cunumpy.kernels import PyccelKernel from feectools.ddm.partition import factorint from line_profiler import profile from maybempi import MPI, SerialComm @@ -40,15 +39,14 @@ WeightsParameters, ) from struphy.pic import sampling_kernels, sobol_seq +from struphy.pic.kernels.box_and_sort_particles import box_and_sort_particles +from struphy.pic.kernels.box_particles import box_particles +from struphy.pic.kernels.fill_boxes import fill_boxes +from struphy.pic.kernels.sort_and_fill_boxes import sort_and_fill_boxes from struphy.pic.pushing.kernels.reflect import reflect from struphy.pic.pushing.kernels.sph_mean_velocity_coeffs import sph_mean_velocity_coeffs from struphy.pic.pushing.kernels.sph_viscosity_tensor import sph_viscosity_tensor from struphy.pic.sorting import SortingBoxes -from struphy.pic.sorting_kernels import ( - assign_box_to_each_particle, - assign_particles_to_boxes, - sort_boxed_particles, -) from struphy.pic.sph.kernels.box_based_evaluation_flat import box_based_evaluation_flat from struphy.pic.sph.kernels.box_based_evaluation_meshgrid import box_based_evaluation_meshgrid from struphy.pic.sph.kernels.naive_evaluation_flat import naive_evaluation_flat @@ -65,22 +63,6 @@ logger = logging.getLogger("struphy") -# Box sorting has no CUDA port yet (#694). On the CuPy backend these wrappers copy the -# device arrays to the host, run the pyccel kernel there and copy the declared outputs -# back: an explicit host/device round trip on every call. On NumPy they call the -# kernel directly. -_assign_box_to_each_particle = PyccelKernel(assign_box_to_each_particle, outputs=(0,)) -_assign_particles_to_boxes = PyccelKernel(assign_particles_to_boxes, outputs=(2, 3)) -_sort_boxed_particles = PyccelKernel(sort_boxed_particles, outputs=(0, 1, 2, 4, 5)) - - -def _to_numpy_for_kernel(value): - """Convert CuPy arrays to NumPy for compiled kernel calls.""" - if hasattr(value, "get"): - # This is a CuPy array - return value.get() - return value - ORBIT_POSITIONS = ( (0, "x", "$x$", "physical position x"), @@ -293,9 +275,12 @@ def __init__( else: self._domain_array = domain_decomp[0] self._nprocs = domain_decomp[1] + # domain_array does not change after setup: one host copy for the host-side bookkeeping + # (Python numbers) and the host kernels, also on the CuPy backend + self._domain_array_host = xp.to_numpy(self._domain_array) # total number of cells (equal to mpi_size if no grid) - n_cells = xp.sum(xp.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) * self.num_clones + n_cells = int(np.sum(np.prod(self._domain_array_host[:, 2::3], axis=1, dtype=int))) * self.num_clones # total number of boxes if self.boxes_per_dim is None: @@ -952,13 +937,16 @@ def args_markers(self) -> MarkerArguments | CudaMarkerArguments: return self._args_markers def kernels(self) -> tuple: - """The kernels these particles call when they apply their boundary conditions. + """The kernels these particles call when they apply their boundary conditions and sort into boxes. :func:`~struphy.pic.pushing.kernels.reflect.reflect` if an axis has reflecting boundary - conditions, otherwise none. Used by + conditions, and the box kernels of :mod:`struphy.pic.kernels` if there are sorting boxes. Used by :meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels`. """ - return (reflect,) if self._reflect_axes else () + kernels = (reflect,) if self._reflect_axes else () + if self._sorting_boxes is not None: + kernels += (box_particles, fill_boxes, sort_and_fill_boxes, box_and_sort_particles) + return kernels # ------------------------------------------- # Initial condition and background -> weights @@ -1993,20 +1981,22 @@ def put_particles_in_boxes(self): neighbouring boxes of neighbouring processes are also communicated (as ghost particles).""" self._remove_ghost_particles() - _assign_box_to_each_particle( + max_in_box = box_particles( self.markers, self.holes, self._sorting_boxes.nx, self._sorting_boxes.ny, self._sorting_boxes.nz, - self.domain_array[self.mpi_rank], + self._domain_array_host[self.mpi_rank], + self._sorting_boxes._boxes, + self._sorting_boxes._next_index, + self._sorting_boxes.box_index, ) - - self._check_and_assign_particles_to_boxes() + self._check_box_capacity(max_in_box) if self.sorting_boxes.communicate: self._communicate_boxes() - self._check_and_assign_particles_to_boxes() + self._fill_boxes() self._update_ghost_particles() # if self.verbose: @@ -2022,30 +2012,53 @@ def do_sort(self, use_numpy_argsort=False): """Assign the particles to their sorting boxes and reorder the markers array accordingly, so that markers in the same box occupy contiguous rows. + The box array :attr:`~struphy.pic.sorting.SortingBoxes.boxes` is rebuilt for the sorted rows, + so it stays valid after the sort. + Parameters ---------- use_numpy_argsort : bool If True, sort via :func:`numpy.argsort` on the box column; if False (default), - use the Pyccel kernel :func:`~struphy.pic.sorting_kernels.sort_boxed_particles`. + use the kernel :func:`~struphy.pic.kernels.sort_and_fill_boxes.sort_and_fill_boxes`. """ - nx = self._sorting_boxes.nx - ny = self._sorting_boxes.ny - nz = self._sorting_boxes.nz - nboxes = (nx + 2) * (ny + 2) * (nz + 2) - - self.put_particles_in_boxes() - + boxes = self._sorting_boxes if use_numpy_argsort: + self.put_particles_in_boxes() self._sort_boxed_particles_numpy() + self.update_holes() + self._fill_boxes() + elif boxes.communicate: + # SPH: the ghost particles are exchanged between assigning the boxes and sorting + self.put_particles_in_boxes() + sort_and_fill_boxes( + self._markers, + self._holes, + boxes._boxes, + boxes._next_index, + boxes._cumul_next_index, + boxes._swap_line_1, + boxes._swap_line_2, + boxes.box_index, + ) else: - _sort_boxed_particles( + # assign the boxes, sort and rebuild the box array in one kernel call + # (on the CuPy backend one host round trip of the markers, #694) + self._remove_ghost_particles() + max_in_box = box_and_sort_particles( self._markers, - self._sorting_boxes._swap_line_1, - self._sorting_boxes._swap_line_2, - nboxes + 1, - self._sorting_boxes._next_index, - self._sorting_boxes._cumul_next_index, + self._holes, + boxes.nx, + boxes.ny, + boxes.nz, + self._domain_array_host[self.mpi_rank], + boxes._boxes, + boxes._next_index, + boxes._cumul_next_index, + boxes._swap_line_1, + boxes._swap_line_2, + boxes.box_index, ) + self._check_box_capacity(max_in_box) # The marker rows have just been reordered. The masks are row-based, # so they must be rebuilt before any later use of valid_mks/f_coords. @@ -2523,14 +2536,11 @@ def _n_mks_load_and_Np_per_clone(self): """Return two arrays: 1) an array of sub_comm.size where the i-th entry corresponds to the number of markers drawn on process i, and 2) an array of size num_clones where the i-th entry corresponds to the number of markers on clone i.""" # number of cells on current process - n_cells_loc = xp.prod( - self.domain_array[self.mpi_rank, 2::3], - dtype=int, - ) + n_cells_loc = int(np.prod(self._domain_array_host[self.mpi_rank, 2::3], dtype=int)) # array of number of markers on each process at loading stage if self.clone_config is not None: - _n_cells_clone = xp.sum(xp.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) + _n_cells_clone = int(np.sum(np.prod(self._domain_array_host[:, 2::3], axis=1, dtype=int))) _n_mks_load_tot = self.clone_config.get_Np_clone(self.Np) _ppc = _n_mks_load_tot / _n_cells_clone else: @@ -2631,9 +2641,8 @@ def _initialize_sorting_boxes(self): # split boxes across MPI processes nboxes = [nboxes // nproc for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)] - # check whether this process touches the domain boundary (host-side bookkeeping: - # Python bools, also on the CuPy backend) - x_l, x_r, _, y_l, y_r, _, z_l, z_r, _ = xp.to_numpy(self.domain_array[self.mpi_rank]).tolist() + # check whether this process touches the domain boundary (Python bools, also on the CuPy backend) + x_l, x_r, _, y_l, y_r, _, z_l, z_r, _ = self._domain_array_host[self.mpi_rank].tolist() is_domain_boundary = {} is_domain_boundary["x_m"] = x_l == 0.0 is_domain_boundary["x_p"] = x_r == 1.0 @@ -3113,22 +3122,22 @@ def _sort_boxed_particles_numpy(self): self._markers[:, :] = self._markers[self._argsort_array] - def _check_and_assign_particles_to_boxes(self): - """Check whether the box array has enough columns (detect load imbalance wrt to sorting boxes), - and then assign the particles to boxes.""" - - from cunumpy.xp import array_backend - - if array_backend.backend == "numpy": - bcount = xp.bincount(xp.int64(self.markers_wo_holes[:, -2])) - else: - import cupy as cp - - indices = self.markers_wo_holes[:, -2] - indices = indices.astype(cp.int64) - bcount = cp.bincount(indices) + def _fill_boxes(self): + """Assign the particles to boxes from the box column of the markers, after checking that the box + array has enough columns (see :meth:`_check_box_capacity`).""" + max_in_box = fill_boxes( + self.markers, + self.holes, + self._sorting_boxes._boxes, + self._sorting_boxes._next_index, + self._sorting_boxes.box_index, + ) + self._check_box_capacity(max_in_box) - max_in_box = xp.max(bcount) + def _check_box_capacity(self, max_in_box: int): + """Raise (or abort under MPI) if the largest number of markers in a box, as returned by the box + kernels (counted on the host, a Python int), does not fit into the box array; the kernels then + leave the box array untouched. This detects load imbalance with respect to the sorting boxes.""" # the last column must stay -1 (terminator for the box readers), hence >= if max_in_box >= self._sorting_boxes.boxes.shape[1]: msg = f'Strong load imbalance detected in sorting boxes: \ @@ -3140,13 +3149,6 @@ def _check_and_assign_particles_to_boxes(self): warnings.warn(msg) self.mpi_comm.Abort() - _assign_particles_to_boxes( - self.markers, - self.holes, - self._sorting_boxes._boxes, - self._sorting_boxes._next_index, - ) - def _update_ghost_particles(self): """Refresh :attr:`~struphy.pic.base.Particles.ghost_particles`: a marker is flagged as a ghost particle when its ID column (last column) equals -2, the marker set by @@ -3882,8 +3884,8 @@ def _get_neighbouring_proc(self): periodic2 = self.bc_sph[1] == "periodic" periodic3 = self.bc_sph[2] == "periodic" - # Determine which proc are on which side (setup on the host, also on the CuPy backend) - dd = xp.to_numpy(self.domain_array) + # Determine which proc are on which side (on the host copy, also on the CuPy backend) + dd = self._domain_array_host rank = self.mpi_rank x_l = dd[rank][0] @@ -4472,9 +4474,9 @@ def _compute_neighbor_ranks(self) -> tuple[list[int], list[int]]: non_neighbor_ranks : list[int] Every other rank (excluding this one and ``neighbor_ranks``). """ - # Tiny (mpi_size, 9) array -- brought to the host once so the O(mpi_size) + # Tiny (mpi_size, 9) array -- the host copy, so the O(mpi_size) # comparisons below don't pay a device sync per rank checked. - domain_array_host = _to_numpy_for_kernel(self.domain_array) + domain_array_host = self._domain_array_host own = domain_array_host[self.mpi_rank] tol = 1e-12 diff --git a/src/struphy/pic/kernels/__init__.py b/src/struphy/pic/kernels/__init__.py new file mode 100644 index 000000000..acd747d00 --- /dev/null +++ b/src/struphy/pic/kernels/__init__.py @@ -0,0 +1,10 @@ +"""Box sorting kernels of :class:`~struphy.pic.base.Particles`, one folder per kernel (see ``CUDA_STRATEGY.md``). + +Each folder ```` holds the pyccel kernel ``_kernels.py`` and, once ported (#694), its CUDA version +``_cuda.cu``. The folder's ``__init__.py`` declares the kernel, a :class:`cunumpy.kernels.Kernel` that runs the +version of the active backend:: + + from struphy.pic.kernels.box_particles import box_particles + +The helpers they call are in :mod:`struphy.pic.sorting_kernels`. +""" diff --git a/src/struphy/pic/kernels/box_and_sort_particles/__init__.py b/src/struphy/pic/kernels/box_and_sort_particles/__init__.py new file mode 100644 index 000000000..d29dee097 --- /dev/null +++ b/src/struphy/pic/kernels/box_and_sort_particles/__init__.py @@ -0,0 +1,16 @@ +"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. + +No CUDA version yet (#694): on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is +emitted once), and only the arrays in ``OUTPUTS`` are copied back to the device. +""" + +from cunumpy.kernels import Kernel + +from struphy.utils.cuda_arguments import CUDA_STRUCTS + +# the arguments the kernel writes to, by position: markers, holes, boxes, next_index, cumul_next_index +OUTPUTS = (0, 1, 6, 7, 8) + +box_and_sort_particles = Kernel.from_folder( + __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} +) diff --git a/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py new file mode 100644 index 000000000..394e67628 --- /dev/null +++ b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py @@ -0,0 +1,48 @@ +"""Assign boxes, sort by box and rebuild the box array in one call, see :meth:`~struphy.pic.base.Particles.do_sort`.""" + +from struphy.pic.sorting_kernels import ( + assign_box_to_each_particle, + assign_particles_to_boxes, + count_particles_in_boxes, + find_holes, + sort_boxed_particles, +) + + +def box_and_sort_particles( + markers: "float[:,:]", + holes: "bool[:]", + nx: "int", + ny: "int", + nz: "int", + domain_array: "float[:]", + boxes: "int[:,:]", + next_index: "int[:]", + cumul_next_index: "int[:]", + swap_line_1: "float[:]", + swap_line_2: "float[:]", + box_index: "int" = -2, +) -> "int": + """``box_particles`` followed by ``sort_and_fill_boxes``: write the box of each particle into column ``box_index``, + count the particles per box, reorder the marker rows by box, then update ``holes`` and fill ``boxes`` for the + sorted rows. On the CuPy backend the markers cross between device and host once in each direction. + + Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, nothing is sorted, ``boxes`` is + left untouched and the caller must raise. + """ + assign_box_to_each_particle(markers, holes, nx, ny, nz, domain_array, box_index) + max_in_box = count_particles_in_boxes(markers, holes, next_index, box_index) + if max_in_box < boxes.shape[1]: + sort_boxed_particles( + markers, + swap_line_1, + swap_line_2, + next_index.shape[0], + next_index, + cumul_next_index, + box_index, + ) + find_holes(markers, holes) + # the number of particles per box is unchanged, so they still fit into the box array + assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + return max_in_box diff --git a/src/struphy/pic/kernels/box_particles/__init__.py b/src/struphy/pic/kernels/box_particles/__init__.py new file mode 100644 index 000000000..3ac09129c --- /dev/null +++ b/src/struphy/pic/kernels/box_particles/__init__.py @@ -0,0 +1,16 @@ +"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. + +No CUDA version yet (#694): on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is +emitted once), and only the arrays in ``OUTPUTS`` are copied back to the device. +""" + +from cunumpy.kernels import Kernel + +from struphy.utils.cuda_arguments import CUDA_STRUCTS + +# the arguments the kernel writes to, by position: markers, boxes, next_index +OUTPUTS = (0, 6, 7) + +box_particles = Kernel.from_folder( + __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} +) diff --git a/src/struphy/pic/kernels/box_particles/box_particles_kernels.py b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py new file mode 100644 index 000000000..a28fe3cde --- /dev/null +++ b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py @@ -0,0 +1,32 @@ +"""Assign the box to each particle and fill the box array, see :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`.""" + +from struphy.pic.sorting_kernels import ( + assign_box_to_each_particle, + assign_particles_to_boxes, + count_particles_in_boxes, +) + + +def box_particles( + markers: "float[:,:]", + holes: "bool[:]", + nx: "int", + ny: "int", + nz: "int", + domain_array: "float[:]", + boxes: "int[:,:]", + next_index: "int[:]", + box_index: "int" = -2, +) -> "int": + """Write the box of each particle into column ``box_index`` of the markers, count the particles per box into + ``next_index`` and, if every box fits into a row of ``boxes`` (the last column stays -1 as terminator), write the + particle rows into ``boxes``. + + Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, ``boxes`` is left untouched and + the caller must raise. + """ + assign_box_to_each_particle(markers, holes, nx, ny, nz, domain_array, box_index) + max_in_box = count_particles_in_boxes(markers, holes, next_index, box_index) + if max_in_box < boxes.shape[1]: + assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + return max_in_box diff --git a/src/struphy/pic/kernels/fill_boxes/__init__.py b/src/struphy/pic/kernels/fill_boxes/__init__.py new file mode 100644 index 000000000..c85f8c71a --- /dev/null +++ b/src/struphy/pic/kernels/fill_boxes/__init__.py @@ -0,0 +1,16 @@ +"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. + +No CUDA version yet (#694): on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is +emitted once), and only the arrays in ``OUTPUTS`` are copied back to the device. +""" + +from cunumpy.kernels import Kernel + +from struphy.utils.cuda_arguments import CUDA_STRUCTS + +# the arguments the kernel writes to, by position: boxes, next_index +OUTPUTS = (2, 3) + +fill_boxes = Kernel.from_folder( + __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} +) diff --git a/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py new file mode 100644 index 000000000..1a37465ec --- /dev/null +++ b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py @@ -0,0 +1,26 @@ +"""Fill the box array from the box column of the markers, see :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`.""" + +from struphy.pic.sorting_kernels import ( + assign_particles_to_boxes, + count_particles_in_boxes, +) + + +def fill_boxes( + markers: "float[:,:]", + holes: "bool[:]", + boxes: "int[:,:]", + next_index: "int[:]", + box_index: "int" = -2, +) -> "int": + """Count the particles per box (box number in column ``box_index`` of the markers) into ``next_index`` and, if + every box fits into a row of ``boxes`` (the last column stays -1 as terminator), write the particle rows into + ``boxes``. + + Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, ``boxes`` is left untouched and + the caller must raise. + """ + max_in_box = count_particles_in_boxes(markers, holes, next_index, box_index) + if max_in_box < boxes.shape[1]: + assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + return max_in_box diff --git a/src/struphy/pic/kernels/sort_and_fill_boxes/__init__.py b/src/struphy/pic/kernels/sort_and_fill_boxes/__init__.py new file mode 100644 index 000000000..343f284cb --- /dev/null +++ b/src/struphy/pic/kernels/sort_and_fill_boxes/__init__.py @@ -0,0 +1,16 @@ +"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. + +No CUDA version yet (#694): on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is +emitted once), and only the arrays in ``OUTPUTS`` are copied back to the device. +""" + +from cunumpy.kernels import Kernel + +from struphy.utils.cuda_arguments import CUDA_STRUCTS + +# the arguments the kernel writes to, by position: markers, holes, boxes, next_index, cumul_next_index +OUTPUTS = (0, 1, 2, 3, 4) + +sort_and_fill_boxes = Kernel.from_folder( + __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} +) diff --git a/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py new file mode 100644 index 000000000..2f03213d9 --- /dev/null +++ b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py @@ -0,0 +1,37 @@ +"""Sort the boxed particles and rebuild the box array, see :meth:`~struphy.pic.base.Particles.do_sort`.""" + +from struphy.pic.sorting_kernels import ( + assign_particles_to_boxes, + find_holes, + sort_boxed_particles, +) + + +def sort_and_fill_boxes( + markers: "float[:,:]", + holes: "bool[:]", + boxes: "int[:,:]", + next_index: "int[:]", + cumul_next_index: "int[:]", + swap_line_1: "float[:]", + swap_line_2: "float[:]", + box_index: "int" = -2, +): + """Reorder the marker rows by box, then update ``holes`` and rebuild ``boxes`` for the new rows, so that the box + array stays valid after the sort. + + ``next_index`` must hold the number of particles per box, as left by ``box_particles`` or ``fill_boxes``. The + swap lines are scratch buffers. + """ + sort_boxed_particles( + markers, + swap_line_1, + swap_line_2, + next_index.shape[0], + next_index, + cumul_next_index, + box_index, + ) + find_holes(markers, holes) + # the number of particles per box is unchanged, so they still fit into the box array + assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) diff --git a/src/struphy/pic/sorting.py b/src/struphy/pic/sorting.py index e4521d171..4df93b0e3 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -227,7 +227,6 @@ def _set_boxes(self): n_particles = self._markers_shape[0] n_mkr = int(n_particles / n_box_in) + 1 - # host-side bookkeeping: plain Python numbers, also on the CuPy backend n_cols = round( n_mkr * (1 + 1 / math.sqrt(n_mkr) + self._box_bufsize), ) diff --git a/src/struphy/pic/sorting_kernels.py b/src/struphy/pic/sorting_kernels.py index 46d84007f..4d073e807 100644 --- a/src/struphy/pic/sorting_kernels.py +++ b/src/struphy/pic/sorting_kernels.py @@ -247,3 +247,31 @@ def assign_particles_to_boxes( a = int(markers[p, l + box_index]) boxes[a, next_index[a]] = p next_index[a] += 1 + + +def count_particles_in_boxes( + markers: "float[:,:]", + holes: "bool[:]", + next_index: "int[:]", + box_index: "int" = -2, +) -> "int": + """Count the particles in each box (box number in column ``box_index``) into ``next_index`` and return the + largest count.""" + next_index[:] = 0 + l = markers.shape[1] + for p in range(markers.shape[0]): + if not holes[p]: + a = int(markers[p, l + box_index]) + next_index[a] += 1 + + max_in_box = 0 + for a in range(next_index.shape[0]): + if next_index[a] > max_in_box: + max_in_box = next_index[a] + return max_in_box + + +def find_holes(markers: "float[:,:]", holes: "bool[:]"): + """Mark the empty rows of the marker array (``markers[:, 0] == -1``), as ``Particles.update_holes``.""" + for p in range(markers.shape[0]): + holes[p] = markers[p, 0] == -1.0 diff --git a/src/struphy/pic/tests/test_cuda_parity.py b/src/struphy/pic/tests/test_cuda_parity.py index 160c32bf6..d9c261d3b 100644 --- a/src/struphy/pic/tests/test_cuda_parity.py +++ b/src/struphy/pic/tests/test_cuda_parity.py @@ -31,6 +31,7 @@ "struphy.pic.accumulation.kernels", "struphy.pic.diagnostics.kernels", "struphy.pic.sph.kernels", + "struphy.pic.kernels", "struphy.bsplines.kernels", "struphy.geometry.kernels", "struphy.feec.kernels", diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index 95ce62aff..7aa363362 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -1,15 +1,14 @@ import logging -import subprocess -import sys from time import time import cunumpy as xp +import numpy as np import pytest from maybempi import MPI from struphy import BoundaryParameters, LoadingParameters, SortingParameters, WeightsParameters, domains from struphy.feec.psydac_derham import Derham -from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child from struphy.io.options import DerhamOptions from struphy.pic.particles import Particles6D, ParticlesSPH from struphy.topology.grids import TensorProductGrid @@ -17,6 +16,22 @@ logger = logging.getLogger("struphy") +def assert_boxes_match_markers(particles): + """The box array lists, for every box, exactly the rows of the markers (not holes) whose box column is that box, + in row order and padded with -1, and ``next_index`` holds their number.""" + sorting_boxes = particles.sorting_boxes + boxes = xp.to_numpy(sorting_boxes.boxes) + next_index = xp.to_numpy(sorting_boxes._next_index) + markers = xp.to_numpy(particles.markers) + holes = xp.to_numpy(particles.holes) + box_of_row = markers[:, sorting_boxes.box_index].astype(int) + for box in range(boxes.shape[0]): + rows = np.flatnonzero(~holes & (box_of_row == box)) + assert next_index[box] == rows.size, box + assert np.array_equal(boxes[box, : rows.size], rows), box + assert np.all(boxes[box, rows.size :] == -1), box + + @pytest.mark.parametrize("nx", [8, 70]) @pytest.mark.parametrize("ny", [16, 80]) @pytest.mark.parametrize("nz", [32, 90]) @@ -154,6 +169,26 @@ def test_sorting(num_elements, degree, bcs, mapping, Np): box_markers = particles.markers[:, -2] assert all(box_markers[i] <= box_markers[i + 1] for i in range(len(box_markers) - 1)) + assert_boxes_match_markers(particles) + + +@pytest.mark.parametrize("use_numpy_argsort", [False, True]) +def test_boxes_valid_after_do_sort(use_numpy_argsort): + """After :meth:`do_sort` the box array points to the sorted rows (it is rebuilt after the rows are reordered).""" + particles = Particles6D( + comm_world=MPI.COMM_WORLD, + loading_params=LoadingParameters(Np=3000, seed=1234), + sorting_params=SortingParameters(boxes_per_dim=(12, 12, 12)), + ) + particles.draw_markers(sort=False) + particles.mpi_sort_markers() + rows_before = np.array(particles.markers, copy=True) + + particles.do_sort(use_numpy_argsort=use_numpy_argsort) + + # the rows were really reordered, so a box array from before the sort would be wrong + assert not np.array_equal(particles.markers, rows_before) + assert_boxes_match_markers(particles) @pytest.mark.parametrize("bc", ["periodic", "remove"]) @@ -198,13 +233,14 @@ def test_mpi_sort_markers_on_rank_boundary(bc): def check_sorting_boxes_on_cupy(): - """Create ``Particles6D`` with sorting boxes on the CuPy backend and compare the boxes with NumPy. + """Create ``Particles6D`` with sorting boxes on the CuPy backend, sort the same markers as on NumPy and compare. - Checks the setup (``SortingBoxes``: neighbours and box arrays, also for SPH particles) and one :meth:`do_sort` of the same markers on - both backends. Box sorting has no CUDA port yet (#694); on CuPy it runs the pyccel kernels on the host. Serial - (no MPI communicator). Runs with the real CuPy on a GPU and with cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``). + Checks the setup (``SortingBoxes.neighbours``, also for SPH particles) and one :meth:`do_sort`: the sorted markers, + the box arrays (which must point to the sorted rows) and that the markers cross between device and host once in + each direction. Box sorting has no CUDA port yet (#694); on CuPy the box kernels run on the host. Serial (no MPI + communicator); the multi-rank version is ``check_sorting_boxes_device`` in ``test_sorting_device.py``. Runs with + the real CuPy on a GPU and with cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``). """ - import numpy as np def make_particles(): return Particles6D( @@ -215,33 +251,34 @@ def make_particles(): with xp.use_backend("numpy"): ref = make_particles() ref.draw_markers(sort=False) - ref_boxes = ref.sorting_boxes - ref_setup = { - name: np.array(getattr(ref_boxes, name), copy=True) - for name in ("_neighbours", "_boxes", "_next_index", "_cumul_next_index") - } markers_in = np.array(ref.markers, copy=True) ref.do_sort() + assert_boxes_match_markers(ref) with xp.use_backend("cupy"): particles = make_particles() boxes = particles.sorting_boxes assert xp.get_array_backend(boxes.neighbours) == "cupy" assert xp.get_array_backend(boxes.boxes) == "cupy" - for name, value in ref_setup.items(): - dev = getattr(boxes, name) - assert dev.shape == value.shape and dev.dtype == value.dtype, name - assert np.array_equal(xp.to_numpy(dev), value), name + # the only box array with content before the first sort + assert np.array_equal(xp.to_numpy(boxes.neighbours), ref.sorting_boxes.neighbours) - # the same markers as on NumPy, then sort on both backends + # the same markers as on NumPy, then sort particles.markers[:] = xp.to_cunumpy(markers_in) particles.update_holes() - particles.do_sort() + with xp.profiling.count_transfers() as transfers: + particles.do_sort() assert np.array_equal(xp.to_numpy(particles.markers), ref.markers) for name in ("_boxes", "_next_index", "_cumul_next_index"): - assert np.array_equal(xp.to_numpy(getattr(boxes, name)), getattr(ref_boxes, name)), name + assert np.array_equal(xp.to_numpy(getattr(boxes, name)), getattr(ref.sorting_boxes, name)), name + assert_boxes_match_markers(particles) + + # one host round trip of the markers: assign the boxes, sort and refill the box array in one kernel call + markers_bytes = ref.markers.nbytes + assert sum(e.kind == "to_host" and e.nbytes == markers_bytes for e in transfers.events) == 1, transfers.report() + assert sum(e.kind == "to_device" and e.nbytes == markers_bytes for e in transfers.events) == 1, transfers.report() - # SPH particles also set up the boundary boxes and the neighbouring processes (on the host) + # SPH particles also set up the boundary boxes and the neighbouring processes (on the host copy of domain_array) def make_sph(): return ParticlesSPH( loading_params=LoadingParameters(Np=3000, seed=1234), @@ -256,17 +293,16 @@ def make_sph(): assert sph.sorting_boxes.is_domain_boundary == ref_sph.sorting_boxes.is_domain_boundary assert all(type(v) is bool for v in sph.sorting_boxes.is_domain_boundary.values()) assert sph.sorting_boxes._bnd_boxes_x_m == ref_sph.sorting_boxes._bnd_boxes_x_m - assert (sph._x_m_proc, sph._y_p_proc, sph._x_p_y_p_z_p_proc) == ( - ref_sph._x_m_proc, - ref_sph._y_p_proc, - ref_sph._x_p_y_p_z_p_proc, - ) assert np.array_equal(xp.to_numpy(sph.sorting_boxes.neighbours), ref_sph.sorting_boxes.neighbours) +@pytest.mark.mpi_skip @pytest.mark.skipif(not xp.cupy_available(), reason="CuPy/GPU not available") def test_sorting_boxes_on_cupy(): - """Sorting boxes on the CuPy backend match the NumPy ones (see :func:`check_sorting_boxes_on_cupy`).""" + """Sorting boxes on the CuPy backend match the NumPy ones (see :func:`check_sorting_boxes_on_cupy`). + + Serial; skipped under ``mpirun``, where every rank would run the same check on the same GPU. + """ check_sorting_boxes_on_cupy() @@ -277,13 +313,9 @@ def test_sorting_boxes_on_cupy_fake_cupy(): Runs in a serial subprocess because the fake CuPy must be installed before cunumpy is imported; under ``mpirun`` only rank 0 starts it. """ - if MPI.COMM_WORLD.Get_rank() != 0: - pytest.skip("serial subprocess test, runs on rank 0 only") - code = "from struphy.pic.tests.test_sorting import check_sorting_boxes_on_cupy; check_sorting_boxes_on_cupy()" - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True + run_fake_cupy_child( + "from struphy.pic.tests.test_sorting import check_sorting_boxes_on_cupy; check_sorting_boxes_on_cupy()" ) - assert result.returncode == 0, result.stderr[-4000:] if __name__ == "__main__": diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index dc8fa9462..ac5c7aa88 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -1,4 +1,4 @@ -"""Marker exchange between ranks on the CuPy backend (:meth:`Particles.mpi_sort_markers`). +"""Marker exchange between ranks (:meth:`Particles.mpi_sort_markers`) and sorting boxes on the CuPy backend. The marker rows are exchanged as device buffers with CUDA-aware MPI, and staged through host memory otherwise. Without a GPU, the tests run on cunumpy's fake CuPy, which must be @@ -20,8 +20,9 @@ from cunumpy.kernel_testing import requires_cupy from maybempi import MPI -from struphy import BoundaryParameters, LoadingParameters -from struphy.pic.particles import Particles6D +from struphy import BoundaryParameters, LoadingParameters, SortingParameters +from struphy.pic.particles import Particles6D, ParticlesSPH +from struphy.pic.tests.test_sorting import assert_boxes_match_markers def _fake_cupy_active() -> bool: @@ -189,5 +190,103 @@ def test_mpi_sort_markers_gpu(bc): check_mpi_sort_markers_device(None, bc=bc, count_reads=False) +# divisible by the number of processes in each direction for 1 to 4 ranks +BOXES_PER_DIM = (12, 12, 12) + + +def check_sorting_boxes_device(): + """Sorting boxes on NumPy and CuPy across all ranks: the same setup, and one :meth:`do_sort` of the same markers. + + On a rank > 0 the box kernels and the setup read the row of that rank from the host copy of ``domain_array``, and + SPH particles find their neighbouring processes with a real domain decomposition (for more than one rank). + """ + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + + def make_particles(): + return Particles6D( + comm_world=comm, + loading_params=LoadingParameters(Np=4000, seed=1234), + sorting_params=SortingParameters(boxes_per_dim=BOXES_PER_DIM), + ) + + ref = make_particles() + ref.draw_markers(sort=False) + ref.mpi_sort_markers() + markers_in = ref.markers.copy() + ref.do_sort() + assert_boxes_match_markers(ref) + assert ref.n_mks_loc > 0 + # all markers of this rank lie in its domain, so none is in the out-of-domain box (last row of the box array): + # the boxes were found with the domain_array row of this rank + out_box = ref.sorting_boxes.boxes.shape[0] - 1 + assert np.all(ref.markers[ref.valid_mks, ref.sorting_boxes.box_index] < out_box) + + with xp.use_backend("cupy"): + dev = make_particles() + assert xp.get_array_backend(dev.domain_array) == "cupy" + assert isinstance(dev._domain_array_host, np.ndarray) + assert np.array_equal(dev._domain_array_host, ref.domain_array) + + dev.markers[:] = xp.to_cupy(markers_in) + dev.update_holes() + with xp.profiling.count_transfers() as transfers: + dev.do_sort() + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + for name in ("_boxes", "_next_index", "_cumul_next_index"): + assert np.array_equal(xp.to_numpy(getattr(dev.sorting_boxes, name)), getattr(ref.sorting_boxes, name)) + assert_boxes_match_markers(dev) + + # the markers cross between device and host once in each direction (#694: no CUDA box kernels yet) + markers_bytes = ref.markers.nbytes + assert sum(e.kind == "to_host" and e.nbytes == markers_bytes for e in transfers.events) == 1, transfers.report() + assert sum(e.kind == "to_device" and e.nbytes == markers_bytes for e in transfers.events) == 1, transfers.report() + + def make_sph(): + return ParticlesSPH( + comm_world=comm, + loading_params=LoadingParameters(Np=4000, seed=1234), + sorting_params=SortingParameters(boxes_per_dim=BOXES_PER_DIM), + ) + + ref_sph = make_sph() + with xp.use_backend("cupy"): + dev_sph = make_sph() + + assert dev_sph.sorting_boxes.communicate + is_domain_boundary = dev_sph.sorting_boxes.is_domain_boundary + assert is_domain_boundary == ref_sph.sorting_boxes.is_domain_boundary + assert all(type(v) is bool for v in is_domain_boundary.values()) + assert np.array_equal(xp.to_numpy(dev_sph.sorting_boxes.neighbours), ref_sph.sorting_boxes.neighbours) + for side in ("x_m", "x_p", "y_m", "y_p", "z_m", "z_p"): + name = f"_bnd_boxes_{side}" + assert getattr(dev_sph.sorting_boxes, name) == getattr(ref_sph.sorting_boxes, name), name + + # the neighbouring process in each of the 26 directions + procs = sorted(name for name in vars(ref_sph) if name.endswith("_proc")) + assert len(procs) == 26, procs + dev_procs = {name: getattr(dev_sph, name) for name in procs} + ref_procs = {name: getattr(ref_sph, name) for name in procs} + assert all(p is None or isinstance(p, (int, np.integer)) for p in dev_procs.values()), dev_procs + assert dev_procs == ref_procs + if size > 1: + # a real decomposition: some neighbours are other ranks + assert any(p is not None and p != rank for p in ref_procs.values()), ref_procs + + +@requires_fake_cupy +@pytest.mark.mpi +def test_sorting_boxes_fake_cupy(): + """On the fake CuPy, on every rank: sorting boxes and :meth:`do_sort` give the NumPy result.""" + check_sorting_boxes_device() + + +@requires_cupy +@pytest.mark.mpi +def test_sorting_boxes_gpu(): + """On a GPU (``mpirun -n 2``, one GPU per rank): sorting boxes and :meth:`do_sort` give the NumPy result.""" + check_sorting_boxes_device() + + if __name__ == "__main__": check_mpi_sort_markers_device(None) From 44c996b1e5f46f286827e96989abfb018905a52f Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Oct 2026 18:18:28 +0200 Subject: [PATCH 3/4] Box kernels: import sorting_kernels as a module, not its names The CI container keeps the compiled sorting_kernels from devel, and `struphy compile` imports struphy before it rebuilds the kernels. Importing count_particles_in_boxes/find_holes by name then failed with an ImportError in every job. With `import struphy.pic.sorting_kernels as sorting_kernels` (as the SPH and pusher kernels do), the names are only looked up when a kernel runs, which is after the rebuild. Co-Authored-By: Claude Opus 5.5 --- .../box_and_sort_particles_kernels.py | 19 +++++++------------ .../box_particles/box_particles_kernels.py | 13 +++++-------- .../kernels/fill_boxes/fill_boxes_kernels.py | 10 ++++------ .../sort_and_fill_boxes_kernels.py | 13 +++++-------- 4 files changed, 21 insertions(+), 34 deletions(-) diff --git a/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py index 394e67628..cad32ca93 100644 --- a/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py +++ b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py @@ -1,12 +1,7 @@ """Assign boxes, sort by box and rebuild the box array in one call, see :meth:`~struphy.pic.base.Particles.do_sort`.""" -from struphy.pic.sorting_kernels import ( - assign_box_to_each_particle, - assign_particles_to_boxes, - count_particles_in_boxes, - find_holes, - sort_boxed_particles, -) +# module import (not names): the module may still be an older compiled version until `struphy compile` runs +import struphy.pic.sorting_kernels as sorting_kernels def box_and_sort_particles( @@ -30,10 +25,10 @@ def box_and_sort_particles( Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, nothing is sorted, ``boxes`` is left untouched and the caller must raise. """ - assign_box_to_each_particle(markers, holes, nx, ny, nz, domain_array, box_index) - max_in_box = count_particles_in_boxes(markers, holes, next_index, box_index) + sorting_kernels.assign_box_to_each_particle(markers, holes, nx, ny, nz, domain_array, box_index) + max_in_box = sorting_kernels.count_particles_in_boxes(markers, holes, next_index, box_index) if max_in_box < boxes.shape[1]: - sort_boxed_particles( + sorting_kernels.sort_boxed_particles( markers, swap_line_1, swap_line_2, @@ -42,7 +37,7 @@ def box_and_sort_particles( cumul_next_index, box_index, ) - find_holes(markers, holes) + sorting_kernels.find_holes(markers, holes) # the number of particles per box is unchanged, so they still fit into the box array - assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + sorting_kernels.assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) return max_in_box diff --git a/src/struphy/pic/kernels/box_particles/box_particles_kernels.py b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py index a28fe3cde..00d26e86a 100644 --- a/src/struphy/pic/kernels/box_particles/box_particles_kernels.py +++ b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py @@ -1,10 +1,7 @@ """Assign the box to each particle and fill the box array, see :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`.""" -from struphy.pic.sorting_kernels import ( - assign_box_to_each_particle, - assign_particles_to_boxes, - count_particles_in_boxes, -) +# module import (not names): the module may still be an older compiled version until `struphy compile` runs +import struphy.pic.sorting_kernels as sorting_kernels def box_particles( @@ -25,8 +22,8 @@ def box_particles( Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, ``boxes`` is left untouched and the caller must raise. """ - assign_box_to_each_particle(markers, holes, nx, ny, nz, domain_array, box_index) - max_in_box = count_particles_in_boxes(markers, holes, next_index, box_index) + sorting_kernels.assign_box_to_each_particle(markers, holes, nx, ny, nz, domain_array, box_index) + max_in_box = sorting_kernels.count_particles_in_boxes(markers, holes, next_index, box_index) if max_in_box < boxes.shape[1]: - assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + sorting_kernels.assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) return max_in_box diff --git a/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py index 1a37465ec..91dc53643 100644 --- a/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py +++ b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py @@ -1,9 +1,7 @@ """Fill the box array from the box column of the markers, see :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`.""" -from struphy.pic.sorting_kernels import ( - assign_particles_to_boxes, - count_particles_in_boxes, -) +# module import (not names): the module may still be an older compiled version until `struphy compile` runs +import struphy.pic.sorting_kernels as sorting_kernels def fill_boxes( @@ -20,7 +18,7 @@ def fill_boxes( Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, ``boxes`` is left untouched and the caller must raise. """ - max_in_box = count_particles_in_boxes(markers, holes, next_index, box_index) + max_in_box = sorting_kernels.count_particles_in_boxes(markers, holes, next_index, box_index) if max_in_box < boxes.shape[1]: - assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + sorting_kernels.assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) return max_in_box diff --git a/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py index 2f03213d9..674e47dfc 100644 --- a/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py +++ b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py @@ -1,10 +1,7 @@ """Sort the boxed particles and rebuild the box array, see :meth:`~struphy.pic.base.Particles.do_sort`.""" -from struphy.pic.sorting_kernels import ( - assign_particles_to_boxes, - find_holes, - sort_boxed_particles, -) +# module import (not names): the module may still be an older compiled version until `struphy compile` runs +import struphy.pic.sorting_kernels as sorting_kernels def sort_and_fill_boxes( @@ -23,7 +20,7 @@ def sort_and_fill_boxes( ``next_index`` must hold the number of particles per box, as left by ``box_particles`` or ``fill_boxes``. The swap lines are scratch buffers. """ - sort_boxed_particles( + sorting_kernels.sort_boxed_particles( markers, swap_line_1, swap_line_2, @@ -32,6 +29,6 @@ def sort_and_fill_boxes( cumul_next_index, box_index, ) - find_holes(markers, holes) + sorting_kernels.find_holes(markers, holes) # the number of particles per box is unchanged, so they still fit into the box array - assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) + sorting_kernels.assign_particles_to_boxes(markers, holes, boxes, next_index, box_index) From 992917ea85973ea88ae2e9ea23bdbbb04a5c77dd Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Oct 2026 18:58:28 +0200 Subject: [PATCH 4/4] Tests: stubs expose _domain_array_host, count the pic.kernels catalog - test_neighbor_ranks: the Particles stubs set _domain_array_host, which _compute_neighbor_ranks and _get_neighbouring_proc now read. - test_kernel_backends: the catalog count includes the 4 box kernels of struphy.pic.kernels. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/test_kernel_backends.py | 1 + src/struphy/pic/tests/test_neighbor_ranks.py | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 3fab53103..75e3178d3 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -198,6 +198,7 @@ def test_catalog_signatures(): "pic.accumulation.kernels": 16, "pic.diagnostics.kernels": 10, "pic.sph.kernels": 4, + "pic.kernels": 4, "bsplines.kernels": 4, "geometry.kernels": 4, "feec.kernels": 17, diff --git a/src/struphy/pic/tests/test_neighbor_ranks.py b/src/struphy/pic/tests/test_neighbor_ranks.py index 5edf8b88d..3b72acc49 100644 --- a/src/struphy/pic/tests/test_neighbor_ranks.py +++ b/src/struphy/pic/tests/test_neighbor_ranks.py @@ -18,6 +18,7 @@ def make_stub(comm, periodic_axes=()): ) domain_array, nprocs = Particles._get_domain_decomp(stub, mpi_dims_mask=None) stub.domain_array = domain_array + stub._domain_array_host = domain_array return stub, nprocs @@ -121,7 +122,13 @@ def test_neighbouring_proc_boxes(mpi_size, mpi_dims_mask, bc_sph): domain_array, nprocs = Particles._get_domain_decomp(stub, mpi_dims_mask=mpi_dims_mask) for rank in range(mpi_size): - s = SimpleNamespace(domain_array=domain_array, mpi_rank=rank, mpi_size=mpi_size, bc_sph=list(bc_sph)) + s = SimpleNamespace( + domain_array=domain_array, + _domain_array_host=domain_array, + mpi_rank=rank, + mpi_size=mpi_size, + bc_sph=list(bc_sph), + ) Particles._get_neighbouring_proc(s) ijk = rank_ijk(rank, nprocs) for offs in itertools.product((-1, 0, 1), repeat=3):