BVH/mesh query optimizations — implementation walkthrough

Seven commits on branch ankac/bvh-query-perf (rebased onto upstream main c4675e0c, 2026-08-14) rework the traversal loops of Warp's BVH and mesh queries; the speedups were re-validated against that upstream tip after the rebase (CUDA median 1.31x, CPU median 1.47x, all results bitwise identical). Every change keeps query results bitwise identical to the old code (verified on 28 CUDA + 28 CPU result arrays per step), so all differences below are purely about how the same set of hits is found.

Benchmarks referenced throughout: "scene bench" = 200k queries against a 10x-replicated Stanford bunny (697k triangles) on an L40, per-launch time, lbvh/sah/cubql constructors, leaf_size 1 and 8 (this mirrors the ASV BvhAABBQuery/BvhRayQuery/MeshQuery matrix); "micro bench" = warp/examples/benchmarks/benchmark_bvh.py (128 queries, 10k random boxes — latency-bound, and its large query boxes overlap almost every node, the opposite regime of the scene bench). Percentages are reduction of running time measured when that change landed, on top of the previous changes.

1. Big picture: how a Warp BVH query runs

A kernel calls wp.bvh_query_aabb() / wp.bvh_query_ray() (or the mesh equivalents), then repeatedly calls wp.bvh_query_next() in a while loop; each call returns one overlapping primitive. On the C++ side this maps to a persistent iterator struct plus a next() function that resumes traversal:

The design lineage matters: mesh_query_ray was redesigned by Daniela Hasenbring in e33de231 (GH-1529, June 2026) — packed node payloads, children tested at the parent, near child carried in registers. Sections 4 and 5 generalize that traversal to the closest-point and iterator-style queries (her commits also added the bvh.h payload helpers this branch builds on), and section 7 extends her payload representation to the shared stack. mesh_query_ray itself is therefore untouched here: measured on the scene bench it runs 0.094-0.121 ms (leaf 1) / 0.180-0.240 ms (leaf 8), and three acceleration attempts all measured neutral or worse and were not kept: section 8's read-only-load treatment (neutral — its leaf loop needs every vertex component unconditionally, so it never had the branch-per-component problem), a shrunken-min_t pop-pruning stack in the style of section 4 (8-10% slower — near-first descent finds hits early and the push-time test already culls, so the bookkeeping is pure overhead for coherent rays), and hoisting the per-node fast/robust slab-test selection into a template (neutral — the per-node work is large enough to hide the branch).

Two structural facts drive every design decision here:

• The traversal stack lives in shared memory (32 slots x 4 bytes x 256 threads = 32 KB per block, bvh.h:548). Replacing it with a per-thread local array measured 2.7x slower despite doubling occupancy, so all changes work within its 32-bit slots.
• The query struct persists across next() calls, and codegen inlines everything into the kernel loop. If the struct contains a dynamically-indexed array, the compiler gives up on register promotion and spills the whole query state to local memory (a tried-and-rejected variant measured 4x slower). So the struct holds only scalars plus the pointer to the shared stack (bvh.h:497-513).

2. Summary: what changed, and what each change bought

Per-change impact, measured incrementally (each row on top of the previous ones) on the scene bench unless marked "micro". GPU is the primary target; CPU shares the same header code and is reported per change from a per-commit benchmark walk (sah constructor, biggest movers listed; "~0" means within run-to-run noise of about +/-2%).

Cumulative, all six changes vs main ebcce325 (full per-case chart below; tables in section 10):

3. Packed-leaf primitives now enumerate through scalar cursors instead of re-popping the leaf node (commit 4336fbc6)

Where this lives / what this part does. The leaf-emission step of the iterator queries: bvh_query_next_aabb (bvh.h:675-689), bvh_query_next_ray (bvh.h:604-617), and mesh_query_aabb_next (mesh.h:1997). When a BVH is built with leaf_size > 1, a leaf node covers a range of primitives, and the iterator must hand them to the caller one at a time, one per next() call.

Old behavior. The query struct kept only a primitive_counter. To emit the k-th primitive of a leaf, the code popped the leaf's index from the stack, loaded both 16-byte node halves again, skipped the AABB test, read one primitive, and then pushed the leaf back onto the stack so the next call could repeat the whole dance. An 8-primitive leaf cost 8 pops, 8 pushes and 16 redundant node loads. This made leaf_size=8 slower than leaf_size=1 across the board (ray, lbvh: 0.737 ms vs 0.330 ms) — the opposite of what packed leaves are for.

