diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index e9711804f..de132864f 100644 --- a/CUDA_STRATEGY.md +++ b/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` 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` 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//_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 diff --git a/src/struphy/bsplines/evaluation_kernels_3d.cuh b/src/struphy/bsplines/evaluation_kernels_3d.cuh index 3c4c353ad..ede3e5f40 100644 --- a/src/struphy/bsplines/evaluation_kernels_3d.cuh +++ b/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 _data, - const long long* kind, const long long* pn, Array1D tn1, - Array1D tn2, Array1D 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; } diff --git a/src/struphy/bsplines/evaluation_kernels_3d.py b/src/struphy/bsplines/evaluation_kernels_3d.py index e54f5d2be..d67f69781 100644 --- a/src/struphy/bsplines/evaluation_kernels_3d.py +++ b/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,17 +1160,9 @@ 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 ------- @@ -1182,43 +1170,34 @@ def eval_spline_mpi( 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, @@ -1226,7 +1205,7 @@ def eval_spline_mpi( 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,67 +1273,56 @@ 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, @@ -1381,7 +1330,7 @@ def eval_spline_mpi_tensor_product_fast( span2, span3, _data, - starts, + args_spline.starts, ) diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu index 7e4785dd0..c39ccfdda 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu @@ -1,4 +1,5 @@ #include "struphy/bsplines/evaluation_kernels_3d.cuh" +#include "struphy/kernel_arguments/spline_args.cuh" // Same arguments, in the same order, as the pyccel kernel in eval_spline_mpi_markers_kernels.py. Array views carry // dimensions and strides, so no CUDA-only lengths are needed. @@ -11,18 +12,12 @@ * @param markers Marker coordinates (Np x 3 or wider, any strides); rows flagged with -1 in the first * column are not on the process domain and are skipped. * @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 tn1 Knot vector of V0 along the first axis. - * @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 of V0 + * in each direction, and start indices of the splines on the current process. * @param values Output values S_p = S(*markers[p, :]), one per row; skipped rows are left unchanged. */ extern "C" __global__ void eval_spline_mpi_markers(Array2D markers, Array3D _data, - const long long* kind, long long* pn, Array1D tn1, - Array1D tn2, Array1D tn3, long long* starts, - Array1D values) { + SplineArgs args_spline, Array1D values) { // CUDA-only: the marker row of this thread replaces the pyccel loop variable long long ip = (long long)blockDim.x * blockIdx.x + threadIdx.x; long long Np = markers.shape[0]; @@ -31,6 +26,5 @@ extern "C" __global__ void eval_spline_mpi_markers(Array2D markers, Arra // point not in process domain if (markers(ip, 0) == -1.) return; - values(ip) = eval_spline_mpi(markers(ip, 0), markers(ip, 1), markers(ip, 2), _data, kind, pn, tn1, tn2, tn3, - starts); + values(ip) = eval_spline_mpi(markers(ip, 0), markers(ip, 1), markers(ip, 2), _data, args_spline); } diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py index 8531ad560..9b166e165 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py @@ -2,18 +2,15 @@ from numpy import shape +import struphy.kernel_arguments.spline_args_kernels as spline_args_kernels # do not remove; needed to identify dependencies from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi +from struphy.kernel_arguments.spline_args_kernels import SplineArguments def eval_spline_mpi_markers( markers: "float[:,:]", _data: "float[:,:,:]", - kind: "int[:]", - pn: "int[:]", - tn1: "float[:]", - tn2: "float[:]", - tn3: "float[:]", - starts: "int[:]", + args_spline: "SplineArguments", values: "float[:]", ): """ @@ -28,17 +25,9 @@ def eval_spline_mpi_markers( _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 1D array for spline values S_p = S(*markers[p, :]). @@ -55,10 +44,5 @@ def eval_spline_mpi_markers( markers[ip, 1], markers[ip, 2], _data, - kind, - pn, - tn1, - tn2, - tn3, - starts, + args_spline, ) diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_cuda.cu b/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_cuda.cu index 5880d6790..61cceedc0 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_cuda.cu +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_cuda.cu @@ -1,4 +1,5 @@ #include "struphy/bsplines/evaluation_kernels_3d.cuh" +#include "struphy/kernel_arguments/spline_args.cuh" // Same arguments, in the same order, as the pyccel kernel in eval_spline_mpi_matrix_kernels.py. Array views carry // dimensions and strides, so no CUDA-only lengths are needed. @@ -12,18 +13,13 @@ * @param eta2 Second coordinates of the points, same shape as eta1. * @param eta3 Third coordinates of the points, same shape as eta1. * @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 tn1 Knot vector of V0 along the first axis. - * @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 of V0 + * in each direction, and start indices of the splines on the current process. * @param values Output values, same shape as eta1; entries of flagged points are left unchanged. */ extern "C" __global__ void eval_spline_mpi_matrix(Array3D eta1, Array3D eta2, Array3D eta3, - Array3D _data, const long long* kind, long long* pn, - Array1D tn1, Array1D tn2, Array1D tn3, - long long* starts, Array3D values) { + Array3D _data, SplineArgs args_spline, + Array3D 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; if (ijk >= values.shape[0] * values.shape[1] * values.shape[2]) return; @@ -34,6 +30,5 @@ extern "C" __global__ void eval_spline_mpi_matrix(Array3D eta1, Array3D< // point not in process domain if (eta1(i, j, k) == -1. || eta2(i, j, k) == -1. || eta3(i, j, k) == -1.) return; - values(i, j, k) = eval_spline_mpi(eta1(i, j, k), eta2(i, j, k), eta3(i, j, k), _data, kind, pn, tn1, tn2, tn3, - starts); + values(i, j, k) = eval_spline_mpi(eta1(i, j, k), eta2(i, j, k), eta3(i, j, k), _data, args_spline); } diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_kernels.py b/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_kernels.py index 101d7cc5d..407c9750e 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_kernels.py +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_matrix/eval_spline_mpi_matrix_kernels.py @@ -2,7 +2,9 @@ from numpy import shape +import struphy.kernel_arguments.spline_args_kernels as spline_args_kernels # do not remove; needed to identify dependencies from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi +from struphy.kernel_arguments.spline_args_kernels import SplineArguments def eval_spline_mpi_matrix( @@ -10,12 +12,7 @@ def eval_spline_mpi_matrix( eta2: "float[:,:,:]", eta3: "float[:,:,:]", _data: "float[:,:,:]", - kind: "int[:]", - pn: "int[:]", - tn1: "float[:]", - tn2: "float[:]", - tn3: "float[:]", - starts: "int[:]", + args_spline: "SplineArguments", values: "float[:,:,:]", ): """ @@ -30,17 +27,9 @@ def eval_spline_mpi_matrix( _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,j,k], eta2[i,j,k], eta3[i,j,k]). @@ -63,10 +52,5 @@ def eval_spline_mpi_matrix( eta2[i, j, k], eta3[i, j, k], _data, - kind, - pn, - tn1, - tn2, - tn3, - starts, + args_spline, ) diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_cuda.cu b/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_cuda.cu index 7f17a381c..1fde8a9a3 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_cuda.cu +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_cuda.cu @@ -1,7 +1,8 @@ #include "struphy/bsplines/evaluation_kernels_3d.cuh" +#include "struphy/kernel_arguments/spline_args.cuh" -// Same arguments, in the same order, as the pyccel kernel in eval_spline_mpi_sparse_meshgrid_kernels.py. Array views carry -// dimensions and strides, so no CUDA-only lengths are needed. +// Same arguments, in the same order, as the pyccel kernel in eval_spline_mpi_sparse_meshgrid_kernels.py. Array views +// carry dimensions and strides, so no CUDA-only lengths are needed. /** * Sparse meshgrid evaluation of a distributed spline, as in evaluation_kernels_3d.eval_spline_mpi_sparse_meshgrid. @@ -12,19 +13,13 @@ * @param eta2 Second coordinates, shape (1, n2, 1). * @param eta3 Third coordinates, shape (1, 1, n3). * @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 tn1 Knot vector of V0 along the first axis. - * @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 of V0 + * in each direction, and start indices of the splines on the current process. * @param values Output values, shape (n1, n2, n3); entries of flagged points are left unchanged. */ extern "C" __global__ void eval_spline_mpi_sparse_meshgrid(Array3D eta1, Array3D eta2, Array3D eta3, Array3D _data, - const long long* kind, long long* pn, Array1D tn1, - Array1D tn2, Array1D tn3, - long long* starts, Array3D values) { + SplineArgs args_spline, Array3D 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; if (ijk >= values.shape[0] * values.shape[1] * values.shape[2]) return; @@ -35,6 +30,5 @@ extern "C" __global__ void eval_spline_mpi_sparse_meshgrid(Array3D eta1, // point not in process domain if (eta1(i, 0, 0) == -1. || eta2(0, j, 0) == -1. || eta3(0, 0, k) == -1.) return; - values(i, j, k) = eval_spline_mpi(eta1(i, 0, 0), eta2(0, j, 0), eta3(0, 0, k), _data, kind, pn, tn1, tn2, tn3, - starts); + values(i, j, k) = eval_spline_mpi(eta1(i, 0, 0), eta2(0, j, 0), eta3(0, 0, k), _data, args_spline); } diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_kernels.py b/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_kernels.py index 41bf26141..527993750 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_kernels.py +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_sparse_meshgrid/eval_spline_mpi_sparse_meshgrid_kernels.py @@ -1,6 +1,8 @@ """Sparse meshgrid evaluation of a tensor-product spline, distributed.""" +import struphy.kernel_arguments.spline_args_kernels as spline_args_kernels # do not remove; needed to identify dependencies from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi +from struphy.kernel_arguments.spline_args_kernels import SplineArguments def eval_spline_mpi_sparse_meshgrid( @@ -8,12 +10,7 @@ def eval_spline_mpi_sparse_meshgrid( eta2: "float[:,:,:]", eta3: "float[:,:,:]", _data: "float[:,:,:]", - kind: "int[:]", - pn: "int[:]", - tn1: "float[:]", - tn2: "float[:]", - tn3: "float[:]", - starts: "int[:]", + args_spline: "SplineArguments", values: "float[:,:,:]", ): """ @@ -28,17 +25,9 @@ def eval_spline_mpi_sparse_meshgrid( _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,0,0], eta2[0,j,0], eta3[0,0,k]). @@ -65,10 +54,5 @@ def eval_spline_mpi_sparse_meshgrid( eta2[0, j, 0], eta3[0, 0, k], _data, - kind, - pn, - tn1, - tn2, - tn3, - starts, + args_spline, ) diff --git a/src/struphy/bsplines/tests/test_eval_spline_mpi.py b/src/struphy/bsplines/tests/test_eval_spline_mpi.py index ea67c7437..99f4c850c 100644 --- a/src/struphy/bsplines/tests/test_eval_spline_mpi.py +++ b/src/struphy/bsplines/tests/test_eval_spline_mpi.py @@ -256,6 +256,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import create_equal_random_arrays as cera from struphy.io.options import DerhamOptions + from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.topology.grids import TensorProductGrid comm = MPI.COMM_WORLD @@ -275,6 +276,9 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): x2, x2_psy = cera(derham.V2fem) x3, x3_psy = cera(derham.V3fem) + def args_spline(kind): + return SplineArguments(kind, xp.array(derham.degree), tn1, tn2, tn3, xp.array(x0_psy.starts)) + # Random points in domain of process dom = derham.domain_array[rank] eta1s = xp.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] @@ -303,12 +307,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x0_psy._data, - derham.V0splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V0splines.spline_types_pyccel[0]), ) assert xp.allclose(val, val_mpi) @@ -339,12 +338,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x1_psy[0]._data, - derham.V1splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V1splines.spline_types_pyccel[0]), ) assert xp.allclose(val, val_mpi) @@ -374,12 +368,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x1_psy[1]._data, - derham.V1splines.spline_types_pyccel[1], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V1splines.spline_types_pyccel[1]), ) assert xp.allclose(val, val_mpi) @@ -409,12 +398,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x1_psy[2]._data, - derham.V1splines.spline_types_pyccel[2], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V1splines.spline_types_pyccel[2]), ) assert xp.allclose(val, val_mpi) @@ -445,12 +429,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x2_psy[0]._data, - derham.V2splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V2splines.spline_types_pyccel[0]), ) assert xp.allclose(val, val_mpi) @@ -480,12 +459,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x2_psy[1]._data, - derham.V2splines.spline_types_pyccel[1], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V2splines.spline_types_pyccel[1]), ) assert xp.allclose(val, val_mpi) @@ -515,12 +489,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x2_psy[2]._data, - derham.V2splines.spline_types_pyccel[2], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V2splines.spline_types_pyccel[2]), ) assert xp.allclose(val, val_mpi) @@ -548,12 +517,7 @@ def test_eval_pointwise(num_elements, degree, bcs, n_markers=10): eta2, eta3, x3_psy._data, - derham.V3splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V3splines.spline_types_pyccel[0]), ) assert xp.allclose(val, val_mpi) @@ -589,6 +553,7 @@ def test_eval_tensor_product(num_elements, degree, bcs, n_markers=10): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import create_equal_random_arrays as cera from struphy.io.options import DerhamOptions + from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.topology.grids import TensorProductGrid comm = MPI.COMM_WORLD @@ -606,6 +571,9 @@ def test_eval_tensor_product(num_elements, degree, bcs, n_markers=10): x0, x0_psy = cera(derham.V0fem) x3, x3_psy = cera(derham.V3fem) + def args_spline(kind): + return SplineArguments(kind, xp.array(derham.degree), tn1, tn2, tn3, xp.array(x0_psy.starts)) + # Random points in domain of process dom = derham.domain_array[rank] eta1s = xp.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] @@ -636,12 +604,7 @@ def test_eval_tensor_product(num_elements, degree, bcs, n_markers=10): eta2s, eta3s, x0_psy._data, - derham.V0splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V0splines.spline_types_pyccel[0]), vals_mpi, ) t1 = time.time() @@ -654,12 +617,7 @@ def test_eval_tensor_product(num_elements, degree, bcs, n_markers=10): eta2s, eta3s, x0_psy._data, - derham.V0splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V0splines.spline_types_pyccel[0]), vals_mpi_fast, ) t1 = time.time() @@ -696,12 +654,7 @@ def test_eval_tensor_product(num_elements, degree, bcs, n_markers=10): eta2s, eta3s, x3_psy._data, - derham.V3splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V3splines.spline_types_pyccel[0]), vals_mpi, ) t1 = time.time() @@ -714,12 +667,7 @@ def test_eval_tensor_product(num_elements, degree, bcs, n_markers=10): eta2s, eta3s, x3_psy._data, - derham.V3splines.spline_types_pyccel[0], - xp.array(derham.degree), - tn1, - tn2, - tn3, - xp.array(x0_psy.starts), + args_spline(derham.V3splines.spline_types_pyccel[0]), vals_mpi_fast, ) t1 = time.time() diff --git a/src/struphy/bsplines/tests/test_evaluation_cuda.py b/src/struphy/bsplines/tests/test_evaluation_cuda.py index c8ad23436..7cb2aa89c 100644 --- a/src/struphy/bsplines/tests/test_evaluation_cuda.py +++ b/src/struphy/bsplines/tests/test_evaluation_cuda.py @@ -7,6 +7,9 @@ import numpy as np import pytest +from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments +from struphy.kernel_arguments.spline_args_kernels import SplineArguments + requires_cupy = pytest.mark.skipif(not xp.cupy_available(), reason="CuPy/GPU not available") @@ -26,10 +29,11 @@ def test_evaluation_parity(mode, kind, empty): for backend in ("numpy", "cupy"): with xp.use_backend(backend): data = xp.asarray(coeff)[::2, ::2, ::2] - metadata = ( + args_class = CudaSplineArguments if backend == "cupy" else SplineArguments + args_spline = args_class( xp.asarray(kind, dtype=xp.int64), xp.asarray(degree), - *(xp.asarray(np.repeat(t, 2))[::2] for t in knots), + *(xp.asarray(t) for t in knots), xp.asarray([1, 1, 1], dtype=xp.int64), ) if mode == "markers": @@ -48,7 +52,7 @@ def test_evaluation_parity(mode, kind, empty): coords[1][0, 0, 0] = -1.0 storage = xp.full((axes[0].size, 5, 8), 17.0) out = storage[:, :, ::2] - kernel(*coords, data, *metadata, out, n_threads=out.size) + kernel(*coords, data, args_spline, out, n_threads=out.size) results.append(xp.to_numpy(out).copy()) assert bool(xp.all(storage[..., 1::2] == 17.0)) np.testing.assert_allclose(results[1], results[0], rtol=1e-12, atol=1e-12) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index a870d58b0..cef3335b9 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -46,6 +46,8 @@ from struphy.io.options import DerhamOptions, FieldsBackground, LiteralOptions from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments +from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments +from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator @@ -2336,8 +2338,9 @@ def __init__( # dimensions in each direction self._nbasis = derham.spline_attributes[space_id].nbasis - # arguments of the evaluation kernels, one (kind, pn, tn1, tn2, tn3, starts) per component, - # on the backend of the coefficients and the same for both kernel versions + # arguments of the evaluation kernels, one SplineArguments per component: the pyccel class on NumPy, + # the CUDA class on CuPy + args_class = CudaSplineArguments if xp.get_backend() == "cupy" else SplineArguments degree = np.asarray(derham.degree, dtype=np.int64) if xp.get_backend() == "cupy" and np.any((degree < 1) | (degree > 8)): raise ValueError("CUDA spline degrees must be between 1 and 8.") @@ -2346,8 +2349,8 @@ def __init__( knots = tuple(xp.asarray(np.ascontiguousarray(t, dtype=float)) for t in derham.V0fem.knots) starts = (self.starts,) if isinstance(self._vector_stencil, StencilVector) else self.starts kinds = derham.spline_attributes[self.space_key].spline_types_pyccel - self._args_eval = tuple( - (xp.asarray(kind, dtype=xp.int64), pn, *knots, xp.asarray(start, dtype=xp.int64)) + self._args_spline = tuple( + args_class(xp.asarray(kind, dtype=xp.int64), pn, *knots, xp.asarray(start, dtype=xp.int64)) for kind, start in zip(kinds, starts) ) @@ -2958,7 +2961,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E2, E3, self._vector_stencil._data, - *self._args_eval[0], + self._args_spline[0], tmp, n_threads=tmp.size, ) @@ -2967,7 +2970,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): eval_spline_mpi_markers( markers, self._vector_stencil._data, - *self._args_eval[0], + self._args_spline[0], tmp, n_threads=tmp.size, ) @@ -2978,7 +2981,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E2, E3, self._vector_stencil._data, - *self._args_eval[0], + self._args_spline[0], tmp, n_threads=tmp.size, ) @@ -3018,7 +3021,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E2, E3, self._vector_stencil[n]._data, - *self._args_eval[n], + self._args_spline[n], tmp, n_threads=tmp.size, ) @@ -3027,7 +3030,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): eval_spline_mpi_markers( markers, self._vector_stencil[n]._data, - *self._args_eval[n], + self._args_spline[n], tmp, n_threads=tmp.size, ) @@ -3038,7 +3041,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E2, E3, self._vector_stencil[n]._data, - *self._args_eval[n], + self._args_spline[n], tmp, n_threads=tmp.size, ) diff --git a/src/struphy/kernel_arguments/spline_args.cuh b/src/struphy/kernel_arguments/spline_args.cuh new file mode 100644 index 000000000..b136a710f --- /dev/null +++ b/src/struphy/kernel_arguments/spline_args.cuh @@ -0,0 +1,16 @@ +// Generated by cunumpy.arguments.CudaStruct from the Python definition; do not edit. +#ifndef STRUPHY_SPLINE_ARGS_CUH +#define STRUPHY_SPLINE_ARGS_CUH + +#include "cunumpy/array_view.cuh" + +struct SplineArgs { + long long* kind; + long long* pn; + Array1D tn1; + Array1D tn2; + Array1D tn3; + long long* starts; +}; + +#endif // STRUPHY_SPLINE_ARGS_CUH diff --git a/src/struphy/kernel_arguments/spline_args_cuda.py b/src/struphy/kernel_arguments/spline_args_cuda.py new file mode 100644 index 000000000..5c8c3ac14 --- /dev/null +++ b/src/struphy/kernel_arguments/spline_args_cuda.py @@ -0,0 +1,39 @@ +"""CUDA version of :class:`~struphy.kernel_arguments.spline_args_kernels.SplineArguments`. + +Takes the same constructor arguments and has the same attributes as the pyccel class, but holds CuPy arrays and +is passed to CUDA kernels as one C struct (``kernel_arguments/spline_args.cuh``, generated from +:attr:`CudaSplineArguments.fields`). +""" + +import numpy as np +from cunumpy.arguments import CudaStructArguments + +from struphy.kernel_arguments.pusher_args_cuda import _device_array + + +class CudaSplineArguments(CudaStructArguments): + """CUDA version of :class:`~struphy.kernel_arguments.spline_args_kernels.SplineArguments` (``SplineArgs``). + + The knot vectors are array views (pointer, shape, strides), because ``find_span`` needs the number of knots. + The pyccel scratch arrays ``bn1``, ..., ``bd3`` are per-thread local arrays in the CUDA kernels, so they are + not part of the struct. + """ + + struct_name = "SplineArgs" + fields = ( + ("kind", "long long*"), + ("pn", "long long*"), + ("tn1", "Array1D"), + ("tn2", "Array1D"), + ("tn3", "Array1D"), + ("starts", "long long*"), + ) + + def __init__(self, kind, pn, tn1, tn2, tn3, starts): + self.kind = _device_array("kind", kind, np.int64) + self.pn = _device_array("pn", pn, np.int64) + if bool(((pn < 1) | (pn > 8)).any()): + raise ValueError("CUDA spline degrees must be between 1 and 8.") + self.tn1, self.tn2, self.tn3 = (_device_array("tn", t, np.float64, ndim=1) for t in (tn1, tn2, tn3)) + self.starts = _device_array("starts", starts, np.int64) + self.pack() diff --git a/src/struphy/kernel_arguments/spline_args_kernels.py b/src/struphy/kernel_arguments/spline_args_kernels.py new file mode 100644 index 000000000..e0960bf2d --- /dev/null +++ b/src/struphy/kernel_arguments/spline_args_kernels.py @@ -0,0 +1,48 @@ +# NOTE: This file must use ONLY numpy for pyccel compilation compatibility. +# Backend conversion (NumPy/CuPy) happens at the Python wrapper level. +import numpy as np + + +class SplineArguments: + """Holds the arguments pertaining to one component of a :class:`~struphy.feec.psydac_derham.SplineFunction` + passed to the spline evaluation kernels (``bsplines/kernels/eval_spline_mpi_*``). + + Paramaters + ---------- + 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[int] + Start indices of the splines of this component on the current process. + + The scratch arrays ``bn1``, ..., ``bd3`` hold the B- and D-spline values at the current evaluation point. + """ + + def __init__( + self, + kind: "int[:]", + pn: "int[:]", + tn1: "float[:]", + tn2: "float[:]", + tn3: "float[:]", + starts: "int[:]", + ): + self.kind = kind + self.pn = pn + self.tn1 = tn1 + self.tn2 = tn2 + self.tn3 = tn3 + self.starts = starts + + self.bn1 = np.empty(int(pn[0] + 1), dtype=float) + self.bn2 = np.empty(int(pn[1] + 1), dtype=float) + self.bn3 = np.empty(int(pn[2] + 1), dtype=float) + self.bd1 = np.empty(int(pn[0]), dtype=float) + self.bd2 = np.empty(int(pn[1]), dtype=float) + self.bd3 = np.empty(int(pn[2]), dtype=float) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index abe1ef445..1d3d1566e 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -14,6 +14,8 @@ from struphy.geometry.domains import Cuboid from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaMarkerArguments from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, MarkerArguments +from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments +from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.ode.utils import ButcherTableau N_MARKERS = 129 # not a multiple of the block size @@ -89,20 +91,21 @@ def spline_coefficients(n=3, seed=11): def spline_evaluation_arguments(kind): - """``_data, kind, pn, tn1, tn2, tn3, starts`` of the spline evaluation kernels, for degrees 2, 3, 1 on 8 cells. + """``_data, args_spline`` of the spline evaluation kernels, for degrees 2, 3, 1 on 8 cells. The coefficients cover every span of the knots with start indices 1 on all axes. """ rng = np.random.default_rng(34) degree = np.array([2, 3, 1], dtype=np.int64) knots = [np.r_[np.zeros(p), np.linspace(0, 1, 9), np.ones(p)] for p in degree] - return ( - xp.asarray(rng.normal(size=(16, 18, 20))), + args_class = CudaSplineArguments if xp.get_backend() == "cupy" else SplineArguments + args_spline = args_class( xp.asarray(kind, dtype=np.int64), xp.asarray(degree), *(xp.asarray(t) for t in knots), xp.ones(3, dtype=np.int64), ) + return xp.asarray(rng.normal(size=(16, 18, 20))), args_spline def evaluation_grid(sparse): diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index f2ab111a2..fbd8d9d7d 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -10,6 +10,7 @@ from struphy.geometry.tests import spline_mapping_cases from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments +from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.pic.tests.cuda_emulation import emulate_struct_kernel from struphy.pic.tests.cuda_parity_cases import PARITY_CASES from struphy.pic.tests.kernel_test_args import N_GEOMETRY_DOMAINS @@ -17,7 +18,8 @@ # pyccel argument class name -> its CUDA version CUDA_CLASSES = { - cls.__name__.removeprefix("Cuda"): cls for cls in (CudaMarkerArguments, CudaDerhamArguments, CudaDomainArguments) + cls.__name__.removeprefix("Cuda"): cls + for cls in (CudaMarkerArguments, CudaDerhamArguments, CudaDomainArguments, CudaSplineArguments) } requires_compiler = pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 41e64dade..e373ddea4 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -18,14 +18,22 @@ from struphy.kernel_arguments.local_projectors_args_cuda import CudaLocalProjectorsArguments from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments, MarkerArguments +from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments +from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.pic.tests.kernel_test_args import N_GEOMETRY_DOMAINS -from struphy.utils.cuda_arguments import CUDA_OPTIONS, write_local_projectors_header, write_pusher_header +from struphy.utils.cuda_arguments import ( + CUDA_OPTIONS, + write_local_projectors_header, + write_pusher_header, + write_spline_header, +) N_COLS = 25 MARKER_INDICES = (3, 6, 7, 8, 14, 17, 18, 4) ARGS_DIR = Path(struphy.__file__).parent / "kernel_arguments" HEADER = ARGS_DIR / "pusher_args.cuh" LOCAL_PROJECTORS_HEADER = ARGS_DIR / "local_projectors_args.cuh" +SPLINE_HEADER = ARGS_DIR / "spline_args.cuh" # (CUDA class, pyccel source, pyccel class, struct members that only the CUDA class has, its header) ARGUMENT_PAIRS = ( @@ -39,6 +47,7 @@ set(), LOCAL_PROJECTORS_HEADER, ), + (CudaSplineArguments, "spline_args_kernels.py", "SplineArguments", set(), SPLINE_HEADER), ) @@ -93,6 +102,10 @@ def test_generated_local_projectors_header(tmp_path): assert LOCAL_PROJECTORS_HEADER.read_text() == write_local_projectors_header(tmp_path / "local_projectors_args.cuh") +def test_generated_spline_header(tmp_path): + assert SPLINE_HEADER.read_text() == write_spline_header(tmp_path / "spline_args.cuh") + + @pytest.mark.parametrize("cuda_class, source, class_name, cuda_only, header", ARGUMENT_PAIRS) def test_argument_classes_correspond(cuda_class, source, class_name, cuda_only, header): """Each CUDA argument class mirrors its pyccel class: same constructor, same attributes in the same order.""" @@ -121,6 +134,9 @@ def test_owners_select_pyccel_classes_on_numpy(): assert type(domain.args_domain) is DomainArguments assert type(particles.args_markers) is MarkerArguments assert type(derham.args_derham) is DerhamArguments + for space in ("H1", "Hcurl"): + spline = derham.create_spline_function("f", space) + assert all(type(args) is SplineArguments for args in spline._args_spline) @requires_cupy diff --git a/src/struphy/utils/cuda_arguments.py b/src/struphy/utils/cuda_arguments.py index eacf78e96..e9896016c 100644 --- a/src/struphy/utils/cuda_arguments.py +++ b/src/struphy/utils/cuda_arguments.py @@ -12,10 +12,12 @@ from struphy.kernel_arguments.local_projectors_args_cuda import CudaLocalProjectorsArguments from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments +from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments PUSHER_STRUCTS = tuple(cls.struct for cls in (CudaMarkerArguments, CudaDerhamArguments, CudaDomainArguments)) LOCAL_PROJECTORS_STRUCTS = (CudaLocalProjectorsArguments.struct,) -CUDA_STRUCTS = PUSHER_STRUCTS + LOCAL_PROJECTORS_STRUCTS +SPLINE_STRUCTS = (CudaSplineArguments.struct,) +CUDA_STRUCTS = PUSHER_STRUCTS + LOCAL_PROJECTORS_STRUCTS + SPLINE_STRUCTS CUDA_INCLUDE_DIR = Path(__file__).resolve().parents[2] CUDA_OPTIONS = {"structs": CUDA_STRUCTS, "include_dirs": (CUDA_INCLUDE_DIR,)} @@ -30,6 +32,11 @@ def write_local_projectors_header(path): return write_cuda_header(path, LOCAL_PROJECTORS_STRUCTS, guard="STRUPHY_LOCAL_PROJECTORS_ARGS_CUH") +def write_spline_header(path): + """Generate the committed ABI header ``spline_args.cuh`` from the CUDA argument class.""" + return write_cuda_header(path, SPLINE_STRUCTS, guard="STRUPHY_SPLINE_ARGS_CUH") + + def check_mapping_on_device(kind_map: int, what: str): """Raise on the CuPy backend if `kind_map` is a spline mapping, which has no CUDA version yet.