Device ghost-exchange plans, fused block exchanges, single-reduction PCG (CUDA-aware MPI)

[max-models]

GPU and MPI performance work for struphy's CuPy backend. This PR must be merged into devel-tiny before the struphy stack that uses it (struphy-hub/struphy#782 bumps the submodule to this branch; it is the first of the stacked PRs struphy-hub/struphy#782 to struphy-hub/struphy#789, which split struphy-hub/struphy#772).

Changes

Device ghost exchange

• ddm/basic.py: exchange plans built once from per-region indices (_DeviceRegionExchange), and a fused exchange of all blocks of a BlockVector (_FusedDeviceExchange): one message per neighbour and one gather and one scatter kernel for all blocks.
• blocking_data_exchanger.py, nonblocking_data_exchanger.py: use these plans for device arrays.
• stencil.py: serial ghost update of vectors and matrices with one kernel on the GPU (_serial_ghost_plan); inner_local / inner_comm for fused reductions.
• block.py: the exchanges of all blocks are in flight together; BlockLinearOperator.dot writes the first block of each row directly.

Solvers

• New PConjugateGradientSingleReduction (inverse(A, "pcg_sr", ...), Chronopoulos–Gear): all inner products of an iteration in one reduction, so one Allreduce per iteration instead of two.
• PCG computes r·r and s·r with one reduction (_inner_products).

Kernels

• stencil_dot_3d (CUDA): 8 threads per row, coalesced reads, warp-shuffle sum.
• IdentityOperator.dot with out: one copy (the old out *= 0; out += v also kept NaNs of out).

Fixes added while splitting the struphy branch

• stencil.py imported feectools.utilities.cached_launch, which was never committed: the .gitignore rule *cache* ignored it, so the package did not import. It is written new (CachedLaunch: a launch with the same buffers reuses its prepared arguments and launch shape; only scalars such as the axpy coefficient are converted again; host calls, the fake CuPy, argument objects, streams, debug mode and check_finite go to the kernel unchanged), with .gitignore exceptions, and the stencil axpy and inner kernels use it. test_cached_launch checks it with a stub kernel, without a GPU.
• On cunumpy's fake CuPy (the GPU-less tests) the serial ghost plan and the fused block exchange launched raw kernels, which the fake cannot run. Both now fall back to the slice copies there.
• PCG had gone from one to two host copies per iteration on a single GPU, which test_one_host_copy_per_iteration caught. v·p is a device scalar again. With the fused r·r/s·r reduction PCG is back to one copy per iteration, with the same iterates.
• The emulation parity test skips stencil_dot_3d, because serial emulation cannot run warp shuffles. The GPU parity test covers it.
• On fake CuPy, StencilMatrix uses the previous one-thread-per-entry kernel (stencil_dot_3d_rows.cu) for device arguments. Without it, struphy's tests that emulate every launch (accumulation, polar and end-to-end GPU tests) failed at the first stencil product. test_emulated_stencil_dot_3d_rows compares it with the pyccel kernel.
• test_inner_on_device also covers pcg_sr.

Tests

Run on macOS without a GPU, with the cunumpy commit struphy pins (f996393). feectools/linalg/tests and feectools/ddm pass, apart from 22 failures that devel-tiny shows in the same environment (PETSc and similar setup issues). The same holds for test_mpi_device.py, test_solvers.py and test_block.py under mpirun -n 2. The fake-CuPy tests show one host copy per iteration for cg, pcg, pcg_sr and pbicgstab, with results identical to NumPy. Not tested on a GPU here: the CUDA kernels and the CUDA-aware exchange paths need a run on the cluster.

🤖 Generated with Claude Code