在浏览器上使用GPU计算

本文是关于我使用实验性WebGPU API并与有兴趣使用GPU进行数据并行计算的Web开发人员分享我的旅程。

背景

您可能已经知道，图形处理单元（GPU）是计算机中的电子子系统，最初专用于处理图形。但是，在过去的十年中，它已经发展成为一种更加灵活的体系结构，允许开发人员在利用GPU独特的体系结构的同时，实现多种类型的算法，而不仅仅是渲染3D图形。这些功能称为GPU计算，将GPU用作通用科学计算的协处理器称为通用GPU（GPGPU）编程。

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

WebGPU是一个低级API，例如WebGL。如您所见，它非常强大且非常冗长。但是没关系。我们正在寻找的是性能。

在本文中，我将重点介绍WebGPU的GPU计算部分，老实说，我只是在摸索表面，以便您可以自己开始游戏。我将深入探讨，并在即将发表的文章中介绍WebGPU渲染（画布，纹理等）。

Dogfood：Chrome 78 for macOS现已提供实验性标记，现已提供WebGPU。您可以在启用它chrome://flags/#enable-unsafe-webgpu。API不断变化，目前不安全。由于尚未为WebGPU API实现GPU沙箱，因此可以读取其他进程的GPU数据！不要启用它来浏览网络。

存取GPU

在WebGPU中，访问GPU很容易。调用navigator.gpu.requestAdapter() 将返回JavaScript承诺，该承诺将与GPU适配器异步解析。将此适配器视为图形卡。它可以集成（与CPU在同一芯片上）或离散（通常是性能更高但使用更多功率的PCIe卡）。

拥有GPU适配器后，致电adapter.requestDevice()以获得一个诺言，该诺言将用于执行一些GPU计算的GPU设备上。

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();

这两个功能都带有选项，使您可以具体确定所需的适配器类型（电源偏好）和设备（扩展，限制）。为了简单起见，我们将在本文中使用默认选项。

写缓冲存储器

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

下面的示例向您展示如何写入四个字节以缓冲可从GPU访问的内存。它调用device.createBufferMapped()哪个将获取缓冲区的大小及其使用情况。即使GPUBufferUsage.MAP_WRITE 此特定调用不需要使用标志，也要明确表明我们要写入此缓冲区。它产生一个GPU缓冲区对象及其关联的原始二进制数据缓冲区。

如果您已经玩过字节写入，这是很熟悉的ArrayBuffer。使用a TypedArray并将值复制到其中。

// Get a GPU buffer in a mapped state and an arrayBuffer for writing.
const [gpuBuffer, arrayBuffer] = device.createBufferMapped({
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE
});

// 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缓冲区是在未映射状态下通过sync创建的 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, arrayBuffer] = device.createBufferMapped({
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC
});

// 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。GPUBuffer另一方面，on中的方法 是“无缓冲的”，这意味着它们在被调用时会自动执行。

有了GPU命令编码器后，请copyEncoder.copyBufferToBuffer() 按如下所示进行调用，以将该命令添加到命令队列中，以供以后执行。最后，通过调用完成编码命令，copyEncoder.finish()并将其提交到GPU设备命令队列。队列负责处理device.defaultQueue.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.defaultQueue.submit([copyCommands]);

至此，GPU队列命令已经发送，但不一定执行。要读取第二个GPU缓冲区，请调用gpuReadBuffer.mapReadAsync()。它返回一个promise，ArrayBuffer一旦所有排队的GPU命令都已执行，它将以与第一个GPU缓冲区相同的值进行解析。

// Read buffer.
const copyArrayBuffer = await gpuReadBuffer.mapReadAsync();
console.log(new Uint8Array(copyArrayBuffer));

您可以尝试此示例。

简而言之，关于缓冲存储器操作，您需要记住以下几点：

必须取消映射GPU缓冲区才能在设备队列提交中使用。
映射后，可以使用JavaScript读写GPU缓冲区。
当GPU缓存映射mapReadAsync()，mapWriteAsync()以及 createBufferMapped()被调用。

着色器编程

在GPU上运行的仅执行计算（而不绘制三角形）的程序称为计算着色器。它们由数百个GPU内核（小于CPU内核）并行执行，这些GPU内核共同操作以处理数据。它们的输入和输出是WebGPU中的缓冲区。

为了说明在WebGPU中计算着色器的用法，我们将玩矩阵乘法，这是下面说明的机器学习中的一种常见算法。

简而言之，这是我们要做的：

1. 创建三个GPU缓冲区（两个用于矩阵相乘，一个用于结果矩阵）
2. 描述计算着色器的输入和输出
3. 编译计算着色器代码
4. 设置计算管道
5. 批量提交编码后的命令到GPU
6. 读取结果矩阵GPU缓冲区

GPU缓冲区创建
为了简单起见，矩阵将表示为浮点数列表。第一个元素是行数，第二个元素是列数，其余是矩阵的实际数。

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

const adapter = await navigator.gpu.requestAdapter();
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, arrayBufferFirstMatrix] = device.createBufferMapped({
  size: firstMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
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, arrayBufferSecondMatrix] = device.createBufferMapped({
  size: secondMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
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,
      type: "readonly-storage-buffer"
    },
    {
      binding: 1,
      visibility: GPUShaderStage.COMPUTE,
      type: "readonly-storage-buffer"
    },
    {
      binding: 2,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    }
  ]
});

const bindGroup = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [
    {
      binding: 0,
      resource: {
        buffer: gpuBufferFirstMatrix
      }
    },
    {
      binding: 1,
      resource: {
        buffer: gpuBufferSecondMatrix
      }
    },
    {
      binding: 2,
      resource: {
        buffer: resultMatrixBuffer
      }
    }
  ]
});

计算着色器代码
用于乘法矩阵的计算着色器代码用GLSL编写，GLSL是WebGL中使用的高级着色语言，其语法基于C编程语言。无需赘述，您应该在下面用关键字标记的三个存储缓冲区中找到buffer。该程序将使用 firstMatrix和secondMatrix作为输入（只读）和resultMatrix其输出。

请注意，每个存储缓冲区都有一个binding限定符，该限定符与在绑定组布局和上面声明的绑定组中定义的相同索引相对应。

const computeShaderCode = `#version 450

  layout(std430, set = 0, binding = 0) readonly buffer FirstMatrix {
      vec2 size;
      float numbers[];
  } firstMatrix;

  layout(std430, set = 0, binding = 1) readonly buffer SecondMatrix {
      vec2 size;
      float numbers[];
  } secondMatrix;

  layout(std430, set = 0, binding = 2) buffer ResultMatrix {
      vec2 size;
      float numbers[];
  } resultMatrix;

  void main() {
    resultMatrix.size = vec2(firstMatrix.size.x, secondMatrix.size.y);

    ivec2 resultCell = ivec2(gl_GlobalInvocationID.x, gl_GlobalInvocationID.y);
    float result = 0.0;
    for (int i = 0; i < firstMatrix.size.y; i++) {
      int a = i + resultCell.x * int(firstMatrix.size.y);
      int b = resultCell.y + i * int(secondMatrix.size.y);
      result += firstMatrix.numbers[a] * secondMatrix.numbers[b];
    }

    int index = resultCell.y + resultCell.x * int(secondMatrix.size.y);
    resultMatrix.numbers[index] = result;
  }
`;

管道设置
Chrome中的WebGPU当前使用字节码代替原始的GLSL代码。这意味着我们必须computeShaderCode在运行计算着色器之前进行编译。对我们来说幸运的是，@ webgpu / glslang包使我们能够computeShaderCode 以Chrome中的WebGPU接受的格式进行编译。此字节码格式基于SPIR-V的安全子集。

请注意，在编写WebGPU的着色语言时，“ Web上的GPU” W3C社区组仍未决定。

import glslangModule from 'https://unpkg.com/@webgpu/[email protected]/dist/web-devel/glslang.js';

计算管道是实际描述我们将要执行的计算操作的对象。通过调用创建它device.createComputePipeline()。它包含两个参数：我们之前创建的绑定组布局，以及一个计算阶段，该阶段定义了我们的计算着色器（mainGLSL函数）和使用编译的实际计算着色器模块的入口点glslang.compileGLSL()。

const glslang = await glslangModule();

const computePipeline = device.createComputePipeline({
  layout: device.createPipelineLayout({
    bindGroupLayouts: [bindGroupLayout]
  }),
  computeStage: {
    module: device.createShaderModule({
      code: glslang.compileGLSL(computeShaderCode, "compute")
    }),
    entryPoint: "main"
  }
});

