Add KernelPackage, source_roots and shared_memory_fits

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.<name>`, CUDA and Pyccel compilation, `sort_by_key`,
   `SegmentPlan.sum`, `compact_by_mask`, `morton_keys`, the solvers, `exchange`,

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

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_*` |

src/cunumpy/_device.py

@@ -199,6 +199,34 @@ 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
+        True if ``num_values * itemsize`` bytes fit in one block.
+
+    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,
     *,

src/cunumpy/_kernel_package.py

@@ -0,0 +1,218 @@
+"""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 "<package>/..."``.
+
+    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 ``<package>/kernels/<name>/``
+    (``<name>_pyccel.py``, ``<name>_numpy.py``, ``<name>_cuda.cu`` and
+    ``<name>_metal.py``) declares one instance, e.g. in
+    ``<package>/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
+
+    @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
+        <cunumpy.kernels.KernelCatalog.from_package>`.
+
+        Returns
+        -------
+        dict
+            ``host_suffix``, ``compile_host`` (:attr:`cache`), ``dispatch``,
+            ``block_size`` and ``n_threads_from``.
+        """
+        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
+            <cunumpy.kernels.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
+            The folder's kernel; nothing is compiled yet.
+        """
+        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 : 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
+            ``init_value=0``.
+
+        Returns
+        -------
+        MetalKernel
+            The kernel, compiled by MLX on its first call.
+        """
+        options = {
+            "threadgroup": self.threadgroup,
+            "float64": "cast",
+            "init_value": 0,
+            **options,
+        }
+        return MetalKernel(
+            source, inputs=inputs, outputs=outputs, name=name, header=header, **options
+        )
+
+    def __repr__(self) -> str:
+        return f"KernelPackage({self.name!r})"

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",
 ]

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",
 ]