各位观众老爷，大家好！今天咱们来聊聊WebGPU里的Compute Shaders，这玩意儿能让你在浏览器里玩转通用GPU计算，搞搞数据并行，听起来是不是有点兴奋？别怕，咱们用人话把它掰开了揉碎了，保证你听完能上手写代码。
一、 啥是Compute Shader？它能干啥？
首先，得明白啥是Shader。简单来说，Shader就是一段运行在GPU上的小程序。以前我们主要用Vertex Shader和Fragment Shader来渲染3D图形，控制顶点位置和像素颜色。但是！GPU的强大计算能力不仅仅能用来画画，还能干很多别的事情。于是，Compute Shader就应运而生了。
Compute Shader 就像一个通用的计算引擎，你可以往里面扔数据，告诉它你要做什么计算，然后它会把结果给你吐出来。它最大的特点就是并行计算，可以将一个大的计算任务分解成很多小的任务，同时在GPU的多个核心上执行，大大提高计算效率。
那么，Compute Shader能干啥呢？
- 图像处理： 比如图像模糊、锐化、色彩校正等等，这些操作都可以并行处理每个像素，速度飞快。
- 物理模拟： 模拟粒子运动、流体流动、碰撞检测等等，需要大量的数值计算，Compute Shader可以胜任。
- 机器学习： 训练神经网络、进行推理等等，GPU的并行计算能力是机器学习的利器。
- 数据分析： 处理大规模数据，进行统计分析、数据挖掘等等，Compute Shader可以加速计算过程。
- 自定义算法： 只要你能把算法写成并行的形式，就可以用Compute Shader来实现，发挥你的想象力吧！
二、 WebGPU Compute Shader 的基本概念
在WebGPU里使用Compute Shader，你需要了解几个关键概念：
- Compute Pipeline： 就像流水线一样，定义了Compute Shader的执行流程。它包含了Shader Module（Shader代码）、Compute Stage（指定Shader的入口函数）等信息。
- Shader Module： 包含了用WGSL（WebGPU Shading Language）编写的Shader代码。WGSL是WebGPU专用的Shader语言，语法类似于Rust和GLSL。
- Bind Group Layout： 定义了Shader访问的资源（比如Buffer、Texture）的布局。
- Bind Group： 包含了实际的资源，并按照Bind Group Layout定义的布局绑定到Shader。
- Buffer： 存储数据的内存区域，可以用来作为Shader的输入和输出。
- Texture： 存储图像数据的内存区域，可以用来作为Shader的输入和输出。
- Command Encoder： 用于记录GPU指令，比如创建Buffer、Texture，创建Compute Pipeline，绑定Bind Group，执行Compute Shader等等。
- Compute Pass： 记录一次Compute Shader的执行过程，包含了Compute Pipeline、Bind Group、Dispatch信息等等。
- Dispatch： 告诉GPU要执行多少个Compute Shader实例。WebGPU使用三维的Grid来组织Compute Shader实例，每个实例都有一个唯一的坐标（x, y, z）。
三、 第一个WebGPU Compute Shader程序
咱们先来写一个简单的Compute Shader程序，把一个Buffer里的每个元素都乘以2。
-
HTML结构 (index.html):
<!DOCTYPE html> <html> <head> <title>WebGPU Compute Shader</title> </head> <body> <canvas id="canvas" width="512" height="512"></canvas> <script src="script.js"></script> </body> </html> -
JavaScript代码 (script.js):
async function initWebGPU() { // 1. 获取GPU设备 const adapter = await navigator.gpu.requestAdapter(); const device = await adapter.requestDevice(); // 2. 创建Buffer const bufferSize = 1024; const inputBuffer = device.createBuffer({ size: bufferSize * 4, // 4 bytes per float usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC, mappedAtCreation: true, }); const inputArray = new Float32Array(inputBuffer.getMappedRange()); for (let i = 0; i < bufferSize; i++) { inputArray[i] = i; // 初始化数据 } inputBuffer.unmap(); const outputBuffer = device.createBuffer({ size: bufferSize * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); const readBuffer = device.createBuffer({ size: bufferSize * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, }); // 3. 创建Shader Module const shaderModule = device.createShaderModule({ code: ` @group(0) @binding(0) var<storage, read_write> input : array<f32>; @group(0) @binding(1) var<storage, read_write> output : array<f32>; @compute @workgroup_size(64) fn main(@builtin(global_invocation_id) global_id : vec3<u32>) { let index = global_id.x; if (index < arrayLength(&input)) { output[index] = input[index] * 2.0; } } `, }); // 4. 创建Pipeline Layout const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [ device.createBindGroupLayout({ entries: [ { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage", }, }, { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage", }, }, ], }), ], }); // 5. 创建Compute Pipeline const computePipeline = device.createComputePipeline({ layout: pipelineLayout, compute: { module: shaderModule, entryPoint: "main", }, }); // 6. 创建Bind Group const bindGroup = device.createBindGroup({ layout: computePipeline.getBindGroupLayout(0), entries: [ { binding: 0, resource: { buffer: inputBuffer, }, }, { binding: 1, resource: { buffer: outputBuffer, }, }, ], }); // 7. 创建Command Encoder const commandEncoder = device.createCommandEncoder(); // 8. 创建Compute Pass const computePass = commandEncoder.beginComputePass(); computePass.setPipeline(computePipeline); computePass.setBindGroup(0, bindGroup); // 9. Dispatch Compute Shader const workgroupSize = 64; // 对应Shader里的 @workgroup_size(64) const workgroupCount = Math.ceil(bufferSize / workgroupSize); computePass.dispatchWorkgroups(workgroupCount, 1, 1); // x, y, z computePass.end(); commandEncoder.copyBufferToBuffer(outputBuffer, 0, readBuffer, 0, bufferSize * 4); // 10. 提交命令 const commandBuffer = commandEncoder.finish(); device.queue.submit([commandBuffer]); // 11. 读取结果 await readBuffer.mapAsync(GPUMapMode.READ); const outputArray = new Float32Array(readBuffer.getMappedRange()); // 12. 打印结果 for (let i = 0; i < 10; i++) { console.log(`input[${i}] = ${inputArray[i]}, output[${i}] = ${outputArray[i]}`); } readBuffer.unmap(); } initWebGPU(); -
代码解释：
- 步骤1： 获取GPU设备。这是WebGPU的入口，你需要先拿到一个Adapter，然后通过Adapter获取Device。Device代表了你的GPU设备。
- 步骤2： 创建Buffer。我们创建了两个Buffer：
inputBuffer用于存储输入数据，outputBuffer用于存储输出数据。readBuffer用于将GPU计算结果拷贝到CPU端读取。GPUBufferUsage.STORAGE：表示这个Buffer可以被Shader读写。GPUBufferUsage.COPY_SRC：表示这个Buffer可以作为拷贝源。GPUBufferUsage.COPY_DST：表示这个Buffer可以作为拷贝目标。GPUBufferUsage.MAP_READ：表示这个Buffer可以被CPU读取。mappedAtCreation: true：表示创建Buffer的时候就映射到CPU内存，方便初始化数据。inputBuffer.getMappedRange()：获取Buffer的映射内存区域，返回一个ArrayBuffer。inputBuffer.unmap()：取消Buffer的映射，让GPU可以使用这个Buffer。
- 步骤3： 创建Shader Module。我们用WGSL编写了一个简单的Shader，它把
inputBuffer里的每个元素都乘以2，然后写入outputBuffer。@group(0) @binding(0)：指定这个变量属于Bind Group 0，Binding 0。var<storage, read_write>：定义一个存储类型的变量，可以被Shader读写。array<f32>：定义一个浮点数数组。@compute @workgroup_size(64)：指定这个函数是一个Compute Shader，并且每个Workgroup包含64个Workitem。@builtin(global_invocation_id)：获取当前Workitem的全局ID。arrayLength(&input)：获取数组的长度。
- 步骤4： 创建Pipeline Layout。Pipeline Layout定义了Shader访问资源的布局。在这个例子里，我们定义了一个Bind Group Layout，它包含了两个Buffer：
input和output。GPUShaderStage.COMPUTE：表示这个资源可以被Compute Shader访问。buffer: { type: "storage" }：指定这个资源是一个存储类型的Buffer。
- 步骤5： 创建Compute Pipeline。Compute Pipeline定义了Compute Shader的执行流程。它包含了Shader Module、Compute Stage等信息。
layout: pipelineLayout：指定Pipeline Layout。compute: { module: shaderModule, entryPoint: "main" }：指定Shader Module和入口函数。
- 步骤6： 创建Bind Group。Bind Group包含了实际的资源，并按照Bind Group Layout定义的布局绑定到Shader。
layout: computePipeline.getBindGroupLayout(0)：指定Bind Group Layout。resource: { buffer: inputBuffer }：绑定inputBuffer到Binding 0。resource: { buffer: outputBuffer }：绑定outputBuffer到Binding 1。
- 步骤7： 创建Command Encoder。Command Encoder用于记录GPU指令。
- 步骤8： 创建Compute Pass。Compute Pass记录一次Compute Shader的执行过程。
computePass.setPipeline(computePipeline)：设置Compute Pipeline。computePass.setBindGroup(0, bindGroup)：设置Bind Group。
- 步骤9： Dispatch Compute Shader。Dispatch告诉GPU要执行多少个Compute Shader实例。
workgroupSize = 64：每个Workgroup包含64个Workitem。workgroupCount = Math.ceil(bufferSize / workgroupSize)：计算需要的Workgroup数量。computePass.dispatchWorkgroups(workgroupCount, 1, 1)：执行Compute Shader。
- 步骤10： 提交命令。把Command Buffer提交到GPU队列，让GPU执行。
- 步骤11： 读取结果。把
outputBuffer里的数据拷贝到readBuffer，然后映射到CPU内存，读取结果。 - 步骤12： 打印结果。
四、 WGSL 语法扫盲
WGSL是WebGPU专用的Shader语言，如果你熟悉GLSL或者Rust，那么学习WGSL会很容易。这里简单介绍一些常用的WGSL语法：
-
变量声明：
var<storage, read_write> my_buffer : array<f32>; // 声明一个存储类型的浮点数数组 let my_constant = 10; // 声明一个常量 var my_variable : i32 = 0; // 声明一个可变的整数 -
数据类型：
f32：32位浮点数i32：32位整数u32：32位无符号整数vec2<f32>：2维浮点数向量vec3<f32>：3维浮点数向量vec4<f32>：4维浮点数向量mat4x4<f32>：4×4浮点数矩阵sampler：采样器texture_2d<f32>：2D纹理
-
函数：
fn my_function(x : f32, y : f32) -> f32 { return x + y; } -
控制流：
if (x > 0) { // do something } else { // do something else } for (var i = 0; i < 10; i = i + 1) { // do something } while (x < 10) { // do something x = x + 1; } -
内置函数：
WGSL提供了很多内置函数，比如：
length(v : vec3<f32>) -> f32：计算向量的长度normalize(v : vec3<f32>) -> vec3<f32>：归一化向量dot(v1 : vec3<f32>, v2 : vec3<f32>) -> f32：计算向量的点积cross(v1 : vec3<f32>, v2 : vec3<f32>) -> vec3<f32>：计算向量的叉积textureSample(t : texture_2d<f32>, s : sampler, uv : vec2<f32>) -> vec4<f32>：采样纹理
-
Attributes:
@group(x): 指定变量属于哪个Bind Group@binding(x): 指定变量属于Bind Group中的哪个Binding@compute: 指定函数为Compute Shader@vertex: 指定函数为Vertex Shader@fragment: 指定函数为Fragment Shader@workgroup_size(x, y, z): 指定Workgroup的大小@builtin(xxx): 访问内置变量，例如global_invocation_id
五、 数据并行：Workgroup 和 Workitem
Compute Shader的核心是数据并行。WebGPU使用Workgroup和Workitem来组织并行计算。
- Workitem： 一个Compute Shader实例，它执行Shader代码的最小单元。
- Workgroup： 一组Workitem，它们共享本地内存，可以进行同步。Workgroup的大小由
@workgroup_size(x, y, z)指定。 - Global Invocation ID： 每个Workitem都有一个唯一的全局ID，它是一个三维向量（x, y, z）。你可以通过
@builtin(global_invocation_id)来获取当前Workitem的全局ID。 - Local Invocation ID： 每个Workitem在所属Workgroup内也有一个唯一的本地ID，它也是一个三维向量（x, y, z）。
- Workgroup ID: 每个Workgroup也有一个唯一ID，通过
@builtin(workgroup_id)可以获取.
假设我们有一个1024个元素的数组，我们想用Compute Shader并行处理每个元素。我们可以把数组分成16个Workgroup，每个Workgroup包含64个Workitem。那么：
| 概念 | 描述 | 例子 |
|---|---|---|
| 数据总大小 | 需要处理的数据总量 | 1024 个元素 |
| Workgroup Size | 每个Workgroup包含的Workitem数量 | @workgroup_size(64) 表示每个Workgroup有64个Workitem |
| Workgroup Count | 需要的Workgroup数量 | Math.ceil(1024 / 64) = 16 表示需要16个Workgroup |
| Global ID | 每个Workitem在全局中的唯一ID | 如果global_id.x = 512，那么这个Workitem处理数组的第512个元素 |
| Local ID | 每个Workitem在所属Workgroup内的唯一ID | 如果local_id.x = 10，那么这个Workitem是所属Workgroup的第10个Workitem |
在Shader代码里，你可以使用global_id来访问全局数据，使用local_id来访问Workgroup内的共享数据。
六、 高级技巧
-
使用本地内存 (Workgroup Shared Memory):
Workgroup内的Workitem可以共享本地内存。你可以使用
var<workgroup>来声明本地内存变量。本地内存的访问速度比全局内存快很多，可以用来优化计算性能。@compute @workgroup_size(64) fn main(@builtin(global_invocation_id) global_id : vec3<u32>, @builtin(local_invocation_id) local_id : vec3<u32>) { var<workgroup> shared_data : array<f32, 64>; // 声明一个本地内存数组 shared_data[local_id.x] = input[global_id.x]; // 把全局数据拷贝到本地内存 workgroupBarrier(); // 同步Workgroup内的所有Workitem output[global_id.x] = shared_data[63 - local_id.x]; // 从本地内存读取数据 }workgroupBarrier()函数用于同步Workgroup内的所有Workitem。它可以保证在所有Workitem都执行完shared_data[local_id.x] = input[global_id.x]之后，才能执行output[global_id.x] = shared_data[63 - local_id.x]。 -
使用纹理 (Texture):
除了Buffer，你还可以使用纹理作为Compute Shader的输入和输出。纹理可以存储图像数据，也可以存储其他类型的数据。
@group(0) @binding(0) var texture : texture_2d<f32>; @group(0) @binding(1) var sampler : sampler; @group(0) @binding(2) var<storage, read_write> output : texture_2d<f32>; @compute @workgroup_size(8, 8) fn main(@builtin(global_invocation_id) global_id : vec3<u32>) { let uv = vec2<f32>(f32(global_id.x) / 512.0, f32(global_id.y) / 512.0); let color = textureSample(texture, sampler, uv); textureStore(output, vec2<i32>(global_id.xy), color * 2.0); }在这个例子里，我们从
texture中采样颜色，然后把颜色乘以2，写入output纹理。 -
优化性能：
- 减少内存访问： 尽量使用本地内存，减少对全局内存的访问。
- 合并内存访问： 尽量连续访问内存，避免随机访问。
- 调整Workgroup Size： 根据GPU的架构，选择合适的Workgroup Size。一般来说，Workgroup Size是32或者64的倍数，可以更好地利用GPU的资源。
- 避免分支： 尽量避免在Shader代码中使用分支语句（if/else），因为分支会导致Workitem执行不同的代码路径，降低并行效率。
- 使用内置函数： 尽量使用WGSL提供的内置函数，因为这些函数经过了优化，性能更好。
七、 总结
WebGPU Compute Shader为我们打开了一扇通往GPU通用计算的大门。 掌握这些知识，你就可以在浏览器里实现各种高性能的计算应用。希望今天的讲座能让你对WebGPU Compute Shader有一个初步的了解。 记住，实践是检验真理的唯一标准，多写代码，多做实验，你才能真正掌握这门技术。
好了，今天的分享就到这里，谢谢大家！下次再见！