From aab744824c86caac2407ccaec44c7c7ca7515b69 Mon Sep 17 00:00:00 2001 From: Lesley Lai Date: Thu, 1 Jun 2023 15:58:15 +0800 Subject: [PATCH] Optimize BVH node layout --- src/lib/accelerators/bvh.cpp | 77 +++++++++++++++++++++++------------- src/lib/accelerators/bvh.hpp | 27 +++++++------ src/lib/path_tracer.cu | 8 ++-- 3 files changed, 69 insertions(+), 43 deletions(-) diff --git a/src/lib/accelerators/bvh.cpp b/src/lib/accelerators/bvh.cpp index 59fa89d..576622c 100644 --- a/src/lib/accelerators/bvh.cpp +++ b/src/lib/accelerators/bvh.cpp @@ -3,8 +3,11 @@ #include #include +#include #include +#include + 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> leaves) } } -void populate_linear_bvh(std::vector& linear_bvh, - const CPUBVHNode& node) -{ - if (node.is_leaf()) { - const auto leaf = static_cast(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(node); - - const auto current_index = static_cast(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(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 @@ -149,7 +140,39 @@ auto bvh_from_mesh(const Mesh& mesh) -> std::vector std::vector 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 nodes_to_process; + + auto push = [&](const CPUBVHNode& node) { + nodes_to_process.push(CPUNodeToProcess{ + .node = &node, + .linear_index = static_cast(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(current_node); + const auto left_index = static_cast(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; } \ No newline at end of file diff --git a/src/lib/accelerators/bvh.hpp b/src/lib/accelerators/bvh.hpp index c1acaaf..affa5ee 100644 --- a/src/lib/accelerators/bvh.hpp +++ b/src/lib/accelerators/bvh.hpp @@ -2,28 +2,31 @@ #define CUDA_PATH_TRACER_BVH_HPP #include "../aabb.hpp" +#include #include #include 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; #endif // CUDA_PATH_TRACER_BVH_HPP diff --git a/src/lib/path_tracer.cu b/src/lib/path_tracer.cu index 2216990..dd3f1b7 100644 --- a/src/lib/path_tracer.cu +++ b/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); }