-
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 set of structs (or other type) that represents your communication space and associated types (handles, requests), as well as partial specializations of different structs for each 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);
};
+++