Hardware-accelerated QR decomposition for Apple Silicon, built on top of MLX and Metal, by leveraging the GPU to maximise performance.
Exposes a single function, custom_math::qr_accelerated, that accepts a batched MLX array and returns Q and R — automatically routing to the most efficient Metal kernel for the given matrix dimensions and batch size. See benchmark results for GPU vs CPU timings across a range of matrix shapes and batch sizes.
This library depends on the MLX C++ library. On macOS:
brew install mlx
#include "qr.h"
// a: MLX array of shape [M, N] or [B, M, N] (float32 or auto-cast)
// Returns: {Q, R} where Q is [M, K] and R is [K, N], K = min(M, N)
auto [Q, R] = custom_math::qr_accelerated(a);
The function operates on the default MLX device. Set it to GPU before calling:
#include "qr.h"
#include <mlx/mlx.h>
using namespace mlx::core;
// Build a random 4 x 4 matrix
std::vector<float> data = {
1, 2, 3, 4,
5, 6, 7, 8,
9, 10, 11, 12,
13, 14, 15, 16
};
array A(data.begin(), {4, 4}, float32);
set_default_device(Device::gpu);
auto [Q, R] = custom_math::qr_accelerated(A);
eval({Q, R});
// Q: [4, 4] orthogonal matrix
// R: [4, 4] upper-triangular matrix
// A ≈ Q * R
Given a matrix $A \in \mathbb{R}^{M \times N}$, the QR decomposition factors it as:
\[A = QR\]where $K = \min(M, N)$.
Q $\in \mathbb{R}^{M \times K}$ is a matrix with orthonormal columns. That is, for any two columns $q_i$ and $q_j$:
\[q_i^T q_j = \delta_{ij} = \begin{cases} 1 & \text{if } i = j \\ 0 & \text{if } i \neq j \end{cases}\]which can be stated compactly as $Q^T Q = I_K$. The columns of $Q$ form an orthonormal basis for the column space of $A$.
R $\in \mathbb{R}^{K \times N}$ is upper triangular. Every entry strictly below the main diagonal is zero:
\[R_{ij} = 0 \quad \text{for all } i > j\]This is the thin (or reduced) QR decomposition. The full decomposition extends $Q$ to a square $M \times M$ orthogonal matrix, but the thin form is sufficient to reconstruct $A$ and is more compact when $M > N$.
For example, given:
A = [[ 1, 2, 3 ],
[ 4, 5, 6 ],
[ 7, 8, 9 ]]
the thin QR decomposition yields:
Q = [[-0.123, 0.904, 0.408 ], R = [[-8.124, -9.601, -11.078 ],
[-0.492, 0.301, -0.816 ], [ 0.0, 0.905, 1.809 ],
[-0.862, -0.301, 0.408 ]] [ 0.0, 0.0, 0.0 ]]
One can verify $QR = A$ and $Q^T Q = I_3$.
The decomposition is fundamental to solving linear least-squares problems, performing Gram-Schmidt orthogonalisation, and as the core step in the QR algorithm for computing eigenvalues.
Both shaders build $Q$ and $R$ by successively applying Householder reflections. A Householder reflector is an orthogonal matrix of the form:
\[H = I - \tau v v^T, \quad \tau \in \mathbb{R}, \quad v \in \mathbb{R}^M\]chosen so that $H x = \mu e_k$ — i.e. it zeros out every entry of a column vector $x$ below position $k$, leaving a single scalar $\mu$ on the diagonal. The sign of $\mu$ is chosen to avoid catastrophic cancellation:
\[\mu = -\text{sign}(\alpha)\|x\|_2\]where $\alpha = x_k$ is the pivot element. The reflector vector $v$ is then:
\[v_k = 1, \quad v_i = \frac{x_i}{\alpha - \mu} \text{ for } i > k, \quad \tau = \frac{\mu - \alpha}{\mu}\]Applying $K = \min(M, N)$ reflectors in sequence drives $A$ to upper triangular form:
\[H_K \cdots H_2 H_1 A = R \implies A = H_1 H_2 \cdots H_K R = QR\]Since each $H_i$ is orthogonal, their product $Q = H_1 H_2 \cdots H_K$ is also orthogonal. Rather than forming this product one reflector at a time, both shaders use the Compact WY representation to batch the updates.
For a block of $b$ consecutive Householder reflectors, the product can be written as:
\[H_1 H_2 \cdots H_b = I - Y T Y^T\]where $Y \in \mathbb{R}^{M \times b}$ has the $b$ reflector vectors as its columns, and $T \in \mathbb{R}^{b \times b}$ is an upper triangular matrix constructed recursively:
\[T_{jj} = \tau_j, \quad T_{ij} = -\tau_j \sum_{m=i}^{j-1} T_{im} (y_m^T y_j) \quad \text{for } i < j\]This lets a full block update be expressed as a pair of matrix multiplications:
\[A \leftarrow A - Y \bigl( T^T (Y^T A) \bigr)\]which maps directly onto the AMX matrix coprocessor’s 8×8 simdgroup_matrix tiles.
qr_unblocked (single-kernel, block size $b = 16$)This kernel processes the entire matrix in one GPU dispatch. It operates on matrices stored in column-major format to align with AMX load/store strides, and pads dimensions to multiples of 32 (rows) and 16 (columns) to eliminate in-kernel boundary branching.
For each block of $b = 16$ columns starting at column $s$:
Step 1 — Panel factorisation. For each column $k = 0, \ldots, b-1$ within the block:
simd_sum, then inter-SIMD via shared memory).Step 2 — Form T. The $b \times b$ upper triangular matrix $T$ is built in threadgroup memory using the recursive formula above.
Step 3 — Trailing matrix update. Each SIMD group owns a set of 8-column tiles of the trailing submatrix $A[{:}, s+b{:}]$. Using AMX 8×8 tiles, it computes:
\[Z = Y^T A_\text{trail}, \quad Z \leftarrow T Z, \quad A_\text{trail} \leftarrow A_\text{trail} - Y Z\]Step 4 — Q accumulation. The same WY update is applied to $Q$ (initialised to $I$):
\[Q \leftarrow Q - Y \bigl( T (Y^T Q) \bigr)\]Step 5 — Restore diagonal. The stored $\mu$ values are written back to the diagonal of $A$ (overwriting the temporary $v_k = 1$ sentinel placed there during factorisation).
qr_streaming_amx (multi-kernel, block size $b = 32$)For large matrices ($M$ or $N \geq 512$), a single-kernel dispatch causes Q-accumulation to bottleneck on a single shader multiprocessor. The streaming variant splits the computation across four separate kernel dispatches per block, allowing the GPU scheduler to assign the trailing update across all available cores in parallel.
The block size is widened to $b = 32$ to match the SIMD group width, maximising AMX tile utilisation and reducing the number of host-side dispatch iterations.
Kernel 0 — Preprocess. The input is transposed from row-major to column-major and padded with identity blocks. $Q$ is initialised to $I_{M \times M}$.
Kernel 1 — Panel factorisation. Dispatched with 1 threadgroup per matrix. For each column $k$ in the current block, 1024 threads cooperatively compute $\tau_k$ and the normalised reflector vector, then apply it to the remaining $b - k - 1$ panel columns. The $\tau$ values and diagonal elements of $R$ are written to global memory for use by subsequent kernels.
Kernel 2 — T-matrix construction. Dispatched with 1 threadgroup per matrix. Reads the reflector columns from $A$ and $\tau$ from global memory and builds $T \in \mathbb{R}^{32 \times 32}$ in threadgroup memory using the recursive Compact WY formula. The completed $T$ is written to global memory negated (i.e. $-T$ is stored), so that Kernel 3 can use simdgroup_multiply_accumulate (which adds) rather than needing a subtract path.
Kernel 3 — Grid-parallel trailing update. Dispatched with one threadgroup per 32-column tile of the trailing submatrix. Each threadgroup independently computes:
\[\text{Phase 1:} \quad Z^T = A_\text{trail}^T \cdot Y\] \[\text{Phase 2:} \quad Z_\text{final}^T = Z^T \cdot T\] \[\text{Phase 3:} \quad A_\text{trail} \leftarrow A_\text{trail} + Y \cdot Z_\text{final}^T\]$Y$ and $T$ are loaded into threadgroup memory (L1 cache) once per tile and reused across all AMX sweeps. The same kernel is reused for Q-accumulation by setting a flag that redirects the target pointer from $A$ to $Q$.
Kernel 4 — Haar fix. Ensures the output $Q$ is a uniform sample from the Haar measure on $O(M)$ and that $R$ has non-negative diagonal. For each column $k$ where $R_{kk} < 0$, the signs of column $k$ in both $Q$ and $R$ are flipped. If the resulting $\det(Q) < 0$, the final column is negated to enforce $\det(Q) = +1$, placing $Q$ in $SO(M)$.
Two Metal kernels handle different regimes, with a dispatcher that selects between them at runtime:
qr_unblocked — Standard Householder QR in a single kernel dispatch. Used for smaller matrices where the overhead of multi-pass streaming is not worth it.
qr_streaming_amx — Multi-pass panel factorisation with grid-parallel trailing matrix updates, designed to saturate the GPU for large matrices. Factorises column panels of width 32, computes the T-matrix for each WY representation, then launches a grid of threadgroups for the trailing update.
| Condition | Kernel |
|---|---|
max(M, N) >= 512 |
qr_streaming_amx |
max(M, N) < 512 and batch >= 16 |
qr_unblocked |
max(M, N) >= 128 and batch < 16 |
qr_streaming_amx |
max(M, N) < 128 and batch < 16 |
qr_unblocked |
The crossover points reflect two competing costs: single-SM Q-accumulation becomes a bottleneck above ~512, while multi-kernel launch overhead makes streaming slower than unblocked for small matrices with few batch elements.
Both backends cache compiled MTLComputePipelineState objects and recycle GPU memory workspaces on repeated calls with the same shape, so warm invocations avoid both JIT compilation and OS-level buffer allocation.
xcrun, metal, metallib)TODO: exportable library packaging is planned for a future release.
Clone the repository and build in release mode:
git clone https://github.com/cormaccinnsealach/qr-apple-silicon.git
cd qr-apple-silicon
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --target benchmark_qr
./build/benchmark_qr
The benchmark sweeps every configuration automatically:
| Class | Configs (M × N) | Batch sizes |
|---|---|---|
| Small (M < 512) | 64×64, 128×64, 256×128, 512×256 | 100, 500, 1000, 5000, 10000, 15000 |
| Large (M ≥ 512) | 512×512, 1024×512, 5000×5000 | 1, 8, 16, 32 |
Per-config batch limits are set via the max_batch field in SMALL_CONFIGS at the top of tests/benchmark_qr.cpp. Any batch size above this value is skipped for that config and shown as -- in the table. Set max_batch = 0 to run a config at all batch sizes.
// tests/benchmark_qr.cpp
static const std::vector<BenchConfig> SMALL_CONFIGS = {
{ 64, 64},
{128, 64},
{256, 128, 5000}, // skipped for batch > 5000
{512, 256, 1000}, // skipped for batch > 1000 — raise or set to 0 if memory allows
};
Each configuration runs for 5 timed repetitions (2 warmup discarded) and reports GPU time, CPU time (MLX CPU), speedup, and mean Frobenius reconstruction error ||QR - A||_F for both.