refactor!: add more requirements to the CommunicationSpace concept (#188)

* refactor!: rework & improve communication space concept

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

* chore: format + fix CommSpace typo in unit test

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

* Fix missing MpiSpace

Signed-off-by: Cedric Chevalier <[email protected]>

* Fix format

Signed-off-by: Cedric Chevalier <[email protected]>

* Update documentation

Signed-off-by: Cedric Chevalier <[email protected]>

---------

Signed-off-by: Gabriel Dos Santos <[email protected]>
Signed-off-by: Cedric Chevalier <[email protected]>
Co-authored-by: Cédric Chevalier <[email protected]>

Author: dssgabriel

docs/dev/impl_comm_space.rst

@@ -20,9 +20,8 @@ For example, for the MPI communication space, we define the following:
 
     namespace KokkosComm {
 
-    struct Mpi {
-      static auto world_size() noexcept -> int { /* ... */ }
-      static auto world_rank() noexcept -> int { /* ... */ }
+    struct MpiSpace {
+      ...
     };
 
     template <>
@@ -31,7 +30,6 @@ For example, for the MPI communication space, we define the following:
     } // end KokkosComm
 
 
-Notice that ``Mpi`` has two static methods, but that these methods are not required. The main point is that ``struct Mpi`` exists.
 To let core API functions know that your communication space is something KokkosComm can use to dispatch messages, you also need to declare the ``Impl::is_communication_space`` specialization using the ``CommunicationSpace`` concept.
 
 
@@ -71,7 +69,7 @@ For example, for the MPI communication space request, we define the following:
     namespace KokkosComm {
 
     template <>
-    class Req<Mpi> { /* ... */ };
+    class Req<MpiSpace> { /* ... */ };
 
     } // end KokkosComm
 
@@ -103,7 +101,7 @@ For example, for the MPI communication space, we create a partial specialization
     namespace KokkosComm::Impl {
 
     template <KokkosView RecvView, KokkosExecutionSpace ExecSpace>
-    struct Recv<RecvView, ExecSpace, Mpi> { /* ... */ };
+    struct Recv<RecvView, ExecSpace, MpiSpace> { /* ... */ };
 
     } // end KokkosComm::Impl
 

src/KokkosComm/concepts.hpp

@@ -27,7 +27,15 @@ template <typename T>
 concept KokkosExecutionSpace = Kokkos::is_execution_space_v<T>;
 
 template <typename T>
-concept CommunicationSpace = KokkosComm::Impl::is_communication_space<T>::value;
+concept CommunicationSpace = requires {
+  KokkosComm::Impl::is_communication_space<T>::value;
+  typename T::communication_space;
+  typename T::handle_type;
+  typename T::request_type;
+  typename T::datatype_type;
+  typename T::reduction_op_type;
+  typename T::rank_type;
+};
 
 template <typename T>
 concept ReductionOperator = KokkosComm::Impl::is_reduction_operator<T>::value;

src/KokkosComm/fwd.hpp

@@ -9,18 +9,19 @@
 
 namespace KokkosComm {
 
-// NCCL backend also implicitly declares MPI
 #if defined(KOKKOSCOMM_ENABLE_NCCL)
 namespace Experimental {
-class Nccl;
-}  // namespace Experimental
-class Mpi;
-using DefaultCommunicationSpace  = Experimental::Nccl;
-using FallbackCommunicationSpace = Mpi;
+struct NcclSpace;
+}
+// NCCL backend also declares the MPI space as fallback
+struct MpiSpace;
+
+using DefaultCommunicationSpace  = Experimental::NcclSpace;
+using FallbackCommunicationSpace = MpiSpace;
 #elif defined(KOKKOSCOMM_ENABLE_MPI)
-class Mpi;
-using DefaultCommunicationSpace  = Mpi;
-using FallbackCommunicationSpace = Mpi;
+struct MpiSpace;
+using DefaultCommunicationSpace  = MpiSpace;
+using FallbackCommunicationSpace = MpiSpace;
 #else
 #error at least one communication space must be enabled
 #endif

src/KokkosComm/mpi/allgather.hpp

@@ -17,7 +17,7 @@
 namespace KokkosComm::mpi {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SView, KokkosView RView>
-auto iallgather(const ExecSpace &space, const SView sv, RView rv, MPI_Comm comm) -> Req<Mpi> {
+auto iallgather(const ExecSpace &space, const SView sv, RView rv, MPI_Comm comm) -> Req<MpiSpace> {
   using ST = typename SView::non_const_value_type;
   using RT = typename RView::non_const_value_type;
   static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::iallgather: View value types must be identical");
@@ -29,7 +29,7 @@ auto iallgather(const ExecSpace &space, const SView sv, RView rv, MPI_Comm comm)
   // Sync: Work in space may have been used to produce view data.
   space.fence("fence before non-blocking all-gather");
 
-  Req<Mpi> req;
+  Req<MpiSpace> req;
   // All ranks send/recv same count
   MPI_Iallgather(data_handle(sv), span(sv), Impl::mpi_type_v<ST>, data_handle(rv), span(sv), Impl::mpi_type_v<RT>, comm,
                  &req.mpi_request());

src/KokkosComm/mpi/alltoall.hpp

@@ -18,7 +18,7 @@
 namespace KokkosComm::mpi {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SView, KokkosView RView>
-auto ialltoall(const ExecSpace &space, const SView sv, RView rv, int count, MPI_Comm comm) -> Req<Mpi> {
+auto ialltoall(const ExecSpace &space, const SView sv, RView rv, int count, MPI_Comm comm) -> Req<MpiSpace> {
   using ST = typename SView::non_const_value_type;
   using RT = typename RView::non_const_value_type;
   static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::ialltoall: View value types must be identical");
@@ -30,7 +30,7 @@ auto ialltoall(const ExecSpace &space, const SView sv, RView rv, int count, MPI_
   // Sync: Work in space may have been used to produce view data.
   space.fence("fence before non-blocking all-gather");
 
-  Req<Mpi> req;
+  Req<MpiSpace> req;
   // All ranks send/recv same count
   MPI_Ialltoall(data_handle(sv), count, Impl::mpi_type_v<ST>, data_handle(rv), count, Impl::mpi_type_v<RT>, comm,
                 &req.mpi_request());

src/KokkosComm/mpi/broadcast.hpp

@@ -17,15 +17,15 @@
 namespace KokkosComm::mpi {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView View>
-auto ibroadcast(const ExecSpace& space, View& v, int root, MPI_Comm comm) -> Req<Mpi> {
+auto ibroadcast(const ExecSpace& space, View& v, int root, MPI_Comm comm) -> Req<MpiSpace> {
   using T = typename View::non_const_value_type;
   Kokkos::Tools::pushRegion("KokkosComm::mpi::ibroadcast");
   fail_if(!is_contiguous(v), "KokkosComm::mpi::ibroadcast: unimplemented for non-contiguous views");
 
   // Sync: Work in space may have been used to produce view data.
   space.fence("fence before non-blocking broadcast");
 
-  Req<Mpi> req;
+  Req<MpiSpace> req;
   MPI_Ibcast(data_handle(v), span(v), Impl::mpi_type_v<T>, root, comm, &req.mpi_request());
   req.extend_view_lifetime(v);
 

src/KokkosComm/mpi/channel.hpp

@@ -57,7 +57,7 @@ class Channel {
 
   void wait() {
     Kokkos::Tools::pushRegion("KokkosComm::Channel::wait");
-    std::vector<Req<Mpi>> reqs;
+    std::vector<Req<MpiSpace>> reqs;
     reqs.reserve(send_reqs_.size() + recv_reqs_.size());
     reqs.insert(reqs.end(), send_reqs_.begin(), send_reqs_.end());
     reqs.insert(reqs.end(), recv_reqs_.begin(), recv_reqs_.end());
@@ -66,12 +66,12 @@ class Channel {
   }
 
  private:
-  std::vector<Req<Mpi>> send_reqs_;  // Queue for send requests
-  std::vector<Req<Mpi>> recv_reqs_;  // Queue for receive requests
-  int dest_rank_;                    // Destination rank for send
-  int src_rank_;                     // Source rank for receive
-  int tag_;                          // MPI tag
-  MPI_Comm comm_;                    // MPI communicator
+  std::vector<Req<MpiSpace>> send_reqs_;  // Queue for send requests
+  std::vector<Req<MpiSpace>> recv_reqs_;  // Queue for receive requests
+  int dest_rank_;                         // Destination rank for send
+  int src_rank_;                          // Source rank for receive
+  int tag_;                               // MPI tag
+  MPI_Comm comm_;                         // MPI communicator
 };
 
 }  // namespace KokkosComm

src/KokkosComm/mpi/handle.hpp

@@ -20,10 +20,10 @@ namespace KokkosComm {
 - post-wait
 */
 template <KokkosExecutionSpace ExecSpace>
-class Handle<ExecSpace, Mpi> {
+class Handle<ExecSpace, MpiSpace> {
  public:
   using execution_space = ExecSpace;
-  using transport_type  = Mpi;
+  using transport_type  = MpiSpace;
   using size_type       = int;
 
   explicit Handle(const execution_space &space, MPI_Comm comm) : space_(space), comm_(comm) {}

src/KokkosComm/mpi/irecv.hpp

@@ -18,15 +18,15 @@ namespace Impl {
 
 // Recv implementation for Mpi
 template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
-struct Recv<RecvView, ExecSpace, Mpi> {
-  static Req<Mpi> execute(Handle<ExecSpace, Mpi> &h, const RecvView &rv, int src) {
+struct Recv<RecvView, ExecSpace, MpiSpace> {
+  static Req<MpiSpace> execute(Handle<ExecSpace, MpiSpace> &h, const RecvView &rv, int src) {
     using KCPT   = KokkosComm::PackTraits<RecvView>;
     using Packer = typename KCPT::packer_type;
     using Args   = typename Packer::args_type;
 
     const ExecSpace &space = h.space();
 
-    Req<Mpi> req;
+    Req<MpiSpace> req;
     if (KokkosComm::is_contiguous(rv)) {
       space.fence("fence before irecv");
       MPI_Irecv(KokkosComm::data_handle(rv), KokkosComm::span(rv), mpi_type_v<typename RecvView::value_type>, src,

src/KokkosComm/mpi/isend.hpp

@@ -20,7 +20,7 @@ namespace KokkosComm {
 namespace Impl {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView, mpi::CommunicationMode SendMode>
-Req<Mpi> isend_impl(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest, int tag, SendMode) {
+Req<MpiSpace> isend_impl(Handle<ExecSpace, MpiSpace> &h, const SendView &sv, int dest, int tag, SendMode) {
   auto mpi_isend_fn = [](void *mpi_view, int mpi_count, MPI_Datatype mpi_datatype, int mpi_dest, int mpi_tag,
                          MPI_Comm mpi_comm, MPI_Request *mpi_req) {
     if constexpr (std::is_same_v<SendMode, mpi::CommModeStandard>) {
@@ -34,7 +34,7 @@ Req<Mpi> isend_impl(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest, int
     }
   };
 
-  Req<Mpi> req;
+  Req<MpiSpace> req;
   if (KokkosComm::is_contiguous(sv)) {
     h.space().fence("fence before isend");
     mpi_isend_fn(KokkosComm::data_handle(sv), KokkosComm::span(sv), mpi_type_v<typename SendView::value_type>, dest,
@@ -55,8 +55,8 @@ Req<Mpi> isend_impl(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest, int
 
 // Implementation of KokkosComm::Send
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
-struct Send<SendView, ExecSpace, Mpi> {
-  static Req<Mpi> execute(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest) {
+struct Send<SendView, ExecSpace, MpiSpace> {
+  static Req<MpiSpace> execute(Handle<ExecSpace, MpiSpace> &h, const SendView &sv, int dest) {
     return isend_impl<ExecSpace, SendView>(h, sv, dest, POINTTOPOINT_TAG, mpi::DefaultCommMode{});
   }
 };
@@ -65,12 +65,12 @@ struct Send<SendView, ExecSpace, Mpi> {
 namespace mpi {
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView, CommunicationMode SendMode>
-Req<Mpi> isend(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest, int tag, SendMode) {
+Req<MpiSpace> isend(Handle<ExecSpace, MpiSpace> &h, const SendView &sv, int dest, int tag, SendMode) {
   return KokkosComm::Impl::isend_impl<ExecSpace, SendView>(h, sv, dest, tag, SendMode{});
 }
 
 template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
-Req<Mpi> isend(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest, int tag) {
+Req<MpiSpace> isend(Handle<ExecSpace, MpiSpace> &h, const SendView &sv, int dest, int tag) {
   return isend<ExecSpace, SendView>(h, sv, dest, tag, DefaultCommMode{});
 }