JS `WebGPU Compute Shaders` 高级：通用 GPU 计算与数据并行

各位观众老爷，晚上好！今天咱们来聊聊 WebGPU 里的“重头戏”——Compute Shaders。别怕，虽然听起来高大上，但其实就是让你用显卡（GPU）来算各种各样的东西，不再局限于画三角形和贴图了。简单来说，就是让你的浏览器拥有了“超能力”！

WebGPU Compute Shaders：通用 GPU 计算与数据并行

咱们先来打个比方。想象一下，你有个特别复杂的数学题，让你一个个数着算，得算到猴年马月。但如果你有一大堆小弟（GPU核心），每个人帮你算一部分，是不是就快多了？Compute Shaders 就是让你把这些小弟组织起来，帮你解决问题的“指挥棒”。

1. 什么是 Compute Shaders？

Compute Shaders 是一种特殊的着色器程序，它不参与图形渲染管线，而是直接在 GPU 上执行通用计算任务。它的主要作用就是利用 GPU 的并行计算能力，加速各种算法和数据处理过程。

通用 GPU 计算 (GPGPU): 指的就是用 GPU 来做图形渲染以外的计算任务。
数据并行: 指的是把一个大的数据处理任务分成很多小的子任务，然后让 GPU 上的很多核心同时处理这些子任务。

2. 为什么要用 Compute Shaders？

速度快: GPU 拥有大量的计算核心，非常适合并行计算。
效率高: 对于某些类型的计算任务，GPU 的效率远高于 CPU。
解放 CPU: 把计算密集型的任务交给 GPU，可以减轻 CPU 的负担，让你的网页更加流畅。

3. Compute Shaders 的基本概念

Workgroup: 你可以把 Workgroup 想象成一个“工作组”，里面包含多个 “Workitem”。
Workitem: Workitem 是最小的计算单元，每个 Workitem 都会执行相同的着色器代码，但处理不同的数据。
Dispatch: Dispatch 指的是启动 Compute Shaders 的过程，你需要指定启动多少个 Workgroup。

4. WebGPU Compute Shaders 的代码结构

一个基本的 WebGPU Compute Shader 程序通常包含以下几个部分：

WGSL 代码: 用 WGSL (WebGPU Shading Language) 编写的着色器代码，定义了 Workitem 的计算逻辑。
Pipeline: Compute Pipeline 定义了如何运行 Compute Shader，包括着色器模块、布局等。
Bind Group: Bind Group 用于将数据（例如 Buffer、Texture）绑定到着色器，让着色器可以访问这些数据。
Command Encoder: Command Encoder 用于记录执行 Compute Shader 的命令。

5. 一个简单的 Compute Shader 例子

咱们先来一个最简单的例子：把一个数组里的每个元素都加 1。

5.1 WGSL 代码 (shader.wgsl)

@group(0) @binding(0) var<storage, read_write> data: array<i32>;

@compute @workgroup_size(64) // 定义 Workgroup 的大小
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let index = global_id.x;
  data[index] = data[index] + 1;
}

@group(0) @binding(0): 指定了 Buffer 绑定到哪个 Bind Group 和 Binding Point。
var<storage, read_write> data: array<i32>: 声明了一个存储 Buffer，可以读写。
@compute: 表示这是一个 Compute Shader。
@workgroup_size(64): 定义了 Workgroup 的大小，这里是 64，表示每个 Workgroup 包含 64 个 Workitem。
@builtin(global_invocation_id) global_id: vec3<u32>: 获取当前 Workitem 的全局 ID。
data[index] = data[index] + 1;: 把数组中对应位置的元素加 1。

5.2 JavaScript 代码 (main.js)

async function runComputeShader() {
  // 1. 创建 WebGPU 设备
  const adapter = await navigator.gpu.requestAdapter();
  const device = await adapter.requestDevice();

  // 2. 创建 Buffer
  const arrayLength = 256;
  const dataArray = new Int32Array(arrayLength);
  for (let i = 0; i < arrayLength; ++i) {
    dataArray[i] = i;
  }

  const dataBuffer = device.createBuffer({
    size: dataArray.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    mappedAtCreation: true, // 初始化时映射，可以写入数据
  });
  new Int32Array(dataBuffer.getMappedRange()).set(dataArray);
  dataBuffer.unmap();

  // 3. 创建 Shader Module
  const shaderModule = device.createShaderModule({
    code: `
      @group(0) @binding(0) var<storage, read_write> data: array<i32>;

      @compute @workgroup_size(64)
      fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
        let index = global_id.x;
        data[index] = data[index] + 1;
      }
    `, // 直接内联 WGSL 代码，也可以从文件读取
  });

  // 4. 创建 Pipeline
  const computePipeline = device.createComputePipeline({
    layout: 'auto', // 让 WebGPU 自动推断布局
    compute: {
      module: shaderModule,
      entryPoint: 'main',
    },
  });

  // 5. 创建 Bind Group
  const bindGroup = device.createBindGroup({
    layout: computePipeline.getBindGroupLayout(0),
    entries: [
      {
        binding: 0,
        resource: {
          buffer: dataBuffer,
        },
      },
    ],
  });

  // 6. 创建 Command Encoder
  const commandEncoder = device.createCommandEncoder();

  // 7. 启动 Compute Shader
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup);
  passEncoder.dispatchWorkgroups(arrayLength / 64); // 需要启动多少个 Workgroup
  passEncoder.end();

  // 8. 复制结果到 CPU
  const readBuffer = device.createBuffer({
    size: dataArray.byteLength,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
  });
  commandEncoder.copyBufferToBuffer(dataBuffer, 0, readBuffer, 0, dataArray.byteLength);

  // 9. 提交命令
  const commandBuffer = commandEncoder.finish();
  device.queue.submit([commandBuffer]);

  // 10. 读取结果
  await readBuffer.mapAsync(GPUMapMode.READ);
  const resultArray = new Int32Array(readBuffer.getMappedRange());
  console.log('Result:', resultArray);
  readBuffer.unmap();
}

