cunumpy 0.6.1__tar.gz → 0.6.2__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.6.1 → cunumpy-0.6.2}/PKG-INFO +6 -2
- {cunumpy-0.6.1 → cunumpy-0.6.2}/README.md +3 -1
- {cunumpy-0.6.1 → cunumpy-0.6.2}/pyproject.toml +2 -1
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/LLM_GUIDE.md +5 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/__init__.py +4 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/__init__.pyi +4 -1
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_algorithms.py +46 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_cuda_kernel.py +9 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_dispatch.py +38 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_emulation.py +180 -12
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_fake_cupy.py +26 -2
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_fake_cupy_impl.py +15 -0
- cunumpy-0.6.2/src/cunumpy/_host.py +98 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_kernel.py +130 -22
- cunumpy-0.6.2/src/cunumpy/_metal_kernel.py +258 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_mpi.py +5 -1
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_transfers.py +39 -7
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/algorithms.py +5 -1
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/kernel_testing.py +123 -30
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/kernels.py +9 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/xp.py +5 -1
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/PKG-INFO +6 -2
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/SOURCES.txt +2 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/requires.txt +5 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/setup.cfg +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_device.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_fusion.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_mirror.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_morton.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_philox.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_profiling.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_random_streams.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_scipy_backend.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_staging.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_streams.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/arguments.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/__init__.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/array_view.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/atomic.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/index.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/morton.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/random.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/reduce.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/scan.cuh +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/memory.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/mpi.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/petsc.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/profiling.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/py.typed +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/rng.py +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/dependency_links.txt +0 -0
- {cunumpy-0.6.1 → cunumpy-0.6.2}/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.6.
|
|
3
|
+
Version: 0.6.2
|
|
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
|
|
@@ -30,6 +30,8 @@ Requires-Dist: pre-commit; extra == "docs"
|
|
|
30
30
|
Requires-Dist: pyproject-fmt; extra == "docs"
|
|
31
31
|
Requires-Dist: sphinx; extra == "docs"
|
|
32
32
|
Requires-Dist: sphinx-book-theme; extra == "docs"
|
|
33
|
+
Provides-Extra: metal
|
|
34
|
+
Requires-Dist: mlx; (sys_platform == "darwin" and platform_machine == "arm64") and extra == "metal"
|
|
33
35
|
Provides-Extra: test
|
|
34
36
|
Requires-Dist: coverage; extra == "test"
|
|
35
37
|
Requires-Dist: pytest; extra == "test"
|
|
@@ -64,7 +66,7 @@ never hide a NumPy name:
|
|
|
64
66
|
|
|
65
67
|
| Submodule | Contents |
|
|
66
68
|
|---|---|
|
|
67
|
-
| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, host implementations, `fuse` |
|
|
69
|
+
| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `MetalKernel`, host implementations, `fuse` |
|
|
68
70
|
| `xp.arguments` | CUDA only: `CudaStruct`, `CudaStructArguments`, `CudaArguments` |
|
|
69
71
|
| `xp.cuda` | CUDA only: devices, streams, debug mode, CUDA headers |
|
|
70
72
|
| `xp.rng` | `random_streams`, `get_rng`, `philox_*` |
|
|
@@ -76,6 +78,8 @@ never hide a NumPy name:
|
|
|
76
78
|
| `cunumpy.kernel_testing` | pytest helpers for host/CUDA kernel pairs |
|
|
77
79
|
|
|
78
80
|
Everything except `xp.cuda`, `xp.arguments` and `CudaKernel` works on both backends.
|
|
81
|
+
`MetalKernel` runs Metal kernels on the GPU of an Apple silicon Mac (MLX, float32, NumPy arrays;
|
|
82
|
+
`pip install 'cunumpy[metal]'`, see [the guide](docs/source/kernels/metal-kernel.md)).
|
|
79
83
|
|
|
80
84
|
## Install
|
|
81
85
|
|
|
@@ -24,7 +24,7 @@ never hide a NumPy name:
|
|
|
24
24
|
|
|
25
25
|
| Submodule | Contents |
|
|
26
26
|
|---|---|
|
|
27
|
-
| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, host implementations, `fuse` |
|
|
27
|
+
| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `MetalKernel`, host implementations, `fuse` |
|
|
28
28
|
| `xp.arguments` | CUDA only: `CudaStruct`, `CudaStructArguments`, `CudaArguments` |
|
|
29
29
|
| `xp.cuda` | CUDA only: devices, streams, debug mode, CUDA headers |
|
|
30
30
|
| `xp.rng` | `random_streams`, `get_rng`, `philox_*` |
|
|
@@ -36,6 +36,8 @@ never hide a NumPy name:
|
|
|
36
36
|
| `cunumpy.kernel_testing` | pytest helpers for host/CUDA kernel pairs |
|
|
37
37
|
|
|
38
38
|
Everything except `xp.cuda`, `xp.arguments` and `CudaKernel` works on both backends.
|
|
39
|
+
`MetalKernel` runs Metal kernels on the GPU of an Apple silicon Mac (MLX, float32, NumPy arrays;
|
|
40
|
+
`pip install 'cunumpy[metal]'`, see [the guide](docs/source/kernels/metal-kernel.md)).
|
|
39
41
|
|
|
40
42
|
## Install
|
|
41
43
|
|
|
@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
|
|
|
5
5
|
|
|
6
6
|
[project]
|
|
7
7
|
name = "cunumpy"
|
|
8
|
-
version = "0.6.
|
|
8
|
+
version = "0.6.2"
|
|
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" ]
|
|
@@ -43,6 +43,7 @@ optional-dependencies.docs = [
|
|
|
43
43
|
"sphinx",
|
|
44
44
|
"sphinx-book-theme",
|
|
45
45
|
]
|
|
46
|
+
optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ]
|
|
46
47
|
optional-dependencies.test = [ "coverage", "pytest", "scipy" ]
|
|
47
48
|
optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ]
|
|
48
49
|
urls."Source" = "https://github.com/max-models/cunumpy"
|
|
@@ -72,6 +72,11 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository.
|
|
|
72
72
|
| hand data to SciPy/matplotlib/h5py | `xp.to_numpy(a)` |
|
|
73
73
|
| call an existing NumPy-only kernel with GPU arrays (slow, correct) | `xp.kernels.PyccelKernel(fn, outputs=(...))` |
|
|
74
74
|
| launch a hand-written CUDA C kernel | `xp.kernels.CudaKernel(source, "name")` / `CudaKernel.from_file(path)` |
|
|
75
|
+
| run a Metal (MSL) kernel on an Apple silicon GPU, NumPy float32 in and out (float64 raises; `float64="cast"` computes in float32); not part of `Kernel` dispatch | `xp.kernels.MetalKernel(body, inputs=[...], outputs=[...])(*args, out=arrays, n_threads=n)`; check `xp.kernels.metal_available()` |
|
|
76
|
+
| call host-only code (SciPy, file readers) with arguments of either backend | `xp.host_call(fun, *args)`; `@xp.evaluate_on_host` on a method; `@xp.setup_on_host` on `__init__` |
|
|
77
|
+
| keep the live rows of particle arrays at the front | `n = xp.algorithms.compact_by_mask(alive, markers, weights)` |
|
|
78
|
+
| run CUDA kernel launches on the CPU in a test (fake CuPy) | `with kernel_testing.emulated_launches(): ...`; `kernel_testing.host_buffer(a)` reads a fake array; struct arguments are read through their fields |
|
|
79
|
+
| find the arguments a host kernel writes (copy only those back) | `PyccelKernel(fn, outputs=("out",))` (names work positionally); `xp.kernels.outputs_from_annotations(fn)`; `Kernel.from_folder(..., outputs="annotations")` |
|
|
75
80
|
| host kernel + CUDA port, chosen by backend | `xp.kernels.Kernel(host_fn, cuda_kernel_or_None)` |
|
|
76
81
|
| many kernels in a package, ported incrementally | `xp.kernels.KernelCatalog.from_package(__name__, missing_cuda="fallback")` |
|
|
77
82
|
| host kernels compiled at first call (your compile function), NumPy fallback | `from_package(..., host_suffix="_pyccel", compile_host=my_compile, host_fallback={...})` -> `xp.kernels.CompiledHostKernel` |
|
|
@@ -14,6 +14,7 @@ from cunumpy import (
|
|
|
14
14
|
rng,
|
|
15
15
|
xp,
|
|
16
16
|
)
|
|
17
|
+
from cunumpy._host import evaluate_on_host, host_call, setup_on_host
|
|
17
18
|
from cunumpy._scipy_backend import scipy
|
|
18
19
|
from cunumpy.xp import (
|
|
19
20
|
as_device_array,
|
|
@@ -79,9 +80,11 @@ __all__ = [
|
|
|
79
80
|
"cupy_available",
|
|
80
81
|
"cupy_backend",
|
|
81
82
|
"default_float_dtype",
|
|
83
|
+
"evaluate_on_host",
|
|
82
84
|
"get_array_backend",
|
|
83
85
|
"get_array_module",
|
|
84
86
|
"get_backend",
|
|
87
|
+
"host_call",
|
|
85
88
|
"is_cpu",
|
|
86
89
|
"is_gpu",
|
|
87
90
|
"kernels",
|
|
@@ -95,6 +98,7 @@ __all__ = [
|
|
|
95
98
|
"same_backend",
|
|
96
99
|
"scipy",
|
|
97
100
|
"set_backend",
|
|
101
|
+
"setup_on_host",
|
|
98
102
|
"synchronize",
|
|
99
103
|
"to_cunumpy",
|
|
100
104
|
"to_cupy",
|
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
# Stub file for Pylance/mypy: exposes all numpy symbols so that
|
|
2
2
|
# `import cunumpy as xp` followed by `xp.<Tab>` shows numpy completions.
|
|
3
3
|
# At runtime the real __init__.py dispatches to numpy or cupy via __getattr__.
|
|
4
|
-
from collections.abc import Generator
|
|
4
|
+
from collections.abc import Callable, Generator
|
|
5
5
|
from contextlib import contextmanager
|
|
6
6
|
from typing import Any
|
|
7
7
|
|
|
@@ -38,6 +38,9 @@ def set_backend(backend: str, *, strict: bool = ...) -> None: ...
|
|
|
38
38
|
def require_version(minimum: str) -> None: ...
|
|
39
39
|
def default_float_dtype() -> Any: ...
|
|
40
40
|
def synchronize() -> None: ...
|
|
41
|
+
def host_call(fun: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: ...
|
|
42
|
+
def evaluate_on_host(method: Callable[..., Any]) -> Callable[..., Any]: ...
|
|
43
|
+
def setup_on_host(init: Callable[..., None]) -> Callable[..., None]: ...
|
|
41
44
|
def as_device_array(
|
|
42
45
|
value: Any, dtype: Any = ..., ndim: int | None = ..., *, name: str | None = ...
|
|
43
46
|
) -> Any: ...
|
|
@@ -251,3 +251,49 @@ def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]:
|
|
|
251
251
|
)
|
|
252
252
|
order = xpm.argsort(keys, kind="stable").astype(xpm.int64, copy=False)
|
|
253
253
|
return (keys[order], order, *(array[order] for array in arrays))
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
def compact_by_mask(mask: Any, *arrays: Any) -> int:
|
|
257
|
+
"""Move the rows where `mask` is True to the front of every array, in place.
|
|
258
|
+
|
|
259
|
+
The usual step after particles left the domain or were absorbed: keep the
|
|
260
|
+
live ones at the front of the marker array (and of the arrays that go with
|
|
261
|
+
it) and continue with ``markers[:n]``. The order of the kept rows is
|
|
262
|
+
preserved, so the result is reproducible::
|
|
263
|
+
|
|
264
|
+
n = xp.algorithms.compact_by_mask(alive, markers, weights)
|
|
265
|
+
markers, weights = markers[:n], weights[:n]
|
|
266
|
+
|
|
267
|
+
Parameters
|
|
268
|
+
----------
|
|
269
|
+
mask : array of bool, shape (n,)
|
|
270
|
+
True for the rows to keep.
|
|
271
|
+
*arrays : arrays
|
|
272
|
+
Arrays with ``n`` rows (any further axes), on the backend of `mask`.
|
|
273
|
+
Rows ``[:count]`` hold the kept rows afterwards; the rows after them
|
|
274
|
+
are unspecified, so ignore them (or overwrite them).
|
|
275
|
+
|
|
276
|
+
Returns
|
|
277
|
+
-------
|
|
278
|
+
int
|
|
279
|
+
The number of kept rows. Its value is needed on the host, so on CuPy
|
|
280
|
+
the call synchronizes once per call (counted by
|
|
281
|
+
:func:`~cunumpy.profiling.count_transfers` where it can be seen).
|
|
282
|
+
"""
|
|
283
|
+
assert_same_backend(mask, *arrays)
|
|
284
|
+
xpm = get_array_module(mask)
|
|
285
|
+
mask = xpm.asarray(mask)
|
|
286
|
+
if mask.ndim != 1 or mask.dtype != np.bool_:
|
|
287
|
+
raise TypeError(
|
|
288
|
+
f"mask must be a 1D boolean array, got dtype {mask.dtype}, {mask.ndim}D",
|
|
289
|
+
)
|
|
290
|
+
for i, array in enumerate(arrays):
|
|
291
|
+
if array.ndim < 1 or array.shape[0] != mask.shape[0]:
|
|
292
|
+
raise ValueError(
|
|
293
|
+
f"array {i} has shape {array.shape}, expected {mask.shape[0]} rows",
|
|
294
|
+
)
|
|
295
|
+
rows = xpm.nonzero(mask)[0]
|
|
296
|
+
n_kept = int(rows.size)
|
|
297
|
+
for array in arrays:
|
|
298
|
+
array[:n_kept] = array[rows] # the right side is a copy: no overlap problem
|
|
299
|
+
return n_kept
|
|
@@ -66,6 +66,9 @@ from typing import Any, NamedTuple
|
|
|
66
66
|
|
|
67
67
|
import numpy as np
|
|
68
68
|
|
|
69
|
+
from cunumpy._transfers import _ACTIVE as _COUNTERS
|
|
70
|
+
from cunumpy._transfers import _record_sync
|
|
71
|
+
|
|
69
72
|
__all__ = [
|
|
70
73
|
"DEBUG_OPTIONS",
|
|
71
74
|
"CudaArguments",
|
|
@@ -2421,6 +2424,10 @@ class CudaKernel:
|
|
|
2421
2424
|
if shared_mem < 0:
|
|
2422
2425
|
raise ValueError(f"shared_mem must be non-negative, got {shared_mem}")
|
|
2423
2426
|
values = self.prepare_args(*args)
|
|
2427
|
+
# RawKernel accepts size-one NumPy arrays for structs passed by value,
|
|
2428
|
+
# but not structured NumPy scalars (np.void). Keep their packed bytes
|
|
2429
|
+
# and alignment intact, including array-view pointers and strides.
|
|
2430
|
+
values = tuple(np.asarray(v) if isinstance(v, np.void) else v for v in values)
|
|
2424
2431
|
if 0 in grid_shape:
|
|
2425
2432
|
return
|
|
2426
2433
|
|
|
@@ -2528,6 +2535,8 @@ class CudaKernel:
|
|
|
2528
2535
|
stream = cp.cuda.get_current_stream()
|
|
2529
2536
|
if _is_capturing(stream):
|
|
2530
2537
|
return
|
|
2538
|
+
if _COUNTERS:
|
|
2539
|
+
_record_sync(f"debug synchronization after kernel {self.expression!r}")
|
|
2531
2540
|
try:
|
|
2532
2541
|
stream.synchronize()
|
|
2533
2542
|
except Exception as error:
|
|
@@ -40,6 +40,7 @@ from cunumpy._kernel import (
|
|
|
40
40
|
HostImplementations,
|
|
41
41
|
PyccelKernel,
|
|
42
42
|
get_device_kernel_implementation,
|
|
43
|
+
outputs_from_annotations,
|
|
43
44
|
)
|
|
44
45
|
from cunumpy._transfers import _ACTIVE as _COUNTERS
|
|
45
46
|
from cunumpy._transfers import _record
|
|
@@ -219,6 +220,10 @@ class Kernel:
|
|
|
219
220
|
"host_options are for wrapping a plain callable; configure the "
|
|
220
221
|
"given PyccelKernel directly",
|
|
221
222
|
)
|
|
223
|
+
if host_kernel._parameters is None:
|
|
224
|
+
# compiled kernels have no Python signature: name the parameters
|
|
225
|
+
# from the host function, so that outputs may be given by name
|
|
226
|
+
host_kernel._parameters = self.host_parameters
|
|
222
227
|
self._host_kernel = host_kernel
|
|
223
228
|
self._cuda_kernel = cuda_kernel
|
|
224
229
|
self._name = name if name is not None else host_kernel.name
|
|
@@ -239,6 +244,7 @@ class Kernel:
|
|
|
239
244
|
check_name_length: bool = True,
|
|
240
245
|
missing_cuda: str = "raise",
|
|
241
246
|
host_options: Mapping[str, Any] | None = None,
|
|
247
|
+
outputs: Sequence[int | str] | str | None = None,
|
|
242
248
|
include_dirs: Sequence[str | Path] | None = None,
|
|
243
249
|
dispatch: str = "backend",
|
|
244
250
|
compile_host: Callable[[Any], Any] | None = None,
|
|
@@ -285,6 +291,14 @@ class Kernel:
|
|
|
285
291
|
As for :meth:`KernelCatalog.from_package`.
|
|
286
292
|
missing_cuda, host_options, dispatch
|
|
287
293
|
Passed on to :class:`Kernel`.
|
|
294
|
+
outputs : Sequence[int | str] | "annotations" | None
|
|
295
|
+
The arguments the host kernel writes to (names or indices, see
|
|
296
|
+
:class:`~cunumpy.kernels.PyccelKernel`), so that only those arrays are
|
|
297
|
+
copied back to the device on the fallback path. ``"annotations"``
|
|
298
|
+
reads them from the annotations of the host function: every
|
|
299
|
+
parameter that is not ``Final`` or a scalar is an output (see
|
|
300
|
+
:func:`~cunumpy.kernels.outputs_from_annotations`). A value in
|
|
301
|
+
`host_options` takes precedence.
|
|
288
302
|
include_dirs : Sequence[str | Path] | None
|
|
289
303
|
Include directories of the CUDA kernel, in addition to the folder
|
|
290
304
|
itself; by default the source root of the top-level package (the
|
|
@@ -334,6 +348,19 @@ class Kernel:
|
|
|
334
348
|
test_args = f"{package}.{name}{test_args_suffix}"
|
|
335
349
|
module = importlib.import_module(f"{package}.{name}{host_suffix}")
|
|
336
350
|
python = getattr(module, name)
|
|
351
|
+
if outputs is not None and "outputs" not in (host_options or {}):
|
|
352
|
+
declared = (
|
|
353
|
+
outputs_from_annotations(python)
|
|
354
|
+
if isinstance(outputs, str) and outputs == "annotations"
|
|
355
|
+
else outputs
|
|
356
|
+
)
|
|
357
|
+
if isinstance(outputs, str) and outputs != "annotations":
|
|
358
|
+
raise ValueError(
|
|
359
|
+
f"outputs must be a sequence of names/indices or 'annotations', "
|
|
360
|
+
f"got {outputs!r}",
|
|
361
|
+
)
|
|
362
|
+
if declared is not None:
|
|
363
|
+
host_options = {**(host_options or {}), "outputs": declared}
|
|
337
364
|
loaders: dict[str, Callable[[], Callable[..., Any]]] = {
|
|
338
365
|
"python": lambda: python,
|
|
339
366
|
"pyccel": (
|
|
@@ -669,6 +696,12 @@ class KernelCatalog(Mapping):
|
|
|
669
696
|
host_options: (
|
|
670
697
|
Mapping[str, Any] | Callable[[str], Mapping[str, Any]] | None
|
|
671
698
|
) = None,
|
|
699
|
+
outputs: (
|
|
700
|
+
Sequence[int | str]
|
|
701
|
+
| str
|
|
702
|
+
| Callable[[str], Sequence[int | str] | str | None]
|
|
703
|
+
| None
|
|
704
|
+
) = None,
|
|
672
705
|
include_dirs: Sequence[str | Path] | None = None,
|
|
673
706
|
dispatch: str = "backend",
|
|
674
707
|
compile_host: Callable[[Any], Any] | None = None,
|
|
@@ -713,6 +746,10 @@ class KernelCatalog(Mapping):
|
|
|
713
746
|
Keyword arguments for the :class:`~cunumpy.kernels.PyccelKernel` wrapping each
|
|
714
747
|
host kernel (see :class:`Kernel`): the same for all kernels, or a
|
|
715
748
|
function of the kernel name, e.g. to declare per-kernel ``outputs``.
|
|
749
|
+
outputs : Sequence[int | str] | "annotations" | Callable | None
|
|
750
|
+
The arguments every host kernel writes to, or ``"annotations"`` to
|
|
751
|
+
read them from the annotations of each host function, or a function
|
|
752
|
+
of the kernel name returning either (see :meth:`Kernel.from_folder`).
|
|
716
753
|
include_dirs : Sequence[str | Path] | None
|
|
717
754
|
Include directories of the CUDA kernels, in addition to each
|
|
718
755
|
kernel's own folder. By default the source root of the top-level
|
|
@@ -757,6 +794,7 @@ class KernelCatalog(Mapping):
|
|
|
757
794
|
host_options=(
|
|
758
795
|
host_options(name) if callable(host_options) else host_options
|
|
759
796
|
),
|
|
797
|
+
outputs=outputs(name) if callable(outputs) else outputs,
|
|
760
798
|
include_dirs=include_dirs,
|
|
761
799
|
dispatch=dispatch,
|
|
762
800
|
compile_host=compile_host,
|
|
@@ -28,11 +28,22 @@ thread runs to its next barrier before any thread continues past it, as on a
|
|
|
28
28
|
GPU. Per-block deposits, shared-memory reductions and tiled kernels therefore
|
|
29
29
|
work.
|
|
30
30
|
|
|
31
|
+
Struct parameters take, for each struct, a mapping from field name to value, a
|
|
32
|
+
:class:`~cunumpy.arguments.CudaStructValue`, or any object with an attribute
|
|
33
|
+
per field (e.g. a :class:`~cunumpy.arguments.CudaStructArguments` or a host
|
|
34
|
+
argument class); the arrays of the fields are updated in place like the others.
|
|
35
|
+
|
|
36
|
+
:func:`emulated_launches` makes every :class:`~cunumpy.kernels.CudaKernel` launch
|
|
37
|
+
in a block run through the emulation, on the arrays of the fake CuPy
|
|
38
|
+
(:mod:`cunumpy._fake_cupy`), so that code that launches kernels (a
|
|
39
|
+
:class:`~cunumpy.kernels.Kernel` on the CuPy backend, a propagator) can be
|
|
40
|
+
tested without a GPU.
|
|
41
|
+
|
|
31
42
|
What it does not emulate: concurrency between the barriers (threads run one
|
|
32
43
|
after another, so atomics are plain additions and races never show), warp
|
|
33
44
|
intrinsics (``__shfl_*``, ``__syncwarp``, ``__ballot_sync``, ...; a kernel or an
|
|
34
|
-
included header using them is refused),
|
|
35
|
-
|
|
45
|
+
included header using them is refused), and ``<cupy/complex.cuh>``. Use a GPU
|
|
46
|
+
for those.
|
|
36
47
|
|
|
37
48
|
Floating point: like NVRTC (``--fmad=true`` by default) the C++ compiler may
|
|
38
49
|
fuse ``a * b + c`` into one fused multiply-add, so results can differ from
|
|
@@ -49,21 +60,25 @@ import re
|
|
|
49
60
|
import shutil
|
|
50
61
|
import subprocess
|
|
51
62
|
import tempfile
|
|
52
|
-
from collections.abc import Sequence
|
|
63
|
+
from collections.abc import Generator, Mapping, Sequence
|
|
64
|
+
from contextlib import contextmanager
|
|
53
65
|
from pathlib import Path
|
|
54
66
|
from typing import Any
|
|
55
67
|
|
|
56
68
|
import numpy as np
|
|
57
69
|
|
|
70
|
+
from cunumpy import _fake_cupy
|
|
58
71
|
from cunumpy._cuda_kernel import (
|
|
59
72
|
CudaKernel,
|
|
60
73
|
CudaParameter,
|
|
74
|
+
CudaStructArguments,
|
|
75
|
+
CudaStructValue,
|
|
61
76
|
_scalar_checker,
|
|
62
77
|
_strip_comments,
|
|
63
78
|
cuda_include_dir,
|
|
64
79
|
)
|
|
65
80
|
|
|
66
|
-
__all__ = ["emulate_cuda_kernel", "emulation_compiler"]
|
|
81
|
+
__all__ = ["emulate_cuda_kernel", "emulated_launches", "emulation_compiler"]
|
|
67
82
|
|
|
68
83
|
# CUDA constructs that serial emulation would get wrong
|
|
69
84
|
_UNSUPPORTED = {
|
|
@@ -258,6 +273,76 @@ def _check_supported(kernel: CudaKernel) -> None:
|
|
|
258
273
|
)
|
|
259
274
|
|
|
260
275
|
|
|
276
|
+
def _host(value: Any) -> Any:
|
|
277
|
+
"""The NumPy buffer of a fake CuPy array; anything else as it is."""
|
|
278
|
+
if _fake_cupy.is_active():
|
|
279
|
+
try:
|
|
280
|
+
return _fake_cupy.host_buffer(value)
|
|
281
|
+
except TypeError:
|
|
282
|
+
pass
|
|
283
|
+
return value
|
|
284
|
+
|
|
285
|
+
|
|
286
|
+
def _field_value(struct_name: str, value: Any, field: CudaParameter) -> Any:
|
|
287
|
+
"""The value of the struct field `field` in the struct argument `value`."""
|
|
288
|
+
try:
|
|
289
|
+
if isinstance(value, (CudaStructValue, Mapping)):
|
|
290
|
+
return value[field.name]
|
|
291
|
+
return getattr(value, field.name)
|
|
292
|
+
except (KeyError, AttributeError):
|
|
293
|
+
raise TypeError(
|
|
294
|
+
f"struct argument {struct_name!r} has no value for the field "
|
|
295
|
+
f"{field.name!r} ({type(value).__name__})",
|
|
296
|
+
) from None
|
|
297
|
+
|
|
298
|
+
|
|
299
|
+
def _declaration(param: CudaParameter, name: str) -> str:
|
|
300
|
+
if param.view_ndim is not None:
|
|
301
|
+
return f"{param.ctype} {name}"
|
|
302
|
+
return f"{param.ctype}{'*' if param.pointer else ''} {name}"
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
def _flatten_structs(
|
|
306
|
+
kernel: CudaKernel,
|
|
307
|
+
args: Sequence[Any],
|
|
308
|
+
) -> tuple[CudaKernel, list[Any]]:
|
|
309
|
+
"""A kernel that takes every struct field as a parameter, and its arguments.
|
|
310
|
+
|
|
311
|
+
The wrapper kernel (the source of `kernel` plus a ``__global__`` function with
|
|
312
|
+
one parameter per field) rebuilds the structs and calls `kernel`, so that the
|
|
313
|
+
emulation only has to deal with arrays and scalars.
|
|
314
|
+
"""
|
|
315
|
+
params, builds, call, flat = [], [], [], []
|
|
316
|
+
for p, value in zip(kernel.signature, args):
|
|
317
|
+
if p.struct is None:
|
|
318
|
+
params.append(_declaration(p, p.name))
|
|
319
|
+
call.append(p.name)
|
|
320
|
+
flat.append(value)
|
|
321
|
+
continue
|
|
322
|
+
names = [f"{p.name}__{f.name}" for f in p.struct.fields]
|
|
323
|
+
params += [_declaration(f, n) for f, n in zip(p.struct.fields, names)]
|
|
324
|
+
builds.append(f" {p.struct.name} {p.name}_struct{{{', '.join(names)}}};")
|
|
325
|
+
call.append(f"{p.name}_struct")
|
|
326
|
+
flat += [_field_value(p.name, value, f) for f in p.struct.fields]
|
|
327
|
+
name = f"emulated_{kernel.name}"
|
|
328
|
+
source = (
|
|
329
|
+
kernel.source
|
|
330
|
+
+ f'\nextern "C" __global__ void {name}({", ".join(params)}) {{\n'
|
|
331
|
+
+ "\n".join(builds)
|
|
332
|
+
+ f"\n {kernel.expression}({', '.join(call)});\n}}\n"
|
|
333
|
+
)
|
|
334
|
+
wrapper = CudaKernel(
|
|
335
|
+
source,
|
|
336
|
+
name,
|
|
337
|
+
block_size=kernel.block_size,
|
|
338
|
+
options=[o for o in kernel.options if not o.startswith("-I")],
|
|
339
|
+
include_dirs=kernel.include_dirs,
|
|
340
|
+
source_dir=kernel.source_dir,
|
|
341
|
+
n_threads_from=kernel.n_threads_from,
|
|
342
|
+
)
|
|
343
|
+
return wrapper, flat
|
|
344
|
+
|
|
345
|
+
|
|
261
346
|
def emulate_cuda_kernel(
|
|
262
347
|
kernel: CudaKernel,
|
|
263
348
|
*args: Any,
|
|
@@ -275,8 +360,11 @@ def emulate_cuda_kernel(
|
|
|
275
360
|
kernel : CudaKernel
|
|
276
361
|
The kernel; its signature must be parsed (the default).
|
|
277
362
|
*args
|
|
278
|
-
The kernel arguments with NumPy arrays in place of CuPy arrays
|
|
279
|
-
are
|
|
363
|
+
The kernel arguments with NumPy arrays in place of CuPy arrays (arrays of
|
|
364
|
+
the fake CuPy are used through their host buffer). Arrays are updated in
|
|
365
|
+
place with what the kernel wrote. A struct parameter takes a mapping of
|
|
366
|
+
field names to values, a :class:`~cunumpy.arguments.CudaStructValue`, or
|
|
367
|
+
an object with an attribute per field.
|
|
280
368
|
n_threads, grid, block
|
|
281
369
|
Launch shape, as for :meth:`CudaKernel.__call__`.
|
|
282
370
|
compiler : str | None
|
|
@@ -291,8 +379,8 @@ def emulate_cuda_kernel(
|
|
|
291
379
|
Raises
|
|
292
380
|
------
|
|
293
381
|
NotImplementedError
|
|
294
|
-
If the kernel (or a header it includes) uses warp intrinsics, or
|
|
295
|
-
|
|
382
|
+
If the kernel (or a header it includes) uses warp intrinsics, or has
|
|
383
|
+
complex scalars.
|
|
296
384
|
TypeError
|
|
297
385
|
If an argument does not match its parameter (dtype, dimensions, a
|
|
298
386
|
scalar that does not fit).
|
|
@@ -309,6 +397,19 @@ def emulate_cuda_kernel(
|
|
|
309
397
|
raise TypeError(
|
|
310
398
|
f"kernel {kernel.name!r} takes {len(params)} arguments, got {len(args)}",
|
|
311
399
|
)
|
|
400
|
+
if any(p.struct is not None for p in params):
|
|
401
|
+
wrapper, flat = _flatten_structs(kernel, args)
|
|
402
|
+
return emulate_cuda_kernel(
|
|
403
|
+
wrapper,
|
|
404
|
+
*flat,
|
|
405
|
+
n_threads=n_threads,
|
|
406
|
+
grid=grid,
|
|
407
|
+
block=block,
|
|
408
|
+
compiler=compiler,
|
|
409
|
+
options=options,
|
|
410
|
+
shared_mem=shared_mem,
|
|
411
|
+
)
|
|
412
|
+
args = tuple(_host(a) for a in args)
|
|
312
413
|
compiler = compiler or emulation_compiler()
|
|
313
414
|
if compiler is None:
|
|
314
415
|
raise RuntimeError("emulation needs a C++ compiler (set CXX or install c++)")
|
|
@@ -323,10 +424,6 @@ def emulate_cuda_kernel(
|
|
|
323
424
|
globals_, inits, call_args, writes, arrays = [], [], [], [], []
|
|
324
425
|
for i, (param, value) in enumerate(zip(params, args)):
|
|
325
426
|
name = f"cunumpy_arg{i}"
|
|
326
|
-
if param.struct is not None:
|
|
327
|
-
raise NotImplementedError(
|
|
328
|
-
"emulation does not support struct parameters",
|
|
329
|
-
)
|
|
330
427
|
if param.pointer or param.view_ndim is not None:
|
|
331
428
|
if not isinstance(value, np.ndarray):
|
|
332
429
|
raise TypeError(
|
|
@@ -444,3 +541,74 @@ def emulate_cuda_kernel(
|
|
|
444
541
|
for value, buffer, out in arrays:
|
|
445
542
|
result = np.fromfile(out, dtype=buffer.dtype).reshape(buffer.shape)
|
|
446
543
|
value[...] = result
|
|
544
|
+
|
|
545
|
+
|
|
546
|
+
@contextmanager
|
|
547
|
+
def emulated_launches(
|
|
548
|
+
*,
|
|
549
|
+
compiler: str | None = None,
|
|
550
|
+
options: Sequence[str] = (),
|
|
551
|
+
) -> Generator[None, None, None]:
|
|
552
|
+
"""Run every :class:`~cunumpy.kernels.CudaKernel` launch in the block on the CPU.
|
|
553
|
+
|
|
554
|
+
With the fake CuPy (:func:`cunumpy.kernel_testing.install_fake_cupy`)
|
|
555
|
+
CUDA kernels cannot run. Inside this block a launch is emulated instead
|
|
556
|
+
(:func:`emulate_cuda_kernel`), on the host buffers of the fake CuPy arrays
|
|
557
|
+
it is given, so code that launches kernels runs on a machine without a GPU
|
|
558
|
+
and its device arrays hold the results afterwards::
|
|
559
|
+
|
|
560
|
+
with emulated_launches():
|
|
561
|
+
propagator(dt) # CuPy backend: the CUDA kernels run on the CPU
|
|
562
|
+
|
|
563
|
+
The launch arguments are those of :meth:`CudaKernel.__call__`
|
|
564
|
+
(`stream` is ignored). Argument objects are flattened like in a launch;
|
|
565
|
+
struct values and :class:`~cunumpy.arguments.CudaStructArguments` objects
|
|
566
|
+
stay one argument and are read through their fields. Arrays must be arrays
|
|
567
|
+
of the fake CuPy or NumPy arrays.
|
|
568
|
+
|
|
569
|
+
Launches are serial and slow (each one compiles the kernel as C++), so use it
|
|
570
|
+
on small problems. The limits of :func:`emulate_cuda_kernel` apply.
|
|
571
|
+
|
|
572
|
+
Parameters
|
|
573
|
+
----------
|
|
574
|
+
compiler : str | None
|
|
575
|
+
C++ compiler; by default :func:`emulation_compiler`.
|
|
576
|
+
options : Sequence[str]
|
|
577
|
+
Additional compiler options for every launch, e.g.
|
|
578
|
+
``("-ffp-contract=off",)``.
|
|
579
|
+
"""
|
|
580
|
+
original = CudaKernel.__call__
|
|
581
|
+
|
|
582
|
+
def launch(
|
|
583
|
+
self: CudaKernel,
|
|
584
|
+
*args: Any,
|
|
585
|
+
n_threads: int | Sequence[int] | None = None,
|
|
586
|
+
grid: int | Sequence[int] | None = None,
|
|
587
|
+
block: int | Sequence[int] | None = None,
|
|
588
|
+
shared_mem: int = 0,
|
|
589
|
+
stream: Any = None,
|
|
590
|
+
) -> None:
|
|
591
|
+
flat: list[Any] = []
|
|
592
|
+
for arg in args:
|
|
593
|
+
if isinstance(arg, (CudaStructValue, CudaStructArguments)):
|
|
594
|
+
flat.append(arg)
|
|
595
|
+
elif hasattr(arg, "__cuda_args__"):
|
|
596
|
+
flat.extend(arg.__cuda_args__())
|
|
597
|
+
else:
|
|
598
|
+
flat.append(arg)
|
|
599
|
+
emulate_cuda_kernel(
|
|
600
|
+
self,
|
|
601
|
+
*flat,
|
|
602
|
+
n_threads=n_threads,
|
|
603
|
+
grid=grid,
|
|
604
|
+
block=block,
|
|
605
|
+
compiler=compiler,
|
|
606
|
+
options=options,
|
|
607
|
+
shared_mem=shared_mem,
|
|
608
|
+
)
|
|
609
|
+
|
|
610
|
+
CudaKernel.__call__ = launch # type: ignore[method-assign]
|
|
611
|
+
try:
|
|
612
|
+
yield
|
|
613
|
+
finally:
|
|
614
|
+
CudaKernel.__call__ = original # type: ignore[method-assign]
|
|
@@ -17,7 +17,9 @@ matters for finding host/device bugs:
|
|
|
17
17
|
:class:`~cunumpy.kernels.CudaKernel` work;
|
|
18
18
|
* CUDA kernels cannot run: ``RawKernel`` and friends raise
|
|
19
19
|
``NotImplementedError`` when called, and :func:`cunumpy.kernel_testing.requires_cupy`
|
|
20
|
-
skips tests while the fake is active.
|
|
20
|
+
skips tests while the fake is active. Inside
|
|
21
|
+
:func:`cunumpy.kernel_testing.emulated_launches` a ``CudaKernel`` launch runs
|
|
22
|
+
on the CPU instead, on the host buffers of the arrays (:func:`host_buffer`).
|
|
21
23
|
|
|
22
24
|
Activate it before CuPy or cunumpy's backend is first used, either with the
|
|
23
25
|
environment variable ``CUNUMPY_FAKE_CUPY=1`` (read when cunumpy is imported)
|
|
@@ -35,7 +37,7 @@ import sys
|
|
|
35
37
|
import types
|
|
36
38
|
from pathlib import Path
|
|
37
39
|
|
|
38
|
-
__all__ = ["install", "is_active", "uninstall"]
|
|
40
|
+
__all__ = ["host_buffer", "install", "is_active", "uninstall"]
|
|
39
41
|
|
|
40
42
|
_IMPLEMENTATION = Path(__file__).with_name("_fake_cupy_impl.py")
|
|
41
43
|
_SUBMODULES = ("cupy.cuda", "cupy.cuda.device", "cupy.cuda.runtime", "cupy.linalg")
|
|
@@ -47,6 +49,28 @@ def is_active() -> bool:
|
|
|
47
49
|
return bool(getattr(module, "__cunumpy_fake__", False))
|
|
48
50
|
|
|
49
51
|
|
|
52
|
+
def host_buffer(array: object) -> object:
|
|
53
|
+
"""The NumPy array that holds the data of a fake ``cupy.ndarray``.
|
|
54
|
+
|
|
55
|
+
Not a copy: writing to the buffer changes the fake device array, so a
|
|
56
|
+
kernel emulation (:func:`cunumpy.kernel_testing.emulated_launches`) can
|
|
57
|
+
update the array in place. A test can also read results with it without a
|
|
58
|
+
``.get()``, which is the point of a "device" array being visible in tests.
|
|
59
|
+
|
|
60
|
+
Raises
|
|
61
|
+
------
|
|
62
|
+
TypeError
|
|
63
|
+
If `array` is not an array of the fake CuPy (including when the fake
|
|
64
|
+
is not active).
|
|
65
|
+
"""
|
|
66
|
+
module = sys.modules.get("cupy")
|
|
67
|
+
if is_active() and isinstance(array, module.ndarray):
|
|
68
|
+
return array._a
|
|
69
|
+
raise TypeError(
|
|
70
|
+
f"host_buffer needs an array of the fake CuPy, got {type(array).__name__}",
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
|
|
50
74
|
def install() -> types.ModuleType:
|
|
51
75
|
"""Install the fake ``cupy`` package into ``sys.modules`` and return it.
|
|
52
76
|
|
|
@@ -15,6 +15,14 @@ __cunumpy_fake__ = True
|
|
|
15
15
|
_HOST = _np.ndarray
|
|
16
16
|
|
|
17
17
|
|
|
18
|
+
def _sync(what):
|
|
19
|
+
"""Count a scalar read of a device array (an implicit sync of the real CuPy)."""
|
|
20
|
+
from cunumpy._transfers import _ACTIVE, _record_sync
|
|
21
|
+
|
|
22
|
+
if _ACTIVE:
|
|
23
|
+
_record_sync(what)
|
|
24
|
+
|
|
25
|
+
|
|
18
26
|
def _err(obj, where=""):
|
|
19
27
|
return TypeError(
|
|
20
28
|
f"Unsupported type {type(obj)}{where} (fake CuPy: host arrays/lists are "
|
|
@@ -75,6 +83,8 @@ class ndarray:
|
|
|
75
83
|
# no __array_interface__ etc.: NumPy must not see the host buffer
|
|
76
84
|
if name.startswith("__"):
|
|
77
85
|
raise AttributeError(name)
|
|
86
|
+
if name in ("item", "tolist"):
|
|
87
|
+
_sync(f"ndarray.{name}()")
|
|
78
88
|
attr = getattr(self._a, name)
|
|
79
89
|
if callable(attr):
|
|
80
90
|
return _wrap_callable(attr, strict=True, name=f"ndarray.{name}")
|
|
@@ -101,18 +111,23 @@ class ndarray:
|
|
|
101
111
|
return format(self._a.item() if self._a.ndim == 0 else self._a, spec)
|
|
102
112
|
|
|
103
113
|
def __bool__(self):
|
|
114
|
+
_sync("bool(device array)")
|
|
104
115
|
return bool(self._a)
|
|
105
116
|
|
|
106
117
|
def __int__(self):
|
|
118
|
+
_sync("int(device array)")
|
|
107
119
|
return int(self._a)
|
|
108
120
|
|
|
109
121
|
def __float__(self):
|
|
122
|
+
_sync("float(device array)")
|
|
110
123
|
return float(self._a)
|
|
111
124
|
|
|
112
125
|
def __complex__(self):
|
|
126
|
+
_sync("complex(device array)")
|
|
113
127
|
return complex(self._a)
|
|
114
128
|
|
|
115
129
|
def __index__(self):
|
|
130
|
+
_sync("index of a device array")
|
|
116
131
|
return operator.index(self._a.item() if self._a.ndim == 0 else self._a)
|
|
117
132
|
|
|
118
133
|
__hash__ = None
|