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.
Files changed (52) hide show
  1. {cunumpy-0.6.1 → cunumpy-0.6.2}/PKG-INFO +6 -2
  2. {cunumpy-0.6.1 → cunumpy-0.6.2}/README.md +3 -1
  3. {cunumpy-0.6.1 → cunumpy-0.6.2}/pyproject.toml +2 -1
  4. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/LLM_GUIDE.md +5 -0
  5. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/__init__.py +4 -0
  6. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/__init__.pyi +4 -1
  7. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_algorithms.py +46 -0
  8. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_cuda_kernel.py +9 -0
  9. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_dispatch.py +38 -0
  10. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_emulation.py +180 -12
  11. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_fake_cupy.py +26 -2
  12. {cunumpy-0.6.1 → 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.1 → 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.1 → cunumpy-0.6.2}/src/cunumpy/_mpi.py +5 -1
  17. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_transfers.py +39 -7
  18. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/algorithms.py +5 -1
  19. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/kernel_testing.py +123 -30
  20. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/kernels.py +9 -0
  21. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/xp.py +5 -1
  22. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/PKG-INFO +6 -2
  23. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/SOURCES.txt +2 -0
  24. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/requires.txt +5 -0
  25. {cunumpy-0.6.1 → cunumpy-0.6.2}/setup.cfg +0 -0
  26. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_device.py +0 -0
  27. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_fusion.py +0 -0
  28. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_mirror.py +0 -0
  29. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_morton.py +0 -0
  30. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_philox.py +0 -0
  31. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_profiling.py +0 -0
  32. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_random_streams.py +0 -0
  33. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_scipy_backend.py +0 -0
  34. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_staging.py +0 -0
  35. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/_streams.py +0 -0
  36. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/arguments.py +0 -0
  37. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/__init__.py +0 -0
  38. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/array_view.cuh +0 -0
  39. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/atomic.cuh +0 -0
  40. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/index.cuh +0 -0
  41. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/morton.cuh +0 -0
  42. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/random.cuh +0 -0
  43. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/reduce.cuh +0 -0
  44. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/cuda/include/cunumpy/scan.cuh +0 -0
  45. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/memory.py +0 -0
  46. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/mpi.py +0 -0
  47. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/petsc.py +0 -0
  48. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/profiling.py +0 -0
  49. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/py.typed +0 -0
  50. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy/rng.py +0 -0
  51. {cunumpy-0.6.1 → cunumpy-0.6.2}/src/cunumpy.egg-info/dependency_links.txt +0 -0
  52. {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.1
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.1"
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), structs and ``CudaArguments`` objects
35
- (not supported), and ``<cupy/complex.cuh>``. Use a GPU for those.
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. Arrays
279
- are updated in place with what the kernel wrote.
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 the
295
- kernel has struct parameters or complex scalars.
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