@multiplekex/shallot 0.2.5 → 0.3.0

package/src/standard/raytracing/index.ts

@@ -1,31 +1,39 @@
 import type { Plugin, State } from "../../core";
 import { MAX_ENTITIES, resource } from "../../core";
+import { setTraits } from "../../core/component";
 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 { Camera } from "../render/camera";
+import { Mesh, getMesh, getMeshVersion } from "../render/mesh";
+import { getDefaultAllSurfaces, Surface } 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 { Render, registerShadowSettings } from "../render";
+import { RasterPlugin, RasterResource } from "../raster";
 import { createBLASAtlas, type BLASAtlas } from "./bvh/blas";
 import { createTLASBuffers, createTLASNode, type TLASBuffers } from "./bvh/tlas";
-import { getMesh } from "../render/mesh";
+import { createBLASRefitNode } from "./bvh/refit";
 import { createInstanceNode } from "./instance";
 import { createDepthConvertNode } from "./depth";
 import { compileUberShader } from "./shaders";
 
-export { compileUberShader, compileRTSurface, compileForwardShader } from "./shaders";
+export const Raytracing = {};
+export const Dynamic = {};
+
+export const RaytracingShadowSettings = {
+    softness: [] as number[],
+    samples: [] as number[],
+};
+
+setTraits(RaytracingShadowSettings, {
+    defaults: () => ({
+        softness: 0,
+        samples: 1,
+    }),
+});
+
+export { compileUberShader, compileRTSurface } from "./shaders";
 export { createDepthConvertNode } from "./depth";
 export { createInstanceNode } from "./instance";
 export { generateRay } from "./ray";
@@ -41,205 +49,184 @@ interface RTState {
     instanceCount: GPUBuffer;
     tlas: TLASBuffers;
     blasAtlas: BLASAtlas;
+    blasVersion: number;
+    rendered: boolean;
 }
 
 export const RTResource = resource<RTState>("raytracing");
 
-function createRTForwardNode(config: {
+interface RTRenderConfig {
     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;
+    getBlasAtlas: () => BLASAtlas;
+    getSurfaces: () => SurfaceData[];
+    getRaytracing: () => boolean;
+    setRendered: (value: boolean) => void;
+    acquire?: (message?: string) => (() => void) | undefined;
+}
 
-    let bindGroup0: GPUBindGroup | null = null;
-    let bindGroup1: GPUBindGroup | null = null;
-    let cachedWidth = 0;
-    let cachedHeight = 0;
+interface RTRenderGPU {
+    pipeline: GPUComputePipeline | null;
+    compiling: boolean;
+    bindGroup0: GPUBindGroup | null;
+    bindGroup1: GPUBindGroup | null;
+    cachedWidth: number;
+    cachedHeight: number;
+    cachedBlasAtlas: BLASAtlas | null;
+}
 
-    let skyBindGroup: GPUBindGroup | null = null;
-    let batchBindGroups: GPUBindGroup[] = [];
+async function compileRTPipeline(
+    device: GPUDevice,
+    surfaces: SurfaceData[]
+): Promise<GPUComputePipeline> {
+    const code = compileUberShader(surfaces);
+    const module = device.createShaderModule({ code });
+    return device.createComputePipelineAsync({
+        layout: "auto",
+        compute: { module, entryPoint: "main" },
+    });
+}
 
-    function executeRT(ctx: ExecutionContext): void {
-        const { device, encoder } = ctx;
+function triggerCompile(gpu: RTRenderGPU, device: GPUDevice, config: RTRenderConfig): void {
+    gpu.compiling = true;
+    const release = config.acquire?.("compiling shaders");
+    const surfaces = config.getSurfaces();
+    compileRTPipeline(device, surfaces)
+        .then((p) => {
+            gpu.pipeline = p;
+        })
+        .finally(() => {
+            gpu.compiling = false;
+            release?.();
+        });
+}
 
-        const colorView = ctx.getTextureView("color");
-        const depthView = ctx.getTextureView("depth");
-        const eidView = ctx.getTextureView("eid");
+function executeRTRender(gpu: RTRenderGPU, ctx: ExecutionContext, config: RTRenderConfig): void {
+    const wantRT = config.getRaytracing();
 
-        if (!colorView || !depthView || !eidView || !rtPipeline) {
-            return;
-        }
+    if (!wantRT) {
+        config.setRendered(false);
+        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 } },
-                ],
-            });
+    if (!gpu.pipeline && !gpu.compiling) {
+        triggerCompile(gpu, ctx.device, config);
+    }
 
-            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 } },
-                ],
-            });
+    if (!gpu.pipeline) {
+        config.setRendered(false);
+        return;
+    }
 
-            cachedWidth = width;
-            cachedHeight = height;
-        }
+    const { device, encoder } = ctx;
+
+    const colorView = ctx.getTextureView("color");
+    const depthView = ctx.getTextureView("depth");
+    const eidView = ctx.getTextureView("eid");
 
-        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();
+    if (!colorView || !depthView || !eidView) {
+        config.setRendered(false);
+        return;
     }
 
-    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 colorTexture = ctx.getTexture("color");
+    if (!colorTexture) {
+        config.setRendered(false);
+        return;
+    }
 
-        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 ?? [];
+    const width = colorTexture.width;
+    const height = colorTexture.height;
+
+    const blas = config.getBlasAtlas();
+
+    const sizeChanged = width !== gpu.cachedWidth || height !== gpu.cachedHeight;
+    const blasChanged = blas !== gpu.cachedBlasAtlas;
+
+    if (sizeChanged) {
+        gpu.bindGroup0 = device.createBindGroup({
+            layout: gpu.pipeline.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 } },
+            ],
+        });
+        gpu.cachedWidth = width;
+        gpu.cachedHeight = height;
     }
 
-    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" },
+    if (sizeChanged || blasChanged) {
+        gpu.bindGroup1 = device.createBindGroup({
+            layout: gpu.pipeline.getBindGroupLayout(1),
+            entries: [
+                { binding: 0, resource: { buffer: config.tlasNodes } },
+                { binding: 1, resource: { buffer: config.tlasInstanceIds } },
+                { binding: 2, resource: { buffer: blas.nodesBuffer } },
+                { binding: 3, resource: { buffer: blas.triIdsBuffer } },
+                { binding: 4, resource: { buffer: blas.trianglesBuffer } },
+                { binding: 5, resource: { buffer: blas.metaBuffer } },
+                { binding: 6, resource: { buffer: config.instanceInverses } },
+            ],
         });
+        gpu.cachedBlasAtlas = blas;
     }
 
+    const pass = encoder.beginComputePass();
+    pass.setPipeline(gpu.pipeline);
+    pass.setBindGroup(0, gpu.bindGroup0!);
+    pass.setBindGroup(1, gpu.bindGroup1!);
+    pass.dispatchWorkgroups(Math.ceil(width / 8), Math.ceil(height / 8));
+    pass.end();
+
+    config.setRendered(true);
+}
+
+function createRTRenderNode(config: RTRenderConfig): ComputeNode {
+    const gpu: RTRenderGPU = {
+        pipeline: null,
+        compiling: false,
+        bindGroup0: null,
+        bindGroup1: null,
+        cachedWidth: 0,
+        cachedHeight: 0,
+        cachedBlasAtlas: null,
+    };
+
     return {
-        id: "forward",
+        id: "rt-render",
         inputs: [{ id: "tlas-bvh-nodes", access: "read" as const }],
         outputs: [
-            { id: "color", access: "write" },
+            { id: "rt-output", 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?.();
-                    });
+                triggerCompile(gpu, device, config);
             }
         },
 
         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);
