グローバルワークグループをテクスチャで可視化

main.ts

import shaderCode from "./compute.wgsl?raw"
import { rgba32ToCanvas } from "./texture-to-canvas"

const IMG_SIZE = 128
const FORMAT = "rgba32float"

const adapter = await navigator.gpu.requestAdapter()
if (!adapter) {
  throw new Error("WebGPU cannot be initialized - Adapter not found")
}

const device = await adapter.requestDevice()
device.lost.then(() => {
  throw new Error("WebGPU cannot be initialized - Device has been lost")
})

const shaderModule = device.createShaderModule({ code: shaderCode })

const outputTexture = device.createTexture({
  size: [IMG_SIZE, IMG_SIZE, 1],
  format: FORMAT,
  usage:
    GPUTextureUsage.COPY_SRC |
    GPUTextureUsage.COPY_DST |
    GPUTextureUsage.STORAGE_BINDING |
    GPUTextureUsage.TEXTURE_BINDING
})

const outputPixels = new Float32Array(IMG_SIZE * IMG_SIZE * 4)
const outputBuffer = device.createBuffer({
  size: outputPixels.byteLength,
  usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
})
device.queue.writeBuffer(outputBuffer, 0, outputPixels)

const bindgroupLayout = device.createBindGroupLayout({
  entries: [
    {
      binding: 0,
      visibility: GPUShaderStage.COMPUTE,
      storageTexture: {
        format: FORMAT,
        viewDimension: "2d",
        access: "write-only"
      }
    }
  ]
})
const bindgroup = device.createBindGroup({
  layout: bindgroupLayout,
  entries: [
    {
      binding: 0,
      resource: outputTexture.createView()
    }
  ]
})

const computePipelineLayout = device.createPipelineLayout({
  bindGroupLayouts: [bindgroupLayout]
})
const computePipeline = device.createComputePipeline({
  layout: computePipelineLayout,
  compute: {
    module: shaderModule,
    entryPoint: "main"
  }
})

const commandEncoder = device.createCommandEncoder()

const computePass = commandEncoder.beginComputePass()
computePass.setPipeline(computePipeline)
computePass.setBindGroup(0, bindgroup)
computePass.dispatchWorkgroups(16, 16)
computePass.end()

commandEncoder.copyTextureToBuffer(
  { texture: outputTexture },
  { buffer: outputBuffer, bytesPerRow: IMG_SIZE * 4 * 4 },
  { width: IMG_SIZE, height: IMG_SIZE, depthOrArrayLayers: 1 }
)

device.queue.submit([commandEncoder.finish()])

await outputBuffer.mapAsync(GPUMapMode.READ)
const arrayBuffer = outputBuffer.getMappedRange()
const result = new Float32Array(arrayBuffer)

const $canvas = document.querySelector("canvas")!
rgba32ToCanvas($canvas, result, IMG_SIZE, IMG_SIZE)

outputBuffer.unmap()

compute.wgsl

@group(0) @binding(0) var out_texture: texture_storage_2d<rgba32float, write>;

@compute @workgroup_size(8, 8)
fn main(
  @builtin(global_invocation_id) global_id: vec3u,
  @builtin(workgroup_id) workgroup_id: vec3u,
  @builtin(num_workgroups) workgroup_size: vec3u,
) {
  let position = vec2i(global_id.xy);
  let color = vec4f(vec3f(workgroup_id) / vec3f(workgroup_size), 1.0);
  textureStore(out_texture, position, color);
}

texture-to-canvas.ts

export const rgba32ToCanvas = (canvas: HTMLCanvasElement, pixels: Float32Array, width: number, height: number) => {
  const ctx = canvas.getContext("2d")
  if (!ctx) {
    throw new Error("Failed to get 2D context")
  }
  canvas.width = width
  canvas.height = height

  const imageData = ctx.createImageData(width, height)
  const data = imageData.data
  for (let i = 0; i < pixels.length; i += 4) {
    data[i + 0] = pixels[i + 0] * 255
    data[i + 1] = pixels[i + 1] * 255
    data[i + 2] = pixels[i + 2] * 255
    data[i + 3] = pixels[i + 3] * 255
  }
  ctx.putImageData(imageData, 0, 0)
}