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

この投稿では、試験運用版の WebGPU API について例を挙げて説明し、GPU を使用してデータの並列計算を行う方法について説明します。

François Beaufort
François Beaufort

背景

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

畳み込みニューラル ネットワークやその他のモデルはこのアーキテクチャを利用して GPU 上でより効率的に実行できるため、GPU コンピューティングは最近の機械学習ブームに大きく貢献しています。現在のウェブ プラットフォームには 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 デバイスで解決される Promise を取得します。

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 バッファがマッピングされます。つまり、GPU バッファは CPU が所有し、JavaScript から読み取り/書き込みでアクセスできます。GPU がアクセスできるようにするには、マッピングを解除する必要があります。これは、gpuBuffer.unmap() を呼び出すだけです。

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

バッファメモリを読み取る

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

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

// 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 コマンドが実行されたら、gpuReadBuffer.getMappedRange() を使用してマッピングされた範囲を取得します。これには、キューに入れられたすべての GPU コマンドが実行された後に、最初の GPU バッファと同じ値が含まれます。

// 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 上で実行される、計算のみを行う(三角形を描画しない)プログラムは、コンピューティング シェーダーと呼ばれます。これらは、データを処理するために連携して動作する何百もの GPU コア(CPU コアよりも小さい)によって並列に実行されます。入力と出力は WebGPU のバッファです。

WebGPU でのコンピューティング シェーダーの使い方を説明するために、以下に示す ML で一般的なアルゴリズムである行列乗算を取り上げます。

行列の乗算の図
行列の乗算の図

大まかな流れは次のとおりです。

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

GPU バッファの作成

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

JavaScript とそれに相当する数学的表記法による行列の単純な表現
JavaScript での行列を簡単に表現し、それと等価な数学的表記法で表す

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

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 に固有のものです。バインド グループ レイアウトは、シェーダーに必要な入出力インターフェースを定義します。一方、バインド グループは、シェーダーの実際の入出力データを表します。

以下の例では、バインド グループ レイアウトは、番号付きのエントリ バインディング 01 に読み取り専用ストレージ バッファを 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 つのストレージ バッファを示します。プログラムは、入力として firstMatrixsecondMatrix を使用し、出力として 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() を呼び出して作成します。この関数は、前に作成したバインド グループ レイアウトと、コンピューティング シェーダーのエントリ ポイント(main WGSL 関数)を定義するコンピューティング ステージと、device.createShaderModule() で作成された実際のコンピューティング シェーダー モジュールの 2 つの引数を取ります。

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

コマンドの送信

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

commandEncoder.beginComputePass() を使用して、プログラム可能なコンピューティング パス エンコーダを起動しましょう。これを使用して行列の乗算を行う GPU コマンドをエンコードしますパイプラインを passEncoder.setPipeline(computePipeline) で、バインド グループをインデックス 0 で passEncoder.setBindGroup(0, bindGroup) で設定します。インデックス 0 は WGSL コードの group(0) 装飾に対応します。

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

結果行列セルごとに並列実行
結果行列のセルごとに並列実行

コンピューティング シェーダーのワークグループ グリッドのサイズは、WGSL コードでは (8, 8) です。そのため、最初のマトリックスの行数と 2 番目のマトリックスの列数である「x」と「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.READgpuReadBuffer.mapAsync() を呼び出し、GPU バッファがマッピングされたことを示す返される Promise が解決されるのを待つのと同じくらい簡単です。この時点で、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));

これで完了です。お疲れさまでした。サンプルでプレイできます。

最後のトリック

コードを読みやすくする方法の 1 つは、コンピューティング パイプラインの便利な 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 をフル活用することは明らかな選択です。

GPU と CPU のベンチマーク
GPU と CPU のベンチマーク

この投稿は、WebGPU について検討するための旅の始まりにすぎません。GPU コンピューティングの詳細と、WebGPU でのレンダリング(キャンバス、テクスチャ、サンプラー)の仕組みについては、今後さらに記事をお届けする予定です。