ウェブで GPU コンピューティングを使ってみる

この投稿では、例を通して試験運用版の WebGPU API について説明します。また、GPU を使用してデータの並列計算を開始する際に役立ちます。

François Beaufort

背景

ご存じのとおり、グラフィック プロセッシング ユニット（GPU）は、元々グラフィック処理に特化したコンピュータ内の電子サブシステムです。しかし、過去 10 年間で、GPU はより柔軟なアーキテクチャへと進化しました。これにより、デベロッパーは GPU の独自のアーキテクチャを活用しながら、3D グラフィックのレンダリングだけでなく、さまざまなタイプのアルゴリズムを実装できるようになりました。これらの機能は GPU コンピューティングと呼ばれ、汎用科学コンピューティングのコプロセッサとして GPU を使用することを汎用 GPU（GPGPU）プログラミングと呼びます。

GPU コンピューティングは、畳み込みニューラル ネットワークなどのモデルがアーキテクチャを活用して GPU でより効率的に実行できるため、最近の ML ブームに大きく貢献しています。現在のウェブ プラットフォームには GPU コンピューティング機能がないため、W3C の「GPU for the Web」コミュニティ グループは、現在のほとんどのデバイスで利用可能な最新の GPU API を公開する API を設計しています。この API は WebGPU と呼ばれます。

WebGPU は、WebGL のような低レベルの API です。非常に強力で、冗長なコードになります。パフォーマンスが重要です。

この記事では、WebGPU の GPU コンピューティング部分に焦点を当てます。正直なところ、自分で試していただけるよう、ほんの一部しか説明しません。WebGPU のレンダリング（キャンバス、テクスチャなど）については、今後の記事で詳しく説明します。

GPU にアクセスする

WebGPU では GPU へのアクセスが簡単です。navigator.gpu.requestAdapter() を呼び出すと、GPU アダプタで非同期的に解決される JavaScript Promise が返されます。このアダプターはグラフィック カードのようなものです。統合型（CPU と同じチップ上）とディスクリート型（通常は、パフォーマンスは高いが消費電力が多い PCIe カード）があります。

GPU アダプターを取得したら、adapter.requestDevice() を呼び出して、GPU 計算に使用する GPU デバイスで解決するプロミスを取得します。

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) { return; }
const device = await adapter.requestDevice();

どちらの関数にも、必要なアダプターの種類（電源設定）とデバイス（拡張機能、制限）を指定するためのオプションがあります。わかりやすくするため、この記事ではデフォルトのオプションを使用します。

バッファメモリへの書き込み

JavaScript を使用して GPU のメモリにデータを書き込む方法を見てみましょう。最新のウェブブラウザではサンドボックス モデルが使用されているため、このプロセスは簡単ではありません。

次の例は、GPU からアクセス可能なバッファメモリに 4 バイトを書き込む方法を示しています。また、バッファのサイズと使用量を取得する device.createBuffer() を呼び出します。この特定の呼び出しでは用途フラグ GPUBufferUsage.MAP_WRITE は必要ありませんが、このバッファに書き込むことを明示的に説明します。mappedAtCreation が true に設定されているため、作成時に GPU バッファ オブジェクトがマッピングされます。次に、GPU バッファ メソッド getMappedRange() を呼び出して、関連する未加工バイナリ データ バッファを取得できます。

ArrayBuffer をすでに使用している場合は、バイトの書き込みはよく知られているはずです。TypedArray を使用して値をコピーします。

// Get a GPU buffer in a mapped state and an arrayBuffer for writing.
const gpuBuffer = device.createBuffer({
  mappedAtCreation: true,
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE
});
const arrayBuffer = gpuBuffer.getMappedRange();

// Write bytes to buffer.
new Uint8Array(arrayBuffer).set([0, 1, 2, 3]);

この時点で、GPU バッファはマッピングされます。つまり、CPU によって所有され、JavaScript から読み取り/書き込みでアクセスできるようになります。GPU がアクセスできるようにするには、マッピングを解除する必要があります。これは gpuBuffer.unmap() を呼び出すだけで済みます。

GPU と CPU が同時にメモリにアクセスする競合状態を防ぐには、マッピング済み/マッピングなしのコンセプトが必要です。

