allgather.hpp

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