From d99b6cdd27b0e81c972711659da6beb9beed9e8e Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Sat, 6 May 2023 16:49:56 -0400 Subject: [PATCH 1/2] Introduce Value as a generic storage in a leaf node --- src/ArborX_LinearBVH.hpp | 5 +-- .../ArborX_DetailsDistributedTreeImpl.hpp | 3 +- src/details/ArborX_DetailsHalfTraversal.hpp | 5 ++- .../ArborX_DetailsHappyTreeFriends.hpp | 8 ++--- src/details/ArborX_DetailsNode.hpp | 30 +++++++++--------- .../ArborX_DetailsTreeConstruction.hpp | 14 ++++----- src/details/ArborX_DetailsTreeTraversal.hpp | 16 +++++----- .../ArborX_DetailsTreeVisualization.hpp | 5 ++- src/details/ArborX_MinimumSpanningTree.hpp | 31 +++++++++---------- test/tstDetailsTreeConstruction.cpp | 3 +- 10 files changed, 60 insertions(+), 60 deletions(-) diff --git a/src/ArborX_LinearBVH.hpp b/src/ArborX_LinearBVH.hpp index 18ac21c0c..45755f061 100644 --- a/src/ArborX_LinearBVH.hpp +++ b/src/ArborX_LinearBVH.hpp @@ -85,7 +85,8 @@ class BasicBoundingVolumeHierarchy private: friend struct Details::HappyTreeFriends; - using leaf_node_type = Details::LeafNode; + using value_type = Details::PairIndexVolume; + using leaf_node_type = Details::LeafNode; using internal_node_type = Details::InternalNode; KOKKOS_FUNCTION @@ -98,7 +99,7 @@ class BasicBoundingVolumeHierarchy assert((n == 1 || Details::HappyTreeFriends::getRoot(*this) == n) && "workaround below assumes root is stored as first element"); return (n > 1 ? &_internal_nodes.data()->bounding_volume - : &_leaf_nodes.data()->bounding_volume); + : &_leaf_nodes.data()->value.bounding_volume); } size_type _size{0}; diff --git a/src/details/ArborX_DetailsDistributedTreeImpl.hpp b/src/details/ArborX_DetailsDistributedTreeImpl.hpp index 20d2a540b..d2da7be6c 100644 --- a/src/details/ArborX_DetailsDistributedTreeImpl.hpp +++ b/src/details/ArborX_DetailsDistributedTreeImpl.hpp @@ -458,8 +458,7 @@ struct CallbackWithDistance "compute_reverse_permutation", Kokkos::RangePolicy(exec_space, 0, n), ARBORX_CLASS_LAMBDA(int const i) { - _rev_permute(HappyTreeFriends::getLeafPermutationIndex(_tree, i)) = - i; + _rev_permute(HappyTreeFriends::getValue(_tree, i).index) = i; }); } } diff --git a/src/details/ArborX_DetailsHalfTraversal.hpp b/src/details/ArborX_DetailsHalfTraversal.hpp index 5ed7886c7..3ad45256f 100644 --- a/src/details/ArborX_DetailsHalfTraversal.hpp +++ b/src/details/ArborX_DetailsHalfTraversal.hpp @@ -54,8 +54,7 @@ struct HalfTraversal { auto const predicate = _get_predicate(HappyTreeFriends::getLeafBoundingVolume(_bvh, i)); - auto const leaf_permutation_i = - HappyTreeFriends::getLeafPermutationIndex(_bvh, i); + auto const leaf_permutation_i = HappyTreeFriends::getValue(_bvh, i).index; int node = HappyTreeFriends::getRope(_bvh, i); while (node != ROPE_SENTINEL) @@ -70,7 +69,7 @@ struct HalfTraversal if (is_leaf) { _callback(leaf_permutation_i, - HappyTreeFriends::getLeafPermutationIndex(_bvh, node)); + HappyTreeFriends::getValue(_bvh, node).index); node = HappyTreeFriends::getRope(_bvh, node); } else diff --git a/src/details/ArborX_DetailsHappyTreeFriends.hpp b/src/details/ArborX_DetailsHappyTreeFriends.hpp index d861b8bdd..5af17a55c 100644 --- a/src/details/ArborX_DetailsHappyTreeFriends.hpp +++ b/src/details/ArborX_DetailsHappyTreeFriends.hpp @@ -70,15 +70,15 @@ struct HappyTreeFriends { static_assert( std::is_same_v); - return bvh._leaf_nodes(i).bounding_volume; + decltype(bvh._leaf_nodes(0).value.bounding_volume)>); + return bvh._leaf_nodes(i).value.bounding_volume; } template - static KOKKOS_FUNCTION auto getLeafPermutationIndex(BVH const &bvh, int i) + static KOKKOS_FUNCTION auto const &getValue(BVH const &bvh, int i) { assert(i >= 0 && i < (int)bvh.size()); - return bvh._leaf_nodes(i).permutation_index; + return bvh._leaf_nodes(i).value; } template diff --git a/src/details/ArborX_DetailsNode.hpp b/src/details/ArborX_DetailsNode.hpp index b9d2c61ea..262e76869 100644 --- a/src/details/ArborX_DetailsNode.hpp +++ b/src/details/ArborX_DetailsNode.hpp @@ -17,24 +17,27 @@ #include #include -#include // UINT_MAX #include // std::move -namespace ArborX -{ -namespace Details +namespace ArborX::Details { constexpr int ROPE_SENTINEL = -1; template +struct PairIndexVolume +{ + unsigned index; + BoundingVolume bounding_volume; +}; + +template struct LeafNode { - using bounding_volume_type = BoundingVolume; + using value_type = Value; - unsigned permutation_index = UINT_MAX; int rope = ROPE_SENTINEL; - BoundingVolume bounding_volume; + Value value; }; template @@ -48,14 +51,13 @@ struct InternalNode BoundingVolume bounding_volume; }; -template -KOKKOS_INLINE_FUNCTION constexpr LeafNode -makeLeafNode(unsigned permutation_index, - BoundingVolume bounding_volume) noexcept +template +KOKKOS_INLINE_FUNCTION constexpr LeafNode +makeLeafNode(Value value) noexcept { - return {permutation_index, ROPE_SENTINEL, std::move(bounding_volume)}; + return {ROPE_SENTINEL, std::move(value)}; } -} // namespace Details -} // namespace ArborX + +} // namespace ArborX::Details #endif diff --git a/src/details/ArborX_DetailsTreeConstruction.hpp b/src/details/ArborX_DetailsTreeConstruction.hpp index 826f384dc..abdc59d16 100644 --- a/src/details/ArborX_DetailsTreeConstruction.hpp +++ b/src/details/ArborX_DetailsTreeConstruction.hpp @@ -76,19 +76,18 @@ inline void initializeSingleLeafNode(ExecutionSpace const &space, Nodes const &leaf_nodes) { using Access = AccessTraits; + using Value = typename Nodes::value_type::value_type; + using BoundingVolume = decltype(std::declval().bounding_volume); ARBORX_ASSERT(leaf_nodes.extent(0) == 1); ARBORX_ASSERT(Access::size(primitives) == 1); - using Node = typename Nodes::value_type; - using BoundingVolume = typename Node::bounding_volume_type; - Kokkos::parallel_for( "ArborX::TreeConstruction::initialize_single_leaf", Kokkos::RangePolicy(space, 0, 1), KOKKOS_LAMBDA(int) { BoundingVolume bounding_volume{}; expand(bounding_volume, Access::get(primitives, 0)); - leaf_nodes(0) = makeLeafNode(0, std::move(bounding_volume)); + leaf_nodes(0) = makeLeafNode(Value{(unsigned)0, bounding_volume}); }); } @@ -209,8 +208,9 @@ class GenerateHierarchy expand(bounding_volume, Access::get(_primitives, original_index)); // Initialize leaf node + using Value = typename LeafNodes::value_type::value_type; auto &leaf_node = _leaf_nodes(i); - leaf_node = makeLeafNode(original_index, bounding_volume); + leaf_node = makeLeafNode(Value{original_index, bounding_volume}); // For a leaf node, the range is just one index int range_left = i; @@ -271,7 +271,7 @@ class GenerateHierarchy Kokkos::load_fence(); expand(bounding_volume, right_child_is_leaf - ? _leaf_nodes(right_child).bounding_volume + ? _leaf_nodes(right_child).value.bounding_volume : _internal_nodes(right_child).bounding_volume); } else @@ -294,7 +294,7 @@ class GenerateHierarchy Kokkos::load_fence(); expand(bounding_volume, left_child_is_leaf - ? _leaf_nodes(left_child).bounding_volume + ? _leaf_nodes(left_child).value.bounding_volume : _internal_nodes(left_child).bounding_volume); if (!left_child_is_leaf) diff --git a/src/details/ArborX_DetailsTreeTraversal.hpp b/src/details/ArborX_DetailsTreeTraversal.hpp index ce0cc4377..5535b370d 100644 --- a/src/details/ArborX_DetailsTreeTraversal.hpp +++ b/src/details/ArborX_DetailsTreeTraversal.hpp @@ -81,7 +81,7 @@ struct TreeTraversal HappyTreeFriends::getLeafBoundingVolume(_bvh, root); if (predicate(root_bounding_volume)) { - _callback(predicate, 0); + _callback(predicate, HappyTreeFriends::getValue(_bvh, 0).index); } } @@ -103,7 +103,7 @@ struct TreeTraversal { if (invoke_callback_and_check_early_exit( _callback, predicate, - HappyTreeFriends::getLeafPermutationIndex(_bvh, node))) + HappyTreeFriends::getValue(_bvh, node).index)) return; node = HappyTreeFriends::getRope(_bvh, node); } @@ -218,7 +218,7 @@ struct TreeTraversal if (k < 1) return; - _callback(predicate, 0); + _callback(predicate, HappyTreeFriends::getValue(_bvh, 0).index); } KOKKOS_FUNCTION void operator()(int queryIndex) const @@ -378,9 +378,9 @@ struct TreeTraversal sortHeap(heap.data(), heap.data() + heap.size(), heap.valueComp()); for (decltype(heap.size()) i = 0; i < heap.size(); ++i) { - int const leaf_index = HappyTreeFriends::getLeafPermutationIndex( - _bvh, (heap.data() + i)->first); - _callback(predicate, leaf_index); + _callback( + predicate, + HappyTreeFriends::getValue(_bvh, (heap.data() + i)->first).index); } } }; @@ -440,7 +440,7 @@ struct TreeTraversal::value; if (distance(getGeometry(predicate), root_bounding_volume) != inf) { - _callback(predicate, 0); + _callback(predicate, HappyTreeFriends::getValue(_bvh, 0).index); } } @@ -488,7 +488,7 @@ struct TreeTraversal(space, 0, n), KOKKOS_LAMBDA(int i) { - permute(HappyTreeFriends::getLeafPermutationIndex(tree, i)) = i; + permute(HappyTreeFriends::getValue(tree, i).index) = i; }); Predicates predicates(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, diff --git a/src/details/ArborX_MinimumSpanningTree.hpp b/src/details/ArborX_MinimumSpanningTree.hpp index a2bd8ec48..60d60973c 100644 --- a/src/details/ArborX_MinimumSpanningTree.hpp +++ b/src/details/ArborX_MinimumSpanningTree.hpp @@ -190,8 +190,7 @@ struct FindComponentNearestNeighbors auto const predicate = [label_i = component, &labels = _labels](int j) { return label_i != labels(j); }; - auto const leaf_permutation_i = - HappyTreeFriends::getLeafPermutationIndex(_bvh, i); + auto const leaf_permutation_i = HappyTreeFriends::getValue(_bvh, i).index; DirectedEdge current_best{}; @@ -245,10 +244,10 @@ struct FindComponentNearestNeighbors { if (HappyTreeFriends::isLeaf(_bvh, left_child)) { - float const candidate_dist = _metric( - leaf_permutation_i, - HappyTreeFriends::getLeafPermutationIndex(_bvh, left_child), - distance_left); + float const candidate_dist = + _metric(leaf_permutation_i, + HappyTreeFriends::getValue(_bvh, left_child).index, + distance_left); DirectedEdge const candidate_edge{i, left_child, candidate_dist}; if (candidate_edge < current_best) { @@ -270,10 +269,10 @@ struct FindComponentNearestNeighbors { if (HappyTreeFriends::isLeaf(_bvh, right_child)) { - float const candidate_dist = _metric( - leaf_permutation_i, - HappyTreeFriends::getLeafPermutationIndex(_bvh, right_child), - distance_right); + float const candidate_dist = + _metric(leaf_permutation_i, + HappyTreeFriends::getValue(_bvh, right_child).index, + distance_right); DirectedEdge const candidate_edge{i, right_child, candidate_dist}; if (candidate_edge < current_best) { @@ -493,9 +492,9 @@ void finalizeEdges(ExecutionSpace const &space, BVH const &bvh, Kokkos::RangePolicy(space, 0, n - 1), KOKKOS_LAMBDA(int i) { edges(i).source = - HappyTreeFriends::getLeafPermutationIndex(bvh, edges(i).source); + HappyTreeFriends::getValue(bvh, edges(i).source).index; edges(i).target = - HappyTreeFriends::getLeafPermutationIndex(bvh, edges(i).target); + HappyTreeFriends::getValue(bvh, edges(i).target).index; }); } @@ -552,8 +551,8 @@ void assignVertexParents(ExecutionSpace const &space, Labels const &labels, auto const &edge = out_edges(e); int i = labels(edge.source()); - parents(HappyTreeFriends::getLeafPermutationIndex(bvh, i) + - vertices_offset) = edges_mapping(i); + parents(HappyTreeFriends::getValue(bvh, i).index + vertices_offset) = + edges_mapping(i); }); } @@ -681,8 +680,8 @@ void resetSharedRadii(ExecutionSpace const &space, BVH const &bvh, if (label_i != label_j) { auto const r = - metric(HappyTreeFriends::getLeafPermutationIndex(bvh, i), - HappyTreeFriends::getLeafPermutationIndex(bvh, j), + metric(HappyTreeFriends::getValue(bvh, i).index, + HappyTreeFriends::getValue(bvh, j).index, distance(HappyTreeFriends::getLeafBoundingVolume(bvh, i), HappyTreeFriends::getLeafBoundingVolume(bvh, j))); Kokkos::atomic_min(&radii(label_i), r); diff --git a/test/tstDetailsTreeConstruction.cpp b/test/tstDetailsTreeConstruction.cpp index 4e3f6afbc..0f7570cd9 100644 --- a/test/tstDetailsTreeConstruction.cpp +++ b/test/tstDetailsTreeConstruction.cpp @@ -229,7 +229,8 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(example_tree_construction, DeviceType, // clang-format on BOOST_TEST_MESSAGE("ref = " << ref.str()); - using LeafNode = ArborX::Details::LeafNode; + using LeafNode = ArborX::Details::LeafNode< + ArborX::Details::PairIndexVolume>; using InternalNode = ArborX::Details::InternalNode; Kokkos::View leaf_nodes("Testing::leaf_nodes", 0); From 13c835a9a261132a976663c58e550c2260e855cb Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Wed, 17 May 2023 11:06:26 -0400 Subject: [PATCH 2/2] Introduce LegacyCallbackWrapper --- src/ArborX_LinearBVH.hpp | 9 ++++++--- src/details/ArborX_Callbacks.hpp | 13 +++++++++++++ src/details/ArborX_DetailsHalfTraversal.hpp | 6 +++--- src/details/ArborX_DetailsTreeTraversal.hpp | 17 +++++++---------- src/details/ArborX_DetailsTreeVisualization.hpp | 5 +++-- 5 files changed, 32 insertions(+), 18 deletions(-) diff --git a/src/ArborX_LinearBVH.hpp b/src/ArborX_LinearBVH.hpp index 45755f061..187d342cc 100644 --- a/src/ArborX_LinearBVH.hpp +++ b/src/ArborX_LinearBVH.hpp @@ -47,6 +47,7 @@ class BasicBoundingVolumeHierarchy static_assert(Kokkos::is_memory_space::value); using size_type = typename MemorySpace::size_type; using bounding_volume_type = BoundingVolume; + using value_type = Details::PairIndexVolume; BasicBoundingVolumeHierarchy() = default; // build an empty tree @@ -85,7 +86,6 @@ class BasicBoundingVolumeHierarchy private: friend struct Details::HappyTreeFriends; - using value_type = Details::PairIndexVolume; using leaf_node_type = Details::LeafNode; using internal_node_type = Details::InternalNode; @@ -213,7 +213,8 @@ template template void BasicBoundingVolumeHierarchy::query( ExecutionSpace const &space, Predicates const &predicates, - Callback const &callback, Experimental::TraversalPolicy const &policy) const + Callback const &legacy_callback, + Experimental::TraversalPolicy const &policy) const { static_assert( KokkosExt::is_accessible_from::value); @@ -222,7 +223,9 @@ void BasicBoundingVolumeHierarchy::query( static_assert(KokkosExt::is_accessible_from::value, "Predicates must be accessible from the execution space"); - Details::check_valid_callback(callback, predicates); + Details::check_valid_callback(legacy_callback, predicates); + Details::LegacyCallbackWrapper callback{ + legacy_callback}; using Tag = typename Details::AccessTraitsHelper::tag; std::string profiling_prefix = "ArborX::BVH::query::"; diff --git a/src/details/ArborX_Callbacks.hpp b/src/details/ArborX_Callbacks.hpp index ca39ccb7d..4235f6fe5 100644 --- a/src/details/ArborX_Callbacks.hpp +++ b/src/details/ArborX_Callbacks.hpp @@ -208,6 +208,19 @@ void check_valid_callback(Callback const &callback, Predicates const &) "Callback 'operator()' return type must be void"); } +template +struct LegacyCallbackWrapper +{ + Callback _callback; + + template + KOKKOS_FUNCTION auto operator()(Predicate const &predicate, + Value const &value) const + { + return _callback(predicate, value.index); + } +}; + } // namespace Details } // namespace ArborX diff --git a/src/details/ArborX_DetailsHalfTraversal.hpp b/src/details/ArborX_DetailsHalfTraversal.hpp index 3ad45256f..ef537717d 100644 --- a/src/details/ArborX_DetailsHalfTraversal.hpp +++ b/src/details/ArborX_DetailsHalfTraversal.hpp @@ -12,6 +12,7 @@ #ifndef ARBORX_DETAILS_HALF_TRAVERSAL_HPP #define ARBORX_DETAILS_HALF_TRAVERSAL_HPP +#include // LegacyCallbackWrapper #include #include // ROPE_SENTINEL @@ -25,7 +26,7 @@ struct HalfTraversal { BVH _bvh; PredicateGetter _get_predicate; - Callback _callback; + LegacyCallbackWrapper _callback; template HalfTraversal(ExecutionSpace const &space, BVH const &bvh, @@ -68,8 +69,7 @@ struct HalfTraversal { if (is_leaf) { - _callback(leaf_permutation_i, - HappyTreeFriends::getValue(_bvh, node).index); + _callback(leaf_permutation_i, HappyTreeFriends::getValue(_bvh, node)); node = HappyTreeFriends::getRope(_bvh, node); } else diff --git a/src/details/ArborX_DetailsTreeTraversal.hpp b/src/details/ArborX_DetailsTreeTraversal.hpp index 5535b370d..7d3227cc4 100644 --- a/src/details/ArborX_DetailsTreeTraversal.hpp +++ b/src/details/ArborX_DetailsTreeTraversal.hpp @@ -81,7 +81,7 @@ struct TreeTraversal HappyTreeFriends::getLeafBoundingVolume(_bvh, root); if (predicate(root_bounding_volume)) { - _callback(predicate, HappyTreeFriends::getValue(_bvh, 0).index); + _callback(predicate, HappyTreeFriends::getValue(_bvh, 0)); } } @@ -102,8 +102,7 @@ struct TreeTraversal if (is_leaf) { if (invoke_callback_and_check_early_exit( - _callback, predicate, - HappyTreeFriends::getValue(_bvh, node).index)) + _callback, predicate, HappyTreeFriends::getValue(_bvh, node))) return; node = HappyTreeFriends::getRope(_bvh, node); } @@ -218,7 +217,7 @@ struct TreeTraversal if (k < 1) return; - _callback(predicate, HappyTreeFriends::getValue(_bvh, 0).index); + _callback(predicate, HappyTreeFriends::getValue(_bvh, 0)); } KOKKOS_FUNCTION void operator()(int queryIndex) const @@ -378,9 +377,8 @@ struct TreeTraversal sortHeap(heap.data(), heap.data() + heap.size(), heap.valueComp()); for (decltype(heap.size()) i = 0; i < heap.size(); ++i) { - _callback( - predicate, - HappyTreeFriends::getValue(_bvh, (heap.data() + i)->first).index); + _callback(predicate, + HappyTreeFriends::getValue(_bvh, (heap.data() + i)->first)); } } }; @@ -440,7 +438,7 @@ struct TreeTraversal::value; if (distance(getGeometry(predicate), root_bounding_volume) != inf) { - _callback(predicate, HappyTreeFriends::getValue(_bvh, 0).index); + _callback(predicate, HappyTreeFriends::getValue(_bvh, 0)); } } @@ -487,8 +485,7 @@ struct TreeTraversal - KOKKOS_FUNCTION void operator()(Query const &, int index) const + KOKKOS_FUNCTION void + operator()(Query const &, typename TreeType::value_type const &value) const { - _visitor.visit(_tree, permute(index)); + _visitor.visit(_tree, permute(value.index)); } TreeType _tree;