Optimizing a WebGPU Matmul Kernel using Subgroups

High-performance matrix multiplication is typically handled by vendor-specific SDKs like CUDA or Metal. WebGPU offers a portable, cross-platform alternative that can target multiple platforms – including the web. For example, Chrome’s WebGPU implementation, Dawn, supports backends such as Direct3D 12, Metal, Vulkan, and OpenGL, enabling the same compute code to run across a wide range of devices and operating systems.

In this post, we’ll walk through a series of optimizations to a matrix multiplication kernel written in WebGPU. Interestingly, the biggest performance gains don’t come from reducing the number of operations, but rather from improving how memory is accessed -— particularly how data is loaded, reused, and stored across threads.

For each version of the kernel, we’ll estimate its achieved floating-point performance in GFLOPs/s (billion operations per second). The table below shows benchmark results from my M4 MacBook Air, running in Google Chrome on 4096×4096 matrices.

1. Naive

In WebGPU, compute kernels are executed by thousands of threads running in parallel. For our baseline implementation, we’ll take a straightforward approach: assign one thread to compute each cell of the output matrix. For square matrices of size N×N, this means launching N×N threads – one for every element in the output matrix.

The kernel operates on three buffers:

• matrixA and matrixB are the input matrices.
• matrixC is the output matrix.

All matrices are stored in row-major order, meaning that elements are laid out row by row in memory: all elements of the first row come first, followed by the second row, and so on.

@group(0) @binding(0) var<storage, read> matrixA: array<f32>;
@group(0) @binding(1) var<storage, read> matrixB: array<f32>;
@group(0) @binding(2) var<storage, read_write> matrixC: array<f32>;

// All matrices are square.
const SIZE: u32 = 4096;

fn index(row: u32, col: u32) -> u32 {
    // Row-major order.
    return row * SIZE + col;
}

@compute @workgroup_size(1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  // The global invocation ID represents the two-dimensional
  // index of the current thread. Here, we assign the x-dimension
  // to rows and the y-dimension to columns. Because we're working
  // with square matrices, this choice is arbitrary.
  let row: u32 = gid.x;
  let col: u32 = gid.y;

  // Each thread computes the dot-product of the gid.x-th row of A
  // and the gid.y-th column of B.
  var sum: f32 = 0;
  for (var k: u32 = 0; k < SIZE; k++) {
    sum += matrixA[index(row, k)] * matrixB[index(k, col)];
  }
  matrixC[index(row, col)] = sum;
}

Threads are organized into one-, two- or three-dimensional workgroups. And, when launching a compute pipeline, you specify many workgroups to launch and how many threads each workgroup contains.

In this kernel, we’ll launch NxN two-dimensional workgroups, each of which with 1x1 threads (i.e. a single thread).

const SIZE = 4096;

// Set up the compute pipeline and dispatch workgroups.
const commands = device.createCommandEncoder();
const pass = commands.beginComputePass();
pass.setPipeline(pipeline);
pass.setBindGroup(0, bindGroup); // Attach input and output buffers.
pass.dispatchWorkgroups(SIZE, SIZE);
pass.end();
device.queue.submit([commands.finish()]); // Go!

2. Workgroups

In the example above, we used workgroups of size 1x1, meaning each workgroup computes just one output cell. As a result, we must launch N×N workgroups to cover the entire output matrix. Although this is okay from a correctness standpoint, it’s inefficient because there is significant overhead in launching each workgroup. For the next kernel, we’ll launch fewer but larger workgroups.

How do we pick the right workgroup size? The optimal workgroup size depends heavily on the hardware running the kernel. If workgroups are too small, performance suffers due to the overhead of managing and launching many separate groups — especially when each one computes just a single output element.

But what limits how large a workgroup can be? To answer that, we need to look at how GPU hardware is structured.

GPUs organize compute resources into processing units (often called cores, or streaming multiprocessors in CUDA terminology). Each workgroup runs on a single core, and each core can execute multiple workgroups simultaneously. However, cores have limited on-chip resources — most importantly:

• Registers, which are allocated per thread, and
• Shared memory, which is allocated per workgroup (we’ll discuss this later)

