WebGPU入门实战:在浏览器中实现高性能GPU并行计算

WebGPU是新一代Web图形与计算API,相比WebGL提供了更底层的GPU访问能力和Compute Shader支持,使浏览器中进行大规模并行计算成为可能。在Web性能优化领域,WebGPU为图像处理、物理模拟、机器学习推理等场景提供了远超JavaScript单线程的性能上限。本文从API初始化到Compute Shader实战,完整演示WebGPU的开发流程。

WebGPU环境检测与设备初始化

WebGPU在Chrome 113+、Edge 113+中已默认启用。使用前需要检测浏览器支持并请求GPU设备。

// 检测WebGPU支持
async function initWebGPU() {
  if (!navigator.gpu) {
    throw new Error('当前浏览器不支持WebGPU,请使用Chrome 113+或Edge 113+');
  }

  // 请求GPU适配器
  const adapter = await navigator.gpu.requestAdapter({
    powerPreference: 'high-performance',
  });
  
  if (!adapter) {
    throw new Error('未找到可用的GPU适配器');
  }

  // 查看适配器信息
  const adapterInfo = await adapter.requestAdapterInfo();
  console.log(`GPU厂商: ${adapterInfo.vendor}`);
  console.log(`GPU架构: ${adapterInfo.architecture}`);
  console.log(`设备描述: ${adapterInfo.description}`);

  // 请求GPU设备
  const device = await adapter.requestDevice({
    requiredFeatures: ['timestamp-query'],
    requiredLimits: {
      maxBufferSize: 1 << 30,        // 1GB
      maxStorageBufferBindingSize: 1 << 28, // 256MB
      maxComputeWorkgroupsPerDimension: 65535,
    },
  });

  return { device, adapter };
}

const { device } = await initWebGPU();

Compute Shader基础:WGSL着色器编写

WebGPU使用WGSL(WebGPU Shading Language)编写着色器。以下是一个向量加法的Compute Shader示例:

// 向量加法 Compute Shader (WGSL)
const shaderCode = /* wgsl */ `
  @group(0) @binding(0) var<storage, read> inputA: array<f32>;
  @group(0) @binding(1) var<storage, read> inputB: array<f32>;
  @group(0) @binding(2) var<storage, read_write> output: array<f32>;

  @compute @workgroup_size(64)
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let index = gid.x;
    if (index >= arrayLength(&output)) {
      return;
    }
    output[index] = inputA[index] + inputB[index];
  }
`;

// 创建着色器模块
const shaderModule = device.createShaderModule({
  code: shaderCode,
});

// 创建计算管线
const computePipeline = device.createComputePipeline({
  layout: 'auto',
  compute: {
    module: shaderModule,
    entryPoint: 'main',
  },
});

@workgroup_size(64)定义每个工作组的线程数为64。GPU以工作组为单位调度计算,总线程数 = workgroup_size × workgroup_count。对于包含100万元素的数组,需要 ceil(1000000 / 64) = 15625 个工作组。

缓冲区创建与数据传输

WebGPU的数据传输通过GPUBuffer完成。数据流向为:CPU写入Staging Buffer → 复制到GPU专用Buffer → GPU计算 → 结果复制回Staging Buffer → CPU读取。

// 数据规模
const ARRAY_SIZE = 1_000_000;
const BUFFER_SIZE = ARRAY_SIZE * 4; // f32 = 4 bytes

// 创建输入数据
const dataA = new Float32Array(ARRAY_SIZE);
const dataB = new Float32Array(ARRAY_SIZE);
for (let i = 0; i < ARRAY_SIZE; i++) {
  dataA[i] = Math.random() * 100;
  dataB[i] = Math.random() * 100;
}

// 创建GPU缓冲区
const bufferA = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});

const bufferB = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});

const bufferOutput = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
});

// Staging Buffer用于读取结果
const stagingBuffer = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});

// 写入数据到GPU
device.queue.writeBuffer(bufferA, 0, dataA);
device.queue.writeBuffer(bufferB, 0, dataB);

// 创建绑定组
const bindGroup = device.createBindGroup({
  layout: computePipeline.getBindGroupLayout(0),
  entries: [
    { binding: 0, resource: { buffer: bufferA } },
    { binding: 1, resource: { buffer: bufferB } },
    { binding: 2, resource: { buffer: bufferOutput } },
  ],
});

提交计算任务与结果读取

// 创建命令编码器
const encoder = device.createCommandEncoder();

// 开始计算通道
const computePass = encoder.beginComputePass();
computePass.setPipeline(computePipeline);
computePass.setBindGroup(0, bindGroup);

// 分发计算任务
const workgroupCount = Math.ceil(ARRAY_SIZE / 64);
computePass.dispatchWorkgroups(workgroupCount);
computePass.end();

