Get rid of BatchedQueries header and unify space filling curve projec…

…tions
arborx · Oct 11, 2024 · aaff2d0 · aaff2d0
1 parent 43a55fd
commit aaff2d0
Show file tree

Hide file tree

Showing 10 changed files with 94 additions and 153 deletions.
diff --git a/examples/viz/tree_visualization.cpp b/examples/viz/tree_visualization.cpp
@@ -161,10 +161,10 @@ void viz(std::string const &prefix, std::string const &infile, int n_neighbors)
   performQueries(prefix + "shuffled_", suffix);
 
   // Sort them
-  auto permute = ArborX::Details::BatchedQueries<DeviceType>::
-      sortPredicatesAlongSpaceFillingCurve(ExecutionSpace{},
-                                           ArborX::Experimental::Morton32(),
-                                           bvh.bounds(), queries);
+  auto permute = ArborX::Details::computeSpaceFillingCurvePermutation(
+      ExecutionSpace{},
+      ArborX::Details::PredicateIndexables<decltype(queries)>{queries},
+      ArborX::Experimental::Morton32{}, bvh.bounds());
   ArborX::Details::applyPermutation(ExecutionSpace{}, permute, queries);
   performQueries(prefix + "sorted_", suffix);
 }

diff --git a/src/spatial/ArborX_LinearBVH.hpp b/src/spatial/ArborX_LinearBVH.hpp
@@ -16,7 +16,6 @@
 #include <ArborX_CrsGraphWrapper.hpp>
 #include <detail/ArborX_AccessTraits.hpp>
 #include <detail/ArborX_AttachIndices.hpp>
-#include <detail/ArborX_BatchedQueries.hpp>
 #include <detail/ArborX_Callbacks.hpp>
 #include <detail/ArborX_CrsGraphWrapperImpl.hpp>
 #include <detail/ArborX_IndexableGetter.hpp>
@@ -245,26 +244,19 @@ BoundingVolumeHierarchy<MemorySpace, Value, IndexableGetter, BoundingVolume>::
       "ArborX::BVH::BVH::calculate_scene_bounding_box");
 
   // determine the bounding box of the scene
-  Box<DIM, typename GeometryTraits::coordinate_type_t<BoundingVolume>> bbox{};
+  Box<DIM, typename GeometryTraits::coordinate_type_t<BoundingVolume>>
+      scene_bounding_box{};
   Details::TreeConstruction::calculateBoundingBoxOfTheScene(space, indexables,
-                                                            bbox);
+                                                            scene_bounding_box);
   Kokkos::Profiling::popRegion();
   Kokkos::Profiling::pushRegion("ArborX::BVH::BVH::compute_linear_ordering");
 
-  // Map indexables from multidimensional domain to one-dimensional interval
-  using LinearOrderingValueType =
-      std::invoke_result_t<SpaceFillingCurve, decltype(bbox), indexable_type>;
-  Kokkos::View<LinearOrderingValueType *, MemorySpace> linear_ordering_indices(
-      Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
-                         "ArborX::BVH::BVH::linear_ordering"),
-      size());
-  Details::TreeConstruction::projectOntoSpaceFillingCurve(
-      space, indexables, curve, bbox, linear_ordering_indices);
+  auto linear_ordering_indices = Details::projectOntoSpaceFillingCurve(
+      space, indexables, curve, scene_bounding_box);
 
   Kokkos::Profiling::popRegion();
   Kokkos::Profiling::pushRegion("ArborX::BVH::BVH::sort_linearized_order");
 
-  // Compute the ordering of the indexables along the space-filling curve
   auto permutation_indices =
       Details::sortObjects(space, linear_ordering_indices);
 
@@ -325,15 +317,16 @@ void BoundingVolumeHierarchy<
   if (policy._sort_predicates)
   {
     Kokkos::Profiling::pushRegion(profiling_prefix + "::compute_permutation");
-    using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
+    using Details::expand;
     Box<GeometryTraits::dimension_v<bounding_volume_type>,
         typename GeometryTraits::coordinate_type_t<bounding_volume_type>>
         scene_bounding_box{};
-    using namespace Details;
     expand(scene_bounding_box, bounds());
-    auto permute = Details::BatchedQueries<DeviceType>::
-        sortPredicatesAlongSpaceFillingCurve(space, Experimental::Morton32(),
-                                             scene_bounding_box, predicates);
+
+    auto permute = Details::computeSpaceFillingCurvePermutation(
+        space, Details::PredicateIndexables<Predicates>{predicates},
+        Experimental::Morton32{}, scene_bounding_box);
+
     Kokkos::Profiling::popRegion();
 
     using PermutedPredicates =

diff --git a/src/spatial/detail/ArborX_BatchedQueries.hpp b/src/spatial/detail/ArborX_BatchedQueries.hpp
diff --git a/src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp b/src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp
@@ -13,10 +13,10 @@
 #define ARBORX_DETAIL_CRS_GRAPH_WRAPPER_IMPL_HPP
 
 #include <ArborX_Box.hpp>
-#include <detail/ArborX_BatchedQueries.hpp>
 #include <detail/ArborX_Callbacks.hpp>
 #include <detail/ArborX_PermutedData.hpp>
 #include <detail/ArborX_Predicates.hpp>
+#include <detail/ArborX_SpaceFillingCurves.hpp>
 #include <detail/ArborX_TraversalPolicy.hpp>
 #include <kokkos_ext/ArborX_KokkosExtStdAlgorithms.hpp>
 #include <kokkos_ext/ArborX_KokkosExtViewHelpers.hpp>
@@ -343,9 +343,6 @@ queryDispatch(Tag, Tree const &tree, ExecutionSpace const &space,
               Experimental::TraversalPolicy const &policy =
                   Experimental::TraversalPolicy())
 {
-  using MemorySpace = typename Tree::memory_space;
-  using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
-
   check_valid_callback<typename Tree::value_type>(callback, predicates, out);
 
   std::string profiling_prefix = "ArborX::CrsGraphWrapper::query::";
@@ -388,9 +385,9 @@ queryDispatch(Tag, Tree const &tree, ExecutionSpace const &space,
         scene_bounding_box{};
     using namespace Details;
     expand(scene_bounding_box, tree.bounds());
-    auto permute = Details::BatchedQueries<DeviceType>::
-        sortPredicatesAlongSpaceFillingCurve(space, Experimental::Morton32(),
-                                             scene_bounding_box, predicates);
+    auto permute = computeSpaceFillingCurvePermutation(
+        space, PredicateIndexables<Predicates>{predicates},
+        Experimental::Morton32{}, scene_bounding_box);
     Kokkos::Profiling::popRegion();
 
     queryImpl(space, tree, predicates, callback, out, offset, permute,

diff --git a/src/spatial/detail/ArborX_Predicates.hpp b/src/spatial/detail/ArborX_Predicates.hpp
@@ -215,6 +215,24 @@ KOKKOS_INLINE_FUNCTION constexpr auto attach(Predicate &&pred, Data &&data)
       std::forward<Predicate>(pred), std::forward<Data>(data)};
 }
 
+namespace Details
+{
+template <typename Predicates>
+struct PredicateIndexables
+{
+  Predicates _predicates;
+
+  using memory_space = typename Predicates::memory_space;
+
+  KOKKOS_FUNCTION decltype(auto) operator()(int i) const
+  {
+    return getGeometry(_predicates(i));
+  }
+
+  KOKKOS_FUNCTION auto size() const { return _predicates.size(); }
+};
+} // namespace Details
+
 } // namespace ArborX
 
 #endif
diff --git a/src/spatial/detail/ArborX_SpaceFillingCurves.hpp b/src/spatial/detail/ArborX_SpaceFillingCurves.hpp
@@ -12,11 +12,11 @@
 #ifndef ARBORX_SPACE_FILLING_CURVES_HPP
 #define ARBORX_SPACE_FILLING_CURVES_HPP
 
-#include <ArborX_Box.hpp>
-#include <ArborX_Point.hpp>
+#include <ArborX_GeometryTraits.hpp>
 #include <algorithms/ArborX_Centroid.hpp>
 #include <algorithms/ArborX_TranslateAndScale.hpp>
 #include <detail/ArborX_MortonCode.hpp>
+#include <misc/ArborX_SortUtils.hpp>
 
 #include <Kokkos_DetectionIdiom.hpp>
 #include <Kokkos_Macros.hpp>
@@ -30,20 +30,11 @@ namespace Experimental
 
 struct Morton32
 {
-  template <typename Box, typename Point,
-            std::enable_if_t<GeometryTraits::is_box_v<Box> &&
-                             GeometryTraits::is_point_v<Point>> * = nullptr>
-  KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box, Point p) const
-  {
-    Details::translateAndScale(p, p, scene_bounding_box);
-    return Details::morton32(p);
-  }
-  template <typename Box, typename Geometry,
-            std::enable_if_t<GeometryTraits::is_box_v<Box> &&
-                             !GeometryTraits::is_point_v<Geometry>> * = nullptr>
+  template <typename Box, class Geometry>
   KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box,
                                   Geometry const &geometry) const
   {
+    static_assert(GeometryTraits::is_box_v<Box>);
     using Details::returnCentroid;
     auto p = returnCentroid(geometry);
     Details::translateAndScale(p, p, scene_bounding_box);
@@ -53,20 +44,11 @@ struct Morton32
 
 struct Morton64
 {
-  template <typename Box, typename Point,
-            std::enable_if_t<GeometryTraits::is_box_v<Box> &&
-                             GeometryTraits::is_point_v<Point>> * = nullptr>
-  KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box, Point p) const
-  {
-    Details::translateAndScale(p, p, scene_bounding_box);
-    return Details::morton64(p);
-  }
-  template <typename Box, class Geometry,
-            std::enable_if_t<GeometryTraits::is_box_v<Box> &&
-                             !GeometryTraits::is_point_v<Geometry>> * = nullptr>
+  template <typename Box, class Geometry>
   KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box,
                                   Geometry const &geometry) const
   {
+    static_assert(GeometryTraits::is_box_v<Box>);
     using Details::returnCentroid;
     auto p = returnCentroid(geometry);
     Details::translateAndScale(p, p, scene_bounding_box);
@@ -79,6 +61,50 @@ struct Morton64
 namespace Details
 {
 
+template <typename ExecutionSpace, typename Values, typename SpaceFillingCurve,
+          typename Box>
+inline Kokkos::View<std::invoke_result_t<SpaceFillingCurve, Box,
+                                         std::decay_t<decltype(returnCentroid(
+                                             std::declval<Values>()(0)))>> *,
+                    typename Values::memory_space>
+projectOntoSpaceFillingCurve(ExecutionSpace const &space, Values const &values,
+                             SpaceFillingCurve const &curve,
+                             Box const &scene_bounding_box)
+{
+  using Point = std::decay_t<decltype(returnCentroid(values(0)))>;
+  static_assert(GeometryTraits::is_point_v<Point>);
+  static_assert(GeometryTraits::is_box_v<Box>);
+
+  auto const n = values.size();
+
+  using LinearOrderingValueType =
+      std::invoke_result_t<SpaceFillingCurve, Box, Point>;
+  Kokkos::View<LinearOrderingValueType *, typename Values::memory_space>
+      linear_ordering_indices(
+          Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                             "ArborX::SpaceFillingCurve::linear_ordering"),
+          n);
+  Kokkos::parallel_for(
+      "ArborX::SpaceFillingCurve::project_onto_space_filling_curve",
+      Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) {
+        linear_ordering_indices(i) = curve(scene_bounding_box, values(i));
+      });
+
+  return linear_ordering_indices;
+}
+
+template <typename ExecutionSpace, typename Values, typename SpaceFillingCurve,
+          typename Box>
+inline auto computeSpaceFillingCurvePermutation(ExecutionSpace const &space,
+                                                Values const &values,
+                                                SpaceFillingCurve const &curve,
+                                                Box const &scene_bounding_box)
+{
+  auto linear_ordering_indices =
+      projectOntoSpaceFillingCurve(space, values, curve, scene_bounding_box);
+  return sortObjects(space, linear_ordering_indices);
+}
+
 template <int DIM, class SpaceFillingCurve>
 void check_valid_space_filling_curve(SpaceFillingCurve const &)
 {

diff --git a/src/spatial/detail/ArborX_TreeConstruction.hpp b/src/spatial/detail/ArborX_TreeConstruction.hpp
@@ -15,7 +15,6 @@
 #include <algorithms/ArborX_Expand.hpp>
 #include <algorithms/ArborX_Reducer.hpp>
 #include <detail/ArborX_Node.hpp> // makeLeafNode
-#include <detail/ArborX_SpaceFillingCurves.hpp>
 #include <kokkos_ext/ArborX_KokkosExtArithmeticTraits.hpp>
 #include <misc/ArborX_Exception.hpp>
 
@@ -39,27 +38,6 @@ inline void calculateBoundingBoxOfTheScene(ExecutionSpace const &space,
       GeometryReducer<Box>(scene_bounding_box));
 }
 
