MPIWrapper.hpp

Go to the documentation of this file.

/*
 * STRUMPACK -- STRUctured Matrices PACKage, Copyright (c) 2014, The
 * Regents of the University of California, through Lawrence Berkeley
 * National Laboratory (subject to receipt of any required approvals
 * from the U.S. Dept. of Energy).  All rights reserved.
 *
 * If you have questions about your rights to use or distribute this
 * software, please contact Berkeley Lab's Technology Transfer
 * Department at TTD@lbl.gov.
 *
 * NOTICE. This software is owned by the U.S. Department of Energy. As
 * such, the U.S. Government has been granted for itself and others
 * acting on its behalf a paid-up, nonexclusive, irrevocable,
 * worldwide license in the Software to reproduce, prepare derivative
 * works, and perform publicly and display publicly.  Beginning five
 * (5) years after the date permission to assert copyright is obtained
 * from the U.S. Department of Energy, and subject to any subsequent
 * five (5) year renewals, the U.S. Government is granted for itself
 * and others acting on its behalf a paid-up, nonexclusive,
 * irrevocable, worldwide license in the Software to reproduce,
 * prepare derivative works, distribute copies to the public, perform
 * publicly and display publicly, and to permit others to do so.
 *
 * Developers: Pieter Ghysels, Francois-Henry Rouet, Xiaoye S. Li.
 *             (Lawrence Berkeley National Lab, Computational Research
 *             Division).
 *
 */
#ifndef STRUMPACK_MPI_WRAPPER_HPP
#define STRUMPACK_MPI_WRAPPER_HPP
 
#include <vector>
#include <complex>
#include <cassert>
#include <numeric>
#include <limits>
#include <memory>
#include <utility>
 
#define OMPI_SKIP_MPICXX 1
#include <mpi.h>
 
#include "StrumpackParameters.hpp"
#include "Triplet.hpp"
 
namespace strumpack {
 
  template<typename T> MPI_Datatype mpi_type() { return T::mpi_type(); }
  template<> inline MPI_Datatype mpi_type<char>() { return MPI_CHAR; }
  template<> inline MPI_Datatype mpi_type<bool>() { return MPI_CXX_BOOL; }
  template<> inline MPI_Datatype mpi_type<int>() { return MPI_INT; }
  template<> inline MPI_Datatype mpi_type<long>() { return MPI_LONG; }
  template<> inline MPI_Datatype mpi_type<unsigned long>() { return MPI_UNSIGNED_LONG; }
  template<> inline MPI_Datatype mpi_type<long long int>() { return MPI_LONG_LONG_INT; }
  template<> inline MPI_Datatype mpi_type<float>() { return MPI_FLOAT; }
  template<> inline MPI_Datatype mpi_type<double>() { return MPI_DOUBLE; }
  template<> inline MPI_Datatype mpi_type<std::complex<float>>() { return MPI_CXX_FLOAT_COMPLEX; }
  template<> inline MPI_Datatype mpi_type<std::complex<double>>() { return MPI_CXX_DOUBLE_COMPLEX; }
  template<> inline MPI_Datatype mpi_type<std::pair<int,int>>() { return MPI_2INT; }
 
  template<> inline MPI_Datatype mpi_type<std::pair<long int,long int>>() {
    static MPI_Datatype l_l_mpi_type = MPI_DATATYPE_NULL;
    if (l_l_mpi_type == MPI_DATATYPE_NULL) {
      MPI_Type_contiguous
        (2, strumpack::mpi_type<long int>(), &l_l_mpi_type);
      MPI_Type_commit(&l_l_mpi_type);
    }
    return l_l_mpi_type;
  }
  template<> inline MPI_Datatype mpi_type<std::pair<long long int,long long int>>() {
    static MPI_Datatype ll_ll_mpi_type = MPI_DATATYPE_NULL;
    if (ll_ll_mpi_type == MPI_DATATYPE_NULL) {
      MPI_Type_contiguous
        (2, strumpack::mpi_type<long long int>(), &ll_ll_mpi_type);
      MPI_Type_commit(&ll_ll_mpi_type);
    }
    return ll_ll_mpi_type;
  }
 
