comm.h

Go to the documentation of this file.

 
#ifndef COMM_H
#define COMM_H
 
#include <array>
#include <cstddef>
#include <cstdint>
#include <span>
#include <vector>
#include <type_traits>
 
#ifdef Mallard_HAS_MPI
#include <mpi.h>
#endif
 
namespace comm {
 
class Session {
    public:
        Session(int & argc, char **& argv);
        ~Session();
        Session(const Session &) = delete;
        Session & operator=(const Session &) = delete;
 
    private:
        [[maybe_unused]] bool owns_mpi = false;
};
 
int rank();
 
int size();
 
inline bool is_root() { return rank() == 0; }
 
void barrier();
 
[[noreturn]] void abort(int code);
 
enum class Op { SUM, MIN, MAX };
 
template <typename T>
void allreduce(std::span<T> data, Op op);
 
template <typename T>
T allreduce(T value, Op op) {
    allreduce(std::span<T>(&value, 1), op);
    return value;
}
 
template <typename T, std::size_t N>
std::array<T, N> allreduce(std::array<T, N> values, Op op) {
    allreduce(std::span<T>(values), op);
    return values;
}
 
template <typename T>
std::vector<std::vector<T>> alltoallv(const std::vector<std::vector<T>> & send);
 
template <typename T>
struct Received {
    std::vector<T> data;
    std::vector<uint64_t> offsets;
 
    std::span<const T> from(size_t r) const { return {data.data() + offsets[r], data.data() + offsets[r + 1]}; }
};
 
template <typename T>
Received<T> exchange(std::vector<std::vector<T>> && send);
 
template <typename T>
std::vector<T> allgatherv(const std::vector<T> & local);
 
#ifdef Mallard_HAS_MPI
MPI_Comm world();
#endif
 
} // namespace comm
 
#endif // COMM_H