If a kernel uses a lot of registers per thread, and the workgroup has many threads, the total register usage may exceed what the core can support concurrently. In that case, the GPU must run fewer workgroups at a time, reducing overall parallelism. This situation is known as low occupancy — a state where the hardware is underutilized due to resource pressure. See CUDA’s occupancy calculator for more context.

WebGPU also imposes standard limits on workgroup dimensions. According to the WebGPU specification, implementations must support at least 256 threads in both the X and Y dimensions (maxComputeWorkgroupSizeX and maxComputeWorkgroupSizeY), though many backends exceed this. For example, my laptop supports workgroups of up to 1024×1024. You can check your own system’s limits at webgpureport.org.

To adjust the size of each workgroup, we modify the @workgroup_size attribute in the shader. For example, the following kernel uses a 16×16 workgroup:

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  // The `global_invocation_id` takes care of calculating
  // the overall thread id, regardless of whether they are
  // subdivided into workgroups.
  // To get the workgroup-local thread ID, we'd use
  // `@builtin(local_invocation_id)` instead.
}

With a 16×16 configuration, each workgroup now computes a 16×16 block of the output matrix — meaning we need significantly fewer workgroups overall compared to the naive 1×1 version.

Since our matrix dimensions (4096×4096) divide evenly by 16, we can skip bounds checks in this example. In general, however, you’d want to guard against out-of-bounds accesses when the matrix size isn’t a clean multiple of the workgroup size. To keep this post focused on performance, we’ll omit such edge handling throughout.

In JavaScript, we update the dispatch call accordingly:

const SIZE = 4096;

const workgroupSize = 16;
pass.dispatchWorkgroups(
  Math.ceil(SIZE / workgroupSize),
  Math.ceil(SIZE / workgroupSize)
);

This change alone increases the performance of our kernel from 86.9 GFLOPs/s to 240.1 GFLOPs/s, almost a 3x improvement!

3. Global Memory Coalescing

Our current kernel is very likely limited by memory bandwidth. This means the GPU is spending more time waiting for data to arrive from memory than it is performing actual computation. To address this bottleneck, we once again need to consider the details of how memory access works on modern GPU hardware.

Modern GPUs fetch global memory in aligned chunks — typically 32, 64, or 128 bytes at a time. When multiple threads read from adjacent memory addresses, these accesses can be coalesced into a single memory transaction. This is far more efficient than issuing many small, scattered reads.

Threads on the GPU are grouped into fixed-size units — typically 32 threads — called warps (in CUDA) or subgroups (in WebGPU). Even though WebGPU doesn’t expose subgroups by default, they still matter: memory coalescing only happens between threads within the same subgroup.

To maximize coalescing, we want threads within a subgroup to read contiguous memory. But there’s a catch: WebGPU does not define how 2D thread IDs map to subgroups — this mapping is left up to the implementation.

There are two common possibilities:

• Row-major subgrouping: Threads like (0, 0), (0, 1), (0, 2) belong to the same subgroup.
• Column-major subgrouping: Threads like (0, 0), (1, 0), (2, 0) belong to the same subgroup.

On my MacBook (using Metal under the hood), the second pattern —- column-major -— is used. This affects how memory coalescing plays out.

Let’s look at a problematic scenario in our current kernel. In the following diagram, we look consider the access pattern of four contiguous threads ((2, 4), (3, 4), (4, 4), and (5, 4)) at a point in time when k=1. In other words, we’re looking the second iteration of the for-loop that calculates the dot-product for every cell in the output matrix.

• Each of the four threads in the subgroup access a different element in matrix A: ((2, 1), (3, 1), (4, 1), and (5, 1)). Because these are not contiguous in memory (remember these matrices are stored in row-major order, not column-major order), the GPU cannot coalesce the memory transfers. ❌
• Each of the four threads in the subgroup access the same element (1, 4) in matrix B, so a single memory transfer request is made. –>

But what if we make a small tweak in how we assign thread indices to locations in the output matrix?

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let row: u32 = gid.y;  // Previously "gid.x"
  let col: u32 = gid.x;  // Previously "gid.y"

  // [...] Same as before...
}