-template <typename ExecutionSpace, typename Indexables,
-          typename SpaceFillingCurve, typename Box, typename LinearOrdering>
-inline void projectOntoSpaceFillingCurve(ExecutionSpace const &space,
-                                         Indexables const &indexables,
-                                         SpaceFillingCurve const &curve,
-                                         Box const &scene_bounding_box,
-                                         LinearOrdering linear_ordering_indices)
-{
-  size_t const n = indexables.size();
-  ARBORX_ASSERT(linear_ordering_indices.extent(0) == n);
-  static_assert(
-      std::is_same_v<typename LinearOrdering::value_type,
-                     decltype(curve(scene_bounding_box, indexables(0)))>);
-
-  Kokkos::parallel_for(
-      "ArborX::TreeConstruction::project_primitives_onto_space_filling_curve",
-      Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) {
-        linear_ordering_indices(i) = curve(scene_bounding_box, indexables(i));
-      });
-}
-
 template <typename ExecutionSpace, typename Values, typename IndexableGetter,
           typename Nodes, typename BoundingVolume>
 inline void

diff --git a/test/tstDetailsTreeConstruction.cpp b/test/tstDetailsTreeConstruction.cpp
@@ -14,6 +14,7 @@
 #include <detail/ArborX_IndexableGetter.hpp>
 #include <detail/ArborX_MortonCode.hpp> // expandBits, morton32
 #include <detail/ArborX_Node.hpp>       // ROPE SENTINEL
