feat(core): expose MPI-based allreduce

Added associated unit test.

Signed-off-by: Gabriel Dos Santos <[email protected]>

Author: dssgabriel

src/KokkosComm/mpi/allreduce.hpp

@@ -15,7 +15,8 @@
 #include "impl/error_handling.hpp"
 #include "impl/types.hpp"
 
-namespace KokkosComm::mpi {
+namespace KokkosComm {
+namespace mpi {
 
 template <KokkosView SView, KokkosView RView, KokkosExecutionSpace ExecSpace>
 auto iallreduce(const ExecSpace &space, const SView sv, RView rv, MPI_Op op, MPI_Comm comm) -> Req<MpiSpace> {
@@ -104,4 +105,15 @@ void allreduce(ExecSpace const &space, View const &v, MPI_Op op, MPI_Comm comm)
   Kokkos::Tools::popRegion();
 }
 
-}  // namespace KokkosComm::mpi
+}  // namespace mpi
+namespace Experimental::Impl {
+
+template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp, KokkosExecutionSpace ExecSpace>
+struct AllReduce<SendView, RecvView, RedOp, ExecSpace, MpiSpace> {
+  static auto execute(Handle<ExecSpace, MpiSpace> &h, const SendView &sv, RecvView rv) -> Req<MpiSpace> {
+    return mpi::iallreduce(h.space(), sv, rv, reduction_op<MpiSpace, RedOp>(), h.comm());
+  }
+};
+
+}  // namespace Experimental::Impl
+}  // namespace KokkosComm

unit_tests/test_allreduce.cpp

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#include <gtest/gtest.h>
+#include <mpi.h>
+#include <Kokkos_Core.hpp>
+#include <KokkosComm/KokkosComm.hpp>
+
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+#include "nccl/utils.hpp"
+#endif
+
+namespace {
+
+template <typename T>
+class AllReduce : public testing::Test {
+ public:
+  using Scalar = T;
+};
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+using ScalarTypes = testing::Types<float, double, int, int64_t>;
+#else
+using ScalarTypes =
+    testing::Types<float, double, Kokkos::complex<float>, Kokkos::complex<double>, int, unsigned, int64_t, size_t>;
+#endif
+TYPED_TEST_SUITE(AllReduce, ScalarTypes);
+
+template <typename Scalar>
+auto allreduce_0d() -> void {
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+  using ExecSpace = Kokkos::Cuda;
+  auto nccl_ctx   = test_utils::nccl::Ctx::init();
+  ExecSpace space(nccl_ctx.stream());
+  KokkosComm::Handle<ExecSpace, KokkosComm::NcclSpace> h(space, nccl_ctx.comm());
+#else
+  using ExecSpace = Kokkos::DefaultExecutionSpace;
+  ExecSpace space{};
+  KokkosComm::Handle<ExecSpace, KokkosComm::MpiSpace> h{};
+#endif
+  int rank = h.rank();
+  int size = h.size();
+
+  Kokkos::View<Scalar> sv("sv");
+  Kokkos::View<Scalar> rv("rv");
+
+  // Prepare send buffer
+  Kokkos::parallel_for(
+      Kokkos::RangePolicy(space, 0, sv.extent(0)), KOKKOS_LAMBDA(const int) { sv() = rank; });
+  space.fence();
+  auto req = KokkosComm::Experimental::allreduce(h, sv, rv, KokkosComm::Sum{});
+  KokkosComm::wait(req);
+
+  int errs;
+  Kokkos::parallel_reduce(
+      rv.extent(0), KOKKOS_LAMBDA(const int, int &lsum) { lsum += (rv() != size * (size - 1) / 2); }, errs);
+  EXPECT_EQ(errs, 0);
+}
+
+template <typename Scalar>
+auto allreduce_contig_1d() -> void {
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+  using ExecSpace = Kokkos::Cuda;
+  auto nccl_ctx   = test_utils::nccl::Ctx::init();
+  auto space      = ExecSpace(nccl_ctx.stream());
+  KokkosComm::Handle<ExecSpace, KokkosComm::NcclSpace> h(space, nccl_ctx.comm());
+#else
+  using ExecSpace = Kokkos::DefaultExecutionSpace;
+  auto space      = ExecSpace();
+  KokkosComm::Handle<ExecSpace, KokkosComm::MpiSpace> h(space, MPI_COMM_WORLD);
+#endif
+  int rank = h.rank();
+  int size = h.size();
+
+  int n_contrib = 10;
+  Kokkos::View<Scalar *> sv("sv", n_contrib);
+  Kokkos::View<Scalar *> rv("rv", n_contrib);
+
+  // Prepare send buffer
+  Kokkos::parallel_for(
+      Kokkos::RangePolicy(space, 0, sv.extent(0)), KOKKOS_LAMBDA(const int i) { sv(i) = rank + i; });
+  space.fence();
+  auto req = KokkosComm::Experimental::allreduce(h, sv, rv, KokkosComm::Sum{});
+  KokkosComm::wait(req);
+
+  int errs;
+  // Fill in this reduction which verifies that each element computed the correct value
+  Kokkos::parallel_reduce(
+      rv.extent(0),
+      KOKKOS_LAMBDA(const int i, int &lsum) { lsum += (rv(i) != ((size * (size - 1)) / 2 + (size * i))); }, errs);
+  EXPECT_EQ(errs, 0);
+}
+
+TYPED_TEST(AllReduce, 0D) { allreduce_0d<typename TestFixture::Scalar>(); }
+TYPED_TEST(AllReduce, Contiguous1D) { allreduce_contig_1d<typename TestFixture::Scalar>(); }
+
+}  // namespace