この投稿では、サンプルを使用して試験運用版の WebGPU API を説明し、GPU を使用してデータ並列計算を開始する方法について説明します。
背景
ご存じのとおり、グラフィック プロセッシング ユニット（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 が必要です。2 つ目の GPU バッファは、今回は device.createBuffer() を使用して未マッピングの状態で作成されます。最初の 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 コマンドが実行されたら、gpuReadBuffer.getMappedRange() を使用して、最初の 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 でのコンピューティング シェーダーの使用を説明するために、機械学習で一般的なアルゴリズムである行列乗算を以下に示します。
具体的には、次の手順で対応します。
- 3 つの GPU バッファを作成します（乗算する行列用に 2 つ、結果行列用に 1 つ）。
- コンピューティング シェーダーの入力と出力を記述する
- コンピューティング シェーダー コードをコンパイルする
- コンピューティング パイプラインを設定する
- エンコードされたコマンドを GPU にバッチで送信する
- 結果マトリックスの 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) で設定し、インデックス 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」は 1 つ目のディメンション、2 つ目の引数「y」は 2 つ目のディメンション、最後の引数「z」は 3 つ目のディメンションです。ここでは 3 つ目のディメンションを使用しないため、デフォルトで 1 になります。GPU コンピューティングの世界では、一連のデータに対してカーネル関数を実行するコマンドをエンコードすることをディスパッチと呼びます。
コンピューティング シェーダーのワークグループ グリッドのサイズは、WGSL コードでは (8, 8) です。そのため、1 つ目の行列の行数と 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.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 でのレンダリング（キャンバス、テクスチャ、サンプラー）の仕組みについて詳しく説明する記事がさらに追加される予定です。