Add basic distributed nearest callback test

Co-authored-by: Daniel Arndt <arndtd@ornl.com>
aprokop · May 3, 2024 · 0bc2c62 · 0bc2c62
1 parent 85c4378
commit 0bc2c62
Showing 1 changed file with 54 additions and 0 deletions.
diff --git a/test/tstDistributedTree.cpp b/test/tstDistributedTree.cpp
@@ -31,6 +31,39 @@ namespace tt = boost::test_tools;
 
 using ArborX::PairIndexRank;
 
+struct PairRankIndex
+{
+  int rank;
+  int index;
+
+  friend bool operator==(PairRankIndex lhs, PairRankIndex rhs)
+  {
+    return lhs.index == rhs.index && lhs.rank == rhs.rank;
+  }
+  friend bool operator<(PairRankIndex lhs, PairRankIndex rhs)
+  {
+    return lhs.rank < rhs.rank ||
+           (lhs.rank == rhs.rank && lhs.index < rhs.index);
+  }
+  friend std::ostream &operator<<(std::ostream &stream,
+                                  PairRankIndex const &pair)
+  {
+    return stream << '[' << pair.rank << ',' << pair.index << ']';
+  }
+};
+
+struct DistributedNearestCallback
+{
+  int rank;
+
+  template <typename Predicate, typename Value, typename OutputFunctor>
+  KOKKOS_FUNCTION void operator()(Predicate const &, Value const &value,
+                                  OutputFunctor const &out) const
+  {
+    out({rank, value});
+  }
+};
+
 BOOST_AUTO_TEST_CASE_TEMPLATE(hello_world, DeviceType, ARBORX_DEVICE_TYPES)
 {
   using Tree = ArborX::DistributedTree<typename DeviceType::memory_space>;
@@ -128,6 +161,27 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(hello_world, DeviceType, ARBORX_DEVICE_TYPES)
             {{0, comm_size - 1 - comm_rank}, {1, comm_size - 1 - comm_rank}},
             {0, 2}));
   }
+
+  // Now do the same with callbacks
+  if (comm_rank < comm_size - 1)
+  {
+    ARBORX_TEST_QUERY_TREE_CALLBACK(ExecutionSpace{}, tree, nearest_queries,
+                                    DistributedNearestCallback{comm_rank},
+                                    make_reference_solution<PairRankIndex>(
+                                        {{comm_size - 1 - comm_rank, 0},
+                                         {comm_size - 2 - comm_rank, n - 1},
+                                         {comm_size - 1 - comm_rank, 1}},
+                                        {0, 3}));
+  }
+  else
+  {
+    ARBORX_TEST_QUERY_TREE_CALLBACK(
+        ExecutionSpace{}, tree, nearest_queries,
+        DistributedNearestCallback{comm_rank},
+        make_reference_solution<PairRankIndex>(
+            {{comm_size - 1 - comm_rank, 0}, {comm_size - 1 - comm_rank, 1}},
+            {0, 2}));
+  }
 }
 
 BOOST_AUTO_TEST_CASE_TEMPLATE(empty_tree, DeviceType, ARBORX_DEVICE_TYPES)