Linear algebra

In short: MatMul follows NumPy's rules (batch axes broadcast, vectors gain and lose an axis) and calls Gemm, folding a batched x @ W into one product and spreading many small products across threads. Gemm picks its micro-kernel at run time (AVX-512 12×32, AVX2 6×16, or a portable 6×8 SSE / 8×12 NEON tile) and its path from the shape: one row of A streams B through an axpy, up to 32 rows stream B through register tiles, and anything larger runs the Goto/BLIS driver that packs panels of A and B and splits disjoint tiles of C across the thread pool. PackedB packs a constant weight once. GemmRef is the scalar reference every kernel is tested against. The decompositions work in double on a Matrix and convert back to f32 tensors. LU uses partial pivoting, QR uses Householder reflections, Cholesky needs no pivoting, and eigen and SVD use Jacobi rotations. Norms, condition number, rank and pseudo-inverse are all read off the SVD.

Each header links to the /ml-models page it builds on: matrix multiplication, matrix decompositions, the linear algebra toolkit and vectors and dot products. The tensors these functions take and return are described on Tensors and dtypes. The thread pool that runs Gemm is on Errors, memory and threads.

Matrix: the double-precision working type

Tensors are f32, but factorizations lose digits to cancellation, so each decomposition copies its input into a Matrix, works in double, and converts the result back (matrix.h L4–11). Matrix is deliberately plain: row-major std::vector<double>, no views, no expression templates.

MatMul: NumPy semantics

matmul.h follows NumPy's matmul. The last two axes are the matrix axes, and any leading batch axes broadcast against each other (see broadcasting). A 1-D left operand is treated as a row vector and a 1-D right operand as a column vector, and the added axis is removed from the result. Inputs may be strided, so multiplying by a Transpose view costs no copy.

The first three run through BatchedMatMul (matmul.cc L36–118). It checks f32 and the inner dimension, unsqueezes vector operands, broadcasts the batch shapes, and allocates the output. Then it picks one of three ways to run the batch:

1. Fold the batch into the rows. When B is a single 2-D matrix shared by every batch (the x @ W of a linear layer) and A's rows are evenly spaced across all its batch axes, the batch is one tall matrix of batch·m rows, so one Gemm call does the whole product and B is packed once instead of once per batch (L67–82).
2. Many small products in parallel. Otherwise internal::ForEachRow<3> collects every batch's offsets with broadcast strides, and a MatrixView is built for each operand straight from the tensor's strides, so a broadcast operand is read in place rather than copied. When there are several batches and each product has fewer than 2²² multiply-adds (attention heads), the batches are split across the default pool with a grain of about 2²⁰ multiply-adds per task, and each product then runs on one thread, because a parallel loop nested in a running chunk executes inline (L84–110).
3. One product at a time, each on the threaded Gemm, for large products or a reference run (L111–113).

Finally the added vector axes are squeezed away.

GEMM

gemm.h computes C = α·A·B + β·C on raw f32 memory. Operands are MatrixViews: a pointer, a shape, and a row and column stride. Swapping the strides is a transpose, so Transposed() is free. When β is 0, C is overwritten without being read, so it may start as garbage or NaN. A mismatched inner dimension, or an ldc smaller than C's columns, is a RML_CHECK failure (gemm.cc L19–22).

GemmRef

The reference (gemm.cc L34–47) is the textbook triple loop. It accumulates each dot product in double and then applies α and β. It is slow and exact enough that the tolerance on every faster path reflects only that path's own rounding.

Kernels and dispatch

Gemm is internal::GemmWithKernel(kAuto, …) (gemm.cc L656–659). The kernel enum is GemmKernel { kAuto, kPortable, kAvx2, kAvx512 } (gemm.h L83–96). GemmKernelSupported asks the CPU through __builtin_cpu_supports (AVX2 needs avx2 and fma, AVX-512 needs avx512f; neither exists off x86-64), and BestGemmKernel caches the widest supported one in a function-local static (gemm.cc L577–603). The x86 kernels are compiled with __attribute__((target(...))), so the baseline build still runs on any x86-64. Tests and benchmarks can force a kernel through GemmWithKernel, which RML_CHECKs that the CPU can run it.

GemmWithKernel (gemm.cc L605–650) returns early when m or n is 0 and only scales C by β when k is 0. Then the shape of A picks the path:

The blocked kernel

The blocked driver follows the Goto/BLIS scheme, which the comment at gemm.cc L49–74 lays out. C is cut into kMc×kNc blocks and the shared dimension k into kKc slices (L77–79). Each block of A is packed into Mr-row panels and each slice of B into Nr-column panels. Both are laid out in the order the micro-kernel reads them, and zero-padded at the edges so the kernel never branches on a partial tile.

Each micro-kernel holds an Mr×Nr tile of C in vector registers. For each k step it does one rank-1 update: Nr/width vector loads of B, Mr broadcasts of A, and Mr·Nr/width multiply-adds. The tile size is chosen to fill the register file:

