Collectives#

using rapidsmpf::coll::ReduceOperator = std::function<void(Buffer const *left, Buffer *right)>#

Type alias for the reduction function signature.

A reduction function is a binary operator left \oplus right. The function implementing the operation must update right in place. That is, the result of calling the reduction should be as if we do right <- left \oplus right.

Note

Both buffers are guaranteed to be on the same stream when the function is called by AllReduce.

class AllGather#
#include <allgather.hpp>

AllGather communication service.

The class provides a communication service where each rank contributes data and all ranks receive all inputs on all ranks.

The implementation uses a ring broadcast. Each rank receives a contribution from its left neighbour, forwards the message to its right neighbour (unless at the end of the ring) and then stores the contribution locally. The cost on P ranks if each rank inserts a message of size N is

(P - 1) alpha + N ((P - 1) / P) beta

Per insertion. Where alpha is the network latency and beta the inverse bandwidth. Although the latency term is linear (rather than logarithmic as is the case for Bruck’s algorithm or recursive doubling) MPI implementations typically observe that for large messages ring algorithms perform better since message passing is only nearest neighbour.

Public Types

enum class Ordered : bool#

Tag requesting ordering for extraction.

Values:

enumerator NO#

Extraction is unordered.

enumerator YES#

Extraction is ordered.

Public Functions

void insert(std::uint64_t sequence_number, PackedData &&packed_data)#

Insert packed data into the allgather operation.

Parameters:
  • sequence_number – Local ordered sequence number of the data.

  • packed_data – The data to contribute to the allgather.

void insert_finished()#

Mark that this rank has finished contributing data.

std::vector<PackedData> wait_and_extract(
Ordered ordered = Ordered::YES,
std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}
)#

Wait for completion and extract all gathered data.

Blocks until the allgather operation completes and returns all collected data from all ranks.

Parameters:
  • ordered – If the extracted data should be ordered? if ordered, returned data will be ordered first by rank and then by insertion order on that rank.

  • timeout – Optional maximum duration to wait. Negative values mean no timeout.

Throws:

std::runtime_error – If the timeout is reached.

Returns:

A vector containing packed data from all participating ranks.

AllGather(
std::shared_ptr<Communicator> comm,
OpID op_id,
BufferResource *br,
std::function<void(void)> &&finished_callback = nullptr
)#

Construct a new allgather operation.

Note

It is safe to reuse the op_id as soon as wait_and_extract has completed locally.

Note

The caller promises that inserted buffers are stream-ordered with respect to their own stream, and extracted buffers are likewise guaranteed to be stream- ordered with respect to their own stream.

Parameters:
  • comm – The communicator for communication.

  • op_id – Unique operation identifier for this allgather.

  • brBuffer resource for memory allocation.

  • finished_callback – Optional callback run when partitions are locally finished. The callback is guaranteed to be called by the progress thread exactly once when the allgather is locally ready.

AllGather(AllGather const&) = delete#

Deleted copy constructor.

AllGather &operator=(AllGather const&) = delete#

Deleted copy assignment operator.

AllGather(AllGather&&) = delete#

Deleted move constructor.

AllGather &operator=(AllGather&&) = delete#

Deleted move assignment operator.

inline std::shared_ptr<Communicator> const &comm() const noexcept#

Gets the communicator associated with this AllGather.

Returns:

Shared pointer to communicator.

~AllGather() noexcept#

Destructor.

Note

This operation is logically collective. If an AllGather is locally destructed before waiting to extract, there is no guarantee that in-flight communication will be completed.

ProgressThread::ProgressState event_loop()#

Main event loop for processing allgather operations.

This method is called by the progress thread to handle ongoing communication and data transfers.

Returns:

The current progress state.

class AllReduce#
#include <allreduce.hpp>

AllReduce collective.

The implementation uses a butterfly recursive doubling scheme for message exchange, using no extra memory and O(log P) rounds for P ranks.

The actual reduction is implemented via a type-erased ReduceOperator that is supplied at construction time. Helper factories such as detail::make_host_reduce_operator or detail::make_device_reduce_operator can be used to build range-based reductions over contiguous arrays.

Note

No internal allocations are made. The memory types and sizes of the two provided buffers must match, and the provided reduction operator must be valid for the memory type of the buffers.

Note

The reduction is safe to use with both non-associative and non-commutative reduction operations in the sense that all participating ranks are guaranteed to receive the same answer even if the operator is not associative or commutative.

Note

It is safe to reuse the op_id passed to the AllReduce construction locally as soon as wait_and_extract is complete.

Warning

Behaviour of this object is undefined if it is destructed without first ensuring that wait_and_extract completes successfully.

Public Functions

AllReduce(
std::shared_ptr<Communicator> comm,
std::unique_ptr<Buffer> input,
std::unique_ptr<Buffer> output,
OpID op_id,
ReduceOperator reduce_operator,
std::function<void(void)> finished_callback = nullptr
)#

Construct a new AllReduce operation.

Note

It is safe to reuse the op_id as soon as wait_and_extract has completed locally.

Parameters:
  • comm – The communicator for communication.

  • input – Local data to contribute to the reduction.

  • output – Allocated buffer in which to place reduction result. Must be the same size and memory type as input. Overwritten with the reduction result (values already in the buffer are ignored).

  • op_id – Unique operation identifier for this allreduce.

  • reduce_operator – Type-erased reduction operator to use. See ReduceOperator.

  • finished_callback – Optional callback run once locally when the allreduce is finished and results are ready for extraction.

Throws:

std::invalid_argument – If the input and output buffers do not match appropriately (same size, same memory type).

