Merge pull request #164 from AdRi1t/inplace_all_gather

Author: cwpearson

src/KokkosComm/mpi/allgather.hpp

@@ -50,19 +50,19 @@ void allgather(const SendView &sv, const RecvView &rv, MPI_Comm comm) {
 }
 
 // in-place allgather
-template <KokkosView RecvView>
-void allgather(const RecvView &rv, MPI_Comm comm) {
+template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
+void allgather(const ExecSpace &space, const RecvView &rv, const size_t recvCount, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Mpi::allgather");
 
-  using RT         = KokkosComm::Traits<RecvView>;
   using RecvScalar = typename RecvView::value_type;
 
-  static_assert(RT::rank() <= 1, "allgather for RecvView::rank > 1 not supported");
+  static_assert(KokkosComm::rank<RecvView>() <= 1, "allgather for RecvView::rank > 1 not supported");
 
-  if (!RT::is_contiguous(rv)) {
+  if (!KokkosComm::is_contiguous(rv)) {
     throw std::runtime_error("low-level allgather requires contiguous recv view");
   }
-  MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, RT::data_handle(rv), RT::span(rv),
+  space.fence("fence before allgather");  // work in space may have been used to produce send view data
+  MPI_Allgather(MPI_IN_PLACE, 0 /*ignored*/, MPI_DATATYPE_NULL /*ignored*/, KokkosComm::data_handle(rv), recvCount,
                 KokkosComm::Impl::mpi_type_v<RecvScalar>, comm);
 
   Kokkos::Tools::popRegion();

unit_tests/mpi/test_allgather.cpp

@@ -83,4 +83,34 @@ void test_allgather_1d_contig() {
 
 TYPED_TEST(Allgather, 1D_contig) { test_allgather_1d_contig<typename TestFixture::Scalar>(); }
 
+template <typename Scalar>
+void test_allgather_1d_inplace_contig() {
+  int rank, size;
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &size);
+
+  const int nContrib = 10;
+
+  Kokkos::View<Scalar *> rv("rv", size * nContrib);
+
+  // fill send buffer
+  Kokkos::parallel_for(
+      nContrib, KOKKOS_LAMBDA(const int i) { rv(rank * nContrib + i) = rank + i; });
+
+  KokkosComm::mpi::allgather(Kokkos::DefaultExecutionSpace(), rv, nContrib, MPI_COMM_WORLD);
+
+  int errs;
+  Kokkos::parallel_reduce(
+      rv.extent(0),
+      KOKKOS_LAMBDA(const int &i, int &lsum) {
+        const int src = i / nContrib;
+        const int j   = i % nContrib;
+        lsum += rv(i) != src + j;
+      },
+      errs);
+  EXPECT_EQ(errs, 0);
+}
+
+TYPED_TEST(Allgather, 1D_inplace_contig) { test_allgather_1d_inplace_contig<typename TestFixture::Scalar>(); }
+
 }  // namespace