diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 6c7f1fccb..e3529b9c4 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 @@ -38,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 @@ -67,14 +67,6 @@ POLE_ETA1_MIN = 1e-10 -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"), (1, "y", "$y$", "physical position y"), @@ -286,9 +278,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: @@ -300,7 +295,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 @@ -957,13 +952,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 @@ -2025,20 +2023,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: @@ -2054,30 +2054,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. @@ -2555,14 +2578,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: @@ -2672,14 +2692,9 @@ 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 (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 = {} - 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 @@ -3158,22 +3173,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: \ @@ -3185,13 +3200,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 @@ -3927,8 +3935,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 (on the host copy, also on the CuPy backend) + dd = self._domain_array_host rank = self.mpi_rank x_l = dd[rank][0] @@ -4517,9 +4525,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..cad32ca93 --- /dev/null +++ b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py @@ -0,0 +1,43 @@ +"""Assign boxes, sort by box and rebuild the box array in one call, see :meth:`~struphy.pic.base.Particles.do_sort`.""" + +# 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( + 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. + """ + 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]: + sorting_kernels.sort_boxed_particles( + markers, + swap_line_1, + swap_line_2, + next_index.shape[0], + next_index, + cumul_next_index, + box_index, + ) + sorting_kernels.find_holes(markers, holes) + # the number of particles per box is unchanged, so they still fit into the box array + 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/__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..00d26e86a --- /dev/null +++ b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py @@ -0,0 +1,29 @@ +"""Assign the box to each particle and fill the box array, see :meth:`~struphy.pic.base.Particles.put_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( + 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. + """ + 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]: + 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/__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..91dc53643 --- /dev/null +++ b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py @@ -0,0 +1,24 @@ +"""Fill the box array from the box column of the markers, see :meth:`~struphy.pic.base.Particles.put_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( + 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 = sorting_kernels.count_particles_in_boxes(markers, holes, next_index, box_index) + if max_in_box < boxes.shape[1]: + 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/__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..674e47dfc --- /dev/null +++ b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py @@ -0,0 +1,34 @@ +"""Sort the boxed particles and rebuild the box array, see :meth:`~struphy.pic.base.Particles.do_sort`.""" + +# 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( + 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. + """ + sorting_kernels.sort_boxed_particles( + markers, + swap_line_1, + swap_line_2, + next_index.shape[0], + next_index, + cumul_next_index, + box_index, + ) + sorting_kernels.find_holes(markers, holes) + # the number of particles per box is unchanged, so they still fit into the box array + sorting_kernels.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 74db8c182..4df93b0e3 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING import cunumpy as xp +import numpy as np from struphy.pic.sorting_kernels import flatten_index, initialize_neighbours @@ -234,11 +235,13 @@ def _set_boxes(self): 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/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_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): diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index c3748a722..7aa363362 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -2,18 +2,36 @@ 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, run_fake_cupy_child 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") +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]) @@ -151,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"]) @@ -194,6 +232,92 @@ 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, sort the same markers as on NumPy and compare. + + 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``). + """ + + 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) + 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" + # 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 + particles.markers[:] = xp.to_cunumpy(markers_in) + particles.update_holes() + 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.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 copy of domain_array) + 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 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`). + + Serial; skipped under ``mpirun``, where every rank would run the same check on the same GPU. + """ + 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. + """ + run_fake_cupy_child( + "from struphy.pic.tests.test_sorting import check_sorting_boxes_on_cupy; check_sorting_boxes_on_cupy()" + ) + + if __name__ == "__main__": test_flattening_roundtrip(8, 8, 8, "c_ordering") # test_sorting( 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)