Skip to content

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

Merged
max-models merged 4 commits into
devel-tinyfrom
device-kronecker-solve
Oct 8, 2026
Merged

Kronecker solver on the device (no host copies in CuPy solves)#96
max-models merged 4 commits into
devel-tinyfrom
device-kronecker-solve

Conversation

@max-models

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

Copy link
Copy Markdown
Member

Stack: part 1 of 3 — based on devel-tiny, next: #97 (stencil-6d-views), then #98. Merge this one first.

Companion struphy stack: struphy-hub/struphy#703 → struphy-hub/struphy#706 → struphy-hub/struphy#716; struphy-hub/struphy#703 points the feectools submodule at this branch's head (64b10fa).

Corresponding PR in struphy: struphy-hub/struphy#703 (moves the submodule here)

Summary

KroneckerLinearSolver now solves device data on the device. Its 1D solvers (BandedSolver: LAPACK ?gbtrs; SparseSolver: SuperLU) are host-only, so on the CuPy backend every solve copied the data of each direction to the host and back. struphy's MassMatrixPreconditioner applies such a solve in every CG iteration, for example in the preconditioner of EfieldWeightsCoupling in LinearVlasovAmpereOneSpecies (struphy-hub/struphy#689, part of struphy-hub/struphy#650).

With this PR, a 3D solve on device data makes no host/device transfer. Before, a 3D solve on cunumpy's fake CuPy made 3 counted to_host copies (one per direction) plus 3 uncounted copies back. The host path is unchanged.

Design

Dense inverses of the 1D matrices, applied with one GEMM. The new direct_solvers.DenseInverse inverts a 1D matrix once on the host, using the solver itself: M = solver.solve(I). On device data it then computes rhs @ M, a single matrix product (cuBLAS) over all right-hand sides of a pass. Why this approach:

  • The 1D matrices are small, n ~ tens to hundreds. So M is cheap to store, and an n x n GEMM over m right-hand sides suits a GPU better than m banded triangular solves.
  • Clamped (banded) and periodic matrices are handled the same way. Periodic matrices have corners, so in band storage they already use the full bandwidth.
  • M comes from the host solver itself. Pivoting, transposition (solver.transpose()) and complex dtypes therefore behave exactly as on the host path. Results agree with the host path to round-off: inverse-times-vector and LU solves are both backward stable, and the mass matrices are well conditioned. The tests use 1e-12 relative to the solution.
  • cupyx.scipy.linalg offers dense lu_factor/lu_solve but, as far as I know, no solve_banded (not checked against a CuPy install here). lu_solve would also need the dense matrix, and it runs two triangular solves instead of one GEMM.

Works with any linear 1D solver. The Kronecker passes build M from solver.solve on host arrays. struphy's FFTSolver (circulant mass matrices, scipy.linalg.solve_circulant) therefore works with no struphy change. I checked this in a local script with fake CuPy and assert_no_transfers. A 1D solver that should not be replaced by its dense matrix sets dense_on_device = False. The FFT "solvers" in linalg/fft.py do this, since an FFT is O(n log n); they still stage through the host as before.

Decided by the array, as before. xp.is_gpu(array) selects the dense path, and LAPACK/SuperLU are used otherwise.

Copied to the device once, at setup. A KroneckerLinearSolver with device temporaries (CuPy backend) builds its inverses and copies them to the device in its constructor. transpose() builds a new solver, so the same applies there. solve itself therefore makes no transfer. A standalone BandedSolver/SparseSolver called with device arrays copies its inverse on the first device solve (one to_device) and caches it per dtype.

Changes

  • linalg/direct_solvers.py: new DenseInverse. BandedSolver.solve and SparseSolver.solve use it for device arrays instead of a host round trip; their host branches no longer call to_numpy/asarray.
  • linalg/kron.py: KroneckerSolverSerialPass.solve_pass, also used between the two Alltoallv of the distributed pass, applies the dense inverse to device data. New prepare_device on both passes, called by the KroneckerLinearSolver constructor when its temporaries are on the device.
  • linalg/fft.py: DistributedFFTBase.OneDimSolver.dense_on_device = False.
  • linalg/tests/test_kron_device_solve.py (new), see Testing.
  • CUDA_STRATEGY.md: new section "Kronecker solver on the device"; open questions on FFT on the device and on large 1D matrices.

Testing

