Address review comments

arborx · Jan 17, 2025 · 8a7d92b · 8a7d92b
1 parent c359581
commit 8a7d92b
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 26 deletions.
diff --git a/src/spatial/ArborX_LinearBVH.hpp b/src/spatial/ArborX_LinearBVH.hpp
@@ -233,12 +233,22 @@ BoundingVolumeHierarchy<MemorySpace, Value, IndexableGetter, BoundingVolume>::
   Kokkos::Profiling::popRegion();
   Kokkos::Profiling::pushRegion("ArborX::BVH::BVH::compute_linear_ordering");
 
-  auto linear_ordering_indices = Details::projectOntoSpaceFillingCurve(
-      space, indexables, curve, scene_bounding_box);
+  // Map indexables from multidimensional domain to one-dimensional interval
+  using Details::returnCentroid;
+  using LinearOrderingValueType = std::invoke_result_t<
+      SpaceFillingCurve, decltype(scene_bounding_box),
+      std::decay_t<decltype(returnCentroid(indexables(0)))>>;
+  Kokkos::View<LinearOrderingValueType *, MemorySpace> linear_ordering_indices(
+      Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                         "ArborX::BVH::BVH::linear_ordering"),
+      size());
+  Details::projectOntoSpaceFillingCurve(
+      space, indexables, curve, scene_bounding_box, linear_ordering_indices);
 
   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);
 

diff --git a/src/spatial/detail/ArborX_SpaceFillingCurves.hpp b/src/spatial/detail/ArborX_SpaceFillingCurves.hpp
@@ -30,7 +30,7 @@ namespace Experimental
 
 struct Morton32
 {
-  template <typename Box, class Geometry>
+  template <typename Box, typename Geometry>
   KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box,
                                   Geometry const &geometry) const
   {
@@ -44,7 +44,7 @@ struct Morton32
 
 struct Morton64
 {
-  template <typename Box, class Geometry>
+  template <typename Box, typename Geometry>
   KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box,
                                   Geometry const &geometry) const
   {
@@ -62,35 +62,25 @@ 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>
+          typename Box, typename LinearOrdering>
+inline void
 projectOntoSpaceFillingCurve(ExecutionSpace const &space, Values const &values,
                              SpaceFillingCurve const &curve,
-                             Box const &scene_bounding_box)
+                             Box const &scene_bounding_box,
+                             LinearOrdering &linear_ordering_indices)
 {
   using Point = std::decay_t<decltype(returnCentroid(values(0)))>;
   static_assert(GeometryTraits::is_point_v<Point>);
   static_assert(GeometryTraits::is_box_v<Box>);
+  ARBORX_ASSERT(linear_ordering_indices.size() == values.size());
+  static_assert(std::is_same_v<typename LinearOrdering::value_type,
+                               decltype(curve(scene_bounding_box, values(0)))>);
 
-  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) {
+      Kokkos::RangePolicy(space, 0, values.size()), 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,
@@ -100,8 +90,16 @@ inline auto computeSpaceFillingCurvePermutation(ExecutionSpace const &space,
                                                 SpaceFillingCurve const &curve,
                                                 Box const &scene_bounding_box)
 {
-  auto linear_ordering_indices =
-      projectOntoSpaceFillingCurve(space, values, curve, scene_bounding_box);
+  using Point = std::decay_t<decltype(returnCentroid(values(0)))>;
+  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"),
+          values.size());
+  projectOntoSpaceFillingCurve(space, values, curve, scene_bounding_box,
+                               linear_ordering_indices);
   return sortObjects(space, linear_ordering_indices);
 }
 

diff --git a/test/tstDetailsTreeConstruction.cpp b/test/tstDetailsTreeConstruction.cpp
@@ -92,8 +92,10 @@ 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}}}));
 
-  auto morton_codes = ArborX::Details::projectOntoSpaceFillingCurve(
-      space, boxes, ArborX::Experimental::Morton64(), scene_host);
+  Kokkos::View<unsigned long long *, DeviceType> morton_codes("morton_codes",
+                                                              n);
+  ArborX::Details::projectOntoSpaceFillingCurve(
+      space, boxes, ArborX::Experimental::Morton64(), scene_host, morton_codes);
   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());