  class MPIRequest {
  public:
    MPIRequest() {
      req_ = std::unique_ptr<MPI_Request>(new MPI_Request());
    }
 
    MPIRequest(const MPIRequest&) = delete;
 
    MPIRequest(MPIRequest&&) = default;
 
    MPIRequest& operator=(const MPIRequest&) = delete;
 
    MPIRequest& operator=(MPIRequest&&) = default;
 
    void wait() { MPI_Wait(req_.get(), MPI_STATUS_IGNORE); }
 
  private:
    std::unique_ptr<MPI_Request> req_;
    friend class MPIComm;
  };
 
  inline void wait_all(std::vector<MPIRequest>& reqs) {
    for (auto& r : reqs) r.wait();
    reqs.clear();
  }
 
  inline void wait_all(std::vector<MPI_Request>& reqs) {
    MPI_Waitall(reqs.size(), reqs.data(), MPI_STATUSES_IGNORE);
  }
 
 
  class MPIComm {
  public:
    MPIComm() {}
 
    MPIComm(MPI_Comm c) { duplicate(c); }
 
    MPIComm(const MPIComm& c) { *this = c; }
 
    MPIComm(MPIComm&& c) noexcept { *this = std::move(c); }
 
    virtual ~MPIComm() {
      if (comm_ != MPI_COMM_NULL && comm_ != MPI_COMM_WORLD)
        MPI_Comm_free(&comm_);
    }
 
    MPIComm& operator=(const MPIComm& c) {
      if (this != &c) duplicate(c.comm());
      return *this;
    }
 
    MPIComm& operator=(MPIComm&& c) noexcept {
      comm_ = c.comm_;
      c.comm_ = MPI_COMM_NULL;
      return *this;
    }
 
    MPI_Comm comm() const { return comm_; }
 
    bool is_null() const { return comm_ == MPI_COMM_NULL; }
 
    int rank() const {
      assert(comm_ != MPI_COMM_NULL);
      int r;
      MPI_Comm_rank(comm_, &r);
      return r;
    }
 
    int size() const {
      assert(comm_ != MPI_COMM_NULL);
      int nprocs;
      MPI_Comm_size(comm_, &nprocs);
      return nprocs;
    }
 
    bool is_root() const { return rank() == 0; }
 
    void barrier() const { MPI_Barrier(comm_); }
 
    template<typename T> void
    broadcast(std::vector<T>& sbuf) const {
      MPI_Bcast(sbuf.data(), sbuf.size(), mpi_type<T>(), 0, comm_);
    }
    template<typename T> void
    broadcast_from(std::vector<T>& sbuf, int src) const {
      MPI_Bcast(sbuf.data(), sbuf.size(), mpi_type<T>(), src, comm_);
    }
 
    template<typename T, std::size_t N> void
    broadcast(std::array<T,N>& sbuf) const {
      MPI_Bcast(sbuf.data(), sbuf.size(), mpi_type<T>(), 0, comm_);
    }
 
    template<typename T> void
    broadcast(T& data) const {
      MPI_Bcast(&data, 1, mpi_type<T>(), 0, comm_);
    }
    template<typename T> void
    broadcast_from(T& data, int src) const {
      MPI_Bcast(&data, 1, mpi_type<T>(), src, comm_);
    }
    template<typename T> void
    broadcast(T* sbuf, std::size_t ssize) const {
      MPI_Bcast(sbuf, ssize, mpi_type<T>(), 0, comm_);
    }
 
    template<typename T>
    void all_gather(T* buf, std::size_t rsize) const {
      MPI_Allgather
        (MPI_IN_PLACE, 0, MPI_DATATYPE_NULL,
         buf, rsize, mpi_type<T>(), comm_);
    }
 
    template<typename T>
    void all_gather_v(T* buf, const int* rcnts, const int* displs) const {
      MPI_Allgatherv
        (MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, buf, rcnts, displs,
         mpi_type<T>(), comm_);
    }
 
 
    template<typename T>
    MPIRequest isend(const std::vector<T>& sbuf, int dest, int tag) const {
      MPIRequest req;
      // const_cast is necessary for ancient openmpi version used on Travis
      MPI_Isend(const_cast<T*>(sbuf.data()), sbuf.size(), mpi_type<T>(),
                dest, tag, comm_, req.req_.get());
      return req;
    }
 