-            }
+            executeRTRender(gpu, ctx, config);
         },
     };
 }
 
 export const RaytracingPlugin: Plugin = {
     systems: [],
+    components: {
+        Raytracing,
+        RaytracingShadowSettings,
+        Dynamic,
+    },
     dependencies: [RasterPlugin],
 
     async initialize(state: State) {
@@ -247,6 +234,12 @@ export const RaytracingPlugin: Plugin = {
         const render = Render.from(state);
         if (!compute || !render) return;
 
+        registerShadowSettings({
+            softness: RaytracingShadowSettings.softness,
+            samples: RaytracingShadowSettings.samples,
+            isActive: (s, eid) => s.hasComponent(eid, RaytracingShadowSettings),
+        });
+
         const { device } = compute;
 
         const createPropertyBuffer = (size: number, label?: string) =>
@@ -268,6 +261,8 @@ export const RaytracingPlugin: Plugin = {
             }),
             tlas: createTLASBuffers(device),
             blasAtlas,
+            blasVersion: getMeshVersion(),
+            rendered: false,
         };
 
         state.setResource(RTResource, rtState);
@@ -290,35 +285,75 @@ export const RaytracingPlugin: Plugin = {
             outputs: [{ id: "instance-count", access: "write" }],
             execute(ctx) {
                 if (!getRaytracing()) return;
+
+                const currentMeshVersion = getMeshVersion();
+                if (currentMeshVersion !== rtState.blasVersion) {
+                    rtState.blasAtlas.nodesBuffer.destroy();
+                    rtState.blasAtlas.triIdsBuffer.destroy();
+                    rtState.blasAtlas.metaBuffer.destroy();
+                    rtState.blasAtlas.trianglesBuffer.destroy();
+                    rtState.blasAtlas.shapeAABBs.destroy();
+                    rtState.blasAtlas.treeNodesBuffer.destroy();
+                    rtState.blasAtlas.parentIndicesBuffer.destroy();
+                    rtState.blasAtlas.baseTrianglesBuffer.destroy();
+                    rtState.blasAtlas.boundsFlagsBuffer.destroy();
+                    rtState.blasAtlas = createBLASAtlas(ctx.device, getMesh);
+                    rtState.blasVersion = currentMeshVersion;
+                }
+
                 const meshEntities = state.query([Mesh, WorldTransform]);
-                let meshCount = 0;
+                let entityCount = 0;
                 for (const eid of meshEntities) {
-                    entityIds[meshCount] = eid;
-                    meshCount++;
+                    entityIds[entityCount] = eid;
+                    entityCount++;
                 }
                 ctx.device.queue.writeBuffer(
                     rtState.tlas.entityIds,
                     0,
                     entityIds,
                     0,
-                    Math.max(meshCount, 1)
+                    Math.max(entityCount, 1)
                 );
-                countBuf[0] = meshCount;
+                countBuf[0] = entityCount;
                 ctx.device.queue.writeBuffer(rtState.instanceCount, 0, countBuf);
             },
         };
         compute.graph.add(uploadNode);
 
+        const getDynamicShapes = () => {
+            const result = new Map<number, SurfaceData>();
+            const surfaces = getDefaultAllSurfaces();
+            for (const eid of state.query([Mesh, Dynamic])) {
+                const shapeId = Mesh.shape[eid];
+                const surfaceType = Surface.type[eid] ?? 0;
+                const data = surfaces[surfaceType];
+                if (data?.vertex && !result.has(shapeId)) {
+                    result.set(shapeId, data);
+                }
+            }
+            return result;
+        };
+
+        const blasRefitNode = createBLASRefitNode({
+            getBlasAtlas: () => rtState.blasAtlas,
+            scene: render.scene,
+            sky: render.sky,
+            getDynamicShapes,
+            getRaytracing,
+        });
+        compute.graph.add(blasRefitNode);
+
         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,
+            getShapeAABBs: () => rtState.blasAtlas.shapeAABBs,
             getEntityCount: () => render.entityCount,
             getRaytracing,
+            inputs: [{ id: "blas-nodes", access: "read" }],
         });
         compute.graph.add(instanceNode);
 
@@ -338,43 +373,37 @@ export const RaytracingPlugin: Plugin = {
             })
         );
 
-        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({
+        const rtRenderNode = createRTRenderNode({
             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,
+            getBlasAtlas: () => rtState.blasAtlas,
+            getSurfaces: getDefaultAllSurfaces,
+            getRaytracing,
+            setRendered: (v: boolean) => {
+                rtState.rendered = v;
+                if (rasterState) rasterState.rtRendered = v;
+            },
+            acquire: (message) => Activity.from(state)?.acquire(message),
         });
-        compute.graph.set("forward", forwardNode);
+        compute.graph.add(rtRenderNode);
+
+        const existingForward = compute.graph.nodes.get("forward");
+        if (existingForward) {
+            const wrappedForward: ComputeNode = {
+                id: "forward",
+                inputs: [...existingForward.inputs, { id: "rt-output", access: "read" as const }],
+                outputs: existingForward.outputs,
+                sync: existingForward.sync,
+                prepare: existingForward.prepare,
+                execute: existingForward.execute,
+            };
+            compute.graph.set("forward", wrappedForward);
+        }
     },
 };

