Skip to content

Implementing a `CommunicationSpace`

Carl Pearson edited this page Sep 12, 2024 · 12 revisions

Implementing a CommunicationSpace

#109 introduced the concept of multiple CommunicationSpaces. A CommunicationSpace defines what Kokkos Comm actually does when you call e.g. KokkosComm::send.

Your implementation is a struct or class that represents your communication space. Then you need partial specializations for some associated types (handles, requests), for each core API.

struct representing CommunicationSpace

The first piece is a struct that represents your CommunicationSpace. We're still a bit fuzzy on the interface for this thing, and ultimately, it might be very small. It's main purpose is to serve as a tag for the partial specializations of each API struct.

For example, in the MPI CommunicationSpace, we define the following

namespace KokkosComm {

struct Mpi {
  static int world_size() { ... }
  static int world_rank() { ... }
};

template <>
struct Impl::is_communication_space<KokkosComm::Mpi> : public std::true_type {};

}

You can see Mpi has two static methods, but these methods are not required. The main point is that the struct Mpi exists. You also need this Impl::is_communication_space specialization to tell KokkosComm this thing is a CommunicationSpace

Partial Specialization of Handle

Warning

coming soon...

namespace KokkosComm {
template <KokkosExecutionSpace ExecSpace, CommunicationSpace CommSpace>
class Handle<ExecSpace, CommSpace> {
};
}

Partial Specialization of Req

Warning

coming soon...

namespace KokkosComm {
template <CommunicationSpace CommSpace>
class Req<CommSpace> {
};
}

struct Partial Specialization for each API

The APIs are actually implemented by partial specializations of structs. Conceptually, there is an internal interface that needs to be satisfied for each API:

namespace Impl {
template <KokkosExecutionSpace ExecSpace, KokkosView RecvView, CommunicationSpace CS>
struct Recv<RecvView, ExecSpace, CS> { ... };
}

In the above, CS is a type that represents the CommunicationSpace implementation. For example, for the Mpi CommunicationSpace, we create a partial specialization of that struct template (note, fewer template parameters and the use of Mpi):

namespace Impl {
template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
struct Recv<RecvView, ExecSpace, Mpi> { ... };
}

You need to implement the following 3 things to get a new backend:

Recv Concept

An async message receive

namespace KokkosComm::Impl {
template <KokkosExecutionSpace ExecSpace, KokkosView RecvView, CommunicationSpace CommSpace>
struct Recv<RecvView, ExecSpace, CommSpace> {

  // implement this to actually do the Recv
  static Req<Mpi> execute(Handle<ExecSpace, Mpi> &h, const RecvView &rv, int src);
};
}}

Send Concept

An async message send

namespace KokkosComm::Impl {
template <KokkosExecutionSpace ExecSpace, KokkosView SendView, CommunicationSpace CommSpace>
struct Send<SendView, ExecSpace, CommSpace> {

  // implement this to actually do the Send
  static Req<Mpi> execute(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest);
};
}}

Barrier Concept

A global barrier

template <KokkosExecutionSpace ExecSpace, CommunicationSpace CommSpace>
struct Barrier {

  // implement this to actually do the barrier
  Barrier(Handle<ExecSpace, Mpi> &&h);
};
+++