Loading...
Searching...
No Matches
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
dolfinx::MPI::Comm::Comm
Comm(MPI_Comm comm, bool duplicate=true)
Duplicate communicator and wrap duplicate.
Definition MPI.cpp:12
dolfinx::MPI::Comm::~Comm
~Comm()
Destructor (frees wrapped communicator).
Definition MPI.cpp:36
dolfinx::MPI::Comm::comm
MPI_Comm comm() const noexcept
Return the underlying MPI_Comm object.
Definition MPI.cpp:62
dolfinx::common::Timer
Timer for measuring and logging elapsed time durations.
Definition Timer.h:40
dolfinx::MPI
MPI support functionality.
Definition MPI.h:30
dolfinx::MPI::mpi_t
MPI_Datatype mpi_t
Retrieves the MPI data type associated to the provided type.
Definition MPI.h:280
dolfinx::MPI::distribute_data
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)
Distribute rows of a rectangular data array to ranks where they are required (scalable version).
Definition MPI.h:679
dolfinx::MPI::distribute_to_postoffice
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)
Distribute row data to 'post office' ranks.
Definition MPI.h:313
dolfinx::MPI::compute_graph_edges_nbx
std::vector< int > compute_graph_edges_nbx(MPI_Comm comm, std::span< const int > edges, int tag=static_cast< int >(tag::consensus_nbx))
Determine incoming graph edges using the NBX consensus algorithm.
Definition MPI.cpp:162
dolfinx::MPI::index_owner
constexpr int index_owner(int size, std::size_t index, std::size_t N)
Return which rank owns index in global range [0, N - 1] (inverse of MPI::local_range).
Definition MPI.h:113
dolfinx::MPI::distribute_from_postoffice
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)
Distribute rows of a rectangular data array from post office ranks to ranks where they are required.
Definition MPI.h:464
dolfinx::MPI::tag
tag
MPI communication tags.
Definition MPI.h:33
dolfinx::MPI::compute_graph_edges_pcx
std::vector< int > compute_graph_edges_pcx(MPI_Comm comm, std::span< const int > edges)
Determine incoming graph edges using the PCX consensus algorithm.
Definition MPI.cpp:95
dolfinx::MPI::check_error
void check_error(MPI_Comm comm, int code)
Check MPI error code. If the error code is not equal to MPI_SUCCESS, then std::abort is called.
Definition MPI.cpp:80
dolfinx::MPI::size
int size(MPI_Comm comm)
Definition MPI.cpp:72
dolfinx::MPI::rank
int rank(MPI_Comm comm)
Return process rank for the communicator.
Definition MPI.cpp:64
dolfinx::MPI::local_range
constexpr std::array< std::int64_t, 2 > local_range(int rank, std::int64_t N, int size)
Return local range for the calling process, partitioning the global [0, N - 1] range across all ranks...
Definition MPI.h:89
dolfinx::MPI::dependent_false
Definition MPI.h:268
dolfinx::MPI::mpi_type_mapping
MPI Type.
Definition MPI.h:275