From 4121d302fc0fef8f02a62bbf6021b0520afd9ef2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 30 Sep 2026 22:29:47 +0200 Subject: [PATCH 01/16] Added CudaKernel, Kernel and KernelCatalog for host/CUDA kernels --- .github/workflows/testing.yml | 2 +- CHANGELOG.md | 14 + README.md | 46 +++ docs/source/api.md | 158 +++++++++ pyproject.toml | 5 +- src/cunumpy/__init__.py | 8 + src/cunumpy/__init__.pyi | 6 + src/cunumpy/cuda_kernel.py | 515 +++++++++++++++++++++++++++++ src/cunumpy/dispatch.py | 282 ++++++++++++++++ tests/unit/test_cuda_kernel.py | 309 +++++++++++++++++ tests/unit/test_kernel_dispatch.py | 172 ++++++++++ 11 files changed, 1513 insertions(+), 4 deletions(-) create mode 100644 src/cunumpy/cuda_kernel.py create mode 100644 src/cunumpy/dispatch.py create mode 100644 tests/unit/test_cuda_kernel.py create mode 100644 tests/unit/test_kernel_dispatch.py diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index aa01052..4188463 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -17,7 +17,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.8", "3.10", "3.13"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: # Checkout the repository diff --git a/CHANGELOG.md b/CHANGELOG.md index 3093847..36df17a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Removed +- Support for Python 3.8 and 3.9 (both end-of-life); `cunumpy` now requires Python 3.10 or newer. + +### Changed +- Python 3.14 is supported. +- CI now tests every supported Python version (3.10, 3.11, 3.12, 3.13 and 3.14) instead of 3.8/3.10/3.13. + +### Added +- `xp.CudaKernel`: Wraps a CUDA C kernel (`cupy.RawKernel`, compiled lazily with NVRTC) so it can be called with the same arguments as the host kernel it mirrors, plus `n_threads`. The `extern "C" __global__` signature is parsed once and every call is checked against it: argument count, array dtypes (host arrays raise, they are never copied), and scalars (Python scalars are cast to the declared C types with range checks; lossy or mismatching scalars raise instead of reaching the kernel as silently wrong values). Supports `block_size`, NVRTC `options`, `include_dirs`, `shared_mem`, `stream`, `CudaKernel.from_file()` (`_cuda.cu`), `compile()` and `prepare_args()`; `check_signature=False` skips the checks. +- `xp.CudaArguments`: Base class for argument objects that are flattened into several CUDA kernel arguments; any object with a `__cuda_args__()` method is flattened. +- `xp.parse_cuda_signature(source, name)` and `xp.CudaParameter`: Parse the parameters of a `__global__` function. +- `xp.Kernel`: A host kernel (`PyccelKernel`) and its CUDA counterpart, calling the one matching the active backend. Without a CUDA kernel on the CuPy backend it raises `NotImplementedError` (`missing_cuda="raise"`, default) or falls back to the host kernel with host copies (`missing_cuda="fallback"`). +- `xp.KernelCatalog`: Read-only mapping of `Kernel`s; `KernelCatalog.from_package()` collects them from a package with one folder per kernel (`name/name_kernels.py`, `name/name_cuda.cu`); `without_cuda` lists the kernels still to port. + ## [0.2.0] - 2026-09-28 ### Changed diff --git a/README.md b/README.md index 3df254f..d133519 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,52 @@ The wrapper can also traverse arrays nested in lists, tuples, dictionaries, and selected application objects; see the full [API reference](docs/source/api.md) for `object_modules`, `is_array`, aliasing, and output declarations. +## Write CUDA kernels next to host kernels + +`CudaKernel` wraps a CUDA C kernel (compiled with NVRTC through +`cupy.RawKernel`) so that it is called with the same arguments as the host +kernel it mirrors, plus the number of threads. Arrays are never copied: they +must be CuPy arrays. The `extern "C" __global__` signature is parsed once and +every call is checked against it: Python scalars are cast to the declared C +types, and a wrong argument count, an array of the wrong dtype, or a scalar +that does not fit its type raises instead of silently producing wrong values. + +`Kernel` pairs a host kernel with its CUDA kernel and calls the one matching +the active backend, so kernels can be ported to CUDA one at a time: + +```python +import cunumpy as xp + +AXPY = r""" +extern "C" __global__ +void axpy(double a, const double* x, double* y, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) y[i] += a * x[i]; +} +""" + + +def axpy(a, x, y, n): # host version, e.g. compiled with Pyccel + for i in range(n): + y[i] += a * x[i] + + +kernel = xp.Kernel(axpy, xp.CudaKernel(AXPY, "axpy")) + +with xp.use_backend("cupy"): + x = xp.arange(1000, dtype=xp.float64) + y = xp.zeros(1000) + kernel(2.0, x, y, 1000, n_threads=1000) # runs the CUDA kernel +``` + +On the CuPy backend, a `Kernel` without CUDA kernel raises +`NotImplementedError` (or, with `missing_cuda="fallback"`, runs the host kernel +through `PyccelKernel`, with host copies). Objects implementing +`__cuda_args__()` (see `CudaArguments`) are flattened into several kernel +arguments. `KernelCatalog.from_package()` collects kernel pairs from a package +with one folder per kernel (`name/name_kernels.py` and `name/name_cuda.cu`). +See the [API reference](docs/source/api.md) for details. + ## Pyodide CuNumpy supports the NumPy backend in Pyodide. It does not provide CuPy/CUDA diff --git a/docs/source/api.md b/docs/source/api.md index f7d63e3..2cf90f2 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -353,6 +353,164 @@ lists) are converted back using `is_array`; dictionaries in return values are not recursively converted. On the NumPy path, the original return value and normal Python mutation and exception behavior are preserved. +## `CudaKernel` + +### Constructor + +```python +xp.CudaKernel( + source, + name, + *, + block_size=128, + options=(), + include_dirs=(), + check_signature=True, +) +xp.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs) +``` + +Wraps the `extern "C" __global__` function `name` in the CUDA C `source`. The +kernel is compiled with NVRTC through `cupy.RawKernel` on the first call (or +by `compile()`), and cached. CuPy is imported only then, so kernels can be +created and their signatures parsed without CuPy. + +`from_file` reads the source from a file; the kernel name defaults to the file +name without `suffix` (`axpy_cuda.cu` -> `axpy`), and the directory of the file +is added to the include directories. + +### Parameters + +* `block_size`: threads per block. +* `options`: additional NVRTC options, e.g. `("-std=c++17",)`. +* `include_dirs`: directories for `#include`, passed as `-I`. +* `check_signature`: parse the signature and check every call against it + (default). Raises `ValueError` if the signature cannot be parsed, e.g. with + templates, macros or pointers to pointers in the parameter list; pass + `False` to launch with the arguments as they are, like `cupy.RawKernel`. + +### Calling + +```python +kernel(*args, n_threads, shared_mem=0, stream=None) +``` + +Launches `ceil(n_threads / block_size)` blocks of `block_size` threads on +`stream` (the current stream if `None`); nothing is launched for +`n_threads=0`. The arguments are prepared by `kernel.prepare_args(*args)`: + +* arguments with a `__cuda_args__()` method are replaced by the values it + returns (see `CudaArguments`); +* with a checked signature, the number of arguments must match, and + * pointer parameters take CuPy arrays whose dtype matches the pointed-to type + (any dtype for `void*`); host arrays raise `TypeError`, they are never + copied to the device; + * Python scalars are cast to the declared type: `int` into integer (with a + range check, `OverflowError`), floating-point and complex parameters, + `float` into floating-point and complex parameters, `bool` into boolean + and integer parameters; anything else raises `TypeError`; + * NumPy scalars are passed as they are if their dtype matches, cast if the + cast is safe (e.g. `np.float32` into `double`), and raise `TypeError` + otherwise (e.g. `np.float64` into `float`). + +This matters because `cupy.RawKernel` reads each argument with the size +declared in the signature and does not check types: an integer passed to a +`double` parameter, or a `double` passed to a `float` parameter, arrives as a +wrong value without an error. The checks cost about 0.3 µs per argument (about +10 µs for a kernel with 29 arguments, measured on an H100 node, where the launch +itself costs about as much), which is negligible for kernels that run for +100 µs or more. For very short kernels called in a hot loop, pass +`check_signature=False` once the calls are known to be correct. + +C types are mapped to NumPy dtypes as on Linux (LP64): `int` is `int32`, +`long` and `long long` are `int64`, `float` is `float32`, `double` is +`float64`, `complex` is `complex128`; fixed-width types such as +`int64_t` and `size_t` are supported too. `xp.parse_cuda_signature(source, +name)` returns the parsed parameters (`CudaParameter` tuples of `name`, +`ctype`, `dtype`, `pointer`). + +## `CudaArguments` + +```python +class Particles(xp.CudaArguments): + def __init__(self, positions, velocities): + self.positions = positions + super().__init__(positions, velocities, positions.shape[0]) + +kernel(dt, Particles(x, v), n_threads=x.shape[0]) +``` + +Base class for objects passed to a `CudaKernel` as one argument that stands +for several kernel parameters. `CudaArguments(*values)` stores the values; +`__cuda_args__()` returns them. Subclassing is optional: any object with a +`__cuda_args__()` method returning a tuple is flattened. This lets an +application keep its host argument objects (e.g. Pyccel classes holding NumPy +arrays) and matching device argument objects that reference the same data on +the device, and pass either to the same call. + +## `Kernel` + +```python +xp.Kernel( + host_kernel, + cuda_kernel=None, + *, + name=None, + missing_cuda="raise", + cuda_path=None, +) +``` + +A host kernel (a `PyccelKernel`; other callables are wrapped in one) and its +CUDA counterpart. `kernel.get_kernel()` returns the host kernel on the NumPy +backend and the CUDA kernel on the CuPy backend; call it once at setup to fail +early if a CUDA kernel is missing. `kernel(*args, n_threads=None)` calls the +kernel of the active backend; `n_threads` is required for the CUDA kernel and +ignored by the host kernel. + +Without a CUDA kernel on the CuPy backend, `missing_cuda="raise"` raises +`NotImplementedError` (naming `cuda_path`, if given), and +`missing_cuda="fallback"` calls the host kernel through `PyccelKernel`, which +copies the arrays to the host and back at every call (a `RuntimeWarning` is +emitted once). + +Properties: `name`, `host_kernel`, `cuda_kernel`, `has_cuda`, `missing_cuda`, +`cuda_path`. + +## `KernelCatalog` + +```python +catalog = xp.KernelCatalog.from_package( + package, + *, + host_suffix="_kernels", + cuda_suffix="_cuda.cu", + missing_cuda="raise", + **cuda_options, +) +kernel = catalog["push"] +``` + +A read-only mapping from names to `Kernel` objects. `from_package` scans the +subfolders of `package`: for every folder `` containing the module +`.py`, the function `` of that module is the host +kernel, and `` in the same folder, if present, is the CUDA +kernel (`__global__` function ``). `cuda_options` are passed on to +`CudaKernel.from_file`. Typically called in the package's `__init__.py`: + +```text +my_kernels/ +├── __init__.py # catalog = xp.KernelCatalog.from_package(__name__) +├── push/ +│ ├── push_kernels.py # def push(...): ... +│ └── push_cuda.cu # __global__ void push(...) +└── deposit/ + └── deposit_kernels.py # no CUDA kernel yet +``` + +`catalog.without_cuda` lists the kernels still to port. `KernelCatalog(kernels)` +and `catalog.register(kernel, name=None)` build a catalog by hand. + ## Version `xp.__version__` is the installed package version. When package metadata is diff --git a/pyproject.toml b/pyproject.toml index 794f7b3..411700d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,16 +11,15 @@ readme = "README.md" keywords = [ "python" ] license = { file = "LICENSE.txt" } authors = [ { name = "Max" } ] -requires-python = ">=3.8" +requires-python = ">=3.10" classifiers = [ "Development Status :: 3 - Alpha", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", ] dependencies = [ "array-api-compat", diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 4e9e194..8ac8dff 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -2,6 +2,8 @@ from importlib.metadata import PackageNotFoundError, version from . import xp +from .cuda_kernel import CudaArguments, CudaKernel, CudaParameter, parse_cuda_signature +from .dispatch import Kernel, KernelCatalog from .kernel import PyccelKernel from .xp import ( assert_same_backend, @@ -35,6 +37,11 @@ __version__ = "0.0.0+unknown" __all__ = [ + "CudaArguments", + "CudaKernel", + "CudaParameter", + "Kernel", + "KernelCatalog", "PyccelKernel", "__version__", "assert_same_backend", @@ -51,6 +58,7 @@ "is_gpu", "memory_info", "numpy_backend", + "parse_cuda_signature", "pin_memory", "same_backend", "set_backend", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index d71893d..6cc2098 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -8,6 +8,12 @@ import numpy as np from numpy import * from . import xp as xp +from .cuda_kernel import CudaArguments as CudaArguments +from .cuda_kernel import CudaKernel as CudaKernel +from .cuda_kernel import CudaParameter as CudaParameter +from .cuda_kernel import parse_cuda_signature as parse_cuda_signature +from .dispatch import Kernel as Kernel +from .dispatch import KernelCatalog as KernelCatalog from .kernel import PyccelKernel as PyccelKernel def to_numpy(array: Any) -> np.ndarray: ... diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py new file mode 100644 index 0000000..c4abf53 --- /dev/null +++ b/src/cunumpy/cuda_kernel.py @@ -0,0 +1,515 @@ +"""CUDA kernels (``cupy.RawKernel``) called like their NumPy/Pyccel counterparts. + +:class:`CudaKernel` wraps a CUDA C kernel so that it can be called with the same +arguments as the host kernel it mirrors: + +* argument objects that implement the :class:`CudaArguments` protocol + (a ``__cuda_args__()`` method) are flattened into their device arrays and + scalars, so an object holding several arrays can be passed as one argument; +* the ``extern "C" __global__`` signature is parsed once, and every call is + checked against it: the number of arguments, the dtype of every array, and + every scalar. Python scalars are cast to the declared C type; a scalar that + does not fit the declared type (a ``float`` for an ``int``, an integer out of + range, a NumPy scalar that would lose precision) raises instead of reaching + the kernel as a silently wrong value, which is what ``cupy.RawKernel`` would + do; +* arrays are never converted or copied: they must already be CuPy arrays. + +The number of threads is given at each call (``n_threads``); the kernel is +launched on ``ceil(n_threads / block_size)`` blocks. + +This module imports CuPy only when a kernel is compiled, so it can be imported +(and signatures parsed) without CuPy. +""" + +from __future__ import annotations + +import math +import re +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Callable, NamedTuple, Sequence + +import numpy as np + +__all__ = [ + "CudaArguments", + "CudaKernel", + "CudaParameter", + "parse_cuda_signature", +] + + +class CudaArguments: + """Base class for objects passed to a :class:`CudaKernel` as one argument. + + A :class:`CudaKernel` replaces every argument that has a ``__cuda_args__()`` + method by the values it returns, in order. Subclassing this class is + optional: any object implementing ``__cuda_args__()`` is flattened. + + Parameters + ---------- + *values + The CUDA kernel arguments this object stands for: CuPy arrays and + scalars, in the order of the kernel signature. + + Examples + -------- + >>> class Particles(CudaArguments): + ... def __init__(self, positions, velocities): + ... self.positions = positions + ... super().__init__(positions, velocities, positions.shape[0]) + >>> kernel(dt, Particles(x, v), n_threads=x.shape[0]) # doctest: +SKIP + """ + + def __init__(self, *values: Any) -> None: + self._cuda_args = tuple(values) + + def __cuda_args__(self) -> tuple[Any, ...]: + """The CUDA kernel arguments this object stands for.""" + return self._cuda_args + + +class CudaParameter(NamedTuple): + """One parameter of a CUDA kernel signature. + + Attributes + ---------- + name : str + Parameter name. + ctype : str + Normalized C type without qualifiers or ``*``, e.g. ``"double"``. + dtype : numpy.dtype | None + NumPy dtype of the value (or of the pointed-to elements); ``None`` for + ``void*``. + pointer : bool + Whether the parameter is a pointer (a device array). + """ + + name: str + ctype: str + dtype: np.dtype | None + pointer: bool + + +# C types (after removing qualifiers) and their NumPy dtypes. ``long`` is 64 bit, +# as on Linux (LP64), the platform CUDA runs on in practice. +_CTYPES = { + "bool": np.bool_, + "char": np.int8, + "signed char": np.int8, + "unsigned char": np.uint8, + "short": np.int16, + "short int": np.int16, + "unsigned short": np.uint16, + "unsigned short int": np.uint16, + "int": np.int32, + "signed": np.int32, + "signed int": np.int32, + "unsigned": np.uint32, + "unsigned int": np.uint32, + "long": np.int64, + "long int": np.int64, + "long long": np.int64, + "long long int": np.int64, + "unsigned long": np.uint64, + "unsigned long int": np.uint64, + "unsigned long long": np.uint64, + "unsigned long long int": np.uint64, + "int8_t": np.int8, + "int16_t": np.int16, + "int32_t": np.int32, + "int64_t": np.int64, + "uint8_t": np.uint8, + "uint16_t": np.uint16, + "uint32_t": np.uint32, + "uint64_t": np.uint64, + "size_t": np.uint64, + "ptrdiff_t": np.int64, + "ssize_t": np.int64, + "float": np.float32, + "double": np.float64, + "complex": np.complex64, + "complex": np.complex128, +} + +_QUALIFIERS = {"const", "volatile", "__restrict__", "__restrict", "restrict"} + +_COMPLEX = re.compile(r"(?:(?:thrust|cuda::std)::)?complex\s*<\s*(float|double)\s*>") +_TOKEN = re.compile(r"complex<(?:float|double)>|[A-Za-z_]\w*|\*|\[\s*\]") + + +def _strip_comments(source: str) -> str: + source = re.sub(r"/\*.*?\*/", " ", source, flags=re.DOTALL) + return re.sub(r"//[^\n]*", " ", source) + + +def _parse_parameter(text: str) -> CudaParameter: + text = _COMPLEX.sub(lambda m: f"complex<{m.group(1)}>", text) + tokens = _TOKEN.findall(text) + pointers = sum(1 for t in tokens if t == "*" or t.startswith("[")) + words = [t for t in tokens if t != "*" and not t.startswith("[")] + words = [t for t in words if t not in _QUALIFIERS] + if len(words) < 2: + raise ValueError(f"cannot parse the kernel parameter {text.strip()!r}") + name, ctype = words[-1], " ".join(words[:-1]) + + if ctype == "void" and pointers == 1: + return CudaParameter(name, ctype, None, True) + if pointers > 1 or ctype not in _CTYPES: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: unsupported type " + f"{ctype + '*' * pointers!r}" + ) + return CudaParameter(name, ctype, np.dtype(_CTYPES[ctype]), pointers == 1) + + +def parse_cuda_signature(source: str, name: str) -> tuple[CudaParameter, ...]: + """Parse the parameters of the ``__global__`` function `name` in `source`. + + Parameters + ---------- + source : str + CUDA C source code. + name : str + Name of the ``__global__`` function. + + Returns + ------- + tuple[CudaParameter, ...] + The parameters, in order. + + Raises + ------ + ValueError + If there is no such function, or a parameter has a type that cannot be + checked (e.g. a template parameter, a macro or a pointer to pointer). + """ + code = _strip_comments(source) + match = re.search(r"__global__\s+void\s+" + re.escape(name) + r"\s*\(", code) + if match is None: + raise ValueError(f"no __global__ function {name!r} found in the CUDA source") + + depth, start = 1, match.end() + for pos in range(start, len(code)): + if code[pos] == "(": + depth += 1 + elif code[pos] == ")": + depth -= 1 + if depth == 0: + break + else: + raise ValueError(f"unbalanced parentheses in the signature of {name!r}") + + params = code[start:pos].strip() + if params in ("", "void"): + return () + return tuple(_parse_parameter(p) for p in params.split(",")) + + +def _flatten(args: Sequence[Any]) -> list[Any]: + values: list[Any] = [] + for arg in args: + cuda_args = getattr(arg, "__cuda_args__", None) + if cuda_args is not None: + values.extend(cuda_args()) + else: + values.append(arg) + return values + + +def _describe(param: CudaParameter, index: int) -> str: + ctype = param.ctype + ("*" if param.pointer else "") + return f"argument {index} ({ctype} {param.name})" + + +def _pointer_checker(param: CudaParameter, index: int) -> Callable[[Any], Any]: + """Checker for a pointer parameter: a device array with the right dtype.""" + dtype = param.dtype + + def check(value: Any) -> Any: + # checked on the class: on the instance, CuPy builds the whole interface dict + if not hasattr(type(value), "__cuda_array_interface__") and not hasattr( + value, "__cuda_array_interface__" + ): + raise TypeError( + f"{_describe(param, index)} must be a CuPy array, got " + f"{type(value).__name__}; arrays are never copied to the device" + ) + if dtype is not None and value.dtype != dtype: + raise TypeError( + f"{_describe(param, index)} must have dtype {dtype}, got {value.dtype}" + ) + return value + + return check + + +def _scalar_checker(param: CudaParameter, index: int) -> Callable[[Any], Any]: + """Checker for a scalar parameter: casts to the declared type, or raises.""" + dtype = param.dtype + kind = dtype.kind + scalar_type = dtype.type + if kind in "iu": + info = np.iinfo(dtype) + low, high = int(info.min), int(info.max) + + def cast_int(value: int) -> Any: + if kind in "iu": + if not low <= value <= high: + raise OverflowError( + f"{_describe(param, index)}: {value} is out of range " + f"[{low}, {high}]" + ) + return scalar_type(value) + if kind in "fc": + return scalar_type(value) + raise TypeError( + f"{_describe(param, index)} cannot take a value of type " + f"{type(value).__name__}" + ) + + def check(value: Any) -> Any: + value_type = type(value) + # fast paths for the common cases + if value_type is scalar_type: + return value + if value_type is int: + return cast_int(value) + if value_type is float and kind in "fc": + return scalar_type(value) + + if isinstance(value, np.generic): + if value.dtype == dtype: + return value + if np.can_cast(value.dtype, dtype, casting="safe"): + return scalar_type(value) + raise TypeError( + f"{_describe(param, index)} cannot take a {value.dtype} scalar " + f"without losing information" + ) + if isinstance(value, bool): + if kind in "biu": + return scalar_type(value) + elif isinstance(value, int): + return cast_int(int(value)) + elif isinstance(value, float): + if kind in "fc": + return scalar_type(value) + elif isinstance(value, complex): + if kind == "c": + return scalar_type(value) + raise TypeError( + f"{_describe(param, index)} cannot take a value of type " + f"{value_type.__name__}" + ) + + return check + + +class CudaKernel: + """A CUDA C kernel, compiled with NVRTC through ``cupy.RawKernel``. + + Parameters + ---------- + source : str + CUDA C source code containing the ``extern "C" __global__`` function + `name`. + name : str + Name of the kernel function in `source`. + block_size : int + Number of threads per block. + options : Sequence[str] + Additional NVRTC compiler options, e.g. ``("-std=c++17",)``. + include_dirs : Sequence[str | Path] + Directories searched for ``#include`` files (passed as ``-I``). + check_signature : bool + Parse the kernel signature and check (and cast) every call against it. + Raises ``ValueError`` at construction if the signature cannot be parsed + (e.g. templates or macros in the parameter list); pass False to launch + with the arguments as they are, like ``cupy.RawKernel``. + + Examples + -------- + >>> axpy = CudaKernel(r''' + ... extern "C" __global__ + ... void axpy(double a, const double* x, double* y, int n) { + ... int i = blockDim.x * blockIdx.x + threadIdx.x; + ... if (i < n) y[i] += a * x[i]; + ... }''', "axpy") + >>> axpy(2.0, x, y, x.size, n_threads=x.size) # doctest: +SKIP + """ + + def __init__( + self, + source: str, + name: str, + *, + block_size: int = 128, + options: Sequence[str] = (), + include_dirs: Sequence[str | Path] = (), + check_signature: bool = True, + ) -> None: + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + self._source = source + self._name = name + self._block_size = int(block_size) + self._options = tuple(options) + tuple(f"-I{d}" for d in include_dirs) + self._signature = ( + parse_cuda_signature(source, name) if check_signature else None + ) + # one checker per parameter, built once so that calls stay cheap + self._checkers = ( + None + if self._signature is None + else [ + _pointer_checker(p, i) if p.pointer else _scalar_checker(p, i) + for i, p in enumerate(self._signature) + ] + ) + self._raw_kernel = None + + @classmethod + def from_file( + cls, + path: str | Path, + name: str | None = None, + *, + suffix: str = "_cuda.cu", + **kwargs: Any, + ) -> CudaKernel: + """Load the CUDA source from a file. + + Parameters + ---------- + path : str | Path + Path of the CUDA source file. + name : str | None + Name of the kernel function; by default the file name without + `suffix` (``axpy_cuda.cu`` -> ``axpy``). + suffix : str + File name suffix stripped to get the default kernel name. + **kwargs + Passed on to :class:`CudaKernel`. The directory of the file is + always added to ``include_dirs``. + """ + path = Path(path) + if name is None: + if not path.name.endswith(suffix): + raise ValueError( + f"{path.name} does not end with {suffix!r}; pass the kernel name" + ) + name = path.name[: -len(suffix)] + include_dirs = (path.parent, *kwargs.pop("include_dirs", ())) + return cls(path.read_text(), name, include_dirs=include_dirs, **kwargs) + + def __repr__(self) -> str: + return f"CudaKernel(name={self._name!r}, block_size={self._block_size})" + + @property + def name(self) -> str: + """Name of the kernel function.""" + return self._name + + @property + def source(self) -> str: + """CUDA C source code.""" + return self._source + + @property + def block_size(self) -> int: + """Number of threads per block.""" + return self._block_size + + @property + def options(self) -> tuple[str, ...]: + """NVRTC compiler options, including ``-I`` include directories.""" + return self._options + + @property + def signature(self) -> tuple[CudaParameter, ...] | None: + """The parsed kernel parameters, or None if calls are not checked.""" + return self._signature + + def compile(self) -> Any: + """Compile the kernel now (it is otherwise compiled on the first call). + + Returns + ------- + cupy.RawKernel + The compiled kernel; compiled once and cached (also on disk by CuPy). + """ + if self._raw_kernel is None: + from .xp import cupy_available + + if not cupy_available(): + raise RuntimeError( + f"cannot compile CUDA kernel {self._name!r}: " + "CuPy is not installed or no GPU is available" + ) + import cupy as cp + + self._raw_kernel = cp.RawKernel( + self._source, self._name, options=self._options + ) + return self._raw_kernel + + def prepare_args(self, *args: Any) -> tuple[Any, ...]: + """The arguments as passed to ``cupy.RawKernel``: flattened and checked. + + Argument objects with ``__cuda_args__()`` are flattened. If the signature + is checked, the number of arguments, the dtype of every array and every + scalar are checked, and Python scalars are cast to the declared C types. + + Raises + ------ + TypeError + Wrong number of arguments, a host array or an array of the wrong + dtype for a pointer parameter, or a scalar of an incompatible type. + OverflowError + A Python integer out of range of the declared integer type. + """ + values = _flatten(args) + if self._signature is None: + return tuple(values) + + if len(values) != len(self._signature): + raise TypeError( + f"{self._name}() takes {len(self._signature)} arguments after " + f"flattening argument objects, got {len(values)}" + ) + return tuple([check(v) for check, v in zip(self._checkers, values)]) + + def __call__( + self, + *args: Any, + n_threads: int, + shared_mem: int = 0, + stream: Any = None, + ) -> None: + """Launch the kernel. + + Parameters + ---------- + *args + Kernel arguments: CuPy arrays, scalars and argument objects with + ``__cuda_args__()``, see :meth:`prepare_args`. + n_threads : int + Number of threads to launch, rounded up to a multiple of the block + size; nothing is launched for 0. + shared_mem : int + Dynamic shared memory per block, in bytes. + stream : cupy.cuda.Stream | None + Stream to launch on; the current stream if None. + """ + if n_threads < 0: + raise ValueError(f"n_threads must be non-negative, got {n_threads}") + values = self.prepare_args(*args) + if n_threads == 0: + return + + kernel = self.compile() + grid = (math.ceil(n_threads / self._block_size),) + with stream if stream is not None else nullcontext(): + kernel(grid, (self._block_size,), values, shared_mem=shared_mem) diff --git a/src/cunumpy/dispatch.py b/src/cunumpy/dispatch.py new file mode 100644 index 0000000..942487d --- /dev/null +++ b/src/cunumpy/dispatch.py @@ -0,0 +1,282 @@ +"""Pairs of host and CUDA kernels, chosen by the active backend. + +A :class:`Kernel` holds a host kernel (a :class:`~cunumpy.PyccelKernel`, e.g. a +Pyccel-compiled function) and, optionally, its 1:1 corresponding CUDA kernel +(:class:`~cunumpy.CudaKernel`). It calls the host kernel on the NumPy backend and +the CUDA kernel on the CuPy backend, so a code base can port its kernels to CUDA +one by one. + +:class:`KernelCatalog` collects such pairs from a package laid out with one +folder per kernel:: + + my_kernels/ + ├── __init__.py # catalog = KernelCatalog.from_package(__name__) + ├── push/ + │ ├── push_kernels.py # def push(...): ... (host kernel) + │ └── push_cuda.cu # __global__ void push(...) (CUDA kernel) + └── deposit/ + └── deposit_kernels.py # not ported yet +""" + +from __future__ import annotations + +import importlib +import warnings +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Callable, Iterator + +from .cuda_kernel import CudaKernel +from .kernel import PyccelKernel +from .xp import get_backend + +__all__ = ["Kernel", "KernelCatalog"] + +_MISSING_CUDA = ("raise", "fallback") + + +class Kernel: + """A host kernel and its CUDA counterpart; calls the one matching the backend. + + Parameters + ---------- + host_kernel : PyccelKernel | callable + The host kernel, called on the NumPy backend. A plain callable is + wrapped in a :class:`~cunumpy.PyccelKernel`. + cuda_kernel : CudaKernel | None + The CUDA kernel, called on the CuPy backend; None if it has not been + written yet. + name : str | None + Name of the kernel; defaults to the name of the host kernel. + missing_cuda : {"raise", "fallback"} + What happens on the CuPy backend if there is no CUDA kernel: ``"raise"`` + raises ``NotImplementedError``; ``"fallback"`` calls the host kernel + through :class:`~cunumpy.PyccelKernel`, which copies the arrays to the + host and back at every call (a warning is emitted once). + cuda_path : str | Path | None + Where the CUDA kernel is expected, for the error message if it is missing. + + Notes + ----- + Both kernels take the same arguments, except that the CUDA kernel gets the + number of threads, ``n_threads``, and argument objects in their CUDA form + (see :class:`~cunumpy.CudaArguments`). + """ + + def __init__( + self, + host_kernel: PyccelKernel | Callable[..., Any], + cuda_kernel: CudaKernel | None = None, + *, + name: str | None = None, + missing_cuda: str = "raise", + cuda_path: str | Path | None = None, + ) -> None: + if missing_cuda not in _MISSING_CUDA: + raise ValueError( + f"missing_cuda must be one of {_MISSING_CUDA}, got {missing_cuda!r}" + ) + if cuda_kernel is not None and not isinstance(cuda_kernel, CudaKernel): + raise TypeError( + "cuda_kernel must be a CudaKernel or None, " + f"got {type(cuda_kernel).__name__}" + ) + if not isinstance(host_kernel, PyccelKernel): + host_kernel = PyccelKernel(host_kernel) + self._host_kernel = host_kernel + self._cuda_kernel = cuda_kernel + self._name = name if name is not None else host_kernel.name + self._missing_cuda = missing_cuda + self._cuda_path = None if cuda_path is None else Path(cuda_path) + self._warned = False + + def __repr__(self) -> str: + return ( + f"Kernel(name={self._name!r}, cuda={self._cuda_kernel is not None}, " + f"missing_cuda={self._missing_cuda!r})" + ) + + @property + def name(self) -> str: + """Name of the kernel.""" + return self._name + + @property + def host_kernel(self) -> PyccelKernel: + """The host kernel.""" + return self._host_kernel + + @property + def cuda_kernel(self) -> CudaKernel | None: + """The CUDA kernel, or None if there is none.""" + return self._cuda_kernel + + @property + def has_cuda(self) -> bool: + """Whether there is a CUDA kernel.""" + return self._cuda_kernel is not None + + @property + def missing_cuda(self) -> str: + """What happens on the CuPy backend without CUDA kernel. + + ``"raise"`` or ``"fallback"``, see :class:`Kernel`. + """ + return self._missing_cuda + + @property + def cuda_path(self) -> Path | None: + """Where the CUDA kernel is expected, if known.""" + return self._cuda_path + + def get_kernel(self) -> PyccelKernel | CudaKernel: + """The kernel for the active backend. + + Call this once at setup to fail early if a CUDA kernel is missing. + + Raises + ------ + NotImplementedError + On the CuPy backend, if there is no CUDA kernel and + ``missing_cuda="raise"``. + """ + if get_backend() != "cupy": + return self._host_kernel + if self._cuda_kernel is not None: + return self._cuda_kernel + if self._missing_cuda == "raise": + expected = ( + "" if self._cuda_path is None else f" (expected {self._cuda_path})" + ) + raise NotImplementedError( + f"No CUDA version of kernel {self._name!r}{expected}." + ) + if not self._warned: + warnings.warn( + f"No CUDA version of kernel {self._name!r}: calling the host kernel, " + "which copies its arrays to the host and back at every call.", + RuntimeWarning, + stacklevel=3, + ) + self._warned = True + return self._host_kernel + + def __call__(self, *args: Any, n_threads: int | None = None) -> Any: + """Call the kernel for the active backend. + + Parameters + ---------- + *args + Kernel arguments. + n_threads : int | None + Number of CUDA threads; required when the CUDA kernel is called, + ignored by the host kernel. + """ + kernel = self.get_kernel() + if kernel is self._host_kernel: + return kernel(*args) + if n_threads is None: + raise ValueError( + f"{self._name}: n_threads is required to launch the CUDA kernel" + ) + return kernel(*args, n_threads=n_threads) + + +class KernelCatalog(Mapping): + """Kernels by name. + + A read-only mapping from names to :class:`Kernel` objects, usually built with + :meth:`from_package`; :meth:`register` adds kernels one by one. + + Parameters + ---------- + kernels : Mapping[str, Kernel] | None + Initial kernels. + """ + + def __init__(self, kernels: Mapping[str, Kernel] | None = None) -> None: + self._kernels: dict[str, Kernel] = {} + for name, kernel in (kernels or {}).items(): + self.register(kernel, name=name) + + @classmethod + def from_package( + cls, + package: str, + *, + host_suffix: str = "_kernels", + cuda_suffix: str = "_cuda.cu", + missing_cuda: str = "raise", + **cuda_options: Any, + ) -> KernelCatalog: + """Collect the kernels of a package with one folder per kernel. + + For every subfolder ```` of the package that contains the module + ``.py``, the function ```` of that module is the + host kernel, and ```` in the same folder, if present, + is the CUDA kernel (with a ``__global__`` function ````). + + Parameters + ---------- + package : str + Full name of the package, e.g. ``__name__`` in its ``__init__.py``. + host_suffix : str + Module name suffix of the host kernels. + cuda_suffix : str + File name suffix of the CUDA kernels. + missing_cuda : {"raise", "fallback"} + Passed on to every :class:`Kernel`. + **cuda_options + Passed on to :meth:`CudaKernel.from_file`, e.g. ``block_size`` or + ``include_dirs``. + """ + root = Path(importlib.import_module(package).__file__).parent + kernels = {} + for folder in sorted(p for p in root.iterdir() if p.is_dir()): + name = folder.name + if not (folder / f"{name}{host_suffix}.py").is_file(): + continue + module = importlib.import_module(f"{package}.{name}.{name}{host_suffix}") + cuda_path = folder / f"{name}{cuda_suffix}" + cuda_kernel = ( + CudaKernel.from_file(cuda_path, name, **cuda_options) + if cuda_path.is_file() + else None + ) + kernels[name] = Kernel( + getattr(module, name), + cuda_kernel, + name=name, + missing_cuda=missing_cuda, + cuda_path=cuda_path, + ) + return cls(kernels) + + def register(self, kernel: Kernel, name: str | None = None) -> Kernel: + """Add a kernel under `name` (by default its own name) and return it.""" + if not isinstance(kernel, Kernel): + raise TypeError(f"expected a Kernel, got {type(kernel).__name__}") + name = kernel.name if name is None else name + if name in self._kernels: + raise KeyError(f"a kernel named {name!r} is already registered") + self._kernels[name] = kernel + return kernel + + def __getitem__(self, name: str) -> Kernel: + return self._kernels[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._kernels) + + def __len__(self) -> int: + return len(self._kernels) + + def __repr__(self) -> str: + return ( + f"KernelCatalog({len(self)} kernels, {len(self.without_cuda)} without CUDA)" + ) + + @property + def without_cuda(self) -> list[str]: + """Names of the kernels without a CUDA kernel, i.e. still to port.""" + return [name for name, kernel in self._kernels.items() if not kernel.has_cuda] diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py new file mode 100644 index 0000000..d43fc74 --- /dev/null +++ b/tests/unit/test_cuda_kernel.py @@ -0,0 +1,309 @@ +"""Tests for `cunumpy.CudaKernel` and `cunumpy.parse_cuda_signature`. + +Signature parsing and argument checking run everywhere: a small stand-in for a +device array (`FakeDeviceArray`) takes the place of CuPy arrays. Launching +kernels needs a GPU and is skipped without one. +""" + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import CudaArguments, CudaKernel, parse_cuda_signature + +AXPY = r""" +// y = a * x + y +extern "C" __global__ +void axpy(double a, const double* __restrict__ x, double *y, int n) +{ + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) y[i] += a * x[i]; +} +""" + +ALL_TYPES = r""" +#include +/* several kernels; the one we look for is not the first */ +extern "C" __global__ void other(int n) {} +extern "C" __global__ +void all_types(bool b, char c, unsigned char uc, short s, int i, unsigned u, + long l, long long ll, unsigned long long ull, size_t sz, + int64_t i64, uint32_t u32, float f, double d, + complex cf, thrust::complex cd, + const float* pf, int* __restrict__ pi, void* pv, double arr[]) +{ +} +""" + + +class FakeDeviceArray: + """Enough of a CuPy array for the argument checks: a dtype and the interface.""" + + __cuda_array_interface__ = {} + + def __init__(self, dtype): + self.dtype = np.dtype(dtype) + + +def _skip_without_cupy(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + + +# --------------------------------------------------------------------------- +# signature parsing +# --------------------------------------------------------------------------- + + +def test_parse_axpy(): + params = parse_cuda_signature(AXPY, "axpy") + assert [(p.name, p.ctype, p.pointer) for p in params] == [ + ("a", "double", False), + ("x", "double", True), + ("y", "double", True), + ("n", "int", False), + ] + assert [p.dtype for p in params] == [np.dtype(np.float64)] * 3 + [ + np.dtype(np.int32) + ] + + +def test_parse_all_types(): + params = {p.name: p for p in parse_cuda_signature(ALL_TYPES, "all_types")} + expected = { + "b": np.bool_, + "c": np.int8, + "uc": np.uint8, + "s": np.int16, + "i": np.int32, + "u": np.uint32, + "l": np.int64, + "ll": np.int64, + "ull": np.uint64, + "sz": np.uint64, + "i64": np.int64, + "u32": np.uint32, + "f": np.float32, + "d": np.float64, + "cf": np.complex64, + "cd": np.complex128, + "pf": np.float32, + "pi": np.int32, + "arr": np.float64, + } + for name, dtype in expected.items(): + assert params[name].dtype == np.dtype(dtype), name + assert params["pv"].dtype is None and params["pv"].pointer + assert params["arr"].pointer and not params["d"].pointer + + +def test_parse_no_parameters(): + assert parse_cuda_signature('extern "C" __global__ void f() {}', "f") == () + assert parse_cuda_signature('extern "C" __global__ void f(void) {}', "f") == () + + +def test_parse_errors(): + with pytest.raises(ValueError, match="no __global__ function 'missing'"): + parse_cuda_signature(AXPY, "missing") + with pytest.raises(ValueError, match="unsupported type"): + parse_cuda_signature("__global__ void f(MyStruct s) {}", "f") + with pytest.raises(ValueError, match="unsupported type"): + parse_cuda_signature("__global__ void f(double** p) {}", "f") + # a commented-out kernel is not found + with pytest.raises(ValueError, match="no __global__ function"): + parse_cuda_signature("// __global__ void f(int n) {}", "f") + + +def test_unparsable_signature_can_be_skipped(): + source = "#define ARGS double* x, int n\n__global__ void f(ARGS) {}" + with pytest.raises(ValueError): + CudaKernel(source, "f") + kernel = CudaKernel(source, "f", check_signature=False) + assert kernel.signature is None + values = kernel.prepare_args(1, 2.5, "anything") + assert values == (1, 2.5, "anything") # passed on as they are + + +# --------------------------------------------------------------------------- +# argument checks (no GPU needed) +# --------------------------------------------------------------------------- + + +def test_python_scalars_are_cast_to_the_declared_types(): + kernel = CudaKernel(AXPY, "axpy") + x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) + + a, x_out, y_out, n = kernel.prepare_args(2, x, y, 10) + assert type(a) is np.float64 and a == 2.0 # int into double: cast, not garbage + assert type(n) is np.int32 and n == 10 + assert x_out is x and y_out is y # arrays are passed as they are + + a, _, _, n = kernel.prepare_args(0.5, x, y, True) + assert type(a) is np.float64 and type(n) is np.int32 and n == 1 + + +def test_wrong_scalars_raise(): + kernel = CudaKernel(AXPY, "axpy") + x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) + + with pytest.raises(TypeError, match=r"argument 3 \(int n\)"): + kernel.prepare_args(1.0, x, y, 2.5) # float into int + with pytest.raises(OverflowError, match="out of range"): + kernel.prepare_args(1.0, x, y, 2**31) # overflows int + with pytest.raises(TypeError, match="losing information"): + kernel.prepare_args(1.0, x, y, np.int64(5)) # int64 into int + with pytest.raises(TypeError): + kernel.prepare_args("1.0", x, y, 5) + + +def test_numpy_scalars(): + kernel = CudaKernel(AXPY, "axpy") + x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) + + a, _, _, n = kernel.prepare_args(np.float32(1.5), x, y, np.int16(3)) + assert type(a) is np.float64 and type(n) is np.int32 # safe casts + a, _, _, n = kernel.prepare_args(np.float64(1.5), x, y, np.int32(3)) + assert type(a) is np.float64 and type(n) is np.int32 + + kernel_f = CudaKernel("__global__ void f(float a) {}", "f") + with pytest.raises(TypeError, match="losing information"): + kernel_f.prepare_args(np.float64(1.5)) + assert type(kernel_f.prepare_args(1.5)[0]) is np.float32 + + +def test_array_checks(): + kernel = CudaKernel(AXPY, "axpy") + x = FakeDeviceArray(np.float64) + + with pytest.raises( + TypeError, match=r"argument 2 \(double\* y\) must be a CuPy array" + ): + kernel.prepare_args(1.0, x, np.zeros(3), 3) # host array: never copied + with pytest.raises(TypeError, match="must have dtype float64, got float32"): + kernel.prepare_args(1.0, x, FakeDeviceArray(np.float32), 3) + with pytest.raises(TypeError, match="takes 4 arguments"): + kernel.prepare_args(1.0, x, x) + + void_kernel = CudaKernel("__global__ void f(void* p) {}", "f") + void_kernel.prepare_args(FakeDeviceArray(np.int8)) # any dtype + + +def test_argument_objects_are_flattened(): + class Vectors(CudaArguments): + def __init__(self, x, y, n): + super().__init__(x, y, n) + + class Duck: + def __init__(self, x, y, n): + self.values = (x, y, n) + + def __cuda_args__(self): + return self.values + + kernel = CudaKernel(AXPY, "axpy") + x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) + for args in (Vectors(x, y, 7), Duck(x, y, 7)): + a, x_out, y_out, n = kernel.prepare_args(2.0, args) + assert x_out is x and y_out is y and n == 7 and type(n) is np.int32 + + # the flattened arguments are checked too + with pytest.raises(TypeError, match="takes 4 arguments"): + kernel.prepare_args(2.0, CudaArguments(x, y)) + + +def test_from_file(tmp_path): + path = tmp_path / "axpy_cuda.cu" + path.write_text(AXPY) + kernel = CudaKernel.from_file(path, block_size=64) + assert kernel.name == "axpy" and kernel.block_size == 64 + assert f"-I{tmp_path}" in kernel.options + + other = tmp_path / "saxpy.cu" + other.write_text(AXPY) + with pytest.raises(ValueError, match="does not end with"): + CudaKernel.from_file(other) + assert CudaKernel.from_file(other, "axpy").name == "axpy" + + +def test_launch_argument_validation(): + kernel = CudaKernel(AXPY, "axpy") + with pytest.raises(ValueError, match="non-negative"): + kernel(1.0, n_threads=-1) + with pytest.raises(ValueError, match="block_size"): + CudaKernel(AXPY, "axpy", block_size=0) + # n_threads=0 checks the arguments but launches nothing (works without a GPU) + x = FakeDeviceArray(np.float64) + kernel(1.0, x, x, 0, n_threads=0) + + +def test_compile_without_gpu_raises(): + if xp.cupy_available(): + pytest.skip("a GPU is available") + with pytest.raises(RuntimeError, match="CuPy is not installed"): + CudaKernel(AXPY, "axpy").compile() + + +# --------------------------------------------------------------------------- +# launching (GPU) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("n", [1, 127, 128, 129, 1000]) +def test_axpy_on_gpu(n): + _skip_without_cupy() + import cupy as cp + + kernel = CudaKernel(AXPY, "axpy") + x = cp.arange(n, dtype=cp.float64) + y = cp.ones(n) + kernel(2, x, y, n, n_threads=n) + assert cp.allclose(y, 2 * x + 1) + assert kernel.compile() is kernel.compile() # compiled once + + +def test_argument_object_on_gpu(): + _skip_without_cupy() + import cupy as cp + + n = 50 + x, y = cp.arange(n, dtype=cp.float64), cp.zeros(n) + CudaKernel(AXPY, "axpy")(0.5, CudaArguments(x, y, n), n_threads=n) + assert cp.allclose(y, 0.5 * x) + + +def test_stream_and_include_dirs_on_gpu(tmp_path): + _skip_without_cupy() + import cupy as cp + + (tmp_path / "helpers.cuh").write_text( + "__device__ double twice(double v) { return 2 * v; }\n" + ) + source = r""" + #include "helpers.cuh" + extern "C" __global__ void double_it(double* y, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) y[i] = twice(y[i]); + } + """ + kernel = CudaKernel(source, "double_it", include_dirs=[tmp_path], block_size=32) + y = cp.ones(100) + stream = cp.cuda.Stream() + kernel(y, 100, n_threads=100, stream=stream) + stream.synchronize() + assert cp.all(y == 2) + + +def test_scalars_arrive_correctly_on_gpu(): + """Python scalars reach int/long long/float/double/bool parameters correctly.""" + _skip_without_cupy() + import cupy as cp + + source = r""" + extern "C" __global__ + void write(double* out, int a, long long b, float c, double d, bool e) { + out[0] = a; out[1] = (double)b; out[2] = c; out[3] = d; out[4] = e; + } + """ + out = cp.zeros(5) + CudaKernel(source, "write")(out, 3, 2**40, 1.5, 2, True, n_threads=1) + assert out.get().tolist() == [3.0, float(2**40), 1.5, 2.0, 1.0] diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py new file mode 100644 index 0000000..ddd091f --- /dev/null +++ b/tests/unit/test_kernel_dispatch.py @@ -0,0 +1,172 @@ +"""Tests for `cunumpy.Kernel` (host/CUDA pairs) and `cunumpy.KernelCatalog`. + +The host kernels here are plain Python functions (wrapped in `PyccelKernel`), so +the NumPy-side tests run everywhere; the CuPy-side tests need a GPU. +""" + +import importlib +import sys +import textwrap + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import CudaKernel, Kernel, KernelCatalog, PyccelKernel + +SCALE_CUDA = r""" +extern "C" __global__ void scale(double* x, double factor, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) x[i] *= factor; +} +""" + + +def scale(x, factor, n): + """Host version of the `scale` kernel.""" + for i in range(n): + x[i] *= factor + + +def _skip_without_cupy(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + + +# --------------------------------------------------------------------------- +# Kernel +# --------------------------------------------------------------------------- + + +def test_kernel_on_numpy(): + kernel = Kernel(scale, CudaKernel(SCALE_CUDA, "scale")) + assert isinstance(kernel.host_kernel, PyccelKernel) # callables are wrapped + assert kernel.name == "scale" and kernel.has_cuda + + x = np.ones(4) + with xp.use_backend("numpy"): + assert kernel.get_kernel() is kernel.host_kernel + kernel(x, 3.0, 4, n_threads=4) # n_threads is ignored by the host kernel + kernel(x, 2.0, 4) + assert np.all(x == 6.0) + + +def test_kernel_arguments(): + with pytest.raises(ValueError, match="missing_cuda"): + Kernel(scale, missing_cuda="ignore") + with pytest.raises(TypeError, match="cuda_kernel must be a CudaKernel"): + Kernel(scale, PyccelKernel(scale)) + assert Kernel(scale, name="other").name == "other" + assert not Kernel(scale).has_cuda + + +def test_kernel_on_cupy(): + _skip_without_cupy() + import cupy as cp + + kernel = Kernel(scale, CudaKernel(SCALE_CUDA, "scale")) + x = cp.ones(300) + with xp.use_backend("cupy"): + assert kernel.get_kernel() is kernel.cuda_kernel + kernel(x, 3, 300, n_threads=300) + with pytest.raises(ValueError, match="n_threads is required"): + kernel(x, 3, 300) + assert cp.all(x == 3.0) + + +def test_missing_cuda_raises_on_cupy(): + _skip_without_cupy() + kernel = Kernel(scale, cuda_path="kernels/scale/scale_cuda.cu") + with xp.use_backend("cupy"): + with pytest.raises( + NotImplementedError, + match="No CUDA version of kernel 'scale'.*scale_cuda.cu", + ): + kernel.get_kernel() + + +def test_missing_cuda_fallback_on_cupy(): + """missing_cuda="fallback" calls the host kernel via PyccelKernel (host copies).""" + _skip_without_cupy() + import cupy as cp + + kernel = Kernel(scale, missing_cuda="fallback") + x = cp.ones(5) + with xp.use_backend("cupy"): + with pytest.warns(RuntimeWarning, match="copies its arrays to the host"): + kernel(x, 4.0, 5) + kernel(x, 0.5, 5) # warned only once + assert cp.all(x == 2.0) + + +# --------------------------------------------------------------------------- +# KernelCatalog +# --------------------------------------------------------------------------- + + +@pytest.fixture +def kernel_package(tmp_path, monkeypatch): + """One folder per kernel: `scale` has a CUDA kernel, `shift` has not.""" + root = tmp_path / "demo_kernel_pkg" + for name, body in (("scale", "x[i] *= a"), ("shift", "x[i] += a")): + (root / name).mkdir(parents=True) + (root / name / "__init__.py").write_text("") + (root / name / f"{name}_kernels.py").write_text( + textwrap.dedent( + f""" + def {name}(x, a, n): + for i in range(n): + {body} + """ + ) + ) + (root / "scale" / "scale_cuda.cu").write_text(SCALE_CUDA) + (root / "not_a_kernel").mkdir() + (root / "__init__.py").write_text( + "from cunumpy import KernelCatalog\n\n" + "catalog = KernelCatalog.from_package(__name__)\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + yield importlib.import_module("demo_kernel_pkg").catalog + for module in [m for m in sys.modules if m.startswith("demo_kernel_pkg")]: + del sys.modules[module] + + +def test_catalog_from_package(kernel_package): + catalog = kernel_package + assert list(catalog) == ["scale", "shift"] and len(catalog) == 2 + assert "scale" in catalog and "not_a_kernel" not in catalog + assert catalog.without_cuda == ["shift"] + assert catalog["scale"].cuda_kernel.name == "scale" + assert catalog["shift"].cuda_path.name == "shift_cuda.cu" + + x = np.ones(3) + with xp.use_backend("numpy"): + catalog["scale"](x, 2.0, 3) + catalog["shift"](x, 1.0, 3) + assert np.all(x == 3.0) + + +def test_catalog_on_cupy(kernel_package): + _skip_without_cupy() + import cupy as cp + + x = cp.ones(10) + with xp.use_backend("cupy"): + kernel_package["scale"](x, 5.0, 10, n_threads=10) + with pytest.raises(NotImplementedError, match="shift"): + kernel_package["shift"](x, 1.0, 10, n_threads=10) + assert cp.all(x == 5.0) + + +def test_catalog_register(): + catalog = KernelCatalog() + kernel = catalog.register(Kernel(scale)) + assert catalog["scale"] is kernel and catalog.without_cuda == ["scale"] + with pytest.raises(KeyError, match="already registered"): + catalog.register(Kernel(scale)) + with pytest.raises(TypeError, match="expected a Kernel"): + catalog.register(scale) + catalog.register(Kernel(scale), name="scale_again") + assert dict(catalog).keys() == {"scale", "scale_again"} + assert KernelCatalog({"x": kernel})["x"] is kernel From c545a3e2a6fc3fd9a3ea52fd6707f00bd42c047f Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 30 Sep 2026 22:30:23 +0200 Subject: [PATCH 02/16] Bump version number to 0.2.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 411700d..d90ca1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.2.0" +version = "0.2.1" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] From 61b999d8fb709aec01d6052e051a376526e5baaf Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 30 Sep 2026 23:21:44 +0200 Subject: [PATCH 03/16] Prepare cuda methods (#38) * ruff check --fix * Add CUDA structs, templates, kernel variants, 1D-3D launches, compile_all, host kernel options and MPI device helpers * Test CUDA structs, templates, variants, launch shapes, compile_all and device binding * Document the new CUDA kernel features and MPI device helpers * fixed ruff errors * GPU CI: use CUDA 13.2 consistently and wait for the GitLab pipeline of the pushed commit --- .github/workflows/gpu_ci_trigger.yml | 42 +- .gitlab-ci.yml | 17 +- CHANGELOG.md | 10 + README.md | 52 ++- docs/source/api.md | 250 ++++++++++- src/cunumpy/__init__.py | 21 +- src/cunumpy/__init__.pyi | 10 +- src/cunumpy/cuda_kernel.py | 640 +++++++++++++++++++++++---- src/cunumpy/dispatch.py | 98 +++- src/cunumpy/kernel.py | 3 +- src/cunumpy/xp.py | 89 +++- tests/unit/test_cuda_kernel.py | 372 +++++++++++++++- tests/unit/test_cunumpy.py | 8 +- tests/unit/test_device_binding.py | 75 ++++ tests/unit/test_kernel_dispatch.py | 103 ++++- 15 files changed, 1636 insertions(+), 154 deletions(-) create mode 100644 tests/unit/test_device_binding.py diff --git a/.github/workflows/gpu_ci_trigger.yml b/.github/workflows/gpu_ci_trigger.yml index 4418074..60fa668 100644 --- a/.github/workflows/gpu_ci_trigger.yml +++ b/.github/workflows/gpu_ci_trigger.yml @@ -64,6 +64,7 @@ jobs: # 4. Force push (This automatically starts the GitLab Pipeline) git push -f gitlab HEAD:refs/heads/$TARGET_BRANCH + echo "PUSHED_SHA=$(git rev-parse HEAD)" >> $GITHUB_ENV # 5. Provide the direct link PIPELINE_URL="https://gitlab.mpcdf.mpg.de/maxlin/cunumpy/-/pipelines?ref=$TARGET_BRANCH" @@ -72,10 +73,41 @@ jobs: echo "::notice::View Pipeline: $PIPELINE_URL" - name: Wait for GitLab Pipeline - uses: docker://gitlab/glab:latest + timeout-minutes: 90 env: GITLAB_TOKEN: ${{ secrets.GITLAB_TOKEN }} - GITLAB_HOST: gitlab.mpcdf.mpg.de - with: - entrypoint: glab - args: ci status --live --branch ${{ env.TARGET_BRANCH }} --repo maxlin/cunumpy + run: | + API="https://gitlab.mpcdf.mpg.de/api/v4/projects/maxlin%2Fcunumpy" + AUTH=() + if [ -n "$GITLAB_TOKEN" ]; then AUTH=(--header "PRIVATE-TOKEN: $GITLAB_TOKEN"); fi + + # 1. GitLab creates the pipeline asynchronously after the push: wait until + # the pipeline for the pushed commit exists (asking right away finds none). + PIPELINE="" + for i in $(seq 1 60); do + PIPELINE=$(curl -sf "${AUTH[@]}" "$API/pipelines?ref=$TARGET_BRANCH&sha=$PUSHED_SHA&per_page=1" | jq -r '.[0].id // empty' || true) + if [ -n "$PIPELINE" ]; then break; fi + sleep 5 + done + if [ -z "$PIPELINE" ]; then + echo "::error::No GitLab pipeline for $PUSHED_SHA on $TARGET_BRANCH after 5 minutes" + exit 1 + fi + URL="https://gitlab.mpcdf.mpg.de/maxlin/cunumpy/-/pipelines/$PIPELINE" + echo "::notice::GitLab pipeline: $URL" + + # 2. Follow the pipeline until it has finished. + while true; do + STATUS=$(curl -sf "${AUTH[@]}" "$API/pipelines/$PIPELINE" | jq -r '.status // empty' || true) + case "$STATUS" in + success) + echo "GitLab pipeline $PIPELINE succeeded: $URL" + exit 0 ;; + failed|canceled|skipped|manual|scheduled) + echo "::error::GitLab pipeline $PIPELINE is $STATUS: $URL" + exit 1 ;; + *) + echo "GitLab pipeline $PIPELINE: ${STATUS:-status not available yet}" + sleep 20 ;; + esac + done diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 79d4881..17a6fcb 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -29,18 +29,19 @@ gpu_tests: - git --version - echo "--- Pytest Execution ---" - # The MPCDF image likely has a specific python environment. + # The MPCDF image likely has a specific python environment. # We install our dependencies into the user directory or a virtualenv. - - python3 -m pip install --user cupy-cuda12x - - python3 -m pip install --user nvidia-cublas-cu12 nvidia-cufft-cu12 nvidia-curand-cu12 nvidia-cusolver-cu12 nvidia-cusparse-cu12 + # One CUDA version only: nvhpcsdk/26 provides CUDA 13.2 (its headers are used + # when CuPy compiles kernels with NVRTC), so install CuPy for CUDA 13 with the + # CUDA 13.2 libraries and NVRTC. Mixing CUDA 12 (cupy-cuda12x) with the CUDA 13.2 + # headers fails to compile CuPy's own kernels (e.g. CUB reductions). + - python3 -m pip install --user "cupy-cuda13x[ctk]" "cuda-toolkit==13.2.*" - python3 -m pip install --user -e . - + # Add the user bin to PATH for pytest - export PATH="$HOME/.local/bin:$PATH" - - # Try to find libcublas and other libraries in the HPC environment - - export LD_LIBRARY_PATH=$(find /mpcdf/soft /opt/nvidia -name libcublas.so.12 -exec dirname {} \; 2>/dev/null | head -n 1):$LD_LIBRARY_PATH - + - export ARRAY_BACKEND=cupy + - python3 -c "import cupy; cupy.show_config()" - pytest -xvs . diff --git a/CHANGELOG.md b/CHANGELOG.md index 36df17a..a06fae6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `xp.parse_cuda_signature(source, name)` and `xp.CudaParameter`: Parse the parameters of a `__global__` function. - `xp.Kernel`: A host kernel (`PyccelKernel`) and its CUDA counterpart, calling the one matching the active backend. Without a CUDA kernel on the CuPy backend it raises `NotImplementedError` (`missing_cuda="raise"`, default) or falls back to the host kernel with host copies (`missing_cuda="fallback"`). - `xp.KernelCatalog`: Read-only mapping of `Kernel`s; `KernelCatalog.from_package()` collects them from a package with one folder per kernel (`name/name_kernels.py`, `name/name_cuda.cu`); `without_cuda` lists the kernels still to port. +- `xp.CudaStruct` and `xp.CudaStructValue`: C structs passed to CUDA kernels by value. A `CudaStruct` is defined once from `(field, C type)` pairs; it provides the C `declaration`, the NumPy `dtype` with the C memory layout, and packs values (device arrays as addresses, scalars checked and cast) into a `CudaStructValue` that is passed as one kernel argument. `CudaKernel(..., structs=[...])` checks struct parameters and that a struct definition in the source matches. +- `CudaKernel(..., template_args=...)`: Instantiate C++ function templates (e.g. `template_args=(np.float64, 3)` for `name`); the template parameters are substituted into the checked signature. +- `xp.CudaKernelVariants`: Creates and caches one `CudaKernel` per variant key for generated kernel sources (e.g. per dimension and dtype); `compile_all()` compiles given and existing variants. +- `xp.ctype_of(dtype)`: The C type of a NumPy dtype, e.g. for generating CUDA source. +- 1D to 3D launches: `CudaKernel` accepts a tuple `block_size`, and calls take `n_threads` as an integer or tuple, or an explicit `grid`, plus a per-call `block`; `launch_shape()` returns the `(grid, block)` of a call. Dynamic shared memory (`shared_mem`) and `stream` are passed through by `Kernel` as well. +- Compiling at setup: `CudaKernel.compile()` and `is_compiled`, `Kernel.compile()`, and `KernelCatalog.compile_all()`. +- `Kernel(..., host_options=...)` and `KernelCatalog.from_package(..., host_options=...)`: `PyccelKernel` options (e.g. `object_modules`, `outputs`) for the host kernels, for all kernels or per kernel name; needed for the fallback to find device arrays inside application objects. +- `xp.local_rank()`: The node-local rank from the MPI launcher's environment (Open MPI, MVAPICH2, Intel MPI/MPICH, PMI, Cray PALS, Slurm, `LOCAL_RANK`), available before `MPI_Init`. +- `xp.bind_local_device()`: Selects the GPU `local_rank() % device_count()` and creates its context, before `MPI_Init`, for one-rank-per-GPU MPI programs. +- `xp.synchronize_for_mpi(*arrays)`: Waits for pending work on the current stream before MPI uses device buffers (no-op for host buffers and on the NumPy backend). ## [0.2.0] - 2026-09-28 diff --git a/README.md b/README.md index d133519..837b926 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,22 @@ active CuPy device and `None` on NumPy. `set_device_for_rank(rank)` is a round-robin convenience for MPI layouts where local ranks map contiguously to GPUs. If your scheduler uses a different mapping, select the device directly. +For MPI programs with one rank per GPU, `bind_local_device()` selects the GPU +from the node-local rank that the MPI launcher exports (`local_rank()`), so it +can run before MPI is initialized, as CUDA-aware MPI requires. Before passing +device buffers to MPI, call `synchronize_for_mpi(*buffers)`: kernels run +asynchronously, and MPI would otherwise send a buffer a kernel is still +writing, without an error. + +```python +xp.set_backend("cupy") +xp.bind_local_device() # before MPI_Init +from mpi4py import MPI + +xp.synchronize_for_mpi(send, recv) +MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source) +``` + CuPy caches released allocations in memory pools. This can make process-level GPU memory appear occupied after arrays go out of scope. `free_memory()` asks CuPy to release currently free cached blocks; it does not free memory still @@ -253,11 +269,39 @@ with xp.use_backend("cupy"): On the CuPy backend, a `Kernel` without CUDA kernel raises `NotImplementedError` (or, with `missing_cuda="fallback"`, runs the host kernel -through `PyccelKernel`, with host copies). Objects implementing +through `PyccelKernel`, with host copies; `host_options` configure that +`PyccelKernel`). `KernelCatalog.from_package()` collects kernel pairs from a +package with one folder per kernel (`name/name_kernels.py` and +`name/name_cuda.cu`), and `catalog.compile_all()` compiles all CUDA kernels at +setup. + +Groups of arguments can be passed as one: objects implementing `__cuda_args__()` (see `CudaArguments`) are flattened into several kernel -arguments. `KernelCatalog.from_package()` collects kernel pairs from a package -with one folder per kernel (`name/name_kernels.py` and `name/name_cuda.cu`). -See the [API reference](docs/source/api.md) for details. +arguments, and `CudaStruct` defines a C struct once (its C `declaration` and +the matching memory layout) and packs values into it, which the kernel takes +as one parameter: + +```python +Vec = xp.CudaStruct("Vec", [("data", "double*"), ("n", "int")]) +scale = xp.CudaKernel( + Vec.declaration + + r""" + extern "C" __global__ void scale(Vec v, double a) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < v.n) v.data[i] *= a; + }""", + "scale", + structs=[Vec], +) +scale(Vec(data=y, n=y.size), 0.5, n_threads=y.size) +``` + +Launches can be 1D to 3D (`n_threads=(nx, ny)`, `block_size=(16, 16)`) or use +an explicit `grid`, with dynamic shared memory (`shared_mem`) and a `stream`. +C++ function templates are instantiated with `template_args`, and +`CudaKernelVariants` caches kernels whose source is generated per variant +(e.g. per dimension and dtype). See the [API reference](docs/source/api.md) for +details. ## Pyodide diff --git a/docs/source/api.md b/docs/source/api.md index 2cf90f2..6bbdae0 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -227,6 +227,52 @@ mapping differs: device_id = xp.set_device_for_rank(mpi_rank) ``` +### `local_rank()` + +The rank of the process within its node, read from the environment variables +that MPI launchers export (Open MPI, MVAPICH2, Intel MPI/MPICH, PMI, Cray +PALS, Slurm, `LOCAL_RANK`), or `0` if none is set. The launcher sets them +before `MPI_Init`, so this works before MPI is initialized and without +importing `mpi4py`. + +### `bind_local_device()` + +Selects device `local_rank() % device_count()` for this process and creates its +CUDA context. Returns the device id, or `None` on the NumPy backend or without +devices. Call it before `MPI_Init` (before importing `mpi4py.MPI`), so that a +CUDA-aware MPI sees the right device; otherwise all ranks of a node would use +device 0. If the launcher gives each rank its own device through +`CUDA_VISIBLE_DEVICES`, each process sees one device and selects it: + +```python +import cunumpy as xp + +xp.set_backend("cupy") +xp.bind_local_device() +from mpi4py import MPI # initializes MPI after the device is bound +``` + +Unlike `set_device_for_rank()`, it needs no MPI rank, and it uses the rank +within the node rather than assuming contiguous ranks per node. + +### `synchronize_for_mpi(*arrays)` + +Waits for the work pending on the current stream if at least one of `arrays` +is a CuPy array; `None` entries and host arrays are ignored, so it costs +nothing for host buffers and on the NumPy backend. Call it before every MPI +call that sends or receives device buffers: CuPy launches kernels +asynchronously and MPI knows nothing about CUDA streams, so a buffer that a +kernel is still writing would be sent as it is at that moment, without an +error: + +```python +xp.synchronize_for_mpi(send_buffer, recv_buffer) +comm.Sendrecv(send_buffer, dest, recvbuf=recv_buffer, source=source) +``` + +No synchronization is needed after MPI returns: kernels launched afterwards see +the received data. + ### `memory_info()` Returns `(free_bytes, total_bytes)` reported by the CUDA runtime for the @@ -365,15 +411,19 @@ xp.CudaKernel( block_size=128, options=(), include_dirs=(), + structs=(), + template_args=None, check_signature=True, ) xp.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs) ``` -Wraps the `extern "C" __global__` function `name` in the CUDA C `source`. The -kernel is compiled with NVRTC through `cupy.RawKernel` on the first call (or -by `compile()`), and cached. CuPy is imported only then, so kernels can be -created and their signatures parsed without CuPy. +Wraps the `__global__` function `name` in the CUDA C `source` (declared +`extern "C"`, unless it is a template). The kernel is compiled with NVRTC +through CuPy on the first call (or by `compile()`), and cached, also on disk by +CuPy. CuPy is imported only then, so kernels can be created and their +signatures parsed without CuPy; `compile()` raises `RuntimeError` without a +GPU. `from_file` reads the source from a file; the kernel name defaults to the file name without `suffix` (`axpy_cuda.cu` -> `axpy`), and the directory of the file @@ -381,30 +431,79 @@ is added to the include directories. ### Parameters -* `block_size`: threads per block. +* `block_size`: threads per block, an integer for 1D launches or a tuple of 1 + to 3 integers, e.g. `(16, 16)`; at most 1024 threads in total. * `options`: additional NVRTC options, e.g. `("-std=c++17",)`. * `include_dirs`: directories for `#include`, passed as `-I`. +* `structs`: `CudaStruct` types that the kernel takes as parameters (by + value), see `CudaStruct` below. +* `template_args`: template arguments if `name` is a function template, see + "Templates and generated kernels" below. * `check_signature`: parse the signature and check every call against it (default). Raises `ValueError` if the signature cannot be parsed, e.g. with - templates, macros or pointers to pointers in the parameter list; pass - `False` to launch with the arguments as they are, like `cupy.RawKernel`. + macros or pointers to pointers in the parameter list; pass `False` to launch + with the arguments as they are, like `cupy.RawKernel`. + +Properties: `name`, `expression` (`name`, or the template instantiation such +as `"scale"`), `source`, `block_size`, `options`, `structs`, +`template_args`, `signature`, `is_compiled`. ### Calling ```python -kernel(*args, n_threads, shared_mem=0, stream=None) +kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) +``` + +Launches the kernel on `stream` (the current stream if `None`). The launch +shape is given either by `n_threads` or by `grid`: + +* `n_threads`: number of threads, an integer or a tuple of 1 to 3 integers such + as `(nx, ny)`. The grid is `ceil(n_threads / block)` per dimension. With a 1D + `block_size` and multi-dimensional `n_threads`, the block is + `(block_size, 1, ...)`. +* `grid`: number of blocks per dimension, instead of `n_threads`. +* `block`: block shape for this call, instead of `block_size`. +* `shared_mem`: dynamic shared memory per block in bytes, for + `extern __shared__` arrays. + +Nothing is launched if the grid has a zero dimension (e.g. `n_threads=0`). +`kernel.launch_shape(n_threads=None, *, grid=None, block=None)` returns the +`(grid, block)` a call would use, e.g. to size a per-block output: + +```python +BLOCK_SUM = r""" +extern "C" __global__ void block_sum(const double* x, double* out, int n) { + extern __shared__ double buffer[]; + int i = blockDim.x * blockIdx.x + threadIdx.x; + buffer[threadIdx.x] = i < n ? x[i] : 0.0; + __syncthreads(); + for (int s = blockDim.x / 2; s > 0; s /= 2) { + if (threadIdx.x < s) buffer[threadIdx.x] += buffer[threadIdx.x + s]; + __syncthreads(); + } + if (threadIdx.x == 0) out[blockIdx.x] = buffer[0]; +} +""" +block_sum = xp.CudaKernel(BLOCK_SUM, "block_sum", block_size=128) +(n_blocks,), _ = block_sum.launch_shape(x.size) +partial = xp.zeros(n_blocks) +block_sum(x, partial, x.size, n_threads=x.size, shared_mem=128 * 8) ``` -Launches `ceil(n_threads / block_size)` blocks of `block_size` threads on -`stream` (the current stream if `None`); nothing is launched for -`n_threads=0`. The arguments are prepared by `kernel.prepare_args(*args)`: +In a 2D kernel, use `blockIdx.y`/`threadIdx.y` for the second dimension and +launch with `n_threads=(nx, ny)` and, e.g., `block_size=(16, 16)`. + +### Argument checks + +The arguments are prepared by `kernel.prepare_args(*args)`: * arguments with a `__cuda_args__()` method are replaced by the values it - returns (see `CudaArguments`); + returns (see `CudaArguments` and `CudaStruct` below); * with a checked signature, the number of arguments must match, and * pointer parameters take CuPy arrays whose dtype matches the pointed-to type (any dtype for `void*`); host arrays raise `TypeError`, they are never copied to the device; + * struct parameters take values of that `CudaStruct`; * Python scalars are cast to the declared type: `int` into integer (with a range check, `OverflowError`), floating-point and complex parameters, `float` into floating-point and complex parameters, `bool` into boolean @@ -425,9 +524,87 @@ itself costs about as much), which is negligible for kernels that run for C types are mapped to NumPy dtypes as on Linux (LP64): `int` is `int32`, `long` and `long long` are `int64`, `float` is `float32`, `double` is `float64`, `complex` is `complex128`; fixed-width types such as -`int64_t` and `size_t` are supported too. `xp.parse_cuda_signature(source, -name)` returns the parsed parameters (`CudaParameter` tuples of `name`, -`ctype`, `dtype`, `pointer`). +`int64_t` and `size_t` are supported too. `xp.ctype_of(dtype)` gives the C +type of a dtype (`xp.ctype_of(np.float64) == "double"`), e.g. to generate +source. `xp.parse_cuda_signature(source, name, *, structs=(), +template_args=None)` returns the parsed parameters (`CudaParameter` tuples of +`name`, `ctype`, `dtype`, `pointer`, `struct`). + +### Templates and generated kernels + +A function template is instantiated with `template_args`: C types (or NumPy +dtypes, converted with `ctype_of`) for type parameters, integers or booleans +for non-type parameters. The template parameters are substituted into the +signature, so calls are checked as for any other kernel: + +```python +SCALE = r""" +template +__global__ void scale(T* x, T factor, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) x[i] = factor * x[i] * (T)N; +} +""" +scale_f64 = xp.CudaKernel(SCALE, "scale", template_args=(np.float64, 3)) +scale_f64(x, 2.0, x.size, n_threads=x.size) # instantiation scale +``` + +For kernels whose source is generated per variant (e.g. per number of +dimensions and dtype), `CudaKernelVariants` creates and caches one kernel per +key: + +```python +matvec = xp.CudaKernelVariants( + lambda ndim, dtype: xp.CudaKernel(make_source(ndim, xp.ctype_of(dtype)), "matvec") +) +matvec.get(3, np.float64)(mat, x, out, n_threads=out.size) # created once +matvec.compile_all([(3, np.float64), (3, np.complex128)]) # at setup +``` + +`get(*key)` calls the factory the first time a key is used; `keys()`, +iteration and `len()` give the variants created so far; `compile_all(keys=())` creates the +given variants and compiles all of them. + +## `CudaStruct` + +```python +Particles = xp.CudaStruct( + "Particles", + [("x", "double*"), ("v", "double*"), ("n", "int"), ("charge", "double")], +) +source = Particles.declaration + r""" +extern "C" __global__ void push(Particles p, double dt) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < p.n) p.x[i] += dt * p.charge * p.v[i]; +} +""" +push = xp.CudaKernel(source, "push", structs=[Particles]) +push(Particles(x=x, v=v, n=x.size, charge=-1.0), 0.1, n_threads=x.size) +``` + +A C struct passed to kernels by value. It groups arguments, e.g. all arrays +describing a set of particles, into one kernel parameter, so adding a field +changes one definition instead of every kernel signature. + +`CudaStruct(name, fields)` takes the fields as `(name, C type)` pairs; scalar +fields and pointers to the scalar types above (or `void*`) are supported. + +* `declaration`: the C definition of the struct, to put in the CUDA source + or a header. +* `dtype`: the NumPy structured dtype with the memory layout of the C struct + (C alignment and padding; pointers stored as 64-bit device addresses). +* `fields`: the parsed fields (`CudaParameter` tuples). +* `check_source(source)`: raises `ValueError` if `source` defines the struct + with other fields; a kernel created with `structs=[...]` does this check. +* Calling the struct with keyword arguments, one per field, packs the values: + pointer fields take CuPy arrays of the declared dtype (never copied), scalar + fields are checked and cast like scalar kernel arguments. + +The result is a `CudaStructValue`: it keeps references to the arrays it points +to (the packed struct only holds their addresses, so keep the value alive while +the kernel may run), gives access to the field values with +`value["field"]`, holds the packed struct in `value.packed`, and is flattened +into it when passed to a kernel. ## `CudaArguments` @@ -446,7 +623,8 @@ for several kernel parameters. `CudaArguments(*values)` stores the values; `__cuda_args__()` method returning a tuple is flattened. This lets an application keep its host argument objects (e.g. Pyccel classes holding NumPy arrays) and matching device argument objects that reference the same data on -the device, and pass either to the same call. +the device, and pass either to the same call. A `CudaArguments` object may also +return struct values (`CudaStructValue.packed`) among its values. ## `Kernel` @@ -458,15 +636,23 @@ xp.Kernel( name=None, missing_cuda="raise", cuda_path=None, + host_options=None, ) ``` A host kernel (a `PyccelKernel`; other callables are wrapped in one) and its CUDA counterpart. `kernel.get_kernel()` returns the host kernel on the NumPy backend and the CUDA kernel on the CuPy backend; call it once at setup to fail -early if a CUDA kernel is missing. `kernel(*args, n_threads=None)` calls the -kernel of the active backend; `n_threads` is required for the CUDA kernel and -ignored by the host kernel. +early if a CUDA kernel is missing. + +```python +kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) +``` + +calls the kernel of the active backend. The launch arguments are passed to the +CUDA kernel (`n_threads` or `grid` is required there) and ignored by the host +kernel. `kernel.compile()` compiles the CUDA kernel now and returns whether +there is one. Without a CUDA kernel on the CuPy backend, `missing_cuda="raise"` raises `NotImplementedError` (naming `cuda_path`, if given), and @@ -474,6 +660,13 @@ Without a CUDA kernel on the CuPy backend, `missing_cuda="raise"` raises copies the arrays to the host and back at every call (a `RuntimeWarning` is emitted once). +`host_options` are keyword arguments for the `PyccelKernel` that wraps a plain +callable `host_kernel`, e.g. `{"object_modules": ("my_package.",), "outputs": +(2,)}`. They matter for the fallback: `object_modules` lets it find the device +arrays inside application objects, and `outputs` limits the copies back to the +device. Passing `host_options` together with a `PyccelKernel` raises +`ValueError`; configure that `PyccelKernel` directly. + Properties: `name`, `host_kernel`, `cuda_kernel`, `has_cuda`, `missing_cuda`, `cuda_path`. @@ -486,6 +679,7 @@ catalog = xp.KernelCatalog.from_package( host_suffix="_kernels", cuda_suffix="_cuda.cu", missing_cuda="raise", + host_options=None, **cuda_options, ) kernel = catalog["push"] @@ -495,8 +689,8 @@ A read-only mapping from names to `Kernel` objects. `from_package` scans the subfolders of `package`: for every folder `` containing the module `.py`, the function `` of that module is the host kernel, and `` in the same folder, if present, is the CUDA -kernel (`__global__` function ``). `cuda_options` are passed on to -`CudaKernel.from_file`. Typically called in the package's `__init__.py`: +kernel (`__global__` function ``). Typically called in the package's +`__init__.py`: ```text my_kernels/ @@ -508,8 +702,18 @@ my_kernels/ └── deposit_kernels.py # no CUDA kernel yet ``` -`catalog.without_cuda` lists the kernels still to port. `KernelCatalog(kernels)` -and `catalog.register(kernel, name=None)` build a catalog by hand. +* `host_options`: `PyccelKernel` options for the host kernels (see `Kernel`), + the same for all kernels or a function of the kernel name, e.g. + `lambda name: {"outputs": OUTPUTS[name]}`. +* `cuda_options`: passed on to `CudaKernel.from_file`, e.g. `block_size`, + `include_dirs` or `structs`. + +`catalog.without_cuda` lists the kernels still to port. +`catalog.compile_all()` compiles every CUDA kernel and returns their names; +call it at setup so that the first time step does not pay for compilation +(after the first run, CuPy loads the kernels from its disk cache). +`KernelCatalog(kernels)` and `catalog.register(kernel, name=None)` build a +catalog by hand. ## Version diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 8ac8dff..a46fa66 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -2,11 +2,21 @@ from importlib.metadata import PackageNotFoundError, version from . import xp -from .cuda_kernel import CudaArguments, CudaKernel, CudaParameter, parse_cuda_signature +from .cuda_kernel import ( + CudaArguments, + CudaKernel, + CudaKernelVariants, + CudaParameter, + CudaStruct, + CudaStructValue, + ctype_of, + parse_cuda_signature, +) from .dispatch import Kernel, KernelCatalog from .kernel import PyccelKernel from .xp import ( assert_same_backend, + bind_local_device, cupy_available, default_float_dtype, device_count, @@ -17,6 +27,7 @@ get_rng, is_cpu, is_gpu, + local_rank, memory_info, pin_memory, same_backend, @@ -25,6 +36,7 @@ set_device_for_rank, stream, synchronize, + synchronize_for_mpi, to_cunumpy, to_cupy, to_numpy, @@ -39,12 +51,17 @@ __all__ = [ "CudaArguments", "CudaKernel", + "CudaKernelVariants", "CudaParameter", + "CudaStruct", + "CudaStructValue", "Kernel", "KernelCatalog", "PyccelKernel", "__version__", "assert_same_backend", + "bind_local_device", + "ctype_of", "cupy_available", "cupy_backend", "default_float_dtype", @@ -56,6 +73,7 @@ "get_rng", "is_cpu", "is_gpu", + "local_rank", "memory_info", "numpy_backend", "parse_cuda_signature", @@ -66,6 +84,7 @@ "set_device_for_rank", "stream", "synchronize", + "synchronize_for_mpi", "to_cunumpy", "to_cupy", "to_numpy", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 6cc2098..2006279 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -1,8 +1,9 @@ # Stub file for Pylance/mypy: exposes all numpy symbols so that # `import cunumpy as xp` followed by `xp.` shows numpy completions. # At runtime the real __init__.py dispatches to numpy or cupy via __getattr__. +from collections.abc import Generator from contextlib import contextmanager -from typing import Any, Generator +from typing import Any import numpy as np from numpy import * @@ -10,7 +11,11 @@ from numpy import * from . import xp as xp from .cuda_kernel import CudaArguments as CudaArguments from .cuda_kernel import CudaKernel as CudaKernel +from .cuda_kernel import CudaKernelVariants as CudaKernelVariants from .cuda_kernel import CudaParameter as CudaParameter +from .cuda_kernel import CudaStruct as CudaStruct +from .cuda_kernel import CudaStructValue as CudaStructValue +from .cuda_kernel import ctype_of as ctype_of from .cuda_kernel import parse_cuda_signature as parse_cuda_signature from .dispatch import Kernel as Kernel from .dispatch import KernelCatalog as KernelCatalog @@ -32,6 +37,9 @@ def use_backend(backend: str) -> Generator[None]: ... def set_backend(backend: str) -> None: ... def set_device(device_id: int) -> None: ... def set_device_for_rank(rank: int, devices_per_node: int | None = ...) -> int: ... +def local_rank() -> int: ... +def bind_local_device() -> int | None: ... +def synchronize_for_mpi(*arrays: Any) -> None: ... def device_count() -> int: ... def memory_info() -> tuple[int, int] | None: ... def free_memory() -> None: ... diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index c4abf53..026362b 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -6,17 +6,22 @@ * argument objects that implement the :class:`CudaArguments` protocol (a ``__cuda_args__()`` method) are flattened into their device arrays and scalars, so an object holding several arrays can be passed as one argument; -* the ``extern "C" __global__`` signature is parsed once, and every call is - checked against it: the number of arguments, the dtype of every array, and +* C structs can be passed by value: :class:`CudaStruct` defines the struct once, + generates its C declaration and packs its values; +* the ``__global__`` signature is parsed once, and every call is checked against + it: the number of arguments, the dtype of every array, every struct, and every scalar. Python scalars are cast to the declared C type; a scalar that does not fit the declared type (a ``float`` for an ``int``, an integer out of range, a NumPy scalar that would lose precision) raises instead of reaching the kernel as a silently wrong value, which is what ``cupy.RawKernel`` would do; -* arrays are never converted or copied: they must already be CuPy arrays. +* arrays are never converted or copied: they must already be CuPy arrays; +* C++ function templates are instantiated with ``template_args``, and generated + kernels (one source per variant) are compiled once per variant by + :class:`CudaKernelVariants`. -The number of threads is given at each call (``n_threads``); the kernel is -launched on ``ceil(n_threads / block_size)`` blocks. +The launch shape is given at each call, either as the number of threads +(``n_threads``, in 1 to 3 dimensions) or as an explicit ``grid``. This module imports CuPy only when a kernel is compiled, so it can be imported (and signatures parsed) without CuPy. @@ -26,19 +31,27 @@ import math import re +from collections.abc import Callable, Hashable, Iterable, Iterator, Sequence from contextlib import nullcontext from pathlib import Path -from typing import Any, Callable, NamedTuple, Sequence +from typing import Any, NamedTuple import numpy as np __all__ = [ "CudaArguments", "CudaKernel", + "CudaKernelVariants", "CudaParameter", + "CudaStruct", + "CudaStructValue", + "ctype_of", "parse_cuda_signature", ] +# CUDA limit on the number of threads per block +_MAX_THREADS_PER_BLOCK = 1024 + class CudaArguments: """Base class for objects passed to a :class:`CudaKernel` as one argument. @@ -71,7 +84,7 @@ def __cuda_args__(self) -> tuple[Any, ...]: class CudaParameter(NamedTuple): - """One parameter of a CUDA kernel signature. + """One parameter of a CUDA kernel signature (or one field of a struct). Attributes ---------- @@ -80,16 +93,19 @@ class CudaParameter(NamedTuple): ctype : str Normalized C type without qualifiers or ``*``, e.g. ``"double"``. dtype : numpy.dtype | None - NumPy dtype of the value (or of the pointed-to elements); ``None`` for - ``void*``. + NumPy dtype of the value (or of the pointed-to elements; the structured + dtype for a struct); ``None`` for ``void*``. pointer : bool Whether the parameter is a pointer (a device array). + struct : CudaStruct | None + The struct type, for a struct passed by value. """ name: str ctype: str dtype: np.dtype | None pointer: bool + struct: CudaStruct | None = None # C types (after removing qualifiers) and their NumPy dtypes. ``long`` is 64 bit, @@ -133,27 +149,69 @@ class CudaParameter(NamedTuple): "complex": np.complex128, } +# The C type used for each NumPy dtype, see ctype_of() +_CTYPE_OF = { + np.dtype(np.bool_): "bool", + np.dtype(np.int8): "signed char", + np.dtype(np.uint8): "unsigned char", + np.dtype(np.int16): "short", + np.dtype(np.uint16): "unsigned short", + np.dtype(np.int32): "int", + np.dtype(np.uint32): "unsigned int", + np.dtype(np.int64): "long long", + np.dtype(np.uint64): "unsigned long long", + np.dtype(np.float32): "float", + np.dtype(np.float64): "double", + np.dtype(np.complex64): "complex", + np.dtype(np.complex128): "complex", +} + _QUALIFIERS = {"const", "volatile", "__restrict__", "__restrict", "restrict"} _COMPLEX = re.compile(r"(?:(?:thrust|cuda::std)::)?complex\s*<\s*(float|double)\s*>") _TOKEN = re.compile(r"complex<(?:float|double)>|[A-Za-z_]\w*|\*|\[\s*\]") +def ctype_of(dtype: Any) -> str: + """The C type of a NumPy dtype, e.g. ``ctype_of(np.float64) == "double"``. + + Useful to generate CUDA source or template arguments for a given dtype. + Complex dtypes map to ``complex``/``complex`` (include + ```` in the source). + """ + try: + return _CTYPE_OF[np.dtype(dtype)] + except (KeyError, TypeError): + raise ValueError(f"no C type for dtype {dtype!r}") from None + + def _strip_comments(source: str) -> str: source = re.sub(r"/\*.*?\*/", " ", source, flags=re.DOTALL) return re.sub(r"//[^\n]*", " ", source) -def _parse_parameter(text: str) -> CudaParameter: +def _parse_parameter( + text: str, structs: dict[str, CudaStruct] | None = None +) -> CudaParameter: text = _COMPLEX.sub(lambda m: f"complex<{m.group(1)}>", text) tokens = _TOKEN.findall(text) pointers = sum(1 for t in tokens if t == "*" or t.startswith("[")) words = [t for t in tokens if t != "*" and not t.startswith("[")] words = [t for t in words if t not in _QUALIFIERS] + if words[:1] == ["struct"]: + words = words[1:] if len(words) < 2: raise ValueError(f"cannot parse the kernel parameter {text.strip()!r}") name, ctype = words[-1], " ".join(words[:-1]) + if structs and ctype in structs: + if pointers: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: structs can " + "only be passed by value" + ) + struct = structs[ctype] + return CudaParameter(name, ctype, struct.dtype, False, struct) if ctype == "void" and pointers == 1: return CudaParameter(name, ctype, None, True) if pointers > 1 or ctype not in _CTYPES: @@ -164,7 +222,41 @@ def _parse_parameter(text: str) -> CudaParameter: return CudaParameter(name, ctype, np.dtype(_CTYPES[ctype]), pointers == 1) -def parse_cuda_signature(source: str, name: str) -> tuple[CudaParameter, ...]: +def _split_top_level(text: str) -> list[str]: + """Split at commas that are not inside ``<...>`` (e.g. ``complex``).""" + parts, depth, current = [], 0, [] + for char in text: + if char == "<": + depth += 1 + elif char == ">": + depth -= 1 + elif char == "," and depth == 0: + parts.append("".join(current)) + current = [] + continue + current.append(char) + parts.append("".join(current)) + return parts + + +def _template_arg(value: Any) -> str: + """A template argument as C++ source: a C type for dtypes, else a literal.""" + if isinstance(value, str): + return value + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, np.integer)): + return str(int(value)) + return ctype_of(value) + + +def parse_cuda_signature( + source: str, + name: str, + *, + structs: Iterable[CudaStruct] = (), + template_args: Sequence[Any] | None = None, +) -> tuple[CudaParameter, ...]: """Parse the parameters of the ``__global__`` function `name` in `source`. Parameters @@ -173,6 +265,13 @@ def parse_cuda_signature(source: str, name: str) -> tuple[CudaParameter, ...]: CUDA C source code. name : str Name of the ``__global__`` function. + structs : Iterable[CudaStruct] + Struct types that may appear as parameters (passed by value). If the + source defines a struct of the same name, its fields must match. + template_args : Sequence | None + Template arguments, if `name` is a function template: C types (or NumPy + dtypes, see :func:`ctype_of`) for type parameters, integers or bools for + non-type parameters. They are substituted into the parameter list. Returns ------- @@ -182,11 +281,18 @@ def parse_cuda_signature(source: str, name: str) -> tuple[CudaParameter, ...]: Raises ------ ValueError - If there is no such function, or a parameter has a type that cannot be - checked (e.g. a template parameter, a macro or a pointer to pointer). + If there is no such function, a template is used without (the right + number of) `template_args`, a struct definition in the source does not + match its :class:`CudaStruct`, or a parameter has a type that cannot be + checked (e.g. a macro or a pointer to pointer). """ code = _strip_comments(source) - match = re.search(r"__global__\s+void\s+" + re.escape(name) + r"\s*\(", code) + structs = {s.name: s for s in structs} + for struct in structs.values(): + struct.check_source(code) + + pattern = r"(?:template\s*<(?P