Kronecker solver on the device (no host copies in CuPy solves)

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
 

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

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
 

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.