Skip to content

SOP: GPU Compute with WebGPU

Fresh 🌱

Version: 1.0 | Chrome: 113+

Overview

Compute shaders are WebGPU-only (not in WebGL). They run general-purpose parallel computations on the GPU using hundreds/thousands of threads — ideal for ML inference, physics, image processing, and data transformation.

Performance threshold: GPU compute outperforms CPU at matrix dimensions >256x256.

Compute Pipeline Flowchart

Step-by-Step: Matrix Multiplication

Step 1 — Prepare Data

javascript
// Example: multiply two matrices
const matrixSize = 4; // 4x4 matrices

const matA = new Float32Array([
  1,2,3,4,  5,6,7,8,  9,10,11,12,  13,14,15,16
]);
const matB = new Float32Array([
  1,0,0,0,  0,1,0,0,  0,0,1,0,  0,0,0,1  // identity
]);

Step 2 — Create GPU Buffers

javascript
function createBuffer(device, data, usage) {
  const buffer = device.createBuffer({
    size: data.byteLength,
    usage: usage | GPUBufferUsage.COPY_DST,
    mappedAtCreation: true
  });
  new Float32Array(buffer.getMappedRange()).set(data);
  buffer.unmap();
  return buffer;
}

const bufA = createBuffer(device, matA, GPUBufferUsage.STORAGE);
const bufB = createBuffer(device, matB, GPUBufferUsage.STORAGE);

// Result buffer — STORAGE + COPY_SRC so we can read it back
const resultBuffer = device.createBuffer({
  size: matA.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
});

// Staging buffer for reading results back to CPU
const stagingBuffer = device.createBuffer({
  size: matA.byteLength,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
});

Buffer States

GPU buffers must be unmapped before use in GPU commands. Once mapped (for reading/writing from JS), they cannot be used in commands until unmapped.

Step 3 — Write WGSL Compute Shader

wgsl
// Bind group 0: input matrices and result
@group(0) @binding(0) var<storage, read>       matA   : array<f32>;
@group(0) @binding(1) var<storage, read>       matB   : array<f32>;
@group(0) @binding(2) var<storage, read_write> result : array<f32>;

// Workgroup: 8x8 = 64 threads per workgroup
@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) id : vec3u) {
  let row = id.x;
  let col = id.y;
  let N   = 4u; // matrix dimension

  if (row >= N || col >= N) { return; }

  var sum = 0.0;
  for (var k = 0u; k < N; k++) {
    sum += matA[row * N + k] * matB[k * N + col];
  }
  result[row * N + col] = sum;
}

Step 4 — Create Compute Pipeline

javascript
const shaderModule = device.createShaderModule({ code: wgslCode });

const computePipeline = device.createComputePipeline({
  label: 'Matrix multiply pipeline',
  layout: 'auto',
  compute: {
    module: shaderModule,
    entryPoint: 'main'
  }
});

Step 5 — Create Bind Group

javascript
const bindGroup = device.createBindGroup({
  layout: computePipeline.getBindGroupLayout(0),
  entries: [
    { binding: 0, resource: { buffer: bufA } },
    { binding: 1, resource: { buffer: bufB } },
    { binding: 2, resource: { buffer: resultBuffer } }
  ]
});

Step 6 — Dispatch and Read Results

javascript
// Dispatch compute work
const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);

// Dispatch enough workgroups to cover all elements
const workgroupsX = Math.ceil(matrixSize / 8);
const workgroupsY = Math.ceil(matrixSize / 8);
passEncoder.dispatchWorkgroups(workgroupsX, workgroupsY);
passEncoder.end();

// Copy result to staging buffer for CPU read
commandEncoder.copyBufferToBuffer(
  resultBuffer, 0,
  stagingBuffer, 0,
  matA.byteLength
);

device.queue.submit([commandEncoder.finish()]);

// Read results back to CPU (async)
await stagingBuffer.mapAsync(GPUMapMode.READ);
const resultData = new Float32Array(stagingBuffer.getMappedRange().slice(0));
stagingBuffer.unmap();

console.log('Result matrix:', resultData);

Storage Buffers vs Uniform Buffers

FeatureUniform BufferStorage Buffer
Max size64KB128MB+
Write from shaderNo (read-only)Yes
Runtime-sized arraysNoYes
Atomic operationsNoYes
Use caseConstants, transformsLarge data, ML workloads

Workgroup Sizing Tips

wgsl
// Common workgroup sizes
@compute @workgroup_size(64)         // 1D: 64 threads
@compute @workgroup_size(8, 8)       // 2D: 64 threads (image/matrix processing)
@compute @workgroup_size(4, 4, 4)    // 3D: 64 threads (volume data)
  • Total workgroup size must not exceed device.limits.maxComputeInvocationsPerWorkgroup (usually 256)
  • dispatchWorkgroups(x, y, z) launches x * y * z workgroups

See Also

Based on Chrome for Developers WebGPU Documentation