diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 79aac0f11..a8a1dcc0a 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-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. @@ -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 (#96) + +[#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). + +- **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..c49377314 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,89 @@ 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. + + 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] + 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 +201,7 @@ def transpose(self): obj._space = self._space obj._dtype = self._dtype obj._transposed = not self._transposed + obj._dense_inverses = {} return obj @@ -140,11 +225,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 +250,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 +290,7 @@ def transpose(self): obj._space = self._space obj._splu = self._splu obj._transposed = not self._transposed + obj._dense_inverses = {} return obj @@ -229,19 +314,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) 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"