diff --git a/.github/workflows/gpu_ci_trigger.yml b/.github/workflows/gpu_ci_trigger.yml index 4418074..60fa668 100644 --- a/.github/workflows/gpu_ci_trigger.yml +++ b/.github/workflows/gpu_ci_trigger.yml @@ -64,6 +64,7 @@ jobs: # 4. Force push (This automatically starts the GitLab Pipeline) git push -f gitlab HEAD:refs/heads/$TARGET_BRANCH + echo "PUSHED_SHA=$(git rev-parse HEAD)" >> $GITHUB_ENV # 5. Provide the direct link PIPELINE_URL="https://gitlab.mpcdf.mpg.de/maxlin/cunumpy/-/pipelines?ref=$TARGET_BRANCH" @@ -72,10 +73,41 @@ jobs: echo "::notice::View Pipeline: $PIPELINE_URL" - name: Wait for GitLab Pipeline - uses: docker://gitlab/glab:latest + timeout-minutes: 90 env: GITLAB_TOKEN: ${{ secrets.GITLAB_TOKEN }} - GITLAB_HOST: gitlab.mpcdf.mpg.de - with: - entrypoint: glab - args: ci status --live --branch ${{ env.TARGET_BRANCH }} --repo maxlin/cunumpy + run: | + API="https://gitlab.mpcdf.mpg.de/api/v4/projects/maxlin%2Fcunumpy" + AUTH=() + if [ -n "$GITLAB_TOKEN" ]; then AUTH=(--header "PRIVATE-TOKEN: $GITLAB_TOKEN"); fi + + # 1. GitLab creates the pipeline asynchronously after the push: wait until + # the pipeline for the pushed commit exists (asking right away finds none). + PIPELINE="" + for i in $(seq 1 60); do + PIPELINE=$(curl -sf "${AUTH[@]}" "$API/pipelines?ref=$TARGET_BRANCH&sha=$PUSHED_SHA&per_page=1" | jq -r '.[0].id // empty' || true) + if [ -n "$PIPELINE" ]; then break; fi + sleep 5 + done + if [ -z "$PIPELINE" ]; then + echo "::error::No GitLab pipeline for $PUSHED_SHA on $TARGET_BRANCH after 5 minutes" + exit 1 + fi + URL="https://gitlab.mpcdf.mpg.de/maxlin/cunumpy/-/pipelines/$PIPELINE" + echo "::notice::GitLab pipeline: $URL" + + # 2. Follow the pipeline until it has finished. + while true; do + STATUS=$(curl -sf "${AUTH[@]}" "$API/pipelines/$PIPELINE" | jq -r '.status // empty' || true) + case "$STATUS" in + success) + echo "GitLab pipeline $PIPELINE succeeded: $URL" + exit 0 ;; + failed|canceled|skipped|manual|scheduled) + echo "::error::GitLab pipeline $PIPELINE is $STATUS: $URL" + exit 1 ;; + *) + echo "GitLab pipeline $PIPELINE: ${STATUS:-status not available yet}" + sleep 20 ;; + esac + done diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index aa01052..4188463 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -17,7 +17,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.8", "3.10", "3.13"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: # Checkout the repository diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 79d4881..17a6fcb 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -29,18 +29,19 @@ gpu_tests: - git --version - echo "--- Pytest Execution ---" - # The MPCDF image likely has a specific python environment. + # The MPCDF image likely has a specific python environment. # We install our dependencies into the user directory or a virtualenv. - - python3 -m pip install --user cupy-cuda12x - - python3 -m pip install --user nvidia-cublas-cu12 nvidia-cufft-cu12 nvidia-curand-cu12 nvidia-cusolver-cu12 nvidia-cusparse-cu12 + # One CUDA version only: nvhpcsdk/26 provides CUDA 13.2 (its headers are used + # when CuPy compiles kernels with NVRTC), so install CuPy for CUDA 13 with the + # CUDA 13.2 libraries and NVRTC. Mixing CUDA 12 (cupy-cuda12x) with the CUDA 13.2 + # headers fails to compile CuPy's own kernels (e.g. CUB reductions). + - python3 -m pip install --user "cupy-cuda13x[ctk]" "cuda-toolkit==13.2.*" - python3 -m pip install --user -e . - + # Add the user bin to PATH for pytest - export PATH="$HOME/.local/bin:$PATH" - - # Try to find libcublas and other libraries in the HPC environment - - export LD_LIBRARY_PATH=$(find /mpcdf/soft /opt/nvidia -name libcublas.so.12 -exec dirname {} \; 2>/dev/null | head -n 1):$LD_LIBRARY_PATH - + - export ARRAY_BACKEND=cupy + - python3 -c "import cupy; cupy.show_config()" - pytest -xvs . diff --git a/CHANGELOG.md b/CHANGELOG.md index 3093847..d5e5da3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,61 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Removed +- Support for Python 3.8 and 3.9 (both end-of-life); `cunumpy` now requires Python 3.10 or newer. + +### Changed +- Python 3.14 is supported. +- `KernelCatalog.compile_all(jobs=1)` and `CudaKernelVariants.compile_all(keys, jobs=1)`: With `jobs > 1` the CUDA kernels are compiled in threads (NVRTC releases the GIL); `jobs=None` uses the number of CPUs. All kernels are compiled even if one fails, and the first error is raised afterwards. +- CI now tests every supported Python version (3.10, 3.11, 3.12, 3.13 and 3.14) instead of 3.8/3.10/3.13. +- `CudaKernel.compile()` passes `compile_options()` to CuPy: the given `options` plus `-DCUNUMPY_INCLUDE_HASH=0x` when the source includes header files, so CuPy's kernel cache (keyed on source and options only) is invalidated when an included header changes. `options` still returns the options as given. +- `KernelCatalog.from_package(..., include_dirs=None)` is now an explicit keyword; by default the source root of the top-level package (the directory containing it) is an include directory of every CUDA kernel, in addition to the kernel's own folder, so kernels can `#include "my_pkg/common.cuh"`. + +### Added +- `xp.CudaKernel`: Wraps a CUDA C kernel (`cupy.RawKernel`, compiled lazily with NVRTC) so it can be called with the same arguments as the host kernel it mirrors, plus `n_threads`. The `extern "C" __global__` signature is parsed once and every call is checked against it: argument count, array dtypes (host arrays raise, they are never copied), and scalars (Python scalars are cast to the declared C types with range checks; lossy or mismatching scalars raise instead of reaching the kernel as silently wrong values). Supports `block_size`, NVRTC `options`, `include_dirs`, `shared_mem`, `stream`, `CudaKernel.from_file()` (`_cuda.cu`), `compile()` and `prepare_args()`; `check_signature=False` skips the checks. +- `xp.CudaArguments`: Base class for argument objects that are flattened into several CUDA kernel arguments; any object with a `__cuda_args__()` method is flattened. +- `xp.parse_cuda_signature(source, name)` and `xp.CudaParameter`: Parse the parameters of a `__global__` function. +- `xp.Kernel`: A host kernel (`PyccelKernel`) and its CUDA counterpart, calling the one matching the active backend. Without a CUDA kernel on the CuPy backend it raises `NotImplementedError` (`missing_cuda="raise"`, default) or falls back to the host kernel with host copies (`missing_cuda="fallback"`). +- `xp.KernelCatalog`: Read-only mapping of `Kernel`s; `KernelCatalog.from_package()` collects them from a package with one folder per kernel (`name/name_kernels.py`, `name/name_cuda.cu`); `without_cuda` lists the kernels still to port. +- `xp.CudaStruct` and `xp.CudaStructValue`: C structs passed to CUDA kernels by value. A `CudaStruct` is defined once from `(field, C type)` pairs; it provides the C `declaration`, the NumPy `dtype` with the C memory layout, and packs values (device arrays as addresses, scalars checked and cast) into a `CudaStructValue` that is passed as one kernel argument. `CudaKernel(..., structs=[...])` checks struct parameters and that a struct definition in the source matches. +- `CudaKernel(..., template_args=...)`: Instantiate C++ function templates (e.g. `template_args=(np.float64, 3)` for `name`); the template parameters are substituted into the checked signature. +- `xp.CudaKernelVariants`: Creates and caches one `CudaKernel` per variant key for generated kernel sources (e.g. per dimension and dtype); `compile_all()` compiles given and existing variants. +- `xp.ctype_of(dtype)`: The C type of a NumPy dtype, e.g. for generating CUDA source. +- 1D to 3D launches: `CudaKernel` accepts a tuple `block_size`, and calls take `n_threads` as an integer or tuple, or an explicit `grid`, plus a per-call `block`; `launch_shape()` returns the `(grid, block)` of a call. Dynamic shared memory (`shared_mem`) and `stream` are passed through by `Kernel` as well. +- Compiling at setup: `CudaKernel.compile()` and `is_compiled`, `Kernel.compile()`, and `KernelCatalog.compile_all()`. +- `Kernel(..., host_options=...)` and `KernelCatalog.from_package(..., host_options=...)`: `PyccelKernel` options (e.g. `object_modules`, `outputs`) for the host kernels, for all kernels or per kernel name; needed for the fallback to find device arrays inside application objects. +- `xp.local_rank()`: The node-local rank from the MPI launcher's environment (Open MPI, MVAPICH2, Intel MPI/MPICH, PMI, Cray PALS, Slurm, `LOCAL_RANK`), available before `MPI_Init`. +- `xp.bind_local_device()`: Selects the GPU `local_rank() % device_count()` and creates its context, before `MPI_Init`, for one-rank-per-GPU MPI programs. +- `KernelCatalog.summary()` (also `str(catalog)`) and `KernelCatalog.with_cuda`: One line on the porting status, `"CUDA kernels: 3 of 60 (missing: a, b, c)"` (at most 10 missing names, then `...`), and the names of the ported kernels, mirroring `without_cuda`. +- `CudaKernel.all_from_file(path, **kwargs)` and `xp.cuda_kernel_names(source)`: Load every `__global__` function of a file as a `CudaKernel` (a `dict` by name, sharing the source so CuPy compiles the file once), and list the `__global__` functions of a source string. `KernelCatalog.from_package` still uses only the function `` of `_cuda.cu`; other kernels in the file are ignored by the catalog. +- `xp.synchronize_for_mpi(*arrays)`: Waits for pending work on the current stream before MPI uses device buffers (no-op for host buffers and on the NumPy backend). +- CUDA headers shipped as package data, found by every `CudaKernel` automatically (`xp.cuda_include_dir()` for other compilers): `cunumpy/array_view.cuh` defines the strided views `Array1D`, `Array2D`, `Array3D` (`data`, `shape[]`, `strides[]` in elements, `a(i, j)`, `size()`, bounds checks with `-DCUNUMPY_BOUNDS_CHECK`), so ported pyccel kernels index `markers(ip, j)` instead of computing offsets from hand-passed sizes; `cunumpy/index.cuh` defines `CUNUMPY_THREAD_1D(i, n)` (2D/3D likewise) and the grid-stride loop `CUNUMPY_GRID_STRIDE_1D(i, n)`. +- Array view parameters and struct fields: `CudaKernel` signatures and `CudaStruct` fields accept `Array1D` to `Array3D` of the scalar C types (`CudaParameter.view_ndim`). Such a parameter takes a CuPy array of the declared dtype and number of dimensions, contiguous or not, and is packed into (pointer, shape, strides in elements) with the memory layout of the C struct; dtype and ndim mismatches raise. `CudaStruct.has_views` tells whether the declaration needs `#include "cunumpy/array_view.cuh"`. +- `CudaStruct.from_signature(func, name, *, int_type="long long", scalar_names=None)`: Builds a struct from the pyccel-style annotations of a Python function (typically an argument class's `__init__`; `self` is skipped): `"float[:, :]"` -> `Array2D`, `"int[:]"` -> `Array1D`, `int` -> `int_type`, `float` -> `double`, `bool` -> `bool`; `Final[...]` and `const` are ignored. The Python class becomes the one definition of the arguments on the host and on the device. +- `CudaStruct.to_header(path=None, *, guard=None, includes=())` and `xp.write_cuda_header(path, structs, guard=None, *, includes=())`: Generate a header with include guard, the `cunumpy/array_view.cuh` include when needed, and the struct definitions; the header can be written at setup and a test can assert that the committed file equals the generated one. +- `xp.DeviceMirror(host_array)`: Pairs a host NumPy array owned by another library (e.g. a stencil vector's `_data`) with a lazily allocated device copy. `device` is the CuPy array kernels write into on the CuPy backend, and the host array itself on the NumPy backend, so accumulation code is written once without copies on the CPU. `to_device()` and `to_host()` copy explicitly and in place (the host array keeps its identity; no-ops on NumPy), `zero()` clears the buffer, and `rebind()` follows a reallocation by the owner; a host array whose shape or dtype changed raises `ValueError`. Non-NumPy inputs raise `TypeError`. +- `cunumpy/atomic.cuh`: CUDA header shipped with the package, with `cunumpy_atomic_add(double*|float*, value)` (`atomicAdd`, with a compare-and-swap fallback for `double` before sm_60) and the indexed `cunumpy_atomic_add_2d()` / `cunumpy_atomic_add_3d()` for C-contiguous arrays, for many-threads-to-one-cell accumulation. +- `xp.cuda_include_dir()`: The directory of the shipped CUDA headers; `CudaKernel` adds it to its NVRTC options automatically (once), so sources can `#include `. +- `cunumpy.testing`: Helpers for testing host/CUDA kernel pairs with pytest (pytest is imported only when its objects are used, never by `import cunumpy`). `requires_cupy` is a `skipif` marker for tests that need a GPU, `BACKENDS = ["numpy", pytest.param("cupy", marks=requires_cupy)]` parametrizes a test over the backends, and the `backend` fixture runs a test once per backend with that backend active. +- `cunumpy.testing.assert_kernels_agree(kernel, make_args, *, n_threads=..., rtol=1e-12, atol=0.0, n_calls=1, outputs=None, seed=0)`: Builds the arguments with `make_args(backend, seed)` on the NumPy and the CuPy backend, runs the host and the CUDA kernel of a `Kernel`, and compares the arrays they wrote (the declared `outputs`, or every array argument, including arrays held by argument objects) with `numpy.testing.assert_allclose`, naming the differing argument. Skips the test without a GPU and returns the host arrays. +- `KernelCatalog.parity_cases()`: The `(name, kernel)` pairs of the kernels with a CUDA version, so one test parametrised with them and `assert_kernels_agree` covers a whole catalog. +- `cunumpy.testing.device_function_kernel(header_source, signature, *, name=None, includes=(), n_threads_param="n")`: Generates an elementwise `extern "C" __global__` wrapper around a `__device__` function given its C prototype (pointer parameters are shared, scalar parameters become per-thread arrays, the return value goes into `out`), so device helpers can be tested from Python against their host versions without a hand-written test kernel. +- `xp.mpi_is_cuda_aware(comm=None)`: Collective startup check of whether the MPI library can pass device buffers: each rank exchanges a tiny CuPy array with `Sendrecv` and the ranks agree with `allreduce`; any failure gives `False`. Returns `False` on the NumPy backend and without a functional CuPy, without importing `mpi4py`. +- `xp.require_cuda_aware_mpi(comm=None)`: Raises `RuntimeError`, with hints on obtaining a CUDA-aware build, when `mpi_is_cuda_aware()` fails on the CuPy backend; no-op on NumPy. +- `xp.nvtx_range(name, color=None)`: Context manager and decorator marking a code region as an NVTX range (`cupy.cuda.nvtx.RangePush`/`RangePop`), so it shows up in `nsys`/Nsight next to the kernels it launches; no-op on the NumPy backend or without NVTX. The range is popped when the block raises. +- `xp.timed_region(name, sync=True)` and `xp.Timing`: Context manager timing a code region including the device work it queues: on the CuPy backend it synchronizes before reading the clock (wall-clock timers around launches otherwise measure the launch, not the kernel), and pushes an NVTX range of the same name. Yields a `Timing` with `name`, `elapsed` (seconds) and `synced`. +- `xp.count_transfers()`: Context manager yielding a `TransferCounter` that records every host/device transfer made through cunumpy in the block, with the call site of each: `to_numpy`/`to_cunumpy` of a device array (`to_host`), `to_cupy`/`to_cunumpy` of a host array (`to_device`), `PyccelKernel` calls that copy device arrays to the host (`kernel_conversions`, one per call) and `Kernel` calls that fall back to the host kernel on the CuPy backend (`fallbacks`). `counter.total`, `counter.events` (`TransferEvent(kind, description, where)`) and `counter.report()` (events grouped by kind and call site) let a test verify that a time step does not transfer. Counters nest; calls that do not copy (e.g. `to_numpy` of a NumPy array) are not counted. Transfers that bypass cunumpy (raw `cupy.ndarray.get()`, `cupy.asarray(numpy_array)`, conversions inside other libraries) are not seen. +- `xp.assert_no_transfers()`: Context manager raising `AssertionError` with the counter's report if the block makes a transfer through cunumpy. +- CUDA debug mode: `CudaKernel(..., debug=True)`, `xp.set_cuda_debug(True)`, the context manager `xp.cuda_debug()` or the environment variable `CUNUMPY_CUDA_DEBUG=1` compile kernels with `-lineinfo` and `-DCUNUMPY_BOUNDS_CHECK` (`xp.DEBUG_OPTIONS`; `-G` is not available with NVRTC) and synchronize the stream after every launch, so an asynchronous CUDA error (illegal memory access, launch failure) is raised as a `RuntimeError` naming the kernel and its launch shape instead of surfacing at a later `.get()` or MPI call. `xp.get_cuda_debug()` reads the global setting, which applies to kernels created with `debug=None` at every launch; `CudaKernel.debug`, `debug_active()` and `compile_options()` expose a kernel's setting and the options a compilation uses. +- Header-aware compile cache: `xp.resolve_includes(source, include_dirs, base_dir=None)` lists the `#include "..."` files of a CUDA source recursively (resolved relative to the including file, then in `include_dirs`; cycles, system headers and missing files are ignored), and `xp.include_hash(paths)` is a short digest of their contents. `CudaKernel` exposes `included_headers`, `include_dirs`, `source_dir` (set by `from_file`, or the new `source_dir` keyword) and `compile_options()`. +- `cunumpy.testing`: Helpers for testing host/CUDA kernel pairs with pytest (pytest is imported only when its objects are used, never by `import cunumpy`). `requires_cupy` is a `skipif` marker for tests that need a GPU, `BACKENDS = ["numpy", pytest.param("cupy", marks=requires_cupy)]` parametrizes a test over the backends, and the `backend` fixture runs a test once per backend with that backend active. +- `cunumpy.testing.assert_kernels_agree(kernel, make_args, *, n_threads=..., rtol=1e-12, atol=0.0, n_calls=1, outputs=None, seed=0)`: Builds the arguments with `make_args(backend, seed)` on the NumPy and the CuPy backend, runs the host and the CUDA kernel of a `Kernel`, and compares the arrays they wrote (the declared `outputs`, or every array argument, including arrays held by argument objects) with `numpy.testing.assert_allclose`, naming the differing argument. Skips the test without a GPU and returns the host arrays. +- `KernelCatalog.parity_cases()`: The `(name, kernel)` pairs of the kernels with a CUDA version, so one test parametrised with them and `assert_kernels_agree` covers a whole catalog. +- `cunumpy.testing.device_function_kernel(header_source, signature, *, name=None, includes=(), n_threads_param="n")`: Generates an elementwise `extern "C" __global__` wrapper around a `__device__` function given its C prototype (pointer parameters are shared, scalar parameters become per-thread arrays, the return value goes into `out`), so device helpers can be tested from Python against their host versions without a hand-written test kernel. +- `xp.as_device_array(value, dtype=None, ndim=None, *, name=None)`: The "reference or copy once" rule for building CUDA argument objects: a CuPy array that already has the requested dtype and is C-contiguous is returned unchanged, anything else (tuples, lists, host arrays, other dtypes, non-contiguous views) becomes one C-contiguous device copy. Raises `RuntimeError` on the NumPy backend, so host data is never copied to the device implicitly, and `ValueError` if `ndim` does not match. +- `CudaKernel` pointer parameters and `CudaStruct` pointer fields now also reject non-C-contiguous arrays with `TypeError`: a kernel reads a pointer as a flat buffer, so a view such as `a[:, 0:3]` would give silently wrong results. +- `xp.KernelArguments` and `xp.resolve_host_args(args, kwargs=None)`: Argument objects with a host and a device form. An object whose type defines `__host_args__()` is replaced by its result (the object the host kernel receives, e.g. a Pyccel class of NumPy arrays) by `Kernel` on the host path and by `PyccelKernel` (so the `missing_cuda="fallback"` path works too); on the CuPy backend `CudaKernel` flattens the same object via `__cuda_args__()`. Owners can expose one `kernel_args` property that builds each form lazily, so call sites never branch on the backend and CPU runs never build device arguments. + ## [0.2.0] - 2026-09-28 ### Changed diff --git a/README.md b/README.md index 3df254f..090abdb 100644 --- a/README.md +++ b/README.md @@ -123,6 +123,21 @@ with xp.use_backend("cupy"): result = xp.to_numpy(filtered) # one transfer for a CPU-only consumer ``` +To verify that a block, such as a time step, makes no transfer at all, count +them: `count_transfers()` records every `to_numpy()`, `to_cupy()` and +`to_cunumpy()` call that actually copies, every `PyccelKernel` call that +converts device arrays, and every `Kernel` fallback to the host kernel, with +the call site of each. `assert_no_transfers()` raises with that report if +anything was counted. Only transfers made through CuNumpy are seen; raw +`cupy.ndarray.get()` or `cupy.asarray()` calls need a profiler such as `nsys`. + +```python +with xp.count_transfers() as counter: + propagator(dt) + +assert counter.total == 0, counter.report() +``` + ## Random numbers and dtypes `get_rng(seed)` returns a random generator for the active backend. NumPy and @@ -159,6 +174,34 @@ active CuPy device and `None` on NumPy. `set_device_for_rank(rank)` is a round-robin convenience for MPI layouts where local ranks map contiguously to GPUs. If your scheduler uses a different mapping, select the device directly. +For MPI programs with one rank per GPU, the startup sequence is: + +1. `bind_local_device()` selects the GPU from the node-local rank that the MPI + launcher exports (`local_rank()`) and creates its CUDA context. It runs + before MPI is initialized because a CUDA-aware MPI binds to the device that + is current at `MPI_Init`; without it, every rank of a node would use + device 0. +2. `from mpi4py import MPI` initializes MPI. +3. `require_cuda_aware_mpi()` (or `mpi_is_cuda_aware(comm)`) checks, with one + tiny device `Sendrecv` on every rank, that the MPI library can pass device + buffers at all. Passing CuPy arrays to a plain MPI build segfaults or + silently sends garbage; the check turns that into a clear error at + startup. It is a no-op on the NumPy backend. +4. `synchronize_for_mpi(*buffers)` before every MPI call with device buffers: + kernels run asynchronously, and MPI would otherwise send a buffer a kernel + is still writing, without an error. + +```python +xp.set_backend("cupy") +xp.bind_local_device() # before MPI_Init +from mpi4py import MPI # MPI_Init + +xp.require_cuda_aware_mpi() # once, on all ranks + +xp.synchronize_for_mpi(send, recv) +MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source) +``` + CuPy caches released allocations in memory pools. This can make process-level GPU memory appear occupied after arrays go out of scope. `free_memory()` asks CuPy to release currently free cached blocks; it does not free memory still @@ -179,6 +222,23 @@ xp.synchronize() result = xp.to_numpy(transformed) ``` +Because GPU work is asynchronous, a wall-clock timer around a kernel launch +measures the launch, not the kernel. `timed_region(name)` synchronizes the +device before reading the clock (on NumPy it is a plain timer), and +`nvtx_range(name)` marks a region so it shows up in `nsys`/Nsight; both are +no-ops or plain timers on NumPy, and `nvtx_range` also works as a decorator: + +```python +with xp.timed_region("fft") as timing: + transformed = xp.fft.fft(device) +print(timing.elapsed, timing.synced) + + +@xp.nvtx_range("step") +def step(dt): + ... +``` + ## Use NumPy-only kernels with CuPy arrays `PyccelKernel` adapts a callable that expects NumPy arrays. When conversion is @@ -213,6 +273,197 @@ 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. +## Write CUDA kernels next to host kernels + +`CudaKernel` wraps a CUDA C kernel (compiled with NVRTC through +`cupy.RawKernel`) so that it is called with the same arguments as the host +kernel it mirrors, plus the number of threads. Arrays are never copied: they +must be C-contiguous CuPy arrays. The `extern "C" __global__` signature is +parsed once and every call is checked against it: Python scalars are cast to +the declared C types, and a wrong argument count, an array of the wrong dtype +or a non-contiguous view, or a scalar that does not fit its type raises instead +of silently producing wrong values. + +`Kernel` pairs a host kernel with its CUDA kernel and calls the one matching +the active backend, so kernels can be ported to CUDA one at a time: + +```python +import cunumpy as xp + +AXPY = 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]; +} +""" + + +def axpy(a, x, y, n): # host version, e.g. compiled with Pyccel + for i in range(n): + y[i] += a * x[i] + + +kernel = xp.Kernel(axpy, xp.CudaKernel(AXPY, "axpy")) + +with xp.use_backend("cupy"): + x = xp.arange(1000, dtype=xp.float64) + y = xp.zeros(1000) + kernel(2.0, x, y, 1000, n_threads=1000) # runs the CUDA kernel +``` + +On the CuPy backend, a `Kernel` without CUDA kernel raises +`NotImplementedError` (or, with `missing_cuda="fallback"`, runs the host kernel +through `PyccelKernel`, with host copies; `host_options` configure that +`PyccelKernel`). `KernelCatalog.from_package()` collects kernel pairs from a +package with one folder per kernel (`name/name_kernels.py` and +`name/name_cuda.cu`), and `catalog.compile_all()` compiles all CUDA kernels at +setup. + +Groups of arguments can be passed as one: objects implementing +`__cuda_args__()` (see `CudaArguments`) are flattened into several kernel +arguments, and `CudaStruct` defines a C struct once (its C `declaration` and +the matching memory layout) and packs values into it, which the kernel takes +as one parameter: + +```python +Vec = xp.CudaStruct("Vec", [("data", "double*"), ("n", "int")]) +scale = xp.CudaKernel( + 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", + structs=[Vec], +) +scale(Vec(data=y, n=y.size), 0.5, n_threads=y.size) +``` + +When building such argument objects, `xp.as_device_array(value, dtype, +ndim=None)` applies the "reference or copy once" rule: a CuPy array that +already has the dtype and is C-contiguous is returned as it is, anything else +(a tuple such as `degree = (3, 3, 3)`, a host array, another dtype, a +non-contiguous view) becomes one device copy. Call it once when the object is +built, not per kernel call; on the NumPy backend it raises, so host data is +never copied to the device implicitly. +When the host kernel takes such a group as one object too (e.g. a Pyccel class +holding NumPy arrays), give the group both forms with `KernelArguments`: +`__host_args__()` returns the object for the host kernel, `__cuda_args__()` +the flattened device arguments. `Kernel` and `PyccelKernel` resolve +`__host_args__()` on the host path and `CudaKernel` flattens `__cuda_args__()` +on the CUDA path, so the call site is the same on both backends and each form +can be built lazily on first access (a CPU run never builds device arguments): + +```python +class ParticleArguments(xp.KernelArguments): + def __init__(self, markers): + self.markers = markers + self._host = None + + def __host_args__(self): + if self._host is None: + self._host = MarkerArguments(self.markers) # Pyccel class + return self._host + + def __cuda_args__(self): + return (self.markers, self.markers.shape[0]) + + +kernel(particles.kernel_args, dt, n_threads=n) # host or CUDA kernel + +Kernels ported from pyccel index arrays like `markers[ip, j]`, which needs +shapes and strides rather than bare pointers. The shipped header +`cunumpy/array_view.cuh` (found by every `CudaKernel`) provides the strided +views `Array1D` to `Array3D`; a parameter or struct field of that type +takes a CuPy array, contiguous or not, and indexes `a(i, j)`. The struct can be +generated from the annotations of the pyccel argument class, so the Python +class is the one definition, and written to a header that a test keeps in sync: + +```python +class MarkerArguments: + def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): + ... + +MarkerArgs = xp.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") +MarkerArgs.to_header("marker_args.cuh") # Array2D markers; long long n_markers; ... +push = xp.CudaKernel( + r""" + #include "marker_args.cuh" + #include + extern "C" __global__ void push(MarkerArgs m, double dt) { + CUNUMPY_THREAD_1D(ip, m.n_markers); + if (m.valid(ip)) m.markers(ip, 0) += dt * m.markers(ip, 3); + }""", + "push", + structs=[MarkerArgs], + include_dirs=["."], +) +push(MarkerArgs(markers=markers, n_markers=markers.shape[0], valid=valid), 0.1, + n_threads=markers.shape[0]) +``` + +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. + +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 +caused it. In debug mode, enabled with `xp.set_cuda_debug(True)`, the +context manager `xp.cuda_debug()`, `CudaKernel(..., 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 +the error is raised as a `RuntimeError` naming the kernel and its launch shape. +To find the faulting line and out-of-bounds accesses that do not crash, the +next step is NVIDIA's memory checker: +`CUNUMPY_CUDA_DEBUG=1 compute-sanitizer python -m pytest ...`. + +## Test kernel pairs + +`cunumpy.testing` helps to test the ports with pytest. `assert_kernels_agree` +builds the arguments on both backends, runs the host and the CUDA kernel and +compares the arrays they wrote; with `catalog.parity_cases()`, one +parametrised test covers every ported kernel of a catalog. `BACKENDS` and +`requires_cupy` parametrize tests over the backends, skipping CuPy without a +GPU, and `device_function_kernel` wraps a `__device__` helper in an elementwise +kernel so it can be checked against its host version without writing a test +kernel: + +```python +import pytest +from cunumpy.testing import assert_kernels_agree + + +def make_args(backend, seed): + x = xp.to_cunumpy(np.random.default_rng(seed).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) + +Accumulation kernels often write into a buffer that another library owns on +the host (a stencil vector's `_data`, exchanged over MPI). `DeviceMirror` +pairs that NumPy array with a device copy: `mirror.device` is the CuPy array +on the GPU and the host array itself on the CPU, `to_host()` copies back in +place (the host array keeps its identity) and `zero()` clears the buffer, so +the one transfer per accumulation is explicit. The shipped header +`cunumpy/atomic.cuh` (found automatically, see `cuda_include_dir()`) provides +`cunumpy_atomic_add()` and 2D/3D indexed variants for the many-threads-to-one-cell +writes: + +```python +mirror = xp.DeviceMirror(vector._data) +mirror.zero() +accumulate(markers, mirror.device, n_threads=n_markers) +mirror.to_host() # vector._data holds the result on both backends +``` + ## Pyodide CuNumpy supports the NumPy backend in Pyodide. It does not provide CuPy/CUDA diff --git a/docs/source/api.md b/docs/source/api.md index f7d63e3..645a1de 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -172,6 +172,103 @@ 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. +## 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. + +### `count_transfers()` + +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.count_transfers() as counter: + propagator(dt) + +assert counter.total == 0, counter.report() +``` + +Four kinds of events are recorded: + +* `to_host`: `to_numpy()` (or `to_cunumpy()`) called with a CuPy array; +* `to_device`: `to_cupy()` (or `to_cunumpy()`) called with anything that is not + a CuPy array already; +* `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. The host copies it makes are counted as one `kernel_conversion` + event in addition. + +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`, +`events` (a list of `TransferEvent(kind, description, where)`, where `where` +is the `file:line` of the caller outside CuNumpy) and +`kernel_conversion_calls` (the `kernel_conversion` events). `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) + 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. +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)`, +`float(device_array)`, and implicit conversions inside other libraries are +not counted. Use `nsys` (or CuPy's profiling hooks) to find those. + +### `assert_no_transfers()` + +Context manager that raises `AssertionError` with the counter's `report()` if +the block makes a transfer through CuNumpy. 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.assert_no_transfers(): + propagator(dt) +### `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 becomes one C-contiguous device copy, + `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. + +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.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]) +``` + ## Random numbers and dtype ### `get_rng(seed=None)` @@ -227,6 +324,98 @@ mapping differs: device_id = xp.set_device_for_rank(mpi_rank) ``` +### `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`. + +### `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.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_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. + +### `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.bind_local_device() # 1. select the GPU, before MPI_Init +from mpi4py import MPI # 2. MPI_Init, on the bound device + +xp.require_cuda_aware_mpi() # 3. clear error instead of a segfault later + +xp.synchronize_for_mpi(send, recv) # 4. before every MPI call with device buffers +MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source) +``` + +### `synchronize_for_mpi(*arrays)` + +Waits for the work pending on the current stream if at least one of `arrays` +is a CuPy array; `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.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. + ### `memory_info()` Returns `(free_bytes, total_bytes)` reported by the CUDA runtime for the @@ -240,6 +429,14 @@ 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_include_dir()` + +Returns the directory (as `str`) of the CUDA headers shipped with CuNumpy, +currently `cunumpy/atomic.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. + ### `pin_memory(array)` Copies a host array to page-locked (pinned) host memory. Pinned memory can @@ -269,6 +466,57 @@ 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. +## 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. + +### `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.nvtx_range("push markers"): + kernel(markers, dt, n_threads=n) + + +@xp.nvtx_range("accumulate") +def accumulate(particles, grid): + ... +``` + +### `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.timed_region("push markers") as timing: + kernel(markers, dt, n_threads=n) + +print(f"{timing.name}: {timing.elapsed:.4f} s, synced={timing.synced}") +``` + +### `Timing` + +Dataclass returned by `timed_region()`, with the fields `name` (`str`), +`elapsed` (`float`, `None` until the block exits) and `synced` (`bool`). + ## `PyccelKernel` ### Constructor @@ -353,6 +601,860 @@ 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. +## `CudaKernel` + +### Constructor + +```python +xp.CudaKernel( + source, + name, + *, + block_size=128, + options=(), + include_dirs=(), + source_dir=None, + structs=(), + template_args=None, + check_signature=True, + debug=None, +) +xp.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs) +xp.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. + +`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_kernel_names(source)` lists the `__global__` +functions of a source string (ignoring comments). + +```python +kernels = xp.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`, `structs`, +`template_args`, `signature`, `is_compiled`, `debug`. +as `"scale"`), `source`, `block_size`, `options`, `include_dirs`, +`source_dir`, `included_headers`, `structs`, `template_args`, `signature`, +`is_compiled`. + +### 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.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, like NVRTC does. System headers in angle brackets + 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 quoted includes never touch the file system. + +The two building blocks are available on their own: + +* `xp.resolve_includes(source, include_dirs=(), *, base_dir=None)`: the + resolved header paths of a source, as a list. +* `xp.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. + +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)` returns the +`(grid, block)` a call would use, e.g. to size a per-block output: + +```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.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`). + +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.ctype_of(dtype)` gives the C +type of a dtype (`xp.ctype_of(np.float64) == "double"`), e.g. to generate +source. `xp.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.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_include_dir()` returns their directory (a `str`) for other compilers +(`-I`). + +`cunumpy/array_view.cuh` defines the strided views `Array1D`, `Array2D` +and `Array3D`: `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.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.CudaKernelVariants( + lambda ndim, dtype: xp.CudaKernel(make_source(ndim, xp.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.set_cuda_debug(enabled) +xp.get_cuda_debug() +xp.cuda_debug(enabled=True) # context manager +xp.CudaKernel(..., debug=None) +kernel.debug_active() +kernel.compile_options() +xp.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.set_cuda_debug(True)`, temporarily +with the context manager `xp.cuda_debug()`, or before starting Python with +the environment variable `CUNUMPY_CUDA_DEBUG=1` (`true`, `yes` and `on` work +too); `xp.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 `compile()` 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_debug(): + kernel = xp.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. + +## `CudaStruct` + +```python +Particles = xp.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.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 `Array3D` 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. +* 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.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`. + +### Generating headers + +```python +xp.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.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.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`). + +## `CudaArguments` + +```python +class Particles(xp.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. This lets an +application keep its host argument objects (e.g. Pyccel classes holding NumPy +arrays) and matching device argument objects that reference the same data on +the device, and pass either to the same call. A `CudaArguments` object may also +return struct values (`CudaStructValue.packed`) among its values. + +## `KernelArguments` + +```python +class ParticleArguments(xp.KernelArguments): + def __init__(self, particles): + self._particles = particles + self._host = None + self._cuda = None + + def __host_args__(self): + if self._host is None: # e.g. a Pyccel class holding NumPy arrays + self._host = MarkerArguments(self._particles.markers) + return self._host + + def __cuda_args__(self): + if self._cuda is None: # device arrays and scalars, flattened + markers = self._particles.markers + self._cuda = (markers, markers.shape[0], markers.shape[1]) + return self._cuda + + +class Particles: + @property + def kernel_args(self): + if self._kernel_args is None: + self._kernel_args = ParticleArguments(self) + return self._kernel_args + + +push(particles.kernel_args, dt, n_threads=n) # same call on both backends +``` + +Base class for argument objects that have a host form and a device form. A +group of arrays, e.g. the marker data of a particle species, is typically +passed to the host kernel as one object holding NumPy arrays (a Pyccel class) +and to the CUDA kernel as several device arrays and scalars. `KernelArguments` +lets one object stand for both, so a `Kernel` call never branches on the +backend: + +* `__host_args__()` returns the single object the host kernel receives in that + position. `Kernel` (on the NumPy backend) and `PyccelKernel` (always, so the + `missing_cuda="fallback"` path works with the same objects) replace the + argument by this value. +* `__cuda_args__()` returns the tuple of CUDA kernel arguments the object + stands for, the `CudaArguments` protocol above; `CudaKernel` flattens it. + +Only top-level positional and keyword arguments are resolved, not objects +nested in tuples, lists or dicts. The check is made on the type, like for +`__cuda_args__`: an instance attribute named `__host_args__` (e.g. a stored +object) is not treated as the protocol. Subclassing is optional; both methods +of the base class raise `NotImplementedError`, so a subclass overrides the ones +it supports (a `KernelArguments` without `__cuda_args__` raises when it reaches +a `CudaKernel`). + +In the example above both forms are built lazily on first access and cached, +so a CPU run never builds device arguments and a GPU run never builds the host +object. The owner is responsible for invalidating the cache (setting the +stored forms to `None`, or replacing the `ParticleArguments` object) when its +arrays are replaced, e.g. after resizing, `deepcopy` or unpickling. + +### `resolve_host_args(args, kwargs=None)` + +Returns `(args, kwargs)` with every top-level argument whose type defines a +callable `__host_args__()` replaced by its result; everything else is passed +through untouched. `Kernel` and `PyccelKernel` call it before the host kernel; +it is exported for code that calls host kernels by other means: + +```python +args, kwargs = xp.resolve_host_args((particles.kernel_args, dt), {"out": out}) +host_push(*args, **kwargs) +``` + +## `Kernel` + +```python +xp.Kernel( + host_kernel, + cuda_kernel=None, + *, + name=None, + missing_cuda="raise", + cuda_path=None, + host_options=None, +) +``` + +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 (`n_threads` or `grid` is required there) and ignored by the host +kernel. Arguments implementing `KernelArguments` are replaced by their +`__host_args__()` on the host path and flattened via `__cuda_args__()` on the +CUDA path. `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. + +Properties: `name`, `host_kernel`, `cuda_kernel`, `has_cuda`, `missing_cuda`, +`cuda_path`. + +## `KernelCatalog` + +```python +catalog = xp.KernelCatalog.from_package( + package, + *, + host_suffix="_kernels", + cuda_suffix="_cuda.cu", + missing_cuda="raise", + host_options=None, + include_dirs=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.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`. +* `cuda_options`: passed on to `CudaKernel.from_file`, e.g. `block_size` or + `structs`. + +`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. + +## Testing utilities + +```python +from cunumpy.testing import ( + BACKENDS, + assert_kernels_agree, + device_function_kernel, + requires_cupy, +) +``` + +`cunumpy.testing` holds helpers for testing kernels with pytest. It is not +imported by `import cunumpy` (so `xp.testing` remains NumPy's or CuPy's +`testing` module until `cunumpy.testing` is imported), and it imports pytest +only when one of its pytest objects is used, so `device_function_kernel` works +without pytest. + +### `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.testing import backend`) or the +test module, then take `backend` as a test argument. + +### `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). + +## `DeviceMirror` + +```python +mirror = xp.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()`: 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. + +The transfers are explicit so that one per accumulation is visible and +bounded: + +```python +mirror = xp.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 indexed helpers (also for `float`) 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. + ## Version `xp.__version__` is the installed package version. When package metadata is diff --git a/pyproject.toml b/pyproject.toml index 794f7b3..b76dce0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,22 +5,21 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.2.0" +version = "0.4.0" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] license = { file = "LICENSE.txt" } authors = [ { name = "Max" } ] -requires-python = ">=3.8" +requires-python = ">=3.10" classifiers = [ "Development Status :: 3 - Alpha", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", ] dependencies = [ "array-api-compat", @@ -52,7 +51,7 @@ urls."Source" = "https://github.com/max-models/cunumpy" where = [ "src" ] [tool.setuptools.package-data] -cunumpy = [ "py.typed", "*.pyi" ] +cunumpy = [ "py.typed", "*.pyi", "cuda/include/cunumpy/*.cuh" ] [tool.isort] profile = "black" diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 4e9e194..21e0016 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -2,9 +2,37 @@ from importlib.metadata import PackageNotFoundError, version from . import xp -from .kernel import PyccelKernel +from .cuda_kernel import ( + DEBUG_OPTIONS, + CudaArguments, + CudaKernel, + CudaKernelVariants, + CudaParameter, + CudaStruct, + CudaStructValue, + ctype_of, + cuda_include_dir, + cuda_kernel_names, + include_hash, + parse_cuda_signature, + resolve_includes, + write_cuda_header, +) +from .dispatch import Kernel, KernelCatalog +from .kernel import KernelArguments, PyccelKernel, resolve_host_args +from .mirror import DeviceMirror +from .transfers import ( + TransferCounter, + TransferEvent, + assert_no_transfers, + count_transfers, +) from .xp import ( + Timing, + as_device_array, assert_same_backend, + bind_local_device, + cuda_debug, cupy_available, default_float_dtype, device_count, @@ -12,17 +40,25 @@ get_array_backend, get_array_module, get_backend, + get_cuda_debug, get_rng, is_cpu, is_gpu, + local_rank, memory_info, + mpi_is_cuda_aware, + nvtx_range, pin_memory, + require_cuda_aware_mpi, same_backend, set_backend, + set_cuda_debug, set_device, set_device_for_rank, stream, synchronize, + synchronize_for_mpi, + timed_region, to_cunumpy, to_cupy, to_numpy, @@ -35,9 +71,31 @@ __version__ = "0.0.0+unknown" __all__ = [ + "DEBUG_OPTIONS", + "CudaArguments", + "CudaKernel", + "CudaKernelVariants", + "CudaParameter", + "CudaStruct", + "CudaStructValue", + "DeviceMirror", + "Kernel", + "KernelArguments", + "KernelCatalog", "PyccelKernel", + "Timing", + "TransferCounter", + "TransferEvent", "__version__", + "as_device_array", + "assert_no_transfers", "assert_same_backend", + "bind_local_device", + "count_transfers", + "ctype_of", + "cuda_debug", + "cuda_include_dir", + "cuda_kernel_names", "cupy_available", "cupy_backend", "default_float_dtype", @@ -46,22 +104,35 @@ "get_array_backend", "get_array_module", "get_backend", + "get_cuda_debug", "get_rng", + "include_hash", "is_cpu", "is_gpu", + "local_rank", "memory_info", + "mpi_is_cuda_aware", "numpy_backend", + "nvtx_range", + "parse_cuda_signature", "pin_memory", + "require_cuda_aware_mpi", + "resolve_host_args", + "resolve_includes", "same_backend", "set_backend", + "set_cuda_debug", "set_device", "set_device_for_rank", "stream", "synchronize", + "synchronize_for_mpi", + "timed_region", "to_cunumpy", "to_cupy", "to_numpy", "use_backend", + "write_cuda_header", "xp", ] diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index d71893d..4f5416a 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -1,14 +1,30 @@ # Stub file for Pylance/mypy: exposes all numpy symbols so that # `import cunumpy as xp` followed by `xp.` shows numpy completions. # At runtime the real __init__.py dispatches to numpy or cupy via __getattr__. +from collections.abc import Generator from contextlib import contextmanager -from typing import Any, Generator +from typing import Any import numpy as np from numpy import * from . import xp as xp +from .cuda_kernel import CudaArguments as CudaArguments +from .cuda_kernel import CudaKernel as CudaKernel +from .cuda_kernel import CudaKernelVariants as CudaKernelVariants +from .cuda_kernel import CudaParameter as CudaParameter +from .cuda_kernel import CudaStruct as CudaStruct +from .cuda_kernel import CudaStructValue as CudaStructValue +from .cuda_kernel import ctype_of as ctype_of +from .cuda_kernel import cuda_include_dir as cuda_include_dir +from .cuda_kernel import parse_cuda_signature as parse_cuda_signature +from .cuda_kernel import write_cuda_header as write_cuda_header +from .dispatch import Kernel as Kernel +from .dispatch import KernelCatalog as KernelCatalog from .kernel import PyccelKernel as PyccelKernel +from .mirror import DeviceMirror as DeviceMirror +from .transfers import TransferCounter as TransferCounter +from .transfers import TransferEvent as TransferEvent def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... @@ -26,6 +42,9 @@ def use_backend(backend: str) -> Generator[None]: ... def set_backend(backend: str) -> None: ... def set_device(device_id: int) -> None: ... def set_device_for_rank(rank: int, devices_per_node: int | None = ...) -> int: ... +def local_rank() -> int: ... +def bind_local_device() -> int | None: ... +def synchronize_for_mpi(*arrays: Any) -> None: ... def device_count() -> int: ... def memory_info() -> tuple[int, int] | None: ... def free_memory() -> None: ... @@ -35,6 +54,10 @@ def stream() -> Generator[Any]: ... def get_rng(seed: int | None = ...) -> Any: ... def default_float_dtype() -> Any: ... def synchronize() -> None: ... +@contextmanager +def count_transfers() -> Generator[TransferCounter]: ... +@contextmanager +def assert_no_transfers() -> Generator[TransferCounter]: ... numpy_backend: bool cupy_backend: bool diff --git a/src/cunumpy/cuda/include/cunumpy/array_view.cuh b/src/cunumpy/cuda/include/cunumpy/array_view.cuh new file mode 100644 index 0000000..af3efb9 --- /dev/null +++ b/src/cunumpy/cuda/include/cunumpy/array_view.cuh @@ -0,0 +1,105 @@ +// Strided array views for CUDA kernels, passed by value from Python. +// +// Array1D, Array2D and Array3D describe a (possibly non-contiguous) +// device array the way NumPy/CuPy do: a data pointer, a shape and strides. +// Strides are in ELEMENTS, not bytes, so that `a(i, j)` is +// `data[i * strides[0] + j * strides[1]]`. Elements are accessed with +// `operator()`, which mirrors the `a[i, j]` indexing of the pyccel kernels +// being ported. +// +// The views are created on the Python side by cunumpy (a kernel parameter or a +// CudaStruct field of type `Array2D` takes a CuPy array). They are +// passed by value, so the memory layout must be exactly, for ndim = 1, 2, 3: +// +// T* data; // 8 bytes +// long long shape[ndim]; // ndim * 8 bytes +// long long strides[ndim]; // ndim * 8 bytes +// +// with 8-byte alignment and no padding (sizeof == 8 * (1 + 2 * ndim)). Do not +// add data members, virtual functions or a base class; the static_asserts at +// the end of this file check the size. Member functions do not change the +// layout. +// +// Bounds checks: compile with -DCUNUMPY_BOUNDS_CHECK to check every index +// against the shape (an out-of-bounds index prints a message and traps the +// kernel, which CuPy reports as a CUDA error). Without the macro, indexing is +// unchecked. + +#ifndef CUNUMPY_ARRAY_VIEW_CUH +#define CUNUMPY_ARRAY_VIEW_CUH + +#ifdef CUNUMPY_BOUNDS_CHECK +#define CUNUMPY_CHECK_INDEX(index, axis, extent) \ + do { \ + if ((index) < 0 || (index) >= (extent)) { \ + printf("cunumpy: index %lld is out of bounds for axis %d with size " \ + "%lld\n", \ + (long long)(index), (int)(axis), (long long)(extent)); \ + __trap(); \ + } \ + } while (0) +#else +#define CUNUMPY_CHECK_INDEX(index, axis, extent) ((void)0) +#endif + +template +struct Array1D { + T* data; + long long shape[1]; + long long strides[1]; + + __device__ __forceinline__ T& operator()(long long i) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + return data[i * strides[0]]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { return shape[0]; } +}; + +template +struct Array2D { + T* data; + long long shape[2]; + long long strides[2]; + + __device__ __forceinline__ T& operator()(long long i, long long j) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + CUNUMPY_CHECK_INDEX(j, 1, shape[1]); + return data[i * strides[0] + j * strides[1]]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { + return shape[0] * shape[1]; + } +}; + +template +struct Array3D { + T* data; + long long shape[3]; + long long strides[3]; + + __device__ __forceinline__ T& operator()(long long i, long long j, + long long k) const { + CUNUMPY_CHECK_INDEX(i, 0, shape[0]); + CUNUMPY_CHECK_INDEX(j, 1, shape[1]); + CUNUMPY_CHECK_INDEX(k, 2, shape[2]); + return data[i * strides[0] + j * strides[1] + k * strides[2]]; + } + + // Number of elements. + __device__ __forceinline__ long long size() const { + return shape[0] * shape[1] * shape[2]; + } +}; + +// The layout the Python side packs: pointer, shape, strides, 8-byte aligned. +static_assert(sizeof(Array1D) == 24, "unexpected Array1D layout"); +static_assert(sizeof(Array2D) == 40, "unexpected Array2D layout"); +static_assert(sizeof(Array3D) == 56, "unexpected Array3D layout"); +static_assert(sizeof(Array1D) == 24, "unexpected Array1D layout"); +static_assert(alignof(Array2D) == 8, "unexpected Array2D alignment"); + +#endif // CUNUMPY_ARRAY_VIEW_CUH diff --git a/src/cunumpy/cuda/include/cunumpy/atomic.cuh b/src/cunumpy/cuda/include/cunumpy/atomic.cuh new file mode 100644 index 0000000..43d1e0f --- /dev/null +++ b/src/cunumpy/cuda/include/cunumpy/atomic.cuh @@ -0,0 +1,69 @@ +// cunumpy/atomic.cuh: atomic accumulation helpers for CUDA kernels. +// +// Many threads adding into a few cells (particle-to-grid accumulation) must +// use atomics or lose updates. These helpers wrap atomicAdd for double and +// float, and index C-contiguous 2D/3D arrays passed as bare pointers plus +// their trailing extents. +// +// atomicAdd(double*, double) is a hardware instruction from compute +// capability 6.0 (sm_60) on; for older devices a compare-and-swap loop is +// used, which is correct but slow. +// +// The header directory is added to every CudaKernel's NVRTC options, so: +// #include + +#ifndef CUNUMPY_ATOMIC_CUH +#define CUNUMPY_ATOMIC_CUH + +// Add v to *p atomically; returns the old value of *p. +__device__ __forceinline__ double cunumpy_atomic_add(double* p, double v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600) + unsigned long long* address = reinterpret_cast(p); + unsigned long long old = *address; + unsigned long long assumed; + do { + assumed = old; + old = atomicCAS(address, assumed, + __double_as_longlong(v + __longlong_as_double(assumed))); + } while (assumed != old); + return __longlong_as_double(old); +#else + return atomicAdd(p, v); +#endif +} + +__device__ __forceinline__ float cunumpy_atomic_add(float* p, float v) +{ + return atomicAdd(p, v); +} + +// data[i, j] += v for a C-contiguous array of shape (n0, n1). +__device__ __forceinline__ double cunumpy_atomic_add_2d( + double* data, long long n1, long long i, long long j, double v) +{ + return cunumpy_atomic_add(data + i * n1 + j, v); +} + +__device__ __forceinline__ float cunumpy_atomic_add_2d( + float* data, long long n1, long long i, long long j, float v) +{ + return cunumpy_atomic_add(data + i * n1 + j, v); +} + +// data[i, j, k] += v for a C-contiguous array of shape (n0, n1, n2). +__device__ __forceinline__ double cunumpy_atomic_add_3d( + double* data, long long n1, long long n2, + long long i, long long j, long long k, double v) +{ + return cunumpy_atomic_add(data + (i * n1 + j) * n2 + k, v); +} + +__device__ __forceinline__ float cunumpy_atomic_add_3d( + float* data, long long n1, long long n2, + long long i, long long j, long long k, float v) +{ + return cunumpy_atomic_add(data + (i * n1 + j) * n2 + k, v); +} + +#endif // CUNUMPY_ATOMIC_CUH diff --git a/src/cunumpy/cuda/include/cunumpy/index.cuh b/src/cunumpy/cuda/include/cunumpy/index.cuh new file mode 100644 index 0000000..c06c308 --- /dev/null +++ b/src/cunumpy/cuda/include/cunumpy/index.cuh @@ -0,0 +1,52 @@ +// Thread-index helpers for CUDA kernels. +// +// #include +// +// extern "C" __global__ void axpy(double a, const double* x, double* y, +// long long n) { +// CUNUMPY_THREAD_1D(i, n); // long long i; returns if i >= n +// y[i] += a * x[i]; +// } +// +// Indices are `long long`, so that they can address more than 2^31 elements +// and match the `shape`/`strides` of the views in array_view.cuh. + +#ifndef CUNUMPY_INDEX_CUH +#define CUNUMPY_INDEX_CUH + +// Global thread index along x (y, z); one value per thread. +#define CUNUMPY_GLOBAL_INDEX_X() \ + ((long long)blockDim.x * blockIdx.x + threadIdx.x) +#define CUNUMPY_GLOBAL_INDEX_Y() \ + ((long long)blockDim.y * blockIdx.y + threadIdx.y) +#define CUNUMPY_GLOBAL_INDEX_Z() \ + ((long long)blockDim.z * blockIdx.z + threadIdx.z) + +// Declares `long long i` as the global thread index and returns from the +// kernel if `i >= n`. Use with `n_threads=n` in Python. +#define CUNUMPY_THREAD_1D(i, n) \ + long long i = CUNUMPY_GLOBAL_INDEX_X(); \ + if (i >= (long long)(n)) return + +// Likewise in 2 and 3 dimensions, for `n_threads=(ni, nj)` / `(ni, nj, nk)`. +#define CUNUMPY_THREAD_2D(i, j, ni, nj) \ + long long i = CUNUMPY_GLOBAL_INDEX_X(); \ + long long j = CUNUMPY_GLOBAL_INDEX_Y(); \ + if (i >= (long long)(ni) || j >= (long long)(nj)) return + +#define CUNUMPY_THREAD_3D(i, j, k, ni, nj, nk) \ + long long i = CUNUMPY_GLOBAL_INDEX_X(); \ + long long j = CUNUMPY_GLOBAL_INDEX_Y(); \ + long long k = CUNUMPY_GLOBAL_INDEX_Z(); \ + if (i >= (long long)(ni) || j >= (long long)(nj) || \ + k >= (long long)(nk)) return + +// Grid-stride loop over `i` in [0, n): every thread handles several elements, +// so any launch size works (e.g. `grid=` a fixed number of blocks in Python): +// +// CUNUMPY_GRID_STRIDE_1D(i, n) { y[i] += a * x[i]; } +#define CUNUMPY_GRID_STRIDE_1D(i, n) \ + for (long long i = CUNUMPY_GLOBAL_INDEX_X(); i < (long long)(n); \ + i += (long long)blockDim.x * gridDim.x) + +#endif // CUNUMPY_INDEX_CUH diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py new file mode 100644 index 0000000..7b35593 --- /dev/null +++ b/src/cunumpy/cuda_kernel.py @@ -0,0 +1,1729 @@ +"""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; +* 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.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. +""" + +from __future__ import annotations + +import hashlib +import inspect +import math +import os +import re +import typing +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +from pathlib import Path +from typing import Any, NamedTuple + +import numpy as np + +__all__ = [ + "DEBUG_OPTIONS", + "CudaArguments", + "CudaKernel", + "CudaKernelVariants", + "CudaParameter", + "CudaStruct", + "CudaStructValue", + "ctype_of", + "cuda_include_dir", + "cuda_kernel_names", + "include_hash", + "parse_cuda_signature", + "resolve_includes", + "write_cuda_header", +] + +# CUDA limit on the number of threads per block +_MAX_THREADS_PER_BLOCK = 1024 + +# Headers shipped with cunumpy: #include etc. +_CUDA_INCLUDE_DIR = Path(__file__).resolve().parent / "cuda" / "include" +_ARRAY_VIEW_INCLUDE = '#include "cunumpy/array_view.cuh"' + + +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``, + ``Array2D``, ``Array3D`` views passed by value) and + ``#include "cunumpy/index.cuh"`` (thread-index and grid-stride macros such + as ``CUNUMPY_THREAD_1D(i, n)``). Pass it as ``-I`` to other compilers. + """ + 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.) +DEBUG_OPTIONS = ("-lineinfo", "-DCUNUMPY_BOUNDS_CHECK") + + +class CudaArguments: + """Base class for objects passed to a :class:`CudaKernel` as one argument. + + 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. + + Parameters + ---------- + *values + The CUDA kernel arguments this object stands for: CuPy arrays and + scalars, in the order of the kernel signature. + + Examples + -------- + >>> class Particles(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 + """ + + 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 self._cuda_args + + +class CudaParameter(NamedTuple): + """One parameter of a CUDA kernel signature (or one field of a struct). + + Attributes + ---------- + name : str + Parameter name. + ctype : str + Normalized 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*``. + pointer : bool + Whether the parameter is a pointer (a device array). + struct : CudaStruct | None + The struct type, for a struct passed by value. + view_ndim : int | None + The number of dimensions, for an array view (``Array1D`` to + ``Array3D``, see :func:`cuda_include_dir`) passed by value. + """ + + name: str + ctype: str + dtype: np.dtype | None + pointer: bool + struct: CudaStruct | None = None + view_ndim: int | None = None + + +# C types (after removing qualifiers) and their NumPy dtypes. ``long`` is 64 bit, +# as on Linux (LP64), the platform CUDA runs on in practice. +_CTYPES = { + "bool": np.bool_, + "char": np.int8, + "signed char": np.int8, + "unsigned char": np.uint8, + "short": np.int16, + "short int": np.int16, + "unsigned short": np.uint16, + "unsigned short int": np.uint16, + "int": np.int32, + "signed": np.int32, + "signed int": np.int32, + "unsigned": np.uint32, + "unsigned int": np.uint32, + "long": np.int64, + "long int": np.int64, + "long long": np.int64, + "long long int": np.int64, + "unsigned long": np.uint64, + "unsigned long int": np.uint64, + "unsigned long long": np.uint64, + "unsigned long long int": np.uint64, + "int8_t": np.int8, + "int16_t": np.int16, + "int32_t": np.int32, + "int64_t": np.int64, + "uint8_t": np.uint8, + "uint16_t": np.uint16, + "uint32_t": np.uint32, + "uint64_t": np.uint64, + "size_t": np.uint64, + "ptrdiff_t": np.int64, + "ssize_t": np.int64, + "float": np.float32, + "double": np.float64, + "complex": np.complex64, + "complex": np.complex128, +} + +# The C type used for each NumPy dtype, see ctype_of() +_CTYPE_OF = { + np.dtype(np.bool_): "bool", + np.dtype(np.int8): "signed char", + np.dtype(np.uint8): "unsigned char", + np.dtype(np.int16): "short", + np.dtype(np.uint16): "unsigned short", + np.dtype(np.int32): "int", + np.dtype(np.uint32): "unsigned int", + np.dtype(np.int64): "long long", + np.dtype(np.uint64): "unsigned long long", + np.dtype(np.float32): "float", + np.dtype(np.float64): "double", + np.dtype(np.complex64): "complex", + np.dtype(np.complex128): "complex", +} + +_QUALIFIERS = {"const", "volatile", "__restrict__", "__restrict", "restrict"} + +_COMPLEX = re.compile(r"(?:(?:thrust|cuda::std)::)?complex\s*<\s*(float|double)\s*>") +# Array1D to Array3D (cunumpy/array_view.cuh), T a scalar type of _CTYPES +_VIEW = re.compile(r"\bArray([123])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>") +_TOKEN = re.compile( + r"Array[123]D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>" + r"|[A-Za-z_]\w*|\*|\[\s*\]" +) + + +def _normalize_view(match: re.Match) -> str: + """``Array2D< const double >`` -> ``Array2D``.""" + words = [w for w in match.group(2).split() if w not in _QUALIFIERS] + return f"Array{match.group(1)}D<{' '.join(words)}>" + + +def _view_dtype(ndim: int) -> np.dtype: + """The structured dtype with the C layout of ``ArrayD``.""" + return np.dtype( + [ + ("data", np.uint64), + ("shape", np.int64, (ndim,)), + ("strides", np.int64, (ndim,)), + ], + align=True, + ) + + +def ctype_of(dtype: Any) -> str: + """The C type of a NumPy dtype, e.g. ``ctype_of(np.float64) == "double"``. + + Useful to generate CUDA source or template arguments for a given dtype. + Complex dtypes map to ``complex``/``complex`` (include + ```` in the source). + """ + try: + return _CTYPE_OF[np.dtype(dtype)] + except (KeyError, TypeError): + raise ValueError(f"no C type for dtype {dtype!r}") from None + + +def _strip_comments(source: str) -> str: + source = re.sub(r"/\*.*?\*/", " ", source, flags=re.DOTALL) + return re.sub(r"//[^\n]*", " ", source) + + +# ``#include "name"``: quoted includes are the project's own headers. Angle +# bracket includes are system headers and are not tracked. +_QUOTED_INCLUDE = re.compile(r'^[ \t]*#[ \t]*include[ \t]*"([^"\n]+)"', re.MULTILINE) + + +def _quoted_includes(source: str) -> list[str]: + return _QUOTED_INCLUDE.findall(_strip_comments(source)) + + +def resolve_includes( + source: str, + include_dirs: Iterable[str | Path] = (), + *, + base_dir: str | Path | None = None, +) -> list[Path]: + """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`, in order. Found headers are + scanned in turn, relative to their own directory. Includes in angle + brackets (system headers) and includes that cannot be found are ignored; + NVRTC reports the latter when the kernel is compiled. + + 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. + + Returns + ------- + list[Path] + The resolved header files, each once, in order of first inclusion + (depth first). Empty if the source has no quoted includes; the file + system is not touched in that case. + """ + dirs = tuple(Path(d) for d in include_dirs) + found: list[Path] = [] + seen: set[Path] = set() + + def visit(code: str, directory: Path | None) -> None: + for name in _quoted_includes(code): + candidates = [directory / name] if directory is not None else [] + candidates += [d / name for d in dirs] + for candidate in candidates: + if candidate.is_file(): + path = candidate.resolve() + if path not in seen: + seen.add(path) + found.append(candidate) + visit(path.read_text(errors="replace"), path.parent) + break + + visit(source, None if base_dir is None else Path(base_dir)) + return found + + +def include_hash(paths: Iterable[str | Path]) -> str: + """A short hex digest of the contents of `paths`, 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`. + + Parameters + ---------- + paths : Iterable[str | Path] + Files to hash, e.g. from :func:`resolve_includes`. + + Returns + ------- + str + The first 16 hex digits of the SHA-256 digest. + """ + digest = hashlib.sha256() + for path in paths: + content = Path(path).read_bytes() + digest.update(len(content).to_bytes(8, "little")) + digest.update(content) + return digest.hexdigest()[:16] + + +_GLOBAL_FUNCTION = re.compile(r"__global__\s+void\s+([A-Za-z_]\w*)\s*\(") + + +def cuda_kernel_names(source: str) -> list[str]: + """The names of the ``__global__`` functions defined in `source`, in order. + + Comments are ignored. Templates are included; a function declared more + than once (e.g. a forward declaration) is listed once. + + Parameters + ---------- + source : str + CUDA C source code. + + Returns + ------- + list[str] + The kernel names, in the order of their first appearance. + """ + return list(dict.fromkeys(_GLOBAL_FUNCTION.findall(_strip_comments(source)))) + + +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`. + """ + if jobs is None: + jobs = os.cpu_count() or 1 + if jobs < 1: + raise ValueError(f"jobs must be positive or None, got {jobs}") + if jobs == 1 or len(compilers) <= 1: + for compile in compilers.values(): + compile() + return list(compilers) + + with ThreadPoolExecutor(max_workers=min(jobs, len(compilers))) as pool: + futures = {name: pool.submit(compile) for name, compile in compilers.items()} + compiled, error = [], None + for name, future in futures.items(): + exc = future.exception() + if exc is None: + compiled.append(name) + elif error is None: + error = exc + if error is not None: + raise error + return compiled + + +def _parse_parameter( + text: str, structs: dict[str, CudaStruct] | None = None +) -> CudaParameter: + text = _COMPLEX.sub(lambda m: f"complex<{m.group(1)}>", text) + text = _VIEW.sub(_normalize_view, text) + tokens = _TOKEN.findall(text) + pointers = sum(1 for t in tokens if t == "*" or t.startswith("[")) + words = [t for t in tokens if t != "*" and not t.startswith("[")] + words = [t for t in words if t not in _QUALIFIERS] + if words[:1] == ["struct"]: + words = words[1:] + if len(words) < 2: + raise ValueError(f"cannot parse the kernel parameter {text.strip()!r}") + name, ctype = words[-1], " ".join(words[:-1]) + + if structs and ctype in structs: + if pointers: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: structs can " + "only be passed by value" + ) + struct = structs[ctype] + return CudaParameter(name, ctype, struct.dtype, False, struct) + view = _VIEW.fullmatch(ctype) + if view is not None: + element = view.group(2) + if pointers or element not in _CTYPES: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: array views " + f"take a scalar element type and are passed by value" + ) + ndim = int(view.group(1)) + return CudaParameter(name, ctype, np.dtype(_CTYPES[element]), False, None, ndim) + if ctype == "void" and pointers == 1: + return CudaParameter(name, ctype, None, True) + if pointers > 1 or ctype not in _CTYPES: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: unsupported type " + f"{ctype + '*' * pointers!r}" + ) + return CudaParameter(name, ctype, np.dtype(_CTYPES[ctype]), pointers == 1) + + +def _split_top_level(text: str) -> list[str]: + """Split at commas that are not inside ``<...>`` (e.g. ``complex``).""" + parts, depth, current = [], 0, [] + for char in text: + if char == "<": + depth += 1 + elif char == ">": + depth -= 1 + elif char == "," and depth == 0: + parts.append("".join(current)) + current = [] + continue + current.append(char) + parts.append("".join(current)) + return parts + + +def _template_arg(value: Any) -> str: + """A template argument as C++ source: a C type for dtypes, else a literal.""" + if isinstance(value, str): + return value + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, np.integer)): + return str(int(value)) + return ctype_of(value) + + +def parse_cuda_signature( + source: str, + name: str, + *, + structs: Iterable[CudaStruct] = (), + template_args: Sequence[Any] | None = None, +) -> tuple[CudaParameter, ...]: + """Parse the parameters of the ``__global__`` function `name` in `source`. + + Parameters + ---------- + source : str + 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. + + Returns + ------- + tuple[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). + """ + code = _strip_comments(source) + structs = {s.name: s for s in structs} + for struct in structs.values(): + struct.check_source(code) + + pattern = r"(?:template\s*<(?P