Core

Point-to-point

Send

Warning

This is not a blocking operation despite being named like MPI_Send.

template<KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, CommunicationSpace CommSpace = DefaultCommunicationSpace>
auto send(Handle<ExecSpace, CommSpace> &h, SendView &sv, int dest) -> Request<CommSpace>

Initiates a non-blocking send operation.

Template Parameters:

• SendView – The type of the Kokkos view to send.

• ExecSpace – The execution space to use. Defaults to Kokkos::DefaultExecutionSpace.

• CommSpace – The communication backend to use. Defaults to DefaultCommunicationSpace.

Parameters:

• h – A handle to the execution space and transport mechanism.

• sv – The Kokkos view to send.

• dest – The destination rank.

Returns:

A request object of type Request<CommSpace> representing the non-blocking send operation.

template<KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, CommunicationSpace CommSpace = DefaultCommunicationSpace>
auto send(SendView &sv, int dest) -> Request<CommSpace>

Initiates a non-blocking send operation using a default handle.

Template Parameters:

• SendView – The type of the Kokkos view to send.

• ExecSpace – The execution space to use. Defaults to Kokkos::DefaultExecutionSpace.

• CommSpace – The communication backend to use. Defaults to DefaultCommunicationSpace.

Parameters:

• sv – The Kokkos view to send.

• dest – The destination rank.

Returns:

A request object of type Request<CommSpace> representing the non-blocking send operation.

Example usage:

#include "KokkosComm/KokkosComm.hpp"

// Define the execution space and transport
using ExecSpace = Kokkos::DefaultExecutionSpace;
using CommSpace = DefaultCommunicationSpace;

// Create a Kokkos view
Kokkos::View<double*> data("data", 100);

// Fill the view with some data
Kokkos::parallel_for("fill_data", Kokkos::RangePolicy<ExecSpace>(0, 100), KOKKOS_LAMBDA(int i) {
  data(i) = static_cast<double>(i);
});

// Destination rank
int dest = 1;

// Create a handle
KokkosComm::Handle<> handle; // Same as Handle<Execspace, CommSpace>

// Initiate a non-blocking send with a handle
auto req1 = send(handle, data, dest);

// Initiate a non-blocking send with a default handle
auto req2 = send(data, dest);

// Wait for the requests to complete (assuming a wait function exists)
KokkosComm::wait(req1);
KokkosComm::wait(req2);

Receive

Warning

This is not a blocking operation despite being named like MPI_Recv.

template<KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, CommunicationSpace CommSpace = DefaultCommunicationSpace>
auto recv(Handle<ExecSpace, CommSpace> &h, RecvView &sv, int dest) -> Request<CommSpace>

Initiates a non-blocking receive operation.

Template Parameters:

• RecvView – The type of the Kokkos view for receiving data.

• ExecSpace – The execution space where the operation will be performed. Defaults to Kokkos::DefaultExecutionSpace.

• CommSpace – The communication backend to use. Defaults to DefaultCommunicationSpace.

Parameters:

• h – A handle to the execution space and transport mechanism.

• rv – The Kokkos view where the received data will be stored.

• src – The source rank from which to receive data.

Returns:

A request object of type Request<CommSpace> representing the non-blocking receive operation.

This function initiates a non-blocking receive operation using the specified execution space and transport mechanism. The data will be received into the provided view from the specified source rank and message tag. The function returns a request object that can be used to check the status of the receive operation or to wait for its completion.

template<KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, CommunicationSpace CommSpace = DefaultCommunicationSpace>
auto recv(RecvView &sv, int dest) -> Request<CommSpace>

Initiates a non-blocking receive operation using a default handle.

Template Parameters:

• RecvView – The type of the Kokkos view for receiving data.

• ExecSpace – The execution space where the operation will be performed. Defaults to Kokkos::DefaultExecutionSpace.

• CommSpace – The communication backend to use. Defaults to DefaultCommunicationSpace.

