cunumpy 0.6.0__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.0 → cunumpy-0.6.2}/PKG-INFO +10 -4
- {cunumpy-0.6.0 → cunumpy-0.6.2}/README.md +7 -3
- {cunumpy-0.6.0 → cunumpy-0.6.2}/pyproject.toml +2 -1
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/LLM_GUIDE.md +7 -2
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/__init__.py +4 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/__init__.pyi +4 -1
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_algorithms.py +46 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_cuda_kernel.py +18 -9
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_dispatch.py +38 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_emulation.py +182 -14
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_fake_cupy.py +26 -2
- {cunumpy-0.6.0 → 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.0 → cunumpy-0.6.2}/src/cunumpy/_kernel.py +130 -22
- cunumpy-0.6.2/src/cunumpy/_metal_kernel.py +258 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_mpi.py +5 -1
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_transfers.py +39 -7
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/algorithms.py +5 -1
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/array_view.cuh +148 -4
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/kernel_testing.py +123 -30
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/kernels.py +9 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/xp.py +5 -1
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/PKG-INFO +10 -4
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/SOURCES.txt +2 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/requires.txt +5 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/setup.cfg +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_device.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_fusion.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_mirror.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_morton.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_philox.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_profiling.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_random_streams.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_scipy_backend.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_staging.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_streams.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/arguments.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/__init__.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/atomic.cuh +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/index.cuh +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/morton.cuh +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/random.cuh +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/reduce.cuh +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/scan.cuh +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/memory.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/mpi.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/petsc.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/profiling.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/py.typed +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/rng.py +0 -0
- {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/dependency_links.txt +0 -0
- {cunumpy-0.6.0 → 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
|
|
|
@@ -434,7 +438,7 @@ kernel(particles.args_markers, dt) # host or CUDA kernel
|
|
|
434
438
|
Kernels ported from pyccel index arrays like `markers[ip, j]`, which needs
|
|
435
439
|
shapes and strides rather than bare pointers. The shipped header
|
|
436
440
|
`cunumpy/array_view.cuh` (found by every `CudaKernel`) provides the strided
|
|
437
|
-
views `Array1D<T>` to `
|
|
441
|
+
views `Array1D<T>` to `Array16D<T>`; a parameter or struct field of that type
|
|
438
442
|
takes a CuPy array, contiguous or not, and indexes `a(i, j)`. The struct can be
|
|
439
443
|
generated from the annotations of the pyccel argument class, so the Python
|
|
440
444
|
class is the one definition, and written to a header that a test keeps in sync:
|
|
@@ -444,7 +448,9 @@ class MarkerArguments:
|
|
|
444
448
|
def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ...
|
|
445
449
|
|
|
446
450
|
|
|
447
|
-
MarkerArgs = xp.arguments.CudaStruct.from_signature(
|
|
451
|
+
MarkerArgs = xp.arguments.CudaStruct.from_signature(
|
|
452
|
+
MarkerArguments.__init__, "MarkerArgs"
|
|
453
|
+
)
|
|
448
454
|
MarkerArgs.to_header(
|
|
449
455
|
"marker_args.cuh"
|
|
450
456
|
) # Array2D<double> markers; long long n_markers; ...
|
|
@@ -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
|
|
|
@@ -394,7 +396,7 @@ kernel(particles.args_markers, dt) # host or CUDA kernel
|
|
|
394
396
|
Kernels ported from pyccel index arrays like `markers[ip, j]`, which needs
|
|
395
397
|
shapes and strides rather than bare pointers. The shipped header
|
|
396
398
|
`cunumpy/array_view.cuh` (found by every `CudaKernel`) provides the strided
|
|
397
|
-
views `Array1D<T>` to `
|
|
399
|
+
views `Array1D<T>` to `Array16D<T>`; a parameter or struct field of that type
|
|
398
400
|
takes a CuPy array, contiguous or not, and indexes `a(i, j)`. The struct can be
|
|
399
401
|
generated from the annotations of the pyccel argument class, so the Python
|
|
400
402
|
class is the one definition, and written to a header that a test keeps in sync:
|
|
@@ -404,7 +406,9 @@ class MarkerArguments:
|
|
|
404
406
|
def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ...
|
|
405
407
|
|
|
406
408
|
|
|
407
|
-
MarkerArgs = xp.arguments.CudaStruct.from_signature(
|
|
409
|
+
MarkerArgs = xp.arguments.CudaStruct.from_signature(
|
|
410
|
+
MarkerArguments.__init__, "MarkerArgs"
|
|
411
|
+
)
|
|
408
412
|
MarkerArgs.to_header(
|
|
409
413
|
"marker_args.cuh"
|
|
410
414
|
) # Array2D<double> markers; long long n_markers; ...
|
|
@@ -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` |
|
|
@@ -96,7 +101,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository.
|
|
|
96
101
|
| PETSc solve on device arrays without copies | `xp.petsc.petsc_vec(array)` (CUDA/HIP petsc4py for CuPy arrays); `xp.synchronize()` around PETSc calls |
|
|
97
102
|
| reduction inside a CUDA kernel (energy, max velocity) | `<cunumpy/reduce.cuh>`: `cunumpy_block_sum_to(out, v)`, `cunumpy_block_min/max`, `cunumpy_warp_sum` |
|
|
98
103
|
| kernel writes into a host buffer owned by another library | `xp.memory.DeviceMirror(host_array)` + `<cunumpy/atomic.cuh>` |
|
|
99
|
-
| N-D indexing in CUDA, non-contiguous arrays | `Array1D<T>`..`
|
|
104
|
+
| N-D indexing in CUDA, non-contiguous arrays | `Array1D<T>`..`Array16D<T>` params from `<cunumpy/array_view.cuh>` |
|
|
100
105
|
| one MPI rank per GPU | `bind_local_device()` → `from mpi4py import MPI` → `require_cuda_aware_mpi()` → `synchronize_for_mpi(...)` before each call |
|
|
101
106
|
| timing GPU code | `with xp.profiling.timed_region("name") as t:` → `t.elapsed` |
|
|
102
107
|
| profiler markers | `xp.profiling.nvtx_range("name")` (context manager or decorator) |
|
|
@@ -307,7 +312,7 @@ Shipped CUDA headers (always on the include path):
|
|
|
307
312
|
|
|
308
313
|
```c
|
|
309
314
|
#include <cunumpy/index.cuh> // CUNUMPY_THREAD_1D(i, n) /_2D/_3D, CUNUMPY_GRID_STRIDE_1D(i, n) {...}
|
|
310
|
-
#include <cunumpy/array_view.cuh> // Array1D<T>..
|
|
315
|
+
#include <cunumpy/array_view.cuh> // Array1D<T>..Array16D<T>: data, shape[], strides[] (elements), a(i, j), size()
|
|
311
316
|
#include <cunumpy/atomic.cuh> // cunumpy_atomic_add(double*|float*, v), _2d(data, n1, i, j, v), _3d(...)
|
|
312
317
|
```
|
|
313
318
|
|
|
@@ -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",
|
|
@@ -123,7 +126,7 @@ def cuda_include_dir() -> str:
|
|
|
123
126
|
|
|
124
127
|
:class:`CudaKernel` adds it to the include path automatically, so kernels
|
|
125
128
|
can ``#include "cunumpy/array_view.cuh"`` (strided ``Array1D<T>``
|
|
126
|
-
to ``
|
|
129
|
+
to ``Array16D<T>`` views passed by value) and
|
|
127
130
|
``#include "cunumpy/index.cuh"`` (thread-index and grid-stride macros such
|
|
128
131
|
as ``CUNUMPY_THREAD_1D(i, n)``), ``#include "cunumpy/atomic.cuh"`` (atomic
|
|
129
132
|
adds) and ``#include "cunumpy/reduce.cuh"`` (warp and block reductions).
|
|
@@ -188,7 +191,7 @@ class CudaParameter(NamedTuple):
|
|
|
188
191
|
The struct type, for a struct passed by value.
|
|
189
192
|
view_ndim : int | None
|
|
190
193
|
The number of dimensions, for an array view (``Array1D<T>`` to
|
|
191
|
-
``
|
|
194
|
+
``Array16D<T>`` or ``CArray1D<T>`` to ``CArray16D<T>``, see
|
|
192
195
|
:func:`cuda_include_dir`) passed by value.
|
|
193
196
|
contiguous : bool
|
|
194
197
|
Whether the array view is C-contiguous (``CArray2D<T>``): packed
|
|
@@ -265,11 +268,11 @@ _CTYPE_OF = {
|
|
|
265
268
|
_QUALIFIERS = {"const", "volatile", "__restrict__", "__restrict", "restrict"}
|
|
266
269
|
|
|
267
270
|
_COMPLEX = re.compile(r"(?:(?:thrust|cuda::std)::)?complex\s*<\s*(float|double)\s*>")
|
|
268
|
-
# Array1D<T> to
|
|
271
|
+
# Array1D<T> to Array16D<T> and CArray1D<T> to CArray16D<T> (cunumpy/array_view.cuh),
|
|
269
272
|
# T a scalar type of _CTYPES
|
|
270
|
-
_VIEW = re.compile(r"\b(C?)Array([
|
|
273
|
+
_VIEW = re.compile(r"\b(C?)Array([1-9]|1[0-6])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>")
|
|
271
274
|
_TOKEN = re.compile(
|
|
272
|
-
r"C?Array[
|
|
275
|
+
r"C?Array(?:[1-9]|1[0-6])D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>"
|
|
273
276
|
r"|[A-Za-z_]\w*|\*|\[\s*\]",
|
|
274
277
|
)
|
|
275
278
|
|
|
@@ -905,8 +908,8 @@ def _pyccel_ctype(annotation: Any, scalars: Mapping[str, str], what: str) -> str
|
|
|
905
908
|
ctype = scalars[scalar]
|
|
906
909
|
if ndim == 0:
|
|
907
910
|
return ctype
|
|
908
|
-
if ndim >
|
|
909
|
-
raise ValueError(f"{what}: arrays have at most
|
|
911
|
+
if ndim > 16:
|
|
912
|
+
raise ValueError(f"{what}: arrays have at most 16 dimensions, got {ndim}")
|
|
910
913
|
return f"Array{ndim}D<{ctype}>"
|
|
911
914
|
|
|
912
915
|
|
|
@@ -1075,7 +1078,7 @@ class CudaStruct:
|
|
|
1075
1078
|
``(field name, C type)`` pairs, in order, e.g. ``("x", "double*")`` or
|
|
1076
1079
|
``("n", "int")``. Scalar fields, pointers to the scalar types of
|
|
1077
1080
|
:func:`ctype_of` (or ``void*``), and array views ``Array1D<T>`` to
|
|
1078
|
-
``
|
|
1081
|
+
``Array16D<T>`` of those scalar types (from ``cunumpy/array_view.cuh``,
|
|
1079
1082
|
packed as pointer, shape and strides in elements) are supported.
|
|
1080
1083
|
|
|
1081
1084
|
Examples
|
|
@@ -1848,7 +1851,7 @@ class CudaKernel:
|
|
|
1848
1851
|
The headers shipped with cunumpy (:func:`cuda_include_dir`) are always
|
|
1849
1852
|
found at compile time (see :meth:`compile_options`):
|
|
1850
1853
|
``#include "cunumpy/array_view.cuh"`` gives the ``Array1D<T>`` to
|
|
1851
|
-
``
|
|
1854
|
+
``Array16D<T>`` views, ``#include "cunumpy/index.cuh"`` the thread-index
|
|
1852
1855
|
macros, ``#include "cunumpy/atomic.cuh"`` atomic adds,
|
|
1853
1856
|
``#include "cunumpy/reduce.cuh"`` warp and block reductions.
|
|
1854
1857
|
source_dir : str | Path | None
|
|
@@ -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,
|