Merge pull request #969 from aprokop/pair_index_value

Introduce PairValueIndex and AttachIndices
arborx · Dec 23, 2023 · 8918428 · 8918428
2 parents 5374e84 + 2f317e6
commit 8918428
Show file tree

Hide file tree

Showing 16 changed files with 130 additions and 69 deletions.
diff --git a/benchmarks/brute_force_vs_bvh/brute_force_vs_bvh_timpl.hpp b/benchmarks/brute_force_vs_bvh/brute_force_vs_bvh_timpl.hpp
@@ -90,9 +90,8 @@ static void run_fp(int nprimitives, int nqueries, int nrepeats)
     {
       Kokkos::Timer timer;
       ArborX::BoundingVolumeHierarchy<MemorySpace,
-                                      ArborX::Details::PairIndexVolume<Point>>
-          bvh{space, ArborX::Details::LegacyValues<decltype(primitives), Point>{
-                         primitives}};
+                                      ArborX::PairValueIndex<Point>>
+          bvh{space, ArborX::Experimental::attach_indices(primitives)};
 
       Kokkos::View<int *, ExecutionSpace> indices("Benchmark::indices_ref", 0);
       Kokkos::View<int *, ExecutionSpace> offset("Benchmark::offset_ref", 0);
@@ -110,10 +109,8 @@ static void run_fp(int nprimitives, int nqueries, int nrepeats)
 
     {
       Kokkos::Timer timer;
-      ArborX::BruteForce<MemorySpace, ArborX::Details::PairIndexVolume<Point>>
-          brute{space,
-                ArborX::Details::LegacyValues<decltype(primitives), Point>{
-                    primitives}};
+      ArborX::BruteForce<MemorySpace, ArborX::PairValueIndex<Point>> brute{
+          space, ArborX::Experimental::attach_indices(primitives)};
 
       Kokkos::View<int *, ExecutionSpace> indices("Benchmark::indices", 0);
       Kokkos::View<int *, ExecutionSpace> offset("Benchmark::offset", 0);

diff --git a/benchmarks/dbscan/ArborX_DBSCANVerification.hpp b/benchmarks/dbscan/ArborX_DBSCANVerification.hpp
@@ -13,7 +13,6 @@
 #define ARBORX_DETAILSDBSCANVERIFICATION_HPP
 
 #include <ArborX_DetailsUtils.hpp>
-#include <ArborX_HyperBox.hpp>
 #include <ArborX_LinearBVH.hpp>
 
 #include <Kokkos_Core.hpp>
@@ -312,9 +311,8 @@ bool verifyDBSCAN(ExecutionSpace exec_space, Primitives const &primitives,
   using Point = typename Points::value_type;
   static_assert(GeometryTraits::is_point<Point>{});
 
-  ArborX::BoundingVolumeHierarchy<MemorySpace,
-                                  ArborX::Details::PairIndexVolume<Point>>
-      bvh(exec_space, ArborX::Details::LegacyValues<Points, Point>{points});
+  ArborX::BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<Point>>
+      bvh(exec_space, ArborX::Experimental::attach_indices(points));
 
   auto const predicates = Details::PrimitivesWithRadius<Points>{points, eps};
 

diff --git a/examples/triangle_intersection/triangle_intersection.cpp b/examples/triangle_intersection/triangle_intersection.cpp
@@ -322,9 +322,8 @@ int main()
 
   // Create BVH tree
   ArborX::BoundingVolumeHierarchy<MemorySpace,
-                                  ArborX::Details::PairIndexVolume<Box>> const
-      tree(execution_space,
-           ArborX::Details::LegacyValues<decltype(triangles), Box>{triangles});
+                                  ArborX::PairValueIndex<Box>> const
+      tree(execution_space, ArborX::Experimental::attach_indices(triangles));
 
   // Create the points used for queries
   Points<MemorySpace> points(execution_space);

diff --git a/src/ArborX_BruteForce.hpp b/src/ArborX_BruteForce.hpp
@@ -96,10 +96,10 @@ template <typename MemorySpace>
 class BruteForce<MemorySpace, Details::LegacyDefaultTemplateValue,
                  Details::DefaultIndexableGetter,
                  ExperimentalHyperGeometry::Box<3, float>>
-    : public BruteForce<MemorySpace, Details::PairIndexVolume<Box>,
+    : public BruteForce<MemorySpace, PairValueIndex<Box>,
                         Details::DefaultIndexableGetter, Box>
 {
-  using base_type = BruteForce<MemorySpace, Details::PairIndexVolume<Box>,
+  using base_type = BruteForce<MemorySpace, PairValueIndex<Box>,
                                Details::DefaultIndexableGetter, Box>;
 
 public:

diff --git a/src/ArborX_DBSCAN.hpp b/src/ArborX_DBSCAN.hpp
@@ -278,9 +278,8 @@ dbscan(ExecutionSpace const &exec_space, Primitives const &primitives,
   {
     // Build the tree
     Kokkos::Profiling::pushRegion("ArborX::DBSCAN::tree_construction");
-    ArborX::BoundingVolumeHierarchy<MemorySpace,
-                                    Details::PairIndexVolume<Point>>
-        bvh(exec_space, Details::LegacyValues<Points, Point>{points});
+    BoundingVolumeHierarchy<MemorySpace, PairValueIndex<Point>> bvh(
+        exec_space, Experimental::attach_indices(points));
     Kokkos::Profiling::popRegion();
 
     // Initialize labels after the hierarchy construction to lower memory high
@@ -421,9 +420,8 @@ dbscan(ExecutionSpace const &exec_space, Primitives const &primitives,
                          sorted_cell_indices,
                          permute};
 
-    ArborX::BoundingVolumeHierarchy<MemorySpace, Details::PairIndexVolume<Box>>
-        bvh(exec_space, Details::LegacyValues<decltype(mixed_primitives), Box>{
-                            mixed_primitives});
+    BoundingVolumeHierarchy<MemorySpace, PairValueIndex<Box>> bvh(
+        exec_space, Experimental::attach_indices(mixed_primitives));
 
     Kokkos::Profiling::popRegion();
 

diff --git a/src/ArborX_LinearBVH.hpp b/src/ArborX_LinearBVH.hpp
@@ -141,11 +141,11 @@ template <typename MemorySpace>
 class BoundingVolumeHierarchy<MemorySpace, Details::LegacyDefaultTemplateValue,
                               Details::DefaultIndexableGetter,
                               ExperimentalHyperGeometry::Box<3, float>>
-    : public BoundingVolumeHierarchy<MemorySpace, Details::PairIndexVolume<Box>,
+    : public BoundingVolumeHierarchy<MemorySpace, PairValueIndex<Box>,
                                      Details::DefaultIndexableGetter, Box>
 {
   using base_type =
-      BoundingVolumeHierarchy<MemorySpace, Details::PairIndexVolume<Box>,
+      BoundingVolumeHierarchy<MemorySpace, PairValueIndex<Box>,
                               Details::DefaultIndexableGetter, Box>;
 
 public:

diff --git a/src/details/ArborX_DetailsHappyTreeFriends.hpp b/src/details/ArborX_DetailsHappyTreeFriends.hpp
@@ -12,8 +12,6 @@
 #ifndef ARBORX_DETAILS_HAPPY_TREE_FRIENDS_HPP
 #define ARBORX_DETAILS_HAPPY_TREE_FRIENDS_HPP
 
-#include <ArborX_DetailsNode.hpp>
-
 #include <Kokkos_Macros.hpp>
 #if KOKKOS_VERSION >= 40200
 #include <Kokkos_Assert.hpp> // KOKKOS_ASSERT

diff --git a/src/details/ArborX_DetailsLegacy.hpp b/src/details/ArborX_DetailsLegacy.hpp
@@ -13,7 +13,7 @@
 #define ARBORX_DETAILS_LEGACY_HPP
 
 #include <ArborX_AccessTraits.hpp>
-#include <ArborX_DetailsNode.hpp>
+#include <ArborX_PairValueIndex.hpp>
 
 namespace ArborX::Details
 {
@@ -26,7 +26,8 @@ class LegacyValues
 
 public:
   using memory_space = typename Access::memory_space;
-  using value_type = Details::PairIndexVolume<BoundingVolume>;
+  using index_type = unsigned;
+  using value_type = PairValueIndex<BoundingVolume, index_type>;
   using size_type =
       Kokkos::detected_t<Details::AccessTraitsSizeArchetypeExpression, Access,
                          Primitives>;
@@ -41,13 +42,13 @@ class LegacyValues
     if constexpr (std::is_same_v<BoundingVolume,
                                  typename AccessTraitsHelper<Access>::type>)
     {
-      return value_type{(unsigned)i, Access::get(_primitives, i)};
+      return value_type{Access::get(_primitives, i), (index_type)i};
     }
     else
     {
       BoundingVolume bounding_volume{};
       expand(bounding_volume, Access::get(_primitives, i));
-      return value_type{(unsigned)i, bounding_volume};
+      return value_type{bounding_volume, (index_type)i};
     }
   }
 
@@ -60,16 +61,16 @@ struct LegacyCallbackWrapper
 {
   Callback _callback;
 
-  template <typename Predicate, typename Geometry>
+  template <typename Predicate, typename Value>
   KOKKOS_FUNCTION auto operator()(Predicate const &predicate,
-                                  PairIndexVolume<Geometry> const &value) const
+                                  PairValueIndex<Value> const &value) const
   {
     return _callback(predicate, value.index);
   }
 
-  template <typename Predicate, typename Geometry, typename Output>
+  template <typename Predicate, typename Value, typename Output>
   KOKKOS_FUNCTION auto operator()(Predicate const &predicate,
-                                  PairIndexVolume<Geometry> const &value,
+                                  PairValueIndex<Value> const &value,
                                   Output const &out) const
   {
     return _callback(predicate, value.index, out);
@@ -78,9 +79,9 @@ struct LegacyCallbackWrapper
 
 struct LegacyDefaultCallback
 {
-  template <typename Query, typename Geometry, typename OutputFunctor>
+  template <typename Query, typename Value, typename OutputFunctor>
   KOKKOS_FUNCTION void operator()(Query const &,
-                                  PairIndexVolume<Geometry> const &value,
+                                  PairValueIndex<Value> const &value,
                                   OutputFunctor const &output) const
   {
     output(value.index);

diff --git a/src/details/ArborX_DetailsNode.hpp b/src/details/ArborX_DetailsNode.hpp
@@ -23,13 +23,6 @@ namespace ArborX::Details
 
 constexpr int ROPE_SENTINEL = -1;
 
-template <class BoundingVolume>
-struct PairIndexVolume
-{
-  unsigned index;
-  BoundingVolume bounding_volume;
-};
-
 template <class Value>
 struct LeafNode
 {

diff --git a/src/details/ArborX_IndexableGetter.hpp b/src/details/ArborX_IndexableGetter.hpp
@@ -13,8 +13,8 @@
 #define ARBORX_INDEXABLE_GETTER_HPP
 
 #include <ArborX_AccessTraits.hpp>
-#include <ArborX_DetailsNode.hpp> // PairIndexVolume
 #include <ArborX_GeometryTraits.hpp>
+#include <ArborX_PairValueIndex.hpp>
 
 namespace ArborX::Details
 {
@@ -32,17 +32,17 @@ struct DefaultIndexableGetter
     return geometry;
   }
 
-  template <typename Geometry>
-  KOKKOS_FUNCTION Geometry const &
-  operator()(PairIndexVolume<Geometry> const &pair) const
+  template <typename Value, typename Index>
+  KOKKOS_FUNCTION Value const &
+  operator()(PairValueIndex<Value, Index> const &pair) const
   {
-    return pair.bounding_volume;
+    return pair.value;
   }
 
-  template <typename Geometry>
-  KOKKOS_FUNCTION Geometry operator()(PairIndexVolume<Geometry> &&pair) const
+  template <typename Value, typename Index>
+  KOKKOS_FUNCTION Value operator()(PairValueIndex<Value, Index> &&pair) const
   {
-    return pair.bounding_volume;
+    return pair.value;
   }
 };
 

diff --git a/src/details/ArborX_MinimumSpanningTree.hpp b/src/details/ArborX_MinimumSpanningTree.hpp
@@ -57,8 +57,8 @@ struct MinimumSpanningTree
     auto const n = points.size();
 
     Kokkos::Profiling::pushRegion("ArborX::MST::construction");
-    BoundingVolumeHierarchy<MemorySpace, PairIndexVolume<Point>> bvh(
-        space, Details::LegacyValues<Points, Point>{points});
+    BoundingVolumeHierarchy<MemorySpace, PairValueIndex<Point>> bvh(
+        space, Experimental::attach_indices(points));
     Kokkos::Profiling::popRegion();
 
     if (k > 1)

diff --git a/src/details/ArborX_PairValueIndex.hpp b/src/details/ArborX_PairValueIndex.hpp
@@ -0,0 +1,83 @@
+/****************************************************************************
+ * Copyright (c) 2023 by the ArborX authors                                 *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ArborX library. ArborX is                       *
+ * distributed under a BSD 3-clause license. For the licensing terms see    *
+ * the LICENSE file in the top-level directory.                             *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+
+#ifndef ARBORX_PAIR_VALUE_INDEX_HPP
+#define ARBORX_PAIR_VALUE_INDEX_HPP
+
+#include <ArborX_AccessTraits.hpp>
+
+#include <Kokkos_Macros.hpp>
+
+namespace ArborX
+{
+
+template <typename Value, typename Index = unsigned>
+struct PairValueIndex
+{
+  static_assert(std::is_integral_v<Index>);
+
+  using value_type = Value;
+  using index_type = Index;
+
+  Value value;
+  Index index;
+};
+
+namespace Experimental
+{
+template <typename Values, typename Index>
+class AttachIndices
+{
+private:
+  using Data = Details::AccessValues<Values>;
+
+public:
+  Data _data;
+
+  using memory_space = typename Data::memory_space;
+  using value_type = PairValueIndex<typename Data::value_type, Index>;
+
+  AttachIndices(Values const &values)
+      : _data{values}
+  {}
+
+  KOKKOS_FUNCTION
+  auto operator()(int i) const { return value_type{_data(i), Index(i)}; }
+
+  KOKKOS_FUNCTION
+  auto size() const { return _data.size(); }
+};
+
+template <typename Values, typename Index = unsigned>
+auto attach_indices(Values const &values)
+{
+  return AttachIndices<Values, Index>{values};
+}
+} // namespace Experimental
+
+} // namespace ArborX
+
+template <typename Values, typename Index>
+struct ArborX::AccessTraits<ArborX::Experimental::AttachIndices<Values, Index>,
+                            ArborX::PrimitivesTag>
+{
+  using Self = ArborX::Experimental::AttachIndices<Values, Index>;
+
+  using memory_space = typename Self::memory_space;
+
+  KOKKOS_FUNCTION static auto size(Self const &values) { return values.size(); }
+  KOKKOS_FUNCTION static decltype(auto) get(Self const &values, int i)
+  {
+    return values(i);
+  }
+};
+
+#endif
diff --git a/src/interpolation/ArborX_InterpMovingLeastSquares.hpp b/src/interpolation/ArborX_InterpMovingLeastSquares.hpp
@@ -129,12 +129,8 @@ class MovingLeastSquares
     KOKKOS_ASSERT(0 < num_neighbors_val && num_neighbors_val <= _source_size);
 
     // Organize the source points as a tree
-    using src_coord = typename GeometryTraits::coordinate_type<src_point>::type;
-    using box = ExperimentalHyperGeometry::Box<dimension, src_coord>;
-    using bvh = BoundingVolumeHierarchy<MemorySpace,
-                                        ArborX::Details::PairIndexVolume<box>>;
-    bvh source_tree(
-        space, ArborX::Details::LegacyValues<SourcePoints, box>{source_points});
+    BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<src_point>>
+        source_tree(space, ArborX::Experimental::attach_indices(source_points));
 
     // Create the predicates
     Details::MLSTargetPointsPredicateWrapper<TargetPoints> predicates{
@@ -259,4 +255,4 @@ class MovingLeastSquares
 
 } // namespace ArborX::Interpolation
 
-#endif
+#endif
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -79,7 +79,7 @@ foreach(_test Callbacks Degenerate ManufacturedSolution ComparisonWithBoost)
     "template <class MemorySpace>\n"
     "using ArborX_Legacy_BVH_Box =\n"
     "    LegacyTree<ArborX::BoundingVolumeHierarchy<\n"
-    "        MemorySpace, ArborX::Details::PairIndexVolume<ArborX::Box>,\n"
+    "        MemorySpace, ArborX::PairValueIndex<ArborX::Box>,\n"
     "        ArborX::Details::DefaultIndexableGetter, ArborX::Box>>;\n"
     "#define ARBORX_TEST_TREE_TYPES Tuple<ArborX_BVH_Box, ArborX_Legacy_BVH_Box>\n"
     "#define ARBORX_TEST_DEVICE_TYPES std::tuple<${ARBORX_DEVICE_TYPES}>\n"
@@ -99,7 +99,7 @@ foreach(_test Callbacks Degenerate ManufacturedSolution ComparisonWithBoost)
     "template <class MemorySpace>\n"
     "using ArborX_Legacy_BruteForce_Box =\n"
     "    LegacyTree<ArborX::BruteForce<\n"
-    "        MemorySpace, ArborX::Details::PairIndexVolume<ArborX::Box>,\n"
+    "        MemorySpace, ArborX::PairValueIndex<ArborX::Box>,\n"
     "        ArborX::Details::DefaultIndexableGetter, ArborX::Box>>;\n"
     "#define ARBORX_TEST_TREE_TYPES Tuple<ArborX_BruteForce_Box, ArborX_Legacy_BruteForce_Box>\n"
     "#define ARBORX_TEST_DEVICE_TYPES std::tuple<${ARBORX_DEVICE_TYPES}>\n"
@@ -124,7 +124,7 @@ foreach(_test Callbacks Degenerate ManufacturedSolution ComparisonWithBoost)
       "template <class MemorySpace>\n"
       "using ArborX_Legacy_BVH_${_bounding_volume} =\n"
       "    LegacyTree<ArborX::BoundingVolumeHierarchy<\n"
-      "        MemorySpace, ArborX::Details::PairIndexVolume<${_bounding_volume}>,\n"
+      "        MemorySpace, ArborX::PairValueIndex<${_bounding_volume}>,\n"
       "        ArborX::Details::DefaultIndexableGetter, ${_bounding_volume}>>;\n"
       "#define ARBORX_TEST_TREE_TYPES Tuple<ArborX_Legacy_BVH_${_bounding_volume}>\n"
       "#define ARBORX_TEST_DEVICE_TYPES std::tuple<${ARBORX_DEVICE_TYPES}>\n"

diff --git a/test/tstDetailsMutualReachabilityDistance.cpp b/test/tstDetailsMutualReachabilityDistance.cpp
@@ -32,11 +32,9 @@ auto compute_core_distances(ExecutionSpace exec_space,
 
   ARBORX_ASSERT(points.extent_int(0) >= k);
   using MemorySpace = typename ExecutionSpace::memory_space;
-  ArborX::BoundingVolumeHierarchy<
-      MemorySpace, ArborX::Details::PairIndexVolume<ArborX::Point>>
-      bvh{exec_space,
-          ArborX::Details::LegacyValues<decltype(points), ArborX::Point>{
-              points}};
+  ArborX::BoundingVolumeHierarchy<MemorySpace,
+                                  ArborX::PairValueIndex<ArborX::Point>>
+      bvh{exec_space, ArborX::Experimental::attach_indices(points)};
   Kokkos::View<float *, MemorySpace> distances(
       Kokkos::view_alloc(Kokkos::WithoutInitializing, "Test::core_distances"),
       bvh.size());