[gpu-1] End-to-end GPU test of VlasovAmpereOneSpecies on 1 and 2 ranks

.github/workflows/reusable-scheduled.yml

@@ -107,7 +107,7 @@ jobs:
               TEST_FILES="bsplines/tests/ console/tests/ fields_background/tests/ geometry/tests/ initial/tests/ kinetic_background/tests/ linear_algebra/tests/"
               ;;
             shard-2)
-              TEST_FILES="propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/"
+              TEST_FILES="propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py"
               ;;
             shard-3)
               TEST_FILES="feec/tests/"

.github/workflows/reusable-unit-testing.yml

@@ -39,7 +39,7 @@ jobs:
           - shard: shard-1
             test_files: bsplines/tests/ console/tests/ fields_background/tests/ geometry/tests/ initial/tests/ kinetic_background/tests/ linear_algebra/tests/
           - shard: shard-2
-            test_files: propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/
+            test_files: propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py
           - shard: shard-3
             test_files: feec/tests/
           - shard: shard-4

.github/workflows/test-PR-unit-mpi.yml

@@ -88,3 +88,11 @@ jobs:
           source /struphy_fortran_/env_fortran_/bin/activate
           cd ${{ env.STRUPHY_PATH }}
           mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi polar/tests/test_polar_cupy.py -k mpi_fake
+
+      - name: Run the end-to-end GPU model tests on the fake CuPy
+        # 8 elements of degree 3: at most 2 ranks
+        if: matrix.n-procs <= 2
+        run: |
+          source /struphy_fortran_/env_fortran_/bin/activate
+          cd ${{ env.STRUPHY_PATH }}
+          mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py

CUDA_STRATEGY.md

@@ -31,6 +31,7 @@ Progress is tracked in struphy-hub/struphy#650.
 - [x] **FEEC assembly kernels in kernel folders** (#411) — the 16 assembly kernels of `feec/mass.py`, `feec/boundary_mass.py` and `feec/basis_projection_ops.py` move from `feec/mass_kernels.py` and `feec/basis_projection_kernels.py` into one folder each under `feec/kernels/`; pure restructuring, ready for CUDA ports (see [FEEC assembly kernels](#feec-assembly-kernels-411)).
 - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)).
 - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)).
+- [ ] **End-to-end tests of `LinearVlasovAmpereOneSpecies` and `VlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` and `test_gpu_VlasovAmpereOneSpecies.py` (shared helpers in `models/tests/gpu_e2e.py`) run a small Cuboid case (320 Sobol markers, 8 elements, 3 steps; for `VlasovAmpereOneSpecies` Landau damping with full-f weights and an `e1_v1` binning plot) on NumPy and CuPy, compare energies, field coefficients and markers, and count the transfers of the time loop: in `model.integrate` only the 8-byte residual norm of each Krylov iteration (feectools' convergence test), scalar-sized downloads in the diagnostics, one download per saved dataset in the output. They run on a GPU, or on the fake CuPy with every CUDA launch emulated on the CPU (`CUNUMPY_FAKE_CUPY=1`, also under `mpirun -n 2`; `CudaKernel.compile` is skipped there); without either, `test_fake_cupy_matches_numpy` runs the comparison in a serial child process. Both pass on the fake CuPy on 1 and 2 ranks (not yet on a GPU); on 2 ranks only with the markers MPI-sorted after loading (#764). CI runs them in unit shard-2 (serial and on 2 ranks on NumPy) and on the fake CuPy on 1 and 2 ranks (`test-PR-unit-mpi.yml`). The setup still assembles the mass matrices and vectors with host kernels (`kernel_1d_mat`, `kernel_3d_mat`, `kernel_3d_vec`, whose arrays are copied to the host and back at every call; outside the time loop, so not in the transfer budget); CUDA versions of these are open. Needs the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere`/`vlasov_maxwell` CUDA kernels, `MassMatrixPreconditioner` on CuPy and device inner products, and a feectools release with them.
 - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes.
 - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged).
 
