allreduce.hpp

 
#pragma once
 
#include <algorithm>
#include <atomic>
#include <chrono>
#include <concepts>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <span>
#include <utility>
 
#ifdef __CUDACC__
#include <thrust/execution_policy.h>
#include <thrust/transform.h>
#endif
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/progress_thread.hpp>
#include <rapidsmpf/statistics.hpp>
 
namespace rapidsmpf::coll {
 
using ReduceOperator = std::function<void(Buffer const* left, Buffer* right)>;
 
class AllReduce {
  public:
    AllReduce(
        std::shared_ptr<Communicator> comm,
        std::unique_ptr<Buffer> input,
        std::unique_ptr<Buffer> output,
        OpID op_id,
        ReduceOperator reduce_operator,
        std::function<void(void)> finished_callback = nullptr
    );
 
    [[nodiscard]] std::shared_ptr<Communicator> const& comm() const noexcept {
        return comm_;
    }
 
    AllReduce(AllReduce const&) = delete;
    AllReduce& operator=(AllReduce const&) = delete;
    AllReduce(AllReduce&&) = delete;
    AllReduce& operator=(AllReduce&&) = delete;
 
    ~AllReduce() noexcept;
 
    [[nodiscard]] bool finished() const noexcept;
 
    [[nodiscard]] std::pair<std::unique_ptr<Buffer>, std::unique_ptr<Buffer>>
    wait_and_extract(std::chrono::milliseconds timeout = std::chrono::milliseconds{-1});
 
    [[nodiscard]] bool is_ready() const noexcept;
 
  private:
    enum class Phase : uint8_t {
        StartPreRemainder,
        CompletePreRemainder,
        StartButterfly,
        CompleteButterfly,
        StartPostRemainder,
        CompletePostRemainder,
        Done,
        ResultAvailable
    };
 
    [[nodiscard]] ProgressThread::ProgressState event_loop();
 
    std::shared_ptr<Communicator> comm_{};
    ReduceOperator reduce_operator_;  
    std::unique_ptr<Buffer> in_buffer_{};
    std::unique_ptr<Buffer> out_buffer_{};
    OpID op_id_{};
    std::atomic<Phase> phase_{Phase::StartPreRemainder};
    std::atomic<bool> active_{true};
    std::function<void()>
        finished_callback_;  
 
    mutable std::mutex mutex_;  
    mutable std::condition_variable cv_;  
 
    Rank logical_rank_{-1};
    Rank nearest_pow2_{0};
    Rank non_pow2_remainder_{0};
    Rank stage_mask_{1};
    Rank stage_partner_{-1};
 
    ProgressThread::FunctionID function_id_{};  
 
    std::unique_ptr<Communicator::Future> send_future_{};
    std::unique_ptr<Communicator::Future> recv_future_{};
};
 
namespace detail {
 
template <typename T, typename Op>
struct HostOp {
    Op op;  
 
    void operator()(Buffer const* left, Buffer* right) {
        auto const left_nbytes = left->size;
        RAPIDSMPF_EXPECTS(
            left_nbytes % sizeof(T) == 0,
            "HostOp buffer size must be a multiple of sizeof(T)"
        );
 
        auto const count = left_nbytes / sizeof(T);
        if (count == 0) {
            return;
        }
 
        RAPIDSMPF_EXPECTS(
            left->mem_type() == MemoryType::HOST && right->mem_type() == MemoryType::HOST,
            "HostOp expects host memory"
        );
 
        auto* left_bytes = left->data();
        auto* right_bytes = right->exclusive_data_access();
 
        std::span<T const> left_span{reinterpret_cast<T const*>(left_bytes), count};
        std::span<T> right_span{reinterpret_cast<T*>(right_bytes), count};
 
        std::ranges::transform(left_span, right_span, right_span.begin(), op);
        right->unlock();
    }
};
 
template <typename T, typename Op>
struct DeviceOp {
    Op op;  
 
    void operator()(Buffer const* left, Buffer* right) {
#ifdef __CUDACC__
        auto const left_nbytes = left->size;
        RAPIDSMPF_EXPECTS(
            left_nbytes % sizeof(T) == 0,
            "DeviceOp buffer size must be a multiple of sizeof(T)"
        );
 
        auto const count = left_nbytes / sizeof(T);
        if (count == 0) {
            return;
        }
 
        RAPIDSMPF_EXPECTS(
            left->mem_type() == MemoryType::DEVICE
                && right->mem_type() == MemoryType::DEVICE,
            "DeviceOp expects device memory"
        );
        // Both buffers are guaranteed to be on the same stream by the AllReduce ctor.
        right->write_access([&](std::byte* right_bytes, rmm::cuda_stream_view stream) {
            auto const* left_bytes = reinterpret_cast<std::byte const*>(left->data());
 
            T* right_ptr = reinterpret_cast<T*>(right_bytes);
            T const* left_ptr = reinterpret_cast<T const*>(left_bytes);
 
            thrust::transform(
                thrust::cuda::par_nosync.on(stream.value()),
                left_ptr,
                left_ptr + count,
                right_ptr,
                right_ptr,
                op
            );
        });
#else
        // This should never be reached if DeviceOp is only instantiated with CUDA
        std::ignore = left;
        std::ignore = right;
        RAPIDSMPF_FAIL(
            "DeviceOp::operator() called but CUDA compilation (__CUDACC__) "
            "was not available. DeviceOp requires CUDA/thrust support.",
            std::runtime_error
        );
#endif
    }
};
 
template <typename T, typename Op>
    requires std::invocable<Op, T const&, T const&>
ReduceOperator make_host_reduce_operator(Op op) {
    return HostOp<T, Op>{std::move(op)};
}
 
template <typename T, typename Op>
    requires std::invocable<Op, T const&, T const&>
ReduceOperator make_device_reduce_operator(Op op) {
#ifdef __CUDACC__
    return DeviceOp<T, Op>{std::move(op)};
#else
    std::ignore = op;
 
    RAPIDSMPF_FAIL(
        "make_device_reduce_operator was called from code that was not compiled "
        "with NVCC (__CUDACC__ is not defined).",
        std::runtime_error
    );
#endif
}
 
}  // namespace detail
 
}  // namespace rapidsmpf::coll