Argument class for SplineFunction kernel arguments

CUDA_STRATEGY.md

@@ -61,9 +61,9 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke
 | File | Content |
 |---|---|
 | `cunumpy.kernels`, `cunumpy.arguments`, `cunumpy.cuda` (cunumpy 0.6) | `Kernel`, `KernelCatalog`, `PyccelKernel` and `CudaKernel` in `cunumpy.kernels`; `CudaStructArguments`, `CudaStruct` and `write_cuda_header` in `cunumpy.arguments`; the device runtime (`bind_local_device`, ...) in `cunumpy.cuda`. Struphy uses cunumpy for dispatch and struct packing |
-| `src/struphy/kernel_arguments/` | the argument classes in pairs: the pyccel classes in `pusher_args_kernels.py` / `local_projectors_args_kernels.py` (NumPy backend) and their CUDA versions `Cuda<Name>` in `pusher_args_cuda.py` / `local_projectors_args_cuda.py` (CuPy backend; subclasses of `CudaStructArguments` whose `fields` define the C struct) |
-| `src/struphy/utils/cuda_arguments.py` | `CUDA_STRUCTS`, `CUDA_OPTIONS`, `write_pusher_header()` and `write_local_projectors_header()` |
-| `src/struphy/kernel_arguments/pusher_args.cuh`, `local_projectors_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs`, `DomainArgs` and `LocalProjectorsArgs`, generated from the CUDA classes, using cunumpy array views |
+| `src/struphy/kernel_arguments/` | the argument classes in pairs: the pyccel classes in `pusher_args_kernels.py` / `local_projectors_args_kernels.py` / `spline_args_kernels.py` (NumPy backend) and their CUDA versions `Cuda<Name>` in `pusher_args_cuda.py` / `local_projectors_args_cuda.py` / `spline_args_cuda.py` (CuPy backend; subclasses of `CudaStructArguments` whose `fields` define the C struct) |
+| `src/struphy/utils/cuda_arguments.py` | `CUDA_STRUCTS`, `CUDA_OPTIONS`, `write_pusher_header()`, `write_local_projectors_header()` and `write_spline_header()` |
+| `src/struphy/kernel_arguments/pusher_args.cuh`, `local_projectors_args.cuh`, `spline_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs`, `DomainArgs`, `LocalProjectorsArgs` and `SplineArgs`, generated from the CUDA classes, using cunumpy array views |
 | `src/struphy/geometry/base.py`, `src/struphy/pic/base.py`, `src/struphy/feec/psydac_derham.py` | `Domain.args_domain`, `Particles.args_markers` and `Derham.args_derham` are the pyccel class on NumPy and the CUDA class on CuPy, chosen once at construction; every kernel call goes through a `Kernel` |
 | `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Sixteen have CUDA versions (`linear_vlasov_ampere` the latest), with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) |
 | `src/struphy/geometry/evaluation_kernels.cuh`, `transform_kernels.cuh`, `spline_mappings_kernels.cuh`, `geometry/domains/<name>/<name>_cuda.cuh` | device versions of the mapping helpers for every mapping: `f`/`df` per analytic domain (`kind_map` 10–12, 20–22, 30–32, PR 18) and the spline mappings `spline_3d`, `spline_2d_straight`, `spline_2d_torus` (`kind_map` 0–2, PR 19), the `kind_map` switch, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`, `push`, `tran` |
@@ -554,6 +554,7 @@ The argument classes come in pairs, one for each backend:
 | `pusher_args_kernels.DerhamArguments` | `pusher_args_cuda.CudaDerhamArguments` (`DerhamArgs`) |
 | `pusher_args_kernels.DomainArguments` | `pusher_args_cuda.CudaDomainArguments` (`DomainArgs`) |
 | `local_projectors_args_kernels.LocalProjectorsArguments` | `local_projectors_args_cuda.CudaLocalProjectorsArguments` (`LocalProjectorsArgs`) |
+| `spline_args_kernels.SplineArguments` | `spline_args_cuda.CudaSplineArguments` (`SplineArgs`) |
 
 The two classes of a pair take the same constructor arguments and have the same attributes; the CUDA struct
 adds only derived members that CUDA pointers cannot carry (`n_markers` is also a pyccel attribute; `nt1`,
@@ -580,6 +581,10 @@ step: CUDA versions of `kernel_evaluate_pic`, `kernel_evaluate` and the pull/pus
 `CudaLocalProjectorsArguments` is not used by any CUDA kernel yet; local projectors are rejected on the
 CuPy backend by `Derham`.
 
+`SplineFunction` creates one `SplineArguments` (`CudaSplineArguments` on CuPy) per component, holding `kind`,
+`pn`, `tn1`, `tn2`, `tn3` and `starts`; the three `eval_spline_mpi_*` entry kernels take it as `args_spline`
+instead of six loose arguments (#678), and pass it on to the device helper `eval_spline_mpi`.
+
 ### Kernel folders
 
 Every kernel that Python code calls with argument objects now lives in its own folder, like the pusher and

src/struphy/bsplines/evaluation_kernels_3d.cuh

@@ -1,5 +1,6 @@
 #pragma once
 #include "struphy/kernel_arguments/pusher_args.cuh"
+#include "struphy/kernel_arguments/spline_args.cuh"
 #include "struphy/bsplines/bsplines_kernels.cuh"
 namespace struphy_cuda {
 /** Per-thread outputs of get_spans; spline names match DerhamArguments. */
@@ -148,35 +149,36 @@ __device__ inline void eval_2form_spline_mpi(int span1, int span2, int span3, co
  * @param eta2 Evaluation point along the second axis.
  * @param eta3 Evaluation point along the third axis.
  * @param _data Spline coefficients of the current process (the _data of a StencilVector), 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 tn1 Knot vector of V0 along the first axis (contiguous).
- * @param tn2 Knot vector of V0 along the second axis.
- * @param tn3 Knot vector of V0 along the third axis.
- * @param starts Start indices of the splines on the current process (three entries).
+ * @param args_spline Kind of 1d basis (0 = N-spline, 1 = D-spline), spline degrees (1 to 8) and knot vectors
+ *        (contiguous) of V0 in each direction, and start indices of the splines on the current process.
  * @return value, the value of the spline at (eta1, eta2, eta3).
  *
- * Pyccel allocates bn1, ..., bd3 per call; here they are fields of the thread-local SplineScratch.
+ * Pyccel writes the spline values into bn1, ..., bd3 of args_spline; here they are fields of the thread-local
+ * SplineScratch.
  */
 __device__ inline double eval_spline_mpi(double eta1, double eta2, double eta3, Array3D<double> _data,
-                                         const long long* kind, const long long* pn, Array1D<double> tn1,
-                                         Array1D<double> tn2, Array1D<double> tn3, const long long* starts) {
+                                         const SplineArgs& args_spline) {
     struphy_cuda::SplineScratch scratch;
 
     // get spline values at eta
-    scratch.span1 = struphy_cuda::find_span(tn1.data, tn1.shape[0], pn[0], eta1);
-    scratch.span2 = struphy_cuda::find_span(tn2.data, tn2.shape[0], pn[1], eta2);
-    scratch.span3 = struphy_cuda::find_span(tn3.data, tn3.shape[0], pn[2], eta3);
-    struphy_cuda::b_d_splines_slim(tn1.data, pn[0], eta1, scratch.span1, scratch.bn1, scratch.bd1);
-    struphy_cuda::b_d_splines_slim(tn2.data, pn[1], eta2, scratch.span2, scratch.bn2, scratch.bd2);
-    struphy_cuda::b_d_splines_slim(tn3.data, pn[2], eta3, scratch.span3, scratch.bn3, scratch.bd3);
+    scratch.span1 = struphy_cuda::find_span(args_spline.tn1.data, args_spline.tn1.shape[0], args_spline.pn[0], eta1);
+    scratch.span2 = struphy_cuda::find_span(args_spline.tn2.data, args_spline.tn2.shape[0], args_spline.pn[1], eta2);
+    scratch.span3 = struphy_cuda::find_span(args_spline.tn3.data, args_spline.tn3.shape[0], args_spline.pn[2], eta3);
+    struphy_cuda::b_d_splines_slim(args_spline.tn1.data, args_spline.pn[0], eta1, scratch.span1, scratch.bn1,
+                                   scratch.bd1);
+    struphy_cuda::b_d_splines_slim(args_spline.tn2.data, args_spline.pn[1], eta2, scratch.span2, scratch.bn2,
+                                   scratch.bd2);
+    struphy_cuda::b_d_splines_slim(args_spline.tn3.data, args_spline.pn[2], eta3, scratch.span3, scratch.bn3,
+                                   scratch.bd3);
 
-    const double* b1 = kind[0] == 0 ? scratch.bn1 : scratch.bd1;
-    const double* b2 = kind[1] == 0 ? scratch.bn2 : scratch.bd2;
-    const double* b3 = kind[2] == 0 ? scratch.bn3 : scratch.bd3;
+    const double* b1 = args_spline.kind[0] == 0 ? scratch.bn1 : scratch.bd1;
+    const double* b2 = args_spline.kind[1] == 0 ? scratch.bn2 : scratch.bd2;
+    const double* b3 = args_spline.kind[2] == 0 ? scratch.bn3 : scratch.bd3;
 
-    double value = struphy_cuda::eval_spline_mpi_kernel(pn[0] - kind[0], pn[1] - kind[1], pn[2] - kind[2], b1, b2, b3,
-                                                        scratch.span1, scratch.span2, scratch.span3, _data, starts);
+    double value = struphy_cuda::eval_spline_mpi_kernel(
+        args_spline.pn[0] - args_spline.kind[0], args_spline.pn[1] - args_spline.kind[1],
+        args_spline.pn[2] - args_spline.kind[2], b1, b2, b3, scratch.span1, scratch.span2, scratch.span3, _data,
+        args_spline.starts);
     return value;
 }
 

src/struphy/bsplines/evaluation_kernels_3d.py

@@ -24,7 +24,9 @@
 
 import struphy.bsplines.bsplines_kernels as bsplines_kernels
 import struphy.kernel_arguments.pusher_args_kernels as pusher_args_kernels
+import struphy.kernel_arguments.spline_args_kernels as spline_args_kernels  # do not remove; needed to identify dependencies
 from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments
+from struphy.kernel_arguments.spline_args_kernels import SplineArguments
 
 #################################
 ### Single process evaluation ###
@@ -1140,18 +1142,12 @@ def eval_spline_derivative_mpi_kernel(
     return spline_value
 
 
-@stack_array("bn1", "bn2", "bn3", "bd1", "bd2", "bd3")
 def eval_spline_mpi(
     eta1: float,
     eta2: float,
     eta3: float,
     _data: "float[:,:,:]",
-    kind: "int[:]",
-    pn: "int[:]",
-    tn1: "float[:]",
-    tn2: "float[:]",
-    tn3: "float[:]",
-    starts: "int[:]",
+    args_spline: "SplineArguments",
 ) -> float:
     """
     Point-wise evaluation of a tensor-product spline, distributed.
@@ -1164,69 +1160,52 @@ def eval_spline_mpi(
         _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.
-
-        tn1, tn2, tn3 : array[float]
-            Knot vectors of V0 in each direction.
-
-        starts : array[float]
-            Start indices of splines on current process.
+        args_spline : SplineArguments
+            Kind of 1d basis, spline degrees and knot vectors of V0, and start indices of the splines on the
+            current process.
 
     Returns
     -------
         value : float
             Value of tensor-product spline at point (eta_1, eta_2, eta_3).
     """
 
-    # allocate spline values
-    bn1 = empty(pn[0] + 1, dtype=float)
-    bn2 = empty(pn[1] + 1, dtype=float)
-    bn3 = empty(pn[2] + 1, dtype=float)
-
-    bd1 = empty(pn[0], dtype=float)
-    bd2 = empty(pn[1], dtype=float)
-    bd3 = empty(pn[2], dtype=float)
-
     # get spline values at eta
-    span1 = bsplines_kernels.find_span(tn1, pn[0], eta1)
-    span2 = bsplines_kernels.find_span(tn2, pn[1], eta2)
-    span3 = bsplines_kernels.find_span(tn3, pn[2], eta3)
+    span1 = bsplines_kernels.find_span(args_spline.tn1, args_spline.pn[0], eta1)
+    span2 = bsplines_kernels.find_span(args_spline.tn2, args_spline.pn[1], eta2)
+    span3 = bsplines_kernels.find_span(args_spline.tn3, args_spline.pn[2], eta3)
 
-    bsplines_kernels.b_d_splines_slim(tn1, pn[0], eta1, span1, bn1, bd1)
-    bsplines_kernels.b_d_splines_slim(tn2, pn[1], eta2, span2, bn2, bd2)
-    bsplines_kernels.b_d_splines_slim(tn3, pn[2], eta3, span3, bn3, bd3)
+    bsplines_kernels.b_d_splines_slim(args_spline.tn1, args_spline.pn[0], eta1, span1, args_spline.bn1, args_spline.bd1)
+    bsplines_kernels.b_d_splines_slim(args_spline.tn2, args_spline.pn[1], eta2, span2, args_spline.bn2, args_spline.bd2)
+    bsplines_kernels.b_d_splines_slim(args_spline.tn3, args_spline.pn[2], eta3, span3, args_spline.bn3, args_spline.bd3)
 
-    if kind[0] == 0:
-        b1 = bn1
+    if args_spline.kind[0] == 0:
+        b1 = args_spline.bn1
     else:
-        b1 = bd1
+        b1 = args_spline.bd1
 
-    if kind[1] == 0:
-        b2 = bn2
+    if args_spline.kind[1] == 0:
+        b2 = args_spline.bn2
     else:
-        b2 = bd2
+        b2 = args_spline.bd2
 
-    if kind[2] == 0:
-        b3 = bn3
+    if args_spline.kind[2] == 0:
+        b3 = args_spline.bn3
     else:
-        b3 = bd3
+        b3 = args_spline.bd3
 
     value = eval_spline_mpi_kernel(
-        pn[0] - kind[0],
-        pn[1] - kind[1],
-        pn[2] - kind[2],
+        args_spline.pn[0] - args_spline.kind[0],
+        args_spline.pn[1] - args_spline.kind[1],
+        args_spline.pn[2] - args_spline.kind[2],
         b1,
         b2,
         b3,
         span1,
         span2,
         span3,
         _data,
-        starts,
+        args_spline.starts,
     )
 
     return value
@@ -1237,12 +1216,7 @@ def eval_spline_mpi_tensor_product(
     eta2: "float[:]",
     eta3: "float[:]",
     _data: "float[:,:,:]",
-    kind: "int[:]",
-    pn: "int[:]",
-    tn1: "float[:]",
-    tn2: "float[:]",
-    tn3: "float[:]",
-    starts: "int[:]",
+    args_spline: "SplineArguments",
     values: "float[:,:,:]",
 ):
     """
@@ -1257,17 +1231,9 @@ def eval_spline_mpi_tensor_product(
         _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.
-
-        tn1, tn2, tn3 : array[float]
-            Knot vectors of V0 in each direction.
-
-        starts : array[float]
-            Start indices of splines on current process.
+        args_spline : SplineArguments
+            Kind of 1d basis, spline degrees and knot vectors of V0, and start indices of the splines on the
+            current process.
 
         values : array[float]
             Return array for spline values S_ijk = S(eta1[i], eta2[j], eta3[k]).
@@ -1283,21 +1249,15 @@ def eval_spline_mpi_tensor_product(
                 if eta3[k] == -1.0:
                     continue  # point not in process domain
 
-                values[i, j, k] = eval_spline_mpi(eta1[i], eta2[j], eta3[k], _data, kind, pn, tn1, tn2, tn3, starts)
+                values[i, j, k] = eval_spline_mpi(eta1[i], eta2[j], eta3[k], _data, args_spline)
 
 
-@stack_array("bn1", "bn2", "bn3", "bd1", "bd2", "bd3")
 def eval_spline_mpi_tensor_product_fast(
     eta1: "float[:]",
     eta2: "float[:]",
     eta3: "float[:]",
     _data: "float[:,:,:]",
-    kind: "int[:]",
-    pn: "int[:]",
-    tn1: "float[:]",
-    tn2: "float[:]",
-    tn3: "float[:]",
-    starts: "int[:]",
+    args_spline: "SplineArguments",
     values: "float[:,:,:]",
 ):
     """
@@ -1313,75 +1273,64 @@ def eval_spline_mpi_tensor_product_fast(
         _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.
-
-        tn1, tn2, tn3 : array[float]
-            Knot vectors of V0 in each direction.
-
-        starts : array[float]
-            Start indices of splines on current process.
+        args_spline : SplineArguments
+            Kind of 1d basis, spline degrees and knot vectors of V0, and start indices of the splines on the
+            current process.
 
         values : array[float]
             Return array for spline values S_ijk = S(eta1[i], eta2[j], eta3[k]).
     """
-    # allocate spline values
-    bn1 = empty(pn[0] + 1, dtype=float)
-    bn2 = empty(pn[1] + 1, dtype=float)
-    bn3 = empty(pn[2] + 1, dtype=float)
-
-    bd1 = empty(pn[0], dtype=float)
-    bd2 = empty(pn[1], dtype=float)
-    bd3 = empty(pn[2], dtype=float)
-
     for i in range(len(eta1)):
         if eta1[i] == -1.0:
             continue  # point not in process domain
 
-        span1 = bsplines_kernels.find_span(tn1, pn[0], eta1[i])
-        bsplines_kernels.b_d_splines_slim(tn1, pn[0], eta1[i], span1, bn1, bd1)
-        if kind[0] == 0:
-            b1 = bn1
+        span1 = bsplines_kernels.find_span(args_spline.tn1, args_spline.pn[0], eta1[i])
+        bsplines_kernels.b_d_splines_slim(
+            args_spline.tn1, args_spline.pn[0], eta1[i], span1, args_spline.bn1, args_spline.bd1
+        )
+        if args_spline.kind[0] == 0:
+            b1 = args_spline.bn1
         else:
-            b1 = bd1
+            b1 = args_spline.bd1
 
         for j in range(len(eta2)):
             if eta2[j] == -1.0:
                 continue  # point not in process domain
 
-            span2 = bsplines_kernels.find_span(tn2, pn[1], eta2[j])
-            bsplines_kernels.b_d_splines_slim(tn2, pn[1], eta2[j], span2, bn2, bd2)
-            if kind[1] == 0:
-                b2 = bn2
+            span2 = bsplines_kernels.find_span(args_spline.tn2, args_spline.pn[1], eta2[j])
+            bsplines_kernels.b_d_splines_slim(
+                args_spline.tn2, args_spline.pn[1], eta2[j], span2, args_spline.bn2, args_spline.bd2
+            )
+            if args_spline.kind[1] == 0:
+                b2 = args_spline.bn2
             else:
-                b2 = bd2
+                b2 = args_spline.bd2
 
             for k in range(len(eta3)):
                 if eta3[k] == -1.0:
                     continue  # point not in process domain
 
-                span3 = bsplines_kernels.find_span(tn3, pn[2], eta3[k])
-                bsplines_kernels.b_d_splines_slim(tn3, pn[2], eta3[k], span3, bn3, bd3)
-                if kind[2] == 0:
-                    b3 = bn3
+                span3 = bsplines_kernels.find_span(args_spline.tn3, args_spline.pn[2], eta3[k])
+                bsplines_kernels.b_d_splines_slim(
+                    args_spline.tn3, args_spline.pn[2], eta3[k], span3, args_spline.bn3, args_spline.bd3
+                )
+                if args_spline.kind[2] == 0:
+                    b3 = args_spline.bn3
                 else:
-                    b3 = bd3
+                    b3 = args_spline.bd3
 
                 values[i, j, k] = eval_spline_mpi_kernel(
-                    pn[0] - kind[0],
-                    pn[1] - kind[1],
-                    pn[2] - kind[2],
+                    args_spline.pn[0] - args_spline.kind[0],
+                    args_spline.pn[1] - args_spline.kind[1],
+                    args_spline.pn[2] - args_spline.kind[2],
                     b1,
                     b2,
                     b3,
                     span1,
                     span2,
                     span3,
                     _data,
-                    starts,
+                    args_spline.starts,
                 )