[gpu-1] End-to-end GPU test of VlasovAmpereOneSpecies on 1 and 2 ranks

[max-models]

Solves the following issue(s):

Part of #689, tracked in #650. This PR covers the items of #689 that are left for VlasovAmpereOneSpecies: running on 1 and 2 ranks, no host/device transfers in the time loop, and the diagnostics it needs. It does not close #689, because these runs use the CPU emulation; the first run on a real GPU is #687.

Stack: this is gpu-1 of a new stack. Merge it after #722 (cuda-namespaces), #716, #717, #719 and #765 (MPI sort of the markers after drawing, needed for the 2-rank runs). It merges those branches in, so until they are merged, this diff also shows their commits. gpu-2 (the benchmark script) is based on this branch.

Core changes:

• New test models/tests/test_gpu_VlasovAmpereOneSpecies.py. A small Landau-damping run on NumPy and on CuPy: 8 cells, 320 Sobol markers, 3 steps, full-f weights, and an e1_v1 binning plot so that the diagnostics run on the device too. It compares the energies, the e_field and phi coefficients, and the markers by ID (tolerances explained in compare_runs), and checks the transfer budget of the time loop.
• Shared helpers in models/tests/gpu_e2e.py. The transfer budget, the instrumentation, gathering markers across ranks and the device context, moved out of Fix sobol random numbers on cupy #719's test. test_gpu_LinearVlasovAmpereOneSpecies.py now uses them.
• The GPU tests also run without a GPU. On cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1), every CUDA launch is emulated on the CPU and CudaKernel.compile is skipped. Simulation.compile_cuda_kernels still checks that every kernel of the time loop has a CUDA version. Without a GPU and without the fake CuPy, test_fake_cupy_matches_numpy runs the comparison in a serial child process.
• Transfer budget of model.integrate. Only scalar downloads are allowed: the residual norm of each Krylov iteration, which feectools copies for its convergence test (feectools#98 keeps everything else of the solve on the device). No arrays, no host-to-device copies, no host kernels.
• Bug fix (every backend): VlasovAmpereCoupling updates the ghost regions of the mid-point field before the push. The pusher kernel reads _e_sum._data including its ghost regions, and _e_sum is built with += from the solver output, which does not sync them. Stale ghost regions would give a wrong field near process (and periodic) boundaries.
• Particles. Checking whether MPI is CUDA-aware now happens when the marker array is allocated (in the setup), not on the first marker exchange. Before, that check was the only array transfer inside the time loop on 2 ranks.

Results (fake CuPy, every CUDA launch emulated on the CPU; macOS; not run on a GPU):

test	1 rank	2 ranks
test_gpu_VlasovAmpereOneSpecies.py (NumPy reference + CuPy vs NumPy + transfer budget)	passed	passed
test_gpu_LinearVlasovAmpereOneSpecies.py::test_cupy_matches_numpy_without_transfers	passed	not run

Transfers in the time loop of the CuPy run:

• integrate: one 8-byte residual norm per pcg iteration.
• Diagnostics: one 24-byte copy of the three scalars per update.
• Output: one copy per saved dataset.

Diagnostics: no new CUDA kernels are needed for this model. The scalars use xp reductions and feectools' device inner products, and binning uses xp.histogramdd.

Found along the way (not changed here): now #766 (Poisson stabilization) and #767 (energy diagnostic with the control variate).

• The initial Poisson solve with the default stab_eps=1e-14 puts (Monte-Carlo net charge) / 1e-14 into the constant mode of phi. That makes phi about 1e14, so E = -grad(phi) loses about 5% to cancellation on either backend. The verification tests use this default. The new test uses stab_eps=1e-6, as Fix sobol random numbers on cupy #719 does.
• With control_variate=True, kinetic_energy at t=0 is computed with the full-f weights and afterwards with the delta-f weights, so total_energy jumps after the first step. With full-f weights the total is conserved from step to step to 1e-9, but the t=0 value still differs by about 2e-4 (not investigated).

Review fixes (pushed by @spossann, see the review):

• the e2e tests run in CI: unit shard-2 (serial and 2 ranks) and on the fake CuPy on 1 and 2 ranks (test-PR-unit-mpi.yml); merged draw_markers: send the markers to their processes also without sorting boxes #765, without which they segfault on 2 ranks (draw_markers does not MPI-sort the markers without sorting boxes (off-rank markers in setup accumulations, segfault on 2 ranks) #764);
• MassMatrixPreconditioner._core_diagonal returns a diagonal operator again (with the wrapped band entries);
• test_preconditioner_cupy.py uses cunumpy's emulated_launches and run_fake_cupy_child, and its GPU test emulates its launches on the fake CuPy;
• FFTSolver.solve asserts host input.

Documentation changes:

CUDA_STRATEGY.md: the end-to-end entry now covers both models, the CPU emulation and the transfer budget.

🤖 Generated with Claude Code

[spossann]

Thanks, the design is solid: the transfer budget is well specified, and the tolerances in compare_runs are explained. I reran the e2e tests locally at the PR head (727083eb), in a fresh env with all kernels compiled. Linux, Open MPI 4.1.6, no GPU:

run	result
serial, no fake CuPy (test_fake_cupy_matches_numpy child)	passed
serial, CUNUMPY_FAKE_CUPY=1	passed (both models, budget included)
mpirun -n 2, fake CuPy	segfault on rank 1, in every test, including test_numpy_reference
mpirun -n 2, plain NumPy	segfault
mpirun -n 2, fake CuPy, with an mpi_sort_markers() after draw_markers (patched in the test run only)	passed

Blocking

1. The 2-rank run crashes: markers are never MPI-sorted after loading when there are no sorting boxes. Particles.draw_markers calls mpi_sort_markers() only if self._initialized_sorting and sort, i.e. only with sorting boxes. Both e2e tests use SortingParameters(), so after sobol_standard loading, 80 of the 160 markers of each rank lie outside the rank's domain. The initial Poisson accumulation in post_allocate runs before the first PushEta sorts them, and writes them outside the local stencil arrays. The crash then shows in exchange_assembly_data. On macOS this may have corrupted memory without crashing, so I would not rely on the 2-rank result in the description. The bug is already on devel; see #764. It needs its own fix in draw_markers, and this PR should be rebased on that fix.

2. CI does not run the new tests. models/tests/test_gpu_*.py are top-level files in models/tests/, and no workflow collects that folder: the model workflows run only default_params/ and verification/, and the unit shards don't include models/. The green CI therefore says nothing about this PR's main deliverable. Please add them like the existing fake-CuPy lines in test-PR-unit-mpi.yml (mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi models/tests/test_gpu_*.py), plus a serial shard and console/test.py. A 2-rank job would have caught point 1.

3. MassMatrixPreconditioner._core_diagonal breaks its contract. The base class documents it as "diagonal D of the core operator as diagonal operator", and StiffnessPreconditioner still returns one. It now returns the full mass matrix. That works today only because the single caller passes it to _local_diagonal; the next caller who trusts the name and calls .dot gets M instead of D, with no error. Keep the wrap-around fix and the contract:

def _core_diagonal(self):
    # diagonal of the mass matrix, including the band entries that wrap onto the main diagonal
    # (StencilMatrix.diagonal() keeps only the entry at offset 0)
    M = self._mass_operator.matrix
    return _diagonal_operator(M.domain, _local_diagonal(M))

_stencil_local_diagonal itself looks right to me: the band centre is at p for shift 1, and the sum over wrapped offsets is correct for a tensor product. The underlying gap is in feectools' StencilMatrix.diagonal(), which affects the Jacobi preconditioners there as well; that is worth a feectools issue.

Should fix

4. The ghost-region update in VlasovAmpereCoupling is a bug fix on every backend, not only on CuPy. The pusher reads _e_sum._data including the ghost regions, and += does not sync them. Please list it under bug fixes in the description and changelog.
5. Particles._allocate_marker_array now runs a collective probe (_mpi_cuda_aware). That is safe for the current callers (__init__ and the lazy n_rows), but please add a comment that all ranks of mpi_comm must call it, so that a future per-rank reallocation doesn't hang.
6. feec/tests/test_preconditioner_cupy.py re-implements cunumpy.kernel_testing.emulated_launches (_emulated_launches, _host_buffer, _HostFields, which read fake-CuPy internals via object.__getattribute__(value, "_a")). gpu_e2e.device_context already uses the cunumpy version. That file also starts its subprocess on every rank under mpirun, instead of run_fake_cupy_child, which runs only on rank 0 and was added to avoid exactly the concurrent-children crash.

Minor

7. FFTSolver is host-only now (_space = np.ndarray); an assert isinstance(rhs, np.ndarray) in solve would make a device caller fail loudly instead of crashing deep inside.
8. On the fake CuPy, the setup falls back to host kernels for the mass-matrix assembly (kernel_1d_mat, kernel_3d_vec, RuntimeWarning: No CUDA version of kernel). This is outside the time loop, so the budget is unaffected, but worth listing as open work in CUDA_STRATEGY.md.
9. The two findings in the description (stab_eps=1e-14 puts phi ≈ 1e14; total_energy jumps with the control variate because t=0 uses full-f weights) deserve their own issues.

🤖 Generated with Claude Code

[spossann]

@max-models I pushed two commits to this branch to address my review:

• d356ad9: review fixes.
  • The e2e tests now run in CI: unit shard-2 (serial and 2 ranks, in both workflows and in struphy test unit), and on the fake CuPy on 1 and 2 ranks in test-PR-unit-mpi.yml. 8 elements of degree 3 allow at most 2 ranks.
  • MassMatrixPreconditioner._core_diagonal returns a diagonal operator again, including the wrapped band entries.
  • test_preconditioner_cupy.py uses cunumpy's emulated_launches and run_fake_cupy_child. Its GPU test also failed under CUNUMPY_FAKE_CUPY=1 (it tried to compile CUDA); it now emulates its launches there.
  • FFTSolver.solve asserts host input; added a comment on the collective probe; CUDA_STRATEGY.md updated.
• 1a02d0c: merge of draw_markers: send the markers to their processes also without sorting boxes #765 (draw-markers-mpi-sort), the fix for draw_markers does not MPI-sort the markers without sorting boxes (off-rank markers in setup accumulations, segfault on 2 ranks) #764. Without it, the 2-rank unit job segfaulted in test_numpy_reference once the e2e tests were in CI. So this PR now has to be merged after draw_markers: send the markers to their processes also without sorting boxes #765.

Verified locally on Linux (Open MPI 4.1.6):

• e2e: serial, and 2 ranks on NumPy, pass (3 passed, 2 skipped); fake CuPy on 1 and 2 ranks passes (4 passed, 1 skipped);
• full shard-2 on 2 ranks: 396 passed;
• -m mpi_pic pic/tests on 2 ranks: 41 passed;
• preconditioner and mass-matrix tests: 66 serial, 100 on 2 ranks.

I also updated the description: the ghost-region update in VlasovAmpereCoupling is listed as a bug fix on every backend, and the two side findings are now #766 and #767.

Still yours: please check that the changes fit your plans for gpu-2.

🤖 Generated with Claude Code