Compute Pipelines

The Compute Pipeline and dispatchWorkgroups()

A compute pipeline has one shader stage, a layout and optional override constants; build it with createComputePipelineAsync() at startup so the browser compiles off the main thread. A compute pass records setPipeline(), setBindGroup() and dispatchWorkgroups(x, y, z), whose counts (y and z default to 1) are workgroups: dispatchWorkgroups(4, 2) with @workgroup_size(64) runs 512 invocations:

One dispatch: a grid of workgroups, each a block of invocations that the GPU runs in subgroups
One dispatch: a grid of workgroups, each a block of invocations that the GPU runs in subgroups

The workgroup is the unit that matters: its invocations run on one core, share var<workgroup> memory and can wait for each other at a barrier, while different workgroups can do none of that within a dispatch. Between dispatches, WebGPU makes every storage write of one visible to the next (each dispatch is its own usage scope), so multi-pass algorithms (GPGPU Patterns) synchronize across workgroups by starting another dispatch. A count above maxComputeWorkgroupsPerDimension (65,535 by default) is a validation error.

The experiments in this section run in a compute lab, demos/ch04/compute-lab.html, an extension of Scalars and Vectors's WGSL lab. The shader's first line gives the size and type of out and optionally the dispatch (// out: 4 x u32, dispatch 2); the lab enables subgroups and timestamp-query when present, runs the dispatch five times from a cleared buffer, and prints the start of out and the median GPU time between timestamps written at the start and end of the pass (timestampWrites, then resolveQuerySet(); Timestamp Queries). Two workgroups of 64 running atomicAdd(&out[0], 1) printed out: 128, 0, 0, 0 and gpu: 0.00 ms: Chrome 1 coarsens every timestamp to a multiple of 65,536 ns here (every value read was divisible by 2^16), as the specification requires for privacy, so only passes of a few tenths of a millisecond or more can be compared.

dispatchWorkgroups(4, 2) with @workgroup_size(64): 512 invocations, each painting its own cell by workgroupHTMLLive
<!doctype html>
<style>
  body { margin: 0; background: #f7f4ee; font: 14px system-ui, sans-serif; }
  .stage { position: relative; width: 100%; max-width: 600px; }
  .stage canvas { display: block; width: 100%; }
  .stage canvas + canvas { position: absolute; inset: 0; pointer-events: none; }
</style>
<div class="stage">
  <canvas id="view" width="600" height="330"></canvas>
  <canvas id="labels" width="600" height="330"></canvas>
</div>
<script>
const canvas = document.getElementById('view');
const ink = document.getElementById('labels').getContext('2d');

function showMessage(text) {                     // 2D fallback when WebGPU is missing
  const ctx = canvas.getContext('2d');
  ctx.fillStyle = '#fbeaea'; ctx.fillRect(0, 0, canvas.width, canvas.height);
  ctx.fillStyle = '#8a2b2b'; ctx.font = '18px system-ui, sans-serif'; ctx.textAlign = 'center';
  ctx.fillText(text, canvas.width / 2, canvas.height / 2);
}

// Each invocation records which workgroup ran it, in the cell of its global index.
const code = /* wgsl */ `
struct Out { count: atomic<u32>, cells: array<u32, 512> }
@group(0) @binding(0) var<storage, read_write> out: Out;
@compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3u,
                                     @builtin(local_invocation_index) local: u32,
                                     @builtin(num_workgroups) groups: vec3u) {
  let group = wg.x + wg.y * groups.x;                 // 0..7
  out.cells[group * 64 + local] = group + 1;
  atomicAdd(&out.count, 1);
}`;
const view = /* wgsl */ `
@group(0) @binding(0) var<storage> data: array<u32>;   // word 0: count, then 512 cells
struct Out { @builtin(position) pos: vec4f, @location(0) color: vec3f }
@vertex fn vs(@builtin(vertex_index) v: u32, @builtin(instance_index) i: u32) -> Out {
  let q = vec2f(f32(v & 1), f32(v >> 1));
  let group = i / 64;  let local = i % 64;
  let origin = vec2f(f32(group % 4) * 146, f32(group / 4) * 140);        // workgroups laid out 4 x 2
  let px = vec2f(10, 40) + origin + vec2f(f32(local % 8), f32(local / 8)) * 14 + q * 12;
  let g = f32(data[1 + i]);                                             // 1..8, or 0 if never run
  let hue = g / 8 * 6.28;
  let color = select(vec3f(0.85), 0.55 + 0.35 * cos(vec3f(hue, hue + 2.1, hue + 4.2)), g > 0);
  return Out(vec4f(px.x / 300 - 1, 1 - px.y / 165, 0, 1), color);
}
@fragment fn fs(in: Out) -> @location(0) vec4f { return vec4f(in.color, 1); }`;

async function main() {
  const adapter = await navigator.gpu?.requestAdapter();
  if (!adapter) return showMessage('WebGPU is not available in this browser');
  const device = await adapter.requestDevice();
  const context = canvas.getContext('webgpu');
  const format = navigator.gpu.getPreferredCanvasFormat();
  context.configure({ device, format });
  const B = GPUBufferUsage, SIZE = 4 + 512 * 4;
  const out = device.createBuffer({ size: SIZE, usage: B.STORAGE | B.COPY_SRC });
  const read = device.createBuffer({ size: 4, usage: B.COPY_DST | B.MAP_READ });
  // Built asynchronously at startup, so the browser compiles off the main thread.
  const pipeline = await device.createComputePipelineAsync({ layout: 'auto', compute: { module: device.createShaderModule({ code }) } });
  const module = device.createShaderModule({ code: view });
  const render = device.createRenderPipeline({ layout: 'auto', primitive: { topology: 'triangle-strip' },
    vertex: { module }, fragment: { module, targets: [{ format }] } });

  const encoder = device.createCommandEncoder();
  const pass = encoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: [{ binding: 0, resource: { buffer: out } }] }));
  pass.dispatchWorkgroups(4, 2);                     // 4 x 2 workgroups of 64 = 512 invocations
  pass.end();
  const draw = encoder.beginRenderPass({ colorAttachments: [{ view: context.getCurrentTexture().createView(),
    clearValue: [0.97, 0.96, 0.93, 1], loadOp: 'clear', storeOp: 'store' }] });
  draw.setPipeline(render);
  draw.setBindGroup(0, device.createBindGroup({ layout: render.getBindGroupLayout(0), entries: [{ binding: 0, resource: { buffer: out } }] }));
  draw.draw(4, 512);
  draw.end();
  encoder.copyBufferToBuffer(out, 0, read, 0, 4);
  device.queue.submit([encoder.finish()]);
  await read.mapAsync(GPUMapMode.READ);
  const count = new Uint32Array(read.getMappedRange())[0];
  read.unmap();

  ink.font = '12.5px ui-monospace, monospace'; ink.fillStyle = '#222';
  ink.fillText(`dispatchWorkgroups(4, 2) x @workgroup_size(64): atomicAdd counted ${count} invocations`, 10, 24);
  ink.font = '11px system-ui, sans-serif'; ink.fillStyle = '#444';
  for (let g = 0; g < 8; g++) ink.fillText(`workgroup_id (${g % 4}, ${Math.floor(g / 4)})`, 12 + (g % 4) * 146, 166 + Math.floor(g / 4) * 140);
  ink.fillText('One workgroup runs on one core and can share var<workgroup> memory and barriers; different workgroups cannot.', 10, 322);
}
main();
</script>