Merge branch 'develop' into feature/core-alltoall

Author: dssgabriel

src/KokkosComm/collective.hpp

@@ -28,20 +28,11 @@
 
 namespace KokkosComm::Experimental {
 
-/// Copy the `sv` view on the `root` rank to all ranks' `rv` view.
-///
-/// The `sv` view is only used on the `root` rank and ignored for all other ranks.
-template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
-          CommunicationSpace CommSpace = DefaultCommunicationSpace>
-auto broadcast(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv, int root) -> Req<CommSpace> {
-  return Impl::Broadcast<SendView, RecvView, ExecSpace, CommSpace>::execute(h, sv, rv, root);
-}
-
-/// In-place variant of `broadcast`. Copy the `v` view from the `root` rank to all ranks' `v` view.
+/// Copy the `v` view from the `root` rank to all ranks' `v` view.
 template <KokkosView View, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace = DefaultCommunicationSpace>
 auto broadcast(Handle<ExecSpace, CommSpace>& h, View v, int root) -> Req<CommSpace> {
-  return Impl::Broadcast<View, View, ExecSpace, CommSpace>::execute(h, v, v, root);
+  return Impl::Broadcast<View, ExecSpace, CommSpace>::execute(h, v, root);
 }
 
 /// Copy the `sv` view from each rank to the `rv` view, receiving data from rank `i` at offset

src/KokkosComm/fwd.hpp

@@ -47,7 +47,7 @@ struct Send;
 // Collectives are currently experimental functions
 namespace Experimental::Impl {
 
-template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+template <KokkosView View, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace = DefaultCommunicationSpace>
 struct Broadcast;
 

src/KokkosComm/nccl/broadcast.hpp

@@ -17,25 +17,20 @@ namespace nccl {
 
 namespace KC = KokkosComm;
 
-template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView>
-auto broadcast(const ExecSpace &space, const SendView &sv, const RecvView &rv, int root, ncclComm_t comm) -> Req<Nccl> {
-  using ST = typename SendView::non_const_value_type;
-  using RT = typename RecvView::non_const_value_type;
-  static_assert(std::is_same_v<ST, RT>,
-                "KokkosComm::Experimental::nccl::broadcast: View value types must be identical");
-  static_assert(KC::rank<SendView>() <= 1 and KC::rank<RecvView>() <= 1,
+template <KokkosView View>
+auto broadcast(const Kokkos::Cuda& space, View& v, int root, ncclComm_t comm) -> Req<Nccl> {
+  using T = typename View::non_const_value_type;
+  static_assert(KC::rank<View>() <= 1,
                 "KokkosComm::Experimental::nccl::broadcast: Views with rank higher than 1 are not supported");
   Kokkos::Tools::pushRegion("KokkosComm::Experimental::nccl::broadcast");
 
   Req<Nccl> req{space.cuda_stream()};
-  if (KC::is_contiguous(sv) and KC::is_contiguous(rv)) {
-    ncclBroadcast(KC::data_handle(sv), KC::data_handle(rv), KC::span(sv), Impl::datatype_v<ST>, root, comm,
-                  space.cuda_stream());
+  if (KC::is_contiguous(v)) {
+    ncclBcast(KC::data_handle(v), KC::span(v), Impl::datatype_v<T>, root, comm, space.cuda_stream());
   } else {
     Kokkos::abort("KokkosComm::Experimental::nccl::broadcast: unimplemented for non-contiguous views");
   }
-  req.extend_view_lifetime(sv);
-  req.extend_view_lifetime(rv);
+  req.extend_view_lifetime(v);
 
   Kokkos::Tools::popRegion();
   return req;
@@ -44,10 +39,10 @@ auto broadcast(const ExecSpace &space, const SendView &sv, const RecvView &rv, i
 }  // namespace nccl
 namespace Impl {
 
-template <KokkosView SendView, KokkosView RecvView>
-struct Broadcast<SendView, RecvView, Kokkos::Cuda, Nccl> {
-  static auto execute(Handle<Kokkos::Cuda, Nccl> &h, const SendView sv, RecvView rv, int root) -> Req<Nccl> {
-    return nccl::broadcast(h.space(), sv, rv, root, h.comm());
+template <KokkosView View>
+struct Broadcast<View, Kokkos::Cuda, Nccl> {
+  static auto execute(Handle<Kokkos::Cuda, Nccl>& h, View v, int root) -> Req<Nccl> {
+    return nccl::broadcast(h.space(), v, root, h.comm());
   }
 };
 

unit_tests/nccl/test_broadcast.cpp

@@ -42,20 +42,20 @@ auto broadcast_0d() -> void {
 
   int errs;
   Kokkos::parallel_reduce(
-      v.extent(0), KOKKOS_LAMBDA(const int, int &lsum) { lsum += v() != size; }, errs);
+      v.extent(0), KOKKOS_LAMBDA(const int, int& lsum) { lsum += v() != size; }, errs);
   EXPECT_EQ(errs, 0);
 }
 
 template <typename Scalar>
-auto broadcast_inplace_contig_1d() -> void {
+auto broadcast_contig_1d() -> void {
   auto nccl_ctx = test_utils::nccl::Ctx::init();
   ExecSpace space(nccl_ctx.stream());
   KokkosComm::Handle<ExecSpace, CommSpace> h(space, nccl_ctx.comm());
   int rank = h.rank();
   int size = h.size();
   int root = 0;
 
-  Kokkos::View<Scalar *> v("v", 100);
+  Kokkos::View<Scalar*> v("v", 100);
   if (rank == root) {
     // Prepare broadcast view
     Kokkos::parallel_for(
@@ -67,11 +67,11 @@ auto broadcast_inplace_contig_1d() -> void {
 
   int errs;
   Kokkos::parallel_reduce(
-      v.extent(0), KOKKOS_LAMBDA(const int i, int &lsum) { lsum += (v(i) != size + i); }, errs);
+      v.extent(0), KOKKOS_LAMBDA(const int i, int& lsum) { lsum += (v(i) != size + i); }, errs);
   EXPECT_EQ(errs, 0);
 }
 
-TYPED_TEST(Broadcast, InPlace0D) { broadcast_0d<typename TestFixture::Scalar>(); }
-TYPED_TEST(Broadcast, InPlaceContiguous1D) { broadcast_inplace_contig_1d<typename TestFixture::Scalar>(); }
+TYPED_TEST(Broadcast, 0D) { broadcast_0d<typename TestFixture::Scalar>(); }
+TYPED_TEST(Broadcast, Contiguous1D) { broadcast_contig_1d<typename TestFixture::Scalar>(); }
 
 }  // namespace