diff --git a/CHANGELOG.md b/CHANGELOG.md index 41b0fb1..eb79b49 100644 --- a/CHANGELOG.md +++ b/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//kernels`, or + `_KERNEL_CACHE`), keyed on its source, the OpenMP setting + (`_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. diff --git a/docs/source/api.md b/docs/source/api.md index fdb68db..292125e 100644 --- a/docs/source/api.md +++ b/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_*` | diff --git a/docs/source/guides/array-ordering.md b/docs/source/guides/array-ordering.md index 0312bff..83ff533 100644 --- a/docs/source/guides/array-ordering.md +++ b/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 diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md index 8ec33cd..cc0170b 100644 --- a/docs/source/kernels/dispatch.md +++ b/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//kernels`, or the directory in + `_KERNEL_CACHE`; if that is not writable, to `/-kernels`; +- modules are built with OpenMP; `_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. + +`` 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 diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 0fea615..5d84a7a 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/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")`; `_numba.py`, `_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) | diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index 44f7339..76dd43b 100644 --- a/src/cunumpy/_kernel.py +++ b/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"} diff --git a/src/cunumpy/_pyccel_cache.py b/src/cunumpy/_pyccel_cache.py new file mode 100644 index 0000000..07cec3a --- /dev/null +++ b/src/cunumpy/_pyccel_cache.py @@ -0,0 +1,305 @@ +"""Compile Pyccel kernel modules once and reuse the builds across processes. + +:class:`PyccelCache` builds a module with the ``pyccel compile`` command line +interface, in a subprocess, once. It keys a build on the module source, the +target language, the OpenMP setting, the Pyccel and Python versions and the +platform, and stores it under a per-package cache directory. Later processes, +and other MPI ranks, import the stored extension directly; only the first run +after an edit compiles. + +cunumpy needs no Pyccel to import this module: Pyccel is looked up when a +module is compiled, and a missing Pyccel is an :class:`ImportError` there, which +the host-kernel dispatch records like any other failed compilation. +""" + +from __future__ import annotations + +import hashlib +import importlib.machinery +import importlib.metadata +import importlib.util +import inspect +import os +import platform +import re +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from types import ModuleType + +#: Bumped when the layout of a cached build changes (it was an ``epyccel`` +#: folder before the ``pyccel compile`` command line was used). Part of the key. +_LAYOUT = "pyccel-cli-1" + +_FALSE = ("0", "false", "no", "off") + + +def _uses_openmp_runtime(source: Path) -> bool: + """Whether the module calls the OpenMP runtime (thread ids, thread counts). + + Such a module only compiles with OpenMP, so it is built with it even when + OpenMP is turned off; ``OMP_NUM_THREADS=1`` then runs it on one thread. + """ + return "pyccel.stdlib.internal.openmp" in source.read_text() + + +def _source_path(module: ModuleType) -> Path: + source = inspect.getsourcefile(module) + if source is None: + raise FileNotFoundError(f"No source file for {module.__name__}") + return Path(source) + + +def _writable(directory: Path) -> bool: + try: + directory.mkdir(parents=True, exist_ok=True) + except OSError: + return False + return os.access(directory, os.W_OK | os.X_OK) + + +def _load_extension(directory: Path, stem: str) -> ModuleType | None: + for suffix in importlib.machinery.EXTENSION_SUFFIXES: + path = directory / f"{stem}{suffix}" + if not path.is_file(): + continue + spec = importlib.util.spec_from_file_location(stem, path) + if spec is None or spec.loader is None: + continue + extension = importlib.util.module_from_spec(spec) + spec.loader.exec_module(extension) + return extension + return None + + +def _pyccel_version() -> str: + try: + return importlib.metadata.version("pyccel") + except importlib.metadata.PackageNotFoundError as error: + raise ImportError("pyccel is not installed") from error + + +def _pyccel_command() -> list[str]: + """Return the ``pyccel`` executable of the running environment.""" + beside_python = Path(sys.executable).parent / "pyccel" + found = str(beside_python) if beside_python.is_file() else shutil.which("pyccel") + if found is None: + raise ImportError("the pyccel command line tool is not installed") + return [found] + + +class PyccelCache: + """Compile Pyccel modules on first use and keep the builds on disk. + + One instance per package: builds go to ``~/.cache//kernels``, or to + the directory in the environment variable ``_KERNEL_CACHE``. If that + directory is not writable (a CI container whose home belongs to another + user), builds go to ``/-kernels`` instead, so the code does not + fall back to slow uncompiled kernels. + + Modules are built with OpenMP (``--openmp``), so loops marked + ``#$ omp parallel for`` run on all cores; limit them with + ``OMP_NUM_THREADS`` (under MPI, the cores per rank). ``_OPENMP=0`` + builds them serial, except modules that call the OpenMP runtime. The setting + is part of the key, so serial and OpenMP builds live side by side. + + Builds are published with an atomic rename: processes that compile the + same module at the same time (MPI ranks on a cold cache) never see a + partial build, and the first one to finish wins. + + The instance is callable, so it can be passed as ``compile_host`` to + :meth:`Kernel.from_folder` and :meth:`KernelCatalog.from_package`. + + Parameters + ---------- + name : str + The package name, used for the cache directories. + env_prefix : str, optional + Prefix of the environment variables. By default `name` in upper case + with ``-`` and ``.`` replaced by ``_`` (``"tree-d-3d"`` gives + ``TREE_D_3D_KERNEL_CACHE`` and ``TREE_D_3D_OPENMP``). + language : str, optional + Pyccel's target language. Default ``"c"``. + + Examples + -------- + >>> cache = xp.kernels.PyccelCache("my-sim") + >>> cache.cache_env, cache.openmp_env + ('MY_SIM_KERNEL_CACHE', 'MY_SIM_OPENMP') + + In the package's kernel folder:: + + kernel = xp.kernels.Kernel.from_folder(__name__, compile_host=cache) + """ + + def __init__( + self, name: str, *, env_prefix: str | None = None, language: str = "c" + ) -> None: + if not name: + raise ValueError("PyccelCache needs a package name") + self.name = name + self.env_prefix = ( + env_prefix + if env_prefix is not None + else re.sub(r"[^0-9A-Za-z]", "_", name).upper() + ) + self.language = language + + def __repr__(self) -> str: + return f"PyccelCache({self.name!r}, env_prefix={self.env_prefix!r})" + + @property + def cache_env(self) -> str: + """Name of the environment variable that moves the cache directory.""" + return f"{self.env_prefix}_KERNEL_CACHE" + + @property + def openmp_env(self) -> str: + """Name of the environment variable that turns OpenMP off (``0``).""" + return f"{self.env_prefix}_OPENMP" + + def cache_root(self) -> Path: + """Directory holding the builds (``~/.cache//kernels`` by default).""" + default = Path.home() / ".cache" / self.name / "kernels" + return Path(os.environ.get(self.cache_env, default)) + + def fallback_root(self) -> Path: + """Directory for builds when :meth:`cache_root` is not writable.""" + return Path(tempfile.gettempdir()) / f"{self.name}-kernels" + + def openmp_enabled(self) -> bool: + """Whether modules are built with OpenMP (``_OPENMP=0`` turns it off).""" + return os.environ.get(self.openmp_env, "1").lower() not in _FALSE + + def _openmp_for(self, module: ModuleType) -> bool: + return self.openmp_enabled() or _uses_openmp_runtime(_source_path(module)) + + def build_key(self, module: ModuleType, language: str | None = None) -> str: + """Return the key of `module`'s build, which changes with anything that changes it. + + The module's own source, the language, the OpenMP setting, the Pyccel + and Python versions and the platform. Modules the kernel imports are not + part of it: edit the kernel (or clear the cache) after changing one. + + Parameters + ---------- + module : module + The module to build. + language : str, optional + Pyccel's target language; by default the instance's. + + Returns + ------- + str + 20 hexadecimal digits, the suffix of the build's directory name. + + Raises + ------ + ImportError + If Pyccel is not installed. + """ + digest = hashlib.sha256() + for part in ( + inspect.getsource(module), + language or self.language, + _LAYOUT, + f"openmp={self._openmp_for(module)}", + _pyccel_version(), + sys.version, + platform.platform(), + ): + digest.update(part.encode()) + return digest.hexdigest()[:20] + + def _build(self, module: ModuleType, language: str, staging: Path) -> Path: + """Run ``pyccel compile`` on a copy of the module's file; return the build folder.""" + source = _source_path(module) + stem = module.__name__.rsplit(".", 1)[-1] + sources = staging / "src" + sources.mkdir() + shutil.copy(source, sources / f"{stem}.py") + output = staging / "build" + command = [ + *_pyccel_command(), + "compile", + str(sources / f"{stem}.py"), + "--language", + language, + "--output", + str(output), + ] + if self._openmp_for(module): + command.append("--openmp") + # Pyccel writes lock files (``.lock_acquisition.lock``, ``stc.lock``) to + # the directory it runs in: run it in the staging folder, removed after + # the build, so they never land in the caller's working directory. + result = subprocess.run( + command, capture_output=True, text=True, check=False, cwd=staging + ) + if result.returncode != 0: + raise RuntimeError( + f"{' '.join(command)} failed ({result.returncode}):\n" + f"{result.stdout}{result.stderr}" + ) + return output + + def compile(self, module: ModuleType, *, language: str | None = None) -> ModuleType: + """Return the compiled form of `module`, building it on first use. + + Parameters + ---------- + module : module + A module Pyccel can compile (its functions keep their names). + language : str, optional + Overrides the instance's language for this build. + + Returns + ------- + module + The compiled extension module. + + Raises + ------ + ImportError + If Pyccel is not installed. + RuntimeError + With the output of ``pyccel compile`` if the build fails; callers + decide on a fallback. + PermissionError + If neither the cache nor the fallback directory is writable. + """ + language = language or self.language + stem = module.__name__.rsplit(".", 1)[-1] + name = f"{stem}-{self.build_key(module, language)}" + roots = [self.cache_root(), self.fallback_root()] + for root in roots: + if (root / name).is_dir(): + extension = _load_extension(root / name, stem) + if extension is not None: + return extension + target = next((root / name for root in roots if _writable(root)), None) + if target is None: + raise PermissionError(f"No writable kernel cache among {roots}") + # absolute: Pyccel runs inside it, so a relative cache path must not move + staging = Path( + tempfile.mkdtemp(prefix=f".{stem}-", dir=target.parent) + ).absolute() + try: + build = self._build(module, language, staging) + try: + # Atomic, so concurrent builders never see a partial directory. + os.rename(build, target) + except OSError: + pass # Another process published the same build first. + finally: + shutil.rmtree(staging, ignore_errors=True) + extension = _load_extension(target, stem) + if extension is None: + raise ImportError(f"Compiled module for {module.__name__} not found") + return extension + + def __call__(self, module: ModuleType) -> ModuleType: + """:meth:`compile`, so the instance can be a ``compile_host``.""" + return self.compile(module) diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 58b0eb3..7fd63d3 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -12,10 +12,15 @@ or :func:`use_host_kernel_implementation` (one of ``HOST_IMPLEMENTATIONS``); device dispatch can require CUDA with :func:`set_device_kernel_implementation` or :func:`use_device_kernel_implementation` (``DEVICE_IMPLEMENTATIONS``). -:func:`as_kernel_array` and :func:`kernel_output` bring the arguments of a -kernel to the side of its main array. :func:`fuse` turns an elementwise +:func:`as_kernel_array` and :func:`kernel_output` (several outputs at once: +:func:`kernel_outputs`) bring the arguments of a kernel to the side of its +main array. :func:`fuse` turns an elementwise function into one CuPy kernel. See :doc:`/kernels/overview`. +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`). + 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 float32 arrays. @@ -39,6 +44,7 @@ get_device_kernel_implementation, get_host_kernel_implementation, kernel_output, + kernel_outputs, outputs_from_annotations, set_device_kernel_implementation, set_host_kernel_implementation, @@ -46,6 +52,7 @@ use_host_kernel_implementation, ) from cunumpy._metal_kernel import MetalKernel, metal_available +from cunumpy._pyccel_cache import PyccelCache __all__ = [ "DEVICE_IMPLEMENTATIONS", @@ -57,12 +64,14 @@ "Kernel", "KernelCatalog", "MetalKernel", + "PyccelCache", "PyccelKernel", "as_kernel_array", "fuse", "get_device_kernel_implementation", "get_host_kernel_implementation", "kernel_output", + "kernel_outputs", "metal_available", "outputs_from_annotations", "set_device_kernel_implementation", diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index bec4c58..ebabf44 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -658,6 +658,38 @@ def test_strided_without_dtype_keeps_any_dtype(): assert xp.kernels.as_kernel_array(view, like=np.zeros(1), strided=True) is view +def test_kernel_outputs_yields_one_buffer_per_output(): + grid = np.zeros(3) + direct = np.zeros(4) + strided = np.zeros((2, 8))[:, :5] + converted = np.zeros(4, dtype=np.float32) + with xp.kernels.kernel_outputs( + direct, strided, converted, like=grid, dtype=float, strided=True + ) as (a, b, c): + assert a is direct and b is strided # written directly + assert c is not converted and c.dtype == np.float64 + a[:] = 1.0 + b[:] = 2.0 + c[:] = 3.5 + np.testing.assert_array_equal(direct, 1.0) + np.testing.assert_array_equal(strided, 2.0) + np.testing.assert_array_equal(converted, np.full(4, 3.5, dtype=np.float32)) + + +def test_kernel_outputs_keeps_the_outputs_after_an_error(): + out = np.zeros(4, dtype=np.float32) + outputs = xp.kernels.kernel_outputs(out, like=np.zeros(1), dtype=float) + with pytest.raises(RuntimeError), outputs as (buffer,): + buffer[:] = 1.0 + raise RuntimeError("kernel failed") + np.testing.assert_array_equal(out, 0.0) + + +def test_kernel_outputs_without_outputs_yields_an_empty_tuple(): + with xp.kernels.kernel_outputs(like=np.zeros(1)) as buffers: + assert buffers == () + + def test_kernel_output_writes_into_its_target(): grid = np.zeros(3) out = np.zeros(4) @@ -691,3 +723,16 @@ def test_kernel_arrays_follow_a_device_grid(): assert xp.is_gpu(buffer) buffer[...] = 2.0 assert host_out.tolist() == [2.0] * 4 + + +def test_one_column_selected_by_a_mask_gets_c_order_strides(): + # nodes[:, mask] with one True: shape (2, 1), strides (8, 16); NumPy calls + # it C-contiguous, but Pyccel rejects the length-1 axis's larger stride + nodes = np.zeros((2, 5, 4)) + mask = np.zeros((5, 4), dtype=bool) + mask[2, 1] = True + view = nodes[:, mask] + assert view.flags.c_contiguous and view.strides == (8, 16) + converted = xp.kernels.as_kernel_array(view, like=view, dtype=float) + assert converted.strides == (8, 8) + assert np.shares_memory(converted, view) # a view, not a copy diff --git a/tests/unit/test_pyccel_cache.py b/tests/unit/test_pyccel_cache.py new file mode 100644 index 0000000..ce74dfb --- /dev/null +++ b/tests/unit/test_pyccel_cache.py @@ -0,0 +1,228 @@ +"""PyccelCache: compile once, keep the build on disk, reuse it in later processes.""" + +import importlib.metadata +import importlib.util +import os +import stat +import sys +import tempfile +from types import SimpleNamespace + +import pytest + +import cunumpy as xp +from cunumpy import _pyccel_cache + +SOURCE = """ +def add_one(x: float) -> float: + return x + 1.0 +""" + + +@pytest.fixture +def make_module(tmp_path): + """Write `source` as an importable module and return the module.""" + created = [] + + def make(source=SOURCE, name="tiny_kernel"): + folder = tmp_path / f"src{len(created)}" + folder.mkdir() + path = folder / f"{name}.py" + path.write_text(source) + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + created.append(name) + return module + + yield make + for name in created: + sys.modules.pop(name, None) + + +@pytest.fixture +def cache(tmp_path, monkeypatch): + cache = xp.kernels.PyccelCache("cunumpy-test") + monkeypatch.setenv(cache.cache_env, str(tmp_path / "cache")) + # serial builds: these tests are about the cache, and a runner may have no + # OpenMP runtime (macOS without Homebrew's libomp) + monkeypatch.setenv(cache.openmp_env, "0") + return cache + + +def test_environment_variables_follow_the_package_name(): + cache = xp.kernels.PyccelCache("tree-d-3d") + assert (cache.cache_env, cache.openmp_env) == ( + "TREE_D_3D_KERNEL_CACHE", + "TREE_D_3D_OPENMP", + ) + custom = xp.kernels.PyccelCache("my.sim", env_prefix="SIM") + assert custom.cache_env == "SIM_KERNEL_CACHE" + assert xp.kernels.PyccelCache("my.sim").env_prefix == "MY_SIM" + with pytest.raises(ValueError): + xp.kernels.PyccelCache("") + + +def test_cache_directories(monkeypatch, tmp_path): + cache = xp.kernels.PyccelCache("pkg") + monkeypatch.delenv(cache.cache_env, raising=False) + assert cache.cache_root().parts[-3:] == (".cache", "pkg", "kernels") + monkeypatch.setenv(cache.cache_env, str(tmp_path / "elsewhere")) + assert cache.cache_root() == tmp_path / "elsewhere" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path / "tmp")) + assert cache.fallback_root() == tmp_path / "tmp" / "pkg-kernels" + + +@pytest.mark.parametrize( + ("value", "enabled"), [("1", True), ("0", False), ("off", False)] +) +def test_openmp_setting(monkeypatch, value, enabled): + cache = xp.kernels.PyccelCache("pkg") + monkeypatch.setenv(cache.openmp_env, value) + assert cache.openmp_enabled() is enabled + + +def test_key_changes_with_the_source_and_the_openmp_setting( + make_module, cache, monkeypatch +): + pytest.importorskip("pyccel") + module = make_module() + key = cache.build_key(module) + assert cache.build_key(module) == key + assert cache.build_key(module, "fortran") != key + monkeypatch.setenv(cache.openmp_env, "1") + assert cache.build_key(module) != key + monkeypatch.setenv(cache.openmp_env, "0") + edited = make_module(SOURCE.replace("1.0", "2.0"), name="tiny_kernel_edited") + edited.__name__ = module.__name__ # same name, other source + assert cache.build_key(edited) != key + + +def test_a_module_calling_the_openmp_runtime_keeps_openmp( + make_module, cache, monkeypatch +): + pytest.importorskip("pyccel") + module = make_module( + "from pyccel.stdlib.internal.openmp import omp_get_num_threads\n" + ) + monkeypatch.setenv(cache.openmp_env, "0") + assert cache._openmp_for(module) + + +def test_compiles_once_and_reuses_the_build(make_module, cache, monkeypatch): + pytest.importorskip("pyccel") + module = make_module() + compiled = cache.compile(module) + assert compiled.add_one(1.0) == 2.0 + assert compiled is not module + (build,) = [p for p in cache.cache_root().iterdir() if not p.name.startswith(".")] + assert build.name == f"tiny_kernel-{cache.build_key(module)}" + + def no_build(*args, **kwargs): + raise AssertionError("the cached build should have been loaded") + + # A new instance, as in a later process, loads the stored build. + later = xp.kernels.PyccelCache("cunumpy-test") + monkeypatch.setattr(xp.kernels.PyccelCache, "_build", no_build) + assert later(module).add_one(2.0) == 3.0 + + +def test_no_lock_files_are_left_in_the_working_directory( + make_module, cache, tmp_path, monkeypatch +): + # Pyccel writes .lock_acquisition.lock (and stc.lock for kernels that use + # its STC library) to the directory it runs in + pytest.importorskip("pyccel") + work = tmp_path / "work" + work.mkdir() + monkeypatch.chdir(work) + module = make_module() + assert cache.compile(module).add_one(1.0) == 2.0 + assert list(work.iterdir()) == [] + + +@pytest.mark.parametrize(("openmp", "flag"), [("1", True), ("0", False)]) +def test_pyccel_runs_in_the_staging_folder_with_the_openmp_flag( + make_module, cache, monkeypatch, openmp, flag +): + # no compiler needed: the command is captured instead of run + pytest.importorskip("pyccel") + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return SimpleNamespace(returncode=1, stdout="", stderr="captured") + + monkeypatch.setattr(_pyccel_cache.subprocess, "run", fake_run) + monkeypatch.setenv(cache.openmp_env, openmp) + with pytest.raises(RuntimeError, match="captured"): + cache.compile(make_module()) + ((command, kwargs),) = calls + assert ("--openmp" in command) is flag + staging = kwargs["cwd"] + assert staging.is_absolute() and staging.parent == cache.cache_root() + assert not staging.exists() # removed after the failed build + + +def test_a_relative_cache_directory_works(make_module, tmp_path, monkeypatch): + pytest.importorskip("pyccel") + cache = xp.kernels.PyccelCache("cunumpy-test") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv(cache.cache_env, "relative/cache") + monkeypatch.setenv(cache.openmp_env, "0") + module = make_module() + assert cache.compile(module).add_one(1.0) == 2.0 + assert any((tmp_path / "relative" / "cache").iterdir()) + + +@pytest.mark.skipif( + not hasattr(os, "geteuid") or os.geteuid() == 0, + reason="root ignores directory permissions", +) +def test_unwritable_cache_builds_in_the_temp_dir(tmp_path, monkeypatch, make_module): + pytest.importorskip("pyccel") + cache = xp.kernels.PyccelCache("cunumpy-test") + readonly = tmp_path / "readonly" + readonly.mkdir() + readonly.chmod(stat.S_IRUSR | stat.S_IXUSR) + monkeypatch.setenv(cache.cache_env, str(readonly / "kernels")) + monkeypatch.setenv(cache.openmp_env, "0") + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path / "tmp")) + (tmp_path / "tmp").mkdir() + module = make_module() + try: + assert cache.compile(module).add_one(1.0) == 2.0 + assert any((tmp_path / "tmp" / "cunumpy-test-kernels").iterdir()) + assert cache.compile(module).add_one(2.0) == 3.0 # reuses the fallback build + finally: + readonly.chmod(stat.S_IRWXU) + + +def test_a_failed_build_raises_with_the_compiler_output(make_module, cache): + pytest.importorskip("pyccel") + module = make_module( + "def broken(x: float) -> float:\n return x.missing_attribute\n" + ) + with pytest.raises(RuntimeError, match="pyccel"): + cache.compile(module) + assert not [p for p in cache.cache_root().iterdir() if not p.name.startswith(".")] + + +def test_missing_pyccel_is_an_import_error(make_module, cache, monkeypatch): + module = make_module() + + def not_installed(name): + raise importlib.metadata.PackageNotFoundError(name) + + monkeypatch.setattr(_pyccel_cache.importlib.metadata, "version", not_installed) + with pytest.raises(ImportError): + cache.compile(module) + + +def test_the_cache_is_a_host_compiler(make_module, cache): + pytest.importorskip("pyccel") + module = make_module() + kernel = xp.kernels.CompiledHostKernel(module, "add_one", cache) + assert kernel.compiled + assert kernel(4.0) == 5.0 diff --git a/tests/unit/test_pyccel_kernel.py b/tests/unit/test_pyccel_kernel.py index bf9b92a..43ccfa3 100644 --- a/tests/unit/test_pyccel_kernel.py +++ b/tests/unit/test_pyccel_kernel.py @@ -35,8 +35,10 @@ def _skip_without_cupy(): def kernels(tmp_path_factory): """Compile `pyccel_kernels.py` with pyccel and return the compiled module. - The source is copied into a temporary directory first so that pyccel's - build artefacts (`__pyccel__/`) never land in the repository. + The source is copied into a temporary directory, and pyccel runs there, so + that its build artefacts (`__pyccel__/`) and the lock files it writes to the + working directory (`.lock_acquisition.lock`, `stc.lock`) never land in the + repository. """ require_pyccel = os.environ.get("CUNUMPY_REQUIRE_PYCCEL", "").lower() in { "1", @@ -61,7 +63,9 @@ def kernels(tmp_path_factory): spec.loader.exec_module(module) try: - return pyccel.epyccel(module, language="c") + with pytest.MonkeyPatch.context() as patch: + patch.chdir(build_dir) + return pyccel.epyccel(module, language="c") except Exception as exc: if require_pyccel: raise