@multiplekex/shallot 0.2.5 → 0.3.0

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

@@ -1,6 +1,7 @@
 import type { Vec3, BVHNode, MortonPair, AABB } from "./structs";
-import { LEAF_FLAG, isLeaf, leafIndex } from "./structs";
+import { LEAF_FLAG, isLeaf, leafIndex, INVALID_NODE } from "./structs";
 import type { MeshData } from "../../render/mesh";
+import { meshCount } from "../../render/mesh";
 
 export interface BLASTriangle {
     v0: Vec3;
@@ -14,18 +15,13 @@ export interface BLASTriangle {
 export interface BLASData {
     nodes: BVHNode[];
     sortedTriIds: number[];
+    sortedPairs: MortonPair[];
+    parents: number[];
     aabbMin: Vec3;
     aabbMax: Vec3;
     triCount: number;
 }
 
-export interface BLASMeta {
-    nodeOffset: number;
-    triIdOffset: number;
-    nodeCount: number;
-    triCount: number;
-}
-
 export interface BLASAtlas {
     blasData: Map<number, BLASData>;
     nodesBuffer: GPUBuffer;
@@ -34,6 +30,11 @@ export interface BLASAtlas {
     trianglesBuffer: GPUBuffer;
     triangles: Map<number, BLASTriangle[]>;
     shapeAABBs: GPUBuffer;
+    treeNodesBuffer: GPUBuffer;
+    parentIndicesBuffer: GPUBuffer;
+    baseTrianglesBuffer: GPUBuffer;
+    boundsFlagsBuffer: GPUBuffer;
+    metas: BLASMeta[];
 }
 
 function vec3(x: number, y: number, z: number): Vec3 {
@@ -406,8 +407,6 @@ function propagateBounds(
     const boundsFlags = new Array(n - 1).fill(0);
 
     for (let leafIdx = 0; leafIdx < n; leafIdx++) {
-        const tri = triangles[pairs[leafIdx].triangleId];
-
         let current = leafIdx;
         let isLeafNode = true;
 
@@ -466,6 +465,8 @@ export function buildShapeBLAS(triangles: BLASTriangle[]): BLASData {
         return {
             nodes: [],
             sortedTriIds: [],
+            sortedPairs: [],
+            parents: [],
             aabbMin: vec3(0, 0, 0),
             aabbMax: vec3(0, 0, 0),
             triCount: 0,
@@ -478,6 +479,8 @@ export function buildShapeBLAS(triangles: BLASTriangle[]): BLASData {
         return {
             nodes: [],
             sortedTriIds: [0],
+            sortedPairs: [{ code: 0, triangleId: 0 }],
+            parents: [-1],
             aabbMin: bounds.min,
             aabbMax: bounds.max,
             triCount: 1,
@@ -497,114 +500,293 @@ export function buildShapeBLAS(triangles: BLASTriangle[]): BLASData {
     return {
         nodes,
         sortedTriIds,
+        sortedPairs,
+        parents,
         aabbMin: rootBounds.min,
         aabbMax: rootBounds.max,
         triCount,
     };
 }
 
-function validateBLASBounds(blas: BLASData, triangles: BLASTriangle[]): boolean {
-    if (blas.triCount === 0) return true;
-
-    const epsilon = 1e-5;
+export interface BVH4Node {
+    children: [number, number, number, number];
+    bounds: [AABB, AABB, AABB, AABB];
+}
 
-    for (const tri of triangles) {
-        const v0 = tri.v0;
-        const v1 = vec3Add(v0, tri.e1);
-        const v2 = vec3Add(v0, tri.e2);
+function emptyAABB(): AABB {
+    return { min: vec3(1e30, 1e30, 1e30), max: vec3(-1e30, -1e30, -1e30) };
+}
 
-        for (const v of [v0, v1, v2]) {
-            if (
-                v.x < blas.aabbMin.x - epsilon ||
-                v.y < blas.aabbMin.y - epsilon ||
-                v.z < blas.aabbMin.z - epsilon ||
-                v.x > blas.aabbMax.x + epsilon ||
-                v.y > blas.aabbMax.y + epsilon ||
-                v.z > blas.aabbMax.z + epsilon
-            ) {
-                return false;
-            }
+function computeDepths(nodes: BVHNode[], parents: number[], n: number): number[] {
+    const depths = new Array<number>(nodes.length).fill(0);
+    for (let i = 0; i < nodes.length; i++) {
+        let depth = 0;
+        let current = i;
+        while (current !== 0) {
+            const parent = parents[n + current];
+            if (parent === -1 || parent === undefined) break;
+            current = parent;
+            depth++;
         }
+        depths[i] = depth;
     }
-
-    return true;
+    return depths;
 }
 
-function validateBLASNodes(blas: BLASData): { valid: boolean; errors: string[] } {
-    const errors: string[] = [];
+function getChildBounds(
+    child: number,
+    nodes: BVHNode[],
+    triangles: BLASTriangle[],
+    sortedPairs: MortonPair[]
+): AABB {
+    if (isLeaf(child)) {
+        const tri = triangles[sortedPairs[leafIndex(child)].triangleId];
+        return getTriangleBounds(tri);
+    }
+    const node = nodes[child];
+    return { min: node.min, max: node.max };
+}
 
-    if (blas.triCount <= 1) {
-        return { valid: true, errors: [] };
+export function collapseToBVH4(
+    bvh2Nodes: BVHNode[],
+    triCount: number,
+    sortedPairs: MortonPair[],
+    triangles: BLASTriangle[],
+    parents: number[]
+): BVH4Node[] {
+    if (triCount === 0) {
+        return [];
     }
 
-    const expectedNodes = blas.triCount - 1;
-    if (blas.nodes.length !== expectedNodes) {
-        errors.push(`Expected ${expectedNodes} nodes, got ${blas.nodes.length}`);
+    if (triCount === 1) {
+        const bounds = getTriangleBounds(triangles[sortedPairs[0].triangleId]);
+        return [
+            {
+                children: [LEAF_FLAG | 0, INVALID_NODE, INVALID_NODE, INVALID_NODE],
+                bounds: [bounds, emptyAABB(), emptyAABB(), emptyAABB()],
+            },
+        ];
     }
 
-    for (let i = 0; i < blas.nodes.length; i++) {
-        const node = blas.nodes[i];
+    if (triCount === 2) {
+        const bounds0 = getTriangleBounds(triangles[sortedPairs[0].triangleId]);
+        const bounds1 = getTriangleBounds(triangles[sortedPairs[1].triangleId]);
+        const node = bvh2Nodes[0];
+        return [
+            {
+                children: [node.leftChild, node.rightChild, INVALID_NODE, INVALID_NODE],
+                bounds: [bounds0, bounds1, emptyAABB(), emptyAABB()],
+            },
+        ];
+    }
 
-        if (node.min.x > node.max.x || node.min.y > node.max.y || node.min.z > node.max.z) {
-            errors.push(`Node ${i} has invalid bounds: min > max`);
-        }
+    const n = triCount;
+    const depths = computeDepths(bvh2Nodes, parents, n);
+    const bvh4Nodes: BVH4Node[] = new Array(bvh2Nodes.length);
 
+    for (let i = 0; i < bvh2Nodes.length; i++) {
+        const depth = depths[i];
+        const node = bvh2Nodes[i];
         const left = node.leftChild;
         const right = node.rightChild;
 
-        if (isLeaf(left)) {
-            const idx = leafIndex(left);
-            if (idx >= blas.triCount) {
-                errors.push(`Node ${i} left child leaf index ${idx} >= triCount ${blas.triCount}`);
-            }
+        const out: BVH4Node = {
+            children: [INVALID_NODE, INVALID_NODE, INVALID_NODE, INVALID_NODE],
+            bounds: [emptyAABB(), emptyAABB(), emptyAABB(), emptyAABB()],
+        };
+
+        if ((depth & 1) !== 0) {
+            out.children[0] = left;
+            out.bounds[0] = getChildBounds(left, bvh2Nodes, triangles, sortedPairs);
+            out.children[1] = right;
+            out.bounds[1] = getChildBounds(right, bvh2Nodes, triangles, sortedPairs);
         } else {
-            if (left >= blas.nodes.length) {
-                errors.push(`Node ${i} left child ${left} >= node count ${blas.nodes.length}`);
-            }
-        }
+            let slot = 0;
 
-        if (isLeaf(right)) {
-            const idx = leafIndex(right);
-            if (idx >= blas.triCount) {
-                errors.push(`Node ${i} right child leaf index ${idx} >= triCount ${blas.triCount}`);
+            if (isLeaf(left)) {
+                out.children[slot] = left;
+                out.bounds[slot] = getChildBounds(left, bvh2Nodes, triangles, sortedPairs);
+                slot++;
+            } else {
+                const leftNode = bvh2Nodes[left];
+                out.children[slot] = leftNode.leftChild;
+                out.bounds[slot] = getChildBounds(
+                    leftNode.leftChild,
+                    bvh2Nodes,
+                    triangles,
+                    sortedPairs
+                );
+                slot++;
+                out.children[slot] = leftNode.rightChild;
+                out.bounds[slot] = getChildBounds(
+                    leftNode.rightChild,
+                    bvh2Nodes,
+                    triangles,
+                    sortedPairs
+                );
+                slot++;
             }
-        } else {
-            if (right >= blas.nodes.length) {
-                errors.push(`Node ${i} right child ${right} >= node count ${blas.nodes.length}`);
+
+            if (isLeaf(right)) {
+                out.children[slot] = right;
+                out.bounds[slot] = getChildBounds(right, bvh2Nodes, triangles, sortedPairs);
+            } else {
+                const rightNode = bvh2Nodes[right];
+                out.children[slot] = rightNode.leftChild;
+                out.bounds[slot] = getChildBounds(
+                    rightNode.leftChild,
+                    bvh2Nodes,
+                    triangles,
+                    sortedPairs
+                );
+                slot++;
+                out.children[slot] = rightNode.rightChild;
+                out.bounds[slot] = getChildBounds(
+                    rightNode.rightChild,
+                    bvh2Nodes,
+                    triangles,
+                    sortedPairs
+                );
             }
         }
+
+        bvh4Nodes[i] = out;
     }
 
-    return { valid: errors.length === 0, errors };
+    return bvh4Nodes;
 }
 
-const MAX_SHAPES = 16;
-const TREE_NODE_SIZE = 32;
 const BLAS_TRIANGLE_SIZE = 64;
 
-export interface BLASMetaExtended {
+function createRefitBuffers(
+    device: GPUDevice,
+    blasData: Map<number, BLASData>,
+    trianglesMap: Map<number, BLASTriangle[]>,
+    metas: BLASMeta[],
+    totalTreeNodes: number,
+    totalParentIndices: number
+): { treeNodesBuffer: GPUBuffer; parentIndicesBuffer: GPUBuffer; boundsFlagsBuffer: GPUBuffer } {
+    const treeNodesData = new Uint32Array(Math.max(totalTreeNodes * 8, 8));
+    const treeNodesFloat = new Float32Array(treeNodesData.buffer);
+    const parentIndicesData = new Uint32Array(Math.max(totalParentIndices, 1));
+
+    let treeNodeWriteOffset = 0;
+    let parentWriteOffset = 0;
+
+    for (let shapeId = 0; shapeId < metas.length; shapeId++) {
+        const meta = metas[shapeId];
+        if (meta.triCount === 0) continue;
+        const blas = blasData.get(shapeId);
+        const triangles = trianglesMap.get(shapeId);
+        if (!blas || !triangles) continue;
+
+        const n = blas.triCount;
+        if (n <= 1) {
+            treeNodesData.fill(0, treeNodeWriteOffset * 8, treeNodeWriteOffset * 8 + 8);
+            treeNodeWriteOffset++;
+            parentIndicesData[parentWriteOffset++] = 0;
+        } else {
+            for (let i = 0; i < n - 1; i++) {
+                const bvhNode = blas.nodes[i];
+                const base = (treeNodeWriteOffset + i) * 8;
+                treeNodesFloat[base + 0] = bvhNode.min.x;
+                treeNodesFloat[base + 1] = bvhNode.min.y;
+                treeNodesFloat[base + 2] = bvhNode.min.z;
+                treeNodesData[base + 3] = bvhNode.leftChild >>> 0;
+                treeNodesFloat[base + 4] = bvhNode.max.x;
+                treeNodesFloat[base + 5] = bvhNode.max.y;
+                treeNodesFloat[base + 6] = bvhNode.max.z;
+                treeNodesData[base + 7] = bvhNode.rightChild >>> 0;
+            }
+
+            for (let leafIdx = 0; leafIdx < n; leafIdx++) {
+                const leafNodeIdx = n - 1 + leafIdx;
+                const base = (treeNodeWriteOffset + leafNodeIdx) * 8;
+                const triId = blas.sortedPairs[leafIdx].triangleId;
+                const tri = triangles[triId];
+                const v0 = tri.v0;
+                const v1 = vec3Add(v0, tri.e1);
+                const v2 = vec3Add(v0, tri.e2);
+                const tMin = vec3Min(vec3Min(v0, v1), v2);
+                const tMax = vec3Max(vec3Max(v0, v1), v2);
+                treeNodesFloat[base + 0] = tMin.x;
+                treeNodesFloat[base + 1] = tMin.y;
+                treeNodesFloat[base + 2] = tMin.z;
+                treeNodesData[base + 3] = 0;
+                treeNodesFloat[base + 4] = tMax.x;
+                treeNodesFloat[base + 5] = tMax.y;
+                treeNodesFloat[base + 6] = tMax.z;
+                treeNodesData[base + 7] = 0;
+            }
+
+            for (let i = 0; i < 2 * n; i++) {
+                parentIndicesData[parentWriteOffset + i] = blas.parents[i] >>> 0;
+            }
+
+            treeNodeWriteOffset += 2 * n - 1;
+            parentWriteOffset += 2 * n;
+        }
+    }
+
+    const treeNodesBuffer = device.createBuffer({
+        label: "blas-tree-nodes",
+        size: Math.max(treeNodesData.byteLength, 32),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(treeNodesBuffer, 0, treeNodesData);
+
+    const parentIndicesBuffer = device.createBuffer({
+        label: "blas-parent-indices",
+        size: Math.max(parentIndicesData.byteLength, 4),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(parentIndicesBuffer, 0, parentIndicesData);
+
+    const boundsFlagsBuffer = device.createBuffer({
+        label: "blas-bounds-flags",
+        size: Math.max(totalTreeNodes * 4, 4),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+
+    return { treeNodesBuffer, parentIndicesBuffer, boundsFlagsBuffer };
+}
+
+export interface BLASMeta {
     nodeOffset: number;
     triIdOffset: number;
     triOffset: number;
     triCount: number;
+    treeNodeOffset: number;
+    parentOffset: number;
 }
 
 export function createBLASAtlas(
     device: GPUDevice,
     getMesh: (id: number) => MeshData | undefined
 ): BLASAtlas {
+    const shapeCount = meshCount();
     const blasData = new Map<number, BLASData>();
+    const bvh4Data = new Map<number, BVH4Node[]>();
     const trianglesMap = new Map<number, BLASTriangle[]>();
-    const metas: BLASMetaExtended[] = [];
+    const metas: BLASMeta[] = [];
 
     let totalNodes = 0;
     let totalTriIds = 0;
     let totalTris = 0;
+    let totalTreeNodes = 0;
+    let totalParentIndices = 0;
 
-    for (let shapeId = 0; shapeId < MAX_SHAPES; shapeId++) {
+    for (let shapeId = 0; shapeId < shapeCount; shapeId++) {
         const mesh = getMesh(shapeId);
         if (!mesh || mesh.indexCount === 0) {
-            metas.push({ nodeOffset: 0, triIdOffset: 0, triOffset: 0, triCount: 0 });
+            metas.push({
+                nodeOffset: 0,
+                triIdOffset: 0,
+                triOffset: 0,
+                triCount: 0,
+                treeNodeOffset: 0,
+                parentOffset: 0,
+            });
             continue;
         }
 
@@ -613,45 +795,89 @@ export function createBLASAtlas(
         blasData.set(shapeId, blas);
         trianglesMap.set(shapeId, triangles);
 
+        const bvh4Nodes = collapseToBVH4(
+            blas.nodes,
+            blas.triCount,
+            blas.sortedPairs,
+            triangles,
+            blas.parents
+        );
+        bvh4Data.set(shapeId, bvh4Nodes);
+
+        const n = blas.triCount;
+        const treeNodeCount = n <= 1 ? 1 : 2 * n - 1;
+        const parentCount = n <= 1 ? 1 : 2 * n;
+
         metas.push({
             nodeOffset: totalNodes,
             triIdOffset: totalTriIds,
             triOffset: totalTris,
             triCount: blas.triCount,
+            treeNodeOffset: totalTreeNodes,
+            parentOffset: totalParentIndices,
         });
 
-        totalNodes += blas.nodes.length;
+        totalNodes += bvh4Nodes.length;
         totalTriIds += blas.sortedTriIds.length;
         totalTris += triangles.length;
+        totalTreeNodes += treeNodeCount;
+        totalParentIndices += parentCount;
     }
 
-    const nodesData = new Float32Array(Math.max(totalNodes * 8, 8));
+    const nodesData = new Float32Array(Math.max(totalNodes * 32, 32));
+    const nodesU32 = new Uint32Array(nodesData.buffer);
     const triIdsData = new Uint32Array(Math.max(totalTriIds, 1));
-    const metaData = new Uint32Array(MAX_SHAPES * 4);
+    const metaData = new Uint32Array(shapeCount * 4);
     const trianglesData = new Uint32Array(Math.max(totalTris * 16, 16));
 
     let nodeOffset = 0;
     let triIdOffset = 0;
     let triOffset = 0;
 
-    for (let shapeId = 0; shapeId < MAX_SHAPES; shapeId++) {
+    for (let shapeId = 0; shapeId < shapeCount; shapeId++) {
         const blas = blasData.get(shapeId);
+        const bvh4Nodes = bvh4Data.get(shapeId);
         const triangles = trianglesMap.get(shapeId);
-        if (!blas || !triangles) continue;
+        if (!blas || !triangles || !bvh4Nodes) continue;
+
+        for (const node of bvh4Nodes) {
+            const base = nodeOffset * 32;
+            nodesData[base + 0] = node.bounds[0].min.x;
+            nodesData[base + 1] = node.bounds[0].min.y;
+            nodesData[base + 2] = node.bounds[0].min.z;
+            nodesU32[base + 3] = node.children[0];
+            nodesData[base + 4] = node.bounds[0].max.x;
+            nodesData[base + 5] = node.bounds[0].max.y;
+            nodesData[base + 6] = node.bounds[0].max.z;
+            nodesData[base + 7] = 0;
+
+            nodesData[base + 8] = node.bounds[1].min.x;
+            nodesData[base + 9] = node.bounds[1].min.y;
+            nodesData[base + 10] = node.bounds[1].min.z;
+            nodesU32[base + 11] = node.children[1];
+            nodesData[base + 12] = node.bounds[1].max.x;
+            nodesData[base + 13] = node.bounds[1].max.y;
+            nodesData[base + 14] = node.bounds[1].max.z;
+            nodesData[base + 15] = 0;
+
+            nodesData[base + 16] = node.bounds[2].min.x;
+            nodesData[base + 17] = node.bounds[2].min.y;
+            nodesData[base + 18] = node.bounds[2].min.z;
+            nodesU32[base + 19] = node.children[2];
+            nodesData[base + 20] = node.bounds[2].max.x;
+            nodesData[base + 21] = node.bounds[2].max.y;
+            nodesData[base + 22] = node.bounds[2].max.z;
+            nodesData[base + 23] = 0;
+
+            nodesData[base + 24] = node.bounds[3].min.x;
+            nodesData[base + 25] = node.bounds[3].min.y;
+            nodesData[base + 26] = node.bounds[3].min.z;
+            nodesU32[base + 27] = node.children[3];
+            nodesData[base + 28] = node.bounds[3].max.x;
+            nodesData[base + 29] = node.bounds[3].max.y;
+            nodesData[base + 30] = node.bounds[3].max.z;
+            nodesData[base + 31] = 0;
 
-        for (const node of blas.nodes) {
-            nodesData[nodeOffset * 8 + 0] = node.min.x;
-            nodesData[nodeOffset * 8 + 1] = node.min.y;
-            nodesData[nodeOffset * 8 + 2] = node.min.z;
-            nodesData[nodeOffset * 8 + 3] = new Float32Array(
-                new Uint32Array([node.leftChild]).buffer
-            )[0];
-            nodesData[nodeOffset * 8 + 4] = node.max.x;
-            nodesData[nodeOffset * 8 + 5] = node.max.y;
-            nodesData[nodeOffset * 8 + 6] = node.max.z;
-            nodesData[nodeOffset * 8 + 7] = new Float32Array(
-                new Uint32Array([node.rightChild]).buffer
-            )[0];
             nodeOffset++;
         }
 
@@ -678,6 +904,7 @@ export function createBLASAtlas(
             trianglesData[base + 13] = octEncode(tri.n1);
             trianglesData[base + 14] = octEncode(tri.n2);
             trianglesData[base + 15] = 0;
+
             triOffset++;
         }
     }
@@ -691,7 +918,7 @@ export function createBLASAtlas(
 
     const nodesBuffer = device.createBuffer({
         label: "blas-nodes",
-        size: Math.max(nodesData.byteLength, 32),
+        size: Math.max(nodesData.byteLength, 128),
         usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
     });
     device.queue.writeBuffer(nodesBuffer, 0, nodesData);
@@ -717,8 +944,24 @@ export function createBLASAtlas(
     });
     device.queue.writeBuffer(trianglesBuffer, 0, trianglesData);
 
-    const shapeAABBsData = new Float32Array(MAX_SHAPES * 8);
-    for (let shapeId = 0; shapeId < MAX_SHAPES; shapeId++) {
+    const baseTrianglesBuffer = device.createBuffer({
+        label: "blas-base-triangles",
+        size: Math.max(totalTris * BLAS_TRIANGLE_SIZE, BLAS_TRIANGLE_SIZE),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(baseTrianglesBuffer, 0, trianglesData);
+
+    const { treeNodesBuffer, parentIndicesBuffer, boundsFlagsBuffer } = createRefitBuffers(
+        device,
+        blasData,
+        trianglesMap,
+        metas,
+        totalTreeNodes,
+        totalParentIndices
+    );
+
+    const shapeAABBsData = new Float32Array(shapeCount * 8);
+    for (let shapeId = 0; shapeId < shapeCount; shapeId++) {
         const blas = blasData.get(shapeId);
         const offset = shapeId * 8;
         if (blas) {
@@ -748,5 +991,10 @@ export function createBLASAtlas(
         trianglesBuffer,
         triangles: trianglesMap,
         shapeAABBs,
+        treeNodesBuffer,
+        parentIndicesBuffer,
+        baseTrianglesBuffer,
+        boundsFlagsBuffer,
+        metas,
     };
 }