New behavior. The struct carries the current leaf's primitive range as two scalars (prim_cur, prim_end, bvh.h:507-508). While prim_cur < prim_end, each loop iteration reads one primitive index, tests its bounds, and either returns it or moves on — the stack and the node arrays are never touched again for that leaf.

Why it is faster. Per packed-leaf primitive it removes one shared-memory push, one pop, and two 16-byte node loads, replacing them with two scalar register updates. The same commit also deleted a dead "descend to the first leaf" loop in mesh_query_aabb() init, whose AABB test was guarded by primitive_counter == 0 while the counter started at -1 — it culled nothing and only pre-filled the stack with the whole left spine of the tree.

Impact. leaf_size=8 rays: -21% (lbvh 0.737 -> 0.579 ms, cubql 0.518 -> 0.406 ms); leaf_size=8 AABB queries: -7%; mesh_query_aabb also gained from the dead-init removal (cubql leaf_size=1: -12%, 0.0556 -> 0.0490 ms). leaf_size=1 BVH queries and the micro bench: unchanged (their leaves have one primitive and take the fast path). CPU: the init fix dominates — mesh_aabb -33% to -41% (the spine walk had no latency hiding there); rays -2% to -6%.

4. The four mesh_query_point variants share one shrinking-radius core with a distance-carrying stack (commit 4b4d78c0)

Where this lives / what this part does. mesh_query_point_core() (mesh.h:129-245) now implements the closest-point-on-mesh search used by mesh_query_point, mesh_query_point_sign_parity, mesh_query_point_no_sign and mesh_query_point_sign_winding_number (mesh.h:257, mesh.h:293, mesh.h:320, mesh.h:845). These are single-call queries (not iterators): one thread walks the tree once, keeping the best triangle found so far and shrinking the search radius as it goes.

Old behavior. Four near-identical ~160-line copies of the same loop. Each kept a stack of plain node indices. Processing a node meant: pop its index, load both node halves, recompute the point-to-box distance that had already been computed when the node was pushed, compare against the shrunken radius, and only then load the two children to decide where to go next. Every node's 32 bytes were loaded twice (once as a child of its parent, once when popped), and a pop that the shrinking radius had already invalidated still paid the full load + distance computation.

New behavior. One core, and the stack carries (packed payload, distance) pairs in two function-local arrays (mesh.h:135-136 — local arrays are fine here because this is a single-scope function, not a persistent struct). The child distances computed at the parent are stored with the entry; a pop that the radius has outrun is discarded by one register compare, touching no memory:

The near child never visits the stack at all — it is carried to the next iteration in registers (have_node/node), and the traversal order (nearer child first, ties to the right) is exactly the old one, which is why results stay bit-identical:

Why it is faster. Each node is loaded exactly once instead of ~1.5 times, the pop-time distance recomputation (9 FLOPs plus min/max per node) disappears, radius-pruned pops cost one compare instead of a 32-byte load, and roughly half the stack traffic goes away because near children ride in registers. Two slower shapes were tried and rejected on measurements: a nested descend-loop (leaf_size=8 regressed 4-11% from warp divergence) and a fully flat push-both-children version (12-byte pushes tripled local-memory stack traffic, 20% slower than the old code). The single flat loop with register-carried state was the only shape that won everywhere. It also deleted 655 lines of duplicated code.

Impact. Closest-point queries at leaf_size=1: -14% to -19% unsigned (lbvh 2.474 -> 2.027 ms, sah 2.292 -> 1.846 ms, cubql 1.937 -> 1.665 ms), -6% to -11% signed (the signed variants spend part of their time in the already-optimized ray-cast sign check, diluting the gain); leaf_size=8: -2% to -8%. CPU: -6% to -12% unsigned, -3% to -8% signed.

5. bvh_query_next split into a ray loop and an AABB loop; AABB descends through registers (commit 314644e1)

Where this lives / what this part does. bvh_query_next() (bvh.h:767) is now a one-time dispatch into two specialized flat loops: bvh_query_next_ray (bvh.h:599) and bvh_query_next_aabb (bvh.h:670); mesh_query_aabb_next (mesh.h:1990) uses the AABB shape. Query construction (bvh.h:561) seeds each shape differently.

Old behavior. One loop served both query types, re-checking a runtime is_ray flag (and dragging ray-only t/max_dist bookkeeping) at every node and every primitive test. Traversal itself was uniform test-on-pop: pop an index, load the node, test its AABB, and if it was an internal hit, push both children — every node in the tree costs a full pop / load / test / push round trip through the shared stack, including the child you are about to visit next anyway.

