Cache pyccel kernels

CHANGELOG.md

@@ -8,6 +8,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- `xp.kernels.kernel_outputs(*outs, like=..., dtype=..., strided=...)`: the
+  buffers of several kernel outputs in one block, each written back into its
+  array (as `kernel_output` does) when the block ends without an error.
+- `xp.kernels.PyccelCache(name)`: compiles a Pyccel kernel module once with
+  `pyccel compile` and keeps the build on disk (`~/.cache/<name>/kernels`, or
+  `<NAME>_KERNEL_CACHE`), keyed on its source, the OpenMP setting
+  (`<NAME>_OPENMP`), the Pyccel and Python versions and the platform; a
+  temporary directory if the cache is not writable, and an atomic publish for
+  concurrent builders. Callable, so it can be the `compile_host` of
+  `Kernel.from_folder` and `KernelCatalog.from_package`.
 - Supported CuPy versions: 14.0 or newer (`xp.MIN_CUPY_VERSION`,
   `xp.cupy_version_supported`); an older CuPy warns when the CuPy backend is
   first checked, and `backend_info()` reports the minimum. The GPU CI runs the
@@ -45,6 +55,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   warnings as errors. Docstring style: `docs/source/contributing.md`.
 
 ### Fixed
+- `as_kernel_array` gives a length-1 axis inside a longer one exactly the
+  C-order stride: `nodes[:, mask]` with a single selected column (shape
+  `(2, 1)`, strides `(8, 16)`) is C-contiguous to NumPy but was rejected by
+  Pyccel's order check.
 - Splines and other objects of the fake `cupyx.scipy` can be copied with
   `copy.copy` and `copy.deepcopy` (also inside a copied container), as real
   CuPy splines can.

docs/source/api.md

@@ -44,7 +44,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not
 | Submodule | Backends | Contents |
 |---|---|---|
 | [`cunumpy`](reference/top-level.md) | 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`](reference/kernels.md) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` |
+| [`cunumpy.kernels`](reference/kernels.md) | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, `PyccelCache`, host implementations, `as_kernel_array`, `kernel_output`, `kernel_outputs`, `fuse` |
 | [`cunumpy.arguments`](reference/arguments.md) | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` |
 | [`cunumpy.cuda`](reference/cuda.md) | 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`](reference/rng.md) | both | `random_streams`, `get_rng`, `philox_*` |

docs/source/guides/array-ordering.md

@@ -129,6 +129,7 @@ Some convenience helpers deliberately require C order:
 | `xp.as_device_array` | Returns a C-contiguous device array, copying if needed. |
 | `xp.kernels.as_kernel_array` | Returns a C-contiguous array on the side of `like`; with `strided=True`, also a NumPy array in C order with gaps, unchanged. |
 | `xp.kernels.kernel_output` | Yields a C-contiguous working buffer (with `strided=True`, also a NumPy output in C order with gaps) and copies updates back into the original output when needed. |
+| `xp.kernels.kernel_outputs` | `kernel_output` for several outputs in one block: yields a tuple of buffers. |
 
 "C order with gaps" means positive strides, each at least the extent of the next
 axis: `a[:, :n]`, `a[::2]` or `a[:, ::2]` of a C-contiguous `a`. Pyccel's

docs/source/kernels/dispatch.md

@@ -223,14 +223,11 @@ version when that fails (no Pyccel, no compiler):
 
 ```python
 import cunumpy as xp
-import pyccel
 
 from my_sim.kernels.numpy_versions import NUMPY_VERSIONS  # {"push": push_numpy, ...}
 
-
-def compile_kernels(module):
-    return pyccel.epyccel(module, language="c")  # add a cache in real code
-
+# builds once with `pyccel compile`, kept in ~/.cache/my-sim/kernels for later runs
+compile_kernels = xp.kernels.PyccelCache("my-sim")
 
 catalog = xp.kernels.KernelCatalog.from_package(
     __name__,
@@ -253,9 +250,30 @@ host implementations" above); `catalog["push"].host_kernel.kernel.available("pyc
 reports whether the compiled version builds, and `catalog["push"].selected()`
 which version runs. To test the path of a machine without Pyccel, run the code
 inside `with xp.kernels.use_host_kernel_implementation("numpy"):` (or set
-`CUNUMPY_HOST_KERNEL_IMPLEMENTATION=numpy` for a whole run). Note that `epyccel` compiles again on every call; a
-code that compiles at run time usually keeps the builds in an on-disk cache
-keyed on the module source, so that only the first run after an edit compiles.
+`CUNUMPY_HOST_KERNEL_IMPLEMENTATION=numpy` for a whole run).
+
+`epyccel` compiles again on every call, in every process. `PyccelCache(name)`
+compiles a module once with the `pyccel compile` command and keeps the build
+on disk, keyed on the module source, the OpenMP setting, the Pyccel and Python
+versions and the platform, so only the first run after an edit compiles:
+
+- builds go to `~/.cache/<name>/kernels`, or the directory in
+  `<NAME>_KERNEL_CACHE`; if that is not writable, to `<tmp>/<name>-kernels`;
+- modules are built with OpenMP; `<NAME>_OPENMP=0` builds them serial, except
+  modules that call the OpenMP runtime;
+- builds are published with an atomic rename, so MPI ranks compiling the same
+  module on a cold cache never load a partial build;
+- the key covers the module's own source only: after editing a module the
+  kernel imports, edit the kernel or clear the cache.
+
+`<NAME>` is the name in upper case with `-` and `.` replaced by `_` (pass
+`env_prefix=` to choose it). To build every kernel before a run, for example
+before starting MPI ranks, load the compiled version of each:
+
+```python
+for name in catalog:
+    catalog[name].host_kernel.kernel.get("pyccel")
+```
 
 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
