From c16f12076daf1f8749d938ff54148fac0cea8297 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 25 Sep 2026 23:01:23 +0200 Subject: [PATCH] Added cupy helpers --- CHANGELOG.md | 9 ++ docs/source/api.md | 36 ++++++++ src/cunumpy/__init__.py | 16 ++++ src/cunumpy/__init__.pyi | 9 ++ src/cunumpy/kernel.py | 26 ++++-- src/cunumpy/xp.py | 137 +++++++++++++++++++++++++++++++ tests/unit/test_cunumpy.py | 126 ++++++++++++++++++++++++++++ tests/unit/test_pyccel_kernel.py | 39 +++++++++ 8 files changed, 391 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cc84f7..1c346c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - `xp.get_array_module(array)`: Return the array-api-compat module (`numpy`/`cupy`) matching a given array's own backend, regardless of the process-wide active backend. Mirrors `cupy.get_array_module`, but works in Pyodide and returns array-api-compat modules for consistency with `xp.xp`. +- `xp.get_rng(seed=None)`: Return a `numpy.random.Generator`/`cupy.random.Generator` matching the active backend, without having to branch on the backend yourself. +- `xp.device_count()`: Number of visible CUDA devices (`0` on the NumPy backend or without a functional CuPy/CUDA install), independent of the currently active backend. +- `xp.set_device_for_rank(rank, devices_per_node=None)`: Convenience for one-MPI-rank-per-GPU codes; selects `rank % devices_per_node` (defaulting `devices_per_node` to `device_count()`) via `set_device()` and returns the chosen device id. +- `xp.memory_info()`: `(free, total)` bytes of memory on the active CUDA device, or `None` on the NumPy backend. +- `xp.free_memory()`: Release all free blocks held by CuPy's device and pinned-host memory pools (no-op on the NumPy backend). +- `xp.default_float_dtype()`: Return the active backend's `float64` dtype object, for pinning a portable float precision instead of the backend/platform-dependent `dtype=float`. +- `xp.stream()`: Context manager for a CUDA stream, to overlap host/device transfers with compute (no-op, yielding `None`, on the NumPy backend). +- `xp.pin_memory(array)`: Copy a host array into pinned (page-locked) CUDA host memory for faster transfers. +- `PyccelKernel(..., is_array=...)`: Extension point overriding the default `isinstance(value, np.ndarray)` check used to decide which host values returned by (or reachable from a declared output of) the wrapped kernel are converted back to the device -- for kernels that return/mutate a NumPy subclass or other custom host array type. - Pyodide NumPy support documentation and CI that installs the built wheel in Pyodide's WebAssembly runtime and runs compiler-free tests for arrays, conversions, contexts, and Python kernels without CuPy or Pyccel imports. - `test-compiled` extra for native Pyccel tests. The `test` extra is now compiler-free; `dev` continues to include compiled-test dependencies. - `xp.same_backend(*arrays)`: Return `True` if all given arrays live on the same backend. diff --git a/docs/source/api.md b/docs/source/api.md index 3452596..1f3b004 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -56,6 +56,40 @@ with xp.use_backend("numpy"): ### `synchronize()` Blocks until all preceding GPU operations are complete. This is a no-op when using the NumPy backend. +### `device_count()` +Returns the number of visible CUDA devices. Returns `0` on the NumPy backend or if CuPy/CUDA is unavailable. Independent of the currently active backend. + +### `set_device_for_rank(rank, devices_per_node=None)` +Convenience for one-MPI-rank-per-GPU codes: selects device `rank % devices_per_node` via `set_device()` and returns the device id chosen. `devices_per_node` defaults to `device_count()`. No-op (returns `0`) with no visible devices. + +```python +xp.set_device_for_rank(mpi_rank) # each rank picks its own GPU +``` + +### `memory_info()` +Returns `(free, total)` bytes of memory on the active CUDA device, or `None` on the NumPy backend. + +### `free_memory()` +Releases all free blocks held by CuPy's device and pinned-host memory pools. No-op on the NumPy backend. CuPy caches freed memory rather than returning it to the driver immediately, which can look like a leak in long-running processes. + +### `pin_memory(array)` +Copies a host array into pinned (page-locked) CUDA host memory, which transfers to/from the GPU faster than regular pageable memory. Raises `ImportError` if CuPy is unavailable. + +### `stream()` +Context manager for a CUDA stream, to overlap transfers and compute. No-op (yields `None`) on the NumPy backend. + +```python +with xp.stream() as s: + arr = xp.to_cupy(host_array) # enqueued on the new stream +xp.synchronize() # wait for it before reading results +``` + +### `get_rng(seed=None)` +Returns a `numpy.random.Generator` or `cupy.random.Generator` matching the active backend, so callers don't have to branch on the backend themselves. + +### `default_float_dtype()` +Returns the active backend's `float64` dtype object. NumPy and CuPy resolve Python literals and the bare `dtype=float` spelling to a platform- or backend-dependent default; pass this explicitly when a specific, portable precision matters. + ## Compiled Kernels ### `PyccelKernel(kernel, use_cupy=None, object_modules=(), outputs=None)` @@ -78,6 +112,8 @@ with xp.use_backend("cupy"): Tuples, lists and dicts are traversed recursively. Pass `object_modules` to also traverse the attributes of your own objects, e.g. `object_modules=("struphy.", "feectools.")`; instances from other modules are handed to the kernel untouched. +By default, only `numpy.ndarray` values are recognized as arrays to convert back to the device. Pass `is_array` to recognize a different (or additional) host array type instead, e.g. `is_array=lambda v: isinstance(v, (np.ndarray, np.ma.MaskedArray))`. + #### Declaring outputs By default every array that was copied to the host is copied back afterwards, since the wrapper cannot know which ones the kernel wrote to. Most pyccel kernels write to one `out` argument and only read the rest, so `outputs` lets you skip the needless transfers: diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 53438c5..d0087c5 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -6,13 +6,21 @@ from .xp import ( assert_same_backend, cupy_available, + default_float_dtype, + device_count, + free_memory, get_array_module, get_backend, + get_rng, is_cpu, is_gpu, + memory_info, + pin_memory, same_backend, set_backend, set_device, + set_device_for_rank, + stream, synchronize, to_cunumpy, to_cupy, @@ -31,14 +39,22 @@ "assert_same_backend", "cupy_available", "cupy_backend", + "default_float_dtype", + "device_count", + "free_memory", "get_array_module", "get_backend", + "get_rng", "is_cpu", "is_gpu", + "memory_info", "numpy_backend", + "pin_memory", "same_backend", "set_backend", "set_device", + "set_device_for_rank", + "stream", "synchronize", "to_cunumpy", "to_cupy", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 88a27af..9e4c15e 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -24,6 +24,15 @@ def assert_same_backend(*arrays: Any) -> None: ... def use_backend(backend: str) -> Generator[None]: ... def set_backend(backend: str) -> None: ... def set_device(device_id: int) -> None: ... +def set_device_for_rank(rank: int, devices_per_node: int | None = ...) -> int: ... +def device_count() -> int: ... +def memory_info() -> tuple[int, int] | None: ... +def free_memory() -> None: ... +def pin_memory(array: Any) -> Any: ... +@contextmanager +def stream() -> Generator[Any]: ... +def get_rng(seed: int | None = ...) -> Any: ... +def default_float_dtype() -> Any: ... def synchronize() -> None: ... numpy_backend: bool diff --git a/src/cunumpy/kernel.py b/src/cunumpy/kernel.py index eb7ba69..5258633 100644 --- a/src/cunumpy/kernel.py +++ b/src/cunumpy/kernel.py @@ -41,6 +41,12 @@ class PyccelKernel: Module prefixes (e.g. ``("struphy.", "feectools.")``) whose instances should be traversed attribute-by-attribute when looking for arrays to convert. Objects from other modules are passed through untouched. + is_array : callable, optional + Predicate deciding whether a host-side value returned by the kernel + (or reachable from a declared output) counts as an array to move + back to the device. Defaults to ``isinstance(value, np.ndarray)``. + Override this if the kernel returns/mutates a NumPy subclass or a + custom host array type that should also be converted back to CuPy. outputs : sequence of int or str, optional Which arguments the kernel writes to. Only those are copied back to the device after the call, which avoids pointless device transfers for the @@ -68,11 +74,13 @@ def __init__( kernel: Callable[..., Any], use_cupy: bool | None = None, object_modules: Sequence[str] = (), + is_array: Callable[[Any], bool] | None = None, outputs: Sequence[int | str] | None = None, ) -> None: self._kernel = kernel self._use_cupy = use_cupy self._object_modules = tuple(object_modules) + self._is_array = is_array or (lambda value: isinstance(value, np.ndarray)) if outputs is None: self._outputs: tuple[int | str, ...] | None = None @@ -165,15 +173,14 @@ def _convert_to_numpy( return value - @staticmethod - def _convert_from_numpy(value: Any) -> Any: - """Move NumPy arrays returned by the kernel back to the device.""" - if isinstance(value, np.ndarray): + def _convert_from_numpy(self, value: Any) -> Any: + """Move host arrays returned by the kernel back to the device.""" + if self._is_array(value): return to_cupy(value) if isinstance(value, tuple): - return tuple(PyccelKernel._convert_from_numpy(item) for item in value) + return tuple(self._convert_from_numpy(item) for item in value) if isinstance(value, list): - return [PyccelKernel._convert_from_numpy(item) for item in value] + return [self._convert_from_numpy(item) for item in value] return value def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> None: @@ -183,7 +190,7 @@ def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> N :meth:`_convert_to_numpy`, so that an output declared as a container or an object contributes the arrays nested inside it. """ - if isinstance(value, np.ndarray): + if self._is_array(value): found.add(id(value)) return @@ -341,6 +348,11 @@ def object_modules(self) -> tuple[str, ...]: """Module prefixes whose instances are traversed for arrays.""" return self._object_modules + @property + def is_array(self) -> Callable[[Any], bool]: + """Predicate identifying host values to convert back to the device.""" + return self._is_array + @property def outputs(self) -> tuple[int | str, ...] | None: """Declared output arguments, or ``None`` if every array is copied back.""" diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index e1a95cf..bf081b7 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os import warnings from contextlib import contextmanager @@ -134,6 +136,141 @@ def set_device(device_id: int) -> None: cp.cuda.Device(device_id).use() +def device_count() -> int: + """Number of visible CUDA devices. + + Returns 0 on the NumPy backend, or if CuPy/CUDA is not available. + Independent of the currently active backend -- this reports what + hardware is visible, not what `xp.xp` currently dispatches to. + """ + if not cupy_available(): + return 0 + + import cupy as cp + + try: + return cp.cuda.runtime.getDeviceCount() + except Exception: # noqa: BLE001 - tolerate any driver/runtime failure + return 0 + + +def set_device_for_rank(rank: int, devices_per_node: int | None = None) -> int: + """Select a CUDA device for an MPI rank, round-robin across the node. + + Convenience for one-rank-per-GPU codes: computes + ``device_id = rank % devices_per_node`` and calls `set_device()` with + it. `devices_per_node` defaults to `device_count()`. Returns the + selected device id, or 0 as a no-op if there are no visible devices. + + This assumes ranks map to devices in contiguous blocks per node (i.e. + local rank == ``rank % devices_per_node``); codes with a different + rank-to-device layout should call `set_device()` directly instead. + """ + n = devices_per_node if devices_per_node is not None else device_count() + if n == 0: + return 0 + + device_id = rank % n + set_device(device_id) + return device_id + + +def memory_info() -> tuple[int, int] | None: + """Return `(free, total)` bytes of memory on the active CUDA device. + + Returns `None` on the NumPy backend. Queries the CUDA runtime directly, + so it reflects the whole device rather than just CuPy's memory pool. + """ + if array_backend.backend != "cupy": + return None + + import cupy as cp + + return cp.cuda.runtime.memGetInfo() + + +def free_memory() -> None: + """Release all free blocks held by CuPy's memory pools (no-op on NumPy). + + CuPy caches freed device (and pinned host) memory in pools rather than + returning it to the driver/OS immediately, which can look like a leak + in long-running processes. Call this to give it back. + """ + if array_backend.backend == "cupy": + import cupy as cp + + cp.get_default_memory_pool().free_all_blocks() + cp.get_default_pinned_memory_pool().free_all_blocks() + + +def pin_memory(array: Any) -> Any: + """Copy a host array into pinned (page-locked) CUDA host memory. + + Pinned memory transfers to/from the GPU faster than regular pageable + memory, since the driver can DMA it directly. `array` must already be + on the host (use `to_numpy()` first if it may be on the GPU). Raises + `ImportError` if CuPy is not available. + """ + if not cupy_available(): + raise ImportError("CuPy is not available or not functional.") + + import cupy as cp + + array = np.asarray(array) + mem = cp.cuda.alloc_pinned_memory(array.nbytes) + pinned = np.frombuffer(mem, array.dtype, array.size).reshape(array.shape) + pinned[...] = array + return pinned + + +@contextmanager +def stream() -> Generator[Any, None, None]: + """Context manager for a CUDA stream, to overlap transfers and compute. + + On the CuPy backend, operations issued inside the block are enqueued on + a new, non-blocking stream rather than the default one. Call + `xp.synchronize()` (or the yielded stream's own `.synchronize()`) before + reading results computed inside the block. No-op on the NumPy backend, + where it yields `None`. + """ + if array_backend.backend == "cupy": + import cupy as cp + + with cp.cuda.Stream(non_blocking=True) as s: + yield s + else: + yield None + + +def get_rng(seed: int | None = None) -> Any: + """Return a random Generator matching the active backend. + + NumPy and CuPy both provide `default_rng(seed)`, returning a + `Generator` with a largely-compatible distribution API, but picking the + right one requires branching on the backend -- this does that for you. + """ + if array_backend.backend == "cupy": + import cupy as cp + + return cp.random.default_rng(seed) + + import numpy as numpy_raw + + return numpy_raw.random.default_rng(seed) + + +def default_float_dtype() -> Any: + """Return the active backend's `float64` dtype object. + + NumPy and CuPy resolve Python `int`/`float` literals and the bare + `dtype=float` spelling to a platform- or backend-dependent default + (e.g. NumPy's default integer width differs between Windows and + Linux/macOS). Use ``dtype=xp.default_float_dtype()`` instead of + ``dtype=float`` when a specific, portable precision matters. + """ + return array_backend.xp.float64 + + def synchronize() -> None: """Wait for all kernels in all streams on current device to complete.""" if array_backend.backend == "cupy": diff --git a/tests/unit/test_cunumpy.py b/tests/unit/test_cunumpy.py index a6f9755..33a1ab8 100644 --- a/tests/unit/test_cunumpy.py +++ b/tests/unit/test_cunumpy.py @@ -116,6 +116,132 @@ def test_set_backend(): assert arr2 is not None +def test_device_count_is_zero_without_cupy(): + if xp.cupy_available(): + pytest.skip("CuPy is installed/functional; device_count() may be > 0") + assert xp.device_count() == 0 + + +def test_device_count_matches_cupy_when_available(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + assert xp.device_count() == cp.cuda.runtime.getDeviceCount() + + +def test_set_device_for_rank_is_noop_without_gpus(): + if xp.cupy_available(): + pytest.skip("CuPy is installed/functional") + assert xp.set_device_for_rank(3) == 0 + + +def test_set_device_for_rank_wraps_around_devices_per_node(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + n = xp.device_count() + assert xp.set_device_for_rank(n, devices_per_node=n) == 0 + assert xp.set_device_for_rank(n + 1, devices_per_node=n) == 1 % n + + +def test_memory_info_is_none_on_numpy_backend(): + with xp.use_backend("numpy"): + assert xp.memory_info() is None + + +def test_memory_info_returns_free_and_total_on_cupy(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + with xp.use_backend("cupy"): + free, total = xp.memory_info() + assert 0 <= free <= total + + +def test_free_memory_is_noop_on_numpy_backend(): + with xp.use_backend("numpy"): + xp.free_memory() # must not raise + + +def test_free_memory_releases_cupy_pool(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + with xp.use_backend("cupy"): + _ = xp.to_cupy(np.ones(1_000)) + xp.free_memory() # must not raise + assert cp.get_default_memory_pool().n_free_blocks() == 0 + + +def test_pin_memory_requires_cupy(): + if xp.cupy_available(): + pytest.skip("CuPy is installed/functional") + with pytest.raises(ImportError): + xp.pin_memory(np.ones(3)) + + +def test_pin_memory_round_trips_values(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + arr = np.array([1.0, 2.0, 3.0]) + pinned = xp.pin_memory(arr) + assert np.array_equal(pinned, arr) + + +def test_stream_is_noop_on_numpy_backend(): + with xp.use_backend("numpy"): + with xp.stream() as s: + assert s is None + + +def test_stream_yields_a_cupy_stream_on_cupy_backend(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + with xp.use_backend("cupy"): + with xp.stream() as s: + assert isinstance(s, cp.cuda.Stream) + arr = xp.zeros(10) + assert xp.is_gpu(arr) + xp.synchronize() + + +def test_get_rng_returns_numpy_generator_on_numpy_backend(): + with xp.use_backend("numpy"): + rng = xp.get_rng(42) + assert isinstance(rng, np.random.Generator) + assert rng.random(3).shape == (3,) + + +def test_get_rng_returns_cupy_generator_on_cupy_backend(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + with xp.use_backend("cupy"): + rng = xp.get_rng(42) + assert isinstance(rng, cp.random.Generator) + + +def test_get_rng_is_reproducible_given_a_seed(): + with xp.use_backend("numpy"): + a = xp.get_rng(123).random(5) + b = xp.get_rng(123).random(5) + assert np.array_equal(a, b) + + +def test_default_float_dtype_matches_active_backend(): + with xp.use_backend("numpy"): + assert xp.default_float_dtype() == np.float64 + + if xp.cupy_available(): + import cupy as cp + + with xp.use_backend("cupy"): + assert xp.default_float_dtype() == cp.float64 + + def test_backend_bools(): with xp.use_backend("numpy"): assert xp.numpy_backend is True diff --git a/tests/unit/test_pyccel_kernel.py b/tests/unit/test_pyccel_kernel.py index 1288b6c..993b955 100644 --- a/tests/unit/test_pyccel_kernel.py +++ b/tests/unit/test_pyccel_kernel.py @@ -282,6 +282,45 @@ def kernel(obj): PyccelKernel(kernel)(holder) +# --------------------------------------------------------------------------- +# Custom is_array predicate +# --------------------------------------------------------------------------- + + +class _HostArrayLike: + """Stand-in for a custom host array type that isn't a `np.ndarray`.""" + + def __init__(self, data: np.ndarray) -> None: + self.data = data + + +def test_default_is_array_ignores_custom_array_like_return_value(): + _skip_without_cupy() + + def kernel(): + return _HostArrayLike(np.ones(3)) + + result = PyccelKernel(kernel, use_cupy=True)() + assert isinstance(result, _HostArrayLike) + assert xp.is_cpu(result.data) # not moved back: not a np.ndarray + + +def test_custom_is_array_moves_custom_return_value_back_to_device(): + _skip_without_cupy() + + def kernel(): + return _HostArrayLike(np.ones(3)) + + wrapped = PyccelKernel( + kernel, + use_cupy=True, + is_array=lambda v: isinstance(v, _HostArrayLike), + ) + result = wrapped() + assert xp.is_gpu(result) + assert wrapped.is_array is wrapped._is_array + + # --------------------------------------------------------------------------- # Declared outputs # ---------------------------------------------------------------------------