communicator_mpi.cpp

Go to the documentation of this file.

#include <communicator_quda.h>
 
#define MPI_CHECK(mpi_call)                                                                                            \
  do {                                                                                                                 \
    int status = mpi_call;                                                                                             \
    if (status != MPI_SUCCESS) {                                                                                       \
      char err_string[128];                                                                                            \
      int err_len;                                                                                                     \
      MPI_Error_string(status, err_string, &err_len);                                                                  \
      err_string[127] = '\0';                                                                                          \
      errorQuda("(MPI) %s", err_string);                                                                               \
    }                                                                                                                  \
  } while (0)
 
struct MsgHandle_s {
  MPI_Request request;
 
  MPI_Datatype datatype;
 
  bool custom;
};
 
Communicator::Communicator(int nDim, const int *commDims, QudaCommsMap rank_from_coords, void *map_data,
                           bool user_set_comm_handle_, void *user_comm)
{
  user_set_comm_handle = user_set_comm_handle_;
 
  int initialized;
  MPI_CHECK(MPI_Initialized(&initialized));
 
  if (!initialized) { assert(false); }
 
  if (user_set_comm_handle) {
    MPI_COMM_HANDLE = *((MPI_Comm *)user_comm);
  } else {
    MPI_Comm_dup(MPI_COMM_WORLD, &MPI_COMM_HANDLE);
  }
 
  comm_init(nDim, commDims, rank_from_coords, map_data);
}
 
Communicator::Communicator(Communicator &other, const int *comm_split)
{
  user_set_comm_handle = false;
 
  constexpr int nDim = 4;
 
  quda::CommKey comm_dims_split;
 
  quda::CommKey comm_key_split;
  quda::CommKey comm_color_split;
 
  for (int d = 0; d < nDim; d++) {
    assert(other.comm_dim(d) % comm_split[d] == 0);
    comm_dims_split[d] = other.comm_dim(d) / comm_split[d];
    comm_key_split[d] = other.comm_coord(d) % comm_dims_split[d];
    comm_color_split[d] = other.comm_coord(d) / comm_dims_split[d];
  }
 
  int key = index(nDim, comm_dims_split.data(), comm_key_split.data());
  int color = index(nDim, comm_split, comm_color_split.data());
 
  MPI_CHECK(MPI_Comm_split(other.MPI_COMM_HANDLE, color, key, &MPI_COMM_HANDLE));
  int my_rank_;
  MPI_CHECK(MPI_Comm_rank(MPI_COMM_HANDLE, &my_rank_));
 
  QudaCommsMap func = lex_rank_from_coords_dim_t;
  comm_init(nDim, comm_dims_split.data(), func, comm_dims_split.data());
}
 
Communicator::~Communicator()
{
  comm_finalize();
  if (!user_set_comm_handle) { MPI_Comm_free(&MPI_COMM_HANDLE); }
}
 
void Communicator::comm_gather_hostname(char *hostname_recv_buf)
{
  // determine which GPU this rank will use
  char *hostname = comm_hostname();
  MPI_CHECK(MPI_Allgather(hostname, 128, MPI_CHAR, hostname_recv_buf, 128, MPI_CHAR, MPI_COMM_HANDLE));
}
 
void Communicator::comm_gather_gpuid(int *gpuid_recv_buf)
{
  int gpuid = comm_gpuid();
  MPI_CHECK(MPI_Allgather(&gpuid, 1, MPI_INT, gpuid_recv_buf, 1, MPI_INT, MPI_COMM_HANDLE));
}
 
void Communicator::comm_init(int ndim, const int *dims, QudaCommsMap rank_from_coords, void *map_data)
{
  int initialized;
  MPI_CHECK(MPI_Initialized(&initialized));
 
  if (!initialized) { errorQuda("MPI has not been initialized"); }
 
  MPI_CHECK(MPI_Comm_rank(MPI_COMM_HANDLE, &rank));
  MPI_CHECK(MPI_Comm_size(MPI_COMM_HANDLE, &size));
 
  int grid_size = 1;
  for (int i = 0; i < ndim; i++) { grid_size *= dims[i]; }
  if (grid_size != size) {
    errorQuda("Communication grid size declared via initCommsGridQuda() does not match"
              " total number of MPI ranks (%d != %d)",
              grid_size, size);
  }
 
  comm_init_common(ndim, dims, rank_from_coords, map_data);
}
 
int Communicator::comm_rank(void) { return rank; }
 
int Communicator::comm_size(void) { return size; }
 
MsgHandle *Communicator::comm_declare_send_rank(void *buffer, int rank, int tag, size_t nbytes)
{
  MsgHandle *mh = (MsgHandle *)safe_malloc(sizeof(MsgHandle));
  MPI_CHECK(MPI_Send_init(buffer, nbytes, MPI_BYTE, rank, tag, MPI_COMM_HANDLE, &(mh->request)));
  mh->custom = false;
 
  return mh;
}
 
