WebGPU 行列積の
限界を測る

行列積（Matrix Multiplication / GEMM）は機械学習・グラフィックス・科学計算の中核となる演算です。 このベンチマークでは CPU 実装、WebGPU ナイーブ実装、 WebGPU タイル化実装の3種を、同一の入力に対して計測・比較します。

N×N 行列同士の積 C = A × B は、各要素 C[i][j] が「A の i 行目」と「B の j 列目」の内積で計算されます。 N が大きくなると計算量は O(N³) で爆発的に増えますが、 各要素の計算は完全に独立しているため並列化に向いています。

Tiled 実装の核心は workgroup shared memory です。 GPU は通常、各スレッドが独立したレジスタを持ちますが、ワークグループ内のスレッドは 共有メモリ（var<workgroup>）を通じてデータをやり取りできます。 これがタイル化アルゴリズムを成立させる鍵です。

↑ ワークグループ内 256 スレッドが同じタイルを共有 → グローバルメモリアクセスを 256回 → 1回 に削減

本ベンチマークの中核となる3つの実装を、初学者向けに解説します。

function matmulCPU(A, B, N) {
  const C = new Float32Array(N * N);
  for (let i = 0; i < N; i++) {
    for (let j = 0; j < N; j++) {
      let sum = 0;
      for (let k = 0; k < N; k++) {
        sum += A[i * N + k] * B[k * N + j];
      }
      C[i * N + j] = sum;
    }
  }
  return C;
}

最も基本的な3重ループ実装です。N=512 で約 1.34 億回の演算が必要となり、ブラウザの JavaScript エンジンでは数百ms〜数秒かかります。

// 各スレッドが C[row][col] の 1 要素を担当する
@group(0) @binding(0) var<storage, read>       A: array<f32>;
@group(0) @binding(1) var<storage, read>       B: array<f32>;
@group(0) @binding(2) var<storage, read_write> C: array<f32>;
@group(0) @binding(3) var<uniform>             dim: u32;

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let row = gid.y;
  let col = gid.x;
  if (row >= dim || col >= dim) { return; }

  var sum = 0.0;
  for (var k: u32 = 0u; k < dim; k = k + 1u) {
    sum = sum + A[row * dim + k] * B[k * dim + col];
  }
  C[row * dim + col] = sum;
}

N×N スレッドが同時に走りますが、各スレッドは A・B から 2N 回グローバルメモリを読みます。 同じ行・列を別スレッドが何度も読むため、メモリ帯域が無駄遣いされるのが弱点です。

const TS: u32 = 16u;  // タイルサイズ

// ワークグループ全体で共有するメモリ（全スレッドが同じデータを参照可能）
var<workgroup> tileA: array<array<f32, TS>, TS>;
var<workgroup> tileB: array<array<f32, TS>, TS>;

@compute @workgroup_size(16, 16)
fn main(
  @builtin(global_invocation_id) gid: vec3u,
  @builtin(local_invocation_id)  lid: vec3u
) {
  let row = gid.y;
  let col = gid.x;
  var sum = 0.0;
  let numTiles = (dim + TS - 1u) / TS;

  for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
    // ① タイル分のデータを共有メモリにロード（256スレッドが協力）
    tileA[lid.y][lid.x] = A[row * dim + (t * TS + lid.x)];
    tileB[lid.y][lid.x] = B[(t * TS + lid.y) * dim + col];

    // ② 全スレッドのロード完了を待つ
    workgroupBarrier();

    // ③ 共有メモリ上で内積計算（高速）
    for (var k: u32 = 0u; k < TS; k = k + 1u) {
      sum = sum + tileA[lid.y][k] * tileB[k][lid.x];
    }
    workgroupBarrier();
  }

  if (row < dim && col < dim) {
    C[row * dim + col] = sum;
  }
}

ハイライト部分が最適化の核心です。workgroupBarrier() で「ワークグループ内全スレッドのロード完了」を待ち合わせ、その後で共有メモリから内積計算を行います。 各タイルは1度だけグローバルメモリから読まれ、256スレッドで再利用されます。

// GPU の処理完了を確実に待つには readBuffer 経由が必要
async function measureGPU(N) {
  const A = new Float32Array(N * N);
  const B = new Float32Array(N * N);
  // ... A, B にランダム値をセット、GPU バッファ作成 ...

  const t0 = performance.now();

  const encoder = device.createCommandEncoder();
  const pass = encoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(Math.ceil(N / 16), Math.ceil(N / 16));
  pass.end();

  // 結果バッファを CPU 側に読み戻す（= GPU の処理完了を確実に待つ）
  encoder.copyBufferToBuffer(cBuffer, 0, readBuffer, 0, byteSize);
  device.queue.submit([encoder.finish()]);
  await readBuffer.mapAsync(GPUMapMode.READ);

  const t1 = performance.now();
  readBuffer.unmap();

  const elapsedMs = t1 - t0;
  const gflops = (2 * N ** 3) / (elapsedMs / 1000) / 1e9;
  return { elapsedMs, gflops };
}

ハイライト部分が 計測精度の核心です。submit() は非同期で、単に呼ぶだけでは GPU はまだ計算中です。mapAsync() でバッファを CPU 側にマップする時点で初めて GPU の処理完了が保証され、正確な実行時間が測れます。