Locally (macOS, Open MPI, cunumpy 0.6.1, pyccel kernels compiled; no GPU):

  • New tests, test_kron_device_solve.py:
    • CPU: DenseInverse against the host solvers (banded/sparse, clamped/periodic, transposed, real/complex), and the Kronecker solve against a global np.kron reference (1D to 3D, mixed periodic/clamped, transposed). Serial: 36 passed. mpirun -n 2: 6 passed (distributed passes with Alltoallv).
    • Fake CuPy: run_device_checks runs in a serial subprocess with CUNUMPY_FAKE_CUPY=1, MAYBEMPI=0 and the MPI launcher variables removed. It compares device and host solves (Kronecker and 1D, with and without transposition) and wraps every device solve in assert_no_transfers. 1 test, passed. It fails on the old code with 3 to_host transfers per 3D solve.
    • GPU (requires_cupy): device against host solves under assert_no_transfers, serial (28) and MPI (6). Not run, skipped here (no GPU). They need a run by hand on the H100 before merging.
    • Also checked by hand with a script that is not part of the PR: mpirun -n 2 with fake CuPy, so the distributed passes with device buffers ran with no transfers and matched the host results.
  • feectools/linalg/tests, serial -m "not mpi and not petsc": before 7942 passed / 31 skipped, after 7979 passed / 59 skipped. mpirun -n 2 ... -m "mpi and not petsc" --with-mpi: before 738 passed, after 744 passed / 6 skipped.
  • Full feectools suite, after: serial 9493 passed, 64 skipped, 6 failed. The 6 failures are the known ones in ddm/tests/test_cart_2d.py/test_cart_3d.py, not caused by this PR. mpirun -n 2: 894 passed, 10 skipped.

🤖 Generated with Claude Code

max-models and others added 2 commits October 7, 2026 18:42
@max-models @claude
Kronecker solver on the device (no host copies in CuPy solves)
04dda21 
KroneckerLinearSolver copied the data of every direction to the host and
back in every solve on the CuPy backend, because its 1D solvers
(BandedSolver: LAPACK gbtrs, SparseSolver: SuperLU) are host-only.

- direct_solvers.DenseInverse: the 1D matrix is inverted once on the
  host by the solver itself (M = solver.solve(I)) and applied to device
  data as rhs @ M (one GEMM per pass).
