@multiplekex/shallot 0.2.5 → 0.3.0

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

@@ -195,77 +195,81 @@ interface PrefixPass {
     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 } },
-            ],
-        });
+interface PrefixSumState {
+    passes: PrefixPass[];
+}
 
-        const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [layout] });
+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: "scan", constants: { COUNT: count } },
+                compute: { module, entryPoint: "addBlocks", 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);
-        }
+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);
     }
 }
 
@@ -274,189 +278,182 @@ interface RadixPass {
     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" } },
-            ],
-        });
+interface RadixSortState {
+    passes: RadixPass[];
+    prefixSum: PrefixSumState;
+    workgroups: [number, number];
+}
 
-        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" } },
-            ],
-        });
+export interface RadixSortConfig {
+    keys: GPUBuffer;
+    values: GPUBuffer;
+    count: number;
+}
 
-        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 },
-                            },
+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],
                         }),
-                    ]);
-                    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 } },
-                        ],
+                        compute: {
+                            module: blockSumModule,
+                            entryPoint: "main",
+                            constants: { WG_COUNT: wgCount, BIT: bit, COUNT: count },
+                        },
                     }),
-                },
-                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 } },
-                        ],
+                    device.createComputePipelineAsync({
+                        layout: device.createPipelineLayout({
+                            bindGroupLayouts: [reorderLayout],
+                        }),
+                        compute: {
+                            module: reorderModule,
+                            entryPoint: "main",
+                            constants: { WG_COUNT: wgCount, BIT: bit, COUNT: count },
+                        },
                     }),
-                },
-            });
-        }
-
-        return new RadixSort(passes, prefixSum, workgroups);
+                ]);
+                return { blockSum: blockSumPipeline, reorder: reorderPipeline };
+            })()
+        );
     }
 
-    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);
+    const pipelines = await Promise.all(pipelinePromises);
+    const passes: RadixPass[] = [];
 
-            this.prefixSum.dispatch(pass);
+    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;
 
-            pass.setPipeline(p.reorder.pipeline);
-            pass.setBindGroup(0, p.reorder.bindGroup);
-            pass.dispatchWorkgroups(x, y, 1);
-        }
+        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 } },
+                    ],
+                }),
+            },
+        });
     }
-}
 
-export interface RadixSortConfig {
-    keys: GPUBuffer;
-    values: GPUBuffer;
-    count: number;
+    return { passes, prefixSum, workgroups };
 }
 
-export function createRadixSort(device: GPUDevice, config: RadixSortConfig): Promise<RadixSort> {
-    return RadixSort.create(device, config.keys, config.values, config.count);
+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: RadixSort | null = null;
+    let sort: RadixSortState | null = null;
 
     return {
         id: "radix-sort",
@@ -469,7 +466,7 @@ export function createRadixSortNode(config: RadixSortConfig): ComputeNode {
 
         execute(ctx: ExecutionContext) {
             const pass = ctx.encoder.beginComputePass();
-            sort!.dispatch(pass);
+            dispatchRadixSort(sort!, pass);
             pass.end();
         },
     };