@@ -274,6 +275,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-<n>-<topic>` PR, linked in the table above and
 from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/issues/650)).
 

src/struphy/console/test.py

@@ -60,6 +60,9 @@ def struphy_test(
             f"{LIBPATH}/post_processing/tests/",
             f"{LIBPATH}/propagators/tests/",
             f"{LIBPATH}/simulation/tests/",
+            # end-to-end GPU tests (NumPy reference, and the CuPy comparison on a GPU or the fake CuPy)
+            f"{LIBPATH}/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py",
+            f"{LIBPATH}/models/tests/test_gpu_VlasovAmpereOneSpecies.py",
         ]
 
         if mpi > 1:

src/struphy/feec/preconditioner.py

@@ -332,7 +332,10 @@ def weight_1d(c: int, d: int) -> Callable | xp.ndarray:
         return _kronecker_approximation(mass_operator, bc, weight_1d)
 
     def _core_diagonal(self):
-        return self._mass_operator.matrix.diagonal()
+        # diagonal of the mass matrix, including the band entries that wrap onto the main diagonal
+        # (StencilMatrix.diagonal() keeps only the entry at offset 0)
+        M = self._mass_operator.matrix
+        return _diagonal_operator(M.domain, _local_diagonal(M))
 
     @property
     def dim_reduce(self) -> int | None:
@@ -870,13 +873,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
@@ -1156,7 +1157,7 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]:
     if isinstance(op, StencilDiagonalMatrix):
         return [op._data]
     if isinstance(op, StencilMatrix):
-        return [op.diagonal()._data]
+        return [_stencil_local_diagonal(op)]
     if isinstance(op, KroneckerStencilMatrix):
         # outer product of the diagonals of the (process-local) factors on the local rows
         diags = []
@@ -1165,9 +1166,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()]
@@ -1179,6 +1181,31 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]:
     raise NotImplementedError(f"Diagonal of {type(op).__name__} is not supported.")
 
 
+def _stencil_local_diagonal(M: StencilMatrix) -> xp.ndarray:
+    """
+    Local diagonal of a StencilMatrix, including the band entries that wrap onto the main diagonal.
+
+    In a periodic direction with fewer points than the band width (``2p + 1 > n``, e.g. a single
+    element), the band entries at the offsets ``k = m n`` (``m != 0``) are on the main diagonal too;
+    :meth:`StencilMatrix.diagonal` returns only the entry at ``k = 0``.
+    """
+    V = M.codomain
+    assert all(m == 1 for m in V.shifts) and all(m == 1 for m in M.domain.shifts), "Only shifts 1 are supported."
+    offsets = [
+        [k for k in range(-p, p + 1) if k % n == 0] if periodic else [0]
+        for n, p, periodic in zip(V.npts, M.pads, V.periods)
+    ]
+    diag = M.diagonal()._data
+    if all(o == [0] for o in offsets):
+        return diag
+    idx = M._get_diagonal_indices()
+    rows, cols = idx[: V.ndim], idx[V.ndim :]
+    for ks in itertools.product(*offsets):
+        if any(ks):
+            diag = diag + M._data[rows + tuple(j + k for j, k in zip(cols, ks))]
+    return diag
+
+
 def _diagonal_operator(space: VectorSpace, diags: list[xp.ndarray]) -> LinearOperator:
     """Diagonal operator on ``space`` (Stencil- or BlockVectorSpace) with the given local diagonals."""
     if isinstance(space, BlockVectorSpace):
@@ -1450,27 +1477,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.
@@ -1489,21 +1523,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
 
@@ -1556,7 +1593,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))
 
@@ -1885,26 +1923,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
@@ -1923,13 +1966,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
@@ -1938,6 +1981,7 @@ def solve(self, rhs, out=None, transposed=False):
 
         from scipy.linalg import solve_circulant
 
+        assert isinstance(rhs, np.ndarray), f"FFTSolver solves on the host, got {type(rhs).__name__}."
         assert rhs.T.shape[0] == self._column.size
 
         if out is None:
@@ -1956,19 +2000,19 @@ 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
     -------
     circulant : bool
         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]))