CUDA version of the vlasov_maxwell accumulation

CUDA_STRATEGY.md

@@ -290,7 +290,7 @@ done (✓), in a PR (PR 14), blocked (⏸).
 |---|---|---|---|---|
 | 1 | `push_v_with_efield` | P | `VlasovAmpereCoupling`, `PushVinForceField` | PR 14 |
 | 2 | `charge_density_0form` | A (vector) | initial Poisson solve of these models | PR 14 |
-| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | next (6D views available) |
+| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ✓ (#688, see [notes](#vlasov_maxwell-implementation-notes)) |
 
 **Step 2 – LinearVlasovAmpère/Maxwell (δf)**
 
@@ -552,8 +552,8 @@ and dicts (`utils/kernel_compilation.collect_kernels`). A propagator that calls
 overrides `kernels()`. Diagnostics kernels and the initial solves are not listed; they run during the setup.
 
 Fail fast: on CuPy, if a listed kernel has no CUDA version, `compile_cuda_kernels()` raises
-`NotImplementedError` naming all of them, before anything is compiled (e.g. `VlasovAmpereOneSpecies`:
-`vlasov_maxwell`). On NumPy it does nothing. A test runs a model for one step on NumPy and checks that every
+`NotImplementedError` naming all of them, before anything is compiled (e.g. a model
+that pushes with `push_bxu_Hdiv`). On NumPy it does nothing. A test runs a model for one step on NumPy and checks that every
 `Kernel` called in the time loop is listed.
 
 ## PR 18 implementation notes
@@ -663,6 +663,7 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a
   during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake
   CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI
   workflow); the GPU variant has not run (no GPU here).