@@ -295,6 +313,15 @@ with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer:
     catalog["gather"](convert(positions), convert(field), buffer, n_threads=n)
 ```
 
+For a kernel that writes several arrays, `xp.kernels.kernel_outputs(*outs, like=...,
+dtype=...)` does the same for all of them in one block, which is cheaper at every
+step than nested `kernel_output` blocks:
+
+```python
+with xp.kernels.kernel_outputs(pos, vel, like=field, dtype=float) as (p, v):
+    catalog["push"](p, v, convert(field), dt)
+```
+
 Pyccel host kernels also take arrays in C order with gaps, such as the stored
 markers `storage[:, :n]` of a component-major buffer with spare capacity. Pass
 `strided=True` to both helpers to hand those over without a copy on the host

src/cunumpy/LLM_GUIDE.md

@@ -83,7 +83,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository.
 | host kernels compiled at first call (your compile function), NumPy fallback | `from_package(..., host_suffix="_pyccel", compile_host=my_compile, host_fallback={...})` -> `xp.kernels.CompiledHostKernel` |
 | host arrays reach kernels while CuPy is active | `Kernel(..., dispatch="arrays")` / `from_package(..., dispatch="arrays")`: CUDA only for device arguments |
 | one kernel folder declares its kernel in its own `__init__.py` | `kernel = xp.kernels.Kernel.from_folder(__name__, host_suffix="_pyccel", compile_host=..., dispatch="arrays")`; `<name>_numba.py`, `<name>_numpy.py` in the folder are further host implementations |
-| bring a `dispatch="arrays"` kernel's arguments to the side of the main array | `xp.kernels.as_kernel_array(a, like=grid, dtype=float)`; outputs: `with xp.kernels.kernel_output(out, like=grid, dtype=float) as buf:`; `strided=True` passes NumPy views in C order with gaps (`a[:, :n]`) to Pyccel hosts without a copy |
+| bring a `dispatch="arrays"` kernel's arguments to the side of the main array | `xp.kernels.as_kernel_array(a, like=grid, dtype=float)`; outputs: `with xp.kernels.kernel_output(out, like=grid, dtype=float) as buf:` (several: `with xp.kernels.kernel_outputs(a, b, like=grid, dtype=float) as (buf_a, buf_b):`); `strided=True` passes NumPy views in C order with gaps (`a[:, :n]`) to Pyccel hosts without a copy |
 | choose the host implementation (pyccel/numba/numpy/python) | `xp.kernels.set_host_kernel_implementation("numpy")`, `with xp.kernels.use_host_kernel_implementation(...)`, `CUNUMPY_HOST_KERNEL_IMPLEMENTATION=numpy`; default: first available of pyccel, numba, numpy; `kernel.selected()` |
 | require CUDA for device kernel dispatch | `xp.kernels.set_device_kernel_implementation("cuda")`, `get_device_kernel_implementation()`, `with xp.kernels.use_device_kernel_implementation(...)`, `CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION=cuda`; default `None` preserves `missing_cuda` policy; explicit CUDA rejects host fallback |
 | check host and CUDA kernels take the same parameters | `catalog.check_signatures()` (in a unit test) |

src/cunumpy/_kernel.py

@@ -950,7 +950,16 @@ def _with_unit_axis_strides(array: np.ndarray) -> np.ndarray:
     extent = array.itemsize
     for axis in reversed(range(array.ndim)):
         if array.shape[axis] == 1:
-            strides[axis] = max(strides[axis], extent)
+            # Any stride is valid for a length-1 axis, but Pyccel's C-order check
+            # wants each stride at least the extent of the axes inside it. Inside
+            # a longer axis, use exactly the C stride, so that the outer stride
+            # (e.g. 8 in a (2, 1) array with strides (8, 16)) still covers it;
+            # a leading length-1 axis may keep a larger stride (a row of a
+            # wider array, ``storage[1:2, :n]``).
+            if any(length > 1 for length in array.shape[:axis]):
+                strides[axis] = extent
+            else:
+                strides[axis] = max(strides[axis], extent)
         else:
             extent = strides[axis] * array.shape[axis]
     if tuple(strides) == array.strides:
@@ -1052,6 +1061,18 @@ def _same_memory(buffer: Any, out: Any) -> bool:
     )
 
 
+def _copy_back(buffer: Any, out: Any) -> None:
+    """Write `buffer` into `out`, on its own side, unless they share their memory."""
+    if _same_memory(buffer, out):
+        return
+    if is_gpu(out) and not is_gpu(buffer):
+        out[...] = to_cupy(buffer)
+    elif not is_gpu(out) and is_gpu(buffer):
+        out[...] = to_numpy(buffer)
+    else:
+        out[...] = buffer
+
+
 class _KernelOutput:
     """The context manager of :func:`kernel_output` (a class: cheaper than a generator)."""
 
@@ -1070,16 +1091,8 @@ def __exit__(
         exc: BaseException | None,
         traceback: TracebackType | None,
     ) -> None:
-        if exc_type is not None:
-            return
-        buffer, out = self._buffer, self._out
-        if not _same_memory(buffer, out):
-            if is_gpu(out) and not is_gpu(buffer):
-                out[...] = to_cupy(buffer)
-            elif not is_gpu(out) and is_gpu(buffer):
-                out[...] = to_numpy(buffer)
-            else:
-                out[...] = buffer
+        if exc_type is None:
+            _copy_back(self._buffer, self._out)
 
 
 def kernel_output(
@@ -1118,6 +1131,73 @@ def kernel_output(
     return _KernelOutput(out, as_kernel_array(out, like, dtype, strided=strided))
 
 
+class _KernelOutputs:
+    """The context manager of :func:`kernel_outputs`."""
+
+    __slots__ = ("_buffers", "_outs")
+
+    def __init__(self, outs: tuple[Any, ...], buffers: tuple[Any, ...]) -> None:
+        self._outs = outs
+        self._buffers = buffers
+
+    def __enter__(self) -> tuple[Any, ...]:
+        return self._buffers
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if exc_type is None:
+            for buffer, out in zip(self._buffers, self._outs, strict=True):
+                _copy_back(buffer, out)
+
+
+def kernel_outputs(
+    *outs: Any, like: Any, dtype: Any = None, strided: bool = False
+) -> _KernelOutputs:
+    """Provide the output buffers of several arrays in one block.
+
+    As :func:`~cunumpy.kernels.kernel_output` for each of `outs`, with one
+    context manager for all: cheaper for a kernel that writes several arrays,
+    called at every step. Each buffer is its array itself if
+    :func:`~cunumpy.kernels.as_kernel_array` takes it unchanged, else a
+    converted copy written back into the array when the block ends without an
+    error.
+
+    Parameters
+    ----------
+    *outs : array
+        The arrays the kernel should write.
+    like : array
+        A CuPy array for a device kernel, anything else for a host kernel.
+    dtype : dtype, optional
+        Required dtype of every buffer; by default any.
+    strided : bool, optional
+        As for :func:`~cunumpy.kernels.as_kernel_array`. Default False.
+
+    Returns
+    -------
+    context manager
+        Yields the tuple of buffers to pass to the kernel, in the order of
+        `outs`.
+
+    Examples
+    --------
+    >>> x = np.zeros(2)
+    >>> counts = np.zeros(2, dtype=np.float32)
+    >>> with xp.kernels.kernel_outputs(x, counts, like=x, dtype=float) as (a, b):
+    ...     a[:] = 1.0  # x itself
+    ...     b[:] = 2.5  # a float64 copy, written back into counts
+    >>> x, counts
+    (array([1., 1.]), array([2.5, 2.5], dtype=float32))
+    """
+    return _KernelOutputs(
+        outs, tuple(as_kernel_array(out, like, dtype, strided=strided) for out in outs)
+    )
+
+
 _SCALAR_ANNOTATIONS = {"int", "float", "bool", "complex", "str"}