package/src/standard/raytracing/instance.ts

@@ -1,4 +1,4 @@
-import type { ComputeNode, ExecutionContext } from "../compute";
+import type { ComputeNode, ExecutionContext, ResourceRef } from "../compute";
 
 const WORKGROUP_SIZE = 64;
 
@@ -160,50 +160,62 @@ export interface InstanceConfig {
     matrices: GPUBuffer;
     sizes: GPUBuffer;
     shapes: GPUBuffer;
-    shapeAABBs: GPUBuffer;
     entityCount: GPUBuffer;
     instanceAABBs: GPUBuffer;
     instanceInverses: GPUBuffer;
+    getShapeAABBs: () => GPUBuffer;
     getEntityCount: () => number;
-    getRaytracing?: () => boolean;
+    getRaytracing: () => boolean;
+    inputs?: ResourceRef[];
 }
 
 export function createInstanceNode(config: InstanceConfig): ComputeNode {
     let pipeline: GPUComputePipeline | null = null;
     let bindGroup: GPUBindGroup | null = null;
+    let cachedShapeAABBs: GPUBuffer | null = null;
+
+    function rebuildBindGroup(device: GPUDevice, shapeAABBs: GPUBuffer): void {
+        bindGroup = device.createBindGroup({
+            layout: pipeline!.getBindGroupLayout(0),
+            entries: [
+                { binding: 0, resource: { buffer: config.matrices } },
+                { binding: 1, resource: { buffer: config.sizes } },
+                { binding: 2, resource: { buffer: config.shapes } },
+                { binding: 3, resource: { buffer: shapeAABBs } },
+                { binding: 4, resource: { buffer: config.entityCount } },
+                { binding: 5, resource: { buffer: config.instanceAABBs } },
+                { binding: 6, resource: { buffer: config.instanceInverses } },
+            ],
+        });
+        cachedShapeAABBs = shapeAABBs;
+    }
 
     return {
         id: "instance",
-        inputs: [],
+        inputs: config.inputs ?? [],
         outputs: [
             { id: "instance-aabbs", access: "write" },
             { id: "instance-inverses", access: "write" },
         ],
 
         async prepare(device: GPUDevice) {
-            const module = await device.createShaderModule({ code: shader });
+            const module = device.createShaderModule({ code: shader });
 
             pipeline = await device.createComputePipelineAsync({
                 layout: "auto",
                 compute: { module, entryPoint: "main" },
             });
 
-            bindGroup = device.createBindGroup({
-                layout: pipeline.getBindGroupLayout(0),
-                entries: [
-                    { binding: 0, resource: { buffer: config.matrices } },
-                    { binding: 1, resource: { buffer: config.sizes } },
-                    { binding: 2, resource: { buffer: config.shapes } },
-                    { binding: 3, resource: { buffer: config.shapeAABBs } },
-                    { binding: 4, resource: { buffer: config.entityCount } },
-                    { binding: 5, resource: { buffer: config.instanceAABBs } },
-                    { binding: 6, resource: { buffer: config.instanceInverses } },
-                ],
-            });
+            rebuildBindGroup(device, config.getShapeAABBs());
         },
 
         execute(ctx: ExecutionContext) {
-            if (config.getRaytracing && !config.getRaytracing()) return;
+            if (!config.getRaytracing()) return;
+
+            const currentAABBs = config.getShapeAABBs();
+            if (currentAABBs !== cachedShapeAABBs) {
+                rebuildBindGroup(ctx.device, currentAABBs);
+            }
 
             const workgroups = Math.ceil(config.getEntityCount() / WORKGROUP_SIZE);
 

package/src/standard/raytracing/ray.ts

@@ -1,4 +1,4 @@
-import type { Vec3, Ray } from "./bvh/structs";
+import type { Ray } from "./bvh/structs";
 
 export type { Ray };