    template<typename T>
    void isend(const std::vector<T>& sbuf, int dest, int tag,
               MPI_Request* req) const {
      // const_cast is necessary for ancient openmpi version used on Travis
      MPI_Isend(const_cast<T*>(sbuf.data()), sbuf.size(), mpi_type<T>(),
                dest, tag, comm_, req);
    }
 
    template<typename T>
    void isend(const T* sbuf, std::size_t ssize, int dest,
               int tag, MPI_Request* req) const {
      // const_cast is necessary for ancient openmpi version used on Travis
      MPI_Isend(const_cast<T*>(sbuf), ssize, mpi_type<T>(),
                dest, tag, comm_, req);
    }
 
    template<typename T>
    void isend(const T& buf, int dest, int tag, MPI_Request* req) const {
      // const_cast is necessary for ancient openmpi version used on Travis
      MPI_Isend(const_cast<T*>(&buf), 1, mpi_type<T>(),
                dest, tag, comm_, req);
    }
 
    template<typename T>
    void send(const std::vector<T>& sbuf, int dest, int tag) const {
      // const_cast is necessary for ancient openmpi version used on Travis
      MPI_Send(const_cast<T*>(sbuf.data()), sbuf.size(), mpi_type<T>(), dest, tag, comm_);
    }
 
    template<typename T> std::vector<T> recv(int src, int tag) const {
      MPI_Status stat;
      MPI_Probe(src, tag, comm_, &stat);
      int msgsize;
      MPI_Get_count(&stat, mpi_type<T>(), &msgsize);
      //std::vector<T,NoInit<T>> rbuf(msgsize);
      std::vector<T> rbuf(msgsize);
      MPI_Recv(rbuf.data(), msgsize, mpi_type<T>(), src, tag,
               comm_, MPI_STATUS_IGNORE);
      return rbuf;
    }
 
    template<typename T>
    std::pair<int,std::vector<T>> recv_any_src(int tag) const {
      MPI_Status stat;
      MPI_Probe(MPI_ANY_SOURCE, tag, comm_, &stat);
      int msgsize;
      MPI_Get_count(&stat, mpi_type<T>(), &msgsize);
      std::vector<T> rbuf(msgsize);
      MPI_Recv(rbuf.data(), msgsize, mpi_type<T>(), stat.MPI_SOURCE,
               tag, comm_, MPI_STATUS_IGNORE);
      return {stat.MPI_SOURCE, std::move(rbuf)};
    }
 
    template<typename T> T recv_one(int src, int tag) const {
      T t;
      MPI_Recv(&t, 1, mpi_type<T>(), src, tag, comm_, MPI_STATUS_IGNORE);
      return t;
    }
 
    template<typename T>
    void irecv(const T* rbuf, std::size_t rsize, int src,
               int tag, MPI_Request* req) const {
      // const_cast is necessary for ancient openmpi version used on Travis
      MPI_Irecv(const_cast<T*>(rbuf), rsize, mpi_type<T>(),
                src, tag, comm_, req);
    }
 
    template<typename T> T all_reduce(T t, MPI_Op op) const {
      MPI_Allreduce(MPI_IN_PLACE, &t, 1, mpi_type<T>(), op, comm_);
      return t;
    }
 
    template<typename T> T reduce(T t, MPI_Op op) const {
      if (is_root())
        MPI_Reduce(MPI_IN_PLACE, &t, 1, mpi_type<T>(), op, 0, comm_);
      else MPI_Reduce(&t, &t, 1, mpi_type<T>(), op, 0, comm_);
      return t;
    }
 
    template<typename T> void all_reduce(T* t, int ssize, MPI_Op op) const {
      MPI_Allreduce(MPI_IN_PLACE, t, ssize, mpi_type<T>(), op, comm_);
    }
 
    template<typename T> void all_reduce(std::vector<T>& t, MPI_Op op) const {
      all_reduce(t.data(), t.size(), op);
    }
 
