sparse_alltoall.hpp

 
#pragma once
 
#include <memory>
#include <vector>
 
#include <coro/event.hpp>
#include <coro/task.hpp>
 
#include <rapidsmpf/coll/sparse_alltoall.hpp>
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/memory/packed_data.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
 
namespace rapidsmpf::streaming {
 
class SparseAlltoall {
  public:
    SparseAlltoall(
        std::shared_ptr<Context> ctx,
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        std::vector<Rank> srcs,
        std::vector<Rank> dsts
    );
 
    SparseAlltoall(SparseAlltoall const&) = delete;
    SparseAlltoall& operator=(SparseAlltoall const&) = delete;
    SparseAlltoall(SparseAlltoall&&) = delete;
    SparseAlltoall& operator=(SparseAlltoall&&) = delete;
 
    ~SparseAlltoall() noexcept;
 
    [[nodiscard]] std::shared_ptr<Context> const& ctx() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Communicator> const& comm() const noexcept;
 
    void insert(Rank dst, PackedData&& packed_data);
 
    [[nodiscard]] coro::task<void> insert_finished();
 
    [[nodiscard]] std::vector<PackedData> extract(Rank src);
 
  private:
    coro::event
        event_{};  
    std::shared_ptr<Context> ctx_;  
    coll::SparseAlltoall exchange_;  
};
 
}  // namespace rapidsmpf::streaming