From 6721871a5244849abc7f8dc0d4177304867f86db Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 27 Oct 2026 15:57:43 +0100 Subject: [PATCH 1/2] Add KernelPackage, source_roots and shared_memory_fits --- CHANGELOG.md | 9 + docs/src/content/docs/kernels/dispatch.md | 37 ++++ docs/src/content/docs/reference/overview.md | 2 +- src/cunumpy/_device.py | 27 +++ src/cunumpy/_kernel_package.py | 202 ++++++++++++++++++++ src/cunumpy/cuda/__init__.py | 2 + src/cunumpy/kernels.py | 7 +- tests/unit/test_kernel_package.py | 78 ++++++++ 8 files changed, 362 insertions(+), 2 deletions(-) create mode 100644 src/cunumpy/_kernel_package.py create mode 100644 tests/unit/test_kernel_package.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ab8d227..9154bb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 but not NumPy no longer share builds. ### Added +- `xp.kernels.KernelPackage(name)`: the `PyccelCache` and the `Kernel.from_folder` options + shared by the kernel folders of one package (`host_suffix="_pyccel"`, `dispatch="arrays"`, + `block_size=256`, `n_threads_from="last_axis"`), with `from_folder(package, **overrides)` + and `metal_kernel(...)` (float32 `MetalKernel`s with zeroed outputs), so packages no longer + keep their own copies of these settings. +- `xp.kernels.source_roots(*packages)`: the source roots of installed packages, for + `include_dirs` of a CUDA kernel that includes another package's header. +- `xp.cuda.shared_memory_fits(num_values, itemsize=8)`: whether a per-block buffer fits + `DEFAULT_SHARED_MEMORY_PER_BLOCK`. - scope-profiler regions around the coarse operations (`Kernel` calls as `cunumpy.kernel.`, CUDA and Pyccel compilation, `sort_by_key`, `SegmentPlan.sum`, `compact_by_mask`, `morton_keys`, the solvers, `exchange`, diff --git a/docs/src/content/docs/kernels/dispatch.md b/docs/src/content/docs/kernels/dispatch.md index 4770592..62a5e33 100644 --- a/docs/src/content/docs/kernels/dispatch.md +++ b/docs/src/content/docs/kernels/dispatch.md @@ -286,6 +286,43 @@ Packages that compile their kernels at install time with the `pyccel` command do not need `compile_host`: the compiled modules are imported like the Python ones. +### One set of options per package + +A package with many kernel folders passes the same options to every +`from_folder()`. A `KernelPackage` holds them, with the package's +`PyccelCache`, so they are declared once: + +```python +# my_sim/kernels/__init__.py +KERNEL_PACKAGE = xp.kernels.KernelPackage("my-sim") + +# my_sim/kernels/push/__init__.py +from my_sim.kernels import KERNEL_PACKAGE + +kernel = KERNEL_PACKAGE.from_folder(__name__) + +# my_sim/kernels/push/push_metal.py +_KERNEL = KERNEL_PACKAGE.metal_kernel("push", SOURCE, ["x", "v"], ["x_new"]) +``` + +The defaults are `host_suffix="_pyccel"`, `dispatch="arrays"` (see below), +`block_size=256` and `n_threads_from="last_axis"` (the marker count of a +component-major `(ncomp, N)` array). A folder overrides any of them in its own +call, e.g. `KERNEL_PACKAGE.from_folder(__name__, n_threads_from=rows)`; +`KERNEL_PACKAGE.options` are the same options for +`KernelCatalog.from_package()`. `metal_kernel()` builds a `MetalKernel` with +`threadgroup=256`, `float64="cast"` and `init_value=0`. + +A CUDA source that includes another package's header (`#include +"other_pkg/kernels/cell.cuh"`) needs that package's source root as an include +directory; `xp.kernels.source_roots()` finds it without importing the package: + +```python +kernel = KERNEL_PACKAGE.from_folder( + __name__, include_dirs=xp.kernels.source_roots("my_sim", "other_pkg") +) +``` + ## Choosing the kernel by where the arrays are `Kernel` picks the CUDA kernel on the CuPy backend. Real codes also hand host diff --git a/docs/src/content/docs/reference/overview.md b/docs/src/content/docs/reference/overview.md index ffe525f..f084b3a 100644 --- a/docs/src/content/docs/reference/overview.md +++ b/docs/src/content/docs/reference/overview.md @@ -47,7 +47,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | Submodule | Backends | Contents | |---|---|---| | [`cunumpy`](/cunumpy/api/cunumpy/) | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `as_device_array`, `to_host_async`, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` | -| [`cunumpy.kernels`](/cunumpy/api/cunumpy/kernels/) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, `PyccelCache`, host implementations, `as_kernel_array`, `kernel_output`, `kernel_outputs`, `fuse` | +| [`cunumpy.kernels`](/cunumpy/api/cunumpy/kernels/) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, `PyccelCache`, `KernelPackage`, `source_roots`, host implementations, `as_kernel_array`, `kernel_output`, `kernel_outputs`, `fuse` | | [`cunumpy.arguments`](/cunumpy/api/cunumpy/arguments/) | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` | | [`cunumpy.cuda`](/cunumpy/api/cunumpy/cuda/) | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) | | [`cunumpy.rng`](/cunumpy/api/cunumpy/rng/) | both | `random_streams`, `get_rng`, `philox_*` | diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index d567e00..ffa4db2 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -199,6 +199,33 @@ def memory_info() -> tuple[int, int] | None: """ +def shared_memory_fits(num_values: int, itemsize: int = 8) -> bool: + """Return whether `num_values` values fit the default shared memory of a block. + + The limit is :data:`DEFAULT_SHARED_MEMORY_PER_BLOCK`, which every device + provides without an opt-in, so the answer is the same with or without a GPU. + + Parameters + ---------- + num_values : int + Number of values in the per-block buffer. + itemsize : int, optional + Bytes per value. Default 8 (float64). + + Returns + ------- + bool + + Examples + -------- + >>> xp.cuda.shared_memory_fits(6144) + True + >>> xp.cuda.shared_memory_fits(6145) + False + """ + return num_values * itemsize <= DEFAULT_SHARED_MEMORY_PER_BLOCK + + def max_shared_memory_per_block( device: int | None = None, *, diff --git a/src/cunumpy/_kernel_package.py b/src/cunumpy/_kernel_package.py new file mode 100644 index 0000000..32e0810 --- /dev/null +++ b/src/cunumpy/_kernel_package.py @@ -0,0 +1,202 @@ +"""The kernel settings shared by every kernel folder of one package.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from importlib.util import find_spec +from pathlib import Path +from typing import Any + +from cunumpy._dispatch import Kernel +from cunumpy._metal_kernel import MetalKernel +from cunumpy._pyccel_cache import PyccelCache + + +def source_roots(*packages: str) -> tuple[Path, ...]: + """Return the directories containing `packages`, for ``#include "/..."``. + + A CUDA kernel that includes a header of another package passes these as + ``include_dirs``; the packages are found without importing them. + + Parameters + ---------- + *packages : str + Top-level package names. + + Returns + ------- + tuple of Path + The parent directory of each package, in order. + + Raises + ------ + ModuleNotFoundError + If a package is not installed. + + Examples + -------- + >>> (xp.kernels.source_roots("cunumpy")[0] / "cunumpy").is_dir() + True + """ + roots = [] + for package in packages: + spec = find_spec(package) + if spec is None or spec.origin is None: + raise ModuleNotFoundError(package) + roots.append(Path(spec.origin).resolve().parents[1]) + return tuple(roots) + + +class KernelPackage: + """The kernel settings of a package with one folder per kernel. + + A package whose kernels live in ``/kernels//`` + (``_pyccel.py``, ``_numpy.py``, ``_cuda.cu`` and + ``_metal.py``) declares one instance, e.g. in + ``/kernels/__init__.py``, and builds every kernel with it, so the + folders share one build cache and one set of launch options. Nothing is + compiled when a kernel is built: the Pyccel version on its first host call + (by :attr:`cache`), the CUDA version on its first launch. + + Parameters + ---------- + name : str + The package name, for the build cache (see + :class:`~cunumpy.kernels.PyccelCache`). + env_prefix : str, optional + Prefix of the cache's environment variables, passed on to + :class:`~cunumpy.kernels.PyccelCache`. + host_suffix : str, optional + Module name suffix of the host kernels. Default ``"_pyccel"``. + dispatch : {"arrays", "backend"}, optional + How a kernel picks its version. Default ``"arrays"``: CUDA for CuPy + arguments, the host version for NumPy ones, also while CuPy is active. + block_size : int, optional + CUDA threads per block. Default 256. + n_threads_from : str or callable, optional + The CUDA launch size of a call without ``n_threads``. Default + ``"last_axis"``: the last axis of the first array, the marker count of + a component-major ``(ncomp, N)`` or ``(N,)`` marker array. + threadgroup : int, optional + Threads per threadgroup of the Metal kernels. Default 256. + + Examples + -------- + In ``my_sim/kernels/__init__.py``:: + + KERNEL_PACKAGE = xp.kernels.KernelPackage("my-sim") + + in ``my_sim/kernels/push/__init__.py``:: + + from my_sim.kernels import KERNEL_PACKAGE + + kernel = KERNEL_PACKAGE.from_folder(__name__) + + and in ``my_sim/kernels/push/push_metal.py``:: + + _KERNEL = KERNEL_PACKAGE.metal_kernel("push", SOURCE, ["x", "v"], ["x_new"]) + + >>> package = xp.kernels.KernelPackage("my-sim") + >>> package.cache.cache_env + 'MY_SIM_KERNEL_CACHE' + >>> package.options["n_threads_from"] + 'last_axis' + """ + + def __init__( + self, + name: str, + *, + env_prefix: str | None = None, + host_suffix: str = "_pyccel", + dispatch: str = "arrays", + block_size: int = 256, + n_threads_from: str | Callable[[tuple[Any, ...]], int] = "last_axis", + threadgroup: int = 256, + ) -> None: + self.name = name + #: The on-disk build cache of the package's Pyccel kernels. + self.cache = PyccelCache(name, env_prefix=env_prefix) + self.host_suffix = host_suffix + self.dispatch = dispatch + self.block_size = block_size + self.n_threads_from = n_threads_from + self.threadgroup = threadgroup + + def __repr__(self) -> str: + return f"KernelPackage({self.name!r})" + + @property + def options(self) -> dict[str, Any]: + """The options :meth:`from_folder` passes to ``Kernel.from_folder``. + + The same options fit :meth:`KernelCatalog.from_package + `. + """ + return { + "host_suffix": self.host_suffix, + "compile_host": self.cache, + "dispatch": self.dispatch, + "block_size": self.block_size, + "n_threads_from": self.n_threads_from, + } + + def from_folder(self, package: str, **options: Any) -> Kernel: + """Build the kernel of one kernel folder with the package's options. + + Parameters + ---------- + package : str + Dotted name of the kernel folder (``__name__`` in its ``__init__.py``). + **options + Passed on to :meth:`Kernel.from_folder + `, replacing the package's + :attr:`options` of the same name (e.g. ``n_threads_from`` for a + kernel whose first array is not a per-marker array, or + ``include_dirs=xp.kernels.source_roots(...)``). + + Returns + ------- + Kernel + """ + return Kernel.from_folder(package, **{**self.options, **options}) + + def metal_kernel( + self, + name: str, + source: str, + inputs: Sequence[str], + outputs: Sequence[str], + *, + header: str = "", + **options: Any, + ) -> MetalKernel: + """Build the Metal version of a kernel: float32 on the GPU, outputs zeroed. + + The GPU of an Apple-silicon Mac has no float64, so by default float64 + arguments are cast (``float64="cast"``, a relative error of about 1e-5 + against the CPU versions) and the outputs start at zero + (``init_value=0``), as the accumulating kernels need. + + Parameters + ---------- + name, source, inputs, outputs, header + As for :class:`~cunumpy.kernels.MetalKernel`. + **options + Passed on to :class:`~cunumpy.kernels.MetalKernel`, replacing + ``threadgroup=``:attr:`threadgroup`, ``float64="cast"`` and + ``init_value=0``. + + Returns + ------- + MetalKernel + """ + options = { + "threadgroup": self.threadgroup, + "float64": "cast", + "init_value": 0, + **options, + } + return MetalKernel( + source, inputs=inputs, outputs=outputs, name=name, header=header, **options + ) diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index 8281dc7..1e29033 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -42,6 +42,7 @@ set_cuda_debug, set_device, set_device_for_rank, + shared_memory_fits, stream, ) from cunumpy._streams import ( @@ -80,6 +81,7 @@ "set_cuda_debug", "set_device", "set_device_for_rank", + "shared_memory_fits", "stream", "wait_event", ] diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 7fd63d3..177ce60 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -19,7 +19,9 @@ A :class:`PyccelCache` compiles Pyccel kernel modules on first use and keeps the builds on disk, keyed on their source, for later processes: pass it as -``compile_host`` (see :doc:`/kernels/dispatch`). +``compile_host`` (see :doc:`/kernels/dispatch`). A :class:`KernelPackage` holds +the cache and the launch options shared by the kernel folders of one package; +:func:`source_roots` finds the headers of other packages for ``include_dirs``. A :class:`MetalKernel` runs a Metal Shading Language kernel on the GPU of an Apple silicon Mac through MLX (``pip install 'cunumpy[metal]'``), on NumPy @@ -51,6 +53,7 @@ use_device_kernel_implementation, use_host_kernel_implementation, ) +from cunumpy._kernel_package import KernelPackage, source_roots from cunumpy._metal_kernel import MetalKernel, metal_available from cunumpy._pyccel_cache import PyccelCache @@ -63,6 +66,7 @@ "HostImplementations", "Kernel", "KernelCatalog", + "KernelPackage", "MetalKernel", "PyccelCache", "PyccelKernel", @@ -76,6 +80,7 @@ "outputs_from_annotations", "set_device_kernel_implementation", "set_host_kernel_implementation", + "source_roots", "use_device_kernel_implementation", "use_host_kernel_implementation", ] diff --git a/tests/unit/test_kernel_package.py b/tests/unit/test_kernel_package.py new file mode 100644 index 0000000..cb6919c --- /dev/null +++ b/tests/unit/test_kernel_package.py @@ -0,0 +1,78 @@ +"""KernelPackage: the cache and launch options shared by a package's kernel folders.""" + +import sys +import textwrap + +import pytest + +import cunumpy as xp + +HOST = """ +def add_one(x: "float[:]"): + for i in range(x.shape[0]): + x[i] += 1.0 +""" + + +@pytest.fixture +def kernel_folder(tmp_path, monkeypatch): + """A package ``tiny_pkg`` with one kernel folder ``tiny_pkg.add_one``.""" + folder = tmp_path / "tiny_pkg" / "add_one" + folder.mkdir(parents=True) + (tmp_path / "tiny_pkg" / "__init__.py").write_text("") + (folder / "__init__.py").write_text("") + (folder / "add_one_pyccel.py").write_text(textwrap.dedent(HOST)) + monkeypatch.syspath_prepend(str(tmp_path)) + yield "tiny_pkg.add_one" + for name in [name for name in sys.modules if name.startswith("tiny_pkg")]: + del sys.modules[name] + + +def test_options_share_one_cache(): + package = xp.kernels.KernelPackage("my-sim") + options = package.options + assert options["compile_host"] is package.cache + assert package.cache.cache_env == "MY_SIM_KERNEL_CACHE" + assert options["host_suffix"] == "_pyccel" + assert options["dispatch"] == "arrays" + assert options["block_size"] == 256 + assert options["n_threads_from"] == "last_axis" + assert repr(package) == "KernelPackage('my-sim')" + + +def test_env_prefix_reaches_the_cache(): + package = xp.kernels.KernelPackage("my-sim", env_prefix="SIM") + assert package.cache.openmp_env == "SIM_OPENMP" + + +def test_from_folder_runs_the_python_version(kernel_folder): + package = xp.kernels.KernelPackage("tiny-pkg") + kernel = package.from_folder(kernel_folder, missing_cuda="fallback") + x = xp.zeros(3) + with xp.kernels.use_host_kernel_implementation("python"): + kernel(x) + assert x.tolist() == [1.0, 1.0, 1.0] + + +def test_metal_kernel_defaults_and_overrides(): + package = xp.kernels.KernelPackage("my-sim", threadgroup=128) + kernel = package.metal_kernel("scale", "y[i] = 2 * x[i];", ["x"], ["y"]) + assert kernel.name == "scale" + assert kernel.threadgroup == (128, 1, 1) + assert kernel.float64 == "cast" + assert kernel.init_value == 0 + exact = package.metal_kernel("scale", "y[i] = x[i];", ["x"], ["y"], float64="error") + assert exact.float64 == "error" + + +def test_source_roots_find_packages_without_importing(): + (root,) = xp.kernels.source_roots("cunumpy") + assert (root / "cunumpy" / "__init__.py").is_file() + with pytest.raises(ModuleNotFoundError): + xp.kernels.source_roots("no_such_package_here") + + +def test_shared_memory_fits_the_default_block(): + assert xp.cuda.shared_memory_fits(6144) + assert not xp.cuda.shared_memory_fits(6145) + assert xp.cuda.shared_memory_fits(12288, itemsize=4) From a5f22d30f006dc69caf575894b5bf02234e67447 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 27 Oct 2026 16:15:53 +0100 Subject: [PATCH 2/2] Fix numpydoc sections and statement order in KernelPackage --- src/cunumpy/_device.py | 1 + src/cunumpy/_kernel_package.py | 26 +++++++++++++++++++++----- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index ffa4db2..a113df3 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -215,6 +215,7 @@ def shared_memory_fits(num_values: int, itemsize: int = 8) -> bool: Returns ------- bool + True if ``num_values * itemsize`` bytes fit in one block. Examples -------- diff --git a/src/cunumpy/_kernel_package.py b/src/cunumpy/_kernel_package.py index 32e0810..4550bd5 100644 --- a/src/cunumpy/_kernel_package.py +++ b/src/cunumpy/_kernel_package.py @@ -123,15 +123,18 @@ def __init__( self.n_threads_from = n_threads_from self.threadgroup = threadgroup - def __repr__(self) -> str: - return f"KernelPackage({self.name!r})" - @property def options(self) -> dict[str, Any]: """The options :meth:`from_folder` passes to ``Kernel.from_folder``. The same options fit :meth:`KernelCatalog.from_package `. + + Returns + ------- + dict + ``host_suffix``, ``compile_host`` (:attr:`cache`), ``dispatch``, + ``block_size`` and ``n_threads_from``. """ return { "host_suffix": self.host_suffix, @@ -158,6 +161,7 @@ def from_folder(self, package: str, **options: Any) -> Kernel: Returns ------- Kernel + The folder's kernel; nothing is compiled yet. """ return Kernel.from_folder(package, **{**self.options, **options}) @@ -180,8 +184,16 @@ def metal_kernel( Parameters ---------- - name, source, inputs, outputs, header - As for :class:`~cunumpy.kernels.MetalKernel`. + name : str + Name of the kernel; part of the compiled function name. + source : str + Body of the Metal kernel function. + inputs : sequence of str + Names of the input arrays, in call order. + outputs : sequence of str + Names of the output arrays, in the order passed as ``out``. + header : str, optional + Source before the kernel function: includes, defines, helpers. **options Passed on to :class:`~cunumpy.kernels.MetalKernel`, replacing ``threadgroup=``:attr:`threadgroup`, ``float64="cast"`` and @@ -190,6 +202,7 @@ def metal_kernel( Returns ------- MetalKernel + The kernel, compiled by MLX on its first call. """ options = { "threadgroup": self.threadgroup, @@ -200,3 +213,6 @@ def metal_kernel( return MetalKernel( source, inputs=inputs, outputs=outputs, name=name, header=header, **options ) + + def __repr__(self) -> str: + return f"KernelPackage({self.name!r})"