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.
Files changed (52) hide show
  1. {cunumpy-0.6.0 → cunumpy-0.6.2}/PKG-INFO +10 -4
  2. {cunumpy-0.6.0 → cunumpy-0.6.2}/README.md +7 -3
  3. {cunumpy-0.6.0 → cunumpy-0.6.2}/pyproject.toml +2 -1
  4. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/LLM_GUIDE.md +7 -2
  5. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/__init__.py +4 -0
  6. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/__init__.pyi +4 -1
  7. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_algorithms.py +46 -0
  8. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_cuda_kernel.py +18 -9
  9. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_dispatch.py +38 -0
  10. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_emulation.py +182 -14
  11. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_fake_cupy.py +26 -2
  12. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_fake_cupy_impl.py +15 -0
  13. cunumpy-0.6.2/src/cunumpy/_host.py +98 -0
  14. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_kernel.py +130 -22
  15. cunumpy-0.6.2/src/cunumpy/_metal_kernel.py +258 -0
  16. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_mpi.py +5 -1
  17. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_transfers.py +39 -7
  18. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/algorithms.py +5 -1
  19. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/array_view.cuh +148 -4
  20. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/kernel_testing.py +123 -30
  21. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/kernels.py +9 -0
  22. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/xp.py +5 -1
  23. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/PKG-INFO +10 -4
  24. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/SOURCES.txt +2 -0
  25. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/requires.txt +5 -0
  26. {cunumpy-0.6.0 → cunumpy-0.6.2}/setup.cfg +0 -0
  27. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_device.py +0 -0
  28. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_fusion.py +0 -0
  29. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_mirror.py +0 -0
  30. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_morton.py +0 -0
  31. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_philox.py +0 -0
  32. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_profiling.py +0 -0
  33. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_random_streams.py +0 -0
  34. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_scipy_backend.py +0 -0
  35. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_staging.py +0 -0
  36. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/_streams.py +0 -0
  37. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/arguments.py +0 -0
  38. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/__init__.py +0 -0
  39. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/atomic.cuh +0 -0
  40. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/index.cuh +0 -0
  41. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/morton.cuh +0 -0
  42. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/random.cuh +0 -0
  43. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/reduce.cuh +0 -0
  44. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/scan.cuh +0 -0
  45. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/memory.py +0 -0
  46. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/mpi.py +0 -0
  47. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/petsc.py +0 -0
  48. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/profiling.py +0 -0
  49. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/py.typed +0 -0
  50. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy/rng.py +0 -0
  51. {cunumpy-0.6.0 → cunumpy-0.6.2}/src/cunumpy.egg-info/dependency_links.txt +0 -0
  52. {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.0
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 `Array4D<T>`; a parameter or struct field of that type
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(MarkerArguments.__init__, "MarkerArgs")
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 `Array4D<T>`; a parameter or struct field of that type
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(MarkerArguments.__init__, "MarkerArgs")
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.0"
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>`..`Array4D<T>` params from `<cunumpy/array_view.cuh>` |
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>..Array4D<T>: data, shape[], strides[] (elements), a(i, j), size()
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 ``Array4D<T>`` views passed by value) and
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
- ``Array4D<T>`` or ``CArray1D<T>`` to ``CArray4D<T>``, see
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 Array4D<T> and CArray1D<T> to CArray4D<T> (cunumpy/array_view.cuh),
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([1234])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>")
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[1234]D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>"
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 > 4:
909
- raise ValueError(f"{what}: arrays have at most 4 dimensions, got {ndim}")
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
- ``Array4D<T>`` of those scalar types (from ``cunumpy/array_view.cuh``,
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
- ``Array4D<T>`` views, ``#include "cunumpy/index.cuh"`` the thread-index
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,