Add DistributedTree Nearest query with callback

.jenkins/continuous.groovy

@@ -89,6 +89,8 @@ pipeline {
                                         -D CMAKE_CXX_COMPILER=$KOKKOS_DIR/bin/nvcc_wrapper \
                                         -D CMAKE_CXX_EXTENSIONS=OFF \
                                         -D CMAKE_PREFIX_PATH="$KOKKOS_DIR;$ARBORX_DIR" \
+                                        -D MPIEXEC_PREFLAGS="--allow-run-as-root" \
+                                        -D MPIEXEC_MAX_NUMPROCS=4 \
                                     examples \
                                 '''
                                 sh 'make VERBOSE=1'
@@ -148,6 +150,8 @@ pipeline {
                                         -D CMAKE_CXX_COMPILER=$KOKKOS_DIR/bin/nvcc_wrapper \
                                         -D CMAKE_CXX_EXTENSIONS=OFF \
                                         -D CMAKE_PREFIX_PATH="$KOKKOS_DIR;$ARBORX_DIR" \
+                                        -D MPIEXEC_PREFLAGS="--allow-run-as-root" \
+                                        -D MPIEXEC_MAX_NUMPROCS=4 \
                                     examples \
                                 '''
                                 sh 'make VERBOSE=1'
@@ -205,6 +209,8 @@ pipeline {
                                         -D CMAKE_CXX_COMPILER=clang++ \
                                         -D CMAKE_CXX_EXTENSIONS=OFF \
                                         -D CMAKE_PREFIX_PATH="$KOKKOS_DIR;$ARBORX_DIR" \
+                                        -D MPIEXEC_PREFLAGS="--allow-run-as-root" \
+                                        -D MPIEXEC_MAX_NUMPROCS=4 \
                                     examples \
                                 '''
                                 sh 'make VERBOSE=1'
@@ -266,6 +272,8 @@ pipeline {
                                         -D CMAKE_CXX_COMPILER=clang++ \
                                         -D CMAKE_CXX_EXTENSIONS=OFF \
                                         -D CMAKE_PREFIX_PATH="$KOKKOS_DIR;$ARBORX_DIR" \
+                                        -D MPIEXEC_PREFLAGS="--allow-run-as-root" \
+                                        -D MPIEXEC_MAX_NUMPROCS=4 \
                                     examples \
                                 '''
                                 sh 'make VERBOSE=1'
@@ -329,6 +337,8 @@ pipeline {
                                         -D CMAKE_CXX_EXTENSIONS=OFF \
                                         -D CMAKE_BUILD_TYPE=RelWithDebInfo \
                                         -D CMAKE_PREFIX_PATH="$KOKKOS_DIR;$ARBORX_DIR" \
+                                        -D MPIEXEC_PREFLAGS="--allow-run-as-root" \
+                                        -D MPIEXEC_MAX_NUMPROCS=4 \
                                     examples \
                                 '''
                                 sh 'make VERBOSE=1'
@@ -391,6 +401,8 @@ pipeline {
                                         -D CMAKE_CXX_FLAGS="-Wno-unknown-cuda-version" \
                                         -D CMAKE_PREFIX_PATH="$KOKKOS_DIR;$ARBORX_DIR;$ONE_DPL_DIR" \
                                         -D ONEDPL_PAR_BACKEND=serial \
+                                        -D MPIEXEC_PREFLAGS="--allow-run-as-root" \
+                                        -D MPIEXEC_MAX_NUMPROCS=4 \
                                     examples \
                                 '''
                                 sh 'make VERBOSE=1'

examples/CMakeLists.txt

@@ -13,6 +13,10 @@ add_subdirectory(simple_intersection)
 
 add_subdirectory(molecular_dynamics)
 
+if(ARBORX_ENABLE_MPI)
+  add_subdirectory(distributed_tree)
+endif()
+
 find_package(Boost COMPONENTS program_options)
 if(Boost_FOUND)
   add_subdirectory(viz)

examples/distributed_tree/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_executable(ArborX_DistributedTree_KNNCallback.exe distributed_knn_callback.cpp)
+target_link_libraries(ArborX_DistributedTree_KNNCallback.exe ArborX::ArborX)
+add_test(NAME ArborX_DistributedTree_KNNCallback_Example COMMAND ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} ${MPIEXEC_MAX_NUMPROCS} ${MPIEXEC_PREFLAGS} ./ArborX_DistributedTree_KNNCallback.exe ${MPIEXEC_POSTFLAGS})

examples/distributed_tree/distributed_knn_callback.cpp

@@ -0,0 +1,139 @@
+/****************************************************************************
+ * Copyright (c) 2017-2022 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                                    *
+ ****************************************************************************/
+
+#include <ArborX.hpp>
+
+#include <Kokkos_Core.hpp>
+
+#include <cstdarg>
+#include <cstdio>
+#include <iostream>
+#include <random>
+#include <vector>
+
+#include <mpi.h>
+
+using ExecutionSpace = Kokkos::DefaultExecutionSpace;
+using MemorySpace = ExecutionSpace::memory_space;
+
+namespace Example
+{
+template <class Points>
+struct Nearest
+{
+  Points points;
+  int k;
+  int mpi_rank;
+};
+template <class Points>
+Nearest(Points const &, int, int) -> Nearest<Points>;
+
+struct IndexAndRank
+{
+  int index;
+  int rank;
+};
+
+template <typename DeviceType>
+struct PrintAndInsert
+{
+  Kokkos::View<ArborX::Point *, DeviceType> points;
+  int mpi_rank;
+
+  PrintAndInsert(Kokkos::View<ArborX::Point *, DeviceType> const &points_,
+                 int mpi_rank_)
+      : points(points_)
+      , mpi_rank(mpi_rank_)
+  {}
+
+  template <typename Predicate, typename OutputFunctor>
+  KOKKOS_FUNCTION void operator()(Predicate const &predicate,
+                                  int primitive_index,
+                                  OutputFunctor const &out) const
+  {
+#ifndef KOKKOS_ENABLE_SYCL
+    auto data = ArborX::getData(predicate);
+    auto const &point = points(primitive_index);
+    printf("Intersection for query %d from MPI rank %d on MPI rank %d for "
+           "point %f,%f,%f with index %d\n",
+           data.index, data.rank, mpi_rank, point[0], point[1], point[2],
+           primitive_index);
+#endif
+
+    out({primitive_index, mpi_rank});
+  }
+};
+
+} // namespace Example
+
+template <class Points>
+struct ArborX::AccessTraits<Example::Nearest<Points>, ArborX::PredicatesTag>
+{
+  static KOKKOS_FUNCTION std::size_t size(Example::Nearest<Points> const &x)
+  {
+    return x.points.extent(0);
+  }
+  static KOKKOS_FUNCTION auto get(Example::Nearest<Points> const &x, int i)
+  {
+    return attach(ArborX::nearest(x.points(i), x.k),
+                  Example::IndexAndRank{i, x.mpi_rank});
+  }
+  using memory_space = MemorySpace;
+};
+
+int main(int argc, char *argv[])
+{
+  MPI_Init(&argc, &argv);
+  Kokkos::initialize(argc, argv);
+  {
+    MPI_Comm comm = MPI_COMM_WORLD;
+    int comm_rank;
+    MPI_Comm_rank(comm, &comm_rank);
+    int comm_size;
+    MPI_Comm_size(comm, &comm_size);
+    ArborX::Point lower_left_corner = {static_cast<float>(comm_rank),
+                                       static_cast<float>(comm_rank),
+                                       static_cast<float>(comm_rank)};
+    ArborX::Point center = {static_cast<float>(comm_rank) + .5f,
+                            static_cast<float>(comm_rank) + .5f,
+                            static_cast<float>(comm_rank) + .5f};
+    std::vector points = {lower_left_corner, center};
+    auto points_device = Kokkos::create_mirror_view_and_copy(
+        MemorySpace{},
+        Kokkos::View<ArborX::Point *, Kokkos::HostSpace,
+                     Kokkos::MemoryUnmanaged>(points.data(), points.size()));
+
+    ExecutionSpace exec;
+    ArborX::DistributedTree<MemorySpace> tree(comm, exec, points_device);
+
+    Kokkos::View<Example::IndexAndRank *, MemorySpace> values("values", 0);
+    Kokkos::View<int *, MemorySpace> offsets("offsets", 0);
+    tree.query(exec, Example::Nearest{points_device, 3, comm_rank},
+               Example::PrintAndInsert<MemorySpace>(points_device, comm_rank),
+               values, offsets);
+
+    auto host_values =
+        Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, values);
+    auto host_offsets =
+        Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, offsets);
+    for (unsigned int i = 0; i + 1 < host_offsets.size(); ++i)
+    {
+      std::cout << "Results for query " << i << " on MPI rank " << comm_rank
+                << '\n';
+      for (int j = host_offsets(i); j < host_offsets(i + 1); ++j)
+        std::cout << "point " << host_values(j).index << ", rank "
+                  << host_values(j).rank << std::endl;
+    }
+  }
+  Kokkos::finalize();
+  MPI_Finalize();
+  return 0;
+}

