Added cupy helpers

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.

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:

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",

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

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."""

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":