runComputeShader();

这个例子虽然简单，但涵盖了 Compute Shader 的基本流程：

创建 WebGPU 设备: 获取 GPU 设备。
创建 Buffer: 创建用于存储数据的 Buffer，并初始化数据。
创建 Shader Module: 创建 Shader Module，包含 WGSL 代码。
创建 Pipeline: 创建 Compute Pipeline，指定 Shader Module 和入口点。
创建 Bind Group: 创建 Bind Group，将 Buffer 绑定到 Shader。
创建 Command Encoder: 创建 Command Encoder，用于记录命令。
启动 Compute Shader: 使用 dispatchWorkgroups 启动 Compute Shader，指定启动多少个 Workgroup。
复制结果到 CPU: 创建一个用于读取的buffer，并将结果复制到这个buffer。
提交命令: 提交命令，让 GPU 执行计算。
读取结果: 读取结果，并输出到控制台。

6. 深入理解 Workgroup 和 Workitem

Workgroup 和 Workitem 是 Compute Shader 中最重要的概念。理解它们之间的关系，才能更好地利用 GPU 的并行计算能力。

Workgroup 的大小: @workgroup_size(x, y, z) 定义了 Workgroup 的大小，其中 x, y, z 分别表示 Workgroup 在 X, Y, Z 轴上的 Workitem 数量。

Global ID, Local ID 和 Workgroup ID:

Global ID: 每个 Workitem 都有一个唯一的全局 ID，可以通过 @builtin(global_invocation_id) 获取。
Local ID: 每个 Workitem 在其所属的 Workgroup 中也有一个唯一的局部 ID，可以通过 @builtin(local_invocation_id) 获取。
Workgroup ID: 每个 Workgroup 也有一个唯一的 ID，可以通过 @builtin(workgroup_id) 获取。

可以用下面的表格来更清晰的描述：

ID 类型	描述	获取方式
Global ID	整个 Dispatch 中 Workitem 的唯一 ID。	`@builtin(global_invocation_id)`
Local ID	Workgroup 内部 Workitem 的唯一 ID。	`@builtin(local_invocation_id)`
Workgroup ID	Workgroup 的唯一 ID。	`@builtin(workgroup_id)`
Num Workgroups	Dispatch 中 Workgroup 的数量 (在 JS 中计算)。	N/A (在 JS 中计算)
Workgroup Size	每个 Workgroup 内部 Workitem 的数量 (在 WGSL 中定义，例如 `@workgroup_size(64)`)。	`@workgroup_size` annotation (在 WGSL 中定义)

这些 ID 可以帮助你确定 Workitem 应该处理哪些数据。例如，你可以使用 Global ID 来访问数组中的元素，使用 Local ID 来进行 Workgroup 内部的同步。

7. Workgroup 共享内存

Workgroup 共享内存 (Shared Memory) 是一种特殊的内存区域，只能被同一个 Workgroup 中的 Workitem 访问。它可以用于 Workgroup 内部的数据共享和同步。

7.1 声明共享内存

在 WGSL 代码中，可以使用 var<workgroup> 声明共享内存。

var<workgroup> shared_data: array<f32, 64>; // 声明一个大小为 64 的浮点数数组

7.2 使用共享内存

Workgroup 中的 Workitem 可以通过 Local ID 来访问共享内存。

let local_id = local_invocation_id.x;
shared_data[local_id] = ...; // 写入共享内存
... = shared_data[local_id]; // 读取共享内存

7.3 Workgroup Barrier

由于 Workgroup 中的 Workitem 是并行执行的，因此在访问共享内存时需要进行同步。可以使用 workgroupBarrier() 函数来确保 Workgroup 中的所有 Workitem 都执行到同一个位置。

