Implementing a Workgroup Prefix Sum for Stream Compaction

Stream compaction with an atomic counter gives each survivor a slot; a prefix sum gives each survivor its slot — the same one, every run, in input order. That determinism is what a picking index, a test fixture or a searchable bucket structure needs, and the primitive that provides it is a scan: log₂(n) rounds of shifted addition inside a workgroup, then a three-pass structure to stitch workgroups together. This page is the WGSL for both, the barrier discipline that the scan cannot function without, and the compaction that consumes its output. It is one stage of GPU sorting and prefix sums for spatial data.

Three rounds of a Hillis-Steele scan over eight lanes Three strips of eight lanes each show the scan progressing. The initial row holds the mask values one, zero, one, one, zero, one, zero and one. After the first round, with a stride of one, each lane holds the sum of itself and its immediate neighbour. After the second round, with a stride of two, each lane holds the sum of a window of four. After the third round, with a stride of four, each lane holds the inclusive sum of everything up to and including itself: one, one, two, three, three, four, four and five. LANE 0 1 2 3 4 5 6 7 8 1 0 1 1 0 1 0 1 Initial the mask one flag per element, no communication yet 1 1 2 2 2 3 2 3 Round 2 stride 2 each lane now holds the sum of a window of four 1 1 2 3 3 4 4 5 Round 3 inclusive subtract each lane’s own flag to get the exclusive offsets Every round is separated from the next by two barriers — both are required.
Eight lanes take three rounds; 256 take eight. The round count is the base-two logarithm of the workgroup size, which is why a scan is cheap enough to sit inside a sort.

Runnable reference implementation

The within-workgroup scan is a Hillis–Steele ladder. Each round adds the value from stride lanes back, doubling the stride until it exceeds the workgroup size.

wgsl
const WG : u32 = 256u;
var<workgroup> tile : array<u32, WG>;

@group(0) @binding(0) var<storage, read>       flags      : array<u32>;
@group(0) @binding(1) var<storage, read_write> offsets    : array<u32>;
@group(0) @binding(2) var<storage, read_write> block_sums : array<u32>;

@compute @workgroup_size(256)
fn scan_blocks(@builtin(global_invocation_id) gid : vec3<u32>,
               @builtin(local_invocation_id)  lid : vec3<u32>,
               @builtin(workgroup_id)         wid : vec3<u32>) {
  let i = gid.x;
  let t = lid.x;
  let n = arrayLength(&flags);

  let own = select(0u, flags[i], i < n);        // tail guard: past the end is 0
  tile[t] = own;
  workgroupBarrier();

  for (var stride : u32 = 1u; stride < WG; stride = stride << 1u) {
    var add : u32 = 0u;
    if (t >= stride) { add = tile[t - stride]; }
    workgroupBarrier();          // all reads complete before any write
    tile[t] = tile[t] + add;
    workgroupBarrier();          // all writes complete before the next read
  }

  // tile[t] is now the INCLUSIVE sum. Exclusive = inclusive - own.
  if (i < n) { offsets[i] = tile[t] - own; }

  // The last lane publishes this block's total for the second pass.
  if (t == WG - 1u) { block_sums[wid.x] = tile[t]; }
}

The scatter that consumes it is trivial by comparison, and that is the point: all the coordination happened in the scan, so the write is contention-free.

wgsl
@group(0) @binding(0) var<storage, read>       flags       : array<u32>;
@group(0) @binding(1) var<storage, read>       offsets     : array<u32>;
@group(0) @binding(2) var<storage, read_write> compacted   : array<u32>;

@compute @workgroup_size(256)
fn scatter(@builtin(global_invocation_id) gid : vec3<u32>) {
  let i = gid.x;
  if (i >= arrayLength(&flags)) { return; }
  if (flags[i] == 1u) { compacted[offsets[i]] = i; }
}

Parameter reference

Value Setting here Guidance
Workgroup size 256 The var<workgroup> array is 4 × WG bytes, so 256 costs 1 KiB of the 16 KiB floor. 512 is worth measuring.
Barriers per round 2 Both are required. One is enough on hardware that happens to run a workgroup in lockstep, which makes the bug machine-dependent.
Rounds log₂(WG) = 8 Fixed by the workgroup size; the loop is fully unrollable by the compiler.
Block sums buffer ceil(n / WG) entries Allocate once at start-up for the largest array the pipeline will scan.
Output buffer worst case n Allocated for everything passing, used for however many did.
Stitching workgroup scans into a global one Three passes left to right. The first pass scans each block of 256 elements independently and publishes that block total. The second pass scans the array of block totals, which is a four-hundredth of the size. The third pass adds each scanned block total back into every element of its block, converting local offsets into global ones. GLOBAL SCAN · THREE PASSES Scan blocks 256 at a time publish a block total Scan the totals n / 256 elements recurse if large Add back block total + local offset global offsets Every GPU scan implementation has this shape — recognising it makes them all readable.
The middle pass is tiny, which is why three passes cost far less than three times one. Recursion on the middle pass extends the structure to any array size without new code.

