diff --git a/CHANGELOG.md b/CHANGELOG.md index 495aba5..37aa43c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -117,6 +117,9 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Changed +- The kernels are built with cunumpy's `KernelPackage` + (`simgeometry.kernels.KERNEL_PACKAGE`): `kernels/_build.py`, `kernels/_metal.py` and + `kernels/compile.py` are gone, and the build cache is still `SIMGEOMETRY_KERNEL_CACHE`. - `__version__` is read from the installed package metadata (`importlib.metadata`), not hardcoded, so it can no longer disagree with `pyproject.toml` (the 0.1.1 wheel reported 0.1.0). - One out-of-range rule for owners: per axis, positions below the lower bound belong to the diff --git a/src/simgeometry/kernels/__init__.py b/src/simgeometry/kernels/__init__.py index 4843dea..6699051 100644 --- a/src/simgeometry/kernels/__init__.py +++ b/src/simgeometry/kernels/__init__.py @@ -1 +1,23 @@ -"""Compiled containment kernels for geometric objects.""" +"""simgeometry's compiled kernels, one folder per kernel. + +Every kernel folder builds its kernel with :data:`KERNEL_PACKAGE`:: + + kernel = KERNEL_PACKAGE.from_folder(__name__) + +with the defaults of :class:`cunumpy.kernels.KernelPackage`: the Pyccel version is +compiled on its first host call and kept under ``SIMGEOMETRY_KERNEL_CACHE`` (default +``~/.cache/simgeometry/kernels``; ``SIMGEOMETRY_OPENMP=0`` builds serial), the CUDA version runs +for CuPy arguments and the host version for NumPy ones, one CUDA thread per entry of +the last axis of the first array. Importing a kernel needs neither Pyccel nor CuPy. + +The Metal versions (``/_metal.py``) wrap a kernel built by +``KERNEL_PACKAGE.metal_kernel`` in a function with the arguments of the Pyccel +version. The Apple GPU has no float64, so they compute in float32: expect a relative +error of about 1e-5 against the CPU versions. Import one only where +:func:`cunumpy.kernels.metal_available` is true. +""" + +import cunumpy as xp + +#: The build cache and launch options of simgeometry's kernels. +KERNEL_PACKAGE = xp.kernels.KernelPackage("simgeometry") diff --git a/src/simgeometry/kernels/_build.py b/src/simgeometry/kernels/_build.py deleted file mode 100644 index f703878..0000000 --- a/src/simgeometry/kernels/_build.py +++ /dev/null @@ -1,53 +0,0 @@ -"""The options every kernel folder passes to :meth:`cunumpy.kernels.Kernel.from_folder`. - -The ``__init__.py`` of a kernel folder ``simgeometry/kernels//`` declares its -kernel with these:: - - kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) - -which finds ``_pyccel.py``, ``_numpy.py`` and ``_cuda.cu`` in -the folder. -Nothing is compiled here: the Pyccel version is compiled on its first host call -(cached on disk by :func:`~simgeometry.kernels.compile.compile_cached`), the CUDA -version on its first launch. Importing this module needs neither Pyccel nor CuPy. -""" - -from __future__ import annotations - -from typing import Any - -#: CUDA threads per block of every kernel. -BLOCK_SIZE = 256 - -#: Default dynamic shared memory per block of a CUDA device, in bytes. -SHARED_MEMORY_PER_BLOCK = 48 * 1024 - - -def shared_memory_fits(num_values: int, itemsize: int = 8) -> bool: - """Whether ``num_values`` values fit the default shared memory of a block.""" - return num_values * itemsize <= SHARED_MEMORY_PER_BLOCK - - -def _compile(module: Any) -> Any: - from simgeometry.kernels.compile import compile_cached - - return compile_cached(module) - - -def _marker_count(args: tuple[Any, ...]) -> int: - """Use the last axis of the first array as the CUDA thread count.""" - return int(args[0].shape[-1]) - - -#: ``_pyccel.py`` compiled with Pyccel on first use; CUDA for CuPy -#: arguments, the host version for NumPy ones (also while CuPy is active); -#: without ``n_threads`` one CUDA thread per entry of the last axis of the -#: first array (the marker count of a component-major ``(ncomp, N)`` or -#: ``(N,)`` marker array). CUDA sources include cunumpy's array-view header. -KERNEL_OPTIONS: dict[str, Any] = { - "host_suffix": "_pyccel", - "compile_host": _compile, - "dispatch": "arrays", - "block_size": BLOCK_SIZE, - "n_threads_from": _marker_count, -} diff --git a/src/simgeometry/kernels/_metal.py b/src/simgeometry/kernels/_metal.py deleted file mode 100644 index 10360db..0000000 --- a/src/simgeometry/kernels/_metal.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Helpers of the Metal versions (``/_metal.py``) of the kernels. - -A Metal version is the body of a Metal Shading Language kernel run by -:class:`cunumpy.kernels.MetalKernel` on the GPU of an Apple-silicon Mac, wrapped in -a function with the arguments of the Pyccel version: NumPy arrays in, results -written into the output arrays. The GPU has no float64, so the kernels compute in -float32 (``float64="cast"``): expect a relative error of about 1e-5 against the -CPU versions. Import a version only where :func:`metal_available` is true. -""" - -from __future__ import annotations - -import cunumpy as xp - -#: Threads per threadgroup of every Metal kernel. -THREADGROUP = 256 - - -def metal_available() -> bool: - """Whether the Metal versions can run here (MLX installed, Metal GPU present).""" - return bool(xp.kernels.metal_available()) - - -def make_kernel( - name: str, - source: str, - inputs: list[str], - outputs: list[str], - header: str = "", -) -> xp.kernels.MetalKernel: - """A :class:`~cunumpy.kernels.MetalKernel` computing in float32, outputs zeroed.""" - return xp.kernels.MetalKernel( - source, - inputs=inputs, - outputs=outputs, - name=name, - header=header, - threadgroup=THREADGROUP, - float64="cast", - init_value=0, - ) diff --git a/src/simgeometry/kernels/absorb_polygons_2d/__init__.py b/src/simgeometry/kernels/absorb_polygons_2d/__init__.py index c506694..c14fb3c 100644 --- a/src/simgeometry/kernels/absorb_polygons_2d/__init__.py +++ b/src/simgeometry/kernels/absorb_polygons_2d/__init__.py @@ -6,15 +6,11 @@ version runs one thread per entry of ``rows`` (the third argument). """ -import cunumpy as xp - -from simgeometry.kernels._build import KERNEL_OPTIONS +from simgeometry.kernels import KERNEL_PACKAGE def _row_count(args: tuple) -> int: return int(args[2].shape[0]) -kernel = xp.kernels.Kernel.from_folder( - __name__, **{**KERNEL_OPTIONS, "n_threads_from": _row_count} -) +kernel = KERNEL_PACKAGE.from_folder(__name__, n_threads_from=_row_count) diff --git a/src/simgeometry/kernels/compile.py b/src/simgeometry/kernels/compile.py deleted file mode 100644 index 8f4f16c..0000000 --- a/src/simgeometry/kernels/compile.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Compile simgeometry's Pyccel kernel modules once and reuse the builds across processes. - -The cache is :class:`cunumpy.kernels.PyccelCache`: a kernel module is built with -``pyccel compile`` on its first use, keyed on its source, the Pyccel and Python -versions and the platform, and stored under ``SIMGEOMETRY_KERNEL_CACHE`` (default -``~/.cache/simgeometry/kernels``, or ``/simgeometry-kernels`` if that is not writable). -Later processes import the stored extension directly. - -Kernels are built with OpenMP, so the loops marked ``#$ omp parallel for`` run on -all cores (limit them with ``OMP_NUM_THREADS``; under MPI set it to the cores per -rank). ``SIMGEOMETRY_OPENMP=0`` builds them serial, except the kernels that call the OpenMP -runtime; the setting is part of the cache key. -""" - -from __future__ import annotations - -import cunumpy as xp - -#: The on-disk build cache of simgeometry's kernels. -KERNEL_CACHE = xp.kernels.PyccelCache("simgeometry") - -#: Return the compiled form of a kernel module, building it on first use -#: (:meth:`cunumpy.kernels.PyccelCache.compile`). -compile_cached = KERNEL_CACHE.compile diff --git a/src/simgeometry/kernels/contains_ball/__init__.py b/src/simgeometry/kernels/contains_ball/__init__.py index 3c05d88..077bd3b 100644 --- a/src/simgeometry/kernels/contains_ball/__init__.py +++ b/src/simgeometry/kernels/contains_ball/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from simgeometry.kernels import KERNEL_PACKAGE -from simgeometry.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/simgeometry/kernels/contains_ball/contains_ball_metal.py b/src/simgeometry/kernels/contains_ball/contains_ball_metal.py index 5ae163d..8e66989 100644 --- a/src/simgeometry/kernels/contains_ball/contains_ball_metal.py +++ b/src/simgeometry/kernels/contains_ball/contains_ball_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``contains_ball`` kernel. ``contains_ball_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` is ``(ndim, N)``); see :mod:`simgeometry.kernels._metal`. The test is +arguments. One thread per marker (``positions`` is ``(ndim, N)``); see :mod:`simgeometry.kernels`. The test is done in float32, so a marker within about 1e-7 (relative) of the boundary may be classified differently from the CPU version. """ @@ -11,7 +11,7 @@ import numpy as np from simgeometry._types import Array -from simgeometry.kernels._metal import make_kernel +from simgeometry.kernels import KERNEL_PACKAGE _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -28,7 +28,7 @@ inside_flag[marker] = inside ? 1 : 0; """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "simgeometry_contains_ball", _SOURCE, inputs=["positions", "midpoint", "params"], diff --git a/src/simgeometry/kernels/contains_box/__init__.py b/src/simgeometry/kernels/contains_box/__init__.py index 1bddf2c..22a2e93 100644 --- a/src/simgeometry/kernels/contains_box/__init__.py +++ b/src/simgeometry/kernels/contains_box/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from simgeometry.kernels import KERNEL_PACKAGE -from simgeometry.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/simgeometry/kernels/contains_box/contains_box_metal.py b/src/simgeometry/kernels/contains_box/contains_box_metal.py index f250f86..7ec3906 100644 --- a/src/simgeometry/kernels/contains_box/contains_box_metal.py +++ b/src/simgeometry/kernels/contains_box/contains_box_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``contains_box`` kernel. ``contains_box_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` is ``(ndim, N)``); see :mod:`simgeometry.kernels._metal`. The test is +arguments. One thread per marker (``positions`` is ``(ndim, N)``); see :mod:`simgeometry.kernels`. The test is done in float32, so a marker within about 1e-7 (relative) of the boundary may be classified differently from the CPU version. """ @@ -11,7 +11,7 @@ import numpy as np from simgeometry._types import Array -from simgeometry.kernels._metal import make_kernel +from simgeometry.kernels import KERNEL_PACKAGE _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -27,7 +27,7 @@ inside_flag[marker] = inside ? 1 : 0; """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "simgeometry_contains_box", _SOURCE, inputs=["positions", "lower", "upper", "flags"], diff --git a/src/simgeometry/kernels/contains_polygon_2d/__init__.py b/src/simgeometry/kernels/contains_polygon_2d/__init__.py index aaeab5c..e96fa48 100644 --- a/src/simgeometry/kernels/contains_polygon_2d/__init__.py +++ b/src/simgeometry/kernels/contains_polygon_2d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from simgeometry.kernels import KERNEL_PACKAGE -from simgeometry.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/simgeometry/kernels/contains_polygon_2d/contains_polygon_2d_metal.py b/src/simgeometry/kernels/contains_polygon_2d/contains_polygon_2d_metal.py index 993d289..d6e9d15 100644 --- a/src/simgeometry/kernels/contains_polygon_2d/contains_polygon_2d_metal.py +++ b/src/simgeometry/kernels/contains_polygon_2d/contains_polygon_2d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``contains_polygon_2d`` kernel. ``contains_polygon_2d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` is ``(2, N)``); see :mod:`simgeometry.kernels._metal`. The test is +arguments. One thread per marker (``positions`` is ``(2, N)``); see :mod:`simgeometry.kernels`. The test is done in float32, so a marker within about 1e-7 (relative) of the boundary may be classified differently from the CPU version. """ @@ -11,7 +11,7 @@ import numpy as np from simgeometry._types import Array -from simgeometry.kernels._metal import make_kernel +from simgeometry.kernels import KERNEL_PACKAGE _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -36,7 +36,7 @@ inside_flag[marker] = inside ? 1 : 0; """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "simgeometry_contains_polygon_2d", _SOURCE, inputs=["positions", "vertices", "flags"], diff --git a/uv.lock b/uv.lock index 04213a9..ca15887 100644 --- a/uv.lock +++ b/uv.lock @@ -1114,13 +1114,14 @@ wheels = [ [[package]] name = "cunumpy" version = "0.7.1" -source = { git = "https://github.com/max-models/cunumpy.git?branch=devel#9fb918b130b08a2c159005997186efdb5edacb4f" } +source = { git = "https://github.com/max-models/cunumpy.git?branch=devel#92b38a5a0e00d7ccf03aa39849d6c7eaee946859" } dependencies = [ { name = "array-api-compat" }, { name = "maybempi" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "scope-profiler" }, ] [[package]]