From bdf89df94be4c04efd733a0b5ec6431b192257b5 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 20:18:44 +0200 Subject: [PATCH 1/5] - 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/5] 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/5] 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/5] 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/5] 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