// 将结果复制到Staging Buffer
encoder.copyBufferToBuffer(
  bufferOutput, 0,
  stagingBuffer, 0,
  BUFFER_SIZE
);

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

// 异步读取结果
await stagingBuffer.mapAsync(GPUMapMode.READ);
const result = new Float32Array(stagingBuffer.getMappedRange());

// 验证结果
let correct = 0;
for (let i = 0; i < ARRAY_SIZE; i++) {
  if (Math.abs(result[i] - (dataA[i] + dataB[i])) < 0.001) {
    correct++;
  }
}
console.log(`正确率: ${(correct / ARRAY_SIZE * 100).toFixed(2)}%`);

// 解除映射
stagingBuffer.unmap();

实战:GPU加速图像模糊处理

高斯模糊是图像处理中的常见操作,使用Compute Shader可以大幅加速。以下是一个可配置半径的高斯模糊实现:

const blurShader = /* wgsl */ `
  @group(0) @binding(0) var<storage, read> inputImage: array<vec4<f32>>;
  @group(0) @binding(1) var<storage, read_write> outputImage: array<vec4<f32>>;
  @group(0) @binding(2) var<uniform> params: BlurParams;

  struct BlurParams {
    width: u32,
    height: u32,
    radius: u32,
  };

  fn gaussianWeight(x: i32, sigma: f32) -> f32 {
    let s = 2.0 * sigma * sigma;
    return exp(-f32(x * x) / s) / (3.14159265 * s);
  }

  @compute @workgroup_size(8, 8)
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let x = gid.x;
    let y = gid.y;
    if (x >= params.width || y >= params.height) {
      return;
    }
    
    let idx = y * params.width + x;
    var color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
    var totalWeight = 0.0;
    let sigma = f32(params.radius) / 2.0;
    
    for (var dy = -i32(params.radius); dy <= i32(params.radius); dy++) {
      for (var dx = -i32(params.radius); dx <= i32(params.radius); dx++) {
        let nx = i32(x) + dx;
        let ny = i32(y) + dy;
        if (nx >= 0 && nx < i32(params.width) && ny >= 0 && ny < i32(params.height)) {
          let weight = gaussianWeight(dx, sigma) * gaussianWeight(dy, sigma);
          let nidx = u32(ny) * params.width + u32(nx);
          color += inputImage[nidx] * weight;
          totalWeight += weight;
        }
      }
    }
    
    outputImage[idx] = color / totalWeight;
  }
`;

TypeScript类型定义与工程化封装

在TypeScript实战项目中,建议对WebGPU API做类型安全的封装:

// gpu-compute.ts
interface ComputeKernelOptions {
  shaderCode: string;
  entryPoint?: string;
  workgroupSize?: [number, number, number];
}

class GPUComputeEngine {
  private device: GPUDevice;
  private pipelines: Map<string, GPUComputePipeline> = new Map();

  constructor(device: GPUDevice) {
    this.device = device;
  }

  async createKernel(name: string, options: ComputeKernelOptions) {
    const module = this.device.createShaderModule({
      code: options.shaderCode,
    });

    const pipeline = this.device.createComputePipeline({
      layout: 'auto',
      compute: {
        module,
        entryPoint: options.entryPoint ?? 'main',
      },
    });

    this.pipelines.set(name, pipeline);
    return pipeline;
  }

  async run(
    kernelName: string,
    buffers: GPUBuffer[],
    workgroupCount: [number, number, number]
  ): Promise<void> {
    const pipeline = this.pipelines.get(kernelName);
    if (!pipeline) throw new Error(`Kernel ${kernelName} not found`);

    const bindGroup = this.device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: buffers.map((buffer, i) => ({
        binding: i,
        resource: { buffer },
      })),
    });

    const encoder = this.device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(...workgroupCount);
    pass.end();

    this.device.queue.submit([encoder.finish()]);
  }
}

// 使用示例
const engine = new GPUComputeEngine(device);
await engine.createKernel('vectorAdd', { shaderCode });
await engine.run('vectorAdd', [bufferA, bufferB, bufferOutput], [15625, 1, 1]);

WebGPU的Compute Shader在100万元素向量加法测试中,相比JavaScript循环计算约有50-100倍的性能提升(测试环境:RTX 3060 + Chrome 120)。实际性能增益取决于数据规模和计算复杂度,数据传输开销在小规模计算中可能抵消GPU加速收益。在生产环境中使用WebGPU时,需要添加降级方案(fallback到JavaScript或WebGL),确保不支持的浏览器仍能正常运行。

原创文章,作者:小编,如若转载,请注明出处:https://www.yunthe.com/webgpu-ru-men-shi-zhan-zai-liu-lan-qi-zhong-shi-xian-gao/

(0)
小编小编
上一篇 9小时前
下一篇 9小时前

相关推荐