- KroneckerSolverSerialPass (also used inside the distributed pass) uses
  it for device arrays; the inverses are copied to the device when a
  KroneckerLinearSolver with device temporaries is built. Works for any
  linear 1D solver (e.g. struphy's FFTSolver); the FFT "solvers" opt out
  with dense_on_device = False.
- BandedSolver/SparseSolver.solve use it for device arrays (copied on the
  first device solve). Host path unchanged.
- Tests: CPU (dense inverse vs host solvers, Kronecker solve vs global
  reference, serial and MPI), fake CuPy subprocess with
  assert_no_transfers, GPU tests (serial and MPI, skipped without GPU).
- CUDA_STRATEGY.md: notes for this change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@max-models @claude
CUDA_STRATEGY.md: name the PR (#96)
9cb793b 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@max-models max-models mentioned this pull request Oct 7, 2026
Closed
spossann
spossann reviewed Oct 7, 2026
View reviewed changes
Comment thread feectools/linalg/direct_solvers.py
@max-models @claude
DenseInverse.solve: say why the product is rhs @ M (rows, M = op(A)^-…
64b10fa 
…T, no symmetry assumed)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@max-models max-models mentioned this pull request Oct 7, 2026
Draft
max-models added a commit that referenced this pull request Oct 7, 2026
@max-models @claude
Merge device-kronecker-solve (#96) into stencil-6d-views
d16a6ad 
Stack #96 -> #97: keep both CUDA_STRATEGY.md sections (Kronecker solver first).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
max-models added a commit that referenced this pull request Oct 7, 2026
@max-models @claude
Merge stencil-6d-views (#97) into inner-on-device
a8bde31 
Stack #96 -> #97 -> #98: test_mpi_device.py keeps both new tests.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
max-models added a commit that referenced this pull request Oct 7, 2026
@max-models @claude
Merge stencil-6d-views (#97, now containing #96) into inner-on-device
6d88806 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 7, 2026
Merged
Merged
Closed
Merged
@max-models
max-models added this pull request to stack #99 October 8, 2026 05:42
@max-models
max-models marked this pull request as draft October 8, 2026 05:44
@spossann

spossann commented Oct 8, 2026

Copy link
Copy Markdown
Member

@max-models do you want to merge this before my Kronecker (I am still working on it)?

@max-models

max-models commented Oct 8, 2026

Copy link
Copy Markdown
Member Author

@max-models do you want to merge this before my Kronecker (I am still working on it)?

Sure!

@max-models
max-models marked this pull request as ready for review October 8, 2026 09:26
@max-models
max-models requested a review from spossann October 8, 2026 09:26
@max-models

max-models commented Oct 8, 2026

Copy link
Copy Markdown
Member Author

@max-models do you want to merge this before my Kronecker (I am still working on it)?

This stack (and struphy-hub/struphy#716 on the struphy side) is ready for review.

@max-models
max-models merged commit 7f4bcbd into devel-tiny Oct 8, 2026
9 checks passed
spossann added a commit to struphy-hub/struphy that referenced this pull request Oct 9, 2026
@max-models @claude @spossann
feectools: inner products stay on the device (#716)
9ce5491 
**Stack:** part 3 of 3 — based on #706 (`feectools-stencil-6d-views`,
merge that first, after #703), next: none (top of the stack).

Corresponding update in feectools:
struphy-hub/feectools#98. Mirrors the feectools
stack struphy-hub/feectools#96 →
struphy-hub/feectools#97 →
struphy-hub/feectools#98.
`feectools-stencil-6d-views` is merged into this branch (no rebase), and
the `feectools` submodule now points at the stacked head of
struphy-hub/feectools#98, `6d88806`, which
contains feectools#96 and #97. Against its base (#706) this PR only
moves the submodule `d16a6ad` → `6d88806`. After each feectools PR of
the stack is merged into `devel-tiny`, the `feectools` submodule of the
struphy stack must be moved to the corresponding merged `devel-tiny`
commit (the `pr-feectools-submodule` check fails until then).

---

**Solves the following issue(s):**

Part of #689 (`VlasovAmpereOneSpecies` end to end on the GPU, no
host/device transfers in the time loop), tracked in #650. This PR moves
the `feectools` submodule to struphy-hub/feectools#98 ("Inner products
stay on the device"). On the CuPy backend, every CG iteration of a
struphy solve copied two scalars from the device to the host; after this
change it copies one (the convergence test).

**The `pr-feectools-submodule` check fails until
struphy-hub/feectools#98 is merged** into `devel-tiny`. After that, the
submodule should point at the merge commit.

**Core changes:**

- `feectools` submodule: `d16a6ad` (base #706) → `6d88806` (branch
`inner-on-device` of struphy-hub/feectools#98, stacked on feectools#96
and #97; originally `a15a8e8`). No struphy code changes.
- What changes in feectools, and only on the CuPy backend:
- `StencilVector.inner`/`BlockVector.inner`/`dot_inner` return a 0-d
device array, in serial and with MPI.
  - `axpy` accepts such a scalar without copying it to the host.
- CG, PCG, BiCG, BiCGStab and PBiCGStab keep alpha/beta on the device
and copy only the residual norm, once per iteration.

  On NumPy, results and types are unchanged.

  Copies to the host per iteration, counted under cunumpy's fake CuPy:

  | Solver | before | after |
  | --- | --- | --- |
  | CG | 2 | 1 |
  | PCG | 3 | 1 |
  | BiCG | 4 | 1 |
  | BiCGStab | 6 | 1 |
  | PBiCGStab | 5 | 1 |

- struphy call sites of `.inner` that run on CuPy now get a 0-d device
array. I checked them:
- the scalars in `models/scalars.py` write it into `xp` buffers
(`local_value[0] = ...`), which works on the device;
- the model energy methods (`linear_mhd.py`, `shear_alfven.py`, ...)
return it, and the scalar machinery handles it like the existing
`dot_inner` results;
- `PolarVector.dot` adds a NumPy scalar to it, which works, but polar
splines are not on CuPy yet (#695);
- the multigrid smoothers run their own PCG with `.inner`. It works with
device scalars: the arithmetic stays on the device, and the comparisons
with 0 copy to the host implicitly.
- This removes the first blocker listed in #705 ("CG inner products ...
The Schur solve is outside the transfer guard because of this"). The
remaining copy per CG iteration still counts as a transfer for
`assert_no_transfers`. The transfer guard could then cover the Schur
solve with an allowance of one 8-byte copy per iteration.

**Model-specific changes:**

None.

**Documentation changes:**

None in struphy. The feectools PR documents the change in feectools'
`CUDA_STRATEGY.md` (section "Inner products on the device").

**Testing** (macOS, no GPU; struphy kernels compiled with GNU/Fortran;
**GPU tests not run**):

- feectools (see struphy-hub/feectools#98):
- serial suite: 9456 → 9480 passed, with the same 6 unrelated failures
in `ddm/tests/test_cart_*d.py`;
  - `mpirun -n 2` linalg MPI suite: 738 → 739 passed;
- the new fake-CuPy tests check the copy counts and that the results
equal NumPy's.
- struphy, solver tests on the NumPy backend, with feectools
`devel-tiny` (before) and this branch (after):
- `propagators/tests/test_poisson.py -m "not mpi" -k "not multigrid"`:
40 passed before, 40 passed after;
- `linear_algebra/tests/test_saddlepoint_massmatrices.py`: the first
case (`SaddlePointSolverUzawaNumpy`) passed before and after.

I stopped the remaining saddle-point and multigrid tests: with the
feectools kernels uncompiled in the checkouts, each test took more than
10 minutes. On NumPy, feectools' change is limited to identity helpers
and an equivalent loop structure in BiCGStab, and the feectools solver
tests check the same iteration counts and results.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Stefan Possanner <stefan.possanner@ipp.mpg.de>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@spossann spossann spossann approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@max-models @spossann