allgather.hpp

 
#pragma once
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <limits>
#include <memory>
#include <mutex>
#include <optional>
#include <vector>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/error.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>
#include <rapidsmpf/statistics.hpp>
 
namespace rapidsmpf::coll {
namespace detail {
 
using ChunkID = std::uint64_t;
 
class Chunk {
  private:
    ChunkID id_;  
    std::unique_ptr<std::vector<std::uint8_t>> metadata_;  
    std::unique_ptr<Buffer> data_;  
    std::uint64_t
        data_size_;  
 
    Chunk(
        ChunkID id,
        std::unique_ptr<std::vector<std::uint8_t>> metadata,
        std::unique_ptr<Buffer> data
    );
 
    Chunk(ChunkID id);
 
  public:
    [[nodiscard]] bool is_ready() const noexcept;
 
    [[nodiscard]] MemoryType memory_type() const noexcept;
 
    [[nodiscard]] bool is_finish() const noexcept;
 
    [[nodiscard]] ChunkID id() const noexcept;
 
    [[nodiscard]] ChunkID sequence() const noexcept;
 
    [[nodiscard]] Rank origin() const noexcept;
 
    [[nodiscard]] std::uint64_t data_size() const noexcept;
 
    [[nodiscard]] std::uint64_t metadata_size() const noexcept;
 
    [[nodiscard]] static std::unique_ptr<Chunk> from_packed_data(
        std::uint64_t sequence, Rank origin, PackedData&& packed_data
    );
 
    [[nodiscard]] static std::unique_ptr<Chunk> from_empty(
        std::uint64_t num_local_insertions, Rank origin
    );
 
    [[nodiscard]] PackedData release();
 
    static constexpr std::uint64_t ID_BITS = 38;
    static constexpr std::uint64_t RANK_BITS =
        sizeof(ChunkID) * std::numeric_limits<unsigned char>::digits - ID_BITS;
 
    static constexpr ChunkID chunk_id(std::uint64_t sequence, Rank origin);
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> serialize() const;
 
    [[nodiscard]] static std::unique_ptr<Chunk> deserialize(
        std::vector<std::uint8_t>& data, BufferResource* br
    );
 
    [[nodiscard]] std::unique_ptr<Buffer> release_data_buffer() noexcept;
 
    void attach_data_buffer(std::unique_ptr<Buffer> data);
 
    ~Chunk() = default;
    Chunk(Chunk&&) = default;
    Chunk& operator=(Chunk&&) = default;
    Chunk(Chunk const&) = delete;
    Chunk& operator=(Chunk const&) = delete;
};
 
class PostBox {
  public:
    PostBox() = default;
    ~PostBox() = default;
    PostBox(PostBox const&) = delete;
    PostBox& operator=(PostBox const&) = delete;
    PostBox(PostBox&&) = delete;
    PostBox& operator=(PostBox&&) = delete;
 
    void insert(std::unique_ptr<Chunk> chunk);
 
    void insert(std::vector<std::unique_ptr<Chunk>>&& chunks);
 
    void increment_goalpost(std::uint64_t amount);
 
    [[nodiscard]] bool ready() const noexcept;
 
    [[nodiscard]] std::vector<std::unique_ptr<Chunk>> extract_ready();
 
    [[nodiscard]] std::vector<std::unique_ptr<Chunk>> extract();
 
    [[nodiscard]] bool empty() const noexcept;
 
    [[nodiscard]] std::size_t spill(BufferResource* br, std::size_t amount);
 
  private:
    mutable std::mutex mutex_{};  
    std::vector<std::unique_ptr<Chunk>> chunks_{};  
    std::atomic<std::uint64_t> goalpost_{0};  
};
 
}  // namespace detail
 
class AllGather {
  public:
    void insert(std::uint64_t sequence_number, PackedData&& packed_data);
 
    void insert_finished();
 
    [[nodiscard]] bool finished() const noexcept;
 
    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}
    );
 
    [[nodiscard]] std::vector<PackedData> extract_ready();
 
    AllGather(
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        BufferResource* br,
        std::shared_ptr<Statistics> statistics = Statistics::disabled(),
        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::shared_ptr<Statistics> statistics_;  
    std::function<void(void)> finished_callback_{
        nullptr
    };  
    std::atomic<Rank> finish_counter_;  
    std::atomic<std::uint32_t> nlocal_insertions_;  
    OpID op_id_;  
    std::atomic<bool> locally_finished_{false};  
    std::atomic<bool> active_{true};  
    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