@multiplekex/shallot 0.2.3 → 0.2.5

package/src/standard/raytracing/index.ts

@@ -0,0 +1,380 @@
+import type { Plugin, State } from "../../core";
+import { MAX_ENTITIES, resource } from "../../core";
+import { Compute } from "../compute";
+import type { ComputeNode, ExecutionContext } from "../compute";
+import { WorldTransform } from "../transforms";
+import { Activity } from "../activity";
+import { Camera, Raytracing, getClearColor, hasSkyComponent } from "../render/camera";
+import { Mesh, type Batch } from "../render/mesh";
+import { getDefaultAllSurfaces } from "../render/surface";
+import type { SurfaceData } from "../render/surface";
+import { COLOR_FORMAT } from "../render/scene";
+import { Render, RenderPlugin } from "../render";
+import {
+    RasterPlugin,
+    RasterResource,
+    createRasterPipeline,
+    createSkyPipeline,
+    executeRasterPass,
+    type RasterPassConfig,
+} from "../raster";
+import { createBLASAtlas, type BLASAtlas } from "./bvh/blas";
+import { createTLASBuffers, createTLASNode, type TLASBuffers } from "./bvh/tlas";
+import { getMesh } from "../render/mesh";
+import { createInstanceNode } from "./instance";
+import { createDepthConvertNode } from "./depth";
+import { compileUberShader } from "./shaders";
+
+export { compileUberShader, compileRTSurface, compileForwardShader } from "./shaders";
+export { createDepthConvertNode } from "./depth";
+export { createInstanceNode } from "./instance";
+export { generateRay } from "./ray";
+export type { Ray } from "./ray";
+export { intersectRayTriangle } from "./intersection";
+export type { RayTriangleResult } from "./intersection";
+export { extractTriangles } from "./triangle";
+export type { Triangle, FlatTriangle } from "./triangle";
+
+interface RTState {
+    instanceAABBs: GPUBuffer;
+    instanceInverses: GPUBuffer;
+    instanceCount: GPUBuffer;
+    tlas: TLASBuffers;
+    blasAtlas: BLASAtlas;
+}
+
+export const RTResource = resource<RTState>("raytracing");
+
+function createRTForwardNode(config: {
+    scene: GPUBuffer;
+    sky: GPUBuffer;
+    data: GPUBuffer;
+    matrices: GPUBuffer;
+    sizes: GPUBuffer;
+    indirect: GPUBuffer;
+    getSurfaces: () => SurfaceData[];
+    getRaytracing: () => boolean;
+    getClearColor: () => { r: number; g: number; b: number };
+    getSky: () => boolean;
+    acquire?: (message?: string) => (() => void) | undefined;
+    batches: () => (Batch | null)[];
+    tlasNodes: GPUBuffer;
+    tlasInstanceIds: GPUBuffer;
+    blasNodes: GPUBuffer;
+    blasTriIds: GPUBuffer;
+    blasTriangles: GPUBuffer;
+    blasMeta: GPUBuffer;
+    instanceInverses: GPUBuffer;
+}): ComputeNode {
+    let rasterPipeline: GPURenderPipeline | null = null;
+    let rtPipeline: GPUComputePipeline | null = null;
+    let rtCompiling = false;
+    let skyPipeline: GPURenderPipeline | null = null;
+
+    let bindGroup0: GPUBindGroup | null = null;
+    let bindGroup1: GPUBindGroup | null = null;
+    let cachedWidth = 0;
+    let cachedHeight = 0;
+
+    let skyBindGroup: GPUBindGroup | null = null;
+    let batchBindGroups: GPUBindGroup[] = [];
+
+    function executeRT(ctx: ExecutionContext): void {
+        const { device, encoder } = ctx;
+
+        const colorView = ctx.getTextureView("color");
+        const depthView = ctx.getTextureView("depth");
+        const eidView = ctx.getTextureView("eid");
+
+        if (!colorView || !depthView || !eidView || !rtPipeline) {
+            return;
+        }
+
+        const colorTexture = ctx.getTexture("color");
+        if (!colorTexture) return;
+
+        const width = colorTexture.width;
+        const height = colorTexture.height;
+
+        if (width !== cachedWidth || height !== cachedHeight) {
+            bindGroup0 = device.createBindGroup({
+                layout: rtPipeline.getBindGroupLayout(0),
+                entries: [
+                    { binding: 0, resource: { buffer: config.scene } },
+                    { binding: 1, resource: { buffer: config.data } },
+                    { binding: 2, resource: colorView },
+                    { binding: 3, resource: depthView },
+                    { binding: 4, resource: eidView },
+                    { binding: 5, resource: { buffer: config.sky } },
+                ],
+            });
+
+            bindGroup1 = device.createBindGroup({
+                layout: rtPipeline.getBindGroupLayout(1),
+                entries: [
+                    { binding: 0, resource: { buffer: config.tlasNodes } },
+                    { binding: 1, resource: { buffer: config.tlasInstanceIds } },
+                    { binding: 2, resource: { buffer: config.blasNodes } },
+                    { binding: 3, resource: { buffer: config.blasTriIds } },
+                    { binding: 4, resource: { buffer: config.blasTriangles } },
+                    { binding: 5, resource: { buffer: config.blasMeta } },
+                    { binding: 6, resource: { buffer: config.instanceInverses } },
+                ],
+            });
+
+            cachedWidth = width;
+            cachedHeight = height;
+        }
+
+        const pass = encoder.beginComputePass();
+        pass.setPipeline(rtPipeline);
+        pass.setBindGroup(0, bindGroup0!);
+        pass.setBindGroup(1, bindGroup1!);
+        pass.dispatchWorkgroups(Math.ceil(width / 8), Math.ceil(height / 8));
+        pass.end();
+    }
+
+    function executeRasterFallback(ctx: ExecutionContext): void {
+        if (!rasterPipeline) return;
+
+        const depthView = ctx.getTextureView("depth");
+        if (depthView) {
+            const clearPass = ctx.encoder.beginRenderPass({
+                colorAttachments: [
+                    {
+                        view: depthView as GPUTextureView,
+                        clearValue: { r: 0, g: 0, b: 0, a: 0 },
+                        loadOp: "clear" as const,
+                        storeOp: "store" as const,
+                    },
+                ],
+            });
+            clearPass.end();
+        }
+
+        const passConfig: RasterPassConfig = {
+            scene: config.scene,
+            sky: config.sky,
+            data: config.data,
+            matrices: config.matrices,
+            sizes: config.sizes,
+            indirect: config.indirect,
+            rasterPipeline,
+            skyPipeline,
+            getClearColor: config.getClearColor,
+            getSky: config.getSky,
+            batches: config.batches,
+            skyBindGroup,
+            batchBindGroups,
+        };
+        executeRasterPass(ctx, passConfig);
+        skyBindGroup = passConfig.skyBindGroup ?? null;
+        batchBindGroups = passConfig.batchBindGroups ?? [];
+    }
+
+    async function createRTPipeline(
+        device: GPUDevice,
+        surfaces: SurfaceData[]
+    ): Promise<GPUComputePipeline> {
+        const code = compileUberShader(surfaces);
+        const module = device.createShaderModule({ code });
+        return device.createComputePipelineAsync({
+            layout: "auto",
+            compute: { module, entryPoint: "main" },
+        });
+    }
+
+    return {
+        id: "forward",
+        inputs: [{ id: "tlas-bvh-nodes", access: "read" as const }],
+        outputs: [
+            { id: "color", access: "write" },
+            { id: "depth", access: "write" },
+            { id: "eid", access: "write" },
+        ],
+
+        async prepare(device: GPUDevice) {
+            const surfaces = config.getSurfaces();
+            rasterPipeline = await createRasterPipeline(device, surfaces, COLOR_FORMAT);
+            skyPipeline = await createSkyPipeline(device, COLOR_FORMAT);
+
+            if (config.getRaytracing()) {
+                rtCompiling = true;
+                const release = config.acquire?.("compiling shaders");
+                createRTPipeline(device, surfaces)
+                    .then((p) => {
+                        rtPipeline = p;
+                    })
+                    .finally(() => {
+                        rtCompiling = false;
+                        release?.();
+                    });
+            }
+        },
+
+        execute(ctx: ExecutionContext) {
+            const wantRT = config.getRaytracing();
+
+            if (wantRT && !rtPipeline && !rtCompiling) {
+                rtCompiling = true;
+                const release = config.acquire?.("compiling shaders");
+                const surfaces = config.getSurfaces();
+                createRTPipeline(ctx.device, surfaces)
+                    .then((p) => {
+                        rtPipeline = p;
+                    })
+                    .finally(() => {
+                        rtCompiling = false;
+                        release?.();
+                    });
+            }
+
+            if (wantRT && rtPipeline) {
+                executeRT(ctx);
+            } else {
+                executeRasterFallback(ctx);
+            }
+        },
+    };
+}
+
+export const RaytracingPlugin: Plugin = {
+    systems: [],
+    dependencies: [RasterPlugin],
+
+    async initialize(state: State) {
+        const compute = Compute.from(state);
+        const render = Render.from(state);
+        if (!compute || !render) return;
+
+        const { device } = compute;
+
+        const createPropertyBuffer = (size: number, label?: string) =>
+            device.createBuffer({
+                label: label ?? "property",
+                size,
+                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
+            });
+
+        const blasAtlas = createBLASAtlas(device, getMesh);
+
+        const rtState: RTState = {
+            instanceAABBs: createPropertyBuffer(MAX_ENTITIES * 32, "instanceAABBs"),
+            instanceInverses: createPropertyBuffer(MAX_ENTITIES * 64, "instanceInverses"),
+            instanceCount: device.createBuffer({
+                label: "instanceCount",
+                size: 4,
+                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
+            }),
+            tlas: createTLASBuffers(device),
+            blasAtlas,
+        };
+
+        state.setResource(RTResource, rtState);
+
+        const getRaytracing = () => {
+            for (const eid of state.query([Camera])) {
+                if (Camera.active[eid]) {
+                    return state.hasComponent(eid, Raytracing);
+                }
+            }
+            return false;
+        };
+
+        const entityIds = new Uint32Array(MAX_ENTITIES);
+        const countBuf = new Uint32Array(1);
+
+        const uploadNode: ComputeNode = {
+            id: "rt-upload",
+            inputs: [],
+            outputs: [{ id: "instance-count", access: "write" }],
+            execute(ctx) {
+                if (!getRaytracing()) return;
+                const meshEntities = state.query([Mesh, WorldTransform]);
+                let meshCount = 0;
+                for (const eid of meshEntities) {
+                    entityIds[meshCount] = eid;
+                    meshCount++;
+                }
+                ctx.device.queue.writeBuffer(
+                    rtState.tlas.entityIds,
+                    0,
+                    entityIds,
+                    0,
+                    Math.max(meshCount, 1)
+                );
+                countBuf[0] = meshCount;
+                ctx.device.queue.writeBuffer(rtState.instanceCount, 0, countBuf);
+            },
+        };
+        compute.graph.add(uploadNode);
+
+        const instanceNode = createInstanceNode({
+            matrices: render.matrices,
+            sizes: render.sizes,
+            shapes: render.shapes,
+            shapeAABBs: rtState.blasAtlas.shapeAABBs,
+            entityCount: render.entityCountBuffer,
+            instanceAABBs: rtState.instanceAABBs,
+            instanceInverses: rtState.instanceInverses,
+            getEntityCount: () => render.entityCount,
+            getRaytracing,
+        });
+        compute.graph.add(instanceNode);
+
+        const tlasNode = createTLASNode({
+            instanceAABBs: rtState.instanceAABBs,
+            instanceCount: rtState.instanceCount,
+            tlas: rtState.tlas,
+            getEntityCount: () => render.entityCount,
+            getRaytracing,
+        });
+        compute.graph.add(tlasNode);
+
+        compute.graph.add(
+            createDepthConvertNode({
+                scene: render.scene,
+                getRaytracing,
+            })
+        );
+
+        const getActiveClearColor = () => {
+            for (const eid of state.query([Camera])) {
+                if (Camera.active[eid]) return getClearColor(eid);
+            }
+            return { r: 0, g: 0, b: 0 };
+        };
+
+        const getSky = () => {
+            for (const eid of state.query([Camera])) {
+                if (Camera.active[eid]) return hasSkyComponent(state, eid);
+            }
+            return false;
+        };
+
+        const rasterState = RasterResource.from(state);
+
+        const forwardNode = createRTForwardNode({
+            scene: render.scene,
+            sky: render.sky,
+            data: render.data,
+            matrices: render.matrices,
+            sizes: render.sizes,
+            indirect: rasterState!.indirect,
+            getSurfaces: getDefaultAllSurfaces,
+            getRaytracing,
+            getClearColor: getActiveClearColor,
+            getSky,
+            acquire: (message) => Activity.from(state)?.acquire(message),
+            batches: () => rasterState?.batches ?? [],
+            tlasNodes: rtState.tlas.bvhNodes,
+            tlasInstanceIds: rtState.tlas.instanceIds,
+            blasNodes: rtState.blasAtlas.nodesBuffer,
+            blasTriIds: rtState.blasAtlas.triIdsBuffer,
+            blasTriangles: rtState.blasAtlas.trianglesBuffer,
+            blasMeta: rtState.blasAtlas.metaBuffer,
+            instanceInverses: rtState.instanceInverses,
+        });
+        compute.graph.set("forward", forwardNode);
+    },
+};

package/src/standard/{render → raytracing}/instance.ts

@@ -165,6 +165,7 @@ export interface InstanceConfig {
     instanceAABBs: GPUBuffer;
     instanceInverses: GPUBuffer;
     getEntityCount: () => number;
+    getRaytracing?: () => boolean;
 }
 
 export function createInstanceNode(config: InstanceConfig): ComputeNode {
@@ -202,6 +203,8 @@ export function createInstanceNode(config: InstanceConfig): ComputeNode {
         },
 
         execute(ctx: ExecutionContext) {
+            if (config.getRaytracing && !config.getRaytracing()) return;
+
             const workgroups = Math.ceil(config.getEntityCount() / WORKGROUP_SIZE);
 
             const pass = ctx.encoder.beginComputePass();