From e1debfac399f8f80090eb271754893632b6dcc42 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 4 Oct 2026 22:29:43 +0200 Subject: [PATCH] Use CUNUMPY_BACKEND for environment variable and calculate n_threads based on the shape of the first array --- .gitlab-ci.yml | 3 +- CHANGELOG.md | 7 +- README.md | 9 +- docs/source/api.md | 41 ++++++-- docs/source/best-practices.md | 2 +- docs/source/guides/backends.md | 8 +- docs/source/guides/gpu-devices.md | 2 +- docs/source/installation.md | 10 +- docs/source/kernels/cuda-kernel.md | 18 ++++ docs/source/kernels/dispatch.md | 9 +- docs/source/kernels/overview.md | 2 +- docs/source/kernels/testing.md | 4 +- docs/source/pyodide.md | 2 +- docs/source/quickstart.md | 2 +- docs/source/troubleshooting.md | 4 +- src/cunumpy/LLM_GUIDE.md | 13 ++- src/cunumpy/_cuda_kernel.py | 89 +++++++++++++--- src/cunumpy/_dispatch.py | 3 +- src/cunumpy/_emulation.py | 4 +- src/cunumpy/_fake_cupy.py | 2 +- src/cunumpy/kernel_testing.py | 5 +- src/cunumpy/xp.py | 2 +- tests/unit/test_automatic_launch.py | 122 ++++++++++++++++++++++ tests/unit/test_backend_environment.py | 42 ++++++++ tests/unit/test_cuda_kernel.py | 21 ++++ tests/unit/test_emulation.py | 16 ++- tests/unit/test_kernel_dispatch.py | 8 +- tests/unit/test_kernel_dispatch_arrays.py | 3 + tests/unit/test_kernel_testing.py | 4 +- tests/unit/test_namespaces.py | 2 +- tests/unit/test_porting_helpers.py | 2 +- 31 files changed, 395 insertions(+), 66 deletions(-) create mode 100644 tests/unit/test_automatic_launch.py create mode 100644 tests/unit/test_backend_environment.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index eb5f82f..86f7a54 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -1,7 +1,7 @@ variables: CUDA_IMAGE: "gitlab-registry.mpcdf.mpg.de/mpcdf/ci-module-image/nvhpcsdk_26:2026" PIP_DISABLE_PIP_VERSION_CHECK: "1" - ARRAY_BACKEND: "cupy" + CUNUMPY_BACKEND: "cupy" CUNUMPY_REQUIRE_CUDA: "1" stages: @@ -42,6 +42,7 @@ gpu_sanitizers: compute-sanitizer --tool "$SANITIZER" --error-exitcode 1 --target-processes all python3 -m pytest -q tests/unit/test_cuda_collectives.py + tests/unit/test_automatic_launch.py tests/unit/test_cuda_launch_limits.py tests/unit/test_reusable_streams.py tests/unit/test_staging.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a3a2539..46cfb95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- CUDA launches infer thread counts from the first array by default, including + arrays in argument objects. 1D blocks use rows; multidimensional blocks use + matching leading shape axes. Explicit sizes and callbacks override inference. - Per-device eager CUDA compilation, compiler log streams and explicit `recompile()`. - Launch validation against actual device/kernel dimensions, thread limits and static plus dynamic shared memory, with cached per-device opt-in state. @@ -57,6 +60,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Support for Python 3.8 and 3.9 (both end-of-life); `cunumpy` now requires Python 3.10 or newer. ### Changed +- Renamed the startup backend selector from `ARRAY_BACKEND` to `CUNUMPY_BACKEND`; + update batch-job environments. CuNumpy-owned environment options use `CUNUMPY_*`. - `Kernel(..., dispatch="arrays")` calls the host kernel directly for host arguments, without `PyccelKernel`'s check for device arrays to convert, which converted nothing but ran while CuPy was the active backend (unless the `PyccelKernel` was built with `use_cupy=True`). - `Kernel.check_signature()` (and `KernelCatalog.check_signatures()`) also compares the parameters of every host implementation (numba, NumPy; nothing is compiled) and of a `CompiledHostKernel`'s fallback with those of the host kernel. - `KernelCatalog.from_package(..., compile_host=..., host_fallback=...)` builds `HostImplementations` instead of `CompiledHostKernel`s: `host_fallback` becomes the `"numpy"` implementation, and `_numba.py`/`_numpy.py` files in a kernel folder are picked up. @@ -75,7 +80,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `cunumpy/morton.cuh` and `xp.morton_keys`, `morton_encode`, `morton_decode`, `morton_scales`, `MAX_MORTON_LEVELS`: Morton (Z-order) keys of 2D and 3D points (`uint64`, up to 32 and 21 bits per axis), the same in a kernel and on the host (bit for bit), for sorting particles along a space-filling curve and building quadtrees and octrees from sorted keys. - `xp.sort_by_key(keys, *arrays)`: one stable argsort of `keys` applied to every array; returns the sorted keys, the order and the sorted arrays. - `Kernel.from_folder(package, ...)`: the kernel of one kernel folder, with the options of `KernelCatalog.from_package` (`host_suffix`, `compile_host`, `dispatch`, `include_dirs`, CUDA options...). Every version in the folder is an implementation: `.py` (`"pyccel"`, compiled with `compile_host`, and `"python"`, uncompiled), `_numba.py`, `_numpy.py` and `_cuda.cu`. A folder's own `__init__.py` can declare `kernel = xp.Kernel.from_folder(__name__, ...)`, so code imports the kernel from where it is written; `from_package` now builds each kernel with it. `Kernel.implementations` lists them, `Kernel.selected(device=False)` tells which one a call runs. -- `xp.HostImplementations` and `xp.HOST_IMPLEMENTATIONS`: the host implementations of a kernel, loaded on first use; a call runs the default, the first available of pyccel, numba and NumPy (the uncompiled Python version, with a warning, if none is), or the one chosen with `xp.set_kernel_implementation(name)` / `with xp.use_kernel_implementation(name):` / `CUNUMPY_KERNEL_IMPLEMENTATION=name` (read at import), like `set_backend`/`use_backend`/`ARRAY_BACKEND`; a chosen implementation that a kernel lacks or cannot load raises instead of running another. `xp.get_kernel_implementation()` reads the setting. +- `xp.HostImplementations` and `xp.HOST_IMPLEMENTATIONS`: the host implementations of a kernel, loaded on first use; a call runs the default, the first available of pyccel, numba and NumPy (the uncompiled Python version, with a warning, if none is), or the one chosen with `xp.set_kernel_implementation(name)` / `with xp.use_kernel_implementation(name):` / `CUNUMPY_KERNEL_IMPLEMENTATION=name` (read at import), like `set_backend`/`use_backend`/`CUNUMPY_BACKEND`; a chosen implementation that a kernel lacks or cannot load raises instead of running another. `xp.get_kernel_implementation()` reads the setting. - `xp.as_kernel_array(value, like, dtype=None)` and `xp.kernel_output(out, like, dtype=None)`: bring the arguments of a `dispatch="arrays"` kernel to the side of the main array `like` (CuPy or NumPy, C-contiguous, `dtype`), without a copy when they already fit; `kernel_output` yields the buffer the kernel writes and copies it back into `out` if it had to be converted. - `CompiledHostKernel.fallback` returns the fallback. - `cunumpy/random.cuh` and `xp.philox_uniform`, `philox_uniform2`, `philox_normal`, `philox_normal2`, `philox4x32_10`: counter-based random numbers (Philox4x32-10, passing the Random123 known-answer tests) as a pure function of `(seed, stream, counter)`, the same in a kernel and on the host (uniform numbers bit for bit, normal numbers up to the last bits of the math functions), so kernels that draw random numbers can be compared with their host versions. diff --git a/README.md b/README.md index 2653299..92d709d 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ explanation and examples. ## Choose a backend -CuNumpy starts with NumPy unless `ARRAY_BACKEND=cupy` is set before import. +CuNumpy starts with NumPy unless `CUNUMPY_BACKEND=cupy` is set before import. You can also choose at runtime: ```python @@ -303,7 +303,10 @@ for `object_modules`, `is_array`, aliasing, and output declarations. `CudaKernel` wraps a CUDA C kernel (compiled with NVRTC through `cupy.RawKernel`) so that it is called with the same arguments as the host -kernel it mirrors, plus the number of threads. Arrays are never copied: they +kernel it mirrors. Thread counts default to the first array's leading shape +axes: one thread per row for 1D blocks, matching axes for 2D/3D blocks. Explicit +`n_threads`, `grid`, or a custom `n_threads_from` controls the launch when needed. +Arrays are never copied: they must be C-contiguous CuPy arrays. The `extern "C" __global__` signature is parsed once and every call is checked against it: Python scalars are cast to the declared C types, and a wrong argument count, an array of the wrong dtype @@ -335,7 +338,7 @@ kernel = xp.kernels.Kernel(axpy, xp.cuda.CudaKernel(AXPY, "axpy")) with xp.use_backend("cupy"): x = xp.arange(1000, dtype=xp.float64) y = xp.zeros(1000) - kernel(2.0, x, y, 1000, n_threads=1000) # runs the CUDA kernel + kernel(2.0, x, y, 1000) # infer n_threads = x.shape[0], run the CUDA kernel ``` On the CuPy backend, a `Kernel` without CUDA kernel raises diff --git a/docs/source/api.md b/docs/source/api.md index 861842c..c7c1c4c 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -93,7 +93,7 @@ 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 +The initial backend is NumPy unless `CUNUMPY_BACKEND=cupy` is set before CuNumpy is imported. Other values of this environment variable result in the NumPy default. @@ -870,6 +870,8 @@ xp.cuda.CudaKernel( template_args=None, check_signature=True, debug=None, + n_threads_from="auto", + check_finite=False, ) xp.cuda.CudaKernel.from_file(path, name=None, *, suffix="_cuda.cu", **kwargs) xp.cuda.CudaKernel.all_from_file(path, **kwargs) @@ -991,20 +993,36 @@ shape is given either by `n_threads` or by `grid`: subtracting static storage opts in through `max_dynamic_shared_size_bytes`. Block/grid dimensions and device/kernel thread limits are also checked. -Without `n_threads` and `grid`, a launch uses `n_threads_from(args)` if the -kernel has one (constructor argument and settable property): a function of the -argument tuple, or `"first_array"` for the length of the first array argument -(its first axis), i.e. one thread per marker: +Without `n_threads` and `grid`, the default `n_threads_from="auto"` infers the +launch from the first array argument, skipping scalars and zero-dimensional +arrays. Arrays inside `CudaArguments`, `CudaStructArguments`, and struct values +are searched in argument/field order. The effective block shape determines +how many leading array axes are used: a 1D block uses `shape[0]` (one thread per +row/particle), a 2D block uses `shape[:2]`, and a 3D block uses `shape[:3]`. These +axes map to CUDA x, y, z. A `block=` override changes the inference dimensions. +Missing arrays or insufficient array dimensions raise with instructions to +give an explicit size. ```python -push = xp.cuda.CudaKernel.from_file("push_cuda.cu", n_threads_from="first_array") +push = xp.cuda.CudaKernel.from_file("push_cuda.cu") push(positions, velocities, e_field, dt) # n_threads = positions.shape[0] ``` +Explicit `n_threads` or `grid` always takes precedence. Set `n_threads_from` +(constructor argument or settable property) to a callable such as +`lambda args: args[0].size` for flattened element kernels, or +`lambda args: args[0].shape[::-1]` for kernels whose x index follows columns. +`"first_array"` always uses the first axis; None disables inference and requires +explicit launch sizes. The same defaults apply through `kernels.Kernel` and +`kernel_testing.assert_kernels_agree`, and in CPU emulation. + Nothing is launched if the grid has a zero dimension (e.g. `n_threads=0`). -`kernel.launch_shape(n_threads=None, *, grid=None, block=None)` returns the +`kernel.launch_shape(n_threads=None, *, grid=None, block=None, args=None)` returns the `(grid, block)` a call would use, e.g. to size a per-block output: +Supply `args=(...)` to inspect an automatically inferred launch without running +the kernel, e.g. `kernel.launch_shape(args=(positions, velocities, e_field, dt))`. + ```python BLOCK_SUM = r""" extern "C" __global__ void block_sum(const double* x, double* out, int n) { @@ -1521,8 +1539,9 @@ kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) ``` calls the kernel of the active backend. The launch arguments are passed to the -CUDA kernel (`n_threads` or `grid` is required there) and ignored by the host -kernel. Arguments implementing `KernelArguments` are replaced by their +CUDA kernel and ignored by the host kernel. Omitted sizes use the CUDA kernel's +shape-based default or its configured callback. Arguments implementing +`KernelArguments` are replaced by their `__host_args__()` on the host path and flattened via `__cuda_args__()` on the CUDA path. `kernel.compile()` compiles the CUDA kernel now and returns whether there is one. @@ -1571,8 +1590,8 @@ Properties: `name`, `host_kernel`, `cuda_kernel`, `has_cuda`, `missing_cuda`, * `Kernel.host_parameters()` falls back to the `__pyccel__/.pyi` stub for a pyccel-compiled host function, so `check_signature()` works for compiled kernels. -* `Kernel.__call__` needs no `n_threads` when the CUDA kernel has - `n_threads_from`. +* `Kernel.__call__` infers thread counts by default; custom `n_threads_from` + callbacks are supported, and None requires explicit sizes. ## `kernels.KernelCatalog` diff --git a/docs/source/best-practices.md b/docs/source/best-practices.md index ede9dca..b3295a3 100644 --- a/docs/source/best-practices.md +++ b/docs/source/best-practices.md @@ -7,7 +7,7 @@ A condensed checklist. Each item links to the guide with the reasoning. * Import as `import cunumpy as xp` and always call `xp.(...)`; never `from cunumpy import `. ([Backend-agnostic code](guides/portable-code.md)) -* Choose the backend once, in the entry point, with `ARRAY_BACKEND` or +* Choose the backend once, in the entry point, with `CUNUMPY_BACKEND` or `set_backend()`. Library code never calls `set_backend()`. ([Choosing a backend](guides/backends.md)) * Read back `xp.get_backend()` after requesting CuPy; log it with diff --git a/docs/source/guides/backends.md b/docs/source/guides/backends.md index 83e05c3..c670a94 100644 --- a/docs/source/guides/backends.md +++ b/docs/source/guides/backends.md @@ -6,17 +6,21 @@ where newly created arrays live. ## Select the backend at start-up -The backend is NumPy unless the environment variable `ARRAY_BACKEND=cupy` is +The backend is NumPy unless the environment variable `CUNUMPY_BACKEND=cupy` is set when CuNumpy is first imported: ```bash -ARRAY_BACKEND=cupy python simulate.py +CUNUMPY_BACKEND=cupy python simulate.py ``` This is the least intrusive option for scripts and batch jobs: the code does not change, and a job script decides whether it runs on a GPU. The variable is read once, at import; setting it later in `os.environ` has no effect. +The startup setting is named `CUNUMPY_BACKEND`; migrate existing job scripts +from `ARRAY_BACKEND`. CuNumpy-owned environment options share the `CUNUMPY_` +prefix; the complete list is in [Installation](../installation.md). + To choose from inside the program, for example from a command-line flag, call `set_backend()` once, early, before arrays are created: diff --git a/docs/source/guides/gpu-devices.md b/docs/source/guides/gpu-devices.md index 0edc448..80b8861 100644 --- a/docs/source/guides/gpu-devices.md +++ b/docs/source/guides/gpu-devices.md @@ -30,7 +30,7 @@ The usual alternative is to restrict visibility from outside the process, which also works for libraries that do not know about CuNumpy: ```bash -CUDA_VISIBLE_DEVICES=1 ARRAY_BACKEND=cupy python simulate.py +CUDA_VISIBLE_DEVICES=1 CUNUMPY_BACKEND=cupy python simulate.py ``` For MPI programs with one rank per GPU, use `bind_local_device()` instead diff --git a/docs/source/installation.md b/docs/source/installation.md index 9bd2b42..5d06ee4 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -67,8 +67,16 @@ Tests that need a GPU are skipped automatically where CuPy is not functional. | Variable | Effect | | --- | --- | -| `ARRAY_BACKEND=cupy` | start with the CuPy backend instead of NumPy (read once, at import) | +| `CUNUMPY_BACKEND=cupy` | start with the CuPy backend instead of NumPy (read once, at import) | | `CUNUMPY_CUDA_DEBUG=1` | enable [CUDA debug mode](kernels/debugging.md) for all kernels | +| `CUNUMPY_KERNEL_IMPLEMENTATION=numpy` | choose the host kernel implementation (read at import) | +| `CUNUMPY_MPI=1` / `0` | require MPI / use serial MPI regardless of launcher detection | +| `CUNUMPY_FAKE_CUPY=1` | install the strict CPU stand-in for CuPy for tests | +| `CUNUMPY_REQUIRE_CUDA=1` | require a real usable GPU when starting the test suite (CI guard) | + +Use `CUNUMPY_BACKEND` in job scripts; the former `ARRAY_BACKEND` setting is no +longer read. Standard toolchain/device variables such as `CXX` and +`CUDA_VISIBLE_DEVICES` keep their standard meanings. MPI launchers also export node-local rank variables (`OMPI_COMM_WORLD_LOCAL_RANK`, `SLURM_LOCALID`, ...), which `xp.mpi.local_rank()` reads to pick a GPU per process. diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index df2abe5..6e77948 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/docs/source/kernels/cuda-kernel.md @@ -85,6 +85,24 @@ correct; the kernel then behaves like a raw `RawKernel`. kernel(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) ``` +Omit both `n_threads` and `grid` to infer the launch from the first array's shape. +The default `n_threads_from="auto"` uses `shape[0]` for a 1D block (one thread per +row/particle), `shape[:2]` for a 2D block, and `shape[:3]` for a 3D block. Axes map +to CUDA x, y, z. Scalar arguments and zero-dimensional arrays are skipped; +arrays inside supported argument objects are searched in field/argument order. +Per-call block overrides determine the inference dimensions too. + +```python +axpy(2.0, x, y, x.size) # infer n_threads = x.shape[0] +``` + +Explicit sizes override inference. Set `n_threads_from=lambda args: args[0].size` +for a flattened element kernel, or another callback for a different axis order +or logical work count. `n_threads_from=None` requires explicit sizes; +`"first_array"` always selects only the first axis. `launch_shape(args=(...))` +inspects the inferred grid/block without compilation or a launch. CPU emulation +and host/CUDA parity tests use the same defaults. + * **1D**: `n_threads=n` with the default `block_size=128` (set per kernel with `CudaKernel(..., block_size=256)`). * **2D/3D**: `n_threads=(nx, ny)` and a tuple block, e.g. diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md index 98430c5..d7a19b0 100644 --- a/docs/source/kernels/dispatch.md +++ b/docs/source/kernels/dispatch.md @@ -26,14 +26,15 @@ def axpy_host(a, x, y, n): axpy = xp.kernels.Kernel(axpy_host, xp.cuda.CudaKernel(AXPY, "axpy"), name="axpy") -axpy(2.0, x, y, x.size, n_threads=x.size) +axpy(2.0, x, y, x.size) # infer n_threads = x.shape[0] ``` * On the **NumPy backend** the host kernel is called with the positional arguments; `n_threads`, `grid`, `block`, `shared_mem` and `stream` are ignored. -* On the **CuPy backend** the CUDA kernel is launched; `n_threads` (or `grid`) - is required. +* On the **CuPy backend** the CUDA kernel is launched; thread counts default to + the first array's leading shape axes. Explicit `n_threads` or `grid` overrides + this choice; `n_threads_from=None` requires an explicit size. * A plain host function is wrapped in a [`PyccelKernel`](pyccel-kernel.md); pass `host_options={"outputs": (2,)}` to configure that wrapper, or pass a `PyccelKernel` you built yourself. @@ -179,7 +180,7 @@ xp.kernels.set_kernel_implementation(None) # back to the default ``` or `CUNUMPY_KERNEL_IMPLEMENTATION=numpy` for a whole run (read at import, like -`ARRAY_BACKEND`). A chosen implementation that a kernel does not have, or +`CUNUMPY_BACKEND`). A chosen implementation that a kernel does not have, or cannot load, raises `LookupError` instead of running another one: a benchmark of numba never silently measures NumPy. `kernel.implementations` lists the implementations, `kernel.selected()` names the one a call with host arrays runs diff --git a/docs/source/kernels/overview.md b/docs/source/kernels/overview.md index 775bdac..2ad1ddf 100644 --- a/docs/source/kernels/overview.md +++ b/docs/source/kernels/overview.md @@ -24,7 +24,7 @@ unchanged. ## Which one do I need? -* **"I just want my existing code to run with `ARRAY_BACKEND=cupy`."** Wrap the +* **"I just want my existing code to run with `CUNUMPY_BACKEND=cupy`."** Wrap the host kernels in `PyccelKernel`. Everything works, but every kernel call copies its arrays to the host and back. This is a correct starting point, not a fast one. diff --git a/docs/source/kernels/testing.md b/docs/source/kernels/testing.md index 8f279b5..18b135c 100644 --- a/docs/source/kernels/testing.md +++ b/docs/source/kernels/testing.md @@ -274,7 +274,7 @@ transfer counting, the backend branches of a simulation) runs only with CuPy present. For CI machines without a GPU, cunumpy ships a strict stand-in: ```bash -CUNUMPY_FAKE_CUPY=1 ARRAY_BACKEND=cupy pytest tests/ +CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy pytest tests/ ``` Its arrays live in host memory but are not NumPy arrays: `numpy.asarray(a)` @@ -328,7 +328,7 @@ offsets. are reported as skipped, and `emulate_cuda_kernel` tests check the CUDA kernels' arithmetic (the runner needs a C++ compiler, which Linux images have). -* Run it a second time with `CUNUMPY_FAKE_CUPY=1 ARRAY_BACKEND=cupy`, so the +* Run it a second time with `CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy`, so the CuPy code paths are exercised on the CPU runner too (kernel launches are skipped). * Run the same suite on a GPU runner, optionally with `CUNUMPY_CUDA_DEBUG=1` so diff --git a/docs/source/pyodide.md b/docs/source/pyodide.md index ae51d32..11dc95b 100644 --- a/docs/source/pyodide.md +++ b/docs/source/pyodide.md @@ -31,7 +31,7 @@ await pyodide.runPythonAsync(` `); ``` -NumPy is selected by default when `ARRAY_BACKEND` is unset. Leave it unset or set +NumPy is selected by default when `CUNUMPY_BACKEND` is unset. Leave it unset or set it to `numpy` before importing cuNumPy. `to_numpy` and `to_cunumpy` preserve existing NumPy arrays, including views. `synchronize` and `set_device` are no-ops on this backend. diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index e76a5fd..0ccde3c 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -23,7 +23,7 @@ results are ordinary NumPy (or CuPy) arrays. Select CuPy before the arrays are created, either from the shell: ```bash -ARRAY_BACKEND=cupy python my_script.py +CUNUMPY_BACKEND=cupy python my_script.py ``` or in the program: diff --git a/docs/source/troubleshooting.md b/docs/source/troubleshooting.md index 9abc881..fb436f5 100644 --- a/docs/source/troubleshooting.md +++ b/docs/source/troubleshooting.md @@ -9,8 +9,8 @@ the real error. Typical causes: no CuPy installed, a CuPy wheel for a different CUDA major version, no GPU visible (`CUDA_VISIBLE_DEVICES` empty, a login node without GPUs), or a driver too old for the CUDA runtime. -**`ARRAY_BACKEND=cupy` has no effect.** The variable is read once, when -CuNumpy is first imported. Setting `os.environ["ARRAY_BACKEND"]` after the +**`CUNUMPY_BACKEND=cupy` has no effect.** The variable is read once, when +CuNumpy is first imported. Setting `os.environ["CUNUMPY_BACKEND"]` after the import does nothing; use `xp.set_backend()`. **Code still runs on the CPU after `set_backend("cupy")`.** Look for diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 05b1fa2..e9a1b2d 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -30,7 +30,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. 2. **Do not use `np.` for data that should follow the backend.** `np.zeros` always allocates on the host. Using `numpy` for dtypes (`np.float64`), host-only I/O, and host-side random test data is fine. -3. **Select the backend once, at the program entry point** (`ARRAY_BACKEND=cupy` +3. **Select the backend once, at the program entry point** (`CUNUMPY_BACKEND=cupy` env var before import, or `xp.set_backend("cupy")` early). Library code must not call `set_backend()`. Use `with xp.use_backend(...)` for scoped switches (tests, CPU references). Backend state is process-global and not thread-safe. @@ -84,7 +84,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | shared-memory budget of a block | `xp.cuda.max_shared_memory_per_block()` (48 KiB without a GPU) | | random numbers inside a kernel, equal on the host | `#include `: `cunumpy_uniform(seed, particle_id, step)`; host: `xp.rng.philox_uniform(seed, ids, step)` | | sort points along a Z-curve / quadtree or octree nodes as contiguous ranges | `keys = xp.algorithms.morton_keys(pos, lower, upper, levels)`, `keys, order, pos = xp.algorithms.sort_by_key(keys, pos)`; in a kernel `#include `: `cunumpy_morton_key2(x, y, x0, y0, sx, sy, levels)` with `xp.algorithms.morton_scales(...)` | -| one thread per marker without passing n_threads | `CudaKernel(..., n_threads_from="first_array")` | +| one thread per marker without passing n_threads | default: `CudaKernel(...)` uses the first array's row count for 1D launches | | copy device arrays to the host for output without stalling | `xp.memory.HostStaging(shape, dtype)`: `c = staging.copy(a)` ... `c.result()` | | PIC recipes (compaction, sort by cell, MPI exchange, graphs) | docs guide "Particle codes" | | reproducible random numbers per MPI rank | `xp.rng.random_streams.seed(seed, rank=rank)`, then `xp.rng.random_streams.normal(...)` / `.generator()` | @@ -270,7 +270,7 @@ ks = xp.cuda.CudaKernel.all_from_file("ops.cu") # dict name -> k k(*args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None) k.compile(log_stream=None); k.recompile(log_stream=None) k.is_compiled # successful compilation on the current CUDA device -k.launch_shape(n_threads) -> (grid, block) +k.launch_shape(n_threads=None, *, grid=None, block=None, args=None) -> (grid, block) k.included_headers; k.compile_options(); k.debug_active() xp.cuda.CudaKernelVariants(factory).get(*key); .compile_all(keys, jobs=1) xp.cuda.ctype_of(np.float64) == "double" @@ -295,6 +295,11 @@ xp.cuda.cuda_include_dir() validates actual device/kernel dimensions, threads and static+dynamic shared memory. Use the returned CuPy raw kernel's attributes for resource inspection. * Creating a `CudaKernel` does not import CuPy; compiling needs a GPU. +* Default `n_threads_from="auto"` infers from the first array, including arrays + in supported argument objects: 1D block -> shape[0], 2D -> shape[:2], 3D -> + shape[:3] (x, y, z). Per-call block overrides apply. Explicit n_threads/grid + wins; use a callback for flattened `.size`, reversed axes or another work + count, or None to require explicit sizes. Missing arrays/axes raise. Shipped CUDA headers (always on the include path): @@ -434,7 +439,7 @@ from cunumpy.kernel_testing import parity_cases, check_parity @pytest.mark.parametrize("kernel", parity_cases(catalog)) # skip-marked if no test args def test_parity(kernel): check_parity(kernel) -# without a GPU: CUNUMPY_FAKE_CUPY=1 ARRAY_BACKEND=cupy pytest (fake CuPy: strict host +# without a GPU: CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy pytest (fake CuPy: strict host # stand-in, no kernel launches; fake_cupy_active(); requires_cupy skips) # argument classes with a pyccel host class: one object on both backends diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 1a37f52..c3f931f 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -1843,15 +1843,43 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.pack() -def _first_array_length(args: tuple[Any, ...]) -> int: - """``n_threads_from="first_array"``: the first axis of the first array argument.""" +def _array_shapes_in(args: Sequence[Any]) -> Iterator[tuple[int, ...]]: + """Array shapes in argument order, including supported argument objects.""" for arg in args: shape = getattr(arg, "shape", None) - 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", - ) + if shape is not None and hasattr(arg, "dtype"): + if len(shape) > 0: + yield tuple(shape) + elif isinstance(arg, CudaStructArguments) and hasattr(type(arg), "struct"): + yield from _array_shapes_in( + [getattr(arg, field.name, None) for field in arg.struct.fields] + ) + elif isinstance(arg, CudaStructValue): + yield from _array_shapes_in( + [arg[field.name] for field in arg.struct.fields] + ) + elif hasattr(arg, "__cuda_args__"): + yield from _array_shapes_in(arg.__cuda_args__()) + + +def _first_array_shape(args: Sequence[Any], dimensions: int) -> int | tuple[int, ...]: + shape = next(_array_shapes_in(args), None) + if shape is None: + raise TypeError( + "inferring n_threads needs an array argument; pass n_threads or grid, " + "or set n_threads_from to a callable", + ) + if len(shape) < dimensions: + raise ValueError( + f"cannot infer {dimensions}D n_threads from the first array's shape " + f"{shape}; pass n_threads or grid, or set n_threads_from to a callable", + ) + return int(shape[0]) if dimensions == 1 else shape[:dimensions] + + +def _first_array_length(args: tuple[Any, ...]) -> int: + """``n_threads_from="first_array"``: the first axis of the first array argument.""" + return typing.cast(int, _first_array_shape(args, 1)) def _as_shape(value: int | Sequence[int], what: str) -> tuple[int, ...]: @@ -1933,6 +1961,13 @@ class CudaKernel: (:func:`cunumpy.cuda.set_cuda_debug`, ``CUNUMPY_CUDA_DEBUG``) at every launch; True or False fix it for this kernel. The compile options are fixed when the kernel is compiled. + n_threads_from : {"auto", "first_array"} | callable | None + Default "auto" infers thread counts from the first array's leading + shape axes, matching the block dimensionality (1D: one thread per row). + Arrays in supported argument objects are included. "first_array" + always uses the first axis. A callable receives the positional argument + tuple; None requires an explicit launch size. Explicit `n_threads` or + `grid` overrides inference. Notes ----- @@ -1968,7 +2003,7 @@ def __init__( template_args: Sequence[Any] | None = None, check_signature: bool = True, debug: bool | None = None, - n_threads_from: Callable[[tuple[Any, ...]], Any] | str | None = None, + n_threads_from: Callable[[tuple[Any, ...]], Any] | str | None = "auto", check_finite: bool = False, ) -> None: self._block = self._check_block(_as_shape(block_size, "block_size")) @@ -2207,8 +2242,10 @@ def n_threads_from(self) -> Callable[[tuple[Any, ...]], Any] | None: `n_threads` nor `grid`, and returns `n_threads` (an integer or a tuple), e.g. ``lambda args: args[2].n_markers`` for a kernel whose third argument is a struct argument object with the marker count. - Set it to ``"first_array"`` for the most common case, one thread per - row: the length of the first array argument (its first axis). + Default ``"auto"`` uses the first array's leading axes, matching the + launch block dimensions; 1D launches use its first axis, one thread per + row. Supported argument objects are searched in field/argument order. + ``"first_array"`` always uses its first axis. None disables inference. Settable, also on the ``cuda_kernel`` of a :class:`~cunumpy.kernels.Kernel`. """ return self._n_threads_from @@ -2218,12 +2255,21 @@ def n_threads_from( self, value: Callable[[tuple[Any, ...]], Any] | str | None, ) -> None: - if value == "first_array": + automatic = isinstance(value, str) and value == "auto" + if automatic: + value = self._default_n_threads + elif isinstance(value, str) and value == "first_array": value = _first_array_length if value is not None and not callable(value): - raise TypeError("n_threads_from must be callable, 'first_array' or None") + raise TypeError( + "n_threads_from must be callable, 'auto', 'first_array' or None" + ) + self._automatic_threads = automatic self._n_threads_from = value + def _default_n_threads(self, args: tuple[Any, ...]) -> int | tuple[int, ...]: + return _first_array_shape(args, len(self._block)) + @property def check_finite(self) -> bool: """Whether every launch checks the floating-point arrays for NaN or inf. @@ -2371,16 +2417,24 @@ def launch_shape( *, grid: int | Sequence[int] | None = None, block: int | Sequence[int] | None = None, + args: Sequence[Any] | None = None, ) -> tuple[tuple[int, ...], tuple[int, ...]]: """The ``(grid, block)`` a call with these launch arguments uses. - See :meth:`__call__`. A grid with a zero dimension launches nothing. + Pass `args` to use default thread-count inference when neither + `n_threads` nor `grid` is given. See :meth:`__call__`. A grid with a + zero dimension launches nothing. """ block_shape = ( self._block if block is None else self._check_block(_as_shape(block, "block")) ) + if n_threads is None and grid is None and args is not None: + if self._automatic_threads: + n_threads = _first_array_shape(args, len(block_shape)) + elif self._n_threads_from is not None: + n_threads = self._n_threads_from(tuple(args)) if (n_threads is None) == (grid is None): raise TypeError("pass exactly one of n_threads and grid") @@ -2416,7 +2470,8 @@ def __call__( The launch shape is given either by `n_threads` (the grid is the number of threads divided by the block size, rounded up, per dimension) or by - an explicit `grid`. + an explicit `grid`. With neither, it is inferred from the first array's + leading axes (or the configured `n_threads_from` callback). Parameters ---------- @@ -2447,9 +2502,9 @@ def __call__( Dimensions, threads or static plus dynamic shared memory exceed this device/kernel's limits, or arrays/stream belong to another device. """ - if n_threads is None and grid is None and self._n_threads_from is not None: - n_threads = self._n_threads_from(args) - grid_shape, block_shape = self.launch_shape(n_threads, grid=grid, block=block) + grid_shape, block_shape = self.launch_shape( + n_threads, grid=grid, block=block, args=args + ) shared_mem = operator.index(shared_mem) if shared_mem < 0: raise ValueError(f"shared_mem must be non-negative, got {shared_mem}") diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index 7a22014..2fdcfa6 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -598,7 +598,8 @@ def __call__( n_threads, grid, block, shared_mem, stream Launch configuration of the CUDA kernel, see :meth:`CudaKernel.__call__ `; - `n_threads` (or `grid`) is required when the CUDA kernel is called. + Thread counts are inferred from array shapes by default; + `n_threads` or `grid` overrides that choice. Ignored by the host kernel. """ if self._dispatch == "arrays": diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index 1cbf4d3..499b888 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -312,7 +312,9 @@ def emulate_cuda_kernel( compiler = compiler or emulation_compiler() if compiler is None: raise RuntimeError("emulation needs a C++ compiler (set CXX or install c++)") - grid_shape, block_shape = kernel.launch_shape(n_threads, grid=grid, block=block) + grid_shape, block_shape = kernel.launch_shape( + n_threads, grid=grid, block=block, args=args + ) grid_shape = tuple(grid_shape) + (1,) * (3 - len(grid_shape)) block_shape = tuple(block_shape) + (1,) * (3 - len(block_shape)) diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index 804815c..f8d189d 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -22,7 +22,7 @@ Activate it before CuPy or cunumpy's backend is first used, either with the environment variable ``CUNUMPY_FAKE_CUPY=1`` (read when cunumpy is imported) or by calling :func:`install` first thing in a test session (e.g. in -``conftest.py``). Then ``ARRAY_BACKEND=cupy`` or ``xp.set_backend("cupy")`` +``conftest.py``). Then ``CUNUMPY_BACKEND=cupy`` or ``xp.set_backend("cupy")`` selects the fake like the real thing. Never installed with a real CuPy present (``install`` raises), never shipped diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 50d3622..684b128 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -293,8 +293,9 @@ def assert_kernels_agree( Launch configuration of the CUDA kernel, see :meth:`CudaKernel.__call__ `. `n_threads` may also be a function of the tuple of arguments, e.g. - ``lambda args: args[0].shape[0]``. One of `n_threads` and `grid` is - required unless the CUDA kernel has ``n_threads_from``. + ``lambda args: args[0].shape[0]``. Omitted sizes use the CUDA kernel's + shape-based default or its configured ``n_threads_from``; explicit sizes + are required when inference is disabled. rtol, atol : float Tolerances of ``numpy.testing.assert_allclose``. n_calls : int diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 272c068..93f0d48 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -154,7 +154,7 @@ def use_backend( array_backend = ArrayBackend( backend=( - "cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy" + "cupy" if os.getenv("CUNUMPY_BACKEND", "numpy").lower() == "cupy" else "numpy" ), verbose=False, ) diff --git a/tests/unit/test_automatic_launch.py b/tests/unit/test_automatic_launch.py new file mode 100644 index 0000000..36e73d0 --- /dev/null +++ b/tests/unit/test_automatic_launch.py @@ -0,0 +1,122 @@ +"""Shape-based CUDA launch defaults, without allocating device memory.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy.cuda import CudaArguments, CudaKernel, CudaStructArguments, CudaStructValue + +SOURCE = 'extern "C" __global__ void work() {}' + + +@pytest.mark.parametrize( + "shape,block,expected", + [ + ((1000,), 128, ((8,), (128,))), + ((1000, 6), 128, ((8,), (128,))), + ((0, 6), 128, ((0,), (128,))), + ((35, 19), (8, 4), ((5, 5), (8, 4))), + ((35, 19, 7, 3), (8, 4, 2), ((5, 5, 4), (8, 4, 2))), + ], +) +def test_default_inference_matches_leading_axes_to_block_dimensions( + shape, block, expected +): + kernel = CudaKernel(SOURCE, "work", block_size=block) + args = (2.0, np.empty(shape), np.empty(1)) + assert kernel.launch_shape(args=args) == expected + inferred = kernel.n_threads_from(args) + assert inferred == (shape[0] if isinstance(block, int) else shape[: len(block)]) + + +def test_block_override_changes_inferred_dimensions(): + kernel = CudaKernel(SOURCE, "work") + args = (np.empty((35, 19)),) + assert kernel.launch_shape(args=args, block=(8, 4)) == ((5, 5), (8, 4)) + kernel = CudaKernel(SOURCE, "work", block_size=(8, 4)) + assert kernel.launch_shape(args=args, block=16) == ((3,), (16,)) + + +def test_explicit_sizes_bypass_inference(): + kernel = CudaKernel(SOURCE, "work") + assert kernel.launch_shape(300, args=()) == ((3,), (128,)) + assert kernel.launch_shape(grid=4, args=()) == ((4,), (128,)) + with pytest.raises(TypeError, match="exactly one"): + kernel.launch_shape(300, grid=4, args=()) + + +def test_callbacks_and_explicit_opt_out(): + kernel = CudaKernel(SOURCE, "work", n_threads_from=lambda args: args[0]) + assert kernel.launch_shape(args=(300,)) == ((3,), (128,)) + kernel.n_threads_from = None + with pytest.raises(TypeError, match="exactly one"): + kernel.launch_shape(args=(np.empty(300),)) + kernel.n_threads_from = "auto" + assert kernel.launch_shape(args=(np.empty(300),)) == ((3,), (128,)) + kernel.n_threads_from = "first_array" + assert kernel.launch_shape(args=(np.empty((300, 7)),)) == ((3,), (128,)) + + +def test_invalid_setting_preserves_auto_inference(): + kernel = CudaKernel(SOURCE, "work") + with pytest.raises(TypeError, match="n_threads_from"): + kernel.n_threads_from = "invalid" + assert kernel.launch_shape(args=(np.empty((35, 19)),), block=(8, 4)) == ( + (5, 5), + (8, 4), + ) + + +def test_missing_array_and_insufficient_dimensions_raise_clear_errors(): + kernel = CudaKernel(SOURCE, "work") + with pytest.raises(TypeError, match="needs an array argument"): + kernel.launch_shape(args=(2.0, np.array(3.0))) + with pytest.raises(ValueError, match="cannot infer 2D n_threads"): + kernel.launch_shape(args=(np.empty(3), np.empty((5, 7))), block=(4, 4)) + assert kernel.launch_shape(args=(np.array(3.0), np.empty(300))) == ((3,), (128,)) + + +@pytest.mark.parametrize("kind", ["flattened", "struct_arguments", "struct_value"]) +def test_inference_finds_arrays_in_supported_argument_objects(kind): + array = np.empty((300, 6)) + if kind == "flattened": + args = CudaArguments(3.0, array) + elif kind == "struct_arguments": + + class Arguments(CudaStructArguments): + struct_name = "Arguments" + fields = (("count", "int"), ("data", "Array2D")) + + # Shape inspection does not pack host arrays into CUDA structs. + args = Arguments() + args.count, args.data = 300, array + else: + struct = SimpleNamespace( + fields=[SimpleNamespace(name="count"), SimpleNamespace(name="data")] + ) + args = CudaStructValue(struct, None, {"count": 300, "data": array}) + kernel = CudaKernel(SOURCE, "work") + assert kernel.launch_shape(args=(args, np.empty(1))) == ((3,), (128,)) + + +@pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA") +def test_automatic_multidimensional_launch_on_gpu(): + import cupy as cp + + source = r""" + #include + extern "C" __global__ void fill(Array2D a) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + long long j = blockIdx.y * (long long)blockDim.y + threadIdx.y; + if (i < a.shape[0] && j < a.shape[1]) a(i, j) = 10 * i + j; + } + """ + array = cp.zeros((5, 7), dtype=cp.int64) + kernel = CudaKernel(source, "fill", block_size=(2, 4), options=("-lineinfo",)) + with xp.profiling.assert_no_transfers(): + kernel(array) + np.testing.assert_array_equal( + cp.asnumpy(array), 10 * np.arange(5)[:, None] + np.arange(7) + ) diff --git a/tests/unit/test_backend_environment.py b/tests/unit/test_backend_environment.py new file mode 100644 index 0000000..aa003eb --- /dev/null +++ b/tests/unit/test_backend_environment.py @@ -0,0 +1,42 @@ +"""Backend selection from CuNumpy's prefixed environment setting at import.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +import cunumpy + + +@pytest.mark.parametrize( + "backend,legacy,expected", + [ + (None, None, "numpy"), + ("numpy", None, "numpy"), + ("cupy", None, "cupy"), + ("CuPy", None, "cupy"), + (None, "cupy", "numpy"), + ("numpy", "cupy", "numpy"), + ("cupy", "numpy", "cupy"), + ], +) +def test_prefixed_backend_is_selected_at_import(backend, legacy, expected): + env = {**os.environ, "CUNUMPY_FAKE_CUPY": "1"} + env["PYTHONPATH"] = str(Path(cunumpy.__file__).parents[1]) + env.pop("CUNUMPY_BACKEND", None) + env.pop("ARRAY_BACKEND", None) + if backend is not None: + env["CUNUMPY_BACKEND"] = backend + if legacy is not None: + env["ARRAY_BACKEND"] = legacy + code = ( + "import os, cunumpy as xp\n" + f"assert xp.get_backend() == {expected!r}\n" + "os.environ['CUNUMPY_BACKEND'] = 'numpy'\n" + f"assert xp.get_backend() == {expected!r}\n" + "xp.set_backend('numpy')\n" + "assert xp.get_backend() == 'numpy'\n" + ) + subprocess.run([sys.executable, "-c", code], env=env, check=True) diff --git a/tests/unit/test_cuda_kernel.py b/tests/unit/test_cuda_kernel.py index 68b0841..1309007 100644 --- a/tests/unit/test_cuda_kernel.py +++ b/tests/unit/test_cuda_kernel.py @@ -2224,6 +2224,27 @@ def test_n_threads_from_first_array(recorded): as_option.n_threads_from((1.0, 2)) +def test_default_launch_size_and_overrides(recorded): + kernel, raw = recorded + x = FakeDeviceArray(np.float64, shape=(1000, 6)) + y = FakeDeviceArray(np.float64, shape=(1000, 6)) + kernel(2.0, x, y, 1000) + assert raw.launches == [((8,), (128,), 0)] + kernel(2.0, x, y, 1000, n_threads=10) + kernel(2.0, x, y, 1000, grid=2) + assert raw.launches[-2:] == [((1,), (128,), 0), ((2,), (128,), 0)] + + +def test_empty_default_launch_does_not_compile(recorded, monkeypatch): + kernel, raw = recorded + monkeypatch.setattr(kernel, "compile", lambda: pytest.fail("empty launch compiled")) + x = FakeDeviceArray( + np.float64, shape=(0, 6), flags=SimpleNamespace(c_contiguous=True) + ) + kernel(2.0, x, x, 0) + assert raw.launches == [] + + def test_shared_memory_above_the_default_is_opted_in(recorded): kernel, raw = recorded x, y = FakeDeviceArray(np.float64), FakeDeviceArray(np.float64) diff --git a/tests/unit/test_emulation.py b/tests/unit/test_emulation.py index 2bc3640..1da8a7d 100644 --- a/tests/unit/test_emulation.py +++ b/tests/unit/test_emulation.py @@ -24,7 +24,7 @@ def test_axpy(): rng = np.random.default_rng(0) x, y = rng.random(1000), rng.random(1000) expected = y + 2.5 * x - emulate_cuda_kernel(CudaKernel(AXPY, "axpy"), 2.5, x, y, 1000, n_threads=1000) + emulate_cuda_kernel(CudaKernel(AXPY, "axpy"), 2.5, x, y, 1000) # the compiler may fuse y + a * x into one FMA, as NVRTC does by default np.testing.assert_allclose(y, expected, rtol=1e-15, atol=0) exact = rng.random(1000) @@ -102,6 +102,20 @@ def test_2d_launch(): np.testing.assert_array_equal(a, 10 * np.arange(5)[:, None] + np.arange(7)) +def test_inferred_2d_launch_updates_every_element(): + source = r""" + #include + extern "C" __global__ void fill(Array2D a) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + long long j = blockIdx.y * (long long)blockDim.y + threadIdx.y; + if (i < a.shape[0] && j < a.shape[1]) a(i, j) = 10 * i + j; + } + """ + array = np.zeros((5, 7), dtype=np.int64) + emulate_cuda_kernel(CudaKernel(source, "fill", block_size=(2, 4)), array) + np.testing.assert_array_equal(array, 10 * np.arange(5)[:, None] + np.arange(7)) + + HISTOGRAM = r""" #include "cunumpy/atomic.cuh" extern "C" __global__ diff --git a/tests/unit/test_kernel_dispatch.py b/tests/unit/test_kernel_dispatch.py index 9bebfde..ed88deb 100644 --- a/tests/unit/test_kernel_dispatch.py +++ b/tests/unit/test_kernel_dispatch.py @@ -85,9 +85,11 @@ def test_kernel_on_cupy(): with xp.use_backend("cupy"): assert kernel.get_kernel() is kernel.cuda_kernel kernel(x, 3, 300, n_threads=300) + kernel(x, 2, 300) + kernel.cuda_kernel.n_threads_from = None with pytest.raises(ValueError, match="n_threads is required"): kernel(x, 3, 300) - assert cp.all(x == 3.0) + assert cp.all(x == 6.0) def test_missing_cuda_raises_on_cupy(): @@ -261,9 +263,11 @@ def test_kernel_launch_configuration_on_cupy(): with xp.use_backend("cupy"): kernel(x, 2.0, 300, grid=5, block=64) # 320 threads kernel(x, 2.0, 300, n_threads=300, block=32, stream=cp.cuda.Stream.null) + kernel(x, 2.0, 300, block=32) + kernel.cuda_kernel.n_threads_from = None with pytest.raises(ValueError, match="n_threads is required"): kernel(x, 2.0, 300, block=32) - assert cp.all(x == 4.0) + assert cp.all(x == 8.0) def test_compile(kernel_package_factory): diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index 8a09c27..e28f60e 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -86,6 +86,9 @@ def test_arrays_dispatch_runs_device_arrays_on_the_gpu(fake_gpu): kernel(x, 3.0, 4, n_threads=4) ((name, args, options),) = fake_gpu assert name == "scale" and args[0] is x and options["n_threads"] == 4 + kernel(x, 3.0, 4) + assert fake_gpu[-1][2]["n_threads"] is None + kernel.cuda_kernel.n_threads_from = None with pytest.raises(ValueError, match="n_threads is required"): kernel(x, 3.0, 4) diff --git a/tests/unit/test_kernel_testing.py b/tests/unit/test_kernel_testing.py index af7fc1b..c5147ef 100644 --- a/tests/unit/test_kernel_testing.py +++ b/tests/unit/test_kernel_testing.py @@ -149,7 +149,7 @@ def test_assert_kernels_agree_arguments(): assert_kernels_agree(scale, make_scale_args, n_threads=300) with pytest.raises(ValueError, match="has no CUDA version"): assert_kernels_agree(Kernel(scale), make_scale_args, n_threads=300) - kernel = Kernel(scale, CudaKernel(SCALE_CUDA, "scale")) + kernel = Kernel(scale, CudaKernel(SCALE_CUDA, "scale", n_threads_from=None)) with pytest.raises(TypeError, match="n_threads"): assert_kernels_agree(kernel, make_scale_args) with pytest.raises(ValueError, match="n_calls"): @@ -172,7 +172,7 @@ def test_assert_kernels_agree_skips_without_cupy(): @pytest.mark.skipif(not xp.cupy_available(), reason="CuPy/GPU not available") def test_assert_kernels_agree_on_gpu(): kernel = Kernel(scale, CudaKernel(SCALE_CUDA, "scale")) - host = assert_kernels_agree(kernel, make_scale_args, n_threads=300, n_calls=2) + host = assert_kernels_agree(kernel, make_scale_args, n_calls=2) expected = 4 * np.random.default_rng(0).random(300) np.testing.assert_allclose(host["argument 0"], expected) diff --git a/tests/unit/test_namespaces.py b/tests/unit/test_namespaces.py index a4af07d..28245af 100644 --- a/tests/unit/test_namespaces.py +++ b/tests/unit/test_namespaces.py @@ -121,7 +121,7 @@ def test_switching_the_backend_replaces_the_names(): "xp.set_backend('numpy')\n" "assert type(xp.arange(3)).__module__ == 'numpy'\n" ) - env = {**os.environ, "CUNUMPY_FAKE_CUPY": "1", "ARRAY_BACKEND": "numpy"} + env = {**os.environ, "CUNUMPY_FAKE_CUPY": "1", "CUNUMPY_BACKEND": "numpy"} subprocess.run([sys.executable, "-c", code], check=True, env=env) diff --git a/tests/unit/test_porting_helpers.py b/tests/unit/test_porting_helpers.py index 3e89c33..20be042 100644 --- a/tests/unit/test_porting_helpers.py +++ b/tests/unit/test_porting_helpers.py @@ -625,7 +625,7 @@ def test_require_version(monkeypatch): def test_fake_cupy_in_subprocess(): root = Path(__file__).resolve().parents[2] - env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", ARRAY_BACKEND="cupy") + env = dict(os.environ, CUNUMPY_FAKE_CUPY="1", CUNUMPY_BACKEND="cupy") env["PYTHONPATH"] = os.pathsep.join( p for p in (str(root / "src"), env.get("PYTHONPATH", "")) if p )