diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..068760679 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -270,6 +270,16 @@ 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). The Kronecker preconditioners +(`MassMatrixPreconditioner`) build their 1d mass matrices and 1d solvers +(`FFTSolver` for circulant, `SparseSolver` otherwise) on the host on every backend; the process-local 1d stencil +matrices (factors of the `KroneckerStencilMatrix`) and the diagonal scaling live on the device. Applying them makes +no host copies: `KroneckerLinearSolver` solves device data with dense inverses built once from the host 1d solvers +([feectools#96](https://github.com/struphy-hub/feectools/pull/96)). Tested with fake CuPy, emulated launches and +`assert_no_transfers` (`feec/tests/test_preconditioner_cupy.py`). `StiffnessPreconditioner` builds its 1d data the +same way, but `KroneckerSumSolver` (feectools 0.7.0) does not yet run on the device: its dense "solvers" hold device +matrices, from which `KroneckerLinearSolver` tries to build dense inverses on the host. + 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 d64838dc2..3703ffa7d 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -675,13 +675,11 @@ def _build_approximation(self, bc): stiffness.append(S_d) mass.append(M_d) - local_mass.append(_process_local_matrix_1d(xp.asarray(M_d), comp.coeff_space, d, domain_decomposition)) + local_mass.append(_process_local_matrix_1d(M_d, comp.coeff_space, d, domain_decomposition)) local_stiffness.append( None if S_d is None - else _process_local_matrix_1d( - xp.asarray(coeffs[d] * S_d), comp.coeff_space, d, domain_decomposition - ) + else _process_local_matrix_1d(coeffs[d] * S_d, comp.coeff_space, d, domain_decomposition) ) # stabilization: sigma_c * (M_1 x M_2 x M_3 - F_1 x F_2 x F_3), or the identity @@ -774,9 +772,10 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: # the factor may own more rows than the codomain on this process off = A.codomain.pads[0] * A.codomain.shifts[0] + s - A.codomain.starts[0] diags.append(A._data[off : off + e - s + 1, A.pads[0]]) + # outer product by broadcasting (works on every backend, unlike ufunc.outer) out = diags[0] for d in diags[1:]: - out = xp.multiply.outer(out, d) + out = out[..., None] * d return [out] if isinstance(op, _DenseKroneckerProduct): return [op.local_diagonal()] @@ -1059,27 +1058,34 @@ def _apply_bc_1d(M: StencilMatrix, bc_d: tuple[bool, bool]) -> None: apply_essential_bc_stencil(M, axis=0, ext=ext, order=0, identity=True) -def _solver_1d(M: StencilMatrix, M_arr: xp.ndarray) -> FFTSolver | SparseSolver: - """Direct solver for the 1d matrix M (dense copy M_arr): FFT if circulant, else sparse LU.""" +def _solver_1d(M: StencilMatrix, M_arr: np.ndarray) -> FFTSolver | SparseSolver: + """ + Direct solver for the 1d matrix M (dense host copy M_arr): FFT if circulant, else sparse LU. + + The solver is a host solver on every backend, as expected by KroneckerLinearSolver: on the + CuPy backend, the latter solves device data with the dense inverse it builds from this solver + once, so that applying the preconditioner makes no host copies. + """ if is_circulant(M_arr): return FFTSolver(M_arr) return SparseSolver(M.tosparse()) def _process_local_matrix_1d( - M_arr: xp.ndarray, coeff_space: StencilVectorSpace, d: int, domain_decomposition: DomainDecomposition + M_arr: np.ndarray, coeff_space: StencilVectorSpace, d: int, domain_decomposition: DomainDecomposition ) -> StencilMatrix: """ Process-local 1d factor of a KroneckerStencilMatrix on ``coeff_space``. The factor lives on a 1d space without communicator that owns the same rows (``starts[d]`` to ``ends[d]``) as ``coeff_space`` in direction ``d`` on this process, - as required by KroneckerStencilMatrix. + as required by KroneckerStencilMatrix. The band is filled on the host and copied to the + data of the factor (active backend) in one step. Parameters ---------- - M_arr : xp.ndarray - The global 1d matrix as dense array. + M_arr : np.ndarray + The global 1d matrix as dense array (a device array is copied to the host). coeff_space : StencilVectorSpace The (distributed) 3d coefficient space. @@ -1098,21 +1104,24 @@ def _process_local_matrix_1d( cart_1d = CartDecomposition(domain_decomposition, [n], [[s]], [[e]], [p], [1]) V_local = StencilVectorSpace(cart_1d) M_local = StencilMatrix(V_local, V_local) + M_arr = xp.to_numpy(M_arr) # copy the rows owned by this process: entry (i, j) is stored at row i - s + p and diagonal # index k + p with the offset k = j - i; in periodic directions k is taken nearest to 0 # (in [-n/2, n/2)), so that the main diagonal is at k = 0 also if 2p + 1 > n - rows, cols = xp.nonzero(M_arr) + rows, cols = np.nonzero(M_arr) on_process = (rows >= s) & (rows <= e) rows, cols = rows[on_process], cols[on_process] k = cols - rows if coeff_space.periods[d]: k = (k + n // 2) % n - n // 2 - assert xp.all(xp.abs(k) <= p), "Entries outside the band of the stencil matrix." - M_local._data[rows - s + p, k + p] = M_arr[rows, cols] + assert np.all(np.abs(k) <= p), "Entries outside the band of the stencil matrix." + data = np.zeros(M_local._data.shape, dtype=M_local._data.dtype) + data[rows - s + p, k + p] = M_arr[rows, cols] + M_local._data[...] = xp.asarray(data) # check if stencil matrix was built correctly - assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + assert np.allclose(xp.to_numpy(M_local.toarray())[s : e + 1], M_arr[s : e + 1]) return M_local @@ -1165,7 +1174,8 @@ def _kronecker_approximation( if bc is not None and _has_essential_bc(mass_operator, basis, c, d): _apply_bc_1d(M, bc[d]) - M_arr = M.toarray() + # the dense 1d matrix and the 1d solver are host setup data on every backend + M_arr = xp.to_numpy(M.toarray()) solvercells.append(_solver_1d(M, M_arr)) matrixcells.append(_process_local_matrix_1d(M_arr, femspace_c.coeff_space, d, domain_decomposition)) @@ -1494,26 +1504,31 @@ 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). + A host (NumPy) solver on every backend: the 1d matrices of Kronecker solvers are small + setup data. On the CuPy backend, KroneckerLinearSolver solves device data with the dense + inverse it builds from this solver (by solving for the identity on the host), so ``solve`` + is only called with host arrays. + Parameters ---------- - circmat : xp.ndarray - Generic circulant matrix. + circmat : np.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) from scipy.linalg import solve_circulant # deferred: scipy.linalg is slow to import - self._space = xp.ndarray + self._space = np.ndarray # copy: circmat is still used by the caller (e.g. for the process-local stencil matrix) - self._column = xp.array(circmat[:, 0], copy=True) + self._column = np.array(circmat[:, 0], copy=True) # stabilize a singular matrix once, here, so that all solves use the same matrix try: - solve_circulant(self._column, xp.ones_like(self._column)) - except xp.linalg.LinAlgError: + solve_circulant(self._column, np.ones_like(self._column)) + except np.linalg.LinAlgError: eps = 1e-4 logger.info(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") self._column[0] *= 1.0 + eps @@ -1532,13 +1547,13 @@ def solve(self, rhs, out=None, transposed=False): Parameters ---------- - rhs : xp.ndarray - The right-hand sides to solve for. The vectors are assumed to be given in C-contiguous order, + rhs : np.ndarray + The right-hand sides to solve for (host array). 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. - out : xp.ndarray, optional + out : np.ndarray, optional Output vector. If given, it has to have the same shape and datatype as rhs. transposed : bool @@ -1565,7 +1580,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 ------- @@ -1573,11 +1588,11 @@ 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] # circulant: every row is the previous one shifted by one, i.e. mat[i, j] = mat[0, (j - i) % n] n = mat.shape[0] - idx = (xp.arange(n)[None, :] - xp.arange(n)[:, None]) % n - return bool(xp.allclose(mat, mat[0][idx])) + idx = (np.arange(n)[None, :] - np.arange(n)[:, None]) % n + return bool(np.allclose(mat, mat[0][idx])) 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..b914fcc64 --- /dev/null +++ b/src/struphy/feec/tests/test_preconditioner_cupy.py @@ -0,0 +1,196 @@ +"""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, the diagonal scaling and the vectors they are applied to live on +the device, and feectools' KroneckerLinearSolver solves device data with dense inverses built from the host +1d solvers. These tests build the preconditioners on the CuPy backend, check that applying them makes no +host/device transfers, 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") +# (class name, keyword arguments) +PRECONDITIONERS = ( + ("MassMatrixPreconditioner", {"diagonal_scaling": False}), + ("MassMatrixPreconditioner", {}), + ("MassMatrixPreconditioner", {"weight_reduction": "average"}), + ("MassMatrixPreconditioner", {"dim_reduce": None}), +) +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. + + Applying the preconditioners (``dot``, also in place) makes no host/device transfers. + """ + 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, kwargs in PRECONDITIONERS: + cls = getattr(preconditioner, name) + case = (bcs_name, space, name, kwargs) + with cunumpy.use_backend("numpy"): + expected = _to_host(cls(M_host, **kwargs).dot(v_host)) + with cunumpy.use_backend("cupy"): + pc = cls(M_device, **kwargs) + out_inplace = pc.codomain.zeros() + with cunumpy.profiling.assert_no_transfers(): + out = pc.dot(v_device) + # in-place application + pc.dot(v_device, out=out_inplace) + for block in _blocks(out): + assert cunumpy.is_gpu(block._data), case + assert np.linalg.norm(expected) > 0, case + 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)