workgroupBarrier(); // 等待 Workgroup 中的所有 Workitem 都执行到这里

8. 一个使用 Workgroup 共享内存的例子

咱们来一个稍微复杂一点的例子：计算一个数组的局部和。每个 Workgroup 计算一部分数组的和，然后将结果存储到另一个数组中。

8.1 WGSL 代码 (shader.wgsl)

@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, write> output: array<f32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>,
        @builtin(local_invocation_id) local_id: vec3<u32>) {
  let index = global_id.x;
  var sum: f32 = 0.0;
  for (var i: u32 = 0; i < 16; i = i + 1) {
    if (index + i < arrayLength) {
      sum = sum + input[index + i];
    }
  }
  output[index] = sum;
}

8.2 JavaScript 代码 (main.js)

async function runComputeShader() {
    // 1. 创建 WebGPU 设备
    const adapter = await navigator.gpu.requestAdapter();
    const device = await adapter.requestDevice();

    // 2. 创建输入 Buffer
    const arrayLength = 256;
    const inputArray = new Float32Array(arrayLength);
    for (let i = 0; i < arrayLength; ++i) {
        inputArray[i] = Math.random();
    }

    const inputBuffer = device.createBuffer({
        size: inputArray.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
        mappedAtCreation: true,
    });
    new Float32Array(inputBuffer.getMappedRange()).set(inputArray);
    inputBuffer.unmap();

    // 3. 创建输出 Buffer
    const outputArray = new Float32Array(arrayLength);
    const outputBuffer = device.createBuffer({
        size: outputArray.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    });

    // 4. 创建 Shader Module
    const shaderModule = device.createShaderModule({
        code: `
            const arrayLength: u32 = 256; // 声明常量

            @group(0) @binding(0) var<storage, read> input: array<f32>;
            @group(0) @binding(1) var<storage, write> output: array<f32>;

            @compute @workgroup_size(64)
            fn main(@builtin(global_invocation_id) global_id: vec3<u32>,
                    @builtin(local_invocation_id) local_id: vec3<u32>) {
                let index = global_id.x;
                var sum: f32 = 0.0;
                for (var i: u32 = 0; i < 16; i = i + 1) {
                    if (index + i < arrayLength) {
                        sum = sum + input[index + i];
                    }
                }
                output[index] = sum;
            }
        `,
    });

    // 5. 创建 Pipeline
    const computePipeline = device.createComputePipeline({
        layout: 'auto',
        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 Shader
    const passEncoder = commandEncoder.beginComputePass();
    passEncoder.setPipeline(computePipeline);
    passEncoder.setBindGroup(0, bindGroup);
    passEncoder.dispatchWorkgroups(arrayLength / 64);
    passEncoder.end();

    // 9. 复制结果到 CPU
    const readBuffer = device.createBuffer({
        size: outputArray.byteLength,
        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
    commandEncoder.copyBufferToBuffer(outputBuffer, 0, readBuffer, 0, outputArray.byteLength);

    // 10. 提交命令
    const commandBuffer = commandEncoder.finish();
    device.queue.submit([commandBuffer]);

    // 11. 读取结果
    await readBuffer.mapAsync(GPUMapMode.READ);
    const resultArray = new Float32Array(readBuffer.getMappedRange());
    console.log('Result:', resultArray);
    readBuffer.unmap();
}

runComputeShader();

9. 实际应用场景

Compute Shaders 可以用于各种各样的计算任务，例如：

图像处理: 图像滤波、边缘检测、色彩校正等。
物理模拟: 粒子系统、流体模拟、碰撞检测等。
机器学习: 神经网络训练、模型推理等。
密码学: 哈希计算、加密解密等。
数据分析: 数据排序、数据聚合等。

10. 性能优化技巧

合理选择 Workgroup 大小: Workgroup 大小会影响 GPU 的利用率。通常来说，选择 64 或 256 是一个不错的选择。
减少内存访问: 尽量减少对全局内存的访问，多使用 Workgroup 共享内存。
避免分支: 尽量避免在着色器代码中使用分支语句，因为分支会导致 GPU 的执行效率降低。
使用矢量化操作: 尽量使用矢量化操作，例如 vec4，可以一次处理多个数据。
数据对齐: 确保数据是对齐的，可以提高内存访问效率。

11. 总结

WebGPU Compute Shaders 为我们提供了一种强大的工具，可以利用 GPU 的并行计算能力来加速各种计算任务。虽然学习曲线可能有点陡峭，但是一旦掌握了，你就可以让你的网页拥有“超能力”，轻松应对各种复杂的计算需求。希望今天的讲解能帮助你入门 WebGPU Compute Shaders，开启你的 GPU 计算之旅！

好了，今天的讲座就到这里，感谢各位的观看！下次有机会再和大家一起探讨 WebGPU 的其他技术。溜了溜了~

发表回复 取消回复

发表回复取消回复