バッファメモリの読み取り

次に、GPU バッファを別の GPU バッファにコピーして、読み戻す方法を見てみましょう。

最初の GPU バッファに書き込み、2 番目の GPU バッファにコピーするため、新しい用途フラグ GPUBufferUsage.COPY_SRC が必要です。今度は device.createBuffer() により、2 つ目の GPU バッファがマッピングされていない状態で作成されます。最初の GPU バッファの宛先として使用され、GPU コピー コマンドが実行されると JavaScript で読み取られるため、使用フラグは GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ です。

// Get a GPU buffer in a mapped state and an arrayBuffer for writing.
const gpuWriteBuffer = device.createBuffer({
  mappedAtCreation: true,
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC
});
const arrayBuffer = gpuWriteBuffer.getMappedRange();

// Write bytes to buffer.
new Uint8Array(arrayBuffer).set([0, 1, 2, 3]);

// Unmap buffer so that it can be used later for copy.
gpuWriteBuffer.unmap();

// Get a GPU buffer for reading in an unmapped state.
const gpuReadBuffer = device.createBuffer({
  size: 4,
  usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
});

GPU は独立したコプロセッサであるため、すべての GPU コマンドは非同期で実行されます。そのため、GPU コマンドのリストが作成され、必要に応じてバッチで送信されます。WebGPU では、device.createCommandEncoder() によって返される GPU コマンド エンコーダは、ある時点で GPU に送信される「バッファリングされた」コマンドのバッチを作成する JavaScript オブジェクトです。一方、GPUBuffer のメソッドは「バッファなし」です。つまり、呼び出された時点でアトミックに実行されます。

GPU コマンド エンコーダを取得したら、次のように copyEncoder.copyBufferToBuffer() を呼び出して、このコマンドをコマンドキューに追加し、後で実行できるようにします。最後に、copyEncoder.finish() を呼び出してコマンドのエンコードを終了し、それらを GPU デバイス コマンドキューに送信します。キューは、GPU コマンドを引数として device.queue.submit() を介して送信された送信を処理します。これにより、配列に格納されているすべてのコマンドが順番にアトミックに実行されます。

// Encode commands for copying buffer to buffer.
const copyEncoder = device.createCommandEncoder();
copyEncoder.copyBufferToBuffer(
  gpuWriteBuffer /* source buffer */,
  0 /* source offset */,
  gpuReadBuffer /* destination buffer */,
  0 /* destination offset */,
  4 /* size */
);

// Submit copy commands.
const copyCommands = copyEncoder.finish();
device.queue.submit([copyCommands]);

この時点で、GPU キュー コマンドは送信されていますが、必ずしも実行されているわけではありません。2 番目の GPU バッファを読み取るには、GPUMapMode.READ で gpuReadBuffer.mapAsync() を呼び出します。GPU バッファがマッピングされたときに解決される Promise を返します。次に、キューに追加されたすべての GPU コマンドが実行された後、最初の GPU バッファと同じ値を含む gpuReadBuffer.getMappedRange() でマッピングされた範囲を取得します。

// Read buffer.
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
const copyArrayBuffer = gpuReadBuffer.getMappedRange();
console.log(new Uint8Array(copyArrayBuffer));

こちらのサンプルをお試しください。

要約すると、バッファメモリ オペレーションについて覚えておくべきことは次のとおりです。

• デバイス キューの送信で GPU バッファを使用するには、マッピングを解除する必要があります。
• マッピングすると、GPU バッファを JavaScript で読み書きできるようになります。
• GPU バッファは、mappedAtCreation が true に設定された mapAsync() と createBuffer() が呼び出されるとマッピングされます。

シェーダー プログラミング

GPU で実行され、計算のみを行い（三角形を描画しない）プログラムはコンピューティング シェーダーと呼ばれます。これらは、連携してデータを処理するために（CPU コアよりも小さい）数百の GPU コアによって並列に実行されます。入力と出力は WebGPU のバッファです。

WebGPU でのコンピューティング シェーダーの使用を説明するために、機械学習で一般的なアルゴリズムである行列乗算を以下に示します。

具体的には、次の手順で対応します。

