@multiplekex/shallot 0.1.12 → 0.2.0

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

@@ -0,0 +1,476 @@
+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();
+        },
+    };
+}

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

@@ -0,0 +1,167 @@
+export interface Vec3 {
+    x: number;
+    y: number;
+    z: number;
+}
+
+export interface AABB {
+    min: Vec3;
+    max: Vec3;
+}
+
+export interface Triangle {
+    v0: Vec3;
+    e1: Vec3;
+    e2: Vec3;
+    n0: Vec3;
+    n1: Vec3;
+    n2: Vec3;
+    entityId: number;
+}
+
+export interface BVHNode {
+    min: Vec3;
+    max: Vec3;
+    leftChild: number;
+    rightChild: number;
+}
+
+export interface Ray {
+    origin: Vec3;
+    direction: Vec3;
+}
+
+export interface MortonPair {
+    code: number;
+    triangleId: number;
+}
+
+export interface HitResult {
+    hit: boolean;
+    t: number;
+    entityId: number;
+    u: number;
+    v: number;
+    normal: Vec3;
+    worldPos: Vec3;
+}
+
+export const LEAF_FLAG = 0x80000000;
+
+export function isLeaf(child: number): boolean {
+    return (child & LEAF_FLAG) !== 0;
+}
+
+export function leafIndex(child: number): number {
+    return child & ~LEAF_FLAG;
+}
+
+export const LEAF_FLAG_WGSL = /* wgsl */ `const LEAF_FLAG: u32 = 0x80000000u;`;
+
+export const TRIANGLE_STRUCT_WGSL = /* wgsl */ `
+struct Triangle {
+    v0: vec3<f32>,
+    entityId: u32,
+    e1: vec3<f32>,
+    _pad0: u32,
+    e2: vec3<f32>,
+    _pad1: u32,
+    n0_enc: u32,
+    n1_enc: u32,
+    n2_enc: u32,
+    _pad2: u32,
+}
+
+fn octEncode(n: vec3<f32>) -> u32 {
+    let absSum = abs(n.x) + abs(n.y) + abs(n.z);
+    var vx = n.x / absSum;
+    var vy = n.y / absSum;
+    let vz = n.z / absSum;
+    if (vz < 0.0) {
+        let signX = select(-1.0, 1.0, vx >= 0.0);
+        let signY = select(-1.0, 1.0, vy >= 0.0);
+        let newVx = (1.0 - abs(vy)) * signX;
+        let newVy = (1.0 - abs(vx)) * signY;
+        vx = newVx;
+        vy = newVy;
+    }
+    let x = u32(clamp((vx * 0.5 + 0.5) * 65535.0, 0.0, 65535.0));
+    let y = u32(clamp((vy * 0.5 + 0.5) * 65535.0, 0.0, 65535.0));
+    return (y << 16u) | x;
+}
+
+fn octDecode(enc: u32) -> vec3<f32> {
+    let x = f32(enc & 0xFFFFu) / 65535.0 * 2.0 - 1.0;
+    let y = f32(enc >> 16u) / 65535.0 * 2.0 - 1.0;
+    let z = 1.0 - abs(x) - abs(y);
+    var n: vec3<f32>;
+    if (z < 0.0) {
+        let signX = select(-1.0, 1.0, x >= 0.0);
+        let signY = select(-1.0, 1.0, y >= 0.0);
+        n = vec3<f32>((1.0 - abs(y)) * signX, (1.0 - abs(x)) * signY, z);
+    } else {
+        n = vec3<f32>(x, y, z);
+    }
+    return normalize(n);
+}`;
+
+export const TREE_NODE_STRUCT_WGSL = /* wgsl */ `
+struct TreeNode {
+    minX: f32,
+    minY: f32,
+    minZ: f32,
+    leftChild: u32,
+    maxX: f32,
+    maxY: f32,
+    maxZ: f32,
+    rightChild: u32,
+}`;
+
+export const BVH_NODE_STRUCT_WGSL = /* wgsl */ `
+struct BVHNode {
+    c0_minX: f32, c0_minY: f32, c0_minZ: f32, child0: u32,
+    c0_maxX: f32, c0_maxY: f32, c0_maxZ: f32, _pad0: u32,
+    c1_minX: f32, c1_minY: f32, c1_minZ: f32, child1: u32,
+    c1_maxX: f32, c1_maxY: f32, c1_maxZ: f32, _pad1: u32,
+    c2_minX: f32, c2_minY: f32, c2_minZ: f32, child2: u32,
+    c2_maxX: f32, c2_maxY: f32, c2_maxZ: f32, _pad2: u32,
+    c3_minX: f32, c3_minY: f32, c3_minZ: f32, child3: u32,
+    c3_maxX: f32, c3_maxY: f32, c3_maxZ: f32, _pad3: u32,
+}`;
+
+export const TREE_NODE_SIZE = 32;
+export const BVH_NODE_SIZE = 128;
+export const BLAS_NODE_SIZE = 32;
+export const BLAS_TRIANGLE_SIZE = 64;
+export const INVALID_NODE = 0xffffffff;
+
+export const BLAS_NODE_STRUCT_WGSL = /* wgsl */ `
+struct BLASNode {
+    minX: f32, minY: f32, minZ: f32, leftChild: u32,
+    maxX: f32, maxY: f32, maxZ: f32, rightChild: u32,
+}`;
+
+export const BLAS_TRIANGLE_STRUCT_WGSL = /* wgsl */ `
+struct BLASTriangle {
+    v0: vec3<f32>, _pad0: u32,
+    e1: vec3<f32>, _pad1: u32,
+    e2: vec3<f32>, _pad2: u32,
+    n0_enc: u32, n1_enc: u32, n2_enc: u32, _pad3: u32,
+}`;
+
+export const RAY_STRUCT_WGSL = /* wgsl */ `
+struct Ray {
+    origin: vec3<f32>,
+    direction: vec3<f32>,
+}`;
+
+export const HIT_RESULT_STRUCT_WGSL = /* wgsl */ `
+struct HitResult {
+    hit: bool,
+    t: f32,
+    entityId: u32,
+    u: f32,
+    v: f32,
+    normal: vec3<f32>,
+    worldPos: vec3<f32>,
+}`;