Now, the memory access pattern is quite different. Again, let’s look at four contiguous threads ((2, 1), (2, 3), (2, 4), and (2, 5)) when k=1.

• Each of the four threads in the subgroup access the same element (2, 1) in matrix A, so a single memory transfer request is made.
• Each of the four threads in the subgroup access a different element in matrix B: ((2, 1), (2, 2), (2, 3), and (2, 4)). Because these are contiguous in memory (row-major order), the GPU can coalesce the memory transfers. ✅

Running the benchmark again, we see that the throuput increases from 240.1 GFLOPs/s to 410.9 GFLOPs/s – almost 2x faster by simply changing the order in which we calculate the output matrix!

4. Shared Memory Coalescing

In the previous kernel, we improved performance by coalescing global memory accesses. However, we’re still reading the same input values from global memory many times — especially for values reused across multiple output elements. For example, all threads computing a row of the output matrix C reuse the same row of A, but each thread reloads it independently from slow global memory.

To reduce this redundancy, we can take advantage of workgroup-local (shared) memory. This kind of memory is co-located with the GPU cores and is both much faster and much but much smaller than global memory (~100s of KiB per core). Threads in the same workgroup can collaboratively load data into shared memory, reuse it efficiently, and drastically cut down on costly global reads.

Since entire matrices won’t fit into shared memory, we split them into 16×16 tiles, load one tile at a time, compute partial results, and repeat.

1. Each thread loads one value from A and one from B into shared memory. Collectively, all threads in the workgroup load the entire tiles for A and B.
2. All threads synchronize using workgroupBarrier() to ensure the tile is fully loaded.
3. Each thread performs a partial dot product using the cached tile values.
4. This process repeats across tiles, accumulating the final result for one output cell.

It’s important to note that, for the first time, we’re using multiple threads in a workgroup cooperatively: the tiles populated by a given thread are read by many others. This, however, requires ensuring that all threads are done writing their assigned values to shared memory before other threads start reading them. We do this by invoking workgroupBarrier(), which ensures that all threads in the workgroup reach a particular point in the execution of the kernel before any of them can proceed.

@group(0) @binding(0) var<storage, read> matrixA: array<f32>;
@group(0) @binding(1) var<storage, read> matrixB: array<f32>;
@group(0) @binding(2) var<storage, read_write> matrixC: array<f32>;

// 16x16 tiles of input matrices A and B.
var<workgroup> cacheA: array<array<f32, 16>, 16>;
var<workgroup> cacheB: array<array<f32, 16>, 16>;

// All matrices are square.
const SIZE: u32 = 4096;

fn index(row: u32, col: u32) -> u32 {
  // Row-major order.
  return row * SIZE + col;
}

@compute @workgroup_size(16, 16)
fn main(
  @builtin(global_invocation_id) gid: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let row: u32 = gid.y;
  let col: u32 = gid.x;

  // This variable holds the value of `C` at (`row`, `col`).
  var sum: f32 = 0;

  for (var tileOffset: u32 = 0; tileOffset < SIZE; tileOffset += 16) {
    // Each workgroup populates a 16x16 tile in the output matrix `C`.
    // Each thread in the workgroup is responsible for a single cell of A
    // and a single cell of B.
    cacheA[lid.y][lid.x] = matrixA[index(row, tileOffset + lid.x)];
    cacheB[lid.y][lid.x] = matrixB[index(tileOffset + lid.y, col)];
    workgroupBarrier();

    // Once the tiles are loaded, we do a partial dot-product over the
    // values in the current input tiles. We'll accumulate these partial
    // dot-products over all tiles to compute the final value.
    for (var k: u32 = 0; k < 16; k++) {
      sum += cacheA[lid.y][k] * cacheB[k][lid.x];
    }
    workgroupBarrier();
  }

  matrixC[index(row, col)] = sum;
}

With these changes, we reach 483.9 GFLOPs/s, a moderate improvement over the previous kernel, which achieved 410.9 GFLOPs/s (+18%). The reason why the increase in performance is not as large is likely that the various hierarchical caches in the GPU cores were likely preventing many of the redundant reads from reaching the GPU’s global memory.

