@multiplekex/shallot 0.1.12 → 0.2.0

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

@@ -0,0 +1,752 @@
+import type { Vec3, BVHNode, MortonPair, AABB } from "./structs";
+import { LEAF_FLAG, isLeaf, leafIndex } from "./structs";
+import type { MeshData } from "../mesh";
+
+export interface BLASTriangle {
+    v0: Vec3;
+    e1: Vec3;
+    e2: Vec3;
+    n0: Vec3;
+    n1: Vec3;
+    n2: Vec3;
+}
+
+export interface BLASData {
+    nodes: BVHNode[];
+    sortedTriIds: 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;
+    triIdsBuffer: GPUBuffer;
+    metaBuffer: GPUBuffer;
+    trianglesBuffer: GPUBuffer;
+    triangles: Map<number, BLASTriangle[]>;
+    shapeAABBs: GPUBuffer;
+}
+
+function vec3(x: number, y: number, z: number): Vec3 {
+    return { x, y, z };
+}
+
+function vec3Min(a: Vec3, b: Vec3): Vec3 {
+    return { x: Math.min(a.x, b.x), y: Math.min(a.y, b.y), z: Math.min(a.z, b.z) };
+}
+
+function vec3Max(a: Vec3, b: Vec3): Vec3 {
+    return { x: Math.max(a.x, b.x), y: Math.max(a.y, b.y), z: Math.max(a.z, b.z) };
+}
+
+function vec3Add(a: Vec3, b: Vec3): Vec3 {
+    return { x: a.x + b.x, y: a.y + b.y, z: a.z + b.z };
+}
+
+function vec3Sub(a: Vec3, b: Vec3): Vec3 {
+    return { x: a.x - b.x, y: a.y - b.y, z: a.z - b.z };
+}
+
+function vec3Scale(v: Vec3, s: number): Vec3 {
+    return { x: v.x * s, y: v.y * s, z: v.z * s };
+}
+
+function octEncode(n: Vec3): number {
+    const absSum = Math.abs(n.x) + Math.abs(n.y) + Math.abs(n.z);
+    let vx = n.x / absSum;
+    let vy = n.y / absSum;
+    const vz = n.z / absSum;
+    if (vz < 0) {
+        const signX = vx >= 0 ? 1 : -1;
+        const signY = vy >= 0 ? 1 : -1;
+        const newVx = (1 - Math.abs(vy)) * signX;
+        const newVy = (1 - Math.abs(vx)) * signY;
+        vx = newVx;
+        vy = newVy;
+    }
+    const x = Math.floor(Math.max(0, Math.min(65535, (vx * 0.5 + 0.5) * 65535)));
+    const y = Math.floor(Math.max(0, Math.min(65535, (vy * 0.5 + 0.5) * 65535)));
+    return ((y << 16) | x) >>> 0;
+}
+
+export function extractShapeTriangles(mesh: MeshData): BLASTriangle[] {
+    const triangles: BLASTriangle[] = [];
+    const { vertices, indices, indexCount } = mesh;
+    const stride = 6;
+
+    for (let i = 0; i < indexCount; i += 3) {
+        const i0 = indices[i];
+        const i1 = indices[i + 1];
+        const i2 = indices[i + 2];
+
+        const v0 = vec3(
+            vertices[i0 * stride],
+            vertices[i0 * stride + 1],
+            vertices[i0 * stride + 2]
+        );
+        const v1 = vec3(
+            vertices[i1 * stride],
+            vertices[i1 * stride + 1],
+            vertices[i1 * stride + 2]
+        );
+        const v2 = vec3(
+            vertices[i2 * stride],
+            vertices[i2 * stride + 1],
+            vertices[i2 * stride + 2]
+        );
+
+        const n0 = vec3(
+            vertices[i0 * stride + 3],
+            vertices[i0 * stride + 4],
+            vertices[i0 * stride + 5]
+        );
+        const n1 = vec3(
+            vertices[i1 * stride + 3],
+            vertices[i1 * stride + 4],
+            vertices[i1 * stride + 5]
+        );
+        const n2 = vec3(
+            vertices[i2 * stride + 3],
+            vertices[i2 * stride + 4],
+            vertices[i2 * stride + 5]
+        );
+
+        triangles.push({
+            v0,
+            e1: vec3Sub(v1, v0),
+            e2: vec3Sub(v2, v0),
+            n0,
+            n1,
+            n2,
+        });
+    }
+
+    return triangles;
+}
+
+function computeBounds(triangles: BLASTriangle[]): AABB {
+    if (triangles.length === 0) {
+        return { min: vec3(0, 0, 0), max: vec3(0, 0, 0) };
+    }
+
+    let min = vec3(Infinity, Infinity, Infinity);
+    let max = vec3(-Infinity, -Infinity, -Infinity);
+
+    for (const tri of triangles) {
+        const v0 = tri.v0;
+        const v1 = vec3Add(v0, tri.e1);
+        const v2 = vec3Add(v0, tri.e2);
+        min = vec3Min(min, v0);
+        min = vec3Min(min, v1);
+        min = vec3Min(min, v2);
+        max = vec3Max(max, v0);
+        max = vec3Max(max, v1);
+        max = vec3Max(max, v2);
+    }
+
+    return { min, max };
+}
+
+function expandBits(v: number): number {
+    let x = v & 0x3ff;
+    x = (x | (x << 16)) & 0x030000ff;
+    x = (x | (x << 8)) & 0x0300f00f;
+    x = (x | (x << 4)) & 0x030c30c3;
+    x = (x | (x << 2)) & 0x09249249;
+    return x >>> 0;
+}
+
+function mortonCode3D(x: number, y: number, z: number): number {
+    return ((expandBits(x) << 2) | (expandBits(y) << 1) | expandBits(z)) >>> 0;
+}
+
+function computeMortonCode(centroid: Vec3, bounds: AABB): number {
+    const size = {
+        x: bounds.max.x - bounds.min.x,
+        y: bounds.max.y - bounds.min.y,
+        z: bounds.max.z - bounds.min.z,
+    };
+
+    const safeSize = {
+        x: Math.max(size.x, 1e-6),
+        y: Math.max(size.y, 1e-6),
+        z: Math.max(size.z, 1e-6),
+    };
+
+    const normalized = {
+        x: (centroid.x - bounds.min.x) / safeSize.x,
+        y: (centroid.y - bounds.min.y) / safeSize.y,
+        z: (centroid.z - bounds.min.z) / safeSize.z,
+    };
+
+    const clamped = {
+        x: Math.max(0, Math.min(1, normalized.x)),
+        y: Math.max(0, Math.min(1, normalized.y)),
+        z: Math.max(0, Math.min(1, normalized.z)),
+    };
+
+    const quantized = {
+        x: Math.floor(clamped.x * 1023),
+        y: Math.floor(clamped.y * 1023),
+        z: Math.floor(clamped.z * 1023),
+    };
+
+    return mortonCode3D(quantized.x, quantized.y, quantized.z);
+}
+
+function buildMortonPairs(triangles: BLASTriangle[], bounds: AABB): MortonPair[] {
+    return triangles.map((tri, i) => {
+        const centroid = vec3Add(tri.v0, vec3Scale(vec3Add(tri.e1, tri.e2), 1 / 3));
+        return {
+            code: computeMortonCode(centroid, bounds),
+            triangleId: i,
+        };
+    });
+}
+
+function radixSort(pairs: MortonPair[]): MortonPair[] {
+    const n = pairs.length;
+    if (n === 0) return [];
+
+    let input = [...pairs];
+    let output = new Array<MortonPair>(n);
+
+    for (let pass = 0; pass < 4; pass++) {
+        const bitOffset = pass * 8;
+        const histogram = new Array<number>(256).fill(0);
+
+        for (const pair of input) {
+            const digit = (pair.code >>> bitOffset) & 0xff;
+            histogram[digit]++;
+        }
+
+        let sum = 0;
+        for (let i = 0; i < 256; i++) {
+            const count = histogram[i];
+            histogram[i] = sum;
+            sum += count;
+        }
+
+        for (const pair of input) {
+            const digit = (pair.code >>> bitOffset) & 0xff;
+            output[histogram[digit]] = pair;
+            histogram[digit]++;
+        }
+
+        [input, output] = [output, input];
+    }
+
+    return input;
+}
+
+function clz32(x: number): number {
+    if (x === 0) return 32;
+    return Math.clz32(x >>> 0);
+}
+
+function delta(sortedPairs: MortonPair[], i: number, j: number): number {
+    const n = sortedPairs.length;
+    if (j < 0 || j >= n) {
+        return -1;
+    }
+    const codeI = sortedPairs[i].code >>> 0;
+    const codeJ = sortedPairs[j].code >>> 0;
+    if (codeI === codeJ) {
+        return clz32((i ^ j) >>> 0) + 32;
+    }
+    return clz32((codeI ^ codeJ) >>> 0);
+}
+
+function determineRange(sortedPairs: MortonPair[], i: number): [number, number] {
+    const n = sortedPairs.length;
+
+    if (i === 0) {
+        return [0, n - 1];
+    }
+
+    const deltaLeft = delta(sortedPairs, i, i - 1);
+    const deltaRight = delta(sortedPairs, i, i + 1);
+    const d = deltaRight > deltaLeft ? 1 : -1;
+
+    const deltaMin = Math.min(deltaLeft, deltaRight);
+
+    let lmax = 2;
+    while (delta(sortedPairs, i, i + lmax * d) > deltaMin) {
+        lmax *= 2;
+    }
+
+    let l = 0;
+    let t = Math.floor(lmax / 2);
+    while (t >= 1) {
+        if (delta(sortedPairs, i, i + (l + t) * d) > deltaMin) {
+            l += t;
+        }
+        t = Math.floor(t / 2);
+    }
+
+    const j = i + l * d;
+    const first = Math.min(i, j);
+    const last = Math.max(i, j);
+
+    return [first, last];
+}
+
+function findSplit(sortedPairs: MortonPair[], first: number, last: number): number {
+    const firstCode = sortedPairs[first].code >>> 0;
+    const lastCode = sortedPairs[last].code >>> 0;
+
+    if (firstCode === lastCode) {
+        return Math.floor((first + last) / 2);
+    }
+
+    const deltaNode = clz32((firstCode ^ lastCode) >>> 0);
+
+    let split = first;
+    let stride = last - first;
+
+    do {
+        stride = Math.floor((stride + 1) / 2);
+        const middle = split + stride;
+
+        if (middle < last) {
+            const splitCode = sortedPairs[middle].code >>> 0;
+            const splitDelta = clz32((firstCode ^ splitCode) >>> 0);
+
+            if (splitDelta > deltaNode) {
+                split = middle;
+            }
+        }
+    } while (stride > 1);
+
+    return split;
+}
+
+function buildKarrasTree(sortedPairs: MortonPair[]): {
+    nodes: BVHNode[];
+    parents: number[];
+} {
+    const n = sortedPairs.length;
+
+    if (n === 0) {
+        return { nodes: [], parents: [] };
+    }
+
+    if (n === 1) {
+        return { nodes: [], parents: [-1] };
+    }
+
+    const numInternal = n - 1;
+    const nodes: BVHNode[] = new Array(numInternal);
+    const parents: number[] = new Array(2 * n).fill(-1);
+
+    for (let i = 0; i < numInternal; i++) {
+        const [first, last] = determineRange(sortedPairs, i);
+        const gamma = findSplit(sortedPairs, first, last);
+
+        const leftIsLeaf = Math.min(first, last) === gamma;
+        const rightIsLeaf = Math.max(first, last) === gamma + 1;
+
+        let leftChild: number;
+        let rightChild: number;
+
+        if (leftIsLeaf) {
+            leftChild = (gamma | LEAF_FLAG) >>> 0;
+            parents[gamma] = i;
+        } else {
+            leftChild = gamma;
+            parents[n + gamma] = i;
+        }
+
+        if (rightIsLeaf) {
+            rightChild = ((gamma + 1) | LEAF_FLAG) >>> 0;
+            parents[gamma + 1] = i;
+        } else {
+            rightChild = gamma + 1;
+            parents[n + (gamma + 1)] = i;
+        }
+
+        nodes[i] = {
+            min: vec3(1e30, 1e30, 1e30),
+            max: vec3(-1e30, -1e30, -1e30),
+            leftChild,
+            rightChild,
+        };
+    }
+
+    return { nodes, parents };
+}
+
+function getTriangleBounds(tri: BLASTriangle): AABB {
+    const v0 = tri.v0;
+    const v1 = vec3Add(v0, tri.e1);
+    const v2 = vec3Add(v0, tri.e2);
+    const min = vec3Min(vec3Min(v0, v1), v2);
+    const max = vec3Max(vec3Max(v0, v1), v2);
+    return { min, max };
+}
+
+function propagateBounds(
+    nodes: BVHNode[],
+    triangles: BLASTriangle[],
+    pairs: MortonPair[],
+    parents: number[]
+): void {
+    const n = triangles.length;
+    if (n <= 1) return;
+
+    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;
+
+        for (let iter = 0; iter < 64; iter++) {
+            const parent = isLeafNode ? parents[current] : parents[n + current];
+
+            if (parent === -1 || parent === undefined) {
+                break;
+            }
+
+            const oldFlag = boundsFlags[parent];
+            boundsFlags[parent]++;
+
+            if (oldFlag === 0) {
+                break;
+            }
+
+            const node = nodes[parent];
+            const left = node.leftChild;
+            const right = node.rightChild;
+
+            let leftBounds: AABB;
+            let rightBounds: AABB;
+
+            if (isLeaf(left)) {
+                const leftTri = triangles[pairs[leafIndex(left)].triangleId];
+                leftBounds = getTriangleBounds(leftTri);
+            } else {
+                leftBounds = { min: nodes[left].min, max: nodes[left].max };
+            }
+
+            if (isLeaf(right)) {
+                const rightTri = triangles[pairs[leafIndex(right)].triangleId];
+                rightBounds = getTriangleBounds(rightTri);
+            } else {
+                rightBounds = { min: nodes[right].min, max: nodes[right].max };
+            }
+
+            nodes[parent].min = vec3Min(leftBounds.min, rightBounds.min);
+            nodes[parent].max = vec3Max(leftBounds.max, rightBounds.max);
+
+            current = parent;
+            isLeafNode = false;
+
+            if (parent === 0) {
+                break;
+            }
+        }
+    }
+}
+
+export function buildShapeBLAS(triangles: BLASTriangle[]): BLASData {
+    const triCount = triangles.length;
+
+    if (triCount === 0) {
+        return {
+            nodes: [],
+            sortedTriIds: [],
+            aabbMin: vec3(0, 0, 0),
+            aabbMax: vec3(0, 0, 0),
+            triCount: 0,
+        };
+    }
+
+    const bounds = computeBounds(triangles);
+
+    if (triCount === 1) {
+        return {
+            nodes: [],
+            sortedTriIds: [0],
+            aabbMin: bounds.min,
+            aabbMax: bounds.max,
+            triCount: 1,
+        };
+    }
+
+    const pairs = buildMortonPairs(triangles, bounds);
+    const sortedPairs = radixSort(pairs);
+    const { nodes, parents } = buildKarrasTree(sortedPairs);
+
+    propagateBounds(nodes, triangles, sortedPairs, parents);
+
+    const sortedTriIds = sortedPairs.map((p) => p.triangleId);
+
+    const rootBounds = nodes.length > 0 ? { min: nodes[0].min, max: nodes[0].max } : bounds;
+
+    return {
+        nodes,
+        sortedTriIds,
+        aabbMin: rootBounds.min,
+        aabbMax: rootBounds.max,
+        triCount,
+    };
+}
+
+export function validateBLASBounds(blas: BLASData, triangles: BLASTriangle[]): boolean {
+    if (blas.triCount === 0) return true;
+
+    const epsilon = 1e-5;
+
+    for (const tri of triangles) {
+        const v0 = tri.v0;
+        const v1 = vec3Add(v0, tri.e1);
+        const v2 = vec3Add(v0, tri.e2);
+
+        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;
+            }
+        }
+    }
+
+    return true;
+}
+
+export function validateBLASNodes(blas: BLASData): { valid: boolean; errors: string[] } {
+    const errors: string[] = [];
+
+    if (blas.triCount <= 1) {
+        return { valid: true, errors: [] };
+    }
+
+    const expectedNodes = blas.triCount - 1;
+    if (blas.nodes.length !== expectedNodes) {
+        errors.push(`Expected ${expectedNodes} nodes, got ${blas.nodes.length}`);
+    }
+
+    for (let i = 0; i < blas.nodes.length; i++) {
+        const node = blas.nodes[i];
+
+        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 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}`);
+            }
+        } else {
+            if (left >= blas.nodes.length) {
+                errors.push(`Node ${i} left child ${left} >= node count ${blas.nodes.length}`);
+            }
+        }
+
+        if (isLeaf(right)) {
+            const idx = leafIndex(right);
+            if (idx >= blas.triCount) {
+                errors.push(`Node ${i} right child leaf index ${idx} >= triCount ${blas.triCount}`);
+            }
+        } else {
+            if (right >= blas.nodes.length) {
+                errors.push(`Node ${i} right child ${right} >= node count ${blas.nodes.length}`);
+            }
+        }
+    }
+
+    return { valid: errors.length === 0, errors };
+}
+
+const MAX_SHAPES = 16;
+const TREE_NODE_SIZE = 32;
+const BLAS_TRIANGLE_SIZE = 64;
+
+export interface BLASMetaExtended {
+    nodeOffset: number;
+    triIdOffset: number;
+    triOffset: number;
+    triCount: number;
+}
+
+export function createBLASAtlas(
+    device: GPUDevice,
+    getMesh: (id: number) => MeshData | undefined
+): BLASAtlas {
+    const blasData = new Map<number, BLASData>();
+    const trianglesMap = new Map<number, BLASTriangle[]>();
+    const metas: BLASMetaExtended[] = [];
+
+    let totalNodes = 0;
+    let totalTriIds = 0;
+    let totalTris = 0;
+
+    for (let shapeId = 0; shapeId < MAX_SHAPES; shapeId++) {
+        const mesh = getMesh(shapeId);
+        if (!mesh || mesh.indexCount === 0) {
+            metas.push({ nodeOffset: 0, triIdOffset: 0, triOffset: 0, triCount: 0 });
+            continue;
+        }
+
+        const triangles = extractShapeTriangles(mesh);
+        const blas = buildShapeBLAS(triangles);
+        blasData.set(shapeId, blas);
+        trianglesMap.set(shapeId, triangles);
+
+        metas.push({
+            nodeOffset: totalNodes,
+            triIdOffset: totalTriIds,
+            triOffset: totalTris,
+            triCount: blas.triCount,
+        });
+
+        totalNodes += blas.nodes.length;
+        totalTriIds += blas.sortedTriIds.length;
+        totalTris += triangles.length;
+    }
+
+    const nodesData = new Float32Array(Math.max(totalNodes * 8, 8));
+    const triIdsData = new Uint32Array(Math.max(totalTriIds, 1));
+    const metaData = new Uint32Array(MAX_SHAPES * 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++) {
+        const blas = blasData.get(shapeId);
+        const triangles = trianglesMap.get(shapeId);
+        if (!blas || !triangles) continue;
+
+        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++;
+        }
+
+        for (const triId of blas.sortedTriIds) {
+            triIdsData[triIdOffset++] = triId;
+        }
+
+        for (const tri of triangles) {
+            const base = triOffset * 16;
+            const floatView = new Float32Array(trianglesData.buffer, base * 4, 16);
+            floatView[0] = tri.v0.x;
+            floatView[1] = tri.v0.y;
+            floatView[2] = tri.v0.z;
+            trianglesData[base + 3] = 0;
+            floatView[4] = tri.e1.x;
+            floatView[5] = tri.e1.y;
+            floatView[6] = tri.e1.z;
+            trianglesData[base + 7] = 0;
+            floatView[8] = tri.e2.x;
+            floatView[9] = tri.e2.y;
+            floatView[10] = tri.e2.z;
+            trianglesData[base + 11] = 0;
+            trianglesData[base + 12] = octEncode(tri.n0);
+            trianglesData[base + 13] = octEncode(tri.n1);
+            trianglesData[base + 14] = octEncode(tri.n2);
+            trianglesData[base + 15] = 0;
+            triOffset++;
+        }
+    }
+
+    for (let i = 0; i < metas.length; i++) {
+        metaData[i * 4 + 0] = metas[i].nodeOffset;
+        metaData[i * 4 + 1] = metas[i].triIdOffset;
+        metaData[i * 4 + 2] = metas[i].triOffset;
+        metaData[i * 4 + 3] = metas[i].triCount;
+    }
+
+    const nodesBuffer = device.createBuffer({
+        label: "blas-nodes",
+        size: Math.max(nodesData.byteLength, 32),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(nodesBuffer, 0, nodesData);
+
+    const triIdsBuffer = device.createBuffer({
+        label: "blas-triIds",
+        size: Math.max(triIdsData.byteLength, 4),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(triIdsBuffer, 0, triIdsData);
+
+    const metaBuffer = device.createBuffer({
+        label: "blas-meta",
+        size: metaData.byteLength,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(metaBuffer, 0, metaData);
+
+    const trianglesBuffer = device.createBuffer({
+        label: "blas-triangles",
+        size: Math.max(totalTris * BLAS_TRIANGLE_SIZE, BLAS_TRIANGLE_SIZE),
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(trianglesBuffer, 0, trianglesData);
+
+    const shapeAABBsData = new Float32Array(MAX_SHAPES * 8);
+    for (let shapeId = 0; shapeId < MAX_SHAPES; shapeId++) {
+        const blas = blasData.get(shapeId);
+        const offset = shapeId * 8;
+        if (blas) {
+            shapeAABBsData[offset + 0] = blas.aabbMin.x;
+            shapeAABBsData[offset + 1] = blas.aabbMin.y;
+            shapeAABBsData[offset + 2] = blas.aabbMin.z;
+            shapeAABBsData[offset + 3] = 0;
+            shapeAABBsData[offset + 4] = blas.aabbMax.x;
+            shapeAABBsData[offset + 5] = blas.aabbMax.y;
+            shapeAABBsData[offset + 6] = blas.aabbMax.z;
+            shapeAABBsData[offset + 7] = 0;
+        }
+    }
+
+    const shapeAABBs = device.createBuffer({
+        label: "shape-aabbs",
+        size: shapeAABBsData.byteLength,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(shapeAABBs, 0, shapeAABBsData);
+
+    return {
+        blasData,
+        nodesBuffer,
+        triIdsBuffer,
+        metaBuffer,
+        trianglesBuffer,
+        triangles: trianglesMap,
+        shapeAABBs,
+    };
+}