The packers (gemm.cc L190–224) copy through the views' strides, which is how transposed operands are handled. PackB takes a memcpy fast path when a panel is full and B's columns are contiguous. StoreTile (L226–242) writes α·tile plus β·C on the first k-slice, and adds α·tile on every later slice, so β is applied exactly once.

PackedB: packing a weight once

For more than 32 rows the blocked driver repacks B on every call. A model's constant weight never changes, so PackedB (gemm.h L55–81) packs it once for the kernel chosen when it is made: every k-slice's panels, laid out at offset p0·padded_n as the driver expects (PackAllB, gemm.cc L244–261). GemmPackedB (L683–701) passes the packed buffer to Blocked, or falls back to GemmWithKernel on the original view when A has 32 rows or fewer (where B is streamed, not packed) or k or n is 0. That is why the original B must stay alive and unchanged. GemmPacksB(rows) (L661) says whether packing is worth it.

Its caller is the graph executor's op compiler: a matmul node whose second input is a constant 2-D f32 weight, with enough rows that GemmPacksB holds, gets a kernel that calls MatMulPacked on a PackedB made at load (rml/graph/ops.cc L929–944). rml measured it in two interleaved runs on four threads: torch_lstm_b64 went from 53 / 45 ms to 33 / 37 ms and torch_mlp_b64 from 1.05 / 1.39 ms to 0.82 / 0.82 ms (kernels.md L75–78, same Xeon container).

Threading

For each k-slice the blocked driver (gemm.cc L307–347) runs up to two parallel loops:

1. Unless B was packed ahead of time, B's slice is packed, split by column chunk. Packing is a copy, so it is threaded only when the slice has at least 2¹⁸ elements (kParallelPackElements); otherwise one task packs it all (L310–319). The packed buffer is a std::vector sized for one slice, allocated once per call.
2. The product runs over (row block, column chunk) pairs, which own disjoint tiles of C, so no synchronization is needed (L320–346). Each task packs its A block into a 64-byte-aligned stack buffer, and re-packs only when the row block changes. The cost of re-packing A is 1/kNc of the pair's multiply-adds.

A task's column chunk need not be the kNc = 192 columns B was packed in, because packed panels are contiguous. UnitColumns keeps kNc when there are at least four pairs per thread, so a chunk's B stays in L2. With fewer, a fixed width leaves threads idle: 64 × 1024 is six kNc chunks on four threads, two rounds for 5.3 chunks of work. It then tries every width up to 4·kNc, in NR-column panels, and keeps the one whose slowest thread computes the fewest panels. On a packed 64 × 256 × 1024 weight (BM_GemmPackedB) that takes 104 µs on four threads against 286 on one, and in two interleaved runs on four threads torch_rnn_b64 went from 10.5 / 9.8 ms to 5.2 / 5.0 ms and torch_lstm_b64 from 17.5 / 16.3 ms to 15.6 / 13.0 ms (kernels.md L67–74).

The grain is max(1, kMinTaskMacs / unit_flops), where a unit is one block's multiply-adds for this slice and kMinTaskMacs is 2²⁰. Waking the pool's workers costs tens of microseconds, so every task gets about a million multiply-adds and small products run on the calling thread (L81–86). The streaming paths use the same floor. Because ParallelFor's chunk boundaries depend only on the range, grain and thread count (chunking), and each tile of C is summed by one thread in a fixed order, the result does not change from run to run.

The benchmark

gemm_benchmark.cc is a Google Benchmark target (rml_cc_benchmark in BUILD.bazel L56). It fills A and B from a fixed-seed mt19937 and reports GFLOP/s, counting 2·m·k·n flops per product (L25–29), or GB/s of B read for the matrix-vector case.

rml's docs/kernels.md records these results on its Linux x86_64 development container (4 vCPUs, Xeon at 2.8 GHz with AVX-512, GCC 13 -O3), which is not the Apple Silicon target:

The same doc puts the AVX-512 kernel at 6–7× the SSE-only portable kernel it replaced (16.5 and 60 GFLOP/s for 1024³ on 1 and 4 threads). For the streaming path it gives 80 GFLOP/s for 8 × 256 × 1024 on one core, and 88–89 GFLOP/s at 16 and 32 rows of 864 × 303 against 41–60 packed; at 64 rows packing wins again (72 against 67) (kernels.md L48–63). Run it with bazel run --config=opt //rml/linalg:gemm_benchmark.

LU with partial pivoting

lu.h factors P·A = L·U. At each column, the row with the largest |entry| becomes the pivot, so every multiplier in L is at most 1 in magnitude, which keeps rounding error from growing. Once A is factored, each right-hand side costs two triangular solves, O(n²). The determinant is the product of U's diagonal times the sign of the permutation.

LuFactor (lu.cc L18–50) always succeeds. It stores the result packed into one matrix, with U on and above the diagonal and L's multipliers below it (LuFactors). A pivot no larger than MaxAbs · n · 64 · ε marks the matrix singular, but elimination carries on. LuSolve (L52–72) then returns InvalidArgument for a singular matrix, and Determinant returns 0.

The tensor API (lu.cc L74–115) is LuDecompose (explicit L, U and the permutation), Determinant, Solve (b of shape [n] or [n, k]) and Inverse. The last two return Expected and propagate the singular error with RML_ASSIGN_OR_RETURN.

