Exploring GPU Algorithms with compute-sim

compute-sim started from a practical frustration. The only GPU I had was an old AMD laptop iGPU that AMD's developer tools no longer supported. So I wrote a CPU simulator to experiment with compute shaders instead.

The project runs shader-like Nim code on the CPU. It supplies the GLSL pieces used by compute algorithms: invocation IDs, workgroups, subgroups, shared memory, barriers, and atomics. I don't use it to estimate GPU performance. I use it to check an algorithm and to watch a small group of invocations execute it.

Reduction, scan, and matrix multiplication are the examples I return to most often. None has complicated arithmetic. All three become interesting when the arithmetic is divided among many invocations.

Enough of the execution model

A dispatch is split into workgroups. gl_GlobalInvocationID identifies an invocation across the whole dispatch, while gl_WorkGroupID and gl_LocalInvocationID locate it inside a workgroup.

Workgroups are split again into subgroups. The invocations in a subgroup execute in lockstep and can exchange values through collective operations. gl_SubgroupInvocationID gives the position inside that subgroup. compute-sim uses eight invocations per subgroup by default, though the size is configurable.

The shader itself is a Nim procedure:

proc shader(buffers: ptr Buffers; args: Args) {.computeShader.} =
  # compute code

The {.computeShader.} macro lets the simulator advance the invocations together and intercept operations such as subgroupAdd and barrier. Branches and loops can diverge, then reconverge. The code still resembles the shader I have in mind, but I can run it in a debugger and print what one subgroup is doing.

A reduction using subgroupAdd

Here is the small version of reduction:

proc reduce(b: ptr Buffers; numElements: uint32) {.computeShader.} =
  let gid = gl_GlobalInvocationID.x
  let value = if gid < numElements: b.input[gid] else: 0
  let sum = subgroupAdd(value)

  if gl_SubgroupInvocationID == 0:
    atomicAdd b.atomicSum, sum

Each invocation loads one input. Extra invocations at the end contribute zero. After subgroupAdd, every active invocation has the subgroup's sum, but only lane zero writes it. The write is atomic because other subgroups update the same variable.

This is about as compact as the mapping can be without hiding it in host code. The input index, the collective operation, and the choice of writer are all in the shader.

I often compile this example with -d:debugSubgroup. For the first subgroup, compute-sim prints:

[Add #0] inputs {t0: 0, t1: 1, t2: 2, t3: 3,
                 t4: 4, t5: 5, t6: 6, t7: 7} | sum: 28

That trace is more useful to me than a diagram of a subgroup. It shows the actual values that arrived at the operation. If a branch excluded some lanes, the difference would be visible here.

Scan, one subgroup at a time

A scan preserves the intermediate sums instead of keeping only the final one. [3, 1, 4, 1] becomes [3, 4, 8, 9] with an inclusive scan.

The corresponding scan shader is almost the same as the reduction:

proc inclusiveScan(values: ptr seq[int32]; numElements: uint32) {.computeShader.} =
  let gid = gl_GlobalInvocationID.x
  let value = if gid < numElements: values[gid] else: 0
  let prefix = subgroupInclusiveAdd(value)

  if gid < numElements:
    values[gid] = prefix

Now every lane has a different result. Lane zero gets the first value. Lane one gets the first two, and the last lane gets the subgroup total. subgroupExclusiveAdd shifts those results by one lane and puts zero at the start.

The example uses a single subgroup. For a larger scan, I save the total from each subgroup, scan those totals, and add the offsets back to the original results. I found this easier to reason about after writing the eight-element case: the larger algorithm is another scan wrapped around the first one.

One invocation per matrix cell

Matrix multiplication looks unrelated to scan, but it starts with the same question: what does one invocation own? For A with shape M × K and B with shape K × N, one output cell is a dot product:

C[row, col] = sum(A[row, i] * B[i, col], i = 0 ..< K)

In the first matrix multiplication shader, an invocation owns one cell of C:

