Merge pull request #69 from cwpearson/mpi/allgather-inplace

MPI: in-place allgather for contiguous views
kokkos · Jun 5, 2024 · 05ea6d7 · 05ea6d7
2 parents 26eba36 + 7b6615f
commit 05ea6d7
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/src/impl/KokkosComm_allgather.hpp b/src/impl/KokkosComm_allgather.hpp
@@ -50,6 +50,24 @@ void allgather(const SendView &sv, const RecvView &rv, MPI_Comm comm) {
   Kokkos::Tools::popRegion();
 }
 
+// in-place allgather
+template <KokkosView RecvView>
+void allgather(const RecvView &rv, MPI_Comm comm) {
+  Kokkos::Tools::pushRegion("KokkosComm::Impl::allgather");
+
+  using RT         = KokkosComm::Traits<RecvView>;
+  using RecvScalar = typename RecvView::value_type;
+
+  static_assert(RT::rank() <= 1, "allgather for RecvView::rank > 1 not supported");
+
+  if (!RT::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), mpi_type_v<RecvScalar>, comm);
+
+  Kokkos::Tools::popRegion();
+}
+
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView>
 void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::allgather");