DenoでWebGPUによる計算処理

「Rustのwgpuで計算処理」と同じ計算処理を Deno（2.0.6）で WebGPU を使って実装してみました。

多少の違いはあるものの概ね同じ API で処理を組み立てる事ができ、WGSL（WebGPU Shading Language）はそのまま使えました。¹

Adapter や Device は以下のように取得できるので、特にライブラリ等を使う必要もありませんでした。

GPUAdapterとGPUDeviceの取得例

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

実装

TypeScript による実装結果はこのようになります。

sample.ts

const workGroupSize = 4

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

if (!device) {
    console.error('not found device')
    Deno.exit(0)
}

const input = new Uint32Array([1, 2, 3, 4, 5, 6, 7, 8])

console.log(`input: ${input}`)

// WGSL
const wgsl = `
@group(0) @binding(0)
var<storage, read_write> data: array<u32>;

@compute @workgroup_size(${workGroupSize})
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    data[global_id.x] = 3u * data[global_id.x] + 1u;
}
`

const shader = device.createShaderModule({
    code: wgsl
})

// コンピューティングパイプラインの作成
const pipeline = await device.createComputePipelineAsync({
    layout: 'auto',
    compute: {
        module: shader,
        entryPoint: 'main'
    }
})

// ストレージバッファの作成
const storageBuf = device.createBuffer({
    mappedAtCreation: true,
    size: input.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC
})
// ストレージバッファへの書き込み
const data = new Uint32Array(storageBuf.getMappedRange())
data.set(input)
storageBuf.unmap()

// 出力用バッファ作成
const outputBuf = device.createBuffer({
    size: storageBuf.size,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
})

// バインドグループ作成
const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
        {
            binding: 0,
            resource: {
                buffer: storageBuf
            }
        }
    ]
})

const encoder = device.createCommandEncoder()

const pass = encoder.beginComputePass()
pass.setPipeline(pipeline)
pass.setBindGroup(0, bindGroup)
pass.dispatchWorkgroups(input.length / workGroupSize)
pass.end()

// ストレージバッファの内容を出力用のバッファへコピー
encoder.copyBufferToBuffer(storageBuf, 0, outputBuf, 0, storageBuf.size)

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

await outputBuf.mapAsync(GPUMapMode.READ)

const output = new Uint32Array(outputBuf.getMappedRange())

console.log(`output: ${output}`)

outputBuf.unmap()

wgpu版との主な違いは下記のようになっています。

• バッファの作成と初期化を行う wgpu の create_buffer_init に相当する API が無さそうだったので、バッファ作成後に値を書き込み
• コンピューティングパイプライン作成時の layout に 'auto' を指定
• Uint32Array を使ってバッファを処理
• GPUQueue は GPUDevice.queue で取得

実行

実行結果はこのようになります。

実行結果

$ deno run sample.ts
input: 1,2,3,4,5,6,7,8
output: 4,7,10,13,16,19,22,25

1. WGSL はその名の通り WebGPU のためのものなので当然ではありますが ↩