5. Thread Tiling

In Kernel #2, we saw that launching many threads per workgroup was significantly faster than launching a single thread per workgroup because it amortized the performance cost of each workgroup. Similarly, on this kernel, we’ll calculate multiple output cells per thread. More specifically, each thread will calculate a block of 4x4 cells, and we’ll launch 16x fewer threads. This amortizes the cost of launching each thread and results in overall higher performance.

There’s two important hyperparameters that will significantly impact performance:

• The size of the block of elements computed by a single thread:
  • Too low means that the per-thread overhead doesn’t get amortized across a large number of elements.
  • Too high causes each kernel thread to require an excessive number of registers. Because registers are a finite resource in each core, this can lead to fewer workgroups being launched per core, reducing occupancy making poor use of the GPU hardware resources.
• The number of threads per workgroup:
  • Too low means that the per-workgroup overhead doesn’t get amortized across a large number of elements.
  • Too high causes each workgrouo to require an excessive amount of shared memory. Because shared memory is a finite resource in each core, this can lead to fewer workgroups being launched per core, reducing occupancy making poor use of the GPU hardware resources.

On my M4 MacBook Air, workgroups of size 8x8 with each thread working on blocks of size 4x4 seems to work best. This particular kernel is almost 2x as fast as the previous one, going from 483.9 GFLOPs/s to 951.3 GFLOPs/s.

@group(0) @binding(0) var<storage, read> matrixA: array<f32>;
@group(0) @binding(1) var<storage, read> matrixB: array<f32>;
@group(0) @binding(2) var<storage, read_write> matrixC: array<f32>;

var<workgroup> cacheA: array<array<f32, 32>, 32>;
var<workgroup> cacheB: array<array<f32, 32>, 32>;

const SIZE: u32 = 4096;

fn index(row: u32, col: u32) -> u32 {
  // Row-major order.
  return row * SIZE + col;
}

@compute @workgroup_size(8, 8)
fn main(
  @builtin(global_invocation_id) gid: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let row: u32 = gid.y;
  let col: u32 = gid.x;

  var sum: array<array<f32, 4>, 4>;
  for (var tileOffset: u32 = 0; tileOffset < SIZE; tileOffset += 32) {
    // Each thread loads a block of 4x4 cells of A and B.
    for (var i: u32 = 0; i < 4; i++) {
      for (var j: u32 = 0; j < 4; j++) {
        cacheA[4 * lid.y + i][4 * lid.x + j] =
            matrixA[index(4 * row + i, tileOffset + 4 * lid.x + j)];
        cacheB[4 * lid.y + i][4 * lid.x + j] =
            matrixB[index(tileOffset + 4 * lid.y + i, 4 * col + j)];
      }
    }
    workgroupBarrier();

    // Each thread calculates the dot-products for a block of 4x4 cells in C.
    for (var k: u32 = 0; k < 32; k += 1) {
      for (var i: u32 = 0; i < 4; i++) {
        for (var j: u32 = 0; j < 4; j++) {
          sum[i][j] += cacheA[4 * lid.y + i][k] * cacheB[k][4 * lid.x + j];
        }
      }
    }
    workgroupBarrier();
  }

  // Each thread writes a block of 4x4 cells in C.
  for (var i: u32 = 0; i < 4; i++) {
    for (var j: u32 = 0; j < 4; j++) {
      matrixC[index(4 * row + i, 4 * col + j)] = sum[i][j];
    }
  }
}

6. Thread Tiling (Unrolled)

The previous kernel iterates over every cell assigned to each thread in various for-loops with statically known bounds. Although some compiles compilers will automatically unroll these for performance, others will not. Manually unrolling them can sometimes improve performance due to:

• Fewer instructions per kernel thread (less likely), or
• Linear flow of instructions exposing more parallelism to the instruction scheduler (more likely).

In my benchmarkes, unrolling all loops increases performance from 951.3 GFLOPs/s to 1074.2 GFLOPs/s.

7. Vectorized Buffers

For this kernel, we’ll change the way we read and write from global memory.

