Optimize BVH node layout

src/lib/accelerators/bvh.cpp

@@ -3,8 +3,11 @@
 
 #include <lib/prelude.hpp>
 #include <memory>
+#include <queue>
 #include <span>
 
+#include <fmt/format.h>
+
 namespace {
 
 struct CPUBVHNode {
@@ -15,6 +18,8 @@ struct CPUBVHNode {
   explicit CPUBVHNode(AABB aabb_) : aabb{aabb_} {}
 
   [[nodiscard]] virtual auto is_leaf() const -> bool = 0;
+
+  [[nodiscard]] virtual auto to_linear_node() const -> BVHNode = 0;
 };
 
 struct CPUBVHLeaf : CPUBVHNode {
@@ -26,6 +31,13 @@ struct CPUBVHLeaf : CPUBVHNode {
   }
 
   [[nodiscard]] auto is_leaf() const -> bool final { return true; }
+
+  [[nodiscard]] auto to_linear_node() const -> BVHNode final
+  {
+    return BVHNode{.aabb = aabb,
+                   .first_child_or_primitive = triangle_index_begin,
+                   .primitive_count = 1};
+  }
 };
 
 struct CPUBVHInner : CPUBVHNode {
@@ -40,6 +52,11 @@ struct CPUBVHInner : CPUBVHNode {
   }
 
   [[nodiscard]] auto is_leaf() const -> bool final { return false; }
+
+  [[nodiscard]] auto to_linear_node() const -> BVHNode final
+  {
+    return BVHNode{.aabb = aabb, .primitive_count = 0};
+  }
 };
 
 [[nodiscard]] auto
@@ -114,32 +131,6 @@ cpu_bvh_from_leaves(std::span<std::shared_ptr<CPUBVHLeaf>> leaves)
   }
 }
 
-void populate_linear_bvh(std::vector<BVHNode>& linear_bvh,
-                         const CPUBVHNode& node)
-{
-  if (node.is_leaf()) {
-    const auto leaf = static_cast<const CPUBVHLeaf&>(node);
-    linear_bvh.push_back(BVHNode{
-        .aabb = leaf.aabb,
-        .is_leaf = true,
-        .data = {.leaf = {.triangle_index_begin = leaf.triangle_index_begin}}});
-  } else { // Depth first: always go to left first
-    const auto inner = static_cast<const CPUBVHInner&>(node);
-
-    const auto current_index = static_cast<std::uint32_t>(linear_bvh.size());
-    linear_bvh.push_back(BVHNode{.aabb = inner.aabb, .is_leaf = false});
-    const auto left_index = current_index + 1;
-    populate_linear_bvh(linear_bvh, *inner.left);
-    const auto right_index = static_cast<std::uint32_t>(linear_bvh.size());
-    populate_linear_bvh(linear_bvh, *inner.right);
-
-    linear_bvh[current_index].data.inner = {
-        .left_index = left_index,
-        .right_index = right_index,
-    };
-  }
-}
-
 } // anonymous namespace
 
 auto bvh_from_mesh(const Mesh& mesh) -> std::vector<BVHNode>
@@ -149,7 +140,39 @@ auto bvh_from_mesh(const Mesh& mesh) -> std::vector<BVHNode>
   std::vector<BVHNode> linear_bvh;
   const auto size = mesh.triangle_count() * 2 - 1;
   linear_bvh.reserve(size);
-  populate_linear_bvh(linear_bvh, *root_node);
+
+  // Breath-first flatten the tree
+  struct CPUNodeToProcess {
+    const CPUBVHNode* node = nullptr;
+    std::uint32_t linear_index = 0;
+  };
+  std::queue<CPUNodeToProcess> nodes_to_process;
+
+  auto push = [&](const CPUBVHNode& node) {
+    nodes_to_process.push(CPUNodeToProcess{
+        .node = &node,
+        .linear_index = static_cast<std::uint32_t>(linear_bvh.size())});
+    linear_bvh.push_back(node.to_linear_node());
+  };
+
+  push(*root_node);
+
+  while (!nodes_to_process.empty()) {
+    const auto [current_node, index] = nodes_to_process.front();
+
+    if (!current_node->is_leaf()) {
+      const auto* inner_node = static_cast<const CPUBVHInner*>(current_node);
+      const auto left_index = static_cast<std::uint32_t>(linear_bvh.size());
+
+      // update children indices
+      linear_bvh[index].first_child_or_primitive = left_index;
+
+      push(*inner_node->left);
+      push(*inner_node->right);
+    }
+
+    nodes_to_process.pop();
+  }
 
   return linear_bvh;
 }

src/lib/accelerators/bvh.hpp

@@ -2,28 +2,31 @@
 #define CUDA_PATH_TRACER_BVH_HPP
 
 #include "../aabb.hpp"
+#include <cassert>
 #include <cstdint>
 #include <vector>
 
 struct Mesh;
 
+// A zero `primitive_count` means that we have an inner node.
+// - If it is a leaf, `first_child_or_primitive` is the index of the first
+// triangle
+// - If it is an inner node, `first_child_or_primitive` is the index of
+// the left child
+// And the index of the right child is `first_child_or_primitive + 1`
 struct BVHNode {
   AABB aabb;
-  bool is_leaf = false;
+  std::uint32_t first_child_or_primitive = 0;
+  std::uint32_t primitive_count = 0;
 
-  union {
-    struct {
-      std::uint32_t triangle_index_begin =
-          UINT32_MAX; // Index in the index buffer for the index of the first
-                      // vertex in triangle
-    } leaf;
-    struct {
-      std::uint32_t left_index = UINT32_MAX;  // Index of the left node
-      std::uint32_t right_index = UINT32_MAX; // Index of the right node
-    } inner;
-  } data;
+  [[nodiscard]] constexpr auto is_leaf() const noexcept
+  {
+    return primitive_count != 0;
+  }
 };
 
+static_assert(sizeof(BVHNode) == 32);
+
 auto bvh_from_mesh(const Mesh& mesh) -> std::vector<BVHNode>;
 
 #endif // CUDA_PATH_TRACER_BVH_HPP

src/lib/path_tracer.cu

@@ -54,9 +54,8 @@ __device__ auto ray_mesh_intersection_test(Ray ray, const glm::vec3* positions,
     const std::uint32_t node_index = node_stack.pop();
     const BVHNode node = bvh[node_index];
 
-    if (node.is_leaf) {
-      const std::uint32_t i = node.data.leaf.triangle_index_begin;
-
+    if (node.is_leaf()) {
+      const std::uint32_t i = node.first_child_or_primitive;
       const std::uint32_t index0 = indices[i];
       const std::uint32_t index1 = indices[i + 1];
       const std::uint32_t index2 = indices[i + 2];
@@ -71,7 +70,8 @@ __device__ auto ray_mesh_intersection_test(Ray ray, const glm::vec3* positions,
     } else {
       // Intersect AABB for an inner node
       if (ray_aabb_intersection_test(transformed_ray, node.aabb)) {
-        const auto [left_index, right_index] = node.data.inner;
+        const auto left_index = node.first_child_or_primitive;
+        const auto right_index = left_index + 1;
         node_stack.push(right_index);
         node_stack.push(left_index);
       }