WebGPU:开始在浏览器上使用 GPU 计算

1,398 阅读9分钟

2019年8月28日 发表 — 2022年5月1日 更新

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

背景#

大家熟知的图形处理单元(GPU) 是计算机中电子元件的一部分,最初是专门用于处理图形的。然而,在过去的10年中,GPU不仅仅是被用在渲染3D图形,它已经朝着更灵活的架构发展,利用GPU这种灵活而独特架构,开发人员能够实现各种各样的算法。这些功能被称为GPU计算(GPU Compute),使用GPU作为通用科学计算的协处理器被称为GPGPU编程(general-purpose GPU)。

因为卷积神经网络和其他模型可以利用GPU独特的架构在高效地计算运行,GPU Compute为最近的机器学习热潮做出了重大贡献。由于当前的Web平台缺乏GPU计算功能,W3C的“GPU for the Web”社区组设计了一个 Web API,以便能够运行在当前大多数GPU设备上并暴露现代GPU的原生APIs。这个 API 就是 WebGPU

WebGPU是一个低级API,就像WebGL一样。它非常强大且非常冗长。但没关系不大,我们想要的、最求的是性能。

在本文中,我将重点介绍WebGPU的GPU计算部分,说实话,我只是触及了表面,以便您可以自己开始玩。我将在即将发表的文章中更深入地介绍WebGPU渲染(画布,纹理等)。

WebGPU目前还未正式发布,你可以在Chrome Canary桌面端的实验性功能里使用。在 chrome://flags/#enable-unsafe-webgpu上启用WebGPU。WebGPU 的 API 在不断变化,目前不安全。由于尚未为WebGPU API实现沙盒,因此可以为其他进程读取GPU数据!不要在启用它的情况下浏览网页。

访问显卡#

在WebGPU中访问GPU很容易。调用 navigator.gpu.requestAdapter() 返回一个 JavaScript Promise ,该Promise 会异步处理 GPU 适配器 adapter 。可以把此 adapter 理解为显卡硬件。它可以是集成显卡(与CPU在同一芯片上)或独立显卡(通常是性能更高但功耗更高的PCIe卡)。

拥有 GPU 适配器后,调用 adapter.requestDevice() 以获取 GPU 设备Device的Promise ,用于执行某些 GPU 计算。

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

这两个功能都采用选项,允许您具体说明所需的适配器类型(电源首选项)和设备(扩展,限制)。简单起见,本文中使用的都是默认选项。

写Buffer内存#

让我们看看如何使用 JavaScript 将数据写入 GPU 的内存。此过程并不简单,因为现代 Web 浏览器中使用的沙盒模型。

下面的示例演示如何将四个字节写入 GPU 的可访问的缓冲内存。通过调用 device.createBuffer() 创建需要用到的GPU Buffer,包括了缓存区的大小及其用法。虽然usage属性的默认不需要特别指定为,但是我们指定为 GPUBufferUsage.MAP_WRITE 表明是要写入此 Buffer。由于设置了 mappedAtCreation 为 true 导致GPU在创建缓冲对象的时候就进行了映射。然后,通过调用GPU Buffer的 getMappedRange() 方法,来检索相关联的原始二进制数据 Buffer。

如果您已经用过 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 Buffer被映射,CPU开始管理二进制缓存,并且可以从JavaScript 读/写访问它。为了使GPU可以访问它,必须调用 gpuBuffer.unmap() 取消映射。

需要映射/取消映射的来防止GPU和CPU同时访问内存的引起争用情况。

读Buffer存储器#

现在,让我们看看如何将 GPU Buffer复制到另一个 GPU Buffer并将其读回。

