Reducing Register Pressure in WGSL Spatial Kernels

The exact sub-problem here is the register-file term of the occupancy min: the scheduler co-residents at most workgroups, where is the register file per compute unit and is the registers each invocation needs to hold its live values. WGSL never names a register, so this ceiling is invisible in source — but a spatial kernel that keeps a whole neighbourhood window in a var<function> array, or that carries a dozen live accumulators across a long loop, quietly inflates until the register term becomes the limiter and co-residency collapses. Worse, once exceeds what the physical file holds, the compiler spills to VRAM and every access that was a register read becomes a memory round-trip. This page is the register-side companion to workgroup occupancy optimization for spatial kernels: it shows the WGSL patterns that shrink — fewer live temporaries, no large local arrays, recompute over store — with before-and-after kernels whose only difference is register pressure.

Runnable reference implementation

The single kernel below is a separable-Gaussian density smooth, the second pass of a heatmap pipeline. The first version is the natural, register-hungry way to write it; the second is the register-lean rewrite. Both compute the identical result, but the lean form holds a handful of live values where the heavy form holds an entire window plus a running index, so the register term stops bounding co-residency and occupancy climbs.

wgsl
@group(0) @binding(0) var<storage, read>       density: array<f32>;
@group(0) @binding(1) var<storage, read_write> out:     array<f32>;

const GRID_W: u32 = 2048u;
const GRID_H: u32 = 2048u;
const RADIUS: i32 = 4;
const SIGMA: f32 = 2.0;

fn idx(x: i32, y: i32) -> u32 { return u32(y) * GRID_W + u32(x); }

// -------- BEFORE: register-hungry --------
// Materializes the whole 9-wide tap window and a precomputed weight table into
// local arrays. That is 9 live f32 samples + 9 live f32 weights + indices, all
// live across the accumulation loop — r is large, so few workgroups co-resident.
@compute @workgroup_size(64)
fn smooth_heavy(@builtin(global_invocation_id) gid: vec3<u32>) {
    let x = i32(gid.x % GRID_W);
    let y = i32(gid.x / GRID_W);
    if (gid.x >= GRID_W * GRID_H) { return; }

    var samples: array<f32, 9>;          // 9 live values
    var weights: array<f32, 9>;          // 9 more live values
    var k = 0;
    for (var d = -RADIUS; d <= RADIUS; d++) {
        let sx = clamp(x + d, 0, i32(GRID_W) - 1);
        samples[k] = density[idx(sx, y)];
        weights[k] = exp(-f32(d * d) / (2.0 * SIGMA * SIGMA));
        k++;
    }
    var acc = 0.0;
    var norm = 0.0;
    for (var j = 0; j < 9; j++) { acc += samples[j] * weights[j]; norm += weights[j]; }
    out[idx(x, y)] = acc / norm;
}

// -------- AFTER: register-lean --------
// Fuses sampling, weighting, and accumulation into one pass. Only acc, norm,
// and the loop scalar are live — no local arrays, so r is tiny and many more
// workgroups fit per compute unit. The weight is recomputed, not stored.
@compute @workgroup_size(64)
fn smooth_lean(@builtin(global_invocation_id) gid: vec3<u32>) {
    let x = i32(gid.x % GRID_W);
    let y = i32(gid.x / GRID_W);
    if (gid.x >= GRID_W * GRID_H) { return; }

    var acc = 0.0;
    var norm = 0.0;
    let inv2s2 = 1.0 / (2.0 * SIGMA * SIGMA);       // hoisted loop-invariant scalar
    for (var d = -RADIUS; d <= RADIUS; d++) {
        let sx = clamp(x + d, 0, i32(GRID_W) - 1);
        let w = exp(-f32(d * d) * inv2s2);          // recompute per tap, do not store
        acc += density[idx(sx, y)] * w;
        norm += w;
    }
    out[idx(x, y)] = acc / norm;
}

The lean kernel keeps three live scalars against the heavy kernel’s eighteen array slots plus its index, so on a register-limited part its occupancy — and therefore its throughput on a memory-bound smooth — can rise by a large multiple. Confirm the gain with the sweep in choosing workgroup size for GPU occupancy, because register pressure and workgroup size are two terms in the same occupancy min and move the same curve.

Because WebGPU exposes no register counter, register pressure is diagnosed by its fingerprint on a size sweep rather than read directly. Run the same kernel across candidate workgroup sizes and watch where throughput turns over: a register-limited kernel speeds up to a point and then slows down as the size grows, because a wider workgroup multiplies the per-invocation register demand r · |wg| until fewer workgroups co-resident and, past the physical file, values spill to VRAM. A kernel that is not register-limited keeps climbing or plateaus over the same range. So the test for whether a rewrite is worth doing is to sweep the heavy variant first: if its curve peaks early and falls, register pressure is the limiter and the lean rewrite will pay off; if the curve is flat, the register file was never the binding term and the effort belongs elsewhere. This is why every technique in this guide is directional until measured — the same source change that frees three workgroups of co-residency on a narrow mobile part may do nothing on a desktop with a far larger register file.

A useful mental model is that each invocation carries a working set of live values, and the compiler must fit every resident invocation’s working set into one physical register file at once. Anything that enlarges the working set — a materialized window, a wide accumulator bank, a value kept live across a long loop merely because it was declared early — shrinks the number of invocations that fit. The rewrites below all attack the working set from a different angle: recompute-over-store removes stored values, loop fusion removes materialized intermediates, and narrow live ranges shorten how long each value occupies its slot. None of them change the result the kernel produces; they change only how much silicon each invocation reserves while producing it.

Parameter and technique reference

Technique What it removes Guidance for spatial kernels
Recompute over store A local array of precomputed values Recomputing an exp weight per tap is cheaper than the occupancy lost to a live array; the kernel is memory-bound, so spare ALU is free.
Fuse loops A materialized intermediate buffer in registers Merge the sample-gather and accumulate loops into one so no window array is ever live.
Narrow live ranges Temporaries that outlive their use Declare each let at first use inside the innermost scope; a value live across a long loop pins a register for its whole span.
Hoist invariants Repeated sub-expressions Compute 1/(2σ²) once before the loop — one live scalar replaces a repeated divide, and it does not raise peak pressure.
Avoid large var<function> arrays Per-invocation arrays the compiler cannot keep in registers Any array the compiler cannot prove small spills to VRAM; keep windows implicit, indexed straight from storage.
Prefer vec4 over four scalars Four separately live scalars Packing four channels into one vec4<f32> lets the compiler allocate one vector register slot instead of four scalar ones.
Split fat kernels A single kernel holding many roles’ state at once Two chained passes each with low pressure beat one pass that keeps both working sets live.

The register file is not exposed to WebGPU, so treat these as directional and confirm each with a measured dispatch. The normative rules on WGSL variable scope and var<function> storage are in the WGSL specification; the practical effect on co-residency is developed in workgroup occupancy optimization for spatial kernels.

Failure modes

  • Register spilling silently converts reads to VRAM traffic. A kernel with a large local array compiles and runs correctly but far slower, because the array spilled out of the register file. Detection: the timestamp sweep shows throughput falling as workgroup size grows, where a lean variant keeps climbing. Fix: eliminate the local array (fuse loops, recompute), then re-measure.
  • Recompute made a compute-bound kernel worse. Trading storage for arithmetic backfires when the kernel was already ALU-bound, not memory-bound. Detection: the lean rewrite is no faster or slightly slower, and occupancy was never the limiter. Fix: profile first; apply recompute-over-store only to memory-bound spatial kernels where spare ALU cycles are genuinely idle.
  • A vec4 pack broke coalescing. Packing scalars into a vector changed the storage layout so reads are no longer contiguous, costing more bandwidth than the register saved. Detection: memory throughput drops even as occupancy rises. Fix: keep the storage buffer a packed Structure-of-Arrays per the memory alignment for spatial data buffers rules, and pack only in registers, not in the buffer.
  • Splitting a kernel added a costly round-trip. Breaking one pass into two moved a register-resident intermediate into a VRAM buffer, and the extra read and write outweighed the occupancy gain. Detection: total time across both passes exceeds the fused kernel despite each pass showing higher occupancy. Fix: split only when the intermediate is small or already needed downstream; otherwise reduce pressure within the single kernel.

Backend and Python interop note

Register pressure is a property of the compiled shader, so it is identical whether the kernel runs in a browser or in a headless wgpu-py server driving the same WGSL. There is no Python-side lever — the fix lives entirely in the shader — but a batch server can afford to A/B the heavy and lean variants at start-up over a representative tile, log the median dispatch time of each against the adapter string, and pin whichever wins for the process. Emit the density buffer the smooth consumes as a contiguous C-ordered float32 array so the lean kernel’s straight-from-storage indexing stays coalesced.

Up: Workgroup Occupancy Optimization for Spatial Kernels