DOLFINx
DOLFINx C++ interface
Scatterer.h
1// Copyright (C) 2022 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 "MPI.h"
10#include <algorithm>
11#include <mpi.h>
12#include <span>
13#include <vector>
14
15using namespace dolfinx;
16
17namespace dolfinx::common
18{
19class IndexMap;
20
28{
29public:
35 Scatterer(const IndexMap& map, int bs);
36
55 template <typename T>
56 void scatter_fwd_begin(const std::span<const T>& send_buffer,
57 const std::span<T>& recv_buffer,
58 MPI_Request& request) const
59 {
60 // Return early if there are no incoming or outgoing edges
61 if (_sizes_local.empty() and _sizes_remote.empty())
62 return;
63
64 MPI_Ineighbor_alltoallv(
65 send_buffer.data(), _sizes_local.data(), _displs_local.data(),
66 MPI::mpi_type<T>(), recv_buffer.data(), _sizes_remote.data(),
67 _displs_remote.data(), MPI::mpi_type<T>(), _comm0.comm(), &request);
68 }
69
78 void scatter_fwd_end(MPI_Request& request) const;
79
98 template <typename T, typename Functor>
99 void scatter_fwd_begin(const std::span<const T>& local_data,
100 std::span<T> local_buffer, std::span<T> remote_buffer,
101 Functor pack_fn, MPI_Request& request) const
102 {
103 assert(local_buffer.size() == _local_inds.size());
104 assert(remote_buffer.size() == _remote_inds.size());
105 pack_fn(local_data, _local_inds, local_buffer);
106 scatter_fwd_begin(std::span<const T>(local_buffer), remote_buffer, request);
107 }
108
128 template <typename T, typename Functor>
129 void scatter_fwd_end(const std::span<const T>& remote_buffer,
130 std::span<T> remote_data, Functor unpack_fn,
131 MPI_Request& request) const
132 {
133 assert(remote_buffer.size() == _remote_inds.size());
134 assert(remote_data.size() == _remote_inds.size());
135 scatter_fwd_end(request);
136 unpack_fn(remote_buffer, _remote_inds, remote_data,
137 [](T /*a*/, T b) { return b; });
138 }
139
151 template <typename T>
152 void scatter_fwd(const std::span<const T>& local_data,
153 std::span<T> remote_data) const
154 {
155 MPI_Request request;
156 std::vector<T> local_buffer(local_buffer_size(), 0);
157 std::vector<T> remote_buffer(remote_buffer_size(), 0);
158 auto pack_fn = [](const auto& in, const auto& idx, auto& out)
159 {
160 for (std::size_t i = 0; i < idx.size(); ++i)
161 out[i] = in[idx[i]];
162 };
163 scatter_fwd_begin(local_data, std::span<T>(local_buffer),
164 std::span<T>(remote_buffer), pack_fn, request);
165
166 auto unpack_fn = [](const auto& in, const auto& idx, auto& out, auto op)
167 {
168 for (std::size_t i = 0; i < idx.size(); ++i)
169 out[idx[i]] = op(out[idx[i]], in[i]);
170 };
171
172 scatter_fwd_end(std::span<const T>(remote_buffer), remote_data, unpack_fn,
173 request);
174 }
175
200 template <typename T>
201 void scatter_rev_begin(const std::span<const T>& send_buffer,
202 const std::span<T>& recv_buffer,
203 MPI_Request& request) const
204 {
205 // Return early if there are no incoming or outgoing edges
206 if (_sizes_local.empty() and _sizes_remote.empty())
207 return;
208
209 // Send and receive data
210 MPI_Ineighbor_alltoallv(
211 send_buffer.data(), _sizes_remote.data(), _displs_remote.data(),
212 MPI::mpi_type<T>(), recv_buffer.data(), _sizes_local.data(),
213 _displs_local.data(), MPI::mpi_type<T>(), _comm1.comm(), &request);
214 }
215
224 void scatter_rev_end(MPI_Request& request) const;
225
252 template <typename T, typename Functor>
253 void scatter_rev_begin(const std::span<const T>& remote_data,
254 std::span<T> remote_buffer, std::span<T> local_buffer,
255 Functor pack_fn, MPI_Request& request) const
256 {
257 assert(local_buffer.size() == _local_inds.size());
258 assert(remote_buffer.size() == _remote_inds.size());
259 pack_fn(remote_data, _remote_inds, remote_buffer);
260 scatter_rev_begin(std::span<const T>(remote_buffer), local_buffer, request);
261 }
262
282 template <typename T, typename Functor, typename BinaryOp>
283 void scatter_rev_end(const std::span<const T>& local_buffer,
284 std::span<T> local_data, Functor unpack_fn, BinaryOp op,
285 MPI_Request& request)
286 {
287 assert(local_buffer.size() == _local_inds.size());
288 if (_local_inds.size() > 0)
289 assert(*std::max_element(_local_inds.begin(), _local_inds.end())
290 < std::int32_t(local_data.size()));
291 scatter_rev_end(request);
292 unpack_fn(local_buffer, _local_inds, local_data, op);
293 }
294
297 template <typename T, typename BinaryOp>
298 void scatter_rev(std::span<T> local_data,
299 const std::span<const T>& remote_data, BinaryOp op)
300 {
301 std::vector<T> local_buffer(local_buffer_size(), 0);
302 std::vector<T> remote_buffer(remote_buffer_size(), 0);
303 auto pack_fn = [](const auto& in, const auto& idx, auto& out)
304 {
305 for (std::size_t i = 0; i < idx.size(); ++i)
306 out[i] = in[idx[i]];
307 };
308 auto unpack_fn = [](const auto& in, const auto& idx, auto& out, auto op)
309 {
310 for (std::size_t i = 0; i < idx.size(); ++i)
311 out[idx[i]] = op(out[idx[i]], in[i]);
312 };
313 MPI_Request request;
314 scatter_rev_begin(remote_data, std::span<T>(remote_buffer),
315 std::span<T>(local_buffer), pack_fn, request);
316 scatter_rev_end(std::span<const T>(local_buffer), local_data, unpack_fn, op,
317 request);
318 }
319
323 std::int32_t local_buffer_size() const noexcept;
324
328 std::int32_t remote_buffer_size() const noexcept;
329
333 const std::vector<std::int32_t>& local_indices() const noexcept;
334
337 const std::vector<std::int32_t>& remote_indices() const noexcept;
338
342 int bs() const noexcept;
343
344private:
345 // Block size
346 int _bs;
347
348 // Communicator where the source ranks own the indices in the callers
349 // halo, and the destination ranks 'ghost' indices owned by the
350 // caller. I.e.,
351 // - in-edges (src) are from ranks that own my ghosts
352 // - out-edges (dest) go to ranks that 'ghost' my owned indices
353 dolfinx::MPI::Comm _comm0;
354
355 // Communicator where the source ranks have ghost indices that are
356 // owned by the caller, and the destination ranks are the owners of
357 // indices in the callers halo region. I.e.,
358 // - in-edges (src) are from ranks that 'ghost' my owned indicies
359 // - out-edges (dest) are to the owning ranks of my ghost indices
360 dolfinx::MPI::Comm _comm1;
361
362 // Permutation indices used to pack and unpack ghost data (remote)
363 std::vector<std::int32_t> _remote_inds;
364
365 // Number of remote indices (ghosts) for each neighbor process
366 std::vector<int> _sizes_remote;
367
368 // Displacements of remote data for mpi scatter and gather
369 std::vector<int> _displs_remote;
370
371 // Permutation indices used to pack and unpack local shared data
372 // (owned indices that are shared with other processes). Indices are
373 // grouped by neighbor process.
374 std::vector<std::int32_t> _local_inds;
375
376 // Number of local shared indices per neighbor process
377 std::vector<int> _sizes_local;
378
379 // Displacements of local data for mpi scatter and gather
380 std::vector<int> _displs_local;
381};
382} // namespace dolfinx::common
This class represents the distribution index arrays across processes. An index array is a contiguous ...
Definition: IndexMap.h:64
A Scatterer supports the MPI scattering and gathering of data that is associated with a common::Index...
Definition: Scatterer.h:28
void scatter_fwd_end(const std::span< const T > &remote_buffer, std::span< T > remote_data, Functor unpack_fn, MPI_Request &request) const
Complete a non-blocking send from the local owner to process ranks that have the index as a ghost,...
Definition: Scatterer.h:129
void scatter_fwd(const std::span< const T > &local_data, std::span< T > remote_data) const
Scatter data associated with owned indices to ghosting ranks.
Definition: Scatterer.h:152
std::int32_t remote_buffer_size() const noexcept
Buffer size for remote data (ghosts) used in forward and reverse communication.
Definition: Scatterer.cpp:172
std::int32_t local_buffer_size() const noexcept
Size of buffer for local data (owned and shared) used in forward and reverse communication.
Definition: Scatterer.cpp:167
void scatter_rev_end(MPI_Request &request) const
End the reverse scatter communication.
Definition: Scatterer.cpp:157
void scatter_fwd_begin(const std::span< const T > &local_data, std::span< T > local_buffer, std::span< T > remote_buffer, Functor pack_fn, MPI_Request &request) const
Scatter data associated with owned indices to ghosting ranks.
Definition: Scatterer.h:99
void scatter_rev(std::span< T > local_data, const std::span< const T > &remote_data, BinaryOp op)
Scatter data associated with ghost indices to ranks that own the indices.
Definition: Scatterer.h:298
void scatter_fwd_end(MPI_Request &request) const
Complete a non-blocking send from the local owner to process ranks that have the index as a ghost.
Definition: Scatterer.cpp:147
void scatter_rev_end(const std::span< const T > &local_buffer, std::span< T > local_data, Functor unpack_fn, BinaryOp op, MPI_Request &request)
End the reverse scatter communication, and unpack the received local buffer into local data.
Definition: Scatterer.h:283
const std::vector< std::int32_t > & local_indices() const noexcept
Return a vector of local indices (owned) used to pack/unpack local data. These indices are grouped by...
Definition: Scatterer.cpp:177
void scatter_rev_begin(const std::span< const T > &send_buffer, const std::span< T > &recv_buffer, MPI_Request &request) const
Start a non-blocking send of ghost data to ranks that own the data.
Definition: Scatterer.h:201
void scatter_fwd_begin(const std::span< const T > &send_buffer, const std::span< T > &recv_buffer, MPI_Request &request) const
Start a non-blocking send of owned data to ranks that ghost the data.
Definition: Scatterer.h:56
const std::vector< std::int32_t > & remote_indices() const noexcept
Return a vector of remote indices (ghosts) used to pack/unpack ghost data. These indices are grouped ...
Definition: Scatterer.cpp:182
int bs() const noexcept
The number values (block size) to send per index in the common::IndexMap use to create the scatterer.
Definition: Scatterer.cpp:187
void scatter_rev_begin(const std::span< const T > &remote_data, std::span< T > remote_buffer, std::span< T > local_buffer, Functor pack_fn, MPI_Request &request) const
Scatter data associated with ghost indices to owning ranks.
Definition: Scatterer.h:253
Scatterer(const IndexMap &map, int bs)
Create a scatterer.
Definition: Scatterer.cpp:18
Miscellaneous classes, functions and types.