diff --git a/.github/workflows/gpu_ci_trigger.yml b/.github/workflows/gpu_ci_trigger.yml index 60fa668..7f7a8e7 100644 --- a/.github/workflows/gpu_ci_trigger.yml +++ b/.github/workflows/gpu_ci_trigger.yml @@ -23,10 +23,11 @@ name: Sync to GitLab and Run GPU CI on: push: - branches: [main, devel] + branches: [devel] pull_request: branches: [main, devel] workflow_dispatch: + workflow_call: jobs: sync-and-test: @@ -57,14 +58,14 @@ jobs: SAFE_REF="${SOURCE_REF//\//-}" TARGET_BRANCH="gpu-test-${SAFE_REF}" fi - echo "TARGET_BRANCH=$TARGET_BRANCH" >> $GITHUB_ENV + echo "TARGET_BRANCH=$TARGET_BRANCH" >> "$GITHUB_ENV" # 3. Add GitLab SSH remote git remote add gitlab git@gitlab.mpcdf.mpg.de:maxlin/cunumpy.git # 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 + 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" @@ -84,7 +85,7 @@ jobs: # 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 + for _ 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 diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml new file mode 100644 index 0000000..b13bab3 --- /dev/null +++ b/.github/workflows/macos.yml @@ -0,0 +1,56 @@ +name: Tests (macOS) + +on: + push: + branches: + - main + - devel + pull_request: + branches: + - main + - devel + workflow_dispatch: + +jobs: + build: + # macos-latest is Apple silicon (arm64): checks the NumPy backend, the + # compiled host kernels and the MLX import on the platform of MetalKernel. + runs-on: macos-latest + + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.13"] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: pyproject.toml + + # Pyccel is built from source here (no wheel for macOS arm64) and needs a + # Fortran compiler, which the macOS runners do not have. + - name: Install gfortran + run: | + brew install gcc + gfortran --version + + - name: Install project + run: | + pip install --upgrade pip + pip install ".[test-compiled,metal]" + + # Hosted macOS runners are virtual machines and usually have no Metal GPU, + # in which case the MetalKernel launch tests are skipped. This step shows + # which case this run is in. + - name: Report Metal availability + run: | + python -c "import platform, cunumpy as xp; print(platform.machine(), 'metal_available =', xp.kernels.metal_available())" + + - name: Run tests + run: pytest . -rs diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index eaff1cd..48386c0 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -6,15 +6,25 @@ on: - main jobs: + # Local reusable workflows check out this push's commit. Main validation is + # invoked here, so publishing cannot race independent CPU/GPU workflows. + cpu-tests: + uses: ./.github/workflows/testing.yml + + gpu-tests: + uses: ./.github/workflows/gpu_ci_trigger.yml + secrets: inherit + build-and-publish: + needs: [cpu-tests, gpu-tests] runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.10" @@ -26,6 +36,13 @@ jobs: - name: Build the package run: python -m build + - name: Check distributions and smoke-test the installed wheel + run: | + python -m twine check dist/* + python -m venv "$RUNNER_TEMP/wheel-smoke" + "$RUNNER_TEMP/wheel-smoke/bin/python" -m pip install dist/*.whl + "$RUNNER_TEMP/wheel-smoke/bin/python" -I testing/smoke_wheel.py + - name: Publish to PyPI env: TWINE_USERNAME: __token__ # Use API token diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 4188463..614da82 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -3,14 +3,33 @@ name: Tests on: push: branches: - - main - devel pull_request: branches: - main - devel + workflow_call: jobs: + compiled-fake-cupy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install compilers and compiled-test dependencies + run: | + sudo apt-get update + sudo apt-get install -y gcc gfortran + python -m pip install ".[test-compiled]" + - name: Test compiled host kernels with fake device arrays + env: + CUNUMPY_FAKE_CUPY: "1" + CUNUMPY_BACKEND: cupy + CUNUMPY_REQUIRE_PYCCEL: "1" + run: python -m pytest -q tests/unit/test_pyccel_kernel.py tests/unit/test_cupy.py + build: runs-on: ubuntu-latest @@ -22,17 +41,17 @@ jobs: steps: # Checkout the repository - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 # Set up Python - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} # Cache pip dependencies - name: Cache pip - uses: actions/cache@v3 + uses: actions/cache@v4 with: path: ~/.cache/pip key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }} diff --git a/CHANGELOG.md b/CHANGELOG.md index 7022e86..7499713 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,98 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- Compiled Pyccel regression coverage for F-ordered arrays and column-block + copy-back, run in CPU CI with fake CuPy and required compilation. +- Publishing now requires CPU and GPU validation of the release commit, checks + the built distributions, and smoke-tests the wheel in a fresh environment. + +### Fixed +- The ordering guide identifies 0.6.3 as the first release preserving F order + in CuPy-to-host conversions. + +## [0.6.3] - 2026-10-08 + +### Fixed +- The emulation of kernels with `__syncthreads` compiles with GCC on macOS. +- Preserve Fortran order when copying CuPy arrays to the host through + `to_numpy`, `host_call`, `evaluate_on_host`, and `PyccelKernel` conversions. + +### Added +- `emulate_cuda_kernel` compiles each kernel once, into a shared library that + takes the arguments at run time and runs in the process (ctypes), on the arrays + themselves; later launches with any values and sizes reuse it. Libraries are + cached per process and on disk (`kernel_testing.emulation_cache_dir()`, + `CUNUMPY_EMULATION_CACHE`, default `~/.cache/cunumpy/emulation`). + `kernel_testing.compile_for_emulation(kernel)` builds one without a launch. + `__trap()` (a failed bounds check) still raises `RuntimeError`, but a + segmentation fault now ends the process. +- Inside `emulated_launches()`, `CudaKernel.compile()` (and so `recompile()`, + `CudaKernelVariants.compile_all()` and `KernelCatalog.compile_all()`) builds the + emulation library instead of compiling CUDA, and returns None. Code that compiles + its kernels before the time loop runs on the fake CuPy and still reports compile + errors there; patching `CudaKernel.compile` is no longer needed. +- The emulation compiles inline PTX (`asm(...)`, `asm volatile(...)`) as a trap, so + a kernel with an `asm("trap;")` branch needs no `-Dasm(x)=__trap()` option. +- `kernel_testing.fake_cupy_session()` runs a CuPy-backend program on the CPU + (CuPy backend on the fake CuPy, launches and compilation emulated); + `kernel_testing.device_backend_available()` and the marker + `requires_device_backend` (a GPU or the fake CuPy); + `kernel_testing.run_in_fake_cupy_subprocess(code)` runs code in a serial child + process on the fake CuPy (rank 0 only under MPI) and fails the test with the + signal or exit code and the end of the child's output. +- `profiling.TransferBudget` counts transfers per phase of a program + (`budget.phase(name)`, the decorator `budget.count(name)`, `start()`/`stop()`) + and checks a rule per phase (`require(phase, allow={"to_host": {"max_nbytes": 8}}, + calls=n)`, `check()`, `report()`). `count_transfers(into=counter)` adds to an + existing counter (counted once when nested). +- `TransferEvent` has `blocking` (False for `to_host_async()` and + `HostStaging.copy()`) and `implicit` (True for the scalar reads of the fake CuPy). +- `xp.to_host_async(a)` copies a device scalar or small array to the host on a + separate stream without waiting; the returned `memory.HostCopy` has `ready()` + (never waits) and `result()`. +- `kernel_testing.emulated_launches()` runs every `CudaKernel` launch in a block + on the CPU, on the host buffers of the fake CuPy arrays, so code that launches + kernels can be tested without a GPU. `emulate_cuda_kernel` now accepts struct + parameters (a mapping of field values, a `CudaStructValue`, or an object with + an attribute per field). `kernel_testing.host_buffer(array)` returns the NumPy + array behind a fake CuPy array. +- `CUNUMPY_REQUIRE_CUDA=1` makes the GPU markers of `kernel_testing` fail instead + of skipping: `requires_cupy`, the `cupy` run of the `backend` fixture (which + activates CuPy strictly) and `assert_kernels_agree`. New `kernel_testing.cuda_required()`. +- `count_transfers()` records `sync` events (the host waiting for the device): + `xp.synchronize()`, the waits of the MPI helpers and of the CUDA debug mode, and, + on the fake CuPy, scalar reads of device arrays. They are in `counter.syncs` + and the report but not in `total`; `assert_no_transfers(syncs=True)` rejects them. + The real CuPy's own `float(a)` cannot be observed from Python and is not counted. +- `xp.algorithms.compact_by_mask(mask, *arrays, axis=0)` moves the masked entries + of arrays to the front along `axis`, in place and in order, and returns their + number; `axis=-1` compacts component-major `(ncomp, N)` and `(N,)` arrays together. +- `as_kernel_array(..., strided=True)` and `kernel_output(..., strided=True)` take a + NumPy array in C order with gaps (positive strides, each at least the extent of + the next axis, e.g. `storage[:, :n]`) unchanged for host kernels, instead of + copying it to a C-contiguous array; Pyccel's wrappers take such arrays. +- `CudaKernel(n_threads_from="last_axis")`: one thread per entry of the last axis of + the first array argument, for component-major `(ncomp, N)` marker arrays. +- Kernel outputs: a name in `PyccelKernel(outputs=...)` also finds a positional + argument and an index a keyword argument, using the parameter names of the + function (or the new `parameters=`); a `Kernel` supplies those of its host + function. `xp.kernels.outputs_from_annotations()` reads the outputs from the + annotations (not `Final`, `const` or a scalar), and `Kernel.from_folder()` / + `KernelCatalog.from_package()` take `outputs=` (names, indices or + `"annotations"`). `assert_kernels_agree(outputs=...)` takes parameter names and + `"name.field"` to compare only some fields of a struct argument. +- `xp.kernels.MetalKernel` runs a Metal Shading Language kernel on the GPU of an + Apple silicon Mac through MLX (`pip install 'cunumpy[metal]'`). It takes and + fills NumPy arrays, is float32 only (`float64="cast"` computes float64 data in + float32), and its copies are counted by `xp.profiling.count_transfers()`. + `xp.kernels.metal_available()` tells whether it can run. +- `xp.host_call`, `xp.evaluate_on_host` and `xp.setup_on_host` run host-only code + (SciPy splines, file readers, external libraries) with arguments of either + backend: device arrays are copied to the host, the call runs on the NumPy + backend and array results are copied back. The copies are counted by + `xp.profiling.count_transfers()`. + ### Changed - The launcher detection and the serial MPI stand-in moved to the new package [maybempi](https://github.com/max-models/maybempi), a dependency of cunumpy. diff --git a/README.md b/README.md index 21e5c34..a71001a 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ never hide a NumPy name: | Submodule | Contents | |---|---| -| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, host implementations, `fuse` | +| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `MetalKernel`, host implementations, `fuse` | | `xp.arguments` | CUDA only: `CudaStruct`, `CudaStructArguments`, `CudaArguments` | | `xp.cuda` | CUDA only: devices, streams, debug mode, CUDA headers | | `xp.rng` | `random_streams`, `get_rng`, `philox_*` | @@ -36,6 +36,8 @@ never hide a NumPy name: | `cunumpy.kernel_testing` | pytest helpers for host/CUDA kernel pairs | Everything except `xp.cuda`, `xp.arguments` and `CudaKernel` works on both backends. +`MetalKernel` runs Metal kernels on the GPU of an Apple silicon Mac (MLX, float32, NumPy arrays; +`pip install 'cunumpy[metal]'`, see [the guide](docs/source/kernels/metal-kernel.md)). ## Install diff --git a/docs/source/api.md b/docs/source/api.md index 06971ac..eda5916 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -39,15 +39,15 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | Submodule | Backends | Contents | |---|---|---| -| `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `scipy`, `require_version` | -| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | +| `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` | +| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | | `cunumpy.arguments` | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` | | `cunumpy.cuda` | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) | | `cunumpy.rng` | both | `random_streams`, `get_rng`, `philox_*` | -| `cunumpy.algorithms` | both | `morton_*`, `sort_by_key`, `segment_sum` | +| `cunumpy.algorithms` | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` | | `cunumpy.mpi` | both | `mpi_buffer`, CUDA-aware MPI detection, `local_rank`, `synchronize_for_mpi` | -| `cunumpy.profiling` | both | `timed_region`, `nvtx_range`, `count_transfers`, `assert_no_transfers` | -| `cunumpy.memory` | both | `HostStaging`, `DeviceMirror` | +| `cunumpy.profiling` | both | `timed_region`, `nvtx_range`, `count_transfers`, `assert_no_transfers`, `TransferBudget` | +| `cunumpy.memory` | both | `HostStaging`, `HostCopy`, `DeviceMirror` | | `cunumpy.petsc` | both | `petsc_vec` | | `cunumpy.kernel_testing` | both | pytest helpers for host/CUDA kernel pairs (not imported by `import cunumpy`) | @@ -206,6 +206,10 @@ Converts to a host-side NumPy array. CuPy arrays are copied from device to host. Other array-like inputs are passed through `numpy.asarray`; NumPy arrays may therefore be returned as-is rather than copied. +F-contiguous CuPy inputs keep F order on the host; other CuPy inputs become +C-contiguous. Arbitrary device strides are not preserved. See +[Array ordering and strides](guides/array-ordering.md). + ### `to_cupy(array)` Converts an array-like input to a CuPy array. Raises `ImportError` if CuPy or @@ -224,6 +228,28 @@ assert xp.get_array_backend(normalized) == xp.get_backend() Each conversion returns a suitable array; it does not change the active backend or mutate the source. +### `to_host_async(array)` + +```python +pending = xp.to_host_async(residual_norm) # a device scalar; returns at once +... # queue the next iteration's kernels +if pending.ready() and pending.result() < tol: + break +``` + +Starts copying a device scalar or small array to the host without waiting for +it, so that the host can keep queueing kernels (e.g. a convergence test read one +iteration late, without a sync per iteration). The copy runs on a separate +stream into page-locked memory, after the work queued so far on the current +stream. The returned `memory.HostCopy` has `ready()` (never waits) and +`result()` (waits only if the copy has not finished; a NumPy scalar for a 0-d +array, else a NumPy array). `count_transfers()` records a `to_host` event with +`blocking=False`, and a wait in `result()` as a `sync`. + +On the NumPy backend the value is copied at once and `ready()` is always True; +on the fake CuPy too, and the copy is counted. For large arrays copied +repeatedly, use `memory.HostStaging`, which reuses its buffers. + ### `algorithms.segment_sum(values, keys, n_segments, *, out=None)` `out[k] = sum(values[i] for keys[i] == k)` on the backend of `keys`, with @@ -255,17 +281,47 @@ sparse, negative, and uint64 Morton keys. Starts/stops are int64 half-open indic Empty input returns three empty arrays. Validation and variable-length GPU output may synchronize; prepare boundaries outside repeated operations. -### `algorithms.sort_by_key(keys, *arrays)` +### `algorithms.sort_by_key(keys, *arrays, axis=0)` Stable argsort of the 1D `keys` (CuPy's radix sort on the device), applied to -every array along axis 0, in one call: +every array along `axis`, in one call: ```python keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, charges) +# component-major (ncomp, N) markers: sort along the last axis of each array +keys, order, positions, weights = xp.algorithms.sort_by_key( + keys, positions, weights, axis=-1 +) +``` + +Returns `(keys[order], order, *(take(a, order, axis) for a in arrays))`, +`order` as `int64`. Equal keys keep their order, so the result is reproducible. +On NumPy, integer keys of more than 4096 entries are sorted by an LSD radix +sort on 16-bit digits (one stable `argsort` of `uint16` per digit, as many as +the key range needs), about ten times faster than the stable sort of 64-bit +integers. + +### `algorithms.compact_by_mask(mask, *arrays, axis=0)` + +Moves the rows where the boolean `mask` is True to the front of every array, in +place and in their original order, and returns how many there are. Typical use: +keep the live particles at the front of the marker arrays. + +```python +n = xp.algorithms.compact_by_mask(alive, markers, weights) +markers, weights = markers[:n], weights[:n] ``` -Returns `(keys[order], order, *(a[order] for a in arrays))`, `order` as -`int64`. Equal keys keep their order, so the result is reproducible. +The rows after the first `n` are unspecified. The count is needed on the host, +so on CuPy each call synchronizes once. The mask and the arrays must be on the +same backend. `axis` is the axis of every array that the mask indexes; `-1` +compacts component-major arrays, `(ncomp, N)` next to `(N,)`, along their +marker axis: + +```python +n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1) +positions, weights = positions[:, :n], weights[:n] +``` ## Count transfers @@ -273,7 +329,7 @@ A transfer inside a time loop is the classic performance bug of a GPU port: every step then waits for the device and copies an array. These helpers let a test verify that a block of code does not transfer at all. -### `profiling.count_transfers()` +### `profiling.count_transfers(into=None)` Context manager yielding a `TransferCounter` that records every host/device transfer made through CuNumpy while the block runs, with the call site of @@ -286,7 +342,7 @@ with xp.profiling.count_transfers() as counter: assert counter.total == 0, counter.report() ``` -Five kinds of events are recorded: +Six kinds of events are recorded: * `to_host`: an actual device-to-host copy through conversion, mirror, staging, serial MPI, or host kernel helpers; @@ -298,7 +354,12 @@ Five kinds of events are recorded: * `fallback`: a `Kernel` without CUDA kernel calling its host kernel on the CuPy backend (`missing_cuda="fallback"`), one event per call, naming the kernel. Physical copies and a `kernel_conversion` marker are recorded separately; -* `device_copy`: a device-only dtype/layout conversion through CuNumpy helpers. +* `device_copy`: a device-only dtype/layout conversion through CuNumpy helpers; +* `sync`: the host waited for the device: `xp.synchronize()`, the waits of the MPI + helpers (`synchronize_for_mpi()`, `mpi_buffer()` staging) and of the CUDA debug + mode, and, on the fake CuPy, a scalar read of a device array (`float(a)`, + `int(a)`, `bool(a)`, `a.item()`, `a.tolist()`). They are in `counter.syncs` and in + the report, but not in `total`. Only real transfers count: `to_numpy()` of a NumPy array or `to_cupy()` of a CuPy array records nothing. The counter has the attributes `to_host`, @@ -306,8 +367,12 @@ CuPy array records nothing. The counter has the attributes `to_host`, `device_copies`, `bytes_to_host`, `bytes_to_device`, and `bytes(kind)`. `total` includes physical copies and conversion/fallback markers. Byte totals include only physical copies, so markers do not double-count bytes. -`events` is a list of `TransferEvent(kind, description, where, nbytes=None)`; -`where` is the caller's `file:line` and `nbytes` is None for markers/unknown sizes. +`events` is a list of `TransferEvent(kind, description, where, nbytes=None, +blocking=True, implicit=False)`; `where` is the caller's `file:line` and `nbytes` +is None for markers/unknown sizes. `blocking` is False for the copies that the host +does not wait for (`to_host_async()`, `HostStaging.copy()`), and `implicit` is True +for the scalar reads of the fake CuPy (`float(a)`), which are `sync` events; the +report marks them `[async]` and `[implicit]`. `kernel_conversion_calls` selects the conversion markers. `report()` returns a multi-line string with the events grouped by kind and call site, with counts: @@ -321,22 +386,27 @@ counts: ``` Blocks can be nested; every active counter sees the transfers made inside it. +`count_transfers(counter)` adds the events to an existing counter instead, e.g. to +accumulate over several calls; a counter that is already active is not added +again, so each event is counted once. Like every `contextlib.contextmanager`, it +is also a decorator: `@count_transfers(counter)`. When no counter is active, the instrumentation costs a single check per call. Like the backend selection, the active counters are process-wide state and not thread-safe. **Limitation:** only transfers made through CuNumpy are seen. Raw `cupy.ndarray.get()`, `cupy.asarray(numpy_array)`, `numpy.asarray(cupy_array)`, -`float(device_array)`, forwarded backend calls such as `xp.asarray()`, and -implicit conversions inside other libraries are -not counted. Use `nsys` (or CuPy's profiling hooks) to find those. +forwarded backend calls such as `xp.asarray()`, and implicit conversions inside +other libraries are not counted. Neither is `float(device_array)` (an implicit +sync) with the real CuPy, which cannot be observed from Python; the fake CuPy +counts it. Use `nsys` (or CuPy's profiling hooks) to find those. -### `profiling.assert_no_transfers()` +### `profiling.assert_no_transfers(*, syncs=False)` Context manager that raises `AssertionError` with the counter's `report()` if the block makes a host/device transfer or host fallback through CuNumpy. -Device-only dtype/layout conversions are allowed. It yields the `TransferCounter` -too. An exception raised inside the block propagates as it is: +Device-only dtype/layout conversions are allowed, and so are syncs unless +`syncs=True`. It yields the `TransferCounter` too. An exception raised inside the block propagates as it is: ```python def test_time_step_stays_on_the_device(): @@ -344,6 +414,47 @@ def test_time_step_stays_on_the_device(): propagator(dt) ``` +### `profiling.TransferBudget(*, started=True)` + +Counts the transfers per phase of a program (the time step, the diagnostics, +the output) and checks a rule for each phase: + +```python +budget = xp.profiling.TransferBudget(started=False) +model.integrate = budget.count("integrate")(model.integrate) # a decorator +... # setup: not counted +budget.start() +for step in range(n_steps): + model.integrate(dt) + with budget.phase("output"): # or a context + save(model) + +budget.require("integrate", allow={"to_host": {"max_nbytes": 8}}, calls=n_steps) +budget.require("output", allow={"to_host": {"max_count": n, "max_total_bytes": b}}) +budget.check() # AssertionError with budget.report() if a rule is broken +``` + +`phase(name)` counts its block in phase `name` and yields the phase's +`TransferCounter`; `count(name)` is a decorator that does the same for every +call. The events of a phase accumulate over its calls (`budget[name]`, +`budget.phases`), and `budget.calls[name]` counts the calls. Phases nest: an +event is counted in the innermost phase only, and a phase entered again inside +itself (recursion) counts each event and call once. Until `start()` (with +`started=False`), and after `stop()`, phases run uncounted. + +`require(phase, allow=None, *, ignore=("device_copy", "sync"), calls=None)` +sets the rule of a phase: the event kinds in `allow` are allowed, each with +optional limits, every other kind is forbidden except those in `ignore` (unless +they are in `allow`). The limits are `max_nbytes` (each event; an event of +unknown size breaks it), `max_count` and `max_total_bytes` (the whole phase), +and `blocking` and `implicit` (the events must have this value), e.g. +`{"to_host": {"max_nbytes": 8, "blocking": False}}` to allow only non-blocking +scalar copies, or `{"sync": {"implicit": False}}` with `ignore=("device_copy",)` +to reject the scalar reads of the fake CuPy. `calls` is the number of calls the +phase must have had. `violations()` lists the broken rules, with the offending +events and where they happened; `report()` shows every phase's events and the +broken rules; `check()` raises `AssertionError` with the report. + ### `as_device_array(value, dtype=None, ndim=None, *, name=None)` The "reference or copy once" rule for building CUDA argument objects @@ -375,7 +486,7 @@ class DeviceParticles(xp.arguments.CudaArguments): super().__init__(self.markers, self.degree, self.markers.shape[0]) ``` -### `kernels.as_kernel_array(value, like, dtype=None)`, `kernels.kernel_output(out, like, dtype=None)` +### `kernels.as_kernel_array(value, like, dtype=None, *, strided=False)`, `kernels.kernel_output(out, like, dtype=None, *, strided=False)` For the arguments of a `Kernel` with `dispatch="arrays"`, whose choice follows the arrays. `as_kernel_array` returns `value` on the side of `like` (a CuPy @@ -387,6 +498,13 @@ itself if `as_kernel_array` takes it unchanged, else a converted copy whose contents are written into `out` (on its own side) when the block ends without an error. +With `strided=True` a NumPy array for a host kernel is also taken unchanged +when it is in C order with gaps (positive strides, each at least the extent of +the next axis), e.g. the first `n` columns `storage[:, :n]` of a marker buffer. +Pyccel's wrappers take such arrays without a copy; a host kernel that needs +contiguous memory must not use it. F-ordered, transposed and negative-stride +arrays are still copied, and CuPy arrays are always made C-contiguous. + ```python with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer: gather( @@ -808,8 +926,18 @@ CuPy arrays. * `is_array`: predicate for host array values to convert back to CuPy. The default is `isinstance(value, numpy.ndarray)`. * `outputs`: sequence of arguments the kernel may write to. Entries are - positional indices or keyword names. If omitted, every converted array is + indices or parameter names; a name also finds a positional argument, and an + index a keyword argument, when the parameter names are known (from the + Python signature, or `parameters`). If omitted, every converted array is copied back. +* `parameters`: the names of the positional parameters, or a function that + returns them, for compiled kernels without a Python signature. A `Kernel` + supplies the names of its host function. + +`kernels.outputs_from_annotations(function)` returns the parameters a kernel may +write to, read from its annotations: everything that is not `Final`, `const` or a +scalar. Use it as `outputs=outputs_from_annotations(push)`, or pass +`outputs="annotations"` to `Kernel.from_folder()` / `KernelCatalog.from_package()`. ### Example and output declarations @@ -857,6 +985,51 @@ lists) are converted back using `is_array`; dictionaries in return values are not recursively converted. On the NumPy path, the original return value and normal Python mutation and exception behavior are preserved. +## `kernels.MetalKernel` + +A Metal Shading Language kernel for the GPU of an Apple silicon Mac, run with +[MLX](https://github.com/ml-explore/mlx) (`pip install 'cunumpy[metal]'`). It +takes NumPy arrays and fills the output arrays you pass, so no backend switch +is needed. + +```python +import numpy as np +import cunumpy as xp + +scale = xp.kernels.MetalKernel( + "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];", + inputs=["x", "a"], + outputs=["y"], +) +x = np.arange(8, dtype=np.float32) +y = np.empty_like(x) +scale(x, 2.0, out=y) +``` + +`MetalKernel(source, inputs, outputs, *, name="cunumpy_kernel", header="", +threadgroup=256, float64="error", atomic_outputs=False, init_value=None)` + +- `source` is the body of the kernel function. MLX generates the signature: each + name in `inputs` and `outputs` is a pointer to the flat, row-major data of that + array, so `x[i]` is the flat index. `thread_position_in_grid` and the other + Metal attributes used in the body are added automatically. `header` goes before + the function (includes, defines, helper functions). +- Calling the kernel: `kernel(*inputs, out=array_or_arrays, n_threads=None, + template=None)`. `n_threads` is the total thread count (default: first axis of + the first output). `template` gives compile-time constants, e.g. + `template={"NSTEPS": 200}`. It returns the output array, or a tuple of them. +- Outputs are uninitialized: write every element, pass the old array as an input + too if the kernel reads it, or set `init_value`. +- **float32 only.** The Apple GPU has no float64: a float64 input or output + raises `TypeError`. With `float64="cast"` float64 data is computed in float32 + (a push over 200 steps agreed with float64 to about 3e-5). +- Every call copies the inputs to MLX arrays and the results back, counted as + `to_device` and `to_host` transfers by `count_transfers()`. On a 4M-particle + push these copies were about 7 ms next to an 11.5 ms kernel. +- `xp.kernels.metal_available()` is True if MLX is installed and a Metal GPU is + present. Without them, calling a `MetalKernel` raises `ImportError` or + `RuntimeError`. + ## `kernels.CudaKernel` ### Constructor @@ -1016,7 +1189,8 @@ Explicit `n_threads` or `grid` always takes precedence. Set `n_threads_from` (constructor argument or settable property) to a callable such as `lambda args: args[0].size` for flattened element kernels, or `lambda args: args[0].shape[::-1]` for kernels whose x index follows columns. -`"first_array"` always uses the first axis; None disables inference and requires +`"first_array"` always uses the first axis and `"last_axis"` the last one (one +thread per entry of a component-major `(ncomp, N)` or `(N,)` array); None disables inference and requires explicit launch sizes. The same defaults apply through `kernels.Kernel` and `kernel_testing.assert_kernels_agree`, and in CPU emulation. @@ -1710,7 +1884,11 @@ from cunumpy.kernel_testing import ( assert_kernels_agree, device_function_kernel, emulate_cuda_kernel, + emulated_launches, + fake_cupy_session, requires_cupy, + requires_device_backend, + run_in_fake_cupy_subprocess, ) ``` @@ -1738,6 +1916,15 @@ The `backend` fixture does the same and activates the backend for the test; import it into a `conftest.py` (`from cunumpy.kernel_testing import backend`) or the test module, then take `backend` as a test argument. +### `requires_device_backend`, `device_backend_available()` + +`device_backend_available()` tells whether a CuPy-backend program can run: a +GPU, or the fake CuPy. That is more than `requires_cupy` allows (CUDA kernels +can be launched): on the fake CuPy, launches only run inside +`emulated_launches()` (or `fake_cupy_session()`). `requires_device_backend` +skips tests without either; with `CUNUMPY_REQUIRE_CUDA` set it fails them +instead. + ### `assert_kernels_agree(kernel, make_args, ...)` ```python @@ -1894,17 +2081,28 @@ so that CI without a GPU can compare a kernel with its host version. The kernel source is compiled as C++ (C++17, `CXX` or `c++`; see `emulation_compiler()`) with the CUDA built-ins replaced: `threadIdx`, `blockIdx`, `blockDim`, `gridDim`, atomics (`atomicAdd`, `atomicMin`, ..., -plain operations), `__ldg`, `rsqrt`, `__trap` (aborts). The shipped headers -and the kernel's include directories and `-D` options apply. Then the kernel is -called once per thread, for every block and thread index of the launch shape. +plain operations), `__ldg`, `rsqrt`, `__trap`. The shipped headers and the +kernel's include directories and `-D` options apply. Then the kernel is called +once per thread, for every block and thread index of the launch shape. + +Each kernel is compiled once (for each compiler, options and template +arguments) into a shared library that takes the arguments at run time, so all +later launches, with any values and array sizes, reuse it. Libraries are cached +in the process and on disk, like CuPy's kernel cache: in +`emulation_cache_dir()`, which is `CUNUMPY_EMULATION_CACHE` (`0`: no disk +cache), else `$XDG_CACHE_HOME/cunumpy/emulation` or `~/.cache/cunumpy/emulation`. +The library runs in the Python process, on the arrays themselves. Arguments follow the signature, with NumPy arrays in place of CuPy arrays: pointer and view parameters (`Array1D` to `Array16D`) take arrays of the -declared dtype (and ndim), passed as contiguous copies, so any strides work, -and written back into the given arrays; scalars are checked and cast like in a -launch. Like NVRTC by default, the compiler may fuse `a * b + c` into an FMA, -so compare with NumPy using a tolerance of a few ulp, or pass -`options=("-ffp-contract=off",)` for NumPy's rounding. +declared dtype (and ndim); the kernel writes into them directly (an array that +is not writeable, or whose layout a parameter cannot take, is passed as a +copy and written back), so any strides work and aliased arguments see each +other's writes. Struct parameters take a mapping of field names to values, a +`CudaStructValue`, or an object with an attribute per field. Scalars are +checked and cast like in a launch. Like NVRTC by default, the compiler may fuse +`a * b + c` into an FMA, so compare with NumPy using a tolerance of a few ulp, +or pass `options=("-ffp-contract=off",)` for NumPy's rounding. Block shared memory and `__syncthreads` are emulated: `__shared__` variables are one copy per block (blocks run one after another), `extern __shared__` @@ -1915,12 +2113,74 @@ thread continues past it. Per-block deposits, shared-memory reductions and tiled kernels work. Not emulated: concurrency between barriers (races and atomic ordering never -show), warp intrinsics, struct parameters and complex scalars. A kernel, or a -header it includes, using `__syncwarp`, warp shuffles or votes raises -`NotImplementedError` (serial threads would give wrong results); a kernel that -does not compile, or crashes (an out-of-bounds -index with `-DCUNUMPY_BOUNDS_CHECK`, `__trap()`), raises `RuntimeError` with -the compiler or program output. +show), warp intrinsics, complex scalars and inline PTX. Inline `asm(...)` and +`asm volatile(...)` statements compile, but trap when reached, so a kernel with +a PTX branch that a test never takes (e.g. `asm("trap;")` for an unknown case) +runs without extra options. A kernel, or a header it includes, using +`__syncwarp`, warp shuffles or votes raises `NotImplementedError` (serial +threads would give wrong results). A kernel that does not compile, or traps (an +out-of-bounds index with `-DCUNUMPY_BOUNDS_CHECK`, `__trap()`, inline `asm`), +raises `RuntimeError`. Other crashes, such as a segmentation fault from an +unchecked out-of-bounds index, end the process; run such code in a child +process (`run_in_fake_cupy_subprocess()`, which prints the traceback with +`faulthandler`). + +### `compile_for_emulation(kernel, *, compiler=None, options=())`, `emulation_cache_dir()` + +`compile_for_emulation()` builds (or fetches from the cache) the emulation +library of a kernel without launching it, so compile errors show up early; +it raises `RuntimeError` with the compiler output. A kernel without a parsed +signature is only checked for syntax and type errors (`-fsyntax-only`). +`emulation_cache_dir()` is the disk cache, or None. + +### `emulated_launches(*, compiler=None, options=())` + +```python +with emulated_launches(): + propagator(dt) # CuPy backend = the fake CuPy: the kernels run on the CPU +np.testing.assert_allclose(host_buffer(markers), expected) +``` + +Inside the block every `CudaKernel` launch runs through `emulate_cuda_kernel`, +on the host buffers of the fake CuPy arrays (`host_buffer(array)` is the NumPy +array behind one), with the `compiler` and `options` of the block. Argument +objects are flattened like in a launch. No CUDA is compiled in the block +either: `CudaKernel.compile()` (and so `recompile()`, +`CudaKernelVariants.compile_all()` and `KernelCatalog.compile_all()`) calls +`compile_for_emulation()` and returns None instead of a `cupy.RawKernel`, so a +program that compiles its kernels up front still reports compile errors there. +Kernels the emulation cannot run (warp intrinsics) are skipped by `compile()`. +This holds on a GPU too: the block means "no CUDA here". The original methods +are restored when the block exits, also by an exception. + +### `fake_cupy_session(*, compiler=None, options=())` + +```python +with fake_cupy_session(): + sim.run() # CuPy backend; kernels compiled up front and launched on the CPU +``` + +Everything a CuPy-backend program needs to run on the CPU: activates the CuPy +backend (the fake CuPy) and `emulated_launches(compiler=..., options=...)`. +Raises `RuntimeError` if the fake CuPy is not active. + +### `run_in_fake_cupy_subprocess(code, *, env=None, timeout=None)` + +```python +def test_domain_on_the_fake_cupy(): + run_in_fake_cupy_subprocess("from my_sim.tests import check_domain; check_domain()") +``` + +The fake CuPy must be installed before anything imports cunumpy, so a test +process that already uses cunumpy cannot switch to it. This runs +`python -X faulthandler -c code` in a child process with `CUNUMPY_FAKE_CUPY=1` +and `OMP_NUM_THREADS=1`, without the environment variables of an MPI launcher +and with `MAYBEMPI=0`, so the child runs serially and does not join the +parent's MPI job; `env` adds variables. Under MPI only rank 0 starts the child, +the other ranks skip the test. If the child fails (or runs longer than +`timeout` seconds), the test fails with the exit code or the signal (e.g. +`SIGSEGV`) and the last 50 lines of its stdout and stderr. Returns the +`subprocess.CompletedProcess` otherwise. ## `memory.HostStaging` diff --git a/docs/source/guides/array-ordering.md b/docs/source/guides/array-ordering.md new file mode 100644 index 0000000..0312bff --- /dev/null +++ b/docs/source/guides/array-ordering.md @@ -0,0 +1,152 @@ +# Array ordering and strides + +CuNumpy supports C-ordered (row-major) and F-ordered (column-major) arrays on +both NumPy and CuPy. Choose the layout when allocating an array, and keep it +consistent with the compiled kernels that consume it. + +For a two-dimensional array `a[row, column]`, C order stores each row together; +F order stores each column together. Indexing and shape stay the same. F order +can help code that processes columns of a particle buffer, but measure the +whole workload, including kernels and copies, before choosing a layout. + +## Allocate and inspect + +```python +import cunumpy as xp + +markers = xp.zeros((100, 6), dtype=xp.float64, order="F") +assert markers.flags.f_contiguous + +positions = markers[:, 1:4] +assert positions.flags.f_contiguous +assert positions.shape == (100, 3) +``` + +A block of complete columns of an F-contiguous array remains F-contiguous. +A block of complete rows of a C-contiguous array remains C-contiguous. +Stepped slices such as `markers[::2, :]` generally have neither layout. +Check `a.flags.c_contiguous`, `a.flags.f_contiguous`, and `a.strides` when +debugging a kernel boundary. Strides are measured in bytes. Arrays with +singleton dimensions or no elements can satisfy both contiguity flags. + +Use `xp.asfortranarray(a)` to require F order or `xp.ascontiguousarray(a)` to +require C order. These may allocate a copy; a transpose changes the shape and +indexing meaning and is not a substitute for an order conversion. + +## Host and device conversions + +| Operation | Layout behavior | +| --- | --- | +| `xp.to_numpy(cupy_array)` | F-contiguous input becomes F-contiguous host storage; other input becomes C-contiguous host storage. | +| `xp.to_numpy(numpy_array)` | Uses `numpy.asarray`; an ordinary NumPy array keeps its layout and strides without copying. | +| `xp.to_cupy(numpy_array)` | Uses `cupy.asarray`, preserving C or F order for contiguous input. | +| `xp.to_cunumpy(a)` | Uses the conversion for the active backend. | +| `xp.host_call`, `xp.evaluate_on_host` | Device arguments follow `to_numpy`; returned NumPy arrays follow `to_cupy`. | +| `xp.kernels.PyccelKernel` | Device arguments use the same host conversion as `to_numpy`; declared outputs are copied back into the original device arrays. | + +The device-to-host conversion uses `array.get(order="A")`. This preserves +F-contiguous layout, including complete column blocks, but does not preserve +arbitrary strides or shared storage between distinct views. Calling CuPy's +`array.get()` directly defaults to C order; see the +[CuPy array documentation](https://docs.cupy.dev/en/stable/reference/generated/cupy.ndarray.html). + +The host function itself controls the layout of any new arrays it creates. +Returning a C-ordered result from `host_call` does not make that result F-ordered +just because an input was F-ordered. + +```{note} +CuPy-to-host conversions preserve F order starting with **cuNumPy 0.6.3**. +Earlier releases used `array.get()` and produced C-ordered host copies. +Applications depending on F order at a host-kernel boundary require +`cunumpy>=0.6.3`. +``` + +## Copies, masks, and assignment + +`a.copy()` defaults to C order on both NumPy and CuPy. Use `a.copy(order="K")` +to retain the layout of a C- or F-contiguous array, or `a.copy(order="F")` to +require F order. `order="K"` is not a promise to preserve arbitrary strides. + +Boolean row masks and fancy indexing create new arrays. Do not rely on those +results retaining F order. For example, a boolean row selection from an +F-ordered marker buffer produces a C-ordered result on CuPy. Normalize an +independent result before passing it to a kernel requiring F order: + +```python +alive = markers[:, 0] >= 0 +selected = xp.asfortranarray(markers[alive]) +assert selected.flags.f_contiguous +``` + +Assignment into an existing buffer keeps the destination's layout: + +```python +markers[...] = markers.copy() # C-ordered temporary, F-ordered destination +assert markers.flags.f_contiguous +``` + +Component-major arrays, `(ncomp, N)` with the markers along the last axis, keep +a C-ordered buffer and a contiguous `positions[d]` per component without any +order flag; `compact_by_mask(alive, positions, weights, axis=-1)` compacts them +in place, and the prefix `positions[:, :n]` is C order with gaps, which Pyccel +kernels take as it is (`as_kernel_array(..., strided=True)`). + +Likewise, `xp.algorithms.compact_by_mask(alive, markers)` gathers selected rows +and writes them into the front of the original buffer, preserving that +buffer's layout. Only the returned count of leading rows is valid. However, +`markers[:n_kept]` generally loses F contiguity when it trims rows: the columns +still have the original buffer's stride. Pass the full buffer plus a valid-row +count to a kernel designed for that interface, or make an F-contiguous copy +of the prefix for a kernel that requires one. A stride-aware CUDA array view +can consume the prefix directly. + +Check or normalize arrays created by concatenation, file readers, and other +libraries before passing them across a layout-sensitive boundary. + +## Compiled kernel boundaries + +Pyccel array signatures specify an expected order. A two-dimensional +`"float[:, :]"` argument expects C order; use `"float[:, :](order=F)"` for an +F-ordered argument. Every caller must supply the declared layout, including +temporary arrays and array attributes in argument objects. +`PyccelKernel` preserves F-contiguous device arguments when moving them to +the host; it does not inspect signatures or fix mismatches for you. The same +requirement applies to a `Kernel` using its host fallback without a CUDA version. +See [Host kernels with GPU data](../kernels/pyccel-kernel.md). + +CUDA `Array2D` and other `ArrayND` view arguments carry shape and strides, +so indexing through the view supports C order, F order, and strided slices. +The same applies to array view fields in `CudaStruct` and to emulated launches. +Raw pointer arguments and `CArrayND` views require C-contiguous storage. +Hand-written indexing such as `data[row * n_columns + column]` assumes C order; +use the array view's indexing to handle both layouts. See +[Index macros and array views](../kernels/cuda-kernel.md). + +Some convenience helpers deliberately require C order: + +| Helper | Contract | +| --- | --- | +| `xp.as_device_array` | Returns a C-contiguous device array, copying if needed. | +| `xp.kernels.as_kernel_array` | Returns a C-contiguous array on the side of `like`; with `strided=True`, also a NumPy array in C order with gaps, unchanged. | +| `xp.kernels.kernel_output` | Yields a C-contiguous working buffer (with `strided=True`, also a NumPy output in C order with gaps) and copies updates back into the original output when needed. | + +"C order with gaps" means positive strides, each at least the extent of the next +axis: `a[:, :n]`, `a[::2]` or `a[:, ::2]` of a C-contiguous `a`. Pyccel's +wrappers take such arrays without a copy. They refuse F-ordered arrays with a +`TypeError` and abort the process on negative strides, so `strided=True` still +copies those. + +For an F-order kernel, allocate or normalize its arguments explicitly instead +of using these helpers to prepare F-ordered inputs. + +## Check values and layout separately + +Numerical equality does not check memory order. Tests for layout-sensitive +code should assert the host argument's contiguity inside the called function, +as well as checking values. Use a two-dimensional shape with both dimensions +greater than one so C and F order are distinguishable. + +The fake CuPy supports these conversion checks without a GPU. Its +`kernel_testing.host_buffer(a)` exposes the existing NumPy storage without a +copy, preserving layout and strides. It does not exercise the `to_numpy` +transfer path; test that path separately. See [Testing kernels](../kernels/testing.md). diff --git a/docs/source/guides/data-movement.md b/docs/source/guides/data-movement.md index 07f3f78..21c7603 100644 --- a/docs/source/guides/data-movement.md +++ b/docs/source/guides/data-movement.md @@ -21,6 +21,10 @@ implicitly: every transfer is a visible function call. `to_numpy()`, on CuPy like `to_cupy()`. * None of them modify the source array or change the active backend. +CuPy-to-host conversions preserve F-contiguous layout; other device inputs +become C-contiguous host arrays. See [Array ordering and strides](array-ordering.md) +for conversion rules, copies, and compiled kernel requirements. + ## Transfer at boundaries, not in loops Load or generate data, move it to the device once, run the whole computation @@ -165,3 +169,27 @@ with xp.cuda.stream(): Pinned memory is a limited system resource; use it for large, repeatedly transferred buffers after a profile shows transfers matter. `pin_memory()` requires CuPy. + +## Host-only code: `host_call` + +Some code can only run on the host: a SciPy spline, a file reader, an external +equilibrium code. `xp.host_call(fun, *args, **kwargs)` calls it with arguments of +either backend. Device arrays are copied to the host, `fun` runs on the NumPy +backend, and array results are copied back, once per call. `@xp.evaluate_on_host` +does the same for a method, and `@xp.setup_on_host` runs an `__init__` on the NumPy +backend, so the object holds only host data. The copies are counted by +`count_transfers()`. Use it for setup and diagnostics, not in a time loop. + +```python +values = xp.host_call(spline, x) # x on the device -> values on the device + + +class Equilibrium: + @xp.setup_on_host + def __init__(self, path): + self.spline = read_spline(path) # NumPy and SciPy + + @xp.evaluate_on_host + def pressure(self, x): + return self.spline(x) +``` diff --git a/docs/source/guides/execution-helpers.md b/docs/source/guides/execution-helpers.md index fc67400..07df405 100644 --- a/docs/source/guides/execution-helpers.md +++ b/docs/source/guides/execution-helpers.md @@ -142,10 +142,24 @@ per device; catalog setup catches CUDA compiler failures before timesteps. checks actual block/grid, kernel thread, and static-plus-dynamic shared-memory limits. These execution checks are independent of scope-profiler. -GPU CI requires a real CUDA device (`CUNUMPY_REQUIRE_CUDA=1`) and runs focused +GPU CI requires a real CUDA device (`CUNUMPY_REQUIRE_CUDA=1`; the GPU markers of +`cunumpy.kernel_testing` then fail instead of skipping) and runs focused `memcheck`, `racecheck`, and `synccheck` jobs with nonzero sanitizer error exits. Numerical parity tests are separate from performance comparisons. +## Keep the live particles: `compact_by_mask` + +```python +n = xp.algorithms.compact_by_mask(alive, markers, weights) +markers, weights = markers[:n], weights[:n] +``` + +Moves the rows where the boolean mask is True to the front of every array, in +place and in order, and returns their number. The rows after the first `n` are +unspecified. The count is needed on the host, so on CuPy each call +synchronizes once. For component-major arrays, `(ncomp, N)` next to `(N,)`, +pass `axis=-1` and continue with `positions[:, :n]`. + ## Prepare cell ranges and segment reductions ```python diff --git a/docs/source/index.md b/docs/source/index.md index 801c4d9..49f1fd2 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -47,6 +47,7 @@ quickstart guides/backends guides/portable-code guides/data-movement +guides/array-ordering guides/gpu-devices guides/execution-helpers guides/mpi @@ -63,6 +64,7 @@ array-api-compat kernels/overview kernels/pyccel-kernel kernels/cuda-kernel +kernels/metal-kernel kernels/dispatch kernels/arguments kernels/accumulation diff --git a/docs/source/installation.md b/docs/source/installation.md index e9f9248..7c62fb5 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -39,10 +39,26 @@ print("active backend:", xp.get_backend()) # 'cupy' if the GPU works If `cupy_available()` is `False`, CuNumpy quietly falls back to NumPy when CuPy is requested. See [Troubleshooting](troubleshooting.md) for the usual causes. +## Apple silicon GPUs + +CuPy does not run on Macs. The GPU of an Apple silicon Mac can run +[Metal kernels](kernels/metal-kernel.md) on NumPy float32 arrays through MLX: + +```bash +python -m pip install 'cunumpy[metal]' +``` + +```python +import cunumpy as xp + +print("Metal usable:", xp.kernels.metal_available()) +``` + ## Optional extras | Extra | Installs | Use it for | | --- | --- | --- | +| `cunumpy[metal]` | `mlx` (Apple silicon Macs only) | [`MetalKernel`](kernels/metal-kernel.md) on the Mac GPU | | `cunumpy[test]` | `pytest`, `coverage` | running the test suite, using `cunumpy.kernel_testing` | | `cunumpy[test-compiled]` | the above plus `pyccel` | tests that compile host kernels with Pyccel | | `cunumpy[docs]` | Sphinx, MyST, the book theme | building this documentation | diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index e0b0949..a501a4e 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/docs/source/kernels/cuda-kernel.md @@ -99,7 +99,8 @@ axpy(2.0, x, y, x.size) # infer n_threads = x.shape[0] Explicit sizes override inference. Set `n_threads_from=lambda args: args[0].size` for a flattened element kernel, or another callback for a different axis order or logical work count. `n_threads_from=None` requires explicit sizes; -`"first_array"` always selects only the first axis. `launch_shape(args=(...))` +`"first_array"` always selects only the first axis, `"last_axis"` only the last +one (one thread per marker of component-major `(ncomp, N)` arrays). `launch_shape(args=(...))` inspects the inferred grid/block without compilation or a launch. CPU emulation and host/CUDA parity tests use the same defaults. diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md index 9cd99dd..61d7ce7 100644 --- a/docs/source/kernels/dispatch.md +++ b/docs/source/kernels/dispatch.md @@ -279,6 +279,11 @@ with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer: catalog["gather"](convert(positions), convert(field), buffer, n_threads=n) ``` +Pyccel host kernels also take arrays in C order with gaps, such as the stored +markers `storage[:, :n]` of a component-major buffer with spare capacity. Pass +`strided=True` to both helpers to hand those over without a copy on the host +(see [Array ordering and strides](../guides/array-ordering.md)). + ## Same parameters on both sides The call site is the same for both kernels only if they take the same diff --git a/docs/source/kernels/metal-kernel.md b/docs/source/kernels/metal-kernel.md new file mode 100644 index 0000000..9e0b19d --- /dev/null +++ b/docs/source/kernels/metal-kernel.md @@ -0,0 +1,105 @@ +# Metal kernels on Apple silicon + +`MetalKernel` runs a kernel written in the Metal Shading Language (MSL) on the +GPU of an Apple silicon Mac. It uses [MLX](https://github.com/ml-explore/mlx), +Apple's array framework, to compile and launch the kernel, so no Objective-C or +Xcode project is needed. It is the Mac counterpart of +[`CudaKernel`](cuda-kernel.md), with two differences: + +* it works on **NumPy arrays**: there is no Metal backend, and the kernel copies + its arguments to MLX arrays and the results back into your output arrays; +* it is **float32 only**, because the Apple GPU has no float64. + +Install it with `pip install 'cunumpy[metal]'` (MLX is installed on arm64 Macs +only). `xp.kernels.metal_available()` is `True` when MLX is installed and a Metal +GPU is present. + +## A first kernel + +```python +import numpy as np +import cunumpy as xp + +scale = xp.kernels.MetalKernel( + "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];", + inputs=["x", "a"], + outputs=["y"], +) + +x = np.arange(8, dtype=np.float32) +y = np.empty_like(x) +scale(x, 2.0, out=y) +``` + +The essentials: + +* `source` is only the **body** of the kernel function. MLX writes the + signature from `inputs` and `outputs`: each name is a pointer to the flat, + row-major data of that array, so a 2D array `a` of shape `(n, 3)` is read as + `a[3 * i + j]`. Attributes such as `thread_position_in_grid` are added to the + signature when the body uses them. +* Inputs are passed in the order of `inputs`. A Python `float` or `int` becomes + a one-element `float32` or `int32` array; read it as `a[0]`. +* `out` is one array, or one array per name in `outputs`. Their shapes and dtypes + define the outputs, and they are filled in place. The call returns them. +* `n_threads` is the **total** number of threads, not the number of threadgroups + (the opposite of the CUDA `grid`). It defaults to the first axis of the first + output. Threads beyond the data must return early, as for CUDA, if the thread + count is not a multiple of the threadgroup size. +* Outputs start uninitialized. Write every element, or pass the old array as an + input too when the kernel updates it, or give `init_value=0.0`. +* Compilation happens at the first call and MLX caches the result. + +## Compile-time constants + +`template` gives constants that are known when the kernel is compiled, which +lets the compiler unroll loops. Each name is available in the source: + +```python +push = xp.kernels.MetalKernel( + """ + uint i = thread_position_in_grid.x; + float x = pos[i], v = vel[i]; + for (int k = 0; k < NSTEPS; ++k) { x += v * dt[0]; } + pos_out[i] = x; + """, + inputs=["pos", "vel", "dt"], + outputs=["pos_out"], +) +push(pos, vel, 0.01, out=pos_out, template={"NSTEPS": 200}) +``` + +A different value of `NSTEPS` compiles another variant. `header` takes +`#include`s, `#define`s and helper functions that go before the kernel. + +## float32 and float64 + +An Apple GPU computes in float32. A float64 array passed to a `MetalKernel` +raises `TypeError` that names the argument, so no precision is lost silently. If +float32 is accurate enough for the kernel, create it with `float64="cast"`: +float64 inputs are computed in float32 and float64 outputs are filled from the +float32 results. A particle push over 200 steps agreed with a float64 reference +to about 3e-5, but whether that is acceptable depends on the physics, so check it +against the host kernel with +[`assert_kernels_agree`](testing.md) using a float32 tolerance. + +## Cost of the copies + +Every call copies the inputs to MLX arrays and the results back. The copies are +counted as `to_device` and `to_host` by +[`count_transfers()`](../guides/profiling.md), so they can be found in a profile. +On an M1, 4 million particles (3 float32 coordinates each) took about 5 ms to +copy to the GPU and 2 ms to copy back, next to a push kernel of 11.5 ms. A kernel +that runs once per time step on fresh NumPy arrays therefore pays a noticeable +share of its time in copies, and a kernel with little work per element (an +`a * x + y` update) is no faster than NumPy. The GPU pays off for kernels with +much arithmetic per element, such as pushers and interpolation. + +## What it does not do + +* It is not part of the [`Kernel`](dispatch.md) dispatch: a `Kernel` pairs a host + kernel with a CUDA kernel only, so call a `MetalKernel` directly. +* It does not keep data on the GPU between calls. +* It cannot be tested on GitHub's hosted macOS runners, which are virtual + machines without a Metal GPU. Its launch tests are skipped there and run on a + Mac with `pytest tests/unit/test_metal_kernel.py -rs`. diff --git a/docs/source/kernels/overview.md b/docs/source/kernels/overview.md index 4c892fd..c770a64 100644 --- a/docs/source/kernels/overview.md +++ b/docs/source/kernels/overview.md @@ -16,6 +16,7 @@ unchanged. | --- | --- | --- | | `PyccelKernel` | calls a host kernel with CuPy arrays by copying them to the host and back | [Host kernels with GPU data](pyccel-kernel.md) | | `CudaKernel` | wraps a CUDA C kernel, checks every call against its signature | [Writing CUDA kernels](cuda-kernel.md) | +| `MetalKernel` | runs a Metal kernel on the GPU of an Apple silicon Mac, on NumPy float32 arrays | [Metal kernels on Apple silicon](metal-kernel.md) | | `Kernel` | a host kernel plus its CUDA kernel; calls the one matching the backend | [Pairing host and CUDA kernels](dispatch.md) | | `KernelCatalog` | all `Kernel`s of a package, found by folder convention | [Pairing host and CUDA kernels](dispatch.md) | | `CudaArguments`, `CudaStruct`, `CudaStructArguments` | pass a group of arrays and scalars as one argument | [Kernel arguments and structs](arguments.md) | diff --git a/docs/source/kernels/pyccel-kernel.md b/docs/source/kernels/pyccel-kernel.md index 22db221..05315d6 100644 --- a/docs/source/kernels/pyccel-kernel.md +++ b/docs/source/kernels/pyccel-kernel.md @@ -37,6 +37,11 @@ with xp.use_backend("cupy"): smooth_kernel(field, out) # out is updated on the device ``` +Device-to-host conversion preserves F-contiguous arrays. The host kernel must +accept that order: the wrapper does not adapt arrays to a compiled signature. +See [Array ordering and strides](../guides/array-ordering.md) for Pyccel signatures +and the operations that can change layout. + ## Declare the outputs The wrapper cannot know which arguments the function writes. By default it @@ -50,10 +55,16 @@ xp.kernels.PyccelKernel(solve, outputs=("out",)) # solve(a, b, out=out) xp.kernels.PyccelKernel(norm, outputs=()) # writes nothing ``` -* Positional arguments are declared by index (negative indices count from the - end), keyword arguments by name. The two forms are not interchangeable, - because compiled functions usually do not expose a Python signature that - would map names to positions. +* Arguments are declared by index (negative indices count from the end) or by + name. A name finds the argument also when it is passed positionally, and an + index also when it is passed as a keyword, as long as the parameter names are + known: from the Python signature, from `parameters=[...]`, or, for a kernel in a + `Kernel`, from its host function. For a compiled function without any of those + the two forms are not interchangeable. +* `xp.kernels.outputs_from_annotations(function)` reads the outputs from the + annotations: every parameter that is not `Final`, `const` or a scalar. A + `Kernel.from_folder(..., outputs="annotations")` (and `KernelCatalog.from_package`) + applies it to the host function of each kernel folder. * A container or object declared as output has all its arrays copied back. * **A missing declaration is a silent bug**: if the function writes an argument that is not declared, the device array keeps its old values. When in diff --git a/docs/source/kernels/testing.md b/docs/source/kernels/testing.md index 18b135c..7922c27 100644 --- a/docs/source/kernels/testing.md +++ b/docs/source/kernels/testing.md @@ -95,8 +95,11 @@ Things to know: * **Build random data on the host.** NumPy and CuPy generators produce different sequences from the same seed, so use `numpy.random.default_rng` and convert with `to_cunumpy()`, as above. -* **Which arguments are compared**: `outputs=(2,)` selects them by index; - otherwise the host kernel's declared `outputs` are used, and if there are +* **Which arguments are compared**: `outputs=(2,)` selects them by index and + `outputs=("markers",)` by parameter name (the names of the host function). + `"markers.positions"` compares only that field of a struct or argument + object, leaving out fields the two kernels fill differently (scratch buffers, + for instance); otherwise the host kernel's declared `outputs` are used, and if there are none, every array argument. Arrays held by argument objects (one level deep, e.g. a `CudaArguments` object or a list) are compared too. A `CudaStructArguments` object is read through its struct fields, so its @@ -206,16 +209,77 @@ def test_gather_cuda_arithmetic(): np.testing.assert_allclose(result, expected, rtol=1e-12, atol=1e-14) ``` -Arrays are passed as NumPy arrays (any strides) and written back; scalars are -checked like in a launch. It catches wrong indices, clamping, periodic wrapping +Arrays are passed as NumPy arrays (any strides) and the kernel writes into +them; scalars are checked like in a launch. Each kernel is compiled once into a +shared library and cached (in the process and in `~/.cache/cunumpy/emulation`, +see `emulation_cache_dir()`), so further launches with other values and sizes +cost no compilation. It catches wrong indices, clamping, periodic wrapping and weights, i.e. most porting bugs of gather, scatter and push kernels. Block shared memory and `__syncthreads` are emulated (pass `shared_mem=` for `extern __shared__` arrays), so per-block deposits and shared-memory reductions are covered too. It does not emulate concurrency between barriers or warp intrinsics; kernels using the latter are refused with `NotImplementedError`, so -those still need a GPU run. The compiler may fuse multiply-adds as NVRTC does, +those still need a GPU run. Inline PTX is not emulated either: `asm(...)` and +`asm volatile(...)` compile but trap when reached, so a PTX branch a test never +takes needs no extra options, and reaching it raises `RuntimeError`. The compiler may fuse multiply-adds as NVRTC does, so compare with a tolerance of a few ulp. +### Struct parameters and `emulated_launches()` + +A kernel with struct parameters takes, for each struct, a dictionary of field +names to values, a `CudaStructValue`, or any object with an attribute per field +(a host argument class, a `CudaStructArguments`); the arrays in it are updated in +place: + +```python +emulate_cuda_kernel( + push, + {"markers": markers, "alive": alive, "n": 4}, # the struct Particles + 0.5, + total, + n_threads=4, +) +``` + +To test code that *launches* kernels (a `Kernel` on the CuPy backend, a +propagator), wrap it in `emulated_launches()`. With the fake CuPy (below) every +`CudaKernel` launch in the block then runs through the emulation on the host +buffers of the fake arrays, so the device arrays hold the results: + +```python +from cunumpy.kernel_testing import emulated_launches, host_buffer + +with emulated_launches(): + propagator(dt) # CuPy backend = the fake CuPy +np.testing.assert_allclose(host_buffer(markers), expected) +``` + +`host_buffer(array)` is the NumPy array behind a fake CuPy array (not a copy). +Launches are serial, so use small problems. + +No CUDA is compiled inside the block either: `kernel.compile()`, +`recompile()` and the `compile_all()` methods build the emulation library +instead (and return None), so a program that compiles its kernels before the +time loop runs unchanged and still reports compile errors there. +`fake_cupy_session()` does all of it for a whole program, activating the CuPy +backend too: + +```python +from cunumpy.kernel_testing import fake_cupy_session, requires_device_backend + + +@requires_device_backend # a GPU, or the fake CuPy +def test_simulation_on_the_cupy_backend(): + with fake_cupy_session(): # only with the fake CuPy + sim.run() +``` + +The fake CuPy must be installed before cunumpy is used, so a test process +that runs on NumPy cannot switch to it. `run_in_fake_cupy_subprocess(code)` +runs the code in a serial child process on the fake CuPy (outside the MPI +job, only on rank 0 under MPI) and fails the test with the signal or exit code +and the end of the child's output if it fails. + ## Test `__device__` helpers: `device_function_kernel` Helpers such as B-spline evaluation or coordinate maps are `__device__` @@ -314,6 +378,13 @@ This catches `to_numpy()` calls, `PyccelKernel` conversions and `Kernel` fallbacks that crept into the step. It does not see copies made outside CuNumpy (see [Data movement](../guides/data-movement.md)). +Syncs (the host waiting for the device) are recorded as well, in +`counter.syncs`, and are accepted unless you ask for none: +`assert_no_transfers(syncs=True)`. They include `xp.synchronize()` and the waits +of the MPI helpers; on the fake CuPy also `float(a)`, `int(a)`, `bool(a)`, +`a.item()` and `a.tolist()`, which stall the real CuPy too but cannot be +observed there from Python. + ## Test generated headers When struct headers are generated with `write_cuda_header()` and committed, @@ -324,6 +395,12 @@ offsets. ## CI setup +* With `CUNUMPY_REQUIRE_CUDA=1` the GPU markers fail instead of skipping: + `requires_cupy` (as an error when the test is set up), the `cupy` run of the + `backend` fixture (which also activates CuPy strictly, never falling back to + NumPy) and `assert_kernels_agree`. Set it on the GPU CI job, so a broken CuPy + or driver cannot pass as a set of skipped tests. + * Run the suite on a normal CPU runner: everything on NumPy runs, GPU cases are reported as skipped, and `emulate_cuda_kernel` tests check the CUDA kernels' arithmetic (the runner needs a C++ compiler, which Linux images diff --git a/pyproject.toml b/pyproject.toml index 7dbd7ee..285ef40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ] [project] name = "cunumpy" -version = "0.6.1" +version = "0.6.3" description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`." readme = "README.md" keywords = [ "python" ] @@ -43,6 +43,7 @@ optional-dependencies.docs = [ "sphinx", "sphinx-book-theme", ] +optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] optional-dependencies.test = [ "coverage", "pytest", "scipy" ] optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ] urls."Source" = "https://github.com/max-models/cunumpy" diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 65a6a8b..7adcf62 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -72,21 +72,28 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | hand data to SciPy/matplotlib/h5py | `xp.to_numpy(a)` | | call an existing NumPy-only kernel with GPU arrays (slow, correct) | `xp.kernels.PyccelKernel(fn, outputs=(...))` | | launch a hand-written CUDA C kernel | `xp.kernels.CudaKernel(source, "name")` / `CudaKernel.from_file(path)` | +| run a Metal (MSL) kernel on an Apple silicon GPU, NumPy float32 in and out (float64 raises; `float64="cast"` computes in float32); not part of `Kernel` dispatch | `xp.kernels.MetalKernel(body, inputs=[...], outputs=[...])(*args, out=arrays, n_threads=n)`; check `xp.kernels.metal_available()` | +| call host-only code (SciPy, file readers) with arguments of either backend | `xp.host_call(fun, *args)`; `@xp.evaluate_on_host` on a method; `@xp.setup_on_host` on `__init__` | +| keep the live rows of particle arrays at the front | `n = xp.algorithms.compact_by_mask(alive, markers, weights)`; component-major `(ncomp, N)`: `axis=-1` | +| run CUDA kernel launches on the CPU in a test (fake CuPy) | `with kernel_testing.emulated_launches(): ...` (also makes `kernel.compile()`/`compile_all()` build the emulation instead of CUDA); `with kernel_testing.fake_cupy_session(): ...` adds the CuPy backend; `kernel_testing.host_buffer(a)` reads a fake array; struct arguments are read through their fields | +| run a test on the fake CuPy from a NumPy test process | `kernel_testing.run_in_fake_cupy_subprocess(code)` (serial child, rank 0 only under MPI); `@kernel_testing.requires_device_backend` = a GPU or the fake CuPy | +| find the arguments a host kernel writes (copy only those back) | `PyccelKernel(fn, outputs=("out",))` (names work positionally); `xp.kernels.outputs_from_annotations(fn)`; `Kernel.from_folder(..., outputs="annotations")` | | host kernel + CUDA port, chosen by backend | `xp.kernels.Kernel(host_fn, cuda_kernel_or_None)` | | many kernels in a package, ported incrementally | `xp.kernels.KernelCatalog.from_package(__name__, missing_cuda="fallback")` | | host kernels compiled at first call (your compile function), NumPy fallback | `from_package(..., host_suffix="_pyccel", compile_host=my_compile, host_fallback={...})` -> `xp.kernels.CompiledHostKernel` | | host arrays reach kernels while CuPy is active | `Kernel(..., dispatch="arrays")` / `from_package(..., dispatch="arrays")`: CUDA only for device arguments | | one kernel folder declares its kernel in its own `__init__.py` | `kernel = xp.kernels.Kernel.from_folder(__name__, host_suffix="_pyccel", compile_host=..., dispatch="arrays")`; `_numba.py`, `_numpy.py` in the folder are further host implementations | -| bring a `dispatch="arrays"` kernel's arguments to the side of the main array | `xp.kernels.as_kernel_array(a, like=grid, dtype=float)`; outputs: `with xp.kernels.kernel_output(out, like=grid, dtype=float) as buf:` | +| bring a `dispatch="arrays"` kernel's arguments to the side of the main array | `xp.kernels.as_kernel_array(a, like=grid, dtype=float)`; outputs: `with xp.kernels.kernel_output(out, like=grid, dtype=float) as buf:`; `strided=True` passes NumPy views in C order with gaps (`a[:, :n]`) to Pyccel hosts without a copy | | choose the host implementation (pyccel/numba/numpy/python) | `xp.kernels.set_host_kernel_implementation("numpy")`, `with xp.kernels.use_host_kernel_implementation(...)`, `CUNUMPY_HOST_KERNEL_IMPLEMENTATION=numpy`; default: first available of pyccel, numba, numpy; `kernel.selected()` | | require CUDA for device kernel dispatch | `xp.kernels.set_device_kernel_implementation("cuda")`, `get_device_kernel_implementation()`, `with xp.kernels.use_device_kernel_implementation(...)`, `CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION=cuda`; default `None` preserves `missing_cuda` policy; explicit CUDA rejects host fallback | | check host and CUDA kernels take the same parameters | `catalog.check_signatures()` (in a unit test) | -| test a CUDA kernel's arithmetic without a GPU | `cunumpy.kernel_testing.emulate_cuda_kernel(kernel, *numpy_args, n_threads=n)` (C++ compiler; shared memory and __syncthreads ok, no warp ops; `shared_mem=` for extern shared) | +| test a CUDA kernel's arithmetic without a GPU | `cunumpy.kernel_testing.emulate_cuda_kernel(kernel, *numpy_args, n_threads=n)` (C++ compiler; shared memory and __syncthreads ok, no warp ops; inline asm traps; compiled once per kernel and cached; `shared_mem=` for extern shared) | | shared-memory budget of a block | `xp.cuda.max_shared_memory_per_block()` (48 KiB without a GPU) | | random numbers inside a kernel, equal on the host | `#include `: `cunumpy_uniform(seed, particle_id, step)`; host: `xp.rng.philox_uniform(seed, ids, step)` | | sort points along a Z-curve / quadtree or octree nodes as contiguous ranges | `keys = xp.algorithms.morton_keys(pos, lower, upper, levels)`, `keys, order, pos = xp.algorithms.sort_by_key(keys, pos)`; in a kernel `#include `: `cunumpy_morton_key2(x, y, x0, y0, sx, sy, levels)` with `xp.algorithms.morton_scales(...)` | | one thread per marker without passing n_threads | default: `CudaKernel(...)` uses the first array's row count for 1D launches | | copy device arrays to the host for output without stalling | `xp.memory.HostStaging(shape, dtype)`: `c = staging.copy(a)` ... `c.result()` | +| read a device scalar without a sync (e.g. a convergence test one iteration late) | `h = xp.to_host_async(x)`; `h.ready()` never waits, `h.result()` | | PIC recipes (compaction, sort by cell, MPI exchange, graphs) | docs guide "Particle codes" | | reproducible random numbers per MPI rank | `xp.rng.random_streams.seed(seed, rank=rank)`, then `xp.rng.random_streams.normal(...)` / `.generator()` | | group arrays/scalars into one kernel argument | `xp.arguments.CudaArguments` (flattened), `xp.arguments.CudaStruct` (C struct), `xp.arguments.CudaStructArguments` (C struct as a class); host kernels take their own argument objects, the caller picks one per backend | @@ -101,6 +108,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | timing GPU code | `with xp.profiling.timed_region("name") as t:` → `t.elapsed` | | profiler markers | `xp.profiling.nvtx_range("name")` (context manager or decorator) | | find transfers | `with xp.profiling.count_transfers() as c: ...; print(c.report())` | +| check transfers per phase of a time loop | `b = xp.profiling.TransferBudget()`; `b.count("step")(fn)` or `with b.phase("output"):`; `b.require("step", allow={"to_host": {"max_nbytes": 8}})`; `b.check()` | | debug an illegal memory access | `CUNUMPY_CUDA_DEBUG=1`, then `compute-sanitizer` | | test on both backends | `cunumpy.kernel_testing.BACKENDS`, `backend` fixture, `requires_cupy` | | test CUDA vs host kernel | `cunumpy.kernel_testing.assert_kernels_agree(kernel, make_args, n_threads=...)` | @@ -137,6 +145,9 @@ xp.as_device_array(value, dtype=None, ndim=None, *, name=None) with xp.profiling.count_transfers() as c: ... # c.total, c.to_host, c.to_device, # c.kernel_conversions, c.fallbacks, c.events, c.report() with xp.profiling.assert_no_transfers(): ... # rejects host/device copies and host fallback +h = xp.to_host_async(a) # non-blocking copy of a device scalar: h.ready(), h.result() +budget = xp.profiling.TransferBudget() # per-phase counts: budget.phase(name), + # budget.count(name)(fn), budget.require(...), budget.check() ``` Mirror refreshes, MPI/output staging, argument conversion and kernel output @@ -300,8 +311,10 @@ xp.cuda.cuda_include_dir() * Default `n_threads_from="auto"` infers from the first array, including arrays in supported argument objects: 1D block -> shape[0], 2D -> shape[:2], 3D -> shape[:3] (x, y, z). Per-call block overrides apply. Explicit n_threads/grid - wins; use a callback for flattened `.size`, reversed axes or another work - count, or None to require explicit sizes. Missing arrays/axes raise. + wins; `"last_axis"` uses the first array's last axis (component-major + `(ncomp, N)` marker arrays); use a callback for flattened `.size`, reversed + axes or another work count, or None to require explicit sizes. Missing + arrays/axes raise. Shipped CUDA headers (always on the include path): diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 1125dcc..f4db38d 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -14,7 +14,9 @@ rng, xp, ) +from cunumpy._host import evaluate_on_host, host_call, setup_on_host from cunumpy._scipy_backend import scipy +from cunumpy._staging import to_host_async from cunumpy.xp import ( as_device_array, assert_same_backend, @@ -79,9 +81,11 @@ def require_version(minimum: str) -> None: "cupy_available", "cupy_backend", "default_float_dtype", + "evaluate_on_host", "get_array_backend", "get_array_module", "get_backend", + "host_call", "is_cpu", "is_gpu", "kernels", @@ -95,9 +99,11 @@ def require_version(minimum: str) -> None: "same_backend", "scipy", "set_backend", + "setup_on_host", "synchronize", "to_cunumpy", "to_cupy", + "to_host_async", "to_numpy", "use_backend", "xp", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 52ab855..a432d57 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -1,7 +1,7 @@ # 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 collections.abc import Callable, Generator from contextlib import contextmanager from typing import Any @@ -19,10 +19,12 @@ from cunumpy import profiling as profiling from cunumpy import rng as rng from cunumpy import xp as xp from cunumpy._scipy_backend import scipy as scipy +from cunumpy._staging import HostCopy def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... def to_cunumpy(array: Any) -> Any: ... +def to_host_async(array: Any) -> HostCopy: ... def cupy_available() -> bool: ... def backend_info() -> dict[str, Any]: ... def get_array_module(array: Any) -> Any: ... @@ -38,6 +40,9 @@ def set_backend(backend: str, *, strict: bool = ...) -> None: ... def require_version(minimum: str) -> None: ... def default_float_dtype() -> Any: ... def synchronize() -> None: ... +def host_call(fun: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: ... +def evaluate_on_host(method: Callable[..., Any]) -> Callable[..., Any]: ... +def setup_on_host(init: Callable[..., None]) -> Callable[..., None]: ... def as_device_array( value: Any, dtype: Any = ..., ndim: int | None = ..., *, name: str | None = ... ) -> Any: ... diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index 353c83a..5127b64 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -212,7 +212,42 @@ def segment_sum(values: Any, keys: Any, n_segments: int, *, out: Any = None) -> return SegmentPlan(keys, n_segments).sum(values, out=out) -def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: +#: Integer keys of more entries than this are sorted by :func:`_radix_argsort` on NumPy. +_RADIX_MIN_SIZE = 4096 + + +def _radix_argsort(keys: np.ndarray) -> np.ndarray: + """Stable argsort of integer `keys`: least significant 16 bits first. + + NumPy's stable argsort is a radix sort for 16-bit integers only; wider + integers get a timsort, about ten times slower on a million random keys. + Sorting the 16-bit digits of ``keys - keys.min()`` one after the other, + each pass stable, gives the same order (an LSD radix sort), in as many + passes as the range of the keys needs (two for up to 2**32 cells). + """ + low = keys.min() + span = int(keys.max()) - int(low) + if keys.dtype.kind == "i": + # the span of int64 keys fits uint64; shift them to start at zero + shifted = keys.astype(np.int64, copy=False).view(np.uint64) - np.uint64( + np.int64(low).view(np.uint64) + ) + else: + shifted = keys.astype(np.uint64, copy=False) - np.uint64(low) + order = None + shift = 0 + while True: + digits = shifted if order is None else shifted[order] + digits = (digits >> np.uint64(shift)).astype(np.uint16) + step = np.argsort(digits, kind="stable") + order = step if order is None else order[step] + shift += 16 + if span >> shift == 0: + break + return order.astype(np.int64, copy=False) + + +def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]: """Sort `keys` and reorder every array the same way, in one stable argsort. The usual first step of a particle code on the GPU: sort the particles by @@ -222,20 +257,35 @@ def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, charges) + Component-major marker arrays, ``(ncomp, N)`` next to ``(N,)`` scalars, keep + the markers along the last axis of each, so sort along that one:: + + keys, order, positions, weights = xp.algorithms.sort_by_key( + keys, positions, weights, axis=-1 + ) + + On NumPy, integer keys (cell indices, Morton keys) are sorted by a radix + sort on their 16-bit digits, as on CuPy: two passes for up to ``2**32`` + distinct cells, about ten times faster than NumPy's stable sort of 64-bit + integers. + Parameters ---------- keys : array, shape (n,) The sort keys. *arrays : arrays - Arrays with ``n`` rows, on the backend of `keys`, reordered along - axis 0. + Arrays with ``n`` entries along `axis` (any other axes), on the backend + of `keys`. + axis : int + The axis of every array that `keys` indexes, the first by default; + ``-1`` is the last axis of each array, whatever its number of axes. Returns ------- tuple ``(sorted_keys, order, *sorted_arrays)``: ``order`` (int64) is the permutation, ``sorted_keys = keys[order]``, and each sorted array is - ``array[order]`` (a new array). + ``take(array, order, axis=axis)`` (a new array). """ if get_array_backend(keys) == "cupy": import cupy as xpm # its argsort is a stable radix sort @@ -244,10 +294,94 @@ def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: keys = xpm.asarray(keys) if keys.ndim != 1: raise ValueError(f"keys must be 1D, got shape {keys.shape}") - for array in arrays: - if array.shape[:1] != keys.shape: + axes = [] + for i, array in enumerate(arrays): + if not -array.ndim <= axis < array.ndim: + raise ValueError( + f"array {i} has {array.ndim} axes, so it has no axis {axis}", + ) + array_axis = axis % array.ndim + if array.shape[array_axis] != keys.shape[0]: + raise ValueError( + f"array {i} has shape {array.shape}, expected {keys.shape[0]} " + f"entries along axis {axis}", + ) + axes.append(array_axis) + if xpm is np and keys.dtype.kind in "iu" and keys.size > _RADIX_MIN_SIZE: + order = _radix_argsort(keys) + else: + order = xpm.argsort(keys, kind="stable").astype(xpm.int64, copy=False) + return ( + keys[order], + order, + *( + xpm.take(array, order, axis=array_axis) + for array, array_axis in zip(arrays, axes, strict=True) + ), + ) + + +def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int: + """Move the entries where `mask` is True to the front of every array, in place. + + The usual step after particles left the domain or were absorbed: keep the + live ones at the front of the marker array (and of the arrays that go with + it) and continue with ``markers[:n]``. The order of the kept entries is + preserved, so the result is reproducible:: + + n = xp.algorithms.compact_by_mask(alive, markers, weights) + markers, weights = markers[:n], weights[:n] + + Component-major marker arrays, ``(ncomp, N)`` next to ``(N,)`` scalars, keep + the markers along the last axis of each, so compact that one:: + + n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1) + positions, weights = positions[:, :n], weights[:n] + + Parameters + ---------- + mask : array of bool, shape (n,) + True for the entries to keep. + *arrays : arrays + Arrays with ``n`` entries along `axis` (any other axes), on the backend + of `mask`. Entries ``[:count]`` along `axis` hold the kept ones + afterwards; the entries after them are unspecified, so ignore them (or + overwrite them). + axis : int + The axis of every array that `mask` indexes, the first by default; + ``-1`` is the last axis of each array, whatever its number of axes. + + Returns + ------- + int + The number of kept entries. Its value is needed on the host, so on CuPy + the call synchronizes once per call (counted by + :func:`~cunumpy.profiling.count_transfers` where it can be seen). + """ + assert_same_backend(mask, *arrays) + xpm = get_array_module(mask) + mask = xpm.asarray(mask) + if mask.ndim != 1 or mask.dtype != np.bool_: + raise TypeError( + f"mask must be a 1D boolean array, got dtype {mask.dtype}, {mask.ndim}D", + ) + axes = [] + for i, array in enumerate(arrays): + if not -array.ndim <= axis < array.ndim: + raise ValueError( + f"array {i} has {array.ndim} axes, so it has no axis {axis}", + ) + array_axis = axis % array.ndim + if array.shape[array_axis] != mask.shape[0]: raise ValueError( - f"every array needs {keys.shape[0]} rows, got shape {array.shape}", + f"array {i} has shape {array.shape}, expected {mask.shape[0]} " + f"entries along axis {axis}", ) - order = xpm.argsort(keys, kind="stable").astype(xpm.int64, copy=False) - return (keys[order], order, *(array[order] for array in arrays)) + axes.append(array_axis) + rows = xpm.nonzero(mask)[0] + n_kept = int(rows.size) + for array, array_axis in zip(arrays, axes, strict=True): + front = (slice(None),) * array_axis + (slice(0, n_kept),) + # the right side is a copy: no overlap problem + array[front] = xpm.take(array, rows, axis=array_axis) + return n_kept diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index fafd890..80b203d 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -66,6 +66,9 @@ class is the one definition of the arguments. import numpy as np +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record_sync + __all__ = [ "DEBUG_OPTIONS", "CudaArguments", @@ -1794,6 +1797,20 @@ def _first_array_length(args: tuple[Any, ...]) -> int: return typing.cast(int, _first_array_shape(args, 1)) +def _last_axis_length(args: tuple[Any, ...]) -> int: + """``n_threads_from="last_axis"``: the last axis of the first array argument. + + One thread per entry of a component-major array, ``(ncomp, N)`` or ``(N,)``. + """ + shape = next(_array_shapes_in(args), None) + if shape is None: + raise TypeError( + "inferring n_threads needs an array argument; pass n_threads or grid, " + "or set n_threads_from to a callable", + ) + return int(shape[-1]) + + def _as_shape(value: int | Sequence[int], what: str) -> tuple[int, ...]: shape = (value,) if isinstance(value, (int, np.integer)) else tuple(value) if not 1 <= len(shape) <= 3: @@ -1873,11 +1890,12 @@ class CudaKernel: (:func:`cunumpy.cuda.set_cuda_debug`, ``CUNUMPY_CUDA_DEBUG``) at every launch; True or False fix it for this kernel. The compile options are fixed when the kernel is compiled. - n_threads_from : {"auto", "first_array"} | callable | None + n_threads_from : {"auto", "first_array", "last_axis"} | callable | None Default "auto" infers thread counts from the first array's leading shape axes, matching the block dimensionality (1D: one thread per row). Arrays in supported argument objects are included. "first_array" - always uses the first axis. A callable receives the positional argument + always uses the first axis, "last_axis" the last one (one thread per + entry of a component-major ``(ncomp, N)`` or ``(N,)`` array). A callable receives the positional argument tuple; None requires an explicit launch size. Explicit `n_threads` or `grid` overrides inference. @@ -2157,7 +2175,9 @@ def n_threads_from(self) -> Callable[[tuple[Any, ...]], Any] | None: Default ``"auto"`` uses the first array's leading axes, matching the launch block dimensions; 1D launches use its first axis, one thread per row. Supported argument objects are searched in field/argument order. - ``"first_array"`` always uses its first axis. None disables inference. + ``"first_array"`` always uses its first axis, ``"last_axis"`` its last + one (one thread per entry of a component-major ``(ncomp, N)`` or + ``(N,)`` array). None disables inference. Settable, also on the ``cuda_kernel`` of a :class:`~cunumpy.kernels.Kernel`. """ return self._n_threads_from @@ -2172,9 +2192,12 @@ def n_threads_from( value = self._default_n_threads elif isinstance(value, str) and value == "first_array": value = _first_array_length + elif isinstance(value, str) and value == "last_axis": + value = _last_axis_length if value is not None and not callable(value): raise TypeError( - "n_threads_from must be callable, 'auto', 'first_array' or None" + "n_threads_from must be callable, 'auto', 'first_array', " + "'last_axis' or None" ) self._automatic_threads = automatic self._n_threads_from = value @@ -2421,6 +2444,10 @@ def __call__( if shared_mem < 0: raise ValueError(f"shared_mem must be non-negative, got {shared_mem}") values = self.prepare_args(*args) + # RawKernel accepts size-one NumPy arrays for structs passed by value, + # but not structured NumPy scalars (np.void). Keep their packed bytes + # and alignment intact, including array-view pointers and strides. + values = tuple(np.asarray(v) if isinstance(v, np.void) else v for v in values) if 0 in grid_shape: return @@ -2528,6 +2555,8 @@ def _synchronize_after_launch( stream = cp.cuda.get_current_stream() if _is_capturing(stream): return + if _COUNTERS: + _record_sync(f"debug synchronization after kernel {self.expression!r}") try: stream.synchronize() except Exception as error: diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index ed3e388..305082f 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -40,6 +40,7 @@ HostImplementations, PyccelKernel, get_device_kernel_implementation, + outputs_from_annotations, ) from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import _record @@ -219,6 +220,10 @@ def __init__( "host_options are for wrapping a plain callable; configure the " "given PyccelKernel directly", ) + if host_kernel._parameters is None: + # compiled kernels have no Python signature: name the parameters + # from the host function, so that outputs may be given by name + host_kernel._parameters = self.host_parameters self._host_kernel = host_kernel self._cuda_kernel = cuda_kernel self._name = name if name is not None else host_kernel.name @@ -239,6 +244,7 @@ def from_folder( check_name_length: bool = True, missing_cuda: str = "raise", host_options: Mapping[str, Any] | None = None, + outputs: Sequence[int | str] | str | None = None, include_dirs: Sequence[str | Path] | None = None, dispatch: str = "backend", compile_host: Callable[[Any], Any] | None = None, @@ -285,6 +291,14 @@ def from_folder( As for :meth:`KernelCatalog.from_package`. missing_cuda, host_options, dispatch Passed on to :class:`Kernel`. + outputs : Sequence[int | str] | "annotations" | None + The arguments the host kernel writes to (names or indices, see + :class:`~cunumpy.kernels.PyccelKernel`), so that only those arrays are + copied back to the device on the fallback path. ``"annotations"`` + reads them from the annotations of the host function: every + parameter that is not ``Final`` or a scalar is an output (see + :func:`~cunumpy.kernels.outputs_from_annotations`). A value in + `host_options` takes precedence. include_dirs : Sequence[str | Path] | None Include directories of the CUDA kernel, in addition to the folder itself; by default the source root of the top-level package (the @@ -334,6 +348,19 @@ def from_folder( test_args = f"{package}.{name}{test_args_suffix}" module = importlib.import_module(f"{package}.{name}{host_suffix}") python = getattr(module, name) + if outputs is not None and "outputs" not in (host_options or {}): + declared = ( + outputs_from_annotations(python) + if isinstance(outputs, str) and outputs == "annotations" + else outputs + ) + if isinstance(outputs, str) and outputs != "annotations": + raise ValueError( + f"outputs must be a sequence of names/indices or 'annotations', " + f"got {outputs!r}", + ) + if declared is not None: + host_options = {**(host_options or {}), "outputs": declared} loaders: dict[str, Callable[[], Callable[..., Any]]] = { "python": lambda: python, "pyccel": ( @@ -669,6 +696,12 @@ def from_package( host_options: ( Mapping[str, Any] | Callable[[str], Mapping[str, Any]] | None ) = None, + outputs: ( + Sequence[int | str] + | str + | Callable[[str], Sequence[int | str] | str | None] + | None + ) = None, include_dirs: Sequence[str | Path] | None = None, dispatch: str = "backend", compile_host: Callable[[Any], Any] | None = None, @@ -713,6 +746,10 @@ def from_package( Keyword arguments for the :class:`~cunumpy.kernels.PyccelKernel` wrapping each host kernel (see :class:`Kernel`): the same for all kernels, or a function of the kernel name, e.g. to declare per-kernel ``outputs``. + outputs : Sequence[int | str] | "annotations" | Callable | None + The arguments every host kernel writes to, or ``"annotations"`` to + read them from the annotations of each host function, or a function + of the kernel name returning either (see :meth:`Kernel.from_folder`). include_dirs : Sequence[str | Path] | None Include directories of the CUDA kernels, in addition to each kernel's own folder. By default the source root of the top-level @@ -757,6 +794,7 @@ def from_package( host_options=( host_options(name) if callable(host_options) else host_options ), + outputs=outputs(name) if callable(outputs) else outputs, include_dirs=include_dirs, dispatch=dispatch, compile_host=compile_host, diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index 56d47d0..54ac02c 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -5,9 +5,9 @@ check, and the kernel's index and weight arithmetic go untested. :func:`emulate_cuda_kernel` closes that gap: it compiles the kernel source as C++ with the CUDA built-ins replaced by plain C++ (``threadIdx``, ``blockIdx``, -``blockDim``, ``gridDim``, ``atomicAdd``, ...), calls the kernel once per -thread, serially, on copies of the NumPy arguments, and copies the arrays back, -so the call looks like a launch:: +``blockDim``, ``gridDim``, ``atomicAdd``, ...) into a shared library, and calls +the kernel once per thread, serially, on the NumPy arguments, so the call looks +like a launch:: from cunumpy.kernel_testing import emulate_cuda_kernel @@ -16,9 +16,18 @@ np.testing.assert_allclose(y, 2.0 * x) Arguments follow the kernel signature: NumPy arrays for pointer and array view -parameters (``Array1D`` to ``Array16D``; any strides, they are passed as -contiguous copies), Python or NumPy scalars for scalar parameters (cast and -checked like in a launch). Arrays are written back into the given arrays. +parameters (``Array1D`` to ``Array16D``, any strides), Python or NumPy +scalars for scalar parameters (cast and checked like in a launch). The kernel +writes into the given arrays. + +Each kernel is compiled once: the library takes the arguments at run time +(``cunumpy_launch(args, grid, block, shared_mem)``, loaded with ctypes), so +launches with other values and array sizes reuse it. Libraries are cached in +the process and on disk (:func:`emulation_cache_dir`), keyed on the generated +source, the included headers, the compiler and its options. +:func:`compile_for_emulation` builds one without a launch. The launch runs in +the Python process: ``__trap()`` (a failed bounds check) is reported as a +``RuntimeError``, but a segmentation fault ends the process. Block shared memory and ``__syncthreads`` are emulated: ``__shared__`` variables (and ``extern __shared__`` arrays, sized by ``shared_mem``) are one @@ -28,11 +37,23 @@ GPU. Per-block deposits, shared-memory reductions and tiled kernels therefore work. +Struct parameters take, for each struct, a mapping from field name to value, a +:class:`~cunumpy.arguments.CudaStructValue`, or any object with an attribute +per field (e.g. a :class:`~cunumpy.arguments.CudaStructArguments` or a host +argument class); the arrays of the fields are updated in place like the others. + +:func:`emulated_launches` makes every :class:`~cunumpy.kernels.CudaKernel` launch +(and compilation) in a block run through the emulation, on the arrays of the +fake CuPy (:mod:`cunumpy._fake_cupy`), so that code that launches kernels (a +:class:`~cunumpy.kernels.Kernel` on the CuPy backend, a propagator) can be +tested without a GPU. + What it does not emulate: concurrency between the barriers (threads run one after another, so atomics are plain additions and races never show), warp intrinsics (``__shfl_*``, ``__syncwarp``, ``__ballot_sync``, ...; a kernel or an -included header using them is refused), structs and ``CudaArguments`` objects -(not supported), and ````. Use a GPU for those. +included header using them is refused), ````, and inline PTX: +``asm(...)`` and ``asm volatile(...)`` compile, but trap when reached. Use a +GPU for those. Floating point: like NVRTC (``--fmad=true`` by default) the C++ compiler may fuse ``a * b + c`` into one fused multiply-add, so results can differ from @@ -44,26 +65,40 @@ from __future__ import annotations +import ctypes +import hashlib import os import re import shutil import subprocess import tempfile -from collections.abc import Sequence +import threading +import weakref +from collections.abc import Generator, Mapping, Sequence +from contextlib import contextmanager from pathlib import Path from typing import Any import numpy as np +from cunumpy import _fake_cupy from cunumpy._cuda_kernel import ( CudaKernel, CudaParameter, + CudaStructArguments, + CudaStructValue, _scalar_checker, _strip_comments, cuda_include_dir, ) -__all__ = ["emulate_cuda_kernel", "emulation_compiler"] +__all__ = [ + "compile_for_emulation", + "emulate_cuda_kernel", + "emulated_launches", + "emulation_cache_dir", + "emulation_compiler", +] # CUDA constructs that serial emulation would get wrong _UNSUPPORTED = { @@ -75,6 +110,8 @@ _STUBS = r""" // --- cunumpy emulation: CUDA built-ins as plain C++, one thread at a time --- #include +#include +#include #include #include #include @@ -93,7 +130,17 @@ struct cunumpy_dim3 { unsigned int x, y, z; }; static cunumpy_dim3 threadIdx, blockIdx, blockDim, gridDim; static const int warpSize = 32; -inline void __trap() { abort(); } +// __trap() ends the launch: cunumpy_launch returns 1 (in a coroutine, the hook +// returns to the scheduler first) +static jmp_buf cunumpy_trap_jump; +static int cunumpy_trapped = 0; +static void (*cunumpy_trap_hook)() = 0; +[[noreturn]] inline void __trap() { + cunumpy_trapped = 1; + fflush(stdout); + if (cunumpy_trap_hook) cunumpy_trap_hook(); + longjmp(cunumpy_trap_jump, 1); +} template inline T __ldg(const T* p) { return *p; } inline double rsqrt(double v) { return 1.0 / sqrt(v); } inline float rsqrtf(float v) { return 1.0f / sqrtf(v); } @@ -114,16 +161,11 @@ inline int atomicCAS(int* p, int c, int v) { int o = *p; if (o == c) *p = v; return o; } inline long long __double_as_longlong(double v) { long long r; memcpy(&r, &v, sizeof r); return r; } inline double __longlong_as_double(long long v) { double r; memcpy(&r, &v, sizeof r); return r; } -static void cunumpy_read(const char* path, void* data, size_t bytes) { - FILE* f = fopen(path, "rb"); - if (!f || fread(data, 1, bytes, f) != bytes) { perror(path); exit(2); } - fclose(f); -} -static void cunumpy_write(const char* path, const void* data, size_t bytes) { - FILE* f = fopen(path, "wb"); - if (!f || fwrite(data, 1, bytes, f) != bytes) { perror(path); exit(2); } - fclose(f); -} +// inline PTX cannot run on the CPU: `asm(...)` and `asm volatile(...)` trap +static int cunumpy_inline_asm; +#define cunumpy_inline_asm(...) __trap() +#define asm cunumpy_inline_asm +#define volatile(...) , __trap() // --------------------------------------------------------------------------- """ @@ -139,15 +181,14 @@ inline void __syncthreads() { swapcontext(&cunumpy_threads[cunumpy_current].ctx, &cunumpy_scheduler); } +static void cunumpy_coroutine_trap() { + cunumpy_threads[cunumpy_current].done = 1; + setcontext(&cunumpy_scheduler); +} """ -_MAIN_SERIAL = r""" -{globals} -int main() {{ -{inits} - cunumpy_dynamic_shared = calloc({shared_mem} + 16, 1); - gridDim = {{{gx}u, {gy}u, {gz}u}}; - blockDim = {{{bx}u, {by}u, {bz}u}}; +_RUN_SERIAL = r""" +static void cunumpy_run() {{ for (unsigned int bz = 0; bz < gridDim.z; ++bz) for (unsigned int by = 0; by < gridDim.y; ++by) for (unsigned int bx = 0; bx < gridDim.x; ++bx) @@ -158,26 +199,20 @@ threadIdx = {{tx, ty, tz}}; {call}; }} -{writes} - return 0; }} """ -_MAIN_COROUTINES = r""" -{globals} +_RUN_COROUTINES = r""" static void cunumpy_entry() {{ {call}; cunumpy_threads[cunumpy_current].done = 1; }} -int main() {{ -{inits} - cunumpy_dynamic_shared = calloc({shared_mem} + 16, 1); - gridDim = {{{gx}u, {gy}u, {gz}u}}; - blockDim = {{{bx}u, {by}u, {bz}u}}; +static void cunumpy_run() {{ const unsigned int n = blockDim.x * blockDim.y * blockDim.z; cunumpy_threads = (cunumpy_thread*)calloc(n, sizeof(cunumpy_thread)); for (unsigned int t = 0; t < n; ++t) cunumpy_threads[t].stack = (char*)malloc(CUNUMPY_STACK_BYTES); + cunumpy_trap_hook = cunumpy_coroutine_trap; for (unsigned int bz = 0; bz < gridDim.z; ++bz) for (unsigned int by = 0; by < gridDim.y; ++by) for (unsigned int bx = 0; bx < gridDim.x; ++bx) {{ @@ -201,12 +236,37 @@ threadIdx = {{t % blockDim.x, (t / blockDim.x) % blockDim.y, t / (blockDim.x * blockDim.y)}}; swapcontext(&cunumpy_scheduler, &cunumpy_threads[t].ctx); + if (cunumpy_trapped) goto done; if (!cunumpy_threads[t].done) running = 1; }} }} }} -{writes} - return 0; +done: + for (unsigned int t = 0; t < n; ++t) free(cunumpy_threads[t].stack); + free(cunumpy_threads); + cunumpy_threads = 0; + cunumpy_trap_hook = 0; +}} +""" + +# the entry point of the library: the arguments are pointers to the values +# (scalars), to the data (pointers) or to {data, shape..., strides...} as 64-bit +# integers (array views) +_LAUNCH = r""" +{globals} +{run} +extern "C" int cunumpy_launch(void** cunumpy_args, const unsigned int* g, + const unsigned int* b, unsigned long long shared_mem) {{ +{inits} + cunumpy_trapped = 0; + cunumpy_dynamic_shared = calloc(shared_mem + 16, 1); + gridDim = {{g[0], g[1], g[2]}}; + blockDim = {{b[0], b[1], b[2]}}; + if (setjmp(cunumpy_trap_jump) == 0) cunumpy_run(); + free(cunumpy_dynamic_shared); + cunumpy_dynamic_shared = 0; + fflush(stdout); + return cunumpy_trapped; }} """ @@ -219,17 +279,11 @@ def emulation_compiler() -> str | None: return shutil.which("c++") -def _scalar_literal(param: CudaParameter, index: int, value: Any) -> str: - """A C++ literal of `value`, checked and cast like a scalar kernel argument.""" - cast = _scalar_checker(param, index)(value) - kind = np.dtype(param.dtype).kind - if kind == "b": - return "true" if cast else "false" - if kind in "iu": - return f"(({param.ctype}){int(cast)}{'ULL' if kind == 'u' else 'LL'})" - if kind == "f": - return f"(({param.ctype}){float(cast).hex()})" # exact - raise NotImplementedError(f"emulation does not support {param.ctype} scalars") +def _scalar_value(param: CudaParameter, index: int, value: Any) -> np.ndarray: + """`value` as a 0-d array of the C type, checked and cast like a scalar argument.""" + if np.dtype(param.dtype).kind not in "biuf": + raise NotImplementedError(f"emulation does not support {param.ctype} scalars") + return np.asarray(_scalar_checker(param, index)(value), dtype=param.dtype) def _code(kernel: CudaKernel) -> str: @@ -258,6 +312,459 @@ def _check_supported(kernel: CudaKernel) -> None: ) +def _host(value: Any) -> Any: + """The NumPy buffer of a fake CuPy array; anything else as it is.""" + if _fake_cupy.is_active(): + try: + return _fake_cupy.host_buffer(value) + except TypeError: + pass + return value + + +def _field_value(struct_name: str, value: Any, field: CudaParameter) -> Any: + """The value of the struct field `field` in the struct argument `value`.""" + try: + if isinstance(value, (CudaStructValue, Mapping)): + return value[field.name] + return getattr(value, field.name) + except (KeyError, AttributeError): + raise TypeError( + f"struct argument {struct_name!r} has no value for the field " + f"{field.name!r} ({type(value).__name__})", + ) from None + + +def _declaration(param: CudaParameter, name: str) -> str: + if param.view_ndim is not None: + return f"{param.ctype} {name}" + return f"{param.ctype}{'*' if param.pointer else ''} {name}" + + +# wrapper kernels of the kernels with struct parameters, see _struct_wrapper +_WRAPPERS: weakref.WeakKeyDictionary[CudaKernel, CudaKernel] = ( + weakref.WeakKeyDictionary() +) + + +def _struct_wrapper(kernel: CudaKernel) -> CudaKernel: + """A kernel that takes every struct field of `kernel` as a parameter. + + The wrapper kernel (the source of `kernel` plus a ``__global__`` function with + one parameter per field) rebuilds the structs and calls `kernel`, so that the + emulation only has to deal with arrays and scalars. Its arguments are + :func:`_struct_fields`. Made once per kernel. + """ + wrapper = _WRAPPERS.get(kernel) + if wrapper is not None: + return wrapper + params, builds, call = [], [], [] + for p in kernel.signature: + if p.struct is None: + params.append(_declaration(p, p.name)) + call.append(p.name) + continue + names = [f"{p.name}__{f.name}" for f in p.struct.fields] + params += [_declaration(f, n) for f, n in zip(p.struct.fields, names)] + builds.append(f" {p.struct.name} {p.name}_struct{{{', '.join(names)}}};") + call.append(f"{p.name}_struct") + name = f"emulated_{kernel.name}" + source = ( + kernel.source + + f'\nextern "C" __global__ void {name}({", ".join(params)}) {{\n' + + "\n".join(builds) + + f"\n {kernel.expression}({', '.join(call)});\n}}\n" + ) + wrapper = CudaKernel( + source, + name, + block_size=kernel.block_size, + options=[o for o in kernel.options if not o.startswith("-I")], + include_dirs=kernel.include_dirs, + source_dir=kernel.source_dir, + n_threads_from=kernel.n_threads_from, + ) + _WRAPPERS[kernel] = wrapper + return wrapper + + +def _struct_fields(kernel: CudaKernel, args: Sequence[Any]) -> list[Any]: + """The arguments of :func:`_struct_wrapper`: each struct replaced by its fields.""" + flat = [] + for p, value in zip(kernel.signature, args): + if p.struct is None: + flat.append(value) + else: + flat += [_field_value(p.name, value, f) for f in p.struct.fields] + return flat + + +def _emulation_source(kernel: CudaKernel) -> tuple[str, bool]: + """The CUDA stubs and the kernel source as C++, without the launch function. + + Returns + ------- + tuple[str, bool] + The source, and whether the threads of a block run as coroutines (the + kernel calls ``__syncthreads``). + """ + kernel_source = kernel.source.replace('extern "C"', "") + kernel_source = _EXTERN_SHARED.sub( + r"\1* \2 = (\1*)cunumpy_dynamic_shared;", + kernel_source, + ) + coroutines = "__syncthreads" in _code(kernel) + if coroutines: + # ucontext needs _XOPEN_SOURCE before any system header (and macOS then + # hides the rest of its C library, which GCC's needs) + source = ( + "#define _XOPEN_SOURCE 700\n" + "#ifdef __APPLE__\n#define _DARWIN_C_SOURCE\n#endif\n" + + _STUBS + + _COROUTINES + ) + else: + source = _STUBS + "inline void __syncthreads() {}\n" + return source + kernel_source, coroutines + + +def _element_type(param: CudaParameter) -> str: + """The C type of the elements of a pointer or array view parameter.""" + if param.dtype is None: + return "unsigned char" + if param.view_ndim is not None: + return re.match(r"C?Array\d+D<(.*)>", param.ctype).group(1) + return param.ctype + + +def _library_source(kernel: CudaKernel) -> str: + """The C++ source of the emulation library of `kernel` (see :data:`_LAUNCH`).""" + source, coroutines = _emulation_source(kernel) + globals_, inits, call_args = [], [], [] + for i, param in enumerate(kernel.signature): + name = f"cunumpy_arg{i}" + if param.view_ndim is not None: + element = _element_type(param) + n = param.view_ndim + shape = ", ".join(f"m[{1 + k}]" for k in range(n)) + members = f"{{{shape}}}" + if not param.contiguous: + strides = ", ".join(f"m[{1 + n + k}]" for k in range(n)) + members += f", {{{strides}}}" + globals_.append(f"static {param.ctype} {name};") + inits.append( + f" {{ const long long* m = (const long long*)cunumpy_args[{i}];\n" + f" {name} = {param.ctype}{{({element}*)(intptr_t)m[0], {members}}}; }}", + ) + elif param.pointer: + element = _element_type(param) + globals_.append(f"static {element}* {name};") + inits.append(f" {name} = ({element}*)cunumpy_args[{i}];") + else: + globals_.append(f"static {param.ctype} {name};") + inits.append(f" {name} = *(const {param.ctype}*)cunumpy_args[{i}];") + call_args.append(name) + run = _RUN_COROUTINES if coroutines else _RUN_SERIAL + return source + _LAUNCH.format( + globals="\n".join(globals_), + run=run.format(call=f"{kernel.expression}({', '.join(call_args)})"), + inits="\n".join(inits), + ) + + +def _compile_command( + kernel: CudaKernel, + compiler: str, + options: Sequence[str], +) -> list[str]: + """The compiler, include directories and options for the emulation of `kernel`.""" + include_dirs = [cuda_include_dir(), *map(str, kernel.include_dirs)] + if kernel.source_dir is not None: + include_dirs.insert(0, str(kernel.source_dir)) + defines = [o for o in kernel.options if o.startswith("-D")] + return [ + compiler, + "-std=c++17", + "-w", + *(f"-I{d}" for d in include_dirs), + *defines, + *options, + ] + + +def emulation_cache_dir() -> Path | None: + """Where the emulation libraries are cached on disk, or None for no disk cache. + + ``CUNUMPY_EMULATION_CACHE`` sets the directory (``0`` or empty: no disk + cache); the default is ``$XDG_CACHE_HOME/cunumpy/emulation``, else + ``~/.cache/cunumpy/emulation``. + """ + value = os.environ.get("CUNUMPY_EMULATION_CACHE") + if value is not None: + if value.strip().lower() in {"", "0", "false", "no", "off"}: + return None + return Path(value).expanduser() + base = os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache" + return Path(base) / "cunumpy" / "emulation" + + +class _Library: + """A loaded emulation library; launches of it are serialized (global state).""" + + def __init__(self, path: Path) -> None: + self._dll = ctypes.CDLL(str(path)) + self.launch = self._dll.cunumpy_launch + self.launch.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint), + ctypes.POINTER(ctypes.c_uint), + ctypes.c_ulonglong, + ] + self.launch.restype = ctypes.c_int + self.lock = threading.Lock() + + +# libraries loaded in this process, and one lock per key so that a kernel is +# built once even when compile_all() builds in threads +_LIBRARIES: dict[str, _Library] = {} +_BUILDING: dict[str, threading.Lock] = {} +_BUILDING_LOCK = threading.Lock() + + +def _build(command: list[str], source: str, out: Path, kernel: CudaKernel) -> None: + with tempfile.TemporaryDirectory(prefix="cunumpy-emulation-") as tmp: + cpp = Path(tmp) / "kernel.cpp" + cpp.write_text(source) + built = subprocess.run( + [*command, str(cpp), "-o", str(out)], + capture_output=True, + text=True, + check=False, + ) + if built.returncode: + raise RuntimeError( + f"kernel {kernel.name!r} does not compile for emulation:\n{built.stderr[:4000]}", + ) + + +def _library(kernel: CudaKernel, compiler: str, options: Sequence[str]) -> _Library: + """The emulation library of `kernel`: from this process, the disk cache, or built.""" + command = [*_compile_command(kernel, compiler, options), "-O1", "-shared", "-fPIC"] + source = _library_source(kernel) + digest = hashlib.sha256() + for part in (source, _code(kernel), *command): + digest.update(part.encode()) + digest.update(b"\0") + key = digest.hexdigest()[:32] + library = _LIBRARIES.get(key) + if library is not None: + return library + with _BUILDING_LOCK: + lock = _BUILDING.setdefault(key, threading.Lock()) + with lock: + library = _LIBRARIES.get(key) + if library is not None: + return library + cache = emulation_cache_dir() + if cache is not None: + try: + cache.mkdir(parents=True, exist_ok=True) + except OSError: + cache = None + if cache is not None and os.access(cache, os.W_OK): + path = cache / f"{key}.so" + if not path.exists(): + partial = cache / f"{key}.{os.getpid()}.{threading.get_ident()}.tmp" + try: + _build(command, source, partial, kernel) + os.replace(partial, path) + finally: + partial.unlink(missing_ok=True) + library = _Library(path) + else: + with tempfile.TemporaryDirectory(prefix="cunumpy-emulation-") as tmp: + path = Path(tmp) / f"{key}.so" + _build(command, source, path, kernel) + library = _Library(path) # loaded: the file may go + _LIBRARIES[key] = library + return library + + +# per kernel: whether it passed _check_supported, and its libraries by +# (compiler, options), so that a launch does not hash the source and headers +_SUPPORTED: weakref.WeakSet[CudaKernel] = weakref.WeakSet() +_PREPARED: weakref.WeakKeyDictionary[CudaKernel, dict[tuple, _Library]] = ( + weakref.WeakKeyDictionary() +) + + +def _check_supported_once(kernel: CudaKernel) -> None: + if kernel not in _SUPPORTED: + _check_supported(kernel) + _SUPPORTED.add(kernel) + + +def _prepared( + kernel: CudaKernel, + compiler: str, + options: Sequence[str], + *, + refresh: bool = False, +) -> _Library: + """The library of `kernel`, looked up once per kernel object (or again with `refresh`).""" + libraries = _PREPARED.setdefault(kernel, {}) + key = (compiler, tuple(options)) + if refresh or key not in libraries: + libraries[key] = _library(kernel, compiler, options) + return libraries[key] + + +def _resolve_compiler(compiler: str | None) -> str: + compiler = compiler or emulation_compiler() + if compiler is None: + raise RuntimeError("emulation needs a C++ compiler (set CXX or install c++)") + return compiler + + +# syntax checks that passed, for kernels without a parsed signature +_CHECKED: set[tuple[str, ...]] = set() + + +def _check_syntax(kernel: CudaKernel, compiler: str, options: Sequence[str]) -> None: + command = [*_compile_command(kernel, compiler, options), "-fsyntax-only"] + key = (_code(kernel), kernel.expression, *command) + if key in _CHECKED: + return + source, _ = _emulation_source(kernel) + # the address instantiates a kernel template + source += f"\nstatic void cunumpy_check() {{ (void)&{kernel.expression}; }}\n" + with tempfile.TemporaryDirectory(prefix="cunumpy-emulation-") as tmp: + cpp = Path(tmp) / "kernel.cpp" + cpp.write_text(source) + built = subprocess.run( + [*command, str(cpp)], + capture_output=True, + text=True, + check=False, + ) + if built.returncode: + raise RuntimeError( + f"kernel {kernel.name!r} does not compile for emulation:\n{built.stderr[:4000]}", + ) + _CHECKED.add(key) + + +def compile_for_emulation( + kernel: CudaKernel, + *, + compiler: str | None = None, + options: Sequence[str] = (), +) -> None: + """Build the emulation library of `kernel` now, without a launch. + + :func:`emulate_cuda_kernel` builds it on the first launch otherwise; this + surfaces compile errors early, like :meth:`CudaKernel.compile + ` does on a GPU (inside + :func:`emulated_launches`, ``compile()`` calls this). The library is cached + (see :func:`emulate_cuda_kernel`). A kernel without a parsed signature is + only checked for syntax and type errors (``-fsyntax-only``). + + Parameters + ---------- + kernel : CudaKernel + The kernel. + compiler : str | None + C++ compiler; by default :func:`emulation_compiler`. + options : Sequence[str] + Additional compiler options, as for :func:`emulate_cuda_kernel`. + + Raises + ------ + NotImplementedError + If the kernel uses what the emulation cannot run (warp intrinsics). + RuntimeError + If there is no C++ compiler, or the kernel does not compile. + """ + _check_supported(kernel) + compiler = _resolve_compiler(compiler) + if kernel.signature is None: + _check_syntax(kernel, compiler, options) + return + if any(p.struct is not None for p in kernel.signature): + kernel = _struct_wrapper(kernel) + _prepared(kernel, compiler, options, refresh=True) + + +def _pack( + params: Sequence[CudaParameter], + args: Sequence[Any], +) -> tuple[Any, list[Any], list[tuple[np.ndarray, np.ndarray]]]: + """The ``void*`` arguments of ``cunumpy_launch``. + + Returns the argument array, the objects it points into (to keep alive), and + the (array, copy) pairs whose copy must be written back after the launch. + """ + argv = (ctypes.c_void_p * max(len(params), 1))() + keep: list[Any] = [] + write_back: list[tuple[np.ndarray, np.ndarray]] = [] + for i, (param, value) in enumerate(zip(params, args)): + if not (param.pointer or param.view_ndim is not None): + scalar = _scalar_value(param, i, value) + keep.append(scalar) + argv[i] = scalar.ctypes.data + continue + if not isinstance(value, np.ndarray): + raise TypeError( + f"argument {i} ({param.name}) must be a NumPy array, got " + f"{type(value).__name__}", + ) + if param.dtype is not None and value.dtype != param.dtype: + raise TypeError( + f"argument {i} ({param.name}) must have dtype " + f"{np.dtype(param.dtype)}, got {value.dtype}", + ) + if param.view_ndim is not None and value.ndim != param.view_ndim: + raise TypeError( + f"argument {i} ({param.name}) must be a {param.view_ndim}D " + f"array, got {value.ndim}D", + ) + if param.contiguous and not value.flags.c_contiguous: + # as on the GPU: a copy would drop what the kernel writes + raise TypeError( + f"argument {i} ({param.name}) must be C-contiguous for {param.ctype}", + ) + # the kernel works on the array itself, unless it is not writeable or + # its layout cannot be passed (a pointer needs C order, a view strides + # in whole elements): then on a copy, written back afterwards + strided = param.view_ndim is not None and not param.contiguous + usable = ( + value.flags.writeable + and value.flags.aligned + and ( + value.flags.c_contiguous + or (strided and all(s % value.itemsize == 0 for s in value.strides)) + ) + ) + array = value if usable else np.ascontiguousarray(value) + if array is not value and value.flags.writeable: + write_back.append((value, array)) + keep.append(array) + if param.view_ndim is None: + argv[i] = array.ctypes.data + continue + meta = np.array( + [ + array.ctypes.data, + *array.shape, + *(s // array.itemsize for s in array.strides), + ], + dtype=np.int64, + ) + keep.append(meta) + argv[i] = meta.ctypes.data + return argv, keep, write_back + + def emulate_cuda_kernel( kernel: CudaKernel, *args: Any, @@ -270,13 +777,22 @@ def emulate_cuda_kernel( ) -> None: """Run `kernel` on the CPU, serially, as if it were launched with `args`. + The kernel is compiled once (for each `compiler` and `options`) into a + shared library that takes the arguments at run time, so later launches, + with any argument values and array sizes, reuse it. Libraries are cached in + the process and on disk (:func:`emulation_cache_dir`). The launch runs in + this process, on the arrays themselves. + Parameters ---------- kernel : CudaKernel The kernel; its signature must be parsed (the default). *args - The kernel arguments with NumPy arrays in place of CuPy arrays. Arrays - are updated in place with what the kernel wrote. + The kernel arguments with NumPy arrays in place of CuPy arrays (arrays of + the fake CuPy are used through their host buffer). Arrays are updated in + place with what the kernel wrote. A struct parameter takes a mapping of + field names to values, a :class:`~cunumpy.arguments.CudaStructValue`, or + an object with an attribute per field. n_threads, grid, block Launch shape, as for :meth:`CudaKernel.__call__`. compiler : str | None @@ -291,156 +807,151 @@ def emulate_cuda_kernel( Raises ------ NotImplementedError - If the kernel (or a header it includes) uses warp intrinsics, or the - kernel has struct parameters or complex scalars. + If the kernel (or a header it includes) uses warp intrinsics, or has + complex scalars. TypeError If an argument does not match its parameter (dtype, dimensions, a scalar that does not fit). RuntimeError - If there is no C++ compiler, or the kernel does not compile or crashes. + If there is no C++ compiler, the kernel does not compile, or it calls + ``__trap()`` (a failed bounds check, inline ``asm``). Other crashes + (e.g. a segmentation fault from an unchecked out-of-bounds index) end + the process, as the launch runs in it. """ if kernel.signature is None: raise TypeError( "emulation needs a parsed kernel signature (check_signature=True)", ) - _check_supported(kernel) + _check_supported_once(kernel) params = kernel.signature if len(args) != len(params): raise TypeError( f"kernel {kernel.name!r} takes {len(params)} arguments, got {len(args)}", ) - compiler = compiler or emulation_compiler() - if compiler is None: - raise RuntimeError("emulation needs a C++ compiler (set CXX or install c++)") + if any(p.struct is not None for p in params): + return emulate_cuda_kernel( + _struct_wrapper(kernel), + *_struct_fields(kernel, args), + n_threads=n_threads, + grid=grid, + block=block, + compiler=compiler, + options=options, + shared_mem=shared_mem, + ) + if shared_mem < 0: + raise ValueError(f"shared_mem must be non-negative, got {shared_mem}") + args = tuple(_host(a) for a in args) + compiler = _resolve_compiler(compiler) grid_shape, block_shape = kernel.launch_shape( n_threads, grid=grid, block=block, args=args ) grid_shape = tuple(grid_shape) + (1,) * (3 - len(grid_shape)) block_shape = tuple(block_shape) + (1,) * (3 - len(block_shape)) + argv, keep, write_back = _pack(params, args) + library = _prepared(kernel, compiler, options) + with library.lock: + trapped = library.launch( + argv, + (ctypes.c_uint * 3)(*grid_shape), + (ctypes.c_uint * 3)(*block_shape), + int(shared_mem), + ) + del keep + for value, copy in write_back: + value[...] = copy + if trapped: + raise RuntimeError( + f"kernel {kernel.name!r} crashed in emulation: it called __trap() " + "(e.g. a failed bounds check or inline asm, which is not emulated)", + ) - with tempfile.TemporaryDirectory(prefix="cunumpy-emulation-") as tmp: - tmp_path = Path(tmp) - globals_, inits, call_args, writes, arrays = [], [], [], [], [] - for i, (param, value) in enumerate(zip(params, args)): - name = f"cunumpy_arg{i}" - if param.struct is not None: - raise NotImplementedError( - "emulation does not support struct parameters", - ) - if param.pointer or param.view_ndim is not None: - if not isinstance(value, np.ndarray): - raise TypeError( - f"argument {i} ({param.name}) must be a NumPy array, got " - f"{type(value).__name__}", - ) - if param.dtype is not None and value.dtype != param.dtype: - raise TypeError( - f"argument {i} ({param.name}) must have dtype " - f"{np.dtype(param.dtype)}, got {value.dtype}", - ) - if param.view_ndim is not None and value.ndim != param.view_ndim: - raise TypeError( - f"argument {i} ({param.name}) must be a {param.view_ndim}D " - f"array, got {value.ndim}D", - ) - if param.contiguous and not value.flags.c_contiguous: - # as on the GPU: a copy would drop what the kernel writes - raise TypeError( - f"argument {i} ({param.name}) must be C-contiguous for " - f"{param.ctype}", - ) - buffer = np.ascontiguousarray(value) - path = tmp_path / f"{name}.bin" - buffer.tofile(path) - ctype = param.ctype if param.dtype is not None else "unsigned char" - element = ( - re.match(r"C?Array\d+D<(.*)>", ctype).group(1) - if param.view_ndim is not None - else ctype - ) - globals_.append(f"static {element}* {name};") - inits.append( - f" {name} = ({element}*)malloc({max(buffer.nbytes, 1)});\n" - f' cunumpy_read("{path}", {name}, {buffer.nbytes});', - ) - if param.view_ndim is not None: - shape = ", ".join(f"{n}LL" for n in buffer.shape) - strides = ", ".join( - f"{s // buffer.itemsize}LL" for s in buffer.strides - ) - members = f"{{{shape}}}" - if not param.contiguous: - members += f", {{{strides}}}" - globals_.append(f"static {ctype} {name}_view;") - inits.append(f" {name}_view = {ctype}{{{name}, {members}}};") - call_args.append(f"{name}_view") - else: - call_args.append(name) - out = tmp_path / f"{name}.out" - writes.append(f' cunumpy_write("{out}", {name}, {buffer.nbytes});') - arrays.append((value, buffer, out)) + +@contextmanager +def emulated_launches( + *, + compiler: str | None = None, + options: Sequence[str] = (), +) -> Generator[None, None, None]: + """Run every :class:`~cunumpy.kernels.CudaKernel` launch in the block on the CPU. + + With the fake CuPy (:func:`cunumpy.kernel_testing.install_fake_cupy`) + CUDA kernels cannot run. Inside this block a launch is emulated instead + (:func:`emulate_cuda_kernel`), on the host buffers of the fake CuPy arrays + it is given, so code that launches kernels runs on a machine without a GPU + and its device arrays hold the results afterwards:: + + with emulated_launches(): + propagator(dt) # CuPy backend: the CUDA kernels run on the CPU + + The launch arguments are those of :meth:`CudaKernel.__call__` + (`stream` is ignored). Argument objects are flattened like in a launch; + struct values and :class:`~cunumpy.arguments.CudaStructArguments` objects + stay one argument and are read through their fields. Arrays must be arrays + of the fake CuPy or NumPy arrays. + + No CUDA compilation happens in the block either: :meth:`CudaKernel.compile` + (and so ``recompile()``, ``CudaKernelVariants.compile_all()`` and + ``KernelCatalog.compile_all()``) builds the emulation library instead + (:func:`compile_for_emulation`), so compile errors still show up before + the first launch, and returns None instead of a ``cupy.RawKernel``. Kernels + the emulation cannot run (warp intrinsics) are skipped by ``compile()``; + launching them raises. This holds on a real GPU too: the block means "no + CUDA here". + + Launches are serial, so use it on small problems. Each kernel is compiled + once (see :func:`emulate_cuda_kernel`). The limits of + :func:`emulate_cuda_kernel` apply. + + Parameters + ---------- + compiler : str | None + C++ compiler; by default :func:`emulation_compiler`. + options : Sequence[str] + Additional compiler options for every launch, e.g. + ``("-ffp-contract=off",)``. + """ + original_call = CudaKernel.__call__ + original_compile = CudaKernel.compile + + def launch( + self: CudaKernel, + *args: Any, + n_threads: int | Sequence[int] | None = None, + grid: int | Sequence[int] | None = None, + block: int | Sequence[int] | None = None, + shared_mem: int = 0, + stream: Any = None, + ) -> None: + flat: list[Any] = [] + for arg in args: + if isinstance(arg, (CudaStructValue, CudaStructArguments)): + flat.append(arg) + elif hasattr(arg, "__cuda_args__"): + flat.extend(arg.__cuda_args__()) else: - call_args.append(_scalar_literal(param, i, value)) - - if shared_mem < 0: - raise ValueError(f"shared_mem must be non-negative, got {shared_mem}") - kernel_source = kernel.source.replace('extern "C"', "") - kernel_source = _EXTERN_SHARED.sub( - r"\1* \2 = (\1*)cunumpy_dynamic_shared;", - kernel_source, + flat.append(arg) + emulate_cuda_kernel( + self, + *flat, + n_threads=n_threads, + grid=grid, + block=block, + compiler=compiler, + options=options, + shared_mem=shared_mem, ) - coroutines = "__syncthreads" in _code(kernel) - if coroutines: - # ucontext needs _XOPEN_SOURCE before any system header - source = "#define _XOPEN_SOURCE 700\n" + _STUBS + _COROUTINES - main = _MAIN_COROUTINES - else: - source = _STUBS + "inline void __syncthreads() {}\n" - main = _MAIN_SERIAL - source += kernel_source - source += main.format( - globals="\n".join(globals_), - inits="\n".join(inits), - shared_mem=int(shared_mem), - gx=grid_shape[0], - gy=grid_shape[1], - gz=grid_shape[2], - bx=block_shape[0], - by=block_shape[1], - bz=block_shape[2], - call=f"{kernel.expression}({', '.join(call_args)})", - writes="\n".join(writes), - ) - cpp = tmp_path / "kernel.cpp" - cpp.write_text(source) - exe = tmp_path / "kernel" - include_dirs = [cuda_include_dir(), *map(str, kernel.include_dirs)] - if kernel.source_dir is not None: - include_dirs.insert(0, str(kernel.source_dir)) - defines = [o for o in kernel.options if o.startswith("-D")] - command = [ - compiler, - "-std=c++17", - "-O1", - "-w", - *(f"-I{d}" for d in include_dirs), - *defines, - *options, - str(cpp), - "-o", - str(exe), - ] - built = subprocess.run(command, capture_output=True, text=True, check=False) - if built.returncode: - raise RuntimeError( - f"kernel {kernel.name!r} does not compile for emulation:\n{built.stderr[:4000]}", - ) - ran = subprocess.run([str(exe)], capture_output=True, text=True, check=False) - if ran.returncode: - raise RuntimeError( - f"kernel {kernel.name!r} crashed in emulation (exit {ran.returncode}):\n" - f"{ran.stdout[-2000:]}{ran.stderr[-2000:]}", - ) - for value, buffer, out in arrays: - result = np.fromfile(out, dtype=buffer.dtype).reshape(buffer.shape) - value[...] = result + + def compile(self: CudaKernel, *, log_stream: Any = None) -> None: + try: + compile_for_emulation(self, compiler=compiler, options=options) + except NotImplementedError: + pass # GPU only; a launch raises + + CudaKernel.__call__ = launch # type: ignore[method-assign] + CudaKernel.compile = compile # type: ignore[method-assign] + try: + yield + finally: + CudaKernel.__call__ = original_call # type: ignore[method-assign] + CudaKernel.compile = original_compile # type: ignore[method-assign] diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index 6f84e4a..47d3e09 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -17,7 +17,9 @@ :class:`~cunumpy.kernels.CudaKernel` work; * CUDA kernels cannot run: ``RawKernel`` and friends raise ``NotImplementedError`` when called, and :func:`cunumpy.kernel_testing.requires_cupy` - skips tests while the fake is active. + skips tests while the fake is active. Inside + :func:`cunumpy.kernel_testing.emulated_launches` a ``CudaKernel`` launch runs + on the CPU instead, on the host buffers of the arrays (:func:`host_buffer`). Activate it before CuPy or cunumpy's backend is first used, either with the environment variable ``CUNUMPY_FAKE_CUPY=1`` (read when cunumpy is imported) @@ -35,7 +37,7 @@ import types from pathlib import Path -__all__ = ["install", "is_active", "uninstall"] +__all__ = ["host_buffer", "install", "is_active", "uninstall"] _IMPLEMENTATION = Path(__file__).with_name("_fake_cupy_impl.py") _SUBMODULES = ("cupy.cuda", "cupy.cuda.device", "cupy.cuda.runtime", "cupy.linalg") @@ -47,6 +49,28 @@ def is_active() -> bool: return bool(getattr(module, "__cunumpy_fake__", False)) +def host_buffer(array: object) -> object: + """The NumPy array that holds the data of a fake ``cupy.ndarray``. + + Not a copy: writing to the buffer changes the fake device array, so a + kernel emulation (:func:`cunumpy.kernel_testing.emulated_launches`) can + update the array in place. A test can also read results with it without a + ``.get()``, which is the point of a "device" array being visible in tests. + + Raises + ------ + TypeError + If `array` is not an array of the fake CuPy (including when the fake + is not active). + """ + module = sys.modules.get("cupy") + if is_active() and isinstance(array, module.ndarray): + return array._a + raise TypeError( + f"host_buffer needs an array of the fake CuPy, got {type(array).__name__}", + ) + + def install() -> types.ModuleType: """Install the fake ``cupy`` package into ``sys.modules`` and return it. diff --git a/src/cunumpy/_fake_cupy_impl.py b/src/cunumpy/_fake_cupy_impl.py index 0a4a13a..7f6fc00 100644 --- a/src/cunumpy/_fake_cupy_impl.py +++ b/src/cunumpy/_fake_cupy_impl.py @@ -15,6 +15,14 @@ _HOST = _np.ndarray +def _sync(what): + """Count a scalar read of a device array (an implicit sync of the real CuPy).""" + from cunumpy._transfers import _ACTIVE, _record_sync + + if _ACTIVE: + _record_sync(what, implicit=True) + + def _err(obj, where=""): return TypeError( f"Unsupported type {type(obj)}{where} (fake CuPy: host arrays/lists are " @@ -75,6 +83,8 @@ def __getattr__(self, name): # no __array_interface__ etc.: NumPy must not see the host buffer if name.startswith("__"): raise AttributeError(name) + if name in ("item", "tolist"): + _sync(f"ndarray.{name}()") attr = getattr(self._a, name) if callable(attr): return _wrap_callable(attr, strict=True, name=f"ndarray.{name}") @@ -101,18 +111,23 @@ def __format__(self, spec): return format(self._a.item() if self._a.ndim == 0 else self._a, spec) def __bool__(self): + _sync("bool(device array)") return bool(self._a) def __int__(self): + _sync("int(device array)") return int(self._a) def __float__(self): + _sync("float(device array)") return float(self._a) def __complex__(self): + _sync("complex(device array)") return complex(self._a) def __index__(self): + _sync("index of a device array") return operator.index(self._a.item() if self._a.ndim == 0 else self._a) __hash__ = None diff --git a/src/cunumpy/_host.py b/src/cunumpy/_host.py new file mode 100644 index 0000000..d208273 --- /dev/null +++ b/src/cunumpy/_host.py @@ -0,0 +1,98 @@ +"""Run host-only code (SciPy, file readers, external libraries) with arguments of any backend.""" + +import functools +from collections.abc import Callable +from typing import Any + +import numpy as np + +from cunumpy import xp + + +def _on_device(arg: Any) -> bool: + """Whether ``arg`` is (or contains, for tuples and lists) a device array.""" + if isinstance(arg, (tuple, list)): + return any(_on_device(a) for a in arg) + return xp.is_gpu(arg) + + +def _to_host(arg: Any) -> Any: + """Copy device arrays in ``arg`` (an array, or a tuple/list of them) to the host.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_host(a) for a in arg) + return xp.to_numpy(arg) if xp.is_gpu(arg) else arg + + +def _to_device(arg: Any) -> Any: + """Copy NumPy arrays in ``arg`` (an array, or a tuple/list of them) to the device.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_device(a) for a in arg) + return xp.to_cupy(arg) if isinstance(arg, np.ndarray) else arg + + +def host_call(fun: Callable, *args: Any, **kwargs: Any) -> Any: + """Call a host-only function with arguments of any backend. + + Device (CuPy) arrays among the arguments are copied to the host, `fun` runs + on the NumPy backend, and its array results are copied back to the device, + once per call. Without device arguments `fun` is called directly (on the + NumPy backend this is a plain call), so the result lives where the + arguments live. The copies are counted by `xp.profiling.count_transfers()`. + + Parameters + ---------- + fun + The host-only function (a SciPy spline, an external code, ...). + *args, **kwargs + Arguments of `fun`: arrays (or tuples/lists of arrays) of either + backend, or scalars. + + Returns + ------- + The result of `fun` (an array, a tuple/list of arrays or a scalar), with + arrays on the device if any argument was on the device. + + Examples + -------- + >>> values = xp.host_call(spline, x) # x on the device -> values on the device + """ + device = _on_device(args) or _on_device(list(kwargs.values())) + if xp.get_backend() == "numpy" and not device: + return fun(*args, **kwargs) + with xp.use_backend("numpy"): + out = fun(*_to_host(args), **{k: _to_host(v) for k, v in kwargs.items()}) + return _to_device(out) if device else out + + +def evaluate_on_host(method: Callable) -> Callable: + """Decorator for methods that can only be evaluated on the host (see `host_call`). + + Examples + -------- + >>> class Equilibrium: + ... @xp.evaluate_on_host + ... def pressure(self, x): + ... return scipy_spline(x) + """ + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + return host_call(method, self, *args, **kwargs) + + return wrapper + + +def setup_on_host(init: Callable) -> Callable: + """Decorator for an ``__init__`` whose setup is host-only. + + `init` runs on the NumPy backend, so the object holds only host data (NumPy + arrays, SciPy splines, floats) on either backend. Host-only parts of its + evaluation can then go through `host_call` or `evaluate_on_host`. + """ + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + with xp.use_backend("numpy"): + init(self, *args, **kwargs) + + return wrapper diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index a87c2de..db2ed74 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -17,6 +17,7 @@ import copy import importlib +import inspect import os import warnings from collections.abc import Callable, Generator, Iterator, Mapping, Sequence @@ -73,11 +74,21 @@ class PyccelKernel: interpolate = PyccelKernel(some_interpolation_kernel, outputs=(5,)) interpolate(x, y, z, basis, coeffs, out) # `out` is argument 5 - Pyccel-compiled kernels are builtins with no introspectable signature, - so an index and a name are *not* interchangeable: declare the form you - actually call with. An empty sequence declares that the kernel writes to - none of its arguments. By default (``None``) every converted array is - copied back, which is always correct but does more work. + A name also finds the argument when it is passed positionally, and an + index also finds a keyword argument, if the parameter names of the + kernel are known: from its Python signature, or from `parameters` + (Pyccel-compiled kernels are builtins with no introspectable + signature; a :class:`~cunumpy.kernels.Kernel` supplies the names of + its host function). Without parameter names an index and a name are + *not* interchangeable: declare the form you actually call with. An + empty sequence declares that the kernel writes to none of its + arguments. By default (``None``) every converted array is copied back, + which is always correct but does more work. + parameters : sequence of str or callable, optional + The names of the positional parameters of `kernel`, or a function that + returns them (or None if they are unknown), for resolving `outputs` + names and indices. By default they are read from the signature of + `kernel`, where it has one. Examples -------- @@ -93,8 +104,10 @@ def __init__( object_modules: Sequence[str] = (), is_array: Callable[[Any], bool] | None = None, outputs: Sequence[int | str] | None = None, + parameters: Sequence[str] | Callable[[], Sequence[str] | None] | None = None, ) -> None: self._kernel = kernel + self._parameters = parameters self._use_cupy = use_cupy self._object_modules = tuple(object_modules) self._is_array = is_array or (lambda value: isinstance(value, np.ndarray)) @@ -122,6 +135,23 @@ def __repr__(self) -> str: f"outputs={self._outputs!r})" ) + def _parameter_names(self) -> Sequence[str] | None: + """The names of the positional parameters, or None if they are unknown.""" + source = self._parameters + if source is not None: + return source() if callable(source) else tuple(source) + function = getattr(self._kernel, "python", self._kernel) + try: + signature = inspect.signature(function) + except (TypeError, ValueError): + return None + positional = ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ) + names = [p.name for p in signature.parameters.values() if p.kind in positional] + return names or None + def _convert_to_numpy( self, value: Any, @@ -261,31 +291,61 @@ def _output_host_arrays( IndexError, KeyError If a declared output does not correspond to an argument of this call -- typically because an argument declared by index was passed - as a keyword, or vice versa. + as a keyword, or vice versa, and the parameter names are unknown. """ found: set[int] = set() seen: set[int] = set() + names = self._parameter_names() if self._outputs else None for entry in self._outputs or (): if isinstance(entry, int): index = entry + len(args_np) if entry < 0 else entry - if not 0 <= index < len(args_np): - raise IndexError( - f"{self.name}() was declared with output argument " - f"{entry}, but was called with {len(args_np)} " - "positional argument(s). Note that an output passed as " - "a keyword must be declared by name, not by index.", - ) - self._collect_host_arrays(args_np[index], found, seen) - else: - if entry not in kwargs_np: - raise KeyError( - f"{self.name}() was declared with output argument " - f"{entry!r}, but no such keyword argument was passed. " - "Note that an output passed positionally must be " - "declared by index, not by name.", - ) + if 0 <= index < len(args_np): + self._collect_host_arrays(args_np[index], found, seen) + continue + if ( + names is not None + and 0 <= entry < len(names) + and names[entry] in kwargs_np + ): + self._collect_host_arrays(kwargs_np[names[entry]], found, seen) + continue + raise IndexError( + f"{self.name}() was declared with output argument " + f"{entry}, but was called with {len(args_np)} " + "positional argument(s)" + + ( + "" + if names is not None + else ". Note that an output passed as a keyword must " + "be declared by name, not by index (the parameter " + "names of the kernel are unknown)." + ), + ) + if entry in kwargs_np: self._collect_host_arrays(kwargs_np[entry], found, seen) + continue + if ( + names is not None + and entry in names + and names.index(entry) + < len( + args_np, + ) + ): + self._collect_host_arrays(args_np[names.index(entry)], found, seen) + continue + raise KeyError( + f"{self.name}() was declared with output argument " + f"{entry!r}, but " + + ( + "no argument of that name was passed" + if names is not None + else "no such keyword argument was passed. Note that an " + "output passed positionally must be declared by index, not " + "by name (the parameter names of the kernel are unknown)." + ), + ) return found @@ -729,7 +789,52 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: return kernel(*args, **kwargs) -def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: +def _c_ordered_with_gaps(array: np.ndarray) -> bool: + """Whether `array` is laid out in C order, possibly with gaps between entries. + + Every stride is positive and at least the extent of the next axis, e.g. the + first ``n`` columns ``a[:, :n]`` of a C-contiguous ``a`` or every other + entry ``a[::2]``. Axes of length 1 are ignored, as their stride is never used. + """ + extent = array.itemsize + for length, stride in reversed(list(zip(array.shape, array.strides, strict=True))): + if length == 1: + continue + if stride < extent: + return False + extent = stride * length + return True + + +def _with_unit_axis_strides(array: np.ndarray) -> np.ndarray: + """`array`, or a view of it, whose axes of length 1 have strides in C order. + + NumPy never steps along an axis of length 1, so it leaves any stride there: + ``a[:, order]`` of a ``(1, n)`` array has strides ``(8, 8)`` and still counts + as C-contiguous. Pyccel's wrappers do use that stride to check the array + against its shape, and abort the process when it is smaller than the + extent of the next axes, so give such an axis the stride a fresh C-ordered + array would have (a view, no copy). + """ + if array.size == 0: + return array + strides = list(array.strides) + extent = array.itemsize + for axis in reversed(range(array.ndim)): + if array.shape[axis] == 1: + strides[axis] = max(strides[axis], extent) + else: + extent = strides[axis] * array.shape[axis] + if tuple(strides) == array.strides: + return array + return np.lib.stride_tricks.as_strided( + array, strides=strides, writeable=array.flags.writeable + ) + + +def as_kernel_array( + value: Any, like: Any, dtype: Any = None, *, strided: bool = False +) -> Any: """`value` as an array the kernel chosen for `like` takes. On the device of `like` (a CuPy array if `like` is one, a NumPy array @@ -745,6 +850,15 @@ def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: For an array the kernel writes, use :func:`kernel_output`, which copies a converted array back. + + With ``strided=True`` a NumPy array for a host kernel is also taken + unchanged when it is in C order with gaps: positive strides, each at least + the extent of the next axis, such as the first ``n`` columns + ``storage[:, :n]`` of a component-major marker buffer, which C-contiguity + would copy on every call. Pyccel's wrappers take such arrays without a copy + (they refuse F order and abort on negative strides, which are still copied); + a host kernel that needs contiguous memory, e.g. one using raw pointers, + must not use it. CuPy arrays are made C-contiguous either way. """ if is_gpu(like): import cupy @@ -759,11 +873,38 @@ def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: nbytes=_nbytes(result), ) return result - return np.ascontiguousarray(to_numpy(value), dtype=dtype) + value = to_numpy(value) + if ( + strided + and value.ndim > 0 + and (dtype is None or value.dtype == np.dtype(dtype)) + and _c_ordered_with_gaps(value) + ): + return _with_unit_axis_strides(value) + return _with_unit_axis_strides(np.ascontiguousarray(value, dtype=dtype)) + + +def _same_memory(buffer: Any, out: Any) -> bool: + """Whether `buffer` is `out` or a view of exactly its entries (no copy back needed).""" + if buffer is out: + return True + return ( + isinstance(buffer, np.ndarray) + and isinstance(out, np.ndarray) + and buffer.shape == out.shape + and buffer.dtype == out.dtype + and buffer.ctypes.data == out.ctypes.data + and all( + length == 1 or a == b + for length, a, b in zip(out.shape, buffer.strides, out.strides, strict=True) + ) + ) @contextmanager -def kernel_output(out: Any, like: Any, dtype: Any = None) -> Generator[Any]: +def kernel_output( + out: Any, like: Any, dtype: Any = None, *, strided: bool = False +) -> Generator[Any]: """The buffer a kernel chosen for `like` writes, copied back into `out` after it. Yields `out` itself if :func:`as_kernel_array` takes it unchanged (then the @@ -773,13 +914,64 @@ def kernel_output(out: Any, like: Any, dtype: Any = None) -> Generator[Any]: with xp.kernels.kernel_output(result, like=grid, dtype=float) as buffer: gather(convert(positions), grid, buffer, ...) + + `strided` is passed on to :func:`as_kernel_array`: with True, a host kernel + writes directly into a NumPy `out` in C order with gaps. """ - buffer = as_kernel_array(out, like, dtype) + buffer = as_kernel_array(out, like, dtype, strided=strided) yield buffer - if buffer is not out: + if not _same_memory(buffer, out): if is_gpu(out) and not is_gpu(buffer): out[...] = to_cupy(buffer) elif not is_gpu(out) and is_gpu(buffer): out[...] = to_numpy(buffer) else: out[...] = buffer + + +_SCALAR_ANNOTATIONS = {"int", "float", "bool", "complex", "str"} + + +def outputs_from_annotations(function: Callable[..., Any]) -> tuple[str, ...] | None: + """The names of the parameters a kernel may write to, read from its annotations. + + Pyccel kernels mark what they only read with ``Final``: ``x: "Final[float[:]]"``. + A parameter is an output unless it is annotated ``Final`` (or ``const``) or + has the annotation of a scalar (``int``, ``float``, ``bool``, ``complex``, + ``str``); an annotated array or an argument object may be written to. Use the + result as the `outputs` of a :class:`PyccelKernel` so that only those arrays + are copied back to the device:: + + PyccelKernel(push, outputs=outputs_from_annotations(push)) + + Returns + ------- + tuple[str, ...] | None + The parameter names, in order, or None if `function` has no signature or + none of its parameters is annotated (so nothing can be said). + """ + try: + parameters = inspect.signature(function).parameters.values() + except (TypeError, ValueError): + return None + if all(p.annotation is inspect.Parameter.empty for p in parameters): + return None + outputs = [] + for p in parameters: + if isinstance(p.annotation, str): + text = p.annotation + elif isinstance(p.annotation, type): + text = p.annotation.__name__ + else: + text = str(p.annotation) + text = text.replace("typing.", "").strip().strip("'\"") + if p.annotation is inspect.Parameter.empty: + outputs.append(p.name) # nothing is known: assume it may be written + elif ( + text.startswith(("Final[", "const ", "Final ")) + or text in _SCALAR_ANNOTATIONS + ): + continue + else: + outputs.append(p.name) + return tuple(outputs) diff --git a/src/cunumpy/_metal_kernel.py b/src/cunumpy/_metal_kernel.py new file mode 100644 index 0000000..cbf45d6 --- /dev/null +++ b/src/cunumpy/_metal_kernel.py @@ -0,0 +1,258 @@ +"""Metal kernels for the GPU of Apple silicon Macs, run through MLX.""" + +from __future__ import annotations + +import importlib +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _nbytes, _record + +_FLOAT64 = np.dtype(np.float64) +_UNSUPPORTED = (np.dtype(np.complex128),) + + +def _mlx() -> Any: + """Import ``mlx.core``, with an actionable error if it is not usable.""" + try: + mx = importlib.import_module("mlx.core") + except ImportError as error: + raise ImportError( + "MetalKernel needs MLX on an Apple silicon Mac: pip install 'cunumpy[metal]'", + ) from error + if not mx.metal.is_available(): + raise RuntimeError("MetalKernel needs a Metal GPU, but MLX reports none") + return mx + + +def metal_available() -> bool: + """Whether `MetalKernel` can run here (MLX installed and a Metal GPU present).""" + try: + _mlx() + except (ImportError, RuntimeError): + return False + return True + + +class MetalKernel: + """A Metal Shading Language kernel for the GPU of Apple silicon, run with MLX. + + The arrays are NumPy arrays: they are copied to MLX arrays for the launch + and the results are written back into the output arrays you pass, like a + :class:`~cunumpy.kernels.PyccelKernel` writes into its output arguments. No + backend switch is needed. The copies are counted by + :func:`~cunumpy.profiling.count_transfers`. + + Parameters + ---------- + source : str + Body of the kernel function (MSL). MLX generates the signature from + `inputs` and `outputs`: every name is a pointer to the flat, row-major + data of that array (``const device T*`` for inputs, ``device T*`` for + outputs), and ``thread_position_in_grid`` and the other Metal + attributes used in the body are added to the signature automatically. + Names from `template` are available as compile-time constants. + inputs : Sequence[str] + Names of the input arrays, in the order they are passed to the call. + outputs : Sequence[str] + Names of the output arrays, in the order they are passed as `out`. + name : str + Name of the kernel; part of the compiled function name. + header : str + Source placed before the kernel function: includes, ``#define``\\ s and + helper functions. + threadgroup : int | Sequence[int] + Threads per threadgroup: an integer, or 1 to 3 integers. + float64 : {"error", "cast"} + The GPU has no float64. With "error" (the default) a float64 input or + output raises ``TypeError``. With "cast" float64 inputs are computed + in float32 and float64 outputs are filled from float32 results, which + is what to pick when float32 precision is enough. + atomic_outputs : bool + Declare the outputs as ``device atomic*`` for atomic updates. + init_value : float | None + Value the outputs are filled with before the launch. Outputs are + otherwise uninitialized: write every element, or pass the old array as + an input as well to read its values. + + Notes + ----- + The kernel is compiled by MLX at its first call and cached. All inputs + and outputs are made C-contiguous, so ``a[i]`` in the source is the flat + index of the NumPy array. + + Examples + -------- + >>> scale = xp.kernels.MetalKernel( + ... "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];", + ... inputs=["x", "a"], + ... outputs=["y"], + ... ) + >>> y = np.empty(n, dtype=np.float32) + >>> scale(x, np.float32([2.0]), out=y) + """ + + def __init__( + self, + source: str, + inputs: Sequence[str], + outputs: Sequence[str], + *, + name: str = "cunumpy_kernel", + header: str = "", + threadgroup: int | Sequence[int] = 256, + float64: str = "error", + atomic_outputs: bool = False, + init_value: float | None = None, + ) -> None: + if float64 not in ("error", "cast"): + raise ValueError(f"float64 must be 'error' or 'cast', not {float64!r}") + if not outputs: + raise ValueError("a MetalKernel needs at least one output") + names = [*inputs, *outputs] + if len(set(names)) != len(names): + raise ValueError(f"input and output names must be distinct: {names}") + self.source = source + self.inputs = tuple(inputs) + self.outputs = tuple(outputs) + self.name = name + self.header = header + self.threadgroup = self._as_triple(threadgroup, "threadgroup", default=1) + self.float64 = float64 + self.atomic_outputs = atomic_outputs + self.init_value = init_value + self._kernel: Any = None + + def __repr__(self) -> str: + return ( + f"MetalKernel({self.name!r}, inputs={self.inputs}, outputs={self.outputs})" + ) + + @staticmethod + def _as_triple(value: int | Sequence[int], what: str, default: int) -> tuple: + values = (int(value),) if isinstance(value, (int, np.integer)) else tuple(value) + if not 1 <= len(values) <= 3 or any(v < 1 for v in values): + raise ValueError(f"{what} must be 1 to 3 positive integers, got {value!r}") + return values + (default,) * (3 - len(values)) + + @staticmethod + def _as_input(value: Any) -> np.ndarray: + """A NumPy array; Python scalars become float32 or int32 arrays of shape (1,).""" + if isinstance(value, (bool, int)): + return np.array([value], dtype=np.int32) + if isinstance(value, float): + return np.array([value], dtype=np.float32) + array = np.asarray(value) + return array.reshape(1) if array.ndim == 0 else array + + def _compiled(self, mx: Any) -> Any: + if self._kernel is None: + self._kernel = mx.fast.metal_kernel( + name=self.name, + input_names=list(self.inputs), + output_names=list(self.outputs), + source=self.source, + header=self.header, + ensure_row_contiguous=True, + atomic_outputs=self.atomic_outputs, + ) + return self._kernel + + def _device_dtype(self, dtype: np.dtype, role: str, name: str) -> np.dtype: + if dtype == _FLOAT64 and self.float64 == "cast": + return np.dtype(np.float32) + if dtype == _FLOAT64 or dtype in _UNSUPPORTED: + raise TypeError( + f"{role} {name!r} of MetalKernel {self.name!r} has dtype {dtype}, " + "which the Apple GPU does not support; use float32, or " + "float64='cast' to compute in float32", + ) + return dtype + + def __call__( + self, + *args: Any, + out: Any, + n_threads: int | Sequence[int] | None = None, + template: Mapping[str, Any] | None = None, + ) -> Any: + """Launch the kernel. + + Parameters + ---------- + *args : numpy.ndarray or scalar + The inputs, in the order of `inputs`. Python scalars become + 1-element float32 or int32 arrays; read them as ``a[0]`` in the source. + out : numpy.ndarray or Sequence[numpy.ndarray] + The output arrays, in the order of `outputs`; their shapes and + dtypes define the outputs and they are filled in place. + n_threads : int | Sequence[int] | None + Total number of threads (1 to 3 dimensions), not threadgroups. The + default is the first axis of the first output. + template : Mapping[str, int | bool | numpy.dtype] | None + Compile-time constants; each name is a constant in the source. A + different value compiles another variant. + + Returns + ------- + The output array, or a tuple of them if there are several. + """ + mx = _mlx() + if len(args) != len(self.inputs): + raise TypeError( + f"MetalKernel {self.name!r} takes {len(self.inputs)} input(s) " + f"{self.inputs}, got {len(args)}", + ) + outs = (out,) if isinstance(out, np.ndarray) else tuple(out or ()) + if len(outs) != len(self.outputs) or not all( + isinstance(o, np.ndarray) for o in outs + ): + raise TypeError( + f"out must be {len(self.outputs)} NumPy array(s) for {self.outputs}", + ) + + host_inputs = [self._as_input(a) for a in args] + for name, array in zip(self.inputs, host_inputs): + self._device_dtype(array.dtype, "input", name) + device_dtypes = [ + self._device_dtype(o.dtype, "output", name) + for name, o in zip(self.outputs, outs) + ] + if n_threads is None: + n_threads = outs[0].shape[0] if outs[0].ndim else 1 + grid = self._as_triple(n_threads, "n_threads", default=1) + + mlx_inputs = [] + for array in host_inputs: + if array.dtype == _FLOAT64: + array = array.astype(np.float32) + mlx_inputs.append(mx.array(np.ascontiguousarray(array))) + if _COUNTERS: + _record( + "to_device", + f"MetalKernel {self.name!r} input ({_describe(array)})", + nbytes=_nbytes(array), + ) + + results = self._compiled(mx)( + inputs=mlx_inputs, + template=list((template or {}).items()), + grid=grid, + threadgroup=self.threadgroup, + output_shapes=[o.shape for o in outs], + output_dtypes=[getattr(mx, dtype.name) for dtype in device_dtypes], + init_value=self.init_value, + ) + mx.eval(*results) + for result, host in zip(results, outs): + np.copyto(host, np.asarray(result), casting="same_kind") + if _COUNTERS: + _record( + "to_host", + f"MetalKernel {self.name!r} output ({_describe(host)})", + nbytes=_nbytes(host), + ) + return outs[0] if len(outs) == 1 else outs diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 8dfaf98..fd46415 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -13,7 +13,7 @@ from maybempi import get_mpi from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _describe, _nbytes, _record +from cunumpy._transfers import _describe, _nbytes, _record, _record_sync from cunumpy.xp import array_backend, cupy_available, to_numpy _logger = logging.getLogger(__name__) @@ -45,6 +45,8 @@ def synchronize_for_mpi(*arrays: Any, stream: Any = None, event: Any = None) -> if isinstance(event, HostEvent) or isinstance(stream, HostStream): raise TypeError("device buffers require a CUDA producer stream or event") + if _COUNTERS: + _record_sync("synchronize_for_mpi()") if event is not None: event.synchronize() return @@ -262,6 +264,8 @@ def mpi_buffer( # Ensure MPI's host buffer can be reused immediately on context exit. if transfer_array is not array: array[...] = transfer_array + if _COUNTERS: + _record_sync("mpi_buffer() staging for recv") cp.cuda.get_current_stream().synchronize() diff --git a/src/cunumpy/_staging.py b/src/cunumpy/_staging.py index 1fd8b7d..aa2c277 100644 --- a/src/cunumpy/_staging.py +++ b/src/cunumpy/_staging.py @@ -37,9 +37,9 @@ import numpy as np from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _nbytes, _record +from cunumpy._transfers import _describe, _nbytes, _record, _record_sync -__all__ = ["HostStaging", "StagedCopy"] +__all__ = ["HostCopy", "HostStaging", "StagedCopy", "to_host_async"] # substituted in tests that have no GPU _is_device_array = array_api_compat.is_cupy_array @@ -278,6 +278,7 @@ def copy(self, array: Any, *, stream: Any = None, event: Any = None) -> StagedCo "to_host", f"HostStaging.copy({self.shape} {self.dtype})", nbytes=_nbytes(array), + blocking=False, ) finally: # Retain completion even if a copy failed after enqueuing work. @@ -290,3 +291,93 @@ def synchronize(self) -> None: if slot.event is not None: with slot.context(): slot.event.synchronize() + + +class HostCopy: + """A device-to-host copy started by :func:`~cunumpy.to_host_async`.""" + + def __init__(self, host: np.ndarray, event: Any = None, source: Any = None) -> None: + self._host = host + self._event = event + self._source = source # the device array stays alive until the copy is done + + def ready(self) -> bool: + """Whether the copy has finished (never waits).""" + return self._event is None or bool(self._event.done) + + def result(self) -> Any: + """The value on the host; waits for the copy only if it has not finished. + + A NumPy scalar for a 0-d array, else a NumPy array (in page-locked memory + on the GPU). A wait is counted as a ``sync`` by + :func:`~cunumpy.profiling.count_transfers`. + """ + if self._event is not None: + if not self._event.done: + if _COUNTERS: + _record_sync("to_host_async(...).result()") + self._event.synchronize() + self._event = self._source = None + return self._host[()] if self._host.ndim == 0 else self._host + + +# one copy stream per device +_COPY_STREAMS: dict[int, Any] = {} + + +def to_host_async(array: Any) -> HostCopy: + """Start copying `array` (a device scalar or small array) to the host; do not wait. + + The copy runs on a separate stream into page-locked memory, after the work + queued so far on the current stream, so the host can queue more kernels + while it is in flight. :meth:`HostCopy.ready` tells whether it has + finished (never waits), :meth:`HostCopy.result` returns the value:: + + pending = xp.to_host_async(residual_norm) + ... # queue the next iteration + if pending.ready() and pending.result() < tol: + break + + The value is that of the array when the queued work is done: later kernels + may overwrite the array. For large arrays copied repeatedly, use + :class:`~cunumpy.memory.HostStaging`, which reuses its buffers. + + On the NumPy backend (a host array) the value is copied at once and + ``ready()`` is always True. On the fake CuPy too, and the copy is counted. + :func:`~cunumpy.profiling.count_transfers` records a ``to_host`` event with + ``blocking=False``. + """ + if not _is_device_array(array): + return HostCopy(np.array(array)) + from cunumpy import _fake_cupy + + nbytes = _nbytes(array) + if _fake_cupy.is_active(): + host = array.get() + event = None + source = None + else: + cp = _cupy() + with cp.cuda.Device(array.device.id): + source = cp.ascontiguousarray(array) # on the current stream + host = _empty_pinned(source.shape, source.dtype) + queued = cp.cuda.get_current_stream().record() + stream = _COPY_STREAMS.get(array.device.id) + if stream is None: + stream = cp.cuda.Stream(non_blocking=True) + _COPY_STREAMS[array.device.id] = stream + stream.wait_event(queued) + try: + # blocking=False (CuPy >= 13): return once the copy is enqueued + source.get(stream=stream, out=host, blocking=False) + except TypeError: # older CuPy + source.get(stream=stream, out=host) + event = stream.record() + if _COUNTERS: + _record( + "to_host", + f"to_host_async({_describe(array)})", + nbytes=nbytes, + blocking=False, + ) + return HostCopy(host, event, source) diff --git a/src/cunumpy/_transfers.py b/src/cunumpy/_transfers.py index f2d642e..3b6d6a5 100644 --- a/src/cunumpy/_transfers.py +++ b/src/cunumpy/_transfers.py @@ -26,7 +26,13 @@ device arrays to the host (and back), one event per call; * ``fallback``: a :class:`~cunumpy.kernels.Kernel` without CUDA kernel calling its host kernel on the CuPy backend (``missing_cuda="fallback"``), one event per call; -* ``device_copy``: device-only dtype/layout conversions in CuNumpy helpers. +* ``device_copy``: device-only dtype/layout conversions in CuNumpy helpers; +* ``sync``: the host waited for the device: :func:`~cunumpy.synchronize`, the + waits of the MPI helpers and of the CUDA debug mode, and, on the fake CuPy, + a scalar read of a device array (``float(a)``, ``int(a)``, ``bool(a)``, + ``a.item()``, ``a.tolist()``). Syncs are listed in :attr:`TransferCounter.syncs` + and the report but not in :attr:`~TransferCounter.total`, and + :func:`assert_no_transfers` accepts them unless called with ``syncs=True``. Mirror refreshes, argument conversions, staging, serial MPI and kernel output copy-back are also counted. Each physical host/device copy is recorded with its @@ -37,7 +43,9 @@ Limitations ----------- Only transfers made *through cunumpy* are seen. Raw ``cupy.ndarray.get()``, -``cupy.asarray(numpy_array)``, ``numpy.asarray(cupy_array)``, ``float(device_array)``, +``cupy.asarray(numpy_array)``, ``numpy.asarray(cupy_array)``, ``float(device_array)`` +(an implicit sync of the real CuPy, which cannot be observed from Python; the fake +CuPy reports it), forwarded backend operations such as ``xp.asarray`` and implicit conversions inside other libraries are not counted; use ``nsys`` (or CuPy's own profiling hooks) to find those. @@ -49,15 +57,17 @@ from __future__ import annotations +import functools import math import os import sys -from collections.abc import Generator +from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from typing import Any __all__ = [ + "TransferBudget", "TransferCounter", "TransferEvent", "assert_no_transfers", @@ -65,7 +75,14 @@ ] #: Event kinds, in the order they are reported. -KINDS = ("to_host", "to_device", "kernel_conversion", "fallback", "device_copy") +KINDS = ( + "to_host", + "to_device", + "kernel_conversion", + "fallback", + "device_copy", + "sync", +) # The currently active counters, innermost last. Instrumented code checks # ``if _ACTIVE:`` before doing any work, so the overhead of an inactive counter @@ -92,15 +109,31 @@ class TransferEvent: The call site outside cunumpy, as ``"file:line"``. nbytes : int | None Payload bytes of a physical copy. None for markers or unknown sizes. + blocking : bool + Whether the host waited for the copy. False for the copies started by + :func:`~cunumpy.to_host_async`. + implicit : bool + Whether the event was not asked for explicitly: a scalar read of a + device array (``float(a)``) on the fake CuPy, recorded as a ``sync``. """ kind: str description: str where: str nbytes: int | None = None + blocking: bool = True + implicit: bool = False + + @property + def label(self) -> str: + """The description, marked ``[async]`` or ``[implicit]`` where it applies.""" + marks = ("" if self.blocking else " [async]") + ( + " [implicit]" if self.implicit else "" + ) + return self.description + marks def __str__(self) -> str: - return f"{self.where}: {self.kind}: {self.description}" + return f"{self.where}: {self.kind}: {self.label}" class TransferCounter: @@ -151,10 +184,15 @@ def fallbacks(self) -> int: """Number of `Kernel` calls that fell back to the host kernel on CuPy.""" return self.count("fallback") + @property + def syncs(self) -> int: + """Number of times the host waited for the device (see the module documentation).""" + return self.count("sync") + @property def total(self) -> int: - """Number of observations, including conversion/fallback markers.""" - return len(self.events) + """Number of observations, including conversion/fallback markers (not syncs).""" + return sum(1 for event in self.events if event.kind != "sync") def bytes(self, kind: str) -> int: """Known bytes copied for `kind`; markers and unknown sizes add zero.""" @@ -186,7 +224,7 @@ def report(self) -> str: lines.append(f" {kind} ({len(events)}):") grouped: dict[tuple[str, str], int] = {} for event in events: - key = (event.where, event.description) + key = (event.where, event.label) grouped[key] = grouped.get(key, 0) + 1 for (where, description), n in grouped.items(): times = f" (x{n})" if n > 1 else "" @@ -205,13 +243,25 @@ def _caller() -> str: return "" -def _record(kind: str, description: str, *, nbytes: int | None = None) -> None: +def _record( + kind: str, + description: str, + *, + nbytes: int | None = None, + blocking: bool = True, + implicit: bool = False, +) -> None: """Record a transfer in every active counter (call only ``if _ACTIVE:``).""" - event = TransferEvent(kind, description, _caller(), nbytes) + event = TransferEvent(kind, description, _caller(), nbytes, blocking, implicit) for counter in _ACTIVE: counter._add(event) +def _record_sync(description: str, *, implicit: bool = False) -> None: + """Record that the host waits for the device (call only ``if _ACTIVE:``).""" + _record("sync", description, implicit=implicit) + + def _describe(array: Any) -> str: """``shape=..., dtype=...`` of an array, or the type name otherwise.""" shape = getattr(array, "shape", None) @@ -247,7 +297,9 @@ def _is_device_copy(source: Any, result: Any) -> bool: @contextmanager -def count_transfers() -> Generator[TransferCounter, None, None]: +def count_transfers( + into: TransferCounter | None = None, +) -> Generator[TransferCounter, None, None]: """Count the host/device transfers made through cunumpy in the block. Yields a :class:`TransferCounter` that records every ``to_numpy``, @@ -258,7 +310,11 @@ def count_transfers() -> Generator[TransferCounter, None, None]: NumPy array. Blocks can be nested; each active counter sees the transfers made inside - it. Transfers that bypass cunumpy (raw ``cupy.ndarray.get()``, + it. With `into`, the events are added to that counter (e.g. to accumulate + over several calls); a counter that is already active is not added again, + so nested blocks with the same counter count each event once. Like any + context manager made with :func:`contextlib.contextmanager`, it is also a + decorator: ``@count_transfers(counter)``. Transfers that bypass cunumpy (raw ``cupy.ndarray.get()``, ``cupy.asarray(numpy_array)``, conversions inside other libraries) are not seen; see the module documentation. @@ -268,7 +324,10 @@ def count_transfers() -> Generator[TransferCounter, None, None]: ... propagator(dt) >>> assert counter.total == 0, counter.report() """ - counter = TransferCounter() + counter = TransferCounter() if into is None else into + if any(active is counter for active in _ACTIVE): + yield counter + return _ACTIVE.append(counter) try: yield counter @@ -277,9 +336,15 @@ def count_transfers() -> Generator[TransferCounter, None, None]: @contextmanager -def assert_no_transfers() -> Generator[TransferCounter, None, None]: +def assert_no_transfers( + *, syncs: bool = False +) -> Generator[TransferCounter, None, None]: """Raise ``AssertionError`` if the block makes a transfer through cunumpy. + Syncs (the host waiting for the device) are accepted unless `syncs` is + True, e.g. ``assert_no_transfers(syncs=True)`` for a time step that must + never stall on the device. + A :func:`count_transfers` block that, on exit, raises with the counter's :meth:`~TransferCounter.report` if a host/device copy or host conversion/ fallback was counted. Device-only conversions are allowed. Only checked if @@ -292,8 +357,211 @@ def assert_no_transfers() -> Generator[TransferCounter, None, None]: """ with count_transfers() as counter: yield counter - if any(event.kind != "device_copy" for event in counter.events): + ignored = ("device_copy",) if syncs else ("device_copy", "sync") + if any(event.kind not in ignored for event in counter.events): raise AssertionError( "host/device transfers inside a block that must not transfer:\n" + counter.report(), ) + + +class _PhaseRouter(TransferCounter): + """The counter a :class:`TransferBudget` keeps active: it adds every event + to the innermost phase of the budget only.""" + + def __init__(self, budget: TransferBudget) -> None: + super().__init__() + self._budget = budget + + def _add(self, event: TransferEvent) -> None: + self._budget[self._budget._stack[-1]]._add(event) + + +_LIMITS = ("max_nbytes", "max_count", "max_total_bytes", "blocking", "implicit") + + +@dataclass +class _Rule: + allow: dict[str, dict[str, Any]] + ignore: tuple[str, ...] + calls: int | None + + +class TransferBudget: + """Transfers counted per phase of a program, checked against rules per phase. + + A time loop typically has phases with different budgets: the time step must + not copy arrays between host and device at all, the diagnostics may copy a + few scalars to the host, the output each saved array once. A budget counts + the transfers of each phase (with :func:`count_transfers`) and checks them:: + + budget = TransferBudget(started=False) + model.integrate = budget.count("integrate")(model.integrate) + ... # setup: not counted + budget.start() + for step in range(n_steps): + model.integrate(dt) + with budget.phase("output"): + save(model) + budget.require("integrate", allow={"to_host": dict(max_nbytes=8)}, calls=n_steps) + budget.require("output", allow={"to_host": dict(max_count=n, max_total_bytes=b)}) + budget.check() # AssertionError with report() if a rule is broken + + Phases nest: an event is counted in the innermost phase only, and a phase + entered again inside itself (recursion) counts each event and call once. + + Parameters + ---------- + started : bool + Whether to count from the start; otherwise phases run uncounted until + :meth:`start`. + + Attributes + ---------- + phases : dict[str, TransferCounter] + The events of each phase, accumulated over its calls. + calls : dict[str, int] + How many times each phase was entered while counting. + """ + + def __init__(self, *, started: bool = True) -> None: + self.phases: dict[str, TransferCounter] = {} + self.calls: dict[str, int] = {} + self.started = started + self._rules: dict[str, _Rule] = {} + self._stack: list[str] = [] + self._router = _PhaseRouter(self) + + def __getitem__(self, phase: str) -> TransferCounter: + """The counter of `phase` (empty if it has not run).""" + return self.phases.setdefault(phase, TransferCounter()) + + def start(self) -> None: + """Count the phases from now on.""" + self.started = True + + def stop(self) -> None: + """Stop counting; phases run uncounted until :meth:`start`.""" + self.started = False + + @contextmanager + def phase(self, name: str) -> Generator[TransferCounter, None, None]: + """Count the transfers of the block in phase `name` (accumulated).""" + counter = self[name] + if not self.started: + yield counter + return + if name not in self._stack: + self.calls[name] = self.calls.get(name, 0) + 1 + self._stack.append(name) + try: + with count_transfers(into=self._router): + yield counter + finally: + self._stack.pop() + + def count(self, name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """A decorator: every call of the function is counted in phase `name`.""" + + def decorate(function: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(function) + def counted(*args: Any, **kwargs: Any) -> Any: + with self.phase(name): + return function(*args, **kwargs) + + return counted + + return decorate + + def require( + self, + phase: str, + allow: Mapping[str, Mapping[str, Any] | None] | None = None, + *, + ignore: Sequence[str] = ("device_copy", "sync"), + calls: int | None = None, + ) -> None: + """Set the rule of `phase`, checked by :meth:`check`. + + Parameters + ---------- + phase : str + The phase. + allow : Mapping[str, Mapping | None] | None + The event kinds the phase may have, each with its limits (None or + ``{}``: any number). Every other kind is forbidden, except those in + `ignore`. Limits: ``max_nbytes`` (of each event; events of unknown + size break it), ``max_count`` and ``max_total_bytes`` (over the + whole phase), ``blocking`` and ``implicit`` (the events must have + this value of :attr:`TransferEvent.blocking` / ``implicit``), e.g. + ``{"to_host": dict(max_nbytes=8, blocking=False)}``. + ignore : Sequence[str] + Kinds that are not checked unless they are in `allow`: by default + device-only copies and syncs (the host waiting for the device). + calls : int | None + The number of times the phase must have been entered. + """ + allow = {kind: dict(limits or {}) for kind, limits in (allow or {}).items()} + for kind, limits in allow.items(): + if kind not in KINDS: + raise ValueError(f"unknown transfer kind {kind!r}; kinds: {KINDS}") + unknown = set(limits) - set(_LIMITS) + if unknown: + raise ValueError( + f"unknown limits {sorted(unknown)} for {kind!r}; limits: {_LIMITS}" + ) + self._rules[phase] = _Rule(allow, tuple(ignore), calls) + + def violations(self) -> list[str]: + """The broken rules, one line each, with the offending events.""" + found = [] + for name, rule in self._rules.items(): + events = self[name].events + n_calls = self.calls.get(name, 0) + if rule.calls is not None and n_calls != rule.calls: + found.append(f"{name}: {n_calls} call(s), expected {rule.calls}") + for event in events: + if event.kind not in rule.allow: + if event.kind not in rule.ignore: + found.append(f"{name}: not allowed: {event}") + continue + limits = rule.allow[event.kind] + limit = limits.get("max_nbytes") + if limit is not None and (event.nbytes is None or event.nbytes > limit): + found.append( + f"{name}: {event.nbytes} bytes > max_nbytes={limit}: {event}" + ) + for flag in ("blocking", "implicit"): + if flag in limits and getattr(event, flag) != limits[flag]: + found.append( + f"{name}: {flag}={getattr(event, flag)} not allowed: {event}" + ) + for kind, limits in rule.allow.items(): + of_kind = [e for e in events if e.kind == kind] + limit = limits.get("max_count") + if limit is not None and len(of_kind) > limit: + found.append(f"{name}: {len(of_kind)} {kind} > max_count={limit}") + limit = limits.get("max_total_bytes") + total = sum(e.nbytes or 0 for e in of_kind) + if limit is not None and total > limit: + found.append( + f"{name}: {total} bytes of {kind} > max_total_bytes={limit}" + ) + return found + + def report(self) -> str: + """The events of every phase (see :meth:`TransferCounter.report`) and the broken rules.""" + lines = [] + for name, counter in self.phases.items(): + lines.append(f"{name} ({self.calls.get(name, 0)} call(s)):") + lines += [f" {line}" for line in counter.report().splitlines()] + found = self.violations() + if found: + lines.append("broken rules:") + lines += [f" {line}" for line in found] + return "\n".join(lines) + + def check(self) -> None: + """Raise ``AssertionError`` with :meth:`report` if a rule is broken.""" + if self.violations(): + raise AssertionError("transfer budget exceeded:\n" + self.report()) diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py index c98421f..6ffc26c 100644 --- a/src/cunumpy/algorithms.py +++ b/src/cunumpy/algorithms.py @@ -1,18 +1,21 @@ """Array algorithms missing from NumPy/CuPy, on either backend. Morton (Z-order) keys (the same as ``cunumpy/morton.cuh`` computes in a -kernel), a stable sort of several arrays by one key, and sums per key:: +kernel), a stable sort of several arrays by one key, sums per key, and the compaction of +the live rows of particle arrays:: import cunumpy as xp keys = xp.algorithms.morton_keys(positions, lower, upper, levels) keys, order, positions = xp.algorithms.sort_by_key(keys, positions) charge = xp.algorithms.segment_sum(q, cell, n_cells) + n = xp.algorithms.compact_by_mask(alive, positions, charges) """ from cunumpy._algorithms import ( SegmentPlan, cell_offsets, + compact_by_mask, segment_boundaries, segment_sum, sort_by_key, @@ -29,6 +32,7 @@ "MAX_MORTON_LEVELS", "SegmentPlan", "cell_offsets", + "compact_by_mask", "morton_decode", "morton_encode", "morton_keys", diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index c26f3f5..a83c4ee 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -15,6 +15,11 @@ still run on the fake CuPy of :mod:`cunumpy._fake_cupy` (:func:`install_fake_cupy`, or ``CUNUMPY_FAKE_CUPY=1``); :func:`fake_cupy_active` tells whether it is in use, and ``requires_cupy`` skips the tests that launch kernels then. +:func:`fake_cupy_session` runs a whole CuPy-backend program on it (launches and +compilation emulated, see :func:`emulated_launches`), :data:`requires_device_backend` +skips tests that need a GPU or the fake CuPy, and :func:`run_in_fake_cupy_subprocess` +runs code in a child process on the fake CuPy from a test process that cannot +switch to it. A catalog's parity tests need no code per kernel when each kernel folder holds ``_test_args.py`` with ``make_args(backend, seed)`` (and @@ -49,8 +54,13 @@ def test_parity(name, kernel): from __future__ import annotations +import contextlib +import os import re -from collections.abc import Callable, Sequence +import signal +import subprocess +import sys +from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Any import array_api_compat @@ -67,7 +77,14 @@ def test_parity(name, kernel): _strip_comments, ) from cunumpy._dispatch import Kernel -from cunumpy._emulation import emulate_cuda_kernel, emulation_compiler +from cunumpy._emulation import ( + compile_for_emulation, + emulate_cuda_kernel, + emulated_launches, + emulation_cache_dir, + emulation_compiler, +) +from cunumpy._fake_cupy import host_buffer from cunumpy.xp import cupy_available, get_backend, to_numpy, use_backend # the pytest objects are created on first access, see __getattr__ @@ -76,16 +93,28 @@ def test_parity(name, kernel): "assert_kernels_agree", "backend", # noqa: F822 "check_parity", + "compile_for_emulation", + "cuda_required", + "device_backend_available", "device_function_kernel", "emulate_cuda_kernel", + "emulated_launches", + "emulation_cache_dir", "emulation_compiler", "fake_cupy_active", + "fake_cupy_session", + "host_buffer", "install_fake_cupy", "parity_cases", "requires_cupy", # noqa: F822 + "requires_device_backend", # noqa: F822 + "run_in_fake_cupy_subprocess", ] SKIP_REASON = "CuPy/GPU not available" +DEVICE_SKIP_REASON = ( + "neither a GPU nor the fake CuPy (CUNUMPY_FAKE_CUPY=1) is available" +) FAKE_SKIP_REASON = "the fake CuPy cannot run CUDA kernels" @@ -109,6 +138,185 @@ def _can_launch() -> bool: return cupy_available() and not fake_cupy_active() +def cuda_required() -> bool: + """Whether ``CUNUMPY_REQUIRE_CUDA`` demands a real GPU (the CI guard). + + Then tests that need a GPU fail instead of being skipped where there is + none (:data:`requires_cupy`, :func:`assert_kernels_agree`, the ``cupy`` + parameter of :func:`backend`), so that a CI job on a GPU machine cannot + pass silently because CuPy or the driver is broken. + """ + return os.environ.get("CUNUMPY_REQUIRE_CUDA", "").lower() in {"1", "true", "yes"} + + +def _skip_or_fail(reason: str) -> None: + """``pytest.skip``, or ``pytest.fail`` if a GPU is required.""" + if cuda_required(): + _pytest().fail(f"CUNUMPY_REQUIRE_CUDA is set, but {reason}", pytrace=False) + _pytest().skip(reason) + + +def cuda_gate() -> bool: + """Condition of ``requires_cupy`` under ``CUNUMPY_REQUIRE_CUDA``: fail, or False.""" + if not _can_launch(): + reason = FAKE_SKIP_REASON if fake_cupy_active() else SKIP_REASON + _pytest().fail(f"CUNUMPY_REQUIRE_CUDA is set, but {reason}", pytrace=False) + return False + + +def device_backend_available() -> bool: + """Whether a CuPy-backend program can run here: a GPU, or the fake CuPy. + + Unlike :data:`requires_cupy` (CUDA kernels can be launched), this is also + True on the fake CuPy, where launches only work inside + :func:`emulated_launches` (see :func:`fake_cupy_session`). The marker + :data:`requires_device_backend` skips tests that need it. + """ + return fake_cupy_active() or cupy_available() + + +def device_backend_gate() -> bool: + """Condition of ``requires_device_backend`` under ``CUNUMPY_REQUIRE_CUDA``.""" + if not device_backend_available(): + _pytest().fail( + f"CUNUMPY_REQUIRE_CUDA is set, but {DEVICE_SKIP_REASON}", pytrace=False + ) + return False + + +@contextlib.contextmanager +def fake_cupy_session( + *, + compiler: str | None = None, + options: Sequence[str] = (), +) -> Iterator[None]: + """Run a CuPy-backend program on the CPU, on the fake CuPy. + + Activates the CuPy backend (the fake CuPy) and emulates every CUDA launch + and compilation in the block (:func:`emulated_launches`, which takes + `compiler` and `options`):: + + with fake_cupy_session(): + sim.run() # kernels compiled up front and launched, all on the CPU + + Raises + ------ + RuntimeError + If the fake CuPy is not active (:func:`install_fake_cupy`, + ``CUNUMPY_FAKE_CUPY=1``; or use :func:`run_in_fake_cupy_subprocess`). + """ + if not fake_cupy_active(): + raise RuntimeError( + "fake_cupy_session() needs the fake CuPy: set CUNUMPY_FAKE_CUPY=1 or call " + "install_fake_cupy() before cunumpy is used, or use " + "run_in_fake_cupy_subprocess()", + ) + with ( + use_backend("cupy", strict=True), + emulated_launches(compiler=compiler, options=options), + ): + yield + + +#: Prefixes of the environment variables through which an MPI launcher (Open +#: MPI, MPICH/Hydra, Intel MPI, Slurm, ...) hands a process its place in the job. +MPI_LAUNCHER_PREFIXES = ( + "OMPI_", + "PMIX_", + "PMI_", + "HYDRA_", + "MPIR_", + "I_MPI_", + "SLURM_", + "MV2_", + "MPI_LOCALRANKID", + "ALPS_APP_PE", + "PALS_", +) + + +def _tail(text: str, lines: int = 50) -> str: + return "\n".join(text.splitlines()[-lines:]) + + +def run_in_fake_cupy_subprocess( + code: str, + *, + env: Mapping[str, str] | None = None, + timeout: float | None = None, +) -> subprocess.CompletedProcess: + """Run `code` in a serial child Python process on the fake CuPy; fail the test if it fails. + + The fake CuPy must be installed before anything imports cunumpy, so a test + process that already uses cunumpy cannot switch to it; the child starts + with ``CUNUMPY_FAKE_CUPY=1``. It runs ``python -X faulthandler -c code`` + (a crash prints the Python traceback) with ``OMP_NUM_THREADS=1`` and + without the variables of an MPI launcher, so it does not join the MPI job + of the parent (``MAYBEMPI=0``). Under MPI only rank 0 starts the child and + the other ranks skip the test: concurrent children are not needed for a + serial check, and have crashed external libraries. + + Parameters + ---------- + code : str + Python source to run. + env : Mapping[str, str] | None + Additional environment variables for the child (e.g. ``PYTHONPATH``). + timeout : float | None + Seconds after which the child is killed and the test fails. + + Returns + ------- + subprocess.CompletedProcess + The finished child (exit code 0), with its stdout and stderr. + """ + pytest = _pytest() + from cunumpy.mpi import get_mpi + + if get_mpi().COMM_WORLD.Get_rank() != 0: + pytest.skip("serial check in a child process, runs on MPI rank 0") + child_env = { + name: value + for name, value in os.environ.items() + if not name.startswith(MPI_LAUNCHER_PREFIXES) + } + child_env.update(CUNUMPY_FAKE_CUPY="1", OMP_NUM_THREADS="1", MAYBEMPI="0") + child_env.update(env or {}) + command = [sys.executable, "-X", "faulthandler", "-c", code] + try: + result = subprocess.run( + command, + env=child_env, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired as error: + out, err = ( + t.decode(errors="replace") if isinstance(t, bytes) else (t or "") + for t in (error.stdout, error.stderr) + ) + pytest.fail( + f"child process timed out after {timeout} s\n" + f"--- stdout (end) ---\n{_tail(out)}\n--- stderr (end) ---\n{_tail(err)}", + pytrace=False, + ) + if result.returncode != 0: + rc = result.returncode + try: + how = f"signal {signal.Signals(-rc).name}" if rc < 0 else f"exit code {rc}" + except ValueError: + how = f"signal {-rc}" + pytest.fail( + f"child process failed with {how}\n" + f"--- stdout (end) ---\n{_tail(result.stdout)}\n" + f"--- stderr (end) ---\n{_tail(result.stderr)}", + pytrace=False, + ) + return result + + # pytest objects, built on first use so that importing this module does not # import pytest (see __getattr__ below) _LAZY: dict[str, Any] = {} @@ -126,23 +334,50 @@ def _pytest() -> Any: def _build_lazy() -> None: pytest = _pytest() - requires_cupy = pytest.mark.skipif( - not _can_launch(), - reason=FAKE_SKIP_REASON if fake_cupy_active() else SKIP_REASON, - ) + if cuda_required() and not _can_launch(): + # a string condition is evaluated when the test is set up; it fails the + # test (instead of skipping it) because there is no real GPU. With a + # working GPU the marker is the ordinary one below. + requires_cupy = pytest.mark.skipif( + "__import__('cunumpy.kernel_testing', fromlist=['_']).cuda_gate()", + reason=SKIP_REASON, + ) + else: + requires_cupy = pytest.mark.skipif( + not _can_launch(), + reason=FAKE_SKIP_REASON if fake_cupy_active() else SKIP_REASON, + ) + if cuda_required() and not device_backend_available(): + requires_device_backend = pytest.mark.skipif( + "__import__('cunumpy.kernel_testing', fromlist=['_']).device_backend_gate()", + reason=DEVICE_SKIP_REASON, + ) + else: + requires_device_backend = pytest.mark.skipif( + not device_backend_available(), reason=DEVICE_SKIP_REASON + ) backends = ["numpy", pytest.param("cupy", marks=requires_cupy)] @pytest.fixture(params=backends) def backend(request): - """Run the test once per backend, with that backend active.""" - with use_backend(request.param): + """Run the test once per backend, with that backend active. + + With ``CUNUMPY_REQUIRE_CUDA`` set, the ``cupy`` run fails if the CuPy + backend cannot be activated, instead of silently running on NumPy. + """ + with use_backend(request.param, strict=cuda_required()): yield request.param - _LAZY.update(requires_cupy=requires_cupy, BACKENDS=backends, backend=backend) + _LAZY.update( + requires_cupy=requires_cupy, + requires_device_backend=requires_device_backend, + BACKENDS=backends, + backend=backend, + ) def __getattr__(name: str) -> Any: - if name in ("requires_cupy", "BACKENDS", "backend"): + if name in ("requires_cupy", "requires_device_backend", "BACKENDS", "backend"): if not _LAZY: _build_lazy() return _LAZY[name] @@ -190,9 +425,43 @@ def _arrays_in(value: Any, name: str, found: dict[str, Any], depth: int) -> None _arrays_in(item, f"{name}.{attr}", found, depth - 1) +def _resolve_output( + entry: Any, + n_args: int, + parameters: Sequence[str] | None, +) -> tuple[int, str | None]: + """The argument index and the field filter (or None) of an `outputs` entry.""" + if isinstance(entry, bool) or not isinstance(entry, (int, str)): + raise TypeError( + "outputs entries must be argument indices (int) or names (str, " + f"optionally 'name.field'), got {entry!r}", + ) + field = None + if isinstance(entry, str): + head, _, field = entry.partition(".") + field = field or None + if head.lstrip("-").isdigit(): + entry = int(head) + elif parameters is not None and head in parameters: + entry = list(parameters).index(head) + else: + known = "unknown" if parameters is None else sorted(parameters) + raise KeyError( + f"output {head!r} is not a parameter of the kernel (parameters: " + f"{known}); give an index if the names are unknown", + ) + index = entry + n_args if entry < 0 else entry + if not 0 <= index < n_args: + raise IndexError( + f"output argument {entry} does not exist: there are {n_args} arguments", + ) + return index, field + + def _collect_arrays( args: Sequence[Any], - outputs: Sequence[int] | None = None, + outputs: Sequence[int | str] | None = None, + parameters: Sequence[str] | None = None, ) -> dict[str, Any]: """The arrays among `args` (or among the arguments `outputs`), by name. @@ -205,22 +474,29 @@ def _collect_arrays( through its struct fields, ``"argument ."``, so that its arrays get the names of the attributes of the host argument object it mirrors, also when the fields are properties. + + An entry of `outputs` is an argument index, or the name of a parameter (in + `parameters`), optionally followed by ``.`` to compare only that + field (attribute) of a struct or argument object, e.g. ``"markers.positions"``. """ - indices = range(len(args)) if outputs is None else outputs + entries = range(len(args)) if outputs is None else outputs found: dict[str, Any] = {} - for entry in indices: - if not isinstance(entry, int) or isinstance(entry, bool): - raise TypeError( - "outputs must be positional argument indices (kernels take " - f"positional arguments only), got {entry!r}", - ) - index = entry + len(args) if entry < 0 else entry - if not 0 <= index < len(args): - raise IndexError( - f"output argument {entry} does not exist: there are {len(args)} " - "arguments", - ) - _arrays_in(args[index], f"argument {index}", found, depth=2) + for entry in entries: + index, field = _resolve_output(entry, len(args), parameters) + local: dict[str, Any] = {} + _arrays_in(args[index], f"argument {index}", local, depth=2) + if field is not None: + prefix = f"argument {index}.{field}" + local = { + name: array + for name, array in local.items() + if name == prefix or name.startswith((prefix + ".", prefix + "[")) + } + if not local: + raise KeyError( + f"output {entry!r}: argument {index} has no array field {field!r}", + ) + found.update(local) return found @@ -264,7 +540,7 @@ def assert_kernels_agree( rtol: float = 1e-12, atol: float = 0.0, n_calls: int = 1, - outputs: Sequence[int] | None = None, + outputs: Sequence[int | str] | None = None, seed: int = 0, ) -> dict[str, np.ndarray]: """Check that the host and CUDA versions of `kernel` compute the same. @@ -301,9 +577,12 @@ def assert_kernels_agree( n_calls : int How many times the kernel is called on each backend (e.g. to test a kernel that accumulates). - outputs : Sequence[int] | None - Indices of the arguments to compare (negative indices count from the - end), like ``PyccelKernel(outputs=...)``. By default the ``outputs`` + outputs : Sequence[int | str] | None + The arguments to compare, like ``PyccelKernel(outputs=...)``: indices + (negative indices count from the end) or parameter names. A name with + ``.`` (``"markers.positions"``) compares only that field of a + struct or argument object, leaving the other fields (e.g. buffers the + two kernels fill differently) out. 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; for a tuple, list, dict or object argument (e.g. a ``CudaArguments`` object), the arrays @@ -328,7 +607,7 @@ def assert_kernels_agree( Notes ----- The test is skipped with ``pytest.skip`` if CuPy or a GPU is not available, - or if the fake CuPy is active. + or if the fake CuPy is active; with ``CUNUMPY_REQUIRE_CUDA=1`` it fails instead. """ if not isinstance(kernel, Kernel): raise TypeError(f"expected a Kernel, got {type(kernel).__name__}") @@ -341,10 +620,11 @@ def assert_kernels_agree( if outputs is None: outputs = kernel.host_kernel.outputs if fake_cupy_active(): - _pytest().skip(FAKE_SKIP_REASON) + _skip_or_fail(FAKE_SKIP_REASON) if not cupy_available(): - _pytest().skip(SKIP_REASON) + _skip_or_fail(SKIP_REASON) + parameters = kernel.host_parameters() results = {} for backend in ("numpy", "cupy"): with use_backend(backend): @@ -354,7 +634,7 @@ def assert_kernels_agree( launch = n_threads(args) if callable(n_threads) else n_threads for _ in range(n_calls): kernel(*args, n_threads=launch, grid=grid, block=block) - results[backend] = _collect_arrays(args, outputs) + results[backend] = _collect_arrays(args, outputs, parameters) host = {name: to_numpy(a) for name, a in results["numpy"].items()} _compare_results(host, results["cupy"], rtol, atol, kernel.name) diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 6d4d644..30aef3f 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -17,6 +17,10 @@ Device dispatch can require CUDA with :func:`set_device_kernel_implementation` or temporarily with :func:`use_device_kernel_implementation`. +A :class:`MetalKernel` runs a Metal Shading Language kernel on the GPU of an +Apple silicon Mac through MLX (``pip install 'cunumpy[metal]'``), on NumPy +float32 arrays. + Argument objects for CUDA kernels (:class:`~cunumpy.arguments.CudaStruct`, ...) are in :mod:`cunumpy.arguments`, the device runtime (streams, devices, debug mode) in :mod:`cunumpy.cuda`, and the pytest helpers for kernel pairs in @@ -36,11 +40,13 @@ get_device_kernel_implementation, get_host_kernel_implementation, kernel_output, + outputs_from_annotations, set_device_kernel_implementation, set_host_kernel_implementation, use_device_kernel_implementation, use_host_kernel_implementation, ) +from cunumpy._metal_kernel import MetalKernel, metal_available __all__ = [ "DEVICE_IMPLEMENTATIONS", @@ -51,12 +57,15 @@ "HostImplementations", "Kernel", "KernelCatalog", + "MetalKernel", "PyccelKernel", "as_kernel_array", "fuse", "get_device_kernel_implementation", "get_host_kernel_implementation", "kernel_output", + "metal_available", + "outputs_from_annotations", "set_device_kernel_implementation", "set_host_kernel_implementation", "use_device_kernel_implementation", diff --git a/src/cunumpy/memory.py b/src/cunumpy/memory.py index 7fb400b..5fe2b5e 100644 --- a/src/cunumpy/memory.py +++ b/src/cunumpy/memory.py @@ -17,6 +17,6 @@ """ from cunumpy._mirror import DeviceMirror -from cunumpy._staging import HostStaging, StagedCopy +from cunumpy._staging import HostCopy, HostStaging, StagedCopy -__all__ = ["DeviceMirror", "HostStaging", "StagedCopy"] +__all__ = ["DeviceMirror", "HostCopy", "HostStaging", "StagedCopy"] diff --git a/src/cunumpy/profiling.py b/src/cunumpy/profiling.py index 662e797..a96c435 100644 --- a/src/cunumpy/profiling.py +++ b/src/cunumpy/profiling.py @@ -3,7 +3,8 @@ :func:`timed_region` times a block (synchronizing the device first and last), :class:`nvtx_range` marks it for Nsight (a no-op without NVTX), and :func:`count_transfers` / :func:`assert_no_transfers` count the copies between -host and device that cunumpy makes inside a block:: +host and device that cunumpy makes inside a block, and :class:`TransferBudget` +counts them per phase of a program and checks a rule for each phase:: import cunumpy as xp @@ -15,6 +16,7 @@ from cunumpy._profiling import Timing, nvtx_range, timed_region from cunumpy._transfers import ( + TransferBudget, TransferCounter, TransferEvent, assert_no_transfers, @@ -23,6 +25,7 @@ __all__ = [ "Timing", + "TransferBudget", "TransferCounter", "TransferEvent", "assert_no_transfers", diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 93f0d48..bb94eb5 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -20,6 +20,7 @@ _is_device_copy, _nbytes, _record, + _record_sync, ) if os.environ.get("CUNUMPY_FAKE_CUPY", "").strip().lower() in ("1", "true", "yes"): @@ -248,6 +249,8 @@ def default_float_dtype() -> Any: def synchronize() -> None: """Wait for all kernels in all streams on current device to complete.""" if array_backend.backend == "cupy": + if _COUNTERS: + _record_sync("synchronize()") try: import cupy as cp @@ -266,7 +269,7 @@ def synchronize() -> None: def _to_numpy(array: Any) -> np.ndarray: """`to_numpy` without transfer counting, for internal use.""" if get_array_backend(array) == "cupy": - return array.get() + return array.get(order="A") return np.asarray(array) @@ -286,6 +289,7 @@ def to_numpy(array: Any) -> np.ndarray: A CuPy array is copied to the host, which `count_transfers()` counts as a ``to_host`` transfer; anything else is passed through `numpy.asarray`. + Fortran-contiguous CuPy arrays keep F order; other CuPy arrays use C order. """ result = _to_numpy(array) if _COUNTERS and get_array_backend(array) == "cupy": diff --git a/tests/unit/pyccel_kernels.py b/tests/unit/pyccel_kernels.py index 7e6ab09..b9d20d4 100644 --- a/tests/unit/pyccel_kernels.py +++ b/tests/unit/pyccel_kernels.py @@ -20,6 +20,13 @@ def scale_inplace(x: "float[:]", factor: float): x[i] = x[i] * factor +def scale_fortran_inplace(x: "float[:, :](order=F)", factor: float): # noqa: F821 + """Scale an F-ordered matrix, including complete column-block views.""" + for j in range(x.shape[1]): + for i in range(x.shape[0]): + x[i, j] = x[i, j] * factor + + def dot(x: "float[:]", y: "float[:]") -> float: """Return the dot product of `x` and `y` (a scalar return value).""" result = 0.0 diff --git a/tests/unit/test_automatic_launch.py b/tests/unit/test_automatic_launch.py index 4668a18..4ff15b4 100644 --- a/tests/unit/test_automatic_launch.py +++ b/tests/unit/test_automatic_launch.py @@ -60,6 +60,17 @@ def test_callbacks_and_explicit_opt_out(): assert kernel.launch_shape(args=(np.empty((300, 7)),)) == ((3,), (128,)) +def test_last_axis_launches_one_thread_per_marker_of_component_major_arrays(): + kernel = CudaKernel(SOURCE, "work", n_threads_from="last_axis") + # (ncomp, N) positions first: N threads, not ncomp + assert kernel.n_threads_from((2.0, np.empty((3, 1000)), np.empty(5))) == 1000 + assert kernel.launch_shape(args=(np.empty((3, 1000)),)) == ((8,), (128,)) + # a 1D per-marker array, and a 0D array before it is skipped + assert kernel.n_threads_from((np.array(1.0), np.empty(300))) == 300 + with pytest.raises(TypeError, match="needs an array argument"): + kernel.launch_shape(args=(2.0,)) + + def test_invalid_setting_preserves_auto_inference(): kernel = CudaKernel(SOURCE, "work") with pytest.raises(TypeError, match="n_threads_from"): diff --git a/tests/unit/test_compact_by_mask.py b/tests/unit/test_compact_by_mask.py new file mode 100644 index 0000000..7796fee --- /dev/null +++ b/tests/unit/test_compact_by_mask.py @@ -0,0 +1,78 @@ +"""compact_by_mask keeps the masked rows, in order, at the front of every array.""" + +import numpy as np +import pytest + +import cunumpy as xp + + +def test_rows_are_moved_to_the_front_in_order_for_every_array(): + markers = np.arange(12.0).reshape(4, 3) + weights = np.array([1.0, 2.0, 3.0, 4.0]) + ids = np.array([10, 11, 12, 13]) + alive = np.array([True, False, True, True]) + n = xp.algorithms.compact_by_mask(alive, markers, weights, ids) + assert n == 3 + np.testing.assert_array_equal(markers[:n], [[0, 1, 2], [6, 7, 8], [9, 10, 11]]) + np.testing.assert_array_equal(weights[:n], [1.0, 3.0, 4.0]) + np.testing.assert_array_equal(ids[:n], [10, 12, 13]) + + +def test_all_true_none_true_and_empty(): + a = np.arange(4) + assert xp.algorithms.compact_by_mask(np.ones(4, bool), a) == 4 + np.testing.assert_array_equal(a, np.arange(4)) + assert xp.algorithms.compact_by_mask(np.zeros(4, bool), a) == 0 + assert xp.algorithms.compact_by_mask(np.zeros(0, bool), np.zeros(0)) == 0 + + +def test_matches_boolean_indexing_on_random_masks(): + rng = np.random.default_rng(1) + for _ in range(20): + a = rng.random((50, 2)) + mask = rng.random(50) < 0.4 + expected = a[mask] + n = xp.algorithms.compact_by_mask(mask, a) + np.testing.assert_array_equal(a[:n], expected) + + +def test_argument_errors(): + with pytest.raises(TypeError, match="boolean"): + xp.algorithms.compact_by_mask(np.array([1, 0]), np.zeros(2)) + with pytest.raises(ValueError, match="2 entries along axis 0"): + xp.algorithms.compact_by_mask(np.array([True, False]), np.zeros(3)) + + +def test_last_axis_compacts_component_major_and_scalar_arrays_together(): + positions = np.arange(12.0).reshape(3, 4) # (ncomp, N) + weights = np.array([1.0, 2.0, 3.0, 4.0]) # (N,) + alive = np.array([False, True, True, False]) + n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1) + assert n == 2 + np.testing.assert_array_equal(positions[:, :n], [[1, 2], [5, 6], [9, 10]]) + np.testing.assert_array_equal(weights[:n], [2.0, 3.0]) + + +def test_axis_matches_boolean_indexing_on_random_masks(): + rng = np.random.default_rng(2) + for axis in (1, -1): + a = rng.random((3, 40)) + mask = rng.random(40) < 0.5 + expected = a[:, mask] + n = xp.algorithms.compact_by_mask(mask, a, axis=axis) + np.testing.assert_array_equal(a[:, :n], expected) + + +def test_compacting_a_column_prefix_view_keeps_the_storage(): + storage = np.arange(20.0).reshape(2, 10) + view = storage[:, :6] # rows further apart than the view is wide + n = xp.algorithms.compact_by_mask(np.arange(6) % 2 == 0, view, axis=-1) + np.testing.assert_array_equal(storage[:, :n], [[0, 2, 4], [10, 12, 14]]) + np.testing.assert_array_equal(storage[:, 6:], [[6, 7, 8, 9], [16, 17, 18, 19]]) + + +def test_axis_errors(): + with pytest.raises(ValueError, match="3 entries along axis -1"): + xp.algorithms.compact_by_mask(np.ones(3, bool), np.zeros((3, 2)), axis=-1) + with pytest.raises(ValueError, match="no axis 1"): + xp.algorithms.compact_by_mask(np.ones(3, bool), np.zeros(3), axis=1) diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 0ef8c8c..b268f36 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -2306,6 +2306,7 @@ def __init__(self): def __call__(self, grid, block, args, shared_mem=0): self.launches.append((grid, block, shared_mem)) + self.args = args @pytest.fixture @@ -2333,6 +2334,28 @@ def recorded(monkeypatch): return kernel, raw +def test_launch_passes_structs_as_size_one_numpy_arrays(recorded, monkeypatch): + kernel, raw = recorded + packed = np.zeros((), dtype=[("data", np.uintp), ("shape", np.int64, (1,))]) + packed["data"] = 0x1000 + packed["shape"] = (7,) + scalar = np.float64(2.0) + device_array = FakeDeviceArray(np.float64, shape=(7,)) + monkeypatch.setattr( + kernel, "prepare_args", lambda *args: (packed[()], scalar, device_array) + ) + + kernel(n_threads=1) + + struct_arg, scalar_arg, pointer_arg = raw.args + assert isinstance(struct_arg, np.ndarray) + assert struct_arg.size == 1 + assert struct_arg.dtype == packed.dtype + assert struct_arg.tobytes() == packed.tobytes() + assert scalar_arg is scalar + assert pointer_arg is device_array + + def test_n_threads_from_first_array(recorded): kernel, raw = recorded kernel.n_threads_from = "first_array" diff --git a/tests/unit/test_cupy.py b/tests/unit/test_cupy.py index a8e804d..f03baaf 100644 --- a/tests/unit/test_cupy.py +++ b/tests/unit/test_cupy.py @@ -2,6 +2,50 @@ import pytest import cunumpy as xp +from cunumpy.kernels import PyccelKernel + + +@pytest.mark.parametrize("order", ["C", "F"]) +@pytest.mark.parametrize("column_block", [False, True]) +@pytest.mark.parametrize( + "conversion", ["to_numpy", "host_call", "evaluate_on_host", "pyccel_kernel"] +) +def test_host_conversion_preserves_order(order, column_block, conversion): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + + expected = np.array(np.arange(30.0).reshape(5, 6), order=order) + with xp.use_backend("cupy"): + device = xp.to_cupy(expected) + if column_block: + device = device[:, 1:4] + expected = expected[:, 1:4] + + def check_host(array): + assert isinstance(array, np.ndarray) + np.testing.assert_array_equal(array, expected) + assert array.flags.f_contiguous == (order == "F") + assert array.flags.c_contiguous == (order == "C") + return array + + if conversion == "to_numpy": + check_host(xp.to_numpy(device)) + else: + if conversion == "host_call": + result = xp.host_call(check_host, device) + elif conversion == "evaluate_on_host": + + class Model: + @xp.evaluate_on_host + def evaluate(self, array): + return check_host(array) + + result = Model().evaluate(device) + else: + result = PyccelKernel(check_host)(device) + + assert xp.is_gpu(result) + check_host(xp.to_numpy(result)) def test_to_cupy_available(): diff --git a/tests/unit/test_emulated_launches.py b/tests/unit/test_emulated_launches.py new file mode 100644 index 0000000..ac97ade --- /dev/null +++ b/tests/unit/test_emulated_launches.py @@ -0,0 +1,253 @@ +"""Struct parameters in the emulation, and `emulated_launches` on the fake CuPy.""" + +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +from cunumpy.arguments import CudaStruct +from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler +from cunumpy.kernels import CudaKernel + +pytestmark = pytest.mark.skipif( + emulation_compiler() is None, + reason="no C++ compiler for the emulation", +) + +PARTICLES = CudaStruct( + "Particles", + [("markers", "Array2D"), ("alive", "bool*"), ("n", "int")], +) +SOURCE = ( + '#include "cunumpy/array_view.cuh"\n' + + PARTICLES.declaration + + r""" +extern "C" __global__ +void push(Particles p, double dt, double* total) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < p.n && p.alive[i]) { + p.markers(i, 0) += dt * p.markers(i, 1); + total[0] += p.markers(i, 0); + } +} +""" +) + + +def particles(): + markers = np.arange(12.0).reshape(4, 3) + return markers, np.array([True, True, False, True]) + + +def expected(markers, alive, dt): + out = markers.copy() + out[alive, 0] += dt * out[alive, 1] + return out + + +@pytest.mark.parametrize("kind", ["mapping", "object"]) +def test_struct_given_as_mapping_or_object(kind): + markers, alive = particles() + want = expected(markers, alive, 0.5) + total = np.zeros(1) + fields = {"markers": markers, "alive": alive, "n": 4} + value = fields if kind == "mapping" else type("Args", (), fields)() + push = CudaKernel(SOURCE, "push", structs=[PARTICLES]) + emulate_cuda_kernel(push, value, 0.5, total, n_threads=4) + np.testing.assert_allclose(markers, want) + assert total[0] == pytest.approx(want[alive, 0].sum()) + + +def test_struct_launch_shape_is_inferred_from_the_first_array(): + markers, alive = particles() + push = CudaKernel(SOURCE, "push", structs=[PARTICLES]) + emulate_cuda_kernel( + push, {"markers": markers, "alive": alive, "n": 4}, 1.0, np.zeros(1) + ) + np.testing.assert_allclose(markers, expected(*particles(), 1.0)) + + +def test_struct_with_a_missing_field_is_a_type_error(): + markers, alive = particles() + push = CudaKernel(SOURCE, "push", structs=[PARTICLES]) + with pytest.raises(TypeError, match="no value for the field 'n'"): + emulate_cuda_kernel( + push, {"markers": markers, "alive": alive}, 1.0, np.zeros(1), n_threads=4 + ) + + +SCRIPT = r""" +import numpy as np +import cunumpy as xp +from cunumpy.arguments import CudaStruct, CudaStructArguments +from cunumpy.kernel_testing import emulated_launches, host_buffer +from cunumpy.kernels import CudaKernel, Kernel + +Particles = CudaStruct("Particles", [("x", "double*"), ("n", "int")]) +source = Particles.declaration + ''' +extern "C" __global__ void scale(Particles p, double f) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < p.n) p.x[i] *= f; +}''' +scale = CudaKernel(source, "scale", structs=[Particles]) + + +class Args(CudaStructArguments): + struct_name = "Particles" + fields = (("x", "double*"), ("n", "int")) + + def __init__(self, x): + self.x = x + self.n = x.shape[0] + self.pack() + + +x = xp.asarray(np.arange(5.0)) +assert xp.get_backend() == "cupy" +try: + scale(Particles(x=x, n=5), 2.0, n_threads=5) +except NotImplementedError: + pass +else: + raise AssertionError("the fake CuPy must refuse a launch outside the block") + +with emulated_launches(): + scale(Particles(x=x, n=5), 2.0, n_threads=5) + scale(Args(x), 3.0, n_threads=5) + host_view = host_buffer(x) +np.testing.assert_allclose(host_buffer(x), np.arange(5.0) * 6.0) +assert host_view is host_buffer(x) # not a copy + +# a Kernel dispatches to the CUDA kernel on the CuPy backend +def host_scale(args, f): + raise AssertionError("the CUDA kernel must run") + +kernel = Kernel(host_scale, scale) +with emulated_launches(): + kernel(Args(x), 0.5, n_threads=5) +np.testing.assert_allclose(host_buffer(x), np.arange(5.0) * 3.0) + +try: + host_buffer(np.zeros(2)) +except TypeError: + pass +else: + raise AssertionError("host_buffer must refuse a NumPy array") +print("emulated launches OK") +""" + + +def test_emulated_launches_on_the_fake_cupy(): + root = Path(__file__).resolve().parents[2] + env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", CUNUMPY_BACKEND="cupy") + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(root / "src"), env.get("PYTHONPATH", "")) if p + ) + env.pop("CUNUMPY_CUDA_DEBUG", None) + result = subprocess.run( + [sys.executable, "-c", SCRIPT], + env=env, + capture_output=True, + text=True, + check=False, + cwd=str(root), + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "emulated launches OK" in result.stdout + + +COMPILE_SCRIPT = r""" +import numpy as np +import cunumpy as xp +from cunumpy.kernel_testing import emulated_launches, host_buffer +from cunumpy.kernels import CudaKernel, CudaKernelVariants, Kernel, KernelCatalog + +original = CudaKernel.compile +source = ''' +template +__global__ void power(double* x, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) { double v = x[i]; for (int k = 1; k < P; ++k) x[i] *= v; } +}''' +variants = CudaKernelVariants(lambda p: CudaKernel(source, "power", template_args=[p])) +scale = CudaKernel(''' +extern "C" __global__ void scale(double* x, double f, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) x[i] *= f; +}''', "scale") +catalog = KernelCatalog({"scale": Kernel(lambda x, f, n: None, scale)}) + +with emulated_launches(): + assert scale.compile() is None + assert scale.recompile() is None + variants.compile_all([(2,), (3,)], jobs=2) + assert catalog.compile_all() == ["scale"] + x = xp.asarray(np.arange(4.0)) + scale(x, 2.0, 4, n_threads=4) + variants.get(2)(x, 4, n_threads=4) +np.testing.assert_allclose(host_buffer(x), (2.0 * np.arange(4.0)) ** 2) +assert CudaKernel.compile is original + +# a compile error shows up in compile(), without a launch +broken = CudaKernel('extern "C" __global__ void k(int n) { undefined_call(n); }', "k") +try: + with emulated_launches(): + broken.compile() +except RuntimeError as error: + assert "does not compile" in str(error), error +else: + raise AssertionError("compile() must report the compile error") +assert CudaKernel.compile is original + +# kernels the emulation cannot run are skipped by compile() +warp = CudaKernel('''extern "C" __global__ void w(double* x) { + x[0] = __shfl_down_sync(0xffffffff, x[0], 1); }''', "w") +with emulated_launches(): + warp.compile() + +# outside the block, the fake CuPy cannot compile +try: + scale.compile() +except (RuntimeError, NotImplementedError): + pass +else: + raise AssertionError("compile() outside the block must not emulate") +print("emulated compile OK") +""" + + +def _run_on_fake_cupy(script): + root = Path(__file__).resolve().parents[2] + env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", CUNUMPY_BACKEND="cupy") + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(root / "src"), env.get("PYTHONPATH", "")) if p + ) + env.pop("CUNUMPY_CUDA_DEBUG", None) + return subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + check=False, + cwd=str(root), + ) + + +def test_compile_inside_emulated_launches_on_the_fake_cupy(): + result = _run_on_fake_cupy(COMPILE_SCRIPT) + assert result.returncode == 0, result.stdout + result.stderr + assert "emulated compile OK" in result.stdout + + +def test_compile_and_call_are_restored_after_an_exception(): + from cunumpy.kernel_testing import emulated_launches + + call, compile_ = CudaKernel.__call__, CudaKernel.compile + with pytest.raises(KeyError), emulated_launches(): + assert CudaKernel.compile is not compile_ + raise KeyError("boom") + assert CudaKernel.compile is compile_ + assert CudaKernel.__call__ is call diff --git a/tests/unit/test_emulation.py b/tests/unit/test_emulation.py index c5f28f6..87fbac3 100644 --- a/tests/unit/test_emulation.py +++ b/tests/unit/test_emulation.py @@ -1,5 +1,7 @@ """Tests for `cunumpy.kernel_testing.emulate_cuda_kernel`: CUDA kernels run on the CPU.""" +import weakref + import numpy as np import pytest @@ -456,3 +458,174 @@ def test_high_dimensional_bounds_check(contiguous): n_threads=1, options=("-DCUNUMPY_BOUNDS_CHECK",), ) + + +@pytest.fixture +def fresh_cache(tmp_path, monkeypatch): + """An empty disk cache and no libraries loaded; yields the list of compiler calls.""" + import subprocess + + from cunumpy import _emulation + + monkeypatch.setenv("CUNUMPY_EMULATION_CACHE", str(tmp_path / "cache")) + monkeypatch.setattr(_emulation, "_LIBRARIES", {}) + monkeypatch.setattr(_emulation, "_PREPARED", weakref.WeakKeyDictionary()) + calls = [] + run = subprocess.run + + def counting_run(command, *args, **kwargs): + calls.append(command) + return run(command, *args, **kwargs) + + monkeypatch.setattr(_emulation.subprocess, "run", counting_run) + return calls + + +def test_a_kernel_is_compiled_once_for_all_launches(fresh_cache): + kernel = CudaKernel(AXPY, "axpy") + for n, a in [(10, 2.0), (1000, -0.5), (3, 7.0)]: + x, y = np.arange(float(n)), np.ones(n) + emulate_cuda_kernel(kernel, a, x, y, n, options=("-ffp-contract=off",)) + np.testing.assert_array_equal(y, 1.0 + a * np.arange(float(n))) + # an equal kernel object shares the library + emulate_cuda_kernel( + CudaKernel(AXPY, "axpy"), 1.0, x, y, 3, options=("-ffp-contract=off",) + ) + assert len(fresh_cache) == 1 + + +def test_options_and_template_arguments_rebuild(fresh_cache): + x = np.array([1.0, 2.0, 3.0], dtype=np.float32) + for k in (2, 3, 2): + emulate_cuda_kernel( + CudaKernel(TEMPLATE, "power", template_args=(np.float32, k)), x, 3 + ) + assert x.tolist() == [1.0, 2.0**12, 3.0**12] + assert len(fresh_cache) == 2 + emulate_cuda_kernel( + CudaKernel(TEMPLATE, "power", template_args=(np.float32, 2)), + x, + 3, + options=("-DUNUSED=1",), + ) + assert len(fresh_cache) == 3 + + +def test_libraries_are_reused_from_the_disk_cache(fresh_cache, monkeypatch): + from cunumpy import _emulation + + kernel = CudaKernel(AXPY, "axpy") + y = np.zeros(4) + emulate_cuda_kernel(kernel, 1.0, np.ones(4), y, 4) + # a new process + monkeypatch.setattr(_emulation, "_LIBRARIES", {}) + monkeypatch.setattr(_emulation, "_PREPARED", weakref.WeakKeyDictionary()) + emulate_cuda_kernel(kernel, 1.0, np.ones(4), y, 4) + assert len(fresh_cache) == 1 + np.testing.assert_array_equal(y, 2.0) + assert len(list(_emulation.emulation_cache_dir().glob("*.so"))) == 1 + + +def test_without_disk_cache(fresh_cache, monkeypatch): + from cunumpy import _emulation + + monkeypatch.setenv("CUNUMPY_EMULATION_CACHE", "0") + assert _emulation.emulation_cache_dir() is None + y = np.zeros(4) + emulate_cuda_kernel(CudaKernel(AXPY, "axpy"), 1.0, np.ones(4), y, 4) + np.testing.assert_array_equal(y, 1.0) + + +def test_compile_for_emulation_builds_without_a_launch(fresh_cache): + from cunumpy.kernel_testing import compile_for_emulation + + compile_for_emulation(CudaKernel(AXPY, "axpy")) + y = np.zeros(2) + emulate_cuda_kernel(CudaKernel(AXPY, "axpy"), 1.0, np.ones(2), y, 2) + assert len(fresh_cache) == 1 + with pytest.raises(RuntimeError, match="does not compile"): + compile_for_emulation( + CudaKernel('extern "C" __global__ void k(int n) { n = ; }', "k") + ) + unparsed = CudaKernel(AXPY, "axpy", check_signature=False) + compile_for_emulation(unparsed) # a syntax check only + + +def test_arrays_are_used_in_place_and_aliases_see_each_other(): + source = r""" + extern "C" __global__ void shift(const double* a, double* b, int n) { + if (blockIdx.x == 0 && threadIdx.x == 0) + for (int i = 1; i < n; ++i) b[i] = a[i - 1]; + }""" + x = np.arange(5.0) + emulate_cuda_kernel(CudaKernel(source, "shift"), x, x, 5, n_threads=1) + np.testing.assert_array_equal(x, 0.0) # serial: each write is read next + readonly = np.arange(5.0) + readonly.flags.writeable = False + out = np.zeros(5) + emulate_cuda_kernel(CudaKernel(source, "shift"), readonly, out, 5, n_threads=1) + np.testing.assert_array_equal(out, [0, 0, 1, 2, 3]) + + +def test_trap_in_a_kernel_with_barriers_is_reported(): + source = r""" + extern "C" __global__ void k(double* x, int bad) { + __shared__ double s[4]; + s[threadIdx.x] = x[threadIdx.x]; + __syncthreads(); + if (bad && threadIdx.x == 2) __trap(); + x[threadIdx.x] = s[3 - threadIdx.x]; + }""" + kernel = CudaKernel(source, "k") + with pytest.raises(RuntimeError, match="crashed"): + emulate_cuda_kernel(kernel, np.arange(4.0), 1, grid=1, block=4) + x = np.arange(4.0) + emulate_cuda_kernel(kernel, x, 0, grid=1, block=4) # still usable + np.testing.assert_array_equal(x, [3, 2, 1, 0]) + + +INLINE_ASM = r""" +extern "C" __global__ void map(double* x, int kind) { + if (kind == 0) { + x[0] = 1.0; + } else if (kind == 1) { + asm volatile("trap;"); + } else { + asm("trap;"); // unknown kind + } +} +""" + + +def test_inline_asm_is_trapped(): + kernel = CudaKernel(INLINE_ASM, "map") + x = np.zeros(1) + emulate_cuda_kernel(kernel, x, 0, n_threads=1) # no extra options needed + assert x[0] == 1.0 + for kind in (1, 2): + with pytest.raises(RuntimeError, match="__trap"): + emulate_cuda_kernel(kernel, x, kind, n_threads=1) + # the option struphy used before still works + emulate_cuda_kernel(kernel, x, 0, n_threads=1, options=("-Dasm(x)=__trap()",)) + + +def test_compile_for_emulation_sees_changed_headers(fresh_cache, tmp_path): + from cunumpy.kernel_testing import compile_for_emulation + + header = tmp_path / "factor.cuh" + header.write_text("#define FACTOR 2.0\n") + source = r""" + #include "factor.cuh" + extern "C" __global__ void k(double* x) { + if (blockIdx.x == 0 && threadIdx.x == 0) x[0] *= FACTOR; + }""" + kernel = CudaKernel(source, "k", include_dirs=[tmp_path]) + x = np.ones(1) + emulate_cuda_kernel(kernel, x, n_threads=1) + header.write_text("#define FACTOR 3.0\n") + emulate_cuda_kernel(kernel, x, n_threads=1) # launches reuse the library + assert x[0] == 4.0 + compile_for_emulation(kernel) # as kernel.recompile() in emulated_launches + emulate_cuda_kernel(kernel, x, n_threads=1) + assert x[0] == 12.0 + assert len(fresh_cache) == 2 diff --git a/tests/unit/test_execution_transfers.py b/tests/unit/test_execution_transfers.py index da5ba0a..aa9806e 100644 --- a/tests/unit/test_execution_transfers.py +++ b/tests/unit/test_execution_transfers.py @@ -31,10 +31,10 @@ def set(self, host): cp.copies.append("set") np.copyto(self.array, host) - def get(self, out=None, stream=None): + def get(self, out=None, stream=None, order="C"): cp.copies.append(("get", stream)) if out is None: - return self.array.copy() + return self.array.copy(order=order) np.copyto(out, self.array) return out diff --git a/tests/unit/test_fake_cupy_programs.py b/tests/unit/test_fake_cupy_programs.py new file mode 100644 index 0000000..8309f9f --- /dev/null +++ b/tests/unit/test_fake_cupy_programs.py @@ -0,0 +1,102 @@ +"""Tests for the helpers that run whole CuPy-backend programs on the fake CuPy. + +`device_backend_available`, `requires_device_backend`, `fake_cupy_session` and +`run_in_fake_cupy_subprocess` of `cunumpy.kernel_testing`. +""" + +from pathlib import Path + +import pytest + +import cunumpy as xp +from cunumpy.kernel_testing import ( + device_backend_available, + emulation_compiler, + fake_cupy_active, + fake_cupy_session, + run_in_fake_cupy_subprocess, +) + +SRC = str(Path(__file__).resolve().parents[2] / "src") + + +def run(code, **kwargs): + return run_in_fake_cupy_subprocess(code, env={"PYTHONPATH": SRC}, **kwargs) + + +def test_device_backend_available(): + assert device_backend_available() == (fake_cupy_active() or xp.cupy_available()) + result = run( + "from cunumpy.kernel_testing import device_backend_available, requires_device_backend\n" + "assert device_backend_available()\n" + "assert not requires_device_backend.args[0]\n" + "print('available')" + ) + assert "available" in result.stdout + + +@pytest.mark.skipif(fake_cupy_active(), reason="the fake CuPy is active here") +def test_fake_cupy_session_needs_the_fake_cupy(): + with pytest.raises(RuntimeError, match="CUNUMPY_FAKE_CUPY"), fake_cupy_session(): + pass + + +@pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") +def test_fake_cupy_session_runs_a_cupy_program(): + code = r""" +import numpy as np +import cunumpy as xp +from cunumpy.kernel_testing import fake_cupy_session +from cunumpy.kernels import CudaKernel, KernelCatalog, Kernel + +scale = CudaKernel(''' +extern "C" __global__ void scale(double* x, double f, int n) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < n) x[i] *= f; +}''', "scale") +kernel = Kernel(lambda x, f, n: None, scale) +before = xp.get_backend() # CUNUMPY_BACKEND may already select CuPy +with fake_cupy_session(): + assert xp.get_backend() == "cupy" + KernelCatalog({"scale": kernel}).compile_all() + x = xp.arange(4.0) + kernel(x, 3.0, 4, n_threads=4) + print(xp.to_numpy(x).tolist()) +assert xp.get_backend() == before +""" + assert "[0.0, 3.0, 6.0, 9.0]" in run(code).stdout + + +def test_child_runs_serially_on_the_fake_cupy(monkeypatch): + monkeypatch.setenv("OMPI_MCA_test_variable", "1") + result = run( + "import os\n" + "assert os.environ['CUNUMPY_FAKE_CUPY'] == '1'\n" + "assert os.environ['OMP_NUM_THREADS'] == '1'\n" + "assert 'OMPI_MCA_test_variable' not in os.environ\n" + "from cunumpy.kernel_testing import fake_cupy_active\n" + "assert fake_cupy_active()\n" + "print('serial')" + ) + assert result.returncode == 0 + assert "serial" in result.stdout + + +def test_a_failing_child_fails_the_test_with_the_end_of_its_output(): + code = "print('\\n'.join(f'line {i}' for i in range(100)))\nraise SystemExit(3)" + with pytest.raises(pytest.fail.Exception) as info: + run(code) + message = str(info.value) + assert "exit code 3" in message + assert "line 99" in message and "line 49" not in message + + +def test_a_crashing_child_reports_the_signal(): + code = "import os, signal\nos.kill(os.getpid(), signal.SIGSEGV)" + with pytest.raises(pytest.fail.Exception, match="signal SIGSEGV"): + run(code) + + +def test_a_child_that_hangs_times_out(): + with pytest.raises(pytest.fail.Exception, match="timed out"): + run("import time\nprint('started', flush=True)\ntime.sleep(60)", timeout=2) diff --git a/tests/unit/test_host.py b/tests/unit/test_host.py new file mode 100644 index 0000000..11456c1 --- /dev/null +++ b/tests/unit/test_host.py @@ -0,0 +1,82 @@ +"""host_call, evaluate_on_host and setup_on_host, with a stand-in for device arrays.""" + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import _host + + +class FakeDevice: + """Stands in for a CuPy array: the host helpers only convert it.""" + + def __init__(self, array): + self.array = np.asarray(array) + + +@pytest.fixture +def device(monkeypatch): + monkeypatch.setattr(_host.xp, "is_gpu", lambda a: isinstance(a, FakeDevice)) + monkeypatch.setattr( + _host.xp, "to_numpy", lambda a: a.array if isinstance(a, FakeDevice) else a + ) + monkeypatch.setattr(_host.xp, "to_cupy", FakeDevice) + + +def test_host_arrays_are_a_plain_call(): + x = np.arange(3.0) + out = xp.host_call(lambda a, scale: a * scale, x, scale=2.0) + np.testing.assert_array_equal(out, 2 * x) + + +def test_device_arguments_run_on_numpy_and_return_on_device(device): + seen = {} + + def fun(a, b, scale=1.0): + seen["backend"] = xp.get_backend() + seen["types"] = (type(a), type(b[0])) + return a + b[0], 3.0 + + a, b = FakeDevice([1.0, 2.0]), [FakeDevice([10.0, 20.0])] + arr, scalar = xp.host_call(fun, a, b, scale=2.0) + assert seen["backend"] == "numpy" and seen["types"] == (np.ndarray, np.ndarray) + assert isinstance(arr, FakeDevice) and scalar == 3.0 + np.testing.assert_array_equal(arr.array, [11.0, 22.0]) + + +def test_device_keyword_argument_is_detected(device): + out = xp.host_call(lambda x=None: x * 2, x=FakeDevice([1.0])) + assert isinstance(out, FakeDevice) + + +def test_evaluate_on_host_binds_self(device): + class Model: + offset = 1.0 + + @xp.evaluate_on_host + def f(self, x): + """Docstring.""" + assert isinstance(x, np.ndarray) + return x + self.offset + + assert Model.f.__doc__ == "Docstring." + out = Model().f(FakeDevice([1.0])) + np.testing.assert_array_equal(out.array, [2.0]) + + +def test_setup_on_host_runs_init_on_numpy_backend(): + class Model: + @xp.setup_on_host + def __init__(self, n): + self.backend = xp.get_backend() + self.data = xp.arange(n) + + model = Model(3) + assert model.backend == "numpy" and isinstance(model.data, np.ndarray) + + +def test_copies_are_counted(monkeypatch): + """With the real conversions, only to_numpy/to_cupy on device arrays count.""" + with xp.profiling.count_transfers() as counter: + xp.host_call(lambda a: a + 1, np.zeros(2)) + assert counter.total == 0 diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index a04f14b..bec4c58 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -594,6 +594,70 @@ def test_as_kernel_array_on_the_host(): assert ints.dtype == np.float64 and isinstance(ints, np.ndarray) +@pytest.mark.parametrize( + ("name", "make"), + [ + ("column prefix", lambda: np.zeros((3, 10))[:, :6]), + ("every other entry", lambda: np.zeros(12)[::2]), + ("every other row", lambda: np.zeros((6, 4))[::2]), + ("every other column", lambda: np.zeros((3, 8))[:, ::2]), + ("one row of a prefix", lambda: np.zeros((3, 10))[1:2, :6]), + ], +) +def test_strided_takes_c_ordered_arrays_with_gaps(name, make): + grid = np.zeros(3) + view = make() + assert not view.flags.c_contiguous or name == "one row of a prefix" + assert ( + xp.kernels.as_kernel_array(view, like=grid, dtype=float, strided=True) is view + ) + converted = xp.kernels.as_kernel_array(view, like=grid, dtype=float) + assert converted.flags.c_contiguous # the default still copies + with xp.kernels.kernel_output(view, like=grid, dtype=float, strided=True) as out: + assert out is view + + +@pytest.mark.parametrize( + ("name", "make"), + [ + ("F order", lambda: np.asfortranarray(np.zeros((3, 4)))), + ("transposed", lambda: np.zeros((4, 3)).T), + ("negative stride", lambda: np.zeros(5)[::-1]), + ("negative row stride", lambda: np.zeros((3, 4))[::-1]), + ("other dtype", lambda: np.zeros((3, 10), dtype=np.float32)[:, :6]), + ], +) +def test_strided_still_copies_other_layouts(name, make): + grid = np.zeros(3) + value = make() + converted = xp.kernels.as_kernel_array(value, like=grid, dtype=float, strided=True) + assert converted is not value + assert converted.flags.c_contiguous and converted.dtype == np.float64 + np.testing.assert_array_equal(converted, value) + + +@pytest.mark.parametrize("strided", [False, True]) +def test_single_row_from_fancy_indexing_gets_c_order_strides(strided): + """``a[:, order]`` of a ``(1, n)`` array has strides (8, 8), which Pyccel aborts on.""" + grid = np.zeros(3) + row = np.arange(5.0)[None, :][:, [4, 3, 2, 1, 0]] + assert row.strides == (8, 8) and row.flags.c_contiguous + converted = xp.kernels.as_kernel_array(row, like=grid, dtype=float, strided=strided) + assert converted.strides == (40, 8) + assert np.shares_memory(converted, row) # a view, not a copy + with xp.kernels.kernel_output( + row, like=grid, dtype=float, strided=strided + ) as buffer: + assert buffer.strides == (40, 8) + buffer[0, 0] = -1.0 + assert row[0, 0] == -1.0 + + +def test_strided_without_dtype_keeps_any_dtype(): + view = np.zeros((2, 8), dtype=np.int32)[:, :5] + assert xp.kernels.as_kernel_array(view, like=np.zeros(1), strided=True) is view + + def test_kernel_output_writes_into_its_target(): grid = np.zeros(3) out = np.zeros(4) diff --git a/tests/unit/test_kernel_outputs.py b/tests/unit/test_kernel_outputs.py new file mode 100644 index 0000000..358f0ac --- /dev/null +++ b/tests/unit/test_kernel_outputs.py @@ -0,0 +1,156 @@ +"""Kernel outputs by name, from annotations and from a kernel folder; struct fields in tests.""" + +import importlib +import sys +import textwrap +from typing import Final + +import numpy as np +import pytest + +from cunumpy.kernel_testing import _collect_arrays +from cunumpy.kernels import Kernel, PyccelKernel, outputs_from_annotations + + +def solve(a, b, out): + out[:] = a + b + + +def host_arrays(kernel, args, kwargs=None): + """The ids of the arrays the declared outputs of `kernel` reach.""" + return kernel._output_host_arrays(list(args), dict(kwargs or {})) + + +def test_output_name_finds_a_positional_argument(): + a, b, out = np.zeros(2), np.zeros(2), np.zeros(2) + kernel = PyccelKernel(solve, outputs=("out",)) + assert host_arrays(kernel, (a, b, out)) == {id(out)} + assert host_arrays(kernel, (a, b), {"out": out}) == {id(out)} + + +def test_output_index_finds_a_keyword_argument(): + a, b, out = np.zeros(2), np.zeros(2), np.zeros(2) + kernel = PyccelKernel(solve, outputs=(2,)) + assert host_arrays(kernel, (a, b, out)) == {id(out)} + assert host_arrays(kernel, (a, b), {"out": out}) == {id(out)} + + +def test_unknown_parameter_names_keep_the_old_rules(): + def hidden(*args): # no named parameters, like a compiled kernel + pass + + a, out = np.zeros(2), np.zeros(2) + with pytest.raises(KeyError, match="declared by index"): + host_arrays(PyccelKernel(hidden, outputs=("out",)), (a, out)) + with pytest.raises(IndexError, match="declared by name"): + host_arrays(PyccelKernel(hidden, outputs=(1,)), (a,), {"out": out}) + named = PyccelKernel(hidden, outputs=("out",), parameters=["a", "out"]) + assert host_arrays(named, (a, out)) == {id(out)} + + +def test_unknown_name_is_reported(): + kernel = PyccelKernel(solve, outputs=("missing",)) + with pytest.raises(KeyError, match="no argument of that name"): + host_arrays(kernel, (np.zeros(1),) * 3) + + +def test_outputs_from_annotations(): + def push( + markers: "float[:, :]", + weights: "Final[float[:]]", + dt: "float", + n: int, + args: "MarkerArgs", # noqa: F821 + scratch: Final[np.ndarray], + ): + pass + + assert outputs_from_annotations(push) == ("markers", "args") + + def unannotated(a, b): + pass + + assert outputs_from_annotations(unannotated) is None + assert outputs_from_annotations(len) is None # no signature to read + + +@pytest.fixture +def folder_kernel(tmp_path, monkeypatch): + folder = tmp_path / "output_pkg" / "scale" + folder.mkdir(parents=True) + (tmp_path / "output_pkg" / "__init__.py").write_text("") + (folder / "__init__.py").write_text("") + (folder / "scale_kernels.py").write_text( + textwrap.dedent(""" + def scale(x: "float[:]", factor: "Final[float[:]]", a: "float", n: "int"): + for i in range(n): + x[i] *= factor[i] * a + """), + ) + monkeypatch.syspath_prepend(str(tmp_path)) + yield "output_pkg.scale" + for module in [m for m in sys.modules if m.startswith("output_pkg")]: + del sys.modules[module] + importlib.invalidate_caches() + + +def test_from_folder_reads_outputs_from_annotations(folder_kernel): + kernel = Kernel.from_folder(folder_kernel, outputs="annotations") + assert kernel.host_kernel.outputs == ("x",) + x, factor = np.ones(3), np.full(3, 2.0) + assert host_arrays(kernel.host_kernel, (x, factor, 3.0, 3)) == {id(x)} + + +def test_from_folder_outputs_by_name_use_the_host_parameters(folder_kernel): + kernel = Kernel.from_folder(folder_kernel, outputs=("x",)) + x, factor = np.ones(3), np.full(3, 2.0) + assert host_arrays(kernel.host_kernel, (x, factor, 3.0, 3)) == {id(x)} + assert host_arrays(kernel.host_kernel, (), {"x": x}) == {id(x)} + + +def test_host_options_take_precedence_over_outputs(folder_kernel): + kernel = Kernel.from_folder( + folder_kernel, outputs="annotations", host_options={"outputs": ("factor",)} + ) + assert kernel.host_kernel.outputs == ("factor",) + + +def test_from_folder_rejects_an_unknown_outputs_string(folder_kernel): + with pytest.raises(ValueError, match="annotations"): + Kernel.from_folder(folder_kernel, outputs="all") + + +# --- comparing only some fields of a struct argument ------------------------- + + +class Markers: + def __init__(self): + self.positions = np.zeros((3, 2)) + self.buffer = np.ones(5) + self.n = 3 + + +def test_collect_arrays_by_parameter_name_and_field(): + markers, weights = Markers(), np.zeros(3) + args = (markers, weights) + names = ["markers", "weights"] + + both = _collect_arrays(args, ("markers",), names) + assert set(both) == {"argument 0.positions", "argument 0.buffer"} + + only = _collect_arrays(args, ("markers.positions", "weights"), names) + assert set(only) == {"argument 0.positions", "argument 1"} + assert only["argument 0.positions"] is markers.positions + + assert set(_collect_arrays(args, ("0.buffer",), names)) == {"argument 0.buffer"} + assert set(_collect_arrays(args, (-1,), names)) == {"argument 1"} + + +def test_collect_arrays_field_and_name_errors(): + args = (Markers(),) + with pytest.raises(KeyError, match="no array field 'velocities'"): + _collect_arrays(args, ("markers.velocities",), ["markers"]) + with pytest.raises(KeyError, match="not a parameter"): + _collect_arrays(args, ("particles",), ["markers"]) + with pytest.raises(KeyError, match="names are unknown|unknown"): + _collect_arrays(args, ("markers",), None) diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index ed01e9c..dceb1f7 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -127,8 +127,10 @@ def __init__(self, x, values): with pytest.raises(IndexError, match="output argument 6 does not exist"): _collect_arrays(args, outputs=(6,)) - with pytest.raises(TypeError, match="positional argument indices"): + with pytest.raises(KeyError, match="not a parameter of the kernel"): _collect_arrays(args, outputs=("out",)) + with pytest.raises(TypeError, match="indices \\(int\\) or names"): + _collect_arrays(args, outputs=(1.5,)) def test_compare_results(): diff --git a/tests/unit/test_metal_kernel.py b/tests/unit/test_metal_kernel.py new file mode 100644 index 0000000..6d92e07 --- /dev/null +++ b/tests/unit/test_metal_kernel.py @@ -0,0 +1,145 @@ +"""MetalKernel: argument checking everywhere, launches on Apple silicon with MLX.""" + +import numpy as np +import pytest + +import cunumpy as xp + +MetalKernel = xp.kernels.MetalKernel +needs_metal = pytest.mark.skipif( + not xp.kernels.metal_available(), reason="needs MLX and a Metal GPU" +) + +AXPY = "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i] + b[i];" + + +def axpy(**kwargs): + return MetalKernel(AXPY, inputs=["x", "a", "b"], outputs=["y"], **kwargs) + + +def test_constructor_validates_names_and_options(): + with pytest.raises(ValueError, match="at least one output"): + MetalKernel("", inputs=["x"], outputs=[]) + with pytest.raises(ValueError, match="distinct"): + MetalKernel("", inputs=["x"], outputs=["x"]) + with pytest.raises(ValueError, match="float64"): + MetalKernel("", inputs=["x"], outputs=["y"], float64="ignore") + with pytest.raises(ValueError, match="threadgroup"): + MetalKernel("", inputs=["x"], outputs=["y"], threadgroup=(1, 2, 3, 4)) + + +def test_missing_mlx_or_gpu_gives_a_clear_error(monkeypatch): + import cunumpy._metal_kernel as module + + def unavailable(): + raise ImportError("no mlx") + + monkeypatch.setattr(module, "_mlx", unavailable) + assert not xp.kernels.metal_available() + with pytest.raises(ImportError): + axpy()(np.zeros(1, np.float32), 1.0, np.zeros(1, np.float32), out=np.zeros(1)) + + +@needs_metal +def test_wrong_argument_counts_and_outputs_are_rejected(): + kernel = axpy() + x = np.zeros(4, np.float32) + with pytest.raises(TypeError, match="takes 3 input"): + kernel(x, out=x.copy()) + with pytest.raises(TypeError, match="out must be 1 NumPy"): + kernel(x, 1.0, x, out=[x, x]) + with pytest.raises(TypeError, match="out must be 1 NumPy"): + kernel(x, 1.0, x, out=None) + + +@needs_metal +def test_float64_is_rejected_unless_cast(): + x = np.arange(4.0) + with pytest.raises(TypeError, match="does not support"): + axpy()(x.astype(np.float32), 2.0, x, out=np.empty(4, np.float32)) + with pytest.raises(TypeError, match="output 'y'"): + axpy()(x.astype(np.float32), 2.0, x.astype(np.float32), out=np.empty(4)) + + +@needs_metal +def test_axpy_with_python_scalar_and_non_contiguous_input(): + x = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::2] + assert not x.flags.c_contiguous + b = np.ones(x.shape, np.float32) + y = np.empty(x.shape, np.float32) + result = axpy()(x, 2.0, b, out=y, n_threads=x.size) + assert result is y + np.testing.assert_allclose(y, 2.0 * x + 1.0) + + +@needs_metal +def test_float64_cast_computes_in_float32_and_fills_float64_output(): + x = np.linspace(0, 1, 100) + y = np.empty(100) + axpy(float64="cast")(x, 3.0, x, out=y) + np.testing.assert_allclose(y, 4.0 * x, rtol=1e-6) + assert y.dtype == np.float64 + + +@needs_metal +def test_several_outputs_template_and_init_value(): + kernel = MetalKernel( + "uint i = thread_position_in_grid.x;" + "if (i < N) { s[i] = x[i] + 1; d[i] = x[i] - 1; }", + inputs=["x"], + outputs=["s", "d"], + init_value=-99.0, + ) + x = np.arange(4, dtype=np.float32) + s, d = np.empty(4, np.float32), np.empty(4, np.float32) + out = kernel(x, out=(s, d), n_threads=4, template={"N": 3}) + assert out == (s, d) + np.testing.assert_array_equal(s, [1, 2, 3, -99]) + np.testing.assert_array_equal(d, [-1, 0, 1, -99]) + + +@needs_metal +def test_push_matches_float64_reference(): + n, steps, dt, B = 1000, 50, 0.01, 1.5 + c, s = np.cos(dt * B), np.sin(dt * B) + rng = np.random.default_rng(0) + pos = rng.random((n, 3)).astype(np.float32) + vel = rng.standard_normal((n, 3)).astype(np.float32) + kernel = MetalKernel( + """ + uint i = thread_position_in_grid.x; + float px = pos[3*i], py = pos[3*i+1], pz = pos[3*i+2]; + float vx = vel[3*i], vy = vel[3*i+1], vz = vel[3*i+2]; + for (int k = 0; k < NSTEPS; ++k) { + px += vx*p[2]; py += vy*p[2]; pz += vz*p[2]; + float nx = p[0]*vx + p[1]*vy; vy = -p[1]*vx + p[0]*vy; vx = nx; + } + pos_out[3*i] = px; pos_out[3*i+1] = py; pos_out[3*i+2] = pz; + vel_out[3*i] = vx; vel_out[3*i+1] = vy; vel_out[3*i+2] = vz; + """, + inputs=["pos", "vel", "p"], + outputs=["pos_out", "vel_out"], + ) + pos_out, vel_out = np.empty_like(pos), np.empty_like(vel) + kernel( + pos, + vel, + np.array([c, s, dt], np.float32), + out=(pos_out, vel_out), + n_threads=n, + template={"NSTEPS": steps}, + ) + p, v = pos.astype(float), vel.astype(float) + for _ in range(steps): + p += v * dt + v[:, :2] = np.stack([c * v[:, 0] + s * v[:, 1], -s * v[:, 0] + c * v[:, 1]], 1) + np.testing.assert_allclose(pos_out, p, atol=1e-4) + np.testing.assert_allclose(vel_out, v, atol=1e-4) + + +@needs_metal +def test_copies_are_counted(): + x = np.zeros(8, np.float32) + with xp.profiling.count_transfers() as counter: + axpy()(x, 1.0, x, out=np.empty(8, np.float32)) + assert counter.to_device == 3 and counter.to_host == 1 diff --git a/tests/unit/test_morton.py b/tests/unit/test_morton.py index bb8282e..3c1bb8a 100644 --- a/tests/unit/test_morton.py +++ b/tests/unit/test_morton.py @@ -131,8 +131,44 @@ def test_sort_by_key_is_stable_and_reorders_all_arrays(): def test_sort_by_key_validates_shapes(): with pytest.raises(ValueError, match="1D"): xp.algorithms.sort_by_key(np.zeros((2, 2))) - with pytest.raises(ValueError, match="3 rows"): + with pytest.raises(ValueError, match="expected 3 entries"): xp.algorithms.sort_by_key(np.zeros(3), np.zeros(4)) + with pytest.raises(ValueError, match="no axis 1"): + xp.algorithms.sort_by_key(np.zeros(3), np.zeros(3), axis=1) + + +@pytest.mark.parametrize( + ("dtype", "low", "high"), + [ + (np.int64, 0, 300), # one 16-bit pass, many equal keys + (np.int64, 0, 400_000), # two passes: cells of a 3D grid + (np.int64, -(2**40), 2**40), # negative keys, three passes + (np.int32, -5, 70_000), + (np.uint64, 0, 2**63), # Morton keys, four passes + ], +) +def test_radix_sort_of_integer_keys_matches_the_stable_argsort(dtype, low, high): + keys = np.random.default_rng(5).integers(low, high, 50_000, dtype=dtype) + _, order, sorted_ids = xp.algorithms.sort_by_key(keys, np.arange(keys.size)) + expected = np.argsort(keys, kind="stable") + assert order.dtype == np.int64 + np.testing.assert_array_equal(order, expected) + np.testing.assert_array_equal(sorted_ids, expected) + + +def test_sort_by_key_along_the_last_axis_of_component_major_arrays(): + keys = np.array([2, 0, 1, 0], dtype=np.int64) + positions = np.arange(12.0).reshape(3, 4) + weights = np.array([10.0, 11.0, 12.0, 13.0]) + _, order, sorted_positions, sorted_weights = xp.algorithms.sort_by_key( + keys, + positions, + weights, + axis=-1, + ) + assert order.tolist() == [1, 3, 2, 0] + np.testing.assert_array_equal(sorted_positions, positions[:, order]) + np.testing.assert_array_equal(sorted_weights, weights[order]) def test_header_is_shipped(): diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index c6ac552..baaef36 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -532,7 +532,10 @@ def test_require_version(monkeypatch): assert isinstance(buf, np.ndarray) and buf.tolist() == [1.0, 2.0] buf[:] = [5.0, 6.0] assert xp.to_numpy(d).tolist() == [5.0, 6.0] -assert sorted(e.kind for e in counter.events) == ["to_device", "to_host"] +assert sorted(e.kind for e in counter.events if e.kind != "sync") == [ + "to_device", + "to_host", +] with xp.mpi.mpi_buffer(d, cuda_aware=True) as buf: assert buf is d producer = xp.cuda.create_stream() @@ -550,6 +553,7 @@ def test_require_version(monkeypatch): def test_fake_cupy_in_subprocess(): root = Path(__file__).resolve().parents[2] env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", CUNUMPY_BACKEND="cupy") + env.pop("CUNUMPY_REQUIRE_CUDA", None) # the fake CuPy skips by design env["PYTHONPATH"] = os.pathsep.join( p for p in (str(root / "src"), env.get("PYTHONPATH", "")) if p ) diff --git a/tests/unit/test_pyccel_kernel.py b/tests/unit/test_pyccel_kernel.py index 43113b8..bf9b92a 100644 --- a/tests/unit/test_pyccel_kernel.py +++ b/tests/unit/test_pyccel_kernel.py @@ -6,12 +6,14 @@ run everywhere; * end-to-end tests, which compile `pyccel_kernels.py` with pyccel and drive the compiled kernels through `PyccelKernel`. They are skipped when pyccel (or a - working compiler) is unavailable. + working compiler) is unavailable, unless `CUNUMPY_REQUIRE_PYCCEL=1` makes + missing dependencies and compilation failures errors in compiled-test CI. -The CuPy-side assertions can only be exercised on a machine with a GPU; the -NumPy-side assertions run everywhere. +The CuPy-side assertions run on a GPU or with `CUNUMPY_FAKE_CUPY=1`; the +NumPy-side assertions run everywhere with the compiled-test dependencies. """ +import os import shutil from pathlib import Path @@ -36,7 +38,15 @@ def kernels(tmp_path_factory): The source is copied into a temporary directory first so that pyccel's build artefacts (`__pyccel__/`) never land in the repository. """ - pyccel = pytest.importorskip("pyccel", reason="pyccel is not installed") + require_pyccel = os.environ.get("CUNUMPY_REQUIRE_PYCCEL", "").lower() in { + "1", + "true", + "yes", + } + if require_pyccel: + import pyccel + else: + pyccel = pytest.importorskip("pyccel", reason="pyccel is not installed") import importlib.util import sys @@ -52,7 +62,9 @@ def kernels(tmp_path_factory): try: return pyccel.epyccel(module, language="c") - except Exception as exc: # noqa: BLE001 - no compiler / broken toolchain + except Exception as exc: + if require_pyccel: + raise pytest.skip(f"pyccel could not compile the example kernels: {exc}") finally: sys.modules.pop(source.stem, None) @@ -423,7 +435,7 @@ def test_misdeclared_output_index_raises(): def test_misdeclared_output_name_raises(): wrapped = PyccelKernel(lambda out: None, use_cupy=True, outputs=("nope",)) - with pytest.raises(KeyError, match="no such keyword argument"): + with pytest.raises(KeyError, match="no argument of that name"): wrapped(np.zeros(2)) @@ -449,6 +461,30 @@ def test_outputs_is_ignored_on_the_numpy_path(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +@pytest.mark.parametrize("column_block", [False, True]) +def test_compiled_fortran_order_and_copyback(kernels, backend, column_block): + if backend == "cupy": + _skip_without_cupy() + + expected = np.array(np.arange(30.0).reshape(5, 6), order="F") + with xp.use_backend(backend, strict=True): + storage = xp.array(expected, order="F", copy=True) + target = storage[:, 1:4] if column_block else storage + assert target.flags.f_contiguous and not target.flags.c_contiguous + + PyccelKernel(kernels.scale_fortran_inplace, outputs=(0,))(target, 3.0) + + if column_block: + expected[:, 1:4] *= 3.0 + else: + expected *= 3.0 + # Check the parent as well: copy-back must update the view without + # changing the adjacent columns or replacing the caller's storage. + np.testing.assert_array_equal(xp.to_numpy(storage), expected) + assert storage.flags.f_contiguous and target.flags.f_contiguous + + @pytest.mark.parametrize("backend", ["numpy", "cupy"]) def test_compiled_axpy(kernels, backend): if backend == "cupy": diff --git a/tests/unit/test_require_cuda.py b/tests/unit/test_require_cuda.py new file mode 100644 index 0000000..8325267 --- /dev/null +++ b/tests/unit/test_require_cuda.py @@ -0,0 +1,66 @@ +"""With CUNUMPY_REQUIRE_CUDA=1 the GPU markers of kernel_testing fail (as errors) instead of skipping.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +import cunumpy as xp + +pytestmark = pytest.mark.skipif( + xp.cupy_available(), reason="checks the behavior without a GPU" +) + +TESTS = """ +import pytest +from cunumpy.kernel_testing import backend, requires_cupy + + +@requires_cupy +def test_marked(): + pass + + +def test_backend(backend): + pass +""" + + +def run(tmp_path, require): + (tmp_path / "test_gpu.py").write_text(TESTS) + root = Path(__file__).resolve().parents[2] + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(root / "src"), env.get("PYTHONPATH", "")) if p + ) + env.pop("CUNUMPY_REQUIRE_CUDA", None) + env.pop("CUNUMPY_FAKE_CUPY", None) + env.pop("CUNUMPY_BACKEND", None) + if require: + env["CUNUMPY_REQUIRE_CUDA"] = "1" + return subprocess.run( + [sys.executable, "-m", "pytest", "-p", "no:cacheprovider", "-rA", "-q"], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + check=False, + ) + + +def test_markers_skip_without_a_gpu(tmp_path): + result = run(tmp_path, require=False) + assert result.returncode == 0, result.stdout + assert "SKIPPED" in result.stdout and "ERROR" not in result.stdout + + +def test_markers_error_when_a_gpu_is_required(tmp_path): + result = run(tmp_path, require=True) + assert result.returncode != 0, result.stdout + assert "CUNUMPY_REQUIRE_CUDA is set" in result.stdout + # the condition of the marker raises while the test is set up: an error + assert "ERROR test_gpu.py::test_marked" in result.stdout + assert "ERROR test_gpu.py::test_backend[cupy]" in result.stdout + assert "PASSED test_gpu.py::test_backend[numpy]" in result.stdout diff --git a/tests/unit/test_staging.py b/tests/unit/test_staging.py index 78f4ff3..cc0e613 100644 --- a/tests/unit/test_staging.py +++ b/tests/unit/test_staging.py @@ -320,3 +320,46 @@ def test_staged_result_restores_callers_device_on_gpu(): assert cp.cuda.runtime.getDevice() == 1 with pytest.raises(ValueError, match="bound to another"): staging.copy(cp.zeros(3)) + + +def test_to_host_async_of_a_host_array_is_ready_at_once(): + a = np.array(2.5) + copy = xp.to_host_async(a) + a[...] = 0.0 + assert copy.ready() + assert copy.result() == 2.5 + assert isinstance(copy.result(), np.floating) + with xp.profiling.count_transfers() as counter: + xp.to_host_async(np.zeros(3)) + assert counter.events == [] + + +def test_to_host_async_copies_on_its_own_stream_without_waiting( + fake_device, monkeypatch +): + from cunumpy import _fake_cupy + + log = fake_device + monkeypatch.setattr(_fake_cupy, "is_active", lambda: False) + monkeypatch.setattr(staging_module, "_COPY_STREAMS", {}) + staging_module._cupy().ascontiguousarray = lambda a: a + norm = np.array(4.0).view(DeviceArray) + with xp.profiling.count_transfers() as counter: + copy = xp.to_host_async(norm) + assert log == [ + ("record", "compute#1"), # after the kernels queued so far + ("wait", "staging", "compute#1"), + ("get", "staging"), + ("record", "staging#1"), + ] + assert not copy.ready() + assert counter.syncs == 0 + assert copy.result() == 4.0 # waits: counted as a sync + assert copy.ready() + assert copy.result() == 4.0 + (to_host,) = [e for e in counter.events if e.kind == "to_host"] + assert to_host.nbytes == 8 and not to_host.blocking + assert "[async]" in counter.report() + assert counter.syncs == 1 + xp.to_host_async(norm) + assert log.count(("wait", "staging", "compute#2")) == 1 # the stream is reused diff --git a/tests/unit/test_sync_counting.py b/tests/unit/test_sync_counting.py new file mode 100644 index 0000000..d0db679 --- /dev/null +++ b/tests/unit/test_sync_counting.py @@ -0,0 +1,76 @@ +"""count_transfers counts syncs: synchronize(), and scalar reads on the fake CuPy.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +import cunumpy as xp +from cunumpy._transfers import TransferCounter, _record_sync + + +def test_sync_events_are_listed_but_not_in_the_total(): + with xp.profiling.count_transfers() as counter: + _record_sync("something") + assert counter.syncs == 1 and counter.total == 0 + assert "sync (1)" in counter.report() + + +def test_assert_no_transfers_accepts_syncs_unless_asked(): + with xp.profiling.assert_no_transfers(): + _record_sync("something") + with ( + pytest.raises(AssertionError, match="sync"), + xp.profiling.assert_no_transfers(syncs=True), + ): + _record_sync("something") + + +def test_numpy_backend_synchronize_is_not_a_sync(): + with xp.profiling.count_transfers() as counter: + xp.synchronize() + assert counter.syncs == 0 + assert isinstance(counter, TransferCounter) + + +SCRIPT = r""" +import cunumpy as xp + +a = xp.asarray([1.0, 2.0, 3.0]) +with xp.profiling.count_transfers() as counter: + float(a[0]); int(a[1]); bool(a[2]); a.sum().item(); a.tolist() + xp.synchronize() + b = a * 2 # device work: no sync +assert counter.syncs == 6, counter.report() +assert counter.total == 0 +with xp.profiling.assert_no_transfers(): + float(a[0]) +try: + with xp.profiling.assert_no_transfers(syncs=True): + float(a[0]) +except AssertionError: + pass +else: + raise AssertionError("syncs=True must reject a scalar read") +print("sync counting OK") +""" + + +def test_scalar_reads_and_synchronize_are_counted_on_the_fake_cupy(): + root = Path(__file__).resolve().parents[2] + env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", CUNUMPY_BACKEND="cupy") + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(root / "src"), env.get("PYTHONPATH", "")) if p + ) + result = subprocess.run( + [sys.executable, "-c", SCRIPT], + env=env, + capture_output=True, + text=True, + check=False, + cwd=str(root), + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "sync counting OK" in result.stdout diff --git a/tests/unit/test_transfer_budget.py b/tests/unit/test_transfer_budget.py new file mode 100644 index 0000000..1ce7870 --- /dev/null +++ b/tests/unit/test_transfer_budget.py @@ -0,0 +1,176 @@ +"""Tests for `xp.profiling.TransferBudget` and nesting of `count_transfers`.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from cunumpy._transfers import _record, _record_sync +from cunumpy.profiling import TransferBudget, TransferCounter, count_transfers + + +def download(nbytes=8, blocking=True): + _record("to_host", f"download({nbytes})", nbytes=nbytes, blocking=blocking) + + +def test_count_transfers_into_a_counter_and_as_a_decorator(): + counter = TransferCounter() + + @count_transfers(counter) + def step(): + download() + + step() + step() + assert counter.to_host == 2 + with count_transfers(counter), count_transfers(counter): # nested: counted once + download() + assert counter.to_host == 3 + with count_transfers() as outer, count_transfers() as inner: + download() + assert outer.to_host == inner.to_host == 1 + + +def test_phases_accumulate_and_count_calls(): + budget = TransferBudget() + + @budget.count("integrate") + def integrate(): + download() + + for _ in range(3): + integrate() + with budget.phase("output") as output: + download(800) + assert budget["integrate"].to_host == 3 + assert budget.calls == {"integrate": 3, "output": 1} + assert output.bytes_to_host == 800 + assert "integrate (3 call(s))" in budget.report() + + +def test_counting_starts_with_start(): + budget = TransferBudget(started=False) + with budget.phase("integrate"): + download() # setup + budget.start() + with budget.phase("integrate"): + download() + budget.stop() + with budget.phase("integrate"): + download() + assert budget["integrate"].to_host == 1 + assert budget.calls == {"integrate": 1} + + +def test_nested_phases_count_each_event_once(): + budget = TransferBudget() + + @budget.count("integrate") + def integrate(depth): + download() + if depth: + integrate(depth - 1) # recursion: the same phase + with budget.phase("diagnostics"): + download(16) + + integrate(1) + assert budget["integrate"].to_host == 2 + assert budget["diagnostics"].to_host == 1 # not also in "integrate" + assert budget.calls == {"integrate": 1, "diagnostics": 1} + + +def test_rules_that_pass(): + budget = TransferBudget() + for _ in range(2): + with budget.phase("integrate"): + download(8) + _record("device_copy", "astype", nbytes=80) + _record_sync("synchronize()") + with budget.phase("output"): + download(400) + download(400) + budget.require("integrate", allow={"to_host": {"max_nbytes": 8}}, calls=2) + budget.require( + "output", allow={"to_host": {"max_count": 2, "max_total_bytes": 800}} + ) + budget.require("never_run", calls=0) + assert budget.violations() == [] + budget.check() + + +def test_rules_that_fail_report_the_events_and_where(): + budget = TransferBudget() + with budget.phase("integrate"): + download(80) + _record("fallback", "Kernel 'push' ran its host kernel") + with budget.phase("output"): + for _ in range(3): + download(400) + budget.require("integrate", allow={"to_host": {"max_nbytes": 8}}, calls=2) + budget.require( + "output", allow={"to_host": {"max_count": 2, "max_total_bytes": 1000}} + ) + with pytest.raises(AssertionError) as info: + budget.check() + report = str(info.value) + assert "integrate: 1 call(s), expected 2" in report + assert "80 bytes > max_nbytes=8" in report + assert "not allowed" in report and "fallback" in report + assert "3 to_host > max_count=2" in report + assert "1200 bytes of to_host > max_total_bytes=1000" in report + assert f"{Path(__file__).name}:" in report # where it happened + + +def test_blocking_and_implicit_events_are_told_apart(): + budget = TransferBudget() + with budget.phase("solve"): + download(8, blocking=False) + _record_sync("float(a)", implicit=True) + budget.require( + "solve", + allow={"to_host": {"blocking": False}, "sync": {"implicit": False}}, + ) + found = budget.violations() + assert len(found) == 1 and "implicit=True not allowed" in found[0] + with budget.phase("solve"): + download(8) # blocking + assert any("blocking=True not allowed" in v for v in budget.violations()) + + +def test_unknown_kinds_and_limits_are_refused(): + budget = TransferBudget() + with pytest.raises(ValueError, match="unknown transfer kind"): + budget.require("x", allow={"to_hots": None}) + with pytest.raises(ValueError, match="unknown limits"): + budget.require("x", allow={"to_host": {"max_bytes": 8}}) + + +def test_fake_cupy_scalar_reads_are_implicit_syncs(): + code = """ +import cunumpy as xp + +with xp.use_backend("cupy"), xp.profiling.count_transfers() as counter: + a = xp.ones(3) + float(a.sum()) + xp.synchronize() + copy = xp.to_host_async(a.sum()) + assert copy.ready() and copy.result() == 3.0 +syncs = [e for e in counter.events if e.kind == "sync"] +assert [e.implicit for e in syncs] == [True, False], syncs +(event,) = [e for e in counter.events if e.kind == "to_host"] +assert not event.blocking and event.nbytes == 8 +print("ok") +""" + root = Path(__file__).resolve().parents[2] + env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", PYTHONPATH=str(root / "src")) + result = subprocess.run( + [sys.executable, "-c", code], + env=env, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "ok" in result.stdout diff --git a/tests/unit/test_transfers.py b/tests/unit/test_transfers.py index 773cc11..f72493e 100644 --- a/tests/unit/test_transfers.py +++ b/tests/unit/test_transfers.py @@ -42,8 +42,8 @@ def shape(self): def dtype(self): return self.data.dtype - def get(self): - return self.data.copy() + def get(self, order="C"): + return self.data.copy(order=order) def __setitem__(self, key, value): self.data[key] = value.data if isinstance(value, _FakeDeviceArray) else value @@ -84,7 +84,7 @@ def test_empty_counter(): assert counter.events == [] and counter.kernel_conversion_calls == [] assert counter.report().startswith("0 transfer(s) through cunumpy") assert repr(counter) == ( - "TransferCounter(to_host=0, to_device=0, kernel_conversion=0, fallback=0, device_copy=0)" + "TransferCounter(to_host=0, to_device=0, kernel_conversion=0, fallback=0, device_copy=0, sync=0)" ) @@ -189,7 +189,7 @@ def test_report_groups_events_by_kind_and_call_site(fake_device): lines = report.splitlines() assert lines[0] == ( "4 transfer(s) through cunumpy " - "(3 to_host, 1 to_device, 0 kernel_conversion, 0 fallback, 0 device_copy)" + "(3 to_host, 1 to_device, 0 kernel_conversion, 0 fallback, 0 device_copy, 0 sync)" ) assert " to_host (3):" in lines assert " to_device (1):" in lines