@multiplekex/shallot 0.2.4 → 0.3.0

package/src/standard/raytracing/bvh/radix.ts

@@ -0,0 +1,473 @@
+import type { ComputeNode, ExecutionContext } from "../../compute";
+
+const WG_X = 16;
+const WG_Y = 16;
+const WG_SIZE = WG_X * WG_Y;
+const ITEMS_PER_WG = 2 * WG_SIZE;
+
+const blockSumShader = /* wgsl */ `
+@group(0) @binding(0) var<storage, read> input: array<u32>;
+@group(0) @binding(1) var<storage, read_write> localSums: array<u32>;
+@group(0) @binding(2) var<storage, read_write> blockSums: array<u32>;
+
+override WG_COUNT: u32;
+override BIT: u32;
+override COUNT: u32;
+
+var<workgroup> wgData: array<u32, 2 * (${WG_SIZE} + 1)>;
+
+@compute @workgroup_size(${WG_X}, ${WG_Y}, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) wdim: vec3<u32>,
+    @builtin(local_invocation_index) tid: u32,
+) {
+    let workgroup = wid.x + wid.y * wdim.x;
+    let base = workgroup * ${WG_SIZE}u;
+    let gid = base + tid;
+
+    let val = select(input[gid], 0u, gid >= COUNT);
+    let bits = (val >> BIT) & 0x3;
+
+    var sums = array<u32, 4>(0, 0, 0, 0);
+    var lastThread = 0xffffffffu;
+
+    if (workgroup < WG_COUNT) {
+        lastThread = min(${WG_SIZE}u, COUNT - base) - 1;
+    }
+
+    let stride = ${WG_SIZE}u + 1;
+    var swap = 0u;
+    var inOff = tid;
+    var outOff = tid + stride;
+
+    for (var b = 0u; b < 4; b++) {
+        let mask = select(0u, 1u, bits == b);
+        wgData[inOff + 1] = mask;
+        workgroupBarrier();
+
+        var sum = 0u;
+        for (var off = 1u; off < ${WG_SIZE}u; off *= 2) {
+            if (tid >= off) {
+                sum = wgData[inOff] + wgData[inOff - off];
+            } else {
+                sum = wgData[inOff];
+            }
+            wgData[outOff] = sum;
+            outOff = inOff;
+            swap = stride - swap;
+            inOff = tid + swap;
+            workgroupBarrier();
+        }
+
+        sums[b] = sum;
+
+        if (tid == lastThread) {
+            blockSums[b * WG_COUNT + workgroup] = sum + mask;
+        }
+
+        outOff = inOff;
+        swap = stride - swap;
+        inOff = tid + swap;
+    }
+
+    if (gid < COUNT) {
+        localSums[gid] = sums[bits];
+    }
+}
+`;
+
+const reorderShader = /* wgsl */ `
+@group(0) @binding(0) var<storage, read> inKeys: array<u32>;
+@group(0) @binding(1) var<storage, read_write> outKeys: array<u32>;
+@group(0) @binding(2) var<storage, read> localSums: array<u32>;
+@group(0) @binding(3) var<storage, read> blockSums: array<u32>;
+@group(0) @binding(4) var<storage, read> inVals: array<u32>;
+@group(0) @binding(5) var<storage, read_write> outVals: array<u32>;
+
+override WG_COUNT: u32;
+override BIT: u32;
+override COUNT: u32;
+
+@compute @workgroup_size(${WG_X}, ${WG_Y}, 1)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) wdim: vec3<u32>,
+    @builtin(local_invocation_index) tid: u32,
+) {
+    let workgroup = wid.x + wid.y * wdim.x;
+    let gid = workgroup * ${WG_SIZE}u + tid;
+
+    if (gid >= COUNT) { return; }
+
+    let k = inKeys[gid];
+    let v = inVals[gid];
+    let bits = (k >> BIT) & 0x3;
+    let dst = blockSums[bits * WG_COUNT + workgroup] + localSums[gid];
+
+    outKeys[dst] = k;
+    outVals[dst] = v;
+}
+`;
+
+const prefixSumShader = /* wgsl */ `
+@group(0) @binding(0) var<storage, read_write> data: array<u32>;
+@group(0) @binding(1) var<storage, read_write> blockSums: array<u32>;
+
+override COUNT: u32;
+
+var<workgroup> temp: array<u32, ${ITEMS_PER_WG * 2}>;
+
+@compute @workgroup_size(${WG_X}, ${WG_Y}, 1)
+fn scan(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) wdim: vec3<u32>,
+    @builtin(local_invocation_index) tid: u32,
+) {
+    let workgroup = wid.x + wid.y * wdim.x;
+    let base = workgroup * ${WG_SIZE}u;
+    let gid = base + tid;
+    let eid = gid * 2;
+
+    temp[tid * 2] = select(data[eid], 0u, eid >= COUNT);
+    temp[tid * 2 + 1] = select(data[eid + 1], 0u, eid + 1 >= COUNT);
+
+    var offset = 1u;
+    for (var d = ${ITEMS_PER_WG}u >> 1; d > 0; d >>= 1) {
+        workgroupBarrier();
+        if (tid < d) {
+            let ai = offset * (tid * 2 + 1) - 1;
+            let bi = offset * (tid * 2 + 2) - 1;
+            temp[bi] += temp[ai];
+        }
+        offset *= 2;
+    }
+
+    if (tid == 0) {
+        blockSums[workgroup] = temp[${ITEMS_PER_WG}u - 1];
+        temp[${ITEMS_PER_WG}u - 1] = 0;
+    }
+
+    for (var d = 1u; d < ${ITEMS_PER_WG}u; d *= 2) {
+        offset >>= 1;
+        workgroupBarrier();
+        if (tid < d) {
+            let ai = offset * (tid * 2 + 1) - 1;
+            let bi = offset * (tid * 2 + 2) - 1;
+            let t = temp[ai];
+            temp[ai] = temp[bi];
+            temp[bi] += t;
+        }
+    }
+    workgroupBarrier();
+
+    if (eid < COUNT) { data[eid] = temp[tid * 2]; }
+    if (eid + 1 < COUNT) { data[eid + 1] = temp[tid * 2 + 1]; }
+}
+
+@compute @workgroup_size(${WG_X}, ${WG_Y}, 1)
+fn addBlocks(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(num_workgroups) wdim: vec3<u32>,
+    @builtin(local_invocation_index) tid: u32,
+) {
+    let workgroup = wid.x + wid.y * wdim.x;
+    let eid = (workgroup * ${WG_SIZE}u + tid) * 2;
+
+    if (eid >= COUNT) { return; }
+
+    let sum = blockSums[workgroup];
+    data[eid] += sum;
+    if (eid + 1 < COUNT) { data[eid + 1] += sum; }
+}
+`;
+
+function dispatchSize(device: GPUDevice, count: number): [number, number] {
+    const max = device.limits.maxComputeWorkgroupsPerDimension;
+    if (count <= max) return [count, 1];
+    const x = Math.ceil(Math.sqrt(count));
+    return [x, Math.ceil(count / x)];
+}
+
+interface PrefixPass {
+    pipeline: GPUComputePipeline;
+    bindGroup: GPUBindGroup;
+    dispatch: [number, number];
+}
+
+interface PrefixSumState {
+    passes: PrefixPass[];
+}
+
+async function buildPrefixPasses(
+    device: GPUDevice,
+    module: GPUShaderModule,
+    data: GPUBuffer,
+    count: number,
+    passes: PrefixPass[]
+): Promise<void> {
+    const wgCount = Math.ceil(count / ITEMS_PER_WG);
+    const dispatch = dispatchSize(device, wgCount);
+
+    const blockSums = device.createBuffer({
+        size: Math.max(wgCount * 4, 4),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+
+    const layout = device.createBindGroupLayout({
+        entries: [
+            { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+            { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        ],
+    });
+
+    const bindGroup = device.createBindGroup({
+        layout,
+        entries: [
+            { binding: 0, resource: { buffer: data } },
+            { binding: 1, resource: { buffer: blockSums } },
+        ],
+    });
+
+    const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [layout] });
+
+    passes.push({
+        pipeline: await device.createComputePipelineAsync({
+            layout: pipelineLayout,
+            compute: { module, entryPoint: "scan", constants: { COUNT: count } },
+        }),
+        bindGroup,
+        dispatch,
+    });
+
+    if (wgCount > 1) {
+        await buildPrefixPasses(device, module, blockSums, wgCount, passes);
+
+        passes.push({
+            pipeline: await device.createComputePipelineAsync({
+                layout: pipelineLayout,
+                compute: { module, entryPoint: "addBlocks", constants: { COUNT: count } },
+            }),
+            bindGroup,
+            dispatch,
+        });
+    }
+}
+
+async function createPrefixSum(
+    device: GPUDevice,
+    data: GPUBuffer,
+    count: number
+): Promise<PrefixSumState> {
+    const passes: PrefixPass[] = [];
+    const module = device.createShaderModule({ code: prefixSumShader });
+    await buildPrefixPasses(device, module, data, count, passes);
+    return { passes };
+}
+
+function dispatchPrefixSum(state: PrefixSumState, pass: GPUComputePassEncoder): void {
+    for (const p of state.passes) {
+        pass.setPipeline(p.pipeline);
+        pass.setBindGroup(0, p.bindGroup);
+        pass.dispatchWorkgroups(p.dispatch[0], p.dispatch[1], 1);
+    }
+}
+
+interface RadixPass {
+    blockSum: { pipeline: GPUComputePipeline; bindGroup: GPUBindGroup };
+    reorder: { pipeline: GPUComputePipeline; bindGroup: GPUBindGroup };
+}
+
+interface RadixSortState {
+    passes: RadixPass[];
+    prefixSum: PrefixSumState;
+    workgroups: [number, number];
+}
+
+export interface RadixSortConfig {
+    keys: GPUBuffer;
+    values: GPUBuffer;
+    count: number;
+}
+
+export async function createRadixSort(
+    device: GPUDevice,
+    config: RadixSortConfig
+): Promise<RadixSortState> {
+    const { keys, values, count } = config;
+    const wgCount = Math.ceil(count / WG_SIZE);
+    const workgroups = dispatchSize(device, wgCount);
+
+    const tmpKeys = device.createBuffer({
+        size: count * 4,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+    const tmpVals = device.createBuffer({
+        size: count * 4,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+    const localSums = device.createBuffer({
+        size: count * 4,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+    const blockSums = device.createBuffer({
+        size: 4 * wgCount * 4,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+    });
+
+    const prefixSum = await createPrefixSum(device, blockSums, 4 * wgCount);
+
+    const blockSumModule = device.createShaderModule({ code: blockSumShader });
+    const reorderModule = device.createShaderModule({ code: reorderShader });
+
+    const blockSumLayout = device.createBindGroupLayout({
+        entries: [
+            {
+                binding: 0,
+                visibility: GPUShaderStage.COMPUTE,
+                buffer: { type: "read-only-storage" },
+            },
+            { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+            { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        ],
+    });
+
+    const reorderLayout = device.createBindGroupLayout({
+        entries: [
+            {
+                binding: 0,
+                visibility: GPUShaderStage.COMPUTE,
+                buffer: { type: "read-only-storage" },
+            },
+            { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+            {
+                binding: 2,
+                visibility: GPUShaderStage.COMPUTE,
+                buffer: { type: "read-only-storage" },
+            },
+            {
+                binding: 3,
+                visibility: GPUShaderStage.COMPUTE,
+                buffer: { type: "read-only-storage" },
+            },
+            {
+                binding: 4,
+                visibility: GPUShaderStage.COMPUTE,
+                buffer: { type: "read-only-storage" },
+            },
+            { binding: 5, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        ],
+    });
+
+    const pipelinePromises: Promise<{
+        blockSum: GPUComputePipeline;
+        reorder: GPUComputePipeline;
+    }>[] = [];
+
+    for (let bit = 0; bit < 32; bit += 2) {
+        pipelinePromises.push(
+            (async () => {
+                const [blockSumPipeline, reorderPipeline] = await Promise.all([
+                    device.createComputePipelineAsync({
+                        layout: device.createPipelineLayout({
+                            bindGroupLayouts: [blockSumLayout],
+                        }),
+                        compute: {
+                            module: blockSumModule,
+                            entryPoint: "main",
+                            constants: { WG_COUNT: wgCount, BIT: bit, COUNT: count },
+                        },
+                    }),
+                    device.createComputePipelineAsync({
+                        layout: device.createPipelineLayout({
+                            bindGroupLayouts: [reorderLayout],
+                        }),
+                        compute: {
+                            module: reorderModule,
+                            entryPoint: "main",
+                            constants: { WG_COUNT: wgCount, BIT: bit, COUNT: count },
+                        },
+                    }),
+                ]);
+                return { blockSum: blockSumPipeline, reorder: reorderPipeline };
+            })()
+        );
+    }
+
+    const pipelines = await Promise.all(pipelinePromises);
+    const passes: RadixPass[] = [];
+
+    for (let i = 0; i < 16; i++) {
+        const bit = i * 2;
+        const even = bit % 4 === 0;
+        const inK = even ? keys : tmpKeys;
+        const inV = even ? values : tmpVals;
+        const outK = even ? tmpKeys : keys;
+        const outV = even ? tmpVals : values;
+
+        passes.push({
+            blockSum: {
+                pipeline: pipelines[i].blockSum,
+                bindGroup: device.createBindGroup({
+                    layout: blockSumLayout,
+                    entries: [
+                        { binding: 0, resource: { buffer: inK } },
+                        { binding: 1, resource: { buffer: localSums } },
+                        { binding: 2, resource: { buffer: blockSums } },
+                    ],
+                }),
+            },
+            reorder: {
+                pipeline: pipelines[i].reorder,
+                bindGroup: device.createBindGroup({
+                    layout: reorderLayout,
+                    entries: [
+                        { binding: 0, resource: { buffer: inK } },
+                        { binding: 1, resource: { buffer: outK } },
+                        { binding: 2, resource: { buffer: localSums } },
+                        { binding: 3, resource: { buffer: blockSums } },
+                        { binding: 4, resource: { buffer: inV } },
+                        { binding: 5, resource: { buffer: outV } },
+                    ],
+                }),
+            },
+        });
+    }
+
+    return { passes, prefixSum, workgroups };
+}
+
+export function dispatchRadixSort(state: RadixSortState, pass: GPUComputePassEncoder): void {
+    const [x, y] = state.workgroups;
+    for (const p of state.passes) {
+        pass.setPipeline(p.blockSum.pipeline);
+        pass.setBindGroup(0, p.blockSum.bindGroup);
+        pass.dispatchWorkgroups(x, y, 1);
+
+        dispatchPrefixSum(state.prefixSum, pass);
+
+        pass.setPipeline(p.reorder.pipeline);
+        pass.setBindGroup(0, p.reorder.bindGroup);
+        pass.dispatchWorkgroups(x, y, 1);
+    }
+}
+
+export function createRadixSortNode(config: RadixSortConfig): ComputeNode {
+    let sort: RadixSortState | null = null;
+
+    return {
+        id: "radix-sort",
+        inputs: [],
+        outputs: [],
+
+        async prepare(device: GPUDevice) {
+            sort = await createRadixSort(device, config);
+        },
+
+        execute(ctx: ExecutionContext) {
+            const pass = ctx.encoder.beginComputePass();
+            dispatchRadixSort(sort!, pass);
+            pass.end();
+        },
+    };
+}