使用WebGPU进行数学计算

创建日期:2024-07-14
更新日期:2025-01-12

compute-pipeline.html

<!doctype html>

<html lang="zh-cn">
  <head>
    <meta charset="utf-8" />
    <title>计算管线</title>
    <link rel="shortcut icon" href="/favicon.ico" />
    <meta
      name="viewport"
      content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0"
    />
    <script id="shader" type="text/wgsl">
      @group(0) @binding(0)
      var<storage, read_write> output: array<f32>;

      @compute @workgroup_size(64)
      fn main(
        @builtin(global_invocation_id) global_id: vec3u,
        @builtin(local_invocation_id) local_id: vec3u,
      ) {
        // Avoid accessing the buffer out of bounds
        if (global_id.x >= 1000) {
          return;
        }
        output[global_id.x] = f32(global_id.x) * 1000. + f32(local_id.x);
      }
    </script>
  </head>

  <body>
    <script>
      const BUFFER_SIZE = 1000

      async function init() {
        if (!navigator.gpu) {
          throw Error('WebGPU not supported')
        }
        const adapter = await navigator.gpu.requestAdapter()
        if (!adapter) {
          throw Error("Couldn't request WebGPU adapter.")
        }
        const device = await adapter.requestDevice()

        const shaders = document.querySelector('#shader').innerText
        const shaderModule = device.createShaderModule({
          code: shaders
        })

        const output = device.createBuffer({
          size: BUFFER_SIZE,
          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
        })
        const stagingBuffer = device.createBuffer({
          size: BUFFER_SIZE,
          usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
        })

        const bindGroupLayout = device.createBindGroupLayout({
          entries: [
            {
              binding: 0,
              visibility: GPUShaderStage.COMPUTE,
              buffer: {
                type: 'storage'
              }
            }
          ]
        })
        const bindGroup = device.createBindGroup({
          layout: bindGroupLayout,
          entries: [
            {
              binding: 0,
              resource: {
                buffer: output
              }
            }
          ]
        })
        const computePipeline = device.createComputePipeline({
          layout: device.createPipelineLayout({
            bindGroupLayouts: [bindGroupLayout]
          }),
          compute: {
            module: shaderModule,
            entryPoint: 'main'
          }
        })

        const commandEncoder = device.createCommandEncoder()
        const passEncoder = commandEncoder.beginComputePass()
        passEncoder.setPipeline(computePipeline)
        passEncoder.setBindGroup(0, bindGroup)
        passEncoder.dispatchWorkgroups(Math.ceil(BUFFER_SIZE / 64))
        passEncoder.end()

        commandEncoder.copyBufferToBuffer(output, 0, stagingBuffer, 0, BUFFER_SIZE)
        device.queue.submit([commandEncoder.finish()])

        await stagingBuffer.mapAsync(GPUMapMode.READ, 0, BUFFER_SIZE)
        const copyArrayBuffer = stagingBuffer.getMappedRange(0, BUFFER_SIZE)
        const data = copyArrayBuffer.slice()
        stagingBuffer.unmap()
        console.log(new Float32Array(data))
      }

      function start() {
        init()
      }
      start()
    </script>
  </body>
</html>