sparse_alltoall.hpp

 
#pragma once
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/coll/utils.hpp>
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/packed_data.hpp>
#include <rapidsmpf/progress_thread.hpp>
 
namespace rapidsmpf::coll {
 
class SparseAlltoall {
  public:
    SparseAlltoall(
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        BufferResource* br,
        std::vector<Rank> srcs,
        std::vector<Rank> dsts,
        std::function<void()>&& finished_callback = nullptr
    );
 
    ~SparseAlltoall() noexcept;
 
    SparseAlltoall(SparseAlltoall const&) = delete;
    SparseAlltoall& operator=(SparseAlltoall const&) = delete;
    SparseAlltoall(SparseAlltoall&&) = delete;
    SparseAlltoall& operator=(SparseAlltoall&&) = delete;
 
    [[nodiscard]] std::shared_ptr<Communicator> const& comm() const noexcept;
 
    void insert(Rank dst, PackedData&& packed_data);
 
    void insert_finished();
 
    void wait(std::chrono::milliseconds timeout = std::chrono::milliseconds{-1});
 
    [[nodiscard]] std::vector<PackedData> extract(Rank src);
 
  private:
    struct SourceState {
        std::uint64_t expected_count{0};
        std::uint64_t received_count{0};
        std::vector<std::unique_ptr<detail::Chunk>> chunks{};
        std::vector<std::unique_ptr<detail::Chunk>> incoming{};
 
        [[nodiscard]] bool ready() const noexcept {
            return expected_count > 0 && expected_count == received_count;
        }
    };
 
    void send_ready_messages();
    void receive_metadata_messages();
    void receive_data_messages();
    void complete_data_messages();
    [[nodiscard]] bool containers_empty() const;
    [[nodiscard]] ProgressThread::ProgressState event_loop();
 
    std::shared_ptr<Communicator> comm_;
    BufferResource* br_;
    std::vector<Rank> srcs_;
    std::vector<Rank> dsts_;
    std::unordered_map<Rank, std::atomic<std::uint64_t>> next_ordinal_per_dst_;
 
    mutable std::mutex mutex_;
    std::condition_variable cv_;
    OpID op_id_;
    std::atomic<bool> locally_finished_{false};
    bool can_extract_{false};
 
    detail::PostBox outgoing_{};
    std::vector<std::unique_ptr<detail::Chunk>> receive_posted_;
    std::vector<std::unique_ptr<Communicator::Future>> receive_futures_;
    std::vector<std::unique_ptr<Communicator::Future>> fire_and_forget_;
    std::unordered_map<Rank, SourceState> source_states_;
    std::function<void()> finished_callback_;
    ProgressThread::FunctionID function_id_;
};
 
}  // namespace rapidsmpf::coll