From 1845bc652446ce46eebc825597db01d14fec0dd8 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:13:49 +0200 Subject: [PATCH] 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)