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 "local_range.h"
#include "log.h"
#include "sort.h"
#include "types.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <complex>
#include <concepts>
#include <cstdint>
#include <iterator>
#include <numeric>
#include <ranges>
#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 : int
{
  consensus_pcx = 1200,
  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 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));

std::pair<std::vector<int>, std::vector<int>>
compute_graph_edges_nbx(MPI_Comm comm, std::span<const int> edges0, int tag0,
                        std::span<const int> edges1, int tag1);

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

template <std::ranges::contiguous_range U>
std::vector<std::ranges::range_value_t<U>>
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 <std::ranges::contiguous_range U>
std::vector<std::ranges::range_value_t<U>>
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>
MPI_Datatype mpi_datatype()
{
  if constexpr (std::same_as<T, float>)
    return MPI_FLOAT;
  else if constexpr (std::same_as<T, double>)
    return MPI_DOUBLE;
  else if constexpr (std::same_as<T, std::complex<float>>)
    return MPI_C_FLOAT_COMPLEX;
  else if constexpr (std::same_as<T, std::complex<double>>)
    return MPI_C_DOUBLE_COMPLEX;
  else if constexpr (std::same_as<T, std::int8_t>)
    return MPI_INT8_T;
  else if constexpr (std::same_as<T, std::int16_t>)
    return MPI_INT16_T;
  else if constexpr (std::same_as<T, std::int32_t>)
    return MPI_INT32_T;
  else if constexpr (std::same_as<T, std::int64_t>)
    return MPI_INT64_T;
  else if constexpr (std::same_as<T, std::uint8_t>)
    return MPI_UINT8_T;
  else if constexpr (std::same_as<T, std::uint16_t>)
    return MPI_UINT16_T;
  else if constexpr (std::same_as<T, std::uint32_t>)
    return MPI_UINT32_T;
  else if constexpr (std::same_as<T, std::uint64_t>)
    return MPI_UINT64_T;
  else
    static_assert(dependent_false<T>::value,
                  "No MPI datatype registered for this type.");
}

template <typename T>
MPI_Datatype mpi_t = mpi_datatype<T>();
 
//---------------------------------------------------------------------------
namespace impl
{
std::tuple<std::vector<int>, std::vector<std::int32_t>,
           std::vector<std::int32_t>>
postoffice_plan(int size, int rank, std::int32_t shape0_local,
                std::int64_t shape0, std::int64_t rank_offset);

template <std::ranges::contiguous_range U>
std::pair<std::vector<std::int32_t>, std::vector<std::ranges::range_value_t<U>>>
postoffice_exchange(MPI_Comm comm, const U& x,
                    std::array<std::int64_t, 2> shape, std::int64_t rank_offset,
                    std::span<const int> dest,
                    std::span<const std::int32_t> num_items_per_dest0,
                    std::span<const std::int32_t> pos_to_neigh_rank,
                    std::span<const int> src)
{
  using T = std::ranges::range_value_t<U>;
 
  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];
 
  // 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> num_items_per_dest(num_items_per_dest0.begin(),
                                               num_items_per_dest0.end());
  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);
 
  // Convert to local indices
  const std::int64_t r0 = common::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](std::int64_t idx) { return idx - r0; });
 
  return {index_local, recv_buffer_data};
}
} // namespace impl
 
template <std::ranges::contiguous_range U>
std::pair<std::vector<std::int32_t>, std::vector<std::ranges::range_value_t<U>>>
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());
  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)");
 
  const int size = dolfinx::MPI::size(comm);
  const int rank = dolfinx::MPI::rank(comm);
  auto [dest, num_items_per_dest, pos_to_neigh_rank]
      = impl::postoffice_plan(size, rank, shape0_local, shape[0], rank_offset);
 
  // 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: {}",
      src.size());
 
  auto result
      = impl::postoffice_exchange(comm, x, shape, rank_offset, dest,
                                  num_items_per_dest, pos_to_neigh_rank, src);
  spdlog::debug("Completed send data to post offices.");
  return result;
}
//---------------------------------------------------------------------------
template <std::ranges::contiguous_range U>
std::vector<std::ranges::range_value_t<U>>
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 = std::ranges::range_value_t<U>;
 
  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, and 1. determine which post
  //    office ranks hold the data I need (indices) -- these are
  //    independent local computations, so the two NBX consensus
  //    rounds they each need (below) are run concurrently in a single
  //    overlapped round rather than back-to-back.
 
  auto [send_dest, num_items_per_send_dest, pos_to_neigh_rank]
      = impl::postoffice_plan(size, rank, shape0_local, shape[0], rank_offset);
 
  // Build (src, global index, position) for each entry in 'indices'
  // not held locally, then sort. Locally-held entries are read
  // directly below -- skipping them here avoids a wasted round trip.
  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::int64_t idx = indices[i];
    if (idx >= rank_offset and idx < rank_offset + shape0_local)
      continue;
    if (int src = dolfinx::MPI::index_owner(size, idx, shape[0]); src != rank)
      src_to_index.push_back({src, idx, i});
  }
 
  // Radix sort on the rank alone (not the full tuple) -- only
  // grouping by rank matters below; order within a group doesn't.
  {
    std::vector<std::int32_t> perm(src_to_index.size());
    std::iota(perm.begin(), perm.end(), 0);
    dolfinx::radix_sort(perm, [&src_to_index](std::int32_t i)
                        { return std::get<0>(src_to_index[i]); });
    std::vector<std::tuple<int, std::int64_t, std::int32_t>> sorted(
        src_to_index.size());
    for (std::size_t i = 0; i < perm.size(); ++i)
      sorted[i] = src_to_index[perm[i]];
    src_to_index = std::move(sorted);
  }
 
  // 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<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::ranges::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::ranges::distance(it, it1));
      it = it1;
    }
  }
 
  // Determine, in one overlapped NBX round, (0) the post office ranks
  // that hold data for me to receive (post office send round) and (1)
  // the 'delivery' destination ranks that want data from me (my
  // request round)
  auto [post_src, dest] = dolfinx::MPI::compute_graph_edges_nbx(
      comm, send_dest, static_cast<int>(tag::consensus_nbx), src,
      static_cast<int>(tag::consensus_nbx) + 1);
  spdlog::info(
      "Neighbourhood destination ranks from post office in "
      "distribute_data (rank, num dests, num dests/mpi_size): {}, {}, {}",
      rank, dest.size(), static_cast<double>(dest.size()) / size);
 
  // Send receive x data to post office (only for rows that need to be
  // communicated)
  auto [post_indices, post_x] = impl::postoffice_exchange(
      comm, x, {shape[0], shape[1]}, rank_offset, send_dest,
      num_items_per_send_dest, pos_to_neigh_rank, post_src);
  assert(post_indices.size() == post_x.size() / shape[1]);
 
  // 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() == static_cast<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
      = common::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] < static_cast<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::int32_t i = 0; i < recv_disp.back(); ++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] * i));
    }
    else
    {
      // Take from my 'post bag'
      std::int64_t 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] * i));
    }
  }
 
  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
      std::int64_t 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 (std::int32_t pos = index_pos_to_buffer[i]; pos != -1)
    {
      // In my received post: index_pos_to_buffer[i] != -1 iff
      // index_owner would say this rank isn't the owner -- avoids
      // recomputing it.
      std::copy_n(std::next(recv_buffer_data.begin(), shape[1] * pos), shape[1],
                  std::next(x_new.begin(), shape[1] * i));
    }
    else
    {
      // In my post office bag
      std::int64_t local_index = index - postoffice_range[0];
      std::int32_t bag_pos = post_indices_map[local_index];
      assert(bag_pos != -1);
      std::copy_n(std::next(post_x.begin(), shape[1] * bag_pos), shape[1],
                  std::next(x_new.begin(), shape[1] * i));
    }
  }
 
  return x_new;
}
//---------------------------------------------------------------------------
template <std::ranges::contiguous_range U>
std::vector<std::ranges::range_value_t<U>>
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;
 
  std::int64_t shape0 = 0;
  int 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 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