Merge pull request #1038 from aprokop/simplify_predicates_construction

aprokop · web-flow · commit bf6f89b7d3d0 · 2024-04-05T11:57:54.000-04:00
Add helper functions to construct predicates
diff --git a/benchmarks/dbscan/ArborX_DBSCANVerification.hpp b/benchmarks/dbscan/ArborX_DBSCANVerification.hpp
@@ -314,7 +314,8 @@ bool verifyDBSCAN(ExecutionSpace exec_space, Primitives const &primitives,
   ArborX::BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<Point>>
       bvh(exec_space, ArborX::Experimental::attach_indices(points));
 
-  auto const predicates = Details::PrimitivesWithRadius<Points>{points, eps};
+  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);
diff --git a/benchmarks/distributed_tree_driver/distributed_tree_driver.cpp b/benchmarks/distributed_tree_driver/distributed_tree_driver.cpp
@@ -170,52 +170,6 @@ class TimeMonitor
   }
 };
 
-template <typename DeviceType>
-struct NearestNeighborsSearches
-{
-  Kokkos::View<ArborX::Point *, DeviceType> points;
-  int k;
-};
-template <typename DeviceType>
-struct RadiusSearches
-{
-  Kokkos::View<ArborX::Point *, DeviceType> points;
-  double radius;
-};
-
-template <typename DeviceType>
-struct ArborX::AccessTraits<RadiusSearches<DeviceType>, ArborX::PredicatesTag>
-{
-  using memory_space = typename DeviceType::memory_space;
-  static KOKKOS_FUNCTION std::size_t
-  size(RadiusSearches<DeviceType> const &pred)
-  {
-    return pred.points.extent(0);
-  }
-  static KOKKOS_FUNCTION auto get(RadiusSearches<DeviceType> const &pred,
-                                  std::size_t i)
-  {
-    return ArborX::intersects(ArborX::Sphere{pred.points(i), pred.radius});
-  }
-};
-
-template <typename DeviceType>
-struct ArborX::AccessTraits<NearestNeighborsSearches<DeviceType>,
-                            ArborX::PredicatesTag>
-{
-  using memory_space = typename DeviceType::memory_space;
-  static KOKKOS_FUNCTION std::size_t
-  size(NearestNeighborsSearches<DeviceType> const &pred)
-  {
-    return pred.points.extent(0);
-  }
-  static KOKKOS_FUNCTION auto
-  get(NearestNeighborsSearches<DeviceType> const &pred, std::size_t i)
-  {
-    return ArborX::nearest(pred.points(i), pred.k);
-  }
-};
-
 namespace bpo = boost::program_options;
 
 template <class NO>
@@ -421,8 +375,8 @@ int main_(std::vector<std::string> const &args, MPI_Comm const comm)
     knn->start();
     distributed_tree.query(
         ExecutionSpace{},
-        NearestNeighborsSearches<DeviceType>{random_queries, n_neighbors},
-        values, offsets);
+        ArborX::Experimental::make_nearest(random_queries, n_neighbors), values,
+        offsets);
     knn->stop();
 
     if (comm_rank == 0)
@@ -457,9 +411,10 @@ int main_(std::vector<std::string> const &args, MPI_Comm const comm)
     auto radius = time_monitor.getNewTimer("radius");
     MPI_Barrier(comm);
     radius->start();
-    distributed_tree.query(ExecutionSpace{},
-                           RadiusSearches<DeviceType>{random_queries, r},
-                           values, offsets);
+    distributed_tree.query(
+        ExecutionSpace{},
+        ArborX::Experimental::make_intersects(random_queries, r), values,
+        offsets);
     radius->stop();
 
     if (comm_rank == 0)
diff --git a/examples/molecular_dynamics/example_molecular_dynamics.cpp b/examples/molecular_dynamics/example_molecular_dynamics.cpp
@@ -34,7 +34,7 @@ struct ArborX::AccessTraits<Neighbors<MemorySpace>, ArborX::PredicatesTag>
   }
   static KOKKOS_FUNCTION auto get(Neighbors<MemorySpace> const &x, size_type i)
   {
-    return attach(intersects(Sphere{x._particles(i), x._radius}), (int)i);
+    return intersects(Sphere{x._particles(i), x._radius});
   }
 };
 
@@ -119,7 +119,9 @@ int main(int argc, char *argv[])
 
   Kokkos::View<int *, MemorySpace> indices("Example::indices", 0);
   Kokkos::View<int *, MemorySpace> offsets("Example::offsets", 0);
