本文通过示例探索了实验性 WebGPU API，并帮助您开始使用 GPU 执行数据并行计算。
背景
您可能已经知道，图形处理器 (GPU) 是计算机中的一种电子子系统，最初专门用于处理图形。然而，在过去的 10 年中，它已发展为更灵活的架构，允许开发者实现多种类型的算法，而不仅仅是渲染 3D 图形，同时利用 GPU 的独特架构。这些功能称为 GPU 计算，将 GPU 用作协处理器进行通用科学计算称为通用 GPU (GPGPU) 编程。
GPU 计算对近期机器学习热潮的兴起做出了重要贡献，因为卷积神经网络和其他模型可以利用该架构在 GPU 上更高效地运行。由于当前 Web 平台缺乏 GPU 计算功能，W3C 的“Web GPU”社区小组正在设计一个 API，用于公开大多数当前设备上可用的现代 GPU API。此 API 称为 WebGPU。
WebGPU 是一种低级 API，类似于 WebGL。它非常强大且非常详尽，您很快就会发现这一点。不过没关系。我们注重的是效果。
在本文中，我将重点介绍 WebGPU 的 GPU 计算部分。说实话，我只是触及了浅层，因此您可以开始自行体验游戏。在后续文章中，我将深入介绍 WebGPU 渲染（画布、纹理等）。
访问 GPU
在 WebGPU 中，访问 GPU 非常简单。调用 navigator.gpu.requestAdapter() 会返回一个 JavaScript promise，该 promise 将异步解析为 GPU 适配器。请将此适配器视为显卡。它可以集成（与 CPU 在同一芯片上），也可以离散（通常是性能更高但耗电量更高的 PCIe 卡）。
有了 GPU 适配器后，请调用 adapter.requestDevice() 来获取可通过 GPU 设备进行解析的 promise，您将使用该 GPU 设备进行一些 GPU 计算。
const adapter = await navigator.gpu.requestAdapter();
if (!adapter) { return; }
const device = await adapter.requestDevice();
这两个函数都提供选项，让您可以指定所需的适配器类型（电源偏好设置）和设备（扩展程序、限制）。为简单起见，我们将使用本文中的默认选项。
写入缓冲区内存
我们来看看如何使用 JavaScript 将数据写入 GPU 的内存。由于现代网络浏览器中使用的沙盒模型，此过程并不简单。
以下示例展示了如何将四个字节写入可从 GPU 访问的缓冲区内存。它会调用 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 缓冲区写入数据，并希望将其复制到第二个 GPU 缓冲区，因此需要使用新的使用标志 GPUBufferUsage.COPY_SRC。这次，第二个 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 命令编码器是一个 JavaScript 对象，用于构建一批将在某个时间点发送到 GPU 的“缓冲”命令。另一方面，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 队列命令已发送，但未必会执行。如需读取第二个 GPU 缓冲区，请使用 GPUMapMode.READ 调用 gpuReadBuffer.mapAsync()。它会返回一个 promise，该 promise 会在 GPU 缓冲区映射时解析。然后，使用 gpuReadBuffer.getMappedRange() 获取映射范围，该范围包含与第一个 GPU 缓冲区相同的值，并且在所有队列的 GPU 命令都已执行后才会执行。
// Read buffer.
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
const copyArrayBuffer = gpuReadBuffer.getMappedRange();
console.log(new Uint8Array(copyArrayBuffer));
您可以试用此示例。
简而言之，关于缓冲区内存操作，您需要记住以下几点：
- GPU 缓冲区必须取消映射，才能在设备队列提交中使用。
- 映射后，您可以使用 JavaScript 读取和写入 GPU 缓冲区。
- 在调用
mapAsync()和createBuffer()（并将mappedAtCreation设置为 true）时，系统会映射 GPU 缓冲区。
着色器编程
在 GPU 上运行且仅执行计算（而不绘制三角形）的程序称为计算着色器。它们由数百个 GPU 核心（比 CPU 核心小）并行执行，这些核心协同工作来处理数据。它们的输入和输出是 WebGPU 中的缓冲区。
为了说明如何在 WebGPU 中使用计算着色器，我们将使用矩阵乘法（机器学习中常见的算法，如下所示）进行演示。
简而言之，我们将执行以下操作：
- 创建三个 GPU 缓冲区（两个用于矩阵相乘，一个用于结果矩阵）
- 描述计算着色器的输入和输出
- 编译计算着色器代码
- 设置计算流水线
- 向 GPU 批量提交已编码的命令
- 读取结果矩阵 GPU 缓冲区
GPU 缓冲区创建
为简单起见，矩阵将表示为浮点数列表。第一个元素是行数，第二个元素是列数，其余元素是矩阵的实际数。
由于我们需要在计算着色器中存储和检索数据，因此这三个 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 处有一个计算着色器的存储缓冲区。另一方面，为此绑定组布局定义的绑定组会将 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> 标识的三个存储缓冲区。该程序将使用 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() 创建它。它需要两个参数：我们之前创建的绑定组布局，以及定义计算着色器（main WGSL 函数）入口点的计算阶段和使用 device.createShaderModule() 创建的实际计算着色器模块。
const computePipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [bindGroupLayout]
}),
compute: {
module: shaderModule,
entryPoint: "main"
}
});
提交命令
使用三个 GPU 缓冲区和具有 bind 组布局的计算流水线实例化 bind 组后，接下来就可以使用它们了。
我们将使用 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”是第一个维度，第二个参数“y”是第二个维度，最后一个参数“z”是第三个维度，默认值为 1，因为我们在这里不需要它。在 GPU 计算领域，对一组数据执行内核函数的命令称为调度。
在 WGSL 代码中，计算着色器的工作组网格大小为 (8, 8)。因此，“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()。然后，创建一个 GPU 缓冲区，以用作使用 copyBufferToBuffer 复制结果矩阵缓冲区的目标。最后，使用 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()，然后等待返回的 promise 解析，这表示 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 中工作的文章，敬请期待。