Parameters:

• rv – The Kokkos view where the received data will be stored.

• src – The source rank from which to receive data.

Returns:

A request object of type Request<CommSpace> representing the non-blocking receive operation.

Example usage:

#include "KokkosComm/KokkosComm.hpp"

// Define the execution space and transport
using ExecSpace = Kokkos::DefaultExecutionSpace;
using CommSpace = DefaultCommunicationSpace;

// Source rank
int src = 1;

// Create a handle
KokkosComm::Handle<> handle; // Same as Handle<Execspace, CommSpace>

// Allocate a view to receive the data
Kokkos::View<double*> data("recv_view", 100);

// Initiate a non-blocking receive with a handle
auto req1 = recv(handle, data, src);

// Initiate a non-blocking receive with a default handle
auto req2 = recv(data, src);

// Wait for the requests to complete (assuming a wait function exists)
KokkosComm::wait(req1);
KokkosComm::wait(req2);

Collectives

Important

Collective operations act element-wise on the input Views. Multi-dimensional Views are treated as a logically flattened sequence of values, and the reduction is applied over that sequence. All participating Views must have identical extents; mismatched shapes result in undefined behavior.

The reduction operator must be associative, but ordering of partial combinations is not guaranteed, and the operation is not required to be commutative.

template<KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, CommunicationSpace CommSpace = DefaultCommunicationSpace>
auto barrier(Handle<ExecSpace, CommSpace> &&h) -> void

A function to create a barrier using the given execution space and transport handle.

Template Parameters:

• ExecSpace – The execution space where the operation will be performed. Defaults to Kokkos::DefaultExecutionSpace.

• CommSpace – The communication backend to use. Defaults to DefaultCommunicationSpace.

Parameters:

h – A handle of type Handle<ExecSpace, CommSpace> to be forwarded to the barrier implementation.

Related types

template<CommunicationSpace C = DefaultCommSpace>
class Request

Template class for request wrappers of different communication space types.

Request objects are move-only: copy construction and copy assignment are explicitly deleted. There is always exactly one owner of a Request and its associated callbacks. This design ensures it is impossible for the same callback to be executed more than once.

Template Parameters:

C – The communication backend to use. Defaults to DefaultCommunicationSpace.

using communication_space = C

using request_type = C::request_type

using rank_type = C::rank_type

Common interfaces

Both specializations share the following interface:

Request(const Request&) = delete

Copy constructor is deleted because a Request can only be moved.

auto operator=(const Request&) -> Request& = delete

Copy assignment operator is deleted because a Request can only be moved.

Request(Request&&) = default

Move constructor is defaulted.

auto operator=(Request&&) -> Request& = default

Move assignment operator is defaulted.

template<KokkosView V>
auto extend_view_lifetime(const V &view) -> void

Captures a Kokkos View to extend its lifetime until the request completes. Has no effect on unmanaged Views.

Template Parameters:

V – A Kokkos View type.

Parameters:

view – The view whose lifetime should be extended.

auto add_callback(std::function<void()> &&cb) -> void

Registers a callback to be invoked after the request completes.

Parameters:

cb – The callback function to register.

auto wait() -> void

Blocks until the associated operation completes. Executes all registered callbacks upon completion.

auto test() -> bool

Non-blocking query for the completion of the associated operation. Executes all registered callbacks if the operation has completed.

Returns:

true if the request has completed, false otherwise.

auto wait(Request &request) -> void

Free function overload. Waits on request until the associated operation completes.

Parameters:

request – A reference to the request to wait on.

auto wait(Request &&request) -> void

Free function overload. Waits on an r-value request, consuming it upon completion.

Parameters:

request – An r-value reference to the request to wait on.

auto wait_all(std::span<Request> requests) -> void

Waits for the completion of all requests in requests.

Parameters:

requests – The list of requests to complete.

auto wait_any(std::span<Request> requests) -> std::optional<rank_type>

