diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index ba8fb9d3d..a45f283ae 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -13,7 +13,7 @@ The work is split into small PRs that can be reviewed and merged one at a time. - [x] **PR 5: `Domain` on the GPU** — domain arguments are selected and stored at domain construction; CUDA arguments reference device arrays, and deepcopy/unpickling rebuilds the arguments from the copied or restored arrays. - [x] **PR 6: `Particles` on the GPU** — `Particles` can be created on the CuPy backend, and `Particles.args_markers` is selected as the CUDA or Pyccel argument bundle at construction. - [x] **PR 7: `Derham` on the GPU** — `Derham` can be created on the CuPy backend, plus `Derham.cuda_args_derham`. -- [ ] **PR 8: Shared CUDA headers for the argument classes** — one `.cuh` per argument class instead of long flat kernel signatures. +- [x] **PR 8: Shared CUDA header for the argument classes** — one struct per argument class in `kernel_arguments/pusher_args.cuh`, passed by value, instead of long flat kernel signatures. - [ ] **PR 9: One folder per kernel, starting with `pic/pushing`** — pure refactor, no behaviour change. - [ ] **PR 10: Device versions of helper kernels** — B-spline evaluation, mapping evaluation (per domain), small linear algebra, as `__device__` functions in `.cuh` headers. - [ ] **PR 11: First real CUDA kernel** — `push_eta_stage` with a pyccel/CUDA parity test and an end-to-end run on the GPU. @@ -41,12 +41,13 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke - **No silent CPU fallback on the GPU.** A kernel without a CUDA version raises an error on the GPU backend. Falling back would mean copying data to the host and back at every call. - **Small steps.** Every PR keeps the CPU code path working and tested. -## Current state (PR 7) +## Current state (PR 8) | File | Content | |---|---| -| `src/struphy/utils/kernel_backends.py` | `is_cuda_backend()`, `CudaKernel` (wraps a `cupy.RawKernel`, compiled lazily; expands `Argument.get_cuda_args()` and takes `n_threads`), `Kernel` and `KernelCatalog` for backend selection and discovery | -| `src/struphy/utils/cuda_arguments.py` | `Argument` contract plus `CudaMarkerArguments`, `CudaDerhamArguments` and `CudaDomainArguments`; CUDA arrays are stored individually and returned in signature order by `get_cuda_args()` | +| `src/struphy/utils/kernel_backends.py` | `is_cuda_backend()`, `CudaKernel` (wraps a `cupy.RawKernel`, compiled lazily with the struphy headers on the include path; expands `Argument.get_cuda_args()` and takes `n_threads`), `Kernel` and `KernelCatalog` for backend selection and discovery | +| `src/struphy/utils/cuda_arguments.py` | `Argument` base class plus `CudaMarkerArguments`, `CudaDerhamArguments` and `CudaDomainArguments`; each references its device arrays and packs them once into its C struct, which `get_cuda_args()` returns | +| `src/struphy/kernel_arguments/pusher_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs` and `DomainArgs` that CUDA kernels take in place of the pyccel argument classes | | `src/struphy/geometry/base.py` | `Domain.args_domain` is selected once at construction; CUDA domains use device arrays, while direct Pyccel geometry calls retain a host argument bundle | | `src/struphy/pic/base.py` | `Particles` arrays and `args_markers` use the backend selected at construction; direct Pyccel methods retain a private host bundle | | `src/struphy/pic/tests/test_kernel_backends.py` | the demo kernel pair `push_eta_linear` (pyccel function compiled with `epyccel` at test time, CUDA source string) and tests on both backends | @@ -57,7 +58,7 @@ Things we learned in the proof of concept: - `Particles.args_markers` is the CUDA or Pyccel bundle selected at construction. Existing direct Pyccel methods use a private host bundle. `Derham` still needs its CUDA creation work (PR 7). - `Domain.args_domain` returns the argument type selected at domain construction; on CuPy it is a `CudaDomainArguments` object. Direct Pyccel methods on `Domain` use a separate internal host bundle. - `cupy.RawKernel` accepts only device arrays (host arrays raise) and does **not** check the kernel signature. Each argument is read with the size declared in the signature, so Python `int`/`float` arrive correctly in `int`/`double` parameters, but a wrongly typed scalar (e.g. an integer for a `double`, or a value that overflows an `int`) gives a wrong value **without an error**. Casting Python scalars in `CudaKernel` does not prevent this, so it is not done; see the follow-up in [Open questions](#open-questions). -- Flattening the argument classes at each call (joining their `values`) costs well under 1 µs, compared to about 70 µs for launching the kernel. +- Flattening the argument classes at each call costs well under 1 µs, compared to about 70 µs for launching the kernel. Since PR 8 each class is one struct, packed once when the argument object is created. - `struphy compile` compiles every `.py` file whose name contains `kernels`. Non-pyccel modules must not contain `kernels` in their name; `.cu` files are ignored by it. - On an H100, the demo kernel pushes 10⁶ markers in about 0.13 ms per step. @@ -100,7 +101,7 @@ kernel = catalog["push_eta_stage"] # Kernel: pyccel or CUDA depending on the ba - `CudaKernel.from_file(path)` reads `_cuda.cu`; the kernel name is taken from the file name. The kernel is compiled lazily on first call. CuPy caches compiled kernels on disk (`~/.cupy/kernel_cache`), so the compile cost is paid once per machine. - Add `"**/*.cu"` and `"**/*.cuh"` to `[tool.setuptools.package-data]` in `pyproject.toml`. -- Later (PR 10), when the first shared header is needed: headers are found through NVRTC include paths (`cupy.RawModule(code=..., options=("-I",))`), so a `.cu` file can `#include "struphy/bsplines/bsplines_kernels.cuh"`. +- Shared headers are found through the NVRTC include path (done in PR 8: `CudaKernel` compiles with `-I`), so a `.cu` file can `#include "struphy/kernel_arguments/pusher_args.cuh"` or, later, `"struphy/bsplines/bsplines_kernels.cuh"`. ### PR 3: Kernel catalog @@ -156,13 +157,16 @@ kernel = catalog["push_eta_stage"] # Kernel: pyccel or CUDA depending on the ba - Field evaluation (`SplineFunction.__call__`, ...) still calls pyccel kernels with the coefficients, which are device arrays on CuPy. It needs CUDA evaluation kernels (PR 10+). - Tests: `feec/tests/test_derham_gpu.py`. Without a GPU, a strict host stand-in for CuPy (rejects host/device mixing, cannot run kernels; not part of the repository) was used, on top of struphy-hub/feectools#85. With it, a `Derham` created on the "CuPy" backend matches the NumPy one on 1, 2 and 4 MPI processes. -### PR 8: Argument structs in shared headers +### PR 8: Argument structs in a shared header (complete) -Today every CUDA kernel repeats the full flat signature (26 parameters for markers and domain alone). CuPy does not check it, so adding a field to `MarkerArguments` would shift all following arguments of all CUDA kernels **without an error**. +Before, every CUDA kernel repeated the full flat signature (26 parameters for markers and domain alone). CuPy does not check it, so adding a field to `MarkerArguments` would have shifted all following arguments of all CUDA kernels **without an error**. -- Define `struct MarkerArgs { double* markers; bool* valid_mks; int n_markers; ... };` etc. in `kernel_arguments/pusher_args.cuh`, and pass one struct per argument class. -- To check first: how to pass a struct to a `cupy.RawKernel` (e.g. as a NumPy structured scalar with pointer fields). If this does not work well, keep the flat signature, but generate it from the Python class so that it is defined in one place. -- A test compares the struct layout (field names, types, order) with the Python class. +- `kernel_arguments/pusher_args.cuh` defines `struct MarkerArgs`, `DerhamArgs` and `DomainArgs`. A kernel takes them by value in the place of the pyccel argument classes, e.g. `void push_eta_linear(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain)`, and reads `args_markers.markers`, `args_markers.n_markers`, ... The member names are the attribute names of the pyccel classes; the only CUDA-specific member is `MarkerArgs.n_cols` (pyccel takes `markers.shape[1]`). +- Passing a struct: CuPy passes a NumPy scalar by value, copying `itemsize` bytes. Each `Argument` subclass lists its members as `fields = (("double*", "markers"), ...)`, from which `struct_dtype()` builds a NumPy structured dtype with `align=True` (C alignment and padding). Pointer members are `uint64` holding `array.data.ptr`. The struct (a `numpy.void`) is packed once, in the constructor, so a kernel call does no extra work. +- The struct holds device addresses: it is repacked when an argument object is deepcopied or unpickled (`__getstate__`/`__setstate__`), and owners that reallocate an array must rebuild their argument object, as before. +- Scalar members are checked when packing: a `float` for an `int` member raises `TypeError` and a value that does not fit raises `OverflowError`, instead of arriving truncated or wrapped around (NumPy assignment alone truncates `1.7` to `1`). +- Tests: the header is parsed and compared with `fields` (names, C types, order) without a GPU; on a GPU, an NVRTC-compiled kernel reports `sizeof` and the member offsets, which are compared with the dtype. The demo kernels use the structs. On the host, `offsetof`/`sizeof` from a C++ compiler agree with the dtypes (x86-64/arm64 lay these structs out like CUDA). +- Fixed along the way: the GPU-only tests still used the `values` attribute that was replaced by `get_cuda_args()`. ### PR 9: One folder per kernel @@ -211,7 +215,7 @@ Port the kernels in the order the target models need them, so that complete mode ## Open questions -- **Scalar types.** Scalars are not checked against the kernel signature (see [Current state](#current-state-pr-1)). `CudaKernel` could read the parameter types from the `extern "C"` signature once, when it is created, and cast each scalar to its declared type or raise if it does not fit (e.g. a Python `float` for an `int`, or an overflowing integer). This could go together with PR 8. +- **Scalar types.** Since PR 8 the members of the argument structs are checked when they are packed. Scalars passed directly to a kernel (e.g. `dt`, `stage`) are still not checked against the kernel signature (see [Current state](#current-state-pr-8)). `CudaKernel` could read the parameter types from the `extern "C"` signature once, when it is created, and cast each scalar to its declared type or raise if it does not fit (e.g. a Python `float` for an `int`, or an overflowing integer). - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. - **MPI + GPUs.** One GPU per MPI rank (`cunumpy.set_device(rank % n_gpus)`), and GPU-aware MPI for the marker exchange, so markers do not go through the host. diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index ab4838638..092ac8b57 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -1135,26 +1135,18 @@ def test_cuda_args_domain(mapping): assert isinstance(args, CudaDomainArguments) assert domain.args_domain is args # built once - kind_map, params, degree, t1, t2, t3, ind1, ind2, ind3, cx, cy, cz = args.values - assert int(kind_map) == domain.kind_map + assert args.kind_map == domain.kind_map # no copies of arrays that already have the right dtype and layout - assert t1 is domain.T[0] and ind3 is domain.indN[2] + assert args.t1 is domain.T[0] and args.ind3 is domain.indN[2] + # the struct holds the device addresses of these arrays + (struct,) = args.get_cuda_args() + assert struct["kind_map"] == domain.kind_map host = domain._pyccel_args_domain - for dev, ref in ( - (params, host.params), - (degree, host.degree), - (t1, host.t1), - (t2, host.t2), - (t3, host.t3), - (ind1, host.ind1), - (ind2, host.ind2), - (ind3, host.ind3), - (cx, host.cx), - (cy, host.cy), - (cz, host.cz), - ): - assert (cunumpy.to_numpy(dev) == ref).all() + for name in ("params", "degree", "t1", "t2", "t3", "ind1", "ind2", "ind3", "cx", "cy", "cz"): + dev = getattr(args, name) + assert struct[name] == dev.data.ptr, name + assert (cunumpy.to_numpy(dev) == getattr(host, name)).all(), name @requires_cupy @@ -1172,7 +1164,8 @@ def test_domain_deepcopy_and_pickle_on_cupy(mapping): assert (other.args_domain.params == domain.args_domain.params).all() other_cuda = other.args_domain assert other_cuda is not cuda_args - assert other_cuda.values[3] is other.T[0] + assert other_cuda.t1 is other.T[0] + assert other_cuda.get_cuda_args()[0]["t1"] == other.T[0].data.ptr if __name__ == "__main__": diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh new file mode 100644 index 000000000..02a32c5c2 --- /dev/null +++ b/src/struphy/kernel_arguments/pusher_args.cuh @@ -0,0 +1,57 @@ +// CUDA versions of the argument classes in pusher_args_kernels.py, one struct per class. +// +// CUDA kernels take these structs by value, in the place of the pyccel argument classes, e.g. +// +// #include "struphy/kernel_arguments/pusher_args.cuh" +// +// extern "C" __global__ +// void push_eta_stage(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, ...) +// +// The structs are filled on the host by the classes in struphy/utils/cuda_arguments.py, whose `fields` list the +// members below in the same order and with the same types; test_cuda_argument_structs checks that they agree. +// The member names are the attribute names of the pyccel classes. Pointers are device pointers. +#pragma once + +// CUDA version of MarkerArguments (struphy.utils.cuda_arguments.CudaMarkerArguments). +struct MarkerArgs { + double* markers; // (n_markers, n_cols), row-major + bool* valid_mks; // (n_markers,), true for markers that are neither holes nor ghosts + int n_markers; + int n_cols; + int Np; + int vdim; + int weight_idx; + int first_diagnostics_idx; + int first_init_idx; + int first_shift_idx; + int residual_idx; + int first_free_idx; + int mu_idx; + long long* bc_type; // (3,) +}; + +// CUDA version of DerhamArguments (struphy.utils.cuda_arguments.CudaDerhamArguments). +// The scratch arrays of the pyccel class (bn1, ..., bd3) are local arrays in the kernels. +struct DerhamArgs { + long long* pn; // (3,) + double* tn1; + double* tn2; + double* tn3; + long long* starts; // (3,) +}; + +// CUDA version of DomainArguments (struphy.utils.cuda_arguments.CudaDomainArguments). +struct DomainArgs { + int kind_map; + double* params; + long long* degree; // (3,) + double* t1; + double* t2; + double* t3; + long long* ind1; // (number of mapping grid cells, degree + 1) + long long* ind2; + long long* ind3; + double* cx; // control points + double* cy; + double* cz; +}; diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 276878275..215af37d4 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -6,16 +6,24 @@ import importlib import inspect +import re import sys +from pathlib import Path import cunumpy import numpy as np import pytest from cunumpy import PyccelKernel +import struphy from struphy.geometry.domains import Cuboid from struphy.kernel_arguments.pusher_args_kernels import DomainArguments, MarkerArguments -from struphy.utils.cuda_arguments import CudaDomainArguments, CudaMarkerArguments +from struphy.utils.cuda_arguments import ( + C_TYPES, + CudaDerhamArguments, + CudaDomainArguments, + CudaMarkerArguments, +) from struphy.utils.kernel_backends import CudaKernel, Kernel, KernelCatalog, is_cuda_backend requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") @@ -51,49 +59,79 @@ def push_eta_linear( markers[ip, 2] += dt * markers[ip, 5] -# Arguments: (dt, stage, CudaMarkerArguments, CudaDomainArguments), see struphy.utils.cuda_arguments. -CUDA_ARGS = r""" - double dt, int stage, - double* markers, bool* valid_mks, int n_markers, int n_cols, - int Np, int vdim, int weight_idx, int first_diagnostics_idx, int first_init_idx, - int first_shift_idx, int residual_idx, int first_free_idx, int mu_idx, long long* bc_type, - int kind_map, double* params, long long* degree, - double* t1, double* t2, double* t3, - long long* ind1, long long* ind2, long long* ind3, - double* cx, double* cy, double* cz -""" +# Same arguments as the pyccel kernel; the argument classes are the structs of pusher_args.cuh. +PUSH_ETA_LINEAR_SRC = r""" +#include "struphy/kernel_arguments/pusher_args.cuh" -PUSH_ETA_LINEAR_SRC = f""" extern "C" __global__ -void push_eta_linear({CUDA_ARGS}) -{{ +void push_eta_linear(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain) +{ int ip = blockDim.x * blockIdx.x + threadIdx.x; // only do something if particle is valid (i.e. not a hole or ghost) - if (ip >= n_markers || !valid_mks[ip]) return; + if (ip >= args_markers.n_markers || !args_markers.valid_mks[ip]) return; - double* mk = markers + (long long)ip * n_cols; + double* mk = args_markers.markers + (long long)ip * args_markers.n_cols; mk[0] += dt * mk[3]; mk[1] += dt * mk[4]; mk[2] += dt * mk[5]; -}} +} """ -# writes the scalar arguments into the markers, to check that they arrive with the right types -WRITE_SCALARS_SRC = f""" +# writes scalar arguments and struct members into the markers, to check that they arrive with the right types +WRITE_SCALARS_SRC = r""" +#include "struphy/kernel_arguments/pusher_args.cuh" + extern "C" __global__ -void write_scalars({CUDA_ARGS}) -{{ +void write_scalars(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain) +{ int ip = blockDim.x * blockIdx.x + threadIdx.x; - if (ip >= n_markers) return; + if (ip >= args_markers.n_markers) return; - double* mk = markers + (long long)ip * n_cols; + double* mk = args_markers.markers + (long long)ip * args_markers.n_cols; mk[0] = dt; mk[1] = stage; - mk[2] = n_cols; - mk[3] = first_init_idx; - mk[4] = mu_idx; - mk[5] = kind_map; + mk[2] = args_markers.n_cols; + mk[3] = args_markers.first_init_idx; + mk[4] = args_markers.mu_idx; + mk[5] = args_domain.kind_map; + mk[6] = args_markers.bc_type[2]; + mk[7] = args_domain.t3[1]; +} +""" + +STRUCT_CLASSES = [CudaMarkerArguments, CudaDerhamArguments, CudaDomainArguments] +HEADER = Path(struphy.__file__).parent / "kernel_arguments" / "pusher_args.cuh" + + +def header_structs() -> dict: + """The structs of pusher_args.cuh, as {name: ((C type, member), ...)} in declaration order.""" + text = re.sub(r"//[^\n]*", "", HEADER.read_text()) + structs = {} + for name, body in re.findall(r"struct\s+(\w+)\s*\{(.*?)\};", text, flags=re.S): + members = re.findall(r"([A-Za-z_][\w ]*?\s*\**)\s*(\w+)\s*;", body) + structs[name] = tuple( + (" ".join(ctype.replace("*", " *").split()).replace(" *", "*"), m) for ctype, m in members + ) + return structs + + +def layout_kernel_source(cls) -> str: + """CUDA kernel writing sizeof and (offsetof, sizeof) of each member of the struct of cls into an int64 array.""" + # NVRTC has no standard headers (no offsetof), so offsets are taken from a local struct + lines = [f"{cls.struct_name} s;", f"out[0] = sizeof({cls.struct_name});"] + for i, (_, name) in enumerate(cls.fields): + lines.append(f"out[{2 * i + 1}] = (char*)&s.{name} - (char*)&s;") + lines.append(f"out[{2 * i + 2}] = sizeof(s.{name});") + body = "\n ".join(lines) + return f""" +#include "struphy/kernel_arguments/pusher_args.cuh" + +extern "C" __global__ +void struct_layout(long long* out) +{{ + if (blockDim.x * blockIdx.x + threadIdx.x != 0) return; + {body} }} """ @@ -212,30 +250,100 @@ def test_cuda_kernel_updates_device_array_in_place(kernel): assert args_markers.markers is markers assert markers.data.ptr == ptr - assert args_markers.values[0] is markers + assert args_markers.get_cuda_args()[0]["markers"] == ptr @requires_cupy def test_cuda_scalar_arguments(): - """Python scalars (not cast) and the flattened argument classes arrive in the CUDA kernel correctly and in order.""" + """Python scalars (not cast) and the struct members arrive in the CUDA kernel correctly and in order.""" write_scalars = CudaKernel(WRITE_SCALARS_SRC, "write_scalars") with cunumpy.use_backend("cupy"): args_markers, args_domain = make_arguments(10) + args_markers.bc_type[2] = 7 # read through the pointer in the struct write_scalars(0.25, 3, args_markers, args_domain, n_threads=10) - row = cunumpy.to_numpy(args_markers.markers)[0, :6] + row = cunumpy.to_numpy(args_markers.markers)[0, :8] first_pusher_idx, mu_idx = MARKER_INDICES[3], MARKER_INDICES[7] - assert np.array_equal(row, [0.25, 3, N_COLS, first_pusher_idx, mu_idx, Cuboid().kind_map]) + t3 = cunumpy.to_numpy(args_domain.t3)[1] + assert np.array_equal(row, [0.25, 3, N_COLS, first_pusher_idx, mu_idx, Cuboid().kind_map, 7, t3]) @requires_cupy def test_cuda_domain_arguments_reference_domain_arrays(): + """The struct holds the device addresses of the domain's own arrays.""" with cunumpy.use_backend("cupy"): domain = Cuboid() args = domain.args_domain assert isinstance(args, CudaDomainArguments) - assert len(args.values) == 12 - assert args.values[3] is domain.T[0] and args.values[8] is domain.indN[2] and args.values[9] is domain.cx + (struct,) = args.get_cuda_args() + assert struct["kind_map"] == domain.kind_map + assert struct["t1"] == domain.T[0].data.ptr and struct["ind3"] == domain.indN[2].data.ptr + assert struct["cx"] == domain.cx.data.ptr + + +def test_cuda_argument_structs_match_header(): + """The fields of the CUDA argument classes are the members of the structs in pusher_args.cuh (names, C types, order).""" + structs = header_structs() + assert sorted(structs) == sorted(cls.struct_name for cls in STRUCT_CLASSES) + for cls in STRUCT_CLASSES: + assert structs[cls.struct_name] == cls.fields, cls.struct_name + assert all(ctype in C_TYPES for ctype, _ in cls.fields) + + +def test_cuda_struct_members_are_pyccel_attributes(): + """1:1 correspondence: the struct members are named like the attributes of the pyccel argument classes. + + Only ``n_cols`` is CUDA-specific: pyccel kernels take it from ``markers.shape[1]``. + """ + text = (Path(struphy.__file__).parent / "kernel_arguments" / "pusher_args_kernels.py").read_text() + for cls in STRUCT_CLASSES: + for _, name in cls.fields: + assert f"self.{name} =" in text or name == "n_cols", f"{cls.struct_name}.{name}" + + +@requires_cupy +@pytest.mark.parametrize("cls", STRUCT_CLASSES, ids=lambda cls: cls.struct_name) +def test_cuda_struct_layout(cls): + """The NumPy dtype of each struct has the memory layout NVRTC gives the C struct (size, offsets, member sizes).""" + with cunumpy.use_backend("cupy"): + out = cunumpy.zeros(2 * len(cls.fields) + 1, dtype=np.int64) + CudaKernel(layout_kernel_source(cls), "struct_layout")(out, n_threads=1) + layout = cunumpy.to_numpy(out) + dtype = cls.struct_dtype() + expected = [dtype.itemsize] + for _, name in cls.fields: + field_dtype, offset = dtype.fields[name][:2] + expected += [offset, field_dtype.itemsize] + assert layout.tolist() == expected + + +@requires_cupy +def test_cuda_struct_follows_copies(): + """Deepcopies and unpickled copies repack the struct with the addresses of their own arrays.""" + import copy + import pickle + + with cunumpy.use_backend("cupy"): + args_markers, _ = make_arguments(10) + for other in (copy.deepcopy(args_markers), pickle.loads(pickle.dumps(args_markers))): + assert other.markers is not args_markers.markers + (struct,) = other.get_cuda_args() + assert struct["markers"] == other.markers.data.ptr and struct["bc_type"] == other.bc_type.data.ptr + assert struct["mu_idx"] == args_markers.mu_idx + + +@requires_cupy +def test_cuda_struct_scalars_are_checked(): + """Scalars are checked when the struct is packed: no silent truncation or wrap-around in the kernel.""" + import cupy as cp + + markers, valid_mks, bc_type = cp.zeros((10, N_COLS)), cp.ones(10, dtype=bool), cp.zeros(3, dtype=int) + indices = list(MARKER_INDICES) + with pytest.raises(TypeError): + CudaMarkerArguments(markers, valid_mks, 10.0, *indices, bc_type) + with pytest.raises(OverflowError, match="Np"): + CudaMarkerArguments(markers, valid_mks, 2**31, *indices, bc_type) + CudaMarkerArguments(markers, valid_mks, np.int64(10), *indices, bc_type) # NumPy integers are fine @requires_cupy diff --git a/src/struphy/utils/cuda_arguments.py b/src/struphy/utils/cuda_arguments.py index 70c54edaa..15c9da450 100644 --- a/src/struphy/utils/cuda_arguments.py +++ b/src/struphy/utils/cuda_arguments.py @@ -2,22 +2,94 @@ The compiled classes in :mod:`struphy.kernel_arguments.pusher_args_kernels` only accept NumPy arrays. The classes here take the same constructor arguments, keep references to **CuPy** arrays -(no copies, other arrays raise) and flatten them into the arguments of a ``cupy.RawKernel``. -The order of :attr:`values` is the corresponding part of the CUDA kernel signature. +(no copies, other arrays raise) and pack them into one C struct each, which CUDA kernels take by value. +The structs are declared in ``struphy/kernel_arguments/pusher_args.cuh``; :attr:`Argument.fields` lists +their members in declaration order, with the C types. """ -from abc import ABC, abstractmethod +import operator +from abc import ABC import numpy as np +# NumPy types of the struct members, by C type; pointers are passed as device addresses +C_TYPES = { + "int": np.int32, + "double": np.float64, + "double*": np.uint64, + "bool*": np.uint64, + "long long*": np.uint64, +} + class Argument(ABC): - """Base class for objects that provide arguments to a CUDA kernel.""" + """Base class for objects that are passed to a CUDA kernel as one C struct. + + A subclass names the struct (:attr:`struct_name`), lists its members (:attr:`fields`) and stores each member + as an attribute of the same name, then calls :meth:`_pack` at the end of its constructor. The struct is packed + once; kernel calls pass it as it is. + """ + + struct_name: str + """Name of the C struct in ``pusher_args.cuh``.""" + + fields: tuple[tuple[str, str], ...] + """``(C type, name)`` of each struct member, in declaration order.""" + + @classmethod + def struct_dtype(cls) -> np.dtype: + """NumPy dtype with the memory layout of the C struct (C alignment and padding). + + Returns + ------- + numpy.dtype + Structured dtype, one field per struct member. + """ + return np.dtype( + {"names": [name for _, name in cls.fields], "formats": [C_TYPES[ctype] for ctype, _ in cls.fields]}, + align=True, + ) + + def _pack(self): + """Pack the members into the struct; pointer members hold the device address of the array attribute. + + Scalars are checked against the C type of their member: a non-integer for an ``int`` or a value that + does not fit raises, instead of arriving in the kernel truncated or wrapped around. + """ + struct = np.zeros((), dtype=self.struct_dtype()) + for ctype, name in self.fields: + value = getattr(self, name) + if ctype.endswith("*"): + struct[name] = value.data.ptr + elif ctype == "int": + value = operator.index(value) # raises TypeError for floats + info = np.iinfo(C_TYPES[ctype]) + if not info.min <= value <= info.max: + raise OverflowError(f"{name} = {value} does not fit into a C {ctype}.") + struct[name] = value + else: + struct[name] = float(value) + self._struct = struct[()] + + def __getstate__(self): + # the struct holds device addresses, which are not valid for the arrays of a copy + state = self.__dict__.copy() + state.pop("_struct", None) + return state + + def __setstate__(self, state): + self.__dict__.update(state) + self._pack() - @abstractmethod def get_cuda_args(self) -> tuple: - """Return this object's arguments in CUDA kernel signature order.""" - raise NotImplementedError + """Return this object's arguments in CUDA kernel signature order: the packed struct. + + Returns + ------- + tuple + One ``numpy.void`` with the bytes of the C struct. + """ + return (self._struct,) def _cupy_array(name: str, arr, dtype): @@ -49,9 +121,7 @@ def _cupy_array(name: str, arr, dtype): class CudaMarkerArguments(Argument): """CUDA version of :class:`~struphy.kernel_arguments.pusher_args_kernels.MarkerArguments`. - CUDA signature of :attr:`values`: ``double* markers, bool* valid_mks, int n_markers, int n_cols, int Np, - int vdim, int weight_idx, int first_diagnostics_idx, int first_init_idx, int first_shift_idx, - int residual_idx, int first_free_idx, int mu_idx, long long* bc_type`` + Passed to CUDA kernels as ``struct MarkerArgs``, see :attr:`fields`. Parameters ---------- @@ -101,11 +171,26 @@ class CudaMarkerArguments(Argument): n_markers : int Number of rows of ``markers``, e.g. for the number of CUDA threads. - - values : tuple - Flat CUDA kernel arguments, see the signature above. """ + struct_name = "MarkerArgs" + fields = ( + ("double*", "markers"), + ("bool*", "valid_mks"), + ("int", "n_markers"), + ("int", "n_cols"), + ("int", "Np"), + ("int", "vdim"), + ("int", "weight_idx"), + ("int", "first_diagnostics_idx"), + ("int", "first_init_idx"), + ("int", "first_shift_idx"), + ("int", "residual_idx"), + ("int", "first_free_idx"), + ("int", "mu_idx"), + ("long long*", "bc_type"), + ) + def __init__( self, markers, @@ -124,44 +209,25 @@ def __init__( self.markers = _cupy_array("markers", markers, np.float64) self.valid_mks = _cupy_array("valid_mks", valid_mks, np.bool_) self.n_markers = markers.shape[0] - self.n_cols = np.int32(markers.shape[1]) - self.Np = np.int32(Np) - self.vdim = np.int32(vdim) - self.weight_idx = np.int32(weight_idx) - self.first_diagnostics_idx = np.int32(first_diagnostics_idx) - self.first_pusher_idx = np.int32(first_pusher_idx) - self.first_shift_idx = np.int32(first_shift_idx) - self.residual_idx = np.int32(residual_idx) - self.first_free_idx = np.int32(first_free_idx) - self.mu_idx = np.int32(mu_idx) + self.n_cols = markers.shape[1] + self.Np = Np + self.vdim = vdim + self.weight_idx = weight_idx + self.first_diagnostics_idx = first_diagnostics_idx + self.first_init_idx = first_pusher_idx + self.first_shift_idx = first_shift_idx + self.residual_idx = residual_idx + self.first_free_idx = first_free_idx + self.mu_idx = mu_idx self.bc_type = _cupy_array("bc_type", bc_type, np.int64) - - def get_cuda_args(self) -> tuple: - return ( - self.markers, - self.valid_mks, - np.int32(self.n_markers), - self.n_cols, - self.Np, - self.vdim, - self.weight_idx, - self.first_diagnostics_idx, - self.first_pusher_idx, - self.first_shift_idx, - self.residual_idx, - self.first_free_idx, - self.mu_idx, - self.bc_type, - ) + self._pack() class CudaDerhamArguments(Argument): """CUDA version of :class:`~struphy.kernel_arguments.pusher_args_kernels.DerhamArguments`. - CUDA signature of :meth:`get_cuda_args`: ``long long* pn, double* tn1, double* tn2, double* tn3, long long* starts`` - - The scratch arrays of the pyccel class (``bn1``, ..., ``bd3``) are not part of it; CUDA kernels use - per-thread local arrays instead. + Passed to CUDA kernels as ``struct DerhamArgs``, see :attr:`fields`. The scratch arrays of the pyccel class + (``bn1``, ..., ``bd3``) are not part of it; CUDA kernels use per-thread local arrays instead. Parameters ---------- @@ -175,20 +241,26 @@ class CudaDerhamArguments(Argument): Start indices (current MPI process) of :class:`~struphy.feec.psydac_derham.Derham` (int64). """ + struct_name = "DerhamArgs" + fields = ( + ("long long*", "pn"), + ("double*", "tn1"), + ("double*", "tn2"), + ("double*", "tn3"), + ("long long*", "starts"), + ) + def __init__(self, pn, tn1, tn2, tn3, starts): self.pn = _cupy_array("pn", pn, np.int64) self.tn1, self.tn2, self.tn3 = (_cupy_array("tn", t, np.float64) for t in (tn1, tn2, tn3)) self.starts = _cupy_array("starts", starts, np.int64) - - def get_cuda_args(self) -> tuple: - return (self.pn, self.tn1, self.tn2, self.tn3, self.starts) + self._pack() class CudaDomainArguments(Argument): """CUDA version of :class:`~struphy.kernel_arguments.pusher_args_kernels.DomainArguments`. - CUDA signature of :attr:`values`: ``int kind_map, double* params, long long* degree, double* t1, - double* t2, double* t3, long long* ind1, long long* ind2, long long* ind3, double* cx, double* cy, double* cz`` + Passed to CUDA kernels as ``struct DomainArgs``, see :attr:`fields`. Parameters ---------- @@ -209,33 +281,29 @@ class CudaDomainArguments(Argument): cx, cy, cz : cupy.ndarray[float] Spline coefficients (control points) of the mapping. - - Attributes - ---------- - values : tuple - Flat CUDA kernel arguments, see the signature above. """ + struct_name = "DomainArgs" + fields = ( + ("int", "kind_map"), + ("double*", "params"), + ("long long*", "degree"), + ("double*", "t1"), + ("double*", "t2"), + ("double*", "t3"), + ("long long*", "ind1"), + ("long long*", "ind2"), + ("long long*", "ind3"), + ("double*", "cx"), + ("double*", "cy"), + ("double*", "cz"), + ) + def __init__(self, kind_map: int, params, degree, t1, t2, t3, ind1, ind2, ind3, cx, cy, cz): - self.kind_map = np.int32(kind_map) + self.kind_map = kind_map self.params = _cupy_array("params", params, np.float64) self.degree = _cupy_array("degree", degree, np.int64) self.t1, self.t2, self.t3 = (_cupy_array("t", t, np.float64) for t in (t1, t2, t3)) self.ind1, self.ind2, self.ind3 = (_cupy_array("ind", ind, np.int64) for ind in (ind1, ind2, ind3)) self.cx, self.cy, self.cz = (_cupy_array("c", c, np.float64) for c in (cx, cy, cz)) - - def get_cuda_args(self) -> tuple: - return ( - self.kind_map, - self.params, - self.degree, - self.t1, - self.t2, - self.t3, - self.ind1, - self.ind2, - self.ind3, - self.cx, - self.cy, - self.cz, - ) + self._pack() diff --git a/src/struphy/utils/kernel_backends.py b/src/struphy/utils/kernel_backends.py index c6a8816ef..6c22641b4 100644 --- a/src/struphy/utils/kernel_backends.py +++ b/src/struphy/utils/kernel_backends.py @@ -6,6 +6,9 @@ Both kernels are called with the same arguments; the CUDA kernel takes the CUDA versions of the argument classes (:mod:`struphy.utils.cuda_arguments`), which reference arrays on the device, and the number of threads ``n_threads``. No arrays are converted or copied at call time. + +CUDA sources can include struphy headers relative to the parent folder of the ``struphy`` package, +e.g. ``#include "struphy/kernel_arguments/pusher_args.cuh"`` for the argument structs. """ import importlib @@ -17,6 +20,9 @@ from struphy.utils.cuda_arguments import Argument +INCLUDE_DIR = Path(__file__).resolve().parents[2] +"""NVRTC include path of the CUDA kernels: the folder that contains the ``struphy`` package.""" + def is_cuda_backend() -> bool: """Whether the active cunumpy backend is CuPy. @@ -32,8 +38,9 @@ def is_cuda_backend() -> bool: class CudaKernel: """A ``cupy.RawKernel``, the CUDA counterpart of a pyccel kernel. - The kernel is compiled on the first call. At each call, the CUDA argument classes are replaced by their - ``values``; all other arguments are passed to the ``cupy.RawKernel`` as they are. Arrays must be CuPy + The kernel is compiled on the first call, with :data:`INCLUDE_DIR` on the include path. At each call, the + CUDA argument classes are replaced by their structs; all other arguments are passed to the ``cupy.RawKernel`` + as they are. Arrays must be CuPy arrays (``cupy`` raises otherwise); they are never converted or copied. Python ``int`` and ``float`` arrive correctly in ``int`` and ``double`` parameters; scalars are not checked against the kernel signature. @@ -96,7 +103,7 @@ def __call__(self, *args, n_threads: int): if self._raw_kernel is None: import cupy as cp - self._raw_kernel = cp.RawKernel(self._source, self.name) + self._raw_kernel = cp.RawKernel(self._source, self.name, options=(f"-I{INCLUDE_DIR}",)) values = [] for arg in args: