Task 5 of 5

One Kernel, Any Size

Every kernel so far had its size welded on: output: [16, 16], loop to 16. Real code multiplies whatever matrices show up. gpu.js has three switches for that: dynamicOutput: true lets you call kernel.setOutput([n, n]) before each run, dynamicArguments: true lets the input arrays change size between calls, and loopMaxIterations raises the safety cap so the loop bound can be a runtime argument instead of a constant.

Pass the size in as a plain number, loop k < size, and one kernel object serves an 8×8 and a 48×48 multiply back to back. This is the payoff of the module: the naive triple loop from task 2, now packaged as a function that scales.

Goal: make multiply(a, b) work for any square size up to 64 using a single kernel — verify it on the 8×8 and 48×48 pairs provided.

Requirements

Hint 1 — why the cap?

On the GPU backend a loop bound that isn't a compile-time constant becomes

for (i = 0; i < LOOP_MAX; i++) {
  if (!(i < size)) break;
  // …
}

in the shader — loopMaxIterations is that LOOP_MAX. Set it to the largest size you'll ever pass: 64 here.

Hint 2 — sizing per call

Inside multiply — which is async, because a kernel call is awaited:

const n = a.length;
matmul.setOutput([n, n]);
return await matmul(a, b, n);

— set the launch shape first, then invoke with the size as the last argument. Its callers then await multiply(…) in turn.

Hint 3 — the kernel
function (a, b, size) {
  let sum = 0;
  for (let k = 0; k < size; k++) {
    sum += a[this.thread.y][k] * b[k][this.thread.x];
  }
  return sum;
}

with options

{
  dynamicOutput: true,
  dynamicArguments: true,
  loopMaxIterations: 64,
}

Same idea elsewhere

Shipping one kernel that covers a size range is standard practice everywhere: CUDA kernels take M, N, K as launch parameters and pick grid dimensions at call time, WebGPU dispatches a runtime-computed number of workgroups, and Metal binds sizes through a constant buffer. Compile once, launch at any size — exactly what you just built.

All tasks in Matrix Multiply

  1. One Cell, One Dot Product
  2. The Full Grid: Matrix × Matrix
  3. Rectangular: Three Different Sizes
  4. Transpose: Swap the Axes
  5. One Kernel, Any Size

This page is an interactive exercise — the editor, the GPU runner and your saved progress need JavaScript. The text above is the full brief.