metal-linalg 2.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- metal_linalg-2.2.0/.gitignore +24 -0
- metal_linalg-2.2.0/CHANGELOG.md +162 -0
- metal_linalg-2.2.0/CMakeLists.txt +327 -0
- metal_linalg-2.2.0/CONTRIBUTING.md +203 -0
- metal_linalg-2.2.0/LICENSE +21 -0
- metal_linalg-2.2.0/PKG-INFO +106 -0
- metal_linalg-2.2.0/README.md +532 -0
- metal_linalg-2.2.0/cmake/EmbedMetallib.cmake +24 -0
- metal_linalg-2.2.0/cmake/MetalLinalgConfig.cmake.in +9 -0
- metal_linalg-2.2.0/include/metal_linalg/c_api.h +138 -0
- metal_linalg-2.2.0/include/metal_linalg/core.h +451 -0
- metal_linalg-2.2.0/include/metal_linalg/device.h +17 -0
- metal_linalg-2.2.0/include/metal_linalg/eigh.h +57 -0
- metal_linalg-2.2.0/include/metal_linalg/metal_linalg.h +8 -0
- metal_linalg-2.2.0/include/metal_linalg/qr.h +38 -0
- metal_linalg-2.2.0/include/metal_linalg/svd.h +57 -0
- metal_linalg-2.2.0/pyproject.toml +53 -0
- metal_linalg-2.2.0/python/CMakeLists.txt +76 -0
- metal_linalg-2.2.0/python/README.md +87 -0
- metal_linalg-2.2.0/python/check_wheel.py +61 -0
- metal_linalg-2.2.0/python/metal_linalg/__init__.py +166 -0
- metal_linalg-2.2.0/python/nanobind_abi.py +28 -0
- metal_linalg-2.2.0/python/src/bindings.cpp +122 -0
- metal_linalg-2.2.0/shaders/Eigh_BlockJacobi.metal +342 -0
- metal_linalg-2.2.0/shaders/Eigh_Jacobi.metal +236 -0
- metal_linalg-2.2.0/shaders/QR_Streaming_AMX_Complete.metal +962 -0
- metal_linalg-2.2.0/shaders/QR_Streaming_AMX_Reduced.metal +993 -0
- metal_linalg-2.2.0/shaders/QR_Unblocked.metal +833 -0
- metal_linalg-2.2.0/shaders/Svd_BlockJacobi.metal +344 -0
- metal_linalg-2.2.0/shaders/Svd_Jacobi.metal +339 -0
- metal_linalg-2.2.0/shaders/block_jacobi_common.h +58 -0
- metal_linalg-2.2.0/shaders/eigh_jacobi_common.h +245 -0
- metal_linalg-2.2.0/shaders/prebuilt/Eigh_BlockJacobi.metallib +0 -0
- metal_linalg-2.2.0/shaders/prebuilt/Eigh_Jacobi.metallib +0 -0
- metal_linalg-2.2.0/shaders/prebuilt/QR_Streaming_AMX_Complete.metallib +0 -0
- metal_linalg-2.2.0/shaders/prebuilt/QR_Streaming_AMX_Reduced.metallib +0 -0
- metal_linalg-2.2.0/shaders/prebuilt/QR_Unblocked.metallib +0 -0
- metal_linalg-2.2.0/shaders/prebuilt/Svd_BlockJacobi.metallib +0 -0
- metal_linalg-2.2.0/shaders/prebuilt/Svd_Jacobi.metallib +0 -0
- metal_linalg-2.2.0/src/c_api.cpp +183 -0
- metal_linalg-2.2.0/src/device.mm +62 -0
- metal_linalg-2.2.0/src/eigh.mm +529 -0
- metal_linalg-2.2.0/src/eigh_block_jacobi.mm +359 -0
- metal_linalg-2.2.0/src/metal_runtime.h +183 -0
- metal_linalg-2.2.0/src/metal_runtime.mm +245 -0
- metal_linalg-2.2.0/src/mlx_api.cpp +226 -0
- metal_linalg-2.2.0/src/qr.mm +206 -0
- metal_linalg-2.2.0/src/qr_cpu.mm +79 -0
- metal_linalg-2.2.0/src/qr_streaming_amx_complete.mm +227 -0
- metal_linalg-2.2.0/src/qr_streaming_amx_reduced.mm +256 -0
- metal_linalg-2.2.0/src/qr_unblocked.mm +188 -0
- metal_linalg-2.2.0/src/shaders.h +35 -0
- metal_linalg-2.2.0/src/svd.mm +667 -0
- metal_linalg-2.2.0/src/svd_block_jacobi.mm +362 -0
- metal_linalg-2.2.0/src/tuned/eigh.inc +6 -0
- metal_linalg-2.2.0/src/tuned/qr.inc +6 -0
- metal_linalg-2.2.0/src/tuned/svd.inc +4 -0
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
# Build directories
|
|
2
|
+
build*/
|
|
3
|
+
cmake-build-*/
|
|
4
|
+
.build/
|
|
5
|
+
.swiftpm/
|
|
6
|
+
Package.resolved
|
|
7
|
+
|
|
8
|
+
# Compiled Metal shaders (generated by CMake from the .metal sources)
|
|
9
|
+
*.air
|
|
10
|
+
*.metallib
|
|
11
|
+
!shaders/prebuilt/*.metallib
|
|
12
|
+
|
|
13
|
+
# IDE
|
|
14
|
+
.idea/
|
|
15
|
+
.vscode/
|
|
16
|
+
.DS_Store
|
|
17
|
+
|
|
18
|
+
# Python bytecode
|
|
19
|
+
__pycache__/
|
|
20
|
+
|
|
21
|
+
# Written by tuning/tune_{qr,eigh,svd}.py (committed copies live in docs/results/)
|
|
22
|
+
qr-tune-results/
|
|
23
|
+
eigh-tune-results/
|
|
24
|
+
svd-tune-results/
|
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
# Changes
|
|
2
|
+
|
|
3
|
+
## 2.2.0 (2026-10-02)
|
|
4
|
+
|
|
5
|
+
- **The Python package is on PyPI**: `pip install metal-linalg` installs a
|
|
6
|
+
prebuilt wheel (Apple Silicon, macOS 14+, Python 3.10 to 3.14) and the MLX
|
|
7
|
+
it was built against, which it pins exactly (`mlx==0.32.3`). Every release
|
|
8
|
+
builds, checks and tests the wheels and publishes them
|
|
9
|
+
([wheels.yml](.github/workflows/wheels.yml)). Building from source no
|
|
10
|
+
longer needs `--no-build-isolation`: the build fetches the pinned MLX and
|
|
11
|
+
nanobind itself.
|
|
12
|
+
- The Python build accepts nanobind 3's ABI, which the pip MLX uses from
|
|
13
|
+
0.32.3 (nanobind 3.0.1); it still builds against MLX built with nanobind
|
|
14
|
+
2.x, such as Homebrew's.
|
|
15
|
+
- The Python tests check residuals with CPU matmuls: MLX 0.32.3 multiplies
|
|
16
|
+
float32 on the M5's GPU to only about 1e-2, which failed checks of results
|
|
17
|
+
that are accurate to 1e-6.
|
|
18
|
+
|
|
19
|
+
## 2.1.0 (2026-10-02)
|
|
20
|
+
|
|
21
|
+
- **QR has a CPU path**, like the eigensolver and the SVD: LAPACK's `sgeqrf`
|
|
22
|
+
and `sorgqr`, for lone and small-batch calls that do not pay for a GPU
|
|
23
|
+
launch (on an M5 Pro, 4 matrices of 64×64 take 0.13 ms on the CPU against
|
|
24
|
+
1.04 ms on the GPU). `QrPolicy` gains the GPU-or-CPU boundary
|
|
25
|
+
(`gpu_max_k`, `gpu_min_batch_times_k`, `gpu_min_batch`), with the
|
|
26
|
+
`QR_GPU_*` and `QR_DEVICE` environment overrides; `QrBackend::cpu`,
|
|
27
|
+
`qr_gpu_backend()`, `qr_uses_gpu()` and `detail::qr_cpu` are new, and the C,
|
|
28
|
+
Python and Swift APIs follow. The QR sweep times the CPU and the tuner fits
|
|
29
|
+
the boundary; a device measured before this keeps sending every QR call to
|
|
30
|
+
the GPU until it is measured again. See [docs/qr.md](docs/qr.md).
|
|
31
|
+
- The QR guide documents the backends' sign conventions for R's diagonal,
|
|
32
|
+
which differ, and how to normalise them.
|
|
33
|
+
|
|
34
|
+
## 2.0.1 (2026-10-01)
|
|
35
|
+
|
|
36
|
+
- metal-linalg is licensed under the MIT licence ([LICENSE](LICENSE)).
|
|
37
|
+
|
|
38
|
+
## 2.0.0 (2026-10-01)
|
|
39
|
+
|
|
40
|
+
The project is renamed from `qr-apple-silicon` to **metal-linalg**, since it
|
|
41
|
+
now covers three decompositions, and is packaged as a library.
|
|
42
|
+
|
|
43
|
+
### New
|
|
44
|
+
|
|
45
|
+
- **Symmetric eigensolver**: `eigh_accelerated`, `eigvalsh_accelerated`. Cyclic
|
|
46
|
+
Jacobi in two kernels, a whole-matrix kernel (simd and threadgroup modes)
|
|
47
|
+
and a block kernel that spreads one matrix over the GPU, with a per-device
|
|
48
|
+
route to MLX's CPU `eigh` where that is faster. See [docs/eigh.md](docs/eigh.md).
|
|
49
|
+
- **Thin SVD**: `svd_accelerated`, `svdvals_accelerated`. One-sided Jacobi in a
|
|
50
|
+
whole-matrix and a block kernel, each optionally after this library's QR for
|
|
51
|
+
tall input, with a per-device route to the CPU. See [docs/svd.md](docs/svd.md).
|
|
52
|
+
- **Per-device routing policies** for all three solvers, measured on an Apple
|
|
53
|
+
M1 (QR, eigh) and an Apple M5 Pro (all three), with environment and
|
|
54
|
+
programmatic overrides.
|
|
55
|
+
- **Measuring a Mac is one command**, `python3 tuning/run.py`: it checks the
|
|
56
|
+
machine, builds, tests, measures all three decompositions and writes a
|
|
57
|
+
uniquely named submission (`docs/results/<device>/<date>-<random>/`), so any
|
|
58
|
+
number of people with the same Mac can contribute.
|
|
59
|
+
The library's per-device tables (`src/tuned/`) are generated from every
|
|
60
|
+
run submitted for each device, by `tuning/generate_tables.py`, which a
|
|
61
|
+
GitHub Action runs on each results pull request (to validate it and show the
|
|
62
|
+
effect) and after each merge (to apply it). See [docs/tuning.md](docs/tuning.md). For QR
|
|
63
|
+
this replaces the fixed rule of 1.0 with a crossover on the row count alone,
|
|
64
|
+
measured over square, tall, wide and near-square shapes
|
|
65
|
+
([study](docs/studies/qr-routing-apple-m1.md)).
|
|
66
|
+
- **Python package** `metal_linalg` (`python/`, `pyproject.toml`): `qr`, `eigh`,
|
|
67
|
+
`eigvalsh`, `svd`, `svdvals` on `mlx.core` arrays, the routing queries and
|
|
68
|
+
policies. Compiled against the installed MLX and sharing its arrays without
|
|
69
|
+
copying; see [python/README.md](python/README.md).
|
|
70
|
+
- **A core without MLX**, on plain float buffers: `<metal_linalg/core.h>`
|
|
71
|
+
(`core::qr`, `core::eigh`, `core::svd` and every backend). The MLX API is
|
|
72
|
+
now a thin layer over it, with the same names and signatures as before.
|
|
73
|
+
`-DMETAL_LINALG_WITH_MLX=OFF` builds the core alone.
|
|
74
|
+
- **A C API**, `<metal_linalg/c_api.h>`: the decompositions, routing queries
|
|
75
|
+
and policies behind C types, with status codes and per-thread error
|
|
76
|
+
messages. See [docs/c-api.md](docs/c-api.md).
|
|
77
|
+
- **A Swift package**: `MetalLinalg` on `[Float]`, and `MetalLinalgMLX` on
|
|
78
|
+
mlx-swift's `MLXArray`, built on the C API; the shaders reach it through
|
|
79
|
+
C23 `#embed` of `shaders/prebuilt/`. See [docs/swift.md](docs/swift.md).
|
|
80
|
+
- **Objective-C**: a guide, [docs/objective-c.md](docs/objective-c.md), and
|
|
81
|
+
`examples/objc_quickstart.mm`; Objective-C can use the MLX API from `.mm`
|
|
82
|
+
files or the C API from anywhere.
|
|
83
|
+
- Every measurement run records the exact Mac (e.g. "MacBook Pro (16-inch, M5
|
|
84
|
+
Pro)"), and its power and thermal state after each part of the run; a
|
|
85
|
+
device's combined summary compares the runs by machine, to find outliers.
|
|
86
|
+
- [CONTRIBUTING.md](CONTRIBUTING.md): measuring a Mac and sending the results as a
|
|
87
|
+
pull request (with the GitHub CLI, with git alone, or as a zip on an issue).
|
|
88
|
+
- [docs/reading-reports.md](docs/reading-reports.md) explains every number in
|
|
89
|
+
a measurement report, with worked examples from the M5 Pro run.
|
|
90
|
+
- `<metal_linalg/device.h>`: `device_name()`, `gpu_core_count()`.
|
|
91
|
+
- `<metal_linalg/metal_linalg.h>`, which includes everything.
|
|
92
|
+
- `qr_backend(m, n, batch)`, like `eigh_backend` and `svd_backend`.
|
|
93
|
+
- `sweep_qr --policy`, like the other two sweeps.
|
|
94
|
+
- `examples/`: five self-checking programs (a quick start, orthonormal bases
|
|
95
|
+
with QR, PCA with eigh, the nearest orthogonal matrix with the SVD, and the
|
|
96
|
+
routing queries), built with the tests and run by `ctest`.
|
|
97
|
+
|
|
98
|
+
### Packaging
|
|
99
|
+
|
|
100
|
+
- The compiled shaders are embedded in the library, so an installed
|
|
101
|
+
`libmetal_linalg` is self-contained: nothing is looked up on disk at run
|
|
102
|
+
time and consumers need no Metal compiler. Without the compiler the build
|
|
103
|
+
uses the metallibs committed under `shaders/prebuilt/`.
|
|
104
|
+
- A CMake package: `find_package(MetalLinalg)` and
|
|
105
|
+
`metal_linalg::metal_linalg`. As a subproject (`add_subdirectory`,
|
|
106
|
+
`FetchContent`) it builds static and installs nothing.
|
|
107
|
+
- A Homebrew formula, in its own tap, [c0rmac/homebrew-metal-linalg](https://github.com/c0rmac/homebrew-metal-linalg):
|
|
108
|
+
`brew tap c0rmac/metal-linalg`, `brew trust c0rmac/metal-linalg`, then `brew install metal-linalg`.
|
|
109
|
+
- **Releases are automatic**: every update to `main` that changes the library
|
|
110
|
+
publishes the next version (a tag, a GitHub release with its source tarball)
|
|
111
|
+
and points the Homebrew formula at it. See CONTRIBUTING.md, "Releases".
|
|
112
|
+
|
|
113
|
+
### Breaking changes
|
|
114
|
+
|
|
115
|
+
| 1.x | 2.0 |
|
|
116
|
+
|---|---|
|
|
117
|
+
| namespace `custom_math` | `metal_linalg` |
|
|
118
|
+
| `#include "qr.h"`, `"qr_detail.h"` | `#include <metal_linalg/qr.h>` (the `detail` backends are in it) |
|
|
119
|
+
| `DispatchPolicy` | `QrPolicy` |
|
|
120
|
+
| `dispatch_policy()`, `dispatch_policy_source()`, `set_dispatch_policy()` | `qr_policy()`, `qr_policy_source()`, `set_qr_policy()` |
|
|
121
|
+
| `m_crossover_for_batch(p, batch)` | `qr_backend(m, n, batch)` |
|
|
122
|
+
| CMake target `qr_metal` | `metal_linalg::metal_linalg` |
|
|
123
|
+
|
|
124
|
+
### Changed
|
|
125
|
+
|
|
126
|
+
- **The CPU paths call LAPACK directly**: `ssyevd` for eigh and `sgesdd` for
|
|
127
|
+
the SVD (Accelerate, after a thin QR for tall input), instead of MLX's CPU
|
|
128
|
+
`eigh` and `svd`. On an M5 Pro the eigensolver's CPU path is as fast as
|
|
129
|
+
before and the SVD's is as fast on square and wide shapes and 5-25% faster
|
|
130
|
+
on tall ones. Non-finite input gives NaN on the CPU too, for that matrix
|
|
131
|
+
alone, as on the GPU. The routing tables were measured against the old
|
|
132
|
+
paths and are due to be remeasured.
|
|
133
|
+
- The eigensolver's tuning sweep times the library's own CPU path, as the
|
|
134
|
+
SVD's already did, rather than MLX's.
|
|
135
|
+
|
|
136
|
+
### Fixes
|
|
137
|
+
|
|
138
|
+
- **QR was not scale-invariant.** The kernels compare squared column norms
|
|
139
|
+
with an absolute threshold, so entries around 1e-3 lost accuracy, below 1e-5
|
|
140
|
+
the factorisation failed, and above 1e+18 it returned NaN. Every matrix is
|
|
141
|
+
now scaled by an exact power of two on the way in.
|
|
142
|
+
- **QR discarded real data for nearly dependent columns**: the reflection
|
|
143
|
+
threshold (1e-7 on a squared norm) treated column tails shorter than 3e-4 of
|
|
144
|
+
the matrix's scale as zero. It is now 1e-30, which only guards the division.
|
|
145
|
+
- **The streaming QR backends could write past their buffers.** Their cached
|
|
146
|
+
workspaces were keyed by the padded shape but sized by the exact one, so a
|
|
147
|
+
call whose shape padded like an earlier, smaller one (1000×1000, then
|
|
148
|
+
1024×1024) reused buffers too small for it. They are now keyed by the exact
|
|
149
|
+
shape.
|
|
150
|
+
- **Every call leaked Metal objects**: the library was built without ARC and
|
|
151
|
+
without an autorelease pool, so each call kept its command buffers and a
|
|
152
|
+
buffer wrapper alive for the life of the thread, which mattered in long
|
|
153
|
+
loops. It is now built with ARC, and every GPU call drains its own pool.
|
|
154
|
+
- **Transposed and other strided views were read as their untransposed
|
|
155
|
+
buffer**: an unevaluated MLX array reports itself contiguous, so contiguity
|
|
156
|
+
is now checked after evaluation.
|
|
157
|
+
|
|
158
|
+
## 1.0.0
|
|
159
|
+
|
|
160
|
+
`qr-apple-silicon`: batched QR on Apple GPUs (`custom_math::qr_accelerated`),
|
|
161
|
+
with single-threadgroup and grid-parallel Householder backends and a fixed
|
|
162
|
+
rule choosing between them.
|
|
@@ -0,0 +1,327 @@
|
|
|
1
|
+
cmake_minimum_required(VERSION 3.25)
|
|
2
|
+
if (NOT CMAKE_HOST_SYSTEM_NAME STREQUAL "Darwin")
|
|
3
|
+
message(FATAL_ERROR "metal-linalg runs on Apple Silicon Macs only (Metal and Accelerate); "
|
|
4
|
+
"it cannot be built on ${CMAKE_HOST_SYSTEM_NAME}.")
|
|
5
|
+
endif ()
|
|
6
|
+
project(metal_linalg VERSION 2.2.0 LANGUAGES CXX OBJCXX
|
|
7
|
+
DESCRIPTION "QR, symmetric eigendecomposition and SVD on Apple GPUs, for MLX")
|
|
8
|
+
|
|
9
|
+
set(CMAKE_CXX_STANDARD 20)
|
|
10
|
+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
|
11
|
+
set(CMAKE_CXX_EXTENSIONS OFF)
|
|
12
|
+
|
|
13
|
+
# Standalone (top level) it is a shared library with install rules, which is
|
|
14
|
+
# what Homebrew ships. As a subproject (add_subdirectory / FetchContent) it is
|
|
15
|
+
# a static library folded into the parent's binary, with no install rules, so
|
|
16
|
+
# a parent that builds it from source gets the shaders and routing tables of
|
|
17
|
+
# that checkout and nothing else to ship.
|
|
18
|
+
option(METAL_LINALG_SHARED "Build libmetal_linalg as a shared library" ${PROJECT_IS_TOP_LEVEL})
|
|
19
|
+
option(METAL_LINALG_INSTALL "Generate install and CMake package rules" ${PROJECT_IS_TOP_LEVEL})
|
|
20
|
+
option(METAL_LINALG_BUILD_TESTS "Build the tests, benchmarks and tuning harnesses" ${PROJECT_IS_TOP_LEVEL})
|
|
21
|
+
option(METAL_LINALG_BUILD_EXAMPLES "Build the examples (run as tests)" ${METAL_LINALG_BUILD_TESTS})
|
|
22
|
+
option(METAL_LINALG_BUILD_PYTHON "Build the Python extension (pyproject.toml turns this on)" OFF)
|
|
23
|
+
option(METAL_LINALG_USE_PREBUILT_SHADERS
|
|
24
|
+
"Use shaders/prebuilt/*.metallib instead of compiling the .metal sources" OFF)
|
|
25
|
+
# Without MLX the library is the buffer core alone (core.h, c_api.h): no MLX
|
|
26
|
+
# API, and nothing that needs MLX (tests, benchmarks, examples, Python).
|
|
27
|
+
option(METAL_LINALG_WITH_MLX "Build the MLX API over the core" ON)
|
|
28
|
+
if (NOT METAL_LINALG_WITH_MLX)
|
|
29
|
+
foreach (_opt METAL_LINALG_BUILD_TESTS METAL_LINALG_BUILD_EXAMPLES METAL_LINALG_BUILD_PYTHON)
|
|
30
|
+
if (${_opt})
|
|
31
|
+
message(STATUS "metal-linalg: ${_opt} needs MLX; turned off")
|
|
32
|
+
set(${_opt} OFF CACHE BOOL "" FORCE)
|
|
33
|
+
endif ()
|
|
34
|
+
endforeach ()
|
|
35
|
+
endif ()
|
|
36
|
+
|
|
37
|
+
# The Python extension must use the very MLX the `mlx` Python package loads, so
|
|
38
|
+
# look where that package says first: a pip-installed mlx carries its own, and
|
|
39
|
+
# for Homebrew's the search falls through to the Homebrew prefix.
|
|
40
|
+
if (METAL_LINALG_BUILD_PYTHON)
|
|
41
|
+
find_package(Python 3.10 REQUIRED COMPONENTS Interpreter Development.Module)
|
|
42
|
+
execute_process(COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
|
|
43
|
+
OUTPUT_VARIABLE _mlx_py_root OUTPUT_STRIP_TRAILING_WHITESPACE
|
|
44
|
+
RESULT_VARIABLE _mlx_py_rc ERROR_QUIET)
|
|
45
|
+
if (NOT _mlx_py_rc EQUAL 0)
|
|
46
|
+
message(FATAL_ERROR "metal-linalg: mlx is not installed for ${Python_EXECUTABLE}. "
|
|
47
|
+
"See python/README.md.")
|
|
48
|
+
endif ()
|
|
49
|
+
list(PREPEND CMAKE_PREFIX_PATH "${_mlx_py_root}")
|
|
50
|
+
endif ()
|
|
51
|
+
|
|
52
|
+
if (METAL_LINALG_WITH_MLX)
|
|
53
|
+
find_package(MLX REQUIRED)
|
|
54
|
+
endif ()
|
|
55
|
+
include(GNUInstallDirs)
|
|
56
|
+
|
|
57
|
+
# =============================================================================
|
|
58
|
+
# Metal shaders
|
|
59
|
+
# =============================================================================
|
|
60
|
+
# Each .metal is compiled to a metallib and the metallib's bytes are compiled
|
|
61
|
+
# into the library, so an installed libmetal_linalg is self-contained: nothing
|
|
62
|
+
# is looked up on disk at run time and consumers need no Metal toolchain.
|
|
63
|
+
#
|
|
64
|
+
# Building from source needs the Metal shader compiler. From Xcode 26 it is a
|
|
65
|
+
# separate download (xcodebuild -downloadComponent MetalToolchain); where it is
|
|
66
|
+
# missing, the metallibs committed under shaders/prebuilt/ are used instead.
|
|
67
|
+
# They are the same AIR whatever GPU compiled them. After editing a shader on
|
|
68
|
+
# a machine that has the compiler, `cmake --build build --target
|
|
69
|
+
# update_prebuilt_shaders` refreshes that directory.
|
|
70
|
+
|
|
71
|
+
execute_process(COMMAND xcrun -sdk macosx metal --version
|
|
72
|
+
RESULT_VARIABLE _metal_rc OUTPUT_QUIET ERROR_QUIET)
|
|
73
|
+
if (_metal_rc EQUAL 0)
|
|
74
|
+
set(METAL_LINALG_HAVE_METAL_COMPILER ON)
|
|
75
|
+
else ()
|
|
76
|
+
set(METAL_LINALG_HAVE_METAL_COMPILER OFF)
|
|
77
|
+
endif ()
|
|
78
|
+
if (NOT METAL_LINALG_HAVE_METAL_COMPILER AND NOT METAL_LINALG_USE_PREBUILT_SHADERS)
|
|
79
|
+
message(STATUS "metal-linalg: no Metal shader compiler (xcodebuild -downloadComponent "
|
|
80
|
+
"MetalToolchain installs it); using shaders/prebuilt")
|
|
81
|
+
set(METAL_LINALG_USE_PREBUILT_SHADERS ON)
|
|
82
|
+
endif ()
|
|
83
|
+
|
|
84
|
+
set(_ml_shader_dir "${CMAKE_CURRENT_BINARY_DIR}/shaders")
|
|
85
|
+
file(MAKE_DIRECTORY "${_ml_shader_dir}")
|
|
86
|
+
|
|
87
|
+
# metal_linalg_add_shader(<name> [FLAGS ...] [DEPENDS <header> ...])
|
|
88
|
+
# Compiles shaders/<name>.metal (or copies the prebuilt metallib) and embeds
|
|
89
|
+
# it as the C symbols metal_linalg_<name>_metallib / _len; see src/shaders.h.
|
|
90
|
+
function(metal_linalg_add_shader name)
|
|
91
|
+
cmake_parse_arguments(_arg "" "" "FLAGS;DEPENDS" ${ARGN})
|
|
92
|
+
set(_src "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${name}.metal")
|
|
93
|
+
set(_air "${_ml_shader_dir}/${name}.air")
|
|
94
|
+
set(_lib "${_ml_shader_dir}/${name}.metallib")
|
|
95
|
+
set(_cpp "${_ml_shader_dir}/${name}_metallib.cpp")
|
|
96
|
+
|
|
97
|
+
if (METAL_LINALG_USE_PREBUILT_SHADERS)
|
|
98
|
+
set(_pre "${CMAKE_CURRENT_SOURCE_DIR}/shaders/prebuilt/${name}.metallib")
|
|
99
|
+
if (NOT EXISTS "${_pre}")
|
|
100
|
+
message(FATAL_ERROR "metal-linalg: no Metal shader compiler and no prebuilt ${name}.metallib")
|
|
101
|
+
endif ()
|
|
102
|
+
add_custom_command(
|
|
103
|
+
OUTPUT "${_lib}"
|
|
104
|
+
COMMAND ${CMAKE_COMMAND} -E copy_if_different "${_pre}" "${_lib}"
|
|
105
|
+
DEPENDS "${_pre}"
|
|
106
|
+
COMMENT "Using prebuilt ${name}.metallib"
|
|
107
|
+
VERBATIM)
|
|
108
|
+
else ()
|
|
109
|
+
set(_deps "${_src}")
|
|
110
|
+
foreach (_d IN LISTS _arg_DEPENDS)
|
|
111
|
+
list(APPEND _deps "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${_d}")
|
|
112
|
+
endforeach ()
|
|
113
|
+
# Re-runs whenever the source or a header it includes changes, so an
|
|
114
|
+
# edited shader can never be masked by a stale metallib.
|
|
115
|
+
add_custom_command(
|
|
116
|
+
OUTPUT "${_lib}"
|
|
117
|
+
COMMAND xcrun -sdk macosx metal ${_arg_FLAGS} -c "${_src}" -o "${_air}"
|
|
118
|
+
COMMAND xcrun -sdk macosx metallib "${_air}" -o "${_lib}"
|
|
119
|
+
DEPENDS ${_deps}
|
|
120
|
+
BYPRODUCTS "${_air}"
|
|
121
|
+
COMMENT "Compiling ${name}.metal -> ${name}.metallib"
|
|
122
|
+
VERBATIM)
|
|
123
|
+
endif ()
|
|
124
|
+
|
|
125
|
+
add_custom_command(
|
|
126
|
+
OUTPUT "${_cpp}"
|
|
127
|
+
COMMAND ${CMAKE_COMMAND} -DINPUT=${_lib} -DOUTPUT=${_cpp}
|
|
128
|
+
-DSYMBOL=metal_linalg_${name}_metallib
|
|
129
|
+
-P "${CMAKE_CURRENT_SOURCE_DIR}/cmake/EmbedMetallib.cmake"
|
|
130
|
+
DEPENDS "${_lib}" "${CMAKE_CURRENT_SOURCE_DIR}/cmake/EmbedMetallib.cmake"
|
|
131
|
+
COMMENT "Embedding ${name}.metallib"
|
|
132
|
+
VERBATIM)
|
|
133
|
+
|
|
134
|
+
set_property(GLOBAL APPEND PROPERTY METAL_LINALG_EMBEDDED_SOURCES "${_cpp}")
|
|
135
|
+
set_property(GLOBAL APPEND PROPERTY METAL_LINALG_METALLIBS "${_lib}")
|
|
136
|
+
endfunction()
|
|
137
|
+
|
|
138
|
+
metal_linalg_add_shader(QR_Unblocked)
|
|
139
|
+
metal_linalg_add_shader(QR_Streaming_AMX_Reduced)
|
|
140
|
+
metal_linalg_add_shader(QR_Streaming_AMX_Complete)
|
|
141
|
+
|
|
142
|
+
# The Jacobi rotations need c^2 + s^2 == 1 to roundoff and the non-finite
|
|
143
|
+
# check needs IEEE semantics, neither of which fast-math guarantees. The
|
|
144
|
+
# kernels are barrier- and memory-bound, so precise arithmetic costs nothing
|
|
145
|
+
# measurable.
|
|
146
|
+
metal_linalg_add_shader(Eigh_Jacobi FLAGS -fno-fast-math DEPENDS eigh_jacobi_common.h)
|
|
147
|
+
metal_linalg_add_shader(Eigh_BlockJacobi FLAGS -fno-fast-math DEPENDS eigh_jacobi_common.h block_jacobi_common.h)
|
|
148
|
+
metal_linalg_add_shader(Svd_Jacobi FLAGS -fno-fast-math DEPENDS eigh_jacobi_common.h)
|
|
149
|
+
metal_linalg_add_shader(Svd_BlockJacobi FLAGS -fno-fast-math DEPENDS eigh_jacobi_common.h block_jacobi_common.h)
|
|
150
|
+
|
|
151
|
+
get_property(_ml_embedded GLOBAL PROPERTY METAL_LINALG_EMBEDDED_SOURCES)
|
|
152
|
+
get_property(_ml_metallibs GLOBAL PROPERTY METAL_LINALG_METALLIBS)
|
|
153
|
+
|
|
154
|
+
if (NOT METAL_LINALG_USE_PREBUILT_SHADERS)
|
|
155
|
+
add_custom_target(update_prebuilt_shaders
|
|
156
|
+
COMMAND ${CMAKE_COMMAND} -E copy_if_different ${_ml_metallibs}
|
|
157
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/shaders/prebuilt/"
|
|
158
|
+
DEPENDS ${_ml_metallibs}
|
|
159
|
+
COMMENT "Refreshing shaders/prebuilt/ from this build"
|
|
160
|
+
VERBATIM)
|
|
161
|
+
endif ()
|
|
162
|
+
|
|
163
|
+
# =============================================================================
|
|
164
|
+
# The library
|
|
165
|
+
# =============================================================================
|
|
166
|
+
|
|
167
|
+
if (METAL_LINALG_SHARED)
|
|
168
|
+
set(_ml_kind SHARED)
|
|
169
|
+
else ()
|
|
170
|
+
set(_ml_kind STATIC)
|
|
171
|
+
endif ()
|
|
172
|
+
add_library(metal_linalg ${_ml_kind}
|
|
173
|
+
src/device.mm
|
|
174
|
+
src/metal_runtime.mm
|
|
175
|
+
src/qr.mm
|
|
176
|
+
src/qr_unblocked.mm
|
|
177
|
+
src/qr_streaming_amx_reduced.mm
|
|
178
|
+
src/qr_streaming_amx_complete.mm
|
|
179
|
+
src/qr_cpu.mm
|
|
180
|
+
src/eigh.mm
|
|
181
|
+
src/eigh_block_jacobi.mm
|
|
182
|
+
src/svd.mm
|
|
183
|
+
src/svd_block_jacobi.mm
|
|
184
|
+
src/c_api.cpp
|
|
185
|
+
${_ml_embedded}
|
|
186
|
+
)
|
|
187
|
+
if (METAL_LINALG_WITH_MLX)
|
|
188
|
+
target_sources(metal_linalg PRIVATE src/mlx_api.cpp)
|
|
189
|
+
endif ()
|
|
190
|
+
add_library(metal_linalg::metal_linalg ALIAS metal_linalg)
|
|
191
|
+
|
|
192
|
+
# An absolute install name, so a consumer linked against the installed
|
|
193
|
+
# library finds it without an rpath; Homebrew passes its opt path.
|
|
194
|
+
set_target_properties(metal_linalg PROPERTIES
|
|
195
|
+
VERSION ${PROJECT_VERSION}
|
|
196
|
+
SOVERSION ${PROJECT_VERSION_MAJOR}
|
|
197
|
+
EXPORT_NAME metal_linalg
|
|
198
|
+
INSTALL_NAME_DIR "${CMAKE_INSTALL_FULL_LIBDIR}"
|
|
199
|
+
)
|
|
200
|
+
target_compile_features(metal_linalg PUBLIC cxx_std_20)
|
|
201
|
+
target_include_directories(metal_linalg
|
|
202
|
+
PUBLIC
|
|
203
|
+
$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
|
|
204
|
+
$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
|
|
205
|
+
PRIVATE
|
|
206
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/src"
|
|
207
|
+
)
|
|
208
|
+
# The Objective-C++ sources are written for ARC: Metal objects are released
|
|
209
|
+
# when their last reference goes, not leaked per call.
|
|
210
|
+
target_compile_options(metal_linalg PRIVATE $<$<COMPILE_LANGUAGE:OBJCXX>:-fobjc-arc>)
|
|
211
|
+
|
|
212
|
+
# The public headers only need MLX; the frameworks are the library's own, but
|
|
213
|
+
# a static build hands them to the consumer's link line.
|
|
214
|
+
set(_ml_frameworks "-framework Metal" "-framework MetalPerformanceShaders" "-framework Foundation"
|
|
215
|
+
"-framework IOKit" "-framework Accelerate")
|
|
216
|
+
target_link_libraries(metal_linalg PRIVATE ${_ml_frameworks})
|
|
217
|
+
if (METAL_LINALG_WITH_MLX)
|
|
218
|
+
target_link_libraries(metal_linalg PUBLIC mlx)
|
|
219
|
+
endif ()
|
|
220
|
+
if (NOT METAL_LINALG_SHARED)
|
|
221
|
+
target_link_libraries(metal_linalg INTERFACE ${_ml_frameworks})
|
|
222
|
+
endif ()
|
|
223
|
+
|
|
224
|
+
# =============================================================================
|
|
225
|
+
# Install and export: find_package(MetalLinalg) -> metal_linalg::metal_linalg
|
|
226
|
+
# =============================================================================
|
|
227
|
+
|
|
228
|
+
if (METAL_LINALG_INSTALL)
|
|
229
|
+
install(TARGETS metal_linalg
|
|
230
|
+
EXPORT MetalLinalgTargets
|
|
231
|
+
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
|
|
232
|
+
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
|
|
233
|
+
RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
|
|
234
|
+
)
|
|
235
|
+
install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
|
|
236
|
+
|
|
237
|
+
set(_ml_cmake_dir "${CMAKE_INSTALL_LIBDIR}/cmake/MetalLinalg")
|
|
238
|
+
install(EXPORT MetalLinalgTargets
|
|
239
|
+
FILE MetalLinalgTargets.cmake
|
|
240
|
+
NAMESPACE metal_linalg::
|
|
241
|
+
DESTINATION "${_ml_cmake_dir}"
|
|
242
|
+
)
|
|
243
|
+
include(CMakePackageConfigHelpers)
|
|
244
|
+
configure_package_config_file(cmake/MetalLinalgConfig.cmake.in
|
|
245
|
+
"${CMAKE_CURRENT_BINARY_DIR}/MetalLinalgConfig.cmake"
|
|
246
|
+
INSTALL_DESTINATION "${_ml_cmake_dir}"
|
|
247
|
+
)
|
|
248
|
+
write_basic_package_version_file(
|
|
249
|
+
"${CMAKE_CURRENT_BINARY_DIR}/MetalLinalgConfigVersion.cmake"
|
|
250
|
+
VERSION ${PROJECT_VERSION}
|
|
251
|
+
COMPATIBILITY SameMajorVersion
|
|
252
|
+
)
|
|
253
|
+
install(FILES
|
|
254
|
+
"${CMAKE_CURRENT_BINARY_DIR}/MetalLinalgConfig.cmake"
|
|
255
|
+
"${CMAKE_CURRENT_BINARY_DIR}/MetalLinalgConfigVersion.cmake"
|
|
256
|
+
DESTINATION "${_ml_cmake_dir}"
|
|
257
|
+
)
|
|
258
|
+
endif ()
|
|
259
|
+
|
|
260
|
+
# =============================================================================
|
|
261
|
+
# Tests, benchmarks and tuning harnesses
|
|
262
|
+
# =============================================================================
|
|
263
|
+
|
|
264
|
+
if (METAL_LINALG_BUILD_TESTS)
|
|
265
|
+
include(CTest)
|
|
266
|
+
|
|
267
|
+
function(metal_linalg_add_tool name)
|
|
268
|
+
add_executable(${name} ${ARGN})
|
|
269
|
+
target_link_libraries(${name} PRIVATE metal_linalg mlx)
|
|
270
|
+
endfunction()
|
|
271
|
+
|
|
272
|
+
metal_linalg_add_tool(test_qr tests/test_qr.cpp)
|
|
273
|
+
metal_linalg_add_tool(test_eigh tests/test_eigh.cpp)
|
|
274
|
+
metal_linalg_add_tool(test_svd tests/test_svd.cpp)
|
|
275
|
+
add_test(NAME qr_correctness COMMAND test_qr)
|
|
276
|
+
add_test(NAME eigh_correctness COMMAND test_eigh)
|
|
277
|
+
add_test(NAME svd_correctness COMMAND test_svd)
|
|
278
|
+
|
|
279
|
+
# The buffer core on its own, without MLX, and the C API compiled as C.
|
|
280
|
+
add_executable(test_core tests/test_core.cpp)
|
|
281
|
+
target_link_libraries(test_core PRIVATE metal_linalg)
|
|
282
|
+
add_test(NAME core_correctness COMMAND test_core)
|
|
283
|
+
enable_language(C)
|
|
284
|
+
add_executable(test_c_api tests/test_c_api.c)
|
|
285
|
+
set_target_properties(test_c_api PROPERTIES C_STANDARD 99 C_STANDARD_REQUIRED ON)
|
|
286
|
+
target_link_libraries(test_c_api PRIVATE metal_linalg)
|
|
287
|
+
add_test(NAME c_api COMMAND test_c_api)
|
|
288
|
+
|
|
289
|
+
metal_linalg_add_tool(benchmark_qr benchmarks/benchmark_qr.cpp)
|
|
290
|
+
metal_linalg_add_tool(benchmark_eigh benchmarks/benchmark_eigh.cpp)
|
|
291
|
+
metal_linalg_add_tool(benchmark_svd benchmarks/benchmark_svd.cpp)
|
|
292
|
+
|
|
293
|
+
# Tuning harnesses: run these when porting to a different Apple GPU, since
|
|
294
|
+
# core count and per-core throughput set every routing crossover; see
|
|
295
|
+
# docs/tuning.md.
|
|
296
|
+
metal_linalg_add_tool(sweep_qr tuning/sweep_qr.cpp)
|
|
297
|
+
metal_linalg_add_tool(sweep_eigh tuning/sweep_eigh.cpp)
|
|
298
|
+
metal_linalg_add_tool(sweep_svd tuning/sweep_svd.cpp)
|
|
299
|
+
|
|
300
|
+
metal_linalg_add_tool(probe_occupancy tuning/probe_occupancy.mm)
|
|
301
|
+
target_include_directories(probe_occupancy PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src")
|
|
302
|
+
target_link_libraries(probe_occupancy PRIVATE "-framework Metal" "-framework Foundation")
|
|
303
|
+
endif ()
|
|
304
|
+
|
|
305
|
+
# =============================================================================
|
|
306
|
+
# Examples: small self-checking programs, also run by ctest
|
|
307
|
+
# =============================================================================
|
|
308
|
+
|
|
309
|
+
if (METAL_LINALG_BUILD_EXAMPLES)
|
|
310
|
+
include(CTest)
|
|
311
|
+
foreach (_ex quickstart orthonormal_bases pca nearest_orthogonal routing)
|
|
312
|
+
add_executable(example_${_ex} examples/${_ex}.cpp)
|
|
313
|
+
target_link_libraries(example_${_ex} PRIVATE metal_linalg)
|
|
314
|
+
add_test(NAME example_${_ex} COMMAND example_${_ex})
|
|
315
|
+
endforeach ()
|
|
316
|
+
add_executable(example_objc_quickstart examples/objc_quickstart.mm)
|
|
317
|
+
target_link_libraries(example_objc_quickstart PRIVATE metal_linalg "-framework Foundation")
|
|
318
|
+
add_test(NAME example_objc_quickstart COMMAND example_objc_quickstart)
|
|
319
|
+
endif ()
|
|
320
|
+
|
|
321
|
+
# =============================================================================
|
|
322
|
+
# Python extension: metal_linalg._core (see python/README.md)
|
|
323
|
+
# =============================================================================
|
|
324
|
+
|
|
325
|
+
if (METAL_LINALG_BUILD_PYTHON)
|
|
326
|
+
add_subdirectory(python)
|
|
327
|
+
endif ()
|