Switch to using Iota instead of PairValueIndex in examples and benchm…

…arks
arborx · Oct 14, 2024 · 805036d · 805036d
1 parent 1648842
commit 805036d
Show file tree

Hide file tree

Showing 10 changed files with 50 additions and 40 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
@@ -75,43 +75,43 @@ template <int DIM, typename FloatingPoint>
 static void run_fp(int nprimitives, int nqueries, int nrepeats)
 {
   ExecutionSpace space{};
+
   Placeholder<DIM, FloatingPoint> primitives{nprimitives};
   Placeholder<DIM, FloatingPoint> predicates{nqueries};
+  using Point = ArborX::Point<DIM, FloatingPoint>;
 
   for (int i = 0; i < nrepeats; i++)
   {
     [[maybe_unused]] unsigned int out_count;
     {
       Kokkos::Timer timer;
-      ArborX::BoundingVolumeHierarchy bvh{
-          space, ArborX::Experimental::attach_indices(primitives)};
+      ArborX::BoundingVolumeHierarchy bvh{space, primitives};
 
-      Kokkos::View<int *, ExecutionSpace> indices("Benchmark::indices_ref", 0);
+      Kokkos::View<Point *, ExecutionSpace> values("Benchmark::values_ref", 0);
       Kokkos::View<int *, ExecutionSpace> offset("Benchmark::offset_ref", 0);
-      bvh.query(space, predicates, indices, offset);
+      bvh.query(space, predicates, values, offset);
 
       space.fence();
       double time = timer.seconds();
       if (i == 0)
         printf("Collisions: %.5f\n",
-               (float)(indices.extent(0)) / (nprimitives * nqueries));
+               (float)(values.extent(0)) / (nprimitives * nqueries));
       printf("Time BVH  : %lf\n", time);
-      out_count = indices.extent(0);
+      out_count = values.extent(0);
     }
 
     {
       Kokkos::Timer timer;
-      ArborX::BruteForce brute{
-          space, ArborX::Experimental::attach_indices(primitives)};
+      ArborX::BruteForce brute{space, primitives};
 
-      Kokkos::View<int *, ExecutionSpace> indices("Benchmark::indices", 0);
+      Kokkos::View<Point *, ExecutionSpace> values("Benchmark::values", 0);
       Kokkos::View<int *, ExecutionSpace> offset("Benchmark::offset", 0);
-      brute.query(space, predicates, indices, offset);
+      brute.query(space, predicates, values, offset);
 
       space.fence();
       double time = timer.seconds();
       printf("Time BF   : %lf\n", time);
-      assert(out_count == indices.extent(0));
+      assert(out_count == values.extent(0));
     }
   }
 }

diff --git a/benchmarks/bvh_driver/benchmark_registration.hpp b/benchmarks/bvh_driver/benchmark_registration.hpp
@@ -14,6 +14,7 @@
 
 #include <ArborXBenchmark_PointClouds.hpp>
 #include <ArborX_Point.hpp>
+#include <detail/ArborX_Iota.hpp>
 #include <detail/ArborX_Predicates.hpp>
 
 #include <Kokkos_Core.hpp>
@@ -135,7 +136,10 @@ auto makeTree(ExecutionSpace const &space, Primitives const &primitives)
   if constexpr (is_boost_rtree_v<TreeType>)
     return TreeType(space, primitives);
   else
-    return TreeType(space, ArborX::Experimental::attach_indices(primitives));
+    return TreeType(space,
+                    ArborX::Experimental::Iota<typename TreeType::memory_space>{
+                        primitives.size()},
+                    primitives);
 }
 
 template <typename DeviceType>

diff --git a/benchmarks/bvh_driver/bvh_driver.cpp b/benchmarks/bvh_driver/bvh_driver.cpp
@@ -51,9 +51,10 @@ struct BenchmarkRegistration
 
 template <typename ExecutionSpace>
 using BVHBenchmarkRegistration = BenchmarkRegistration<
