WebGPU计算着色器使用入门

Number of views 92

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

背景

你可能已经知道,图形处理单元(GPU)是计算机内部的一个电子子系统,最早是专门为处理图像而设计的。但是,在过去十年中,它逐渐发展成一种更灵活的架构,使得开发者不仅可以用来渲染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() 即可获取一个 Promise,该 Promise 将解析为您用于执行某些 GPU 计算的 GPU 设备。

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

这两个函数都接受可选的参数,让你能够具体指定你想要的适配器(比如电源偏好)和设备(比如扩展功能、限制条件)类型。为了简单起见,本文中我们将使用默认选项。

写入缓冲区内存

我们来看看如何使用 JavaScript 将数据写入 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()方法就行了。

映射和未映射这两个概念主要是为了避免CPU和GPU同时访问同一块内存时可能出现的冲突问题。换句话说,为了避免它们“撞车”,我们需要在适当的时候切换缓冲区的状态:当需要通过JavaScript写入数据时就映射它,写完之后为了让GPU处理这些数据就解除映射。这样就能确保数据处理既安全又高效。

读取缓冲区内存

现在来看看如何把一个GPU缓冲区的内容复制到另一个GPU缓冲区,并且再读取回来。

因为我们是先把数据写入第一个GPU缓冲区,然后想把它复制到第二个GPU缓冲区,所以需要给第一个缓冲区添加一个新的使用标志GPUBufferUsage.COPY_SRC。至于第二个GPU缓冲区,这次我们用device.createBuffer()创建它时让它处于未映射状态。它的使用标志是GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,这意味着它将作为第一个缓冲区的复制目标,并且在GPU执行完复制命令后,我们可以通过JavaScript读取它。

简单来说,就是先设置好第一个缓冲区以便它可以被复制出去,然后创建第二个缓冲区准备接收这些数据并且之后可以被读取。这样做就能实现从一个GPU缓冲区向另一个缓冲区复制数据并随后读取的功能。

// 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()结束命令编码,并通过device.queue.submit()把这些命令提交给GPU设备的命令队列。这个队列负责处理通过device.queue.submit()传入的GPU命令数组,按顺序原子性地执行所有存储在数组中的命令。

简单来说,就是先准备好一个命令编码器,然后用它来添加你想让GPU执行的命令(比如复制缓冲区),之后完成编码并把这一批命令提交给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缓冲区的内容,你需要调用 gpuReadBuffer.mapAsync()并指定模式为 GPUMapMode.READ。这个方法会返回一个Promise,在GPU缓冲区成功映射后这个Promise就会被resolve。接下来,使用 gpuReadBuffer.getMappedRange()获取映射区域,一旦所有排队的GPU命令都执行完毕,这个映射区域里就会包含和第一个GPU缓冲区相同的数据。

换句话说,就是先异步请求映射你想读取的GPU缓冲区,并等待它准备好。一旦缓冲区映射好了,你就可以通过获取映射范围来访问里面的数据了,这样就能确保你在读取数据之前,所有之前的GPU命令都已经执行完成。

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

以下是上述过程的完整代码:

(async () => {
  if (!("gpu" in navigator)) {
    console.log(
      "WebGPU is not supported. Enable chrome://flags/#enable-unsafe-webgpu flag."
    );
    return;
  }

  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) {
    console.log("Failed to get GPU adapter.");
    return;
  }
  const device = await adapter.requestDevice();

  // 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({
    mappedAtCreation: false,
    size: 4,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
  });

  // 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]);

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

  console.log(new Uint8Array(copyArrayBuffer));
})();

简单来说,关于缓冲区内存操作,你需要记住以下几点:

  • 要在设备队列提交中使用GPU缓冲区时,必须先解除它们的映射状态。
  • 当缓冲区处于映射状态时,可以通过JavaScript对其进行读写操作。
  • 当你调用mapAsync()或者在创建缓冲区时通过createBuffer()并将mappedAtCreation设置为true,GPU缓冲区就会被映射。

着色器程序

在GPU上运行的只进行计算(而不是绘制三角形)的程序叫做计算着色器。这些计算着色器会在数百个GPU核心(这些核心比CPU核心小)上并行执行,这些核心协同工作来处理数据。在WebGPU中,它们的输入和输出都是缓冲区。

为了展示WebGPU中计算着色器的使用,我们将通过下面的例子来玩转矩阵乘法,这是一个在机器学习中常见的算法。

简单来说,计算着色器就是在GPU上专门用来做计算的小程序,不负责画图。它们能在很多个小的核心上一起快速运行,处理大量数据。在WebGPU里,你可以把数据放进缓冲区,然后让计算着色器对这些数据进行操作。我们接下来会用矩阵乘法的例子来具体说明怎么在WebGPU里使用计算着色器。