由于我们正在写入第一个 GPU Buffer,且我们希望将其复制到第二个 GPU Buffer,因此usage属性需要一个新的标识 GPUBufferUsage.COPY_SRC。第二个 GPU Buffer是在未映射状态下使用创建的,usage属性值是 GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ ,它将作为第一个GPU Buffer的复制目的地,一旦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中,GPU命令编码器是通过device.createCommandEncoder() 创建并返回的一个JavaScript对象,它构建了一批“缓冲”命令,这些命令将在某个时候发送到GPU。另一方面,GPU Buffer的方法是“无缓冲的”,这意味着它们在被调用时以原子方式执行。

获得 GPU 命令编码器后,按下面代码所示,调用 copyEncoder.copyBufferToBuffer(),将此命令添加到命令队列中,以GPU便稍后执行。最后,通过调用 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 Buffer的内容,需要先以GPUMapMode.READ 为参数调用.mapAsync()方法。它返回一个Promise,当 GPU Buffer被映射后,就可以在第二个GPU Buffer 上通过使用 getMappedRange() 获取映射范围,得到执行和第一个 GPU Buffer相同的值。

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

简而言之,以下是您需要记住的有关Buffer内存操作的内容:

  • 必须取消映射 GPU Buffer才能在GPU设备队列提交中使用。
  • 当被映射后,GPU Buffer可以通过 JavaScript 读取和写入。
  • 在调用mapAsync()、createBuffer() 或者 把mappedAtCreation 设置为 true 时 GPU Buffer 进行映射。

着色器编程#

在 GPU 上运行且仅执行计算(不绘制三角形)的程序称为计算着色器(compute shaders)。它们由数百个GPU内核(小于CPU内核)并行运算处理数据。它们的输入和输出是WebGPU Buffers。

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

简而言之,这就是我们要做的事情:

  • 创建三个 GPU 缓冲区(两个用于存放相乘的矩阵,一个用于存放结果矩阵)
  • 描述计算着色器的输入和输出
  • 编译计算着色器代码
  • 设置计算管道
  • 将编码的命令批量提交到 GPU
  • 读取 GPU 缓冲区中的结果矩阵

创建 GPU 缓冲区#

为简单起见,矩阵将表示为浮点数列表。

第一个元素是行数,第二个元素是列数,其余元素是矩阵的实际数。

三个 GPU 缓冲区是存储缓冲区,因为我们需要在计算着色器中存储和检索数据。这就解释了为什么 GPU 缓冲区使用标志包括所有这些缓冲区。结果矩阵使用标志也有,因为一旦所有GPU队列命令全部执行完毕,它将被复制到另一个缓冲区进行读取。

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

// 第一个矩阵
const firstMatrix = new Float32Array([2 /* rows */, 4 /* columns */, 1, 2, 3, 4, 5, 6, 7, 8]);
const gpuBufferFirstMatrix = device.createBuffer({
    // 创建 GPU Buffer
    mappedAtCreation: true, // 自动映射
    size: firstMatrix.byteLength, // 和第一个数组一样的大小,字节长度
    usage: GPUBufferUsage.STORAGE // 标识这个 Buffer 只用作存值
});
// 获取 GPU Buffer 显存范围
const arrayBufferFirstMatrix = gpuBufferFirstMatrix.getMappedRange();
// 写入数据
new Float32Array(arrayBufferFirstMatrix).set(firstMatrix);
// 解除映射,以便 CPU 可以操作 Buffer
gpuBufferFirstMatrix.unmap();

// 第二个矩阵
const secondMatrix = new Float32Array([4 /* rows */, 2 /* columns */, 1, 2, 3, 4, 5, 6, 7, 8]);
const gpuBufferSecondMatrix = device.createBuffer({
    // 创建 GPU Buffer
    mappedAtCreation: true, // 自动映射
    size: secondMatrix.byteLength, // 和第二个数组一样的大小,字节长度
    usage: GPUBufferUsage.STORAGE // 标识这个 Buffer 只用作存值
});
// 获取 GPU Buffer 显存范围
const arrayBufferSecondMatrix = gpuBufferSecondMatrix.getMappedRange();
// 写入数据
new Float32Array(arrayBufferSecondMatrix).set(secondMatrix);
// 解除映射,以便 CPU 可以操作 Buffer
gpuBufferSecondMatrix.unmap();

