各位观众老爷，晚上好！我是你们的老朋友，今晚咱们聊聊 WebGPU 里那些让人头大，但又不得不学的硬骨头：Compute Shaders、Shared Memory、Atomic Operations，以及 Synchronization Barriers。保证让你们听完之后，感觉像是打通了任督二脉，功力大增！
WebGPU Compute Shaders：GPU 上的搬砖工
首先，咱们得明白 Compute Shaders 是个啥。简单来说，它就是让 GPU 不光干渲染的活儿，还能干一些通用计算的活儿。想象一下，你雇了一群 GPU，让它们帮你算算数学题、处理图像、或者跑跑物理模拟，是不是很爽？
Compute Shaders 用一种叫做 GLSL 的语言来编写（WebGPU 里用的是 WGSL，但本质差不多）。它运行在 GPU 的每个计算单元上，并行处理数据。这就像雇了一堆搬砖工，每个人搬一块砖，效率杠杠的。
Shared Memory：工友们的小仓库
光有搬砖工还不行，得给他们提供一个存放临时数据的地方，这就是 Shared Memory（共享内存）。每个 Compute Shader 的工作组（Workgroup）都拥有一块共享内存。这个内存对同一个工作组内的所有 Shader 实例都是可见的。你可以把它想象成工地上工友们共用的一个小仓库，大家可以往里面放东西，也可以从里面拿东西。
Atomic Operations：谁先抢到算谁的！
好，现在工地上有个问题：好几个人同时想往小仓库里放东西，或者同时想拿东西，怎么办？如果处理不当，就会出现数据竞争，导致结果出错。这时候，Atomic Operations（原子操作）就派上用场了。
原子操作是一种不可分割的操作，它可以保证在同一时间只有一个 Shader 实例能够访问某个特定的内存位置。这就像给小仓库加了个锁，每次只允许一个人进去操作，操作完再把锁交给下一个人。
常见的原子操作包括：
atomicAdd(memory, value)：原子加法，将value加到memory指向的内存位置。atomicSub(memory, value)：原子减法，将value从memory指向的内存位置减去。atomicMin(memory, value)：原子最小值，将memory指向的内存位置的值更新为min(memory, value)。atomicMax(memory, value)：原子最大值，将memory指向的内存位置的值更新为max(memory, value)。atomicExchange(memory, value)：原子交换，将memory指向的内存位置的值与value交换。atomicCompareExchangeWeak(memory, expected, desired)：原子比较并交换，如果memory指向的内存位置的值等于expected，则将其更新为desired。atomicLoad(memory): 原子加载，从memory指向的内存位置原子地加载值。atomicStore(memory, value): 原子存储，将value原子地存储到memory指向的内存位置。
Synchronization Barriers：排队，一个一个来！
原子操作解决了单个内存位置的并发访问问题，但有时候我们需要保证整个工作组内的所有 Shader 实例都完成了某个操作，才能进行下一步操作。这时候，Synchronization Barriers（同步屏障）就派上用场了。
同步屏障就像工地上的一声哨响，所有工人都听到哨响后，必须停下手里的活儿，等待其他人也停下来，然后才能一起进行下一步工作。
在 WebGPU 的 WGSL 中，我们使用 workgroupBarrier() 函数来实现同步屏障。
实战演练：用 Compute Shader 计算数组的和
光说不练假把式，咱们来一个实际的例子：用 Compute Shader 计算一个数组的和。
1. WGSL 代码 (sum.wgsl):
struct Data {
numbers: array<f32>,
}
@group(0) @binding(0) var<storage, read> input: Data;
@group(0) @binding(1) var<storage, write> output: Data;
var<workgroup> shared_data : array<f32, 256>; // 假设工作组大小为 256
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) local_id : vec3<u32>,
@builtin(workgroup_id) workgroup_id : vec3<u32>) {
let index = local_id.x;
shared_data[index] = input.numbers[workgroup_id.x * 256 + index];
workgroupBarrier(); // 确保所有线程都将数据加载到共享内存
var i = 1u;
loop {
if (i >= 256u) {
break;
}
if (index % (i * 2) == 0) {
shared_data[index] = shared_data[index] + shared_data[index + i];
}
i = i * 2;
workgroupBarrier(); // 确保每次归约都完成
}
if (index == 0) {
output.numbers[workgroup_id.x] = shared_data[0];
}
}解释：
struct Data定义了输入和输出数据的结构，包含一个浮点数数组。@group(0) @binding(0) var<storage, read> input: Data;定义了输入缓冲区，只读权限。@group(0) @binding(1) var<storage, write> output: Data;定义了输出缓冲区，只写权限。var<workgroup> shared_data : array<f32, 256>;定义了共享内存，大小为 256 个浮点数。@compute @workgroup_size(256)指定这是一个 Compute Shader，并且每个工作组的大小为 256。local_id.x是当前 Shader 实例在工作组内的 ID，范围是 0 到 255。workgroup_id.x是当前工作组的 ID。shared_data[index] = input.numbers[workgroup_id.x * 256 + index];将输入数组中的一部分数据加载到共享内存中。每个工作组处理 256 个数据。workgroupBarrier();确保所有 Shader 实例都完成了数据加载，才能进行下一步操作。- 后面的循环是归约求和的过程。它将共享内存中的数据两两相加，直到只剩下一个值，就是这个工作组的和。
if (index == 0)只有 ID 为 0 的 Shader 实例将结果写入输出缓冲区。output.numbers[workgroup_id.x] = shared_data[0];将每个工作组的和写入输出缓冲区。
2. JavaScript 代码 (index.js):
async function runComputeShader() {
const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
const shaderModule = device.createShaderModule({
code: await (await fetch('sum.wgsl')).text(),
});
const inputData = new Float32Array(1024);
for (let i = 0; i < inputData.length; ++i) {
inputData[i] = i + 1; // 初始化数据
}
const inputBuffer = device.createBuffer({
size: inputData.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
mappedAtCreation: true,
});
new Float32Array(inputBuffer.getMappedRange()).set(inputData);
inputBuffer.unmap();
const outputData = new Float32Array(inputData.length / 256); // 每个工作组输出一个值
const outputBuffer = device.createBuffer({
size: outputData.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
mappedAtCreation: false,
});
const stagingBuffer = device.createBuffer({
size: outputData.byteLength,
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});
const bindGroupLayout = device.createBindGroupLayout({
entries: [
{
binding: 0,
visibility: GPUShaderStage.COMPUTE,
buffer: {
type: 'read-only-storage',
},
},
{
binding: 1,
visibility: GPUShaderStage.COMPUTE,
buffer: {
type: 'storage',
},
},
],
});
const bindGroup = device.createBindGroup({
layout: bindGroupLayout,
entries: [
{
binding: 0,
resource: {
buffer: inputBuffer,
},
},
{
binding: 1,
resource: {
buffer: outputBuffer,
},
},
],
});
const computePipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [bindGroupLayout],
}),
compute: {
module: shaderModule,
entryPoint: 'main',
},
});
const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);
passEncoder.dispatchWorkgroups(inputData.length / 256); // 启动 4 个工作组
passEncoder.end();
commandEncoder.copyBufferToBuffer(
outputBuffer,
0,
stagingBuffer,
0,
outputData.byteLength
);
device.queue.submit([commandEncoder.finish()]);
await stagingBuffer.mapAsync(GPUMapMode.READ, 0, outputData.byteLength);
const copyArrayBuffer = stagingBuffer.getMappedRange(0, outputData.byteLength);
const result = new Float32Array(copyArrayBuffer);
stagingBuffer.unmap();
let sum = 0;
for (let i = 0; i < result.length; i++) {
sum += result[i];
}
console.log("Sum of array elements:", sum); // 打印最终结果
}
runComputeShader();解释：
- 这段代码创建了 WebGPU 设备和 Shader 模块。
- 创建了输入缓冲区
inputBuffer，并将数据写入。 - 创建了输出缓冲区
outputBuffer。 - 创建了 Bind Group Layout 和 Bind Group，将缓冲区绑定到 Shader。
- 创建了 Compute Pipeline。
- 创建 Command Encoder 和 Compute Pass，设置 Pipeline 和 Bind Group，并启动 Compute Shader。
- 将输出缓冲区的数据复制到 staging buffer。
- 映射 staging buffer，读取结果，并打印到控制台。
- 将每个工作组的结果加起来，得到最终的总和。
3. HTML (index.html):
<!DOCTYPE html>
<html>
<head>
<title>WebGPU Compute Shader Example</title>
</head>
<body>
<script src="index.js"></script>
</body>
</html>运行：
将这三个文件放在同一个目录下，然后在浏览器中打开 index.html，你将在控制台中看到数组的总和。
注意事项：
- 工作组大小：
workgroup_size必须与 Shader 代码中的shared_data大小一致。 - 数据对齐： 在 WebGPU 中，数据需要按照一定的规则对齐。如果数据没有对齐，可能会导致性能问题或者错误。
- 错误处理： 在实际开发中，需要添加错误处理机制，以便在出现错误时能够及时发现并解决问题。
总结：
通过这个例子，我们学习了如何使用 Compute Shaders、Shared Memory、Atomic Operations 和 Synchronization Barriers 来计算数组的总和。虽然这个例子很简单，但它涵盖了 WebGPU Compute Shader 的基本概念和用法。
一些更高级的用法和需要注意的点：
| 特性/概念 | 描述 | 示例或注意事项 |
|---|---|---|
| 复杂数据结构 | 在 WGSL 中，可以定义复杂的结构体，包含数组、矩阵、向量等，并在 Compute Shader 中使用。 | struct MyStruct { values: array<f32, 16>; matrix: mat4x4; }; 确保结构体成员的内存布局符合预期，可能需要使用 @location 和 @offset 来控制。 |
| 多级归约 | 当数组非常大时，一个工作组无法处理所有数据。可以采用多级归约的方式，先在每个工作组内进行归约，然后将每个工作组的结果再进行归约。 | 第一级归约在每个工作组内完成，结果写入全局内存；第二级归约启动另一个 Compute Shader，读取全局内存中的结果，并进行最终归约。 |
| 动态共享内存 | 某些情况下，你可能需要在运行时确定共享内存的大小。这在 WebGPU 中是不直接支持的，但可以通过一些技巧来实现类似的效果，例如使用全局内存模拟共享内存。 | 使用全局内存模拟共享内存的性能通常不如真正的共享内存。 |
| Atomic 指针 | WGSL 中，原子操作只能作用于 atomic<T> 类型。如果你想对存储缓冲区中的数据进行原子操作，你需要先将数据加载到 atomic<T> 类型的变量中，然后进行操作，最后再将结果写回存储缓冲区。 |
let atomic_value : atomic = atomicLoad(&buffer[index]); atomicAdd(&atomic_value, 1u); buffer[index] = atomicLoad(&atomic_value); |
| 调试 | WebGPU 的调试工具相对有限。可以使用 console.log 在 Shader 代码中打印信息，但这会影响性能。一些浏览器提供了 WebGPU 的调试器，可以用来查看 Shader 的执行过程和内存状态。 |
Chrome 的 GPU Internals (chrome://gpu) 提供了一些 WebGPU 的调试信息。 |
| 性能优化 | 优化 Compute Shader 的性能是一个复杂的问题。需要考虑的因素包括：工作组大小、数据局部性、内存访问模式、指令调度等。 | 尽量使用连续的内存访问，减少分支语句，利用向量化指令，避免数据竞争。 |
| WebGPU 限制 | WebGPU 有一些限制，例如最大工作组大小、最大缓冲区大小等。在编写 Compute Shader 时，需要注意这些限制。 | 最大工作组大小通常为 (256, 256, 64)，具体取决于 GPU。 |
| 错误处理 | 在 WebGPU 中，错误处理非常重要。需要捕获 Shader 编译错误、Pipeline 创建错误、缓冲区创建错误等，并进行适当的处理。 | 使用 device.createShaderModule 时，检查 GPUShaderModuleCompilationError 错误。 |
| 跨平台兼容性 | 不同的 GPU 和操作系统对 WebGPU 的支持程度可能不同。在开发 WebGPU 应用时，需要考虑跨平台兼容性问题。 | 使用 Dawn 等跨平台库可以简化跨平台开发。 |
最后，送给大家一句鸡汤：
学习 WebGPU Compute Shader 就像攀登高峰，虽然路途艰辛，但只要坚持不懈，终将看到美丽的风景！
今天的讲座就到这里，感谢大家的观看！如果大家觉得有用，记得点赞、收藏、转发哦！下次再见！