GPU 计算使用入门 (Web)

这篇博文通过示例探讨了实验性 WebGPU API,并可帮助您开始使用 GPU 执行数据并行计算。

François Beaufort
François Beaufort

背景

您可能已经了解,图形处理器 (GPU) 是计算机内的一个电子子系统,最初专门用于处理图形。然而,在过去的 10 年里,它已经发展为一种更灵活的架构,允许开发者实现多种类型的算法,而不仅仅是渲染 3D 图形,同时利用 GPU 的独特架构。这些功能称为 GPU 计算,使用 GPU 作为协处理器进行通用科学计算称为通用 GPU (GPGPU) 编程。

GPU 计算为最近的机器学习热潮做出了巨大贡献,因为卷积神经网络和其他模型可以利用该架构在 GPU 上更高效地运行。由于当前 Web 平台缺乏 GPU 计算功能,因此 W3C 的“Web GPU”社区小组正在设计一个 API,用于公开在大多数当前设备上可用的新式 GPU API。该 API 称为 WebGPU

WebGPU 是类似于 WebGL 的低层级 API。正如你看到的那样,它功能强大且十分冗长不过没关系。我们看重的是性能。

在本文中,我将重点介绍 WebGPU 的 GPU 计算部分,老实说,我只是触及了一些皮毛,让您可以自行开始游戏。我会在后续文章中更深入地介绍 WebGPU 渲染(画布、纹理等)。

访问 GPU

在 WebGPU 中访问 GPU 非常简单。调用 navigator.gpu.requestAdapter() 会返回一个 JavaScript promise,该 promise 将使用 GPU 适配器异步解析。可以将此适配器视为显卡。它可以是集成卡(与 CPU 位于同一芯片上),也可以是离散卡(通常是性能更高但耗电量更高的 PCIe 卡)。

有了 GPU 适配器后,请调用 adapter.requestDevice() 以获取将用于执行一些 GPU 计算的 GPU 设备解析的 promise。

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

这两个函数都采用一些选项,让您能够指定所需的适配器类型(电源偏好设置)和设备(扩展程序、限制)。为简单起见,在本文中,我们将使用默认选项。

写入缓冲区内存

我们来看看如何使用 JavaScript 将数据写入 GPU 内存。这个过程并不简单,因为现代网络浏览器中使用了沙盒模型。

以下示例展示了如何将 4 个字节写入可从 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。这次使用 device.createBuffer() 以未映射状态创建第二个 GPU 缓冲区。其用法标志为 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 设备命令队列。队列负责处理通过 device.queue.submit() 以 GPU 命令作为参数完成的提交。这将以原子方式按顺序执行存储在数组中的所有命令。

// 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()。它会返回一个在映射 GPU 缓冲区时可以解析的 promise。然后,使用 gpuReadBuffer.getMappedRange() 获取映射范围,在执行完所有已加入队列的 GPU 命令后,该范围包含与第一个 GPU 缓冲区相同的值。

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

您可以试用此示例

简而言之,对于缓冲区内存操作,您需要注意以下几点:

  • GPU 缓冲区必须未映射才能在设备队列提交中使用。
  • 进行映射后,可以使用 JavaScript 读取和写入 GPU 缓冲区。
  • 当调用 mappedAtCreation 设置为 true 的 mapAsync()createBuffer() 时,系统会映射 GPU 缓冲区。

着色器编程

在 GPU 上运行的仅执行计算(不绘制三角形)的程序称为计算着色器。这些任务由数百个 GPU 核心(小于 CPU 核心)并行执行,这些核心一起运行以处理数据。它们的输入和输出是 WebGPU 中的缓冲区。

为了说明在 WebGPU 中使用计算着色器的情况,我们将使用矩阵乘法,这是一种机器学习中的常见算法,如下图所示。

矩阵乘法图
矩阵乘法图

简而言之,我们将执行以下操作:

  1. 创建三个 GPU 缓冲区(两个用于相乘矩阵,一个用于结果矩阵)
  2. 描述计算着色器的输入和输出
  3. 编译计算着色器代码
  4. 设置计算流水线
  5. 将编码命令批量提交至 GPU
  6. 读取结果矩阵 GPU 缓冲区

GPU 缓冲区创建

为简单起见,矩阵将表示为一个浮点数列表。第一个元素是行数,第二个元素是列数,其余元素是矩阵的实际数。

以 JavaScript 表示的矩阵及其等效数学符号的简单表示法
以 JavaScript 表示的矩阵及其等效的数学符号

三个 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。绑定组布局定义着色器预期的输入/输出接口,而绑定组表示着色器的实际输入/输出数据。

在以下示例中,绑定组布局需要在编号条目绑定 01 处有两个只读存储缓冲区,并在 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
      }
    }
  ]
});

计算着色器代码

用于矩阵乘法的计算着色器代码是用 WebGPU 着色器语言 WGSL 编写的,可以轻松转换为 SPIR-V。无需详细介绍,下面列出了三个用 var<storage> 标识的存储缓冲区。该程序将使用 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() 创建的实际计算着色器模块的计算阶段。

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

提交命令

使用三个 GPU 缓冲区和具有绑定组布局的计算流水线实例化绑定组后,就可以开始使用它们了。

我们使用 commandEncoder.beginComputePass() 启动可编程计算传递编码器。我们将使用它对将执行矩阵乘法的 GPU 命令进行编码。使用 passEncoder.setPipeline(computePipeline) 设置其流水线,使用 passEncoder.setBindGroup(0, bindGroup) 设置索引 0 处的绑定组。索引 0 对应于 WGSL 代码中的 group(0) 装饰。

现在,我们来讨论一下这个计算着色器如何在 GPU 上运行。我们的目标是逐步为结果矩阵的每个单元并行执行此程序。例如,对于大小为 16x32 的结果矩阵,如需对执行命令进行编码,我们可以对 @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()。然后,使用 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() 并等待返回的 promise 进行解析(指示 GPU 缓冲区现已映射)一样简单。此时,可以使用 gpuReadBuffer.getMappedRange() 获取映射范围。

矩阵乘法结果
矩阵乘法结果

在我们的代码中,开发者工具 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: [

性能发现结果

那么,与在 CPU 上运行矩阵乘法相比,在 GPU 上运行矩阵乘法有什么区别?为了一探究竟,我编写了刚才针对 CPU 进行描述的程序。如下图所示,当矩阵大小大于 256x256 时,充分利用 GPU 的能力似乎是一个显而易见的选择。

GPU 基准测试与 CPU 基准测试
GPU 与 CPU 基准比较

这篇文章只是我探索 WebGPU 的起点。我们很快会发布更多文章,深入探讨 GPU 计算以及 WebGPU 中的渲染(画布、纹理、采样器)的工作原理。