From 340fffa264c45052e783e79c3cb97b21ab36898e Mon Sep 17 00:00:00 2001 From: Max Yang Date: Fri, 21 Aug 2026 17:11:49 -0700 Subject: [PATCH 01/11] Sort indices by morton index to try and speed up closest point kernels --- .../detail/DistributedClosestPointImpl.hpp | 64 +++++++++++++++++-- 1 file changed, 60 insertions(+), 4 deletions(-) diff --git a/src/axom/quest/detail/DistributedClosestPointImpl.hpp b/src/axom/quest/detail/DistributedClosestPointImpl.hpp index 5e7d427a99..2c39b06d40 100644 --- a/src/axom/quest/detail/DistributedClosestPointImpl.hpp +++ b/src/axom/quest/detail/DistributedClosestPointImpl.hpp @@ -8,8 +8,6 @@ #include "axom/config.hpp" #include "axom/core.hpp" -#include "axom/core/NumericLimits.hpp" -#include "axom/core/execution/runtime_policy.hpp" #include "axom/slic.hpp" #include "axom/primal.hpp" #include "axom/spin.hpp" @@ -1058,6 +1056,8 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl /// Create an ArrayView in ExecSpace that is compatible with queryPts PointArray execPoints(queryPts, m_allocatorID); auto query_pts = execPoints.view(); + auto query_order = mortonSortQueryPoints(query_pts, qPtCount); + auto query_order_view = query_order.view(); const double sqDistThreshold = m_sqDistanceThreshold; auto it = m_bvh->getTraverser(); const int rank = m_rank; @@ -1074,7 +1074,8 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl axom::ReduceMax maxSqDistance(currentMaxSqDistance); axom::for_all( qPtCount, - AXOM_LAMBDA(std::int32_t idx) mutable { + AXOM_LAMBDA(std::int32_t sorted_idx) { + const auto idx = query_order_view[sorted_idx]; PointType qpt = query_pts[idx]; MinCandidate curr_min {}; @@ -1131,7 +1132,8 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl AXOM_ANNOTATE_SCOPE("ComputeClosestPoints"); axom::for_all( qPtCount, - AXOM_LAMBDA(std::int32_t idx) mutable { + AXOM_LAMBDA(std::int32_t sorted_idx) { + const auto idx = query_order_view[sorted_idx]; PointType qpt = query_pts[idx]; MinCandidate curr_min {}; @@ -1210,6 +1212,60 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl } private: + /*! \brief Returns query point indices ordered by their Morton codes. */ + axom::Array mortonSortQueryPoints(const axom::ArrayView& queryPoints, + axom::IndexType queryPointCount) const + { + axom::Array queryOrder(queryPointCount, queryPointCount, m_allocatorID); + if(queryPointCount == 0) + { + return queryOrder; + } + + PointType minPoint; + PointType inverseExtent; + for(int dim = 0; dim < DIM; ++dim) + { + axom::ReduceMin minCoord(axom::numeric_limits::max()); + axom::ReduceMax maxCoord(axom::numeric_limits::lowest()); + axom::for_all( + queryPointCount, + AXOM_LAMBDA(axom::IndexType idx) { + minCoord.min(queryPoints[idx][dim]); + maxCoord.max(queryPoints[idx][dim]); + }); + + minPoint[dim] = minCoord.get(); + const double extent = maxCoord.get() - minPoint[dim]; + inverseExtent[dim] = extent > 0.0 ? 1.0 / extent : 0.0; + } + + axom::Array mortonCodes(queryPointCount, queryPointCount, m_allocatorID); + auto morton_codes = mortonCodes.view(); + auto query_order = queryOrder.view(); + axom::for_all( + queryPointCount, + AXOM_LAMBDA(axom::IndexType idx) { + constexpr int bits_per_dimension = 32 / DIM; + constexpr double coordinate_scale = 1 << bits_per_dimension; + constexpr double coordinate_max = coordinate_scale - 1.0; + + primal::Point gridPoint; + for(int dim = 0; dim < DIM; ++dim) + { + const double coordinate = (queryPoints[idx][dim] - minPoint[dim]) * inverseExtent[dim]; + gridPoint[dim] = static_cast( + axom::utilities::clampVal(coordinate * coordinate_scale, 0.0, coordinate_max)); + } + + morton_codes[idx] = spin::convertPointToMorton(gridPoint); + query_order[idx] = idx; + }); + + axom::stable_sort_pairs(morton_codes, query_order); + return queryOrder; + } + /*! @brief Object point coordindates array. From 12b7823f3f4ab6638a8a1ea0bf5842c6b2433222 Mon Sep 17 00:00:00 2001 From: Max Yang Date: Mon, 24 Aug 2026 15:59:29 -0700 Subject: [PATCH 02/11] Add a BVH2Node struct for improved HIP performance HIP architectures appear to do temporal coalescing; that is, memory accesses for array-of-struct types that would be individually strided can be coalesced so long as the loads are all issued before the next wait ("s_waitcnt vmcnt(*)"). In some cases, that makes loading the whole struct faster than through a struct-of-arrays layout. We should revisit this when we get a chance on Nvidia platforms. --- src/axom/spin/CMakeLists.txt | 1 + src/axom/spin/internal/linear_bvh/BVHNode.hpp | 39 +++++++ .../spin/internal/linear_bvh/bvh_traverse.hpp | 26 +++-- .../spin/internal/linear_bvh/bvh_vtkio.hpp | 13 +-- src/axom/spin/policy/LinearBVH.hpp | 105 ++++++------------ 5 files changed, 92 insertions(+), 92 deletions(-) create mode 100644 src/axom/spin/internal/linear_bvh/BVHNode.hpp diff --git a/src/axom/spin/CMakeLists.txt b/src/axom/spin/CMakeLists.txt index d696f649a8..9466772220 100644 --- a/src/axom/spin/CMakeLists.txt +++ b/src/axom/spin/CMakeLists.txt @@ -31,6 +31,7 @@ set( spin_headers UniformGrid.hpp ## internal + internal/linear_bvh/BVHNode.hpp internal/linear_bvh/RadixTree.hpp internal/linear_bvh/build_radix_tree.hpp internal/linear_bvh/bvh_traverse.hpp diff --git a/src/axom/spin/internal/linear_bvh/BVHNode.hpp b/src/axom/spin/internal/linear_bvh/BVHNode.hpp new file mode 100644 index 0000000000..5536e12846 --- /dev/null +++ b/src/axom/spin/internal/linear_bvh/BVHNode.hpp @@ -0,0 +1,39 @@ +// Copyright (c) 2017-2025, Lawrence Livermore National Security, LLC and +// other Axom Project Developers. See the top-level LICENSE file for details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +#ifndef Axom_Spin_BVHNode_HH +#define Axom_Spin_BVHNode_HH + +#include "axom/primal/geometry/BoundingBox.hpp" + +namespace axom +{ +namespace spin +{ +namespace internal +{ +namespace linear_bvh +{ + +/*! + * \brief Node structure for a 2-wide BVH tree. + */ +template +struct BVH2Node +{ + using BoxType = primal::BoundingBox; + + BoxType left; + BoxType right; + std::int32_t left_child; + std::int32_t right_child; +}; + +} // namespace linear_bvh +} // namespace internal +} // namespace spin +} // namespace axom + +#endif diff --git a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp index 9de97b7962..7bfae3aeb2 100644 --- a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp +++ b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp @@ -10,6 +10,7 @@ #include "axom/core/Macros.hpp" // for AXOM_HOST_DEVICE #include "axom/core/Types.hpp" // for axom types #include "axom/slic.hpp" // for SLIC macros +#include "axom/spin/internal/linear_bvh/BVHNode.hpp" #include // For template magic #include @@ -106,16 +107,15 @@ inline bool leaf_node(const std::int32_t& nodeIdx) { return (nodeIdx < 0); } * */ template -AXOM_HOST_DEVICE inline void bvh_traverse( - axom::ArrayView> inner_nodes, - axom::ArrayView inner_node_children, - axom::ArrayView leaf_nodes, - const PrimitiveType& p, - InBinCheck&& B, - LeafAction&& A, - TraversePref&& Comp) +AXOM_HOST_DEVICE inline void bvh_traverse(axom::ArrayView> inner_nodes, + axom::ArrayView leaf_nodes, + const PrimitiveType& p, + InBinCheck&& B, + LeafAction&& A, + TraversePref&& Comp) { using BBoxType = primal::BoundingBox; + using BVHNode = BVH2Node; // setup stack constexpr std::int32_t STACK_SIZE = 64; @@ -132,14 +132,16 @@ AXOM_HOST_DEVICE inline void bvh_traverse( // Traverse until we hit a leaf node or the barrier. while(!leaf_node(current_node)) { - BBoxType left_bin = inner_nodes[current_node + 0]; - BBoxType right_bin = inner_nodes[current_node + 1]; + BVHNode curr_node; + curr_node = inner_nodes[current_node]; + BBoxType left_bin = curr_node.left; + BBoxType right_bin = curr_node.right; const bool in_left = left_bin.isValid() ? invoke_InBinCheck(B, p, left_bin, current_node + 0) : false; const bool in_right = right_bin.isValid() ? invoke_InBinCheck(B, p, right_bin, current_node + 1) : false; - std::int32_t l_child = inner_node_children[current_node + 0]; - std::int32_t r_child = inner_node_children[current_node + 1]; + std::int32_t l_child = curr_node.left_child; + std::int32_t r_child = curr_node.right_child; bool swap = Comp(left_bin, right_bin, p); if(!in_left && !in_right) diff --git a/src/axom/spin/internal/linear_bvh/bvh_vtkio.hpp b/src/axom/spin/internal/linear_bvh/bvh_vtkio.hpp index 1b745d9e66..35d7188c74 100644 --- a/src/axom/spin/internal/linear_bvh/bvh_vtkio.hpp +++ b/src/axom/spin/internal/linear_bvh/bvh_vtkio.hpp @@ -152,8 +152,7 @@ void write_root(const primal::BoundingBox& root, //------------------------------------------------------------------------------ template -void write_recursive(ArrayView> inner_nodes, - ArrayView inner_node_children, +void write_recursive(ArrayView> inner_nodes, std::int32_t current_node, std::int32_t level, std::int32_t& numPoints, @@ -163,12 +162,12 @@ void write_recursive(ArrayView> inne std::ostringstream& levels) { // STEP 0: get the flat BVH bounding boxes - primal::BoundingBox l_box = inner_nodes[current_node + 0]; - primal::BoundingBox r_box = inner_nodes[current_node + 1]; + primal::BoundingBox l_box = inner_nodes[current_node].left; + primal::BoundingBox r_box = inner_nodes[current_node].right; // STEP 1: extract children information - std::int32_t l_child = inner_node_children[current_node + 0]; - std::int32_t r_child = inner_node_children[current_node + 1]; + std::int32_t l_child = inner_nodes[current_node].left_child; + std::int32_t r_child = inner_nodes[current_node].right_child; write_box(l_box, numPoints, numBins, nodes, cells); levels << level << std::endl; @@ -179,7 +178,6 @@ void write_recursive(ArrayView> inne if(l_child > -1) { write_recursive(inner_nodes, - inner_node_children, l_child, level + 1, numPoints, @@ -193,7 +191,6 @@ void write_recursive(ArrayView> inne if(r_child > -1) { write_recursive(inner_nodes, - inner_node_children, r_child, level + 1, numPoints, diff --git a/src/axom/spin/policy/LinearBVH.hpp b/src/axom/spin/policy/LinearBVH.hpp index 7db3f14ca5..702a3b4d0b 100644 --- a/src/axom/spin/policy/LinearBVH.hpp +++ b/src/axom/spin/policy/LinearBVH.hpp @@ -17,6 +17,7 @@ #include "axom/primal/geometry/Vector.hpp" // linear bvh includes +#include "axom/spin/internal/linear_bvh/BVHNode.hpp" #include "axom/spin/internal/linear_bvh/RadixTree.hpp" #include "axom/spin/internal/linear_bvh/build_radix_tree.hpp" #include "axom/spin/internal/linear_bvh/bvh_traverse.hpp" @@ -37,6 +38,9 @@ namespace policy { namespace lbvh = internal::linear_bvh; +template +using BVH2Node = lbvh::BVH2Node; + /* * \brief Interface for a BVH tree through a traversal operation (which * searches a tree based on a user-provided predicate) or a reduce @@ -70,12 +74,11 @@ class LinearBVHTraverser public: using BoxType = primal::BoundingBox; using PointType = primal::Point; + using BVHNode = BVH2Node; - LinearBVHTraverser(axom::ArrayView bboxes, - axom::ArrayView inner_node_children, + LinearBVHTraverser(axom::ArrayView nodes, axom::ArrayView leaf_nodes) - : m_inner_nodes(bboxes) - , m_inner_node_children(inner_node_children) + : m_inner_nodes(nodes) , m_leaf_nodes(leaf_nodes) { } @@ -93,13 +96,7 @@ class LinearBVHTraverser return sqDistL > sqDistR; }; - lbvh::bvh_traverse(m_inner_nodes, - m_inner_node_children, - m_leaf_nodes, - p, - predicate, - leaf_action, - traversePref); + lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, predicate, leaf_action, traversePref); } /* @@ -119,13 +116,7 @@ class LinearBVHTraverser return false; }; - lbvh::bvh_traverse(m_inner_nodes, - m_inner_node_children, - m_leaf_nodes, - p, - predicate, - leaf_action, - noTraversePref); + lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, predicate, leaf_action, noTraversePref); } /*! @@ -144,7 +135,8 @@ class LinearBVHTraverser int allocatorID = axom::getDefaultAllocatorID()) const { // Make a field over all of the nodes (the return field). - axom::Array reducedField(m_inner_nodes.size(), m_inner_nodes.size(), allocatorID); + const auto num_node_slots = 2 * m_inner_nodes.size(); + axom::Array reducedField(num_node_slots, num_node_slots, allocatorID); if constexpr(std::is_same_v) { @@ -191,7 +183,13 @@ class LinearBVHTraverser axom::ArrayView node_data, std::int32_t current_node) const { - auto child_index = m_inner_node_children[current_node]; + const auto& node = m_inner_nodes[current_node / 2]; + auto child_index = current_node % 2 == 0 ? node.left_child : node.right_child; + + if(child_index >= 0) + { + child_index *= 2; + } // Check if node is a leaf if(child_index < 0) @@ -209,8 +207,7 @@ class LinearBVHTraverser node_data[current_node] = node_data[child_index + 0] + node_data[child_index + 1]; } - axom::ArrayView m_inner_nodes; // BVH bins including leafs - axom::ArrayView m_inner_node_children; + axom::ArrayView m_inner_nodes; // BVH bins including leafs axom::ArrayView m_leaf_nodes; // leaf data }; @@ -230,6 +227,7 @@ class LinearBVH public: using TraverserType = LinearBVHTraverser; using BoundingBoxType = primal::BoundingBox; + using BVHNode = BVH2Node; LinearBVH() = default; @@ -269,27 +267,23 @@ class LinearBVH TraverserType getTraverserImpl() const { - return TraverserType(m_inner_nodes.view(), m_inner_node_children.view(), m_leaf_nodes.view()); + return TraverserType(m_inner_nodes.view(), m_leaf_nodes.view()); } private: void allocate(std::int32_t size, int allocID) { AXOM_ANNOTATE_SCOPE("LinearBVH::allocate"); - IndexType numInnerNodes = (size - 1) * 2; + IndexType numInnerNodes = size - 1; // Need to allocate this uninitialized, since primal::BoundingBox is // considered non-trivially-copyable on GCC 4.9.3 - m_inner_nodes = axom::Array(axom::ArrayOptions::Uninitialized {}, - numInnerNodes, - numInnerNodes, - allocID); - m_inner_node_children = axom::Array(numInnerNodes, numInnerNodes, allocID); + m_inner_nodes = + axom::Array(axom::ArrayOptions::Uninitialized {}, numInnerNodes, numInnerNodes, allocID); m_leaf_nodes = axom::Array(size, size, allocID); } bool m_initialized {false}; - axom::Array m_inner_nodes; // BVH bins including leafs - axom::Array m_inner_node_children; + axom::Array m_inner_nodes; // BVH bins including leafs axom::Array m_leaf_nodes; // leaf data primal::BoundingBox m_bounds; }; @@ -334,7 +328,6 @@ void LinearBVH::buildImpl(const BoxIndexable boxes, const auto inner_aabb_ptr = radix_tree.m_inner_aabbs.view(); const auto bvh_inner_nodes = m_inner_nodes.view(); - const auto bvh_inner_node_children = m_inner_node_children.view(); AXOM_ANNOTATE_BEGIN("emit_bvh_parents"); for_all( @@ -351,8 +344,6 @@ void LinearBVH::buildImpl(const BoxIndexable boxes, else { l_aabb = inner_aabb_ptr[lchild]; - // do the offset now - lchild *= 2; } std::int32_t rchild = rchildren_ptr[node]; @@ -364,16 +355,13 @@ void LinearBVH::buildImpl(const BoxIndexable boxes, else { r_aabb = inner_aabb_ptr[rchild]; - // do the offset now - rchild *= 2; } - const std::int32_t out_offset = node * 2; - bvh_inner_nodes[out_offset + 0] = l_aabb; - bvh_inner_nodes[out_offset + 1] = r_aabb; + bvh_inner_nodes[node].left = l_aabb; + bvh_inner_nodes[node].right = r_aabb; - bvh_inner_node_children[out_offset + 0] = lchild; - bvh_inner_node_children[out_offset + 1] = rchild; + bvh_inner_nodes[node].left_child = lchild; + bvh_inner_nodes[node].right_child = rchild; }); AXOM_ANNOTATE_END("emit_bvh_parents"); @@ -400,7 +388,6 @@ axom::Array LinearBVH::findCandidatesImp SLIC_ASSERT(m_initialized); const auto inner_nodes = m_inner_nodes.view(); - const auto inner_node_children = m_inner_node_children.view(); const auto leaf_nodes = m_leaf_nodes.view(); auto noTraversePref = [] AXOM_HOST_DEVICE(const BoundingBoxType&, @@ -421,13 +408,7 @@ axom::Array LinearBVH::findCandidatesImp auto leafAction = [&count](std::int32_t AXOM_UNUSED_PARAM(current_node), const std::int32_t* AXOM_UNUSED_PARAM(leaf_nodes)) { count++; }; - lbvh::bvh_traverse(inner_nodes, - inner_node_children, - leaf_nodes, - primitive, - predicate, - leafAction, - noTraversePref); + lbvh::bvh_traverse(inner_nodes, leaf_nodes, primitive, predicate, leafAction, noTraversePref); counts[i] = count; total_count_reduce += count; @@ -461,13 +442,7 @@ axom::Array LinearBVH::findCandidatesImp offset++; }; - lbvh::bvh_traverse(inner_nodes, - inner_node_children, - leaf_nodes, - obj, - predicate, - leafAction, - noTraversePref); + lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, predicate, leafAction, noTraversePref); }); AXOM_ANNOTATE_END("PASS[2]:fill_traversal"); @@ -491,13 +466,7 @@ axom::Array LinearBVH::findCandidatesImp current_offset++; }; - lbvh::bvh_traverse(inner_nodes, - inner_node_children, - leaf_nodes, - obj, - predicate, - leafAction, - noTraversePref); + lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, predicate, leafAction, noTraversePref); counts[i] = matching_leaves; }); AXOM_ANNOTATE_END("PASS[1]:fill_traversal"); @@ -530,15 +499,7 @@ void LinearBVH::writeVtkFileImpl(const std::string& // STEP 2: traverse the BVH and dump each bin constexpr std::int32_t ROOT = 0; - lbvh::write_recursive(m_inner_nodes, - m_inner_node_children, - ROOT, - 1, - numPoints, - numBins, - nodes, - cells, - levels); + lbvh::write_recursive(m_inner_nodes, ROOT, 1, numPoints, numBins, nodes, cells, levels); // STEP 3: write nodes ofs << "POINTS " << numPoints << " double\n"; From db86043b834d6b36b1df5d23c64f32b530681c86 Mon Sep 17 00:00:00 2001 From: Max Yang Date: Wed, 17 Dec 2025 15:46:29 -0800 Subject: [PATCH 03/11] BVH: add support for shared memory-based stack --- .../spin/internal/linear_bvh/bvh_traverse.hpp | 166 ++++++++++++++++-- src/axom/spin/policy/LinearBVH.hpp | 20 ++- 2 files changed, 163 insertions(+), 23 deletions(-) diff --git a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp index 7bfae3aeb2..4c59290eae 100644 --- a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp +++ b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp @@ -70,6 +70,136 @@ namespace linear_bvh AXOM_HOST_DEVICE inline bool leaf_node(const std::int32_t& nodeIdx) { return (nodeIdx < 0); } +struct BVHStack +{ +public: + constexpr static std::int32_t STACK_SIZE = 64; + constexpr static std::int32_t BARRIER = -2000000000; + + using LocalStack = std::int32_t[STACK_SIZE]; + AXOM_HOST_DEVICE BVHStack() { } + + AXOM_HOST_DEVICE void setLocalStack(LocalStack& local_stack) + { + stack_ptr = 0; + stack = &(local_stack[0]); + stack[stack_ptr] = BARRIER; + } + + AXOM_HOST_DEVICE std::int32_t pop() + { + std::int32_t top = stack[stack_ptr]; + stack_ptr--; + return top; + } + + AXOM_HOST_DEVICE void push(std::int32_t value) + { + stack_ptr++; + stack[stack_ptr] = value; + } + +private: + std::int32_t stack_ptr {-1}; + std::int32_t* stack {nullptr}; +}; + +template +struct SharedBVHStack +{ +public: + constexpr static std::int32_t CHUNK_SIZE = 4; + constexpr static std::int32_t STACK_SIZE = 16; + constexpr static std::int32_t SHMEM_SIZE_PER_THREAD = CHUNK_SIZE * 2; + constexpr static std::int32_t BARRIER = -2000000000; + + struct Chunk + { + std::int32_t values[CHUNK_SIZE]; + }; + + using LocalStack = Chunk[STACK_SIZE]; + + AXOM_DEVICE static int* Get_Shared_Mem_Buffer() + { + __shared__ int shmem_buf[SHMEM_SIZE_PER_THREAD * BlockSize]; + return &(shmem_buf[0]); + } + + AXOM_DEVICE SharedBVHStack() + : s_block_dim(blockDim.x) + , s_thread_id(threadIdx.x) + , s_stack(SharedBVHStack::Get_Shared_Mem_Buffer()) + { } + + AXOM_HOST_DEVICE void setLocalStack(LocalStack& local_stack) { g_stack = &(local_stack[0]); } + + AXOM_DEVICE std::int32_t pop() + { + // Shared is empty, try to refill a chunk. + if(s_ptr == 0 && g_ptr > 0) + { + --g_ptr; + Chunk g_top_chunk = g_stack[g_ptr]; + for(int i = 0; i < CHUNK_SIZE; i++) + { + shared_stack(s_ptr + i) = g_top_chunk.values[i]; + } + s_ptr += CHUNK_SIZE; + } + if(s_ptr > 0) + { + // Can pop directly from shared. + --s_ptr; + std::int32_t top = shared_stack(s_ptr); + return top; + } + else + { + // Empty stack, return barrier. + return BARRIER; + } + } + + AXOM_DEVICE void push(std::int32_t value) + { + if(s_ptr == 2 * CHUNK_SIZE) + { + // At capacity. Take bottom values and push onto global memory stack. + Chunk g_bottom_chunk; + for(int i = 0; i < CHUNK_SIZE; i++) + { + g_bottom_chunk.values[i] = shared_stack(i); + } + g_stack[g_ptr] = g_bottom_chunk; + g_ptr++; + // Move remaining stack values down. + for(int i = 0; i < CHUNK_SIZE; i++) + { + shared_stack(i) = shared_stack(i + CHUNK_SIZE); + } + s_ptr -= CHUNK_SIZE; + } + assert(s_ptr < 2 * CHUNK_SIZE); + // Push value onto shared stack. + shared_stack(s_ptr) = value; + s_ptr++; + } + +private: + AXOM_DEVICE std::int32_t& shared_stack(int index) + { + return s_stack[index * s_block_dim + s_thread_id]; + } + + std::int16_t s_block_dim; + std::int16_t s_thread_id; + std::int32_t s_ptr {0}; + std::int32_t* s_stack; + std::int32_t g_ptr {0}; + Chunk* g_stack; +}; + /*! * \brief Generic BVH traversal routine. * @@ -106,10 +236,17 @@ inline bool leaf_node(const std::int32_t& nodeIdx) { return (nodeIdx < 0); } * device and unified memory. * */ -template +template AXOM_HOST_DEVICE inline void bvh_traverse(axom::ArrayView> inner_nodes, axom::ArrayView leaf_nodes, const PrimitiveType& p, + TraverseStack& stack, InBinCheck&& B, LeafAction&& A, TraversePref&& Comp) @@ -118,16 +255,13 @@ AXOM_HOST_DEVICE inline void bvh_traverse(axom::ArrayView; // setup stack - constexpr std::int32_t STACK_SIZE = 64; - constexpr std::int32_t BARRIER = -2000000000; - std::int32_t todo[STACK_SIZE]; - std::int32_t stackptr = 0; - todo[stackptr] = BARRIER; + typename TraverseStack::LocalStack local_mem; + stack.setLocalStack(local_mem); std::int32_t found_leaf = 0; std::int32_t current_node = 0; - while(current_node != BARRIER) + while(current_node != TraverseStack::BARRIER) { // Traverse until we hit a leaf node or the barrier. while(!leaf_node(current_node)) @@ -147,8 +281,7 @@ AXOM_HOST_DEVICE inline void bvh_traverse(axom::ArrayView sqDistR; }; - lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, predicate, leaf_action, traversePref); + lbvh::BVHStack stack; + + lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, stack, predicate, leaf_action, traversePref); } /* @@ -116,7 +118,9 @@ class LinearBVHTraverser return false; }; - lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, predicate, leaf_action, noTraversePref); + lbvh::BVHStack stack; + + lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, stack, predicate, leaf_action, noTraversePref); } /*! @@ -408,7 +412,9 @@ axom::Array LinearBVH::findCandidatesImp auto leafAction = [&count](std::int32_t AXOM_UNUSED_PARAM(current_node), const std::int32_t* AXOM_UNUSED_PARAM(leaf_nodes)) { count++; }; - lbvh::bvh_traverse(inner_nodes, leaf_nodes, primitive, predicate, leafAction, noTraversePref); + lbvh::BVHStack stack; + + lbvh::bvh_traverse(inner_nodes, leaf_nodes, primitive, stack, predicate, leafAction, noTraversePref); counts[i] = count; total_count_reduce += count; @@ -442,7 +448,9 @@ axom::Array LinearBVH::findCandidatesImp offset++; }; - lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, predicate, leafAction, noTraversePref); + lbvh::BVHStack stack; + + lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, stack, predicate, leafAction, noTraversePref); }); AXOM_ANNOTATE_END("PASS[2]:fill_traversal"); @@ -466,7 +474,9 @@ axom::Array LinearBVH::findCandidatesImp current_offset++; }; - lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, predicate, leafAction, noTraversePref); + lbvh::BVHStack stack; + + lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, stack, predicate, leafAction, noTraversePref); counts[i] = matching_leaves; }); AXOM_ANNOTATE_END("PASS[1]:fill_traversal"); From 8426dc35d9b5fa9377cf3ba2c8219636d986c76d Mon Sep 17 00:00:00 2001 From: Max Yang Date: Mon, 24 Aug 2026 20:33:39 -0700 Subject: [PATCH 04/11] Add block size parameter to axom::execution_space --- src/axom/core/execution/internal/cuda_exec.hpp | 4 ++++ src/axom/core/execution/internal/hip_exec.hpp | 4 ++++ src/axom/core/execution/internal/omp_exec.hpp | 2 ++ src/axom/core/execution/internal/seq_exec.hpp | 2 ++ 4 files changed, 12 insertions(+) diff --git a/src/axom/core/execution/internal/cuda_exec.hpp b/src/axom/core/execution/internal/cuda_exec.hpp index 16b30fdf5c..a47e75afbc 100644 --- a/src/axom/core/execution/internal/cuda_exec.hpp +++ b/src/axom/core/execution/internal/cuda_exec.hpp @@ -55,6 +55,8 @@ struct execution_space> using atomic_policy = RAJA::cuda_atomic; using sync_policy = RAJA::cuda_synchronize; + static constexpr int BlockSize = BLOCK_SIZE; + static constexpr MemorySpace memory_space = MemorySpace::Device; AXOM_HOST_DEVICE static constexpr bool async() noexcept { return false; } @@ -97,6 +99,8 @@ struct execution_space> using atomic_policy = RAJA::cuda_atomic; using sync_policy = RAJA::cuda_synchronize; + static constexpr int BlockSize = BLOCK_SIZE; + static constexpr MemorySpace memory_space = MemorySpace::Device; AXOM_HOST_DEVICE static constexpr bool async() noexcept { return true; } diff --git a/src/axom/core/execution/internal/hip_exec.hpp b/src/axom/core/execution/internal/hip_exec.hpp index 87418778d7..6ff9bf9bdb 100644 --- a/src/axom/core/execution/internal/hip_exec.hpp +++ b/src/axom/core/execution/internal/hip_exec.hpp @@ -53,6 +53,8 @@ struct execution_space> using atomic_policy = RAJA::hip_atomic; using sync_policy = RAJA::hip_synchronize; + static constexpr int BlockSize = BLOCK_SIZE; + static constexpr MemorySpace memory_space = MemorySpace::Device; AXOM_HOST_DEVICE static constexpr bool async() noexcept { return false; } @@ -95,6 +97,8 @@ struct execution_space> using atomic_policy = RAJA::hip_atomic; using sync_policy = RAJA::hip_synchronize; + static constexpr int BlockSize = BLOCK_SIZE; + static constexpr MemorySpace memory_space = MemorySpace::Device; AXOM_HOST_DEVICE static constexpr bool async() noexcept { return true; } diff --git a/src/axom/core/execution/internal/omp_exec.hpp b/src/axom/core/execution/internal/omp_exec.hpp index 3da39efb96..5192fdee47 100644 --- a/src/axom/core/execution/internal/omp_exec.hpp +++ b/src/axom/core/execution/internal/omp_exec.hpp @@ -41,6 +41,8 @@ struct execution_space using atomic_policy = RAJA::omp_atomic; using sync_policy = RAJA::omp_synchronize; + static constexpr int BlockSize = 1; + #ifdef AXOM_USE_UMPIRE static constexpr MemorySpace memory_space = MemorySpace::Host; #else diff --git a/src/axom/core/execution/internal/seq_exec.hpp b/src/axom/core/execution/internal/seq_exec.hpp index ce9bac758b..ad3592f7a6 100644 --- a/src/axom/core/execution/internal/seq_exec.hpp +++ b/src/axom/core/execution/internal/seq_exec.hpp @@ -51,6 +51,8 @@ struct execution_space using sync_policy = void; + static constexpr int BlockSize = 1; + #ifdef AXOM_USE_UMPIRE static constexpr MemorySpace memory_space = MemorySpace::Host; #else From b180bf6e16250fac78046d1fb6b3199d14757084 Mon Sep 17 00:00:00 2001 From: Max Yang Date: Mon, 24 Aug 2026 20:39:48 -0700 Subject: [PATCH 05/11] BVH: add shared memory traversal interface --- .../detail/DistributedClosestPointImpl.hpp | 4 +- src/axom/spin/policy/LinearBVH.hpp | 62 ++++++++++++------- 2 files changed, 42 insertions(+), 24 deletions(-) diff --git a/src/axom/quest/detail/DistributedClosestPointImpl.hpp b/src/axom/quest/detail/DistributedClosestPointImpl.hpp index 2c39b06d40..f965e3c863 100644 --- a/src/axom/quest/detail/DistributedClosestPointImpl.hpp +++ b/src/axom/quest/detail/DistributedClosestPointImpl.hpp @@ -1107,7 +1107,7 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl return sqDist <= curr_min.sqDist && sqDist <= sqDistThreshold; }; - it.traverse_tree(qpt, checkMinDist, traversePredicate); + it.template traverseTreeShared(qpt, checkMinDist, traversePredicate); if(curr_min.rank == rank) { @@ -1167,7 +1167,7 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl }; // Traverse the tree, searching for the point with minimum distance. - it.traverse_tree(qpt, checkMinDist, traversePredicate); + it.template traverseTreeShared(qpt, checkMinDist, traversePredicate); // If modified, update the fields that changed if(curr_min.rank == rank) diff --git a/src/axom/spin/policy/LinearBVH.hpp b/src/axom/spin/policy/LinearBVH.hpp index 7369fa03b2..11accd3a74 100644 --- a/src/axom/spin/policy/LinearBVH.hpp +++ b/src/axom/spin/policy/LinearBVH.hpp @@ -82,18 +82,18 @@ class LinearBVHTraverser , m_leaf_nodes(leaf_nodes) { } - template - AXOM_HOST_DEVICE void traverse_tree(const PointType& p, + /* + * Functors \a leaf_action and \a predicate should access only memory compatible + * with the execution space. For example, GPU execution should access + * only device and unified memory. + */ + template + AXOM_HOST_DEVICE void traverse_tree(const Primitive& p, LeafAction&& leaf_action, Predicate&& predicate) const { - auto traversePref = [](const BoxType& l, const BoxType& r, const PointType& p) { - double sqDistL = primal::squared_distance(p, l.getCentroid()); - // If the right bbox is not valid, return max. Otherwise, the invalid right - // bbox might actually win when we should ignore it. - double sqDistR = r.isValid() ? primal::squared_distance(p, r.getCentroid()) - : axom::numerics::floating_point_limits::max(); - return sqDistL > sqDistR; + auto traversePref = [](const BoxType& l, const BoxType& r, const Primitive& p) { + return LinearBVHTraverser::traverseClosestFirst(l, r, p); }; lbvh::BVHStack stack; @@ -101,24 +101,22 @@ class LinearBVHTraverser lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, stack, predicate, leaf_action, traversePref); } - /* - * Functors \a leaf_action and \a predicate should access only memory compatible - * with the execution space. For example, GPU execution should access - * only device and unified memory. - */ - template - AXOM_HOST_DEVICE void traverse_tree(const Primitive& p, - LeafAction&& leaf_action, - Predicate&& predicate) const + template + AXOM_HOST_DEVICE void traverseTreeShared(const Primitive& p, + LeafAction&& leaf_action, + Predicate&& predicate) const { auto noTraversePref = [](const BoxType& l, const BoxType& r, const Primitive& p) { - AXOM_UNUSED_VAR(l); - AXOM_UNUSED_VAR(r); - AXOM_UNUSED_VAR(p); - return false; + return LinearBVHTraverser::traverseClosestFirst(l, r, p); }; + constexpr int BlockSize = axom::execution_space::BlockSize; + +#ifdef AXOM_DEVICE_CODE + lbvh::SharedBVHStack stack; +#else lbvh::BVHStack stack; +#endif lbvh::bvh_traverse(m_inner_nodes, m_leaf_nodes, p, stack, predicate, leaf_action, noTraversePref); } @@ -175,6 +173,26 @@ class LinearBVHTraverser } private: + template + AXOM_HOST_DEVICE static bool traverseClosestFirst(const BoxType& l, + const BoxType& r, + const PrimitiveType& p) + { + if constexpr(std::is_same_v) + { + double sqDistL = primal::squared_distance(p, l.getCentroid()); + // If the right bbox is not valid, return max. Otherwise, the invalid right + // bbox might actually win when we should ignore it. + double sqDistR = r.isValid() ? primal::squared_distance(p, r.getCentroid()) + : axom::numerics::floating_point_limits::max(); + return sqDistL > sqDistR; + } + else + { + return false; + } + } + /*! * \brief This is a helper method used in reduce_tree. * From aa6d92a5507e80987c1dd026c4553170ec5576f5 Mon Sep 17 00:00:00 2001 From: Max Yang Date: Mon, 24 Aug 2026 20:40:25 -0700 Subject: [PATCH 06/11] GPU optimizations for squared_distance(Point, BoundingBox) Use fmin/fmax to perform clamping operation. This avoids extra instructions that are generated to perform a generic ternary-based clamp operation. --- .../primal/operators/squared_distance.hpp | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/src/axom/primal/operators/squared_distance.hpp b/src/axom/primal/operators/squared_distance.hpp index e2efd9fbc1..009349fc81 100644 --- a/src/axom/primal/operators/squared_distance.hpp +++ b/src/axom/primal/operators/squared_distance.hpp @@ -85,16 +85,28 @@ AXOM_HOST_DEVICE inline double squared_distance(const Point& P, return axom::numerics::floating_point_limits::max(); } - if(B.contains(P)) + // compute closest point to the box + Point cp; + if constexpr(std::is_floating_point_v) { - return 0; + for(int i = 0; i < NDIMS; ++i) + { + cp[i] = fmax(B.getMin()[i], fmin(P[i], B.getMax()[i])); + } + } + else + { + for(int i = 0; i < NDIMS; ++i) + { + cp[i] = clampVal(P[i], B.getMin()[i], B.getMax()[i]); + } } - // compute closest point to the box - Point cp; - for(int i = 0; i < NDIMS; ++i) + // if clamped point is the same as the original point, our point + // was already in the bounding box + if(cp == P) { - cp[i] = clampVal(P[i], B.getMin()[i], B.getMax()[i]); + return 0; } // return squared distance to the closest point From 90cb8c4b1ab5ce932e39b8bcc5560c3323b8f58a Mon Sep 17 00:00:00 2001 From: Max Yang Date: Mon, 31 Aug 2026 17:24:18 -0700 Subject: [PATCH 07/11] BVH: use device-side traverser interface in findCandidatesImpl --- src/axom/spin/policy/LinearBVH.hpp | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/src/axom/spin/policy/LinearBVH.hpp b/src/axom/spin/policy/LinearBVH.hpp index 11accd3a74..862dec1e9c 100644 --- a/src/axom/spin/policy/LinearBVH.hpp +++ b/src/axom/spin/policy/LinearBVH.hpp @@ -409,12 +409,7 @@ axom::Array LinearBVH::findCandidatesImp SLIC_ERROR_IF(counts.size() != numObjs, "counts length not equal to numObjs"); SLIC_ASSERT(m_initialized); - const auto inner_nodes = m_inner_nodes.view(); - const auto leaf_nodes = m_leaf_nodes.view(); - - auto noTraversePref = [] AXOM_HOST_DEVICE(const BoundingBoxType&, - const BoundingBoxType&, - const PrimitiveType&) { return false; }; + TraverserType tree_view = this->getTraverserImpl(); #if defined(AXOM_USE_RAJA) // STEP 1: count number of candidates for each query point @@ -430,9 +425,7 @@ axom::Array LinearBVH::findCandidatesImp auto leafAction = [&count](std::int32_t AXOM_UNUSED_PARAM(current_node), const std::int32_t* AXOM_UNUSED_PARAM(leaf_nodes)) { count++; }; - lbvh::BVHStack stack; - - lbvh::bvh_traverse(inner_nodes, leaf_nodes, primitive, stack, predicate, leafAction, noTraversePref); + tree_view.traverse_tree(primitive, leafAction, predicate); counts[i] = count; total_count_reduce += count; @@ -466,9 +459,7 @@ axom::Array LinearBVH::findCandidatesImp offset++; }; - lbvh::BVHStack stack; - - lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, stack, predicate, leafAction, noTraversePref); + tree_view.traverse_tree(obj, leafAction, predicate); }); AXOM_ANNOTATE_END("PASS[2]:fill_traversal"); @@ -492,9 +483,8 @@ axom::Array LinearBVH::findCandidatesImp current_offset++; }; - lbvh::BVHStack stack; + tree_view.traverse_tree(obj, leafAction, predicate); - lbvh::bvh_traverse(inner_nodes, leaf_nodes, obj, stack, predicate, leafAction, noTraversePref); counts[i] = matching_leaves; }); AXOM_ANNOTATE_END("PASS[1]:fill_traversal"); From df7d91007d7056e79b82bf9b11c4e95bebbb4870 Mon Sep 17 00:00:00 2001 From: Max Yang Date: Mon, 31 Aug 2026 17:51:35 -0700 Subject: [PATCH 08/11] BVH: fix child slot index layouts used by GWN reduction --- src/axom/spin/internal/linear_bvh/bvh_traverse.hpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp index 4c59290eae..dec867a1cc 100644 --- a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp +++ b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp @@ -270,10 +270,13 @@ AXOM_HOST_DEVICE inline void bvh_traverse(axom::ArrayView Date: Mon, 31 Aug 2026 17:55:13 -0700 Subject: [PATCH 09/11] Fixup: guard SharedBVHStack on non-GPU platforms --- src/axom/spin/internal/linear_bvh/bvh_traverse.hpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp index dec867a1cc..691ff94e3f 100644 --- a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp +++ b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp @@ -104,6 +104,7 @@ struct BVHStack std::int32_t* stack {nullptr}; }; +#if defined(AXOM_USE_HIP) || defined(AXOM_USE_CUDA) template struct SharedBVHStack { @@ -199,6 +200,7 @@ struct SharedBVHStack std::int32_t g_ptr {0}; Chunk* g_stack; }; +#endif /*! * \brief Generic BVH traversal routine. From a561e5cf40a21ed8842887538633b3096b1787e2 Mon Sep 17 00:00:00 2001 From: Max Yang Date: Tue, 1 Sep 2026 08:26:31 -0700 Subject: [PATCH 10/11] Document internal BVH stack classes --- src/axom/spin/internal/linear_bvh/bvh_traverse.hpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp index 691ff94e3f..6153881b06 100644 --- a/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp +++ b/src/axom/spin/internal/linear_bvh/bvh_traverse.hpp @@ -70,6 +70,12 @@ namespace linear_bvh AXOM_HOST_DEVICE inline bool leaf_node(const std::int32_t& nodeIdx) { return (nodeIdx < 0); } +/*! + * \brief Implements a simple FIFO stack of integers for stackful BVH traversal. + * + * Local storage is allocated external to this class in an attempt to keep + * this class stored in registers on the GPU. + */ struct BVHStack { public: @@ -105,6 +111,14 @@ struct BVHStack }; #if defined(AXOM_USE_HIP) || defined(AXOM_USE_CUDA) +/*! + * \brief GPU-only class for a FIFO stack of integers. + * + * This version uses shared memory as an LRU cache to store the most-recently used + * entries. When the stack fills to capacity, we take a "chunk" of the oldest 4 integers + * and push them to the "local stack." Conversely, we can refill the shared memory stack + * on a pop() operation by popping a chunk of 4 integers from the local stack. + */ template struct SharedBVHStack { From bf4072b4545a9ff467ba69d128e6845f362cd45e Mon Sep 17 00:00:00 2001 From: Max Yang Date: Wed, 2 Sep 2026 11:27:37 -0700 Subject: [PATCH 11/11] Silence build warning --- src/axom/spin/policy/LinearBVH.hpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/axom/spin/policy/LinearBVH.hpp b/src/axom/spin/policy/LinearBVH.hpp index 862dec1e9c..26c3a537cc 100644 --- a/src/axom/spin/policy/LinearBVH.hpp +++ b/src/axom/spin/policy/LinearBVH.hpp @@ -115,6 +115,7 @@ class LinearBVHTraverser #ifdef AXOM_DEVICE_CODE lbvh::SharedBVHStack stack; #else + AXOM_UNUSED_VAR(BlockSize); lbvh::BVHStack stack; #endif