inline std::shared_ptr<Communicator> const &comm() const noexcept#

Gets the communicator associated with this AllReduce.

Returns:

Shared pointer to communicator.

~AllReduce() noexcept#

Destructor.

Note

This operation is logically collective. If an AllReduce is locally destructed before wait_and_extract is called, there is no guarantee that in-flight communication will be completed.

bool finished() const noexcept#

Check if the allreduce operation has completed.

Returns:

True if all data and finish messages from all ranks have been received and locally reduced.

std::pair<std::unique_ptr<Buffer>, std::unique_ptr<Buffer>> wait_and_extract(
std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}
)#

Wait for completion and extract the reduced data.

Blocks until the allreduce operation completes and returns the globally reduced result.

This method is destructive and can only be called once. The first call extracts the buffers provided to the AllReduce constructor. Subsequent calls will throw std::runtime_error because the underlying data has already been consumed.

Note

The streams of the Buffers may change in an implementation-defined way while owned by the AllReduce object, if you need to launch new stream-ordered work on a Buffer you obtain from this function, you must obtain the correct stream from the Buffer itself.

Warning

There may be outstanding stream-ordered work reading from the first Buffer when this function returns (not tracked by the buffer’s Buffer::latest_write_event()). If you want to pass it to a non-stream-ordered API that writes to the buffer you must synchronise its stream first.

Parameters:

timeout – Optional maximum duration to wait. Negative values mean no timeout.

Throws:

std::runtime_error – If the timeout is reached or if this method is called more than once.

Returns:

A pair of the two Buffers passed to the constructor. The first Buffer contains an implementation-defined value, the second Buffer contains the final reduced result.

bool is_ready() const noexcept#

Check if reduced results are ready for extraction.

This returns true once the underlying allgather has completed and, if wait_and_extract has not yet been called, indicates that calling it would not block.

Returns:

True if the allreduce operation has completed and results are ready for extraction, false otherwise.

class SparseAlltoall#
#include <sparse_alltoall.hpp>

Sparse all-to-all collective over explicit source and destination peer sets.

Each rank may send zero or more messages to ranks listed in dsts and receives zero or more messages from ranks listed in srcs. Sender order is defined by the local order of calls to insert(dst, ...) for each destination rank.

This object is logically collective over the communicator and identified by op_id. Local extraction is only valid after wait() has completed.

Public Functions

SparseAlltoall(
std::shared_ptr<Communicator> comm,
OpID op_id,
BufferResource *br,
std::vector<Rank> srcs,
std::vector<Rank> dsts,
std::function<void()> &&finished_callback = nullptr
)#

Construct a sparse all-to-all collective instance.

Note

It is safe to reuse the op_id as soon as wait has completed locally or the finished_callback has been invoked.

Note

The caller promises that inserted buffers are stream-ordered with respect to their own stream, and extracted buffers are likewise guaranteed to be stream- ordered with respect to their own stream.

Note

Collectively the src and dst pairs of participating ranks must be consistent (not checked for). That is if rank-A advertises that rank-B is in its dst set, rank-B must advertise that rank-A is in its src set. If we ever need to relax this restriction we could have each rank advertise its send set and bootstrap the two-sided information using the non-blocking consensus algorithm of Hoefler, Siebert, and Lumsdaine, ACM SIGPLAN (2010), https://dl.acm.org/doi/10.1145/1837853.1693476.

Parameters:
  • commCommunicator for the collective.

  • op_id – Collective operation identifier.

  • brBuffer resource used for allocations.

  • srcs – Ranks this rank will receive from.

  • dsts – Ranks this rank will send to.

  • finished_callback – Optional callback invoked exactly once when the collective is locally complete. The callback should be fast and non-blocking. Ideally it should only be used to signal a thread to do the actual work of extraction. Note in particular that the callback should not extract any data.

Throws:
  • std::out_of_range – If either srcs or dsts have invalid values. All source and destination ranks must be in [0, ..., comm->nranks()), and not equal to the current rank.

  • std::invalid_argument – If the rank lists are not unique.

  • std::logic_error – If the communicator or buffer resource pointers are null.

std::shared_ptr<Communicator> const &comm() const noexcept#

Gets the communicator associated with this SparseAlltoall.

Returns:

Shared pointer to communicator.

void insert(Rank dst, PackedData &&packed_data)#

Insert data to send to a destination rank.

The order the destination rank obtains the sent data is given by the insertion order on the send side. If inserting concurrently to the same destination, the caller must establish a total order of the insertions, otherwise the reconstruction order on the receive side is unspecified.

Note

Concurrent insertion by multiple threads is supported.

Note

the caller must ensure that insert_finished() is called after all insert() calls have completed.

Parameters:
  • dst – Destination rank. Must be present in the constructor’s dsts.

  • packed_data – Packed payload and metadata to send.

void insert_finished()#

Indicate that no more data will be inserted for any destination.

Must be called exactly once.

Note

If multiple threads are insert()ing, you must establish a happens-before relationship between the completion of all insert()s and the final call to insert_finished().

void wait(
std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}
)#

Wait for local completion.

Parameters:

timeout – Optional timeout. Negative values mean no timeout.

Throws:

std::runtime_error – If the timeout is reached.

std::vector<PackedData> extract(Rank src)#

Extract all received messages from a source rank.

The returned vector is ordered by the sender’s local insertion order.

Note

Concurrent extraction is supported, behaviour is undefined if two threads attempt to extract data from the same source.

Parameters:

src – Source rank. Must be present in the constructor’s srcs.

Throws:

std::logic_error – If extracting before the collective is complete.

Returns:

All messages received from src.