MPI.h

// Copyright (C) 2007-2023 Magnus Vikstrøm, Garth N. Wells and Paul T. Kühner
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Timer.h"
#include "log.h"
#include "types.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <complex>
#include <cstdint>
#include <numeric>
#include <set>
#include <span>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
#define MPICH_IGNORE_CXX_SEEK 1
#include <mpi.h>

namespace dolfinx::MPI
{
enum class tag : std::uint16_t
{
  consensus_pcx = 1200,
  consensus_pex = 1201,
  consensus_nbx = 1202,
};

class Comm
{
public:
  explicit Comm(MPI_Comm comm, bool duplicate = true);

  Comm(const Comm& comm) noexcept;

  Comm(Comm&& comm) noexcept;
 
  // Disable copy assignment operator
  Comm& operator=(const Comm& comm) = delete;

  Comm& operator=(Comm&& comm) noexcept;

  ~Comm();

  MPI_Comm comm() const noexcept;
 
private:
  // MPI communicator
  MPI_Comm _comm;
};

int rank(MPI_Comm comm);

int size(MPI_Comm comm);

void check_error(MPI_Comm comm, int code);

constexpr std::array<std::int64_t, 2> local_range(int rank, std::int64_t N,
                                                  int size)
{
  assert(rank >= 0);
  assert(N >= 0);
  assert(size > 0);
 
  // Compute number of items per rank and remainder
  const std::int64_t n = N / size;
  const std::int64_t r = N % size;
 
  // Compute local range
  if (rank < r)
    return {rank * (n + 1), rank * (n + 1) + n + 1};
  else
    return {rank * n + r, rank * n + r + n};
}

constexpr int index_owner(int size, std::size_t index, std::size_t N)
{
  assert(index < N);
 
  // Compute number of items per rank and remainder
  const std::size_t n = N / size;
  const std::size_t r = N % size;
 
  if (index < r * (n + 1))
  {
    // First r ranks own n + 1 indices
    return index / (n + 1);
  }
  else
  {
    // Remaining ranks own n indices
    return r + (index - r * (n + 1)) / n;
  }
}

std::vector<int> compute_graph_edges_pcx(MPI_Comm comm,
                                         std::span<const int> edges);

std::vector<int>
compute_graph_edges_nbx(MPI_Comm comm, std::span<const int> edges,
                        int tag = static_cast<int>(tag::consensus_nbx));

template <typename U>
std::pair<std::vector<std::int32_t>,
          std::vector<typename std::remove_reference_t<typename U::value_type>>>
distribute_to_postoffice(MPI_Comm comm, const U& x,
                         std::array<std::int64_t, 2> shape,
                         std::int64_t rank_offset);

template <typename U>
std::vector<typename std::remove_reference_t<typename U::value_type>>
distribute_from_postoffice(MPI_Comm comm, std::span<const std::int64_t> indices,
                           const U& x, std::array<std::int64_t, 2> shape,
                           std::int64_t rank_offset);

template <typename U>
std::vector<typename std::remove_reference_t<typename U::value_type>>
distribute_data(MPI_Comm comm0, std::span<const std::int64_t> indices,
                MPI_Comm comm1, const U& x, int shape1);
 
template <typename T>
struct dependent_false : std::false_type
{
};


template <typename T>
struct mpi_type_mapping;

template <typename T>
MPI_Datatype mpi_t = mpi_type_mapping<T>::type;

#define MAP_TO_MPI_TYPE(cpp_t, mpi_t)                                          \
  template <>                                                                  \
  struct mpi_type_mapping<cpp_t>                                               \
  {                                                                            \
    static inline MPI_Datatype type = mpi_t;                                   \
  };

MAP_TO_MPI_TYPE(float, MPI_FLOAT)
MAP_TO_MPI_TYPE(double, MPI_DOUBLE)
MAP_TO_MPI_TYPE(std::complex<float>, MPI_C_FLOAT_COMPLEX)
MAP_TO_MPI_TYPE(std::complex<double>, MPI_C_DOUBLE_COMPLEX)
MAP_TO_MPI_TYPE(std::int8_t, MPI_INT8_T)
MAP_TO_MPI_TYPE(std::int16_t, MPI_INT16_T)
MAP_TO_MPI_TYPE(std::int32_t, MPI_INT32_T)
MAP_TO_MPI_TYPE(std::int64_t, MPI_INT64_T)
MAP_TO_MPI_TYPE(std::uint8_t, MPI_UINT8_T)
MAP_TO_MPI_TYPE(std::uint16_t, MPI_UINT16_T)
MAP_TO_MPI_TYPE(std::uint32_t, MPI_UINT32_T)
MAP_TO_MPI_TYPE(std::uint64_t, MPI_UINT64_T)
 
//---------------------------------------------------------------------------
template <typename U>
std::pair<std::vector<std::int32_t>,
          std::vector<typename std::remove_reference_t<typename U::value_type>>>
distribute_to_postoffice(MPI_Comm comm, const U& x,
                         std::array<std::int64_t, 2> shape,
                         std::int64_t rank_offset)
{
  assert(rank_offset >= 0 or x.empty());
  using T = typename std::remove_reference_t<typename U::value_type>;
 
  const int size = dolfinx::MPI::size(comm);
  const int rank = dolfinx::MPI::rank(comm);
  assert(x.size() % shape[1] == 0);
  const std::int32_t shape0_local = x.size() / shape[1];
 
  spdlog::debug("Sending data to post offices (distribute_to_postoffice)");
 
  // Post office ranks will receive data from this rank
  std::vector<int> row_to_dest(shape0_local);
  for (std::int32_t i = 0; i < shape0_local; ++i)
  {
    int dest = MPI::index_owner(size, i + rank_offset, shape[0]);
    row_to_dest[i] = dest;
  }
 
  // Build list of (dest, positions) for each row that doesn't belong to
  // this rank, then sort
  std::vector<std::array<std::int32_t, 2>> dest_to_index;
  dest_to_index.reserve(shape0_local);
  for (std::int32_t i = 0; i < shape0_local; ++i)
  {
    std::size_t idx = i + rank_offset;
    if (int dest = MPI::index_owner(size, idx, shape[0]); dest != rank)
      dest_to_index.push_back({dest, i});
  }
  std::ranges::sort(dest_to_index);
 
  // Build list of neighbour src ranks and count number of items (rows
  // of x) to receive from each src post office (by neighbourhood rank)
  std::vector<int> dest;
  std::vector<std::int32_t> num_items_per_dest,
      pos_to_neigh_rank(shape0_local, -1);
  {
    auto it = dest_to_index.begin();
    while (it != dest_to_index.end())
    {
      const int neigh_rank = dest.size();
 
      // Store global rank
      dest.push_back((*it)[0]);
 
      // Find iterator to next global rank
      auto it1
          = std::find_if(it, dest_to_index.end(),
                         [r = dest.back()](auto& idx) { return idx[0] != r; });
 
      // Store number of items for current rank
      num_items_per_dest.push_back(std::distance(it, it1));
 
      // Map from local x index to local destination rank
      for (auto e = it; e != it1; ++e)
        pos_to_neigh_rank[(*e)[1]] = neigh_rank;
 
      // Advance iterator
      it = it1;
    }
  }
 
  // Determine source ranks
  const std::vector<int> src = MPI::compute_graph_edges_nbx(comm, dest);
  spdlog::info(
      "Number of neighbourhood source ranks in distribute_to_postoffice: {}",
      static_cast<int>(src.size()));
 
  // Create neighbourhood communicator for sending data to post offices
  MPI_Comm neigh_comm;
  int err = MPI_Dist_graph_create_adjacent(
      comm, src.size(), src.data(), MPI_UNWEIGHTED, dest.size(), dest.data(),
      MPI_UNWEIGHTED, MPI_INFO_NULL, false, &neigh_comm);
  dolfinx::MPI::check_error(comm, err);
 
  // Compute send displacements
  std::vector<std::int32_t> send_disp{0};
  std::partial_sum(num_items_per_dest.begin(), num_items_per_dest.end(),
                   std::back_inserter(send_disp));
 
  // Pack send buffers
  std::vector<T> send_buffer_data(shape[1] * send_disp.back());
  std::vector<std::int64_t> send_buffer_index(send_disp.back());
  {
    std::vector<std::int32_t> send_offsets = send_disp;
    for (std::int32_t i = 0; i < shape0_local; ++i)
    {
      if (int neigh_dest = pos_to_neigh_rank[i]; neigh_dest != -1)
      {
        std::size_t pos = send_offsets[neigh_dest];
        send_buffer_index[pos] = i + rank_offset;
        std::copy_n(std::next(x.begin(), i * shape[1]), shape[1],
                    std::next(send_buffer_data.begin(), shape[1] * pos));
        ++send_offsets[neigh_dest];
      }
    }
  }
 
  // Send number of items to post offices (destination) that I will be
  // sending
  std::vector<int> num_items_recv(src.size());
  num_items_per_dest.reserve(1);
  num_items_recv.reserve(1);
  err = MPI_Neighbor_alltoall(num_items_per_dest.data(), 1, MPI_INT,
                              num_items_recv.data(), 1, MPI_INT, neigh_comm);
  dolfinx::MPI::check_error(comm, err);
 
  // Prepare receive displacement and buffers
  std::vector<std::int32_t> recv_disp(num_items_recv.size() + 1, 0);
  std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
                   std::next(recv_disp.begin()));
 
