mpi_wrapper.hpp

Go to the documentation of this file.

/*
  mptensor - Parallel Library for Tensor Network Methods
 
  Copyright 2016 Satoshi Morita
 
  mptensor is free software: you can redistribute it and/or modify it
  under the terms of the GNU Lesser General Public License as
  published by the Free Software Foundation, either version 3 of the
  License, or (at your option) any later version.
 
  mptensor is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
 
  You should have received a copy of the GNU Lesser General Public
  License along with mptensor.  If not, see
  <https://www.gnu.org/licenses/>.
*/
 
#ifndef _MPI_WRAPPER_HPP_
#define _MPI_WRAPPER_HPP_
 
#ifndef _NO_MPI
 
#include <mpi.h>
 
#include <vector>
 
#include "mptensor/complex.hpp"
 
namespace mptensor {
namespace mpi {
 
template <typename C>
inline MPI_Datatype mpi_datatype = MPI_DATATYPE_NULL;
 
// clang-format off
template <> inline MPI_Datatype mpi_datatype<char>                = MPI_CHAR;
template <> inline MPI_Datatype mpi_datatype<signed char>         = MPI_SIGNED_CHAR;
template <> inline MPI_Datatype mpi_datatype<unsigned char>       = MPI_UNSIGNED_CHAR;
template <> inline MPI_Datatype mpi_datatype<short>               = MPI_SHORT;
template <> inline MPI_Datatype mpi_datatype<unsigned short>      = MPI_UNSIGNED_SHORT;
template <> inline MPI_Datatype mpi_datatype<int>                 = MPI_INT;
template <> inline MPI_Datatype mpi_datatype<unsigned int>        = MPI_UNSIGNED;
template <> inline MPI_Datatype mpi_datatype<long int>            = MPI_LONG;
template <> inline MPI_Datatype mpi_datatype<unsigned long int>   = MPI_UNSIGNED_LONG;
template <> inline MPI_Datatype mpi_datatype<long long int>       = MPI_LONG_LONG;
template <> inline MPI_Datatype mpi_datatype<unsigned long long int> = MPI_UNSIGNED_LONG_LONG;
template <> inline MPI_Datatype mpi_datatype<double>              = MPI_DOUBLE;
template <> inline MPI_Datatype mpi_datatype<complex>             = MPI_DOUBLE_COMPLEX;
// clang-format on
 
 
template <typename C>
inline C allreduce_sum(C val, const MPI_Comm& comm) {
  C recv;
  MPI_Allreduce(&val, &recv, 1, mpi_datatype<C>, MPI_SUM, comm);
  return recv;
};
 
 
template <typename C>
inline std::vector<C> allreduce_vec(const std::vector<C>& vec,
                                    const MPI_Comm& comm) {
  size_t n = vec.size();
  std::vector<C> recv(n);
  MPI_Allreduce(const_cast<C*>(vec.data()), recv.data(), static_cast<int>(n),
                mpi_datatype<C>, MPI_SUM, comm);
  return recv;
};
 
 
template <typename C>
inline void allreduce(const C* sendbuf, C* recvbuf, int count, MPI_Op op,
                      const MPI_Comm& comm) {
  MPI_Allreduce(const_cast<C*>(sendbuf), recvbuf, count, mpi_datatype<C>, op,
                comm);
};
 
 
template <typename C>
inline void sendrecv(const C* sendbuf, int sendcount, int dest, int sendtag,
                     C* recvbuf, int recvcount, int source, int recvtag,
                     const MPI_Comm& comm) {
  MPI_Sendrecv(const_cast<C*>(sendbuf), sendcount, mpi_datatype<C>, dest,
               sendtag, recvbuf, recvcount, mpi_datatype<C>, source, recvtag,
               comm, MPI_STATUS_IGNORE);
};
 
 
template <typename C>
inline void sendrecv(const std::vector<C>& send_vec, int dest, int sendtag,
                     std::vector<C>& recv_vec, int source, int recvtag,
                     const MPI_Comm& comm) {
  MPI_Sendrecv(const_cast<C*>(send_vec.data()),
               static_cast<int>(send_vec.size()), mpi_datatype<C>, dest,
               sendtag, recv_vec.data(), static_cast<int>(recv_vec.size()),
               mpi_datatype<C>, source, recvtag, comm, MPI_STATUS_IGNORE);
};
 
 
template <typename C>
inline void alltoall(const C* sendbuf, int sendcount, C* recvbuf, int recvcount,
                     const MPI_Comm& comm) {
  MPI_Alltoall(const_cast<C*>(sendbuf), sendcount, mpi_datatype<C>, recvbuf,
               recvcount, mpi_datatype<C>, comm);
};
 
 
template <typename C>
inline void alltoallv(const C* sendbuf, const int* sendcounts,
                      const int* sdispls, C* recvbuf, const int* recvcounts,
                      const int* rdispls, const MPI_Comm& comm) {
  MPI_Alltoallv(const_cast<C*>(sendbuf), const_cast<int*>(sendcounts),
                const_cast<int*>(sdispls), mpi_datatype<C>, recvbuf,
                const_cast<int*>(recvcounts), const_cast<int*>(rdispls),
                mpi_datatype<C>, comm);
};
 
 
template <typename C>
inline void bcast(C* buffer, int count, int root, const MPI_Comm& comm) {
  MPI_Bcast(buffer, count, mpi_datatype<C>, root, comm);
};
 
}  // namespace mpi
}  // namespace mptensor
 
#endif  // _NO_MPI
#endif  // _MPI_WRAPPER_HPP_