Merge pull request #1053 from aprokop/brute_force_nearest

Implement nearest query for BruteForce

Author: aprokop

src/ArborX_BruteForce.hpp

@@ -237,18 +237,12 @@ void BruteForce<MemorySpace, Value, IndexableGetter, BoundingVolume>::query(
   Predicates predicates{user_predicates}; // NOLINT
 
   using Tag = typename Predicates::value_type::Tag;
-  static_assert(std::is_same<Tag, Details::SpatialPredicateTag>{},
-                "nearest query not implemented yet");
-
-  Kokkos::Profiling::pushRegion("ArborX::BruteForce::query::spatial");
 
   Details::BruteForceImpl::query(
-      space, predicates, _values,
+      Tag{}, space, predicates, _values,
       Details::Indexables<decltype(_values), IndexableGetter>{
           _values, _indexable_getter},
       callback);
-
-  Kokkos::Profiling::popRegion();
 }
 
 } // namespace ArborX

src/details/ArborX_DetailsBruteForceImpl.hpp

@@ -14,13 +14,17 @@
 
 #include <ArborX_DetailsAlgorithms.hpp> // expand
 #include <ArborX_DetailsKokkosExtMinMaxOperations.hpp>
+#include <ArborX_DetailsKokkosExtStdAlgorithms.hpp>
+#include <ArborX_DetailsKokkosExtViewHelpers.hpp>
+#include <ArborX_DetailsNearestBufferProvider.hpp>
+#include <ArborX_DetailsPriorityQueue.hpp>
 #include <ArborX_Exception.hpp>
+#include <ArborX_Predicates.hpp>
 
 #include <Kokkos_Core.hpp>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 
-namespace ArborX
-{
-namespace Details
+namespace ArborX::Details
 {
 struct BruteForceImpl
 {
@@ -48,10 +52,12 @@ struct BruteForceImpl
 
   template <class ExecutionSpace, class Predicates, class Values,
             class Indexables, class Callback>
-  static void query(ExecutionSpace const &space, Predicates const &predicates,
-                    Values const &values, Indexables const &indexables,
-                    Callback const &callback)
+  static void query(SpatialPredicateTag, ExecutionSpace const &space,
+                    Predicates const &predicates, Values const &values,
+                    Indexables const &indexables, Callback const &callback)
   {
+    Kokkos::Profiling::ScopedRegion guard("ArborX::BruteForce::query::spatial");
+
     using TeamPolicy = Kokkos::TeamPolicy<ExecutionSpace>;
     using PredicateType = typename Predicates::value_type;
     using IndexableType = std::decay_t<decltype(indexables(0))>;
@@ -136,8 +142,80 @@ struct BruteForceImpl
               });
         });
   }
+
+  template <class ExecutionSpace, class Predicates, class Values,
+            class Indexables, class Callback>
+  static void query(NearestPredicateTag, ExecutionSpace const &space,
+                    Predicates const &predicates, Values const &values,
+                    Indexables const &indexables, Callback const &callback)
+  {
+    Kokkos::Profiling::ScopedRegion guard("ArborX::BruteForce::query::nearest");
+
+    using MemorySpace = typename Values::memory_space;
+
+    int const n_indexables = values.size();
+    int const n_predicates = predicates.size();
+
+    NearestBufferProvider<MemorySpace> buffer_provider(space, predicates);
+
+    Kokkos::parallel_for(
+        "ArborX::BruteForce::query::nearest::"
+        "check_all_predicates_against_all_indexables",
+        Kokkos::RangePolicy<ExecutionSpace>(space, 0, n_predicates),
+        KOKKOS_LAMBDA(int i) {
+          auto const &predicate = predicates(i);
+          auto const k = getK(predicate);
+          auto const buffer = buffer_provider(i);
+
+          if (k < 1)
+            return;
+
+          using PairIndexDistance =
+              typename NearestBufferProvider<MemorySpace>::PairIndexDistance;
+          struct CompareDistance
+          {
+            KOKKOS_INLINE_FUNCTION bool
+            operator()(PairIndexDistance const &lhs,
+                       PairIndexDistance const &rhs) const
+            {
+              return lhs.second < rhs.second;
+            }
+          };
+
+          PriorityQueue<PairIndexDistance, CompareDistance,
+                        UnmanagedStaticVector<PairIndexDistance>>
+              heap(UnmanagedStaticVector<PairIndexDistance>(buffer.data(),
+                                                            buffer.size()));
+
+          // Nodes with a distance that exceed that radius can safely be
+          // discarded. Initialize the radius to infinity and tighten it once k
+          // neighbors have been found.
+          auto radius = KokkosExt::ArithmeticTraits::infinity<float>::value;
+
+          int j = 0;
+          for (; j < n_indexables && j < k; ++j)
+          {
+            auto const distance = predicate.distance(indexables(j));
+            heap.push(Kokkos::make_pair(j, distance));
+          }
+          for (; j < n_indexables; ++j)
+          {
+            auto const distance = predicate.distance(indexables(j));
+            if (distance < radius)
+            {
+              heap.popPush(Kokkos::make_pair(j, distance));
+              radius = heap.top().second;
+            }
+          }
+
+          // Match the logic in TreeTraversal and do the sorting
+          sortHeap(heap.data(), heap.data() + heap.size(), heap.valueComp());
+          for (decltype(heap.size()) i = 0; i < heap.size(); ++i)
+            callback(predicate, values((heap.data() + i)->first));
+        });
+  }
 };
