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 updaterightin place. That is, the result of calling the reduction should be as if we doright <- 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
Pranks if each rank inserts a message of sizeNis(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
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( )#
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.
- 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_idas soon aswait_and_extracthas 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.
br – Buffer 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.
-
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
AllGatheris locally destructed beforewaiting 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.
-
void insert(std::uint64_t sequence_number, PackedData &&packed_data)#
-
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 forPranks.The actual reduction is implemented via a type-erased
ReduceOperatorthat is supplied at construction time. Helper factories such asdetail::make_host_reduce_operatorordetail::make_device_reduce_operatorcan 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_idpassed to theAllReduceconstruction locally as soon aswait_and_extractis complete.Warning
Behaviour of this object is undefined if it is destructed without first ensuring that
wait_and_extractcompletes successfully.Public Functions
- 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_idas soon aswait_and_extracthas 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
AllReduceis locally destructed beforewait_and_extractis 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
AllReduceconstructor. 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
AllReduceobject, if you need to launch new stream-ordered work on aBufferyou obtain from this function, you must obtain the correct stream from theBufferitself.Warning
There may be outstanding stream-ordered work reading from the first
Bufferwhen this function returns (not tracked by the buffer’sBuffer::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 firstBuffercontains an implementation-defined value, the secondBuffercontains 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_extracthas 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
dstsand receives zero or more messages from ranks listed insrcs. Sender order is defined by the local order of calls toinsert(dst, ...)for each destination rank.This object is logically collective over the communicator and identified by
op_id. Local extraction is only valid afterwait()has completed.Public Functions
- 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_idas soon aswaithas completed locally or thefinished_callbackhas 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:
comm – Communicator for the collective.
op_id – Collective operation identifier.
br – Buffer 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
srcsordstshave 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 allinsert()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 allinsert()s and the final call toinsert_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.