    template<typename T> void reduce(T* t, int ssize, MPI_Op op) const {
      if (is_root())
        MPI_Reduce(MPI_IN_PLACE, t, ssize, mpi_type<T>(), op, 0, comm_);
      else MPI_Reduce(t, t, ssize, mpi_type<T>(), op, 0, comm_);
    }
 
    template<typename T>
    void all_to_all(const T* sbuf, int scnt, T* rbuf) const {
      MPI_Alltoall
        (sbuf, scnt, mpi_type<T>(), rbuf, scnt, mpi_type<T>(), comm_);
    }
 
    template<typename T, typename A=std::allocator<T>> std::vector<T,A>
    all_to_allv(const T* sbuf, int* scnts, int* sdispls,
                int* rcnts, int* rdispls) const {
      std::size_t rsize = 0;
      for (int p=0; p<size(); p++)
        rsize += rcnts[p];
      std::vector<T,A> rbuf(rsize);
      MPI_Alltoallv
        (sbuf, scnts, sdispls, mpi_type<T>(),
         rbuf.data(), rcnts, rdispls, mpi_type<T>(), comm_);
      return rbuf;
    }
 
    template<typename T> void
    all_to_allv(const T* sbuf, int* scnts, int* sdispls,
                T* rbuf, int* rcnts, int* rdispls) const {
      MPI_Alltoallv
        (sbuf, scnts, sdispls, mpi_type<T>(),
         rbuf, rcnts, rdispls, mpi_type<T>(), comm_);
    }
 
    template<typename T, typename A=std::allocator<T>> void
    all_to_all_v(std::vector<std::vector<T>>& sbuf, std::vector<T,A>& rbuf,
                 std::vector<T*>& pbuf) const {
      all_to_all_v(sbuf, rbuf, pbuf, mpi_type<T>());
    }
 
    template<typename T, typename A=std::allocator<T>> std::vector<T,A>
    all_to_all_v(std::vector<std::vector<T>>& sbuf) const {
      std::vector<T,A> rbuf;
      std::vector<T*> pbuf;
      all_to_all_v(sbuf, rbuf, pbuf, mpi_type<T>());
      return rbuf;
    }
 
    template<typename T, typename A=std::allocator<T>> void
    all_to_all_v(std::vector<std::vector<T>>& sbuf, std::vector<T,A>& rbuf,
                 std::vector<T*>& pbuf, const MPI_Datatype Ttype) const {
      assert(sbuf.size() == std::size_t(size()));
      auto P = size();
      std::unique_ptr<int[]> iwork(new int[4*P]);
      auto ssizes = iwork.get();
      auto rsizes = ssizes + P;
      auto sdispl = ssizes + 2*P;
      auto rdispl = ssizes + 3*P;
      for (int p=0; p<P; p++) {
        if (sbuf[p].size() >
            static_cast<std::size_t>(std::numeric_limits<int>::max())) {
          std::cerr << "# ERROR: 32bit integer overflow in all_to_all_v!!"
                    << std::endl;
          MPI_Abort(comm_, 1);
        }
        ssizes[p] = sbuf[p].size();
      }
      MPI_Alltoall
        (ssizes, 1, mpi_type<int>(), rsizes, 1, mpi_type<int>(), comm_);
      std::size_t totssize = std::accumulate(ssizes, ssizes+P, std::size_t(0)),
        totrsize = std::accumulate(rsizes, rsizes+P, std::size_t(0));
      if (totrsize >
          static_cast<std::size_t>(std::numeric_limits<int>::max()) ||
          totssize >
          static_cast<std::size_t>(std::numeric_limits<int>::max())) {
        // This case will probably cause an overflow in the
        // rdispl/sdispl elements. Here we do the all_to_all_v
        // manually by just using Isend/Irecv. This might be slower
        // than splitting into multiple calls to MPI_Alltoallv
        // (although it avoids a copy from the sbuf).
        rbuf.resize(totrsize);
        std::unique_ptr<MPI_Request[]> reqs(new MPI_Request[2*P]);
        std::size_t displ = 0;
        pbuf.resize(P);
        for (int p=0; p<P; p++) {
          pbuf[p] = rbuf.data() + displ;
          MPI_Irecv(pbuf[p], rsizes[p], Ttype, p, 0, comm_, reqs.get()+p);
          displ += rsizes[p];
        }
        for (int p=0; p<P; p++)
          MPI_Isend
            (sbuf[p].data(), ssizes[p], Ttype, p, 0, comm_, reqs.get()+P+p);
        MPI_Waitall(2*P, reqs.get(), MPI_STATUSES_IGNORE);
        std::vector<std::vector<T>>().swap(sbuf);
      } else {
        std::unique_ptr<T[]> sendbuf_(new T[totssize]);
        auto sendbuf = sendbuf_.get();
        sdispl[0] = rdispl[0] = 0;
        for (int p=1; p<P; p++) {
          sdispl[p] = sdispl[p-1] + ssizes[p-1];
          rdispl[p] = rdispl[p-1] + rsizes[p-1];
        }
        for (int p=0; p<P; p++)
          std::copy(sbuf[p].begin(), sbuf[p].end(), sendbuf+sdispl[p]);
        std::vector<std::vector<T>>().swap(sbuf);
        rbuf.resize(totrsize);
        MPI_Alltoallv(sendbuf, ssizes, sdispl, Ttype,
                      rbuf.data(), rsizes, rdispl, Ttype, comm_);
        pbuf.resize(P);
        for (int p=0; p<P; p++)
          pbuf[p] = rbuf.data() + rdispl[p];
      }
    }
 
    MPIComm sub(int P0, int P, int stride=1) const {
      if (is_null() || size() == 1)
        return MPIComm(MPI_COMM_NULL);
      assert(P0 + P <= size());
      MPIComm sub_comm;
      std::vector<int> sub_ranks(P);
      for (int i=0; i<P; i++)
        sub_ranks[i] = P0 + i*stride;
      MPI_Group group, sub_group;
      MPI_Comm_group(comm_, &group);                          // get group from comm
      MPI_Group_incl(group, P, sub_ranks.data(), &sub_group); // group ranks [P0,P0+P) into sub_group
      MPI_Comm_create(comm_, sub_group, &sub_comm.comm_);     // create new sub_comm
      MPI_Group_free(&group);
      MPI_Group_free(&sub_group);
      return sub_comm;
    }
 
    // return MPI_COMM_SELF??? or MPI_COMM_NULL if not rank??
    MPIComm sub_self(int p) const {
      if (is_null()) return MPIComm(MPI_COMM_NULL);
      MPIComm c0;
      MPI_Group group, sub_group;
      MPI_Comm_group(comm_, &group);
      MPI_Group_incl(group, 1, &p, &sub_group);
      MPI_Comm_create(comm_, sub_group, &c0.comm_);
      MPI_Group_free(&group);
      MPI_Group_free(&sub_group);
      return c0;
    }
 
    static void control_start(const std::string& name) {
      MPI_Pcontrol(1, name.c_str());
    }
    static void control_stop(const std::string& name) {
      MPI_Pcontrol(-1, name.c_str());
    }
 
    static bool initialized() {
      int flag;
      MPI_Initialized(&flag);
      return static_cast<bool>(flag);
    }
 
  private:
    MPI_Comm comm_ = MPI_COMM_WORLD;
 
    void duplicate(MPI_Comm c) {
      if (c == MPI_COMM_NULL) comm_ = c;
      else MPI_Comm_dup(c, &comm_);
    }
  };
 
 
  inline int mpi_rank(MPI_Comm c=MPI_COMM_WORLD) {
    assert(c != MPI_COMM_NULL);
    int rank;
    MPI_Comm_rank(c, &rank);
    return rank;
  }
 
  inline int mpi_nprocs(MPI_Comm c=MPI_COMM_WORLD) {
    assert(c != MPI_COMM_NULL);
    int nprocs;
    MPI_Comm_size(c, &nprocs);
    return nprocs;
  }
 
} // end namespace strumpack
 
#endif // STRUMPACK_MPI_WRAPPER_HPP