allgather.hpp

 
#pragma once
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <vector>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <rapidsmpf/coll/utils.hpp>
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/packed_data.hpp>
#include <rapidsmpf/memory/spill_manager.hpp>
#include <rapidsmpf/progress_thread.hpp>
 
namespace rapidsmpf::coll {
 
class AllGather {
  public:
    void insert(std::uint64_t sequence_number, PackedData&& packed_data);
 
    void insert_finished();
 
    enum class Ordered : bool {
        NO,  
        YES,  
    };
 
    [[nodiscard]] std::vector<PackedData> wait_and_extract(
        Ordered ordered = Ordered::YES,
        std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}
    );
 
    AllGather(
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        BufferResource* br,
        std::function<void(void)>&& finished_callback = nullptr
    );
 
    AllGather(AllGather const&) = delete;
    AllGather& operator=(AllGather const&) = delete;
    AllGather(AllGather&&) = delete;
    AllGather& operator=(AllGather&&) = delete;
 
    [[nodiscard]] std::shared_ptr<Communicator> const& comm() const noexcept {
        return comm_;
    }
 
    ~AllGather() noexcept;
 
    ProgressThread::ProgressState event_loop();
 
  private:
    void insert(std::unique_ptr<detail::Chunk> chunk);
 
    void mark_finish(std::uint64_t expected_chunks) noexcept;
 
    void wait(std::chrono::milliseconds timeout = std::chrono::milliseconds{-1});
 
    std::size_t spill(std::optional<std::size_t> amount = std::nullopt);
 
    std::shared_ptr<Communicator> comm_;  
    BufferResource* br_;  
    std::function<void(void)> finished_callback_{
        nullptr
    };  
    std::atomic<Rank> finish_counter_;  
    std::atomic<std::uint32_t> nlocal_insertions_;  
    std::atomic<std::uint64_t> extraction_goalpost_{
        0
    };  
    OpID op_id_;  
    std::atomic<bool> locally_finished_{false};  
    bool can_extract_{false};  
    mutable std::mutex mutex_;  
    std::condition_variable cv_;  
    detail::PostBox inserted_{};  
    detail::PostBox for_extraction_{};  
    ProgressThread::FunctionID function_id_{};  
    SpillManager::SpillFunctionID spill_function_id_{};  
    // We track remote finishes separately from the finish_counter_ above since the path
    // through the event loop state machine for a local finish marker is slightly
    // different from a remote finish marker.
    Rank remote_finish_counter_;
    std::uint64_t num_expected_messages_{0};
    std::uint64_t num_received_messages_{0};
    std::vector<std::unique_ptr<detail::Chunk>> to_receive_{};
    std::vector<std::unique_ptr<Communicator::Future>> fire_and_forget_{};
    std::vector<std::unique_ptr<detail::Chunk>> sent_posted_{};
    std::vector<std::unique_ptr<Communicator::Future>> sent_futures_{};
    std::vector<std::unique_ptr<detail::Chunk>> receive_posted_{};
    std::vector<std::unique_ptr<Communicator::Future>> receive_futures_{};
};
 
}  // namespace rapidsmpf::coll