Reduce the number of search points in general phase of distributed DBSCAN

src/cluster/ArborX_DistributedDBSCAN.hpp

@@ -76,10 +76,8 @@ void dbscan(MPI_Comm comm, ExecutionSpace const &space,
 
   // Step 2: do local DBSCAN
   auto local_labels =
-      dbscan(space,
-             Details::UnifiedPoints<Points, decltype(ghost_points)>{
-                 points, ghost_points},
-             eps, core_min_size, params);
+      dbscan(space, Details::UnifiedPoints{points, ghost_points}, eps,
+             core_min_size, params);
 
   // Step 3: convert local labels to global
   Kokkos::View<long long *, MemorySpace> rank_offsets(prefix + "rank_offsets",
@@ -147,34 +145,51 @@ void dbscan(MPI_Comm comm, ExecutionSpace const &space,
   Kokkos::resize(space, ghost_labels, num_compressed);
   Details::communicateNeighborDataBack(comm, space, ghost_ranks, ghost_ids,
                                        ghost_labels);
+  Details::KokkosExt::sortByKey(space, ghost_ids, ghost_labels);
+  Kokkos::View<int *, MemorySpace> ghost_offsets(
+      Kokkos::view_alloc(space, prefix + "ghost_offsets"), 0);
+  Details::computeOffsetsInOrderedView(space, ghost_ids, ghost_offsets);
   Kokkos::resize(ghost_ranks, 0); // free space
 
   // Step 5: process multi-labeled indices
   Kokkos::View<Details::MergePair *, MemorySpace> local_merge_pairs(
       prefix + "local_merge_pairs", 0);
-  if (is_special_case)
+  if (ghost_offsets.size() > 1)
   {
-    Details::computeMergePairs(space, Details::CCSCorePoints{}, labels,
-                               ghost_ids, ghost_labels, local_merge_pairs);
+    if (is_special_case)
+    {
+      Details::computeMergePairs(space, Details::CCSCorePoints{}, labels,
+                                 ghost_offsets, ghost_ids, ghost_labels,
+                                 local_merge_pairs);
+    }
+    else
+    {
+      // As we are treating local DBSCAN as a black box, we always do this,
+      // even if it may be unnecessary (e.g., FDBSCAN).
+      BoundingVolumeHierarchy bvh(space,
+                                  Details::UnifiedPoints{points, ghost_points});
+
+      // Find the number of neighbors only for points that appear in the ghost
+      // ids as we only need to resolve the labels of those points. We know
+      // that none of them are noise (the received labels are never -1 because
+      // we filter those before communicating). Some of them may be border
+      // points, but it should be a small fraction, and filtering them out is
+      // not efficient.
+      Kokkos::View<int *, MemorySpace> num_neigh(prefix + "num_neighbors",
+                                                 n_local);
+      bvh.query(space,
+                Details::PointsRequiringResolution{points, ghost_offsets,
+                                                   ghost_ids, eps},
+                Details::CountUpToN<MemorySpace>{num_neigh, core_min_size});
+
+      Details::computeMergePairs(
+          space,
+          Details::DBSCANCorePoints<MemorySpace>{num_neigh, core_min_size},
+          labels, ghost_offsets, ghost_ids, ghost_labels, local_merge_pairs);
+    }
+    sortAndFilterMergePairs(space, local_merge_pairs);
   }
-  else
-  {
-    BoundingVolumeHierarchy bvh(
-        space, Details::UnifiedPoints<Points, decltype(ghost_points)>{
-                   points, ghost_points});
-
-    Kokkos::View<int *, MemorySpace> num_neigh(prefix + "num_neighbors",
-                                               n_local);
-    bvh.query(space,
-              Experimental::attach_indices(
-                  Experimental::make_intersects(points, eps)),
-              Details::CountUpToN<MemorySpace>{num_neigh, core_min_size});
-
-    Details::computeMergePairs(
-        space, Details::DBSCANCorePoints<MemorySpace>{num_neigh, core_min_size},
-        labels, ghost_ids, ghost_labels, local_merge_pairs);
-  }
-  sortAndFilterMergePairs(space, local_merge_pairs);
+  Kokkos::resize(ghost_points, 0); // free space
   Kokkos::resize(ghost_ids, 0);    // free space
   Kokkos::resize(ghost_labels, 0); // free space
 

src/cluster/detail/ArborX_DistributedDBSCANHelpers.hpp

@@ -15,6 +15,7 @@
 #include <ArborX_Box.hpp>
 #include <ArborX_BruteForce.hpp>
 #include <detail/ArborX_Distributor.hpp>
+#include <detail/ArborX_Predicates.hpp>
 #include <detail/ArborX_TreeConstruction.hpp>
 #include <kokkos_ext/ArborX_KokkosExtKernelStdAlgorithms.hpp>
 #include <kokkos_ext/ArborX_KokkosExtViewHelpers.hpp>
@@ -35,6 +36,31 @@ struct UnifiedPoints
   GhostPoints _ghost_points;
 };
 
+template <typename Points, typename GhostOffsets, typename GhostIds,
+          typename Coordinate>
+struct PointsRequiringResolution
+{
+  Points _points;
+  GhostOffsets _ghost_offsets;
+  GhostIds _ghost_ids;
+  Coordinate _eps;
+};
+
+// FIXME_CLANG(Clang<17): Clang 16 or earlier does not support aggregate
+// initialization type deduction
+// https://github.com/llvm/llvm-project/issues/54050
+#if defined(__clang__) && (__clang_major__ < 17)
+template <typename Points, typename GhostPoints>
+KOKKOS_DEDUCTION_GUIDE UnifiedPoints(Points, GhostPoints)
+    -> UnifiedPoints<Points, GhostPoints>;
+
+template <typename Points, typename GhostOffsets, typename GhostIds,
+          typename Coordinate>
+KOKKOS_DEDUCTION_GUIDE PointsRequiringResolution(Points, GhostOffsets, GhostIds,
+                                                 Coordinate)
+    -> PointsRequiringResolution<Points, GhostOffsets, GhostIds, Coordinate>;
+#endif
+
 } // namespace ArborX::Details
 
 template <typename Points, typename GhostPoints>
@@ -55,6 +81,27 @@ struct ArborX::AccessTraits<ArborX::Details::UnifiedPoints<Points, GhostPoints>>
   }
 };
 
