Merge pull request #1182 from aprokop/apiv1_v2

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));
     }
   }
 }

benchmarks/bvh_driver/benchmark_registration.hpp

@@ -32,6 +32,24 @@ struct is_boost_rtree<BoostExt::RTree<Geometry>> : std::true_type
 template <typename Geometry>
 inline constexpr bool is_boost_rtree_v = is_boost_rtree<Geometry>::value;
 
+template <typename MemorySpace>
+struct Iota
+{
+  static_assert(Kokkos::is_memory_space_v<MemorySpace>);
+  using memory_space = MemorySpace;
+  int _n;
+};
+
+template <typename MemorySpace>
+struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
+{
+  using Self = Iota<MemorySpace>;
+
+  using memory_space = typename Self::memory_space;
+  static KOKKOS_FUNCTION size_t size(Self const &self) { return self._n; }
+  static KOKKOS_FUNCTION auto get(Self const &, int i) { return i; }
+};
+
 struct Spec
 {
   using PointCloudType = ArborXBenchmark::PointCloudType;
@@ -135,7 +153,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,
+                    Iota<typename TreeType::memory_space>{
+                        static_cast<int>(primitives.size())},
+                    primitives);
 }
 
 template <typename DeviceType>

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)
 {

benchmarks/dbscan/ArborX_DBSCANVerification.hpp

@@ -276,6 +276,16 @@ bool verifyClusters(ExecutionSpace const &exec_space, IndicesView indices,
   });
 }
 
+struct IndexOnlyCallback
+{
+  template <typename Query, typename Value, typename Output>
+  KOKKOS_FUNCTION auto operator()(Query const &, Value const &value,
+                                  Output const &out) const
+  {
+    out(value.index);
+  }
+};
+
 template <typename ExecutionSpace, typename Primitives, typename LabelsView>
 bool verifyDBSCAN(ExecutionSpace exec_space, Primitives const &primitives,
                   float eps, int core_min_size, LabelsView const &labels)
@@ -301,19 +311,18 @@ bool verifyDBSCAN(ExecutionSpace exec_space, Primitives const &primitives,
   ArborX::BoundingVolumeHierarchy bvh(
       exec_space, ArborX::Experimental::attach_indices(points));
 
-  auto const predicates = ArborX::Experimental::attach_indices(
-      ArborX::Experimental::make_intersects(points, eps));
-
   Kokkos::View<int *, MemorySpace> indices("ArborX::DBSCAN::indices", 0);
   Kokkos::View<int *, MemorySpace> offset("ArborX::DBSCAN::offset", 0);
-  ArborX::query(bvh, exec_space, predicates, indices, offset);
+  bvh.query(exec_space, ArborX::Experimental::make_intersects(points, eps),
+            IndexOnlyCallback{}, indices, offset);
 
   auto passed = Details::verifyClusters(exec_space, indices, offset, labels,
                                         core_min_size);
   Kokkos::Profiling::popRegion();
 
   return passed;
 }
+
 } // namespace Details
 } // namespace ArborX
 

examples/access_traits/example_cuda_access_traits.cpp

@@ -79,19 +79,20 @@ int main(int argc, char *argv[])
   cudaMemcpyAsync(d_a, a.data(), sizeof(a), cudaMemcpyHostToDevice, stream);
 
   Kokkos::Cuda cuda{stream};
-  ArborX::BoundingVolumeHierarchy bvh{
-      cuda, ArborX::Experimental::attach_indices(PointCloud{d_a, d_a, d_a, N})};
+  ArborX::BoundingVolumeHierarchy bvh{cuda, PointCloud{d_a, d_a, d_a, N}};
 
-  Kokkos::View<int *, Kokkos::CudaSpace> indices("Example::indices", 0);
+  Kokkos::View<ArborX::Point<3> *, Kokkos::CudaSpace> points("Example::points",
+                                                             0);
   Kokkos::View<int *, Kokkos::CudaSpace> offset("Example::offset", 0);
-  ArborX::query(bvh, cuda, Spheres{d_a, d_a, d_a, d_a, N}, indices, offset);
+  ArborX::query(bvh, cuda, Spheres{d_a, d_a, d_a, d_a, N}, points, offset);
 
   Kokkos::parallel_for(
-      "Example::print_indices", Kokkos::RangePolicy(cuda, 0, N),
+      "Example::print_points", Kokkos::RangePolicy(cuda, 0, N),
       KOKKOS_LAMBDA(int i) {
         for (int j = offset(i); j < offset(i + 1); ++j)
         {
-          printf("%i %i\n", i, indices(j));
+          printf("%i: (%.1f, %.1f, %.1f)\n", i, points(j)[0], points(j)[1],
+                 points(j)[2]);
         }
       });
 

examples/brute_force/example_brute_force.cpp

@@ -23,13 +23,31 @@ struct Dummy
 using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 using MemorySpace = ExecutionSpace::memory_space;
 
-template <>
-struct ArborX::AccessTraits<Dummy, ArborX::PrimitivesTag>
+template <typename MemorySpace>
+struct Iota
 {
+  static_assert(Kokkos::is_memory_space_v<MemorySpace>);
   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)
+  int _n;
+};
+
+template <typename MemorySpace>
+struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
+{
+  using Self = Iota<MemorySpace>;
+
+  using memory_space = typename Self::memory_space;
+  static KOKKOS_FUNCTION size_t size(Self const &self) { return self._n; }
+  static KOKKOS_FUNCTION auto get(Self const &, int i) { return i; }
+};
+
+struct DummyIndexableGetter
+{
+  int count;
+
+  using memory_space = MemorySpace;
+  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 +87,13 @@ int main(int argc, char *argv[])
   int nprimitives = 5;
   int npredicates = 5;
 
-  Dummy primitives{nprimitives};
+  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 +106,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);

