II.5: Metal GPU backend

The CPU backends are about as fast as a CPU gets for this: SIMD packs 16 floats into one instruction (II.3), threads spread the work over every core (II.4). A ten-core laptop CPU has, very roughly, ten powerful arithmetic units running at once.

A GPU has thousands of small ones. Matmul, where every output element is an independent dot product, is the canonical workload they were built for. This chapter writes a fourth Backend, Metal, that runs matmul on the GPU through Apple's Metal API: a compute kernel written in Metal Shading Language, command buffers to launch it, and Apple Silicon's unified memory so the CPU and GPU share the same bytes without copying.

Like every backend before it, Metal changes only where matmul runs, not what it computes. Like Parallel, it uses a size threshold to decide when the GPU is worth the trip and when to fall back to the CPU.

What "offload to the GPU" actually means

For a generalist, here's the model. The GPU is a separate processor. It does not share the CPU's instruction stream or its program. To make it do work you must, every time:

Make sure the data the GPU needs lives in memory the GPU can read. On a discrete graphics card this means copying data across the PCIe bus into the card's own memory, which is slow. On Apple Silicon the CPU and GPU share one physical pool of RAM (unified memory), so "making the data available" can be free, which matters a lot here.
Tell the GPU which program to run. A GPU program is called a kernel (or shader). You write it in a small C-like language (for Metal, Metal Shading Language (MSL)) and the driver compiles it for the GPU.
Launch the kernel over a grid of threads. The GPU runs the same kernel on thousands of threads simultaneously, each with a different (x, y) coordinate. For matmul, you launch one thread per output element; each thread computes its own C[row][col].
Wait for it to finish, then read the results back.

Steps 1 and 4 (getting data to and from the GPU, and the round-trip of asking it to do something) are latency. They cost a fixed amount of time no matter how small the actual computation is. The GPU only pays off when step 3, the kernel, is doing enough work to dwarf that fixed overhead. That trade-off drives every design decision in this chapter, and it's the same trade-off as II.4's MIN_ROWS_FOR_PARALLEL threshold, just with a bigger fixed cost.

Three crates for talking to Metal

We don't reimplement Metal; we bind to Apple's system framework through three crates:

Cargo.tomlTOML

[dependencies]
regex = "1"
rayon = "1"
mtl-gpu = "1.0"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
mtl-foundation = "1.0.1"
mtl-sys = "1.0.1"

mtl-sys is the raw, unsafe layer: the Objective-C messaging machinery (Metal's API is Objective-C under the hood). mtl-foundation and mtl-gpu are safer Rust wrappers over the device, command queue, buffers, and pipelines. We mostly use mtl-gpu, dropping to mtl-sys once for one operation the wrapper doesn't expose.

The Metal code lives in a new submodule, src/backend/metal/, split three ways: shaders.rs (the kernel source), context.rs (device and dispatch plumbing), and backend.rs (the Backend impl). We'll take them in that order.

The matmul kernel

src/backend/metal/shaders.rs holds the GPU program as a Rust string constant, MSL source that the Metal driver compiles at startup:

src/backend/metal/shaders.rsRUST

pub const SHADERS: &str = r#"
#include <metal_stdlib>
using namespace metal;
 
constant uint TS = 16;
 
kernel void matmul_fp32_fp32(
    device const float* a [[buffer(0)]],
    device const float* b [[buffer(1)]],
    device       float* c [[buffer(2)]],
    constant     uint&  m [[buffer(3)]],
    constant     uint&  n [[buffer(4)]],
    constant     uint&  p [[buffer(5)]],
    uint2 gid   [[thread_position_in_grid]],
    uint2 lid   [[thread_position_in_threadgroup]]
) {
    threadgroup float tA[16][16];
    threadgroup float tB[16][16];
 
    uint row = gid.x;
    uint col = gid.y;
    float sum = 0.0f;
 
    uint num_tiles = (n + TS - 1) / TS;
 
    for (uint t = 0; t < num_tiles; t++) {
        uint a_col = t * TS + lid.y;
        uint b_row = t * TS + lid.x;
 
        tA[lid.x][lid.y] = (row < m && a_col < n) ? a[row * n + a_col] : 0.0f;
        tB[lid.x][lid.y] = (b_row < n && col < p) ? b[b_row * p + col] : 0.0f;
 
        threadgroup_barrier(mem_flags::mem_threadgroup);
 
        for (uint j = 0; j < TS; j++) {
            sum += tA[lid.x][j] * tB[j][lid.y];
        }
 
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
 
    if (row < m && col < p) {
        c[row * p + col] = sum;
    }
}
"#;

This computes C = A · B with A being m × n and B being n × p. The kernel keyword marks matmul_fp32_fp32 as a GPU entry point. Its parameters are bound to numbered slots: [[buffer(0)]] through [[buffer(5)]] are how Rust passes the three matrices and the three dimensions in. [[thread_position_in_grid]] and [[thread_position_in_threadgroup]] are filled in by the hardware. Every thread sees a different gid, its coordinate in the launch grid.

The basic idea: one GPU thread per output element. Thread at grid position (row, col) computes C[row][col]. With one thread per element, the GPU can have thousands of output elements in flight at once.

The kernel is a tiled matmul rather than the naive version, and the reason is the same memory story as everywhere else in Act 2. The naive approach has each thread stream a full row of A and a full column of B straight from main memory, and adjacent threads re-read overlapping data, hammering the slow path. So the kernel works in 16×16 tiles. The threads in a 16 × 16 threadgroup cooperate: each thread loads one element of A and one of B into the small, fast threadgroup-shared arrays tA and tB, all 256 threads doing it in parallel. A threadgroup_barrier makes every thread wait until the tile is fully loaded. Then each thread does 16 multiply-adds reading only from the fast shared tiles. Slide to the next tile and repeat. Each value of A and B is read from main memory once per tile instead of once per thread, the classic GPU matmul optimization. The (row < m && ...) guards handle matrices whose dimensions aren't exact multiples of 16, zero-padding the edge tiles.

The Metal context

src/backend/metal/context.rs owns the GPU connection (the device, the command queue, the compiled kernel) and the function that dispatches a matmul. It opens with kernel registration and two helpers: