【译】在Web中使用GPU Compute

原文链接：Get started with GPU Compute on the Web

背景

正如我们所知，图形处理单元（GPU）是计算机中的一个电子子系统，最初专门用于处理图形。然而，在过去10年中，它已经发展成为一种更灵活的架构，利用GPU的独特架构，开发人员可以实现多种类型的算法，而不仅仅是渲染3D图形。这些功能称为GPU Compute，使用GPU作为通用科学计算的协处理器称为general-purpose GPU（GPGPU）编程。

GPU Compute为最近的机器学习热潮做出了重大贡献，例如卷积神经网络和其他模型可以利用此功能在GPU上更高效地运行。由于当前的Web平台缺乏GPU Compute功能，W3C的“GPU for the Web”社区组正在设计一个API来暴露出GPU API，以便在大多数设备上使用，此API称为WebGPU。

WebGPU是一种底层API，例如WebGL。正如我们所看到的，它非常强大且非常详尽。但没关系，我们关注的是性能。

在本文中，我将重点关注WebGPU的GPU Compute部分，说实话，我会讲的浅显易懂，这样大家就可以玩出花样。我会在即将发表的文章中介绍并深入探讨WebGPU渲染（画布，纹理等）。

WebGPU目前在macOS的Chrome 78中可通过experimental flag开启使用。您可以在chrome://flags/ #enable-unsafe-webgpu中启用它。API会经常变动，目前不太安全。由于目前还没有为WebGPU API实现GPU沙盒，因此可以读取其他进程的GPU数据！因此请勿在启用WebGPU下模式下浏览网页。

访问GPU

在WebGPU中访问GPU很容易。调用navigator.gpu.requestAdapter()会返回一个JavaScript Promise，经过解析异步返回一个GPU adapter（GPU适配器）。我们可以将此适配器视为显卡。它可以是集成显卡（与CPU在同一芯片上）或独立显卡（通常是性能更高但功耗更大的PCIe卡）。

一旦获得GPU适配器后，就可以调用adapter.requestDevice()并返回一个将使用GPU设备进行解析的promise，使用它来进行一些GPU计算。

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

这两个函数都允许传入所需要的适配器类型（电源首选项）和设备（扩展，限制）的参数。为简单起见，我们将使用本文中的默认选项。

写buffer

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

下面的示例显示了如何将四个字节写入可从GPU访问的buffer memory。它调用device.createBufferMappedAsync()来得到缓冲区的大小并表示用来干嘛。即使这个API调用不需要使用标志GPUBufferUsage.MAP_WRITE，也会表明我们要写入此buffer。生成的promise将返回GPU buffer对象及其关联的原始二进制数据array buffer。

如果你已经玩过ArrayBuffer，那么对写字节必定很熟悉;使用TypedArray并将值复制进去。

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

// Write bytes to buffer.
new Uint8Array(arrayBuffer).set([0, 1, 2, 3]);

此时，GPU buffer被映射，这意味着它由CPU拥有，并且可以从JavaScript读取/写入。因此GPU可以访问它，它必须是未映射的，这就像调用gpuBuffer.unmap()一样简单。

映射或未映射的概念就是来防止GPU和CPU同时访问内存的竞争条件。

读buffer

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

由于我们在第一个GPU缓冲区中写入并且我们想将其复制到第二个GPU缓冲区，因此需要新的使用标志GPUBufferUsage.COPY_SRC。使用同步device.createBuffer（）以未映射状态创建第二个GPU缓冲区。它的用法标志是GPUBufferUsage.COPY_DST |GPUBufferUsage.MAP_READ，因为它将用作第一个GPU缓冲区的目标，并在执行GPU复制命令后在JavaScript中读取。

// Get a GPU buffer and an arrayBuffer for writing.
// Upon success the GPU buffer is returned in the mapped state.
const [gpuWriteBuffer, arrayBuffer] = await device.createBufferMappedAsync({
  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中，device.createCommandEncoder（）返回的GPU命令编码器是构建一批“缓冲”命令的JavaScript对象，这些命令将在某个时刻发送到GPU。另一方面，GPUBuffer上的方法是“无缓冲的”，这意味着它们在被调用时以原子方式执行。

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

此时，已发送GPU队列命令，但不一定执行。要读取第二个GPU缓冲区，请调用gpuReadBuffer.mapReadAsync（）。一旦执行了所有排队的GPU命令，它将返回一个承诺，该承诺将使用包含与第一个GPU缓冲区相同的值的ArrayBuffer来解析。

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

可以戳示例

简而言之，这是关于缓冲存储器操作需要记住的内容：

• 必须取消映射GPU缓冲区才能在设备队列提交中使用。

• 映射后，可以使用JavaScript读取和写入GPU缓冲区。

• 调用mapReadAsync（），mapWriteAsync（），createBufferMappedAsync（）和createBufferMapped（）时，将映射GPU缓冲区。

着色编程

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

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

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

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

GPU buffer创建

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

三个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] = await device.createBufferMappedAsync({
  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] = await device.createBufferMappedAsync({
  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({
  bindings: [
    {
      binding: 0,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    },
    {
      binding: 1,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    },
    {
      binding: 2,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    }
  ]
});

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

计算着色器代码

用于乘法矩阵的计算着色器代码用GLSL编写，GLSL是WebGL中使用的高级着色语言，其具有基于C编程语言的语法。在不详细说明的情况下，您应该在下面找到标有关键字缓冲区的三个存储缓冲区。该程序将使用firstMatrix和secondMatrix作为输入，并使用resultMatrix作为其输出。

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

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包允许我们以Chrome中的WebGPU接受的格式编译computeShaderCode。此字节码格式基于SPIR-V的安全子集。

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

import glslangModule from 'https://unpkg.com/@webgpu/glslang/web/glslang.js';

计算管道是实际描述我们将要执行的计算操作的对象。通过调用device.createComputePipeline（）创建它。它有两个参数：我们之前创建的绑定组布局，以及定义计算着色器（主GLSL函数）的入口点的计算阶段和使用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.setPindline（computePipeline）设置其管道，使用passEncoder.setBindGroup（0，bindGroup）在索引0处设置其绑定组。索引0对应于GLSL代码中的set = 0限定符。

现在，我们来谈谈这个计算着色器将如何在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（）完成编码命令，并通过使用GPU命令调用device.getQueue（）。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.getQueue().submit([gpuCommands]);

读取结果矩阵

读取结果矩阵就像调用gpuReadBuffer.mapReadAsync（）并记录生成的promise返回的ArrayBuffer一样简单。

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

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

示例请戳这里

性能调查

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

这篇文章只是我探索WebGPU之旅的开始。很快就会有更多文章介绍GPU Compute中更深入的潜力以及渲染（画布，纹理，采样器）在WebGPU中的工作方式。