diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml
index 2cdc674..02b58a7 100644
--- a/.github/workflows/static_analysis.yml
+++ b/.github/workflows/static_analysis.yml
@@ -74,6 +74,17 @@ jobs:
pip install ruff
ruff check src/ tests/
+ numpydoc:
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout the code
+ uses: actions/checkout@v4
+
+ - name: Docstring sections with numpydoc
+ run: |
+ pip install . numpydoc
+ python -m numpydoc lint $(git ls-files 'src/cunumpy/*.py' 'src/cunumpy/**/*.py' | grep -v _fake_cupy)
+
pylint:
runs-on: ubuntu-latest
continue-on-error: true
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index 614da82..ab6fa28 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -72,11 +72,15 @@ jobs:
run: |
pytest .
+ - name: Run docstring examples
+ run: |
+ pytest --doctest-modules src/cunumpy
+
- name: Test docs build
run: |
pip install ".[docs]"
cd docs
make clean
- make html
+ make html SPHINXOPTS="-W --keep-going"
cd ..
ls docs/build/html/index.html
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7499713..298b86e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,10 +12,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
copy-back, run in CPU CI with fake CuPy and required compilation.
- Publishing now requires CPU and GPU validation of the release commit, checks
the built distributions, and smoke-tests the wheel in a fresh environment.
+- NumPy-format docstrings with runnable examples for the whole public API. The
+ API reference is generated from them (`docs/source/reference/`); the
+ hand-written `api.md` is now a short overview, and the CUDA headers have
+ their own page. CI checks the docstrings (ruff `D` rules, numpydoc
+ validation), runs their examples as doctests, and builds the docs with
+ warnings as errors. Docstring style: `docs/source/contributing.md`.
### Fixed
- The ordering guide identifies 0.6.3 as the first release preserving F order
in CuPy-to-host conversions.
+- `CudaStruct.from_signature` accepts quoted pyccel annotations
+ (`markers: "float[:, :]"`) in modules with `from __future__ import annotations`.
## [0.6.3] - 2026-10-08
diff --git a/README.md b/README.md
index a71001a..b1fc5fa 100644
--- a/README.md
+++ b/README.md
@@ -299,8 +299,9 @@ positional argument 0 is written, avoiding unnecessary copy-back for
read-only inputs. For a keyword call, declare the keyword name, such as
`outputs=("out",)`. A wrong declaration can leave GPU output values stale.
The wrapper can also traverse arrays nested in lists, tuples, dictionaries,
-and selected application objects; see the full [API reference](docs/source/api.md)
-for `object_modules`, `is_array`, aliasing, and output declarations.
+and selected application objects; see [Host kernels with GPU
+data](docs/source/kernels/pyccel-kernel.md) for `object_modules`, `is_array`,
+aliasing, and output declarations.
## Write CUDA kernels next to host kernels
@@ -435,8 +436,9 @@ Launches can be 1D to 3D (`n_threads=(nx, ny)`, `block_size=(16, 16)`) or use
an explicit `grid`, with dynamic shared memory (`shared_mem`) and a `stream`.
C++ function templates are instantiated with `template_args`, and
`CudaKernelVariants` caches kernels whose source is generated per variant
-(e.g. per dimension and dtype). See the [API reference](docs/source/api.md) for
-details.
+(e.g. per dimension and dtype). See [Writing CUDA
+kernels](docs/source/kernels/cuda-kernel.md) and [CUDA
+headers](docs/source/kernels/cuda-headers.md) for details.
Kernels run asynchronously, so a CUDA error (an illegal memory access, say)
normally surfaces at a later `.get()` or MPI call, far from the kernel that
@@ -523,8 +525,9 @@ published at :
* [Worked examples](docs/source/examples/index.md),
[best practices](docs/source/best-practices.md),
[troubleshooting](docs/source/troubleshooting.md),
- [Pyodide](docs/source/pyodide.md) and the
- [API reference](docs/source/api.md).
+ [Pyodide](docs/source/pyodide.md), the
+ [API overview](docs/source/api.md) and the
+ [API reference](docs/source/reference/index.md).
### For AI coding assistants
diff --git a/conftest.py b/conftest.py
new file mode 100644
index 0000000..a162124
--- /dev/null
+++ b/conftest.py
@@ -0,0 +1,17 @@
+"""Doctest setup: examples assume ``import cunumpy as xp`` and ``import numpy as np``."""
+
+import numpy as np
+import pytest
+
+import cunumpy as xp
+
+
+@pytest.fixture(autouse=True)
+def _doctest_namespace(request, doctest_namespace):
+ if not isinstance(request.node, pytest.DoctestItem):
+ yield # leave the backend of the regular tests alone
+ return
+ doctest_namespace["np"] = np
+ doctest_namespace["xp"] = xp
+ with xp.use_backend("numpy"):
+ yield
diff --git a/docs/source/api.md b/docs/source/api.md
index eda5916..fb00929 100644
--- a/docs/source/api.md
+++ b/docs/source/api.md
@@ -1,8 +1,11 @@
-# API reference
+# API overview
CuNumpy exports the active NumPy-like namespace and a set of helpers for
backend selection, array inspection and conversion, hardware control, and
-host-only kernels. Most examples use `import cunumpy as xp`.
+host-only kernels. Most examples use `import cunumpy as xp`. This page shows how
+the package is laid out; every function and class is documented in the [API
+reference](reference/index.md), generated from the docstrings, and the guides
+explain how to use them together.
## NumPy-like namespace
@@ -17,9 +20,10 @@ At runtime, NumPy-like attributes such as `array`, `sum`, `fft`, and `linalg`
are those of the currently selected `array-api-compat` NumPy or CuPy module.
They are copied into the `cunumpy` namespace and replaced when the backend
changes, so `xp.sum` costs the same as `numpy.sum` (a switch between NumPy and
-CuPy takes some tens of microseconds; avoid switching inside a hot loop). CuNumpy does not wrap every operation individually. The available
-operations and some details can therefore vary with the installed NumPy and
-CuPy versions. In normal use, access those operations through the top-level
+CuPy takes some tens of microseconds; avoid switching inside a hot loop).
+CuNumpy does not wrap every operation individually. The available operations
+and some details can therefore vary with the installed NumPy and CuPy
+versions. In normal use, access those operations through the top-level
`cunumpy` namespace, commonly imported as `xp`.
For an explanation of what the compatibility module does and why CuNumpy
@@ -39,2452 +43,27 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not
| Submodule | Backends | Contents |
|---|---|---|
-| `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` |
-| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` |
-| `cunumpy.arguments` | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` |
-| `cunumpy.cuda` | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) |
-| `cunumpy.rng` | both | `random_streams`, `get_rng`, `philox_*` |
-| `cunumpy.algorithms` | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` |
-| `cunumpy.mpi` | both | `mpi_buffer`, CUDA-aware MPI detection, `local_rank`, `synchronize_for_mpi` |
-| `cunumpy.profiling` | both | `timed_region`, `nvtx_range`, `count_transfers`, `assert_no_transfers`, `TransferBudget` |
-| `cunumpy.memory` | both | `HostStaging`, `HostCopy`, `DeviceMirror` |
-| `cunumpy.petsc` | both | `petsc_vec` |
-| `cunumpy.kernel_testing` | both | pytest helpers for host/CUDA kernel pairs (not imported by `import cunumpy`) |
+| [`cunumpy`](reference/top-level.md) | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `as_device_array`, `to_host_async`, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` |
+| [`cunumpy.kernels`](reference/kernels.md) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` |
+| [`cunumpy.arguments`](reference/arguments.md) | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` |
+| [`cunumpy.cuda`](reference/cuda.md) | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) |
+| [`cunumpy.rng`](reference/rng.md) | both | `random_streams`, `get_rng`, `philox_*` |
+| [`cunumpy.algorithms`](reference/algorithms.md) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` |
+| [`cunumpy.mpi`](reference/mpi.md) | both | `mpi_buffer`, `MPIStaging`, CUDA-aware MPI detection, `synchronize_for_mpi`, and the MPI stand-ins of maybempi (`get_mpi`, `local_rank`, ...) |
+| [`cunumpy.profiling`](reference/profiling.md) | both | `timed_region`, `nvtx_range`, `count_transfers`, `assert_no_transfers`, `TransferBudget` |
+| [`cunumpy.memory`](reference/memory.md) | both | `HostStaging`, `HostCopy`, `DeviceMirror` |
+| [`cunumpy.petsc`](reference/petsc.md) | both | `petsc_vec` |
+| [`cunumpy.kernel_testing`](reference/kernel-testing.md) | both | pytest helpers for host/CUDA kernel pairs (not imported by `import cunumpy`) |
"Both" means the functions work on NumPy and CuPy arrays; the functions of
`cunumpy.cuda` do nothing (or return `None`/`0`) on the NumPy backend.
+The CUDA headers shipped for hand-written kernels (`cunumpy/atomic.cuh`,
+`cunumpy/reduce.cuh`, `cunumpy/random.cuh`, ...) have no Python API; they are
+documented in [CUDA headers](kernels/cuda-headers.md).
+
Each name has one import path, through its submodule: `xp.kernels.CudaKernel`,
`xp.arguments.CudaStruct`, `xp.mpi.mpi_buffer`, `xp.rng.random_streams`. The top
level of `cunumpy` is the NumPy/CuPy namespace plus backend selection and array
conversion, so it never hides a NumPy or CuPy name (`xp.fuse` is CuPy's `fuse`;
cunumpy's is `xp.kernels.fuse`). Modules starting with `_` are private.
-
-## Version
-
-### `require_version(minimum)`
-
-Raises `ImportError` if the installed cunumpy is older than `minimum`
-(`xp.require_version("0.4.0")`). Only the numeric parts are compared; nothing
-is checked when the version is unknown (not installed as a package).
-
-## Backend selection
-
-### `set_backend(backend, *, strict=False)`
-
-Selects the process-wide backend used for new NumPy-like operations. Supported
-values are `"numpy"` and `"cupy"`:
-
-```python
-xp.set_backend("cupy")
-values = xp.arange(10)
-```
-
-Requesting CuPy selects it only when CuPy and its CUDA runtime are functional;
-otherwise CuNumpy falls back to NumPy. Check `get_backend()` to inspect the
-effective selection. Changing the selection does not move arrays that have
-already been created.
-
-With `strict=True`, requesting unavailable CuPy raises `RuntimeError` with the
-availability failure reason, preserving the previous backend and namespace.
-`use_backend(backend, strict=True)` offers the same guarantee.
-
-### `backend_info()`
-
-Returns a JSON-compatible dictionary with the selected backend, cached CUDA
-availability and failure reason, dependency versions, and current device/runtime
-information when available. Inspection errors are included in the result. It
-does not change the backend or initialize MPI.
-
-The initial backend is NumPy unless `CUNUMPY_BACKEND=cupy` is set before CuNumpy
-is imported. Other values of this environment variable result in the NumPy
-default.
-
-Backend state is shared process-wide and `set_backend()` is not thread-safe.
-Concurrent tasks that change the backend may race.
-
-### `get_backend()`
-
-Returns the active global backend name, either `"numpy"` or `"cupy"`. This is
-the getter paired with `set_backend()`:
-
-```python
-xp.set_backend("numpy")
-assert xp.get_backend() == "numpy"
-```
-
-### `use_backend(backend)`
-
-Context manager that temporarily selects a backend and restores the previous
-backend on exit, even if the block raises an exception:
-
-```python
-with xp.use_backend("numpy"):
- reference = xp.zeros(10)
-```
-
-As backend selection is shared process state, this context manager is suited
-to sequential use rather than concurrent backend switching.
-
-### `numpy_backend`, `cupy_backend`
-
-Boolean properties indicating whether the currently selected global backend
-is NumPy or CuPy:
-
-```python
-if xp.cupy_backend:
- print("new arrays are being created on the GPU")
-```
-
-For a string value, prefer `get_backend()`.
-
-## Inspect arrays and select operations
-
-### `get_array_backend(array)`
-
-Returns `"numpy"` or `"cupy"` according to the given array's type. It reports
-the array's location, not the active global selection:
-
-```python
-xp.set_backend("numpy")
-values_gpu = xp.to_cupy([1, 2, 3])
-assert xp.get_backend() == "numpy"
-assert xp.get_array_backend(values_gpu) == "cupy"
-```
-
-### `get_array_module(array)`
-
-Returns the `array-api-compat` module matching the array: the NumPy module for
-a NumPy array or the CuPy module for a CuPy array. This supports functions
-that dispatch based on their input rather than global state:
-
-```python
-def standardize(values):
- array_xp = xp.get_array_module(values)
- mean = array_xp.mean(values)
- scale = array_xp.std(values)
- return (values - mean) / scale
-```
-
-The returned module is an `array-api-compat` module, consistent with the
-active module used by CuNumpy. It is not necessarily identical to importing
-raw `numpy` or raw `cupy`.
-
-### `is_cpu(array)`, `is_gpu(array)`
-
-Return booleans indicating whether an array is a NumPy (CPU) or CuPy (GPU)
-array. They are convenience checks equivalent to comparing
-`get_array_backend(array)` with `"numpy"` or `"cupy"`.
-
-### `same_backend(*arrays)`
-
-Returns `True` if all provided arrays have the same backend. Zero or one
-argument is considered to match:
-
-```python
-if xp.same_backend(position, velocity):
- update(position, velocity)
-```
-
-### `assert_same_backend(*arrays)`
-
-Raises `TypeError` with the detected backend names if arrays do not all share
-a backend. Use it at API boundaries to provide a clear error before a mixed
-NumPy/CuPy operation fails deeper in a library:
-
-```python
-def combine(left, right):
- xp.assert_same_backend(left, right)
- return left + right
-```
-
-## Convert arrays
-
-### `to_numpy(array)`
-
-Converts to a host-side NumPy array. CuPy arrays are copied from device to
-host. Other array-like inputs are passed through `numpy.asarray`; NumPy arrays
-may therefore be returned as-is rather than copied.
-
-F-contiguous CuPy inputs keep F order on the host; other CuPy inputs become
-C-contiguous. Arbitrary device strides are not preserved. See
-[Array ordering and strides](guides/array-ordering.md).
-
-### `to_cupy(array)`
-
-Converts an array-like input to a CuPy array. Raises `ImportError` if CuPy or
-CUDA is unavailable or not functional. The source is not modified.
-
-### `to_cunumpy(array)`
-
-Converts to the currently active backend. This is convenient at an API
-boundary when an input should be normalized to the configured backend:
-
-```python
-normalized = xp.to_cunumpy(input_array)
-assert xp.get_array_backend(normalized) == xp.get_backend()
-```
-
-Each conversion returns a suitable array; it does not change the active
-backend or mutate the source.
-
-### `to_host_async(array)`
-
-```python
-pending = xp.to_host_async(residual_norm) # a device scalar; returns at once
-... # queue the next iteration's kernels
-if pending.ready() and pending.result() < tol:
- break
-```
-
-Starts copying a device scalar or small array to the host without waiting for
-it, so that the host can keep queueing kernels (e.g. a convergence test read one
-iteration late, without a sync per iteration). The copy runs on a separate
-stream into page-locked memory, after the work queued so far on the current
-stream. The returned `memory.HostCopy` has `ready()` (never waits) and
-`result()` (waits only if the copy has not finished; a NumPy scalar for a 0-d
-array, else a NumPy array). `count_transfers()` records a `to_host` event with
-`blocking=False`, and a wait in `result()` as a `sync`.
-
-On the NumPy backend the value is copied at once and `ready()` is always True;
-on the fake CuPy too, and the copy is counted. For large arrays copied
-repeatedly, use `memory.HostStaging`, which reuses its buffers.
-
-### `algorithms.segment_sum(values, keys, n_segments, *, out=None)`
-
-`out[k] = sum(values[i] for keys[i] == k)` on the backend of `keys`, with
-one CUDA accumulation launch for all trailing components. `values` has shape
-`(n, ...)`; a negative key drops the row; integer keys must be smaller than
-`n_segments`. The result has shape `(n_segments, ...)`, keeps a floating-point
-or complex dtype, and is `float64` otherwise. Optional `out` is overwritten,
-must match shape/dtype, be writable and C-contiguous, and must not alias values.
-Mixed backends raise. CUDA float16 accumulates in float32; floating-point atomic
-accumulation order is not deterministic.
-
-### `algorithms.SegmentPlan(keys, n_segments)`
-
-Copies and validates keys once, then `plan.sum(values, out=None)` reuses them.
-Changing the original keys does not change the plan. GPU preparation may
-synchronize; repeated sums avoid key validation and per-column scalar reads.
-A GPU plan requires its device current and values/output on that device.
-
-### `algorithms.cell_offsets(sorted_cells, n_cells)`
-
-Returns `n_cells+1` int64 offsets on the input backend. Cell k occupies the
-half-open slice `[offsets[k], offsets[k+1])`. Missing cells have equal endpoints.
-Input must be sorted integer IDs in `[0, n_cells)`; filter negative IDs first.
-
-### `algorithms.segment_boundaries(sorted_keys)`
-
-Returns `(unique_keys, starts, stops)` for runs of sorted integer keys. Supports
-sparse, negative, and uint64 Morton keys. Starts/stops are int64 half-open indices.
-Empty input returns three empty arrays. Validation and variable-length GPU
-output may synchronize; prepare boundaries outside repeated operations.
-
-### `algorithms.sort_by_key(keys, *arrays, axis=0)`
-
-Stable argsort of the 1D `keys` (CuPy's radix sort on the device), applied to
-every array along `axis`, in one call:
-
-```python
-keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, charges)
-# component-major (ncomp, N) markers: sort along the last axis of each array
-keys, order, positions, weights = xp.algorithms.sort_by_key(
- keys, positions, weights, axis=-1
-)
-```
-
-Returns `(keys[order], order, *(take(a, order, axis) for a in arrays))`,
-`order` as `int64`. Equal keys keep their order, so the result is reproducible.
-On NumPy, integer keys of more than 4096 entries are sorted by an LSD radix
-sort on 16-bit digits (one stable `argsort` of `uint16` per digit, as many as
-the key range needs), about ten times faster than the stable sort of 64-bit
-integers.
-
-### `algorithms.compact_by_mask(mask, *arrays, axis=0)`
-
-Moves the rows where the boolean `mask` is True to the front of every array, in
-place and in their original order, and returns how many there are. Typical use:
-keep the live particles at the front of the marker arrays.
-
-```python
-n = xp.algorithms.compact_by_mask(alive, markers, weights)
-markers, weights = markers[:n], weights[:n]
-```
-
-The rows after the first `n` are unspecified. The count is needed on the host,
-so on CuPy each call synchronizes once. The mask and the arrays must be on the
-same backend. `axis` is the axis of every array that the mask indexes; `-1`
-compacts component-major arrays, `(ncomp, N)` next to `(N,)`, along their
-marker axis:
-
-```python
-n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1)
-positions, weights = positions[:, :n], weights[:n]
-```
-
-## Count transfers
-
-A transfer inside a time loop is the classic performance bug of a GPU port:
-every step then waits for the device and copies an array. These helpers let a
-test verify that a block of code does not transfer at all.
-
-### `profiling.count_transfers(into=None)`
-
-Context manager yielding a `TransferCounter` that records every host/device
-transfer made through CuNumpy while the block runs, with the call site of
-each:
-
-```python
-with xp.profiling.count_transfers() as counter:
- propagator(dt)
-
-assert counter.total == 0, counter.report()
-```
-
-Six kinds of events are recorded:
-
-* `to_host`: an actual device-to-host copy through conversion, mirror, staging,
- serial MPI, or host kernel helpers;
-* `to_device`: an actual host-to-device copy through those helpers, including
- `as_device_array()` and `kernel_output()` copy-back;
-* `kernel_conversion`: a `PyccelKernel` call that copied device arrays to the
- host (and back), one event per call, naming the kernel and the number of
- arrays converted;
-* `fallback`: a `Kernel` without CUDA kernel calling its host kernel on the
- CuPy backend (`missing_cuda="fallback"`), one event per call, naming the
- kernel. Physical copies and a `kernel_conversion` marker are recorded separately;
-* `device_copy`: a device-only dtype/layout conversion through CuNumpy helpers;
-* `sync`: the host waited for the device: `xp.synchronize()`, the waits of the MPI
- helpers (`synchronize_for_mpi()`, `mpi_buffer()` staging) and of the CUDA debug
- mode, and, on the fake CuPy, a scalar read of a device array (`float(a)`,
- `int(a)`, `bool(a)`, `a.item()`, `a.tolist()`). They are in `counter.syncs` and in
- the report, but not in `total`.
-
-Only real transfers count: `to_numpy()` of a NumPy array or `to_cupy()` of a
-CuPy array records nothing. The counter has the attributes `to_host`,
-`to_device`, `kernel_conversions`, `fallbacks` (counts per kind), `total`,
-`device_copies`, `bytes_to_host`, `bytes_to_device`, and `bytes(kind)`.
-`total` includes physical copies and conversion/fallback markers. Byte totals
-include only physical copies, so markers do not double-count bytes.
-`events` is a list of `TransferEvent(kind, description, where, nbytes=None,
-blocking=True, implicit=False)`; `where` is the caller's `file:line` and `nbytes`
-is None for markers/unknown sizes. `blocking` is False for the copies that the host
-does not wait for (`to_host_async()`, `HostStaging.copy()`), and `implicit` is True
-for the scalar reads of the fake CuPy (`float(a)`), which are `sync` events; the
-report marks them `[async]` and `[implicit]`.
-`kernel_conversion_calls` selects the conversion markers. `report()` returns
-a multi-line string with the events grouped by kind and call site, with
-counts:
-
-```text
-4 transfer(s) through cunumpy (3 to_host, 1 to_device, 0 kernel_conversion, 0 fallback, 0 device_copy)
- to_host (3):
- /home/me/sim/diagnostics.py:42: to_numpy(shape=(100000,), dtype=float64) (x3)
- to_device (1):
- /home/me/sim/setup.py:17: to_cupy(shape=(100000,), dtype=float64)
-```
-
-Blocks can be nested; every active counter sees the transfers made inside it.
-`count_transfers(counter)` adds the events to an existing counter instead, e.g. to
-accumulate over several calls; a counter that is already active is not added
-again, so each event is counted once. Like every `contextlib.contextmanager`, it
-is also a decorator: `@count_transfers(counter)`.
-When no counter is active, the instrumentation costs a single check per call.
-Like the backend selection, the active counters are process-wide state and
-not thread-safe.
-
-**Limitation:** only transfers made through CuNumpy are seen. Raw
-`cupy.ndarray.get()`, `cupy.asarray(numpy_array)`, `numpy.asarray(cupy_array)`,
-forwarded backend calls such as `xp.asarray()`, and implicit conversions inside
-other libraries are not counted. Neither is `float(device_array)` (an implicit
-sync) with the real CuPy, which cannot be observed from Python; the fake CuPy
-counts it. Use `nsys` (or CuPy's profiling hooks) to find those.
-
-### `profiling.assert_no_transfers(*, syncs=False)`
-
-Context manager that raises `AssertionError` with the counter's `report()` if
-the block makes a host/device transfer or host fallback through CuNumpy.
-Device-only dtype/layout conversions are allowed, and so are syncs unless
-`syncs=True`. It yields the `TransferCounter` too. An exception raised inside the block propagates as it is:
-
-```python
-def test_time_step_stays_on_the_device():
- with xp.profiling.assert_no_transfers():
- propagator(dt)
-```
-
-### `profiling.TransferBudget(*, started=True)`
-
-Counts the transfers per phase of a program (the time step, the diagnostics,
-the output) and checks a rule for each phase:
-
-```python
-budget = xp.profiling.TransferBudget(started=False)
-model.integrate = budget.count("integrate")(model.integrate) # a decorator
-... # setup: not counted
-budget.start()
-for step in range(n_steps):
- model.integrate(dt)
- with budget.phase("output"): # or a context
- save(model)
-
-budget.require("integrate", allow={"to_host": {"max_nbytes": 8}}, calls=n_steps)
-budget.require("output", allow={"to_host": {"max_count": n, "max_total_bytes": b}})
-budget.check() # AssertionError with budget.report() if a rule is broken
-```
-
-`phase(name)` counts its block in phase `name` and yields the phase's
-`TransferCounter`; `count(name)` is a decorator that does the same for every
-call. The events of a phase accumulate over its calls (`budget[name]`,
-`budget.phases`), and `budget.calls[name]` counts the calls. Phases nest: an
-event is counted in the innermost phase only, and a phase entered again inside
-itself (recursion) counts each event and call once. Until `start()` (with
-`started=False`), and after `stop()`, phases run uncounted.
-
-`require(phase, allow=None, *, ignore=("device_copy", "sync"), calls=None)`
-sets the rule of a phase: the event kinds in `allow` are allowed, each with
-optional limits, every other kind is forbidden except those in `ignore` (unless
-they are in `allow`). The limits are `max_nbytes` (each event; an event of
-unknown size breaks it), `max_count` and `max_total_bytes` (the whole phase),
-and `blocking` and `implicit` (the events must have this value), e.g.
-`{"to_host": {"max_nbytes": 8, "blocking": False}}` to allow only non-blocking
-scalar copies, or `{"sync": {"implicit": False}}` with `ignore=("device_copy",)`
-to reject the scalar reads of the fake CuPy. `calls` is the number of calls the
-phase must have had. `violations()` lists the broken rules, with the offending
-events and where they happened; `report()` shows every phase's events and the
-broken rules; `check()` raises `AssertionError` with the report.
-
-### `as_device_array(value, dtype=None, ndim=None, *, name=None)`
-
-The "reference or copy once" rule for building CUDA argument objects
-(`CudaArguments` subclasses, `CudaStruct` values). Call it once when the
-argument object is built, never per kernel call:
-
-* a CuPy array that already has `dtype` (any dtype if `dtype` is `None`) and
- is C-contiguous is returned unchanged, the same object without a copy, so
- kernels write into the caller's array;
-* anything else is converted to a C-contiguous device array using
- `cupy.ascontiguousarray(cupy.asarray(value, dtype))`: tuples and lists
- (`degree = (3, 3, 3)`), host NumPy arrays (one explicit transfer at build
- time), device arrays of another dtype, and non-contiguous views.
-
-Dtype and layout conversion can require separate device copies; each copy
-through this helper appears in transfer accounting.
-
-The result passes the pointer checks of `CudaKernel` and `CudaStruct`. On the
-NumPy backend it raises `RuntimeError`: device arguments are only built when
-running on CuPy, and host data is never copied to the device implicitly. If
-`ndim` is given and the result has another number of dimensions, it raises
-`ValueError`; `name` is the argument name used in error messages.
-
-```python
-class DeviceParticles(xp.arguments.CudaArguments):
- def __init__(self, markers, degree):
- self.markers = xp.as_device_array(markers, np.float64, ndim=2, name="markers")
- self.degree = xp.as_device_array(degree, np.int32, ndim=1, name="degree")
- super().__init__(self.markers, self.degree, self.markers.shape[0])
-```
-
-### `kernels.as_kernel_array(value, like, dtype=None, *, strided=False)`, `kernels.kernel_output(out, like, dtype=None, *, strided=False)`
-
-For the arguments of a `Kernel` with `dispatch="arrays"`, whose choice follows
-the arrays. `as_kernel_array` returns `value` on the side of `like` (a CuPy
-array if `like` is one, a NumPy array otherwise), C-contiguous and with `dtype`
-(any if `None`): `value` itself if it already is such an array, else a conversion,
-moved across if needed (counted by `count_transfers()`). `kernel_output` is a
-context manager yielding the buffer for an array the kernel writes: `out`
-itself if `as_kernel_array` takes it unchanged, else a converted copy whose
-contents are written into `out` (on its own side) when the block ends without
-an error.
-
-With `strided=True` a NumPy array for a host kernel is also taken unchanged
-when it is in C order with gaps (positive strides, each at least the extent of
-the next axis), e.g. the first `n` columns `storage[:, :n]` of a marker buffer.
-Pyccel's wrappers take such arrays without a copy; a host kernel that needs
-contiguous memory must not use it. F-ordered, transposed and negative-stride
-arrays are still copied, and CuPy arrays are always made C-contiguous.
-
-```python
-with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer:
- gather(
- xp.kernels.as_kernel_array(positions, like=field, dtype=float), field, buffer
- )
-```
-
-## Random numbers and dtype
-
-### `rng.get_rng(seed=None)`
-
-Returns a NumPy or CuPy `Generator` matching the active backend:
-
-```python
-rng = xp.rng.get_rng(seed=7)
-samples = rng.uniform(size=100)
-```
-
-The generator APIs are similar, but seeds do not guarantee identical random
-sequences across NumPy and CuPy.
-
-### `rng.random_streams`
-
-```python
-xp.rng.random_streams.seed(42, rank=comm.Get_rank(), bit_generator="PCG64")
-v = xp.rng.random_streams.normal(0.0, v_th, (n, 3))
-rng = xp.rng.random_streams.generator() # numpy or cupy Generator
-own = xp.rng.random_streams.make_generator(seed) # a component's own generator
-```
-
-One seeded random generator per process and backend, for reproducible MPI
-runs: after `seed(value, rank)` every draw of the process comes from the stream
-`(value, rank)` (a NumPy `SeedSequence` with the rank as spawn key), so the same
-seed and number of ranks give the same results and the ranks' streams are
-independent. `seed(None)` (or no seed) seeds from the operating system.
-`bit_generator` selects the NumPy bit generator (`MT19937`, `PCG64`,
-`PCG64DXSM`, `Philox`, `SFC64`); the CuPy generator uses CuPy's default.
-`seed` also seeds NumPy's global state (and CuPy's on the CuPy backend) for code
-that calls `np.random.*` directly.
-
-`generator(backend=None)` returns the process generator of a backend (the
-active one by default), created on first use. `make_generator(seed=None,
-backend=None)` returns a separate generator for a component with a seed of its
-own, and the process generator otherwise. `random`, `standard_normal`,
-`normal` and `uniform` draw from the process generator (or `rng=`); `normal`
-and `uniform` fall back to `standard_normal` and `random` for CuPy generators
-without those methods. `xp.rng.RandomStreams()` makes an independent instance.
-
-### `default_float_dtype()`
-
-Returns the active backend module's `float64` dtype object. Pass it to array
-creation when code requires an explicit precision:
-
-```python
-values = xp.asarray([0.1, 0.2], dtype=xp.default_float_dtype())
-```
-
-## CUDA devices and memory
-
-### `cupy_available()`
-
-Returns whether CuPy can be imported and reports itself functional. The result
-is cached for the process. This checks availability, not whether every
-GPU-specific operation will succeed later.
-
-### `cuda.device_count()`
-
-Returns the number of visible CUDA devices. Returns zero when CuPy/CUDA is
-unavailable or querying the runtime fails. This is independent of the active
-backend, so it may return a positive number while NumPy is selected.
-
-### `cuda.set_device(device_id)`
-
-Selects a CUDA device when CuPy is active; it is a no-op on NumPy. The device
-must be valid for the current CUDA process.
-
-### `cuda.set_device_for_rank(rank, devices_per_node=None)`
-
-Selects a device using `rank % devices_per_node` and returns its ID. If
-`devices_per_node` is omitted, it uses `device_count()`. When no devices are
-visible, it returns `0` without selecting a device. This helper assumes
-contiguous rank-to-device mapping on each node, suitable for a common
-one-rank-per-GPU MPI layout. Use `set_device()` directly when the scheduler's
-mapping differs:
-
-```python
-device_id = xp.cuda.set_device_for_rank(mpi_rank)
-```
-
-### `mpi.get_mpi(use_mpi=None)`, `mpi.launched_under_mpi()`
-
-Importing `mpi4py.MPI` starts MPI (`MPI_Init`), which takes time and makes
-every collective cost something even on one process. `launched_under_mpi()`
-tells, from the environment that `mpirun`/`mpiexec`/`srun` set up and without
-importing mpi4py, whether the process belongs to an MPI job
-(`MAYBEMPI=1`/`0` overrides it). These functions and the stand-in come from
-[maybempi](https://max-models.github.io/maybempi/) and are re-exported in `xp.mpi`. `get_mpi()` returns `mpi4py.MPI` then, and
-otherwise the serial stand-in, so that the same code runs with and without
-MPI:
-
-```python
-MPI = xp.mpi.get_mpi() # decided once per process
-comm = MPI.COMM_WORLD
-comm.Allreduce(MPI.IN_PLACE, rho, op=MPI.SUM) # nothing to do on one process
-n_total = comm.allreduce(n_local) # n_local itself
-if xp.mpi.is_serial(MPI):
- ... # a serial run
-```
-
-`get_mpi(True)` imports mpi4py (`ImportError` if missing), `get_mpi(False)`
-returns the stand-in. Launched under MPI without mpi4py installed, `get_mpi()`
-warns (`RuntimeWarning`) and returns the stand-in.
-
-### `mpi.SerialMPI`, `mpi.SerialComm`
-
-The stand-in for `mpi4py.MPI` and its communicators of size 1.
-`SerialComm` returns (object methods) or copies (buffer methods) what one rank
-gets: `bcast`, `allreduce`, `reduce` and `scan` return their argument,
-`gather`/`allgather` return `[x]`, `scatter([x])` returns `x`;
-`Allreduce`/`Reduce`/`Allgather`/`Gather`/`Scatter`/`Alltoall`/`Scan` copy the
-send buffer into the receive buffer (nothing with `IN_PLACE`), the vector
-forms (`Allgatherv`, `Gatherv`, `Scatterv`) use the displacement of rank 0,
-`Bcast` and `Barrier` do nothing, and `sendrecv`/`Sendrecv` work to and from
-rank 0 or `PROC_NULL`. Buffers are NumPy or CuPy arrays, or mpi4py buffer
-specifications (`[array, MPI.DOUBLE]`). Non-blocking versions return completed
-requests. Any other method raises `AttributeError`, so a missing feature shows
-instead of doing nothing. `SerialMPI` has `COMM_WORLD`, `COMM_SELF`,
-`COMM_NULL` (false), `IN_PLACE`, the reduction operations, the common
-datatypes (`isinstance(MPI.DOUBLE, MPI.Datatype)` holds), `PROC_NULL`, `ROOT`,
-`ANY_SOURCE`, `UNDEFINED`, `Comm`/`Intracomm`, `Request`, `Prequest`, `Status`,
-`Wtime()` and `Is_initialized()` (False).
-
-### `mpi.local_rank()`
-
-The rank of the process within its node, read from the environment variables
-that MPI launchers export (Open MPI, MVAPICH2, Intel MPI/MPICH, PMI, Cray
-PALS, Slurm, `LOCAL_RANK`), or `0` if none is set. The launcher sets them
-before `MPI_Init`, so this works before MPI is initialized and without
-importing `mpi4py`.
-
-### `cuda.bind_local_device()`
-
-Selects device `local_rank() % device_count()` for this process and creates its
-CUDA context. Returns the device id, or `None` on the NumPy backend or without
-devices. Call it before `MPI_Init` (before importing `mpi4py.MPI`), so that a
-CUDA-aware MPI sees the right device; otherwise all ranks of a node would use
-device 0. If the launcher gives each rank its own device through
-`CUDA_VISIBLE_DEVICES`, each process sees one device and selects it:
-
-```python
-import cunumpy as xp
-
-xp.set_backend("cupy")
-xp.cuda.bind_local_device()
-from mpi4py import MPI # initializes MPI after the device is bound
-```
-
-Unlike `set_device_for_rank()`, it needs no MPI rank, and it uses the rank
-within the node rather than assuming contiguous ranks per node.
-
-### `mpi.mpi_is_cuda_aware(comm=None, *, method="probe")`
-
-Checks whether the MPI library can send and receive device (CuPy) buffers,
-which needs a CUDA-aware MPI build; with a plain build, passing a CuPy array
-to MPI segfaults or silently sends garbage. Returns `False` on the NumPy
-backend and without a functional CuPy, without importing `mpi4py`: the
-question only makes sense with device buffers. `comm` defaults to
-`mpi4py.MPI.COMM_WORLD`, and `mpi4py` is imported only then.
-
-The check is collective: every rank of `comm` must call it, and all ranks get
-the same result. Each rank sends a tiny device array to rank
-`(rank + 1) % size` and receives from `(rank - 1) % size` with `Sendrecv`
-(after `synchronize_for_mpi()`; with a single rank, it sends to itself), checks
-the received values, and the ranks combine their outcomes with
-`allreduce(op=LAND)`. Any exception in the exchange, on any rank, gives
-`False`. Only `method="probe"` exists: `mpi4py` does not expose the library
-query (`MPIX_Query_cuda_support`) and the library version string is not a
-reliable indicator.
-
-An MPI library that is not CUDA-aware may also read the device address as a
-host address and crash the process. A segfault inside this call therefore
-means the same thing as `False`. Call it once at startup, after
-`bind_local_device()` and `MPI_Init`, before any communication of device
-buffers.
-
-### `mpi.require_cuda_aware_mpi(comm=None)`
-
-Raises `RuntimeError`, explaining how to get a CUDA-aware build (Open MPI
-`--with-cuda`, MPICH with a CUDA-enabled UCX, the site's CUDA-aware MPI
-module), if `mpi_is_cuda_aware(comm)` returns `False` on the CuPy backend.
-No-op on the NumPy backend. The complete startup sequence for one rank per
-GPU:
-
-```python
-import cunumpy as xp
-
-xp.set_backend("cupy")
-xp.cuda.bind_local_device() # 1. select the GPU, before MPI_Init
-from mpi4py import MPI # 2. MPI_Init, on the bound device
-
-xp.mpi.require_cuda_aware_mpi() # 3. clear error instead of a segfault later
-
-xp.mpi.synchronize_for_mpi(send, recv) # 4. before every MPI call with device buffers
-MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source)
-```
-
-### `mpi.synchronize_for_mpi(*arrays, stream=None, event=None)`
-
-Waits for all work on devices owning the CuPy arrays, or only the explicit
-producer stream/event when supplied (pass at most one). The dependency must
-cover all supplied device buffers. `None` entries and host arrays are ignored, so it costs
-nothing for host buffers and on the NumPy backend. Call it before every MPI
-call that sends or receives device buffers: CuPy launches kernels
-asynchronously and MPI knows nothing about CUDA streams, so a buffer that a
-kernel is still writing would be sent as it is at that moment, without an
-error:
-
-```python
-xp.mpi.synchronize_for_mpi(send_buffer, recv_buffer)
-comm.Sendrecv(send_buffer, dest, recvbuf=recv_buffer, source=source)
-```
-
-No synchronization is needed after MPI returns: kernels launched afterwards see
-the received data.
-
-### `mpi.mpi_buffer(array, *, send=True, recv=False, cuda_aware=None, staging=None, stream=None, event=None)`
-
-Context manager yielding the buffer to pass to MPI for `array`: a host array
-unchanged; a device array unchanged (after `synchronize_for_mpi`) when MPI is
-CUDA-aware; otherwise a pinned host staging buffer, filled from the device
-before the block (`send`) and copied back after it (`recv`), both counted by
-`count_transfers()`. `cuda_aware=None` uses the answer recorded by
-`mpi_is_cuda_aware()` or `set_mpi_cuda_aware()`; without one, a device array
-raises `RuntimeError`.
-
-Pass `stream` or `event` for producer synchronization. Receive-only device
-staging also waits for preceding device work. Receive copy-back completes before
-the context releases host storage. With nonblocking MPI, call `request.Wait()`
-inside the context before reading/releasing the buffer.
-
-### `mpi.MPIStaging(shape, dtype)`
-
-Caller-owned reusable host staging storage, pinned when available and allocated
-lazily. Use `staging.buffer(array, **options)` or
-`mpi_buffer(array, staging=staging, **options)`. Device staging is bound to a
-shape, dtype, and device, and rejects overlapping uses. Host arrays still pass
-through unchanged. Use separate staging instances for simultaneously active
-send and receive buffers.
-
-Non-C-contiguous device arrays use a reusable device packing buffer; received
-values are written back into the original view, preserving its storage.
-
-### `mpi.set_mpi_cuda_aware(value)`, `mpi.get_mpi_cuda_aware()`
-
-Record (or read) whether MPI can take device buffers, for `mpi_buffer()`.
-`mpi_is_cuda_aware()` records its own result; set it by hand when the answer
-is known otherwise, or `None` to forget it.
-
-### `cuda.memory_info()`
-
-Returns `(free_bytes, total_bytes)` reported by the CUDA runtime for the
-active device, or `None` on NumPy. The values cover the device, not only
-allocations owned by CuPy.
-
-### `cuda.free_memory()`
-
-Releases currently free blocks in CuPy's device and pinned-host memory pools.
-It is a no-op on NumPy. It does not release blocks still referenced by live
-arrays. CuPy normally caches freed allocations for reuse, so cached memory
-does not necessarily indicate a leak.
-
-### `cuda.max_shared_memory_per_block(device=None, *, opt_in=False)`
-
-The bytes of shared memory a block may use on the current (or given) device,
-from the device attributes, e.g. to decide whether a per-block copy of a grid
-fits. `opt_in=True` gives the larger limit of newer GPUs, which a kernel uses
-only after setting `max_dynamic_shared_size_bytes` on its compiled
-`cupy.RawKernel`. Without CuPy it returns `DEFAULT_SHARED_MEMORY_PER_BLOCK`
-(48 KiB, which every CUDA device provides).
-
-### `cuda.cuda_include_dir()`
-
-Returns the directory (as `str`) of the CUDA headers shipped with CuNumpy,
-`cunumpy/array_view.cuh`, `cunumpy/atomic.cuh`, `cunumpy/index.cuh` and `cunumpy/reduce.cuh`. `CudaKernel` adds it to its NVRTC options as
-`-I` automatically (and only once), so kernel sources can write
-`#include ` without configuration. Use it to pass the
-same headers to other compilers.
-
-### `cuda.pin_memory(array)`
-
-Copies a host array to page-locked (pinned) host memory. Pinned memory can
-improve host/device transfer throughput in suitable asynchronous workloads.
-Raises `ImportError` if CuPy is unavailable. If the input may be a CuPy array,
-first transfer it with `to_numpy()`.
-
-### `synchronize()`
-
-Waits for queued work on the current CUDA device to finish. This is useful
-before reading asynchronously computed results from host code. It is a no-op
-on NumPy.
-
-### `cuda.stream(existing=None)`
-
-Context manager that creates a non-blocking CuPy stream and yields it. Work
-issued in the block is enqueued on that stream. On NumPy, yields `None` and
-does nothing:
-
-```python
-with xp.cuda.stream() as work_stream:
- result = xp.to_cupy(host_values) * 2
-
-work_stream.synchronize() # on CuPy; the yielded value is None on NumPy
-```
-
-Do not call methods on the yielded value without checking the backend. Use
-`xp.synchronize()` for code that should work on both backends.
-
-Pass `existing` to select a previously allocated stream for the context.
-
-### `cuda.create_stream(non_blocking=True)`, `cuda.create_event(timing=False)`
-
-Create reusable CuPy streams and events. On CPU they return synchronous
-`HostStream`/`HostEvent` objects supporting context selection, `.synchronize()`,
-`.done`, recording and waits. Events disable timing by default. CUDA streams
-belong to the device current at construction; keep that device current while
-using them.
-
-### `cuda.record_event(event=None, *, stream=None)`, `cuda.wait_event(event, *, stream=None)`
-
-Record an existing or newly created completion event on the producer stream,
-then enqueue a wait on the consumer stream without blocking the CPU. Omitting
-the stream uses the current CUDA stream. All consumers of an event recording
-must enqueue their waits before that event is re-recorded. See the
-[execution helpers guide](guides/execution-helpers.md) for examples.
-
-## Profiling
-
-CUDA kernels run asynchronously: a wall-clock timer around a launch measures
-the launch, not the kernel, and regions of an application profiler are not
-visible to `nsys`. These helpers address both; they are no-ops (or plain
-timers) on NumPy, so instrumented code runs unchanged on both backends.
-
-### `profiling.nvtx_range(name, color=None)`
-
-Context manager and decorator marking a code region as an NVTX range. On CuPy
-it calls `cupy.cuda.nvtx.RangePush(name)` on entry and `RangePop()` on exit
-(also when the block raises), so the region appears on the `nsys`/Nsight
-timeline next to the kernels launched inside it. `color` is an optional index
-into NVTX's colour table (the `id_color` argument of `RangePush`). On NumPy,
-or if NVTX is not available in the CuPy build, it does nothing. The same
-instance may be nested or re-entered, e.g. as the decorator of a recursive
-function.
-
-```python
-with xp.profiling.nvtx_range("push markers"):
- kernel(markers, dt, n_threads=n)
-
-
-@xp.profiling.nvtx_range("accumulate")
-def accumulate(particles, grid): ...
-```
-
-### `profiling.timed_region(name, *, sync=True)`
-
-Context manager timing a code region, including the device work it queues.
-It yields a `Timing` object whose `elapsed` (seconds, from
-`time.perf_counter`) is set when the block exits, also when it raises. On
-CuPy it synchronizes the device on entry, so earlier queued work is not
-charged to the region, and, if `sync` is true, again on exit before reading
-the clock; `synced` records whether that happened. It also pushes an
-`nvtx_range()` of the same name. On NumPy it is a plain timer and `synced`
-is `False`. With `sync=False` only the host time is measured.
-
-```python
-with xp.profiling.timed_region("push markers") as timing:
- kernel(markers, dt, n_threads=n)
-
-print(f"{timing.name}: {timing.elapsed:.4f} s, synced={timing.synced}")
-```
-
-### `profiling.Timing`
-
-Dataclass returned by `timed_region()`, with the fields `name` (`str`),
-`elapsed` (`float`, `None` until the block exits) and `synced` (`bool`).
-
-## `kernels.PyccelKernel`
-
-### Constructor
-
-```python
-xp.kernels.PyccelKernel(
- kernel,
- use_cupy=None,
- object_modules=(),
- is_array=None,
- outputs=None,
-)
-```
-
-Wraps a callable expecting host NumPy arrays so it can be used with CuPy
-arrays. This can adapt a Pyccel-compiled kernel or an ordinary Python
-callable. CuNumpy does not compile the callable or import Pyccel.
-
-When no conversion is needed, the original callable is invoked directly. If
-conversion is needed, CuNumpy recursively replaces CuPy arrays in supported
-arguments with NumPy host copies, calls the kernel, copies in-place updates
-back to the corresponding CuPy arrays, and converts returned NumPy arrays to
-CuPy arrays.
-
-### Parameters
-
-* `kernel`: callable that accepts the host-side arguments.
-* `use_cupy`: `None` (default) chooses conversion for each call when the
- global backend is CuPy or a CuPy array is present. `True` forces conversion;
- `False` disables it.
-* `object_modules`: module prefixes whose instances should be shallow-copied
- and traversed by attributes. For example,
- `object_modules=("my_project.",)`.
-* `is_array`: predicate for host array values to convert back to CuPy. The
- default is `isinstance(value, numpy.ndarray)`.
-* `outputs`: sequence of arguments the kernel may write to. Entries are
- indices or parameter names; a name also finds a positional argument, and an
- index a keyword argument, when the parameter names are known (from the
- Python signature, or `parameters`). If omitted, every converted array is
- copied back.
-* `parameters`: the names of the positional parameters, or a function that
- returns them, for compiled kernels without a Python signature. A `Kernel`
- supplies the names of its host function.
-
-`kernels.outputs_from_annotations(function)` returns the parameters a kernel may
-write to, read from its annotations: everything that is not `Final`, `const` or a
-scalar. Use it as `outputs=outputs_from_annotations(push)`, or pass
-`outputs="annotations"` to `Kernel.from_folder()` / `KernelCatalog.from_package()`.
-
-### Example and output declarations
-
-```python
-def scale_and_shift(scale, values, out):
- out[:] = scale * values + 1
- return out
-
-
-kernel = xp.kernels.PyccelKernel(scale_and_shift, outputs=(2,))
-
-with xp.use_backend("cupy"):
- values = xp.arange(8, dtype=xp.float64)
- out = xp.empty_like(values)
- result = kernel(2.0, values, out)
-```
-
-`outputs=(2,)` marks only `out` for copy-back. Read-only `values` is copied
-to the host for the call but not transferred back. Declare a positional
-argument by index (negative indices count from the end), and a keyword
-argument by its name, such as `outputs=("out",)` for `kernel(..., out=out)`.
-The forms are not interchangeable because compiled builtins may not expose a
-Python signature. Invalid declarations raise `IndexError` or `KeyError`.
-
-An incorrect declaration is a correctness bug: if the kernel writes an
-argument that was not declared, the device array will not receive that
-update. Use `outputs=()` only when no converted input is mutated. Containers
-and selected objects can be declared as outputs; every nested supported
-array is then copied back. If an array is reachable through multiple paths,
-declaring any path that includes it is sufficient.
-
-### Supported containers, aliasing, and returns
-
-Tuples, lists, and dictionaries are traversed recursively. Objects are
-traversed only when their class module starts with one of the configured
-`object_modules` prefixes; those objects are shallow-copied, and their
-attributes are converted on the copy. Other objects are passed to the kernel
-unchanged.
-
-The wrapper memoizes conversions within a call. If the same device array is
-passed more than once or appears inside a supported container, the host kernel
-sees the same NumPy array object, preserving aliasing. Supported reference
-cycles terminate safely. Returned NumPy arrays (and arrays inside tuples or
-lists) are converted back using `is_array`; dictionaries in return values are
-not recursively converted. On the NumPy path, the original return value and
-normal Python mutation and exception behavior are preserved.
-
-## `kernels.MetalKernel`
-
-A Metal Shading Language kernel for the GPU of an Apple silicon Mac, run with
-[MLX](https://github.com/ml-explore/mlx) (`pip install 'cunumpy[metal]'`). It
-takes NumPy arrays and fills the output arrays you pass, so no backend switch
-is needed.
-
-```python
-import numpy as np
-import cunumpy as xp
-
-scale = xp.kernels.MetalKernel(
- "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];",
- inputs=["x", "a"],
- outputs=["y"],
-)
-x = np.arange(8, dtype=np.float32)
-y = np.empty_like(x)
-scale(x, 2.0, out=y)
-```
-
-`MetalKernel(source, inputs, outputs, *, name="cunumpy_kernel", header="",
-threadgroup=256, float64="error", atomic_outputs=False, init_value=None)`
-
-- `source` is the body of the kernel function. MLX generates the signature: each
- name in `inputs` and `outputs` is a pointer to the flat, row-major data of that
- array, so `x[i]` is the flat index. `thread_position_in_grid` and the other
- Metal attributes used in the body are added automatically. `header` goes before
- the function (includes, defines, helper functions).
-- Calling the kernel: `kernel(*inputs, out=array_or_arrays, n_threads=None,
- template=None)`. `n_threads` is the total thread count (default: first axis of
- the first output). `template` gives compile-time constants, e.g.
- `template={"NSTEPS": 200}`. It returns the output array, or a tuple of them.
-- Outputs are uninitialized: write every element, pass the old array as an input
- too if the kernel reads it, or set `init_value`.
-- **float32 only.** The Apple GPU has no float64: a float64 input or output
- raises `TypeError`. With `float64="cast"` float64 data is computed in float32
- (a push over 200 steps agreed with float64 to about 3e-5).
-- Every call copies the inputs to MLX arrays and the results back, counted as
- `to_device` and `to_host` transfers by `count_transfers()`. On a 4M-particle
- push these copies were about 7 ms next to an 11.5 ms kernel.
-- `xp.kernels.metal_available()` is True if MLX is installed and a Metal GPU is
- present. Without them, calling a `MetalKernel` raises `ImportError` or
- `RuntimeError`.
-
-## `kernels.CudaKernel`
-
-### Constructor
-
-```python
-xp.kernels.CudaKernel(
- source,
- name,
- *,
- block_size=128,
- options=(),
- include_dirs=(),
- source_dir=None,
- structs=(),
- template_args=None,
- check_signature=True,
- debug=None,
- n_threads_from="auto",
- check_finite=False,
-)
-xp.kernels.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs)
-xp.kernels.CudaKernel.all_from_file(path, **kwargs)
-```
-
-Wraps the `__global__` function `name` in the CUDA C `source` (declared
-`extern "C"`, unless it is a template). The kernel is compiled with NVRTC
-through CuPy on the first call (or by `compile()`), and cached, also on disk by
-CuPy. CuPy is imported only then, so kernels can be created and their
-signatures parsed without CuPy; `compile()` raises `RuntimeError` without a
-GPU. `compile(log_stream=None)` compiles eagerly and returns the CuPy raw kernel;
-repeated calls reuse successful state on the current device. `is_compiled`
-reports that device's state. `recompile(log_stream=None)` refreshes headers and
-options on the current device; finish in-flight launches before rebuilding.
-Compiler failures remain retryable and a writable `log_stream` receives compiler
-output. Catalog compilation therefore reports errors during setup.
-
-`from_file` reads the source from a file; the kernel name defaults to the file
-name without `suffix` (`axpy_cuda.cu` -> `axpy`), and the directory of the file
-is added to the include directories and is the `source_dir`.
-
-`all_from_file` loads every `__global__` function of a file, for files that
-group several small kernels, and returns a `dict` of kernels by name in the
-order of the source. The kernels share the source and options, so CuPy
-compiles the file once. `xp.cuda.cuda_kernel_names(source)` lists the `__global__`
-functions of a source string (ignoring comments).
-
-```python
-kernels = xp.kernels.CudaKernel.all_from_file("small_kernels.cu", block_size=64)
-kernels["scale"](x, 2.0, x.size, n_threads=x.size)
-kernels["shift"](x, 1.0, x.size, n_threads=x.size)
-```
-
-### Parameters
-
-* `block_size`: threads per block, an integer for 1D launches or a tuple of 1
- to 3 integers, e.g. `(16, 16)`; at most 1024 threads in total.
-* `options`: additional NVRTC options, e.g. `("-std=c++17",)`.
-* `include_dirs`: directories for `#include`, passed as `-I`. The headers
- shipped with cunumpy (see "CUDA headers and array views" below) are always
- found.
-* `source_dir`: the directory the source was read from, where
- `#include "..."` files are looked up first (set by `from_file`).
-* `structs`: `CudaStruct` types that the kernel takes as parameters (by
- value), see `CudaStruct` below.
-* `template_args`: template arguments if `name` is a function template, see
- "Templates and generated kernels" below.
-* `check_signature`: parse the signature and check every call against it
- (default). Raises `ValueError` if the signature cannot be parsed, e.g. with
- macros or pointers to pointers in the parameter list; pass `False` to launch
- with the arguments as they are, like `cupy.RawKernel`.
-* `debug`: `None` (default) follows the global debug setting, `True`/`False`
- fix it for this kernel, see "Debugging" below.
-
-Properties: `name`, `expression` (`name`, or the template instantiation such
-as `"scale"`), `source`, `block_size`, `options`, `include_dirs`,
-`source_dir`, `included_headers`, `structs`, `template_args`, `signature`,
-`is_compiled`, `debug`.
-
-### Included headers and the compile cache
-
-CuPy caches compiled kernels on disk (`~/.cupy/kernel_cache`), keyed on the
-source string and the compiler options only: a file pulled in through
-`#include "..."` is not part of the key, so editing a shared `.cuh` header
-would not recompile the kernels that include it. `CudaKernel` therefore
-resolves the quoted includes of its source when it compiles and adds a define
-with a hash of their contents to the options:
-
-```python
-kernel = xp.kernels.CudaKernel.from_file("push/push_cuda.cu", include_dirs=[src_root])
-kernel.included_headers # (Path('push/helpers.cuh'), Path('.../common.cuh'))
-kernel.options # ('-Ipush', '-I')
-kernel.compile_options() # options + ('-DCUNUMPY_INCLUDE_HASH=0x3f9a...',)
-```
-
-* `included_headers`: the header files the source includes with
- `#include "name"`, recursively, each once in order of first inclusion. A
- name is looked up relative to the including file (`source_dir` for the
- kernel source, the header's own directory for nested includes), then in
- `include_dirs` in order, then in cunumpy's header directory, like NVRTC
- does. cunumpy's shipped headers are tracked also when included in angle
- brackets (`#include `), so upgrading cunumpy with a
- changed header recompiles the kernels that use it. Other angle-bracket
- (system) headers and includes that cannot be found are ignored (NVRTC
- reports the latter). Recomputed at every access, so it follows the files on
- disk.
-* `compile_options()`: the options passed to CuPy at compile time: `options`
- plus `-DCUNUMPY_INCLUDE_HASH=0x` if the source includes any header,
- where the hash covers the contents of `included_headers` (not their paths).
- A changed header gives another define, hence another cache entry. Sources
- without includes never touch the file system.
-
-The two building blocks are available on their own:
-
-* `xp.cuda.resolve_includes(source, include_dirs=(), *, base_dir=None,
- angle_dirs=())`: the resolved header paths of a source, as a list;
- `angle_dirs` are searched last and also for `#include `.
-* `xp.cuda.include_hash(paths)`: the first 16 hex digits of the SHA-256 digest of
- the contents of the files, in order.
-
-### Calling
-
-```python
-kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None)
-```
-
-Launches the kernel on `stream` (the current stream if `None`). The launch
-shape is given either by `n_threads` or by `grid`:
-
-* `n_threads`: number of threads, an integer or a tuple of 1 to 3 integers such
- as `(nx, ny)`. The grid is `ceil(n_threads / block)` per dimension. With a 1D
- `block_size` and multi-dimensional `n_threads`, the block is
- `(block_size, 1, ...)`.
-* `grid`: number of blocks per dimension, instead of `n_threads`.
-* `block`: block shape for this call, instead of `block_size`.
-* `shared_mem`: dynamic shared memory per block in bytes, for
- `extern __shared__` arrays. Static plus dynamic storage is checked against
- the device's opt-in limit. Dynamic storage above the default allowance after
- subtracting static storage opts in through `max_dynamic_shared_size_bytes`.
- Block/grid dimensions and device/kernel thread limits are also checked.
-
-Without `n_threads` and `grid`, the default `n_threads_from="auto"` infers the
-launch from the first array argument, skipping scalars and zero-dimensional
-arrays. Arrays inside `CudaArguments`, `CudaStructArguments`, and struct values
-are searched in argument/field order. The effective block shape determines
-how many leading array axes are used: a 1D block uses `shape[0]` (one thread per
-row/particle), a 2D block uses `shape[:2]`, and a 3D block uses `shape[:3]`. These
-axes map to CUDA x, y, z. A `block=` override changes the inference dimensions.
-Missing arrays or insufficient array dimensions raise with instructions to
-give an explicit size.
-
-```python
-push = xp.kernels.CudaKernel.from_file("push_cuda.cu")
-push(positions, velocities, e_field, dt) # n_threads = positions.shape[0]
-```
-
-Explicit `n_threads` or `grid` always takes precedence. Set `n_threads_from`
-(constructor argument or settable property) to a callable such as
-`lambda args: args[0].size` for flattened element kernels, or
-`lambda args: args[0].shape[::-1]` for kernels whose x index follows columns.
-`"first_array"` always uses the first axis and `"last_axis"` the last one (one
-thread per entry of a component-major `(ncomp, N)` or `(N,)` array); None disables inference and requires
-explicit launch sizes. The same defaults apply through `kernels.Kernel` and
-`kernel_testing.assert_kernels_agree`, and in CPU emulation.
-
-Nothing is launched if the grid has a zero dimension (e.g. `n_threads=0`).
-`kernel.launch_shape(n_threads=None, *, grid=None, block=None, args=None)` returns the
-`(grid, block)` a call would use, e.g. to size a per-block output:
-
-Supply `args=(...)` to inspect an automatically inferred launch without running
-the kernel, e.g. `kernel.launch_shape(args=(positions, velocities, e_field, dt))`.
-
-```python
-BLOCK_SUM = r"""
-extern "C" __global__ void block_sum(const double* x, double* out, int n) {
- extern __shared__ double buffer[];
- int i = blockDim.x * blockIdx.x + threadIdx.x;
- buffer[threadIdx.x] = i < n ? x[i] : 0.0;
- __syncthreads();
- for (int s = blockDim.x / 2; s > 0; s /= 2) {
- if (threadIdx.x < s) buffer[threadIdx.x] += buffer[threadIdx.x + s];
- __syncthreads();
- }
- if (threadIdx.x == 0) out[blockIdx.x] = buffer[0];
-}
-"""
-block_sum = xp.kernels.CudaKernel(BLOCK_SUM, "block_sum", block_size=128)
-(n_blocks,), _ = block_sum.launch_shape(x.size)
-partial = xp.zeros(n_blocks)
-block_sum(x, partial, x.size, n_threads=x.size, shared_mem=128 * 8)
-```
-
-In a 2D kernel, use `blockIdx.y`/`threadIdx.y` for the second dimension and
-launch with `n_threads=(nx, ny)` and, e.g., `block_size=(16, 16)`.
-
-### Argument checks
-
-The arguments are prepared by `kernel.prepare_args(*args)`:
-
-* arguments with a `__cuda_args__()` method are replaced by the values it
- returns (see `CudaArguments` and `CudaStruct` below);
-* with a checked signature, the number of arguments must match, and
- * pointer parameters take C-contiguous CuPy arrays whose dtype matches the
- pointed-to type (any dtype for `void*`); host arrays raise `TypeError`,
- they are never copied to the device, and so do non-contiguous views such
- as `a[:, 0:3]`, which the kernel would read as a flat buffer (build the
- arrays with `as_device_array()` or `cupy.ascontiguousarray()`);
- * struct parameters take values of that `CudaStruct`;
- * array view parameters (`Array2D`, see "CUDA headers and array
- views" below) take CuPy arrays of the declared dtype and number of
- dimensions, contiguous or not, and are packed into (pointer, shape,
- strides in elements);
- * Python scalars are cast to the declared type: `int` into integer (with a
- range check, `OverflowError`), floating-point and complex parameters,
- `float` into floating-point and complex parameters, `bool` into boolean
- and integer parameters; anything else raises `TypeError`;
- * NumPy scalars are passed as they are if their dtype matches, cast if the
- cast is safe (e.g. `np.float32` into `double`), and raise `TypeError`
- otherwise (e.g. `np.float64` into `float`). NumPy integer scalars are
- checked by value, like Python ints: `np.int64(5)` fits an `int`
- parameter, `np.int64(2**31)` raises `OverflowError`.
-
-This matters because `cupy.RawKernel` reads each argument with the size
-declared in the signature and does not check types: an integer passed to a
-`double` parameter, or a `double` passed to a `float` parameter, arrives as a
-wrong value without an error. The checks cost about 0.3 µs per argument (about
-10 µs for a kernel with 29 arguments, measured on an H100 node, where the launch
-itself costs about as much), which is negligible for kernels that run for
-100 µs or more. For very short kernels called in a hot loop, pass
-`check_signature=False` once the calls are known to be correct.
-
-C types are mapped to NumPy dtypes as on Linux (LP64): `int` is `int32`,
-`long` and `long long` are `int64`, `float` is `float32`, `double` is
-`float64`, `complex` is `complex128`; fixed-width types such as
-`int64_t` and `size_t` are supported too. `xp.cuda.ctype_of(dtype)` gives the C
-type of a dtype (`xp.cuda.ctype_of(np.float64) == "double"`), e.g. to generate
-source. `xp.cuda.parse_cuda_signature(source, name, *, structs=(),
-template_args=None)` returns the parsed parameters (`CudaParameter` tuples of
-`name`, `ctype`, `dtype`, `pointer`, `struct`, `view_ndim`).
-
-### CUDA headers and array views
-
-```python
-SCALE_COLUMN = r"""
-#include
-#include
-extern "C" __global__
-void scale_column(Array2D a, long long column, double factor) {
- CUNUMPY_THREAD_1D(i, a.shape[0]); // long long i; returns if i >= shape[0]
- a(i, column) *= factor;
-}
-"""
-scale_column = xp.kernels.CudaKernel(SCALE_COLUMN, "scale_column")
-view = markers[::2, 1:5] # non-contiguous is fine
-scale_column(view, 1, 10.0, n_threads=view.shape[0])
-```
-
-cunumpy ships CUDA headers that every `CudaKernel` finds automatically;
-`xp.cuda.cuda_include_dir()` returns their directory (a `str`) for other compilers
-(`-I`).
-
-`cunumpy/array_view.cuh` defines the strided views `Array1D` to
-`Array16D` (for example, a 4D view can describe a 3D grid of vector
-components `(nx, ny, nz, ncomp)`): `T* data`, `long long shape[ndim]`, `long long
-strides[ndim]` (in elements, not bytes), `operator()(i, j, ...)` returning a
-reference to the element, and `size()`. A kernel indexes `a(i, j)` like the
-pyccel kernel it is ported from indexes `a[i, j]`, without hand-passed sizes.
-Compiling with `options=("-DCUNUMPY_BOUNDS_CHECK",)` checks every index against
-the shape (an out-of-bounds index prints a message and traps the kernel).
-
-A kernel parameter or a `CudaStruct` field of type `ArrayD`, for the
-scalar C types above, takes a CuPy array of that dtype and number of dimensions
-(dtype and ndim mismatches raise `TypeError`), contiguous or not: it is packed
-by value into pointer, shape and strides with the memory layout of the C
-struct (8-byte aligned, `sizeof == 8 * (1 + 2 * ndim)`; the header checks this
-with `static_assert`). `CudaParameter.view_ndim` is the number of dimensions of
-such a parameter.
-
-`cunumpy/index.cuh` defines `CUNUMPY_THREAD_1D(i, n)` (declares `long long i`
-as the global thread index and returns if `i >= n`), `CUNUMPY_THREAD_2D(i, j,
-ni, nj)`, `CUNUMPY_THREAD_3D(i, j, k, ni, nj, nk)` and the grid-stride loop
-`CUNUMPY_GRID_STRIDE_1D(i, n) { ... }`.
-
-### Templates and generated kernels
-
-A function template is instantiated with `template_args`: C types (or NumPy
-dtypes, converted with `ctype_of`) for type parameters, integers or booleans
-for non-type parameters. The template parameters are substituted into the
-signature, so calls are checked as for any other kernel:
-
-```python
-SCALE = r"""
-template
-__global__ void scale(T* x, T factor, int n) {
- int i = blockDim.x * blockIdx.x + threadIdx.x;
- if (i < n) x[i] = factor * x[i] * (T)N;
-}
-"""
-scale_f64 = xp.kernels.CudaKernel(SCALE, "scale", template_args=(np.float64, 3))
-scale_f64(x, 2.0, x.size, n_threads=x.size) # instantiation scale
-```
-
-For kernels whose source is generated per variant (e.g. per number of
-dimensions and dtype), `CudaKernelVariants` creates and caches one kernel per
-key:
-
-```python
-matvec = xp.kernels.CudaKernelVariants(
- lambda ndim, dtype: xp.kernels.CudaKernel(
- make_source(ndim, xp.cuda.ctype_of(dtype)), "matvec"
- )
-)
-matvec.get(3, np.float64)(mat, x, out, n_threads=out.size) # created once
-matvec.compile_all([(3, np.float64), (3, np.complex128)]) # at setup
-```
-
-`get(*key)` calls the factory the first time a key is used; `keys()`,
-iteration and `len()` give the variants created so far; `compile_all(keys=(), jobs=1)`
-creates the given variants and compiles all of them, `jobs` at a time in
-threads (see `KernelCatalog.compile_all`).
-
-### Debugging
-
-```python
-xp.cuda.set_cuda_debug(enabled)
-xp.cuda.get_cuda_debug()
-xp.cuda.cuda_debug(enabled=True) # context manager
-xp.kernels.CudaKernel(..., debug=None)
-kernel.debug_active()
-kernel.compile_options()
-xp.cuda.DEBUG_OPTIONS # ("-lineinfo", "-DCUNUMPY_BOUNDS_CHECK")
-```
-
-Kernel launches are asynchronous: a CUDA error such as an illegal memory
-access or a launch failure is reported by the next operation that
-synchronizes (a `.get()`, an MPI call, ...), which may be far from the kernel
-that caused it. In debug mode, a `CudaKernel`
-
-* is compiled with `-lineinfo` (source line information for
- `compute-sanitizer` and profilers) and `-DCUNUMPY_BOUNDS_CHECK` (bounds
- checks in cunumpy's array views, and available to your own `#ifdef`s),
- unless the option is already among its `options`. `-G` (device debug
- symbols) is not added, because NVRTC does not support it;
-* synchronizes the stream after every launch (the `stream` passed, else the
- current one), so an error is raised at the launch that caused it, as a
- `RuntimeError` that names the kernel and its grid and block, with the CuPy
- error chained.
-
-Debug mode is enabled globally with `xp.cuda.set_cuda_debug(True)`, temporarily
-with the context manager `xp.cuda.cuda_debug()`, or before starting Python with
-the environment variable `CUNUMPY_CUDA_DEBUG=1` (`true`, `yes` and `on` work
-too); `xp.cuda.get_cuda_debug()` returns the current setting. A kernel created with
-`debug=None` (the default) reads the global setting at every launch, so
-enabling it also affects kernels created earlier; `debug=True` or
-`debug=False` fix the mode for one kernel. Only the compile options are fixed
-at compile time: a kernel compiled before debug mode was enabled keeps its
-options, so call `recompile()` after enabling, or create the kernels after
-enabling. `kernel.debug_active()` tells whether debug mode applies to a
-kernel now, and `kernel.compile_options()` returns the options a compilation
-now would use.
-
-```python
-with xp.cuda.cuda_debug():
- kernel = xp.kernels.CudaKernel(SOURCE, "kernel")
- kernel(
- x, y, n, n_threads=n
- ) # RuntimeError: CUDA error after launching kernel 'kernel' ...
-```
-
-The `RuntimeError` says which kernel failed, not where. The next step is
-NVIDIA's memory checker, which reports the faulting source line (thanks to
-`-lineinfo`) and also finds out-of-bounds accesses that do not crash:
-
-```bash
-CUNUMPY_CUDA_DEBUG=1 compute-sanitizer python -m pytest tests/unit/test_my_kernel.py
-```
-
-Note that after an illegal memory access the CUDA context is unusable; the
-process (or the pytest run) has to be restarted.
-
-## `arguments.CudaStruct`
-
-```python
-Particles = xp.arguments.CudaStruct(
- "Particles",
- [("x", "double*"), ("v", "double*"), ("n", "int"), ("charge", "double")],
-)
-source = (
- Particles.declaration
- + r"""
-extern "C" __global__ void push(Particles p, double dt) {
- int i = blockDim.x * blockIdx.x + threadIdx.x;
- if (i < p.n) p.x[i] += dt * p.charge * p.v[i];
-}
-"""
-)
-push = xp.kernels.CudaKernel(source, "push", structs=[Particles])
-push(Particles(x=x, v=v, n=x.size, charge=-1.0), 0.1, n_threads=x.size)
-```
-
-A C struct passed to kernels by value. It groups arguments, e.g. all arrays
-describing a set of particles, into one kernel parameter, so adding a field
-changes one definition instead of every kernel signature.
-
-`CudaStruct(name, fields)` takes the fields as `(name, C type)` pairs; scalar
-fields, pointers to the scalar types above (or `void*`), and array views
-`Array1D` to `Array16D` of those scalar types (see "CUDA headers and
-array views") are supported.
-
-* `declaration`: the C definition of the struct, to put in the CUDA source
- or a header. A struct with array view fields (`has_views`) needs
- `#include "cunumpy/array_view.cuh"` before it; `to_header()` adds it.
-* `dtype`: the NumPy structured dtype with the memory layout of the C struct
- (C alignment and padding; pointers stored as 64-bit device addresses).
-* `fields`: the parsed fields (`CudaParameter` tuples).
-* `check_source(source)`: raises `ValueError` if `source` defines the struct
- with other fields; a kernel created with `structs=[...]` does this check.
-* `verify_layout(include=None, *, include_dirs=(), options=())` (needs CuPy):
- compiles and runs a one-thread kernel that reports `sizeof`, `alignof` and
- every field offset as the CUDA compiler lays the struct out, and raises
- `ValueError` listing the differences from `dtype`. Returns the measured
- layout as a dict. With `include` (a header file name or `#include` line,
- found in `include_dirs`) the struct is defined by that header instead of
- `declaration`, which checks a hand-written or generated header. Call it once
- per struct in a GPU test, and on every new platform (e.g. ROCm).
- `layout_source(include=None)` returns the kernel source.
-* Calling the struct with keyword arguments, one per field, packs the values:
- pointer fields take C-contiguous CuPy arrays of the declared dtype (never
- copied), array view fields take CuPy arrays of the declared dtype and number
- of dimensions (contiguous or not), scalar fields are checked and cast like
- scalar kernel arguments.
-
-The result is a `CudaStructValue`: it keeps references to the arrays it points
-to (the packed struct only holds their addresses, so keep the value alive while
-the kernel may run), gives access to the field values with
-`value["field"]`, holds the packed struct in `value.packed`, and is flattened
-into it when passed to a kernel.
-
-### Structs from Python annotations
-
-```python
-class MarkerArguments: # the pyccel argument class, e.g. in struphy
- def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ...
-
-
-MarkerArgs = xp.arguments.CudaStruct.from_signature(
- MarkerArguments.__init__, "MarkerArgs"
-)
-print(MarkerArgs.declaration)
-# struct MarkerArgs {
-# Array2D markers;
-# long long n_markers;
-# Array1D valid;
-# };
-```
-
-`CudaStruct.from_signature(func, name, *, int_type="long long",
-scalar_names=None)` builds the struct from the annotated parameters of `func`
-(one field per parameter, in order; `self` is skipped), so that the Python
-class is the one definition of the arguments on the host and on the device.
-Annotations are written in the pyccel style, as strings or real types:
-`"float"`/`float` -> `double`, `"int"`/`int` -> `int_type` (`"long long"` by
-default, since pyccel integers are 64-bit), `"bool"`/`bool` -> `bool`, NumPy
-scalar types such as `np.float32` -> `float`, and an array `"float[:, :]"` ->
-`Array2D` (1 to 3 dimensions; `Final[...]` and `const` are ignored).
-`scalar_names` adds or changes mappings from annotation scalar names to C
-types, e.g. `{"float": "float"}` for single precision. A parameter without
-annotation, or with an annotation that cannot be mapped, raises `ValueError`.
-
-### Structs from a pyccel source file
-
-`CudaStruct.from_pyccel_class(source, class_name, name=None, *, int_type="long long",
-scalar_names=None, exclude=(), attribute_names=True)` builds a struct from the
-annotated `__init__` of a class in a `.py` file (or a source string), parsed
-with `ast` and never imported, for classes whose module is compiled by pyccel.
-Fields are named after the attributes the parameters are stored in
-(`self. = `; `attribute_names=False` keeps the parameter
-names), parameters in `exclude` are skipped, and the mappings are those of
-`from_signature`. Raises `ValueError` if the class or its `__init__` is missing
-or an annotation cannot be mapped.
-
-### Generating headers
-
-```python
-xp.arguments.write_cuda_header("pusher_args.cuh", [MarkerArgs, DomainArgs])
-```
-
-`struct.to_header(path=None, *, guard=None, includes=())` returns the struct
-definition as a header: an include guard (`_CUH` by default),
-`#include "cunumpy/array_view.cuh"` if the struct has array view fields, the
-`includes` (file names or `#include` lines), and the definition. With `path`
-the header is also written. `xp.arguments.write_cuda_header(path, structs, guard=None,
-*, includes=())` writes several structs to one header (the guard defaults to
-the file name, `pusher_args.cuh` -> `PUSHER_ARGS_CUH`) and returns the source.
-
-The pattern: write the header once (at setup, or in a script), commit it next
-to the kernels that `#include` it, and keep it in sync with a test:
-
-```python
-def test_pusher_args_header_is_up_to_date():
- generated = xp.arguments.write_cuda_header(
- tmp_path / "pusher_args.cuh", [MarkerArgs, DomainArgs]
- )
- assert Path("kernels/pusher_args.cuh").read_text() == generated
-```
-
-Kernels created with `structs=[MarkerArgs, ...]` also check a definition in
-their own source against the Python definition (`check_source`).
-
-## `arguments.CudaStructArguments`
-
-```python
-class MarkerArguments(xp.arguments.CudaStructArguments):
- struct_name = "MarkerArgs"
- fields = (("markers", "Array2D"), ("valid", "bool*"), ("n_markers", "int"))
-
- def __init__(self, markers, valid):
- self.markers = markers
- self.valid = valid
- self.n_markers = markers.shape[0]
- self.pack()
-
-
-push = xp.kernels.CudaKernel(source, "push", structs=[MarkerArguments.struct])
-push(MarkerArguments(markers, valid), dt, n_threads=markers.shape[0])
-```
-
-Base class for argument objects that are one C struct: the class form of
-`CudaStruct`. A subclass sets `struct_name` and `fields` (`(field name, C type)`
-pairs, as for `CudaStruct`), stores every field as an attribute of the same
-name, and calls `pack()` at the end of its constructor. The `CudaStruct` is
-built once per subclass when the class is defined and is the class attribute
-`struct` (for `structs=[...]`, `declaration`, `to_header()`,
-`write_cuda_header()`); invalid field types raise when the class is defined.
-
-* `pack()` packs the field attributes, with the checks of `CudaStruct`
- (C-contiguous CuPy arrays of the declared dtype, range-checked scalars). A
- field without an attribute raises `AttributeError`.
-* The packed struct always matches the current attributes: at every use
- (`packed`, `__cuda_args__()`, so at every launch) the device address, shape
- and strides of each array field and the value of each scalar field are
- compared with what was packed, and the struct is packed again if anything
- changed. Fields may be properties that read an owner's current arrays, so a
- resized array is picked up at the next launch. Calling `pack()` again is
- never needed.
-* `packed` is the packed struct (`numpy.void`); `__cuda_args__()` returns
- `(packed,)`, so a `CudaKernel` receives the struct.
-* Copies (`copy.copy`, `copy.deepcopy`) and unpickled objects are packed again
- from their own arrays; the packed struct is not part of the pickled state.
-* A subclass that sets neither `struct_name` nor `fields` is an intermediate
- base class (its instances cannot be packed); setting only one raises
- `TypeError`. Subclasses of a complete class inherit its struct.
-
-## `arguments.CudaArguments`
-
-```python
-class Particles(xp.arguments.CudaArguments):
- def __init__(self, positions, velocities):
- self.positions = positions
- super().__init__(positions, velocities, positions.shape[0])
-
-
-kernel(dt, Particles(x, v), n_threads=x.shape[0])
-```
-
-Base class for objects passed to a `CudaKernel` as one argument that stands
-for several kernel parameters. `CudaArguments(*values)` stores the values;
-`__cuda_args__()` returns them. Subclassing is optional: any object with a
-`__cuda_args__()` method returning a tuple is flattened. Host kernels
-receive their own argument objects (e.g. Pyccel classes holding NumPy arrays)
-unchanged: the caller passes the host or the CUDA object, cunumpy never
-converts one into the other. A `CudaArguments` object may also return struct
-values (`CudaStructValue.packed`) among its values.
-
-## `kernels.Kernel`
-
-```python
-xp.kernels.Kernel(
- host_kernel,
- cuda_kernel=None,
- *,
- name=None,
- missing_cuda="raise",
- cuda_path=None,
- host_options=None,
- dispatch="backend",
-)
-```
-
-A host kernel (a `PyccelKernel`; other callables are wrapped in one) and its
-CUDA counterpart. `kernel.get_kernel()` returns the host kernel on the NumPy
-backend and the CUDA kernel on the CuPy backend; call it once at setup to fail
-early if a CUDA kernel is missing.
-
-```python
-kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None)
-```
-
-calls the kernel of the active backend. The launch arguments are passed to the
-CUDA kernel and ignored by the host kernel. Omitted sizes use the CUDA kernel's
-shape-based default or its configured callback. Arguments are passed to the
-selected kernel as they are (a `CudaKernel` flattens `__cuda_args__()`
-objects). `kernel.compile()` compiles the CUDA kernel now and returns whether
-there is one.
-
-Without a CUDA kernel on the CuPy backend, `missing_cuda="raise"` raises
-`NotImplementedError` (naming `cuda_path`, if given), and
-`missing_cuda="fallback"` calls the host kernel through `PyccelKernel`, which
-copies the arrays to the host and back at every call (a `RuntimeWarning` is
-emitted once).
-
-`host_options` are keyword arguments for the `PyccelKernel` that wraps a plain
-callable `host_kernel`, e.g. `{"object_modules": ("my_package.",), "outputs":
-(2,)}`. They matter for the fallback: `object_modules` lets it find the device
-arrays inside application objects, and `outputs` limits the copies back to the
-device. Passing `host_options` together with a `PyccelKernel` raises
-`ValueError`; configure that `PyccelKernel` directly.
-
-`dispatch` decides which kernel a call runs. `"backend"` (default): the CUDA
-kernel on the CuPy backend, the host kernel on the NumPy backend.
-`"arrays"`: the CUDA kernel if any top-level argument lives on the GPU (a CuPy
-array, or a CUDA argument object: one with `__cuda_args__()`, such as a
-`CudaArguments` or a struct value), else the host
-kernel, whatever the backend. Use `"arrays"` in codes that hand host arrays to
-kernels while CuPy is active (diagnostics, MPI staging, CPU fallbacks): those
-calls then run the host kernel instead of failing in the CUDA argument checks,
-calling the host function directly (no device-array conversion).
-`missing_cuda` applies to device arguments without a CUDA kernel.
-
-`kernel.check_signature()` checks that the host and CUDA kernels (and the
-fallback of a `CompiledHostKernel`) take the same parameters in the same order (the names of the Python host function, or of the
-uncompiled Python version of a `CompiledHostKernel`, against the parsed
-`__global__` signature) and raises `ValueError` showing both lists otherwise.
-It does nothing without a CUDA kernel, with `check_signature=False`, or when
-the host kernel has no Python signature (a compiled function).
-`kernel.host_parameters()` returns the host names, or None.
-
-Properties: `name`, `host_kernel`, `cuda_kernel`, `has_cuda`, `missing_cuda`,
-`cuda_path`, `dispatch`.
-
-### Test arguments and compiled host kernels
-
-* `Kernel(..., test_args="pkg.push.push_test_args")`,
- `Kernel.test_args_module` and `Kernel.test_args` (the module, imported on
- first access): the test-arguments module of the kernel, set by
- `KernelCatalog.from_package()`; see `check_parity` below.
-* `Kernel.host_parameters()` falls back to the `__pyccel__/.pyi` stub
- for a pyccel-compiled host function, so `check_signature()` works for
- compiled kernels.
-* `Kernel.__call__` infers thread counts by default; custom `n_threads_from`
- callbacks are supported, and None requires explicit sizes.
-
-## `kernels.KernelCatalog`
-
-```python
-catalog = xp.kernels.KernelCatalog.from_package(
- package,
- *,
- host_suffix="_kernels",
- cuda_suffix="_cuda.cu",
- missing_cuda="raise",
- host_options=None,
- include_dirs=None,
- dispatch="backend",
- compile_host=None,
- host_fallback=None,
- **cuda_options,
-)
-kernel = catalog["push"]
-```
-
-A read-only mapping from names to `Kernel` objects. `from_package` scans the
-subfolders of `package`: for every folder `` containing the module
-`.py`, the function `` of that module is the host
-kernel, and `` in the same folder, if present, is the CUDA
-kernel (`__global__` function ``). Other `__global__` functions in that
-file are ignored by the catalog; they can be loaded with
-`CudaKernel.all_from_file`. Typically called in the package's `__init__.py`:
-
-```text
-my_kernels/
-├── __init__.py # catalog = xp.kernels.KernelCatalog.from_package(__name__)
-├── push/
-│ ├── push_kernels.py # def push(...): ...
-│ └── push_cuda.cu # __global__ void push(...)
-└── deposit/
- └── deposit_kernels.py # no CUDA kernel yet
-```
-
-* `host_options`: `PyccelKernel` options for the host kernels (see `Kernel`),
- the same for all kernels or a function of the kernel name, e.g.
- `lambda name: {"outputs": OUTPUTS[name]}`.
-* `include_dirs`: include directories of the CUDA kernels, in addition to
- each kernel's own folder. By default the source root of the top-level
- package (the directory containing it), so that a kernel of
- `my_pkg.kernels` can `#include "my_pkg/common.cuh"`. Headers found this
- way take part in the compile cache key, see "Included headers and the
- compile cache" under `CudaKernel`.
-* `dispatch`: passed on to every `Kernel` (`"backend"` or `"arrays"`).
-* `compile_host`: your function that compiles a host kernel module (cunumpy
- does not compile anything itself), e.g. a wrapper around `pyccel.epyccel`
- with a cache. Each host kernel is then a `CompiledHostKernel` (below),
- compiled on its first call. Without it, the plain Python functions are
- called.
-* `host_fallback`: for kernels whose compilation fails, a callable with the
- same arguments (e.g. a vectorized NumPy version), given as a mapping from
- names or a function of the name. Without one, a failed compilation runs the
- uncompiled Python function, with a warning.
-* `cuda_options`: passed on to `CudaKernel.from_file`, e.g. `block_size` or
- `structs`.
-
-A Pyccel package with the layout `/_pyccel.py` and
-`/_cuda.cu`, compiled host kernels and dispatch by argument:
-
-```python
-catalog = xp.kernels.KernelCatalog.from_package(
- __name__,
- host_suffix="_pyccel",
- dispatch="arrays",
- compile_host=my_pkg.compile_kernels, # e.g. pyccel.epyccel with a cache
- host_fallback=NUMPY_VERSIONS, # {"gather": gather_numpy, ...}
-)
-```
-
-`catalog.check_signatures()` runs `check_signature()` on every kernel and
-raises one `ValueError` listing every kernel whose host and CUDA parameters
-differ; call it in a unit test of a ported package.
-
-`catalog.without_cuda` lists the kernels still to port, `catalog.with_cuda`
-the ported ones. `catalog.summary()` (also `str(catalog)`) is one line on the
-porting status, e.g. for a `--status` command:
-`"CUDA kernels: 3 of 60 (missing: a, b, c)"`; at most `max_missing=10` names
-are listed before `...`.
-
-`catalog.compile_all(jobs=1)` compiles every CUDA kernel and returns their
-names; call it at setup so that the first time step does not pay for
-compilation (after the first run, CuPy loads the kernels from its disk cache).
-With `jobs > 1` the kernels are compiled in that many threads (NVRTC releases
-the GIL; all threads use the current device), `jobs=None` uses the number of
-CPUs. All kernels are compiled even if one fails; the first error is raised
-afterwards.
-
-`catalog.parity_cases()` returns the `(name, kernel)` pairs of the kernels
-that have a CUDA kernel, for a parametrised parity test (see "Testing
-utilities"). `KernelCatalog(kernels)` and `catalog.register(kernel, name=None)`
-build a catalog by hand.
-
-## `kernels.Kernel.from_folder`
-
-```python
-# my_sim/kernels/push/__init__.py
-kernel = xp.kernels.Kernel.from_folder(
- __name__, host_suffix="_pyccel", dispatch="arrays", compile_host=compile_kernels
-)
-kernel.implementations # ("pyccel", "numpy", "python", "cuda")
-kernel.selected() # "pyccel": what a call with host arrays runs now
-```
-
-The kernel of one kernel folder `package` (its dotted name, `__name__` in its
-`__init__.py`). Each version of `` in the folder is an implementation:
-`.py` gives `"pyccel"` (compiled with `compile_host` on
-first use, as it is without one) and `"python"` (uncompiled), `_numba.py`
-gives `"numba"`, `_numpy.py` gives `"numpy"`, and `` the
-CUDA kernel. `extra_implementations` adds host implementations that are not
-files, as loaders by name (`{"numpy": lambda: push_numpy}`). Takes the options of
-`KernelCatalog.from_package` for one kernel: `host_suffix`, `cuda_suffix`,
-`test_args_suffix`, `check_name_length`, `missing_cuda`, `host_options`,
-`include_dirs`, `dispatch`, `compile_host` and CUDA options such as
-`block_size` or `n_threads_from`. Raises `FileNotFoundError` if the folder has
-no host kernel module and `ModuleNotFoundError` if `package` is not a package.
-`kernel.selected(device=True)` names the implementation for device arguments.
-
-## `kernels.HostImplementations`, `kernels.set_host_kernel_implementation`
-
-```python
-host = xp.kernels.HostImplementations(
- "push",
- {"pyccel": load_compiled, "numpy": lambda: push_numpy, "python": lambda: push},
-)
-host(*args) # the default implementation
-xp.kernels.set_host_kernel_implementation("numpy") # every kernel: like xp.set_backend
-with xp.kernels.use_host_kernel_implementation("python"): # like xp.use_backend
- host(*args)
-```
-
-The host implementations of one kernel (names from `xp.kernels.HOST_IMPLEMENTATIONS`:
-`"pyccel"`, `"numba"`, `"numpy"`, `"python"`; `"python"` is required), each
-given as a loader that returns the function or raises if it is unavailable.
-Loaded on first use; `available(name)` loads and reports, `get(name)` returns it
-or raises `LookupError` (missing, or failed to load with the error as cause),
-`errors` maps names to load errors, `names` lists them, `python` is the
-uncompiled function, `build()` loads the default now. A call runs
-`selected()`: the implementation set with `set_host_kernel_implementation(name)` (or
-`use_host_kernel_implementation`, or the environment variable
-`CUNUMPY_HOST_KERNEL_IMPLEMENTATION` read at import), which raises if the kernel
-lacks it or cannot load it, else the default: the first available of pyccel,
-numba and NumPy, and else `"python"` with a `RuntimeWarning` (once).
-`get_host_kernel_implementation()` reads the setting; `None` is the default. The
-setting is global, not per thread, and applies to host calls only.
-
-## Device kernel implementation selection
-
-```python
-xp.kernels.set_device_kernel_implementation("cuda")
-xp.kernels.get_device_kernel_implementation() # "cuda"
-with xp.kernels.use_device_kernel_implementation(None):
- push(positions, velocities, dt) # automatic selection
-xp.kernels.set_device_kernel_implementation(None)
-```
-
-`xp.kernels.DEVICE_IMPLEMENTATIONS` is currently `("cuda",)`. The setter accepts
-`"cuda"` or `None` (automatic selection); unsupported values raise `ValueError`
-without changing the setting. The getter reports the requested setting, so
-it returns `None` in automatic mode even when a call would use CUDA.
-`CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION` initializes the setting at import;
-an unset or empty value means automatic selection. Environment values are
-case-insensitive and stripped of whitespace; unsupported values fail at import.
-
-Explicit `"cuda"` requires a CUDA implementation for device dispatch: missing
-implementations raise `LookupError`, even with `missing_cuda="fallback"`.
-Automatic mode preserves that per-kernel fallback policy. This affects
-`Kernel` and `KernelCatalog` device dispatch; it does not switch the array
-backend, alter host calls, or affect direct `CudaKernel`/CuPy RawKernel calls.
-The context manager restores the previous setting even after an exception.
-The setting is global, not per thread.
-
-## `kernels.CompiledHostKernel`
-
-```python
-kernel = xp.kernels.CompiledHostKernel(
- my_kernels_module, "push", compiler, fallback=push_numpy
-)
-kernel(*args)
-```
-
-A host kernel compiled on its first call, for `KernelCatalog.from_package(...,
-compile_host=...)` or by hand. cunumpy does not compile anything itself:
-`compiler(module)` is your function returning the compiled form of the module
-(with the same function names), e.g. a wrapper around `pyccel.epyccel` with an
-on-disk cache, or one that imports modules compiled ahead of time with the
-`pyccel` command. If compilation fails, the kernel calls `fallback`, or else
-the uncompiled Python function with a `RuntimeWarning`. `kernel.compiled`
-builds and reports whether that worked (so that callers can choose another
-path), `kernel.error` is the exception of a failed build, `kernel.python` the
-uncompiled function, `kernel.fallback` the fallback and `kernel.build()`
-compiles now.
-
-## Testing utilities
-
-```python
-from cunumpy.kernel_testing import (
- BACKENDS,
- assert_kernels_agree,
- device_function_kernel,
- emulate_cuda_kernel,
- emulated_launches,
- fake_cupy_session,
- requires_cupy,
- requires_device_backend,
- run_in_fake_cupy_subprocess,
-)
-```
-
-`cunumpy.kernel_testing` holds helpers for testing kernels with pytest. It is
-not imported by `import cunumpy`, and it imports pytest only when one of its
-pytest objects is used, so `device_function_kernel` works without pytest.
-It is not called `cunumpy.testing`, because a submodule of that name would
-replace NumPy's `xp.testing` once imported.
-
-### `requires_cupy`, `BACKENDS`, `backend`
-
-`requires_cupy` is `pytest.mark.skipif(not cupy_available(), reason="CuPy/GPU
-not available")`, for tests that need a GPU. `BACKENDS` is
-`["numpy", pytest.param("cupy", marks=requires_cupy)]`, so a test parametrised
-with it runs on NumPy everywhere and on CuPy where a GPU is available:
-
-```python
-@pytest.mark.parametrize("backend", BACKENDS)
-def test_norm(backend):
- with xp.use_backend(backend):
- assert xp.linalg.norm(xp.ones(4)) == 2.0
-```
-
-The `backend` fixture does the same and activates the backend for the test;
-import it into a `conftest.py` (`from cunumpy.kernel_testing import backend`) or the
-test module, then take `backend` as a test argument.
-
-### `requires_device_backend`, `device_backend_available()`
-
-`device_backend_available()` tells whether a CuPy-backend program can run: a
-GPU, or the fake CuPy. That is more than `requires_cupy` allows (CUDA kernels
-can be launched): on the fake CuPy, launches only run inside
-`emulated_launches()` (or `fake_cupy_session()`). `requires_device_backend`
-skips tests without either; with `CUNUMPY_REQUIRE_CUDA` set it fails them
-instead.
-
-### `assert_kernels_agree(kernel, make_args, ...)`
-
-```python
-assert_kernels_agree(
- kernel,
- make_args,
- *,
- n_threads=None,
- grid=None,
- block=None,
- rtol=1e-12,
- atol=0.0,
- n_calls=1,
- outputs=None,
- seed=0,
-)
-```
-
-Checks that the host and the CUDA version of a `Kernel` compute the same. For
-each backend, `"numpy"` then `"cupy"`, the backend is activated with
-`use_backend`, the positional arguments are built with `make_args(backend,
-seed)` (a tuple or list; kernels take positional arguments only), the kernel is
-called `n_calls` times (with `n_threads`, `grid` and `block` on CuPy), and the
-arrays among the arguments are collected. The CUDA results are copied to the
-host and compared with the host results using `numpy.testing.assert_allclose`
-with `rtol` and `atol`; the `AssertionError` names the argument that differs.
-The test is skipped (`pytest.skip`) without a GPU, and `ValueError` is raised
-for a kernel without CUDA version. The host arrays are returned by argument
-name (`"argument 0"`, `"argument 3.x"`) for further checks.
-
-`make_args` runs with the backend active, so arrays created through `cunumpy`
-land on it. NumPy and CuPy generators do not produce the same random sequence
-from one seed, so build random data on the host and convert it:
-
-```python
-def make_args(backend, seed):
- x = xp.to_cunumpy(np.random.default_rng(seed).random(1000))
- return (x, 2.0, x.size)
-
-
-def test_scale():
- assert_kernels_agree(catalog["scale"], make_args, n_threads=1000)
-```
-
-`outputs` selects the arguments to compare by index (negative indices count
-from the end), like `PyccelKernel(outputs=...)`; by default the `outputs`
-declared by the host kernel are used, and if it declares none, every argument.
-An argument that is an array is compared directly. For a tuple, list, dict or
-object argument (e.g. a `CudaArguments` object), the arrays it holds one level
-deep are compared, plus the arrays in a container attribute of an object.
-
-Together with `KernelCatalog.parity_cases()`, one test covers a catalog:
-
-```python
-MAKE_ARGS = {"scale": make_scale_args, "push": make_push_args}
-
-
-@pytest.mark.parametrize("name, kernel", catalog.parity_cases())
-def test_parity(name, kernel):
- assert_kernels_agree(kernel, MAKE_ARGS[name], n_threads=1000)
-```
-
-### `device_function_kernel(header_source, signature, ...)`
-
-```python
-device_function_kernel(
- header_source,
- signature,
- *,
- name=None,
- includes=(),
- n_threads_param="n",
- out_param="out",
- **cuda_kernel_options,
-)
-```
-
-Generates an elementwise `extern "C" __global__` kernel that calls a
-`__device__` function once per thread and returns it as a `CudaKernel`, so
-device helpers (B-spline evaluation, mapping evaluation, small linear algebra)
-can be run from Python on many inputs at once and compared with their host
-versions. `header_source` is the CUDA source defining the function (or the
-content of its header; `includes` adds `#include` lines before it, with quotes,
-or with angle brackets for `""`), and `signature` is its C
-prototype, e.g. `"int find_span(const double* t, int p, double eta)"`.
-Additional keyword arguments such as `include_dirs` and `block_size` go to
-`CudaKernel`.
-
-The generated kernel takes the parameters of the function in their order,
-followed by the output array and the number of elements:
-
-* a pointer parameter is kept as it is and passed unchanged to every call (an
- array shared by all threads);
-* a scalar parameter `T x` becomes a device array `const T* x` of length `n`,
- and thread `i` calls the function with `x[i]`;
-* the return value of thread `i` is stored in `out[i]` (`R* out`, with `R` the
- return type); a `void` function has no `out`;
-* `int n` is the number of elements; threads `i >= n` do nothing.
-
-The prototype above gives:
-
-```c
-extern "C" __global__ void find_span_kernel(
- const double* t, const int* p, const double* eta, int* out, int n)
-{
- int i = blockDim.x * blockIdx.x + threadIdx.x;
- if (i >= n) return;
- out[i] = find_span(t, p[i], eta[i]);
-}
-```
-
-```python
-find_span = device_function_kernel(
- BSPLINES_CUH, "int find_span(const double* t, int p, double eta)"
-)
-find_span(t, p, eta, spans, eta.size, n_threads=eta.size)
-```
-
-The kernel is named `_kernel` unless `name` is given. Scalar
-parameters and return types are those `CudaKernel` supports; a struct or
-pointer return type, an unsupported parameter type, or a parameter named like
-`out_param` or `n_threads_param` raises `ValueError` (rename the generated
-parameter in that case).
-
-### `parity_cases(catalog)`, `check_parity(kernel, **overrides)`
-
-`parity_cases(catalog)` returns one `pytest.param(kernel, id=name)` per kernel
-of `catalog.parity_cases()`; a kernel without a test-arguments module is
-marked `skip` with a reason naming the missing `_test_args.py`.
-`check_parity(kernel)` runs `assert_kernels_agree` with the module's
-`make_args` and the settings of `TEST_ARGS_SETTINGS` (`N_THREADS`, `GRID`,
-`BLOCK`, `RTOL`, `ATOL`, `N_CALLS`, `OUTPUTS`, `SEED`), overridden by keyword
-arguments. Raises `ValueError` without a module and `TypeError` without a
-callable `make_args`.
-
-### `install_fake_cupy()`, `fake_cupy_active()`
-
-Install the fake CuPy of `cunumpy._fake_cupy` (a strict host stand-in for CuPy
-for machines without a GPU; also installed by `CUNUMPY_FAKE_CUPY=1` when
-cunumpy is imported), and tell whether it is active. `requires_cupy` and
-`assert_kernels_agree` skip while it is. `install_fake_cupy()` raises if the
-real CuPy was imported already or cunumpy already checked for CuPy.
-
-### `emulate_cuda_kernel(kernel, *args, n_threads=None, grid=None, block=None, compiler=None, options=(), shared_mem=0)`
-
-```python
-x, y = rng.random(1000), np.zeros(1000)
-emulate_cuda_kernel(axpy, 2.0, x, y, 1000, n_threads=1000)
-np.testing.assert_allclose(y, 2.0 * x, rtol=1e-15)
-```
-
-Runs a `CudaKernel` on the CPU, serially, as if it were launched with `args`,
-so that CI without a GPU can compare a kernel with its host version. The
-kernel source is compiled as C++ (C++17, `CXX` or `c++`; see
-`emulation_compiler()`) with the CUDA built-ins replaced: `threadIdx`,
-`blockIdx`, `blockDim`, `gridDim`, atomics (`atomicAdd`, `atomicMin`, ...,
-plain operations), `__ldg`, `rsqrt`, `__trap`. The shipped headers and the
-kernel's include directories and `-D` options apply. Then the kernel is called
-once per thread, for every block and thread index of the launch shape.
-
-Each kernel is compiled once (for each compiler, options and template
-arguments) into a shared library that takes the arguments at run time, so all
-later launches, with any values and array sizes, reuse it. Libraries are cached
-in the process and on disk, like CuPy's kernel cache: in
-`emulation_cache_dir()`, which is `CUNUMPY_EMULATION_CACHE` (`0`: no disk
-cache), else `$XDG_CACHE_HOME/cunumpy/emulation` or `~/.cache/cunumpy/emulation`.
-The library runs in the Python process, on the arrays themselves.
-
-Arguments follow the signature, with NumPy arrays in place of CuPy arrays:
-pointer and view parameters (`Array1D` to `Array16D`) take arrays of the
-declared dtype (and ndim); the kernel writes into them directly (an array that
-is not writeable, or whose layout a parameter cannot take, is passed as a
-copy and written back), so any strides work and aliased arguments see each
-other's writes. Struct parameters take a mapping of field names to values, a
-`CudaStructValue`, or an object with an attribute per field. Scalars are
-checked and cast like in a launch. Like NVRTC by default, the compiler may fuse
-`a * b + c` into an FMA, so compare with NumPy using a tolerance of a few ulp,
-or pass `options=("-ffp-contract=off",)` for NumPy's rounding.
-
-Block shared memory and `__syncthreads` are emulated: `__shared__` variables
-are one copy per block (blocks run one after another), `extern __shared__`
-arrays point into a buffer of `shared_mem` bytes, and in a kernel that calls
-`__syncthreads` the threads of a block run as coroutines (POSIX `ucontext`,
-each with its own stack), so that every thread reaches a barrier before any
-thread continues past it. Per-block deposits, shared-memory reductions and
-tiled kernels work.
-
-Not emulated: concurrency between barriers (races and atomic ordering never
-show), warp intrinsics, complex scalars and inline PTX. Inline `asm(...)` and
-`asm volatile(...)` statements compile, but trap when reached, so a kernel with
-a PTX branch that a test never takes (e.g. `asm("trap;")` for an unknown case)
-runs without extra options. A kernel, or a header it includes, using
-`__syncwarp`, warp shuffles or votes raises `NotImplementedError` (serial
-threads would give wrong results). A kernel that does not compile, or traps (an
-out-of-bounds index with `-DCUNUMPY_BOUNDS_CHECK`, `__trap()`, inline `asm`),
-raises `RuntimeError`. Other crashes, such as a segmentation fault from an
-unchecked out-of-bounds index, end the process; run such code in a child
-process (`run_in_fake_cupy_subprocess()`, which prints the traceback with
-`faulthandler`).
-
-### `compile_for_emulation(kernel, *, compiler=None, options=())`, `emulation_cache_dir()`
-
-`compile_for_emulation()` builds (or fetches from the cache) the emulation
-library of a kernel without launching it, so compile errors show up early;
-it raises `RuntimeError` with the compiler output. A kernel without a parsed
-signature is only checked for syntax and type errors (`-fsyntax-only`).
-`emulation_cache_dir()` is the disk cache, or None.
-
-### `emulated_launches(*, compiler=None, options=())`
-
-```python
-with emulated_launches():
- propagator(dt) # CuPy backend = the fake CuPy: the kernels run on the CPU
-np.testing.assert_allclose(host_buffer(markers), expected)
-```
-
-Inside the block every `CudaKernel` launch runs through `emulate_cuda_kernel`,
-on the host buffers of the fake CuPy arrays (`host_buffer(array)` is the NumPy
-array behind one), with the `compiler` and `options` of the block. Argument
-objects are flattened like in a launch. No CUDA is compiled in the block
-either: `CudaKernel.compile()` (and so `recompile()`,
-`CudaKernelVariants.compile_all()` and `KernelCatalog.compile_all()`) calls
-`compile_for_emulation()` and returns None instead of a `cupy.RawKernel`, so a
-program that compiles its kernels up front still reports compile errors there.
-Kernels the emulation cannot run (warp intrinsics) are skipped by `compile()`.
-This holds on a GPU too: the block means "no CUDA here". The original methods
-are restored when the block exits, also by an exception.
-
-### `fake_cupy_session(*, compiler=None, options=())`
-
-```python
-with fake_cupy_session():
- sim.run() # CuPy backend; kernels compiled up front and launched on the CPU
-```
-
-Everything a CuPy-backend program needs to run on the CPU: activates the CuPy
-backend (the fake CuPy) and `emulated_launches(compiler=..., options=...)`.
-Raises `RuntimeError` if the fake CuPy is not active.
-
-### `run_in_fake_cupy_subprocess(code, *, env=None, timeout=None)`
-
-```python
-def test_domain_on_the_fake_cupy():
- run_in_fake_cupy_subprocess("from my_sim.tests import check_domain; check_domain()")
-```
-
-The fake CuPy must be installed before anything imports cunumpy, so a test
-process that already uses cunumpy cannot switch to it. This runs
-`python -X faulthandler -c code` in a child process with `CUNUMPY_FAKE_CUPY=1`
-and `OMP_NUM_THREADS=1`, without the environment variables of an MPI launcher
-and with `MAYBEMPI=0`, so the child runs serially and does not join the
-parent's MPI job; `env` adds variables. Under MPI only rank 0 starts the child,
-the other ranks skip the test. If the child fails (or runs longer than
-`timeout` seconds), the test fails with the exit code or the signal (e.g.
-`SIGSEGV`) and the last 50 lines of its stdout and stderr. Returns the
-`subprocess.CompletedProcess` otherwise.
-
-## `memory.HostStaging`
-
-```python
-staging = xp.memory.HostStaging(rho.shape, rho.dtype, buffers=2)
-copy = staging.copy(rho) # returns at once; rho may be overwritten
-...
-if copy.ready():
- h5file["rho"] = copy.result() # a NumPy array
-```
-
-Copies device arrays to page-locked host buffers in the background, so output
-overlaps the next time steps. `copy(array)` snapshots the array on the device
-(on the current stream, after the kernels that wrote it) and copies the
-snapshot to the next of `buffers` pinned host buffers on its own stream. It
-waits only if that buffer's previous copy has not finished, so the program
-runs at most `buffers` copies ahead. The returned `StagedCopy` has `ready()`
-(never waits) and `result()` (waits, returns the host buffer, valid until the
-buffer is reused `buffers` copies later; a stale result raises
-`RuntimeError`). `staging.synchronize()` waits for all copies. Host arrays and
-the NumPy backend copy at once. Device copies are counted by
-`count_transfers()`. The arrays must have the staging shape and dtype.
-
-`copy(array, *, stream=None, event=None)` accepts an explicit producer stream or
-event (at most one). An event queues a wait before snapshotting on the current
-stream. Later writes on another stream must wait for the snapshot/copy;
-`copy.result()` is a conservative completion point. Keep the source's device
-current when submitting copies. Storage binds to that device and rejects another
-device. `ready()`/`result()` temporarily select the owning device and restore the
-caller's device. The first GPU use of CPU-initialized storage allocates fresh
-pinned slots; previous completed CPU handles keep their snapshots. `shape` and
-`dtype` are read-only.
-
-## `memory.DeviceMirror`
-
-```python
-mirror = xp.memory.DeviceMirror(host_array)
-```
-
-Pairs a host NumPy array that another library owns and keeps using on the host
-(for example a stencil vector's `_data` that is exchanged over MPI) with a
-device copy of the same shape and dtype, for accumulation kernels that must
-write into that buffer. `host_array` must be a `numpy.ndarray`; anything else
-raises `TypeError`.
-
-* `device`: the array kernels write into. On the CuPy backend it is a CuPy
- array, allocated on first access as a copy of the host (the only implicit
- transfer). On the NumPy backend it is the host array itself, so the same
- code runs without any copy on the CPU.
-* `to_device()`: copies the host array into the existing device array;
- `to_host(stream=None, event=None)`: copies the device array into the host array, in place, so the
- host array keeps its identity and the owning library sees the new values.
- Both are no-ops on the NumPy backend.
-* `zero()`: zeroes the device array (allocating it empty if needed), or the
- host array on the NumPy backend.
-* `rebind(host_array)`: follows a reallocation by the owner; the device array
- is kept if shape and dtype are unchanged. If the host array's shape or dtype
- changed without a `rebind()`, `device`, `to_device()` and `to_host()` raise
- `ValueError`.
-* `host`, `shape`, `dtype` properties. `to_device()`, `to_host()`, `zero()`
- and `rebind()` return the mirror, for chaining.
-
-Make the mirror's device current before using its device storage. A mismatched
-device raises before copying or zeroing. Pass an explicit producer `stream=` or
-`event=` to `to_host()` when production happened outside the current stream;
-the host copy is complete on return. Initial uploads and refreshes in both
-directions appear in transfer accounting, with payload byte counts. Switching
-to NumPy uses the host array; explicitly refresh with `to_device()` after CPU
-changes before resuming GPU work.
-
-The transfers are explicit so that one per accumulation is visible and
-bounded:
-
-```python
-mirror = xp.memory.DeviceMirror(vector._data)
-mirror.zero()
-accumulate(markers, mirror.device, n_threads=n_markers) # a Kernel
-mirror.to_host() # vector._data now holds the result, same object
-```
-
-### `cunumpy/atomic.cuh`
-
-A CUDA header shipped with the package (found through `cuda_include_dir()`,
-which `CudaKernel` adds automatically) for the many-threads-to-one-cell writes
-of accumulation kernels:
-
-```c
-#include
-
-double cunumpy_atomic_add(double* p, double v); // *p += v, returns old *p
-float cunumpy_atomic_add(float* p, float v);
-double cunumpy_atomic_add_2d(double* data, long long n1,
- long long i, long long j, double v);
-double cunumpy_atomic_add_3d(double* data, long long n1, long long n2,
- long long i, long long j, long long k, double v);
-```
-
-The scalar and indexed helpers also support `int`, `unsigned int`, `long long`,
-and `unsigned long long`. Signed 64-bit addition uses CUDA's unsigned 64-bit
-atomic and modulo-2^64 arithmetic. The indexed helpers address C-contiguous arrays of shape
-`(n0, n1)` and `(n0, n1, n2)`. They wrap `atomicAdd`, a hardware instruction
-for `double` from compute capability 6.0 (sm_60) on; older devices use a
-compare-and-swap loop.
-
-### `cunumpy/random.cuh` and `rng.philox_uniform`
-
-Counter-based random numbers (Philox4x32-10, as in Random123 and cuRAND): a
-pure function of a key and a counter, with no generator state, so each thread
-draws from `(seed, stream, counter)`, e.g. `(seed, particle id, step)`, and
-the host computes the same numbers:
-
-```c
-#include
-
-cunumpy_u32x4 cunumpy_philox4x32_10(cunumpy_u32x4 ctr, unsigned int key0, unsigned int key1);
-double cunumpy_uniform(seed, stream, counter); // [0, 1), 53 bits
-void cunumpy_uniform2(seed, stream, counter, &u0, &u1); // two from one call
-double cunumpy_normal(seed, stream, counter); // Box-Muller
-void cunumpy_normal2(seed, stream, counter, &z0, &z1);
-```
-
-(`seed`, `stream` and `counter` are `unsigned long long`.)
-
-```python
-ids = xp.arange(n, dtype=xp.uint64)
-u0, u1 = xp.rng.philox_uniform2(seed, ids, step) # == cunumpy_uniform2 in thread i
-z0, z1 = xp.rng.philox_normal2(seed, ids, step)
-words = xp.rng.philox4x32_10(counter_words, key0, key1) # the raw generator
-```
-
-`xp.rng.philox_uniform`, `philox_uniform2`, `philox_normal`, `philox_normal2` and
-`philox4x32_10` broadcast their arguments and return NumPy or CuPy arrays,
-matching the inputs. The uniform numbers equal the kernel's bit for bit; the
-normal numbers can differ in the last bits (`log`, `sqrt`, `sin`, `cos` on the
-GPU are not the host's). The generator passes the Random123 known-answer
-tests. Use a different `counter` for every random decision of a step.
-
-### `cunumpy/morton.cuh` and `algorithms.morton_keys`
-
-Morton (Z-order) keys: the bits of a point's integer cell coordinates,
-interleaved into one `uint64`. Sorted by key, nearby points are nearby in
-memory, and the points of every node of a quadtree (2D) or octree (3D) on the
-same box form a contiguous range, the starting point of tree builds on the
-GPU.
-
-```python
-keys = xp.algorithms.morton_keys(
- positions, lower, upper, levels
-) # (n, 2|3) -> (n,) uint64
-keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
-node = keys >> np.uint64(ndim * (levels - level)) # node index at `level`
-cells = xp.algorithms.morton_decode(node, ndim) # its integer coordinates
-key = xp.algorithms.morton_encode(ix, iy) # from integer cells
-scales = xp.algorithms.morton_scales(
- lower, upper, levels
-) # 2**levels / (upper - lower)
-```
-
-```c
-#include
-
-unsigned long long cunumpy_morton_key2(x, y, lower_x, lower_y, scale_x, scale_y, levels);
-unsigned long long cunumpy_morton_key3(x, y, z, lower_x, ..., scale_x, ..., levels);
-unsigned long long cunumpy_morton_encode2(ix, iy); // and _encode3(ix, iy, iz)
-unsigned long long cunumpy_morton_cell(x, lower, scale, levels);
-unsigned long long cunumpy_morton_spread2(v); // and _compact2, _spread3, _compact3
-```
-
-`levels` is the number of bits per axis, at most 32 in 2D and 21 in 3D
-(`xp.algorithms.MAX_MORTON_LEVELS`). Axis 0 is the lowest bit of every group of `ndim`
-bits; the top group is the child of the root. The cell along an axis is
-`floor((x - lower) * scale)` clipped to `[0, 2**levels - 1]`: points on a cell
-boundary go to the upper cell, points outside the box to the nearest face, and
-`lower > upper` reverses the axis. Given the `morton_scales` of the host, the
-kernel functions return the host keys bit for bit. All host functions run on
-NumPy and CuPy arrays.
-
-### `cunumpy/reduce.cuh`
-
-Warp- and block-level reductions for hand-written kernels: in-kernel
-diagnostics (energy, momentum, total charge, the maximum velocity for a CFL
-check) and combining values in a block before one atomic write.
-
-```c
-#include
-
-T cunumpy_warp_sum(T v, unsigned mask = 0xffffffffu); // also min/max
-T cunumpy_block_sum(T v); T cunumpy_block_min(T v); T cunumpy_block_max(T v);
-void cunumpy_block_sum_to(T* out, T v); // *out += block sum, one atomic per block
-int cunumpy_block_thread(); // linear thread index in a 1D-3D block
-int cunumpy_block_threads(); // threads per block
-```
-
-`T` is `int`, `unsigned`, `long long`, `unsigned long long`, `float` or
-`double`; `block_sum_to` supports the same types. Every
-thread gets the result. Rules: every thread of the block calls block
-functions (no early `return`; threads without a value pass the identity, e.g.
-`0.0` for a sum). Warp functions support arbitrary nonzero masks: every named
-lane calls with the same mask, and only named lanes participate. The default
-requires all 32 lanes. Block functions handle partial warps and any legal block
-size. The block functions use 32 values of static shared memory per
-type and may be called several times in a kernel.
-
-```c
-extern "C" __global__ void kinetic_energy(const double* v, long long n,
- double mass, double* energy) {
- long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x;
- double e = i < n ? 0.5 * mass * v[i] * v[i] : 0.0;
- cunumpy_block_sum_to(energy, e); // zero *energy before the launch
-}
-```
-
-### `cunumpy/scan.cuh`
-
-Warp/block prefix sums for compaction and binning:
-
-```c
-#include
-T cunumpy_warp_inclusive_sum(T v, unsigned mask = 0xffffffffu);
-T cunumpy_warp_exclusive_sum(T v, unsigned mask = 0xffffffffu);
-T cunumpy_block_inclusive_sum(T v);
-T cunumpy_block_exclusive_sum(T v);
-```
-
-Same shuffle arithmetic types and collective participation rules as reductions.
-Warp order follows increasing participating lane IDs, including sparse masks.
-Block order follows linear thread index (x fastest), with partial warps supported.
-Exclusive scans start at zero. Block scans use 32 shared values per type and
-specialization; repeated calls are supported. These are block-local operations;
-a global compaction needs a separate pass to combine block totals.
-
-## `scipy`
-
-```python
-A = xp.scipy.sparse.csr_matrix((data, (rows, cols)), shape=(n, n))
-x, info = xp.scipy.sparse.linalg.cg(A, b)
-rho_k = xp.scipy.fft.rfftn(rho)
-```
-
-SciPy for the active backend: `scipy` on NumPy, `cupyx.scipy` on CuPy. The
-forwarded subpackages are those `cupyx.scipy` has (`SUBMODULES`): `fft`,
-`fftpack`, `interpolate`, `linalg`, `ndimage`, `signal`, `sparse`,
-`sparse.csgraph`, `sparse.linalg`, `spatial`, `special`, `stats`. Names are
-looked up at every access, so a backend switch takes effect immediately
-(Python caches the imports). Nothing is imported until a name is used; SciPy
-is not a dependency of cunumpy.
-
-* A name missing on the active backend (`cupyx.scipy` covers part of SciPy)
- raises `AttributeError` naming the backend. Keyword arguments can differ
- too (SciPy's `cg(..., rtol=)` is `tol=` in CuPy).
-* `xp.scipy.special.available("erfcx")` checks a name without raising.
-* `xp.scipy.sparse.linalg.resolve()` returns the module itself.
-* A missing SciPy (NumPy backend) or CuPy raises `ImportError` with the module
- it needs.
-
-Sparse matrices assembled on the host move to the device once, with the
-constructor of the device type: `xp.scipy.sparse.csr_matrix(host_matrix)` on
-the CuPy backend copies a SciPy matrix; `matrix.get()` copies back.
-
-## `kernels.fuse(function=None, *, kernel_name=None)`
-
-```python
-@xp.kernels.fuse
-def pressure(rho, T, gamma):
- return (gamma - 1.0) * rho * T
-```
-
-Compiles an elementwise function into one kernel with `cupy.fuse` when it is
-called with a CuPy array (positional or keyword), with the CuPy backend active
-while it is traced, so `xp.exp` etc. resolve to CuPy ufuncs; other calls run
-the function as it is. The fused kernel is created on first use and reused.
-`kernel_name` names it in profilers (default: the function name). The function
-must be elementwise in the sense of `cupy.fuse`: arithmetic, comparisons,
-ufuncs, `xp.where`, and supported reductions as the last operation; no Python
-control flow on array values or indexing. Test the CuPy path: a function
-`cupy.fuse` cannot trace raises at its first call with CuPy arrays.
-
-## `petsc.petsc_vec(array, comm=None)`
-
-```python
-b_vec = xp.petsc.petsc_vec(b) # b: NumPy or CuPy array, shared, never copied
-x_vec = xp.petsc.petsc_vec(x)
-xp.synchronize()
-ksp.solve(b_vec, x_vec) # PETSc writes into x
-xp.synchronize()
-```
-
-A `petsc4py.PETSc.Vec` that shares the memory of a C-contiguous array of
-`PETSc.ScalarType`, through DLPack: a `seq`/`mpi` vector for a NumPy array, a
-`seqcuda`/`mpicuda` (or HIP) vector for a CuPy array. The vector keeps a
-reference to the array. With several processes, `array` is this process's
-part of the vector (`comm`, default `COMM_WORLD`).
-
-* Another dtype raises `TypeError` and a non-contiguous array `ValueError`
- (both would need a copy).
-* A CuPy array with a petsc4py built without CUDA/HIP raises `RuntimeError`
- instead of PETSc working on a host copy.
-* Synchronize between CuPy and PETSc work on the same memory (they may use
- different streams).
-* For the solve to stay on the GPU, the matrix must be a GPU type too
- (`aijcusparse`, or `-mat_type aijcusparse -vec_type cuda`).
-
-## Version
-
-`xp.__version__` is the installed package version. When package metadata is
-not available (for example, some source-tree imports), it is
-`"0.0.0+unknown"`.
diff --git a/docs/source/conf.py b/docs/source/conf.py
index e6a5abb..ffc2437 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -17,9 +17,54 @@
"nbsphinx",
"sphinx.ext.mathjax",
"sphinx.ext.autodoc",
+ "sphinx.ext.intersphinx",
+ "numpydoc", # NumPy-format docstrings (see contributing.md)
"myst_parser", # enable Markdown support
]
+# The API reference (reference/*.md) is generated from the docstrings.
+autodoc_typehints = "none" # types are written in the docstrings
+autodoc_member_order = "bysource"
+nitpicky = True # with -W in CI, a broken cross-reference fails the build
+numpydoc_show_class_members = False
+numpydoc_xref_param_type = False
+intersphinx_mapping = {
+ "python": ("https://docs.python.org/3", None),
+ "numpy": ("https://numpy.org/doc/stable/", None),
+}
+
+
+def _defining_comment(name, obj):
+ """The ``#:`` comment of a constant that a public module re-exports."""
+ import sys
+
+ from sphinx.pycode import ModuleAnalyzer
+
+ for modname, module in list(sys.modules.items()):
+ if not modname.startswith("cunumpy._") or getattr(module, name, None) is not obj:
+ continue
+ try:
+ docs = ModuleAnalyzer.for_module(modname).find_attr_docs()
+ except Exception:
+ continue
+ if ("", name) in docs:
+ return list(docs[("", name)])
+ return None
+
+
+def _reexported_constant_doc(app, what, name, obj, options, lines):
+ # constants imported from a private module are listed with ``autodata`` on
+ # the reference pages; their ``#:`` comment is in the defining module
+ if what == "data":
+ comment = _defining_comment(name.rsplit(".", 1)[-1], obj)
+ if comment:
+ lines[:] = comment
+
+
+def setup(app):
+ app.connect("autodoc-process-docstring", _reexported_constant_doc)
+
+
exclude_patterns = []
# Recognize both .rst and .md
@@ -35,7 +80,6 @@
html_theme = "sphinx_book_theme"
-html_static_path = ["_static"]
templates_path = ["_templates"]
diff --git a/docs/source/contributing.md b/docs/source/contributing.md
new file mode 100644
index 0000000..a04e414
--- /dev/null
+++ b/docs/source/contributing.md
@@ -0,0 +1,88 @@
+# Writing docstrings
+
+Docstrings use the [NumPy format](https://numpydoc.readthedocs.io/en/latest/format.html)
+and are the source of the [API reference](reference/index.md). Keep them short:
+help text, not a guide. Longer explanations go in a guide page that the
+docstring links to.
+
+```python
+def cell_offsets(sorted_cells, n_cells):
+ """Return the offsets of each cell into sorted cell IDs.
+
+ Cell ``k`` occupies ``[offsets[k], offsets[k + 1])``. The result stays on
+ the backend of `sorted_cells`; validation may synchronize on CUDA.
+
+ Parameters
+ ----------
+ sorted_cells : array of int
+ Cell IDs in ``[0, n_cells)``, sorted in nondecreasing order.
+ n_cells : int
+ Number of cells.
+
+ Returns
+ -------
+ array of int64
+ ``n_cells + 1`` offsets.
+
+ Raises
+ ------
+ ValueError
+ If `sorted_cells` is not sorted or has an ID outside ``[0, n_cells)``.
+
+ See Also
+ --------
+ segment_boundaries : The runs of sorted keys that are present.
+
+ Examples
+ --------
+ >>> xp.algorithms.cell_offsets(xp.asarray([0, 0, 2]), 3)
+ array([0, 2, 2, 3])
+ """
+```
+
+## Rules
+
+- **Summary**: one line, imperative mood ("Return", "Count", "Select"),
+ ending with a period. Properties and attributes use a noun phrase ("The
+ device of the array.").
+- **Extended summary**: at most about three sentences, only when needed. Keep
+ what is specific to cunumpy: which backend the result lives on, whether the
+ call copies, and whether it synchronizes the GPU.
+- **Sections**, in this order, empty ones left out: Parameters, Returns (or
+ Yields), Raises, See Also, Notes, Examples.
+- **Types** are short and conceptual (`array of int`, `str`, `int, optional`,
+ `CudaKernel`), not a copy of the annotation. Write `, optional` for
+ parameters that have a default and say the default in the description when
+ it is not obvious.
+- **Raises**: only errors a caller is expected to handle (bad arguments,
+ missing CUDA), not every internal check.
+- **Classes**: document the constructor parameters in the class docstring;
+ `__init__` has no docstring. Public attributes go in an Attributes section.
+- **Properties and trivial methods** (`start`, `stop`, `__getitem__`): one
+ line, no sections.
+- **Private helpers** (names starting with `_`): at most a one-line summary.
+- **Cross-references** use the public path: `` :class:`~cunumpy.kernels.CudaKernel` ``,
+ not `cunumpy._cuda_kernel.CudaKernel`. Parameter names in single
+ backticks, literal code in double backticks.
+
+## Examples
+
+- Every public function and class has an Examples section; methods have one
+ when the use is not obvious.
+- One example, two to five lines. The examples assume `import cunumpy as xp`
+ and `import numpy as np` and run on the NumPy backend (see the root
+ `conftest.py`), so don't import those again.
+- Examples are tested: `pytest --doctest-modules src/cunumpy`. Lines that need
+ a GPU, MPI, Pyccel or Metal get `# doctest: +SKIP`, as do lines that use
+ undefined names for illustration. Prefer real, runnable examples.
+- Don't paste long CUDA sources; a kernel of a few lines is fine, otherwise
+ link the [kernel guides](kernels/overview.md).
+
+## Checking
+
+```bash
+ruff check src # D rules: summary line, imperative mood, ...
+python -m numpydoc lint src/cunumpy/*.py # sections, documented parameters
+python -m pytest --doctest-modules src/cunumpy
+cd docs && make html
+```
diff --git a/docs/source/guides/backends.md b/docs/source/guides/backends.md
index c670a94..01ff114 100644
--- a/docs/source/guides/backends.md
+++ b/docs/source/guides/backends.md
@@ -15,7 +15,10 @@ CUNUMPY_BACKEND=cupy python simulate.py
This is the least intrusive option for scripts and batch jobs: the code does
not change, and a job script decides whether it runs on a GPU. The variable is
-read once, at import; setting it later in `os.environ` has no effect.
+read once, at import; setting it later in `os.environ` has no effect. The value
+is case-insensitive (`CUPY` works too), and any other value selects NumPy.
+Requesting CuPy this way falls back to NumPy, like `set_backend("cupy")`, when
+CuPy or CUDA is not usable.
The startup setting is named `CUNUMPY_BACKEND`; migrate existing job scripts
from `ARRAY_BACKEND`. CuNumpy-owned environment options share the `CUNUMPY_`
diff --git a/docs/source/guides/data-movement.md b/docs/source/guides/data-movement.md
index 21c7603..27f867a 100644
--- a/docs/source/guides/data-movement.md
+++ b/docs/source/guides/data-movement.md
@@ -109,9 +109,74 @@ def test_step_stays_on_device():
The counter only sees transfers made through CuNumpy helpers. Forwarded backend
operations such as `xp.asarray(host)`, raw `cupy.asarray(host)`,
-`device_array.get()`, `float(device_scalar)` and conversions inside other
-libraries are invisible to it; use `nsys` to find those (see [Timing and
-profiling](profiling.md)).
+`numpy.asarray(device_array)`, `device_array.get()` and conversions inside other
+libraries are invisible to it. So is `float(device_scalar)`, an implicit sync,
+with the real CuPy: it cannot be observed from Python (the fake CuPy of the
+[test helpers](../kernels/testing.md) records it as a `sync` event marked
+`[implicit]`). Use `nsys` or CuPy's profiling hooks to find those (see [Timing
+and profiling](profiling.md)).
+
+Besides the copies, the counter records `sync` events, where the host waited
+for the device (`xp.synchronize()`, the MPI helpers, the CUDA debug mode). They
+are listed in the report and counted in `counter.syncs`, not in `total`. The
+report marks copies the host did not wait for (`to_host_async()`,
+`HostStaging.copy()`) as `[async]`. Blocks can be nested, and
+`count_transfers(counter)` adds the events of another block to an existing
+counter, e.g. to accumulate over several calls.
+
+## Transfer rules per phase: `TransferBudget`
+
+A real program may copy in some phases (diagnostics, output) but not in others
+(the time step). `TransferBudget` counts the transfers per named phase and
+checks a rule for each:
+
+```python
+budget = xp.profiling.TransferBudget(started=False)
+model.integrate = budget.count("integrate")(model.integrate) # a decorator
+... # setup: not counted
+budget.start()
+for step in range(n_steps):
+ model.integrate(dt)
+ with budget.phase("output"): # or a context
+ save(model)
+
+budget.require("integrate", allow={"to_host": {"max_nbytes": 8}}, calls=n_steps)
+budget.require("output", allow={"to_host": {"max_count": n, "max_total_bytes": b}})
+budget.check() # AssertionError with budget.report() if a rule is broken
+```
+
+`phase(name)` yields the phase's `TransferCounter` (`budget[name]`); the events
+of a phase accumulate over its calls and `budget.calls[name]` counts the calls.
+Phases nest: an event is counted in the innermost phase only. With
+`started=False` phases run uncounted until `start()`, and again after `stop()`.
+
+`require()` allows the event kinds listed in `allow`, each with optional limits,
+and forbids every other kind except those in `ignore` (`device_copy` and `sync`
+by default). Next to the size and count limits, `blocking` and `implicit`
+require a value of the events' fields: `{"to_host": {"max_nbytes": 8,
+"blocking": False}}` allows only non-blocking scalar copies, and
+`{"sync": {"implicit": False}}` with `ignore=("device_copy",)` rejects the scalar
+reads of the fake CuPy. `calls` is the number of calls the phase must have had.
+`violations()` lists the broken rules with the offending events and where they
+happened, and `report()` shows every phase's events and the broken rules.
+
+## Read a value without waiting: `to_host_async`
+
+```python
+pending = xp.to_host_async(residual_norm) # a device scalar; returns at once
+... # queue the next iteration's kernels
+if pending.ready() and pending.result() < tol:
+ break
+```
+
+A convergence test read one iteration late needs no sync per iteration: the
+copy runs on its own stream into page-locked memory, after the work queued so
+far. `result()` returns a NumPy scalar for a 0-d array, else a NumPy array, and
+waits only if the copy has not finished; that wait is counted as a `sync`. On
+the NumPy backend, and on the fake CuPy, the value is copied at once and
+`ready()` is always true; on the fake CuPy the copy is still counted. For large
+arrays copied repeatedly, such as output snapshots, use `memory.HostStaging`
+(see [Particle codes](particle-codes.md)), which reuses its buffers.
## Build device arguments once: `as_device_array`
diff --git a/docs/source/guides/mpi.md b/docs/source/guides/mpi.md
index afa80c7..9d1b56f 100644
--- a/docs/source/guides/mpi.md
+++ b/docs/source/guides/mpi.md
@@ -58,8 +58,10 @@ Importing `mpi4py.MPI` initializes MPI by default. Do it after step 1.
To run the same program serially without starting MPI, get the module from
`xp.mpi.get_mpi()` instead of importing it: it returns `mpi4py.MPI` when the
process was started by `mpirun`/`mpiexec`/`srun`, and otherwise a serial
-stand-in whose `COMM_WORLD` has size 1 (see the
-[API reference](../api.md)):
+stand-in whose `COMM_WORLD` has size 1. `get_mpi()`, `launched_under_mpi()`,
+`local_rank()` and the serial stand-in come from
+[maybempi](https://max-models.github.io/maybempi/), which documents them, and
+are re-exported in `xp.mpi`:
```python
MPI = xp.mpi.get_mpi()
@@ -130,9 +132,12 @@ with (
A host array is yielded as it is. A device array is yielded as it is (after
`synchronize_for_mpi`) when MPI is CUDA-aware, and otherwise replaced by a
pinned host copy: filled from the device before the block when `send=True`,
-copied back into the device array after the block when `recv=True`. The
-copies are counted by `count_transfers()`, so a GPU run with a plain MPI
-build is visible in the transfer report. Without a recorded answer (no
+copied back into the device array after the block when `recv=True`. A
+receive-only device buffer also waits for preceding device work before MPI
+writes into the host copy. The copies are counted by `count_transfers()`, so a
+GPU run with a plain MPI build is visible in the transfer report. To reuse
+pinned staging buffers instead of allocating them per call, pass an
+`xp.mpi.MPIStaging` (see [Execution helpers](execution-helpers.md)). Without a recorded answer (no
`mpi_is_cuda_aware()` call and no `set_mpi_cuda_aware()`), a device array
raises instead of guessing.
diff --git a/docs/source/guides/particle-codes.md b/docs/source/guides/particle-codes.md
index e3c468b..fe4a84f 100644
--- a/docs/source/guides/particle-codes.md
+++ b/docs/source/guides/particle-codes.md
@@ -45,7 +45,8 @@ result independent of how the markers were ordered before.
the reordering of several arrays in one call.
For a tree code, or for better locality in 2D and 3D, sort by Morton key
-(`xp.algorithms.morton_keys`, see the API page) instead of by cell: the markers of every
+(`xp.algorithms.morton_keys`, see [CUDA headers](../kernels/cuda-headers.md) for
+the key layout and the kernel version) instead of by cell: the markers of every
quadtree or octree node are then a contiguous range of the sorted arrays.
## Deposit without atomics: sort, then reduce
@@ -152,7 +153,12 @@ for step in range(n_steps):
`copy()` snapshots the array on the device and copies the snapshot to pinned
host memory on its own stream, so the next steps overwrite `rho` while the copy
-runs. See `HostStaging` in the [API reference](../api.md).
+runs. With `buffers=2` (the default) one copy runs while the previous result is
+written out; `copy()` waits only if the buffer it reuses is still being copied.
+A `result()` is the staging buffer itself, valid until it is reused, and a
+stale result raises `RuntimeError`. See
+{class}`~cunumpy.memory.HostStaging` and [Execution
+helpers](execution-helpers.md) for producer streams and devices.
## Run a time step as a CUDA graph
diff --git a/docs/source/index.md b/docs/source/index.md
index 49f1fd2..fdcf08d 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -27,7 +27,8 @@ CUDA version for GPU use. Start with the [Quickstart](quickstart.md).
* **Porting compiled kernels to CUDA**: start at [Porting kernels to the
GPU](kernels/overview.md).
* **Complete programs**: [Worked examples](examples/index.md).
-* **Looking up a function**: the [API reference](api.md).
+* **Looking up a function**: the [API reference](reference/index.md), and
+ the [API overview](api.md) for the module layout.
* **Using an AI assistant**: point it at [the guide for AI
assistants](ai-assistants.md), which ships with the package as
`cunumpy/LLM_GUIDE.md`.
@@ -64,6 +65,7 @@ array-api-compat
kernels/overview
kernels/pyccel-kernel
kernels/cuda-kernel
+kernels/cuda-headers
kernels/metal-kernel
kernels/dispatch
kernels/arguments
@@ -88,4 +90,6 @@ troubleshooting
pyodide
ai-assistants
api
+reference/index
+contributing
```
diff --git a/docs/source/installation.md b/docs/source/installation.md
index 7c62fb5..6eb5b94 100644
--- a/docs/source/installation.md
+++ b/docs/source/installation.md
@@ -91,6 +91,10 @@ Tests that need a GPU are skipped automatically where CuPy is not functional.
| `CUNUMPY_FAKE_CUPY=1` | install the strict CPU stand-in for CuPy for tests |
| `CUNUMPY_REQUIRE_CUDA=1` | require a real usable GPU when starting the test suite (CI guard) |
+The two kernel implementation variables are case-insensitive and stripped of
+whitespace; an unset or empty value means the default (automatic) selection,
+and an unsupported value raises `ValueError` when CuNumpy is imported.
+
Use `CUNUMPY_BACKEND` in job scripts; the former `ARRAY_BACKEND` setting is no
longer read. Standard toolchain/device variables such as `CXX` and
`CUDA_VISIBLE_DEVICES` keep their standard meanings.
diff --git a/docs/source/kernels/arguments.md b/docs/source/kernels/arguments.md
index 6c53e39..709deb6 100644
--- a/docs/source/kernels/arguments.md
+++ b/docs/source/kernels/arguments.md
@@ -48,6 +48,9 @@ particles = DeviceParticles(x, v)
push(0.1, particles, n_threads=particles.positions.shape[0]) # -> push(0.1, x, v, n)
```
+The values may include packed structs (`CudaStructValue.packed`) for the
+struct parameters of the kernel.
+
Build the object once and reuse it. `as_device_array()` references existing
device arrays of the right dtype and copies everything else exactly once (see
[Data movement](../guides/data-movement.md), section "Build device arguments once").
@@ -207,7 +210,7 @@ struct MarkerArgs {
Mappings: `float` to `double`, `int` to `long long` (Pyccel integers are
64-bit; change with `int_type=`), `bool` to `bool`, NumPy scalar types to their
-C types, `"float[:, :]"` to `Array2D` (1 to 3 dimensions). `Final[...]`
+C types, `"float[:, :]"` to `Array2D` (1 to 16 dimensions). `Final[...]`
and `const` are ignored. `scalar_names={"float": "float"}` switches to single
precision. Parameters without a mappable annotation raise `ValueError`.
diff --git a/docs/source/kernels/cuda-headers.md b/docs/source/kernels/cuda-headers.md
new file mode 100644
index 0000000..6c2115f
--- /dev/null
+++ b/docs/source/kernels/cuda-headers.md
@@ -0,0 +1,215 @@
+# CUDA headers
+
+CuNumpy ships CUDA headers for hand-written kernels. Every `CudaKernel` finds
+them without configuration: it adds their directory to its NVRTC options as
+`-I` (once), so a kernel source can write `#include `.
+`xp.cuda.cuda_include_dir()` returns that directory (a `str`) to pass the same
+headers to other compilers. The shipped headers take part in the compile cache
+key also when included in angle brackets, so upgrading CuNumpy with a changed
+header recompiles the kernels that use it (see "Headers and the compile cache"
+in [Writing CUDA kernels](cuda-kernel.md)).
+
+| Header | Provides |
+| --- | --- |
+| `cunumpy/index.cuh` | thread-index macros |
+| `cunumpy/array_view.cuh` | strided views `Array1D` to `Array16D` |
+| `cunumpy/atomic.cuh` | atomic adds, also into 2D/3D arrays |
+| `cunumpy/reduce.cuh` | warp and block reductions |
+| `cunumpy/scan.cuh` | warp and block prefix sums |
+| `cunumpy/random.cuh` | counter-based random numbers, equal to `xp.rng.philox_*` |
+| `cunumpy/morton.cuh` | Morton (Z-order) keys, equal to `xp.algorithms.morton_*` |
+
+## `cunumpy/index.cuh`
+
+```c
+#include
+
+CUNUMPY_THREAD_1D(i, n); // long long i = global x index; return if i >= n
+CUNUMPY_THREAD_2D(i, j, ni, nj); // 2D launch, n_threads=(ni, nj)
+CUNUMPY_THREAD_3D(i, j, k, ni, nj, nk);
+CUNUMPY_GRID_STRIDE_1D(i, n) { ... } // grid-stride loop; launch with any grid
+CUNUMPY_GLOBAL_INDEX_X() // also _Y(), _Z(): the global index as long long
+```
+
+The `CUNUMPY_THREAD_*` macros return early, so do not use them before a block
+collective of `reduce.cuh` or `scan.cuh`, which every thread of the block must
+reach.
+
+## `cunumpy/array_view.cuh`
+
+Strided views `Array1D` to `Array16D`: `T* data`, `long long
+shape[ndim]`, `long long strides[ndim]` (in elements, not bytes),
+`operator()(i, j, ...)` returning a reference to the element, and `size()`. A
+kernel indexes `a(i, j)` like the Pyccel kernel it is ported from indexes
+`a[i, j]`, without hand-passed sizes. A 4D view can, for example, describe a 3D
+grid of vector components `(nx, ny, nz, ncomp)`. `CArray1D` to
+`CArray16D` are C-contiguous views without strides; they convert to the
+strided view of the same dimension.
+
+A kernel parameter or a `CudaStruct` field of type `ArrayND`, for the
+scalar C types `CudaKernel` supports, takes a CuPy array of that dtype and
+number of dimensions (a mismatch raises `TypeError`), contiguous or not. It is
+packed by value into pointer, shape and strides with the memory layout of the
+C struct: 8-byte aligned, `sizeof == 8 * (1 + 2 * ndim)` (`8 * (1 + ndim)` for
+a `CArrayND`); the header checks this with `static_assert`.
+`CudaParameter.view_ndim` is the number of dimensions of such a parameter.
+
+Compiling with `options=("-DCUNUMPY_BOUNDS_CHECK",)`, or in [debug
+mode](debugging.md), checks every index against the shape: an out-of-bounds
+index prints a message and traps the kernel. Without the macro, indexing is
+unchecked.
+
+## `cunumpy/atomic.cuh`
+
+For the many-threads-to-one-cell writes of accumulation kernels (see
+[Accumulation kernels](accumulation.md)):
+
+```c
+#include
+
+double cunumpy_atomic_add(double* p, double v); // *p += v, returns old *p
+float cunumpy_atomic_add(float* p, float v);
+double cunumpy_atomic_add_2d(double* data, long long n1,
+ long long i, long long j, double v);
+double cunumpy_atomic_add_3d(double* data, long long n1, long long n2,
+ long long i, long long j, long long k, double v);
+```
+
+The scalar and indexed helpers also support `int`, `unsigned int`, `long long`
+and `unsigned long long`, returning the old value. Signed 64-bit addition uses
+CUDA's unsigned 64-bit atomic with modulo-2^64 arithmetic. The indexed helpers
+address C-contiguous arrays of shape `(n0, n1)` and `(n0, n1, n2)`. They wrap
+`atomicAdd`, a hardware instruction for `double` from compute capability 6.0
+(sm_60) on; older devices use a compare-and-swap loop.
+
+## `cunumpy/reduce.cuh`
+
+Warp- and block-level reductions for hand-written kernels: in-kernel
+diagnostics (energy, momentum, total charge, the maximum velocity for a CFL
+check) and combining values in a block before one atomic write.
+
+```c
+#include
+
+T cunumpy_warp_sum(T v, unsigned mask = 0xffffffffu); // also _min, _max
+T cunumpy_block_sum(T v); T cunumpy_block_min(T v); T cunumpy_block_max(T v);
+void cunumpy_block_sum_to(T* out, T v); // *out += block sum, one atomic per block
+int cunumpy_block_thread(); // linear thread index in a 1D-3D block
+int cunumpy_block_threads(); // threads per block
+```
+
+`T` is `int`, `unsigned`, `long long`, `unsigned long long`, `float` or
+`double`; `block_sum_to` supports the same types. Every participating thread
+gets the result.
+
+* Every thread of the block calls the block functions: no early `return`
+ (and no `CUNUMPY_THREAD_1D` before them); threads without a value pass the
+ identity, e.g. `0.0` for a sum.
+* Block functions handle partial warps and any legal block shape, with x
+ fastest. They use 32 values of static shared memory per type and may be
+ called several times in a kernel.
+* Warp functions take an optional nonzero lane mask, also a sparse one: only
+ the named lanes participate, and all of them call with the same mask. The
+ default requires all 32 lanes. These rules follow [NVIDIA's warp intrinsic
+ constraints](https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html).
+
+```c
+extern "C" __global__ void kinetic_energy(const double* v, long long n,
+ double mass, double* energy) {
+ long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x;
+ double e = i < n ? 0.5 * mass * v[i] * v[i] : 0.0;
+ cunumpy_block_sum_to(energy, e); // zero *energy before the launch
+}
+```
+
+## `cunumpy/scan.cuh`
+
+Warp and block prefix sums, for compaction and binning:
+
+```c
+#include
+
+T cunumpy_warp_inclusive_sum(T v, unsigned mask = 0xffffffffu);
+T cunumpy_warp_exclusive_sum(T v, unsigned mask = 0xffffffffu);
+T cunumpy_block_inclusive_sum(T v);
+T cunumpy_block_exclusive_sum(T v);
+```
+
+The types and participation rules are those of `reduce.cuh`. Warp scans follow
+increasing lane IDs of the participating lanes, also for sparse masks; block
+scans follow the linear thread index (x fastest), with partial warps
+supported. Exclusive scans start at zero. Block scans use 32 shared values per
+type and specialization and may be called repeatedly. They are block-local: a
+global compaction needs another pass that combines the block totals (see the
+example in [Execution helpers](../guides/execution-helpers.md)).
+
+## `cunumpy/random.cuh`
+
+Counter-based random numbers (Philox4x32-10, as in Random123 and cuRAND): a
+pure function of a key and a counter, with no generator state. Each thread
+draws from `(seed, stream, counter)`, e.g. `(seed, particle id, step)`, and the
+host computes the same numbers:
+
+```c
+#include
+
+cunumpy_u32x4 cunumpy_philox4x32_10(cunumpy_u32x4 ctr, unsigned int key0, unsigned int key1);
+double cunumpy_uniform(seed, stream, counter); // [0, 1), 53 bits
+void cunumpy_uniform2(seed, stream, counter, &u0, &u1); // two from one call
+double cunumpy_normal(seed, stream, counter); // Box-Muller
+void cunumpy_normal2(seed, stream, counter, &z0, &z1);
+```
+
+`seed`, `stream` and `counter` are `unsigned long long`. On the host:
+
+```python
+ids = xp.arange(n, dtype=xp.uint64)
+u0, u1 = xp.rng.philox_uniform2(seed, ids, step) # == cunumpy_uniform2 in thread i
+z0, z1 = xp.rng.philox_normal2(seed, ids, step)
+words = xp.rng.philox4x32_10(counter_words, key0, key1) # the raw generator
+```
+
+`xp.rng.philox_uniform`, `philox_uniform2`, `philox_normal`, `philox_normal2`
+and `philox4x32_10` broadcast their arguments and return NumPy or CuPy arrays,
+matching the inputs. The uniform numbers equal the kernel's bit for bit; the
+normal numbers can differ in the last bits (`log`, `sqrt`, `sin` and `cos` on
+the GPU are not the host's). The generator passes the Random123 known-answer
+tests. Use a different `counter` for every random decision of a step.
+
+## `cunumpy/morton.cuh`
+
+Morton (Z-order) keys: the bits of a point's integer cell coordinates,
+interleaved into one `uint64`. Sorted by key, nearby points are nearby in
+memory, and the points of every node of a quadtree (2D) or octree (3D) on the
+same box form a contiguous range, the starting point of tree builds on the
+GPU. On the host:
+
+```python
+keys = xp.algorithms.morton_keys(positions, lower, upper, levels) # (n, 2|3) -> (n,) uint64
+keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
+node = keys >> np.uint64(ndim * (levels - level)) # node index at `level`
+cells = xp.algorithms.morton_decode(node, ndim) # its integer coordinates
+key = xp.algorithms.morton_encode(ix, iy) # from integer cells
+scales = xp.algorithms.morton_scales(lower, upper, levels) # 2**levels / (upper - lower)
+```
+
+In a kernel:
+
+```c
+#include
+
+unsigned long long cunumpy_morton_key2(x, y, lower_x, lower_y, scale_x, scale_y, levels);
+unsigned long long cunumpy_morton_key3(x, y, z, lower_x, ..., scale_x, ..., levels);
+unsigned long long cunumpy_morton_encode2(ix, iy); // and _encode3(ix, iy, iz)
+unsigned long long cunumpy_morton_cell(x, lower, scale, levels);
+unsigned long long cunumpy_morton_spread2(v); // and _compact2, _spread3, _compact3
+```
+
+`levels` is the number of bits per axis, at most 32 in 2D and 21 in 3D
+(`xp.algorithms.MAX_MORTON_LEVELS`). Axis 0 is the lowest bit of every group of
+`ndim` bits; the top group is the child of the root. The cell along an axis is
+`floor((x - lower) * scale)` clipped to `[0, 2**levels - 1]`: points on a cell
+boundary go to the upper cell, points outside the box to the nearest face, and
+`lower > upper` reverses the axis. Given the `morton_scales` of the host, the
+kernel functions return the host keys bit for bit. All host functions run on
+NumPy and CuPy arrays.
diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md
index a501a4e..c0f4e71 100644
--- a/docs/source/kernels/cuda-kernel.md
+++ b/docs/source/kernels/cuda-kernel.md
@@ -57,9 +57,21 @@ every call:
| `Array1D` ... `Array16D` | CuPy array of dtype `T` and that ndim, contiguous or not | wrong dtype or ndim |
| a `CudaStruct` type | a value of that struct | anything else |
-C types map to NumPy dtypes as on 64-bit Linux: `int` is `int32`, `long` and
-`long long` are `int64`, `float` is `float32`, `double` is `float64`.
-`xp.cuda.ctype_of(np.float64)` returns `"double"`, useful when generating source.
+C types map to NumPy dtypes as on 64-bit Linux (LP64): `int` is `int32`, `long`
+and `long long` are `int64`, `float` is `float32`, `double` is `float64`,
+`complex` is `complex128`; fixed-width types such as `int64_t` and
+`size_t` work too. `xp.cuda.ctype_of(np.float64)` returns `"double"`, useful
+when generating source, and `xp.cuda.parse_cuda_signature()` returns the parsed
+parameters.
+
+Scalars in detail: a Python `int` is cast into integer (range-checked),
+floating-point and complex parameters, a `float` into floating-point and
+complex parameters, a `bool` into boolean and integer parameters. A NumPy
+scalar is passed as it is if its dtype matches, cast if the cast is safe
+(`np.float32` into `double`), and rejected otherwise (`np.float64` into
+`float`); NumPy integers are checked by value like Python ints, so
+`np.int64(5)` fits an `int` parameter and `np.int64(2**31)` raises
+`OverflowError`.
A wrong argument count, dtype or layout raises before anything is launched,
with the parameter name in the message:
@@ -75,9 +87,13 @@ because the kernel would read them as a flat buffer. Make them contiguous with
`xp.ascontiguousarray()` (a copy), or use an array view parameter (below),
which carries the strides.
-The checks cost about 0.3 µs per argument. For tiny kernels called millions of
-times, `check_signature=False` disables them once the calls are known to be
-correct; the kernel then behaves like a raw `RawKernel`.
+The checks cost about 0.3 µs per argument: about 10 µs for a kernel with 29
+arguments, measured on an H100 node, where the launch itself costs about as
+much. That is negligible for kernels that run for 100 µs or more. For tiny
+kernels called millions of times, `check_signature=False` disables them once the
+calls are known to be correct (also needed for signatures the parser cannot
+read, e.g. with macros or pointers to pointers); the kernel then behaves like a
+raw `RawKernel`.
## Launch configuration
@@ -185,7 +201,10 @@ Pass extra include directories with `include_dirs=[...]` and NVRTC flags with
| `` | Morton (Z-order) keys `cunumpy_morton_key2(x, y, ...)`, `_key3`, equal to `xp.algorithms.morton_keys` on the host |
| `` | counter-based random numbers `cunumpy_uniform(seed, stream, counter)`, `cunumpy_normal2(...)`, equal to `xp.rng.philox_uniform` on the host |
-`xp.cuda.cuda_include_dir()` returns their directory for use with other compilers.
+`` and `` add warp and block reductions
+and prefix sums. [CUDA headers](cuda-headers.md) documents every function of
+these headers. `xp.cuda.cuda_include_dir()` returns their directory for use
+with other compilers.
CuPy's disk cache is keyed on the source string and the options only, so
editing an included header would normally *not* trigger a recompile.
@@ -193,7 +212,17 @@ editing an included header would normally *not* trigger a recompile.
hashes their contents, and adds `-DCUNUMPY_INCLUDE_HASH=0x...` to the options,
so a changed header produces a new cache entry. `kernel.included_headers` lists
the resolved files and `kernel.compile_options()` the options actually used.
-System includes in angle brackets are not hashed.
+
+A quoted include is looked up like NVRTC does: relative to the including file
+(`source_dir` for the kernel source, the header's own directory for nested
+includes), then in `include_dirs` in order, then in CuNumpy's header directory.
+CuNumpy's own headers are tracked also when included in angle brackets
+(`#include `), so upgrading CuNumpy with a changed header
+recompiles the kernels that use it. Other system includes in angle brackets,
+and includes that cannot be found (NVRTC reports those), are not hashed. The
+list is recomputed at every access, so it follows the files on disk; a source
+without includes never touches the file system. `xp.cuda.resolve_includes()`
+and `xp.cuda.include_hash()` are the two building blocks.
## Index macros and array views
@@ -224,7 +253,9 @@ An `ArrayND` parameter takes a CuPy array of dtype `T` and `N` dimensions; it
is passed by value as pointer, shape and strides (in elements). `a(i, j)`
returns a reference, `a.shape[k]` and `a.size()` give the extents. Compiling
with `-DCUNUMPY_BOUNDS_CHECK` (automatic in [debug mode](debugging.md)) checks
-every index against the shape and traps with a message on violation.
+every index against the shape and traps with a message on violation. The
+packed layout of the views, the contiguous `CArrayND` variants and the view
+fields of a `CudaStruct` are described in [CUDA headers](cuda-headers.md).
The index macros remove the boilerplate of thread index computation:
diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md
index 61d7ce7..8ec33cd 100644
--- a/docs/source/kernels/dispatch.md
+++ b/docs/source/kernels/dispatch.md
@@ -200,12 +200,18 @@ with xp.kernels.use_device_kernel_implementation(None):
xp.kernels.set_device_kernel_implementation(None) # restore automatic selection
```
-`CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION=cuda` sets the same choice at import.
-Unsupported values raise `ValueError`. An explicit CUDA choice raises
+`CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION=cuda` sets the same choice at import (the
+value is case-insensitive and stripped of whitespace; an unsupported value makes
+the import fail). The setter raises `ValueError` for unsupported values and
+keeps the previous setting. The getter reports the requested setting, so it
+returns `None` in automatic mode even when calls use CUDA. An explicit CUDA choice raises
`LookupError` for a missing device implementation, including kernels configured
with `missing_cuda="fallback"`. The default (`None`, or an unset/empty environment
-variable) preserves that fallback policy. This setting controls device dispatch;
-the array backend and host implementation are selected independently.
+variable) preserves that fallback policy. This setting controls the device
+dispatch of `Kernel` and `KernelCatalog` only: the array backend and host
+implementation are selected independently, and direct `CudaKernel` (or CuPy
+`RawKernel`) calls are not affected. Like the other selections it is global,
+not per thread.
## Compiled Pyccel host kernels
@@ -229,11 +235,19 @@ def compile_kernels(module):
catalog = xp.kernels.KernelCatalog.from_package(
__name__,
host_suffix="_pyccel", # push/push_pyccel.py next to push/push_cuda.cu
+ dispatch="arrays", # see "Choosing the kernel by where the arrays are"
compile_host=compile_kernels,
host_fallback=NUMPY_VERSIONS,
)
```
+Without `compile_host` the plain Python functions are called; without
+`host_fallback` a kernel whose compilation fails runs its uncompiled Python
+function, with a warning. `Kernel.from_folder()` takes `compile_host` too (pass
+a fallback there as `extra_implementations={"numpy": ...}`), and
+`CompiledHostKernel(module, name, compiler, fallback=...)` wraps a single host
+kernel by hand.
+
`host_fallback` becomes each kernel's `"numpy"` implementation (see "Several
host implementations" above); `catalog["push"].host_kernel.kernel.available("pyccel")`
reports whether the compiled version builds, and `catalog["push"].selected()`
@@ -263,7 +277,9 @@ catalog["gather"](host_positions, host_field, host_result) # host kernel, also
The CUDA kernel runs if any top-level argument is on the GPU: a CuPy array, or
a CUDA argument object (`CudaArguments`, `CudaStructArguments`, a struct
-value). Host arguments go to the host function directly, without conversion.
+value). Host arguments go to the host function directly: the call bypasses the
+device-array conversion of `PyccelKernel`, so it costs nothing extra. On device
+arguments without a CUDA kernel, `missing_cuda` applies as usual.
The arguments of one call must then all be on one side, with the dtype and
layout the kernels take. `xp.kernels.as_kernel_array(value, like, dtype)` brings an
diff --git a/docs/source/kernels/pyccel-kernel.md b/docs/source/kernels/pyccel-kernel.md
index 05315d6..fba552f 100644
--- a/docs/source/kernels/pyccel-kernel.md
+++ b/docs/source/kernels/pyccel-kernel.md
@@ -87,11 +87,12 @@ other modules are passed through unchanged.
Other details:
-* **Aliasing is preserved.** If the same device array appears twice, the host
- function sees one host array twice.
+* **Aliasing is preserved.** If the same device array appears twice, also
+ inside containers, the host function sees one host array twice. Reference
+ cycles among supported containers and objects are handled.
* **Return values**: NumPy arrays (also inside returned tuples and lists) are
converted to CuPy; `is_array` customizes which return values count as arrays
- (for NumPy subclasses).
+ (for NumPy subclasses). Dictionaries in return values are not converted.
* **`use_cupy`** forces conversion on (`True`) or off (`False`); the default
`None` decides per call.
diff --git a/docs/source/kernels/testing.md b/docs/source/kernels/testing.md
index 7922c27..ed6a1f6 100644
--- a/docs/source/kernels/testing.md
+++ b/docs/source/kernels/testing.md
@@ -16,7 +16,9 @@ from cunumpy.kernel_testing import (
```
`cunumpy.kernel_testing` is not imported by `import cunumpy`, and it imports pytest only
-when one of its pytest objects is used.
+when one of its pytest objects is used, so `device_function_kernel` and the
+emulation work without pytest. It is not called `cunumpy.testing`, because a
+submodule of that name would replace NumPy's `xp.testing` once imported.
## Run a test on both backends
@@ -222,7 +224,34 @@ intrinsics; kernels using the latter are refused with `NotImplementedError`, so
those still need a GPU run. Inline PTX is not emulated either: `asm(...)` and
`asm volatile(...)` compile but trap when reached, so a PTX branch a test never
takes needs no extra options, and reaching it raises `RuntimeError`. The compiler may fuse multiply-adds as NVRTC does,
-so compare with a tolerance of a few ulp.
+so compare with a tolerance of a few ulp, or pass `options=("-ffp-contract=off",)`
+for NumPy's rounding.
+
+How the emulation works, for when a result surprises you:
+
+* The source is compiled as C++17 with `CXX` (else `c++`, see
+ `emulation_compiler()`), with `threadIdx`, `blockIdx`, `blockDim`, `gridDim`,
+ the atomics (`atomicAdd`, `atomicMin`, ..., as plain operations), `__ldg`,
+ `rsqrt` and `__trap` replaced. The shipped headers, the kernel's include
+ directories and its `-D` options apply. `CUNUMPY_EMULATION_CACHE` sets the
+ disk cache (`0`: none).
+* The library runs in the Python process, on the arrays themselves. An array
+ that is not writeable, or whose layout a parameter cannot take, is passed as
+ a copy and written back, so aliased arguments see each other's writes.
+* `__shared__` variables exist once per block (blocks run one after another),
+ `extern __shared__` arrays point into a buffer of `shared_mem` bytes, and in a
+ kernel that calls `__syncthreads` the threads of a block run as coroutines
+ (POSIX `ucontext`, each with its own stack), so every thread reaches a barrier
+ before any continues past it.
+* Not emulated: concurrency between barriers (races and atomic ordering never
+ show), warp intrinsics (`__syncwarp`, shuffles, votes: `NotImplementedError`,
+ also from an included header), complex scalars and inline PTX.
+* A kernel that does not compile, or traps (an out-of-bounds index with
+ `-DCUNUMPY_BOUNDS_CHECK`, `__trap()`, inline `asm`), raises `RuntimeError`.
+ Other crashes, such as a segmentation fault from an unchecked out-of-bounds
+ index, end the process: run such tests through
+ `run_in_fake_cupy_subprocess()`, which prints the traceback with
+ `faulthandler`.
### Struct parameters and `emulated_launches()`
@@ -268,7 +297,7 @@ backend too:
from cunumpy.kernel_testing import fake_cupy_session, requires_device_backend
-@requires_device_backend # a GPU, or the fake CuPy
+@requires_device_backend # a GPU, or the fake CuPy (device_backend_available())
def test_simulation_on_the_cupy_backend():
with fake_cupy_session(): # only with the fake CuPy
sim.run()
@@ -320,9 +349,25 @@ def test_find_span_matches_host():
np.testing.assert_array_equal(xp.to_numpy(out), expected)
```
+The prototype above generates this kernel, named `_kernel` unless
+`name=` is given:
+
+```c
+extern "C" __global__ void find_span_kernel(
+ const double* t, const int* p, const double* eta, int* out, int n)
+{
+ int i = blockDim.x * blockIdx.x + threadIdx.x;
+ if (i >= n) return;
+ out[i] = find_span(t, p[i], eta[i]);
+}
+```
+
In the generated kernel, pointer parameters are passed unchanged to every
thread (shared data), scalar parameters become per-thread arrays, the return
-value of thread `i` goes to `out[i]`, and `n` is the number of elements. A
+value of thread `i` goes to `out[i]` (a `void` function has no `out`), and `n`
+is the number of elements. A parameter of the function named `out` or `n`
+raises `ValueError`; rename the generated ones with `out_param=` and
+`n_threads_param=`. A
struct parameter, by value or by `const` reference (`const DomainArgs& d`),
is passed through unchanged as well; give the struct types in `structs=`
and pass a `CudaStructArguments` object or a packed value, so helpers that
diff --git a/docs/source/reference/algorithms.md b/docs/source/reference/algorithms.md
new file mode 100644
index 0000000..1f1d9fb
--- /dev/null
+++ b/docs/source/reference/algorithms.md
@@ -0,0 +1,15 @@
+# Algorithms: `cunumpy.algorithms`
+
+Morton keys, sorting by key, segment sums and compaction.
+
+```{eval-rst}
+.. automodule:: cunumpy.algorithms
+ :members:
+ :special-members: __call__
+```
+
+## Constants
+
+```{eval-rst}
+.. autodata:: cunumpy.algorithms.MAX_MORTON_LEVELS
+```
diff --git a/docs/source/reference/arguments.md b/docs/source/reference/arguments.md
new file mode 100644
index 0000000..a6fe098
--- /dev/null
+++ b/docs/source/reference/arguments.md
@@ -0,0 +1,9 @@
+# Kernel arguments: `cunumpy.arguments`
+
+Structs and argument objects passed to CUDA kernels.
+
+```{eval-rst}
+.. automodule:: cunumpy.arguments
+ :members:
+ :special-members: __call__
+```
diff --git a/docs/source/reference/cuda.md b/docs/source/reference/cuda.md
new file mode 100644
index 0000000..d3127e2
--- /dev/null
+++ b/docs/source/reference/cuda.md
@@ -0,0 +1,16 @@
+# CUDA runtime: `cunumpy.cuda`
+
+Devices, memory, streams, events, debug mode and CUDA source tools. These do nothing on the NumPy backend.
+
+```{eval-rst}
+.. automodule:: cunumpy.cuda
+ :members:
+ :special-members: __call__
+```
+
+## Constants
+
+```{eval-rst}
+.. autodata:: cunumpy.cuda.DEBUG_OPTIONS
+.. autodata:: cunumpy.cuda.DEFAULT_SHARED_MEMORY_PER_BLOCK
+```
diff --git a/docs/source/reference/index.md b/docs/source/reference/index.md
new file mode 100644
index 0000000..edf425b
--- /dev/null
+++ b/docs/source/reference/index.md
@@ -0,0 +1,22 @@
+# API reference
+
+Generated from the docstrings. Each name has one import path, through the
+module it is listed under (`xp.kernels.CudaKernel`, `xp.rng.random_streams`).
+For an overview of the module layout and the NumPy-like namespace, see the
+[API overview](../api.md).
+
+```{toctree}
+:maxdepth: 1
+
+top-level
+kernels
+arguments
+cuda
+rng
+algorithms
+mpi
+profiling
+memory
+petsc
+kernel-testing
+```
diff --git a/docs/source/reference/kernel-testing.md b/docs/source/reference/kernel-testing.md
new file mode 100644
index 0000000..c34a1ee
--- /dev/null
+++ b/docs/source/reference/kernel-testing.md
@@ -0,0 +1,9 @@
+# Kernel testing: `cunumpy.kernel_testing`
+
+Pytest helpers for host/CUDA kernel pairs. Not imported by `import cunumpy`.
+
+```{eval-rst}
+.. automodule:: cunumpy.kernel_testing
+ :members:
+ :special-members: __call__
+```
diff --git a/docs/source/reference/kernels.md b/docs/source/reference/kernels.md
new file mode 100644
index 0000000..a465618
--- /dev/null
+++ b/docs/source/reference/kernels.md
@@ -0,0 +1,16 @@
+# Kernels: `cunumpy.kernels`
+
+Kernels that run on either backend: dispatch, host implementations, CUDA and Metal kernels, fusion.
+
+```{eval-rst}
+.. automodule:: cunumpy.kernels
+ :members:
+ :special-members: __call__
+```
+
+## Constants
+
+```{eval-rst}
+.. autodata:: cunumpy.kernels.HOST_IMPLEMENTATIONS
+.. autodata:: cunumpy.kernels.DEVICE_IMPLEMENTATIONS
+```
diff --git a/docs/source/reference/memory.md b/docs/source/reference/memory.md
new file mode 100644
index 0000000..818abba
--- /dev/null
+++ b/docs/source/reference/memory.md
@@ -0,0 +1,9 @@
+# Memory: `cunumpy.memory`
+
+Pinned host staging, asynchronous host copies and device mirrors.
+
+```{eval-rst}
+.. automodule:: cunumpy.memory
+ :members:
+ :special-members: __call__
+```
diff --git a/docs/source/reference/mpi.md b/docs/source/reference/mpi.md
new file mode 100644
index 0000000..afa6be9
--- /dev/null
+++ b/docs/source/reference/mpi.md
@@ -0,0 +1,16 @@
+# MPI: `cunumpy.mpi`
+
+MPI buffers, CUDA-aware MPI detection and synchronization.
+
+`cunumpy.mpi` also re-exports the MPI stand-ins of
+[maybempi](https://max-models.github.io/maybempi/), documented there:
+`get_mpi`, `is_serial`, `launched_under_mpi`, `local_rank`,
+`OVERRIDE_VARIABLE`, `SerialMPI`, `SerialComm`, `SerialRequest` and
+`SerialStatus`.
+
+```{eval-rst}
+.. automodule:: cunumpy.mpi
+ :members:
+ :special-members: __call__
+ :exclude-members: get_mpi, is_serial, launched_under_mpi, local_rank, OVERRIDE_VARIABLE, SerialMPI, SerialComm, SerialRequest, SerialStatus
+```
diff --git a/docs/source/reference/petsc.md b/docs/source/reference/petsc.md
new file mode 100644
index 0000000..39fa44a
--- /dev/null
+++ b/docs/source/reference/petsc.md
@@ -0,0 +1,9 @@
+# PETSc: `cunumpy.petsc`
+
+PETSc vectors that share memory with arrays.
+
+```{eval-rst}
+.. automodule:: cunumpy.petsc
+ :members:
+ :special-members: __call__
+```
diff --git a/docs/source/reference/profiling.md b/docs/source/reference/profiling.md
new file mode 100644
index 0000000..a3da2d7
--- /dev/null
+++ b/docs/source/reference/profiling.md
@@ -0,0 +1,9 @@
+# Profiling: `cunumpy.profiling`
+
+Timed regions, NVTX ranges and host/device transfer counting.
+
+```{eval-rst}
+.. automodule:: cunumpy.profiling
+ :members:
+ :special-members: __call__
+```
diff --git a/docs/source/reference/rng.md b/docs/source/reference/rng.md
new file mode 100644
index 0000000..16c57b0
--- /dev/null
+++ b/docs/source/reference/rng.md
@@ -0,0 +1,16 @@
+# Random numbers: `cunumpy.rng`
+
+Reproducible random streams and counter-based Philox generators.
+
+```{eval-rst}
+.. automodule:: cunumpy.rng
+ :members:
+ :special-members: __call__
+```
+
+## Constants
+
+```{eval-rst}
+.. autodata:: cunumpy.rng.BIT_GENERATORS
+.. autodata:: cunumpy.rng.random_streams
+```
diff --git a/docs/source/reference/top-level.md b/docs/source/reference/top-level.md
new file mode 100644
index 0000000..52b6248
--- /dev/null
+++ b/docs/source/reference/top-level.md
@@ -0,0 +1,15 @@
+# Top level: `cunumpy`
+
+The NumPy/CuPy namespace (see the [API overview](../api.md)) plus backend selection, array inspection and conversion, and host execution helpers.
+
+```{eval-rst}
+.. automodule:: cunumpy
+ :members:
+ :special-members: __call__
+```
+
+## Constants
+
+```{eval-rst}
+.. autodata:: cunumpy.scipy
+```
diff --git a/pyproject.toml b/pyproject.toml
index 285ef40..635a1f5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,7 @@ optional-dependencies.docs = [
"myst-parser",
"nbconvert",
"nbsphinx",
+ "numpydoc",
"jupyterlab",
"pre-commit",
"pyproject-fmt",
@@ -58,8 +59,33 @@ cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ]
# agent worktrees and scratch scripts are local copies, not part of the project
extend-exclude = [ ".claude" ]
# formatting: `ruff format`; import sorting: the isort rules of `ruff check`
-lint.extend-select = [ "I", "TID252" ]
+lint.extend-select = [ "D", "I", "TID252" ]
+# D105/D107: magic methods and __init__ (constructor parameters go in the class docstring)
+lint.extend-ignore = [ "D105", "D107" ]
lint.isort.known-first-party = [ "cunumpy" ]
+lint.per-file-ignores."_fake_cupy*.py" = [ "D" ]
+lint.per-file-ignores."conftest.py" = [ "D" ]
+lint.per-file-ignores."docs/**" = [ "D" ]
+lint.per-file-ignores."tests/**" = [ "D" ]
+lint.pydocstyle.convention = "numpy"
[tool.ruff.lint.flake8-tidy-imports]
ban-relative-imports = "all"
+
+# `python -m numpydoc lint src/cunumpy/*.py`: the public API (names re-exported by a public
+# module) has full numpydoc sections; private helpers only need a one-line summary.
+[tool.numpydoc_validation]
+checks = [
+ "all",
+ "ES01", # extended summary is optional (docstrings are kept short)
+ "EX01", # examples are required for public functions and classes, not every method
+ "SA01", # See Also is optional
+ "GL08", # undocumented objects are caught by ruff's D rules instead
+]
+# one-line docstrings (properties, trivial methods) need no Returns section
+override_RT01 = [ '\A[^\n]*\Z' ]
+exclude = [
+ '\._[A-Za-z]\w*$', # private functions, classes and methods
+ '\.__\w+__$', # dunder methods
+ '_fake_cupy',
+]
diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py
index f4db38d..a202070 100644
--- a/src/cunumpy/__init__.py
+++ b/src/cunumpy/__init__.py
@@ -1,3 +1,16 @@
+"""NumPy and CuPy behind one namespace.
+
+``import cunumpy as xp`` gives the functions of the active backend's
+``array-api-compat`` module (``xp.zeros``, ``xp.linalg``, ...), replaced when
+the backend changes, plus backend selection, array inspection and conversion,
+host execution helpers (:func:`host_call`), and :data:`scipy`. Everything else
+is in a submodule (:mod:`cunumpy.kernels`, :mod:`cunumpy.cuda`, ...). See
+:doc:`/api` for the module layout and :doc:`/guides/backends`.
+
+``xp.numpy_backend`` and ``xp.cupy_backend`` are booleans telling whether the
+active backend is NumPy or CuPy (prefer :func:`get_backend` for the name).
+"""
+
# cunumpy/__init__.py
import re as _re
from importlib.metadata import PackageNotFoundError, version
@@ -38,6 +51,7 @@
)
try:
+ #: The installed version of cunumpy; ``"0.0.0+unknown"`` without package metadata.
__version__ = version("cunumpy")
except PackageNotFoundError:
__version__ = "0.0.0+unknown"
@@ -48,18 +62,27 @@ def _version_key(text: str) -> tuple[int, ...]:
def require_version(minimum: str) -> None:
- """Raise ``ImportError`` if this cunumpy is older than `minimum`.
-
- For projects that depend on a feature of a given release, as a clearer
- error than an ``AttributeError`` later::
-
- import cunumpy as xp
-
- xp.require_version("0.4.0")
-
- Only the numeric part of the versions is compared (``0.4.0`` and
- ``0.4.0.dev1`` compare equal). Nothing is checked when the installed
- version is unknown (cunumpy not installed as a package).
+ """Raise ``ImportError`` if the installed cunumpy is older than `minimum`.
+
+ A clearer error than a later ``AttributeError`` for code that needs a
+ feature of a given release. Only the numeric part is compared (``0.4.0``
+ and ``0.4.0.dev1`` are equal); nothing is checked when the installed
+ version is unknown (cunumpy not installed as a package, ``xp.__version__``
+ is then ``"0.0.0+unknown"``).
+
+ Parameters
+ ----------
+ minimum : str
+ The oldest acceptable version, e.g. ``"0.4.0"``.
+
+ Raises
+ ------
+ ImportError
+ If the installed version is older than `minimum`.
+
+ Examples
+ --------
+ >>> xp.require_version("0.4.0")
"""
if __version__.startswith("0.0.0+unknown"):
return
@@ -111,12 +134,7 @@ def require_version(minimum: str) -> None:
def __getattr__(name: str):
- """Resolve names that are not in the namespace: cunumpy. -> cunumpy.xp..
-
- The public names of the active backend are copied into this namespace (see
- `_sync_backend_namespace`), so this only runs for names missing from the
- backend's ``__all__`` and for ``numpy_backend``/``cupy_backend``.
- """
+ """Resolve ``numpy_backend``, ``cupy_backend`` and backend names not copied here."""
if name == "numpy_backend":
return xp.numpy_backend
if name == "cupy_backend":
diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py
index 5127b64..b5cd42c 100644
--- a/src/cunumpy/_algorithms.py
+++ b/src/cunumpy/_algorithms.py
@@ -36,12 +36,39 @@ def _sorted_keys(keys: Any) -> tuple[Any, Any]:
def segment_boundaries(sorted_keys: Any) -> tuple[Any, Any, Any]:
- """Return ``(unique_keys, starts, stops)`` for runs of sorted integer keys.
+ """Return the runs of equal keys in sorted integer keys.
- Starts/stops are int64 indices and describe half-open slices of the input.
- Sparse, negative, and uint64 Morton keys are supported. All results stay on
- the input backend. Validation and variable-length GPU output may synchronize;
- prepare once and reuse the boundaries in repeated operations.
+ Sparse, negative and uint64 Morton keys are supported. The results stay on
+ the backend of `sorted_keys`. Validation and the variable-length output may
+ synchronize on CUDA: compute the boundaries once and reuse them.
+
+ Parameters
+ ----------
+ sorted_keys : array of int, shape (n,)
+ Keys sorted in nondecreasing order.
+
+ Returns
+ -------
+ unique_keys : array of int
+ The distinct keys.
+ starts, stops : array of int64
+ Run ``i`` is the half-open slice ``[starts[i], stops[i])`` of the input.
+
+ Raises
+ ------
+ ValueError
+ If `sorted_keys` is not 1D or not sorted.
+ TypeError
+ If `sorted_keys` is not an integer array.
+
+ See Also
+ --------
+ cell_offsets : Dense offsets for every cell, also the empty ones.
+
+ Examples
+ --------
+ >>> xp.algorithms.segment_boundaries(xp.asarray([3, 3, 5, 9, 9, 9]))
+ (array([3, 5, 9]), array([0, 2, 3]), array([2, 3, 6]))
"""
xpm, keys = _sorted_keys(sorted_keys)
if keys.size == 0:
@@ -54,11 +81,37 @@ def segment_boundaries(sorted_keys: Any) -> tuple[Any, Any, Any]:
def cell_offsets(sorted_cells: Any, n_cells: int) -> Any:
- """Dense int64 offsets into sorted cell IDs in ``[0, n_cells)``.
+ """Return the offsets of each cell into sorted cell IDs.
+
+ Cell ``k`` occupies ``[offsets[k], offsets[k + 1])``; empty cells have equal
+ offsets. The result stays on the backend of `sorted_cells`; validation may
+ synchronize on CUDA. Filter out negative (invalid) cell IDs first.
+
+ Parameters
+ ----------
+ sorted_cells : array of int
+ Cell IDs in ``[0, n_cells)``, sorted in nondecreasing order.
+ n_cells : int
+ Number of cells.
- Cell k occupies ``[offsets[k], offsets[k+1])``; empty cells have equal
- offsets. Returns n_cells+1 entries on the input backend. Filter negative
- (invalid) cell IDs before calling. Validation may synchronize on CUDA.
+ Returns
+ -------
+ array of int64
+ ``n_cells + 1`` offsets.
+
+ Raises
+ ------
+ ValueError
+ If `sorted_cells` is not sorted or has an ID outside ``[0, n_cells)``.
+
+ See Also
+ --------
+ segment_boundaries : The runs of sorted keys that are present.
+
+ Examples
+ --------
+ >>> xp.algorithms.cell_offsets(xp.asarray([0, 0, 2]), 3)
+ array([0, 2, 2, 3])
"""
n_cells = _segment_count(n_cells)
xpm, cells = _sorted_keys(sorted_cells)
@@ -95,12 +148,37 @@ def _sum_kernel(dtype: Any) -> Any:
class SegmentPlan:
- """Snapshot and validate segment keys once, then reuse :meth:`sum`.
+ """Validated segment keys, reused by repeated :meth:`sum` calls.
- A negative key drops its row; other keys must be below `n_segments`.
- Keys are copied so later caller mutation cannot invalidate the plan. A GPU
- plan is bound to its CUDA device. Setup may synchronize, but repeated sums
- do not read GPU reductions back into Python or validate keys per column.
+ The keys are copied, so later changes to the caller's array do not affect the
+ plan. Setup may synchronize on CUDA; repeated sums do not validate the keys
+ again or read GPU values back to the host. A GPU plan is bound to its CUDA
+ device, which must be current when it is created and used.
+
+ Parameters
+ ----------
+ keys : array of int, shape (n,)
+ Segment of every row; a negative key drops the row.
+ n_segments : int
+ Number of segments; every key must be smaller.
+
+ Raises
+ ------
+ ValueError
+ If `keys` is not 1D, a key is ``>= n_segments``, or the keys' CUDA device
+ is not current.
+ TypeError
+ If `keys` is not an integer array.
+
+ See Also
+ --------
+ segment_sum : The same for a single use of the keys.
+
+ Examples
+ --------
+ >>> plan = xp.algorithms.SegmentPlan(xp.asarray([0, 1, 0, -1]), 2)
+ >>> plan.sum(xp.asarray([1.0, 2.0, 3.0, 4.0]))
+ array([4., 2.])
"""
def __init__(self, keys: Any, n_segments: int) -> None:
@@ -120,17 +198,37 @@ def __init__(self, keys: Any, n_segments: int) -> None:
@property
def n_segments(self) -> int:
- """Number of output segments, fixed when the plan is prepared."""
+ """Number of output segments, fixed when the plan is created."""
return self._n_segments
def sum(self, values: Any, *, out: Any = None) -> Any:
- """Sum rows into ``(n_segments, *values.shape[1:])``; overwrite `out`.
-
- Floating/complex dtypes are preserved; integer/bool inputs produce
- float64. CUDA supports float16/32/64 and complex64/128 outputs; float16
- accumulates in a float32 workspace. Other CUDA accumulation
- uses atomics and its floating-point order is not deterministic. `out`
- must have the exact shape/dtype, be C-contiguous, and not alias values.
+ """Sum the rows of `values` per segment.
+
+ Floating-point and complex dtypes are kept; integer and bool values give
+ float64. CUDA supports float16/32/64 and complex64/128 (float16 accumulates
+ in float32) and sums all components in one launch with atomics, so the
+ floating-point order is not deterministic.
+
+ Parameters
+ ----------
+ values : array, shape (n, ...)
+ One row per key, on the backend of the keys.
+ out : array, optional
+ Output to overwrite: exact shape and dtype, writable, C-contiguous, and
+ not aliasing `values`.
+
+ Returns
+ -------
+ array, shape (n_segments, ...)
+ The sums, `out` if given.
+
+ Raises
+ ------
+ ValueError
+ If the shapes do not match, `out` is unsuitable, or the CUDA device is
+ not the plan's.
+ TypeError
+ If `values` is not numeric, or its dtype is not supported on CUDA.
"""
assert_same_backend(self._keys, values)
xpm = self._xp
@@ -202,12 +300,34 @@ def sum(self, values: Any, *, out: Any = None) -> Any:
def segment_sum(values: Any, keys: Any, n_segments: int, *, out: Any = None) -> Any:
- """Sum rows per integer key, dropping negatives; optionally overwrite `out`.
+ """Sum the rows of `values` per integer key.
- Accepts arbitrary trailing value dimensions. Keys are validated once per
- call and CUDA reduces all components in one launch, without per-column host
- checks. Use ``SegmentPlan(keys, n_segments).sum(values, out=out)`` when the
- keys are reused, to avoid repeating setup and its GPU synchronization.
+ ``out[k]`` is the sum of the rows ``i`` with ``keys[i] == k``; negative keys
+ drop their row. The result is on the backend of `keys`; mixed backends raise.
+ Use :class:`SegmentPlan` when the keys are reused, to avoid repeating the
+ validation and its GPU synchronization.
+
+ Parameters
+ ----------
+ values : array, shape (n, ...)
+ Values with any trailing dimensions.
+ keys : array of int, shape (n,)
+ Segment of every row.
+ n_segments : int
+ Number of segments; every key must be smaller.
+ out : array, optional
+ Output to overwrite, see :meth:`SegmentPlan.sum`.
+
+ Returns
+ -------
+ array, shape (n_segments, ...)
+ The sums, float64 for integer or bool values.
+
+ Examples
+ --------
+ >>> keys = xp.asarray([0, 1, 0, -1])
+ >>> xp.algorithms.segment_sum(xp.asarray([1.0, 2.0, 3.0, 4.0]), keys, 2)
+ array([4., 2.])
"""
return SegmentPlan(keys, n_segments).sum(values, out=out)
@@ -217,14 +337,7 @@ def segment_sum(values: Any, keys: Any, n_segments: int, *, out: Any = None) ->
def _radix_argsort(keys: np.ndarray) -> np.ndarray:
- """Stable argsort of integer `keys`: least significant 16 bits first.
-
- NumPy's stable argsort is a radix sort for 16-bit integers only; wider
- integers get a timsort, about ten times slower on a million random keys.
- Sorting the 16-bit digits of ``keys - keys.min()`` one after the other,
- each pass stable, gives the same order (an LSD radix sort), in as many
- passes as the range of the keys needs (two for up to 2**32 cells).
- """
+ """Argsort integer `keys` stably, by an LSD radix sort on 16-bit digits."""
low = keys.min()
span = int(keys.max()) - int(low)
if keys.dtype.kind == "i":
@@ -250,42 +363,40 @@ def _radix_argsort(keys: np.ndarray) -> np.ndarray:
def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
"""Sort `keys` and reorder every array the same way, in one stable argsort.
- The usual first step of a particle code on the GPU: sort the particles by
- cell index or Morton key (:func:`cunumpy.algorithms.morton_keys`), then work on
- contiguous ranges. The sort is stable, so equal keys keep their order and
- the result is reproducible::
-
- keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, charges)
-
- Component-major marker arrays, ``(ncomp, N)`` next to ``(N,)`` scalars, keep
- the markers along the last axis of each, so sort along that one::
-
- keys, order, positions, weights = xp.algorithms.sort_by_key(
- keys, positions, weights, axis=-1
- )
-
- On NumPy, integer keys (cell indices, Morton keys) are sorted by a radix
- sort on their 16-bit digits, as on CuPy: two passes for up to ``2**32``
- distinct cells, about ten times faster than NumPy's stable sort of 64-bit
- integers.
+ The usual first step of a particle code: sort the particles by cell index or
+ Morton key (:func:`morton_keys`), then work on contiguous ranges. Equal keys
+ keep their order, so the result is reproducible. On CuPy the argsort is a
+ radix sort; on NumPy, integer keys of more than 4096 entries get an LSD
+ radix sort too (one stable ``uint16`` argsort per 16-bit digit the key range
+ needs), about ten times faster than NumPy's stable sort of 64-bit integers.
Parameters
----------
keys : array, shape (n,)
The sort keys.
- *arrays : arrays
- Arrays with ``n`` entries along `axis` (any other axes), on the backend
- of `keys`.
- axis : int
- The axis of every array that `keys` indexes, the first by default;
- ``-1`` is the last axis of each array, whatever its number of axes.
+ *arrays : array
+ Arrays with ``n`` entries along `axis`, on the backend of `keys`.
+ axis : int, optional
+ The axis of every array that `keys` indexes, 0 by default; ``-1`` is the
+ last axis of each array, for component-major ``(ncomp, n)`` markers next
+ to ``(n,)`` arrays.
Returns
-------
- tuple
- ``(sorted_keys, order, *sorted_arrays)``: ``order`` (int64) is the
- permutation, ``sorted_keys = keys[order]``, and each sorted array is
- ``take(array, order, axis=axis)`` (a new array).
+ tuple of arrays
+ ``(keys[order], order, *sorted_arrays)``, with the int64 permutation
+ ``order`` and each array taken along `axis` (new arrays).
+
+ Raises
+ ------
+ ValueError
+ If `keys` is not 1D or an array does not have ``n`` entries along `axis`.
+
+ Examples
+ --------
+ >>> keys = xp.asarray([2, 0, 1, 0])
+ >>> xp.algorithms.sort_by_key(keys, xp.asarray([10.0, 11.0, 12.0, 13.0]))
+ (array([0, 0, 1, 2]), array([1, 3, 2, 0]), array([11., 13., 12., 10.]))
"""
if get_array_backend(keys) == "cupy":
import cupy as xpm # its argsort is a stable radix sort
@@ -324,39 +435,42 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int:
"""Move the entries where `mask` is True to the front of every array, in place.
- The usual step after particles left the domain or were absorbed: keep the
- live ones at the front of the marker array (and of the arrays that go with
- it) and continue with ``markers[:n]``. The order of the kept entries is
- preserved, so the result is reproducible::
-
- n = xp.algorithms.compact_by_mask(alive, markers, weights)
- markers, weights = markers[:n], weights[:n]
-
- Component-major marker arrays, ``(ncomp, N)`` next to ``(N,)`` scalars, keep
- the markers along the last axis of each, so compact that one::
-
- n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1)
- positions, weights = positions[:, :n], weights[:n]
+ The usual step after particles left the domain: keep the live ones at the
+ front, in their original order, and continue with ``markers[:n]``. The
+ entries after the first ``n`` are unspecified. The count is needed on the
+ host, so on CuPy each call synchronizes once (seen by
+ :func:`~cunumpy.profiling.count_transfers`).
Parameters
----------
mask : array of bool, shape (n,)
True for the entries to keep.
- *arrays : arrays
- Arrays with ``n`` entries along `axis` (any other axes), on the backend
- of `mask`. Entries ``[:count]`` along `axis` hold the kept ones
- afterwards; the entries after them are unspecified, so ignore them (or
- overwrite them).
- axis : int
- The axis of every array that `mask` indexes, the first by default;
- ``-1`` is the last axis of each array, whatever its number of axes.
+ *arrays : array
+ Arrays with ``n`` entries along `axis`, on the backend of `mask`;
+ modified in place.
+ axis : int, optional
+ The axis of every array that `mask` indexes, 0 by default; ``-1`` is the
+ last axis of each array, for component-major ``(ncomp, n)`` markers next
+ to ``(n,)`` arrays.
Returns
-------
int
- The number of kept entries. Its value is needed on the host, so on CuPy
- the call synchronizes once per call (counted by
- :func:`~cunumpy.profiling.count_transfers` where it can be seen).
+ The number of kept entries.
+
+ Raises
+ ------
+ TypeError
+ If `mask` is not a 1D boolean array.
+ ValueError
+ If an array does not have ``n`` entries along `axis`.
+
+ Examples
+ --------
+ >>> x = xp.asarray([1.0, 2.0, 3.0, 4.0])
+ >>> n = xp.algorithms.compact_by_mask(xp.asarray([True, False, True, False]), x)
+ >>> x[:n]
+ array([1., 3.])
"""
assert_same_backend(mask, *arrays)
xpm = get_array_module(mask)
diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py
index 80b203d..2b95972 100644
--- a/src/cunumpy/_cuda_kernel.py
+++ b/src/cunumpy/_cuda_kernel.py
@@ -1,49 +1,18 @@
"""CUDA kernels (``cupy.RawKernel``) called like their NumPy/Pyccel counterparts.
-:class:`CudaKernel` wraps a CUDA C kernel so that it can be called with the same
-arguments as the host kernel it mirrors:
-
-* argument objects that implement the :class:`CudaArguments` protocol
- (a ``__cuda_args__()`` method) are flattened into their device arrays and
- scalars, so an object holding several arrays can be passed as one argument;
-* C structs can be passed by value: :class:`CudaStruct` defines the struct once,
- generates its C declaration and packs its values;
-* the ``__global__`` signature is parsed once, and every call is checked against
- it: the number of arguments, the dtype of every array, every struct, and
- every scalar. Python scalars are cast to the declared C type; a scalar that
- does not fit the declared type (a ``float`` for an ``int``, an integer out of
- range, a NumPy scalar that would lose precision) raises instead of reaching
- the kernel as a silently wrong value, which is what ``cupy.RawKernel`` would
- do;
-* arrays are never converted or copied: they must already be C-contiguous
- CuPy arrays (build them once with :func:`cunumpy.as_device_array`);
-* strided array views: a parameter or struct field of type ``Array2D``
- (from the shipped header ``cunumpy/array_view.cuh``, see
- :func:`cuda_include_dir`) takes a 2D CuPy array, contiguous or not, and
- receives its pointer, shape and strides, so that kernels index ``a(i, j)``
- like the pyccel kernels they are ported from; ``CArray2D`` is the
- C-contiguous variant (shape only, ``a(i, j)`` is ``data[i * shape[1] + j]``),
- which takes C-contiguous arrays only and raises for other views;
-* C++ function templates are instantiated with ``template_args``, and generated
- kernels (one source per variant) are compiled once per variant by
- :class:`CudaKernelVariants`.
-
-The launch shape is given at each call, either as the number of threads
-(``n_threads``, in 1 to 3 dimensions) or as an explicit ``grid``.
-
-Argument structs can be generated from the annotations of a Python class
-(:meth:`CudaStruct.from_signature`) and written to a header
-(:meth:`CudaStruct.to_header`, :func:`write_cuda_header`), so that the Python
-class is the one definition of the arguments.
-
-In debug mode (``debug=True``, ``xp.cuda.set_cuda_debug(True)`` or the environment
-variable ``CUNUMPY_CUDA_DEBUG=1``) kernels are compiled with ``-lineinfo`` and
-``-DCUNUMPY_BOUNDS_CHECK``, and every launch is synchronized so that an
-asynchronous CUDA error is raised, as a ``RuntimeError`` naming the kernel, at
-the launch that caused it.
-
-This module imports CuPy only when a kernel is compiled, so it can be imported
-(and signatures parsed) without CuPy.
+:class:`~cunumpy.kernels.CudaKernel` parses the ``__global__`` signature once
+and checks every call against it: argument count, array dtypes and
+contiguity, structs, and scalars (cast to the declared C type, or raising
+instead of reaching the kernel as a wrong value). Arrays are never copied to
+the device. Argument objects (:class:`~cunumpy.arguments.CudaArguments`) are
+flattened, C structs (:class:`~cunumpy.arguments.CudaStruct`) are passed by
+value, and ``Array2D`` parameters take strided CuPy arrays.
+:class:`~cunumpy.kernels.CudaKernelVariants` caches generated kernels per
+variant.
+
+CuPy is imported only when a kernel is compiled, so kernels can be created and
+signatures parsed without CuPy. See :doc:`/kernels/cuda-kernel`,
+:doc:`/kernels/arguments` and :doc:`/kernels/debugging`.
"""
from __future__ import annotations
@@ -122,80 +91,101 @@ class _CompiledKernel:
def cuda_include_dir() -> str:
- """The directory of the CUDA headers shipped with cunumpy.
-
- :class:`CudaKernel` adds it to the include path automatically, so kernels
- can ``#include "cunumpy/array_view.cuh"`` (strided ``Array1D``
- to ``Array16D`` views passed by value) and
- ``#include "cunumpy/index.cuh"`` (thread-index and grid-stride macros such
- as ``CUNUMPY_THREAD_1D(i, n)``), ``#include "cunumpy/atomic.cuh"`` (atomic
- adds) and ``#include "cunumpy/reduce.cuh"`` (warp and block reductions).
- Pass it as ``-I`` to other compilers.
+ """Return the directory of the CUDA headers shipped with cunumpy.
+
+ :class:`~cunumpy.kernels.CudaKernel` always finds these headers:
+ ``cunumpy/array_view.cuh`` (strided ``Array1D`` to ``Array16D``
+ views), ``cunumpy/index.cuh`` (thread-index macros such as
+ ``CUNUMPY_THREAD_1D(i, n)``), ``cunumpy/atomic.cuh`` (atomic adds),
+ ``cunumpy/reduce.cuh`` and ``cunumpy/scan.cuh`` (warp and block
+ reductions and prefix sums), ``cunumpy/random.cuh`` (Philox random
+ numbers) and ``cunumpy/morton.cuh`` (Morton keys). See
+ :doc:`/kernels/cuda-headers`.
+
+ Returns
+ -------
+ str
+ The directory, to pass as ``-I`` to other compilers.
+
+ Examples
+ --------
+ >>> import os
+ >>> os.path.isfile(os.path.join(xp.cuda.cuda_include_dir(), "cunumpy", "atomic.cuh"))
+ True
"""
return str(_CUDA_INCLUDE_DIR)
-#: NVRTC options added in debug mode: source line information for
-#: ``compute-sanitizer``/``nsys``, and bounds checks in the array views.
-#: (``-G`` is not among them: NVRTC does not support it.)
+#: NVRTC options added in debug mode: line information for
+#: ``compute-sanitizer``/``nsys`` and bounds checks in the array views
+#: (no ``-G``: NVRTC does not support it).
DEBUG_OPTIONS = ("-lineinfo", "-DCUNUMPY_BOUNDS_CHECK")
class CudaArguments:
- """Base class for objects passed to a :class:`CudaKernel` as one argument.
+ """Base class for objects passed to a CUDA kernel as several arguments.
- A :class:`CudaKernel` replaces every argument that has a ``__cuda_args__()``
- method by the values it returns, in order. Subclassing this class is
- optional: any object implementing ``__cuda_args__()`` is flattened.
+ A :class:`~cunumpy.kernels.CudaKernel` replaces every argument that has a
+ ``__cuda_args__()`` method by the values it returns, in order. Subclassing
+ is optional: any object with ``__cuda_args__()`` is flattened. See
+ :doc:`/kernels/arguments`.
Parameters
----------
*values
- The CUDA kernel arguments this object stands for: CuPy arrays and
- scalars, in the order of the kernel signature.
+ The kernel arguments this object stands for (CuPy arrays and
+ scalars), in the order of the kernel signature.
Examples
--------
- >>> class Particles(CudaArguments):
+ >>> class Particles(xp.arguments.CudaArguments):
... def __init__(self, positions, velocities):
- ... self.positions = positions
... super().__init__(positions, velocities, positions.shape[0])
- >>> kernel(dt, Particles(x, v), n_threads=x.shape[0]) # doctest: +SKIP
+ >>> Particles(np.zeros(3), np.ones(3)).__cuda_args__()
+ (array([0., 0., 0.]), array([1., 1., 1.]), 3)
+ >>> kernel(dt, Particles(x, v)) # doctest: +SKIP
"""
def __init__(self, *values: Any) -> None:
self._cuda_args = tuple(values)
def __cuda_args__(self) -> tuple[Any, ...]:
- """The CUDA kernel arguments this object stands for."""
+ """Return the kernel arguments this object stands for."""
return self._cuda_args
class CudaParameter(NamedTuple):
- """One parameter of a CUDA kernel signature (or one field of a struct).
+ """One parameter of a CUDA kernel signature, or one field of a struct.
+
+ Returned by :func:`~cunumpy.cuda.parse_cuda_signature`.
Attributes
----------
name : str
Parameter name.
ctype : str
- Normalized C type without qualifiers or ``*``, e.g. ``"double"`` or
+ C type without qualifiers or ``*``, e.g. ``"double"`` or
``"Array2D"``.
- dtype : numpy.dtype | None
- NumPy dtype of the value (or of the pointed-to elements, or of the
- elements of an array view; the structured dtype for a struct);
- ``None`` for ``void*``.
+ dtype : numpy.dtype or None
+ Dtype of the value, of the pointed-to or viewed elements, or the
+ structured dtype of a struct; None for ``void*``.
pointer : bool
Whether the parameter is a pointer (a device array).
- struct : CudaStruct | None
+ struct : CudaStruct or None
The struct type, for a struct passed by value.
- view_ndim : int | None
- The number of dimensions, for an array view (``Array1D`` to
- ``Array16D`` or ``CArray1D`` to ``CArray16D``, see
- :func:`cuda_include_dir`) passed by value.
+ view_ndim : int or None
+ Number of dimensions of an array view (``Array1D`` to
+ ``Array16D``, ``CArray1D`` to ``CArray16D``).
contiguous : bool
Whether the array view is C-contiguous (``CArray2D``): packed
- without strides, and only C-contiguous arrays are accepted.
+ without strides, takes C-contiguous arrays only.
+
+ Examples
+ --------
+ >>> source = 'extern "C" __global__ void f(const double* x, int n) {}'
+ >>> x, n = xp.cuda.parse_cuda_signature(source, "f")
+ >>> x.ctype, x.dtype, x.pointer
+ ('double', dtype('float64'), True)
"""
name: str
@@ -278,16 +268,13 @@ class CudaParameter(NamedTuple):
def _normalize_view(match: re.Match) -> str:
- """``Array2D< const double >`` -> ``Array2D``."""
+ """Normalize a view type: ``Array2D< const double >`` -> ``Array2D``."""
words = [w for w in match.group(3).split() if w not in _QUALIFIERS]
return f"{match.group(1)}Array{match.group(2)}D<{' '.join(words)}>"
def _view_dtype(ndim: int, contiguous: bool = False) -> np.dtype:
- """The structured dtype with the C layout of ``ArrayD``.
-
- ``CArrayD`` (`contiguous`) has no strides.
- """
+ """Return the structured dtype of ``ArrayD`` (no strides if `contiguous`)."""
fields = [("data", np.uint64), ("shape", np.int64, (ndim,))]
if not contiguous:
fields.append(("strides", np.int64, (ndim,)))
@@ -295,11 +282,31 @@ def _view_dtype(ndim: int, contiguous: bool = False) -> np.dtype:
def ctype_of(dtype: Any) -> str:
- """The C type of a NumPy dtype, e.g. ``ctype_of(np.float64) == "double"``.
+ """Return the C type of a NumPy dtype.
- Useful to generate CUDA source or template arguments for a given dtype.
- Complex dtypes map to ``complex``/``complex`` (include
+ Useful to generate CUDA source or template arguments for a dtype. Complex
+ dtypes map to ``complex``/``complex`` (include
```` in the source).
+
+ Parameters
+ ----------
+ dtype : dtype-like
+ A boolean, integer, floating-point or complex dtype.
+
+ Returns
+ -------
+ str
+ The C type, e.g. ``"double"`` or ``"long long"``.
+
+ Raises
+ ------
+ ValueError
+ If the dtype has no C type.
+
+ Examples
+ --------
+ >>> xp.cuda.ctype_of(np.float64), xp.cuda.ctype_of(np.int64)
+ ('double', 'long long')
"""
try:
return _CTYPE_OF[np.dtype(dtype)]
@@ -322,7 +329,7 @@ def _strip_comments(source: str) -> str:
def _includes(source: str) -> list[tuple[str, bool]]:
- """``(name, quoted)`` of every ``#include`` in `source`, in order."""
+ """Return ``(name, quoted)`` of every ``#include`` in `source`, in order."""
return [
(quoted or angle, bool(quoted))
for quoted, angle in _INCLUDE.findall(_strip_comments(source))
@@ -336,37 +343,38 @@ def resolve_includes(
base_dir: str | Path | None = None,
angle_dirs: Iterable[str | Path] = (),
) -> list[Path]:
- """The header files a CUDA source includes, recursively.
+ r"""Return the header files a CUDA source includes, recursively.
- Scans `source` (comments removed) for ``#include "name"`` and resolves each
- name like NVRTC does: relative to `base_dir` (the directory of the
- including file), then in `include_dirs`, then in `angle_dirs`, in order.
- Found headers are scanned in turn, relative to their own directory.
- Includes in angle brackets (``#include ``) are system headers and
- ignored, unless they are found in `angle_dirs`. Includes that cannot be
- found are ignored; NVRTC reports them when the kernel is compiled.
+ Each ``#include "name"`` (comments ignored) is resolved like NVRTC does:
+ in `base_dir`, then `include_dirs`, then `angle_dirs`; found headers are
+ scanned in turn, relative to their own directory. ``#include `` is
+ ignored unless found in `angle_dirs`, and so are includes that cannot be
+ found (NVRTC reports them at compile time).
Parameters
----------
source : str
CUDA C source code.
- include_dirs : Iterable[str | Path]
- Directories searched for included files, in order (the ``-I`` options).
- base_dir : str | Path | None
- Directory of the file `source` was read from, searched first; None if
- the source is not from a file.
- angle_dirs : Iterable[str | Path]
- Directories whose headers are tracked also when included in angle
- brackets, searched last; :class:`CudaKernel` passes
- :func:`cuda_include_dir`, so that ``#include ``
- is tracked.
+ include_dirs : iterable of str or Path, optional
+ Directories searched in order (the ``-I`` options).
+ base_dir : str or Path, optional
+ Directory of the file `source` was read from, searched first.
+ angle_dirs : iterable of str or Path, optional
+ Directories searched last, also for ``#include ``;
+ :class:`~cunumpy.kernels.CudaKernel` passes
+ :func:`~cunumpy.cuda.cuda_include_dir`.
Returns
-------
- list[Path]
- The resolved header files, each once, in order of first inclusion
- (depth first). Empty if the source has no includes to track; the file
- system is not touched in that case.
+ list of Path
+ The headers, each once, in depth-first order of first inclusion.
+ Empty, without touching the file system, if there is nothing to track.
+
+ Examples
+ --------
+ >>> source = '#include \n#include '
+ >>> [p.name for p in xp.cuda.resolve_includes(source, angle_dirs=[xp.cuda.cuda_include_dir()])]
+ ['index.cuh']
"""
angle = tuple(Path(d) for d in angle_dirs)
dirs = tuple(Path(d) for d in include_dirs)
@@ -395,21 +403,29 @@ def visit(code: str, directory: Path | None) -> None:
def include_hash(paths: Iterable[str | Path]) -> str:
- """A short hex digest of the contents of `paths`, in order.
+ """Return a short hex digest of the contents of files, in order.
- Only the file contents count, not their locations: moving a header does not
- change the hash, editing it does. Used to make CuPy's kernel cache key
- depend on the included headers, see :meth:`CudaKernel.compile_options`.
+ Only the contents count: moving a header keeps the hash, editing it
+ changes it. Makes CuPy's kernel cache key depend on the included headers,
+ see :meth:`CudaKernel.compile_options
+ `.
Parameters
----------
- paths : Iterable[str | Path]
- Files to hash, e.g. from :func:`resolve_includes`.
+ paths : iterable of str or Path
+ Files to hash, e.g. from :func:`~cunumpy.cuda.resolve_includes`.
Returns
-------
str
The first 16 hex digits of the SHA-256 digest.
+
+ Examples
+ --------
+ >>> import os
+ >>> header = os.path.join(xp.cuda.cuda_include_dir(), "cunumpy", "index.cuh")
+ >>> len(xp.cuda.include_hash([header]))
+ 16
"""
digest = hashlib.sha256()
for path in paths:
@@ -423,10 +439,10 @@ def include_hash(paths: Iterable[str | Path]) -> str:
def cuda_kernel_names(source: str) -> list[str]:
- """The names of the ``__global__`` functions defined in `source`, in order.
+ """Return the names of the ``__global__`` functions in a CUDA source.
- Comments are ignored. Templates are included; a function declared more
- than once (e.g. a forward declaration) is listed once.
+ Comments are ignored, templates included; a function declared more than
+ once is listed once.
Parameters
----------
@@ -435,8 +451,17 @@ def cuda_kernel_names(source: str) -> list[str]:
Returns
-------
- list[str]
- The kernel names, in the order of their first appearance.
+ list of str
+ The kernel names, in order of first appearance.
+
+ Examples
+ --------
+ >>> source = '''
+ ... extern "C" __global__ void scale(double* x) {}
+ ... // __global__ void old(double* x) {}
+ ... extern "C" __global__ void shift(double* x) {}'''
+ >>> xp.cuda.cuda_kernel_names(source)
+ ['scale', 'shift']
"""
return list(dict.fromkeys(_GLOBAL_FUNCTION.findall(_strip_comments(source))))
@@ -445,18 +470,7 @@ def _compile_in_threads(
compilers: Mapping[Hashable, Callable[[], Any]],
jobs: int | None,
) -> list[Hashable]:
- """Run the `compilers` (name -> compile function), `jobs` at a time.
-
- With ``jobs=1`` they run one after the other in the calling thread; with
- ``jobs=None`` as many threads as CPUs are used. All compilers are run even
- if one fails; the first exception (in the order of `compilers`) is raised
- afterwards.
-
- Returns
- -------
- list
- The names whose compiler succeeded, in the order of `compilers`.
- """
+ """Run the `compilers` `jobs` at a time; raise the first error after all ran."""
if jobs is None:
jobs = os.cpu_count() or 1
if jobs < 1:
@@ -557,7 +571,7 @@ def _split_top_level(text: str) -> list[str]:
def _template_arg(value: Any) -> str:
- """A template argument as C++ source: a C type for dtypes, else a literal."""
+ """Return a template argument as C++ source: a C type for dtypes, else a literal."""
if isinstance(value, str):
return value
if isinstance(value, bool):
@@ -574,7 +588,10 @@ def parse_cuda_signature(
structs: Iterable[CudaStruct] = (),
template_args: Sequence[Any] | None = None,
) -> tuple[CudaParameter, ...]:
- """Parse the parameters of the ``__global__`` function `name` in `source`.
+ """Parse the parameters of a ``__global__`` function.
+
+ C types map to dtypes as on Linux (LP64): ``int`` is int32, ``long`` and
+ ``long long`` are int64.
Parameters
----------
@@ -582,26 +599,34 @@ def parse_cuda_signature(
CUDA C source code.
name : str
Name of the ``__global__`` function.
- structs : Iterable[CudaStruct]
- Struct types that may appear as parameters (passed by value). If the
- source defines a struct of the same name, its fields must match.
- template_args : Sequence | None
- Template arguments, if `name` is a function template: C types (or NumPy
- dtypes, see :func:`ctype_of`) for type parameters, integers or bools for
- non-type parameters. They are substituted into the parameter list.
+ structs : iterable of CudaStruct, optional
+ Struct types that may appear as parameters (by value). A definition
+ of the struct in `source` must match.
+ template_args : sequence, optional
+ Template arguments if `name` is a function template: C types or
+ dtypes (see :func:`~cunumpy.cuda.ctype_of`) for type parameters,
+ ints or bools for non-type parameters.
Returns
-------
- tuple[CudaParameter, ...]
+ tuple of CudaParameter
The parameters, in order.
Raises
------
ValueError
- If there is no such function, a template is used without (the right
- number of) `template_args`, a struct definition in the source does not
- match its :class:`CudaStruct`, or a parameter has a type that cannot be
- checked (e.g. a macro or a pointer to pointer).
+ If there is no such function, `template_args` do not fit, a struct
+ definition does not match its :class:`~cunumpy.arguments.CudaStruct`,
+ or a parameter type cannot be checked (e.g. a macro or ``double**``).
+
+ Examples
+ --------
+ >>> source = '''
+ ... template
+ ... __global__ void scale(T* x, T factor, int n) {}'''
+ >>> [(p.name, p.ctype, p.pointer) for p in xp.cuda.parse_cuda_signature(
+ ... source, "scale", template_args=[np.float32])]
+ [('x', 'float', True), ('factor', 'float', False), ('n', 'int', False)]
"""
code = _strip_comments(source)
structs = {s.name: s for s in structs}
@@ -654,11 +679,7 @@ def _describe(param: CudaParameter, index: int) -> str:
def _current_device_id() -> int | None:
- """The id of the current CUDA device, or None if CuPy has not been imported.
-
- CuPy is looked up in ``sys.modules`` and never imported here: without it
- there are no device arrays whose device could be checked.
- """
+ """Return the current CUDA device id, or None if CuPy was never imported."""
cp = sys.modules.get("cupy")
if cp is None:
return None
@@ -717,13 +738,7 @@ def check(value: Any) -> Any:
def _view_checker(param: CudaParameter, index: int) -> Callable[[Any], Any]:
- """Checker for an array view: packs (pointer, shape, strides) of a device array.
-
- Strides are converted from bytes to elements; the array need not be
- contiguous. A contiguous view (``CArray2D``) packs (pointer, shape) and
- takes C-contiguous arrays only: it is never copied, since what the kernel
- writes into a copy would be lost.
- """
+ """Checker for an array view: packs pointer, shape and strides (in elements)."""
ndim = param.view_ndim
contiguous = param.contiguous
dtype = _view_dtype(ndim, contiguous)
@@ -853,7 +868,7 @@ def _checker(param: CudaParameter, index: int) -> Callable[[Any], Any]:
def _field_dtype(field: CudaParameter) -> np.dtype:
- """The dtype of a struct field: pointers are stored as device addresses."""
+ """Return the dtype of a struct field (pointers are stored as addresses)."""
if field.pointer:
return np.dtype(np.uint64)
if field.view_ndim is not None:
@@ -882,7 +897,7 @@ def _field_dtype(field: CudaParameter) -> np.dtype:
def _pyccel_ctype(annotation: Any, scalars: Mapping[str, str], what: str) -> str:
- """The C type for a pyccel-style annotation, e.g. ``"float[:, :]"``."""
+ """Return the C type of a pyccel-style annotation, e.g. ``"float[:, :]"``."""
if annotation is inspect.Parameter.empty:
raise ValueError(f"{what} has no type annotation")
if typing.get_origin(annotation) is typing.Final:
@@ -890,7 +905,10 @@ def _pyccel_ctype(annotation: Any, scalars: Mapping[str, str], what: str) -> str
if isinstance(annotation, typing.ForwardRef):
annotation = annotation.__forward_arg__
if isinstance(annotation, str):
- match = _PYCCEL_ANNOTATION.match(annotation.strip())
+ # with `from __future__ import annotations`, 'float[:]' arrives quoted:
+ # "'float[:]'" (and Final['float[:]'] as one string), as in _annotation_text
+ text = annotation.replace("'", "").replace('"', "")
+ match = _PYCCEL_ANNOTATION.match(text.strip())
if match is None:
raise ValueError(f"{what}: cannot parse the annotation {annotation!r}")
scalar, dims = match.group("scalar"), match.group("dims")
@@ -918,7 +936,7 @@ def _apply_contiguous(
contiguous: bool | Iterable[str],
what: str,
) -> list[tuple[str, str]]:
- """Make array fields ``CArrayD``: all of them, or those named."""
+ """Make all array fields, or the named ones, ``CArrayD``."""
if contiguous is True:
names = {f for f, ctype in fields if ctype.startswith("Array")}
elif contiguous is False:
@@ -971,30 +989,46 @@ def write_cuda_header(
*,
includes: Iterable[str] = (),
) -> str:
- """Write the declarations of several structs to one header file.
+ """Write the definitions of several structs to one header file.
The header has an include guard, ``#include "cunumpy/array_view.cuh"`` if
- a struct has array view fields, then the struct definitions in order.
- Generate the header at build or test time from the Python definitions,
- and commit it next to the kernels; a test can regenerate it and compare
- (see :meth:`CudaStruct.to_header`).
+ a struct has array view fields, any `includes`, then the structs in
+ order. Commit it next to the kernels and let a test regenerate and
+ compare it (see :meth:`CudaStruct.to_header`).
Parameters
----------
- path : str | Path
+ path : str or Path
File to write.
- structs : Iterable[CudaStruct]
- The structs, in the order they are declared.
- guard : str | None
+ structs : iterable of CudaStruct
+ The structs, in declaration order.
+ guard : str, optional
Include guard macro; by default from the file name
(``pusher_args.cuh`` -> ``PUSHER_ARGS_CUH``).
- includes : Iterable[str]
- Additional headers to include, as file names or ``#include`` lines.
+ includes : iterable of str, optional
+ Additional headers, as file names or ``#include`` lines.
Returns
-------
str
The header source that was written.
+
+ Examples
+ --------
+ >>> import os, tempfile
+ >>> Vec = xp.arguments.CudaStruct("Vec", [("data", "double*"), ("n", "int")])
+ >>> path = os.path.join(tempfile.mkdtemp(), "vec.cuh")
+ >>> print(xp.arguments.write_cuda_header(path, [Vec]), end="")
+ // Generated by cunumpy.arguments.CudaStruct from the Python definition; do not edit.
+ #ifndef VEC_CUH
+ #define VEC_CUH
+
+ struct Vec {
+ double* data;
+ int n;
+ };
+
+ #endif // VEC_CUH
"""
path = Path(path)
source = _header_source(tuple(structs), guard or _header_guard(path.name), includes)
@@ -1003,7 +1037,7 @@ def write_cuda_header(
def _read_python_source(source: str | Path) -> tuple[str, str]:
- """`(text, description)` of a Python source given as a path or as code."""
+ """Return ``(text, description)`` of a Python source given as a path or as code."""
if isinstance(source, Path) or (
"\n" not in source and source.strip().endswith(".py")
):
@@ -1013,7 +1047,7 @@ def _read_python_source(source: str | Path) -> tuple[str, str]:
def _find_init(tree: ast.Module, class_name: str, where: str) -> ast.FunctionDef:
- """The ``__init__`` function of the class `class_name` in a parsed module."""
+ """Return the ``__init__`` of the class `class_name` in a parsed module."""
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == class_name:
for item in node.body:
@@ -1024,7 +1058,7 @@ def _find_init(tree: ast.Module, class_name: str, where: str) -> ast.FunctionDef
def _stored_parameters(init: ast.FunctionDef) -> dict[str, str]:
- """``{parameter: attribute}`` for the ``self. = `` statements."""
+ """Return ``{parameter: attribute}`` for ``self. = ``."""
stored: dict[str, str] = {}
for node in ast.walk(init):
targets: list[ast.expr] = []
@@ -1047,7 +1081,7 @@ def _stored_parameters(init: ast.FunctionDef) -> dict[str, str]:
def _annotation_text(annotation: ast.expr) -> str:
- """The pyccel-style annotation string of a parsed annotation."""
+ """Return the pyccel-style annotation string of a parsed annotation."""
if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str):
return annotation.value
text = ast.unparse(annotation)
@@ -1058,50 +1092,44 @@ def _annotation_text(annotation: ast.expr) -> str:
class CudaStruct:
"""A C struct type passed to CUDA kernels by value.
- Defines the struct once: :attr:`declaration` is its C definition, to put in
- the CUDA source (or a header), and calling the struct packs values into it.
- The packed value has the memory layout of the C struct (the NumPy structured
- dtype with C alignment), pointers are stored as device addresses. Kernels
- that are told about the struct (``CudaKernel(..., structs=[...])``) check
- that they get a value of the right struct, and that a definition of the
- struct in their source matches.
-
- Instead of a long flat parameter list, a group of arguments (e.g. all arrays
- describing particles) becomes one struct parameter; adding a field then
- changes one definition instead of every kernel signature.
+ Defines the struct once: :attr:`declaration` is its C definition, and
+ calling the struct packs values into a :class:`CudaStructValue` with the C
+ memory layout (pointers stored as device addresses). Kernels created with
+ ``structs=[...]`` check that they get a value of this struct and that a
+ definition in their source matches. See :doc:`/kernels/arguments`.
Parameters
----------
name : str
Name of the struct type in C.
- fields : Sequence[tuple[str, str]]
- ``(field name, C type)`` pairs, in order, e.g. ``("x", "double*")`` or
- ``("n", "int")``. Scalar fields, pointers to the scalar types of
- :func:`ctype_of` (or ``void*``), and array views ``Array1D`` to
- ``Array16D`` of those scalar types (from ``cunumpy/array_view.cuh``,
- packed as pointer, shape and strides in elements) are supported.
+ fields : sequence of (str, str)
+ ``(field name, C type)`` pairs in order, e.g. ``("x", "double*")``.
+ Scalars, pointers to scalars (or ``void*``), and array views
+ ``Array1D`` to ``Array16D`` (packed as pointer, shape and
+ strides in elements) or C-contiguous ``CArray1D`` to
+ ``CArray16D`` are supported.
+
+ Raises
+ ------
+ ValueError
+ If the name is not a C identifier, a field type is not supported or
+ field names repeat.
Examples
--------
- >>> Vec = CudaStruct("Vec", [("data", "double*"), ("n", "int")])
+ >>> Vec = xp.arguments.CudaStruct("Vec", [("data", "double*"), ("n", "int")])
+ >>> print(Vec.declaration, end="")
+ struct Vec {
+ double* data;
+ int n;
+ };
>>> source = Vec.declaration + r'''
... extern "C" __global__ void scale(Vec v, double a) {
... int i = blockDim.x * blockIdx.x + threadIdx.x;
... if (i < v.n) v.data[i] *= a;
... }'''
- >>> scale = CudaKernel(source, "scale", structs=[Vec])
- >>> scale(Vec(data=x, n=x.size), 2.0, n_threads=x.size) # doctest: +SKIP
-
- With array views, the kernel indexes like the pyccel kernel it mirrors:
-
- >>> Markers = CudaStruct("Markers", [("markers", "Array2D"), ("n", "int")])
- >>> source = '#include "cunumpy/array_view.cuh"\n' + Markers.declaration + r'''
- ... extern "C" __global__ void push(Markers m, double dt) {
- ... int ip = blockDim.x * blockIdx.x + threadIdx.x;
- ... if (ip < m.n) m.markers(ip, 0) += dt * m.markers(ip, 3);
- ... }'''
- >>> push = CudaKernel(source, "push", structs=[Markers])
- >>> push(Markers(markers=markers, n=markers.shape[0]), 0.1, n_threads=markers.shape[0]) # doctest: +SKIP
+ >>> scale = xp.kernels.CudaKernel(source, "scale", structs=[Vec])
+ >>> scale(Vec(data=x, n=x.size), 2.0) # doctest: +SKIP
"""
def __init__(self, name: str, fields: Sequence[tuple[str, str]]) -> None:
@@ -1131,50 +1159,53 @@ def from_signature(
) -> CudaStruct:
"""Build the struct from the annotated parameters of a Python function.
- One field per parameter of `func` (``self`` is skipped), in order,
- with the C type given by the parameter's annotation, written in the
- pyccel style: ``"float"`` -> ``double``, ``"int"`` -> `int_type`,
- ``"bool"`` -> ``bool``, and an array ``"float[:, :]"`` ->
- ``Array2D`` (``Final[...]`` and ``const`` are ignored). The
- annotations may also be the real types ``int``, ``float``, ``bool``
- (or NumPy scalar types such as ``np.float32``). Typically `func` is the
- ``__init__`` of an argument class, so that the class is the one
- definition of the arguments on the host and on the device.
+ One field per parameter of `func` (``self`` skipped), in order, typed
+ by its pyccel-style annotation: ``float`` -> ``double``, ``int`` ->
+ `int_type`, ``bool`` -> ``bool``, ``float[:, :]`` ->
+ ``Array2D`` (``Final[...]`` and ``const`` ignored). Real
+ types (``int``, ``np.float32``, ...) work too. Typically `func` is the
+ ``__init__`` of the host argument class, so that one class defines
+ the arguments on host and device.
Parameters
----------
- func : Callable
+ func : callable
The function whose parameters define the fields.
name : str
Name of the struct type in C.
- int_type : str
- C type for ``int`` annotations: pyccel integers are 64 bit, so
- ``"long long"`` by default; ``"int"`` for 32 bit.
- scalar_names : Mapping[str, str] | None
- Additional (or changed) mappings from annotation scalar names to
- C types, e.g. ``{"float": "float"}`` for single precision.
- contiguous : bool | Iterable[str]
- Array fields that become C-contiguous views (``CArray2D``
- instead of ``Array2D``): ``True`` for all of them, or their
- names. Packing such a field raises for a non-contiguous array.
+ int_type : str, optional
+ C type for ``int``: ``"long long"`` (default, pyccel integers are
+ 64 bit) or ``"int"``.
+ scalar_names : mapping of str to str, optional
+ Additional or changed scalar mappings, e.g. ``{"float": "float"}``
+ for single precision.
+ contiguous : bool or iterable of str, optional
+ Array fields that become C-contiguous views (``CArray2D``):
+ True for all, or their names. Packing them raises for a
+ non-contiguous array.
+
+ Returns
+ -------
+ CudaStruct
+ The struct.
Raises
------
ValueError
- A parameter without annotation, or with an annotation that cannot
- be mapped (an unknown scalar, more than 3 dimensions).
+ If a parameter has no annotation, or one that cannot be mapped
+ (an unknown scalar, more than 16 dimensions).
Examples
--------
>>> class MarkerArguments:
- ... def __init__(self, markers: "float[:, :]", n_markers: "int", valid: "bool[:]"):
- ... ...
- >>> MarkerArgs = CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs")
- >>> print(MarkerArgs.declaration)
+ ... def __init__(self, markers: "float[:, :]", n_markers: int): ...
+ >>> MarkerArgs = xp.arguments.CudaStruct.from_signature(
+ ... MarkerArguments.__init__, "MarkerArgs"
+ ... )
+ >>> print(MarkerArgs.declaration, end="")
struct MarkerArgs {
Array2D markers;
long long n_markers;
- Array1D valid;
};
"""
scalars = dict(_PYCCEL_SCALARS, int=int_type)
@@ -1205,49 +1236,58 @@ def from_pyccel_class(
attribute_names: bool = True,
contiguous: bool | Iterable[str] = False,
) -> CudaStruct:
- """Build the struct from the ``__init__`` of a class in a Python source file.
-
- Like :meth:`from_signature`, but for a class whose module is compiled
- by pyccel: importing such a module gives the compiled class, whose
- ``__init__`` has no Python signature. The ``.py`` source is parsed
- with :mod:`ast` instead and is never imported or executed.
+ """Build the struct from the ``__init__`` of a class in a Python source.
- One field per parameter of ``__init__`` (``self`` skipped), in order.
- A parameter that ``__init__`` stores as ``self. = ``
- gives a field named after the attribute (`attribute_names`), so that
- the struct members are the attribute names the host kernels use, also
- when the constructor parameter is called differently.
+ Like :meth:`from_signature`, for a class whose module pyccel compiles
+ (the compiled ``__init__`` has no Python signature): the source is
+ parsed with :mod:`ast`, never imported or executed. A parameter stored
+ as ``self. = `` gives a field named after the
+ attribute, as the host kernels use it.
Parameters
----------
- source : str | Path
- Path of the ``.py`` file, or the source code itself (a string that
- contains a newline or does not end with ``.py``).
+ source : str or Path
+ Path of the ``.py`` file, or the source code (a string with a
+ newline or not ending in ``.py``).
class_name : str
Name of the class in the source.
- name : str | None
+ name : str, optional
Name of the struct type in C; `class_name` by default.
- int_type, scalar_names, contiguous
- As for :meth:`from_signature` (`contiguous` takes field names, i.e.
- attribute names when `attribute_names` is set).
- exclude : Sequence[str]
+ int_type : str, optional
+ As for :meth:`from_signature`.
+ scalar_names : mapping of str to str, optional
+ As for :meth:`from_signature`.
+ exclude : sequence of str, optional
Parameter or attribute names that do not become fields, e.g.
- scratch arrays the host class allocates for itself.
- attribute_names : bool
- Name the fields after the attributes the parameters are stored in
- (default); False keeps the parameter names.
+ scratch arrays of the host class.
+ attribute_names : bool, optional
+ Name the fields after the attributes (default), else after the
+ parameters.
+ contiguous : bool or iterable of str, optional
+ As for :meth:`from_signature`, with field names.
+
+ Returns
+ -------
+ CudaStruct
+ The struct.
Raises
------
ValueError
- The class or its ``__init__`` is not found, or an annotation
- cannot be mapped (see :meth:`from_signature`).
+ If the class or its ``__init__`` is not found, or an annotation
+ cannot be mapped.
Examples
--------
- >>> MarkerArgs = CudaStruct.from_pyccel_class(
- ... "kernel_arguments/pusher_args_kernels.py", "MarkerArguments", "MarkerArgs"
- ... ) # doctest: +SKIP
+ >>> source = '''
+ ... class MarkerArguments:
+ ... def __init__(self, markers: 'float[:, :]', n: 'int'):
+ ... self.markers = markers
+ ... self.n_markers = n
+ ... '''
+ >>> MarkerArgs = xp.arguments.CudaStruct.from_pyccel_class(source, "MarkerArguments")
+ >>> [f.name for f in MarkerArgs.fields]
+ ['markers', 'n_markers']
"""
text, where = _read_python_source(source)
init = _find_init(ast.parse(text, filename=where), class_name, where)
@@ -1281,7 +1321,7 @@ def name(self) -> str:
@property
def fields(self) -> tuple[CudaParameter, ...]:
- """The fields, in order."""
+ """The fields, as :class:`~cunumpy.cuda.CudaParameter` tuples in order."""
return self._fields
@property
@@ -1296,12 +1336,7 @@ def has_views(self) -> bool:
@property
def declaration(self) -> str:
- """The C definition of the struct, to include in the CUDA source.
-
- A struct with array view fields needs
- ``#include "cunumpy/array_view.cuh"`` before the definition (see
- :attr:`has_views`); :meth:`to_header` adds it.
- """
+ """The C definition, to put after ``cunumpy/array_view.cuh`` if :attr:`has_views`."""
lines = [
f" {f.ctype}{'*' if f.pointer else ''} {f.name};" for f in self._fields
]
@@ -1314,30 +1349,32 @@ def to_header(
guard: str | None = None,
includes: Iterable[str] = (),
) -> str:
- """The struct definition as a header file, with include guard.
+ """Return the struct definition as a header, with include guard.
The header includes ``cunumpy/array_view.cuh`` if the struct has array
- view fields, then any `includes`, then the definition. A generated
- header committed next to the kernels stays in sync with the Python
- definition through a test::
+ view fields, then any `includes`, then the definition. A test keeps a
+ committed header in sync with the Python definition::
def test_header_is_up_to_date():
assert Path("pusher_args.cuh").read_text() == MarkerArgs.to_header()
Parameters
----------
- path : str | Path | None
+ path : str or Path, optional
If given, the header is also written to this file.
- guard : str | None
- Include guard macro; by default ``_CUH`` from the struct
- name.
- includes : Iterable[str]
- Additional headers to include, as file names or ``#include`` lines.
+ guard : str, optional
+ Include guard macro; ``_CUH`` by default.
+ includes : iterable of str, optional
+ Additional headers, as file names or ``#include`` lines.
Returns
-------
str
The header source.
+
+ See Also
+ --------
+ cunumpy.arguments.write_cuda_header : Several structs in one header.
"""
source = _header_source(
(self,),
@@ -1349,10 +1386,15 @@ def test_header_is_up_to_date():
return source
def check_source(self, source: str) -> None:
- """Check a definition of this struct in `source` against its fields.
+ """Check a definition of this struct in a CUDA source against its fields.
- Does nothing if `source` does not define the struct (e.g. because it is
- in a header).
+ Does nothing if `source` does not define the struct (e.g. it is in a
+ header).
+
+ Parameters
+ ----------
+ source : str
+ CUDA C source code.
Raises
------
@@ -1382,17 +1424,22 @@ def check_source(self, source: str) -> None:
)
def layout_source(self, include: str | None = None) -> str:
- """The CUDA source of the kernel used by :meth:`verify_layout`.
+ """Return the CUDA source of the kernel used by :meth:`verify_layout`.
The kernel ``cunumpy_layout_(unsigned long long* out)`` writes
- ``sizeof``, ``alignof`` and the offset of every field (in field order)
- of the struct, as the compiler lays it out.
+ ``sizeof``, ``alignof`` and the offset of every field, as the compiler
+ lays out the struct.
Parameters
----------
- include : str | None
- Header that defines the struct, as a file name or an ``#include``
- line; by default the struct is defined by :attr:`declaration`.
+ include : str, optional
+ Header that defines the struct, as a file name or ``#include``
+ line; by default :attr:`declaration` is used.
+
+ Returns
+ -------
+ str
+ The CUDA source.
"""
if include is None:
lines = ['#include "cunumpy/array_view.cuh"'] if self.has_views else []
@@ -1423,38 +1470,41 @@ def verify_layout(
include_dirs: Sequence[str | Path] = (),
options: Sequence[str] = (),
) -> dict[str, int]:
- """Check the struct layout of the CUDA compiler against :attr:`dtype`.
+ """Check the CUDA compiler's struct layout against :attr:`dtype`.
- Compiles and runs a one-thread kernel (:meth:`layout_source`) that
- reports the size, alignment and field offsets of the struct, and
- compares them with the NumPy dtype that values are packed into. A
- difference means every kernel taking the struct reads some fields at
- the wrong place, without any error. Run it once per struct in a GPU
- test, especially with a hand-written or generated header (`include`)
- and on a new compiler or platform (e.g. ROCm).
+ Compiles and runs a one-thread kernel (:meth:`layout_source`) and
+ compares size, alignment and field offsets with the dtype values are
+ packed into; a difference would make kernels silently read fields at
+ the wrong place. Run it once per struct in a GPU test, especially with
+ a hand-written header or a new compiler or platform (e.g. ROCm).
Parameters
----------
- include : str | None
+ include : str, optional
Header that defines the struct (file name or ``#include`` line),
- found in `include_dirs`; by default :attr:`declaration` is compiled.
- include_dirs : Sequence[str | Path]
+ found in `include_dirs`; by default :attr:`declaration`.
+ include_dirs : sequence of str or Path, optional
Directories searched for `include`.
- options : Sequence[str]
+ options : sequence of str, optional
Additional compiler options.
Returns
-------
- dict[str, int]
+ dict of str to int
``"sizeof"``, ``"alignof"`` and the offset of each field, by name.
Raises
------
ValueError
- If the compiled layout differs from :attr:`dtype` (the message lists
- each difference).
+ If the compiled layout differs from :attr:`dtype` (the message
+ lists each difference).
RuntimeError
If CuPy is not available.
+
+ Examples
+ --------
+ >>> Vec.verify_layout() # doctest: +SKIP
+ {'sizeof': 16, 'alignof': 8, 'data': 0, 'n': 8}
"""
from cunumpy.xp import cupy_available, to_numpy
@@ -1491,18 +1541,29 @@ def verify_layout(
return layout
def __call__(self, **values: Any) -> CudaStructValue:
- """Pack values into the struct.
+ """Pack field values, given by name, into a struct value.
Pointer fields take C-contiguous CuPy arrays of the declared dtype
- (never copied), array view fields take CuPy arrays of the declared
- dtype and number of dimensions (contiguous or not; their pointer, shape
- and strides are packed), scalar fields are checked and cast like scalar kernel
- arguments.
+ (never copied), array view fields CuPy arrays of the declared dtype
+ and ndim (contiguous or not), and scalars are checked and cast like
+ scalar kernel arguments.
+
+ Parameters
+ ----------
+ **values
+ One value per field.
+
+ Returns
+ -------
+ CudaStructValue
+ The packed value.
Raises
------
TypeError
- A missing or unknown field, or a value of the wrong type.
+ If a field is missing or unknown, or a value has the wrong type.
+ OverflowError
+ If an integer is out of range of its field type.
"""
missing = [f.name for f in self._fields if f.name not in values]
unknown = [k for k in values if k not in self._checkers]
@@ -1521,10 +1582,25 @@ def __call__(self, **values: Any) -> CudaStructValue:
class CudaStructValue:
"""A value of a :class:`CudaStruct`, passed to kernels as one argument.
- Holds the packed struct and references to the arrays it points to (the
- struct itself only stores device addresses), and flattens into the packed
- struct through ``__cuda_args__()``. Field values are available as
- ``value["name"]``.
+ Created by calling the struct. Keeps the arrays it points to alive (the
+ packed struct only stores device addresses) and flattens into the packed
+ struct through ``__cuda_args__()``. Field values are ``value["name"]``.
+
+ Parameters
+ ----------
+ struct : CudaStruct
+ The struct type.
+ packed : numpy.void
+ The packed struct.
+ values : dict
+ The field values, by name.
+
+ Examples
+ --------
+ >>> Params = xp.arguments.CudaStruct("Params", [("dt", "double"), ("n", "int")])
+ >>> value = Params(dt=0.1, n=10)
+ >>> value["n"], value.packed["n"]
+ (10, np.int32(10))
"""
def __init__(self, struct: CudaStruct, packed: np.void, values: dict) -> None:
@@ -1549,77 +1625,57 @@ def __getitem__(self, field: str) -> Any:
return self._values[field]
def __cuda_args__(self) -> tuple[np.void]:
+ """Return the packed struct, the one kernel argument this value stands for."""
return (self._packed,)
class CudaStructArguments(CudaArguments):
- """Base class for argument objects that are passed to CUDA kernels as one C struct.
-
- The class form of :class:`CudaStruct`: a subclass names the struct
- (:attr:`struct_name`) and lists its fields (:attr:`fields`), stores every
- field as an attribute of the same name, and calls :meth:`pack` at the end
- of its constructor. The object is then passed to a :class:`CudaKernel` as
- it is and arrives as the packed struct.
-
- The :class:`CudaStruct` is built once per subclass, when the class is
- defined, and is available as :attr:`struct` (for ``CudaKernel(...,
- structs=[...])``, :attr:`CudaStruct.declaration` and
- :meth:`CudaStruct.to_header`). Packing checks every field like a kernel
- argument: pointers need C-contiguous CuPy arrays of the declared dtype
- (never copied), scalars are range-checked and cast.
-
- The packed struct always reflects the current field attributes: at every
- use (:attr:`packed`, :meth:`__cuda_args__`, so at every kernel launch) the
- device address, shape and strides of every array field and the value of
- every scalar field are compared with those that were packed, and the
- struct is packed again if any changed. A field may therefore be a
- property that reads the owner's current array, so that resizing the
- owner's arrays never leaves the struct pointing at freed device memory
- (see the second example). Copies (``copy.copy``, ``copy.deepcopy``) and
- unpickled objects are packed again from their own arrays.
-
- A subclass that sets neither :attr:`struct_name` nor :attr:`fields` is an
- intermediate base class; a subclass that sets only one of them raises
- ``TypeError``.
+ """Base class for argument objects passed to CUDA kernels as one C struct.
+
+ The class form of :class:`CudaStruct`: a subclass sets ``struct_name``
+ and ``fields``, stores every field as an attribute of the same name
+ and calls :meth:`pack` at the end of its constructor. The object is
+ passed to a :class:`~cunumpy.kernels.CudaKernel` as is and arrives as the
+ packed struct. See :doc:`/kernels/arguments`.
+
+ The struct is repacked whenever a field changed (another array, shape,
+ strides or scalar value), checked at every use, so a field may be a
+ property reading an owner's current array (second example). Copies and
+ unpickled objects are packed from their own arrays. A subclass setting
+ neither ``struct_name`` nor ``fields`` is an intermediate base
+ class; setting only one raises ``TypeError``.
Attributes
----------
struct_name : str
Name of the C struct type.
- fields : Sequence[tuple[str, str]]
- ``(field name, C type)`` pairs, in declaration order, as for
- :class:`CudaStruct`.
+ fields : sequence of (str, str)
+ ``(field name, C type)`` pairs, as for :class:`CudaStruct`.
struct : CudaStruct
- The struct type, built from :attr:`struct_name` and :attr:`fields`.
+ The struct type, built once when the subclass is defined.
Examples
--------
- >>> class MarkerArguments(CudaStructArguments):
+ >>> class MarkerArguments(xp.arguments.CudaStructArguments):
... struct_name = "MarkerArgs"
- ... fields = (
- ... ("markers", "Array2D"),
- ... ("valid", "bool*"),
- ... ("n_markers", "int"),
- ... )
+ ... fields = (("markers", "Array2D"), ("n_markers", "int"))
...
- ... def __init__(self, markers, valid):
+ ... def __init__(self, markers):
... self.markers = markers
- ... self.valid = valid
... self.n_markers = markers.shape[0]
... self.pack()
- >>> print(MarkerArguments.struct.declaration)
+ >>> print(MarkerArguments.struct.declaration, end="")
struct MarkerArgs {
Array2D markers;
- bool* valid;
int n_markers;
};
- >>> push = CudaKernel(source, "push", structs=[MarkerArguments.struct]) # doctest: +SKIP
- >>> push(MarkerArguments(markers, valid), 0.1, n_threads=markers.shape[0]) # doctest: +SKIP
+ >>> push = xp.kernels.CudaKernel(source, "push", structs=[MarkerArguments.struct]) # doctest: +SKIP
+ >>> push(MarkerArguments(markers), 0.1) # doctest: +SKIP
Fields as properties follow the arrays of an owner object, also after the
- owner replaced them (e.g. when it resized its marker array):
+ owner replaced them (e.g. resized its marker array):
- >>> class ParticleArguments(CudaStructArguments):
+ >>> class ParticleArguments(xp.arguments.CudaStructArguments):
... struct_name = "ParticleArgs"
... fields = (("markers", "Array2D"), ("n_markers", "int"))
...
@@ -1655,16 +1711,14 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
def pack(self) -> None:
"""Pack the field attributes into the struct.
- Called at the end of the constructor, so that invalid field values
- raise there. Afterwards the struct is packed again automatically when
- a field changes (see the class documentation); calling this method
- again is never needed, but harmless.
+ Call it at the end of the constructor, so that invalid values raise
+ there; later changes are repacked automatically.
Raises
------
TypeError
- If the class does not define a struct, or a field value has the
- wrong type (see :meth:`CudaStruct.__call__`).
+ If the class defines no struct, or a field value has the wrong
+ type (see :meth:`CudaStruct.__call__ `).
AttributeError
If a field has no attribute of the same name.
"""
@@ -1678,7 +1732,7 @@ def pack(self) -> None:
self._packed_state = _field_state(struct, values)
def _field_values(self, struct: CudaStruct) -> dict[str, Any]:
- """The current value of every field attribute, by field name."""
+ """Return the current value of every field attribute, by field name."""
values = {}
for field in struct.fields:
try:
@@ -1692,11 +1746,7 @@ def _field_values(self, struct: CudaStruct) -> dict[str, Any]:
@property
def packed(self) -> np.void:
- """The packed struct, with the memory layout of the C struct.
-
- Packed on first use, and again whenever a field attribute changed
- since the last packing (a different array, or a different scalar).
- """
+ """The packed struct, repacked first if a field attribute changed."""
value = self.__dict__.get("_struct_value")
if value is None:
self.pack()
@@ -1709,7 +1759,7 @@ def packed(self) -> np.void:
return self._struct_value.packed
def __cuda_args__(self) -> tuple[np.void]:
- """The packed struct, as the one kernel argument this object stands for."""
+ """Return the packed struct, the one kernel argument this object stands for."""
return (self.packed,)
def __getstate__(self) -> dict[str, Any]:
@@ -1726,13 +1776,7 @@ def __setstate__(self, state: dict[str, Any]) -> None:
def _field_state(struct: CudaStruct, values: Mapping[str, Any]) -> tuple[Any, ...]:
- """What the packed struct depends on, to detect changed field attributes.
-
- For an array field the device address, shape and strides (the address
- alone for a pointer field); for a scalar field its value. A value that is
- neither (e.g. a host array in a pointer field) is identified by its id,
- so replacing it triggers a repack, which then raises the type error.
- """
+ """Return what the packed struct depends on, to detect changed fields."""
state = []
for field in struct.fields:
value = values[field.name]
@@ -1759,7 +1803,7 @@ def _is_device_array(value: Any) -> bool:
def _array_shapes_in(args: Sequence[Any]) -> Iterator[tuple[int, ...]]:
- """Array shapes in argument order, including supported argument objects."""
+ """Yield array shapes in argument order, including those in argument objects."""
for arg in args:
shape = getattr(arg, "shape", None)
if shape is not None and hasattr(arg, "dtype"):
@@ -1793,15 +1837,12 @@ def _first_array_shape(args: Sequence[Any], dimensions: int) -> int | tuple[int,
def _first_array_length(args: tuple[Any, ...]) -> int:
- """``n_threads_from="first_array"``: the first axis of the first array argument."""
+ """Return the first axis of the first array (``n_threads_from="first_array"``)."""
return typing.cast(int, _first_array_shape(args, 1))
def _last_axis_length(args: tuple[Any, ...]) -> int:
- """``n_threads_from="last_axis"``: the last axis of the first array argument.
-
- One thread per entry of a component-major array, ``(ncomp, N)`` or ``(N,)``.
- """
+ """Return the last axis of the first array (``n_threads_from="last_axis"``)."""
shape = next(_array_shapes_in(args), None)
if shape is None:
raise TypeError(
@@ -1819,12 +1860,7 @@ def _as_shape(value: int | Sequence[int], what: str) -> tuple[int, ...]:
def _device_arrays_in(args: Sequence[Any]) -> Iterator[tuple[str, Any]]:
- """``(label, array)`` for every device array among kernel arguments.
-
- Arrays are found directly, in the fields of :class:`CudaStructArguments`
- objects and struct values, and in the values of other ``__cuda_args__()``
- objects.
- """
+ """Yield ``(label, array)`` for every device array among the arguments."""
for index, arg in enumerate(args):
label = f"argument {index}"
if _is_device_array(arg):
@@ -1848,75 +1884,81 @@ def _device_arrays_in(args: Sequence[Any]) -> Iterator[tuple[str, Any]]:
class CudaKernel:
"""A CUDA C kernel, compiled with NVRTC through CuPy.
+ The ``__global__`` signature is parsed at construction (no GPU needed)
+ and every call is checked against it: arrays must be C-contiguous CuPy
+ arrays of the declared dtype (never copied), scalars are cast to the
+ declared C type or raise, argument objects with ``__cuda_args__()`` are
+ flattened. The kernel is compiled on the first call (or by
+ :meth:`compile`) and cached per device, also on disk by CuPy. See
+ :doc:`/kernels/cuda-kernel`.
+
Parameters
----------
source : str
- CUDA C source code containing the ``__global__`` function `name`
- (``extern "C"`` unless it is a template).
+ CUDA C source with the ``__global__`` function `name` (``extern "C"``
+ unless it is a template).
name : str
- Name of the kernel function in `source`.
- block_size : int | Sequence[int]
- Threads per block: an integer for 1D launches, or 1 to 3 integers
- (at most 1024 threads in total).
- options : Sequence[str]
- Additional NVRTC compiler options, e.g. ``("-std=c++17",)``.
- include_dirs : Sequence[str | Path]
- Directories searched for ``#include`` files (passed as ``-I``).
- The headers shipped with cunumpy (:func:`cuda_include_dir`) are always
- found at compile time (see :meth:`compile_options`):
- ``#include "cunumpy/array_view.cuh"`` gives the ``Array1D`` to
- ``Array16D`` views, ``#include "cunumpy/index.cuh"`` the thread-index
- macros, ``#include "cunumpy/atomic.cuh"`` atomic adds,
- ``#include "cunumpy/reduce.cuh"`` warp and block reductions.
- source_dir : str | Path | None
- Directory the source was read from (set by :meth:`from_file`), where
- ``#include "..."`` files are looked up first.
- structs : Iterable[CudaStruct]
- Struct types passed to the kernel by value.
- template_args : Sequence | None
+ Name of the kernel function.
+ block_size : int or sequence of int, optional
+ Threads per block, 1 to 3 dimensions, at most 1024 in total; 128 by
+ default.
+ options : sequence of str, optional
+ Additional NVRTC options, e.g. ``("-std=c++17",)``.
+ include_dirs : sequence of str or Path, optional
+ Directories for ``#include`` (``-I``). cunumpy's own headers
+ (:func:`~cunumpy.cuda.cuda_include_dir`) are always found.
+ source_dir : str or Path, optional
+ Directory the source was read from, searched first for
+ ``#include "..."`` (set by :meth:`from_file`).
+ structs : iterable of CudaStruct, optional
+ Struct types the kernel takes by value.
+ template_args : sequence, optional
Template arguments if `name` is a function template, e.g.
- ``("double", 3)`` or ``(np.float64, 3)``; the kernel is then the
- instantiation ``name``.
- check_signature : bool
- Parse the kernel signature and check (and cast) every call against it.
- Raises ``ValueError`` at construction if the signature cannot be parsed
- (e.g. macros in the parameter list); pass False to launch with the
- arguments as they are, like ``cupy.RawKernel``.
- debug : bool | None
- Debug mode: compile with :data:`DEBUG_OPTIONS` (``-lineinfo`` and
- ``-DCUNUMPY_BOUNDS_CHECK``) and synchronize after every launch, so
- that an asynchronous CUDA error is raised as a ``RuntimeError`` naming
- this kernel. None (the default) follows the global setting
- (:func:`cunumpy.cuda.set_cuda_debug`, ``CUNUMPY_CUDA_DEBUG``) at every
- launch; True or False fix it for this kernel. The compile options are
- fixed when the kernel is compiled.
- n_threads_from : {"auto", "first_array", "last_axis"} | callable | None
- Default "auto" infers thread counts from the first array's leading
- shape axes, matching the block dimensionality (1D: one thread per row).
- Arrays in supported argument objects are included. "first_array"
- always uses the first axis, "last_axis" the last one (one thread per
- entry of a component-major ``(ncomp, N)`` or ``(N,)`` array). A callable receives the positional argument
- tuple; None requires an explicit launch size. Explicit `n_threads` or
- `grid` overrides inference.
+ ``(np.float64, 3)`` for ``name``.
+ check_signature : bool, optional
+ Check and cast every call against the signature (default). False
+ launches with the arguments as they are, like ``cupy.RawKernel``.
+ debug : bool, optional
+ Debug mode: compile with :data:`~cunumpy.cuda.DEBUG_OPTIONS` and
+ synchronize after every launch, raising a ``RuntimeError`` naming
+ this kernel. None (default) follows
+ :func:`~cunumpy.cuda.set_cuda_debug` at every launch. See
+ :doc:`/kernels/debugging`.
+ n_threads_from : {"auto", "first_array", "last_axis"}, callable or None, optional
+ Launch size when a call gives neither `n_threads` nor `grid`. "auto"
+ (default): the leading axes of the first array argument (also inside
+ argument objects), one per block dimension; "first_array": its first
+ axis; "last_axis": its last axis (component-major ``(ncomp, N)``); a
+ callable of the argument tuple, e.g. ``lambda args: args[0].size``;
+ None: no inference. Settable later as :attr:`n_threads_from`.
+ check_finite : bool, optional
+ After every launch, synchronize and raise ``RuntimeError`` if a
+ floating-point or complex array argument (struct fields included)
+ holds NaN or inf. Costs a pass over the arrays: for debugging only.
+
+ Raises
+ ------
+ ValueError
+ If the signature cannot be parsed (e.g. macros in the parameter
+ list) while `check_signature` is set, or `block_size` is invalid.
Notes
-----
- CuPy caches compiled kernels on disk, keyed on the source and the compiler
- options, but not on the files pulled in by ``#include "..."``. At compile
- time the headers are resolved (:attr:`included_headers`) and a define with
- the hash of their contents is added to the options
- (:meth:`compile_options`), so editing a header recompiles the kernel; this
- includes cunumpy's own headers (``#include ``), so an
- upgrade that changes them recompiles too.
+ CuPy's disk cache is keyed on the source and options, not on included
+ files. :meth:`compile_options` therefore adds a define with the hash of
+ the :attr:`included_headers`, so editing a header (also one of cunumpy's)
+ recompiles the kernel.
Examples
--------
- >>> axpy = CudaKernel(r'''
+ >>> axpy = xp.kernels.CudaKernel(r'''
... extern "C" __global__
... void axpy(double a, const double* x, double* y, int n) {
... int i = blockDim.x * blockIdx.x + threadIdx.x;
... if (i < n) y[i] += a * x[i];
... }''', "axpy")
+ >>> axpy.launch_shape(1000)
+ ((8,), (128,))
>>> axpy(2.0, x, y, x.size, n_threads=x.size) # doctest: +SKIP
"""
@@ -1974,20 +2016,34 @@ def from_file(
suffix: str = "_cuda.cu",
**kwargs: Any,
) -> CudaKernel:
- """Load the CUDA source from a file.
+ """Create a kernel from a CUDA source file.
Parameters
----------
- path : str | Path
+ path : str or Path
Path of the CUDA source file.
- name : str | None
+ name : str, optional
Name of the kernel function; by default the file name without
`suffix` (``axpy_cuda.cu`` -> ``axpy``).
- suffix : str
- File name suffix stripped to get the default kernel name.
+ suffix : str, optional
+ File name suffix stripped to get the default name.
**kwargs
- Passed on to :class:`CudaKernel`. The directory of the file is
- always added to ``include_dirs`` and is the ``source_dir``.
+ Passed on to :class:`CudaKernel`. The file's directory is added
+ to `include_dirs` and is the `source_dir`.
+
+ Returns
+ -------
+ CudaKernel
+ The kernel.
+
+ Raises
+ ------
+ ValueError
+ If `name` is not given and the file name does not end in `suffix`.
+
+ Examples
+ --------
+ >>> push = xp.kernels.CudaKernel.from_file("push/push_cuda.cu") # doctest: +SKIP
"""
path = Path(path)
if name is None:
@@ -2002,30 +2058,34 @@ def from_file(
@classmethod
def all_from_file(cls, path: str | Path, **kwargs: Any) -> dict[str, CudaKernel]:
- """Load every ``__global__`` function of a file as a kernel.
+ """Create one kernel per ``__global__`` function of a file.
- For files that group several small kernels. The kernels share the
- source (and the compile options), so CuPy compiles the file once and
- the kernels are functions of the same compiled module.
+ The kernels share the source and options, so CuPy compiles the file
+ once.
Parameters
----------
- path : str | Path
+ path : str or Path
Path of the CUDA source file.
**kwargs
- Passed on to :class:`CudaKernel` for every kernel. The directory of
- the file is always added to ``include_dirs``.
+ Passed on to :class:`CudaKernel` for every kernel. The file's
+ directory is added to `include_dirs`.
Returns
-------
- dict[str, CudaKernel]
- One kernel per ``__global__`` function, by name, in the order of
- the source (see :func:`cuda_kernel_names`).
+ dict of str to CudaKernel
+ The kernels by name, in source order (see
+ :func:`~cunumpy.cuda.cuda_kernel_names`).
Raises
------
ValueError
If the file defines no ``__global__`` function.
+
+ Examples
+ --------
+ >>> kernels = xp.kernels.CudaKernel.all_from_file("small_kernels.cu") # doctest: +SKIP
+ >>> kernels["scale"](x, 2.0, x.size) # doctest: +SKIP
"""
path = Path(path)
source = path.read_text()
@@ -2059,7 +2119,7 @@ def name(self) -> str:
@property
def expression(self) -> str:
- """The compiled function: `name`, or its template instantiation."""
+ """The compiled function: `name`, or its instantiation such as ``"scale"``."""
if self._template_args is None:
return self._name
args = ", ".join(_template_arg(a) for a in self._template_args)
@@ -2072,16 +2132,12 @@ def source(self) -> str:
@property
def block_size(self) -> int | tuple[int, ...]:
- """Threads per block: an integer for 1D blocks, else a tuple."""
+ """The threads per block: an integer for 1D blocks, else a tuple."""
return self._block[0] if len(self._block) == 1 else self._block
@property
def options(self) -> tuple[str, ...]:
- """NVRTC compiler options as given, including ``-I`` include directories.
-
- The debug options and the header hash define are not part of them;
- they are added at compile time, see :meth:`compile_options`.
- """
+ """NVRTC options as given, with ``-I`` for `include_dirs` (see :meth:`compile_options`)."""
return self._options
@property
@@ -2090,11 +2146,13 @@ def debug(self) -> bool | None:
return self._debug
def debug_active(self) -> bool:
- """Whether debug mode applies to this kernel now.
+ """Return whether debug mode applies to this kernel now.
- The kernel's own setting if it was created with ``debug=True`` or
- ``debug=False``, else the global setting (:func:`cunumpy.cuda.get_cuda_debug`),
- read at the time of the call.
+ Returns
+ -------
+ bool
+ The kernel's own `debug` setting, else the current global one
+ (:func:`~cunumpy.cuda.get_cuda_debug`).
"""
if self._debug is not None:
return self._debug
@@ -2104,7 +2162,7 @@ def debug_active(self) -> bool:
@property
def include_dirs(self) -> tuple[Path, ...]:
- """Directories searched for ``#include`` files."""
+ """The directories searched for ``#include`` files."""
return self._include_dirs
@property
@@ -2114,16 +2172,7 @@ def source_dir(self) -> Path | None:
@property
def included_headers(self) -> tuple[Path, ...]:
- """The header files the source includes, recursively.
-
- Quoted includes (``#include "..."``) are resolved in `source_dir`,
- `include_dirs` and cunumpy's header directory
- (:func:`cuda_include_dir`); cunumpy's shipped headers are tracked
- also when included in angle brackets (``#include ``).
- Resolved at every access (see :func:`resolve_includes`), so the result
- follows the files on disk. Empty if the source includes no project or
- cunumpy headers.
- """
+ """The header files the source includes, resolved at every access (see :func:`~cunumpy.cuda.resolve_includes`)."""
return tuple(
resolve_includes(
self._source,
@@ -2134,18 +2183,18 @@ def included_headers(self) -> tuple[Path, ...]:
)
def compile_options(self) -> tuple[str, ...]:
- """The NVRTC options a compilation now would use.
-
- :attr:`options`, followed by ``-I`` for cunumpy's own header directory
- (:func:`cunumpy.cuda.cuda_include_dir`) unless already present, then
- :data:`DEBUG_OPTIONS` (``-lineinfo`` and
- ``-DCUNUMPY_BOUNDS_CHECK``) if :meth:`debug_active` and they are not
- already among the options (``-G`` is not added: NVRTC does not support
- it), and, if the source includes header files,
- ``-DCUNUMPY_INCLUDE_HASH=0x`` with the hash of the contents of
- :attr:`included_headers` (see :func:`include_hash`). CuPy keys its
- kernel cache on the options, so a changed header means a recompile,
- also for a header shipped with cunumpy that changed in an upgrade.
+ """Return the NVRTC options a compilation now would use.
+
+ These are :attr:`options`, ``-I`` for cunumpy's header directory,
+ :data:`~cunumpy.cuda.DEBUG_OPTIONS` if :meth:`debug_active`, and, if
+ the source includes headers, ``-DCUNUMPY_INCLUDE_HASH=0x`` of
+ their contents (:func:`~cunumpy.cuda.include_hash`). CuPy keys its
+ kernel cache on the options, so a changed header means a recompile.
+
+ Returns
+ -------
+ tuple of str
+ The options.
"""
options = self._options
# cunumpy's own headers (, ...) are always found
@@ -2166,20 +2215,7 @@ def structs(self) -> tuple[CudaStruct, ...]:
@property
def n_threads_from(self) -> Callable[[tuple[Any, ...]], Any] | None:
- """Default launch size: a function of the positional arguments, or None.
-
- Called with the tuple of arguments of a launch that gives neither
- `n_threads` nor `grid`, and returns `n_threads` (an integer or a
- tuple), e.g. ``lambda args: args[2].n_markers`` for a kernel whose
- third argument is a struct argument object with the marker count.
- Default ``"auto"`` uses the first array's leading axes, matching the
- launch block dimensions; 1D launches use its first axis, one thread per
- row. Supported argument objects are searched in field/argument order.
- ``"first_array"`` always uses its first axis, ``"last_axis"`` its last
- one (one thread per entry of a component-major ``(ncomp, N)`` or
- ``(N,)`` array). None disables inference.
- Settable, also on the ``cuda_kernel`` of a :class:`~cunumpy.kernels.Kernel`.
- """
+ """The launch size inference: a callable of the argument tuple, or None; settable as in the constructor."""
return self._n_threads_from
@n_threads_from.setter
@@ -2207,15 +2243,7 @@ def _default_n_threads(self, args: tuple[Any, ...]) -> int | tuple[int, ...]:
@property
def check_finite(self) -> bool:
- """Whether every launch checks the floating-point arrays for NaN or inf.
-
- After the launch (synchronized), every floating-point or complex array
- among the arguments, including the array fields of struct argument
- objects, is scanned, and a non-finite value raises ``RuntimeError``
- naming the kernel and the argument. Costs a synchronization and one
- pass over the arrays per launch; for debugging (e.g. a pusher writing
- NaN velocities), not for production. Settable.
- """
+ """Whether every launch synchronizes and raises on NaN or inf in its float arrays (debugging; settable)."""
return self._check_finite
@check_finite.setter
@@ -2238,21 +2266,22 @@ def is_compiled(self) -> bool:
return bool(self._compiled) and _current_device() in self._compiled
def compile(self, *, log_stream: Any = None) -> Any:
- """Compile the kernel now (it is otherwise compiled on the first call).
+ """Compile the kernel on the current device now, instead of at the first call.
+
+ Compiled once per device and cached (also on disk by CuPy) with the
+ :meth:`compile_options` of now: a kernel compiled before debug mode
+ was enabled keeps its options (see :meth:`recompile`). A failed
+ compilation can be retried.
- The options are :meth:`compile_options`, evaluated now: a kernel
- compiled before debug mode was enabled keeps its options. Call
- :meth:`recompile` to refresh headers/options explicitly. Compilation
- and launch resources are cached separately on each CUDA device;
- failed compilation remains retryable. ``log_stream`` receives compiler
- output (a writable file object, or None).
+ Parameters
+ ----------
+ log_stream : file-like, optional
+ Writable file object that receives the compiler output.
Returns
-------
cupy.RawKernel
- The compiled kernel; compiled once and cached (also on disk by CuPy,
- keyed on the source and :meth:`compile_options`). No launch is needed
- to surface compiler errors.
+ The compiled kernel.
Raises
------
@@ -2298,10 +2327,19 @@ def compile(self, *, log_stream: Any = None) -> Any:
return self._compiled[device].raw
def recompile(self, *, log_stream: Any = None) -> Any:
- """Recompile on the current device using current headers/debug options.
+ """Recompile on the current device with the current headers and options.
- Other devices keep their compiled kernels. A failed rebuild remains
- uncompiled and can be retried; in-flight launches must finish first.
+ Other devices keep their kernels; in-flight launches must finish first.
+
+ Parameters
+ ----------
+ log_stream : file-like, optional
+ Writable file object that receives the compiler output.
+
+ Returns
+ -------
+ cupy.RawKernel
+ The compiled kernel.
"""
from cunumpy.xp import cupy_available
@@ -2310,24 +2348,39 @@ def recompile(self, *, log_stream: Any = None) -> Any:
return self.compile(log_stream=log_stream)
def prepare_args(self, *args: Any) -> tuple[Any, ...]:
- """The arguments as passed to ``cupy.RawKernel``: flattened and checked.
+ """Return the arguments as passed to ``cupy.RawKernel``: flattened and checked.
+
+ Argument objects with ``__cuda_args__()`` are flattened. With a
+ checked signature, the argument count, every array's dtype and
+ contiguity, structs and scalars are checked, scalars are cast to the
+ declared C types and arrays for view parameters (``Array2D``)
+ are packed into pointer, shape and strides.
+
+ Parameters
+ ----------
+ *args
+ The kernel arguments.
- Argument objects with ``__cuda_args__()`` (including struct values) are
- flattened. If the signature is checked, the number of arguments, the
- dtype and C-contiguity of every array, every struct and every scalar
- are checked, Python scalars are cast to the declared C types, and arrays
- for array view parameters (``Array2D``) are packed into
- (pointer, shape, strides).
+ Returns
+ -------
+ tuple
+ The values to launch with.
Raises
------
TypeError
- Wrong number of arguments, a host array, an array of the wrong
- dtype or a non-contiguous array for a pointer parameter, an array of
- the wrong dtype or number of dimensions for an array view, a value
- of the wrong struct, or a scalar of an incompatible type.
+ If the argument count is wrong, an array is a host array, has the
+ wrong dtype, contiguity or ndim, a struct value is of another
+ struct, or a scalar has an incompatible type.
OverflowError
- A Python integer out of range of the declared integer type.
+ If an integer is out of range of the declared type.
+
+ Examples
+ --------
+ >>> scale = xp.kernels.CudaKernel(
+ ... 'extern "C" __global__ void scale(double a, int n) {}', "scale")
+ >>> scale.prepare_args(2, 10)
+ (np.float64(2.0), np.int32(10))
"""
values: list[Any] = []
for arg in args:
@@ -2354,11 +2407,37 @@ def launch_shape(
block: int | Sequence[int] | None = None,
args: Sequence[Any] | None = None,
) -> tuple[tuple[int, ...], tuple[int, ...]]:
- """The ``(grid, block)`` a call with these launch arguments uses.
+ """Return the ``(grid, block)`` a call with these launch arguments uses.
+
+ Parameters
+ ----------
+ n_threads : int or sequence of int, optional
+ Number of threads, as for :meth:`~cunumpy.kernels.CudaKernel.__call__`.
+ grid : int or sequence of int, optional
+ Number of blocks, instead of `n_threads`.
+ block : int or sequence of int, optional
+ Block shape, instead of :attr:`block_size`.
+ args : sequence, optional
+ Call arguments, to infer the launch size (:attr:`n_threads_from`)
+ when neither `n_threads` nor `grid` is given.
+
+ Returns
+ -------
+ tuple of (tuple of int, tuple of int)
+ The grid and block shapes. A grid with a zero dimension launches
+ nothing.
- Pass `args` to use default thread-count inference when neither
- `n_threads` nor `grid` is given. See :meth:`__call__`. A grid with a
- zero dimension launches nothing.
+ Raises
+ ------
+ TypeError
+ If not exactly one of `n_threads` and `grid` is given or inferred.
+
+ Examples
+ --------
+ >>> kernel = xp.kernels.CudaKernel(
+ ... 'extern "C" __global__ void f(double* a) {}', "f", block_size=(16, 16))
+ >>> kernel.launch_shape((100, 50))
+ ((7, 4), (16, 16))
"""
block_shape = (
self._block
@@ -2401,41 +2480,38 @@ def __call__(
shared_mem: int = 0,
stream: Any = None,
) -> None:
- """Launch the kernel.
+ """Launch the kernel (asynchronously, unless in debug mode).
- The launch shape is given either by `n_threads` (the grid is the number
- of threads divided by the block size, rounded up, per dimension) or by
- an explicit `grid`. With neither, it is inferred from the first array's
- leading axes (or the configured `n_threads_from` callback).
+ The grid is `n_threads` divided by the block, rounded up, or an
+ explicit `grid`; with neither, the size is inferred
+ (:attr:`n_threads_from`).
Parameters
----------
*args
- Kernel arguments: CuPy arrays, scalars and argument objects with
- ``__cuda_args__()``, see :meth:`prepare_args`.
- n_threads : int | Sequence[int] | None
- Number of threads in 1 to 3 dimensions, e.g. ``n`` or ``(nx, ny)``.
- With a 1D block size, the block is ``(block_size, 1, ...)``.
- grid : int | Sequence[int] | None
+ Kernel arguments: CuPy arrays, scalars and argument objects, see
+ :meth:`prepare_args`.
+ n_threads : int or sequence of int, optional
+ Number of threads in 1 to 3 dimensions, e.g. ``(nx, ny)``. With a
+ 1D block size the block is ``(block_size, 1, ...)``.
+ grid : int or sequence of int, optional
Number of blocks in 1 to 3 dimensions, instead of `n_threads`.
- block : int | Sequence[int] | None
- Block shape for this call, instead of the kernel's `block_size`.
- shared_mem : int
- Dynamic shared memory per block, in bytes (``extern __shared__``).
- stream : cupy.cuda.Stream | None
- Stream to launch on; the current stream if None.
+ block : int or sequence of int, optional
+ Block shape for this call, instead of :attr:`block_size`.
+ shared_mem : int, optional
+ Dynamic shared memory per block in bytes (``extern __shared__``).
+ stream : cupy.cuda.Stream, optional
+ Stream to launch on; the current stream by default.
Raises
------
RuntimeError
- In debug mode (see :meth:`debug_active`), an asynchronous CUDA
- error found when synchronizing the stream after the launch, e.g. an
- illegal memory access; the CuPy error is chained. Without debug
- mode, such an error surfaces at a later synchronization (a
- ``.get()``, an MPI call, ...), not necessarily in this kernel.
+ In debug mode, a CUDA error found by synchronizing after the
+ launch (CuPy's error chained); otherwise such errors surface at a
+ later synchronization. Also a NaN or inf with :attr:`check_finite`.
ValueError
- Dimensions, threads or static plus dynamic shared memory exceed
- this device/kernel's limits, or arrays/stream belong to another device.
+ If the launch exceeds the device or kernel limits (dimensions,
+ threads, shared memory), or arrays or stream are on another device.
"""
grid_shape, block_shape = self.launch_shape(
n_threads, grid=grid, block=block, args=args
@@ -2542,13 +2618,7 @@ def _synchronize_after_launch(
grid: tuple[int, ...],
block: tuple[int, ...],
) -> None:
- """Wait for the launch and re-raise a CUDA error naming this kernel.
-
- Skipped while the stream is being captured into a CUDA graph: the
- launch is only recorded then, and synchronizing would invalidate the
- capture. Errors of the captured kernels surface when the graph is
- launched (in debug mode, synchronize after ``graph.launch()``).
- """
+ """Wait for the launch and re-raise a CUDA error naming this kernel (not while graph-capturing)."""
if stream is None:
import cupy as cp
@@ -2585,24 +2655,26 @@ def _is_capturing(stream: Any) -> bool:
class CudaKernelVariants:
- """Kernels generated per variant (e.g. per dtype and dimension), compiled once.
+ """Kernels generated per variant (e.g. per dtype and dimension), created once.
- For kernels whose source is generated for each variant, e.g. a stencil
- product for ``ndim`` in 1 to 3 and several dtypes. The `factory` is called
- the first time a variant is requested; the kernel is cached per key.
+ For kernels whose source is generated per variant. The `factory` is called
+ the first time a key is requested, and the kernel is cached per key.
Parameters
----------
- factory : Callable[..., CudaKernel]
- Creates the kernel for a variant from its key, e.g.
- ``lambda ndim, dtype: CudaKernel(make_source(ndim, ctype_of(dtype)), "f")``.
+ factory : callable
+ Returns the :class:`CudaKernel` for a variant, called with the key.
Examples
--------
- >>> matvec = CudaKernelVariants(
- ... lambda ndim, dtype: CudaKernel(source(ndim, ctype_of(dtype)), "matvec")
- ... )
- >>> matvec.get(3, np.float64)(mat, x, out, n_threads=out.size) # doctest: +SKIP
+ >>> def make_scale(dtype):
+ ... ctype = xp.cuda.ctype_of(dtype)
+ ... return xp.kernels.CudaKernel(
+ ... f'extern "C" __global__ void scale({ctype}* x, {ctype} a) {{}}', "scale")
+ >>> scale = xp.kernels.CudaKernelVariants(make_scale)
+ >>> scale.get(np.float32).signature[1].ctype
+ 'float'
+ >>> scale.compile_all([(np.float64,)]) # doctest: +SKIP
"""
def __init__(self, factory: Callable[..., CudaKernel]) -> None:
@@ -2613,7 +2685,23 @@ def __repr__(self) -> str:
return f"CudaKernelVariants({len(self._kernels)} variants)"
def get(self, *key: Hashable) -> CudaKernel:
- """The kernel for the variant `key`, created on first use."""
+ """Return the kernel for a variant, created on first use.
+
+ Parameters
+ ----------
+ *key
+ The variant key, passed to the factory.
+
+ Returns
+ -------
+ CudaKernel
+ The kernel.
+
+ Raises
+ ------
+ TypeError
+ If the factory does not return a :class:`CudaKernel`.
+ """
kernel = self._kernels.get(key)
if kernel is None:
kernel = self._factory(*key)
@@ -2632,7 +2720,13 @@ def __iter__(self) -> Iterator[tuple[Hashable, ...]]:
return iter(list(self._kernels))
def keys(self) -> list[tuple[Hashable, ...]]:
- """The keys of the variants created so far."""
+ """Return the keys of the variants created so far.
+
+ Returns
+ -------
+ list of tuple
+ The keys.
+ """
return list(self._kernels)
def compile_all(
@@ -2643,14 +2737,15 @@ def compile_all(
) -> None:
"""Compile the given variants (created if needed) and all existing ones.
+ All are compiled even if one fails; the first error is raised after.
+
Parameters
----------
- keys : Iterable[Sequence]
- Keys of variants to create and compile now, e.g. ``[(3, np.float64)]``.
- jobs : int | None
+ keys : iterable of sequence, optional
+ Keys of variants to create now, e.g. ``[(3, np.float64)]``.
+ jobs : int or None, optional
Number of variants compiled at a time, in threads (NVRTC releases
- the GIL); None for the number of CPUs. See
- :meth:`KernelCatalog.compile_all `.
+ the GIL); None for the number of CPUs. 1 by default.
"""
for key in keys:
self.get(*key)
diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py
index b1f6cc5..77c72cb 100644
--- a/src/cunumpy/_device.py
+++ b/src/cunumpy/_device.py
@@ -14,7 +14,24 @@
def set_device(device_id: int) -> None:
- """Select the active CUDA device for the current process (no-op on NumPy)."""
+ """Select the active CUDA device of this process.
+
+ Does nothing on the NumPy backend.
+
+ Parameters
+ ----------
+ device_id : int
+ A device id valid for this process.
+
+ See Also
+ --------
+ set_device_for_rank : Select the device from an MPI rank.
+ bind_local_device : Select the device from the node-local rank.
+
+ Examples
+ --------
+ >>> xp.cuda.set_device(0) # no-op on NumPy
+ """
if array_backend.backend == "cupy":
import cupy as cp
@@ -22,11 +39,19 @@ def set_device(device_id: int) -> None:
def device_count() -> int:
- """Number of visible CUDA devices.
+ """Return the number of visible CUDA devices.
+
+ Reports the visible hardware, independent of the active backend: it can be
+ positive while NumPy is selected.
+
+ Returns
+ -------
+ int
+ The device count, or 0 if CuPy/CUDA is unavailable or the query fails.
- Returns 0 on the NumPy backend, or if CuPy/CUDA is not available.
- Independent of the currently active backend -- this reports what
- hardware is visible, not what `xp.xp` currently dispatches to.
+ Examples
+ --------
+ >>> n = xp.cuda.device_count() # 0 on a machine without GPUs
"""
if not cupy_available():
return 0
@@ -40,16 +65,34 @@ def device_count() -> int:
def set_device_for_rank(rank: int, devices_per_node: int | None = None) -> int:
- """Select a CUDA device for an MPI rank, round-robin across the node.
+ """Select device ``rank % devices_per_node`` for an MPI rank.
- Convenience for one-rank-per-GPU codes: computes
- ``device_id = rank % devices_per_node`` and calls `set_device()` with
- it. `devices_per_node` defaults to `device_count()`. Returns the
- selected device id, or 0 as a no-op if there are no visible devices.
+ For one-rank-per-GPU codes whose ranks fill each node in contiguous blocks;
+ otherwise use :func:`bind_local_device` or :func:`set_device`.
- This assumes ranks map to devices in contiguous blocks per node (i.e.
- local rank == ``rank % devices_per_node``); codes with a different
- rank-to-device layout should call `set_device()` directly instead.
+ Parameters
+ ----------
+ rank : int
+ The MPI rank of this process.
+ devices_per_node : int, optional
+ Devices per node; :func:`device_count` by default.
+
+ Returns
+ -------
+ int
+ ``rank % devices_per_node``, selected on the CuPy backend only (on
+ NumPy the id is returned without selecting anything). 0, with nothing
+ selected, if the number of devices per node is 0, e.g. when
+ `devices_per_node` is omitted and no device is visible.
+
+ See Also
+ --------
+ bind_local_device : Uses the node-local rank instead.
+
+ Examples
+ --------
+ >>> xp.cuda.set_device_for_rank(3, devices_per_node=4)
+ 3
"""
n = devices_per_node if devices_per_node is not None else device_count()
if n == 0:
@@ -61,20 +104,23 @@ def set_device_for_rank(rank: int, devices_per_node: int | None = None) -> int:
def bind_local_device() -> int | None:
- """Bind this process to one GPU of its node, by node-local rank.
+ """Bind this process to device ``local_rank() % device_count()`` of its node.
- Selects device ``local_rank() % device_count()`` and creates its CUDA
- context. Call it before ``MPI_Init`` (i.e. before importing
- ``mpi4py.MPI``), so that CUDA-aware MPI sees the right device. Without it,
- every rank on a node would use device 0. If the launcher already restricts
- each rank to its own device with ``CUDA_VISIBLE_DEVICES``, every process
- sees a single device and selects it.
+ Also creates the CUDA context. Call it before ``MPI_Init`` (before importing
+ ``mpi4py.MPI``) so that CUDA-aware MPI sees the right device; otherwise every
+ rank of a node uses device 0. With ``CUDA_VISIBLE_DEVICES`` set per rank, each
+ process sees and selects one device. Needs no MPI rank, unlike
+ :func:`set_device_for_rank`.
Returns
-------
- int | None
- The selected device id, or None on the NumPy backend or if no device
- is available.
+ int or None
+ The selected device id, or None on the NumPy backend or without devices.
+
+ Examples
+ --------
+ >>> print(xp.cuda.bind_local_device()) # on NumPy
+ None
"""
if array_backend.backend != "cupy":
return None
@@ -91,10 +137,20 @@ def bind_local_device() -> int | None:
def memory_info() -> tuple[int, int] | None:
- """Return `(free, total)` bytes of memory on the active CUDA device.
+ """Return the free and total memory of the active CUDA device.
+
+ Queries the CUDA runtime, so it covers the whole device, not only CuPy's
+ memory pool.
+
+ Returns
+ -------
+ tuple of int or None
+ ``(free, total)`` in bytes, or None on the NumPy backend.
- Returns `None` on the NumPy backend. Queries the CUDA runtime directly,
- so it reflects the whole device rather than just CuPy's memory pool.
+ Examples
+ --------
+ >>> print(xp.cuda.memory_info()) # on NumPy
+ None
"""
if array_backend.backend != "cupy":
return None
@@ -115,26 +171,30 @@ def max_shared_memory_per_block(
*,
opt_in: bool = False,
) -> int:
- """Bytes of shared memory a block of a CUDA kernel may use on `device`.
+ """Return the bytes of shared memory a CUDA block may use on `device`.
- Use it to decide whether a per-block buffer (e.g. a copy of a small grid
- for a deposit) fits, instead of a hard-coded limit.
+ Use it to decide whether a per-block buffer (e.g. a copy of a small grid for
+ a deposit) fits, instead of a hard-coded limit.
Parameters
----------
- device : int | None
+ device : int, optional
CUDA device id; the current device by default.
- opt_in : bool
- The larger limit a kernel can opt in to on newer GPUs (e.g. 99 or 227
- KiB). Using more than the default 48 KiB needs the kernel attribute
- ``max_dynamic_shared_size_bytes`` set on the compiled ``cupy.RawKernel``
- (``kernel.compile()``).
+ opt_in : bool, optional
+ Return the larger limit a kernel can opt in to on newer GPUs (e.g. 99 or
+ 227 KiB). Using more than 48 KiB needs ``max_dynamic_shared_size_bytes``
+ set on the compiled ``cupy.RawKernel``.
Returns
-------
int
- The limit in bytes; :data:`DEFAULT_SHARED_MEMORY_PER_BLOCK` if CuPy is
- not available (so that code choosing a GPU strategy runs everywhere).
+ The limit in bytes; :data:`DEFAULT_SHARED_MEMORY_PER_BLOCK` without CuPy,
+ so that code choosing a GPU strategy runs everywhere.
+
+ Examples
+ --------
+ >>> xp.cuda.max_shared_memory_per_block() # without CuPy
+ 49152
"""
if not cupy_available():
return DEFAULT_SHARED_MEMORY_PER_BLOCK
@@ -146,11 +206,15 @@ def max_shared_memory_per_block(
def free_memory() -> None:
- """Release all free blocks held by CuPy's memory pools (no-op on NumPy).
+ """Release the free blocks of CuPy's device and pinned memory pools.
- CuPy caches freed device (and pinned host) memory in pools rather than
- returning it to the driver/OS immediately, which can look like a leak
- in long-running processes. Call this to give it back.
+ CuPy caches freed memory for reuse, which can look like a leak in
+ long-running processes. Blocks still used by live arrays are kept. Does
+ nothing on the NumPy backend.
+
+ Examples
+ --------
+ >>> xp.cuda.free_memory() # no-op on NumPy
"""
if array_backend.backend == "cupy":
import cupy as cp
@@ -160,12 +224,28 @@ def free_memory() -> None:
def pin_memory(array: Any) -> Any:
- """Copy a host array into pinned (page-locked) CUDA host memory.
+ """Copy a host array into pinned (page-locked) host memory.
+
+ Pinned memory transfers to and from the GPU faster than pageable memory.
+
+ Parameters
+ ----------
+ array : array_like
+ A host array; use :func:`~cunumpy.to_numpy` first if it may be on the GPU.
+
+ Returns
+ -------
+ numpy.ndarray
+ A copy of `array` backed by pinned memory.
+
+ Raises
+ ------
+ ImportError
+ If CuPy is not available.
- Pinned memory transfers to/from the GPU faster than regular pageable
- memory, since the driver can DMA it directly. `array` must already be
- on the host (use `to_numpy()` first if it may be on the GPU). Raises
- `ImportError` if CuPy is not available.
+ Examples
+ --------
+ >>> pinned = xp.cuda.pin_memory(np.zeros(1000)) # doctest: +SKIP
"""
if not cupy_available():
raise ImportError("CuPy is not available or not functional.")
@@ -181,16 +261,31 @@ def pin_memory(array: Any) -> Any:
@contextmanager
def stream(existing: Any = None) -> Generator[Any, None, None]:
- """Context manager for a CUDA stream, to overlap transfers and compute.
+ """Run the work issued in a ``with`` block on a CUDA stream.
- On the CuPy backend, operations issued inside the block are enqueued on
- a new, non-blocking stream rather than the default one. Call
- `xp.synchronize()` (or the yielded stream's own `.synchronize()`) before
- reading results computed inside the block. No-op on the NumPy backend,
- where it yields `None`.
+ On CuPy, yields a new non-blocking stream (or selects `existing`); call
+ :func:`~cunumpy.synchronize` or the stream's ``synchronize()`` before reading
+ results computed in the block. On NumPy, yields None and does nothing, so
+ portable code calls :func:`~cunumpy.synchronize` rather than methods of the
+ yielded stream.
- Pass a reusable stream from :func:`cunumpy.cuda.create_stream` as `existing`
- to select it temporarily; this also works with its synchronous CPU equivalent.
+ Parameters
+ ----------
+ existing : stream, optional
+ A reusable stream from :func:`create_stream` to select instead of a new
+ one; also works with its CPU equivalent :class:`HostStream`.
+
+ Yields
+ ------
+ stream or None
+ The selected stream, or None on the NumPy backend without `existing`.
+
+ Examples
+ --------
+ >>> with xp.cuda.stream() as s:
+ ... y = xp.ones(3) * 2
+ >>> print(s) # on NumPy
+ None
"""
if existing is not None:
with existing:
@@ -211,11 +306,7 @@ def stream(existing: Any = None) -> Generator[Any, None, None]:
def _debug_from_env(value: str | None) -> bool:
- """Whether the value of ``CUNUMPY_CUDA_DEBUG`` enables the debug mode.
-
- ``"1"``, ``"true"``, ``"yes"`` and ``"on"`` (any case, surrounding
- whitespace ignored) enable it; anything else, including unset, does not.
- """
+ """Whether a ``CUNUMPY_CUDA_DEBUG`` value enables the debug mode."""
if value is None:
return False
return value.strip().lower() in _CUDA_DEBUG_TRUE
@@ -227,30 +318,70 @@ def _debug_from_env(value: str | None) -> bool:
def set_cuda_debug(enabled: bool) -> None:
"""Enable or disable the CUDA debug mode globally.
- In debug mode, `CudaKernel`s created with ``debug=None`` (the default)
- compile with ``-lineinfo`` and ``-DCUNUMPY_BOUNDS_CHECK``, and synchronize
- the stream after every launch, re-raising an asynchronous CUDA error as a
- ``RuntimeError`` naming the kernel that caused it. The setting is read at
- every launch, so it also applies to kernels created earlier; only their
- compile options are fixed once they are compiled (call ``compile()`` again
- or create the kernels after enabling debug mode). Initialised from the
- environment variable ``CUNUMPY_CUDA_DEBUG`` (``1``, ``true``, ``yes`` or
- ``on``) at import.
+ In debug mode, a :class:`~cunumpy.kernels.CudaKernel` created with
+ ``debug=None`` (the default) compiles with ``-lineinfo`` and
+ ``-DCUNUMPY_BOUNDS_CHECK`` and synchronizes after every launch, re-raising an
+ asynchronous CUDA error as a ``RuntimeError`` naming the kernel. The setting
+ is read at every launch, but compile options are fixed once a kernel is
+ compiled. Initialized at import from ``CUNUMPY_CUDA_DEBUG`` (``1``, ``true``,
+ ``yes`` or ``on``). See :doc:`/kernels/debugging`.
+
+ Parameters
+ ----------
+ enabled : bool
+ Whether to enable the debug mode.
+
+ See Also
+ --------
+ cuda_debug : Enable it temporarily.
+
+ Examples
+ --------
+ >>> xp.cuda.set_cuda_debug(True)
+ >>> xp.cuda.get_cuda_debug()
+ True
+ >>> xp.cuda.set_cuda_debug(False)
"""
global _cuda_debug
_cuda_debug = bool(enabled)
def get_cuda_debug() -> bool:
- """Whether the CUDA debug mode is enabled globally, see `set_cuda_debug`."""
+ """Return whether the CUDA debug mode is enabled globally.
+
+ Returns
+ -------
+ bool
+ The setting of :func:`set_cuda_debug`.
+
+ Examples
+ --------
+ >>> xp.cuda.get_cuda_debug()
+ False
+ """
return _cuda_debug
@contextmanager
def cuda_debug(enabled: bool = True) -> Generator[None, None, None]:
- """Temporarily enable (or disable) the CUDA debug mode.
+ """Enable (or disable) the CUDA debug mode inside a ``with`` block.
- Restores the previous setting on exit, see `set_cuda_debug`.
+ Restores the previous setting on exit, see :func:`set_cuda_debug`.
+
+ Parameters
+ ----------
+ enabled : bool, optional
+ The setting inside the block, True by default.
+
+ Yields
+ ------
+ None
+
+ Examples
+ --------
+ >>> with xp.cuda.cuda_debug():
+ ... xp.cuda.get_cuda_debug()
+ True
"""
previous = _cuda_debug
set_cuda_debug(enabled)
diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py
index 305082f..d7bb198 100644
--- a/src/cunumpy/_dispatch.py
+++ b/src/cunumpy/_dispatch.py
@@ -1,14 +1,10 @@
"""Pairs of host and CUDA kernels, chosen by the backend or by the arguments.
-A :class:`Kernel` holds a host kernel (a :class:`~cunumpy.kernels.PyccelKernel`, e.g. a
-Pyccel-compiled function) and, optionally, its 1:1 corresponding CUDA kernel
-(:class:`~cunumpy.kernels.CudaKernel`). It calls the host kernel on the NumPy backend and
-the CUDA kernel on the CuPy backend (or, with ``dispatch="arrays"``, the CUDA
-kernel for device arguments and the host kernel for host arguments), so a code
-base can port its kernels to CUDA one by one.
-
-:class:`KernelCatalog` collects such pairs from a package laid out with one
-folder per kernel::
+A :class:`~cunumpy.kernels.Kernel` holds a host kernel (a
+:class:`~cunumpy.kernels.PyccelKernel`) and, optionally, its CUDA counterpart
+(a :class:`~cunumpy.kernels.CudaKernel`), so a code base can port its kernels
+to CUDA one by one. :class:`~cunumpy.kernels.KernelCatalog` collects such
+pairs from a package with one folder per kernel (see :doc:`/kernels/dispatch`)::
my_kernels/
├── __init__.py # catalog = KernelCatalog.from_package(__name__)
@@ -56,27 +52,18 @@
def _on_device(arg: Any) -> bool:
- """Whether a kernel argument lives on the GPU.
-
- A CuPy array, or a CUDA argument object (one with ``__cuda_args__``, e.g. a
- :class:`~cunumpy.arguments.CudaArguments` or a struct value).
- """
+ """Whether `arg` is a CuPy array or a CUDA argument object (``__cuda_args__``)."""
return _is_device_array(arg) or callable(getattr(type(arg), "__cuda_args__", None))
-#: Longest name a Fortran compiler accepts. pyccel names the wrapper module of a
-#: host kernel module ``bind_c_``, so a kernel module name longer than
-#: 63 - len("bind_c_") characters cannot be compiled with the Fortran backend.
+#: Longest name a Fortran compiler accepts. Pyccel names the wrapper module of a
+#: host kernel module ``bind_c_``, so a module name longer than
+#: ``63 - len("bind_c_")`` characters cannot be compiled with the Fortran backend.
FORTRAN_NAME_LIMIT = 63
def _pyccel_stub_parameters(function: Any) -> list[str] | None:
- """Parameter names of a pyccel-compiled function, from its ``.pyi`` stub.
-
- pyccel writes ``__pyccel__/.pyi`` next to the compiled extension
- module; the compiled function itself has no Python signature. Returns
- None if the stub or the function is not found.
- """
+ """Read the parameters of a Pyccel-compiled function from its ``.pyi`` stub."""
module = getattr(function, "__self__", None) # extension functions: the module
if not isinstance(module, ModuleType):
name = getattr(function, "__module__", None)
@@ -100,7 +87,7 @@ def _pyccel_stub_parameters(function: Any) -> list[str] | None:
def _import_loader(module: str, name: str) -> Callable[[], Callable[..., Any]]:
- """Loads function `name` of `module` on first use (the import may fail)."""
+ """Return a loader of function `name` of `module` (imported on first use)."""
return lambda: getattr(importlib.import_module(module), name)
@@ -112,7 +99,7 @@ def _fallback_loader(
),
name: str,
) -> dict[str, Callable[[], Callable[..., Any]]]:
- """`from_package`'s `host_fallback` for kernel `name`, as its NumPy implementation."""
+ """Return the `host_fallback` of kernel `name` as its ``"numpy"`` loader."""
fallback = (
host_fallback(name)
if callable(host_fallback)
@@ -122,7 +109,7 @@ def _fallback_loader(
def _positional_parameters(function: Any) -> list[str] | None:
- """Names of the positional parameters of `function`, or None without a signature."""
+ """Return the positional parameter names of `function`, or None if unknown."""
try:
signature = inspect.signature(function)
except (TypeError, ValueError):
@@ -135,7 +122,7 @@ def _positional_parameters(function: Any) -> list[str] | None:
def _source_root(package: str) -> Path:
- """The directory containing the top-level package of `package`."""
+ """Return the directory containing the top-level package of `package`."""
top = importlib.import_module(package.partition(".")[0])
if top.__file__ is not None:
return Path(top.__file__).parent.parent
@@ -143,50 +130,67 @@ def _source_root(package: str) -> Path:
class Kernel:
- """A host kernel and its CUDA counterpart; calls the one matching the backend.
+ """A host kernel and its CUDA counterpart; a call runs the one matching the backend.
+
+ See :doc:`/kernels/dispatch` for the porting workflow.
Parameters
----------
- host_kernel : PyccelKernel | callable
- The host kernel, called on the NumPy backend. A plain callable is
- wrapped in a :class:`~cunumpy.kernels.PyccelKernel`.
- cuda_kernel : CudaKernel | None
- The CUDA kernel, called on the CuPy backend; None if it has not been
- written yet.
- name : str | None
- Name of the kernel; defaults to the name of the host kernel.
- missing_cuda : {"raise", "fallback"}
- What happens on the CuPy backend if there is no CUDA kernel: ``"raise"``
- raises ``NotImplementedError``; ``"fallback"`` calls the host kernel
- through :class:`~cunumpy.kernels.PyccelKernel`, which copies the arrays to the
- host and back at every call (a warning is emitted once).
- cuda_path : str | Path | None
- Where the CUDA kernel is expected, for the error message if it is missing.
- host_options : Mapping[str, Any] | None
- Keyword arguments for the :class:`~cunumpy.kernels.PyccelKernel` that wraps a
- plain callable `host_kernel`, e.g. ``{"object_modules": ("my_pkg.",),
- "outputs": (2,)}``; they matter for the fallback on the CuPy backend.
- Not allowed if `host_kernel` already is a ``PyccelKernel``.
- dispatch : {"backend", "arrays"}
- How a call chooses between the two kernels. ``"backend"`` (default):
- the CUDA kernel on the CuPy backend, the host kernel on the NumPy
- backend. ``"arrays"``: the CUDA kernel if any top-level argument lives
- on the GPU (a CuPy array, or a device-only argument object such as a
- :class:`~cunumpy.arguments.CudaArguments` or a struct value), else the host
- kernel, whatever the backend. Use ``"arrays"`` when a code deliberately
- hands host arrays to kernels while CuPy is active (diagnostics, MPI
- staging, CPU fallbacks): the host kernel then runs on the host arrays
- instead of the CUDA kernel rejecting them.
+ host_kernel : PyccelKernel or callable
+ The host kernel. A plain callable is wrapped in a
+ :class:`~cunumpy.kernels.PyccelKernel`.
+ cuda_kernel : CudaKernel, optional
+ The CUDA kernel; None if it has not been written yet.
+ name : str, optional
+ Name of the kernel; by default the name of the host kernel.
+ missing_cuda : {"raise", "fallback"}, optional
+ Behavior on the device without CUDA kernel: ``"raise"`` (default)
+ raises ``NotImplementedError``; ``"fallback"`` calls the host kernel,
+ which copies the arrays to the host and back at every call (a
+ ``RuntimeWarning`` is emitted once).
+ cuda_path : str or Path, optional
+ Where the CUDA kernel is expected, named in the error if it is missing.
+ host_options : mapping, optional
+ Keyword arguments for the :class:`~cunumpy.kernels.PyccelKernel`
+ wrapping a plain callable, e.g. ``{"outputs": (2,)}``; they matter
+ for the fallback. Not allowed with a `host_kernel` that already is a
+ ``PyccelKernel``.
+ dispatch : {"backend", "arrays"}, optional
+ ``"backend"`` (default): the CUDA kernel on the CuPy backend, the
+ host kernel on the NumPy backend. ``"arrays"``: the CUDA kernel if
+ any top-level argument is on the GPU (a CuPy array or an object with
+ ``__cuda_args__``), else the host kernel called directly, whatever the
+ backend; for codes that hand host arrays to kernels while CuPy is
+ active (diagnostics, MPI staging).
+ test_args : str, optional
+ Dotted name of the test-arguments module of the kernel (see
+ :func:`~cunumpy.kernel_testing.check_parity`).
+
+ Raises
+ ------
+ ValueError
+ If `dispatch` or `missing_cuda` is unknown, or `host_options` is
+ given with a ``PyccelKernel``.
+ TypeError
+ If `cuda_kernel` is not a :class:`~cunumpy.kernels.CudaKernel`.
Notes
-----
Both kernels take the same arguments, except that the CUDA kernel gets the
- launch shape (``n_threads`` or ``grid``) and argument objects in their CUDA
- form (see :class:`~cunumpy.arguments.CudaArguments` and :class:`~cunumpy.arguments.CudaStruct`).
- Argument objects are passed as they are: the caller passes the host
- argument object (e.g. a pyccel class) on the host and the CUDA one (a
+ launch size and argument objects in their CUDA form: the caller passes
+ the host object (e.g. a Pyccel class) on the host and the CUDA one (e.g. a
:class:`~cunumpy.arguments.CudaStructArguments`) on the device; see
:doc:`/kernels/arguments`.
+
+ Examples
+ --------
+ >>> def scale(x, out):
+ ... out[:] = 2 * x
+ >>> kernel = xp.kernels.Kernel(scale)
+ >>> out = np.empty(3)
+ >>> kernel(np.arange(3.0), out)
+ >>> out
+ array([0., 2., 4.])
"""
def __init__(
@@ -253,72 +257,78 @@ def from_folder(
) = None,
**cuda_options: Any,
) -> Kernel:
- """The kernel of one kernel folder, the unit of :meth:`KernelCatalog.from_package`.
-
- `package` is the dotted name of the folder ````. Every file of the
- folder that defines a version of the kernel (a function ````, or
- the ``__global__`` function ````) is one implementation:
-
- * ``.py``: ``"pyccel"``, compiled with `compile_host`
- (as it is without one), and ``"python"``, the same file uncompiled;
- * ``_numba.py``: ``"numba"`` (the function decorated there, e.g.
- with ``numba.njit``; unavailable if the import fails);
- * ``_numpy.py``: ``"numpy"``;
- * ````: the CUDA kernel.
-
- The host implementations form a :class:`~cunumpy.kernels.HostImplementations`:
- a call runs the one set with :func:`~cunumpy.kernels.set_host_kernel_implementation`,
- or by default the first available of pyccel, numba and NumPy. The
- folder's own ``__init__.py`` can declare its kernel with this method, so
- that the kernel is imported from where it is written::
-
- # my_sim/kernels/push/__init__.py
- kernel = xp.kernels.Kernel.from_folder(
- __name__, host_suffix="_pyccel", compile_host=compile, dispatch="arrays"
- )
+ """Build the kernel of one kernel folder.
- # anywhere in the code
- from my_sim.kernels.push import kernel as push
-
- Nothing is compiled here: the CUDA source is parsed, the host module
- imported.
+ Each version of ```` in the folder is an implementation:
+ ``.py`` gives ``"pyccel"`` (compiled with
+ `compile_host` on first use) and ``"python"`` (uncompiled),
+ ``_numba.py`` gives ``"numba"``, ``_numpy.py`` gives
+ ``"numpy"`` and ```` the CUDA kernel. The host ones
+ form a :class:`~cunumpy.kernels.HostImplementations`. Nothing is
+ compiled here: the host module is imported, the CUDA source parsed.
Parameters
----------
package : str
Dotted name of the kernel folder (``__name__`` in its ``__init__.py``).
- host_suffix, cuda_suffix, test_args_suffix, check_name_length
- As for :meth:`KernelCatalog.from_package`.
- missing_cuda, host_options, dispatch
- Passed on to :class:`Kernel`.
- outputs : Sequence[int | str] | "annotations" | None
- The arguments the host kernel writes to (names or indices, see
- :class:`~cunumpy.kernels.PyccelKernel`), so that only those arrays are
- copied back to the device on the fallback path. ``"annotations"``
- reads them from the annotations of the host function: every
- parameter that is not ``Final`` or a scalar is an output (see
- :func:`~cunumpy.kernels.outputs_from_annotations`). A value in
- `host_options` takes precedence.
- include_dirs : Sequence[str | Path] | None
- Include directories of the CUDA kernel, in addition to the folder
- itself; by default the source root of the top-level package (the
- directory containing it).
- compile_host : Callable | None
- Returns the compiled form of the ```` module
- (cunumpy does not compile anything itself, e.g. a wrapper around
- ``pyccel.epyccel`` with a cache), called on first use of the pyccel
- implementation; if it raises, that implementation is unavailable.
- extra_implementations : Mapping | None
- Loaders of host implementations that are not files of the folder,
- by implementation name, e.g. ``{"numpy": lambda: push_numpy}``.
+ host_suffix : str, optional
+ Module name suffix of the host kernel.
+ cuda_suffix : str, optional
+ File name suffix of the CUDA kernel.
+ test_args_suffix : str or None, optional
+ Module name suffix of the test arguments; None disables them.
+ check_name_length : bool, optional
+ Warn if the module name is too long for Pyccel's Fortran backend.
+ missing_cuda : {"raise", "fallback"}, optional
+ Passed on to :class:`~cunumpy.kernels.Kernel`.
+ host_options : mapping, optional
+ Passed on to :class:`~cunumpy.kernels.Kernel`.
+ outputs : sequence of int or str, or "annotations", optional
+ The arguments the host kernel writes to, copied back to the device
+ on the fallback path. ``"annotations"`` takes every parameter that
+ is not ``Final`` or a scalar (see
+ :func:`~cunumpy.kernels.outputs_from_annotations`). An
+ ``"outputs"`` entry of `host_options` takes precedence.
+ include_dirs : sequence of str or Path, optional
+ Include directories of the CUDA kernel besides the folder; by
+ default the directory containing the top-level package.
+ dispatch : {"backend", "arrays"}, optional
+ Passed on to :class:`~cunumpy.kernels.Kernel`.
+ compile_host : callable, optional
+ Returns the compiled form of the host module (e.g. a wrapper of
+ ``pyccel.epyccel`` with a cache); cunumpy compiles nothing itself.
+ If it raises, the ``"pyccel"`` implementation is unavailable.
+ extra_implementations : mapping, optional
+ Loaders of further host implementations by name, e.g.
+ ``{"numpy": lambda: push_numpy}``.
**cuda_options
- Passed on to :meth:`CudaKernel.from_file`, e.g. ``block_size`` or
- ``n_threads_from``.
+ Passed on to :meth:`CudaKernel.from_file
+ `, e.g. ``block_size``.
+
+ Returns
+ -------
+ Kernel
+ The kernel of the folder.
Raises
------
+ ModuleNotFoundError
+ If `package` is not a package.
FileNotFoundError
If the folder has no host kernel module.
+
+ See Also
+ --------
+ KernelCatalog.from_package : All kernel folders of a package.
+
+ Examples
+ --------
+ >>> # my_sim/kernels/push/__init__.py
+ >>> kernel = xp.kernels.Kernel.from_folder( # doctest: +SKIP
+ ... __name__, host_suffix="_pyccel", dispatch="arrays"
+ ... )
+ >>> kernel.implementations # doctest: +SKIP
+ ('pyccel', 'python', 'cuda')
"""
spec = importlib.util.find_spec(package)
if spec is None or not spec.submodule_search_locations:
@@ -425,10 +435,7 @@ def has_cuda(self) -> bool:
@property
def missing_cuda(self) -> str:
- """What happens on the CuPy backend without CUDA kernel.
-
- ``"raise"`` or ``"fallback"``, see :class:`Kernel`.
- """
+ """Behavior on the device without CUDA kernel: ``"raise"`` or ``"fallback"``."""
return self._missing_cuda
@property
@@ -443,11 +450,7 @@ def dispatch(self) -> str:
@property
def test_args_module(self) -> str | None:
- """Dotted name of the module with the test arguments of this kernel, or None.
-
- Set by :meth:`KernelCatalog.from_package` for a kernel folder that
- contains ``_test_args.py``; see :func:`cunumpy.kernel_testing.check_parity`.
- """
+ """Dotted name of the test-arguments module of the kernel, or None."""
return self._test_args_module
@property
@@ -458,13 +461,16 @@ def test_args(self) -> ModuleType | None:
return self._test_args
def host_parameters(self) -> list[str] | None:
- """The parameter names of the host kernel, or None if they are unknown.
+ """Return the parameter names of the host kernel.
- Read from the Python function (for a :class:`~cunumpy.kernels.CompiledHostKernel`,
- its uncompiled Python version), or, for a pyccel-compiled function
- without a Python signature, from the ``__pyccel__/.pyi`` stub
- pyccel writes next to the extension module. None if neither is
- available.
+ Read from the Python function (the uncompiled version of a
+ :class:`~cunumpy.kernels.CompiledHostKernel`), or else from the
+ ``__pyccel__/.pyi`` stub Pyccel writes next to a compiled module.
+
+ Returns
+ -------
+ list of str or None
+ The names, or None if neither source is available.
"""
function = self._host_kernel.kernel
if isinstance(function, (CompiledHostKernel, HostImplementations)):
@@ -477,13 +483,12 @@ def host_parameters(self) -> list[str] | None:
def check_signature(self) -> None:
"""Check that all implementations of the kernel take the same parameters.
- Compares the parameter names of the host function with those of the
- parsed ``__global__`` signature, with those of the other host
- implementations of a :class:`~cunumpy.kernels.HostImplementations` (numba,
- NumPy; nothing is compiled, an implementation that fails to import is
- skipped) and with the fallback of a :class:`~cunumpy.kernels.CompiledHostKernel`.
- A side is skipped without a CUDA kernel, without a parsed CUDA signature
- (``check_signature=False``), or without a Python signature.
+ Compares the host parameter names with the parsed ``__global__``
+ signature, the numba and NumPy implementations of a
+ :class:`~cunumpy.kernels.HostImplementations` (nothing is compiled;
+ one that fails to import is skipped) and the fallback of a
+ :class:`~cunumpy.kernels.CompiledHostKernel`. A side without a
+ signature is skipped.
Raises
------
@@ -520,11 +525,11 @@ def check_signature(self) -> None:
)
@property
- def implementations(self) -> tuple[str, ...]:
- """Names of the implementations, e.g. ``("pyccel", "numpy", "python", "cuda")``.
+ def implementations(self) -> tuple[str, ...]: # numpydoc ignore=RT01
+ """The names of the implementations, e.g. ``("pyccel", "python", "cuda")``.
- A host kernel that is not a :class:`~cunumpy.kernels.HostImplementations` counts
- as ``"host"``.
+ A host kernel that is not a :class:`~cunumpy.kernels.HostImplementations`
+ counts as ``"host"``.
"""
function = self._host_kernel.kernel
host = (
@@ -533,16 +538,33 @@ def implementations(self) -> tuple[str, ...]:
return (*host, "cuda") if self._cuda_kernel is not None else host
def selected(self, device: bool = False) -> str:
- """The implementation a call with host (or `device`) arguments runs now.
-
- For host arguments: the setting of
- :func:`~cunumpy.kernels.set_host_kernel_implementation` or the default (loads it), or
- ``"host"`` for a host kernel that is not a
- :class:`~cunumpy.kernels.HostImplementations`. For device arguments ``"cuda"``,
- or ``"host"`` if there is no CUDA kernel and ``missing_cuda="fallback"``.
- Explicit CUDA selection rejects missing CUDA implementations with
- ``LookupError``, including when host fallback is configured.
- Useful to check that a run does not use a slow path.
+ """Return the implementation a call runs now, e.g. to rule out a slow path.
+
+ Parameters
+ ----------
+ device : bool, optional
+ Ask for device arguments instead of host arguments.
+
+ Returns
+ -------
+ str
+ For host arguments the implementation chosen by
+ :func:`~cunumpy.kernels.set_host_kernel_implementation` or the
+ default (loaded now), or ``"host"`` for a plain host kernel. For
+ device arguments ``"cuda"``, or ``"host"`` for the fallback.
+
+ Raises
+ ------
+ LookupError
+ With `device`, if CUDA is explicitly selected but missing.
+ NotImplementedError
+ With `device`, if there is no CUDA kernel and
+ ``missing_cuda="raise"``.
+
+ Examples
+ --------
+ >>> xp.kernels.Kernel(lambda x: None, name="noop").selected()
+ 'host'
"""
if device:
return "cuda" if self._device_kernel() is self._cuda_kernel else "host"
@@ -552,10 +574,16 @@ def selected(self, device: bool = False) -> str:
)
def get_kernel(self) -> PyccelKernel | CudaKernel:
- """The kernel for the active backend.
+ """Return the kernel for the active backend.
Call this once at setup to fail early if a CUDA kernel is missing.
+ Returns
+ -------
+ PyccelKernel or CudaKernel
+ The host kernel on the NumPy backend; on the CuPy backend the CUDA
+ kernel, or the host kernel with ``missing_cuda="fallback"``.
+
Raises
------
LookupError
@@ -616,18 +644,27 @@ def __call__(
shared_mem: int = 0,
stream: Any = None,
) -> Any:
- """Call the kernel for the active backend.
+ """Call the host or CUDA kernel, as chosen by `dispatch`.
Parameters
----------
*args
- Kernel arguments, as the selected kernel takes them.
+ Kernel arguments, passed as they are.
n_threads, grid, block, shared_mem, stream
Launch configuration of the CUDA kernel, see
:meth:`CudaKernel.__call__ `;
- Thread counts are inferred from array shapes by default;
- `n_threads` or `grid` overrides that choice.
- Ignored by the host kernel.
+ by default inferred from the arguments. Ignored by the host kernel.
+
+ Returns
+ -------
+ object
+ What the selected kernel returns.
+
+ Raises
+ ------
+ ValueError
+ If the CUDA kernel runs without `n_threads`, `grid` or
+ ``n_threads_from``.
"""
if self._dispatch == "arrays":
on_device = any(_on_device(arg) for arg in args)
@@ -667,15 +704,22 @@ def __call__(
class KernelCatalog(Mapping):
- """Kernels by name.
+ """A read-only mapping from names to :class:`~cunumpy.kernels.Kernel` objects.
- A read-only mapping from names to :class:`Kernel` objects, usually built with
- :meth:`from_package`; :meth:`register` adds kernels one by one.
+ Usually built with :meth:`from_package`; :meth:`register` adds kernels one
+ by one. ``str(catalog)`` is :meth:`summary`. See :doc:`/kernels/dispatch`.
Parameters
----------
- kernels : Mapping[str, Kernel] | None
+ kernels : mapping of str to Kernel, optional
Initial kernels.
+
+ Examples
+ --------
+ >>> kernel = xp.kernels.Kernel(lambda x: None, name="push")
+ >>> catalog = xp.kernels.KernelCatalog({"push": kernel})
+ >>> print(catalog)
+ CUDA kernels: 0 of 1 (missing: push)
"""
def __init__(self, kernels: Mapping[str, Kernel] | None = None) -> None:
@@ -714,67 +758,68 @@ def from_package(
) -> KernelCatalog:
"""Collect the kernels of a package with one folder per kernel.
- For every subfolder ```` of the package that contains the module
- ``.py``, the function ```` of that module is the
- host kernel, and ```` in the same folder, if present,
- is the CUDA kernel (with a ``__global__`` function ````). Other
- ``__global__`` functions in that file are ignored by the catalog; load
- them with :meth:`CudaKernel.all_from_file
- `.
+ Every subfolder ```` containing ``.py`` gives
+ a kernel (see :meth:`Kernel.from_folder
+ `): the function ```` of that
+ module is the host kernel, the ``__global__`` function ```` of
+ ````, if present, the CUDA kernel. Other
+ ``__global__`` functions are ignored; load them with
+ :meth:`CudaKernel.all_from_file `.
Parameters
----------
package : str
Full name of the package, e.g. ``__name__`` in its ``__init__.py``.
- host_suffix : str
+ host_suffix : str, optional
Module name suffix of the host kernels.
- cuda_suffix : str
+ cuda_suffix : str, optional
File name suffix of the CUDA kernels.
- test_args_suffix : str | None
- Module name suffix of the test arguments: ``.py``
- in the kernel's folder, if present, is recorded as
- :attr:`Kernel.test_args_module` (imported only when a test asks for
- it, see :func:`cunumpy.kernel_testing.check_parity`). None disables this.
- check_name_length : bool
- Warn about a kernel whose module name is too long for the Fortran
- backend of pyccel: the wrapper module ``bind_c_``
- must fit :data:`FORTRAN_NAME_LIMIT` (63) characters, so with the
- default suffix a kernel name has at most 48 characters.
- missing_cuda : {"raise", "fallback"}
- Passed on to every :class:`Kernel`.
- host_options : Mapping | Callable[[str], Mapping] | None
- Keyword arguments for the :class:`~cunumpy.kernels.PyccelKernel` wrapping each
- host kernel (see :class:`Kernel`): the same for all kernels, or a
- function of the kernel name, e.g. to declare per-kernel ``outputs``.
- outputs : Sequence[int | str] | "annotations" | Callable | None
- The arguments every host kernel writes to, or ``"annotations"`` to
- read them from the annotations of each host function, or a function
- of the kernel name returning either (see :meth:`Kernel.from_folder`).
- include_dirs : Sequence[str | Path] | None
- Include directories of the CUDA kernels, in addition to each
- kernel's own folder. By default the source root of the top-level
- package (the directory containing it), so that a kernel in
- ``my_pkg.kernels`` can ``#include "my_pkg/common.cuh"``.
- dispatch : {"backend", "arrays"}
- Passed on to every :class:`Kernel`: choose the CUDA kernel by the
- active backend or by where the arguments live.
- compile_host : Callable | None
- Compiles a host kernel module, e.g.
- a function wrapping ``pyccel.epyccel`` (cunumpy does not compile
- anything itself). Each host kernel then is a
- :class:`~cunumpy.kernels.CompiledHostKernel`: compiled on its first
- call, falling back to
- `host_fallback`, or to the uncompiled Python function with a
- warning, if compilation fails. By default the Python function is
- called as it is.
- host_fallback : Mapping | Callable | None
- The fallback per kernel name, for a failed compilation (e.g. a
- vectorized NumPy implementation with the same arguments): a mapping
- from names to callables, or a function of the name returning a
- callable or None.
+ test_args_suffix : str or None, optional
+ Module name suffix of the test arguments: a module
+ ``.py`` in a kernel folder is recorded as
+ :attr:`Kernel.test_args_module
+ `. None disables this.
+ check_name_length : bool, optional
+ Warn about a kernel whose Pyccel Fortran wrapper module
+ ``bind_c_`` exceeds 63 characters (kernel names
+ of at most 48 characters with the default suffix).
+ missing_cuda : {"raise", "fallback"}, optional
+ Passed on to every :class:`~cunumpy.kernels.Kernel`.
+ host_options : mapping or callable, optional
+ :class:`~cunumpy.kernels.PyccelKernel` options of the host kernels,
+ the same for all or a function of the kernel name.
+ outputs : sequence, "annotations" or callable, optional
+ The arguments every host kernel writes to, or a function of the
+ kernel name returning them (see :meth:`Kernel.from_folder
+ `).
+ include_dirs : sequence of str or Path, optional
+ Include directories of the CUDA kernels besides each kernel's
+ folder; by default the directory containing the top-level package,
+ so ``#include "my_pkg/common.cuh"`` works.
+ dispatch : {"backend", "arrays"}, optional
+ Passed on to every :class:`~cunumpy.kernels.Kernel`.
+ compile_host : callable, optional
+ Compiles a host kernel module (e.g. a wrapper of ``pyccel.epyccel``);
+ cunumpy compiles nothing itself. Each host kernel then is compiled
+ on first use and falls back to `host_fallback`, or to the Python
+ function with a warning, if compilation fails.
+ host_fallback : mapping or callable, optional
+ The fallback for a failed compilation by kernel name (e.g. a
+ vectorized NumPy version): a mapping or a function of the name.
**cuda_options
- Passed on to :meth:`CudaKernel.from_file`, e.g. ``block_size`` or
- ``structs``.
+ Passed on to :meth:`CudaKernel.from_file
+ `, e.g. ``block_size``.
+
+ Returns
+ -------
+ KernelCatalog
+ The kernels, by folder name.
+
+ Examples
+ --------
+ >>> # my_kernels/__init__.py
+ >>> catalog = xp.kernels.KernelCatalog.from_package(__name__) # doctest: +SKIP
+ >>> catalog["push"](x, out, n_threads=len(x)) # doctest: +SKIP
"""
root = Path(importlib.import_module(package).__file__).parent
if include_dirs is None:
@@ -804,7 +849,25 @@ def from_package(
return cls(kernels)
def register(self, kernel: Kernel, name: str | None = None) -> Kernel:
- """Add a kernel under `name` (by default its own name) and return it."""
+ """Add a kernel to the catalog.
+
+ Parameters
+ ----------
+ kernel : Kernel
+ The kernel.
+ name : str, optional
+ Its name in the catalog; by default the kernel's own name.
+
+ Returns
+ -------
+ Kernel
+ `kernel` itself.
+
+ Raises
+ ------
+ KeyError
+ If a kernel of that name is already registered.
+ """
if not isinstance(kernel, Kernel):
raise TypeError(f"expected a Kernel, got {type(kernel).__name__}")
name = kernel.name if name is None else name
@@ -831,12 +894,12 @@ def __str__(self) -> str:
return self.summary()
def summary(self, max_missing: int = 10) -> str:
- """One line on the porting status, e.g. for a ``--status`` command.
+ """Return one line on the porting status, e.g. for a ``--status`` command.
Parameters
----------
- max_missing : int
- How many of the kernels without CUDA kernel to name; the rest is
+ max_missing : int, optional
+ How many kernels without CUDA kernel to name (default 10); the rest is
shortened to ``...``.
Returns
@@ -856,39 +919,43 @@ def summary(self, max_missing: int = 10) -> str:
@property
def with_cuda(self) -> list[str]:
- """Names of the kernels with a CUDA kernel."""
+ """The names of the kernels with a CUDA kernel."""
return [name for name, kernel in self._kernels.items() if kernel.has_cuda]
@property
def without_cuda(self) -> list[str]:
- """Names of the kernels without a CUDA kernel, i.e. still to port."""
+ """The names of the kernels without a CUDA kernel, i.e. still to port."""
return [name for name, kernel in self._kernels.items() if not kernel.has_cuda]
def parity_cases(self) -> list[tuple[str, Kernel]]:
- """The ``(name, kernel)`` pairs of the kernels that have a CUDA kernel.
+ """Return the ``(name, kernel)`` pairs of the kernels with a CUDA kernel.
For a parametrised parity test of the whole catalog with
- :func:`cunumpy.kernel_testing.assert_kernels_agree`::
+ :func:`~cunumpy.kernel_testing.assert_kernels_agree`::
@pytest.mark.parametrize("name, kernel", catalog.parity_cases())
def test_parity(name, kernel):
assert_kernels_agree(kernel, make_args[name], n_threads=1000)
+
+ Returns
+ -------
+ list of tuple of (str, Kernel)
+ The pairs, in catalog order.
"""
return [
(name, kernel) for name, kernel in self._kernels.items() if kernel.has_cuda
]
def check_signatures(self) -> None:
- """Check every kernel's host and CUDA parameters (:meth:`Kernel.check_signature`).
+ """Check the parameters of every kernel, e.g. in a unit test.
- A cheap test for a ported package: call it once, e.g. in a unit test,
- to catch a CUDA kernel whose parameters are missing, extra or in
- another order than those of its host kernel.
+ Runs :meth:`Kernel.check_signature
+ ` on every kernel.
Raises
------
ValueError
- Listing every kernel whose two signatures differ.
+ Listing every kernel whose implementations take different parameters.
"""
problems = []
for kernel in self._kernels.values():
@@ -905,21 +972,20 @@ def check_signatures(self) -> None:
def compile_all(self, jobs: int | None = 1) -> list[str]:
"""Compile all CUDA kernels now, e.g. at setup instead of in the first step.
- Compilation is cached on disk by CuPy, so after the first run this mostly
- loads the compiled kernels.
+ CuPy caches compilations on disk, so later runs mostly load them. All
+ kernels are compiled even if one fails; the first error is raised
+ afterwards.
Parameters
----------
- jobs : int | None
- Number of kernels compiled at a time. With ``jobs > 1`` the kernels
- are compiled in threads (NVRTC compilation releases the GIL, and
- all threads use the current device); None uses the number of CPUs.
- All kernels are compiled even if one fails; the first error is
- raised afterwards.
+ jobs : int or None, optional
+ Number of kernels compiled at a time, in threads (NVRTC releases
+ the GIL; all threads use the current device); None uses the number
+ of CPUs.
Returns
-------
- list[str]
+ list of str
Names of the compiled kernels.
Raises
diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py
index 54ac02c..f2c0c02 100644
--- a/src/cunumpy/_emulation.py
+++ b/src/cunumpy/_emulation.py
@@ -1,66 +1,23 @@
"""Run a CUDA kernel on the CPU, one thread after another, for tests without a GPU.
-A ported kernel is usually checked against its host version on a GPU
-(:func:`cunumpy.kernel_testing.assert_kernels_agree`). Without one, CI cannot run that
-check, and the kernel's index and weight arithmetic go untested.
-:func:`emulate_cuda_kernel` closes that gap: it compiles the kernel source as
-C++ with the CUDA built-ins replaced by plain C++ (``threadIdx``, ``blockIdx``,
-``blockDim``, ``gridDim``, ``atomicAdd``, ...) into a shared library, and calls
-the kernel once per thread, serially, on the NumPy arguments, so the call looks
-like a launch::
-
- from cunumpy.kernel_testing import emulate_cuda_kernel
-
- y = np.zeros(1000)
- emulate_cuda_kernel(axpy, 2.0, x, y, 1000, n_threads=1000)
- np.testing.assert_allclose(y, 2.0 * x)
-
-Arguments follow the kernel signature: NumPy arrays for pointer and array view
-parameters (``Array1D`` to ``Array16D``, any strides), Python or NumPy
-scalars for scalar parameters (cast and checked like in a launch). The kernel
-writes into the given arrays.
-
-Each kernel is compiled once: the library takes the arguments at run time
-(``cunumpy_launch(args, grid, block, shared_mem)``, loaded with ctypes), so
-launches with other values and array sizes reuse it. Libraries are cached in
-the process and on disk (:func:`emulation_cache_dir`), keyed on the generated
-source, the included headers, the compiler and its options.
-:func:`compile_for_emulation` builds one without a launch. The launch runs in
-the Python process: ``__trap()`` (a failed bounds check) is reported as a
-``RuntimeError``, but a segmentation fault ends the process.
-
-Block shared memory and ``__syncthreads`` are emulated: ``__shared__``
-variables (and ``extern __shared__`` arrays, sized by ``shared_mem``) are one
-copy per block, and in a kernel that calls ``__syncthreads`` the threads of a
-block run as coroutines (POSIX ``ucontext``, each with its own stack): every
-thread runs to its next barrier before any thread continues past it, as on a
-GPU. Per-block deposits, shared-memory reductions and tiled kernels therefore
-work.
-
-Struct parameters take, for each struct, a mapping from field name to value, a
-:class:`~cunumpy.arguments.CudaStructValue`, or any object with an attribute
-per field (e.g. a :class:`~cunumpy.arguments.CudaStructArguments` or a host
-argument class); the arrays of the fields are updated in place like the others.
-
-:func:`emulated_launches` makes every :class:`~cunumpy.kernels.CudaKernel` launch
-(and compilation) in a block run through the emulation, on the arrays of the
-fake CuPy (:mod:`cunumpy._fake_cupy`), so that code that launches kernels (a
-:class:`~cunumpy.kernels.Kernel` on the CuPy backend, a propagator) can be
-tested without a GPU.
-
-What it does not emulate: concurrency between the barriers (threads run one
-after another, so atomics are plain additions and races never show), warp
-intrinsics (``__shfl_*``, ``__syncwarp``, ``__ballot_sync``, ...; a kernel or an
-included header using them is refused), ````, and inline PTX:
-``asm(...)`` and ``asm volatile(...)`` compile, but trap when reached. Use a
-GPU for those.
-
-Floating point: like NVRTC (``--fmad=true`` by default) the C++ compiler may
-fuse ``a * b + c`` into one fused multiply-add, so results can differ from
-NumPy's in the last bit; compare with a tolerance of a few ulp, or pass
-``options=("-ffp-contract=off",)`` for NumPy's rounding.
-
-It needs a C++17 compiler: ``CXX`` from the environment, else ``c++``.
+:func:`~cunumpy.kernel_testing.emulate_cuda_kernel` compiles the kernel source
+as C++17 (``CXX``, else ``c++``) with the CUDA built-ins (``threadIdx``,
+``atomicAdd``, ``__syncthreads``, ``__shared__``, ...) replaced by plain C++,
+into a shared library loaded with ctypes, and calls the kernel once per thread
+on NumPy arrays, so that CI without a GPU can check a port against its host
+kernel (:func:`~cunumpy.kernel_testing.assert_kernels_agree` needs a GPU).
+:func:`~cunumpy.kernel_testing.emulated_launches` routes every
+:class:`~cunumpy.kernels.CudaKernel` launch through it, on the arrays of the
+fake CuPy (:mod:`cunumpy._fake_cupy`).
+
+Libraries are cached in the process and on disk
+(:func:`~cunumpy.kernel_testing.emulation_cache_dir`). Threads of a block that
+calls ``__syncthreads`` run as coroutines (POSIX ``ucontext``), so barriers
+and block shared memory work. Not emulated: concurrency between barriers
+(atomics are plain operations, races never show), warp intrinsics (refused),
+```` and inline PTX (``asm`` traps when reached). The
+compiler may fuse multiply-adds like NVRTC: compare with a few ulp, or pass
+``options=("-ffp-contract=off",)``. See :doc:`/kernels/testing`.
"""
from __future__ import annotations
@@ -272,7 +229,19 @@
def emulation_compiler() -> str | None:
- """The C++ compiler :func:`emulate_cuda_kernel` uses, or None if there is none."""
+ """Return the C++ compiler the emulation uses: ``CXX``, else ``c++``.
+
+ Returns
+ -------
+ str or None
+ Path of the compiler, or None if there is none.
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import emulation_compiler
+ >>> emulation_compiler() # doctest: +SKIP
+ '/usr/bin/c++'
+ """
compiler = os.environ.get("CXX")
if compiler and shutil.which(compiler):
return shutil.which(compiler)
@@ -280,14 +249,14 @@ def emulation_compiler() -> str | None:
def _scalar_value(param: CudaParameter, index: int, value: Any) -> np.ndarray:
- """`value` as a 0-d array of the C type, checked and cast like a scalar argument."""
+ """Return `value` as a 0-d array of the C type, checked like a scalar argument."""
if np.dtype(param.dtype).kind not in "biuf":
raise NotImplementedError(f"emulation does not support {param.ctype} scalars")
return np.asarray(_scalar_checker(param, index)(value), dtype=param.dtype)
def _code(kernel: CudaKernel) -> str:
- """The kernel source and the headers it includes, without comments."""
+ """Return the kernel source and its included headers, without comments."""
texts = [kernel.source]
for header in kernel.included_headers:
try:
@@ -313,7 +282,7 @@ def _check_supported(kernel: CudaKernel) -> None:
def _host(value: Any) -> Any:
- """The NumPy buffer of a fake CuPy array; anything else as it is."""
+ """Return the NumPy buffer of a fake CuPy array, anything else unchanged."""
if _fake_cupy.is_active():
try:
return _fake_cupy.host_buffer(value)
@@ -323,7 +292,7 @@ def _host(value: Any) -> Any:
def _field_value(struct_name: str, value: Any, field: CudaParameter) -> Any:
- """The value of the struct field `field` in the struct argument `value`."""
+ """Return the value of struct field `field` in the struct argument `value`."""
try:
if isinstance(value, (CudaStructValue, Mapping)):
return value[field.name]
@@ -348,13 +317,7 @@ def _declaration(param: CudaParameter, name: str) -> str:
def _struct_wrapper(kernel: CudaKernel) -> CudaKernel:
- """A kernel that takes every struct field of `kernel` as a parameter.
-
- The wrapper kernel (the source of `kernel` plus a ``__global__`` function with
- one parameter per field) rebuilds the structs and calls `kernel`, so that the
- emulation only has to deal with arrays and scalars. Its arguments are
- :func:`_struct_fields`. Made once per kernel.
- """
+ """Return a cached wrapper kernel taking each struct field of `kernel` as a parameter."""
wrapper = _WRAPPERS.get(kernel)
if wrapper is not None:
return wrapper
@@ -389,7 +352,7 @@ def _struct_wrapper(kernel: CudaKernel) -> CudaKernel:
def _struct_fields(kernel: CudaKernel, args: Sequence[Any]) -> list[Any]:
- """The arguments of :func:`_struct_wrapper`: each struct replaced by its fields."""
+ """Return the arguments for :func:`_struct_wrapper`, structs replaced by their fields."""
flat = []
for p, value in zip(kernel.signature, args):
if p.struct is None:
@@ -400,14 +363,7 @@ def _struct_fields(kernel: CudaKernel, args: Sequence[Any]) -> list[Any]:
def _emulation_source(kernel: CudaKernel) -> tuple[str, bool]:
- """The CUDA stubs and the kernel source as C++, without the launch function.
-
- Returns
- -------
- tuple[str, bool]
- The source, and whether the threads of a block run as coroutines (the
- kernel calls ``__syncthreads``).
- """
+ """Return the C++ source (stubs and kernel) and whether threads run as coroutines."""
kernel_source = kernel.source.replace('extern "C"', "")
kernel_source = _EXTERN_SHARED.sub(
r"\1* \2 = (\1*)cunumpy_dynamic_shared;",
@@ -429,7 +385,7 @@ def _emulation_source(kernel: CudaKernel) -> tuple[str, bool]:
def _element_type(param: CudaParameter) -> str:
- """The C type of the elements of a pointer or array view parameter."""
+ """Return the C element type of a pointer or array view parameter."""
if param.dtype is None:
return "unsigned char"
if param.view_ndim is not None:
@@ -438,7 +394,7 @@ def _element_type(param: CudaParameter) -> str:
def _library_source(kernel: CudaKernel) -> str:
- """The C++ source of the emulation library of `kernel` (see :data:`_LAUNCH`)."""
+ """Return the C++ source of the emulation library of `kernel`."""
source, coroutines = _emulation_source(kernel)
globals_, inits, call_args = [], [], []
for i, param in enumerate(kernel.signature):
@@ -477,7 +433,7 @@ def _compile_command(
compiler: str,
options: Sequence[str],
) -> list[str]:
- """The compiler, include directories and options for the emulation of `kernel`."""
+ """Return the compile command (without the source) for the emulation of `kernel`."""
include_dirs = [cuda_include_dir(), *map(str, kernel.include_dirs)]
if kernel.source_dir is not None:
include_dirs.insert(0, str(kernel.source_dir))
@@ -493,11 +449,22 @@ def _compile_command(
def emulation_cache_dir() -> Path | None:
- """Where the emulation libraries are cached on disk, or None for no disk cache.
+ """Return the directory of the on-disk cache of emulation libraries.
- ``CUNUMPY_EMULATION_CACHE`` sets the directory (``0`` or empty: no disk
+ ``CUNUMPY_EMULATION_CACHE`` sets it (``0``, ``off`` or empty: no disk
cache); the default is ``$XDG_CACHE_HOME/cunumpy/emulation``, else
``~/.cache/cunumpy/emulation``.
+
+ Returns
+ -------
+ Path or None
+ The directory, or None for no disk cache.
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import emulation_cache_dir
+ >>> emulation_cache_dir() # doctest: +SKIP
+ PosixPath('/home/me/.cache/cunumpy/emulation')
"""
value = os.environ.get("CUNUMPY_EMULATION_CACHE")
if value is not None:
@@ -548,7 +515,7 @@ def _build(command: list[str], source: str, out: Path, kernel: CudaKernel) -> No
def _library(kernel: CudaKernel, compiler: str, options: Sequence[str]) -> _Library:
- """The emulation library of `kernel`: from this process, the disk cache, or built."""
+ """Return the emulation library of `kernel`: from this process, the disk cache, or built."""
command = [*_compile_command(kernel, compiler, options), "-O1", "-shared", "-fPIC"]
source = _library_source(kernel)
digest = hashlib.sha256()
@@ -611,7 +578,7 @@ def _prepared(
*,
refresh: bool = False,
) -> _Library:
- """The library of `kernel`, looked up once per kernel object (or again with `refresh`)."""
+ """Return the library of `kernel`, looked up once per kernel object (or with `refresh`)."""
libraries = _PREPARED.setdefault(kernel, {})
key = (compiler, tuple(options))
if refresh or key not in libraries:
@@ -660,30 +627,42 @@ def compile_for_emulation(
compiler: str | None = None,
options: Sequence[str] = (),
) -> None:
- """Build the emulation library of `kernel` now, without a launch.
+ """Build (or fetch from the cache) the emulation library of `kernel` now.
- :func:`emulate_cuda_kernel` builds it on the first launch otherwise; this
- surfaces compile errors early, like :meth:`CudaKernel.compile
- ` does on a GPU (inside
- :func:`emulated_launches`, ``compile()`` calls this). The library is cached
- (see :func:`emulate_cuda_kernel`). A kernel without a parsed signature is
- only checked for syntax and type errors (``-fsyntax-only``).
+ Surfaces compile errors before the first launch, like
+ :meth:`CudaKernel.compile ` on a GPU
+ (inside :func:`~cunumpy.kernel_testing.emulated_launches`, ``compile()``
+ calls this). A kernel without a parsed signature is only checked for
+ syntax and type errors (``-fsyntax-only``).
Parameters
----------
kernel : CudaKernel
The kernel.
- compiler : str | None
- C++ compiler; by default :func:`emulation_compiler`.
- options : Sequence[str]
- Additional compiler options, as for :func:`emulate_cuda_kernel`.
+ compiler : str, optional
+ C++ compiler; by default
+ :func:`~cunumpy.kernel_testing.emulation_compiler`.
+ options : sequence of str, optional
+ Additional compiler options.
Raises
------
NotImplementedError
- If the kernel uses what the emulation cannot run (warp intrinsics).
+ If the kernel uses warp intrinsics, which the emulation cannot run.
RuntimeError
- If there is no C++ compiler, or the kernel does not compile.
+ If there is no C++ compiler, or the kernel does not compile (the
+ message holds the compiler output).
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import compile_for_emulation
+ >>> axpy = xp.kernels.CudaKernel(
+ ... 'extern "C" __global__ void axpy(double a, const double* x, double* y,'
+ ... " int n) { int i = blockIdx.x * blockDim.x + threadIdx.x;"
+ ... " if (i < n) y[i] += a * x[i]; }",
+ ... "axpy",
+ ... )
+ >>> compile_for_emulation(axpy)
"""
_check_supported(kernel)
compiler = _resolve_compiler(compiler)
@@ -699,11 +678,7 @@ def _pack(
params: Sequence[CudaParameter],
args: Sequence[Any],
) -> tuple[Any, list[Any], list[tuple[np.ndarray, np.ndarray]]]:
- """The ``void*`` arguments of ``cunumpy_launch``.
-
- Returns the argument array, the objects it points into (to keep alive), and
- the (array, copy) pairs whose copy must be written back after the launch.
- """
+ """Return the ``void*`` arguments of ``cunumpy_launch``, objects to keep alive, and copies to write back."""
argv = (ctypes.c_void_p * max(len(params), 1))()
keep: list[Any] = []
write_back: list[tuple[np.ndarray, np.ndarray]] = []
@@ -774,35 +749,37 @@ def emulate_cuda_kernel(
compiler: str | None = None,
options: Sequence[str] = (),
shared_mem: int = 0,
-) -> None:
+) -> None: # numpydoc ignore=RT01
"""Run `kernel` on the CPU, serially, as if it were launched with `args`.
- The kernel is compiled once (for each `compiler` and `options`) into a
- shared library that takes the arguments at run time, so later launches,
- with any argument values and array sizes, reuse it. Libraries are cached in
- the process and on disk (:func:`emulation_cache_dir`). The launch runs in
- this process, on the arrays themselves.
+ The kernel is compiled once per `compiler` and `options` into a library
+ that takes the arguments at run time, so later launches with any values
+ and sizes reuse it (cached in the process and in
+ :func:`~cunumpy.kernel_testing.emulation_cache_dir`). The launch runs in
+ this process and writes into the given arrays. A segmentation fault (e.g.
+ an unchecked out-of-bounds index) ends the process.
Parameters
----------
kernel : CudaKernel
- The kernel; its signature must be parsed (the default).
- *args
- The kernel arguments with NumPy arrays in place of CuPy arrays (arrays of
- the fake CuPy are used through their host buffer). Arrays are updated in
- place with what the kernel wrote. A struct parameter takes a mapping of
- field names to values, a :class:`~cunumpy.arguments.CudaStructValue`, or
- an object with an attribute per field.
- n_threads, grid, block
- Launch shape, as for :meth:`CudaKernel.__call__`.
- compiler : str | None
- C++ compiler; by default :func:`emulation_compiler`.
- options : Sequence[str]
- Additional compiler options, e.g. ``("-DMY_FLAG=1",)``. ``-D`` options
- of the kernel are passed on as well.
- shared_mem : int
+ The kernel, with a parsed signature (the default).
+ *args : array or scalar
+ The kernel arguments, with NumPy arrays (any strides; fake CuPy arrays
+ are used through their host buffer) in place of CuPy arrays. A struct
+ parameter takes a mapping of field names to values, a
+ :class:`~cunumpy.arguments.CudaStructValue`, or an object with an
+ attribute per field.
+ n_threads, grid, block : int or sequence of int, optional
+ Launch shape, as for :class:`~cunumpy.kernels.CudaKernel` calls.
+ compiler : str, optional
+ C++ compiler; by default
+ :func:`~cunumpy.kernel_testing.emulation_compiler`.
+ options : sequence of str, optional
+ Additional compiler options, e.g. ``("-ffp-contract=off",)``. The
+ kernel's own ``-D`` options apply too.
+ shared_mem : int, optional
Dynamic shared memory per block in bytes, for ``extern __shared__``
- arrays, as in a launch.
+ arrays. Default 0.
Raises
------
@@ -811,12 +788,24 @@ def emulate_cuda_kernel(
complex scalars.
TypeError
If an argument does not match its parameter (dtype, dimensions, a
- scalar that does not fit).
+ scalar that does not fit), or the signature is not parsed.
RuntimeError
If there is no C++ compiler, the kernel does not compile, or it calls
- ``__trap()`` (a failed bounds check, inline ``asm``). Other crashes
- (e.g. a segmentation fault from an unchecked out-of-bounds index) end
- the process, as the launch runs in it.
+ ``__trap()`` (a failed bounds check, inline ``asm``).
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import emulate_cuda_kernel
+ >>> axpy = xp.kernels.CudaKernel(
+ ... 'extern "C" __global__ void axpy(double a, const double* x, double* y,'
+ ... " int n) { int i = blockIdx.x * blockDim.x + threadIdx.x;"
+ ... " if (i < n) y[i] += a * x[i]; }",
+ ... "axpy",
+ ... )
+ >>> y = np.zeros(3)
+ >>> emulate_cuda_kernel(axpy, 2.0, np.arange(3.0), y, 3, n_threads=3)
+ >>> y
+ array([0., 2., 4.])
"""
if kernel.signature is None:
raise TypeError(
@@ -875,41 +864,42 @@ def emulated_launches(
) -> Generator[None, None, None]:
"""Run every :class:`~cunumpy.kernels.CudaKernel` launch in the block on the CPU.
- With the fake CuPy (:func:`cunumpy.kernel_testing.install_fake_cupy`)
- CUDA kernels cannot run. Inside this block a launch is emulated instead
- (:func:`emulate_cuda_kernel`), on the host buffers of the fake CuPy arrays
- it is given, so code that launches kernels runs on a machine without a GPU
- and its device arrays hold the results afterwards::
-
- with emulated_launches():
- propagator(dt) # CuPy backend: the CUDA kernels run on the CPU
-
- The launch arguments are those of :meth:`CudaKernel.__call__`
- (`stream` is ignored). Argument objects are flattened like in a launch;
- struct values and :class:`~cunumpy.arguments.CudaStructArguments` objects
- stay one argument and are read through their fields. Arrays must be arrays
- of the fake CuPy or NumPy arrays.
-
- No CUDA compilation happens in the block either: :meth:`CudaKernel.compile`
- (and so ``recompile()``, ``CudaKernelVariants.compile_all()`` and
- ``KernelCatalog.compile_all()``) builds the emulation library instead
- (:func:`compile_for_emulation`), so compile errors still show up before
- the first launch, and returns None instead of a ``cupy.RawKernel``. Kernels
- the emulation cannot run (warp intrinsics) are skipped by ``compile()``;
- launching them raises. This holds on a real GPU too: the block means "no
- CUDA here".
-
- Launches are serial, so use it on small problems. Each kernel is compiled
- once (see :func:`emulate_cuda_kernel`). The limits of
- :func:`emulate_cuda_kernel` apply.
+ Each launch goes through :func:`~cunumpy.kernel_testing.emulate_cuda_kernel`
+ on NumPy arrays or the host buffers of fake CuPy arrays, so code that
+ launches kernels (a :class:`~cunumpy.kernels.Kernel` on the CuPy backend)
+ runs without a GPU. Argument objects are flattened as in a launch;
+ `stream` is ignored. ``CudaKernel.compile()`` (and so the ``compile_all()``
+ methods) builds the emulation library instead and returns None, skipping
+ kernels with warp intrinsics. This holds on a GPU too: the block means
+ "no CUDA here". The methods are restored on exit. Not thread-safe: it
+ patches the class. Launches are serial, so keep problems small.
Parameters
----------
- compiler : str | None
- C++ compiler; by default :func:`emulation_compiler`.
- options : Sequence[str]
- Additional compiler options for every launch, e.g.
- ``("-ffp-contract=off",)``.
+ compiler : str, optional
+ C++ compiler; by default
+ :func:`~cunumpy.kernel_testing.emulation_compiler`.
+ options : sequence of str, optional
+ Additional compiler options for every launch.
+
+ Yields
+ ------
+ None
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import emulated_launches
+ >>> axpy = xp.kernels.CudaKernel(
+ ... 'extern "C" __global__ void axpy(double a, const double* x, double* y,'
+ ... " int n) { int i = blockIdx.x * blockDim.x + threadIdx.x;"
+ ... " if (i < n) y[i] += a * x[i]; }",
+ ... "axpy",
+ ... )
+ >>> y = np.zeros(3)
+ >>> with emulated_launches():
+ ... axpy(2.0, np.arange(3.0), y, 3, n_threads=3)
+ >>> y
+ array([0., 2., 4.])
"""
original_call = CudaKernel.__call__
original_compile = CudaKernel.compile
diff --git a/src/cunumpy/_fusion.py b/src/cunumpy/_fusion.py
index 68541b8..8ff8a67 100644
--- a/src/cunumpy/_fusion.py
+++ b/src/cunumpy/_fusion.py
@@ -1,26 +1,9 @@
-"""``xp.kernels.fuse``: elementwise functions as one GPU kernel, plain calls on the host.
-
-A chain of elementwise operations (a pressure from density and temperature,
-fluxes and limiters of a fluid update, a Maxwellian at many velocities) runs
-as one kernel per operation on the GPU, each writing a temporary array: it is
-limited by memory bandwidth. ``cupy.fuse`` compiles the whole chain into one
-kernel. :func:`fuse` applies it when the function is called with CuPy arrays
-and calls the function as it is otherwise, so the same code runs on both
-backends::
-
- @xp.kernels.fuse
- def pressure(rho, T, gamma):
- return (gamma - 1.0) * rho * T
-
- p = pressure(rho, T, 5.0 / 3.0)
-
-The function must be elementwise: arithmetic, comparisons, ufuncs such as
-``xp.exp``, ``xp.sqrt``, ``xp.where`` and reductions that ``cupy.fuse``
-supports (``xp.sum`` as the last operation). On the CuPy path the CuPy
-backend is active while the function is traced, so ``xp.*`` names resolve to
-CuPy's ufuncs. Functions that ``cupy.fuse`` cannot trace (Python control
-flow on array values, indexing, wrappers that are not ufuncs) raise when the
-fused function is first called with CuPy arrays; test the CuPy path.
+"""Elementwise functions as one GPU kernel, plain calls on the host.
+
+On the GPU a chain of elementwise operations runs one kernel per operation,
+each writing a temporary array. :func:`~cunumpy.kernels.fuse` compiles the
+chain into one kernel with ``cupy.fuse`` when called with CuPy arrays and
+calls the function as it is otherwise, so the same code runs on both backends.
"""
from __future__ import annotations
@@ -49,12 +32,7 @@ def _cupy_fuse(function: Callable[..., Any], kernel_name: str | None) -> Any:
def _typed_scalars(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
- """Give Python scalars the dtype they would take next to the arrays.
-
- ``cupy.fuse`` types a Python scalar on its own: ``gamma - 1.0`` with
- ``gamma=5/3`` runs in float16. Eagerly, NumPy 2 and CuPy promote it with
- the arrays (float64 arrays: float64), so cast it to that dtype first.
- """
+ """Cast Python scalars to the dtype they promote to with the array arguments."""
dtypes = [a.dtype for a in (*args, *kwargs.values()) if hasattr(a, "dtype")]
def typed(a: Any) -> Any:
@@ -75,24 +53,36 @@ def fuse(
) -> F | Callable[[F], F]:
"""Fuse an elementwise function into one kernel when called with CuPy arrays.
- Usable as ``@xp.kernels.fuse`` or ``@xp.kernels.fuse(kernel_name="pressure")``.
+ Usable as ``@xp.kernels.fuse`` or ``@xp.kernels.fuse(kernel_name=...)``.
+ The function must be elementwise in the sense of ``cupy.fuse``:
+ arithmetic, comparisons, ufuncs, ``xp.where`` and supported reductions as
+ the last operation; no Python control flow on array values, no indexing.
+ The CuPy backend is active while it is traced, so ``xp.exp`` etc. resolve
+ to CuPy ufuncs. A function ``cupy.fuse`` cannot trace raises at its first
+ call with CuPy arrays: test the CuPy path.
Parameters
----------
- function : callable
- The elementwise function.
- kernel_name : str | None
- Name of the generated kernel (shown by profilers); by default the
- function's name.
+ function : callable, optional
+ The elementwise function; omitted when used as ``fuse(kernel_name=...)``.
+ kernel_name : str, optional
+ Name of the kernel in profilers; by default the function's name.
Returns
-------
callable
- A function with the same signature. If any positional or keyword
- argument is a CuPy array, it calls ``cupy.fuse(function)`` (created on
- first use, with the CuPy backend active) with Python scalars cast to
- the dtype they promote to with the array arguments; otherwise it calls
- `function` itself.
+ A function with the same signature. If any argument is a CuPy array
+ it calls the fused kernel (created on first call and reused), with
+ Python scalars cast to the dtype they promote to with the arrays;
+ otherwise it calls `function` itself. Without `function`, a decorator.
+
+ Examples
+ --------
+ >>> @xp.kernels.fuse
+ ... def pressure(rho, T, gamma):
+ ... return (gamma - 1.0) * rho * T
+ >>> pressure(np.array([1.0, 2.0]), np.array([3.0, 3.0]), 2.0)
+ array([3., 6.])
"""
if function is None:
return lambda f: fuse(f, kernel_name=kernel_name)
diff --git a/src/cunumpy/_host.py b/src/cunumpy/_host.py
index d208273..8ba25b3 100644
--- a/src/cunumpy/_host.py
+++ b/src/cunumpy/_host.py
@@ -1,4 +1,4 @@
-"""Run host-only code (SciPy, file readers, external libraries) with arguments of any backend."""
+"""Run host-only code (SciPy, file readers, external libraries) with arguments of either backend."""
import functools
from collections.abc import Callable
@@ -31,17 +31,17 @@ def _to_device(arg: Any) -> Any:
def host_call(fun: Callable, *args: Any, **kwargs: Any) -> Any:
- """Call a host-only function with arguments of any backend.
+ """Call a host-only function with arguments of either backend.
Device (CuPy) arrays among the arguments are copied to the host, `fun` runs
on the NumPy backend, and its array results are copied back to the device,
- once per call. Without device arguments `fun` is called directly (on the
- NumPy backend this is a plain call), so the result lives where the
- arguments live. The copies are counted by `xp.profiling.count_transfers()`.
+ once per call (counted by :func:`~cunumpy.profiling.count_transfers`).
+ Without device arguments `fun` is called directly, so the result lives
+ where the arguments live. See :doc:`/guides/data-movement`.
Parameters
----------
- fun
+ fun : callable
The host-only function (a SciPy spline, an external code, ...).
*args, **kwargs
Arguments of `fun`: arrays (or tuples/lists of arrays) of either
@@ -49,12 +49,19 @@ def host_call(fun: Callable, *args: Any, **kwargs: Any) -> Any:
Returns
-------
- The result of `fun` (an array, a tuple/list of arrays or a scalar), with
- arrays on the device if any argument was on the device.
+ object
+ The result of `fun` (an array, a tuple/list of arrays or a scalar),
+ with its arrays on the device if any argument was on the device.
+
+ See Also
+ --------
+ evaluate_on_host : The same for a method.
Examples
--------
- >>> values = xp.host_call(spline, x) # x on the device -> values on the device
+ >>> xp.host_call(np.cumsum, xp.asarray([1, 2, 3]))
+ array([1, 3, 6])
+ >>> values = xp.host_call(spline, x_device) # doctest: +SKIP
"""
device = _on_device(args) or _on_device(list(kwargs.values()))
if xp.get_backend() == "numpy" and not device:
@@ -65,14 +72,30 @@ def host_call(fun: Callable, *args: Any, **kwargs: Any) -> Any:
def evaluate_on_host(method: Callable) -> Callable:
- """Decorator for methods that can only be evaluated on the host (see `host_call`).
+ """Make a host-only method callable with arguments of either backend.
+
+ Each call goes through :func:`host_call`: device arrays are copied to the
+ host, the method runs on the NumPy backend, and array results are copied
+ back if any argument was on the device.
+
+ Parameters
+ ----------
+ method : callable
+ The method, ``method(self, *args, **kwargs)``.
+
+ Returns
+ -------
+ callable
+ The wrapped method.
Examples
--------
- >>> class Equilibrium:
+ >>> class Profile:
... @xp.evaluate_on_host
- ... def pressure(self, x):
- ... return scipy_spline(x)
+ ... def cumulative(self, x):
+ ... return np.cumsum(x)
+ >>> Profile().cumulative(xp.asarray([1, 2, 3]))
+ array([1, 3, 6])
"""
@functools.wraps(method)
@@ -83,11 +106,30 @@ def wrapper(self, *args, **kwargs):
def setup_on_host(init: Callable) -> Callable:
- """Decorator for an ``__init__`` whose setup is host-only.
+ """Run an ``__init__`` on the NumPy backend.
- `init` runs on the NumPy backend, so the object holds only host data (NumPy
- arrays, SciPy splines, floats) on either backend. Host-only parts of its
- evaluation can then go through `host_call` or `evaluate_on_host`.
+ The object then holds only host data (NumPy arrays, SciPy splines, floats)
+ on either backend; host-only parts of its evaluation can go through
+ :func:`host_call` or :func:`evaluate_on_host`.
+
+ Parameters
+ ----------
+ init : callable
+ The ``__init__`` method.
+
+ Returns
+ -------
+ callable
+ The wrapped ``__init__``.
+
+ Examples
+ --------
+ >>> class Grid:
+ ... @xp.setup_on_host
+ ... def __init__(self, n):
+ ... self.x = xp.linspace(0.0, 1.0, n)
+ >>> xp.get_array_backend(Grid(3).x)
+ 'numpy'
"""
@functools.wraps(init)
diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py
index db2ed74..16de9c4 100644
--- a/src/cunumpy/_kernel.py
+++ b/src/cunumpy/_kernel.py
@@ -1,16 +1,15 @@
-"""Interface for calling Pyccel-compiled kernels with CuPy arrays.
-
-Kernels generated by `pyccel `_ are compiled
-C/Fortran routines that only understand NumPy (host) arrays. :class:`PyccelKernel`
-wraps such a kernel so that it can be called transparently with CuPy (device)
-arrays: the arguments are copied to the host before the call, in-place updates
-made by the kernel are copied back to the device afterwards, and any arrays
-returned by the kernel are moved back to the device.
-
-On the NumPy backend the wrapper is a no-op and the kernel is called directly.
-Ordinary Python callables are supported, including in Pyodide. This module
-neither imports Pyccel nor compiles kernels; compilation, if desired, is the
-caller's responsibility.
+"""Host kernels: calling them with CuPy arrays, and choosing their implementation.
+
+:class:`~cunumpy.kernels.PyccelKernel` wraps a host kernel (a
+`Pyccel `_-compiled function or any callable
+taking NumPy arrays) so that it can be called with CuPy arrays: inputs are
+copied to the host, written arrays copied back. On the NumPy backend the
+kernel is called directly. This module neither imports Pyccel nor compiles
+anything. It also holds the host/device implementation settings,
+:class:`~cunumpy.kernels.HostImplementations`,
+:class:`~cunumpy.kernels.CompiledHostKernel` and the argument helpers
+:func:`~cunumpy.kernels.as_kernel_array` and
+:func:`~cunumpy.kernels.kernel_output`. See :doc:`/kernels/pyccel-kernel`.
"""
from __future__ import annotations
@@ -45,56 +44,65 @@
class PyccelKernel:
- """Call a NumPy callable or Pyccel-compiled kernel with NumPy or CuPy arrays.
+ """Call a host kernel (Pyccel-compiled or any NumPy callable) with CuPy arrays.
+
+ When conversion is needed, CuPy arrays in the arguments (also inside
+ lists, tuples, dicts and objects from `object_modules`) are copied to the
+ host, the kernel runs, the arrays it may have written are copied back into
+ the CuPy arrays, and returned NumPy arrays are moved to the device. A
+ device array passed several times is copied once. Otherwise the kernel is
+ called directly. See :doc:`/kernels/pyccel-kernel`.
Parameters
----------
kernel : callable
- The pyccelized kernel (or any callable expecting NumPy arrays).
+ The host kernel, taking NumPy arrays.
use_cupy : bool, optional
- Force host/device conversion on (``True``) or off (``False``). By
- default (``None``) it is decided at call time: conversion happens when
- the active backend is CuPy or when a CuPy array is passed in.
+ Force conversion on or off. By default (None) it is decided per call:
+ on when the backend is CuPy or a CuPy array is passed.
object_modules : sequence of str, optional
- Module prefixes (e.g. ``("struphy.", "feectools.")``) whose instances
- should be traversed attribute-by-attribute when looking for arrays to
- convert. Objects from other modules are passed through untouched.
+ Module prefixes (e.g. ``("my_project.",)``) whose instances are
+ shallow-copied and searched attribute by attribute for arrays. Other
+ objects are passed unchanged.
is_array : callable, optional
- Predicate deciding whether a host-side value returned by the kernel
- (or reachable from a declared output) counts as an array to move
- back to the device. Defaults to ``isinstance(value, np.ndarray)``.
- Override this if the kernel returns/mutates a NumPy subclass or a
- custom host array type that should also be converted back to CuPy.
+ Predicate for host values to move back to the device (returned, or
+ reachable from a declared output). Default:
+ ``isinstance(value, np.ndarray)``.
outputs : sequence of int or str, optional
- Which arguments the kernel writes to. Only those are copied back to the
- device after the call, which avoids pointless device transfers for the
- (usually much larger) read-only inputs. Positional arguments are named
- by index, keyword arguments by name::
-
- interpolate = PyccelKernel(some_interpolation_kernel, outputs=(5,))
- interpolate(x, y, z, basis, coeffs, out) # `out` is argument 5
-
- A name also finds the argument when it is passed positionally, and an
- index also finds a keyword argument, if the parameter names of the
- kernel are known: from its Python signature, or from `parameters`
- (Pyccel-compiled kernels are builtins with no introspectable
- signature; a :class:`~cunumpy.kernels.Kernel` supplies the names of
- its host function). Without parameter names an index and a name are
- *not* interchangeable: declare the form you actually call with. An
- empty sequence declares that the kernel writes to none of its
- arguments. By default (``None``) every converted array is copied back,
- which is always correct but does more work.
+ The arguments the kernel writes: indices (negative from the end) for
+ positional arguments, names for keyword arguments. Only these are
+ copied back. A name also finds a positional argument, and an index a
+ keyword argument, only when the parameter names are known (Python
+ signature or `parameters`). ``()`` means no output; None (default)
+ copies back every converted array.
parameters : sequence of str or callable, optional
- The names of the positional parameters of `kernel`, or a function that
- returns them (or None if they are unknown), for resolving `outputs`
- names and indices. By default they are read from the signature of
- `kernel`, where it has one.
+ Names of the positional parameters of `kernel`, or a function returning
+ them (or None), to resolve `outputs`; needed for compiled kernels,
+ which have no Python signature. Default: from the signature.
+
+ Raises
+ ------
+ TypeError
+ If `outputs` is a bare int or str, or has entries of another type.
+
+ See Also
+ --------
+ outputs_from_annotations : Read `outputs` from the annotations.
+
+ Notes
+ -----
+ A call raises IndexError or KeyError if a declared output does not match
+ an argument of that call.
Examples
--------
- >>> from cunumpy.kernels import PyccelKernel
- >>> kernel = PyccelKernel(my_pyccelized_function)
- >>> kernel(out, x, y) # `out`, `x`, `y` may be NumPy or CuPy arrays
+ >>> def scale(a, x, out):
+ ... out[:] = a * x
+ >>> kernel = xp.kernels.PyccelKernel(scale, outputs=(2,))
+ >>> out = np.zeros(3)
+ >>> kernel(2.0, np.arange(3.0), out) # NumPy or CuPy arrays
+ >>> out
+ array([0., 2., 4.])
"""
def __init__(
@@ -136,7 +144,7 @@ def __repr__(self) -> str:
)
def _parameter_names(self) -> Sequence[str] | None:
- """The names of the positional parameters, or None if they are unknown."""
+ """Return the names of the positional parameters, or None if unknown."""
source = self._parameters
if source is not None:
return source() if callable(source) else tuple(source)
@@ -158,20 +166,7 @@ def _convert_to_numpy(
converted: list[tuple[Any, np.ndarray]],
memo: dict[int, Any],
) -> Any:
- """Recursively replace CuPy arrays in `value` by host copies.
-
- Every replacement is appended to `converted` as a
- ``(device_array, host_copy)`` pair.
-
- `memo` maps ``id(original) -> converted`` and is shared across all
- arguments of a single call. It serves two purposes: a device array
- reachable by several paths is copied to the host exactly once (so the
- kernel sees one shared array, as the caller intended, and the write-back
- happens once), and reference cycles terminate instead of recursing
- forever. Everything traversed here stays reachable from the caller's
- arguments for the duration of the call, so the `id` keys cannot be
- reused by unrelated objects.
- """
+ """Replace CuPy arrays in `value` by host copies, once each (`memo` by id)."""
key = id(value)
if key in memo:
return memo[key]
@@ -247,12 +242,7 @@ def _copy_to_device(self, value: Any) -> Any:
return result
def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> None:
- """Record the id of every host array reachable from `value`.
-
- Runs over the *converted* arguments, using the same traversal rules as
- :meth:`_convert_to_numpy`, so that an output declared as a container or
- an object contributes the arrays nested inside it.
- """
+ """Record the id of every host array reachable from `value`."""
if self._is_array(value):
found.add(id(value))
return
@@ -284,15 +274,7 @@ def _output_host_arrays(
args_np: list[Any],
kwargs_np: dict[str, Any],
) -> set[int]:
- """Ids of the host arrays reachable from the declared output arguments.
-
- Raises
- ------
- IndexError, KeyError
- If a declared output does not correspond to an argument of this
- call -- typically because an argument declared by index was passed
- as a keyword, or vice versa, and the parameter names are unknown.
- """
+ """Return the ids of the host arrays reachable from the declared outputs."""
found: set[int] = set()
seen: set[int] = set()
@@ -350,12 +332,7 @@ def _output_host_arrays(
return found
def _contains_cupy(self, value: Any, seen: set[int] | None = None) -> bool:
- """Whether `value` holds a CuPy array, following the same traversal
- rules as :meth:`_convert_to_numpy`.
-
- `seen` tracks already-visited containers so that reference cycles
- terminate.
- """
+ """Return whether `value` holds a CuPy array (traversed like the conversion)."""
if _is_device_array(value):
return True
@@ -430,7 +407,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
@property
def name(self) -> str:
- """Name of the wrapped kernel."""
+ """The name of the wrapped kernel."""
return getattr(self._kernel, "__name__", type(self._kernel).__name__)
@property
@@ -447,21 +424,22 @@ def use_cupy(self) -> bool:
@property
def object_modules(self) -> tuple[str, ...]:
- """Module prefixes whose instances are traversed for arrays."""
+ """The module prefixes whose instances are searched for arrays."""
return self._object_modules
@property
def is_array(self) -> Callable[[Any], bool]:
- """Predicate identifying host values to convert back to the device."""
+ """The predicate for host values to move back to the device."""
return self._is_array
@property
def outputs(self) -> tuple[int | str, ...] | None:
- """Declared output arguments, or ``None`` if every array is copied back."""
+ """The declared output arguments, or None if every array is copied back."""
return self._outputs
-#: The host implementations a kernel folder can have, see `HostImplementations`.
+#: The names of the host implementations a kernel can have, see
+#: :class:`~cunumpy.kernels.HostImplementations`.
HOST_IMPLEMENTATIONS = ("pyccel", "numba", "numpy", "python")
@@ -480,33 +458,74 @@ def _check_implementation(name: str | None) -> str | None:
def set_host_kernel_implementation(name: str | None) -> None:
- """Choose the host implementation every kernel runs, like :func:`set_backend`.
-
- ``"pyccel"``, ``"numba"``, ``"numpy"`` or ``"python"`` (the uncompiled
- source of the pyccel version); None restores the default (see
- :class:`HostImplementations`). A kernel without the chosen implementation,
- or whose chosen implementation is unavailable (e.g. pyccel failed to
- compile), raises instead of running another one. CUDA kernels are not
- affected: device arrays always run the CUDA version. The environment
- variable ``CUNUMPY_HOST_KERNEL_IMPLEMENTATION`` (read when cunumpy is imported)
- sets it for a whole run.
+ """Choose the host implementation every kernel runs, like ``set_backend``.
+
+ A kernel that lacks the chosen implementation, or cannot load it (e.g.
+ Pyccel failed to compile), raises instead of running another one. Device
+ calls are not affected. The setting is global, not per thread;
+ ``CUNUMPY_HOST_KERNEL_IMPLEMENTATION`` sets it at import.
+
+ Parameters
+ ----------
+ name : str or None
+ One of :data:`~cunumpy.kernels.HOST_IMPLEMENTATIONS` (``"python"`` is
+ the uncompiled source of the Pyccel version), or None for the default
+ rule of :class:`~cunumpy.kernels.HostImplementations`.
+
+ Raises
+ ------
+ ValueError
+ If `name` is not a known implementation.
+
+ Examples
+ --------
+ >>> xp.kernels.set_host_kernel_implementation("numpy")
+ >>> xp.kernels.get_host_kernel_implementation()
+ 'numpy'
+ >>> xp.kernels.set_host_kernel_implementation(None)
"""
global _KERNEL_IMPLEMENTATION
_KERNEL_IMPLEMENTATION = _check_implementation(name)
def get_host_kernel_implementation() -> str | None:
- """The host implementation set with :func:`set_host_kernel_implementation`, or None."""
+ """Return the host implementation that is set, or None for the default.
+
+ Returns
+ -------
+ str or None
+ The name set with :func:`~cunumpy.kernels.set_host_kernel_implementation`.
+
+ Examples
+ --------
+ >>> xp.kernels.get_host_kernel_implementation() is None
+ True
+ """
return _KERNEL_IMPLEMENTATION
@contextmanager
def use_host_kernel_implementation(name: str | None) -> Iterator[None]:
- """Temporarily choose the host implementation, like :func:`use_backend`.
+ """Choose the host implementation inside a ``with`` block, like ``use_backend``.
- For tests and benchmarks, e.g. ``with xp.kernels.use_host_kernel_implementation("numpy"):``
- to run the code path of a machine without pyccel. The setting is global,
- not per thread.
+ For tests and benchmarks, e.g. to run the code path of a machine without
+ Pyccel. The setting is global, not per thread; the previous one is
+ restored on exit.
+
+ Parameters
+ ----------
+ name : str or None
+ As for :func:`~cunumpy.kernels.set_host_kernel_implementation`.
+
+ Yields
+ ------
+ None
+
+ Examples
+ --------
+ >>> with xp.kernels.use_host_kernel_implementation("python"):
+ ... xp.kernels.get_host_kernel_implementation()
+ 'python'
"""
global _KERNEL_IMPLEMENTATION
previous = _KERNEL_IMPLEMENTATION
@@ -517,6 +536,7 @@ def use_host_kernel_implementation(name: str | None) -> Iterator[None]:
_KERNEL_IMPLEMENTATION = previous
+#: The names of the device implementations a kernel can have.
DEVICE_IMPLEMENTATIONS = ("cuda",)
@@ -537,24 +557,72 @@ def _check_device_implementation(name: str | None) -> str | None:
def set_device_kernel_implementation(name: str | None) -> None:
"""Choose the device implementation for dispatched kernels.
- Currently only ``"cuda"`` is supported; ``None`` restores automatic selection.
- Explicit CUDA selection raises if a kernel has no CUDA implementation, even
- with ``missing_cuda="fallback"``. This does not switch the array backend or
- affect host calls or direct CudaKernel calls. The import-time environment
- variable ``CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION`` initializes this setting.
+ With ``"cuda"``, a :class:`~cunumpy.kernels.Kernel` without a CUDA
+ implementation raises :class:`LookupError` on device arguments, even with
+ ``missing_cuda="fallback"``. Does not switch the array backend or affect
+ host calls or direct :class:`~cunumpy.kernels.CudaKernel` calls. Global,
+ not per thread; ``CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION`` sets it at import.
+
+ Parameters
+ ----------
+ name : str or None
+ ``"cuda"`` (see :data:`~cunumpy.kernels.DEVICE_IMPLEMENTATIONS`), or
+ None for automatic selection.
+
+ Raises
+ ------
+ ValueError
+ If `name` is not supported; the setting is unchanged.
+
+ Examples
+ --------
+ >>> xp.kernels.set_device_kernel_implementation("cuda")
+ >>> xp.kernels.get_device_kernel_implementation()
+ 'cuda'
+ >>> xp.kernels.set_device_kernel_implementation(None)
"""
global _DEVICE_KERNEL_IMPLEMENTATION
_DEVICE_KERNEL_IMPLEMENTATION = _check_device_implementation(name)
def get_device_kernel_implementation() -> str | None:
- """Return the requested device implementation, or None for automatic selection."""
+ """Return the requested device implementation, or None for automatic selection.
+
+ Returns
+ -------
+ str or None
+ The setting; None in automatic mode even when calls use CUDA.
+
+ Examples
+ --------
+ >>> xp.kernels.get_device_kernel_implementation() is None
+ True
+ """
return _DEVICE_KERNEL_IMPLEMENTATION
@contextmanager
def use_device_kernel_implementation(name: str | None) -> Iterator[None]:
- """Temporarily choose the device implementation; global, not per thread."""
+ """Choose the device implementation inside a ``with`` block.
+
+ The setting is global, not per thread; the previous one is restored on
+ exit, also after an exception.
+
+ Parameters
+ ----------
+ name : str or None
+ As for :func:`~cunumpy.kernels.set_device_kernel_implementation`.
+
+ Yields
+ ------
+ None
+
+ Examples
+ --------
+ >>> with xp.kernels.use_device_kernel_implementation("cuda"):
+ ... xp.kernels.get_device_kernel_implementation()
+ 'cuda'
+ """
global _DEVICE_KERNEL_IMPLEMENTATION
previous = _DEVICE_KERNEL_IMPLEMENTATION
set_device_kernel_implementation(name)
@@ -565,23 +633,44 @@ def use_device_kernel_implementation(name: str | None) -> Iterator[None]:
class HostImplementations:
- """The host implementations of one kernel, run by name or by the default rule.
+ """Hold the host implementations of one kernel and call the selected one.
- Each implementation is loaded on first use (a pyccel build, an import) and
- may be unavailable (no compiler, numba not installed); a failed load is
- remembered with its exception. A call runs the implementation chosen with
- :func:`set_host_kernel_implementation`, which must exist and load, or else the
- default: the first available of ``"pyccel"``, ``"numba"`` and ``"numpy"``,
- and as a last resort ``"python"``, with a warning (correct, but slow).
- Built by :meth:`Kernel.from_folder` from the files of a kernel folder.
+ Each implementation is loaded on first use (a Pyccel build, an import); a
+ failed load is remembered in `errors`. A call runs the implementation set
+ with :func:`~cunumpy.kernels.set_host_kernel_implementation`, or else the
+ first available of ``"pyccel"``, ``"numba"`` and ``"numpy"``, and as a
+ last resort ``"python"`` with a :class:`RuntimeWarning` (once). Built by
+ :meth:`Kernel.from_folder `.
Parameters
----------
name : str
Name of the kernel.
- loaders : Mapping[str, Callable[[], Callable]]
- Per implementation name (one of ``HOST_IMPLEMENTATIONS``), a function
- returning the implementation, or raising if it is unavailable.
+ loaders : mapping of str to callable
+ For each implementation name (from
+ :data:`~cunumpy.kernels.HOST_IMPLEMENTATIONS`), a function returning
+ the implementation or raising if it is unavailable. ``"python"`` is
+ required.
+
+ Attributes
+ ----------
+ errors : dict of str to Exception
+ The load errors, by implementation name.
+
+ Raises
+ ------
+ ValueError
+ If `loaders` has an unknown name or no ``"python"`` entry.
+
+ Examples
+ --------
+ >>> def push(x):
+ ... return x + 1
+ >>> host = xp.kernels.HostImplementations(
+ ... "push", {"numpy": lambda: push, "python": lambda: push}
+ ... )
+ >>> host.selected(), host(1)
+ ('numpy', 2)
"""
def __init__(
@@ -617,7 +706,18 @@ def python(self) -> Callable[..., Any]:
return self.get("python")
def available(self, name: str) -> bool:
- """Whether implementation `name` exists and loads (loads it now)."""
+ """Return whether implementation `name` exists and loads (loads it now).
+
+ Parameters
+ ----------
+ name : str
+ Implementation name.
+
+ Returns
+ -------
+ bool
+ False if it is missing or failed to load (see `errors`).
+ """
if name not in self._loaders:
return False
try:
@@ -627,7 +727,17 @@ def available(self, name: str) -> bool:
return True
def get(self, name: str) -> Callable[..., Any]:
- """Implementation `name`, loaded on first use.
+ """Return implementation `name`, loading it on first use.
+
+ Parameters
+ ----------
+ name : str
+ Implementation name.
+
+ Returns
+ -------
+ callable
+ The implementation.
Raises
------
@@ -660,7 +770,17 @@ def get(self, name: str) -> Callable[..., Any]:
return function
def selected(self) -> str:
- """The implementation a call runs now (loads the default on first use)."""
+ """Return the name of the implementation a call runs now.
+
+ Loads the default on first use and keeps it.
+
+ Returns
+ -------
+ str
+ The name set with
+ :func:`~cunumpy.kernels.set_host_kernel_implementation`, or the
+ default.
+ """
if _KERNEL_IMPLEMENTATION is not None:
return _KERNEL_IMPLEMENTATION
if self._default is None:
@@ -680,7 +800,13 @@ def selected(self) -> str:
return self._default
def build(self) -> str:
- """Load the default implementation now, e.g. compile at setup; its name."""
+ """Load the default implementation now, e.g. to compile at setup.
+
+ Returns
+ -------
+ str
+ The name of the selected implementation.
+ """
return self.selected()
def __call__(self, *args: Any, **kwargs: Any) -> Any:
@@ -688,31 +814,40 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
class CompiledHostKernel:
- """A host kernel compiled on its first call, with a fallback.
+ """Call a host kernel that is compiled on its first call, with a fallback.
- cunumpy does not compile anything itself: `compiler` is the caller's
- function that turns the kernel module into a compiled one, e.g. a wrapper
- around ``pyccel.epyccel`` with a cache, or a function importing modules
- compiled ahead of time.
+ cunumpy compiles nothing itself: `compiler` is the caller's function, e.g. a
+ wrapper around ``pyccel.epyccel`` with a cache, or one importing modules
+ compiled ahead of time. Used by ``KernelCatalog.from_package(...,
+ compile_host=...)``.
Parameters
----------
- module : str | ModuleType
+ module : str or module
The module that defines the kernel, or its name.
name : str
Name of the kernel function in `module`.
- compiler : Callable[[ModuleType], ModuleType]
- Returns the compiled form of the module (with the same function names).
- fallback : Callable | None
- Called instead when compilation fails (e.g. a vectorized NumPy
- implementation with the same arguments). Without one, the uncompiled
- Python function is called, with a warning: correct, but slow.
+ compiler : callable
+ Takes the module and returns its compiled form (same function names).
+ fallback : callable, optional
+ Called when compilation fails, e.g. a NumPy version with the same
+ arguments. Without one, the uncompiled Python function runs, with a
+ :class:`RuntimeWarning` (once).
+
+ Attributes
+ ----------
+ error : Exception or None
+ The exception of a failed build.
- Notes
- -----
- :attr:`compiled` builds the kernel and reports whether that worked, so
- that callers can choose another code path; :attr:`error` keeps the
- exception of a failed build.
+ Examples
+ --------
+ >>> def no_compiler(module):
+ ... raise RuntimeError("no Pyccel")
+ >>> kernel = xp.kernels.CompiledHostKernel(
+ ... "math", "sqrt", no_compiler, fallback=np.sqrt
+ ... )
+ >>> kernel(4.0), kernel.compiled
+ (np.float64(2.0), False)
"""
def __init__(
@@ -757,7 +892,13 @@ def fallback(self) -> Callable[..., Any] | None:
return self._fallback
def build(self) -> Callable[..., Any] | None:
- """Compile now (once); the compiled function, or None if that failed."""
+ """Compile now, once; a failure is kept in `error`.
+
+ Returns
+ -------
+ callable or None
+ The compiled function, or None if compilation failed.
+ """
if not self._built:
self._built = True
try:
@@ -790,12 +931,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
def _c_ordered_with_gaps(array: np.ndarray) -> bool:
- """Whether `array` is laid out in C order, possibly with gaps between entries.
-
- Every stride is positive and at least the extent of the next axis, e.g. the
- first ``n`` columns ``a[:, :n]`` of a C-contiguous ``a`` or every other
- entry ``a[::2]``. Axes of length 1 are ignored, as their stride is never used.
- """
+ """Return whether `array` is in C order, possibly with gaps (e.g. ``a[:, :n]``)."""
extent = array.itemsize
for length, stride in reversed(list(zip(array.shape, array.strides, strict=True))):
if length == 1:
@@ -807,15 +943,7 @@ def _c_ordered_with_gaps(array: np.ndarray) -> bool:
def _with_unit_axis_strides(array: np.ndarray) -> np.ndarray:
- """`array`, or a view of it, whose axes of length 1 have strides in C order.
-
- NumPy never steps along an axis of length 1, so it leaves any stride there:
- ``a[:, order]`` of a ``(1, n)`` array has strides ``(8, 8)`` and still counts
- as C-contiguous. Pyccel's wrappers do use that stride to check the array
- against its shape, and abort the process when it is smaller than the
- extent of the next axes, so give such an axis the stride a fresh C-ordered
- array would have (a view, no copy).
- """
+ """Return `array`, or a view, with C-order strides on its length-1 axes."""
if array.size == 0:
return array
strides = list(array.strides)
@@ -835,30 +963,42 @@ def _with_unit_axis_strides(array: np.ndarray) -> np.ndarray:
def as_kernel_array(
value: Any, like: Any, dtype: Any = None, *, strided: bool = False
) -> Any:
- """`value` as an array the kernel chosen for `like` takes.
-
- On the device of `like` (a CuPy array if `like` is one, a NumPy array
- otherwise), C-contiguous and with `dtype` (any dtype if None): `value`
- itself when it already is such an array, else one copy, moved to the other
- side if needed (counted by `count_transfers()`). For input arrays of a
- :class:`Kernel` with ``dispatch="arrays"``, whose choice follows the arrays:
- pass the main array (e.g. the grid) as `like`, and the kernel gets
- arguments all on one side::
-
- convert = functools.partial(xp.kernels.as_kernel_array, like=grid, dtype=float)
- deposit(convert(positions), convert(weights), grid, ...)
-
- For an array the kernel writes, use :func:`kernel_output`, which copies a
- converted array back.
-
- With ``strided=True`` a NumPy array for a host kernel is also taken
- unchanged when it is in C order with gaps: positive strides, each at least
- the extent of the next axis, such as the first ``n`` columns
- ``storage[:, :n]`` of a component-major marker buffer, which C-contiguity
- would copy on every call. Pyccel's wrappers take such arrays without a copy
- (they refuse F order and abort on negative strides, which are still copied);
- a host kernel that needs contiguous memory, e.g. one using raw pointers,
- must not use it. CuPy arrays are made C-contiguous either way.
+ """Return `value` as an input array for the kernel chosen for `like`.
+
+ For a :class:`~cunumpy.kernels.Kernel` with ``dispatch="arrays"``: pass the
+ main array (e.g. the grid) as `like` so that all arguments end up on one
+ side. Returns `value` itself if it already fits, else one copy (moved
+ across if needed, counted by ``count_transfers()``). For arrays the kernel
+ writes, use :func:`~cunumpy.kernels.kernel_output`.
+
+ Parameters
+ ----------
+ value : array_like
+ The array to convert.
+ like : array
+ A CuPy array for a device kernel, anything else for a host kernel.
+ dtype : dtype, optional
+ Required dtype; by default any.
+ strided : bool, optional
+ Also take a NumPy array unchanged when it is in C order with gaps
+ (positive strides, each at least the extent of the next axis), such as
+ ``storage[:, :n]``. Pyccel's wrappers accept these without a copy; a
+ host kernel that needs contiguous memory must not use it. Default
+ False. CuPy arrays are always made C-contiguous.
+
+ Returns
+ -------
+ array
+ A C-contiguous (or, with `strided`, C-ordered) array on the side of
+ `like`, with `dtype`.
+
+ Examples
+ --------
+ >>> grid = np.zeros((2, 3))
+ >>> xp.kernels.as_kernel_array(grid, like=grid) is grid
+ True
+ >>> xp.kernels.as_kernel_array([1, 2], like=grid, dtype=float)
+ array([1., 2.])
"""
if is_gpu(like):
import cupy
@@ -885,7 +1025,7 @@ def as_kernel_array(
def _same_memory(buffer: Any, out: Any) -> bool:
- """Whether `buffer` is `out` or a view of exactly its entries (no copy back needed)."""
+ """Return whether `buffer` is `out` or a view of exactly its entries."""
if buffer is out:
return True
return (
@@ -905,18 +1045,35 @@ def _same_memory(buffer: Any, out: Any) -> bool:
def kernel_output(
out: Any, like: Any, dtype: Any = None, *, strided: bool = False
) -> Generator[Any]:
- """The buffer a kernel chosen for `like` writes, copied back into `out` after it.
+ """Provide the output buffer for the kernel chosen for `like`, copied into `out`.
- Yields `out` itself if :func:`as_kernel_array` takes it unchanged (then the
- kernel writes into it directly), else a converted copy whose contents are
- written into `out` when the block ends without an error (moved back to the
- side of `out`)::
+ The buffer is `out` itself if :func:`~cunumpy.kernels.as_kernel_array`
+ takes it unchanged, else a converted copy whose contents are written into
+ `out` (on its own side) when the block ends without an error.
- with xp.kernels.kernel_output(result, like=grid, dtype=float) as buffer:
- gather(convert(positions), grid, buffer, ...)
+ Parameters
+ ----------
+ out : array
+ The array the kernel should write.
+ like : array
+ A CuPy array for a device kernel, anything else for a host kernel.
+ dtype : dtype, optional
+ Required dtype; by default any.
+ strided : bool, optional
+ As for :func:`~cunumpy.kernels.as_kernel_array`. Default False.
+
+ Yields
+ ------
+ array
+ The buffer to pass to the kernel.
- `strided` is passed on to :func:`as_kernel_array`: with True, a host kernel
- writes directly into a NumPy `out` in C order with gaps.
+ Examples
+ --------
+ >>> out = np.zeros(2, dtype=np.float32)
+ >>> with xp.kernels.kernel_output(out, like=out, dtype=float) as buffer:
+ ... buffer[:] = 1.5 # the kernel writes a float64 copy
+ >>> out
+ array([1.5, 1.5], dtype=float32)
"""
buffer = as_kernel_array(out, like, dtype, strided=strided)
yield buffer
@@ -933,22 +1090,30 @@ def kernel_output(
def outputs_from_annotations(function: Callable[..., Any]) -> tuple[str, ...] | None:
- """The names of the parameters a kernel may write to, read from its annotations.
+ """Return the parameters a kernel may write to, read from its annotations.
- Pyccel kernels mark what they only read with ``Final``: ``x: "Final[float[:]]"``.
- A parameter is an output unless it is annotated ``Final`` (or ``const``) or
- has the annotation of a scalar (``int``, ``float``, ``bool``, ``complex``,
- ``str``); an annotated array or an argument object may be written to. Use the
- result as the `outputs` of a :class:`PyccelKernel` so that only those arrays
- are copied back to the device::
+ A parameter is an output unless it is annotated ``Final`` (Pyccel's mark
+ for read-only, e.g. ``"Final[float[:]]"``), ``const``, or a scalar type
+ (``int``, ``float``, ``bool``, ``complex``, ``str``). Use the result as the
+ `outputs` of a :class:`~cunumpy.kernels.PyccelKernel`.
- PyccelKernel(push, outputs=outputs_from_annotations(push))
+ Parameters
+ ----------
+ function : callable
+ The Python kernel function.
Returns
-------
- tuple[str, ...] | None
- The parameter names, in order, or None if `function` has no signature or
- none of its parameters is annotated (so nothing can be said).
+ tuple of str or None
+ The parameter names, in order, or None if `function` has no signature
+ or no annotated parameter.
+
+ Examples
+ --------
+ >>> def push(x: "Final[float[:]]", dt: float, v: "float[:]"):
+ ... v[:] += dt * x
+ >>> xp.kernels.outputs_from_annotations(push)
+ ('v',)
"""
try:
parameters = inspect.signature(function).parameters.values()
diff --git a/src/cunumpy/_metal_kernel.py b/src/cunumpy/_metal_kernel.py
index cbf45d6..b011d3a 100644
--- a/src/cunumpy/_metal_kernel.py
+++ b/src/cunumpy/_metal_kernel.py
@@ -1,4 +1,7 @@
-"""Metal kernels for the GPU of Apple silicon Macs, run through MLX."""
+"""Metal kernels for the GPU of Apple silicon Macs, run through MLX.
+
+See :doc:`/kernels/metal-kernel`.
+"""
from __future__ import annotations
@@ -29,7 +32,18 @@ def _mlx() -> Any:
def metal_available() -> bool:
- """Whether `MetalKernel` can run here (MLX installed and a Metal GPU present)."""
+ """Check whether a :class:`~cunumpy.kernels.MetalKernel` can run here.
+
+ Returns
+ -------
+ bool
+ True if MLX is installed and reports a Metal GPU.
+
+ Examples
+ --------
+ >>> xp.kernels.metal_available() # doctest: +SKIP
+ True
+ """
try:
_mlx()
except (ImportError, RuntimeError):
@@ -40,49 +54,44 @@ def metal_available() -> bool:
class MetalKernel:
"""A Metal Shading Language kernel for the GPU of Apple silicon, run with MLX.
- The arrays are NumPy arrays: they are copied to MLX arrays for the launch
- and the results are written back into the output arrays you pass, like a
- :class:`~cunumpy.kernels.PyccelKernel` writes into its output arguments. No
- backend switch is needed. The copies are counted by
- :func:`~cunumpy.profiling.count_transfers`.
+ Takes NumPy arrays, no backend switch needed: every call copies the inputs
+ to MLX arrays and the results back into the output arrays you pass,
+ counted by :func:`~cunumpy.profiling.count_transfers`. The kernel is
+ compiled by MLX at its first call and cached. Needs ``cunumpy[metal]``;
+ see :doc:`/kernels/metal-kernel`.
Parameters
----------
source : str
- Body of the kernel function (MSL). MLX generates the signature from
- `inputs` and `outputs`: every name is a pointer to the flat, row-major
- data of that array (``const device T*`` for inputs, ``device T*`` for
- outputs), and ``thread_position_in_grid`` and the other Metal
- attributes used in the body are added to the signature automatically.
- Names from `template` are available as compile-time constants.
- inputs : Sequence[str]
- Names of the input arrays, in the order they are passed to the call.
- outputs : Sequence[str]
- Names of the output arrays, in the order they are passed as `out`.
- name : str
+ Body of the kernel function. MLX generates the signature: each name
+ in `inputs` and `outputs` is a pointer to the flat, row-major data of
+ that array, and ``thread_position_in_grid`` and the other Metal
+ attributes used in the body are added automatically.
+ inputs : sequence of str
+ Names of the input arrays, in call order.
+ outputs : sequence of str
+ Names of the output arrays, in the order passed as `out`.
+ name : str, optional
Name of the kernel; part of the compiled function name.
- header : str
- Source placed before the kernel function: includes, ``#define``\\ s and
- helper functions.
- threadgroup : int | Sequence[int]
- Threads per threadgroup: an integer, or 1 to 3 integers.
- float64 : {"error", "cast"}
- The GPU has no float64. With "error" (the default) a float64 input or
- output raises ``TypeError``. With "cast" float64 inputs are computed
- in float32 and float64 outputs are filled from float32 results, which
- is what to pick when float32 precision is enough.
- atomic_outputs : bool
+ header : str, optional
+ Source before the kernel function: includes, defines, helper functions.
+ threadgroup : int or sequence of int, optional
+ Threads per threadgroup, 1 to 3 integers.
+ float64 : {"error", "cast"}, optional
+ The Apple GPU has no float64. ``"error"`` (default) raises
+ ``TypeError`` for a float64 input or output; ``"cast"`` computes
+ float64 data in float32.
+ atomic_outputs : bool, optional
Declare the outputs as ``device atomic*`` for atomic updates.
- init_value : float | None
- Value the outputs are filled with before the launch. Outputs are
- otherwise uninitialized: write every element, or pass the old array as
- an input as well to read its values.
+ init_value : float, optional
+ Value the outputs are filled with before the launch. Otherwise they
+ are uninitialized: write every element, or pass the old array as an
+ input too to read it.
- Notes
- -----
- The kernel is compiled by MLX at its first call and cached. All inputs
- and outputs are made C-contiguous, so ``a[i]`` in the source is the flat
- index of the NumPy array.
+ Raises
+ ------
+ ValueError
+ If `float64` is unknown, there is no output, or names repeat.
Examples
--------
@@ -91,8 +100,9 @@ class MetalKernel:
... inputs=["x", "a"],
... outputs=["y"],
... )
- >>> y = np.empty(n, dtype=np.float32)
- >>> scale(x, np.float32([2.0]), out=y)
+ >>> y = np.empty(4, dtype=np.float32)
+ >>> scale(np.arange(4, dtype=np.float32), 2.0, out=y) # doctest: +SKIP
+ array([0., 2., 4., 6.], dtype=float32)
"""
def __init__(
@@ -140,7 +150,7 @@ def _as_triple(value: int | Sequence[int], what: str, default: int) -> tuple:
@staticmethod
def _as_input(value: Any) -> np.ndarray:
- """A NumPy array; Python scalars become float32 or int32 arrays of shape (1,)."""
+ """Return `value` as a NumPy array; Python scalars become 1-element arrays."""
if isinstance(value, (bool, int)):
return np.array([value], dtype=np.int32)
if isinstance(value, float):
@@ -179,26 +189,37 @@ def __call__(
n_threads: int | Sequence[int] | None = None,
template: Mapping[str, Any] | None = None,
) -> Any:
- """Launch the kernel.
+ """Launch the kernel; blocks until the results are copied back.
Parameters
----------
*args : numpy.ndarray or scalar
The inputs, in the order of `inputs`. Python scalars become
- 1-element float32 or int32 arrays; read them as ``a[0]`` in the source.
- out : numpy.ndarray or Sequence[numpy.ndarray]
+ 1-element float32 or int32 arrays; read them as ``a[0]``.
+ out : numpy.ndarray or sequence of numpy.ndarray
The output arrays, in the order of `outputs`; their shapes and
- dtypes define the outputs and they are filled in place.
- n_threads : int | Sequence[int] | None
- Total number of threads (1 to 3 dimensions), not threadgroups. The
- default is the first axis of the first output.
- template : Mapping[str, int | bool | numpy.dtype] | None
- Compile-time constants; each name is a constant in the source. A
- different value compiles another variant.
+ dtypes define the outputs, and they are filled in place.
+ n_threads : int or sequence of int, optional
+ Total number of threads (1 to 3 dimensions), not threadgroups; by
+ default the first axis of the first output.
+ template : mapping, optional
+ Compile-time constants (int, bool or dtype) by name, e.g.
+ ``{"NSTEPS": 200}``; a new value compiles another variant.
Returns
-------
- The output array, or a tuple of them if there are several.
+ numpy.ndarray or tuple of numpy.ndarray
+ The output array, or a tuple of them if there are several.
+
+ Raises
+ ------
+ TypeError
+ If the number of inputs or outputs is wrong, or a dtype is not
+ supported on the GPU.
+ ImportError
+ If MLX is not installed.
+ RuntimeError
+ If MLX reports no Metal GPU.
"""
mx = _mlx()
if len(args) != len(self.inputs):
diff --git a/src/cunumpy/_mirror.py b/src/cunumpy/_mirror.py
index fbcd650..44192f6 100644
--- a/src/cunumpy/_mirror.py
+++ b/src/cunumpy/_mirror.py
@@ -1,15 +1,10 @@
"""Host buffers owned by other libraries, mirrored on the device.
-Accumulation kernels write into buffers that another library owns and keeps
-using on the host (e.g. a stencil vector's ``_data`` that is then exchanged
-over MPI). On the GPU the kernel writes into a device array instead, and one
-transfer per accumulation is unavoidable; :class:`DeviceMirror` makes that
-transfer explicit and keeps the host array's identity, so the owning library
-still sees its own buffer.
-
-Every host/device copy of a mirror goes through ``xp.to_cupy`` or the
-mirror's own ``to_device()``/``to_host()`` calls, so that transfers are easy
-to find (and to account for) in application code.
+Accumulation kernels may have to write into a buffer that another library owns
+and keeps using on the host (e.g. a stencil vector's ``_data`` exchanged over
+MPI). :class:`~cunumpy.memory.DeviceMirror` makes the one transfer per
+accumulation explicit and keeps the host array's identity. See
+:doc:`/guides/data-movement`.
"""
from __future__ import annotations
@@ -28,21 +23,17 @@
class DeviceMirror:
"""A host NumPy array paired with a lazily created device copy.
- On the CuPy backend, `device` is a CuPy array of the same shape and dtype
- as the host array, allocated (and filled from the host) on first use.
- On the NumPy backend, `device` is the host array itself, so a kernel that
- is given ``mirror.device`` writes straight into the host buffer, without
- any copy.
-
- Transfers are explicit: `to_device()` copies host to device and
- `to_host()` copies device to host, in place, so the host array keeps its
- identity. Both are no-ops on the NumPy backend. `zero()` clears the
- buffer the kernel writes into.
+ On the CuPy backend, :attr:`device` is a CuPy array of the same shape and
+ dtype, allocated as a copy of the host on first use. On the NumPy backend
+ it is the host array itself, so kernels write straight into it. Transfers
+ are explicit and in place (:meth:`to_device`, :meth:`to_host`, no-ops on
+ NumPy) and counted by :func:`~cunumpy.profiling.count_transfers`. Make the
+ mirror's CUDA device current before using it.
Parameters
----------
host : numpy.ndarray
- The host buffer, owned by the caller (or by another library).
+ The host buffer, owned by the caller or another library.
Raises
------
@@ -51,12 +42,10 @@ class DeviceMirror:
Examples
--------
- >>> data = np.zeros((4, 3)) # owned by another library
- >>> mirror = DeviceMirror(data)
- >>> mirror.zero() # doctest: +SKIP
- >>> accumulate(mirror.device, n_threads=n) # doctest: +SKIP
- >>> mirror.to_host() # doctest: +SKIP
- >>> mirror.host is data
+ >>> data = np.ones((4, 3)) # owned by another library
+ >>> mirror = xp.memory.DeviceMirror(data)
+ >>> accumulate(mirror.zero().device, n_threads=n) # doctest: +SKIP
+ >>> mirror.to_host().host is data # same object, new values
True
"""
@@ -98,12 +87,7 @@ def dtype(self) -> np.dtype:
@property
def device(self) -> Any:
- """The array kernels write into: a CuPy array, or the host on NumPy.
-
- On the CuPy backend the device array is allocated on first access, as
- a copy of the host array. On the NumPy backend this is the host array
- itself.
- """
+ """The array kernels write into: CuPy (copied from the host on first use), or the host on NumPy."""
if not _cupy_backend():
return self._host
self._check_bound()
@@ -137,9 +121,9 @@ def _check_bound(self) -> None:
def rebind(self, host: np.ndarray) -> DeviceMirror:
"""Bind to a new host array, e.g. after its owner reallocated it.
- The device buffer is kept if the shape and dtype are unchanged (it
- is not refreshed: call `to_device()` for that), otherwise it is
- dropped and allocated again on the next use.
+ The device array is kept if shape and dtype are unchanged (not
+ refreshed: call :meth:`to_device`), otherwise it is allocated again on
+ the next use.
Parameters
----------
@@ -159,12 +143,18 @@ def rebind(self, host: np.ndarray) -> DeviceMirror:
return self
def to_device(self) -> DeviceMirror:
- """Copy the host array to the device (no-op on the NumPy backend).
+ """Copy the host array into the device array (no-op on the NumPy backend).
+
+ Returns
+ -------
+ DeviceMirror
+ `self`, for chaining.
Raises
------
ValueError
- If the host array was reallocated with another shape or dtype.
+ If the host array was reallocated with another shape or dtype
+ (call :meth:`rebind`).
"""
self._check_bound()
if not _cupy_backend():
@@ -184,21 +174,31 @@ def to_device(self) -> DeviceMirror:
return self
def to_host(self, *, stream: Any = None, event: Any = None) -> DeviceMirror:
- """Copy the device array into the host array, in place.
+ """Copy the device array into the host array, in place, and wait for it.
- The host array keeps its identity, so a library that holds the buffer
- sees the new values. No-op on the NumPy backend, and if no device
- array has been created yet.
+ The host array keeps its identity, so the library that holds it sees
+ the new values. No-op on the NumPy backend and before the device array
+ exists.
- Make the mirror's device current. Pass ``stream=`` to copy on a known
- producer stream, or ``event=`` to wait for production on the current
- stream. With neither, only the current stream's work is ordered.
- The method blocks until the host copy is complete.
+ Parameters
+ ----------
+ stream : cupy.cuda.Stream, optional
+ Producer stream to copy on; by default only the current stream's
+ work is ordered before the copy.
+ event : cupy.cuda.Event, optional
+ Producer event the current stream waits for before the copy. Pass
+ at most one of `stream` and `event`.
+
+ Returns
+ -------
+ DeviceMirror
+ `self`, for chaining.
Raises
------
ValueError
- If the host array was reallocated with another shape or dtype.
+ If the host array was reallocated with another shape or dtype
+ (call :meth:`rebind`), or for both `stream` and `event`.
"""
if stream is not None and event is not None:
raise ValueError("pass only one of stream and event")
@@ -235,7 +235,13 @@ def to_host(self, *, stream: Any = None, event: Any = None) -> DeviceMirror:
return self
def zero(self) -> DeviceMirror:
- """Zero the array kernels write into (the device array, or the host on NumPy)."""
+ """Zero the array kernels write into (the device array, or the host on NumPy).
+
+ Returns
+ -------
+ DeviceMirror
+ `self`, for chaining.
+ """
self._check_bound()
if not _cupy_backend():
self._host.fill(0)
diff --git a/src/cunumpy/_morton.py b/src/cunumpy/_morton.py
index 30e15ff..4e20a8a 100644
--- a/src/cunumpy/_morton.py
+++ b/src/cunumpy/_morton.py
@@ -1,23 +1,19 @@
-"""Morton (Z-order) keys, the same on host and device.
-
-The host side of ``cunumpy/morton.cuh``. A Morton key interleaves the bits of
-the integer cell coordinates of a point, so sorting points by their keys
-orders them along a Z-shaped space-filling curve: points close in space end up
-close in memory (better locality for gathers and neighbour loops), and the
-points of every node of a quadtree (2D) or octree (3D) on the same box are a
-contiguous range of the sorted array::
-
- scales = xp.algorithms.morton_scales(lower, upper, levels)
- keys = xp.algorithms.morton_keys(positions, lower, upper, levels) # uint64, one per point
- keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, charges)
- node = keys >> np.uint64(ndim * (levels - level)) # node index at `level`
-
-Bit layout: with ``levels`` bits per axis the key has ``ndim * levels`` bits;
-axis 0 is the lowest bit of every group of ``ndim`` bits. The top group is the
-child of the root a point lies in, the next group the child of that child, and
-so on. In a kernel, ``cunumpy_morton_key2`` / ``cunumpy_morton_key3`` with the
-:func:`morton_scales` of the host return exactly the keys of
-:func:`morton_keys` (the cell index is ``floor((x - lower) * scale)`` in both).
+"""Morton (Z-order) keys, the same on the host and in ``cunumpy/morton.cuh``.
+
+A Morton key interleaves the bits of the integer cell coordinates of a point.
+Sorting by it orders points along a Z-shaped curve: points close in space end
+up close in memory, and the points of every quadtree (2D) or octree (3D) node
+on the box are a contiguous range::
+
+ keys = xp.algorithms.morton_keys(positions, lower, upper, levels)
+ keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
+ node = keys >> np.uint64(ndim * (levels - level)) # node index at `level`
+
+With ``levels`` bits per axis the key has ``ndim * levels`` bits; axis 0 is
+the lowest bit of every group of ``ndim`` bits, and the top group is the child
+of the root. In a kernel, ``cunumpy_morton_key2`` / ``cunumpy_morton_key3``
+with the :func:`morton_scales` of the host return the keys of
+:func:`morton_keys`.
"""
from __future__ import annotations
@@ -91,14 +87,29 @@ def morton_encode(*cells: Any) -> Any:
Parameters
----------
- *cells : arrays of non-negative integers
+ *cells : array of int
One array per axis (2 or 3 of them), broadcast against each other;
- values must fit in 32 bits (2D) or 21 bits (3D).
+ non-negative, fitting in 32 bits (2D) or 21 bits (3D).
Returns
-------
array of uint64
- The keys, bit ``ndim * b + a`` holding bit ``b`` of axis ``a``.
+ The keys, bit ``ndim * b + a`` holding bit ``b`` of axis ``a``, on the
+ backend of `cells`.
+
+ Raises
+ ------
+ ValueError
+ If there are not 2 or 3 arrays.
+
+ See Also
+ --------
+ morton_decode : The inverse.
+
+ Examples
+ --------
+ >>> xp.algorithms.morton_encode(xp.asarray([1, 0, 3]), xp.asarray([0, 1, 3]))
+ array([ 1, 2, 15], dtype=uint64)
"""
ndim = len(cells)
_check_levels(ndim, 1)
@@ -111,12 +122,31 @@ def morton_encode(*cells: Any) -> Any:
def morton_decode(keys: Any, ndim: int) -> tuple[Any, ...]:
- """The integer cell coordinates of Morton keys, the inverse of :func:`morton_encode`.
+ """Return the integer cell coordinates of Morton keys.
+
+ The inverse of :func:`morton_encode`.
+
+ Parameters
+ ----------
+ keys : array of uint64
+ Morton keys.
+ ndim : int
+ Number of axes, 2 or 3.
Returns
-------
tuple of arrays of uint64
- One array per axis, shaped like `keys`.
+ One array per axis, shaped like `keys`, on its backend.
+
+ Raises
+ ------
+ ValueError
+ If `ndim` is not 2 or 3.
+
+ Examples
+ --------
+ >>> xp.algorithms.morton_decode(xp.asarray([1, 2, 15], dtype=np.uint64), 2)
+ (array([1, 0, 3], dtype=uint64), array([0, 1, 3], dtype=uint64))
"""
_check_levels(ndim, 1)
xpm = _module(keys)
@@ -125,10 +155,33 @@ def morton_decode(keys: Any, ndim: int) -> tuple[Any, ...]:
def morton_scales(lower: Sequence[float], upper: Sequence[float], levels: int) -> Any:
- """Cells per unit length of every axis, ``2**levels / (upper - lower)``.
+ """Return the cells per unit length of every axis, ``2**levels / (upper - lower)``.
+
+ Pass them to ``cunumpy_morton_key2`` / ``cunumpy_morton_key3`` in a kernel so
+ that it computes the keys of :func:`morton_keys`.
- The numbers to pass to ``cunumpy_morton_key2`` / ``_key3`` in a kernel so
- that it computes the keys of :func:`morton_keys`. A float64 NumPy array.
+ Parameters
+ ----------
+ lower, upper : sequence of float
+ Corners of the box, one value per axis (2 or 3).
+ levels : int
+ Bits per axis: 1 to 32 in 2D, 1 to 21 in 3D.
+
+ Returns
+ -------
+ numpy.ndarray of float64
+ One scale per axis, always on the host.
+
+ Raises
+ ------
+ ValueError
+ If the corners have different lengths, a dimension other than 2 or 3, or
+ are equal on an axis, or if `levels` is out of range.
+
+ Examples
+ --------
+ >>> xp.algorithms.morton_scales([0.0, 0.0], [1.0, 2.0], 4)
+ array([16., 8.])
"""
lower = np.asarray(lower, dtype=np.float64)
upper = np.asarray(upper, dtype=np.float64)
@@ -149,13 +202,12 @@ def morton_keys(
upper: Sequence[float],
levels: int,
) -> Any:
- """Morton keys of points in the box ``[lower, upper]``, ``levels`` bits per axis.
+ """Return the Morton keys of points in the box ``[lower, upper]``.
The cell of a point along an axis is ``floor((x - lower) * scale)`` with the
- :func:`morton_scales`, clipped to ``[0, 2**levels - 1]``, so points on or
- outside the box get the cell at the nearest face. A point exactly on a cell
- boundary belongs to the upper cell. ``lower > upper`` on an axis reverses
- that axis (cell 0 at ``lower``).
+ :func:`morton_scales`, clipped to ``[0, 2**levels - 1]``: points on or outside
+ the box get the cell at the nearest face, and a point on a cell boundary
+ belongs to the upper cell. ``lower > upper`` on an axis reverses that axis.
Parameters
----------
@@ -169,7 +221,22 @@ def morton_keys(
Returns
-------
array of uint64, shape (n,)
- On the backend of `positions`.
+ The keys, on the backend of `positions`.
+
+ Raises
+ ------
+ ValueError
+ If the shapes do not match or `levels` is out of range.
+
+ See Also
+ --------
+ sort_by_key : Sort points by their keys.
+
+ Examples
+ --------
+ >>> positions = xp.asarray([[0.1, 0.1], [0.9, 0.1], [0.9, 0.9]])
+ >>> xp.algorithms.morton_keys(positions, [0.0, 0.0], [1.0, 1.0], levels=2)
+ array([ 0, 5, 15], dtype=uint64)
"""
xpm = _module(positions)
positions = xpm.asarray(positions)
diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py
index fd46415..e382423 100644
--- a/src/cunumpy/_mpi.py
+++ b/src/cunumpy/_mpi.py
@@ -22,19 +22,39 @@
def synchronize_for_mpi(*arrays: Any, stream: Any = None, event: Any = None) -> None:
"""Wait for pending device work before MPI reads or writes `arrays`.
- CuPy launches kernels asynchronously; MPI does not know about CUDA streams.
- Passing a device buffer to MPI while a kernel is still writing it sends
- whatever is in memory at that moment -- silently wrong data, no error. Call
- this before every MPI call that uses device buffers. It synchronizes the
- devices owning the arrays if no producer `stream` or `event` is supplied.
- An explicit dependency waits only for that producer. Host
- buffers and the NumPy backend cost nothing. (After MPI returns, no
- synchronization is needed: kernels launched later see the received data.)
+ CuPy launches kernels asynchronously and MPI knows nothing about CUDA
+ streams: a device buffer that a kernel is still writing is sent as it is,
+ silently wrong. Call this before every MPI call with device buffers. No
+ synchronization is needed after MPI returns. Host arrays and the NumPy
+ backend cost nothing.
Parameters
----------
- *arrays
- The buffers about to be passed to MPI; ``None`` entries are ignored.
+ *arrays : array or None
+ The buffers about to be passed to MPI; None and host arrays are
+ ignored.
+ stream : cupy.cuda.Stream, optional
+ Producer stream to wait for, instead of the whole devices owning
+ `arrays`. It must cover all the device buffers.
+ event : cupy.cuda.Event, optional
+ Producer event to wait for instead; pass at most one of `stream` and
+ `event`.
+
+ Raises
+ ------
+ ValueError
+ If both `stream` and `event` are given.
+ TypeError
+ If device buffers come with a host stream or event.
+
+ See Also
+ --------
+ mpi_buffer : Synchronizes, and stages through the host if needed.
+
+ Examples
+ --------
+ >>> send, recv = np.ones(4), np.empty(4)
+ >>> xp.mpi.synchronize_for_mpi(send, recv) # no-op for host arrays
"""
if stream is not None and event is not None:
raise ValueError("pass only one of stream and event")
@@ -66,23 +86,51 @@ def synchronize_for_mpi(*arrays: Any, stream: Any = None, event: Any = None) ->
def set_mpi_cuda_aware(value: bool | None) -> None:
- """Tell `mpi_buffer()` whether MPI can take device buffers.
+ """Tell :func:`mpi_buffer` whether MPI can take device buffers.
+
+ :func:`mpi_is_cuda_aware` records the result of its probe itself; call
+ this when the answer is known otherwise (e.g. from the cluster
+ documentation, or to force host staging in a test).
- `mpi_is_cuda_aware()` records its result itself; call this instead when
- the answer is known otherwise (e.g. from the cluster documentation, or to
- force host staging for a test). ``None`` forgets the setting.
+ Parameters
+ ----------
+ value : bool or None
+ Whether MPI is CUDA-aware; None forgets the setting.
+
+ See Also
+ --------
+ get_mpi_cuda_aware : Read the setting.
+
+ Examples
+ --------
+ >>> xp.mpi.set_mpi_cuda_aware(False) # always stage device buffers
+ >>> xp.mpi.get_mpi_cuda_aware()
+ False
+ >>> xp.mpi.set_mpi_cuda_aware(None)
"""
global _MPI_CUDA_AWARE
_MPI_CUDA_AWARE = None if value is None else bool(value)
def get_mpi_cuda_aware() -> bool | None:
- """The recorded answer of `mpi_is_cuda_aware()`/`set_mpi_cuda_aware()`, or None."""
+ """Return whether MPI was recorded as CUDA-aware.
+
+ Returns
+ -------
+ bool or None
+ The answer recorded by :func:`mpi_is_cuda_aware` or
+ :func:`set_mpi_cuda_aware`; None if nothing was recorded.
+
+ Examples
+ --------
+ >>> xp.mpi.get_mpi_cuda_aware() is None
+ True
+ """
return _MPI_CUDA_AWARE
def _pinned_or_host_empty(shape: tuple[int, ...], dtype: Any) -> np.ndarray:
- """A host buffer for staging, pinned when CuPy can allocate pinned memory."""
+ """Return a host buffer for staging, pinned when CuPy can allocate pinned memory."""
try:
import cupy as cp
@@ -96,12 +144,37 @@ def _pinned_or_host_empty(shape: tuple[int, ...], dtype: Any) -> np.ndarray:
class MPIStaging:
- """Reusable host storage for blocking MPI exchanges of device arrays.
+ """Reusable host storage for :func:`mpi_buffer` staging of device arrays.
+
+ Make one per buffer that is in use at the same time (e.g. one for sending,
+ one for receiving). The host storage is allocated on first use, pinned
+ when available, and bound to one shape, dtype and CUDA device; only one
+ :func:`mpi_buffer` context may use it at a time. Non-C-contiguous device
+ arrays go through a reusable device packing buffer, and received values
+ are written back into the original view. Host arrays pass through
+ unchanged.
+
+ Parameters
+ ----------
+ shape : int or tuple of int
+ Shape of the staged arrays.
+ dtype : dtype-like
+ Their dtype (not object).
- Allocate once per send/receive buffer. Storage is allocated lazily, pinned
- when available, and bound to one shape, dtype, and CUDA device. Only one
- context may use it at a time. For nonblocking MPI, call request.Wait() inside
- the context: the buffer must not be released or copied back while MPI uses it.
+ Raises
+ ------
+ ValueError
+ For a negative shape.
+ TypeError
+ For an object dtype.
+
+ Examples
+ --------
+ >>> staging = xp.mpi.MPIStaging((3,), np.float64)
+ >>> rho = np.ones(3)
+ >>> with staging.buffer(rho, recv=True) as buf:
+ ... buf is rho # host arrays are not staged
+ True
"""
def __init__(self, shape: int | tuple[int, ...], dtype: Any) -> None:
@@ -131,11 +204,24 @@ def dtype(self) -> np.dtype:
return self._dtype
def buffer(self, array: Any, **kwargs: Any) -> Any:
- """Equivalent to ``mpi_buffer(array, staging=self, **kwargs)``."""
+ """Return ``mpi_buffer(array, staging=self, **kwargs)``.
+
+ Parameters
+ ----------
+ array : array
+ The buffer of the MPI call.
+ **kwargs
+ Further arguments of :func:`mpi_buffer`.
+
+ Returns
+ -------
+ context manager
+ Yields the buffer to pass to MPI.
+ """
return mpi_buffer(array, staging=self, **kwargs)
def _transfer_array(self, array: Any) -> Any:
- """Reuse C-order device storage when the MPI array has another layout."""
+ """Return `array`, or reusable C-order device storage for another layout."""
if array.flags.c_contiguous:
return array
if self._packed is None:
@@ -174,50 +260,61 @@ def mpi_buffer(
stream: Any = None,
event: Any = None,
) -> Generator[Any, None, None]:
- """The buffer to hand to MPI for `array`: the array itself, or a host copy.
-
- One MPI call site for both backends and both kinds of MPI builds::
-
- with xp.mpi.mpi_buffer(markers_out) as sendbuf, xp.mpi.mpi_buffer(
- markers_in, send=False, recv=True
- ) as recvbuf:
- comm.Sendrecv(sendbuf, dest, recvbuf=recvbuf, source=source)
-
- * A host array (NumPy backend, or a NumPy array on the CuPy backend) is
- yielded unchanged.
- * A device array with CUDA-aware MPI is yielded unchanged after
- `synchronize_for_mpi()`, so MPI reads what the kernels wrote.
- * A device array without CUDA-aware MPI is staged through a host buffer
- (pinned memory when available): with `send`, the array is copied to the
- host first (counted as a ``to_host`` transfer by `count_transfers()`);
- with `recv`, the host buffer is copied back into the array when the
- block ends (a ``to_device`` transfer). The device array itself is never
- given to MPI.
+ """Yield the buffer to hand to MPI for `array`: the array itself, or a host copy.
+
+ One MPI call site for both backends and both kinds of MPI builds. A host
+ array is yielded unchanged. A device array is yielded unchanged after
+ :func:`synchronize_for_mpi` when MPI is CUDA-aware; otherwise it is staged
+ through a host buffer (pinned when available), copied from the device
+ before the block (`send`) and back after it (`recv`), both counted by
+ :func:`~cunumpy.profiling.count_transfers`. The copy-back has finished
+ when the block ends. With nonblocking MPI, call ``request.Wait()`` inside
+ the block.
Parameters
----------
- array
+ array : array
The buffer of the MPI call: a NumPy or CuPy array.
- send : bool
+ send : bool, optional
Whether MPI reads the buffer (copy device to host before the block).
- recv : bool
+ recv : bool, optional
Whether MPI writes the buffer (copy host to device after the block).
- cuda_aware : bool | None
+ cuda_aware : bool, optional
Whether MPI can take device buffers. None uses the answer recorded by
- `mpi_is_cuda_aware()` or `set_mpi_cuda_aware()`.
- staging : MPIStaging | None
- Reusable storage for host staging. For nonblocking MPI, wait for the
- request inside the context before releasing this storage.
- stream, event
- Optional producer dependency (only one); without it all work on the
- array's device is synchronized before MPI accesses the buffer.
+ :func:`mpi_is_cuda_aware` or :func:`set_mpi_cuda_aware`.
+ staging : MPIStaging, optional
+ Reusable host storage; a new one is allocated for each call otherwise.
+ stream : cupy.cuda.Stream, optional
+ Producer stream to wait for; all work on the array's device by
+ default.
+ event : cupy.cuda.Event, optional
+ Producer event to wait for instead; pass at most one of `stream` and
+ `event`.
+
+ Yields
+ ------
+ array
+ The buffer to pass to MPI.
Raises
------
RuntimeError
For a device array when `cuda_aware` is None and nothing was recorded:
- call `mpi_is_cuda_aware(comm)` (collective) once at startup, or
- `set_mpi_cuda_aware()`.
+ call :func:`mpi_is_cuda_aware` (collective) once at startup, or
+ :func:`set_mpi_cuda_aware`.
+
+ See Also
+ --------
+ MPIStaging : Reusable staging storage.
+
+ Examples
+ --------
+ >>> MPI = xp.mpi.get_mpi() # the serial stand-in without an MPI launcher
+ >>> rho = np.ones(3)
+ >>> with xp.mpi.mpi_buffer(rho, recv=True) as buf:
+ ... MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM)
+ >>> rho
+ array([1., 1., 1.])
"""
if not array_api_compat.is_cupy_array(array):
yield array
@@ -275,11 +372,7 @@ def _device_buffers_in_use() -> bool:
def _mpi_probe_buffers(rank: int, n: int = 4) -> tuple[Any, Any]:
- """Device send and receive buffers for the CUDA-aware MPI probe.
-
- The send buffer of rank ``r`` holds ``r + arange(n)``, so the receiver can
- verify that the values came from the expected source.
- """
+ """Device send/receive buffers for the CUDA-aware probe; rank ``r`` sends ``r + arange(n)``."""
import cupy as cp
send = cp.arange(n, dtype=cp.float64) + rank
@@ -290,41 +383,49 @@ def _mpi_probe_buffers(rank: int, n: int = 4) -> tuple[Any, Any]:
def mpi_is_cuda_aware(comm: Any = None, *, method: str = "probe") -> bool:
"""Check whether the MPI library can send and receive device buffers.
- Passing CuPy arrays to MPI requires a CUDA-aware MPI build (e.g. Open MPI
+ Passing CuPy arrays to MPI needs a CUDA-aware MPI build (e.g. Open MPI
with ``--with-cuda``); with a plain build the call segfaults or silently
- sends garbage. This is a collective call: every rank of `comm` must call
- it, and all ranks receive the same result.
+ sends garbage. Collective: every rank of `comm` must call it, and all get
+ the same result, which is recorded for :func:`mpi_buffer`. Call it once
+ at startup, after :func:`~cunumpy.cuda.bind_local_device` and
+ ``MPI_Init``.
Parameters
----------
- comm
- The communicator to check; ``None`` means ``COMM_WORLD`` of
- :func:`~cunumpy.mpi.get_mpi` (mpi4py under an MPI launcher, else the
- serial stand-in; only on the CuPy backend).
- method
- Only ``"probe"`` is available: each rank sends a tiny device buffer to
- rank ``(rank + 1) % size`` and receives from ``(rank - 1) % size`` with
- ``Sendrecv`` (with one rank, to itself), checks the received values,
- and the ranks agree with ``allreduce(op=LAND)``. Any exception in the
- exchange, on any rank, gives ``False``. ``mpi4py`` does not expose the
- library's own query (``MPIX_Query_cuda_support``), and the library
- version string is not a reliable indicator, so no ``"query"`` method
- is offered.
+ comm : Comm, optional
+ The communicator; ``COMM_WORLD`` of ``cunumpy.mpi.get_mpi()`` by
+ default.
+ method : {"probe"}, optional
+ Each rank sends a tiny device buffer to rank ``(rank + 1) % size``
+ with ``Sendrecv``, checks what it received, and the ranks agree with
+ ``allreduce(op=LAND)``; any exception gives False. mpi4py does
+ not expose ``MPIX_Query_cuda_support``, so there is no other method.
Returns
-------
bool
- ``True`` if all ranks exchanged a device buffer successfully. ``False``
- on the NumPy backend and without a functional CuPy, without touching
- MPI: the question only makes sense with device buffers.
+ True if all ranks exchanged a device buffer correctly. False on the
+ NumPy backend and without a functional CuPy, without touching MPI.
+
+ Raises
+ ------
+ ValueError
+ For a `method` other than ``"probe"``.
+
+ See Also
+ --------
+ require_cuda_aware_mpi : Raise with instructions instead of returning False.
Notes
-----
- Not every failure is an exception: an MPI library that is not CUDA-aware
- may read the device address as a host address and crash the process. A
- segfault inside this call therefore also means that the MPI build is not
- CUDA-aware. Call it once at startup, after `bind_local_device()` and
- ``MPI_Init``, before any communication of device buffers.
+ A library that is not CUDA-aware may read the device address as a host
+ address and crash the process: a segfault inside this call means the same
+ as False. See :doc:`/guides/mpi`.
+
+ Examples
+ --------
+ >>> xp.mpi.mpi_is_cuda_aware() # NumPy backend: no device buffers
+ False
"""
if method != "probe":
raise ValueError(
@@ -360,17 +461,26 @@ def mpi_is_cuda_aware(comm: Any = None, *, method: str = "probe") -> bool:
def require_cuda_aware_mpi(comm: Any = None) -> None:
- """Raise if device buffers cannot be passed to MPI (no-op on NumPy).
+ """Raise if device buffers cannot be passed to MPI.
- Runs `mpi_is_cuda_aware()` and raises `RuntimeError` with instructions
- when it returns ``False`` on the CuPy backend. Collective: every rank of
- `comm` must call it. On the NumPy backend, or without a functional CuPy,
- it returns without touching MPI.
+ Runs :func:`mpi_is_cuda_aware` on the CuPy backend. Collective: every rank
+ of `comm` must call it. On the NumPy backend, or without a functional
+ CuPy, it returns without touching MPI.
Parameters
----------
- comm
- The communicator to check; ``None`` means ``get_mpi().COMM_WORLD``.
+ comm : Comm, optional
+ The communicator; ``COMM_WORLD`` of ``cunumpy.mpi.get_mpi()`` by
+ default.
+
+ Raises
+ ------
+ RuntimeError
+ If MPI is not CUDA-aware, explaining how to get a CUDA-aware build.
+
+ Examples
+ --------
+ >>> xp.mpi.require_cuda_aware_mpi() # no-op on the NumPy backend
"""
if not _device_buffers_in_use():
return
diff --git a/src/cunumpy/_philox.py b/src/cunumpy/_philox.py
index 583ada0..853c3e4 100644
--- a/src/cunumpy/_philox.py
+++ b/src/cunumpy/_philox.py
@@ -1,21 +1,14 @@
"""Counter-based random numbers (Philox4x32-10), the same on host and device.
-The host side of ``cunumpy/random.cuh``: :func:`philox_uniform` returns, for
-every (seed, stream, counter), exactly the doubles that ``cunumpy_uniform`` /
-``cunumpy_uniform2`` return in a kernel, so a kernel that samples random
-numbers can be compared with its host version element by element::
-
- ids = xp.arange(n, dtype=xp.uint64) # one stream per particle
- u0, u1 = xp.rng.philox_uniform2(seed, ids, step) # what each GPU thread draws
- z0, z1 = xp.rng.philox_normal2(seed, ids, step)
-
-The numbers are a pure function of the key (``seed``, 64 bit) and the 128-bit
-counter (``counter`` and ``stream``, 64 bit each): no state, independent of the
-launch shape. Arguments broadcast like NumPy arrays; the result is a NumPy or a
-CuPy array, matching the inputs (NumPy for Python ints). Uniform numbers are
-bit-identical to the device ones; normal numbers (Box-Muller with ``log``,
-``sqrt``, ``sin``, ``cos``) can differ in the last bits, since the device math
-functions are not the host's.
+The host side of ``cunumpy/random.cuh``: for every (seed, stream, counter) the
+``philox_*`` functions return the numbers a kernel draws, so a kernel can be
+compared with its host version element by element. The numbers are a pure
+function of the 64-bit key (``seed``) and the 128-bit counter (``counter``,
+``stream``): no state, independent of the launch shape. Arguments broadcast;
+the result is a NumPy or CuPy array matching the inputs (NumPy for Python
+ints). Uniform numbers are bit-identical to the device ones; normal numbers
+can differ in the last bits, since the device math functions are not the
+host's.
"""
from __future__ import annotations
@@ -49,7 +42,10 @@ def _module(*values: Any) -> Any:
def philox4x32_10(counter: Any, key0: Any, key1: Any) -> Any:
- """Philox4x32-10 of a ``(..., 4)`` array of uint32 counter words and a 2x32-bit key.
+ """Return Philox4x32-10 of uint32 counter words and a 2x32-bit key.
+
+ The raw generator, as ``cunumpy_philox4x32_10`` in ``cunumpy/random.cuh``
+ computes it; it passes the Random123 known-answer tests.
Parameters
----------
@@ -61,7 +57,16 @@ def philox4x32_10(counter: Any, key0: Any, key1: Any) -> Any:
Returns
-------
array of uint32, shape (..., 4)
- The random words, as ``cunumpy_philox4x32_10`` computes them.
+ The random words, on the backend of the inputs.
+
+ See Also
+ --------
+ philox_uniform2 : Uniform doubles from (seed, stream, counter).
+
+ Examples
+ --------
+ >>> xp.rng.philox4x32_10(np.zeros(4, dtype=np.uint32), 0, 0)
+ array([1713891541, 3781805453, 3159862348, 2600524760], dtype=uint32)
"""
xp = _module(counter, key0, key1)
ctr = xp.asarray(counter, dtype=xp.uint64)
@@ -105,21 +110,39 @@ def _to_uniform(xp: Any, hi: Any, lo: Any) -> Any:
def philox_uniform2(seed: Any, stream: Any, counter: Any) -> tuple[Any, Any]:
- """Two uniform doubles in [0, 1) per (seed, stream, counter), as ``cunumpy_uniform2``.
+ """Return two uniform doubles in [0, 1) per (seed, stream, counter).
+
+ Bit-identical to ``cunumpy_uniform2`` in a kernel (53 random bits each).
+ Use a different `counter` for every random decision of a step.
Parameters
----------
- seed : int or array
- The key (64 bit).
- stream : int or array
- The stream, e.g. a particle id (64 bit).
- counter : int or array
- The draw index, e.g. the time step (64 bit).
+ seed : int or array of uint64
+ The key.
+ stream : int or array of uint64
+ The stream, e.g. a particle id.
+ counter : int or array of uint64
+ The draw index, e.g. the time step.
Returns
-------
- tuple of arrays
- ``(u0, u1)``, broadcast to the shape of the arguments.
+ u0 : array of float64
+ The first number, broadcast to the shape of the arguments, on their
+ backend.
+ u1 : array of float64
+ The second number.
+
+ See Also
+ --------
+ philox_uniform : The first of the two numbers.
+ philox_normal2 : Normal numbers from the same draw.
+
+ Examples
+ --------
+ >>> ids = xp.arange(3, dtype=xp.uint64) # one stream per particle
+ >>> u0, u1 = xp.rng.philox_uniform2(42, ids, 0)
+ >>> u0
+ array([0.61295988, 0.01005884, 0.3984137 ])
"""
xp = _module(seed, stream, counter)
r = _random_bits(seed, stream, counter)
@@ -127,15 +150,62 @@ def philox_uniform2(seed: Any, stream: Any, counter: Any) -> tuple[Any, Any]:
def philox_uniform(seed: Any, stream: Any, counter: Any) -> Any:
- """One uniform double in [0, 1) per (seed, stream, counter), as ``cunumpy_uniform``."""
+ """Return one uniform double in [0, 1) per (seed, stream, counter).
+
+ Bit-identical to ``cunumpy_uniform`` in a kernel: the first number of
+ :func:`philox_uniform2`.
+
+ Parameters
+ ----------
+ seed : int or array of uint64
+ The key.
+ stream : int or array of uint64
+ The stream, e.g. a particle id.
+ counter : int or array of uint64
+ The draw index, e.g. the time step.
+
+ Returns
+ -------
+ array of float64
+ Broadcast to the shape of the arguments, on their backend.
+
+ Examples
+ --------
+ >>> xp.rng.philox_uniform(42, xp.arange(3, dtype=xp.uint64), 0)
+ array([0.61295988, 0.01005884, 0.3984137 ])
+ """
return philox_uniform2(seed, stream, counter)[0]
def philox_normal2(seed: Any, stream: Any, counter: Any) -> tuple[Any, Any]:
- """Two standard normal doubles per (seed, stream, counter), as ``cunumpy_normal2``.
+ """Return two standard normal doubles per (seed, stream, counter).
+
+ Box-Muller of :func:`philox_uniform2`, as ``cunumpy_normal2`` in a kernel;
+ equal to the device numbers up to the last bits of ``log``, ``sqrt``,
+ ``sin`` and ``cos``.
- Box-Muller of :func:`philox_uniform2`; equal to the device numbers up to the
- last bits of the math functions.
+ Parameters
+ ----------
+ seed : int or array of uint64
+ The key.
+ stream : int or array of uint64
+ The stream, e.g. a particle id.
+ counter : int or array of uint64
+ The draw index, e.g. the time step.
+
+ Returns
+ -------
+ z0 : array of float64
+ The first number, broadcast to the shape of the arguments, on their
+ backend.
+ z1 : array of float64
+ The second number.
+
+ Examples
+ --------
+ >>> z0, z1 = xp.rng.philox_normal2(42, xp.arange(1000, dtype=xp.uint64), 0)
+ >>> z0.shape, z0.dtype
+ ((1000,), dtype('float64'))
"""
xp = _module(seed, stream, counter)
u0, u1 = philox_uniform2(seed, stream, counter)
@@ -145,5 +215,28 @@ def philox_normal2(seed: Any, stream: Any, counter: Any) -> tuple[Any, Any]:
def philox_normal(seed: Any, stream: Any, counter: Any) -> Any:
- """One standard normal double per (seed, stream, counter), as ``cunumpy_normal``."""
+ """Return one standard normal double per (seed, stream, counter).
+
+ As ``cunumpy_normal`` in a kernel: the first number of
+ :func:`philox_normal2`.
+
+ Parameters
+ ----------
+ seed : int or array of uint64
+ The key.
+ stream : int or array of uint64
+ The stream, e.g. a particle id.
+ counter : int or array of uint64
+ The draw index, e.g. the time step.
+
+ Returns
+ -------
+ array of float64
+ Broadcast to the shape of the arguments, on their backend.
+
+ Examples
+ --------
+ >>> xp.rng.philox_normal(42, xp.arange(4, dtype=xp.uint64), 0).shape
+ (4,)
+ """
return philox_normal2(seed, stream, counter)[0]
diff --git a/src/cunumpy/_profiling.py b/src/cunumpy/_profiling.py
index 586109a..c3ae902 100644
--- a/src/cunumpy/_profiling.py
+++ b/src/cunumpy/_profiling.py
@@ -13,11 +13,7 @@
def _nvtx_module() -> ModuleType | None:
- """Return ``cupy.cuda.nvtx`` on the CuPy backend, else None.
-
- None is also returned when NVTX is not available in this CuPy build, so
- callers can degrade to a no-op instead of failing.
- """
+ """Return ``cupy.cuda.nvtx`` on the CuPy backend with NVTX, else None."""
if array_backend.backend != "cupy":
return None
try:
@@ -29,30 +25,34 @@ def _nvtx_module() -> ModuleType | None:
class nvtx_range(ContextDecorator):
"""Mark a code region as an NVTX range, visible in ``nsys`` and Nsight.
- On the CuPy backend the block is wrapped in ``cupy.cuda.nvtx.RangePush``
- / ``RangePop``, so the region appears on the profiler timeline next to
- the kernels it launches. On the NumPy backend, or if NVTX is not
- available, it is a no-op. The range is popped when the block raises.
-
- It can also be used as a decorator, and the same instance can be nested
- or re-entered (e.g. on a recursive function).
+ A context manager and decorator. On the CuPy backend it calls
+ ``cupy.cuda.nvtx.RangePush`` / ``RangePop`` (also when the block raises),
+ so the region shows on the timeline above the kernels it launches; it
+ does not synchronize. On NumPy, or without NVTX, it does nothing. One
+ instance may be nested or re-entered (e.g. on a recursive function). See
+ :doc:`/guides/profiling`.
Parameters
----------
name : str
Name shown in the profiler.
color : int, optional
- Index into NVTX's colour table (``id_color`` of ``RangePush``).
- ``None`` uses the default colour.
+ Index into NVTX's colour table (``id_color`` of ``RangePush``); None
+ uses the default colour.
+
+ See Also
+ --------
+ timed_region : Time a region (also pushes an NVTX range).
Examples
--------
>>> with xp.profiling.nvtx_range("push markers"):
- ... kernel(markers, dt, n_threads=n)
-
- >>> @xp.profiling.nvtx_range("accumulate")
- ... def accumulate(...):
- ... ...
+ ... x = xp.ones(3) * 2.0
+ >>> @xp.profiling.nvtx_range("step")
+ ... def step(x):
+ ... return 2.0 * x
+ >>> step(xp.ones(2))
+ array([2., 2.])
"""
def __init__(self, name: str, color: int | None = None) -> None:
@@ -81,18 +81,24 @@ def __repr__(self) -> str:
@dataclass
class Timing:
- """Result of a `timed_region()` block.
+ """The result of a :func:`timed_region` block.
Attributes
----------
name : str
Name of the region.
- elapsed : float | None
- Wall-clock seconds spent in the block; ``None`` until the block
- exits.
+ elapsed : float or None
+ Wall-clock seconds spent in the block; None until the block exits.
synced : bool
Whether the device was synchronized before the clock was read, i.e.
whether `elapsed` includes the queued device work.
+
+ Examples
+ --------
+ >>> with xp.profiling.timed_region("solve") as timing:
+ ... pass
+ >>> timing.name, timing.synced
+ ('solve', False)
"""
name: str
@@ -104,34 +110,33 @@ class Timing:
def timed_region(name: str, *, sync: bool = True) -> Generator[Timing, None, None]:
"""Time a code region, including the device work it queues.
- CUDA kernels run asynchronously: a wall-clock timer around a launch
- measures the launch, not the kernel. On the CuPy backend this context
- manager synchronizes the device on entry (so earlier queued work is not
- charged to the region) and, if `sync` is true, again on exit before the
- clock is read, so the measured time includes the kernels launched in
- the block. It also pushes an `nvtx_range()` of the same name, so the
- region shows in ``nsys``. On the NumPy backend it only times the block.
- The time is recorded when the block raises as well.
+ CUDA kernels run asynchronously, so a plain timer measures only the
+ launch. On the CuPy backend the device is synchronized on entry (earlier
+ work is not charged to the region) and, if `sync` is true, on exit before
+ the clock is read. The region is also an :class:`nvtx_range` of the same
+ name. On NumPy it is a plain timer. The time is recorded when the block
+ raises as well. See :doc:`/guides/profiling`.
Parameters
----------
name : str
Name of the region (also the NVTX range name).
- sync : bool, default True
- Synchronize the device before reading the clock on exit. With
- ``False`` the time is the host time only, as for a plain timer.
+ sync : bool, optional
+ Synchronize the device on exit before reading the clock (default).
+ With False only the host time (the launch overhead) is measured.
Yields
------
Timing
- `elapsed` is set (in seconds, from ``time.perf_counter``) when the
- block exits; `synced` tells whether the device was synchronized.
+ `elapsed` (seconds, from ``time.perf_counter``) is set when the block
+ exits; `synced` tells whether the device was synchronized.
Examples
--------
- >>> with xp.profiling.timed_region("push markers") as timing:
- ... kernel(markers, dt, n_threads=n)
- >>> print(f"{timing.name}: {timing.elapsed:.3f} s (synced={timing.synced})")
+ >>> with xp.profiling.timed_region("solve") as timing:
+ ... x = xp.linalg.solve(xp.eye(3), xp.ones(3))
+ >>> timing.elapsed >= 0.0, timing.synced
+ (True, False)
"""
timing = Timing(name=str(name))
on_gpu = array_backend.backend == "cupy"
diff --git a/src/cunumpy/_random_streams.py b/src/cunumpy/_random_streams.py
index 002ee86..f3a3e1b 100644
--- a/src/cunumpy/_random_streams.py
+++ b/src/cunumpy/_random_streams.py
@@ -1,31 +1,11 @@
"""Reproducible random numbers for MPI programs on either backend.
-A simulation that draws random numbers in many places (initial loading,
-injection, collisions) is reproducible only if every draw comes from a seeded
-generator, and an MPI run needs a different stream on every rank.
-:data:`random_streams` is one generator per process and backend for that::
-
- import cunumpy as xp
-
- xp.rng.random_streams.seed(42, rank=comm.Get_rank()) # once, at start-up
- v = xp.rng.random_streams.normal(0.0, v_th, (n, 3)) # anywhere afterwards
- rng = xp.rng.random_streams.generator() # the Generator itself
-
-Each rank draws the stream ``(seed, rank)`` (a NumPy ``SeedSequence`` with the
-rank as spawn key), so a run with the same seed and the same number of ranks
-reproduces its results exactly, and the ranks' streams are independent. The
-generator of each backend is created on first use from the same stream: a
-``numpy.random.Generator`` with the chosen bit generator, or a
-``cupy.random.Generator`` (CuPy's default bit generator).
-
-Components with a seed of their own (e.g. a source configured with a ``seed``)
-get a separate generator with :meth:`RandomStreams.make_generator`; without one
-they share the process generator. Without :meth:`~RandomStreams.seed`, or with
-``seed(None)``, the stream is seeded from the operating system.
-
-The draw functions work with NumPy and CuPy generators alike; CuPy's
-``Generator`` lacks some of NumPy's methods, and the missing ones are derived
-from ``random`` and ``standard_normal``.
+:data:`random_streams` is one seeded generator per process and backend. Each
+rank draws the stream ``(seed, rank)`` (a NumPy ``SeedSequence`` with the rank
+as spawn key), so the same seed and number of ranks reproduce a run exactly
+and the ranks' streams are independent. The generator of each backend is
+created on first use from that stream. See :doc:`/guides/mpi` for the use in
+MPI programs.
"""
from __future__ import annotations
@@ -45,7 +25,26 @@
class RandomStreams:
- """One seeded random generator per process and backend; see :mod:`cunumpy.rng`."""
+ """One seeded random generator per process and backend.
+
+ Without :meth:`seed`, or with ``seed(None)``, the stream is seeded from the
+ operating system. The draw methods work with NumPy and CuPy generators
+ alike; :data:`random_streams` is the process-wide instance, and a new
+ instance is independent of it.
+
+ See Also
+ --------
+ get_rng : A fresh, unshared generator for the active backend.
+
+ Examples
+ --------
+ >>> streams = xp.rng.RandomStreams()
+ >>> streams.seed(42, rank=1)
+ >>> streams
+ RandomStreams(entropy=42, spawn_key=(1,), bit_generator='PCG64')
+ >>> streams.uniform(0.0, 1.0, 3)
+ array([0.57002868, 0.9724905 , 0.50377011])
+ """
def __init__(self) -> None:
self._sequence: np.random.SeedSequence | None = None
@@ -69,23 +68,27 @@ def seed(
"""Seed all draws of this process with the stream ``(value, rank)``.
Replaces the generators of all backends. NumPy's global random state
- (and, on the CuPy backend, CuPy's) is seeded from the same stream, for
- code that still calls ``np.random.*`` or ``xp.random.*`` directly.
+ (and CuPy's on the CuPy backend) is seeded from the same stream, for
+ code that calls ``np.random.*`` directly.
Parameters
----------
- value : int | None
+ value : int or None
The seed; None seeds from the operating system.
- rank : int
+ rank : int, optional
MPI rank of this process, so that ranks draw independent streams.
- bit_generator : str | None
+ bit_generator : str, optional
The NumPy bit generator, one of :data:`BIT_GENERATORS`; ``PCG64``
- (NumPy's default) if None. CuPy always uses its own default.
+ by default. CuPy always uses its own default.
Raises
------
ValueError
For an unknown `bit_generator`.
+
+ Examples
+ --------
+ >>> xp.rng.random_streams.seed(42, rank=comm.Get_rank()) # doctest: +SKIP
"""
if bit_generator is not None and bit_generator not in BIT_GENERATORS:
raise ValueError(
@@ -114,7 +117,23 @@ def _backend_sequence(self, backend: str) -> np.random.SeedSequence:
)
def generator(self, backend: str | None = None) -> Any:
- """The process generator of `backend` (the active backend by default)."""
+ """Return the process generator of a backend, created on first use.
+
+ Parameters
+ ----------
+ backend : {"numpy", "cupy"}, optional
+ The backend; the active one by default.
+
+ Returns
+ -------
+ numpy.random.Generator or cupy.random.Generator
+ The shared generator of this process and backend.
+
+ Raises
+ ------
+ ValueError
+ For a `backend` other than ``"numpy"`` or ``"cupy"``.
+ """
backend = get_backend() if backend is None else backend
if backend not in _BACKEND_KEYS:
raise ValueError(f"backend must be 'numpy' or 'cupy', got {backend!r}")
@@ -135,15 +154,20 @@ def make_generator(
seed: int | None = None,
backend: str | None = None,
) -> Any:
- """A generator for a component: its own if it has a seed, else the process one.
+ """Return a generator for a component: its own if it has a seed, else the shared one.
Parameters
----------
- seed : int | None
- The component's own seed; None shares :meth:`generator`.
- backend : str | None
+ seed : int, optional
+ The component's own seed; None returns :meth:`generator`.
+ backend : {"numpy", "cupy"}, optional
``"numpy"`` for a component that draws on the host whatever the
active backend; the active backend by default.
+
+ Returns
+ -------
+ numpy.random.Generator or cupy.random.Generator
+ A new generator seeded with `seed`, or the process generator.
"""
if seed is None:
return self.generator(backend)
@@ -158,7 +182,20 @@ def _rng(self, rng: Any) -> Any:
return self.generator() if rng is None else rng
def random(self, size: int | tuple[int, ...] | None = None, rng: Any = None) -> Any:
- """Uniform samples in [0, 1) from `rng` (the process generator by default)."""
+ """Draw uniform samples in [0, 1).
+
+ Parameters
+ ----------
+ size : int or tuple of int, optional
+ Output shape; None gives a scalar.
+ rng : Generator, optional
+ The generator to draw from; the process generator by default.
+
+ Returns
+ -------
+ array or float
+ On the backend of the generator.
+ """
return self._rng(rng).random(size=size)
def standard_normal(
@@ -166,7 +203,20 @@ def standard_normal(
size: int | tuple[int, ...] | None = None,
rng: Any = None,
) -> Any:
- """Standard normal samples from `rng` (the process generator by default)."""
+ """Draw standard normal samples.
+
+ Parameters
+ ----------
+ size : int or tuple of int, optional
+ Output shape; None gives a scalar.
+ rng : Generator, optional
+ The generator to draw from; the process generator by default.
+
+ Returns
+ -------
+ array or float
+ On the backend of the generator.
+ """
return self._rng(rng).standard_normal(size=size)
def normal(
@@ -176,7 +226,24 @@ def normal(
size: int | tuple[int, ...] | None = None,
rng: Any = None,
) -> Any:
- """Normal samples from `rng` (the process generator by default)."""
+ """Draw normal samples, also from CuPy generators without ``normal``.
+
+ Parameters
+ ----------
+ loc : float or array, optional
+ Mean.
+ scale : float or array, optional
+ Standard deviation.
+ size : int or tuple of int, optional
+ Output shape; None gives a scalar.
+ rng : Generator, optional
+ The generator to draw from; the process generator by default.
+
+ Returns
+ -------
+ array or float
+ On the backend of the generator.
+ """
rng = self._rng(rng)
if hasattr(rng, "normal"):
return rng.normal(loc=loc, scale=scale, size=size)
@@ -189,7 +256,22 @@ def uniform(
size: int | tuple[int, ...] | None = None,
rng: Any = None,
) -> Any:
- """Uniform samples in [low, high) from `rng` (the process generator by default)."""
+ """Draw uniform samples in [low, high), also from CuPy generators without ``uniform``.
+
+ Parameters
+ ----------
+ low, high : float or array, optional
+ The bounds, 0 and 1 by default.
+ size : int or tuple of int, optional
+ Output shape; None gives a scalar.
+ rng : Generator, optional
+ The generator to draw from; the process generator by default.
+
+ Returns
+ -------
+ array or float
+ On the backend of the generator.
+ """
rng = self._rng(rng)
if hasattr(rng, "uniform"):
return rng.uniform(low=low, high=high, size=size)
@@ -201,11 +283,25 @@ def uniform(
def get_rng(seed: int | None = None) -> Any:
- """Return a random Generator matching the active backend.
+ """Return a new random generator for the active backend.
+
+ ``default_rng(seed)`` of NumPy or CuPy. The APIs are similar, but the same
+ seed does not give the same numbers on both backends.
+
+ Parameters
+ ----------
+ seed : int, optional
+ The seed; None seeds from the operating system.
+
+ Returns
+ -------
+ numpy.random.Generator or cupy.random.Generator
+ A generator independent of :data:`random_streams`.
- NumPy and CuPy both provide `default_rng(seed)`, returning a
- `Generator` with a largely-compatible distribution API, but picking the
- right one requires branching on the backend -- this does that for you.
+ Examples
+ --------
+ >>> xp.rng.get_rng(seed=7).random(2)
+ array([0.62509547, 0.8972138 ])
"""
if get_backend() == "cupy":
import cupy as cp
diff --git a/src/cunumpy/_scipy_backend.py b/src/cunumpy/_scipy_backend.py
index 64a0826..3e00026 100644
--- a/src/cunumpy/_scipy_backend.py
+++ b/src/cunumpy/_scipy_backend.py
@@ -1,25 +1,19 @@
"""SciPy for the active backend: ``xp.scipy`` is SciPy or ``cupyx.scipy``.
-Field solvers and fluid codes need more than array functions: sparse matrices
-and their iterative solvers, FFTs, special functions, image filters.
``xp.scipy`` forwards to :mod:`scipy` on the NumPy backend and to
-:mod:`cupyx.scipy` on the CuPy backend, so this code runs on both::
+:mod:`cupyx.scipy` on the CuPy backend (sparse matrices and solvers, FFTs,
+special functions, ...), so this code runs on both::
A = xp.scipy.sparse.csr_matrix((data, (rows, cols)), shape=(n, n))
x, info = xp.scipy.sparse.linalg.cg(A, b)
- phi_k = xp.scipy.fft.rfftn(rho)
f = xp.scipy.special.erf(v / v_th)
-The module is resolved at every attribute access, so switching the backend
-(:func:`cunumpy.set_backend`) takes effect immediately; imports are cached by
-Python, so the lookup is cheap. Neither SciPy nor CuPy is imported until a
-name is used.
-
-``cupyx.scipy`` covers only part of SciPy. A name that the active backend's
-module does not have raises ``AttributeError`` saying which backend lacks it;
-:meth:`ScipyNamespace.available` checks a name without raising. Keyword
-arguments can differ as well: SciPy's ``cg`` takes ``rtol`` since SciPy 1.12,
-CuPy's still takes ``tol``.
+Names are resolved at every access, so a backend switch takes effect at once;
+nothing is imported until a name is used, and SciPy is not a dependency of
+cunumpy. ``cupyx.scipy`` covers only part of SciPy, and keyword arguments can
+differ (SciPy's ``cg`` takes ``rtol``, CuPy's ``tol``). A SciPy sparse matrix
+moves to the device once with ``xp.scipy.sparse.csr_matrix(host_matrix)`` on
+the CuPy backend, and back with ``matrix.get()``. See :doc:`/guides/solvers`.
"""
from __future__ import annotations
@@ -57,13 +51,28 @@
class ScipyNamespace:
- """SciPy (sub)package of the active backend; see :data:`cunumpy.scipy`.
+ """A SciPy (sub)package of the active backend; see :data:`cunumpy.scipy`.
+
+ Attribute access forwards to the module of the active backend. A name the
+ backend's module lacks raises ``AttributeError`` naming the backend; a
+ missing SciPy (NumPy backend) or CuPy raises ``ImportError``.
Parameters
----------
- path : str
- Dotted path below the SciPy root, e.g. ``"sparse.linalg"``; empty for
- the root itself.
+ path : str, optional
+ Dotted path below the SciPy root, one of ``SUBMODULES``, e.g.
+ ``"sparse.linalg"``; empty for the root itself.
+
+ Raises
+ ------
+ ValueError
+ If `path` is not a forwarded subpackage.
+
+ Examples
+ --------
+ >>> xp.scipy.sparse
+ .sparse>
+ >>> x, info = xp.scipy.sparse.linalg.cg(A, b) # doctest: +SKIP
"""
def __init__(self, path: str = "") -> None:
@@ -82,7 +91,7 @@ def _name(self, root: str) -> str:
return f"{root}.{self._path}" if self._path else root
def resolve(self) -> ModuleType:
- """The module this namespace stands for on the active backend.
+ """Return the module this namespace stands for on the active backend.
Returns
-------
@@ -123,9 +132,22 @@ def __getattr__(self, name: str) -> Any:
) from None
def available(self, name: str) -> bool:
- """Whether `name` exists in this namespace on the active backend.
+ """Return whether `name` exists in this namespace on the active backend.
+
+ Parameters
+ ----------
+ name : str
+ A function, class or subpackage name.
+
+ Returns
+ -------
+ bool
+ False also if the backend's SciPy is not installed; never raises.
- Never raises; False also if the backend's SciPy is not installed.
+ Examples
+ --------
+ >>> xp.scipy.special.available("erfcx") # doctest: +SKIP
+ True
"""
child = f"{self._path}.{name}" if self._path else name
if child in SUBMODULES:
diff --git a/src/cunumpy/_staging.py b/src/cunumpy/_staging.py
index aa2c277..6fdfa36 100644
--- a/src/cunumpy/_staging.py
+++ b/src/cunumpy/_staging.py
@@ -1,30 +1,10 @@
"""Copy device arrays to the host in the background, for output that should not stall the GPU.
-Writing a field or the markers to HDF5 every few steps needs host arrays, and
``array.get()`` waits for the GPU and then for the copy, while no kernel runs.
-:class:`HostStaging` overlaps the copy with the next time steps::
-
- staging = xp.memory.HostStaging(rho.shape, rho.dtype) # once
- for step in range(n_steps):
- advance(...)
- if step % output_every == 0:
- pending.append((step, staging.copy(rho))) # returns at once
- while pending and pending[0][1].ready():
- s, copy = pending.pop(0)
- h5file[f"rho/{s}"] = copy.result() # a NumPy array
-
-Each :meth:`HostStaging.copy` first snapshots the array on the device (on the
-current stream, after the kernels that wrote it, so the next steps may
-overwrite the array), then copies the snapshot to a page-locked host buffer on
-its own stream. The buffers are reused in turn: with ``buffers=2``, a third
-copy waits until the first one has finished, so the program never runs more
-than ``buffers`` copies ahead of the output. :meth:`StagedCopy.result` waits for
-its copy and returns the host buffer, valid until that buffer is reused
-(``buffers`` copies later); a stale result raises instead of returning another
-step's data. Copy it (``result().copy()``) to keep it longer.
-
-Host arrays (and the NumPy backend) are copied at once, into the same buffers,
-so the code is the same on both backends.
+:class:`~cunumpy.memory.HostStaging` (for large arrays copied repeatedly) and
+:func:`~cunumpy.to_host_async` (for scalars and small arrays) overlap the copy
+with the next kernels. Host arrays (and the NumPy backend) are copied at once,
+so the code is the same on both backends. See :doc:`/guides/particle-codes`.
"""
from __future__ import annotations
@@ -74,7 +54,18 @@ def context(self) -> Any:
class StagedCopy:
- """A copy started by :meth:`HostStaging.copy`."""
+ """A background copy started by :meth:`HostStaging.copy`.
+
+ Not created directly. Its host buffer is valid until the staging reuses
+ it, `buffers` copies later.
+
+ Parameters
+ ----------
+ slot : object
+ The staging buffer (internal).
+ generation : int
+ The buffer's use count when this copy started.
+ """
def __init__(self, slot: _Slot, generation: int) -> None:
self._slot = slot
@@ -89,7 +80,7 @@ def _check(self) -> None:
)
def ready(self) -> bool:
- """Whether the copy has finished (never waits)."""
+ """Whether the copy has finished (never waits); raises if the buffer was reused."""
self._check()
if self._slot.event is None:
return True
@@ -99,7 +90,11 @@ def ready(self) -> bool:
def result(self) -> np.ndarray:
"""Wait for the copy and return the host array.
- The array is the staging buffer itself: valid until the buffer is reused.
+ Returns
+ -------
+ numpy.ndarray
+ The staging buffer itself, valid until it is reused; keep
+ ``result().copy()`` to hold it longer.
Raises
------
@@ -116,22 +111,50 @@ def result(self) -> np.ndarray:
class HostStaging:
"""Page-locked host buffers that device arrays are copied to in the background.
+ :meth:`copy` snapshots the array on the device, so later steps may
+ overwrite it, then copies the snapshot to the next of `buffers` pinned host
+ buffers on its own stream. It waits only if that buffer's previous copy is
+ still running, so the program runs at most `buffers` copies ahead. Host
+ arrays are copied at once into the same buffers. Device copies are counted
+ by :func:`~cunumpy.profiling.count_transfers`.
+
Parameters
----------
- shape : tuple[int, ...]
+ shape : int or tuple of int
Shape of the arrays to copy.
dtype : dtype-like
- Their dtype.
- buffers : int
- Number of buffers (and device snapshots) used in turn: how many copies
- may be in flight at once. 2 (double buffering) lets one copy run while
- the previous result is written out.
+ Their dtype (not object).
+ buffers : int, optional
+ Number of buffers (and device snapshots) used in turn: how many
+ copies may be in flight. The default 2 lets one copy run while the
+ previous result is written out.
+
+ Raises
+ ------
+ ValueError
+ If `buffers` < 1 or `shape` is negative.
+ TypeError
+ For an object dtype.
+
+ See Also
+ --------
+ cunumpy.to_host_async : One-off copy of a scalar or small array.
Notes
-----
- GPU storage is bound to the first source device; make it current when
- submitting copies. A CPU-initialized instance allocates new pinned slots
- on its first GPU copy; earlier CPU results retain their completed snapshots.
+ GPU storage is bound to the first source device; make that device current
+ when submitting copies. The first GPU copy of an instance used on the host
+ allocates new pinned buffers; earlier host results keep theirs.
+
+ Examples
+ --------
+ >>> staging = xp.memory.HostStaging((2, 3), np.float64)
+ >>> copy = staging.copy(np.ones((2, 3))) # returns at once on the GPU
+ >>> copy.ready()
+ True
+ >>> copy.result()
+ array([[1., 1., 1.],
+ [1., 1., 1.]])
"""
def __init__(
@@ -199,22 +222,28 @@ def copy(self, array: Any, *, stream: Any = None, event: Any = None) -> StagedCo
Parameters
----------
- array : cupy.ndarray | numpy.ndarray
- An array of the staging shape and dtype. Later writes on the
- snapshot's stream are ordered after it; another stream must wait
- for the snapshot/copy before overwriting the source.
- stream : cupy.cuda.Stream | None
- Producer stream on which to snapshot, current stream by default.
- event : cupy.cuda.Event | None
- Producer completion to wait for on the current stream before the
- snapshot. Pass at most one of stream and event. Host copies ignore
- these dependencies because host input is assumed ready.
+ array : array
+ A NumPy or CuPy array of the staging shape and dtype. Later writes
+ on the snapshot's stream are ordered after the snapshot; another
+ stream must wait for it before overwriting the array.
+ stream : cupy.cuda.Stream, optional
+ Producer stream to snapshot on; the current stream by default.
+ event : cupy.cuda.Event, optional
+ Producer event the current stream waits for before the snapshot.
+ Pass at most one of `stream` and `event`; host copies ignore both.
Returns
-------
StagedCopy
``ready()`` tells whether the copy has finished, ``result()``
waits for it and returns the host array.
+
+ Raises
+ ------
+ ValueError
+ For another shape or dtype, both `stream` and `event`, or a
+ device other than the current one or the one the staging is
+ bound to.
"""
if stream is not None and event is not None:
raise ValueError("pass only one of stream and event")
@@ -294,7 +323,19 @@ def synchronize(self) -> None:
class HostCopy:
- """A device-to-host copy started by :func:`~cunumpy.to_host_async`."""
+ """A device-to-host copy started by :func:`~cunumpy.to_host_async`.
+
+ Not created directly.
+
+ Parameters
+ ----------
+ host : numpy.ndarray
+ The host buffer the copy writes into.
+ event : cupy.cuda.Event, optional
+ Recorded after the copy; None for a finished copy.
+ source : array, optional
+ The device array, kept alive until the copy is done.
+ """
def __init__(self, host: np.ndarray, event: Any = None, source: Any = None) -> None:
self._host = host
@@ -306,11 +347,15 @@ def ready(self) -> bool:
return self._event is None or bool(self._event.done)
def result(self) -> Any:
- """The value on the host; waits for the copy only if it has not finished.
+ """Return the value on the host, waiting only if the copy has not finished.
- A NumPy scalar for a 0-d array, else a NumPy array (in page-locked memory
- on the GPU). A wait is counted as a ``sync`` by
+ A wait is counted as a ``sync`` by
:func:`~cunumpy.profiling.count_transfers`.
+
+ Returns
+ -------
+ numpy scalar or numpy.ndarray
+ A scalar for a 0-d array, else an array (page-locked on the GPU).
"""
if self._event is not None:
if not self._event.done:
@@ -326,26 +371,37 @@ def result(self) -> Any:
def to_host_async(array: Any) -> HostCopy:
- """Start copying `array` (a device scalar or small array) to the host; do not wait.
+ """Start copying a device scalar or small array to the host without waiting.
The copy runs on a separate stream into page-locked memory, after the work
- queued so far on the current stream, so the host can queue more kernels
- while it is in flight. :meth:`HostCopy.ready` tells whether it has
- finished (never waits), :meth:`HostCopy.result` returns the value::
-
- pending = xp.to_host_async(residual_norm)
- ... # queue the next iteration
- if pending.ready() and pending.result() < tol:
- break
-
- The value is that of the array when the queued work is done: later kernels
- may overwrite the array. For large arrays copied repeatedly, use
- :class:`~cunumpy.memory.HostStaging`, which reuses its buffers.
-
- On the NumPy backend (a host array) the value is copied at once and
- ``ready()`` is always True. On the fake CuPy too, and the copy is counted.
- :func:`~cunumpy.profiling.count_transfers` records a ``to_host`` event with
- ``blocking=False``.
+ queued so far on the current stream, so the host can keep queueing kernels
+ (e.g. a convergence test read one iteration late). The value is that of
+ the array when the queued work is done, even if later kernels overwrite
+ it. A host array is copied at once. Counted as a ``to_host`` transfer with
+ ``blocking=False`` by :func:`~cunumpy.profiling.count_transfers`.
+
+ Parameters
+ ----------
+ array : array
+ A NumPy or CuPy array or scalar.
+
+ Returns
+ -------
+ HostCopy
+ ``ready()`` tells whether the copy has finished (never waits),
+ ``result()`` returns the value.
+
+ See Also
+ --------
+ cunumpy.memory.HostStaging : For large arrays copied repeatedly.
+
+ Examples
+ --------
+ >>> pending = xp.to_host_async(np.float64(3.5)) # e.g. a residual norm
+ >>> pending.ready()
+ True
+ >>> pending.result()
+ np.float64(3.5)
"""
if not _is_device_array(array):
return HostCopy(np.array(array))
diff --git a/src/cunumpy/_streams.py b/src/cunumpy/_streams.py
index 706214f..ab6e245 100644
--- a/src/cunumpy/_streams.py
+++ b/src/cunumpy/_streams.py
@@ -1,4 +1,4 @@
-"""Reusable CUDA streams/events with synchronous CPU equivalents."""
+"""Reusable CUDA streams and events, with synchronous CPU equivalents."""
from __future__ import annotations
@@ -8,22 +8,56 @@
class HostEvent:
- """Already-completed CPU event, matching CuPy's completion interface."""
+ """Already-completed CPU event with the completion interface of ``cupy.cuda.Event``.
+
+ Returned by :func:`create_event` on the NumPy backend.
+
+ Examples
+ --------
+ >>> event = xp.cuda.create_event() # a HostEvent on NumPy
+ >>> event.done
+ True
+ """
@property
def done(self) -> bool:
+ """Whether the event has completed; always True."""
return True
def record(self, stream: Any = None) -> None:
+ """Record the event on a host stream; does nothing.
+
+ Parameters
+ ----------
+ stream : HostStream, optional
+ The stream to record on.
+
+ Raises
+ ------
+ TypeError
+ If `stream` is not a :class:`HostStream`.
+ """
if stream is not None and not isinstance(stream, HostStream):
raise TypeError("a host event requires a host stream")
def synchronize(self) -> None:
- pass
+ """Wait for the event; returns at once."""
class HostStream:
- """Synchronous CPU stream; reusable and nestable as a context manager."""
+ """Synchronous CPU stream with the interface of ``cupy.cuda.Stream``.
+
+ Returned by :func:`create_stream` on the NumPy backend; reusable and nestable
+ as a context manager.
+
+ Examples
+ --------
+ >>> s = xp.cuda.create_stream() # a HostStream on NumPy
+ >>> with s:
+ ... y = xp.ones(3)
+ >>> s.record().done
+ True
+ """
def __enter__(self) -> Any:
return self
@@ -33,26 +67,69 @@ def __exit__(self, *exc: object) -> None:
@property
def done(self) -> bool:
+ """Whether all work on the stream has completed; always True."""
return True
def synchronize(self) -> None:
- pass
+ """Wait for the work on the stream; returns at once."""
def record(self, event: Any = None) -> HostEvent:
+ """Record an event on this stream.
+
+ Parameters
+ ----------
+ event : HostEvent, optional
+ The event to record; a new one by default.
+
+ Returns
+ -------
+ HostEvent
+ The recorded event.
+ """
event = HostEvent() if event is None else event
event.record(self)
return event
def wait_event(self, event: Any) -> None:
+ """Order later work after `event`; does nothing.
+
+ Parameters
+ ----------
+ event : HostEvent
+ The event to wait for.
+
+ Raises
+ ------
+ TypeError
+ If `event` is not a :class:`HostEvent`.
+ """
if not isinstance(event, HostEvent):
raise TypeError("a host stream requires a host event")
def create_stream(*, non_blocking: bool = True) -> Any:
- """Allocate a reusable CuPy stream, or a synchronous :class:`HostStream`.
-
- CUDA streams belong to the device current at construction. Keep the same
- device current while selecting the stream and submitting work to it.
+ """Allocate a reusable CUDA stream, or a :class:`HostStream` on NumPy.
+
+ A CUDA stream belongs to the device current at construction; keep that device
+ current while selecting the stream and submitting work to it.
+
+ Parameters
+ ----------
+ non_blocking : bool, optional
+ Create a stream that does not synchronize with the default stream; True
+ by default.
+
+ Returns
+ -------
+ cupy.cuda.Stream or HostStream
+ The stream, usable with :func:`stream`.
+
+ Examples
+ --------
+ >>> s = xp.cuda.create_stream()
+ >>> with xp.cuda.stream(s):
+ ... y = xp.ones(3) * 2
+ >>> s.synchronize()
"""
if get_backend() == "numpy":
return HostStream()
@@ -62,10 +139,24 @@ def create_stream(*, non_blocking: bool = True) -> Any:
def create_event(*, timing: bool = False) -> Any:
- """Allocate a reusable completion event; CPU events are always complete.
+ """Allocate a reusable completion event, or a :class:`HostEvent` on NumPy.
+
+ Host events are always complete and do not measure time.
+
+ Parameters
+ ----------
+ timing : bool, optional
+ Enable CUDA timing, for CuPy's elapsed-time functions; False by default.
- CUDA timing is disabled by default. Pass `timing=True` for events used
- with CuPy's elapsed-time functions; host events do not measure time.
+ Returns
+ -------
+ cupy.cuda.Event or HostEvent
+ The event.
+
+ Examples
+ --------
+ >>> xp.cuda.create_event().done
+ True
"""
if get_backend() == "numpy":
return HostEvent()
@@ -75,10 +166,32 @@ def create_event(*, timing: bool = False) -> Any:
def record_event(event: Any = None, *, stream: Any = None) -> Any:
- """Record completion on `stream` (current stream by default).
-
- Re-recording an event replaces its completion point: enqueue all consumers
- of a recording before reusing the event for a subsequent producer.
+ """Record a completion event on `stream`.
+
+ Re-recording an event replaces its completion point: enqueue all waits on
+ one recording before reusing the event. See :doc:`/guides/execution-helpers`.
+
+ Parameters
+ ----------
+ event : event, optional
+ The event to record; a new one from :func:`create_event` by default.
+ stream : stream, optional
+ The producer stream; the current stream by default.
+
+ Returns
+ -------
+ event
+ The recorded event.
+
+ See Also
+ --------
+ wait_event : Make another stream wait for the event.
+
+ Examples
+ --------
+ >>> producer, consumer = xp.cuda.create_stream(), xp.cuda.create_stream()
+ >>> event = xp.cuda.record_event(stream=producer)
+ >>> xp.cuda.wait_event(event, stream=consumer)
"""
if event is None:
event = HostEvent() if isinstance(stream, HostStream) else create_event()
@@ -87,7 +200,24 @@ def record_event(event: Any = None, *, stream: Any = None) -> Any:
def wait_event(event: Any, *, stream: Any = None) -> None:
- """Order future work on `stream` after `event`, without blocking the CPU."""
+ """Order future work on `stream` after `event`, without blocking the CPU.
+
+ Parameters
+ ----------
+ event : event
+ An event recorded with :func:`record_event`.
+ stream : stream, optional
+ The consumer stream; the current stream by default.
+
+ See Also
+ --------
+ record_event : Record the event to wait for.
+
+ Examples
+ --------
+ >>> event = xp.cuda.record_event(stream=xp.cuda.create_stream())
+ >>> xp.cuda.wait_event(event)
+ """
if stream is None:
if isinstance(event, HostEvent):
return
diff --git a/src/cunumpy/_transfers.py b/src/cunumpy/_transfers.py
index 3b6d6a5..d419023 100644
--- a/src/cunumpy/_transfers.py
+++ b/src/cunumpy/_transfers.py
@@ -1,58 +1,16 @@
"""Counting host/device transfers made through cunumpy.
-Moving data between the host and the device is the classic performance bug of
-a GPU port: a single ``to_numpy()`` inside a time loop makes every step wait for
-the device and copy an array. :func:`count_transfers` records every transfer
-that goes through cunumpy while the block runs, together with the call site
-that caused it, so a test can verify that a time step does not transfer at
-all::
-
- with xp.profiling.count_transfers() as counter:
- propagator(dt)
- assert counter.total == 0, counter.report()
-
-or, to reject host/device copies and host execution on the GPU backend::
-
- with xp.profiling.assert_no_transfers():
- propagator(dt)
-
-Counted are
-
-* ``to_host``: :func:`~cunumpy.to_numpy` (and :func:`~cunumpy.to_cunumpy`)
- called with a device array;
-* ``to_device``: :func:`~cunumpy.to_cupy` (and :func:`~cunumpy.to_cunumpy`)
- called with anything that is not already a device array;
-* ``kernel_conversion``: a :class:`~cunumpy.kernels.PyccelKernel` call that copied
- device arrays to the host (and back), one event per call;
-* ``fallback``: a :class:`~cunumpy.kernels.Kernel` without CUDA kernel calling its host
- kernel on the CuPy backend (``missing_cuda="fallback"``), one event per call;
-* ``device_copy``: device-only dtype/layout conversions in CuNumpy helpers;
-* ``sync``: the host waited for the device: :func:`~cunumpy.synchronize`, the
- waits of the MPI helpers and of the CUDA debug mode, and, on the fake CuPy,
- a scalar read of a device array (``float(a)``, ``int(a)``, ``bool(a)``,
- ``a.item()``, ``a.tolist()``). Syncs are listed in :attr:`TransferCounter.syncs`
- and the report but not in :attr:`~TransferCounter.total`, and
- :func:`assert_no_transfers` accepts them unless called with ``syncs=True``.
-
-Mirror refreshes, argument conversions, staging, serial MPI and kernel output
-copy-back are also counted. Each physical host/device copy is recorded with its
-payload size; conversion/fallback markers have no byte count. ``total`` counts
-all observations, including markers. ``assert_no_transfers`` allows device-only
-copies and rejects host/device copies and host fallback/conversion markers.
-
-Limitations
------------
-Only transfers made *through cunumpy* are seen. Raw ``cupy.ndarray.get()``,
-``cupy.asarray(numpy_array)``, ``numpy.asarray(cupy_array)``, ``float(device_array)``
-(an implicit sync of the real CuPy, which cannot be observed from Python; the fake
-CuPy reports it),
-forwarded backend operations such as ``xp.asarray`` and implicit conversions
-inside other libraries are not counted; use ``nsys``
-(or CuPy's own profiling hooks) to find those.
-
-Like the backend selection, the set of active counters is process-wide state
-and not thread-safe: counters started in one thread see the transfers of every
-thread.
+:func:`count_transfers` records every transfer that goes through cunumpy while
+a block runs, with the call site that caused it, :func:`assert_no_transfers`
+rejects them, and :class:`TransferBudget` checks them per phase of a program.
+All are public in :mod:`cunumpy.profiling`; see :doc:`/guides/profiling`.
+
+Only transfers made *through cunumpy* are seen: raw ``cupy.ndarray.get()``,
+``cupy.asarray(numpy_array)``, ``numpy.asarray(cupy_array)``, forwarded
+backend operations such as ``xp.asarray`` and conversions inside other
+libraries are not, nor is ``float(device_array)`` on the real CuPy (the fake
+CuPy reports it); use ``nsys`` for those. Like the backend selection, the set
+of active counters is process-wide and not thread-safe.
"""
from __future__ import annotations
@@ -95,26 +53,34 @@
@dataclass(frozen=True)
class TransferEvent:
- """One recorded transfer.
+ """One transfer recorded by :func:`count_transfers`.
Attributes
----------
kind : str
- One of ``"to_host"``, ``"to_device"``, ``"kernel_conversion"`` or
- ``"fallback"``, or ``"device_copy"``.
+ ``"to_host"``, ``"to_device"``, ``"kernel_conversion"``,
+ ``"fallback"``, ``"device_copy"`` or ``"sync"`` (see
+ :func:`count_transfers`).
description : str
What was transferred, e.g. ``"to_numpy(shape=(1000,), dtype=float64)"``
or ``"PyccelKernel 'push': 3 array(s) copied to the host"``.
where : str
The call site outside cunumpy, as ``"file:line"``.
- nbytes : int | None
- Payload bytes of a physical copy. None for markers or unknown sizes.
+ nbytes : int or None
+ Payload bytes of a physical copy; None for markers or unknown sizes.
blocking : bool
- Whether the host waited for the copy. False for the copies started by
- :func:`~cunumpy.to_host_async`.
+ Whether the host waited for the copy. False for the copies of
+ :func:`~cunumpy.to_host_async` and
+ :meth:`HostStaging.copy `.
implicit : bool
Whether the event was not asked for explicitly: a scalar read of a
device array (``float(a)``) on the fake CuPy, recorded as a ``sync``.
+
+ Examples
+ --------
+ >>> event = xp.profiling.TransferEvent("to_host", "x", "step.py:12", nbytes=8)
+ >>> print(event)
+ step.py:12: to_host: x
"""
kind: str
@@ -137,12 +103,23 @@ def __str__(self) -> str:
class TransferCounter:
- """Transfers recorded while a :func:`count_transfers` block runs.
+ """The transfers recorded while a :func:`count_transfers` block runs.
+
+ The properties count the events by kind; :meth:`report` lists them by
+ call site. Usually created by :func:`count_transfers`.
Attributes
----------
- events : list[TransferEvent]
+ events : list of TransferEvent
Every recorded transfer, in order.
+
+ Examples
+ --------
+ >>> counter = xp.profiling.TransferCounter()
+ >>> event = xp.profiling.TransferEvent("to_host", "x", "step.py:12", nbytes=8)
+ >>> counter.events.append(event)
+ >>> counter.to_host, counter.bytes_to_host, counter.total
+ (1, 8, 1)
"""
def __init__(self) -> None:
@@ -156,46 +133,68 @@ def _add(self, event: TransferEvent) -> None:
self.events.append(event)
def count(self, kind: str) -> int:
- """Number of events of `kind`."""
+ """Return the number of events of one kind.
+
+ Parameters
+ ----------
+ kind : str
+ An event kind, e.g. ``"to_host"``.
+
+ Returns
+ -------
+ int
+ The number of events of `kind`.
+ """
return sum(1 for event in self.events if event.kind == kind)
@property
def to_host(self) -> int:
- """Number of device-to-host copies (`to_numpy`, `to_cunumpy`)."""
+ """Number of device-to-host copies (``to_host`` events)."""
return self.count("to_host")
@property
def to_device(self) -> int:
- """Number of host-to-device copies (`to_cupy`, `to_cunumpy`)."""
+ """Number of host-to-device copies (``to_device`` events)."""
return self.count("to_device")
@property
def kernel_conversions(self) -> int:
- """Number of `PyccelKernel` calls that copied device arrays to the host."""
+ """Number of ``PyccelKernel`` calls that copied device arrays to the host."""
return self.count("kernel_conversion")
@property
def kernel_conversion_calls(self) -> list[TransferEvent]:
- """The `PyccelKernel` conversion events: kernel name, arrays, call site."""
+ """The ``kernel_conversion`` events: kernel name, arrays, call site."""
return [event for event in self.events if event.kind == "kernel_conversion"]
@property
def fallbacks(self) -> int:
- """Number of `Kernel` calls that fell back to the host kernel on CuPy."""
+ """Number of ``Kernel`` calls that fell back to the host kernel on CuPy."""
return self.count("fallback")
@property
def syncs(self) -> int:
- """Number of times the host waited for the device (see the module documentation)."""
+ """Number of times the host waited for the device (``sync`` events)."""
return self.count("sync")
@property
def total(self) -> int:
- """Number of observations, including conversion/fallback markers (not syncs)."""
+ """Number of events except syncs, including conversion/fallback markers."""
return sum(1 for event in self.events if event.kind != "sync")
def bytes(self, kind: str) -> int:
- """Known bytes copied for `kind`; markers and unknown sizes add zero."""
+ """Return the known bytes copied by the events of one kind.
+
+ Parameters
+ ----------
+ kind : str
+ An event kind, e.g. ``"to_host"``.
+
+ Returns
+ -------
+ int
+ The sum of `nbytes`; markers and unknown sizes add zero.
+ """
return sum(e.nbytes or 0 for e in self.events if e.kind == kind)
@property
@@ -210,11 +209,18 @@ def bytes_to_device(self) -> int:
@property
def device_copies(self) -> int:
- """Device-only conversions recorded by CuNumpy helpers."""
+ """Number of device-only conversions (``device_copy`` events)."""
return self.count("device_copy")
def report(self) -> str:
- """A readable summary: events grouped by kind and call site, with counts."""
+ """Return a readable summary, grouped by kind and call site.
+
+ Returns
+ -------
+ str
+ One header line with the counts per kind, then one line per call
+ site and description, e.g. ``"step.py:12: to_numpy(...) (x3)"``.
+ """
counts = ", ".join(f"{self.count(kind)} {kind}" for kind in KINDS)
lines = [f"{self.total} transfer(s) through cunumpy ({counts})"]
for kind in KINDS:
@@ -233,7 +239,7 @@ def report(self) -> str:
def _caller() -> str:
- """The innermost ``file:line`` on the stack that is outside cunumpy."""
+ """Return the innermost ``file:line`` on the stack outside cunumpy."""
frame = sys._getframe(1)
while frame is not None:
filename = frame.f_code.co_filename
@@ -302,27 +308,55 @@ def count_transfers(
) -> Generator[TransferCounter, None, None]:
"""Count the host/device transfers made through cunumpy in the block.
- Yields a :class:`TransferCounter` that records every ``to_numpy``,
- ``to_cupy`` and ``to_cunumpy`` call that actually copies, every
- :class:`~cunumpy.kernels.PyccelKernel` call that converts device arrays and every
- :class:`~cunumpy.kernels.Kernel` fallback to the host kernel, with the call site of
- each. Nothing is counted for calls that do not copy, e.g. ``to_numpy`` of a
- NumPy array.
-
- Blocks can be nested; each active counter sees the transfers made inside
- it. With `into`, the events are added to that counter (e.g. to accumulate
- over several calls); a counter that is already active is not added again,
- so nested blocks with the same counter count each event once. Like any
- context manager made with :func:`contextlib.contextmanager`, it is also a
- decorator: ``@count_transfers(counter)``. Transfers that bypass cunumpy (raw ``cupy.ndarray.get()``,
- ``cupy.asarray(numpy_array)``, conversions inside other libraries) are not
- seen; see the module documentation.
+ Records every transfer cunumpy makes while the block runs, with its call
+ site. Calls that do not copy record nothing (e.g. ``to_numpy`` of a NumPy
+ array); with no counter active, a call costs one truth test. Blocks can
+ be nested; each active counter sees the transfers made inside it. Like any
+ :func:`contextlib.contextmanager`, it is also a decorator. Transfers that
+ bypass cunumpy (raw ``cupy.ndarray.get()``, conversions inside other
+ libraries) are not seen. Process-wide state, not thread-safe.
+
+ Parameters
+ ----------
+ into : TransferCounter, optional
+ Add the events to this counter instead of a new one (e.g. to
+ accumulate over calls). A counter that is already active is not added
+ again, so each event is counted once.
+
+ Yields
+ ------
+ TransferCounter
+ The counter, filled while the block runs.
+
+ See Also
+ --------
+ assert_no_transfers : Fail if the block transfers.
+ TransferBudget : Count and check transfers per phase.
+
+ Notes
+ -----
+ The event kinds are:
+
+ * ``to_host`` / ``to_device``: a physical copy between host and device
+ (conversions such as :func:`~cunumpy.to_numpy` and
+ :func:`~cunumpy.to_cupy`, mirrors, staging, serial MPI, kernel output
+ copy-back), with its size in bytes;
+ * ``kernel_conversion``: a :class:`~cunumpy.kernels.PyccelKernel` call that
+ copied device arrays to the host and back, one marker per call;
+ * ``fallback``: a :class:`~cunumpy.kernels.Kernel` without CUDA version
+ that ran its host kernel on the CuPy backend, one marker per call;
+ * ``device_copy``: a device-only dtype/layout conversion;
+ * ``sync``: the host waited for the device (:func:`~cunumpy.synchronize`,
+ the MPI helpers, the CUDA debug mode and, on the fake CuPy, scalar
+ reads such as ``float(a)``). Not included in
+ :attr:`TransferCounter.total`.
Examples
--------
>>> with xp.profiling.count_transfers() as counter:
- ... propagator(dt)
- >>> assert counter.total == 0, counter.report()
+ ... y = xp.to_numpy(xp.ones(3)) # a NumPy array: no copy
+ >>> counter.total
+ 0
"""
counter = TransferCounter() if into is None else into
if any(active is counter for active in _ACTIVE):
@@ -339,21 +373,38 @@ def count_transfers(
def assert_no_transfers(
*, syncs: bool = False
) -> Generator[TransferCounter, None, None]:
- """Raise ``AssertionError`` if the block makes a transfer through cunumpy.
+ """Fail if the block makes a host/device transfer through cunumpy.
- Syncs (the host waiting for the device) are accepted unless `syncs` is
- True, e.g. ``assert_no_transfers(syncs=True)`` for a time step that must
- never stall on the device.
+ A :func:`count_transfers` block that, on exit, rejects host/device copies
+ and host conversion/fallback markers. Device-only conversions are allowed.
+ Only checked if the block exits normally; an exception raised inside
+ propagates as it is.
- A :func:`count_transfers` block that, on exit, raises with the counter's
- :meth:`~TransferCounter.report` if a host/device copy or host conversion/
- fallback was counted. Device-only conversions are allowed. Only checked if
- the block exits normally; an exception raised inside propagates as it is.
+ Parameters
+ ----------
+ syncs : bool, optional
+ Also reject syncs (the host waiting for the device), e.g. for a time
+ step that must never stall. Default False.
+
+ Yields
+ ------
+ TransferCounter
+ The counter of the block.
+
+ Raises
+ ------
+ AssertionError
+ If a transfer was counted; the message is the counter's
+ :meth:`~TransferCounter.report`.
+
+ See Also
+ --------
+ count_transfers : Count without failing.
Examples
--------
>>> with xp.profiling.assert_no_transfers():
- ... propagator(dt)
+ ... y = 2.0 * xp.ones(3)
"""
with count_transfers() as counter:
yield counter
@@ -366,8 +417,7 @@ def assert_no_transfers(
class _PhaseRouter(TransferCounter):
- """The counter a :class:`TransferBudget` keeps active: it adds every event
- to the innermost phase of the budget only."""
+ """Counter that adds every event to the innermost phase of a budget."""
def __init__(self, budget: TransferBudget) -> None:
super().__init__()
@@ -388,40 +438,44 @@ class _Rule:
class TransferBudget:
- """Transfers counted per phase of a program, checked against rules per phase.
-
- A time loop typically has phases with different budgets: the time step must
- not copy arrays between host and device at all, the diagnostics may copy a
- few scalars to the host, the output each saved array once. A budget counts
- the transfers of each phase (with :func:`count_transfers`) and checks them::
-
- budget = TransferBudget(started=False)
- model.integrate = budget.count("integrate")(model.integrate)
- ... # setup: not counted
- budget.start()
- for step in range(n_steps):
- model.integrate(dt)
- with budget.phase("output"):
- save(model)
- budget.require("integrate", allow={"to_host": dict(max_nbytes=8)}, calls=n_steps)
- budget.require("output", allow={"to_host": dict(max_count=n, max_total_bytes=b)})
- budget.check() # AssertionError with report() if a rule is broken
-
- Phases nest: an event is counted in the innermost phase only, and a phase
- entered again inside itself (recursion) counts each event and call once.
+ """Count transfers per phase of a program and check a rule for each phase.
+
+ E.g. the time step must not copy between host and device, the diagnostics
+ may copy a few scalars, the output each saved array once. Mark phases with
+ :meth:`phase` or :meth:`count`, set rules with :meth:`require` and call
+ :meth:`check`. Events of a phase accumulate over its calls. Phases nest:
+ an event is counted in the innermost phase only, and a phase re-entered
+ inside itself (recursion) counts each event and call once.
Parameters
----------
- started : bool
- Whether to count from the start; otherwise phases run uncounted until
- :meth:`start`.
+ started : bool, optional
+ Count from the start (default); with False, phases run uncounted
+ until :meth:`start`.
Attributes
----------
- phases : dict[str, TransferCounter]
+ phases : dict of str to TransferCounter
The events of each phase, accumulated over its calls.
- calls : dict[str, int]
+ calls : dict of str to int
How many times each phase was entered while counting.
+ started : bool
+ Whether phases are counted.
+
+ See Also
+ --------
+ count_transfers : Count the transfers of one block.
+
+ Examples
+ --------
+ >>> budget = xp.profiling.TransferBudget()
+ >>> for step in range(3):
+ ... with budget.phase("step"):
+ ... x = 2.0 * xp.ones(4)
+ >>> budget.require("step", allow={"to_host": {"max_nbytes": 8}}, calls=3)
+ >>> budget.check()
+ >>> budget.violations()
+ []
"""
def __init__(self, *, started: bool = True) -> None:
@@ -433,7 +487,7 @@ def __init__(self, *, started: bool = True) -> None:
self._router = _PhaseRouter(self)
def __getitem__(self, phase: str) -> TransferCounter:
- """The counter of `phase` (empty if it has not run)."""
+ """Return the counter of `phase` (empty if it has not run)."""
return self.phases.setdefault(phase, TransferCounter())
def start(self) -> None:
@@ -446,7 +500,18 @@ def stop(self) -> None:
@contextmanager
def phase(self, name: str) -> Generator[TransferCounter, None, None]:
- """Count the transfers of the block in phase `name` (accumulated)."""
+ """Count the transfers of the block in a phase.
+
+ Parameters
+ ----------
+ name : str
+ The phase; its events accumulate over calls.
+
+ Yields
+ ------
+ TransferCounter
+ The counter of the phase (``budget[name]``).
+ """
counter = self[name]
if not self.started:
yield counter
@@ -461,7 +526,18 @@ def phase(self, name: str) -> Generator[TransferCounter, None, None]:
self._stack.pop()
def count(self, name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
- """A decorator: every call of the function is counted in phase `name`."""
+ """Return a decorator that counts every call of a function in a phase.
+
+ Parameters
+ ----------
+ name : str
+ The phase.
+
+ Returns
+ -------
+ callable
+ The decorator, e.g. ``f = budget.count("integrate")(f)``.
+ """
def decorate(function: Callable[..., Any]) -> Callable[..., Any]:
@functools.wraps(function)
@@ -481,25 +557,29 @@ def require(
ignore: Sequence[str] = ("device_copy", "sync"),
calls: int | None = None,
) -> None:
- """Set the rule of `phase`, checked by :meth:`check`.
+ """Set the rule of a phase, checked by :meth:`check`.
Parameters
----------
phase : str
The phase.
- allow : Mapping[str, Mapping | None] | None
+ allow : mapping, optional
The event kinds the phase may have, each with its limits (None or
- ``{}``: any number). Every other kind is forbidden, except those in
- `ignore`. Limits: ``max_nbytes`` (of each event; events of unknown
- size break it), ``max_count`` and ``max_total_bytes`` (over the
- whole phase), ``blocking`` and ``implicit`` (the events must have
- this value of :attr:`TransferEvent.blocking` / ``implicit``), e.g.
- ``{"to_host": dict(max_nbytes=8, blocking=False)}``.
- ignore : Sequence[str]
- Kinds that are not checked unless they are in `allow`: by default
- device-only copies and syncs (the host waiting for the device).
- calls : int | None
+ ``{}``: any number); every other kind is forbidden unless in
+ `ignore`. Limits: ``max_nbytes`` (each event; an unknown size
+ breaks it), ``max_count`` and ``max_total_bytes`` (whole phase),
+ ``blocking`` and ``implicit`` (required value of the event field),
+ e.g. ``{"to_host": {"max_nbytes": 8, "blocking": False}}``.
+ ignore : sequence of str, optional
+ Kinds not checked unless in `allow`; by default ``device_copy``
+ and ``sync``.
+ calls : int, optional
The number of times the phase must have been entered.
+
+ Raises
+ ------
+ ValueError
+ If `allow` has an unknown kind or limit.
"""
allow = {kind: dict(limits or {}) for kind, limits in (allow or {}).items()}
for kind, limits in allow.items():
@@ -513,7 +593,14 @@ def require(
self._rules[phase] = _Rule(allow, tuple(ignore), calls)
def violations(self) -> list[str]:
- """The broken rules, one line each, with the offending events."""
+ """Return the broken rules.
+
+ Returns
+ -------
+ list of str
+ One line per broken rule, with the offending event and its call
+ site.
+ """
found = []
for name, rule in self._rules.items():
events = self[name].events
@@ -550,7 +637,14 @@ def violations(self) -> list[str]:
return found
def report(self) -> str:
- """The events of every phase (see :meth:`TransferCounter.report`) and the broken rules."""
+ """Return the events of every phase and the broken rules.
+
+ Returns
+ -------
+ str
+ :meth:`TransferCounter.report` of each phase, then the
+ :meth:`violations`.
+ """
lines = []
for name, counter in self.phases.items():
lines.append(f"{name} ({self.calls.get(name, 0)} call(s)):")
@@ -562,6 +656,12 @@ def report(self) -> str:
return "\n".join(lines)
def check(self) -> None:
- """Raise ``AssertionError`` with :meth:`report` if a rule is broken."""
+ """Raise if a rule is broken.
+
+ Raises
+ ------
+ AssertionError
+ If a rule is broken; the message is :meth:`report`.
+ """
if self.violations():
raise AssertionError("transfer budget exceeded:\n" + self.report())
diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py
index 6ffc26c..c13cebc 100644
--- a/src/cunumpy/algorithms.py
+++ b/src/cunumpy/algorithms.py
@@ -1,15 +1,16 @@
"""Array algorithms missing from NumPy/CuPy, on either backend.
Morton (Z-order) keys (the same as ``cunumpy/morton.cuh`` computes in a
-kernel), a stable sort of several arrays by one key, sums per key, and the compaction of
-the live rows of particle arrays::
-
- import cunumpy as xp
+kernel, at most ``MAX_MORTON_LEVELS`` bits per axis), a stable sort of several
+arrays by one key, sums per key, and the compaction of the live rows of
+particle arrays::
keys = xp.algorithms.morton_keys(positions, lower, upper, levels)
keys, order, positions = xp.algorithms.sort_by_key(keys, positions)
charge = xp.algorithms.segment_sum(q, cell, n_cells)
n = xp.algorithms.compact_by_mask(alive, positions, charges)
+
+See :doc:`/guides/particle-codes`.
"""
from cunumpy._algorithms import (
diff --git a/src/cunumpy/arguments.py b/src/cunumpy/arguments.py
index 5ef2de1..b2ccfb7 100644
--- a/src/cunumpy/arguments.py
+++ b/src/cunumpy/arguments.py
@@ -7,9 +7,6 @@
Kernels receive argument objects as they are; cunumpy never converts one form
into the other::
- import cunumpy as xp
-
-
class CudaMarkerArguments(xp.arguments.CudaStructArguments):
struct_name = "MarkerArgs"
fields = (("markers", "Array2D"), ("n_markers", "int"))
@@ -22,7 +19,7 @@ def __init__(self, markers, n_markers):
:class:`CudaArguments` flattens a group into several kernel parameters,
:class:`CudaStruct` defines a C struct passed by value, and
:func:`write_cuda_header` writes struct definitions to a header. The kernel
-classes are in :mod:`cunumpy.kernels`.
+classes are in :mod:`cunumpy.kernels`. See :doc:`/kernels/arguments`.
"""
from cunumpy._cuda_kernel import (
diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py
index 5a026fc..2ec9d24 100644
--- a/src/cunumpy/cuda/__init__.py
+++ b/src/cunumpy/cuda/__init__.py
@@ -1,25 +1,18 @@
-"""CUDA device runtime, reusable streams/events, and CUDA source tools.
+"""CUDA devices, memory, streams, debug mode and CUDA source tools.
The device functions (:func:`set_device`, :func:`memory_info`, :func:`stream`,
-...) do nothing (or return ``None``/``0``) on the NumPy backend::
+...) do nothing, or return ``None`` or ``0``, on the NumPy backend, so the same
+code runs on both::
- import cunumpy as xp
-
- kernel = xp.kernels.CudaKernel(source, "push")
with xp.cuda.stream():
kernel(positions, velocities, dt, n_threads=n)
-The CUDA headers shipped with cunumpy (``cunumpy/atomic.cuh``,
-``cunumpy/random.cuh``, ...) are in :func:`cuda_include_dir`;
-:func:`parse_cuda_signature` and the other source tools inspect CUDA sources.
-
-Reusable :func:`create_stream` and :func:`create_event` return synchronous host
-equivalents on NumPy, supporting the same recording and completion interface.
-
-The kernel classes (:class:`~cunumpy.kernels.CudaKernel`,
-:class:`~cunumpy.kernels.PyccelKernel`, :class:`~cunumpy.kernels.Kernel`, ...)
-are in :mod:`cunumpy.kernels`, the argument objects
-(:class:`~cunumpy.arguments.CudaStructArguments`, ...) in :mod:`cunumpy.arguments`.
+:func:`create_stream` and :func:`create_event` return synchronous
+:class:`HostStream` and :class:`HostEvent` objects on NumPy. The CUDA headers
+shipped with cunumpy (``cunumpy/atomic.cuh``, ...) are in
+:func:`cuda_include_dir`; :func:`parse_cuda_signature` and the other source
+tools inspect CUDA sources. The kernel classes are in :mod:`cunumpy.kernels`,
+the argument objects in :mod:`cunumpy.arguments`. See :doc:`/guides/gpu-devices`.
Importing this module makes ``xp.cuda`` refer to it instead of ``cupy.cuda``;
use ``import cupy; cupy.cuda`` for CuPy's module.
diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py
index a83c4ee..1ba0706 100644
--- a/src/cunumpy/kernel_testing.py
+++ b/src/cunumpy/kernel_testing.py
@@ -1,55 +1,47 @@
-"""Helpers for testing host/CUDA kernel pairs with pytest.
-
-A code base that ports its kernels to CUDA one by one needs the same test for
-every kernel: build the arguments on both backends, run the host kernel and the
-CUDA kernel, and compare what they wrote. This module provides that test
-(:func:`assert_kernels_agree`), the pytest markers to parametrize tests over
-the backends (:data:`BACKENDS`, :data:`requires_cupy`, the :func:`backend`
-fixture), :func:`device_function_kernel`, which wraps a ``__device__``
-function in an elementwise ``__global__`` kernel so that device helpers can be
-tested from Python without a hand-written test kernel, and
-:func:`emulate_cuda_kernel` (from the private module ``cunumpy._emulation``), which runs a CUDA
-kernel on the CPU, one thread after another, so that its arithmetic can be
-checked against the host kernel in CI without a GPU. Without a GPU, the CuPy
-code paths of a program (argument objects, conversions, backend branches) can
-still run on the fake CuPy of :mod:`cunumpy._fake_cupy` (:func:`install_fake_cupy`,
-or ``CUNUMPY_FAKE_CUPY=1``); :func:`fake_cupy_active` tells whether it is in
-use, and ``requires_cupy`` skips the tests that launch kernels then.
-:func:`fake_cupy_session` runs a whole CuPy-backend program on it (launches and
-compilation emulated, see :func:`emulated_launches`), :data:`requires_device_backend`
-skips tests that need a GPU or the fake CuPy, and :func:`run_in_fake_cupy_subprocess`
-runs code in a child process on the fake CuPy from a test process that cannot
-switch to it.
-
-A catalog's parity tests need no code per kernel when each kernel folder
-holds ``_test_args.py`` with ``make_args(backend, seed)`` (and
-``N_THREADS``): :func:`parity_cases` and :func:`check_parity` drive
-:func:`assert_kernels_agree` from these modules.
-
-The module imports pytest only when one of its pytest objects is used, so it
-can be imported (e.g. for :func:`device_function_kernel`) without pytest, and
-``import cunumpy`` never imports pytest.
+"""Pytest helpers for testing host/CUDA kernel pairs.
+
+Not imported by ``import cunumpy``; pytest is imported only when one of the
+pytest objects below is used, so e.g. :func:`device_function_kernel` works
+without it. See :doc:`/kernels/testing`.
+
+* Compare host and CUDA kernels: :func:`assert_kernels_agree`, and for a
+ catalog :func:`parity_cases` / :func:`check_parity`.
+* Run on both backends (built on first access):
+
+ ``BACKENDS``
+ ``["numpy", pytest.param("cupy", marks=requires_cupy)]``, to
+ parametrize a test over the backends.
+ ``backend``
+ A pytest fixture that runs the test once per backend, with it active.
+ ``requires_cupy``
+ A ``skipif`` marker for tests that launch CUDA kernels (needs a GPU,
+ not the fake CuPy).
+ ``requires_device_backend``
+ A ``skipif`` marker for tests that need a GPU or the fake CuPy, see
+ :func:`device_backend_available`.
+
+ With ``CUNUMPY_REQUIRE_CUDA=1`` (:func:`cuda_required`) they fail
+ instead of skipping.
+* Test ``__device__`` helpers: :func:`device_function_kernel`.
+* Run CUDA kernels on the CPU, without a GPU (from ``cunumpy._emulation``):
+ :func:`emulate_cuda_kernel`, :func:`emulated_launches`,
+ :func:`compile_for_emulation`, :func:`emulation_compiler` and
+ :func:`emulation_cache_dir`.
+* Run CuPy-backend code on the fake CuPy, a strict host stand-in for CuPy
+ (also installed by ``CUNUMPY_FAKE_CUPY=1``): :func:`install_fake_cupy`,
+ :func:`fake_cupy_active`, :func:`fake_cupy_session`,
+ :func:`run_in_fake_cupy_subprocess` and :func:`host_buffer` (the NumPy
+ array behind a fake CuPy array, not a copy).
Examples
--------
-One parametrised test covers every kernel of a catalog that has a CUDA
-version::
+One parametrized test covers every kernel of a catalog::
- import pytest
- from cunumpy.kernel_testing import assert_kernels_agree
+ from cunumpy.kernel_testing import parity_cases, check_parity
- from my_kernels import catalog
-
-
- def make_args(backend, seed):
- rng = xp.rng.get_rng(seed) # the backend is active: arrays land on it
- x = rng.random(1000)
- return (x, 2.0, x.size)
-
-
- @pytest.mark.parametrize("name, kernel", catalog.parity_cases())
- def test_parity(name, kernel):
- assert_kernels_agree(kernel, make_args, n_threads=1000)
+ @pytest.mark.parametrize("kernel", parity_cases(catalog))
+ def test_parity(kernel):
+ check_parity(kernel)
"""
from __future__ import annotations
@@ -119,16 +111,46 @@ def test_parity(name, kernel):
def fake_cupy_active() -> bool:
- """Whether the fake CuPy (:mod:`cunumpy._fake_cupy`) stands in for CuPy."""
+ """Tell whether the fake CuPy stands in for CuPy in this process.
+
+ Returns
+ -------
+ bool
+ True after :func:`install_fake_cupy` or with ``CUNUMPY_FAKE_CUPY=1``.
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import fake_cupy_active
+ >>> fake_cupy_active() # doctest: +SKIP
+ False
+ """
return _fake_cupy.is_active()
def install_fake_cupy() -> Any:
- """Install the fake CuPy for this process; see :mod:`cunumpy._fake_cupy`.
+ """Install the fake CuPy, a strict host stand-in for CuPy, in this process.
- Call it before the first backend use (e.g. at the top of ``conftest.py``),
- or set ``CUNUMPY_FAKE_CUPY=1`` in the environment instead. Returns the
- fake ``cupy`` module.
+ Its arrays live in host memory but behave like CuPy arrays (e.g.
+ ``numpy.asarray`` rejects them); CUDA kernels cannot run on it except in
+ :func:`emulated_launches`. Call it before the first backend use (e.g. at
+ the top of ``conftest.py``), or set ``CUNUMPY_FAKE_CUPY=1`` instead.
+ Calling it again is a no-op.
+
+ Returns
+ -------
+ module
+ The fake ``cupy`` module.
+
+ Raises
+ ------
+ RuntimeError
+ If the real CuPy was imported already, or cunumpy already checked
+ for CuPy.
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import install_fake_cupy
+ >>> install_fake_cupy() # doctest: +SKIP
"""
return _fake_cupy.install()
@@ -139,12 +161,23 @@ def _can_launch() -> bool:
def cuda_required() -> bool:
- """Whether ``CUNUMPY_REQUIRE_CUDA`` demands a real GPU (the CI guard).
+ """Tell whether ``CUNUMPY_REQUIRE_CUDA`` demands a real GPU (the CI guard).
+
+ If so, tests that need a GPU fail instead of being skipped (``requires_cupy``,
+ ``requires_device_backend``, :func:`assert_kernels_agree`, the ``cupy``
+ run of the ``backend`` fixture), so that a GPU CI job cannot pass
+ silently because CuPy or the driver is broken.
- Then tests that need a GPU fail instead of being skipped where there is
- none (:data:`requires_cupy`, :func:`assert_kernels_agree`, the ``cupy``
- parameter of :func:`backend`), so that a CI job on a GPU machine cannot
- pass silently because CuPy or the driver is broken.
+ Returns
+ -------
+ bool
+ True if ``CUNUMPY_REQUIRE_CUDA`` is ``1``, ``true`` or ``yes``.
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import cuda_required
+ >>> cuda_required() # doctest: +SKIP
+ False
"""
return os.environ.get("CUNUMPY_REQUIRE_CUDA", "").lower() in {"1", "true", "yes"}
@@ -157,7 +190,13 @@ def _skip_or_fail(reason: str) -> None:
def cuda_gate() -> bool:
- """Condition of ``requires_cupy`` under ``CUNUMPY_REQUIRE_CUDA``: fail, or False."""
+ """Return the ``requires_cupy`` condition under ``CUNUMPY_REQUIRE_CUDA``.
+
+ Returns
+ -------
+ bool
+ False (do not skip) if CUDA kernels can run; otherwise the test fails.
+ """
if not _can_launch():
reason = FAKE_SKIP_REASON if fake_cupy_active() else SKIP_REASON
_pytest().fail(f"CUNUMPY_REQUIRE_CUDA is set, but {reason}", pytrace=False)
@@ -165,18 +204,36 @@ def cuda_gate() -> bool:
def device_backend_available() -> bool:
- """Whether a CuPy-backend program can run here: a GPU, or the fake CuPy.
+ """Tell whether a CuPy-backend program can run here: a GPU, or the fake CuPy.
+
+ Unlike ``requires_cupy`` (CUDA kernels can be launched), this is also true
+ on the fake CuPy, where launches only work inside :func:`emulated_launches`
+ (see :func:`fake_cupy_session`). The marker ``requires_device_backend``
+ skips tests where it is false.
+
+ Returns
+ -------
+ bool
+ True with a working CuPy or with the fake CuPy active.
- Unlike :data:`requires_cupy` (CUDA kernels can be launched), this is also
- True on the fake CuPy, where launches only work inside
- :func:`emulated_launches` (see :func:`fake_cupy_session`). The marker
- :data:`requires_device_backend` skips tests that need it.
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import device_backend_available
+ >>> device_backend_available() # doctest: +SKIP
+ True
"""
return fake_cupy_active() or cupy_available()
def device_backend_gate() -> bool:
- """Condition of ``requires_device_backend`` under ``CUNUMPY_REQUIRE_CUDA``."""
+ """Return the ``requires_device_backend`` condition under ``CUNUMPY_REQUIRE_CUDA``.
+
+ Returns
+ -------
+ bool
+ False (do not skip) if a device backend is available; otherwise the
+ test fails.
+ """
if not device_backend_available():
_pytest().fail(
f"CUNUMPY_REQUIRE_CUDA is set, but {DEVICE_SKIP_REASON}", pytrace=False
@@ -193,17 +250,31 @@ def fake_cupy_session(
"""Run a CuPy-backend program on the CPU, on the fake CuPy.
Activates the CuPy backend (the fake CuPy) and emulates every CUDA launch
- and compilation in the block (:func:`emulated_launches`, which takes
- `compiler` and `options`)::
+ and compilation in the block with :func:`emulated_launches`.
- with fake_cupy_session():
- sim.run() # kernels compiled up front and launched, all on the CPU
+ Parameters
+ ----------
+ compiler : str, optional
+ The C++ compiler of the emulation; default :func:`emulation_compiler`.
+ options : sequence of str, optional
+ Extra compiler options, e.g. ``("-ffp-contract=off",)``.
+
+ Yields
+ ------
+ None
+ The block runs on the fake CuPy with emulated launches.
Raises
------
RuntimeError
If the fake CuPy is not active (:func:`install_fake_cupy`,
``CUNUMPY_FAKE_CUPY=1``; or use :func:`run_in_fake_cupy_subprocess`).
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import fake_cupy_session
+ >>> with fake_cupy_session(): # doctest: +SKIP
+ ... sim.run() # kernels compiled and launched, all on the CPU
"""
if not fake_cupy_active():
raise RuntimeError(
@@ -245,30 +316,39 @@ def run_in_fake_cupy_subprocess(
env: Mapping[str, str] | None = None,
timeout: float | None = None,
) -> subprocess.CompletedProcess:
- """Run `code` in a serial child Python process on the fake CuPy; fail the test if it fails.
+ """Run code in a serial child process on the fake CuPy; fail the test if it fails.
The fake CuPy must be installed before anything imports cunumpy, so a test
- process that already uses cunumpy cannot switch to it; the child starts
- with ``CUNUMPY_FAKE_CUPY=1``. It runs ``python -X faulthandler -c code``
- (a crash prints the Python traceback) with ``OMP_NUM_THREADS=1`` and
- without the variables of an MPI launcher, so it does not join the MPI job
- of the parent (``MAYBEMPI=0``). Under MPI only rank 0 starts the child and
- the other ranks skip the test: concurrent children are not needed for a
- serial check, and have crashed external libraries.
+ process that already uses cunumpy cannot switch to it. The child runs
+ ``python -X faulthandler -c code`` with ``CUNUMPY_FAKE_CUPY=1``,
+ ``OMP_NUM_THREADS=1``, ``MAYBEMPI=0`` and without the variables of an MPI
+ launcher, so it does not join the parent's MPI job. Under MPI only rank 0
+ starts the child; the other ranks skip the test.
Parameters
----------
code : str
Python source to run.
- env : Mapping[str, str] | None
+ env : mapping of str to str, optional
Additional environment variables for the child (e.g. ``PYTHONPATH``).
- timeout : float | None
+ timeout : float, optional
Seconds after which the child is killed and the test fails.
Returns
-------
subprocess.CompletedProcess
The finished child (exit code 0), with its stdout and stderr.
+
+ Notes
+ -----
+ If the child fails or times out, the test fails (``pytest.fail``) with the
+ exit code or signal (e.g. ``SIGSEGV``) and the last 50 lines of its stdout
+ and stderr.
+
+ Examples
+ --------
+ >>> from cunumpy.kernel_testing import run_in_fake_cupy_subprocess
+ >>> run_in_fake_cupy_subprocess("import my_sim; my_sim.check()") # doctest: +SKIP
"""
pytest = _pytest()
from cunumpy.mpi import get_mpi
@@ -364,6 +444,16 @@ def backend(request):
With ``CUNUMPY_REQUIRE_CUDA`` set, the ``cupy`` run fails if the CuPy
backend cannot be activated, instead of silently running on NumPy.
+
+ Parameters
+ ----------
+ request : pytest.FixtureRequest
+ The pytest request; ``request.param`` is the backend.
+
+ Yields
+ ------
+ str
+ The active backend, ``"numpy"`` or ``"cupy"``.
"""
with use_backend(request.param, strict=cuda_required()):
yield request.param
@@ -430,7 +520,7 @@ def _resolve_output(
n_args: int,
parameters: Sequence[str] | None,
) -> tuple[int, str | None]:
- """The argument index and the field filter (or None) of an `outputs` entry."""
+ """Return the argument index and field filter (or None) of an `outputs` entry."""
if isinstance(entry, bool) or not isinstance(entry, (int, str)):
raise TypeError(
"outputs entries must be argument indices (int) or names (str, "
@@ -463,22 +553,7 @@ def _collect_arrays(
outputs: Sequence[int | str] | None = None,
parameters: Sequence[str] | None = None,
) -> dict[str, Any]:
- """The arrays among `args` (or among the arguments `outputs`), by name.
-
- An argument that is an array is named ``"argument "``; arrays found one
- level deep, in a tuple, list or dict argument or in the attributes of an
- argument object (e.g. a ``CudaArguments`` object), are named
- ``"argument []"`` or ``"argument .