image1738736509786.png

简单来说,我们接下来要做的是:

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

这就是整个过程的简要步骤。通过这些步骤,我们可以利用WebGPU执行矩阵乘法等计算任务。

创建 GPU 缓冲区

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

image1738737029051.png

要创建的三个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标识的存储缓冲区。程序会把firstMatrix和secondMatrix作为输入,而resultMatrix作为输出。

需要注意的是,每个存储缓冲区都有一个绑定装饰符(binding decoration),这个装饰符对应的索引号与上面提到的绑定组布局和绑定组中定义的一致。换句话说,就是这些缓冲区在代码里怎么指定的,在实际设置的时候就得对应上,这样才能确保数据能正确地流入和流出着色器。

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()来创建它。这个方法需要两个参数:一个是之前创建的绑定组布局,另一个是计算阶段,它定义了计算着色器的入口点(即WGSL的主函数)以及使用device.createShaderModule()创建的实际计算着色器模块。

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

提交命令

上面我们使用3个GPU缓冲区实例化一个绑定组,并用一个绑定组布局创建了一个计算管线之后,现在是时候使用它们了。

首先,我们通过 commandEncoder.beginComputePass()启动一个可编程的计算通道编码器。我们将用它来编码执行矩阵乘法的GPU命令。设置计算管线为 passEncoder.setPipeline(computePipeline),并设置索引为0的绑定组为 passEncoder.setBindGroup(0, bindGroup)。这里的索引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计算领域,将一个内核函数在一组数据上执行的命令编码称为 调度(dispatching)

image1738739240555.png

我们的计算着色器在WGSL代码中设定的工作组大小是 (8, 8)。因此,对于第一个矩阵的行数 "x"和第二个矩阵的列数 "y",我们需要分别除以8来确定需要多少个工作组。这样我们就可以通过 passEncoder.dispatchWorkgroups(firstMatrix[0] / 8, secondMatrix[1] / 8)来分派一个计算调用。这里 dispatchWorkgroups()的参数就是需要运行的工作组网格的数量。

正如上面的图所示,每个着色器都能访问一个唯一的内置对象 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()完成命令编码,并通过调用 device.queue.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.queue.submit([gpuCommands]);

读取结果矩阵

读取结果矩阵就像调用 gpuReadBuffer.mapAsync()并使用 GPUMapMode.READ模式一样简单,然后等待返回的 Promise完成,这表示GPU缓冲区现在已经被映射。此时,你可以通过调用 gpuReadBuffer.getMappedRange()来获取映射区域。

image1738739927712.png在我们的代码中,最终在开发者工具的JavaScript控制台中打印出的结果是 "2, 2, 50, 60, 114, 140"

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

恭喜你!你做到了,下面是完整代码:

(async () => {
  if (!("gpu" in navigator)) {
    console.log(
      "WebGPU is not supported. Enable chrome://flags/#enable-unsafe-webgpu flag."
    );
    return;
  }

  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) {
    console.log("Failed to get GPU 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
  });

  // 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"
        }
      }
    ]
  });

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

  // Compute shader code

  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, read_write> resultMatrix : Matrix;

      @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;
      }
    `
  });
  
  // Pipeline setup

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

  // Commands submission

  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();

  // 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]);

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

最后一个小技巧

使你的代码更易读的一个方法是使用计算管线的便捷方法 getBindGroupLayout从着色器模块推断绑定组布局。这个小技巧省去了手动创建自定义绑定组布局和在计算管线中指定管线布局的需要,如下所示。

对于前面的例子,有一个使用 getBindGroupLayout方法的示例图解可供参考。

换句话说,就是你可以利用计算管线里的 getBindGroupLayout方法来自动获取绑定组布局,而不需要自己手动去定义它。这样可以让代码变得更简洁易读。对于之前的例子,有一个图解可以帮你更好地理解这一点。

// Pipeline setup
  
  const computePipeline = device.createComputePipeline({
    layout: "auto",
    compute: {
      module: shaderModule,
      entryPoint: "main"
    }
  });


  // Bind group

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

性能对比

那么,在GPU上运行矩阵乘法与在CPU上运行相比有何不同呢?为了找出答案,我编写了上述描述的程序用于CPU。如你在下面的图表中所见,当矩阵的大小超过256x256时,充分利用GPU的性能似乎是一个明显的选择。

image1738740572322.png这篇文章只是我探索WebGPU旅程的开始。敬请期待更多即将发布的文章,这些文章将会深入探讨GPU计算以及WebGPU中渲染(画布、纹理、采样器)的工作原理。

0 Answers