Skip to content

MassMatrixPreconditioner on the CuPy backend - #717

Draft
max-models wants to merge 3 commits into
develfrom
mass-preconditioner-cupy
Draft

max-models wants to merge 3 commits into
develfrom
mass-preconditioner-cupy

Conversation

@max-models

@max-models max-models commented Oct 7, 2026 •

Copy link
Copy Markdown
Member

Solves the following issue(s):

Part of #689 (LinearVlasovAmpereOneSpecies on the GPU), tracked in #650. EfieldWeightsCoupling uses MassMatrixPreconditioner by default, and the preconditioner could not even be created on the CuPy backend.

The failure (reproduced with cunumpy's fake CuPy, CUDA launches emulated on the CPU):

File "src/struphy/feec/preconditioner.py", in _solver_1d
    if is_circulant(M_arr):
File "src/struphy/feec/preconditioner.py", in is_circulant
    assert isinstance(mat, xp.ndarray)
AssertionError

StencilMatrix.toarray() returns a host (NumPy) array on every backend, because feectools builds it through a scipy COO matrix. is_circulant and FFTSolver.__init__ asserted xp.ndarray, which is cupy.ndarray on CuPy. The next steps would have failed as well: xp.nonzero(M_arr) on a host array and writing host values into the device _data of the process-local stencil matrix. With diagonal scaling (MassMatrixDiagonalPreconditioner), xp.multiply.outer in _local_diagonal is not available on every backend.

Core changes (on top of the refactor of #731; rebased on devel with feectools 0.7.0):

  • The 1d setup data stay on the host, on every backend. These are the dense 1d mass matrices (_kronecker_approximation) and the 1d solvers (_solver_1d: FFTSolver for circulant matrices, SparseSolver otherwise). KroneckerLinearSolver expects host 1d solvers. Since feectools#96 (in 0.7.0), it solves device data with dense inverses that it builds once from these host solvers, so applying the preconditioners makes no host copies.
  • Only the process-local 1d stencil matrices go to the device. They are the factors of the KroneckerStencilMatrix. _process_local_matrix_1d fills the band in a host array and copies it to M_local._data in one step. StiffnessPreconditioner passes its host 1d matrices to it directly.
  • is_circulant accepts host or device matrices and checks on the host.
  • FFTSolver is a host solver: it keeps its circulant column on the host and accepts a host or device matrix. Its solve is only called with host arrays (by DenseInverse on CuPy), so the host round trip of the earlier version of this PR is gone.
  • _local_diagonal (diagonal scaling): outer product by broadcasting instead of xp.multiply.outer.
  • NumPy path: same arithmetic as before.

Model-specific changes:

None. Still missing for #689: CG inner products returning host scalars, and an H100 run.

StiffnessPreconditioner (#731) builds its 1d data the same way after this PR, but does not run on CuPy yet: KroneckerSumSolver (feectools 0.7.0) gives KroneckerLinearSolver dense "solvers" that hold device matrices, from which it tries to build dense inverses on the host (numpy @ cupy fails). With KroneckerSumSolver._DenseApply.dense_on_device = False in feectools, StiffnessPreconditioner (grad/curl/div, diagonal_scaling=False) gives the NumPy result on fake CuPy with no transfers. Follow-up in feectools.

Documentation changes:

CUDA_STRATEGY.md: a short note in the feectools section on where the preconditioner data live, the device Kronecker solve, and the StiffnessPreconditioner gap.

Tests: new feec/tests/test_preconditioner_cupy.py.

  • check_preconditioners_on_cupy builds MassMatrixPreconditioner (defaults, and weight_reduction="average" with diagonal_scaling=True) and MassMatrixDiagonalPreconditioner for M0, M1 and M2 (Colella, 6x5x4 elements, degrees (2, 3, 1)) on both backends. It covers periodic bcs and clamped bcs (Dirichlet in eta1/eta3, so the boundary-operator path and both FFTSolver and SparseSolver are used).
  • It applies them to the same random vector, dot and dot(out=...) inside cunumpy.profiling.assert_no_transfers(), checks that the results live on the device, and compares with NumPy at rtol=1e-12.
  • Without a GPU, the check runs in a subprocess on cunumpy's fake CuPy (clean serial environment). There, every CudaKernel launch is emulated on the CPU through a small private _emulated_launches(). It is a minimal version of emulated_launches() from CUDA version of the linear_vlasov_ampere accumulation #705 and can be replaced by it once CUDA version of the linear_vlasov_ampere accumulation #705 is merged.
  • On a GPU, the same check runs directly (skipped here).

Testing (macOS, no GPU; feectools 0.7.0 at devel-tiny 5385af9, pyccel kernels compiled with GNU/Fortran; GPU tests not run):

result
test_preconditioner_cupy.py (fake CuPy) 2 passed, 2 skipped (GPU)
test_preconditioner_transpose.py, test_stiffness_preconditioner.py, test_mass_matrices.py::test_reduced_weight_1d, ::test_mass_preconditioner, ::test_mass_preconditioner_array_weights_mpi (serial) 49 passed

Not run locally: test_mass_preconditioner_polar and the MPI runs. CI covers them.

🤖 Generated with Claude Code

The mass-matrix preconditioners could not be created on the CuPy backend:
the assembled 1d mass matrices come back from StencilMatrix.toarray() as
host (NumPy) arrays, and is_circulant/FFTSolver asserted xp.ndarray
(cupy.ndarray on CuPy). The stencil-matrix fill also mixed xp.nonzero
with host arrays.

The 1d matrices and solvers are now host setup data on every backend
(what KroneckerLinearSolver expects, and what feectools#96 builds its
device inverses from); only the process-local 1d stencil matrices of the
KroneckerStencilMatrix are copied to the device once. The duplicated 1d
setup of both preconditioners moves into _solver_and_local_matrix_1d.
FFTSolver accepts host or device matrices and solves device right-hand
sides on the host. The NumPy path computes the same as before.

Adds a fake-CuPy test (subprocess, CUDA launches emulated on the CPU)
comparing both preconditioners for M0/M1/M2 (periodic and clamped) with
NumPy, and the same check on a GPU.

Part of #689 and #650.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
max-models and others added 2 commits October 9, 2026 15:48
…ss-preconditioner-cupy

Conflict in src/struphy/feec/preconditioner.py: #731's refactor (KroneckerPreconditioner base class,
module-level helpers, StiffnessPreconditioner) is kept; the CuPy changes of this branch are
re-applied on top of it:

- _kronecker_approximation, _solver_1d, is_circulant, FFTSolver: the dense 1d matrices and the 1d
  solvers are host (NumPy) setup data on every backend. KroneckerLinearSolver (feectools#96, in
  0.7.0) solves device data with dense inverses built from these host solvers, so FFTSolver no
  longer needs the host round trip of the previous version of this branch.
- _process_local_matrix_1d: band filled on the host, copied to the factor's data in one step;
  StiffnessPreconditioner passes its host 1d matrices directly.
- _local_diagonal: outer product by broadcasting instead of ufunc.outer (not available on every
  backend), so diagonal scaling (MassMatrixDiagonalPreconditioner) works on CuPy.

test_preconditioner_cupy.py also covers weight_reduction="average" with diagonal_scaling and
checks that applying the preconditioners makes no host/device transfers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MassMatrixDiagonalPreconditioner was removed on devel (#739); the CuPy test now checks
MassMatrixPreconditioner with and without diagonal scaling and with dim_reduce=None instead.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant