Choosing Workgroup Size for GPU Occupancy

The exact sub-problem here is that @workgroup_size is a compile-time constant baked into a WGSL entry point, yet the size that maximizes occupancy is a property of the silicon you happen to be running on — 256 on a 32-wide desktop NVIDIA part, often 64 on a narrow mobile Adreno or an integrated Intel GPU. Hard-code one number and you either under-fill wide hardware, leaving compute units idle waiting on storage reads, or over-subscribe the register file on narrow hardware, forcing the scheduler to co-resident fewer workgroups. Neither is visible in source; both show up only as a slower dispatch. This page is the empirical companion to workgroup occupancy optimization for spatial kernels: rather than reason about occupancy, it specializes one shader with an override constant and measures every candidate size on the actual device, returning the winner per tier.

Runnable reference implementation

The kernel is a single spatial scatter written once, with its workgroup width exposed as a WGSL override. The TypeScript harness compiles one pipeline per candidate size by supplying the constants map at pipeline creation, times each through a warmed timestamp query set, and returns the size with the lowest median dispatch time. Nothing about the kernel body changes between candidates, so the sweep isolates the effect of size alone.

// scatter.wgsl — one source, specialized per device class via the override.
override WG_SIZE: u32 = 256u;

@group(0) @binding(0) var<storage, read>       positions: array<vec2<f32>>;
@group(0) @binding(1) var<storage, read_write> grid:      array<atomic<u32>>;
@group(0) @binding(2) var<uniform>             extent:    vec4<f32>; // min_x,min_y,max_x,max_y

const GRID_W: u32 = 2048u;
const GRID_H: u32 = 2048u;

@compute @workgroup_size(WG_SIZE)
fn scatter(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= arrayLength(&positions)) { return; }      // guard the ragged tail workgroup
    let p = positions[i];
    let span = extent.zw - extent.xy;
    let n = (p - extent.xy) / span;                    // normalize to [0,1) in projected space
    if (n.x < 0.0 || n.x >= 1.0 || n.y < 0.0 || n.y >= 1.0) { return; } // drop, do not clamp
    let cx = u32(n.x * f32(GRID_W));
    let cy = u32(n.y * f32(GRID_H));
    atomicAdd(&grid[cy * GRID_W + cx], 1u);            // one coalesced read, one scatter
}

interface SweepResult {
  size: number;
  medianNanos: number;
}

// Build one pipeline per candidate size by injecting WG_SIZE as a pipeline constant.
async function sweepWorkgroupSize(
  device: GPUDevice,
  module: GPUShaderModule,          // compiled from scatter.wgsl
  layout: GPUPipelineLayout,
  bindGroup: GPUBindGroup,
  pointCount: number,
  candidates: number[],             // e.g. [32, 64, 128, 256] — subgroup multiples
  samples = 32,
): Promise<SweepResult[]> {
  // Timestamp plumbing: two writes (begin, end) per dispatch.
  const querySet = device.createQuerySet({ type: 'timestamp', count: 2 });
  const resolve = device.createBuffer({
    size: 16, usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC,
  });
  const readback = device.createBuffer({
    size: 16, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
  });

  const results: SweepResult[] = [];
  for (const size of candidates) {
    const pipeline = device.createComputePipeline({
      layout,
      compute: { module, entryPoint: 'scatter', constants: { WG_SIZE: size } },
    });
    const workgroups = Math.ceil(pointCount / size); // keep the whole device saturated

    const times: number[] = [];
    for (let s = 0; s < samples; s++) {
      const enc = device.createCommandEncoder();
      const pass = enc.beginComputePass({
        timestampWrites: { querySet, beginningOfPassWriteIndex: 0, endOfPassWriteIndex: 1 },
      });
      pass.setPipeline(pipeline);
      pass.setBindGroup(0, bindGroup);
      pass.dispatchWorkgroups(workgroups);
      pass.end();
      enc.resolveQuerySet(querySet, 0, 2, resolve, 0);
      enc.copyBufferToBuffer(resolve, 0, readback, 0, 16);
      device.queue.submit([enc.finish()]);

      await readback.mapAsync(GPUMapMode.READ);
      const t = new BigUint64Array(readback.getMappedRange());
      if (s >= 2) times.push(Number(t[1] - t[0])); // discard first two warm-up runs
      readback.unmap();
    }
    times.sort((a, b) => a - b);
    results.push({ size, medianNanos: times[times.length >> 1] });
  }
  return results.sort((a, b) => a.medianNanos - b.medianNanos);
}