  // Send/receive global indices
  std::vector<std::int64_t> recv_buffer_index(recv_disp.back());
  err = MPI_Neighbor_alltoallv(
      send_buffer_index.data(), num_items_per_dest.data(), send_disp.data(),
      MPI_INT64_T, recv_buffer_index.data(), num_items_recv.data(),
      recv_disp.data(), MPI_INT64_T, neigh_comm);
  dolfinx::MPI::check_error(comm, err);
 
  // Send/receive data (x)
  MPI_Datatype compound_type;
  MPI_Type_contiguous(shape[1], dolfinx::MPI::mpi_t<T>, &compound_type);
  MPI_Type_commit(&compound_type);
  std::vector<T> recv_buffer_data(shape[1] * recv_disp.back());
  err = MPI_Neighbor_alltoallv(
      send_buffer_data.data(), num_items_per_dest.data(), send_disp.data(),
      compound_type, recv_buffer_data.data(), num_items_recv.data(),
      recv_disp.data(), compound_type, neigh_comm);
  dolfinx::MPI::check_error(comm, err);
  err = MPI_Type_free(&compound_type);
  dolfinx::MPI::check_error(comm, err);
  err = MPI_Comm_free(&neigh_comm);
  dolfinx::MPI::check_error(comm, err);
 
