mpi_common.h

#ifndef CAFFE2_MPI_MPI_COMMON_H_
#define CAFFE2_MPI_MPI_COMMON_H_

#include <mpi.h>
#include <mutex>

#include "caffe2/core/common.h"
#include "caffe2/core/logging.h"

namespace caffe2 {

inline void CheckInitializedMPI() {
  int flag;
  MPI_Initialized(&flag);
  CAFFE_ENFORCE(flag, "MPI does not seem to have been initialized.");
}

template <typename T> class MPIDataTypeWrapper;

#define MPI_DATATYPE_WRAPPER(c_type, mpi_type)                                 \
  template<> class MPIDataTypeWrapper<c_type> {                                \
   public:                                                                     \
    inline static MPI_Datatype type() { return mpi_type; }                     \
  };

MPI_DATATYPE_WRAPPER(char, MPI_CHAR)
MPI_DATATYPE_WRAPPER(float, MPI_FLOAT)
MPI_DATATYPE_WRAPPER(double, MPI_DOUBLE)
// Note(Yangqing): as necessary, add more specializations.
#undef MPI_DATATYPE_WRAPPER

// For all Caffe MPI calls, we will wrap it inside an MPI mutex lock guard.
CAFFE2_API std::mutex& MPIMutex();

#define MPI_CHECK(condition)                                 \
  do {                                                       \
    std::lock_guard<std::mutex> guard(::caffe2::MPIMutex()); \
    int error = (condition);                                 \
    CAFFE_ENFORCE(                                           \
        error == MPI_SUCCESS,                                \
        "Caffe2 MPI Error at: ",                             \
        __FILE__,                                            \
        ":",                                                 \
        __LINE__,                                            \
        ": ",                                                \
        error);                                              \
  } while (0)

CAFFE2_API MPI_Comm GlobalMPIComm();

CAFFE2_API void SetGlobalMPIComm(MPI_Comm new_comm);

CAFFE2_API int MPICommSize(MPI_Comm comm);

CAFFE2_API int MPICommRank(MPI_Comm comm);

class MPICommonWorldWrapper {
 public:
  explicit MPICommonWorldWrapper(
      MPI_Comm src_comm = MPI_COMM_NULL,
      int color = 0,
      int rank = -1) {
    if (src_comm == MPI_COMM_NULL) {
      src_comm = GlobalMPIComm();
    }
    if (rank == -1) {
      MPI_CHECK(MPI_Comm_rank(src_comm, &rank));
    }
    MPI_CHECK(MPI_Comm_split(src_comm, color, rank, &comm_));
    MPI_CHECK(MPI_Comm_size(comm_, &size_));
    MPI_CHECK(MPI_Comm_rank(comm_, &rank_));
  }

  ~MPICommonWorldWrapper() {
    int ret;
    MPI_CHECK(MPI_Finalized(&ret));
    if (!ret) {
      MPI_Comm_free(&comm_);
    }
  }

  inline MPI_Comm comm() const {
    return comm_;
  }
  inline int size() const {
    return size_;
  }
  inline int rank() const {
    return rank_;
  }

 private:
  MPI_Comm comm_;
  int size_;
  int rank_;
};

void MPISetupPeers(
    const int replicas,
    const string& role,
    const string& job_path);
}  // namespace caffe2

#endif  // CAFFE2_MPI_MPI_COMMON_H_