communicator.hpp

 
#pragma once
 
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <ostream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include <rapidsmpf/communicator/logger.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/progress_thread.hpp>
 
namespace rapidsmpf {
 
using Rank = std::int32_t;
 
using OpID = std::int32_t;
 
using StageID = std::int32_t;
 
class Tag {
  public:
    using StorageT = std::int32_t;
 
    static constexpr int stage_id_bits{3};
 
    static constexpr StorageT stage_id_mask{(1 << stage_id_bits) - 1};
 
    static constexpr int op_id_bits{20};
 
    static constexpr StorageT op_id_mask{
        ((1 << (op_id_bits + stage_id_bits)) - 1) ^ stage_id_mask
    };
 
    constexpr Tag(OpID const op, StageID const stage)
        : tag_{
              (static_cast<StorageT>(op) << stage_id_bits) | static_cast<StorageT>(stage)
          } {
        RAPIDSMPF_EXPECTS(
            stage >= 0 && stage < (1 << stage_id_bits),
            "Invalid stage value",
            std::overflow_error
        );
        RAPIDSMPF_EXPECTS(
            op >= 0 && op < (1 << op_id_bits), "Invalid OpID value", std::overflow_error
        );
    }
 
    [[nodiscard]] static constexpr std::size_t bit_length() noexcept {
        return op_id_bits + stage_id_bits;
    }
 
    [[nodiscard]] static constexpr StorageT max_value() noexcept {
        return (1 << bit_length()) - 1;
    }
 
    constexpr operator StorageT() const noexcept {
        return tag_;
    }
 
    [[nodiscard]] constexpr OpID op() const noexcept {
        return (tag_ & op_id_mask) >> stage_id_bits;
    }
 
    [[nodiscard]] constexpr StageID stage() const noexcept {
        return tag_ & stage_id_mask;
    }
 
  private:
    StorageT const tag_;
};
 
class Communicator {
  public:
    class Future {
      public:
        Future() = default;
        virtual ~Future() noexcept = default;
        Future(Future&&) = default;  
        Future& operator=(Future&&) = default;
        Future(Future const&) = delete;  
        Future& operator=(Future const&) = delete;  
    };
 
  protected:
    Communicator() = default;
 
  public:
    virtual ~Communicator() noexcept = default;
 
    [[nodiscard]] virtual Rank rank() const = 0;
 
    [[nodiscard]] virtual Rank nranks() const = 0;
 
    [[nodiscard]] virtual std::unique_ptr<Future> send(
        std::unique_ptr<std::vector<std::uint8_t>> msg, Rank rank, Tag tag
    ) = 0;
 
    [[nodiscard]] virtual std::unique_ptr<Future> send(
        std::unique_ptr<Buffer> msg, Rank rank, Tag tag
    ) = 0;
 
    [[nodiscard]] virtual std::unique_ptr<Future> recv(
        Rank rank, Tag tag, std::unique_ptr<Buffer> recv_buffer
    ) = 0;
 
    [[nodiscard]] virtual std::unique_ptr<Future> recv_sync_host_data(
        Rank rank, Tag tag, std::unique_ptr<std::vector<std::uint8_t>> synced_buffer
    ) = 0;
 
    [[nodiscard]] virtual std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank>
    recv_any(Tag tag) = 0;
 
    [[nodiscard]] virtual std::unique_ptr<std::vector<std::uint8_t>> recv_from(
        Rank src, Tag tag
    ) = 0;
 
    [[nodiscard]] virtual std::
        pair<std::vector<std::unique_ptr<Future>>, std::vector<std::size_t>>
        test_some(std::vector<std::unique_ptr<Future>>& future_vector) = 0;
 
    std::vector<std::size_t> virtual test_some(
        std::unordered_map<std::size_t, std::unique_ptr<Communicator::Future>> const&
            future_map
    ) = 0;
 
    [[nodiscard]] virtual bool test(std::unique_ptr<Communicator::Future>& future) = 0;
 
    [[nodiscard]] virtual std::vector<std::unique_ptr<Buffer>> wait_all(
        std::vector<std::unique_ptr<Communicator::Future>>&& futures
    ) = 0;
 
    [[nodiscard]] virtual std::unique_ptr<Buffer> wait(
        std::unique_ptr<Future> future
    ) = 0;
 
    [[nodiscard]] std::unique_ptr<Buffer> virtual release_data(
        std::unique_ptr<Communicator::Future> future
    ) = 0;
 
    [[nodiscard]] std::
        unique_ptr<std::vector<std::uint8_t>> virtual release_sync_host_data(
            std::unique_ptr<Communicator::Future> future
        ) = 0;
 
    [[nodiscard]] virtual std::shared_ptr<Logger> const& logger() = 0;
 
    [[nodiscard]] virtual std::shared_ptr<ProgressThread> const&
    progress_thread() const = 0;
 
    [[nodiscard]] virtual std::string str() const = 0;
};
 
#ifdef RAPIDSMPF_HAVE_UCXX
constexpr bool COMM_HAVE_UCXX = true;
#else
constexpr bool COMM_HAVE_UCXX = false;
#endif
 
#ifdef RAPIDSMPF_HAVE_MPI
constexpr bool COMM_HAVE_MPI = true;
#else
constexpr bool COMM_HAVE_MPI = false;
#endif
 
inline std::ostream& operator<<(std::ostream& os, Communicator const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf