diff --git a/.gitignore b/.gitignore index 136d46db9..19c0df4aa 100644 --- a/.gitignore +++ b/.gitignore @@ -101,6 +101,8 @@ src/struphy/io/out/ src/struphy/state.yml src/struphy/io/inp/params_* *.bin +bench_gpu/out_* +bench_gpu/ # models list bin/ diff --git a/feectools b/feectools index ce78b9bb2..ef8695040 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit ce78b9bb2cbeeed0dfe34327900df5f3fc1b7608 +Subproject commit ef86950402a667a9ef77b49e8868f6e581e20dac diff --git a/pyproject.toml b/pyproject.toml index 955937c55..3d69386e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,9 @@ phys = [ "gvec>=1.1.0, <=1.5.0", "desc-opt<=0.17.1", ] +cuda = [ + "cupy-cuda12x==13.6.0", +] dev = [ "struphy[mpi]", "notebook", @@ -137,6 +140,7 @@ struphy = "struphy.console.main:struphy" ] struphy = [ "compile_struphy.mk", + "**/*.cu", ] [tool.autopep8] @@ -182,4 +186,5 @@ markers = [ "hybrid", "single", "mpi_pic", + "needs_host_kernels", ] diff --git a/setup/modules.pitagora.sh b/setup/modules.pitagora.sh index 085db7430..cf8632ad7 100644 --- a/setup/modules.pitagora.sh +++ b/setup/modules.pitagora.sh @@ -5,12 +5,22 @@ python/3.11.7" # openmpi/4.1.6--gcc--12.3.0 MODULES_GCC="gcc/12.3.0 \ +openmpi/4.1.6--gcc--12.3.0 \ python/3.11.7 \ hdf5/1.14.3--gcc--12.3.0 \ cmake/3.27.9 \ netcdf-fortran/4.6.1--gcc--12.3.0 \ netlib-scalapack/2.2.0--openmpi--4.1.6--gcc--12.3.0-ucx1.20" +# On the Booster (GPU) partition, ARRAY_BACKEND=cupy runs need libnvrtc.so.12 for +# cupy's RawKernel/JIT compilation -- otherwise every cupy import fails as soon as +# it touches the GPU (e.g. `xp.tri()` at struphy import time). SLURM_JOB_PARTITION +# is only set inside a submitted job, so this is a no-op on the DCGP (CPU) partition +# or outside SLURM. +if [[ "${SLURM_JOB_PARTITION:-}" == *boost* ]]; then + MODULES_INTEL="$MODULES_INTEL cuda/12.6" + MODULES_GCC="$MODULES_GCC cuda/12.6" +fi # For GVEC # Should be fixed so it works with both gcc and intel diff --git a/src/struphy/__init__.py b/src/struphy/__init__.py index bbf73f1e2..526e62d51 100644 --- a/src/struphy/__init__.py +++ b/src/struphy/__init__.py @@ -4,6 +4,33 @@ import logging.config import os +# HDF5's file locking relies on flock(), which is unreliable/unsupported on +# parallel filesystems such as Lustre or GPFS (common on HPC clusters) and +# causes spurious `BlockingIOError: Unable to synchronously open file` errors. +# Disable it unless the user has explicitly configured it. +os.environ.setdefault("HDF5_USE_FILE_LOCKING", "FALSE") + +# mpi4py defaults to requesting MPI_THREAD_MULTIPLE (thread level 3) from +# MPI_Init_thread. On at least one cluster this repo runs on (Pitagora's Booster +# partition, OpenMPI 4.1.6 + UCX 1.20), the UCX worker does not support that level, +# which OpenMPI reports at every multi-rank run ("UCP worker does not support +# MPI_THREAD_MULTIPLE" / "failed to init ucx" / hcoll init failure) and works +# around by making hcoll (its GPU-aware collective component) fail to initialize, +# silently falling back to a different, working collective implementation. Struphy +# only ever calls MPI from the main Python thread (CuPy's internal CUDA driver +# threads don't touch MPI), so requesting the weaker MPI_THREAD_FUNNELED guarantee +# instead is sufficient and avoids the warnings -- but it ALSO lets hcoll +# successfully initialize where it previously failed to, and hcoll's own +# Alltoallv implementation on this cluster then segfaults +# (hmca_bcol_ucx_p2p_alltoallv_pairwise_chunk_progress) the first time it's +# actually used, something the failed init was silently protecting us from. +# hcoll must therefore stay disabled explicitly alongside the thread-level +# change, not just left to fail its own init. Both must be set before mpi4py.MPI +# is imported anywhere (thread level can't change after MPI_Init), and only if +# the user hasn't already configured them themselves. +os.environ.setdefault("MPI4PY_RC_THREAD_LEVEL", "funneled") +os.environ.setdefault("OMPI_MCA_coll_hcoll_enable", "0") + from feectools.ddm.mpi import mpi as MPI from struphy.utils.mpi_launch import launched_under_mpi diff --git a/src/struphy/conftest.py b/src/struphy/conftest.py index 4bddafb21..a707b7bd7 100644 --- a/src/struphy/conftest.py +++ b/src/struphy/conftest.py @@ -1,9 +1,17 @@ import logging +import os import pytest from struphy import set_logging_level +# Tests marked "needs_host_kernels" call a Pyccel kernel that is not yet +# ported for the CuPy backend (e.g. FEEC mass-matrix/basis-projection +# assembly, particle-to-grid accumulation). Under ARRAY_BACKEND=cupy they are +# skipped rather than run to failure, so a GPU CI run reports the state of +# the actually-ported code paths instead of drowning in known gaps. +_NEEDS_HOST_KERNELS_SKIP_REASON = "needs_host_kernels: not yet ported to the CuPy backend (ARRAY_BACKEND=cupy)" + def set_logging_level_pytest(config): level_name = str(config.getoption("--logging-level")).upper() @@ -25,6 +33,15 @@ def pytest_configure(config): set_logging_level_pytest(config) +def pytest_collection_modifyitems(config, items): + if os.environ.get("ARRAY_BACKEND") != "cupy": + return + skip_host_only = pytest.mark.skip(reason=_NEEDS_HOST_KERNELS_SKIP_REASON) + for item in items: + if "needs_host_kernels" in item.keywords: + item.add_marker(skip_host_only) + + def pytest_addoption(parser): parser.addoption("--with-desc", action="store_true") parser.addoption("--vrbose", action="store_true") diff --git a/src/struphy/console/format.py b/src/struphy/console/format.py index f2670a9bb..c70a712f0 100644 --- a/src/struphy/console/format.py +++ b/src/struphy/console/format.py @@ -1211,7 +1211,15 @@ def confirm_formatting(python_files, linters, yes): ) print("\n") if not yes: - ans = input("Format files (Y/n)?\n") + try: + ans = input("Format files (Y/n)?\n") + except EOFError: + # stdin isn't an interactive terminal (piped, scripted, non-tty + # shell, ...) -- input() can't prompt, so there's no way to get + # a real answer. Fail safe (don't format) instead of crashing + # with a raw traceback, and point at the flag that avoids this. + print("\nNo interactive terminal to confirm on. Exiting... (use --yes/-y to skip this prompt)") + sys.exit(1) if ans.lower() not in ("y", "yes", ""): print("Exiting...") sys.exit(1) diff --git a/src/struphy/cuda.py b/src/struphy/cuda.py new file mode 100644 index 000000000..d18571f92 --- /dev/null +++ b/src/struphy/cuda.py @@ -0,0 +1,63 @@ +"""Helpers for loading CUDA C sources used by CuPy RawKernel wrappers.""" + +from functools import lru_cache +from pathlib import Path +from typing import Sequence + + +@lru_cache(maxsize=None) +def load_cuda_source(module_file: str, source_name: str) -> str: + """Load a CUDA C source fragment stored alongside its Python wrapper.""" + path = Path(module_file).with_name("cuda") / source_name + return path.read_text(encoding="utf-8") + + +class CudaKernel: + """A lazily-compiled, cached ``cupy.RawKernel``. + + Every hand-written CUDA replacement in ``struphy.pic.*_cuda`` / + ``struphy.feec.*_cuda`` used to repeat the same boilerplate at each call + site: a module-level ``_foo_kernel = None`` sentinel, a + ``_get_foo_kernel()`` function that imports ``cupy`` and compiles the + ``RawKernel`` the first time it's needed (so importing these modules + under ``ARRAY_BACKEND=numpy`` never touches CuPy), and caches it back into + the global. This class replaces that boilerplate with one declaration: + + _foo_kernel = CudaKernel(_FOO_SRC, "foo_cuda") + + made once at module level, right next to the ``*_gpu`` function it backs. + Compilation is still deferred to first call (``import cupy`` only happens + inside :meth:`__call__`), and the compiled kernel is cached on the + instance -- identical behavior to the old pattern, but now every real GPU + kernel a module launches shows up as a ``CudaKernel(...)`` at module + scope, so `grep -n "CudaKernel("` (or just reading the top of the file) + tells you exactly which functions do device work and which don't. + """ + + __slots__ = ("_source", "_name", "_kernel") + + def __init__(self, source: str, name: str) -> None: + self._source = source + self._name = name + self._kernel = None + + def __call__(self, grid, block, args) -> None: + self._compiled()(grid, block, args) + + def _compiled(self): + if self._kernel is None: + import cupy as cp + + kernel = cp.RawKernel(self._source, self._name) + kernel.compile() + self._kernel = kernel + return self._kernel + + +def launch_1d(kernel: CudaKernel, n: int, args: Sequence, threads: int = 256) -> None: + """Launch ``kernel`` over a 1-D grid with one thread per element of a + length-``n`` array (markers, indices, quadrature points, ...) -- the + launch geometry shared by every kernel in ``struphy.pic.*_cuda`` / + ``struphy.feec.*_cuda``.""" + blocks = (n + threads - 1) // threads + kernel((blocks,), (threads,), tuple(args)) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 5bf9e957c..a566550c8 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -1,6 +1,7 @@ import logging import cunumpy as xp +import numpy as np from feectools.api.essential_bc import apply_essential_bc_stencil from feectools.ddm.cart import CartDecomposition, DomainDecomposition from feectools.ddm.mpi import MockComm @@ -260,7 +261,7 @@ def fun(e): M_local = StencilMatrix(V_local, V_local) - row_indices, col_indices = xp.nonzero(M_arr) + row_indices, col_indices = np.nonzero(M_arr) for row_i, col_i in zip(row_indices, col_indices): # only consider row indices on process @@ -273,7 +274,7 @@ def fun(e): ] = M_arr[row_i, col_i] # check if stencil matrix was built correctly - assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + assert np.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) matrixcells += [M_local.copy()] # ======================================================================================================= @@ -625,7 +626,7 @@ def __init__(self, mass_operator, apply_bc=True): M_local = StencilMatrix(V_local, V_local) - row_indices, col_indices = xp.nonzero(M_arr) + row_indices, col_indices = np.nonzero(M_arr) for row_i, col_i in zip(row_indices, col_indices): # only consider row indices on process @@ -638,7 +639,7 @@ def __init__(self, mass_operator, apply_bc=True): ] = M_arr[row_i, col_i] # check if stencil matrix was built correctly - assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + assert np.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) matrixcells += [M_local.copy()] # ======================================================================================================= @@ -911,10 +912,12 @@ class FFTSolver(BandedSolver): """ def __init__(self, circmat): - assert isinstance(circmat, xp.ndarray) + # circmat comes from StencilMatrix.toarray(), always host numpy; the + # underlying solve() also calls scipy's solve_circulant, host-only. + assert isinstance(circmat, np.ndarray) assert is_circulant(circmat) - self._space = xp.ndarray + self._space = np.ndarray self._column = circmat[:, 0] # -------------------------------------- @@ -979,13 +982,15 @@ def is_circulant(mat): Whether the matrix is circulant (=True) or not (=False). """ - assert isinstance(mat, xp.ndarray) + # mat comes from StencilMatrix.toarray(), which always densifies to a + # host numpy.ndarray regardless of the active backend. + assert isinstance(mat, np.ndarray) assert len(mat.shape) == 2 assert mat.shape[0] == mat.shape[1] if mat.shape[0] > 1: for i in range(mat.shape[0] - 1): - circulant = xp.allclose(mat[i, :], xp.roll(mat[i + 1, :], -1)) + circulant = np.allclose(mat[i, :], np.roll(mat[i + 1, :], -1)) if not circulant: return circulant else: diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 101bcfd12..8da190c4b 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1925,11 +1925,12 @@ def _get_domain_array(self): else: nproc = 1 - # send buffer - dom_arr_loc = xp.zeros(9, dtype=float) + # send buffer -- host-resident: mpi4py's Allgather needs a real + # buffer-protocol array, not a CuPy array, regardless of backend. + dom_arr_loc = np.zeros(9, dtype=float) # main array (receive buffers) - dom_arr = xp.zeros(nproc * 9, dtype=float) + dom_arr = np.zeros(nproc * 9, dtype=float) # Get global starts and ends of domain decomposition gl_s = self.domain_decomposition.starts diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index 98c18b288..3d3a73afe 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -108,7 +108,6 @@ def __init__( self.bd2 = np.empty(int(pn[1]), dtype=float) self.bd3 = np.empty(int(pn[2]), dtype=float) - class DomainArguments: """Holds the mandatory arguments pertaining to :class:`~struphy.geometry.base.Domain` passed to particle pusher kernels. diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index cb8e5aa7b..9e742b282 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,8 +2,10 @@ import os import warnings from abc import ABCMeta, abstractmethod +from contextlib import contextmanager import h5py +import numpy as np import scipy.special as sp try: @@ -69,12 +71,134 @@ class Intracomm: logger = logging.getLogger("struphy") +def _array_types(): + """Array classes a marker-column setter may legitimately be handed. + + Marker data lives on the active backend (device under CuPy), so a + setter must accept that backend's array type as well as NumPy's -- + plain NumPy is still valid input and gets converted on assignment. + """ + if xp.cupy_backend: + import cupy as cp + + return (np.ndarray, cp.ndarray) + return (np.ndarray,) + + +_ARRAY_TYPES = _array_types() + + 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 + """Convert CuPy arrays to NumPy for compiled kernel calls. + + xp.is_gpu, not xp.to_numpy: some callers pass plain Python scalars (e.g. + self.Np, self.vdim) that must reach MarkerArguments unchanged, and + xp.to_numpy would wrap those into 0-d NumPy arrays via np.asarray -- not + the plain int MarkerArguments' typed fields expect. xp.is_gpu leaves + anything that isn't actually a CuPy array untouched, matching the original + hasattr(value, "get") passthrough behaviour exactly. + """ + return value.get() if xp.is_gpu(value) else value + + +def _dev(*arrays): + """Convert host (marker) coordinate arrays to the active array backend, + for feeding into equilibrium/domain/perturbation functions that follow + the global backend rather than the (always host-resident) markers.""" + out = tuple(xp.to_cunumpy(a) for a in arrays) + return out[0] if len(out) == 1 else out + + +def _pinned_zeros(shape, dtype=float): + """Allocate a zeroed NumPy array, backed by page-locked ("pinned") host + memory when the active backend is CuPy. + + The returned object is a plain ``numpy.ndarray`` in every respect + (Pyccel kernels, aliasing with ``args_markers.markers``, etc. all work + exactly as with a regular allocation) — pinning only changes how fast the + *host* memory can later be DMA'd to/from the device. Pageable memory + (the default) transfers the full markers array at roughly PCIe-over-copy + speed (~90 ms for 137 MiB, measured); pinned memory reaches near the + PCIe link's true bandwidth (~5 ms for the same array), which is what + makes it worthwhile to bounce the per-step vectorized bookkeeping + (:meth:`Particles._find_outside_particles`, the column-block resets in + :class:`~struphy.pic.pushing.pusher.Pusher`) through the device. + + Under the NumPy backend nothing is ever transferred, so pinning would + only tie up a scarcer resource for no benefit; a plain allocation is + used instead. + """ + if not xp.cupy_backend: + return np.zeros(shape, dtype=dtype) + import cupy as cp + + size = int(np.prod(shape)) + nbytes = size * np.dtype(dtype).itemsize + mem = cp.cuda.alloc_pinned_memory(nbytes) + arr = np.frombuffer(mem, dtype=dtype, count=size).reshape(shape) + arr[:] = 0 + return arr + + +class _HostMarkerMirror: + """Explicit host mirror of a device-resident marker array. + + Under ``ARRAY_BACKEND=cupy`` the marker array (and its companion boolean + masks) live on the device — that is where every ported CUDA particle + kernel reads and writes them, and where the vectorized bookkeeping + (boundary conditions, hole tracking, sorting) runs, so no transfer is + needed in the hot path at all. + + A handful of *cold-path* consumers are still compiled, host-only Pyccel + kernels (SPH evaluation, some diagnostics/accumulation kernels that have + no CUDA port yet). Those need a real ``numpy.ndarray`` that they can read + -- and in some cases write -- in place. This class owns that host buffer + and makes the crossing explicit rather than implicit: callers wrap the + Pyccel call in :meth:`Particles.host_markers`, which copies device->host + on entry and (when ``write=True``) host->device on exit. + + The buffer is allocated once, in pinned memory, and reused; the mirror + object handed to ``MarkerArguments`` therefore stays identity-stable for + the lifetime of the :class:`Particles` instance, exactly as the old + always-host array did. + """ + + __slots__ = ("_pairs", "_depth") + + def __init__(self): + # list of [device_array, host_array]; host buffers are identity-stable + self._pairs = [] + self._depth = 0 + + def add(self, device_array): + """Register a device array and return its identity-stable host buffer.""" + host = _pinned_zeros(device_array.shape, dtype=device_array.dtype) + self._pairs.append([device_array, host]) + return host + + def rebind(self, old_device_array, new_device_array): + """Point the mirror at a new device array (after a resize/realloc). + + Returns the host buffer for the new array, reallocating it only if + the shape actually changed. + """ + for pair in self._pairs: + if pair[0] is old_device_array: + pair[0] = new_device_array + if new_device_array.shape != pair[1].shape: + pair[1] = _pinned_zeros(new_device_array.shape, dtype=new_device_array.dtype) + return pair[1] + return self.add(new_device_array) + + def pull(self): + """device -> host, for every registered array.""" + for device, host in self._pairs: + device.get(out=host) + + def push(self): + """host -> device (in place, so device array identities are kept).""" + for device, host in self._pairs: + device.set(host) ORBIT_POSITIONS = ( @@ -286,7 +410,7 @@ def __init__( if domain_decomp is None: self._domain_array, self._nprocs = self._get_domain_decomp(self.sorting_params.dims_mask) else: - self._domain_array = domain_decomp[0] + self._domain_array = xp.to_numpy(domain_decomp[0]) self._nprocs = domain_decomp[1] # total number of cells (equal to mpi_size if no grid) @@ -402,7 +526,10 @@ def __init__( # if self.loading_params["moments"] is None and not isinstance(self, ParticlesSPH) and isinstance(self.bckgr_params, dict): self._generate_sampling_moments() - # create buffers for mpi_sort_markers + # create buffers for mpi_sort_markers -- marker-row-indexed, so they + # live on the same backend as the markers themselves (device under + # CuPy). The actual mpi4py Alltoall/Isend/Irecv calls further below + # take host buffers and convert explicitly at that point. self._sorting_etas = xp.zeros((self.markers.shape[0], 3), dtype=float) self._is_on_proc_domain = xp.zeros((self.markers.shape[0], 3), dtype=bool) self._can_stay = xp.zeros(self.markers.shape[0], dtype=bool) @@ -699,6 +826,19 @@ def domain_array(self): """ return self._domain_array + @property + def domain_array_dev(self): + """:attr:`domain_array` on the active backend. + + The domain decomposition is small, fixed metadata built once on the + host, but it is repeatedly compared against marker positions, which + live on the device under CuPy. Cached here so those comparisons do + not re-upload it on every call. + """ + if getattr(self, "_domain_array_dev", None) is None: + self._domain_array_dev = xp.asarray(self._domain_array) + return self._domain_array_dev + @property def mpi_dims_mask(self): """3-list | tuple; True if the dimension is to be used in the domain decomposition (=default for each dimension). @@ -775,15 +915,29 @@ def ghost_particles(self): self._ghost_particles = self.markers[:, -1] == -2.0 return self._ghost_particles + @property + def has_ghost_particles(self): + """Whether any row of the markers array may currently be a ghost particle. + + Ghost particles only ever enter through :meth:`_communicate_boxes` (the SPH + ghost-box layer): that is the sole path that writes ``-2`` into the ID column, + both for the rows this rank marks and sends and for the rows it receives. + Every other model -- all the PIC ones -- never creates one, so this stays False + for the whole run and lets the hot marker-sorting path skip the full-array + passes that would only ever find nothing. See :meth:`_update_ghost_particles` + and :meth:`_remove_ghost_particles`. + """ + return getattr(self, "_has_ghost_particles", False) + @property def markers_wo_holes(self): """Array holding the marker information, excluding holes. The i-th row holds the i-th marker info.""" - return self.markers[~self.holes] + return self.markers[np.nonzero(~self.holes)[0]] @property def markers_wo_holes_and_ghost(self): """Array holding the marker information, excluding holes and ghosts (only valid markers). The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks] + return self.markers[self._valid_row_idx] @property def lost_markers(self): @@ -799,13 +953,31 @@ def n_lost_markers(self): def valid_mks(self): """Array of booleans stating if an entry in the markers array is a true local particle (not a hole or ghost).""" if not hasattr(self, "_valid_mks"): - self._valid_mks = ~xp.logical_or(self.holes, self.ghost_particles) + self._valid_mks = ~np.logical_or(self.holes, self.ghost_particles) return self._valid_mks + @property + def _valid_row_idx(self): + """Integer row indices where :attr:`valid_mks` is True. + + Used (instead of ``markers[self.valid_mks, colslice]``) by the + read-only marker-column properties below: slicing columns first -- + cheap, since a plain-slice column index is a view -- and only then + gathering rows via integer fancy indexing is measurably faster + (~35-60% at Np=1e6, measured) than boolean-masking the full-width + row range directly, which is what made ``update_scalar_quantities()`` + cost about as much as the (CUDA-accelerated) ``model.integrate()`` + step itself at large Np, on both backends equally -- this is a plain + host/NumPy indexing cost, unrelated to ARRAY_BACKEND, since + ``markers`` is always host-resident (see + ``ISSUE_cupy_particles_never_pushed.md``). + """ + return np.nonzero(self.valid_mks)[0] + @property def n_mks_loc(self): """Number of valid markers on process (without holes and ghosts).""" - return xp.count_nonzero(self.valid_mks) + return np.count_nonzero(self.valid_mks) @property def n_mks_on_each_proc(self): @@ -830,89 +1002,89 @@ def n_mks_global(self): @property def positions(self): """Array holding the marker positions in logical space. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["pos"]] + return self.markers[:, self.index["pos"]][self._valid_row_idx] @positions.setter def positions(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc, 3) self._markers[self.valid_mks, self.index["pos"]] = new @property def velocities(self): """Array holding the marker velocities in logical space. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["vel"]] + return self.markers[:, self.index["vel"]][self._valid_row_idx] @velocities.setter def velocities(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc, self.vdim), f"{self.n_mks_loc =} and {self.vdim =} but {new.shape =}" self._markers[self.valid_mks, self.index["vel"]] = new @property def phasespace_coords(self): """Array holding the marker positions and velocities in logical space. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["coords"]] + return self.markers[:, self.index["coords"]][self._valid_row_idx] @phasespace_coords.setter def phasespace_coords(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc, 3 + self.vdim) self._markers[self.valid_mks, self.index["coords"]] = new @property def weights(self): """Array holding the current marker weights. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["weights"]] + return self.markers[:, self.index["weights"]][self._valid_row_idx] @weights.setter def weights(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["weights"]] = new @property def sampling_density_values(self): """Array holding the current marker 0form sampling density s0. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["s0"]] + return self.markers[:, self.index["s0"]][self._valid_row_idx] @sampling_density_values.setter def sampling_density_values(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["s0"]] = new @property def weights0(self): """Array holding the initial marker weights. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["w0"]] + return self.markers[:, self.index["w0"]][self._valid_row_idx] @weights0.setter def weights0(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["w0"]] = new @property def marker_ids(self): """Array holding the marker id's on the current process.""" - return self.markers[self.valid_mks, self.index["ids"]] + return self.markers[:, self.index["ids"]][self._valid_row_idx] @marker_ids.setter def marker_ids(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["ids"]] = new @property def f_coords(self): """Coordinates of the distribution function.""" - return self.markers[self.valid_mks, self.f_coords_index] + return self.markers[:, self.f_coords_index][self._valid_row_idx] @f_coords.setter def f_coords(self, new): - assert isinstance(new, xp.ndarray) - self.markers[self.valid_mks, self.f_coords_index] = new + assert isinstance(new, _ARRAY_TYPES) + self.markers[:, self.f_coords_index][self._valid_row_idx] = new @property def f_jacobian_coords(self): @@ -924,7 +1096,7 @@ def f_jacobian_coords(self): @f_jacobian_coords.setter def f_jacobian_coords(self, new): - assert isinstance(new, xp.ndarray) + assert isinstance(new, _ARRAY_TYPES) if isinstance(self.f_jacobian_coords_index, list): self.markers[ xp.ix_( @@ -937,9 +1109,65 @@ def f_jacobian_coords(self, new): @property def args_markers(self) -> MarkerArguments: - """Collection of mandatory arguments for pusher kernels.""" + """Collection of mandatory arguments for pusher kernels. + + Note + ---- + Under ``ARRAY_BACKEND=cupy`` the arrays inside are the *host mirror* + of the device-resident markers, not the markers themselves. They are + only valid inside a :meth:`host_markers` block — outside one they + hold whatever the last such block left behind. Every compiled + (Pyccel) kernel call that takes ``args_markers`` must therefore be + wrapped; CUDA kernels take :attr:`markers` directly instead and need + no wrapping. + """ return self._args_markers + @contextmanager + def host_markers(self, *, write: bool = True): + """Make :attr:`args_markers` valid for a compiled, host-only (Pyccel) + kernel call, and write any changes back to the device afterwards. + + Under the NumPy backend this is a no-op: ``args_markers`` already + aliases the marker array, so there is nothing to copy either way. + + Under CuPy the markers live on the device (that is the whole point + -- the ported CUDA kernels and the vectorized bookkeeping never + transfer them), so a kernel that can only run on the host needs an + explicit crossing. Entering copies device->host; leaving copies + host->device, unless ``write=False`` marks the call read-only (the + common case for accumulation/evaluation kernels, which only ever + read markers and write into grid arrays), in which case the + write-back is skipped. + + Nested blocks are reference-counted, so the transfer happens once + per outermost block. A nested ``write=True`` upgrades the outer + block, never the reverse. + + Parameters + ---------- + write : bool + Whether the wrapped kernel may modify the marker array. Passing + ``False`` when it actually does modify markers silently discards + those changes, so only use it for kernels verified read-only. + """ + if self._markers_mirror is None: + yield self._args_markers + return + + mirror = self._markers_mirror + if mirror._depth == 0: + mirror.pull() + self._markers_mirror_write = False + mirror._depth += 1 + self._markers_mirror_write = self._markers_mirror_write or write + try: + yield self._args_markers + finally: + mirror._depth -= 1 + if mirror._depth == 0 and self._markers_mirror_write: + mirror.push() + # ------------------------------------------- # Initial condition and background -> weights # ------------------------------------------- @@ -1153,12 +1381,12 @@ def draw_markers( self._markers[n_mks_load_loc:] = -1.0 # number of holes and markers on process - self.update_holes() + self.update_holes(update_valid_mks=False) self._update_ghost_particles() # cumulative sum of number of markers on each process at loading stage. - n_mks_load_cum_sum = xp.cumsum(self.n_mks_load) - Np_per_clone_cum_sum = xp.cumsum(self.Np_per_clone) + n_mks_load_cum_sum = np.cumsum(self.n_mks_load) + Np_per_clone_cum_sum = np.cumsum(self.Np_per_clone) _first_marker_id = (Np_per_clone_cum_sum - self.Np_per_clone)[self.clone_id] + ( n_mks_load_cum_sum - self.n_mks_load )[self._mpi_rank] @@ -1184,7 +1412,7 @@ def draw_markers( self._load_tesselation() if isinstance(self, ParticlesSPH): self._set_initial_condition() - self.velocities = xp.array(self.u_init(self.positions)).T + self.velocities = self.u_init(self.positions).T # set markers ID in last column self.marker_ids = _first_marker_id + xp.arange(n_mks_load_loc, dtype=float) else: @@ -1197,7 +1425,7 @@ def draw_markers( # set seed _seed = self.loading_params.seed if _seed is not None: - xp.random.seed(_seed) + np.random.seed(_seed) # counting integers num_loaded_particles_loc = 0 # number of particles alreday loaded (local) @@ -1208,22 +1436,22 @@ def draw_markers( while num_loaded_particles_glob < int(self.Np): # Generate a chunk of random particles num_to_add_glob = min(chunk_size, int(self.Np) - num_loaded_particles_glob) - temp = xp.random.rand(num_to_add_glob, 3 + self.vdim) + temp = np.random.rand(num_to_add_glob, 3 + self.vdim) # check which particles are on the current process domain - is_on_proc_domain = xp.logical_and( - temp[:, :3] > self.domain_array[self.mpi_rank, 0::3], - temp[:, :3] < self.domain_array[self.mpi_rank, 1::3], + is_on_proc_domain = np.logical_and( + temp[:, :3] > self._domain_array[self.mpi_rank, 0::3], + temp[:, :3] < self._domain_array[self.mpi_rank, 1::3], ) - valid_idx = xp.nonzero(xp.all(is_on_proc_domain, axis=1))[0] + valid_idx = np.nonzero(np.all(is_on_proc_domain, axis=1))[0] valid_particles = temp[valid_idx] - valid_particles = xp.array_split(valid_particles, self.num_clones)[self.clone_id] + valid_particles = np.array_split(valid_particles, self.num_clones)[self.clone_id] num_valid = valid_particles.shape[0] # Add the valid particles to the phasespace_coords array self._markers[ num_loaded_particles_loc : num_loaded_particles_loc + num_valid, : 3 + self.vdim, - ] = valid_particles + ] = xp.asarray(valid_particles) num_loaded_particles_glob += num_to_add_glob num_loaded_particles_loc += num_valid @@ -1233,7 +1461,7 @@ def draw_markers( # set new n_mks_load self.gather_scalar_in_subcomm_array(num_loaded_particles_loc, out=self.n_mks_load) n_mks_load_loc = self.n_mks_load[self.mpi_rank] - n_mks_load_cum_sum = xp.cumsum(self.n_mks_load) + n_mks_load_cum_sum = np.cumsum(self.n_mks_load) # set new holes in markers array to -1 self._markers[num_loaded_particles_loc:] = -1.0 @@ -1259,10 +1487,12 @@ def draw_markers( 1000 + (n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank] // 64, ) - sampling_kernels.set_particles_symmetric_3d_3v( - temp_markers, - self.markers, - ) + # compiled host-only sampler; fills marker rows in place + with self.host_markers(write=True) as args_markers: + sampling_kernels.set_particles_symmetric_3d_3v( + _to_numpy_for_kernel(temp_markers), + args_markers.markers, + ) # 4. Wrong specification else: @@ -1273,7 +1503,7 @@ def draw_markers( # initial velocities - SPH case: v(0) = u(x(0)) for given velocity u(x) if isinstance(self, ParticlesSPH): self._set_initial_condition() - self.velocities = xp.array(self.u_init(self.positions)).T + self.velocities = self.u_init(self.positions).T else: # inverse transform sampling in velocity space # Avoid exact 0 or 1 from low-discrepancy sequences: erfinv(±1) @@ -1291,22 +1521,16 @@ def draw_markers( # Particles6D: (1d Maxwellian, 1d Maxwellian, 1d Maxwellian) if isinstance(self, Particles6D): + # sp is plain scipy.special: host-only, so it needs a host + # copy of the velocities and its result converted back to the + # active backend before mixing with v_th/u_mean. self.velocities = ( - sp.erfinv( - 2 * self.velocities - 1, - ) - * xp.sqrt(2) - * v_th - + u_mean + _dev(sp.erfinv(2 * _to_numpy_for_kernel(self.velocities) - 1)) * xp.sqrt(2) * v_th + u_mean ) # Particles5D: (1d Maxwellian, muB0-Maxwellian as volume-form) elif isinstance(self, Particles5D): self._markers[:n_mks_load_loc, 3] = ( - sp.erfinv( - 2 * self.velocities[:, 0] - 1, - ) - * xp.sqrt(2) - * v_th[0] + _dev(sp.erfinv(2 * _to_numpy_for_kernel(self.velocities[:, 0]) - 1)) * xp.sqrt(2) * v_th[0] + u_mean[0] ) @@ -1322,20 +1546,12 @@ def draw_markers( # Particles5Dvperp: (1d Maxwellian, polar Maxwellian as volume-form) elif isinstance(self, Particles5Dvperp): self._markers[:n_mks_load_loc, 3] = ( - sp.erfinv( - 2 * self.velocities[:, 0] - 1, - ) - * xp.sqrt(2) - * v_th[0] + _dev(sp.erfinv(2 * _to_numpy_for_kernel(self.velocities[:, 0]) - 1)) * xp.sqrt(2) * v_th[0] + u_mean[0] ) self._markers[:n_mks_load_loc, 4] = ( - xp.sqrt( - -xp.log(1.0 - self.velocities[:, 1]), - ) - * xp.sqrt(2) - * v_th[1] + xp.sqrt(-xp.log(1.0 - self.velocities[:, 1])) * xp.sqrt(2) * v_th[1] ) # v_perp is a polar velocity coordinate and must be >= 0. @@ -1381,8 +1597,8 @@ def draw_markers( # check if all particle positions are inside the unit cube [0, 1]^3 n_mks_load_loc = self.n_mks_load[self._mpi_rank] - assert xp.all(~self.holes[:n_mks_load_loc]) - assert xp.all(self.holes[n_mks_load_loc:]) + assert np.all(~self.holes[:n_mks_load_loc]) + assert np.all(self.holes[n_mks_load_loc:]) if self._initialized_sorting and sort: logger.info("\nSorting the markers after initial draw") @@ -1431,7 +1647,7 @@ def initialize_weights( else: assert self.domain is not None, "A domain is needed to initialize weights." - if xp.size(self.markers_wo_holes_and_ghost) == 0: + if np.size(self.markers_wo_holes_and_ghost) == 0: return # set initial condition @@ -1447,7 +1663,9 @@ def initialize_weights( # if isinstance(self.f_init, CanonicalMaxwellian): # self.save_constants_of_motion() - # evaluate initial distribution function + # evaluate initial distribution function. Marker-derived inputs + # and the field/background functions now agree on the backend + # (both follow xp), so no conversion is needed either way. if isinstance(self, ParticlesSPH): f_init = self.f_init(self.positions) else: @@ -1455,15 +1673,16 @@ def initialize_weights( # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: - f_init /= self.domain.jacobian_det(self.positions) + f_init = f_init / self.domain.jacobian_det(self.positions) if self.is_volume_form[1]: - f_init /= self.f_init.velocity_jacobian_det( - *self.f_jacobian_coords.T, - ) + f_init = f_init / self.f_init.velocity_jacobian_det(*self.f_jacobian_coords.T) # compute s0 and save at vdim + 4 - self.sampling_density_values = self.s0(*self.phasespace_coords.T, flat_eval=True) + self.sampling_density_values = self.s0( + *self.phasespace_coords.T, + flat_eval=True, + ) # compute w0 and save at vdim + 5 self.weights0 = f_init / self.sampling_density_values / self.Np @@ -1492,7 +1711,7 @@ def update_weights(self): """ from struphy.pic.particles import ParticlesSPH - if xp.size(self.markers_wo_holes_and_ghost) == 0: + if np.size(self.markers_wo_holes_and_ghost) == 0: return if isinstance(self, ParticlesSPH): @@ -1505,10 +1724,10 @@ def update_weights(self): # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: - f0 /= self.domain.jacobian_det(self.positions) + f0 = f0 / self.domain.jacobian_det(self.positions) if self.is_volume_form[1]: - f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords.T) + f0 = f0 / self.f0.velocity_jacobian_det(*self.f_jacobian_coords.T) self.weights = self.weights0 - f0 / self.sampling_density_values / self.Np @@ -1546,7 +1765,13 @@ def binning( The reconstructed delta-f distribution function. """ - assert xp.count_nonzero(components) == len(bin_edges) + # np.histogramdd below needs host bin edges, like the markers-derived + # sample/weights arrays it's called with (see ISSUE_cupy_particles_never_pushed.md); + # accept xp-typed input defensively so callers on the active backend + # don't need to know that. + bin_edges = tuple(_to_numpy_for_kernel(be) for be in bin_edges) + + assert np.count_nonzero(components) == len(bin_edges) # volume of a bin bin_vol = 1.0 @@ -1578,28 +1803,36 @@ def binning( _weights = self.weights * self.Np * multiplier if divide_by_jac: - _weights /= self.domain.jacobian_det(self.positions, remove_outside=False) + jac_det = self.domain.jacobian_det(self.positions, remove_outside=False) + _weights = _weights / jac_det # _weights /= self.velocity_jacobian_det(*self.phasespace_coords.T) - _weights0 /= self.domain.jacobian_det(self.positions, remove_outside=False) + _weights0 = _weights0 / jac_det # _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords.T) - f_slice = xp.histogramdd( - self.markers_wo_holes_and_ghost[:, slicing], - bins=bin_edges, - weights=_weights0, + # numpy.histogramdd has no CuPy equivalent, so the binning itself is + # done on the host; the inputs are brought across explicitly here. + _binned_coords = _to_numpy_for_kernel(self.markers_wo_holes_and_ghost[:, slicing]) + + f_slice = np.histogramdd( + _binned_coords, + bins=[_to_numpy_for_kernel(be) for be in bin_edges], + weights=_to_numpy_for_kernel(_weights0), )[0] - df_slice = xp.histogramdd( - self.markers_wo_holes_and_ghost[:, slicing], - bins=bin_edges, - weights=_weights, + df_slice = np.histogramdd( + _binned_coords, + bins=[_to_numpy_for_kernel(be) for be in bin_edges], + weights=_to_numpy_for_kernel(_weights), )[0] f_slice /= self.Np * bin_vol df_slice /= self.Np * bin_vol - return f_slice, df_slice + # np.histogramdd forces f_slice/df_slice to be host arrays regardless + # of the active backend; convert back so callers get results on the + # same backend as everything else this class returns. + return xp.asarray(f_slice), xp.asarray(df_slice) def show_distribution_function(self, components: list[bool], bin_edges: list[np.ndarray], do_plot=False): """ @@ -1737,15 +1970,30 @@ def mpi_sort_markers( remove_ghost : bool Remove ghost particles before send. """ + # No self._Barrier() here (there used to be one): _remove_ghost_particles and + # apply_kinetic_bc below are both purely local -- neither touches self.mpi_comm + # -- so there is no cross-rank dependency for a barrier to protect at this + # point. The barrier previously here looks like it was compensating for + # _sendrecv_markers not waiting on its Isend requests (see there); now that it + # does, buffer reuse across mpi_sort_markers calls is safe without it too (see + # the trailing comment at the end of this method). if remove_ghost: self._remove_ghost_particles() - self._Barrier() - # before sorting, apply kinetic bc if apply_bc: self.apply_kinetic_bc() + # With a single MPI rank there is no destination calculation or + # communication to perform. Keep the bookkeeping updates that the + # exchange path performs below (boundary conditions may have created + # holes), but avoid materialising the O(n_markers) sorting masks and + # destination arrays just to discover that every marker stays local. + if self.mpi_size == 1: + self.update_holes(update_valid_mks=False) + self._update_ghost_particles() + return + if isinstance(alpha, int) or isinstance(alpha, float): alpha = (alpha, alpha, alpha) @@ -1764,27 +2012,32 @@ def mpi_sort_markers( # send and receive markers self._sendrecv_markers(recv_info, hole_inds_after_send) - # new holes and new number of holes and markers on process - self.update_holes() - - # refresh ghost mask: received markers may land in rows that previously held - # ghost particles. update_holes alone recomputes valid_mks from a stale - # _ghost_particles mask, which would wrongly exclude these incoming real markers. + # new holes and new number of holes and markers on process. update_valid_mks + # deferred to _update_ghost_particles below (see its parameter docstring) -- + # this alone recomputing valid_mks would use a stale _ghost_particles mask + # anyway (received markers may land in rows that previously held ghosts). + self.update_holes(update_valid_mks=False) self._update_ghost_particles() # check if all markers are on the right process after sorting if do_test: all_on_right_proc = xp.all( xp.logical_and( - self.positions > self.domain_array[self.mpi_rank, 0::3], - self.positions < self.domain_array[self.mpi_rank, 1::3], + self.positions > self.domain_array_dev[self.mpi_rank, 0::3], + self.positions < self.domain_array_dev[self.mpi_rank, 1::3], ), ) assert all_on_right_proc # assert self.phasespace_coords.size > 0, f'No particles on process {self.mpi_rank}, please rebalance, aborting ...' - self._Barrier() + # No trailing self._Barrier() here either (there used to be one): every Isend + # and Irecv issued by _sendrecv_markers is now waited on before it returns, so + # this rank's send/recv buffers are already safe to overwrite on the next call + # without an extra rendezvous. Cross-round message ordering between any given + # pair of ranks is additionally guaranteed by MPI itself (non-overtaking + # messages for the same source/dest/tag), so a later round's Isend can't be + # mistaken for an earlier one even without this barrier. @profile @ProfileManager.profile("apply_kinetic_bc") @@ -1816,44 +2069,83 @@ def apply_kinetic_bc(self, newton=False): self._particle_refilling() self._markers[self._is_outside, :-1] = -1.0 - self._n_lost_markers += len(xp.nonzero(self._is_outside)[0]) + self._n_lost_markers += len(np.nonzero(self._is_outside)[0]) if self._periodic_axes: self._eta_bc_buf[:] = self.markers[:, :3] + # Computed once and reused for every axis below, instead of recomputing + # inside _find_outside_particles on each iteration: nothing in this loop + # body creates a hole or a ghost particle (it only wraps positions and + # updates the shift buffer), so the holes/ghost_particles masks -- and + # therefore this derived mask -- cannot change between axes. Measured + # ~9.5% faster on CuPy for the 3-axis periodic loop at Np_local=12.5M. + periodic_not_hole_or_ghost = ~(self.holes | self.ghost_particles) + + # Pass 1 -- locate the outside markers on every periodic axis, keeping only + # the (small) index arrays. The masks themselves cannot be held across axes: + # _find_outside_particles writes the shared _is_outside_left/_is_outside_right + # buffers, which the next axis overwrites. + # + # Reverted from a branchless/xp.where + unconditional-elementwise version: + # measured on a real 50M-marker GPU run, that version was a net loss -- on + # each call only a small fraction of markers are ever actually outside + # (holes/ghosts and in-range markers are the overwhelming majority), so the + # sparse, index-based writes below plus the early exit (skipping this axis + # entirely when nothing is outside) touch far less memory than an + # unconditional dense pass over all n_rows markers, even accounting for the + # nonzero sync the indices cost. Avoiding a device sync is not free if the + # alternative is doing O(n_rows) dense work every call instead of O(outside + # markers) sparse work most calls skip entirely. + shift_col_0 = self.first_pusher_idx + 3 + self.vdim + periodic_outside = {} for axis in self._periodic_axes: - outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf) + outside_inds = self._find_outside_particles( + axis, + eta=self._eta_bc_buf, + not_hole_or_ghost=periodic_not_hole_or_ghost, + ) if len(outside_inds) == 0: continue + periodic_outside[axis] = ( + outside_inds, + xp.nonzero(self._is_outside_right)[0], + xp.nonzero(self._is_outside_left)[0], + ) + + # Zero the shift columns of exactly the axes that had markers outside -- the + # same set the per-axis loop below writes, so this is not a behaviour change + # (an axis with nothing outside keeps its column untouched, as before). The + # point is to do it in ONE dense pass over contiguous columns instead of one + # pass per axis: `markers[:, c] = 0.0` is a strided write over all n_rows and + # was the single most expensive operation in this function -- 1.32 ms per axis + # at Np_local=12.5M on an H100, against 0.03 ms for the sparse writes it + # exists to prepare. Hoisting it out of the loop measured 7.49 ms -> 4.93 ms + # per call for the 3-axis periodic case, with the no-markers-outside path + # unchanged (2.61 ms -> 2.69 ms, i.e. within noise) because it is still + # skipped entirely when `periodic_outside` is empty. + if periodic_outside and not newton: + active = sorted(periodic_outside) + if active[-1] - active[0] + 1 == len(active): + self.markers[:, shift_col_0 + active[0] : shift_col_0 + active[-1] + 1] = 0.0 + else: + # non-contiguous set of periodic axes: no single slice covers them + for axis in active: + self.markers[:, shift_col_0 + axis] = 0.0 + + # Pass 2 -- wrap the positions and set the shift for the alpha-weighted + # mid-point computation. + for axis, (outside_inds, outside_right_inds, outside_left_inds) in periodic_outside.items(): self.markers[outside_inds, axis] = self.markers[outside_inds, axis] % 1.0 - # set shift for alpha-weighted mid-point computation - outside_right_inds = xp.nonzero(self._is_outside_right)[0] - outside_left_inds = xp.nonzero(self._is_outside_left)[0] if newton: - self.markers[ - outside_right_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] += 1.0 - self.markers[ - outside_left_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] += -1.0 + self.markers[outside_right_inds, shift_col_0 + axis] += 1.0 + self.markers[outside_left_inds, shift_col_0 + axis] += -1.0 else: - self.markers[ - :, - self.first_pusher_idx + 3 + self.vdim + axis, - ] = 0.0 - self.markers[ - outside_right_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] = 1.0 - self.markers[ - outside_left_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] = -1.0 + self.markers[outside_right_inds, shift_col_0 + axis] = 1.0 + self.markers[outside_left_inds, shift_col_0 + axis] = -1.0 # put all coordinate inside the unit cube (avoid wrong Jacobian evaluations) outside_inds_per_axis = {} @@ -1874,21 +2166,34 @@ def apply_kinetic_bc(self, newton=False): for axis in self._reflect_axes: if len(outside_inds_per_axis[axis]) == 0: continue - # flip velocity - reflect( - self.markers, - self.domain.args_domain, - outside_inds_per_axis[axis], - axis, - ) + # flip velocity via the compiled host-only kernel, through the + # marker host mirror. + with self.host_markers(write=True) as args_markers: + reflect( + args_markers.markers, + self.domain.args_domain, + _to_numpy_for_kernel(outside_inds_per_axis[axis]), + axis, + ) - def update_holes(self): + def update_holes(self, update_valid_mks: bool = True): """Recompute the :attr:`~struphy.pic.base.Particles.holes` mask (rows with ``markers[:, 0] == -1``) and, from it, refresh :attr:`~struphy.pic.base.Particles.valid_mks`. Must be called after any operation that creates, removes or moves markers - (e.g. sorting, boundary conditions, refilling), since holes are tracked per row index.""" + (e.g. sorting, boundary conditions, refilling), since holes are tracked per row index. + + Parameters + ---------- + update_valid_mks : bool + Set to False when the caller is about to call :meth:`_update_ghost_particles` + immediately afterwards (which also refreshes ``valid_mks``, from fresh + holes *and* fresh ghosts) -- refreshing it here first would just be + overwritten a moment later with the ghost mask still stale, a full-array + pass wasted on every call to a very hot path (called every substep from + :meth:`mpi_sort_markers`).""" self._holes[:] = self.markers[:, 0] == -1.0 - self._update_valid_mks() + if update_valid_mks: + self._update_valid_mks() def set_velocities_comp(self, velocity, comp): """Set one or several velocity components to the same constant value, for all valid markers. @@ -1914,14 +2219,17 @@ 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( - self.markers, - self.holes, - self._sorting_boxes.nx, - self._sorting_boxes.ny, - self._sorting_boxes.nz, - self.domain_array[self.mpi_rank], - ) + # compiled host-only kernel; writes the box index into markers[:, -2] + # in place, through the marker host mirror. + with self.host_markers(write=True) as args_markers: + assign_box_to_each_particle( + args_markers.markers, + _to_numpy_for_kernel(self.holes), + self._sorting_boxes.nx, + self._sorting_boxes.ny, + self._sorting_boxes.nz, + self.domain_array[self.mpi_rank], + ) self._check_and_assign_particles_to_boxes() @@ -1939,15 +2247,24 @@ def put_particles_in_boxes(self): @profile @ProfileManager.profile("do_sort") - def do_sort(self, use_numpy_argsort=False): + def do_sort(self, use_numpy_argsort=None): """Assign the particles to their sorting boxes and reorder the markers array accordingly, so that markers in the same box occupy contiguous rows. 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_numpy_argsort : bool, optional + If True, sort via :func:`numpy.argsort` on the box column; if False, + use the Pyccel kernel :func:`~struphy.pic.sorting_kernels.sort_boxed_particles` + (a sequential cycle-sort). Default (None) picks the argsort path under the + CuPy backend and the Pyccel kernel under NumPy: both kernels are host-only + compiled/vectorized code -- ``self._markers`` is always host-resident (see + ``ISSUE_cupy_particles_never_pushed.md``) -- so there is nothing here for CUDA + to accelerate, but the vectorized argsort is faster than the cycle-sort on + plain CPU too (~25% at Np=2*10**6, measured), so it is the better default + whenever it's already needed anyway (i.e. under CuPy, to also avoid a redundant + code path). It stays opt-in rather than the universal default to avoid changing + existing NumPy-backend behaviour/tests. """ nx = self._sorting_boxes.nx ny = self._sorting_boxes.ny @@ -1956,23 +2273,30 @@ def do_sort(self, use_numpy_argsort=False): self.put_particles_in_boxes() + if use_numpy_argsort is None: + use_numpy_argsort = xp.cupy_backend + if use_numpy_argsort: self._sort_boxed_particles_numpy() else: - sort_boxed_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, - ) + # compiled host-only cycle-sort; reorders marker rows in place + with self.host_markers(write=True) as args_markers: + sort_boxed_particles( + args_markers.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, + ) # 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. - self.update_holes() + # _update_ghost_particles already refreshes valid_mks from fresh holes and + # fresh ghosts -- update_holes redoing it first, then a third explicit call + # redoing it again right after, were both pure repeats of the same result. + self.update_holes(update_valid_mks=False) self._update_ghost_particles() - self._update_valid_mks() def eval_density( self, @@ -2090,23 +2414,25 @@ def eval_velocity( func = PyccelKernel(eval_kernels_sph.sph_mean_velocity_coeffs) - func( - alpha=xp.array((0.0, 0.0, 0.0)), - column_nr=first_free_idx, - comps=comps, - args_markers=self.args_markers, - args_domain=self.domain.args_domain, - boxes=self.sorting_boxes.boxes, - neighbours=self.sorting_boxes.neighbours, - holes=self.holes, - periodic1=self.boundary_params.bc_sph[0] == "periodic", - periodic2=self.boundary_params.bc_sph[1] == "periodic", - periodic3=self.boundary_params.bc_sph[2] == "periodic", - kernel_type=self.ker_dct()[kernel_type], - h1=h1, - h2=h2, - h3=h3, - ) + # compiled host-only SPH kernel; writes a marker column in place + with self.host_markers(write=True) as _args_markers: + func( + alpha=xp.array((0.0, 0.0, 0.0)), + column_nr=first_free_idx, + comps=comps, + args_markers=_args_markers, + args_domain=self.domain.args_domain, + boxes=self.sorting_boxes.boxes, + neighbours=self.sorting_boxes.neighbours, + holes=self.holes, + periodic1=self.boundary_params.bc_sph[0] == "periodic", + periodic2=self.boundary_params.bc_sph[1] == "periodic", + periodic3=self.boundary_params.bc_sph[2] == "periodic", + kernel_type=self.ker_dct()[kernel_type], + h1=h1, + h2=h2, + h3=h3, + ) v1 = self._eval_sph( eta1, @@ -2202,45 +2528,49 @@ def eval_div_viscosity( # 1st kernel func = PyccelKernel(eval_kernels_sph.sph_mean_velocity_coeffs) comps = xp.array((0, 1, 2)) - func( - alpha=xp.array((0.0, 0.0, 0.0)), - column_nr=first_free_idx, - comps=comps, - args_markers=self.args_markers, - args_domain=self.domain.args_domain, - boxes=self.sorting_boxes.boxes, - neighbours=self.sorting_boxes.neighbours, - holes=self.holes, - periodic1=self.boundary_params.bc_sph[0] == "periodic", - periodic2=self.boundary_params.bc_sph[1] == "periodic", - periodic3=self.boundary_params.bc_sph[2] == "periodic", - kernel_type=self.ker_dct()[kernel_type], - h1=h1, - h2=h2, - h3=h3, - ) + # compiled host-only SPH kernel; writes a marker column in place + with self.host_markers(write=True) as _args_markers: + func( + alpha=xp.array((0.0, 0.0, 0.0)), + column_nr=first_free_idx, + comps=comps, + args_markers=_args_markers, + args_domain=self.domain.args_domain, + boxes=self.sorting_boxes.boxes, + neighbours=self.sorting_boxes.neighbours, + holes=self.holes, + periodic1=self.boundary_params.bc_sph[0] == "periodic", + periodic2=self.boundary_params.bc_sph[1] == "periodic", + periodic3=self.boundary_params.bc_sph[2] == "periodic", + kernel_type=self.ker_dct()[kernel_type], + h1=h1, + h2=h2, + h3=h3, + ) # 2nd kernel func = PyccelKernel(eval_kernels_sph.sph_viscosity_tensor) comps = xp.arange(9) - func( - alpha=xp.array((0.0, 0.0, 0.0)), - column_nr=first_free_idx + 3, - comps=comps, - args_markers=self.args_markers, - args_domain=self.domain.args_domain, - boxes=self.sorting_boxes.boxes, - neighbours=self.sorting_boxes.neighbours, - holes=self.holes, - periodic1=self.boundary_params.bc_sph[0] == "periodic", - periodic2=self.boundary_params.bc_sph[1] == "periodic", - periodic3=self.boundary_params.bc_sph[2] == "periodic", - kernel_type=self.ker_dct()[kernel_type], - h1=h1, - h2=h2, - h3=h3, - mu=mu, - ) + # compiled host-only SPH kernel; writes marker columns in place + with self.host_markers(write=True) as _args_markers: + func( + alpha=xp.array((0.0, 0.0, 0.0)), + column_nr=first_free_idx + 3, + comps=comps, + args_markers=_args_markers, + args_domain=self.domain.args_domain, + boxes=self.sorting_boxes.boxes, + neighbours=self.sorting_boxes.neighbours, + holes=self.holes, + periodic1=self.boundary_params.bc_sph[0] == "periodic", + periodic2=self.boundary_params.bc_sph[1] == "periodic", + periodic3=self.boundary_params.bc_sph[2] == "periodic", + kernel_type=self.ker_dct()[kernel_type], + h1=h1, + h2=h2, + h3=h3, + mu=mu, + ) # grid evaluation gamma = [] @@ -2299,7 +2629,7 @@ def gather_scalar_in_subcomm_array(self, scalar: int, out: xp.ndarray = None): The returned array (optional). """ if out is None: - _tmp = xp.zeros(self.mpi_size, dtype=int) + _tmp = np.zeros(self.mpi_size, dtype=int) else: assert out.size == self.mpi_size _tmp = out @@ -2327,7 +2657,7 @@ def gather_scalar_in_intercomm_array(self, scalar: int, out: xp.ndarray = None): The returned array (optional). """ if out is None: - _tmp = xp.zeros(self.num_clones, dtype=int) + _tmp = np.zeros(self.num_clones, dtype=int) else: assert out.size == self.num_clones _tmp = out @@ -2351,7 +2681,7 @@ def gather_scalar_in_intercomm_array(self, scalar: int, out: xp.ndarray = None): def _update_valid_mks(self): """Refresh :attr:`~struphy.pic.base.Particles.valid_mks`: a row is a valid marker if and only if it is neither a hole nor a ghost particle.""" - self._valid_mks[:] = ~xp.logical_or(self.holes, self.ghost_particles) + self._valid_mks[:] = ~np.logical_or(self.holes, self.ghost_particles) def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): """ @@ -2377,7 +2707,7 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): if mpi_dims_mask is None: mpi_dims_mask = [True, True, True] - dom_arr = xp.zeros((self.mpi_size, 9), dtype=float) + dom_arr = np.zeros((self.mpi_size, 9), dtype=float) # factorize mpi size factors = factorint(self.mpi_size) @@ -2401,10 +2731,13 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): mm = (mm + 1) % 3 nprocs[mm] *= fac - assert xp.prod(nprocs) == self.mpi_size + # nprocs is a plain 3-element Python list (process counts, not + # physics data); np.prod handles it on both backends, unlike + # xp.prod which (on CuPy) requires an actual ndarray input. + assert np.prod(nprocs) == self.mpi_size # domain decomposition - breaks = [xp.linspace(0.0, 1.0, nproc + 1) for nproc in nprocs] + breaks = [np.linspace(0.0, 1.0, nproc + 1) for nproc in nprocs] # fill domain array for n in range(self.mpi_size): @@ -2492,15 +2825,25 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return + # The marker array and every array indexed by marker row live on the + # active backend: on the device under CuPy, where the ported CUDA + # particle kernels and all the vectorized bookkeeping below (boundary + # conditions, hole/ghost tracking, sorting) operate on them directly, + # with no per-call host<->device transfer. The remaining host-only + # Pyccel kernels reach them through the explicit host mirror + # (see :class:`_HostMarkerMirror` and :meth:`host_markers`). self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float) - # allocate auxiliary arrays self._holes = xp.zeros(self.n_rows, dtype=bool) self._ghost_particles = xp.zeros(self.n_rows, dtype=bool) self._valid_mks = xp.zeros(self.n_rows, dtype=bool) - self._is_outside_right = xp.zeros(self.n_rows, dtype=bool) - self._is_outside_left = xp.zeros(self.n_rows, dtype=bool) - self._is_outside = xp.zeros(self.n_rows, dtype=bool) + # _is_outside_right/_is_outside_left/_is_outside are views into one + # buffer, so the three masks stay contiguous for the combined + # comparison in _find_outside_particles. + self._is_outside_buf = xp.zeros((3, self.n_rows), dtype=bool) + self._is_outside_right = self._is_outside_buf[0] + self._is_outside_left = self._is_outside_buf[1] + self._is_outside = self._is_outside_buf[2] # contiguous scratch copy of markers[:, :3], refreshed once per apply_kinetic_bc # boundary-condition-type loop (see there) instead of re-striding into the full # (n_rows, n_cols) row-major marker array once per axis. @@ -2510,10 +2853,24 @@ def _allocate_marker_array(self, dry_run: bool = False): self._n_lost_markers = 0 self._lost_markers = xp.zeros((int(self.n_rows * 0.5), 10), dtype=float) + # Host mirror for the compiled, host-only Pyccel kernels. Under the + # NumPy backend there is nothing to mirror -- markers already are a + # host array -- so args_markers keeps aliasing it directly and + # host_markers() is a no-op. + if xp.cupy_backend: + self._markers_mirror = _HostMarkerMirror() + markers_for_kernels = self._markers_mirror.add(self._markers) + valid_mks_for_kernels = self._markers_mirror.add(self._valid_mks) + else: + self._markers_mirror = None + markers_for_kernels = self._markers + valid_mks_for_kernels = self._valid_mks + # arguments for kernels + self._args_markers = MarkerArguments( - _to_numpy_for_kernel(self.markers), - _to_numpy_for_kernel(self.valid_mks), + markers_for_kernels, + valid_mks_for_kernels, _to_numpy_for_kernel(self.Np), _to_numpy_for_kernel(self.vdim), _to_numpy_for_kernel(self.index["weights"]), @@ -2525,6 +2882,9 @@ def _allocate_marker_array(self, dry_run: bool = False): _to_numpy_for_kernel(self.mu_idx), ) + if "cupy": + self._args_markers = transform() + def _initialize_sorting_boxes(self): """Initializes the sorting boxes. @@ -2696,6 +3056,14 @@ def _set_initial_condition(self): # TODO: add other velocity components def _f_init(*etas, flat_eval=False): + # etas may come from a host-resident marker property (e.g. + # Particles.positions, always NumPy -- see + # ISSUE_cupy_particles_never_pushed.md), while self.f0.n0/ + # _density are generic field-evaluation callables that follow + # the active array backend; convert on entry so both callers + # that already pass backend-native points (a no-op then) and + # ones that pass host marker data work correctly. + etas = tuple(_dev(e) for e in etas) if len(etas) == 1: if _density is None: out = self.f0.n0(etas[0]) @@ -2722,6 +3090,8 @@ def _f_init(*etas, flat_eval=False): return out def _u_init(*etas, flat_eval=False): + # see _f_init above for why this conversion is needed. + etas = tuple(_dev(e) for e in etas) if len(etas) == 1: out = self.f0.uv(etas[0]) if _u1 is not None: @@ -2785,7 +3155,7 @@ def _load_external( dtype=float, ) self._mpi_comm.Recv(recvbuf, source=0, tag=123) - self._markers[:n_mks_load_loc, :] = recvbuf + self._markers[:n_mks_load_loc, :] = xp.asarray(recvbuf) def _load_restart(self): """Load markers from restart .hdf5 file.""" @@ -2804,7 +3174,7 @@ def _load_restart(self): data = DataContainer(data_path, comm=self.mpi_comm) with h5py.File(data.file_path, "a") as file: - self._markers[:, :] = file["restart/" + self.loading_params.restart_key][-1, :, :] + self._markers[:, :] = xp.asarray(file["restart/" + self.loading_params.restart_key][-1, :, :]) def _load_tesselation(self, n_quad: int = 1): """ @@ -2822,22 +3192,23 @@ def _load_tesselation(self, n_quad: int = 1): sorting_boxes=self.sorting_boxes, ) eta1, eta2, eta3 = self.tesselation.draw_markers() - self._markers[: eta1.size, 0] = eta1 - self._markers[: eta2.size, 1] = eta2 - self._markers[: eta3.size, 2] = eta3 + eta1, eta2, eta3 = _to_numpy_for_kernel(eta1), _to_numpy_for_kernel(eta2), _to_numpy_for_kernel(eta3) + self._markers[: eta1.size, 0] = xp.asarray(eta1) + self._markers[: eta2.size, 1] = xp.asarray(eta2) + self._markers[: eta3.size, 2] = xp.asarray(eta3) self._update_valid_mks() def _reset_marker_ids(self): """Reset the marker ids (last column in marker array) according to the current distribution of particles. The first marker on rank 0 gets the id '0', the last marker on the last rank gets the id 'n_mks_global - 1'.""" - n_mks_proc_cumsum = xp.cumsum(self.n_mks_on_each_proc) - n_mks_clone_cumsum = xp.cumsum(self.n_mks_on_each_clone) + n_mks_proc_cumsum = np.cumsum(self.n_mks_on_each_proc) + n_mks_clone_cumsum = np.cumsum(self.n_mks_on_each_clone) first_marker_id = (n_mks_clone_cumsum - self.n_mks_on_each_clone)[self.clone_id] + ( n_mks_proc_cumsum - self.n_mks_on_each_proc )[self.mpi_rank] - self.marker_ids = first_marker_id + xp.arange(self.n_mks_loc, dtype=int) + self.marker_ids = first_marker_id + np.arange(self.n_mks_loc, dtype=int) - def _find_outside_particles(self, axis, eta=None): + def _find_outside_particles(self, axis, eta=None, not_hole_or_ghost=None): """Find markers whose ``axis``-th logical coordinate lies outside ``[0, 1]`` (holes and ghost particles are excluded), updating :attr:`_is_outside_left`/:attr:`_is_outside_right`/:attr:`_is_outside` accordingly. @@ -2853,25 +3224,35 @@ def _find_outside_particles(self, axis, eta=None): correct but slower, since a single-column slice of the row-major ``markers`` array is strided (see the comment in :meth:`apply_kinetic_bc`). + not_hole_or_ghost : xp.ndarray[bool], optional + Pre-computed ``~(self.holes | self.ghost_particles)``, for callers that loop + over several axes without any hole/ghost-changing operation in between (e.g. + the periodic-axis loop in :meth:`apply_kinetic_bc`, which only wraps + positions) -- recomputing this per axis is redundant there. If ``None``, + computed fresh (correct default for callers that may create holes/ghosts + between axes, e.g. the remove-axis loop). + Returns ------- - outside_inds : xp.ndarray[int] + outside_inds : numpy.ndarray[int] Row indices of the markers that are outside the logical unit cube. """ + # Runs on whichever backend the markers live on -- device under + # CuPy, with no transfer: markers, the holes/ghost masks and the + # _is_outside_* views are all allocated with xp (see + # _allocate_marker_array). col = self.markers[:, axis] if eta is None else eta[:, axis] - - # determine particles outside of the logical unit cube - self._is_outside_right[:] = col > 1.0 - self._is_outside_left[:] = col < 0.0 - - self._is_outside_right[self.holes] = False - self._is_outside_right[self.ghost_particles] = False - self._is_outside_left[self.holes] = False - self._is_outside_left[self.ghost_particles] = False - - self._is_outside[:] = xp.logical_or( + if not_hole_or_ghost is None: + not_hole_or_ghost = ~(self.holes | self.ghost_particles) + + xp.greater(col, 1.0, out=self._is_outside_right) + self._is_outside_right &= not_hole_or_ghost + xp.less(col, 0.0, out=self._is_outside_left) + self._is_outside_left &= not_hole_or_ghost + xp.logical_or( self._is_outside_right, self._is_outside_left, + out=self._is_outside, ) # indices or particles that are outside of the logical unit cube @@ -2899,12 +3280,12 @@ def _particle_refilling(self): for kind in self.bc_refill: # sorting out particles which are out of the domain if kind == "inner": - outside_inds = xp.nonzero(self._is_outside_left)[0] + outside_inds = np.nonzero(self._is_outside_left)[0] self.markers[outside_inds, 0] = 1e-4 r_loss = self.domain.params["a1"] else: - outside_inds = xp.nonzero(self._is_outside_right)[0] + outside_inds = np.nonzero(self._is_outside_right)[0] self.markers[outside_inds, 0] = 1 - 1e-4 r_loss = 1.0 @@ -3018,31 +3399,32 @@ def _gyro_transfer(self, outside_inds): return xp.logical_and(1.0 > gc_etas[0], gc_etas[0] > 0.0) def _sort_boxed_particles_numpy(self): - """Sort the particles by box using numpy.argsort.""" + """Sort the particles by box using numpy.argsort. + + ``_argsort_array`` lives on the same backend as the markers, so both + the argsort and the subsequent gather run entirely on the device + under CuPy. + """ sorting_axis = self._sorting_boxes.box_index if not hasattr(self, "_argsort_array"): self._argsort_array = xp.zeros(self.markers.shape[0], dtype=int) self._argsort_array[:] = self._markers[:, sorting_axis].argsort() + # gather into a temporary: an in-place fancy-index self-assignment is + # not safe (source and destination overlap). 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 + # markers may be CuPy-resident under the active backend, but np.bincount + # dispatches to cupy's implementation via the array API protocol, so + # this works unconditionally without an explicit host round-trip. + bcount = np.bincount(self.markers_wo_holes[:, -2].astype(np.int64)) - 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) - - max_in_box = xp.max(bcount) + max_in_box = np.max(bcount) if max_in_box > self._sorting_boxes.boxes.shape[1]: warnings.warn( f'Strong load imbalance detected in sorting boxes: \ @@ -3052,27 +3434,46 @@ def _check_and_assign_particles_to_boxes(self): ) self.mpi_comm.Abort() - assign_particles_to_boxes( - self.markers, - self.holes, - self._sorting_boxes._boxes, - self._sorting_boxes._next_index, - ) + # compiled host-only kernel; markers are read-only here, but boxes/ + # next_index are fully overwritten, so round-trip them through host + # buffers and write the results back onto the active backend. + boxes = _to_numpy_for_kernel(self._sorting_boxes._boxes).copy() + next_index = _to_numpy_for_kernel(self._sorting_boxes._next_index).copy() + with self.host_markers(write=False) as args_markers: + assign_particles_to_boxes( + args_markers.markers, + _to_numpy_for_kernel(self.holes), + boxes, + next_index, + ) + self._sorting_boxes._boxes[:, :] = xp.asarray(boxes) + self._sorting_boxes._next_index[:] = xp.asarray(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 :meth:`_prepare_ghost_particles`/:meth:`_sendrecv_markers_boxes` for SPH ghost-box particles received from a neighbouring process.""" - self._ghost_particles[:] = self.markers[:, -1] == -2.0 + # markers[:, -1] is a strided read over every row (~0.5 ms at Np_local=12.5M on + # an H100). When no ghost particle has been created since the last removal the + # answer is known to be all-False and the mask already holds it, so only the + # cheap valid_mks refresh is needed -- holes may still have changed. + if self.has_ghost_particles: + self._ghost_particles[:] = self.markers[:, -1] == -2.0 self._update_valid_mks() def _remove_ghost_particles(self): """Discard all current ghost particles: turn their marker-array rows into new holes (so the space can be reused before the next SPH ghost-box update).""" - self._update_ghost_particles() - new_holes = xp.nonzero(self.ghost_particles) - self._markers[new_holes] = -1.0 + # Skip the ghost-specific work (a strided full-array compare plus a nonzero, + # which also forces a device sync) when no ghost exists to remove. update_holes + # still runs unconditionally: callers rely on it refreshing holes, which other + # operations do change. + if self.has_ghost_particles: + self._update_ghost_particles() + new_holes = np.nonzero(self.ghost_particles) + self._markers[new_holes] = -1.0 + self._has_ghost_particles = False self.update_holes() def _prepare_ghost_particles(self): @@ -3311,16 +3712,20 @@ def _mirror_particles( if "x_m" in arr_name and is_domain_boundary["x_m"]: arr[:, 0] *= -1.0 if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: - boundary_values = self.f_init( - *arr[:, :3].T, - flat_eval=True, + boundary_values = _to_numpy_for_kernel( + self.f_init( + *_dev(*arr[:, :3].T), + flat_eval=True, + ), ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = ( -boundary_values - / self.s0( - *arr[:, :3].T, - flat_eval=True, - remove_holes=False, + / _to_numpy_for_kernel( + self.s0( + *_dev(*arr[:, :3].T), + flat_eval=True, + remove_holes=False, + ), ) / self.Np ) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right @@ -3338,16 +3743,20 @@ def _mirror_particles( elif "x_p" in arr_name and is_domain_boundary["x_p"]: arr[:, 0] = 2.0 - arr[:, 0] if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: - boundary_values = self.f_init( - *arr[:, :3].T, - flat_eval=True, + boundary_values = _to_numpy_for_kernel( + self.f_init( + *_dev(*arr[:, :3].T), + flat_eval=True, + ), ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = ( -boundary_values - / self.s0( - *arr[:, :3].T, - flat_eval=True, - remove_holes=False, + / _to_numpy_for_kernel( + self.s0( + *_dev(*arr[:, :3].T), + flat_eval=True, + remove_holes=False, + ), ) / self.Np ) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right @@ -3367,16 +3776,20 @@ def _mirror_particles( if "y_m" in arr_name and is_domain_boundary["y_m"]: arr[:, 1] *= -1.0 if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: - boundary_values = self.f_init( - *arr[:, :3].T, - flat_eval=True, + boundary_values = _to_numpy_for_kernel( + self.f_init( + *_dev(*arr[:, :3].T), + flat_eval=True, + ), ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = ( -boundary_values - / self.s0( - *arr[:, :3].T, - flat_eval=True, - remove_holes=False, + / _to_numpy_for_kernel( + self.s0( + *_dev(*arr[:, :3].T), + flat_eval=True, + remove_holes=False, + ), ) / self.Np ) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right @@ -3394,16 +3807,20 @@ def _mirror_particles( elif "y_p" in arr_name and is_domain_boundary["y_p"]: arr[:, 1] = 2.0 - arr[:, 1] if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: - boundary_values = self.f_init( - *arr[:, :3].T, - flat_eval=True, + boundary_values = _to_numpy_for_kernel( + self.f_init( + *_dev(*arr[:, :3].T), + flat_eval=True, + ), ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = ( -boundary_values - / self.s0( - *arr[:, :3].T, - flat_eval=True, - remove_holes=False, + / _to_numpy_for_kernel( + self.s0( + *_dev(*arr[:, :3].T), + flat_eval=True, + remove_holes=False, + ), ) / self.Np ) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right @@ -3423,16 +3840,20 @@ def _mirror_particles( if "z_m" in arr_name and is_domain_boundary["z_m"]: arr[:, 2] *= -1.0 if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: - boundary_values = self.f_init( - *arr[:, :3].T, - flat_eval=True, + boundary_values = _to_numpy_for_kernel( + self.f_init( + *_dev(*arr[:, :3].T), + flat_eval=True, + ), ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = ( -boundary_values - / self.s0( - *arr[:, :3].T, - flat_eval=True, - remove_holes=False, + / _to_numpy_for_kernel( + self.s0( + *_dev(*arr[:, :3].T), + flat_eval=True, + remove_holes=False, + ), ) / self.Np ) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right @@ -3450,16 +3871,20 @@ def _mirror_particles( elif "z_p" in arr_name and is_domain_boundary["z_p"]: arr[:, 2] = 2.0 - arr[:, 2] if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: - boundary_values = self.f_init( - *arr[:, :3].T, - flat_eval=True, + boundary_values = _to_numpy_for_kernel( + self.f_init( + *_dev(*arr[:, :3].T), + flat_eval=True, + ), ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = ( -boundary_values - / self.s0( - *arr[:, :3].T, - flat_eval=True, - remove_holes=False, + / _to_numpy_for_kernel( + self.s0( + *_dev(*arr[:, :3].T), + flat_eval=True, + remove_holes=False, + ), ) / self.Np ) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right @@ -3492,10 +3917,13 @@ def _determine_markers_in_box(self, list_boxes): """ indices = [] for i in list_boxes: - indices += list(self._sorting_boxes._boxes[i][self._sorting_boxes._boxes[i] != -1]) + box_row = _to_numpy_for_kernel(self._sorting_boxes._boxes[i]) + indices += list(box_row[box_row != -1]) - indices = xp.array(indices, dtype=int) - markers_in_box = self.markers[indices] + # Box membership is host bookkeeping; the gathered rows are handed to + # the mpi4py box-communication path below, which needs host buffers. + indices = np.array(indices, dtype=int) + markers_in_box = _to_numpy_for_kernel(self.markers[xp.asarray(indices)]) return markers_in_box def _get_destinations_box(self): @@ -3504,156 +3932,164 @@ def _get_destinations_box(self): :meth:`_get_neighbouring_proc`), accumulating, per destination rank, the number of markers to send (:attr:`_send_info_box`) and the markers themselves (:attr:`_send_list_box`, used by :meth:`_self_communication_boxes` and - :meth:`_sendrecv_markers_boxes`).""" - self._send_info_box = xp.zeros(self.mpi_size, dtype=int) - self._send_list_box = [xp.zeros((0, self.n_cols))] * self.mpi_size + :meth:`_sendrecv_markers_boxes`). + + This method and the rest of the box-communication subsystem below are + host-resident throughout, like the analogous MPI marker-sort methods + (:meth:`_sendrecv_determine_mtbs` etc.): they operate on rows of + ``self.markers`` (always host, see ISSUE_cupy_particles_never_pushed.md) + and feed counts/buffers straight into mpi4py Alltoall/Isend/Irecv, + which need host-readable arguments regardless of backend. + """ + self._send_info_box = np.zeros(self.mpi_size, dtype=int) + self._send_list_box = [np.zeros((0, self.n_cols))] * self.mpi_size # Faces # if self._x_m_proc is not None: self._send_info_box[self._x_m_proc] += len(self._markers_x_m) - self._send_list_box[self._x_m_proc] = xp.concatenate((self._send_list_box[self._x_m_proc], self._markers_x_m)) + self._send_list_box[self._x_m_proc] = np.concatenate((self._send_list_box[self._x_m_proc], self._markers_x_m)) # if self._x_p_proc is not None: self._send_info_box[self._x_p_proc] += len(self._markers_x_p) - self._send_list_box[self._x_p_proc] = xp.concatenate((self._send_list_box[self._x_p_proc], self._markers_x_p)) + self._send_list_box[self._x_p_proc] = np.concatenate((self._send_list_box[self._x_p_proc], self._markers_x_p)) # if self._y_m_proc is not None: self._send_info_box[self._y_m_proc] += len(self._markers_y_m) - self._send_list_box[self._y_m_proc] = xp.concatenate((self._send_list_box[self._y_m_proc], self._markers_y_m)) + self._send_list_box[self._y_m_proc] = np.concatenate((self._send_list_box[self._y_m_proc], self._markers_y_m)) # if self._y_p_proc is not None: self._send_info_box[self._y_p_proc] += len(self._markers_y_p) - self._send_list_box[self._y_p_proc] = xp.concatenate((self._send_list_box[self._y_p_proc], self._markers_y_p)) + self._send_list_box[self._y_p_proc] = np.concatenate((self._send_list_box[self._y_p_proc], self._markers_y_p)) # if self._z_m_proc is not None: self._send_info_box[self._z_m_proc] += len(self._markers_z_m) - self._send_list_box[self._z_m_proc] = xp.concatenate((self._send_list_box[self._z_m_proc], self._markers_z_m)) + self._send_list_box[self._z_m_proc] = np.concatenate((self._send_list_box[self._z_m_proc], self._markers_z_m)) # if self._z_p_proc is not None: self._send_info_box[self._z_p_proc] += len(self._markers_z_p) - self._send_list_box[self._z_p_proc] = xp.concatenate((self._send_list_box[self._z_p_proc], self._markers_z_p)) + self._send_list_box[self._z_p_proc] = np.concatenate((self._send_list_box[self._z_p_proc], self._markers_z_p)) # x-y edges # if self._x_m_y_m_proc is not None: self._send_info_box[self._x_m_y_m_proc] += len(self._markers_x_m_y_m) - self._send_list_box[self._x_m_y_m_proc] = xp.concatenate( + self._send_list_box[self._x_m_y_m_proc] = np.concatenate( (self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m), ) # if self._x_m_y_p_proc is not None: self._send_info_box[self._x_m_y_p_proc] += len(self._markers_x_m_y_p) - self._send_list_box[self._x_m_y_p_proc] = xp.concatenate( + self._send_list_box[self._x_m_y_p_proc] = np.concatenate( (self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p), ) # if self._x_p_y_m_proc is not None: self._send_info_box[self._x_p_y_m_proc] += len(self._markers_x_p_y_m) - self._send_list_box[self._x_p_y_m_proc] = xp.concatenate( + self._send_list_box[self._x_p_y_m_proc] = np.concatenate( (self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m), ) # if self._x_p_y_p_proc is not None: self._send_info_box[self._x_p_y_p_proc] += len(self._markers_x_p_y_p) - self._send_list_box[self._x_p_y_p_proc] = xp.concatenate( + self._send_list_box[self._x_p_y_p_proc] = np.concatenate( (self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p), ) # x-z edges # if self._x_m_z_m_proc is not None: self._send_info_box[self._x_m_z_m_proc] += len(self._markers_x_m_z_m) - self._send_list_box[self._x_m_z_m_proc] = xp.concatenate( + self._send_list_box[self._x_m_z_m_proc] = np.concatenate( (self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m), ) # if self._x_m_z_p_proc is not None: self._send_info_box[self._x_m_z_p_proc] += len(self._markers_x_m_z_p) - self._send_list_box[self._x_m_z_p_proc] = xp.concatenate( + self._send_list_box[self._x_m_z_p_proc] = np.concatenate( (self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p), ) # if self._x_p_z_m_proc is not None: self._send_info_box[self._x_p_z_m_proc] += len(self._markers_x_p_z_m) - self._send_list_box[self._x_p_z_m_proc] = xp.concatenate( + self._send_list_box[self._x_p_z_m_proc] = np.concatenate( (self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m), ) # if self._x_p_z_p_proc is not None: self._send_info_box[self._x_p_z_p_proc] += len(self._markers_x_p_z_p) - self._send_list_box[self._x_p_z_p_proc] = xp.concatenate( + self._send_list_box[self._x_p_z_p_proc] = np.concatenate( (self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p), ) # y-z edges # if self._y_m_z_m_proc is not None: self._send_info_box[self._y_m_z_m_proc] += len(self._markers_y_m_z_m) - self._send_list_box[self._y_m_z_m_proc] = xp.concatenate( + self._send_list_box[self._y_m_z_m_proc] = np.concatenate( (self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m), ) # if self._y_m_z_p_proc is not None: self._send_info_box[self._y_m_z_p_proc] += len(self._markers_y_m_z_p) - self._send_list_box[self._y_m_z_p_proc] = xp.concatenate( + self._send_list_box[self._y_m_z_p_proc] = np.concatenate( (self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p), ) # if self._y_p_z_m_proc is not None: self._send_info_box[self._y_p_z_m_proc] += len(self._markers_y_p_z_m) - self._send_list_box[self._y_p_z_m_proc] = xp.concatenate( + self._send_list_box[self._y_p_z_m_proc] = np.concatenate( (self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m), ) # if self._y_p_z_p_proc is not None: self._send_info_box[self._y_p_z_p_proc] += len(self._markers_y_p_z_p) - self._send_list_box[self._y_p_z_p_proc] = xp.concatenate( + self._send_list_box[self._y_p_z_p_proc] = np.concatenate( (self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p), ) # corners # if self._x_m_y_m_z_m_proc is not None: self._send_info_box[self._x_m_y_m_z_m_proc] += len(self._markers_x_m_y_m_z_m) - self._send_list_box[self._x_m_y_m_z_m_proc] = xp.concatenate( + self._send_list_box[self._x_m_y_m_z_m_proc] = np.concatenate( (self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m), ) # if self._x_m_y_m_z_p_proc is not None: self._send_info_box[self._x_m_y_m_z_p_proc] += len(self._markers_x_m_y_m_z_p) - self._send_list_box[self._x_m_y_m_z_p_proc] = xp.concatenate( + self._send_list_box[self._x_m_y_m_z_p_proc] = np.concatenate( (self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p), ) # if self._x_m_y_p_z_m_proc is not None: self._send_info_box[self._x_m_y_p_z_m_proc] += len(self._markers_x_m_y_p_z_m) - self._send_list_box[self._x_m_y_p_z_m_proc] = xp.concatenate( + self._send_list_box[self._x_m_y_p_z_m_proc] = np.concatenate( (self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m), ) # if self._x_m_y_p_z_p_proc is not None: self._send_info_box[self._x_m_y_p_z_p_proc] += len(self._markers_x_m_y_p_z_p) - self._send_list_box[self._x_m_y_p_z_p_proc] = xp.concatenate( + self._send_list_box[self._x_m_y_p_z_p_proc] = np.concatenate( (self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p), ) # if self._x_p_y_m_z_m_proc is not None: self._send_info_box[self._x_p_y_m_z_m_proc] += len(self._markers_x_p_y_m_z_m) - self._send_list_box[self._x_p_y_m_z_m_proc] = xp.concatenate( + self._send_list_box[self._x_p_y_m_z_m_proc] = np.concatenate( (self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m), ) # if self._x_p_y_m_z_p_proc is not None: self._send_info_box[self._x_p_y_m_z_p_proc] += len(self._markers_x_p_y_m_z_p) - self._send_list_box[self._x_p_y_m_z_p_proc] = xp.concatenate( + self._send_list_box[self._x_p_y_m_z_p_proc] = np.concatenate( (self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p), ) # if self._x_p_y_p_z_m_proc is not None: self._send_info_box[self._x_p_y_p_z_m_proc] += len(self._markers_x_p_y_p_z_m) - self._send_list_box[self._x_p_y_p_z_m_proc] = xp.concatenate( + self._send_list_box[self._x_p_y_p_z_m_proc] = np.concatenate( (self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m), ) # if self._x_p_y_p_z_p_proc is not None: self._send_info_box[self._x_p_y_p_z_p_proc] += len(self._markers_x_p_y_p_z_p) - self._send_list_box[self._x_p_y_p_z_p_proc] = xp.concatenate( + self._send_list_box[self._x_p_y_p_z_p_proc] = np.concatenate( (self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p), ) @@ -3663,7 +4099,7 @@ def _self_communication_boxes(self): if self._send_info_box[self.mpi_rank] > 0: self.update_holes() - holes_inds = xp.nonzero(self.holes)[0] + holes_inds = np.nonzero(self.holes)[0] if holes_inds.size < self._send_info_box[self.mpi_rank]: warnings.warn( @@ -3685,9 +4121,9 @@ def _self_communication_boxes(self): # self.update_holes() # self._update_ghost_particles() # self._update_valid_mks() - # holes_inds = xp.nonzero(self.holes)[0] + # holes_inds = np.nonzero(self.holes)[0] - self.markers[holes_inds[xp.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[self.mpi_rank] + self.markers[holes_inds[np.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[self.mpi_rank] def _sendrecv_all_to_all_boxes(self): """ @@ -3695,7 +4131,7 @@ def _sendrecv_all_to_all_boxes(self): for the communication of particles in boundary boxes. """ - self._recv_info_box = xp.zeros(self.mpi_comm.Get_size(), dtype=int) + self._recv_info_box = np.zeros(self.mpi_comm.Get_size(), dtype=int) self.mpi_comm.Alltoall(self._send_info_box, self._recv_info_box) @@ -3706,8 +4142,8 @@ def _sendrecv_markers_boxes(self): """ # i-th entry holds the number (not the index) of the first hole to be filled by data from process i - first_hole = xp.cumsum(self._recv_info_box) - self._recv_info_box - hole_inds = xp.nonzero(self._holes)[0] + first_hole = np.cumsum(self._recv_info_box) - self._recv_info_box + hole_inds = np.nonzero(self._holes)[0] # Initialize send and receive commands reqs = [] recvbufs = [] @@ -3718,7 +4154,7 @@ def _sendrecv_markers_boxes(self): else: self.mpi_comm.Isend(data, dest=i, tag=self.mpi_comm.Get_rank()) - recvbufs += [xp.zeros((N_recv, self._markers.shape[1]), dtype=float)] + recvbufs += [np.zeros((N_recv, self._markers.shape[1]), dtype=float)] reqs += [self.mpi_comm.Irecv(recvbufs[-1], source=i, tag=i)] # Wait for buffer, then put markers into holes @@ -3741,7 +4177,7 @@ def _sendrecv_markers_boxes(self): self.mpi_comm.Abort() # exit() - self._markers[hole_inds[first_hole[i] + xp.arange(self._recv_info_box[i])]] = recvbufs[i] + self._markers[hole_inds[first_hole[i] + np.arange(self._recv_info_box[i])]] = recvbufs[i] test_reqs.pop() reqs[i] = None @@ -4176,15 +4612,26 @@ def _communicate_boxes(self): # n_ghosts = xp.count_nonzero(self.ghost_particles) # logger.info(f"before communicate_boxes: {self.mpi_rank = }, {n_valid = } {n_holes = }, {n_ghosts = }") + # This is the one path that turns rows into ghost particles (it writes -2 into + # the ID column of the outgoing ghost markers, and receives rows already + # carrying it), so it is the one place that arms the flag. It must be set + # before the _update_ghost_particles call below, which is gated on it. + self._has_ghost_particles = True + self._prepare_ghost_particles() self._get_destinations_box() self._self_communication_boxes() - self.update_holes() + # Both update_holes calls below defer valid_mks to the final + # _update_ghost_particles call, which is always the last of these to run + # (whichever update_holes call precedes it) and refreshes it from fresh + # holes and fresh ghosts in one pass, instead of every intermediate call + # redoing the same full-array work with a ghost mask that's about to change. + self.update_holes(update_valid_mks=False) if self.mpi_comm is not None: self._Barrier() self._sendrecv_all_to_all_boxes() self._sendrecv_markers_boxes() - self.update_holes() + self.update_holes(update_valid_mks=False) self._update_ghost_particles() # if verbose: @@ -4307,26 +4754,28 @@ def _eval_sph( func = PyccelKernel(naive_evaluation_flat) elif len(_shp) == 3: func = PyccelKernel(naive_evaluation_meshgrid) - func( - self.args_markers, - eta1, - eta2, - eta3, - self.holes, - periodic1, - periodic2, - periodic3, - index, - ker_id, - h1, - h2, - h3, - out, - ) + with self.host_markers(write=False) as args_markers: + func( + args_markers, + eta1, + eta2, + eta3, + self.holes, + periodic1, + periodic2, + periodic3, + index, + ker_id, + h1, + h2, + h3, + out, + ) return out ### MPI comm for domain decomposition ### + @ProfileManager.profile("_sendrecv_determine_mtbs") def _sendrecv_determine_mtbs( self, alpha: list | tuple | xp.ndarray = (1.0, 1.0, 1.0), @@ -4350,34 +4799,83 @@ def _sendrecv_determine_mtbs( sorting_etas : array[float] Eta-values of shape (n_send, :) according to which the sorting is performed. """ - # position that determines the sorting (including periodic shift of boundary conditions) - if not isinstance(alpha, xp.ndarray): - alpha = xp.array(alpha, dtype=float) - assert alpha.size == 3 - assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) - bi = self.first_pusher_idx - xp.mod( - alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) - + (1.0 - alpha) * self.markers[:, bi : bi + 3], - 1.0, - out=self._sorting_etas, - ) - - # check which particles are on the current process domain - self._is_on_proc_domain = xp.logical_and( - self._sorting_etas > self.domain_array[self.mpi_rank, 0::3], - self._sorting_etas < self.domain_array[self.mpi_rank, 1::3], + # Fast path: alpha == 1 collapses alpha*(A + B) + (1 - alpha)*C to exactly A + B + # (1*x == x and 0*C == 0 for the finite phase-space values C holds in practice, + # so this is not an approximation -- verified bit-for-bit identical to the + # general formula below). Checked in plain Python against the raw, unconverted + # argument -- most callers pass the default (a Python tuple/float), so this + # costs nothing when it doesn't apply and never touches a device array or + # forces a sync just to decide. Only an xp.ndarray alpha (the dynamic, + # per-kernel case in pusher.py) skips the check and always takes the general + # path below, since its value isn't known without a sync anyway. + alpha_is_one = (isinstance(alpha, (int, float)) and alpha == 1.0) or ( + not isinstance(alpha, xp.ndarray) and hasattr(alpha, "__iter__") and all(a == 1.0 for a in alpha) ) + bi = self.first_pusher_idx + if alpha_is_one: + # Measured ~2.2x faster than the general formula below on CuPy at + # Np_local=12.5M (2.9ms -> 1.3ms): the general path is 2 strided column + # reads, an add, 2 multiplies and a second add -- several separate kernel + # launches over non-contiguous (strided) column slices. This fast path + # keeps only the one unavoidable add. + _y = self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3] + else: + # position that determines the sorting (including periodic shift of + # boundary conditions). Runs on the backend the markers live on; alpha is + # a 3-element weighting, not physics data, so it is converted to match. + alpha = xp.asarray(alpha, dtype=float) + assert alpha.size == 3 + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + _y = ( + alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) + + (1.0 - alpha) * self.markers[:, bi : bi + 3] + ) - # to stay on the current process, all three columns must be True. - # Reducing over a size-3 trailing axis of an array with many rows is slow - # This is faster - # Improvement of approximately 10x (both numpy and cupy) - self._can_stay = self._is_on_proc_domain[:, 0] & self._is_on_proc_domain[:, 1] & self._is_on_proc_domain[:, 2] + # y - floor(y), not xp.mod(y, 1.0): mathematically identical for a modulus of 1 + # (verified bit-for-bit equal), but xp.mod dispatches to a true floating-point + # remainder (fmod-like). On NumPy that fmod path is ~15x slower than plain + # floor+subtract (313k x 3: ~15ms vs ~1ms), making this single line the + # dominant cost of the whole function before this fix. On CuPy it is the other + # way around: cp.mod is one fused kernel launch, while floor+subtract is two, + # and kernel-launch overhead dominates at this array size (~0.07ms vs ~0.17ms) + # -- so this is backend-conditional rather than a universal fix. xp.floor + # (unlike xp.subtract/xp.mod) does not accept out= in the array-api-compat + # wrapper, so the NumPy path gets its own temporary for it. + if xp.cupy_backend: + xp.mod(_y, 1.0, out=self._sorting_etas) + else: + xp.subtract(_y, xp.floor(_y), out=self._sorting_etas) + + if xp.cupy_backend: + # Keep the established CuPy path unchanged for now; its kernel + # characteristics differ from NumPy's allocation costs. + self._is_on_proc_domain = xp.logical_and( + self._sorting_etas > self.domain_array_dev[self.mpi_rank, 0::3], + self._sorting_etas < self.domain_array_dev[self.mpi_rank, 1::3], + ) + self._can_stay[:] = ( + self._is_on_proc_domain[:, 0] & self._is_on_proc_domain[:, 1] & self._is_on_proc_domain[:, 2] + ) + else: + # Build only the one-dimensional result needed by the exchange + # path; retaining a temporary (n_markers, 3) boolean array adds + # another full marker-sized allocation and memory pass on NumPy. + bounds = self.domain_array_dev[self.mpi_rank] + eta = self._sorting_etas + self._can_stay[:] = ( + (eta[:, 0] > bounds[0]) + & (eta[:, 0] < bounds[1]) + & (eta[:, 1] > bounds[3]) + & (eta[:, 1] < bounds[4]) + & (eta[:, 2] > bounds[6]) + & (eta[:, 2] < bounds[7]) + ) - # holes and ghosts can stay, too - self._can_stay[self.holes] = True - self._can_stay[self.ghost_particles] = True + # holes and ghosts can stay, too. One merged boolean-mask assignment instead + # of two separate ones -- measured ~23% faster on CuPy at Np_local=12.5M + # (0.40ms -> 0.30ms): setting the same rows to True twice costs a second + # kernel launch + mask read for no benefit, since "True" is idempotent. + self._can_stay[self.holes | self.ghost_particles] = True # True values can stay on the process, False must be sent, already empty rows (-1) cannot be sent send_inds = xp.nonzero(~self._can_stay)[0] @@ -4386,6 +4884,7 @@ def _sendrecv_determine_mtbs( return hole_inds_after_send, send_inds + @ProfileManager.profile("_compute_neighbor_ranks") def _compute_neighbor_ranks(self) -> tuple[list[int], list[int]]: """Split every other rank into geometric neighbours of this rank's sub-domain box and everyone else, for :meth:`_sendrecv_get_destinations`. @@ -4436,6 +4935,7 @@ def _compute_neighbor_ranks(self) -> tuple[list[int], list[int]]: return neighbor_ranks, non_neighbor_ranks + @ProfileManager.profile("_sendrecv_get_destinations") def _sendrecv_get_destinations(self, send_inds): """ Determine to which process particles have to be sent. @@ -4451,12 +4951,16 @@ def _sendrecv_get_destinations(self, send_inds): """ # One entry for each process - send_info = xp.zeros(self.mpi_size, dtype=int) - - # Gathered once and reused for every rank below, instead of re-gathering - # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every - # iteration of the rank loop (as the previous version did). - candidates = self.markers[send_inds] + send_info = np.zeros(self.mpi_size, dtype=int) + + # etas_to_send is gathered once and reused for every rank below (cheap: 3 + # columns wide). The full marker payload is deliberately NOT gathered into a + # "candidates" intermediate here: that would gather all of send_inds' rows + # (every column) once, only to re-gather a subset of that copy per matched + # rank below -- two gather passes moving comparable total data. Composing the + # index instead (self.markers[send_inds[matched]] per rank, below) does the + # same total row selection in one gather pass per rank instead of two overall + # -- measured ~30% faster for this function on NumPy at Np=1e6/4 ranks. etas_to_send = self._sorting_etas[send_inds] # Reset every rank's send buffer to empty first. The neighbour/non-neighbour @@ -4466,7 +4970,7 @@ def _sendrecv_get_destinations(self, send_inds): # non-empty buffer from a previous call (send/recv size would then disagree # with send_info, which is always correct since it defaults to 0 above). empty_local = xp.empty(0, dtype=int) - empty_rows = candidates[:0] + empty_rows = self._markers[:0] for i in range(self.mpi_size): self._send_to_i[i] = empty_local self._send_list[i] = empty_rows @@ -4479,6 +4983,14 @@ def _sendrecv_get_destinations(self, send_inds): # everyone (a marker moved further than one sub-domain this step), the leftover # few are checked against every other rank in the second pass, so this changes # only how many ranks get checked in the common case, not correctness. + # A batched, per-pass variant of this loop (one broadcasted compare + one + # argsort per pass instead of one xp.nonzero per rank) was tried here to cut + # CuPy device syncs from O(n_group) to O(1) per pass. Measured net loss on + # BOTH backends at mpi_size=4 (NumPy: no sync cost to amortize against the + # extra broadcast memory traffic; CuPy: too few ranks per group -- 3 here -- + # for the eliminated syncs to outweigh the broadcast/argsort/searchsorted + # overhead). Might still win at much larger rank counts (more neighbours per + # group), but not re-added without measuring that regime first. remaining = xp.arange(send_inds.shape[0]) for rank_group in (self._neighbor_ranks, self._non_neighbor_ranks): if remaining.size == 0: @@ -4488,16 +5000,15 @@ def _sendrecv_get_destinations(self, send_inds): still_remaining = xp.ones(remaining.shape[0], dtype=bool) for i in rank_group: conds = xp.logical_and( - etas_remaining > self.domain_array[i, 0::3], - etas_remaining < self.domain_array[i, 1::3], + etas_remaining > self.domain_array_dev[i, 0::3], + etas_remaining < self.domain_array_dev[i, 1::3], ) - - matched_local = xp.nonzero(xp.all(conds, axis=1))[0] + matched_local = xp.nonzero(conds[:, 0] & conds[:, 1] & conds[:, 2])[0] matched = remaining[matched_local] self._send_to_i[i] = matched send_info[i] = matched.size - self._send_list[i] = candidates[matched] + self._send_list[i] = self._markers[send_inds[matched]] still_remaining[matched_local] = False @@ -4505,6 +5016,7 @@ def _sendrecv_get_destinations(self, send_inds): return send_info + @ProfileManager.profile("_sendrecv_all_to_all") def _sendrecv_all_to_all(self, send_info): """ Distribute info on how many markers will be sent/received to/from each process via all-to-all. @@ -4520,12 +5032,13 @@ def _sendrecv_all_to_all(self, send_info): Amount of marticles to be received from i-th process. """ - recv_info = xp.zeros(self.mpi_size, dtype=int) + recv_info = np.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) return recv_info + @ProfileManager.profile("_sendrecv_markers") def _sendrecv_markers(self, recv_info, hole_inds_after_send): """ Use non-blocking communication. In-place modification of markers @@ -4540,7 +5053,15 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): """ # i-th entry holds the number (not the index) of the first hole to be filled by data from process i - first_hole = xp.cumsum(recv_info) - recv_info + first_hole = np.cumsum(recv_info) - recv_info + + # Send requests are waited on below (unlike the previous version, which never + # stored or waited on them): otherwise nothing guarantees the send buffer + # (self._send_list[i]) is safe to overwrite before the underlying transfer has + # actually completed -- MPI is free to defer completing a large Isend until the + # receiver posts a matching receive (rendezvous protocol), so a fire-and-forget + # Isend is not necessarily done just because this call returns. + send_reqs = [] # Initialize send and receive commands for i, (data, N_recv) in enumerate(zip(self._send_list, list(recv_info))): @@ -4548,34 +5069,41 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): self._reqs[i] = None self._recvbufs[i] = None else: - self.mpi_comm.Isend(data, dest=i, tag=self.mpi_rank) - - self._recvbufs[i] = xp.zeros((N_recv, self.markers.shape[1]), dtype=float) + send_reqs.append(self.mpi_comm.Isend(data, dest=i, tag=self.mpi_rank)) + + # xp.empty, not xp.zeros: Irecv below overwrites the buffer completely + # (exactly N_recv rows, matching what the sender sent -- see the + # send/recv size cross-check in _sendrecv_get_destinations/_sendrecv_ + # all_to_all), so zero-initializing it first is wasted work. Under CuPy + # this is a device buffer, so mpi4py receives straight onto the GPU + # (CUDA-aware BTL/UCX) instead of into host memory -- see + # _sendrecv_get_destinations for the matching send side. + self._recvbufs[i] = xp.empty((N_recv, self.markers.shape[1]), dtype=float) self._reqs[i] = self.mpi_comm.Irecv(self._recvbufs[i], source=i, tag=i) - # Wait for buffer, then put markers into holes - test_reqs = [False] * (recv_info.size - 1) - while len(test_reqs) > 0: - # loop over all receive requests - for i, req in enumerate(self._reqs): - if req is None: - continue - else: - # check if data has been received - if req.Test(): - if hole_inds_after_send.size < first_hole[i] + recv_info[i]: - warnings.warn( - f'Strong load imbalance detected: \ + # Block on every receive at once via the MPI library's own wait, instead of a + # tight Python-level req.Test() poll loop -- lets the MPI progress engine (not + # our interpreter) do the waiting, and avoids repeatedly re-scanning still- + # pending requests from Python. + recv_ranks = [i for i, req in enumerate(self._reqs) if req is not None] + if recv_ranks: + MPI.Request.Waitall([self._reqs[i] for i in recv_ranks]) + + for i in recv_ranks: + if hole_inds_after_send.size < first_hole[i] + recv_info[i]: + warnings.warn( + f'Strong load imbalance detected: \ number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ is smaller than number of incoming particles ({first_hole[i] + recv_info[i]}). \ Increasing the value of "bufsize" in the markers parameters for the next run.', - ) - self.mpi_comm.Abort() + ) + self.mpi_comm.Abort() - self.markers[hole_inds_after_send[first_hole[i] + xp.arange(recv_info[i])]] = self._recvbufs[i] + self.markers[hole_inds_after_send[first_hole[i] + np.arange(recv_info[i])]] = self._recvbufs[i] + self._reqs[i] = None - test_reqs.pop() - self._reqs[i] = None + if send_reqs: + MPI.Request.Waitall(send_reqs) class Tesselation: @@ -4634,9 +5162,11 @@ def __init__( self._rank = comm.Get_rank() assert domain_array is not None + # tile/box geometry bookkeeping (small, ndim-sized), always host -- + # domain_array (when given) is already host, see _get_domain_decomp. if domain_array is None: - self._starts = xp.zeros(3) - self._ends = xp.ones(3) + self._starts = np.zeros(3) + self._ends = np.ones(3) else: self._starts = domain_array[self.rank, 0::3] self._ends = domain_array[self.rank, 1::3] @@ -4659,9 +5189,9 @@ def __init__( if n_boxes == 1: self._dims_mask = [True] * 3 else: - self._dims_mask = xp.array(self.boxes_per_dim) > 1 + self._dims_mask = np.array(self.boxes_per_dim) > 1 - min_tiles = 2 ** xp.count_nonzero(self.dims_mask) + min_tiles = 2 ** np.count_nonzero(self.dims_mask) assert self.tiles_pb >= min_tiles, ( f"At least {min_tiles} tiles per sorting box is enforced, but you have {self.tiles_pb}!" ) @@ -4689,20 +5219,25 @@ def get_tiles(self): # logger.info(f'{factors_vec = }') # logger.info(f'{self.dims_mask = }') - # tiles in one sorting box - self._nt_per_dim = xp.array([1, 1, 1]) - _ids = xp.nonzero(self._dims_mask)[0] + # tiles in one sorting box -- geometry bookkeeping, always host (see + # note in __init__); nt below must be a plain int for xp.linspace's + # num= argument. + self._nt_per_dim = np.array([1, 1, 1]) + _ids = np.nonzero(self._dims_mask)[0] for fac in factors_vec: _nt = self.nt_per_dim[self._dims_mask] - d = _ids[xp.argmin(_nt)] + d = _ids[np.argmin(_nt)] self._nt_per_dim[d] *= fac # logger.info(f'{_nt = }, {d = }, {self.nt_per_dim = }') - assert xp.prod(self.nt_per_dim) == self.tiles_pb + assert np.prod(self.nt_per_dim) == self.tiles_pb - # tiles between [0, box_width] in each direction - self._tile_breaks = [xp.linspace(0.0, bw, nt + 1) for bw, nt in zip(self.box_widths, self.nt_per_dim)] - self._tile_midpoints = [(xp.roll(tbs, -1)[:-1] + tbs[:-1]) / 2 for tbs in self.tile_breaks] + # tiles between [0, box_width] in each direction -- geometry, always + # host (like the rest of this method); the only legitimate device + # computation in this class is fun() in cell_averages, which already + # crosses via _dev()/_to_numpy_for_kernel at that one boundary. + self._tile_breaks = [np.linspace(0.0, bw, int(nt) + 1) for bw, nt in zip(self.box_widths, self.nt_per_dim)] + self._tile_midpoints = [(np.roll(tbs, -1)[:-1] + tbs[:-1]) / 2 for tbs in self.tile_breaks] self._tile_volume = 1.0 for tb in self.tile_breaks: self._tile_volume *= tb[1] @@ -4717,8 +5252,8 @@ def draw_markers(self): 1d arrays of logical-space marker coordinates, one entry per tile (length :attr:`n_tiles`).""" _, eta1 = self._tile_output_arrays() - eta2 = xp.zeros_like(eta1) - eta3 = xp.zeros_like(eta1) + eta2 = np.zeros_like(eta1) + eta3 = np.zeros_like(eta1) nt_x, nt_y, nt_z = self.nt_per_dim @@ -4729,7 +5264,7 @@ def draw_markers(self): for k in range(self.boxes_per_dim[2]): z_midpoints = self._get_midpoints(k, 2) - xx, yy, zz = xp.meshgrid( + xx, yy, zz = np.meshgrid( x_midpoints, y_midpoints, z_midpoints, @@ -4768,10 +5303,17 @@ def _get_quad_pts(self, n_quad=None): self._tile_quad_pts = [] self._tile_quad_wts = [] for nq, tb in zip(n_quad, self.tile_breaks): - pts_loc, wts_loc = xp.polynomial.legendre.leggauss(nq) + # cupy has no polynomial.legendre; this is tiny host-scale math + # (n_quad eigenvalues), and quadrature_grid below converts the + # result to the active backend via xp.asarray, which (unlike + # the reverse direction) is always safe. + pts_loc, wts_loc = np.polynomial.legendre.leggauss(nq) + # quadrature_grid always converts its output to the active + # backend (xp.asarray internally); convert straight back so the + # rest of this class stays host, like tile_breaks above. pts, wts = quadrature_grid(tb[:2], pts_loc, wts_loc) - self._tile_quad_pts += [pts[0]] - self._tile_quad_wts += [wts[0]] + self._tile_quad_pts += [_to_numpy_for_kernel(pts[0])] + self._tile_quad_wts += [_to_numpy_for_kernel(wts[0])] def cell_averages(self, fun, n_quad=None): """Compute the cell average of ``fun`` over every tile on the current process, @@ -4808,14 +5350,19 @@ def cell_averages(self, fun, n_quad=None): for k in range(self.boxes_per_dim[2]): z_pts = self._get_box_quad_pts(k, 2) - xx, yy, zz = xp.meshgrid( + # meshgrid stays host (x_pts/y_pts/z_pts are host, see + # _get_box_quad_pts); fun() is evaluated on the active + # backend via _dev(), the one legitimate device boundary + # in this class, and converted straight back for + # tile_int_kernel (a compiled, host-only Pyccel kernel). + xx, yy, zz = np.meshgrid( x_pts.flatten(), y_pts.flatten(), z_pts.flatten(), indexing="ij", ) - fun_vals = fun(xx, yy, zz) + fun_vals = _to_numpy_for_kernel(fun(*_dev(xx, yy, zz))) sampling_kernels.tile_int_kernel( fun_vals, @@ -4842,8 +5389,11 @@ def _tile_output_arrays(self): on the current process (i.e. the first array tiled over all sorting boxes). """ # self._quad_pts = [xp.zeros((nt, nq)).flatten() for nt, nq in zip(self.nt_per_dim, self.tile_quad_pts)] - single_box_out = xp.zeros(self.nt_per_dim) - out = xp.tile(single_box_out, self.boxes_per_dim) + # host, like the rest of this class (see get_tiles); feeds + # tile_int_kernel (cell_averages) and self._markers (draw_markers), + # both host-only. + single_box_out = np.zeros(self.nt_per_dim) + out = np.tile(single_box_out, self.boxes_per_dim) return single_box_out, out def _get_midpoints(self, i: int, dim: int): @@ -4886,7 +5436,7 @@ def _get_box_quad_pts(self, i: int, dim: int): xl = self.starts[dim] + i * self.box_widths[dim] x_tile_breaks = xl + self.tile_breaks[dim][:-1] x_tile_pts = self.tile_quad_pts[dim] - x_pts = xp.tile(x_tile_breaks, (x_tile_pts.size, 1)).T + x_tile_pts + x_pts = np.tile(x_tile_breaks, (x_tile_pts.size, 1)).T + x_tile_pts return x_pts @property diff --git a/src/struphy/pic/pushing/cuda/pusher_kernels_cuda/_general_geometry_src.cu b/src/struphy/pic/pushing/cuda/pusher_kernels_cuda/_general_geometry_src.cu new file mode 100644 index 000000000..aedd4e1be --- /dev/null +++ b/src/struphy/pic/pushing/cuda/pusher_kernels_cuda/_general_geometry_src.cu @@ -0,0 +1,373 @@ +#define MAXP 8 + +__device__ void matrix_inv_dev(const double* a, double* b) +{ + double det_a = a[0]*(a[4]*a[8] - a[5]*a[7]) + - a[1]*(a[3]*a[8] - a[5]*a[6]) + + a[2]*(a[3]*a[7] - a[4]*a[6]); + + b[0] = (a[4]*a[8] - a[7]*a[5]) / det_a; + b[1] = (a[7]*a[2] - a[1]*a[8]) / det_a; + b[2] = (a[1]*a[5] - a[4]*a[2]) / det_a; + b[3] = (a[5]*a[6] - a[8]*a[3]) / det_a; + b[4] = (a[8]*a[0] - a[2]*a[6]) / det_a; + b[5] = (a[2]*a[3] - a[5]*a[0]) / det_a; + b[6] = (a[3]*a[7] - a[6]*a[4]) / det_a; + b[7] = (a[6]*a[1] - a[0]*a[7]) / det_a; + b[8] = (a[0]*a[4] - a[3]*a[1]) / det_a; +} + +// c = a^T @ v (used for DF^-T @ e_form in push_v_with_efield_general below). +__device__ void matvecT_dev(const double* a, const double* v, double* out) +{ + out[0] = a[0]*v[0] + a[3]*v[1] + a[6]*v[2]; + out[1] = a[1]*v[0] + a[4]*v[1] + a[7]*v[2]; + out[2] = a[2]*v[0] + a[5]*v[1] + a[8]*v[2]; +} + +// df_out is row-major 3x3 (df_out[3*i+j] = dF_i/deta_j), matching +// struphy.geometry.mappings_kernels.cuboid_df / colella_df exactly. +__device__ void cuboid_df_dev(const double* params, double* df_out) +{ + // params = (l1, r1, l2, r2, l3, r3) + for (int k = 0; k < 9; k++) df_out[k] = 0.0; + df_out[0] = params[1] - params[0]; + df_out[4] = params[3] - params[2]; + df_out[8] = params[5] - params[4]; +} + +__device__ void colella_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (Lx, Ly, alpha, Lz) + const double lx = params[0], ly = params[1], alpha = params[2], lz = params[3]; + const double twopi = 6.283185307179586; + const double s1 = sin(twopi * eta1), c1 = cos(twopi * eta1); + const double s2 = sin(twopi * eta2), c2 = cos(twopi * eta2); + + df_out[0] = lx * (1.0 + alpha * c1 * s2 * twopi); + df_out[1] = lx * alpha * s1 * c2 * twopi; + df_out[2] = 0.0; + df_out[3] = ly * alpha * c1 * s2 * twopi; + df_out[4] = ly * (1.0 + alpha * s1 * c2 * twopi); + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} + +__device__ void orthogonal_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (Lx, Ly, alpha, Lz) + const double lx = params[0], ly = params[1], alpha = params[2], lz = params[3]; + const double twopi = 6.283185307179586; + + for (int k = 0; k < 9; k++) df_out[k] = 0.0; + df_out[0] = lx * (1.0 + alpha * cos(twopi * eta1) * twopi); + df_out[4] = ly * (1.0 + alpha * cos(twopi * eta2) * twopi); + df_out[8] = lz; +} + +__device__ void hollow_cyl_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (a1, a2, Lz, poc); faithful port of + // struphy.geometry.mappings_kernels.hollow_cyl_df, including its + // existing df_out[0,0]/df_out[1,0] not dividing eta2's argument by poc + // (unlike f_out and every other entry here) -- not "fixed" here, since + // this is a port, not a bugfix. + const double a1 = params[0], a2 = params[1], lz = params[2], poc = params[3]; + const double twopi = 6.283185307179586; + const double da = a2 - a1; + const double r = a1 + eta1 * da; + + df_out[0] = da * cos(twopi * eta2); + df_out[1] = -twopi / poc * r * sin(twopi * eta2 / poc); + df_out[2] = 0.0; + df_out[3] = da * sin(twopi * eta2); + df_out[4] = twopi / poc * r * cos(twopi * eta2 / poc); + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} + +__device__ void powered_ellipse_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (rx, ry, Lz, s) + const double rx = params[0], ry = params[1], lz = params[2], s = params[3]; + const double twopi = 6.283185307179586; + const double c2 = cos(twopi * eta2), s2 = sin(twopi * eta2); + const double e_sm1 = pow(eta1, s - 1.0); + const double e_s = pow(eta1, s); + + df_out[0] = e_sm1 * rx * c2; + df_out[1] = -twopi * e_s * rx * s2; + df_out[2] = 0.0; + df_out[3] = e_sm1 * ry * s2; + df_out[4] = twopi * e_s * ry * c2; + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} + +__device__ void hollow_torus_df_dev(double eta1, double eta2, double eta3, const double* params, double* df_out) +{ + // params = (a1, a2, R0, sfl, pol_period, tor_period) + const double a1 = params[0], a2 = params[1], r0 = params[2]; + const double sfl = params[3], pol_period = params[4], tor_period = params[5]; + const double pi = 3.14159265358979323846; + const double twopi = 6.283185307179586; + const double da = a2 - a1; + + if (sfl == 1.0) { + const double r = a1 + da * eta1; + const double eps = r / r0; + const double eps_p = da / r0; + const double tpe = tan(pi * eta2); + const double cpe = cos(pi * eta2); + const double tpe_p = pi / (cpe * cpe); + const double g = sqrt((1.0 + eps) / (1.0 - eps)); + const double g_p = 1.0 / (2.0 * g) * (eps_p * (1.0 - eps) + (1.0 + eps) * eps_p) / ((1.0 - eps) * (1.0 - eps)); + const double theta = 2.0 * atan(g * tpe); + const double denom = 1.0 + (g * tpe) * (g * tpe); + const double dtheta_deta1 = 2.0 / denom * g_p * tpe; + const double dtheta_deta2 = 2.0 / denom * g * tpe_p; + const double ct = cos(theta), st = sin(theta); + const double cf = cos(twopi * eta3 / tor_period), sf = sin(twopi * eta3 / tor_period); + + df_out[0] = (da * ct - r * st * dtheta_deta1) * cf; + df_out[1] = -r * st * dtheta_deta2 * cf; + df_out[2] = -twopi / tor_period * (r * ct + r0) * sf; + + df_out[3] = (da * ct - r * st * dtheta_deta1) * (-1.0) * sf; + df_out[4] = -r * st * dtheta_deta2 * (-1.0) * sf; + df_out[5] = twopi / tor_period * (r * ct + r0) * (-1.0) * cf; + + df_out[6] = da * st + r * ct * dtheta_deta1; + df_out[7] = r * ct * dtheta_deta2; + df_out[8] = 0.0; + } else { + const double r = a1 + eta1 * da; + const double cp = cos(twopi * eta2 / pol_period), sp = sin(twopi * eta2 / pol_period); + const double cf = cos(twopi * eta3 / tor_period), sf = sin(twopi * eta3 / tor_period); + + df_out[0] = da * cp * cf; + df_out[1] = -twopi / pol_period * r * sp * cf; + df_out[2] = -twopi / tor_period * (r * cp + r0) * sf; + + df_out[3] = da * cp * (-1.0) * sf; + df_out[4] = -twopi / pol_period * r * sp * (-1.0) * sf; + df_out[5] = (r * cp + r0) * (-1.0) * cf * twopi / tor_period; + + df_out[6] = da * sp; + df_out[7] = r * cp * twopi / pol_period; + df_out[8] = 0.0; + } +} + +__device__ void shafranov_shift_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (rx, ry, Lz, delta) + const double rx = params[0], ry = params[1], lz = params[2], de = params[3]; + const double twopi = 6.283185307179586; + const double c2 = cos(twopi * eta2), s2 = sin(twopi * eta2); + + df_out[0] = rx * c2 - 2.0 * eta1 * rx * de; + df_out[1] = -twopi * (eta1 * rx) * s2; + df_out[2] = 0.0; + df_out[3] = ry * s2; + df_out[4] = twopi * (eta1 * ry) * c2; + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} + +__device__ void shafranov_sqrt_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (rx, ry, Lz, delta) + const double rx = params[0], ry = params[1], lz = params[2], de = params[3]; + const double twopi = 6.283185307179586; + const double c2 = cos(twopi * eta2), s2 = sin(twopi * eta2); + + df_out[0] = rx * c2 - 0.5 / sqrt(eta1) * rx * de; + df_out[1] = -twopi * (eta1 * rx) * s2; + df_out[2] = 0.0; + df_out[3] = ry * s2; + df_out[4] = twopi * (eta1 * ry) * c2; + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} + +__device__ void shafranov_dshaped_df_dev(double eta1, double eta2, const double* params, double* df_out) +{ + // params = (R0, Lz, delta_x, delta_y, delta_gs, epsilon_gs, kappa_gs) + const double r0 = params[0], lz = params[1], dx = params[2], dy = params[3]; + const double dg = params[4], eg = params[5], kg = params[6]; + const double pi = 3.14159265358979323846; + const double twopi = 6.283185307179586; + const double asin_dg = asin(dg); + const double s2 = sin(twopi * eta2), c2 = cos(twopi * eta2); + const double phase = eta1 * s2 * asin_dg + twopi * eta2; + + df_out[0] = r0 * ( + -2.0 * dx * eta1 + - eg * eta1 * s2 * asin_dg * sin(phase) + + eg * cos(phase) + ); + df_out[1] = -r0 * eg * eta1 * (twopi * eta1 * c2 * asin_dg + twopi) * sin(phase); + df_out[2] = 0.0; + df_out[3] = r0 * (-2.0 * dy * eta1 + eg * kg * s2); + df_out[4] = twopi * r0 * eg * eta1 * kg * c2; + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} + +// Returns 1 if kind_map is supported and df_out was filled, 0 otherwise. +__device__ int df_dispatch_dev(int kind_map, double eta1, double eta2, double eta3, + const double* params, double* df_out) +{ + if (kind_map == 10) { cuboid_df_dev(params, df_out); return 1; } + if (kind_map == 11) { orthogonal_df_dev(eta1, eta2, params, df_out); return 1; } + if (kind_map == 12) { colella_df_dev(eta1, eta2, params, df_out); return 1; } + if (kind_map == 20) { hollow_cyl_df_dev(eta1, eta2, params, df_out); return 1; } + if (kind_map == 21) { powered_ellipse_df_dev(eta1, eta2, params, df_out); return 1; } + if (kind_map == 22) { hollow_torus_df_dev(eta1, eta2, eta3, params, df_out); return 1; } + if (kind_map == 30) { shafranov_shift_df_dev(eta1, eta2, params, df_out); return 1; } + if (kind_map == 31) { shafranov_sqrt_df_dev(eta1, eta2, params, df_out); return 1; } + if (kind_map == 32) { shafranov_dshaped_df_dev(eta1, eta2, params, df_out); return 1; } + return 0; +} + +__device__ int find_span_dev(const double* t, int p, int len_t, double eta) +{ + int low = p; + int high = len_t - 1 - p; + + if (eta <= t[low]) return low; + if (eta >= t[high]) return high - 1; + + int span = (low + high) / 2; + while (eta < t[span] || eta >= t[span + 1]) { + if (eta < t[span]) high = span; + else low = span; + span = (low + high) / 2; + } + return span; +} + +// Same as pusher_kernels_cuda.py's push_v_with_efield_cuboid's b_d_splines_dev, +// duplicated here because each cp.RawKernel source string is compiled +// independently (no cross-source linking). +__device__ void b_d_splines_dev(const double* t, int p, double eta, int span, double* bn, double* bd) +{ + double left[MAXP]; + double right[MAXP]; + int pd = p - 1; + + for (int i = 0; i <= p; i++) bn[i] = 0.0; + for (int i = 0; i < p; i++) bd[i] = 0.0; + bn[0] = 1.0; + + for (int j = 0; j < p; j++) { + left[j] = eta - t[span - j]; + right[j] = t[span + 1 + j] - eta; + double saved = 0.0; + + if (j == p - 1) { + for (int il = 0; il <= pd; il++) { + bd[pd - il] = (double)p / (t[span - il + p] - t[span - il]) * bn[pd - il]; + } + } + + for (int r = 0; r <= j; r++) { + double temp = bn[r] / (right[r] + left[j - r]); + bn[r] = saved + right[r] * temp; + saved = left[j - r] * temp; + } + bn[j + 1] = saved; + } +} + +extern "C" __global__ +void push_v_with_efield_general( + double* markers, + const int n_cols, + const int n_markers, + const int p1, const int p2, const int p3, + const double* tn1, const int len_tn1, + const double* tn2, const int len_tn2, + const double* tn3, const int len_tn3, + const int start0, const int start1, const int start2, + const double* e1_1, const int n2x1, const int n3x1, + const double* e1_2, const int n2x2, const int n3x2, + const double* e1_3, const int n2x3, const int n3x3, + const int kind_map, + const double* params, + const double dt_const) +{ + int ip = blockIdx.x * blockDim.x + threadIdx.x; + if (ip >= n_markers) return; + + double* row = markers + (size_t)ip * n_cols; + if (row[0] == -1.0 || row[n_cols - 1] == -2.0) return; + + const double eta1 = row[0], eta2 = row[1], eta3 = row[2]; + + double bn1[MAXP + 1], bd1[MAXP]; + double bn2[MAXP + 1], bd2[MAXP]; + double bn3[MAXP + 1], bd3[MAXP]; + + const int span1 = find_span_dev(tn1, p1, len_tn1, eta1); + const int span2 = find_span_dev(tn2, p2, len_tn2, eta2); + const int span3 = find_span_dev(tn3, p3, len_tn3, eta3); + + b_d_splines_dev(tn1, p1, eta1, span1, bn1, bd1); + b_d_splines_dev(tn2, p2, eta2, span2, bn2, bd2); + b_d_splines_dev(tn3, p3, eta3, span3, bn3, bd3); + + double e_form[3] = {0.0, 0.0, 0.0}; + for (int il1 = 0; il1 < p1; il1++) { + int i1 = span1 + il1 - start0; + for (int il2 = 0; il2 <= p2; il2++) { + int i2 = span2 + il2 - start1; + for (int il3 = 0; il3 <= p3; il3++) { + int i3 = span3 + il3 - start2; + e_form[0] += e1_1[(size_t)i1 * n2x1 * n3x1 + (size_t)i2 * n3x1 + i3] * bd1[il1] * bn2[il2] * bn3[il3]; + } + } + } + for (int il1 = 0; il1 <= p1; il1++) { + int i1 = span1 + il1 - start0; + for (int il2 = 0; il2 < p2; il2++) { + int i2 = span2 + il2 - start1; + for (int il3 = 0; il3 <= p3; il3++) { + int i3 = span3 + il3 - start2; + e_form[1] += e1_2[(size_t)i1 * n2x2 * n3x2 + (size_t)i2 * n3x2 + i3] * bn1[il1] * bd2[il2] * bn3[il3]; + } + } + } + for (int il1 = 0; il1 <= p1; il1++) { + int i1 = span1 + il1 - start0; + for (int il2 = 0; il2 <= p2; il2++) { + int i2 = span2 + il2 - start1; + for (int il3 = 0; il3 < p3; il3++) { + int i3 = span3 + il3 - start2; + e_form[2] += e1_3[(size_t)i1 * n2x3 * n3x3 + (size_t)i2 * n3x3 + i3] * bn1[il1] * bn2[il2] * bd3[il3]; + } + } + } + + double dfm[9], dfinv[9], dfinvT_e[3]; + df_dispatch_dev(kind_map, eta1, eta2, eta3, params, dfm); + matrix_inv_dev(dfm, dfinv); + matvecT_dev(dfinv, e_form, dfinvT_e); + + row[3] += dt_const * dfinvT_e[0]; + row[4] += dt_const * dfinvT_e[1]; + row[5] += dt_const * dfinvT_e[2]; +} diff --git a/src/struphy/pic/pushing/cuda/pusher_kernels_cuda/_push_v_efield_cuboid_src.cu b/src/struphy/pic/pushing/cuda/pusher_kernels_cuda/_push_v_efield_cuboid_src.cu new file mode 100644 index 000000000..33b40e175 --- /dev/null +++ b/src/struphy/pic/pushing/cuda/pusher_kernels_cuda/_push_v_efield_cuboid_src.cu @@ -0,0 +1,152 @@ +#define MAXP 8 + +__device__ int find_span_dev(const double* t, int p, int len_t, double eta) +{ + int low = p; + int high = len_t - 1 - p; + + if (eta <= t[low]) return low; + if (eta >= t[high]) return high - 1; + + int span = (low + high) / 2; + while (eta < t[span] || eta >= t[span + 1]) { + if (eta < t[span]) high = span; + else low = span; + span = (low + high) / 2; + } + return span; +} + +// Combined N-spline (bn, p+1 values) and D-spline (bd, p values) evaluation, +// matching struphy.bsplines.bsplines_kernels.b_d_splines_slim exactly. +__device__ void b_d_splines_dev(const double* t, int p, double eta, int span, double* bn, double* bd) +{ + double left[MAXP]; + double right[MAXP]; + int pd = p - 1; + + for (int i = 0; i <= p; i++) bn[i] = 0.0; + for (int i = 0; i < p; i++) bd[i] = 0.0; + bn[0] = 1.0; + + for (int j = 0; j < p; j++) { + left[j] = eta - t[span - j]; + right[j] = t[span + 1 + j] - eta; + double saved = 0.0; + + if (j == p - 1) { + for (int il = 0; il <= pd; il++) { + bd[pd - il] = (double)p / (t[span - il + p] - t[span - il]) * bn[pd - il]; + } + } + + for (int r = 0; r <= j; r++) { + double temp = bn[r] / (right[r] + left[j - r]); + bn[r] = saved + right[r] * temp; + saved = left[j - r] * temp; + } + bn[j + 1] = saved; + } +} + +extern "C" __global__ +void push_v_with_efield_cuboid( + double* markers, + const int n_cols, + const int n_markers, + const int p1, + const int p2, + const int p3, + const double* tn1, + const int len_tn1, + const double* tn2, + const int len_tn2, + const double* tn3, + const int len_tn3, + const int start0, + const int start1, + const int start2, + const double* e1_1, + const int n2x1, + const int n3x1, + const double* e1_2, + const int n2x2, + const int n3x2, + const double* e1_3, + const int n2x3, + const int n3x3, + const double sx, + const double sy, + const double sz, + const double dt_const) +{ + int ip = blockIdx.x * blockDim.x + threadIdx.x; + if (ip >= n_markers) return; + + double* row = markers + (size_t)ip * n_cols; + + // skip holes and ghost/boundary particles, matching Particles.valid_mks + if (row[0] == -1.0 || row[n_cols - 1] == -2.0) return; + + const double eta1 = row[0]; + const double eta2 = row[1]; + const double eta3 = row[2]; + + double bn1[MAXP + 1], bd1[MAXP]; + double bn2[MAXP + 1], bd2[MAXP]; + double bn3[MAXP + 1], bd3[MAXP]; + + const int span1 = find_span_dev(tn1, p1, len_tn1, eta1); + const int span2 = find_span_dev(tn2, p2, len_tn2, eta2); + const int span3 = find_span_dev(tn3, p3, len_tn3, eta3); + + b_d_splines_dev(tn1, p1, eta1, span1, bn1, bd1); + b_d_splines_dev(tn2, p2, eta2, span2, bn2, bd2); + b_d_splines_dev(tn3, p3, eta3, span3, bn3, bd3); + + // e_form[0]: D-spline in direction 1, N-splines in directions 2, 3 + double e_form0 = 0.0; + for (int il1 = 0; il1 < p1; il1++) { + int i1 = span1 + il1 - start0; + for (int il2 = 0; il2 <= p2; il2++) { + int i2 = span2 + il2 - start1; + for (int il3 = 0; il3 <= p3; il3++) { + int i3 = span3 + il3 - start2; + e_form0 += e1_1[(size_t)i1 * n2x1 * n3x1 + (size_t)i2 * n3x1 + i3] * bd1[il1] * bn2[il2] * bn3[il3]; + } + } + } + + // e_form[1]: N-spline in direction 1, D-spline in direction 2, N-spline in direction 3 + double e_form1 = 0.0; + for (int il1 = 0; il1 <= p1; il1++) { + int i1 = span1 + il1 - start0; + for (int il2 = 0; il2 < p2; il2++) { + int i2 = span2 + il2 - start1; + for (int il3 = 0; il3 <= p3; il3++) { + int i3 = span3 + il3 - start2; + e_form1 += e1_2[(size_t)i1 * n2x2 * n3x2 + (size_t)i2 * n3x2 + i3] * bn1[il1] * bd2[il2] * bn3[il3]; + } + } + } + + // e_form[2]: N-splines in directions 1, 2, D-spline in direction 3 + double e_form2 = 0.0; + for (int il1 = 0; il1 <= p1; il1++) { + int i1 = span1 + il1 - start0; + for (int il2 = 0; il2 <= p2; il2++) { + int i2 = span2 + il2 - start1; + for (int il3 = 0; il3 < p3; il3++) { + int i3 = span3 + il3 - start2; + e_form2 += e1_3[(size_t)i1 * n2x3 * n3x3 + (size_t)i2 * n3x3 + i3] * bn1[il1] * bn2[il2] * bd3[il3]; + } + } + } + + // Cartesian E-field is DF^-T @ e_form; for Cuboid, DF is diag(sx^-1, sy^-1, sz^-1) + // so DF^-T is diag(sx, sy, sz) -- same convention as push_eta_stage_cuboid's scale. + row[3] += dt_const * sx * e_form0; + row[4] += dt_const * sy * e_form1; + row[5] += dt_const * sz * e_form2; +} + diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index a224a5dcf..b77058467 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -2,7 +2,8 @@ import logging -import cunumpy as xp +import cunumpy +import numpy as np from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from line_profiler import profile @@ -10,6 +11,11 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles +from struphy.pic.pushing.pusher_kernels_cuda import ( + SUPPORTED_GENERAL_KIND_MAPS, + push_v_with_efield_cuboid_gpu, + push_v_with_efield_general_gpu, +) logger = logging.getLogger("struphy") @@ -138,7 +144,7 @@ def __init__( comps = ker_args[2] # check marker array column number - assert isinstance(comps, xp.ndarray) + assert isinstance(comps, np.ndarray) assert column_nr + comps.size < particles.n_cols, ( f"{column_nr + comps.size} not smaller than {particles.n_cols =}; not enough columns in marker array !!" ) @@ -150,7 +156,7 @@ def __init__( comps = ker_args[3] # check marker array column number - assert isinstance(comps, xp.ndarray) + assert isinstance(comps, np.ndarray) assert column_nr + comps.size < particles.n_cols, ( f"{column_nr + comps.size} not smaller than {particles.n_cols =}; not enough columns in marker array !!" ) @@ -162,7 +168,9 @@ def __init__( self._region_name = "pusher: " + self.kernel.name self._kernel_region_names = {} - self._residuals = xp.zeros(self.particles.markers.shape[0]) + # marker-row-indexed, so they live on the same backend as the markers + # (device under CuPy) -- see Particles._allocate_marker_array + self._residuals = cunumpy.zeros(self.particles.markers.shape[0]) self._converged_loc = self._residuals == 1.0 self._not_converged_loc = self._residuals == 0.0 @@ -171,6 +179,66 @@ def __init__( else: self._box_comm = False + # hand-written CUDA replacement for push_v_with_efield's per-marker + # math on a Cuboid domain. It only swaps out the inner kernel call + # (see the "push markers" branch in _push()), so it stays correct + # alongside unmodified apply_kinetic_bc/mpi_sort_markers/update_holes + # for multi-rank runs. + self._gpu_v_efield_cuboid = ( + cunumpy.cupy_backend and kernel.name == "push_v_with_efield" and args_domain.kind_map == 10 + ) + if self._gpu_v_efield_cuboid: + import cupy as cp + + l1, r1, l2, r2, l3, r3 = (float(p) for p in args_domain.params[:6]) + self._gpu_v_efield_scale = (1.0 / (r1 - l1), 1.0 / (r2 - l2), 1.0 / (r3 - l3)) + + args_derham, e1_1, e1_2, e1_3, const = args_kernel + self._gpu_v_efield_const = float(const) + self._gpu_v_efield_pn = tuple(int(p) for p in args_derham.pn) + self._gpu_v_efield_starts = tuple(int(s) for s in args_derham.starts) + # knot vectors are tiny host arrays; cache them on the device once + self._gpu_v_efield_tn1 = cp.asarray(args_derham.tn1, dtype=cp.float64) + self._gpu_v_efield_tn2 = cp.asarray(args_derham.tn2, dtype=cp.float64) + self._gpu_v_efield_tn3 = cp.asarray(args_derham.tn3, dtype=cp.float64) + # FE coefficients are already device-resident CuPy arrays under the + # CuPy backend (StencilVector allocates via cunumpy's xp) and are + # never reassigned after PushVinEfield.allocate() builds them, so + # these references stay valid and need no per-call transfer. + self._gpu_v_efield_e1_1 = e1_1 + self._gpu_v_efield_e1_2 = e1_2 + self._gpu_v_efield_e1_3 = e1_3 + + # general (non-Cuboid) CUDA replacement for push_v_with_efield: same + # B-spline evaluation as _gpu_v_efield_cuboid, but with DF(eta) + # evaluated per marker instead of assumed constant-diagonal. + self._gpu_v_efield_general = ( + cunumpy.cupy_backend + and kernel.name == "push_v_with_efield" + and not self._gpu_v_efield_cuboid + and args_domain.kind_map in SUPPORTED_GENERAL_KIND_MAPS + ) + if self._gpu_v_efield_general: + import cupy as cp + + self._gpu_v_efield_general_kind_map = int(args_domain.kind_map) + self._gpu_v_efield_general_params = cp.asarray( + np.asarray(args_domain.params, dtype=float), dtype=cp.float64 + ) + + args_derham, e1_1, e1_2, e1_3, const = args_kernel + self._gpu_v_efield_general_const = float(const) + self._gpu_v_efield_general_pn = tuple(int(p) for p in args_derham.pn) + self._gpu_v_efield_general_starts = tuple(int(s) for s in args_derham.starts) + self._gpu_v_efield_general_tn1 = cp.asarray(args_derham.tn1, dtype=cp.float64) + self._gpu_v_efield_general_tn2 = cp.asarray(args_derham.tn2, dtype=cp.float64) + self._gpu_v_efield_general_tn3 = cp.asarray(args_derham.tn3, dtype=cp.float64) + # FE coefficients are already device-resident under CuPy, see + # _gpu_v_efield_cuboid above. + self._gpu_v_efield_general_e1_1 = e1_1 + self._gpu_v_efield_general_e1_2 = e1_2 + self._gpu_v_efield_general_e1_3 = e1_3 + @profile def __call__(self, dt: float): """ @@ -188,6 +256,14 @@ def _kernel_region(self, kernel) -> str: self._kernel_region_names[id(kernel)] = name return name + def _run_marker_column_kernel(self, ker, alpha, column_nr, comps, add_args): + """Run one init/eval kernel (they write a marker column in place).""" + with ( + ProfileManager.profile_region(self._kernel_region(ker)), + self.particles.host_markers(write=True) as args_markers, + ): + ker(alpha, column_nr, comps, args_markers, self._args_domain, *add_args) + def _push(self, dt: float): """Body of :meth:`__call__`, see there.""" @@ -206,6 +282,8 @@ def _push(self, dt: float): init_slice = slice(first_pusher_idx, first_shift_idx) shift_slice = slice(first_shift_idx, residual_idx) + # Runs in place on whichever backend the markers live on -- device + # under CuPy, with no transfer (see Particles._allocate_marker_array). # save initial phase space coordinates markers[:, init_slice] = markers[:, : 3 + vdim] @@ -225,15 +303,13 @@ def _push(self, dt: float): comps = ker_args[2] add_args = ker_args[3] - with ProfileManager.profile_region(self._kernel_region(ker)): - ker( - xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), - column_nr, - comps, - self.particles.args_markers, - self._args_domain, - *add_args, - ) + self._run_marker_column_kernel( + ker, + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + column_nr, + comps, + add_args, + ) # update boxes if self._box_comm: @@ -242,7 +318,7 @@ def _push(self, dt: float): # start stages (e.g. n_stages=4 for RK4) for stage in range(self.n_stages): # start iteration (maxiter=1 for explicit schemes) - n_not_converged = xp.empty(1, dtype=int) + n_not_converged = np.empty(1, dtype=int) n_not_converged[0] = self.particles.n_mks_loc k = 0 @@ -273,15 +349,7 @@ def _push(self, dt: float): ) # evaluate - with ProfileManager.profile_region(self._kernel_region(ker)): - ker( - alpha, - column_nr, - comps, - self.particles.args_markers, - self._args_domain, - *add_args, - ) + self._run_marker_column_kernel(ker, alpha, column_nr, comps, add_args) # update boxes if self._box_comm: @@ -296,14 +364,53 @@ def _push(self, dt: float): ) # push markers - with ProfileManager.profile_region("kernel: " + self.kernel.name): - self.kernel( - dt, - stage, - self.particles.args_markers, - self._args_domain, - *self._args_kernel, - ) + if self._gpu_v_efield_cuboid: + with ProfileManager.profile_region("kernel: " + self.kernel.name + " [cuda]"): + push_v_with_efield_cuboid_gpu( + markers, + self.particles.n_cols, + self._gpu_v_efield_pn, + self._gpu_v_efield_tn1, + self._gpu_v_efield_tn2, + self._gpu_v_efield_tn3, + self._gpu_v_efield_starts, + self._gpu_v_efield_e1_1, + self._gpu_v_efield_e1_2, + self._gpu_v_efield_e1_3, + self._gpu_v_efield_scale, + dt * self._gpu_v_efield_const, + ) + elif self._gpu_v_efield_general: + with ProfileManager.profile_region("kernel: " + self.kernel.name + " [cuda general]"): + push_v_with_efield_general_gpu( + markers, + self.particles.n_cols, + self._gpu_v_efield_general_pn, + self._gpu_v_efield_general_tn1, + self._gpu_v_efield_general_tn2, + self._gpu_v_efield_general_tn3, + self._gpu_v_efield_general_starts, + self._gpu_v_efield_general_e1_1, + self._gpu_v_efield_general_e1_2, + self._gpu_v_efield_general_e1_3, + self._gpu_v_efield_general_kind_map, + self._gpu_v_efield_general_params, + dt * self._gpu_v_efield_general_const, + ) + else: + # no CUDA port for this kernel: fall back to the compiled + # host-only one, which pushes markers in place + with ( + ProfileManager.profile_region("kernel: " + self.kernel.name), + self.particles.host_markers(write=True) as args_markers, + ): + self.kernel( + dt, + stage, + args_markers, + self._args_domain, + *self._args_kernel, + ) self.particles.apply_kinetic_bc(newton=self._newton) self.particles.update_holes() @@ -315,13 +422,15 @@ def _push(self, dt: float): # compute number of non-converged particles (maxiter=1 for explicit schemes) if self.maxiter > 1: self._residuals[:] = markers[:, residual_idx] - max_res = xp.max(self._residuals) + max_res = float(cunumpy.max(self._residuals)) if max_res < 0.0: max_res = None self._converged_loc[:] = self._residuals < self._tol self._not_converged_loc[:] = ~self._converged_loc - n_not_converged[0] = xp.count_nonzero( - self._not_converged_loc, + # n_not_converged is a host buffer: it is passed straight + # into an mpi4py Allreduce below. + n_not_converged[0] = int( + cunumpy.count_nonzero(self._not_converged_loc), ) logger.debug( diff --git a/src/struphy/pic/pushing/pusher_kernels_cuda.py b/src/struphy/pic/pushing/pusher_kernels_cuda.py new file mode 100644 index 000000000..f1b4f2331 --- /dev/null +++ b/src/struphy/pic/pushing/pusher_kernels_cuda.py @@ -0,0 +1,215 @@ +"""Hand-written CUDA replacements for select pusher kernels, used only under +``ARRAY_BACKEND=cupy``. + +Unlike the generic Pyccel kernels in :mod:`~struphy.pic.pushing.pusher_kernels` +(which operate on plain host NumPy arrays regardless of backend), the kernels +here are real ``cupy.RawKernel`` CUDA source, executed directly on the GPU. +They are deliberately narrow: each one reproduces the exact arithmetic of one +Pyccel kernel, specialized for one :class:`~struphy.geometry.domains.Domain` +whose Jacobian is cheap enough that hand-specializing pays off. + +Currently covered: :func:`~struphy.pic.pushing.pusher_kernels.push_v_with_efield`, +for the Cuboid domain (:func:`push_v_with_efield_cuboid_gpu`) and, more +generally, for any domain in :data:`SUPPORTED_GENERAL_KIND_MAPS` +(:func:`push_v_with_efield_general_gpu`). This one does need a real +(small-degree) tensor-product B-spline evaluation -- the electric field is a +1-form FEEC spline, not a constant -- so these functions port ``find_span`` +and the combined N-/D-spline basis recursion +(:func:`~struphy.bsplines.bsplines_kernels.b_d_splines_slim`) to device code +alongside the local stencil sum +(:func:`~struphy.bsplines.evaluation_kernels_3d.eval_spline_mpi_kernel`). +Basis arrays are sized to a compile-time ``MAXP`` (spline degree 8), which +comfortably covers Struphy's usual degrees. The FE coefficient arrays +(``e1_1``, ``e1_2``, ``e1_3``) are the raw ``._data`` of the field's +:class:`~feectools.linalg.stencil.StencilVector` components; under the CuPy +backend these already live on the device (``StencilVector`` allocates via +``cunumpy``'s array-backend-aware ``xp``) and are never reassigned after +:meth:`~struphy.propagators.push_vin_efield.PushVinEfield.allocate` runs, so +they are passed straight through with no transfer at all -- only the marker +array round-trips through the device, exactly once per call. + +Every ``*_gpu`` function here is a thin wrapper around exactly one +:class:`~struphy.cuda.CudaKernel` (declared at module level, right above the +function that launches it): the kernel is compiled once, on first use, and +the function itself only builds the argument tuple and calls +:func:`~struphy.cuda.launch_1d`. If a function in this module does *not* sit +next to a ``CudaKernel``, it does not touch the GPU. +""" + +from struphy.cuda import CudaKernel, launch_1d, load_cuda_source + +_PUSH_V_EFIELD_CUBOID_SRC = load_cuda_source(__file__, "pusher_kernels_cuda/_push_v_efield_cuboid_src.cu") +_push_v_efield_cuboid_kernel = CudaKernel(_PUSH_V_EFIELD_CUBOID_SRC, "push_v_with_efield_cuboid") + + +def push_v_with_efield_cuboid_gpu( + markers, + n_cols: int, + pn: tuple[int, int, int], + tn1_dev, + tn2_dev, + tn3_dev, + starts: tuple[int, int, int], + e1_1_dev, + e1_2_dev, + e1_3_dev, + scale: tuple[float, float, float], + dt_const: float, +): + """GPU replacement for one call of + :func:`~struphy.pic.pushing.pusher_kernels.push_v_with_efield`, restricted + to the :class:`~struphy.geometry.domains.Cuboid` domain. + + ``markers`` is the host marker array and is round-tripped through the + device once. ``tn1_dev``, ``tn2_dev``, ``tn3_dev`` (knot vectors) and + ``e1_1_dev``, ``e1_2_dev``, ``e1_3_dev`` (FE coefficients of the 1-form + E-field) are expected to already be CuPy arrays resident on the device -- + callers should cache them once rather than converting on every call, see + :class:`~struphy.pic.pushing.pusher.Pusher`. + """ + import numpy as np + + n_markers = markers.shape[0] + launch_1d( + _push_v_efield_cuboid_kernel, + n_markers, + ( + markers, + np.int32(n_cols), + np.int32(n_markers), + np.int32(pn[0]), + np.int32(pn[1]), + np.int32(pn[2]), + tn1_dev, + np.int32(tn1_dev.shape[0]), + tn2_dev, + np.int32(tn2_dev.shape[0]), + tn3_dev, + np.int32(tn3_dev.shape[0]), + np.int32(starts[0]), + np.int32(starts[1]), + np.int32(starts[2]), + e1_1_dev, + np.int32(e1_1_dev.shape[1]), + np.int32(e1_1_dev.shape[2]), + e1_2_dev, + np.int32(e1_2_dev.shape[1]), + np.int32(e1_2_dev.shape[2]), + e1_3_dev, + np.int32(e1_3_dev.shape[1]), + np.int32(e1_3_dev.shape[2]), + np.float64(scale[0]), + np.float64(scale[1]), + np.float64(scale[2]), + np.float64(dt_const), + ), + ) + + +# ============================================================================ +# General (non-Cuboid-restricted) domain support +# ============================================================================ +# +# push_v_with_efield_cuboid_gpu above hardcodes Cuboid's Jacobian (a constant +# diagonal matrix, precomputed on the host as `scale`) directly into the +# marker update, which is what makes it fast but restricts it to +# `kind_map == 10`. Everything else about it -- the B-spline evaluation -- is +# already fully general (arbitrary degree, arbitrary non-uniform knot vector; +# nothing there assumes Cuboid). +# +# push_v_with_efield_general_gpu below drops the constant-Jacobian +# assumption: it evaluates DF(eta) (and its inverse) per marker, per call, on +# the device, matching the general struphy.geometry.evaluation_kernels.df / +# struphy.linear_algebra.linalg_kernels dispatch that +# struphy.pic.pushing.pusher_kernels.push_v_with_efield uses on the CPU. This +# is genuinely more per-marker work (a Jacobian evaluation instead of a +# lookup), but still embarrassingly parallel across markers, so it remains a +# good GPU fit. +# +# All analytic (closed-form) mappings in struphy.geometry.mappings_kernels +# are implemented: Cuboid (10), Orthogonal (11), Colella (12), +# HollowCylinder (20), PoweredEllipticCylinder (21), HollowTorus (22, +# including both its straight-field-line and equal-angle branches), +# ShafranovShiftCylinder (30), ShafranovSqrtCylinder (31) and +# ShafranovDshapedCylinder (32) -- see SUPPORTED_GENERAL_KIND_MAPS. +# +# NOT implemented: kind_map 0/1/2 (spline_3d / spline_2d_straight / +# spline_2d_torus), where the domain mapping F itself is an IGA B-spline +# volume (control points args.cx/cy/cz) rather than a closed-form function -- +# evaluating DF there means differentiating that spline (basis_funs_1st_der / +# a derivative-spline evaluation, not just the tensor-product sum this file +# already has for FEEC fields), which is a separate, larger piece of work. +# Callers must check kind_map themselves (see Pusher._gpu_v_efield_general in +# pusher.py) and fall back to the host Pyccel kernel for anything else -- +# this function does not raise on an unsupported kind_map, it is simply not +# wired up for one. + +_GENERAL_GEOMETRY_SRC = load_cuda_source(__file__, "pusher_kernels_cuda/_general_geometry_src.cu") + +_push_v_efield_general_kernel = CudaKernel(_GENERAL_GEOMETRY_SRC, "push_v_with_efield_general") + +#: kind_map values df_dispatch_dev supports. Callers should check membership +#: before dispatching to push_v_with_efield_general_gpu. +SUPPORTED_GENERAL_KIND_MAPS = (10, 11, 12, 20, 21, 22, 30, 31, 32) + + +def push_v_with_efield_general_gpu( + markers, + n_cols: int, + pn: tuple[int, int, int], + tn1_dev, + tn2_dev, + tn3_dev, + starts: tuple[int, int, int], + e1_1_dev, + e1_2_dev, + e1_3_dev, + kind_map: int, + params_dev, + dt_const: float, +): + """GPU replacement for one call of + :func:`~struphy.pic.pushing.pusher_kernels.push_v_with_efield`, for any + domain in :data:`SUPPORTED_GENERAL_KIND_MAPS`. See + :func:`push_v_with_efield_cuboid_gpu` for the argument conventions + (``tn*_dev``/``e1_*_dev`` are expected to already be device-resident); + ``params_dev`` is the domain's mapping-parameter array + (``args_domain.params``), expected to already be a small CuPy array + (cheap to keep device-resident; callers should cache it once). + """ + import numpy as np + + n_markers = markers.shape[0] + launch_1d( + _push_v_efield_general_kernel, + n_markers, + ( + markers, + np.int32(n_cols), + np.int32(n_markers), + np.int32(pn[0]), + np.int32(pn[1]), + np.int32(pn[2]), + tn1_dev, + np.int32(tn1_dev.shape[0]), + tn2_dev, + np.int32(tn2_dev.shape[0]), + tn3_dev, + np.int32(tn3_dev.shape[0]), + np.int32(starts[0]), + np.int32(starts[1]), + np.int32(starts[2]), + e1_1_dev, + np.int32(e1_1_dev.shape[1]), + np.int32(e1_1_dev.shape[2]), + e1_2_dev, + np.int32(e1_2_dev.shape[1]), + np.int32(e1_2_dev.shape[2]), + e1_3_dev, + np.int32(e1_3_dev.shape[1]), + np.int32(e1_3_dev.shape[2]), + np.int32(kind_map), + params_dev, + np.float64(dt_const), + ), + ) diff --git a/src/struphy/pic/sorting.py b/src/struphy/pic/sorting.py index 940264130..f99978cc1 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -1,4 +1,5 @@ import logging +import math try: from mpi4py.MPI import Intracomm @@ -9,9 +10,15 @@ class Intracomm: import cunumpy as xp +from cunumpy import PyccelKernel from struphy.pic.sorting_kernels import flatten_index, initialize_neighbours +# initialize_neighbours writes into an array that may be CuPy-resident under +# the active backend; flatten_index only ever takes plain ints, so it does +# not need wrapping. +initialize_neighbours = PyccelKernel(initialize_neighbours) + logger = logging.getLogger("struphy") @@ -230,7 +237,7 @@ def _set_boxes(self): n_particles = self._markers_shape[0] n_mkr = int(n_particles / n_box_in) + 1 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) diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index 8a1ec6122..51e5b8387 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -1,4 +1,5 @@ import logging +import math import pytest from cunumpy import PyccelKernel @@ -47,6 +48,7 @@ ], ) @pytest.mark.parametrize("num_clones", [1, 2]) +@pytest.mark.needs_host_kernels def test_accum_poisson(num_elements, degree, bcs, mapping, num_clones, Np=10000, show_plot: bool = False): r"""Test that AccumulatorVector provides an MC approximation of the L2 projection RHS. @@ -115,7 +117,7 @@ def test_accum_poisson(num_elements, degree, bcs, mapping, num_clones, Np=10000, params = { "grid": {"num_elements": num_elements}, - "kinetic": {"test_particles": {"markers": {"Np": Np, "ppc": Np / xp.prod(num_elements)}}}, + "kinetic": {"test_particles": {"markers": {"Np": Np, "ppc": Np / math.prod(num_elements)}}}, } grid = TensorProductGrid(num_elements=num_elements) @@ -328,6 +330,7 @@ def test_accum_poisson(num_elements, degree, bcs, mapping, num_clones, Np=10000, (None, None, None), ], ) +@pytest.mark.needs_host_kernels def test_accum_div_u_weak_1form(num_elements, degree, bcs, Np=10000, show_plot: bool = False): r"""Test that AccumulatorVector with kernel :func:`~struphy.pic.accumulation.accum_kernels.div_u_weak_1form` provides an MC approximation of the L2 projection RHS into V1 (Hcurl). diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 3751a3e7f..b0426488c 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -115,10 +115,12 @@ def test_draw(num_elements, degree, bcs, mapping, ppc=10): logger.info("Number of particles w/wo holes on each process after sorting : ") logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.markers.shape[0]}") - # are all markers in the correct domain? + # are all markers in the correct domain? domain_array is host-resident + # (fixed decomposition metadata); markers may live on the device under + # CuPy, so bring the (tiny) domain bounds to the same backend to compare. conds = xp.logical_and( - particles.markers[:, :3] > derham.domain_array[rank, 0::3], - particles.markers[:, :3] < derham.domain_array[rank, 1::3], + particles.markers[:, :3] > xp.asarray(derham.domain_array[rank, 0::3]), + particles.markers[:, :3] < xp.asarray(derham.domain_array[rank, 1::3]), ) holes = particles.markers[:, 0] == -1.0 stay = xp.all(conds, axis=1) diff --git a/src/struphy/pic/tests/test_mat_vec_filler.py b/src/struphy/pic/tests/test_mat_vec_filler.py index e0bdf4026..9ab3bbe33 100644 --- a/src/struphy/pic/tests/test_mat_vec_filler.py +++ b/src/struphy/pic/tests/test_mat_vec_filler.py @@ -1,6 +1,7 @@ import logging import cunumpy as xp +import numpy as np import pytest logger = logging.getLogger("struphy") @@ -16,6 +17,7 @@ (None, ("free", "free"), ("free", "free")), ], ) +@pytest.mark.needs_host_kernels def test_particle_to_mat_kernels(num_elements, degree, bcs, n_markers=1): """This test assumes a single particle and verifies a) if the correct indices are non-zero in _data @@ -47,12 +49,14 @@ def test_particle_to_mat_kernels(num_elements, degree, bcs, n_markers=1): logger.info(f"\nnum_elements={num_elements}, degree={degree}, bcs={bcs}\n") # DR attributes - pn = xp.array(DR.degree) + # plain int metadata (polynomial degrees) -- kept host-side so downstream + # index/span arithmetic doesn't leak CuPy scalars into arange() etc. + pn = np.array(DR.degree) tn1, tn2, tn3 = DR.V0fem.knots starts1 = {} - starts1["v0"] = xp.array(DR.V0.starts) + starts1["v0"] = np.array(DR.V0.starts) comm.Barrier() sleep(0.02 * (rank + 1)) @@ -124,6 +128,9 @@ def test_particle_to_mat_kernels(num_elements, degree, bcs, n_markers=1): eta3s = xp.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] for eta1, eta2, eta3 in zip(eta1s, eta2s, eta3s): + # the compiled bsplines_kernels functions below require native + # Python floats; eta1s/eta2s/eta3s may be CuPy-resident. + eta1, eta2, eta3 = float(eta1), float(eta2), float(eta3) comm.Barrier() sleep(0.02 * (rank + 1)) logger.info(f"rank {rank} | eta1 = {eta1}") @@ -137,14 +144,15 @@ def test_particle_to_mat_kernels(num_elements, degree, bcs, n_markers=1): span2 = bsp.find_span(tn2, DR.degree[1], eta2) span3 = bsp.find_span(tn3, DR.degree[2], eta3) - # non-zero spline values at eta - bn1 = xp.empty(DR.degree[0] + 1, dtype=float) - bn2 = xp.empty(DR.degree[1] + 1, dtype=float) - bn3 = xp.empty(DR.degree[2] + 1, dtype=float) + # non-zero spline values at eta -- output buffers for the compiled + # bsplines_kernels functions, which require host numpy arrays. + bn1 = np.empty(DR.degree[0] + 1, dtype=float) + bn2 = np.empty(DR.degree[1] + 1, dtype=float) + bn3 = np.empty(DR.degree[2] + 1, dtype=float) - bd1 = xp.empty(DR.degree[0], dtype=float) - bd2 = xp.empty(DR.degree[1], dtype=float) - bd3 = xp.empty(DR.degree[2], dtype=float) + bd1 = np.empty(DR.degree[0], dtype=float) + bd2 = np.empty(DR.degree[1], dtype=float) + bd3 = np.empty(DR.degree[2], dtype=float) bsp.b_d_splines_slim(tn1, DR.degree[0], eta1, span1, bn1, bd1) bsp.b_d_splines_slim(tn2, DR.degree[1], eta2, span2, bn2, bd2) @@ -155,10 +163,11 @@ def test_particle_to_mat_kernels(num_elements, degree, bcs, n_markers=1): ie2 = span2 - pn[1] ie3 = span3 - pn[2] - # global indices of non-vanishing B- and D-splines (no modulo) - glob_n1 = xp.arange(ie1, ie1 + pn[0] + 1) - glob_n2 = xp.arange(ie2, ie2 + pn[1] + 1) - glob_n3 = xp.arange(ie3, ie3 + pn[2] + 1) + # global indices of non-vanishing B- and D-splines (no modulo) -- pure + # host-side index bookkeeping used below for Python set comparisons. + glob_n1 = np.arange(ie1, ie1 + pn[0] + 1) + glob_n2 = np.arange(ie2, ie2 + pn[1] + 1) + glob_n3 = np.arange(ie3, ie3 + pn[2] + 1) glob_d1 = glob_n1[:-1] glob_d2 = glob_n2[:-1] @@ -184,10 +193,10 @@ def test_particle_to_mat_kernels(num_elements, degree, bcs, n_markers=1): # local column indices in _data of non-vanishing B- and D-splines, as sets for comparison cols = [{}, {}, {}] for n in range(3): - cols[n]["NN"] = set(xp.arange(2 * pn[n] + 1)) - cols[n]["ND"] = set(xp.arange(2 * pn[n])) - cols[n]["DN"] = set(xp.arange(1, 2 * pn[n] + 1)) - cols[n]["DD"] = set(xp.arange(1, 2 * pn[n])) + cols[n]["NN"] = set(np.arange(2 * pn[n] + 1)) + cols[n]["ND"] = set(np.arange(2 * pn[n])) + cols[n]["DN"] = set(np.arange(1, 2 * pn[n] + 1)) + cols[n]["DD"] = set(np.arange(1, 2 * pn[n])) # testing vector-valued spaces spaces_vector = ["v1", "v2"] diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index fb139de89..27a5b92d0 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -627,6 +627,7 @@ def test_push_bxu_Hdiv_pauli(num_elements, degree, bcs, mapping, show_plots=Fals ) def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): import cunumpy as xp + import numpy as np from feectools.ddm.mpi import mpi as MPI from struphy import BoundaryParameters, LoadingParameters, WeightsParameters, domains @@ -699,12 +700,14 @@ def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): pusher_psy(dt) - n_mks_load = xp.zeros(size, dtype=int) + # MPI communication buffers must be host-resident regardless of the + # active array backend (mpi4py has no CuPy awareness here). + n_mks_load = np.zeros(size, dtype=int) - comm.Allgather(xp.array(xp.shape(particles.markers)[0]), n_mks_load) + comm.Allgather(np.array(xp.shape(particles.markers)[0]), n_mks_load) - sendcounts = xp.zeros(size, dtype=int) - displacements = xp.zeros(size, dtype=int) + sendcounts = np.zeros(size, dtype=int) + displacements = np.zeros(size, dtype=int) accum_sendcounts = 0.0 for i in range(size): @@ -712,10 +715,13 @@ def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): displacements[i] = accum_sendcounts accum_sendcounts += sendcounts[i] - all_particles_psy = xp.zeros((int(accum_sendcounts) * 3,), dtype=float) + all_particles_psy = np.zeros((int(accum_sendcounts) * 3,), dtype=float) comm.Barrier() - comm.Allgatherv(xp.array(particles.markers[:, :3]), [all_particles_psy, sendcounts, displacements, MPI.DOUBLE]) + comm.Allgatherv( + np.ascontiguousarray(xp.to_numpy(particles.markers[:, :3])), + [all_particles_psy, sendcounts, displacements, MPI.DOUBLE], + ) comm.Barrier() diff --git a/src/struphy/pic/tests/test_set_zero_velocity.py b/src/struphy/pic/tests/test_set_zero_velocity.py index 2eb563859..729222954 100644 --- a/src/struphy/pic/tests/test_set_zero_velocity.py +++ b/src/struphy/pic/tests/test_set_zero_velocity.py @@ -130,6 +130,7 @@ def test_set_zero_velocity_mpi(mapping, comp: int, show_plot=False): """ import cunumpy as xp + import numpy as np from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI from matplotlib import pyplot as plt @@ -181,10 +182,13 @@ def test_set_zero_velocity_mpi(mapping, comp: int, show_plot=False): if comm is None: mpi_result = binned_result else: - mpi_result = xp.zeros_like(binned_result) + # mpi4py needs host buffers regardless of the active backend. + host_binned = [xp.to_numpy(b) for b in binned_result] + host_result = [np.zeros_like(b) for b in host_binned] for i in range(3): - comm.Allreduce(binned_result[i], mpi_result[i], op=MPI.SUM) + comm.Allreduce(host_binned[i], host_result[i], op=MPI.SUM) comm.Barrier() + mpi_result = [xp.asarray(r) for r in host_result] # tests if show_plot and rank == 0: diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index d536004b7..4bac61630 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -500,10 +500,10 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela logger.info(f"{Np =}, {ppb =}, {diff =}") if tesselation: - fit = xp.polyfit(xp.log(ppbs), xp.log(err_vec), 1) + fit = xp.polyfit(xp.log(xp.array(ppbs)), xp.log(xp.array(err_vec)), 1) xvec = ppbs else: - fit = xp.polyfit(xp.log(Nps), xp.log(err_vec), 1) + fit = xp.polyfit(xp.log(xp.array(Nps)), xp.log(xp.array(err_vec)), 1) xvec = Nps if show_plot and rank == 0: @@ -623,9 +623,9 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat err_vec += [diff] if tesselation: - fit = xp.polyfit(xp.log(h_vec[1:5]), xp.log(err_vec[1:5]), 1) + fit = xp.polyfit(xp.log(xp.array(h_vec[1:5])), xp.log(xp.array(err_vec[1:5])), 1) else: - fit = xp.polyfit(xp.log(h_vec[:-2]), xp.log(err_vec[:-2]), 1) + fit = xp.polyfit(xp.log(xp.array(h_vec[:-2])), xp.log(xp.array(err_vec[:-2])), 1) if show_plot and rank == 0: plt.figure(figsize=(12, 8)) @@ -908,10 +908,10 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation # fig.savefig(f"2d_sph_{Np}_{ppb}.png") if tesselation: - fit = xp.polyfit(xp.log(ppbs), xp.log(err_vec), 1) + fit = xp.polyfit(xp.log(xp.array(ppbs)), xp.log(xp.array(err_vec)), 1) xvec = ppbs else: - fit = xp.polyfit(xp.log(Nps), xp.log(err_vec), 1) + fit = xp.polyfit(xp.log(xp.array(Nps)), xp.log(xp.array(err_vec)), 1) xvec = Nps if show_plot and rank == 0: diff --git a/src/struphy/propagators/push_eta.py b/src/struphy/propagators/push_eta.py index de8bfe595..9f9a4c7a2 100644 --- a/src/struphy/propagators/push_eta.py +++ b/src/struphy/propagators/push_eta.py @@ -91,8 +91,13 @@ def options(self, new): @profile def allocate(self): # get kernel + + # Old kernel = PyccelKernel(pusher_kernels.push_eta_stage) + # New (returns a PyccelKernel) + kernel = pusher_kernels.push_eta_stage.kernel + # define algorithm butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: diff --git a/src/struphy/propagators/tests/test_curl_curl.py b/src/struphy/propagators/tests/test_curl_curl.py index b3205cdcc..96bbcfbb9 100644 --- a/src/struphy/propagators/tests/test_curl_curl.py +++ b/src/struphy/propagators/tests/test_curl_curl.py @@ -28,6 +28,10 @@ logger = logging.getLogger("struphy") set_logging_level(logging.INFO) +# curl-curl FEEC solve: mass-matrix assembly/preconditioning not yet ported +# to the CuPy backend. +pytestmark = pytest.mark.needs_host_kernels + comm = MPI.COMM_WORLD rank = comm.Get_rank() plt.rcParams.update({"font.size": 22}) @@ -57,9 +61,11 @@ def test_convergence_1d( # Test over spline degree and grid resolution Nels = [2**n for n in range(Nmin, Nmax + 1)] - e1 = 0.0 - e2 = 0.0 - e3 = 0.0 + # xp.meshgrid (unlike numpy's) requires actual arrays, not plain floats, + # for whichever of e1/e2/e3 isn't replaced by e below. + e1 = xp.array([0.0]) + e2 = xp.array([0.0]) + e3 = xp.array([0.0]) e = xp.linspace(0.0, 1.0, 64) bcs = (None, None, None) @@ -309,7 +315,9 @@ def scalar_current(a, b): Nels = [2**n for n in range(Nmin, Nmax + 1)] e = xp.linspace(0.0, 1.0, 64) - egrid = [0.0, 0.0, 0.0] + # xp.meshgrid (unlike numpy's) requires actual arrays, not plain floats, + # for whichever entries aren't replaced by e below. + egrid = [xp.array([0.0]), xp.array([0.0]), xp.array([0.0])] for idx in space["coords"]: egrid[idx] = e e1, e2, e3 = egrid diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index a67ec2406..d78705be0 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -20,6 +20,10 @@ logger = logging.getLogger("struphy") set_logging_level(logging.INFO) +# Gyrokinetic Poisson FEEC solve: mass-matrix assembly/preconditioning not +# yet ported to the CuPy backend. +pytestmark = pytest.mark.needs_host_kernels + comm = MPI.COMM_WORLD rank = comm.Get_rank() # plt.rcParams.update({'font.size': 22}) diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 55ca649fd..15704eda9 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -30,6 +30,10 @@ logger = logging.getLogger("struphy") +# Poisson FEEC solve: mass-matrix assembly/preconditioning not yet ported to +# the CuPy backend. +pytestmark = pytest.mark.needs_host_kernels + comm = MPI.COMM_WORLD rank = comm.Get_rank() plt.rcParams.update({"font.size": 22})