-    ExecutionSpace,
-    ArborX::BoundingVolumeHierarchy<typename ExecutionSpace::memory_space,
-                                    ArborX::PairValueIndex<ArborX::Point<3>>>>;
+    ExecutionSpace, ArborX::BoundingVolumeHierarchy<
+                        typename ExecutionSpace::memory_space, int,
+                        Kokkos::View<ArborX::Point<3> *,
+                                     typename ExecutionSpace::memory_space>>>;
 
 void register_bvh_benchmarks(Spec const &spec)
 {

diff --git a/benchmarks/dbscan/ArborX_DBSCANVerification.hpp b/benchmarks/dbscan/ArborX_DBSCANVerification.hpp
@@ -13,6 +13,7 @@
 #define ARBORX_DETAILSDBSCANVERIFICATION_HPP
 
 #include <ArborX_LinearBVH.hpp>
+#include <detail/ArborX_Iota.hpp>
 #include <kokkos_ext/ArborX_KokkosExtViewHelpers.hpp>
 
 #include <Kokkos_Core.hpp>
@@ -299,7 +300,8 @@ bool verifyDBSCAN(ExecutionSpace exec_space, Primitives const &primitives,
   static_assert(GeometryTraits::is_point_v<Point>);
 
   ArborX::BoundingVolumeHierarchy bvh(
-      exec_space, ArborX::Experimental::attach_indices(points));
+      exec_space, ArborX::Experimental::Iota<MemorySpace>{points.size()},
+      points);
 
   auto const predicates = ArborX::Experimental::attach_indices(
       ArborX::Experimental::make_intersects(points, eps));

diff --git a/examples/brute_force/example_brute_force.cpp b/examples/brute_force/example_brute_force.cpp
@@ -23,13 +23,13 @@ struct Dummy
 using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 using MemorySpace = ExecutionSpace::memory_space;
 
-template <>
-struct ArborX::AccessTraits<Dummy, ArborX::PrimitivesTag>
+struct DummyIndexableGetter
 {
+  int count;
+
   using memory_space = MemorySpace;
-  using size_type = typename MemorySpace::size_type;
-  static KOKKOS_FUNCTION size_type size(Dummy const &d) { return d.count; }
-  static KOKKOS_FUNCTION auto get(Dummy const &, size_type i)
+  KOKKOS_FUNCTION auto size() const { return count; }
+  KOKKOS_FUNCTION auto operator()(int i) const
   {
     return ArborX::Point{(float)i, (float)i, (float)i};
   }
@@ -69,13 +69,13 @@ int main(int argc, char *argv[])
   int nprimitives = 5;
   int npredicates = 5;
 
-  Dummy primitives{nprimitives};
+  ArborX::Experimental::Iota<MemorySpace> primitives{nprimitives};
+  DummyIndexableGetter indexable_getter{nprimitives};
   Dummy predicates{npredicates};
 
   unsigned int out_count;
   {
-    ArborX::BoundingVolumeHierarchy bvh{
-        space, ArborX::Experimental::attach_indices(primitives)};
+    ArborX::BoundingVolumeHierarchy bvh{space, primitives, indexable_getter};
 
     Kokkos::View<int *, ExecutionSpace> indices("Example::indices_ref", 0);
     Kokkos::View<int *, ExecutionSpace> offset("Example::offset_ref", 0);
@@ -88,8 +88,7 @@ int main(int argc, char *argv[])
   }
 
   {
-    ArborX::BruteForce brute{space,
-                             ArborX::Experimental::attach_indices(primitives)};
+    ArborX::BruteForce brute{space, primitives, indexable_getter};
 
     Kokkos::View<int *, ExecutionSpace> indices("Example::indices", 0);
     Kokkos::View<int *, ExecutionSpace> offset("Example::offset", 0);

diff --git a/examples/simple_intersection/example_intersection.cpp b/examples/simple_intersection/example_intersection.cpp
@@ -51,7 +51,7 @@ int main(int argc, char *argv[])
   ExecutionSpace space;
 
   ArborX::BoundingVolumeHierarchy const tree(
-      space, ArborX::Experimental::attach_indices(boxes));
+      space, ArborX::Experimental::Iota<MemorySpace>{boxes.size()}, boxes);
 
   // The query will resize indices and offsets accordingly
   Kokkos::View<int *, MemorySpace> indices("Example::indices", 0);

diff --git a/src/ArborX.hpp b/src/ArborX.hpp
@@ -23,6 +23,7 @@
 #include <ArborX_LinearBVH.hpp>
 #include <ArborX_Point.hpp>
 #include <ArborX_Sphere.hpp>
+#include <detail/ArborX_Iota.hpp>
 #include <detail/ArborX_PredicateHelpers.hpp>
 #include <detail/ArborX_Predicates.hpp>
 #include <misc/ArborX_Exception.hpp>

diff --git a/src/distributed/ArborX_DistributedTree.hpp b/src/distributed/ArborX_DistributedTree.hpp
@@ -17,6 +17,7 @@
 #include <detail/ArborX_AccessTraits.hpp>
 #include <detail/ArborX_DistributedTreeNearest.hpp>
 #include <detail/ArborX_DistributedTreeSpatial.hpp>
+#include <detail/ArborX_Iota.hpp>
 #include <detail/ArborX_PairValueIndex.hpp>
 #include <kokkos_ext/ArborX_KokkosExtStdAlgorithms.hpp>
 
@@ -36,8 +37,9 @@ class DistributedTreeBase
   using MemorySpace = typename BottomTree::memory_space;
   using BoundingVolume = typename BottomTree::bounding_volume_type;
   using TopTree =
-      BoundingVolumeHierarchy<MemorySpace, PairValueIndex<BoundingVolume, int>,
-                              Details::DefaultIndexableGetter, BoundingVolume>;
+      BoundingVolumeHierarchy<MemorySpace, int,
+                              Kokkos::View<BoundingVolume *, MemorySpace>,
+                              BoundingVolume>;
 
   using bottom_tree_type = BottomTree;
   using top_tree_type = TopTree;
@@ -228,7 +230,8 @@ DistributedTreeBase<BottomTree>::DistributedTreeBase(
 #endif
 
   // Build top tree with attached ranks
-  _top_tree = TopTree{space, Experimental::attach_indices<int>(volumes)};
+  _top_tree =
+      TopTree{space, Experimental::Iota<MemorySpace>{volumes.size()}, volumes};
 
   Kokkos::Profiling::popRegion();
   Kokkos::Profiling::pushRegion("ArborX::DistributedTree::DistributedTree::"

diff --git a/src/spatial/detail/ArborX_Iota.hpp b/src/spatial/detail/ArborX_Iota.hpp
@@ -18,7 +18,7 @@
 
 #include <type_traits>
 
-namespace ArborX
+namespace ArborX::Experimental
 {
 
 template <typename MemorySpace, typename Index = int>
@@ -35,11 +35,13 @@ struct Iota
       : _n(n)
   {}
 };
+} // namespace ArborX::Experimental
 
 template <typename MemorySpace>
-struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
+struct ArborX::AccessTraits<ArborX::Experimental::Iota<MemorySpace>,
+                            ArborX::PrimitivesTag>
 {
-  using Self = Iota<MemorySpace>;
+  using Self = ArborX::Experimental::Iota<MemorySpace>;
 
   using memory_space = typename Self::memory_space;
   static KOKKOS_FUNCTION size_t size(Self const &self) { return self._n; }
@@ -49,6 +51,4 @@ struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
   }
 };
 
-} // namespace ArborX
-
 #endif
diff --git a/test/tstQueryTreeIntersectsKDOP.cpp b/test/tstQueryTreeIntersectsKDOP.cpp
@@ -13,6 +13,7 @@
 #include <ArborX_KDOP.hpp>
 #include <ArborX_LinearBVH.hpp>
 #include <ArborX_Point.hpp>
+#include <detail/ArborX_Iota.hpp>
 
 #include <Kokkos_Core.hpp>
 
@@ -28,8 +29,6 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(intersects_kdop, DeviceType, ARBORX_DEVICE_TYPES)
   using MemorySpace = typename DeviceType::memory_space;
 
   using Point = ArborX::Point<3>;
-  using Tree = ArborX::BoundingVolumeHierarchy<MemorySpace,
-                                               ArborX::PairValueIndex<Point>>;
 
   std::vector<Point> primitives = {
       {{0, 0, 0}}, // 0
@@ -46,12 +45,13 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(intersects_kdop, DeviceType, ARBORX_DEVICE_TYPES)
       {{0, 0, 2}}, // 11
       {{0, 0, 3}}, // 12
   };
-  Tree const tree(
+  ArborX::BoundingVolumeHierarchy const tree(
       ExecutionSpace{},
-      ArborX::Experimental::attach_indices(Kokkos::create_mirror_view_and_copy(
+      ArborX::Experimental::Iota<MemorySpace>{primitives.size()},
+      Kokkos::create_mirror_view_and_copy(
           MemorySpace{},
           Kokkos::View<Point *, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>(
-              primitives.data(), primitives.size()))));
+              primitives.data(), primitives.size())));
 
   // (0,0,0)->(1,2,3) box with (0,0,0)--(0,0,3) edge cut off
   ArborX::Experimental::KDOP<3, 18> x;