GPU compute shaders

Deep dive

How it works

In: a compiled pipeline, bound buffers, a workgroup count and a workgroup size · Out: modified buffers once the queue has finished · State: pipeline and buffers persist; workgroup memory lives only while its group runs · Cost: about 1-3 µs per dispatch when batched (measured), then bandwidth or arithmetic bound

1. Pipeline. WGSL, Metal Shading Language or SPIR-V is compiled to GPU code once; cache the pipeline (pipeline-state caching).
2. Bindings. Buffers bind by index: WebGPU bind groups, Metal setBuffer, Vulkan descriptor sets.
3. Encode. Record pipeline, bindings and dispatch(groups) into a command buffer. Workgroup size is fixed in the shader (@workgroup_size, LocalSize) or, in Metal, at encode time (threadsPerThreadgroup).
4. Run. Each workgroup gets its own var<workgroup> / threadgroup / shared memory; invocations run in subgroups (SIMD-groups: 32 wide on Apple, 4-128 in the WGSL spec). global_id = group_id × size + local_id.
5. Barrier. workgroupBarrier(), threadgroup_barrier() or barrier() orders shared-memory accesses; all invocations must reach it.
6. Combine groups. Write one partial per group, reduce in a second dispatch (or use atomics). Vulkan dispatches may overlap, so add a pipeline barrier; Metal's default encoder runs them in order; WebGPU makes each dispatch of a pass its own usage scope.
7. Submit and wait. Queue submit, then a fence, waitUntilCompleted or mapAsync; readback costs a round trip.

Code

rustc --edition 2021 -O. One dispatch emulated on the CPU: invocations are threads, Barrier is workgroupBarrier; a tree reduction and a tiled matmul as a shader holds them, asserted against references. Removing the first matmul barrier made the asserts fail in my run.

use std::sync::atomic::{AtomicU32, Ordering::Relaxed};
use std::sync::Barrier;

struct Buf(Vec<AtomicU32>); // f32 bits in atomics, so threads can share it safely
impl Buf {
    fn new(v: &[f32]) -> Self { Buf(v.iter().map(|x| AtomicU32::new(x.to_bits())).collect()) }
    fn get(&self, i: usize) -> f32 { f32::from_bits(self.0[i].load(Relaxed)) }
    fn set(&self, i: usize, x: f32) { self.0[i].store(x.to_bits(), Relaxed) }
    fn vec(&self) -> Vec<f32> { (0..self.0.len()).map(|i| self.get(i)).collect() }
}
struct Inv<'a> { group: [usize; 2], local: [usize; 2], groups: [usize; 2],
                 shared: &'a Buf, barrier: &'a Barrier }

/// dispatchWorkgroups(groups) with @workgroup_size(size). Groups never wait for each other: run in turn.
fn dispatch(groups: [usize; 2], size: [usize; 2], shared_len: usize, kernel: impl Fn(&Inv) + Sync) {
    for gy in 0..groups[1] { for gx in 0..groups[0] {
        let shared = Buf::new(&vec![0.0; shared_len]); // fresh var<workgroup> per group
        let barrier = Barrier::new(size[0] * size[1]);
        std::thread::scope(|s| for ly in 0..size[1] { for lx in 0..size[0] {
            let (shared, barrier, kernel) = (&shared, &barrier, &kernel);
            s.spawn(move || kernel(&Inv { group: [gx, gy], local: [lx, ly], groups, shared, barrier }));
        }});
    }}
}

const W: usize = 64; // workgroup size of the reduction
// one partial sum per workgroup: grid-stride loop, then a tree in shared memory
fn reduce(x: &Buf, n: usize, groups: usize) -> Vec<f32> {
    let out = Buf::new(&vec![0.0; groups]);
    dispatch([groups, 1], [W, 1], W, |t| {
        let (lid, stride) = (t.local[0], W * t.groups[0]);
        let mut acc = 0.0;
        let mut i = t.group[0] * W + lid;
        while i < n { acc += x.get(i); i += stride; }
        t.shared.set(lid, acc);
        t.barrier.wait(); // workgroupBarrier()
        let mut s = W / 2;
        while s > 0 {
            if lid < s { t.shared.set(lid, t.shared.get(lid) + t.shared.get(lid + s)); }
            t.barrier.wait(); // outside the `if`: every invocation reaches it
            s /= 2;
        }
        if lid == 0 { out.set(t.group[0], t.shared.get(0)); }
    });
    out.vec()
}

const T: usize = 8; // tile edge; workgroup is T x T
// C (m x n) = A (m x k) * B (k x n), row-major; m, n, k multiples of T
fn matmul(a: &Buf, b: &Buf, m: usize, n: usize, k: usize) -> Vec<f32> {
    let c = Buf::new(&vec![0.0; m * n]);
    dispatch([n / T, m / T], [T, T], 2 * T * T, |t| {
        let (lx, ly) = (t.local[0], t.local[1]);
        let (row, col) = (t.group[1] * T + ly, t.group[0] * T + lx);
        let mut acc = 0.0;
        for p in 0..k / T {
            t.shared.set(ly * T + lx, a.get(row * k + p * T + lx));
            t.shared.set(T * T + ly * T + lx, b.get((p * T + ly) * n + col));
            t.barrier.wait(); // tile complete before anyone reads it
            for i in 0..T { acc += t.shared.get(ly * T + i) * t.shared.get(T * T + i * T + lx); }
            t.barrier.wait(); // everyone done before the next overwrite
        }
        c.set(row * n + col, acc);
    });
    c.vec()
}