// 结果矩阵
const resultMatrixBufferSize = Float32Array.BYTES_PER_ELEMENT * (2 + firstMatrix[0] * secondMatrix[1]);
const resultMatrixBuffer = device.createBuffer({
    size: resultMatrixBufferSize, // 和结果数组一样的大小,字节长度
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC // 标识这个 Buffer 只用作存值和赋值对象
});

绑定组布局和绑定组#

绑定组布局和绑定组是 WebGPU 中特有的概念。绑定组布局定义了着色器所需的输入/输出类型接口,绑定组是着色器输入/输出的实际数据。

在下面的示例中,绑定组布局定义了两个只读的缓冲区存储(用编号0、1绑定),以及一个普通缓冲区存储(用编号2绑定)。另一方面,与这个绑定组布局对应的绑定组分别将三个矩阵的GPU Buffer绑定到 0、1、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 标识的三个缓冲区存储 .该程序将使用firstMatrix、secondMatrix和resultMatrix作为输入和输出。

请注意,每个缓冲区存储都使用了一个 binding 该修饰符,与上面的绑定组布局和绑定组中定义的索引对应。

const shaderModule = device.createShaderModule({
  code: `
    struct Matrix {
      size : vec2<f32>,
      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, write> resultMatrix : Matrix;

    @stage(compute) @workgroup_size(8, 8)
    fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
      // 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()来创建它。它需要两个参数:我们之前创建的绑定组布局和计算着色器(指明作色器代码入口),以及定义计算着色器入口点(WGSL 函数)的计算阶段以及使用 创建的实际计算着色器模块。

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

命令提交#

用三个 GPU Buffer和具有绑定组布局的计算管道实例化绑定组后,是时候使用它们了。

让我们使用commandEncoder.beginComputePass()启动可编程计算传递编码器。我们将使用它来编码GPU 命令执行矩阵乘法 。使用passEncoder.setPipeline(computePipeline)设置计算管道,并绑定组,索引 0 对应于 WGSL 代码中的group(0)。

现在,让我们来谈谈这个计算着色器将如何在 GPU 上运行。我们的目标是逐步为结果矩阵的每个单元格并行执行此程序。例如,对于大小为 2 x 4 的结果矩阵,我们将调用passEncoder.dispatchWorkgroups(2, 4)来对执行命令进行编码。第一个参数“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 将计算结果复制到目标Buffer。最后,使用copyEncoder.finish()完成命令编码,并通过调用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返回的resolve值,此时GPU Buffer已完成映射,可以通过.getMappedRange()读取 GPU Buffer 的映射范围中的值。

运行我们的代码,会在浏览器 DevTools 控制台中打印的结果:“2, 2, 50, 60, 114, 140”。

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

最后一招#

使代码更易于阅读的一种方法是使用计算管道的 getBindGroupLayout 方法便捷地从着色器模块推断绑定组布局。这种方式无需创建自定义绑定组布局,以及无需在计算管道中指定管道布局。

这里有一个示例 available

性能调查结果#

最后,在 GPU 上运行矩阵乘法与在 CPU 上运行矩阵乘法二者的性能对比如何呢?为了找到答案,我为CPU编写了刚刚的程序。正如下图中看到的,当矩阵的大小大于256×256时,使用GPU是一个显而易见的选择。

这篇文章只是我探索WebGPU之旅的开始。敬请期待更多文章很快将更深入地介绍 GPU 计算以及在 WebGPU 中进行渲染(画布、纹理、采样器)。

阅读文章和使用GPU有疑问的,可以通过在Stack Overflow上提问或浏览其他开发人员提出的问题列表来获得帮助。

翻译自:web.dev/gpu-compute…