MsgHandle *Communicator::comm_declare_recv_rank(void *buffer, int rank, int tag, size_t nbytes)
{
  MsgHandle *mh = (MsgHandle *)safe_malloc(sizeof(MsgHandle));
  MPI_CHECK(MPI_Recv_init(buffer, nbytes, MPI_BYTE, rank, tag, MPI_COMM_HANDLE, &(mh->request)));
  mh->custom = false;
 
  return mh;
}
 
MsgHandle *Communicator::comm_declare_send_displaced(void *buffer, const int displacement[], size_t nbytes)
{
  Topology *topo = comm_default_topology();
  int ndim = comm_ndim(topo);
  check_displacement(displacement, ndim);
 
  int rank = comm_rank_displaced(topo, displacement);
 
  int tag = 0;
  for (int i = ndim - 1; i >= 0; i--) tag = tag * 4 * max_displacement + displacement[i] + max_displacement;
  tag = tag >= 0 ? tag : 2 * pow(4 * max_displacement, ndim) + tag;
 
  MsgHandle *mh = (MsgHandle *)safe_malloc(sizeof(MsgHandle));
  MPI_CHECK(MPI_Send_init(buffer, nbytes, MPI_BYTE, rank, tag, MPI_COMM_HANDLE, &(mh->request)));
  mh->custom = false;
 
  return mh;
}
 
MsgHandle *Communicator::comm_declare_receive_displaced(void *buffer, const int displacement[], size_t nbytes)
{
  Topology *topo = comm_default_topology();
  int ndim = comm_ndim(topo);
  check_displacement(displacement, ndim);
 
  int rank = comm_rank_displaced(topo, displacement);
 
  int tag = 0;
  for (int i = ndim - 1; i >= 0; i--) tag = tag * 4 * max_displacement - displacement[i] + max_displacement;
  tag = tag >= 0 ? tag : 2 * pow(4 * max_displacement, ndim) + tag;
 
  MsgHandle *mh = (MsgHandle *)safe_malloc(sizeof(MsgHandle));
  MPI_CHECK(MPI_Recv_init(buffer, nbytes, MPI_BYTE, rank, tag, MPI_COMM_HANDLE, &(mh->request)));
  mh->custom = false;
 
  return mh;
}
 
MsgHandle *Communicator::comm_declare_strided_send_displaced(void *buffer, const int displacement[], size_t blksize,
                                                             int nblocks, size_t stride)
{
  Topology *topo = comm_default_topology();
  int ndim = comm_ndim(topo);
  check_displacement(displacement, ndim);
 
  int rank = comm_rank_displaced(topo, displacement);
 
  int tag = 0;
  for (int i = ndim - 1; i >= 0; i--) tag = tag * 4 * max_displacement + displacement[i] + max_displacement;
  tag = tag >= 0 ? tag : 2 * pow(4 * max_displacement, ndim) + tag;
 
  MsgHandle *mh = (MsgHandle *)safe_malloc(sizeof(MsgHandle));
 
  // create a new strided MPI type
  MPI_CHECK(MPI_Type_vector(nblocks, blksize, stride, MPI_BYTE, &(mh->datatype)));
  MPI_CHECK(MPI_Type_commit(&(mh->datatype)));
  mh->custom = true;
 
  MPI_CHECK(MPI_Send_init(buffer, 1, mh->datatype, rank, tag, MPI_COMM_HANDLE, &(mh->request)));
 
  return mh;
}
 
MsgHandle *Communicator::comm_declare_strided_receive_displaced(void *buffer, const int displacement[], size_t blksize,
                                                                int nblocks, size_t stride)
{
  Topology *topo = comm_default_topology();
  int ndim = comm_ndim(topo);
  check_displacement(displacement, ndim);
 
  int rank = comm_rank_displaced(topo, displacement);
 
  int tag = 0;
  for (int i = ndim - 1; i >= 0; i--) tag = tag * 4 * max_displacement - displacement[i] + max_displacement;
  tag = tag >= 0 ? tag : 2 * pow(4 * max_displacement, ndim) + tag;
 
  MsgHandle *mh = (MsgHandle *)safe_malloc(sizeof(MsgHandle));
 
  // create a new strided MPI type
  MPI_CHECK(MPI_Type_vector(nblocks, blksize, stride, MPI_BYTE, &(mh->datatype)));
  MPI_CHECK(MPI_Type_commit(&(mh->datatype)));
  mh->custom = true;
 
  MPI_CHECK(MPI_Recv_init(buffer, 1, mh->datatype, rank, tag, MPI_COMM_HANDLE, &(mh->request)));
 
  return mh;
}
 
