Performing a Data‑Parallel Reduction with WebGPU Storage Buffers

Useful answer

You can compute a parallel reduction (e.g., sum of an array) entirely on the GPU by using a WebGPU compute pipeline that reads from and writes to a storage buffer. The steps are: allocate a GPUBuffer with STORAGE usage, copy your data into it, bind the buffer to a compute shader, dispatch workgroups that each produce a partial sum, run a second pass to combine the partials, copy the final result to a mappable buffer, and read it back on the CPU.

Worked configuration

The example below shows a two‑pass reduction: first pass reduces chunks of 256 elements per workgroup into a temporary array; second pass sums that temporary array into a single value. All code is written in JavaScript (WebGPU API) and WGSL.

1. Create buffers

// Assume `device` is a valid GPUDevice
const inputSize = 1 << 20; // 1,048,576 floats
const inputData = new Float32Array(inputSize);
// fill `inputData` with your values here

// Input buffer (read‑only for the shader)
const inputBuffer = device.createBuffer({
  size: inputData.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(inputBuffer, 0, inputData);

// Workgroup size for the first pass
const WORKGROUP_SIZE = 256;
const workgroupCount = Math.ceil(inputSize / WORKGROUP_SIZE);
// Temporary buffer to hold partial sums (one per workgroup)
const partialSize = workgroupCount * 4; // float per group
const partialBuffer = device.createBuffer({
  size: partialSize,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});

// Final result buffer (single float)
const outputBuffer = device.createBuffer({
  size: 4,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
});

// Staging buffer to read back the result
const readbackBuffer = device.createBuffer({
  size: 4,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});

2. Define the bind group layout and pipeline

const bindGroupLayout = device.createBindGroupLayout({
  entries: [
    { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
    { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
    { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
  ],
});

const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] });

// First pass shader (reduce chunk → partial)
const reduceModule = device.createShaderModule({
  code: `
    @group(0) @binding(0) var input: array;
    @group(0) @binding(1) var partial: array;
    @group(0) @binding(2) var flags: array;

    @compute @workgroup_size(256)
    fn main(@builtin(global_invocation_id) gid: vec3) {
      let idx = gid.x;
      let stride = workgroup_count_x * 256u;
      var acc = 0.0;
      for (var i = idx; i < arrayLength(&input); i += stride) {
        acc += input[i];
      }
      // Work‑group reduction using shared memory
      var shared: array;
      shared[gid.x] = acc;
      workgroupBarrier();
      for (var offset = 128u; offset > 0; offset = offset >> 1) {
        if (gid.x < offset) {
          shared[gid.x] += shared[gid.x + offset];
        }
        workgroupBarrier();
      }
      if (gid.x == 0) {
        partial[gid.y] = shared[0]; // store partial sum for this workgroup
      }
    }
  `,
});

// Second pass shader (sum partials → final)
const sumModule = device.createShaderModule({
  code: `
    @group(0) @binding(0) var partial: array;
    @group(0) @binding(1) var output: array;

    @compute @workgroup_size(256)
    fn main(@builtin(global_invocation_id) gid: vec3) {
      let idx = gid.x;
      if (idx >= arrayLength(&partial)) { return; }
      var acc = 0.0;
      for (var i = idx; i < arrayLength(&partial); i += 256u) {
        acc += partial[i];
      }
      // Simple tree‑reduce in one workgroup (assumes partials fit)
      var shared: array;
      shared[gid.x] = acc;
      workgroupBarrier();
      for (var offset = 128u; offset > 0; offset = offset >> 1) {
        if (gid.x < offset) {
          shared[gid.x] += shared[gid.x + offset];
        }
        workgroupBarrier();
      }
      if (gid.x == 0) {
        output[0] = shared[0];
      }
    }
  `,
});

const reducePipeline = device.createComputePipeline({
  layout: pipelineLayout,
  compute: { module: reduceModule, entryPoint: "main" },
});

const sumPipeline = device.createComputePipeline({
  layout: pipelineLayout,
  compute: { module: sumModule, entryPoint: "main" },
});

3. Encode commands

function encodePass(pipeline, bindings) {
  const commandEncoder = device.createCommandEncoder();
  const pass = commandEncoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindings);
  // Dispatch enough workgroups to cover the buffer
  pass.dispatchWorkgroups(Math.ceil(bindings[1].size / 4 / 256), 1, 1);
  pass.end();
  return commandEncoder.finish();
}

// Bind groups for first pass
const bindGroup1 = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [
    { binding: 0, resource: { buffer: inputBuffer } },
    { binding: 1, resource: { buffer: partialBuffer } },
    { binding: 2, resource: { buffer: device.createBuffer({ size: 4, usage: GPUBufferUsage.STORAGE }) } }, // dummy flags buffer
  ],
});

// Bind groups for second pass
const bindGroup2 = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [
    { binding: 0, resource: { buffer: partialBuffer } },
    { binding: 1, resource: { buffer: outputBuffer } },
    { binding: 2, resource: { buffer: device.createBuffer({ size: 4, usage: GPUBufferUsage.STORAGE }) } }, // dummy
  ],
});

const cmds1 = encodePass(reducePipeline, bindGroup1);
const cmds2 = encodePass(sumPipeline, bindGroup2);

// Copy result to staging buffer for reading
const copyEncoder = device.createCommandEncoder();
copyEncoder.copyBufferToBuffer(outputBuffer, 0, readbackBuffer, 0, 4);
const copyCmd = copyEncoder.finish();

device.queue.submit([cmds1, cmds2, copyCmd]);

4. Read the result

await readbackBuffer.mapAsync(GPUMapMode.READ);
const result = new Float32Array(readbackBuffer.getMappedRange());
console.log('GPU sum:', result[0]);
readbackBuffer.unmap();

Limits and verification

• Maximum storage buffer size: Check device.limits.maxStorageBufferBindingSize. If your input or partial buffer exceeds this, the pipeline creation will fail with a validation error.
• Offset alignment: Buffer offsets used in bind groups must be multiples of 256 bytes. When creating buffers, ensure the size is a multiple of 256 or use aligned offsets.
• Workgroup dimensions: The product of workgroup counts in each dimension must not exceed device.limits.maxComputeWorkgroupCountPerDimension. For large inputs, dispatch multiple passes or increase workgroup size.
• Validation: Wrap submission in device.pushErrorScope('validation') and popErrorScope() to catch mismatched usage (e.g., forgetting STORAGE flag or providing a read‑only buffer where read_write is expected).

Common mistakes

• Omitting GPUBufferUsage.STORAGE on a buffer that the shader declares as storage; the pipeline creation will reject with an "incompatible access mode" error.
• Mapping a buffer that is still in use by the GPU. Always wait for queue submission to complete (via the returned promise from device.queue.submit) before calling mapAsync.
• Using the wrong WGSL layout qualifier. Storage buffers default to std430; if you declare [[block]] structures, ensure they match std430 rules or use [[align(16)]] as needed.
• Missing a storage texture or buffer barrier between passes. In the example we reuse the same buffer without an explicit barrier because we submit separate command encoders; if you encode multiple passes in a single pass, you need textureBarrier or storageBarrier.
• Dispatching more workgroups than the device allows in a single dimension, causing GPUValidationError: dispatch workgroup count exceeds limits. Split the work into multiple dispatches or reduce workgroup size.

Practical way to check the result

After mapping the staging buffer, compare the GPU‑produced sum to a CPU‑calculated reference (e.g., inputData.reduce((a,v)=>a+v,0)). If they match within floating‑point tolerance, the reduction succeeded.