Waits for the completion of at least one request in requests.

Parameters:

requests – The list of requests to poll.

Returns:

The index of the completed request, or std::nullopt if requests is empty.

auto test(Request &request) -> bool

Free function overload. Queries request for completion of the associated operation.

Parameters:

request – A reference to the request to query.

Returns:

true if the request has completed, false otherwise.

MPI specialization

template<>
class Request<MpiSpace>

Request specialization for the MpiSpace communication space. Wraps an MPI_Request handle.

using communication_space = MpiSpace

using request_type = MpiSpace::request_type

using rank_type = MpiSpace::rank_type

explicit Request(request_type request = MPI_REQUEST_NULL)

Constructs a Request from an MPI_Request handle.

Parameters:

request – The MPI_Request to encapsulate. Defaults to MPI_REQUEST_NULL.

auto request() noexcept -> request_type&

auto request() const noexcept -> const request_type&

Returns:

A reference to the underlying MPI_Request object.

auto request_ptr() noexcept -> request_type*

auto request_ptr() const noexcept -> const request_type*

Returns:

A pointer to the underlying MPI_Request object.

Note

Both wait_all and wait_any copy the underlying MPI_Request objects into an intermediate container before calling MPI_Waitall and MPI_Waitany, respectively, which incurs an allocation overhead.

NCCL specialization

template<>
class Request<Experimental::NcclSpace>

Request specialization for the Experimental::NcclSpace communication space. Wraps a cudaEvent_t handle to track the completion of CUDA stream operations.

using communication_space = Experimental::NcclSpace

using request_type = Experimental::NcclSpace::request_type

using rank_type = Experimental::NcclSpace::rank_type

explicit Request()

Constructs an empty Request with a null event handle.

~Request() noexcept

Destructor. Destroys the underlying cudaEvent_t if one has been created.

auto capture_stream_state(cudaStream_t stream) noexcept -> void

Records a CUDA event on stream to capture its current state for completion tracking. If a cudaEvent_t was previously created on this request, it is destroyed first.

Parameters:

stream – The CUDA stream whose state to capture.

auto request() noexcept -> request_type&

auto request() const noexcept -> const request_type&

Returns:

A reference to the underlying cudaEvent_t object.

auto request_ptr() noexcept -> request_type*

auto request_ptr() const noexcept -> const request_type*

Returns:

A pointer to the underlying cudaEvent_t object.

Note

Both wait_all and wait_any use active polling loops rather than blocking synchronization. While this increases CPU utilization, it avoids the overhead of spawning threads or completing requests sequentially.

Utility

Warning

Non-system data types (i.e. the data types not natively supported by the communication space) are not convertible. This notably includes user-defined types.

template<CommunicationSpace C, typename T>
auto datatype() -> C::datatype_type

Converts a type T to its communication space C equivalent representation.

When C is:

• MpiSpace, returns the corresponding MPI_Datatype type.

• NcclSpace, returns the corresponding ncclDataType_t type.

Template Parameters:

• C – The target communication space backend to use for data type conversion.

• T – The C++-native data type to convert from.

Returns:

The communication space representation of the C++-native data type.

template<CommunicationSpace C, KokkosView V>
auto datatype_for([[maybe_unused]] const V &view) -> C::datatype_type

Template Parameters:

• C – The target communication space backend to use for data type conversion.

• V – A Kokkos View type.

Parameters:

view – The Kokkos View to convert the value type from.

Returns:

The communication space representation of the Kokkos View value type.

template<CommunicationSpace C, KokkosView V>
auto datatype_for([[maybe_unused]] C &&comm, [[maybe_unused]] const V &view) -> C::datatype_type

Template Parameters:

• C – The target communication space backend to use for data type conversion.

• V – A Kokkos View type.

Parameters:

• comm – A communication space object, immediately consumed.

• view – The Kokkos View to convert the value type from.

Returns:

The communication space representation of the Kokkos View value type.