feat(core): expose MPI-based all-gather

Also add associated unit test.

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

Author: dssgabriel

src/KokkosComm/collective.hpp

@@ -9,6 +9,12 @@
 
 #include "fwd.hpp"
 #include "reduction_op.hpp"
+#if defined(KOKKOSCOMM_ENABLE_MPI)
+#include "mpi/mpi_space.hpp"
+#include "mpi/handle.hpp"
+#include "mpi/req.hpp"
+#include "mpi/allgather.hpp"
+#endif
 #if defined(KOKKOSCOMM_ENABLE_NCCL)
 #include "nccl/nccl_space.hpp"
 #include "nccl/handle.hpp"

src/KokkosComm/mpi/allgather.hpp

@@ -14,7 +14,8 @@
 #include "impl/types.hpp"
 #include "impl/error_handling.hpp"
 
-namespace KokkosComm::mpi {
+namespace KokkosComm {
+namespace mpi {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SView, KokkosView RView>
 auto iallgather(const ExecSpace& space, const SView sv, RView rv, MPI_Comm comm) -> Req<Mpi> {
@@ -92,3 +93,14 @@ void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, M
 }
 
 }  // namespace KokkosComm::mpi
+namespace Experimental::Impl {
+
+template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace>
+struct AllGather<SendView, RecvView, ExecSpace, Mpi> {
+  static auto execute(Handle<ExecSpace, Mpi>& h, const SendView sv, RecvView rv) -> Req<Mpi> {
+    return mpi::iallgather(h.space(), sv, rv, h.comm());
+  }
+};
+
+}  // namespace Experimental::Impl
+}  // namespace KokkosComm

unit_tests/test_allgather.cpp

@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#include <gtest/gtest.h>
+#include <Kokkos_Core.hpp>
+#include <KokkosComm/KokkosComm.hpp>
+
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+#include "nccl/utils.hpp"
+#endif
+
+namespace {
+
+template <typename T>
+class AllGather : 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(AllGather, ScalarTypes);
+
+template <typename Scalar>
+auto allgather_0d() -> void {
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+  using ExecSpace = Kokkos::Cuda;
+  auto nccl_ctx   = test_utils::nccl::Ctx::init();
+  KokkosComm::Handle<ExecSpace, KokkosComm::Nccl> h(ExecSpace(), nccl_ctx.comm());
+#else
+  using ExecSpace = Kokkos::DefaultExecutionSpace;
+  KokkosComm::Handle<Kokkos::DefaultExecutionSpace, KokkosComm::Mpi> h{};
+#endif
+  int rank = h.rank();
+  int size = h.size();
+
+  Kokkos::View<Scalar> sv("sv");
+  Kokkos::View<Scalar *> rv("rv", size);
+
+  // Prepare send view, 1 element per sender: their rank
+  Kokkos::parallel_for(
+      Kokkos::RangePolicy(ExecSpace(), 0, sv.extent(0)), KOKKOS_LAMBDA(const int) { sv() = rank; });
+  // Using the same execution space for both operations lets us not need an explicit `fence`
+  auto req = KokkosComm::Experimental::allgather(h, sv, rv);
+  KokkosComm::wait(req);
+
+  int errs;
+  Kokkos::parallel_reduce(
+      rv.extent(0), KOKKOS_LAMBDA(const int src, int &lsum) { lsum += rv(src) != src; }, errs);
+  EXPECT_EQ(errs, 0);
+}
+
+template <typename Scalar>
+auto allgather_contig_1d() -> void {
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+  using ExecSpace = Kokkos::Cuda;
+  auto nccl_ctx   = test_utils::nccl::Ctx::init();
+  KokkosComm::Handle<ExecSpace, KokkosComm::Nccl> h(ExecSpace(), nccl_ctx.comm());
+#else
+  using ExecSpace = Kokkos::DefaultExecutionSpace;
+  KokkosComm::Handle<Kokkos::DefaultExecutionSpace, KokkosComm::Mpi> h{};
+#endif
+  int rank = h.rank();
+  int size = h.size();
+
+  const int n_contrib = 100;
+  Kokkos::View<Scalar *> sv("sv", n_contrib);
+  Kokkos::View<Scalar *> rv("rv", size * n_contrib);
+
+  // Prepare send view
+  Kokkos::parallel_for(
+      Kokkos::RangePolicy(ExecSpace(), 0, sv.extent(0)), KOKKOS_LAMBDA(const int i) { sv(i) = rank + i; });
+  // Using the same execution space for both operations lets us not need an explicit `fence`
+  auto req = KokkosComm::Experimental::allgather(h, sv, rv);
+  KokkosComm::wait(req);
+
+  int errs;
+  Kokkos::parallel_reduce(
+      rv.extent(0),
+      KOKKOS_LAMBDA(const int i, int &lsum) {
+        const int src = i / n_contrib;
+        const int j   = i % n_contrib;
+        lsum += rv(i) != src + j;
+      },
+      errs);
+  EXPECT_EQ(errs, 0);
+}
+
+TYPED_TEST(AllGather, 0D) { allgather_0d<typename TestFixture::Scalar>(); }
+TYPED_TEST(AllGather, Contiguous1D) { allgather_contig_1d<typename TestFixture::Scalar>(); }
+
+}  // namespace