Merge pull request #329 from dalg24/new_query_overload

arborx · Aug 7, 2020 · 779f43d · 779f43d
2 parents ffed0fe + 6f3dc47
commit 779f43d
Show file tree

Hide file tree

Showing 7 changed files with 155 additions and 88 deletions.
diff --git a/examples/callback/example_callback.cpp b/examples/callback/example_callback.cpp
@@ -132,14 +132,15 @@ int main(int argc, char *argv[])
 
   {
     // EXPERIMENTAL
-    // TODO replace with BVH::query(ExecutionSpace, Predicates, Callback) when
-    // new overload is added
     Kokkos::View<int, ExecutionSpace, Kokkos::MemoryTraits<Kokkos::Atomic>> c(
         "counter");
 
-    ArborX::Details::traverse(
-        ExecutionSpace{}, bvh, FirstOctant{},
-        KOKKOS_LAMBDA(int i, int j) { printf("%d %d %d\n", ++c(), i, j); });
+#ifndef __NVCC__
+    bvh.query(ExecutionSpace{}, FirstOctant{},
+              KOKKOS_LAMBDA(auto /*predicate*/, int j) {
+                printf("%d %d %d\n", ++c(), -1, j);
+              });
+#endif
   }
 
   return 0;

diff --git a/src/ArborX_LinearBVH.hpp b/src/ArborX_LinearBVH.hpp
@@ -55,10 +55,8 @@ class BoundingVolumeHierarchy
   KOKKOS_FUNCTION
   bounding_volume_type bounds() const noexcept { return _bounds; }
 
-  template <typename ExecutionSpace, typename Predicates,
-            typename CallbackOrView, typename View, typename... Args>
+  template <typename ExecutionSpace, typename Predicates, typename... Args>
   void query(ExecutionSpace const &space, Predicates const &predicates,
-             CallbackOrView &&callback_or_view, View &&view,
              Args &&... args) const
   {
     Details::check_valid_access_traits(PredicatesTag{}, predicates);
@@ -67,15 +65,8 @@ class BoundingVolumeHierarchy
                                                 ExecutionSpace>::value,
                   "Predicates must be accessible from the execution space");
 
-    Details::check_valid_callback_if_first_argument_is_not_a_view(
-        callback_or_view, predicates, view);
-
-    using Tag = typename Details::AccessTraitsHelper<Access>::tag;
-
-    Details::BoundingVolumeHierarchyImpl::queryDispatch(
-        Tag{}, *this, space, predicates,
-        std::forward<CallbackOrView>(callback_or_view),
-        std::forward<View>(view), std::forward<Args>(args)...);
+    Details::BoundingVolumeHierarchyImpl::query(space, *this, predicates,
+                                                std::forward<Args>(args)...);
   }
 
 private:

diff --git a/src/details/ArborX_DetailsBoundingVolumeHierarchyImpl.hpp b/src/details/ArborX_DetailsBoundingVolumeHierarchyImpl.hpp
@@ -322,7 +322,6 @@ queryDispatch(NearestPredicateTag, BVH const &bvh, ExecutionSpace const &space,
                          distances(i) = out(i).second;
                        });
 }
-} // namespace BoundingVolumeHierarchyImpl
 
 template <typename Callback, typename Predicates, typename OutputView>
 std::enable_if_t<!Kokkos::is_view<Callback>{} &&
@@ -353,6 +352,51 @@ check_valid_callback_if_first_argument_is_not_a_view(View const &,
   // do nothing
 }
 
+template <typename ExecutionSpace, typename BVH, typename Predicates,
+          typename CallbackOrView, typename View, typename... Args>
+inline std::enable_if_t<Kokkos::is_view<std::decay_t<View>>{}>
+query(ExecutionSpace const &space, BVH const &bvh, Predicates const &predicates,
+      CallbackOrView &&callback_or_view, View &&view, Args &&... args)
+{
+  check_valid_callback_if_first_argument_is_not_a_view(callback_or_view,
+                                                       predicates, view);
+
+  using Access = AccessTraits<Predicates, Traits::PredicatesTag>;
+  using Tag = typename AccessTraitsHelper<Access>::tag;
+
+  queryDispatch(Tag{}, bvh, space, predicates,
+                std::forward<CallbackOrView>(callback_or_view),
+                std::forward<View>(view), std::forward<Args>(args)...);
+}
+
+template <typename ExecutionSpace, typename BVH, typename Predicates,
+          typename Callback>
+inline void query(ExecutionSpace const &space, BVH const &bvh,
+                  Predicates const &predicates, Callback const &callback,
+                  Experimental::TraversalPolicy const &policy =
+                      Experimental::TraversalPolicy())
+{
+  // TODO check signature of the callback
+  if (policy._sort_predicates)
+  {
+    Kokkos::Profiling::pushRegion("ArborX:BVH:compute_permutation");
+    using MemorySpace = typename BVH::memory_space;
+    using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
+    auto permute =
+        Details::BatchedQueries<DeviceType>::sortQueriesAlongZOrderCurve(
+            space, bvh.bounds(), predicates);
+    Kokkos::Profiling::popRegion();
+
+    using PermutedPredicates = PermutedData<Predicates, decltype(permute)>;
+    traverse(space, bvh, PermutedPredicates{predicates, permute}, callback);
+  }
+  else
+  {
+    traverse(space, bvh, predicates, callback);
+  }
+}
+
+} // namespace BoundingVolumeHierarchyImpl
 } // namespace Details
 } // namespace ArborX