examples/simple_intersection/example_intersection.cpp

@@ -50,18 +50,20 @@ int main(int argc, char *argv[])
 
   ExecutionSpace space;
 
+  using Value = ArborX::PairValueIndex<Box>;
+
   ArborX::BoundingVolumeHierarchy const tree(
       space, ArborX::Experimental::attach_indices(boxes));
 
-  // The query will resize indices and offsets accordingly
-  Kokkos::View<int *, MemorySpace> indices("Example::indices", 0);
+  // The query will resize values and offsets accordingly
+  Kokkos::View<Value *, MemorySpace> values("Example::values", 0);
   Kokkos::View<int *, MemorySpace> offsets("Example::offsets", 0);
-  tree.query(space, queries, indices, offsets);
+  tree.query(space, queries, values, offsets);
 
   auto offsets_host =
       Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, offsets);
-  auto indices_host =
-      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, indices);
+  auto values_host =
+      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, values);
 
   // Expected output:
   //   offsets: 0 1 2 6
@@ -70,10 +72,10 @@ int main(int argc, char *argv[])
   std::cout << "offsets: ";
   std::copy(offsets_host.data(), offsets_host.data() + offsets.size(),
             std::ostream_iterator<int>(std::cout, " "));
-  std::cout << "\nindices: ";
-  std::copy(indices_host.data(), indices_host.data() + indices.size(),
-            std::ostream_iterator<int>(std::cout, " "));
-  std::cout << "\n";
+  std::cout << "\nindices:";
+  for (int i = 0; i < values_host.extent_int(0); ++i)
+    std::cout << " " << values_host(i).index;
+  std::cout << '\n';
 
   return 0;
 }

src/distributed/detail/ArborX_DistributedTreeNearest.hpp

@@ -52,7 +52,8 @@ void DistributedTreeImpl::phaseI(ExecutionSpace const &space, Tree const &tree,
   // Find the k nearest local trees.
   Kokkos::View<int *, MemorySpace> offset(prefix + "::offset", 0);
   Kokkos::View<int *, MemorySpace> nearest_ranks(prefix + "::nearest_ranks", 0);
-  tree._top_tree.query(space, predicates, nearest_ranks, offset);
+  tree._top_tree.query(space, predicates, DistributedTree::IndexOnlyCallback{},
+                       nearest_ranks, offset);
 
   // Accumulate total leave count in the local trees until it reaches k which
   // is the number of neighbors queried for.  Stop if local trees get
@@ -139,7 +140,8 @@ void DistributedTreeImpl::phaseII(ExecutionSpace const &space, Tree const &tree,
   tree._top_tree.query(space,
                        WithinDistanceFromPredicates<Predicates, Distances>{
                            predicates, distances},
-                       nearest_ranks, offset);
+                       DistributedTree::IndexOnlyCallback{}, nearest_ranks,
+                       offset);
 
   auto const &bottom_tree = tree._bottom_tree;
   using BottomTree = std::decay_t<decltype(bottom_tree)>;

src/distributed/detail/ArborX_DistributedTreeSpatial.hpp

@@ -50,7 +50,8 @@ DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
 
   Kokkos::View<int *, MemorySpace> intersected_ranks(
       "ArborX::DistributedTree::query::spatial::intersected_ranks", 0);
-  top_tree.query(space, predicates, intersected_ranks, offset);
+  top_tree.query(space, predicates, DistributedTree::IndexOnlyCallback{},
+                 intersected_ranks, offset);
 
   DistributedTree::forwardQueriesAndCommunicateResults(
       tree.getComm(), space, tree._bottom_tree, predicates, callback,
@@ -81,7 +82,8 @@ void DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
   Kokkos::View<int *, MemorySpace> intersected_ranks(
       prefix + "::intersected_ranks", 0);
   Kokkos::View<int *, MemorySpace> offset(prefix + "::offset", 0);
-  top_tree.query(space, predicates, intersected_ranks, offset);
+  top_tree.query(space, predicates, DistributedTree::IndexOnlyCallback{},
+                 intersected_ranks, offset);
 
   using Query = typename Predicates::value_type;
   Kokkos::View<Query *, MemorySpace> fwd_predicates(prefix + "::fwd_predicates",

src/distributed/detail/ArborX_DistributedTreeUtils.hpp

@@ -14,6 +14,7 @@
 #include <ArborX_Config.hpp>
 
 #include <detail/ArborX_Distributor.hpp>
+#include <detail/ArborX_PairValueIndex.hpp>
 #include <kokkos_ext/ArborX_KokkosExtSort.hpp>
 #include <kokkos_ext/ArborX_KokkosExtStdAlgorithms.hpp>
 #include <kokkos_ext/ArborX_KokkosExtViewHelpers.hpp>
@@ -335,6 +336,17 @@ void filterResults(ExecutionSpace const &space, Predicates const &queries,
   offset = new_offset;
 }
 
+struct IndexOnlyCallback
+{
+  template <typename Query, typename Value, typename Index, typename Output>
+  KOKKOS_FUNCTION auto operator()(Query const &,
+                                  PairValueIndex<Value, Index> const &value,
+                                  Output const &out) const
+  {
+    out(value.index);
+  }
+};
+
 } // namespace ArborX::Details::DistributedTree
 
 #endif