WebGPU workgroup shared memory for parallel reduction

Summing a large array without leaving the GPU

Suppose you have a million 32-bit floats in a GPU buffer and need their total. The obvious options are slow: copy the buffer back to the CPU, or launch a kernel that contends on a global atomic. WebGPU offers a middle path: workgroup shared memory. Threads in the same workgroup can load a tile into shared memory, synchronize with a barrier, and reduce the tile to a single partial sum using descending strides. A second launch then reduces those partials.

A minimal reduction kernel

This WGSL compute shader loads 128 elements per workgroup into shared memory, reduces them in place, and writes one partial result per workgroup.

// Run in a browser context with WebGPU enabled.
// @group(0) @binding(0) is the input storage buffer.
// @group(0) @binding(1) is the output partials buffer.
@group(0) @binding(0) var<storage, read> inputData: array<f32>;
@group(0) @binding(1) var<storage, read_write> partials: array<f32>;

@compute @workgroup_size(128)
fn main(@builtin(global_invocation_id) gid: vec3<u32>,
        @builtin(local_invocation_id) lid: vec3<u32>,
        @builtin(num_workgroups) numGroups: vec3<u32>) {
    let index = gid.x;
    let local = lid.x;

    var<shared> memory: array<f32, 128>;

    if (index < arrayLength(inputData)) {
        memory[local] = inputData[index];
    } else {
        memory[local] = 0.0;
    }

    barrier();

    let n = 128;
    for (var stride = n / 2; stride > 0; stride = stride / 2) {
        if (local < stride) {
            memory[local] = memory[local] + memory[local + stride];
        }
        barrier();
    }

    if (local == 0) {
        partials[gid.x / 128] = memory[0];
    }
}

Host setup and device checks

Run this in a secure context, such as a module script or DevTools console on an origin served over HTTPS. You need a GPU device and two buffers: one for input, one for partial outputs.

// JavaScript: host side
if (navigator.gpu === undefined || navigator.gpu.status !== 'available') {
    throw new Error('WebGPU is not available');
}

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();

const maxWorkgroup = device.limits.maxComputeWorkgroupSize[0];
if (maxWorkgroup < 128) {
    console.warn('Requested workgroup size 128 exceeds limit', maxWorkgroup);
}

// Create shader module, pipeline, and bind groups here.
// layout: 'auto' lets the browser derive the layout from shader bindings.

Call device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } }). If the browser throws a ValidationError, check that the shader binding order matches the bind group layout you provide. Once a pipeline layout is created from a bind group layout, entry count and types are fixed; you cannot swap buffers between bindings at runtime without recreating the pipeline.

Limits and trade-offs

Workgroup shared memory is typically capped at 16 KB per workgroup. A 128-element f32 tile uses 512 bytes, which fits comfortably. If you widen the tile to 1024 elements or store vec4 values, you can exceed the cap and pipeline creation will fail. Reduce the workgroup size, split the tile across multiple launches, or use a narrower element type.

Browser support is uneven. Safari Technology Preview has trailed Chrome and Edge on some WebGPU features, so query device.features before shipping code that depends on newer texture or compute capabilities. Shared memory itself is widely supported, but always verify limits on target devices.

Verifying the result

After dispatching the reduction kernel, read the partials buffer back and compare it against a CPU reference implementation. A mismatch usually indicates a missed barrier(), an off-by-one stride, or an out-of-bounds load. Confirm no validation errors appear in the browser console before treating the output as correct.

Closing

WebGPU shared memory makes classic GPU reduction practical without global atomics or CPU round trips. Start with a 128-thread workgroup, keep shared memory under the 16 KB cap, lock binding order at pipeline creation time, and validate against a CPU reference on every target browser.