WebGPU のコンピュートシェーダで 10 万粒子を動かす

粒子を大量に動かすとき、CPU で座標を計算して毎フレーム GPU に送るやり方はすぐ頭打ちになります。 10 万個ともなると、転送そのものがボトルネックになるからです。

コンピュートシェーダを使えば、座標を GPU 上に置いたまま更新できます。CPU からは 「更新しろ」という指示だけを送ればよく、毎フレームの転送がまるごと消えます。

ストレージバッファのレイアウト

まず粒子の状態を置くバッファを用意します。位置と速度をひとつの構造体にまとめました。

struct Particle {
  pos: vec2f,
  vel: vec2f,
}

@group(0) @binding(0) var<storage, read_write> particles: array<Particle>;

ここで気をつけるのが アラインメント です。vec2f は 8 バイト境界に揃えられるので、 この構造体はちょうど 16 バイトになります。もし vec3f を混ぜると 16 バイト境界に 切り上げられ、意図しない隙間が空きます。JS 側で Float32Array に詰めるときのオフセットが ずれる原因はたいていこれです。

構造体のサイズは、必ず WGSL 側の想定と JS 側の書き込みで突き合わせてください。 ずれていても実行時エラーは出ず、粒子が静かに散らかるだけなので気づきにくいです。

ワークグループの大きさ

次に更新カーネルです。

@compute @workgroup_size(64)
fn update(@builtin(global_invocation_id) gid: vec3u) {
  let i = gid.x;
  if (i >= arrayLength(&particles)) { return; }

  var p = particles[i];
  p.vel += vec2f(0.0, -0.0004);
  p.pos += p.vel;
  if (p.pos.y < -1.0) { p.pos.y = 1.0; }
  particles[i] = p;
}

workgroup_size は 64 にしています。この値は「GPU が一度に走らせる最小単位の倍数」に しておくのが基本で、多くの環境で 32 か 64 が実行単位です。16 のような半端な値にすると 実行ユニットが遊びます。

冒頭の範囲チェックは省略できません。ディスパッチ数は切り上げで計算するため、 粒子数が 64 の倍数でないとき、末尾のスレッドが配列の外を触りにいきます。

const groups = Math.ceil(count / 64);
pass.dispatchWorkgroups(groups);

描画とのつなぎ方

更新したバッファは、そのまま頂点シェーダから読めます。バッファを CPU に戻す必要はありません。

同じバッファを「コンピュートパスで書き込み、レンダーパスで読み込み」という順に使うので、 バッファの usage には STORAGE と VERTEX の両方を立てておきます。ここを忘れると パイプライン作成時に弾かれます。

実測

手元の環境では、10 万粒子で 60fps を維持できました。同じ処理を CPU 側の requestAnimationFrame で回した場合は 1 万粒子あたりでフレーム落ちが始まったので、 転送を消した効果は素直に出ています。

デモページで実際のパラメータを動かせるようにしてあります。粒子数を上げていくと、 どのあたりから崩れるかが環境ごとに違うのが見えて面白いはずです。