src/details/ArborX_DetailsDistributedTreeImpl.hpp

@@ -209,6 +209,15 @@ struct DistributedTreeImpl
                       &distances);
   }
 
+  template <typename DistributedTree, typename ExecutionSpace,
+            typename Predicates, typename OutputView, typename OffsetView,
+            typename Callback>
+  static std::enable_if_t<Kokkos::is_view<OutputView>{} &&
+                          Kokkos::is_view<OffsetView>{}>
+  queryDispatch(NearestPredicateTag, DistributedTree const &tree,
+                ExecutionSpace const &space, Predicates const &queries,
+                Callback const &callback, OutputView &out, OffsetView &offset);
+
   template <typename DistributedTree, typename ExecutionSpace,
             typename Predicates, typename IndicesAndRanks, typename Offset>
   static std::enable_if_t<Kokkos::is_view<IndicesAndRanks>{} &&
@@ -217,7 +226,6 @@ struct DistributedTreeImpl
                 ExecutionSpace const &space, Predicates const &queries,
                 IndicesAndRanks &values, Offset &offset)
   {
-    // FIXME avoid zipping when distributed nearest callbacks become available
     Kokkos::View<int *, ExecutionSpace> indices(
         "ArborX::DistributedTree::query::nearest::indices", 0);
     Kokkos::View<int *, ExecutionSpace> ranks(
@@ -358,8 +366,8 @@ void DistributedTreeImpl<DeviceType>::deviseStrategy(
 
   // Accumulate total leave count in the local trees until it reaches k which
   // is the number of neighbors queried for.  Stop if local trees get
-  // empty because it means that they are no more leaves and there is no point
-  // on forwarding queries to leafless trees.
+  // empty because that means that there are no more leaves and there is no
+  // point in forwarding queries to leafless trees.
   using Access = AccessTraits<Predicates, PredicatesTag>;
   auto const n_queries = Access::size(queries);
   Kokkos::View<int *, DeviceType> new_offset(
@@ -607,7 +615,6 @@ DistributedTreeImpl<DeviceType>::queryDispatchImpl(
       Kokkos::Profiling::popRegion();
     }
   }
-
   Kokkos::Profiling::popRegion();
 }
 
@@ -671,6 +678,141 @@ DistributedTreeImpl<DeviceType>::queryDispatch(
   Kokkos::Profiling::popRegion();
 }
 
+template <typename Query>
+struct QueriesWithIndices
+{
+  Query query;
+  int query_id;
+  int primitive_index;
+};
+
+template <typename DeviceType>
+template <typename DistributedTree, typename ExecutionSpace,
+          typename Predicates, typename OutputView, typename OffsetView,
+          typename Callback>
+std::enable_if_t<Kokkos::is_view<OutputView>{} && Kokkos::is_view<OffsetView>{}>
+DistributedTreeImpl<DeviceType>::queryDispatch(
+    NearestPredicateTag, DistributedTree const &tree,
+    ExecutionSpace const &space, Predicates const &queries,
+    Callback const &callback, OutputView &out, OffsetView &offset)
+{
+  using Indices = Kokkos::View<int *, ExecutionSpace>;
+  Indices indices("ArborX::DistributedTree::query::nearest::indices", 0);
+  Kokkos::View<int *, ExecutionSpace> ranks(
+      "ArborX::DistributedTree::query::nearest::ranks", 0);
+
+  // The strategy for executing the callback only for the primitives that are
+  // closest for every predicate is as follows:
+  // - Find the ranks and indices for the nearest queries using the overload not
+  //   taking a callback.
+  // - Send the predicate-primitive pairs to the process where the match was
+  //   found.
+  // - Execute the callback on the process owning the primitives.
+  // - Send the result back to the process owning the predicates.
+
+  // Find the ranks and indices for the nearest queries using the overload not
+  // taking a callback.
+  queryDispatchImpl(NearestPredicateTag{}, tree, space, queries, indices,
+                    offset, ranks);
+
+  Kokkos::Profiling::pushRegion(
+      "ArborX::DistributedTree::query::nearest::postprocess_callback");
+
+  // Send the predicate-primitive pairs to the process where the match was
+  // found.
+  auto comm = tree.getComm();
+  int comm_rank;
+  MPI_Comm_rank(comm, &comm_rank);
+
+  using Access = AccessTraits<Predicates, PredicatesTag>;
+  using Query = typename AccessTraitsHelper<Access>::type;
+
+  Kokkos::View<QueriesWithIndices<Query> *, typename DeviceType::memory_space>
+      exported_queries_with_indices(
+          Kokkos::view_alloc(
+              space, Kokkos::WithoutInitializing,
+              "ArborX::DistributedTree::query::exported_queries_with_indices"),
+          ranks.size());
+  Kokkos::parallel_for(
+      "ArborX::DistributedTree::query::zip_queries_and_primitives",
+      Kokkos::RangePolicy<ExecutionSpace>(space, 0, Access::size(queries)),
+      KOKKOS_LAMBDA(int i) {
+        using index_type = typename OffsetView::value_type;
+        for (index_type i = offset(q); i < offset(q + 1); ++i)
+          exported_queries_with_indices(i) = {Access::get(queries, q), q,
+                                              indices(i)};
+      });
+
+  Distributor<DeviceType> distributor(comm);
+  auto const n_imports = distributor.createFromSends(space, ranks);
+
+  Kokkos::View<QueriesWithIndices<Query> *, typename DeviceType::memory_space>
+      imported_queries_with_indices(
+          Kokkos::view_alloc(
+              space, Kokkos::WithoutInitializing,
+              "ArborX::DistributedTree::query::imported_queries_with_indices"),
+          n_imports);
+
+  sendAcrossNetwork(space, distributor, exported_queries_with_indices,
+                    imported_queries_with_indices);
+
+  // Execute the callback on the process owning the primitives.
+  OutputView remote_out(
+      Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                         "ArborX::DistributedTree::query::remote_out"),
+      n_imports);
+  KokkosExt::reallocWithoutInitializing(space, indices, n_imports);
+  Kokkos::parallel_for(
+      "ArborX::DistributedTree::query::execute_callbacks",
+      Kokkos::RangePolicy<ExecutionSpace>(space, 0,
+                                          imported_queries_with_indices.size()),
+      KOKKOS_LAMBDA(int i) {
+        callback(imported_queries_with_indices(i).query,
+                 imported_queries_with_indices(i).primitive_index,
+                 [&](typename OutputView::value_type const &value) {
+                   remote_out(i) = value;
+                   indices(i) = imported_queries_with_indices(i).query_id;
+                 });
+      });
+
+  // Send the result back to the process owning the predicates.
+  Distributor<DeviceType> back_distributor(comm);
+  auto const &dest = distributor.get_sources();
+  auto const &off = distributor.get_source_offsets();
+
+  Kokkos::View<int const *, Kokkos::HostSpace> host_destinations(dest.data(),
+                                                                 dest.size());
+  Kokkos::View<int const *, Kokkos::HostSpace> host_offsets(off.data(),
+                                                            off.size());
+  typename DeviceType::memory_space memory_space;
+  Kokkos::View<int *, DeviceType> destinations(
+      Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                         "ArborX::DistributedTree::query::destinations"),
+      dest.size());
+  Kokkos::deep_copy(space, destinations, host_destinations);
+  Kokkos::View<int *, DeviceType> offsets(
+      Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                         "ArborX::DistributedTree::query::offsets"),
+      off.size());
+  Kokkos::deep_copy(space, offsets, host_offsets);
+  auto const n_imports_back =
+      back_distributor.createFromSends(space, destinations, offsets);
+  KokkosExt::reallocWithoutInitializing(space, out, n_imports_back);
+  Kokkos::View<int *, DeviceType> query_ids(
+      Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                         "ArborX::DistributedTree::query::nearest::query_ids"),
+      n_imports_back);
+
+  // FIXME does combining communication here help?
+  sendAcrossNetwork(space, back_distributor, remote_out, out);
+  sendAcrossNetwork(space, back_distributor, indices, query_ids);
+
+  auto const permutation = ArborX::Details::sortObjects(space, query_ids);
+  ArborX::Details::applyPermutation(space, permutation, out);
+
+  Kokkos::Profiling::popRegion();
+}
+
 template <typename DeviceType>
 template <typename ExecutionSpace, typename View, typename... OtherViews>
 void DistributedTreeImpl<DeviceType>::sortResults(ExecutionSpace const &space,