@multiplekex/shallot 0.2.5 → 0.3.0

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

@@ -0,0 +1,711 @@
+import type { ComputeNode, ExecutionContext } from "../../compute";
+import type { BLASAtlas } from "./blas";
+import type { SurfaceData } from "../../render/surface";
+import { compileVertexBody } from "../../render/surface/shaders";
+import { SCENE_STRUCT_WGSL, SKY_STRUCT_WGSL } from "../../render/surface/structs";
+import {
+    BLAS_TRIANGLE_STRUCT_WGSL,
+    TREE_NODE_STRUCT_WGSL,
+    BVH_NODE_STRUCT_WGSL,
+    LEAF_FLAG_WGSL,
+} from "./structs";
+
+const WORKGROUP_SIZE = 64;
+
+const OCT_DECODE_WGSL = /* wgsl */ `
+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);
+}`;
+
+function compileDisplacementShader(vertexCode: string): string {
+    const vertexBody = compileVertexBody(vertexCode);
+
+    return /* wgsl */ `
+${SCENE_STRUCT_WGSL}
+${SKY_STRUCT_WGSL}
+${BLAS_TRIANGLE_STRUCT_WGSL}
+${OCT_DECODE_WGSL}
+
+struct Params {
+    triOffset: u32,
+    triCount: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> scene: Scene;
+@group(0) @binding(1) var<uniform> sky: Sky;
+@group(0) @binding(2) var<storage, read> baseTriangles: array<BLASTriangle>;
+@group(0) @binding(3) var<storage, read_write> triangles: array<BLASTriangle>;
+@group(0) @binding(4) var<uniform> params: Params;
+
+fn displaceVertex(worldPos: vec3<f32>, normal: vec3<f32>, eid: u32) -> vec3<f32> {
+    ${vertexBody}
+}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    if (gid.x >= params.triCount) { return; }
+    let idx = params.triOffset + gid.x;
+    let base = baseTriangles[idx];
+
+    let v0 = base.v0;
+    let v1 = v0 + base.e1;
+    let v2 = v0 + base.e2;
+    let n0 = octDecode(base.n0_enc);
+    let n1 = octDecode(base.n1_enc);
+    let n2 = octDecode(base.n2_enc);
+
+    let dv0 = displaceVertex(v0, n0, 0u);
+    let dv1 = displaceVertex(v1, n1, 0u);
+    let dv2 = displaceVertex(v2, n2, 0u);
+
+    var tri: BLASTriangle;
+    tri.v0 = dv0;
+    tri.e1 = dv1 - dv0;
+    tri.e2 = dv2 - dv0;
+    tri._pad0 = 0u;
+    tri._pad1 = 0u;
+    tri._pad2 = 0u;
+    tri.n0_enc = base.n0_enc;
+    tri.n1_enc = base.n1_enc;
+    tri.n2_enc = base.n2_enc;
+    tri._pad3 = 0u;
+    triangles[idx] = tri;
+}
+`;
+}
+
+const propagateShader = /* wgsl */ `
+${BLAS_TRIANGLE_STRUCT_WGSL}
+${LEAF_FLAG_WGSL}
+
+struct PropagateParams {
+    triOffset: u32,
+    triIdOffset: u32,
+    treeNodeOffset: u32,
+    parentOffset: u32,
+    triCount: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<storage, read> triangles: array<BLASTriangle>;
+@group(0) @binding(1) var<storage, read> triIds: array<u32>;
+@group(0) @binding(2) var<storage, read_write> treeNodesRaw: array<atomic<u32>>;
+@group(0) @binding(3) var<storage, read_write> boundsFlags: array<atomic<u32>>;
+@group(0) @binding(4) var<storage, read> parentIndices: array<u32>;
+@group(0) @binding(5) var<uniform> params: PropagateParams;
+
+fn isLeaf(child: u32) -> bool {
+    return (child & LEAF_FLAG) != 0u;
+}
+
+fn leafIndex(child: u32) -> u32 {
+    return child & ~LEAF_FLAG;
+}
+
+fn getTriangleBounds(leafIdx: u32) -> array<vec3<f32>, 2> {
+    let triId = triIds[params.triIdOffset + leafIdx];
+    let tri = triangles[params.triOffset + triId];
+    let v0 = tri.v0;
+    let v1 = v0 + tri.e1;
+    let v2 = v0 + tri.e2;
+    return array<vec3<f32>, 2>(min(min(v0, v1), v2), max(max(v0, v1), v2));
+}
+
+fn nodeBase(idx: u32) -> u32 {
+    return (params.treeNodeOffset + idx) * 8u;
+}
+
+fn readChildBounds(childIdx: u32) -> array<vec3<f32>, 2> {
+    let base = nodeBase(childIdx);
+    let minX = bitcast<f32>(atomicLoad(&treeNodesRaw[base + 0u]));
+    let minY = bitcast<f32>(atomicLoad(&treeNodesRaw[base + 1u]));
+    let minZ = bitcast<f32>(atomicLoad(&treeNodesRaw[base + 2u]));
+    let maxX = bitcast<f32>(atomicLoad(&treeNodesRaw[base + 4u]));
+    let maxY = bitcast<f32>(atomicLoad(&treeNodesRaw[base + 5u]));
+    let maxZ = bitcast<f32>(atomicLoad(&treeNodesRaw[base + 6u]));
+    return array<vec3<f32>, 2>(vec3(minX, minY, minZ), vec3(maxX, maxY, maxZ));
+}
+
+fn writeBounds(nodeIdx: u32, minB: vec3<f32>, maxB: vec3<f32>) {
+    let base = nodeBase(nodeIdx);
+    atomicStore(&treeNodesRaw[base + 0u], bitcast<u32>(minB.x));
+    atomicStore(&treeNodesRaw[base + 1u], bitcast<u32>(minB.y));
+    atomicStore(&treeNodesRaw[base + 2u], bitcast<u32>(minB.z));
+    atomicStore(&treeNodesRaw[base + 4u], bitcast<u32>(maxB.x));
+    atomicStore(&treeNodesRaw[base + 5u], bitcast<u32>(maxB.y));
+    atomicStore(&treeNodesRaw[base + 6u], bitcast<u32>(maxB.z));
+}
+
+fn readLeftChild(nodeIdx: u32) -> u32 {
+    return atomicLoad(&treeNodesRaw[nodeBase(nodeIdx) + 3u]);
+}
+
+fn readRightChild(nodeIdx: u32) -> u32 {
+    return atomicLoad(&treeNodesRaw[nodeBase(nodeIdx) + 7u]);
+}
+
+fn writeLeafBounds(leafIdx: u32, minB: vec3<f32>, maxB: vec3<f32>) {
+    let n = params.triCount;
+    let leafNodeIdx = n - 1u + leafIdx;
+    writeBounds(leafNodeIdx, minB, maxB);
+}
+
+fn getParent(nodeIdx: u32, isLeafNode: bool) -> u32 {
+    let n = params.triCount;
+    if (isLeafNode) {
+        return parentIndices[params.parentOffset + nodeIdx];
+    } else {
+        return parentIndices[params.parentOffset + n + nodeIdx];
+    }
+}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let n = params.triCount;
+    let leafIdx = gid.x;
+
+    if (leafIdx >= n) { return; }
+
+    let bounds = getTriangleBounds(leafIdx);
+    writeLeafBounds(leafIdx, bounds[0], bounds[1]);
+
+    if (n <= 1u) { return; }
+
+    var current = leafIdx;
+    var isLeafNode = true;
+
+    for (var iter = 0u; iter < 64u; iter++) {
+        let parent = getParent(current, isLeafNode);
+
+        let oldFlag = atomicAdd(&boundsFlags[params.treeNodeOffset + parent], 1u);
+        if (oldFlag == 0u) { return; }
+
+        let left = readLeftChild(parent);
+        let right = readRightChild(parent);
+
+        var leftMin: vec3<f32>;
+        var leftMax: vec3<f32>;
+        var rightMin: vec3<f32>;
+        var rightMax: vec3<f32>;
+
+        if (isLeaf(left)) {
+            let leftBounds = getTriangleBounds(leafIndex(left));
+            leftMin = leftBounds[0];
+            leftMax = leftBounds[1];
+        } else {
+            let leftBounds = readChildBounds(left);
+            leftMin = leftBounds[0];
+            leftMax = leftBounds[1];
+        }
+
+        if (isLeaf(right)) {
+            let rightBounds = getTriangleBounds(leafIndex(right));
+            rightMin = rightBounds[0];
+            rightMax = rightBounds[1];
+        } else {
+            let rightBounds = readChildBounds(right);
+            rightMin = rightBounds[0];
+            rightMax = rightBounds[1];
+        }
+
+        writeBounds(parent, min(leftMin, rightMin), max(leftMax, rightMax));
+
+        current = parent;
+        isLeafNode = false;
+
+        if (parent == 0u) { break; }
+    }
+}
+`;
+
+const collapseShader = /* wgsl */ `
+${TREE_NODE_STRUCT_WGSL}
+${BVH_NODE_STRUCT_WGSL}
+${LEAF_FLAG_WGSL}
+
+const INVALID_NODE: u32 = 0xFFFFFFFFu;
+const MAX_DEPTH_ITERS: u32 = 32u;
+
+struct CollapseParams {
+    treeNodeOffset: u32,
+    parentOffset: u32,
+    nodeOffset: u32,
+    triCount: u32,
+}
+
+@group(0) @binding(0) var<storage, read> treeNodes: array<TreeNode>;
+@group(0) @binding(1) var<storage, read> parentIndices: array<u32>;
+@group(0) @binding(2) var<storage, read_write> bvhNodes: array<BVHNode>;
+@group(0) @binding(3) var<uniform> params: CollapseParams;
+
+fn isLeaf(child: u32) -> bool {
+    return (child & LEAF_FLAG) != 0u;
+}
+
+fn leafIndex(child: u32) -> u32 {
+    return child & ~LEAF_FLAG;
+}
+
+fn getDepth(nodeIdx: u32) -> u32 {
+    let n = params.triCount;
+    var depth = 0u;
+    var current = nodeIdx;
+    for (var iter = 0u; iter < MAX_DEPTH_ITERS; iter++) {
+        if (current == 0u) { break; }
+        current = parentIndices[params.parentOffset + n + current];
+        depth++;
+    }
+    return depth;
+}
+
+fn getChildBounds(child: u32) -> array<vec3<f32>, 2> {
+    let n = params.triCount;
+    if (isLeaf(child)) {
+        let leafNodeIdx = n - 1u + leafIndex(child);
+        let node = treeNodes[params.treeNodeOffset + leafNodeIdx];
+        return array<vec3<f32>, 2>(
+            vec3<f32>(node.minX, node.minY, node.minZ),
+            vec3<f32>(node.maxX, node.maxY, node.maxZ)
+        );
+    } else {
+        let node = treeNodes[params.treeNodeOffset + child];
+        return array<vec3<f32>, 2>(
+            vec3<f32>(node.minX, node.minY, node.minZ),
+            vec3<f32>(node.maxX, node.maxY, node.maxZ)
+        );
+    }
+}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let n = params.triCount;
+    let nodeIdx = gid.x;
+
+    if (n == 1u) {
+        if (nodeIdx == 0u) {
+            var out: BVHNode;
+            out.child0 = 0u | LEAF_FLAG;
+            out.child1 = INVALID_NODE;
+            out.child2 = INVALID_NODE;
+            out.child3 = INVALID_NODE;
+
+            let bounds = getChildBounds(0u | LEAF_FLAG);
+            out.c0_minX = bounds[0].x; out.c0_minY = bounds[0].y; out.c0_minZ = bounds[0].z;
+            out.c0_maxX = bounds[1].x; out.c0_maxY = bounds[1].y; out.c0_maxZ = bounds[1].z;
+            out.c1_minX = 1e30; out.c1_minY = 1e30; out.c1_minZ = 1e30;
+            out.c1_maxX = -1e30; out.c1_maxY = -1e30; out.c1_maxZ = -1e30;
+            out.c2_minX = 1e30; out.c2_minY = 1e30; out.c2_minZ = 1e30;
+            out.c2_maxX = -1e30; out.c2_maxY = -1e30; out.c2_maxZ = -1e30;
+            out.c3_minX = 1e30; out.c3_minY = 1e30; out.c3_minZ = 1e30;
+            out.c3_maxX = -1e30; out.c3_maxY = -1e30; out.c3_maxZ = -1e30;
+
+            bvhNodes[params.nodeOffset] = out;
+        }
+        return;
+    }
+
+    if (nodeIdx >= n - 1u) { return; }
+
+    let depth = getDepth(nodeIdx);
+    let node = treeNodes[params.treeNodeOffset + nodeIdx];
+    let left = node.leftChild;
+    let right = node.rightChild;
+
+    var out: BVHNode;
+    out.child0 = INVALID_NODE; out.child1 = INVALID_NODE;
+    out.child2 = INVALID_NODE; out.child3 = INVALID_NODE;
+    out.c0_minX = 1e30; out.c0_minY = 1e30; out.c0_minZ = 1e30;
+    out.c0_maxX = -1e30; out.c0_maxY = -1e30; out.c0_maxZ = -1e30;
+    out.c1_minX = 1e30; out.c1_minY = 1e30; out.c1_minZ = 1e30;
+    out.c1_maxX = -1e30; out.c1_maxY = -1e30; out.c1_maxZ = -1e30;
+    out.c2_minX = 1e30; out.c2_minY = 1e30; out.c2_minZ = 1e30;
+    out.c2_maxX = -1e30; out.c2_maxY = -1e30; out.c2_maxZ = -1e30;
+    out.c3_minX = 1e30; out.c3_minY = 1e30; out.c3_minZ = 1e30;
+    out.c3_maxX = -1e30; out.c3_maxY = -1e30; out.c3_maxZ = -1e30;
+
+    if ((depth & 1u) != 0u) {
+        out.child0 = left;
+        let bounds0 = getChildBounds(left);
+        out.c0_minX = bounds0[0].x; out.c0_minY = bounds0[0].y; out.c0_minZ = bounds0[0].z;
+        out.c0_maxX = bounds0[1].x; out.c0_maxY = bounds0[1].y; out.c0_maxZ = bounds0[1].z;
+
+        out.child1 = right;
+        let bounds1 = getChildBounds(right);
+        out.c1_minX = bounds1[0].x; out.c1_minY = bounds1[0].y; out.c1_minZ = bounds1[0].z;
+        out.c1_maxX = bounds1[1].x; out.c1_maxY = bounds1[1].y; out.c1_maxZ = bounds1[1].z;
+
+        bvhNodes[params.nodeOffset + nodeIdx] = out;
+        return;
+    }
+
+    if (isLeaf(left)) {
+        out.child0 = left;
+        let bounds = getChildBounds(left);
+        out.c0_minX = bounds[0].x; out.c0_minY = bounds[0].y; out.c0_minZ = bounds[0].z;
+        out.c0_maxX = bounds[1].x; out.c0_maxY = bounds[1].y; out.c0_maxZ = bounds[1].z;
+    } else {
+        let leftNode = treeNodes[params.treeNodeOffset + left];
+        let ll = leftNode.leftChild;
+        let lr = leftNode.rightChild;
+
+        out.child0 = ll;
+        let bounds0 = getChildBounds(ll);
+        out.c0_minX = bounds0[0].x; out.c0_minY = bounds0[0].y; out.c0_minZ = bounds0[0].z;
+        out.c0_maxX = bounds0[1].x; out.c0_maxY = bounds0[1].y; out.c0_maxZ = bounds0[1].z;
+
+        out.child1 = lr;
+        let bounds1 = getChildBounds(lr);
+        out.c1_minX = bounds1[0].x; out.c1_minY = bounds1[0].y; out.c1_minZ = bounds1[0].z;
+        out.c1_maxX = bounds1[1].x; out.c1_maxY = bounds1[1].y; out.c1_maxZ = bounds1[1].z;
+    }
+
+    if (isLeaf(right)) {
+        out.child2 = right;
+        let bounds = getChildBounds(right);
+        out.c2_minX = bounds[0].x; out.c2_minY = bounds[0].y; out.c2_minZ = bounds[0].z;
+        out.c2_maxX = bounds[1].x; out.c2_maxY = bounds[1].y; out.c2_maxZ = bounds[1].z;
+    } else {
+        let rightNode = treeNodes[params.treeNodeOffset + right];
+        let rl = rightNode.leftChild;
+        let rr = rightNode.rightChild;
+
+        out.child2 = rl;
+        let bounds2 = getChildBounds(rl);
+        out.c2_minX = bounds2[0].x; out.c2_minY = bounds2[0].y; out.c2_minZ = bounds2[0].z;
+        out.c2_maxX = bounds2[1].x; out.c2_maxY = bounds2[1].y; out.c2_maxZ = bounds2[1].z;
+
+        out.child3 = rr;
+        let bounds3 = getChildBounds(rr);
+        out.c3_minX = bounds3[0].x; out.c3_minY = bounds3[0].y; out.c3_minZ = bounds3[0].z;
+        out.c3_maxX = bounds3[1].x; out.c3_maxY = bounds3[1].y; out.c3_maxZ = bounds3[1].z;
+    }
+
+    bvhNodes[params.nodeOffset + nodeIdx] = out;
+}
+`;
+
+const aabbShader = /* wgsl */ `
+struct ShapeAABB {
+    minX: f32, minY: f32, minZ: f32, _pad0: u32,
+    maxX: f32, maxY: f32, maxZ: f32, _pad1: u32,
+}
+
+${TREE_NODE_STRUCT_WGSL}
+
+struct AABBParams {
+    treeNodeOffset: u32,
+    shapeId: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<storage, read> treeNodes: array<TreeNode>;
+@group(0) @binding(1) var<storage, read_write> shapeAABBs: array<ShapeAABB>;
+@group(0) @binding(2) var<uniform> params: AABBParams;
+
+@compute @workgroup_size(1)
+fn main() {
+    let root = treeNodes[params.treeNodeOffset];
+    var aabb: ShapeAABB;
+    aabb.minX = root.minX;
+    aabb.minY = root.minY;
+    aabb.minZ = root.minZ;
+    aabb._pad0 = 0u;
+    aabb.maxX = root.maxX;
+    aabb.maxY = root.maxY;
+    aabb.maxZ = root.maxZ;
+    aabb._pad1 = 0u;
+    shapeAABBs[params.shapeId] = aabb;
+}
+`;
+
+interface DisplacementPipeline {
+    pipeline: GPUComputePipeline;
+    bindGroup: GPUBindGroup;
+}
+
+export interface BLASRefitConfig {
+    getBlasAtlas: () => BLASAtlas;
+    scene: GPUBuffer;
+    sky: GPUBuffer;
+    getDynamicShapes: () => Map<number, SurfaceData>;
+    getRaytracing: () => boolean;
+}
+
+interface RefitGPU {
+    propagatePipeline: GPUComputePipeline;
+    propagateLayout: GPUBindGroupLayout;
+    collapsePipeline: GPUComputePipeline;
+    collapseLayout: GPUBindGroupLayout;
+    aabbPipeline: GPUComputePipeline;
+    aabbLayout: GPUBindGroupLayout;
+    displacementPipelines: Map<string, DisplacementPipeline>;
+    displacementBindGroupLayout: GPUBindGroupLayout;
+    displacementPipelineLayout: GPUPipelineLayout;
+    displaceParams: GPUBuffer;
+    propagateParams: GPUBuffer;
+    collapseParams: GPUBuffer;
+    aabbParams: GPUBuffer;
+    propagateBindGroup: GPUBindGroup | null;
+    collapseBindGroup: GPUBindGroup | null;
+    aabbBindGroup: GPUBindGroup | null;
+    cachedAtlas: BLASAtlas | null;
+}
+
+export function createBLASRefitNode(config: BLASRefitConfig): ComputeNode {
+    let gpu: RefitGPU | null = null;
+
+    return {
+        id: "blas-refit",
+        inputs: [],
+        outputs: [{ id: "blas-nodes", access: "write" }],
+
+        async prepare(device: GPUDevice) {
+            const createParams = (label: string, size: number) =>
+                device.createBuffer({
+                    label,
+                    size,
+                    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+                });
+
+            const displaceParams = createParams("blas-refit-displace-params", 16);
+            const propagateParams = createParams("blas-refit-propagate-params", 32);
+            const collapseParams = createParams("blas-refit-collapse-params", 16);
+            const aabbParams = createParams("blas-refit-aabb-params", 16);
+
+            const displacementBindGroupLayout = device.createBindGroupLayout({
+                entries: [
+                    { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+                    { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+                    {
+                        binding: 2,
+                        visibility: GPUShaderStage.COMPUTE,
+                        buffer: { type: "read-only-storage" },
+                    },
+                    { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+                    { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+                ],
+            });
+            const displacementPipelineLayout = device.createPipelineLayout({
+                bindGroupLayouts: [displacementBindGroupLayout],
+            });
+
+            const propagateModule = device.createShaderModule({ code: propagateShader });
+            const propagatePipeline = await device.createComputePipelineAsync({
+                layout: "auto",
+                compute: { module: propagateModule, entryPoint: "main" },
+            });
+
+            const collapseModule = device.createShaderModule({ code: collapseShader });
+            const collapsePipeline = await device.createComputePipelineAsync({
+                layout: "auto",
+                compute: { module: collapseModule, entryPoint: "main" },
+            });
+
+            const aabbModule = device.createShaderModule({ code: aabbShader });
+            const aabbPipeline = await device.createComputePipelineAsync({
+                layout: "auto",
+                compute: { module: aabbModule, entryPoint: "main" },
+            });
+
+            gpu = {
+                propagatePipeline,
+                propagateLayout: propagatePipeline.getBindGroupLayout(0),
+                collapsePipeline,
+                collapseLayout: collapsePipeline.getBindGroupLayout(0),
+                aabbPipeline,
+                aabbLayout: aabbPipeline.getBindGroupLayout(0),
+                displacementPipelines: new Map(),
+                displacementBindGroupLayout,
+                displacementPipelineLayout,
+                displaceParams,
+                propagateParams,
+                collapseParams,
+                aabbParams,
+                propagateBindGroup: null,
+                collapseBindGroup: null,
+                aabbBindGroup: null,
+                cachedAtlas: null,
+            };
+        },
+
+        execute(ctx: ExecutionContext) {
+            if (!gpu) return;
+            if (!config.getRaytracing()) return;
+
+            const dynamicShapes = config.getDynamicShapes();
+            if (dynamicShapes.size === 0) return;
+
+            const { device, encoder } = ctx;
+            const atlas = config.getBlasAtlas();
+
+            if (atlas !== gpu.cachedAtlas) {
+                gpu.displacementPipelines.clear();
+
+                gpu.propagateBindGroup = device.createBindGroup({
+                    layout: gpu.propagateLayout,
+                    entries: [
+                        { binding: 0, resource: { buffer: atlas.trianglesBuffer } },
+                        { binding: 1, resource: { buffer: atlas.triIdsBuffer } },
+                        { binding: 2, resource: { buffer: atlas.treeNodesBuffer } },
+                        { binding: 3, resource: { buffer: atlas.boundsFlagsBuffer } },
+                        { binding: 4, resource: { buffer: atlas.parentIndicesBuffer } },
+                        { binding: 5, resource: { buffer: gpu.propagateParams } },
+                    ],
+                });
+
+                gpu.collapseBindGroup = device.createBindGroup({
+                    layout: gpu.collapseLayout,
+                    entries: [
+                        { binding: 0, resource: { buffer: atlas.treeNodesBuffer } },
+                        { binding: 1, resource: { buffer: atlas.parentIndicesBuffer } },
+                        { binding: 2, resource: { buffer: atlas.nodesBuffer } },
+                        { binding: 3, resource: { buffer: gpu.collapseParams } },
+                    ],
+                });
+
+                gpu.aabbBindGroup = device.createBindGroup({
+                    layout: gpu.aabbLayout,
+                    entries: [
+                        { binding: 0, resource: { buffer: atlas.treeNodesBuffer } },
+                        { binding: 1, resource: { buffer: atlas.shapeAABBs } },
+                        { binding: 2, resource: { buffer: gpu.aabbParams } },
+                    ],
+                });
+
+                gpu.cachedAtlas = atlas;
+            }
+
+            for (const [shapeId, surfaceData] of dynamicShapes) {
+                if (!surfaceData.vertex) continue;
+
+                const meta = atlas.metas[shapeId];
+                if (!meta || meta.triCount === 0) continue;
+
+                const vertexKey = surfaceData.vertex;
+                if (!gpu.displacementPipelines.has(vertexKey)) {
+                    const code = compileDisplacementShader(vertexKey);
+                    const module = device.createShaderModule({ code });
+                    const pipeline = device.createComputePipeline({
+                        layout: gpu.displacementPipelineLayout,
+                        compute: { module, entryPoint: "main" },
+                    });
+
+                    const bindGroup = device.createBindGroup({
+                        layout: gpu.displacementBindGroupLayout,
+                        entries: [
+                            { binding: 0, resource: { buffer: config.scene } },
+                            { binding: 1, resource: { buffer: config.sky } },
+                            { binding: 2, resource: { buffer: atlas.baseTrianglesBuffer } },
+                            { binding: 3, resource: { buffer: atlas.trianglesBuffer } },
+                            { binding: 4, resource: { buffer: gpu.displaceParams } },
+                        ],
+                    });
+
+                    gpu.displacementPipelines.set(vertexKey, { pipeline, bindGroup });
+                }
+
+                const displacement = gpu.displacementPipelines.get(vertexKey)!;
+
+                device.queue.writeBuffer(
+                    gpu.displaceParams,
+                    0,
+                    new Uint32Array([meta.triOffset, meta.triCount, 0, 0])
+                );
+
+                const displacePass = encoder.beginComputePass();
+                displacePass.setPipeline(displacement.pipeline);
+                displacePass.setBindGroup(0, displacement.bindGroup);
+                displacePass.dispatchWorkgroups(Math.ceil(meta.triCount / WORKGROUP_SIZE));
+                displacePass.end();
+
+                if (meta.triCount <= 1) {
+                    device.queue.writeBuffer(
+                        gpu.aabbParams,
+                        0,
+                        new Uint32Array([meta.treeNodeOffset, shapeId, 0, 0])
+                    );
+
+                    const aabbPass = encoder.beginComputePass();
+                    aabbPass.setPipeline(gpu.aabbPipeline);
+                    aabbPass.setBindGroup(0, gpu.aabbBindGroup!);
+                    aabbPass.dispatchWorkgroups(1);
+                    aabbPass.end();
+                    continue;
+                }
+
+                const n = meta.triCount;
+                const internalNodes = n - 1;
+
+                const flagsByteOffset = meta.treeNodeOffset * 4;
+                const flagsByteSize = internalNodes * 4;
+                encoder.clearBuffer(atlas.boundsFlagsBuffer, flagsByteOffset, flagsByteSize);
+
+                device.queue.writeBuffer(
+                    gpu.propagateParams,
+                    0,
+                    new Uint32Array([
+                        meta.triOffset,
+                        meta.triIdOffset,
+                        meta.treeNodeOffset,
+                        meta.parentOffset,
+                        meta.triCount,
+                        0,
+                        0,
+                        0,
+                    ])
+                );
+
+                const propagatePass = encoder.beginComputePass();
+                propagatePass.setPipeline(gpu.propagatePipeline);
+                propagatePass.setBindGroup(0, gpu.propagateBindGroup!);
+                propagatePass.dispatchWorkgroups(Math.ceil(n / WORKGROUP_SIZE));
+                propagatePass.end();
+
+                device.queue.writeBuffer(
+                    gpu.collapseParams,
+                    0,
+                    new Uint32Array([
+                        meta.treeNodeOffset,
+                        meta.parentOffset,
+                        meta.nodeOffset,
+                        meta.triCount,
+                    ])
+                );
+
+                const collapsePass = encoder.beginComputePass();
+                collapsePass.setPipeline(gpu.collapsePipeline);
+                collapsePass.setBindGroup(0, gpu.collapseBindGroup!);
+                collapsePass.dispatchWorkgroups(Math.ceil(internalNodes / WORKGROUP_SIZE));
+                collapsePass.end();
+
+                device.queue.writeBuffer(
+                    gpu.aabbParams,
+                    0,
+                    new Uint32Array([meta.treeNodeOffset, shapeId, 0, 0])
+                );
+
+                const aabbPass = encoder.beginComputePass();
+                aabbPass.setPipeline(gpu.aabbPipeline);
+                aabbPass.setBindGroup(0, gpu.aabbBindGroup!);
+                aabbPass.dispatchWorkgroups(1);
+                aabbPass.end();
+            }
+        },
+    };
+}