From 48246710ae3f1df4477f0c3ce24f884796fbac17 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 5 Aug 2026 11:57:58 +0200 Subject: [PATCH 1/2] Added a PyccelKernel class --- .gitignore | 5 + CHANGELOG.md | 5 + docs/source/api.md | 26 +++ pyproject.toml | 2 +- src/cunumpy/__init__.py | 2 + src/cunumpy/__init__.pyi | 1 + src/cunumpy/kernel.py | 223 +++++++++++++++++++ tests/unit/pyccel_kernels.py | 37 ++++ tests/unit/test_pyccel_kernel.py | 367 +++++++++++++++++++++++++++++++ 9 files changed, 667 insertions(+), 1 deletion(-) create mode 100644 src/cunumpy/kernel.py create mode 100644 tests/unit/pyccel_kernels.py create mode 100644 tests/unit/test_pyccel_kernel.py diff --git a/.gitignore b/.gitignore index eaf73dd..282246a 100644 --- a/.gitignore +++ b/.gitignore @@ -172,3 +172,8 @@ cython_debug/ # Unit test output junit/ + +# Pyccel build artefacts and lock files +__pyccel__/ +.lock_acquisition.lock +stc.lock diff --git a/CHANGELOG.md b/CHANGELOG.md index 779b8fb..1a75486 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ All notable changes to the `cunumpy` library are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- `xp.PyccelKernel`: Wraps a kernel compiled with [pyccel](https://github.com/pyccel/pyccel) (which only accepts NumPy arrays) so it can be called with CuPy arrays. Arguments are copied to the host before the call, in-place kernel updates are copied back to the device, and returned arrays are moved back to the device. On the NumPy backend the kernel is called directly, without conversion. Tuples, lists and dicts are traversed recursively; pass `object_modules=("your_package.",)` to also traverse the attributes of your own objects. Conversion is identity-aware: an array reachable by several paths (passed twice, or both directly and as an object attribute) becomes a single host array, so the kernel sees the aliasing the caller intended and in-place updates are not lost; reference cycles are handled rather than recursed into. + ## [0.1.3] - 2026-07-31 ### Added diff --git a/docs/source/api.md b/docs/source/api.md index 232518d..29378d9 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -46,3 +46,29 @@ with xp.use_backend("numpy"): ### `synchronize()` Blocks until all preceding GPU operations are complete. This is a no-op when using the NumPy backend. + +## Compiled Kernels + +### `PyccelKernel(kernel, use_cupy=None, object_modules=())` +Wraps a kernel compiled with [pyccel](https://github.com/pyccel/pyccel) — which only accepts NumPy arrays — so that it can be called with CuPy arrays as well. + +On the CuPy backend the arguments are copied to the host before the call, any in-place updates the kernel makes are copied back to the device afterwards, and arrays returned by the kernel are moved back to the device. On the NumPy backend the kernel is called directly, without any conversion. + +```python +import cunumpy as xp +from my_package.kernels import axpy # pyccelized kernel + +axpy = xp.PyccelKernel(axpy) + +with xp.use_backend("cupy"): + x = xp.arange(10, dtype=xp.float64) + y = xp.ones(10, dtype=xp.float64) + out = xp.zeros(10, dtype=xp.float64) + axpy(2.0, x, y, out) # `out` is updated in place, on the GPU +``` + +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. + +Conversion is identity-aware: an array reachable by several paths — passed as two arguments, or both directly and as an attribute of a traversed object — becomes a single array on the host, so the kernel sees the aliasing the caller intended and no in-place update is lost on the way back. Reference cycles are handled rather than recursed into. + +Set `use_cupy` explicitly to force conversion on or off. By default it is decided per call: conversion happens when the active backend is CuPy, or when a CuPy array is passed in. diff --git a/pyproject.toml b/pyproject.toml index 4b8ab64..48ab7b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ optional-dependencies.docs = [ "sphinx", "sphinx-book-theme", ] -optional-dependencies.test = [ "coverage", "pytest" ] +optional-dependencies.test = [ "coverage", "pyccel", "pytest" ] urls."Source" = "https://github.com/max-models/cunumpy" [tool.setuptools.packages.find] diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index a5f42ab..73dce6b 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -2,6 +2,7 @@ from importlib.metadata import PackageNotFoundError, version from . import xp +from .kernel import PyccelKernel from .xp import ( cupy_available, get_backend, @@ -22,6 +23,7 @@ __version__ = "0.0.0+unknown" __all__ = [ + "PyccelKernel", "__version__", "cupy_available", "cupy_backend", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 8379680..a7e237b 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -8,6 +8,7 @@ import numpy as np from numpy import * from . import xp as xp +from .kernel import PyccelKernel as PyccelKernel def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... diff --git a/src/cunumpy/kernel.py b/src/cunumpy/kernel.py new file mode 100644 index 0000000..639d210 --- /dev/null +++ b/src/cunumpy/kernel.py @@ -0,0 +1,223 @@ +"""Interface for calling Pyccel-compiled kernels with CuPy arrays. + +Kernels generated by `pyccel `_ are compiled +C/Fortran routines that only understand NumPy (host) arrays. :class:`PyccelKernel` +wraps such a kernel so that it can be called transparently with CuPy (device) +arrays: the arguments are copied to the host before the call, in-place updates +made by the kernel are copied back to the device afterwards, and any arrays +returned by the kernel are moved back to the device. + +On the NumPy backend the wrapper is a no-op and the kernel is called directly. +""" + +from __future__ import annotations + +import copy +from typing import Any, Callable, Sequence + +import array_api_compat +import numpy as np + +from .xp import _cupy_backend, to_cupy, to_numpy + +__all__ = ["PyccelKernel"] + + +class PyccelKernel: + """Call a Pyccel-compiled kernel with NumPy or CuPy arrays. + + Parameters + ---------- + kernel : callable + The pyccelized kernel (or any callable expecting NumPy arrays). + use_cupy : bool, optional + Force host/device conversion on (``True``) or off (``False``). By + default (``None``) it is decided at call time: conversion happens when + the active backend is CuPy or when a CuPy array is passed in. + object_modules : sequence of str, optional + 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. + + Examples + -------- + >>> from cunumpy.kernel import PyccelKernel + >>> kernel = PyccelKernel(my_pyccelized_function) + >>> kernel(out, x, y) # `out`, `x`, `y` may be NumPy or CuPy arrays + """ + + def __init__( + self, + kernel: Callable[..., Any], + use_cupy: bool | None = None, + object_modules: Sequence[str] = (), + ) -> None: + self._kernel = kernel + self._use_cupy = use_cupy + self._object_modules = tuple(object_modules) + + def __repr__(self) -> str: + return f"PyccelKernel(kernel={self.name!r}, use_cupy={self.use_cupy!r})" + + def _convert_to_numpy( + self, + value: Any, + converted: list[tuple[Any, np.ndarray]], + memo: dict[int, Any], + ) -> Any: + """Recursively replace CuPy arrays in `value` by host copies. + + Every replacement is appended to `converted` as a + ``(device_array, host_copy)`` pair. + + `memo` maps ``id(original) -> converted`` and is shared across all + arguments of a single call. It serves two purposes: a device array + reachable by several paths is copied to the host exactly once (so the + kernel sees one shared array, as the caller intended, and the write-back + happens once), and reference cycles terminate instead of recursing + forever. Everything traversed here stays reachable from the caller's + arguments for the duration of the call, so the `id` keys cannot be + reused by unrelated objects. + """ + key = id(value) + if key in memo: + return memo[key] + + if array_api_compat.is_cupy_array(value): + value_np = to_numpy(value) + memo[key] = value_np + converted.append((value, value_np)) + return value_np + + if isinstance(value, tuple): + # A tuple cannot be memoized before its items are converted, but it + # can only take part in a cycle through a mutable container, and + # those are memoized before they are filled in below. + value_np = tuple( + self._convert_to_numpy(item, converted, memo) for item in value + ) + memo[key] = value_np + return value_np + + if isinstance(value, list): + value_np = [] + memo[key] = value_np + value_np.extend( + self._convert_to_numpy(item, converted, memo) for item in value + ) + return value_np + + if isinstance(value, dict): + value_np = {} + memo[key] = value_np + for k, v in value.items(): + value_np[k] = self._convert_to_numpy(v, converted, memo) + return value_np + + if hasattr(value, "__dict__") and value.__class__.__module__.startswith( + self._object_modules + ): + # Shallow-copy the object so the caller's instance keeps pointing at + # its device arrays; only the copy holds the host views. + value_np = copy.copy(value) + memo[key] = value_np + for name, attr in vars(value).items(): + setattr(value_np, name, self._convert_to_numpy(attr, converted, memo)) + return value_np + + 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): + return to_cupy(value) + if isinstance(value, tuple): + return tuple(PyccelKernel._convert_from_numpy(item) for item in value) + if isinstance(value, list): + return [PyccelKernel._convert_from_numpy(item) for item in value] + return value + + def _contains_cupy(self, value: Any, seen: set[int] | None = None) -> bool: + """Whether `value` holds a CuPy array, following the same traversal + rules as :meth:`_convert_to_numpy`. + + `seen` tracks already-visited containers so that reference cycles + terminate. + """ + if array_api_compat.is_cupy_array(value): + return True + + if seen is None: + seen = set() + if id(value) in seen: + return False + + if isinstance(value, (tuple, list)): + seen.add(id(value)) + return any(self._contains_cupy(item, seen) for item in value) + + if isinstance(value, dict): + seen.add(id(value)) + return any(self._contains_cupy(item, seen) for item in value.values()) + + if hasattr(value, "__dict__") and value.__class__.__module__.startswith( + self._object_modules + ): + seen.add(id(value)) + return any(self._contains_cupy(attr, seen) for attr in vars(value).values()) + + return False + + def _needs_conversion(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: + if self._use_cupy is not None: + return self._use_cupy + if _cupy_backend(): + return True + # The backend is NumPy, but individual CuPy arrays may still have been + # passed in explicitly. + return any(self._contains_cupy(value) for value in (*args, *kwargs.values())) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + if not self._needs_conversion(args, kwargs): + return self._kernel(*args, **kwargs) + + # Convert CuPy arrays in args/kwargs to NumPy arrays on the host. The + # memo is shared across args and kwargs so that an array passed more + # than once stays a single array on the host too. + converted: list[tuple[Any, np.ndarray]] = [] + memo: dict[int, Any] = {} + args_np = [self._convert_to_numpy(x, converted, memo) for x in args] + kwargs_np = { + k: self._convert_to_numpy(v, converted, memo) for k, v in kwargs.items() + } + + result = self._kernel(*args_np, **kwargs_np) + + # Copy in-place kernel updates back to the device arrays. + for device_array, host_array in converted: + device_array[...] = to_cupy(host_array) + + return self._convert_from_numpy(result) + + @property + def name(self) -> str: + """Name of the wrapped kernel.""" + return getattr(self._kernel, "__name__", type(self._kernel).__name__) + + @property + def kernel(self) -> Callable[..., Any]: + """The wrapped kernel.""" + return self._kernel + + @property + def use_cupy(self) -> bool: + """Whether calls currently convert between device and host arrays.""" + if self._use_cupy is not None: + return self._use_cupy + return _cupy_backend() + + @property + def object_modules(self) -> tuple[str, ...]: + """Module prefixes whose instances are traversed for arrays.""" + return self._object_modules diff --git a/tests/unit/pyccel_kernels.py b/tests/unit/pyccel_kernels.py new file mode 100644 index 0000000..8ab2726 --- /dev/null +++ b/tests/unit/pyccel_kernels.py @@ -0,0 +1,37 @@ +"""Example kernels written for compilation with pyccel. + +This module is plain Python annotated with pyccel's array type hints, so it can +be imported and run as-is, or compiled to C/Fortran with +``pyccel.epyccel(pyccel_kernels, language="c")``. The tests in +`test_pyccel_kernel.py` compile it and drive the compiled kernels through +:class:`cunumpy.PyccelKernel`. +""" + + +def axpy(a: float, x: "float[:]", y: "float[:]", out: "float[:]"): + """Write ``a * x + y`` into `out` (in-place, no return value).""" + for i in range(x.shape[0]): + out[i] = a * x[i] + y[i] + + +def scale_inplace(x: "float[:]", factor: float): + """Multiply `x` by `factor`, in place.""" + for i in range(x.shape[0]): + x[i] = x[i] * factor + + +def dot(x: "float[:]", y: "float[:]") -> float: + """Return the dot product of `x` and `y` (a scalar return value).""" + result = 0.0 + for i in range(x.shape[0]): + result += x[i] * y[i] + return result + + +def matvec(mat: "float[:,:]", vec: "float[:]", out: "float[:]"): + """Write ``mat @ vec`` into `out` (2D input, in-place output).""" + for i in range(mat.shape[0]): + acc = 0.0 + for j in range(mat.shape[1]): + acc += mat[i, j] * vec[j] + out[i] = acc diff --git a/tests/unit/test_pyccel_kernel.py b/tests/unit/test_pyccel_kernel.py new file mode 100644 index 0000000..b42ee7c --- /dev/null +++ b/tests/unit/test_pyccel_kernel.py @@ -0,0 +1,367 @@ +"""Tests for `cunumpy.PyccelKernel`. + +Two groups of tests live here: + +* conversion tests, which use plain Python callables as stand-in kernels and + run everywhere; +* end-to-end tests, which compile `pyccel_kernels.py` with pyccel and drive the + compiled kernels through `PyccelKernel`. They are skipped when pyccel (or a + working compiler) is unavailable. + +The CuPy-side assertions can only be exercised on a machine with a GPU; the +NumPy-side assertions run everywhere. +""" + +import shutil +from pathlib import Path + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import PyccelKernel + +KERNEL_SOURCE = Path(__file__).parent / "pyccel_kernels.py" + + +def _skip_without_cupy(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + + +@pytest.fixture(scope="module") +def kernels(tmp_path_factory): + """Compile `pyccel_kernels.py` with pyccel and return the compiled module. + + The source is copied into a temporary directory first so that pyccel's + build artefacts (`__pyccel__/`) never land in the repository. + """ + pyccel = pytest.importorskip("pyccel", reason="pyccel is not installed") + + import importlib.util + import sys + + build_dir = tmp_path_factory.mktemp("pyccel_build") + source = build_dir / KERNEL_SOURCE.name + shutil.copy(KERNEL_SOURCE, source) + + spec = importlib.util.spec_from_file_location(source.stem, source) + module = importlib.util.module_from_spec(spec) + sys.modules[source.stem] = module + spec.loader.exec_module(module) + + try: + return pyccel.epyccel(module, language="c") + except Exception as exc: # noqa: BLE001 - no compiler / broken toolchain + pytest.skip(f"pyccel could not compile the example kernels: {exc}") + finally: + sys.modules.pop(source.stem, None) + + +# --------------------------------------------------------------------------- +# Conversion behaviour (no pyccel needed) +# --------------------------------------------------------------------------- + + +def test_name_kernel_and_repr(): + def my_kernel(x): + return x + + wrapped = PyccelKernel(my_kernel, use_cupy=False) + + assert wrapped.name == "my_kernel" + assert wrapped.kernel is my_kernel + assert wrapped.use_cupy is False + assert repr(wrapped) == "PyccelKernel(kernel='my_kernel', use_cupy=False)" + + +def test_numpy_backend_calls_kernel_unchanged(): + """On the NumPy backend the arguments must reach the kernel untouched.""" + seen = {} + + def kernel(x, *, scale): + seen["x"] = x + seen["scale"] = scale + return x * scale + + arr = np.arange(4, dtype=float) + with xp.use_backend("numpy"): + result = PyccelKernel(kernel)(arr, scale=2.0) + + assert seen["x"] is arr + assert seen["scale"] == 2.0 + assert np.array_equal(result, arr * 2.0) + + +def test_use_cupy_follows_active_backend(): + wrapped = PyccelKernel(lambda: None) + + with xp.use_backend("numpy"): + assert wrapped.use_cupy is False + + _skip_without_cupy() + + with xp.use_backend("cupy"): + assert wrapped.use_cupy is True + + +def test_explicit_use_cupy_overrides_backend(): + wrapped = PyccelKernel(lambda: None, use_cupy=False) + + _skip_without_cupy() + + with xp.use_backend("cupy"): + assert wrapped.use_cupy is False + + +def test_kernel_receives_numpy_arrays_when_given_cupy_arrays(): + _skip_without_cupy() + + seen = {} + + def kernel(x, y): + seen["types"] = (type(x), type(y)) + + gpu = xp.to_cupy(np.arange(4, dtype=float)) + PyccelKernel(kernel)(gpu, y=gpu) + + assert seen["types"] == (np.ndarray, np.ndarray) + + +def test_inplace_updates_are_copied_back_to_device(): + _skip_without_cupy() + + def kernel(out): + out[:] = 42.0 + + gpu = xp.to_cupy(np.zeros(5)) + PyccelKernel(kernel)(gpu) + + assert np.array_equal(xp.to_numpy(gpu), np.full(5, 42.0)) + + +def test_returned_arrays_are_moved_back_to_device(): + _skip_without_cupy() + + def kernel(x): + return x + 1.0, x.sum(), None + + gpu = xp.to_cupy(np.arange(3, dtype=float)) + arr, total, nothing = PyccelKernel(kernel)(gpu) + + assert xp.is_gpu(arr) + assert np.array_equal(xp.to_numpy(arr), np.arange(3, dtype=float) + 1.0) + assert total == 3.0 + assert nothing is None + + +def test_nested_containers_are_converted_and_written_back(): + _skip_without_cupy() + + def kernel(pair, mapping): + pair[0][:] = 1.0 + pair[1][0][:] = 2.0 + mapping["a"][:] = 3.0 + assert all(isinstance(a, np.ndarray) for a in (pair[0], pair[1][0])) + assert isinstance(mapping["a"], np.ndarray) + + first = xp.to_cupy(np.zeros(3)) + second = xp.to_cupy(np.zeros(3)) + third = xp.to_cupy(np.zeros(3)) + + PyccelKernel(kernel)((first, [second]), {"a": third}) + + assert np.array_equal(xp.to_numpy(first), np.full(3, 1.0)) + assert np.array_equal(xp.to_numpy(second), np.full(3, 2.0)) + assert np.array_equal(xp.to_numpy(third), np.full(3, 3.0)) + + +def test_object_attributes_are_converted_when_module_is_listed(): + """Instances of listed modules are traversed; others are passed through.""" + _skip_without_cupy() + + class Container: + def __init__(self, data): + self.data = data + self.label = "unchanged" + + holder = Container(xp.to_cupy(np.zeros(3))) + + def kernel(obj): + assert isinstance(obj.data, np.ndarray) + assert obj.label == "unchanged" + obj.data[:] = 7.0 + + # `Container` is defined in this test module, so use its own module name. + PyccelKernel(kernel, object_modules=(Container.__module__,))(holder) + + assert xp.is_gpu(holder.data), "the caller's object must keep its device array" + assert np.array_equal(xp.to_numpy(holder.data), np.full(3, 7.0)) + + +def test_aliased_arguments_stay_a_single_host_array(): + """One device array passed twice must become one host array, so the + kernel's in-place updates are not lost when copying back.""" + _skip_without_cupy() + + def kernel(a, b): + assert a is b + a += 1.0 + b += 1.0 + + shared = xp.to_cupy(np.zeros(3)) + PyccelKernel(kernel)(shared, shared) + + assert np.array_equal(xp.to_numpy(shared), np.full(3, 2.0)) + + +def test_aliasing_between_object_attribute_and_argument(): + _skip_without_cupy() + + class Container: + def __init__(self, data): + self.data = data + + holder = Container(xp.to_cupy(np.zeros(3))) + + def kernel(arr, obj): + assert arr is obj.data + obj.data[:] = 5.0 + + PyccelKernel(kernel, object_modules=(Container.__module__,))(holder.data, holder) + + assert np.array_equal(xp.to_numpy(holder.data), np.full(3, 5.0)) + + +def test_reference_cycles_do_not_recurse_forever(): + _skip_without_cupy() + + cyclic = [xp.to_cupy(np.zeros(2))] + cyclic.append(cyclic) + + def kernel(items): + assert items[1] is items, "the cycle must be preserved" + items[0][:] = 4.0 + + PyccelKernel(kernel)(cyclic) + + assert np.array_equal(xp.to_numpy(cyclic[0]), np.full(2, 4.0)) + + +def test_dict_argument_is_not_mistaken_for_a_device_array(): + """Dict-like objects expose `.get`; detection must not rely on that.""" + seen = {} + + def kernel(cfg, arr): + seen["cfg"] = cfg + arr[:] = 8.0 + + arr = np.zeros(2) + PyccelKernel(kernel)({"n": 3}, arr) + + assert seen["cfg"] == {"n": 3} + assert np.array_equal(arr, np.full(2, 8.0)) + + +def test_unlisted_objects_are_passed_through_untouched(): + _skip_without_cupy() + + class Container: + def __init__(self, data): + self.data = data + + holder = Container(xp.to_cupy(np.zeros(3))) + + def kernel(obj): + assert obj is holder + assert xp.is_gpu(obj.data) + + PyccelKernel(kernel)(holder) + + +# --------------------------------------------------------------------------- +# End-to-end with real pyccel-compiled kernels +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +def test_compiled_axpy(kernels, backend): + if backend == "cupy": + _skip_without_cupy() + + with xp.use_backend(backend): + a = 2.5 + x = xp.asarray(np.arange(6, dtype=np.float64)) + y = xp.asarray(np.ones(6, dtype=np.float64)) + out = xp.zeros(6, dtype=np.float64) + + PyccelKernel(kernels.axpy)(a, x, y, out) + + assert np.allclose(xp.to_numpy(out), a * np.arange(6) + 1.0) + + +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +def test_compiled_scale_inplace(kernels, backend): + if backend == "cupy": + _skip_without_cupy() + + with xp.use_backend(backend): + x = xp.asarray(np.arange(4, dtype=np.float64)) + + PyccelKernel(kernels.scale_inplace)(x, 3.0) + + assert xp.get_backend(x) == backend + assert np.allclose(xp.to_numpy(x), 3.0 * np.arange(4)) + + +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +def test_compiled_dot_returns_scalar(kernels, backend): + if backend == "cupy": + _skip_without_cupy() + + with xp.use_backend(backend): + x = xp.asarray(np.arange(5, dtype=np.float64)) + y = xp.asarray(np.arange(5, dtype=np.float64)) + + result = PyccelKernel(kernels.dot)(x, y) + + assert isinstance(result, float) + assert result == pytest.approx(float(np.arange(5) @ np.arange(5))) + + +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +def test_compiled_matvec(kernels, backend): + if backend == "cupy": + _skip_without_cupy() + + mat_np = np.arange(6, dtype=np.float64).reshape(3, 2) + vec_np = np.array([1.0, 2.0]) + + with xp.use_backend(backend): + mat = xp.asarray(mat_np) + vec = xp.asarray(vec_np) + out = xp.zeros(3, dtype=np.float64) + + PyccelKernel(kernels.matvec)(mat, vec, out) + + assert np.allclose(xp.to_numpy(out), mat_np @ vec_np) + + +def test_compiled_kernel_matches_pure_python(kernels): + """The compiled kernel and its Python original must agree.""" + import importlib.util + + spec = importlib.util.spec_from_file_location("pyccel_kernels_ref", KERNEL_SOURCE) + reference = importlib.util.module_from_spec(spec) + spec.loader.exec_module(reference) + + x = np.random.rand(10) + y = np.random.rand(10) + + expected = np.zeros(10) + reference.axpy(1.5, x, y, expected) + + got = np.zeros(10) + PyccelKernel(kernels.axpy, use_cupy=False)(1.5, x, y, got) + + assert np.allclose(got, expected) From c7f7d3b7ebb66cd3050e769f319e887d8a7d3279 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 5 Aug 2026 13:19:03 +0200 Subject: [PATCH 2/2] Added outputs --- CHANGELOG.md | 2 +- docs/source/api.md | 20 +++- src/cunumpy/kernel.py | 125 +++++++++++++++++++++++- tests/unit/test_pyccel_kernel.py | 160 ++++++++++++++++++++++++++++++- 4 files changed, 302 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a75486..9be463e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added -- `xp.PyccelKernel`: Wraps a kernel compiled with [pyccel](https://github.com/pyccel/pyccel) (which only accepts NumPy arrays) so it can be called with CuPy arrays. Arguments are copied to the host before the call, in-place kernel updates are copied back to the device, and returned arrays are moved back to the device. On the NumPy backend the kernel is called directly, without conversion. Tuples, lists and dicts are traversed recursively; pass `object_modules=("your_package.",)` to also traverse the attributes of your own objects. Conversion is identity-aware: an array reachable by several paths (passed twice, or both directly and as an object attribute) becomes a single host array, so the kernel sees the aliasing the caller intended and in-place updates are not lost; reference cycles are handled rather than recursed into. +- `xp.PyccelKernel`: Wraps a kernel compiled with [pyccel](https://github.com/pyccel/pyccel) (which only accepts NumPy arrays) so it can be called with CuPy arrays. Arguments are copied to the host before the call, in-place kernel updates are copied back to the device, and returned arrays are moved back to the device. On the NumPy backend the kernel is called directly, without conversion. Tuples, lists and dicts are traversed recursively; pass `object_modules=("your_package.",)` to also traverse the attributes of your own objects. Pass `outputs=(5,)` (indices for positional arguments, names for keyword arguments) to declare which arguments the kernel writes to, so only those are copied back to the device instead of every converted array; `outputs=()` declares none. Conversion is identity-aware: an array reachable by several paths (passed twice, or both directly and as an object attribute) becomes a single host array, so the kernel sees the aliasing the caller intended and in-place updates are not lost; reference cycles are handled rather than recursed into. ## [0.1.3] - 2026-07-31 diff --git a/docs/source/api.md b/docs/source/api.md index 29378d9..178d6df 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -49,7 +49,7 @@ Blocks until all preceding GPU operations are complete. This is a no-op when usi ## Compiled Kernels -### `PyccelKernel(kernel, use_cupy=None, object_modules=())` +### `PyccelKernel(kernel, use_cupy=None, object_modules=(), outputs=None)` Wraps a kernel compiled with [pyccel](https://github.com/pyccel/pyccel) — which only accepts NumPy arrays — so that it can be called with CuPy arrays as well. On the CuPy backend the arguments are copied to the host before the call, any in-place updates the kernel makes are copied back to the device afterwards, and arrays returned by the kernel are moved back to the device. On the NumPy backend the kernel is called directly, without any conversion. @@ -69,6 +69,24 @@ 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. +#### 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: + +```python +interpolate = xp.PyccelKernel(some_interpolation_kernel, outputs=(5,)) + +interpolate(x, y, z, basis, coeffs, out) # `out` is argument 5 +``` + +Only the declared arguments are copied back; `basis` and `coeffs` make the trip to the host and no further. Containers and traversed objects may be declared too — every array nested inside them is copied back. An array that is *also* reachable from a declared output (e.g. passed as both an input and the output) is still copied back. + +Declare positional arguments by index (negatives count from the end) and keyword arguments by name, e.g. `outputs=("out",)` for `interpolate(..., out=out)`. The two are not interchangeable: pyccel-compiled kernels are builtins with no introspectable signature, so the wrapper cannot map a name onto a position. A declaration that matches no argument of the call raises `IndexError`/`KeyError` rather than silently copying nothing back. + +`outputs=()` declares that the kernel writes to none of its arguments. Leaving `outputs` unset keeps the always-correct default. Note that a *wrong* declaration is a silent-wrong-answer bug: an argument the kernel writes to but that you did not declare keeps its stale values on the GPU. + +#### Aliasing and cycles + Conversion is identity-aware: an array reachable by several paths — passed as two arguments, or both directly and as an attribute of a traversed object — becomes a single array on the host, so the kernel sees the aliasing the caller intended and no in-place update is lost on the way back. Reference cycles are handled rather than recursed into. Set `use_cupy` explicitly to force conversion on or off. By default it is decided per call: conversion happens when the active backend is CuPy, or when a CuPy array is passed in. diff --git a/src/cunumpy/kernel.py b/src/cunumpy/kernel.py index 639d210..88fb819 100644 --- a/src/cunumpy/kernel.py +++ b/src/cunumpy/kernel.py @@ -38,6 +38,20 @@ 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. + 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 + (usually much larger) read-only inputs. Positional arguments are named + by index, keyword arguments by name:: + + interpolate = PyccelKernel(some_interpolation_kernel, outputs=(5,)) + interpolate(x, y, z, basis, coeffs, out) # `out` is argument 5 + + Pyccel-compiled kernels are builtins with no introspectable signature, + so an index and a name are *not* interchangeable: declare the form you + actually call with. An empty sequence declares that the kernel writes to + none of its arguments. By default (``None``) every converted array is + copied back, which is always correct but does more work. Examples -------- @@ -51,13 +65,34 @@ def __init__( kernel: Callable[..., Any], use_cupy: bool | None = None, object_modules: Sequence[str] = (), + outputs: Sequence[int | str] | None = None, ) -> None: self._kernel = kernel self._use_cupy = use_cupy self._object_modules = tuple(object_modules) + if outputs is None: + self._outputs: tuple[int | str, ...] | None = None + else: + if isinstance(outputs, (int, str)): + raise TypeError( + "outputs must be a sequence of argument indices/names, " + f"not a bare {type(outputs).__name__} " + f"(did you mean outputs=({outputs!r},)?)" + ) + for entry in outputs: + if not isinstance(entry, (int, str)) or isinstance(entry, bool): + raise TypeError( + "outputs entries must be argument indices (int) or " + f"names (str), got {entry!r}" + ) + self._outputs = tuple(outputs) + def __repr__(self) -> str: - return f"PyccelKernel(kernel={self.name!r}, use_cupy={self.use_cupy!r})" + return ( + f"PyccelKernel(kernel={self.name!r}, use_cupy={self.use_cupy!r}, " + f"outputs={self._outputs!r})" + ) def _convert_to_numpy( self, @@ -138,6 +173,77 @@ def _convert_from_numpy(value: Any) -> Any: return [PyccelKernel._convert_from_numpy(item) for item in value] return value + def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> None: + """Record the id of every host array reachable from `value`. + + Runs over the *converted* arguments, using the same traversal rules as + :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): + found.add(id(value)) + return + + if id(value) in seen: + return + + if isinstance(value, (tuple, list)): + seen.add(id(value)) + for item in value: + self._collect_host_arrays(item, found, seen) + return + + if isinstance(value, dict): + seen.add(id(value)) + for item in value.values(): + self._collect_host_arrays(item, found, seen) + return + + if hasattr(value, "__dict__") and value.__class__.__module__.startswith( + self._object_modules + ): + seen.add(id(value)) + for attr in vars(value).values(): + self._collect_host_arrays(attr, found, seen) + + def _output_host_arrays( + self, args_np: list[Any], kwargs_np: dict[str, Any] + ) -> set[int]: + """Ids of the host arrays reachable from the declared output arguments. + + Raises + ------ + IndexError, KeyError + If a declared output does not correspond to an argument of this + call -- typically because an argument declared by index was passed + as a keyword, or vice versa. + """ + found: set[int] = set() + seen: set[int] = set() + + for entry in self._outputs or (): + if isinstance(entry, int): + index = entry + len(args_np) if entry < 0 else entry + if not 0 <= index < len(args_np): + raise IndexError( + f"{self.name}() was declared with output argument " + f"{entry}, but was called with {len(args_np)} " + "positional argument(s). Note that an output passed as " + "a keyword must be declared by name, not by index." + ) + self._collect_host_arrays(args_np[index], found, seen) + else: + if entry not in kwargs_np: + raise KeyError( + f"{self.name}() was declared with output argument " + f"{entry!r}, but no such keyword argument was passed. " + "Note that an output passed positionally must be " + "declared by index, not by name." + ) + self._collect_host_arrays(kwargs_np[entry], found, seen) + + return found + def _contains_cupy(self, value: Any, seen: set[int] | None = None) -> bool: """Whether `value` holds a CuPy array, following the same traversal rules as :meth:`_convert_to_numpy`. @@ -192,11 +298,21 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: k: self._convert_to_numpy(v, converted, memo) for k, v in kwargs.items() } + # Which arrays the kernel may have written to is resolved before the + # call, so a mis-declared output is reported even if the kernel itself + # would have raised first. + writeable = ( + None + if self._outputs is None + else self._output_host_arrays(args_np, kwargs_np) + ) + result = self._kernel(*args_np, **kwargs_np) # Copy in-place kernel updates back to the device arrays. for device_array, host_array in converted: - device_array[...] = to_cupy(host_array) + if writeable is None or id(host_array) in writeable: + device_array[...] = to_cupy(host_array) return self._convert_from_numpy(result) @@ -221,3 +337,8 @@ def use_cupy(self) -> bool: def object_modules(self) -> tuple[str, ...]: """Module prefixes whose instances are traversed for arrays.""" return self._object_modules + + @property + def outputs(self) -> tuple[int | str, ...] | None: + """Declared output arguments, or ``None`` if every array is copied back.""" + return self._outputs diff --git a/tests/unit/test_pyccel_kernel.py b/tests/unit/test_pyccel_kernel.py index b42ee7c..1288b6c 100644 --- a/tests/unit/test_pyccel_kernel.py +++ b/tests/unit/test_pyccel_kernel.py @@ -72,7 +72,10 @@ def my_kernel(x): assert wrapped.name == "my_kernel" assert wrapped.kernel is my_kernel assert wrapped.use_cupy is False - assert repr(wrapped) == "PyccelKernel(kernel='my_kernel', use_cupy=False)" + assert wrapped.outputs is None + assert repr(wrapped) == ( + "PyccelKernel(kernel='my_kernel', use_cupy=False, outputs=None)" + ) def test_numpy_backend_calls_kernel_unchanged(): @@ -279,6 +282,124 @@ def kernel(obj): PyccelKernel(kernel)(holder) +# --------------------------------------------------------------------------- +# Declared outputs +# --------------------------------------------------------------------------- + + +def test_outputs_restricts_write_back_to_declared_arguments(): + """Inputs must not be copied back, even if the kernel writes to them.""" + _skip_without_cupy() + + def kernel(x, y, out): + x[:] = 999.0 # a stray write to an input + out[:] = x[0] + y[0] + + x = xp.to_cupy(np.ones(3)) + y = xp.to_cupy(np.full(3, 2.0)) + out = xp.to_cupy(np.zeros(3)) + + PyccelKernel(kernel, outputs=(2,))(x, y, out) + + assert np.array_equal(xp.to_numpy(x), np.ones(3)) + assert np.array_equal(xp.to_numpy(out), np.full(3, 1001.0)) + + +def test_outputs_accepts_negative_indices_and_keyword_names(): + _skip_without_cupy() + + def kernel(x, out): + out[:] = 7.0 + + positional = xp.to_cupy(np.zeros(2)) + PyccelKernel(kernel, outputs=(-1,))(xp.to_cupy(np.ones(2)), positional) + assert np.array_equal(xp.to_numpy(positional), np.full(2, 7.0)) + + keyword = xp.to_cupy(np.zeros(2)) + PyccelKernel(kernel, outputs=("out",))(xp.to_cupy(np.ones(2)), out=keyword) + assert np.array_equal(xp.to_numpy(keyword), np.full(2, 7.0)) + + +def test_empty_outputs_copies_nothing_back(): + _skip_without_cupy() + + def kernel(out): + out[:] = 5.0 + + arr = xp.to_cupy(np.zeros(2)) + PyccelKernel(kernel, outputs=())(arr) + + assert np.array_equal(xp.to_numpy(arr), np.zeros(2)) + + +def test_outputs_traverse_nested_containers_and_objects(): + _skip_without_cupy() + + class Container: + def __init__(self, data): + self.data = data + + nested = xp.to_cupy(np.zeros(2)) + held = xp.to_cupy(np.zeros(2)) + + def kernel(inp, pack, obj): + pack[0]["m"][:] = 1.0 + obj.data[:] = 2.0 + + PyccelKernel(kernel, object_modules=(Container.__module__,), outputs=(1, 2))( + xp.to_cupy(np.ones(2)), [{"m": nested}], Container(held) + ) + + assert np.array_equal(xp.to_numpy(nested), np.full(2, 1.0)) + assert np.array_equal(xp.to_numpy(held), np.full(2, 2.0)) + + +def test_array_aliased_into_an_output_is_written_back(): + """An input that is also the output must still come back.""" + _skip_without_cupy() + + def kernel(inp, out): + out[:] = 6.0 + + shared = xp.to_cupy(np.zeros(2)) + PyccelKernel(kernel, outputs=(1,))(shared, shared) + + assert np.array_equal(xp.to_numpy(shared), np.full(2, 6.0)) + + +def test_misdeclared_output_index_raises(): + """`use_cupy=True` exercises the conversion path without needing a GPU: + NumPy arguments need no conversion, so nothing is sent to a device.""" + wrapped = PyccelKernel(lambda out: None, use_cupy=True, outputs=(5,)) + + with pytest.raises(IndexError, match="positional argument"): + wrapped(np.zeros(2)) + + +def test_misdeclared_output_name_raises(): + wrapped = PyccelKernel(lambda out: None, use_cupy=True, outputs=("nope",)) + + with pytest.raises(KeyError, match="no such keyword argument"): + wrapped(np.zeros(2)) + + +@pytest.mark.parametrize("bad", [5, "out", (None,), (1.5,), (True,)]) +def test_invalid_outputs_rejected_at_construction(bad): + with pytest.raises(TypeError): + PyccelKernel(lambda: None, outputs=bad) + + +def test_outputs_is_ignored_on_the_numpy_path(): + """Without conversion there is no copy-back to skip: the kernel writes + straight into the caller's arrays.""" + arr = np.zeros(2) + + with xp.use_backend("numpy"): + PyccelKernel(lambda out: out.__setitem__(slice(None), 4.0), outputs=())(arr) + + assert np.array_equal(arr, np.full(2, 4.0)) + + # --------------------------------------------------------------------------- # End-to-end with real pyccel-compiled kernels # --------------------------------------------------------------------------- @@ -347,6 +468,43 @@ def test_compiled_matvec(kernels, backend): assert np.allclose(xp.to_numpy(out), mat_np @ vec_np) +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +def test_compiled_axpy_with_declared_output(kernels, backend): + """The `outputs=` form from the issue: only `out` is copied back.""" + if backend == "cupy": + _skip_without_cupy() + + axpy = PyccelKernel(kernels.axpy, outputs=(3,)) + + with xp.use_backend(backend): + x = xp.asarray(np.arange(6, dtype=np.float64)) + y = xp.asarray(np.ones(6, dtype=np.float64)) + out = xp.zeros(6, dtype=np.float64) + + axpy(2.5, x, y, out) + + assert np.allclose(xp.to_numpy(out), 2.5 * np.arange(6) + 1.0) + # The inputs are untouched either way, but assert it explicitly since + # they are the arrays whose copy-back we skipped. + assert np.allclose(xp.to_numpy(x), np.arange(6)) + assert np.allclose(xp.to_numpy(y), np.ones(6)) + + +def test_compiled_scale_inplace_needs_its_argument_declared(kernels): + """`scale_inplace` writes to argument 0; declaring no outputs loses that + update on the GPU, which is exactly what the declaration is for.""" + _skip_without_cupy() + + with xp.use_backend("cupy"): + declared = xp.asarray(np.arange(4, dtype=np.float64)) + PyccelKernel(kernels.scale_inplace, outputs=(0,))(declared, 3.0) + assert np.allclose(xp.to_numpy(declared), 3.0 * np.arange(4)) + + undeclared = xp.asarray(np.arange(4, dtype=np.float64)) + PyccelKernel(kernels.scale_inplace, outputs=())(undeclared, 3.0) + assert np.allclose(xp.to_numpy(undeclared), np.arange(4)) + + def test_compiled_kernel_matches_pure_python(kernels): """The compiled kernel and its Python original must agree.""" import importlib.util