First, instead of treating the input as a buffer of f32s, we’ll treat them as buffers of vec4<f32>, where each element in the buffer represents 4 elements in the matrix. This has the benefit of allowing the kernel to issue a single instruction to load/store 4 elements at a time from global/shared memory.

And second, we’ll transpose input matrix B. Together with treating the matrix as chunks of 4 elements, this allows calculating the partial dot-products using WGSL’s built-in dot function, which can sometimes translate into more efficient machine code.

These two changes increase performance from 1074.2 GFLOPs/s to 1349.4 GFLOPs/s, a nice 25% boost.

// All of these were previously buffers of type `array<f32>`.
@group(0) @binding(0) var<storage, read> matrixA: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read> matrixB: array<vec4<f32>>;
@group(0) @binding(2) var<storage, read_write> matrixC: array<vec4<f32>>;

// The shared memory is also composed of arrays of `vec<f32>` instead of `f32`.
var<workgroup> cacheA: array<array<vec4<f32>, 8>, 32>;
var<workgroup> cacheB: array<array<vec4<f32>, 8>, 32>;

const SIZE: u32 = 4096;

@compute @workgroup_size(8, 8)
fn main(
  @builtin(global_invocation_id) gid: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  var sum: array<array<f32, 4>, 4>;
  for (var tileOffset: u32 = 0; tileOffset < SIZE; tileOffset += 32) {
    // Populate shared memory, 4 cells at a time.
    cacheA[4 * lid.y + 0][lid.x] = matrixA[(4 * gid.y + 0) * SIZE / 4 + tileOffset/4 + lid.x];
    cacheA[4 * lid.y + 1][lid.x] = matrixA[(4 * gid.y + 1) * SIZE / 4 + tileOffset/4 + lid.x];
    cacheA[4 * lid.y + 2][lid.x] = matrixA[(4 * gid.y + 2) * SIZE / 4 + tileOffset/4 + lid.x];
    cacheA[4 * lid.y + 3][lid.x] = matrixA[(4 * gid.y + 3) * SIZE / 4 + tileOffset/4 + lid.x];
    cacheB[4 * lid.x + 0][lid.y] = matrixB[(4 * gid.x + 0) * SIZE / 4 + tileOffset/4 + lid.y];
    cacheB[4 * lid.x + 1][lid.y] = matrixB[(4 * gid.x + 1) * SIZE / 4 + tileOffset/4 + lid.y];
    cacheB[4 * lid.x + 2][lid.y] = matrixB[(4 * gid.x + 2) * SIZE / 4 + tileOffset/4 + lid.y];
    cacheB[4 * lid.x + 3][lid.y] = matrixB[(4 * gid.x + 3) * SIZE / 4 + tileOffset/4 + lid.y];
    workgroupBarrier();

    for (var k: u32 = 0; k < 8; k += 1) {
      let cacheA0: vec4<f32> = cacheA[4 * lid.y + 0][k];
      let cacheA1: vec4<f32> = cacheA[4 * lid.y + 1][k];
      let cacheA2: vec4<f32> = cacheA[4 * lid.y + 2][k];
      let cacheA3: vec4<f32> = cacheA[4 * lid.y + 3][k];

      let cacheB0: vec4<f32> = cacheB[4 * lid.x + 0][k];
      let cacheB1: vec4<f32> = cacheB[4 * lid.x + 1][k];
      let cacheB2: vec4<f32> = cacheB[4 * lid.x + 2][k];
      let cacheB3: vec4<f32> = cacheB[4 * lid.x + 3][k];

      // Perform all the dot-products in chunks of 4 elements.
      sum[0][0] += dot(cacheA0, cacheB0);
      sum[0][1] += dot(cacheA0, cacheB1);
      sum[0][2] += dot(cacheA0, cacheB2);
      sum[0][3] += dot(cacheA0, cacheB3);
      sum[1][0] += dot(cacheA1, cacheB0);
      sum[1][1] += dot(cacheA1, cacheB1);
      sum[1][2] += dot(cacheA1, cacheB2);
      sum[1][3] += dot(cacheA1, cacheB3);
      sum[2][0] += dot(cacheA2, cacheB0);
      sum[2][1] += dot(cacheA2, cacheB1);
      sum[2][2] += dot(cacheA2, cacheB2);
      sum[2][3] += dot(cacheA2, cacheB3);
      sum[3][0] += dot(cacheA3, cacheB0);
      sum[3][1] += dot(cacheA3, cacheB1);
      sum[3][2] += dot(cacheA3, cacheB2);
      sum[3][3] += dot(cacheA3, cacheB3);
    }
    workgroupBarrier();
  }

  matrixC[(4 * gid.y + 0) * SIZE / 4 + gid.x] = vec4(sum[0][0], sum[0][1], sum[0][2], sum[0][3]);
  matrixC[(4 * gid.y + 1) * SIZE / 4 + gid.x] = vec4(sum[1][0], sum[1][1], sum[1][2], sum[1][3]);
  matrixC[(4 * gid.y + 2) * SIZE / 4 + gid.x] = vec4(sum[2][0], sum[2][1], sum[2][2], sum[2][3]);
  matrixC[(4 * gid.y + 3) * SIZE / 4 + gid.x] = vec4(sum[3][0], sum[3][1], sum[3][2], sum[3][3]);
}

