allgather.hpp

 
#pragma once
 
#include <cstdint>
#include <memory>
#include <vector>
 
#include <rapidsmpf/allgather/allgather.hpp>
#include <rapidsmpf/buffer/packed_data.hpp>
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/streaming/chunks/packed_data.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
 
#include <coro/event.hpp>
#include <coro/task.hpp>
 
namespace rapidsmpf::streaming {
 
class AllGather {
  public:
    using Ordered = rapidsmpf::allgather::AllGather::Ordered;
    AllGather(std::shared_ptr<Context> ctx, OpID op_id);
 
    AllGather(AllGather const&) = delete;
    AllGather& operator=(AllGather const&) = delete;
    AllGather(AllGather&&) = delete;
    AllGather& operator=(AllGather&&) = delete;
 
    ~AllGather();
 
    [[nodiscard]] std::shared_ptr<Context> ctx() const noexcept;
 
    void insert(std::uint64_t sequence_number, PackedDataChunk&& chunk);
 
    void insert_finished();
 
    coro::task<std::vector<PackedDataChunk>> extract_all(Ordered ordered = Ordered::YES);
 
  private:
    coro::event
        event_{};  
    std::shared_ptr<Context> ctx_;  
    allgather::AllGather gatherer_;  
};
 
namespace node {
 
Node allgather(
    std::shared_ptr<Context> ctx,
    std::shared_ptr<Channel> ch_in,
    std::shared_ptr<Channel> ch_out,
    OpID op_id,
    AllGather::Ordered ordered = AllGather::Ordered::YES
);
}  // namespace node
}  // namespace rapidsmpf::streaming