Merge pull request #60 from cwpearson/mpi/reduce-low

mpi: contiguous reduce wrapper
kokkos · May 29, 2024 · a31c457 · a31c457
2 parents 4cbcfcc + 4eb2322
commit a31c457
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/impl/KokkosComm_reduce.hpp b/src/impl/KokkosComm_reduce.hpp
@@ -26,6 +26,22 @@
 #include "KokkosComm_types.hpp"
 
 namespace KokkosComm::Impl {
+
+template <KokkosView SendView, KokkosView RecvView>
+void reduce(const SendView &sv, const RecvView &rv, MPI_Op op, int root, MPI_Comm comm) {
+  Kokkos::Tools::pushRegion("KokkosComm::Impl::reduce");
+  using SPT = KokkosComm::PackTraits<SendView>;
+  using RPT = KokkosComm::PackTraits<RecvView>;
+
+  if (SPT::is_contiguous(sv) && RPT::is_contiguous(rv)) {
+    using SendScalar = typename SendView::non_const_value_type;
+    MPI_Reduce(SPT::data_handle(sv), RPT::data_handle(rv), SPT::span(sv), mpi_type_v<SendScalar>, op, root, comm);
+  } else {
+    throw std::runtime_error("only contiguous views supported for low-level reduce");
+  }
+  Kokkos::Tools::popRegion();
+}
+
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView>
 void reduce(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Op op, int root, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::reduce");