The winner is results[0].size. Persist it keyed by adapter description so the sweep runs once per device class, not once per session — the first-run cost is a few milliseconds, but re-running it every page load is waste the user pays for.

Read the shape of the returned array, not just its first element, because the shape tells you why a size won and whether the choice is stable. A curve that climbs steeply from 32 to a plateau at 128 and 256 is the signature of a memory-bound kernel escaping an underfilled regime: below the plateau the compute unit cannot hide storage latency, and any size on the plateau is safe. A curve that peaks at 128 and then regresses at 256 is the register or shared-memory limiter biting — the larger size reserves so much per-workgroup state that fewer workgroups stay co-resident — and it warns that a small change to the kernel could move the peak, so pin the size and re-sweep whenever the shader body changes. A curve that is flat across every candidate means occupancy is not the bottleneck at all; the kernel is compute-bound, and no size will help until the arithmetic per invocation comes down. Treating the sweep as a single winner throws away this diagnosis; keep the whole array in your logs so a later regression on one device tier is explicable rather than mysterious.

The margin between neighbouring candidates also matters for how aggressively to specialize. When the plateau is broad and flat — 96, 128, 160, and 192 all within a few percent — a single portable size like 128 serves every tier and the per-class cache buys little. When the winner beats its neighbours by a wide margin, or the winning size differs sharply between a desktop and a mobile adapter, the per-class cache earns its keep and a shared default would leave real throughput on the table. Let the measured spread, not a rule of thumb, decide whether one size ships everywhere or each device class gets its own.

Parameter reference

For a 2D grid-indexed kernel, sweep (x, y) pairs whose product stays within the invocation cap — (8,8), (16,16), (32,8) — rather than 1D widths, because the access pattern of a neighbourhood read favours a squarer tile. The full occupancy model behind these numbers, including the shared-memory and register terms, is in workgroup occupancy optimization for spatial kernels, and the timing plumbing is shared with frame profiling with timestamp queries.

Failure modes

• The sweep picks a size that is fast once, slow in production. Timing a cold pipeline folds shader compilation into the measurement, favouring whichever size happened to compile first. Detection: the winning size changes run to run for identical input. Fix: discard the first two submissions per candidate (the harness does this) and take the median, never the minimum.
• A candidate over the invocation cap crashes the sweep. Including 512 on a device whose maxComputeInvocationsPerWorkgroup is 256 throws at createComputePipeline. Detection: a GPUValidationError naming the limit, before any dispatch runs. Fix: clamp candidates to device.limits.maxComputeInvocationsPerWorkgroup and, for 2D, to the per-axis caps.
• The chosen size regresses on an untested device class. A 256-wide winner from a desktop sweep under-fills nothing but over-subscribes a narrow mobile register file, halving throughput there. Detection: median dispatch time regresses only on one tier of hardware. Fix: run the sweep on each shipped device class and cache a size per class, falling back through the browser support fallback strategies where WebGPU is absent.
• Non-multiple candidates blur the result. Sweeping [100, 200, 300] mixes lane-masking waste into the timing, so the fastest of them is not the fastest possible size. Detection: the winner beats its neighbours by an implausibly small margin. Fix: generate candidates as whole subgroup multiples, as the reference does.

Backend and Python interop note

When a headless render server drives the same kernels through the wgpu-py bindings, it exposes the identical constants surface at pipeline creation, so the sweep ports directly. Run it once at server start, log the winning size against the adapter string, and pin it for the process — a batch tile renderer has no user waiting, so it can afford a longer, more thorough sweep than a browser session and should cache the result to disk keyed by GPU model.

Related

• Workgroup Occupancy Optimization for Spatial Kernels — the occupancy model and shared-memory and register limits behind the size you are choosing.
• Reducing Register Pressure in WGSL Spatial Kernels — how to widen the size window by freeing the register term.
• Frame Profiling with Timestamp Queries — the timestamp query mechanism the sweep harness is built on.
• Optimizing Workgroup Sizes for Vector Geometry Filtering — the same sweep applied to a filter pass rather than a scatter.

Up: Workgroup Occupancy Optimization for Spatial Kernels