Inner products stay on the device

[max-models]

Stack: part 3 of 3 — based on #97 (stencil-6d-views, merge that first, after #96), next: none (top of the stack).

stencil-6d-views (which contains #96) is merged into this branch (merge commits a8bde31 and 6d88806, no rebase); the only conflict was feectools/linalg/tests/test_mpi_device.py, where both new tests are kept (test_fewer_diagonals_dot_and_transpose_match_global_reference from #97 and test_inner_result_is_a_copy_of_the_reduction_buffer from this PR). Tests run on the merged head 6d88806 (macOS, no GPU): test_kron_device_solve.py, test_inner_on_device.py, test_device_matvec.py, test_mpi_device.py serial (82 passed, 54 skipped), mpirun -n 2 ... --with-mpi test_mpi_device.py test_kron_device_solve.py (51 passed, 34 skipped), test_cuda_parity.py (4 passed, 49 skipped, GPU).

Companion struphy PR: struphy-hub/struphy#716 (part 3 of the struphy stack struphy-hub/struphy#703 → struphy-hub/struphy#706 → struphy-hub/struphy#716); its feectools submodule points at this branch's new head 6d88806, which contains #96 and #97.

---

Solves the following issue(s):

Part of the CUDA work tracked in struphy-hub/struphy#650, for running a model end to end on the GPU (struphy-hub/struphy#689: no host/device transfers in the time loop). #88 moved the inner reduction to the device, but in the serial case it still returned a NumPy scalar, so every CG iteration of a struphy solve copied two scalars from the device to the host. No issue closed.

Companion struphy PR (moves the feectools submodule to this branch): struphy-hub/struphy#716

Core changes

1. inner returns a device scalar for device vectors

StencilVectorSpace.inner (and so StencilVector.inner, BlockVectorSpace.inner, BlockVector.inner and dot_inner) returns a 0-d CuPy array when the data are on the device:

• serial case: no to_numpy any more;
• MPI case: unchanged Allreduce on the device buffers after cunumpy.mpi.synchronize_for_mpi, and the result stays on the device as well;
• the result is a copy (on the device) of the reduction buffer: the buffer belongs to the first vector and the next inner product with that vector overwrites it. Before, the MPI case on the device returned x._dot_recv_data[0], a view of that buffer.

On NumPy the result is a NumPy scalar, as before (the same value and type).

2. axpy with a device scalar

The axpy kernel takes alpha by value, so float(a) would copy a device scalar to the host and wait for the device. StencilVectorSpace.axpy with a 0-d device array computes y += a * x with array operations instead (also for the interface data); Python and NumPy scalars still go through the kernel. BlockVectorSpace.axpy delegates to it.

3. Solvers keep their scalars on the device (linalg/solvers.py)

The design: one 8-byte copy per iteration, for the convergence test; the step sizes stay on the device.

• CG, PCG, BiCG, BiCGStab, PBiCGStab: alpha, beta, omega, rho are 0-d device arrays; arithmetic between them and axpy/*= with them run on the device. The residual norm is copied explicitly with a helper _host (an identity on NumPy) in the convergence test and for get_info(), which holds host scalars as before.
• BiCGStab tested the residual twice per iteration (at the end of an iteration and again, redundantly, at the start of the next one). It now tests it once per iteration (and once before the first). Same iteration count.
• MINRES, LSMR and the Uzawa solver do their scalar recurrences (Givens rotations, norm estimates) in Python on the host; they copy every inner product with _host, i.e. as before.
• GMRES: the first residual norm is copied, the Arnoldi inner products and norms stay on the device (_sqrt). Its Givens rotations still call .item() per entry (unchanged).
• On NumPy every helper returns its argument unchanged; the only cost is one xp.is_gpu check (about 0.3 µs) per call.

I did not implement "test the residual every k iterations": it would remove most of the remaining synchronizations, but it changes iteration counts compared with NumPy, and after convergence the extra iterations can divide by a zero residual. It is listed as an open question in CUDA_STRATEGY.md.

Copies to the host per iteration, before → after

Counted under cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1) on a 1D SPD stencil matrix (40 points, degree 2). Device kernel launches run the host kernel on the fake device buffers. Two counts are used: cunumpy.profiling.count_transfers().to_host, and all device → host copies, including implicit float()/bool()/.item() conversions that count_transfers does not see. These are counted by wrapping the fake cupy.ndarray methods. Both counts agree below, except for GMRES.

Solver	before	after
CG	2	1
PCG	3	1
BiCG	4	1
BiCGStab	6	1
PBiCGStab	5	1
MINRES	3	3
LSMR	3	3
GMRES (15 iterations, all copies)	405	136

The iteration counts and solutions are identical to NumPy (bitwise in this test).

With MPI (mpirun -n 2, fake CuPy, Open MPI reading the host memory behind the fake device arrays), I checked with an ad-hoc script that is not committed: the parallel inner product returns a 0-d device array with the global value, and distributed CG makes 1 copy per iteration and converges to the exact solution.

Tests

New linalg/tests/test_inner_on_device.py (+ inner_on_device_child.py, the fake-CuPy child process):

• NumPy: inner returns NumPy scalars (stencil and block), and get_info() holds host scalars for every solver;
• fake CuPy: inner returns a 0-d device array that is independent of the reduction buffer, axpy with it makes no copy, the five Krylov solvers make exactly 1 copy per iteration and none to the device, and all eight solvers give the same iterations and solution as NumPy;
• GPU (skipped here): the same on CuPy with count_transfers.

New MPI test in test_mpi_device.py: the result of the reduction survives the next inner product with the same vector, and it is a NumPy scalar on NumPy or a 0-d device array on CuPy.

	before (devel-tiny)	after
pytest feectools -m "not mpi and not petsc"	9456 passed, 6 failed, 36 skipped	9480 passed, 6 failed, 42 skipped
mpirun -n 2 pytest feectools/linalg -m "mpi and not petsc" --with-mpi	738 passed	739 passed

The 6 failures are the same on both (ddm/tests/test_cart_2d.py, test_cart_3d.py without MPI) and are unrelated to this PR. The 6 new skips are the GPU tests.

GPU tests were not run (no GPU here). The device paths were exercised only through the fake CuPy.

Documentation changes

CUDA_STRATEGY.md: new section "Inner products on the device" (with the table above). The inner open question is updated, and a new open question covers testing the residual less often.

🤖 Generated with Claude Code