communicator.hpp

 
#pragma once
 
#include <cstdlib>
#include <memory>
#include <mutex>
#include <sstream>
#include <thread>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/buffer/buffer.hpp>
#include <rapidsmpf/buffer/resource.hpp>
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/error.hpp>
 
namespace rapidsmpf {
 
using Rank = std::int32_t;
 
using OpID = std::uint8_t;
 
using StageID = std::uint8_t;
 
class Tag {
  public:
    using StorageT = std::int32_t;
 
    static constexpr int stage_id_bits{sizeof(StageID) * 8};
 
    static constexpr StorageT stage_id_mask{(1 << stage_id_bits) - 1};
 
    static constexpr int op_id_bits{sizeof(OpID) * 8};
 
    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)
          } {}
 
    [[nodiscard]] static constexpr 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;  
    };
 
    class Logger {
      public:
        enum class LOG_LEVEL : std::uint32_t {
            NONE = 0,  
            PRINT,  
            WARN,  
            INFO,  
            DEBUG,  
            TRACE  
        };
 
        static constexpr std::array<char const*, 6> LOG_LEVEL_NAMES{
            "NONE", "PRINT", "WARN", "INFO", "DEBUG", "TRACE"
        };
 
        static constexpr const char* level_name(LOG_LEVEL level) {
            auto index = static_cast<std::size_t>(level);
            return index < LOG_LEVEL_NAMES.size() ? LOG_LEVEL_NAMES[index] : "UNKNOWN";
        }
 
        Logger(Communicator* comm, config::Options options);
        virtual ~Logger() noexcept = default;
 
        LOG_LEVEL verbosity_level() const {
            return level_;
        }
 
        template <typename... Args>
        void log(LOG_LEVEL level, Args const&... args) {
            if (static_cast<std::uint32_t>(level_) < static_cast<std::uint32_t>(level)) {
                return;
            }
            std::ostringstream ss;
            (ss << ... << args);
            do_log(level, std::move(ss));
        }
 
        template <typename... Args>
        void print(Args const&... args) {
            log(LOG_LEVEL::PRINT, std::forward<Args const&>(args)...);
        }
 
        template <typename... Args>
        void warn(Args const&... args) {
            log(LOG_LEVEL::WARN, std::forward<Args const&>(args)...);
        }
 
        template <typename... Args>
        void info(Args const&... args) {
            log(LOG_LEVEL::INFO, std::forward<Args const&>(args)...);
        }
 
        template <typename... Args>
        void debug(Args const&... args) {
            log(LOG_LEVEL::DEBUG, std::forward<Args const&>(args)...);
        }
 
        template <typename... Args>
        void trace(Args const&... args) {
            log(LOG_LEVEL::TRACE, std::forward<Args const&>(args)...);
        }
 
      protected:
        virtual std::uint32_t get_thread_id() {
            auto const tid = std::this_thread::get_id();
 
            // To avoid large IDs, we map the thread ID to an unique counter.
            auto const [name, inserted] =
                thread_id_names.insert({tid, thread_id_names_counter});
            if (inserted) {
                ++thread_id_names_counter;
            }
            return name->second;
        }
 
        virtual void do_log(LOG_LEVEL level, std::ostringstream&& ss) {
            std::ostringstream full_log_msg;
            full_log_msg << "[" << level_name(level) << ":" << comm_->rank() << ":"
                         << get_thread_id() << "] " << ss.str();
            std::lock_guard<std::mutex> lock(mutex_);
            std::cout << full_log_msg.str() << std::endl;
        }
 
        Communicator* get_communicator() const {
            return comm_;
        }
 
      private:
        std::mutex mutex_;
        Communicator* comm_;
        LOG_LEVEL const level_;
 
        std::uint32_t thread_id_names_counter{0};
 
        std::unordered_map<std::thread::id, std::uint32_t> thread_id_names;
    };
 
  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<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<uint8_t>> synced_buffer
    ) = 0;
 
    [[nodiscard]] virtual std::pair<std::unique_ptr<std::vector<uint8_t>>, Rank> recv_any(
        Tag tag
    ) = 0;
 
    [[nodiscard]] virtual std::unique_ptr<std::vector<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 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<uint8_t>> virtual release_sync_host_data(
        std::unique_ptr<Communicator::Future> future
    ) = 0;
 
    [[nodiscard]] virtual Logger& logger() = 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