Build the kernels with cunumpy's KernelPackage

CHANGELOG.md

@@ -5,6 +5,13 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ## [Unreleased]
 
+### Changed
+
+- The kernels are built with cunumpy's `KernelPackage` (`mpifield.kernels.KERNEL_PACKAGE`):
+  `kernels/_build.py`, `kernels/_metal.py` and `kernels/compile.py` are gone, and the build
+  cache is still `MPIFIELD_KERNEL_CACHE`. `shared_memory_fits` is
+  `xp.cuda.shared_memory_fits`; the catalog no longer exports it.
+
 ### Added
 
 - `Field`, moved out of picnic (`picnic.field`): a field on the nodes of a uniform 1D, 2D or

src/mpifield/kernels/__init__.py

@@ -1 +1,23 @@
-"""Compiled linear gather and scatter kernels (see :mod:`mpifield.kernels.catalog`)."""
+"""mpifield's compiled kernels, one folder per kernel (see :mod:`mpifield.kernels.catalog`).
+
+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 ``MPIFIELD_KERNEL_CACHE`` (default
+``~/.cache/mpifield/kernels``; ``MPIFIELD_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 (``<name>/<name>_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 mpifield's kernels.
+KERNEL_PACKAGE = xp.kernels.KernelPackage("mpifield")

src/mpifield/kernels/_metal_grid.py

@@ -6,7 +6,7 @@
 ``quadratic_cell.cuh`` for the quadratic ones). A scatter accumulates
 into a zeroed float32 grid with float atomics (no 64-bit atomics on the Apple GPU),
 which the host adds to the caller's grid in float64: the summation order, hence the
-last bits, differ from the CPU. See :mod:`mpifield.kernels._metal`.
+last bits, differ from the CPU. See :mod:`mpifield.kernels`.
 """
 
 from __future__ import annotations
@@ -15,12 +15,11 @@
 import numpy as np
 
 from mpifield._types import Array
