DOLFINx 0.10.0.0
DOLFINx C++ interface
Loading...
Searching...
No Matches
Scatterer.h
1// Copyright (C) 2022-2025 Igor Baratta and Garth N. Wells
2//
3// This file is part of DOLFINx (https://www.fenicsproject.org)
4//
5// SPDX-License-Identifier: LGPL-3.0-or-later
6
7#pragma once
8
9#include "IndexMap.h"
10#include "MPI.h"
11#include "sort.h"
12#include <algorithm>
13#include <concepts>
14#include <functional>
15#include <memory>
16#include <mpi.h>
17#include <numeric>
18#include <span>
19#include <type_traits>
20#include <vector>
21
22namespace dolfinx::common
23{
40template <class Container = std::vector<std::int32_t>>
41 requires std::is_integral_v<typename Container::value_type>
43{
44public:
46 using container_type = Container;
47
55 Scatterer(const IndexMap& map, int bs)
56 : _src(map.src().begin(), map.src().end()),
57 _dest(map.dest().begin(), map.dest().end()),
58 _sizes_remote(_src.size(), 0), _displs_remote(_src.size() + 1),
59 _sizes_local(_dest.size()), _displs_local(_dest.size() + 1)
60
61 {
62 if (dolfinx::MPI::size(map.comm()) == 1)
63 return;
64
65 int ierr;
66
67 // Check that src and dest ranks are unique and sorted
68 assert(std::ranges::is_sorted(_src));
69 assert(std::ranges::is_sorted(_dest));
70
71 // Create communicators with directed edges:
72 // (0) owner -> ghost,
73 // (1) ghost -> owner
74 MPI_Comm comm0;
75 ierr = MPI_Dist_graph_create_adjacent(
76 map.comm(), _src.size(), _src.data(), MPI_UNWEIGHTED, _dest.size(),
77 _dest.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm0);
78 _comm0 = dolfinx::MPI::Comm(comm0, false);
80
81 MPI_Comm comm1;
82 ierr = MPI_Dist_graph_create_adjacent(
83 map.comm(), _dest.size(), _dest.data(), MPI_UNWEIGHTED, _src.size(),
84 _src.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm1);
85 _comm1 = dolfinx::MPI::Comm(comm1, false);
87
88 // Build permutation array that sorts ghost indices by owning rank
89 std::span owners = map.owners();
90 std::vector<std::int32_t> perm(owners.size());
91 std::iota(perm.begin(), perm.end(), 0);
92 dolfinx::radix_sort(perm, [&owners](auto index) { return owners[index]; });
93
94 // Sort (i) ghost indices and (ii) ghost index owners by rank
95 // (using perm array)
96 std::span ghosts = map.ghosts();
97 std::vector<int> owners_sorted(owners.size());
98 std::vector<std::int64_t> ghosts_sorted(owners.size());
99 std::ranges::transform(perm, owners_sorted.begin(),
100 [&owners](auto idx) { return owners[idx]; });
101 std::ranges::transform(perm, ghosts_sorted.begin(),
102 [&ghosts](auto idx) { return ghosts[idx]; });
103
104 // For data associated with ghost indices, packed by owning
105 // (neighbourhood) rank, compute sizes and displacements. I.e., when
106 // sending ghost index data from this rank to the owning ranks,
107 // disp[i] is the first entry in the buffer sent to neighbourhood
108 // rank i, and disp[i + 1] - disp[i] is the number of values sent to
109 // rank i.
110 assert(_sizes_remote.size() == _src.size());
111 assert(_displs_remote.size() == _src.size() + 1);
112 auto begin = owners_sorted.begin();
113 for (std::size_t i = 0; i < _src.size(); i++)
114 {
115 auto upper = std::upper_bound(begin, owners_sorted.end(), _src[i]);
116 std::size_t num_ind = std::distance(begin, upper);
117 _displs_remote[i + 1] = _displs_remote[i] + num_ind;
118 _sizes_remote[i] = num_ind;
119 begin = upper;
120 }
121
122 // For data associated with owned indices that are ghosted by other
123 // ranks, compute the size and displacement arrays. When sending
124 // data associated with ghost indices to the owner, these size and
125 // displacement arrays are for the receive buffer.
126
127 // Compute sizes and displacements of local data (how many local
128 // elements to be sent/received grouped by neighbors)
129 assert(_sizes_local.size() == _dest.size());
130 assert(_displs_local.size() == _dest.size() + 1);
131 _sizes_remote.reserve(1);
132 _sizes_local.reserve(1);
133 ierr = MPI_Neighbor_alltoall(_sizes_remote.data(), 1, MPI_INT32_T,
134 _sizes_local.data(), 1, MPI_INT32_T,
135 _comm1.comm());
136 dolfinx::MPI::check_error(_comm1.comm(), ierr);
137
138 std::partial_sum(_sizes_local.begin(), _sizes_local.end(),
139 std::next(_displs_local.begin()));
140
141 assert((std::int32_t)ghosts_sorted.size() == _displs_remote.back());
142 assert((std::int32_t)ghosts_sorted.size() == _displs_remote.back());
143
144 // Send ghost global indices to owning rank, and receive owned
145 // indices that are ghosts on other ranks
146 std::vector<std::int64_t> recv_buffer(_displs_local.back(), 0);
147 ierr = MPI_Neighbor_alltoallv(
148 ghosts_sorted.data(), _sizes_remote.data(), _displs_remote.data(),
149 MPI_INT64_T, recv_buffer.data(), _sizes_local.data(),
150 _displs_local.data(), MPI_INT64_T, _comm1.comm());
151 dolfinx::MPI::check_error(_comm1.comm(), ierr);
152
153 const std::array<std::int64_t, 2> range = map.local_range();
154#ifndef NDEBUG
155 // Check that all received indice are within the owned range
156 std::ranges::for_each(recv_buffer, [range](auto idx)
157 { assert(idx >= range[0] and idx < range[1]); });
158#endif
159
160 {
161 // Scale sizes and displacements by block size
162 for (auto& x : {std::ref(_sizes_local), std::ref(_displs_local),
163 std::ref(_sizes_remote), std::ref(_displs_remote)})
164 {
165 std::ranges::transform(x.get(), x.get().begin(),
166 [bs](auto e) { return e *= bs; });
167 }
168 }
169
170 {
171 // Expand local indices using block size and convert it from
172 // global to local numbering
173 std::vector<typename container_type::value_type> idx(recv_buffer.size()
174 * bs);
175 std::int64_t offset = range[0] * bs;
176 for (std::size_t i = 0; i < recv_buffer.size(); i++)
177 for (int j = 0; j < bs; j++)
178 idx[i * bs + j] = (recv_buffer[i] * bs + j) - offset;
179 _local_inds = std::move(idx);
180 }
181
182 {
183 // Expand remote indices using block size
184 std::vector<typename container_type::value_type> idx(perm.size() * bs);
185 for (std::size_t i = 0; i < perm.size(); i++)
186 for (int j = 0; j < bs; j++)
187 idx[i * bs + j] = perm[i] * bs + j;
188 _remote_inds = std::move(idx);
189 }
190 }
191
217 template <typename T>
218 void scatter_fwd_begin(const T* send_buffer, T* recv_buffer,
219 MPI_Request& request) const
220 {
221 // Return early if there are no incoming or outgoing edges
222 if (_sizes_local.empty() and _sizes_remote.empty())
223 return;
224
225 int ierr = MPI_Ineighbor_alltoallv(
226 send_buffer, _sizes_local.data(), _displs_local.data(),
227 dolfinx::MPI::mpi_t<T>, recv_buffer, _sizes_remote.data(),
228 _displs_remote.data(), dolfinx::MPI::mpi_t<T>, _comm0.comm(), &request);
229 dolfinx::MPI::check_error(_comm0.comm(), ierr);
230 }
231
245 template <typename T>
246 void scatter_fwd_begin(const T* send_buffer, T* recv_buffer,
247 std::span<MPI_Request> requests) const
248 {
249 if (requests.size() != _dest.size() + _src.size())
250 {
251 throw std::runtime_error(
252 "Point-to-point scatterer has wrong number of MPI_Requests.");
253 }
254
255 // Return early if there are no incoming or outgoing edges
256 if (_sizes_local.empty() and _sizes_remote.empty())
257 return;
258
259 for (std::size_t i = 0; i < _src.size(); ++i)
260 {
261 int ierr = MPI_Irecv(recv_buffer + _displs_remote[i], _sizes_remote[i],
262 dolfinx::MPI::mpi_t<T>, _src[i], MPI_ANY_TAG,
263 _comm0.comm(), &requests[i]);
264 dolfinx::MPI::check_error(_comm0.comm(), ierr);
265 }
266
267 for (std::size_t i = 0; i < _dest.size(); ++i)
268 {
269 int ierr = MPI_Isend(send_buffer + _displs_local[i], _sizes_local[i],
270 dolfinx::MPI::mpi_t<T>, _dest[i], 0, _comm0.comm(),
271 &requests[i + _src.size()]);
272 dolfinx::MPI::check_error(_comm0.comm(), ierr);
273 }
274 }
275
302 template <typename T>
303 void scatter_rev_begin(const T* send_buffer, T* recv_buffer,
304 MPI_Request& request) const
305 {
306 // Return early if there are no incoming or outgoing edges
307 if (_sizes_local.empty() and _sizes_remote.empty())
308 return;
309
310 int ierr = MPI_Ineighbor_alltoallv(
311 send_buffer, _sizes_remote.data(), _displs_remote.data(), MPI::mpi_t<T>,
312 recv_buffer, _sizes_local.data(), _displs_local.data(), MPI::mpi_t<T>,
313 _comm1.comm(), &request);
314 dolfinx::MPI::check_error(_comm1.comm(), ierr);
315 }
316
330 template <typename T>
331 void scatter_rev_begin(const T* send_buffer, T* recv_buffer,
332 std::span<MPI_Request> requests) const
333 {
334 if (requests.size() != _dest.size() + _src.size())
335 {
336 throw std::runtime_error(
337 "Point-to-point scatterer has wrong number of MPI_Requests.");
338 }
339
340 // Return early if there are no incoming or outgoing edges
341 if (_sizes_local.empty() and _sizes_remote.empty())
342 return;
343
344 // Start non-blocking send from this process to ghost owners
345 for (std::size_t i = 0; i < _dest.size(); i++)
346 {
347 int ierr = MPI_Irecv(recv_buffer + _displs_local[i], _sizes_local[i],
348 dolfinx::MPI::mpi_t<T>, _dest[i], MPI_ANY_TAG,
349 _comm0.comm(), &requests[i]);
350 dolfinx::MPI::check_error(_comm0.comm(), ierr);
351 }
352
353 // Start non-blocking receive from neighbor process for which an
354 // owned index is a ghost
355 for (std::size_t i = 0; i < _src.size(); i++)
356 {
357 int ierr = MPI_Isend(send_buffer + _displs_remote[i], _sizes_remote[i],
358 dolfinx::MPI::mpi_t<T>, _src[i], 0, _comm0.comm(),
359 &requests[i + _dest.size()]);
360 dolfinx::MPI::check_error(_comm0.comm(), ierr);
361 }
362 }
363
371 void scatter_end(std::span<MPI_Request> requests) const
372 {
373 // Return early if there are no incoming or outgoing edges
374 if (_sizes_local.empty() and _sizes_remote.empty())
375 return;
376
377 // Wait for communication to complete
378 MPI_Waitall(requests.size(), requests.data(), MPI_STATUS_IGNORE);
379 }
380
388 void scatter_end(MPI_Request& request) const
389 {
390 scatter_end(std::span<MPI_Request>(&request, 1));
391 }
392
423 const container_type& local_indices() const noexcept { return _local_inds; }
424
451 const container_type& remote_indices() const noexcept { return _remote_inds; }
452
457 std::size_t num_p2p_requests() { return _dest.size() + _src.size(); }
458
459private:
460 // Communicator where the source ranks own the indices in the callers
461 // halo, and the destination ranks 'ghost' indices owned by the
462 // caller. I.e.,
463 // - in-edges (src) are from ranks that own my ghosts
464 // - out-edges (dest) go to ranks that 'ghost' my owned indices
465 dolfinx::MPI::Comm _comm0{MPI_COMM_NULL};
466
467 // Communicator where the source ranks have ghost indices that are
468 // owned by the caller, and the destination ranks are the owners of
469 // indices in the callers halo region. I.e.,
470 // - in-edges (src) are from ranks that 'ghost' my owned indices
471 // - out-edges (dest) are to the owning ranks of my ghost indices
472 dolfinx::MPI::Comm _comm1{MPI_COMM_NULL};
473
474 // Set of ranks that own ghosts
475 // FIXME: Should we store the index map instead?
476 std::vector<int> _src;
477
478 // Set of ranks ghost owned indices
479 // FIXME: Should we store the index map instead?
480 std::vector<int> _dest;
481
482 // Permutation indices used to pack and unpack ghost data (remote)
483 container_type _remote_inds;
484
485 // Number of remote indices (ghosts) for each neighbor process
486 std::vector<int> _sizes_remote;
487
488 // Displacements of remote data for mpi scatter and gather
489 std::vector<int> _displs_remote;
490
491 // Permutation indices used to pack and unpack local shared data
492 // (owned indices that are shared with other processes). Indices are
493 // grouped by neighbor process.
494 container_type _local_inds;
495
496 // Number of local shared indices per neighbor process
497 std::vector<int> _sizes_local;
498
499 // Displacements of local data for mpi scatter and gather
500 std::vector<int> _displs_local;
501};
502} // namespace dolfinx::common
A duplicate MPI communicator and manage lifetime of the communicator.
Definition MPI.h:42
Definition IndexMap.h:97
std::span< const int > owners() const
The ranks that own each ghost index.
Definition IndexMap.h:221
std::array< std::int64_t, 2 > local_range() const noexcept
Range of indices (global) owned by this process.
Definition IndexMap.cpp:926
std::span< const std::int64_t > ghosts() const noexcept
Definition IndexMap.cpp:940
MPI_Comm comm() const
Return the MPI communicator that the map is defined on.
Definition IndexMap.cpp:1005
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:246
void scatter_end(std::span< MPI_Request > requests) const
Complete non-blocking MPI point-to-point sends.
Definition Scatterer.h:371
std::size_t num_p2p_requests()
Number of required MPI_Requests for point-to-point communication.
Definition Scatterer.h:457
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:303
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:218
const container_type & local_indices() const noexcept
Array of indices for packing/unpacking owned data to/from a send/receive buffer.
Definition Scatterer.h:423
const container_type & remote_indices() const noexcept
Array of indices for packing/unpacking ghost data to/from a send/receive buffer.
Definition Scatterer.h:451
void scatter_end(MPI_Request &request) const
Complete a non-blocking MPI neighbourhood collective send.
Definition Scatterer.h:388
Container container_type
Container type used to store local and remote indices.
Definition Scatterer.h:46
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:331
Scatterer(const IndexMap &map, int bs)
Create a scatterer for data with a layout described by an IndexMap, and with a block size.
Definition Scatterer.h:55
MPI_Datatype mpi_t
Retrieves the MPI data type associated to the provided type.
Definition MPI.h:280
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
int size(MPI_Comm comm)
Definition MPI.cpp:72
Miscellaneous classes, functions and types.
Definition dolfinx_common.h:8
constexpr __radix_sort radix_sort
Radix sort.
Definition sort.h:170