Repository navigation
MassMatrixPreconditioner on the CuPy backend - #717
Draft
max-models wants to merge 3 commits into
Draft
max-models wants to merge 3 commits into
max-models wants to merge 3 commits into
Conversation
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>
…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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Solves the following issue(s):
Part of #689 (
LinearVlasovAmpereOneSpecieson the GPU), tracked in #650.EfieldWeightsCouplingusesMassMatrixPreconditionerby 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):
StencilMatrix.toarray()returns a host (NumPy) array on every backend, because feectools builds it through a scipy COO matrix.is_circulantandFFTSolver.__init__assertedxp.ndarray, which iscupy.ndarrayon CuPy. The next steps would have failed as well:xp.nonzero(M_arr)on a host array and writing host values into the device_dataof the process-local stencil matrix. With diagonal scaling (MassMatrixDiagonalPreconditioner),xp.multiply.outerin_local_diagonalis not available on every backend.Core changes (on top of the refactor of #731; rebased on
develwith feectools 0.7.0):_kronecker_approximation) and the 1d solvers (_solver_1d:FFTSolverfor circulant matrices,SparseSolverotherwise).KroneckerLinearSolverexpects 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.KroneckerStencilMatrix._process_local_matrix_1dfills the band in a host array and copies it toM_local._datain one step.StiffnessPreconditionerpasses its host 1d matrices to it directly.is_circulantaccepts host or device matrices and checks on the host.FFTSolveris a host solver: it keeps its circulant column on the host and accepts a host or device matrix. Itssolveis only called with host arrays (byDenseInverseon CuPy), so the host round trip of the earlier version of this PR is gone._local_diagonal(diagonal scaling): outer product by broadcasting instead ofxp.multiply.outer.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) givesKroneckerLinearSolverdense "solvers" that hold device matrices, from which it tries to build dense inverses on the host (numpy @ cupyfails). WithKroneckerSumSolver._DenseApply.dense_on_device = Falsein 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 theStiffnessPreconditionergap.Tests: new
feec/tests/test_preconditioner_cupy.py.check_preconditioners_on_cupybuildsMassMatrixPreconditioner(defaults, andweight_reduction="average"withdiagonal_scaling=True) andMassMatrixDiagonalPreconditionerfor 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 bothFFTSolverandSparseSolverare used).dotanddot(out=...)insidecunumpy.profiling.assert_no_transfers(), checks that the results live on the device, and compares with NumPy atrtol=1e-12.CudaKernellaunch is emulated on the CPU through a small private_emulated_launches(). It is a minimal version ofemulated_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.Testing (macOS, no GPU; feectools 0.7.0 at devel-tiny 5385af9, pyccel kernels compiled with GNU/Fortran; GPU tests not run):
test_preconditioner_cupy.py(fake CuPy)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)Not run locally:
test_mass_preconditioner_polarand the MPI runs. CI covers them.🤖 Generated with Claude Code