Correct allgather in place

Signed-off-by: Adrien Taberner <[email protected]>

Author: AdRi1t

src/KokkosComm/mpi/allgather.hpp

@@ -50,8 +50,8 @@ 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 RecvScalar = typename RecvView::value_type;
@@ -61,7 +61,8 @@ void allgather(const RecvView &rv, MPI_Comm comm) {
   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, KokkosComm::data_handle(rv), KokkosComm::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

@@ -95,9 +95,9 @@ void test_allgather_1d_inplace_contig() {
 
   // fill send buffer
   Kokkos::parallel_for(
-      nContrib, KOKKOS_LAMBDA(const int i) { rv(i + rank * nContrib) = rank + i; });
+      nContrib, KOKKOS_LAMBDA(const int i) { rv(rank * nContrib + i) = rank + i; });
 
-  KokkosComm::mpi::allgather(rv, MPI_COMM_WORLD);
+  KokkosComm::mpi::allgather(Kokkos::DefaultExecutionSpace(), rv, nContrib, MPI_COMM_WORLD);
 
   int errs;
   Kokkos::parallel_reduce(