feat(core): expose MPI-based broadcast

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/broadcast.hpp"
+#endif
 #if defined(KOKKOSCOMM_ENABLE_NCCL)
 #include "nccl/nccl_space.hpp"
 #include "nccl/handle.hpp"

src/KokkosComm/mpi/broadcast.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 View>
 auto ibroadcast(const ExecSpace& space, View& v, int root, MPI_Comm comm) -> Req<Mpi> {
@@ -56,4 +57,15 @@ void broadcast(ExecSpace const& space, View const& v, int root, MPI_Comm comm) {
   Kokkos::Tools::popRegion();
 }
 
-}  // namespace KokkosComm::mpi
+}  // namespace mpi
+namespace Experimental::Impl {
+
+template <KokkosView View, KokkosExecutionSpace ExecSpace>
+struct Broadcast<View, ExecSpace, Mpi> {
+  static auto execute(Handle<ExecSpace, Mpi>& h, View& v, int root) -> Req<Mpi> {
+    return KokkosComm::mpi::ibroadcast(h.space(), v, root, h.comm());
+  }
+};
+
+}  // namespace Experimental::Impl
+}  // namespace KokkosComm

unit_tests/test_broadcast.cpp

@@ -0,0 +1,90 @@
+// 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 Broadcast : 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(Broadcast, ScalarTypes);
+
+template <typename Scalar>
+auto broadcast_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();
+  int root = 0;
+
+  Kokkos::View<Scalar> v("v");
+  if (rank == root) {
+    // Prepare broadcast view
+    Kokkos::parallel_for(
+        Kokkos::RangePolicy(ExecSpace(), 0, v.extent(0)), KOKKOS_LAMBDA(const int) { v() = size; });
+  }
+  // Using the same execution space for both operations lets us not need an explicit `fence`
+  auto req = KokkosComm::Experimental::broadcast(h, v, root);
+  KokkosComm::wait(req);
+
+  int errs;
+  Kokkos::parallel_reduce(
+      v.extent(0), KOKKOS_LAMBDA(const int, int& lsum) { lsum += v() != size; }, errs);
+  EXPECT_EQ(errs, 0);
+}
+
+template <typename Scalar>
+auto broadcast_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();
+  int root = 0;
+
+  Kokkos::View<Scalar*> v("v", 100);
+  if (rank == root) {
+    // Prepare broadcast view
+    Kokkos::parallel_for(
+        Kokkos::RangePolicy(ExecSpace(), 0, v.extent(0)), KOKKOS_LAMBDA(const int i) { v(i) = size + i; });
+  }
+  // Using the same execution space for both operations lets us not need an explicit `fence`
+  auto req = KokkosComm::Experimental::broadcast(h, v, root);
+  KokkosComm::wait(req);
+
+  int errs;
+  Kokkos::parallel_reduce(
+      v.extent(0), KOKKOS_LAMBDA(const int i, int& lsum) { lsum += (v(i) != size + i); }, errs);
+  EXPECT_EQ(errs, 0);
+}
+
+TYPED_TEST(Broadcast, 0D) { broadcast_0d<typename TestFixture::Scalar>(); }
+TYPED_TEST(Broadcast, Contiguous1D) { broadcast_contig_1d<typename TestFixture::Scalar>(); }
+
+}  // namespace