New behavior — AABB loop. When an internal node is processed, both children are loaded and tested at the parent. The surviving near child's packed payload (child indices + leaf flag, 63 bits) stays in registers (cur_node/have_node, bvh.h:512-513) and is processed by the very next iteration with no stack traffic and no reload; only the far child is pushed. Stack entries are therefore pre-tested — pops skip the AABB re-test:

New behavior — ray loop. Rays keep the classic test-on-pop shape (bvh.h:609) but with intersect_ray_aabb called directly, no is_ray flag checks. This asymmetry is deliberate and measured: a ray traveling through the scene overlaps both children of most nodes it visits, so the register-descent scheme degenerates (every node routes through the stack anyway, plus overhead) and measured ~25% slower for rays — while for the small query boxes of collision detection, most nodes have only one overlapping child and the register chain wins big.

Why it is faster. For AABB queries, loop iterations drop to roughly the number of overlapping nodes (children that fail the test are rejected at the parent and never enter the loop — previously every child cost a full iteration), miss-children never touch the stack, and the near-child chain runs entirely in registers. Removing the per-node is_ray branch and the dead t/max_dist bookkeeping from the AABB path measured a further ~20% — the AABB test itself is only ~6 compares, so the bookkeeping was a real fraction of the loop body.

Impact. Scene-bench AABB queries at leaf_size=1: -21% to -27% (bvh_aabb lbvh 0.0646 -> 0.0485 ms, cubql 0.0497 -> 0.0364 ms; mesh_aabb lbvh 0.0646 -> 0.0482 ms); leaf_size=8: -17%. Rays: unchanged (by design). CPU: bvh_aabb -32% to -35%, mesh_aabb -14% to -19%, rays -7% to -9% (fewer stack round-trips are cache-friendly there too). One casualty: the micro bench's AABB single-thread row regressed +42% (0.120 -> 0.171 ms) because with its huge query boxes nearly every node overlaps, so nearly every node took the far-child path — fixed by sections 6 and 7.

6. Popped nodes are processed in the same loop iteration (commit 5a1073f7)

Where this lives / what this part does. The pop step of the AABB loops (bvh.h:691, and the same shape in mesh_query_aabb_next and mesh_query_point_core).

Old behavior (as first landed in section 5). The loop spent one iteration just popping and loading a stack entry into cur_node, then continued, and the next iteration processed it. Harmless when stack pops are rare (sparse overlap), but in the dense both-children-overlap regime nearly every node arrives via the stack, so nearly every node cost two iterations.

New behavior. The pop falls through into the processing in the same iteration:

Why it is faster. Loop iterations are not free even when their body is small: each one re-evaluates the leaf-cursor check and the loop branches, and under SIMT a warp pays for the union of its lanes' iterations. Halving the trip count in the dense regime recovers exactly that.

Impact. Micro bench AABB single: -19% (0.171 -> 0.139 ms). Scene bench: within noise (pops are rare there). CPU: -1% to -3%.

7. Far children are pushed as tagged payload pairs, so pops touch no memory (commit 6d3be29e)