-from mpifield.kernels._metal import THREADGROUP, make_kernel
+from mpifield.kernels import KERNEL_PACKAGE
 
 __all__ = [
     "CELL_HEADER",
     "QUADRATIC_HEADER",
-    "make_kernel",
     "make_scatter_kernel",
     "relative",
     "scatter",
@@ -94,16 +93,8 @@ def make_scatter_kernel(
     name: str, source: str, inputs: list[str], header: str = ""
 ) -> xp.kernels.MetalKernel:
     """A kernel adding into the zeroed float32 output ``delta`` with float atomics."""
-    return xp.kernels.MetalKernel(
-        source,
-        inputs=inputs,
-        outputs=["delta"],
-        name=name,
-        header=header,
-        threadgroup=THREADGROUP,
-        float64="cast",
-        atomic_outputs=True,
-        init_value=0,
+    return KERNEL_PACKAGE.metal_kernel(
+        name, source, inputs, ["delta"], header=header, atomic_outputs=True
     )
 
 

src/mpifield/kernels/catalog.py

@@ -2,7 +2,7 @@
 
 Every kernel ``<name>`` lives in its own folder::
 
-    mpifield/kernels/<name>/__init__.py        kernel = xp.kernels.Kernel.from_folder(...)
+    mpifield/kernels/<name>/__init__.py        kernel = KERNEL_PACKAGE.from_folder(__name__)
     mpifield/kernels/<name>/<name>_pyccel.py   def <name>(...)               CPU, compiled with Pyccel
     mpifield/kernels/<name>/<name>_cuda.cu     __global__ void <name>(...)   GPU, compiled with NVRTC
     mpifield/kernels/<name>/<name>_numpy.py    def <name>(...)               CPU, when Pyccel is unavailable
@@ -17,10 +17,10 @@
 runs the CUDA version when the arrays live on the GPU (CuPy) and the compiled Pyccel
 version otherwise (``dispatch="arrays"``). The kernel object is a
 :class:`cunumpy.kernels.Kernel`, built in the folder's ``__init__.py`` by
-:meth:`cunumpy.kernels.Kernel.from_folder` with the options of
-:mod:`mpifield.kernels._build`. Both versions are compiled on first use: importing
-this module needs neither Pyccel nor CuPy. CUDA sources include mpifield's header
-relative to the source root: ``#include "mpifield/kernels/linear_cell.cuh"``.
+``KERNEL_PACKAGE.from_folder`` (see :mod:`mpifield.kernels`). Both versions are
+compiled on first use: importing this module needs neither Pyccel nor CuPy. CUDA
+sources include mpifield's header relative to the source root:
+``#include "mpifield/kernels/linear_cell.cuh"``.
 
 The kernels are the linear (``*_linear_*``) and quadratic (``*_quadratic_*``) deposits
 and gathers of :mod:`mpifield.weighting`. The quadratic ones, and the 1D and 2D linear
@@ -39,23 +39,9 @@
 from collections.abc import Callable
 
 import cunumpy as xp
+from cunumpy.kernels import metal_available
 
-from mpifield.kernels._build import (
-    BLOCK_SIZE,
-    SHARED_MEMORY_PER_BLOCK,
-    shared_memory_fits,
-)
-from mpifield.kernels._metal import metal_available
-
-__all__ = [
-    "BLOCK_SIZE",
-    "KERNELS",
-    "KERNEL_NAMES",
-    "METAL_KERNEL_NAMES",
-    "SHARED_MEMORY_PER_BLOCK",
-    "metal_kernel",
-    "shared_memory_fits",
-]
+__all__ = ["KERNELS", "KERNEL_NAMES", "METAL_KERNEL_NAMES", "metal_kernel"]
 
 #: The kernels, by name.
 KERNEL_NAMES = (

src/mpifield/kernels/gather_linear_1d/__init__.py

@@ -5,8 +5,6 @@
 (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
 """
 
-import cunumpy as xp
+from mpifield.kernels import KERNEL_PACKAGE
 
-from mpifield.kernels._build import KERNEL_OPTIONS
-
-kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)
+kernel = KERNEL_PACKAGE.from_folder(__name__)

src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py

@@ -9,7 +9,8 @@
 import numpy as np
 
 from mpifield._types import Array
-from mpifield.kernels._metal_grid import CELL_HEADER, make_kernel, relative
+from mpifield.kernels import KERNEL_PACKAGE
+from mpifield.kernels._metal_grid import CELL_HEADER, relative
 
 _SOURCE = """
     uint marker = thread_position_in_grid.x;
@@ -25,7 +26,7 @@
     }
 """
 
-_KERNEL = make_kernel(
+_KERNEL = KERNEL_PACKAGE.metal_kernel(
     "mpifield_gather_linear_1d",
     _SOURCE,
     inputs=["rel", "field", "grid"],

src/mpifield/kernels/gather_linear_2d/__init__.py

@@ -5,8 +5,6 @@
 (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
 """
 
-import cunumpy as xp
+from mpifield.kernels import KERNEL_PACKAGE
 
-from mpifield.kernels._build import KERNEL_OPTIONS
-
-kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)
+kernel = KERNEL_PACKAGE.from_folder(__name__)

src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py

@@ -7,7 +7,8 @@
 from __future__ import annotations
 
 from mpifield._types import Array
-from mpifield.kernels._metal_grid import CELL_HEADER, make_kernel, relative
+from mpifield.kernels import KERNEL_PACKAGE
+from mpifield.kernels._metal_grid import CELL_HEADER, relative
 
 _SOURCE = """
     uint marker = thread_position_in_grid.x;
@@ -31,7 +32,7 @@
     }
 """
 
-_KERNEL = make_kernel(
+_KERNEL = KERNEL_PACKAGE.metal_kernel(
     "mpifield_gather_linear_2d",
     _SOURCE,
     inputs=["rel", "field", "grid"],

src/mpifield/kernels/gather_linear_3d/__init__.py

@@ -5,8 +5,6 @@
 (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
 """
 
-import cunumpy as xp
+from mpifield.kernels import KERNEL_PACKAGE
 
-from mpifield.kernels._build import KERNEL_OPTIONS
-
-kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)
+kernel = KERNEL_PACKAGE.from_folder(__name__)

src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py

@@ -7,7 +7,8 @@
 from __future__ import annotations
 
 from mpifield._types import Array
-from mpifield.kernels._metal_grid import CELL_HEADER, make_kernel, relative
+from mpifield.kernels import KERNEL_PACKAGE
+from mpifield.kernels._metal_grid import CELL_HEADER, relative
 
 _SOURCE = """
     uint marker = thread_position_in_grid.x;
@@ -34,7 +35,7 @@
     }
 """
 
-_KERNEL = make_kernel(
+_KERNEL = KERNEL_PACKAGE.metal_kernel(
     "mpifield_gather_linear_3d",
     _SOURCE,
     inputs=["rel", "field", "grid"],

src/mpifield/kernels/gather_quadratic_1d/__init__.py

@@ -5,8 +5,6 @@
 (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
 """
 
-import cunumpy as xp
+from mpifield.kernels import KERNEL_PACKAGE
 
-from mpifield.kernels._build import KERNEL_OPTIONS
-
-kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)
+kernel = KERNEL_PACKAGE.from_folder(__name__)

src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py

@@ -8,7 +8,8 @@
 from __future__ import annotations
 
 from mpifield._types import Array
-from mpifield.kernels._metal_grid import QUADRATIC_HEADER, make_kernel, relative
+from mpifield.kernels import KERNEL_PACKAGE
+from mpifield.kernels._metal_grid import QUADRATIC_HEADER, relative
 
 _SOURCE = """
     uint marker = thread_position_in_grid.x;
@@ -26,7 +27,7 @@
     }
 """
 
-_KERNEL = make_kernel(
+_KERNEL = KERNEL_PACKAGE.metal_kernel(
     "mpifield_gather_quadratic_1d",
     _SOURCE,
     inputs=["rel", "field", "grid"],

src/mpifield/kernels/gather_quadratic_2d/__init__.py

@@ -5,8 +5,6 @@
 (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
 """
 
-import cunumpy as xp
+from mpifield.kernels import KERNEL_PACKAGE
 
-from mpifield.kernels._build import KERNEL_OPTIONS
-
-kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)
+kernel = KERNEL_PACKAGE.from_folder(__name__)

src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py

@@ -8,7 +8,8 @@
 from __future__ import annotations
 
 from mpifield._types import Array
-from mpifield.kernels._metal_grid import QUADRATIC_HEADER, make_kernel, relative
+from mpifield.kernels import KERNEL_PACKAGE
+from mpifield.kernels._metal_grid import QUADRATIC_HEADER, relative
 
 _SOURCE = """
     uint marker = thread_position_in_grid.x;
@@ -31,7 +32,7 @@
     }
 """
 
-_KERNEL = make_kernel(
+_KERNEL = KERNEL_PACKAGE.metal_kernel(
     "mpifield_gather_quadratic_2d",
     _SOURCE,
     inputs=["rel", "field", "grid"],