@multiplekex/shallot 0.2.4 → 0.3.0

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

@@ -1,476 +0,0 @@
-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];
-}
-
-class PrefixSum {
-    private constructor(private passes: PrefixPass[]) {}
-
-    static async create(device: GPUDevice, data: GPUBuffer, count: number): Promise<PrefixSum> {
-        const passes: PrefixPass[] = [];
-        const module = device.createShaderModule({ code: prefixSumShader });
-        await PrefixSum.build(device, module, data, count, passes);
-        return new PrefixSum(passes);
-    }
-
-    private static async build(
-        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 PrefixSum.build(device, module, blockSums, wgCount, passes);
-
-            passes.push({
-                pipeline: await device.createComputePipelineAsync({
-                    layout: pipelineLayout,
-                    compute: { module, entryPoint: "addBlocks", constants: { COUNT: count } },
-                }),
-                bindGroup,
-                dispatch,
-            });
-        }
-    }
-
-    dispatch(pass: GPUComputePassEncoder): void {
-        for (const p of this.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 };
-}
-
-class RadixSort {
-    private constructor(
-        private passes: RadixPass[],
-        private prefixSum: PrefixSum,
-        private workgroups: [number, number]
-    ) {}
-
-    static async create(
-        device: GPUDevice,
-        keys: GPUBuffer,
-        values: GPUBuffer,
-        count: number
-    ): Promise<RadixSort> {
-        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 PrefixSum.create(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 new RadixSort(passes, prefixSum, workgroups);
-    }
-
-    dispatch(pass: GPUComputePassEncoder): void {
-        const [x, y] = this.workgroups;
-        for (const p of this.passes) {
-            pass.setPipeline(p.blockSum.pipeline);
-            pass.setBindGroup(0, p.blockSum.bindGroup);
-            pass.dispatchWorkgroups(x, y, 1);
-
-            this.prefixSum.dispatch(pass);
-
-            pass.setPipeline(p.reorder.pipeline);
-            pass.setBindGroup(0, p.reorder.bindGroup);
-            pass.dispatchWorkgroups(x, y, 1);
-        }
-    }
-}
-
-export interface RadixSortConfig {
-    keys: GPUBuffer;
-    values: GPUBuffer;
-    count: number;
-}
-
-export function createRadixSort(device: GPUDevice, config: RadixSortConfig): Promise<RadixSort> {
-    return RadixSort.create(device, config.keys, config.values, config.count);
-}
-
-export function createRadixSortNode(config: RadixSortConfig): ComputeNode {
-    let sort: RadixSort | 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();
-            sort!.dispatch(pass);
-            pass.end();
-        },
-    };
-}