diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index f41ea0446..eb2cf47dc 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` | @@ -540,6 +540,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`, @@ -566,6 +567,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). The device helper `eval_spline_mpi` keeps its flat signature. + ### 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/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..7bb5024e0 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,6 @@ 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.kind, + args_spline.pn, args_spline.tn1, args_spline.tn2, args_spline.tn3, args_spline.starts); } 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..22291157e 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,10 @@ def eval_spline_mpi_markers( markers[ip, 1], markers[ip, 2], _data, - kind, - pn, - tn1, - tn2, - tn3, - starts, + args_spline.kind, + args_spline.pn, + args_spline.tn1, + args_spline.tn2, + args_spline.tn3, + args_spline.starts, ) 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..674f15aa7 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,7 @@ 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.kind, args_spline.pn, + args_spline.tn1, args_spline.tn2, args_spline.tn3, args_spline.starts); } 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..abb2f903a 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,10 @@ def eval_spline_mpi_matrix( eta2[i, j, k], eta3[i, j, k], _data, - kind, - pn, - tn1, - tn2, - tn3, - starts, + args_spline.kind, + args_spline.pn, + args_spline.tn1, + args_spline.tn2, + args_spline.tn3, + args_spline.starts, ) 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..0994a7c94 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,7 @@ 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.kind, args_spline.pn, + args_spline.tn1, args_spline.tn2, args_spline.tn3, args_spline.starts); } 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..0fe4116db 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,10 @@ def eval_spline_mpi_sparse_meshgrid( eta2[0, j, 0], eta3[0, 0, k], _data, - kind, - pn, - tn1, - tn2, - tn3, - starts, + args_spline.kind, + args_spline.pn, + args_spline.tn1, + args_spline.tn2, + args_spline.tn3, + args_spline.starts, ) 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 e0bb22981..2f67ce133 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 @@ -2312,8 +2314,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.") @@ -2322,8 +2325,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) ) @@ -2911,7 +2914,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, ) @@ -2920,7 +2923,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, ) @@ -2931,7 +2934,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, ) @@ -2971,7 +2974,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, ) @@ -2980,7 +2983,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, ) @@ -2991,7 +2994,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..de1bd2257 --- /dev/null +++ b/src/struphy/kernel_arguments/spline_args_cuda.py @@ -0,0 +1,37 @@ +"""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. + """ + + 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..2ecd020fd --- /dev/null +++ b/src/struphy/kernel_arguments/spline_args_kernels.py @@ -0,0 +1,38 @@ +# NOTE: This file must use ONLY numpy for pyccel compilation compatibility. +# Backend conversion (NumPy/CuPy) happens at the Python wrapper level. + + +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. + """ + + 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 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.