From d26dc3a17506b09fed56c475665aec4c3847dc76 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 1 Oct 2026 21:44:41 +0200 Subject: [PATCH 1/9] Repack CudaStructArguments when a field changes --- CHANGELOG.md | 3 +- docs/source/api.md | 10 ++- docs/source/kernels/arguments.md | 29 ++++++- docs/source/kernels/debugging.md | 5 +- src/cunumpy/LLM_GUIDE.md | 5 +- src/cunumpy/cuda_kernel.py | 126 +++++++++++++++++++++++++++---- tests/unit/test_cuda_kernel.py | 111 ++++++++++++++++++++++++++- 7 files changed, 261 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f96452..c75c833 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,8 @@ 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.CudaStructArguments`: Base class for argument objects that are passed to CUDA kernels as one C struct (the class form of `CudaStruct`). A subclass sets `struct_name` and `fields`, stores each field as an attribute and calls `pack()`; the `CudaStruct` is built once per class (`cls.struct`), `__cuda_args__()` returns the packed struct, and copies and unpickled objects are packed again from their own arrays. +- `xp.CudaStructArguments`: Base class for argument objects that are passed to CUDA kernels as one C struct (the class form of `CudaStruct`). A subclass sets `struct_name` and `fields`, stores each field as an attribute and calls `pack()`; the `CudaStruct` is built once per class (`cls.struct`), and `__cuda_args__()` returns the packed struct. The struct is packed again automatically at the next use when a field attribute changed (another array address, shape or strides, or another scalar value), so fields can be properties that follow an owner's resized arrays; copies and unpickled objects are packed again from their own arrays. +- Debug mode no longer synchronizes after a launch while the stream is being captured into a CUDA graph, which would invalidate the capture. - `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. diff --git a/docs/source/api.md b/docs/source/api.md index 4ab5e54..f363a27 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -1061,8 +1061,14 @@ built once per subclass when the class is defined and is the class attribute * `pack()` packs the field attributes, with the checks of `CudaStruct` (C-contiguous CuPy arrays of the declared dtype, range-checked scalars). A - field without an attribute raises `AttributeError`. Call it again after - replacing an array attribute: the packed struct holds device addresses. + field without an attribute raises `AttributeError`. +* The packed struct always matches the current attributes: at every use + (`packed`, `__cuda_args__()`, so at every launch) the device address, shape + and strides of each array field and the value of each scalar field are + compared with what was packed, and the struct is packed again if anything + changed. Fields may be properties that read an owner's current arrays, so a + resized array is picked up at the next launch. Calling `pack()` again is + never needed. * `packed` is the packed struct (`numpy.void`); `__cuda_args__()` returns `(packed,)`, so a `CudaKernel` receives the struct. * Copies (`copy.copy`, `copy.deepcopy`) and unpickled objects are packed again diff --git a/docs/source/kernels/arguments.md b/docs/source/kernels/arguments.md index 973ecdf..e890d52 100644 --- a/docs/source/kernels/arguments.md +++ b/docs/source/kernels/arguments.md @@ -170,8 +170,33 @@ push(args, dt, n_threads=args.n_markers) * The `CudaStruct` is built when the class is defined (`CudaMarkerArguments.struct`), so a bad field type raises at import, not at the first launch. -* `pack()` checks every field like `CudaStruct` does. Call it again after - replacing an array attribute, e.g. after the marker array was resized. +* `pack()` checks every field like `CudaStruct` does. +* The struct is packed again automatically, at the next launch, when a field + attribute changed: another array (address, shape or strides) or another + scalar value. Make the fields properties when the arrays belong to another + object that replaces them, e.g. a particle container that resizes its marker + array: + + ```python + class CudaMarkerArguments(xp.CudaStructArguments): + struct_name = "MarkerArgs" + fields = (("markers", "Array2D"), ("n_markers", "int")) + + def __init__(self, particles): + self._particles = particles + self.pack() + + @property + def markers(self): + return self._particles.markers + + @property + def n_markers(self): + return self._particles.markers.shape[0] + ``` + + Without the properties, the object keeps the old array alive and kernels + keep working on it: assign the new array to the attribute instead. * Copies and unpickled objects are packed again from their own arrays, so a `deepcopy` never points at the device memory of the original. * When the same call site must also reach a host kernel, pair it with the host diff --git a/docs/source/kernels/debugging.md b/docs/source/kernels/debugging.md index e952e69..4f0e143 100644 --- a/docs/source/kernels/debugging.md +++ b/docs/source/kernels/debugging.md @@ -30,7 +30,10 @@ In debug mode a `CudaKernel`: and shape, then traps); * synchronizes after every launch, so a failure raises at the launch that caused it, as a `RuntimeError` naming the kernel and its grid and block, with - the CUDA error chained. + the CUDA error chained. While the stream is being captured into a CUDA graph + (`stream.begin_capture()`), the synchronization is skipped, since it would + invalidate the capture; errors of captured kernels surface when the graph is + launched, so synchronize after `graph.launch()` to see them. ```python with xp.cuda_debug(): diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 95b50f2..8ba1aed 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -248,14 +248,15 @@ class A(xp.CudaStructArguments): # the struct as a class; A.struct is th fields = (("x", "double*"), ("n", "int")) def __init__(self, x): self.x, self.n = x, x.shape[0] - self.pack() # again after replacing an array; copies repack + self.pack() # repacks itself when a field changes; copies repack xp.CudaKernel(S.declaration + src, "k", structs=[S]) xp.resolve_host_args(args, kwargs) ``` Only top-level arguments are resolved. Cache both forms lazily and invalidate them when the underlying arrays are replaced. A packed struct holds device -addresses: re-pack after replacing an array. +addresses: re-pack after replacing an array (a `CudaStructArguments` does this +itself at the next launch; make its fields properties to follow an owner's arrays). `DeviceMirror`: diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index 4a4c65f..93442c7 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -1168,9 +1168,15 @@ class CudaStructArguments(CudaArguments): argument: pointers need C-contiguous CuPy arrays of the declared dtype (never copied), scalars are range-checked and cast. - The packed struct holds device addresses. Call :meth:`pack` again after - replacing an array attribute. Copies (``copy.copy``, ``copy.deepcopy``) - and unpickled objects are packed again from their own arrays. + The packed struct always reflects the current field attributes: at every + use (:attr:`packed`, :meth:`__cuda_args__`, so at every kernel launch) the + device address, shape and strides of every array field and the value of + every scalar field are compared with those that were packed, and the + struct is packed again if any changed. A field may therefore be a + property that reads the owner's current array, so that resizing the + owner's arrays never leaves the struct pointing at freed device memory + (see the second example). Copies (``copy.copy``, ``copy.deepcopy``) and + unpickled objects are packed again from their own arrays. A subclass that sets neither :attr:`struct_name` nor :attr:`fields` is an intermediate base class; a subclass that sets only one of them raises @@ -1209,6 +1215,25 @@ class CudaStructArguments(CudaArguments): }; >>> push = CudaKernel(source, "push", structs=[MarkerArguments.struct]) # doctest: +SKIP >>> push(MarkerArguments(markers, valid), 0.1, n_threads=markers.shape[0]) # doctest: +SKIP + + Fields as properties follow the arrays of an owner object, also after the + owner replaced them (e.g. when it resized its marker array): + + >>> class ParticleArguments(CudaStructArguments): + ... struct_name = "ParticleArgs" + ... fields = (("markers", "Array2D"), ("n_markers", "int")) + ... + ... def __init__(self, particles): + ... self._particles = particles + ... self.pack() + ... + ... @property + ... def markers(self): + ... return self._particles.markers + ... + ... @property + ... def n_markers(self): + ... return self._particles.markers.shape[0] """ struct_name: str @@ -1230,8 +1255,10 @@ def __init_subclass__(cls, **kwargs: Any) -> None: def pack(self) -> None: """Pack the field attributes into the struct. - Called at the end of the constructor, and again after an array - attribute has been replaced. + Called at the end of the constructor, so that invalid field values + raise there. Afterwards the struct is packed again automatically when + a field changes (see the class documentation); calling this method + again is never needed, but harmless. Raises ------ @@ -1246,6 +1273,12 @@ def pack(self) -> None: raise TypeError( f"{type(self).__qualname__} does not define struct_name and fields" ) + values = self._field_values(struct) + self._struct_value = struct(**values) + self._packed_state = _field_state(struct, values) + + def _field_values(self, struct: CudaStruct) -> dict[str, Any]: + """The current value of every field attribute, by field name.""" values = {} for field in struct.fields: try: @@ -1255,13 +1288,24 @@ def pack(self) -> None: f"{type(self).__qualname__} has no attribute {field.name!r} " f"for the field of struct {struct.name}" ) from None - self._struct_value = struct(**values) + return values @property def packed(self) -> np.void: - """The packed struct, with the memory layout of the C struct; packs on first use.""" - if self.__dict__.get("_struct_value") is None: + """The packed struct, with the memory layout of the C struct. + + Packed on first use, and again whenever a field attribute changed + since the last packing (a different array, or a different scalar). + """ + value = self.__dict__.get("_struct_value") + if value is None: self.pack() + return self._struct_value.packed + struct = value.struct + values = self._field_values(struct) + if _field_state(struct, values) != self._packed_state: + self._struct_value = struct(**values) + self._packed_state = _field_state(struct, values) return self._struct_value.packed def __cuda_args__(self) -> tuple[np.void]: @@ -1273,6 +1317,7 @@ def __getstate__(self) -> dict[str, Any]: # which a copy (or another process) does not share state = self.__dict__.copy() state.pop("_struct_value", None) + state.pop("_packed_state", None) return state def __setstate__(self, state: dict[str, Any]) -> None: @@ -1280,6 +1325,32 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.pack() +def _field_state(struct: CudaStruct, values: Mapping[str, Any]) -> tuple[Any, ...]: + """What the packed struct depends on, to detect changed field attributes. + + For an array field the device address, shape and strides (the address + alone for a pointer field); for a scalar field its value. A value that is + neither (e.g. a host array in a pointer field) is identified by its id, + so replacing it triggers a repack, which then raises the type error. + """ + state = [] + for field in struct.fields: + value = values[field.name] + if field.pointer or field.view_ndim is not None: + ptr = getattr(getattr(value, "data", None), "ptr", None) + if ptr is None: + state.append(("id", id(value))) + elif field.pointer: + state.append(ptr) + else: + state.append((ptr, tuple(value.shape), tuple(value.strides))) + elif isinstance(value, (bool, int, float, complex, np.generic)): + state.append((type(value), value)) + else: + state.append(("id", id(value))) + return tuple(state) + + def _as_shape(value: int | Sequence[int], what: str) -> tuple[int, ...]: shape = (value,) if isinstance(value, (int, np.integer)) else tuple(value) if not 1 <= len(shape) <= 3: @@ -1769,23 +1840,46 @@ def __call__( def _synchronize_after_launch( self, stream: Any, grid: tuple[int, ...], block: tuple[int, ...] ) -> None: - """Wait for the launch and re-raise a CUDA error naming this kernel.""" - import cupy as cp + """Wait for the launch and re-raise a CUDA error naming this kernel. + Skipped while the stream is being captured into a CUDA graph: the + launch is only recorded then, and synchronizing would invalidate the + capture. Errors of the captured kernels surface when the graph is + launched (in debug mode, synchronize after ``graph.launch()``). + """ + if stream is None: + import cupy as cp + + stream = cp.cuda.get_current_stream() + if _is_capturing(stream): + return try: - if stream is None: - stream = cp.cuda.get_current_stream() stream.synchronize() - except ( - cp.cuda.runtime.CUDARuntimeError, - cp.cuda.driver.CUDADriverError, - ) as error: + except Exception as error: + import cupy as cp + + if not isinstance( + error, + (cp.cuda.runtime.CUDARuntimeError, cp.cuda.driver.CUDADriverError), + ): + raise raise RuntimeError( f"CUDA error after launching kernel {self.expression!r} with " f"grid {grid} and block {block}: {error}" ) from error +def _is_capturing(stream: Any) -> bool: + """Whether `stream` is being captured into a CUDA graph.""" + is_capturing = getattr(stream, "is_capturing", None) + if is_capturing is None: + return False + try: + return bool(is_capturing()) + except Exception: # noqa: BLE001 -- e.g. the legacy null stream, which cannot capture + return False + + class CudaKernelVariants: """Kernels generated per variant (e.g. per dtype and dimension), compiled once. diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 592e13b..f8c6e70 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -1710,12 +1710,71 @@ class Derived(ParticleArguments): # inherits the struct of its parent assert Derived.struct is ParticleArguments.struct -def test_struct_arguments_repack_after_replacing_an_array(): +def test_struct_arguments_repack_when_a_field_changes(): args = _particle_arguments(ptr=0x100) + packed = args.packed + assert args.packed is packed # nothing changed: no repacking + args.x = FakeDeviceArray(np.float64, ptr=0x200, shape=(3,)) - assert args.packed["x"] == 0x100 # still the old address - args.pack() - assert args.packed["x"] == 0x200 + assert args.packed["x"] == 0x200 # a new array: repacked at the next use + assert args.__cuda_args__() == (args.packed,) + + args.charge = 5.0 + assert args.packed["charge"] == 5.0 # a scalar changed: repacked too + args.charge = 5 # equal value of another type: repacked, same result + assert args.packed["charge"] == 5.0 + + args.x = np.zeros(3) # an invalid value raises at the next use + with pytest.raises(TypeError, match="must be a CuPy array"): + args.packed # noqa: B018 + + +class Owner: + """An object owning a marker array that it replaces when it grows.""" + + def __init__(self, n, ptr=0x100): + self.markers = FakeDeviceArray(np.float64, ptr=ptr, shape=(n, 4)) + + def grow(self, n, ptr): + self.markers = FakeDeviceArray(np.float64, ptr=ptr, shape=(n, 4)) + + +class OwnerArguments(xp.CudaStructArguments): + struct_name = "OwnerArgs" + fields = (("markers", "Array2D"), ("n_markers", "int")) + + def __init__(self, owner): + self._owner = owner + self.pack() + + @property + def markers(self): + return self._owner.markers + + @property + def n_markers(self): + return self._owner.markers.shape[0] + + +def test_struct_arguments_follow_the_arrays_of_an_owner(): + owner = Owner(3) + args = OwnerArguments(owner) + assert args.packed["markers"]["data"] == 0x100 + assert args.packed["n_markers"] == 3 + + owner.grow(8, ptr=0x900) + assert args.packed["markers"]["data"] == 0x900 + assert args.packed["markers"]["shape"].tolist() == [8, 4] + assert args.packed["n_markers"] == 8 + + +def test_struct_arguments_repack_when_a_view_changes_shape(): + # a view of the same allocation with fewer rows: same address, new shape + owner = Owner(8, ptr=0x100) + args = OwnerArguments(owner) + owner.markers = FakeDeviceArray(np.float64, ptr=0x100, shape=(5, 4)) + assert args.packed["markers"]["shape"].tolist() == [5, 4] + assert args.packed["n_markers"] == 5 def test_struct_arguments_are_packed_again_when_copied(): @@ -1772,3 +1831,47 @@ def test_struct_arguments_on_gpu(): assert int(size.get()[0]) == ParticleArguments.struct.dtype.itemsize assert out.get().tolist() == [n, 2.0, 42.0, 1.5] assert cp.all(args.x[1::2] == 1.0) and cp.all(args.x[::2] == 0.0) + + +def test_debug_synchronization_is_skipped_while_capturing(): + kernel = CudaKernel(AXPY, "axpy") + calls = [] + + class Stream: + def __init__(self, capturing): + self.capturing = capturing + + def is_capturing(self): + return self.capturing + + def synchronize(self): + calls.append(self.capturing) + + class LegacyStream(Stream): + def is_capturing(self): + raise RuntimeError("not supported on the legacy stream") + + kernel._synchronize_after_launch(Stream(True), (1,), (128,)) + assert calls == [] + kernel._synchronize_after_launch(Stream(False), (1,), (128,)) + kernel._synchronize_after_launch(LegacyStream(False), (1,), (128,)) + assert calls == [False, False] + + +def test_debug_kernel_in_a_cuda_graph(): + _skip_without_cupy() + import cupy as cp + + n = 1000 + x, y = cp.ones(n), cp.zeros(n) + kernel = CudaKernel(AXPY, "axpy", debug=True) + kernel(2.0, x, y, n, n_threads=n) # compile outside the capture + stream = cp.cuda.Stream(non_blocking=True) + with stream: + stream.begin_capture() + kernel(2.0, x, y, n, n_threads=n, stream=stream) + graph = stream.end_capture() + graph.launch(stream) + graph.launch(stream) + stream.synchronize() + assert cp.all(y == 6.0) From 326460e33f424d6ce3c94cad4ed946c1c149c7b5 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 1 Oct 2026 23:36:08 +0200 Subject: [PATCH 2/9] Verify layout and parrity naming bug --- CHANGELOG.md | 4 ++ docs/source/api.md | 9 +++ docs/source/kernels/arguments.md | 15 +++++ docs/source/kernels/testing.md | 5 +- src/cunumpy/LLM_GUIDE.md | 1 + src/cunumpy/cuda_kernel.py | 109 +++++++++++++++++++++++++++++++ src/cunumpy/testing.py | 20 +++++- tests/unit/test_cuda_kernel.py | 55 ++++++++++++++++ tests/unit/test_testing.py | 50 ++++++++++++++ 9 files changed, 266 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c75c833..47f8b21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- `xp.testing.assert_kernels_agree` reads `CudaStructArguments` objects and struct values through their struct fields, so their arrays get the same names as the attributes of the host argument object (before, arrays behind properties were named after the private attribute holding the owner, and the comparison failed with "do not have the same array arguments"). + ### Removed - Support for Python 3.8 and 3.9 (both end-of-life); `cunumpy` now requires Python 3.10 or newer. @@ -27,6 +30,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `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.CudaStructArguments`: Base class for argument objects that are passed to CUDA kernels as one C struct (the class form of `CudaStruct`). A subclass sets `struct_name` and `fields`, stores each field as an attribute and calls `pack()`; the `CudaStruct` is built once per class (`cls.struct`), and `__cuda_args__()` returns the packed struct. The struct is packed again automatically at the next use when a field attribute changed (another array address, shape or strides, or another scalar value), so fields can be properties that follow an owner's resized arrays; copies and unpickled objects are packed again from their own arrays. - Debug mode no longer synchronizes after a launch while the stream is being captured into a CUDA graph, which would invalidate the capture. +- `CudaStruct.verify_layout(include=None, *, include_dirs=(), options=())` and `CudaStruct.layout_source()`: Compile and run a one-thread kernel that reports `sizeof`, `alignof` and the field offsets of the struct as the CUDA compiler lays it out, and raise `ValueError` if they differ from the NumPy `dtype` that values are packed into; with `include` the struct is defined by a header instead of `declaration`. - `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. diff --git a/docs/source/api.md b/docs/source/api.md index f363a27..b6a078b 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -967,6 +967,15 @@ array views") are supported. * `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. +* `verify_layout(include=None, *, include_dirs=(), options=())` (needs CuPy): + compiles and runs a one-thread kernel that reports `sizeof`, `alignof` and + every field offset as the CUDA compiler lays the struct out, and raises + `ValueError` listing the differences from `dtype`. Returns the measured + layout as a dict. With `include` (a header file name or `#include` line, + found in `include_dirs`) the struct is defined by that header instead of + `declaration`, which checks a hand-written or generated header. Call it once + per struct in a GPU test, and on every new platform (e.g. ROCm). + `layout_source(include=None)` returns the kernel source. * Calling the struct with keyword arguments, one per field, packs the values: pointer fields take C-contiguous CuPy arrays of the declared dtype (never copied), array view fields take CuPy arrays of the declared dtype and number diff --git a/docs/source/kernels/arguments.md b/docs/source/kernels/arguments.md index e890d52..8adb91e 100644 --- a/docs/source/kernels/arguments.md +++ b/docs/source/kernels/arguments.md @@ -203,6 +203,21 @@ push(args, dt, n_threads=args.n_markers) argument object in a `KernelArguments`: `__host_args__()` returns the host object, `__cuda_args__()` returns `cuda_args.__cuda_args__()`. +### Check the layout against the compiler + +Values are packed with the NumPy dtype of the struct. If the compiler lays the +struct out differently (a hand-edited header, a `#pragma pack`, another +compiler such as hiprtc), kernels read fields at the wrong offsets without any +error. One GPU test per struct catches it: + +```python +def test_marker_args_layout(): + CudaMarkerArguments.struct.verify_layout() # the Python declaration + CudaMarkerArguments.struct.verify_layout( # the committed header + "marker_args.cuh", include_dirs=["kernels"] + ) +``` + ### Generate the struct from the host argument class If the host kernels already use an annotated argument class (Pyccel style), the diff --git a/docs/source/kernels/testing.md b/docs/source/kernels/testing.md index 145a168..672d174 100644 --- a/docs/source/kernels/testing.md +++ b/docs/source/kernels/testing.md @@ -98,7 +98,10 @@ Things to know: * **Which arguments are compared**: `outputs=(2,)` selects them by index; otherwise the host kernel's declared `outputs` are used, and if there are none, every array argument. Arrays held by argument objects (one level deep, - e.g. a `CudaArguments` object or a list) are compared too. + e.g. a `CudaArguments` object or a list) are compared too. A + `CudaStructArguments` object is read through its struct fields, so its + arrays get the names of the host argument object's attributes + (`argument 0.markers`), also when the fields are properties. * **Tolerances**: the default `rtol=1e-12` suits deterministic kernels. Kernels with atomics or a different summation order need looser tolerances, for example `rtol=1e-10, atol=1e-14`. diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 8ba1aed..fc2ed04 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -240,6 +240,7 @@ class Args(xp.KernelArguments): # one object, host form + device form S = xp.CudaStruct("S", [("x", "double*"), ("n", "long long"), ("a", "Array2D")]) S.declaration; S.dtype; S.to_header(path); value = S(x=..., n=..., a=...) +S.verify_layout() # GPU test: compiler layout == S.dtype (also verify_layout("hdr.cuh")) S = xp.CudaStruct.from_signature(Cls.__init__, "S", int_type="long long") xp.write_cuda_header("args.cuh", [S1, S2]) diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index 93442c7..b60f0ce 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -1090,6 +1090,115 @@ def check_source(self, source: str) -> None: f"match its CudaStruct:\n{self.declaration}" ) + def layout_source(self, include: str | None = None) -> str: + """The CUDA source of the kernel used by :meth:`verify_layout`. + + The kernel ``cunumpy_layout_(unsigned long long* out)`` writes + ``sizeof``, ``alignof`` and the offset of every field (in field order) + of the struct, as the compiler lays it out. + + Parameters + ---------- + include : str | None + Header that defines the struct, as a file name or an ``#include`` + line; by default the struct is defined by :attr:`declaration`. + """ + if include is None: + lines = ['#include "cunumpy/array_view.cuh"'] if self.has_views else [] + lines.append(self.declaration) + else: + line = include.strip() + lines = [line if line.startswith("#") else f'#include "{line}"'] + offsets = "\n".join( + f" out[{i + 2}] = (unsigned long long)((const char*)&s.{f.name}" + " - (const char*)&s);" + for i, f in enumerate(self._fields) + ) + lines.append( + f'extern "C" __global__ void cunumpy_layout_{self._name}(' + "unsigned long long* out) {\n" + f" {self._name} s;\n" + f" out[0] = sizeof({self._name});\n" + f" out[1] = alignof({self._name});\n" + f"{offsets}\n" + "}\n" + ) + return "\n".join(lines) + + def verify_layout( + self, + include: str | None = None, + *, + include_dirs: Sequence[str | Path] = (), + options: Sequence[str] = (), + ) -> dict[str, int]: + """Check the struct layout of the CUDA compiler against :attr:`dtype`. + + Compiles and runs a one-thread kernel (:meth:`layout_source`) that + reports the size, alignment and field offsets of the struct, and + compares them with the NumPy dtype that values are packed into. A + difference means every kernel taking the struct reads some fields at + the wrong place, without any error. Run it once per struct in a GPU + test, especially with a hand-written or generated header (`include`) + and on a new compiler or platform (e.g. ROCm). + + Parameters + ---------- + include : str | None + Header that defines the struct (file name or ``#include`` line), + found in `include_dirs`; by default :attr:`declaration` is compiled. + include_dirs : Sequence[str | Path] + Directories searched for `include`. + options : Sequence[str] + Additional compiler options. + + Returns + ------- + dict[str, int] + ``"sizeof"``, ``"alignof"`` and the offset of each field, by name. + + Raises + ------ + ValueError + If the compiled layout differs from :attr:`dtype` (the message lists + each difference). + RuntimeError + If CuPy is not available. + """ + from .xp import cupy_available + + if not cupy_available(): + raise RuntimeError("verify_layout() compiles a CUDA kernel and needs CuPy") + import cupy as cp + + kernel = CudaKernel( + self.layout_source(include), + f"cunumpy_layout_{self._name}", + include_dirs=include_dirs, + options=options, + structs=[self], + block_size=1, + ) + out = cp.zeros(len(self._fields) + 2, dtype=cp.uint64) + kernel(out, n_threads=1) + measured = [int(v) for v in out.get()] + layout = {"sizeof": measured[0], "alignof": measured[1]} + layout.update({f.name: n for f, n in zip(self._fields, measured[2:])}) + + expected = {"sizeof": self._dtype.itemsize, "alignof": self._dtype.alignment} + expected.update({f.name: self._dtype.fields[f.name][1] for f in self._fields}) + differences = [ + f"{key}: compiler {layout[key]}, dtype {expected[key]}" + for key in expected + if layout[key] != expected[key] + ] + if differences: + raise ValueError( + f"the compiled layout of struct {self._name!r} differs from its " + "CudaStruct dtype:\n " + "\n ".join(differences) + ) + return layout + def __call__(self, **values: Any) -> CudaStructValue: """Pack values into the struct. diff --git a/src/cunumpy/testing.py b/src/cunumpy/testing.py index 4a602c1..f0e896c 100644 --- a/src/cunumpy/testing.py +++ b/src/cunumpy/testing.py @@ -47,6 +47,8 @@ def test_parity(name, kernel): from .cuda_kernel import ( CudaKernel, CudaParameter, + CudaStructArguments, + CudaStructValue, _parse_parameter, _split_top_level, _strip_comments, @@ -119,6 +121,18 @@ def _arrays_in(value: Any, name: str, found: dict[str, Any], depth: int) -> None found[name] = value elif depth == 0: return + elif isinstance(value, (CudaStructArguments, CudaStructValue)) and hasattr( + value, "struct" + ): + # by field name, like the attributes of the host argument object; the + # fields of a CudaStructArguments may be properties (not in vars()) + for field in value.struct.fields: + item = ( + value[field.name] + if isinstance(value, CudaStructValue) + else getattr(value, field.name) + ) + _arrays_in(item, f"{name}.{field.name}", found, depth - 1) elif isinstance(value, (tuple, list)): for i, item in enumerate(value): _arrays_in(item, f"{name}[{i}]", found, depth - 1) @@ -139,7 +153,11 @@ def _collect_arrays( level deep, in a tuple, list or dict argument or in the attributes of an argument object (e.g. a ``CudaArguments`` object), are named ``"argument []"`` or ``"argument ."``, and arrays in a - container attribute of an object ``"argument .[]"``. + container attribute of an object ``"argument .[]"``. A + :class:`~cunumpy.CudaStructArguments` object or a struct value is read + through its struct fields, ``"argument ."``, so that its arrays + get the names of the attributes of the host argument object it mirrors, + also when the fields are properties. """ indices = range(len(args)) if outputs is None else outputs found: dict[str, Any] = {} diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index f8c6e70..a4c932a 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -1875,3 +1875,58 @@ def test_debug_kernel_in_a_cuda_graph(): graph.launch(stream) stream.synchronize() assert cp.all(y == 6.0) + + +# --------------------------------------------------------------------------- +# struct layout checked against the compiler +# --------------------------------------------------------------------------- + + +def test_layout_source_reports_size_alignment_and_offsets(): + struct = CudaStruct( + "LayoutArgs", [("markers", "Array2D"), ("valid", "bool*"), ("n", "int")] + ) + source = struct.layout_source() + assert source.startswith('#include "cunumpy/array_view.cuh"') + assert struct.declaration in source + (param,) = parse_cuda_signature(source, "cunumpy_layout_LayoutArgs") + assert param.pointer and param.dtype == np.dtype(np.uint64) + assert "out[0] = sizeof(LayoutArgs);" in source + assert "out[1] = alignof(LayoutArgs);" in source + for i, name in enumerate(["markers", "valid", "n"]): + assert f"out[{i + 2}] = (unsigned long long)((const char*)&s.{name}" in source + + # a header instead of the declaration: the struct is not defined in the source + from_header = struct.layout_source("pkg/layout_args.cuh") + assert from_header.startswith('#include "pkg/layout_args.cuh"') + assert "struct LayoutArgs {" not in from_header + assert struct.layout_source("#include ").startswith( + "#include " + ) + # no views: no array_view include + assert "array_view" not in PARTICLES.layout_source() + + +def test_verify_layout_needs_cupy(monkeypatch): + monkeypatch.setattr(xp.xp, "cupy_available", lambda: False) + with pytest.raises(RuntimeError, match="needs CuPy"): + PARTICLES.verify_layout() + + +def test_verify_layout_on_gpu(tmp_path): + _skip_without_cupy() + layout = PARTICLES.verify_layout() + assert layout["sizeof"] == PARTICLES.dtype.itemsize + assert layout["charge"] == PARTICLES.dtype.fields["charge"][1] + + views = CudaStruct("Views", [("n", "int"), ("a", "Array2D"), ("b", "bool")]) + views.verify_layout() + + # a header that drifted from the Python definition + header = tmp_path / "drifted.cuh" + header.write_text( + '#include "cunumpy/array_view.cuh"\n' + "struct Views { int n; bool b; Array2D a; };\n" # b moved before a + ) + with pytest.raises(ValueError, match="differs from its CudaStruct dtype"): + views.verify_layout("drifted.cuh", include_dirs=[tmp_path]) diff --git a/tests/unit/test_testing.py b/tests/unit/test_testing.py index 14f2cf0..4c79813 100644 --- a/tests/unit/test_testing.py +++ b/tests/unit/test_testing.py @@ -304,3 +304,53 @@ def test_device_function_kernel_on_gpu(): spans = cp.empty(3, dtype=cp.int32) find_span(t, p, eta, spans, 3, n_threads=3) assert spans.get().tolist() == [2, 3, 3] + + +def test_struct_arguments_are_compared_by_field_name(): + """A CudaStructArguments object (fields may be properties) gets the host names.""" + from cunumpy.testing import _collect_arrays + + class Owner: + def __init__(self): + self.markers = np.zeros((3, 4)) + self.weights = np.ones(3) + + class HostArguments: # e.g. a Pyccel class + def __init__(self, owner): + self.markers = owner.markers + self.weights = owner.weights + self.n = 3 + + class DeviceArguments(xp.CudaStructArguments): + struct_name = "OwnerArgs" + fields = (("markers", "Array2D"), ("weights", "double*"), ("n", "int")) + + def __init__(self, owner): + self._owner = owner # not packed: no device arrays in this test + + @property + def markers(self): + return self._owner.markers + + @property + def weights(self): + return self._owner.weights + + n = 3 + + owner = Owner() + host = _collect_arrays((1.0, HostArguments(owner))) + device = _collect_arrays((1.0, DeviceArguments(owner))) + assert ( + sorted(host) == sorted(device) == ["argument 1.markers", "argument 1.weights"] + ) + assert device["argument 1.markers"] is owner.markers + + struct = DeviceArguments.struct + value = xp.CudaStructValue( + struct, np.zeros((), struct.dtype)[()], vars(owner) | {"n": 3} + ) + assert sorted(_collect_arrays((value,))) == [ + "argument 0.markers", + "argument 0.weights", + ] From d168359705c46e6bec420e7c6c369836bf969e85 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 1 Oct 2026 23:59:57 +0200 Subject: [PATCH 3/9] Add xp.scipy, xp.fuse, xp.petsc_vec and reduce.cuh --- CHANGELOG.md | 4 + docs/source/api.md | 106 ++++++++++++- docs/source/guides/portable-code.md | 5 +- docs/source/guides/solvers.md | 141 ++++++++++++++++++ docs/source/index.md | 1 + docs/source/kernels/accumulation.md | 2 + src/cunumpy/LLM_GUIDE.md | 4 + src/cunumpy/__init__.py | 6 + src/cunumpy/__init__.pyi | 3 + src/cunumpy/cuda/include/cunumpy/reduce.cuh | 154 +++++++++++++++++++ src/cunumpy/cuda_kernel.py | 7 +- src/cunumpy/fusion.py | 90 +++++++++++ src/cunumpy/petsc.py | 122 +++++++++++++++ src/cunumpy/scipy_backend.py | 157 ++++++++++++++++++++ tests/unit/test_cuda_kernel.py | 68 +++++++++ tests/unit/test_fusion.py | 78 ++++++++++ tests/unit/test_petsc.py | 133 +++++++++++++++++ tests/unit/test_scipy_backend.py | 104 +++++++++++++ 18 files changed, 1180 insertions(+), 5 deletions(-) create mode 100644 docs/source/guides/solvers.md create mode 100644 src/cunumpy/cuda/include/cunumpy/reduce.cuh create mode 100644 src/cunumpy/fusion.py create mode 100644 src/cunumpy/petsc.py create mode 100644 src/cunumpy/scipy_backend.py create mode 100644 tests/unit/test_fusion.py create mode 100644 tests/unit/test_petsc.py create mode 100644 tests/unit/test_scipy_backend.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 47f8b21..4127124 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `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.CudaStructArguments`: Base class for argument objects that are passed to CUDA kernels as one C struct (the class form of `CudaStruct`). A subclass sets `struct_name` and `fields`, stores each field as an attribute and calls `pack()`; the `CudaStruct` is built once per class (`cls.struct`), and `__cuda_args__()` returns the packed struct. The struct is packed again automatically at the next use when a field attribute changed (another array address, shape or strides, or another scalar value), so fields can be properties that follow an owner's resized arrays; copies and unpickled objects are packed again from their own arrays. - Debug mode no longer synchronizes after a launch while the stream is being captured into a CUDA graph, which would invalidate the capture. +- `xp.scipy`: SciPy for the active backend, `scipy` on NumPy and `cupyx.scipy` on CuPy, resolved at every access (`fft`, `fftpack`, `interpolate`, `linalg`, `ndimage`, `signal`, `sparse`, `sparse.csgraph`, `sparse.linalg`, `spatial`, `special`, `stats`). A name missing on the active backend raises `AttributeError` naming the backend; `available(name)` checks without raising; `resolve()` returns the module. SciPy stays an optional dependency. +- `xp.fuse`: Decorator that compiles an elementwise function into one kernel with `cupy.fuse` when it is called with CuPy arrays (with the CuPy backend active while tracing) and calls it unchanged otherwise. +- `xp.petsc_vec(array, comm=None)`: A `petsc4py` vector sharing the memory of a NumPy or CuPy array through DLPack (a CUDA or HIP vector for a CuPy array), never copying; raises instead of falling back to a host copy when petsc4py has no GPU support, or when the dtype or layout would need a copy. +- `cunumpy/reduce.cuh`: Warp and block reductions for CUDA kernels (`cunumpy_warp_sum/min/max`, `cunumpy_block_sum/min/max`, and `cunumpy_block_sum_to(out, v)` with one atomic add per block), for in-kernel diagnostics and per-block accumulation. - `CudaStruct.verify_layout(include=None, *, include_dirs=(), options=())` and `CudaStruct.layout_source()`: Compile and run a one-thread kernel that reports `sizeof`, `alignof` and the field offsets of the struct as the CUDA compiler lays it out, and raise `ValueError` if they differ from the NumPy `dtype` that values are packed into; with `include` the struct is defined by a header instead of `declaration`. - `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. diff --git a/docs/source/api.md b/docs/source/api.md index b6a078b..f8498b4 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -434,7 +434,7 @@ does not necessarily indicate a leak. ### `cuda_include_dir()` Returns the directory (as `str`) of the CUDA headers shipped with CuNumpy, -`cunumpy/array_view.cuh`, `cunumpy/atomic.cuh` and `cunumpy/index.cuh`. `CudaKernel` adds it to its NVRTC options as +`cunumpy/array_view.cuh`, `cunumpy/atomic.cuh`, `cunumpy/index.cuh` and `cunumpy/reduce.cuh`. `CudaKernel` adds it to its NVRTC options as `-I` automatically (and only once), so kernel sources can write `#include ` without configuration. Use it to pass the same headers to other compilers. @@ -1507,6 +1507,110 @@ The indexed helpers (also for `float`) address C-contiguous arrays of shape for `double` from compute capability 6.0 (sm_60) on; older devices use a compare-and-swap loop. +### `cunumpy/reduce.cuh` + +Warp- and block-level reductions for hand-written kernels: in-kernel +diagnostics (energy, momentum, total charge, the maximum velocity for a CFL +check) and combining values in a block before one atomic write. + +```c +#include + +T cunumpy_warp_sum(T v); T cunumpy_warp_min(T v); T cunumpy_warp_max(T v); +T cunumpy_block_sum(T v); T cunumpy_block_min(T v); T cunumpy_block_max(T v); +void cunumpy_block_sum_to(T* out, T v); // *out += block sum, one atomic per block +int cunumpy_block_thread(); // linear thread index in a 1D-3D block +int cunumpy_block_threads(); // threads per block +``` + +`T` is `int`, `unsigned`, `long long`, `unsigned long long`, `float` or +`double` (`block_sum_to`: `double`, `float`, `int`, `unsigned long long`). Every +thread gets the result. Rules: every thread of the block calls block +functions and all 32 lanes call warp functions (no early `return`; threads +without a value pass the identity, e.g. `0.0` for a sum); the block size is a +multiple of 32. The block functions use 32 values of static shared memory per +type and may be called several times in a kernel. + +```c +extern "C" __global__ void kinetic_energy(const double* v, long long n, + double mass, double* energy) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + double e = i < n ? 0.5 * mass * v[i] * v[i] : 0.0; + cunumpy_block_sum_to(energy, e); // zero *energy before the launch +} +``` + +## `scipy` + +```python +A = xp.scipy.sparse.csr_matrix((data, (rows, cols)), shape=(n, n)) +x, info = xp.scipy.sparse.linalg.cg(A, b) +rho_k = xp.scipy.fft.rfftn(rho) +``` + +SciPy for the active backend: `scipy` on NumPy, `cupyx.scipy` on CuPy. The +forwarded subpackages are those `cupyx.scipy` has (`SUBMODULES`): `fft`, +`fftpack`, `interpolate`, `linalg`, `ndimage`, `signal`, `sparse`, +`sparse.csgraph`, `sparse.linalg`, `spatial`, `special`, `stats`. Names are +looked up at every access, so a backend switch takes effect immediately +(Python caches the imports). Nothing is imported until a name is used; SciPy +is not a dependency of cunumpy. + +* A name missing on the active backend (`cupyx.scipy` covers part of SciPy) + raises `AttributeError` naming the backend. Keyword arguments can differ + too (SciPy's `cg(..., rtol=)` is `tol=` in CuPy). +* `xp.scipy.special.available("erfcx")` checks a name without raising. +* `xp.scipy.sparse.linalg.resolve()` returns the module itself. +* A missing SciPy (NumPy backend) or CuPy raises `ImportError` with the module + it needs. + +Sparse matrices assembled on the host move to the device once, with the +constructor of the device type: `xp.scipy.sparse.csr_matrix(host_matrix)` on +the CuPy backend copies a SciPy matrix; `matrix.get()` copies back. + +## `fuse(function=None, *, kernel_name=None)` + +```python +@xp.fuse +def pressure(rho, T, gamma): + return (gamma - 1.0) * rho * T +``` + +Compiles an elementwise function into one kernel with `cupy.fuse` when it is +called with a CuPy array (positional or keyword), with the CuPy backend active +while it is traced, so `xp.exp` etc. resolve to CuPy ufuncs; other calls run +the function as it is. The fused kernel is created on first use and reused. +`kernel_name` names it in profilers (default: the function name). The function +must be elementwise in the sense of `cupy.fuse`: arithmetic, comparisons, +ufuncs, `xp.where`, and supported reductions as the last operation; no Python +control flow on array values or indexing. Test the CuPy path: a function +`cupy.fuse` cannot trace raises at its first call with CuPy arrays. + +## `petsc_vec(array, comm=None)` + +```python +b_vec = xp.petsc_vec(b) # b: NumPy or CuPy array, shared, never copied +x_vec = xp.petsc_vec(x) +xp.synchronize() +ksp.solve(b_vec, x_vec) # PETSc writes into x +xp.synchronize() +``` + +A `petsc4py.PETSc.Vec` that shares the memory of a C-contiguous array of +`PETSc.ScalarType`, through DLPack: a `seq`/`mpi` vector for a NumPy array, a +`seqcuda`/`mpicuda` (or HIP) vector for a CuPy array. The vector keeps a +reference to the array. With several processes, `array` is this process's +part of the vector (`comm`, default `COMM_WORLD`). + +* Another dtype raises `TypeError` and a non-contiguous array `ValueError` + (both would need a copy). +* A CuPy array with a petsc4py built without CUDA/HIP raises `RuntimeError` + instead of PETSc working on a host copy. +* Synchronize between CuPy and PETSc work on the same memory (they may use + different streams). +* For the solve to stay on the GPU, the matrix must be a GPU type too + (`aijcusparse`, or `-mat_type aijcusparse -vec_type cuda`). + ## Version `xp.__version__` is the installed package version. When package metadata is diff --git a/docs/source/guides/portable-code.md b/docs/source/guides/portable-code.md index abc66ca..02e04e2 100644 --- a/docs/source/guides/portable-code.md +++ b/docs/source/guides/portable-code.md @@ -130,8 +130,9 @@ of the installed CuPy version. The common differences: device array on CuPy. Keep it as an array, or call `float()` and accept the synchronization. * SciPy functions accept only NumPy arrays; CuPy has its own `cupyx.scipy`. - Convert with `to_numpy()` at that boundary, or branch on - `xp.get_array_backend(a)`. + Use `xp.scipy`, which is the one for the active backend (see + [Solvers and fluid updates](solvers.md)); for functions `cupyx.scipy` lacks, + convert with `to_numpy()` at that boundary. * Plotting, HDF5 and most I/O libraries need host arrays: `to_numpy()` first. When a function genuinely needs a different implementation per backend, branch diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md new file mode 100644 index 0000000..058e7e3 --- /dev/null +++ b/docs/source/guides/solvers.md @@ -0,0 +1,141 @@ +# Solvers and fluid updates + +Particle pushes and deposits are only part of a plasma code. Field solves need +sparse matrices, iterative solvers and FFTs; fluid (MHD) updates are long +chains of pointwise operations; diagnostics reduce over all cells or +particles. This page covers the three tools for that: `xp.scipy`, `xp.fuse` +and `xp.petsc_vec`, plus in-kernel reductions with `cunumpy/reduce.cuh`. + +## SciPy on both backends: `xp.scipy` + +`xp.scipy` is SciPy on the NumPy backend and `cupyx.scipy` on the CuPy +backend, so solver code is written once: + +```python +import numpy as np + +import cunumpy as xp + + +def poisson_1d(rho, dx): + """Solve -phi'' = rho with phi = 0 at both ends (second-order FD).""" + n = rho.size + main = xp.full(n, 2.0 / dx**2) + off = xp.full(n - 1, -1.0 / dx**2) + A = xp.scipy.sparse.diags([off, main, off], [-1, 0, 1], format="csr") + phi, info = xp.scipy.sparse.linalg.cg(A, rho) + assert info == 0 + return phi + + +def periodic_poisson(rho, length): + """Solve -phi'' = rho on a periodic grid with FFTs.""" + n = rho.size + k = 2 * np.pi * xp.fft.rfftfreq(n, d=length / n) + rho_k = xp.scipy.fft.rfft(rho) + phi_k = xp.where(k == 0, 0.0, rho_k / xp.where(k == 0, 1.0, k**2)) + return xp.scipy.fft.irfft(phi_k, n) +``` + +* The subpackages forwarded are the ones `cupyx.scipy` has: `fft`, `linalg`, + `ndimage`, `signal`, `sparse`, `sparse.linalg`, `sparse.csgraph`, + `spatial`, `special`, `stats`, `interpolate`, `fftpack`. +* `cupyx.scipy` has only part of SciPy. A missing name raises + `AttributeError` saying which backend lacks it; check with + `xp.scipy.special.available("erfcx")` where a fallback is possible. +* Keyword arguments can differ between the two: SciPy's `cg` takes `rtol` + (since SciPy 1.12), CuPy's takes `tol`. Pass tolerances only after checking + both, or through a small wrapper. +* Assemble a sparse matrix once and move it to the device once: + `A = xp.scipy.sparse.csr_matrix(host_matrix)` on the CuPy backend copies a + SciPy matrix to the device. Do not rebuild matrices in the time loop. +* Special functions for distribution functions and cross sections + (`erf`, `erfc`, Bessel functions `i0`, `i1`, `k0`, ...) are in + `xp.scipy.special` on both backends. + +## Fused elementwise updates: `xp.fuse` + +On the GPU, `(gamma - 1) * (E - 0.5 * rho * u**2)` runs as five kernels, each +reading and writing a full temporary array. `xp.fuse` turns the whole +function into one kernel with `cupy.fuse` when it is called with CuPy arrays, +and calls it unchanged with NumPy arrays: + +```python +@xp.fuse +def pressure(rho, mom, energy, gamma): + u = mom / rho + return (gamma - 1.0) * (energy - 0.5 * rho * u * u) + + +@xp.fuse +def maxwellian(v, n, u, v_th): + return n / (xp.sqrt(2.0 * np.pi) * v_th) * xp.exp(-0.5 * ((v - u) / v_th) ** 2) + + +p = pressure(rho, mom, energy, 5.0 / 3.0) +``` + +Memory-bound chains like these typically get several times faster. The +function must be elementwise: arithmetic, comparisons, ufuncs (`xp.exp`, +`xp.sqrt`, ...), `xp.where`; no `if` on array values, no indexing. Test the +CuPy path (`cunumpy.testing.BACKENDS`): `cupy.fuse` reports a function it +cannot trace at the first call with CuPy arrays. + +## PETSc without copies: `xp.petsc_vec` + +When the field solve goes through PETSc, `xp.petsc_vec(array)` wraps an +array as a PETSc vector that uses the array's memory, a CUDA vector for a +CuPy array. The deposit writes into `rho`, PETSc reads it and writes `phi`, +the gather reads `phi`, and no data leaves the GPU: + +```python +from petsc4py import PETSc + +rho = xp.zeros(n_local) +phi = xp.zeros(n_local) +rho_vec, phi_vec = xp.petsc_vec(rho), xp.petsc_vec(phi) + +A = assemble_laplacian() # a PETSc Mat +A.setType("aijcusparse") # keep the matrix on the GPU too +ksp = PETSc.KSP().create() +ksp.setOperators(A) + +for step in range(n_steps): + deposit(rho) + xp.synchronize() # CuPy finished writing rho + ksp.solve(rho_vec, phi_vec) + xp.synchronize() # PETSc finished writing phi + gather(phi) +``` + +This needs a petsc4py built with CUDA (or HIP) support; otherwise +`petsc_vec` raises for CuPy arrays rather than let PETSc work on a host copy. +The array must be C-contiguous and of `PETSc.ScalarType`, since anything else +would need a copy. With MPI, each process passes its local part and the +vector's communicator. + +## Reductions inside kernels: `cunumpy/reduce.cuh` + +Outside kernels, `xp.sum` and friends are the right tool. Inside a +hand-written kernel, for a diagnostic computed in the same pass as the push, +or to combine values in a block before one atomic write, include the header: + +```c +#include + +extern "C" __global__ void push_and_energy(double* x, double* v, const double* E, + long long n, double qm_dt, double dt, + double half_m, double* energy) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + double e = 0.0; + if (i < n) { + v[i] += qm_dt * E[i]; + x[i] += dt * v[i]; + e = half_m * v[i] * v[i]; + } + cunumpy_block_sum_to(energy, e); // every thread calls it: no early return +} +``` + +The block size must be a multiple of 32, and every thread must reach the +reduction. See the API reference for the warp and block functions. diff --git a/docs/source/index.md b/docs/source/index.md index 54d3585..7e6a318 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -49,6 +49,7 @@ guides/portable-code guides/data-movement guides/gpu-devices guides/mpi +guides/solvers guides/profiling array-api-compat ``` diff --git a/docs/source/kernels/accumulation.md b/docs/source/kernels/accumulation.md index 4ccad9d..1359f50 100644 --- a/docs/source/kernels/accumulation.md +++ b/docs/source/kernels/accumulation.md @@ -111,5 +111,7 @@ host kernel writes into it directly, and `to_host()` does nothing. contents; then call `to_device()` first if the host side changed. * Atomics on one hot cell serialize. If most particles hit few cells, consider sorting particles by cell or accumulating per block in shared memory first. + For a single value (a total charge, an energy), `cunumpy_block_sum_to` from + `cunumpy/reduce.cuh` makes one atomic add per block instead of one per thread. * Floating-point atomics make the summation order non-deterministic. Results differ between runs in the last bits; compare with a tolerance in tests. diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index fc2ed04..19c41c3 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -68,6 +68,10 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | many kernels in a package, ported incrementally | `xp.KernelCatalog.from_package(__name__, missing_cuda="fallback")` | | group arrays/scalars into one kernel argument | `xp.CudaArguments` (device only), `xp.KernelArguments` (host object + device tuple), `xp.CudaStruct` (C struct), `xp.CudaStructArguments` (C struct as a class) | | CUDA struct from a Pyccel argument class | `xp.CudaStruct.from_signature(Cls.__init__, "Name")`, `xp.write_cuda_header(...)` | +| SciPy (sparse, sparse.linalg, fft, special, ndimage, ...) on either backend | `xp.scipy..` (SciPy or `cupyx.scipy`); `xp.scipy.special.available(name)` | +| chain of elementwise operations as one GPU kernel | `@xp.fuse` (`cupy.fuse` for CuPy arrays, plain call otherwise) | +| PETSc solve on device arrays without copies | `xp.petsc_vec(array)` (CUDA/HIP petsc4py for CuPy arrays); `xp.synchronize()` around PETSc calls | +| reduction inside a CUDA kernel (energy, max velocity) | ``: `cunumpy_block_sum_to(out, v)`, `cunumpy_block_min/max`, `cunumpy_warp_sum` | | kernel writes into a host buffer owned by another library | `xp.DeviceMirror(host_array)` + `` | | N-D indexing in CUDA, non-contiguous arrays | `Array1D`..`Array3D` params from `` | | one MPI rank per GPU | `bind_local_device()` → `from mpi4py import MPI` → `require_cuda_aware_mpi()` → `synchronize_for_mpi(...)` before each call | diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 04d8b39..6d017ac 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -20,8 +20,11 @@ write_cuda_header, ) from .dispatch import Kernel, KernelCatalog +from .fusion import fuse from .kernel import KernelArguments, PyccelKernel, resolve_host_args from .mirror import DeviceMirror +from .petsc import petsc_vec +from .scipy_backend import scipy from .transfers import ( TransferCounter, TransferEvent, @@ -103,6 +106,7 @@ "default_float_dtype", "device_count", "free_memory", + "fuse", "get_array_backend", "get_array_module", "get_backend", @@ -117,11 +121,13 @@ "numpy_backend", "nvtx_range", "parse_cuda_signature", + "petsc_vec", "pin_memory", "require_cuda_aware_mpi", "resolve_host_args", "resolve_includes", "same_backend", + "scipy", "set_backend", "set_cuda_debug", "set_device", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index cf24ab4..879765e 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -22,8 +22,11 @@ from .cuda_kernel import parse_cuda_signature as parse_cuda_signature from .cuda_kernel import write_cuda_header as write_cuda_header from .dispatch import Kernel as Kernel from .dispatch import KernelCatalog as KernelCatalog +from .fusion import fuse as fuse from .kernel import PyccelKernel as PyccelKernel from .mirror import DeviceMirror as DeviceMirror +from .petsc import petsc_vec as petsc_vec +from .scipy_backend import scipy as scipy from .transfers import TransferCounter as TransferCounter from .transfers import TransferEvent as TransferEvent diff --git a/src/cunumpy/cuda/include/cunumpy/reduce.cuh b/src/cunumpy/cuda/include/cunumpy/reduce.cuh new file mode 100644 index 0000000..2371fa1 --- /dev/null +++ b/src/cunumpy/cuda/include/cunumpy/reduce.cuh @@ -0,0 +1,154 @@ +// cunumpy/reduce.cuh: warp- and block-level reductions for CUDA kernels. +// +// Diagnostics computed inside a kernel (kinetic energy, momentum, total +// charge, a maximum velocity for the CFL condition) and accumulation that +// combines values in a block before writing to global memory all need the +// same reduction: shuffles within each warp, then shared memory across the +// warps of a block. These helpers implement it once: +// +// #include +// +// extern "C" __global__ void kinetic_energy(const double* v, long long n, +// double mass, double* energy) { +// long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; +// double e = i < n ? 0.5 * mass * v[i] * v[i] : 0.0; // no early return +// cunumpy_block_sum_to(energy, e); // one atomic add per block +// } +// +// Rules: +// * Every thread of the block must call a block function, and all 32 lanes of +// the warp a warp function: do not return early. Threads without a value +// pass the identity (0 for a sum, the largest value for a minimum, ...). +// * The number of threads per block must be a multiple of 32 (CudaKernel's +// default block size is 128). +// * The result is returned to every thread (warp functions: every lane). +// * Supported types are those of the shuffle intrinsics: int, unsigned, +// long long, unsigned long long, float, double. +// +// The warp size is 32, as on all NVIDIA GPUs. +// +// The header directory is added to every CudaKernel's NVRTC options. + +#ifndef CUNUMPY_REDUCE_CUH +#define CUNUMPY_REDUCE_CUH + +#include "cunumpy/atomic.cuh" + +#define CUNUMPY_WARP_SIZE 32 +#define CUNUMPY_FULL_WARP_MASK 0xffffffffu + +// The binary operations; fill() gives the value of the lanes of the last +// step that have no partial result of their own. +struct cunumpy_sum_op { + template + __device__ __forceinline__ T operator()(T a, T b) const { return a + b; } + template + __device__ __forceinline__ static T fill(const T*) { return T(0); } +}; + +struct cunumpy_min_op { + template + __device__ __forceinline__ T operator()(T a, T b) const { return b < a ? b : a; } + template + __device__ __forceinline__ static T fill(const T* partial) { return partial[0]; } +}; + +struct cunumpy_max_op { + template + __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; } + template + __device__ __forceinline__ static T fill(const T* partial) { return partial[0]; } +}; + +// Linear index of the thread in its block, and the number of threads per block +// (for 1D to 3D blocks). +__device__ __forceinline__ int cunumpy_block_thread() +{ + return threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z); +} + +__device__ __forceinline__ int cunumpy_block_threads() +{ + return blockDim.x * blockDim.y * blockDim.z; +} + +// Reduce v over the 32 lanes of the warp; every lane gets the result. +template +__device__ __forceinline__ T cunumpy_warp_reduce(T v, Op op) +{ + for (int mask = CUNUMPY_WARP_SIZE / 2; mask > 0; mask /= 2) { + v = op(v, __shfl_xor_sync(CUNUMPY_FULL_WARP_MASK, v, mask)); + } + return v; +} + +template +__device__ __forceinline__ T cunumpy_warp_sum(T v) { return cunumpy_warp_reduce(v, cunumpy_sum_op()); } + +template +__device__ __forceinline__ T cunumpy_warp_min(T v) { return cunumpy_warp_reduce(v, cunumpy_min_op()); } + +template +__device__ __forceinline__ T cunumpy_warp_max(T v) { return cunumpy_warp_reduce(v, cunumpy_max_op()); } + +// Reduce v over all threads of the block; every thread gets the result. +// Uses 32 values of static shared memory per type and synchronizes the block +// (__syncthreads) three times; it may be called several times in a kernel. +template +__device__ T cunumpy_block_reduce(T v, Op op) +{ + __shared__ T partial[CUNUMPY_WARP_SIZE]; + const int thread = cunumpy_block_thread(); + const int lane = thread % CUNUMPY_WARP_SIZE; + const int warp = thread / CUNUMPY_WARP_SIZE; + const int n_warps = (cunumpy_block_threads() + CUNUMPY_WARP_SIZE - 1) / CUNUMPY_WARP_SIZE; + + v = cunumpy_warp_reduce(v, op); + __syncthreads(); // a previous call may still be reading partial + if (lane == 0) partial[warp] = v; + __syncthreads(); + if (warp == 0) { + v = lane < n_warps ? partial[lane] : Op::fill(partial); + v = cunumpy_warp_reduce(v, op); + if (lane == 0) partial[0] = v; + } + __syncthreads(); + return partial[0]; +} + +template +__device__ T cunumpy_block_sum(T v) { return cunumpy_block_reduce(v, cunumpy_sum_op()); } + +template +__device__ T cunumpy_block_min(T v) { return cunumpy_block_reduce(v, cunumpy_min_op()); } + +template +__device__ T cunumpy_block_max(T v) { return cunumpy_block_reduce(v, cunumpy_max_op()); } + +// *out += sum of v over the block, with one atomic add per block (thread 0). +// For a sum over the whole grid, zero *out before the launch. +__device__ __forceinline__ void cunumpy_block_sum_to(double* out, double v) +{ + const double total = cunumpy_block_sum(v); + if (cunumpy_block_thread() == 0) cunumpy_atomic_add(out, total); +} + +__device__ __forceinline__ void cunumpy_block_sum_to(float* out, float v) +{ + const float total = cunumpy_block_sum(v); + if (cunumpy_block_thread() == 0) cunumpy_atomic_add(out, total); +} + +__device__ __forceinline__ void cunumpy_block_sum_to(int* out, int v) +{ + const int total = cunumpy_block_sum(v); + if (cunumpy_block_thread() == 0) atomicAdd(out, total); +} + +__device__ __forceinline__ void cunumpy_block_sum_to(unsigned long long* out, unsigned long long v) +{ + const unsigned long long total = cunumpy_block_sum(v); + if (cunumpy_block_thread() == 0) atomicAdd(out, total); +} + +#endif // CUNUMPY_REDUCE_CUH diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index b60f0ce..2ee3364 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -93,7 +93,9 @@ def cuda_include_dir() -> str: can ``#include "cunumpy/array_view.cuh"`` (strided ``Array1D``, ``Array2D``, ``Array3D`` views passed by value) and ``#include "cunumpy/index.cuh"`` (thread-index and grid-stride macros such - as ``CUNUMPY_THREAD_1D(i, n)``). Pass it as ``-I`` to other compilers. + as ``CUNUMPY_THREAD_1D(i, n)``), ``#include "cunumpy/atomic.cuh"`` (atomic + adds) and ``#include "cunumpy/reduce.cuh"`` (warp and block reductions). + Pass it as ``-I`` to other compilers. """ return str(_CUDA_INCLUDE_DIR) @@ -1488,7 +1490,8 @@ class CudaKernel: found at compile time (see :meth:`compile_options`): ``#include "cunumpy/array_view.cuh"`` gives the ``Array1D`` to ``Array3D`` views, ``#include "cunumpy/index.cuh"`` the thread-index - macros, ``#include "cunumpy/atomic.cuh"`` atomic adds. + macros, ``#include "cunumpy/atomic.cuh"`` atomic adds, + ``#include "cunumpy/reduce.cuh"`` warp and block reductions. source_dir : str | Path | None Directory the source was read from (set by :meth:`from_file`), where ``#include "..."`` files are looked up first. diff --git a/src/cunumpy/fusion.py b/src/cunumpy/fusion.py new file mode 100644 index 0000000..4713deb --- /dev/null +++ b/src/cunumpy/fusion.py @@ -0,0 +1,90 @@ +"""``xp.fuse``: elementwise functions as one GPU kernel, plain calls on the host. + +A chain of elementwise operations (a pressure from density and temperature, +fluxes and limiters of a fluid update, a Maxwellian at many velocities) runs +as one kernel per operation on the GPU, each writing a temporary array: it is +limited by memory bandwidth. ``cupy.fuse`` compiles the whole chain into one +kernel. :func:`fuse` applies it when the function is called with CuPy arrays +and calls the function as it is otherwise, so the same code runs on both +backends:: + + @xp.fuse + def pressure(rho, T, gamma): + return (gamma - 1.0) * rho * T + + p = pressure(rho, T, 5.0 / 3.0) + +The function must be elementwise: arithmetic, comparisons, ufuncs such as +``xp.exp``, ``xp.sqrt``, ``xp.where`` and reductions that ``cupy.fuse`` +supports (``xp.sum`` as the last operation). On the CuPy path the CuPy +backend is active while the function is traced, so ``xp.*`` names resolve to +CuPy's ufuncs. Functions that ``cupy.fuse`` cannot trace (Python control +flow on array values, indexing, wrappers that are not ufuncs) raise when the +fused function is first called with CuPy arrays; test the CuPy path. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from typing import Any, TypeVar + +import array_api_compat + +from .xp import use_backend + +__all__ = ["fuse"] + +F = TypeVar("F", bound=Callable[..., Any]) + +# substituted in tests that have no GPU +_is_device_array = array_api_compat.is_cupy_array + + +def _cupy_fuse(function: Callable[..., Any], kernel_name: str | None) -> Any: + import cupy + + return cupy.fuse(kernel_name=kernel_name)(function) + + +def fuse( + function: F | None = None, *, kernel_name: str | None = None +) -> F | Callable[[F], F]: + """Fuse an elementwise function into one kernel when called with CuPy arrays. + + Usable as ``@xp.fuse`` or ``@xp.fuse(kernel_name="pressure")``. + + Parameters + ---------- + function : callable + The elementwise function. + kernel_name : str | None + Name of the generated kernel (shown by profilers); by default the + function's name. + + Returns + ------- + callable + A function with the same signature. If any positional or keyword + argument is a CuPy array, it calls ``cupy.fuse(function)`` (created on + first use, with the CuPy backend active); otherwise it calls + `function` itself. + """ + if function is None: + return lambda f: fuse(f, kernel_name=kernel_name) + + name = kernel_name if kernel_name is not None else function.__name__ + fused = None + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal fused + if not any(_is_device_array(a) for a in (*args, *kwargs.values())): + return function(*args, **kwargs) + with use_backend("cupy"): + if fused is None: + fused = _cupy_fuse(function, name) + return fused(*args, **kwargs) + + wrapper.__wrapped__ = function + return wrapper diff --git a/src/cunumpy/petsc.py b/src/cunumpy/petsc.py new file mode 100644 index 0000000..a5ca0a9 --- /dev/null +++ b/src/cunumpy/petsc.py @@ -0,0 +1,122 @@ +"""PETSc vectors sharing memory with NumPy or CuPy arrays (no copies). + +Field solvers often go through PETSc (KSP), while the rest of a GPU code keeps +its data in CuPy arrays. Copying the right-hand side to the host and the +solution back at every solve is the largest transfer of such a time step. +:func:`petsc_vec` wraps an array as a PETSc vector that uses the array's +memory: a host vector for a NumPy array, a CUDA (or HIP) vector for a CuPy +array, which needs a petsc4py built with CUDA (or HIP) support:: + + b = xp.zeros(n) # filled by the deposit kernel + phi = xp.zeros(n) # the solution, read by the gather kernel + b_vec, phi_vec = xp.petsc_vec(b), xp.petsc_vec(phi) + ... + xp.synchronize() # CuPy work on b done before PETSc reads it + ksp.solve(b_vec, phi_vec) # writes into phi + xp.synchronize() # PETSc done before CuPy reads phi + +For the solve to stay on the GPU, the matrix must be a GPU type as well +(``mat.setType("aijcusparse")``, or ``-mat_type aijcusparse -vec_type cuda`` +in the PETSc options); otherwise PETSc copies the vectors to the host for the +matrix products. + +petsc4py is imported only when :func:`petsc_vec` is called. +""" + +from __future__ import annotations + +from typing import Any + +import array_api_compat +import numpy as np + +__all__ = ["petsc_vec"] + +# substituted in tests that have no GPU +_is_device_array = array_api_compat.is_cupy_array + +_DEVICE_VEC_TYPES = ("cuda", "hip") + + +def _petsc() -> Any: + try: + from petsc4py import PETSc + except ImportError as error: + raise ImportError( + "xp.petsc_vec needs petsc4py (pip install petsc4py)" + ) from error + return PETSc + + +def petsc_vec(array: Any, comm: Any = None) -> Any: + """A PETSc vector that shares the memory of `array`. + + Parameters + ---------- + array : numpy.ndarray | cupy.ndarray + C-contiguous array of PETSc's scalar type (``PETSc.ScalarType``, + usually float64). A multi-dimensional array is seen by PETSc in + row-major order, as ``array.ravel()``. The array is never copied. + comm : mpi4py.MPI.Comm | PETSc.Comm | None + Communicator of the vector; PETSc's default (``COMM_WORLD``) if None. + With several processes, `array` is this process's part. + + Returns + ------- + petsc4py.PETSc.Vec + A sequential or MPI vector (``seq``/``mpi`` for a NumPy array, + ``seqcuda``/``mpicuda`` or ``seqhip``/``mpihip`` for a CuPy array). + It keeps a reference to `array`, which therefore stays alive as long + as the vector. + + Raises + ------ + TypeError + If `array` has another dtype than ``PETSc.ScalarType``, or is not an + array. + ValueError + If `array` is not C-contiguous. + RuntimeError + If `array` is a CuPy array and petsc4py has no CUDA or HIP support + (PETSc would otherwise silently work on a host copy). + ImportError + If petsc4py is not installed. + + Notes + ----- + PETSc and CuPy may run on different streams: synchronize + (:func:`cunumpy.synchronize`) before PETSc reads an array that CuPy wrote, + and before CuPy reads a vector that PETSc wrote. + """ + PETSc = _petsc() + device = _is_device_array(array) + if not (device or isinstance(array, np.ndarray)): + raise TypeError( + f"petsc_vec takes a NumPy or CuPy array, got {type(array).__name__}" + ) + scalar = np.dtype(PETSc.ScalarType) + if array.dtype != scalar: + raise TypeError( + f"petsc_vec needs an array of PETSc's scalar type {scalar}, got " + f"{array.dtype} (convert it once, outside the time loop)" + ) + if not array.flags.c_contiguous: + raise ValueError("petsc_vec needs a C-contiguous array (no copy is made)") + + try: + vec = PETSc.Vec().createWithDLPack(array, comm=comm) + except PETSc.Error as error: + if device: + raise RuntimeError( + "petsc4py could not wrap the CuPy array; it needs a PETSc built " + "with CUDA or HIP support (--with-cuda / --with-hip)" + ) from error + raise + if device and not any(t in vec.getType() for t in _DEVICE_VEC_TYPES): + vec.destroy() + raise RuntimeError( + f"petsc4py created a {vec.getType()!r} vector for a CuPy array; it " + "needs a PETSc built with CUDA or HIP support" + ) + vec.setAttr("cunumpy_array", array) # the vector does not own the memory + return vec diff --git a/src/cunumpy/scipy_backend.py b/src/cunumpy/scipy_backend.py new file mode 100644 index 0000000..94bb74b --- /dev/null +++ b/src/cunumpy/scipy_backend.py @@ -0,0 +1,157 @@ +"""SciPy for the active backend: ``xp.scipy`` is SciPy or ``cupyx.scipy``. + +Field solvers and fluid codes need more than array functions: sparse matrices +and their iterative solvers, FFTs, special functions, image filters. +``xp.scipy`` forwards to :mod:`scipy` on the NumPy backend and to +:mod:`cupyx.scipy` on the CuPy backend, so this code runs on both:: + + A = xp.scipy.sparse.csr_matrix((data, (rows, cols)), shape=(n, n)) + x, info = xp.scipy.sparse.linalg.cg(A, b) + phi_k = xp.scipy.fft.rfftn(rho) + f = xp.scipy.special.erf(v / v_th) + +The module is resolved at every attribute access, so switching the backend +(:func:`cunumpy.set_backend`) takes effect immediately; imports are cached by +Python, so the lookup is cheap. Neither SciPy nor CuPy is imported until a +name is used. + +``cupyx.scipy`` covers only part of SciPy. A name that the active backend's +module does not have raises ``AttributeError`` saying which backend lacks it; +:meth:`ScipyNamespace.available` checks a name without raising. Keyword +arguments can differ as well: SciPy's ``cg`` takes ``rtol`` since SciPy 1.12, +CuPy's still takes ``tol``. +""" + +from __future__ import annotations + +import importlib +from types import ModuleType +from typing import Any + +from .xp import get_backend + +__all__ = ["SUBMODULES", "ScipyNamespace", "scipy"] + +#: The SciPy subpackages forwarded by ``xp.scipy``: those that ``cupyx.scipy`` +#: provides as well. +SUBMODULES = ( + "fft", + "fftpack", + "interpolate", + "linalg", + "ndimage", + "signal", + "sparse", + "sparse.csgraph", + "sparse.linalg", + "spatial", + "special", + "stats", +) + +_ROOTS = {"numpy": "scipy", "cupy": "cupyx.scipy"} +_INSTALL = { + "numpy": "SciPy is not installed (pip install scipy)", + "cupy": "CuPy is not installed or does not provide it", +} + + +class ScipyNamespace: + """SciPy (sub)package of the active backend; see :mod:`cunumpy.scipy_backend`. + + Parameters + ---------- + path : str + Dotted path below the SciPy root, e.g. ``"sparse.linalg"``; empty for + the root itself. + """ + + def __init__(self, path: str = "") -> None: + if path and path not in SUBMODULES: + raise ValueError( + f"scipy.{path} is not forwarded; forwarded subpackages: " + f"{', '.join(SUBMODULES)}" + ) + self._path = path + self._children: dict[str, ScipyNamespace] = {} + + def __repr__(self) -> str: + return f"')}>" + + def _name(self, root: str) -> str: + return f"{root}.{self._path}" if self._path else root + + def resolve(self) -> ModuleType: + """The module this namespace stands for on the active backend. + + Returns + ------- + ModuleType + E.g. ``scipy.sparse.linalg`` or ``cupyx.scipy.sparse.linalg``. + + Raises + ------ + ImportError + If SciPy (NumPy backend) or CuPy (CuPy backend) is not installed. + """ + backend = get_backend() + name = self._name(_ROOTS[backend]) + try: + return importlib.import_module(name) + except ImportError as error: + raise ImportError( + f"xp.scipy on the {backend} backend needs {name}: {_INSTALL[backend]}" + ) from error + + def __getattr__(self, name: str) -> Any: + if name.startswith("__"): + raise AttributeError(name) + child = f"{self._path}.{name}" if self._path else name + if child in SUBMODULES: + if name not in self._children: + self._children[name] = ScipyNamespace(child) + return self._children[name] + module = self.resolve() + try: + return getattr(module, name) + except AttributeError: + other = "cupy" if get_backend() == "numpy" else "numpy" + raise AttributeError( + f"{module.__name__} has no attribute {name!r}: it is not available " + f"on the {get_backend()} backend (it may exist in " + f"{self._name(_ROOTS[other])})" + ) from None + + def available(self, name: str) -> bool: + """Whether `name` exists in this namespace on the active backend. + + Never raises; False also if the backend's SciPy is not installed. + """ + child = f"{self._path}.{name}" if self._path else name + if child in SUBMODULES: + try: + ScipyNamespace(child).resolve() + except ImportError: + return False + return True + try: + return hasattr(self.resolve(), name) + except ImportError: + return False + + def __dir__(self) -> list[str]: + prefix = f"{self._path}." if self._path else "" + children = [ + s[len(prefix) :] + for s in SUBMODULES + if s.startswith(prefix) and "." not in s[len(prefix) :] + ] + try: + names = dir(self.resolve()) + except ImportError: + names = [] + return sorted(set(names) | set(children) | {"available", "resolve"}) + + +#: SciPy of the active backend, available as ``xp.scipy``. +scipy = ScipyNamespace() diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index a4c932a..5684b5c 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -1930,3 +1930,71 @@ def test_verify_layout_on_gpu(tmp_path): ) with pytest.raises(ValueError, match="differs from its CudaStruct dtype"): views.verify_layout("drifted.cuh", include_dirs=[tmp_path]) + + +# --------------------------------------------------------------------------- +# cunumpy/reduce.cuh +# --------------------------------------------------------------------------- + +REDUCE_SOURCE = r""" +#include "cunumpy/reduce.cuh" +extern "C" __global__ +void reductions(const double* x, long long n, double* sum, unsigned long long* count, + double* block_min, double* block_max, double* warp_sum) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + double v = i < n ? x[i] : 0.0; // no early return: every thread takes part + cunumpy_block_sum_to(sum, v); + cunumpy_block_sum_to(count, i < n ? 1ull : 0ull); + double lo = cunumpy_block_min(i < n ? v : 1e300); + double hi = cunumpy_block_max(i < n ? v : -1e300); + if (threadIdx.x == 0) { block_min[blockIdx.x] = lo; block_max[blockIdx.x] = hi; } + double w = cunumpy_warp_sum(v); + if (threadIdx.x % 32 == 0) warp_sum[i / 32] = w; +} +""" + + +def test_reduce_header_is_shipped(): + header = Path(cuda_include_dir()) / "cunumpy" / "reduce.cuh" + text = header.read_text() + assert "#ifndef CUNUMPY_REDUCE_CUH" in text and "#endif" in text + for name in ( + "cunumpy_warp_sum", + "cunumpy_warp_min", + "cunumpy_warp_max", + "cunumpy_block_sum", + "cunumpy_block_min", + "cunumpy_block_max", + "cunumpy_block_sum_to", + ): + assert f"{name}(" in text, name + # the kernel's include resolves to the shipped header, which includes atomic.cuh + headers = resolve_includes(REDUCE_SOURCE, [cuda_include_dir()]) + assert [p.name for p in headers] == ["reduce.cuh", "atomic.cuh"] + CudaKernel(REDUCE_SOURCE, "reductions") # the signature parses + + +@pytest.mark.parametrize("block_size", [32, 128, 1024]) +def test_reductions_on_gpu(block_size): + _skip_without_cupy() + import cupy as cp + + n = 5000 # not a multiple of the block size: the last block is partial + x = cp.asarray(np.random.default_rng(1).normal(size=n)) + n_blocks = -(-n // block_size) + total, count = cp.zeros(1), cp.zeros(1, dtype=cp.uint64) + lo, hi = cp.zeros(n_blocks), cp.zeros(n_blocks) + warp = cp.zeros(n_blocks * block_size // 32) + kernel = CudaKernel(REDUCE_SOURCE, "reductions", block_size=block_size) + kernel(x, n, total, count, lo, hi, warp, n_threads=n) + + host = cp.asnumpy(x) + assert abs(float(total[0]) - host.sum()) < 1e-10 * n + assert int(count[0]) == n + padded = np.concatenate([host, np.full(n_blocks * block_size - n, np.nan)]) + blocks = padded.reshape(n_blocks, block_size) + np.testing.assert_array_equal(cp.asnumpy(lo), np.nanmin(blocks, axis=1)) + np.testing.assert_array_equal(cp.asnumpy(hi), np.nanmax(blocks, axis=1)) + np.testing.assert_allclose( + cp.asnumpy(warp), np.nan_to_num(padded).reshape(-1, 32).sum(axis=1), rtol=1e-12 + ) diff --git a/tests/unit/test_fusion.py b/tests/unit/test_fusion.py new file mode 100644 index 0000000..24ebc06 --- /dev/null +++ b/tests/unit/test_fusion.py @@ -0,0 +1,78 @@ +"""Tests for `xp.fuse`: cupy.fuse for CuPy arrays, a plain call otherwise.""" + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import fusion + + +def pressure(rho, T, gamma): + return (gamma - 1.0) * rho * xp.exp(T) + + +def test_host_arrays_call_the_function(): + fused = xp.fuse(pressure) + rho, T = np.full(4, 2.0), np.zeros(4) + np.testing.assert_array_equal(fused(rho, T, 3.0), pressure(rho, T, 3.0)) + assert fused.__name__ == "pressure" and fused.__wrapped__ is pressure + assert "fuse" in xp.__all__ + + +def test_decorator_forms(): + @xp.fuse + def double(x): + return 2.0 * x + + @xp.fuse(kernel_name="triple_kernel") + def triple(x): + return 3.0 * x + + assert double(np.ones(2)).tolist() == [2.0, 2.0] + assert triple(np.ones(2)).tolist() == [3.0, 3.0] + + +class FakeDeviceArray: + def __init__(self, values): + self.values = np.asarray(values) + + +def test_device_arrays_use_cupy_fuse_once(monkeypatch): + created = [] + + def fake_cupy_fuse(function, kernel_name): + created.append(kernel_name) + + def fused(*args, **kwargs): + assert xp.get_backend() in ("cupy", "numpy") # numpy: no CuPy installed + return ("fused", function.__name__, len(args), sorted(kwargs)) + + return fused + + monkeypatch.setattr(fusion, "_cupy_fuse", fake_cupy_fuse) + monkeypatch.setattr( + fusion, "_is_device_array", lambda a: isinstance(a, FakeDeviceArray) + ) + + @xp.fuse(kernel_name="p") + def p(rho, T, gamma=1.0): + raise AssertionError("not called with device arrays") + + rho = FakeDeviceArray([1.0]) + assert p(rho, 0.0) == ("fused", "p", 2, []) + assert p(1.0, 0.0, gamma=rho) == ("fused", "p", 2, ["gamma"]) # keyword array + assert created == ["p"] # fused once, reused + with pytest.raises(AssertionError, match="not called"): + p(np.ones(1), 0.0) # host arrays: the function itself + + +def test_fuse_on_gpu(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + fused = xp.fuse(pressure) + rho, T = cp.full(1000, 2.0), cp.linspace(0.0, 1.0, 1000) + with xp.use_backend("cupy"): + expected = pressure(rho, T, 5.0 / 3.0) + cp.testing.assert_allclose(fused(rho, T, 5.0 / 3.0), expected, rtol=1e-14) diff --git a/tests/unit/test_petsc.py b/tests/unit/test_petsc.py new file mode 100644 index 0000000..b2a6eb9 --- /dev/null +++ b/tests/unit/test_petsc.py @@ -0,0 +1,133 @@ +"""Tests for `xp.petsc_vec`: PETSc vectors sharing the memory of an array.""" + +import sys +import types + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import petsc + + +def _petsc(): + petsc4py = pytest.importorskip("petsc4py") + petsc4py.init() + from petsc4py import PETSc + + return PETSc + + +def test_host_vector_shares_memory(): + PETSc = _petsc() + a = np.zeros(6, dtype=PETSc.ScalarType) + vec = xp.petsc_vec(a, comm=PETSc.COMM_SELF) + assert vec.getSize() == 6 and vec.getType() == "seq" + vec.set(3.0) + assert a.tolist() == [3.0] * 6 # PETSc wrote into the array + a[0] = -1.0 + assert vec.getValue(0) == -1.0 # and reads from it + assert vec.getAttr("cunumpy_array") is a # keeps the array alive + + +def test_multidimensional_arrays_are_unrolled(): + PETSc = _petsc() + a = np.arange(6, dtype=PETSc.ScalarType).reshape(2, 3) + vec = xp.petsc_vec(a, comm=PETSc.COMM_SELF) + assert vec.getArray().tolist() == a.ravel().tolist() + + +def test_ksp_solve_writes_into_the_array(): + PETSc = _petsc() + n = 10 + A = PETSc.Mat().createAIJ([n, n], nnz=3, comm=PETSc.COMM_SELF) + for i in range(n): + A.setValue(i, i, 2.0) + if i > 0: + A.setValue(i, i - 1, -1.0) + if i < n - 1: + A.setValue(i, i + 1, -1.0) + A.assemble() + b, x = np.ones(n), np.zeros(n) + ksp = PETSc.KSP().create(comm=PETSc.COMM_SELF) + ksp.setOperators(A) + ksp.setType("cg") + ksp.getPC().setType("none") + ksp.setTolerances(rtol=1e-12) + ksp.solve( + xp.petsc_vec(b, comm=PETSc.COMM_SELF), xp.petsc_vec(x, comm=PETSc.COMM_SELF) + ) + residual = 2 * x - np.r_[0.0, x[:-1]] - np.r_[x[1:], 0.0] - 1.0 + assert np.abs(residual).max() < 1e-9 + + +def test_rejects_arrays_that_would_need_a_copy(): + PETSc = _petsc() + other = np.float32 if np.dtype(PETSc.ScalarType) != np.float32 else np.float64 + with pytest.raises(TypeError, match="scalar type"): + xp.petsc_vec(np.zeros(4, dtype=other)) + with pytest.raises(ValueError, match="C-contiguous"): + xp.petsc_vec(np.zeros((4, 4))[:, 0]) + with pytest.raises(TypeError, match="NumPy or CuPy array"): + xp.petsc_vec([0.0, 1.0]) + + +class FakeDeviceArray: + dtype = np.dtype(np.float64) + flags = types.SimpleNamespace(c_contiguous=True) + + +def _fake_petsc(monkeypatch, vec_type=None, error=False): + class Error(Exception): + pass + + class Vec: + destroyed = False + + def createWithDLPack(self, array, comm=None): + if error: + raise Error("no CUDA") + return self + + def getType(self): + return vec_type + + def destroy(self): + Vec.destroyed = True + + def setAttr(self, name, value): + self.attr = (name, value) + + PETSc = types.SimpleNamespace(Vec=Vec, Error=Error, ScalarType=np.float64) + package = types.ModuleType("petsc4py") + package.PETSc = PETSc + monkeypatch.setitem(sys.modules, "petsc4py", package) + monkeypatch.setitem(sys.modules, "petsc4py.PETSc", PETSc) + monkeypatch.setattr( + petsc, "_is_device_array", lambda a: isinstance(a, FakeDeviceArray) + ) + return Vec + + +@pytest.mark.parametrize("vec_type", ["seqcuda", "mpicuda", "seqhip"]) +def test_device_arrays_give_device_vectors(monkeypatch, vec_type): + _fake_petsc(monkeypatch, vec_type) + array = FakeDeviceArray() + vec = xp.petsc_vec(array) + assert vec.attr == ("cunumpy_array", array) + + +def test_device_arrays_need_a_gpu_petsc(monkeypatch): + Vec = _fake_petsc(monkeypatch, "seq") # PETSc without CUDA made a host vector + with pytest.raises(RuntimeError, match="created a 'seq' vector for a CuPy array"): + xp.petsc_vec(FakeDeviceArray()) + assert Vec.destroyed + _fake_petsc(monkeypatch, error=True) + with pytest.raises(RuntimeError, match="built with CUDA or HIP support"): + xp.petsc_vec(FakeDeviceArray()) + + +def test_missing_petsc4py(monkeypatch): + monkeypatch.setitem(sys.modules, "petsc4py", None) + with pytest.raises(ImportError, match="needs petsc4py"): + xp.petsc_vec(np.zeros(3)) diff --git a/tests/unit/test_scipy_backend.py b/tests/unit/test_scipy_backend.py new file mode 100644 index 0000000..4d3fe29 --- /dev/null +++ b/tests/unit/test_scipy_backend.py @@ -0,0 +1,104 @@ +"""Tests for `xp.scipy`: SciPy or cupyx.scipy, by the active backend.""" + +import sys +import types + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import scipy_backend +from cunumpy.scipy_backend import SUBMODULES, ScipyNamespace + + +@pytest.fixture +def fake_cupyx(monkeypatch): + """A fake `cupyx.scipy` with `special.erf` and `sparse.linalg.cg`, as backend.""" + modules = {} + for name in ("cupyx", "cupyx.scipy", "cupyx.scipy.special", "cupyx.scipy.sparse"): + modules[name] = types.ModuleType(name) + modules["cupyx.scipy.sparse.linalg"] = types.ModuleType("cupyx.scipy.sparse.linalg") + modules["cupyx.scipy.special"].erf = "device erf" + modules["cupyx.scipy.sparse.linalg"].cg = "device cg" + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setattr(scipy_backend, "get_backend", lambda: "cupy") + return modules + + +def test_scipy_is_exported(): + assert isinstance(xp.scipy, ScipyNamespace) + assert "scipy" in xp.__all__ + assert set(dir(xp.scipy)) >= {"sparse", "special", "fft", "available", "resolve"} + assert set(dir(xp.scipy.sparse)) >= {"linalg", "csgraph"} + + +def test_unknown_subpackage(): + with pytest.raises(ValueError, match="not forwarded"): + ScipyNamespace("optimize") + assert all(ScipyNamespace(name) for name in SUBMODULES) + + +def test_numpy_backend_forwards_to_scipy(): + pytest.importorskip("scipy") + import scipy.sparse.linalg + import scipy.special + + assert xp.scipy.special.erf is scipy.special.erf + assert xp.scipy.sparse.linalg.cg is scipy.sparse.linalg.cg + assert xp.scipy.sparse.linalg.resolve() is scipy.sparse.linalg + assert xp.scipy.sparse.linalg is xp.scipy.sparse.linalg # cached namespaces + + n = 20 + A = xp.scipy.sparse.diags( + [np.full(n - 1, -1.0), np.full(n, 2.0), np.full(n - 1, -1.0)], + [-1, 0, 1], + format="csr", + ) + x, info = xp.scipy.sparse.linalg.cg(A, np.ones(n), rtol=1e-12) + assert info == 0 and np.allclose(A @ x, 1.0) + + +def test_missing_names_say_which_backend_lacks_them(): + pytest.importorskip("scipy") + with pytest.raises(AttributeError, match=r"not available on the numpy backend"): + xp.scipy.special.no_such_function # noqa: B018 + assert xp.scipy.special.available("erf") + assert not xp.scipy.special.available("no_such_function") + assert xp.scipy.available("sparse") + with pytest.raises(AttributeError): + xp.scipy.__wrapped__ # noqa: B018 -- dunder names are not forwarded + + +def test_cupy_backend_forwards_to_cupyx(fake_cupyx): + assert xp.scipy.special.erf == "device erf" + assert xp.scipy.sparse.linalg.cg == "device cg" + assert xp.scipy.special.resolve() is fake_cupyx["cupyx.scipy.special"] + with pytest.raises( + AttributeError, match=r"cupy backend \(it may exist in scipy\.special\)" + ): + xp.scipy.special.erfcx # noqa: B018 + # a subpackage cupyx does not provide + assert not xp.scipy.available("ndimage") + with pytest.raises(ImportError, match="needs cupyx.scipy.ndimage"): + xp.scipy.ndimage.resolve() + + +def test_backend_switch_takes_effect_immediately(fake_cupyx, monkeypatch): + pytest.importorskip("scipy") + import scipy.special + + special = xp.scipy.special + assert special.erf == "device erf" + monkeypatch.setattr(scipy_backend, "get_backend", lambda: "numpy") + assert special.erf is scipy.special.erf + + +def test_missing_scipy(monkeypatch): + monkeypatch.setitem(sys.modules, "scipy", None) # import scipy raises ImportError + monkeypatch.setitem(sys.modules, "scipy.special", None) + with pytest.raises( + ImportError, match=r"needs scipy.special: SciPy is not installed" + ): + xp.scipy.special.erf # noqa: B018 + assert not xp.scipy.special.available("erf") From 9b508ca741378989a2ba385539e3e48ec0654515 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 2 Oct 2026 00:00:20 +0200 Subject: [PATCH 4/9] Formatting --- src/cunumpy/cuda_kernel.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index 2ee3364..16bb290 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -1988,7 +1988,9 @@ def _is_capturing(stream: Any) -> bool: return False try: return bool(is_capturing()) - except Exception: # noqa: BLE001 -- e.g. the legacy null stream, which cannot capture + except ( + Exception + ): # noqa: BLE001 -- e.g. the legacy null stream, which cannot capture return False From 792f28dcf36b46f694f18737392be22ed544935c Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 2 Oct 2026 00:17:15 +0200 Subject: [PATCH 5/9] Hash shipped headers in the kernel cache key; add scipy to test extra --- CHANGELOG.md | 2 + docs/source/api.md | 17 +++++--- pyproject.toml | 2 +- src/cunumpy/cuda_kernel.py | 71 +++++++++++++++++++++++++--------- tests/unit/test_cuda_kernel.py | 53 +++++++++++++++++++++++++ tests/unit/test_mirror.py | 5 ++- 6 files changed, 123 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4127124..50f8ac5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Fixed +- `CudaKernel`'s header hash (`-DCUNUMPY_INCLUDE_HASH`) now covers the headers shipped with cunumpy (`cunumpy/atomic.cuh`, `reduce.cuh`, ...), also when included in angle brackets. Before, an upgrade of cunumpy that changed one of them left CuPy's kernel cache serving the kernel compiled with the old header. `resolve_includes(..., angle_dirs=...)` tracks angle-bracket includes found in the given directories. - `xp.testing.assert_kernels_agree` reads `CudaStructArguments` objects and struct values through their struct fields, so their arrays get the same names as the attributes of the host argument object (before, arrays behind properties were named after the private attribute holding the owner, and the comparison failed with "do not have the same array arguments"). ### Removed @@ -15,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - Python 3.14 is supported. +- The `test` extra installs SciPy, so the `xp.scipy` tests run in CI instead of being skipped. - `KernelCatalog.compile_all(jobs=1)` and `CudaKernelVariants.compile_all(keys, jobs=1)`: With `jobs > 1` the CUDA kernels are compiled in threads (NVRTC releases the GIL); `jobs=None` uses the number of CPUs. All kernels are compiled even if one fails, and the first error is raised afterwards. - 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. - `CudaKernel.compile()` passes `compile_options()` to CuPy: the given `options` plus `-DCUNUMPY_INCLUDE_HASH=0x` when the source includes header files, so CuPy's kernel cache (keyed on source and options only) is invalidated when an included header changes. `options` still returns the options as given. diff --git a/docs/source/api.md b/docs/source/api.md index f8498b4..91b1a01 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -694,19 +694,24 @@ kernel.compile_options() # options + ('-DCUNUMPY_INCLUDE_HASH=0x3f9a...',) `#include "name"`, recursively, each once in order of first inclusion. A name is looked up relative to the including file (`source_dir` for the kernel source, the header's own directory for nested includes), then in - `include_dirs` in order, like NVRTC does. System headers in angle brackets - and includes that cannot be found are ignored (NVRTC reports the latter). - Recomputed at every access, so it follows the files on disk. + `include_dirs` in order, then in cunumpy's header directory, like NVRTC + does. cunumpy's shipped headers are tracked also when included in angle + brackets (`#include `), so upgrading cunumpy with a + changed header recompiles the kernels that use it. Other angle-bracket + (system) headers and includes that cannot be found are ignored (NVRTC + reports the latter). Recomputed at every access, so it follows the files on + disk. * `compile_options()`: the options passed to CuPy at compile time: `options` plus `-DCUNUMPY_INCLUDE_HASH=0x` if the source includes any header, where the hash covers the contents of `included_headers` (not their paths). A changed header gives another define, hence another cache entry. Sources - without quoted includes never touch the file system. + without includes never touch the file system. The two building blocks are available on their own: -* `xp.resolve_includes(source, include_dirs=(), *, base_dir=None)`: the - resolved header paths of a source, as a list. +* `xp.resolve_includes(source, include_dirs=(), *, base_dir=None, + angle_dirs=())`: the resolved header paths of a source, as a list; + `angle_dirs` are searched last and also for `#include `. * `xp.include_hash(paths)`: the first 16 hex digits of the SHA-256 digest of the contents of the files, in order. diff --git a/pyproject.toml b/pyproject.toml index ba99ffd..07c8b48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ optional-dependencies.docs = [ "sphinx", "sphinx-book-theme", ] -optional-dependencies.test = [ "coverage", "pytest" ] +optional-dependencies.test = [ "coverage", "pytest", "scipy" ] optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ] urls."Source" = "https://github.com/max-models/cunumpy" diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index 16bb290..edb9b33 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -273,12 +273,23 @@ def _strip_comments(source: str) -> str: # ``#include "name"``: quoted includes are the project's own headers. Angle -# bracket includes are system headers and are not tracked. -_QUOTED_INCLUDE = re.compile(r'^[ \t]*#[ \t]*include[ \t]*"([^"\n]+)"', re.MULTILINE) +# bracket includes are system headers and are not tracked, except in the +# directories given as ``angle_dirs`` (cunumpy's shipped headers). +_INCLUDE = re.compile( + r'^[ \t]*#[ \t]*include[ \t]*(?:"([^"\n]+)"|<([^>\n]+)>)', re.MULTILINE +) + + +def _includes(source: str) -> list[tuple[str, bool]]: + """``(name, quoted)`` of every ``#include`` in `source`, in order.""" + return [ + (quoted or angle, bool(quoted)) + for quoted, angle in _INCLUDE.findall(_strip_comments(source)) + ] def _quoted_includes(source: str) -> list[str]: - return _QUOTED_INCLUDE.findall(_strip_comments(source)) + return [name for name, quoted in _includes(source) if quoted] def resolve_includes( @@ -286,15 +297,17 @@ def resolve_includes( include_dirs: Iterable[str | Path] = (), *, base_dir: str | Path | None = None, + angle_dirs: Iterable[str | Path] = (), ) -> list[Path]: """The header files a CUDA source includes, recursively. Scans `source` (comments removed) for ``#include "name"`` and resolves each name like NVRTC does: relative to `base_dir` (the directory of the - including file), then in `include_dirs`, in order. Found headers are - scanned in turn, relative to their own directory. Includes in angle - brackets (system headers) and includes that cannot be found are ignored; - NVRTC reports the latter when the kernel is compiled. + including file), then in `include_dirs`, then in `angle_dirs`, in order. + Found headers are scanned in turn, relative to their own directory. + Includes in angle brackets (``#include ``) are system headers and + ignored, unless they are found in `angle_dirs`. Includes that cannot be + found are ignored; NVRTC reports them when the kernel is compiled. Parameters ---------- @@ -305,22 +318,32 @@ def resolve_includes( base_dir : str | Path | None Directory of the file `source` was read from, searched first; None if the source is not from a file. + angle_dirs : Iterable[str | Path] + Directories whose headers are tracked also when included in angle + brackets, searched last; :class:`CudaKernel` passes + :func:`cuda_include_dir`, so that ``#include `` + is tracked. Returns ------- list[Path] The resolved header files, each once, in order of first inclusion - (depth first). Empty if the source has no quoted includes; the file + (depth first). Empty if the source has no includes to track; the file system is not touched in that case. """ + angle = tuple(Path(d) for d in angle_dirs) dirs = tuple(Path(d) for d in include_dirs) + dirs += tuple(d for d in angle if d not in dirs) found: list[Path] = [] seen: set[Path] = set() def visit(code: str, directory: Path | None) -> None: - for name in _quoted_includes(code): - candidates = [directory / name] if directory is not None else [] - candidates += [d / name for d in dirs] + for name, quoted in _includes(code): + if quoted: + candidates = [directory / name] if directory is not None else [] + candidates += [d / name for d in dirs] + else: + candidates = [d / name for d in angle] for candidate in candidates: if candidate.is_file(): path = candidate.resolve() @@ -1521,7 +1544,9 @@ class CudaKernel: options, but not on the files pulled in by ``#include "..."``. At compile time the headers are resolved (:attr:`included_headers`) and a define with the hash of their contents is added to the options - (:meth:`compile_options`), so editing a header recompiles the kernel. + (:meth:`compile_options`), so editing a header recompiles the kernel; this + includes cunumpy's own headers (``#include ``), so an + upgrade that changes them recompiles too. Examples -------- @@ -1721,15 +1746,22 @@ def source_dir(self) -> Path | None: @property def included_headers(self) -> tuple[Path, ...]: - """The header files the source includes with ``#include "..."``. - - Resolved recursively in `source_dir` and `include_dirs` at every access - (see :func:`resolve_includes`), so the result follows the files on - disk. Empty if the source has no quoted includes. + """The header files the source includes, recursively. + + Quoted includes (``#include "..."``) are resolved in `source_dir`, + `include_dirs` and cunumpy's header directory + (:func:`cuda_include_dir`); cunumpy's shipped headers are tracked + also when included in angle brackets (``#include ``). + Resolved at every access (see :func:`resolve_includes`), so the result + follows the files on disk. Empty if the source includes no project or + cunumpy headers. """ return tuple( resolve_includes( - self._source, self._include_dirs, base_dir=self._source_dir + self._source, + self._include_dirs, + base_dir=self._source_dir, + angle_dirs=(cuda_include_dir(),), ) ) @@ -1744,7 +1776,8 @@ def compile_options(self) -> tuple[str, ...]: it), and, if the source includes header files, ``-DCUNUMPY_INCLUDE_HASH=0x`` with the hash of the contents of :attr:`included_headers` (see :func:`include_hash`). CuPy keys its - kernel cache on the options, so a changed header means a recompile. + kernel cache on the options, so a changed header means a recompile, + also for a header shipped with cunumpy that changed in an upgrade. """ options = self._options # cunumpy's own headers (, ...) are always found diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 5684b5c..2d54e9b 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -1517,6 +1517,59 @@ def test_compile_options_contain_the_header_hash(header_tree): assert _user_options(CudaKernel(INCLUDING_SOURCE, "double_it")) == () +def test_resolve_includes_angle_dirs(tmp_path): + shipped = tmp_path / "shipped" + (shipped / "lib").mkdir(parents=True) + (shipped / "lib" / "a.cuh").write_text('#include "lib/b.cuh"\n') + (shipped / "lib" / "b.cuh").write_text("") + source = "#include \n#include \n" + # angle brackets: system headers, not tracked by default + assert resolve_includes(source, [shipped]) == [] + # in angle_dirs they are, and their quoted includes resolve there too + expected = [shipped / "lib" / "a.cuh", shipped / "lib" / "b.cuh"] + assert resolve_includes(source, angle_dirs=[shipped]) == expected + assert resolve_includes('#include "lib/a.cuh"\n', angle_dirs=[shipped]) == expected + # include_dirs come first for quoted includes + user = tmp_path / "user" + (user / "lib").mkdir(parents=True) + (user / "lib" / "a.cuh").write_text("") + assert resolve_includes('#include "lib/a.cuh"\n', [user], angle_dirs=[shipped]) == [ + user / "lib" / "a.cuh" + ] + + +def test_shipped_headers_are_part_of_the_hash(): + include = Path(cuda_include_dir()) / "cunumpy" + for line in ("#include ", '#include "cunumpy/reduce.cuh"'): + kernel = CudaKernel(line + "\n" + AXPY, "axpy") + assert kernel.included_headers == ( + include / "reduce.cuh", + include / "atomic.cuh", + ) + digest = include_hash(kernel.included_headers) + assert _user_options(kernel) == (f"-DCUNUMPY_INCLUDE_HASH=0x{digest}",) + # a source without includes still gets no define + assert _user_options(CudaKernel(AXPY, "axpy")) == () + + +def test_changed_shipped_header_changes_the_hash(tmp_path, monkeypatch): + # a copy of the shipped headers stands for the installed ones before and + # after an upgrade of cunumpy + import shutil + + from cunumpy import cuda_kernel + + installed = tmp_path / "include" + shutil.copytree(cuda_include_dir(), installed) + monkeypatch.setattr(cuda_kernel, "_CUDA_INCLUDE_DIR", installed) + kernel = CudaKernel("#include \n" + AXPY, "axpy") + before = kernel.compile_options()[-1] + assert before.startswith("-DCUNUMPY_INCLUDE_HASH=0x") + header = installed / "cunumpy" / "atomic.cuh" + header.write_text(header.read_text() + "\n// changed in an upgrade\n") + assert kernel.compile_options()[-1] != before + + def test_editing_a_header_recompiles_on_gpu(header_tree): _skip_without_cupy() import cupy as cp diff --git a/tests/unit/test_mirror.py b/tests/unit/test_mirror.py index d9a0a2e..4ab0a24 100644 --- a/tests/unit/test_mirror.py +++ b/tests/unit/test_mirror.py @@ -147,7 +147,10 @@ def test_cuda_kernel_options_include_cunumpy_headers(): assert kernel.compile_options().count(flag) == 1 assert kernel.compile_options()[0] == "-I/some/dir" kernel = CudaKernel(BIN_ADD, "bin_add", options=["-std=c++17", flag]) - assert kernel.compile_options() == ("-std=c++17", flag) + options = kernel.compile_options() + assert options[:2] == ("-std=c++17", flag) + # the shipped atomic.cuh is part of the header hash + assert len(options) == 3 and options[2].startswith("-DCUNUMPY_INCLUDE_HASH=0x") # --- CuPy backend ------------------------------------------------------------ From cd98f0e240b1ec3c349ba04de54ee0952202dbdb Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 2 Oct 2026 08:55:18 +0200 Subject: [PATCH 6/9] Fix xp.fuse test --- src/cunumpy/fusion.py | 25 ++++++++++++++++++++++++- tests/unit/test_fusion.py | 19 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/src/cunumpy/fusion.py b/src/cunumpy/fusion.py index 4713deb..af72edf 100644 --- a/src/cunumpy/fusion.py +++ b/src/cunumpy/fusion.py @@ -30,6 +30,7 @@ def pressure(rho, T, gamma): from typing import Any, TypeVar import array_api_compat +import numpy from .xp import use_backend @@ -47,6 +48,26 @@ def _cupy_fuse(function: Callable[..., Any], kernel_name: str | None) -> Any: return cupy.fuse(kernel_name=kernel_name)(function) +def _typed_scalars(args: tuple, kwargs: dict) -> tuple[tuple, dict]: + """Give Python scalars the dtype they would take next to the arrays. + + ``cupy.fuse`` types a Python scalar on its own: ``gamma - 1.0`` with + ``gamma=5/3`` runs in float16. Eagerly, NumPy 2 and CuPy promote it with + the arrays (float64 arrays: float64), so cast it to that dtype first. + """ + dtypes = [a.dtype for a in (*args, *kwargs.values()) if hasattr(a, "dtype")] + + def typed(a: Any) -> Any: + if isinstance(a, (int, float, complex)) and not isinstance(a, bool): + return numpy.result_type(*dtypes, a).type(a) + return a + + return ( + tuple(typed(a) for a in args), + {k: typed(v) for k, v in kwargs.items()}, + ) + + def fuse( function: F | None = None, *, kernel_name: str | None = None ) -> F | Callable[[F], F]: @@ -67,7 +88,8 @@ def fuse( callable A function with the same signature. If any positional or keyword argument is a CuPy array, it calls ``cupy.fuse(function)`` (created on - first use, with the CuPy backend active); otherwise it calls + first use, with the CuPy backend active) with Python scalars cast to + the dtype they promote to with the array arguments; otherwise it calls `function` itself. """ if function is None: @@ -81,6 +103,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: nonlocal fused if not any(_is_device_array(a) for a in (*args, *kwargs.values())): return function(*args, **kwargs) + args, kwargs = _typed_scalars(args, kwargs) with use_backend("cupy"): if fused is None: fused = _cupy_fuse(function, name) diff --git a/tests/unit/test_fusion.py b/tests/unit/test_fusion.py index 24ebc06..717ffe9 100644 --- a/tests/unit/test_fusion.py +++ b/tests/unit/test_fusion.py @@ -66,6 +66,25 @@ def p(rho, T, gamma=1.0): p(np.ones(1), 0.0) # host arrays: the function itself +def test_python_scalars_take_the_dtype_of_the_arrays(monkeypatch): + seen = [] + monkeypatch.setattr( + fusion, + "_cupy_fuse", + lambda function, kernel_name: lambda *a, **k: seen.append((a, k)), + ) + monkeypatch.setattr(fusion, "_is_device_array", lambda a: hasattr(a, "dtype")) + + p = xp.fuse(pressure) + p(np.ones(2), 0.5, 5.0 / 3.0) + p(np.ones(2, dtype=np.float32), 0.5, gamma=2) + p(np.ones(2, dtype=np.int64), True, 3) + (a64, _), (a32, k32), (aint, _) = seen + assert a64[1].dtype == a64[2].dtype == np.float64 and a64[2] == 5.0 / 3.0 + assert a32[1].dtype == k32["gamma"].dtype == np.float32 + assert aint[1] is True and aint[2].dtype == np.int64 # bools stay Python + + def test_fuse_on_gpu(): if not xp.cupy_available(): pytest.skip("CuPy not installed or not functional") From 1afc5b4e83e0450cfedd1ceffe925d4cae4a77dd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 2 Oct 2026 08:57:47 +0200 Subject: [PATCH 7/9] fix ruff check, noqa: F403 (re-export numpy for completions) --- src/cunumpy/__init__.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 879765e..6181372 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -6,7 +6,7 @@ from contextlib import contextmanager from typing import Any import numpy as np -from numpy import * +from numpy import * # noqa: F403 (re-export numpy for completions) from . import xp as xp from .cuda_kernel import CudaArguments as CudaArguments From f299e7e9e0819341dcf954e59d7bc9d5bdb3d2b4 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 2 Oct 2026 09:03:52 +0200 Subject: [PATCH 8/9] ruff check --fix --- src/cunumpy/__init__.pyi | 2 +- src/cunumpy/cuda_kernel.py | 2 +- tests/unit/test_cuda_kernel.py | 7 +++---- tests/unit/test_profiling.py | 26 +++++++++++--------------- tests/unit/test_testing.py | 2 +- tests/unit/test_transfers.py | 29 ++++++++++++----------------- 6 files changed, 29 insertions(+), 39 deletions(-) diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 6181372..879765e 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -6,7 +6,7 @@ from contextlib import contextmanager from typing import Any import numpy as np -from numpy import * # noqa: F403 (re-export numpy for completions) +from numpy import * from . import xp as xp from .cuda_kernel import CudaArguments as CudaArguments diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index edb9b33..ace4d07 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -2023,7 +2023,7 @@ def _is_capturing(stream: Any) -> bool: return bool(is_capturing()) except ( Exception - ): # noqa: BLE001 -- e.g. the legacy null stream, which cannot capture + ): return False diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 2d54e9b..e0206d8 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -948,7 +948,7 @@ def unknown(x: "str[:]"): def too_many(x: "float[:, :, :, :]"): pass - def unparsable(x: "float[:](order=F)"): + def unparsable(x: "float[:](order=F)"): # noqa: F821 (deliberately unparsable) pass with pytest.raises( @@ -1289,9 +1289,8 @@ def test_cuda_debug_context_restores(debug_off): assert xp.get_cuda_debug() is True assert xp.get_cuda_debug() is False - with pytest.raises(ValueError): - with xp.cuda_debug(): - raise ValueError + with pytest.raises(ValueError), xp.cuda_debug(): + raise ValueError assert xp.get_cuda_debug() is False # restored after an exception too diff --git a/tests/unit/test_profiling.py b/tests/unit/test_profiling.py index cb4d82d..2c70a25 100644 --- a/tests/unit/test_profiling.py +++ b/tests/unit/test_profiling.py @@ -64,11 +64,10 @@ def test_nvtx_range_repr(): def test_timed_region_on_numpy(): - with xp.use_backend("numpy"): - with xp.timed_region("sleep") as timing: - assert timing.name == "sleep" - assert timing.elapsed is None - time.sleep(0.02) + with xp.use_backend("numpy"), xp.timed_region("sleep") as timing: + assert timing.name == "sleep" + assert timing.elapsed is None + time.sleep(0.02) assert isinstance(timing, xp.Timing) assert timing.elapsed >= 0.02 @@ -77,18 +76,16 @@ def test_timed_region_on_numpy(): def test_timed_region_without_sync_on_numpy(): - with xp.use_backend("numpy"): - with xp.timed_region("no sync", sync=False) as timing: - pass + with xp.use_backend("numpy"), xp.timed_region("no sync", sync=False) as timing: + pass assert timing.elapsed >= 0.0 assert timing.synced is False def test_timed_region_records_time_on_exception(): - with xp.use_backend("numpy"): - with pytest.raises(RuntimeError, match="boom"): - with xp.timed_region("failing") as timing: - raise RuntimeError("boom") + with xp.use_backend("numpy"), pytest.raises(RuntimeError, match="boom"): + with xp.timed_region("failing") as timing: + raise RuntimeError("boom") assert timing.elapsed is not None assert timing.elapsed >= 0.0 @@ -126,9 +123,8 @@ def test_nvtx_range_nested_and_reentrant(fake_nvtx): def test_nvtx_range_pops_on_exception(fake_nvtx): - with pytest.raises(ValueError): - with xp.nvtx_range("failing"): - raise ValueError + with pytest.raises(ValueError), xp.nvtx_range("failing"): + raise ValueError assert fake_nvtx == [("push", "failing", -1), ("pop",)] diff --git a/tests/unit/test_testing.py b/tests/unit/test_testing.py index 4c79813..a716b35 100644 --- a/tests/unit/test_testing.py +++ b/tests/unit/test_testing.py @@ -16,12 +16,12 @@ import cunumpy as xp import cunumpy.testing from cunumpy import CudaArguments, CudaKernel, Kernel, parse_cuda_signature -from cunumpy.testing import backend # noqa: F401 - the fixture is used by name from cunumpy.testing import ( BACKENDS, _collect_arrays, _compare_results, assert_kernels_agree, + backend, # noqa: F401 - the fixture is used by name device_function_kernel, requires_cupy, ) diff --git a/tests/unit/test_transfers.py b/tests/unit/test_transfers.py index 3c2d317..6a6bbdf 100644 --- a/tests/unit/test_transfers.py +++ b/tests/unit/test_transfers.py @@ -135,10 +135,9 @@ def test_to_cupy_of_device_array_is_not_counted(fake_device): def test_to_cunumpy_counts_the_direction_it_delegates_to(fake_device, monkeypatch): device = fake_device(np.zeros(2)) - with xp.count_transfers() as counter: - with xp.use_backend("numpy"): - xp.to_cunumpy(device) # device -> host - xp.to_cunumpy(np.zeros(2)) # already on the host + with xp.count_transfers() as counter, xp.use_backend("numpy"): + xp.to_cunumpy(device) # device -> host + xp.to_cunumpy(np.zeros(2)) # already on the host assert counter.to_host == 1 and counter.to_device == 0 @@ -168,9 +167,8 @@ def test_where_points_at_the_caller_outside_cunumpy(fake_device): def test_where_skips_frames_inside_cunumpy(fake_device): """`to_cunumpy` calls `to_numpy`; the call site is still the test.""" - with xp.count_transfers() as counter: - with xp.use_backend("numpy"): - xp.to_cunumpy(fake_device(np.zeros(1))) + with xp.count_transfers() as counter, xp.use_backend("numpy"): + xp.to_cunumpy(fake_device(np.zeros(1))) (event,) = counter.events assert event.where.startswith(THIS_FILE + ":") @@ -224,9 +222,8 @@ def test_nested_counters_each_see_their_own_block(fake_device): def test_counter_is_removed_when_the_block_raises(fake_device): - with pytest.raises(RuntimeError): - with xp.count_transfers(): - raise RuntimeError + with pytest.raises(RuntimeError), xp.count_transfers(): + raise RuntimeError assert transfers_module._ACTIVE == [] @@ -239,9 +236,8 @@ def test_assert_no_transfers_passes_without_transfers(): def test_assert_no_transfers_raises_with_report(fake_device): - with pytest.raises(AssertionError) as info: - with xp.assert_no_transfers(): - xp.to_numpy(fake_device(np.zeros(3))) + with pytest.raises(AssertionError) as info, xp.assert_no_transfers(): + xp.to_numpy(fake_device(np.zeros(3))) message = str(info.value) assert "1 transfer(s) through cunumpy" in message @@ -251,10 +247,9 @@ def test_assert_no_transfers_raises_with_report(fake_device): def test_assert_no_transfers_lets_exceptions_through(fake_device): - with pytest.raises(ValueError, match="inside"): - with xp.assert_no_transfers(): - xp.to_numpy(fake_device(np.zeros(3))) - raise ValueError("inside") + with pytest.raises(ValueError, match="inside"), xp.assert_no_transfers(): + xp.to_numpy(fake_device(np.zeros(3))) + raise ValueError("inside") # --------------------------------------------------------------------------- From 2d4274a473879690244c2b62b10c709b044eb237 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 2 Oct 2026 09:11:27 +0200 Subject: [PATCH 9/9] ruff fixes --- pyproject.toml | 4 ++++ src/cunumpy/cuda_kernel.py | 7 ++----- tests/unit/test_cuda_kernel.py | 25 ++++++++++++++++--------- tests/unit/test_kernel_dispatch.py | 6 ++++-- tests/unit/test_mirror.py | 3 ++- tests/unit/test_profiling.py | 9 ++++++--- tests/unit/test_testing.py | 2 +- tests/unit/test_transfers.py | 11 +++++++---- 8 files changed, 42 insertions(+), 25 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 07c8b48..43a8bda 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,3 +55,7 @@ cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ] [tool.isort] profile = "black" + +[tool.ruff] +# agent worktrees and scratch scripts are local copies, not part of the project +extend-exclude = [ ".claude" ] diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index ace4d07..ade7725 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -646,8 +646,7 @@ def check(value: Any) -> Any: _check_device_array(param, index, value) if value.ndim != ndim: raise TypeError( - f"{_describe(param, index)} must be a {ndim}D array, got " - f"{value.ndim}D" + f"{_describe(param, index)} must be a {ndim}D array, got {value.ndim}D" ) itemsize = value.dtype.itemsize strides = [s // itemsize for s in value.strides] @@ -2021,9 +2020,7 @@ def _is_capturing(stream: Any) -> bool: return False try: return bool(is_capturing()) - except ( - Exception - ): + except Exception: # noqa: BLE001 -- e.g. the legacy null stream, which cannot capture return False diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index e0206d8..5969b13 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -467,7 +467,9 @@ def test_ctype_of(): ], ) -PUSH_SOURCE = PARTICLES.declaration + r""" +PUSH_SOURCE = ( + PARTICLES.declaration + + r""" extern "C" __global__ void push(Particles p, double dt, double* out, unsigned long long* size) { int i = blockDim.x * blockIdx.x + threadIdx.x; @@ -478,6 +480,7 @@ def test_ctype_of(): if (i < p.n && p.alive[i]) p.x[i] += dt * p.charge; } """ +) def test_struct_layout_and_declaration(): @@ -545,13 +548,13 @@ def test_struct_values(): def test_struct_pointer_fields_must_be_contiguous(): - values = dict( - n=3, - charge=2.0, - alive=FakeDeviceArray(np.bool_), - ids=FakeDeviceArray(np.int64), - weight=0.5, - ) + values = { + "n": 3, + "charge": 2.0, + "alive": FakeDeviceArray(np.bool_), + "ids": FakeDeviceArray(np.int64), + "weight": 0.5, + } view = FakeDeviceArray(np.float64, flags=SimpleNamespace(c_contiguous=False)) with pytest.raises( TypeError, match=r"argument 0 \(double\* x\) must be C-contiguous" @@ -1221,7 +1224,11 @@ def _run_python(code, env=None): [str(Path(xp.__file__).parents[1]), environment.get("PYTHONPATH", "")] ) result = subprocess.run( - [sys.executable, "-c", code], capture_output=True, text=True, env=environment + [sys.executable, "-c", code], + capture_output=True, + text=True, + env=environment, + check=False, ) return result.stdout, result.stderr diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py index a34c8d2..bef2066 100644 --- a/tests/unit/test_kernel_dispatch.py +++ b/tests/unit/test_kernel_dispatch.py @@ -129,11 +129,13 @@ def kernel_package(tmp_path, monkeypatch): 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""" + (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( diff --git a/tests/unit/test_mirror.py b/tests/unit/test_mirror.py index 4ab0a24..1a97892 100644 --- a/tests/unit/test_mirror.py +++ b/tests/unit/test_mirror.py @@ -130,7 +130,8 @@ def test_cuda_include_dir_contains_atomic_header(): assert include_dir.endswith(os.path.join("cuda", "include")) header = os.path.join(include_dir, "cunumpy", "atomic.cuh") assert os.path.isfile(header) - source = open(header).read() + with open(header) as f: + source = f.read() assert "cunumpy_atomic_add(double* p, double v)" in source assert "cunumpy_atomic_add(float* p, float v)" in source assert "cunumpy_atomic_add_2d(" in source diff --git a/tests/unit/test_profiling.py b/tests/unit/test_profiling.py index 2c70a25..6d01102 100644 --- a/tests/unit/test_profiling.py +++ b/tests/unit/test_profiling.py @@ -83,9 +83,12 @@ def test_timed_region_without_sync_on_numpy(): def test_timed_region_records_time_on_exception(): - with xp.use_backend("numpy"), pytest.raises(RuntimeError, match="boom"): - with xp.timed_region("failing") as timing: - raise RuntimeError("boom") + with ( + xp.use_backend("numpy"), + pytest.raises(RuntimeError, match="boom"), + xp.timed_region("failing") as timing, + ): + raise RuntimeError("boom") assert timing.elapsed is not None assert timing.elapsed >= 0.0 diff --git a/tests/unit/test_testing.py b/tests/unit/test_testing.py index a716b35..7ab3778 100644 --- a/tests/unit/test_testing.py +++ b/tests/unit/test_testing.py @@ -74,7 +74,7 @@ def test_backends_and_marker(): assert requires_cupy.args == (not xp.cupy_available(),) assert requires_cupy.kwargs["reason"] == "CuPy/GPU not available" with pytest.raises(AttributeError): - cunumpy.testing.no_such_thing + _ = cunumpy.testing.no_such_thing @pytest.mark.parametrize("backend_name", BACKENDS) diff --git a/tests/unit/test_transfers.py b/tests/unit/test_transfers.py index 6a6bbdf..00cb46c 100644 --- a/tests/unit/test_transfers.py +++ b/tests/unit/test_transfers.py @@ -310,9 +310,9 @@ def shift(x, n): x = np.zeros(2) with xp.count_transfers() as counter: + line = _current_line() + 2 with pytest.warns(RuntimeWarning, match="copies its arrays"): kernel(x, 2) - line = _current_line() - 1 kernel(x, 2) assert np.all(x == 2.0) @@ -389,9 +389,12 @@ def scale(x, factor, n): kernel = Kernel(scale, missing_cuda="fallback") x = cp.ones(3) - with xp.count_transfers() as counter, xp.use_backend("cupy"): - with pytest.warns(RuntimeWarning): - kernel(x, 2.0, 3) + with ( + xp.count_transfers() as counter, + xp.use_backend("cupy"), + pytest.warns(RuntimeWarning), + ): + kernel(x, 2.0, 3) assert cp.all(x == 2.0) assert counter.fallbacks == 1