diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index d9df172..2cdc674 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -28,27 +28,16 @@ jobs: ./cloc --version ./cloc $(git ls-files) - black: + ruff-format: runs-on: ubuntu-latest steps: - name: Checkout the code uses: actions/checkout@v4 - - name: Code formatting with black + - name: Code formatting with ruff run: | - pip install black "black[jupyter]" - black --check src/ - - isort: - runs-on: ubuntu-latest - steps: - - name: Checkout the code - uses: actions/checkout@v4 - - - name: Code formatting with isort - run: | - pip install isort - isort --check src/ + pip install ruff + ruff format --check src/ tests/ mypy: runs-on: ubuntu-latest @@ -80,10 +69,10 @@ jobs: - name: Checkout the code uses: actions/checkout@v4 - - name: Linting with ruff + - name: Linting with ruff (including import sorting) run: | pip install ruff - ruff check src/ + ruff check src/ tests/ pylint: runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index 76a0c91..362b607 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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. - Importing `cunumpy.cuda` makes `xp.cuda` the cunumpy submodule instead of CuPy's `cupy.cuda`; use `import cupy` for the latter. +- The implementation modules are private (`cunumpy._cuda_kernel`, `_dispatch`, `_kernel`, `_fusion`, `_philox`, `_morton`, `_random_streams`, `_staging`, `_mirror`, `_transfers`, `_emulation`, `_scipy_backend`, and the new `_device`, `_mpi`, `_profiling`, `_algorithms` split off `cunumpy.xp`), so every public name has one import path, through the submodules above. `cunumpy.cuda_kernel`, `cunumpy.dispatch` and `cunumpy.kernel` (in 0.3.0) still import, with a `DeprecationWarning`, until 0.6. `cunumpy.xp` keeps the backend selection and array conversion only. + +### Performance +- `xp.` for a NumPy/CuPy function (`xp.zeros`, `xp.add`, ...) is a plain attribute lookup, as fast as `numpy.` (before: about 3 us per access through the module `__getattr__`, 20 times `numpy.add`). The public names of the active backend module are copied into the `cunumpy` namespace and replaced when `set_backend`/`use_backend` change the backend (about 30 us per switch between NumPy and CuPy; nothing when the module stays the same). `dir(xp)` lists them, so IPython and Jupyter complete NumPy names. `ArrayBackend.add_listener(callback)` is called with the new module on every switch. + +### Development +- Formatting is checked with `ruff format` and import order with ruff's isort rules (`ruff check`), on `src/` and `tests/`, instead of black and isort; the `dev` extra installs ruff. ### Fixed - `CudaKernel`'s header hash (`-DCUNUMPY_INCLUDE_HASH`) now covers the headers shipped with cunumpy (`cunumpy/atomic.cuh`, `reduce.cuh`, ...), also when included in angle brackets. Before, an upgrade of cunumpy that changed one of them left CuPy's kernel cache serving the kernel compiled with the old header. `resolve_includes(..., angle_dirs=...)` tracks angle-bracket includes found in the given directories. @@ -52,7 +59,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `CudaKernel(..., check_finite=True)` (also a settable property): after every launch (synchronized) the floating-point and complex arrays among the arguments, including the array fields of struct argument objects, are scanned, and a NaN or inf raises `RuntimeError` naming the kernel and the argument. For debugging a kernel that produces NaN; costs a synchronization and a pass over the arrays per launch. - `cunumpy.testing.device_function_kernel` accepts struct parameters, by value (`DomainArgs d`) or by const reference (`const DomainArgs& d`), for the structs given in `structs=`; they are passed through to every thread unchanged, so device helpers that take argument structs can be tested from Python. - Test arguments next to the kernel: `KernelCatalog.from_package(..., test_args_suffix="_test_args")` records `/_test_args.py` as `Kernel.test_args_module` (imported on first access as `Kernel.test_args`; `Kernel(..., test_args=...)` sets it by hand). The module defines `make_args(backend, seed)` and `N_THREADS` (an integer, a tuple, or a function of the argument tuple) or `GRID`, optionally `BLOCK`, `RTOL`, `ATOL`, `N_CALLS`, `OUTPUTS`, `SEED` (`cunumpy.testing.TEST_ARGS_SETTINGS`). `cunumpy.testing.parity_cases(catalog)` gives one `pytest.param` per kernel with a CUDA version, kernels without a test-arguments module marked `skip` with a reason naming the missing file, and `check_parity(kernel, **overrides)` runs `assert_kernels_agree` with the module's settings. A ported package's parity test is then one parametrized test, and a developer who adds a kernel adds a file, not test code. -- `KernelCatalog.from_package(..., check_name_length=True)` warns about a kernel whose module name does not fit pyccel's Fortran wrapper module `bind_c__kernels` into Fortran's 63-character limit (`cunumpy.dispatch.FORTRAN_NAME_LIMIT`): with the default suffix a kernel name has at most 48 characters. +- `KernelCatalog.from_package(..., check_name_length=True)` warns about a kernel whose module name does not fit pyccel's Fortran wrapper module `bind_c__kernels` into Fortran's 63-character limit (`FORTRAN_NAME_LIMIT`): with the default suffix a kernel name has at most 48 characters. - `Kernel.host_parameters()` reads the parameter names of a pyccel-compiled host kernel from the `__pyccel__/.pyi` stub pyccel writes next to the extension module, so `Kernel.check_signature()` and `KernelCatalog.check_signatures()` also check compiled kernels. - The fake CuPy (`cunumpy._fake_cupy`): a strict host stand-in for CuPy for CI machines without a GPU, installed with the environment variable `CUNUMPY_FAKE_CUPY=1` (read when cunumpy is imported) or `cunumpy.testing.install_fake_cupy()`. Its arrays live in host memory but are not NumPy arrays (`numpy.asarray` raises, as for real CuPy arrays), reject host arrays and lists as CuPy does, have `data.ptr`, `device` and `__cuda_array_interface__`, so argument objects, struct packing, `as_device_array`, `count_transfers` and backend branches run on the CuPy code path on the CPU; CUDA kernels cannot run (`NotImplementedError`). `cunumpy.testing.fake_cupy_active()` tells; `requires_cupy` and `assert_kernels_agree` skip while it is active. - `xp.mpi_buffer(array, *, send=True, recv=False, cuda_aware=None)`: context manager yielding the buffer to hand to MPI: a host array unchanged; a device array unchanged (after `synchronize_for_mpi`) when MPI is CUDA-aware; otherwise a pinned host staging copy, copied from the device before the block (`send`) and back after it (`recv`), counted by `count_transfers()`. `cuda_aware=None` uses the answer recorded by `mpi_is_cuda_aware()` (which now remembers its result) or `xp.set_mpi_cuda_aware()`; `xp.get_mpi_cuda_aware()` reads it. One MPI call site for both backends and both kinds of MPI builds. @@ -61,10 +68,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `Kernel(..., dispatch="arrays")` and `KernelCatalog.from_package(..., dispatch="arrays")`: choose the CUDA kernel when an argument lives on the GPU (a CuPy array or a device-only argument object) and the host kernel otherwise, whatever the backend, for codes that hand host arrays to kernels while CuPy is active. The default, `dispatch="backend"`, is unchanged. - `xp.CompiledHostKernel`: a host kernel compiled on its first call by a compile function the caller provides (cunumpy does not compile anything itself, e.g. a wrapper around `pyccel.epyccel`), falling back to a given callable (or, with a warning, to the uncompiled Python function) when compilation fails. `KernelCatalog.from_package(..., compile_host=..., host_fallback=...)` uses it for every host kernel. - `Kernel.check_signature()`, `Kernel.host_parameters()` and `KernelCatalog.check_signatures()`: check that a host kernel and its CUDA kernel take the same parameters in the same order, listing every kernel that differs. -- `cunumpy.testing.emulate_cuda_kernel(kernel, *args, n_threads=...)` (in `cunumpy.emulation`) and `emulation_compiler()`: run a CUDA kernel on the CPU, one thread after another, on NumPy arrays (any strides, written back), with the kernel compiled as C++ and the CUDA built-ins stubbed, so CPU-only CI can check the arithmetic of ported kernels. Kernels using shared memory, `__syncthreads` or warp intrinsics are refused. +- `cunumpy.kernel_testing.emulate_cuda_kernel(kernel, *args, n_threads=...)` and `emulation_compiler()`: run a CUDA kernel on the CPU, one thread after another, on NumPy arrays (any strides, written back), with the kernel compiled as C++ and the CUDA built-ins stubbed, so CPU-only CI can check the arithmetic of ported kernels. Kernels using shared memory, `__syncthreads` or warp intrinsics are refused. - `Array4D` in `cunumpy/array_view.cuh` and as a kernel parameter, struct field and `from_signature` annotation (`"float[:, :, :, :]"`), e.g. for a 3D grid of vector components. - `xp.max_shared_memory_per_block(device=None, *, opt_in=False)` and `xp.DEFAULT_SHARED_MEMORY_PER_BLOCK`: the shared memory a block may use on a device (48 KiB without a GPU). -- `xp.random_streams` (`cunumpy.random_streams.RandomStreams`): one seeded generator per process and backend, with the stream `(seed, rank)` for MPI runs, a choice of NumPy bit generator, per-component generators and draw functions that work with NumPy and CuPy generators. +- `xp.rng.random_streams` (`cunumpy.rng.RandomStreams`): one seeded generator per process and backend, with the stream `(seed, rank)` for MPI runs, a choice of NumPy bit generator, per-component generators and draw functions that work with NumPy and CuPy generators. - Documentation restructured into getting started, user guide (backends, backend-agnostic code, data movement, devices, MPI, profiling), kernel porting guides (`PyccelKernel`, `CudaKernel`, `Kernel`/`KernelCatalog`, argument objects, accumulation, debugging, testing), worked examples, best practices and troubleshooting pages. - `cunumpy/LLM_GUIDE.md`: a self-contained guide to the API and its rules for AI coding assistants, shipped as package data and rendered in the documentation. - `xp.CudaKernel`: Wraps a CUDA C kernel (`cupy.RawKernel`, compiled lazily with NVRTC) so it can be called with the same arguments as the host kernel it mirrors, plus `n_threads`. The `extern "C" __global__` signature is parsed once and every call is checked against it: argument count, array dtypes (host arrays raise, they are never copied), and scalars (Python scalars are cast to the declared C types with range checks; lossy or mismatching scalars raise instead of reaching the kernel as silently wrong values). Supports `block_size`, NVRTC `options`, `include_dirs`, `shared_mem`, `stream`, `CudaKernel.from_file()` (`_cuda.cu`), `compile()` and `prepare_args()`; `check_signature=False` skips the checks. diff --git a/docs/source/api.md b/docs/source/api.md index 1ef95eb..2bc43af 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -14,8 +14,10 @@ total = xp.sum(values) ``` At runtime, NumPy-like attributes such as `array`, `sum`, `fft`, and `linalg` -are forwarded to the currently selected `array-api-compat` NumPy or CuPy -module. CuNumpy does not wrap every operation individually. The available +are those of the currently selected `array-api-compat` NumPy or CuPy module. +They are copied into the `cunumpy` namespace and replaced when the backend +changes, so `xp.sum` costs the same as `numpy.sum` (a switch between NumPy and +CuPy takes some tens of microseconds; avoid switching inside a hot loop). CuNumpy does not wrap every operation individually. The available operations and some details can therefore vary with the installed NumPy and CuPy versions. In normal use, access those operations through the top-level `cunumpy` namespace, commonly imported as `xp`. diff --git a/pyproject.toml b/pyproject.toml index f6e8447..ac9a151 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,8 +27,7 @@ dependencies = [ ] optional-dependencies.dev = [ - "black[jupyter]", - "isort", + "ruff", "cunumpy[test-compiled,docs]", ] # https://medium.com/@pratikdomadiya123/build-project-documentation-quickly-with-the-sphinx-python-2a9732b66594 @@ -53,9 +52,9 @@ where = [ "src" ] [tool.setuptools.package-data] cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ] -[tool.isort] -profile = "black" - [tool.ruff] # agent worktrees and scratch scripts are local copies, not part of the project extend-exclude = [ ".claude" ] +# formatting: `ruff format`; import sorting: the isort rules of `ruff check` +lint.extend-select = [ "I" ] +lint.isort.known-first-party = [ "cunumpy" ] diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index c82f348..d07b036 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -59,7 +59,8 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. `xp.memory.HostStaging`, `xp.petsc.petsc_vec`; kernel test helpers in `cunumpy.kernel_testing`. The old top-level names (`xp.CudaKernel`) and `cunumpy.testing` are deprecated (removed in 0.6); do not write new code - with them. + with them. Modules starting with `_` (`cunumpy._cuda_kernel`, ...) are + private; never import from them. ## Decision guide @@ -141,20 +142,26 @@ Only transfers through cunumpy are counted (not raw `cupy.asarray`, `.get()`, MPI, accumulation and versions: ```python -xp.mpi.mpi_is_cuda_aware(comm) # collective, once at startup; remembered -with xp.mpi.mpi_buffer(a) as buf: comm.Send(buf, ...) # host array, CUDA-aware device -with xp.mpi.mpi_buffer(a, send=False, recv=True) as buf: ... # array, or pinned staging copy +xp.mpi.mpi_is_cuda_aware(comm) # collective, once at startup; remembered +with xp.mpi.mpi_buffer(a) as buf: + comm.Send(buf, ...) # host array, CUDA-aware device +with xp.mpi.mpi_buffer(a, send=False, recv=True) as buf: + ... # array, or pinned staging copy xp.mpi.set_mpi_cuda_aware(True | False | None), xp.mpi.get_mpi_cuda_aware() -xp.algorithms.segment_sum(values, keys, n_segments) # out[k] = sum(values[keys == k]); keys < 0 dropped -keys, order, a, b = xp.algorithms.sort_by_key(keys, a, b) # stable argsort applied to every array -xp.require_version("0.4.0") # ImportError if cunumpy is older +xp.algorithms.segment_sum( + values, keys, n_segments +) # out[k] = sum(values[keys == k]); keys < 0 dropped +keys, order, a, b = xp.algorithms.sort_by_key( + keys, a, b +) # stable argsort applied to every array +xp.require_version("0.4.0") # ImportError if cunumpy is older ``` Random numbers and dtypes: ```python -rng = xp.rng.get_rng(seed=None) # numpy or cupy Generator for the active backend -xp.default_float_dtype() # float64 of the active backend +rng = xp.rng.get_rng(seed=None) # numpy or cupy Generator for the active backend +xp.default_float_dtype() # float64 of the active backend ``` NumPy and CuPy generators give different sequences for the same seed. For @@ -187,14 +194,18 @@ xp.mpi.synchronize_for_mpi(*buffers) # before every MPI call that touches devic Profiling: ```python -with xp.profiling.timed_region("name", sync=True) as t: ... # t.name, t.elapsed (s), t.synced -with xp.profiling.nvtx_range("name", color=None): ... # also usable as @decorator +with xp.profiling.timed_region("name", sync=True) as t: + ... # t.name, t.elapsed (s), t.synced +with xp.profiling.nvtx_range("name", color=None): + ... # also usable as @decorator ``` `PyccelKernel`: ```python -k = xp.kernels.PyccelKernel(fn, use_cupy=None, object_modules=(), is_array=None, outputs=None) +k = xp.kernels.PyccelKernel( + fn, use_cupy=None, object_modules=(), is_array=None, outputs=None +) k(*args, **kwargs) ``` @@ -272,25 +283,42 @@ function `` (host); optional `pkg//_cuda.cu` defines Argument objects: ```python -class Dev(xp.cuda.CudaArguments): # flattened into several CUDA params - def __init__(self, x, n): super().__init__(x, n) +class Dev(xp.cuda.CudaArguments): # flattened into several CUDA params + def __init__(self, x, n): + super().__init__(x, n) -class Args(xp.kernels.KernelArguments): # one object, host form + device form - def __host_args__(self): return host_object # host kernel gets this - def __cuda_args__(self): return (arr, n, ...) # CUDA kernel gets these, flattened -S = xp.cuda.CudaStruct("S", [("x", "double*"), ("n", "long long"), ("a", "Array2D")]) -S.declaration; S.dtype; S.to_header(path); value = S(x=..., n=..., a=...) -S.verify_layout() # GPU test: compiler layout == S.dtype (also verify_layout("hdr.cuh")) +class Args(xp.kernels.KernelArguments): # one object, host form + device form + def __host_args__(self): + return host_object # host kernel gets this + + def __cuda_args__(self): + return (arr, n, ...) # CUDA kernel gets these, flattened + + +S = xp.cuda.CudaStruct( + "S", [("x", "double*"), ("n", "long long"), ("a", "Array2D")] +) +S.declaration +S.dtype +S.to_header(path) +value = S(x=..., n=..., a=...) +S.verify_layout() # GPU test: compiler layout == S.dtype (also verify_layout("hdr.cuh")) S = xp.cuda.CudaStruct.from_signature(Cls.__init__, "S", int_type="long long") xp.cuda.write_cuda_header("args.cuh", [S1, S2]) -class A(xp.cuda.CudaStructArguments): # the struct as a class; A.struct is the CudaStruct + +class A( + xp.cuda.CudaStructArguments +): # the struct as a class; A.struct is the CudaStruct struct_name = "A" fields = (("x", "double*"), ("n", "int")) + def __init__(self, x): self.x, self.n = x, x.shape[0] - self.pack() # repacks itself when a field changes; copies repack + self.pack() # repacks itself when a field changes; copies repack + + xp.cuda.CudaKernel(S.declaration + src, "k", structs=[S]) xp.kernels.resolve_host_args(args, kwargs) ``` @@ -304,9 +332,11 @@ itself at the next launch; make its fields properties to follow an owner's array ```python m = xp.memory.DeviceMirror(host_numpy_array) # TypeError if not numpy.ndarray -m.device # CuPy copy (lazy) on CuPy; the host array itself on NumPy -m.zero(); m.to_device(); m.to_host() # to_host copies in place; no-ops on NumPy -m.rebind(new_host_array) # after the owner reallocates +m.device # CuPy copy (lazy) on CuPy; the host array itself on NumPy +m.zero() +m.to_device() +m.to_host() # to_host copies in place; no-ops on NumPy +m.rebind(new_host_array) # after the owner reallocates ``` Debugging: @@ -372,13 +402,14 @@ Script entry point: ```python import cunumpy as xp + def main(use_gpu: bool): xp.set_backend("cupy" if use_gpu else "numpy") print("backend:", xp.get_backend()) - data = xp.to_cunumpy(load_host_data()) # one transfer in + data = xp.to_cunumpy(load_host_data()) # one transfer in for _ in range(n_steps): - data = update(data) # no transfers here - save(xp.to_numpy(data)) # one transfer out + data = update(data) # no transfers here + save(xp.to_numpy(data)) # one transfer out ``` Kernel pair: @@ -393,12 +424,15 @@ void scale(double* x, double a, long long n) { } """ + def scale_host(x: "float[:]", a: float, n: int): for i in range(n): x[i] *= a -scale = xp.kernels.Kernel(scale_host, xp.cuda.CudaKernel(SRC, "scale"), - host_options={"outputs": (0,)}) + +scale = xp.kernels.Kernel( + scale_host, xp.cuda.CudaKernel(SRC, "scale"), host_options={"outputs": (0,)} +) scale(x, 2.0, x.size, n_threads=x.size) ``` @@ -409,6 +443,7 @@ def make_args(backend, seed): x = xp.to_cunumpy(np.random.default_rng(seed).random(1000)) return (x, 2.0, x.size) + def test_scale(): assert_kernels_agree(scale, make_args, n_threads=1000) ``` diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index e372fdb..17c4056 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -14,7 +14,7 @@ rng, xp, ) -from .scipy_backend import scipy +from ._scipy_backend import scipy from .xp import ( as_device_array, assert_same_backend, @@ -48,10 +48,6 @@ } _MOVED.pop("BIT_GENERATORS") # never was at the top level -# Importing cunumpy.rng loads the module cunumpy.random_streams, which would -# hide the deprecated top-level name random_streams (the generator). -globals().pop("random_streams", None) - try: __version__ = version("cunumpy") except PackageNotFoundError: @@ -120,10 +116,13 @@ def require_version(minimum: str) -> None: def __getattr__(name: str): - """Set cunumpy. to cunumpy.xp. (NumPy/CuPy). + """Resolve names that are not in the namespace: cunumpy. -> cunumpy.xp.. - Names moved to a submodule in cunumpy 0.5 (see ``_MOVED``) still resolve, - with a ``DeprecationWarning``. + The public names of the active backend are copied into this namespace (see + `_sync_backend_namespace`), so this only runs for names missing from the + backend's ``__all__``, for ``numpy_backend``/``cupy_backend``, and for the + names moved to a submodule in cunumpy 0.5 (see ``_MOVED``), which still + resolve with a ``DeprecationWarning``. """ if name == "numpy_backend": return xp.numpy_backend @@ -139,3 +138,56 @@ def __getattr__(name: str): ) return getattr(globals()[submodule], name) return getattr(xp.xp, name) + + +# `xp.zeros` must be as fast as `numpy.zeros`. A module-level __getattr__ runs +# only after the normal lookup failed, which costs about 3 us per access, so the +# public names of the active backend module are copied into this namespace, and +# replaced whenever the backend changes. cunumpy's own names and the deprecated +# names of _MOVED (e.g. `fuse`, which CuPy also has) are never overwritten. +_OWN_NAMES = frozenset(globals()) +_backend_names: dict[int, dict[str, object]] = {} # id(module) -> names to copy +_switches: dict[tuple[int, int], tuple[tuple[str, ...], dict[str, object]]] = {} +_synced_module: list[int] = [0] # id of the module whose names are in the namespace + + +def _names_of(module) -> dict[str, object]: + names = _backend_names.get(id(module)) + if names is None: + names = { + name: getattr(module, name) + for name in getattr(module, "__all__", ()) + if not name.startswith("_") + and name not in _OWN_NAMES + and name not in _MOVED + and hasattr(module, name) + } + _backend_names[id(module)] = names + return names + + +def _sync_backend_namespace(module) -> None: + new = _names_of(module) + namespace = globals() + old_id = _synced_module[0] + if old_id not in _backend_names: + namespace.update(new) + else: + # per pair of modules: the names to drop and the names whose value + # changes (NumPy and CuPy share dtypes, constants, ...) + key = (old_id, id(module)) + switch = _switches.get(key) + if switch is None: + old = _backend_names[old_id] + switch = _switches[key] = ( + tuple(old.keys() - new.keys()), + {k: v for k, v in new.items() if old.get(k, new) is not v}, + ) + stale, changed = switch + for name in stale: + namespace.pop(name, None) + namespace.update(changed) + _synced_module[0] = id(module) + + +xp.array_backend.add_listener(_sync_backend_namespace) diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 659f3ea..bf57473 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -17,7 +17,7 @@ from . import petsc as petsc from . import profiling as profiling from . import rng as rng from . import xp as xp -from .scipy_backend import scipy as scipy +from ._scipy_backend import scipy as scipy def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py new file mode 100644 index 0000000..ee3ca00 --- /dev/null +++ b/src/cunumpy/_algorithms.py @@ -0,0 +1,105 @@ +"""Sums per key and sorting by key on either backend (see :mod:`cunumpy.algorithms`).""" + +from __future__ import annotations + +from typing import Any + +import array_api_compat.numpy as np + +from .xp import 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)``. + + 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. + + Returns + ------- + array + Shape ``(n_segments,)`` or ``(n_segments, m)``, dtype of `values` for + floating-point and complex values, ``float64`` otherwise. + """ + 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 + ) + 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 + + +def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: + """Sort `keys` and reorder every array the same way, in one stable argsort. + + The usual first step of a particle code on the GPU: sort the particles by + cell index or Morton key (:func:`cunumpy.algorithms.morton_keys`), then work on + contiguous ranges. The sort is stable, so equal keys keep their order and + the result is reproducible:: + + keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, charges) + + Parameters + ---------- + keys : array, shape (n,) + The sort keys. + *arrays : arrays + Arrays with ``n`` rows, on the backend of `keys`, reordered along + axis 0. + + Returns + ------- + tuple + ``(sorted_keys, order, *sorted_arrays)``: ``order`` (int64) is the + permutation, ``sorted_keys = keys[order]``, and each sorted array is + ``array[order]`` (a new array). + """ + if get_array_backend(keys) == "cupy": + import cupy as xpm # its argsort is a stable radix sort + else: + xpm = np + keys = xpm.asarray(keys) + if keys.ndim != 1: + raise ValueError(f"keys must be 1D, got shape {keys.shape}") + for array in arrays: + if array.shape[:1] != keys.shape: + raise ValueError( + 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 new file mode 100644 index 0000000..7b6fb13 --- /dev/null +++ b/src/cunumpy/_cuda_kernel.py @@ -0,0 +1,2529 @@ +"""CUDA kernels (``cupy.RawKernel``) called like their NumPy/Pyccel counterparts. + +:class:`CudaKernel` wraps a CUDA C kernel so that it can be called with the same +arguments as the host kernel it mirrors: + +* argument objects that implement the :class:`CudaArguments` protocol + (a ``__cuda_args__()`` method) are flattened into their device arrays and + scalars, so an object holding several arrays can be passed as one argument; +* C structs can be passed by value: :class:`CudaStruct` defines the struct once, + generates its C declaration and packs its values; +* the ``__global__`` signature is parsed once, and every call is checked against + it: the number of arguments, the dtype of every array, every struct, and + every scalar. Python scalars are cast to the declared C type; a scalar that + does not fit the declared type (a ``float`` for an ``int``, an integer out of + range, a NumPy scalar that would lose precision) raises instead of reaching + the kernel as a silently wrong value, which is what ``cupy.RawKernel`` would + do; +* arrays are never converted or copied: they must already be C-contiguous + CuPy arrays (build them once with :func:`cunumpy.as_device_array`); +* strided array views: a parameter or struct field of type ``Array2D`` + (from the shipped header ``cunumpy/array_view.cuh``, see + :func:`cuda_include_dir`) takes a 2D CuPy array, contiguous or not, and + receives its pointer, shape and strides, so that kernels index ``a(i, j)`` + like the pyccel kernels they are ported from; +* C++ function templates are instantiated with ``template_args``, and generated + kernels (one source per variant) are compiled once per variant by + :class:`CudaKernelVariants`. + +The launch shape is given at each call, either as the number of threads +(``n_threads``, in 1 to 3 dimensions) or as an explicit ``grid``. + +Argument structs can be generated from the annotations of a Python class +(:meth:`CudaStruct.from_signature`) and written to a header +(:meth:`CudaStruct.to_header`, :func:`write_cuda_header`), so that the Python +class is the one definition of the arguments. + +In debug mode (``debug=True``, ``xp.cuda.set_cuda_debug(True)`` or the environment +variable ``CUNUMPY_CUDA_DEBUG=1``) kernels are compiled with ``-lineinfo`` and +``-DCUNUMPY_BOUNDS_CHECK``, and every launch is synchronized so that an +asynchronous CUDA error is raised, as a ``RuntimeError`` naming the kernel, at +the launch that caused it. + +This module imports CuPy only when a kernel is compiled, so it can be imported +(and signatures parsed) without CuPy. +""" + +from __future__ import annotations + +import ast +import hashlib +import inspect +import math +import os +import re +import sys +import typing +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +from pathlib import Path +from typing import Any, NamedTuple + +import numpy as np + +__all__ = [ + "DEBUG_OPTIONS", + "CudaArguments", + "CudaKernel", + "CudaKernelVariants", + "CudaParameter", + "CudaStruct", + "CudaStructArguments", + "CudaStructValue", + "PyccelStructArguments", + "ctype_of", + "cuda_include_dir", + "cuda_kernel_names", + "include_hash", + "parse_cuda_signature", + "resolve_includes", + "write_cuda_header", +] + +# CUDA limit on the number of threads per block +_MAX_THREADS_PER_BLOCK = 1024 + +# Headers shipped with cunumpy: #include etc. +_CUDA_INCLUDE_DIR = Path(__file__).resolve().parent / "cuda" / "include" +_ARRAY_VIEW_INCLUDE = '#include "cunumpy/array_view.cuh"' + + +def cuda_include_dir() -> str: + """The directory of the CUDA headers shipped with cunumpy. + + :class:`CudaKernel` adds it to the include path automatically, so kernels + can ``#include "cunumpy/array_view.cuh"`` (strided ``Array1D`` + to ``Array4D`` views passed by value) and + ``#include "cunumpy/index.cuh"`` (thread-index and grid-stride macros such + as ``CUNUMPY_THREAD_1D(i, n)``), ``#include "cunumpy/atomic.cuh"`` (atomic + adds) and ``#include "cunumpy/reduce.cuh"`` (warp and block reductions). + Pass it as ``-I`` to other compilers. + """ + return str(_CUDA_INCLUDE_DIR) + + +#: NVRTC options added in debug mode: source line information for +#: ``compute-sanitizer``/``nsys``, and bounds checks in the array views. +#: (``-G`` is not among them: NVRTC does not support it.) +DEBUG_OPTIONS = ("-lineinfo", "-DCUNUMPY_BOUNDS_CHECK") + + +class CudaArguments: + """Base class for objects passed to a :class:`CudaKernel` as one argument. + + A :class:`CudaKernel` replaces every argument that has a ``__cuda_args__()`` + method by the values it returns, in order. Subclassing this class is + optional: any object implementing ``__cuda_args__()`` is flattened. + + Parameters + ---------- + *values + The CUDA kernel arguments this object stands for: CuPy arrays and + scalars, in the order of the kernel signature. + + Examples + -------- + >>> class Particles(CudaArguments): + ... def __init__(self, positions, velocities): + ... self.positions = positions + ... super().__init__(positions, velocities, positions.shape[0]) + >>> kernel(dt, Particles(x, v), n_threads=x.shape[0]) # doctest: +SKIP + """ + + def __init__(self, *values: Any) -> None: + self._cuda_args = tuple(values) + + def __cuda_args__(self) -> tuple[Any, ...]: + """The CUDA kernel arguments this object stands for.""" + return self._cuda_args + + +class CudaParameter(NamedTuple): + """One parameter of a CUDA kernel signature (or one field of a struct). + + Attributes + ---------- + name : str + Parameter name. + ctype : str + Normalized C type without qualifiers or ``*``, e.g. ``"double"`` or + ``"Array2D"``. + dtype : numpy.dtype | None + NumPy dtype of the value (or of the pointed-to elements, or of the + elements of an array view; the structured dtype for a struct); + ``None`` for ``void*``. + pointer : bool + Whether the parameter is a pointer (a device array). + struct : CudaStruct | None + The struct type, for a struct passed by value. + view_ndim : int | None + The number of dimensions, for an array view (``Array1D`` to + ``Array4D``, see :func:`cuda_include_dir`) passed by value. + """ + + name: str + ctype: str + dtype: np.dtype | None + pointer: bool + struct: CudaStruct | None = None + view_ndim: int | None = None + + +# C types (after removing qualifiers) and their NumPy dtypes. ``long`` is 64 bit, +# as on Linux (LP64), the platform CUDA runs on in practice. +_CTYPES = { + "bool": np.bool_, + "char": np.int8, + "signed char": np.int8, + "unsigned char": np.uint8, + "short": np.int16, + "short int": np.int16, + "unsigned short": np.uint16, + "unsigned short int": np.uint16, + "int": np.int32, + "signed": np.int32, + "signed int": np.int32, + "unsigned": np.uint32, + "unsigned int": np.uint32, + "long": np.int64, + "long int": np.int64, + "long long": np.int64, + "long long int": np.int64, + "unsigned long": np.uint64, + "unsigned long int": np.uint64, + "unsigned long long": np.uint64, + "unsigned long long int": np.uint64, + "int8_t": np.int8, + "int16_t": np.int16, + "int32_t": np.int32, + "int64_t": np.int64, + "uint8_t": np.uint8, + "uint16_t": np.uint16, + "uint32_t": np.uint32, + "uint64_t": np.uint64, + "size_t": np.uint64, + "ptrdiff_t": np.int64, + "ssize_t": np.int64, + "float": np.float32, + "double": np.float64, + "complex": np.complex64, + "complex": np.complex128, +} + +# The C type used for each NumPy dtype, see ctype_of() +_CTYPE_OF = { + np.dtype(np.bool_): "bool", + np.dtype(np.int8): "signed char", + np.dtype(np.uint8): "unsigned char", + np.dtype(np.int16): "short", + np.dtype(np.uint16): "unsigned short", + np.dtype(np.int32): "int", + np.dtype(np.uint32): "unsigned int", + np.dtype(np.int64): "long long", + np.dtype(np.uint64): "unsigned long long", + np.dtype(np.float32): "float", + np.dtype(np.float64): "double", + np.dtype(np.complex64): "complex", + np.dtype(np.complex128): "complex", +} + +_QUALIFIERS = {"const", "volatile", "__restrict__", "__restrict", "restrict"} + +_COMPLEX = re.compile(r"(?:(?:thrust|cuda::std)::)?complex\s*<\s*(float|double)\s*>") +# Array1D to Array4D (cunumpy/array_view.cuh), T a scalar type of _CTYPES +_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*\]" +) + + +def _normalize_view(match: re.Match) -> str: + """``Array2D< const double >`` -> ``Array2D``.""" + words = [w for w in match.group(2).split() if w not in _QUALIFIERS] + return f"Array{match.group(1)}D<{' '.join(words)}>" + + +def _view_dtype(ndim: int) -> np.dtype: + """The structured dtype with the C layout of ``ArrayD``.""" + return np.dtype( + [ + ("data", np.uint64), + ("shape", np.int64, (ndim,)), + ("strides", np.int64, (ndim,)), + ], + align=True, + ) + + +def ctype_of(dtype: Any) -> str: + """The C type of a NumPy dtype, e.g. ``ctype_of(np.float64) == "double"``. + + Useful to generate CUDA source or template arguments for a given dtype. + Complex dtypes map to ``complex``/``complex`` (include + ```` in the source). + """ + try: + return _CTYPE_OF[np.dtype(dtype)] + except (KeyError, TypeError): + raise ValueError(f"no C type for dtype {dtype!r}") from None + + +def _strip_comments(source: str) -> str: + source = re.sub(r"/\*.*?\*/", " ", source, flags=re.DOTALL) + return re.sub(r"//[^\n]*", " ", source) + + +# ``#include "name"``: quoted includes are the project's own headers. Angle +# bracket includes are system headers and are not tracked, except in the +# directories given as ``angle_dirs`` (cunumpy's shipped headers). +_INCLUDE = re.compile( + r'^[ \t]*#[ \t]*include[ \t]*(?:"([^"\n]+)"|<([^>\n]+)>)', re.MULTILINE +) + + +def _includes(source: str) -> list[tuple[str, bool]]: + """``(name, quoted)`` of every ``#include`` in `source`, in order.""" + return [ + (quoted or angle, bool(quoted)) + for quoted, angle in _INCLUDE.findall(_strip_comments(source)) + ] + + +def _quoted_includes(source: str) -> list[str]: + return [name for name, quoted in _includes(source) if quoted] + + +def resolve_includes( + source: str, + include_dirs: Iterable[str | Path] = (), + *, + base_dir: str | Path | None = None, + angle_dirs: Iterable[str | Path] = (), +) -> list[Path]: + """The header files a CUDA source includes, recursively. + + Scans `source` (comments removed) for ``#include "name"`` and resolves each + name like NVRTC does: relative to `base_dir` (the directory of the + including file), then in `include_dirs`, then in `angle_dirs`, in order. + Found headers are scanned in turn, relative to their own directory. + Includes in angle brackets (``#include ``) are system headers and + ignored, unless they are found in `angle_dirs`. Includes that cannot be + found are ignored; NVRTC reports them when the kernel is compiled. + + Parameters + ---------- + source : str + CUDA C source code. + include_dirs : Iterable[str | Path] + Directories searched for included files, in order (the ``-I`` options). + base_dir : str | Path | None + Directory of the file `source` was read from, searched first; None if + the source is not from a file. + angle_dirs : Iterable[str | Path] + Directories whose headers are tracked also when included in angle + brackets, searched last; :class:`CudaKernel` passes + :func:`cuda_include_dir`, so that ``#include `` + is tracked. + + Returns + ------- + list[Path] + The resolved header files, each once, in order of first inclusion + (depth first). Empty if the source has no includes to track; the file + system is not touched in that case. + """ + angle = tuple(Path(d) for d in angle_dirs) + dirs = tuple(Path(d) for d in include_dirs) + dirs += tuple(d for d in angle if d not in dirs) + found: list[Path] = [] + seen: set[Path] = set() + + def visit(code: str, directory: Path | None) -> None: + for name, quoted in _includes(code): + if quoted: + candidates = [directory / name] if directory is not None else [] + candidates += [d / name for d in dirs] + else: + candidates = [d / name for d in angle] + for candidate in candidates: + if candidate.is_file(): + path = candidate.resolve() + if path not in seen: + seen.add(path) + found.append(candidate) + visit(path.read_text(errors="replace"), path.parent) + break + + visit(source, None if base_dir is None else Path(base_dir)) + return found + + +def include_hash(paths: Iterable[str | Path]) -> str: + """A short hex digest of the contents of `paths`, in order. + + Only the file contents count, not their locations: moving a header does not + change the hash, editing it does. Used to make CuPy's kernel cache key + depend on the included headers, see :meth:`CudaKernel.compile_options`. + + Parameters + ---------- + paths : Iterable[str | Path] + Files to hash, e.g. from :func:`resolve_includes`. + + Returns + ------- + str + The first 16 hex digits of the SHA-256 digest. + """ + digest = hashlib.sha256() + for path in paths: + content = Path(path).read_bytes() + digest.update(len(content).to_bytes(8, "little")) + digest.update(content) + return digest.hexdigest()[:16] + + +_GLOBAL_FUNCTION = re.compile(r"__global__\s+void\s+([A-Za-z_]\w*)\s*\(") + + +def cuda_kernel_names(source: str) -> list[str]: + """The names of the ``__global__`` functions defined in `source`, in order. + + Comments are ignored. Templates are included; a function declared more + than once (e.g. a forward declaration) is listed once. + + Parameters + ---------- + source : str + CUDA C source code. + + Returns + ------- + list[str] + The kernel names, in the order of their first appearance. + """ + return list(dict.fromkeys(_GLOBAL_FUNCTION.findall(_strip_comments(source)))) + + +def _compile_in_threads( + compilers: Mapping[Hashable, Callable[[], Any]], jobs: int | None +) -> list[Hashable]: + """Run the `compilers` (name -> compile function), `jobs` at a time. + + With ``jobs=1`` they run one after the other in the calling thread; with + ``jobs=None`` as many threads as CPUs are used. All compilers are run even + if one fails; the first exception (in the order of `compilers`) is raised + afterwards. + + Returns + ------- + list + The names whose compiler succeeded, in the order of `compilers`. + """ + if jobs is None: + jobs = os.cpu_count() or 1 + if jobs < 1: + raise ValueError(f"jobs must be positive or None, got {jobs}") + if jobs == 1 or len(compilers) <= 1: + for compile in compilers.values(): + compile() + return list(compilers) + + with ThreadPoolExecutor(max_workers=min(jobs, len(compilers))) as pool: + futures = {name: pool.submit(compile) for name, compile in compilers.items()} + compiled, error = [], None + for name, future in futures.items(): + exc = future.exception() + if exc is None: + compiled.append(name) + elif error is None: + error = exc + if error is not None: + raise error + return compiled + + +def _parse_parameter( + text: str, structs: dict[str, CudaStruct] | None = None +) -> CudaParameter: + text = _COMPLEX.sub(lambda m: f"complex<{m.group(1)}>", text) + text = _VIEW.sub(_normalize_view, text) + tokens = _TOKEN.findall(text) + pointers = sum(1 for t in tokens if t == "*" or t.startswith("[")) + words = [t for t in tokens if t != "*" and not t.startswith("[")] + words = [t for t in words if t not in _QUALIFIERS] + if words[:1] == ["struct"]: + words = words[1:] + if len(words) < 2: + raise ValueError(f"cannot parse the kernel parameter {text.strip()!r}") + name, ctype = words[-1], " ".join(words[:-1]) + + if structs and ctype in structs: + if pointers: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: structs can " + "only be passed by value" + ) + struct = structs[ctype] + return CudaParameter(name, ctype, struct.dtype, False, struct) + view = _VIEW.fullmatch(ctype) + if view is not None: + element = view.group(2) + if pointers or element not in _CTYPES: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: array views " + f"take a scalar element type and are passed by value" + ) + ndim = int(view.group(1)) + return CudaParameter(name, ctype, np.dtype(_CTYPES[element]), False, None, ndim) + if ctype == "void" and pointers == 1: + return CudaParameter(name, ctype, None, True) + if pointers > 1 or ctype not in _CTYPES: + raise ValueError( + f"cannot check the kernel parameter {text.strip()!r}: unsupported type " + f"{ctype + '*' * pointers!r}" + ) + return CudaParameter(name, ctype, np.dtype(_CTYPES[ctype]), pointers == 1) + + +def _split_top_level(text: str) -> list[str]: + """Split at commas that are not inside ``<...>`` (e.g. ``complex``).""" + parts, depth, current = [], 0, [] + for char in text: + if char == "<": + depth += 1 + elif char == ">": + depth -= 1 + elif char == "," and depth == 0: + parts.append("".join(current)) + current = [] + continue + current.append(char) + parts.append("".join(current)) + return parts + + +def _template_arg(value: Any) -> str: + """A template argument as C++ source: a C type for dtypes, else a literal.""" + if isinstance(value, str): + return value + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, np.integer)): + return str(int(value)) + return ctype_of(value) + + +def parse_cuda_signature( + source: str, + name: str, + *, + structs: Iterable[CudaStruct] = (), + template_args: Sequence[Any] | None = None, +) -> tuple[CudaParameter, ...]: + """Parse the parameters of the ``__global__`` function `name` in `source`. + + Parameters + ---------- + source : str + CUDA C source code. + name : str + Name of the ``__global__`` function. + structs : Iterable[CudaStruct] + Struct types that may appear as parameters (passed by value). If the + source defines a struct of the same name, its fields must match. + template_args : Sequence | None + Template arguments, if `name` is a function template: C types (or NumPy + dtypes, see :func:`ctype_of`) for type parameters, integers or bools for + non-type parameters. They are substituted into the parameter list. + + Returns + ------- + tuple[CudaParameter, ...] + The parameters, in order. + + Raises + ------ + ValueError + If there is no such function, a template is used without (the right + number of) `template_args`, a struct definition in the source does not + match its :class:`CudaStruct`, or a parameter has a type that cannot be + checked (e.g. a macro or a pointer to pointer). + """ + code = _strip_comments(source) + structs = {s.name: s for s in structs} + for struct in structs.values(): + struct.check_source(code) + + pattern = r"(?:template\s*<(?P