命令提交
在使用我们的三个GPU缓冲区和具有绑定组布局的计算管道实例化绑定组之后，就该使用它们了。

让我们开始使用的可编程计算通过编码器 commandEncoder.beginComputePass()。我们将使用它来编码将执行矩阵乘法的GPU命令。使用设置它的管道 passEncoder.setPipeline(computePipeline)和其绑定组在索引0处 passEncoder.setBindGroup(0, bindGroup)。索引0对应set = 0于GLSL代码中的限定符。

现在，让我们讨论一下此计算着色器将如何在GPU上运行。我们的目标是逐步针对结果矩阵的每个单元并行执行此程序。例如，对于大小为2乘4的结果矩阵，我们将调用 passEncoder.dispatch(2, 4)以对执行命令进行编码。第一个参数“ x”是第一个维度，第二个参数“ y”是第二个维度，最后一个参数“ z”是第三个维度，默认情况下为1，因为我们在这里不需要它。在GPU计算世界中，对在一组数据上执行内核功能的命令进行编码称为调度。

在我们的代码中，“ x”和“ y”将分别是第一个矩阵的行数和第二个矩阵的列数。这样，我们现在可以使用发送调度调用

passEncoder.dispatch(firstMatrix[0], secondMatrix[1])

如上图所示，每个着色器都可以访问一个唯一的 gl_GlobalInvocationID对象，该对象将用于知道要计算哪个结果矩阵像元。

const commandEncoder = device.createCommandEncoder();

const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);
passEncoder.dispatch(firstMatrix[0] /* x */, secondMatrix[1] /* y */);
passEncoder.endPass();

要结束计算过程编码器，请调用passEncoder.endPass()。然后，创建一个GPU缓冲区用作复制目标矩阵缓冲区的目标 copyBufferToBuffer。最后，完成对命令的编码， copyEncoder.finish()并通过调用device.defaultQueue.submit()GPU命令将其提交到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.defaultQueue.submit([gpuCommands]);

读取结果矩阵
读取结果矩阵就像调用gpuReadBuffer.mapReadAsync() 并记录ArrayBuffer所产生的Promise返回的结果一样容易。

在我们的代码中，在DevTools JavaScript控制台中记录的结果是“ 2、2、50、60、114、140”。

// Read buffer.
const arrayBuffer = await gpuReadBuffer.mapReadAsync();
console.log(new Float32Array(arrayBuffer));

恭喜你！你做到了。你可以来个实例看看。

最后一招
使代码更易于阅读的getBindGroupLayout一种方法是使用计算管道的便捷 方法从着色器模块推断绑定组的布局。此技巧使您无需创建自定义绑定组布局并在您的计算管道中指定管道布局，如下所示。

的图示getBindGroupLayout为前一个样本是可用的。

 const computePipeline = device.createComputePipeline({
-  layout: device.createPipelineLayout({
-    bindGroupLayouts: [bindGroupLayout]
-  }),
   computeStage: {

-// Bind group layout and bind group
- const bindGroupLayout = device.createBindGroupLayout({
-   entries: [
-     {
-       binding: 0,
-       visibility: GPUShaderStage.COMPUTE,
-       type: "readonly-storage-buffer"
-     },
-     {
-       binding: 1,
-       visibility: GPUShaderStage.COMPUTE,
-       type: "readonly-storage-buffer"
-     },
-     {
-       binding: 2,
-       visibility: GPUShaderStage.COMPUTE,
-       type: "storage-buffer"
-     }
-   ]
- });
+// Bind group
  const bindGroup = device.createBindGroup({
-  layout: bindGroupLayout,
+  layout: computePipeline.getBindGroupLayout(0 /* index */),
   entries: [

性能测试
图5. GPU与CPU基准测试

那么在GPU上运行矩阵乘法与在CPU上运行矩阵乘法相比又如何呢？为了找出答案，我编写了刚刚针对CPU编写的程序。正如您在下图中所看到的，当矩阵的大小大于256 x 256时，使用GPU的全部功能似乎是一个显而易见的选择。

本文只是我探索WebGPU的旅程的开始。很快将有更多文章发表，这些文章将更深入地介绍GPU Compute，以及有关WebGPU中渲染（画布，纹理，采样器）的工作方式。