+#include <detail/ArborX_SpaceFillingCurves.hpp>
 #include <detail/ArborX_TreeConstruction.hpp>
 #include <kokkos_ext/ArborX_KokkosExtStdAlgorithms.hpp>
 #include <misc/ArborX_SortUtils.hpp> // sortObjects
@@ -92,10 +93,8 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(assign_morton_codes, DeviceType,
   BOOST_TEST(ArborX::Details::equals(
       scene_host, {{{0.0, 0.0, 0.0}}, {{(float)N, (float)N, (float)N}}}));
 
-  Kokkos::View<unsigned long long *, DeviceType> morton_codes("morton_codes",
-                                                              n);
-  ArborX::Details::TreeConstruction::projectOntoSpaceFillingCurve(
-      space, boxes, ArborX::Experimental::Morton64(), scene_host, morton_codes);
+  auto morton_codes = ArborX::Details::projectOntoSpaceFillingCurve(
+      space, boxes, ArborX::Experimental::Morton64(), scene_host);
   auto morton_codes_host = Kokkos::create_mirror_view(morton_codes);
   Kokkos::deep_copy(morton_codes_host, morton_codes);
   BOOST_TEST(morton_codes_host == ref, tt::per_element());

diff --git a/test/tstKokkosToolsAnnotations.cpp b/test/tstKokkosToolsAnnotations.cpp
@@ -45,6 +45,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_bvh_allocations_prefixed, DeviceType,
         std::regex re("^(Testing::"
                       "|ArborX::BVH::"
                       "|ArborX::Sorting::"
+                      "|ArborX::SpaceFillingCurve::"
                       "|Kokkos::SortImpl::BinSortFunctor::"
                       "|Kokkos::Serial::" // unsure what's going on
                       ").*");
@@ -94,6 +95,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_query_allocations_prefixed, DeviceType,
          void const * /*ptr*/, uint64_t /*size*/) {
         std::regex re("^(Testing::"
                       "|ArborX::BVH::query::"
+                      "|ArborX::SpaceFillingCurve::"
                       "|ArborX::NearestBufferProvider::"
                       "|ArborX::TreeTraversal::spatial::"
                       "|ArborX::TreeTraversal::nearest::"