-  index.query(execution_space, Neighbors<MemorySpace>{particles, r},
+  index.query(execution_space,
+              ArborX::Experimental::attach_indices<int>(
+                  Neighbors<MemorySpace>{particles, r}),
               ExcludeSelfCollision{}, indices, offsets);
 
   Kokkos::View<float *[3], MemorySpace> forces(
diff --git a/examples/raytracing/example_raytracing.cpp b/examples/raytracing/example_raytracing.cpp
@@ -107,14 +107,6 @@ struct ArborX::AccessTraits<OrderedIntersectsBased::Rays<MemorySpace>,
 
 namespace IntersectsBased
 {
-/*
- * Storage for the rays and access traits used in the query/traverse.
- */
-template <typename MemorySpace>
-struct Rays
-{
-  Kokkos::View<ArborX::Experimental::Ray *, MemorySpace> _rays;
-};
 
 /*
  * IntersectedCell is a storage container for all intersections between rays and
@@ -177,25 +169,6 @@ struct AccumulateRaySphereIntersections
 };
 } // namespace IntersectsBased
 
-template <typename MemorySpace>
-struct ArborX::AccessTraits<IntersectsBased::Rays<MemorySpace>,
-                            ArborX::PredicatesTag>
-{
-  using memory_space = MemorySpace;
-  using size_type = std::size_t;
-
-  KOKKOS_FUNCTION
-  static size_type size(IntersectsBased::Rays<MemorySpace> const &rays)
-  {
-    return rays._rays.extent(0);
-  }
-  KOKKOS_FUNCTION
-  static auto get(IntersectsBased::Rays<MemorySpace> const &rays, size_type i)
-  {
-    return attach(intersects(rays._rays(i)), (int)i);
-  }
-};
-
 int main(int argc, char *argv[])
 {
   using ExecutionSpace = Kokkos::DefaultExecutionSpace;
@@ -336,7 +309,9 @@ int main(int argc, char *argv[])
                                                             0);
     Kokkos::View<int *> offsets("Example::offsets", 0);
     bvh.query(
-        exec_space, IntersectsBased::Rays<MemorySpace>{rays},
+        exec_space,
+        ArborX::Experimental::attach_indices<int>(
+            ArborX::Experimental::make_intersects(rays)),
         IntersectsBased::AccumulateRaySphereIntersections<MemorySpace>{boxes},
         values, offsets);
 
diff --git a/src/ArborX.hpp b/src/ArborX.hpp
@@ -23,6 +23,7 @@
 #include <ArborX_Exception.hpp>
 #include <ArborX_LinearBVH.hpp>
 #include <ArborX_Point.hpp>
+#include <ArborX_PredicateHelpers.hpp>
 #include <ArborX_Predicates.hpp>
 #include <ArborX_Sphere.hpp>
 // FIXME: we include ArborX_DetailsUtils.hpp only for backward compatibility for
diff --git a/src/ArborX_BruteForce.hpp b/src/ArborX_BruteForce.hpp
@@ -22,6 +22,7 @@
 #include <ArborX_DetailsLegacy.hpp>
 #include <ArborX_IndexableGetter.hpp>
 #include <ArborX_PairValueIndex.hpp>
+#include <ArborX_PredicateHelpers.hpp>
 
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Profiling_ScopedRegion.hpp>
diff --git a/src/ArborX_DBSCAN.hpp b/src/ArborX_DBSCAN.hpp
@@ -24,6 +24,7 @@
 #include <ArborX_HyperBox.hpp>
 #include <ArborX_HyperSphere.hpp>
 #include <ArborX_LinearBVH.hpp>
+#include <ArborX_PredicateHelpers.hpp>
 #include <ArborX_Sphere.hpp>
 
 namespace ArborX
@@ -52,13 +53,6 @@ struct DBSCANCorePoints
   }
 };
 
-template <typename Primitives>
-struct PrimitivesWithRadius
-{
-  Primitives _primitives;
-  float _r;
-};
-
 struct WithinRadiusGetter
 {
   float _r;
@@ -100,31 +94,6 @@ struct MixedBoxPrimitives
 
 } // namespace Details
 
-template <typename Primitives>
-struct AccessTraits<Details::PrimitivesWithRadius<Primitives>, PredicatesTag>
-{
-  using memory_space = typename Primitives::memory_space;
-  using Predicates = Details::PrimitivesWithRadius<Primitives>;
-
-  static KOKKOS_FUNCTION size_t size(Predicates const &w)
-  {
-    return w._primitives.size();
-  }
-  static KOKKOS_FUNCTION auto get(Predicates const &w, size_t i)
-  {
-    auto const &point = w._primitives(i);
-    constexpr int dim =
-        GeometryTraits::dimension_v<std::decay_t<decltype(point)>>;
-    // FIXME reinterpret_cast is dangerous here if access traits return user
-    // point structure (e.g., struct MyPoint { float y; float x; })
-    auto const &hyper_point =
-        reinterpret_cast<ExperimentalHyperGeometry::Point<dim> const &>(point);
-    return attach(
-        intersects(ExperimentalHyperGeometry::Sphere<dim>{hyper_point, w._r}),
-        (int)i);
-  }
-};
-
 template <typename Primitives, typename PermuteFilter>
 struct AccessTraits<Details::PrimitivesWithRadiusReorderedAndFiltered<
                         Primitives, PermuteFilter>,
@@ -315,8 +284,8 @@ dbscan(ExecutionSpace const &exec_space, Primitives const &primitives,
     }
     else
     {
-      auto const predicates =
-          Details::PrimitivesWithRadius<Points>{points, eps};
+      auto const predicates = ArborX::Experimental::attach_indices(
+          ArborX::Experimental::make_intersects(points, eps));
 
       // Determine core points
       Kokkos::Profiling::pushRegion("ArborX::DBSCAN::clusters::num_neigh");
@@ -437,8 +406,8 @@ dbscan(ExecutionSpace const &exec_space, Primitives const &primitives,
       // Perform the queries and build clusters through callback
       using CorePoints = Details::CCSCorePoints;
       Kokkos::Profiling::pushRegion("ArborX::DBSCAN::clusters::query");
-      auto const predicates =
-          Details::PrimitivesWithRadius<Points>{points, eps};
+      auto const predicates = Experimental::attach_indices(
+          Experimental::make_intersects(points, eps));
       bvh.query(exec_space, predicates,
                 Details::FDBSCANDenseBoxCallback<UnionFind, CorePoints, Points,
                                                  decltype(dense_cell_offsets),
@@ -478,8 +447,8 @@ dbscan(ExecutionSpace const &exec_space, Primitives const &primitives,
 
       // Perform the queries and build clusters through callback
       Kokkos::Profiling::pushRegion("ArborX::DBSCAN::clusters::query");
-      auto const predicates =
-          Details::PrimitivesWithRadius<Points>{points, eps};
+      auto const predicates = Experimental::attach_indices(
+          Experimental::make_intersects(points, eps));
       bvh.query(exec_space, predicates,
                 Details::FDBSCANDenseBoxCallback<UnionFind, CorePoints, Points,
                                                  decltype(dense_cell_offsets),
diff --git a/src/ArborX_LinearBVH.hpp b/src/ArborX_LinearBVH.hpp
@@ -29,6 +29,7 @@
 #include <ArborX_HyperBox.hpp>
 #include <ArborX_IndexableGetter.hpp>
 #include <ArborX_PairValueIndex.hpp>
+#include <ArborX_PredicateHelpers.hpp>
 #include <ArborX_SpaceFillingCurves.hpp>
 #include <ArborX_TraversalPolicy.hpp>
 
diff --git a/src/details/ArborX_DetailsFDBSCANDenseBox.hpp b/src/details/ArborX_DetailsFDBSCANDenseBox.hpp
@@ -167,7 +167,7 @@ struct FDBSCANDenseBoxCallback
     }
     else
     {
-      int j = _permute(_num_points_in_dense_cells + (k - _num_dense_cells));
+      auto j = _permute(_num_points_in_dense_cells + (k - _num_dense_cells));
 
       // No need to check the distance here, as the fact that we are inside the
       // callback guarantees that it is <= eps
diff --git a/src/details/ArborX_DetailsMutualReachabilityDistance.hpp b/src/details/ArborX_DetailsMutualReachabilityDistance.hpp
@@ -47,34 +47,6 @@ struct MaxDistance
   }
 };
 
-template <class Primitives>
-struct NearestK
-{
-  Primitives primitives;
-  int k; // including self-collisions
-};
-
-} // namespace Details
-
-template <class Primitives>
-struct AccessTraits<Details::NearestK<Primitives>, PredicatesTag>
-{
-  using memory_space = typename Primitives::memory_space;
-  using size_type = typename memory_space::size_type;
-  static KOKKOS_FUNCTION size_type size(Details::NearestK<Primitives> const &x)
-  {
-    return x.primitives.size();
-  }
-  static KOKKOS_FUNCTION auto get(Details::NearestK<Primitives> const &x,
-                                  size_type i)
-  {
-    return attach(nearest(x.primitives(i), x.k), i);
-  }
-};
-
-namespace Details
-{
-
 template <class CoreDistances>
 struct MutualReachability
 {
diff --git a/src/details/ArborX_MinimumSpanningTree.hpp b/src/details/ArborX_MinimumSpanningTree.hpp
@@ -20,6 +20,7 @@
 #include <ArborX_DetailsTreeNodeLabeling.hpp>
 #include <ArborX_DetailsWeightedEdge.hpp>
 #include <ArborX_LinearBVH.hpp>
+#include <ArborX_PredicateHelpers.hpp>
 
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Profiling_ScopedRegion.hpp>
@@ -66,9 +67,11 @@ struct MinimumSpanningTree
       Kokkos::Profiling::pushRegion("ArborX::MST::compute_core_distances");
       Kokkos::View<float *, MemorySpace> core_distances(
           "ArborX::MST::core_distances", n);
-      bvh.query(space, NearestK<Points>{points, k},
-                MaxDistance<Points, decltype(core_distances)>{points,
-                                                              core_distances});
+      bvh.query(
+          space,
+          Experimental::attach_indices(Experimental::make_nearest(points, k)),
+          MaxDistance<Points, decltype(core_distances)>{points,
+                                                        core_distances});
       Kokkos::Profiling::popRegion();
 
       MutualReachability<decltype(core_distances)> mutual_reachability{
diff --git a/src/details/ArborX_PredicateHelpers.hpp b/src/details/ArborX_PredicateHelpers.hpp
diff --git a/src/interpolation/ArborX_InterpMovingLeastSquares.hpp b/src/interpolation/ArborX_InterpMovingLeastSquares.hpp
diff --git a/test/tstDetailsMutualReachabilityDistance.cpp b/test/tstDetailsMutualReachabilityDistance.cpp
diff --git a/test/tstNeighborList.cpp b/test/tstNeighborList.cpp
diff --git a/test/tstQueryTreeRay.cpp b/test/tstQueryTreeRay.cpp

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ struct ArborX::AccessTraits<Neighbors<MemorySpace>, ArborX::PredicatesTag>`
`34`	`34`	`}`
`35`	`35`	`static KOKKOS_FUNCTION auto get(Neighbors<MemorySpace> const &x, size_type i)`
`36`	`36`	`{`
`37`		`- return attach(intersects(Sphere{x._particles(i), x._radius}), (int)i);`
	`37`	`+ return intersects(Sphere{x._particles(i), x._radius});`
`38`	`38`	`}`
`39`	`39`	`};`
`40`	`40`
`@@ -119,7 +119,9 @@ int main(int argc, char *argv[])`
`119`	`119`
`120`	`120`	`Kokkos::View<int *, MemorySpace> indices("Example::indices", 0);`
`121`	`121`	`Kokkos::View<int *, MemorySpace> offsets("Example::offsets", 0);`
`122`		`- index.query(execution_space, Neighbors<MemorySpace>{particles, r},`
	`122`	`+ index.query(execution_space,`
	`123`	`+ ArborX::Experimental::attach_indices<int>(`
	`124`	`+ Neighbors<MemorySpace>{particles, r}),`
`123`	`125`	`ExcludeSelfCollision{}, indices, offsets);`
`124`	`126`
`125`	`127`	`Kokkos::View<float *[3], MemorySpace> forces(`
Original file line number	Diff line number	Diff line change
`@@ -167,7 +167,7 @@ struct FDBSCANDenseBoxCallback`
`167`	`167`	`}`
`168`	`168`	`else`
`169`	`169`	`{`
`170`		`- int j = _permute(_num_points_in_dense_cells + (k - _num_dense_cells));`
	`170`	`+ auto j = _permute(_num_points_in_dense_cells + (k - _num_dense_cells));`
`171`	`171`
`172`	`172`	`// No need to check the distance here, as the fact that we are inside the`
`173`	`173`	`// callback guarantees that it is <= eps`