allreduce.hpp

 
#pragma once
 
#include <memory>
 
#include <coro/event.hpp>
#include <coro/task.hpp>
 
#include <rapidsmpf/coll/allreduce.hpp>
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
 
namespace rapidsmpf::streaming {
 
class AllReduce {
  public:
    AllReduce(
        std::shared_ptr<Context> ctx,
        std::shared_ptr<Communicator> comm,
        std::unique_ptr<Buffer> input,
        std::unique_ptr<Buffer> output,
        OpID op_id,
        coll::ReduceOperator reduce_operator
    );
 
    AllReduce(AllReduce const&) = delete;
    AllReduce& operator=(AllReduce const&) = delete;
    AllReduce(AllReduce&&) = delete;
    AllReduce& operator=(AllReduce&&) = delete;
 
    ~AllReduce() noexcept;
 
    [[nodiscard]] std::shared_ptr<Context> const& ctx() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Communicator> const& comm() const noexcept;
 
    coro::task<std::pair<std::unique_ptr<Buffer>, std::unique_ptr<Buffer>>> extract();
 
  private:
    coro::event
        event_{};  
    std::shared_ptr<Context> ctx_;  
    coll::AllReduce reducer_;  
};
 
}  // namespace rapidsmpf::streaming