Loading...
Searching...
No Matches
Scatterer.h
// Copyright (C) 2022-2025 Igor Baratta and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once
#include "IndexMap.h"
#include "MPI.h"
#include "sort.h"
#include <algorithm>
#include <concepts>
#include <functional>
#include <memory>
#include <mpi.h>
#include <numeric>
#include <span>
#include <type_traits>
#include <vector>
namespace dolfinx::common
{
template <class Container = std::vector<std::int32_t>>
class Scatterer
{
static_assert(std::is_integral_v<typename Container::value_type>);
template <class>
friend class Scatterer;
public:
using container_type = Container;
Scatterer(const IndexMap& map, int bs)
: _src(map.src().begin(), map.src().end()),
_dest(map.dest().begin(), map.dest().end()),
_sizes_remote(_src.size(), 0), _displs_remote(_src.size() + 1),
_sizes_local(_dest.size()), _displs_local(_dest.size() + 1)
{
if (dolfinx::MPI::size(map.comm()) == 1)
return;
int ierr;
// Check that src and dest ranks are unique and sorted
assert(std::ranges::is_sorted(_src));
assert(std::ranges::is_sorted(_dest));
// Create communicators with directed edges:
// (0) owner -> ghost,
// (1) ghost -> owner
MPI_Comm comm0;
ierr = MPI_Dist_graph_create_adjacent(
map.comm(), _src.size(), _src.data(), MPI_UNWEIGHTED, _dest.size(),
_dest.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm0);
_comm0 = dolfinx::MPI::Comm(comm0, false);
dolfinx::MPI::check_error(map.comm(), ierr);
MPI_Comm comm1;
ierr = MPI_Dist_graph_create_adjacent(
map.comm(), _dest.size(), _dest.data(), MPI_UNWEIGHTED, _src.size(),
_src.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm1);
_comm1 = dolfinx::MPI::Comm(comm1, false);
dolfinx::MPI::check_error(map.comm(), ierr);
// Build permutation array that sorts ghost indices by owning rank
std::span owners = map.owners();
std::vector<std::int32_t> perm(owners.size());
std::iota(perm.begin(), perm.end(), 0);
dolfinx::radix_sort(perm, [&owners](auto index) { return owners[index]; });
// Sort (i) ghost indices and (ii) ghost index owners by rank
// (using perm array)
std::span ghosts = map.ghosts();
std::vector<int> owners_sorted(owners.size());
std::vector<std::int64_t> ghosts_sorted(owners.size());
std::ranges::transform(perm, owners_sorted.begin(),
[&owners](auto idx) { return owners[idx]; });
std::ranges::transform(perm, ghosts_sorted.begin(),
[&ghosts](auto idx) { return ghosts[idx]; });
// For data associated with ghost indices, packed by owning
// (neighbourhood) rank, compute sizes and displacements. I.e., when
// sending ghost index data from this rank to the owning ranks,
// disp[i] is the first entry in the buffer sent to neighbourhood
// rank i, and disp[i + 1] - disp[i] is the number of values sent to
// rank i.
assert(_sizes_remote.size() == _src.size());
assert(_displs_remote.size() == _src.size() + 1);
auto begin = owners_sorted.begin();
for (std::size_t i = 0; i < _src.size(); i++)
{
auto upper = std::upper_bound(begin, owners_sorted.end(), _src[i]);
std::size_t num_ind = std::distance(begin, upper);
_displs_remote[i + 1] = _displs_remote[i] + num_ind;
_sizes_remote[i] = num_ind;
begin = upper;
}
// For data associated with owned indices that are ghosted by other
// ranks, compute the size and displacement arrays. When sending
// data associated with ghost indices to the owner, these size and
// displacement arrays are for the receive buffer.
// Compute sizes and displacements of local data (how many local
// elements to be sent/received grouped by neighbors)
assert(_sizes_local.size() == _dest.size());
assert(_displs_local.size() == _dest.size() + 1);
_sizes_remote.reserve(1);
_sizes_local.reserve(1);
ierr = MPI_Neighbor_alltoall(_sizes_remote.data(), 1, MPI_INT32_T,
_sizes_local.data(), 1, MPI_INT32_T,
_comm1.comm());
dolfinx::MPI::check_error(_comm1.comm(), ierr);
std::partial_sum(_sizes_local.begin(), _sizes_local.end(),
std::next(_displs_local.begin()));
assert((int)ghosts_sorted.size() == _displs_remote.back());
assert((int)ghosts_sorted.size() == _displs_remote.back());
// Send ghost global indices to owning rank, and receive owned
// indices that are ghosts on other ranks
std::vector<std::int64_t> recv_buffer(_displs_local.back(), 0);
ierr = MPI_Neighbor_alltoallv(
ghosts_sorted.data(), _sizes_remote.data(), _displs_remote.data(),
MPI_INT64_T, recv_buffer.data(), _sizes_local.data(),
_displs_local.data(), MPI_INT64_T, _comm1.comm());
dolfinx::MPI::check_error(_comm1.comm(), ierr);
const std::array<std::int64_t, 2> range = map.local_range();
#ifndef NDEBUG
// Check that all received indice are within the owned range
std::ranges::for_each(recv_buffer, [range](auto idx)
{ assert(idx >= range[0] and idx < range[1]); });
#endif
{
// Scale sizes and displacements by block size
for (auto& x : {std::ref(_sizes_local), std::ref(_displs_local),
std::ref(_sizes_remote), std::ref(_displs_remote)})
{
std::ranges::transform(x.get(), x.get().begin(),
[bs](auto e) { return e *= bs; });
}
}
{
// Expand local indices using block size and convert it from
// global to local numbering
std::vector<typename container_type::value_type> idx(recv_buffer.size()
* bs);
std::int64_t offset = range[0] * bs;
for (std::size_t i = 0; i < recv_buffer.size(); i++)
for (int j = 0; j < bs; j++)
idx[i * bs + j] = (recv_buffer[i] * bs + j) - offset;
_local_inds = std::move(idx);
}
{
// Expand remote indices using block size
std::vector<typename container_type::value_type> idx(perm.size() * bs);
for (std::size_t i = 0; i < perm.size(); i++)
for (int j = 0; j < bs; j++)
idx[i * bs + j] = perm[i] * bs + j;
_remote_inds = std::move(idx);
}
}
Scatterer(const Scatterer& scatterer) = default;
template <class U>
Scatterer(const Scatterer<U>& s)
: _comm0(s._comm0), _comm1(s._comm1), _src(s._src), _dest(s._dest),
_remote_inds(s._remote_inds.begin(), s._remote_inds.end()),
_sizes_remote(s._sizes_remote), _displs_remote(s._displs_remote),
_local_inds(s._local_inds.begin(), s._local_inds.end()),
_sizes_local(s._sizes_local), _displs_local(s._displs_local)
{
}
template <typename T>
void scatter_fwd_begin(const T* send_buffer, T* recv_buffer,
MPI_Request& request) const
{
// Return early if there are no incoming or outgoing edges
if (_sizes_local.empty() and _sizes_remote.empty())
return;
int ierr = MPI_Ineighbor_alltoallv(
send_buffer, _sizes_local.data(), _displs_local.data(),
dolfinx::MPI::mpi_t<T>, recv_buffer, _sizes_remote.data(),
_displs_remote.data(), dolfinx::MPI::mpi_t<T>, _comm0.comm(), &request);
dolfinx::MPI::check_error(_comm0.comm(), ierr);
}
template <typename T>
void scatter_fwd_begin(const T* send_buffer, T* recv_buffer,
std::span<MPI_Request> requests) const
{
if (requests.size() != _dest.size() + _src.size())
{
throw std::runtime_error(
"Point-to-point scatterer has wrong number of MPI_Requests.");
}
// Return early if there are no incoming or outgoing edges
if (_sizes_local.empty() and _sizes_remote.empty())
return;
for (std::size_t i = 0; i < _src.size(); ++i)
{
int ierr = MPI_Irecv(recv_buffer + _displs_remote[i], _sizes_remote[i],
dolfinx::MPI::mpi_t<T>, _src[i], MPI_ANY_TAG,
_comm0.comm(), &requests[i]);
dolfinx::MPI::check_error(_comm0.comm(), ierr);
}
for (std::size_t i = 0; i < _dest.size(); ++i)
{
int ierr = MPI_Isend(send_buffer + _displs_local[i], _sizes_local[i],
dolfinx::MPI::mpi_t<T>, _dest[i], 0, _comm0.comm(),
&requests[i + _src.size()]);
dolfinx::MPI::check_error(_comm0.comm(), ierr);
}
}
template <typename T>
void scatter_rev_begin(const T* send_buffer, T* recv_buffer,
MPI_Request& request) const
{
// Return early if there are no incoming or outgoing edges
if (_sizes_local.empty() and _sizes_remote.empty())
return;
int ierr = MPI_Ineighbor_alltoallv(
send_buffer, _sizes_remote.data(), _displs_remote.data(), MPI::mpi_t<T>,
recv_buffer, _sizes_local.data(), _displs_local.data(), MPI::mpi_t<T>,
_comm1.comm(), &request);
dolfinx::MPI::check_error(_comm1.comm(), ierr);
}
template <typename T>
void scatter_rev_begin(const T* send_buffer, T* recv_buffer,
std::span<MPI_Request> requests) const
{
if (requests.size() != _dest.size() + _src.size())
{
throw std::runtime_error(
"Point-to-point scatterer has wrong number of MPI_Requests.");
}
// Return early if there are no incoming or outgoing edges
if (_sizes_local.empty() and _sizes_remote.empty())
return;
// Start non-blocking send from this process to ghost owners
for (std::size_t i = 0; i < _dest.size(); i++)
{
int ierr = MPI_Irecv(recv_buffer + _displs_local[i], _sizes_local[i],
dolfinx::MPI::mpi_t<T>, _dest[i], MPI_ANY_TAG,
_comm0.comm(), &requests[i]);
dolfinx::MPI::check_error(_comm0.comm(), ierr);
}
// Start non-blocking receive from neighbor process for which an
// owned index is a ghost
for (std::size_t i = 0; i < _src.size(); i++)
{
int ierr = MPI_Isend(send_buffer + _displs_remote[i], _sizes_remote[i],
dolfinx::MPI::mpi_t<T>, _src[i], 0, _comm0.comm(),
&requests[i + _dest.size()]);
dolfinx::MPI::check_error(_comm0.comm(), ierr);
}
}
void scatter_end(std::span<MPI_Request> requests) const
{
// Return early if there are no incoming or outgoing edges
if (_sizes_local.empty() and _sizes_remote.empty())
return;
// Wait for communication to complete
MPI_Waitall(requests.size(), requests.data(), MPI_STATUS_IGNORE);
}
void scatter_end(MPI_Request& request) const
{
scatter_end(std::span<MPI_Request>(&request, 1));
}
const container_type& local_indices() const noexcept { return _local_inds; }
const container_type& remote_indices() const noexcept { return _remote_inds; }
std::size_t num_p2p_requests() { return _dest.size() + _src.size(); }
private:
// Communicator where the source ranks own the indices in the callers
// halo, and the destination ranks 'ghost' indices owned by the
// caller. I.e.,
// - in-edges (src) are from ranks that own my ghosts
// - out-edges (dest) go to ranks that 'ghost' my owned indices
dolfinx::MPI::Comm _comm0{MPI_COMM_NULL};
// Communicator where the source ranks have ghost indices that are
// owned by the caller, and the destination ranks are the owners of
// indices in the callers halo region. I.e.,
// - in-edges (src) are from ranks that 'ghost' my owned indices
// - out-edges (dest) are to the owning ranks of my ghost indices
dolfinx::MPI::Comm _comm1{MPI_COMM_NULL};
// Set of ranks that own ghosts
// FIXME: Should we store the index map instead?
std::vector<int> _src;
// Set of ranks ghost owned indices
// FIXME: Should we store the index map instead?
std::vector<int> _dest;
// Permutation indices used to pack and unpack ghost data (remote)
container_type _remote_inds;
// Number of remote indices (ghosts) for each neighbor process
std::vector<int> _sizes_remote;
// Displacements of remote data for mpi scatter and gather
std::vector<int> _displs_remote;
// Permutation indices used to pack and unpack local shared data
// (owned indices that are shared with other processes). Indices are
// grouped by neighbor process.
container_type _local_inds;
// Number of local shared indices per neighbor process
std::vector<int> _sizes_local;
// Displacements of local data for mpi scatter and gather
std::vector<int> _displs_local;
};
} // namespace dolfinx::common
dolfinx::MPI::Comm
A duplicate MPI communicator and manage lifetime of the communicator.
Definition MPI.h:42
dolfinx::common::IndexMap
Definition IndexMap.h:97
dolfinx::common::IndexMap::owners
std::span< const int > owners() const
The ranks that own each ghost index.
Definition IndexMap.h:221
dolfinx::common::IndexMap::local_range
std::array< std::int64_t, 2 > local_range() const noexcept
Range of indices (global) owned by this process.
Definition IndexMap.cpp:926
dolfinx::common::IndexMap::ghosts
std::span< const std::int64_t > ghosts() const noexcept
Definition IndexMap.cpp:940
dolfinx::common::IndexMap::comm
MPI_Comm comm() const
Return the MPI communicator that the map is defined on.
Definition IndexMap.cpp:1005
dolfinx::common::Scatterer::scatter_fwd_begin
void scatter_fwd_begin(const T *send_buffer, T *recv_buffer, std::span< MPI_Request > requests) const
Start a non-blocking send of owned data to ranks that ghost the data using point-to-point MPI communi...
Definition Scatterer.h:282
dolfinx::common::Scatterer::scatter_end
void scatter_end(std::span< MPI_Request > requests) const
Complete non-blocking MPI point-to-point sends.
Definition Scatterer.h:407
dolfinx::common::Scatterer::num_p2p_requests
std::size_t num_p2p_requests()
Number of required MPI_Requests for point-to-point communication.
Definition Scatterer.h:493
dolfinx::common::Scatterer::scatter_rev_begin
void scatter_rev_begin(const T *send_buffer, T *recv_buffer, MPI_Request &request) const
Start a non-blocking send of ghost data to ranks that own the data using MPI neighbourhood collective...
Definition Scatterer.h:339
dolfinx::common::Scatterer::scatter_fwd_begin
void scatter_fwd_begin(const T *send_buffer, T *recv_buffer, MPI_Request &request) const
Start a non-blocking send of owned data to ranks that ghost the data using MPI neighbourhood collecti...
Definition Scatterer.h:254
dolfinx::common::Scatterer::local_indices
const container_type & local_indices() const noexcept
Array of indices for packing/unpacking owned data to/from a send/receive buffer.
Definition Scatterer.h:459
dolfinx::common::Scatterer::remote_indices
const container_type & remote_indices() const noexcept
Array of indices for packing/unpacking ghost data to/from a send/receive buffer.
Definition Scatterer.h:487
dolfinx::common::Scatterer::scatter_end
void scatter_end(MPI_Request &request) const
Complete a non-blocking MPI neighbourhood collective send.
Definition Scatterer.h:424
dolfinx::common::Scatterer::container_type
Container container_type
Container type used to store local and remote indices.
Definition Scatterer.h:57
dolfinx::common::Scatterer::scatter_rev_begin
void scatter_rev_begin(const T *send_buffer, T *recv_buffer, std::span< MPI_Request > requests) const
Start a non-blocking send of ghost data to ranks that own the data using point-to-point MPI communica...
Definition Scatterer.h:367
dolfinx::common::Scatterer::Scatterer
Scatterer(const IndexMap &map, int bs)
Create a scatterer for data with a layout described by an IndexMap and a block size.
Definition Scatterer.h:66
dolfinx::common::Scatterer::Scatterer
Scatterer(const Scatterer< U > &s)
Cast-copy constructor.
Definition Scatterer.h:219
dolfinx::common::Scatterer::Scatterer
Scatterer(const Scatterer &scatterer)=default
Copy constructor.
dolfinx::MPI::mpi_t
MPI_Datatype mpi_t
Retrieves the MPI data type associated to the provided type.
Definition MPI.h:280
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::common
Miscellaneous classes, functions and types.
Definition dolfinx_common.h:8
dolfinx::radix_sort
constexpr __radix_sort radix_sort
Radix sort.
Definition sort.h:170