From aeebf95b2ac28f497bfeb903db6a640e7b58f546 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 18:43:56 +0200 Subject: [PATCH 01/12] feectools: Kronecker solver on the device Moves the feectools submodule to struphy-hub/feectools#96 (branch device-kronecker-solve): KroneckerLinearSolver solves device data with dense inverses of the 1D matrices (one GEMM per direction) instead of copying it to the host for LAPACK/SuperLU, so MassMatrixPreconditioner solves on the CuPy backend make no host/device transfers. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index f97ea5af7..9cb793b5d 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit f97ea5af7c9a21808f7d5adfc95c907c338d74de +Subproject commit 9cb793b5de72af2d8d64bb42f0c9e593bd4499a2 From 39f19847e9194257fccd20beb07f55b6d9fbd709 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 19:16:07 +0200 Subject: [PATCH 02/12] feectools: stencil 3D device kernels with 6D views Point the feectools submodule at struphy-hub/feectools#97 (branch stencil-6d-views): the 3D stencil dot/transpose kernels take the matrix data as 6D views, and StencilMatrix.dot/vdot/transpose no longer raise for matrices with fewer than 2p+1 diagonals. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index f97ea5af7..1958568d1 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit f97ea5af7c9a21808f7d5adfc95c907c338d74de +Subproject commit 1958568d176c18aa7e4082b99257e2385a31b6fc From 081768b8765b99574657b0150238c090d0a0087c Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:05:27 +0200 Subject: [PATCH 03/12] Move the feectools submodule to the head of feectools#96 (64b10fa) Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 9cb793b5d..64b10fa0e 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 9cb793b5de72af2d8d64bb42f0c9e593bd4499a2 +Subproject commit 64b10fa0ecaf5d8dc138163b8df6b7594f46676a From c58ccc048334ef8fef3e47d3c272118061925c29 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:50:20 +0200 Subject: [PATCH 04/12] feectools: inner products stay on the device Move the feectools submodule to struphy-hub/feectools#98: on the CuPy backend inner products return 0-d device arrays and the Krylov solvers copy one scalar per iteration to the host (the convergence test). Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index f97ea5af7..a15a8e80c 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit f97ea5af7c9a21808f7d5adfc95c907c338d74de +Subproject commit a15a8e80ce831076b66887f7636ee2d7ce321be1 From 1845bc652446ce46eebc825597db01d14fec0dd8 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:13:49 +0200 Subject: [PATCH 05/12] MassMatrixPreconditioner on the CuPy backend The mass-matrix preconditioners could not be created on the CuPy backend: the assembled 1d mass matrices come back from StencilMatrix.toarray() as host (NumPy) arrays, and is_circulant/FFTSolver asserted xp.ndarray (cupy.ndarray on CuPy). The stencil-matrix fill also mixed xp.nonzero with host arrays. The 1d matrices and solvers are now host setup data on every backend (what KroneckerLinearSolver expects, and what feectools#96 builds its device inverses from); only the process-local 1d stencil matrices of the KroneckerStencilMatrix are copied to the device once. The duplicated 1d setup of both preconditioners moves into _solver_and_local_matrix_1d. FFTSolver accepts host or device matrices and solves device right-hand sides on the host. The NumPy path computes the same as before. Adds a fake-CuPy test (subprocess, CUDA launches emulated on the CPU) comparing both preconditioners for M0/M1/M2 (periodic and clamped) with NumPy, and the same check on a GPU. Part of #689 and #650. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 8 + src/struphy/feec/preconditioner.py | 214 +++++++++--------- .../feec/tests/test_preconditioner_cupy.py | 183 +++++++++++++++ 3 files changed, 293 insertions(+), 112 deletions(-) create mode 100644 src/struphy/feec/tests/test_preconditioner_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..699540382 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -270,6 +270,14 @@ the `pyproject.toml` pin (today `feectools>=0.3.0, <=0.3.0`) point at that relea deprecated (removed in 0.6): import them from `cunumpy.cuda` and `cunumpy.mpi`, and require `cunumpy>=0.5.0, <0.6` like struphy. +**Mass-matrix preconditioners on CuPy** (`feec/preconditioner.py`, for #689). `MassMatrixPreconditioner` and +`MassMatrixDiagonalPreconditioner` build their 1d mass matrices and 1d solvers (`FFTSolver` for circulant, +`SparseSolver` otherwise) on the host on every backend; only the process-local 1d stencil matrices (factors of the +`KroneckerStencilMatrix`) are copied to the device. `is_circulant` and `FFTSolver` take host or device arrays; +`FFTSolver.solve` solves device data on the host and copies it back, until the device Kronecker solve of +[feectools#96](https://github.com/struphy-hub/feectools/pull/96) builds dense inverses from these host solvers. +Tested with fake CuPy and emulated launches (`feec/tests/test_preconditioner_cupy.py`). + New feectools work for the GPU follows the same pattern: a `cuda--` PR, linked in the table above and from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/issues/650)). diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index ac413ba59..9ae8a273f 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.fem.tensor import TensorFemSpace @@ -228,52 +229,10 @@ def fun(e): identity=True, ) - M_arr = M.toarray() - - # create 1d solver for mass matrix - if is_circulant(M_arr): - solvercells += [FFTSolver(M_arr)] - else: - solvercells += [SparseSolver(M.tosparse())] - - # === NOTE: for KroneckerStencilMatrix being built correctly, 1d matrices must be local to process! === - periodic = femspaces[c].coeff_space.periods[d] - - n = femspaces[c].coeff_space.npts[d] - p = femspaces[c].coeff_space.pads[d] - s = femspaces[c].coeff_space.starts[d] - e = femspaces[c].coeff_space.ends[d] - - cart_decomp_1d = CartDecomposition( - domain_decompos_1d, - [n], - [[s]], - [[e]], - [p], - [1], - ) - - V_local = StencilVectorSpace(cart_decomp_1d) - - M_local = StencilMatrix(V_local, V_local) - - row_indices, col_indices = xp.nonzero(M_arr) - - for row_i, col_i in zip(row_indices, col_indices): - # only consider row indices on process - if row_i in range(V_local.starts[0], V_local.ends[0] + 1): - row_i_loc = row_i - s - - M_local._data[ - row_i_loc + p, - (col_i + p - row_i) % M_arr.shape[1], - ] = 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]) - - matrixcells += [M_local.copy()] - # ======================================================================================================= + # 1d solver (host) and process-local 1d stencil matrix (active backend) + solver_1d, M_local = _solver_and_local_matrix_1d(M, femspaces[c].coeff_space, d, domain_decompos_1d) + solvercells += [solver_1d] + matrixcells += [M_local] if isinstance(self._femspace, TensorFemSpace): matrixblocks += [ @@ -587,52 +546,10 @@ def __init__(self, mass_operator, apply_bc=True): identity=True, ) - M_arr = M.toarray() - - # create 1d solver for mass matrix - if is_circulant(M_arr): - solvercells += [FFTSolver(M_arr)] - else: - solvercells += [SparseSolver(M.tosparse())] - - # === NOTE: for KroneckerStencilMatrix being built correctly, 1d matrices must be local to process! === - periodic = femspaces[c].coeff_space.periods[d] - - n = femspaces[c].coeff_space.npts[d] - p = femspaces[c].coeff_space.pads[d] - s = femspaces[c].coeff_space.starts[d] - e = femspaces[c].coeff_space.ends[d] - - cart_decomp_1d = CartDecomposition( - domain_decompos_1d, - [n], - [[s]], - [[e]], - [p], - [1], - ) - - V_local = StencilVectorSpace(cart_decomp_1d) - - M_local = StencilMatrix(V_local, V_local) - - row_indices, col_indices = xp.nonzero(M_arr) - - for row_i, col_i in zip(row_indices, col_indices): - # only consider row indices on process - if row_i in range(V_local.starts[0], V_local.ends[0] + 1): - row_i_loc = row_i - s - - M_local._data[ - row_i_loc + p, - (col_i + p - row_i) % M_arr.shape[1], - ] = 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]) - - matrixcells += [M_local.copy()] - # ======================================================================================================= + # 1d solver (host) and process-local 1d stencil matrix (active backend) + solver_1d, M_local = _solver_and_local_matrix_1d(M, femspaces[c].coeff_space, d, domain_decompos_1d) + solvercells += [solver_1d] + matrixcells += [M_local] if isinstance(self._femspace, TensorFemSpace): matrixblocks += [ @@ -884,23 +801,89 @@ def dot(self, v, out=None): return out +def _solver_and_local_matrix_1d(M, coeff_space, d, domain_decomposition_1d): + """1d solver and process-local 1d stencil matrix of an assembled 1d mass matrix. + + The 1d matrix is small setup data: it is converted to a dense host (NumPy) array on every backend, and + the 1d solver (:class:`FFTSolver` for circulant matrices, :class:`SparseSolver` otherwise) is a host + solver, as expected by :class:`KroneckerLinearSolver`. Only the process-local stencil matrix, which + becomes a factor of the :class:`KroneckerStencilMatrix`, holds its data on the active backend. + + Parameters + ---------- + M : StencilMatrix + The assembled 1d mass matrix (serial, not distributed). + + coeff_space : StencilVectorSpace + The (distributed) 3d coefficient space of the component. + + d : int + The spatial direction of the 1d matrix. + + domain_decomposition_1d : DomainDecomposition + Domain decomposition of the serial 1d FEM space. + + Returns + ------- + solver : FFTSolver | SparseSolver + Host solver for the 1d matrix. + + M_local : StencilMatrix + The rows ``coeff_space.starts[d]`` to ``coeff_space.ends[d]`` of the 1d matrix as a stencil matrix. + """ + M_arr = xp.to_numpy(M.toarray()) + + if is_circulant(M_arr): + solver = FFTSolver(M_arr) + else: + solver = SparseSolver(M.tosparse()) + + # === NOTE: for KroneckerStencilMatrix being built correctly, 1d matrices must be local to process! === + n = coeff_space.npts[d] + p = coeff_space.pads[d] + s = coeff_space.starts[d] + e = coeff_space.ends[d] + + cart_decomp_1d = CartDecomposition(domain_decomposition_1d, [n], [[s]], [[e]], [p], [1]) + V_local = StencilVectorSpace(cart_decomp_1d) + M_local = StencilMatrix(V_local, V_local) + + # fill the stencil data on the host, then copy it to the active backend once + data = np.zeros(M_local._data.shape, dtype=M_local._data.dtype) + 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 + if row_i in range(V_local.starts[0], V_local.ends[0] + 1): + row_i_loc = row_i - s + data[row_i_loc + p, (col_i + p - row_i) % M_arr.shape[1]] = M_arr[row_i, col_i] + M_local._data[...] = xp.asarray(data) + + # check if stencil matrix was built correctly + assert np.allclose(xp.to_numpy(M_local.toarray())[s : e + 1], M_arr[s : e + 1]) + + return solver, M_local + + class FFTSolver(BandedSolver): """ Solve the equation Ax = b for x, assuming A is a circulant matrix. b can contain multiple right-hand sides (RHS) and is of shape (#RHS, N). + The circulant matrix is kept on the host (it is small 1d setup data, see :class:`KroneckerLinearSolver`). + Right-hand sides on the device (CuPy backend) are solved on the host and copied back. + Parameters ---------- - circmat : xp.ndarray - Generic circulant matrix. + circmat : numpy.ndarray + Generic circulant matrix (a device array is copied to the host). """ def __init__(self, circmat): - assert isinstance(circmat, xp.ndarray) + circmat = xp.to_numpy(circmat) assert is_circulant(circmat) self._space = xp.ndarray - self._column = circmat[:, 0] + self._column = circmat[:, 0].copy() # -------------------------------------- # Abstract interface @@ -909,6 +892,19 @@ def __init__(self, circmat): def space(self): return self._space + def _solve_host(self, rhs, stabilize): + from scipy.linalg import solve_circulant + + try: + return solve_circulant(self._column, rhs.T).T + except np.linalg.LinAlgError: + if not stabilize: + raise + eps = 1e-4 + logger.info(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") + self._column[0] *= 1.0 + eps + return solve_circulant(self._column, rhs.T).T + @profile def solve(self, rhs, out=None, transposed=False): """ @@ -920,7 +916,7 @@ def solve(self, rhs, out=None, transposed=False): The right-hand sides to solve for. The vectors are assumed to be given in C-contiguous order, i.e. if multiple right-hand sides are given, then rhs is a two-dimensional array with the 0-th index denoting the number of the right-hand side, and the 1-st index denoting the element inside - a right-hand side. + a right-hand side. A device array is solved on the host; the result is a device array. out : xp.ndarray, optional Output vector. If given, it has to have the same shape and datatype as rhs. @@ -929,24 +925,18 @@ def solve(self, rhs, out=None, transposed=False): If and only if set to true, we solve against the transposed matrix. (supported by the underlying solver) """ - from scipy.linalg import solve_circulant - assert rhs.T.shape[0] == self._column.size - if out is None: - out = solve_circulant(self._column, rhs.T).T + # a singular matrix is stabilized only for in-place solves (as called by KroneckerLinearSolver) + on_device = xp.is_gpu(rhs) + result = self._solve_host(xp.to_numpy(rhs) if on_device else rhs, stabilize=out is not None) + if out is None: + out = xp.asarray(result) if on_device else result else: assert out.shape == rhs.shape assert out.dtype == rhs.dtype - - try: - out[:] = solve_circulant(self._column, rhs.T).T - except xp.linalg.LinAlgError: - eps = 1e-4 - logger.info(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") - self._column[0] *= 1.0 + eps - out[:] = solve_circulant(self._column, rhs.T).T + out[...] = xp.asarray(result) if xp.is_gpu(out) else result return out @@ -958,7 +948,7 @@ def is_circulant(mat): Parameters ---------- mat : array[float] - The matrix that is checked to be circulant. + The matrix that is checked to be circulant (host or device array; checked on the host). Returns ------- @@ -966,13 +956,13 @@ def is_circulant(mat): Whether the matrix is circulant (=True) or not (=False). """ - assert isinstance(mat, xp.ndarray) + mat = xp.to_numpy(mat) 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 = bool(np.allclose(mat[i, :], np.roll(mat[i + 1, :], -1))) if not circulant: return circulant else: diff --git a/src/struphy/feec/tests/test_preconditioner_cupy.py b/src/struphy/feec/tests/test_preconditioner_cupy.py new file mode 100644 index 000000000..358497d13 --- /dev/null +++ b/src/struphy/feec/tests/test_preconditioner_cupy.py @@ -0,0 +1,183 @@ +"""Mass-matrix preconditioners on the CuPy backend (struphy-hub/struphy#689, part of #650). + +The 1d mass matrices and solvers of the Kronecker preconditioners are host (NumPy) setup data on every +backend; the Kronecker stencil matrices and the vectors they are applied to live on the device. These tests +build the preconditioners on the CuPy backend and compare their action with the NumPy backend: + +* without a GPU, with cunumpy's fake CuPy (host memory, CuPy semantics) in a subprocess, the CUDA kernels + launched on the way (feectools' stencil kernels, struphy's geometry kernels) run by CPU emulation; +* on a GPU (skipped otherwise), directly. +""" + +import contextlib +import subprocess +import sys + +import cunumpy +import numpy as np +import pytest + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") + +SPACES = ("M0", "M1", "M2") +PRECONDITIONERS = ("MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner") +BCS = { + "periodic": (None, None, None), + "clamped": (("dirichlet", "dirichlet"), ("free", "free"), ("free", "dirichlet")), +} + + +def _build(backend, bcs): + """Derham (6 x 5 x 4 elements, degrees (2, 3, 1)) and mass operators on a Colella mapping, on `backend`.""" + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + with cunumpy.use_backend(backend): + options = DerhamOptions(degree=(2, 3, 1), bcs=bcs) + derham = Derham(TensorProductGrid(num_elements=(6, 5, 4)), options, comm=MPI.COMM_WORLD) + return WeightedMassOperators(derham, domains.Colella()) + + +def _blocks(v): + """The stencil vectors of a stencil or block vector.""" + return list(v.blocks) if hasattr(v, "blocks") else [v] + + +def _random_vector(space, rng): + v = space.zeros() + for block in _blocks(v): + block._data[...] = cunumpy.asarray(rng.random(block._data.shape)) + v.update_ghost_regions() + return v + + +def _to_device(v_host, space): + """A vector of `space` (active backend) with the data of the host vector `v_host`.""" + v = space.zeros() + for block, block_host in zip(_blocks(v), _blocks(v_host)): + block._data[...] = cunumpy.asarray(block_host._data) + v.update_ghost_regions() + return v + + +def _to_host(v): + return np.concatenate([cunumpy.to_numpy(block._data).ravel() for block in _blocks(v)]) + + +def check_preconditioners_on_cupy(bcs_name): + """Every preconditioner of :data:`PRECONDITIONERS` for M0, M1 and M2 gives the NumPy result on the CuPy backend.""" + from struphy.feec import preconditioner + + bcs = BCS[bcs_name] + host, device = _build("numpy", bcs), _build("cupy", bcs) + rng = np.random.default_rng(1) + + for space in SPACES: + with cunumpy.use_backend("numpy"): + M_host = getattr(host, space) + v_host = _random_vector(M_host.domain, rng) + with cunumpy.use_backend("cupy"): + M_device = getattr(device, space) + v_device = _to_device(v_host, M_device.domain) + + for name in PRECONDITIONERS: + cls = getattr(preconditioner, name) + with cunumpy.use_backend("numpy"): + expected = _to_host(cls(M_host).dot(v_host)) + with cunumpy.use_backend("cupy"): + pc = cls(M_device) + out = pc.dot(v_device) + for block in _blocks(out): + assert cunumpy.is_gpu(block._data), (bcs_name, space, name) + # in-place application + out_inplace = pc.codomain.zeros() + pc.dot(v_device, out=out_inplace) + assert np.linalg.norm(expected) > 0, (bcs_name, space, name) + np.testing.assert_allclose(_to_host(out), expected, rtol=1e-12, atol=1e-12 * np.abs(expected).max()) + np.testing.assert_allclose(_to_host(out_inplace), expected, rtol=1e-12, atol=1e-12 * np.abs(expected).max()) + + +def _host_buffer(value): + """The host buffer behind a fake CuPy array (the same memory, no copy); other values unchanged.""" + if isinstance(value, np.ndarray) or not hasattr(value, "__cuda_array_interface__"): + return value + return object.__getattribute__(value, "_a") + + +class _HostFields: + """A CUDA argument object whose array attributes are their fake CuPy host buffers.""" + + def __init__(self, args): + self._args = args + + def __getattr__(self, name): + return _host_buffer(getattr(self._args, name)) + + +@contextlib.contextmanager +def _emulated_launches(): + """On cunumpy's fake CuPy, run every ``CudaKernel`` launch by CPU emulation, in place on the fake device arrays. + + A minimal version of ``emulated_launches`` of struphy-hub/struphy#705; use that one once it is merged. + """ + from cunumpy.kernel_testing import emulate_cuda_kernel, fake_cupy_active + from cunumpy.kernels import CudaKernel + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS, emulate_struct_kernel + + assert fake_cupy_active() + original = CudaKernel.__call__ + + def launch(self, *args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None): + if any(p.struct is not None for p in self.signature): + host_args = [_host_buffer(a) if p.struct is None else _HostFields(a) for p, a in zip(self.signature, args)] + emulate_struct_kernel(self, *host_args, n_threads=n_threads) + else: + grid, block = self.launch_shape(n_threads, grid=grid, block=block, args=args) + host_args = [_host_buffer(a) for a in args] + emulate_cuda_kernel( + self, *host_args, grid=grid, block=block, shared_mem=shared_mem, options=EMULATION_OPTIONS + ) + + CudaKernel.__call__ = launch + try: + yield + finally: + CudaKernel.__call__ = original + + +def check_preconditioners_on_fake_cupy(bcs_name): + """:func:`check_preconditioners_on_cupy` with the CUDA kernel launches emulated on the CPU.""" + with _emulated_launches(): + check_preconditioners_on_cupy(bcs_name) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("bcs_name", list(BCS)) +def test_mass_preconditioners_fake_cupy(bcs_name): + """Without a GPU: the preconditioners build and apply on cunumpy's fake CuPy, which rejects host/device mixing. + + Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. + """ + code = ( + "from struphy.feec.tests.test_preconditioner_cupy import check_preconditioners_on_fake_cupy; " + f"check_preconditioners_on_fake_cupy({bcs_name!r})" + ) + result = subprocess.run( + [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr[-4000:] + + +@requires_cupy +@pytest.mark.parametrize("bcs_name", list(BCS)) +def test_mass_preconditioners_on_cupy(bcs_name): + """On a GPU: the preconditioners give the NumPy result on the CuPy backend.""" + check_preconditioners_on_cupy(bcs_name) From afcbb9023e5c7adb6d8a18712b27766bae7cf2ef Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:13:50 +0200 Subject: [PATCH 06/12] End-to-end GPU test for LinearVlasovAmpereOneSpecies Run a small LinearVlasovAmpereOneSpecies case (periodic Cuboid, 8 elements, 320 Sobol markers, 3 steps) on NumPy and CuPy, compare energies, E-field coefficients and markers, and count the host/device transfers of the time loop per phase (none in model.integrate, scalar-sized downloads in the diagnostics, one download per saved dataset in the output). GPU-only part is marked requires_cupy; the NumPy reference runs everywhere and under MPI. Sobol marker loading now generates on the host and copies the result to the active backend: sobol_seq is a scalar loop and passed lists to cupy.transpose. Part of #689. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 1 + .../test_gpu_LinearVlasovAmpereOneSpecies.py | 332 ++++++++++++++++++ src/struphy/pic/base.py | 21 +- src/struphy/pic/sobol_seq.py | 57 +-- 4 files changed, 376 insertions(+), 35 deletions(-) create mode 100644 src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..92490e29d 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -28,6 +28,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). - [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). +- [ ] **End-to-end test of `LinearVlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` runs a small Cuboid case (320 Sobol markers, 8 elements, 3 steps) on NumPy and CuPy, compares energies, E-field coefficients and markers, and counts the transfers of the time loop: none in `model.integrate`, scalar-sized downloads in the diagnostics, one download per saved dataset in the output. GPU-only; fails until the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere` CUDA kernel, `MassMatrixPreconditioner` on CuPy and device inner products are merged and released. - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). diff --git a/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py new file mode 100644 index 000000000..7808bde18 --- /dev/null +++ b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py @@ -0,0 +1,332 @@ +"""End-to-end GPU test of :class:`~struphy.models.LinearVlasovAmpereOneSpecies` (issue #689). + +A small run of the model (the energy-conservation setup of +``verification/test_verif_LinearVlasovAmpereOneSpecies.py`` with fewer markers) is done once on +the NumPy backend and once on the CuPy backend. The test checks that + +1. both runs give the same energies, electric-field coefficients and marker weights/positions, + within tolerances explained in :func:`compare_runs`, and +2. the time loop of the CuPy run does no host/device transfers through cunumpy, apart from an + explicit budget for the diagnostics and the output (see :class:`TransferBudget`). + +The CuPy part needs a GPU (``requires_cupy``) and is skipped otherwise. It is MPI-compatible: under +``mpirun -n 2 pytest --with-mpi`` both runs use the same domain decomposition, and the markers are +gathered and compared by marker ID. + +Only transfers made through cunumpy are counted (``to_numpy``, ``to_cupy``, host fallbacks of +kernels, ...). Implicit synchronizations such as ``float(device_scalar)`` in the convergence check +of the iterative solver are not seen; use ``nsys`` for those. +""" + +import logging +import os +from dataclasses import dataclass, field + +import cunumpy as xp +import h5py +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from cunumpy.profiling import TransferCounter, count_transfers +from maybempi import MPI + +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.io.output_handling import DataContainer +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import LinearVlasovAmpereOneSpecies + +logger = logging.getLogger("struphy") + +AMPLITUDE = 0.1 +R1 = 12.56 # box of length 4*pi, k = 0.5 +NUM_ELEMENTS = (8, 1, 1) +PPC = 40 # 320 markers +DT = 0.1 +N_STEPS = 3 +SOLVER_TOL = 1e-12 + +# transfer kinds that move data between host and device (device_copy is device-only and allowed) +HOST_DEVICE_KINDS = ("to_host", "to_device", "kernel_conversion", "fallback") + + +def make_simulation(out_folders: str, sim_folder: str) -> Simulation: + """The model on a periodic 1D Cuboid, with the active cunumpy backend. + + Deterministic marker loading (``sobol_standard``) so that both backends start from the same + markers: ``pseudo_random`` draws from ``xp.random``, whose streams differ between NumPy and CuPy. + No B0 and no E0, so the time step is ``PushEta`` followed by ``EfieldWeightsCoupling``. No sorting + boxes (the verification test uses them): the model does not need them, and the box sorting + kernels have no CUDA version yet. + """ + model = LinearVlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False, with_E0=False) + + env = EnvironmentOptions(out_folders=out_folders, sim_folder=sim_folder) + time_opts = Time(dt=DT, Tend=round(N_STEPS * DT, 14)) # 3 * 0.1 > 0.3 would add a step + domain = domains.Cuboid(r1=R1) + grid = grids.TensorProductGrid(num_elements=NUM_ELEMENTS) + derham_opts = DerhamOptions(degree=(3, 1, 1)) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=PPC, loading="sobol_standard"), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(), + # no sorting boxes: optional for this model, and box sorting has no device version yet + sorting_params=SortingParameters(), + saving_params=SavingParameters(), + bufsize=0.4, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + # a tight solver tolerance, so that the backends differ by round-off and not by where pcg stops + model.propagators.coupling_Eweights.options = model.propagators.coupling_Eweights.Options( + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + model.initial_poisson.options = model.initial_poisson.Options( + stab_mat="M0", + stab_eps=1e-6, + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + perturbation = perturbations.ModesCos(ls=(1,), amps=(AMPLITUDE,)) + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + return Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domain, + grid=grid, + derham_opts=derham_opts, + ) + + +@dataclass +class TransferBudget: + """Transfers recorded per phase of the time loop (setup before it is not counted). + + ``integrate`` (the propagators) must not transfer at all. ``diagnostics`` (scalars, saved + markers, binned distribution functions, coefficient extraction) may copy small results to the + host: at most a few scalars per call, never a host-to-device copy or a host kernel. ``output`` + (``DataContainer.save_data``) may copy each saved dataset to the host once, nothing more. + """ + + integrate: TransferCounter = field(default_factory=TransferCounter) + diagnostics: TransferCounter = field(default_factory=TransferCounter) + output: TransferCounter = field(default_factory=TransferCounter) + # allowed device-to-host copies by `output`: one per saved dataset, and their bytes + output_allowed_copies: int = 0 + output_allowed_bytes: int = 0 + n_integrate_calls: int = 0 + n_diagnostics_calls: int = 0 + + def check(self, n_scalars: int): + report = ( + f"integrate:\n{self.integrate.report()}\n" + f"diagnostics:\n{self.diagnostics.report()}\n" + f"output:\n{self.output.report()}" + ) + assert self.n_integrate_calls == N_STEPS, report + + # 1. the propagators: nothing between host and device + assert all(e.kind not in HOST_DEVICE_KINDS for e in self.integrate.events), report + + # 2. diagnostics: no host-to-device copies, no host kernels, only scalar-sized downloads + for kind in ("to_device", "kernel_conversion", "fallback"): + assert self.diagnostics.count(kind) == 0, report + assert self.diagnostics.to_host <= n_scalars * self.n_diagnostics_calls, report + assert self.diagnostics.bytes_to_host <= 8 * n_scalars * self.n_diagnostics_calls, report + + # 3. output: each saved dataset once to the host + for kind in ("to_device", "kernel_conversion", "fallback"): + assert self.output.count(kind) == 0, report + assert self.output.to_host <= self.output_allowed_copies, report + assert self.output.bytes_to_host <= self.output_allowed_bytes, report + + +def instrument(sim: Simulation, monkeypatch) -> TransferBudget: + """Count the transfers of the time loop of ``sim.run()`` per phase. + + Counting starts once the HDF5 datasets are created, the last step of the setup in + :meth:`Simulation.run`. The model methods are patched on the instance, ``save_data`` on the class + (the ``DataContainer`` is created inside ``run``). + """ + budget = TransferBudget() + in_loop = [False] + model = sim.model + + def counted(fn, phase: str, calls: str | None = None): + def wrapper(*args, **kwargs): + if not in_loop[0]: + return fn(*args, **kwargs) + with count_transfers() as counter: + out = fn(*args, **kwargs) + getattr(budget, phase).events.extend(counter.events) + if calls is not None: + setattr(budget, calls, getattr(budget, calls) + 1) + return out + + return wrapper + + original_init_datasets = sim._initialize_hdf5_datasets + + def init_datasets(*args, **kwargs): + out = original_init_datasets(*args, **kwargs) + in_loop[0] = True + return out + + monkeypatch.setattr(sim, "_initialize_hdf5_datasets", init_datasets) + monkeypatch.setattr(model, "integrate", counted(model.integrate, "integrate", "n_integrate_calls")) + monkeypatch.setattr( + model, + "update_scalar_quantities", + counted(model.update_scalar_quantities, "diagnostics", "n_diagnostics_calls"), + ) + for name in ("update_markers_to_be_saved", "update_distr_functions"): + monkeypatch.setattr(model, name, counted(getattr(model, name), "diagnostics")) + + original_save_data = DataContainer.save_data + + def save_data(self, keys=None): + if not in_loop[0]: + return original_save_data(self, keys=keys) + for key in self._dset_dict if keys is None else keys: + val = self._dset_dict[key] + budget.output_allowed_copies += 1 + budget.output_allowed_bytes += int(getattr(val, "nbytes", 8)) + with count_transfers() as counter: + original_save_data(self, keys=keys) + budget.output.events.extend(counter.events) + + monkeypatch.setattr(DataContainer, "save_data", save_data) + return budget + + +@dataclass +class RunResult: + """Host copies of what is compared between the backends.""" + + scalars: dict # name -> time series (rank 0 only, else empty) + e_coeffs: np.ndarray # local electric-field coefficients at the end + markers: np.ndarray # all valid markers (gathered), sorted by marker ID + n_scalars: int + budget: TransferBudget + + +def run(backend: str, out_folders: str, monkeypatch) -> RunResult: + """Run the simulation on ``backend`` and copy the results to the host.""" + with xp.use_backend(backend): + sim = make_simulation(out_folders, sim_folder=backend) + budget = instrument(sim, monkeypatch) + sim.run() + monkeypatch.undo() + + model = sim.model + comm = MPI.COMM_WORLD + + e_coeffs = np.asarray(xp.to_numpy(model.em_fields.e_field.spline.vector.toarray())) + + particles = model.kinetic_ions.var.particles + local = np.asarray(xp.to_numpy(particles.markers_wo_holes_and_ghost)) + parts = comm.allgather(local) if comm.Get_size() > 1 else [local] + markers = np.concatenate(parts, axis=0) + markers = markers[np.argsort(markers[:, -1])] + + scalars = {} + if comm.Get_rank() == 0: + with h5py.File(os.path.join(sim.env.path_out, "data", "data_proc0.hdf5"), "r") as f: + for key in model.scalars.dct: + scalars[key] = f["scalar"][key][()].ravel() + + return RunResult( + scalars=scalars, e_coeffs=e_coeffs, markers=markers, n_scalars=len(model.scalars.dct), budget=budget + ) + + +def check_reference(ref: RunResult): + """Physics sanity of the NumPy run: the comparison is only meaningful if the reference is right.""" + assert ref.markers.shape[0] == PPC * int(np.prod(NUM_ELEMENTS)) + if MPI.COMM_WORLD.Get_rank() == 0: + en_E, en_tot = ref.scalars["en_E"], ref.scalars["en_tot"] + assert en_tot.size == N_STEPS + 1 + # energy moves between field and particles ... + assert abs(en_E[-1] - en_E[0]) > 1e-3 * en_tot[0], en_E + # ... and the Crank-Nicolson step conserves the total + drift = np.max(np.abs(en_tot - en_tot[0])) / en_tot[0] + assert drift < 1e-6, f"{en_tot = }, {drift = }" + + +def compare_runs(ref: RunResult, gpu: RunResult): + """Compare the CuPy run against the NumPy run. + + Tolerances: + + * Velocities do not change (no E0, no B0) and positions are pushed explicitly + (``eta += dt * v / r1``, mod 1): the same floating-point operations on both backends, up to + fused multiply-adds on the GPU, so ``rtol = atol = 1e-13``. + * Weights, field coefficients and energies come out of ``EfieldWeightsCoupling``: a + Crank-Nicolson step whose Schur complement is solved by pcg to a relative residual of + ``SOLVER_TOL = 1e-12``, and whose matrix and vector are accumulated over markers (on the GPU + with atomic adds, i.e. in a different order). On this grid of 8 elements the preconditioned + system has a condition number of order 10-100, so the two solutions differ by at most + ~1e-10 relative to their norm after 3 steps; the tolerance ``1e-8`` (relative to the largest + entry) leaves a factor 100 against that and is still far below the change per step (~1e-2). + """ + rtol = 1e-8 + np.testing.assert_array_equal(gpu.markers[:, -1], ref.markers[:, -1]) # same markers + np.testing.assert_allclose(gpu.markers[:, 3:6], ref.markers[:, 3:6], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(gpu.markers[:, 0:3], ref.markers[:, 0:3], rtol=1e-13, atol=1e-13) + + w_ref, w_gpu = ref.markers[:, 6], gpu.markers[:, 6] + np.testing.assert_allclose(w_gpu, w_ref, rtol=0, atol=rtol * np.max(np.abs(w_ref))) + + assert gpu.e_coeffs.shape == ref.e_coeffs.shape + np.testing.assert_allclose(gpu.e_coeffs, ref.e_coeffs, rtol=0, atol=rtol * np.max(np.abs(ref.e_coeffs))) + + if MPI.COMM_WORLD.Get_rank() == 0: + for key, series in ref.scalars.items(): + np.testing.assert_allclose(gpu.scalars[key], series, rtol=rtol, err_msg=key) + + +@pytest.fixture +def out_folders(tmp_path_factory): + """One output folder for all ranks (the one of rank 0; ``tmp_path`` differs per rank).""" + comm = MPI.COMM_WORLD + path = str(tmp_path_factory.mktemp("gpu_e2e_LinearVlasovAmpereOneSpecies")) if comm.Get_rank() == 0 else None + return comm.bcast(path, root=0) + + +def test_numpy_reference(out_folders, monkeypatch): + """The NumPy run that the GPU run is compared against; on NumPy nothing can be transferred.""" + ref = run("numpy", out_folders, monkeypatch) + check_reference(ref) + ref.budget.check(n_scalars=ref.n_scalars) + assert ref.budget.integrate.total == ref.budget.diagnostics.total == ref.budget.output.total == 0 + + +@requires_cupy +def test_cupy_matches_numpy_without_transfers(out_folders, monkeypatch): + """The CuPy run agrees with the NumPy run, and its time loop stays on the device.""" + ref = run("numpy", out_folders, monkeypatch) + check_reference(ref) + gpu = run("cupy", out_folders, monkeypatch) + compare_runs(ref, gpu) + gpu.budget.check(n_scalars=gpu.n_scalars) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-ra"]) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1d08e9116..85f8835db 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1249,10 +1249,13 @@ def draw_markers( # 2. plain sobol numbers with skip of first 1000 numbers elif self.loading == "sobol_standard": - self.phasespace_coords = sobol_seq.i4_sobol_generate( - 3 + self.vdim, - n_mks_load_loc, - 1000 + (n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank], + # generated on the host, then copied to the device (setup only) + self.phasespace_coords = xp.to_cunumpy( + sobol_seq.i4_sobol_generate( + 3 + self.vdim, + int(n_mks_load_loc), + 1000 + int((n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank]), + ) ) # 3. symmetric sobol numbers in all 6 dimensions with skip of first 1000 numbers @@ -1263,10 +1266,12 @@ def draw_markers( # each sobol point yields 64 symmetric markers; round up and truncate the last group n_sobol = -(-self.n_mks_load // 64) - temp_markers = sobol_seq.i4_sobol_generate( - 3 + self.vdim, - n_sobol[self._mpi_rank], - 1000 + (xp.cumsum(n_sobol) - n_sobol)[self._mpi_rank], + temp_markers = xp.to_cunumpy( + sobol_seq.i4_sobol_generate( + 3 + self.vdim, + int(n_sobol[self._mpi_rank]), + 1000 + int((xp.cumsum(n_sobol) - n_sobol)[self._mpi_rank]), + ) ) temp_symmetric = xp.zeros((64 * n_sobol[self._mpi_rank], 3 + self.vdim), dtype=float) diff --git a/src/struphy/pic/sobol_seq.py b/src/struphy/pic/sobol_seq.py index 41a5e1f9c..15e92ff51 100644 --- a/src/struphy/pic/sobol_seq.py +++ b/src/struphy/pic/sobol_seq.py @@ -13,13 +13,16 @@ PYTHON versions by Corrado Chisari Original code is available from http://people.sc.fsu.edu/~jburkardt/py_src/sobol/sobol.html + + Runs on the host (NumPy) on every backend: it is a loop over scalars, and the caller copies the result + to the device. """ from __future__ import division import logging -import cunumpy as xp +import numpy as np logger = logging.getLogger("struphy") @@ -62,7 +65,7 @@ def i4_bit_hi1(n): Output, integer BIT, the number of bits base 2. """ - i = xp.floor(n) + i = np.floor(n) bit = 0 while i > 0: bit += 1 @@ -107,7 +110,7 @@ def i4_bit_lo0(n): Output, integer BIT, the position of the low 1 bit. """ bit = 1 - i = xp.floor(n) + i = np.floor(n) while i != 2 * (i // 2): bit += 1 i //= 2 @@ -125,7 +128,7 @@ def i4_sobol_generate(dim_num, n, skip=1): Output, real R(M,N), the points. """ - r = xp.full((n, dim_num), xp.nan) + r = np.full((n, dim_num), np.nan) for j in range(n): seed = j + skip r[j, 0:dim_num], next_seed = i4_sobol(dim_num, seed) @@ -226,8 +229,8 @@ def i4_sobol(dim_num, seed): seed_save = -1 # Initialize (part of) V. - v = xp.zeros((dim_max, log_max)) - v[0:40, 0] = xp.transpose( + v = np.zeros((dim_max, log_max)) + v[0:40, 0] = np.transpose( [ 1, 1, @@ -272,7 +275,7 @@ def i4_sobol(dim_num, seed): ], ) - v[2:40, 1] = xp.transpose( + v[2:40, 1] = np.transpose( [ 1, 3, @@ -315,7 +318,7 @@ def i4_sobol(dim_num, seed): ], ) - v[3:40, 2] = xp.transpose( + v[3:40, 2] = np.transpose( [ 7, 5, @@ -357,7 +360,7 @@ def i4_sobol(dim_num, seed): ], ) - v[5:40, 3] = xp.transpose( + v[5:40, 3] = np.transpose( [ 1, 7, @@ -397,7 +400,7 @@ def i4_sobol(dim_num, seed): ], ) - v[7:40, 4] = xp.transpose( + v[7:40, 4] = np.transpose( [ 9, 3, @@ -435,15 +438,15 @@ def i4_sobol(dim_num, seed): ], ) - v[13:40, 5] = xp.transpose( + v[13:40, 5] = np.transpose( [37, 33, 7, 5, 11, 39, 63, 27, 17, 15, 23, 29, 3, 21, 13, 31, 25, 9, 49, 33, 19, 29, 11, 19, 27, 15, 25], ) - v[19:40, 6] = xp.transpose( + v[19:40, 6] = np.transpose( [13, 33, 115, 41, 79, 17, 29, 119, 75, 73, 105, 7, 59, 65, 21, 3, 113, 61, 89, 45, 107], ) - v[37:40, 7] = xp.transpose([7, 23, 39]) + v[37:40, 7] = np.transpose([7, 23, 39]) # Set POLY. poly = [ @@ -522,7 +525,7 @@ def i4_sobol(dim_num, seed): # Expand this bit pattern to separate components of the logical array INCLUD. j = poly[i - 1] - includ = xp.zeros(m) + includ = np.zeros(m) for k in range(m, 0, -1): j2 = j // 2 includ[k - 1] = j != 2 * j2 @@ -536,7 +539,7 @@ def i4_sobol(dim_num, seed): for k in range(1, m + 1): l *= 2 if includ[k - 1]: - newv = xp.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1])) + newv = np.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1])) v[i - 1, j - 1] = newv # Multiply columns of V by appropriate power of 2. @@ -547,16 +550,16 @@ def i4_sobol(dim_num, seed): # RECIPD is 1/(common denominator of the elements in V). recipd = 1.0 / (2 * l) - lastq = xp.zeros(dim_num) + lastq = np.zeros(dim_num) - seed = int(xp.floor(seed)) + seed = int(np.floor(seed)) if seed < 0: seed = 0 l = 1 if seed == 0: - lastq = xp.zeros(dim_num) + lastq = np.zeros(dim_num) elif seed == seed_save + 1: # Find the position of the right-hand zero in SEED. @@ -564,12 +567,12 @@ def i4_sobol(dim_num, seed): elif seed <= seed_save: seed_save = 0 - lastq = xp.zeros(dim_num) + lastq = np.zeros(dim_num) for seed_temp in range(int(seed_save), int(seed)): l = i4_bit_lo0(seed_temp) for i in range(1, dim_num + 1): - lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) l = i4_bit_lo0(seed) @@ -577,7 +580,7 @@ def i4_sobol(dim_num, seed): for seed_temp in range(int(seed_save + 1), int(seed)): l = i4_bit_lo0(seed_temp) for i in range(1, dim_num + 1): - lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) l = i4_bit_lo0(seed) @@ -590,10 +593,10 @@ def i4_sobol(dim_num, seed): return # Calculate the new components of QUASI. - quasi = xp.zeros(dim_num) + quasi = np.zeros(dim_num) for i in range(1, dim_num + 1): quasi[i - 1] = lastq[i - 1] * recipd - lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) seed_save = seed seed += 1 @@ -643,11 +646,11 @@ def i4_uniform(a, b, seed): logger.info("I4_UNIFORM - Fatal error!") logger.info(" Input SEED = 0!") - seed = xp.floor(seed) + seed = np.floor(seed) a = round(a) b = round(b) - seed = xp.mod(seed, 2147483647) + seed = np.mod(seed, 2147483647) if seed < 0: seed += 2147483647 @@ -701,7 +704,7 @@ def prime_ge(n): Output, integer P, the smallest prime number that is greater than or equal to N. """ - p = max(xp.ceil(n), 2) + p = max(np.ceil(n), 2) while not is_prime(p): p += 1 @@ -725,7 +728,7 @@ def is_prime(n): return False # All primes >3 are of the form 6n+1 or 6n+5 (6n, 6n+2, 6n+4 are 2-divisible, 6n+3 is 3-divisible) p = 5 - root = int(xp.ceil(xp.sqrt(n))) + root = int(np.ceil(np.sqrt(n))) while p <= root: if n % p == 0 or n % (p + 2) == 0: return False From 85bc754e4f441eeca5118ba983b4a86e2de2d012 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 8 Oct 2026 11:05:51 +0200 Subject: [PATCH 07/12] Update feectools to latest commit in inner-on-device --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 6d888061f..fce01b827 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 6d888061f6a9d824df00633dcf475c5593f69dc1 +Subproject commit fce01b827d6fa95d42699d1b4400092e10888c23 From e40e4d7f1b785eeb06f248758981495246ebc207 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 8 Oct 2026 13:57:59 +0200 Subject: [PATCH 08/12] End-to-end GPU test of VlasovAmpereOneSpecies on 1 and 2 ranks - New test_gpu_VlasovAmpereOneSpecies.py: a small Landau-damping run (full-f weights, e1_v1 binning plot) on NumPy and on CuPy; compares energies, E and phi coefficients and markers, and checks the transfer budget of the time loop. Passes on the fake CuPy with emulated CUDA launches on 1 and 2 ranks. - Shared helpers in models/tests/gpu_e2e.py (transfer budget, instrumentation, gathering the markers, device context); test_gpu_LinearVlasovAmpereOneSpecies.py uses them. On the fake CuPy every CUDA launch is emulated and CudaKernel.compile is skipped, so both tests now also run without a GPU (CUNUMPY_FAKE_CUPY=1, or a serial child process). - Transfer budget of model.integrate: only scalar downloads (the residual norm of each Krylov iteration, which feectools copies for its convergence test), no arrays or host kernels. - Particles: probe whether MPI is CUDA-aware when the marker array is allocated (setup), not on the first marker exchange inside the first time step. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 2 +- src/struphy/models/tests/gpu_e2e.py | 255 ++++++++++++++++++ .../test_gpu_LinearVlasovAmpereOneSpecies.py | 194 ++----------- .../tests/test_gpu_VlasovAmpereOneSpecies.py | 217 +++++++++++++++ src/struphy/pic/base.py | 5 + 5 files changed, 506 insertions(+), 167 deletions(-) create mode 100644 src/struphy/models/tests/gpu_e2e.py create mode 100644 src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 2e775ca7f..91806fc89 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -31,7 +31,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **FEEC assembly kernels in kernel folders** (#411) — the 16 assembly kernels of `feec/mass.py`, `feec/boundary_mass.py` and `feec/basis_projection_ops.py` move from `feec/mass_kernels.py` and `feec/basis_projection_kernels.py` into one folder each under `feec/kernels/`; pure restructuring, ready for CUDA ports (see [FEEC assembly kernels](#feec-assembly-kernels-411)). - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). -- [ ] **End-to-end test of `LinearVlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` runs a small Cuboid case (320 Sobol markers, 8 elements, 3 steps) on NumPy and CuPy, compares energies, E-field coefficients and markers, and counts the transfers of the time loop: none in `model.integrate`, scalar-sized downloads in the diagnostics, one download per saved dataset in the output. GPU-only; fails until the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere` CUDA kernel, `MassMatrixPreconditioner` on CuPy and device inner products are merged and released. +- [ ] **End-to-end tests of `LinearVlasovAmpereOneSpecies` and `VlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` and `test_gpu_VlasovAmpereOneSpecies.py` (shared helpers in `models/tests/gpu_e2e.py`) run a small Cuboid case (320 Sobol markers, 8 elements, 3 steps; for `VlasovAmpereOneSpecies` Landau damping with full-f weights and an `e1_v1` binning plot) on NumPy and CuPy, compare energies, field coefficients and markers, and count the transfers of the time loop: in `model.integrate` only the 8-byte residual norm of each Krylov iteration (feectools' convergence test), scalar-sized downloads in the diagnostics, one download per saved dataset in the output. They run on a GPU, or on the fake CuPy with every CUDA launch emulated on the CPU (`CUNUMPY_FAKE_CUPY=1`, also under `mpirun -n 2`; `CudaKernel.compile` is skipped there); without either, `test_fake_cupy_matches_numpy` runs the comparison in a serial child process. Both pass on the fake CuPy on 1 and 2 ranks (not yet on a GPU). Needs the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere`/`vlasov_maxwell` CUDA kernels, `MassMatrixPreconditioner` on CuPy and device inner products, and a feectools release with them. - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py new file mode 100644 index 000000000..332a48492 --- /dev/null +++ b/src/struphy/models/tests/gpu_e2e.py @@ -0,0 +1,255 @@ +"""Shared helpers of the end-to-end GPU tests of the models (issue #689). + +A test runs a small simulation once on the NumPy backend and once on the CuPy backend, compares the two, and checks +that the time loop of the CuPy run makes no host/device transfers through cunumpy, apart from an explicit budget for +the diagnostics and the output (:class:`TransferBudget`). + +The CuPy run needs a GPU, or cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), on which every CUDA launch is emulated +on the CPU (:func:`device_context`). Without either, a test can still run the comparison in a serial child process +on the fake CuPy (:func:`run_fake_cupy_child`). + +Only transfers made through cunumpy are counted (``to_numpy``, ``to_cupy``, host fallbacks of kernels, ...). +Implicit synchronizations such as ``float(device_scalar)`` are not seen; use ``nsys`` for those. +""" + +import contextlib +import os +from collections.abc import Callable +from dataclasses import dataclass, field + +import cunumpy as xp +import h5py +import numpy as np +import pytest +from cunumpy.kernel_testing import fake_cupy_active +from cunumpy.profiling import TransferCounter, count_transfers +from maybempi import MPI + +from struphy.io.output_handling import DataContainer + +# transfer kinds that move data between host and device (device_copy is device-only and allowed) +HOST_DEVICE_KINDS = ("to_host", "to_device", "kernel_conversion", "fallback") + + +def cupy_backend_available() -> bool: + """Whether the CuPy backend can run here: a GPU, or the fake CuPy (with emulated launches).""" + return fake_cupy_active() or xp.cupy_available() + + +requires_cupy_backend = pytest.mark.skipif( + not cupy_backend_available(), reason="needs a GPU or the fake CuPy (CUNUMPY_FAKE_CUPY=1)" +) + + +@contextlib.contextmanager +def device_context(): + """Context for the CuPy run: on the fake CuPy every CUDA launch is emulated on the CPU, on a GPU nothing. + + On the fake CuPy, :meth:`CudaKernel.compile` (called by ``Simulation.compile_cuda_kernels`` before the time loop) + does nothing: the fake CuPy cannot compile CUDA, and the emulated launches compile their own C++ version. + ``compile_cuda_kernels`` still checks that every kernel of the time loop has a CUDA version. + """ + if not fake_cupy_active(): + yield + return + from unittest import mock + + from cunumpy.kernel_testing import emulated_launches + from cunumpy.kernels import CudaKernel + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + with ( + emulated_launches(options=EMULATION_OPTIONS), + mock.patch.object(CudaKernel, "compile", lambda self, **kw: None), + ): + yield + + +@dataclass +class TransferBudget: + """Transfers recorded per phase of the time loop (the setup before it is not counted). + + ``integrate`` (the propagators) may only copy scalars to the host: the convergence test of the Krylov solvers + reads the 8-byte residual norm once per iteration (feectools keeps everything else of the solve on the device). + No arrays, no host-to-device copies, no host kernels. ``diagnostics`` (scalars, saved markers, binned + distribution functions) may copy small results to the host: at most one scalar per tracked scalar and call, + never a host-to-device copy or a host kernel. ``output`` (``DataContainer.save_data``) may copy each saved + dataset to the host once, nothing more. + """ + + n_steps: int + integrate: TransferCounter = field(default_factory=TransferCounter) + diagnostics: TransferCounter = field(default_factory=TransferCounter) + output: TransferCounter = field(default_factory=TransferCounter) + # allowed device-to-host copies by `output`: one per saved dataset, and their bytes + output_allowed_copies: int = 0 + output_allowed_bytes: int = 0 + n_integrate_calls: int = 0 + n_diagnostics_calls: int = 0 + + def report(self) -> str: + return ( + f"integrate:\n{self.integrate.report()}\n" + f"diagnostics:\n{self.diagnostics.report()}\n" + f"output:\n{self.output.report()}" + ) + + def check(self, n_scalars: int): + report = self.report() + assert self.n_integrate_calls == self.n_steps, report + + # 1. the propagators: only scalars to the host (solver convergence tests), nothing else + for e in self.integrate.events: + if e.kind == "to_host": + assert e.nbytes is not None and e.nbytes <= 8, report + else: + assert e.kind not in HOST_DEVICE_KINDS, report + + # 2. diagnostics: no host-to-device copies, no host kernels, only scalar-sized downloads + for kind in ("to_device", "kernel_conversion", "fallback"): + assert self.diagnostics.count(kind) == 0, report + assert self.diagnostics.to_host <= n_scalars * self.n_diagnostics_calls, report + assert self.diagnostics.bytes_to_host <= 8 * n_scalars * self.n_diagnostics_calls, report + + # 3. output: each saved dataset once to the host + for kind in ("to_device", "kernel_conversion", "fallback"): + assert self.output.count(kind) == 0, report + assert self.output.to_host <= self.output_allowed_copies, report + assert self.output.bytes_to_host <= self.output_allowed_bytes, report + + +class _Patcher: + """Minimal ``monkeypatch``: set attributes, and restore them in reverse order.""" + + def __init__(self): + self._saved = [] + + def setattr(self, obj, name, value): + self._saved.append((obj, name, getattr(obj, name))) + setattr(obj, name, value) + + def undo(self): + while self._saved: + obj, name, value = self._saved.pop() + setattr(obj, name, value) + + +def instrument(sim, patcher, n_steps: int) -> TransferBudget: + """Count the transfers of the time loop of ``sim.run()`` per phase. + + Counting starts once the HDF5 datasets are created, the last step of the setup in :meth:`Simulation.run`. The + model methods are patched on the instance, ``save_data`` on the class (the ``DataContainer`` is created inside + ``run``). + """ + budget = TransferBudget(n_steps=n_steps) + in_loop = [False] + model = sim.model + + def counted(fn, phase: str, calls: str | None = None): + def wrapper(*args, **kwargs): + if not in_loop[0]: + return fn(*args, **kwargs) + with count_transfers() as counter: + out = fn(*args, **kwargs) + getattr(budget, phase).events.extend(counter.events) + if calls is not None: + setattr(budget, calls, getattr(budget, calls) + 1) + return out + + return wrapper + + original_init_datasets = sim._initialize_hdf5_datasets + + def init_datasets(*args, **kwargs): + out = original_init_datasets(*args, **kwargs) + in_loop[0] = True + return out + + patcher.setattr(sim, "_initialize_hdf5_datasets", init_datasets) + patcher.setattr(model, "integrate", counted(model.integrate, "integrate", "n_integrate_calls")) + patcher.setattr( + model, + "update_scalar_quantities", + counted(model.update_scalar_quantities, "diagnostics", "n_diagnostics_calls"), + ) + for name in ("update_markers_to_be_saved", "update_distr_functions"): + patcher.setattr(model, name, counted(getattr(model, name), "diagnostics")) + + original_save_data = DataContainer.save_data + + def save_data(self, keys=None): + if not in_loop[0]: + return original_save_data(self, keys=keys) + for key in self._dset_dict if keys is None else keys: + val = self._dset_dict[key] + budget.output_allowed_copies += 1 + budget.output_allowed_bytes += int(getattr(val, "nbytes", 8)) + with count_transfers() as counter: + original_save_data(self, keys=keys) + budget.output.events.extend(counter.events) + + patcher.setattr(DataContainer, "save_data", save_data) + return budget + + +@dataclass +class RunResult: + """Host copies of what is compared between the backends.""" + + scalars: dict # name -> time series (rank 0 only, else empty) + fields: dict # FEEC variable name -> local coefficients at the end + markers: np.ndarray # all valid markers (gathered from all ranks), sorted by marker ID + n_scalars: int + budget: TransferBudget + + +def gather_markers(particles) -> np.ndarray: + """The valid markers of all ranks on the host, sorted by marker ID (the last column).""" + comm = MPI.COMM_WORLD + local = np.asarray(xp.to_numpy(particles.markers_wo_holes_and_ghost)) + parts = comm.allgather(local) if comm.Get_size() > 1 else [local] + markers = np.concatenate(parts, axis=0) + return markers[np.argsort(markers[:, -1])] + + +def run( + backend: str, + make_simulation: Callable[[str, str], object], + out_folders: str, + n_steps: int, + fields: Callable[[object], dict], + particles: Callable[[object], object], +) -> RunResult: + """Run the simulation made by ``make_simulation(out_folders, sim_folder)`` on ``backend``; copy results to host. + + ``fields(model)`` returns the FEEC variables to compare (name -> ``FEECVariable``), ``particles(model)`` the + ``Particles`` object whose markers are compared. + """ + with xp.use_backend(backend), device_context() if backend == "cupy" else contextlib.nullcontext(): + sim = make_simulation(out_folders, backend) + patcher = _Patcher() + budget = instrument(sim, patcher, n_steps) + try: + sim.run() + finally: + patcher.undo() + + model = sim.model + coeffs = {name: np.asarray(xp.to_numpy(var.spline.vector.toarray())) for name, var in fields(model).items()} + markers = gather_markers(particles(model)) + + scalars = {} + if MPI.COMM_WORLD.Get_rank() == 0: + with h5py.File(os.path.join(sim.env.path_out, "data", "data_proc0.hdf5"), "r") as f: + for key in model.scalars.dct: + scalars[key] = f["scalar"][key][()].ravel() + + return RunResult(scalars=scalars, fields=coeffs, markers=markers, n_scalars=len(model.scalars.dct), budget=budget) + + +def shared_out_folder(tmp_path_factory, name: str) -> str: + """One output folder for all ranks (the one of rank 0; ``tmp_path`` differs per rank).""" + comm = MPI.COMM_WORLD + path = str(tmp_path_factory.mktemp(name)) if comm.Get_rank() == 0 else None + return comm.bcast(path, root=0) diff --git a/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py index 7808bde18..b6a9a6f51 100644 --- a/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py @@ -7,27 +7,21 @@ 1. both runs give the same energies, electric-field coefficients and marker weights/positions, within tolerances explained in :func:`compare_runs`, and 2. the time loop of the CuPy run does no host/device transfers through cunumpy, apart from an - explicit budget for the diagnostics and the output (see :class:`TransferBudget`). + explicit budget for the diagnostics and the output (see :class:`~struphy.models.tests.gpu_e2e.TransferBudget`). -The CuPy part needs a GPU (``requires_cupy``) and is skipped otherwise. It is MPI-compatible: under +The CuPy part needs a GPU, or cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), on which the CUDA launches are emulated +on the CPU, and is skipped otherwise. It is MPI-compatible: under ``mpirun -n 2 pytest --with-mpi`` both runs use the same domain decomposition, and the markers are gathered and compared by marker ID. -Only transfers made through cunumpy are counted (``to_numpy``, ``to_cupy``, host fallbacks of -kernels, ...). Implicit synchronizations such as ``float(device_scalar)`` in the convergence check -of the iterative solver are not seen; use ``nsys`` for those. +The shared parts (transfer budget, instrumentation, gathering the markers) are in +:mod:`struphy.models.tests.gpu_e2e`. """ import logging -import os -from dataclasses import dataclass, field -import cunumpy as xp -import h5py import numpy as np import pytest -from cunumpy.kernel_testing import requires_cupy -from cunumpy.profiling import TransferCounter, count_transfers from maybempi import MPI from struphy import ( @@ -45,9 +39,9 @@ maxwellians, perturbations, ) -from struphy.io.output_handling import DataContainer from struphy.linear_algebra.solver import SolverParameters from struphy.models import LinearVlasovAmpereOneSpecies +from struphy.models.tests.gpu_e2e import RunResult, requires_cupy_backend, run, shared_out_folder logger = logging.getLogger("struphy") @@ -59,9 +53,6 @@ N_STEPS = 3 SOLVER_TOL = 1e-12 -# transfer kinds that move data between host and device (device_copy is device-only and allowed) -HOST_DEVICE_KINDS = ("to_host", "to_device", "kernel_conversion", "fallback") - def make_simulation(out_folders: str, sim_folder: str) -> Simulation: """The model on a periodic 1D Cuboid, with the active cunumpy backend. @@ -115,145 +106,15 @@ def make_simulation(out_folders: str, sim_folder: str) -> Simulation: ) -@dataclass -class TransferBudget: - """Transfers recorded per phase of the time loop (setup before it is not counted). - - ``integrate`` (the propagators) must not transfer at all. ``diagnostics`` (scalars, saved - markers, binned distribution functions, coefficient extraction) may copy small results to the - host: at most a few scalars per call, never a host-to-device copy or a host kernel. ``output`` - (``DataContainer.save_data``) may copy each saved dataset to the host once, nothing more. - """ - - integrate: TransferCounter = field(default_factory=TransferCounter) - diagnostics: TransferCounter = field(default_factory=TransferCounter) - output: TransferCounter = field(default_factory=TransferCounter) - # allowed device-to-host copies by `output`: one per saved dataset, and their bytes - output_allowed_copies: int = 0 - output_allowed_bytes: int = 0 - n_integrate_calls: int = 0 - n_diagnostics_calls: int = 0 - - def check(self, n_scalars: int): - report = ( - f"integrate:\n{self.integrate.report()}\n" - f"diagnostics:\n{self.diagnostics.report()}\n" - f"output:\n{self.output.report()}" - ) - assert self.n_integrate_calls == N_STEPS, report - - # 1. the propagators: nothing between host and device - assert all(e.kind not in HOST_DEVICE_KINDS for e in self.integrate.events), report - - # 2. diagnostics: no host-to-device copies, no host kernels, only scalar-sized downloads - for kind in ("to_device", "kernel_conversion", "fallback"): - assert self.diagnostics.count(kind) == 0, report - assert self.diagnostics.to_host <= n_scalars * self.n_diagnostics_calls, report - assert self.diagnostics.bytes_to_host <= 8 * n_scalars * self.n_diagnostics_calls, report - - # 3. output: each saved dataset once to the host - for kind in ("to_device", "kernel_conversion", "fallback"): - assert self.output.count(kind) == 0, report - assert self.output.to_host <= self.output_allowed_copies, report - assert self.output.bytes_to_host <= self.output_allowed_bytes, report - - -def instrument(sim: Simulation, monkeypatch) -> TransferBudget: - """Count the transfers of the time loop of ``sim.run()`` per phase. - - Counting starts once the HDF5 datasets are created, the last step of the setup in - :meth:`Simulation.run`. The model methods are patched on the instance, ``save_data`` on the class - (the ``DataContainer`` is created inside ``run``). - """ - budget = TransferBudget() - in_loop = [False] - model = sim.model - - def counted(fn, phase: str, calls: str | None = None): - def wrapper(*args, **kwargs): - if not in_loop[0]: - return fn(*args, **kwargs) - with count_transfers() as counter: - out = fn(*args, **kwargs) - getattr(budget, phase).events.extend(counter.events) - if calls is not None: - setattr(budget, calls, getattr(budget, calls) + 1) - return out - - return wrapper - - original_init_datasets = sim._initialize_hdf5_datasets - - def init_datasets(*args, **kwargs): - out = original_init_datasets(*args, **kwargs) - in_loop[0] = True - return out - - monkeypatch.setattr(sim, "_initialize_hdf5_datasets", init_datasets) - monkeypatch.setattr(model, "integrate", counted(model.integrate, "integrate", "n_integrate_calls")) - monkeypatch.setattr( - model, - "update_scalar_quantities", - counted(model.update_scalar_quantities, "diagnostics", "n_diagnostics_calls"), - ) - for name in ("update_markers_to_be_saved", "update_distr_functions"): - monkeypatch.setattr(model, name, counted(getattr(model, name), "diagnostics")) - - original_save_data = DataContainer.save_data - - def save_data(self, keys=None): - if not in_loop[0]: - return original_save_data(self, keys=keys) - for key in self._dset_dict if keys is None else keys: - val = self._dset_dict[key] - budget.output_allowed_copies += 1 - budget.output_allowed_bytes += int(getattr(val, "nbytes", 8)) - with count_transfers() as counter: - original_save_data(self, keys=keys) - budget.output.events.extend(counter.events) - - monkeypatch.setattr(DataContainer, "save_data", save_data) - return budget - - -@dataclass -class RunResult: - """Host copies of what is compared between the backends.""" - - scalars: dict # name -> time series (rank 0 only, else empty) - e_coeffs: np.ndarray # local electric-field coefficients at the end - markers: np.ndarray # all valid markers (gathered), sorted by marker ID - n_scalars: int - budget: TransferBudget - - -def run(backend: str, out_folders: str, monkeypatch) -> RunResult: +def run_model(backend: str, out_folders: str) -> RunResult: """Run the simulation on ``backend`` and copy the results to the host.""" - with xp.use_backend(backend): - sim = make_simulation(out_folders, sim_folder=backend) - budget = instrument(sim, monkeypatch) - sim.run() - monkeypatch.undo() - - model = sim.model - comm = MPI.COMM_WORLD - - e_coeffs = np.asarray(xp.to_numpy(model.em_fields.e_field.spline.vector.toarray())) - - particles = model.kinetic_ions.var.particles - local = np.asarray(xp.to_numpy(particles.markers_wo_holes_and_ghost)) - parts = comm.allgather(local) if comm.Get_size() > 1 else [local] - markers = np.concatenate(parts, axis=0) - markers = markers[np.argsort(markers[:, -1])] - - scalars = {} - if comm.Get_rank() == 0: - with h5py.File(os.path.join(sim.env.path_out, "data", "data_proc0.hdf5"), "r") as f: - for key in model.scalars.dct: - scalars[key] = f["scalar"][key][()].ravel() - - return RunResult( - scalars=scalars, e_coeffs=e_coeffs, markers=markers, n_scalars=len(model.scalars.dct), budget=budget + return run( + backend, + make_simulation, + out_folders, + n_steps=N_STEPS, + fields=lambda model: {"e_field": model.em_fields.e_field}, + particles=lambda model: model.kinetic_ions.var.particles, ) @@ -294,8 +155,9 @@ def compare_runs(ref: RunResult, gpu: RunResult): w_ref, w_gpu = ref.markers[:, 6], gpu.markers[:, 6] np.testing.assert_allclose(w_gpu, w_ref, rtol=0, atol=rtol * np.max(np.abs(w_ref))) - assert gpu.e_coeffs.shape == ref.e_coeffs.shape - np.testing.assert_allclose(gpu.e_coeffs, ref.e_coeffs, rtol=0, atol=rtol * np.max(np.abs(ref.e_coeffs))) + e_ref, e_gpu = ref.fields["e_field"], gpu.fields["e_field"] + assert e_gpu.shape == e_ref.shape + np.testing.assert_allclose(e_gpu, e_ref, rtol=0, atol=rtol * np.max(np.abs(e_ref))) if MPI.COMM_WORLD.Get_rank() == 0: for key, series in ref.scalars.items(): @@ -304,26 +166,26 @@ def compare_runs(ref: RunResult, gpu: RunResult): @pytest.fixture def out_folders(tmp_path_factory): - """One output folder for all ranks (the one of rank 0; ``tmp_path`` differs per rank).""" - comm = MPI.COMM_WORLD - path = str(tmp_path_factory.mktemp("gpu_e2e_LinearVlasovAmpereOneSpecies")) if comm.Get_rank() == 0 else None - return comm.bcast(path, root=0) + return shared_out_folder(tmp_path_factory, "gpu_e2e_LinearVlasovAmpereOneSpecies") -def test_numpy_reference(out_folders, monkeypatch): +def test_numpy_reference(out_folders): """The NumPy run that the GPU run is compared against; on NumPy nothing can be transferred.""" - ref = run("numpy", out_folders, monkeypatch) + ref = run_model("numpy", out_folders) check_reference(ref) ref.budget.check(n_scalars=ref.n_scalars) assert ref.budget.integrate.total == ref.budget.diagnostics.total == ref.budget.output.total == 0 -@requires_cupy -def test_cupy_matches_numpy_without_transfers(out_folders, monkeypatch): - """The CuPy run agrees with the NumPy run, and its time loop stays on the device.""" - ref = run("numpy", out_folders, monkeypatch) +@requires_cupy_backend +def test_cupy_matches_numpy_without_transfers(out_folders): + """The CuPy run agrees with the NumPy run, and its time loop stays on the device. + + On a GPU, or on the fake CuPy (``CUNUMPY_FAKE_CUPY=1``) with every CUDA launch emulated on the CPU. + """ + ref = run_model("numpy", out_folders) check_reference(ref) - gpu = run("cupy", out_folders, monkeypatch) + gpu = run_model("cupy", out_folders) compare_runs(ref, gpu) gpu.budget.check(n_scalars=gpu.n_scalars) diff --git a/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py new file mode 100644 index 000000000..c647847e1 --- /dev/null +++ b/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py @@ -0,0 +1,217 @@ +"""End-to-end GPU test of :class:`~struphy.models.VlasovAmpereOneSpecies` (issue #689). + +A small Landau-damping run (the setup of ``verification/test_verif_VlasovAmpereOneSpecies.py`` with fewer cells and +markers, full-f weights, and a larger perturbation so that energy moves within three steps) is done once on the NumPy +backend and once on the CuPy backend. The test checks that + +1. both runs give the same energies, electric-field coefficients and markers, within the tolerances explained in + :func:`compare_runs`, and +2. the time loop of the CuPy run does no host/device transfers through cunumpy, apart from an explicit budget for + the diagnostics and the output (:class:`~struphy.models.tests.gpu_e2e.TransferBudget`). + +A time step is ``PushEta`` followed by ``VlasovAmpereCoupling`` (no B0). The setup also runs the initial Poisson +solve (``charge_density_0form`` and the ``L2Projector`` with its ``MassMatrixPreconditioner``) on the device. + +The CuPy part needs a GPU, or cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), on which every CUDA launch is emulated on +the CPU. Without either, :func:`test_fake_cupy_matches_numpy` runs the comparison in a serial child process on the +fake CuPy. The tests are MPI-compatible (``mpirun -n 2 pytest --with-mpi``): both runs use the same domain +decomposition, and the markers are gathered and compared by marker ID. +""" + +import logging + +import numpy as np +import pytest +from maybempi import MPI + +from struphy import ( + BinningPlot, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.geometry.tests.test_domain import run_fake_cupy_child +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import VlasovAmpereOneSpecies +from struphy.models.tests.gpu_e2e import ( + RunResult, + cupy_backend_available, + requires_cupy_backend, + run, + shared_out_folder, +) + +logger = logging.getLogger("struphy") + +AMPLITUDE = 0.1 +R1 = 12.56 # box of length 4*pi, k = 0.5 +NUM_ELEMENTS = (8, 1, 1) +PPC = 40 # 320 markers +DT = 0.1 +N_STEPS = 3 +SOLVER_TOL = 1e-12 + + +def make_simulation(out_folders: str, sim_folder: str) -> Simulation: + """The model on a periodic 1D Cuboid, with the active cunumpy backend. + + Deterministic marker loading (``sobol_standard``), so that both backends start from the same markers: + ``pseudo_random`` draws from ``xp.random``, whose streams differ between NumPy and CuPy. No sorting boxes, since + ``SortingBoxes`` does not work on the CuPy backend yet (#726); the model does not need them. Full-f weights: with + the control variate, ``kinetic_energy`` at t=0 is computed with the full-f weights and afterwards with the delta-f + weights, so ``total_energy`` cannot be checked. + """ + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folders, sim_folder=sim_folder, save_restart=False) + time_opts = Time(dt=DT, Tend=round(N_STEPS * DT, 14)) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=PPC, loading="sobol_standard"), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters(), + # the binned distribution function of the verification test, so that the diagnostics run on the device too + saving_params=SavingParameters( + binning_plots=(BinningPlot(slice="e1_v1", n_bins=(16, 16), ranges=((0.0, 1.0), (-5.0, 5.0))),), + ), + bufsize=0.4, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + # tight solver tolerances, so that the backends differ by round-off and not by where pcg stops + model.propagators.coupling_va.options = model.propagators.coupling_va.Options( + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + # stab_eps=1e-6, not the default 1e-14: the Monte-Carlo net charge puts (net charge) / stab_eps into the constant + # mode of phi, so with 1e-14 phi is ~1e14 and E = -grad(phi) loses ~14 digits to cancellation (5% of E here, on + # either backend) + model.initial_poisson.options = model.initial_poisson.Options( + stab_mat="M0", + stab_eps=1e-6, + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + perturbation = perturbations.ModesCos(ls=(1,), amps=(AMPLITUDE,)) + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + return Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domains.Cuboid(r1=R1), + grid=grids.TensorProductGrid(num_elements=NUM_ELEMENTS), + derham_opts=DerhamOptions(degree=(3, 1, 1)), + ) + + +def run_model(backend: str, out_folders: str) -> RunResult: + """Run the simulation on ``backend`` and copy the results to the host.""" + return run( + backend, + make_simulation, + out_folders, + n_steps=N_STEPS, + fields=lambda model: {"e_field": model.em_fields.e_field, "phi": model.em_fields.phi}, + particles=lambda model: model.kinetic_ions.var.particles, + ) + + +def check_reference(ref: RunResult): + """Physics sanity of the NumPy run: the comparison is only meaningful if the reference is right.""" + assert ref.markers.shape[0] == PPC * int(np.prod(NUM_ELEMENTS)) + if MPI.COMM_WORLD.Get_rank() == 0: + en_E = ref.scalars["electric_energy"] + en_tot = ref.scalars["total_energy"] + assert en_tot.size == N_STEPS + 1 + # energy moves between field and particles ... + assert abs(en_E[-1] - en_E[0]) > 1e-2 * abs(en_E[0]), en_E + # ... and the Crank-Nicolson coupling conserves the total from step to step, to the solver tolerance. The value + # at t=0 differs by ~2e-4 relative (the initial diagnostic, not the time stepping), so it is checked loosely. + drift = np.max(np.abs(en_tot[1:] - en_tot[1])) / abs(en_tot[1]) + assert drift < 1e-9, f"{en_tot = }, {drift = }" + assert abs(en_tot[1] - en_tot[0]) < 1e-3 * abs(en_tot[0]), en_tot + + +def compare_runs(ref: RunResult, gpu: RunResult): + """Compare the CuPy run against the NumPy run. + + Velocities, the electric field and the energies come out of ``VlasovAmpereCoupling``, a Crank-Nicolson step whose + Schur complement is solved by pcg to a relative residual of ``SOLVER_TOL = 1e-12``, with the matrix and vector + accumulated over markers (on the GPU with atomic adds, i.e. in another order). On this grid of 8 elements the + preconditioned system has a condition number of order 10-100, so the two solutions differ by ~1e-10 relative to + their norm after three steps. Positions are pushed explicitly with these velocities, so they inherit + ``DT * 1e-10``. The tolerance ``1e-8`` (relative to the largest entry) leaves a factor 100 and is still far below + the change per step (~5e-2 in the electric energy). Full-f weights do not change, so they must agree to round-off. + """ + rtol = 1e-8 + np.testing.assert_array_equal(gpu.markers[:, -1], ref.markers[:, -1]) # same markers + + def close(a, b, name): + np.testing.assert_allclose(a, b, rtol=0, atol=rtol * max(np.max(np.abs(b)), 1e-300), err_msg=name) + + close(gpu.markers[:, 0:3], ref.markers[:, 0:3], "positions") + close(gpu.markers[:, 3:6], ref.markers[:, 3:6], "velocities") + np.testing.assert_allclose(gpu.markers[:, 6], ref.markers[:, 6], rtol=1e-13, atol=0, err_msg="weights") + for name, coeffs in ref.fields.items(): + assert gpu.fields[name].shape == coeffs.shape, name + close(gpu.fields[name], coeffs, name) + + if MPI.COMM_WORLD.Get_rank() == 0: + for key, series in ref.scalars.items(): + np.testing.assert_allclose(gpu.scalars[key], series, rtol=rtol, err_msg=key) + + +def check_on_cupy(out_folders: str): + """The whole comparison, for the fake-CuPy child process of :func:`test_fake_cupy_matches_numpy`.""" + ref = run_model("numpy", out_folders) + check_reference(ref) + gpu = run_model("cupy", out_folders) + compare_runs(ref, gpu) + gpu.budget.check(n_scalars=gpu.n_scalars) + + +@pytest.fixture +def out_folders(tmp_path_factory): + return shared_out_folder(tmp_path_factory, "gpu_e2e_VlasovAmpereOneSpecies") + + +def test_numpy_reference(out_folders): + """The NumPy run that the GPU run is compared against; on NumPy nothing can be transferred.""" + ref = run_model("numpy", out_folders) + check_reference(ref) + ref.budget.check(n_scalars=ref.n_scalars) + assert ref.budget.integrate.total == ref.budget.diagnostics.total == ref.budget.output.total == 0 + + +@requires_cupy_backend +def test_cupy_matches_numpy_without_transfers(out_folders): + """The CuPy run agrees with the NumPy run, and its time loop stays on the device. + + On a GPU, or on the fake CuPy (``CUNUMPY_FAKE_CUPY=1``) with every CUDA launch emulated on the CPU; under + ``mpirun -n 2`` with two ranks. + """ + check_on_cupy(out_folders) + + +@pytest.mark.skipif(cupy_backend_available(), reason="runs in this process (test_cupy_matches_numpy_without_transfers)") +def test_fake_cupy_matches_numpy(out_folders): + """Without a GPU and without ``CUNUMPY_FAKE_CUPY``: the comparison in a serial child process on the fake CuPy.""" + run_fake_cupy_child( + f"from struphy.models.tests.test_gpu_VlasovAmpereOneSpecies import check_on_cupy; check_on_cupy({out_folders!r})" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-ra"]) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 8d44a9e36..975de1c3e 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2593,6 +2593,11 @@ def _allocate_marker_array(self, dry_run: bool = False): self._bc_type, ) + # on the CuPy backend with several ranks, find out now (collectively, in the setup) whether MPI takes device + # buffers; otherwise the probe runs on the first marker exchange, inside the first time step + if xp.get_backend() == "cupy" and self.mpi_size > 1: + self._mpi_cuda_aware() + def _initialize_sorting_boxes(self): """Initializes the sorting boxes. From 4d4b51f35db87fee555282df8364bb38ed165db7 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 08:18:35 +0200 Subject: [PATCH 09/12] Diagonal scaling: include band entries that wrap onto the diagonal In a periodic direction with fewer points than the band width (2p + 1 > n, e.g. a single element in a 1D run), StencilMatrix.diagonal() returns only the entry at offset 0, while the entries at offsets k = m n also lie on the main diagonal. The diagonal D of the mass matrix was therefore wrong, and MassMatrixPreconditioner with diagonal_scaling=True (the default since #739) was no longer exact on the logical cube; in VlasovAmpereCoupling this broke energy conservation in 1D runs. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 30 +++++++++++++++-- src/struphy/feec/tests/test_mass_matrices.py | 34 ++++++++++++++++++++ 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index d64838dc2..a3a67c70b 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -318,7 +318,8 @@ def weight_1d(c: int, d: int) -> Callable | xp.ndarray: return _kronecker_approximation(mass_operator, bc, weight_1d) def _core_diagonal(self): - return self._mass_operator.matrix.diagonal() + # the matrix itself: _local_diagonal also adds the band entries that wrap onto the diagonal + return self._mass_operator.matrix @property def dim_reduce(self) -> int | None: @@ -765,7 +766,7 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: if isinstance(op, StencilDiagonalMatrix): return [op._data] if isinstance(op, StencilMatrix): - return [op.diagonal()._data] + return [_stencil_local_diagonal(op)] if isinstance(op, KroneckerStencilMatrix): # outer product of the diagonals of the (process-local) factors on the local rows diags = [] @@ -788,6 +789,31 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: raise NotImplementedError(f"Diagonal of {type(op).__name__} is not supported.") +def _stencil_local_diagonal(M: StencilMatrix) -> xp.ndarray: + """ + Local diagonal of a StencilMatrix, including the band entries that wrap onto the main diagonal. + + In a periodic direction with fewer points than the band width (``2p + 1 > n``, e.g. a single + element), the band entries at the offsets ``k = m n`` (``m != 0``) are on the main diagonal too; + :meth:`StencilMatrix.diagonal` returns only the entry at ``k = 0``. + """ + V = M.codomain + assert all(m == 1 for m in V.shifts) and all(m == 1 for m in M.domain.shifts), "Only shifts 1 are supported." + offsets = [ + [k for k in range(-p, p + 1) if k % n == 0] if periodic else [0] + for n, p, periodic in zip(V.npts, M.pads, V.periods) + ] + diag = M.diagonal()._data + if all(o == [0] for o in offsets): + return diag + idx = M._get_diagonal_indices() + rows, cols = idx[: V.ndim], idx[V.ndim :] + for ks in itertools.product(*offsets): + if any(ks): + diag = diag + M._data[rows + tuple(j + k for j, k in zip(cols, ks))] + return diag + + def _diagonal_operator(space: VectorSpace, diags: list[xp.ndarray]) -> LinearOperator: """Diagonal operator on ``space`` (Stencil- or BlockVectorSpace) with the given local diagonals.""" if isinstance(space, BlockVectorSpace): diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index 23af76b82..799646edf 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -1742,6 +1742,40 @@ def test_transpose_and_copy(): assert xp.allclose(A.T.toarray(), -A_arr, atol=1e-14) +@pytest.mark.parametrize( + "num_elements, degree", [((16, 1, 1), (3, 1, 1)), ((16, 2, 1), (2, 2, 1)), ((8, 6, 4), (2, 2, 2))] +) +@pytest.mark.parametrize("space", ["M0", "M1", "M2", "M3"]) +def test_mass_preconditioner_diagonal_scaling_periodic_exact(num_elements, degree, space): + """On the periodic logical cube the mass matrices are Kronecker products, so MassMatrixPreconditioner with + diagonal scaling inverts them exactly; also with fewer points than the band width in a direction (one element), + where band entries wrap onto the main diagonal.""" + + import numpy as np + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import MassMatrixPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = domains.Cuboid(r1=12.56) + derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=degree), comm=MPI.COMM_WORLD) + M = getattr(WeightedMassOperators(derham, domain), space) + + rng = np.random.default_rng(0) + v = M.domain.zeros() + for block in v.blocks if hasattr(v, "blocks") else [v]: + block._data[:] = rng.random(block._data.shape) + v.update_ghost_regions() + + w = MassMatrixPreconditioner(M, diagonal_scaling=True).dot(M.dot(v)) + err = ((w - v).inner(w - v) / v.inner(v)) ** 0.5 + assert err < 1e-12, f"{space}, {num_elements=}, {degree=}: |P M v - v| / |v| = {err}" + + if __name__ == "__main__": # test_mass( # num_elements=(32, 32, 32), From a0f1cf35558b5c985728fbad4b60915e85aa9bed Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 08:23:55 +0200 Subject: [PATCH 10/12] VlasovAmpereCoupling: update the ghost regions of the mid-point field before the push The pusher kernel reads the ghost regions of e_sum directly. They were only correct if the Schur solver happened to return its solution with valid ghost regions, which is no longer the case with the diagonal scaling of MassMatrixPreconditioner (default since #739): its last step is a StencilDiagonalMatrix.dot that writes only the owned entries. In 1D runs (one element in e2, e3) every evaluation reads ghost values, and the total energy was no longer conserved (drift ~1e-4 instead of ~1e-12). The other couplings already update the ghost regions before pushing. Co-Authored-By: Claude Opus 5.5 --- src/struphy/propagators/vlasov_ampere_coupling.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/struphy/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 0c8c76414..77a129ae0 100644 --- a/src/struphy/propagators/vlasov_ampere_coupling.py +++ b/src/struphy/propagators/vlasov_ampere_coupling.py @@ -214,6 +214,8 @@ def __call__(self, dt): self._e_sum += self.variables.e.spline.vector self._e_sum += self._e_tmp self._e_sum *= 0.5 + # the pusher kernel reads the ghost regions directly (the solver output may leave them stale) + self._e_sum.update_ghost_regions() # Update velocities self._pusher(dt) From 1d2e28e422e4665ba753c8f869d833197d2acad3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 13 Oct 2026 08:40:12 +0200 Subject: [PATCH 11/12] draw_markers: send the markers to their processes also without sorting boxes (#764) On more than one process, draw_markers MPI-sorted the markers only if sorting boxes were set up. Without boxes, the markers drawn on the whole domain stayed on the wrong process until the first push, and accumulations before it (e.g. the initial Poisson solve of VlasovAmpereOneSpecies) wrote outside the local stencil arrays (segfault on 2 ranks). Now draw_markers always MPI-sorts on several processes; sorting into boxes is unchanged. On one process the behaviour is unchanged. New MPI test test_draw_without_boxes_on_rank_domain (pseudo_random, sobol_standard, sobol_antithetic; sort=True/False): fails on 2 ranks without the fix. Closes #764. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 19 +++++--- src/struphy/pic/tests/test_draw_parallel.py | 50 +++++++++++++++++++++ 2 files changed, 63 insertions(+), 6 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 5e1d217b2..c5c629389 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1146,12 +1146,14 @@ def draw_markers( All needed parameters can be set in the parameter file, see :ref:`params_yml`. - An initial sorting will be performed if sort is given as True (default) and sorting_params were given to the init. + On several processes, the markers are always sent to the process of their domain after drawing + (:meth:`mpi_sort_markers`). They are also sorted into boxes if sort is True (default) and sorting + boxes were set up. Parameters ---------- sort : Bool - Wether to sort the particules in boxes after initial drawing (only if sorting params were passed) + Whether to sort the markers into boxes after initial drawing (only if sorting boxes were set up). """ import scipy.special as sp @@ -1405,10 +1407,15 @@ def draw_markers( assert xp.all(~self.holes[:n_mks_load_loc]) assert xp.all(self.holes[n_mks_load_loc:]) - if self._initialized_sorting and sort: - logger.info("\nSorting the markers after initial draw") - if self.mpi_comm is not None: - self.mpi_sort_markers() + # send the markers to the process of their domain: needed on several processes whether or not there are + # sorting boxes, since accumulations before the first push (e.g. an initial Poisson solve) assume it (#764) + sort_boxes = self._initialized_sorting and sort + if self.mpi_comm is not None and (self.mpi_size > 1 or sort_boxes): + logger.info("\nSending the markers to their processes after initial draw") + self.mpi_sort_markers() + + if sort_boxes: + logger.info("\nSorting the markers into boxes after initial draw") self.do_sort() logger.info("Done.") diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index ff18a3321..09020c588 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -177,6 +177,56 @@ def test_marker_ids(loading, Np=1000): assert n_at_origin == 0 +@pytest.mark.parametrize("loading", ["pseudo_random", "sobol_standard", "sobol_antithetic"]) +@pytest.mark.parametrize("sort", [True, False]) +def test_draw_without_boxes_on_rank_domain(loading, sort, Np=1000): + """Without sorting boxes, every marker is on the domain of its process right after drawing (#764), + for ``sort=True`` and ``sort=False`` (which only concerns the boxes), and no marker is lost.""" + + import cunumpy as xp + from maybempi import MPI + + from struphy import LoadingParameters, domains + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.pic.particles import Particles6D + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + + grid = TensorProductGrid(num_elements=[8, 5, 3]) + derham = Derham(grid, DerhamOptions(degree=[1, 1, 1]), comm=comm) + domain_decomp = (derham.domain_array, derham.domain_decomposition.nprocs) + + loading_params = LoadingParameters( + Np=Np, + seed=1234, + loading=loading, + moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), + ) + + # no sorting_params: no sorting boxes + particles = Particles6D( + comm_world=comm, + domain_decomp=domain_decomp, + loading_params=loading_params, + domain=domains.Cuboid(), + ) + assert particles.sorting_boxes is None + + particles.draw_markers(sort=sort) + + positions = particles.markers_wo_holes[:, :3] + on_rank = xp.all(particles._is_on_domain_of_rank(positions, particles.mpi_rank), axis=1) + n_off_rank = int(xp.sum(~on_rank)) + n_markers = positions.shape[0] + if particles.mpi_size > 1: + n_markers = particles.mpi_comm.allreduce(n_markers) + + assert n_off_rank == 0, f"rank {particles.mpi_rank}: {n_off_rank} markers off the process domain after drawing" + assert n_markers == Np + + if __name__ == "__main__": # test_draw([8, 9, 10], [2, 3, 4], [False, False, True], ['Cuboid', { # 'l1': 1., 'r1': 2., 'l2': 10., 'r2': 20., 'l3': 100., 'r3': 200.}]) From d356ad95f304fe25d1a79a3a64f62366d470edec Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 13 Oct 2026 08:40:13 +0200 Subject: [PATCH 12/12] Review fixes for #734: run the e2e tests in CI, keep the _core_diagonal contract - CI: the end-to-end GPU tests run in unit shard-2 (serial and 2 ranks, both workflows and `struphy test unit`) and on the fake CuPy on 1 and 2 ranks in test-PR-unit-mpi.yml (8 elements of degree 3 allow at most 2 ranks). - MassMatrixPreconditioner._core_diagonal returns a diagonal operator again (with the band entries that wrap onto the diagonal), as documented in the base class. - test_preconditioner_cupy: use cunumpy's emulated_launches and run_fake_cupy_child (rank 0 only) instead of local copies; test_mass_preconditioners_on_cupy emulates its launches on the fake CuPy instead of trying to compile CUDA. - FFTSolver.solve asserts host input; comment on the collective probe in _allocate_marker_array. - CUDA_STRATEGY.md: CI coverage, #764, host kernels left in the setup. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/reusable-scheduled.yml | 2 +- .github/workflows/reusable-unit-testing.yml | 2 +- .github/workflows/test-PR-unit-mpi.yml | 8 ++ CUDA_STRATEGY.md | 2 +- src/struphy/console/test.py | 3 + src/struphy/feec/preconditioner.py | 7 +- .../feec/tests/test_preconditioner_cupy.py | 78 ++++--------------- src/struphy/pic/base.py | 3 +- 8 files changed, 37 insertions(+), 68 deletions(-) diff --git a/.github/workflows/reusable-scheduled.yml b/.github/workflows/reusable-scheduled.yml index c24a03184..e4652bd23 100644 --- a/.github/workflows/reusable-scheduled.yml +++ b/.github/workflows/reusable-scheduled.yml @@ -107,7 +107,7 @@ jobs: TEST_FILES="bsplines/tests/ console/tests/ fields_background/tests/ geometry/tests/ initial/tests/ kinetic_background/tests/ linear_algebra/tests/" ;; shard-2) - TEST_FILES="propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/" + TEST_FILES="propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py" ;; shard-3) TEST_FILES="feec/tests/" diff --git a/.github/workflows/reusable-unit-testing.yml b/.github/workflows/reusable-unit-testing.yml index 2e8f59e7b..89944e67b 100644 --- a/.github/workflows/reusable-unit-testing.yml +++ b/.github/workflows/reusable-unit-testing.yml @@ -39,7 +39,7 @@ jobs: - shard: shard-1 test_files: bsplines/tests/ console/tests/ fields_background/tests/ geometry/tests/ initial/tests/ kinetic_background/tests/ linear_algebra/tests/ - shard: shard-2 - test_files: propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ + test_files: propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py - shard: shard-3 test_files: feec/tests/ - shard: shard-4 diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index 6a5b1c387..1743e82a9 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -88,3 +88,11 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi polar/tests/test_polar_cupy.py -k mpi_fake + + - name: Run the end-to-end GPU model tests on the fake CuPy + # 8 elements of degree 3: at most 2 ranks + if: matrix.n-procs <= 2 + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 61c874781..69f8866b5 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -31,7 +31,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **FEEC assembly kernels in kernel folders** (#411) — the 16 assembly kernels of `feec/mass.py`, `feec/boundary_mass.py` and `feec/basis_projection_ops.py` move from `feec/mass_kernels.py` and `feec/basis_projection_kernels.py` into one folder each under `feec/kernels/`; pure restructuring, ready for CUDA ports (see [FEEC assembly kernels](#feec-assembly-kernels-411)). - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). -- [ ] **End-to-end tests of `LinearVlasovAmpereOneSpecies` and `VlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` and `test_gpu_VlasovAmpereOneSpecies.py` (shared helpers in `models/tests/gpu_e2e.py`) run a small Cuboid case (320 Sobol markers, 8 elements, 3 steps; for `VlasovAmpereOneSpecies` Landau damping with full-f weights and an `e1_v1` binning plot) on NumPy and CuPy, compare energies, field coefficients and markers, and count the transfers of the time loop: in `model.integrate` only the 8-byte residual norm of each Krylov iteration (feectools' convergence test), scalar-sized downloads in the diagnostics, one download per saved dataset in the output. They run on a GPU, or on the fake CuPy with every CUDA launch emulated on the CPU (`CUNUMPY_FAKE_CUPY=1`, also under `mpirun -n 2`; `CudaKernel.compile` is skipped there); without either, `test_fake_cupy_matches_numpy` runs the comparison in a serial child process. Both pass on the fake CuPy on 1 and 2 ranks (not yet on a GPU). Needs the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere`/`vlasov_maxwell` CUDA kernels, `MassMatrixPreconditioner` on CuPy and device inner products, and a feectools release with them. +- [ ] **End-to-end tests of `LinearVlasovAmpereOneSpecies` and `VlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` and `test_gpu_VlasovAmpereOneSpecies.py` (shared helpers in `models/tests/gpu_e2e.py`) run a small Cuboid case (320 Sobol markers, 8 elements, 3 steps; for `VlasovAmpereOneSpecies` Landau damping with full-f weights and an `e1_v1` binning plot) on NumPy and CuPy, compare energies, field coefficients and markers, and count the transfers of the time loop: in `model.integrate` only the 8-byte residual norm of each Krylov iteration (feectools' convergence test), scalar-sized downloads in the diagnostics, one download per saved dataset in the output. They run on a GPU, or on the fake CuPy with every CUDA launch emulated on the CPU (`CUNUMPY_FAKE_CUPY=1`, also under `mpirun -n 2`; `CudaKernel.compile` is skipped there); without either, `test_fake_cupy_matches_numpy` runs the comparison in a serial child process. Both pass on the fake CuPy on 1 and 2 ranks (not yet on a GPU); on 2 ranks only with the markers MPI-sorted after loading (#764). CI runs them in unit shard-2 (serial and on 2 ranks on NumPy) and on the fake CuPy on 1 and 2 ranks (`test-PR-unit-mpi.yml`). The setup still assembles the mass matrices and vectors with host kernels (`kernel_1d_mat`, `kernel_3d_mat`, `kernel_3d_vec`, whose arrays are copied to the host and back at every call; outside the time loop, so not in the transfer budget); CUDA versions of these are open. Needs the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere`/`vlasov_maxwell` CUDA kernels, `MassMatrixPreconditioner` on CuPy and device inner products, and a feectools release with them. - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 454ee1050..0af89df1f 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -60,6 +60,9 @@ def struphy_test( f"{LIBPATH}/post_processing/tests/", f"{LIBPATH}/propagators/tests/", f"{LIBPATH}/simulation/tests/", + # end-to-end GPU tests (NumPy reference, and the CuPy comparison on a GPU or the fake CuPy) + f"{LIBPATH}/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py", + f"{LIBPATH}/models/tests/test_gpu_VlasovAmpereOneSpecies.py", ] if mpi > 1: diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 5918a4c61..97ca621f3 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -332,8 +332,10 @@ def weight_1d(c: int, d: int) -> Callable | xp.ndarray: return _kronecker_approximation(mass_operator, bc, weight_1d) def _core_diagonal(self): - # the matrix itself: _local_diagonal also adds the band entries that wrap onto the diagonal - return self._mass_operator.matrix + # diagonal of the mass matrix, including the band entries that wrap onto the main diagonal + # (StencilMatrix.diagonal() keeps only the entry at offset 0) + M = self._mass_operator.matrix + return _diagonal_operator(M.domain, _local_diagonal(M)) @property def dim_reduce(self) -> int | None: @@ -1979,6 +1981,7 @@ def solve(self, rhs, out=None, transposed=False): from scipy.linalg import solve_circulant + assert isinstance(rhs, np.ndarray), f"FFTSolver solves on the host, got {type(rhs).__name__}." assert rhs.T.shape[0] == self._column.size if out is None: diff --git a/src/struphy/feec/tests/test_preconditioner_cupy.py b/src/struphy/feec/tests/test_preconditioner_cupy.py index b914fcc64..9139e592b 100644 --- a/src/struphy/feec/tests/test_preconditioner_cupy.py +++ b/src/struphy/feec/tests/test_preconditioner_cupy.py @@ -11,15 +11,11 @@ * on a GPU (skipped otherwise), directly. """ -import contextlib -import subprocess -import sys - import cunumpy import numpy as np import pytest -from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") @@ -117,58 +113,14 @@ def check_preconditioners_on_cupy(bcs_name): np.testing.assert_allclose(_to_host(out_inplace), expected, rtol=1e-12, atol=1e-12 * np.abs(expected).max()) -def _host_buffer(value): - """The host buffer behind a fake CuPy array (the same memory, no copy); other values unchanged.""" - if isinstance(value, np.ndarray) or not hasattr(value, "__cuda_array_interface__"): - return value - return object.__getattribute__(value, "_a") - - -class _HostFields: - """A CUDA argument object whose array attributes are their fake CuPy host buffers.""" - - def __init__(self, args): - self._args = args - - def __getattr__(self, name): - return _host_buffer(getattr(self._args, name)) - - -@contextlib.contextmanager -def _emulated_launches(): - """On cunumpy's fake CuPy, run every ``CudaKernel`` launch by CPU emulation, in place on the fake device arrays. - - A minimal version of ``emulated_launches`` of struphy-hub/struphy#705; use that one once it is merged. - """ - from cunumpy.kernel_testing import emulate_cuda_kernel, fake_cupy_active - from cunumpy.kernels import CudaKernel +def check_preconditioners_on_fake_cupy(bcs_name): + """:func:`check_preconditioners_on_cupy` with the CUDA kernel launches emulated on the CPU.""" + from cunumpy.kernel_testing import emulated_launches, fake_cupy_active - from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS, emulate_struct_kernel + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS assert fake_cupy_active() - original = CudaKernel.__call__ - - def launch(self, *args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None): - if any(p.struct is not None for p in self.signature): - host_args = [_host_buffer(a) if p.struct is None else _HostFields(a) for p, a in zip(self.signature, args)] - emulate_struct_kernel(self, *host_args, n_threads=n_threads) - else: - grid, block = self.launch_shape(n_threads, grid=grid, block=block, args=args) - host_args = [_host_buffer(a) for a in args] - emulate_cuda_kernel( - self, *host_args, grid=grid, block=block, shared_mem=shared_mem, options=EMULATION_OPTIONS - ) - - CudaKernel.__call__ = launch - try: - yield - finally: - CudaKernel.__call__ = original - - -def check_preconditioners_on_fake_cupy(bcs_name): - """:func:`check_preconditioners_on_cupy` with the CUDA kernel launches emulated on the CPU.""" - with _emulated_launches(): + with emulated_launches(options=EMULATION_OPTIONS): check_preconditioners_on_cupy(bcs_name) @@ -177,20 +129,22 @@ def check_preconditioners_on_fake_cupy(bcs_name): def test_mass_preconditioners_fake_cupy(bcs_name): """Without a GPU: the preconditioners build and apply on cunumpy's fake CuPy, which rejects host/device mixing. - Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. + Runs in a serial subprocess (on MPI rank 0 only) because the fake CuPy must be installed before cunumpy is imported. """ - code = ( + run_fake_cupy_child( "from struphy.feec.tests.test_preconditioner_cupy import check_preconditioners_on_fake_cupy; " f"check_preconditioners_on_fake_cupy({bcs_name!r})" ) - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True - ) - assert result.returncode == 0, result.stderr[-4000:] @requires_cupy @pytest.mark.parametrize("bcs_name", list(BCS)) def test_mass_preconditioners_on_cupy(bcs_name): - """On a GPU: the preconditioners give the NumPy result on the CuPy backend.""" - check_preconditioners_on_cupy(bcs_name) + """On a GPU, or in this process on the fake CuPy (``CUNUMPY_FAKE_CUPY=1``, launches emulated on the CPU): the + preconditioners give the NumPy result on the CuPy backend.""" + from cunumpy.kernel_testing import fake_cupy_active + + if fake_cupy_active(): + check_preconditioners_on_fake_cupy(bcs_name) + else: + check_preconditioners_on_cupy(bcs_name) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index b0ba2404e..93d8b5920 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2593,7 +2593,8 @@ def _allocate_marker_array(self, dry_run: bool = False): ) # on the CuPy backend with several ranks, find out now (collectively, in the setup) whether MPI takes device - # buffers; otherwise the probe runs on the first marker exchange, inside the first time step + # buffers; otherwise the probe runs on the first marker exchange, inside the first time step. + # The probe is collective: all ranks of mpi_comm must allocate their marker arrays together, never one alone. if xp.get_backend() == "cupy" and self.mpi_size > 1: self._mpi_cuda_aware()