QR and least squares

qr.h computes A = Q·R by Householder reflections. Column j of R is made by one reflection H = I − 2vvᵀ, the mirror that sends the rest of column j onto the first axis without touching earlier columns. Reflections are orthogonal, so they never amplify rounding error. That is why QR, not the normal equations AᵀA·x = Aᵀb (which square the condition number), is the way to fit least squares.

• QrFactor (qr.cc L34–77) returns the reduced form: Q is m×k with orthonormal columns and R is k×n, with k = min(m, n). Each reflection targets −sign(x₀)·‖x‖, so forming v = x − αe₁ never cancels (L50–52). A zero column is skipped.
• Q is built by applying the stored reflectors to the first k columns of the identity, in reverse order (L62–68). Each column is then signed so R's diagonal is non-negative, which makes the factorization unique for a full-rank A (L70–76).
• LeastSquares(a, b) (L85–108) returns the x minimizing ‖a·x − b‖₂ by back-substituting R·x = Qᵀb. It returns InvalidArgument when a is wider than tall, or when a diagonal entry of R is no larger than MaxAbs(R) · m · 64 · ε (rank deficient). For those cases, use PseudoInverse (below).

Cholesky

A symmetric positive definite matrix factors uniquely as A = L·Lᵀ, a kind of matrix square root (cholesky.h L4–12). It costs half as much as LU and needs no pivoting. It fails exactly when A is not positive definite, which makes it both the fastest solver for normal equations and covariance matrices and the standard test for definiteness. Only the lower triangle is read.

CholeskyFactor (cholesky.cc L13–29) fills L column by column. It returns InvalidArgument as soon as a diagonal value is not positive, and the check is written !(diagonal > 0) so a NaN fails it too. CholeskySolve (L31–49) solves L·y = b, then Lᵀ·x = y. The tensor wrappers Cholesky and CholeskySolve follow at L51–64.

Symmetric eigendecomposition

A symmetric A has real eigenvalues and an orthonormal basis of eigenvectors, A = V·diag(w)·Vᵀ. eigen.h finds them with the cyclic Jacobi method. Each step picks an off-diagonal pair (p, q) and rotates rows and columns p and q by the angle that zeroes a_pq. A rotation can refill entries zeroed earlier, but the off-diagonal sum of squares falls with every step, and convergence is quadratic at the end. The product of the rotations is V. The header notes it is slower than tridiagonal QR for large n, but simple and accurate to full relative precision even for tiny eigenvalues.

• SymmetricEigenFactor (eigen.cc L32–95) reads only the lower triangle. It sweeps over every pair until the off-diagonal sum of squares is below ε² times the matrix's total, or 100 sweeps (kMaxSweeps) have run.
• The rotation solves cot 2φ = θ with θ = (a_qq − a_pp) / 2a_pq, taking the smaller root for t = tan φ so |φ| ≤ π/4 (L51–56).
• Eigenvalues come back ascending. Each eigenvector is signed so its largest-magnitude entry is positive, which makes the output deterministic (L79–93).

Worked example: for [[2, 1], [1, 2]], θ = 0, so t = 1 and φ = 45°. One rotation diagonalizes it, giving eigenvalues 1 and 3 with eigenvectors (1, −1)/√2 and (1, 1)/√2. This is the KnownSpectrum test.

SVD

Every m×n matrix factors as A = U·diag(s)·Vᵀ, with s ≥ 0 sorted descending. A rotates (Vᵀ), stretches along the axes (s), and rotates again (U). svd.h uses one-sided (Hestenes) Jacobi, which orthogonalizes A's columns directly. Rotating columns p and q until they are orthogonal is a Jacobi step on AᵀA without ever forming AᵀA, so no precision is lost to squaring. Once all columns are orthogonal, their norms are the singular values, the normalized columns are U, and the accumulated rotations are V.

• TallSvd (svd.cc L61–118) sweeps until no pair needs rotating: a pair is skipped once |γ| ≤ ε·√(αβ), where α and β are the columns' squared norms and γ their dot product. It uses the same angle as the symmetric eigensolver.
• A singular value no larger than s_max · m · ε counts as zero (L98). Its U column cannot be normalized, so CompleteColumn (L38–59) fills it with a unit vector orthogonal to the rest, by Gram-Schmidt on the standard basis. U stays orthonormal even for a rank-deficient A.
• A wide matrix is handled by factoring its transpose and swapping U and V (L122–139). The result is the reduced SVD with k = min(m, n), and each pair is signed so V's largest-magnitude entry is positive.
• SvdDecompose returns u [m, k], s [k] and vt [k, n]; SingularValues returns only s.

Norms, condition number, rank, pseudo-inverse

norms.h measures size. The induced matrix norms measure the most a matrix can stretch a vector. The condition number s_max / s_min bounds how much relative error in b can grow in the solution of A·x = b.

The rank tolerance uses f32 epsilon even though the SVD runs in double, because the input was f32 and differences below its precision are noise (norms.cc L17–21).

Tests

How these run under sanitizers is on Testing.