Evaluate on host

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.

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`) |

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

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.<Tab>` 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: ...

src/cunumpy/_host.py

@@ -0,0 +1,98 @@
+"""Run host-only code (SciPy, file readers, external libraries) with arguments of any backend."""
+
+import functools
+from collections.abc import Callable
+from typing import Any
+
+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

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