refactor(nccl): propagate Communicator changes to NCCL backend

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

Author: dssgabriel

src/KokkosComm/nccl/allgather.hpp

@@ -24,14 +24,16 @@ template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvVi
 auto allgather(const ExecSpace& space, const SendView& sv, const RecvView& rv, ncclComm_t comm) -> Request<NcclSpace> {
   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::allgather: View value types must be identical");
+  static_assert(
+      std::is_same_v<ST, RT>, "KokkosComm::Experimental::nccl::allgather: View value types must be identical"
+  );
   Kokkos::Tools::pushRegion("KokkosComm::Experimental::nccl::allgather");
 
   Request<NcclSpace> req;
   if (KC::is_contiguous(sv) and KC::is_contiguous(rv)) {
-    ncclAllGather(KC::data_handle(sv), KC::data_handle(rv), KC::span(sv), datatype<NcclSpace, ST>(), comm,
-                  space.cuda_stream());
+    ncclAllGather(
+        KC::data_handle(sv), KC::data_handle(rv), KC::span(sv), datatype<NcclSpace, ST>(), comm, space.cuda_stream()
+    );
     req.capture_stream_state(space.cuda_stream());
   } else {
     Kokkos::abort("KokkosComm::Experimental::nccl::allgather: unimplemented for non-contiguous views");
@@ -48,8 +50,8 @@ namespace Impl {
 
 template <KokkosView SendView, KokkosView RecvView>
 struct AllGather<SendView, RecvView, Kokkos::Cuda, NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, NcclSpace>& h, const SendView sv, RecvView rv) -> Request<NcclSpace> {
-    return nccl::allgather(h.space(), sv, rv, h.comm());
+  static auto execute(Communicator<NcclSpace, Kokkos::Cuda>& h, const SendView sv, RecvView rv) -> Request<NcclSpace> {
+    return nccl::allgather(h.exec(), sv, rv, h.comm());
   }
 };
 

src/KokkosComm/nccl/allreduce.hpp

@@ -26,14 +26,16 @@ auto allreduce(const ExecSpace& space, const SendView& sv, const RecvView& rv, n
     -> Request<NcclSpace> {
   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::allreduce: View value types must be identical");
+  static_assert(
+      std::is_same_v<ST, RT>, "KokkosComm::Experimental::nccl::allreduce: View value types must be identical"
+  );
   Kokkos::Tools::pushRegion("KokkosComm::Experimental::nccl::allreduce");
 
   Request<NcclSpace> req;
   if (KC::is_contiguous(sv) and KC::is_contiguous(rv)) {
-    ncclAllReduce(KC::data_handle(sv), KC::data_handle(rv), KC::span(sv), datatype<NcclSpace, ST>(), op, comm,
-                  space.cuda_stream());
+    ncclAllReduce(
+        KC::data_handle(sv), KC::data_handle(rv), KC::span(sv), datatype<NcclSpace, ST>(), op, comm, space.cuda_stream()
+    );
     req.capture_stream_state(space.cuda_stream());
   } else {
     Kokkos::abort("KokkosComm::Experimental::nccl::allreduce: unimplemented for non-contiguous Views");
@@ -50,8 +52,8 @@ namespace Impl {
 
 template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp>
 struct AllReduce<SendView, RecvView, RedOp, Kokkos::Cuda, NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, NcclSpace>& h, const SendView sv, RecvView rv) -> Request<NcclSpace> {
-    return nccl::allreduce(h.space(), sv, rv, reduction_op<NcclSpace, RedOp>(), h.comm());
+  static auto execute(Communicator<NcclSpace, Kokkos::Cuda>& h, const SendView sv, RecvView rv) -> Request<NcclSpace> {
+    return nccl::allreduce(h.exec(), sv, rv, reduction_op<NcclSpace, RedOp>(), h.comm());
   }
 };
 

src/KokkosComm/nccl/alltoall.hpp

@@ -58,9 +58,9 @@ namespace Impl {
 
 template <KokkosView SendView, KokkosView RecvView>
 struct AllToAll<SendView, RecvView, Kokkos::Cuda, NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, NcclSpace>& h, const SendView sv, RecvView rv, int count)
+  static auto execute(Communicator<NcclSpace, Kokkos::Cuda>& h, const SendView sv, RecvView rv, int count)
       -> Request<NcclSpace> {
-    return nccl::alltoall(h.space(), sv, rv, count, h.comm());
+    return nccl::alltoall(h.exec(), sv, rv, count, h.comm());
   }
 };
 

src/KokkosComm/nccl/broadcast.hpp

@@ -23,8 +23,10 @@ namespace KC = KokkosComm;
 template <KokkosView View>
 auto broadcast(const Kokkos::Cuda& space, View& v, int root, ncclComm_t comm) -> Request<NcclSpace> {
   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");
+  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");
 
   Request<NcclSpace> req;
@@ -45,8 +47,8 @@ namespace Impl {
 
 template <KokkosView View>
 struct Broadcast<View, Kokkos::Cuda, NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, NcclSpace>& h, View v, int root) -> Request<NcclSpace> {
-    return nccl::broadcast(h.space(), v, root, h.comm());
+  static auto execute(Communicator<NcclSpace, Kokkos::Cuda>& h, View v, int root) -> Request<NcclSpace> {
+    return nccl::broadcast(h.exec(), v, root, h.comm());
   }
 };
 

src/KokkosComm/nccl/communicator.hpp

@@ -28,7 +28,7 @@ class Communicator<Experimental::NcclSpace, Kokkos::Cuda> {
   /// Defaults `exec` to `Kokkos::Cuda`.
   ///
   /// The returned communicator does not own the underlying handle, and the user is responsible for destroying it.
-  [[nodiscard]] static auto from_raw(communicator_type comm, const execution_space& exec = execution_space{}, ) noexcept
+  [[nodiscard]] static auto from_raw(communicator_type comm, const execution_space& exec = execution_space{}) noexcept
       -> std::optional<Communicator<communication_space, execution_space>> {
     if (comm == nullptr) {
       return std::nullopt;
@@ -42,16 +42,17 @@ class Communicator<Experimental::NcclSpace, Kokkos::Cuda> {
   /// value of `key`. All processes with the same value of `color` join the same communicator.
   /// A process that passes `NCCL_SPLIT_NOCOLOR` as `color` will not join a new communicator and `nullopt` is returned.
   [[nodiscard]] static auto split(
-      const communicator_type comm, int color, int key, const execution_space& exec = execution_space{},
-  ) -> std::optional<Communicator<communication_space, execution_space>> {
+      const communicator_type comm, int color, int key, const execution_space& exec = execution_space{}
+  ) noexcept -> std::optional<Communicator<communication_space, execution_space>> {
     communicator_type new_comm;
     ncclCommSplit(comm, color, key, &new_comm, nullptr);
     if (new_comm == nullptr) {
       return std::nullopt;
     }
     return Communicator<communication_space, execution_space>(new_comm, exec, true);
   }
-  [[nodiscard]] auto split(int color, int key) -> std::optional<Communicator<execution_space, communication_space>> {
+  [[nodiscard]] auto split(int color, int key) noexcept
+      -> std::optional<Communicator<communication_space, execution_space>> {
     return Communicator::split(comm_, color, key, exec_);
   }
 
@@ -64,7 +65,7 @@ class Communicator<Experimental::NcclSpace, Kokkos::Cuda> {
     ncclCommUserRank(comm, &rank);
     return Communicator::split(comm, 0, rank, exec_);
   }
-  [[nodiscard]] auto duplicate() -> std::optional<Communicator<execution_space, communication_space>> {
+  [[nodiscard]] auto duplicate() noexcept -> std::optional<Communicator<communication_space, execution_space>> {
     return Communicator::duplicate(comm_, exec_);
   }
 

src/KokkosComm/nccl/recv.hpp

@@ -31,7 +31,8 @@ auto recv(const ExecSpace& space, RecvView& rv, int peer, ncclComm_t comm) -> Re
     using Packer = typename Impl::PackTraits<RecvView>::packer_type;
     auto pckd_rv = Packer::allocate_packed_for(space, "pckd_rv", rv);
     KC_NCCL_CHECK(
-        ncclRecv(data_handle(pckd_rv.view_), pckd_rv.count_, pckd_rv.datatype_, peer, comm, space.cuda_stream()));
+        ncclRecv(data_handle(pckd_rv.view_), pckd_rv.count_, pckd_rv.datatype_, peer, comm, space.cuda_stream())
+    );
     req.capture_stream_state(space.cuda_stream());
     req.add_callback([space, rv, pckd_rv]() {
       Packer::unpack_into(space, rv, pckd_rv.view_);
@@ -49,9 +50,9 @@ namespace Impl {
 
 template <KokkosView RecvView>
 struct Recv<RecvView, Kokkos::Cuda, Experimental::NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, Experimental::NcclSpace>& h, RecvView sv, int peer)
+  static auto execute(Communicator<Experimental::NcclSpace, Kokkos::Cuda>& h, RecvView sv, int peer)
       -> Request<Experimental::NcclSpace> {
-    return Experimental::nccl::recv(h.space(), sv, peer, h.comm());
+    return Experimental::nccl::recv(h.exec(), sv, peer, h.comm());
   }
 };
 

src/KokkosComm/nccl/reduce.hpp

@@ -22,8 +22,9 @@ namespace KokkosComm::Experimental {
 namespace nccl {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView>
-auto reduce(const ExecSpace& space, const SendView& sv, RecvView& rv, ncclRedOp_t op, int root, int rank,
-            ncclComm_t comm) -> Request<NcclSpace> {
+auto reduce(
+    const ExecSpace& space, const SendView& sv, RecvView& rv, ncclRedOp_t op, int root, int rank, ncclComm_t comm
+) -> Request<NcclSpace> {
   using ST         = typename SendView::non_const_value_type;
   using RT         = typename RecvView::non_const_value_type;
   using SendPacker = typename Impl::PackTraits<SendView>::packer_type;
@@ -34,13 +35,16 @@ auto reduce(const ExecSpace& space, const SendView& sv, RecvView& rv, ncclRedOp_
   Request<NcclSpace> req;
   if (is_contiguous(sv)) {
     if (rank != root and is_contiguous(rv)) {
-      ncclReduce(data_handle(sv), data_handle(rv), span(sv), datatype<NcclSpace, ST>(), op, root, comm,
-                 space.cuda_stream());
+      ncclReduce(
+          data_handle(sv), data_handle(rv), span(sv), datatype<NcclSpace, ST>(), op, root, comm, space.cuda_stream()
+      );
       req.capture_stream_state(space.cuda_stream());
     } else {
       auto pckd_rv = RecvPacker::allocate_packed_for(space, "pckd_rv", rv);
-      ncclReduce(data_handle(sv), data_handle(pckd_rv.view_), span(sv), datatype<NcclSpace, ST>(), op, root, comm,
-                 space.cuda_stream());
+      ncclReduce(
+          data_handle(sv), data_handle(pckd_rv.view_), span(sv), datatype<NcclSpace, ST>(), op, root, comm,
+          space.cuda_stream()
+      );
       req.capture_stream_state(space.cuda_stream());
       req.add_callback([space, rv, pckd_rv]() {
         RecvPacker::unpack_into(space, rv, pckd_rv.view_);
@@ -50,13 +54,17 @@ auto reduce(const ExecSpace& space, const SendView& sv, RecvView& rv, ncclRedOp_
   } else {
     auto pckd_sv = SendPacker::pack(space, "pckd_sv", sv);
     if (rank != root and is_contiguous(rv)) {
-      ncclReduce(data_handle(pckd_sv.view_), data_handle(rv), pckd_sv.count_, pckd_sv.datatype_, op, root, comm,
-                 space.cuda_stream());
+      ncclReduce(
+          data_handle(pckd_sv.view_), data_handle(rv), pckd_sv.count_, pckd_sv.datatype_, op, root, comm,
+          space.cuda_stream()
+      );
       req.capture_stream_state(space.cuda_stream());
     } else {
       auto pckd_rv = RecvPacker::allocate_packed_for(space, "pckd_rv", rv);
-      ncclReduce(data_handle(pckd_sv.view_), data_handle(pckd_rv.view_), pckd_sv.count_, pckd_sv.datatype_, op, root,
-                 comm, space.cuda_stream());
+      ncclReduce(
+          data_handle(pckd_sv.view_), data_handle(pckd_rv.view_), pckd_sv.count_, pckd_sv.datatype_, op, root, comm,
+          space.cuda_stream()
+      );
       req.capture_stream_state(space.cuda_stream());
       req.add_callback([space, rv, pckd_rv]() {
         RecvPacker::unpack_into(space, rv, pckd_rv.view_);
@@ -77,9 +85,9 @@ namespace Impl {
 
 template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp>
 struct Reduce<SendView, RecvView, RedOp, Kokkos::Cuda, NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, NcclSpace>& h, const SendView sv, RecvView rv, int root)
+  static auto execute(Communicator<NcclSpace, Kokkos::Cuda>& h, const SendView sv, RecvView rv, int root)
       -> Request<NcclSpace> {
-    return nccl::reduce(h.space(), sv, rv, reduction_op<NcclSpace, RedOp>(), root, h.rank(), h.comm());
+    return nccl::reduce(h.exec(), sv, rv, reduction_op<NcclSpace, RedOp>(), root, h.rank(), h.comm());
   }
 };
 

src/KokkosComm/nccl/send.hpp

@@ -31,7 +31,8 @@ auto send(const ExecSpace& space, const SendView& sv, int peer, ncclComm_t comm)
     using Packer = typename Impl::PackTraits<SendView>::packer_type;
     auto pckd_sv = Packer::pack(space, "pckd_sv", sv);
     KC_NCCL_CHECK(
-        ncclSend(data_handle(pckd_sv.view_), pckd_sv.count_, pckd_sv.datatype_, peer, comm, space.cuda_stream()));
+        ncclSend(data_handle(pckd_sv.view_), pckd_sv.count_, pckd_sv.datatype_, peer, comm, space.cuda_stream())
+    );
     req.capture_stream_state(space.cuda_stream());
     req.extend_view_lifetime(pckd_sv.view_);
   }
@@ -46,9 +47,9 @@ namespace Impl {
 
 template <KokkosView SendView>
 struct Send<SendView, Kokkos::Cuda, Experimental::NcclSpace> {
-  static auto execute(Handle<Kokkos::Cuda, Experimental::NcclSpace>& h, SendView sv, int peer)
+  static auto execute(Communicator<Experimental::NcclSpace, Kokkos::Cuda>& h, SendView sv, int peer)
       -> Request<Experimental::NcclSpace> {
-    return Experimental::nccl::send(h.space(), sv, peer, h.comm());
+    return Experimental::nccl::send(h.exec(), sv, peer, h.comm());
   }
 };