cunumpy 0.1.2__tar.gz → 0.1.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {cunumpy-0.1.2 → cunumpy-0.1.4}/PKG-INFO +5 -2
- {cunumpy-0.1.2 → cunumpy-0.1.4}/README.md +2 -1
- {cunumpy-0.1.2 → cunumpy-0.1.4}/pyproject.toml +3 -2
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy/__init__.py +22 -8
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy/__init__.pyi +6 -2
- cunumpy-0.1.4/src/cunumpy/kernel.py +344 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy/xp.py +70 -32
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy.egg-info/PKG-INFO +5 -2
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy.egg-info/SOURCES.txt +1 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy.egg-info/requires.txt +2 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/setup.cfg +0 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy/main.py +0 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy/py.typed +0 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy.egg-info/dependency_links.txt +0 -0
- {cunumpy-0.1.2 → cunumpy-0.1.4}/src/cunumpy.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: cunumpy
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.4
|
|
4
4
|
Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
|
|
5
5
|
Author: Max
|
|
6
6
|
Project-URL: Source, https://github.com/max-models/cunumpy
|
|
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
|
|
|
15
15
|
Classifier: Programming Language :: Python :: 3.13
|
|
16
16
|
Requires-Python: >=3.8
|
|
17
17
|
Description-Content-Type: text/markdown
|
|
18
|
+
Requires-Dist: array-api-compat
|
|
18
19
|
Requires-Dist: numpy
|
|
19
20
|
Provides-Extra: dev
|
|
20
21
|
Requires-Dist: black[jupyter]; extra == "dev"
|
|
@@ -32,6 +33,7 @@ Requires-Dist: sphinx; extra == "docs"
|
|
|
32
33
|
Requires-Dist: sphinx-book-theme; extra == "docs"
|
|
33
34
|
Provides-Extra: test
|
|
34
35
|
Requires-Dist: coverage; extra == "test"
|
|
36
|
+
Requires-Dist: pyccel; extra == "test"
|
|
35
37
|
Requires-Dist: pytest; extra == "test"
|
|
36
38
|
|
|
37
39
|
# CuNumpy
|
|
@@ -52,7 +54,8 @@ export ARRAY_BACKEND=cupy
|
|
|
52
54
|
|
|
53
55
|
```python
|
|
54
56
|
import cunumpy as xp
|
|
55
|
-
|
|
57
|
+
|
|
58
|
+
arr = xp.array([1, 2])
|
|
56
59
|
|
|
57
60
|
print(type(arr))
|
|
58
61
|
print(xp.__version__)
|
|
@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
|
|
|
5
5
|
|
|
6
6
|
[project]
|
|
7
7
|
name = "cunumpy"
|
|
8
|
-
version = "0.1.
|
|
8
|
+
version = "0.1.4"
|
|
9
9
|
description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`."
|
|
10
10
|
readme = "README.md"
|
|
11
11
|
keywords = [ "python" ]
|
|
@@ -23,6 +23,7 @@ classifiers = [
|
|
|
23
23
|
"Programming Language :: Python :: 3.13",
|
|
24
24
|
]
|
|
25
25
|
dependencies = [
|
|
26
|
+
"array-api-compat",
|
|
26
27
|
"numpy",
|
|
27
28
|
]
|
|
28
29
|
|
|
@@ -43,7 +44,7 @@ optional-dependencies.docs = [
|
|
|
43
44
|
"sphinx",
|
|
44
45
|
"sphinx-book-theme",
|
|
45
46
|
]
|
|
46
|
-
optional-dependencies.test = [ "coverage", "pytest" ]
|
|
47
|
+
optional-dependencies.test = [ "coverage", "pyccel", "pytest" ]
|
|
47
48
|
urls."Source" = "https://github.com/max-models/cunumpy"
|
|
48
49
|
|
|
49
50
|
[tool.setuptools.packages.find]
|
|
@@ -1,10 +1,15 @@
|
|
|
1
1
|
# cunumpy/__init__.py
|
|
2
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
3
|
+
|
|
2
4
|
from . import xp
|
|
5
|
+
from .kernel import PyccelKernel
|
|
3
6
|
from .xp import (
|
|
7
|
+
cupy_available,
|
|
4
8
|
get_backend,
|
|
5
9
|
is_cpu,
|
|
6
10
|
is_gpu,
|
|
7
11
|
set_backend,
|
|
12
|
+
set_device,
|
|
8
13
|
synchronize,
|
|
9
14
|
to_cunumpy,
|
|
10
15
|
to_cupy,
|
|
@@ -12,19 +17,28 @@ from .xp import (
|
|
|
12
17
|
use_backend,
|
|
13
18
|
)
|
|
14
19
|
|
|
20
|
+
try:
|
|
21
|
+
__version__ = version("cunumpy")
|
|
22
|
+
except PackageNotFoundError:
|
|
23
|
+
__version__ = "0.0.0+unknown"
|
|
24
|
+
|
|
15
25
|
__all__ = [
|
|
16
|
-
"
|
|
17
|
-
"
|
|
18
|
-
"
|
|
19
|
-
"
|
|
26
|
+
"PyccelKernel",
|
|
27
|
+
"__version__",
|
|
28
|
+
"cupy_available",
|
|
29
|
+
"cupy_backend",
|
|
20
30
|
"get_backend",
|
|
21
|
-
"is_gpu",
|
|
22
31
|
"is_cpu",
|
|
23
|
-
"
|
|
32
|
+
"is_gpu",
|
|
33
|
+
"numpy_backend",
|
|
24
34
|
"set_backend",
|
|
35
|
+
"set_device",
|
|
25
36
|
"synchronize",
|
|
26
|
-
"
|
|
27
|
-
"
|
|
37
|
+
"to_cunumpy",
|
|
38
|
+
"to_cupy",
|
|
39
|
+
"to_numpy",
|
|
40
|
+
"use_backend",
|
|
41
|
+
"xp",
|
|
28
42
|
]
|
|
29
43
|
|
|
30
44
|
|
|
@@ -7,18 +7,22 @@ from typing import Any, Generator
|
|
|
7
7
|
import numpy as np
|
|
8
8
|
from numpy import *
|
|
9
9
|
|
|
10
|
-
from . import xp
|
|
10
|
+
from . import xp as xp
|
|
11
|
+
from .kernel import PyccelKernel as PyccelKernel
|
|
11
12
|
|
|
12
13
|
def to_numpy(array: Any) -> np.ndarray: ...
|
|
13
14
|
def to_cupy(array: Any) -> Any: ...
|
|
14
15
|
def to_cunumpy(array: Any) -> Any: ...
|
|
16
|
+
def cupy_available() -> bool: ...
|
|
15
17
|
def get_backend(array: Any) -> str: ...
|
|
16
18
|
def is_gpu(array: Any) -> bool: ...
|
|
17
19
|
def is_cpu(array: Any) -> bool: ...
|
|
18
20
|
@contextmanager
|
|
19
|
-
def use_backend(backend: str) -> Generator[None
|
|
21
|
+
def use_backend(backend: str) -> Generator[None]: ...
|
|
20
22
|
def set_backend(backend: str) -> None: ...
|
|
23
|
+
def set_device(device_id: int) -> None: ...
|
|
21
24
|
def synchronize() -> None: ...
|
|
22
25
|
|
|
23
26
|
numpy_backend: bool
|
|
24
27
|
cupy_backend: bool
|
|
28
|
+
__version__: str
|
|
@@ -0,0 +1,344 @@
|
|
|
1
|
+
"""Interface for calling Pyccel-compiled kernels with CuPy arrays.
|
|
2
|
+
|
|
3
|
+
Kernels generated by `pyccel <https://github.com/pyccel/pyccel>`_ are compiled
|
|
4
|
+
C/Fortran routines that only understand NumPy (host) arrays. :class:`PyccelKernel`
|
|
5
|
+
wraps such a kernel so that it can be called transparently with CuPy (device)
|
|
6
|
+
arrays: the arguments are copied to the host before the call, in-place updates
|
|
7
|
+
made by the kernel are copied back to the device afterwards, and any arrays
|
|
8
|
+
returned by the kernel are moved back to the device.
|
|
9
|
+
|
|
10
|
+
On the NumPy backend the wrapper is a no-op and the kernel is called directly.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import copy
|
|
16
|
+
from typing import Any, Callable, Sequence
|
|
17
|
+
|
|
18
|
+
import array_api_compat
|
|
19
|
+
import numpy as np
|
|
20
|
+
|
|
21
|
+
from .xp import _cupy_backend, to_cupy, to_numpy
|
|
22
|
+
|
|
23
|
+
__all__ = ["PyccelKernel"]
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class PyccelKernel:
|
|
27
|
+
"""Call a Pyccel-compiled kernel with NumPy or CuPy arrays.
|
|
28
|
+
|
|
29
|
+
Parameters
|
|
30
|
+
----------
|
|
31
|
+
kernel : callable
|
|
32
|
+
The pyccelized kernel (or any callable expecting NumPy arrays).
|
|
33
|
+
use_cupy : bool, optional
|
|
34
|
+
Force host/device conversion on (``True``) or off (``False``). By
|
|
35
|
+
default (``None``) it is decided at call time: conversion happens when
|
|
36
|
+
the active backend is CuPy or when a CuPy array is passed in.
|
|
37
|
+
object_modules : sequence of str, optional
|
|
38
|
+
Module prefixes (e.g. ``("struphy.", "feectools.")``) whose instances
|
|
39
|
+
should be traversed attribute-by-attribute when looking for arrays to
|
|
40
|
+
convert. Objects from other modules are passed through untouched.
|
|
41
|
+
outputs : sequence of int or str, optional
|
|
42
|
+
Which arguments the kernel writes to. Only those are copied back to the
|
|
43
|
+
device after the call, which avoids pointless device transfers for the
|
|
44
|
+
(usually much larger) read-only inputs. Positional arguments are named
|
|
45
|
+
by index, keyword arguments by name::
|
|
46
|
+
|
|
47
|
+
interpolate = PyccelKernel(some_interpolation_kernel, outputs=(5,))
|
|
48
|
+
interpolate(x, y, z, basis, coeffs, out) # `out` is argument 5
|
|
49
|
+
|
|
50
|
+
Pyccel-compiled kernels are builtins with no introspectable signature,
|
|
51
|
+
so an index and a name are *not* interchangeable: declare the form you
|
|
52
|
+
actually call with. An empty sequence declares that the kernel writes to
|
|
53
|
+
none of its arguments. By default (``None``) every converted array is
|
|
54
|
+
copied back, which is always correct but does more work.
|
|
55
|
+
|
|
56
|
+
Examples
|
|
57
|
+
--------
|
|
58
|
+
>>> from cunumpy.kernel import PyccelKernel
|
|
59
|
+
>>> kernel = PyccelKernel(my_pyccelized_function)
|
|
60
|
+
>>> kernel(out, x, y) # `out`, `x`, `y` may be NumPy or CuPy arrays
|
|
61
|
+
"""
|
|
62
|
+
|
|
63
|
+
def __init__(
|
|
64
|
+
self,
|
|
65
|
+
kernel: Callable[..., Any],
|
|
66
|
+
use_cupy: bool | None = None,
|
|
67
|
+
object_modules: Sequence[str] = (),
|
|
68
|
+
outputs: Sequence[int | str] | None = None,
|
|
69
|
+
) -> None:
|
|
70
|
+
self._kernel = kernel
|
|
71
|
+
self._use_cupy = use_cupy
|
|
72
|
+
self._object_modules = tuple(object_modules)
|
|
73
|
+
|
|
74
|
+
if outputs is None:
|
|
75
|
+
self._outputs: tuple[int | str, ...] | None = None
|
|
76
|
+
else:
|
|
77
|
+
if isinstance(outputs, (int, str)):
|
|
78
|
+
raise TypeError(
|
|
79
|
+
"outputs must be a sequence of argument indices/names, "
|
|
80
|
+
f"not a bare {type(outputs).__name__} "
|
|
81
|
+
f"(did you mean outputs=({outputs!r},)?)"
|
|
82
|
+
)
|
|
83
|
+
for entry in outputs:
|
|
84
|
+
if not isinstance(entry, (int, str)) or isinstance(entry, bool):
|
|
85
|
+
raise TypeError(
|
|
86
|
+
"outputs entries must be argument indices (int) or "
|
|
87
|
+
f"names (str), got {entry!r}"
|
|
88
|
+
)
|
|
89
|
+
self._outputs = tuple(outputs)
|
|
90
|
+
|
|
91
|
+
def __repr__(self) -> str:
|
|
92
|
+
return (
|
|
93
|
+
f"PyccelKernel(kernel={self.name!r}, use_cupy={self.use_cupy!r}, "
|
|
94
|
+
f"outputs={self._outputs!r})"
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
def _convert_to_numpy(
|
|
98
|
+
self,
|
|
99
|
+
value: Any,
|
|
100
|
+
converted: list[tuple[Any, np.ndarray]],
|
|
101
|
+
memo: dict[int, Any],
|
|
102
|
+
) -> Any:
|
|
103
|
+
"""Recursively replace CuPy arrays in `value` by host copies.
|
|
104
|
+
|
|
105
|
+
Every replacement is appended to `converted` as a
|
|
106
|
+
``(device_array, host_copy)`` pair.
|
|
107
|
+
|
|
108
|
+
`memo` maps ``id(original) -> converted`` and is shared across all
|
|
109
|
+
arguments of a single call. It serves two purposes: a device array
|
|
110
|
+
reachable by several paths is copied to the host exactly once (so the
|
|
111
|
+
kernel sees one shared array, as the caller intended, and the write-back
|
|
112
|
+
happens once), and reference cycles terminate instead of recursing
|
|
113
|
+
forever. Everything traversed here stays reachable from the caller's
|
|
114
|
+
arguments for the duration of the call, so the `id` keys cannot be
|
|
115
|
+
reused by unrelated objects.
|
|
116
|
+
"""
|
|
117
|
+
key = id(value)
|
|
118
|
+
if key in memo:
|
|
119
|
+
return memo[key]
|
|
120
|
+
|
|
121
|
+
if array_api_compat.is_cupy_array(value):
|
|
122
|
+
value_np = to_numpy(value)
|
|
123
|
+
memo[key] = value_np
|
|
124
|
+
converted.append((value, value_np))
|
|
125
|
+
return value_np
|
|
126
|
+
|
|
127
|
+
if isinstance(value, tuple):
|
|
128
|
+
# A tuple cannot be memoized before its items are converted, but it
|
|
129
|
+
# can only take part in a cycle through a mutable container, and
|
|
130
|
+
# those are memoized before they are filled in below.
|
|
131
|
+
value_np = tuple(
|
|
132
|
+
self._convert_to_numpy(item, converted, memo) for item in value
|
|
133
|
+
)
|
|
134
|
+
memo[key] = value_np
|
|
135
|
+
return value_np
|
|
136
|
+
|
|
137
|
+
if isinstance(value, list):
|
|
138
|
+
value_np = []
|
|
139
|
+
memo[key] = value_np
|
|
140
|
+
value_np.extend(
|
|
141
|
+
self._convert_to_numpy(item, converted, memo) for item in value
|
|
142
|
+
)
|
|
143
|
+
return value_np
|
|
144
|
+
|
|
145
|
+
if isinstance(value, dict):
|
|
146
|
+
value_np = {}
|
|
147
|
+
memo[key] = value_np
|
|
148
|
+
for k, v in value.items():
|
|
149
|
+
value_np[k] = self._convert_to_numpy(v, converted, memo)
|
|
150
|
+
return value_np
|
|
151
|
+
|
|
152
|
+
if hasattr(value, "__dict__") and value.__class__.__module__.startswith(
|
|
153
|
+
self._object_modules
|
|
154
|
+
):
|
|
155
|
+
# Shallow-copy the object so the caller's instance keeps pointing at
|
|
156
|
+
# its device arrays; only the copy holds the host views.
|
|
157
|
+
value_np = copy.copy(value)
|
|
158
|
+
memo[key] = value_np
|
|
159
|
+
for name, attr in vars(value).items():
|
|
160
|
+
setattr(value_np, name, self._convert_to_numpy(attr, converted, memo))
|
|
161
|
+
return value_np
|
|
162
|
+
|
|
163
|
+
return value
|
|
164
|
+
|
|
165
|
+
@staticmethod
|
|
166
|
+
def _convert_from_numpy(value: Any) -> Any:
|
|
167
|
+
"""Move NumPy arrays returned by the kernel back to the device."""
|
|
168
|
+
if isinstance(value, np.ndarray):
|
|
169
|
+
return to_cupy(value)
|
|
170
|
+
if isinstance(value, tuple):
|
|
171
|
+
return tuple(PyccelKernel._convert_from_numpy(item) for item in value)
|
|
172
|
+
if isinstance(value, list):
|
|
173
|
+
return [PyccelKernel._convert_from_numpy(item) for item in value]
|
|
174
|
+
return value
|
|
175
|
+
|
|
176
|
+
def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> None:
|
|
177
|
+
"""Record the id of every host array reachable from `value`.
|
|
178
|
+
|
|
179
|
+
Runs over the *converted* arguments, using the same traversal rules as
|
|
180
|
+
:meth:`_convert_to_numpy`, so that an output declared as a container or
|
|
181
|
+
an object contributes the arrays nested inside it.
|
|
182
|
+
"""
|
|
183
|
+
if isinstance(value, np.ndarray):
|
|
184
|
+
found.add(id(value))
|
|
185
|
+
return
|
|
186
|
+
|
|
187
|
+
if id(value) in seen:
|
|
188
|
+
return
|
|
189
|
+
|
|
190
|
+
if isinstance(value, (tuple, list)):
|
|
191
|
+
seen.add(id(value))
|
|
192
|
+
for item in value:
|
|
193
|
+
self._collect_host_arrays(item, found, seen)
|
|
194
|
+
return
|
|
195
|
+
|
|
196
|
+
if isinstance(value, dict):
|
|
197
|
+
seen.add(id(value))
|
|
198
|
+
for item in value.values():
|
|
199
|
+
self._collect_host_arrays(item, found, seen)
|
|
200
|
+
return
|
|
201
|
+
|
|
202
|
+
if hasattr(value, "__dict__") and value.__class__.__module__.startswith(
|
|
203
|
+
self._object_modules
|
|
204
|
+
):
|
|
205
|
+
seen.add(id(value))
|
|
206
|
+
for attr in vars(value).values():
|
|
207
|
+
self._collect_host_arrays(attr, found, seen)
|
|
208
|
+
|
|
209
|
+
def _output_host_arrays(
|
|
210
|
+
self, args_np: list[Any], kwargs_np: dict[str, Any]
|
|
211
|
+
) -> set[int]:
|
|
212
|
+
"""Ids of the host arrays reachable from the declared output arguments.
|
|
213
|
+
|
|
214
|
+
Raises
|
|
215
|
+
------
|
|
216
|
+
IndexError, KeyError
|
|
217
|
+
If a declared output does not correspond to an argument of this
|
|
218
|
+
call -- typically because an argument declared by index was passed
|
|
219
|
+
as a keyword, or vice versa.
|
|
220
|
+
"""
|
|
221
|
+
found: set[int] = set()
|
|
222
|
+
seen: set[int] = set()
|
|
223
|
+
|
|
224
|
+
for entry in self._outputs or ():
|
|
225
|
+
if isinstance(entry, int):
|
|
226
|
+
index = entry + len(args_np) if entry < 0 else entry
|
|
227
|
+
if not 0 <= index < len(args_np):
|
|
228
|
+
raise IndexError(
|
|
229
|
+
f"{self.name}() was declared with output argument "
|
|
230
|
+
f"{entry}, but was called with {len(args_np)} "
|
|
231
|
+
"positional argument(s). Note that an output passed as "
|
|
232
|
+
"a keyword must be declared by name, not by index."
|
|
233
|
+
)
|
|
234
|
+
self._collect_host_arrays(args_np[index], found, seen)
|
|
235
|
+
else:
|
|
236
|
+
if entry not in kwargs_np:
|
|
237
|
+
raise KeyError(
|
|
238
|
+
f"{self.name}() was declared with output argument "
|
|
239
|
+
f"{entry!r}, but no such keyword argument was passed. "
|
|
240
|
+
"Note that an output passed positionally must be "
|
|
241
|
+
"declared by index, not by name."
|
|
242
|
+
)
|
|
243
|
+
self._collect_host_arrays(kwargs_np[entry], found, seen)
|
|
244
|
+
|
|
245
|
+
return found
|
|
246
|
+
|
|
247
|
+
def _contains_cupy(self, value: Any, seen: set[int] | None = None) -> bool:
|
|
248
|
+
"""Whether `value` holds a CuPy array, following the same traversal
|
|
249
|
+
rules as :meth:`_convert_to_numpy`.
|
|
250
|
+
|
|
251
|
+
`seen` tracks already-visited containers so that reference cycles
|
|
252
|
+
terminate.
|
|
253
|
+
"""
|
|
254
|
+
if array_api_compat.is_cupy_array(value):
|
|
255
|
+
return True
|
|
256
|
+
|
|
257
|
+
if seen is None:
|
|
258
|
+
seen = set()
|
|
259
|
+
if id(value) in seen:
|
|
260
|
+
return False
|
|
261
|
+
|
|
262
|
+
if isinstance(value, (tuple, list)):
|
|
263
|
+
seen.add(id(value))
|
|
264
|
+
return any(self._contains_cupy(item, seen) for item in value)
|
|
265
|
+
|
|
266
|
+
if isinstance(value, dict):
|
|
267
|
+
seen.add(id(value))
|
|
268
|
+
return any(self._contains_cupy(item, seen) for item in value.values())
|
|
269
|
+
|
|
270
|
+
if hasattr(value, "__dict__") and value.__class__.__module__.startswith(
|
|
271
|
+
self._object_modules
|
|
272
|
+
):
|
|
273
|
+
seen.add(id(value))
|
|
274
|
+
return any(self._contains_cupy(attr, seen) for attr in vars(value).values())
|
|
275
|
+
|
|
276
|
+
return False
|
|
277
|
+
|
|
278
|
+
def _needs_conversion(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool:
|
|
279
|
+
if self._use_cupy is not None:
|
|
280
|
+
return self._use_cupy
|
|
281
|
+
if _cupy_backend():
|
|
282
|
+
return True
|
|
283
|
+
# The backend is NumPy, but individual CuPy arrays may still have been
|
|
284
|
+
# passed in explicitly.
|
|
285
|
+
return any(self._contains_cupy(value) for value in (*args, *kwargs.values()))
|
|
286
|
+
|
|
287
|
+
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
|
288
|
+
if not self._needs_conversion(args, kwargs):
|
|
289
|
+
return self._kernel(*args, **kwargs)
|
|
290
|
+
|
|
291
|
+
# Convert CuPy arrays in args/kwargs to NumPy arrays on the host. The
|
|
292
|
+
# memo is shared across args and kwargs so that an array passed more
|
|
293
|
+
# than once stays a single array on the host too.
|
|
294
|
+
converted: list[tuple[Any, np.ndarray]] = []
|
|
295
|
+
memo: dict[int, Any] = {}
|
|
296
|
+
args_np = [self._convert_to_numpy(x, converted, memo) for x in args]
|
|
297
|
+
kwargs_np = {
|
|
298
|
+
k: self._convert_to_numpy(v, converted, memo) for k, v in kwargs.items()
|
|
299
|
+
}
|
|
300
|
+
|
|
301
|
+
# Which arrays the kernel may have written to is resolved before the
|
|
302
|
+
# call, so a mis-declared output is reported even if the kernel itself
|
|
303
|
+
# would have raised first.
|
|
304
|
+
writeable = (
|
|
305
|
+
None
|
|
306
|
+
if self._outputs is None
|
|
307
|
+
else self._output_host_arrays(args_np, kwargs_np)
|
|
308
|
+
)
|
|
309
|
+
|
|
310
|
+
result = self._kernel(*args_np, **kwargs_np)
|
|
311
|
+
|
|
312
|
+
# Copy in-place kernel updates back to the device arrays.
|
|
313
|
+
for device_array, host_array in converted:
|
|
314
|
+
if writeable is None or id(host_array) in writeable:
|
|
315
|
+
device_array[...] = to_cupy(host_array)
|
|
316
|
+
|
|
317
|
+
return self._convert_from_numpy(result)
|
|
318
|
+
|
|
319
|
+
@property
|
|
320
|
+
def name(self) -> str:
|
|
321
|
+
"""Name of the wrapped kernel."""
|
|
322
|
+
return getattr(self._kernel, "__name__", type(self._kernel).__name__)
|
|
323
|
+
|
|
324
|
+
@property
|
|
325
|
+
def kernel(self) -> Callable[..., Any]:
|
|
326
|
+
"""The wrapped kernel."""
|
|
327
|
+
return self._kernel
|
|
328
|
+
|
|
329
|
+
@property
|
|
330
|
+
def use_cupy(self) -> bool:
|
|
331
|
+
"""Whether calls currently convert between device and host arrays."""
|
|
332
|
+
if self._use_cupy is not None:
|
|
333
|
+
return self._use_cupy
|
|
334
|
+
return _cupy_backend()
|
|
335
|
+
|
|
336
|
+
@property
|
|
337
|
+
def object_modules(self) -> tuple[str, ...]:
|
|
338
|
+
"""Module prefixes whose instances are traversed for arrays."""
|
|
339
|
+
return self._object_modules
|
|
340
|
+
|
|
341
|
+
@property
|
|
342
|
+
def outputs(self) -> tuple[int | str, ...] | None:
|
|
343
|
+
"""Declared output arguments, or ``None`` if every array is copied back."""
|
|
344
|
+
return self._outputs
|
|
@@ -1,23 +1,51 @@
|
|
|
1
1
|
import os
|
|
2
|
+
import warnings
|
|
2
3
|
from contextlib import contextmanager
|
|
3
4
|
from types import ModuleType
|
|
4
5
|
from typing import TYPE_CHECKING, Any, Generator, Literal
|
|
5
6
|
|
|
6
|
-
import
|
|
7
|
+
import array_api_compat
|
|
8
|
+
import array_api_compat.numpy as np
|
|
7
9
|
|
|
8
10
|
BackendType = Literal["numpy", "cupy"]
|
|
9
11
|
|
|
10
12
|
|
|
13
|
+
_CUPY_AVAILABLE_CACHE = None
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def cupy_available() -> bool:
|
|
17
|
+
"""Check if CuPy is available and functional."""
|
|
18
|
+
global _CUPY_AVAILABLE_CACHE
|
|
19
|
+
if _CUPY_AVAILABLE_CACHE is not None:
|
|
20
|
+
return _CUPY_AVAILABLE_CACHE
|
|
21
|
+
|
|
22
|
+
try:
|
|
23
|
+
import cupy as cp
|
|
24
|
+
|
|
25
|
+
# Check if a GPU is available
|
|
26
|
+
_CUPY_AVAILABLE_CACHE = cp.is_available()
|
|
27
|
+
return _CUPY_AVAILABLE_CACHE
|
|
28
|
+
except Exception: # noqa: BLE001 - tolerate any driver/runtime failure
|
|
29
|
+
_CUPY_AVAILABLE_CACHE = False
|
|
30
|
+
return False
|
|
31
|
+
|
|
32
|
+
|
|
11
33
|
class ArrayBackend:
|
|
34
|
+
"""Holds the process-wide active backend (NumPy or CuPy).
|
|
35
|
+
|
|
36
|
+
Not thread-safe: `set_backend`/`use_backend` mutate this single shared
|
|
37
|
+
instance in place, so concurrent code (threads, async tasks) switching
|
|
38
|
+
backends independently will race. Safe for the typical single-threaded
|
|
39
|
+
script/notebook usage this library targets.
|
|
40
|
+
"""
|
|
41
|
+
|
|
12
42
|
def __init__(
|
|
13
43
|
self,
|
|
14
44
|
backend: BackendType = "numpy",
|
|
15
45
|
verbose: bool = False,
|
|
16
46
|
) -> None:
|
|
17
|
-
|
|
18
|
-
"numpy"
|
|
19
|
-
"cupy",
|
|
20
|
-
], "Array backend must be either 'numpy' or 'cupy'."
|
|
47
|
+
if backend.lower() not in ("numpy", "cupy"):
|
|
48
|
+
raise ValueError("Array backend must be either 'numpy' or 'cupy'.")
|
|
21
49
|
|
|
22
50
|
self._backend: BackendType = "cupy" if backend.lower() == "cupy" else "numpy"
|
|
23
51
|
self._xp: ModuleType = np # Placeholder
|
|
@@ -27,24 +55,23 @@ class ArrayBackend:
|
|
|
27
55
|
|
|
28
56
|
def _load_backend(self, backend: BackendType, verbose: bool = False) -> ModuleType:
|
|
29
57
|
if backend == "cupy":
|
|
30
|
-
|
|
31
|
-
import cupy as cp
|
|
58
|
+
if cupy_available():
|
|
59
|
+
import array_api_compat.cupy as cp
|
|
32
60
|
|
|
61
|
+
self._backend = "cupy"
|
|
33
62
|
return cp
|
|
34
|
-
|
|
63
|
+
else:
|
|
35
64
|
if verbose:
|
|
36
|
-
print(
|
|
65
|
+
print(
|
|
66
|
+
"CuPy not available or not functional. Falling back to NumPy."
|
|
67
|
+
)
|
|
68
|
+
self._backend = "numpy"
|
|
37
69
|
return np
|
|
38
|
-
|
|
70
|
+
self._backend = "numpy"
|
|
71
|
+
return np
|
|
39
72
|
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
def __init_post__(self, verbose: bool = False) -> None:
|
|
43
|
-
# This is now redundant but kept for compatibility if called
|
|
44
|
-
self._xp = self._load_backend(self._backend, verbose)
|
|
45
|
-
assert isinstance(self._xp, ModuleType)
|
|
46
|
-
if verbose:
|
|
47
|
-
print(f"Using {self._xp.__name__} backend.")
|
|
73
|
+
def __repr__(self) -> str:
|
|
74
|
+
return f"ArrayBackend(backend={self._backend!r}, module={self._xp.__name__!r})"
|
|
48
75
|
|
|
49
76
|
@property
|
|
50
77
|
def backend(self) -> BackendType:
|
|
@@ -70,15 +97,12 @@ class ArrayBackend:
|
|
|
70
97
|
self._xp = old_xp
|
|
71
98
|
|
|
72
99
|
|
|
73
|
-
# TODO: Make this configurable via environment variable or config file.
|
|
74
100
|
array_backend = ArrayBackend(
|
|
75
101
|
backend=(
|
|
76
102
|
"cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy"
|
|
77
103
|
),
|
|
78
104
|
verbose=False,
|
|
79
105
|
)
|
|
80
|
-
# Re-run initialization logic properly after backend selection
|
|
81
|
-
array_backend.__init_post__(verbose=False)
|
|
82
106
|
|
|
83
107
|
|
|
84
108
|
def use_backend(backend: BackendType) -> Generator[None, None, None]:
|
|
@@ -102,6 +126,14 @@ def _numpy_backend() -> bool:
|
|
|
102
126
|
return array_backend.backend == "numpy"
|
|
103
127
|
|
|
104
128
|
|
|
129
|
+
def set_device(device_id: int) -> None:
|
|
130
|
+
"""Select the active CUDA device for the current process (no-op on NumPy)."""
|
|
131
|
+
if array_backend.backend == "cupy":
|
|
132
|
+
import cupy as cp
|
|
133
|
+
|
|
134
|
+
cp.cuda.Device(device_id).use()
|
|
135
|
+
|
|
136
|
+
|
|
105
137
|
def synchronize() -> None:
|
|
106
138
|
"""Wait for all kernels in all streams on current device to complete."""
|
|
107
139
|
if array_backend.backend == "cupy":
|
|
@@ -109,13 +141,20 @@ def synchronize() -> None:
|
|
|
109
141
|
import cupy as cp
|
|
110
142
|
|
|
111
143
|
cp.cuda.Device().synchronize()
|
|
112
|
-
except
|
|
144
|
+
except ImportError:
|
|
113
145
|
pass
|
|
146
|
+
except AttributeError as e:
|
|
147
|
+
warnings.warn(
|
|
148
|
+
f"CuPy synchronize() failed unexpectedly, this may indicate a "
|
|
149
|
+
f"CuPy API mismatch: {e}",
|
|
150
|
+
RuntimeWarning,
|
|
151
|
+
stacklevel=2,
|
|
152
|
+
)
|
|
114
153
|
|
|
115
154
|
|
|
116
155
|
def to_numpy(array: Any) -> np.ndarray:
|
|
117
156
|
"""Convert an array to a NumPy array."""
|
|
118
|
-
if
|
|
157
|
+
if get_backend(array) == "cupy":
|
|
119
158
|
return array.get()
|
|
120
159
|
|
|
121
160
|
return np.asarray(array)
|
|
@@ -123,25 +162,24 @@ def to_numpy(array: Any) -> np.ndarray:
|
|
|
123
162
|
|
|
124
163
|
def to_cupy(array: Any) -> Any:
|
|
125
164
|
"""Convert an array to a CuPy array."""
|
|
126
|
-
|
|
127
|
-
|
|
165
|
+
if not cupy_available():
|
|
166
|
+
raise ImportError("CuPy is not available or not functional.")
|
|
128
167
|
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
168
|
+
import array_api_compat.cupy as cp
|
|
169
|
+
|
|
170
|
+
return cp.asarray(array)
|
|
132
171
|
|
|
133
172
|
|
|
134
173
|
def to_cunumpy(array: Any) -> Any:
|
|
135
174
|
"""Convert an array to the currently active backend."""
|
|
136
|
-
if array_backend.backend == "cupy":
|
|
175
|
+
if array_backend.backend == "cupy" and cupy_available():
|
|
137
176
|
return to_cupy(array)
|
|
138
177
|
return to_numpy(array)
|
|
139
178
|
|
|
140
179
|
|
|
141
180
|
def get_backend(array: Any) -> BackendType:
|
|
142
181
|
"""Return 'cupy' or 'numpy' depending on the array type."""
|
|
143
|
-
|
|
144
|
-
return "cupy" if "cupy" in module else "numpy"
|
|
182
|
+
return "cupy" if array_api_compat.is_cupy_array(array) else "numpy"
|
|
145
183
|
|
|
146
184
|
|
|
147
185
|
def is_gpu(array: Any) -> bool:
|
|
@@ -157,7 +195,7 @@ def is_cpu(array: Any) -> bool:
|
|
|
157
195
|
# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
|
|
158
196
|
# This allows us to use autocompletion for xp (i.e., numpy/cupy) as if numpy was imported.
|
|
159
197
|
if TYPE_CHECKING:
|
|
160
|
-
import numpy as xp
|
|
198
|
+
import numpy as xp # noqa: F401 - type-checker-only alias for autocompletion
|
|
161
199
|
else:
|
|
162
200
|
# Use module-level __getattr__ for dynamic xp (Python 3.7+)
|
|
163
201
|
def __getattr__(name):
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: cunumpy
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.4
|
|
4
4
|
Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
|
|
5
5
|
Author: Max
|
|
6
6
|
Project-URL: Source, https://github.com/max-models/cunumpy
|
|
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
|
|
|
15
15
|
Classifier: Programming Language :: Python :: 3.13
|
|
16
16
|
Requires-Python: >=3.8
|
|
17
17
|
Description-Content-Type: text/markdown
|
|
18
|
+
Requires-Dist: array-api-compat
|
|
18
19
|
Requires-Dist: numpy
|
|
19
20
|
Provides-Extra: dev
|
|
20
21
|
Requires-Dist: black[jupyter]; extra == "dev"
|
|
@@ -32,6 +33,7 @@ Requires-Dist: sphinx; extra == "docs"
|
|
|
32
33
|
Requires-Dist: sphinx-book-theme; extra == "docs"
|
|
33
34
|
Provides-Extra: test
|
|
34
35
|
Requires-Dist: coverage; extra == "test"
|
|
36
|
+
Requires-Dist: pyccel; extra == "test"
|
|
35
37
|
Requires-Dist: pytest; extra == "test"
|
|
36
38
|
|
|
37
39
|
# CuNumpy
|
|
@@ -52,7 +54,8 @@ export ARRAY_BACKEND=cupy
|
|
|
52
54
|
|
|
53
55
|
```python
|
|
54
56
|
import cunumpy as xp
|
|
55
|
-
|
|
57
|
+
|
|
58
|
+
arr = xp.array([1, 2])
|
|
56
59
|
|
|
57
60
|
print(type(arr))
|
|
58
61
|
print(xp.__version__)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|