MPI Code Example: Latency-hiding Variable-length Messages on a Toroidal Grid

#include <iostream>
#include <sstream>
#include <vector>
#include <unordered_map>
#include <cmath>
#include <algorithm>

#include "mpi.h"

// some slick book-keeping for cardinal directionality
struct Cardi {

  // cardinal directions
  enum Dir { N, S, E, W, NumDirs };

  // transforms
  static constexpr Dir Opp[] { S, N, W, E };
  static constexpr Dir Cw[] { E, W, S, N };
  static constexpr Dir Ccw[] { W, E, N, S };
  static constexpr int Dim[] { 1, 1, 0, 0 };
  static constexpr int Dx[] { 0, 0, 1, -1 };
  static constexpr int Dy[] { 1, -1, 0, 0 };
  static constexpr int Dd[] { 1, -1, 1, -1 };

};

int main(int argc, char* argv[])
{

  MPI_Init(&argc, &argv);

  // make a square grid out of available processes
  MPI_Comm comm_grid;
  const int ndim = 2;
  const int nprocs = [&](){
    int res;
    MPI_Comm_size(
      MPI_COMM_WORLD,
      &res
    );
    return res;
  }(); // use immediately-invoked lambda for consty goodness
  const std::array<int, ndim> dimension{ {
    static_cast<int>(std::sqrt(nprocs)),
    static_cast<int>(std::sqrt(nprocs))
  } };
  const std::array<int, ndim> periodic{ {1, 1} }; // toroidal wraparound

  MPI_Cart_create(
    MPI_COMM_WORLD,
    ndim,
    dimension.data(),
    periodic.data(),
    1, // allow reordering of world ranks
    &comm_grid
  );

  // get world rank info
  // which MPI proc am I in the ?
  const int w_rank = [&](){
    int res;
    MPI_Comm_rank(
      MPI_COMM_WORLD,
      &res
    );
    return res;
  }(); // use immediately-invoked lambda for consty goodness

  // get comm_grid info
  // where am I (x,y) in the grid?
  const std::array<int, ndim> coords = [&](){
    std::array<int, ndim> res;
    MPI_Cart_coords(
      comm_grid,
      w_rank,
      ndim,
      res.data()
    );
    return res;
  }(); // use immediately-invoked lambda for consty goodness
  // what is my rank in the grid?
  const int g_rank = [&](){
    int res;
    MPI_Cart_rank(
      comm_grid,
      coords.data(),
      &res
    );
    return res;
  }();

  // calculate neighboring ranks
  // store in array, indexed by cardinal direction
  const std::array<int, Cardi::Dir::NumDirs> neighs = [&](){
    std::array<int, Cardi::Dir::NumDirs> res;
    for (size_t i = 0; i < Cardi::Dir::NumDirs; ++i) {
      MPI_Cart_shift(
        comm_grid,
        Cardi::Dim[i],
        Cardi::Dd[i],
        const_cast<int *>(&g_rank),
        &res[i]
      );

    }
    return res;
  }(); // use immediately-invoked lambda for consty goodness

  // GOAL:
  // everyone broadcasts a unique character to each cartesian neighbor
  // send 1 byte North, 2 bytes South, 3 bytes East, and 4 bytes West
  const unsigned char id = 65 + g_rank;

  // prepare for sends
  const std::vector<unsigned char> send_data(Cardi::Dir::NumDirs, id);
  std::array<MPI_Request, Cardi::Dir::NumDirs> send_requests;

  // post non-blocking send requests to each neighbor
  for (size_t i = 0; i < Cardi::Dir::NumDirs; ++i) {
    MPI_Isend(
      send_data.data(),
      i + 1, // message length
      MPI_BYTE, // message data type
      neighs[i], // message destination
      i, // tag represents outgoing direction
      comm_grid, // use the grid communicator group
      &send_requests[i] // use later to check up if send completed
    );
  }

  // whee!
  // DO COMPUTE WORK HERE
  // while messages are en route

  // data structure to store incoming message data
  // incoming direction -> result
  std::unordered_map<int, std::vector<unsigned char>> received;

  // receive messages from each neighbor, in any order
  for (size_t i = 0; i < Cardi::Dir::NumDirs; ++i) {

    // block until we get any message
    MPI_Status status;
    MPI_Probe(
      MPI_ANY_SOURCE,
      MPI_ANY_TAG,
      comm_grid,
      &status
    );

    // how long is the message, who sent it, and from what direction
    const int msg_len = [&](){
      int res;
      MPI_Get_count(
        &status,
        MPI_BYTE,
        &res
      );
      return res;
    }(); // use immediately-invoked lambda for consty goodness
    const int msg_source = status.MPI_SOURCE;
    const int msg_tag = status.MPI_TAG;
    const int incoming_direction = Cardi::Opp[msg_tag];

    // make space for incoming data
    received.emplace(
      std::piecewise_construct,
      std::forward_as_tuple(incoming_direction),
      std::forward_as_tuple(msg_len)
    );

    // copy down data from message
    MPI_Recv(
      received.at(incoming_direction).data(), // where to put it
      msg_len, // how long is the message
      MPI_BYTE, // message data type
      msg_source, // who sent it, makes sure we grab the right message!
      msg_tag, // what it's tagged, makes sure we grab the right message!
      comm_grid, // use the grid communicator group
      MPI_STATUS_IGNORE // yup
    );

  }

  // wait for all sends to complete
  MPI_Waitall(
    send_requests.size(),
    send_requests.data(),
    MPI_STATUSES_IGNORE
  );

  // print my results!
  std::ostringstream oss;
  // who am I...
  oss << "rank " << g_rank << ", id " << id
    << " @ (" << coords[0] << ", " << coords[1] << ")";
  // ..and what did I get?
  for (const auto & [k, v] : received) {
    oss << " [" << k << ": ";
    std::for_each(
      std::begin(v),
      std::end(v),
      [&oss](auto c){ oss << c; }
    );
    oss << "]";
  }
  // print out my info on one line in one go!
  std::cout << oss.str() << std::endl;

  MPI_Finalize();

  return 0;
}