WebGPU Compute Shaders

A compute shader uses the GPU for any work you like, not only for drawing.

There are no triangles and no pixels. You give the GPU a big list of data, and thousands of small threads work on it at the same time.

What Is a Compute Shader?

Vertex and fragment shaders are part of the drawing pipeline. A compute shader stands alone. It reads and writes buffers, and that is all.

Use one when you have the same small job to do on many items:

• Move thousands of particles
• Run physics for cloth, water or crowds
• Process images: blur, edge detection, histograms
• Search and sort large amounts of data

A First Compute Shader

This shader doubles every number in an array:

@group(0) @binding(0) var<storage, read_write> numbers: array<f32>;

@compute @workgroup_size(64)
fn computeMain(@builtin(global_invocation_id) id: vec3u) {
  let index = id.x;
  if (index >= arrayLength(&numbers)) {
    return;
  }
  numbers[index] = numbers[index] * 2.0;
}

Workgroups

Threads run in groups called workgroups. You choose the size of a group in the shader, and how many groups to start in JavaScript.

Think of a school. Each classroom (workgroup) holds 64 pupils (threads). To give one million pupils a desk, you open enough classrooms for all of them.

A size of 64 is a good default. It works well on most GPUs.

Storage Buffers

A storage buffer is a buffer the shader can write to. Uniform buffers (from Lesson 5) are small and read-only. Storage buffers can be very large and writable.

Create one with the GPUBufferUsage.STORAGE flag:

const input = new Float32Array(1_000_000);
for (let index = 0; index < input.length; index++) {
  input[index] = index;
}

const workBuffer = device.createBuffer({
  size: input.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
});
device.queue.writeBuffer(workBuffer, 0, input);

COPY_DST lets us upload the numbers. COPY_SRC lets us copy the results out later.

The Compute Pipeline

A compute pipeline is simpler than a render pipeline. It has one shader stage and no colour targets:

const pipeline = device.createComputePipeline({
  layout: "auto",
  compute: { module: shaderModule, entryPoint: "computeMain" }
});

const bindGroup = device.createBindGroup({
  layout: pipeline.getBindGroupLayout(0),
  entries: [{ binding: 0, resource: { buffer: workBuffer } }]
});

Dispatching Work

Record a compute pass instead of a render pass. Instead of draw(), call dispatchWorkgroups() with the number of workgroups to start.

const WORKGROUP_SIZE = 64;
const workgroupCount = Math.ceil(input.length / WORKGROUP_SIZE);

const encoder = device.createCommandEncoder();
const pass = encoder.beginComputePass();
pass.setPipeline(pipeline);
pass.setBindGroup(0, bindGroup);
pass.dispatchWorkgroups(workgroupCount);
pass.end();

The rule is always the same: workgroups = Math.ceil(count / workgroup size).

Reading Results Back

JavaScript cannot read a storage buffer directly. You copy the results into a second buffer with the MAP_READ flag, then map it:

1. Create a buffer with MAP_READ | COPY_DST.
2. Record copyBufferToBuffer after the compute pass.
3. Submit, then wait for mapAsync(GPUMapMode.READ).
4. Read the data with getMappedRange(), then call unmap().

const resultBuffer = device.createBuffer({
  size: input.byteLength,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
});

encoder.copyBufferToBuffer(workBuffer, 0, resultBuffer, 0, input.byteLength);
device.queue.submit([encoder.finish()]);

await resultBuffer.mapAsync(GPUMapMode.READ);
const output = new Float32Array(resultBuffer.getMappedRange().slice(0));
resultBuffer.unmap();

console.log(output[0], output[1], output[2]);

Buffer Usages

Every buffer says up front how it will be used. These are the flags you will meet most:

Compute and Render Together

The real power comes from keeping data on the GPU. One buffer can be both a STORAGE buffer for a compute shader and a VERTEX buffer for drawing.

Each frame, a compute pass moves the particles, then a render pass draws them with instancing. The data never goes back to JavaScript.

const particleBuffer = device.createBuffer({
  size: particles.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
});

Each particle is a position and a velocity. The compute shader pulls every particle towards a moving point:

struct Particle {
  position: vec2f,
  velocity: vec2f,
}

struct Simulation {
  attractor: vec2f,
  deltaTime: f32,
  aspect: f32,
}

@group(0) @binding(0) var<uniform> simulation: Simulation;
@group(0) @binding(1) var<storage, read_write> particles: array<Particle>;

@compute @workgroup_size(64)
fn computeMain(@builtin(global_invocation_id) id: vec3u) {
  let index = id.x;
  if (index >= arrayLength(&particles)) {
    return;
  }

  var particle = particles[index];
  let toAttractor = simulation.attractor - particle.position;
  let direction = normalize(toAttractor);
  particle.velocity += direction * simulation.deltaTime;
  particle.position += particle.velocity * simulation.deltaTime;
  particles[index] = particle;
}

function frame(now) {
  updateSimulation(now);
  const encoder = device.createCommandEncoder();

  const computePass = encoder.beginComputePass();
  computePass.setPipeline(computePipeline);
  computePass.setBindGroup(0, computeBindGroup);
  computePass.dispatchWorkgroups(Math.ceil(PARTICLE_COUNT / WORKGROUP_SIZE));
  computePass.end();

  const renderPass = encoder.beginRenderPass(renderPassDescriptor);
  renderPass.setPipeline(renderPipeline);
  renderPass.setBindGroup(0, renderBindGroup);
  renderPass.setVertexBuffer(0, particleBuffer);
  renderPass.draw(6, PARTICLE_COUNT);
  renderPass.end();

  device.queue.submit([encoder.finish()]);
  requestAnimationFrame(frame);
}

The vertex buffer uses stepMode: "instance", so each instance reads one particle. The full example also adds a swirl and some drag, and colours fast particles orange and slow ones blue.

Exercises

numbers[index] = numbers[index] * numbers[index];

particle.velocity += vec2f(0.0, -0.5) * simulation.deltaTime;

Summary

• A compute shader runs general-purpose work on the GPU, with no triangles.
• @workgroup_size(64) sets the threads per group. global_invocation_id tells each thread which item is its own.
• Dispatch Math.ceil(count / 64) workgroups, and always add a bounds check.
• Storage buffers are large and writable. Read results back with copyBufferToBuffer and mapAsync.
• Keep data on the GPU: one buffer can be STORAGE for compute and VERTEX for drawing.

Official documentation

Go deeper with the specifications and reference documentation:

• developer.mozilla.org

  MDN: GPUDevice.createComputePipeline()

  Creating a pipeline from a compute shader entry point.

• developer.mozilla.org

  MDN: GPUComputePassEncoder.dispatchWorkgroups()

  Launching a grid of workgroups.

• developer.mozilla.org

  MDN: GPUBuffer.mapAsync()

  Reading results back from the GPU to JavaScript.

• w3.org

  WGSL specification: compute shader workgroups

  How workgroups, invocations and their built-in IDs relate.