op_mpi.hpp

#pragma once
#include <mpi.h>
namespace op {

namespace mpi {
namespace detail {

// default template
template <typename T>
struct mpi_t {
  static MPI_Datatype type() {return MPI_BYTE; }
};

template <>
struct mpi_t<double> {
  static MPI_Datatype type() {return MPI_DOUBLE;}
};

template <>
struct mpi_t<int> {
  static MPI_Datatype type() {return MPI_INT;}
};

template <>
struct mpi_t<unsigned long> {
  static MPI_Datatype type() {return MPI_UNSIGNED_LONG;}
};

template <typename T, typename SFINAE = void>
struct has_data : std::false_type {
};

template <typename T>
struct has_data<T, std::void_t<decltype(std::declval<T>().data())>> : std::true_type {
};

template <typename T, typename SFINAE = void>
struct has_size : std::false_type {
};

template <typename T>
struct has_size<T, std::void_t<decltype(std::declval<T>().size())>> : std::true_type {
};

}  // namespace detail

int getRank(MPI_Comm comm = MPI_COMM_WORLD)
{
  int rank;
  MPI_Comm_rank(comm, &rank);
  return rank;
}

int getNRanks(MPI_Comm comm = MPI_COMM_WORLD)
{
  int nranks;
  MPI_Comm_size(comm, &nranks);
  return nranks;
}

template <typename T>
std::enable_if_t<!(detail::has_data<T>::value && detail::has_size<T>::value), int> Allreduce(
    T& local, T& global, MPI_Op operation, MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Allreduce(&local, &global, 1, mpi::detail::mpi_t<T>::type(), operation, comm);
}

template <typename T>
std::enable_if_t<(detail::has_data<T>::value && detail::has_size<T>::value), int> Allreduce(
    T& local, T& global, MPI_Op operation, MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Allreduce(local.data(), global.data(), local.size(), mpi::detail::mpi_t<typename T::value_type>::type(),
                       operation, comm);
}

template <typename T>
std::enable_if_t<!(detail::has_data<T>::value && detail::has_size<T>::value), int> Broadcast(
    T& buf, int root = 0, MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Bcast(&buf, 1, mpi::detail::mpi_t<T>::type(), root, comm);
}

template <typename T>
std::enable_if_t<(detail::has_data<T>::value && detail::has_size<T>::value), int> Broadcast(
    T& buf, int root = 0, MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Bcast(buf.data(), static_cast<int>(buf.size()), mpi::detail::mpi_t<typename T::value_type>::type(), root,
                   comm);
}

template <typename T>
int Allgatherv(T& buf, T& values_on_rank, std::vector<int>& size_on_rank, std::vector<int>& offsets_on_rank,
               MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Allgatherv(buf.data(), static_cast<int>(buf.size()), detail::mpi_t<typename T::value_type>::type(),
                        values_on_rank.data(), size_on_rank.data(), offsets_on_rank.data(),
                        detail::mpi_t<typename T::value_type>::type(), comm);
}

template <typename T>
int Gatherv(T& buf, T& values_on_rank, std::vector<int>& size_on_rank, std::vector<int>& offsets_on_rank, int root = 0,
            MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Gatherv(buf.data(), static_cast<int>(buf.size()), detail::mpi_t<typename T::value_type>::type(),
                     values_on_rank.data(), size_on_rank.data(), offsets_on_rank.data(),
                     detail::mpi_t<typename T::value_type>::type(), root, comm);
}

template <typename T>
int Scatterv(T& sendbuf, std::vector<int>& variables_per_rank, std::vector<int>& offsets, T& recvbuff, int root = 0,
             MPI_Comm comm = MPI_COMM_WORLD)
{
  // only check the size of the recv buff in debug mode
  assert(static_cast<typename T::size_type>(
             [&]() { return variables_per_rank[static_cast<std::size_t>(mpi::getRank(comm))]; }()) == recvbuff.size());

  return MPI_Scatterv(sendbuf.data(), variables_per_rank.data(), offsets.data(),
                      mpi::detail::mpi_t<typename T::value_type>::type(), recvbuff.data(),
                      static_cast<int>(recvbuff.size()), mpi::detail::mpi_t<typename T::value_type>::type(), root, comm);
}

template <typename T>
int Irecv(T& buf, int send_rank, MPI_Request* request, int tag = 0, MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Irecv(buf.data(), static_cast<int>(buf.size()), mpi::detail::mpi_t<typename T::value_type>::type(),
                   send_rank, tag, comm, request);
}

template <typename T>
int Isend(T& buf, int recv_rank, MPI_Request* request, int tag = 0, MPI_Comm comm = MPI_COMM_WORLD)
{
  return MPI_Isend(buf.data(), static_cast<int>(buf.size()), mpi::detail::mpi_t<typename T::value_type>::type(),
                   recv_rank, tag, comm, request);
}

int Waitall(std::vector<MPI_Request>& requests, std::vector<MPI_Status>& status)
{
  return MPI_Waitall(static_cast<int>(requests.size()), requests.data(), status.data());
}

int CreateAndSetErrorHandler(MPI_Errhandler& newerr, void (*err)(MPI_Comm* comm, int* err, ...),
                             MPI_Comm        comm = MPI_COMM_WORLD)
{
  MPI_Comm_create_errhandler(err, &newerr);
  return MPI_Comm_set_errhandler(comm, newerr);
}

}  // namespace mpi

}  // namespace op