+template <typename Points, typename GhostOffsets, typename GhostIds,
+          typename Coordinate>
+struct ArborX::AccessTraits<ArborX::Details::PointsRequiringResolution<
+    Points, GhostOffsets, GhostIds, Coordinate>>
+{
+  using Self = ArborX::Details::PointsRequiringResolution<Points, GhostOffsets,
+                                                          GhostIds, Coordinate>;
+  using memory_space = typename Points::memory_space;
+
+  static KOKKOS_FUNCTION auto size(Self const &self)
+  {
+    return self._ghost_offsets.size() - 1;
+  }
+  static KOKKOS_FUNCTION auto get(Self const &self, size_t i)
+  {
+    auto const id = self._ghost_ids(self._ghost_offsets(i));
+    return ArborX::attach(
+        ArborX::intersects(ArborX::Sphere(self._points(id), self._eps)), id);
+  }
+};
+
 namespace ArborX::Details
 {
 
@@ -312,54 +359,28 @@ struct MergePair
 };
 
 template <typename ExecutionSpace, typename CorePoints, typename Labels,
-          typename ImportedIds, typename ImportedLabels, typename MergePairs>
+          typename GhostOffsets, typename GhostIds, typename GhostLabels,
+          typename MergePairs>
 void computeMergePairs(ExecutionSpace const &space, CorePoints const &is_core,
-                       Labels &local_labels, ImportedIds const &imported_ids,
-                       ImportedLabels const &imported_labels,
-                       MergePairs &merge_pairs)
+                       Labels &local_labels, GhostOffsets const &ghost_offsets,
+                       GhostIds const &ghost_ids,
+                       GhostLabels const &ghost_labels, MergePairs &merge_pairs)
 {
   std::string prefix = "ArborX::DistributedDBSCAN::computeMergePairs";
   Kokkos::Profiling::ScopedRegion guard(prefix);
   prefix += "::";
 
-  using MemorySpace = typename Labels::memory_space;
-
-  int const n_import = imported_labels.size();
-
-  KokkosExt::sortByKey(space, imported_ids, imported_labels);
-
-  Kokkos::View<int *, MemorySpace> offsets(
-      Kokkos::view_alloc(space, prefix + "offsets"), n_import + 1);
-  int num_offsets;
-  Kokkos::parallel_scan(
-      prefix + "compute_offsets", Kokkos::RangePolicy(space, 0, n_import),
-      KOKKOS_LAMBDA(int i, int &update, bool is_final) {
-        if (i == n_import - 1 || imported_ids(i) != imported_ids(i + 1))
-        {
-          if (is_final)
-            offsets(update + 1) = i + 1;
-          ++update;
-        }
-      },
-      num_offsets);
-  ++num_offsets;
-  Kokkos::resize(space, offsets, num_offsets);
-
-  if (num_offsets < 2)
-  {
-    Kokkos::resize(space, merge_pairs, 0);
-    return;
-  }
+  auto const num_offsets = ghost_offsets.size();
 
   Kokkos::resize(Kokkos::view_alloc(space, Kokkos::WithoutInitializing),
-                 merge_pairs, n_import);
+                 merge_pairs, ghost_labels.size());
   int num_merge_pairs;
   Kokkos::parallel_scan(
       prefix + "process_labels", Kokkos::RangePolicy(space, 0, num_offsets - 1),
       KOKKOS_LAMBDA(int const i, int &update, bool is_final) {
-        auto const begin = offsets(i);
-        auto const end = offsets(i + 1);
-        auto const id = imported_ids(begin);
+        auto const begin = ghost_offsets(i);
+        auto const end = ghost_offsets(i + 1);
+        auto const id = ghost_ids(begin);
         auto local_label = local_labels(id);
         bool const is_local_valid = (local_label != -1);
 
@@ -377,7 +398,7 @@ void computeMergePairs(ExecutionSpace const &space, CorePoints const &is_core,
           {
             // Update local label if it is invalid (all imported labels are
             // valid as we filter out noise before communicating)
-            local_labels(id) = imported_labels(begin);
+            local_labels(id) = ghost_labels(begin);
           }
 
           return;
@@ -387,7 +408,7 @@ void computeMergePairs(ExecutionSpace const &space, CorePoints const &is_core,
         auto min_label = (is_local_valid ? local_label : LLONG_MAX);
         for (int j = begin; j < end; ++j)
         {
-          auto const label_j = imported_labels(j);
+          auto const label_j = ghost_labels(j);
           KOKKOS_ASSERT(label_j != -1);
           min_label = Kokkos::min(label_j, min_label);
         }
@@ -404,7 +425,7 @@ void computeMergePairs(ExecutionSpace const &space, CorePoints const &is_core,
 
         for (int j = begin; j < end; ++j)
         {
-          auto label_j = imported_labels(j);
+          auto label_j = ghost_labels(j);
           if (label_j == min_label)
             continue;