如何实现一个基于WebGPU的通用计算程序？

实现基于WebGPU的通用计算需先获取设备，再创建缓冲区上传数据，编写WGSL计算着色器定义并行逻辑，通过管线和绑定组关联资源，最后提交命令执行并读回结果。

实现一个基于WebGPU的通用计算程序，核心在于利用其计算着色器（compute shader）在GPU上并行执行数据密集型任务。整个流程包括请求设备、创建缓冲区、编写计算逻辑、调度执行等步骤。下面分关键环节说明如何构建这样一个程序。

初始化WebGPU环境

要运行WebGPU代码，首先需要获取WebGPU上下文并请求适配器和设备。

if (!navigator.gpu) {
  throw new Error('WebGPU not supported');
}

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();

设备是后续所有操作的基础，用于创建缓冲区、管线和命令编码器。

准备数据与缓冲区管理

通用计算通常处理大量数据，需将数据上传至GPU内存。这通过创建GPUBuffer完成。

const input1 = new Float32Array([1, 2, 3, 4]);
const input2 = new Float32Array([5, 6, 7, 8]);
const output = new Float32Array(input1.length);

const buffer1 = device.createBuffer({
  size: input1.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
const buffer2 = device.createBuffer({
  size: input2.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
});
const outputBuffer = device.createBuffer({
  size: output.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
});

使用COPY_DST表示可从CPU写入，COPY_SRC表示结果可复制回CPU。

编写计算着色器并创建管线

计算着色器用WGSL（WebGPU Shading Language）编写，定义每个线程的运算逻辑。

例如，实现向量加法：

算家云

高效、便捷的人工智能算力服务平台

37

查看详情

const computeShaderCode = `@group(0) @binding(0) var inputData1 : array<f32>;
@group(0) @binding(1) var inputData2 : array<f32>;
@group(0) @binding(2) var outputData : array<f32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
  let i = gid.x;
  if (i < arrayLength(&inputData1)) {
    outputData[i] = inputData1[i] + inputData2[i];
  }
}`;

const shaderModule = device.createShaderModule({ code: computeShaderCode });
const pipeline = device.createComputePipeline({
  layout: 'auto',
  compute: {
    module: shaderModule,
    entryPoint: 'main'
  }
});

@workgroup_size(64) 表示每个工作组包含64个线程，合理设置能提升性能。

绑定资源并提交计算任务

需要将缓冲区绑定到管线，并通过命令编码器调度执行。

const bindGroup = device.createBindGroup({
  layout: pipeline.getBindGroupLayout(0),
  entries: [
    { binding: 0, resource: { buffer: buffer1 } },
    { binding: 1, resource: { buffer: buffer2 } },
    { binding: 2, resource: { buffer: outputBuffer } }
  ]
});

const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(pipeline);
passEncoder.setBindGroup(0, bindGroup);
passEncoder.dispatchWorkgroups(Math.ceil(input1.length / 64));
passEncoder.end();
// 将输入数据写入GPU缓冲区
device.queue.writeBuffer(buffer1, 0, input1);
device.queue.writeBuffer(buffer2, 0, input2);
// 提交命令
device.queue.submit([commandEncoder.finish()]);

dispatchWorkgroups参数决定启动多少个工作组，应根据数据总量和工作组大小计算。

读取计算结果

由于GPU操作异步，需通过映射缓冲区将结果传回JS。

const readBackBuffer = device.createBuffer({
  size: output.byteLength,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
});
const copyEncoder = device.createCommandEncoder();
copyEncoder.copyBufferToBuffer(outputBuffer, 0, readBackBuffer, 0, output.byteLength);
device.queue.submit([copyEncoder.finish()]);

await readBackBuffer.mapAsync(GPUMapMode.READ);
const result = new Float32Array(readBackBuffer.getMappedRange());
console.log(result); // 输出 [6, 8, 10, 12]

mapAsync确保数据已就绪后再读取，避免竞争条件。

基本上就这些。掌握缓冲区管理、着色器编写和命令调度，就能实现各类并行计算，如图像处理、物理模拟或机器学习推理。关键是理解数据流和同步机制，不复杂但容易忽略细节。

以上就是如何实现一个基于WebGPU的通用计算程序？的详细内容，更多请关注php中文网其它相关文章！