@multiplekex/shallot 0.1.12 → 0.2.0

package/src/standard/render/bvh/tlas.ts

@@ -0,0 +1,886 @@
+import { MAX_ENTITIES } from "../../../core";
+import type { ComputeNode, ExecutionContext } from "../../compute";
+import { createRadixSort } from "./radix";
+import {
+    TREE_NODE_STRUCT_WGSL,
+    BVH_NODE_STRUCT_WGSL,
+    LEAF_FLAG_WGSL,
+    TREE_NODE_SIZE,
+    BVH_NODE_SIZE,
+} from "./structs";
+
+const WORKGROUP_SIZE = 256;
+const MAX_TREE_DEPTH = Math.ceil(Math.log2(MAX_ENTITIES)) + 1;
+
+export interface TLASBuffers {
+    treeNodes: GPUBuffer;
+    bvhNodes: GPUBuffer;
+    mortonCodes: GPUBuffer;
+    instanceIds: GPUBuffer;
+    entityIds: GPUBuffer;
+    sceneBounds: GPUBuffer;
+    parentIndices: GPUBuffer;
+    boundsFlags: GPUBuffer;
+}
+
+export function createTLASBuffers(device: GPUDevice): TLASBuffers {
+    return {
+        treeNodes: device.createBuffer({
+            label: "tlas-tree-nodes",
+            size: 2 * MAX_ENTITIES * TREE_NODE_SIZE,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+        }),
+        bvhNodes: device.createBuffer({
+            label: "tlas-bvh-nodes",
+            size: MAX_ENTITIES * BVH_NODE_SIZE,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
+        }),
+        mortonCodes: device.createBuffer({
+            label: "tlas-morton-codes",
+            size: MAX_ENTITIES * 4,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+        }),
+        instanceIds: device.createBuffer({
+            label: "tlas-instance-ids",
+            size: MAX_ENTITIES * 4,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+        }),
+        entityIds: device.createBuffer({
+            label: "tlas-entity-ids",
+            size: MAX_ENTITIES * 4,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+        }),
+        sceneBounds: device.createBuffer({
+            label: "tlas-scene-bounds",
+            size: 32,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
+        }),
+        parentIndices: device.createBuffer({
+            label: "tlas-parent-indices",
+            size: 2 * MAX_ENTITIES * 4,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+        }),
+        boundsFlags: device.createBuffer({
+            label: "tlas-bounds-flags",
+            size: MAX_ENTITIES * 4,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
+        }),
+    };
+}
+
+const INSTANCE_AABB_STRUCT_WGSL = /* wgsl */ `
+struct InstanceAABB {
+    minX: f32,
+    minY: f32,
+    minZ: f32,
+    _pad0: u32,
+    maxX: f32,
+    maxY: f32,
+    maxZ: f32,
+    _pad1: u32,
+}`;
+
+const SCENE_BOUNDS_STRUCT_WGSL = /* wgsl */ `
+struct SceneBounds {
+    minX: atomic<i32>,
+    minY: atomic<i32>,
+    minZ: atomic<i32>,
+    _pad0: u32,
+    maxX: atomic<i32>,
+    maxY: atomic<i32>,
+    maxZ: atomic<i32>,
+    _pad1: u32,
+}`;
+
+const SCENE_BOUNDS_READ_STRUCT_WGSL = /* wgsl */ `
+struct SceneBounds {
+    minX: i32,
+    minY: i32,
+    minZ: i32,
+    _pad0: u32,
+    maxX: i32,
+    maxY: i32,
+    maxZ: i32,
+    _pad1: u32,
+}`;
+
+const FLOAT_INT_CONVERSION_WGSL = /* wgsl */ `
+fn floatToSortableInt(f: f32) -> i32 {
+    let bits = bitcast<i32>(f);
+    let mask = (bits >> 31) & 0x7FFFFFFF;
+    return bits ^ mask;
+}
+
+fn sortableIntToFloat(i: i32) -> f32 {
+    let mask = (i >> 31) & 0x7FFFFFFF;
+    return bitcast<f32>(i ^ mask);
+}`;
+
+const MORTON_CODE_WGSL = /* wgsl */ `
+fn expandBits(v: u32) -> u32 {
+    var x = v & 0x3ffu;
+    x = (x | (x << 16u)) & 0x030000ffu;
+    x = (x | (x << 8u)) & 0x0300f00fu;
+    x = (x | (x << 4u)) & 0x030c30c3u;
+    x = (x | (x << 2u)) & 0x09249249u;
+    return x;
+}
+
+fn mortonCode(x: u32, y: u32, z: u32) -> u32 {
+    return (expandBits(x) << 2u) | (expandBits(y) << 1u) | expandBits(z);
+}`;
+
+const CLZ_WGSL = /* wgsl */ `
+fn clz(x: u32) -> u32 {
+    if (x == 0u) { return 32u; }
+    var n = 0u;
+    var v = x;
+    if ((v & 0xffff0000u) == 0u) { n += 16u; v <<= 16u; }
+    if ((v & 0xff000000u) == 0u) { n += 8u; v <<= 8u; }
+    if ((v & 0xf0000000u) == 0u) { n += 4u; v <<= 4u; }
+    if ((v & 0xc0000000u) == 0u) { n += 2u; v <<= 2u; }
+    if ((v & 0x80000000u) == 0u) { n += 1u; }
+    return n;
+}`;
+
+const LEAF_FUNCTIONS_WGSL = /* wgsl */ `
+fn isLeaf(child: u32) -> bool {
+    return (child & LEAF_FLAG) != 0u;
+}
+
+fn leafIndex(child: u32) -> u32 {
+    return child & ~LEAF_FLAG;
+}`;
+
+const boundsShader = /* wgsl */ `
+${INSTANCE_AABB_STRUCT_WGSL}
+${SCENE_BOUNDS_STRUCT_WGSL}
+
+@group(0) @binding(0) var<storage, read> instanceAABBs: array<InstanceAABB>;
+@group(0) @binding(1) var<storage, read> instanceCount: array<u32>;
+@group(0) @binding(2) var<storage, read_write> sceneBounds: SceneBounds;
+@group(0) @binding(3) var<storage, read> entityIds: array<u32>;
+
+var<workgroup> sharedMin: array<vec3<f32>, ${WORKGROUP_SIZE}>;
+var<workgroup> sharedMax: array<vec3<f32>, ${WORKGROUP_SIZE}>;
+
+${FLOAT_INT_CONVERSION_WGSL}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+    let count = instanceCount[0];
+    let tid = gid.x;
+    let localId = lid.x;
+
+    var localMin = vec3<f32>(1e30, 1e30, 1e30);
+    var localMax = vec3<f32>(-1e30, -1e30, -1e30);
+
+    if (tid < count) {
+        let eid = entityIds[tid];
+        let aabb = instanceAABBs[eid];
+        localMin = vec3<f32>(aabb.minX, aabb.minY, aabb.minZ);
+        localMax = vec3<f32>(aabb.maxX, aabb.maxY, aabb.maxZ);
+    }
+
+    sharedMin[localId] = localMin;
+    sharedMax[localId] = localMax;
+    workgroupBarrier();
+
+    for (var stride = ${WORKGROUP_SIZE}u / 2u; stride > 0u; stride >>= 1u) {
+        if (localId < stride) {
+            sharedMin[localId] = min(sharedMin[localId], sharedMin[localId + stride]);
+            sharedMax[localId] = max(sharedMax[localId], sharedMax[localId + stride]);
+        }
+        workgroupBarrier();
+    }
+
+    if (localId == 0u) {
+        let wgMin = sharedMin[0];
+        let wgMax = sharedMax[0];
+
+        atomicMin(&sceneBounds.minX, floatToSortableInt(wgMin.x));
+        atomicMin(&sceneBounds.minY, floatToSortableInt(wgMin.y));
+        atomicMin(&sceneBounds.minZ, floatToSortableInt(wgMin.z));
+        atomicMax(&sceneBounds.maxX, floatToSortableInt(wgMax.x));
+        atomicMax(&sceneBounds.maxY, floatToSortableInt(wgMax.y));
+        atomicMax(&sceneBounds.maxZ, floatToSortableInt(wgMax.z));
+    }
+}
+`;
+
+const mortonShader = /* wgsl */ `
+${INSTANCE_AABB_STRUCT_WGSL}
+${SCENE_BOUNDS_READ_STRUCT_WGSL}
+
+@group(0) @binding(0) var<storage, read> instanceAABBs: array<InstanceAABB>;
+@group(0) @binding(1) var<storage, read> instanceCount: array<u32>;
+@group(0) @binding(2) var<storage, read> sceneBounds: SceneBounds;
+@group(0) @binding(3) var<storage, read_write> mortonCodes: array<u32>;
+@group(0) @binding(4) var<storage, read_write> instanceIds: array<u32>;
+@group(0) @binding(5) var<storage, read> entityIds: array<u32>;
+
+${FLOAT_INT_CONVERSION_WGSL}
+${MORTON_CODE_WGSL}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let tid = gid.x;
+    if (tid >= ${MAX_ENTITIES}u) { return; }
+
+    let count = instanceCount[0];
+    if (tid >= count) {
+        mortonCodes[tid] = 0xFFFFFFFFu;
+        instanceIds[tid] = 0u;
+        return;
+    }
+
+    let eid = entityIds[tid];
+    let aabb = instanceAABBs[eid];
+    let centroid = vec3<f32>(
+        (aabb.minX + aabb.maxX) * 0.5,
+        (aabb.minY + aabb.maxY) * 0.5,
+        (aabb.minZ + aabb.maxZ) * 0.5
+    );
+
+    let boundsMin = vec3<f32>(
+        sortableIntToFloat(sceneBounds.minX),
+        sortableIntToFloat(sceneBounds.minY),
+        sortableIntToFloat(sceneBounds.minZ)
+    );
+    let boundsMax = vec3<f32>(
+        sortableIntToFloat(sceneBounds.maxX),
+        sortableIntToFloat(sceneBounds.maxY),
+        sortableIntToFloat(sceneBounds.maxZ)
+    );
+
+    let size = boundsMax - boundsMin;
+    let safeSize = max(size, vec3<f32>(1e-6, 1e-6, 1e-6));
+
+    let normalized = (centroid - boundsMin) / safeSize;
+    let clamped = clamp(normalized, vec3<f32>(0.0), vec3<f32>(1.0));
+
+    let quantized = vec3<u32>(clamped * 1023.0);
+
+    mortonCodes[tid] = mortonCode(quantized.x, quantized.y, quantized.z);
+    instanceIds[tid] = eid;
+}
+`;
+
+const treeShader = /* wgsl */ `
+${TREE_NODE_STRUCT_WGSL}
+${LEAF_FLAG_WGSL}
+
+@group(0) @binding(0) var<storage, read> mortonCodes: array<u32>;
+@group(0) @binding(1) var<storage, read> instanceCount: array<u32>;
+@group(0) @binding(2) var<storage, read_write> treeNodes: array<TreeNode>;
+@group(0) @binding(3) var<storage, read_write> parentIndices: array<u32>;
+
+${CLZ_WGSL}
+
+fn delta(i: i32, j: i32, n: i32) -> i32 {
+    if (j < 0 || j >= n) {
+        return -1;
+    }
+    let codeI = mortonCodes[i];
+    let codeJ = mortonCodes[j];
+    if (codeI == codeJ) {
+        return i32(clz(u32(i) ^ u32(j))) + 32;
+    }
+    return i32(clz(codeI ^ codeJ));
+}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let n = i32(instanceCount[0]);
+    let i = i32(gid.x);
+
+    if (i >= n - 1) {
+        return;
+    }
+
+    var first: i32;
+    var last: i32;
+
+    if (i == 0) {
+        first = 0;
+        last = n - 1;
+    } else {
+        let d = select(-1, 1, delta(i, i + 1, n) > delta(i, i - 1, n));
+
+        let deltaMin = delta(i, i - d, n);
+
+        var lmax = 2;
+        for (var iter = 0; iter < ${MAX_TREE_DEPTH}; iter++) {
+            if (delta(i, i + lmax * d, n) <= deltaMin) { break; }
+            lmax *= 2;
+        }
+
+        var l = 0;
+        var t = lmax / 2;
+        for (var iter2 = 0; iter2 < ${MAX_TREE_DEPTH}; iter2++) {
+            if (t < 1) { break; }
+            if (delta(i, i + (l + t) * d, n) > deltaMin) {
+                l += t;
+            }
+            t /= 2;
+        }
+
+        let j = i + l * d;
+        first = min(i, j);
+        last = max(i, j);
+    }
+
+    let firstCode = mortonCodes[first];
+    let lastCode = mortonCodes[last];
+
+    var gamma: i32;
+    if (firstCode == lastCode) {
+        gamma = (first + last) / 2;
+    } else {
+        let deltaNode = i32(clz(firstCode ^ lastCode));
+
+        var split = first;
+        var stride = last - first;
+
+        for (var iter3 = 0; iter3 < ${MAX_TREE_DEPTH}; iter3++) {
+            stride = (stride + 1) / 2;
+            let middle = split + stride;
+
+            if (middle < last) {
+                let splitCode = mortonCodes[middle];
+                let splitDelta = i32(clz(firstCode ^ splitCode));
+
+                if (splitDelta > deltaNode) {
+                    split = middle;
+                }
+            }
+
+            if (stride <= 1) {
+                break;
+            }
+        }
+
+        gamma = split;
+    }
+
+    let leftIsLeaf = first == gamma;
+    let rightIsLeaf = last == gamma + 1;
+
+    var node: TreeNode;
+    node.minX = 1e30;
+    node.minY = 1e30;
+    node.minZ = 1e30;
+    node.maxX = -1e30;
+    node.maxY = -1e30;
+    node.maxZ = -1e30;
+
+    if (leftIsLeaf) {
+        node.leftChild = u32(gamma) | LEAF_FLAG;
+        parentIndices[u32(gamma)] = u32(i);
+    } else {
+        node.leftChild = u32(gamma);
+        parentIndices[u32(n) + u32(gamma)] = u32(i);
+    }
+
+    if (rightIsLeaf) {
+        node.rightChild = u32(gamma + 1) | LEAF_FLAG;
+        parentIndices[u32(gamma + 1)] = u32(i);
+    } else {
+        node.rightChild = u32(gamma + 1);
+        parentIndices[u32(n) + u32(gamma + 1)] = u32(i);
+    }
+
+    treeNodes[i] = node;
+}
+`;
+
+const propagateShader = /* wgsl */ `
+${INSTANCE_AABB_STRUCT_WGSL}
+${LEAF_FLAG_WGSL}
+
+const BOUNDS_SENTINEL: u32 = 0x7f800000u;
+
+@group(0) @binding(0) var<storage, read> instanceAABBs: array<InstanceAABB>;
+@group(0) @binding(1) var<storage, read> instanceIds: array<u32>;
+@group(0) @binding(2) var<storage, read> instanceCount: array<u32>;
+@group(0) @binding(3) var<storage, read_write> bvhNodesRaw: array<atomic<u32>>;
+@group(0) @binding(4) var<storage, read_write> boundsFlags: array<atomic<u32>>;
+@group(0) @binding(5) var<storage, read> parentIndices: array<u32>;
+
+${LEAF_FUNCTIONS_WGSL}
+
+fn getInstanceBounds(leafIdx: u32) -> array<vec3<f32>, 2> {
+    let eid = instanceIds[leafIdx];
+    let aabb = instanceAABBs[eid];
+    return array<vec3<f32>, 2>(
+        vec3<f32>(aabb.minX, aabb.minY, aabb.minZ),
+        vec3<f32>(aabb.maxX, aabb.maxY, aabb.maxZ)
+    );
+}
+
+fn getParent(nodeIdx: u32, isLeafNode: bool, n: u32) -> u32 {
+    if (isLeafNode) {
+        return parentIndices[nodeIdx];
+    } else {
+        return parentIndices[n + nodeIdx];
+    }
+}
+
+fn nodeBase(idx: u32) -> u32 {
+    return idx * 8u;
+}
+
+fn readChildBounds(childIdx: u32) -> array<vec3<f32>, 2> {
+    let base = nodeBase(childIdx);
+    let minX = bitcast<f32>(atomicLoad(&bvhNodesRaw[base + 0u]));
+    let minY = bitcast<f32>(atomicLoad(&bvhNodesRaw[base + 1u]));
+    let minZ = bitcast<f32>(atomicLoad(&bvhNodesRaw[base + 2u]));
+    let maxX = bitcast<f32>(atomicLoad(&bvhNodesRaw[base + 4u]));
+    let maxY = bitcast<f32>(atomicLoad(&bvhNodesRaw[base + 5u]));
+    let maxZ = bitcast<f32>(atomicLoad(&bvhNodesRaw[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(&bvhNodesRaw[base + 0u], bitcast<u32>(minB.x));
+    atomicStore(&bvhNodesRaw[base + 1u], bitcast<u32>(minB.y));
+    atomicStore(&bvhNodesRaw[base + 2u], bitcast<u32>(minB.z));
+    atomicStore(&bvhNodesRaw[base + 4u], bitcast<u32>(maxB.x));
+    atomicStore(&bvhNodesRaw[base + 5u], bitcast<u32>(maxB.y));
+    atomicStore(&bvhNodesRaw[base + 6u], bitcast<u32>(maxB.z));
+}
+
+fn readLeftChild(nodeIdx: u32) -> u32 {
+    return atomicLoad(&bvhNodesRaw[nodeBase(nodeIdx) + 3u]);
+}
+
+fn readRightChild(nodeIdx: u32) -> u32 {
+    return atomicLoad(&bvhNodesRaw[nodeBase(nodeIdx) + 7u]);
+}
+
+fn writeLeafBounds(leafIdx: u32, n: u32, minB: vec3<f32>, maxB: vec3<f32>) {
+    let leafNodeIdx = n - 1u + leafIdx;
+    let base = leafNodeIdx * 8u;
+    atomicStore(&bvhNodesRaw[base + 0u], bitcast<u32>(minB.x));
+    atomicStore(&bvhNodesRaw[base + 1u], bitcast<u32>(minB.y));
+    atomicStore(&bvhNodesRaw[base + 2u], bitcast<u32>(minB.z));
+    atomicStore(&bvhNodesRaw[base + 4u], bitcast<u32>(maxB.x));
+    atomicStore(&bvhNodesRaw[base + 5u], bitcast<u32>(maxB.y));
+    atomicStore(&bvhNodesRaw[base + 6u], bitcast<u32>(maxB.z));
+}
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let n = instanceCount[0];
+    let leafIdx = gid.x;
+
+    if (leafIdx >= n) {
+        return;
+    }
+
+    let bounds = getInstanceBounds(leafIdx);
+    writeLeafBounds(leafIdx, n, bounds[0], bounds[1]);
+
+    var current = leafIdx;
+    var isLeafNode = true;
+
+    for (var iter = 0u; iter < 64u; iter++) {
+        let parent = getParent(current, isLeafNode, n);
+
+        let oldFlag = atomicAdd(&boundsFlags[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 = getInstanceBounds(leafIndex(left));
+            leftMin = leftBounds[0];
+            leftMax = leftBounds[1];
+        } else {
+            let leftBounds = readChildBounds(left);
+            leftMin = leftBounds[0];
+            leftMax = leftBounds[1];
+        }
+
+        if (isLeaf(right)) {
+            let rightBounds = getInstanceBounds(leafIndex(right));
+            rightMin = rightBounds[0];
+            rightMax = rightBounds[1];
+        } else {
+            let rightBounds = readChildBounds(right);
+            rightMin = rightBounds[0];
+            rightMax = rightBounds[1];
+        }
+
+        let newMin = min(leftMin, rightMin);
+        let newMax = max(leftMax, rightMax);
+
+        writeBounds(parent, newMin, newMax);
+
+        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;
+
+@group(0) @binding(0) var<storage, read> treeNodes: array<TreeNode>;
+@group(0) @binding(1) var<storage, read> instanceCount: array<u32>;
+@group(0) @binding(2) var<storage, read> parentIndices: array<u32>;
+@group(0) @binding(3) var<storage, read_write> bvhNodes: array<BVHNode>;
+
+fn isLeaf(child: u32) -> bool {
+    return (child & LEAF_FLAG) != 0u;
+}
+
+fn leafIndex(child: u32) -> u32 {
+    return child & ~LEAF_FLAG;
+}
+
+fn getDepth(nodeIdx: u32, n: u32) -> u32 {
+    var depth = 0u;
+    var current = nodeIdx;
+    for (var iter = 0u; iter < ${MAX_TREE_DEPTH}u; iter++) {
+        if (current == 0u) { break; }
+        current = parentIndices[n + current];
+        depth++;
+    }
+    return depth;
+}
+
+fn getChildBounds(child: u32, n: u32) -> array<vec3<f32>, 2> {
+    if (isLeaf(child)) {
+        let leafNodeIdx = n - 1u + leafIndex(child);
+        let node = treeNodes[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[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 = instanceCount[0];
+    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, n);
+            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[0] = out;
+        }
+        return;
+    }
+
+    if (nodeIdx >= n - 1u) {
+        return;
+    }
+
+    let depth = getDepth(nodeIdx, n);
+    let node = treeNodes[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, n);
+        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, n);
+        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[nodeIdx] = out;
+        return;
+    }
+
+    if (isLeaf(left)) {
+        out.child0 = left;
+        let bounds = getChildBounds(left, n);
+        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[left];
+        let ll = leftNode.leftChild;
+        let lr = leftNode.rightChild;
+
+        out.child0 = ll;
+        let bounds0 = getChildBounds(ll, n);
+        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, n);
+        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, n);
+        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[right];
+        let rl = rightNode.leftChild;
+        let rr = rightNode.rightChild;
+
+        out.child2 = rl;
+        let bounds2 = getChildBounds(rl, n);
+        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, n);
+        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[nodeIdx] = out;
+}
+`;
+
+interface TLASPipelines {
+    bounds: GPUComputePipeline;
+    morton: GPUComputePipeline;
+    tree: GPUComputePipeline;
+    propagate: GPUComputePipeline;
+    collapse: GPUComputePipeline;
+}
+
+interface TLASBindGroups {
+    bounds: GPUBindGroup;
+    morton: GPUBindGroup;
+    tree: GPUBindGroup;
+    propagate: GPUBindGroup;
+    collapse: GPUBindGroup;
+}
+
+export interface TLASConfig {
+    instanceAABBs: GPUBuffer;
+    instanceCount: GPUBuffer;
+    tlas: TLASBuffers;
+    getEntityCount: () => number;
+}
+
+export function createTLASNode(config: TLASConfig): ComputeNode {
+    let pipelines: TLASPipelines | null = null;
+    let bindGroups: TLASBindGroups | null = null;
+    let radixSort: Awaited<ReturnType<typeof createRadixSort>> | null = null;
+
+    return {
+        id: "tlas",
+        sync: true,
+        inputs: [
+            { id: "instance-aabbs", access: "read" },
+            { id: "instance-count", access: "read" },
+        ],
+        outputs: [
+            { id: "tlas-bvh-nodes", access: "write" },
+            { id: "tlas-morton-codes", access: "write" },
+            { id: "tlas-instance-ids", access: "write" },
+        ],
+
+        async prepare(device: GPUDevice) {
+            const [boundsModule, mortonModule, treeModule, propagateModule, collapseModule] =
+                await Promise.all([
+                    device.createShaderModule({ code: boundsShader }),
+                    device.createShaderModule({ code: mortonShader }),
+                    device.createShaderModule({ code: treeShader }),
+                    device.createShaderModule({ code: propagateShader }),
+                    device.createShaderModule({ code: collapseShader }),
+                ]);
+
+            const [bounds, morton, tree, propagate, collapse, sort] = await Promise.all([
+                device.createComputePipelineAsync({
+                    layout: "auto",
+                    compute: { module: boundsModule, entryPoint: "main" },
+                }),
+                device.createComputePipelineAsync({
+                    layout: "auto",
+                    compute: { module: mortonModule, entryPoint: "main" },
+                }),
+                device.createComputePipelineAsync({
+                    layout: "auto",
+                    compute: { module: treeModule, entryPoint: "main" },
+                }),
+                device.createComputePipelineAsync({
+                    layout: "auto",
+                    compute: { module: propagateModule, entryPoint: "main" },
+                }),
+                device.createComputePipelineAsync({
+                    layout: "auto",
+                    compute: { module: collapseModule, entryPoint: "main" },
+                }),
+                createRadixSort(device, {
+                    keys: config.tlas.mortonCodes,
+                    values: config.tlas.instanceIds,
+                    count: MAX_ENTITIES,
+                }),
+            ]);
+
+            pipelines = { bounds, morton, tree, propagate, collapse };
+            radixSort = sort;
+
+            bindGroups = {
+                bounds: device.createBindGroup({
+                    layout: pipelines.bounds.getBindGroupLayout(0),
+                    entries: [
+                        { binding: 0, resource: { buffer: config.instanceAABBs } },
+                        { binding: 1, resource: { buffer: config.instanceCount } },
+                        { binding: 2, resource: { buffer: config.tlas.sceneBounds } },
+                        { binding: 3, resource: { buffer: config.tlas.entityIds } },
+                    ],
+                }),
+                morton: device.createBindGroup({
+                    layout: pipelines.morton.getBindGroupLayout(0),
+                    entries: [
+                        { binding: 0, resource: { buffer: config.instanceAABBs } },
+                        { binding: 1, resource: { buffer: config.instanceCount } },
+                        { binding: 2, resource: { buffer: config.tlas.sceneBounds } },
+                        { binding: 3, resource: { buffer: config.tlas.mortonCodes } },
+                        { binding: 4, resource: { buffer: config.tlas.instanceIds } },
+                        { binding: 5, resource: { buffer: config.tlas.entityIds } },
+                    ],
+                }),
+                tree: device.createBindGroup({
+                    layout: pipelines.tree.getBindGroupLayout(0),
+                    entries: [
+                        { binding: 0, resource: { buffer: config.tlas.mortonCodes } },
+                        { binding: 1, resource: { buffer: config.instanceCount } },
+                        { binding: 2, resource: { buffer: config.tlas.treeNodes } },
+                        { binding: 3, resource: { buffer: config.tlas.parentIndices } },
+                    ],
+                }),
+                propagate: device.createBindGroup({
+                    layout: pipelines.propagate.getBindGroupLayout(0),
+                    entries: [
+                        { binding: 0, resource: { buffer: config.instanceAABBs } },
+                        { binding: 1, resource: { buffer: config.tlas.instanceIds } },
+                        { binding: 2, resource: { buffer: config.instanceCount } },
+                        { binding: 3, resource: { buffer: config.tlas.treeNodes } },
+                        { binding: 4, resource: { buffer: config.tlas.boundsFlags } },
+                        { binding: 5, resource: { buffer: config.tlas.parentIndices } },
+                    ],
+                }),
+                collapse: device.createBindGroup({
+                    layout: pipelines.collapse.getBindGroupLayout(0),
+                    entries: [
+                        { binding: 0, resource: { buffer: config.tlas.treeNodes } },
+                        { binding: 1, resource: { buffer: config.instanceCount } },
+                        { binding: 2, resource: { buffer: config.tlas.parentIndices } },
+                        { binding: 3, resource: { buffer: config.tlas.bvhNodes } },
+                    ],
+                }),
+            };
+        },
+
+        execute(ctx: ExecutionContext) {
+            const { device, encoder } = ctx;
+
+            const workgroups = Math.ceil(MAX_ENTITIES / WORKGROUP_SIZE);
+
+            const initBounds = new Int32Array([
+                0x7f7fffff, 0x7f7fffff, 0x7f7fffff, 0, 0x80800000, 0x80800000, 0x80800000, 0,
+            ]);
+            device.queue.writeBuffer(config.tlas.sceneBounds, 0, initBounds);
+
+            encoder.clearBuffer(config.tlas.boundsFlags);
+            encoder.clearBuffer(config.tlas.parentIndices);
+
+            const boundsPass = encoder.beginComputePass();
+            boundsPass.setPipeline(pipelines!.bounds);
+            boundsPass.setBindGroup(0, bindGroups!.bounds);
+            boundsPass.dispatchWorkgroups(workgroups);
+            boundsPass.end();
+
+            const mortonPass = encoder.beginComputePass();
+            mortonPass.setPipeline(pipelines!.morton);
+            mortonPass.setBindGroup(0, bindGroups!.morton);
+            mortonPass.dispatchWorkgroups(workgroups);
+            mortonPass.end();
+
+            const sortPass = encoder.beginComputePass();
+            radixSort!.dispatch(sortPass);
+            sortPass.end();
+
+            const treePass = encoder.beginComputePass();
+            treePass.setPipeline(pipelines!.tree);
+            treePass.setBindGroup(0, bindGroups!.tree);
+            treePass.dispatchWorkgroups(Math.ceil((MAX_ENTITIES - 1) / WORKGROUP_SIZE));
+            treePass.end();
+
+            const propagatePass = encoder.beginComputePass();
+            propagatePass.setPipeline(pipelines!.propagate);
+            propagatePass.setBindGroup(0, bindGroups!.propagate);
+            propagatePass.dispatchWorkgroups(workgroups);
+            propagatePass.end();
+
+            const collapsePass = encoder.beginComputePass();
+            collapsePass.setPipeline(pipelines!.collapse);
+            collapsePass.setBindGroup(0, bindGroups!.collapse);
+            collapsePass.dispatchWorkgroups(Math.ceil((MAX_ENTITIES - 1) / WORKGROUP_SIZE));
+            collapsePass.end();
+        },
+    };
+}