metal-linalg for PyTorch

QR, symmetric eigendecomposition and SVD for batches of matrices on Apple GPUs, on PyTorch tensors: on an M5 Pro, 4.9-30x faster than torch.linalg for the same calls, against whichever of its CPU and MPS paths is quicker (the measurements).

import torch
import metal_linalg_torch as mlt

a = torch.randn(1000, 64, 32, device="mps")   # a batch of 1000 matrices; "cpu" works too

Q, R = mlt.qr(a)                     # Q (1000, 64, 32), R (1000, 32, 32)
U, S, Vh = mlt.svd(a)                # thin: U (1000, 64, 32), S (1000, 32), Vh (1000, 32, 32)
L, V = mlt.eigh(a.mT @ a)            # L ascending (1000, 32), V (1000, 32, 32)

The functions take the arguments of their torch.linalg namesakes and return the same result types (Q, R = ..., result.eigenvalues, ...), on the device of their input. Each call is routed to the fastest Metal kernel for its shape and batch, or to LAPACK on the CPU, by a policy measured on the Mac it runs on:

mlt.device_name()                      # 'Apple M5 Pro'
mlt.eigh_policy_source()               # 'tuned:Apple M5 Pro'
mlt.eigh_backend(32, 4096)             # 'ql': 4096 matrices of 32x32 go to the GPU
mlt.svd_backend(4096, 4096)            # 'bidiag': GPU bidiagonalization, LAPACK's solver
mlt.svdvals_backend(4096, 4096)        # 'band': the two-stage reduction, values alone
mlt.set_eigh_policy(gpu_min_batch=1)   # override the measured policy

The kernels, the routing and the measurements behind them are described in the main README.

Installing

pip install metal-linalg-torch

On an Apple Silicon Mac with macOS 14 or later, Python 3.10 or later and PyTorch 2.4 or later. This is the PyTorch package; for MLX arrays install metal-linalg instead. The two are independent: this one does not install or load MLX.

One wheel serves every PyTorch and every Python: the package calls the library's C API (on plain float buffers) and is not compiled against either, so upgrading torch never breaks it.

To build from source instead (needs Xcode's command line tools):

pip install ./python-torch                                                      # in a clone
pip install "git+https://github.com/c0rmac/metal-linalg.git#subdirectory=python-torch"

Functions

A is [..., M, N] with any number of batch dimensions, on "cpu" or "mps". What differs from torch.linalg:

• float32. The kernels compute in float32. Other real dtypes (float64, float16, bfloat16, integers) are converted, and results are float32; complex input raises TypeError.
• Thin SVD factors only. svd defaults to full_matrices=False (torch's default is True), and full_matrices=True on a non-square matrix raises NotImplementedError. QR has all three of torch's modes (since 2.17.0); mode="r" is differentiable here (Q is computed for the gradient when A requires grad), where torch's raises.
• No exceptions for bad matrices. A matrix with a NaN or an infinity gives NaN results for that matrix alone, leaving the rest of its batch intact, where torch raises. cholesky is the exception: it raises for a matrix that is not positive definite, as torch's does.

Autograd and torch.compile

Underneath, the functions are custom operators, torch.ops.metal_linalg.{qr, eigh, eigvalsh, svd, svdvals, cholesky, lu_factor, solve, inv, solve_triangular}, with fake implementations and the backward formulas torch.linalg uses:

A = torch.randn(64, 32, 32, device="mps", requires_grad=True)
S = mlt.svdvals(A)
S.sum().backward()                       # A.grad: the gradient of the nuclear norm

f = torch.compile(lambda x: mlt.eigh(x @ x.mT).eigenvalues.sum())
f(torch.randn(8, 16, 16))

As in torch, the gradients of eigh and svd are defined only for distinct eigenvalues or singular values (they divide by their differences), and are not unique for a loss that depends on the signs of the vectors. Second derivatives are not supported. eigvalsh and svdvals compute the vectors as well when their input requires grad, since the gradient needs them.

Performance

Against torch.linalg on an M5 Pro with PyTorch 2.13 (conda-forge's, its CPU LAPACK from Accelerate) and metal-linalg 2.17 (each cell the best of ten calls over two to four runs of benchmarks/benchmark_torch.py; the same tensors on MPS for torch's MPS path and for this package, which uses them in place, see MPS tensors):

It is ahead on every row. Of these calls PyTorch 2.13 runs only QR on the GPU for MPS tensors, and this is 7x faster there (the QR batch runs on the library's Householder kernels); its SVD takes as long on MPS as on the CPU, and eigh, eigvalsh and svdvals have no MPS kernels and go through its CPU fallback. Against those, 13-30x for the other batches of small matrices (the SVD of 256 matrices of 128×64 is a QR and then golub_kahan on the GPU; it and the SVD of 4096 of 32×32 run on the GPU and the CPU at once), 4.9x for eigh of one 2048×2048 and 10x for the SVD of one 4096×4096 with its vectors, and 10-14x for its eigenvalues or singular values alone (the last three by a two-stage reduction). Which backend a shape gets on your Mac: mlt.svd_backend(m, n, batch) and its siblings.

The functions new in 2.18.0, where PyTorch 2.13 has MPS kernels of its own (python benchmarks/benchmark_torch.py --new, best of ten calls, two runs):

Ahead of torch's faster path on every row, by 1.8-3.9x (the batch of 32×32 Cholesky factorizations is on this library's CPU path, whose spotrf calls are faster than torch's; the rest on the GPU).

MPS tensors

An MPS tensor is used in place. On Apple Silicon PyTorch keeps MPS tensors in Metal buffers in shared storage, which the CPU can address too: the library reads its input there and writes its results into new MPS tensors, with no copy to the CPU and back. Its kernels run on a Metal command queue of its own, so a call first waits for the work PyTorch has queued on the GPU (torch.mps.synchronize()), which may still be writing the input, and returns when its results are written. mlt.mps_in_place() says whether this applies; where it does not (a torch that keeps MPS tensors in private storage), or with METAL_LINALG_TORCH_MPS_COPY=1, an MPS tensor is copied to the CPU and the results back. Since 2.16 the tensors' own Metal buffers are also handed to the library for the call (metal_linalg_know_buffer), so that its GPU kernels use them rather than wrapping the memory in new buffers, whose pages the GPU maps on first use.

Up to 2.13 every MPS call made those copies. They cost most where the decomposition costs least: 9-12x for a handful of matrices, 1.2-1.8x for large batches, a few percent for one large matrix (M5 Pro, PyTorch 2.13, best of five with the two alternating, benchmarks/benchmark_torch.py --mps-ab):

A CPU tensor is used in place too, and its results stay on the CPU, though the work still runs on the GPU where that is faster.

Calls are thread-safe; they are serialised, and release the GIL.

Calibration

Each Mac's routing comes from measurements of that Mac. On one that has none (or older ones), the import raises a CalibrationWarning once per decomposition, and mlt.calibration_status() says where each stands. The library works either way, with settings estimated from a measured Mac and published benchmarks, which lean toward the CPU; measuring the Mac takes one command and improves it for everyone with that Mac: see which Macs are measured and how to contribute. METAL_LINALG_NO_CALIBRATION_NOTICE=1 silences the warning.