mpi.hpp

 
#pragma once
 
#include <cstdlib>
#include <memory>
#include <vector>
 
#include <mpi.h>
 
#include <rmm/device_buffer.hpp>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/progress_thread.hpp>
 
namespace rapidsmpf {
 
namespace mpi {
 
void init(int* argc, char*** argv);
 
bool is_initialized();
 
#define RAPIDSMPF_MPI(call) \
    rapidsmpf::mpi::detail::check_mpi_error((call), __FILE__, __LINE__)
 
namespace detail {
void check_mpi_error(int error_code, char const* file, int line);
}  // namespace detail
}  // namespace mpi
 
class MPI final : public Communicator {
  public:
    class Future : public Communicator::Future {
        friend class MPI;
 
      public:
        Future(MPI_Request req, std::unique_ptr<Buffer> data_buffer)
            : req_{req}, data_buffer_{std::move(data_buffer)} {}
 
        Future(
            MPI_Request req, std::unique_ptr<std::vector<std::uint8_t>> synced_host_data
        )
            : req_{std::move(req)}, synced_host_data_{std::move(synced_host_data)} {}
 
        ~Future() noexcept override = default;
 
      private:
        MPI_Request req_;
        // TODO: these buffers are mutually exclusive and looks similar to
        // Buffer::storage_.
        std::unique_ptr<Buffer> data_buffer_;
        // Dedicated storage for host data that is valid at the time of construction.
        std::unique_ptr<std::vector<std::uint8_t>> synced_host_data_;
    };
 
    MPI(MPI_Comm comm,
        std::shared_ptr<ProgressThread> progress_thread,
        std::shared_ptr<Logger> logger);
 
    ~MPI() noexcept override = default;
 
    [[nodiscard]] Rank rank() const override {
        return rank_;
    }
 
    [[nodiscard]] Rank nranks() const override {
        return nranks_;
    }
 
    [[nodiscard]] std::unique_ptr<Communicator::Future> send(
        std::unique_ptr<std::vector<std::uint8_t>> msg, Rank rank, Tag tag
    ) override;
 
    // clang-format off
    // clang-format on
    [[nodiscard]] std::unique_ptr<Communicator::Future> send(
        std::unique_ptr<Buffer> msg, Rank rank, Tag tag
    ) override;
 
    [[nodiscard]] std::unique_ptr<Communicator::Future> recv(
        Rank rank, Tag tag, std::unique_ptr<Buffer> recv_buffer
    ) override;
 
    // clang-format off
    // clang-format on
    [[nodiscard]] std::unique_ptr<Communicator::Future> recv_sync_host_data(
        Rank rank, Tag tag, std::unique_ptr<std::vector<std::uint8_t>> synced_buffer
    ) override;
 
    [[nodiscard]] std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> recv_any(
        Tag tag
    ) override;
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> recv_from(
        Rank src, Tag tag
    ) override;
    std::pair<
        std::vector<std::unique_ptr<Communicator::Future>>,
        std::vector<std::size_t>>
    test_some(std::vector<std::unique_ptr<Communicator::Future>>& future_vector) override;
 
    // clang-format off
    // clang-format on
    std::vector<std::size_t> test_some(
        std::unordered_map<std::size_t, std::unique_ptr<Communicator::Future>> const&
            future_map
    ) override;
 
    bool test(std::unique_ptr<Communicator::Future>& future) override;
    std::vector<std::unique_ptr<Buffer>> wait_all(
        std::vector<std::unique_ptr<Communicator::Future>>&& futures
    ) override;
 
    [[nodiscard]] std::unique_ptr<Buffer> wait(
        std::unique_ptr<Communicator::Future> future
    ) override;
 
    [[nodiscard]] std::unique_ptr<Buffer> release_data(
        std::unique_ptr<Communicator::Future> future
    ) override;
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> release_sync_host_data(
        std::unique_ptr<Communicator::Future> future
    ) override;
 
    [[nodiscard]] std::shared_ptr<Logger> const& logger() override {
        return logger_;
    }
 
    [[nodiscard]] std::shared_ptr<ProgressThread> const&
    progress_thread() const override {
        return progress_thread_;
    }
 
    [[nodiscard]] std::string str() const override;
 
  private:
    MPI_Comm comm_;
    Rank rank_;
    Rank nranks_;
    std::shared_ptr<Logger> logger_;
    std::shared_ptr<ProgressThread> progress_thread_;
};
 
 
}  // namespace rapidsmpf