From 9e904087d6b88a507dd2279ae4a419b5ca7d02fa Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 21:48:14 +0200 Subject: [PATCH 1/3] Added xp.host_call, xp.evaluate_on_host and xp.setup_on_host --- CHANGELOG.md | 7 +++ docs/source/api.md | 2 +- src/cunumpy/__init__.py | 4 ++ src/cunumpy/__init__.pyi | 5 ++- src/cunumpy/_host.py | 97 ++++++++++++++++++++++++++++++++++++++++ tests/unit/test_host.py | 82 +++++++++++++++++++++++++++++++++ 6 files changed, 195 insertions(+), 2 deletions(-) create mode 100644 src/cunumpy/_host.py create mode 100644 tests/unit/test_host.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7022e86..8772d2f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- `xp.host_call`, `xp.evaluate_on_host` and `xp.setup_on_host` run host-only code + (SciPy splines, file readers, external libraries) with arguments of either + backend: device arrays are copied to the host, the call runs on the NumPy + backend and array results are copied back. The copies are counted by + `xp.profiling.count_transfers()`. + ### Changed - The launcher detection and the serial MPI stand-in moved to the new package [maybempi](https://github.com/max-models/maybempi), a dependency of cunumpy. diff --git a/docs/source/api.md b/docs/source/api.md index 06971ac..bc79c13 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -39,7 +39,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | Submodule | Backends | Contents | |---|---|---| -| `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `scipy`, `require_version` | +| `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` | | `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | | `cunumpy.arguments` | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` | | `cunumpy.cuda` | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) | diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 1125dcc..09a469a 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -14,6 +14,7 @@ rng, xp, ) +from cunumpy._host import evaluate_on_host, host_call, setup_on_host from cunumpy._scipy_backend import scipy from cunumpy.xp import ( as_device_array, @@ -79,9 +80,11 @@ def require_version(minimum: str) -> None: "cupy_available", "cupy_backend", "default_float_dtype", + "evaluate_on_host", "get_array_backend", "get_array_module", "get_backend", + "host_call", "is_cpu", "is_gpu", "kernels", @@ -95,6 +98,7 @@ def require_version(minimum: str) -> None: "same_backend", "scipy", "set_backend", + "setup_on_host", "synchronize", "to_cunumpy", "to_cupy", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 52ab855..8e1f7b8 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -1,7 +1,7 @@ # Stub file for Pylance/mypy: exposes all numpy symbols so that # `import cunumpy as xp` followed by `xp.` shows numpy completions. # At runtime the real __init__.py dispatches to numpy or cupy via __getattr__. -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager from typing import Any @@ -38,6 +38,9 @@ def set_backend(backend: str, *, strict: bool = ...) -> None: ... def require_version(minimum: str) -> None: ... def default_float_dtype() -> Any: ... def synchronize() -> None: ... +def host_call(fun: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: ... +def evaluate_on_host(method: Callable[..., Any]) -> Callable[..., Any]: ... +def setup_on_host(init: Callable[..., None]) -> Callable[..., None]: ... def as_device_array( value: Any, dtype: Any = ..., ndim: int | None = ..., *, name: str | None = ... ) -> Any: ... diff --git a/src/cunumpy/_host.py b/src/cunumpy/_host.py new file mode 100644 index 0000000..6040ae6 --- /dev/null +++ b/src/cunumpy/_host.py @@ -0,0 +1,97 @@ +"""Run host-only code (SciPy, file readers, external libraries) with arguments of any backend.""" + +import functools +from typing import Any, Callable + +import numpy as np + +from cunumpy import xp + + +def _on_device(arg: Any) -> bool: + """Whether ``arg`` is (or contains, for tuples and lists) a device array.""" + if isinstance(arg, (tuple, list)): + return any(_on_device(a) for a in arg) + return xp.is_gpu(arg) + + +def _to_host(arg: Any) -> Any: + """Copy device arrays in ``arg`` (an array, or a tuple/list of them) to the host.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_host(a) for a in arg) + return xp.to_numpy(arg) if xp.is_gpu(arg) else arg + + +def _to_device(arg: Any) -> Any: + """Copy NumPy arrays in ``arg`` (an array, or a tuple/list of them) to the device.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_device(a) for a in arg) + return xp.to_cupy(arg) if isinstance(arg, np.ndarray) else arg + + +def host_call(fun: Callable, *args: Any, **kwargs: Any) -> Any: + """Call a host-only function with arguments of any backend. + + Device (CuPy) arrays among the arguments are copied to the host, `fun` runs + on the NumPy backend, and its array results are copied back to the device, + once per call. Without device arguments `fun` is called directly (on the + NumPy backend this is a plain call), so the result lives where the + arguments live. The copies are counted by `xp.profiling.count_transfers()`. + + Parameters + ---------- + fun + The host-only function (a SciPy spline, an external code, ...). + *args, **kwargs + Arguments of `fun`: arrays (or tuples/lists of arrays) of either + backend, or scalars. + + Returns + ------- + The result of `fun` (an array, a tuple/list of arrays or a scalar), with + arrays on the device if any argument was on the device. + + Examples + -------- + >>> values = xp.host_call(spline, x) # x on the device -> values on the device + """ + device = _on_device(args) or _on_device(list(kwargs.values())) + if xp.get_backend() == "numpy" and not device: + return fun(*args, **kwargs) + with xp.use_backend("numpy"): + out = fun(*_to_host(args), **{k: _to_host(v) for k, v in kwargs.items()}) + return _to_device(out) if device else out + + +def evaluate_on_host(method: Callable) -> Callable: + """Decorator for methods that can only be evaluated on the host (see `host_call`). + + Examples + -------- + >>> class Equilibrium: + ... @xp.evaluate_on_host + ... def pressure(self, x): + ... return scipy_spline(x) + """ + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + return host_call(method, self, *args, **kwargs) + + return wrapper + + +def setup_on_host(init: Callable) -> Callable: + """Decorator for an ``__init__`` whose setup is host-only. + + `init` runs on the NumPy backend, so the object holds only host data (NumPy + arrays, SciPy splines, floats) on either backend. Host-only parts of its + evaluation can then go through `host_call` or `evaluate_on_host`. + """ + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + with xp.use_backend("numpy"): + init(self, *args, **kwargs) + + return wrapper diff --git a/tests/unit/test_host.py b/tests/unit/test_host.py new file mode 100644 index 0000000..11456c1 --- /dev/null +++ b/tests/unit/test_host.py @@ -0,0 +1,82 @@ +"""host_call, evaluate_on_host and setup_on_host, with a stand-in for device arrays.""" + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import _host + + +class FakeDevice: + """Stands in for a CuPy array: the host helpers only convert it.""" + + def __init__(self, array): + self.array = np.asarray(array) + + +@pytest.fixture +def device(monkeypatch): + monkeypatch.setattr(_host.xp, "is_gpu", lambda a: isinstance(a, FakeDevice)) + monkeypatch.setattr( + _host.xp, "to_numpy", lambda a: a.array if isinstance(a, FakeDevice) else a + ) + monkeypatch.setattr(_host.xp, "to_cupy", FakeDevice) + + +def test_host_arrays_are_a_plain_call(): + x = np.arange(3.0) + out = xp.host_call(lambda a, scale: a * scale, x, scale=2.0) + np.testing.assert_array_equal(out, 2 * x) + + +def test_device_arguments_run_on_numpy_and_return_on_device(device): + seen = {} + + def fun(a, b, scale=1.0): + seen["backend"] = xp.get_backend() + seen["types"] = (type(a), type(b[0])) + return a + b[0], 3.0 + + a, b = FakeDevice([1.0, 2.0]), [FakeDevice([10.0, 20.0])] + arr, scalar = xp.host_call(fun, a, b, scale=2.0) + assert seen["backend"] == "numpy" and seen["types"] == (np.ndarray, np.ndarray) + assert isinstance(arr, FakeDevice) and scalar == 3.0 + np.testing.assert_array_equal(arr.array, [11.0, 22.0]) + + +def test_device_keyword_argument_is_detected(device): + out = xp.host_call(lambda x=None: x * 2, x=FakeDevice([1.0])) + assert isinstance(out, FakeDevice) + + +def test_evaluate_on_host_binds_self(device): + class Model: + offset = 1.0 + + @xp.evaluate_on_host + def f(self, x): + """Docstring.""" + assert isinstance(x, np.ndarray) + return x + self.offset + + assert Model.f.__doc__ == "Docstring." + out = Model().f(FakeDevice([1.0])) + np.testing.assert_array_equal(out.array, [2.0]) + + +def test_setup_on_host_runs_init_on_numpy_backend(): + class Model: + @xp.setup_on_host + def __init__(self, n): + self.backend = xp.get_backend() + self.data = xp.arange(n) + + model = Model(3) + assert model.backend == "numpy" and isinstance(model.data, np.ndarray) + + +def test_copies_are_counted(monkeypatch): + """With the real conversions, only to_numpy/to_cupy on device arrays count.""" + with xp.profiling.count_transfers() as counter: + xp.host_call(lambda a: a + 1, np.zeros(2)) + assert counter.total == 0 From 8b7db2f94a7fbf4833d4f7737c81343bcc5f4d1a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 21:48:54 +0200 Subject: [PATCH 2/3] Formatting --- tests/unit/test_kernel_testing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index ed01e9c..212613d 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -17,12 +17,12 @@ import cunumpy.kernel_testing from cunumpy.arguments import CudaArguments from cunumpy.cuda import parse_cuda_signature +from cunumpy.kernel_testing import backend # noqa: F401 - the fixture is used by name from cunumpy.kernel_testing import ( BACKENDS, _collect_arrays, _compare_results, assert_kernels_agree, - backend, # noqa: F401 - the fixture is used by name device_function_kernel, requires_cupy, ) From 3a30cce2d5b3183a487f700acb4ca9ed9221176a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 21:52:04 +0200 Subject: [PATCH 3/3] Formatting --- src/cunumpy/_host.py | 3 ++- tests/unit/test_kernel_testing.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/cunumpy/_host.py b/src/cunumpy/_host.py index 6040ae6..d208273 100644 --- a/src/cunumpy/_host.py +++ b/src/cunumpy/_host.py @@ -1,7 +1,8 @@ """Run host-only code (SciPy, file readers, external libraries) with arguments of any backend.""" import functools -from typing import Any, Callable +from collections.abc import Callable +from typing import Any import numpy as np diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index 212613d..ed01e9c 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -17,12 +17,12 @@ import cunumpy.kernel_testing from cunumpy.arguments import CudaArguments from cunumpy.cuda import parse_cuda_signature -from cunumpy.kernel_testing import backend # noqa: F401 - the fixture is used by name from cunumpy.kernel_testing import ( BACKENDS, _collect_arrays, _compare_results, assert_kernels_agree, + backend, # noqa: F401 - the fixture is used by name device_function_kernel, requires_cupy, )