@multiplekex/shallot 0.2.5 → 0.3.0

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

@@ -58,53 +58,6 @@ export function leafIndex(child: number): number {
 
 export const LEAF_FLAG_WGSL = /* wgsl */ `const LEAF_FLAG: u32 = 0x80000000u;`;
 
-export const TRIANGLE_STRUCT_WGSL = /* wgsl */ `
-struct Triangle {
-    v0: vec3<f32>,
-    entityId: u32,
-    e1: vec3<f32>,
-    _pad0: u32,
-    e2: vec3<f32>,
-    _pad1: u32,
-    n0_enc: u32,
-    n1_enc: u32,
-    n2_enc: u32,
-    _pad2: u32,
-}
-
-fn octEncode(n: vec3<f32>) -> u32 {
-    let absSum = abs(n.x) + abs(n.y) + abs(n.z);
-    var vx = n.x / absSum;
-    var vy = n.y / absSum;
-    let vz = n.z / absSum;
-    if (vz < 0.0) {
-        let signX = select(-1.0, 1.0, vx >= 0.0);
-        let signY = select(-1.0, 1.0, vy >= 0.0);
-        let newVx = (1.0 - abs(vy)) * signX;
-        let newVy = (1.0 - abs(vx)) * signY;
-        vx = newVx;
-        vy = newVy;
-    }
-    let x = u32(clamp((vx * 0.5 + 0.5) * 65535.0, 0.0, 65535.0));
-    let y = u32(clamp((vy * 0.5 + 0.5) * 65535.0, 0.0, 65535.0));
-    return (y << 16u) | x;
-}
-
-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);
-}`;
-
 export const TREE_NODE_STRUCT_WGSL = /* wgsl */ `
 struct TreeNode {
     minX: f32,
@@ -131,16 +84,8 @@ struct BVHNode {
 
 export const TREE_NODE_SIZE = 32;
 export const BVH_NODE_SIZE = 128;
-export const BLAS_NODE_SIZE = 32;
-export const BLAS_TRIANGLE_SIZE = 64;
 export const INVALID_NODE = 0xffffffff;
 
-export const BLAS_NODE_STRUCT_WGSL = /* wgsl */ `
-struct BLASNode {
-    minX: f32, minY: f32, minZ: f32, leftChild: u32,
-    maxX: f32, maxY: f32, maxZ: f32, rightChild: u32,
-}`;
-
 export const BLAS_TRIANGLE_STRUCT_WGSL = /* wgsl */ `
 struct BLASTriangle {
     v0: vec3<f32>, _pad0: u32,

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

@@ -1,6 +1,6 @@
 import { MAX_ENTITIES } from "../../../core";
 import type { ComputeNode, ExecutionContext } from "../../compute";
-import { createRadixSort } from "./radix";
+import { createRadixSort, dispatchRadixSort } from "./radix";
 import {
     TREE_NODE_STRUCT_WGSL,
     BVH_NODE_STRUCT_WGSL,
@@ -724,10 +724,156 @@ export interface TLASConfig {
     getRaytracing?: () => boolean;
 }
 
+interface TLASGPU {
+    pipelines: TLASPipelines;
+    bindGroups: TLASBindGroups;
+    radixSort: Awaited<ReturnType<typeof createRadixSort>>;
+}
+
+async function prepareTLAS(device: GPUDevice, config: TLASConfig): Promise<TLASGPU> {
+    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, radixSort] = 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,
+        }),
+    ]);
+
+    const pipelines = { bounds, morton, tree, propagate, collapse };
+
+    const bindGroups: TLASBindGroups = {
+        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 } },
+            ],
+        }),
+    };
+
+    return { pipelines, bindGroups, radixSort };
+}
+
+const initBounds = new Int32Array([
+    0x7f7fffff, 0x7f7fffff, 0x7f7fffff, 0, 0x80800000, 0x80800000, 0x80800000, 0,
+]);
+
+function executeTLAS(gpu: TLASGPU, ctx: ExecutionContext, config: TLASConfig): void {
+    const { device, encoder } = ctx;
+    const workgroups = Math.ceil(MAX_ENTITIES / WORKGROUP_SIZE);
+
+    device.queue.writeBuffer(config.tlas.sceneBounds, 0, initBounds);
+    encoder.clearBuffer(config.tlas.boundsFlags);
+    encoder.clearBuffer(config.tlas.parentIndices);
+
+    const boundsPass = encoder.beginComputePass();
+    boundsPass.setPipeline(gpu.pipelines.bounds);
+    boundsPass.setBindGroup(0, gpu.bindGroups.bounds);
+    boundsPass.dispatchWorkgroups(workgroups);
+    boundsPass.end();
+
+    const mortonPass = encoder.beginComputePass();
+    mortonPass.setPipeline(gpu.pipelines.morton);
+    mortonPass.setBindGroup(0, gpu.bindGroups.morton);
+    mortonPass.dispatchWorkgroups(workgroups);
+    mortonPass.end();
+
+    const sortPass = encoder.beginComputePass();
+    dispatchRadixSort(gpu.radixSort, sortPass);
+    sortPass.end();
+
+    const treePass = encoder.beginComputePass();
+    treePass.setPipeline(gpu.pipelines.tree);
+    treePass.setBindGroup(0, gpu.bindGroups.tree);
+    treePass.dispatchWorkgroups(Math.ceil((MAX_ENTITIES - 1) / WORKGROUP_SIZE));
+    treePass.end();
+
+    const propagatePass = encoder.beginComputePass();
+    propagatePass.setPipeline(gpu.pipelines.propagate);
+    propagatePass.setBindGroup(0, gpu.bindGroups.propagate);
+    propagatePass.dispatchWorkgroups(workgroups);
+    propagatePass.end();
+
+    const collapsePass = encoder.beginComputePass();
+    collapsePass.setPipeline(gpu.pipelines.collapse);
+    collapsePass.setBindGroup(0, gpu.bindGroups.collapse);
+    collapsePass.dispatchWorkgroups(Math.ceil((MAX_ENTITIES - 1) / WORKGROUP_SIZE));
+    collapsePass.end();
+}
+
 export function createTLASNode(config: TLASConfig): ComputeNode {
-    let pipelines: TLASPipelines | null = null;
-    let bindGroups: TLASBindGroups | null = null;
-    let radixSort: Awaited<ReturnType<typeof createRadixSort>> | null = null;
+    let gpu: TLASGPU | null = null;
 
     return {
         id: "tlas",
@@ -743,147 +889,13 @@ export function createTLASNode(config: TLASConfig): ComputeNode {
         ],
 
         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 } },
