diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 2430ee09c..0e30c2d5e 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1236,7 +1236,10 @@ def draw_markers( # counting integers num_loaded_particles_loc = 0 # number of particles alreday loaded (local) num_loaded_particles_glob = 0 # number of particles already loaded (each clone) - chunk_size = 10000 # TODO: number of particle chunk + # markers drawn per chunk: large chunks keep the number of iterations (each a few small kernels and + # a device sync on the GPU) low, e.g. 16 instead of 1600 for 16M markers; a chunk of temporaries + # takes ~(3 + vdim + 4) * 8 bytes per marker (~50 MB). The NumPy draws do not depend on it. + chunk_size = 2**20 # Total number of markers to draw (sum over all clones) while num_loaded_particles_glob < int(self.Np): @@ -1552,6 +1555,12 @@ def initialize_weights( else: self.weights = self.weights0 + def _f_coords_of_valid(self, quantity, rows): + """Coordinates of the distribution function at the valid markers, as a list of arrays. + ``quantity(j)`` gives row ``j`` of the marker buffer; subclasses whose :attr:`f_coords` differs + from the stored rows override this.""" + return [quantity(j) for j in rows] + @profile def update_weights(self): """ @@ -1561,8 +1570,16 @@ def update_weights(self): """ from struphy.pic.particles import ParticlesSPH - if self.n_mks_loc == 0: + # the indices of the valid markers, found once: each property (positions, weights0, ...) would select them + # again with the boolean mask (and n_mks_loc counts them with one more pass and a host read) + idx = xp.flatnonzero(self.valid_mks) + if idx.size == 0: return + markers = self._markers + + def quantity(j): + """Row j of the marker buffer (one quantity) at the valid markers.""" + return xp.take(markers[j], idx) if isinstance(self, ParticlesSPH): f0 = self.f0.n0(self.positions) @@ -1570,7 +1587,9 @@ def update_weights(self): # in case of CanonicalMaxwellian, evaluate constants_of_motion # if isinstance(self.f0, CanonicalMaxwellian): # self.save_constants_of_motion() - f0 = self.f0(*self.f_coords) + coords = self.f_coords_index + rows = range(markers.shape[0])[coords] if isinstance(coords, slice) else np.atleast_1d(coords).tolist() + f0 = self.f0(*self._f_coords_of_valid(quantity, rows)) # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: @@ -1579,7 +1598,7 @@ def update_weights(self): if self.is_volume_form[1]: f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords) - self.weights = self.weights0 - f0 / self.sampling_density_values / self.Np + markers[self.index["weights"], idx] = quantity(self.index["w0"]) - f0 / quantity(self.index["s0"]) / self.Np @profile def binning( @@ -1823,10 +1842,25 @@ def mpi_sort_markers( (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ + # Ghost particles (-2 in the last row) only exist with sorting boxes (SPH ghost boxes, + # _prepare_ghost_particles); without them the ghost mask stays all False and its + # bookkeeping (full passes over the buffer, a nonzero) can be skipped. + ghosts_possible = self._initialized_sorting + + # On the GPU the markers go through cunumpy.mpi.exchange, which waits for every send and + # receive (_sendrecv_markers_device): no barriers are needed around the exchange. The host path + # does not wait for its sends, so it keeps them. + barriers = xp.get_array_backend(self._markers) != "cupy" + if remove_ghost: - self._remove_ghost_particles() + if ghosts_possible: + self._remove_ghost_particles() + else: + # what _remove_ghost_particles leaves when there are no ghosts: refreshed holes + self.update_holes() - self._Barrier() + if barriers: + self._Barrier() # before sorting, apply kinetic bc if apply_bc: @@ -1835,24 +1869,29 @@ def mpi_sort_markers( if isinstance(alpha, int) or isinstance(alpha, float): alpha = (alpha, alpha, alpha) - # create new markers_to_be_sent array and make corresponding holes in markers array - hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha) + if self._use_device_sort(): + # destinations in one kernel pass, one gather per destination rank (see _mpi_sort_device) + n_emptied, recv_info = self._mpi_sort_device(alpha) + else: + # create new markers_to_be_sent array and make corresponding holes in markers array + hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha) - # determine where to send markers_to_be_sent - send_info = self._sendrecv_get_destinations(send_inds) + # determine where to send markers_to_be_sent + send_info = self._sendrecv_get_destinations(send_inds) - # set new holes in markers array to -1 - self._markers[:, send_inds] = -1.0 + # set new holes in markers array to -1 + self._markers[:, send_inds] = -1.0 - # transpose send_info - recv_info = self._sendrecv_all_to_all(send_info) + # transpose send_info + recv_info = self._sendrecv_all_to_all(send_info) - # send and receive markers - self._sendrecv_markers(recv_info, hole_inds_after_send) + # send and receive markers + self._sendrecv_markers(recv_info, hole_inds_after_send) + n_emptied = int(send_inds.size) # the sent markers are holes now, the received ones filled the lowest holes (host counts, no device read) n_received = int(np.sum(recv_info)) - int(recv_info[self.mpi_rank]) - self._track_markers(n_emptied=int(send_inds.size), n_filled=n_received) + self._track_markers(n_emptied=n_emptied, n_filled=n_received) # new holes and new number of holes and markers on process self.update_holes() @@ -1860,7 +1899,8 @@ def mpi_sort_markers( # refresh ghost mask: received markers may land in columns 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. - self._update_ghost_particles() + if ghosts_possible: + self._update_ghost_particles() # check if all markers are on the right process after sorting if do_test: @@ -1874,7 +1914,8 @@ def mpi_sort_markers( assert all_on_right_proc # assert self.phasespace_coords.size > 0, f'No particles on process {self.mpi_rank}, please rebalance, aborting ...' - self._Barrier() + if barriers: + self._Barrier() @profile @ProfileManager.profile("apply_kinetic_bc") @@ -4576,10 +4617,19 @@ def _sendrecv_determine_mtbs( n = self.n_markers_used bi = self.first_pusher_idx sorting_etas = self._sorting_etas[:, :n] - sorting_etas[:] = ( - alpha * (self.markers[:3, :n] + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n]) - + (1.0 - alpha) * self.markers[bi : bi + 3, :n] - ) + if isinstance(alpha, float) and alpha == 1.0: + # the common case: alpha * (pos + shift) + 0 * init is exactly pos + shift for finite values, + # computed with one kernel instead of five + xp.add( + self.markers[:3, :n], + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n], + out=sorting_etas, + ) + else: + sorting_etas[:] = ( + alpha * (self.markers[:3, :n] + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n]) + + (1.0 - alpha) * self.markers[bi : bi + 3, :n] + ) # for non-periodic axes, eta = 1 is the right end of the domain and must not be wrapped to eta = 0 non_periodic_axes = [axis for axis in range(3) if axis not in self._periodic_axes] @@ -4610,6 +4660,151 @@ def _sendrecv_determine_mtbs( return hole_inds_after_send, send_inds + def _use_device_sort(self) -> bool: + """Whether :meth:`mpi_sort_markers` uses :meth:`_mpi_sort_device`: markers on a GPU (not on cunumpy's fake CuPy, + which cannot launch kernels). ``STRUPHY_DEVICE_SORT=1`` or ``0`` forces it on (also on the fake CuPy, e.g. in + tests with emulated launches) or off.""" + forced = os.environ.get("STRUPHY_DEVICE_SORT", "").strip().lower() + if forced in ("0", "false", "no"): + return False + if xp.get_array_backend(self._markers) != "cupy": + return False + if forced in ("1", "true", "yes"): + return True + from cunumpy.kernel_testing import fake_cupy_active + + return not fake_cupy_active() + + def _mpi_sort_device(self, alpha) -> tuple[int, np.ndarray]: + """The exchange of :meth:`mpi_sort_markers` for markers on the GPU. + + The same result as :meth:`_sendrecv_determine_mtbs`, :meth:`_sendrecv_get_destinations` and + :meth:`_sendrecv_markers` (the same markers in the same columns), with fewer passes over the marker buffer: + + - one kernel finds the destination of every marker (:func:`~struphy.pic.mpi_sort_cuda.sort_destinations`) + and counts the markers per destination, instead of a pass per direction and per rank; + - the leaving markers are listed once, stably sorted by destination (so the markers to each rank keep their + order, as before) and gathered once per destination rank. + + Markers that leave towards no rank's domain (e.g. a NaN position) become holes and are not sent, as before. + + Parameters + ---------- + alpha : tuple + Weights of the sorting position, see :meth:`mpi_sort_markers`. + + Returns + ------- + n_emptied : int + Number of markers that left (became holes), for :meth:`_track_markers`. + recv_info : numpy.ndarray + Number of markers received from each rank. + """ + from struphy.pic import mpi_sort_cuda + + alpha_host = np.array(alpha, dtype=float).ravel() + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + + if not hasattr(self, "_dev_sort_bufs"): + ranks = np.array(self._neighbor_ranks + self._non_neighbor_ranks, dtype=np.int32) + periodic = np.array([1 if axis in self._periodic_axes else 0 for axis in range(3)], dtype=np.int32) + bounds = mpi_sort_cuda.domain_bounds(self._domain_array_host) + self._dev_sort_bufs = { + "ranks": xp.asarray(ranks), + "periodic": xp.asarray(periodic), + "bounds": xp.asarray(bounds), + "dest": xp.empty(self.n_markers_max, dtype=np.int32), + "counts": xp.zeros(self.mpi_size + 2, dtype=np.uint64), + "alpha": xp.zeros(3, dtype=float), + } + bufs = self._dev_sort_bufs + use_alpha = not np.all(alpha_host == 1.0) + if use_alpha: + bufs["alpha"][:] = xp.asarray(alpha_host) + + n = self.n_markers_used + bi = self.first_pusher_idx + dest = bufs["dest"][:n] + counts_dev = bufs["counts"] + counts_dev.fill(0) + mpi_sort_cuda.sort_destinations( + self._markers, + self.holes, + self.ghost_particles, + bi + 3 + self.vdim, + bi, + bufs["alpha"], + use_alpha, + bufs["periodic"], + bufs["bounds"], + bufs["ranks"], + self.mpi_rank, + n, + bufs["dest"], + counts_dev, + ) + # markers per destination (lost first): a read of 8 bytes per rank to size the messages, like the one the + # nonzero() of the host path makes to size its result (marker data stays on the device) + counts = xp.to_numpy(counts_dev).astype(int) + n_lost = int(counts[0]) + send_info = counts[2:].copy() + + # the leaving markers, sorted by destination (stable: ascending columns per destination); lost markers first + leaving = dest != -1 + leave_inds = xp.flatnonzero(leaving) + n_leaving = n_lost + int(send_info.sum()) + if n_leaving: + leave_sorted = leave_inds[xp.argsort(dest[leave_inds], kind="stable")] + else: + leave_sorted = leave_inds + + # the columns that are holes after the send: the old holes (all columns from n_markers_used on are holes) + # and the leaving markers; the received markers fill the lowest of them + holes_after = self.holes.copy() + xp.logical_or(holes_after[:n], leaving, out=holes_after[:n]) + hole_inds_after_send = xp.flatnonzero(holes_after) + + # one contiguous buffer (n_cols, count) per destination rank, in the order of the host path + first_send = np.cumsum(send_info) - send_info + sends = [] + for rank in range(self.mpi_size): + if rank == self.mpi_rank or send_info[rank] == 0: + self._send_list[rank] = self._markers[:, :0] + continue + start = n_lost + int(first_send[rank]) + cols = leave_sorted[start : start + int(send_info[rank])] + self._send_list[rank] = xp.ascontiguousarray(self._markers[:, cols]) + sends.append((self._send_list[rank], rank, self.mpi_rank)) + + # the leaving markers are holes now (also the lost ones) + if n_leaving: + self._markers[:, leave_inds] = -1.0 + + recv_info = self._sendrecv_all_to_all(send_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f"Strong load imbalance detected: number of holes ({hole_inds_after_send.size}) on rank " + f"{self.mpi_rank} is smaller than number of incoming particles ({n_recv}). " + 'Increase "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + first_recv = np.cumsum(recv_info) - recv_info + receives = [] + for rank in range(self.mpi_size): + if rank != self.mpi_rank and recv_info[rank] > 0: + receives.append((xp.empty((self.n_cols, int(recv_info[rank])), dtype=self._markers.dtype), rank, rank)) + + xp.mpi.exchange(self.mpi_comm, sends=sends, receives=receives, cuda_aware=self._mpi_cuda_aware()) + + for recvbuf, rank, _ in receives: + start = int(first_recv[rank]) + self._markers[:, hole_inds_after_send[start : start + int(recv_info[rank])]] = recvbuf + + return n_leaving, recv_info + 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`. diff --git a/src/struphy/pic/mpi_sort_cuda.cu b/src/struphy/pic/mpi_sort_cuda.cu new file mode 100644 index 000000000..ff9c3d1e5 --- /dev/null +++ b/src/struphy/pic/mpi_sort_cuda.cu @@ -0,0 +1,79 @@ +// Destinations of the markers in Particles.mpi_sort_markers on the GPU (struphy.pic.mpi_sort_cuda). +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +// The domain test of Particles._is_on_domain_of_rank for one rank: [left, right) in every direction. +__device__ inline bool on_domain(const double* eta, const double* bounds, int r) { + for (int a = 0; a < 3; ++a) { + const double left = bounds[(r * 3 + a) * 2], right = bounds[(r * 3 + a) * 2 + 1]; + if (!(eta[a] >= left && eta[a] < right)) return false; + } + return true; +} + +/** + * The rank each marker has to go to, as Particles._sendrecv_determine_mtbs and _sendrecv_get_destinations decide it. + * + * Marker ip (a column of the buffer, ip < n_used) is sorted by eta = pos + shift (use_alpha == 0) or + * alpha * (pos + shift) + (1 - alpha) * init per direction (rows 0-2, shift_row.., init_row..), reduced modulo 1 as + * numpy.mod does; on a non-periodic axis eta = 1 stays 1. It stays (dest -1) if it is a hole or a ghost or lies in + * the domain of this rank, else it goes to the rank whose domain contains eta (bounds: [left, right) per rank and + * direction, the right end 1 shifted to the next float, see domain_bounds), or it is lost (dest -2) if no domain + * contains it (e.g. a NaN position). The domains do not overlap, so the rank is unique. + * + * One thread per marker (n_threads = n_used). counts[dest + 2] counts the leaving markers (dest -2 or a rank). + * + * @param markers Marker buffer (n_cols x n_markers_max, one column per marker). + * @param holes, ghosts Hole and ghost masks (n_markers_max). + * @param shift_row, init_row First rows of the boundary shifts and of the initial positions. + * @param alpha, use_alpha The three weights, used if use_alpha (else alpha = 1, eta = pos + shift). + * @param periodic Per direction, 1 if periodic. + * @param bounds (n_ranks, 3, 2): left and right end of each rank's domain per direction. + * @param ranks The other ranks in the order they are checked (neighbours first), n_others entries. + * @param n_others Number of entries of ranks. + * @param rank This rank. + * @param n_used Number of markers (columns) that can hold markers. + * @param dest Output (n_markers_max): -1 stay, -2 lost, else the destination rank (written for ip < n_used). + * @param counts Output (n_ranks + 2), added to: leaving markers per destination (dest + 2). + */ +extern "C" __global__ void sort_destinations(Array2D markers, const bool* holes, const bool* ghosts, + int shift_row, int init_row, const double* alpha, int use_alpha, + const int* periodic, + const double* bounds, const int* ranks, int n_others, int rank, + long long n_used, int* dest, unsigned long long* counts) { + CUNUMPY_THREAD_1D(ip, n_used); + if (holes[ip] || ghosts[ip]) { + dest[ip] = -1; + return; + } + + double eta[3]; + for (int a = 0; a < 3; ++a) { + const double moved = markers(a, ip) + markers(shift_row + a, ip); + double e = use_alpha ? alpha[a] * moved + (1.0 - alpha[a]) * markers(init_row + a, ip) : moved; + const bool right_end = !periodic[a] && e == 1.0; + // numpy.mod(e, 1.0): the remainder with the sign of the divisor + double r = fmod(e, 1.0); + if (r != 0.0) { + if (r < 0.0) r += 1.0; + } else { + r = 0.0; + } + eta[a] = right_end ? 1.0 : r; + } + + if (on_domain(eta, bounds, rank)) { + dest[ip] = -1; + return; + } + int d = -2; + for (int k = 0; k < n_others; ++k) { + if (on_domain(eta, bounds, ranks[k])) { + d = ranks[k]; + break; + } + } + dest[ip] = d; + cunumpy_atomic_add(&counts[d + 2], 1ULL); +} diff --git a/src/struphy/pic/mpi_sort_cuda.py b/src/struphy/pic/mpi_sort_cuda.py new file mode 100644 index 000000000..a862d8831 --- /dev/null +++ b/src/struphy/pic/mpi_sort_cuda.py @@ -0,0 +1,99 @@ +"""The destinations of :meth:`~struphy.pic.base.Particles.mpi_sort_markers` on the GPU (``mpi_sort_cuda.cu``). + +:meth:`~struphy.pic.base.Particles._mpi_sort_device` finds the rank of every leaving marker with one kernel +(:func:`sort_destinations`) instead of a pass over the buffer per rank and per direction. +""" + +from functools import cache +from pathlib import Path + +import numpy as np + + +@cache +def _kernel(): + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("mpi_sort_cuda.cu"), "sort_destinations") + + +def domain_bounds(domain_array) -> np.ndarray: + """The ends of each rank's domain per direction, as :meth:`~struphy.pic.base.Particles._is_on_domain_of_rank` tests + them: an array (n_ranks, 3, 2) of [left, right), with the right end 1 shifted to the next float (eta = 1 belongs to + the last rank). + + Parameters + ---------- + domain_array : array + (n_ranks, 9) host array of :attr:`~struphy.pic.base.Particles.domain_array` (left, right, cells per direction). + """ + domain_array = np.asarray(domain_array, dtype=float) + left = domain_array[:, 0::3] + right = domain_array[:, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + return np.ascontiguousarray(np.stack([left, right], axis=-1)) + + +def sort_destinations( + markers, + holes, + ghosts, + shift_row: int, + init_row: int, + alpha, + use_alpha: bool, + periodic, + bounds, + ranks, + rank: int, + n_used: int, + dest, + counts, +): + """Launch the kernel ``sort_destinations`` (one thread per marker, see ``mpi_sort_cuda.cu``). + + Parameters + ---------- + markers : cupy.ndarray + Marker buffer (n_cols, n_markers_max). + holes, ghosts : cupy.ndarray + Hole and ghost masks (bool, n_markers_max). + shift_row, init_row : int + First rows of the boundary shifts and of the initial positions. + alpha : cupy.ndarray + The three sorting weights (float64, 3). + use_alpha : bool + Whether to use them; else the sorting position is position + shift (alpha = 1, computed exactly). + periodic : cupy.ndarray + Per direction, 1 if periodic (int32, 3). + bounds : cupy.ndarray + :func:`domain_bounds` on the device (float64). + ranks : cupy.ndarray + The other ranks, neighbours first (int32). + rank, n_used : int + This rank and the number of markers (columns) that can hold markers. + dest : cupy.ndarray + Output, int32 (n_markers_max): -1 stay, -2 lost, else the destination rank. + counts : cupy.ndarray + Output, uint64 (n_ranks + 2), added to: leaving markers per destination + 2. + """ + if n_used == 0: + return + _kernel()( + markers, + holes, + ghosts, + int(shift_row), + int(init_row), + alpha, + int(bool(use_alpha)), + periodic, + bounds, + ranks, + int(ranks.size), + int(rank), + int(n_used), + dest, + counts, + n_threads=int(n_used), + ) diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 00854437c..4691e6547 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -395,6 +395,10 @@ def f_coords(self): def f_coords(self, new): Particles6D.f_coords.fset(self, new) + def _f_coords_of_valid(self, quantity, rows): + # the control variate is evaluated at the gyrotropic velocity of f_coords, not at the stored v_par b0 + return list(self.f_coords) + def _maxwellians_in(f): """The Maxwellians of a kinetic background (summands of a :class:`~struphy.kinetic_background.base.SumKineticBackground`).""" diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index 4fd52bba0..e66af628c 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -11,6 +11,8 @@ default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. """ +import contextlib +import os import sys from contextlib import contextmanager @@ -89,12 +91,31 @@ def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.nd return new -def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): +def check_mpi_sort_markers_device( + cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True, device_sort: bool = False +): """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` - force the device-buffer or the host-staged exchange. + force the device-buffer or the host-staged exchange. With ``device_sort`` the CuPy markers are sorted with the + destination kernel (:meth:`Particles._mpi_sort_device`, launches emulated on the fake CuPy), which reads the + counts per destination (one copy to the host per sort). """ + if device_sort: + import os + + from cunumpy.kernel_testing import emulated_launches + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + os.environ["STRUPHY_DEVICE_SORT"] = "1" + try: + with emulated_launches(options=EMULATION_OPTIONS): + check_mpi_sort_markers_device(cuda_aware, bc, count_reads=False) + finally: + del os.environ["STRUPHY_DEVICE_SORT"] + return + comm = MPI.COMM_WORLD rank, size = comm.Get_rank(), comm.Get_size() @@ -138,19 +159,23 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", with on_device(): dev.markers[:] = xp.to_cupy(new) dev.update_holes() - with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + reads_ctx = _count_host_reads() if count_reads else contextlib.nullcontext([]) + with xp.profiling.count_transfers() as transfers, reads_ctx as reads: dev.mpi_sort_markers() + # the destination kernel reads its counts per destination (8 bytes per rank) once per sort + with on_device(): + counts_read = 1 if dev._use_device_sort() else 0 if cuda_aware: - assert transfers.total == 0, transfers.report() + assert transfers.total == counts_read, transfers.report() else: # exchange stages each non-empty send and receive buffer separately n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[1] > 0) n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[1] > 0) - assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_host == n_dest + counts_read, transfers.report() assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") - assert sent_bytes == n_send * ref.n_cols * 8 + assert sent_bytes == n_send * ref.n_cols * 8 + counts_read * (size + 2) * 8 received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") assert received_bytes == n_recv * ref.n_cols * 8 @@ -181,6 +206,15 @@ def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): check_mpi_sort_markers_device(cuda_aware, bc=bc) +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_device_kernel_fake_cupy(cuda_aware, bc): + """The destination kernel (Particles._mpi_sort_device) against the NumPy sort, launches emulated.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc, device_sort=True) + + @requires_cupy @pytest.mark.mpi @pytest.mark.parametrize("bc", ["periodic", "remove"]) diff --git a/src/struphy/propagators/push_vxb.py b/src/struphy/propagators/push_vxb.py index f68e7a787..c48fdcac6 100644 --- a/src/struphy/propagators/push_vxb.py +++ b/src/struphy/propagators/push_vxb.py @@ -152,12 +152,16 @@ def allocate(self): @profile def __call__(self, dt): - # sum up total magnetic field in the tensor-product space (the variable is extracted first) - b_full: BlockVector = self.b2_0.copy(out=self._b_full) - if self.b2_var is not None: - b_full += self._E2T.dot(self.b2_var.spline.vector, out=self._tmp) - b_full.update_ghost_regions() - b_full /= self._epsilon + # sum up total magnetic field in the tensor-product space (the variable is extracted first); without a + # magnetic field variable it is the equilibrium field alone, which does not change: computed once, so that the + # step has no ghost exchange (an MPI synchronization point) before the kernel + if self.b2_var is not None or not getattr(self, "_b_full_static_done", False): + b_full: BlockVector = self.b2_0.copy(out=self._b_full) + if self.b2_var is not None: + b_full += self._E2T.dot(self.b2_var.spline.vector, out=self._tmp) + b_full.update_ghost_regions() + b_full /= self._epsilon + self._b_full_static_done = self.b2_var is None # call pusher kernel self._pusher(dt)