Field evaluation and diagnostics on the device

CUDA_STRATEGY.md

@@ -27,6 +27,7 @@ Progress is tracked in struphy-hub/struphy#650.
 - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)).
 - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines were host-only at this stage; their current CuPy support is described in [Polar splines on CuPy](#polar-splines-on-cupy-695). See also the [PR 19 implementation notes](#pr-19-implementation-notes).
 - [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)).
+- [ ] **Diagnostics on the device** (struphy-hub/struphy#697) — scalar quantities, marker binning, saved markers and spline evaluation for output run on the device; the scalars reach the host in one copy per update (`Scalars.to_host`), every other saved dataset in one counted copy per output step (`DataContainer.save_data`), and the time state is host bookkeeping. CUDA version of `eval_spline_mpi_tensor_product_fixed` (`SplineFunction.eval_tp_fixed_loc`), moved into its own folder. Marker diagnostics kernels (`pic/diagnostics/kernels`) and SPH kernel density plots are follow-ups (see [Diagnostics on the device](#diagnostics-on-the-device-697)).
 - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)).
 - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)).
 - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes.
@@ -829,6 +830,41 @@ the [porting order](#porting-order); part of #688).
   ghost-region flags from #746 are preserved. feectools' own halo exchanges still require CUDA-aware MPI
   on real GPUs; staging of the polar collectives does not replace that separate dependency.
 - **Tests.** `polar/tests/test_polar_cupy.py` compares NumPy and CuPy for `IGAPolarCylinder` and `IGAPolarTorus`:
+  extraction and DOF extraction operators and their transposes, `grad`/`curl`/`div` and transposes, `PolarVector`
+  arithmetic, the polar projectors, polar mass matrices (GPU only) and polar `SplineFunction` coefficients
+  (`extract_coeffs`, restart). Without a GPU it runs on the fake CuPy with two substitutes: a dense stand-in for
+  `cupyx.scipy.sparse` (`set_device_sparse_module`) and the host versions of the feectools kernels (the fake cannot
+  launch CUDA kernels). The mass matrices need struphy's CUDA geometry kernels and are compared on the GPU only.
+  Field evaluation (`SplineFunction.__call__`) is not covered: it is not ported to CuPy for any space yet.
+
+## Diagnostics on the device (#697)
+
+What a CuPy run copies to the host per output step, and where:
+
+- **Scalars** (`models/scalars.py`). The `value` of every scalar in a `Scalars` container is a view into one device
+  array of the container; `Scalars.update()` ends with `to_host()`, one `xp.to_numpy` of that array (8 bytes per
+  scalar). `FunctionScalar*` store the result of their function without `float()` (a device sync per scalar before),
+  so a device reduction stays on the device. The output (`Simulation._initialize_hdf5_datasets`) and
+  `print_scalar_quantities` read `Scalars.host_value(name)`. `BilinearEnergyFEEC` and `VolumeFormEnergyFEEC` call
+  feectools' inner products, which copy their result to the host until struphy-hub/feectools#98 (#716) is in.
+- **Output** (`io/output_handling.py`). `DataContainer._as_numpy_array` is `xp.to_numpy`: the one place where device
+  data is copied for saving, one counted copy per saved device dataset (FEEC coefficients, saved markers, binned
+  distribution functions). Before, it used `.get()`, which `count_transfers` does not see.
+- **Time state** (`simulation/sim.py`). `time_state` is NumPy on both backends: the time loop reads it with `float()`
+  every step, which was a device sync per read on CuPy.
+- **Binning** (`Particles.binning`). Failed on CuPy (`xp.count_nonzero` of a list); now gathers the valid markers once
+  and runs `xp.histogramdd` on the device. Saved markers (`update_markers_to_be_saved`) gather once, without counting
+  the mask first.
+- **Spline evaluation.** `SplineFunction.__call__` compared the process bounds as device scalars (six syncs per
+  meshgrid evaluation); the bounds are cached on the host once, and the flagging writes are masked copies.
+  `eval_tp_fixed_loc` (thermal energies of the variational models) called the pyccel kernel with device arrays; its
+  kernel `eval_spline_mpi_tensor_product_fixed` is now in `bsplines/kernels/` with a CUDA version and parity cases.
+  `Derham.prepare_eval_tp_fixed` (setup) computes on the host and copies the spans and basis values to the device once.
+- **Not counted, still there**: boolean-mask gathers (`markers[valid_mks]`, also in `KineticEnergyPIC`) wait for the
+  device to know the result size; one per gather.
+- **Follow-ups**: the marker diagnostics kernels in `pic/diagnostics/kernels/` (5D energies, magnetic moments,
+  canonical momenta, `thermal_energy`; used by the guiding-center and hybrid models) have no CUDA version; SPH kernel
+  density plots need the SPH kernels (step 6).
   extraction/DOF operators, derivatives and transposes, vector arithmetic, all three projector solvers,
   transposed and boundary-condition solves, mass matrices/preconditioners, callable and sampled weights,
   and spline coefficient extraction/restart. Repeated Kronecker and preconditioner applications are checked

src/struphy/bsplines/evaluation_kernels_3d.py

@@ -1334,82 +1334,6 @@ def eval_spline_mpi_tensor_product_fast(
                 )
 
 
-@stack_array("b1", "b2", "b3")
-def eval_spline_mpi_tensor_product_fixed(
-    span1s: "int[:]",
-    span2s: "int[:]",
-    span3s: "int[:]",
-    b1s: "float[:,:]",
-    b2s: "float[:,:]",
-    b3s: "float[:,:]",
-    _data: "float[:,:,:]",
-    kind: "int[:]",
-    pn: "int[:]",
-    starts: "int[:]",
-    values: "float[:,:,:]",
-):
-    """
-    Tensor-product evaluation of a tensor-product spline, distributed,
-    and optimized for a fixed grid (spans and spline values have been pre-evaluated).
-
-    Parameters
-    ----------
-        span1s, span2s, span3s : array[int]
-            Knot span indices.
-
-        b1s, b2s, b3s : array[float]
-            Values of p+1 non-zero basis functions at evaluation points.
-
-        _data : array[float]
-            The spline coefficients c_ijk.
-
-        kind : array[int]
-            Kind of 1d basis in each direction: 0 = N-spline, 1 = D-spline.
-
-        pn : array[int]
-            Spline degrees of V0 in each direction.
-
-        starts : array[float]
-            Start indices of splines on current process.
-
-        values : array[float]
-            Return array for spline values S_ijk where ijk are the flattened indices corresponding to spans.
-    """
-
-    # allocate spline values
-    b1 = empty(pn[0] - kind[0] + 1, dtype=float)
-    b2 = empty(pn[1] - kind[1] + 1, dtype=float)
-    b3 = empty(pn[2] - kind[2] + 1, dtype=float)
-
-    ni = span1s.size
-    nj = span2s.size
-    nk = span3s.size
-
-    for i in range(ni):
-        span1 = span1s[i]
-        b1[:] = b1s[i, :]
-        for j in range(nj):
-            span2 = span2s[j]
-            b2[:] = b2s[j, :]
-            for k in range(nk):
-                span3 = span3s[k]
-                b3[:] = b3s[k, :]
-
-                values[i, j, k] = eval_spline_mpi_kernel(
-                    pn[0] - kind[0],
-                    pn[1] - kind[1],
-                    pn[2] - kind[2],
-                    b1,
-                    b2,
-                    b3,
-                    span1,
-                    span2,
-                    span3,
-                    _data,
-                    starts,
-                )
-
-
 ###########################################
 ### Distributed Derham field evaluation ###
 ###########################################

src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/__init__.py

@@ -0,0 +1,7 @@
+"""The kernel of this folder: pyccel in ``<name>_kernels.py``, CUDA in ``<name>_cuda.cu`` once ported."""
+
+from cunumpy.kernels import Kernel
+
+from struphy.utils.cuda_arguments import CUDA_STRUCTS
+
+eval_spline_mpi_tensor_product_fixed = Kernel.from_folder(__name__, structs=CUDA_STRUCTS)

src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_cuda.cu

@@ -0,0 +1,53 @@
+#include "struphy/bsplines/evaluation_kernels_3d.cuh"
+
+// Same arguments, in the same order, as the pyccel kernel in eval_spline_mpi_tensor_product_fixed_kernels.py. Array
+// views carry dimensions and strides, so no CUDA-only lengths are needed.
+
+/**
+ * Tensor-product evaluation of a distributed spline on a fixed grid, as in
+ * eval_spline_mpi_tensor_product_fixed_kernels.eval_spline_mpi_tensor_product_fixed.
+ *
+ * One thread per entry (i, j, k) of values. The knot spans and the non-zero basis values at the grid points are
+ * pre-evaluated (SplineFunction.eval_tp_fixed_loc, Derham.prepare_eval_tp_fixed), so the thread only sums the
+ * contributions of the local coefficients.
+ *
+ * @param span1s Knot span indices of the grid points along the first axis (ni entries).
+ * @param span2s Knot span indices along the second axis (nj entries).
+ * @param span3s Knot span indices along the third axis (nk entries).
+ * @param b1s Values of the non-zero basis functions along the first axis, shape (ni, pn[0] - kind[0] + 1).
+ * @param b2s Values of the non-zero basis functions along the second axis, shape (nj, pn[1] - kind[1] + 1).
+ * @param b3s Values of the non-zero basis functions along the third axis, shape (nk, pn[2] - kind[2] + 1).
+ * @param _data Spline coefficients of the current process, any strides.
+ * @param kind Kind of 1d basis in each direction (three entries): 0 = N-spline, 1 = D-spline.
+ * @param pn Spline degrees of V0 in each direction (three entries, 1 to 8).
+ * @param starts Start indices of the splines on the current process (three entries).
+ * @param values Output spline values, shape (ni, nj, nk).
+ *
+ * Pyccel copies the rows b1s[i, :], b2s[j, :], b3s[k, :] into b1, b2, b3 before calling eval_spline_mpi_kernel; here
+ * b1, b2, b3 are thread-local arrays of MAX_SPLINE_DEGREE + 1 entries, since the rows of a view may be strided.
+ */
+extern "C" __global__ void eval_spline_mpi_tensor_product_fixed(Array1D<long long> span1s, Array1D<long long> span2s,
+                                                                Array1D<long long> span3s, Array2D<double> b1s,
+                                                                Array2D<double> b2s, Array2D<double> b3s,
+                                                                Array3D<double> _data, const long long* kind,
+                                                                const long long* pn, const long long* starts,
+                                                                Array3D<double> values) {
+    // CUDA-only: flat index of this thread, split into the pyccel loop variables i, j, k
+    long long ijk = (long long)blockDim.x * blockIdx.x + threadIdx.x;
+    long long ni = span1s.shape[0];
+    long long nj = span2s.shape[0];
+    long long nk = span3s.shape[0];
+    if (ijk >= ni * nj * nk) return;
+    long long k = ijk % nk;
+    long long j = (ijk / nk) % nj;
+    long long i = ijk / (nj * nk);
+
+    double b1[struphy_cuda::MAX_SPLINE_DEGREE + 1], b2[struphy_cuda::MAX_SPLINE_DEGREE + 1],
+        b3[struphy_cuda::MAX_SPLINE_DEGREE + 1];
+    for (int il = 0; il <= pn[0] - kind[0]; ++il) b1[il] = b1s(i, il);
+    for (int il = 0; il <= pn[1] - kind[1]; ++il) b2[il] = b2s(j, il);
+    for (int il = 0; il <= pn[2] - kind[2]; ++il) b3[il] = b3s(k, il);
+
+    values(i, j, k) = struphy_cuda::eval_spline_mpi_kernel(pn[0] - kind[0], pn[1] - kind[1], pn[2] - kind[2], b1, b2,
+                                                           b3, span1s(i), span2s(j), span3s(k), _data, starts);
+}

src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_kernels.py

@@ -0,0 +1,82 @@
+"""Tensor-product evaluation of a distributed spline on a fixed grid (spans and basis values pre-evaluated)."""
+
+from numpy import empty
+from pyccel.decorators import stack_array
+
+from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi_kernel
+
+
+@stack_array("b1", "b2", "b3")
+def eval_spline_mpi_tensor_product_fixed(
+    span1s: "int[:]",
+    span2s: "int[:]",
+    span3s: "int[:]",
+    b1s: "float[:,:]",
+    b2s: "float[:,:]",
+    b3s: "float[:,:]",
+    _data: "float[:,:,:]",
+    kind: "int[:]",
+    pn: "int[:]",
+    starts: "int[:]",
+    values: "float[:,:,:]",
+):
+    """
+    Tensor-product evaluation of a tensor-product spline, distributed,
+    and optimized for a fixed grid (spans and spline values have been pre-evaluated).
+
+    Parameters
+    ----------
+        span1s, span2s, span3s : array[int]
+            Knot span indices.
+
+        b1s, b2s, b3s : array[float]
+            Values of p+1 non-zero basis functions at evaluation points.
+
+        _data : array[float]
+            The spline coefficients c_ijk.
+
+        kind : array[int]
+            Kind of 1d basis in each direction: 0 = N-spline, 1 = D-spline.
+
+        pn : array[int]
+            Spline degrees of V0 in each direction.
+
+        starts : array[float]
+            Start indices of splines on current process.
+
+        values : array[float]
+            Return array for spline values S_ijk where ijk are the flattened indices corresponding to spans.
+    """
+
+    # allocate spline values
+    b1 = empty(pn[0] - kind[0] + 1, dtype=float)
+    b2 = empty(pn[1] - kind[1] + 1, dtype=float)
+    b3 = empty(pn[2] - kind[2] + 1, dtype=float)
+
+    ni = span1s.size
+    nj = span2s.size
+    nk = span3s.size
+
+    for i in range(ni):
+        span1 = span1s[i]
+        b1[:] = b1s[i, :]
+        for j in range(nj):
+            span2 = span2s[j]
+            b2[:] = b2s[j, :]
+            for k in range(nk):
+                span3 = span3s[k]
+                b3[:] = b3s[k, :]
+
+                values[i, j, k] = eval_spline_mpi_kernel(
+                    pn[0] - kind[0],
+                    pn[1] - kind[1],
+                    pn[2] - kind[2],
+                    b1,
+                    b2,
+                    b3,
+                    span1,
+                    span2,
+                    span3,
+                    _data,
+                    starts,
+                )

src/struphy/bsplines/tests/test_eval_spline_mpi.py

@@ -699,7 +699,8 @@ def test_eval_tensor_product_grid(num_elements, degree, bcs, n_markers=10):
 
     import time
 
-    from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi_tensor_product_fixed, evaluate_tensor_product
+    from struphy.bsplines.evaluation_kernels_3d import evaluate_tensor_product
+    from struphy.bsplines.kernels.eval_spline_mpi_tensor_product_fixed import eval_spline_mpi_tensor_product_fixed
     from struphy.feec.basis_projection_ops import prepare_projection_of_basis
     from struphy.feec.psydac_derham import Derham
     from struphy.feec.utilities import create_equal_random_arrays as cera