The two barriers, and why one is not enough

The loop body reads tile[t - stride] and then writes tile[t]. Both operations touch memory other lanes in the workgroup are also touching, and the two barriers separate them into phases.

Without the first barrier, a lane that has already written its round-k value can have that value read by a neighbour still executing round k-1, so the neighbour adds a term from the wrong round. Without the second, a lane can begin round k+1 and read a slot its neighbour has not yet written for round k.

What makes this dangerous rather than merely wrong is that it often works. Hardware that executes a whole workgroup in lockstep — which is common for a workgroup no larger than the wavefront — produces the correct answer with no barriers at all, so a scan written and tested at workgroup size 32 on one vendor can fail at 256 on another. The failure is a wrong offset for some elements, which shows up as a compacted array with duplicates and gaps rather than as an error.

The version above is deliberately the conservative one: two barriers, unconditional, outside any if. WGSL requires barriers to be in uniform control flow, so a barrier inside the if (t >= stride) block would be invalid — a rule that exists precisely because the alternative is undefined behaviour.

What the scan costs, and when it is worth it

A three-pass global scan over four million elements is roughly 1.4 times the cost of a single streaming pass over the same data, which is a surprisingly small multiple for something that computes a global dependency. The reason is that the middle pass touches a four-hundredth of the array and the other two are pure bandwidth.

Against an atomic counter, which does the same job in one pass, the scan is therefore about twice the cost. That is the whole trade, and it is worth being concrete about what the extra pass buys. Determinism: the same input produces the same output ordering on every run and every device, which makes a compacted array something a test can assert on. Order preservation: survivors appear in input order, so a compacted index list can be binary-searched if the input was sorted. And no contention: an atomic on one address serialises when the survivor rate is high, so the gap narrows exactly in the case where most elements pass.

The rule that follows is a simple one. Use the atomic when the output is an unordered set that only the GPU will read — particles, point survivors, anything drawn as a cloud. Use the scan when the output is a structure: a bucket table another pass indexes into, a picking index the CPU will search, or anything that appears in a regression test.

There is also a hybrid worth knowing about. A per-workgroup atomic that claims a block of slots, followed by a within-workgroup scan to distribute them, gets most of the scan’s ordering guarantee at close to the atomic’s cost — the ordering is preserved within a block but not across blocks, which is enough for many bucket structures and not enough for a reproducible index.

Failure modes

  • The compacted array has duplicates and gaps. A missing or misplaced barrier, so some offsets are wrong. Check: the last offset plus the last flag must equal the total count.
  • Everything works at workgroup size 64 and fails at 256. The same bug, hidden by lockstep execution at the smaller size. Always test at the size you ship.
  • The tail workgroup corrupts the result. No tail guard, so out-of-bounds reads contributed zeros in the wrong places — or, worse, the block sum for the final partial block was never written.
  • Offsets are off by one. The scan produced an inclusive sum and the scatter treated it as exclusive. Subtract the element’s own flag, or use a down-sweep formulation.
  • GPUValidationError: barrier in non-uniform control flow. A workgroupBarrier() inside a conditional. Restructure so the barrier is unconditional and the conditional only guards the value.
Scan or atomic counter? A comparison of a prefix sum and an atomic counter across four properties for stream compaction. The scan produces output in input order and the atomic does not. The scan produces the same result on every run and the atomic does not. The scan takes three passes and the atomic takes one. The scan has no contention while the atomic contends on a single address, which matters when the survivor rate is high. COMPACTION · SCAN vs ATOMIC Prefix sum Atomic counter Output order input order arbitrary Reproducible yes no Passes three one Contention none one hot address Needing reproducibility is the deciding question — everything else is a tie-break.
The atomic wins on cost and the scan wins on determinism. For point clouds the atomic is right; for anything another pass will index into or a test will assert on, the scan is worth its two extra passes.

Backend / Python interop note

There is no server-side half to a scan — it is a GPU primitive with no wire format — but there is a testing story that belongs on the Python side, and it is unusually cheap.

python
import numpy as np

def reference_scan(flags: np.ndarray) -> np.ndarray:
    """Exclusive prefix sum: what the GPU must produce, bit for bit."""
    return np.concatenate(([0], np.cumsum(flags[:-1]))).astype(np.uint32)

def reference_compaction(flags: np.ndarray) -> np.ndarray:
    return np.flatnonzero(flags).astype(np.uint32)

Generating a fixture of a few thousand flags, running both the reference and the GPU implementation, and asserting exact equality catches every barrier and off-by-one bug in this page — and it catches them at a size small enough to print when it fails. It is also worth asserting the property rather than only the values. The invariant a scan has to satisfy is that the last offset plus the last flag equals the total number of set flags, and that every offset is non-decreasing. Those two checks are three lines, they hold for any input, and they fail loudly for exactly the class of bug — a barrier in the wrong place — that a hand-written expected array would only catch on the particular fixture that happens to trigger it.

Two fixture shapes are worth including specifically: one whose length is an exact multiple of the workgroup size, and one that is not, because the tail-guard bugs only appear in the second.