proc multiply(b: ptr Buffers; dims: Dimensions) {.computeShader.} =
  let row = gl_GlobalInvocationID.y
  let col = gl_GlobalInvocationID.x

  if row < dims.m and col < dims.n:
    var sum = 0'f32
    for i in 0..<dims.k:
      sum += b.a[row * dims.k + i] * b.b[i * dims.n + col]
    b.c[row * dims.n + col] = sum

There is nothing to synchronize. Every invocation has its own accumulator and writes to a different output cell. The rows and columns are covered by the dispatch; the dot product remains a normal loop.

I keep this version around because its indexing is easy to verify. It also makes the redundant loads hard to miss. Two invocations in the same output row read the same row of A. Two in the same output column read the same column of B.

Writing the reduction tree explicitly

subgroupAdd is a good way to begin, but it hides the reduction itself. The next version has the workgroup perform the reduction in shared memory:

proc reduceWorkgroup(b: ptr Buffers; shared: ptr seq[int32];
                     numElements: uint32) {.computeShader.} =
  let localId = gl_LocalInvocationID.x
  let gid = gl_GlobalInvocationID.x
  shared[localId] = if gid < numElements: b.input[gid] else: 0

  memoryBarrier()
  barrier()

  var stride = gl_WorkGroupSize.x div 2
  while stride > 0:
    if localId < stride:
      shared[localId] += shared[localId + stride]
    memoryBarrier()
    barrier()
    stride = stride div 2

  if localId == 0:
    b.partialSums[gl_WorkGroupID.x] = shared[0]

For eight invocations, the strides are four, two, and one:

8 values  ->  4 pair sums  ->  2 sums  ->  1 workgroup sum
stride 4      stride 2         stride 1

The number of additions halves each time. Barrier participation does not. Even an invocation that skipped the addition must reach the barrier before the next stride begins. memoryBarrier() orders the shared-memory accesses; barrier() waits for the rest of the workgroup. If the barrier were placed inside if localId < stride, the inactive invocations would never reach it.

Each workgroup leaves one partial sum. Running the same idea again on those partial sums finishes a larger reduction. I enjoy this version for the shrinking pattern in shared: after each barrier, half of the remaining entries stop mattering.

Loading matrix tiles together

Those repeated loads are why I wrote a tiled version. A workgroup loads a small square from each input matrix into shared memory, then uses those values for several output cells.

Each invocation contributes one value from A and one from B:

for tile in 0..<((dims.k + dims.tileSize - 1) div dims.tileSize):
  let aCol = tile * dims.tileSize + localCol
  let bRow = tile * dims.tileSize + localRow
  let sharedIndex = localRow * dims.tileSize + localCol

  shared.a[sharedIndex] =
    if row < dims.m and aCol < dims.k: b.a[row * dims.k + aCol] else: 0
  shared.b[sharedIndex] =
    if bRow < dims.k and col < dims.n: b.b[bRow * dims.n + col] else: 0

  memoryBarrier()
  barrier()

  for i in 0..<dims.tileSize:
    sum += shared.a[localRow * dims.tileSize + i] *
           shared.b[i * dims.tileSize + localCol]

  memoryBarrier()
  barrier()

The first barrier separates loading from reading. The second separates reading from the next load. Matrix edges are padded with zero so those invocations follow the same barriers as everyone else.

An element placed in shared.a is reused across a row of output invocations. An element in shared.b is reused down a column. I find that detail more revealing than the usual description of tiling as an optimization. The dot product has not changed at all. What changed is the route its operands take to reach it.

These experiments changed how I approach parallel code. I start with the direct version, even when a subgroup operation does most of the work, then make one part explicit at a time. Once the ownership and result are correct, I introduce shared memory and synchronization, keeping each barrier tied to something concrete: data becoming ready to read or storage becoming safe to reuse. I find that more helpful than trying to understand an optimized version all at once.

The complete programs are available in this gist.

Next, I would replace subgroupInclusiveAdd with a shuffle-based scan, then extend it across several subgroups. After that, the blocked reduction introduces coarsening and subgroup shuffles without becoming too difficult to follow. I would save the single-pass version for later, when cross-workgroup coordination can be studied on its own.