void Communicator::comm_free(MsgHandle *&mh)
{
  MPI_CHECK(MPI_Request_free(&(mh->request)));
  if (mh->custom) MPI_CHECK(MPI_Type_free(&(mh->datatype)));
  host_free(mh);
  mh = nullptr;
}
 
void Communicator::comm_start(MsgHandle *mh) { MPI_CHECK(MPI_Start(&(mh->request))); }
 
void Communicator::comm_wait(MsgHandle *mh) { MPI_CHECK(MPI_Wait(&(mh->request), MPI_STATUS_IGNORE)); }
 
int Communicator::comm_query(MsgHandle *mh)
{
  int query;
  MPI_CHECK(MPI_Test(&(mh->request), &query, MPI_STATUS_IGNORE));
 
  return query;
}
 
void Communicator::comm_allreduce(double *data)
{
  if (!comm_deterministic_reduce()) {
    double recvbuf;
    MPI_CHECK(MPI_Allreduce(data, &recvbuf, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_HANDLE));
    *data = recvbuf;
  } else {
    const size_t n = comm_size();
    double *recv_buf = (double *)safe_malloc(n * sizeof(double));
    MPI_CHECK(MPI_Allgather(data, 1, MPI_DOUBLE, recv_buf, 1, MPI_DOUBLE, MPI_COMM_HANDLE));
    *data = deterministic_reduce(recv_buf, n);
    host_free(recv_buf);
  }
}
 
void Communicator::comm_allreduce_max(double *data)
{
  double recvbuf;
  MPI_CHECK(MPI_Allreduce(data, &recvbuf, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_HANDLE));
  *data = recvbuf;
}
 
void Communicator::comm_allreduce_min(double *data)
{
  double recvbuf;
  MPI_CHECK(MPI_Allreduce(data, &recvbuf, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_HANDLE));
  *data = recvbuf;
}
 
void Communicator::comm_allreduce_array(double *data, size_t size)
{
  if (!comm_deterministic_reduce()) {
    double *recvbuf = new double[size];
    MPI_CHECK(MPI_Allreduce(data, recvbuf, size, MPI_DOUBLE, MPI_SUM, MPI_COMM_HANDLE));
    memcpy(data, recvbuf, size * sizeof(double));
    delete[] recvbuf;
  } else {
    size_t n = comm_size();
    double *recv_buf = new double[size * n];
    MPI_CHECK(MPI_Allgather(data, size, MPI_DOUBLE, recv_buf, size, MPI_DOUBLE, MPI_COMM_HANDLE));
 
    double *recv_trans = new double[size * n];
    for (size_t i = 0; i < n; i++) {
      for (size_t j = 0; j < size; j++) { recv_trans[j * n + i] = recv_buf[i * size + j]; }
    }
 
    for (size_t i = 0; i < size; i++) { data[i] = deterministic_reduce(recv_trans + i * n, n); }
 
    delete[] recv_buf;
    delete[] recv_trans;
  }
}
 
void Communicator::comm_allreduce_max_array(double *data, size_t size)
{
  double *recvbuf = new double[size];
  MPI_CHECK(MPI_Allreduce(data, recvbuf, size, MPI_DOUBLE, MPI_MAX, MPI_COMM_HANDLE));
  memcpy(data, recvbuf, size * sizeof(double));
  delete[] recvbuf;
}
 
void Communicator::comm_allreduce_int(int *data)
{
  int recvbuf;
  MPI_CHECK(MPI_Allreduce(data, &recvbuf, 1, MPI_INT, MPI_SUM, MPI_COMM_HANDLE));
  *data = recvbuf;
}
 
void Communicator::comm_allreduce_xor(uint64_t *data)
{
  if (sizeof(uint64_t) != sizeof(unsigned long)) errorQuda("unsigned long is not 64-bit");
  uint64_t recvbuf;
  MPI_CHECK(MPI_Allreduce(data, &recvbuf, 1, MPI_UNSIGNED_LONG, MPI_BXOR, MPI_COMM_HANDLE));
  *data = recvbuf;
}
 
void Communicator::comm_broadcast(void *data, size_t nbytes)
{
  MPI_CHECK(MPI_Bcast(data, (int)nbytes, MPI_BYTE, 0, MPI_COMM_HANDLE));
}
 
void Communicator::comm_barrier(void) { MPI_CHECK(MPI_Barrier(MPI_COMM_HANDLE)); }
 
void Communicator::comm_abort_(int status) { MPI_Abort(MPI_COMM_WORLD, status); }
 
int Communicator::comm_rank_global()
{
  int rank;
  MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &rank));
  return rank;
}