From abd4ad740db52a2acae8af8b7234e734a2ded7ee Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:36:10 +0200 Subject: [PATCH 1/2] Added MetalKernel --- .github/workflows/macos.yml | 49 ++++++ CHANGELOG.md | 5 + docs/source/api.md | 47 +++++- pyproject.toml | 1 + src/cunumpy/_metal_kernel.py | 258 ++++++++++++++++++++++++++++++++ src/cunumpy/kernels.py | 7 + tests/unit/test_metal_kernel.py | 145 ++++++++++++++++++ 7 files changed, 511 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/macos.yml create mode 100644 src/cunumpy/_metal_kernel.py create mode 100644 tests/unit/test_metal_kernel.py diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml new file mode 100644 index 0000000..6256349 --- /dev/null +++ b/.github/workflows/macos.yml @@ -0,0 +1,49 @@ +name: Tests (macOS) + +on: + push: + branches: + - main + - devel + pull_request: + branches: + - main + - devel + workflow_dispatch: + +jobs: + build: + # macos-latest is Apple silicon (arm64): checks the NumPy backend, the + # compiled host kernels and the MLX import on the platform of MetalKernel. + runs-on: macos-latest + + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.13"] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: pyproject.toml + + - name: Install project + run: | + pip install --upgrade pip + pip install ".[test-compiled,metal]" + + # Hosted macOS runners are virtual machines and usually have no Metal GPU, + # in which case the MetalKernel launch tests are skipped. This step shows + # which case this run is in. + - name: Report Metal availability + run: | + python -c "import platform, cunumpy as xp; print(platform.machine(), 'metal_available =', xp.kernels.metal_available())" + + - name: Run tests + run: pytest . -rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 8772d2f..778b957 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- `xp.kernels.MetalKernel` runs a Metal Shading Language kernel on the GPU of an + Apple silicon Mac through MLX (`pip install 'cunumpy[metal]'`). It takes and + fills NumPy arrays, is float32 only (`float64="cast"` computes float64 data in + float32), and its copies are counted by `xp.profiling.count_transfers()`. + `xp.kernels.metal_available()` tells whether it can run. - `xp.host_call`, `xp.evaluate_on_host` and `xp.setup_on_host` run host-only code (SciPy splines, file readers, external libraries) with arguments of either backend: device arrays are copied to the host, the call runs on the NumPy diff --git a/docs/source/api.md b/docs/source/api.md index bc79c13..fba1154 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -40,7 +40,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | Submodule | Backends | Contents | |---|---|---| | `cunumpy` | both | NumPy/CuPy namespace, backend selection, array inspection and conversion, `synchronize`, `host_call`, `evaluate_on_host`, `setup_on_host`, `scipy`, `require_version` | -| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | +| `cunumpy.kernels` | both | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `CudaKernelVariants`, `MetalKernel`, `metal_available`, host implementations, `as_kernel_array`, `kernel_output`, `fuse` | | `cunumpy.arguments` | CUDA only | `CudaArguments`, `CudaStruct`, `CudaStructArguments`, `CudaStructValue`, `write_cuda_header` | | `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` | both | `random_streams`, `get_rng`, `philox_*` | @@ -857,6 +857,51 @@ lists) are converted back using `is_array`; dictionaries in return values are not recursively converted. On the NumPy path, the original return value and normal Python mutation and exception behavior are preserved. +## `kernels.MetalKernel` + +A Metal Shading Language kernel for the GPU of an Apple silicon Mac, run with +[MLX](https://github.com/ml-explore/mlx) (`pip install 'cunumpy[metal]'`). It +takes NumPy arrays and fills the output arrays you pass, so no backend switch +is needed. + +```python +import numpy as np +import cunumpy as xp + +scale = xp.kernels.MetalKernel( + "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];", + inputs=["x", "a"], + outputs=["y"], +) +x = np.arange(8, dtype=np.float32) +y = np.empty_like(x) +scale(x, 2.0, out=y) +``` + +`MetalKernel(source, inputs, outputs, *, name="cunumpy_kernel", header="", +threadgroup=256, float64="error", atomic_outputs=False, init_value=None)` + +- `source` is the body of the kernel function. MLX generates the signature: each + name in `inputs` and `outputs` is a pointer to the flat, row-major data of that + array, so `x[i]` is the flat index. `thread_position_in_grid` and the other + Metal attributes used in the body are added automatically. `header` goes before + the function (includes, defines, helper functions). +- Calling the kernel: `kernel(*inputs, out=array_or_arrays, n_threads=None, + template=None)`. `n_threads` is the total thread count (default: first axis of + the first output). `template` gives compile-time constants, e.g. + `template={"NSTEPS": 200}`. It returns the output array, or a tuple of them. +- Outputs are uninitialized: write every element, pass the old array as an input + too if the kernel reads it, or set `init_value`. +- **float32 only.** The Apple GPU has no float64: a float64 input or output + raises `TypeError`. With `float64="cast"` float64 data is computed in float32 + (a push over 200 steps agreed with float64 to about 3e-5). +- Every call copies the inputs to MLX arrays and the results back, counted as + `to_device` and `to_host` transfers by `count_transfers()`. On a 4M-particle + push these copies were about 7 ms next to an 11.5 ms kernel. +- `xp.kernels.metal_available()` is True if MLX is installed and a Metal GPU is + present. Without them, calling a `MetalKernel` raises `ImportError` or + `RuntimeError`. + ## `kernels.CudaKernel` ### Constructor diff --git a/pyproject.toml b/pyproject.toml index 7dbd7ee..8346503 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ optional-dependencies.docs = [ "sphinx", "sphinx-book-theme", ] +optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] optional-dependencies.test = [ "coverage", "pytest", "scipy" ] optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ] urls."Source" = "https://github.com/max-models/cunumpy" diff --git a/src/cunumpy/_metal_kernel.py b/src/cunumpy/_metal_kernel.py new file mode 100644 index 0000000..cbf45d6 --- /dev/null +++ b/src/cunumpy/_metal_kernel.py @@ -0,0 +1,258 @@ +"""Metal kernels for the GPU of Apple silicon Macs, run through MLX.""" + +from __future__ import annotations + +import importlib +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _nbytes, _record + +_FLOAT64 = np.dtype(np.float64) +_UNSUPPORTED = (np.dtype(np.complex128),) + + +def _mlx() -> Any: + """Import ``mlx.core``, with an actionable error if it is not usable.""" + try: + mx = importlib.import_module("mlx.core") + except ImportError as error: + raise ImportError( + "MetalKernel needs MLX on an Apple silicon Mac: pip install 'cunumpy[metal]'", + ) from error + if not mx.metal.is_available(): + raise RuntimeError("MetalKernel needs a Metal GPU, but MLX reports none") + return mx + + +def metal_available() -> bool: + """Whether `MetalKernel` can run here (MLX installed and a Metal GPU present).""" + try: + _mlx() + except (ImportError, RuntimeError): + return False + return True + + +class MetalKernel: + """A Metal Shading Language kernel for the GPU of Apple silicon, run with MLX. + + The arrays are NumPy arrays: they are copied to MLX arrays for the launch + and the results are written back into the output arrays you pass, like a + :class:`~cunumpy.kernels.PyccelKernel` writes into its output arguments. No + backend switch is needed. The copies are counted by + :func:`~cunumpy.profiling.count_transfers`. + + Parameters + ---------- + source : str + Body of the kernel function (MSL). MLX generates the signature from + `inputs` and `outputs`: every name is a pointer to the flat, row-major + data of that array (``const device T*`` for inputs, ``device T*`` for + outputs), and ``thread_position_in_grid`` and the other Metal + attributes used in the body are added to the signature automatically. + Names from `template` are available as compile-time constants. + inputs : Sequence[str] + Names of the input arrays, in the order they are passed to the call. + outputs : Sequence[str] + Names of the output arrays, in the order they are passed as `out`. + name : str + Name of the kernel; part of the compiled function name. + header : str + Source placed before the kernel function: includes, ``#define``\\ s and + helper functions. + threadgroup : int | Sequence[int] + Threads per threadgroup: an integer, or 1 to 3 integers. + float64 : {"error", "cast"} + The GPU has no float64. With "error" (the default) a float64 input or + output raises ``TypeError``. With "cast" float64 inputs are computed + in float32 and float64 outputs are filled from float32 results, which + is what to pick when float32 precision is enough. + atomic_outputs : bool + Declare the outputs as ``device atomic*`` for atomic updates. + init_value : float | None + Value the outputs are filled with before the launch. Outputs are + otherwise uninitialized: write every element, or pass the old array as + an input as well to read its values. + + Notes + ----- + The kernel is compiled by MLX at its first call and cached. All inputs + and outputs are made C-contiguous, so ``a[i]`` in the source is the flat + index of the NumPy array. + + Examples + -------- + >>> scale = xp.kernels.MetalKernel( + ... "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];", + ... inputs=["x", "a"], + ... outputs=["y"], + ... ) + >>> y = np.empty(n, dtype=np.float32) + >>> scale(x, np.float32([2.0]), out=y) + """ + + def __init__( + self, + source: str, + inputs: Sequence[str], + outputs: Sequence[str], + *, + name: str = "cunumpy_kernel", + header: str = "", + threadgroup: int | Sequence[int] = 256, + float64: str = "error", + atomic_outputs: bool = False, + init_value: float | None = None, + ) -> None: + if float64 not in ("error", "cast"): + raise ValueError(f"float64 must be 'error' or 'cast', not {float64!r}") + if not outputs: + raise ValueError("a MetalKernel needs at least one output") + names = [*inputs, *outputs] + if len(set(names)) != len(names): + raise ValueError(f"input and output names must be distinct: {names}") + self.source = source + self.inputs = tuple(inputs) + self.outputs = tuple(outputs) + self.name = name + self.header = header + self.threadgroup = self._as_triple(threadgroup, "threadgroup", default=1) + self.float64 = float64 + self.atomic_outputs = atomic_outputs + self.init_value = init_value + self._kernel: Any = None + + def __repr__(self) -> str: + return ( + f"MetalKernel({self.name!r}, inputs={self.inputs}, outputs={self.outputs})" + ) + + @staticmethod + def _as_triple(value: int | Sequence[int], what: str, default: int) -> tuple: + values = (int(value),) if isinstance(value, (int, np.integer)) else tuple(value) + if not 1 <= len(values) <= 3 or any(v < 1 for v in values): + raise ValueError(f"{what} must be 1 to 3 positive integers, got {value!r}") + return values + (default,) * (3 - len(values)) + + @staticmethod + def _as_input(value: Any) -> np.ndarray: + """A NumPy array; Python scalars become float32 or int32 arrays of shape (1,).""" + if isinstance(value, (bool, int)): + return np.array([value], dtype=np.int32) + if isinstance(value, float): + return np.array([value], dtype=np.float32) + array = np.asarray(value) + return array.reshape(1) if array.ndim == 0 else array + + def _compiled(self, mx: Any) -> Any: + if self._kernel is None: + self._kernel = mx.fast.metal_kernel( + name=self.name, + input_names=list(self.inputs), + output_names=list(self.outputs), + source=self.source, + header=self.header, + ensure_row_contiguous=True, + atomic_outputs=self.atomic_outputs, + ) + return self._kernel + + def _device_dtype(self, dtype: np.dtype, role: str, name: str) -> np.dtype: + if dtype == _FLOAT64 and self.float64 == "cast": + return np.dtype(np.float32) + if dtype == _FLOAT64 or dtype in _UNSUPPORTED: + raise TypeError( + f"{role} {name!r} of MetalKernel {self.name!r} has dtype {dtype}, " + "which the Apple GPU does not support; use float32, or " + "float64='cast' to compute in float32", + ) + return dtype + + def __call__( + self, + *args: Any, + out: Any, + n_threads: int | Sequence[int] | None = None, + template: Mapping[str, Any] | None = None, + ) -> Any: + """Launch the kernel. + + Parameters + ---------- + *args : numpy.ndarray or scalar + The inputs, in the order of `inputs`. Python scalars become + 1-element float32 or int32 arrays; read them as ``a[0]`` in the source. + out : numpy.ndarray or Sequence[numpy.ndarray] + The output arrays, in the order of `outputs`; their shapes and + dtypes define the outputs and they are filled in place. + n_threads : int | Sequence[int] | None + Total number of threads (1 to 3 dimensions), not threadgroups. The + default is the first axis of the first output. + template : Mapping[str, int | bool | numpy.dtype] | None + Compile-time constants; each name is a constant in the source. A + different value compiles another variant. + + Returns + ------- + The output array, or a tuple of them if there are several. + """ + mx = _mlx() + if len(args) != len(self.inputs): + raise TypeError( + f"MetalKernel {self.name!r} takes {len(self.inputs)} input(s) " + f"{self.inputs}, got {len(args)}", + ) + outs = (out,) if isinstance(out, np.ndarray) else tuple(out or ()) + if len(outs) != len(self.outputs) or not all( + isinstance(o, np.ndarray) for o in outs + ): + raise TypeError( + f"out must be {len(self.outputs)} NumPy array(s) for {self.outputs}", + ) + + host_inputs = [self._as_input(a) for a in args] + for name, array in zip(self.inputs, host_inputs): + self._device_dtype(array.dtype, "input", name) + device_dtypes = [ + self._device_dtype(o.dtype, "output", name) + for name, o in zip(self.outputs, outs) + ] + if n_threads is None: + n_threads = outs[0].shape[0] if outs[0].ndim else 1 + grid = self._as_triple(n_threads, "n_threads", default=1) + + mlx_inputs = [] + for array in host_inputs: + if array.dtype == _FLOAT64: + array = array.astype(np.float32) + mlx_inputs.append(mx.array(np.ascontiguousarray(array))) + if _COUNTERS: + _record( + "to_device", + f"MetalKernel {self.name!r} input ({_describe(array)})", + nbytes=_nbytes(array), + ) + + results = self._compiled(mx)( + inputs=mlx_inputs, + template=list((template or {}).items()), + grid=grid, + threadgroup=self.threadgroup, + output_shapes=[o.shape for o in outs], + output_dtypes=[getattr(mx, dtype.name) for dtype in device_dtypes], + init_value=self.init_value, + ) + mx.eval(*results) + for result, host in zip(results, outs): + np.copyto(host, np.asarray(result), casting="same_kind") + if _COUNTERS: + _record( + "to_host", + f"MetalKernel {self.name!r} output ({_describe(host)})", + nbytes=_nbytes(host), + ) + return outs[0] if len(outs) == 1 else outs diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 6d4d644..46de0a0 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -17,6 +17,10 @@ Device dispatch can require CUDA with :func:`set_device_kernel_implementation` or temporarily with :func:`use_device_kernel_implementation`. +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. + Argument objects for CUDA kernels (:class:`~cunumpy.arguments.CudaStruct`, ...) are in :mod:`cunumpy.arguments`, the device runtime (streams, devices, debug mode) in :mod:`cunumpy.cuda`, and the pytest helpers for kernel pairs in @@ -41,6 +45,7 @@ use_device_kernel_implementation, use_host_kernel_implementation, ) +from cunumpy._metal_kernel import MetalKernel, metal_available __all__ = [ "DEVICE_IMPLEMENTATIONS", @@ -51,12 +56,14 @@ "HostImplementations", "Kernel", "KernelCatalog", + "MetalKernel", "PyccelKernel", "as_kernel_array", "fuse", "get_device_kernel_implementation", "get_host_kernel_implementation", "kernel_output", + "metal_available", "set_device_kernel_implementation", "set_host_kernel_implementation", "use_device_kernel_implementation", diff --git a/tests/unit/test_metal_kernel.py b/tests/unit/test_metal_kernel.py new file mode 100644 index 0000000..6d92e07 --- /dev/null +++ b/tests/unit/test_metal_kernel.py @@ -0,0 +1,145 @@ +"""MetalKernel: argument checking everywhere, launches on Apple silicon with MLX.""" + +import numpy as np +import pytest + +import cunumpy as xp + +MetalKernel = xp.kernels.MetalKernel +needs_metal = pytest.mark.skipif( + not xp.kernels.metal_available(), reason="needs MLX and a Metal GPU" +) + +AXPY = "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i] + b[i];" + + +def axpy(**kwargs): + return MetalKernel(AXPY, inputs=["x", "a", "b"], outputs=["y"], **kwargs) + + +def test_constructor_validates_names_and_options(): + with pytest.raises(ValueError, match="at least one output"): + MetalKernel("", inputs=["x"], outputs=[]) + with pytest.raises(ValueError, match="distinct"): + MetalKernel("", inputs=["x"], outputs=["x"]) + with pytest.raises(ValueError, match="float64"): + MetalKernel("", inputs=["x"], outputs=["y"], float64="ignore") + with pytest.raises(ValueError, match="threadgroup"): + MetalKernel("", inputs=["x"], outputs=["y"], threadgroup=(1, 2, 3, 4)) + + +def test_missing_mlx_or_gpu_gives_a_clear_error(monkeypatch): + import cunumpy._metal_kernel as module + + def unavailable(): + raise ImportError("no mlx") + + monkeypatch.setattr(module, "_mlx", unavailable) + assert not xp.kernels.metal_available() + with pytest.raises(ImportError): + axpy()(np.zeros(1, np.float32), 1.0, np.zeros(1, np.float32), out=np.zeros(1)) + + +@needs_metal +def test_wrong_argument_counts_and_outputs_are_rejected(): + kernel = axpy() + x = np.zeros(4, np.float32) + with pytest.raises(TypeError, match="takes 3 input"): + kernel(x, out=x.copy()) + with pytest.raises(TypeError, match="out must be 1 NumPy"): + kernel(x, 1.0, x, out=[x, x]) + with pytest.raises(TypeError, match="out must be 1 NumPy"): + kernel(x, 1.0, x, out=None) + + +@needs_metal +def test_float64_is_rejected_unless_cast(): + x = np.arange(4.0) + with pytest.raises(TypeError, match="does not support"): + axpy()(x.astype(np.float32), 2.0, x, out=np.empty(4, np.float32)) + with pytest.raises(TypeError, match="output 'y'"): + axpy()(x.astype(np.float32), 2.0, x.astype(np.float32), out=np.empty(4)) + + +@needs_metal +def test_axpy_with_python_scalar_and_non_contiguous_input(): + x = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::2] + assert not x.flags.c_contiguous + b = np.ones(x.shape, np.float32) + y = np.empty(x.shape, np.float32) + result = axpy()(x, 2.0, b, out=y, n_threads=x.size) + assert result is y + np.testing.assert_allclose(y, 2.0 * x + 1.0) + + +@needs_metal +def test_float64_cast_computes_in_float32_and_fills_float64_output(): + x = np.linspace(0, 1, 100) + y = np.empty(100) + axpy(float64="cast")(x, 3.0, x, out=y) + np.testing.assert_allclose(y, 4.0 * x, rtol=1e-6) + assert y.dtype == np.float64 + + +@needs_metal +def test_several_outputs_template_and_init_value(): + kernel = MetalKernel( + "uint i = thread_position_in_grid.x;" + "if (i < N) { s[i] = x[i] + 1; d[i] = x[i] - 1; }", + inputs=["x"], + outputs=["s", "d"], + init_value=-99.0, + ) + x = np.arange(4, dtype=np.float32) + s, d = np.empty(4, np.float32), np.empty(4, np.float32) + out = kernel(x, out=(s, d), n_threads=4, template={"N": 3}) + assert out == (s, d) + np.testing.assert_array_equal(s, [1, 2, 3, -99]) + np.testing.assert_array_equal(d, [-1, 0, 1, -99]) + + +@needs_metal +def test_push_matches_float64_reference(): + n, steps, dt, B = 1000, 50, 0.01, 1.5 + c, s = np.cos(dt * B), np.sin(dt * B) + rng = np.random.default_rng(0) + pos = rng.random((n, 3)).astype(np.float32) + vel = rng.standard_normal((n, 3)).astype(np.float32) + kernel = MetalKernel( + """ + uint i = thread_position_in_grid.x; + float px = pos[3*i], py = pos[3*i+1], pz = pos[3*i+2]; + float vx = vel[3*i], vy = vel[3*i+1], vz = vel[3*i+2]; + for (int k = 0; k < NSTEPS; ++k) { + px += vx*p[2]; py += vy*p[2]; pz += vz*p[2]; + float nx = p[0]*vx + p[1]*vy; vy = -p[1]*vx + p[0]*vy; vx = nx; + } + pos_out[3*i] = px; pos_out[3*i+1] = py; pos_out[3*i+2] = pz; + vel_out[3*i] = vx; vel_out[3*i+1] = vy; vel_out[3*i+2] = vz; + """, + inputs=["pos", "vel", "p"], + outputs=["pos_out", "vel_out"], + ) + pos_out, vel_out = np.empty_like(pos), np.empty_like(vel) + kernel( + pos, + vel, + np.array([c, s, dt], np.float32), + out=(pos_out, vel_out), + n_threads=n, + template={"NSTEPS": steps}, + ) + p, v = pos.astype(float), vel.astype(float) + for _ in range(steps): + p += v * dt + v[:, :2] = np.stack([c * v[:, 0] + s * v[:, 1], -s * v[:, 0] + c * v[:, 1]], 1) + np.testing.assert_allclose(pos_out, p, atol=1e-4) + np.testing.assert_allclose(vel_out, v, atol=1e-4) + + +@needs_metal +def test_copies_are_counted(): + x = np.zeros(8, np.float32) + with xp.profiling.count_transfers() as counter: + axpy()(x, 1.0, x, out=np.empty(8, np.float32)) + assert counter.to_device == 3 and counter.to_host == 1 From 1ae1a5cad3fd7afb2641f51617f34b32c3fb335c Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:45:08 +0200 Subject: [PATCH 2/2] Install gfortran --- .github/workflows/macos.yml | 7 ++ README.md | 4 +- docs/source/index.md | 1 + docs/source/installation.md | 16 +++++ docs/source/kernels/metal-kernel.md | 105 ++++++++++++++++++++++++++++ docs/source/kernels/overview.md | 1 + src/cunumpy/LLM_GUIDE.md | 1 + 7 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 docs/source/kernels/metal-kernel.md diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml index 6256349..b13bab3 100644 --- a/.github/workflows/macos.yml +++ b/.github/workflows/macos.yml @@ -33,6 +33,13 @@ jobs: cache: pip cache-dependency-path: pyproject.toml + # Pyccel is built from source here (no wheel for macOS arm64) and needs a + # Fortran compiler, which the macOS runners do not have. + - name: Install gfortran + run: | + brew install gcc + gfortran --version + - name: Install project run: | pip install --upgrade pip diff --git a/README.md b/README.md index 21e5c34..a71001a 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ never hide a NumPy name: | Submodule | Contents | |---|---| -| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, host implementations, `fuse` | +| `xp.kernels` | `Kernel`, `KernelCatalog`, `PyccelKernel`, `CudaKernel`, `MetalKernel`, host implementations, `fuse` | | `xp.arguments` | CUDA only: `CudaStruct`, `CudaStructArguments`, `CudaArguments` | | `xp.cuda` | CUDA only: devices, streams, debug mode, CUDA headers | | `xp.rng` | `random_streams`, `get_rng`, `philox_*` | @@ -36,6 +36,8 @@ never hide a NumPy name: | `cunumpy.kernel_testing` | pytest helpers for host/CUDA kernel pairs | Everything except `xp.cuda`, `xp.arguments` and `CudaKernel` works on both backends. +`MetalKernel` runs Metal kernels on the GPU of an Apple silicon Mac (MLX, float32, NumPy arrays; +`pip install 'cunumpy[metal]'`, see [the guide](docs/source/kernels/metal-kernel.md)). ## Install diff --git a/docs/source/index.md b/docs/source/index.md index 801c4d9..aed8c8f 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -63,6 +63,7 @@ array-api-compat kernels/overview kernels/pyccel-kernel kernels/cuda-kernel +kernels/metal-kernel kernels/dispatch kernels/arguments kernels/accumulation diff --git a/docs/source/installation.md b/docs/source/installation.md index e9f9248..7c62fb5 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -39,10 +39,26 @@ print("active backend:", xp.get_backend()) # 'cupy' if the GPU works If `cupy_available()` is `False`, CuNumpy quietly falls back to NumPy when CuPy is requested. See [Troubleshooting](troubleshooting.md) for the usual causes. +## Apple silicon GPUs + +CuPy does not run on Macs. The GPU of an Apple silicon Mac can run +[Metal kernels](kernels/metal-kernel.md) on NumPy float32 arrays through MLX: + +```bash +python -m pip install 'cunumpy[metal]' +``` + +```python +import cunumpy as xp + +print("Metal usable:", xp.kernels.metal_available()) +``` + ## Optional extras | Extra | Installs | Use it for | | --- | --- | --- | +| `cunumpy[metal]` | `mlx` (Apple silicon Macs only) | [`MetalKernel`](kernels/metal-kernel.md) on the Mac GPU | | `cunumpy[test]` | `pytest`, `coverage` | running the test suite, using `cunumpy.kernel_testing` | | `cunumpy[test-compiled]` | the above plus `pyccel` | tests that compile host kernels with Pyccel | | `cunumpy[docs]` | Sphinx, MyST, the book theme | building this documentation | diff --git a/docs/source/kernels/metal-kernel.md b/docs/source/kernels/metal-kernel.md new file mode 100644 index 0000000..9e0b19d --- /dev/null +++ b/docs/source/kernels/metal-kernel.md @@ -0,0 +1,105 @@ +# Metal kernels on Apple silicon + +`MetalKernel` runs a kernel written in the Metal Shading Language (MSL) on the +GPU of an Apple silicon Mac. It uses [MLX](https://github.com/ml-explore/mlx), +Apple's array framework, to compile and launch the kernel, so no Objective-C or +Xcode project is needed. It is the Mac counterpart of +[`CudaKernel`](cuda-kernel.md), with two differences: + +* it works on **NumPy arrays**: there is no Metal backend, and the kernel copies + its arguments to MLX arrays and the results back into your output arrays; +* it is **float32 only**, because the Apple GPU has no float64. + +Install it with `pip install 'cunumpy[metal]'` (MLX is installed on arm64 Macs +only). `xp.kernels.metal_available()` is `True` when MLX is installed and a Metal +GPU is present. + +## A first kernel + +```python +import numpy as np +import cunumpy as xp + +scale = xp.kernels.MetalKernel( + "uint i = thread_position_in_grid.x; y[i] = a[0] * x[i];", + inputs=["x", "a"], + outputs=["y"], +) + +x = np.arange(8, dtype=np.float32) +y = np.empty_like(x) +scale(x, 2.0, out=y) +``` + +The essentials: + +* `source` is only the **body** of the kernel function. MLX writes the + signature from `inputs` and `outputs`: each name is a pointer to the flat, + row-major data of that array, so a 2D array `a` of shape `(n, 3)` is read as + `a[3 * i + j]`. Attributes such as `thread_position_in_grid` are added to the + signature when the body uses them. +* Inputs are passed in the order of `inputs`. A Python `float` or `int` becomes + a one-element `float32` or `int32` array; read it as `a[0]`. +* `out` is one array, or one array per name in `outputs`. Their shapes and dtypes + define the outputs, and they are filled in place. The call returns them. +* `n_threads` is the **total** number of threads, not the number of threadgroups + (the opposite of the CUDA `grid`). It defaults to the first axis of the first + output. Threads beyond the data must return early, as for CUDA, if the thread + count is not a multiple of the threadgroup size. +* Outputs start uninitialized. Write every element, or pass the old array as an + input too when the kernel updates it, or give `init_value=0.0`. +* Compilation happens at the first call and MLX caches the result. + +## Compile-time constants + +`template` gives constants that are known when the kernel is compiled, which +lets the compiler unroll loops. Each name is available in the source: + +```python +push = xp.kernels.MetalKernel( + """ + uint i = thread_position_in_grid.x; + float x = pos[i], v = vel[i]; + for (int k = 0; k < NSTEPS; ++k) { x += v * dt[0]; } + pos_out[i] = x; + """, + inputs=["pos", "vel", "dt"], + outputs=["pos_out"], +) +push(pos, vel, 0.01, out=pos_out, template={"NSTEPS": 200}) +``` + +A different value of `NSTEPS` compiles another variant. `header` takes +`#include`s, `#define`s and helper functions that go before the kernel. + +## float32 and float64 + +An Apple GPU computes in float32. A float64 array passed to a `MetalKernel` +raises `TypeError` that names the argument, so no precision is lost silently. If +float32 is accurate enough for the kernel, create it with `float64="cast"`: +float64 inputs are computed in float32 and float64 outputs are filled from the +float32 results. A particle push over 200 steps agreed with a float64 reference +to about 3e-5, but whether that is acceptable depends on the physics, so check it +against the host kernel with +[`assert_kernels_agree`](testing.md) using a float32 tolerance. + +## Cost of the copies + +Every call copies the inputs to MLX arrays and the results back. The copies are +counted as `to_device` and `to_host` by +[`count_transfers()`](../guides/profiling.md), so they can be found in a profile. +On an M1, 4 million particles (3 float32 coordinates each) took about 5 ms to +copy to the GPU and 2 ms to copy back, next to a push kernel of 11.5 ms. A kernel +that runs once per time step on fresh NumPy arrays therefore pays a noticeable +share of its time in copies, and a kernel with little work per element (an +`a * x + y` update) is no faster than NumPy. The GPU pays off for kernels with +much arithmetic per element, such as pushers and interpolation. + +## What it does not do + +* It is not part of the [`Kernel`](dispatch.md) dispatch: a `Kernel` pairs a host + kernel with a CUDA kernel only, so call a `MetalKernel` directly. +* It does not keep data on the GPU between calls. +* It cannot be tested on GitHub's hosted macOS runners, which are virtual + machines without a Metal GPU. Its launch tests are skipped there and run on a + Mac with `pytest tests/unit/test_metal_kernel.py -rs`. diff --git a/docs/source/kernels/overview.md b/docs/source/kernels/overview.md index 4c892fd..c770a64 100644 --- a/docs/source/kernels/overview.md +++ b/docs/source/kernels/overview.md @@ -16,6 +16,7 @@ unchanged. | --- | --- | --- | | `PyccelKernel` | calls a host kernel with CuPy arrays by copying them to the host and back | [Host kernels with GPU data](pyccel-kernel.md) | | `CudaKernel` | wraps a CUDA C kernel, checks every call against its signature | [Writing CUDA kernels](cuda-kernel.md) | +| `MetalKernel` | runs a Metal kernel on the GPU of an Apple silicon Mac, on NumPy float32 arrays | [Metal kernels on Apple silicon](metal-kernel.md) | | `Kernel` | a host kernel plus its CUDA kernel; calls the one matching the backend | [Pairing host and CUDA kernels](dispatch.md) | | `KernelCatalog` | all `Kernel`s of a package, found by folder convention | [Pairing host and CUDA kernels](dispatch.md) | | `CudaArguments`, `CudaStruct`, `CudaStructArguments` | pass a group of arrays and scalars as one argument | [Kernel arguments and structs](arguments.md) | diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 65a6a8b..54ad69d 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -72,6 +72,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | hand data to SciPy/matplotlib/h5py | `xp.to_numpy(a)` | | call an existing NumPy-only kernel with GPU arrays (slow, correct) | `xp.kernels.PyccelKernel(fn, outputs=(...))` | | launch a hand-written CUDA C kernel | `xp.kernels.CudaKernel(source, "name")` / `CudaKernel.from_file(path)` | +| run a Metal (MSL) kernel on an Apple silicon GPU, NumPy float32 in and out (float64 raises; `float64="cast"` computes in float32); not part of `Kernel` dispatch | `xp.kernels.MetalKernel(body, inputs=[...], outputs=[...])(*args, out=arrays, n_threads=n)`; check `xp.kernels.metal_available()` | | host kernel + CUDA port, chosen by backend | `xp.kernels.Kernel(host_fn, cuda_kernel_or_None)` | | many kernels in a package, ported incrementally | `xp.kernels.KernelCatalog.from_package(__name__, missing_cuda="fallback")` | | 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` |