@multiplekex/shallot 0.2.4 → 0.3.0

package/src/standard/raytracing/index.ts

@@ -0,0 +1,409 @@
+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 } from "../render/camera";
+import { Mesh, getMesh, getMeshVersion } from "../render/mesh";
+import { getDefaultAllSurfaces, Surface } from "../render/surface";
+import type { SurfaceData } from "../render/surface";
+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 { createBLASRefitNode } from "./bvh/refit";
+import { createInstanceNode } from "./instance";
+import { createDepthConvertNode } from "./depth";
+import { compileUberShader } 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";
+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;
+    blasVersion: number;
+    rendered: boolean;
+}
+
+export const RTResource = resource<RTState>("raytracing");
+
+interface RTRenderConfig {
+    scene: GPUBuffer;
+    sky: GPUBuffer;
+    data: GPUBuffer;
+    tlasNodes: GPUBuffer;
+    tlasInstanceIds: GPUBuffer;
+    instanceInverses: GPUBuffer;
+    getBlasAtlas: () => BLASAtlas;
+    getSurfaces: () => SurfaceData[];
+    getRaytracing: () => boolean;
+    setRendered: (value: boolean) => void;
+    acquire?: (message?: string) => (() => void) | undefined;
+}
+
+interface RTRenderGPU {
+    pipeline: GPUComputePipeline | null;
+    compiling: boolean;
+    bindGroup0: GPUBindGroup | null;
+    bindGroup1: GPUBindGroup | null;
+    cachedWidth: number;
+    cachedHeight: number;
+    cachedBlasAtlas: BLASAtlas | null;
+}
+
+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 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?.();
+        });
+}
+
+function executeRTRender(gpu: RTRenderGPU, ctx: ExecutionContext, config: RTRenderConfig): void {
+    const wantRT = config.getRaytracing();
+
+    if (!wantRT) {
+        config.setRendered(false);
+        return;
+    }
+
+    if (!gpu.pipeline && !gpu.compiling) {
+        triggerCompile(gpu, ctx.device, config);
+    }
+
+    if (!gpu.pipeline) {
+        config.setRendered(false);
+        return;
+    }
+
+    const { device, encoder } = ctx;
+
+    const colorView = ctx.getTextureView("color");
+    const depthView = ctx.getTextureView("depth");
+    const eidView = ctx.getTextureView("eid");
+
+    if (!colorView || !depthView || !eidView) {
+        config.setRendered(false);
+        return;
+    }
+
+    const colorTexture = ctx.getTexture("color");
+    if (!colorTexture) {
+        config.setRendered(false);
+        return;
+    }
+
+    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;
+    }
+
+    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: "rt-render",
+        inputs: [{ id: "tlas-bvh-nodes", access: "read" as const }],
+        outputs: [
+            { id: "rt-output", access: "write" },
+            { id: "depth", access: "write" },
+        ],
+
+        async prepare(device: GPUDevice) {
+            if (config.getRaytracing()) {
+                triggerCompile(gpu, device, config);
+            }
+        },
+
+        execute(ctx: ExecutionContext) {
+            executeRTRender(gpu, ctx, config);
+        },
+    };
+}
+
+export const RaytracingPlugin: Plugin = {
+    systems: [],
+    components: {
+        Raytracing,
+        RaytracingShadowSettings,
+        Dynamic,
+    },
+    dependencies: [RasterPlugin],
+
+    async initialize(state: State) {
+        const compute = Compute.from(state);
+        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) =>
+            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,
+            blasVersion: getMeshVersion(),
+            rendered: false,
+        };
+
+        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 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 entityCount = 0;
+                for (const eid of meshEntities) {
+                    entityIds[entityCount] = eid;
+                    entityCount++;
+                }
+                ctx.device.queue.writeBuffer(
+                    rtState.tlas.entityIds,
+                    0,
+                    entityIds,
+                    0,
+                    Math.max(entityCount, 1)
+                );
+                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,
+            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);
+
+        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 rasterState = RasterResource.from(state);
+
+        const rtRenderNode = createRTRenderNode({
+            scene: render.scene,
+            sky: render.sky,
+            data: render.data,
+            tlasNodes: rtState.tlas.bvhNodes,
+            tlasInstanceIds: rtState.tlas.instanceIds,
+            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.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/{render → 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,48 +160,63 @@ export interface InstanceConfig {
     matrices: GPUBuffer;
     sizes: GPUBuffer;
     shapes: GPUBuffer;
-    shapeAABBs: GPUBuffer;
     entityCount: GPUBuffer;
     instanceAABBs: GPUBuffer;
     instanceInverses: GPUBuffer;
+    getShapeAABBs: () => GPUBuffer;
     getEntityCount: () => number;
+    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()) return;
+
+            const currentAABBs = config.getShapeAABBs();
+            if (currentAABBs !== cachedShapeAABBs) {
+                rebuildBindGroup(ctx.device, currentAABBs);
+            }
+
             const workgroups = Math.ceil(config.getEntityCount() / WORKGROUP_SIZE);
 
             const pass = ctx.encoder.beginComputePass();

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

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