Switch to using functors for scene reduction

arborx · May 4, 2024 · 63fc08e · 63fc08e
1 parent b05b5f7
commit 63fc08e
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 11 deletions.
diff --git a/src/details/ArborX_DetailsBruteForceImpl.hpp b/src/details/ArborX_DetailsBruteForceImpl.hpp
@@ -28,6 +28,32 @@ namespace ArborX::Details
 {
 struct BruteForceImpl
 {
+
+  template <typename Values, typename IndexableGetter, typename Nodes,
+            typename BoundingVolume>
+  struct SceneReductionFunctor
+  {
+    Values _values;
+    IndexableGetter _indexable_getter;
+    Nodes _nodes;
+
+    KOKKOS_FUNCTION void init(BoundingVolume &volume) const
+    {
+      volume = BoundingVolume{};
+    }
+    KOKKOS_FUNCTION void operator()(int i, BoundingVolume &update) const
+    {
+      using Details::expand;
+      _nodes(i) = _values(i);
+      expand(update, _indexable_getter(_nodes(i)));
+    }
+    KOKKOS_FUNCTION void join(BoundingVolume &result,
+                              BoundingVolume const &update) const
+    {
+      expand(result, update);
+    }
+  };
+
   template <class ExecutionSpace, class Values, class IndexableGetter,
             class Nodes, class Bounds>
   static void initializeBoundingVolumesAndReduceBoundsOfTheScene(
@@ -39,15 +65,9 @@ struct BruteForceImpl
         "ArborX::BruteForce::BruteForce::"
         "initialize_values_and_reduce_bounds",
         Kokkos::RangePolicy<ExecutionSpace>(space, 0, values.size()),
-        KOKKOS_LAMBDA(int i, Bounds &update) {
-          nodes(i) = values(i);
-
-          using Details::expand;
-          Bounds bounding_volume{};
-          expand(bounding_volume, indexable_getter(nodes(i)));
-          update += bounding_volume;
-        },
-        Kokkos::Sum<Bounds>{bounds});
+        SceneReductionFunctor<Values, IndexableGetter, Nodes, Bounds>{
+            values, indexable_getter, nodes},
+        bounds);
   }
 
   template <class ExecutionSpace, class Predicates, class Values,

diff --git a/src/details/ArborX_DetailsTreeConstruction.hpp b/src/details/ArborX_DetailsTreeConstruction.hpp
@@ -22,6 +22,27 @@
 namespace ArborX::Details::TreeConstruction
 {
 
+template <typename Indexables, typename BoundingVolume>
+struct SceneReductionFunctor
+{
+  Indexables _indexables;
+
+  KOKKOS_FUNCTION void init(BoundingVolume &volume) const
+  {
+    volume = BoundingVolume{};
+  }
+  KOKKOS_FUNCTION void operator()(int i, BoundingVolume &update) const
+  {
+    using Details::expand;
+    expand(update, _indexables(i));
+  }
+  KOKKOS_FUNCTION void join(BoundingVolume &result,
+                            BoundingVolume const &update) const
+  {
+    expand(result, update);
+  }
+};
+
 template <typename ExecutionSpace, typename Indexables, typename Box>
 inline void calculateBoundingBoxOfTheScene(ExecutionSpace const &space,
                                            Indexables const &indexables,
@@ -30,8 +51,7 @@ inline void calculateBoundingBoxOfTheScene(ExecutionSpace const &space,
   Kokkos::parallel_reduce(
       "ArborX::TreeConstruction::calculate_bounding_box_of_the_scene",
       Kokkos::RangePolicy<ExecutionSpace>(space, 0, indexables.size()),
-      KOKKOS_LAMBDA(int i, Box &update) { expand(update, indexables(i)); },
-      Kokkos::Sum<Box>{scene_bounding_box});
+      SceneReductionFunctor<Indexables, Box>{indexables}, scene_bounding_box);
 }
 
 template <typename ExecutionSpace, typename Indexables,