-                    ],
-                }),
-            };
+            gpu = await prepareTLAS(device, config);
         },
 
         execute(ctx: ExecutionContext) {
             if (config.getRaytracing && !config.getRaytracing()) return;
-
-            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();
+            if (!gpu) return;
+            executeTLAS(gpu, ctx, config);
         },
     };
 }

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

@@ -1,11 +1,6 @@
-import {
-    BVH_NODE_STRUCT_WGSL,
-    LEAF_FLAG_WGSL,
-    BLAS_NODE_STRUCT_WGSL,
-    BLAS_TRIANGLE_STRUCT_WGSL,
-} from "./structs";
-
-export const OCT_DECODE_WGSL = /* wgsl */ `
+import { BVH_NODE_STRUCT_WGSL, LEAF_FLAG_WGSL, BLAS_TRIANGLE_STRUCT_WGSL } from "./structs";
+
+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;
@@ -32,15 +27,13 @@ ${OCT_DECODE_WGSL}
 `;
 
 export const TLAS_BLAS_STRUCTS = /* wgsl */ `
-${BLAS_NODE_STRUCT_WGSL}
-
 ${BLAS_TRIANGLE_STRUCT_WGSL}
 `;
 
 export const TLAS_BLAS_BINDINGS = /* wgsl */ `
 @group(1) @binding(0) var<storage, read> tlasNodes: array<BVHNode>;
 @group(1) @binding(1) var<storage, read> tlasInstanceIds: array<u32>;
-@group(1) @binding(2) var<storage, read> blasNodes: array<BLASNode>;
+@group(1) @binding(2) var<storage, read> blasNodes: array<BVHNode>;
 @group(1) @binding(3) var<storage, read> blasTriIds: array<u32>;
 @group(1) @binding(4) var<storage, read> blasTriangles: array<BLASTriangle>;
 @group(1) @binding(5) var<storage, read> blasMeta: array<u32>;
@@ -142,7 +135,7 @@ fn traceBLAS(
     nodeOffset: u32,
     triIdOffset: u32,
     triOffset: u32,
-    triCount_: u32,
+    totalTriCount: u32,
     maxT: f32
 ) -> HitResult {
     var closest: HitResult;
@@ -154,17 +147,7 @@ fn traceBLAS(
     closest.normal = vec3(0.0, 1.0, 0.0);
     closest.worldPos = vec3(0.0);
 
-    if (triCount_ == 0u) {
-        return closest;
-    }
-
-    if (triCount_ == 1u) {
-        let triIdx = blasTriIds[triIdOffset];
-        let tri = blasTriangles[triOffset + triIdx];
-        let hit = intersectBLASTriangle(ray, tri);
-        if (hit.hit && hit.t < maxT) {
-            return hit;
-        }
+    if (totalTriCount == 0u) {
         return closest;
     }
 
@@ -177,7 +160,7 @@ fn traceBLAS(
     stackPtr++;
 
     var iterations = 0u;
-    let maxIterations = min(triCount_ * 3u, 10000u);
+    let maxIterations = min(totalTriCount * 3u, 10000u);
 
     while (stackPtr > 0u && iterations < maxIterations) {
         iterations++;
@@ -185,50 +168,75 @@ fn traceBLAS(
         let localIdx = stack[stackPtr];
         let node = blasNodes[nodeOffset + localIdx];
 
-        let leftChild = node.leftChild;
-        let rightChild = node.rightChild;
+        var children: array<u32, 4>;
+        var dists: array<f32, 4>;
 
-        if (leftChild != INVALID_NODE) {
-            if (isLeaf(leftChild)) {
-                let leafIdx = leafIndex(leftChild);
-                let triIdx = blasTriIds[triIdOffset + leafIdx];
-                let tri = blasTriangles[triOffset + triIdx];
-                let hit = intersectBLASTriangle(ray, tri);
-                if (hit.hit && hit.t < closest.t) {
-                    closest = hit;
-                }
-            } else {
-                let leftNode = blasNodes[nodeOffset + leftChild];
-                let leftMin = vec3(leftNode.minX, leftNode.minY, leftNode.minZ);
-                let leftMax = vec3(leftNode.maxX, leftNode.maxY, leftNode.maxZ);
-                let leftDist = intersectAABBDist(ray.origin, invDir, leftMin, leftMax);
-
-                if (leftDist < closest.t && stackPtr < MAX_STACK_DEPTH) {
-                    stack[stackPtr] = leftChild;
-                    stackPtr++;
-                }
+        children[0] = node.child0;
+        children[1] = node.child1;
+        children[2] = node.child2;
+        children[3] = node.child3;
+
+        dists[0] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c0_minX, node.c0_minY, node.c0_minZ),
+                vec3(node.c0_maxX, node.c0_maxY, node.c0_maxZ)),
+            1e30,
+            children[0] == INVALID_NODE
+        );
+        dists[1] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c1_minX, node.c1_minY, node.c1_minZ),
+                vec3(node.c1_maxX, node.c1_maxY, node.c1_maxZ)),
+            1e30,
+            children[1] == INVALID_NODE
+        );
+        dists[2] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c2_minX, node.c2_minY, node.c2_minZ),
+                vec3(node.c2_maxX, node.c2_maxY, node.c2_maxZ)),
+            1e30,
+            children[2] == INVALID_NODE
+        );
+        dists[3] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c3_minX, node.c3_minY, node.c3_minZ),
+                vec3(node.c3_maxX, node.c3_maxY, node.c3_maxZ)),
+            1e30,
+            children[3] == INVALID_NODE
+        );
+
+        for (var i = 1u; i < 4u; i++) {
+            let keyDist = dists[i];
+            let keyChild = children[i];
+            var j = i;
+            while (j > 0u && dists[j - 1u] > keyDist) {
+                dists[j] = dists[j - 1u];
+                children[j] = children[j - 1u];
+                j--;
             }
+            dists[j] = keyDist;
+            children[j] = keyChild;
         }
 
-        if (rightChild != INVALID_NODE) {
-            if (isLeaf(rightChild)) {
-                let leafIdx = leafIndex(rightChild);
+        for (var i = 3i; i >= 0i; i--) {
+            let child = children[i];
+            let dist = dists[i];
+
+            if (child == INVALID_NODE || dist >= closest.t) {
+                continue;
+            }
+
+            if (isLeaf(child)) {
+                let leafIdx = leafIndex(child);
                 let triIdx = blasTriIds[triIdOffset + leafIdx];
                 let tri = blasTriangles[triOffset + triIdx];
                 let hit = intersectBLASTriangle(ray, tri);
                 if (hit.hit && hit.t < closest.t) {
                     closest = hit;
                 }
-            } else {
-                let rightNode = blasNodes[nodeOffset + rightChild];
-                let rightMin = vec3(rightNode.minX, rightNode.minY, rightNode.minZ);
-                let rightMax = vec3(rightNode.maxX, rightNode.maxY, rightNode.maxZ);
-                let rightDist = intersectAABBDist(ray.origin, invDir, rightMin, rightMax);
-
-                if (rightDist < closest.t && stackPtr < MAX_STACK_DEPTH) {
-                    stack[stackPtr] = rightChild;
-                    stackPtr++;
-                }
+            } else if (stackPtr < MAX_STACK_DEPTH) {
+                stack[stackPtr] = child;
+                stackPtr++;
             }
         }
     }
@@ -334,9 +342,9 @@ fn trace(ray: Ray) -> HitResult {
                 let nodeOffset = blasMeta[shapeId * 4u];
                 let triIdOffset = blasMeta[shapeId * 4u + 1u];
                 let triOffset = blasMeta[shapeId * 4u + 2u];
-                let triCount_ = blasMeta[shapeId * 4u + 3u];
+                let totalTriCount = blasMeta[shapeId * 4u + 3u];
 
-                if (triCount_ == 0u) {
+                if (totalTriCount == 0u) {
                     continue;
                 }
 
@@ -345,7 +353,7 @@ fn trace(ray: Ray) -> HitResult {
                 objRay.origin = (invMatrix * vec4(ray.origin, 1.0)).xyz;
                 objRay.direction = (invMatrix * vec4(ray.direction, 0.0)).xyz;
 
-                let blasHit = traceBLAS(objRay, nodeOffset, triIdOffset, triOffset, triCount_, closest.t);
+                let blasHit = traceBLAS(objRay, nodeOffset, triIdOffset, triOffset, totalTriCount, closest.t);
 
                 if (blasHit.hit && blasHit.t < closest.t) {
                     closest = blasHit;
@@ -440,9 +448,9 @@ fn traceAnyHit(ray: Ray, tMax: f32) -> bool {
                 let nodeOffset = blasMeta[shapeId * 4u];
                 let triIdOffset = blasMeta[shapeId * 4u + 1u];
                 let triOffset = blasMeta[shapeId * 4u + 2u];
-                let triCount_ = blasMeta[shapeId * 4u + 3u];
+                let totalTriCount = blasMeta[shapeId * 4u + 3u];
 
-                if (triCount_ == 0u) {
+                if (totalTriCount == 0u) {
                     continue;
                 }
 
@@ -451,7 +459,7 @@ fn traceAnyHit(ray: Ray, tMax: f32) -> bool {
                 objRay.origin = (invMatrix * vec4(ray.origin, 1.0)).xyz;
                 objRay.direction = (invMatrix * vec4(ray.direction, 0.0)).xyz;
 
-                let blasHit = traceBLAS(objRay, nodeOffset, triIdOffset, triOffset, triCount_, tMax);
+                let blasHit = traceBLAS(objRay, nodeOffset, triIdOffset, triOffset, totalTriCount, tMax);
                 if (blasHit.hit && blasHit.t < tMax) {
                     return true;
                 }