Hello WebGPU —— Compute Shader基础

寒璃
2,518 阅读5分钟

「这是我参与2022首次更文挑战的第8天，活动详情查看：2022首次更文挑战」。

前言

也许你已经知道，GPU在过去的10年中朝着更加灵活的架构发展并且允许开发者在其中实现一些算法，它的作用已经不仅仅是渲染3D图形了。这种能力我们称其为GPU计算，使用GPU进行通用计算的编程方式我们称其为GPGPU编程（General-Purpose GPU Programming）。

那么，今天我们暂时放下渲染部分的内容，进入GPU计算管线的知识中来。现在让开始今天的旅程吧~

目标

今天我们要实现的例子是使用GPU进行矩阵计算。比如：

[12345678]×[12345678]\begin{bmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \end{bmatrix} \times \begin{bmatrix} 1 & 2 \\ 3 & 4 \\ 5 & 6 \\ 7 & 8 \\ \end{bmatrix}

我们使用GPU进行矩阵运算的步骤可以简单的概括为以下几步：

  1. 创建3个GPUBuffer （其中两个用于保存输入矩阵，一个用于保存矩阵相乘的结果）
  2. 编写计算shader的代码
  3. 建立计算管线
  4. 提交计算命令到GPU
  5. 从GPUBuffer中读取结果

让我们按照上面的步骤进行：

Coding

创建GPUBuffer

首先，我们先准备矩阵数据：

const firstMatrix = new Float32Array([
            2, 4,
            1, 2, 3, 4,
            5, 6, 7, 8
        ]);

const secondMatrix = new Float32Array([
            4, 2,
            1, 2,
            3, 4,
            5, 6,
            7, 8
        ])

上面的矩阵数据中，前两位表示矩阵的大小，比如第一个矩阵中的2，4表示该矩阵是2x4大小的矩阵。

紧接着，我们创建2个 GPUBuffer 并且将数据写入 GPUBuffer 中；再准备一个用于保存结果的 GPUBuffer

const gpuBufferFirstMatrix = device.createBuffer({
    mappedAtCreation: true,
    size: firstMatrix.byteLength,
    usage: GPUBufferUsage.STORAGE,
});
const arrayBufferFirstMatrix = gpuBufferFirstMatrix.getMappedRange();
new Float32Array(arrayBufferFirstMatrix).set(firstMatrix);
gpuBufferFirstMatrix.unmap();


const gpuBufferSecondMatrix = device.createBuffer({
    mappedAtCreation: true,
    size: secondMatrix.byteLength,
    usage: GPUBufferUsage.STORAGE,
});
const arrayBufferSecondMatrix = gpuBufferSecondMatrix.getMappedRange();
new Float32Array(arrayBufferSecondMatrix).set(secondMatrix);
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时为不同用途的Buffer指定不同的 usage，对于输入矩阵，我们使用 GPUBufferUsage.STORAGE即可，但是用于保存结果的矩阵Buffer还需要加上 GPUBufferUsage.COPY_SRC，因为我们需要从其中将数据拷贝出来，以方便我们从GPUBuffer中进行读取。

编写Shader程序

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(2, 2)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    if (global_id.x >= u32(firstMatrix.size.x) || global_id.y >= u32(secondMatrix.size.y)) {
        return;
    }

    resultMatrix.size = vec2<f32>(firstMatrix.size.x, secondMatrix.size.y);

    let resultCell = vec2<u32>(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;
}

这里说明一下，形如 [[stage(vertex)]] 这样风格的代码，根据WebGPU的最新标准，现在推荐采用 @stage(vertex) 这样的写法，所以从本篇文章开始，WGSL的代码会采用最新的代码风格。

我们解释一下上述的代码： var<storage, read> firstMatrix: Matrix 其中 storage 表示的是 address space, 除了 storage 的类型以外还有 uniform, workgroup 等，read 表示的是 access_mode 表示该变量的读写类型，有： read，write，read_write 三种模式。

接下来，我们注意到 @workgroup_size(2, 2)这一注解。

workgroup

workgroup是一系列的能够同时执行 compute shader 程序的“线程”（这里我们可以将其理解为线程），并且能够共享在 workgroup 地址空间中的变量。换句话说，我们可以将 workgroup 理解为是 “线程组” 的概念。

该“线程组”具有3个维度 workgroup(i, j, k)，可以理解为是一个3维数组。数组中的每一个单元都会执行一次 compute shader，而且这数组中的每个执行单元都是并行执行的。

在 compute shader 中，可以获取到“线程组”中第几个执行单元的ID，也就是上述代码中的 global_invocation_id。其ID的计算方式如下：

i + (j * workgroup_size_x) + (k * workgroup_size_x * workgroup_size_y)

当我们调用了 dispatch API后，WebGPU会开始执行 compute shader 程序，dispatch 命令指定了 "dispatch size"，它是一个3维的数组 [group_cout_x, group_count_y, group_count_z] ，它表示的有多少个 workgroup 应该被执行。

所以，对于执行单元的每次调用，都具有唯一的ID：

  • 0 ≤ CSi < workgroup_size_x × group_count_x
  • 0 ≤ CSj < workgroup_size_y × group_count_y
  • 0 ≤ CSk < workgroup_size_z × group_count_z

后续的矩阵乘法的代码部分，就不过多解释了。

建立计算管线


const computePipeline = device.createComputePipeline({
    compute: {
        module: device.createShaderModule({
            code: computeShader,
        }),
        entryPoint: 'main',
    },
});
const bindGroup = device.createBindGroup({
    layout: computePipeline.getBindGroupLayout(0),
    entries: [
        {
            binding: 0,
            resource: {
                buffer: gpuBufferFirstMatrix,
            },
        },
        {
            binding: 1,
            resource: {
                buffer: gpuBufferSecondMatrix,
            },
        },
        {
            binding: 2,
            resource: {
                buffer: resultMatrixBuffer,
            },
        },
    ],
});

与渲染管线不同的是，计算管线不需要指定 Texture，深度模板测试的配置等，这里计算管线仅仅指定了一个shader 程序。

紧接着，我们与渲染流程中一样，创建了 bindGroup。

提交计算命令到GPU


const commandEncoder = this.device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(this.pipeline);
passEncoder.setBindGroup(0, this.bindGroup);
const x = Math.ceil(firstMatrix[0] / 2);
const y = Math.ceil(secondMatrix[1] / 2);
passEncoder.dispatch(x, y);
passEncoder.endPass();

在 dispatch之前的代码不再过多赘述，这里说明一下 dispatch 中传入的 x 和 y 是什么意思，刚刚在解释 workgroup 的时候说过了，它们表示的是有多少个 workgroup 需要被执行，我们现在只需要1个workgroup在GPU中执行，所以我们传入(1, 1)即可。

从GPUBuffer中读取结果

现在我们计算的结果被保存在 resultMatrixBuffer 中，现在我们需要将数据从其中读取出来。首先我们还需要创建一个Buffer来将 resultMatrixBuffer 拷贝出来。

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

这里我们需要指定 usage 为 GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ。然后我们可以将命令提交到GPU中。

提交了计算命令后，我们可以从GPU中获取结果。通过 mapAsync 这个API可以读取到GPU中的数据。

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

最后我们可以得到如下的结果： image.png

前两位表示的最终矩阵的大小，后面的数字表示的是矩阵中的数字。

总结

今天我们学习了如何使用WebGPU进行计算，大致的流程与使用GPU进行渲染的流程大同小异，其中比较难以理解的点在于 workgroup上。还有如何读取 GPUBuffer 中的数据，以上都是我们需要掌握的点，接下来我们会开始使用WebGPU开始进行一些稍微复杂一点的计算。

如果你觉得本文不错，请为作者点个赞哦~ 您的点赞就是作者更新的动力~！

avatar
文章
阅读
粉丝