-
Notifications
You must be signed in to change notification settings - Fork 10
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.
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
Warning
coming soon...
namespace KokkosComm {
template <KokkosExecutionSpace ExecSpace, CommunicationSpace CommSpace>
class Handle<ExecSpace, CommSpace> {
};
}
Warning
coming soon...
namespace KokkosComm {
template <CommunicationSpace CommSpace>
class Req<CommSpace> {
};
}
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:
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);
};
}}
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);
};
}}
A global barrier
template <KokkosExecutionSpace ExecSpace, CommunicationSpace CommSpace>
struct Barrier {
// implement this to actually do the barrier
Barrier(Handle<ExecSpace, Mpi> &&h);
};
+++