1. 3 つの GPU バッファを作成します（乗算する行列用に 2 つ、結果行列用に 1 つ）。
2. コンピューティング シェーダーの入力と出力を記述する
3. コンピューティング シェーダー コードをコンパイルする
4. コンピューティング パイプラインを設定する
5. エンコードされたコマンドを GPU にバッチで送信する
6. 結果マトリックスの GPU バッファを読み取る

GPU バッファの作成

簡潔にするために、行列は浮動小数点数のリストとして表されます。最初の要素は行数、2 番目の要素は列数、残りは行列の実際の数です。

3 つの GPU バッファはストレージ バッファです。コンピューティング シェーダーでデータを保存して取得する必要があるためです。これが、GPU バッファ使用フラグにすべての GPUBufferUsage.STORAGE が含まれている理由です。結果マトリックスの使用フラグにも GPUBufferUsage.COPY_SRC があります。これは、すべての GPU キュー コマンドが実行されると、読み取りのために別のバッファにコピーされるためです。

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) { return; }
const device = await adapter.requestDevice();


// First Matrix

const firstMatrix = new Float32Array([
  2 /* rows */, 4 /* columns */,
  1, 2, 3, 4,
  5, 6, 7, 8
]);

const gpuBufferFirstMatrix = device.createBuffer({
  mappedAtCreation: true,
  size: firstMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
const arrayBufferFirstMatrix = gpuBufferFirstMatrix.getMappedRange();
new Float32Array(arrayBufferFirstMatrix).set(firstMatrix);
gpuBufferFirstMatrix.unmap();


// Second Matrix

const secondMatrix = new Float32Array([
  4 /* rows */, 2 /* columns */,
  1, 2,
  3, 4,
  5, 6,
  7, 8
]);

const gpuBufferSecondMatrix = device.createBuffer({
  mappedAtCreation: true,
  size: secondMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
const arrayBufferSecondMatrix = gpuBufferSecondMatrix.getMappedRange();
new Float32Array(arrayBufferSecondMatrix).set(secondMatrix);
gpuBufferSecondMatrix.unmap();


// Result Matrix

const resultMatrixBufferSize = Float32Array.BYTES_PER_ELEMENT * (2 + firstMatrix[0] * secondMatrix[1]);
const resultMatrixBuffer = device.createBuffer({
  size: resultMatrixBufferSize,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
});

バインド グループのレイアウトとバインド グループ

バインディング グループのレイアウトとバインディング グループのコンセプトは、WebGPU に固有のものです。バインディング グループ レイアウトは、シェーダーが想定する入出力インターフェースを定義します。バインディング グループは、シェーダーの実際の入出力データを表します。

以下の例のバインド グループ レイアウトは、番号付きのエントリ バインディング 0、1 に 2 つの読み取り専用ストレージ バッファと、コンピューティング シェーダーの 2 のストレージ バッファを想定しています。一方、このバインディング グループ レイアウト用に定義されたバインディング グループは、GPU バッファをエントリに関連付けます。gpuBufferFirstMatrix はバインディング 0 に、gpuBufferSecondMatrix はバインディング 1 に、resultMatrixBuffer はバインディング 2 に関連付けます。

const bindGroupLayout = device.createBindGroupLayout({
  entries: [
    {
      binding: 0,
      visibility: GPUShaderStage.COMPUTE,
      buffer: {
        type: "read-only-storage"
      }
    },
    {
      binding: 1,
      visibility: GPUShaderStage.COMPUTE,
      buffer: {
        type: "read-only-storage"
      }
    },
    {
      binding: 2,
      visibility: GPUShaderStage.COMPUTE,
      buffer: {
        type: "storage"
      }
    }
  ]
});

const bindGroup = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [
    {
      binding: 0,
      resource: {
        buffer: gpuBufferFirstMatrix
      }
    },
    {
      binding: 1,
      resource: {
        buffer: gpuBufferSecondMatrix
      }
    },
    {
      binding: 2,
      resource: {
        buffer: resultMatrixBuffer
      }
    }
  ]
});

コンピューティング シェーダー コード

行列の乗算のコンピューティング シェーダー コードは、WGSL（WebGPU シェーダー言語）で記述されています。これは簡単に SPIR-V に変換できます。詳細は省略しますが、var<storage> で識別される 3 つのストレージ バッファが以下に示されています。このプログラムは、firstMatrix と secondMatrix を入力として、resultMatrix を出力として使用します。

各ストレージ バッファには、上記で宣言されたバインディング グループ レイアウトとバインディング グループで定義された同じインデックスに対応する binding デコレーションが使用されます。

const shaderModule = device.createShaderModule({
  code: `
    struct Matrix {
      size : vec2f,
      numbers: array<f32>,
    }

    @group(0) @binding(0) var<storage, read> firstMatrix : Matrix;
    @group(0) @binding(1) var<storage, read> secondMatrix : Matrix;
    @group(0) @binding(2) var<storage, read_write> resultMatrix : Matrix;

    @compute @workgroup_size(8, 8)
    fn main(@builtin(global_invocation_id) global_id : vec3u) {
      // Guard against out-of-bounds work group sizes
      if (global_id.x >= u32(firstMatrix.size.x) || global_id.y >= u32(secondMatrix.size.y)) {
        return;
      }

      resultMatrix.size = vec2(firstMatrix.size.x, secondMatrix.size.y);

      let resultCell = vec2(global_id.x, global_id.y);
      var result = 0.0;
      for (var i = 0u; i < u32(firstMatrix.size.y); i = i + 1u) {
        let a = i + resultCell.x * u32(firstMatrix.size.y);
        let b = resultCell.y + i * u32(secondMatrix.size.y);
        result = result + firstMatrix.numbers[a] * secondMatrix.numbers[b];
      }

      let index = resultCell.y + resultCell.x * u32(secondMatrix.size.y);
      resultMatrix.numbers[index] = result;
    }
  `
});

パイプラインの設定

コンピューティング パイプラインは、実行するコンピューティング オペレーションを実際に記述するオブジェクトです。device.createComputePipeline() を呼び出して作成します。2 つの引数（前述のバインディング グループ レイアウトと、コンピューティング シェーダーのエント ryPoint（main WGSL 関数）と device.createShaderModule() で作成された実際のコンピューティング シェーダー モジュールを定義するコンピューティング ステージ）を受け取ります。

const computePipeline = device.createComputePipeline({
  layout: device.createPipelineLayout({
    bindGroupLayouts: [bindGroupLayout]
  }),
  compute: {
    module: shaderModule,
    entryPoint: "main"
  }
});

コマンドの送信

3 つの GPU バッファとバインド グループ レイアウトのコンピューティング パイプラインを使用してバインド グループをインスタンス化したら、それらを使用します。

commandEncoder.beginComputePass() を使用して、プログラム可能なコンピューティング パス エンコーダを開始しましょう。行列乗算を行う GPU コマンドを エンコードするために使用しますパイプラインを passEncoder.setPipeline(computePipeline) で設定し、バインド グループを passEncoder.setBindGroup(0, bindGroup) でインデックス 0 に置きます。インデックス 0 は、WGSL コードの group(0) デコレーションに対応しています。

次に、このコンピューティング シェーダーが GPU で実行される仕組みについて説明します。目標は、結果行列の各セルに対して、このプログラムをステップごとに並列に実行することです。たとえば、サイズが 16 x 32 の結果行列の場合、@workgroup_size(8, 8) で実行コマンドをエンコードするには、passEncoder.dispatchWorkgroups(2, 4) または passEncoder.dispatchWorkgroups(16 / 8, 32 / 8) を呼び出します。最初の引数「x」は 1 つ目のディメンション、2 つ目の引数「y」は 2 つ目のディメンション、最後の引数「z」は 3 つ目のディメンションです。ここでは 3 つ目のディメンションを使用しないため、デフォルトで 1 になります。GPU コンピューティングの世界では、一連のデータに対してカーネル関数を実行するコマンドをエンコードすることをディスパッチと呼びます。

コンピューティング シェーダーのワークグループ グリッドのサイズは、WGSL コードでは (8, 8) です。そのため、1 つ目の行列の行数である「x」と 2 つ目の行列の列数である「y」は、それぞれ 8 で除算されます。これで、passEncoder.dispatchWorkgroups(firstMatrix[0] / 8, secondMatrix[1] / 8) でコンピューティング呼び出しをディスパッチできるようになりました。実行するワークグループ グリッドの数は dispatchWorkgroups() 引数です。

上の図に示すように、各シェーダーは、計算すべき結果行列セルの判別に使用される一意の builtin(global_invocation_id) オブジェクトにアクセスできます。

const commandEncoder = device.createCommandEncoder();

const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);
const workgroupCountX = Math.ceil(firstMatrix[0] / 8);
const workgroupCountY = Math.ceil(secondMatrix[1] / 8);
passEncoder.dispatchWorkgroups(workgroupCountX, workgroupCountY);
passEncoder.end();

コンピューティング パス エンコーダを終了するには、passEncoder.end() を呼び出します。次に、copyBufferToBuffer を使用して、結果マトリックス バッファのコピー先として使用する GPU バッファを作成します。最後に、copyEncoder.finish() でエンコード コマンドを終了し、GPU コマンドで device.queue.submit() を呼び出して、GPU デバイスキューに送信します。

// Get a GPU buffer for reading in an unmapped state.
const gpuReadBuffer = device.createBuffer({
  size: resultMatrixBufferSize,
  usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
});

// Encode commands for copying buffer to buffer.
commandEncoder.copyBufferToBuffer(
  resultMatrixBuffer /* source buffer */,
  0 /* source offset */,
  gpuReadBuffer /* destination buffer */,
  0 /* destination offset */,
  resultMatrixBufferSize /* size */
);

// Submit GPU commands.
const gpuCommands = commandEncoder.finish();
device.queue.submit([gpuCommands]);

結果マトリックスの読み取り

結果マトリックスの読み取りは、GPUMapMode.READ で gpuReadBuffer.mapAsync() を呼び出し、GPU バッファがマッピングされたことを示す戻りプロミスの解決を待つだけです。この時点では、gpuReadBuffer.getMappedRange() を使用してマッピングされた範囲を取得できます。

このコードでは、DevTools JavaScript コンソールに記録される結果は「2, 2, 50, 60, 114, 140」です。

// Read buffer.
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
const arrayBuffer = gpuReadBuffer.getMappedRange();
console.log(new Float32Array(arrayBuffer));

これで、これで完了です。サンプルを試すことができます。

最後のトリック

コードを読みやすくする方法の一つは、コンピューティング パイプラインの便利な getBindGroupLayout メソッドを使用して、シェーダー モジュールからバインディング グループ レイアウトを推論することです。この方法では、カスタム バインド グループ レイアウトを作成したり、コンピューティング パイプラインでパイプライン レイアウトを指定したりする必要がなくなります。

上記のサンプルの getBindGroupLayout の図はこちらで確認できます。

 const computePipeline = device.createComputePipeline({
-  layout: device.createPipelineLayout({
-    bindGroupLayouts: [bindGroupLayout]
-  }),
   compute: {

-// Bind group layout and bind group
- const bindGroupLayout = device.createBindGroupLayout({
-   entries: [
-     {
-       binding: 0,
-       visibility: GPUShaderStage.COMPUTE,
-       buffer: {
-         type: "read-only-storage"
-       }
-     },
-     {
-       binding: 1,
-       visibility: GPUShaderStage.COMPUTE,
-       buffer: {
-         type: "read-only-storage"
-       }
-     },
-     {
-       binding: 2,
-       visibility: GPUShaderStage.COMPUTE,
-       buffer: {
-         type: "storage"
-       }
-     }
-   ]
- });
+// Bind group
  const bindGroup = device.createBindGroup({
-  layout: bindGroupLayout,
+  layout: computePipeline.getBindGroupLayout(0 /* index */),
   entries: [

パフォーマンスに関する検出結果

GPU で行列乗算を実行する場合と CPU で実行する場合の違いは何でしょうか。それを調べるために、先ほど説明した CPU 用のプログラムを作成しました。下のグラフに示すように、行列のサイズが 256 x 256 より大きい場合は、GPU の全機能を使用するのが当然の選択肢です。

この記事は、WebGPU の探索の始まりにすぎません。今後、GPU コンピューティングと WebGPU でのレンダリング（キャンバス、テクスチャ、サンプラー）の仕組みについて詳しく説明する記事が公開される予定です。