fn main() {
    let n = 10_000; // integer values: f32 sums stay exact
    let v: Vec<f32> = (0..n).map(|i| ((i * 37 % 101) as f32) - 50.0).collect();
    let parts = reduce(&Buf::new(&v), n, 7);
    assert_eq!(parts.iter().sum::<f32>(), v.iter().sum::<f32>());
    assert_eq!(parts, reduce(&Buf::new(&v), n, 7)); // fixed order: bit-identical rerun
    let (m, nn, k) = (32, 24, 16);
    let av: Vec<f32> = (0..m * k).map(|i| (i % 7) as f32 - 3.0).collect();
    let bv: Vec<f32> = (0..k * nn).map(|i| (i % 5) as f32 - 2.0).collect();
    let c = matmul(&Buf::new(&av), &Buf::new(&bv), m, nn, k);
    for r in 0..m { for q in 0..nn {
        assert_eq!(c[r * nn + q], (0..k).map(|p| av[r * k + p] * bv[p * nn + q]).sum::<f32>());
    }}
    println!("reduce and tiled matmul match the references");
}

Complexity and constants

• Reduction: n loads and adds, log₂W barrier rounds per group. Tiled GEMM, tile edge T: 2mnk/T global loads against 2mnk naive; shared memory 2T² floats per group.
• Limits (specs): WebGPU defaults maxComputeInvocationsPerWorkgroup 256 (compatibility mode 128), maxComputeWorkgroupStorageSize 16,384 B, maxStorageBufferBindingSize 128 MiB, maxBindGroups 4. Vulkan minimums: maxComputeSharedMemorySize 16,384 B, 128 invocations, 65,535 groups per dimension.
• Measured on this Mac (Apple M4 Pro, Metal, my own untuned kernels, median of 10-20 GPU-timed runs, SIMD width 32): sum of 2²⁶ f32 1.1-1.2 ms, 220-240 GB/s, tree and simd_sum equal. GEMM 2048² f32: naive 35 ms (about 490 GFLOP/s), 16×16 shared-memory tiles 19 ms (about 900), simdgroup_matrix 8×8, no shared memory, 11.6 ms (about 1,480). Dispatch plus wait 70-170 µs; 100-1,000 per command buffer 1.3-3.4 µs each.

Efficiency

• Cost: a reduction is bandwidth-bound: 268 MB in about 1.2 ms (measured, M4 Pro). A GEMM is arithmetic-bound: 17.2 GFLOP in 11.6-35 ms by kernel (measured). A batched launch is 1-3 µs, with a wait 70-170 µs (measured).
• Levers:
  • batch dispatches in one command buffer — 70-170 µs → 1-3 µs each (measured) — many small kernels
  • shared-memory tiling — 35 → 19 ms, 1.8× (measured) — GEMM, convolution, stencils
  • subgroup matrix ops (simdgroup_matrix, Vulkan cooperative matrices) — 19 → 11.6 ms (measured, Metal) — GEMM on GPUs that expose them
  • fuse kernels — one memory pass fewer per op — elementwise chains
• Naive → standard → state of the art:
  • one invocation per output, every load from global memory · 35 ms, 2048² GEMM (measured)
  • shared-memory tiles, two barriers per tile · 19 ms (measured)
  • cooperative matrices over shared-memory tiles, as in llama.cpp's mul_mm.comp · not measured here

Implementation notes

• Query the subgroup size (subgroup_size in WGSL) instead of hardcoding 32.
• Grid-stride loops decouple grid from n and keep neighbours on neighbouring addresses; split a 1D grid across 2D beyond 65,535 groups.
• Metal dispatchThreads takes any grid size; elsewhere guard i < n.

Pitfalls

• A barrier in a branch only some invocations take is a bug: WGSL makes the author responsible for uniformity; Metal requires every thread to reach it, though on Apple silicon exited threads are skipped.
• A float sum depends on group count and order: bit-identical only if both are fixed (reproducible reduction).
• A workgroup size copied from CUDA (1,024) exceeds WebGPU's default 256.
• Waiting after every dispatch costs hundreds of µs (measured).

Variants and alternatives

In production

• MLX: Metal kernels incl. steel/gemm on simdgroup_matrix, reduce, scan, sort, FFT.
• llama.cpp: Metal mul_mm stages tiles in threadgroup memory and uses simdgroup_float8x8; Vulkan mul_mm.comp uses shared tiles and optional cooperative matrices; sum_rows.comp is a barrier tree.
• PyTorch MPS: hand-written .metal kernels beside MPSGraph, which matmul uses in places.
• ONNX Runtime Web (WebGPU execution provider) and WebLLM (WebGPU, TVM runtime): models in the browser.
• Burn: Metal, Vulkan and WebGPU backends built on CubeCL.

Implementations

Sources

• W3C, WebGPU, limits table. spec
• W3C, WebGPU Shading Language. spec
• Apple, Metal Shading Language Specification 4.1. pdf · threadgroup sizes
• Khronos, Vulkan Specification, limits and synchronization. spec · vkCmdDispatch
• Harris, Optimizing Parallel Reduction in CUDA, NVIDIA. slides
• Code read: MLX kernels · llama.cpp Metal · llama.cpp Vulkan · PyTorch MPS · ONNX Runtime Web WebGPU · CubeCL