8. Subgroup Multiply

Within each workgroup, threads are launched in bundles called warps (in CUDA) or SIMD-groups (in Metal). At the hardware level, these threads are closer together and can access each other’s registers without going through shared or global memory. Additionally, there’s new APIs to use threads within each bundle to collaborately calculate small chunks of matrix multiplication.

Unfortunately, there are sigificant differences in how each individual vendor has implemented these primitives, and the API WebGPU has not yet been standardized. However, the talented folks that build Dawn have implemented a tentative API for experimentation.

• subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(p, offset, col_major, stride)
  • uses all threads in the subgroup to collectively load a 8x8 subtile of a left matrix at offset *p, advancing the pointer stride elements to reach the next column/row depending on the value of col_major. The analogous function exists for right matrices.
• subgroupMatrixMultiplyAccumulate(a, b, c)
  • uses all threads in the subgroup to collectively calculate a @ b + c, where a, b, and c, are of type subgroup_matrix_left, subgroup_matrix_right, and subgroup_matrix_result, respectively.
• subgroupMatrixStore(p, offset, result, col_major, stride)
  • uses all threads in the subgroup to collectively store result in buffer p at a given offset, advancing the pointer stride elements to reach the next column/row depending on the value of col_major.

Although CUDA specifes that all warps are 32 threads wide, other GPU vendors do not. In fact, many of them only determine the subgroup size after compiling the kernel. Because of this, WebGPU can only expose the minimum and maximum subgroup size before compiling the kernel. For my MacBook Air, the maximum subgroup size is 64, although, in every test I’ve run, the subgroups that end up being invoked have size 32. For the purpose of this benchmark, I’ll hardcode this in the kernel. For real production use-cases, we’d check the subgroup size during kernel execution and use different specializations for each one.

At a high level, the algorithm is as follows:

• Each workroup is of size 128x1. This will give us 4 subgroups, each of which with 32 threads.
• Within each workgroup, we first load tiles of A and B into shared memory, just like before.
• For each subgroup, we will calculate partial sums of 4 matrices, each of which has a size of 8x8 (required for Metal). We do this by loading two 8x8 sub-matrices from A’s shared cache, and two other 8x8 sub-matrices from B’s shared cache. We use them to accumulate the partial products for the 4 sub-matrices in the subgroup, and then continue the loop by loading the next sub-matrices from A and B.
• Finally, we store all aggregate subgroup sub-matrices back into global memory.

Using subgroups increases performance from 1349.4 GFLOPs/s to 1715.5 GFLOPs/s, a nice 27% improvement. There’s a couple of caveats, however:

• As we mentioned before, this API is not yet standardized or thoroughly tested. That said, ONNX has implemented support for it and got excellent results (~3x speed boost in some key benchmarks).
• For now, using subgroup-based matrix multiplication requires running Chrome Canary and enabling the flag for #enable-unsafe-webgpu.
• We also need to disable Chrome’s chromium_disable_uniformity_analysis. Although subgroup operations require that all threads in the subgroup use the same invocation parameters, this check seems to be forcing all threads in the workgroup to use the same parameters, which is overly restrictive.