Where this lives / what this part does. The stack entry encoding of the AABB loops: helpers at bvh.h:377-387, push site at bvh.h:753, pop site at bvh.h:695. The shared stack has 32-bit slots and cannot be widened (its 32 KB per block is already the occupancy limiter, and a 64-bit stack would double it past the hardware's static shared-memory limit).

Old behavior (as first landed in section 5). The far child was pushed as its node index. On pop, recovering its child pointers meant re-loading both 16-byte node halves — 32 bytes of memory traffic to recover 9 bytes of payload that had already been in registers when the node was tested at its parent. In the dense regime this reload sat on the critical latency path of almost every node.

New behavior. The far child's packed payload itself goes onto the stack as two tagged 32-bit slots. Node indices are only 31 bits, so bit 31 of the top slot cleanly distinguishes a payload pair (bit set) from a plain index entry (bit clear); when fewer than two slots remain, the push falls back to the index form rather than losing the entry:

Why it is faster. A pop becomes two shared-memory reads plus a few shifts — no global-memory access at all, nothing for the latency chain to wait on. The cost is one extra shared slot per pending far child, which only matters when the stack is nearly full (then the index fallback kicks in; worst-case depth behavior is unchanged).

Impact. Micro bench AABB single: -16% (0.139 -> 0.117 ms), landing 3% below the original baseline — the section-5 regression fully paid back. Scene bench AABB: a further -3% to -6% (lbvh leaf1 0.0485 -> 0.0458 ms). CPU: +4% to +5% on AABB queries — the two-slot push and unpack are pure overhead where the stack already sits in L1; net CPU is still -33% to -51% vs baseline, and gating the pair encoding to CUDA would be a straightforward refinement if the CPU delta matters.

8. Primitive-level reads go through the read-only data path, and item bounds load before testing (commit 5c366e47)

Where this lives / what this part does. The bvh_load_int / bvh_load_vec3 helpers (bvh.h:254) and every primitive-level read in the query hot paths: leaf emission in both bvh loops (bvh.h:676-681, bvh.h:604), mesh_query_aabb_next (mesh.h:2010), and the triangle fetch of the closest-point leaf loop (mesh.h:172-180).

Old behavior. This change came from reading the generated SASS. Node loads were already ideal (LDG.E.128.CONSTANT — one vectorized read-only load per 16-byte half, thanks to the pre-existing __ldg in bvh_load_node). But every primitive-level read — primitive_indices, item_lowers/uppers, mesh.indices, mesh.points — compiled to a generic LD.E: those arrays are reached through the BVH/Mesh descriptor pointer, so the compiler cannot prove they are global memory, and generic loads probe the shared/local address windows and bypass the read-only cache. The closest-point leaf loop performed 13 such loads per triangle. Worse, in the AABB emit path the short-circuit || inside intersect_aabb_aabb over lazily-loaded components serialized into a dependent load -> test -> branch round per component — up to 6 rounds per primitive where one would do.

New behavior. All primitive-level reads go through two tiny helpers (__ldg on CUDA, plain loads on CPU — bit-identical values either way), and the item bounds are loaded into locals before the test:

Why it is faster. __ldg is side-effect-free, so the compiler hoists all six bound components into parallel loads and folds the AABB test into a single predicate chain (the same shape the node test already had) — the per-primitive latency chain shrinks from ~6 serial load+branch rounds to one parallel load batch plus one branch. The read-only path also caches better under heavy leaf traffic.

Impact. leaf_size=8 AABB queries: a further -18% to -26% (bvh_aabb lbvh r=0.002: 0.0728 -> 0.0598 ms, r=0.008: 0.0917 -> 0.0724 ms; cubql: 0.0600 -> 0.0466 / 0.0805 -> 0.0596 ms). leaf_size=1, rays, closest-point and the micro bench: unchanged (single-primitive leaves skip the emit loop entirely, and the closest-point loop was already latency-hidden and math-bound). CPU: -2% to -3% on AABB queries from the eager bound loads (__ldg compiles to a plain load on CPU).

9. Pair pushes are bounded by the tree depth so results can never be dropped (commit fed67817)

Where this lives / what this part does. bvh_query_pair_limit() (bvh.h:397-415), the pair-push sites in both AABB loops, and the builders: every constructor now records the tree depth (root = depth 1), the device builders behind a device pointer (BVH::max_depth_ptr) updated by an atomicMax in bvh.cu:427-430 so that in-place graph-captured rebuilds keep it current.

Old behavior (as introduced in section 7). Warp's constructors hard-terminate tree construction at the stack depth (bvh.cpp:534, bvh.cu:443, bvh_cubql.cpp:171), which makes the classic one-slot-per-entry stack exactly safe: pending entries never exceed 31. The two-slot pair encoding silently broke that calibration — ~15 pairs plus fallback slots cover only ~17 pending far children where a legally-deep tree can demand 31. On such trees the traversal did not overflow; it silently dropped far children. A degenerate tree of diagonally exponentially spaced boxes (long shared morton prefixes plus a duplicate cluster) reproduced one missing query result with the SAH constructor.

New behavior. Pair pushes stop once the stack holds more than 64 - 2 * max_depth slots and fall back to the index form. Pending entries never exceed max_depth - 1 and live pairs never exceed 33 - max_depth, so slot usage provably stays within the 32-slot stack for every constructor-produced tree. Unknown depths (and grouped host builds, whose depth counter restarts per group) disable pairs entirely — byte-for-byte the old exactly-safe behavior. Shallow trees, where the pairs matter, keep the full speedup.

Impact. Correctness, not speed: benchmark_bvh.py AABB single stays at 0.117 ms (its 10k-box tree is shallow, so pairs stay active), the scene bench is unchanged within noise on both devices, and a new degenerate-deep-tree regression test (test_bvh.py, all constructors, both devices, leaf sizes 1 and 4) fails without the fix and passes with it.

10. Cumulative impact

Final numbers vs the pre-series baseline (min of 20 reps; full tables in results/baseline_*_times.txt vs results/final_*_times.txt — note the final CUDA table predates section 7, whose extra leaf_size=8 gains are in results/v15_cuda_times.txt).

11. Findings

12. Test coverage