feat(mpi): add non-blocking reduce (mpi::ireduce)

Prepare for exposure in core API. Also add unit test.

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Author: dssgabriel

src/KokkosComm/mpi/reduce.hpp

@@ -9,12 +9,68 @@
 #include <KokkosComm/concepts.hpp>
 #include <KokkosComm/traits.hpp>
 #include <KokkosComm/datatype.hpp>
+#include "mpi_space.hpp"
+#include "req.hpp"
 
-#include "impl/pack_traits.hpp"
 #include "impl/error_handling.hpp"
+#include "impl/pack_traits.hpp"
 
 namespace KokkosComm::mpi {
 
+template <KokkosExecutionSpace ExecSpace, KokkosView SView, KokkosView RView>
+auto ireduce(const ExecSpace &space, const SView sv, RView rv, MPI_Op op, int root, MPI_Comm comm) -> Req<MpiSpace> {
+  using ST   = typename SView::non_const_value_type;
+  using RT   = typename RView::non_const_value_type;
+  using SPkr = typename PackTraits<SView>::packer_type;
+  using RPkr = typename PackTraits<RView>::packer_type;
+  static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::ireduce: View value types must be identical");
+  static_assert(rank<SView>() <= 1 and rank<RView>() <= 1,
+                "KokkosComm::mpi::ireduce: Views with rank higher than 1 are not supported");
+  Kokkos::Tools::pushRegion("KokkosComm::mpi::ireduce");
+
+  const int rank = [=]() {
+    int _r;
+    MPI_Comm_rank(comm, &_r);
+    return _r;
+  }();
+
+  Req<MpiSpace> req;
+  if (is_contiguous(sv)) {
+    if (rank == root and not is_contiguous(rv)) {
+      auto pkd_rv = RPkr::allocate_packed_for(space, "KC::mpi::ireduce c_sv pkd_rv", rv);
+      space.fence("fence allocation before MPI call");
+      MPI_Ireduce(data_handle(sv), data_handle(pkd_rv.view), span(sv), datatype<MpiSpace, ST>(), op, root, comm,
+                  &req.mpi_request());
+      RPkr::unpack_into(space, rv, pkd_rv.view);
+      req.extend_view_lifetime(pkd_rv.view);
+    } else {
+      space.fence("fence before MPI call");
+      MPI_Ireduce(data_handle(sv), data_handle(rv), span(sv), datatype<MpiSpace, ST>(), op, root, comm,
+                  &req.mpi_request());
+    }
+  } else {
+    auto send_args = SPkr::pack(space, sv);
+    if (rank == root and not is_contiguous(rv)) {
+      auto pkd_rv = RPkr::allocate_packed_for(space, "KC::mpi::ireduce nc_sv pkd_rv", rv);
+      space.fence("fence allocation before MPI call");
+      MPI_Ireduce(data_handle(send_args.view), data_handle(pkd_rv.view), send_args.count, send_args.datatype, op, root,
+                  comm, &req.mpi_request());
+      RPkr::unpack_into(space, rv, pkd_rv.view);  // no-op
+      req.extend_view_lifetime(pkd_rv.view);
+    } else {
+      space.fence("fence before MPI call");
+      MPI_Ireduce(data_handle(send_args.view), data_handle(rv), send_args.count, datatype<MpiSpace, ST>(), op, root,
+                  comm, &req.mpi_request());
+    }
+    req.extend_view_lifetime(send_args.view);
+  }
+  req.extend_view_lifetime(sv);
+  req.extend_view_lifetime(rv);
+
+  Kokkos::Tools::popRegion();
+  return req;
+}
+
 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::mpi::reduce");

unit_tests/CMakeLists.txt

@@ -125,6 +125,13 @@ if(KOKKOSCOMM_ENABLE_MPI)
     FILES test_main.cpp mpi/test_reduce.cpp
     LIBRARIES KokkosComm::KokkosComm
   )
+  kc_add_unit_test(
+    test.mpi.ireduce
+    MPI
+    NUM_PES 2
+    FILES test_main.cpp mpi/test_ireduce.cpp
+    LIBRARIES KokkosComm::KokkosComm
+  )
   kc_add_unit_test(
     test.mpi.allreduce
     MPI

unit_tests/mpi/test_ireduce.cpp

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#include <gtest/gtest.h>
+
+#include <KokkosComm/KokkosComm.hpp>
+
+namespace {
+
+template <typename T>
+class NonBlockingReduce : public testing::Test {
+ public:
+  using Scalar = T;
+};
+
+using ScalarTypes = ::testing::Types<int, int64_t, float, double, Kokkos::complex<float>, Kokkos::complex<double>>;
+TYPED_TEST_SUITE(NonBlockingReduce, ScalarTypes);
+
+template <typename Scalar>
+void test_ireduce_1d_contig() {
+  int rank, size;
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &size);
+  int root = 0;
+
+  Kokkos::View<Scalar *> sv("sv", 10);
+  Kokkos::View<Scalar *> rv("rv", 10);
+  // if (rank == root) {
+  //   Kokkos::resize(rv, sv.extent(0));
+  // }
+
+  // Prepare send buffer
+  Kokkos::parallel_for(
+      sv.extent(0), KOKKOS_LAMBDA(int i) { sv(i) = rank + i; });
+
+  KokkosComm::mpi::ireduce(Kokkos::DefaultExecutionSpace(), sv, rv, MPI_SUM, root, MPI_COMM_WORLD);
+
+  if (root == rank) {
+    int errs;
+    Kokkos::parallel_reduce(
+        rv.extent(0),
+        KOKKOS_LAMBDA(int i, int &lsum) {
+          Scalar acc = 0;
+          for (int r = 0; r < size; ++r) {
+            acc += r + i;
+          }
+          lsum += rv(i) != acc;
+        },
+        errs);
+    ASSERT_EQ(errs, 0);
+  }
+}
+
+TYPED_TEST(NonBlockingReduce, 1D_contig) { test_ireduce_1d_contig<typename TestFixture::Scalar>(); }
+
+}  // namespace