From bdf89df94be4c04efd733a0b5ec6431b192257b5 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:18:44 +0200 Subject: [PATCH 1/9] - Reusable streams/events and MPI staging buffers + more --- CHANGELOG.md | 17 ++ README.md | 10 +- docs/source/api.md | 128 +++++++-- docs/source/guides/execution-helpers.md | 162 ++++++++++++ docs/source/guides/mpi.md | 5 +- docs/source/guides/solvers.md | 5 +- docs/source/index.md | 1 + src/cunumpy/LLM_GUIDE.md | 36 +++ src/cunumpy/__init__.py | 2 + src/cunumpy/__init__.pyi | 5 +- src/cunumpy/_algorithms.py | 244 ++++++++++++++---- src/cunumpy/_device.py | 10 +- src/cunumpy/_fake_cupy_impl.py | 37 ++- src/cunumpy/_mpi.py | 137 ++++++++-- src/cunumpy/_streams.py | 97 +++++++ src/cunumpy/algorithms.py | 11 +- src/cunumpy/cuda/__init__.py | 21 +- src/cunumpy/cuda/include/cunumpy/atomic.cuh | 45 +++- src/cunumpy/cuda/include/cunumpy/reduce.cuh | 63 ++++- src/cunumpy/cuda/include/cunumpy/scan.cuh | 84 ++++++ src/cunumpy/mpi.py | 2 + src/cunumpy/xp.py | 80 +++++- tests/unit/test_backend_diagnostics.py | 81 ++++++ tests/unit/test_cuda_collectives.py | 271 ++++++++++++++++++++ tests/unit/test_mpi_staging.py | 194 ++++++++++++++ tests/unit/test_porting_helpers.py | 8 + tests/unit/test_reusable_streams.py | 92 +++++++ tests/unit/test_segment_plans.py | 159 ++++++++++++ 28 files changed, 1883 insertions(+), 124 deletions(-) create mode 100644 docs/source/guides/execution-helpers.md create mode 100644 src/cunumpy/_streams.py create mode 100644 src/cunumpy/cuda/include/cunumpy/scan.cuh create mode 100644 tests/unit/test_backend_diagnostics.py create mode 100644 tests/unit/test_cuda_collectives.py create mode 100644 tests/unit/test_mpi_staging.py create mode 100644 tests/unit/test_reusable_streams.py create mode 100644 tests/unit/test_segment_plans.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c1850d..96e35c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- Reusable `cuda.create_stream`, `create_event`, `record_event`, and `wait_event`, + with synchronous host equivalents; `cuda.stream(existing)` selects an existing stream. +- `mpi.MPIStaging` reuses host staging storage and rejects overlapping device uses. +- `algorithms.cell_offsets`, `segment_boundaries`, and `SegmentPlan` for prepared grouping. +- Strict backend selection via `set_backend`/`use_backend(..., strict=True)` and + JSON-compatible `backend_info()` diagnostics. +- `cunumpy/scan.cuh` inclusive/exclusive warp and block prefix sums, mask-aware + warp reductions, partial-warp block reductions, and integer atomic-add helpers. + +### Execution helpers +- `segment_sum(..., out=...)` supports arbitrary trailing component dimensions; + CUDA reduces components in one accumulation launch. Prepared plans avoid + repeated key validation and GPU scalar reads. Floating-point order may vary. +- MPI producer synchronization accepts a stream/event, or waits for all work on + the buffer's device. Receive staging waits for copy-back before releasing storage. + ### Changed (breaking, with deprecation) - The helpers moved from the top level of `cunumpy` to submodules, so that the top level is the NumPy/CuPy namespace plus backend selection and array conversion, and no helper hides a NumPy or CuPy name (`xp.fuse` hid `cupy.fuse`): `cunumpy.cuda` (CUDA only: `CudaKernel`, `CudaKernelVariants`, `CudaStruct*`, `CudaArguments`, `CudaParameter`, header tools, debug mode, device selection and memory, `stream`, `pin_memory`), `cunumpy.kernels` (`Kernel`, `KernelCatalog`, `PyccelKernel`, `KernelArguments`, `PyccelStructArguments`, host implementations, `as_kernel_array`, `kernel_output`, `fuse`), `cunumpy.rng` (`random_streams`, `RandomStreams`, `get_rng`, `philox_*`), `cunumpy.algorithms` (`morton_*`, `sort_by_key`, `segment_sum`), `cunumpy.mpi` (`mpi_buffer`, CUDA-aware MPI, `local_rank`, `synchronize_for_mpi`), `cunumpy.profiling` (`timed_region`, `Timing`, `nvtx_range`, transfer counting), `cunumpy.memory` (`HostStaging`, `StagedCopy`, `DeviceMirror`) and `cunumpy.petsc` (`petsc_vec`). All are imported by `import cunumpy`. The old top-level names still work and raise a `DeprecationWarning` naming the new place; they will be removed in 0.6. - `cunumpy.testing` is now `cunumpy.kernel_testing`. Once imported, `cunumpy.testing` replaced NumPy's `xp.testing`, so `xp.testing.assert_allclose` failed in every test that ran after an `import cunumpy.testing`. `cunumpy.testing` still works, with a `DeprecationWarning`, until 0.6. diff --git a/README.md b/README.md index 59bb42f..f1bdfe6 100644 --- a/README.md +++ b/README.md @@ -27,8 +27,8 @@ never hide a NumPy name: | `xp.cuda` | CUDA only: `CudaKernel`, `CudaStruct`, CUDA headers, devices, streams | | `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, host implementations, `fuse` | | `xp.rng` | `random_streams`, `get_rng`, `philox_*` | -| `xp.algorithms` | `morton_*`, `sort_by_key`, `segment_sum` | -| `xp.mpi` | `mpi_buffer`, CUDA-aware MPI | +| `xp.algorithms` | `morton_*`, `sort_by_key`, `cell_offsets`, `segment_boundaries`, `segment_sum`, `SegmentPlan` | +| `xp.mpi` | `mpi_buffer`, reusable `MPIStaging`, CUDA-aware MPI | | `xp.profiling` | `timed_region`, `nvtx_range`, `count_transfers` | | `xp.memory` | `HostStaging`, `DeviceMirror` | | `xp.petsc` | `petsc_vec` | @@ -72,6 +72,12 @@ but unavailable or not functional, CuNumpy falls back to NumPy. Always check `get_backend()` when the effective backend matters, such as when reporting configuration or deciding whether GPU-specific work will happen. +Use `xp.set_backend("cupy", strict=True)` to raise when CUDA is unavailable, +preserving the previous backend. `xp.backend_info()` returns structured backend, +dependency, and CUDA diagnostics. Reusable streams/events, MPI staging, cell +ranges, and prepared reductions are described in the +[execution helpers guide](docs/source/guides/execution-helpers.md). + Use `use_backend()` for a temporary selection. It restores the previous selection when the block exits, including when an exception is raised: diff --git a/docs/source/api.md b/docs/source/api.md index d6df722..cca787a 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -67,7 +67,7 @@ is checked when the version is unknown (not installed as a package). ## Backend selection -### `set_backend(backend)` +### `set_backend(backend, *, strict=False)` Selects the process-wide backend used for new NumPy-like operations. Supported values are `"numpy"` and `"cupy"`: @@ -82,6 +82,17 @@ otherwise CuNumpy falls back to NumPy. Check `get_backend()` to inspect the effective selection. Changing the selection does not move arrays that have already been created. +With `strict=True`, requesting unavailable CuPy raises `RuntimeError` with the +availability failure reason, preserving the previous backend and namespace. +`use_backend(backend, strict=True)` offers the same guarantee. + +### `backend_info()` + +Returns a JSON-compatible dictionary with the selected backend, cached CUDA +availability and failure reason, dependency versions, and current device/runtime +information when available. Inspection errors are included in the result. It +does not change the backend or initialize MPI. + The initial backend is NumPy unless `ARRAY_BACKEND=cupy` is set before CuNumpy is imported. Other values of this environment variable result in the NumPy default. @@ -210,14 +221,36 @@ 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. -### `algorithms.segment_sum(values, keys, n_segments)` +### `algorithms.segment_sum(values, keys, n_segments, *, out=None)` `out[k] = sum(values[i] for keys[i] == k)` on the backend of `keys`, with -`bincount` under the hood: the reduction step of a sort-then-reduce -accumulation. `values` has shape `(n,)` or `(n, m)` (columns summed -separately); a negative key drops the value; keys must be smaller than -`n_segments`. The result keeps a floating-point or complex dtype and is -`float64` otherwise. +one CUDA accumulation launch for all trailing components. `values` has shape +`(n, ...)`; a negative key drops the row; integer keys must be smaller than +`n_segments`. The result has shape `(n_segments, ...)`, keeps a floating-point +or complex dtype, and is `float64` otherwise. Optional `out` is overwritten, +must match shape/dtype, be writable and C-contiguous, and must not alias values. +Mixed backends raise. CUDA float16 accumulates in float32; floating-point atomic +accumulation order is not deterministic. + +### `algorithms.SegmentPlan(keys, n_segments)` + +Copies and validates keys once, then `plan.sum(values, out=None)` reuses them. +Changing the original keys does not change the plan. GPU preparation may +synchronize; repeated sums avoid key validation and per-column scalar reads. +A GPU plan requires its device current and values/output on that device. + +### `algorithms.cell_offsets(sorted_cells, n_cells)` + +Returns `n_cells+1` int64 offsets on the input backend. Cell k occupies the +half-open slice `[offsets[k], offsets[k+1])`. Missing cells have equal endpoints. +Input must be sorted integer IDs in `[0, n_cells)`; filter negative IDs first. + +### `algorithms.segment_boundaries(sorted_keys)` + +Returns `(unique_keys, starts, stops)` for runs of sorted integer keys. Supports +sparse, negative, and uint64 Morton keys. Starts/stops are int64 half-open indices. +Empty input returns three empty arrays. Validation and variable-length GPU +output may synchronize; prepare boundaries outside repeated operations. ### `algorithms.sort_by_key(keys, *arrays)` @@ -545,10 +578,11 @@ xp.mpi.synchronize_for_mpi(send, recv) # 4. before every MPI call with device b MPI.COMM_WORLD.Sendrecv(send, dest, recvbuf=recv, source=source) ``` -### `mpi.synchronize_for_mpi(*arrays)` +### `mpi.synchronize_for_mpi(*arrays, stream=None, event=None)` -Waits for the work pending on the current stream if at least one of `arrays` -is a CuPy array; `None` entries and host arrays are ignored, so it costs +Waits for all work on devices owning the CuPy arrays, or only the explicit +producer stream/event when supplied (pass at most one). The dependency must +cover all supplied device buffers. `None` entries and host arrays are ignored, so it costs nothing for host buffers and on the NumPy backend. Call it before every MPI call that sends or receives device buffers: CuPy launches kernels asynchronously and MPI knows nothing about CUDA streams, so a buffer that a @@ -563,7 +597,7 @@ comm.Sendrecv(send_buffer, dest, recvbuf=recv_buffer, source=source) No synchronization is needed after MPI returns: kernels launched afterwards see the received data. -### `mpi.mpi_buffer(array, *, send=True, recv=False, cuda_aware=None)` +### `mpi.mpi_buffer(array, *, send=True, recv=False, cuda_aware=None, staging=None, stream=None, event=None)` Context manager yielding the buffer to pass to MPI for `array`: a host array unchanged; a device array unchanged (after `synchronize_for_mpi`) when MPI is @@ -573,6 +607,23 @@ before the block (`send`) and copied back after it (`recv`), both counted by `mpi_is_cuda_aware()` or `set_mpi_cuda_aware()`; without one, a device array raises `RuntimeError`. +Pass `stream` or `event` for producer synchronization. Receive-only device +staging also waits for preceding device work. Receive copy-back completes before +the context releases host storage. With nonblocking MPI, call `request.Wait()` +inside the context before reading/releasing the buffer. + +### `mpi.MPIStaging(shape, dtype)` + +Caller-owned reusable host staging storage, pinned when available and allocated +lazily. Use `staging.buffer(array, **options)` or +`mpi_buffer(array, staging=staging, **options)`. Device staging is bound to a +shape, dtype, and device, and rejects overlapping uses. Host arrays still pass +through unchanged. Use separate staging instances for simultaneously active +send and receive buffers. + +Non-C-contiguous device arrays use a reusable device packing buffer; received +values are written back into the original view, preserving its storage. + ### `mpi.set_mpi_cuda_aware(value)`, `mpi.get_mpi_cuda_aware()` Record (or read) whether MPI can take device buffers, for `mpi_buffer()`. @@ -622,7 +673,7 @@ Waits for queued work on the current CUDA device to finish. This is useful before reading asynchronously computed results from host code. It is a no-op on NumPy. -### `cuda.stream()` +### `cuda.stream(existing=None)` Context manager that creates a non-blocking CuPy stream and yields it. Work issued in the block is enqueued on that stream. On NumPy, yields `None` and @@ -638,6 +689,24 @@ work_stream.synchronize() # on CuPy; the yielded value is None on NumPy Do not call methods on the yielded value without checking the backend. Use `xp.synchronize()` for code that should work on both backends. +Pass `existing` to select a previously allocated stream for the context. + +### `cuda.create_stream(non_blocking=True)`, `cuda.create_event(timing=False)` + +Create reusable CuPy streams and events. On CPU they return synchronous +`HostStream`/`HostEvent` objects supporting context selection, `.synchronize()`, +`.done`, recording and waits. Events disable timing by default. CUDA streams +belong to the device current at construction; keep that device current while +using them. + +### `cuda.record_event(event=None, *, stream=None)`, `cuda.wait_event(event, *, stream=None)` + +Record an existing or newly created completion event on the producer stream, +then enqueue a wait on the consumer stream without blocking the CPU. Omitting +the stream uses the current CUDA stream. All consumers of an event recording +must enqueue their waits before that event is re-recorded. See the +[execution helpers guide](guides/execution-helpers.md) for examples. + ## Profiling CUDA kernels run asynchronously: a wall-clock timer around a launch measures @@ -1939,7 +2008,9 @@ double cunumpy_atomic_add_3d(double* data, long long n1, long long n2, long long i, long long j, long long k, double v); ``` -The indexed helpers (also for `float`) address C-contiguous arrays of shape +The scalar and indexed helpers also support `int`, `unsigned int`, `long long`, +and `unsigned long long`. Signed 64-bit addition uses CUDA's unsigned 64-bit +atomic and modulo-2^64 arithmetic. The indexed helpers address C-contiguous arrays of shape `(n0, n1)` and `(n0, n1, n2)`. They wrap `atomicAdd`, a hardware instruction for `double` from compute capability 6.0 (sm_60) on; older devices use a compare-and-swap loop. @@ -2022,7 +2093,7 @@ check) and combining values in a block before one atomic write. ```c #include -T cunumpy_warp_sum(T v); T cunumpy_warp_min(T v); T cunumpy_warp_max(T v); +T cunumpy_warp_sum(T v, unsigned mask = 0xffffffffu); // also min/max T cunumpy_block_sum(T v); T cunumpy_block_min(T v); T cunumpy_block_max(T v); void cunumpy_block_sum_to(T* out, T v); // *out += block sum, one atomic per block int cunumpy_block_thread(); // linear thread index in a 1D-3D block @@ -2030,11 +2101,13 @@ int cunumpy_block_threads(); // threads per block ``` `T` is `int`, `unsigned`, `long long`, `unsigned long long`, `float` or -`double` (`block_sum_to`: `double`, `float`, `int`, `unsigned long long`). Every +`double`; `block_sum_to` supports the same types. Every thread gets the result. Rules: every thread of the block calls block -functions and all 32 lanes call warp functions (no early `return`; threads -without a value pass the identity, e.g. `0.0` for a sum); the block size is a -multiple of 32. The block functions use 32 values of static shared memory per +functions (no early `return`; threads without a value pass the identity, e.g. +`0.0` for a sum). Warp functions support arbitrary nonzero masks: every named +lane calls with the same mask, and only named lanes participate. The default +requires all 32 lanes. Block functions handle partial warps and any legal block +size. The block functions use 32 values of static shared memory per type and may be called several times in a kernel. ```c @@ -2046,6 +2119,25 @@ extern "C" __global__ void kinetic_energy(const double* v, long long n, } ``` +### `cunumpy/scan.cuh` + +Warp/block prefix sums for compaction and binning: + +```c +#include +T cunumpy_warp_inclusive_sum(T v, unsigned mask = 0xffffffffu); +T cunumpy_warp_exclusive_sum(T v, unsigned mask = 0xffffffffu); +T cunumpy_block_inclusive_sum(T v); +T cunumpy_block_exclusive_sum(T v); +``` + +Same shuffle arithmetic types and collective participation rules as reductions. +Warp order follows increasing participating lane IDs, including sparse masks. +Block order follows linear thread index (x fastest), with partial warps supported. +Exclusive scans start at zero. Block scans use 32 shared values per type and +specialization; repeated calls are supported. These are block-local operations; +a global compaction needs a separate pass to combine block totals. + ## `scipy` ```python diff --git a/docs/source/guides/execution-helpers.md b/docs/source/guides/execution-helpers.md new file mode 100644 index 0000000..4aa6dc9 --- /dev/null +++ b/docs/source/guides/execution-helpers.md @@ -0,0 +1,162 @@ +# Reusable execution and grouping helpers + +Allocate streams, staging storage, and grouping plans during setup, then reuse +them in timesteps. CPU equivalents keep the same call sites usable on both backends. + +## Require CUDA and inspect the environment + +```python +import json +import cunumpy as xp + +xp.set_backend("cupy", strict=True) +print(json.dumps(xp.backend_info(), indent=2)) +``` + +Strict selection raises with the cached availability failure reason when CUDA +cannot be used, preserving the previous backend and array functions. +`use_backend("cupy", strict=True)` behaves the same way. Without strict selection, +the existing fallback to NumPy is preserved. Diagnostics include dependency +versions and current device/runtime information; they do not initialize MPI. + +## Reuse streams and events + +```python +producer = xp.cuda.create_stream() +consumer = xp.cuda.create_stream() +produced = xp.cuda.create_event() +consumed = xp.cuda.create_event() +values = xp.zeros(1000) +result = xp.zeros_like(values) +xp.synchronize() # finish initialization before using nonblocking streams + +for step in range(3): + with xp.cuda.stream(producer): + if step: + xp.cuda.wait_event(consumed, stream=producer) + values.fill(step + 1) + xp.cuda.record_event(produced, stream=producer) + with xp.cuda.stream(consumer): + xp.cuda.wait_event(produced, stream=consumer) + result += values + xp.cuda.record_event(consumed, stream=consumer) + +consumed.synchronize() +host_result = xp.to_numpy(result) +``` + +CUDA waits order consumer work without blocking the CPU. `.done` checks completion +and `.synchronize()` waits. Enqueue all consumers of an event recording before +re-recording it. Retain the arrays until their queued work finishes. + +CPU factories return synchronous `HostStream`/`HostEvent` objects with the same +completion interface. Legacy `cuda.stream()` still yields `None` on CPU. Keep a +CUDA stream's device current while using it. Events disable timing by default; +performance measurement and reporting belong in scope-profiler. + +## Reuse MPI staging buffers + +```python +MPI = xp.mpi.get_mpi() +comm = MPI.COMM_WORLD +send = xp.arange(100, dtype=xp.float64) +recv = xp.zeros_like(send) +send_staging = xp.mpi.MPIStaging(send.shape, send.dtype) +recv_staging = xp.mpi.MPIStaging(recv.shape, recv.dtype) + +with send_staging.buffer(send, cuda_aware=False) as sendbuf, recv_staging.buffer( + recv, send=False, recv=True, cuda_aware=False +) as recvbuf: + comm.Sendrecv(sendbuf, dest=0, recvbuf=recvbuf, source=0) +``` + +This example exchanges with rank 0; use application-specific peer ranks for a +distributed exchange. Each staging object lazily allocates one reusable host +buffer, pinned when available. Copies appear in transfer accounting, and receive +copy-back completes before storage is released. Host arrays pass through unchanged. +Use separate objects for simultaneous send/receive contexts: overlapping device +staging use raises. + +Fortran-order and strided device arrays use an additional C-order device packing +buffer, also allocated once and reused. Receive data is scattered back into the +original array view before the context exits. + +`mpi_buffer(array, staging=...)` is equivalent. Device staging has a fixed shape, +dtype, and device. For nonblocking MPI, call **`request.Wait()` inside the context** +before storage is released or received data is copied back. Requests are not +tracked automatically. + +Pass `stream=producer` or `event=completion` to wait for that explicit producer. +Without either, CuNumpy synchronizes all work on the buffer's device. +`synchronize_for_mpi(*arrays, stream=..., event=...)` follows the same rule; pass +at most one dependency, covering all reads/writes of the supplied buffers. + +## Prepare cell ranges and segment reductions + +```python +cells = xp.asarray([2, 0, 2, 4, 0], dtype=xp.int64) +values = xp.asarray([1.0, 2.0, 3.0, 4.0, 5.0]) +cells, order, values = xp.algorithms.sort_by_key(cells, values) +offsets = xp.algorithms.cell_offsets(cells, n_cells=6) +# [0, 2, 2, 4, 4, 5, 5]; cell k uses [offsets[k], offsets[k+1]). +unique, starts, stops = xp.algorithms.segment_boundaries(cells) +# unique=[0, 2, 4], starts=[0, 2, 4], stops=[2, 4, 5]. +plan = xp.algorithms.SegmentPlan(cells, n_segments=6) +totals = xp.empty(6) +plan.sum(values, out=totals) # [7, 0, 4, 0, 4, 0] +``` + +Dense offsets require sorted integer IDs in `[0, n_cells)`; filter negative IDs +first. Sparse run boundaries also accept negative IDs and uint64 Morton keys. +Results remain on the input backend. Setup validation and variable-length GPU +results may synchronize: prepare outside tight loops and pass GPU offsets to +device kernels instead of reading them individually in Python. + +`SegmentPlan` snapshots/validates keys once and drops negative-key rows. Repeated +sums overwrite their output and accept arbitrary trailing component dimensions. +CUDA reduces all components in one accumulation launch without per-column scalar +reads. Reusing `out` avoids output allocation; noncontiguous values, dtype +conversion, and float16 accumulation can still require temporary storage. + +Floating/complex output dtypes match the input; integer/bool values produce +float64. CUDA float16 accumulates in float32. Floating-point atomic order can +vary, so compare with tolerances. Output must match shape/dtype, be C-contiguous, +and not alias values. A GPU plan requires its device current and all arrays on +that device. Mixed backends raise. Occasional calls can use +`segment_sum(values, keys, n_segments, out=...)` with the same semantics. + +## CUDA scans, masked reductions, and integer atomics + +```c +#include +#include + +extern "C" __global__ void compact_flags(const int* keep, long long n, + int* local_offsets, int* block_counts) { + int t = cunumpy_block_thread(); + long long i = (long long)blockIdx.x * cunumpy_block_threads() + t; + int flag = i < n ? (keep[i] != 0) : 0; + int offset = cunumpy_block_exclusive_sum(flag); + int count = cunumpy_block_sum(flag); + if (i < n) local_offsets[i] = offset; + if (t == 0) block_counts[blockIdx.x] = count; +} +``` + +Offsets are block-local; global compaction needs another prefix sum of block +counts and a scatter pass. Inclusive/exclusive warp/block sums are provided. +Block collectives accept partial warps and any legal shape, with x fastest. +Every block thread must participate, including padding threads; do not use an +early-return indexing macro before a collective. + +Warp reductions and scans accept optional nonzero masks, including sparse masks. +Only named lanes participate and all must call with the same mask. Scan order +follows increasing named lane IDs. The default requires all 32 lanes. These rules +follow [NVIDIA's warp intrinsic constraints](https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html). + +`cunumpy_atomic_add` and its indexed 2D/3D variants also accept `int`, `unsigned int`, +`long long`, and `unsigned long long`, returning the old value. Signed 64-bit +addition uses CUDA's unsigned instruction with modulo-2^64 arithmetic. + +Small fixed-size matrix helpers remain deferred: this repository currently has +no repeated CUDA consumers establishing a shared routine or storage convention. diff --git a/docs/source/guides/mpi.md b/docs/source/guides/mpi.md index c355fd7..2173bba 100644 --- a/docs/source/guides/mpi.md +++ b/docs/source/guides/mpi.md @@ -87,8 +87,9 @@ CUDA-aware. CuPy kernels run asynchronously, and MPI knows nothing about CUDA streams. A buffer that a kernel is still writing would be sent as it is at that moment. -`synchronize_for_mpi(*buffers)` waits for the current stream if any argument is -a CuPy array, and costs nothing otherwise: +`synchronize_for_mpi(*buffers)` waits for each device represented by the CuPy +arrays, and costs nothing for host arrays. Pass `stream=` or `event=` to wait +only for a known producer dependency: ```python def exchange_halo(field, comm, left, right): diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index f8f8290..88e14bc 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -137,5 +137,6 @@ extern "C" __global__ void push_and_energy(double* x, double* v, const double* E } ``` -The block size must be a multiple of 32, and every thread must reach the -reduction. See the API reference for the warp and block functions. +Every thread must reach the reduction; block functions also support partial +warps. Warp functions accept explicit masks for subsets of lanes. See the API +reference for the warp and block functions and prefix scans. diff --git a/docs/source/index.md b/docs/source/index.md index fc72edf..86c7b8a 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -48,6 +48,7 @@ guides/backends guides/portable-code guides/data-movement guides/gpu-devices +guides/execution-helpers guides/mpi guides/particle-codes guides/solvers diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index a94cbe7..50589f7 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -110,6 +110,8 @@ Backend and inspection: ```python xp.set_backend("numpy" | "cupy") # global; falls back to numpy if cupy unusable +xp.set_backend("cupy", strict=True) # unavailable CUDA raises; prior backend preserved +xp.backend_info() # JSON-compatible dependencies/CUDA diagnostics xp.get_backend() -> "numpy" | "cupy" # active backend with xp.use_backend("numpy"): ... # temporary, exception-safe xp.numpy_backend, xp.cupy_backend # bools for the active backend @@ -161,6 +163,40 @@ keys, order, a, b = xp.algorithms.sort_by_key( xp.require_version("0.4.0") # ImportError if cunumpy is older ``` +Reusable helpers (allocate during setup): + +```python +producer, consumer = xp.cuda.create_stream(), xp.cuda.create_stream() +event = xp.cuda.create_event() # CPU: already-completed HostEvent +with xp.cuda.stream(producer): + ... # produce data; retain arrays until completion + xp.cuda.record_event(event, stream=producer) +xp.cuda.wait_event(event, stream=consumer) # future consumer work waits, CPU does not +staging = xp.mpi.MPIStaging(a.shape, a.dtype) +with staging.buffer(a, cuda_aware=False, recv=True) as buf: + ... # blocking MPI, or request.Wait() BEFORE exiting +offsets = xp.algorithms.cell_offsets(sorted_cells, n_cells) +unique, starts, stops = xp.algorithms.segment_boundaries(sorted_keys) +plan = xp.algorithms.SegmentPlan(keys, n_segments) # copied, validated keys +plan.sum(values, out=out) # overwrite; arbitrary trailing dimensions +``` + +CUDA streams/plans require their device current. Re-record an event only after +its consumers enqueue their waits. Staging shape/dtype/device are fixed; use +separate objects for concurrent contexts. MPI synchronization accepts `stream=` +or `event=`; without either it waits for all work on the buffer devices. +Segment keys must be integer; negative keys are dropped. Dense offsets need +sorted nonnegative cell IDs; sparse boundaries accept negative/uint64 keys. +Preparation may synchronize; repeated plan sums avoid scalar reads. `out` must +match shape/dtype, be C-contiguous, and not alias values. Mixed backends raise. +CUDA floating-point accumulation uses atomics and may require tolerances. + +`cunumpy/scan.cuh` supplies `cunumpy_{warp,block}_{inclusive,exclusive}_sum`. +Warp scans and `cunumpy_warp_{sum,min,max}` accept a nonzero mask: every named +lane calls with the same mask. Block collectives handle partial warps, but every +thread must participate (no early return). `cunumpy_atomic_add` and its indexed +variants support int32/uint32/int64/uint64 as well as float/double. + Random numbers and dtypes: ```python diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index f8f94c7..becffd4 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -18,6 +18,7 @@ from .xp import ( as_device_array, assert_same_backend, + backend_info, cupy_available, default_float_dtype, get_array_backend, @@ -98,6 +99,7 @@ def require_version(minimum: str) -> None: "algorithms", "as_device_array", "assert_same_backend", + "backend_info", "cuda", "cupy_available", "cupy_backend", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index bf57473..eadad06 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -23,6 +23,7 @@ def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... def to_cunumpy(array: Any) -> Any: ... def cupy_available() -> bool: ... +def backend_info() -> dict[str, Any]: ... def get_array_module(array: Any) -> Any: ... def get_backend() -> str: ... def get_array_backend(array: Any) -> str: ... @@ -31,8 +32,8 @@ def is_cpu(array: Any) -> bool: ... def same_backend(*arrays: Any) -> bool: ... def assert_same_backend(*arrays: Any) -> None: ... @contextmanager -def use_backend(backend: str) -> Generator[None]: ... -def set_backend(backend: str) -> None: ... +def use_backend(backend: str, *, strict: bool = ...) -> Generator[None]: ... +def set_backend(backend: str, *, strict: bool = ...) -> None: ... def require_version(minimum: str) -> None: ... def default_float_dtype() -> Any: ... def synchronize() -> None: ... diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index ee3ca00..87538c4 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -2,66 +2,212 @@ from __future__ import annotations +import math +import operator from typing import Any import array_api_compat.numpy as np -from .xp import get_array_backend, get_array_module +from .xp import assert_same_backend, get_array_backend, get_array_module -def segment_sum(values: Any, keys: Any, n_segments: int) -> Any: - """Sum `values` per key: ``out[k] = sum(values[i] for keys[i] == k)``. +def _integer_keys(keys: Any) -> tuple[Any, Any]: + xpm = get_array_module(keys) + keys = xpm.asarray(keys) + if keys.ndim != 1: + raise ValueError(f"keys must be 1D, got shape {keys.shape}") + if keys.dtype.kind not in "iu": + raise TypeError("keys must have an integer dtype") + return xpm, keys - The reduction step of a sort-then-reduce accumulation (particles binned to - cells, contributions summed per cell), on either backend, with - ``bincount`` under the hood. For a 2D `values` the columns are summed - separately (one bincount per column). - Parameters - ---------- - values : array - Shape ``(n,)`` or ``(n, m)``, on the backend of `keys`. - keys : array - Integer segment of every value, shape ``(n,)``; a negative key drops - the value (e.g. a particle outside the grid). - n_segments : int - Number of segments; keys must be smaller than it. +def _segment_count(n_segments: int) -> int: + n_segments = operator.index(n_segments) + if n_segments < 0: + raise ValueError("n_segments must be non-negative") + return n_segments - Returns - ------- - array - Shape ``(n_segments,)`` or ``(n_segments, m)``, dtype of `values` for - floating-point and complex values, ``float64`` otherwise. + +def _sorted_keys(keys: Any) -> tuple[Any, Any]: + xpm, keys = _integer_keys(keys) + if bool((keys[1:] < keys[:-1]).any()): + raise ValueError("keys must be sorted in nondecreasing order") + return xpm, keys + + +def segment_boundaries(sorted_keys: Any) -> tuple[Any, Any, Any]: + """Return ``(unique_keys, starts, stops)`` for runs of sorted integer keys. + + Starts/stops are int64 indices and describe half-open slices of the input. + Sparse, negative, and uint64 Morton keys are supported. All results stay on + the input backend. Validation and variable-length GPU output may synchronize; + prepare once and reuse the boundaries in repeated operations. """ - xpm = get_array_module(keys) - keys = xpm.asarray(keys) - values = xpm.asarray(values) - if keys.ndim != 1 or values.shape[:1] != keys.shape: - raise ValueError( - f"keys must be 1D with one entry per value, got keys {keys.shape} and " - f"values {values.shape}" - ) - if values.ndim not in (1, 2): - raise ValueError(f"values must be 1D or 2D, got shape {values.shape}") - if bool((keys >= n_segments).any()): - raise ValueError(f"keys must be smaller than n_segments={n_segments}") - valid = keys >= 0 - if not bool(valid.all()): - keys = keys[valid] - values = values[valid] - out_dtype = values.dtype if values.dtype.kind in "fc" else np.dtype(np.float64) - if values.ndim == 1: - if values.dtype.kind == "c": - real = xpm.bincount(keys, weights=values.real, minlength=n_segments) - imag = xpm.bincount(keys, weights=values.imag, minlength=n_segments) - return (real + 1j * imag).astype(out_dtype, copy=False) - return xpm.bincount(keys, weights=values, minlength=n_segments).astype( - out_dtype, copy=False + xpm, keys = _sorted_keys(sorted_keys) + if keys.size == 0: + return keys.copy(), xpm.empty(0, dtype=np.int64), xpm.empty(0, dtype=np.int64) + starts = xpm.concatenate( + (xpm.zeros(1, dtype=np.int64), xpm.nonzero(keys[1:] != keys[:-1])[0] + 1) + ).astype(np.int64, copy=False) + stops = xpm.concatenate((starts[1:], xpm.full(1, keys.size, dtype=np.int64))) + return keys[starts], starts, stops + + +def cell_offsets(sorted_cells: Any, n_cells: int) -> Any: + """Dense int64 offsets into sorted cell IDs in ``[0, n_cells)``. + + Cell k occupies ``[offsets[k], offsets[k+1])``; empty cells have equal + offsets. Returns n_cells+1 entries on the input backend. Filter negative + (invalid) cell IDs before calling. Validation may synchronize on CUDA. + """ + n_cells = _segment_count(n_cells) + xpm, cells = _sorted_keys(sorted_cells) + if bool(((cells < 0) | (cells >= n_cells)).any()): + raise ValueError("cell IDs must be in [0, n_cells)") + return xpm.searchsorted( + cells, xpm.arange(n_cells + 1, dtype=np.int64), side="left" + ).astype(np.int64, copy=False) + + +_SUM_KERNELS: dict[str, Any] = {} + + +def _sum_kernel(dtype: Any) -> Any: + name = np.dtype(dtype).name + if name not in _SUM_KERNELS: + from ._cuda_kernel import CudaKernel + + ctype = "float" if name == "float32" else "double" + _SUM_KERNELS[name] = CudaKernel( + "#include \n" + 'extern "C" __global__ void segment_sum_components(' + f"const long long* keys, const {ctype}* values, {ctype}* out, " + "long long n, long long width) {\n" + "long long i = (long long)blockDim.x * blockIdx.x + threadIdx.x;\n" + "if (i < n * width) { long long key = keys[i / width];\n" + "if (key >= 0) cunumpy_atomic_add(out + key * width + i % width, values[i]); }\n" + "}", + "segment_sum_components", ) - out = xpm.empty((n_segments, values.shape[1]), dtype=out_dtype) - for j in range(values.shape[1]): - out[:, j] = segment_sum(values[:, j], keys, n_segments) - return out + return _SUM_KERNELS[name] + + +class SegmentPlan: + """Snapshot and validate segment keys once, then reuse :meth:`sum`. + + A negative key drops its row; other keys must be below `n_segments`. + Keys are copied so later caller mutation cannot invalidate the plan. A GPU + plan is bound to its CUDA device. Setup may synchronize, but repeated sums + do not read GPU reductions back into Python or validate keys per column. + """ + + def __init__(self, keys: Any, n_segments: int) -> None: + self._n_segments = _segment_count(n_segments) + self._xp, keys = _integer_keys(keys) + if bool((keys >= self.n_segments).any()): + raise ValueError(f"keys must be smaller than n_segments={n_segments}") + self._backend = get_array_backend(keys) + self._device = keys.device.id if self._backend == "cupy" else None + if self._device is not None: + import cupy as cp + + if cp.cuda.Device().id != self._device: + raise ValueError("segment keys require their CUDA device current") + self._keys = self._xp.array(keys, dtype=np.int64, copy=True) + self._valid = None if self._device is not None else self._keys >= 0 + + @property + def n_segments(self) -> int: + """Number of output segments, fixed when the plan is prepared.""" + return self._n_segments + + def sum(self, values: Any, *, out: Any = None) -> Any: + """Sum rows into ``(n_segments, *values.shape[1:])``; overwrite `out`. + + Floating/complex dtypes are preserved; integer/bool inputs produce + float64. CUDA supports float16/32/64 and complex64/128 outputs; float16 + accumulates in a float32 workspace. Other CUDA accumulation + uses atomics and its floating-point order is not deterministic. `out` + must have the exact shape/dtype, be C-contiguous, and not alias values. + """ + assert_same_backend(self._keys, values) + xpm = self._xp + values = xpm.asarray(values) + if values.ndim < 1 or values.shape[0] != self._keys.size: + raise ValueError("keys must have one entry per value row") + if values.dtype.kind not in "biufc": + raise TypeError("values must have a numeric dtype") + if self._device is not None: + import cupy as cp + + if values.device.id != self._device or cp.cuda.Device().id != self._device: + raise ValueError( + "segment plan and values require their CUDA device current" + ) + dtype = values.dtype if values.dtype.kind in "fc" else np.dtype(np.float64) + if self._device is not None and np.dtype(dtype).name not in ( + "float16", + "float32", + "float64", + "complex64", + "complex128", + ): + raise TypeError("CUDA segment sums require float16/32/64 or complex64/128") + shape = (self.n_segments, *values.shape[1:]) + if out is None: + out = xpm.empty(shape, dtype=dtype) + else: + assert_same_backend(self._keys, out) + if out.shape != shape or out.dtype != dtype or not out.flags.c_contiguous: + raise ValueError( + "out must have the exact shape/dtype and be C-contiguous" + ) + if not getattr(out.flags, "writeable", True): + raise ValueError("out must be writable") + if self._device is not None and out.device.id != self._device: + raise ValueError("out must be on the segment plan's CUDA device") + if xpm.may_share_memory(out, values): + raise ValueError("out must not alias values") + out.fill(0) + if self._device is None: + np.add.at( + out, + self._keys[self._valid], + values[self._valid].astype(dtype, copy=False), + ) + elif values.size and self.n_segments: + half = np.dtype(dtype) == np.dtype(np.float16) + packed = xpm.ascontiguousarray(values, dtype=np.float32 if half else dtype) + if np.dtype(dtype).kind == "c": + real_dtype = np.float32 if np.dtype(dtype).itemsize == 8 else np.float64 + packed, target = packed.view(real_dtype), out.view(real_dtype) + else: + target = xpm.zeros(shape, dtype=np.float32) if half else out + width = math.prod(values.shape[1:]) * ( + 2 if np.dtype(dtype).kind == "c" else 1 + ) + _sum_kernel(packed.dtype)( + self._keys, + packed, + target, + self._keys.size, + width, + n_threads=self._keys.size * width, + ) + if half: + out[...] = target + return out + + +def segment_sum(values: Any, keys: Any, n_segments: int, *, out: Any = None) -> Any: + """Sum rows per integer key, dropping negatives; optionally overwrite `out`. + + Accepts arbitrary trailing value dimensions. Keys are validated once per + call and CUDA reduces all components in one launch, without per-column host + checks. Use ``SegmentPlan(keys, n_segments).sum(values, out=out)`` when the + keys are reused, to avoid repeating setup and its GPU synchronization. + """ + return SegmentPlan(keys, n_segments).sum(values, out=out) def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index a0c6f9a..f9316eb 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -178,7 +178,7 @@ def pin_memory(array: Any) -> Any: @contextmanager -def stream() -> Generator[Any, None, None]: +def stream(existing: Any = None) -> Generator[Any, None, None]: """Context manager for a CUDA stream, to overlap transfers and compute. On the CuPy backend, operations issued inside the block are enqueued on @@ -186,8 +186,14 @@ def stream() -> Generator[Any, None, None]: `xp.synchronize()` (or the yielded stream's own `.synchronize()`) before reading results computed inside the block. No-op on the NumPy backend, where it yields `None`. + + Pass a reusable stream from :func:`cunumpy.cuda.create_stream` as `existing` + to select it temporarily; this also works with its synchronous CPU equivalent. """ - if array_backend.backend == "cupy": + if existing is not None: + with existing: + yield existing + elif array_backend.backend == "cupy": import cupy as cp with cp.cuda.Stream(non_blocking=True) as s: diff --git a/src/cunumpy/_fake_cupy_impl.py b/src/cunumpy/_fake_cupy_impl.py index 7226868..682c07d 100644 --- a/src/cunumpy/_fake_cupy_impl.py +++ b/src/cunumpy/_fake_cupy_impl.py @@ -37,8 +37,13 @@ def __array__(self, *args, **kwargs): "`.get()` to construct a NumPy array explicitly." ) - def get(self, *args, **kwargs): - return self._a.copy() + def get(self, stream=None, order="C", out=None, blocking=True): + if out is not None: + if out.shape != self._a.shape or out.dtype != self._a.dtype: + raise ValueError("out must match the device array shape and dtype") + _np.copyto(out, self._a) + return out + return self._a.copy(order=order) def set(self, arr, *args, **kwargs): self._a[...] = arr @@ -417,6 +422,21 @@ def attributes(self): return {"MaxSharedMemoryPerBlock": 48 * 1024} +class _Event: + def __init__(self, **kwargs): + pass + + @property + def done(self): + return True + + def record(self, stream=None): + pass + + def synchronize(self): + pass + + class _Stream(_Device): def __init__(self, *args, **kwargs): super().__init__() @@ -424,6 +444,18 @@ def __init__(self, *args, **kwargs): def __enter__(self): return self + def record(self, event=None): + event = _Event() if event is None else event + event.record(self) + return event + + def wait_event(self, event): + pass + + @property + def done(self): + return True + _Stream.null = _Stream() @@ -435,6 +467,7 @@ def _alloc_pinned_memory(nbytes): cuda = types.ModuleType("cupy.cuda") cuda.Device = _Device cuda.Stream = _Stream +cuda.Event = _Event cuda.get_current_stream = lambda: _Stream.null cuda.alloc_pinned_memory = _alloc_pinned_memory cuda.device = types.ModuleType("cupy.cuda.device") diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 53f03ca..cbb6eb6 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import operator from collections.abc import Generator from contextlib import contextmanager from typing import Any @@ -18,14 +19,15 @@ _logger = logging.getLogger(__name__) -def synchronize_for_mpi(*arrays: Any) -> None: +def synchronize_for_mpi(*arrays: Any, stream: Any = None, event: Any = None) -> None: """Wait for pending device work before MPI reads or writes `arrays`. CuPy launches kernels asynchronously; MPI does not know about CUDA streams. Passing a device buffer to MPI while a kernel is still writing it sends whatever is in memory at that moment -- silently wrong data, no error. Call this before every MPI call that uses device buffers. It synchronizes the - current stream only if at least one of `arrays` is a CuPy array, so host + devices owning the arrays if no producer `stream` or `event` is supplied. + An explicit dependency waits only for that producer. Host buffers and the NumPy backend cost nothing. (After MPI returns, no synchronization is needed: kernels launched later see the received data.) @@ -34,12 +36,26 @@ def synchronize_for_mpi(*arrays: Any) -> None: *arrays The buffers about to be passed to MPI; ``None`` entries are ignored. """ - if not any(array_api_compat.is_cupy_array(a) for a in arrays if a is not None): + if stream is not None and event is not None: + raise ValueError("pass only one of stream and event") + devices = {a.device.id for a in arrays if array_api_compat.is_cupy_array(a)} + if not devices: + return + from ._streams import HostEvent, HostStream + + if isinstance(event, HostEvent) or isinstance(stream, HostStream): + raise TypeError("device buffers require a CUDA producer stream or event") + if event is not None: + event.synchronize() + return + if stream is not None: + stream.synchronize() return import cupy as cp - cp.cuda.get_current_stream().synchronize() + for device in devices: + cp.cuda.Device(device).synchronize() # the result of the last mpi_is_cuda_aware() probe (or of set_mpi_cuda_aware()), @@ -77,6 +93,74 @@ def _pinned_or_host_empty(shape: tuple[int, ...], dtype: Any) -> np.ndarray: return np.empty(shape, dtype=dtype) +class MPIStaging: + """Reusable host storage for blocking MPI exchanges of device arrays. + + Allocate once per send/receive buffer. Storage is allocated lazily, pinned + when available, and bound to one shape, dtype, and CUDA device. Only one + context may use it at a time. For nonblocking MPI, call request.Wait() inside + the context: the buffer must not be released or copied back while MPI uses it. + """ + + def __init__(self, shape: int | tuple[int, ...], dtype: Any) -> None: + self._shape = ( + (operator.index(shape),) + if isinstance(shape, (int, np.integer)) + else tuple(operator.index(n) for n in shape) + ) + if any(n < 0 for n in self.shape): + raise ValueError("staging shape must be non-negative") + self._dtype = np.dtype(dtype) + if self.dtype.hasobject: + raise TypeError("MPI staging does not support object dtype") + self._host: np.ndarray | None = None + self._packed: Any = None + self._device: int | None = None + self._active = False + + @property + def shape(self) -> tuple[int, ...]: + """The fixed staging shape.""" + return self._shape + + @property + def dtype(self) -> np.dtype: + """The fixed staging dtype.""" + return self._dtype + + def buffer(self, array: Any, **kwargs: Any) -> Any: + """Equivalent to ``mpi_buffer(array, staging=self, **kwargs)``.""" + return mpi_buffer(array, staging=self, **kwargs) + + def _transfer_array(self, array: Any) -> Any: + """Reuse C-order device storage when the MPI array has another layout.""" + if array.flags.c_contiguous: + return array + if self._packed is None: + import cupy as cp + + self._packed = cp.empty(self.shape, dtype=self.dtype) + return self._packed + + @contextmanager + def _lease(self, array: Any) -> Generator[np.ndarray, None, None]: + if self._active: + raise RuntimeError("MPI staging buffer is already in use") + if tuple(array.shape) != self.shape or np.dtype(array.dtype) != self.dtype: + raise ValueError("MPI staging shape and dtype must match the array") + device = array.device.id + if self._device is not None and device != self._device: + raise ValueError("MPI staging buffer is bound to another CUDA device") + self._device = device + if self._host is None: + self._host = _pinned_or_host_empty(self.shape, self.dtype) + self._active = True + try: + yield self._host + finally: + self._active = False + + @contextmanager def mpi_buffer( array: Any, @@ -84,6 +168,9 @@ def mpi_buffer( send: bool = True, recv: bool = False, cuda_aware: bool | None = None, + staging: MPIStaging | None = None, + stream: Any = None, + event: Any = None, ) -> Generator[Any, None, None]: """The buffer to hand to MPI for `array`: the array itself, or a host copy. @@ -116,6 +203,12 @@ def mpi_buffer( cuda_aware : bool | None Whether MPI can take device buffers. None uses the answer recorded by `mpi_is_cuda_aware()` or `set_mpi_cuda_aware()`. + staging : MPIStaging | None + Reusable storage for host staging. For nonblocking MPI, wait for the + request inside the context before releasing this storage. + stream, event + Optional producer dependency (only one); without it all work on the + array's device is synchronized before MPI accesses the buffer. Raises ------ @@ -136,20 +229,32 @@ def mpi_buffer( "xp.mpi.set_mpi_cuda_aware(True/False), or pass cuda_aware=" ) if cuda_aware: - synchronize_for_mpi(array) + synchronize_for_mpi(array, stream=stream, event=event) yield array return - host = _pinned_or_host_empty(tuple(array.shape), array.dtype) - if send: - if _COUNTERS: - _record("to_host", f"mpi_buffer({_describe(array)}) staging for send") - synchronize_for_mpi(array) - host[...] = array.get() - yield host - if recv: - if _COUNTERS: - _record("to_device", f"mpi_buffer({_describe(array)}) staging for recv") - array.set(host) + if staging is None: + staging = MPIStaging(tuple(array.shape), array.dtype) + import cupy as cp + + with cp.cuda.Device(array.device.id), staging._lease(array) as host: + # Even receive-only buffers may still be read/written by a GPU kernel. + synchronize_for_mpi(array, stream=stream, event=event) + transfer_array = staging._transfer_array(array) + if send: + if _COUNTERS: + _record("to_host", f"mpi_buffer({_describe(array)}) staging for send") + if transfer_array is not array: + transfer_array[...] = array + transfer_array.get(out=host) + yield host + if recv: + if _COUNTERS: + _record("to_device", f"mpi_buffer({_describe(array)}) staging for recv") + # Ensure MPI's host buffer can be reused immediately on context exit. + transfer_array.set(host) + if transfer_array is not array: + array[...] = transfer_array + cp.cuda.get_current_stream().synchronize() def _mpi_module() -> Any: diff --git a/src/cunumpy/_streams.py b/src/cunumpy/_streams.py new file mode 100644 index 0000000..1d1f2db --- /dev/null +++ b/src/cunumpy/_streams.py @@ -0,0 +1,97 @@ +"""Reusable CUDA streams/events with synchronous CPU equivalents.""" + +from __future__ import annotations + +from typing import Any + +from .xp import get_backend + + +class HostEvent: + """Already-completed CPU event, matching CuPy's completion interface.""" + + @property + def done(self) -> bool: + return True + + def record(self, stream: Any = None) -> None: + if stream is not None and not isinstance(stream, HostStream): + raise TypeError("a host event requires a host stream") + + def synchronize(self) -> None: + pass + + +class HostStream: + """Synchronous CPU stream; reusable and nestable as a context manager.""" + + def __enter__(self) -> Any: + return self + + def __exit__(self, *exc: object) -> None: + pass + + @property + def done(self) -> bool: + return True + + def synchronize(self) -> None: + pass + + def record(self, event: Any = None) -> HostEvent: + event = HostEvent() if event is None else event + event.record(self) + return event + + def wait_event(self, event: Any) -> None: + if not isinstance(event, HostEvent): + raise TypeError("a host stream requires a host event") + + +def create_stream(*, non_blocking: bool = True) -> Any: + """Allocate a reusable CuPy stream, or a synchronous :class:`HostStream`. + + CUDA streams belong to the device current at construction. Keep the same + device current while selecting the stream and submitting work to it. + """ + if get_backend() == "numpy": + return HostStream() + import cupy as cp + + return cp.cuda.Stream(non_blocking=non_blocking) + + +def create_event(*, timing: bool = False) -> Any: + """Allocate a reusable completion event; CPU events are always complete. + + CUDA timing is disabled by default. Pass `timing=True` for events used + with CuPy's elapsed-time functions; host events do not measure time. + """ + if get_backend() == "numpy": + return HostEvent() + import cupy as cp + + return cp.cuda.Event(disable_timing=not timing) + + +def record_event(event: Any = None, *, stream: Any = None) -> Any: + """Record completion on `stream` (current stream by default). + + Re-recording an event replaces its completion point: enqueue all consumers + of a recording before reusing the event for a subsequent producer. + """ + if event is None: + event = HostEvent() if isinstance(stream, HostStream) else create_event() + event.record(stream) + return event + + +def wait_event(event: Any, *, stream: Any = None) -> None: + """Order future work on `stream` after `event`, without blocking the CPU.""" + if stream is None: + if isinstance(event, HostEvent): + return + import cupy as cp + + stream = cp.cuda.get_current_stream() + stream.wait_event(event) diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py index 1790967..b974e03 100644 --- a/src/cunumpy/algorithms.py +++ b/src/cunumpy/algorithms.py @@ -10,7 +10,13 @@ charge = xp.algorithms.segment_sum(q, cell, n_cells) """ -from ._algorithms import segment_sum, sort_by_key +from ._algorithms import ( + SegmentPlan, + cell_offsets, + segment_boundaries, + segment_sum, + sort_by_key, +) from ._morton import ( MAX_MORTON_LEVELS, morton_decode, @@ -21,10 +27,13 @@ __all__ = [ "MAX_MORTON_LEVELS", + "SegmentPlan", + "cell_offsets", "morton_decode", "morton_encode", "morton_keys", "morton_scales", + "segment_boundaries", "segment_sum", "sort_by_key", ] diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index fc43bd1..2aca624 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -1,6 +1,6 @@ -"""CUDA-only parts of cunumpy: writing and launching CUDA kernels, and the GPU. +"""CUDA kernels, device helpers, and reusable stream/event interfaces. -Everything here is only useful with CuPy and a CUDA device. The kernel classes +The kernel classes (:class:`CudaKernel`, :class:`CudaStruct`, ...) need CuPy to launch; the device functions (:func:`set_device`, :func:`memory_info`, :func:`stream`, ...) do nothing (or return ``None``/``0``) on the NumPy backend:: @@ -14,6 +14,9 @@ The CUDA headers shipped with cunumpy (``cunumpy/atomic.cuh``, ``cunumpy/random.cuh``, ...) are in :func:`cuda_include_dir`. +Reusable :func:`create_stream` and :func:`create_event` return synchronous host +equivalents on NumPy, supporting the same recording and completion interface. + Backend-neutral kernel tools (:class:`~cunumpy.kernels.Kernel`, :class:`~cunumpy.kernels.PyccelKernel`, ...) are in :mod:`cunumpy.kernels`. @@ -53,6 +56,14 @@ set_device_for_rank, stream, ) +from .._streams import ( + HostEvent, + HostStream, + create_event, + create_stream, + record_event, + wait_event, +) __all__ = [ "DEBUG_OPTIONS", @@ -64,7 +75,11 @@ "CudaStruct", "CudaStructArguments", "CudaStructValue", + "HostEvent", + "HostStream", "bind_local_device", + "create_event", + "create_stream", "ctype_of", "cuda_debug", "cuda_include_dir", @@ -77,10 +92,12 @@ "memory_info", "parse_cuda_signature", "pin_memory", + "record_event", "resolve_includes", "set_cuda_debug", "set_device", "set_device_for_rank", "stream", + "wait_event", "write_cuda_header", ] diff --git a/src/cunumpy/cuda/include/cunumpy/atomic.cuh b/src/cunumpy/cuda/include/cunumpy/atomic.cuh index 43d1e0f..0e8f6a6 100644 --- a/src/cunumpy/cuda/include/cunumpy/atomic.cuh +++ b/src/cunumpy/cuda/include/cunumpy/atomic.cuh @@ -1,8 +1,9 @@ // cunumpy/atomic.cuh: atomic accumulation helpers for CUDA kernels. // // Many threads adding into a few cells (particle-to-grid accumulation) must -// use atomics or lose updates. These helpers wrap atomicAdd for double and -// float, and index C-contiguous 2D/3D arrays passed as bare pointers plus +// use atomics or lose updates. These helpers wrap atomicAdd for double, float, +// signed/unsigned 32-bit and 64-bit integers, and index C-contiguous 2D/3D arrays +// passed as bare pointers plus // their trailing extents. // // atomicAdd(double*, double) is a hardware instruction from compute @@ -38,6 +39,46 @@ __device__ __forceinline__ float cunumpy_atomic_add(float* p, float v) return atomicAdd(p, v); } +__device__ __forceinline__ int cunumpy_atomic_add(int* p, int v) +{ + return atomicAdd(p, v); +} + +__device__ __forceinline__ unsigned int cunumpy_atomic_add(unsigned int* p, unsigned int v) +{ + return atomicAdd(p, v); +} + +__device__ __forceinline__ unsigned long long cunumpy_atomic_add( + unsigned long long* p, unsigned long long v) +{ + return atomicAdd(p, v); +} + +// Signed 64-bit addition through CUDA's unsigned 64-bit atomic. Arithmetic +// follows the unsigned instruction's modulo-2^64 behavior. +__device__ __forceinline__ long long cunumpy_atomic_add(long long* p, long long v) +{ + return (long long)atomicAdd(reinterpret_cast(p), + (unsigned long long)v); +} + +// Typed indexed overloads also serve integer counters and occupancy arrays. +template +__device__ __forceinline__ T cunumpy_atomic_add_2d( + T* data, long long n1, long long i, long long j, T v) +{ + return cunumpy_atomic_add(data + i * n1 + j, v); +} + +template +__device__ __forceinline__ T cunumpy_atomic_add_3d( + T* data, long long n1, long long n2, + long long i, long long j, long long k, T v) +{ + return cunumpy_atomic_add(data + (i * n1 + j) * n2 + k, v); +} + // data[i, j] += v for a C-contiguous array of shape (n0, n1). __device__ __forceinline__ double cunumpy_atomic_add_2d( double* data, long long n1, long long i, long long j, double v) diff --git a/src/cunumpy/cuda/include/cunumpy/reduce.cuh b/src/cunumpy/cuda/include/cunumpy/reduce.cuh index 2371fa1..06cf58b 100644 --- a/src/cunumpy/cuda/include/cunumpy/reduce.cuh +++ b/src/cunumpy/cuda/include/cunumpy/reduce.cuh @@ -16,11 +16,11 @@ // } // // Rules: -// * Every thread of the block must call a block function, and all 32 lanes of -// the warp a warp function: do not return early. Threads without a value +// * Every thread of the block must call a block function, and every lane named +// in the mask a warp function with the same mask: do not return early. +// Warp functions default to all 32 lanes. Threads without a value // pass the identity (0 for a sum, the largest value for a minimum, ...). -// * The number of threads per block must be a multiple of 32 (CudaKernel's -// default block size is 128). +// * Block functions also support partial warps (any legal block size). // * The result is returned to every thread (warp functions: every lane). // * Supported types are those of the shuffle intrinsics: int, unsigned, // long long, unsigned long long, float, double. @@ -72,24 +72,50 @@ __device__ __forceinline__ int cunumpy_block_threads() return blockDim.x * blockDim.y * blockDim.z; } -// Reduce v over the 32 lanes of the warp; every lane gets the result. +// Reduce v over the named lanes of the warp; every participating lane gets the result. template -__device__ __forceinline__ T cunumpy_warp_reduce(T v, Op op) +__device__ __forceinline__ T cunumpy_warp_reduce( + T v, Op op, unsigned mask = CUNUMPY_FULL_WARP_MASK) { - for (int mask = CUNUMPY_WARP_SIZE / 2; mask > 0; mask /= 2) { - v = op(v, __shfl_xor_sync(CUNUMPY_FULL_WARP_MASK, v, mask)); + if (mask == CUNUMPY_FULL_WARP_MASK) { + for (int offset = CUNUMPY_WARP_SIZE / 2; offset > 0; offset /= 2) + v = op(v, __shfl_xor_sync(mask, v, offset)); + return v; } - return v; + // Broadcast only from participating lanes; works for arbitrary sparse masks. + // The mask must be nonzero and include the caller's lane. + unsigned remaining = mask; + int source = __ffs(remaining) - 1; + T result = __shfl_sync(mask, v, source); + remaining &= remaining - 1; + while (remaining) { + source = __ffs(remaining) - 1; + result = op(result, __shfl_sync(mask, v, source)); + remaining &= remaining - 1; + } + return result; } template -__device__ __forceinline__ T cunumpy_warp_sum(T v) { return cunumpy_warp_reduce(v, cunumpy_sum_op()); } +__device__ __forceinline__ T cunumpy_warp_sum(T v, unsigned mask = CUNUMPY_FULL_WARP_MASK) +{ return cunumpy_warp_reduce(v, cunumpy_sum_op(), mask); } template -__device__ __forceinline__ T cunumpy_warp_min(T v) { return cunumpy_warp_reduce(v, cunumpy_min_op()); } +__device__ __forceinline__ T cunumpy_warp_min(T v, unsigned mask = CUNUMPY_FULL_WARP_MASK) +{ return cunumpy_warp_reduce(v, cunumpy_min_op(), mask); } template -__device__ __forceinline__ T cunumpy_warp_max(T v) { return cunumpy_warp_reduce(v, cunumpy_max_op()); } +__device__ __forceinline__ T cunumpy_warp_max(T v, unsigned mask = CUNUMPY_FULL_WARP_MASK) +{ return cunumpy_warp_reduce(v, cunumpy_max_op(), mask); } + +// Mask of the physically present lanes of the caller's warp. All threads of +// the block must reach the collective; this is not a divergent-branch mask. +__device__ __forceinline__ unsigned cunumpy_block_warp_mask() +{ + const int warp = cunumpy_block_thread() / CUNUMPY_WARP_SIZE; + const int lanes = cunumpy_block_threads() - warp * CUNUMPY_WARP_SIZE; + return lanes >= CUNUMPY_WARP_SIZE ? CUNUMPY_FULL_WARP_MASK : (1u << lanes) - 1u; +} // Reduce v over all threads of the block; every thread gets the result. // Uses 32 values of static shared memory per type and synchronizes the block @@ -103,13 +129,14 @@ __device__ T cunumpy_block_reduce(T v, Op op) const int warp = thread / CUNUMPY_WARP_SIZE; const int n_warps = (cunumpy_block_threads() + CUNUMPY_WARP_SIZE - 1) / CUNUMPY_WARP_SIZE; - v = cunumpy_warp_reduce(v, op); + const unsigned mask = cunumpy_block_warp_mask(); + v = cunumpy_warp_reduce(v, op, mask); __syncthreads(); // a previous call may still be reading partial if (lane == 0) partial[warp] = v; __syncthreads(); if (warp == 0) { v = lane < n_warps ? partial[lane] : Op::fill(partial); - v = cunumpy_warp_reduce(v, op); + v = cunumpy_warp_reduce(v, op, mask); if (lane == 0) partial[0] = v; } __syncthreads(); @@ -151,4 +178,12 @@ __device__ __forceinline__ void cunumpy_block_sum_to(unsigned long long* out, un if (cunumpy_block_thread() == 0) atomicAdd(out, total); } +// Also covers unsigned int and signed long long counters. +template +__device__ __forceinline__ void cunumpy_block_sum_to(T* out, T v) +{ + const T total = cunumpy_block_sum(v); + if (cunumpy_block_thread() == 0) cunumpy_atomic_add(out, total); +} + #endif // CUNUMPY_REDUCE_CUH diff --git a/src/cunumpy/cuda/include/cunumpy/scan.cuh b/src/cunumpy/cuda/include/cunumpy/scan.cuh new file mode 100644 index 0000000..8ad852c --- /dev/null +++ b/src/cunumpy/cuda/include/cunumpy/scan.cuh @@ -0,0 +1,84 @@ +// Prefix sums for compaction and binning. Scan order is increasing linear +// thread index (x fastest), or increasing lane index among a warp mask's bits. +// Every named lane must call a warp function with the same nonzero mask; +// every thread must call a block function, including padding threads (value 0). +// Supports CUDA shuffle arithmetic types, arbitrary masks and partial warps. +// Block scans use 32 shared values per type; calls must be collective. +#ifndef CUNUMPY_SCAN_CUH +#define CUNUMPY_SCAN_CUH + +#include "cunumpy/reduce.cuh" + +template +__device__ __forceinline__ T cunumpy_warp_inclusive_sum( + T v, unsigned mask = CUNUMPY_FULL_WARP_MASK) +{ + const int lane = cunumpy_block_thread() % CUNUMPY_WARP_SIZE; + if ((mask & (mask + 1u)) == 0u) { + // Full warp or contiguous prefix of lanes. + for (int offset = 1; offset < CUNUMPY_WARP_SIZE; offset *= 2) { + T previous = __shfl_up_sync(mask, v, offset); + if (lane >= offset) v += previous; + } + return v; + } + T result = T(0); + unsigned remaining = mask; + while (remaining) { + const int source = __ffs(remaining) - 1; + T item = __shfl_sync(mask, v, source); + if (source <= lane) result += item; + remaining &= remaining - 1; + } + return result; +} + +template +__device__ __forceinline__ T cunumpy_warp_exclusive_sum( + T v, unsigned mask = CUNUMPY_FULL_WARP_MASK) +{ + const int lane = cunumpy_block_thread() % CUNUMPY_WARP_SIZE; + const T inclusive = cunumpy_warp_inclusive_sum(v, mask); + const unsigned before = mask & ((1u << lane) - 1u); + // Shuffle for every participating lane, including the first (whose result + // is discarded); subtraction would lose precision for large current values. + const int source = before ? 31 - __clz(before) : lane; + const T previous = __shfl_sync(mask, inclusive, source); + return before ? previous : T(0); +} + +template +__device__ T cunumpy_block_scan_sum(T v) +{ + __shared__ T partial[CUNUMPY_WARP_SIZE]; + const int thread = cunumpy_block_thread(); + const int lane = thread % CUNUMPY_WARP_SIZE; + const int warp = thread / CUNUMPY_WARP_SIZE; + const int n_warps = (cunumpy_block_threads() + CUNUMPY_WARP_SIZE - 1) / CUNUMPY_WARP_SIZE; + const unsigned mask = cunumpy_block_warp_mask(); + const T inclusive = cunumpy_warp_inclusive_sum(v, mask); + const unsigned before = mask & ((1u << lane) - 1u); + const int source = before ? 31 - __clz(before) : lane; + const T previous = __shfl_sync(mask, inclusive, source); + const T local = Exclusive ? (before ? previous : T(0)) : inclusive; + __syncthreads(); // protect readers from a previous scan invocation + if (lane == 31 - __clz(mask)) partial[warp] = inclusive; + __syncthreads(); + if (warp == 0) { + const T total = cunumpy_warp_inclusive_sum( + lane < n_warps ? partial[lane] : T(0), mask); + if (lane < n_warps) partial[lane] = total; + } + __syncthreads(); + return local + (warp ? partial[warp - 1] : T(0)); +} + +template +__device__ T cunumpy_block_inclusive_sum(T v) +{ return cunumpy_block_scan_sum(v); } + +template +__device__ T cunumpy_block_exclusive_sum(T v) +{ return cunumpy_block_scan_sum(v); } + +#endif // CUNUMPY_SCAN_CUH diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index da454cc..b9d8206 100644 --- a/src/cunumpy/mpi.py +++ b/src/cunumpy/mpi.py @@ -22,6 +22,7 @@ """ from ._mpi import ( + MPIStaging, get_mpi_cuda_aware, mpi_buffer, mpi_is_cuda_aware, @@ -42,6 +43,7 @@ __all__ = [ "OVERRIDE_VARIABLE", + "MPIStaging", "SerialComm", "SerialMPI", "SerialRequest", diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 661bd58..6dafcf4 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -2,9 +2,11 @@ import logging import os +import sys import warnings from collections.abc import Callable, Generator from contextlib import contextmanager +from importlib.metadata import PackageNotFoundError, version from types import ModuleType from typing import TYPE_CHECKING, Any, Literal @@ -26,11 +28,12 @@ _CUPY_AVAILABLE_CACHE = None +_CUPY_UNAVAILABLE_REASON: str | None = None def cupy_available() -> bool: """Check if CuPy is available and functional.""" - global _CUPY_AVAILABLE_CACHE + global _CUPY_AVAILABLE_CACHE, _CUPY_UNAVAILABLE_REASON if _CUPY_AVAILABLE_CACHE is not None: return _CUPY_AVAILABLE_CACHE @@ -39,9 +42,13 @@ def cupy_available() -> bool: # Check if a GPU is available _CUPY_AVAILABLE_CACHE = cp.is_available() + _CUPY_UNAVAILABLE_REASON = ( + None if _CUPY_AVAILABLE_CACHE else "CuPy reports no usable CUDA device" + ) return _CUPY_AVAILABLE_CACHE - except Exception: # noqa: BLE001 - tolerate any driver/runtime failure + except Exception as error: # noqa: BLE001 - tolerate driver/runtime failure _CUPY_AVAILABLE_CACHE = False + _CUPY_UNAVAILABLE_REASON = f"{type(error).__name__}: {error}" return False @@ -110,19 +117,26 @@ def _set(self, backend: BackendType, module: ModuleType) -> None: for listener in self._listeners: listener(module) - def set(self, backend: BackendType) -> None: + def set(self, backend: BackendType, *, strict: bool = False) -> None: """Select `backend` (falls back to NumPy if CuPy is not functional).""" if backend not in ("numpy", "cupy"): raise ValueError("Array backend must be either 'numpy' or 'cupy'.") + if strict and backend == "cupy" and not cupy_available(): + raise RuntimeError( + "Cannot select the CuPy backend: " + + (_CUPY_UNAVAILABLE_REASON or "CuPy/CUDA is unavailable") + ) module = self._load_backend(backend) # sets self._backend to the effective one self._set(self._backend, module) @contextmanager - def use_backend(self, backend: BackendType) -> Generator[None, None, None]: + def use_backend( + self, backend: BackendType, *, strict: bool = False + ) -> Generator[None, None, None]: """Temporarily change the backend.""" old_backend = self._backend old_xp = self._xp - self.set(backend) + self.set(backend, strict=strict) try: yield finally: @@ -137,14 +151,60 @@ def use_backend(self, backend: BackendType) -> Generator[None, None, None]: ) -def use_backend(backend: BackendType) -> Generator[None, None, None]: +def use_backend( + backend: BackendType, *, strict: bool = False +) -> Generator[None, None, None]: """Temporarily change the backend.""" - return array_backend.use_backend(backend) + return array_backend.use_backend(backend, strict=strict) + + +def set_backend(backend: BackendType, *, strict: bool = False) -> None: + """Select a backend; with `strict`, unavailable CUDA raises without switching.""" + array_backend.set(backend, strict=strict) + +def backend_info() -> dict[str, Any]: + """Return JSON-compatible backend, dependency, and CUDA diagnostics. + + Does not change the backend or import MPI. CUDA availability is cached, as + in :func:`cupy_available`. Device inspection errors are reported in the + result rather than hiding the otherwise useful CPU/dependency information. + """ + versions = {} + for package in ("cunumpy", "numpy", "array-api-compat"): + try: + versions[package] = version(package) + except PackageNotFoundError: + versions[package] = None + available = bool(cupy_available()) + versions["cupy"] = getattr(sys.modules.get("cupy"), "__version__", None) + info: dict[str, Any] = { + "backend": get_backend(), + "cupy_available": available, + "cuda_unavailable_reason": None if available else _CUPY_UNAVAILABLE_REASON, + "versions": versions, + "device": None, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + if available: + try: + import cupy as cp -def set_backend(backend: BackendType) -> None: - """Set the backend globally.""" - array_backend.set(backend) + versions["cupy"] = getattr(cp, "__version__", None) + dev = cp.cuda.Device() + properties = cp.cuda.runtime.getDeviceProperties(dev.id) + name = properties["name"] + info["device"] = { + "id": int(dev.id), + "name": name.decode() if isinstance(name, bytes) else str(name), + "compute_capability": str(dev.compute_capability), + "visible_devices": int(cp.cuda.runtime.getDeviceCount()), + "driver_version": int(cp.cuda.runtime.driverGetVersion()), + "runtime_version": int(cp.cuda.runtime.runtimeGetVersion()), + } + except Exception as error: # noqa: BLE001 - diagnostics must remain usable + info["cuda_inspection_error"] = f"{type(error).__name__}: {error}" + return info def get_backend() -> BackendType: diff --git a/tests/unit/test_backend_diagnostics.py b/tests/unit/test_backend_diagnostics.py new file mode 100644 index 0000000..8ffe8bb --- /dev/null +++ b/tests/unit/test_backend_diagnostics.py @@ -0,0 +1,81 @@ +"""Strict selection is transactional; diagnostic reports survive CUDA failures.""" + +import json +import sys +from types import SimpleNamespace + +import pytest + +import cunumpy as xp +from cunumpy import xp as backend + + +def test_strict_failure_preserves_backend_and_namespace(monkeypatch): + with xp.use_backend("numpy"): + zeros = xp.zeros + monkeypatch.setattr(backend, "_CUPY_AVAILABLE_CACHE", False) + monkeypatch.setattr(backend, "_CUPY_UNAVAILABLE_REASON", "driver mismatch") + with pytest.raises(RuntimeError, match="driver mismatch"): + xp.set_backend("cupy", strict=True) + with ( + pytest.raises(RuntimeError, match="driver mismatch"), + xp.use_backend("cupy", strict=True), + ): + pytest.fail("unavailable backend was entered") + assert xp.get_backend() == "numpy" and xp.zeros is zeros + xp.set_backend("cupy") # preserve the existing fallback + assert xp.get_backend() == "numpy" + info = xp.backend_info() + assert info["cuda_unavailable_reason"] == "driver mismatch" + assert info["device"] is None + json.dumps(info) + + +def test_availability_keeps_exception_details(monkeypatch): + def unavailable(): + raise RuntimeError("bad driver") + + monkeypatch.setattr(backend, "_CUPY_AVAILABLE_CACHE", None) + monkeypatch.setattr(backend, "_CUPY_UNAVAILABLE_REASON", None) + monkeypatch.setitem(sys.modules, "cupy", SimpleNamespace(is_available=unavailable)) + assert not backend.cupy_available() + assert xp.backend_info()["cuda_unavailable_reason"] == "RuntimeError: bad driver" + + +def test_diagnostics_with_cuda_and_failed_inspection(monkeypatch): + runtime = SimpleNamespace( + getDeviceProperties=lambda device: {"name": b"test GPU"}, + getDeviceCount=lambda: 2, + driverGetVersion=lambda: 13000, + runtimeGetVersion=lambda: 12000, + ) + cp = SimpleNamespace( + __version__="test", + cuda=SimpleNamespace( + Device=lambda: SimpleNamespace(id=1, compute_capability="80"), + runtime=runtime, + ), + ) + monkeypatch.setitem(sys.modules, "cupy", cp) + monkeypatch.setattr(backend, "_CUPY_AVAILABLE_CACHE", True) + info = xp.backend_info() + assert info["device"]["id"] == 1 and info["device"]["name"] == "test GPU" + assert info["versions"]["cupy"] == "test" + assert info["cuda_unavailable_reason"] is None + json.dumps(info) + del runtime.getDeviceProperties + info = xp.backend_info() + assert info["cupy_available"] and "cuda_inspection_error" in info + monkeypatch.setitem(sys.modules, "cupy", None) + info = xp.backend_info() + assert "ModuleNotFoundError" in info["cuda_inspection_error"] + + +@pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") +def test_strict_cuda_selection_restores_namespace(): + with xp.use_backend("numpy"): + zeros = xp.zeros + with xp.use_backend("cupy", strict=True): + assert xp.get_backend() == "cupy" + assert xp.zeros is not zeros + assert xp.zeros is zeros diff --git a/tests/unit/test_cuda_collectives.py b/tests/unit/test_cuda_collectives.py new file mode 100644 index 0000000..cd52c0b --- /dev/null +++ b/tests/unit/test_cuda_collectives.py @@ -0,0 +1,271 @@ +"""Hardware tests for scans/masked reductions, plus CPU-emulated integer atomics.""" + +import subprocess + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy.cuda import CudaKernel +from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler + +requires_cuda = pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") + +TYPED_SOURCE = r""" +#include +template +__global__ void typed_scan(const T* x, T* inclusive, T* exclusive, T* total) { + int t = cunumpy_block_thread(); + int i = blockIdx.x * cunumpy_block_threads() + t; + inclusive[i] = cunumpy_block_inclusive_sum(x[i]); + exclusive[i] = cunumpy_block_exclusive_sum(x[i]); + cunumpy_block_sum_to(total, x[i]); +} +""" + + +@requires_cuda +@pytest.mark.parametrize( + "dtype", [np.int32, np.uint32, np.int64, np.uint64, np.float32, np.float64] +) +def test_scans_and_block_atomic_totals_for_shuffle_types(dtype): + import cupy as cp + + host = np.arange(70, dtype=dtype) + data = cp.asarray(host) + inc, exc, total = cp.zeros_like(data), cp.zeros_like(data), cp.zeros(1, dtype=dtype) + kernel = CudaKernel( + TYPED_SOURCE, "typed_scan", template_args=(dtype,), block_size=35 + ) + kernel(data, inc, exc, total, grid=2) + expected = np.cumsum(host.reshape(2, 35), axis=1) + np.testing.assert_array_equal(cp.asnumpy(inc).reshape(2, 35), expected) + np.testing.assert_array_equal( + cp.asnumpy(exc).reshape(2, 35), + np.concatenate((np.zeros((2, 1)), expected[:, :-1]), axis=1), + ) + np.testing.assert_array_equal(cp.asnumpy(total), [host.sum()]) + + +BLOCK_SOURCE = r""" +#include +extern "C" __global__ void collectives(const double* x, double* inclusive, + double* exclusive, double* again, double* sums, double* mins, double* maxs) { + int t = cunumpy_block_thread(); + int count = cunumpy_block_threads(); + int i = blockIdx.x * count + t; + double v = x[i]; + inclusive[i] = cunumpy_block_inclusive_sum(v); + exclusive[i] = cunumpy_block_exclusive_sum(v); + again[i] = cunumpy_block_inclusive_sum(v * 2); + double sum = cunumpy_block_sum(v); + double lo = cunumpy_block_min(v); + double hi = cunumpy_block_max(v); + if (t == 0) { sums[blockIdx.x] = sum; mins[blockIdx.x] = lo; maxs[blockIdx.x] = hi; } +} +""" + + +@requires_cuda +@pytest.mark.parametrize( + "block", [1, 17, 31, 32, 33, 64, 127, 128, 1024, (7, 5), (3, 4, 3)] +) +def test_partial_warps_multidimensional_blocks_and_repeated_collectives(block): + import cupy as cp + + count = int(np.prod(block)) + host = np.arange(3 * count, dtype=float) % 11 - 5 + data = cp.asarray(host) + inc, exc, again = (cp.empty_like(data) for _ in range(3)) + sums, lo, hi = (cp.empty(3) for _ in range(3)) + kernel = CudaKernel(BLOCK_SOURCE, "collectives", block_size=block) + kernel(data, inc, exc, again, sums, lo, hi, grid=3) + rows = host.reshape(3, count) + expected_inc = np.cumsum(rows, axis=1) + expected_exc = np.concatenate((np.zeros((3, 1)), expected_inc[:, :-1]), axis=1) + np.testing.assert_array_equal(cp.asnumpy(inc).reshape(3, count), expected_inc) + np.testing.assert_array_equal(cp.asnumpy(exc).reshape(3, count), expected_exc) + np.testing.assert_array_equal(cp.asnumpy(again).reshape(3, count), expected_inc * 2) + for got, expected in ((sums, rows.sum(1)), (lo, rows.min(1)), (hi, rows.max(1))): + np.testing.assert_array_equal(cp.asnumpy(got), expected) + + +MASKED_SOURCE = r""" +#include +extern "C" __global__ void masked(const double* x, unsigned int mask, + double* sums, double* mins, double* maxs, double* inclusive, double* exclusive) { + int lane = threadIdx.x; + if (mask & (1u << lane)) { + double v = x[lane]; + sums[lane] = cunumpy_warp_sum(v, mask); + mins[lane] = cunumpy_warp_min(v, mask); + maxs[lane] = cunumpy_warp_max(v, mask); + inclusive[lane] = cunumpy_warp_inclusive_sum(v, mask); + exclusive[lane] = cunumpy_warp_exclusive_sum(v, mask); + } +} +""" + + +@requires_cuda +@pytest.mark.parametrize( + "mask", + [1, 1 << 31, 0x80000001, 0x55555555, 0xAAAAAAAA, 0x17, 0x7FFFFFFF, 0xFFFFFFFF], +) +def test_arbitrary_sparse_masks(mask): + import cupy as cp + + host = np.arange(32, dtype=float) - 10 + active = np.array([i for i in range(32) if mask & (1 << i)]) + out = [cp.full(32, -999.0) for _ in range(5)] + kernel = CudaKernel(MASKED_SOURCE, "masked", block_size=32) + kernel(cp.asarray(host), mask, *out, n_threads=32) + selected = host[active] + expected = [ + selected.sum(), + selected.min(), + selected.max(), + selected.cumsum(), + np.concatenate(([0.0], selected.cumsum()[:-1])), + ] + inactive = np.setdiff1d(np.arange(32), active) + for got, want in zip(out, expected): + result = cp.asnumpy(got) + np.testing.assert_array_equal( + result[active], np.broadcast_to(want, active.shape) + ) + np.testing.assert_array_equal(result[inactive], -999) + + +@requires_cuda +@pytest.mark.parametrize("mask", [0xFFFFFFFF, 0x80000001]) +def test_exclusive_scan_preserves_small_previous_values(mask): + import cupy as cp + + host = np.zeros(32) + active = [i for i in range(32) if mask & (1 << i)] + host[active[0]], host[active[1]] = 1.0, 1e20 + out = [cp.zeros(32) for _ in range(5)] + CudaKernel(MASKED_SOURCE, "masked", block_size=32)( + cp.asarray(host), mask, *out, n_threads=32 + ) + assert float(out[-1][active[1]]) == 1.0 + + +ATOMICS_SOURCE = r""" +#include +extern "C" __global__ void integer_counts(int* a, unsigned int* b, + long long* c, unsigned long long* d, long long* previous, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + cunumpy_atomic_add_2d(a, 3, i % 2, i % 3, 1); + cunumpy_atomic_add_3d(b, 2, 3, 0, i % 2, i % 3, 1u); + previous[i] = cunumpy_atomic_add(c, -1ll); + cunumpy_atomic_add(d, 1ull); + } +} +""" + + +@pytest.mark.parametrize("backend", ["emulation", "cuda"]) +def test_integer_atomic_counts_and_returned_old_values(backend): + n = 71 + if backend == "emulation": + if emulation_compiler() is None: + pytest.skip("requires a C++ compiler") + arrays = [ + np.zeros((2, 3), np.int32), + np.zeros((1, 2, 3), np.uint32), + np.zeros(1, np.int64), + np.zeros(1, np.uint64), + np.zeros(n, np.int64), + ] + kernel = CudaKernel(ATOMICS_SOURCE, "integer_counts") + emulate_cuda_kernel(kernel, *arrays, n, n_threads=n) + else: + if not xp.cupy_available(): + pytest.skip("requires CUDA") + import cupy as cp + + arrays = [ + cp.zeros((2, 3), cp.int32), + cp.zeros((1, 2, 3), cp.uint32), + cp.zeros(1, cp.int64), + cp.zeros(1, cp.uint64), + cp.zeros(n, cp.int64), + ] + CudaKernel(ATOMICS_SOURCE, "integer_counts")(*arrays, n, n_threads=n) + arrays = [cp.asnumpy(a) for a in arrays] + expected = np.zeros((2, 3), dtype=int) + for i in range(n): + expected[i % 2, i % 3] += 1 + np.testing.assert_array_equal(arrays[0], expected) + np.testing.assert_array_equal(arrays[1][0], expected) + assert arrays[2][0] == -n and arrays[3][0] == n + np.testing.assert_array_equal(np.sort(arrays[4]), np.sort(-np.arange(n))) + + +def test_scan_header_is_resolved_and_emulation_refuses_warp_intrinsics(): + kernel = CudaKernel(BLOCK_SOURCE, "collectives") + assert [p.name for p in kernel.included_headers] == [ + "scan.cuh", + "reduce.cuh", + "atomic.cuh", + ] + if emulation_compiler() is None: + pytest.skip("requires a C++ compiler") + with pytest.raises(NotImplementedError, match="warp shuffles"): + emulate_cuda_kernel(kernel, *(np.zeros(32) for _ in range(7)), n_threads=32) + + +def test_collective_headers_compile_with_all_supported_arithmetic_types(): + # A C++ syntax/instantiation check only. Hardware tests above establish warp + # semantics; identity stubs must never be used as a correctness emulator. + from cunumpy._emulation import _STUBS + from cunumpy.cuda import cuda_include_dir + + compiler = emulation_compiler() + if compiler is None: + pytest.skip("requires a C++ compiler") + stubs = r""" +inline int __ffs(unsigned v) { return v ? __builtin_ctz(v) + 1 : 0; } +inline int __clz(unsigned v) { return __builtin_clz(v); } +inline void __syncthreads() {} +template T __shfl_sync(unsigned, T v, int) { return v; } +template T __shfl_up_sync(unsigned, T v, int) { return v; } +template T __shfl_xor_sync(unsigned, T v, int) { return v; } +""" + instantiate = r""" +template void check_type(T v) { + cunumpy_block_inclusive_sum(v); + cunumpy_block_exclusive_sum(v); + cunumpy_warp_exclusive_sum(v, 0x55555555u); + cunumpy_block_sum(v); cunumpy_block_min(v); cunumpy_block_max(v); + T out = T(0); cunumpy_block_sum_to(&out, v); +} +void check_all() { + check_type(1); check_type(1u); check_type(1ll); check_type(1ull); + check_type(1.0f); check_type(1.0); +} +""" + result = subprocess.run( + [ + compiler, + "-std=c++17", + "-fsyntax-only", + "-x", + "c++", + "-I" + cuda_include_dir(), + "-", + ], + input=_STUBS + + stubs + + BLOCK_SOURCE + + MASKED_SOURCE + + ATOMICS_SOURCE + + instantiate, + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/unit/test_mpi_staging.py b/tests/unit/test_mpi_staging.py new file mode 100644 index 0000000..655cfc8 --- /dev/null +++ b/tests/unit/test_mpi_staging.py @@ -0,0 +1,194 @@ +"""Host staging reuse and lifetime rules, without requiring an MPI installation.""" + +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import _mpi +from cunumpy.mpi import MPIStaging, mpi_buffer + + +@pytest.fixture +def device(monkeypatch): + log = [] + + class Device: + def __init__(self, id=0): + self.id = id + + def __enter__(self): + return self + + def __exit__(self, *exc): + pass + + def synchronize(self): + log.append(("device_sync", self.id)) + + class Array: + def __init__(self, data, id=0): + self.values = np.asarray(data) + self.shape, self.dtype = self.values.shape, self.values.dtype + self.flags = self.values.flags + self.device = Device(id) + + def get(self, out=None): + log.append(("get", id(out))) + np.copyto(out, self.values) + return out + + def set(self, host): + log.append(("set", id(host))) + self.values[:] = host + + def __setitem__(self, index, value): + self.values[index] = value.values + + def allocate_device(shape, dtype): + log.append(("allocate_device",)) + return Array(np.empty(shape, dtype=dtype)) + + stream = SimpleNamespace(synchronize=lambda: log.append(("stream_sync",))) + monkeypatch.setitem( + sys.modules, + "cupy", + SimpleNamespace( + cuda=SimpleNamespace(Device=Device, get_current_stream=lambda: stream), + empty=allocate_device, + ), + ) + monkeypatch.setattr( + _mpi.array_api_compat, "is_cupy_array", lambda a: isinstance(a, Array) + ) + + def allocate(shape, dtype): + log.append(("allocate",)) + return np.empty(shape, dtype=dtype) + + monkeypatch.setattr(_mpi, "_pinned_or_host_empty", allocate) + return Array, log + + +def test_staging_reuses_storage_counts_copies_and_waits_for_copyback(device): + Array, log = device + data = Array([1.0, 2.0]) + staging = MPIStaging(2, np.float64) + buffers = [] + with xp.profiling.count_transfers() as counter: + for _ in range(3): + with staging.buffer(data, recv=True, cuda_aware=False) as host: + buffers.append(host) + host += 1 + assert all(host is buffers[0] for host in buffers) + assert log.count(("allocate",)) == 1 + assert log[-1] == ("stream_sync",) + assert counter.to_host == counter.to_device == 3 + np.testing.assert_array_equal(data.values, [4, 5]) + + +def test_strided_arrays_reuse_device_packing_and_preserve_unselected_rows(device): + Array, log = device + base = np.arange(24.0).reshape(8, 3) + expected = base.copy() + expected[::2] += 2 + data = Array(base[::2]) + staging = MPIStaging(data.shape, data.dtype) + for _ in range(2): + with staging.buffer(data, recv=True, cuda_aware=False) as host: + host += 1 + np.testing.assert_array_equal(base, expected) + assert log.count(("allocate",)) == log.count(("allocate_device",)) == 1 + + +def test_staging_rejects_overlapping_uses_and_recovers_after_exception(device): + Array, _ = device + data = Array([1.0, 2.0]) + staging = MPIStaging(2, np.float64) + with ( + pytest.raises(RuntimeError, match="already in use"), + mpi_buffer(data, staging=staging, cuda_aware=False), + staging.buffer(data, cuda_aware=False), + ): + pass + with staging.buffer(data, cuda_aware=False) as host: + np.testing.assert_array_equal(host, [1, 2]) + with ( + pytest.raises(ValueError, match="shape and dtype"), + staging.buffer(Array([1.0]), cuda_aware=False), + ): + pass + with ( + pytest.raises(ValueError, match="another CUDA device"), + staging.buffer(Array([1.0, 2.0], id=1), cuda_aware=False), + ): + pass + + +def test_explicit_producer_dependencies_and_receive_only(device): + Array, log = device + data = Array([1.0, 2.0]) + producer = SimpleNamespace(synchronize=lambda: log.append(("producer",))) + with mpi_buffer( + data, send=False, recv=True, cuda_aware=False, event=producer + ) as host: + host[:] = [9, 10] + assert ("producer",) in log + assert not any(entry[0] == "device_sync" for entry in log) + np.testing.assert_array_equal(data.values, [9, 10]) + log.clear() + with mpi_buffer(data, cuda_aware=True, stream=producer) as buf: + assert buf is data + assert log == [("producer",)] + with pytest.raises(ValueError, match="only one"): + _mpi.synchronize_for_mpi(data, stream=producer, event=producer) + with pytest.raises(TypeError, match="CUDA producer"): + _mpi.synchronize_for_mpi(data, event=xp.cuda.HostEvent()) + + +def test_default_sync_waits_for_each_device(device): + Array, log = device + _mpi.synchronize_for_mpi(Array([1], id=0), Array([2], id=1), None) + assert sorted(log) == [("device_sync", 0), ("device_sync", 1)] + + +def test_host_arrays_are_passed_through_and_invalid_staging_is_rejected(): + data = np.zeros(2) + with MPIStaging(2, data.dtype).buffer(data) as buf: + assert buf is data + with pytest.raises(ValueError, match="non-negative"): + MPIStaging((-1,), float) + with pytest.raises(TypeError, match="object"): + MPIStaging(2, object) + + +@pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") +@pytest.mark.parametrize("layout", ["C", "F", "strided"]) +def test_real_device_staging_reuses_host_storage_after_nondefault_producer(layout): + import cupy as cp + + data = cp.zeros((4, 5)) + if layout == "F": + data = cp.asfortranarray(data) + elif layout == "strided": + data = cp.zeros((8, 5))[::2] + staging = MPIStaging(data.shape, data.dtype) + producer = cp.cuda.Stream(non_blocking=True) + cp.cuda.Device().synchronize() + first, packed = None, None + for value in (3, 8): + with producer: + data.fill(value) + with staging.buffer(data, recv=True, cuda_aware=False, stream=producer) as host: + if first is None: + first = host + assert host is first + np.testing.assert_array_equal(host, np.full((4, 5), value)) + host[:] += 2 + np.testing.assert_array_equal(cp.asnumpy(data), np.full((4, 5), value + 2)) + if layout != "C": + if packed is None: + packed = staging._packed + assert staging._packed is packed diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index dbcb34c..21f1376 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -584,6 +584,14 @@ def test_require_version(monkeypatch): assert sorted(e.kind for e in counter.events) == ["to_device", "to_host"] with xp.mpi.mpi_buffer(d, cuda_aware=True) as buf: assert buf is d +producer = xp.cuda.create_stream() +event = xp.cuda.create_event() +with xp.cuda.stream(producer): + assert xp.cuda.record_event(event, stream=producer) is event +xp.cuda.wait_event(event, stream=producer) +assert producer.done and event.done +assert producer.record().done +assert xp.backend_info()["versions"]["cupy"] == "0.0.0+cunumpy-fake" print("fake cupy OK") """ diff --git a/tests/unit/test_reusable_streams.py b/tests/unit/test_reusable_streams.py new file mode 100644 index 0000000..c179fe1 --- /dev/null +++ b/tests/unit/test_reusable_streams.py @@ -0,0 +1,92 @@ +"""Reusable completion primitives and cross-stream producer dependencies.""" + +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import _streams + + +def test_host_stream_and_event_can_be_reused_and_nested(): + with xp.use_backend("numpy"): + stream = xp.cuda.create_stream() + event = xp.cuda.create_event() + for _ in range(3): + with xp.cuda.stream(stream) as selected, stream: + assert selected is stream + assert xp.cuda.record_event(event, stream=stream) is event + stream.wait_event(event) + xp.cuda.wait_event(event) + assert stream.record().done + assert stream.done and event.done + event.synchronize() + stream.synchronize() + with xp.cuda.stream() as legacy: + assert legacy is None + + +def test_cuda_factories_and_event_ordering(monkeypatch): + calls = [] + + class Stream: + def __init__(self, **kwargs): + calls.append(("stream", kwargs)) + + def wait_event(self, event): + calls.append(("wait", event)) + + class Event: + def __init__(self, **kwargs): + calls.append(("event", kwargs)) + + def record(self, stream=None): + calls.append(("record", stream)) + + current = Stream() + monkeypatch.setattr(_streams, "get_backend", lambda: "cupy") + monkeypatch.setitem( + sys.modules, + "cupy", + SimpleNamespace( + cuda=SimpleNamespace( + Stream=Stream, Event=Event, get_current_stream=lambda: current + ) + ), + ) + stream = xp.cuda.create_stream(non_blocking=False) + event = xp.cuda.create_event() + assert xp.cuda.record_event(event, stream=stream) is event + xp.cuda.wait_event(event) + assert calls[-1] == ("wait", event) + assert ("stream", {"non_blocking": False}) in calls + assert ("event", {"disable_timing": True}) in calls + xp.cuda.create_event(timing=True) + assert calls[-1] == ("event", {"disable_timing": False}) + assert isinstance(xp.cuda.record_event(stream=stream), Event) + + +@pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") +def test_gpu_events_order_reused_producer_and_consumer(): + import cupy as cp + + with xp.use_backend("cupy", strict=True): + producer, consumer = xp.cuda.create_stream(), xp.cuda.create_stream() + ready, consumed = xp.cuda.create_event(), xp.cuda.create_event() + data, result = cp.zeros(1000), cp.zeros(1000) + cp.cuda.Device().synchronize() + for step in range(3): + with producer: + if step: + xp.cuda.wait_event(consumed, stream=producer) + data.fill(step + 1) + xp.cuda.record_event(ready, stream=producer) + with consumer: + xp.cuda.wait_event(ready, stream=consumer) + result += data + xp.cuda.record_event(consumed, stream=consumer) + consumed.synchronize() + assert consumed.done + np.testing.assert_array_equal(cp.asnumpy(result), np.full(1000, 6)) diff --git a/tests/unit/test_segment_plans.py b/tests/unit/test_segment_plans.py new file mode 100644 index 0000000..867c69e --- /dev/null +++ b/tests/unit/test_segment_plans.py @@ -0,0 +1,159 @@ +"""Grouping and reusable reductions on CPU and (when available) CUDA.""" + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy.algorithms import ( + SegmentPlan, + cell_offsets, + segment_boundaries, + segment_sum, +) +from cunumpy.kernel_testing import BACKENDS + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_dense_and_sparse_boundaries(backend): + with xp.use_backend(backend, strict=True): + cells = xp.asarray([0, 0, 2, 4, 4], dtype=xp.int64) + np.testing.assert_array_equal( + xp.to_numpy(cell_offsets(cells, 6)), [0, 2, 2, 3, 3, 5, 5] + ) + unique, starts, stops = segment_boundaries(cells) + for got, expected in zip( + (unique, starts, stops), ([0, 2, 4], [0, 2, 3], [2, 3, 5]) + ): + assert xp.get_array_backend(got) == backend + np.testing.assert_array_equal(xp.to_numpy(got), expected) + large = xp.asarray([2**63 + 1, 2**63 + 1, 2**64 - 1], dtype=xp.uint64) + unique, starts, stops = segment_boundaries(large) + np.testing.assert_array_equal( + xp.to_numpy(unique), np.array([2**63 + 1, 2**64 - 1], dtype=np.uint64) + ) + np.testing.assert_array_equal(xp.to_numpy(starts), [0, 2]) + np.testing.assert_array_equal(xp.to_numpy(stops), [2, 3]) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_empty_and_invalid_grouping(backend): + with xp.use_backend(backend, strict=True): + empty = xp.empty(0, dtype=xp.int64) + np.testing.assert_array_equal(xp.to_numpy(cell_offsets(empty, 3)), [0, 0, 0, 0]) + assert all(a.size == 0 for a in segment_boundaries(empty)) + np.testing.assert_array_equal(xp.to_numpy(cell_offsets(empty, 0)), [0]) + for bad in ([2, 0], [-1, 0], [0, 3]): + with pytest.raises(ValueError): + cell_offsets(xp.asarray(bad), 3) + with pytest.raises(TypeError, match="integer"): + segment_boundaries(xp.asarray([1.5])) + with pytest.raises(ValueError, match="1D"): + segment_boundaries(xp.zeros((2, 2), dtype=xp.int64)) + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + "dtype", + [ + np.float16, + np.float32, + np.float64, + np.complex64, + np.complex128, + np.int64, + np.bool_, + ], +) +def test_prepared_reductions_and_reusable_output(backend, dtype): + keys_host = np.array([0, 2, 0, -1, 2, 0]) + values_host = np.arange(36).reshape(6, 2, 3).astype(dtype) + if np.dtype(dtype).kind == "c": + values_host += 1j * values_host + output_dtype = dtype if np.dtype(dtype).kind in "fc" else np.float64 + expected = np.zeros((4, 2, 3), dtype=output_dtype) + np.add.at(expected, keys_host[keys_host >= 0], values_host[keys_host >= 0]) + with xp.use_backend(backend, strict=True): + keys = xp.array(keys_host, copy=True) + plan = SegmentPlan(keys, 4) + keys[:] = 99 # the plan owns a validated snapshot + values = xp.asarray(values_host) + out = xp.full(expected.shape, -1, dtype=output_dtype) + assert plan.sum(values, out=out) is out + np.testing.assert_allclose(xp.to_numpy(out), expected, rtol=1e-5) + plan.sum(values, out=out) # overwrite, rather than accumulate prior sums + np.testing.assert_allclose(xp.to_numpy(out), expected, rtol=1e-5) + # Non-contiguous trailing values are accepted. + got = segment_sum(values[:, :, ::2], xp.asarray(keys_host), 4) + np.testing.assert_allclose(xp.to_numpy(got), expected[:, :, ::2], rtol=1e-5) + with pytest.raises(AttributeError): + plan.n_segments = 1 + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_empty_dropped_rows_and_output_validation(backend): + with xp.use_backend(backend, strict=True): + got = segment_sum(xp.empty((0, 2)), xp.empty(0, dtype=xp.int64), 3) + np.testing.assert_array_equal(xp.to_numpy(got), np.zeros((3, 2))) + got = segment_sum(xp.asarray([np.nan, np.inf]), xp.asarray([-1, -2]), 0) + assert got.shape == (0,) + plan = SegmentPlan(xp.asarray([0, 1]), 2) + values = xp.ones(2) + for out in ( + values, + xp.zeros(3), + xp.zeros(2, dtype=xp.float32), + xp.zeros(4)[::2], + ): + with pytest.raises(ValueError): + plan.sum(values, out=out) + with pytest.raises(ValueError, match="one entry"): + plan.sum(xp.ones(3)) + with pytest.raises(ValueError, match="non-negative"): + SegmentPlan(xp.asarray([0]), -1) + with pytest.raises(ValueError, match="smaller"): + SegmentPlan(xp.asarray([2]), 2) + + +def test_cpu_readonly_out_is_rejected(): + out = np.zeros(2) + out.flags.writeable = False + with pytest.raises(ValueError, match="writable"): + SegmentPlan(np.array([0, 1]), 2).sum(np.ones(2), out=out) + + +@pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") +def test_gpu_plan_avoids_scalar_sync_and_rejects_host_values(): + import cupy as cp + + plan = SegmentPlan(cp.array([0, -1, 1]), 2) + values = cp.arange(12.0).reshape(3, 4) + with pytest.raises(TypeError, match="mismatched"): + plan.sum(np.ones((3, 4))) + # A scalar-read-free sum can be captured; setup/compilation is outside capture. + out = cp.zeros((2, 4)) + plan.sum(values, out=out) + stream = cp.cuda.Stream(non_blocking=True) + cp.cuda.Device().synchronize() + with stream: + stream.begin_capture() + plan.sum(values, out=out) + graph = stream.end_capture() + graph.launch(stream) + stream.synchronize() + np.testing.assert_array_equal(cp.asnumpy(out), [[0, 1, 2, 3], [8, 9, 10, 11]]) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_fused_sum_kernel_matches_reference_in_cpu_emulation(dtype): + from cunumpy._algorithms import _sum_kernel + from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler + + if emulation_compiler() is None: + pytest.skip("requires a C++ compiler") + keys = np.array([0, 2, -1, 0, 2], dtype=np.int64) + values = np.arange(30, dtype=dtype).reshape(5, 6) + out = np.zeros((3, 6), dtype=dtype) + emulate_cuda_kernel(_sum_kernel(dtype), keys, values, out, 5, 6, n_threads=30) + expected = np.zeros_like(out) + np.add.at(expected, keys[keys >= 0], values[keys >= 0]) + np.testing.assert_array_equal(out, expected) From 8bd83fbe9b1851dbb0ab6998d0dd9e5e7fa18f72 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:20:38 +0200 Subject: [PATCH 2/9] formatting --- src/cunumpy/__init__.py | 36 +++------------- src/cunumpy/_cuda_kernel.py | 9 ++-- src/cunumpy/_dispatch.py | 8 +--- src/cunumpy/_emulation.py | 9 +--- src/cunumpy/_mpi.py | 3 +- src/cunumpy/algorithms.py | 18 ++------ src/cunumpy/cuda/__init__.py | 52 ++++++----------------- src/cunumpy/kernel_testing.py | 12 ++---- src/cunumpy/kernels.py | 18 +++----- src/cunumpy/mpi.py | 25 +++-------- src/cunumpy/profiling.py | 8 +--- src/cunumpy/rng.py | 12 ++---- tests/unit/test_cuda_kernel.py | 24 +++-------- tests/unit/test_kernel_dispatch.py | 15 ++----- tests/unit/test_kernel_dispatch_arrays.py | 9 +--- tests/unit/test_kernel_testing.py | 14 +++--- tests/unit/test_porting_helpers.py | 3 +- tests/unit/test_segment_plans.py | 8 +--- 18 files changed, 73 insertions(+), 210 deletions(-) diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index becffd4..9440a50 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -3,37 +3,13 @@ import warnings as _warnings from importlib.metadata import PackageNotFoundError, version -from . import ( - algorithms, - cuda, - kernels, - memory, - mpi, - petsc, - profiling, - rng, - xp, -) +from . import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp from ._scipy_backend import scipy -from .xp import ( - as_device_array, - assert_same_backend, - backend_info, - cupy_available, - default_float_dtype, - get_array_backend, - get_array_module, - get_backend, - is_cpu, - is_gpu, - same_backend, - set_backend, - synchronize, - to_cunumpy, - to_cupy, - to_numpy, - use_backend, -) +from .xp import (as_device_array, assert_same_backend, backend_info, + cupy_available, default_float_dtype, get_array_backend, + get_array_module, get_backend, is_cpu, is_gpu, same_backend, + set_backend, synchronize, to_cunumpy, to_cupy, to_numpy, + use_backend) # Names that were at the top level before cunumpy 0.5, and the submodule each # moved to. They still resolve (with a DeprecationWarning) until cunumpy 0.6. diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 7b6fb13..f6f66a9 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -54,7 +54,8 @@ class is the one definition of the arguments. import re import sys import typing -from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence +from collections.abc import (Callable, Hashable, Iterable, Iterator, Mapping, + Sequence) from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from pathlib import Path @@ -2384,10 +2385,8 @@ def _opt_in_shared_memory(self, kernel: Any, shared_mem: int) -> None: attribute ``max_dynamic_shared_size_bytes``; it is set once (and again for a larger request) up to the device's opt-in limit. """ - from ._device import ( - DEFAULT_SHARED_MEMORY_PER_BLOCK, - max_shared_memory_per_block, - ) + from ._device import (DEFAULT_SHARED_MEMORY_PER_BLOCK, + max_shared_memory_per_block) if shared_mem <= DEFAULT_SHARED_MEMORY_PER_BLOCK: return diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index 4f9ea93..7b719be 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -35,12 +35,8 @@ import array_api_compat from ._cuda_kernel import CudaKernel, _compile_in_threads -from ._kernel import ( - CompiledHostKernel, - HostImplementations, - PyccelKernel, - resolve_host_args, -) +from ._kernel import (CompiledHostKernel, HostImplementations, PyccelKernel, + resolve_host_args) from ._transfers import _ACTIVE as _COUNTERS from ._transfers import _record from .xp import get_backend diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index 020b096..cda8343 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -55,13 +55,8 @@ import numpy as np -from ._cuda_kernel import ( - CudaKernel, - CudaParameter, - _scalar_checker, - _strip_comments, - cuda_include_dir, -) +from ._cuda_kernel import (CudaKernel, CudaParameter, _scalar_checker, + _strip_comments, cuda_include_dir) __all__ = ["emulate_cuda_kernel", "emulation_compiler"] diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index cbb6eb6..ba006cf 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -11,7 +11,8 @@ import array_api_compat import array_api_compat.numpy as np -from ._mpi_serial import _LOCAL_RANK_VARIABLES, local_rank # noqa: F401 - re-exported +from ._mpi_serial import (_LOCAL_RANK_VARIABLES, # noqa: F401 - re-exported + local_rank) from ._transfers import _ACTIVE as _COUNTERS from ._transfers import _describe, _record from .xp import array_backend, cupy_available, to_numpy diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py index b974e03..330310f 100644 --- a/src/cunumpy/algorithms.py +++ b/src/cunumpy/algorithms.py @@ -10,20 +10,10 @@ charge = xp.algorithms.segment_sum(q, cell, n_cells) """ -from ._algorithms import ( - SegmentPlan, - cell_offsets, - segment_boundaries, - segment_sum, - sort_by_key, -) -from ._morton import ( - MAX_MORTON_LEVELS, - morton_decode, - morton_encode, - morton_keys, - morton_scales, -) +from ._algorithms import (SegmentPlan, cell_offsets, segment_boundaries, + segment_sum, sort_by_key) +from ._morton import (MAX_MORTON_LEVELS, morton_decode, morton_encode, + morton_keys, morton_scales) __all__ = [ "MAX_MORTON_LEVELS", diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index 2aca624..dcb1654 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -24,46 +24,18 @@ use ``import cupy; cupy.cuda`` for CuPy's module. """ -from .._cuda_kernel import ( - DEBUG_OPTIONS, - CudaArguments, - CudaKernel, - CudaKernelVariants, - CudaParameter, - CudaStruct, - CudaStructArguments, - CudaStructValue, - ctype_of, - cuda_include_dir, - cuda_kernel_names, - include_hash, - parse_cuda_signature, - resolve_includes, - write_cuda_header, -) -from .._device import ( - DEFAULT_SHARED_MEMORY_PER_BLOCK, - bind_local_device, - cuda_debug, - device_count, - free_memory, - get_cuda_debug, - max_shared_memory_per_block, - memory_info, - pin_memory, - set_cuda_debug, - set_device, - set_device_for_rank, - stream, -) -from .._streams import ( - HostEvent, - HostStream, - create_event, - create_stream, - record_event, - wait_event, -) +from .._cuda_kernel import (DEBUG_OPTIONS, CudaArguments, CudaKernel, + CudaKernelVariants, CudaParameter, CudaStruct, + CudaStructArguments, CudaStructValue, ctype_of, + cuda_include_dir, cuda_kernel_names, include_hash, + parse_cuda_signature, resolve_includes, + write_cuda_header) +from .._device import (DEFAULT_SHARED_MEMORY_PER_BLOCK, bind_local_device, + cuda_debug, device_count, free_memory, get_cuda_debug, + max_shared_memory_per_block, memory_info, pin_memory, + set_cuda_debug, set_device, set_device_for_rank, stream) +from .._streams import (HostEvent, HostStream, create_event, create_stream, + record_event, wait_event) __all__ = [ "DEBUG_OPTIONS", diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 961d437..6db0037 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -57,15 +57,9 @@ def test_parity(name, kernel): import numpy as np from . import _fake_cupy -from ._cuda_kernel import ( - CudaKernel, - CudaParameter, - CudaStructArguments, - CudaStructValue, - _parse_parameter, - _split_top_level, - _strip_comments, -) +from ._cuda_kernel import (CudaKernel, CudaParameter, CudaStructArguments, + CudaStructValue, _parse_parameter, _split_top_level, + _strip_comments) from ._dispatch import Kernel from ._emulation import emulate_cuda_kernel, emulation_compiler from .xp import cupy_available, get_backend, to_numpy, use_backend diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 9149fb5..39f47b6 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -22,19 +22,11 @@ from ._cuda_kernel import PyccelStructArguments from ._dispatch import Kernel, KernelCatalog from ._fusion import fuse -from ._kernel import ( - HOST_IMPLEMENTATIONS, - CompiledHostKernel, - HostImplementations, - KernelArguments, - PyccelKernel, - as_kernel_array, - get_kernel_implementation, - kernel_output, - resolve_host_args, - set_kernel_implementation, - use_kernel_implementation, -) +from ._kernel import (HOST_IMPLEMENTATIONS, CompiledHostKernel, + HostImplementations, KernelArguments, PyccelKernel, + as_kernel_array, get_kernel_implementation, + kernel_output, resolve_host_args, + set_kernel_implementation, use_kernel_implementation) __all__ = [ "HOST_IMPLEMENTATIONS", diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index b9d8206..6cc34ac 100644 --- a/src/cunumpy/mpi.py +++ b/src/cunumpy/mpi.py @@ -21,25 +21,12 @@ mpi4py is imported only by the functions that need it. """ -from ._mpi import ( - MPIStaging, - get_mpi_cuda_aware, - mpi_buffer, - mpi_is_cuda_aware, - require_cuda_aware_mpi, - set_mpi_cuda_aware, - synchronize_for_mpi, -) -from ._mpi_serial import ( - OVERRIDE_VARIABLE, - SerialComm, - SerialMPI, - SerialRequest, - SerialStatus, - get_mpi, - launched_under_mpi, - local_rank, -) +from ._mpi import (MPIStaging, get_mpi_cuda_aware, mpi_buffer, + mpi_is_cuda_aware, require_cuda_aware_mpi, + set_mpi_cuda_aware, synchronize_for_mpi) +from ._mpi_serial import (OVERRIDE_VARIABLE, SerialComm, SerialMPI, + SerialRequest, SerialStatus, get_mpi, + launched_under_mpi, local_rank) __all__ = [ "OVERRIDE_VARIABLE", diff --git a/src/cunumpy/profiling.py b/src/cunumpy/profiling.py index 7fb3ee7..2b18a8a 100644 --- a/src/cunumpy/profiling.py +++ b/src/cunumpy/profiling.py @@ -14,12 +14,8 @@ """ from ._profiling import Timing, nvtx_range, timed_region -from ._transfers import ( - TransferCounter, - TransferEvent, - assert_no_transfers, - count_transfers, -) +from ._transfers import (TransferCounter, TransferEvent, assert_no_transfers, + count_transfers) __all__ = [ "Timing", diff --git a/src/cunumpy/rng.py b/src/cunumpy/rng.py index 69ca9c1..94f9c93 100644 --- a/src/cunumpy/rng.py +++ b/src/cunumpy/rng.py @@ -16,14 +16,10 @@ ``random`` module. """ -from ._philox import ( - philox4x32_10, - philox_normal, - philox_normal2, - philox_uniform, - philox_uniform2, -) -from ._random_streams import BIT_GENERATORS, RandomStreams, get_rng, random_streams +from ._philox import (philox4x32_10, philox_normal, philox_normal2, + philox_uniform, philox_uniform2) +from ._random_streams import (BIT_GENERATORS, RandomStreams, get_rng, + random_streams) __all__ = [ "BIT_GENERATORS", diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 71583ea..522a836 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -15,20 +15,11 @@ import cunumpy as xp import cunumpy._device as device_module from cunumpy import as_device_array -from cunumpy.cuda import ( - CudaArguments, - CudaKernel, - CudaKernelVariants, - CudaStruct, - CudaStructValue, - ctype_of, - cuda_include_dir, - cuda_kernel_names, - include_hash, - parse_cuda_signature, - resolve_includes, - write_cuda_header, -) +from cunumpy.cuda import (CudaArguments, CudaKernel, CudaKernelVariants, + CudaStruct, CudaStructValue, ctype_of, + cuda_include_dir, cuda_kernel_names, include_hash, + parse_cuda_signature, resolve_includes, + write_cuda_header) AXPY = r""" // y = a * x + y @@ -468,9 +459,7 @@ def test_ctype_of(): ], ) -PUSH_SOURCE = ( - PARTICLES.declaration - + r""" +PUSH_SOURCE = PARTICLES.declaration + r""" extern "C" __global__ void push(Particles p, double dt, double* out, unsigned long long* size) { int i = blockDim.x * blockIdx.x + threadIdx.x; @@ -481,7 +470,6 @@ def test_ctype_of(): if (i < p.n && p.alive[i]) p.x[i] += dt * p.charge; } """ -) def test_struct_layout_and_declaration(): diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py index ce113b1..64c9c04 100644 --- a/tests/unit/test_kernel_dispatch.py +++ b/tests/unit/test_kernel_dispatch.py @@ -15,13 +15,8 @@ import cunumpy as xp from cunumpy.cuda import CudaKernel -from cunumpy.kernels import ( - Kernel, - KernelArguments, - KernelCatalog, - PyccelKernel, - resolve_host_args, -) +from cunumpy.kernels import (Kernel, KernelArguments, KernelCatalog, + PyccelKernel, resolve_host_args) SCALE_CUDA = r""" extern "C" __global__ void scale(double* x, double factor, int n) { @@ -129,13 +124,11 @@ def kernel_package(tmp_path, monkeypatch): for name, body in (("scale", "x[i] *= a"), ("shift", "x[i] += a")): (root / name).mkdir(parents=True) (root / name / "__init__.py").write_text("") - (root / name / f"{name}_kernels.py").write_text( - textwrap.dedent(f""" + (root / name / f"{name}_kernels.py").write_text(textwrap.dedent(f""" def {name}(x, a, n): for i in range(n): {body} - """) - ) + """)) (root / "scale" / "scale_cuda.cu").write_text(SCALE_CUDA) (root / "not_a_kernel").mkdir() (root / "__init__.py").write_text( diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index 0373c96..063c287 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -12,13 +12,8 @@ import cunumpy as xp from cunumpy import _dispatch as dispatch_module from cunumpy.cuda import CudaArguments, CudaKernel -from cunumpy.kernels import ( - CompiledHostKernel, - HostImplementations, - Kernel, - KernelArguments, - KernelCatalog, -) +from cunumpy.kernels import (CompiledHostKernel, HostImplementations, Kernel, + KernelArguments, KernelCatalog) SCALE_CUDA = r""" extern "C" __global__ void scale(double* x, double factor, int n) { diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index cd2cf39..e3313ff 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -16,15 +16,11 @@ import cunumpy as xp import cunumpy.kernel_testing from cunumpy.cuda import CudaArguments, CudaKernel, parse_cuda_signature -from cunumpy.kernel_testing import ( - BACKENDS, - _collect_arrays, - _compare_results, - assert_kernels_agree, - backend, # noqa: F401 - the fixture is used by name - device_function_kernel, - requires_cupy, -) +from cunumpy.kernel_testing import \ + backend # noqa: F401 - the fixture is used by name +from cunumpy.kernel_testing import (BACKENDS, _collect_arrays, + _compare_results, assert_kernels_agree, + device_function_kernel, requires_cupy) from cunumpy.kernels import Kernel SCALE_CUDA = r""" diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index 21f1376..2615c25 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -26,7 +26,8 @@ import cunumpy.kernel_testing from cunumpy._dispatch import FORTRAN_NAME_LIMIT, _pyccel_stub_parameters from cunumpy.cuda import CudaKernel, CudaStruct -from cunumpy.kernel_testing import check_parity, device_function_kernel, parity_cases +from cunumpy.kernel_testing import (check_parity, device_function_kernel, + parity_cases) from cunumpy.kernels import Kernel, KernelCatalog, PyccelStructArguments diff --git a/tests/unit/test_segment_plans.py b/tests/unit/test_segment_plans.py index 867c69e..312bd6d 100644 --- a/tests/unit/test_segment_plans.py +++ b/tests/unit/test_segment_plans.py @@ -4,12 +4,8 @@ import pytest import cunumpy as xp -from cunumpy.algorithms import ( - SegmentPlan, - cell_offsets, - segment_boundaries, - segment_sum, -) +from cunumpy.algorithms import (SegmentPlan, cell_offsets, segment_boundaries, + segment_sum) from cunumpy.kernel_testing import BACKENDS From 3ae61e3a85293053f1a6b2181604355e9f521704 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:22:31 +0200 Subject: [PATCH 3/9] Added trailing commas --- README.md | 18 ++- docs/source/api.md | 100 ++++++++------ docs/source/examples/particle-pusher.md | 37 ++++-- docs/source/examples/portable-script.md | 4 +- docs/source/guides/backends.md | 4 +- docs/source/guides/data-movement.md | 10 +- docs/source/guides/execution-helpers.md | 7 +- docs/source/guides/gpu-devices.md | 6 +- docs/source/guides/mpi.md | 20 ++- docs/source/guides/profiling.md | 5 +- docs/source/guides/solvers.md | 8 +- docs/source/index.md | 2 +- docs/source/kernels/accumulation.md | 8 +- docs/source/kernels/arguments.md | 30 +++-- docs/source/kernels/cuda-kernel.md | 4 +- docs/source/kernels/debugging.md | 8 +- docs/source/kernels/dispatch.md | 12 +- docs/source/kernels/pyccel-kernel.md | 8 +- docs/source/kernels/testing.md | 18 ++- docs/source/quickstart.md | 4 +- src/cunumpy/__init__.py | 2 +- src/cunumpy/_algorithms.py | 12 +- src/cunumpy/_cuda_kernel.py | 136 ++++++++++++-------- src/cunumpy/_device.py | 4 +- src/cunumpy/_dispatch.py | 23 ++-- src/cunumpy/_emulation.py | 27 ++-- src/cunumpy/_fake_cupy.py | 4 +- src/cunumpy/_fake_cupy_impl.py | 27 ++-- src/cunumpy/_fusion.py | 4 +- src/cunumpy/_kernel.py | 41 +++--- src/cunumpy/_mirror.py | 4 +- src/cunumpy/_morton.py | 6 +- src/cunumpy/_mpi.py | 14 +- src/cunumpy/_mpi_serial.py | 2 +- src/cunumpy/_philox.py | 3 +- src/cunumpy/_random_streams.py | 18 ++- src/cunumpy/_scipy_backend.py | 6 +- src/cunumpy/_staging.py | 9 +- src/cunumpy/_transfers.py | 2 +- src/cunumpy/kernel_testing.py | 28 ++-- src/cunumpy/petsc.py | 10 +- src/cunumpy/xp.py | 19 ++- tests/unit/test_array_api_compat_backend.py | 3 +- tests/unit/test_benchmarks.py | 6 +- tests/unit/test_cuda_collectives.py | 19 ++- tests/unit/test_cuda_kernel.py | 112 +++++++++++----- tests/unit/test_emulation.py | 39 ++++-- tests/unit/test_fusion.py | 4 +- tests/unit/test_kernel_dispatch.py | 14 +- tests/unit/test_kernel_dispatch_arrays.py | 35 +++-- tests/unit/test_kernel_testing.py | 25 ++-- tests/unit/test_mirror.py | 5 +- tests/unit/test_morton.py | 21 ++- tests/unit/test_mpi_serial.py | 20 ++- tests/unit/test_mpi_staging.py | 10 +- tests/unit/test_petsc.py | 4 +- tests/unit/test_philox.py | 7 +- tests/unit/test_porting_helpers.py | 46 ++++--- tests/unit/test_profiling.py | 2 +- tests/unit/test_pyccel_kernel.py | 4 +- tests/unit/test_reusable_streams.py | 6 +- tests/unit/test_scipy_backend.py | 6 +- tests/unit/test_segment_plans.py | 9 +- tests/unit/test_staging.py | 8 +- tests/unit/test_transfers.py | 7 +- 65 files changed, 719 insertions(+), 407 deletions(-) diff --git a/README.md b/README.md index f1bdfe6..b2dfe0a 100644 --- a/README.md +++ b/README.md @@ -259,8 +259,7 @@ print(timing.elapsed, timing.synced) @xp.profiling.nvtx_range("step") -def step(dt): - ... +def step(dt): ... ``` ## Use NumPy-only kernels with CuPy arrays @@ -408,11 +407,13 @@ class is the one definition, and written to a header that a test keeps in sync: ```python class MarkerArguments: - def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): - ... + def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ... + MarkerArgs = xp.cuda.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") -MarkerArgs.to_header("marker_args.cuh") # Array2D markers; long long n_markers; ... +MarkerArgs.to_header( + "marker_args.cuh" +) # Array2D markers; long long n_markers; ... push = xp.cuda.CudaKernel( r""" #include "marker_args.cuh" @@ -425,8 +426,11 @@ push = xp.cuda.CudaKernel( structs=[MarkerArgs], include_dirs=["."], ) -push(MarkerArgs(markers=markers, n_markers=markers.shape[0], valid=valid), 0.1, - n_threads=markers.shape[0]) +push( + MarkerArgs(markers=markers, n_markers=markers.shape[0], valid=valid), + 0.1, + n_threads=markers.shape[0], +) ``` Launches can be 1D to 3D (`n_threads=(nx, ny)`, `block_size=(16, 16)`) or use diff --git a/docs/source/api.md b/docs/source/api.md index cca787a..46a777f 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -377,7 +377,9 @@ an error. ```python with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer: - gather(xp.kernels.as_kernel_array(positions, like=field, dtype=float), field, buffer) + gather( + xp.kernels.as_kernel_array(positions, like=field, dtype=float), field, buffer + ) ``` ## Random numbers and dtype @@ -399,7 +401,7 @@ sequences across NumPy and CuPy. ```python xp.rng.random_streams.seed(42, rank=comm.Get_rank(), bit_generator="PCG64") v = xp.rng.random_streams.normal(0.0, v_th, (n, 3)) -rng = xp.rng.random_streams.generator() # numpy or cupy Generator +rng = xp.rng.random_streams.generator() # numpy or cupy Generator own = xp.rng.random_streams.make_generator(seed) # a component's own generator ``` @@ -473,12 +475,12 @@ otherwise the serial stand-in, so that the same code runs with and without MPI: ```python -MPI = xp.mpi.get_mpi() # decided once per process +MPI = xp.mpi.get_mpi() # decided once per process comm = MPI.COMM_WORLD -comm.Allreduce(MPI.IN_PLACE, rho, op=MPI.SUM) # nothing to do on one process -n_total = comm.allreduce(n_local) # n_local itself +comm.Allreduce(MPI.IN_PLACE, rho, op=MPI.SUM) # nothing to do on one process +n_total = comm.allreduce(n_local) # n_local itself if isinstance(MPI, xp.mpi.SerialMPI): - ... # a serial run + ... # a serial run ``` `get_mpi(True)` imports mpi4py (`ImportError` if missing), `get_mpi(False)` @@ -731,8 +733,7 @@ with xp.profiling.nvtx_range("push markers"): @xp.profiling.nvtx_range("accumulate") -def accumulate(particles, grid): - ... +def accumulate(particles, grid): ... ``` ### `profiling.timed_region(name, *, sync=True)` @@ -804,6 +805,7 @@ def scale_and_shift(scale, values, out): out[:] = scale * values + 1 return out + kernel = xp.kernels.PyccelKernel(scale_and_shift, outputs=(2,)) with xp.use_backend("cupy"): @@ -924,8 +926,8 @@ with a hash of their contents to the options: ```python kernel = xp.cuda.CudaKernel.from_file("push/push_cuda.cu", include_dirs=[src_root]) -kernel.included_headers # (Path('push/helpers.cuh'), Path('.../common.cuh')) -kernel.options # ('-Ipush', '-I') +kernel.included_headers # (Path('push/helpers.cuh'), Path('.../common.cuh')) +kernel.options # ('-Ipush', '-I') kernel.compile_options() # options + ('-DCUNUMPY_INCLUDE_HASH=0x3f9a...',) ``` @@ -1070,7 +1072,7 @@ void scale_column(Array2D a, long long column, double factor) { } """ scale_column = xp.cuda.CudaKernel(SCALE_COLUMN, "scale_column") -view = markers[::2, 1:5] # non-contiguous is fine +view = markers[::2, 1:5] # non-contiguous is fine scale_column(view, 1, 10.0, n_threads=view.shape[0]) ``` @@ -1125,10 +1127,12 @@ key: ```python matvec = xp.cuda.CudaKernelVariants( - lambda ndim, dtype: xp.cuda.CudaKernel(make_source(ndim, xp.cuda.ctype_of(dtype)), "matvec") + lambda ndim, dtype: xp.cuda.CudaKernel( + make_source(ndim, xp.cuda.ctype_of(dtype)), "matvec" + ) ) matvec.get(3, np.float64)(mat, x, out, n_threads=out.size) # created once -matvec.compile_all([(3, np.float64), (3, np.complex128)]) # at setup +matvec.compile_all([(3, np.float64), (3, np.complex128)]) # at setup ``` `get(*key)` calls the factory the first time a key is used; `keys()`, @@ -1141,7 +1145,7 @@ threads (see `KernelCatalog.compile_all`). ```python xp.cuda.set_cuda_debug(enabled) xp.cuda.get_cuda_debug() -xp.cuda.cuda_debug(enabled=True) # context manager +xp.cuda.cuda_debug(enabled=True) # context manager xp.cuda.CudaKernel(..., debug=None) kernel.debug_active() kernel.compile_options() @@ -1179,7 +1183,9 @@ now would use. ```python with xp.cuda.cuda_debug(): kernel = xp.cuda.CudaKernel(SOURCE, "kernel") - kernel(x, y, n, n_threads=n) # RuntimeError: CUDA error after launching kernel 'kernel' ... + kernel( + x, y, n, n_threads=n + ) # RuntimeError: CUDA error after launching kernel 'kernel' ... ``` The `RuntimeError` says which kernel failed, not where. The next step is @@ -1200,12 +1206,15 @@ Particles = xp.cuda.CudaStruct( "Particles", [("x", "double*"), ("v", "double*"), ("n", "int"), ("charge", "double")], ) -source = Particles.declaration + r""" +source = ( + Particles.declaration + + r""" extern "C" __global__ void push(Particles p, double dt) { int i = blockDim.x * blockIdx.x + threadIdx.x; if (i < p.n) p.x[i] += dt * p.charge * p.v[i]; } """ +) push = xp.cuda.CudaKernel(source, "push", structs=[Particles]) push(Particles(x=x, v=v, n=x.size, charge=-1.0), 0.1, n_threads=x.size) ``` @@ -1251,9 +1260,9 @@ into it when passed to a kernel. ### Structs from Python annotations ```python -class MarkerArguments: # the pyccel argument class, e.g. in struphy - def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): - ... +class MarkerArguments: # the pyccel argument class, e.g. in struphy + def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ... + MarkerArgs = xp.cuda.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") print(MarkerArgs.declaration) @@ -1308,7 +1317,9 @@ to the kernels that `#include` it, and keep it in sync with a test: ```python def test_pusher_args_header_is_up_to_date(): - generated = xp.cuda.write_cuda_header(tmp_path / "pusher_args.cuh", [MarkerArgs, DomainArgs]) + generated = xp.cuda.write_cuda_header( + tmp_path / "pusher_args.cuh", [MarkerArgs, DomainArgs] + ) assert Path("kernels/pusher_args.cuh").read_text() == generated ``` @@ -1328,6 +1339,7 @@ class MarkerArguments(xp.cuda.CudaStructArguments): self.n_markers = markers.shape[0] self.pack() + push = xp.cuda.CudaKernel(source, "push", structs=[MarkerArguments.struct]) push(MarkerArguments(markers, valid), dt, n_threads=markers.shape[0]) ``` @@ -1385,6 +1397,7 @@ class Particles(xp.cuda.CudaArguments): self.positions = positions super().__init__(positions, velocities, positions.shape[0]) + kernel(dt, Particles(x, v), n_threads=x.shape[0]) ``` @@ -1646,10 +1659,11 @@ build a catalog by hand. ```python # my_sim/kernels/push/__init__.py -kernel = xp.kernels.Kernel.from_folder(__name__, host_suffix="_pyccel", dispatch="arrays", - compile_host=compile_kernels) -kernel.implementations # ("pyccel", "numpy", "python", "cuda") -kernel.selected() # "pyccel": what a call with host arrays runs now +kernel = xp.kernels.Kernel.from_folder( + __name__, host_suffix="_pyccel", dispatch="arrays", compile_host=compile_kernels +) +kernel.implementations # ("pyccel", "numpy", "python", "cuda") +kernel.selected() # "pyccel": what a call with host arrays runs now ``` The kernel of one kernel folder `package` (its dotted name, `__name__` in its @@ -1669,11 +1683,13 @@ no host kernel module and `ModuleNotFoundError` if `package` is not a package. ## `kernels.HostImplementations`, `kernels.set_kernel_implementation` ```python -host = xp.kernels.HostImplementations("push", {"pyccel": load_compiled, "numpy": lambda: push_numpy, - "python": lambda: push}) -host(*args) # the default implementation -xp.kernels.set_kernel_implementation("numpy") # every kernel: like xp.set_backend -with xp.kernels.use_kernel_implementation("python"): # like xp.use_backend +host = xp.kernels.HostImplementations( + "push", + {"pyccel": load_compiled, "numpy": lambda: push_numpy, "python": lambda: push}, +) +host(*args) # the default implementation +xp.kernels.set_kernel_implementation("numpy") # every kernel: like xp.set_backend +with xp.kernels.use_kernel_implementation("python"): # like xp.use_backend host(*args) ``` @@ -1695,7 +1711,9 @@ setting is global, not per thread, and applies to host calls only. ## `kernels.CompiledHostKernel` ```python -kernel = xp.kernels.CompiledHostKernel(my_kernels_module, "push", compiler, fallback=push_numpy) +kernel = xp.kernels.CompiledHostKernel( + my_kernels_module, "push", compiler, fallback=push_numpy +) kernel(*args) ``` @@ -1860,7 +1878,9 @@ extern "C" __global__ void find_span_kernel( ``` ```python -find_span = device_function_kernel(BSPLINES_CUH, "int find_span(const double* t, int p, double eta)") +find_span = device_function_kernel( + BSPLINES_CUH, "int find_span(const double* t, int p, double eta)" +) find_span(t, p, eta, spans, eta.size, n_threads=eta.size) ``` @@ -2057,12 +2077,16 @@ same box form a contiguous range, the starting point of tree builds on the GPU. ```python -keys = xp.algorithms.morton_keys(positions, lower, upper, levels) # (n, 2|3) -> (n,) uint64 +keys = xp.algorithms.morton_keys( + positions, lower, upper, levels +) # (n, 2|3) -> (n,) uint64 keys, order, positions = xp.algorithms.sort_by_key(keys, positions) -node = keys >> np.uint64(ndim * (levels - level)) # node index at `level` -cells = xp.algorithms.morton_decode(node, ndim) # its integer coordinates -key = xp.algorithms.morton_encode(ix, iy) # from integer cells -scales = xp.algorithms.morton_scales(lower, upper, levels) # 2**levels / (upper - lower) +node = keys >> np.uint64(ndim * (levels - level)) # node index at `level` +cells = xp.algorithms.morton_decode(node, ndim) # its integer coordinates +key = xp.algorithms.morton_encode(ix, iy) # from integer cells +scales = xp.algorithms.morton_scales( + lower, upper, levels +) # 2**levels / (upper - lower) ``` ```c @@ -2187,10 +2211,10 @@ control flow on array values or indexing. Test the CuPy path: a function ## `petsc.petsc_vec(array, comm=None)` ```python -b_vec = xp.petsc.petsc_vec(b) # b: NumPy or CuPy array, shared, never copied +b_vec = xp.petsc.petsc_vec(b) # b: NumPy or CuPy array, shared, never copied x_vec = xp.petsc.petsc_vec(x) xp.synchronize() -ksp.solve(b_vec, x_vec) # PETSc writes into x +ksp.solve(b_vec, x_vec) # PETSc writes into x xp.synchronize() ``` diff --git a/docs/source/examples/particle-pusher.md b/docs/source/examples/particle-pusher.md index 664b6c3..8d8bf2b 100644 --- a/docs/source/examples/particle-pusher.md +++ b/docs/source/examples/particle-pusher.md @@ -56,7 +56,9 @@ def push(x: "float[:]", v: "float[:]", n: int, dt: float, length: float): ```python # pic/kernels/deposit/deposit_kernels.py -def deposit(x: "float[:]", rho: "float[:]", n: int, dx: float, n_cells: int, weight: float): +def deposit( + x: "float[:]", rho: "float[:]", n: int, dx: float, n_cells: int, weight: float +): for i in range(n): cell = min(int(x[i] / dx), n_cells - 1) rho[cell] += weight / dx @@ -64,8 +66,15 @@ def deposit(x: "float[:]", rho: "float[:]", n: int, dx: float, n_cells: int, wei ```python # pic/kernels/accelerate/accelerate_kernels.py -def accelerate(x: "float[:]", v: "float[:]", e: "float[:]", n: int, dx: float, - n_cells: int, qm_dt: float): +def accelerate( + x: "float[:]", + v: "float[:]", + e: "float[:]", + n: int, + dx: float, + n_cells: int, + qm_dt: float, +): for i in range(n): cell = min(int(x[i] / dx), n_cells - 1) v[i] += qm_dt * e[cell] @@ -122,11 +131,19 @@ class Simulation: def step(self, dt): self.rho[:] = 0.0 - catalog["deposit"](self.x, self.rho, self.n, self.dx, self.n_cells, - self.weight, n_threads=self.n) + catalog["deposit"]( + self.x, + self.rho, + self.n, + self.dx, + self.n_cells, + self.weight, + n_threads=self.n, + ) self.solve_field() - catalog["accelerate"](self.x, self.v, self.e, self.n, self.dx, - self.n_cells, -dt, n_threads=self.n) + catalog["accelerate"]( + self.x, self.v, self.e, self.n, self.dx, self.n_cells, -dt, n_threads=self.n + ) catalog["push"](self.x, self.v, self.n, dt, self.length, n_threads=self.n) def field_energy(self): @@ -142,7 +159,7 @@ from pic.kernels import catalog from pic.simulation import Simulation xp.set_backend("cupy") -print(catalog.summary()) # CUDA kernels: 0 of 3 (missing: accelerate, deposit, push) +print(catalog.summary()) # CUDA kernels: 0 of 3 (missing: accelerate, deposit, push) sim = Simulation() with xp.profiling.count_transfers() as counter: @@ -307,7 +324,9 @@ with xp.profiling.timed_region("100 steps") as timing: for _ in range(100): with xp.profiling.nvtx_range("step"): sim.step(0.05) -print(f"{timing.elapsed / 100 * 1e3:.2f} ms/step, field energy {sim.field_energy():.4e}") +print( + f"{timing.elapsed / 100 * 1e3:.2f} ms/step, field energy {sim.field_energy():.4e}" +) ``` The same script, without `set_backend("cupy")`, runs the Pyccel-compiled (or, diff --git a/docs/source/examples/portable-script.md b/docs/source/examples/portable-script.md index 1e0c26b..9a185ea 100644 --- a/docs/source/examples/portable-script.md +++ b/docs/source/examples/portable-script.md @@ -51,7 +51,9 @@ def main(): args = parser.parse_args() xp.set_backend("cupy" if args.gpu else "numpy") - print(f"backend={xp.get_backend()} devices={xp.cuda.device_count()} cunumpy={xp.__version__}") + print( + f"backend={xp.get_backend()} devices={xp.cuda.device_count()} cunumpy={xp.__version__}" + ) dx = 1.0 / args.n dt = 0.2 * dx**2 # stable for the explicit scheme diff --git a/docs/source/guides/backends.md b/docs/source/guides/backends.md index 05ba108..83e05c3 100644 --- a/docs/source/guides/backends.md +++ b/docs/source/guides/backends.md @@ -84,8 +84,8 @@ xp.set_backend("cupy") on_gpu = xp.arange(4) xp.set_backend("numpy") -print(xp.get_backend()) # 'numpy': what xp.* creates now -print(xp.get_array_backend(on_gpu)) # 'cupy': where this array lives +print(xp.get_backend()) # 'numpy': what xp.* creates now +print(xp.get_array_backend(on_gpu)) # 'cupy': where this array lives ``` Two families of functions answer the two questions: diff --git a/docs/source/guides/data-movement.md b/docs/source/guides/data-movement.md index 9a87d21..ad598e0 100644 --- a/docs/source/guides/data-movement.md +++ b/docs/source/guides/data-movement.md @@ -33,10 +33,10 @@ import cunumpy as xp xp.set_backend("cupy") -signal = xp.to_cunumpy(np.load("signal.npy")) # one host-to-device copy +signal = xp.to_cunumpy(np.load("signal.npy")) # one host-to-device copy spectrum = xp.abs(xp.fft.rfft(signal)) ** 2 -peak = int(xp.argmax(spectrum)) # one tiny device-to-host copy -np.save("spectrum.npy", xp.to_numpy(spectrum)) # one device-to-host copy +peak = int(xp.argmax(spectrum)) # one tiny device-to-host copy +np.save("spectrum.npy", xp.to_numpy(spectrum)) # one device-to-host copy ``` Typical boundaries are file I/O, plotting, calls into host-only libraries @@ -47,9 +47,9 @@ An anti-pattern to avoid: ```python for step in range(n_steps): - state = xp.to_cupy(state) # copied up every step + state = xp.to_cupy(state) # copied up every step state = advance(state, dt) - state = xp.to_numpy(state) # and down again + state = xp.to_numpy(state) # and down again if step % 100 == 0: write_output(state) ``` diff --git a/docs/source/guides/execution-helpers.md b/docs/source/guides/execution-helpers.md index 4aa6dc9..819ca22 100644 --- a/docs/source/guides/execution-helpers.md +++ b/docs/source/guides/execution-helpers.md @@ -64,9 +64,10 @@ recv = xp.zeros_like(send) send_staging = xp.mpi.MPIStaging(send.shape, send.dtype) recv_staging = xp.mpi.MPIStaging(recv.shape, recv.dtype) -with send_staging.buffer(send, cuda_aware=False) as sendbuf, recv_staging.buffer( - recv, send=False, recv=True, cuda_aware=False -) as recvbuf: +with ( + send_staging.buffer(send, cuda_aware=False) as sendbuf, + recv_staging.buffer(recv, send=False, recv=True, cuda_aware=False) as recvbuf, +): comm.Sendrecv(sendbuf, dest=0, recvbuf=recvbuf, source=0) ``` diff --git a/docs/source/guides/gpu-devices.md b/docs/source/guides/gpu-devices.md index 7cd0db8..0edc448 100644 --- a/docs/source/guides/gpu-devices.md +++ b/docs/source/guides/gpu-devices.md @@ -22,7 +22,7 @@ On a multi-GPU workstation, pick the device before creating arrays: ```python xp.set_backend("cupy") print("GPUs:", xp.cuda.device_count()) -xp.cuda.set_device(1) # arrays created from now on live on GPU 1 +xp.cuda.set_device(1) # arrays created from now on live on GPU 1 values = xp.zeros(10**6) ``` @@ -81,10 +81,10 @@ Work on one stream runs in order; work on different streams may overlap. ```python with xp.cuda.stream() as s: - device = xp.to_cupy(pinned_host) # copy and compute queued on s + device = xp.to_cupy(pinned_host) # copy and compute queued on s result = xp.fft.fft(device) -xp.synchronize() # wait before using the result elsewhere +xp.synchronize() # wait before using the result elsewhere ``` On NumPy, `stream()` yields `None`, so do not call methods on the yielded value diff --git a/docs/source/guides/mpi.md b/docs/source/guides/mpi.md index 2173bba..afa80c7 100644 --- a/docs/source/guides/mpi.md +++ b/docs/source/guides/mpi.md @@ -10,19 +10,23 @@ CuNumpy provides one helper per step. import cunumpy as xp xp.set_backend("cupy") -xp.cuda.bind_local_device() # 1. pick this rank's GPU, before MPI_Init +xp.cuda.bind_local_device() # 1. pick this rank's GPU, before MPI_Init -from mpi4py import MPI # 2. MPI_Init happens here +from mpi4py import MPI # 2. MPI_Init happens here -xp.mpi.require_cuda_aware_mpi() # 3. fail clearly if MPI cannot take GPU buffers +xp.mpi.require_cuda_aware_mpi() # 3. fail clearly if MPI cannot take GPU buffers comm = MPI.COMM_WORLD send = xp.full(1000, comm.rank, dtype=xp.float64) recv = xp.empty_like(send) xp.mpi.synchronize_for_mpi(send, recv) # 4. before every MPI call on device buffers -comm.Sendrecv(send, dest=(comm.rank + 1) % comm.size, - recvbuf=recv, source=(comm.rank - 1) % comm.size) +comm.Sendrecv( + send, + dest=(comm.rank + 1) % comm.size, + recvbuf=recv, + source=(comm.rank - 1) % comm.size, +) ``` The same file runs on the NumPy backend: `bind_local_device()` returns `None`, @@ -116,7 +120,10 @@ does the right thing for each case, so the MPI call is written once: ```python xp.mpi.mpi_is_cuda_aware(comm) # once at startup; the answer is remembered -with xp.mpi.mpi_buffer(send_r) as sendbuf, xp.mpi.mpi_buffer(recv_l, send=False, recv=True) as recvbuf: +with ( + xp.mpi.mpi_buffer(send_r) as sendbuf, + xp.mpi.mpi_buffer(recv_l, send=False, recv=True) as recvbuf, +): comm.Sendrecv(sendbuf, dest=right, recvbuf=recvbuf, source=left) ``` @@ -163,6 +170,7 @@ Print the binding once at start-up to catch mapping errors early: ```python device = xp.cuda.bind_local_device() from mpi4py import MPI + print(f"rank {MPI.COMM_WORLD.rank}: local rank {xp.mpi.local_rank()}, device {device}") ``` diff --git a/docs/source/guides/profiling.md b/docs/source/guides/profiling.md index 7b9e4b3..77463f0 100644 --- a/docs/source/guides/profiling.md +++ b/docs/source/guides/profiling.md @@ -25,7 +25,7 @@ microseconds regardless of the problem size: ```python start = time.perf_counter() -solve(field) # only queues the kernels +solve(field) # only queues the kernels elapsed = time.perf_counter() - start # wrong on the GPU ``` @@ -48,8 +48,7 @@ with xp.profiling.nvtx_range("push markers"): @xp.profiling.nvtx_range("time step") -def step(state, dt): - ... +def step(state, dt): ... ``` On CuPy this pushes an NVTX range, so the region appears as a labelled bar on diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index 88e14bc..4d6f84a 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -95,16 +95,16 @@ rho = xp.zeros(n_local) phi = xp.zeros(n_local) rho_vec, phi_vec = xp.petsc.petsc_vec(rho), xp.petsc.petsc_vec(phi) -A = assemble_laplacian() # a PETSc Mat -A.setType("aijcusparse") # keep the matrix on the GPU too +A = assemble_laplacian() # a PETSc Mat +A.setType("aijcusparse") # keep the matrix on the GPU too ksp = PETSc.KSP().create() ksp.setOperators(A) for step in range(n_steps): deposit(rho) - xp.synchronize() # CuPy finished writing rho + xp.synchronize() # CuPy finished writing rho ksp.solve(rho_vec, phi_vec) - xp.synchronize() # PETSc finished writing phi + xp.synchronize() # PETSc finished writing phi gather(phi) ``` diff --git a/docs/source/index.md b/docs/source/index.md index 86c7b8a..801c4d9 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -10,7 +10,7 @@ reference. ```python import cunumpy as xp -xp.set_backend("cupy") # falls back to NumPy without a usable GPU +xp.set_backend("cupy") # falls back to NumPy without a usable GPU values = xp.arange(1_000) print(xp.sum(values**2), xp.get_backend()) ``` diff --git a/docs/source/kernels/accumulation.md b/docs/source/kernels/accumulation.md index 8ebcc0b..bf0953c 100644 --- a/docs/source/kernels/accumulation.md +++ b/docs/source/kernels/accumulation.md @@ -91,7 +91,7 @@ w = xp.to_cunumpy(np.full(10_000, 1e-4)) rho.zero() deposit(x, w, rho.device, x.size, rho_host.size, 1.0 / 64, n_threads=x.size) -rho.to_host() # rho_host now holds the charge density, on both backends +rho.to_host() # rho_host now holds the charge density, on both backends # the host library continues with rho_host, e.g. an MPI Allreduce ``` @@ -107,9 +107,9 @@ cell. The last step is `xp.algorithms.segment_sum(values, keys, n_segments)`, on backend: ```python -cell = ix + nx * (iy + ny * iz) # (n_particles,), -1 for outside -weights = compute_weights(markers) # (n_particles, 8), one per corner -rho_cells = xp.algorithms.segment_sum(weights, cell, nx * ny * nz) # (n_cells, 8) +cell = ix + nx * (iy + ny * iz) # (n_particles,), -1 for outside +weights = compute_weights(markers) # (n_particles, 8), one per corner +rho_cells = xp.algorithms.segment_sum(weights, cell, nx * ny * nz) # (n_cells, 8) ``` Negative keys drop the value; a 2D `values` is summed column by column. Measure diff --git a/docs/source/kernels/arguments.md b/docs/source/kernels/arguments.md index 2463f4b..5e4e80f 100644 --- a/docs/source/kernels/arguments.md +++ b/docs/source/kernels/arguments.md @@ -30,8 +30,12 @@ import cunumpy as xp class DeviceParticles(xp.cuda.CudaArguments): def __init__(self, positions, velocities): - self.positions = xp.as_device_array(positions, np.float64, ndim=2, name="positions") - self.velocities = xp.as_device_array(velocities, np.float64, ndim=2, name="velocities") + self.positions = xp.as_device_array( + positions, np.float64, ndim=2, name="positions" + ) + self.velocities = xp.as_device_array( + velocities, np.float64, ndim=2, name="velocities" + ) super().__init__(self.positions, self.velocities, self.positions.shape[0]) @@ -108,12 +112,15 @@ Particles = xp.cuda.CudaStruct( [("x", "double*"), ("v", "double*"), ("n", "long long"), ("charge", "double")], ) -PUSH = Particles.declaration + r""" +PUSH = ( + Particles.declaration + + r""" extern "C" __global__ void push(Particles p, double dt) { long long i = blockDim.x * (long long)blockIdx.x + threadIdx.x; if (i < p.n) p.x[i] += dt * p.charge * p.v[i]; } """ +) push = xp.cuda.CudaKernel(PUSH, "push", structs=[Particles]) value = Particles(x=x, v=v, n=x.size, charge=-1.0) @@ -155,14 +162,18 @@ class CudaMarkerArguments(xp.cuda.CudaStructArguments): def __init__(self, markers, valid_mks, weight_idx): self.markers = xp.as_device_array(markers, np.float64, ndim=2, name="markers") - self.valid_mks = xp.as_device_array(valid_mks, np.bool_, ndim=1, name="valid_mks") + self.valid_mks = xp.as_device_array( + valid_mks, np.bool_, ndim=1, name="valid_mks" + ) self.n_markers, self.n_cols = self.markers.shape self.weight_idx = weight_idx self.pack() xp.cuda.write_cuda_header("kernels/marker_args.cuh", [CudaMarkerArguments.struct]) -push = xp.cuda.CudaKernel.from_file("kernels/push_cuda.cu", structs=[CudaMarkerArguments.struct]) +push = xp.cuda.CudaKernel.from_file( + "kernels/push_cuda.cu", structs=[CudaMarkerArguments.struct] +) args = CudaMarkerArguments(markers, valid_mks, weight_idx=6) push(args, dt, n_threads=args.n_markers) @@ -254,7 +265,7 @@ error. One GPU test per struct catches it: ```python def test_marker_args_layout(): CudaMarkerArguments.struct.verify_layout() # the Python declaration - CudaMarkerArguments.struct.verify_layout( # the committed header + CudaMarkerArguments.struct.verify_layout( # the committed header "marker_args.cuh", include_dirs=["kernels"] ) ``` @@ -267,8 +278,7 @@ the arguments on host and device: ```python class MarkerArguments: - def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): - ... + def __init__(self, markers: "float[:, :]", n_markers: int, valid: "bool[:]"): ... MarkerArgs = xp.cuda.CudaStruct.from_signature(MarkerArguments.__init__, "MarkerArgs") @@ -333,7 +343,9 @@ from pathlib import Path def test_marker_args_header_is_up_to_date(tmp_path): - generated = xp.cuda.write_cuda_header(tmp_path / "marker_args.cuh", [MarkerArgs, DomainArgs]) + generated = xp.cuda.write_cuda_header( + tmp_path / "marker_args.cuh", [MarkerArgs, DomainArgs] + ) assert Path("kernels/marker_args.cuh").read_text() == generated ``` diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index fa444f8..fd93217 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/docs/source/kernels/cuda-kernel.md @@ -65,7 +65,7 @@ A wrong argument count, dtype or layout raises before anything is launched, with the parameter name in the message: ```python -axpy(2, x, y, x.size, n_threads=x.size) # fine: 2 is cast to double +axpy(2, x, y, x.size, n_threads=x.size) # fine: 2 is cast to double axpy(2.0, xp.to_numpy(x), y, x.size, n_threads=x.size) # TypeError: argument 1 (double* x) must be a CuPy array, got ndarray; arrays are never copied to the device ``` @@ -190,7 +190,7 @@ void scale_column(Array2D a, long long column, double factor) { scale_column = xp.cuda.CudaKernel(SCALE_COLUMN, "scale_column") markers = xp.zeros((1000, 7)) -view = markers[::2, 1:5] # non-contiguous view is fine +view = markers[::2, 1:5] # non-contiguous view is fine scale_column(view, 1, 10.0, n_threads=view.shape[0]) ``` diff --git a/docs/source/kernels/debugging.md b/docs/source/kernels/debugging.md index 16e363b..73753ba 100644 --- a/docs/source/kernels/debugging.md +++ b/docs/source/kernels/debugging.md @@ -16,10 +16,12 @@ CUNUMPY_CUDA_DEBUG=1 python simulate.py # whole process ``` ```python -xp.cuda.set_cuda_debug(True) # globally, from now on -with xp.cuda.cuda_debug(): # for a block +xp.cuda.set_cuda_debug(True) # globally, from now on +with xp.cuda.cuda_debug(): # for a block ... -xp.cuda.CudaKernel(src, "push", debug=True) # one kernel, regardless of the global setting +xp.cuda.CudaKernel( + src, "push", debug=True +) # one kernel, regardless of the global setting ``` In debug mode a `CudaKernel`: diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md index 1bbec2a..98430c5 100644 --- a/docs/source/kernels/dispatch.md +++ b/docs/source/kernels/dispatch.md @@ -172,10 +172,10 @@ none is. Device arrays run the CUDA kernel. To choose, use the same pattern as for the array backend: ```python -xp.kernels.set_kernel_implementation("numpy") # like xp.set_backend +xp.kernels.set_kernel_implementation("numpy") # like xp.set_backend with xp.kernels.use_kernel_implementation("numba"): # like xp.use_backend push(positions, velocities, dt) -xp.kernels.set_kernel_implementation(None) # back to the default +xp.kernels.set_kernel_implementation(None) # back to the default ``` or `CUNUMPY_KERNEL_IMPLEMENTATION=numpy` for a whole run (read at import, like @@ -207,7 +207,7 @@ def compile_kernels(module): catalog = xp.kernels.KernelCatalog.from_package( __name__, - host_suffix="_pyccel", # push/push_pyccel.py next to push/push_cuda.cu + host_suffix="_pyccel", # push/push_pyccel.py next to push/push_cuda.cu compile_host=compile_kernels, host_fallback=NUMPY_VERSIONS, ) @@ -282,8 +282,8 @@ function's, and lists every kernel that differs, e.g. print(catalog.summary()) # CUDA kernels: 2 of 3 (missing: sort) -catalog.without_cuda # ['sort'] -catalog.with_cuda # ['deposit', 'push'] +catalog.without_cuda # ['sort'] +catalog.with_cuda # ['deposit', 'push'] ``` `summary()` is handy for a `--status` command line flag or the start-up log. @@ -292,7 +292,7 @@ At setup, on the GPU backend, compile everything at once: ```python if xp.cupy_backend: - catalog.compile_all(jobs=8) # threads; jobs=None uses all CPUs + catalog.compile_all(jobs=8) # threads; jobs=None uses all CPUs ``` All kernels are compiled even if one fails; the first error is raised diff --git a/docs/source/kernels/pyccel-kernel.md b/docs/source/kernels/pyccel-kernel.md index 8ce73dc..918fcf7 100644 --- a/docs/source/kernels/pyccel-kernel.md +++ b/docs/source/kernels/pyccel-kernel.md @@ -44,10 +44,10 @@ copies back *every* converted array, which is correct but doubles the transfers. `outputs` lists the arguments that may be written: ```python -xp.kernels.PyccelKernel(smooth, outputs=(1,)) # positional argument 1 -xp.kernels.PyccelKernel(update, outputs=(0, -1)) # first and last argument -xp.kernels.PyccelKernel(solve, outputs=("out",)) # solve(a, b, out=out) -xp.kernels.PyccelKernel(norm, outputs=()) # writes nothing +xp.kernels.PyccelKernel(smooth, outputs=(1,)) # positional argument 1 +xp.kernels.PyccelKernel(update, outputs=(0, -1)) # first and last argument +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 diff --git a/docs/source/kernels/testing.md b/docs/source/kernels/testing.md index 25b2426..8f279b5 100644 --- a/docs/source/kernels/testing.md +++ b/docs/source/kernels/testing.md @@ -45,7 +45,7 @@ from cunumpy.kernel_testing import backend # noqa: F401 ```python def test_energy_is_conserved(backend): - state = make_state() # arrays land on the active backend + state = make_state() # arrays land on the active backend e0 = energy(state) for _ in range(100): step(state, 1e-3) @@ -187,10 +187,20 @@ def test_gather_cuda_arithmetic(): rng = np.random.default_rng(0) positions, field = rng.random((500, 2)), rng.normal(size=(17, 9, 2)) expected, result = np.zeros((500, 2)), np.zeros((500, 2)) - catalog["gather"].host_kernel(positions, field, expected, 0.0, 0.0, 0.06, 0.11, 17, 9) + catalog["gather"].host_kernel( + positions, field, expected, 0.0, 0.0, 0.06, 0.11, 17, 9 + ) emulate_cuda_kernel( catalog["gather"].cuda_kernel, - positions, field, result, 0.0, 0.0, 0.06, 0.11, 17, 9, + positions, + field, + result, + 0.0, + 0.0, + 0.06, + 0.11, + 17, + 9, n_threads=500, ) np.testing.assert_allclose(result, expected, rtol=1e-12, atol=1e-14) @@ -295,7 +305,7 @@ real CuPy is importable. def test_time_step_has_no_transfers(): with xp.use_backend("cupy"): state = make_state() - step(state, 1e-3) # warm-up: compilation, allocations + step(state, 1e-3) # warm-up: compilation, allocations with xp.profiling.assert_no_transfers(): step(state, 1e-3) ``` diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index 48e74e4..e76a5fd 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -31,7 +31,7 @@ or in the program: ```python xp.set_backend("cupy") a = xp.linspace(0.0, 1.0, 5) # now a cupy.ndarray on the GPU -print(xp.get_backend()) # 'cupy', or 'numpy' if no usable GPU +print(xp.get_backend()) # 'cupy', or 'numpy' if no usable GPU ``` If CuPy cannot be used, CuNumpy falls back to NumPy, so the same script runs on @@ -56,7 +56,7 @@ Existing arrays never move by themselves. Convert at boundaries such as file output or plotting: ```python -host = xp.to_numpy(a) # always a NumPy array on the host +host = xp.to_numpy(a) # always a NumPy array on the host device = xp.to_cupy(host) # a CuPy array (needs a GPU) active = xp.to_cunumpy(host) # whatever backend is active ``` diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 9440a50..e6f6ecb 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -66,7 +66,7 @@ def require_version(minimum: str) -> None: if _version_key(__version__) < _version_key(minimum): raise ImportError( f"cunumpy {minimum} or newer is required, but {__version__} is " - "installed: pip install --upgrade cunumpy" + "installed: pip install --upgrade cunumpy", ) diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index 87538c4..5f55b5a 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -47,7 +47,7 @@ def segment_boundaries(sorted_keys: Any) -> tuple[Any, Any, Any]: if keys.size == 0: return keys.copy(), xpm.empty(0, dtype=np.int64), xpm.empty(0, dtype=np.int64) starts = xpm.concatenate( - (xpm.zeros(1, dtype=np.int64), xpm.nonzero(keys[1:] != keys[:-1])[0] + 1) + (xpm.zeros(1, dtype=np.int64), xpm.nonzero(keys[1:] != keys[:-1])[0] + 1), ).astype(np.int64, copy=False) stops = xpm.concatenate((starts[1:], xpm.full(1, keys.size, dtype=np.int64))) return keys[starts], starts, stops @@ -65,7 +65,9 @@ def cell_offsets(sorted_cells: Any, n_cells: int) -> Any: if bool(((cells < 0) | (cells >= n_cells)).any()): raise ValueError("cell IDs must be in [0, n_cells)") return xpm.searchsorted( - cells, xpm.arange(n_cells + 1, dtype=np.int64), side="left" + cells, + xpm.arange(n_cells + 1, dtype=np.int64), + side="left", ).astype(np.int64, copy=False) @@ -142,7 +144,7 @@ def sum(self, values: Any, *, out: Any = None) -> Any: if values.device.id != self._device or cp.cuda.Device().id != self._device: raise ValueError( - "segment plan and values require their CUDA device current" + "segment plan and values require their CUDA device current", ) dtype = values.dtype if values.dtype.kind in "fc" else np.dtype(np.float64) if self._device is not None and np.dtype(dtype).name not in ( @@ -160,7 +162,7 @@ def sum(self, values: Any, *, out: Any = None) -> Any: assert_same_backend(self._keys, out) if out.shape != shape or out.dtype != dtype or not out.flags.c_contiguous: raise ValueError( - "out must have the exact shape/dtype and be C-contiguous" + "out must have the exact shape/dtype and be C-contiguous", ) if not getattr(out.flags, "writeable", True): raise ValueError("out must be writable") @@ -245,7 +247,7 @@ def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: for array in arrays: if array.shape[:1] != keys.shape: raise ValueError( - f"every array needs {keys.shape[0]} rows, got shape {array.shape}" + f"every array needs {keys.shape[0]} rows, got shape {array.shape}", ) order = xpm.argsort(keys, kind="stable").astype(xpm.int64, copy=False) return (keys[order], order, *(array[order] for array in arrays)) diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index f6f66a9..33c0bc2 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -236,7 +236,7 @@ class CudaParameter(NamedTuple): _VIEW = re.compile(r"\bArray([1234])D\s*<((?:[^<>]|complex<[^<>]*>)+?)>") _TOKEN = re.compile( r"Array[1234]D<[^<>]*(?:<[^<>]*>[^<>]*)?>|complex<(?:float|double)>" - r"|[A-Za-z_]\w*|\*|\[\s*\]" + r"|[A-Za-z_]\w*|\*|\[\s*\]", ) @@ -280,7 +280,8 @@ def _strip_comments(source: str) -> str: # bracket includes are system headers and are not tracked, except in the # directories given as ``angle_dirs`` (cunumpy's shipped headers). _INCLUDE = re.compile( - r'^[ \t]*#[ \t]*include[ \t]*(?:"([^"\n]+)"|<([^>\n]+)>)', re.MULTILINE + r'^[ \t]*#[ \t]*include[ \t]*(?:"([^"\n]+)"|<([^>\n]+)>)', + re.MULTILINE, ) @@ -409,7 +410,8 @@ def cuda_kernel_names(source: str) -> list[str]: def _compile_in_threads( - compilers: Mapping[Hashable, Callable[[], Any]], jobs: int | None + compilers: Mapping[Hashable, Callable[[], Any]], + jobs: int | None, ) -> list[Hashable]: """Run the `compilers` (name -> compile function), `jobs` at a time. @@ -447,7 +449,8 @@ def _compile_in_threads( def _parse_parameter( - text: str, structs: dict[str, CudaStruct] | None = None + text: str, + structs: dict[str, CudaStruct] | None = None, ) -> CudaParameter: text = _COMPLEX.sub(lambda m: f"complex<{m.group(1)}>", text) text = _VIEW.sub(_normalize_view, text) @@ -465,7 +468,7 @@ def _parse_parameter( if pointers: raise ValueError( f"cannot check the kernel parameter {text.strip()!r}: structs can " - "only be passed by value" + "only be passed by value", ) struct = structs[ctype] return CudaParameter(name, ctype, struct.dtype, False, struct) @@ -475,7 +478,7 @@ def _parse_parameter( if pointers or element not in _CTYPES: raise ValueError( f"cannot check the kernel parameter {text.strip()!r}: array views " - f"take a scalar element type and are passed by value" + f"take a scalar element type and are passed by value", ) ndim = int(view.group(1)) return CudaParameter(name, ctype, np.dtype(_CTYPES[element]), False, None, ndim) @@ -484,7 +487,7 @@ def _parse_parameter( if pointers > 1 or ctype not in _CTYPES: raise ValueError( f"cannot check the kernel parameter {text.strip()!r}: unsupported type " - f"{ctype + '*' * pointers!r}" + f"{ctype + '*' * pointers!r}", ) return CudaParameter(name, ctype, np.dtype(_CTYPES[ctype]), pointers == 1) @@ -581,12 +584,14 @@ def parse_cuda_signature( if template_args is None or len(template_args) != len(template_params): raise ValueError( f"{name!r} is a template with {len(template_params)} parameters; " - f"pass them as template_args" + f"pass them as template_args", ) for words, value in zip(template_params, template_args): if words[0] in ("typename", "class"): params = re.sub( - r"\b" + re.escape(words[-1]) + r"\b", _template_arg(value), params + r"\b" + re.escape(words[-1]) + r"\b", + _template_arg(value), + params, ) elif template_args: raise ValueError(f"{name!r} is not a template, but template_args were given") @@ -620,16 +625,17 @@ def _check_device_array(param: CudaParameter, index: int, value: Any) -> None: """Raise unless `value` is a device array of the declared dtype on the current device.""" # checked on the class: on the instance, CuPy builds the whole interface dict if not hasattr(type(value), "__cuda_array_interface__") and not hasattr( - value, "__cuda_array_interface__" + value, + "__cuda_array_interface__", ): raise TypeError( f"{_describe(param, index)} must be a CuPy array, got " - f"{type(value).__name__}; arrays are never copied to the device" + f"{type(value).__name__}; arrays are never copied to the device", ) if param.dtype is not None and value.dtype != param.dtype: raise TypeError( f"{_describe(param, index)} must have dtype {param.dtype}, got " - f"{value.dtype}" + f"{value.dtype}", ) # an array on another GPU: the kernel would read a foreign address, which # neither cupy.RawKernel nor the struct packing notices @@ -640,7 +646,7 @@ def _check_device_array(param: CudaParameter, index: int, value: Any) -> None: raise ValueError( f"{_describe(param, index)} is on CUDA device {device_id}, but the " f"current device is {current}; kernels only take arrays of the " - "current device (see cunumpy.cuda.bind_local_device)" + "current device (see cunumpy.cuda.bind_local_device)", ) @@ -656,7 +662,7 @@ def check(value: Any) -> Any: raise TypeError( f"{_describe(param, index)} must be C-contiguous: a non-contiguous " f"view (e.g. a[:, 0:3]) would be read as a flat buffer; use " - f"cupy.ascontiguousarray or cunumpy.as_device_array" + f"cupy.ascontiguousarray or cunumpy.as_device_array", ) return value @@ -676,14 +682,14 @@ def check(value: Any) -> Any: _check_device_array(param, index, value) if value.ndim != ndim: raise TypeError( - f"{_describe(param, index)} must be a {ndim}D array, got {value.ndim}D" + f"{_describe(param, index)} must be a {ndim}D array, got {value.ndim}D", ) itemsize = value.dtype.itemsize strides = [s // itemsize for s in value.strides] if any(s * itemsize != stride for s, stride in zip(strides, value.strides)): raise TypeError( f"{_describe(param, index)}: strides {tuple(value.strides)} are not " - f"multiples of the element size {itemsize}" + f"multiples of the element size {itemsize}", ) packed = np.zeros((), dtype=dtype) packed["data"] = value.data.ptr @@ -708,14 +714,14 @@ def cast_int(value: int) -> Any: if not low <= value <= high: raise OverflowError( f"{_describe(param, index)}: {value} is out of range " - f"[{low}, {high}]" + f"[{low}, {high}]", ) return scalar_type(value) if kind in "fc": return scalar_type(value) raise TypeError( f"{_describe(param, index)} cannot take a value of type " - f"{type(value).__name__}" + f"{type(value).__name__}", ) def check(value: Any) -> Any: @@ -738,7 +744,7 @@ def check(value: Any) -> Any: return scalar_type(value) raise TypeError( f"{_describe(param, index)} cannot take a {value.dtype} scalar " - f"without losing information" + f"without losing information", ) # bool is a subclass of int, but must not be cast like one is_bool = isinstance(value, bool) @@ -752,7 +758,7 @@ def check(value: Any) -> Any: return scalar_type(value) raise TypeError( f"{_describe(param, index)} cannot take a value of type " - f"{value_type.__name__}" + f"{value_type.__name__}", ) return check @@ -767,7 +773,7 @@ def check(value: Any) -> Any: return value raise TypeError( f"{_describe(param, index)} must be a value of struct {param.ctype} " - f"(created with that CudaStruct), got {type(value).__name__}" + f"(created with that CudaStruct), got {type(value).__name__}", ) return check @@ -808,7 +814,7 @@ def _field_dtype(field: CudaParameter) -> np.dtype: } _PYCCEL_ANNOTATION = re.compile( r"^(?:(?:typing\.)?Final\s*\[\s*)?(?:const\s+)?(?P\w+)\s*" - r"(?:\[(?P[\s:,]*)\])?\s*\]?$" + r"(?:\[(?P[\s:,]*)\])?\s*\]?$", ) @@ -834,7 +840,7 @@ def _pyccel_ctype(annotation: Any, scalars: Mapping[str, str], what: str) -> str raise ValueError(f"{what}: unsupported annotation {annotation!r}") if scalar not in scalars: raise ValueError( - f"{what}: unsupported scalar type {scalar!r} in {annotation!r}" + f"{what}: unsupported scalar type {scalar!r} in {annotation!r}", ) ctype = scalars[scalar] if ndim == 0: @@ -850,7 +856,9 @@ def _header_guard(name: str) -> str: def _header_source( - structs: Sequence[CudaStruct], guard: str, includes: Iterable[str] + structs: Sequence[CudaStruct], + guard: str, + includes: Iterable[str], ) -> str: includes = list(includes) if any(f.view_ndim is not None for s in structs for f in s.fields): @@ -1239,7 +1247,9 @@ def test_header_is_up_to_date(): The header source. """ source = _header_source( - (self,), guard or _header_guard(f"{self._name}_cuh"), includes + (self,), + guard or _header_guard(f"{self._name}_cuh"), + includes, ) if path is not None: Path(path).write_text(source) @@ -1258,7 +1268,9 @@ def check_source(self, source: str) -> None: """ code = _strip_comments(source) match = re.search( - r"struct\s+" + re.escape(self._name) + r"\s*\{(.*?)\}", code, re.DOTALL + r"struct\s+" + re.escape(self._name) + r"\s*\{(.*?)\}", + code, + re.DOTALL, ) if match is None: return @@ -1267,13 +1279,13 @@ def check_source(self, source: str) -> None: found = [_parse_parameter(m) for m in members] except ValueError as exc: raise ValueError( - f"cannot compare the definition of struct {self._name!r}: {exc}" + f"cannot compare the definition of struct {self._name!r}: {exc}", ) from None key = [(f.name, f.ctype, f.pointer) for f in self._fields] if [(f.name, f.ctype, f.pointer) for f in found] != key: raise ValueError( f"the definition of struct {self._name!r} in the CUDA source does not " - f"match its CudaStruct:\n{self.declaration}" + f"match its CudaStruct:\n{self.declaration}", ) def layout_source(self, include: str | None = None) -> str: @@ -1307,7 +1319,7 @@ def layout_source(self, include: str | None = None) -> str: f" out[0] = sizeof({self._name});\n" f" out[1] = alignof({self._name});\n" f"{offsets}\n" - "}\n" + "}\n", ) return "\n".join(lines) @@ -1381,7 +1393,7 @@ def verify_layout( if differences: raise ValueError( f"the compiled layout of struct {self._name!r} differs from its " - "CudaStruct dtype:\n " + "\n ".join(differences) + "CudaStruct dtype:\n " + "\n ".join(differences), ) return layout @@ -1404,7 +1416,7 @@ def __call__(self, **values: Any) -> CudaStructValue: if missing or unknown: raise TypeError( f"struct {self._name}: missing fields {missing}, unknown fields " - f"{unknown}" + f"{unknown}", ) packed = np.zeros((), dtype=self._dtype) for field in self._fields: @@ -1543,7 +1555,7 @@ def __init_subclass__(cls, **kwargs: Any) -> None: return # an intermediate base class, or a subclass of a complete one if not (has_name and has_fields): raise TypeError( - f"{cls.__qualname__} must define both struct_name and fields" + f"{cls.__qualname__} must define both struct_name and fields", ) cls.struct = CudaStruct(cls.struct_name, cls.fields) @@ -1566,7 +1578,7 @@ def pack(self) -> None: struct = getattr(type(self), "struct", None) if struct is None: raise TypeError( - f"{type(self).__qualname__} does not define struct_name and fields" + f"{type(self).__qualname__} does not define struct_name and fields", ) values = self._field_values(struct) self._struct_value = struct(**values) @@ -1581,7 +1593,7 @@ def _field_values(self, struct: CudaStruct) -> dict[str, Any]: except AttributeError: raise AttributeError( f"{type(self).__qualname__} has no attribute {field.name!r} " - f"for the field of struct {struct.name}" + f"for the field of struct {struct.name}", ) from None return values @@ -1648,7 +1660,8 @@ def _field_state(struct: CudaStruct, values: Mapping[str, Any]) -> tuple[Any, .. def _is_device_array(value: Any) -> bool: return hasattr(type(value), "__cuda_array_interface__") or hasattr( - value, "__cuda_array_interface__" + value, + "__cuda_array_interface__", ) @@ -1726,7 +1739,7 @@ def _host_field_names(self) -> tuple[str, ...]: struct = getattr(type(self), "struct", None) if struct is None: raise TypeError( - f"{type(self).__qualname__} does not define struct_name and fields" + f"{type(self).__qualname__} does not define struct_name and fields", ) return tuple(field.name for field in struct.fields) @@ -1736,7 +1749,7 @@ def __host_args__(self) -> Any: if host_class is None: raise TypeError( f"{type(self).__qualname__}.host_class is not set: the class of " - "the host argument object (e.g. the pyccel class) is required" + "the host argument object (e.g. the pyccel class) is required", ) names = self._host_field_names() values = [getattr(self, name) for name in names] @@ -1754,7 +1767,7 @@ def __host_args__(self) -> Any: f"{type(self).__qualname__}.{name} is a device array: there " "is no host form on the CuPy backend. Call the CUDA kernel, " "or set host_copies = True for a read-only host evaluation " - "from host copies" + "from host copies", ) from .xp import to_numpy @@ -1796,7 +1809,7 @@ def _first_array_length(args: tuple[Any, ...]) -> int: if shape is not None and len(shape) > 0 and hasattr(arg, "dtype"): return int(shape[0]) raise TypeError( - "n_threads_from='first_array' needs an array argument; pass n_threads" + "n_threads_from='first_array' needs an array argument; pass n_threads", ) @@ -1932,7 +1945,10 @@ def __init__( self._template_args = None if template_args is None else tuple(template_args) self._signature = ( parse_cuda_signature( - source, name, structs=self._structs, template_args=self._template_args + source, + name, + structs=self._structs, + template_args=self._template_args, ) if check_signature else None @@ -1973,7 +1989,7 @@ def from_file( if name is None: if not path.name.endswith(suffix): raise ValueError( - f"{path.name} does not end with {suffix!r}; pass the kernel name" + f"{path.name} does not end with {suffix!r}; pass the kernel name", ) name = path.name[: -len(suffix)] include_dirs = (path.parent, *kwargs.pop("include_dirs", ())) @@ -2025,7 +2041,7 @@ def _check_block(block: tuple[int, ...]) -> tuple[int, ...]: if math.prod(block) > _MAX_THREADS_PER_BLOCK: raise ValueError( f"a block has at most {_MAX_THREADS_PER_BLOCK} threads, got " - f"{block} = {math.prod(block)}" + f"{block} = {math.prod(block)}", ) return block @@ -2110,7 +2126,7 @@ def included_headers(self) -> tuple[Path, ...]: self._include_dirs, base_dir=self._source_dir, angle_dirs=(cuda_include_dir(),), - ) + ), ) def compile_options(self) -> tuple[str, ...]: @@ -2160,7 +2176,8 @@ def n_threads_from(self) -> Callable[[tuple[Any, ...]], Any] | None: @n_threads_from.setter def n_threads_from( - self, value: Callable[[tuple[Any, ...]], Any] | str | None + self, + value: Callable[[tuple[Any, ...]], Any] | str | None, ) -> None: if value == "first_array": value = _first_array_length @@ -2223,14 +2240,16 @@ def compile(self) -> Any: if not cupy_available(): raise RuntimeError( f"cannot compile CUDA kernel {self.expression!r}: " - "CuPy is not installed or no GPU is available" + "CuPy is not installed or no GPU is available", ) import cupy as cp options = self.compile_options() if self._template_args is None: self._raw_kernel = cp.RawKernel( - self._source, self._name, options=options + self._source, + self._name, + options=options, ) else: module = cp.RawModule( @@ -2274,7 +2293,7 @@ def prepare_args(self, *args: Any) -> tuple[Any, ...]: if len(values) != len(self._signature): raise TypeError( f"{self._name}() takes {len(self._signature)} arguments after " - f"flattening argument objects, got {len(values)}" + f"flattening argument objects, got {len(values)}", ) return tuple([check(v) for check, v in zip(self._checkers, values)]) @@ -2310,7 +2329,7 @@ def launch_shape( if len(block_shape) != 1: raise ValueError( f"block {block_shape} and n_threads {threads} have different " - "numbers of dimensions" + "numbers of dimensions", ) block_shape = block_shape + (1,) * (len(threads) - 1) grid_shape = tuple(math.ceil(n / b) for n, b in zip(threads, block_shape)) @@ -2394,7 +2413,7 @@ def _opt_in_shared_memory(self, kernel: Any, shared_mem: int) -> None: if shared_mem > limit: raise ValueError( f"kernel {self.expression!r}: shared_mem={shared_mem} bytes exceeds " - f"the {limit} bytes a block may use on this device" + f"the {limit} bytes a block may use on this device", ) kernel.max_dynamic_shared_size_bytes = shared_mem self._shared_mem_opt_in = shared_mem @@ -2410,11 +2429,14 @@ def _check_finite_arrays(self, args: tuple[Any, ...]) -> None: if not bool(cp.isfinite(array).all()): raise RuntimeError( f"kernel {self.expression!r} left a NaN or inf in {label} " - f"(dtype {array.dtype}, shape {tuple(array.shape)})" + f"(dtype {array.dtype}, shape {tuple(array.shape)})", ) def _synchronize_after_launch( - self, stream: Any, grid: tuple[int, ...], block: tuple[int, ...] + self, + stream: Any, + grid: tuple[int, ...], + block: tuple[int, ...], ) -> None: """Wait for the launch and re-raise a CUDA error naming this kernel. @@ -2441,7 +2463,7 @@ def _synchronize_after_launch( raise raise RuntimeError( f"CUDA error after launching kernel {self.expression!r} with " - f"grid {grid} and block {block}: {error}" + f"grid {grid} and block {block}: {error}", ) from error @@ -2491,7 +2513,7 @@ def get(self, *key: Hashable) -> CudaKernel: kernel = self._factory(*key) if not isinstance(kernel, CudaKernel): raise TypeError( - f"the factory must return a CudaKernel, got {type(kernel).__name__}" + f"the factory must return a CudaKernel, got {type(kernel).__name__}", ) self._kernels[key] = kernel return kernel @@ -2508,7 +2530,10 @@ def keys(self) -> list[tuple[Hashable, ...]]: return list(self._kernels) def compile_all( - self, keys: Iterable[Sequence[Hashable]] = (), *, jobs: int | None = 1 + self, + keys: Iterable[Sequence[Hashable]] = (), + *, + jobs: int | None = 1, ) -> None: """Compile the given variants (created if needed) and all existing ones. @@ -2524,5 +2549,6 @@ def compile_all( for key in keys: self.get(*key) _compile_in_threads( - {key: kernel.compile for key, kernel in self._kernels.items()}, jobs + {key: kernel.compile for key, kernel in self._kernels.items()}, + jobs, ) diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index f9316eb..79a094c 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -111,7 +111,9 @@ def memory_info() -> tuple[int, int] | None: def max_shared_memory_per_block( - device: int | None = None, *, opt_in: bool = False + device: int | None = None, + *, + opt_in: bool = False, ) -> int: """Bytes of shared memory a block of a CUDA kernel may use on `device`. diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index 7b719be..4782b77 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -62,7 +62,7 @@ def _on_device(arg: Any) -> bool: return True kind = type(arg) return callable(getattr(kind, "__cuda_args__", None)) and not callable( - getattr(kind, "__host_args__", None) + getattr(kind, "__host_args__", None), ) @@ -208,19 +208,19 @@ def __init__( self._dispatch = dispatch if missing_cuda not in _MISSING_CUDA: raise ValueError( - f"missing_cuda must be one of {_MISSING_CUDA}, got {missing_cuda!r}" + f"missing_cuda must be one of {_MISSING_CUDA}, got {missing_cuda!r}", ) if cuda_kernel is not None and not isinstance(cuda_kernel, CudaKernel): raise TypeError( "cuda_kernel must be a CudaKernel or None, " - f"got {type(cuda_kernel).__name__}" + f"got {type(cuda_kernel).__name__}", ) if not isinstance(host_kernel, PyccelKernel): host_kernel = PyccelKernel(host_kernel, **(host_options or {})) elif host_options: raise ValueError( "host_options are for wrapping a plain callable; configure the " - "given PyccelKernel directly" + "given PyccelKernel directly", ) self._host_kernel = host_kernel self._cuda_kernel = cuda_kernel @@ -316,7 +316,7 @@ def from_folder( name = package.rpartition(".")[2] if not (folder / f"{name}{host_suffix}.py").is_file(): raise FileNotFoundError( - f"kernel folder {folder} has no host kernel {name}{host_suffix}.py" + f"kernel folder {folder} has no host kernel {name}{host_suffix}.py", ) wrapper = f"bind_c_{name}{host_suffix}" if check_name_length and len(wrapper) > FORTRAN_NAME_LIMIT: @@ -348,7 +348,8 @@ def from_folder( for implementation in ("numba", "numpy"): if (folder / f"{name}_{implementation}.py").is_file(): loaders[implementation] = _import_loader( - f"{package}.{name}_{implementation}", name + f"{package}.{name}_{implementation}", + name, ) loaders.update(extra_implementations or {}) host = HostImplementations(name, loaders) @@ -473,7 +474,7 @@ def check_signature(self) -> None: if host != cuda: raise ValueError( f"kernel {self._name!r}: the host kernel takes " - f"({', '.join(host)}), the CUDA kernel ({', '.join(cuda)})" + f"({', '.join(host)}), the CUDA kernel ({', '.join(cuda)})", ) function = self._host_kernel.kernel others: dict[str, Any] = {} @@ -491,7 +492,7 @@ def check_signature(self) -> None: which = "its fallback" if name == "fallback" else f"the {name} version" raise ValueError( f"kernel {self._name!r}: the host kernel takes " - f"({', '.join(host)}), {which} ({', '.join(parameters)})" + f"({', '.join(host)}), {which} ({', '.join(parameters)})", ) @property @@ -548,7 +549,7 @@ def _device_kernel(self) -> PyccelKernel | CudaKernel: "" if self._cuda_path is None else f" (expected {self._cuda_path})" ) raise NotImplementedError( - f"No CUDA version of kernel {self._name!r}{expected}." + f"No CUDA version of kernel {self._name!r}{expected}.", ) if not self._warned: warnings.warn( @@ -622,7 +623,7 @@ def __call__( if n_threads is None and grid is None and kernel.n_threads_from is None: raise ValueError( f"{self._name}: n_threads is required to launch the CUDA kernel " - "(or pass grid, or set cuda_kernel.n_threads_from)" + "(or pass grid, or set cuda_kernel.n_threads_from)", ) return kernel( *args, @@ -856,7 +857,7 @@ def check_signatures(self) -> None: if problems: raise ValueError( "host and CUDA kernels take different parameters:\n " - + "\n ".join(problems) + + "\n ".join(problems), ) def compile_all(self, jobs: int | None = 1) -> list[str]: diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index cda8343..a90f364 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -239,7 +239,7 @@ def _code(kernel: CudaKernel) -> str: _EXTERN_SHARED = re.compile( - r"extern\s+__shared__\s+(?:__align__\(\s*\d+\s*\)\s+)?([\w:<>\s]+?)\s+(\w+)\s*\[\s*\]\s*;" + r"extern\s+__shared__\s+(?:__align__\(\s*\d+\s*\)\s+)?([\w:<>\s]+?)\s+(\w+)\s*\[\s*\]\s*;", ) @@ -249,7 +249,7 @@ def _check_supported(kernel: CudaKernel) -> None: if re.search(pattern, code): raise NotImplementedError( f"kernel {kernel.name!r} uses {what}, which emulation cannot run " - "correctly one thread at a time; test it on a GPU" + "correctly one thread at a time; test it on a GPU", ) @@ -296,13 +296,13 @@ def emulate_cuda_kernel( """ if kernel.signature is None: raise TypeError( - "emulation needs a parsed kernel signature (check_signature=True)" + "emulation needs a parsed kernel signature (check_signature=True)", ) _check_supported(kernel) params = kernel.signature if len(args) != len(params): raise TypeError( - f"kernel {kernel.name!r} takes {len(params)} arguments, got {len(args)}" + f"kernel {kernel.name!r} takes {len(params)} arguments, got {len(args)}", ) compiler = compiler or emulation_compiler() if compiler is None: @@ -318,23 +318,23 @@ def emulate_cuda_kernel( name = f"cunumpy_arg{i}" if param.struct is not None: raise NotImplementedError( - "emulation does not support struct parameters" + "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__}" + 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}" + 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" + f"array, got {value.ndim}D", ) buffer = np.ascontiguousarray(value) path = tmp_path / f"{name}.bin" @@ -348,7 +348,7 @@ def emulate_cuda_kernel( globals_.append(f"static {element}* {name};") inits.append( f" {name} = ({element}*)malloc({max(buffer.nbytes, 1)});\n" - f' cunumpy_read("{path}", {name}, {buffer.nbytes});' + f' cunumpy_read("{path}", {name}, {buffer.nbytes});', ) if param.view_ndim is not None: shape = ", ".join(f"{n}LL" for n in buffer.shape) @@ -357,7 +357,7 @@ def emulate_cuda_kernel( ) globals_.append(f"static {ctype} {name}_view;") inits.append( - f" {name}_view = {ctype}{{{name}, {{{shape}}}, {{{strides}}}}};" + f" {name}_view = {ctype}{{{name}, {{{shape}}}, {{{strides}}}}};", ) call_args.append(f"{name}_view") else: @@ -372,7 +372,8 @@ def emulate_cuda_kernel( 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 + r"\1* \2 = (\1*)cunumpy_dynamic_shared;", + kernel_source, ) coroutines = "__syncthreads" in _code(kernel) if coroutines: @@ -418,13 +419,13 @@ def emulate_cuda_kernel( 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]}" + 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:]}" + f"{ran.stdout[-2000:]}{ran.stderr[-2000:]}", ) for value, buffer, out in arrays: result = np.fromfile(out, dtype=buffer.dtype).reshape(buffer.shape) diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index c1a58ee..804815c 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -66,13 +66,13 @@ def install() -> types.ModuleType: return existing raise RuntimeError( "the real CuPy is already imported; the fake CuPy must be installed " - "before CuPy (set CUNUMPY_FAKE_CUPY=1 or call install() first)" + "before CuPy (set CUNUMPY_FAKE_CUPY=1 or call install() first)", ) xp = sys.modules.get("cunumpy.xp") if xp is not None and getattr(xp, "_CUPY_AVAILABLE_CACHE", None) is not None: raise RuntimeError( "cunumpy has already checked for CuPy; install the fake CuPy before " - "the first backend use (set CUNUMPY_FAKE_CUPY=1 in the environment)" + "the first backend use (set CUNUMPY_FAKE_CUPY=1 in the environment)", ) module = types.ModuleType("cupy") module.__file__ = str(_IMPLEMENTATION) diff --git a/src/cunumpy/_fake_cupy_impl.py b/src/cunumpy/_fake_cupy_impl.py index 682c07d..5a043d0 100644 --- a/src/cunumpy/_fake_cupy_impl.py +++ b/src/cunumpy/_fake_cupy_impl.py @@ -18,7 +18,7 @@ def _err(obj, where=""): return TypeError( f"Unsupported type {type(obj)}{where} (fake CuPy: host arrays/lists are " - "not accepted)" + "not accepted)", ) @@ -34,7 +34,7 @@ def __init__(self, a): def __array__(self, *args, **kwargs): raise TypeError( "Implicit conversion to a NumPy array is not allowed. Please use " - "`.get()` to construct a NumPy array explicitly." + "`.get()` to construct a NumPy array explicitly.", ) def get(self, stream=None, order="C", out=None, blocking=True): @@ -157,7 +157,8 @@ def _wrap(x): if isinstance(x, _HOST): return ndarray(x) if isinstance(x, _np.generic) and not isinstance( - x, (_np.str_, _np.bytes_, _np.void) + x, + (_np.str_, _np.bytes_, _np.void), ): return ndarray(_np.asarray(x)) if isinstance(x, tuple): @@ -329,21 +330,27 @@ def _module_attr(name, src=_np, prefix="cupy"): if name in _SEQUENCE: return _wrap_callable(attr, strict=True, name=f"{prefix}.{name}") return _wrap_callable( - attr, strict=True, name=f"{prefix}.{name}", first_is_data=True + attr, + strict=True, + name=f"{prefix}.{name}", + first_is_data=True, ) def asarray(a, dtype=None, order=None, **kwargs): return ndarray( _np.array( - _unwrap(a), dtype=dtype, order=order or "K", copy=kwargs.pop("copy", None) - ) + _unwrap(a), + dtype=dtype, + order=order or "K", + copy=kwargs.pop("copy", None), + ), ) def array(a, dtype=None, copy=True, order="K", ndmin=0, **kwargs): return ndarray( - _np.array(_unwrap(a), dtype=dtype, copy=copy, order=order, ndmin=ndmin) + _np.array(_unwrap(a), dtype=dtype, copy=copy, order=order, ndmin=ndmin), ) @@ -506,7 +513,9 @@ def __getattr__(self, name): random = types.ModuleType("cupy.random") random.default_rng = lambda seed=None: _RNGProxy(_np.random.default_rng(_unwrap(seed))) random.__getattr__ = lambda attr: _wrap_callable( - getattr(_np.random, attr), strict=False, name=f"cupy.random.{attr}" + getattr(_np.random, attr), + strict=False, + name=f"cupy.random.{attr}", ) for _m in (cuda, cuda.device, cuda.runtime, linalg, fft, random): @@ -532,7 +541,7 @@ def __getattr__(self, name): "is_available", "bool_", "fuse", - } + }, ) diff --git a/src/cunumpy/_fusion.py b/src/cunumpy/_fusion.py index 7c4cd23..2299b4e 100644 --- a/src/cunumpy/_fusion.py +++ b/src/cunumpy/_fusion.py @@ -69,7 +69,9 @@ def typed(a: Any) -> Any: def fuse( - function: F | None = None, *, kernel_name: str | None = None + function: F | None = None, + *, + kernel_name: str | None = None, ) -> F | Callable[[F], F]: """Fuse an elementwise function into one kernel when called with CuPy arrays. diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index c9869e2..a05ca6e 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -108,14 +108,14 @@ def __host_args__(self) -> Any: """The object passed to the host kernel in place of this one.""" raise NotImplementedError( f"{type(self).__name__} does not provide host kernel arguments " - "(implement __host_args__)" + "(implement __host_args__)", ) def __cuda_args__(self) -> tuple[Any, ...]: """The CUDA kernel arguments this object stands for.""" raise NotImplementedError( f"{type(self).__name__} does not provide CUDA kernel arguments " - "(implement __cuda_args__)" + "(implement __cuda_args__)", ) @@ -129,7 +129,8 @@ def _host_args(value: Any) -> Any: def resolve_host_args( - args: Sequence[Any], kwargs: Mapping[str, Any] | None = None + args: Sequence[Any], + kwargs: Mapping[str, Any] | None = None, ) -> tuple[tuple[Any, ...], dict[str, Any]]: """Replace :class:`KernelArguments` objects by their host form. @@ -235,13 +236,13 @@ def __init__( raise TypeError( "outputs must be a sequence of argument indices/names, " f"not a bare {type(outputs).__name__} " - f"(did you mean outputs=({outputs!r},)?)" + f"(did you mean outputs=({outputs!r},)?)", ) for entry in outputs: if not isinstance(entry, (int, str)) or isinstance(entry, bool): raise TypeError( "outputs entries must be argument indices (int) or " - f"names (str), got {entry!r}" + f"names (str), got {entry!r}", ) self._outputs = tuple(outputs) @@ -307,7 +308,7 @@ def _convert_to_numpy( return value_np if hasattr(value, "__dict__") and value.__class__.__module__.startswith( - self._object_modules + self._object_modules, ): # Shallow-copy the object so the caller's instance keeps pointing at # its device arrays; only the copy holds the host views. @@ -356,14 +357,16 @@ def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> N return if hasattr(value, "__dict__") and value.__class__.__module__.startswith( - self._object_modules + self._object_modules, ): seen.add(id(value)) for attr in vars(value).values(): self._collect_host_arrays(attr, found, seen) def _output_host_arrays( - self, args_np: list[Any], kwargs_np: dict[str, Any] + self, + args_np: list[Any], + kwargs_np: dict[str, Any], ) -> set[int]: """Ids of the host arrays reachable from the declared output arguments. @@ -385,7 +388,7 @@ def _output_host_arrays( 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." + "a keyword must be declared by name, not by index.", ) self._collect_host_arrays(args_np[index], found, seen) else: @@ -394,7 +397,7 @@ def _output_host_arrays( 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." + "declared by index, not by name.", ) self._collect_host_arrays(kwargs_np[entry], found, seen) @@ -424,7 +427,7 @@ def _contains_cupy(self, value: Any, seen: set[int] | None = None) -> bool: return any(self._contains_cupy(item, seen) for item in value.values()) if hasattr(value, "__dict__") and value.__class__.__module__.startswith( - self._object_modules + self._object_modules, ): seen.add(id(value)) return any(self._contains_cupy(attr, seen) for attr in vars(value).values()) @@ -521,13 +524,13 @@ def _check_implementation(name: str | None) -> str | None: if name is not None and name not in HOST_IMPLEMENTATIONS: raise ValueError( f"kernel implementation must be one of {HOST_IMPLEMENTATIONS} or None, " - f"got {name!r}" + f"got {name!r}", ) return name _KERNEL_IMPLEMENTATION: str | None = _check_implementation( - os.environ.get("CUNUMPY_KERNEL_IMPLEMENTATION", "").strip().lower() or None + os.environ.get("CUNUMPY_KERNEL_IMPLEMENTATION", "").strip().lower() or None, ) @@ -590,13 +593,15 @@ class HostImplementations: """ def __init__( - self, name: str, loaders: Mapping[str, Callable[[], Callable[..., Any]]] + self, + name: str, + loaders: Mapping[str, Callable[[], Callable[..., Any]]], ) -> None: unknown = set(loaders) - set(HOST_IMPLEMENTATIONS) if unknown: raise ValueError( f"kernel {name!r}: unknown implementations {sorted(unknown)}, " - f"expected names from {HOST_IMPLEMENTATIONS}" + f"expected names from {HOST_IMPLEMENTATIONS}", ) if "python" not in loaders: raise ValueError(f"kernel {name!r}: the 'python' implementation is needed") @@ -644,12 +649,12 @@ def get(self, name: str) -> Callable[..., Any]: if name not in self._loaders: raise LookupError( f"kernel {self.__name__!r} has no {name!r} implementation " - f"(it has {', '.join(self.names)})" + f"(it has {', '.join(self.names)})", ) if name in self.errors: raise LookupError( f"the {name!r} implementation of kernel {self.__name__!r} is " - f"unavailable: {self.errors[name]!r}" + f"unavailable: {self.errors[name]!r}", ) from self.errors[name] try: function = self._loaders[name]() @@ -657,7 +662,7 @@ def get(self, name: str) -> Callable[..., Any]: self.errors[name] = error raise LookupError( f"the {name!r} implementation of kernel {self.__name__!r} is " - f"unavailable: {error!r}" + f"unavailable: {error!r}", ) from error self._loaded[name] = function return function diff --git a/src/cunumpy/_mirror.py b/src/cunumpy/_mirror.py index 89dfcb4..d10294b 100644 --- a/src/cunumpy/_mirror.py +++ b/src/cunumpy/_mirror.py @@ -69,7 +69,7 @@ def _check_host(host: Any) -> np.ndarray: if not isinstance(host, np.ndarray): raise TypeError( "DeviceMirror needs a NumPy array as host buffer, got " - f"{type(host).__name__}" + f"{type(host).__name__}", ) return host @@ -118,7 +118,7 @@ def _check_bound(self) -> None: "the host array was reallocated: DeviceMirror was bound to " f"shape {self._shape} and dtype {self._dtype}, but the host now " f"has shape {self._host.shape} and dtype {self._host.dtype}; " - "call rebind(host) with the new array" + "call rebind(host) with the new array", ) def rebind(self, host: np.ndarray) -> DeviceMirror: diff --git a/src/cunumpy/_morton.py b/src/cunumpy/_morton.py index 54da383..de39881 100644 --- a/src/cunumpy/_morton.py +++ b/src/cunumpy/_morton.py @@ -65,7 +65,7 @@ def _check_levels(ndim: int, levels: int) -> None: if not 1 <= levels <= MAX_MORTON_LEVELS[ndim]: raise ValueError( f"levels must be between 1 and {MAX_MORTON_LEVELS[ndim]} in {ndim}D, " - f"got {levels}" + f"got {levels}", ) @@ -135,7 +135,7 @@ def morton_scales(lower: Sequence[float], upper: Sequence[float], levels: int) - if lower.shape != upper.shape or lower.ndim != 1: raise ValueError( f"lower and upper must be sequences of equal length, got shapes " - f"{lower.shape} and {upper.shape}" + f"{lower.shape} and {upper.shape}", ) _check_levels(lower.shape[0], levels) if np.any(upper == lower): @@ -183,7 +183,7 @@ def morton_keys( cells = [] for axis in range(ndim): cell = xpm.floor( - (positions[:, axis] - float(lower[axis])) * float(scales[axis]) + (positions[:, axis] - float(lower[axis])) * float(scales[axis]), ) cells.append(xpm.clip(cell, 0.0, top).astype(xpm.uint64)) return morton_encode(*cells) diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index ba006cf..835ca95 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -11,8 +11,8 @@ import array_api_compat import array_api_compat.numpy as np -from ._mpi_serial import (_LOCAL_RANK_VARIABLES, # noqa: F401 - re-exported - local_rank) +from ._mpi_serial import _LOCAL_RANK_VARIABLES # noqa: F401 - re-exported +from ._mpi_serial import local_rank from ._transfers import _ACTIVE as _COUNTERS from ._transfers import _describe, _record from .xp import array_backend, cupy_available, to_numpy @@ -88,7 +88,7 @@ def _pinned_or_host_empty(shape: tuple[int, ...], dtype: Any) -> np.ndarray: nbytes = int(np.prod(shape, dtype=np.int64)) * np.dtype(dtype).itemsize mem = cp.cuda.alloc_pinned_memory(max(nbytes, 1)) return np.frombuffer(mem, dtype, int(np.prod(shape, dtype=np.int64))).reshape( - shape + shape, ) except Exception: # noqa: BLE001 - no CuPy, no pinned memory, the fake CuPy, ... return np.empty(shape, dtype=dtype) @@ -227,7 +227,7 @@ def mpi_buffer( raise RuntimeError( "mpi_buffer(): it is not known whether MPI can take device buffers; " "call xp.mpi.mpi_is_cuda_aware(comm) once at startup (every rank), or " - "xp.mpi.set_mpi_cuda_aware(True/False), or pass cuda_aware=" + "xp.mpi.set_mpi_cuda_aware(True/False), or pass cuda_aware=", ) if cuda_aware: synchronize_for_mpi(array, stream=stream, event=event) @@ -265,7 +265,7 @@ def _mpi_module() -> Any: except ImportError as e: raise ImportError( "mpi4py is required for the CUDA-aware MPI check: install it, or " - "pass a communicator explicitly." + "pass a communicator explicitly.", ) from e return MPI @@ -329,7 +329,7 @@ def mpi_is_cuda_aware(comm: Any = None, *, method: str = "probe") -> bool: if method != "probe": raise ValueError( f"Unknown method {method!r}; only 'probe' is available (mpi4py does " - "not expose a reliable CUDA support query)." + "not expose a reliable CUDA support query).", ) if not _device_buffers_in_use(): return False @@ -383,5 +383,5 @@ def require_cuda_aware_mpi(comm: Any = None) -> None: "--with-device=ch4:ucx on a CUDA-enabled UCX; on clusters, load the " "CUDA-aware MPI module (its name differs per site) and rebuild mpi4py " "against it. Alternatively, copy the buffers to the host with " - "xp.to_numpy() before every MPI call." + "xp.to_numpy() before every MPI call.", ) diff --git a/src/cunumpy/_mpi_serial.py b/src/cunumpy/_mpi_serial.py index 88f1473..5cf29a2 100644 --- a/src/cunumpy/_mpi_serial.py +++ b/src/cunumpy/_mpi_serial.py @@ -244,7 +244,7 @@ def _copy(source: Any, target: Any, offset: int = 0) -> None: if offset + source.size > flat.size: raise ValueError( f"receive buffer too small: {flat.size} elements for {source.size} " - f"at offset {offset}" + f"at offset {offset}", ) flat[offset : offset + source.size] = source diff --git a/src/cunumpy/_philox.py b/src/cunumpy/_philox.py index c674c14..583ada0 100644 --- a/src/cunumpy/_philox.py +++ b/src/cunumpy/_philox.py @@ -93,7 +93,8 @@ def _random_bits(seed: Any, stream: Any, counter: Any) -> Any: mask, shift = xp.uint64(_MASK32), xp.uint64(32) seed, stream, counter = xp.broadcast_arrays(seed, stream, counter) words = xp.stack( - [counter & mask, counter >> shift, stream & mask, stream >> shift], axis=-1 + [counter & mask, counter >> shift, stream & mask, stream >> shift], + axis=-1, ) return philox4x32_10(words, seed & mask, seed >> shift) diff --git a/src/cunumpy/_random_streams.py b/src/cunumpy/_random_streams.py index 13d6b14..66daa3a 100644 --- a/src/cunumpy/_random_streams.py +++ b/src/cunumpy/_random_streams.py @@ -61,7 +61,10 @@ def __repr__(self) -> str: ) def seed( - self, value: int | None, rank: int = 0, bit_generator: str | None = None + self, + value: int | None, + rank: int = 0, + bit_generator: str | None = None, ) -> None: """Seed all draws of this process with the stream ``(value, rank)``. @@ -87,11 +90,12 @@ def seed( if bit_generator is not None and bit_generator not in BIT_GENERATORS: raise ValueError( f"Unknown bit generator {bit_generator!r}; use one of " - f"{', '.join(BIT_GENERATORS)}" + f"{', '.join(BIT_GENERATORS)}", ) self._bit_generator = bit_generator or "PCG64" self._sequence = np.random.SeedSequence( - entropy=None if value is None else int(value), spawn_key=(int(rank),) + entropy=None if value is None else int(value), + spawn_key=(int(rank),), ) self._generators.clear() legacy = int(self._sequence.generate_state(1)[0]) @@ -127,7 +131,9 @@ def generator(self, backend: str | None = None) -> Any: return self._generators[backend] def make_generator( - self, seed: int | None = None, backend: str | None = None + self, + seed: int | None = None, + backend: str | None = None, ) -> Any: """A generator for a component: its own if it has a seed, else the process one. @@ -156,7 +162,9 @@ def random(self, size: int | tuple[int, ...] | None = None, rng: Any = None) -> return self._rng(rng).random(size=size) def standard_normal( - self, size: int | tuple[int, ...] | None = None, rng: Any = None + self, + size: int | tuple[int, ...] | None = None, + rng: Any = None, ) -> Any: """Standard normal samples from `rng` (the process generator by default).""" return self._rng(rng).standard_normal(size=size) diff --git a/src/cunumpy/_scipy_backend.py b/src/cunumpy/_scipy_backend.py index e72fc68..1515563 100644 --- a/src/cunumpy/_scipy_backend.py +++ b/src/cunumpy/_scipy_backend.py @@ -70,7 +70,7 @@ def __init__(self, path: str = "") -> None: if path and path not in SUBMODULES: raise ValueError( f"scipy.{path} is not forwarded; forwarded subpackages: " - f"{', '.join(SUBMODULES)}" + f"{', '.join(SUBMODULES)}", ) self._path = path self._children: dict[str, ScipyNamespace] = {} @@ -100,7 +100,7 @@ def resolve(self) -> ModuleType: return importlib.import_module(name) except ImportError as error: raise ImportError( - f"xp.scipy on the {backend} backend needs {name}: {_INSTALL[backend]}" + f"xp.scipy on the {backend} backend needs {name}: {_INSTALL[backend]}", ) from error def __getattr__(self, name: str) -> Any: @@ -119,7 +119,7 @@ def __getattr__(self, name: str) -> Any: raise AttributeError( f"{module.__name__} has no attribute {name!r}: it is not available " f"on the {get_backend()} backend (it may exist in " - f"{self._name(_ROOTS[other])})" + f"{self._name(_ROOTS[other])})", ) from None def available(self, name: str) -> bool: diff --git a/src/cunumpy/_staging.py b/src/cunumpy/_staging.py index 25ae1fa..9faac25 100644 --- a/src/cunumpy/_staging.py +++ b/src/cunumpy/_staging.py @@ -75,7 +75,7 @@ def _check(self) -> None: raise RuntimeError( "this staged copy's host buffer was reused by a later copy; call " "result() before starting more copies than there are buffers, or " - "keep result().copy()" + "keep result().copy()", ) def ready(self) -> bool: @@ -115,7 +115,10 @@ class HostStaging: """ def __init__( - self, shape: tuple[int, ...] | int, dtype: Any, buffers: int = 2 + self, + shape: tuple[int, ...] | int, + dtype: Any, + buffers: int = 2, ) -> None: if buffers < 1: raise ValueError(f"buffers must be at least 1, got {buffers}") @@ -165,7 +168,7 @@ def copy(self, array: Any) -> StagedCopy: if tuple(array.shape) != self.shape or np.dtype(array.dtype) != self.dtype: raise ValueError( f"HostStaging for {self.shape} {self.dtype} got an array of shape " - f"{tuple(array.shape)} and dtype {array.dtype}" + f"{tuple(array.shape)} and dtype {array.dtype}", ) device = _is_device_array(array) slot = self._slot(device) diff --git a/src/cunumpy/_transfers.py b/src/cunumpy/_transfers.py index 14c1fcb..509b63e 100644 --- a/src/cunumpy/_transfers.py +++ b/src/cunumpy/_transfers.py @@ -238,5 +238,5 @@ def assert_no_transfers() -> Generator[TransferCounter, None, None]: if counter.total: raise AssertionError( "host/device transfers inside a block that must not transfer:\n" - + counter.report() + + counter.report(), ) diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 6db0037..2db7b2e 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -113,7 +113,7 @@ def _pytest() -> Any: import pytest except ImportError: # pragma: no cover - pytest is installed in the test suite raise ImportError( - "cunumpy.kernel_testing needs pytest for this feature: pip install pytest" + "cunumpy.kernel_testing needs pytest for this feature: pip install pytest", ) from None return pytest @@ -150,7 +150,7 @@ def __getattr__(name: str) -> Any: def _is_array(value: Any) -> bool: return array_api_compat.is_numpy_array(value) or array_api_compat.is_cupy_array( - value + value, ) @@ -161,7 +161,8 @@ def _arrays_in(value: Any, name: str, found: dict[str, Any], depth: int) -> None elif depth == 0: return elif isinstance(value, (CudaStructArguments, CudaStructValue)) and hasattr( - value, "struct" + value, + "struct", ): # by field name, like the attributes of the host argument object; the # fields of a CudaStructArguments may be properties (not in vars()) @@ -184,7 +185,8 @@ def _arrays_in(value: Any, name: str, found: dict[str, Any], depth: int) -> None def _collect_arrays( - args: Sequence[Any], outputs: Sequence[int] | None = None + args: Sequence[Any], + outputs: Sequence[int] | None = None, ) -> dict[str, Any]: """The arrays among `args` (or among the arguments `outputs`), by name. @@ -204,13 +206,13 @@ def _collect_arrays( 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}" + 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" + "arguments", ) _arrays_in(args[index], f"argument {index}", found, depth=2) return found @@ -234,7 +236,7 @@ def _compare_results( if host.keys() != device.keys(): raise AssertionError( f"{kernel_name}: the host and CUDA calls do not have the same array " - f"arguments: host {sorted(host)}, CUDA {sorted(device)}" + f"arguments: host {sorted(host)}, CUDA {sorted(device)}", ) for name, expected in host.items(): np.testing.assert_allclose( @@ -392,7 +394,7 @@ def test_parity(kernel): marks = ( pytest.mark.skip( reason=f"no test arguments for {name!r}: add {name}_test_args.py " - "with make_args(backend, seed) and N_THREADS to its folder" + "with make_args(backend, seed) and N_THREADS to its folder", ), ) cases.append(pytest.param(kernel, id=name, marks=marks)) @@ -426,7 +428,7 @@ def check_parity(kernel: Kernel, **overrides: Any) -> dict[str, np.ndarray]: if module is None: raise ValueError( f"kernel {kernel.name!r} has no test-arguments module: add " - f"{kernel.name}_test_args.py with make_args(backend, seed) to its folder" + f"{kernel.name}_test_args.py with make_args(backend, seed) to its folder", ) make_args = getattr(module, "make_args", None) if not callable(make_args): @@ -477,11 +479,11 @@ def _parse_prototype( except ValueError: raise ValueError( f"unsupported return type {result_text!r} of {name!r}: a scalar " - "type (or void) is required" + "type (or void) is required", ) from None if result.pointer: raise ValueError( - f"{name!r} returns a pointer; only scalar results can be collected" + f"{name!r} returns a pointer; only scalar results can be collected", ) params_text = match.group("params").strip() params = [] @@ -493,7 +495,7 @@ def _parse_prototype( if by_reference and param.struct is None: raise ValueError( f"unsupported parameter {text!r} of {name!r}: only structs can " - "be passed by reference" + "be passed by reference", ) params.append((text, param)) return result, name, params @@ -591,7 +593,7 @@ def device_function_kernel( raise ValueError( f"the parameter {param.name!r} of {function!r} clashes with the " f"generated parameter of that name; pass another out_param or " - "n_threads_param" + "n_threads_param", ) wrapper_params, call_args = [], [] diff --git a/src/cunumpy/petsc.py b/src/cunumpy/petsc.py index b4df2ce..7b575d7 100644 --- a/src/cunumpy/petsc.py +++ b/src/cunumpy/petsc.py @@ -43,7 +43,7 @@ def _petsc() -> Any: from petsc4py import PETSc except ImportError as error: raise ImportError( - "xp.petsc.petsc_vec needs petsc4py (pip install petsc4py)" + "xp.petsc.petsc_vec needs petsc4py (pip install petsc4py)", ) from error return PETSc @@ -92,13 +92,13 @@ def petsc_vec(array: Any, comm: Any = None) -> Any: device = _is_device_array(array) if not (device or isinstance(array, np.ndarray)): raise TypeError( - f"petsc_vec takes a NumPy or CuPy array, got {type(array).__name__}" + f"petsc_vec takes a NumPy or CuPy array, got {type(array).__name__}", ) scalar = np.dtype(PETSc.ScalarType) if array.dtype != scalar: raise TypeError( f"petsc_vec needs an array of PETSc's scalar type {scalar}, got " - f"{array.dtype} (convert it once, outside the time loop)" + f"{array.dtype} (convert it once, outside the time loop)", ) if not array.flags.c_contiguous: raise ValueError("petsc_vec needs a C-contiguous array (no copy is made)") @@ -109,14 +109,14 @@ def petsc_vec(array: Any, comm: Any = None) -> Any: if device: raise RuntimeError( "petsc4py could not wrap the CuPy array; it needs a PETSc built " - "with CUDA or HIP support (--with-cuda / --with-hip)" + "with CUDA or HIP support (--with-cuda / --with-hip)", ) from error raise if device and not any(t in vec.getType() for t in _DEVICE_VEC_TYPES): vec.destroy() raise RuntimeError( f"petsc4py created a {vec.getType()!r} vector for a CuPy array; it " - "needs a PETSc built with CUDA or HIP support" + "needs a PETSc built with CUDA or HIP support", ) vec.setAttr("cunumpy_array", array) # the vector does not own the memory return vec diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 6dafcf4..5a56f1c 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -86,7 +86,7 @@ def _load_backend(self, backend: BackendType, verbose: bool = False) -> ModuleTy else: if verbose: print( - "CuPy not available or not functional. Falling back to NumPy." + "CuPy not available or not functional. Falling back to NumPy.", ) self._backend = "numpy" return np @@ -124,14 +124,17 @@ def set(self, backend: BackendType, *, strict: bool = False) -> None: if strict and backend == "cupy" and not cupy_available(): raise RuntimeError( "Cannot select the CuPy backend: " - + (_CUPY_UNAVAILABLE_REASON or "CuPy/CUDA is unavailable") + + (_CUPY_UNAVAILABLE_REASON or "CuPy/CUDA is unavailable"), ) module = self._load_backend(backend) # sets self._backend to the effective one self._set(self._backend, module) @contextmanager def use_backend( - self, backend: BackendType, *, strict: bool = False + self, + backend: BackendType, + *, + strict: bool = False, ) -> Generator[None, None, None]: """Temporarily change the backend.""" old_backend = self._backend @@ -152,7 +155,9 @@ def use_backend( def use_backend( - backend: BackendType, *, strict: bool = False + backend: BackendType, + *, + strict: bool = False, ) -> Generator[None, None, None]: """Temporarily change the backend.""" return array_backend.use_backend(backend, strict=strict) @@ -344,7 +349,7 @@ def as_device_array( raise RuntimeError( f"{what}: the active backend is {array_backend.backend!r}; device " "arguments are only built on the CuPy backend, and host data is never " - "copied to the device implicitly (build host arguments instead)" + "copied to the device implicitly (build host arguments instead)", ) import cupy as cp @@ -360,7 +365,7 @@ def as_device_array( if ndim is not None and result.ndim != ndim: raise ValueError( f"{what} must have {ndim} dimension(s), got {result.ndim} " - f"(shape {result.shape})" + f"(shape {result.shape})", ) return result @@ -432,7 +437,7 @@ def assert_same_backend(*arrays: Any) -> None: backends = [get_array_backend(array) for array in arrays] raise TypeError( f"Arrays are on mismatched backends: {backends}. Use " - "xp.to_cunumpy()/xp.to_numpy()/xp.to_cupy() to align them first." + "xp.to_cunumpy()/xp.to_numpy()/xp.to_cupy() to align them first.", ) diff --git a/tests/unit/test_array_api_compat_backend.py b/tests/unit/test_array_api_compat_backend.py index 5a3dab6..5339ab3 100644 --- a/tests/unit/test_array_api_compat_backend.py +++ b/tests/unit/test_array_api_compat_backend.py @@ -51,7 +51,8 @@ def test_to_cupy_returns_real_cupy_ndarray(): @pytest.mark.parametrize( - "dtype", [np.float32, np.float64, np.int32, np.int64, np.complex64] + "dtype", + [np.float32, np.float64, np.int32, np.int64, np.complex64], ) def test_round_trip_preserves_values_and_dtype(dtype): _skip_without_cupy() diff --git a/tests/unit/test_benchmarks.py b/tests/unit/test_benchmarks.py index 4f4e165..e89a0e9 100644 --- a/tests/unit/test_benchmarks.py +++ b/tests/unit/test_benchmarks.py @@ -6,7 +6,8 @@ @pytest.mark.skipif( - not xp.cupy_available(), reason="CuPy/GPU not available or not functional" + not xp.cupy_available(), + reason="CuPy/GPU not available or not functional", ) def test_benchmark_matmul(): """Benchmark matrix multiplication to show CuPy performance gain.""" @@ -49,7 +50,8 @@ def test_benchmark_matmul(): @pytest.mark.skipif( - not xp.cupy_available(), reason="CuPy/GPU not available or not functional" + not xp.cupy_available(), + reason="CuPy/GPU not available or not functional", ) def test_benchmark_fft(): """Benchmark FFT performance.""" diff --git a/tests/unit/test_cuda_collectives.py b/tests/unit/test_cuda_collectives.py index cd52c0b..30b0fc5 100644 --- a/tests/unit/test_cuda_collectives.py +++ b/tests/unit/test_cuda_collectives.py @@ -26,7 +26,8 @@ @requires_cuda @pytest.mark.parametrize( - "dtype", [np.int32, np.uint32, np.int64, np.uint64, np.float32, np.float64] + "dtype", + [np.int32, np.uint32, np.int64, np.uint64, np.float32, np.float64], ) def test_scans_and_block_atomic_totals_for_shuffle_types(dtype): import cupy as cp @@ -35,7 +36,10 @@ def test_scans_and_block_atomic_totals_for_shuffle_types(dtype): data = cp.asarray(host) inc, exc, total = cp.zeros_like(data), cp.zeros_like(data), cp.zeros(1, dtype=dtype) kernel = CudaKernel( - TYPED_SOURCE, "typed_scan", template_args=(dtype,), block_size=35 + TYPED_SOURCE, + "typed_scan", + template_args=(dtype,), + block_size=35, ) kernel(data, inc, exc, total, grid=2) expected = np.cumsum(host.reshape(2, 35), axis=1) @@ -68,7 +72,8 @@ def test_scans_and_block_atomic_totals_for_shuffle_types(dtype): @requires_cuda @pytest.mark.parametrize( - "block", [1, 17, 31, 32, 33, 64, 127, 128, 1024, (7, 5), (3, 4, 3)] + "block", + [1, 17, 31, 32, 33, 64, 127, 128, 1024, (7, 5), (3, 4, 3)], ) def test_partial_warps_multidimensional_blocks_and_repeated_collectives(block): import cupy as cp @@ -132,7 +137,8 @@ def test_arbitrary_sparse_masks(mask): for got, want in zip(out, expected): result = cp.asnumpy(got) np.testing.assert_array_equal( - result[active], np.broadcast_to(want, active.shape) + result[active], + np.broadcast_to(want, active.shape), ) np.testing.assert_array_equal(result[inactive], -999) @@ -147,7 +153,10 @@ def test_exclusive_scan_preserves_small_previous_values(mask): host[active[0]], host[active[1]] = 1.0, 1e20 out = [cp.zeros(32) for _ in range(5)] CudaKernel(MASKED_SOURCE, "masked", block_size=32)( - cp.asarray(host), mask, *out, n_threads=32 + cp.asarray(host), + mask, + *out, + n_threads=32, ) assert float(out[-1][active[1]]) == 1.0 diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 522a836..b2af002 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -73,7 +73,7 @@ def __init__(self, dtype, ptr=0x1000, shape=(1,), strides=None, flags=None): stride *= n self.strides = tuple(strides) c_contiguous = self.strides == tuple( - np.zeros(self.shape, dtype=self.dtype).strides + np.zeros(self.shape, dtype=self.dtype).strides, ) self.flags = SimpleNamespace(c_contiguous=c_contiguous) if flags is not None: @@ -103,7 +103,7 @@ def test_parse_axpy(): ("n", "int", False), ] assert [p.dtype for p in params] == [np.dtype(np.float64)] * 3 + [ - np.dtype(np.int32) + np.dtype(np.int32), ] @@ -215,7 +215,8 @@ def test_array_checks(): x = FakeDeviceArray(np.float64) with pytest.raises( - TypeError, match=r"argument 2 \(double\* y\) must be a CuPy array" + TypeError, + match=r"argument 2 \(double\* y\) must be a CuPy array", ): kernel.prepare_args(1.0, x, np.zeros(3), 3) # host array: never copied with pytest.raises(TypeError, match="must have dtype float64, got float32"): @@ -231,10 +232,13 @@ def test_non_contiguous_arrays_are_rejected(): kernel = CudaKernel(AXPY, "axpy") x = FakeDeviceArray(np.float64) view = FakeDeviceArray( - np.float64, shape=(3, 2), flags=SimpleNamespace(c_contiguous=False) + np.float64, + shape=(3, 2), + flags=SimpleNamespace(c_contiguous=False), ) with pytest.raises( - TypeError, match=r"argument 2 \(double\* y\) must be C-contiguous" + TypeError, + match=r"argument 2 \(double\* y\) must be C-contiguous", ): kernel.prepare_args(1.0, x, view, 3) # the message says why, and how to fix it @@ -396,7 +400,7 @@ def test_stream_and_include_dirs_on_gpu(tmp_path): import cupy as cp (tmp_path / "helpers.cuh").write_text( - "__device__ double twice(double v) { return 2 * v; }\n" + "__device__ double twice(double v) { return 2 * v; }\n", ) source = r""" #include "helpers.cuh" @@ -546,7 +550,8 @@ def test_struct_pointer_fields_must_be_contiguous(): } view = FakeDeviceArray(np.float64, flags=SimpleNamespace(c_contiguous=False)) with pytest.raises( - TypeError, match=r"argument 0 \(double\* x\) must be C-contiguous" + TypeError, + match=r"argument 0 \(double\* x\) must be C-contiguous", ): PARTICLES(x=view, **values) @@ -613,7 +618,11 @@ def test_struct_on_gpu(): ) out, size = cp.zeros(4), cp.zeros(1, dtype=cp.uint64) CudaKernel(PUSH_SOURCE, "push", structs=[PARTICLES])( - value, 0.5, out, size, n_threads=n + value, + 0.5, + out, + size, + n_threads=n, ) assert int(size.get()[0]) == PARTICLES.dtype.itemsize # same layout as in C assert out.get().tolist() == [n, 2.0, 42.0, 1.5] @@ -718,7 +727,8 @@ def test_view_parameters_pack_pointer_shape_and_strides(): def test_view_parameter_errors(): kernel = CudaKernel(SCALE_COLUMN, "scale_column") with pytest.raises( - TypeError, match=r"argument 0 \(Array2D a\) must be a CuPy" + TypeError, + match=r"argument 0 \(Array2D a\) must be a CuPy", ): kernel.prepare_args(np.zeros((2, 2)), 1, 2.0) with pytest.raises(TypeError, match="must have dtype float64, got float32"): @@ -727,7 +737,9 @@ def test_view_parameter_errors(): kernel.prepare_args(FakeDeviceArray(np.float64, shape=(2,)), 1, 2.0) with pytest.raises(TypeError, match="not multiples of the element size"): kernel.prepare_args( - FakeDeviceArray(np.float64, shape=(2, 2), strides=(20, 8)), 1, 2.0 + FakeDeviceArray(np.float64, shape=(2, 2), strides=(20, 8)), + 1, + 2.0, ) @@ -826,7 +838,9 @@ def test_bounds_check_traps_on_gpu(): import cupy as cp kernel = CudaKernel( - SCALE_COLUMN, "scale_column", options=("-DCUNUMPY_BOUNDS_CHECK",) + SCALE_COLUMN, + "scale_column", + options=("-DCUNUMPY_BOUNDS_CHECK",), ) a = cp.ones((4, 3)) kernel(a, 1, 2.0, n_threads=4) @@ -944,7 +958,8 @@ def unparsable(x: "float[:](order=F)"): # noqa: F821 (deliberately unparsable) pass with pytest.raises( - ValueError, match="parameter 'x' of .*missing.* no type annotation" + ValueError, + match="parameter 'x' of .*missing.* no type annotation", ): CudaStruct.from_signature(missing, "A") with pytest.raises(ValueError, match="unsupported scalar type 'str'"): @@ -975,7 +990,8 @@ def test_to_header(tmp_path): assert '#include "cunumpy/array_view.cuh"\n#include "defs.cuh"\n' in written # the pattern: a committed header equals the generated one assert path.read_text() == struct.to_header( - guard="STRUPHY_MARKER_ARGS", includes=["defs.cuh"] + guard="STRUPHY_MARKER_ARGS", + includes=["defs.cuh"], ) # a struct without views does not include the array view header plain = PARTICLES.to_header() @@ -983,25 +999,29 @@ def test_to_header(tmp_path): # the generated header defines the struct as the kernel expects struct.check_source(header) CudaKernel( - header + 'extern "C" __global__ void f(MarkerArgs a) {}', "f", structs=[struct] + header + 'extern "C" __global__ void f(MarkerArgs a) {}', + "f", + structs=[struct], ) def test_write_cuda_header(tmp_path): path = tmp_path / "pusher_args.cuh" header = write_cuda_header( - path, [MARKERS, PARTICLES], includes=["#include "] + path, + [MARKERS, PARTICLES], + includes=["#include "], ) assert path.read_text() == header assert header.startswith( "// Generated by cunumpy.cuda.CudaStruct from the Python definition; do not edit.\n" "#ifndef PUSHER_ARGS_CUH\n#define PUSHER_ARGS_CUH\n\n" - '#include "cunumpy/array_view.cuh"\n#include \n\n' + '#include "cunumpy/array_view.cuh"\n#include \n\n', ) assert header.index("struct Markers") < header.index("struct Particles") assert header.endswith("#endif // PUSHER_ARGS_CUH\n") assert write_cuda_header(path, [PARTICLES], guard="G") == PARTICLES.to_header( - guard="G" + guard="G", ) for struct in (MARKERS, PARTICLES): struct.check_source(header) @@ -1100,8 +1120,10 @@ def test_variants_on_gpu(): variants = CudaKernelVariants( lambda ndim, dtype: CudaKernel( - _generated_source(ndim, ctype_of(dtype)), "fill", block_size=1 - ) + _generated_source(ndim, ctype_of(dtype)), + "fill", + block_size=1, + ), ) variants.compile_all([(2, np.float64), (1, np.int32)], jobs=2) assert all(variants.get(*key).is_compiled for key in variants) @@ -1210,7 +1232,7 @@ def _run_python(code, env=None): environment = {**os.environ, **(env or {})} environment["PYTHONPATH"] = os.pathsep.join( - [str(Path(xp.__file__).parents[1]), environment.get("PYTHONPATH", "")] + [str(Path(xp.__file__).parents[1]), environment.get("PYTHONPATH", "")], ) result = subprocess.run( [sys.executable, "-c", code], @@ -1256,7 +1278,8 @@ def test_debug_from_env(value, expected): @pytest.mark.parametrize( - "value, expected", [("1", "True"), ("on", "True"), ("0", "False"), (None, "False")] + "value, expected", + [("1", "True"), ("on", "True"), ("0", "False"), (None, "False")], ) def test_debug_from_env_at_import(monkeypatch, value, expected): """The environment variable is read when cunumpy.xp is imported.""" @@ -1389,8 +1412,9 @@ def test_out_of_bounds_write_on_gpu(debug): _skip_without_cupy() output = "".join( _run_python( - OUT_OF_BOUNDS.replace("DEBUG", str(debug)), env={"CUNUMPY_CUDA_DEBUG": "0"} - ) + OUT_OF_BOUNDS.replace("DEBUG", str(debug)), + env={"CUNUMPY_CUDA_DEBUG": "0"}, + ), ) assert "no error" not in output, output if debug: @@ -1424,7 +1448,7 @@ def header_tree(tmp_path): (tmp_path / "sub").mkdir() (tmp_path / "double_it_cuda.cu").write_text(INCLUDING_SOURCE) (tmp_path / "b.cuh").write_text( - '#include "sub/c.cuh"\n__device__ double twice(double v) { return FACTOR * v; }\n' + '#include "sub/c.cuh"\n__device__ double twice(double v) { return FACTOR * v; }\n', ) (tmp_path / "sub" / "c.cuh").write_text("#define FACTOR 2\n") return tmp_path @@ -1529,7 +1553,7 @@ def test_resolve_includes_angle_dirs(tmp_path): (user / "lib").mkdir(parents=True) (user / "lib" / "a.cuh").write_text("") assert resolve_includes('#include "lib/a.cuh"\n', [user], angle_dirs=[shipped]) == [ - user / "lib" / "a.cuh" + user / "lib" / "a.cuh", ] @@ -1708,7 +1732,9 @@ def test_struct_arguments_check_their_fields(): ) with pytest.raises(TypeError, match="must be a CuPy array"): ParticleArguments( - np.zeros(3), FakeDeviceArray(np.bool_), FakeDeviceArray(np.int64) + np.zeros(3), + FakeDeviceArray(np.bool_), + FakeDeviceArray(np.int64), ) args = _particle_arguments() args.n = 2.5 # a float for the int field n @@ -1728,7 +1754,8 @@ def __init__(self, x): self.pack() with pytest.raises( - AttributeError, match="no attribute 'n' for the field of struct Incomplete" + AttributeError, + match="no attribute 'n' for the field of struct Incomplete", ): Incomplete(FakeDeviceArray(np.float64)) @@ -1870,11 +1897,18 @@ def test_struct_arguments_on_gpu(): alive = cp.ones(n, dtype=bool) alive[::2] = False args = ParticleArguments( - cp.zeros(n), alive, cp.array([7, 42], dtype=cp.int64), weight=1.5 + cp.zeros(n), + alive, + cp.array([7, 42], dtype=cp.int64), + weight=1.5, ) out, size = cp.zeros(4), cp.zeros(1, dtype=cp.uint64) CudaKernel(PUSH_SOURCE, "push", structs=[ParticleArguments.struct])( - args, 0.5, out, size, n_threads=n + args, + 0.5, + out, + size, + n_threads=n, ) assert int(size.get()[0]) == ParticleArguments.struct.dtype.itemsize assert out.get().tolist() == [n, 2.0, 42.0, 1.5] @@ -1932,7 +1966,8 @@ def test_debug_kernel_in_a_cuda_graph(): def test_layout_source_reports_size_alignment_and_offsets(): struct = CudaStruct( - "LayoutArgs", [("markers", "Array2D"), ("valid", "bool*"), ("n", "int")] + "LayoutArgs", + [("markers", "Array2D"), ("valid", "bool*"), ("n", "int")], ) source = struct.layout_source() assert source.startswith('#include "cunumpy/array_view.cuh"') @@ -1949,7 +1984,7 @@ def test_layout_source_reports_size_alignment_and_offsets(): assert from_header.startswith('#include "pkg/layout_args.cuh"') assert "struct LayoutArgs {" not in from_header assert struct.layout_source("#include ").startswith( - "#include " + "#include ", ) # no views: no array_view include assert "array_view" not in PARTICLES.layout_source() @@ -1974,7 +2009,7 @@ def test_verify_layout_on_gpu(tmp_path): header = tmp_path / "drifted.cuh" header.write_text( '#include "cunumpy/array_view.cuh"\n' - "struct Views { int n; bool b; Array2D a; };\n" # b moved before a + "struct Views { int n; bool b; Array2D a; };\n", # b moved before a ) with pytest.raises(ValueError, match="differs from its CudaStruct dtype"): views.verify_layout("drifted.cuh", include_dirs=[tmp_path]) @@ -2044,7 +2079,9 @@ def test_reductions_on_gpu(block_size): np.testing.assert_array_equal(cp.asnumpy(lo), np.nanmin(blocks, axis=1)) np.testing.assert_array_equal(cp.asnumpy(hi), np.nanmax(blocks, axis=1)) np.testing.assert_allclose( - cp.asnumpy(warp), np.nan_to_num(padded).reshape(-1, 32).sum(axis=1), rtol=1e-12 + cp.asnumpy(warp), + np.nan_to_num(padded).reshape(-1, 32).sum(axis=1), + rtol=1e-12, ) @@ -2065,7 +2102,10 @@ def test_array4d_parameters_pack_pointer_shape_and_strides(): kernel = CudaKernel(VIEW_4D, "scale_4d") # every second component of a (2, 3, 4, 6) grid: a non-contiguous view a = FakeDeviceArray( - np.float64, ptr=0x40, shape=(2, 3, 4, 3), strides=(576, 192, 48, 16) + np.float64, + ptr=0x40, + shape=(2, 3, 4, 3), + strides=(576, 192, 48, 16), ) packed, _, _ = kernel.prepare_args(a, 2.0, 5) assert packed["data"] == 0x40 @@ -2159,7 +2199,9 @@ def test_n_threads_from_first_array(recorded): def test_shared_memory_above_the_default_is_opted_in(recorded, monkeypatch): kernel, raw = recorded monkeypatch.setattr( - device_module, "max_shared_memory_per_block", lambda opt_in=False: 100_000 + device_module, + "max_shared_memory_per_block", + lambda opt_in=False: 100_000, ) x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) kernel(1.0, x, y, 1, n_threads=1, shared_mem=40_000) # below 48 KiB: no setup diff --git a/tests/unit/test_emulation.py b/tests/unit/test_emulation.py index b6bfb64..ff21718 100644 --- a/tests/unit/test_emulation.py +++ b/tests/unit/test_emulation.py @@ -7,7 +7,8 @@ from cunumpy.kernel_testing import emulate_cuda_kernel, emulation_compiler pytestmark = pytest.mark.skipif( - emulation_compiler() is None, reason="no C++ compiler for the emulation" + emulation_compiler() is None, + reason="no C++ compiler for the emulation", ) AXPY = r""" @@ -56,7 +57,12 @@ def test_strided_view_written_back_into_the_callers_array(): markers = np.arange(24.0).reshape(8, 3) every_second_row = markers[::2] # a non-contiguous view emulate_cuda_kernel( - CudaKernel(COLUMN, "scale_column"), every_second_row, 1, 10.0, grid=1, block=2 + CudaKernel(COLUMN, "scale_column"), + every_second_row, + 1, + 10.0, + grid=1, + block=2, ) expected = np.arange(24.0).reshape(8, 3) expected[::2, 1] *= 10.0 @@ -139,7 +145,12 @@ def test_arrays_are_checked(): kernel = CudaKernel(AXPY, "axpy") with pytest.raises(TypeError, match="dtype float64"): emulate_cuda_kernel( - kernel, 1.0, np.ones(3, np.float32), np.ones(3), 3, n_threads=3 + kernel, + 1.0, + np.ones(3, np.float32), + np.ones(3), + 3, + n_threads=3, ) with pytest.raises(TypeError, match="NumPy array"): emulate_cuda_kernel(kernel, 1.0, [1.0], np.ones(3), 3, n_threads=3) @@ -250,7 +261,9 @@ def test_dynamic_shared_memory_per_block_deposit(): def test_barrier_orders_writes_before_reads(): x = np.arange(128.0) emulate_cuda_kernel( - CudaKernel(REVERSE, "reverse_blocks", block_size=64), x, n_threads=128 + CudaKernel(REVERSE, "reverse_blocks", block_size=64), + x, + n_threads=128, ) expected = np.concatenate([np.arange(64.0)[::-1], np.arange(64.0, 128.0)[::-1]]) np.testing.assert_array_equal(x, expected) @@ -271,7 +284,10 @@ def test_barrier_orders_writes_before_reads(): def test_threads_leaving_before_a_barrier_do_not_hang(): x = np.ones(40) emulate_cuda_kernel( - CudaKernel(EARLY_EXIT, "early", block_size=32), x, 40, n_threads=40 + CudaKernel(EARLY_EXIT, "early", block_size=32), + x, + 40, + n_threads=40, ) assert x.tolist() == [2.0] * 40 @@ -301,12 +317,14 @@ def test_2d_blocks_with_shared_tiles(): def test_compile_errors_and_crashes_are_reported(): broken = CudaKernel( - 'extern "C" __global__ void k(int n) { undefined_call(n); }', "k" + 'extern "C" __global__ void k(int n) { undefined_call(n); }', + "k", ) with pytest.raises(RuntimeError, match="does not compile"): emulate_cuda_kernel(broken, 1, n_threads=1) trap = CudaKernel( - 'extern "C" __global__ void k(int n) { if (n > 0) __trap(); }', "k" + 'extern "C" __global__ void k(int n) { if (n > 0) __trap(); }', + "k", ) with pytest.raises(RuntimeError, match="crashed"): emulate_cuda_kernel(trap, 1, n_threads=1) @@ -317,5 +335,10 @@ def test_bounds_checks_from_the_view_header(): a = np.ones((4, 2)) with pytest.raises(RuntimeError, match="crashed"): emulate_cuda_kernel( - kernel, a, 5, 2.0, n_threads=4, options=("-DCUNUMPY_BOUNDS_CHECK",) + kernel, + a, + 5, + 2.0, + n_threads=4, + options=("-DCUNUMPY_BOUNDS_CHECK",), ) diff --git a/tests/unit/test_fusion.py b/tests/unit/test_fusion.py index e50c25e..6bcea1b 100644 --- a/tests/unit/test_fusion.py +++ b/tests/unit/test_fusion.py @@ -51,7 +51,9 @@ def fused(*args, **kwargs): monkeypatch.setattr(fusion, "_cupy_fuse", fake_cupy_fuse) monkeypatch.setattr( - fusion, "_is_device_array", lambda a: isinstance(a, FakeDeviceArray) + fusion, + "_is_device_array", + lambda a: isinstance(a, FakeDeviceArray), ) @xp.kernels.fuse(kernel_name="p") diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py index 64c9c04..96d0506 100644 --- a/tests/unit/test_kernel_dispatch.py +++ b/tests/unit/test_kernel_dispatch.py @@ -124,16 +124,18 @@ def kernel_package(tmp_path, monkeypatch): for name, body in (("scale", "x[i] *= a"), ("shift", "x[i] += a")): (root / name).mkdir(parents=True) (root / name / "__init__.py").write_text("") - (root / name / f"{name}_kernels.py").write_text(textwrap.dedent(f""" + (root / name / f"{name}_kernels.py").write_text( + textwrap.dedent(f""" def {name}(x, a, n): for i in range(n): {body} - """)) + """), + ) (root / "scale" / "scale_cuda.cu").write_text(SCALE_CUDA) (root / "not_a_kernel").mkdir() (root / "__init__.py").write_text( "from cunumpy.kernels import KernelCatalog\n\n" - "catalog = KernelCatalog.from_package(__name__)\n" + "catalog = KernelCatalog.from_package(__name__)\n", ) monkeypatch.syspath_prepend(str(tmp_path)) yield importlib.import_module("demo_kernel_pkg").catalog @@ -239,7 +241,7 @@ def test_catalog_host_options(kernel_package_factory): assert all(catalog[name].host_kernel.outputs == (0,) for name in catalog) catalog = kernel_package_factory( - host_options=lambda name: {"outputs": (0,) if name == "scale" else ()} + host_options=lambda name: {"outputs": (0,) if name == "scale" else ()}, ) assert catalog["scale"].host_kernel.outputs == (0,) assert catalog["shift"].host_kernel.outputs == () @@ -375,10 +377,10 @@ def test_catalog_kernel_includes_from_the_source_root(tmp_path, monkeypatch): (root / "kernels" / "__init__.py").write_text("") (root / "kernels" / "scale" / "__init__.py").write_text("") (root / "kernels" / "scale" / "scale_kernels.py").write_text( - "def scale(x, a, n):\n pass\n" + "def scale(x, a, n):\n pass\n", ) (root / "kernels" / "scale" / "scale_cuda.cu").write_text( - '#include "demo_include_pkg/common.cuh"\n' + SCALE_CUDA + '#include "demo_include_pkg/common.cuh"\n' + SCALE_CUDA, ) monkeypatch.syspath_prepend(str(tmp_path)) try: diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index 063c287..a0ca981 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -41,7 +41,9 @@ def fake_gpu(monkeypatch): CUDA launches are recorded instead of run.""" monkeypatch.setattr(dispatch_module, "get_backend", lambda: "cupy") monkeypatch.setattr( - dispatch_module, "_is_device_array", lambda a: isinstance(a, FakeDeviceArray) + dispatch_module, + "_is_device_array", + lambda a: isinstance(a, FakeDeviceArray), ) launches = [] @@ -133,13 +135,16 @@ def reordered(factor, x, n): ... with pytest.raises(ValueError, match="host kernel takes"): Kernel( - reordered, CudaKernel(SCALE_CUDA, "scale"), name="scale" + reordered, + CudaKernel(SCALE_CUDA, "scale"), + name="scale", ).check_signature() # nothing to compare: no CUDA kernel, an unparsed signature, a compiled builtin Kernel(renamed).check_signature() Kernel( - renamed, CudaKernel(SCALE_CUDA, "scale", check_signature=False) + renamed, + CudaKernel(SCALE_CUDA, "scale", check_signature=False), ).check_signature() compiled = CompiledFunction() Kernel(compiled, CudaKernel(SCALE_CUDA, "scale"), name="scale").check_signature() @@ -194,7 +199,10 @@ def failing(mod): fallback_calls = [] with_fallback = CompiledHostKernel( - module, "double", failing, fallback=fallback_calls.append + module, + "double", + failing, + fallback=fallback_calls.append, ) with_fallback("x") assert fallback_calls == ["x"] and not with_fallback.compiled @@ -219,7 +227,7 @@ def pyccel_style_package(tmp_path, monkeypatch): (root / "scale" / "scale_pyccel.py").write_text( "def scale(x: 'float[:]', factor: float, n: int):\n" " for i in range(n):\n" - " x[i] *= factor\n" + " x[i] *= factor\n", ) (root / "scale" / "scale_cuda.cu").write_text(SCALE_CUDA) monkeypatch.syspath_prepend(str(tmp_path)) @@ -307,16 +315,16 @@ def self_declaring_package(tmp_path, monkeypatch): (root / "scale" / "scale_pyccel.py").write_text( "def scale(x: 'float[:]', factor: float, n: int):\n" " for i in range(n):\n" - " x[i] *= factor\n" + " x[i] *= factor\n", ) (root / "scale" / "scale_numpy.py").write_text( "CALLS = []\n\n" "def scale(x, factor, n):\n" " CALLS.append(n)\n" - " x[:n] *= factor\n" + " x[:n] *= factor\n", ) (root / "scale" / "scale_numba.py").write_text( - "import a_jit_library_that_is_not_installed\n" + "import a_jit_library_that_is_not_installed\n", ) (root / "scale" / "scale_cuda.cu").write_text(SCALE_CUDA) (root / "scale" / "__init__.py").write_text( @@ -328,7 +336,7 @@ def self_declaring_package(tmp_path, monkeypatch): "kernel = xp.kernels.Kernel.from_folder(\n" " __name__, host_suffix='_pyccel', dispatch='arrays',\n" " compile_host=_compile, n_threads_from='first_array',\n" - ")\n" + ")\n", ) monkeypatch.syspath_prepend(str(tmp_path)) yield importlib.import_module("demo_folder_pkg.scale") @@ -405,7 +413,8 @@ def no_pyccel(): with_numpy("x") assert used == ["x"] only_python = HostImplementations( - "scale", {"pyccel": no_pyccel, "python": lambda: scale} + "scale", + {"pyccel": no_pyccel, "python": lambda: scale}, ) x = np.ones(2) with pytest.warns(RuntimeWarning, match="uncompiled Python version"): @@ -446,7 +455,8 @@ def test_kernel_implementation_environment_variable(): def test_arrays_dispatch_calls_the_host_kernel_without_conversion( - fake_gpu, monkeypatch + fake_gpu, + monkeypatch, ): # CuPy is active, but host arguments never go through PyccelKernel's # device-to-host conversion: the choice already says they are host arrays @@ -474,7 +484,8 @@ def swapped(x, n, factor): kernel = Kernel( HostImplementations( - "scale", {"python": lambda: scale, "numpy": lambda: swapped} + "scale", + {"python": lambda: scale, "numpy": lambda: swapped}, ), CudaKernel(SCALE_CUDA, "scale"), ) diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index e3313ff..6ae8418 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -174,10 +174,12 @@ def test_assert_kernels_agree_on_gpu(): # the declared outputs of the host kernel are used by default declared = Kernel( - scale, CudaKernel(SCALE_CUDA, "scale"), host_options={"outputs": (0,)} + scale, + CudaKernel(SCALE_CUDA, "scale"), + host_options={"outputs": (0,)}, ) assert list(assert_kernels_agree(declared, make_scale_args, n_threads=300)) == [ - "argument 0" + "argument 0", ] wrong = Kernel(scale_wrong, CudaKernel(SCALE_CUDA, "scale"), name="scale") @@ -208,7 +210,8 @@ def test_assert_kernels_agree_restores_backend(): def test_device_function_kernel_source(): kernel = device_function_kernel( - FIND_SPAN, "int find_span(const double* t, int p, double eta)" + FIND_SPAN, + "int find_span(const double* t, int p, double eta)", ) assert kernel.name == "find_span_kernel" assert FIND_SPAN in kernel.source @@ -242,7 +245,7 @@ def test_device_function_kernel_options(): ) assert kernel.name == "fill_all" and kernel.block_size == 64 assert kernel.source.startswith( - '#include "helpers.cuh"\n#include \n' + '#include "helpers.cuh"\n#include \n', ) assert ( 'extern "C" __global__ void fill_all(' @@ -272,7 +275,9 @@ def test_device_function_kernel_errors(): device_function_kernel("", "double f(double out)") # the clash is resolved with another generated name kernel = device_function_kernel( - "", "double f(const double* x, int n)", n_threads_param="size" + "", + "double f(const double* x, int n)", + n_threads_param="size", ) assert ( "f_kernel(const double* x, const int* n, double* out, int size)" @@ -285,7 +290,8 @@ def test_device_function_kernel_on_gpu(): import cupy as cp sq = device_function_kernel( - "__device__ double sq(double x) { return x * x; }", "double sq(double x)" + "__device__ double sq(double x) { return x * x; }", + "double sq(double x)", ) x = cp.arange(1000, dtype=cp.float64) out = cp.empty(1000) @@ -293,7 +299,8 @@ def test_device_function_kernel_on_gpu(): assert cp.allclose(out, x * x) find_span = device_function_kernel( - FIND_SPAN, "int find_span(const double* t, int p, double eta)" + FIND_SPAN, + "int find_span(const double* t, int p, double eta)", ) t = cp.asarray([0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0]) eta = cp.asarray([0.1, 0.6, 0.9]) @@ -345,7 +352,9 @@ def weights(self): struct = DeviceArguments.struct value = xp.cuda.CudaStructValue( - struct, np.zeros((), struct.dtype)[()], vars(owner) | {"n": 3} + struct, + np.zeros((), struct.dtype)[()], + vars(owner) | {"n": 3}, ) assert sorted(_collect_arrays((value,))) == [ "argument 0.markers", diff --git a/tests/unit/test_mirror.py b/tests/unit/test_mirror.py index 5daa24b..769fba2 100644 --- a/tests/unit/test_mirror.py +++ b/tests/unit/test_mirror.py @@ -15,7 +15,8 @@ from cunumpy.memory import DeviceMirror requires_gpu = pytest.mark.skipif( - not xp.cupy_available(), reason="CuPy/GPU not available or not functional" + not xp.cupy_available(), + reason="CuPy/GPU not available or not functional", ) BIN_ADD = r""" @@ -247,7 +248,7 @@ def test_atomic_add_2d_matches_bincount(): weights = rng.random(n) flat = rows * shape[1] + cols expected = np.bincount(flat, weights=weights, minlength=np.prod(shape)).reshape( - shape + shape, ) out = np.zeros(shape) diff --git a/tests/unit/test_morton.py b/tests/unit/test_morton.py index 318fad4..89f3003 100644 --- a/tests/unit/test_morton.py +++ b/tests/unit/test_morton.py @@ -60,7 +60,7 @@ def test_keys_cells_and_clipping(): [0.999, 0.5], [1.0, 1.0], # upper face: the last cell [-5.0, 7.0], # outside: the nearest face - ] + ], ) keys = xp.algorithms.morton_keys(positions, [0.0, 0.0], [1.0, 1.0], levels) cells = [(0, 0), (1, 0), (7, 4), (7, 7), (0, 7)] @@ -71,7 +71,10 @@ def test_reversed_axis(): # lower > upper on y: cell 0 at the top, like a quadtree with y < mid as # its second quadrant bit keys = xp.algorithms.morton_keys( - np.array([[0.2, 0.9], [0.2, 0.1]]), [0, 1], [1, 0], 1 + np.array([[0.2, 0.9], [0.2, 0.1]]), + [0, 1], + [1, 0], + 1, ) assert keys.tolist() == [0, 2] @@ -94,7 +97,8 @@ def test_sorted_keys_make_tree_nodes_contiguous(): positions = rng.random((2000, 2)) levels = 10 keys, _, sorted_positions = xp.algorithms.sort_by_key( - xp.algorithms.morton_keys(positions, [0, 0], [1, 1], levels), positions + xp.algorithms.morton_keys(positions, [0, 0], [1, 1], levels), + positions, ) for level in (1, 2, 3): node = keys >> np.uint64(2 * (levels - level)) @@ -111,7 +115,9 @@ def test_sort_by_key_is_stable_and_reorders_all_arrays(): ids = np.arange(5) rows = np.arange(10.0).reshape(5, 2) sorted_keys, order, sorted_ids, sorted_rows = xp.algorithms.sort_by_key( - keys, ids, rows + keys, + ids, + rows, ) assert order.dtype == np.int64 assert order.tolist() == [1, 3, 2, 0, 4] @@ -167,7 +173,9 @@ def _cases(ndim): levels = xp.algorithms.MAX_MORTON_LEVELS[ndim] # points on cell boundaries, where rounding would show edges = np.array(lower) + np.arange(5)[:, None] / xp.algorithms.morton_scales( - lower, upper, levels + lower, + upper, + levels, ) positions = np.ascontiguousarray(np.vstack([positions, edges])) return positions, lower, upper, levels @@ -217,7 +225,8 @@ def test_cupy_arrays_stay_on_the_device(): keys = xp.algorithms.morton_keys(cp.asarray(positions), lower, upper, levels) assert isinstance(keys, cp.ndarray) np.testing.assert_array_equal( - cp.asnumpy(keys), xp.algorithms.morton_keys(positions, lower, upper, levels) + cp.asnumpy(keys), + xp.algorithms.morton_keys(positions, lower, upper, levels), ) sorted_keys, order, _ = xp.algorithms.sort_by_key(keys, cp.asarray(positions)) assert isinstance(order, cp.ndarray) diff --git a/tests/unit/test_mpi_serial.py b/tests/unit/test_mpi_serial.py index e233c1a..f441411 100644 --- a/tests/unit/test_mpi_serial.py +++ b/tests/unit/test_mpi_serial.py @@ -64,8 +64,10 @@ def test_slurm_batch_script_is_not_an_mpi_launch(clean_env): @pytest.mark.parametrize( ("value", "launcher", "expected"), - [("1", False, True), ("on", False, True), ("0", True, False), ("no", True, False), - ("maybe", True, True), ("maybe", False, False)], + [ + ("1", False, True), ("on", False, True), ("0", True, False), ("no", True, False), + ("maybe", True, True), ("maybe", False, False), + ], ) # fmt: skip def test_override(clean_env, value, launcher, expected): clean_env.setenv("CUNUMPY_MPI", value) @@ -261,7 +263,8 @@ def test_module_functions(): def test_constants_used_by_struphy_and_feectools(): assert isinstance(MPI.DOUBLE, MPI.Datatype) and not isinstance( - MPI.SUM, MPI.Datatype + MPI.SUM, + MPI.Datatype, ) assert isinstance(MPI.LOR, MPI.Op) assert MPI._typedict[np.dtype(np.float64).char] is MPI.DOUBLE @@ -274,8 +277,15 @@ def test_constants_used_by_struphy_and_feectools(): assert (MPI.Intracomm | None) is not None -@pytest.mark.parametrize("backend", ["numpy", pytest.param("cupy", marks=pytest.mark.skipif( - not xp.cupy_available(), reason="CuPy/GPU not available"))]) # fmt: skip +@pytest.mark.parametrize( + "backend", [ + "numpy", pytest.param( + "cupy", marks=pytest.mark.skipif( + not xp.cupy_available(), reason="CuPy/GPU not available", + ), + ), + ], +) # fmt: skip def test_buffers_on_either_backend(backend): with xp.use_backend(backend): send = xp.arange(4.0) diff --git a/tests/unit/test_mpi_staging.py b/tests/unit/test_mpi_staging.py index 655cfc8..30dbc3a 100644 --- a/tests/unit/test_mpi_staging.py +++ b/tests/unit/test_mpi_staging.py @@ -61,7 +61,9 @@ def allocate_device(shape, dtype): ), ) monkeypatch.setattr( - _mpi.array_api_compat, "is_cupy_array", lambda a: isinstance(a, Array) + _mpi.array_api_compat, + "is_cupy_array", + lambda a: isinstance(a, Array), ) def allocate(shape, dtype): @@ -132,7 +134,11 @@ def test_explicit_producer_dependencies_and_receive_only(device): data = Array([1.0, 2.0]) producer = SimpleNamespace(synchronize=lambda: log.append(("producer",))) with mpi_buffer( - data, send=False, recv=True, cuda_aware=False, event=producer + data, + send=False, + recv=True, + cuda_aware=False, + event=producer, ) as host: host[:] = [9, 10] assert ("producer",) in log diff --git a/tests/unit/test_petsc.py b/tests/unit/test_petsc.py index 7a328bc..294dee2 100644 --- a/tests/unit/test_petsc.py +++ b/tests/unit/test_petsc.py @@ -105,7 +105,9 @@ def setAttr(self, name, value): monkeypatch.setitem(sys.modules, "petsc4py", package) monkeypatch.setitem(sys.modules, "petsc4py.PETSc", PETSc) monkeypatch.setattr( - petsc, "_is_device_array", lambda a: isinstance(a, FakeDeviceArray) + petsc, + "_is_device_array", + lambda a: isinstance(a, FakeDeviceArray), ) return Vec diff --git a/tests/unit/test_philox.py b/tests/unit/test_philox.py index 9a91fba..b9c13a3 100644 --- a/tests/unit/test_philox.py +++ b/tests/unit/test_philox.py @@ -66,7 +66,9 @@ def test_streams_counters_and_seeds_differ(): def test_broadcasting(): u = xp.rng.philox_uniform( - np.uint64(3), np.arange(4, dtype=np.uint64)[:, None], np.arange(5) + np.uint64(3), + np.arange(4, dtype=np.uint64)[:, None], + np.arange(5), ) assert u.shape == (4, 5) assert u[2, 3] == xp.rng.philox_uniform(3, 2, 3) @@ -131,5 +133,6 @@ def run(kernel, *args, n_threads): du0, _ = xp.rng.philox_uniform2(7, cp.arange(10, dtype=cp.uint64), 1) assert isinstance(du0, cp.ndarray) np.testing.assert_array_equal( - du0.get(), xp.rng.philox_uniform2(7, np.arange(10, dtype=np.uint64), 1)[0] + du0.get(), + xp.rng.philox_uniform2(7, np.arange(10, dtype=np.uint64), 1)[0], ) diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index 2615c25..c695992 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -99,7 +99,8 @@ def test_arrays_on_another_device_are_rejected(monkeypatch): kernel.prepare_args(FakeDeviceArray(np.float64), 2.0, 4) # no device attribute kernel.prepare_args(FakeDeviceArray(np.float64, device=0), 2.0, 4) with pytest.raises( - ValueError, match="on CUDA device 1, but the current device is 0" + ValueError, + match="on CUDA device 1, but the current device is 0", ): kernel.prepare_args(FakeDeviceArray(np.float64, device=1), 2.0, 4) # struct pointer fields go through the same check @@ -112,7 +113,9 @@ def test_device_check_is_skipped_without_cupy(monkeypatch): monkeypatch.delitem(sys.modules, "cupy", raising=False) assert cuda_kernel_module._current_device_id() is None CudaKernel(SCALE, "scale").prepare_args( - FakeDeviceArray(np.float64, device=7), 2.0, 4 + FakeDeviceArray(np.float64, device=7), + 2.0, + 4, ) @@ -123,7 +126,9 @@ def test_device_check_is_skipped_without_cupy(monkeypatch): def test_from_pyccel_class_from_source_and_file(tmp_path): struct = CudaStruct.from_pyccel_class( - PYCCEL_SOURCE, "MarkerArguments", "MarkerArgs" + PYCCEL_SOURCE, + "MarkerArguments", + "MarkerArgs", ) assert struct.name == "MarkerArgs" # fields are named after the attributes the parameters are stored in @@ -140,18 +145,25 @@ def test_from_pyccel_class_from_source_and_file(tmp_path): assert from_file.name == "MarkerArguments" assert from_file.dtype == struct.dtype same = CudaStruct.from_pyccel_class( - str(path), "MarkerArguments", "A", int_type="int" + str(path), + "MarkerArguments", + "A", + int_type="int", ) assert [f.ctype for f in same.fields][2] == "int" def test_from_pyccel_class_options(): keep = CudaStruct.from_pyccel_class( - PYCCEL_SOURCE, "MarkerArguments", attribute_names=False + PYCCEL_SOURCE, + "MarkerArguments", + attribute_names=False, ) assert [f.name for f in keep.fields][3] == "first_pusher_idx" fewer = CudaStruct.from_pyccel_class( - PYCCEL_SOURCE, "MarkerArguments", exclude=("bc_type", "first_pusher_idx") + PYCCEL_SOURCE, + "MarkerArguments", + exclude=("bc_type", "first_pusher_idx"), ) assert [f.name for f in fewer.fields] == ["markers", "valid_mks", "Np"] final = CudaStruct.from_pyccel_class(PYCCEL_SOURCE, "DerhamArguments", "DerhamArgs") @@ -168,7 +180,8 @@ def test_from_pyccel_class_errors(): CudaStruct.from_pyccel_class("class Empty:\n pass\n", "Empty") with pytest.raises(ValueError, match="no type annotation"): CudaStruct.from_pyccel_class( - "class A:\n def __init__(self, x):\n pass\n", "A" + "class A:\n def __init__(self, x):\n pass\n", + "A", ) @@ -312,7 +325,7 @@ def __array__(self, dtype=None, copy=None): def test_device_arrays_in_struct_arguments(): args = MarkerArguments(FakeDeviceArray(np.float64, shape=(5, 3)), 5) found = dict( - cuda_kernel_module._device_arrays_in((1.0, args, FakeDeviceArray("f8"))) + cuda_kernel_module._device_arrays_in((1.0, args, FakeDeviceArray("f8"))), ) assert set(found) == {"argument 1.markers", "argument 2"} @@ -339,14 +352,14 @@ def helper_package(tmp_path, monkeypatch): (root / name).mkdir(parents=True) (root / name / "__init__.py").write_text("") (root / name / f"{name}_kernels.py").write_text( - f"def {name}(x, factor, n):\n for i in range(n):\n x[i] *= factor\n" + f"def {name}(x, factor, n):\n for i in range(n):\n x[i] *= factor\n", ) (root / "scale" / "scale_cuda.cu").write_text(SCALE) (root / "scale" / "scale_test_args.py").write_text( "import numpy as np\nimport cunumpy as xp\n\nN_THREADS = 300\nRTOL = 1e-10\n\n\n" "def make_args(backend, seed):\n" " x = xp.to_cunumpy(np.random.default_rng(seed).random(300))\n" - " return (x, 2.0, x.size)\n" + " return (x, 2.0, x.size)\n", ) (root / "__init__.py").write_text("") monkeypatch.syspath_prepend(str(tmp_path)) @@ -387,10 +400,10 @@ def test_parity_cases_marks_kernels_without_test_args(helper_package): root.mkdir() (root / "__init__.py").write_text("") (root / "with_cuda_no_args_kernels.py").write_text( - "def with_cuda_no_args(x, factor, n):\n pass\n" + "def with_cuda_no_args(x, factor, n):\n pass\n", ) (root / "with_cuda_no_args_cuda.cu").write_text( - SCALE.replace("scale", "with_cuda_no_args") + SCALE.replace("scale", "with_cuda_no_args"), ) catalog = KernelCatalog.from_package("helper_kernel_pkg") cases = parity_cases(catalog) @@ -427,7 +440,7 @@ def test_host_parameters_from_pyccel_stub(tmp_path): "from pyccel.decorators import low_level\n\n" "@low_level('push')\n" "def push(dt : 'float', stage : 'int', markers : 'float64[:,:](order=C)') -> None:\n" - " ...\n" + " ...\n", ) module = ModuleType("push_kernels") module.__file__ = str(so) @@ -457,7 +470,9 @@ def test_device_function_kernel_struct_parameters(): "{ return d.params[0] * x; }\n" ) kernel = device_function_kernel( - source, "double scale_x(const DomainArgs& d, double x)", structs=[domain] + source, + "double scale_x(const DomainArgs& d, double x)", + structs=[domain], ) params = [(p.name, p.ctype, p.struct is not None) for p in kernel.signature] assert params == [ @@ -483,7 +498,8 @@ def test_segment_sum(): keys = np.array([0, 2, 0, -1, 2]) values = np.array([1.0, 2.0, 3.0, 100.0, 4.0]) np.testing.assert_array_equal( - xp.algorithms.segment_sum(values, keys, 4), [4.0, 0.0, 6.0, 0.0] + xp.algorithms.segment_sum(values, keys, 4), + [4.0, 0.0, 6.0, 0.0], ) columns = np.stack([values, -values], axis=1) out = xp.algorithms.segment_sum(columns, keys, 3) diff --git a/tests/unit/test_profiling.py b/tests/unit/test_profiling.py index 58ca4cb..66c143b 100644 --- a/tests/unit/test_profiling.py +++ b/tests/unit/test_profiling.py @@ -22,7 +22,7 @@ def fake_nvtx(monkeypatch): nvtx = ModuleType("cupy.cuda.nvtx") nvtx.RangePush = lambda message, id_color=-1: calls.append( - ("push", message, id_color) + ("push", message, id_color), ) nvtx.RangePop = lambda: calls.append(("pop",)) cuda = ModuleType("cupy.cuda") diff --git a/tests/unit/test_pyccel_kernel.py b/tests/unit/test_pyccel_kernel.py index 8fd39f7..43113b8 100644 --- a/tests/unit/test_pyccel_kernel.py +++ b/tests/unit/test_pyccel_kernel.py @@ -389,7 +389,9 @@ def kernel(inp, pack, obj): obj.data[:] = 2.0 PyccelKernel(kernel, object_modules=(Container.__module__,), outputs=(1, 2))( - xp.to_cupy(np.ones(2)), [{"m": nested}], Container(held) + xp.to_cupy(np.ones(2)), + [{"m": nested}], + Container(held), ) assert np.array_equal(xp.to_numpy(nested), np.full(2, 1.0)) diff --git a/tests/unit/test_reusable_streams.py b/tests/unit/test_reusable_streams.py index c179fe1..a2a05d3 100644 --- a/tests/unit/test_reusable_streams.py +++ b/tests/unit/test_reusable_streams.py @@ -52,8 +52,10 @@ def record(self, stream=None): "cupy", SimpleNamespace( cuda=SimpleNamespace( - Stream=Stream, Event=Event, get_current_stream=lambda: current - ) + Stream=Stream, + Event=Event, + get_current_stream=lambda: current, + ), ), ) stream = xp.cuda.create_stream(non_blocking=False) diff --git a/tests/unit/test_scipy_backend.py b/tests/unit/test_scipy_backend.py index 9cef3e2..01f9777 100644 --- a/tests/unit/test_scipy_backend.py +++ b/tests/unit/test_scipy_backend.py @@ -75,7 +75,8 @@ def test_cupy_backend_forwards_to_cupyx(fake_cupyx): assert xp.scipy.sparse.linalg.cg == "device cg" assert xp.scipy.special.resolve() is fake_cupyx["cupyx.scipy.special"] with pytest.raises( - AttributeError, match=r"cupy backend \(it may exist in scipy\.special\)" + AttributeError, + match=r"cupy backend \(it may exist in scipy\.special\)", ): xp.scipy.special.erfcx # noqa: B018 # a subpackage cupyx does not provide @@ -98,7 +99,8 @@ def test_missing_scipy(monkeypatch): monkeypatch.setitem(sys.modules, "scipy", None) # import scipy raises ImportError monkeypatch.setitem(sys.modules, "scipy.special", None) with pytest.raises( - ImportError, match=r"needs scipy.special: SciPy is not installed" + ImportError, + match=r"needs scipy.special: SciPy is not installed", ): xp.scipy.special.erf # noqa: B018 assert not xp.scipy.special.available("erf") diff --git a/tests/unit/test_segment_plans.py b/tests/unit/test_segment_plans.py index 312bd6d..57b3164 100644 --- a/tests/unit/test_segment_plans.py +++ b/tests/unit/test_segment_plans.py @@ -14,18 +14,21 @@ def test_dense_and_sparse_boundaries(backend): with xp.use_backend(backend, strict=True): cells = xp.asarray([0, 0, 2, 4, 4], dtype=xp.int64) np.testing.assert_array_equal( - xp.to_numpy(cell_offsets(cells, 6)), [0, 2, 2, 3, 3, 5, 5] + xp.to_numpy(cell_offsets(cells, 6)), + [0, 2, 2, 3, 3, 5, 5], ) unique, starts, stops = segment_boundaries(cells) for got, expected in zip( - (unique, starts, stops), ([0, 2, 4], [0, 2, 3], [2, 3, 5]) + (unique, starts, stops), + ([0, 2, 4], [0, 2, 3], [2, 3, 5]), ): assert xp.get_array_backend(got) == backend np.testing.assert_array_equal(xp.to_numpy(got), expected) large = xp.asarray([2**63 + 1, 2**63 + 1, 2**64 - 1], dtype=xp.uint64) unique, starts, stops = segment_boundaries(large) np.testing.assert_array_equal( - xp.to_numpy(unique), np.array([2**63 + 1, 2**64 - 1], dtype=np.uint64) + xp.to_numpy(unique), + np.array([2**63 + 1, 2**64 - 1], dtype=np.uint64), ) np.testing.assert_array_equal(xp.to_numpy(starts), [0, 2]) np.testing.assert_array_equal(xp.to_numpy(stops), [2, 3]) diff --git a/tests/unit/test_staging.py b/tests/unit/test_staging.py index dd4fb80..679e02a 100644 --- a/tests/unit/test_staging.py +++ b/tests/unit/test_staging.py @@ -97,10 +97,14 @@ def empty(shape, dtype): fake = types.SimpleNamespace(cuda=cuda, empty=empty) monkeypatch.setattr(staging_module, "_cupy", lambda: fake) monkeypatch.setattr( - staging_module, "_empty_pinned", lambda s, dtype: np.empty(s, dtype) + staging_module, + "_empty_pinned", + lambda s, dtype: np.empty(s, dtype), ) monkeypatch.setattr( - staging_module, "_is_device_array", lambda a: isinstance(a, DeviceArray) + staging_module, + "_is_device_array", + lambda a: isinstance(a, DeviceArray), ) return DeviceArray.log diff --git a/tests/unit/test_transfers.py b/tests/unit/test_transfers.py index c437bf0..a431c69 100644 --- a/tests/unit/test_transfers.py +++ b/tests/unit/test_transfers.py @@ -23,7 +23,8 @@ THIS_FILE = str(Path(__file__)) requires_cupy = pytest.mark.skipif( - not xp.cupy_available(), reason="CuPy not installed or not functional" + not xp.cupy_available(), + reason="CuPy not installed or not functional", ) @@ -58,7 +59,9 @@ def get_array_backend(array): monkeypatch.setattr(xp_module, "get_array_backend", get_array_backend) monkeypatch.setattr(xp_module, "_to_cupy", _FakeDeviceArray) monkeypatch.setattr( - kernel_module, "_is_device_array", lambda v: isinstance(v, _FakeDeviceArray) + kernel_module, + "_is_device_array", + lambda v: isinstance(v, _FakeDeviceArray), ) monkeypatch.setattr(kernel_module, "_device_to_host", lambda v: v.get()) monkeypatch.setattr(kernel_module, "_host_to_device", _FakeDeviceArray) From ac7444da83a3d7c0a9ec900c71485037d7ffc42e Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:33:02 +0200 Subject: [PATCH 4/9] ruff format --- src/cunumpy/__init__.py | 24 ++++++++--- src/cunumpy/_cuda_kernel.py | 9 ++-- src/cunumpy/_dispatch.py | 8 +++- src/cunumpy/_emulation.py | 9 +++- src/cunumpy/_mpi.py | 5 ++- src/cunumpy/algorithms.py | 18 ++++++-- src/cunumpy/cuda/__init__.py | 52 +++++++++++++++++------ src/cunumpy/kernel_testing.py | 12 ++++-- src/cunumpy/kernels.py | 18 +++++--- src/cunumpy/mpi.py | 25 ++++++++--- src/cunumpy/profiling.py | 8 +++- src/cunumpy/rng.py | 12 ++++-- tests/unit/test_cuda_kernel.py | 24 ++++++++--- tests/unit/test_kernel_dispatch.py | 9 +++- tests/unit/test_kernel_dispatch_arrays.py | 9 +++- tests/unit/test_kernel_testing.py | 14 +++--- tests/unit/test_porting_helpers.py | 3 +- tests/unit/test_segment_plans.py | 8 +++- 18 files changed, 197 insertions(+), 70 deletions(-) diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index e6f6ecb..b0e9da4 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -5,11 +5,25 @@ from . import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp from ._scipy_backend import scipy -from .xp import (as_device_array, assert_same_backend, backend_info, - cupy_available, default_float_dtype, get_array_backend, - get_array_module, get_backend, is_cpu, is_gpu, same_backend, - set_backend, synchronize, to_cunumpy, to_cupy, to_numpy, - use_backend) +from .xp import ( + as_device_array, + assert_same_backend, + backend_info, + cupy_available, + default_float_dtype, + get_array_backend, + get_array_module, + get_backend, + is_cpu, + is_gpu, + same_backend, + set_backend, + synchronize, + to_cunumpy, + to_cupy, + to_numpy, + use_backend, +) # Names that were at the top level before cunumpy 0.5, and the submodule each # moved to. They still resolve (with a DeprecationWarning) until cunumpy 0.6. diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 33c0bc2..344d88a 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -54,8 +54,7 @@ class is the one definition of the arguments. import re import sys import typing -from collections.abc import (Callable, Hashable, Iterable, Iterator, Mapping, - Sequence) +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from pathlib import Path @@ -2404,8 +2403,10 @@ def _opt_in_shared_memory(self, kernel: Any, shared_mem: int) -> None: attribute ``max_dynamic_shared_size_bytes``; it is set once (and again for a larger request) up to the device's opt-in limit. """ - from ._device import (DEFAULT_SHARED_MEMORY_PER_BLOCK, - max_shared_memory_per_block) + from ._device import ( + DEFAULT_SHARED_MEMORY_PER_BLOCK, + max_shared_memory_per_block, + ) if shared_mem <= DEFAULT_SHARED_MEMORY_PER_BLOCK: return diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index 4782b77..ecbcb91 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -35,8 +35,12 @@ import array_api_compat from ._cuda_kernel import CudaKernel, _compile_in_threads -from ._kernel import (CompiledHostKernel, HostImplementations, PyccelKernel, - resolve_host_args) +from ._kernel import ( + CompiledHostKernel, + HostImplementations, + PyccelKernel, + resolve_host_args, +) from ._transfers import _ACTIVE as _COUNTERS from ._transfers import _record from .xp import get_backend diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index a90f364..7ce3282 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -55,8 +55,13 @@ import numpy as np -from ._cuda_kernel import (CudaKernel, CudaParameter, _scalar_checker, - _strip_comments, cuda_include_dir) +from ._cuda_kernel import ( + CudaKernel, + CudaParameter, + _scalar_checker, + _strip_comments, + cuda_include_dir, +) __all__ = ["emulate_cuda_kernel", "emulation_compiler"] diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 835ca95..6b1b158 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -11,8 +11,9 @@ import array_api_compat import array_api_compat.numpy as np -from ._mpi_serial import _LOCAL_RANK_VARIABLES # noqa: F401 - re-exported -from ._mpi_serial import local_rank +from ._mpi_serial import ( + _LOCAL_RANK_VARIABLES, # noqa: F401 - re-exported +) from ._transfers import _ACTIVE as _COUNTERS from ._transfers import _describe, _record from .xp import array_backend, cupy_available, to_numpy diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py index 330310f..b974e03 100644 --- a/src/cunumpy/algorithms.py +++ b/src/cunumpy/algorithms.py @@ -10,10 +10,20 @@ charge = xp.algorithms.segment_sum(q, cell, n_cells) """ -from ._algorithms import (SegmentPlan, cell_offsets, segment_boundaries, - segment_sum, sort_by_key) -from ._morton import (MAX_MORTON_LEVELS, morton_decode, morton_encode, - morton_keys, morton_scales) +from ._algorithms import ( + SegmentPlan, + cell_offsets, + segment_boundaries, + segment_sum, + sort_by_key, +) +from ._morton import ( + MAX_MORTON_LEVELS, + morton_decode, + morton_encode, + morton_keys, + morton_scales, +) __all__ = [ "MAX_MORTON_LEVELS", diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index dcb1654..2aca624 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -24,18 +24,46 @@ use ``import cupy; cupy.cuda`` for CuPy's module. """ -from .._cuda_kernel import (DEBUG_OPTIONS, CudaArguments, CudaKernel, - CudaKernelVariants, CudaParameter, CudaStruct, - CudaStructArguments, CudaStructValue, ctype_of, - cuda_include_dir, cuda_kernel_names, include_hash, - parse_cuda_signature, resolve_includes, - write_cuda_header) -from .._device import (DEFAULT_SHARED_MEMORY_PER_BLOCK, bind_local_device, - cuda_debug, device_count, free_memory, get_cuda_debug, - max_shared_memory_per_block, memory_info, pin_memory, - set_cuda_debug, set_device, set_device_for_rank, stream) -from .._streams import (HostEvent, HostStream, create_event, create_stream, - record_event, wait_event) +from .._cuda_kernel import ( + DEBUG_OPTIONS, + CudaArguments, + CudaKernel, + CudaKernelVariants, + CudaParameter, + CudaStruct, + CudaStructArguments, + CudaStructValue, + ctype_of, + cuda_include_dir, + cuda_kernel_names, + include_hash, + parse_cuda_signature, + resolve_includes, + write_cuda_header, +) +from .._device import ( + DEFAULT_SHARED_MEMORY_PER_BLOCK, + bind_local_device, + cuda_debug, + device_count, + free_memory, + get_cuda_debug, + max_shared_memory_per_block, + memory_info, + pin_memory, + set_cuda_debug, + set_device, + set_device_for_rank, + stream, +) +from .._streams import ( + HostEvent, + HostStream, + create_event, + create_stream, + record_event, + wait_event, +) __all__ = [ "DEBUG_OPTIONS", diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 2db7b2e..4236582 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -57,9 +57,15 @@ def test_parity(name, kernel): import numpy as np from . import _fake_cupy -from ._cuda_kernel import (CudaKernel, CudaParameter, CudaStructArguments, - CudaStructValue, _parse_parameter, _split_top_level, - _strip_comments) +from ._cuda_kernel import ( + CudaKernel, + CudaParameter, + CudaStructArguments, + CudaStructValue, + _parse_parameter, + _split_top_level, + _strip_comments, +) from ._dispatch import Kernel from ._emulation import emulate_cuda_kernel, emulation_compiler from .xp import cupy_available, get_backend, to_numpy, use_backend diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 39f47b6..9149fb5 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -22,11 +22,19 @@ from ._cuda_kernel import PyccelStructArguments from ._dispatch import Kernel, KernelCatalog from ._fusion import fuse -from ._kernel import (HOST_IMPLEMENTATIONS, CompiledHostKernel, - HostImplementations, KernelArguments, PyccelKernel, - as_kernel_array, get_kernel_implementation, - kernel_output, resolve_host_args, - set_kernel_implementation, use_kernel_implementation) +from ._kernel import ( + HOST_IMPLEMENTATIONS, + CompiledHostKernel, + HostImplementations, + KernelArguments, + PyccelKernel, + as_kernel_array, + get_kernel_implementation, + kernel_output, + resolve_host_args, + set_kernel_implementation, + use_kernel_implementation, +) __all__ = [ "HOST_IMPLEMENTATIONS", diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index 6cc34ac..b9d8206 100644 --- a/src/cunumpy/mpi.py +++ b/src/cunumpy/mpi.py @@ -21,12 +21,25 @@ mpi4py is imported only by the functions that need it. """ -from ._mpi import (MPIStaging, get_mpi_cuda_aware, mpi_buffer, - mpi_is_cuda_aware, require_cuda_aware_mpi, - set_mpi_cuda_aware, synchronize_for_mpi) -from ._mpi_serial import (OVERRIDE_VARIABLE, SerialComm, SerialMPI, - SerialRequest, SerialStatus, get_mpi, - launched_under_mpi, local_rank) +from ._mpi import ( + MPIStaging, + get_mpi_cuda_aware, + mpi_buffer, + mpi_is_cuda_aware, + require_cuda_aware_mpi, + set_mpi_cuda_aware, + synchronize_for_mpi, +) +from ._mpi_serial import ( + OVERRIDE_VARIABLE, + SerialComm, + SerialMPI, + SerialRequest, + SerialStatus, + get_mpi, + launched_under_mpi, + local_rank, +) __all__ = [ "OVERRIDE_VARIABLE", diff --git a/src/cunumpy/profiling.py b/src/cunumpy/profiling.py index 2b18a8a..7fb3ee7 100644 --- a/src/cunumpy/profiling.py +++ b/src/cunumpy/profiling.py @@ -14,8 +14,12 @@ """ from ._profiling import Timing, nvtx_range, timed_region -from ._transfers import (TransferCounter, TransferEvent, assert_no_transfers, - count_transfers) +from ._transfers import ( + TransferCounter, + TransferEvent, + assert_no_transfers, + count_transfers, +) __all__ = [ "Timing", diff --git a/src/cunumpy/rng.py b/src/cunumpy/rng.py index 94f9c93..69ca9c1 100644 --- a/src/cunumpy/rng.py +++ b/src/cunumpy/rng.py @@ -16,10 +16,14 @@ ``random`` module. """ -from ._philox import (philox4x32_10, philox_normal, philox_normal2, - philox_uniform, philox_uniform2) -from ._random_streams import (BIT_GENERATORS, RandomStreams, get_rng, - random_streams) +from ._philox import ( + philox4x32_10, + philox_normal, + philox_normal2, + philox_uniform, + philox_uniform2, +) +from ._random_streams import BIT_GENERATORS, RandomStreams, get_rng, random_streams __all__ = [ "BIT_GENERATORS", diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index b2af002..83eed65 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -15,11 +15,20 @@ import cunumpy as xp import cunumpy._device as device_module from cunumpy import as_device_array -from cunumpy.cuda import (CudaArguments, CudaKernel, CudaKernelVariants, - CudaStruct, CudaStructValue, ctype_of, - cuda_include_dir, cuda_kernel_names, include_hash, - parse_cuda_signature, resolve_includes, - write_cuda_header) +from cunumpy.cuda import ( + CudaArguments, + CudaKernel, + CudaKernelVariants, + CudaStruct, + CudaStructValue, + ctype_of, + cuda_include_dir, + cuda_kernel_names, + include_hash, + parse_cuda_signature, + resolve_includes, + write_cuda_header, +) AXPY = r""" // y = a * x + y @@ -463,7 +472,9 @@ def test_ctype_of(): ], ) -PUSH_SOURCE = PARTICLES.declaration + r""" +PUSH_SOURCE = ( + PARTICLES.declaration + + r""" extern "C" __global__ void push(Particles p, double dt, double* out, unsigned long long* size) { int i = blockDim.x * blockIdx.x + threadIdx.x; @@ -474,6 +485,7 @@ def test_ctype_of(): if (i < p.n && p.alive[i]) p.x[i] += dt * p.charge; } """ +) def test_struct_layout_and_declaration(): diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py index 96d0506..9bebfde 100644 --- a/tests/unit/test_kernel_dispatch.py +++ b/tests/unit/test_kernel_dispatch.py @@ -15,8 +15,13 @@ import cunumpy as xp from cunumpy.cuda import CudaKernel -from cunumpy.kernels import (Kernel, KernelArguments, KernelCatalog, - PyccelKernel, resolve_host_args) +from cunumpy.kernels import ( + Kernel, + KernelArguments, + KernelCatalog, + PyccelKernel, + resolve_host_args, +) SCALE_CUDA = r""" extern "C" __global__ void scale(double* x, double factor, int n) { diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index a0ca981..8a09c27 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -12,8 +12,13 @@ import cunumpy as xp from cunumpy import _dispatch as dispatch_module from cunumpy.cuda import CudaArguments, CudaKernel -from cunumpy.kernels import (CompiledHostKernel, HostImplementations, Kernel, - KernelArguments, KernelCatalog) +from cunumpy.kernels import ( + CompiledHostKernel, + HostImplementations, + Kernel, + KernelArguments, + KernelCatalog, +) SCALE_CUDA = r""" extern "C" __global__ void scale(double* x, double factor, int n) { diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index 6ae8418..af7fc1b 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -16,11 +16,15 @@ import cunumpy as xp import cunumpy.kernel_testing from cunumpy.cuda import CudaArguments, CudaKernel, parse_cuda_signature -from cunumpy.kernel_testing import \ - backend # noqa: F401 - the fixture is used by name -from cunumpy.kernel_testing import (BACKENDS, _collect_arrays, - _compare_results, assert_kernels_agree, - device_function_kernel, requires_cupy) +from cunumpy.kernel_testing import ( + BACKENDS, + _collect_arrays, + _compare_results, + assert_kernels_agree, + backend, # noqa: F401 - the fixture is used by name + device_function_kernel, + requires_cupy, +) from cunumpy.kernels import Kernel SCALE_CUDA = r""" diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index c695992..dbdfaa8 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -26,8 +26,7 @@ import cunumpy.kernel_testing from cunumpy._dispatch import FORTRAN_NAME_LIMIT, _pyccel_stub_parameters from cunumpy.cuda import CudaKernel, CudaStruct -from cunumpy.kernel_testing import (check_parity, device_function_kernel, - parity_cases) +from cunumpy.kernel_testing import check_parity, device_function_kernel, parity_cases from cunumpy.kernels import Kernel, KernelCatalog, PyccelStructArguments diff --git a/tests/unit/test_segment_plans.py b/tests/unit/test_segment_plans.py index 57b3164..6301d5c 100644 --- a/tests/unit/test_segment_plans.py +++ b/tests/unit/test_segment_plans.py @@ -4,8 +4,12 @@ import pytest import cunumpy as xp -from cunumpy.algorithms import (SegmentPlan, cell_offsets, segment_boundaries, - segment_sum) +from cunumpy.algorithms import ( + SegmentPlan, + cell_offsets, + segment_boundaries, + segment_sum, +) from cunumpy.kernel_testing import BACKENDS From 21cba12592fc6983ca5ca48e82d44b0b35fea090 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:43:21 +0200 Subject: [PATCH 5/9] Fix broken import --- src/cunumpy/_device.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index 79a094c..d8c1686 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -9,7 +9,7 @@ import array_api_compat.numpy as np -from ._mpi import local_rank +from ._mpi_serial import local_rank from .xp import array_backend, cupy_available From 12df656a785365a570339026d839b7161cb23d19 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:48:45 +0200 Subject: [PATCH 6/9] Use absolute imports --- pyproject.toml | 5 ++++- src/cunumpy/__init__.py | 6 +++--- src/cunumpy/__init__.pyi | 20 ++++++++++---------- src/cunumpy/_algorithms.py | 4 ++-- src/cunumpy/_cuda_kernel.py | 10 +++++----- src/cunumpy/_device.py | 4 ++-- src/cunumpy/_dispatch.py | 10 +++++----- src/cunumpy/_emulation.py | 2 +- src/cunumpy/_fusion.py | 2 +- src/cunumpy/_kernel.py | 6 +++--- src/cunumpy/_mirror.py | 2 +- src/cunumpy/_morton.py | 2 +- src/cunumpy/_mpi.py | 10 +++++----- src/cunumpy/_profiling.py | 2 +- src/cunumpy/_random_streams.py | 2 +- src/cunumpy/_scipy_backend.py | 2 +- src/cunumpy/_staging.py | 4 ++-- src/cunumpy/_streams.py | 2 +- src/cunumpy/algorithms.py | 4 ++-- src/cunumpy/cuda/__init__.py | 6 +++--- src/cunumpy/cuda_kernel.py | 2 +- src/cunumpy/dispatch.py | 2 +- src/cunumpy/kernel.py | 2 +- src/cunumpy/kernel_testing.py | 10 +++++----- src/cunumpy/kernels.py | 8 ++++---- src/cunumpy/memory.py | 4 ++-- src/cunumpy/mpi.py | 4 ++-- src/cunumpy/profiling.py | 4 ++-- src/cunumpy/rng.py | 9 +++++++-- src/cunumpy/testing.py | 2 +- src/cunumpy/xp.py | 6 +++--- 31 files changed, 83 insertions(+), 75 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ac9a151..60f12bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,5 +56,8 @@ cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ] # agent worktrees and scratch scripts are local copies, not part of the project extend-exclude = [ ".claude" ] # formatting: `ruff format`; import sorting: the isort rules of `ruff check` -lint.extend-select = [ "I" ] +lint.extend-select = [ "I", "TID252" ] lint.isort.known-first-party = [ "cunumpy" ] + +[tool.ruff.lint.flake8-tidy-imports] +ban-relative-imports = "all" diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index b0e9da4..24297ab 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -3,9 +3,9 @@ import warnings as _warnings from importlib.metadata import PackageNotFoundError, version -from . import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp -from ._scipy_backend import scipy -from .xp import ( +from cunumpy import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp +from cunumpy._scipy_backend import scipy +from cunumpy.xp import ( as_device_array, assert_same_backend, backend_info, diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index eadad06..ad45ff0 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -8,16 +8,16 @@ from typing import Any import numpy as np from numpy import * -from . import algorithms as algorithms -from . import cuda as cuda -from . import kernels as kernels -from . import memory as memory -from . import mpi as mpi -from . import petsc as petsc -from . import profiling as profiling -from . import rng as rng -from . import xp as xp -from ._scipy_backend import scipy as scipy +from cunumpy import algorithms as algorithms +from cunumpy import cuda as cuda +from cunumpy import kernels as kernels +from cunumpy import memory as memory +from cunumpy import mpi as mpi +from cunumpy import petsc as petsc +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 def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index 5f55b5a..353c83a 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -8,7 +8,7 @@ import array_api_compat.numpy as np -from .xp import assert_same_backend, get_array_backend, get_array_module +from cunumpy.xp import assert_same_backend, get_array_backend, get_array_module def _integer_keys(keys: Any) -> tuple[Any, Any]: @@ -77,7 +77,7 @@ def cell_offsets(sorted_cells: Any, n_cells: int) -> Any: def _sum_kernel(dtype: Any) -> Any: name = np.dtype(dtype).name if name not in _SUM_KERNELS: - from ._cuda_kernel import CudaKernel + from cunumpy._cuda_kernel import CudaKernel ctype = "float" if name == "float32" else "double" _SUM_KERNELS[name] = CudaKernel( diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 344d88a..e022bed 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -1362,7 +1362,7 @@ def verify_layout( RuntimeError If CuPy is not available. """ - from .xp import cupy_available + from cunumpy.xp import cupy_available if not cupy_available(): raise RuntimeError("verify_layout() compiles a CUDA kernel and needs CuPy") @@ -1768,7 +1768,7 @@ def __host_args__(self) -> Any: "or set host_copies = True for a read-only host evaluation " "from host copies", ) - from .xp import to_numpy + from cunumpy.xp import to_numpy value = to_numpy(value) host_values.append(value) @@ -2093,7 +2093,7 @@ def debug_active(self) -> bool: """ if self._debug is not None: return self._debug - from ._device import get_cuda_debug + from cunumpy._device import get_cuda_debug return get_cuda_debug() @@ -2234,7 +2234,7 @@ def compile(self) -> Any: If CuPy or a GPU is not available. """ if self._raw_kernel is None: - from .xp import cupy_available + from cunumpy.xp import cupy_available if not cupy_available(): raise RuntimeError( @@ -2403,7 +2403,7 @@ def _opt_in_shared_memory(self, kernel: Any, shared_mem: int) -> None: attribute ``max_dynamic_shared_size_bytes``; it is set once (and again for a larger request) up to the device's opt-in limit. """ - from ._device import ( + from cunumpy._device import ( DEFAULT_SHARED_MEMORY_PER_BLOCK, max_shared_memory_per_block, ) diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index d8c1686..bb13907 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -9,8 +9,8 @@ import array_api_compat.numpy as np -from ._mpi_serial import local_rank -from .xp import array_backend, cupy_available +from cunumpy._mpi_serial import local_rank +from cunumpy.xp import array_backend, cupy_available def set_device(device_id: int) -> None: diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index ecbcb91..7a22014 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -34,16 +34,16 @@ import array_api_compat -from ._cuda_kernel import CudaKernel, _compile_in_threads -from ._kernel import ( +from cunumpy._cuda_kernel import CudaKernel, _compile_in_threads +from cunumpy._kernel import ( CompiledHostKernel, HostImplementations, PyccelKernel, resolve_host_args, ) -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _record -from .xp import get_backend +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record +from cunumpy.xp import get_backend __all__ = ["Kernel", "KernelCatalog"] diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index 7ce3282..1cbf4d3 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -55,7 +55,7 @@ import numpy as np -from ._cuda_kernel import ( +from cunumpy._cuda_kernel import ( CudaKernel, CudaParameter, _scalar_checker, diff --git a/src/cunumpy/_fusion.py b/src/cunumpy/_fusion.py index 2299b4e..68541b8 100644 --- a/src/cunumpy/_fusion.py +++ b/src/cunumpy/_fusion.py @@ -32,7 +32,7 @@ def pressure(rho, T, gamma): import array_api_compat import numpy -from .xp import use_backend +from cunumpy.xp import use_backend __all__ = ["fuse"] diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index a05ca6e..1c74c1a 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -34,9 +34,9 @@ import array_api_compat import numpy as np -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _record -from .xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record +from cunumpy.xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy __all__ = ["CompiledHostKernel", "KernelArguments", "PyccelKernel", "resolve_host_args"] diff --git a/src/cunumpy/_mirror.py b/src/cunumpy/_mirror.py index d10294b..147735c 100644 --- a/src/cunumpy/_mirror.py +++ b/src/cunumpy/_mirror.py @@ -18,7 +18,7 @@ import numpy as np -from .xp import _cupy_backend, to_cupy +from cunumpy.xp import _cupy_backend, to_cupy __all__ = ["DeviceMirror"] diff --git a/src/cunumpy/_morton.py b/src/cunumpy/_morton.py index de39881..30e15ff 100644 --- a/src/cunumpy/_morton.py +++ b/src/cunumpy/_morton.py @@ -27,7 +27,7 @@ import numpy as np -from ._philox import _module +from cunumpy._philox import _module __all__ = [ "MAX_MORTON_LEVELS", diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 6b1b158..9f9541f 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -11,12 +11,12 @@ import array_api_compat import array_api_compat.numpy as np -from ._mpi_serial import ( +from cunumpy._mpi_serial import ( _LOCAL_RANK_VARIABLES, # noqa: F401 - re-exported ) -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _describe, _record -from .xp import array_backend, cupy_available, to_numpy +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _record +from cunumpy.xp import array_backend, cupy_available, to_numpy _logger = logging.getLogger(__name__) @@ -43,7 +43,7 @@ def synchronize_for_mpi(*arrays: Any, stream: Any = None, event: Any = None) -> devices = {a.device.id for a in arrays if array_api_compat.is_cupy_array(a)} if not devices: return - from ._streams import HostEvent, HostStream + from cunumpy._streams import HostEvent, HostStream if isinstance(event, HostEvent) or isinstance(stream, HostStream): raise TypeError("device buffers require a CUDA producer stream or event") diff --git a/src/cunumpy/_profiling.py b/src/cunumpy/_profiling.py index e6402f7..586109a 100644 --- a/src/cunumpy/_profiling.py +++ b/src/cunumpy/_profiling.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from types import ModuleType -from .xp import array_backend, synchronize +from cunumpy.xp import array_backend, synchronize def _nvtx_module() -> ModuleType | None: diff --git a/src/cunumpy/_random_streams.py b/src/cunumpy/_random_streams.py index 66daa3a..002ee86 100644 --- a/src/cunumpy/_random_streams.py +++ b/src/cunumpy/_random_streams.py @@ -34,7 +34,7 @@ import numpy as np -from .xp import get_backend +from cunumpy.xp import get_backend __all__ = ["BIT_GENERATORS", "RandomStreams", "random_streams"] diff --git a/src/cunumpy/_scipy_backend.py b/src/cunumpy/_scipy_backend.py index 1515563..64a0826 100644 --- a/src/cunumpy/_scipy_backend.py +++ b/src/cunumpy/_scipy_backend.py @@ -28,7 +28,7 @@ from types import ModuleType from typing import Any -from .xp import get_backend +from cunumpy.xp import get_backend __all__ = ["SUBMODULES", "ScipyNamespace", "scipy"] diff --git a/src/cunumpy/_staging.py b/src/cunumpy/_staging.py index 9faac25..b465c9b 100644 --- a/src/cunumpy/_staging.py +++ b/src/cunumpy/_staging.py @@ -34,8 +34,8 @@ import array_api_compat import numpy as np -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _record +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record __all__ = ["HostStaging", "StagedCopy"] diff --git a/src/cunumpy/_streams.py b/src/cunumpy/_streams.py index 1d1f2db..706214f 100644 --- a/src/cunumpy/_streams.py +++ b/src/cunumpy/_streams.py @@ -4,7 +4,7 @@ from typing import Any -from .xp import get_backend +from cunumpy.xp import get_backend class HostEvent: diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py index b974e03..c98421f 100644 --- a/src/cunumpy/algorithms.py +++ b/src/cunumpy/algorithms.py @@ -10,14 +10,14 @@ charge = xp.algorithms.segment_sum(q, cell, n_cells) """ -from ._algorithms import ( +from cunumpy._algorithms import ( SegmentPlan, cell_offsets, segment_boundaries, segment_sum, sort_by_key, ) -from ._morton import ( +from cunumpy._morton import ( MAX_MORTON_LEVELS, morton_decode, morton_encode, diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index 2aca624..412d548 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -24,7 +24,7 @@ use ``import cupy; cupy.cuda`` for CuPy's module. """ -from .._cuda_kernel import ( +from cunumpy._cuda_kernel import ( DEBUG_OPTIONS, CudaArguments, CudaKernel, @@ -41,7 +41,7 @@ resolve_includes, write_cuda_header, ) -from .._device import ( +from cunumpy._device import ( DEFAULT_SHARED_MEMORY_PER_BLOCK, bind_local_device, cuda_debug, @@ -56,7 +56,7 @@ set_device_for_rank, stream, ) -from .._streams import ( +from cunumpy._streams import ( HostEvent, HostStream, create_event, diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index cda7fac..516e70d 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -3,6 +3,6 @@ The public names are in :mod:`cunumpy.cuda`. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy._cuda_kernel", "cunumpy.cuda") diff --git a/src/cunumpy/dispatch.py b/src/cunumpy/dispatch.py index ceb788d..8781b5b 100644 --- a/src/cunumpy/dispatch.py +++ b/src/cunumpy/dispatch.py @@ -3,6 +3,6 @@ The public names are in :mod:`cunumpy.kernels`. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy._dispatch", "cunumpy.kernels") diff --git a/src/cunumpy/kernel.py b/src/cunumpy/kernel.py index 2780d36..d0b8717 100644 --- a/src/cunumpy/kernel.py +++ b/src/cunumpy/kernel.py @@ -3,6 +3,6 @@ The public names are in :mod:`cunumpy.kernels`. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy._kernel", "cunumpy.kernels") diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 4236582..50d3622 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -56,8 +56,8 @@ def test_parity(name, kernel): import array_api_compat import numpy as np -from . import _fake_cupy -from ._cuda_kernel import ( +from cunumpy import _fake_cupy +from cunumpy._cuda_kernel import ( CudaKernel, CudaParameter, CudaStructArguments, @@ -66,9 +66,9 @@ def test_parity(name, kernel): _split_top_level, _strip_comments, ) -from ._dispatch import Kernel -from ._emulation import emulate_cuda_kernel, emulation_compiler -from .xp import cupy_available, get_backend, to_numpy, use_backend +from cunumpy._dispatch import Kernel +from cunumpy._emulation import emulate_cuda_kernel, emulation_compiler +from cunumpy.xp import cupy_available, get_backend, to_numpy, use_backend # the pytest objects are created on first access, see __getattr__ __all__ = [ diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 9149fb5..7f276c4 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -19,10 +19,10 @@ :mod:`cunumpy.kernel_testing`. """ -from ._cuda_kernel import PyccelStructArguments -from ._dispatch import Kernel, KernelCatalog -from ._fusion import fuse -from ._kernel import ( +from cunumpy._cuda_kernel import PyccelStructArguments +from cunumpy._dispatch import Kernel, KernelCatalog +from cunumpy._fusion import fuse +from cunumpy._kernel import ( HOST_IMPLEMENTATIONS, CompiledHostKernel, HostImplementations, diff --git a/src/cunumpy/memory.py b/src/cunumpy/memory.py index 27c91e3..7fb400b 100644 --- a/src/cunumpy/memory.py +++ b/src/cunumpy/memory.py @@ -16,7 +16,7 @@ :mod:`cunumpy.cuda`. """ -from ._mirror import DeviceMirror -from ._staging import HostStaging, StagedCopy +from cunumpy._mirror import DeviceMirror +from cunumpy._staging import HostStaging, StagedCopy __all__ = ["DeviceMirror", "HostStaging", "StagedCopy"] diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index b9d8206..f5a82b7 100644 --- a/src/cunumpy/mpi.py +++ b/src/cunumpy/mpi.py @@ -21,7 +21,7 @@ mpi4py is imported only by the functions that need it. """ -from ._mpi import ( +from cunumpy._mpi import ( MPIStaging, get_mpi_cuda_aware, mpi_buffer, @@ -30,7 +30,7 @@ set_mpi_cuda_aware, synchronize_for_mpi, ) -from ._mpi_serial import ( +from cunumpy._mpi_serial import ( OVERRIDE_VARIABLE, SerialComm, SerialMPI, diff --git a/src/cunumpy/profiling.py b/src/cunumpy/profiling.py index 7fb3ee7..662e797 100644 --- a/src/cunumpy/profiling.py +++ b/src/cunumpy/profiling.py @@ -13,8 +13,8 @@ step() """ -from ._profiling import Timing, nvtx_range, timed_region -from ._transfers import ( +from cunumpy._profiling import Timing, nvtx_range, timed_region +from cunumpy._transfers import ( TransferCounter, TransferEvent, assert_no_transfers, diff --git a/src/cunumpy/rng.py b/src/cunumpy/rng.py index 69ca9c1..4f31873 100644 --- a/src/cunumpy/rng.py +++ b/src/cunumpy/rng.py @@ -16,14 +16,19 @@ ``random`` module. """ -from ._philox import ( +from cunumpy._philox import ( philox4x32_10, philox_normal, philox_normal2, philox_uniform, philox_uniform2, ) -from ._random_streams import BIT_GENERATORS, RandomStreams, get_rng, random_streams +from cunumpy._random_streams import ( + BIT_GENERATORS, + RandomStreams, + get_rng, + random_streams, +) __all__ = [ "BIT_GENERATORS", diff --git a/src/cunumpy/testing.py b/src/cunumpy/testing.py index 8cd33d0..9f3593a 100644 --- a/src/cunumpy/testing.py +++ b/src/cunumpy/testing.py @@ -4,6 +4,6 @@ ``xp.testing.assert_allclose`` failed after any ``import cunumpy.testing``. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy.kernel_testing", "cunumpy.kernel_testing") diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 5a56f1c..3da3621 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -13,12 +13,12 @@ import array_api_compat import array_api_compat.numpy as np -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _describe, _record +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _record if os.environ.get("CUNUMPY_FAKE_CUPY", "").strip().lower() in ("1", "true", "yes"): # tests without a GPU: a strict host stand-in for CuPy, see cunumpy._fake_cupy - from ._fake_cupy import install as _install_fake_cupy + from cunumpy._fake_cupy import install as _install_fake_cupy _install_fake_cupy() From 90ac2302711cf02a28d41e7c2207ac8c968f3f09 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 21:29:50 +0200 Subject: [PATCH 7/9] Add profiling helpers --- .gitlab-ci.yml | 57 +++-- CHANGELOG.md | 9 + README.md | 13 +- docs/source/api.md | 67 +++-- docs/source/guides/data-movement.md | 40 ++- docs/source/guides/execution-helpers.md | 54 ++++ docs/source/kernels/cuda-kernel.md | 17 +- src/cunumpy/LLM_GUIDE.md | 32 ++- src/cunumpy/_cuda_kernel.py | 199 +++++++++++---- src/cunumpy/_fake_cupy_impl.py | 3 + src/cunumpy/_kernel.py | 38 ++- src/cunumpy/_mirror.py | 60 ++++- src/cunumpy/_mpi.py | 18 +- src/cunumpy/_mpi_serial.py | 9 + src/cunumpy/_staging.py | 139 ++++++++--- src/cunumpy/_transfers.py | 69 +++++- src/cunumpy/xp.py | 55 ++++- tests/conftest.py | 23 ++ tests/unit/test_backend_numerics.py | 25 ++ tests/unit/test_benchmarks.py | 81 ------ tests/unit/test_cuda_kernel.py | 25 +- tests/unit/test_cuda_launch_limits.py | 315 ++++++++++++++++++++++++ tests/unit/test_execution_transfers.py | 221 +++++++++++++++++ tests/unit/test_mirror.py | 24 ++ tests/unit/test_porting_helpers.py | 23 +- tests/unit/test_staging.py | 163 ++++++++++++ tests/unit/test_transfers.py | 15 +- 27 files changed, 1513 insertions(+), 281 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/unit/test_backend_numerics.py delete mode 100644 tests/unit/test_benchmarks.py create mode 100644 tests/unit/test_cuda_launch_limits.py create mode 100644 tests/unit/test_execution_transfers.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 17a6fcb..eb5f82f 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -1,12 +1,13 @@ variables: - # Use the specialized MPCDF HPC image CUDA_IMAGE: "gitlab-registry.mpcdf.mpg.de/mpcdf/ci-module-image/nvhpcsdk_26:2026" PIP_DISABLE_PIP_VERSION_CHECK: "1" + ARRAY_BACKEND: "cupy" + CUNUMPY_REQUIRE_CUDA: "1" stages: - test -gpu_tests: +.gpu_job: stage: test image: ${CUDA_IMAGE} tags: @@ -16,32 +17,34 @@ gpu_tests: - module load python-waterboa/2025.06 - module load nvhpcsdk/26 - module load fftw-serial/3.3.10 - script: - - echo "--- CUDA Sanity Check ---" - nvidia-smi - - - echo "--- Detect Compute Capability ---" - - nvidia-smi --query-gpu=compute_cap --format=csv,noheader - - - echo "--- Tool Versions ---" - - cmake --version - - python3 --version - - git --version - - - echo "--- Pytest Execution ---" - # The MPCDF image likely has a specific python environment. - # We install our dependencies into the user directory or a virtualenv. - # One CUDA version only: nvhpcsdk/26 provides CUDA 13.2 (its headers are used - # when CuPy compiles kernels with NVRTC), so install CuPy for CUDA 13 with the - # CUDA 13.2 libraries and NVRTC. Mixing CUDA 12 (cupy-cuda12x) with the CUDA 13.2 - # headers fails to compile CuPy's own kernels (e.g. CUB reductions). + # Keep NVRTC, CUDA libraries and headers on the same toolkit generation. - python3 -m pip install --user "cupy-cuda13x[ctk]" "cuda-toolkit==13.2.*" - - python3 -m pip install --user -e . - - # Add the user bin to PATH for pytest + - python3 -m pip install --user -e ".[test]" - export PATH="$HOME/.local/bin:$PATH" - - - export ARRAY_BACKEND=cupy - python3 -c "import cupy; cupy.show_config()" - - - pytest -xvs . + +gpu_tests: + extends: .gpu_job + script: + - python3 -m pytest -q tests + +gpu_sanitizers: + extends: .gpu_job + timeout: 45 minutes + parallel: + matrix: + - SANITIZER: [memcheck, racecheck, synccheck] + script: + - compute-sanitizer --version + # These tests intentionally contain no illegal-access demonstration kernels. + - >- + compute-sanitizer --tool "$SANITIZER" --error-exitcode 1 + --target-processes all python3 -m pytest -q + tests/unit/test_cuda_collectives.py + tests/unit/test_cuda_launch_limits.py + tests/unit/test_reusable_streams.py + tests/unit/test_staging.py + tests/unit/test_mirror.py + tests/unit/test_mpi_staging.py + tests/unit/test_segment_plans.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 96e35c7..5a2e8b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- Per-device eager CUDA compilation, compiler log streams and explicit `recompile()`. +- Launch validation against actual device/kernel dimensions, thread limits and + static plus dynamic shared memory, with cached per-device opt-in state. +- Transfer payload byte counts, mirror-refresh and kernel-output accounting, and + separate device-only conversion observations. +- Producer `stream=`/`event=` dependencies for output staging and mirror downloads; + device-bound retained storage and CPU-to-pinned staging upgrades. +- GPU CI requires real hardware and runs focused Compute Sanitizer memory, + shared-memory race and synchronization checks. - Reusable `cuda.create_stream`, `create_event`, `record_event`, and `wait_event`, with synchronous host equivalents; `cuda.stream(existing)` selects an existing stream. - `mpi.MPIStaging` reuses host staging storage and rejects overlapping device uses. diff --git a/README.md b/README.md index b2dfe0a..370e333 100644 --- a/README.md +++ b/README.md @@ -149,17 +149,20 @@ with xp.use_backend("cupy"): To verify that a block, such as a time step, makes no transfer at all, count them: `count_transfers()` records every `to_numpy()`, `to_cupy()` and -`to_cunumpy()` call that actually copies, every `PyccelKernel` call that -converts device arrays, and every `Kernel` fallback to the host kernel, with -the call site of each. `assert_no_transfers()` raises with that report if -anything was counted. Only transfers made through CuNumpy are seen; raw +`to_cunumpy()` call that actually copies, mirror/staging refreshes, argument +conversion and kernel output copy-back, with call sites and payload byte counts. +Host kernel conversions and fallbacks have separate explanatory markers. +`assert_no_transfers()` rejects host/device movement and permits device-only +conversions. Only CuNumpy execution/conversion helpers are counted; forwarded +backend calls such as `xp.asarray()` and raw `cupy.ndarray.get()` or `cupy.asarray()` calls need a profiler such as `nsys`. ```python with xp.profiling.count_transfers() as counter: propagator(dt) -assert counter.total == 0, counter.report() +assert counter.to_host == counter.to_device == 0, counter.report() +print(counter.bytes_to_host, counter.bytes_to_device) ``` ## Random numbers and dtypes diff --git a/docs/source/api.md b/docs/source/api.md index 46a777f..6fd171f 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -283,30 +283,34 @@ with xp.profiling.count_transfers() as counter: assert counter.total == 0, counter.report() ``` -Four kinds of events are recorded: +Five kinds of events are recorded: -* `to_host`: `to_numpy()` (or `to_cunumpy()`) called with a CuPy array; -* `to_device`: `to_cupy()` (or `to_cunumpy()`) called with anything that is not - a CuPy array already; +* `to_host`: an actual device-to-host copy through conversion, mirror, staging, + serial MPI, or host kernel helpers; +* `to_device`: an actual host-to-device copy through those helpers, including + `as_device_array()` and `kernel_output()` copy-back; * `kernel_conversion`: a `PyccelKernel` call that copied device arrays to the host (and back), one event per call, naming the kernel and the number of arrays converted; * `fallback`: a `Kernel` without CUDA kernel calling its host kernel on the CuPy backend (`missing_cuda="fallback"`), one event per call, naming the - kernel. The host copies it makes are counted as one `kernel_conversion` - event in addition. + kernel. Physical copies and a `kernel_conversion` marker are recorded separately; +* `device_copy`: a device-only dtype/layout conversion through CuNumpy helpers. Only real transfers count: `to_numpy()` of a NumPy array or `to_cupy()` of a CuPy array records nothing. The counter has the attributes `to_host`, `to_device`, `kernel_conversions`, `fallbacks` (counts per kind), `total`, -`events` (a list of `TransferEvent(kind, description, where)`, where `where` -is the `file:line` of the caller outside CuNumpy) and -`kernel_conversion_calls` (the `kernel_conversion` events). `report()` returns +`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. +`kernel_conversion_calls` selects the conversion markers. `report()` returns a multi-line string with the events grouped by kind and call site, with counts: ```text -4 transfer(s) through cunumpy (3 to_host, 1 to_device, 0 kernel_conversion, 0 fallback) +4 transfer(s) through cunumpy (3 to_host, 1 to_device, 0 kernel_conversion, 0 fallback, 0 device_copy) to_host (3): /home/me/sim/diagnostics.py:42: to_numpy(shape=(100000,), dtype=float64) (x3) to_device (1): @@ -320,13 +324,15 @@ not thread-safe. **Limitation:** only transfers made through CuNumpy are seen. Raw `cupy.ndarray.get()`, `cupy.asarray(numpy_array)`, `numpy.asarray(cupy_array)`, -`float(device_array)`, and implicit conversions inside other libraries are +`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. ### `profiling.assert_no_transfers()` Context manager that raises `AssertionError` with the counter's `report()` if -the block makes a transfer through CuNumpy. It yields the `TransferCounter` +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: ```python @@ -871,7 +877,12 @@ Wraps the `__global__` function `name` in the CUDA C `source` (declared through CuPy on the first call (or by `compile()`), and cached, also on disk by CuPy. CuPy is imported only then, so kernels can be created and their signatures parsed without CuPy; `compile()` raises `RuntimeError` without a -GPU. +GPU. `compile(log_stream=None)` compiles eagerly and returns the CuPy raw kernel; +repeated calls reuse successful state on the current device. `is_compiled` +reports that device's state. `recompile(log_stream=None)` refreshes headers and +options on the current device; finish in-flight launches before rebuilding. +Compiler failures remain retryable and a writable `log_stream` receives compiler +output. Catalog compilation therefore reports errors during setup. `from_file` reads the source from a file; the kernel name defaults to the file name without `suffix` (`axpy_cuda.cu` -> `axpy`), and the directory of the file @@ -972,10 +983,10 @@ shape is given either by `n_threads` or by `grid`: * `grid`: number of blocks per dimension, instead of `n_threads`. * `block`: block shape for this call, instead of `block_size`. * `shared_mem`: dynamic shared memory per block in bytes, for - `extern __shared__` arrays. Above 48 KiB (the limit every device has) the - compiled kernel's `max_dynamic_shared_size_bytes` is raised to `shared_mem` - once, up to the device's opt-in limit (`xp.cuda.max_shared_memory_per_block( - opt_in=True)`); a larger request raises `ValueError` before the launch. + `extern __shared__` arrays. Static plus dynamic storage is checked against + the device's opt-in limit. Dynamic storage above the default allowance after + subtracting static storage opts in through `max_dynamic_shared_size_bytes`. + Block/grid dimensions and device/kernel thread limits are also checked. Without `n_threads` and `grid`, a launch uses `n_threads_from(args)` if the kernel has one (constructor argument and settable property): a function of the @@ -1175,7 +1186,7 @@ too); `xp.cuda.get_cuda_debug()` returns the current setting. A kernel created w enabling it also affects kernels created earlier; `debug=True` or `debug=False` fix the mode for one kernel. Only the compile options are fixed at compile time: a kernel compiled before debug mode was enabled keeps its -options, so call `compile()` after enabling, or create the kernels after +options, so call `recompile()` after enabling, or create the kernels after enabling. `kernel.debug_active()` tells whether debug mode applies to a kernel now, and `kernel.compile_options()` returns the options a compilation now would use. @@ -1972,6 +1983,16 @@ buffer is reused `buffers` copies later; a stale result raises the NumPy backend copy at once. Device copies are counted by `count_transfers()`. The arrays must have the staging shape and dtype. +`copy(array, *, stream=None, event=None)` accepts an explicit producer stream or +event (at most one). An event queues a wait before snapshotting on the current +stream. Later writes on another stream must wait for the snapshot/copy; +`copy.result()` is a conservative completion point. Keep the source's device +current when submitting copies. Storage binds to that device and rejects another +device. `ready()`/`result()` temporarily select the owning device and restore the +caller's device. The first GPU use of CPU-initialized storage allocates fresh +pinned slots; previous completed CPU handles keep their snapshots. `shape` and +`dtype` are read-only. + ## `memory.DeviceMirror` ```python @@ -1989,7 +2010,7 @@ raises `TypeError`. transfer). On the NumPy backend it is the host array itself, so the same code runs without any copy on the CPU. * `to_device()`: copies the host array into the existing device array; - `to_host()`: copies the device array into the host array, in place, so the + `to_host(stream=None, event=None)`: copies the device array into the host array, in place, so the host array keeps its identity and the owning library sees the new values. Both are no-ops on the NumPy backend. * `zero()`: zeroes the device array (allocating it empty if needed), or the @@ -2001,6 +2022,14 @@ raises `TypeError`. * `host`, `shape`, `dtype` properties. `to_device()`, `to_host()`, `zero()` and `rebind()` return the mirror, for chaining. +Make the mirror's device current before using its device storage. A mismatched +device raises before copying or zeroing. Pass an explicit producer `stream=` or +`event=` to `to_host()` when production happened outside the current stream; +the host copy is complete on return. Initial uploads and refreshes in both +directions appear in transfer accounting, with payload byte counts. Switching +to NumPy uses the host array; explicitly refresh with `to_device()` after CPU +changes before resuming GPU work. + The transfers are explicit so that one per accumulation is visible and bounded: diff --git a/docs/source/guides/data-movement.md b/docs/source/guides/data-movement.md index ad598e0..07f3f78 100644 --- a/docs/source/guides/data-movement.md +++ b/docs/source/guides/data-movement.md @@ -59,8 +59,8 @@ Move the conversions out of the loop and convert inside the `if` only. ## Count the transfers The classic performance bug of a GPU port is a transfer that sneaks into the -time loop. `count_transfers()` records every copy made through CuNumpy in a -block, with the file and line that caused it: +time loop. `count_transfers()` records copies made through CuNumpy's conversion +and execution helpers, with the file and line that caused them: ```python with xp.profiling.count_transfers() as counter: @@ -72,23 +72,29 @@ print(counter.report()) ``` ```text -20 transfer(s) through cunumpy (10 to_host, 10 to_device, 0 kernel_conversion, 0 fallback) +20 transfer(s) through cunumpy (10 to_host, 10 to_device, 0 kernel_conversion, 0 fallback, 0 device_copy) to_host (10): /home/me/sim/diagnostics.py:42: to_numpy(shape=(100000,), dtype=float64) (x10) to_device (10): /home/me/sim/step.py:17: to_cupy(shape=(100000,), dtype=float64) (x10) ``` -Four kinds of events are recorded: `to_host`, `to_device`, -`kernel_conversion` (a [`PyccelKernel`](../kernels/pyccel-kernel.md) that copied -device arrays to the host and back) and `fallback` (a -[`Kernel`](../kernels/dispatch.md) without CUDA version running its host kernel -on the GPU backend). Calls that do not copy, such as `to_numpy()` of a NumPy -array, are not counted, so a `count_transfers()` block on the NumPy backend -reports zero. +Physical copies are recorded as `to_host`, `to_device`, or `device_copy` +(device-only dtype/layout conversions). Mirror refreshes, staging, and kernel +output copies are included. Each copy has an `event.nbytes` payload size; +`counter.bytes_to_host`, `counter.bytes_to_device`, and `counter.bytes(kind)` +sum the known sizes. -In tests, `assert_no_transfers()` turns this into a check that fails with the -report: +`kernel_conversion` marks a [`PyccelKernel`](../kernels/pyccel-kernel.md) call +that copied device arrays to the host and back, and `fallback` marks a +[`Kernel`](../kernels/dispatch.md) without a CUDA version running its host kernel +on the GPU backend. Physical copies are recorded separately from these markers, +which add no bytes. `counter.total` includes markers; `counter.to_host + +counter.to_device` counts physical host/device copies. Reference-only calls, +such as `to_numpy()` of a NumPy array, record nothing. + +In tests, `assert_no_transfers()` rejects host/device copies and host execution +on the GPU backend, with the report below. It allows device-only conversions: ```python def test_step_stays_on_device(): @@ -97,7 +103,8 @@ def test_step_stays_on_device(): step(state, dt) ``` -The counter only sees transfers made through CuNumpy. Raw `cupy.asarray(host)`, +The counter only sees transfers made through CuNumpy helpers. Forwarded backend +operations such as `xp.asarray(host)`, raw `cupy.asarray(host)`, `device_array.get()`, `float(device_scalar)` and conversions inside other libraries are invisible to it; use `nsys` to find those (see [Timing and profiling](profiling.md)). @@ -136,6 +143,13 @@ place. On the NumPy backend `mirror.device` *is* the host array and the copies are no-ops. See [Accumulation kernels](../kernels/accumulation.md) for the full pattern. +Both directions appear in transfer accounting. A retained device buffer is +bound to its CUDA device; make that device current before using it. Pass +`mirror.to_host(stream=producer)` or `mirror.to_host(event=completion)` when +the kernel ran on another stream. The refresh blocks until the host data is +ready. After CPU work changes the host buffer, call `mirror.to_device()` before +resuming GPU work. + ## Pinned memory Host-to-device copies from page-locked ("pinned") host memory are faster and diff --git a/docs/source/guides/execution-helpers.md b/docs/source/guides/execution-helpers.md index 819ca22..fc67400 100644 --- a/docs/source/guides/execution-helpers.md +++ b/docs/source/guides/execution-helpers.md @@ -92,6 +92,60 @@ Without either, CuNumpy synchronizes all work on the buffer's device. `synchronize_for_mpi(*arrays, stream=..., event=...)` follows the same rule; pass at most one dependency, covering all reads/writes of the supplied buffers. +## Order output snapshots and mirror refreshes + +```python +staging = xp.memory.HostStaging(field.shape, field.dtype, buffers=2) +copy = staging.copy(field, stream=producer) +# Alternatively: copy = staging.copy(field, event=produced) +host_field = copy.result() # complete host data for an output writer +``` + +The snapshot runs on the supplied producer stream, or on the current stream +after waiting for the supplied event. Without either, the current stream must +already cover production. Subsequent writes on that stream are ordered after +the snapshot; another stream must wait before overwriting the source. Waiting +for `copy.result()` is a conservative way to establish completion. + +Make the source device current when submitting copies. The staging instance +binds to that device and rejects another. Completion checks temporarily select +the owning device and restore the caller's device. CPU-initialized storage gets +fresh pinned slots at first GPU use; earlier CPU snapshots stay valid. Reusing +a slot still invalidates its earlier handle, so copy a returned host buffer if +the writer must retain it. + +`DeviceMirror.to_host(stream=producer)` and `.to_host(event=produced)` provide the +same explicit producer choices for synchronous refreshes into another library's +host array. Make the mirror's device current. After CPU work changes its host +array, explicitly call `.to_device()` before resuming GPU work. + +## Check copies and prepare kernels + +```python +with xp.profiling.count_transfers() as copies: + mirror.to_device() + mirror.to_host() +print(copies.bytes_to_device, copies.bytes_to_host) +``` + +Transfers through execution/conversion helpers carry payload byte counts, +including mirror refreshes, MPI/output staging, argument conversion, and kernel +output copy-back. Conversion/fallback markers describe calls and add no bytes. +`total` includes these markers; use `to_host + to_device` for physical boundary +copy counts. Device-only conversions are recorded separately as `device_copy` +and are permitted by `assert_no_transfers()`. Forwarded backend operations such +as `xp.asarray()` and external library copies are outside this accounting. + +`kernel.compile(log_stream=log)` compiles eagerly and caches successful state +per device; catalog setup catches CUDA compiler failures before timesteps. +`kernel.recompile()` explicitly refreshes headers/debug options. Each launch +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 +`memcheck`, `racecheck`, and `synccheck` jobs with nonzero sanitizer error exits. +Numerical parity tests are separate from performance comparisons. + ## Prepare cell ranges and segment reductions ```python diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index fd93217..df2abe5 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/docs/source/kernels/cuda-kernel.md @@ -98,6 +98,11 @@ kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) * **Stream**: `stream=s` queues the launch on a stream, see [Devices, memory and streams](../guides/gpu-devices.md). +Launches validate block/grid dimensions, threads per block against both device +and compiled-kernel limits, and static plus dynamic shared memory. Dynamic +storage above the default allowance opts in where supported. Queries and opt-in +state are cached per device. Invalid configurations raise before launching. + A zero-sized launch (`n_threads=0`) launches nothing. `launch_shape()` returns the `(grid, block)` a call would use, which is how to size per-block outputs: @@ -248,8 +253,16 @@ The first call of each kernel compiles it, which takes from a fraction of a second to several seconds. Call `kernel.compile()` (or `catalog.compile_all(jobs=8)` for a whole catalog, in parallel threads) during setup so the first time step is not slower than the rest and compilation errors -appear before the simulation starts. `kernel.is_compiled` tells whether it has -happened. Compiling requires a GPU (`RuntimeError` otherwise). +appear before the simulation starts. `kernel.is_compiled` reports successful +compilation on the current device. Failed compilation can be retried. Compiling +requires a GPU (`RuntimeError` otherwise). Parallel catalog compilation preserves +the device selected by its caller. + +`kernel.compile(log_stream=log)` accepts a writable compiler log stream. +`kernel.recompile(log_stream=log)` refreshes headers and debug options on the +current device; other devices keep their compiled versions. Finish in-flight +launches before rebuilding. Scope-profiler can inspect the returned CuPy kernel's +existing `attributes`; CuNumpy does not add a separate resource-reporting API. ## Tips for writing kernels diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 50589f7..a07f73c 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -131,15 +131,19 @@ xp.to_numpy(a) # -> numpy.ndarray; copies only if a is CuPy xp.to_cupy(a) # -> cupy.ndarray; copies only if a is not CuPy; ImportError w/o CuPy xp.to_cunumpy(a) # -> array of the active backend xp.as_device_array(value, dtype=None, ndim=None, *, name=None) - # contiguous CuPy array of dtype: returned as is; else one copy; + # contiguous CuPy array of dtype: returned as is; else convert; # RuntimeError on the NumPy backend 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(): ... # AssertionError with report if anything copied +with xp.profiling.assert_no_transfers(): ... # rejects host/device copies and host fallback ``` -Only transfers through cunumpy are counted (not raw `cupy.asarray`, `.get()`, -`float(device_scalar)`, or `DeviceMirror.to_host()/to_device()`). +Mirror refreshes, MPI/output staging, argument conversion and kernel output +copy-back are counted with `event.nbytes`, `c.bytes_to_host`, and +`c.bytes_to_device`. Conversion/fallback markers add no bytes; `c.total` +includes markers. Device-only conversions are separate `device_copy` events +and are allowed by `assert_no_transfers()`. Forwarded `xp.asarray`, raw CuPy +copies, and external library/scalar conversions are not counted. MPI, accumulation and versions: @@ -264,7 +268,9 @@ k = xp.cuda.CudaKernel(source, name, *, block_size=128, options=(), include_dirs k = xp.cuda.CudaKernel.from_file("push/push_cuda.cu") # name "push" ks = xp.cuda.CudaKernel.all_from_file("ops.cu") # dict name -> kernel k(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) -k.compile(); k.is_compiled; k.launch_shape(n_threads) -> (grid, block) +k.compile(log_stream=None); k.recompile(log_stream=None) +k.is_compiled # successful compilation on the current CUDA device +k.launch_shape(n_threads) -> (grid, block) k.included_headers; k.compile_options(); k.debug_active() xp.cuda.CudaKernelVariants(factory).get(*key); .compile_all(keys, jobs=1) xp.cuda.ctype_of(np.float64) == "double" @@ -375,10 +381,20 @@ m = xp.memory.DeviceMirror(host_numpy_array) # TypeError if not numpy.ndarray m.device # CuPy copy (lazy) on CuPy; the host array itself on NumPy m.zero() m.to_device() -m.to_host() # to_host copies in place; no-ops on NumPy +m.to_host(stream=None, event=None) # copies in place; no-op on NumPy; pass at most one dependency m.rebind(new_host_array) # after the owner reallocates ``` +Make a retained mirror's CUDA device current before using it. `to_host()` +blocks until the copy completes; supply the producer stream or event when work +ran elsewhere. Refresh with `to_device()` after CPU work modifies the host. +`HostStaging.copy(a, stream=None, event=None)` snapshots on the producer stream +(or current stream after an event wait). Later writes must be ordered after +the snapshot. Staging binds to the first source device and upgrades ordinary +CPU slots to pinned slots on first GPU use. `ready()`/`result()` temporarily +select the owning device and restore the caller's device. Copy the returned +host buffer to retain it beyond slot reuse. + Debugging: ```python @@ -508,8 +524,8 @@ def test_scale(): * Comparing CPU and GPU results with exact equality where the GPU uses atomics or a different reduction order; use a tolerance. * Assuming `xp.random.seed(s)` gives the same numbers on both backends. -* Expecting `DeviceMirror` copies or raw CuPy conversions to show up in - `count_transfers()`. +* Expecting raw CuPy or forwarded backend operations to show up in + `count_transfers()`; mirror copies and execution helpers are counted. * Writing code that requires CuPy at import time. `import cunumpy` and creating `CudaKernel` objects work without CuPy; import `cupy` lazily, only on the GPU path, if you need it at all. diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index e022bed..1a37f52 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -50,6 +50,7 @@ class is the one definition of the arguments. import hashlib import inspect import math +import operator import os import re import sys @@ -57,6 +58,7 @@ class is the one definition of the arguments. from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext +from dataclasses import dataclass from pathlib import Path from typing import Any, NamedTuple @@ -84,6 +86,32 @@ class is the one definition of the arguments. # CUDA limit on the number of threads per block _MAX_THREADS_PER_BLOCK = 1024 +# Runtime queries are setup work, never repeated for a prepared launch. +_DEVICE_LIMITS: dict[int, dict[str, Any]] = {} + + +def _current_device() -> int: + import cupy as cp + + return cp.cuda.runtime.getDevice() + + +def _device_limits(device: int) -> dict[str, Any]: + if device not in _DEVICE_LIMITS: + import cupy as cp + + _DEVICE_LIMITS[device] = cp.cuda.runtime.getDeviceProperties(device) + return _DEVICE_LIMITS[device] + + +@dataclass +class _CompiledKernel: + raw: Any + limits: dict[str, Any] + attributes: dict[str, int] + dynamic_shared: int + + # Headers shipped with cunumpy: #include etc. _CUDA_INCLUDE_DIR = Path(__file__).resolve().parent / "cuda" / "include" _ARRAY_VIEW_INCLUDE = '#include "cunumpy/array_view.cuh"' @@ -433,8 +461,22 @@ def _compile_in_threads( compile() return list(compilers) + from cunumpy.xp import cupy_available + + device = _current_device() if cupy_available() else None + + def on_device(compiler: Callable[[], Any]) -> Any: + if device is None: + return compiler() + import cupy as cp + + with cp.cuda.Device(device): + return compiler() + with ThreadPoolExecutor(max_workers=min(jobs, len(compilers))) as pool: - futures = {name: pool.submit(compile) for name, compile in compilers.items()} + futures = { + name: pool.submit(on_device, compile) for name, compile in compilers.items() + } compiled, error = [], None for name, future in futures.items(): exc = future.exception() @@ -1362,7 +1404,7 @@ def verify_layout( RuntimeError If CuPy is not available. """ - from cunumpy.xp import cupy_available + from cunumpy.xp import cupy_available, to_numpy if not cupy_available(): raise RuntimeError("verify_layout() compiles a CUDA kernel and needs CuPy") @@ -1378,7 +1420,7 @@ def verify_layout( ) out = cp.zeros(len(self._fields) + 2, dtype=cp.uint64) kernel(out, n_threads=1) - measured = [int(v) for v in out.get()] + measured = [int(v) for v in to_numpy(out)] layout = {"sizeof": measured[0], "alignof": measured[1]} layout.update({f.name: n for f, n in zip(self._fields, measured[2:])}) @@ -1931,8 +1973,6 @@ def __init__( ) -> None: self._block = self._check_block(_as_shape(block_size, "block_size")) self.n_threads_from = n_threads_from - # dynamic shared memory the compiled kernel is set up for (see __call__) - self._shared_mem_opt_in = 48 * 1024 self._check_finite = bool(check_finite) self._debug = None if debug is None else bool(debug) self._source = source @@ -1958,7 +1998,7 @@ def __init__( if self._signature is None else [_checker(p, i) for i, p in enumerate(self._signature)] ) - self._raw_kernel = None + self._compiled: dict[int, _CompiledKernel] = {} @classmethod def from_file( @@ -2213,51 +2253,80 @@ def signature(self) -> tuple[CudaParameter, ...] | None: @property def is_compiled(self) -> bool: - """Whether the kernel has been compiled in this process.""" - return self._raw_kernel is not None + """Whether compilation succeeded on the current CUDA device.""" + return bool(self._compiled) and _current_device() in self._compiled - def compile(self) -> Any: + def compile(self, *, log_stream: Any = None) -> Any: """Compile the kernel now (it is otherwise compiled on the first call). The options are :meth:`compile_options`, evaluated now: a kernel - compiled before debug mode was enabled keeps its options. + compiled before debug mode was enabled keeps its options. Call + :meth:`recompile` to refresh headers/options explicitly. Compilation + and launch resources are cached separately on each CUDA device; + failed compilation remains retryable. ``log_stream`` receives compiler + output (a writable file object, or None). Returns ------- cupy.RawKernel The compiled kernel; compiled once and cached (also on disk by CuPy, - keyed on the source and :meth:`compile_options`). + keyed on the source and :meth:`compile_options`). No launch is needed + to surface compiler errors. Raises ------ RuntimeError If CuPy or a GPU is not available. """ - if self._raw_kernel is None: - from cunumpy.xp import cupy_available + from cunumpy.xp import cupy_available - if not cupy_available(): - raise RuntimeError( - f"cannot compile CUDA kernel {self.expression!r}: " - "CuPy is not installed or no GPU is available", - ) - import cupy as cp + if not cupy_available(): + raise RuntimeError( + f"cannot compile CUDA kernel {self.expression!r}: " + "CuPy is not installed or no GPU is available", + ) + import cupy as cp + device = _current_device() + if device not in self._compiled: options = self.compile_options() if self._template_args is None: - self._raw_kernel = cp.RawKernel( + raw = cp.RawKernel( self._source, self._name, options=options, ) + raw.compile(log_stream=log_stream) else: module = cp.RawModule( code=self._source, options=options, name_expressions=[self.expression], ) - self._raw_kernel = module.get_function(self.expression) - return self._raw_kernel + module.compile(log_stream=log_stream) + raw = module.get_function(self.expression) + limits = _device_limits(device) + attributes = dict(raw.attributes) + static = max(0, attributes.get("shared_size_bytes", 0)) + dynamic = attributes.get("max_dynamic_shared_size_bytes", -1) + if dynamic < 0: + dynamic = max(0, int(limits["sharedMemPerBlock"]) - static) + dynamic = min(dynamic, max(0, int(limits["sharedMemPerBlock"]) - static)) + # Do not publish failed compilation or incomplete setup as compiled. + self._compiled[device] = _CompiledKernel(raw, limits, attributes, dynamic) + return self._compiled[device].raw + + def recompile(self, *, log_stream: Any = None) -> Any: + """Recompile on the current device using current headers/debug options. + + Other devices keep their compiled kernels. A failed rebuild remains + uncompiled and can be retried; in-flight launches must finish first. + """ + from cunumpy.xp import cupy_available + + if cupy_available(): + self._compiled.pop(_current_device(), None) + return self.compile(log_stream=log_stream) def prepare_args(self, *args: Any) -> tuple[Any, ...]: """The arguments as passed to ``cupy.RawKernel``: flattened and checked. @@ -2331,7 +2400,7 @@ def launch_shape( "numbers of dimensions", ) block_shape = block_shape + (1,) * (len(threads) - 1) - grid_shape = tuple(math.ceil(n / b) for n, b in zip(threads, block_shape)) + grid_shape = tuple((n + b - 1) // b for n, b in zip(threads, block_shape)) return grid_shape, block_shape def __call__( @@ -2374,10 +2443,14 @@ def __call__( illegal memory access; the CuPy error is chained. Without debug mode, such an error surfaces at a later synchronization (a ``.get()``, an MPI call, ...), not necessarily in this kernel. + ValueError + Dimensions, threads or static plus dynamic shared memory exceed + this device/kernel's limits, or arrays/stream belong to another device. """ if n_threads is None and grid is None and self._n_threads_from is not None: n_threads = self._n_threads_from(args) grid_shape, block_shape = self.launch_shape(n_threads, grid=grid, block=block) + shared_mem = operator.index(shared_mem) if shared_mem < 0: raise ValueError(f"shared_mem must be non-negative, got {shared_mem}") values = self.prepare_args(*args) @@ -2385,8 +2458,23 @@ def __call__( return kernel = self.compile() - if shared_mem > self._shared_mem_opt_in: - self._opt_in_shared_memory(kernel, shared_mem) + device = _current_device() + state = self._compiled[device] + self._validate_launch(state, device, grid_shape, block_shape, shared_mem) + if stream is not None and getattr(stream, "device_id", device) not in ( + device, + -1, + ): + raise ValueError( + f"kernel {self.expression!r}: stream belongs to another device" + ) + for label, array in _device_arrays_in(args): + array_device = getattr(getattr(array, "device", None), "id", None) + if array_device is not None and array_device != device: + raise ValueError( + f"kernel {self.expression!r} on CUDA device {device}: " + f"{label} belongs to CUDA device {array_device}", + ) debug = self.debug_active() with stream if stream is not None else nullcontext(): kernel(grid_shape, block_shape, values, shared_mem=shared_mem) @@ -2395,29 +2483,50 @@ def __call__( if self._check_finite: self._check_finite_arrays(args) - def _opt_in_shared_memory(self, kernel: Any, shared_mem: int) -> None: - """Allow `shared_mem` bytes of dynamic shared memory above the default limit. - - Up to :data:`~cunumpy.cuda.DEFAULT_SHARED_MEMORY_PER_BLOCK` (48 KiB) every - device launches without setup. Above it, newer GPUs need the kernel - attribute ``max_dynamic_shared_size_bytes``; it is set once (and again - for a larger request) up to the device's opt-in limit. - """ - from cunumpy._device import ( - DEFAULT_SHARED_MEMORY_PER_BLOCK, - max_shared_memory_per_block, + def _validate_launch( + self, + state: _CompiledKernel, + device: int, + grid: tuple[int, ...], + block: tuple[int, ...], + shared_mem: int, + ) -> None: + limits, attributes = state.limits, state.attributes + name = limits.get("name", "unknown") + if isinstance(name, bytes): + name = name.decode(errors="replace") + context = f"kernel {self.expression!r} on CUDA device {device} ({name})" + for label, shape, maximum in ( + ("block", block, limits["maxThreadsDim"]), + ("grid", grid, limits["maxGridSize"]), + ): + for axis, (requested, limit) in enumerate(zip(shape, maximum)): + if requested > limit: + raise ValueError( + f"{context}: {label}[{axis}]={requested} exceeds the {limit} limit", + ) + threads_limit = int(limits["maxThreadsPerBlock"]) + kernel_limit = attributes.get("max_threads_per_block", -1) + if kernel_limit > 0: + threads_limit = min(threads_limit, kernel_limit) + if math.prod(block) > threads_limit: + raise ValueError( + f"{context}: block {block} has {math.prod(block)} threads, " + f"exceeding the {threads_limit} device/kernel limit", + ) + static = max(0, attributes.get("shared_size_bytes", 0)) + maximum_shared = max( + int(limits["sharedMemPerBlock"]), + int(limits.get("sharedMemPerBlockOptin", 0)), ) - - if shared_mem <= DEFAULT_SHARED_MEMORY_PER_BLOCK: - return - limit = max_shared_memory_per_block(opt_in=True) - if shared_mem > limit: + if shared_mem + static > maximum_shared: raise ValueError( - f"kernel {self.expression!r}: shared_mem={shared_mem} bytes exceeds " - f"the {limit} bytes a block may use on this device", + f"{context}: static shared memory {static} + shared_mem={shared_mem} " + f"exceeds the {maximum_shared} bytes a block may use on this device", ) - kernel.max_dynamic_shared_size_bytes = shared_mem - self._shared_mem_opt_in = shared_mem + if shared_mem > state.dynamic_shared: + state.raw.max_dynamic_shared_size_bytes = shared_mem + state.dynamic_shared = shared_mem def _check_finite_arrays(self, args: tuple[Any, ...]) -> None: """Raise if a floating-point array among `args` holds NaN or inf.""" diff --git a/src/cunumpy/_fake_cupy_impl.py b/src/cunumpy/_fake_cupy_impl.py index 5a043d0..0a4a13a 100644 --- a/src/cunumpy/_fake_cupy_impl.py +++ b/src/cunumpy/_fake_cupy_impl.py @@ -401,6 +401,9 @@ def __init__(self, *args, **kwargs): def __call__(self, *args, **kwargs): raise NotImplementedError("fake CuPy cannot run CUDA kernels") + def compile(self, log_stream=None): + raise NotImplementedError("fake CuPy cannot compile CUDA kernels") + def get_function(self, name): return _NoGPU() diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index 1c74c1a..a7fe6a7 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -35,7 +35,7 @@ import numpy as np from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _record +from cunumpy._transfers import _describe, _nbytes, _record from cunumpy.xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy __all__ = ["CompiledHostKernel", "KernelArguments", "PyccelKernel", "resolve_host_args"] @@ -278,6 +278,12 @@ def _convert_to_numpy( if _is_device_array(value): value_np = _device_to_host(value) + if _COUNTERS: + _record( + "to_host", + f"PyccelKernel {self.name!r} input ({_describe(value)})", + nbytes=_nbytes(value_np), + ) memo[key] = value_np converted.append((value, value_np)) return value_np @@ -323,13 +329,23 @@ def _convert_to_numpy( def _convert_from_numpy(self, value: Any) -> Any: """Move host arrays returned by the kernel back to the device.""" if self._is_array(value): - return _host_to_device(value) + return self._copy_to_device(value) if isinstance(value, tuple): return tuple(self._convert_from_numpy(item) for item in value) if isinstance(value, list): return [self._convert_from_numpy(item) for item in value] return value + def _copy_to_device(self, value: Any) -> Any: + result = _host_to_device(value) + if _COUNTERS: + _record( + "to_device", + f"PyccelKernel {self.name!r} output ({_describe(value)})", + nbytes=_nbytes(value), + ) + return result + def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> None: """Record the id of every host array reachable from `value`. @@ -479,7 +495,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: # Copy in-place kernel updates back to the device arrays. for device_array, host_array in converted: if writeable is None or id(host_array) in writeable: - device_array[...] = _host_to_device(host_array) + device_array[...] = self._copy_to_device(host_array) return self._convert_from_numpy(result) @@ -819,7 +835,14 @@ def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: if not is_gpu(value): value = to_cupy(value) - return cupy.ascontiguousarray(value, dtype=dtype) + result = cupy.ascontiguousarray(value, dtype=dtype) + if _COUNTERS and result is not value: + _record( + "device_copy", + f"as_kernel_array({_describe(value)})", + nbytes=_nbytes(result), + ) + return result return np.ascontiguousarray(to_numpy(value), dtype=dtype) @@ -838,4 +861,9 @@ def kernel_output(out: Any, like: Any, dtype: Any = None) -> Generator[Any]: buffer = as_kernel_array(out, like, dtype) yield buffer if buffer is not out: - out[...] = buffer if is_gpu(out) or not is_gpu(buffer) else to_numpy(buffer) + 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 diff --git a/src/cunumpy/_mirror.py b/src/cunumpy/_mirror.py index 147735c..fbcd650 100644 --- a/src/cunumpy/_mirror.py +++ b/src/cunumpy/_mirror.py @@ -18,6 +18,8 @@ import numpy as np +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _nbytes, _record from cunumpy.xp import _cupy_backend, to_cupy __all__ = ["DeviceMirror"] @@ -104,13 +106,24 @@ def device(self) -> Any: """ if not _cupy_backend(): return self._host + self._check_bound() + self._check_device() if self._device is None: - self._check_bound() # The only host -> device allocation; it goes through to_cupy so # that it is visible wherever transfers are accounted for. self._device = to_cupy(self._host) return self._device + def _check_device(self) -> None: + if self._device is not None: + import cupy as cp + + if cp.cuda.runtime.getDevice() != self._device.device.id: + raise ValueError( + f"DeviceMirror is bound to CUDA device {self._device.device.id}; " + "make that device current before using the mirror", + ) + def _check_bound(self) -> None: """Raise if the host array no longer has the shape/dtype it was bound with.""" if self._host.shape != self._shape or self._host.dtype != self._dtype: @@ -156,37 +169,78 @@ def to_device(self) -> DeviceMirror: self._check_bound() if not _cupy_backend(): return self + self._check_device() if self._device is None: _ = self.device # allocates as a copy of the host, through to_cupy else: # host -> device copy into the existing device buffer self._device.set(self._host) + if _COUNTERS: + _record( + "to_device", + f"DeviceMirror.to_device({_describe(self._host)})", + nbytes=_nbytes(self._host), + ) return self - def to_host(self) -> DeviceMirror: + def to_host(self, *, stream: Any = None, event: Any = None) -> DeviceMirror: """Copy the device array into the host array, in place. The host array keeps its identity, so a library that holds the buffer sees the new values. No-op on the NumPy backend, and if no device array has been created yet. + Make the mirror's device current. Pass ``stream=`` to copy on a known + producer stream, or ``event=`` to wait for production on the current + stream. With neither, only the current stream's work is ordered. + The method blocks until the host copy is complete. + Raises ------ ValueError If the host array was reallocated with another shape or dtype. """ + if stream is not None and event is not None: + raise ValueError("pass only one of stream and event") self._check_bound() if not _cupy_backend() or self._device is None: return self + self._check_device() + from cunumpy._streams import HostEvent, HostStream + + if isinstance(event, HostEvent) or isinstance(stream, HostStream): + raise TypeError("device copies require a CUDA producer stream or event") + import cupy as cp + + if stream is not None and getattr( + stream, "device_id", self._device.device.id + ) not in ( + self._device.device.id, + -1, + ): + raise ValueError("DeviceMirror producer stream belongs to another device") + if event is not None: + cp.cuda.get_current_stream().wait_event(event) # device -> host copy into the existing host buffer - self._device.get(out=self._host) + if stream is None: + self._device.get(out=self._host) + else: + self._device.get(out=self._host, stream=stream) + if _COUNTERS: + _record( + "to_host", + f"DeviceMirror.to_host({_describe(self._host)})", + nbytes=_nbytes(self._host), + ) return self def zero(self) -> DeviceMirror: """Zero the array kernels write into (the device array, or the host on NumPy).""" + self._check_bound() if not _cupy_backend(): self._host.fill(0) return self + self._check_device() if self._device is None: self._check_bound() import cupy as cp diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 9f9541f..1ed6387 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -15,7 +15,7 @@ _LOCAL_RANK_VARIABLES, # noqa: F401 - re-exported ) from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _describe, _record +from cunumpy._transfers import _describe, _nbytes, _record from cunumpy.xp import array_backend, cupy_available, to_numpy _logger = logging.getLogger(__name__) @@ -243,17 +243,25 @@ def mpi_buffer( synchronize_for_mpi(array, stream=stream, event=event) transfer_array = staging._transfer_array(array) if send: - if _COUNTERS: - _record("to_host", f"mpi_buffer({_describe(array)}) staging for send") if transfer_array is not array: transfer_array[...] = array transfer_array.get(out=host) + if _COUNTERS: + _record( + "to_host", + f"mpi_buffer({_describe(array)}) staging for send", + nbytes=_nbytes(array), + ) yield host if recv: + transfer_array.set(host) if _COUNTERS: - _record("to_device", f"mpi_buffer({_describe(array)}) staging for recv") + _record( + "to_device", + f"mpi_buffer({_describe(array)}) staging for recv", + nbytes=_nbytes(array), + ) # Ensure MPI's host buffer can be reused immediately on context exit. - transfer_array.set(host) if transfer_array is not array: array[...] = transfer_array cp.cuda.get_current_stream().synchronize() diff --git a/src/cunumpy/_mpi_serial.py b/src/cunumpy/_mpi_serial.py index 5cf29a2..855f21f 100644 --- a/src/cunumpy/_mpi_serial.py +++ b/src/cunumpy/_mpi_serial.py @@ -236,7 +236,11 @@ def _copy(source: Any, target: Any, offset: int = 0) -> None: if source is _IN_PLACE or source is None or target is None: return if hasattr(source, "get") and not hasattr(target, "get"): + from cunumpy._transfers import _ACTIVE, _nbytes, _record + source = source.get() + if _ACTIVE: + _record("to_host", "SerialComm receive from device", nbytes=_nbytes(source)) if not target.flags.c_contiguous: raise ValueError("the receive buffer must be C-contiguous") flat = target.reshape(-1) @@ -247,6 +251,11 @@ def _copy(source: Any, target: Any, offset: int = 0) -> None: f"at offset {offset}", ) flat[offset : offset + source.size] = source + if not hasattr(source, "get") and hasattr(target, "get"): + from cunumpy._transfers import _ACTIVE, _nbytes, _record + + if _ACTIVE: + _record("to_device", "SerialComm receive from host", nbytes=_nbytes(source)) def _check_rank(rank: int, what: str) -> None: diff --git a/src/cunumpy/_staging.py b/src/cunumpy/_staging.py index b465c9b..1fd8b7d 100644 --- a/src/cunumpy/_staging.py +++ b/src/cunumpy/_staging.py @@ -29,13 +29,15 @@ from __future__ import annotations +import operator +from contextlib import nullcontext from typing import Any import array_api_compat import numpy as np from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _record +from cunumpy._transfers import _nbytes, _record __all__ = ["HostStaging", "StagedCopy"] @@ -56,11 +58,19 @@ def _empty_pinned(shape: tuple[int, ...], dtype: Any) -> np.ndarray: class _Slot: - def __init__(self, host: np.ndarray) -> None: + def __init__(self, host: np.ndarray, device_id: int | None = None) -> None: self.host = host self.device: Any = None # the snapshot, allocated on the first device copy self.event: Any = None # recorded after the copy to the host self.generation = 0 + self.device_id = device_id + + def context(self) -> Any: + return ( + nullcontext() + if self.device_id is None + else _cupy().cuda.Device(self.device_id) + ) class StagedCopy: @@ -81,7 +91,10 @@ def _check(self) -> None: def ready(self) -> bool: """Whether the copy has finished (never waits).""" self._check() - return self._slot.event is None or bool(self._slot.event.done) + if self._slot.event is None: + return True + with self._slot.context(): + return bool(self._slot.event.done) def result(self) -> np.ndarray: """Wait for the copy and return the host array. @@ -95,7 +108,8 @@ def result(self) -> np.ndarray: """ self._check() if self._slot.event is not None: - self._slot.event.synchronize() + with self._slot.context(): + self._slot.event.synchronize() return self._slot.host @@ -112,6 +126,12 @@ class HostStaging: Number of buffers (and device snapshots) used in turn: how many copies may be in flight at once. 2 (double buffering) lets one copy run while the previous result is written out. + + Notes + ----- + GPU storage is bound to the first source device; make it current when + submitting copies. A CPU-initialized instance allocates new pinned slots + on its first GPU copy; earlier CPU results retain their completed snapshots. """ def __init__( @@ -120,14 +140,24 @@ def __init__( dtype: Any, buffers: int = 2, ) -> None: + buffers = operator.index(buffers) if buffers < 1: raise ValueError(f"buffers must be at least 1, got {buffers}") - self.shape = (shape,) if isinstance(shape, int) else tuple(shape) - self.dtype = np.dtype(dtype) + self._shape = ( + (operator.index(shape),) + if isinstance(shape, (int, np.integer)) + else tuple(operator.index(n) for n in shape) + ) + if any(n < 0 for n in self.shape): + raise ValueError("staging shape must be non-negative") + self._dtype = np.dtype(dtype) + if self.dtype.hasobject: + raise TypeError("HostStaging does not support object dtype") self._slots: list[_Slot] = [] self._n_buffers = buffers self._next = 0 self._stream: Any = None + self._device_id: int | None = None def __repr__(self) -> str: return ( @@ -135,29 +165,50 @@ def __repr__(self) -> str: f"buffers={self._n_buffers})" ) + @property + def shape(self) -> tuple[int, ...]: + """The fixed shape of each snapshot.""" + return self._shape + + @property + def dtype(self) -> np.dtype: + """The fixed dtype of each snapshot.""" + return self._dtype + @property def buffers(self) -> int: """Number of buffers used in turn.""" return self._n_buffers def _slot(self, device: bool) -> _Slot: - if not self._slots: + if not self._slots or (device and self._slots[0].device_id is None): allocate = _empty_pinned if device else np.empty + # On the first GPU copy, replace ordinary host storage with pinned + # slots. Existing CPU handles keep their own completed snapshots. self._slots = [ - _Slot(allocate(self.shape, self.dtype)) for _ in range(self._n_buffers) + _Slot(allocate(self.shape, self.dtype), self._device_id) + for _ in range(self._n_buffers) ] + self._next = 0 slot = self._slots[self._next] self._next = (self._next + 1) % self._n_buffers return slot - def copy(self, array: Any) -> StagedCopy: + def copy(self, array: Any, *, stream: Any = None, event: Any = None) -> StagedCopy: """Start copying `array` to the host and return at once. Parameters ---------- array : cupy.ndarray | numpy.ndarray - An array of the staging shape and dtype. It may be overwritten - right after this call: a device array is snapshotted first. + An array of the staging shape and dtype. Later writes on the + snapshot's stream are ordered after it; another stream must wait + for the snapshot/copy before overwriting the source. + stream : cupy.cuda.Stream | None + Producer stream on which to snapshot, current stream by default. + event : cupy.cuda.Event | None + Producer completion to wait for on the current stream before the + snapshot. Pass at most one of stream and event. Host copies ignore + these dependencies because host input is assumed ready. Returns ------- @@ -165,45 +216,77 @@ def copy(self, array: Any) -> StagedCopy: ``ready()`` tells whether the copy has finished, ``result()`` waits for it and returns the host array. """ + if stream is not None and event is not None: + raise ValueError("pass only one of stream and event") if tuple(array.shape) != self.shape or np.dtype(array.dtype) != self.dtype: raise ValueError( f"HostStaging for {self.shape} {self.dtype} got an array of shape " f"{tuple(array.shape)} and dtype {array.dtype}", ) device = _is_device_array(array) + if device: + from cunumpy._streams import HostEvent, HostStream + + if isinstance(event, HostEvent) or isinstance(stream, HostStream): + raise TypeError("device copies require a CUDA producer stream or event") + cp = _cupy() + device_id = array.device.id + if cp.cuda.runtime.getDevice() != device_id: + raise ValueError("make the source array's CUDA device current") + if self._device_id is not None and self._device_id != device_id: + raise ValueError("HostStaging is bound to another CUDA device") + if stream is not None and getattr(stream, "device_id", device_id) not in ( + device_id, + -1, + ): + raise ValueError( + "HostStaging producer stream belongs to another device" + ) + self._device_id = device_id slot = self._slot(device) if slot.event is not None: - slot.event.synchronize() # the buffer's previous copy is done + with slot.context(): + slot.event.synchronize() # the buffer's previous copy is done slot.generation += 1 if not device: np.copyto(slot.host, np.asarray(array)) slot.event = None return StagedCopy(slot, slot.generation) - cp = _cupy() if self._stream is None: self._stream = cp.cuda.Stream(non_blocking=True) - if slot.device is None: - slot.device = cp.empty(self.shape, dtype=self.dtype) - if _COUNTERS: - _record("to_host", f"HostStaging.copy({self.shape} {self.dtype})") - # snapshot on the producer's stream, after the kernels that wrote `array` - slot.device[...] = array - snapshot_done = cp.cuda.get_current_stream().record() + # Snapshot on the producer stream, or wait for a supplied event on the + # current stream. Later writes must be ordered after this snapshot. + producer = cp.cuda.get_current_stream() if stream is None else stream + with producer: + if event is not None: + producer.wait_event(event) + if slot.device is None: + slot.device = cp.empty(self.shape, dtype=self.dtype) + slot.device[...] = array + snapshot_done = producer.record() # copy the snapshot to the pinned buffer on the staging stream self._stream.wait_event(snapshot_done) try: - # blocking=False (CuPy >= 13): return once the copy is enqueued - slot.device.get(stream=self._stream, out=slot.host, blocking=False) - except ( - TypeError - ): # older CuPy: a copy on a stream into pinned memory is asynchronous - slot.device.get(stream=self._stream, out=slot.host) - slot.event = self._stream.record() + try: + # blocking=False (CuPy >= 13): return once the copy is enqueued + slot.device.get(stream=self._stream, out=slot.host, blocking=False) + except TypeError: # older CuPy + slot.device.get(stream=self._stream, out=slot.host) + if _COUNTERS: + _record( + "to_host", + f"HostStaging.copy({self.shape} {self.dtype})", + nbytes=_nbytes(array), + ) + finally: + # Retain completion even if a copy failed after enqueuing work. + slot.event = self._stream.record() return StagedCopy(slot, slot.generation) def synchronize(self) -> None: """Wait for all copies in flight.""" for slot in self._slots: if slot.event is not None: - slot.event.synchronize() + with slot.context(): + slot.event.synchronize() diff --git a/src/cunumpy/_transfers.py b/src/cunumpy/_transfers.py index 509b63e..b7f4b98 100644 --- a/src/cunumpy/_transfers.py +++ b/src/cunumpy/_transfers.py @@ -11,7 +11,7 @@ propagator(dt) assert counter.total == 0, counter.report() -or, equivalently:: +or, to reject host/device copies and host execution on the GPU backend:: with xp.profiling.assert_no_transfers(): propagator(dt) @@ -25,13 +25,21 @@ * ``kernel_conversion``: a :class:`~cunumpy.kernels.PyccelKernel` call that copied device arrays to the host (and back), one event per call; * ``fallback``: a :class:`~cunumpy.kernels.Kernel` without CUDA kernel calling its host - kernel on the CuPy backend (``missing_cuda="fallback"``), one event per call. + kernel on the CuPy backend (``missing_cuda="fallback"``), one event per call; +* ``device_copy``: device-only dtype/layout conversions in CuNumpy helpers. + +Mirror refreshes, argument conversions, staging, serial MPI and kernel output +copy-back are also counted. Each physical host/device copy is recorded with its +payload size; conversion/fallback markers have no byte count. ``total`` counts +all observations, including markers. ``assert_no_transfers`` allows device-only +copies and rejects host/device copies and host fallback/conversion markers. Limitations ----------- Only transfers made *through cunumpy* are seen. Raw ``cupy.ndarray.get()``, ``cupy.asarray(numpy_array)``, ``numpy.asarray(cupy_array)``, ``float(device_array)``, -and implicit conversions inside other libraries are not counted; use ``nsys`` +forwarded backend operations such as ``xp.asarray`` and implicit conversions +inside other libraries are not counted; use ``nsys`` (or CuPy's own profiling hooks) to find those. Like the backend selection, the set of active counters is process-wide state @@ -41,6 +49,7 @@ from __future__ import annotations +import math import os import sys from collections.abc import Generator @@ -56,7 +65,7 @@ ] #: Event kinds, in the order they are reported. -KINDS = ("to_host", "to_device", "kernel_conversion", "fallback") +KINDS = ("to_host", "to_device", "kernel_conversion", "fallback", "device_copy") # The currently active counters, innermost last. Instrumented code checks # ``if _ACTIVE:`` before doing any work, so the overhead of an inactive counter @@ -75,17 +84,20 @@ class TransferEvent: ---------- kind : str One of ``"to_host"``, ``"to_device"``, ``"kernel_conversion"`` or - ``"fallback"``. + ``"fallback"``, or ``"device_copy"``. description : str What was transferred, e.g. ``"to_numpy(shape=(1000,), dtype=float64)"`` or ``"PyccelKernel 'push': 3 array(s) copied to the host"``. where : str The call site outside cunumpy, as ``"file:line"``. + nbytes : int | None + Payload bytes of a physical copy. None for markers or unknown sizes. """ kind: str description: str where: str + nbytes: int | None = None def __str__(self) -> str: return f"{self.where}: {self.kind}: {self.description}" @@ -141,9 +153,28 @@ def fallbacks(self) -> int: @property def total(self) -> int: - """Number of recorded transfers of all kinds.""" + """Number of observations, including conversion/fallback markers.""" return len(self.events) + def bytes(self, kind: str) -> int: + """Known bytes copied for `kind`; markers and unknown sizes add zero.""" + return sum(e.nbytes or 0 for e in self.events if e.kind == kind) + + @property + def bytes_to_host(self) -> int: + """Known payload bytes copied from device to host.""" + return self.bytes("to_host") + + @property + def bytes_to_device(self) -> int: + """Known payload bytes copied from host to device.""" + return self.bytes("to_device") + + @property + def device_copies(self) -> int: + """Device-only conversions recorded by CuNumpy helpers.""" + return self.count("device_copy") + def report(self) -> str: """A readable summary: events grouped by kind and call site, with counts.""" counts = ", ".join(f"{self.count(kind)} {kind}" for kind in KINDS) @@ -174,9 +205,9 @@ def _caller() -> str: return "" -def _record(kind: str, description: str) -> None: +def _record(kind: str, description: str, *, nbytes: int | None = None) -> None: """Record a transfer in every active counter (call only ``if _ACTIVE:``).""" - event = TransferEvent(kind, description, _caller()) + event = TransferEvent(kind, description, _caller(), nbytes) for counter in _ACTIVE: counter._add(event) @@ -190,6 +221,23 @@ def _describe(array: Any) -> str: return f"shape={tuple(shape)}, dtype={dtype}" +def _nbytes(array: Any) -> int | None: + """Array payload size, without copying data or reading a device scalar.""" + shape, dtype = getattr(array, "shape", None), getattr(array, "dtype", None) + if shape is None or dtype is None: + return None + return math.prod(shape) * dtype.itemsize + + +def _device_pointer(array: Any) -> int | None: + pointer = getattr(getattr(array, "data", None), "ptr", None) + if pointer is not None: + return pointer + interface = getattr(array, "__cuda_array_interface__", {}) + data = interface.get("data") + return None if data is None else data[0] + + @contextmanager def count_transfers() -> Generator[TransferCounter, None, None]: """Count the host/device transfers made through cunumpy in the block. @@ -225,7 +273,8 @@ def assert_no_transfers() -> Generator[TransferCounter, None, None]: """Raise ``AssertionError`` if the block makes a transfer through cunumpy. A :func:`count_transfers` block that, on exit, raises with the counter's - :meth:`~TransferCounter.report` if anything was counted. Only checked if + :meth:`~TransferCounter.report` if a host/device copy or host conversion/ + fallback was counted. Device-only conversions are allowed. Only checked if the block exits normally; an exception raised inside propagates as it is. Examples @@ -235,7 +284,7 @@ def assert_no_transfers() -> Generator[TransferCounter, None, None]: """ with count_transfers() as counter: yield counter - if counter.total: + if any(event.kind != "device_copy" for event in counter.events): raise AssertionError( "host/device transfers inside a block that must not transfer:\n" + counter.report(), diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 3da3621..c0f872b 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -14,7 +14,7 @@ import array_api_compat.numpy as np from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _describe, _record +from cunumpy._transfers import _describe, _device_pointer, _nbytes, _record if os.environ.get("CUNUMPY_FAKE_CUPY", "").strip().lower() in ("1", "true", "yes"): # tests without a GPU: a strict host stand-in for CuPy, see cunumpy._fake_cupy @@ -281,20 +281,33 @@ 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`. """ + result = _to_numpy(array) if _COUNTERS and get_array_backend(array) == "cupy": - _record("to_host", f"to_numpy({_describe(array)})") - return _to_numpy(array) + _record("to_host", f"to_numpy({_describe(array)})", nbytes=_nbytes(result)) + return result def to_cupy(array: Any) -> Any: """Convert an array to a CuPy array. - Anything that is not a CuPy array already is copied to the device, which - `count_transfers()` counts as a ``to_device`` transfer. + Host input is copied to the device and counted as a ``to_device`` transfer. + CUDA-array-interface inputs may be referenced without a copy; device-only + copies are recorded separately as ``device_copy``. """ - if _COUNTERS and get_array_backend(array) != "cupy": - _record("to_device", f"to_cupy({_describe(array)})") - return _to_cupy(array) + result = _to_cupy(array) + if _COUNTERS: + device = get_array_backend(array) == "cupy" or hasattr( + array, "__cuda_array_interface__" + ) + if not device: + _record("to_device", f"to_cupy({_describe(array)})", nbytes=_nbytes(result)) + elif _device_pointer(array) is not None and _device_pointer( + array + ) != _device_pointer(result): + _record( + "device_copy", f"to_cupy({_describe(array)})", nbytes=_nbytes(result) + ) + return result def as_device_array( @@ -304,18 +317,20 @@ def as_device_array( *, name: str | None = None, ) -> Any: - """Reference `value` on the device, or make one device copy of it. + """Reference `value` on the device, or convert it to a contiguous device array. The "reference or copy once" rule for building CUDA argument objects (`CudaArguments` subclasses, `CudaStruct` values): call it once when the argument object is built, never per kernel call. A CuPy array that already has the requested `dtype` (any dtype if `dtype` is None) and is C-contiguous is returned unchanged, the same object without a copy, so kernels write - into the caller's array. Anything else is converted with one device copy, + into the caller's array. Anything else is converted with ``cupy.ascontiguousarray(cupy.asarray(value, dtype))``: a tuple or list (e.g. ``degree = (3, 3, 3)``), a host NumPy array (one explicit transfer at build time), a device array of another dtype, or a non-contiguous view. The result passes the pointer checks of `CudaKernel` and `CudaStruct`. + Dtype and layout conversion can require separate device copies; each copy + through this helper is visible in transfer accounting. Raises on the NumPy backend: device argument objects are only built when running on CuPy, and host data is never copied to the device implicitly. @@ -361,7 +376,25 @@ def as_device_array( ): result = value else: - result = cp.ascontiguousarray(cp.asarray(value, dtype=dtype)) + converted = cp.asarray(value, dtype=dtype) + if _COUNTERS: + device_only = isinstance(value, cp.ndarray) or hasattr( + value, + "__cuda_array_interface__", + ) + if not device_only or _device_pointer(value) != _device_pointer(converted): + _record( + "device_copy" if device_only else "to_device", + f"as_device_array({_describe(value)})", + nbytes=_nbytes(converted), + ) + result = cp.ascontiguousarray(converted) + if _COUNTERS and result is not converted: + _record( + "device_copy", + f"as_device_array({_describe(converted)}) layout", + nbytes=_nbytes(result), + ) if ndim is not None and result.ndim != ndim: raise ValueError( f"{what} must have {ndim} dimension(s), got {result.ndim} " diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..cb4ce45 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +"""CUDA CI must fail if its hardware is unavailable rather than skip coverage.""" + +import os + +import pytest + + +def pytest_configure(config): + if os.environ.get("CUNUMPY_REQUIRE_CUDA", "").lower() not in {"1", "true", "yes"}: + return + import cunumpy as xp + + try: + xp.set_backend("cupy", strict=True) + import cupy as cp + + if getattr(cp, "__cunumpy_fake__", False): + raise RuntimeError("fake CuPy does not provide CUDA hardware coverage") + cp.zeros(1).get() # exercise allocation, execution and synchronization + except Exception as error: + raise pytest.UsageError( + f"CUDA CI requires a usable real GPU: {error}" + ) from error diff --git a/tests/unit/test_backend_numerics.py b/tests/unit/test_backend_numerics.py new file mode 100644 index 0000000..59ab52c --- /dev/null +++ b/tests/unit/test_backend_numerics.py @@ -0,0 +1,25 @@ +"""Backend numerical parity; performance comparisons belong in scope-profiler.""" + +import numpy as np +import pytest + +import cunumpy as xp + +pytestmark = pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") + + +def test_matmul_matches_numpy(): + rng = np.random.default_rng(123) + a = rng.normal(size=(64, 96)).astype(np.float32) + b = rng.normal(size=(96, 32)).astype(np.float32) + with xp.use_backend("cupy", strict=True): + actual = xp.to_numpy(xp.asarray(a) @ xp.asarray(b)) + np.testing.assert_allclose(actual, a @ b, rtol=2e-5, atol=2e-5) + + +def test_fft_matches_numpy(): + rng = np.random.default_rng(123) + values = (rng.normal(size=1024) + 1j * rng.normal(size=1024)).astype(np.complex64) + with xp.use_backend("cupy", strict=True): + actual = xp.to_numpy(xp.fft.fft(xp.asarray(values))) + np.testing.assert_allclose(actual, np.fft.fft(values), rtol=2e-5, atol=2e-5) diff --git a/tests/unit/test_benchmarks.py b/tests/unit/test_benchmarks.py deleted file mode 100644 index e89a0e9..0000000 --- a/tests/unit/test_benchmarks.py +++ /dev/null @@ -1,81 +0,0 @@ -import time - -import pytest - -import cunumpy as xp - - -@pytest.mark.skipif( - not xp.cupy_available(), - reason="CuPy/GPU not available or not functional", -) -def test_benchmark_matmul(): - """Benchmark matrix multiplication to show CuPy performance gain.""" - size = 2000 - - # --- Benchmark NumPy --- - with xp.use_backend("numpy"): - a_np = xp.random.rand(size, size).astype(xp.float32) - b_np = xp.random.rand(size, size).astype(xp.float32) - - start_np = time.perf_counter() - _ = a_np @ b_np - # No sync needed for NumPy as it is synchronous - end_np = time.perf_counter() - t_np = end_np - start_np - - # --- Benchmark CuPy --- - with xp.use_backend("cupy"): - a_cp = xp.random.rand(size, size).astype(xp.float32) - b_cp = xp.random.rand(size, size).astype(xp.float32) - - # Warm up - _ = a_cp @ b_cp - xp.synchronize() - - start_cp = time.perf_counter() - _ = a_cp @ b_cp - xp.synchronize() # CRITICAL for benchmarking GPU - end_cp = time.perf_counter() - t_cp = end_cp - start_cp - - print(f"\n[Benchmark] Size: {size}x{size}") - print(f"NumPy time: {t_np:.4f}s") - print(f"CuPy time: {t_cp:.4f}s") - print(f"Speedup: {t_np / t_cp:.2f}x") - - # On a real GPU (A100/A30), CuPy should be significantly faster - # We use a conservative threshold of 1.5x for the test to pass on various hardware - assert t_cp < t_np, f"CuPy ({t_cp:.4f}s) was not faster than NumPy ({t_np:.4f}s)" - - -@pytest.mark.skipif( - not xp.cupy_available(), - reason="CuPy/GPU not available or not functional", -) -def test_benchmark_fft(): - """Benchmark FFT performance.""" - size = 2**22 # ~4 million elements - - with xp.use_backend("numpy"): - data_np = xp.random.rand(size).astype(xp.complex64) - start = time.perf_counter() - _ = xp.fft.fft(data_np) - t_np = time.perf_counter() - start - - with xp.use_backend("cupy"): - data_cp = xp.random.rand(size).astype(xp.complex64) - # Warm up - _ = xp.fft.fft(data_cp) - xp.synchronize() - - start = time.perf_counter() - _ = xp.fft.fft(data_cp) - xp.synchronize() - t_cp = time.perf_counter() - start - - print(f"\n[Benchmark] FFT Size: {size}") - print(f"NumPy time: {t_np:.4f}s") - print(f"CuPy time: {t_cp:.4f}s") - print(f"Speedup: {t_np / t_cp:.2f}x") - assert t_cp < t_np diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 83eed65..68b0841 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -13,7 +13,7 @@ import pytest import cunumpy as xp -import cunumpy._device as device_module +import cunumpy._cuda_kernel as cuda_module from cunumpy import as_device_array from cunumpy.cuda import ( CudaArguments, @@ -72,6 +72,7 @@ class FakeDeviceArray: def __init__(self, dtype, ptr=0x1000, shape=(1,), strides=None, flags=None): self.dtype = np.dtype(dtype) + self.device = SimpleNamespace(id=0) self.data = SimpleNamespace(ptr=ptr) self.shape = tuple(shape) self.ndim = len(self.shape) @@ -2186,6 +2187,21 @@ def recorded(monkeypatch): """A CudaKernel whose launches are recorded instead of run.""" kernel = CudaKernel(AXPY, "axpy", block_size=128) raw = RecordingRawKernel() + limits = { + "name": b"recording GPU", + "maxThreadsPerBlock": 1024, + "maxThreadsDim": (1024, 1024, 64), + "maxGridSize": (2**31 - 1, 65535, 65535), + "sharedMemPerBlock": 48 * 1024, + "sharedMemPerBlockOptin": 100_000, + } + kernel._compiled[0] = cuda_module._CompiledKernel( + raw, + limits, + {"shared_size_bytes": 0, "max_threads_per_block": 1024}, + 48 * 1024, + ) + monkeypatch.setattr(cuda_module, "_current_device", lambda: 0) monkeypatch.setattr(kernel, "compile", lambda: raw) monkeypatch.setattr(kernel, "debug_active", lambda: False) return kernel, raw @@ -2208,13 +2224,8 @@ def test_n_threads_from_first_array(recorded): as_option.n_threads_from((1.0, 2)) -def test_shared_memory_above_the_default_is_opted_in(recorded, monkeypatch): +def test_shared_memory_above_the_default_is_opted_in(recorded): kernel, raw = recorded - monkeypatch.setattr( - device_module, - "max_shared_memory_per_block", - lambda opt_in=False: 100_000, - ) x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) kernel(1.0, x, y, 1, n_threads=1, shared_mem=40_000) # below 48 KiB: no setup assert raw.max_dynamic_shared_size_bytes == 48 * 1024 diff --git a/tests/unit/test_cuda_launch_limits.py b/tests/unit/test_cuda_launch_limits.py new file mode 100644 index 0000000..7a1120e --- /dev/null +++ b/tests/unit/test_cuda_launch_limits.py @@ -0,0 +1,315 @@ +"""Compilation/launch contracts with fake runtime limits and real GPU coverage.""" + +import io +import sys +import threading +import types + +import pytest + +import cunumpy as xp +from cunumpy import _cuda_kernel as implementation +from cunumpy.cuda import CudaKernel +from cunumpy.kernels import Kernel, KernelCatalog + +EMPTY = 'extern "C" __global__ void empty() {}' +requires_gpu = pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") + + +@pytest.fixture +def runtime(monkeypatch): + current = threading.local() + current.device = 0 + raw_kernels, queries, compilations = [], [], [] + properties = { + i: { + "name": f"test GPU {i}".encode(), + "maxThreadsDim": (1024, 1024, 64), + "maxGridSize": (2**31 - 1, 65535, 65535), + "maxThreadsPerBlock": 1024, + "sharedMemPerBlock": 49152, + "sharedMemPerBlockOptin": 98304, + } + for i in range(2) + } + + def device_id(): + return getattr(current, "device", 0) + + class Device: + def __init__(self, device): + self.id = device + + def __enter__(self): + self.previous = device_id() + current.device = self.id + return self + + def __exit__(self, *exc): + current.device = self.previous + + class Raw: + def __init__(self, source, name, options=()): + self.source, self.name, self.options = source, name, options + self.device = device_id() + self.launches = [] + self.attributes = { + "max_threads_per_block": 256, + "shared_size_bytes": 4096, + "max_dynamic_shared_size_bytes": 49152, + } + self.max_dynamic_shared_size_bytes = 49152 + raw_kernels.append(self) + + def compile(self, log_stream=None): + compilations.append((self.device, self.name, self.options)) + if "bad_syntax" in self.source: + raise RuntimeError("NVRTC rejected the source") + if log_stream is not None: + log_stream.write("compiled\n") + + def __call__(self, grid, block, args, shared_mem=0): + assert device_id() == self.device + self.launches.append((grid, block, shared_mem)) + + class Module: + def __init__(self, code, options, name_expressions): + self.raw = Raw(code, name_expressions[0], options) + self.compiled = False + + def compile(self, log_stream=None): + self.raw.compile(log_stream) + self.compiled = True + + def get_function(self, name): + assert self.compiled + assert name == self.raw.name + return self.raw + + def get_properties(device): + queries.append(device) + return properties[device] + + cp = types.SimpleNamespace( + RawKernel=Raw, + RawModule=Module, + cuda=types.SimpleNamespace( + Device=Device, + runtime=types.SimpleNamespace( + getDevice=device_id, getDeviceProperties=get_properties + ), + ), + ) + monkeypatch.setitem(sys.modules, "cupy", cp) + monkeypatch.setattr(xp.xp, "cupy_available", lambda: True) + monkeypatch.setattr(implementation, "_DEVICE_LIMITS", {}) + return types.SimpleNamespace( + cp=cp, + current=current, + properties=properties, + raw=raw_kernels, + queries=queries, + compilations=compilations, + ) + + +def test_compile_is_eager_cached_and_device_specific(runtime): + kernel = CudaKernel(EMPTY, "empty") + assert not kernel.is_compiled + log = io.StringIO() + first = kernel.compile(log_stream=log) + assert log.getvalue() == "compiled\n" + assert kernel.is_compiled + assert kernel.compile() is first + with runtime.cp.cuda.Device(1): + assert not kernel.is_compiled + second = kernel.compile() + assert second is not first + assert kernel.is_compiled + assert kernel.compile() is first + assert [c[0] for c in runtime.compilations] == [0, 1] + assert runtime.queries == [0, 1] + + +def test_failed_compile_is_not_cached_and_can_be_retried(runtime, monkeypatch): + kernel = CudaKernel(EMPTY, "empty") + original = runtime.cp.RawKernel.compile + + def fail(self, log_stream=None): + raise RuntimeError("compiler failed") + + monkeypatch.setattr(runtime.cp.RawKernel, "compile", fail) + with pytest.raises(RuntimeError, match="compiler failed"): + kernel.compile() + assert not kernel.is_compiled + monkeypatch.setattr(runtime.cp.RawKernel, "compile", original) + assert kernel.compile() is runtime.raw[-1] + assert kernel.is_compiled + + +def test_recompile_refreshes_headers_and_debug_on_only_current_device( + runtime, tmp_path +): + header = tmp_path / "value.cuh" + header.write_text("#define VALUE 1\n") + kernel = CudaKernel( + '#include "value.cuh"\n' + EMPTY, "empty", include_dirs=[tmp_path] + ) + first = kernel.compile() + with runtime.cp.cuda.Device(1): + second = kernel.compile() + header.write_text("#define VALUE 2\n") + with xp.cuda.cuda_debug(): + rebuilt = kernel.recompile() + assert rebuilt is not first + assert rebuilt.options != first.options + assert "-lineinfo" in rebuilt.options + with runtime.cp.cuda.Device(1): + assert kernel.compile() is second + + +def test_template_is_compiled_before_function_is_returned(runtime): + kernel = CudaKernel( + "template __global__ void empty() {}", + "empty", + template_args=["double"], + ) + assert kernel.compile().name == "empty" + assert len(runtime.compilations) == 1 + + +def test_catalog_fails_during_setup_and_workers_keep_selected_device(runtime): + good = CudaKernel(EMPTY, "empty") + bad = CudaKernel(EMPTY + "\nbad_syntax", "empty") + catalog = KernelCatalog( + {"good": Kernel(lambda: None, good), "bad": Kernel(lambda: None, bad)} + ) + with runtime.cp.cuda.Device(1): + with pytest.raises(RuntimeError, match="NVRTC rejected"): + catalog.compile_all(jobs=2) + assert good.is_compiled + assert not bad.is_compiled + assert all(c[0] == 1 for c in runtime.compilations) + + +@pytest.mark.parametrize( + "launch, expected", + [ + ({"grid": (1, 65536)}, "grid\\[1\\]=65536"), + ({"grid": 1, "block": (1, 1, 65)}, "block\\[2\\]=65"), + ({"grid": 1, "block": 512}, "256 device/kernel limit"), + ({"grid": 1, "shared_mem": 98304 - 4096 + 1}, "static shared memory 4096"), + ], +) +def test_actual_limits_reject_invalid_launch_before_execution( + runtime, launch, expected +): + kernel = CudaKernel(EMPTY, "empty") + with pytest.raises(ValueError, match=expected) as error: + kernel(**launch) + assert "test GPU 0" in str(error.value) + assert runtime.raw[0].launches == [] + + +def test_device_limits_can_be_stricter_than_kernel_limits(runtime): + runtime.properties[0]["maxThreadsPerBlock"] = 64 + with pytest.raises(ValueError, match="64 device/kernel limit"): + CudaKernel(EMPTY, "empty")(grid=1) + + +def test_static_memory_requires_opt_in_even_with_large_initial_dynamic_attribute( + runtime, +): + kernel = CudaKernel(EMPTY, "empty") + raw = kernel.compile() + state = kernel._compiled[0] + assert state.dynamic_shared == 45056 + kernel(grid=1, shared_mem=49152 - 4096 + 1) + assert raw.max_dynamic_shared_size_bytes == 45057 + + +def test_opt_in_accounts_for_static_memory_and_is_per_device(runtime): + kernel = CudaKernel(EMPTY, "empty") + kernel(grid=1, shared_mem=45057) + first = runtime.raw[-1] + assert first.max_dynamic_shared_size_bytes == 45057 + kernel(grid=1, shared_mem=60000) + kernel(grid=1, shared_mem=50000) + assert first.max_dynamic_shared_size_bytes == 60000 + with runtime.cp.cuda.Device(1): + kernel(grid=1, shared_mem=45057) + assert runtime.raw[-1].max_dynamic_shared_size_bytes == 45057 + assert len(runtime.queries) == 2 + assert len(runtime.compilations) == 2 + + +def test_stream_on_another_device_is_rejected(runtime): + with pytest.raises(ValueError, match="stream belongs to another device"): + CudaKernel(EMPTY, "empty")(grid=1, stream=types.SimpleNamespace(device_id=1)) + + +def test_empty_launch_does_not_compile(runtime): + CudaKernel(EMPTY, "empty")(n_threads=0) + assert runtime.compilations == [] + + +@requires_gpu +def test_invalid_cuda_fails_at_compile_not_first_launch(): + kernel = CudaKernel( + 'extern "C" __global__ void broken() { missing_symbol(); }', "broken" + ) + with pytest.raises(Exception, match="missing_symbol"): + kernel.compile() + assert not kernel.is_compiled + + +@requires_gpu +def test_kernel_launch_bounds_are_enforced(): + kernel = CudaKernel( + 'extern "C" __global__ __launch_bounds__(64) void limited() {}', + "limited", + check_signature=False, + block_size=64, + ) + raw = kernel.compile() + assert raw.attributes["max_threads_per_block"] == 64 + with pytest.raises(ValueError, match="64 device/kernel limit"): + kernel(grid=1, block=128) + kernel(grid=1) + xp.synchronize() + + +@requires_gpu +def test_same_kernel_compiles_on_multiple_devices(): + import cupy as cp + + if cp.cuda.runtime.getDeviceCount() < 2: + pytest.skip("requires two devices") + kernel = CudaKernel(EMPTY, "empty") + for device in (0, 1, 0): + with cp.cuda.Device(device): + kernel(grid=1) + assert kernel.is_compiled + cp.cuda.get_current_stream().synchronize() + + +@requires_gpu +def test_static_shared_memory_is_included_in_actual_gpu_limit(): + import cupy as cp + + source = r"""extern "C" __global__ void scratch(const double* x, double* out) { + __shared__ double values[128]; + values[threadIdx.x] = x[threadIdx.x]; + __syncthreads(); + if (threadIdx.x == 0) out[0] = values[127]; + }""" + kernel = CudaKernel(source, "scratch") + raw = kernel.compile() + static = raw.attributes["shared_size_bytes"] + assert static >= 128 * 8 + properties = cp.cuda.runtime.getDeviceProperties(cp.cuda.runtime.getDevice()) + limit = max( + properties["sharedMemPerBlock"], properties.get("sharedMemPerBlockOptin", 0) + ) + with pytest.raises(ValueError, match="static shared memory"): + kernel(cp.ones(128), cp.empty(1), n_threads=128, shared_mem=limit - static + 1) diff --git a/tests/unit/test_execution_transfers.py b/tests/unit/test_execution_transfers.py new file mode 100644 index 0000000..d483dd5 --- /dev/null +++ b/tests/unit/test_execution_transfers.py @@ -0,0 +1,221 @@ +"""Physical copies, fallback markers, and retained-buffer ownership without CUDA.""" + +import sys +import types + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import _mirror as mirrors +from cunumpy.memory import DeviceMirror + + +@pytest.fixture +def device(monkeypatch): + cp = types.SimpleNamespace(current=0, copies=[], waits=[]) + + class Array: + def __init__(self, data): + self.array = np.asarray(data) + self.data = types.SimpleNamespace(ptr=self.array.ctypes.data) + self.device = types.SimpleNamespace(id=cp.current) + self.shape, self.dtype, self.flags = ( + self.array.shape, + self.array.dtype, + self.array.flags, + ) + self.ndim = self.array.ndim + + def set(self, host): + cp.copies.append("set") + np.copyto(self.array, host) + + def get(self, out=None, stream=None): + cp.copies.append(("get", stream)) + if out is None: + return self.array.copy() + np.copyto(out, self.array) + return out + + def fill(self, value): + self.array.fill(value) + + def __setitem__(self, index, value): + self.array[index] = value.array if isinstance(value, Array) else value + + def asarray(value, dtype=None): + if isinstance(value, Array): + if dtype is None or value.dtype == np.dtype(dtype): + return value + value = value.array + return Array(np.array(value, dtype=dtype, copy=True)) + + def contiguous(value, dtype=None): + value = asarray(value, dtype=dtype) + return ( + value + if value.flags.c_contiguous + else Array(np.ascontiguousarray(value.array)) + ) + + cp.ndarray = Array + cp.asarray = asarray + cp.ascontiguousarray = contiguous + cp.zeros = lambda shape, dtype: Array(np.zeros(shape, dtype=dtype)) + cp.cuda = types.SimpleNamespace( + runtime=types.SimpleNamespace(getDevice=lambda: cp.current), + get_current_stream=lambda: types.SimpleNamespace(wait_event=cp.waits.append), + ) + monkeypatch.setitem(sys.modules, "cupy", cp) + monkeypatch.setattr(xp.xp.array_backend, "_backend", "cupy") + monkeypatch.setattr(xp.xp, "cupy_available", lambda: True) + monkeypatch.setattr(mirrors, "_cupy_backend", lambda: xp.get_backend() == "cupy") + monkeypatch.setattr(xp.xp, "_to_cupy", asarray) + monkeypatch.setattr( + xp.xp, + "get_array_backend", + lambda a: "cupy" if isinstance(a, Array) else "numpy", + ) + return cp + + +def test_mirror_initialization_and_refreshes_count_each_copy_once(device): + host = np.arange(4.0) + mirror = DeviceMirror(host) + with xp.profiling.count_transfers() as counter: + mirror.to_device() + mirror.to_device() + mirror.to_host() + assert counter.to_device == 2 and counter.to_host == 1 + assert counter.bytes_to_device == 2 * host.nbytes + assert counter.bytes_to_host == host.nbytes + assert [e.nbytes for e in counter.events] == [host.nbytes] * 3 + assert mirror.host is host + + +def test_mirror_refresh_breaks_no_transfer_assertion(device): + mirror = DeviceMirror(np.ones(2)) + mirror.to_device() + with ( + pytest.raises(AssertionError, match="DeviceMirror.to_device"), + xp.profiling.assert_no_transfers(), + ): + mirror.to_device() + with ( + pytest.raises(AssertionError, match="DeviceMirror.to_host"), + xp.profiling.assert_no_transfers(), + ): + mirror.to_host() + + +def test_zero_allocates_without_transferring_initial_host_contents(device): + mirror = DeviceMirror(np.ones(3)) + with xp.profiling.assert_no_transfers() as counter: + mirror.zero() + assert mirror.device is not mirror.host + assert counter.total == 0 + np.testing.assert_array_equal(mirror.device.array, np.zeros(3)) + + +@pytest.mark.parametrize("operation", ["device", "to_device", "to_host", "zero"]) +def test_mirror_rejects_another_current_device_before_touching_storage( + device, operation +): + mirror = DeviceMirror(np.ones(3)) + mirror.to_device() + device.current = 1 + with pytest.raises(ValueError, match="bound to CUDA device 0"): + getattr(mirror, operation)() if operation != "device" else mirror.device + assert device.copies == [] + + +def test_mirror_explicit_dependencies_and_wrong_stream(device): + mirror = DeviceMirror(np.ones(2)) + mirror.to_device() + producer = types.SimpleNamespace(device_id=0) + mirror.to_host(stream=producer) + assert device.copies[-1] == ("get", producer) + event = object() + mirror.to_host(event=event) + assert device.waits == [event] + with pytest.raises(ValueError, match="only one"): + mirror.to_host(stream=producer, event=event) + with pytest.raises(ValueError, match="another device"): + mirror.to_host(stream=types.SimpleNamespace(device_id=1)) + with pytest.raises(TypeError, match="CUDA producer"): + mirror.to_host(event=xp.cuda.HostEvent()) + + +def test_cpu_backend_refresh_does_not_touch_retained_device_buffer(device): + mirror = DeviceMirror(np.ones(2)) + mirror.to_device() + with xp.use_backend("numpy"), xp.profiling.assert_no_transfers(): + assert mirror.device is mirror.host + mirror.zero() + mirror.to_device().to_host() + assert device.copies == [] + mirror.to_device() + np.testing.assert_array_equal(mirror.device.array, [0.0, 0.0]) + + +def test_host_argument_conversion_counts_bytes_and_identity_does_not(device): + with xp.profiling.count_transfers() as counter: + array = xp.as_device_array([1, 2, 3], dtype=np.float32) + assert xp.as_device_array(array, dtype=np.float32) is array + assert counter.to_device == 1 and counter.bytes_to_device == 12 + with ( + pytest.raises(AssertionError, match="as_device_array"), + xp.profiling.assert_no_transfers(), + ): + xp.as_device_array(np.ones(2)) + + +def test_device_only_dtype_and_layout_copies_are_distinguished(device): + source = device.ndarray(np.zeros((4, 3))[:, ::2]) + with xp.profiling.assert_no_transfers() as counter: + packed = xp.as_device_array(source) + cast = xp.as_device_array(packed, dtype=np.float32) + assert packed.flags.c_contiguous and cast.dtype == np.float32 + assert counter.to_host == counter.to_device == 0 + assert counter.device_copies == 2 + assert counter.bytes("device_copy") == packed.array.nbytes + cast.array.nbytes + + +def test_failed_transfer_is_not_recorded(device, monkeypatch): + def fail(value): + raise RuntimeError("copy failed") + + monkeypatch.setattr(xp.xp, "_to_cupy", fail) + with ( + xp.profiling.count_transfers() as counter, + pytest.raises(RuntimeError, match="copy failed"), + ): + xp.to_cupy(np.zeros(2)) + assert counter.total == 0 + + +def test_kernel_output_counts_host_back_to_device_copy(device): + out = device.ndarray(np.zeros(4)) + with ( + xp.profiling.count_transfers() as counter, + xp.kernels.kernel_output(out, like=np.zeros(4)) as buffer, + ): + buffer[:] = 3.0 + assert counter.to_host == counter.to_device == 1 + assert counter.bytes_to_host == counter.bytes_to_device == out.array.nbytes + np.testing.assert_array_equal(out.array, np.full(4, 3.0)) + + +def test_interoperable_device_reference_does_not_report_host_transfer(device): + array = device.ndarray(np.ones(3)) + + class Exporter: + def __init__(self): + self.__cuda_array_interface__ = {"data": (array.data.ptr, False)} + + exported = Exporter() + device.asarray = lambda value, dtype=None: array + with xp.profiling.assert_no_transfers() as counter: + assert xp.as_device_array(exported) is array + assert counter.total == 0 diff --git a/tests/unit/test_mirror.py b/tests/unit/test_mirror.py index 769fba2..4b67b6e 100644 --- a/tests/unit/test_mirror.py +++ b/tests/unit/test_mirror.py @@ -19,6 +19,30 @@ reason="CuPy/GPU not available or not functional", ) + +@requires_gpu +@pytest.mark.parametrize("dependency", ["stream", "event"]) +def test_mirror_to_host_after_nondefault_producer(dependency): + import cupy as cp + + host = np.zeros(1000) + with xp.use_backend("cupy"): + mirror = DeviceMirror(host) + mirror.zero() + xp.synchronize() + producer = cp.cuda.Stream(non_blocking=True) + event = cp.cuda.Event(disable_timing=True) + with producer: + mirror.device.fill(7.0) + event.record() + with xp.profiling.count_transfers() as counter: + mirror.to_host( + **{dependency: producer if dependency == "stream" else event} + ) + np.testing.assert_array_equal(host, np.full(1000, 7.0)) + assert counter.bytes_to_host == host.nbytes + + BIN_ADD = r""" #include diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index dbdfaa8..3e89c33 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -274,18 +274,29 @@ def __init__(self): # --------------------------------------------------------------------------- -def _launching(kernel): +def _launching(kernel, monkeypatch): """Replace compilation by a recorder of (grid, block) launches.""" launches = [] - kernel.compile = lambda: ( - lambda grid, block, values, shared_mem: launches.append(grid) + raw = lambda grid, block, values, shared_mem: launches.append(grid) + kernel.compile = lambda: raw + kernel._compiled[0] = cuda_kernel_module._CompiledKernel( + raw, + { + "maxThreadsDim": (1024, 1024, 64), + "maxGridSize": (2**31 - 1, 65535, 65535), + "maxThreadsPerBlock": 1024, + "sharedMemPerBlock": 49152, + }, + {"max_threads_per_block": 1024, "shared_size_bytes": 0}, + 49152, ) + monkeypatch.setattr(cuda_kernel_module, "_current_device", lambda: 0) return launches -def test_n_threads_from(): +def test_n_threads_from(monkeypatch): kernel = CudaKernel(SCALE, "scale", n_threads_from=lambda args: args[2]) - launches = _launching(kernel) + launches = _launching(kernel, monkeypatch) kernel(FakeDeviceArray(np.float64), 2.0, 300) assert launches == [(3,)] kernel(FakeDeviceArray(np.float64), 2.0, 300, n_threads=10) # explicit wins @@ -300,7 +311,7 @@ def test_n_threads_from(): def test_check_finite(monkeypatch): kernel = CudaKernel(SCALE, "scale", check_finite=True) assert kernel.check_finite - _launching(kernel) + _launching(kernel, monkeypatch) kernel._synchronize_after_launch = lambda *a: None fake_cupy = ModuleType("cupy") fake_cupy.isfinite = np.isfinite diff --git a/tests/unit/test_staging.py b/tests/unit/test_staging.py index 679e02a..78f4ff3 100644 --- a/tests/unit/test_staging.py +++ b/tests/unit/test_staging.py @@ -1,6 +1,7 @@ """Tests for `xp.memory.HostStaging`: background copies of device arrays to the host.""" import types +from contextlib import nullcontext import numpy as np import pytest @@ -49,6 +50,10 @@ class DeviceArray(np.ndarray): log = None + @property + def device(self): + return types.SimpleNamespace(id=getattr(self, "_device_id", 0)) + def get(self, stream=None, out=None, blocking=True): assert blocking is False # an asynchronous copy DeviceArray.log.append(("get", stream.name)) @@ -70,6 +75,13 @@ class FakeStream: def __init__(self, name): self.name = name self.count = 0 + self.device_id = 0 + + def __enter__(self): + return self + + def __exit__(self, *exc): + pass def record(self): self.count += 1 @@ -88,6 +100,8 @@ def fake_device(monkeypatch): cuda = types.SimpleNamespace( Stream=lambda non_blocking=False: FakeStream("staging"), get_current_stream=lambda: current, + Device=lambda device: nullcontext(), + runtime=types.SimpleNamespace(getDevice=lambda: 0), ) def empty(shape, dtype): @@ -145,6 +159,101 @@ def test_copies_are_counted_as_transfers(fake_device): with xp.profiling.count_transfers() as counter: staging.copy(np.zeros(2).view(DeviceArray)) assert counter.total == 1 + assert counter.bytes_to_host == 16 + + +def test_explicit_producer_stream_and_event_order_snapshot(fake_device): + log = fake_device + source = np.ones(3).view(DeviceArray) + producer = FakeStream("producer") + staging = HostStaging(3, source.dtype) + staging.copy(source, stream=producer) + assert ("record", "producer#1") in log + log.clear() + event = FakeEvent("produced") + staging.copy(source, event=event) + assert log.index(("wait", "compute", "produced")) < log.index( + ("record", "compute#1") + ) + assert ("synchronize", "produced") not in log + with pytest.raises(ValueError, match="only one"): + staging.copy(source, stream=producer, event=event) + + +def test_cpu_initialized_staging_upgrades_to_pinned_and_preserves_old_copy( + fake_device, monkeypatch +): + allocations = [] + + def pinned(shape, dtype): + allocations.append(shape) + return np.empty(shape, dtype) + + monkeypatch.setattr(staging_module, "_empty_pinned", pinned) + staging = HostStaging(2, np.float64) + old = staging.copy(np.array([1.0, 2.0])) + device_copy = staging.copy(np.array([3.0, 4.0]).view(DeviceArray)) + assert len(allocations) == staging.buffers + np.testing.assert_array_equal(old.result(), [1.0, 2.0]) + np.testing.assert_array_equal(device_copy.result(), [3.0, 4.0]) + assert old.result() is not device_copy.result() + + +def test_device_binding_is_checked_before_reusing_any_buffer(fake_device, monkeypatch): + staging = HostStaging(2, np.float64) + source = np.ones(2).view(DeviceArray) + copy = staging.copy(source) + source._device_id = 1 + with pytest.raises(ValueError, match="device current"): + staging.copy(source) + monkeypatch.setattr(staging_module._cupy().cuda.runtime, "getDevice", lambda: 1) + with pytest.raises(ValueError, match="bound to another"): + staging.copy(source) + assert copy.result().tolist() == [1.0, 1.0] + + +def test_wrong_stream_and_host_dependencies_are_rejected(fake_device): + staging = HostStaging(2, np.float64) + source = np.ones(2).view(DeviceArray) + stream = FakeStream("wrong") + stream.device_id = 1 + with pytest.raises(ValueError, match="stream belongs"): + staging.copy(source, stream=stream) + with pytest.raises(TypeError, match="CUDA producer"): + staging.copy(source, event=xp.cuda.HostEvent()) + + +def test_failed_copy_keeps_completion_before_slot_reuse(fake_device, monkeypatch): + staging = HostStaging(1, np.float64, buffers=1) + source = np.ones(1).view(DeviceArray) + original = DeviceArray.get + + def fail(self, **kwargs): + raise RuntimeError("copy failed after enqueue") + + monkeypatch.setattr(DeviceArray, "get", fail) + with pytest.raises(RuntimeError, match="copy failed"): + staging.copy(source) + monkeypatch.setattr(DeviceArray, "get", original) + result = staging.copy(source) + assert ("synchronize", "staging#1") in fake_device + assert result.result().tolist() == [1.0] + + +def test_host_copy_waits_for_previous_device_slot(fake_device): + staging = HostStaging(1, np.float64, buffers=1) + old = staging.copy(np.ones(1).view(DeviceArray)) + result = staging.copy(np.array([2.0])) + assert ("synchronize", "staging#1") in fake_device + with pytest.raises(RuntimeError, match="reused"): + old.result() + assert result.result().tolist() == [2.0] + + +@pytest.mark.parametrize("shape,dtype", [((-1,), float), ((2,), object)]) +def test_invalid_staging_storage(shape, dtype): + with pytest.raises((ValueError, TypeError)): + HostStaging(shape, dtype) def test_on_gpu(): @@ -157,3 +266,57 @@ def test_on_gpu(): copy = staging.copy(data) data *= 0.0 # overwritten right away on the compute stream np.testing.assert_array_equal(copy.result(), np.arange(1000.0)) + + +@pytest.mark.parametrize("dependency", ["stream", "event"]) +def test_nondefault_producer_on_gpu(dependency): + if not xp.cupy_available(): + pytest.skip("requires CUDA") + import cupy as cp + + source = cp.zeros(1000) + cp.cuda.get_current_stream().synchronize() + producer = cp.cuda.Stream(non_blocking=True) + produced = cp.cuda.Event(disable_timing=True) + with producer: + source.fill(7.0) + produced.record() + staging = HostStaging(source.shape, source.dtype, buffers=1) + host_copy = staging.copy(np.full(1000, 2.0)) + copy = staging.copy( + source, **{dependency: producer if dependency == "stream" else produced} + ) + np.testing.assert_array_equal(host_copy.result(), np.full(1000, 2.0)) + np.testing.assert_array_equal(copy.result(), np.full(1000, 7.0)) + # Copy completion makes the source safe to overwrite on another stream. + source.fill(9.0) + cp.cuda.get_current_stream().synchronize() + again = staging.copy(source) + with pytest.raises(RuntimeError, match="reused"): + copy.result() + np.testing.assert_array_equal(again.result(), np.full(1000, 9.0)) + + +def test_staging_storage_description_cannot_change(): + staging = HostStaging(2, np.float64) + with pytest.raises(AttributeError): + staging.shape = (4,) + with pytest.raises(AttributeError): + staging.dtype = np.dtype(np.float32) + + +def test_staged_result_restores_callers_device_on_gpu(): + if not xp.cupy_available(): + pytest.skip("requires CUDA") + import cupy as cp + + if cp.cuda.runtime.getDeviceCount() < 2: + pytest.skip("requires two devices") + with cp.cuda.Device(0): + staging = HostStaging(3, np.float64) + copy = staging.copy(cp.arange(3.0)) + with cp.cuda.Device(1): + np.testing.assert_array_equal(copy.result(), np.arange(3.0)) + assert cp.cuda.runtime.getDevice() == 1 + with pytest.raises(ValueError, match="bound to another"): + staging.copy(cp.zeros(3)) diff --git a/tests/unit/test_transfers.py b/tests/unit/test_transfers.py index a431c69..773cc11 100644 --- a/tests/unit/test_transfers.py +++ b/tests/unit/test_transfers.py @@ -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)" + "TransferCounter(to_host=0, to_device=0, kernel_conversion=0, fallback=0, device_copy=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)" + "(3 to_host, 1 to_device, 0 kernel_conversion, 0 fallback, 0 device_copy)" ) assert " to_host (3):" in lines assert " to_device (1):" in lines @@ -278,7 +278,9 @@ def scale(x, y, factor): assert np.array_equal(x.data, np.full(3, 18.0)) # 1 * 2 * 3 * 3 (aliased) assert np.array_equal(y.data, np.full(3, 2.0)) assert counter.kernel_conversions == 2 - assert counter.total == 2, "the copies inside the call are not counted twice" + assert counter.total == 8 # six copies and two conversion markers + assert counter.to_host == counter.to_device == 3 + assert counter.bytes_to_host == counter.bytes_to_device == 3 * 3 * 8 first, second = counter.kernel_conversion_calls assert first.kind == "kernel_conversion" assert first.description == ( @@ -377,8 +379,9 @@ def scale(x, factor): PyccelKernel(scale)(x, 2.0) assert cp.all(x == 2.0) - assert counter.kernel_conversions == 1 and counter.total == 1 - assert counter.events[0].description == ( + assert counter.kernel_conversions == 1 and counter.total == 3 + assert counter.bytes_to_host == counter.bytes_to_device == x.nbytes + assert counter.kernel_conversion_calls[0].description == ( "PyccelKernel 'scale': 1 device array(s) copied to the host" ) @@ -403,7 +406,7 @@ def scale(x, factor, n): assert cp.all(x == 2.0) assert counter.fallbacks == 1 assert counter.kernel_conversions == 1, "the fallback converts on the host" - assert counter.total == 2 + assert counter.total == 4 # two copies, conversion and fallback markers @requires_cupy From 98e42057f18355788bb691228151857f46286af1 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 21:40:04 +0200 Subject: [PATCH 8/9] Formatting + bugfixes --- CHANGELOG.md | 2 ++ README.md | 3 ++- docs/source/api.md | 7 +++++-- docs/source/kernels/debugging.md | 7 ++++--- src/cunumpy/LLM_GUIDE.md | 6 +++++- src/cunumpy/_kernel.py | 4 ++-- src/cunumpy/_transfers.py | 8 ++++++++ src/cunumpy/xp.py | 16 ++++++++++------ tests/unit/test_cuda_launch_limits.py | 14 ++++++++++++++ tests/unit/test_execution_transfers.py | 20 ++++++++++++++++++++ 10 files changed, 72 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a2e8b6..6e70370 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Development - Formatting is checked with `ruff format` and import order with ruff's isort rules (`ruff check`), on `src/` and `tests/`, instead of black and isort; the `dev` extra installs ruff. +- GPU matrix/FFT correctness tests check numerical parity without timing assertions; + performance comparisons belong in scope-profiler. ### Fixed - `CudaKernel`'s header hash (`-DCUNUMPY_INCLUDE_HASH`) now covers the headers shipped with cunumpy (`cunumpy/atomic.cuh`, `reduce.cuh`, ...), also when included in angle brackets. Before, an upgrade of cunumpy that changed one of them left CuPy's kernel cache serving the kernel compiled with the old header. `resolve_includes(..., angle_dirs=...)` tracks angle-bracket includes found in the given directories. diff --git a/README.md b/README.md index 370e333..2653299 100644 --- a/README.md +++ b/README.md @@ -371,7 +371,8 @@ When building such argument objects, `xp.as_device_array(value, dtype, ndim=None)` applies the "reference or copy once" rule: a CuPy array that already has the dtype and is C-contiguous is returned as it is, anything else (a tuple such as `degree = (3, 3, 3)`, a host array, another dtype, a -non-contiguous view) becomes one device copy. Call it once when the object is +non-contiguous view) is converted; dtype and layout changes can require separate +device copies. Call it once when the object is built, not per kernel call; on the NumPy backend it raises, so host data is never copied to the device implicitly. When the host kernel takes such a group as one object too (e.g. a Pyccel class diff --git a/docs/source/api.md b/docs/source/api.md index 6fd171f..861842c 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -350,11 +350,14 @@ argument object is built, never per kernel call: * a CuPy array that already has `dtype` (any dtype if `dtype` is `None`) and is C-contiguous is returned unchanged, the same object without a copy, so kernels write into the caller's array; -* anything else becomes one C-contiguous device copy, +* anything else is converted to a C-contiguous device array using `cupy.ascontiguousarray(cupy.asarray(value, dtype))`: tuples and lists (`degree = (3, 3, 3)`), host NumPy arrays (one explicit transfer at build time), device arrays of another dtype, and non-contiguous views. +Dtype and layout conversion can require separate device copies; each copy +through this helper appears in transfer accounting. + The result passes the pointer checks of `CudaKernel` and `CudaStruct`. On the NumPy backend it raises `RuntimeError`: device arguments are only built when running on CuPy, and host data is never copied to the device implicitly. If @@ -374,7 +377,7 @@ class DeviceParticles(xp.cuda.CudaArguments): For the arguments of a `Kernel` with `dispatch="arrays"`, whose choice follows the arrays. `as_kernel_array` returns `value` on the side of `like` (a CuPy array if `like` is one, a NumPy array otherwise), C-contiguous and with `dtype` -(any if `None`): `value` itself if it already is such an array, else one copy, +(any if `None`): `value` itself if it already is such an array, else a conversion, moved across if needed (counted by `count_transfers()`). `kernel_output` is a context manager yielding the buffer for an array the kernel writes: `out` itself if `as_kernel_array` takes it unchanged, else a converted copy whose diff --git a/docs/source/kernels/debugging.md b/docs/source/kernels/debugging.md index 73753ba..4333eb1 100644 --- a/docs/source/kernels/debugging.md +++ b/docs/source/kernels/debugging.md @@ -46,8 +46,9 @@ with xp.cuda.cuda_debug(): Compile options are fixed when a kernel is compiled. A kernel that was already compiled before debug mode was enabled keeps its options (the synchronization -still applies). Enable debug mode before the kernels are first called, or set -`CUNUMPY_CUDA_DEBUG=1` in the environment. `kernel.debug_active()` and +still applies). Enable debug mode before the kernels are first called, call +`kernel.recompile()` after enabling it, or set `CUNUMPY_CUDA_DEBUG=1` in the +environment. Finish in-flight launches before rebuilding. `kernel.debug_active()` and `kernel.compile_options()` show what applies to a kernel. Use `-DCUNUMPY_BOUNDS_CHECK` in your own code too: @@ -108,7 +109,7 @@ on for the kernel under suspicion, not in production. | correct on small inputs, wrong on large ones | `int` overflow in index computations; use `long long`. | | results differ slightly between runs | floating-point atomics in a different order; compare with a tolerance. | | results occasionally garbage after MPI | missing `synchronize_for_mpi()` before the MPI call. | -| a header change has no effect | the header is included with angle brackets (not hashed) or the kernel object was compiled before the change; see "Headers and the compile cache" in [Writing CUDA kernels](cuda-kernel.md). | +| a header change has no effect | the kernel object was compiled before the change; call `kernel.recompile()`. Headers must resolve through the configured include directories to be hashed; see "Headers and the compile cache" in [Writing CUDA kernels](cuda-kernel.md). | After an illegal memory access the CUDA context of the process is unusable: all later CUDA calls fail. Restart the process (or the pytest run) after fixing the diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index a07f73c..05b1fa2 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -290,6 +290,10 @@ xp.cuda.cuda_include_dir() `ArrayND` params: CuPy arrays of dtype T and ndim N, any strides. * Compiled lazily with NVRTC on first call; cached on disk by CuPy; quoted `#include "..."` headers are hashed into the options so edits recompile. + `compile()` compiles eagerly and reuses successful per-device state; + `recompile()` refreshes that state after header/debug changes. Each launch + validates actual device/kernel dimensions, threads and static+dynamic shared + memory. Use the returned CuPy raw kernel's attributes for resource inspection. * Creating a `CudaKernel` does not import CuPy; compiling needs a GPU. Shipped CUDA headers (always on the include path): @@ -381,7 +385,7 @@ m = xp.memory.DeviceMirror(host_numpy_array) # TypeError if not numpy.ndarray m.device # CuPy copy (lazy) on CuPy; the host array itself on NumPy m.zero() m.to_device() -m.to_host(stream=None, event=None) # copies in place; no-op on NumPy; pass at most one dependency +m.to_host(stream=None, event=None) # in place; pass at most one dependency m.rebind(new_host_array) # after the owner reallocates ``` diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index a7fe6a7..d626273 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -35,7 +35,7 @@ import numpy as np from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _describe, _nbytes, _record +from cunumpy._transfers import _describe, _is_device_copy, _nbytes, _record from cunumpy.xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy __all__ = ["CompiledHostKernel", "KernelArguments", "PyccelKernel", "resolve_host_args"] @@ -836,7 +836,7 @@ def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: if not is_gpu(value): value = to_cupy(value) result = cupy.ascontiguousarray(value, dtype=dtype) - if _COUNTERS and result is not value: + if _COUNTERS and _is_device_copy(value, result): _record( "device_copy", f"as_kernel_array({_describe(value)})", diff --git a/src/cunumpy/_transfers.py b/src/cunumpy/_transfers.py index b7f4b98..f2d642e 100644 --- a/src/cunumpy/_transfers.py +++ b/src/cunumpy/_transfers.py @@ -238,6 +238,14 @@ def _device_pointer(array: Any) -> int | None: return None if data is None else data[0] +def _is_device_copy(source: Any, result: Any) -> bool: + """Distinguish a newly allocated conversion from a view of the same storage.""" + if result is source: + return False + pointer = _device_pointer(source) + return pointer is None or pointer != _device_pointer(result) + + @contextmanager def count_transfers() -> Generator[TransferCounter, None, None]: """Count the host/device transfers made through cunumpy in the block. diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index c0f872b..272c068 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -14,7 +14,13 @@ import array_api_compat.numpy as np from cunumpy._transfers import _ACTIVE as _COUNTERS -from cunumpy._transfers import _describe, _device_pointer, _nbytes, _record +from cunumpy._transfers import ( + _describe, + _device_pointer, + _is_device_copy, + _nbytes, + _record, +) if os.environ.get("CUNUMPY_FAKE_CUPY", "").strip().lower() in ("1", "true", "yes"): # tests without a GPU: a strict host stand-in for CuPy, see cunumpy._fake_cupy @@ -301,9 +307,7 @@ def to_cupy(array: Any) -> Any: ) if not device: _record("to_device", f"to_cupy({_describe(array)})", nbytes=_nbytes(result)) - elif _device_pointer(array) is not None and _device_pointer( - array - ) != _device_pointer(result): + elif _device_pointer(array) is not None and _is_device_copy(array, result): _record( "device_copy", f"to_cupy({_describe(array)})", nbytes=_nbytes(result) ) @@ -382,14 +386,14 @@ def as_device_array( value, "__cuda_array_interface__", ) - if not device_only or _device_pointer(value) != _device_pointer(converted): + if not device_only or _is_device_copy(value, converted): _record( "device_copy" if device_only else "to_device", f"as_device_array({_describe(value)})", nbytes=_nbytes(converted), ) result = cp.ascontiguousarray(converted) - if _COUNTERS and result is not converted: + if _COUNTERS and _is_device_copy(converted, result): _record( "device_copy", f"as_device_array({_describe(converted)}) layout", diff --git a/tests/unit/test_cuda_launch_limits.py b/tests/unit/test_cuda_launch_limits.py index 7a1120e..4834bc5 100644 --- a/tests/unit/test_cuda_launch_limits.py +++ b/tests/unit/test_cuda_launch_limits.py @@ -248,6 +248,20 @@ def test_stream_on_another_device_is_rejected(runtime): CudaKernel(EMPTY, "empty")(grid=1, stream=types.SimpleNamespace(device_id=1)) +@pytest.mark.parametrize("wrapped", [False, True]) +def test_array_on_another_device_is_rejected_before_launch(runtime, wrapped): + array = types.SimpleNamespace( + __cuda_array_interface__={}, device=types.SimpleNamespace(id=1) + ) + argument = ( + types.SimpleNamespace(__cuda_args__=lambda: (array,)) if wrapped else array + ) + kernel = CudaKernel(EMPTY, "empty", check_signature=False) + with pytest.raises(ValueError, match="belongs to CUDA device 1"): + kernel(argument, grid=1) + assert runtime.raw[0].launches == [] + + def test_empty_launch_does_not_compile(runtime): CudaKernel(EMPTY, "empty")(n_threads=0) assert runtime.compilations == [] diff --git a/tests/unit/test_execution_transfers.py b/tests/unit/test_execution_transfers.py index d483dd5..da5ba0a 100644 --- a/tests/unit/test_execution_transfers.py +++ b/tests/unit/test_execution_transfers.py @@ -219,3 +219,23 @@ def __init__(self): with xp.profiling.assert_no_transfers() as counter: assert xp.as_device_array(exported) is array assert counter.total == 0 + + +@pytest.mark.parametrize("helper", ["as_device_array", "as_kernel_array"]) +def test_reference_only_layout_view_is_not_counted_as_copy(device, helper): + scalar = device.ndarray(np.array(3.0)) + view = device.ndarray(scalar.array.reshape(1)) + device.ascontiguousarray = lambda value, dtype=None: view + if helper == "as_device_array": + # An interoperable exporter takes the conversion path. The resulting + # one-dimensional view still references the original scalar's storage. + exported = types.SimpleNamespace( + __cuda_array_interface__={"data": (scalar.data.ptr, False)} + ) + device.asarray = lambda value, dtype=None: scalar + convert, value, kwargs = xp.as_device_array, exported, {} + else: + convert, value, kwargs = xp.kernels.as_kernel_array, scalar, {"like": scalar} + with xp.profiling.count_transfers() as counter: + assert convert(value, **kwargs) is view + assert counter.total == 0 From 1fd0f62ef7e9ddfbaf6b7eca1d8b8c9409704b5d Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 21:53:45 +0200 Subject: [PATCH 9/9] Fixed the min/max race in src/cunumpy/cuda/include/cunumpy/reduce.cuh --- CHANGELOG.md | 2 + src/cunumpy/cuda/include/cunumpy/reduce.cuh | 3 ++ tests/unit/test_cuda_collectives.py | 56 ++++++++++++++++++++- tests/unit/test_emulation.py | 2 +- 4 files changed, 60 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6e70370..a3a2539 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 performance comparisons belong in scope-profiler. ### Fixed +- Block min/max reductions synchronize warp shared-memory reads before lane 0 + overwrites the result slot, fixing Compute Sanitizer racecheck hazards. - `CudaKernel`'s header hash (`-DCUNUMPY_INCLUDE_HASH`) now covers the headers shipped with cunumpy (`cunumpy/atomic.cuh`, `reduce.cuh`, ...), also when included in angle brackets. Before, an upgrade of cunumpy that changed one of them left CuPy's kernel cache serving the kernel compiled with the old header. `resolve_includes(..., angle_dirs=...)` tracks angle-bracket includes found in the given directories. - `xp.testing.assert_kernels_agree` reads `CudaStructArguments` objects and struct values through their struct fields, so their arrays get the same names as the attributes of the host argument object (before, arrays behind properties were named after the private attribute holding the owner, and the comparison failed with "do not have the same array arguments"). diff --git a/src/cunumpy/cuda/include/cunumpy/reduce.cuh b/src/cunumpy/cuda/include/cunumpy/reduce.cuh index 06cf58b..ddcc422 100644 --- a/src/cunumpy/cuda/include/cunumpy/reduce.cuh +++ b/src/cunumpy/cuda/include/cunumpy/reduce.cuh @@ -136,6 +136,9 @@ __device__ T cunumpy_block_reduce(T v, Op op) __syncthreads(); if (warp == 0) { v = lane < n_warps ? partial[lane] : Op::fill(partial); + // Min/max padding lanes read partial[0], which lane 0 overwrites below. + // Shuffles do not order shared memory: finish every warp read first. + __syncwarp(mask); v = cunumpy_warp_reduce(v, op, mask); if (lane == 0) partial[0] = v; } diff --git a/tests/unit/test_cuda_collectives.py b/tests/unit/test_cuda_collectives.py index 30b0fc5..3abfdc1 100644 --- a/tests/unit/test_cuda_collectives.py +++ b/tests/unit/test_cuda_collectives.py @@ -83,7 +83,9 @@ def test_partial_warps_multidimensional_blocks_and_repeated_collectives(block): data = cp.asarray(host) inc, exc, again = (cp.empty_like(data) for _ in range(3)) sums, lo, hi = (cp.empty(3) for _ in range(3)) - kernel = CudaKernel(BLOCK_SOURCE, "collectives", block_size=block) + kernel = CudaKernel( + BLOCK_SOURCE, "collectives", block_size=block, options=("-lineinfo",) + ) kernel(data, inc, exc, again, sums, lo, hi, grid=3) rows = host.reshape(3, count) expected_inc = np.cumsum(rows, axis=1) @@ -95,6 +97,54 @@ def test_partial_warps_multidimensional_blocks_and_repeated_collectives(block): np.testing.assert_array_equal(cp.asnumpy(got), expected) +REPEATED_EXTREMA_SOURCE = r""" +#include +template +__global__ void repeated_extrema(const T* x, T* minima, T* maxima, int rounds) { + int t = cunumpy_block_thread(); + int count = cunumpy_block_threads(); + int i = blockIdx.x * count + t; + int stride = gridDim.x * count; + for (int r = 0; r < rounds; ++r) { + T v = ((r & 1) ? -x[i] : x[i]) + T(r); + T lo = cunumpy_block_min(v); + T hi = cunumpy_block_max(v); + minima[r * stride + i] = lo; + maxima[r * stride + i] = hi; + } +} +""" + + +@requires_cuda +@pytest.mark.parametrize("dtype", [np.int32, np.float64]) +@pytest.mark.parametrize("block", [17, 32, 33, 128, 1024, (7, 5)]) +def test_repeated_extrema_are_returned_to_every_thread(dtype, block): + import cupy as cp + + count, blocks, rounds = int(np.prod(block)), 3, 8 + host = np.random.default_rng(123).integers(-100, 100, blocks * count).astype(dtype) + data = cp.asarray(host) + minima, maxima = (cp.empty((rounds, data.size), dtype=dtype) for _ in range(2)) + kernel = CudaKernel( + REPEATED_EXTREMA_SOURCE, + "repeated_extrema", + block_size=block, + template_args=(dtype,), + options=("-lineinfo",), + ) + kernel(data, minima, maxima, rounds, grid=blocks) + rows = host.reshape(blocks, count) + values = np.stack([(rows if r % 2 == 0 else -rows) + r for r in range(rounds)]) + for actual, expected in ( + (minima, values.min(axis=2)), + (maxima, values.max(axis=2)), + ): + np.testing.assert_array_equal( + cp.asnumpy(actual), np.repeat(expected, count, axis=1) + ) + + MASKED_SOURCE = r""" #include extern "C" __global__ void masked(const double* x, unsigned int mask, @@ -223,7 +273,7 @@ def test_scan_header_is_resolved_and_emulation_refuses_warp_intrinsics(): ] if emulation_compiler() is None: pytest.skip("requires a C++ compiler") - with pytest.raises(NotImplementedError, match="warp shuffles"): + with pytest.raises(NotImplementedError, match="warp (shuffles|synchronization)"): emulate_cuda_kernel(kernel, *(np.zeros(32) for _ in range(7)), n_threads=32) @@ -240,6 +290,7 @@ def test_collective_headers_compile_with_all_supported_arithmetic_types(): inline int __ffs(unsigned v) { return v ? __builtin_ctz(v) + 1 : 0; } inline int __clz(unsigned v) { return __builtin_clz(v); } inline void __syncthreads() {} +inline void __syncwarp(unsigned) {} template T __shfl_sync(unsigned, T v, int) { return v; } template T __shfl_up_sync(unsigned, T v, int) { return v; } template T __shfl_xor_sync(unsigned, T v, int) { return v; } @@ -270,6 +321,7 @@ def test_collective_headers_compile_with_all_supported_arithmetic_types(): input=_STUBS + stubs + BLOCK_SOURCE + + REPEATED_EXTREMA_SOURCE + MASKED_SOURCE + ATOMICS_SOURCE + instantiate, diff --git a/tests/unit/test_emulation.py b/tests/unit/test_emulation.py index ff21718..2bc3640 100644 --- a/tests/unit/test_emulation.py +++ b/tests/unit/test_emulation.py @@ -182,7 +182,7 @@ def test_warp_shuffles_in_an_included_header_are_refused(): " double s = cunumpy_block_sum(1.0); if (threadIdx.x == 0) out[0] = s; }", "k", ) - with pytest.raises(NotImplementedError, match="warp shuffles"): + with pytest.raises(NotImplementedError, match="warp (shuffles|synchronization)"): emulate_cuda_kernel(kernel, np.zeros(1), 1, n_threads=32)