如何复用预创建的WGSL计算管线处理不同输入？

要解决每次重建管线导致的性能问题，核心就是把管线创建、缓冲区初始化这些高成本操作从循环里抽出来，只在启动时执行一次，后续计算仅更新输入数据并重新提交命令即可。

具体实现步骤 #

1. 初始化阶段（仅执行一次）

   • 创建GPU适配器、设备、着色器模块、计算管线。
   • 创建所有需要的缓冲区（包括固定数据缓冲区、可变输入缓冲区、结果缓冲区）。
   • 创建绑定组（BindGroup），绑定所有缓冲区到管线的绑定点。
   • 写入固定不变的数据（比如你的sample和table）到对应缓冲区，无需重复写入。

2. 计算阶段（可重复执行多次）

   • 将新的输入数据写入可变输入缓冲区。
   • 新建命令编码器，重新录制计算命令（设置管线、绑定组、调度工作组、复制结果）。
   • 提交命令队列，等待并读取计算结果。

优化后的代码示例 #

// ---------------------- 初始化阶段（只执行一次） ----------------------
let adapter = await navigator.gpu.requestAdapter();
let device = await adapter.requestDevice();
const LEN = 100; // 根据实际场景设置工作组数量

// 创建WGSL着色器模块
const shaderModule = device.createShaderModule({
  code: `
    @group(0) @binding(0) var<storage, read_write> sample: array<u32, 720>;
    @group(0) @binding(1) var<storage, read_write> table: array<array<u32, 720>>;
    @group(0) @binding(2) var<storage, read_write> result: array<u32>;

    @compute @workgroup_size(1,1,1) fn computeThis (@builtin(global_invocation_id) id: vec3<u32>) {
      var diff : u32 = 0;
      for (var i : u32 = 0; i < 720; i++) {
        diff += (table[id.x][i] - sample[i])*(table[id.x][i] - sample[i]);
      }
      result[id.x] = diff;
    }
  `,
});

// 创建计算管线（仅创建一次）
const computePipeline = device.createComputePipeline({
  layout: 'auto',
  compute: { module: shaderModule }
});

// 创建缓冲区：固定数据缓冲区 + 可变输入/输出缓冲区
const sampleBuffer = device.createBuffer({
  size: sample.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
const tableBuffer = device.createBuffer({
  size: table.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
const inputBuffer = device.createBuffer({
  size: input.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
const resultBuffer = device.createBuffer({
  size: input.byteLength,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
});

// 创建绑定组（仅创建一次，缓冲区实例不变即可复用）
const bindGroup = device.createBindGroup({
  layout: computePipeline.getBindGroupLayout(0),
  entries: [
    { binding: 0, resource: { buffer: sampleBuffer }},
    { binding: 1, resource: { buffer: tableBuffer }},
    { binding: 2, resource: { buffer: inputBuffer }}
  ]
});

// 写入固定数据（仅执行一次）
device.queue.writeBuffer(sampleBuffer, 0, sample);
device.queue.writeBuffer(tableBuffer, 0, table);

// ---------------------- 重复执行的计算函数 ----------------------
async function runComputeWithNewInput(newInput) {
  // 更新新的输入数据到缓冲区
  device.queue.writeBuffer(inputBuffer, 0, newInput);

  // 录制新的计算命令
  const encoder = device.createCommandEncoder();
  const computePass = encoder.beginComputePass();
  computePass.setPipeline(computePipeline);
  computePass.setBindGroup(0, bindGroup);
  computePass.dispatchWorkgroups(LEN, 1, 1);
  computePass.end();

  // 将计算结果复制到可映射的结果缓冲区
  encoder.copyBufferToBuffer(inputBuffer, 0, resultBuffer, 0, resultBuffer.size);

  // 提交命令队列
  device.queue.submit([encoder.finish()]);

  // 读取并返回结果
  await resultBuffer.mapAsync(GPUMapMode.READ);
  const result = new Uint32Array(resultBuffer.getMappedRange().slice());
  resultBuffer.unmap();

  return result;
}

关键注意事项 #

• 管线与绑定组复用：ComputePipeline和BindGroup仅需创建一次，它们依赖的是缓冲区的对象实例而非内容，只要缓冲区实例不更换，绑定组就能一直复用。
• 缓冲区复用：所有缓冲区只初始化一次，每次计算仅通过queue.writeBuffer更新可变输入的内容，避免重复创建缓冲区的开销。
• 命令缓冲区重建：命令缓冲区是一次性的，每次计算都需要重新创建并录制命令，确保基于最新的输入数据执行计算。
• 固定数据一次性写入：如果sample和table是静态数据，初始化时写入一次即可，无需每次计算重复操作，减少内存拷贝成本。

内容的提问来源于stack exchange，提问作者James