+
 ## `linear_vlasov_ampere` implementation notes
 
 The first matrix accumulation on the GPU (part of #688, tracked in #650): `linear_vlasov_ampere`, the accumulation of
@@ -705,3 +706,25 @@ The first matrix accumulation on the GPU (part of #688, tracked in #650): `linea
   with pyccel to round-off and catches a transposed `DF`. `m_v_fill_b_v1_symm` is compared with pyccel through a
   wrapper kernel (`test_v1_symm_filler` on a GPU, `test_emulated_v1_symm_filler` by emulation). GPU tests have not run
   on an H100.
+
+## `vlasov_maxwell` implementation notes
+
+The accumulation of `VlasovAmpereCoupling` in `VlasovAmpereOneSpecies` and `VlasovMaxwellOneSpecies` (step 1, #3 of
+the [porting order](#porting-order); part of #688).
+
+- **Kernel.** `vlasov_maxwell_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row; `df`,
+  `matrix_inv`, `transpose`, `matrix_matrix`, `matrix_vector`, then `m_v_fill_b_v1_symm` with
+  `A_p = w_p DF^{-1} DF^{-T}` and `B_p = w_p DF^{-1} v_p`. As in pyccel, only holes are skipped: boundary particles
+  (`markers[ip, -1] == -2`) are accumulated (unlike `linear_vlasov_ampere`).
+- **Device helpers.** `fill_mat`, `fill_mat_vec` (`filler_kernels.cuh`) and `m_v_fill_b_v1_symm`
+  (`particle_to_mat_kernels.cuh`) with atomic adds into `Array6D<double>`/`Array3D<double>` data, byte-identical to the
+  ones of the `linear_vlasov_ampere` PR (#705), as are `v1_symm_accumulation_data()` and the `m_v_fill_b_v1_symm`
+  wrapper tests.
+- **Tests.** `PARITY_CASES["vlasov_maxwell"]`: 129 markers (a hole, a boundary particle) in Cuboid, Colella,
+  HollowTorus, ShafranovDshapedCylinder and a 3d spline mapping; the CPU emulation agrees with pyccel to 7e-12 at
+  entries up to 6e5 (spline mapping) and fails for a transposed `DF`, skipped boundary particles or unskipped holes.
+  `rtol = 1e-12`, `atol = 1e-8` for the GPU's atomic summation order. GPU tests have not run on an H100.
+- **Still missing for the models end to end on the GPU.** `MassMatrixPreconditioner` (the default of
+  `VlasovAmpereCoupling`, `MaxwellWeakAmpere` and the initial `PoissonSolve`) cannot be created on CuPy, CG inner products
+  return host scalars, the FEEC operators (`curl`, `grad`, mass matrices, Schur solves) need the feectools CUDA stack,
+  and the first H100 run (#687) is pending.

src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu

@@ -0,0 +1,78 @@
+#include "struphy/geometry/evaluation_kernels.cuh"
+#include "struphy/linear_algebra/linalg_kernels.cuh"
+#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh"
+
+/**
+ * Accumulate the matrix and vector of the full-orbit Vlasov-Maxwell (Vlasov-Ampere) coupling into V1, as in
+ * vlasov_maxwell_kernels.vlasov_maxwell.
+ *
+ * Filling functions:
+ *     A_p^{mu, nu} = w_p G^{-1}_{mu, nu}(eta_p),   G^{-1} = DF^{-1} DF^{-T},
+ *     B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu;
+ * one thread per marker row.
+ *
+ * @param args_markers Marker buffer (n_markers x n_cols, row-major); the weight is in column 6, as in pyccel.
+ * @param args_derham Spline degrees (1 to 8), knots and start indices.
+ * @param args_domain Mapping arguments; every spline and analytic mapping.
+ * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any
+ *              strides); written with atomic additions.
+ * @param mat12 (1,2)-block data; written atomically.
+ * @param mat13 (1,3)-block data; written atomically.
+ * @param mat22 (2,2)-block data; written atomically.
+ * @param mat23 (2,3)-block data; written atomically.
+ * @param mat33 (3,3)-block data; written atomically.
+ * @param vec1 First component of the V1 vector data; written atomically.
+ * @param vec2 Second component of the V1 vector data; written atomically.
+ * @param vec3 Third component of the V1 vector data; written atomically.
+ *
+ * Holes (markers[ip, 0] == -1) are skipped; boundary particles are not, as in pyccel. The pyccel loop runs over all
+ * rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation order
+ * (and the last bits of the result) differs from the serial pyccel loop.
+ */
+extern "C" __global__ void vlasov_maxwell(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain,
+                                          Array6D<double> mat11, Array6D<double> mat12, Array6D<double> mat13,
+                                          Array6D<double> mat22, Array6D<double> mat23, Array6D<double> mat33,
+                                          Array3D<double> vec1, Array3D<double> vec2, Array3D<double> vec3) {
+    // CUDA-only: the marker row of this thread replaces the pyccel loop variable
+    int ip = blockDim.x * blockIdx.x + threadIdx.x;
+    if (ip >= args_markers.n_markers) return;
+
+    // allocate for metric coeffs (3 x 3 matrices as nine row-major entries)
+    double dfm[9], df_inv[9], df_inv_t[9], g_inv[9];
+
+    // allocate for filling
+    double v[3], df_inv_times_v[3], filling_m[9], filling_v[3];
+
+    // only do something if particle is a "true" particle (i.e. not a hole)
+    if (args_markers.markers(ip, 0) == -1.) return;
+
+    // marker positions
+    double eta1 = args_markers.markers(ip, 0);
+    double eta2 = args_markers.markers(ip, 1);
+    double eta3 = args_markers.markers(ip, 2);
+
+    // evaluate Jacobian, result in dfm
+    struphy_cuda::df(eta1, eta2, eta3, args_domain, dfm);
+
+    // compute shifted and stretched velocity
+    v[0] = args_markers.markers(ip, 3);
+    v[1] = args_markers.markers(ip, 4);
+    v[2] = args_markers.markers(ip, 5);
+
+    // filling functions
+    struphy_cuda::matrix_inv(dfm, df_inv);
+    struphy_cuda::transpose(df_inv, df_inv_t);
+    struphy_cuda::matrix_matrix(df_inv, df_inv_t, g_inv);
+    struphy_cuda::matrix_vector(df_inv, v, df_inv_times_v);
+
+    // filling_m = w_p * DF^{-1} * DF^{-T}
+    for (int k = 0; k < 9; ++k) filling_m[k] = args_markers.markers(ip, 6) * g_inv[k];
+
+    // filling_v = w_p * DF^{-1} * \V
+    for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_times_v[k];
+
+    // call the appropriate matvec filler
+    struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, mat33,
+                                     filling_m[0], filling_m[1], filling_m[2], filling_m[4], filling_m[5],
+                                     filling_m[8], vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2]);
+}

src/struphy/pic/tests/cuda_parity_cases.py

@@ -93,6 +93,14 @@ def linear_vlasov_ampere_args(domain):
     )
 
 
+def vlasov_maxwell_args(domain):
+    """129 markers (a hole in row 0, a boundary particle in row 1, which is accumulated as in pyccel) in domain
+    `domain` of geometry_domain(); zeroed V1 data."""
+    args_markers, _ = marker_arguments((0, 0, 0))
+    args_markers.markers[0, 0] = -1.0
+    return (args_markers, derham_arguments(), geometry_domain(domain).args_domain, *v1_symm_accumulation_data())
+
+
 # ---------------------------------------------------------------- spline evaluation
 
 
@@ -284,6 +292,10 @@ def size_of(index, per_thread=1):
     # Cuboid, Colella and HollowTorus (non-diagonal DF) and a 3d spline mapping, by index into geometry_domain();
     # atomics add in another order than the serial loop, and the entries reach 1e4 (atol for cancelling sums)
     "linear_vlasov_ampere": ParityCases((0, 2, 5, 12), linear_vlasov_ampere_args, rtol=1e-12, atol=1e-10),
+    # Cuboid, then Colella, HollowTorus, ShafranovDshapedCylinder (non-diagonal DF) and a 3d spline mapping, by index into
+    # geometry_domain(); atomics add in another order than the serial loop, and G^{-1} of the spline mapping makes
+    # entries of 6e5 (atol for cancelling sums)
+    "vlasov_maxwell": ParityCases((0, 2, 5, 9, 12), vlasov_maxwell_args, rtol=1e-12, atol=1e-8),
     "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args),
     "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)),
     "eval_spline_mpi_sparse_meshgrid": ParityCases(

src/struphy/pic/tests/test_cuda_parity.py

@@ -11,6 +11,7 @@
 
 from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form
 from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere
+from struphy.pic.accumulation.kernels.vlasov_maxwell import vlasov_maxwell
 from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage
 from struphy.pic.pushing.kernels.push_v_with_efield import push_v_with_efield
 from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic
@@ -118,5 +119,6 @@ def test_vlasov_kernel_coverage():
         push_weights_with_efield_lin_va,
         charge_density_0form,
         linear_vlasov_ampere,
+        vlasov_maxwell,
     ):
         assert kernel.has_cuda, kernel.name

src/struphy/simulation/tests/test_compile_cuda_kernels.py

@@ -119,8 +119,8 @@ def test_kernels_of_a_small_model(tmp_path):
     }
     assert names(sim.model.propagators.push_eta.kernels()) == {"push_eta_stage"}
     assert names(sim.model.propagators.coupling_va.kernels()) == {"vlasov_maxwell", "push_v_with_efield"}
-    # the accumulation kernel vlasov_maxwell has no CUDA version yet: this model cannot run on CuPy
-    assert names(missing_cuda(kernels)) == {"vlasov_maxwell"}
+    # every kernel of this model has a CUDA version (vlasov_maxwell since #713)
+    assert missing_cuda(kernels) == []
 
 
 def test_kernels_of_reflecting_particles(tmp_path):
@@ -236,7 +236,8 @@ def test_compile_cuda_kernels_fake_cupy():
 
 @requires_cupy
 def test_compile_cuda_kernels_on_gpu(tmp_path):
-    """On a GPU: the kernels of Vlasov are compiled at setup, VlasovAmpereOneSpecies fails fast."""
+    """On a GPU: the kernels of Vlasov and VlasovAmpereOneSpecies are compiled at setup, a kernel without CUDA
+    version fails fast."""
     with cunumpy.use_backend("cupy"):
         sim = make_vlasov(tmp_path)
         sim.allocate()
@@ -246,5 +247,10 @@ def test_compile_cuda_kernels_on_gpu(tmp_path):
 
         sim = make_vlasov_ampere(tmp_path / "va")
         sim.allocate()
-        with pytest.raises(NotImplementedError, match="vlasov_maxwell"):
+        compiled = sim.compile_cuda_kernels()
+        assert compiled == sim.kernels()
+        assert "vlasov_maxwell" in names(compiled)
+
+        sim.kernels = lambda: (push_eta_stage, push_bxu_Hdiv)
+        with pytest.raises(NotImplementedError, match="push_bxu_Hdiv"):
             sim.compile_cuda_kernels()