All that said, Dawn seems to be making steady progress! I’m excited about this and will be following the GitHub issue tracking support for WebGPU subgroup-based matrix multiplication.

enable subgroups;
enable chromium_experimental_subgroup_matrix;
enable chromium_disable_uniformity_analysis;

@group(0) @binding(0) var<storage, read> matrixA: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read> matrixB: array<vec4<f32>>;
@group(0) @binding(2) var<storage, read_write> matrixC: array<vec4<f32>>;

var<workgroup> cacheA: array<f32, 1024>;
var<workgroup> cacheB: array<f32, 1024>;
var<workgroup> cacheC: array<f32, 1024>;

@compute @workgroup_size(128, 1)
fn main(
    @builtin(local_invocation_index) lidx: u32,
    @builtin(workgroup_id) wid: vec3<u32>,
    @builtin(subgroup_id) sid: u32,
    @builtin(subgroup_size) ss: u32) {
  if (ss != 32) { return; }

  let wg_origin_row: u32 = 32 * wid.y;
  let wg_origin_col: u32 = 32 * wid.x;

  let store_local_row: u32 = (lidx * 8) / 32;
  let store_local_col: u32 = (lidx * 8) % 32;

  let store_row: u32 = wg_origin_row + store_local_row;
  let store_col: u32 = wg_origin_col + store_local_col;

  let sg_local_row: u32 = 16 * (sid % 2);
  let sg_local_col: u32 = 16 * (sid / 2);

  let sg_row: u32 = wg_origin_row + sg_local_row;
  let sg_col: u32 = wg_origin_col + sg_local_col;

  var sgMatrixC00: subgroup_matrix_result<f32, 8, 8>;
  var sgMatrixC01: subgroup_matrix_result<f32, 8, 8>;
  var sgMatrixC10: subgroup_matrix_result<f32, 8, 8>;
  var sgMatrixC11: subgroup_matrix_result<f32, 8, 8>;

  for (var k: u32 = 0; k < SIZE; k += 32) {
    let load_a_row: u32 = store_row;
    let load_a_col: u32 = store_local_col + k;
    let load_b_row: u32 = store_local_row + k;
    let load_b_col: u32 = store_col;

    let idxA: u32 = SIZE * load_a_row + load_a_col;
    let idxB: u32 = SIZE * load_b_row + load_b_col;

    let mA0: vec4<f32> = matrixA[(idxA / 4) + 0];
    let mA1: vec4<f32> = matrixA[(idxA / 4) + 1];
    let mB0: vec4<f32> = matrixB[(idxB / 4) + 0];
    let mB1: vec4<f32> = matrixB[(idxB / 4) + 1];

    cacheA[32 * store_local_row + store_local_col + 0] = mA0.x;
    cacheA[32 * store_local_row + store_local_col + 1] = mA0.y;
    cacheA[32 * store_local_row + store_local_col + 2] = mA0.z;
    cacheA[32 * store_local_row + store_local_col + 3] = mA0.w;
    cacheA[32 * store_local_row + store_local_col + 4] = mA1.x;
    cacheA[32 * store_local_row + store_local_col + 5] = mA1.y;
    cacheA[32 * store_local_row + store_local_col + 6] = mA1.z;
    cacheA[32 * store_local_row + store_local_col + 7] = mA1.w;

    cacheB[32 * store_local_row + store_local_col + 0] = mB0.x;
    cacheB[32 * store_local_row + store_local_col + 1] = mB0.y;
    cacheB[32 * store_local_row + store_local_col + 2] = mB0.z;
    cacheB[32 * store_local_row + store_local_col + 3] = mB0.w;
    cacheB[32 * store_local_row + store_local_col + 4] = mB1.x;
    cacheB[32 * store_local_row + store_local_col + 5] = mB1.y;
    cacheB[32 * store_local_row + store_local_col + 6] = mB1.z;
    cacheB[32 * store_local_row + store_local_col + 7] = mB1.w;

    workgroupBarrier();

    for (var q: u32 = 0; q < 32; q += 8) {
      let sgMatrixA0: subgroup_matrix_left<f32, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(
          &cacheA, /*offset=*/32 * (sg_local_row + 0) + q, /*col_major=*/false, /*stride=*/32);
      let sgMatrixA1: subgroup_matrix_left<f32, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(
          &cacheA, /*offset=*/32 * (sg_local_row + 8) + q, /*col_major=*/false, /*stride=*/32);

      let sgMatrixB0: subgroup_matrix_right<f32, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(
          &cacheB, /*offset=*/32 * q + (sg_local_col + 0), /*col_major=*/false, /*stride=*/32);
      let sgMatrixB1: subgroup_matrix_right<f32, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(
          &cacheB, /*offset=*/32 * q + (sg_local_col + 8), /*col_major=*/false, /*stride=*/32);

      sgMatrixC00 = subgroupMatrixMultiplyAccumulate(sgMatrixA0, sgMatrixB0, sgMatrixC00);
      sgMatrixC01 = subgroupMatrixMultiplyAccumulate(sgMatrixA0, sgMatrixB1, sgMatrixC01);
      sgMatrixC10 = subgroupMatrixMultiplyAccumulate(sgMatrixA1, sgMatrixB0, sgMatrixC10);
      sgMatrixC11 = subgroupMatrixMultiplyAccumulate(sgMatrixA1, sgMatrixB1, sgMatrixC11);
    }

    workgroupBarrier();
  }

  subgroupMatrixStore(&cacheC, /*offset=*/32 * (sg_local_row + 0) + (sg_local_col + 0), sgMatrixC00, /*col_major=*/false, /*stride=*/32);
  subgroupMatrixStore(&cacheC, /*offset=*/32 * (sg_local_row + 0) + (sg_local_col + 8), sgMatrixC01, /*col_major=*/false, /*stride=*/32);
  subgroupMatrixStore(&cacheC, /*offset=*/32 * (sg_local_row + 8) + (sg_local_col + 0), sgMatrixC10, /*col_major=*/false, /*stride=*/32);
  subgroupMatrixStore(&cacheC, /*offset=*/32 * (sg_local_row + 8) + (sg_local_col + 8), sgMatrixC11, /*col_major=*/false, /*stride=*/32);

  workgroupBarrier();

  let idxC: u32 = SIZE * store_row + store_col;
  matrixC[(idxC / 4) + 0] = vec4<f32>(
      cacheC[32 * store_local_row + store_local_col + 0],
      cacheC[32 * store_local_row + store_local_col + 1],
      cacheC[32 * store_local_row + store_local_col + 2],
      cacheC[32 * store_local_row + store_local_col + 3]);
  matrixC[(idxC / 4) + 1] = vec4<f32>(
      cacheC[32 * store_local_row + store_local_col + 4],
      cacheC[32 * store_local_row + store_local_col + 5],
      cacheC[32 * store_local_row + store_local_col + 6],
      cacheC[32 * store_local_row + store_local_col + 7]);
}

Final Thoughts

And there you have it! We took a naive matrix multiplication kernel running at 87 GFLOPs/s and optimized it all the way up to 1715 GFLOPs/s — a solid 20x performance improvement through careful attention to how GPUs handle memory and computation.

What’s particularly interesting is that most of these gains didn’t come from algorithmic improvements or reducing operation counts. Instead, the biggest wins came from understanding and working with the hardware: coalescing memory accesses, utilizing shared memory effectively, and coordinating threads to maximize throughput.

WebGPU is still evolving, and some of the more advanced features like subgroup matrix operations are currently behind experimental flags. But the ability to write portable, high-performance compute shaders that run across different platforms and devices is genuinely exciting for the future of web-based computing.

Further Reading

• Nvidia’s programming guide for CUDA
• Simon Boehm’s excellent post on “How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance”
• Zach Nussbaum’s (also excellent) post for WebGPU: “Optimizing a WebGPU Matmul Kernel for 1TFLOP+ Performance”