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 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.

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

struct Handle

Warning

coming soon...

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

struct 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);
};
+++