From 04dda211d11cbd38ca0c88953c7569b3b0e8d3f6 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 18:42:59 +0200 Subject: [PATCH 1/4] Kronecker solver on the device (no host copies in CuPy solves) KroneckerLinearSolver copied the data of every direction to the host and back in every solve on the CuPy backend, because its 1D solvers (BandedSolver: LAPACK gbtrs, SparseSolver: SuperLU) are host-only. - direct_solvers.DenseInverse: the 1D matrix is inverted once on the host by the solver itself (M = solver.solve(I)) and applied to device data as rhs @ M (one GEMM per pass). - KroneckerSolverSerialPass (also used inside the distributed pass) uses it for device arrays; the inverses are copied to the device when a KroneckerLinearSolver with device temporaries is built. Works for any linear 1D solver (e.g. struphy's FFTSolver); the FFT "solvers" opt out with dense_on_device = False. - BandedSolver/SparseSolver.solve use it for device arrays (copied on the first device solve). Host path unchanged. - Tests: CPU (dense inverse vs host solvers, Kronecker solve vs global reference, serial and MPI), fake CuPy subprocess with assert_no_transfers, GPU tests (serial and MPI, skipped without GPU). - CUDA_STRATEGY.md: notes for this change. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 40 ++- feectools/linalg/direct_solvers.py | 127 ++++++-- feectools/linalg/fft.py | 3 + feectools/linalg/kron.py | 45 +++ .../linalg/tests/test_kron_device_solve.py | 284 ++++++++++++++++++ 5 files changed, 476 insertions(+), 23 deletions(-) create mode 100644 feectools/linalg/tests/test_kron_device_solve.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 79aac0f11..2d5751588 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -61,7 +61,8 @@ feectools works when cunumpy's backend is CuPy, without device kernels: `cunumpy.kernels.PyccelKernel`, so they accept CuPy arrays (copied to the host and back). - Host-only metadata stays on NumPy: MPI and index bookkeeping in `ddm` and `fem.partitioning`, Kronecker solver sizes, index arithmetic with Python ints. -- Host-only libraries (LAPACK/SuperLU, SciPy FFT, SciPy sparse) get host copies per array, not by global backend. +- Host-only libraries (LAPACK/SuperLU, SciPy FFT, SciPy sparse) get host copies per array, not by global backend + (the Kronecker solver no longer does, see [Kronecker solver on the device](#kronecker-solver-on-the-device-device-kronecker-solve)). - The 1D collocation matrices of the global projectors are built vectorized (element-wise indexing was one device round trip per entry: 334 s of a 348 s Derham setup on the GPU). - Bug fix on both backends: `StencilMatrix._update_ghost_regions_serial` uses a ghost region `pads * shifts` wide. @@ -134,6 +135,39 @@ CuPy. There are no backend branches and no host staging at these call sites any - `psydac-accelerate` compiles every `*_kernels.py`, so the new folders are compiled like the old modules; `.cu` files are shipped as package data. +## Kronecker solver on the device (device-kronecker-solve) + +`KroneckerLinearSolver` solves device data on the device. Before, its 1D solvers (`BandedSolver`, LAPACK +`?gbtrs`; `SparseSolver`, SuperLU) copied the data of every direction to the host and back in every solve, i.e. in +every CG iteration preconditioned by struphy's `MassMatrixPreconditioner` (struphy-hub/struphy#650, #689). + +- **Dense inverses of the 1D matrices.** `direct_solvers.DenseInverse` inverts a 1D matrix once on the host, with + the solver itself (`M = solver.solve(I)`, so pivoting, transposition and the solver's own quirks are those of + the host path), and applies it as `rhs @ M`: one GEMM (cuBLAS) over all right-hand sides of a pass. The 1D + matrices are small (n ~ tens to hundreds), clamped (banded) or periodic (with corners, so banded storage has + full bandwidth anyway), and an n x n GEMM over m right-hand sides is a better fit for a GPU than m banded + triangular solves. `cupyx.scipy.linalg` offers dense `lu_factor`/`lu_solve` but, to our knowledge, no + `solve_banded`; `lu_solve` would need the dense matrix too and run two triangular solves instead of one GEMM. +- **Any linear 1D solver.** The Kronecker passes build the inverse from `solver.solve` on host arrays, so + struphy's `FFTSolver` (circulant mass matrices, `scipy.linalg.solve_circulant`) works without changes. A 1D + solver that must not be replaced by its dense matrix sets `dense_on_device = False`: the FFT "solvers" of + `linalg/fft.py` do (an FFT is O(n log n)); they still stage through the host. +- **Decided by the array**, as before: `KroneckerSolverSerialPass.solve_pass` (also used inside the distributed + pass, between the two `Alltoallv`) and `BandedSolver`/`SparseSolver.solve` use the dense inverse for + `xp.is_gpu(array)` and LAPACK/SuperLU otherwise. The host path is unchanged. +- **Copied to the device once, at setup.** A `KroneckerLinearSolver` whose temporaries are on the device (CuPy + backend) builds the inverses and copies them to the device in its constructor (and in `transpose()`, which + builds a new solver), so `solve` makes no transfer. A standalone `BandedSolver`/`SparseSolver` called with + device arrays copies its inverse on the first device solve (one `to_device`), cached per dtype. +- **Accuracy.** Inverse-times-vector and LU solve are both backward stable; for the well-conditioned mass + matrices the results agree to round-off (tests use 1e-12 relative to the solution). +- **Tests** (`linalg/tests/test_kron_device_solve.py`): `DenseInverse` against the host solvers (banded and + sparse, clamped and periodic, transposed, real and complex) and the Kronecker solve against a global + reference, serial and with MPI, on the CPU; the device path with cunumpy's fake CuPy in a serial subprocess + (clean environment, `MAYBEMPI=0`), including `assert_no_transfers` around `solve`; on a GPU (`requires_cupy`) + device against host solves under `assert_no_transfers`, serial and with MPI. Before this change a 3D solve on + the fake CuPy made 3 `to_host` copies (one per direction) and 3 uncounted copies back. + ## Kernel folders ``` @@ -188,6 +222,10 @@ Conventions, as in struphy: the serial case (in the parallel case it goes into the MPI reduction). To be measured on the H100. - **Interface matrices** (`StencilInterfaceMatrix`) and the remaining stencil kernels (`stencil2coo`, ...) still use `PyccelKernel` with host copies. +- **FFT on the device.** `DistributedFFT` and friends still stage their data through the host (SciPy FFT); they + could use `cupyx.scipy.fft` for device data. +- **Large 1D matrices.** The dense inverse costs n^2 memory and O(m n^2) per pass. Fine for the mass matrices of + struphy's preconditioners; for 1D sizes in the thousands a batched banded solve on the device would be cheaper. ## MPI through maybempi diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index f7c79418a..aefeb403f 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -4,13 +4,14 @@ # for full license details. # #---------------------------------------------------------------------------# from abc import abstractmethod +import numpy as np import cunumpy as xp from cunumpy.xp import array_backend from scipy.sparse import spmatrix, dia_matrix from feectools.linalg.basic import LinearSolver -__all__ = ('to_bnd', 'BandedSolver', 'SparseSolver') +__all__ = ('to_bnd', 'DenseInverse', 'BandedSolver', 'SparseSolver') #=============================================================================== def to_bnd(A): @@ -28,6 +29,84 @@ def to_bnd(A): return A_bnd, la, ua +#=============================================================================== +class DenseInverse: + """ + The inverse of a 1D solver's matrix as a dense matrix, for solves on the device. + + LAPACK and SuperLU run on the host only. For device data the matrix is inverted + once on the host, by the solver itself (it solves for the identity), and the inverse + is applied with one matrix product per solve (a cuBLAS GEMM on CuPy). A device solve + then makes no host/device copy. The 1D matrices of Kronecker solvers are small + (n ~ tens to hundreds), so the n x n inverse is cheap to store, and one GEMM over all + right-hand sides is faster on a GPU than banded triangular solves. + + Right-hand sides are rows, as in `BandedSolver.solve`: for ``rhs`` of shape + ``(m, n)`` the solution is ``rhs @ M`` with ``M = solver.solve(I)``, whose row k is + the solution for the k-th unit vector, i.e. ``M = op(A)^{-T}``. This holds for any + linear 1D solver with the row convention, transposed or not. + + Parameters + ---------- + solver : LinearSolver + A 1D solver whose ``solve`` accepts host (NumPy) arrays of shape ``(m, n)``. + + n : int + The size of the 1D matrix. + + dtype : dtype + The dtype of the right-hand sides. + """ + def __init__(self, solver, n, dtype): + eye = np.eye(int(n), dtype=dtype) + # host bookkeeping at setup: the host solver applied to the identity + self._host = np.ascontiguousarray(xp.to_numpy(solver.solve(eye))) + self._device = None + + @property + def shape(self): + return self._host.shape + + @property + def host_matrix(self): + """``M`` on the host.""" + return self._host + + def device_matrix(self): + """``M`` on the device, copied there once (on the first call).""" + if self._device is None: + self._device = xp.to_cupy(self._host) + return self._device + + def matrix_for(self, array): + """``M`` on the device for a device array, on the host otherwise.""" + return self.device_matrix() if xp.is_gpu(array) else self._host + + def solve(self, rhs, out=None): + """ + Solves for the right-hand sides ``rhs`` (rows) where they live, with one matrix product. + + ``out`` may be ``rhs`` (in-place solve). + """ + assert rhs.shape[-1] == self._host.shape[0] + result = rhs @ self.matrix_for(rhs) + if out is None: + return result + assert out.shape == rhs.shape + out[...] = result + return out + +#=============================================================================== +def _device_inverse(solver, n, dtype): + """The `DenseInverse` of a 1D solver for device solves, built on first use and cached per dtype.""" + cache = getattr(solver, '_dense_inverses', None) + if cache is None: + cache = solver._dense_inverses = {} + key = np.dtype(dtype) + if key not in cache: + cache[key] = DenseInverse(solver, n, key) + return cache[key] + #=============================================================================== class BandedSolver(LinearSolver): """ @@ -117,6 +196,7 @@ def transpose(self): obj._space = self._space obj._dtype = self._dtype obj._transposed = not self._transposed + obj._dense_inverses = {} return obj @@ -140,11 +220,19 @@ def solve(self, rhs, out=None): transposed = self._transposed + # LAPACK is host-only: device data is solved with the dense inverse on the + # device (see DenseInverse). Decided by the array itself, not the global + # backend, since host arrays may be passed on the CuPy backend too. + if xp.is_gpu(rhs): + if out is not None: + assert out.shape == rhs.shape + assert out.dtype == rhs.dtype + return _device_inverse(self, self._bmat.shape[1], rhs.dtype).solve(rhs, out=out) + if out is None: - # LAPACK is host-only: solve on the host, return on the caller's backend. - preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, xp.to_numpy(rhs).T, + preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, rhs.T, self._ipiv, trans=transposed) - out = xp.asarray(preout.T) if xp.is_gpu(rhs) else preout.T + out = preout.T else: assert out.shape == rhs.shape @@ -157,16 +245,7 @@ def solve(self, rhs, out=None): # TODO: handle non-contiguous views? # we want FORTRAN-contiguous data (default is assumed to be C contiguous). - # LAPACK is host-only: a device array is solved in a host copy. Decided by - # the array itself, not the global backend, since host arrays may be passed - # on the CuPy backend too. - if xp.is_gpu(out): - out_cpu = xp.to_numpy(out) - _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out_cpu.T, self._ipiv, overwrite_b=True, - trans=transposed) - out[...] = xp.asarray(out_cpu) - else: - _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True, + _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True, trans=transposed) return out @@ -206,6 +285,7 @@ def transpose(self): obj._space = self._space obj._splu = self._splu obj._transposed = not self._transposed + obj._dense_inverses = {} return obj @@ -229,19 +309,22 @@ def solve(self, rhs, out=None): assert rhs.T.shape[0] == self._splu.shape[1] transposed = self._transposed + # SuperLU is host-only: device data is solved with the dense inverse on the + # device (see DenseInverse); decided by the array, not the global backend. + if xp.is_gpu(rhs): + if out is not None: + assert out.shape == rhs.shape + assert out.dtype == rhs.dtype + return _device_inverse(self, self._splu.shape[1], rhs.dtype).solve(rhs, out=out) + if out is None: - # SuperLU is host-only: solve on the host, return on the caller's backend. - out = self._splu.solve(xp.to_numpy(rhs).T, trans='T' if transposed else 'N').T - if xp.is_gpu(rhs): - out = xp.asarray(out) + out = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T else: assert out.shape == rhs.shape assert out.dtype == rhs.dtype - # currently no in-place solve exposed. SuperLU is host-only; decided by - # the arrays themselves, not the global backend. - result = self._splu.solve(xp.to_numpy(rhs).T, trans='T' if transposed else 'N').T - out[:] = xp.asarray(result) if xp.is_gpu(out) else result + # currently no in-place solve exposed. + out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T return out diff --git a/feectools/linalg/fft.py b/feectools/linalg/fft.py index bec639d2c..8c5c06bd5 100644 --- a/feectools/linalg/fft.py +++ b/feectools/linalg/fft.py @@ -36,6 +36,9 @@ class OneDimSolver(LinearSolver): function : Callable The given function. """ + # KroneckerLinearSolver: keep applying the function (not its dense matrix) to device data + dense_on_device = False + def __init__(self, function): self._function = function diff --git a/feectools/linalg/kron.py b/feectools/linalg/kron.py index eb98c99a3..e2d76d771 100644 --- a/feectools/linalg/kron.py +++ b/feectools/linalg/kron.py @@ -9,6 +9,7 @@ from feectools.linalg.basic import LinearOperator, LinearSolver from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix +from feectools.linalg.direct_solvers import DenseInverse __all__ = ('KroneckerStencilMatrix', 'KroneckerLinearSolver', @@ -439,6 +440,13 @@ def __init__(self, V, W, solvers): # for now: allocate temporary arrays here (can be removed later) self._temp1, self._temp2 = self._allocate_temps() + + # Temporaries on the device (CuPy backend): the 1D solves will run on the + # device, so build their dense inverses and copy them to the device now, + # not inside the first solve (see KroneckerSolverSerialPass). + if xp.is_gpu(self._temp1): + for solver_pass in self._solver_passes: + solver_pass.prepare_device(self._dtype) def _setup_solvers(self): """ @@ -670,6 +678,32 @@ def __init__(self, solver, nglobal, mglobal): self._datasize = nglobal*mglobal self._solver = solver self._view = None + # dense inverses for device solves, per dtype (see DenseInverse) + self._dense_inverses = {} + + def uses_dense_inverse(self): + """ + Whether device data is solved with the dense inverse of the 1D solver. + + True for every 1D solver unless it sets ``dense_on_device = False`` + (as the FFT "solvers" of `feectools.linalg.fft` do; their dense matrix + would replace an O(n log n) transform by an O(n^2) product). + """ + return getattr(self._solver, 'dense_on_device', True) + + def dense_inverse(self, dtype): + """The `DenseInverse` of the 1D solver for right-hand sides of type ``dtype``, built once.""" + key = np.dtype(dtype) + inverse = self._dense_inverses.get(key) + if inverse is None: + inverse = DenseInverse(self._solver, self._dimrhs, key) + self._dense_inverses[key] = inverse + return inverse + + def prepare_device(self, dtype): + """Builds the dense inverse and copies it to the device ahead of the first device solve.""" + if self.uses_dense_inverse(): + self.dense_inverse(dtype).device_matrix() def required_memory(self): """ @@ -695,6 +729,13 @@ def solve_pass(self, workmem, tempmem): view = workmem[:self._datasize] view.shape = (int(self._numrhs), int(self._dimrhs)) + # LAPACK/SuperLU (and the solvers of other libraries) are host-only: + # device data is solved with the precomputed dense inverse, one matrix + # product on the device without host copies. Decided by the array. + if xp.is_gpu(view) and self.uses_dense_inverse(): + self.dense_inverse(view.dtype).solve(view, out=view) + return + # call solver in in-place mode self._solver.solve(view, out=view) @@ -837,6 +878,10 @@ def required_memory(self): """ return max(self._datasize, self._localsize) + def prepare_device(self, dtype): + """See `KroneckerSolverSerialPass.prepare_device`.""" + self._serialsolver.prepare_device(dtype) + def _blocked_to_contiguous(self, blocked, contiguous): """ Copies from a blocked view to a contiguous view. diff --git a/feectools/linalg/tests/test_kron_device_solve.py b/feectools/linalg/tests/test_kron_device_solve.py new file mode 100644 index 000000000..1131d7c94 --- /dev/null +++ b/feectools/linalg/tests/test_kron_device_solve.py @@ -0,0 +1,284 @@ +"""Kronecker solves of device data stay on the device. + +`KroneckerLinearSolver` (and `BandedSolver`/`SparseSolver` called with device arrays) solve +device data with the dense inverses of the 1D matrices (`DenseInverse`), one matrix product on +the device, instead of copying the data to the host for LAPACK/SuperLU. These tests check + +* on the CPU: `DenseInverse` gives the host solvers' results (clamped and periodic matrices, + transposed, real and complex), and the Kronecker solve matches a global reference, serial and + distributed (``mpirun -n 2 python -m pytest -m mpi --with-mpi``); +* with cunumpy's fake CuPy (host memory, CuPy semantics, in a subprocess): device solves give + the host results and make no host/device transfers; +* on a GPU (skipped otherwise): the same, serial and with MPI. +""" +import importlib.util +import os +import subprocess +import sys +from functools import reduce + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from cunumpy.profiling import assert_no_transfers +from maybempi import MPI +from scipy.sparse import csr_matrix + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.linalg.direct_solvers import BandedSolver, DenseInverse, SparseSolver +from feectools.linalg.kron import KroneckerLinearSolver +from feectools.linalg.stencil import StencilVector, StencilVectorSpace + +# Dense inverse times rhs vs. LU solve: both are backward stable, the matrices are well conditioned. +RTOL = 1e-12 + +# (npts, pads, periods): clamped, periodic and mixed directions, 1D to 3D +CASES = [ + ((12,), (2,), (False,)), + ((10,), (3,), (True,)), + ((12, 9), (2, 1), (False, True)), + ((8, 10, 7), (3, 2, 1), (True, False, False)), + ((8, 6, 9), (2, 3, 2), (True, True, True)), +] +SOLVER_KINDS = ("banded", "sparse") + + +# =============================================================================== +def matrix_1d(n, p, periodic, dtype=float): + """A non-symmetric, diagonally dominant matrix with the band of a degree-p mass matrix. + + Periodic matrices have corners (cyclic band), as the 1D mass matrices of periodic directions. + """ + A = np.zeros((n, n), dtype=dtype) + for i in range(n): + for d in range(-p, p + 1): + j = i + d + if periodic: + j %= n + elif not 0 <= j < n: + continue + A[i, j] += 1.0 / (1 + abs(d)) * (1 + 0.1 * np.sin(i + 2 * j)) + if np.dtype(dtype).kind == "c": + A[i, j] += 0.05j * np.cos(3 * i - j) + A[i, i] += 2 * p + 2 + return A + + +def banded_solver(A): + """`BandedSolver` of a dense matrix (band storage as in `to_bnd`).""" + n = A.shape[0] + rows, cols = np.nonzero(A) + la = int(max(0, (rows - cols).max())) + ua = int(max(0, (cols - rows).max())) + ab = np.zeros((1 + ua + 2 * la, n), dtype=A.dtype) + for i, j in zip(rows, cols): + ab[la + ua + i - j, j] = A[i, j] + return BandedSolver(ua, la, ab) + + +def solver_1d(A, kind): + return banded_solver(A) if kind == "banded" else SparseSolver(csr_matrix(A)) + + +def make_space(npts, pads, periods, dtype=float, comm=None): + ndim = len(npts) + D = DomainDecomposition(list(npts), periods=list(periods), comm=comm) + gs, ge = [], [] + for axis in range(ndim): + ee = np.array(D.global_element_ends[axis]).copy() + ee[-1] = npts[axis] - 1 + ge.append(ee) + gs.append(np.array([0] + (ee[:-1] + 1).tolist())) + C = CartDecomposition(D, list(npts), gs, ge, pads=list(pads), shifts=[1] * ndim) + return StencilVectorSpace(C, dtype=dtype) + + +def owned_slices(V): + return tuple(slice(int(s), int(e) + 1) for s, e in zip(V.starts, V.ends)) + + +def scatter(V, glob): + """This rank's part of a global array, in a new vector on the active backend.""" + v = StencilVector(V) + v[owned_slices(V)] = xp.asarray(glob[owned_slices(V)]) + v.update_ghost_regions() + return v + + +def global_rhs(npts, dtype=float): + """A deterministic global right-hand side, independent of the decomposition.""" + idx = np.indices(npts).reshape(len(npts), -1) + b = np.cos(1.0 + idx.T @ np.arange(1, len(npts) + 1)).reshape(npts) + if np.dtype(dtype).kind == "c": + b = b + 1j * np.sin(idx.T @ np.arange(2, len(npts) + 2)).reshape(npts) + return b.astype(dtype) + + +def reference_solution(mats, b, transposed=False): + K = reduce(np.kron, [A.T if transposed else A for A in mats]) + return np.linalg.solve(K, b.ravel()).reshape(b.shape) + + +def kron_solve(npts, pads, periods, kind, *, transposed=False, dtype=float, comm=None, check_transfers=False): + """Solves the case on the active backend; returns this rank's part (host), and the global rhs and matrices.""" + mats = [matrix_1d(n, p, per, dtype) for n, p, per in zip(npts, pads, periods)] + V = make_space(npts, pads, periods, dtype, comm) + S = KroneckerLinearSolver(V, V, [solver_1d(A, kind) for A in mats]) + if transposed: + S = S.transpose() + b = global_rhs(npts, dtype) + rhs = scatter(V, b) + x = StencilVector(V) + if check_transfers: + with assert_no_transfers(): + S.solve(rhs, out=x) + else: + S.solve(rhs, out=x) + return xp.to_numpy(x[owned_slices(V)]), b, mats, owned_slices(V) + + +def assert_close(actual, desired, what=""): + scale = np.abs(desired).max() + np.testing.assert_allclose(actual, desired, rtol=0, atol=RTOL * scale, err_msg=what) + + +def compare_device_with_host(npts, pads, periods, kind, transposed=False, dtype=float, comm=None): + """Device solve (CuPy backend, no transfers) vs host solve (NumPy backend) vs global reference.""" + with xp.use_backend("numpy"): + host, b, mats, owned = kron_solve(npts, pads, periods, kind, transposed=transposed, dtype=dtype, comm=comm) + with xp.use_backend("cupy"): + device, _, _, _ = kron_solve(npts, pads, periods, kind, transposed=transposed, dtype=dtype, comm=comm, + check_transfers=True) + what = f"{npts=} {periods=} {kind=} {transposed=} {dtype=}" + assert_close(device, host, what) + assert_close(device, reference_solution(mats, b, transposed)[owned], what) + + +def compare_1d_device_with_host(kind, periodic, transposed=False, dtype=float): + """`BandedSolver`/`SparseSolver` on device arrays: host results, no transfers after the first solve.""" + A = matrix_1d(11, 2, periodic, dtype) + solver = solver_1d(A, kind) + if transposed: + solver = solver.transpose() + B = np.stack([global_rhs((11,), dtype) * (k + 1) for k in range(4)]) + host = solver.solve(B) + with xp.use_backend("cupy"): + B_dev = xp.to_cupy(B) + solver.solve(B_dev) # first device solve copies the inverse to the device + with assert_no_transfers(): + out_dev = solver.solve(B_dev) + inplace_dev = B_dev.copy() + solver.solve(inplace_dev, out=inplace_dev) + for result in (out_dev, inplace_dev): + assert xp.is_gpu(result) + assert_close(xp.to_numpy(result), host, f"{kind=} {periodic=} {transposed=} {dtype=}") + + +def run_device_checks(): + """All device checks; run in a process whose CuPy backend is the real or fake CuPy.""" + for npts, pads, periods in CASES: + for kind in SOLVER_KINDS: + for transposed in (False, True): + compare_device_with_host(npts, pads, periods, kind, transposed) + compare_device_with_host((9, 8), (2, 2), (True, False), "banded", dtype=complex) + for kind in SOLVER_KINDS: + for periodic in (False, True): + for transposed in (False, True): + compare_1d_device_with_host(kind, periodic, transposed) + compare_1d_device_with_host("sparse", True, dtype=complex) + print("device checks passed") + + +# =============================================================================== +# CPU tests +# =============================================================================== +@pytest.mark.parametrize("kind", SOLVER_KINDS) +@pytest.mark.parametrize("periodic", [False, True]) +@pytest.mark.parametrize("transposed", [False, True]) +@pytest.mark.parametrize("dtype", [float, complex]) +def test_dense_inverse_matches_host_solver(kind, periodic, transposed, dtype): + A = matrix_1d(13, 3, periodic, dtype) + solver = solver_1d(A, kind) + if transposed: + solver = solver.transpose() + op = A.T if transposed else A + B = np.stack([global_rhs((13,), dtype) * (k + 1) + k for k in range(5)]) + + inverse = DenseInverse(solver, 13, dtype) + assert inverse.shape == (13, 13) + host = solver.solve(B.copy()) + assert_close(host, np.linalg.solve(op, B.T).T) + assert_close(inverse.solve(B), host) + assert_close(inverse.solve(B[2]), host[2]) # a single right-hand side + inplace = B.copy() + assert inverse.solve(inplace, out=inplace) is inplace + assert_close(inplace, host) + + +@pytest.mark.parametrize("npts, pads, periods", CASES) +@pytest.mark.parametrize("kind", SOLVER_KINDS) +@pytest.mark.parametrize("transposed", [False, True]) +def test_kronecker_solve_matches_reference(npts, pads, periods, kind, transposed): + with xp.use_backend("numpy"): + x, b, mats, owned = kron_solve(npts, pads, periods, kind, transposed=transposed) + assert_close(x, reference_solution(mats, b, transposed)[owned]) + + +@pytest.mark.mpi +@pytest.mark.parametrize("npts, pads, periods", [c for c in CASES if len(c[0]) > 1]) +@pytest.mark.parametrize("kind", SOLVER_KINDS) +def test_kronecker_solve_matches_reference_mpi(npts, pads, periods, kind): + with xp.use_backend("numpy"): + x, b, mats, owned = kron_solve(npts, pads, periods, kind, comm=MPI.COMM_WORLD) + assert_close(x, reference_solution(mats, b)[owned]) + + +_MPI_LAUNCHER_PREFIXES = ("OMPI_", "PMIX_", "PMI_", "HYDRA_", "I_MPI_", "SLURM_") + + +def _real_cupy_installed(): + module = sys.modules.get("cupy") + if module is not None: + return not getattr(module, "__cunumpy_fake__", False) + return importlib.util.find_spec("cupy") is not None + + +@pytest.mark.skipif(_real_cupy_installed(), + reason="CuPy is installed: the fake CuPy cannot be used, the GPU tests cover this") +def test_device_solve_on_fake_cupy(): + """The device path with cunumpy's fake CuPy, in a serial child process with a clean environment.""" + env = {k: v for k, v in os.environ.items() if not k.startswith(_MPI_LAUNCHER_PREFIXES)} + env.pop("CUNUMPY_BACKEND", None) + env.update(MAYBEMPI="0", CUNUMPY_FAKE_CUPY="1") + code = "from feectools.linalg.tests.test_kron_device_solve import run_device_checks; run_device_checks()" + proc = subprocess.run([sys.executable, "-c", code], env=env, capture_output=True, text=True, timeout=900) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "device checks passed" in proc.stdout + + +# =============================================================================== +# GPU tests +# =============================================================================== +@requires_cupy +@pytest.mark.parametrize("npts, pads, periods", CASES) +@pytest.mark.parametrize("kind", SOLVER_KINDS) +@pytest.mark.parametrize("transposed", [False, True]) +def test_device_kronecker_solve_gpu(npts, pads, periods, kind, transposed): + compare_device_with_host(npts, pads, periods, kind, transposed) + + +@requires_cupy +@pytest.mark.parametrize("kind", SOLVER_KINDS) +@pytest.mark.parametrize("periodic", [False, True]) +@pytest.mark.parametrize("transposed", [False, True]) +def test_device_1d_solve_gpu(kind, periodic, transposed): + compare_1d_device_with_host(kind, periodic, transposed) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("npts, pads, periods", [c for c in CASES if len(c[0]) > 1]) +@pytest.mark.parametrize("kind", SOLVER_KINDS) +def test_device_kronecker_solve_gpu_mpi(npts, pads, periods, kind): + compare_device_with_host(npts, pads, periods, kind, comm=MPI.COMM_WORLD) From 9cb793b5de72af2d8d64bb42f0c9e593bd4499a2 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 18:43:45 +0200 Subject: [PATCH 2/4] CUDA_STRATEGY.md: name the PR (#96) Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 2d5751588..a8a1dcc0a 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -62,7 +62,7 @@ feectools works when cunumpy's backend is CuPy, without device kernels: - Host-only metadata stays on NumPy: MPI and index bookkeeping in `ddm` and `fem.partitioning`, Kronecker solver sizes, index arithmetic with Python ints. - Host-only libraries (LAPACK/SuperLU, SciPy FFT, SciPy sparse) get host copies per array, not by global backend - (the Kronecker solver no longer does, see [Kronecker solver on the device](#kronecker-solver-on-the-device-device-kronecker-solve)). + (the Kronecker solver no longer does, see [Kronecker solver on the device](#kronecker-solver-on-the-device-96)). - The 1D collocation matrices of the global projectors are built vectorized (element-wise indexing was one device round trip per entry: 334 s of a 348 s Derham setup on the GPU). - Bug fix on both backends: `StencilMatrix._update_ghost_regions_serial` uses a ghost region `pads * shifts` wide. @@ -135,9 +135,9 @@ CuPy. There are no backend branches and no host staging at these call sites any - `psydac-accelerate` compiles every `*_kernels.py`, so the new folders are compiled like the old modules; `.cu` files are shipped as package data. -## Kronecker solver on the device (device-kronecker-solve) +## Kronecker solver on the device (#96) -`KroneckerLinearSolver` solves device data on the device. Before, its 1D solvers (`BandedSolver`, LAPACK +[#96](https://github.com/struphy-hub/feectools/pull/96): `KroneckerLinearSolver` solves device data on the device. Before, its 1D solvers (`BandedSolver`, LAPACK `?gbtrs`; `SparseSolver`, SuperLU) copied the data of every direction to the host and back in every solve, i.e. in every CG iteration preconditioned by struphy's `MassMatrixPreconditioner` (struphy-hub/struphy#650, #689). From 64b10fa0ecaf5d8dc138163b8df6b7594f46676a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:04:24 +0200 Subject: [PATCH 3/4] DenseInverse.solve: say why the product is rhs @ M (rows, M = op(A)^-T, no symmetry assumed) Co-Authored-By: Claude Opus 5.5 --- feectools/linalg/direct_solvers.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index aefeb403f..c49377314 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -86,6 +86,11 @@ def solve(self, rhs, out=None): """ Solves for the right-hand sides ``rhs`` (rows) where they live, with one matrix product. + The right-hand sides are the rows of ``rhs``, as in `BandedSolver.solve`, so the + solutions are the rows of ``rhs @ op(A)^{-T}``; ``M`` stores ``op(A)^{-T}`` (no + symmetry is assumed). Multiplying from the right keeps the row-contiguous work + arrays of `KroneckerLinearSolver` as they are: one GEMM, no transpose copy. + ``out`` may be ``rhs`` (in-place solve). """ assert rhs.shape[-1] == self._host.shape[0] From ebe78de4f2e4c0d0df97c4ac3a45c30433bdeb52 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 8 Oct 2026 11:03:31 +0200 Subject: [PATCH 4/4] Bump version to 0.6.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f378db48e..e80c4f4cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "feectools" -version = "0.5.0" +version = "0.6.0" description = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies." readme = "README.md" requires-python = ">= 3.10"