  spdlog::debug("Completed send data to post offices.");
 
  // Convert to local indices
  const std::int64_t r0 = MPI::local_range(rank, shape[0], size)[0];
  std::vector<std::int32_t> index_local(recv_buffer_index.size());
  std::ranges::transform(recv_buffer_index, index_local.begin(),
                         [r0](auto idx) { return idx - r0; });
 
  return {index_local, recv_buffer_data};
}
//---------------------------------------------------------------------------
template <typename U>
std::vector<typename std::remove_reference_t<typename U::value_type>>
distribute_from_postoffice(MPI_Comm comm, std::span<const std::int64_t> indices,
                           const U& x, std::array<std::int64_t, 2> shape,
                           std::int64_t rank_offset)
{
  assert(rank_offset >= 0 or x.empty());
  using T = typename std::remove_reference_t<typename U::value_type>;
 
  common::Timer timer("Distribute row-wise data (scalable)");
  assert(shape[1] > 0);
 
  const int size = dolfinx::MPI::size(comm);
  const int rank = dolfinx::MPI::rank(comm);
  assert(x.size() % shape[1] == 0);
  const std::int64_t shape0_local = x.size() / shape[1];
 
  // 0. Send x data to/from post offices
 
  // Send receive x data to post office (only for rows that need to be
  // communicated)
  auto [post_indices, post_x] = dolfinx::MPI::distribute_to_postoffice(
      comm, x, {shape[0], shape[1]}, rank_offset);
  assert(post_indices.size() == post_x.size() / shape[1]);
 
  // 1. Send request to post office ranks for data
 
  // Build list of (src, global index, global, index position) for each
  // entry in 'indices' that doesn't belong to this rank, then sort
  std::vector<std::tuple<int, std::int64_t, std::int32_t>> src_to_index;
  for (std::size_t i = 0; i < indices.size(); ++i)
  {
    std::size_t idx = indices[i];
    if (int src = dolfinx::MPI::index_owner(size, idx, shape[0]); src != rank)
      src_to_index.push_back({src, idx, i});
  }
  std::ranges::sort(src_to_index);
 
  // Build list is neighbour src ranks and count number of items (rows
  // of x) to receive from each src post office (by neighbourhood rank)
  std::vector<std::int32_t> num_items_per_src;
  std::vector<int> src;
  {
    auto it = src_to_index.begin();
    while (it != src_to_index.end())
    {
      src.push_back(std::get<0>(*it));
      auto it1
          = std::find_if(it, src_to_index.end(), [r = src.back()](auto& idx)
                         { return std::get<0>(idx) != r; });
      num_items_per_src.push_back(std::distance(it, it1));
      it = it1;
    }
  }
 
  // Determine 'delivery' destination ranks (ranks that want data from
  // me)
  const std::vector<int> dest
      = dolfinx::MPI::compute_graph_edges_nbx(comm, src);
  spdlog::info(
      "Neighbourhood destination ranks from post office in "
      "distribute_data (rank, num dests, num dests/mpi_size): {}, {}, {}",
      rank, static_cast<int>(dest.size()),
      static_cast<double>(dest.size()) / size);
 
  // Create neighbourhood communicator for sending data to post offices
  // (src), and receiving data form my send my post office
  MPI_Comm neigh_comm0;
  int err = MPI_Dist_graph_create_adjacent(
      comm, dest.size(), dest.data(), MPI_UNWEIGHTED, src.size(), src.data(),
      MPI_UNWEIGHTED, MPI_INFO_NULL, false, &neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  // Communicate number of requests to each source
  std::vector<int> num_items_recv(dest.size());
  num_items_per_src.reserve(1);
  num_items_recv.reserve(1);
  err = MPI_Neighbor_alltoall(num_items_per_src.data(), 1, MPI_INT,
                              num_items_recv.data(), 1, MPI_INT, neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  // Prepare send/receive displacements
  std::vector<std::int32_t> send_disp{0};
  std::partial_sum(num_items_per_src.begin(), num_items_per_src.end(),
                   std::back_inserter(send_disp));
  std::vector<std::int32_t> recv_disp = {0};
  std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
                   std::back_inserter(recv_disp));
 
  // Pack my requested indices (global) in send buffer ready to send to
  // post offices
  assert(send_disp.back() == (int)src_to_index.size());
  std::vector<std::int64_t> send_buffer_index(src_to_index.size());
  std::ranges::transform(src_to_index, send_buffer_index.begin(),
                         [](auto x) { return std::get<1>(x); });
 
  // Prepare the receive buffer
  std::vector<std::int64_t> recv_buffer_index(recv_disp.back());
  err = MPI_Neighbor_alltoallv(
      send_buffer_index.data(), num_items_per_src.data(), send_disp.data(),
      MPI_INT64_T, recv_buffer_index.data(), num_items_recv.data(),
      recv_disp.data(), MPI_INT64_T, neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  err = MPI_Comm_free(&neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  // 2. Send data (rows of x) from post office back to requesting ranks
  //    (transpose of the preceding communication pattern operation)
 
  // Build map from local index to post_indices position. Set to -1 for
  // data that was already on this rank and was therefore was not
  // sent/received via a postoffice.
  const std::array<std::int64_t, 2> postoffice_range
      = dolfinx::MPI::local_range(rank, shape[0], size);
  std::vector<std::int32_t> post_indices_map(
      postoffice_range[1] - postoffice_range[0], -1);
  for (std::size_t i = 0; i < post_indices.size(); ++i)
  {
    assert(post_indices[i] < (int)post_indices_map.size());
    post_indices_map[post_indices[i]] = i;
  }
 
  // Build send buffer
  std::vector<T> send_buffer_data(shape[1] * recv_disp.back());
  for (std::size_t p = 0; p < recv_disp.size() - 1; ++p)
  {
    int offset = recv_disp[p];
    for (std::int32_t i = recv_disp[p]; i < recv_disp[p + 1]; ++i)
    {
      std::int64_t index = recv_buffer_index[i];
      if (index >= rank_offset and index < (rank_offset + shape0_local))
      {
        // I already had this index before any communication
        std::int32_t local_index = index - rank_offset;
        std::copy_n(std::next(x.begin(), shape[1] * local_index), shape[1],
                    std::next(send_buffer_data.begin(), shape[1] * offset));
      }
      else
      {
        // Take from my 'post bag'
        auto local_index = index - postoffice_range[0];
        std::int32_t pos = post_indices_map[local_index];
        assert(pos != -1);
        std::copy_n(std::next(post_x.begin(), shape[1] * pos), shape[1],
                    std::next(send_buffer_data.begin(), shape[1] * offset));
      }
 
      ++offset;
    }
  }
 
  err = MPI_Dist_graph_create_adjacent(
      comm, src.size(), src.data(), MPI_UNWEIGHTED, dest.size(), dest.data(),
      MPI_UNWEIGHTED, MPI_INFO_NULL, false, &neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  MPI_Datatype compound_type0;
  MPI_Type_contiguous(shape[1], dolfinx::MPI::mpi_t<T>, &compound_type0);
  MPI_Type_commit(&compound_type0);
 
  std::vector<T> recv_buffer_data(shape[1] * send_disp.back());
  err = MPI_Neighbor_alltoallv(
      send_buffer_data.data(), num_items_recv.data(), recv_disp.data(),
      compound_type0, recv_buffer_data.data(), num_items_per_src.data(),
      send_disp.data(), compound_type0, neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  err = MPI_Type_free(&compound_type0);
  dolfinx::MPI::check_error(comm, err);
  err = MPI_Comm_free(&neigh_comm0);
  dolfinx::MPI::check_error(comm, err);
 
  std::vector<std::int32_t> index_pos_to_buffer(indices.size(), -1);
  for (std::size_t i = 0; i < src_to_index.size(); ++i)
    index_pos_to_buffer[std::get<2>(src_to_index[i])] = i;
 
  // Extra data to return
  std::vector<T> x_new(shape[1] * indices.size());
  for (std::size_t i = 0; i < indices.size(); ++i)
  {
    const std::int64_t index = indices[i];
    if (index >= rank_offset and index < (rank_offset + shape0_local))
    {
      // Had data from the start in x
      auto local_index = index - rank_offset;
      std::copy_n(std::next(x.begin(), shape[1] * local_index), shape[1],
                  std::next(x_new.begin(), shape[1] * i));
    }
    else
    {
      if (int src = dolfinx::MPI::index_owner(size, index, shape[0]);
          src == rank)
      {
        // In my post office bag
        auto local_index = index - postoffice_range[0];
        std::int32_t pos = post_indices_map[local_index];
        assert(pos != -1);
        std::copy_n(std::next(post_x.begin(), shape[1] * pos), shape[1],
                    std::next(x_new.begin(), shape[1] * i));
      }
      else
      {
        // In my received post
        std::int32_t pos = index_pos_to_buffer[i];
        assert(pos != -1);
        std::copy_n(std::next(recv_buffer_data.begin(), shape[1] * pos),
                    shape[1], std::next(x_new.begin(), shape[1] * i));
      }
    }
  }
 
  return x_new;
}
//---------------------------------------------------------------------------
template <typename U>
std::vector<typename std::remove_reference_t<typename U::value_type>>
distribute_data(MPI_Comm comm0, std::span<const std::int64_t> indices,
                MPI_Comm comm1, const U& x, int shape1)
{
  assert(shape1 > 0);
  assert(x.size() % shape1 == 0);
  const std::int64_t shape0_local = x.size() / shape1;
 
  int err;
  std::int64_t shape0 = 0;
  err = MPI_Allreduce(&shape0_local, &shape0, 1, MPI_INT64_T, MPI_SUM, comm0);
  dolfinx::MPI::check_error(comm0, err);
 
  std::int64_t rank_offset = -1;
  if (comm1 != MPI_COMM_NULL)
  {
    rank_offset = 0;
    err = MPI_Exscan(&shape0_local, &rank_offset, 1,
                     dolfinx::MPI::mpi_t<std::int64_t>, MPI_SUM, comm1);
    dolfinx::MPI::check_error(comm1, err);
  }
  else
  {
    rank_offset = -1;
    if (!x.empty())
      throw std::runtime_error("Non-empty data on null MPI communicator");
  }
 
  return distribute_from_postoffice(comm0, indices, x, {shape0, shape1},
                                    rank_offset);
}
//---------------------------------------------------------------------------
 
} // namespace dolfinx::MPI