-} // namespace Details
-} // namespace ArborX
+
+} // namespace ArborX::Details
 
 #endif

src/details/ArborX_DetailsNearestBufferProvider.hpp

@@ -0,0 +1,78 @@
+/****************************************************************************
+ * Copyright (c) 2017-2023 by the ArborX authors                            *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ArborX library. ArborX is                       *
+ * distributed under a BSD 3-clause license. For the licensing terms see    *
+ * the LICENSE file in the top-level directory.                             *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+#ifndef ARBORX_DETAILS_NEAREST_BUFFER_PROVIDER_HPP
+#define ARBORX_DETAILS_NEAREST_BUFFER_PROVIDER_HPP
+
+#include <ArborX_DetailsKokkosExtStdAlgorithms.hpp>
+#include <ArborX_DetailsKokkosExtViewHelpers.hpp>
+
+#include <Kokkos_Core.hpp>
+
+namespace ArborX::Details
+{
+
+template <typename MemorySpace>
+struct NearestBufferProvider
+{
+  static_assert(Kokkos::is_memory_space_v<MemorySpace>);
+
+  using PairIndexDistance = Kokkos::pair<int, float>;
+
+  Kokkos::View<PairIndexDistance *, MemorySpace> _buffer;
+  Kokkos::View<int *, MemorySpace> _offset;
+
+  NearestBufferProvider() = default;
+
+  template <typename ExecutionSpace, typename Predicates>
+  NearestBufferProvider(ExecutionSpace const &space,
+                        Predicates const &predicates)
+      : _buffer("ArborX::NearestBufferProvider::buffer", 0)
+      , _offset("ArborX::NearestBufferProvider::offset", 0)
+  {
+    allocateBuffer(space, predicates);
+  }
+
+  KOKKOS_FUNCTION auto operator()(int i) const
+  {
+    return Kokkos::subview(_buffer,
+                           Kokkos::make_pair(_offset(i), _offset(i + 1)));
+  }
+
+  // Enclosing function for an extended __host__ __device__ lambda cannot have
+  // private or protected access within its class
+#ifndef KOKKOS_COMPILER_NVCC
+private:
+#endif
+  template <typename ExecutionSpace, typename Predicates>
+  void allocateBuffer(ExecutionSpace const &space, Predicates const &predicates)
+  {
+    auto const n_queries = predicates.size();
+
+    KokkosExt::reallocWithoutInitializing(space, _offset, n_queries + 1);
+
+    Kokkos::parallel_for(
+        "ArborX::NearestBufferProvider::scan_queries_for_numbers_of_neighbors",
+        Kokkos::RangePolicy<ExecutionSpace>(space, 0, n_queries),
+        KOKKOS_CLASS_LAMBDA(int i) { _offset(i) = getK(predicates(i)); });
+    KokkosExt::exclusive_scan(space, _offset, _offset, 0);
+    int const buffer_size = KokkosExt::lastElement(space, _offset);
+    // Allocate buffer over which to perform heap operations in the nearest
+    // query to store nearest nodes found so far.
+    // It is not possible to anticipate how much memory to allocate since the
+    // number of nearest neighbors k is only known at runtime.
+
+    KokkosExt::reallocWithoutInitializing(space, _buffer, buffer_size);
+  }
+};
+
+} // namespace ArborX::Details
+
+#endif

src/details/ArborX_DetailsTreeTraversal.hpp

@@ -16,6 +16,7 @@
 #include <ArborX_DetailsKokkosExtArithmeticTraits.hpp>
 #include <ArborX_DetailsKokkosExtStdAlgorithms.hpp>
 #include <ArborX_DetailsKokkosExtViewHelpers.hpp>
+#include <ArborX_DetailsNearestBufferProvider.hpp>
 #include <ArborX_DetailsNode.hpp> // ROPE_SENTINEL
 #include <ArborX_DetailsPriorityQueue.hpp>
 #include <ArborX_DetailsStack.hpp>
@@ -128,48 +129,7 @@ struct TreeTraversal<BVH, Predicates, Callback, NearestPredicateTag>
   Predicates _predicates;
   Callback _callback;
 
-  using Buffer = Kokkos::View<Kokkos::pair<int, float> *, MemorySpace>;
-  using Offset = Kokkos::View<int *, MemorySpace>;
-  struct BufferProvider
-  {
-    Buffer _buffer;
-    Offset _offset;
-
-    KOKKOS_FUNCTION auto operator()(int i) const
-    {
-      auto const *offset_ptr = &_offset(i);
-      return Kokkos::subview(_buffer,
-                             Kokkos::make_pair(*offset_ptr, *(offset_ptr + 1)));
-    }
-  };
-
-  BufferProvider _buffer;
-
-  template <typename ExecutionSpace>
-  void allocateBuffer(ExecutionSpace const &space)
-  {
-    auto const n_queries = _predicates.size();
-
-    Offset offset(Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
-                                     "ArborX::TreeTraversal::nearest::offset"),
-                  n_queries + 1);
-    Kokkos::parallel_for(
-        "ArborX::TreeTraversal::nearest::"
-        "scan_queries_for_numbers_of_neighbors",
-        Kokkos::RangePolicy<ExecutionSpace>(space, 0, n_queries),
-        KOKKOS_CLASS_LAMBDA(int i) { offset(i) = getK(_predicates(i)); });
-    KokkosExt::exclusive_scan(space, offset, offset, 0);
-    int const buffer_size = KokkosExt::lastElement(space, offset);
-    // Allocate buffer over which to perform heap operations in
-    // TreeTraversal::nearestQuery() to store nearest leaf nodes found so far.
-    // It is not possible to anticipate how much memory to allocate since the
-    // number of nearest neighbors k is only known at runtime.
-
-    Buffer buffer(Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
-                                     "ArborX::TreeTraversal::nearest::buffer"),
-                  buffer_size);
-    _buffer = BufferProvider{buffer, offset};
-  }
+  NearestBufferProvider<MemorySpace> _buffer;
 
   template <typename ExecutionSpace>
   TreeTraversal(ExecutionSpace const &space, BVH const &bvh,
@@ -192,7 +152,7 @@ struct TreeTraversal<BVH, Predicates, Callback, NearestPredicateTag>
     }
     else
     {
-      allocateBuffer(space);
+      _buffer = NearestBufferProvider<MemorySpace>(space, predicates);
 
       Kokkos::parallel_for(
           "ArborX::TreeTraversal::nearest",
@@ -226,17 +186,8 @@ struct TreeTraversal<BVH, Predicates, Callback, NearestPredicateTag>
     if (k < 1)
       return;
 
-    // Nodes with a distance that exceed that radius can safely be
-    // discarded. Initialize the radius to infinity and tighten it once k
-    // neighbors have been found.
-    auto radius = KokkosExt::ArithmeticTraits::infinity<float>::value;
-
-    using PairIndexDistance = Kokkos::pair<int, float>;
-    static_assert(
-        std::is_same<typename decltype(buffer)::value_type,
-                     PairIndexDistance>::value,
-        "Type of the elements stored in the buffer passed as argument to "
-        "TreeTraversal::nearestQuery is not right");
+    using PairIndexDistance =
+        typename NearestBufferProvider<MemorySpace>::PairIndexDistance;
     struct CompareDistance
     {
       KOKKOS_INLINE_FUNCTION bool operator()(PairIndexDistance const &lhs,
@@ -281,6 +232,11 @@ struct TreeTraversal<BVH, Predicates, Callback, NearestPredicateTag>
     float distance_right = 0.f;
     float distance_node = 0.f;
 
+    // Nodes with a distance that exceed that radius can safely be
+    // discarded. Initialize the radius to infinity and tighten it once k
+    // neighbors have been found.
+    auto radius = KokkosExt::ArithmeticTraits::infinity<float>::value;
+
     do
     {
       bool traverse_left = false;

test/CMakeLists.txt

@@ -106,7 +106,6 @@ foreach(_test Callbacks Degenerate ManufacturedSolution ComparisonWithBoost)
     "        ArborX::Details::DefaultIndexableGetter, ArborX::Box>>;\n"
     "#define ARBORX_TEST_TREE_TYPES Tuple<ArborX_BruteForce_Box, ArborX_Legacy_BruteForce_Box>\n"
     "#define ARBORX_TEST_DEVICE_TYPES std::tuple<${ARBORX_DEVICE_TYPES}>\n"
-    "#define ARBORX_TEST_DISABLE_NEAREST_QUERY\n"
     "#define ARBORX_TEST_DISABLE_CALLBACK_EARLY_EXIT\n"
     "#include <tstQueryTree${_test}.cpp>\n"
   )

test/tstKokkosToolsAnnotations.cpp

@@ -89,6 +89,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::NearestBufferProvider::"
                       "|ArborX::TreeTraversal::spatial::"
                       "|ArborX::TreeTraversal::nearest::"
                       "|ArborX::CrsGraphWrapper::"

test/tstKokkosToolsDistributedAnnotations.cpp

@@ -73,6 +73,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(
                       "|ArborX::DistributedTree::query::"
                       "|ArborX::Distributor::"
                       "|ArborX::BVH::query::"
+                      "|ArborX::NearestBufferProvider::"
                       "|ArborX::TreeTraversal::spatial::"
                       "|ArborX::TreeTraversal::nearest::"
                       "|ArborX::CrsGraphWrapper::"