cunumpy 0.1.3__tar.gz → 0.1.4__tar.gz

{cunumpy-0.1.3 → cunumpy-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cunumpy
-Version: 0.1.3
+Version: 0.1.4
 Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
 Author: Max
 Project-URL: Source, https://github.com/max-models/cunumpy
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+Requires-Dist: array-api-compat
 Requires-Dist: numpy
 Provides-Extra: dev
 Requires-Dist: black[jupyter]; extra == "dev"
@@ -32,6 +33,7 @@ Requires-Dist: sphinx; extra == "docs"
 Requires-Dist: sphinx-book-theme; extra == "docs"
 Provides-Extra: test
 Requires-Dist: coverage; extra == "test"
+Requires-Dist: pyccel; extra == "test"
 Requires-Dist: pytest; extra == "test"
 
 # CuNumpy

{cunumpy-0.1.3 → cunumpy-0.1.4}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
 
 [project]
 name = "cunumpy"
-version = "0.1.3"
+version = "0.1.4"
 description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`."
 readme = "README.md"
 keywords = [ "python" ]
@@ -23,6 +23,7 @@ classifiers = [
   "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
+  "array-api-compat",
   "numpy",
 ]
 
@@ -43,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]

{cunumpy-0.1.3 → cunumpy-0.1.4}/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 @@ except PackageNotFoundError:
     __version__ = "0.0.0+unknown"
 
 __all__ = [
+    "PyccelKernel",
     "__version__",
     "cupy_available",
     "cupy_backend",

{cunumpy-0.1.3 → cunumpy-0.1.4}/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: ...

cunumpy-0.1.4/src/cunumpy/kernel.py

@@ -0,0 +1,344 @@
+"""Interface for calling Pyccel-compiled kernels with CuPy arrays.
+
+Kernels generated by `pyccel <https://github.com/pyccel/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.
+    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
+    --------
+    >>> 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] = (),
+        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}, "
+            f"outputs={self._outputs!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 _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`.
+
+        `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()
+        }
+
+        # 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:
+            if writeable is None or id(host_array) in writeable:
+                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
+
+    @property
+    def outputs(self) -> tuple[int | str, ...] | None:
+        """Declared output arguments, or ``None`` if every array is copied back."""
+        return self._outputs

{cunumpy-0.1.3 → cunumpy-0.1.4}/src/cunumpy/xp.py

@@ -4,7 +4,8 @@ from contextlib import contextmanager
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, Generator, Literal
 
-import numpy as np
+import array_api_compat
+import array_api_compat.numpy as np
 
 BackendType = Literal["numpy", "cupy"]
 
@@ -55,7 +56,7 @@ class ArrayBackend:
     def _load_backend(self, backend: BackendType, verbose: bool = False) -> ModuleType:
         if backend == "cupy":
             if cupy_available():
-                import cupy as cp
+                import array_api_compat.cupy as cp
 
                 self._backend = "cupy"
                 return cp
@@ -66,10 +67,8 @@ class ArrayBackend:
                     )
                 self._backend = "numpy"
                 return np
-        import numpy as np_mod
-
         self._backend = "numpy"
-        return np_mod
+        return np
 
     def __repr__(self) -> str:
         return f"ArrayBackend(backend={self._backend!r}, module={self._xp.__name__!r})"
@@ -166,7 +165,7 @@ def to_cupy(array: Any) -> Any:
     if not cupy_available():
         raise ImportError("CuPy is not available or not functional.")
 
-    import cupy as cp
+    import array_api_compat.cupy as cp
 
     return cp.asarray(array)
 
@@ -180,8 +179,7 @@ def to_cunumpy(array: Any) -> Any:
 
 def get_backend(array: Any) -> BackendType:
     """Return 'cupy' or 'numpy' depending on the array type."""
-    module = getattr(type(array), "__module__", "")
-    return "cupy" if "cupy" in module else "numpy"
+    return "cupy" if array_api_compat.is_cupy_array(array) else "numpy"
 
 
 def is_gpu(array: Any) -> bool:

{cunumpy-0.1.3 → cunumpy-0.1.4}/src/cunumpy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cunumpy
-Version: 0.1.3
+Version: 0.1.4
 Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
 Author: Max
 Project-URL: Source, https://github.com/max-models/cunumpy
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+Requires-Dist: array-api-compat
 Requires-Dist: numpy
 Provides-Extra: dev
 Requires-Dist: black[jupyter]; extra == "dev"
@@ -32,6 +33,7 @@ Requires-Dist: sphinx; extra == "docs"
 Requires-Dist: sphinx-book-theme; extra == "docs"
 Provides-Extra: test
 Requires-Dist: coverage; extra == "test"
+Requires-Dist: pyccel; extra == "test"
 Requires-Dist: pytest; extra == "test"
 
 # CuNumpy

{cunumpy-0.1.3 → cunumpy-0.1.4}/src/cunumpy.egg-info/SOURCES.txt

@@ -2,6 +2,7 @@ README.md
 pyproject.toml
 src/cunumpy/__init__.py
 src/cunumpy/__init__.pyi
+src/cunumpy/kernel.py
 src/cunumpy/main.py
 src/cunumpy/py.typed
 src/cunumpy/xp.py

{cunumpy-0.1.3 → cunumpy-0.1.4}/src/cunumpy.egg-info/requires.txt

@@ -1,3 +1,4 @@
+array-api-compat
 numpy
 
 [dev]
@@ -18,4 +19,5 @@ sphinx-book-theme
 
 [test]
 coverage
+pyccel
 pytest