50 : _bs(
bs), _remote_inds(0, alloc), _local_inds(0, alloc),
51 _src(map.src().begin(), map.src().end()),
52 _dest(map.dest().begin(), map.dest().end())
58 assert(std::ranges::is_sorted(_src));
59 assert(std::ranges::is_sorted(_dest));
65 MPI_Dist_graph_create_adjacent(
66 map.
comm(), _src.size(), _src.data(), MPI_UNWEIGHTED, _dest.size(),
67 _dest.data(), MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &comm0);
71 MPI_Dist_graph_create_adjacent(
72 map.
comm(), _dest.size(), _dest.data(), MPI_UNWEIGHTED, _src.size(),
73 _src.data(), MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &comm1);
77 std::span owners = map.
owners();
78 std::vector<std::int32_t> perm(owners.size());
79 std::iota(perm.begin(), perm.end(), 0);
84 std::span ghosts = map.
ghosts();
85 std::vector<int> owners_sorted(owners.size());
86 std::vector<std::int64_t> ghosts_sorted(owners.size());
87 std::ranges::transform(perm, owners_sorted.begin(),
88 [&owners](
auto idx) { return owners[idx]; });
89 std::ranges::transform(perm, ghosts_sorted.begin(),
90 [&ghosts](
auto idx) { return ghosts[idx]; });
98 _sizes_remote.resize(_src.size(), 0);
99 _displs_remote.resize(_src.size() + 1, 0);
100 std::vector<std::int32_t>::iterator begin = owners_sorted.begin();
101 for (std::size_t i = 0; i < _src.size(); i++)
103 auto upper = std::upper_bound(begin, owners_sorted.end(), _src[i]);
104 int num_ind = std::distance(begin, upper);
105 _displs_remote[i + 1] = _displs_remote[i] + num_ind;
106 _sizes_remote[i] = num_ind;
117 _sizes_local.resize(_dest.size());
118 _displs_local.resize(_sizes_local.size() + 1);
119 _sizes_remote.reserve(1);
120 _sizes_local.reserve(1);
121 MPI_Neighbor_alltoall(_sizes_remote.data(), 1, MPI_INT32_T,
122 _sizes_local.data(), 1, MPI_INT32_T, _comm1.comm());
123 std::partial_sum(_sizes_local.begin(), _sizes_local.end(),
124 std::next(_displs_local.begin()));
126 assert((std::int32_t)ghosts_sorted.size() == _displs_remote.back());
127 assert((std::int32_t)ghosts_sorted.size() == _displs_remote.back());
131 std::vector<std::int64_t> recv_buffer(_displs_local.back(), 0);
132 MPI_Neighbor_alltoallv(ghosts_sorted.data(), _sizes_remote.data(),
133 _displs_remote.data(), MPI_INT64_T,
134 recv_buffer.data(), _sizes_local.data(),
135 _displs_local.data(), MPI_INT64_T, _comm1.comm());
137 const std::array<std::int64_t, 2> range = map.
local_range();
140 std::ranges::for_each(recv_buffer, [range](
auto idx)
141 { assert(idx >= range[0] and idx < range[1]); });
146 auto rescale = [](
auto& x,
int bs)
148 std::ranges::transform(x, x.begin(), [
bs](
auto e) { return e *= bs; });
150 rescale(_sizes_local,
bs);
151 rescale(_displs_local,
bs);
152 rescale(_sizes_remote,
bs);
153 rescale(_displs_remote,
bs);
158 _local_inds = std::vector<std::int32_t, allocator_type>(
159 recv_buffer.size() * _bs, alloc);
160 std::int64_t offset = range[0] * _bs;
161 for (std::size_t i = 0; i < recv_buffer.size(); i++)
162 for (
int j = 0; j < _bs; j++)
163 _local_inds[i * _bs + j] = (recv_buffer[i] * _bs + j) - offset;
167 = std::vector<std::int32_t, allocator_type>(perm.size() * _bs, alloc);
168 for (std::size_t i = 0; i < perm.size(); i++)
169 for (
int j = 0; j < _bs; j++)
170 _remote_inds[i * _bs + j] = perm[i] * _bs + j;
195 std::span<T> recv_buffer,
196 std::span<MPI_Request> requests,
200 if (_sizes_local.empty() and _sizes_remote.empty())
207 assert(requests.size() == std::size_t(1));
208 MPI_Ineighbor_alltoallv(send_buffer.data(), _sizes_local.data(),
210 recv_buffer.data(), _sizes_remote.data(),
212 _comm0.comm(), requests.data());
217 assert(requests.size() == _dest.size() + _src.size());
218 for (std::size_t i = 0; i < _src.size(); i++)
220 MPI_Irecv(recv_buffer.data() + _displs_remote[i], _sizes_remote[i],
225 for (std::size_t i = 0; i < _dest.size(); i++)
227 MPI_Isend(send_buffer.data() + _displs_local[i], _sizes_local[i],
229 &requests[i + _src.size()]);
234 throw std::runtime_error(
"Scatter::type not recognized");
338 std::span<T> remote_data)
const
340 std::vector<MPI_Request> requests(1, MPI_REQUEST_NULL);
343 auto pack_fn = [](
auto&& in,
auto&& idx,
auto&& out)
345 for (std::size_t i = 0; i < idx.size(); ++i)
349 std::span<T>(remote_buffer), pack_fn,
350 std::span<MPI_Request>(requests));
352 auto unpack_fn = [](
auto&& in,
auto&& idx,
auto&& out,
auto op)
354 for (std::size_t i = 0; i < idx.size(); ++i)
355 out[idx[i]] = op(out[idx[i]], in[i]);
358 scatter_fwd_end(std::span<const T>(remote_buffer), remote_data, unpack_fn,
359 std::span<MPI_Request>(requests));
390 std::span<T> recv_buffer,
391 std::span<MPI_Request> requests,
395 if (_sizes_local.empty() and _sizes_remote.empty())
404 assert(requests.size() == 1);
405 MPI_Ineighbor_alltoallv(
406 send_buffer.data(), _sizes_remote.data(), _displs_remote.data(),
408 _displs_local.data(),
MPI::mpi_t<T>, _comm1.comm(), &requests[0]);
413 assert(requests.size() == _dest.size() + _src.size());
415 for (std::size_t i = 0; i < _dest.size(); i++)
417 MPI_Irecv(recv_buffer.data() + _displs_local[i], _sizes_local[i],
424 for (std::size_t i = 0; i < _src.size(); i++)
426 MPI_Isend(send_buffer.data() + _displs_remote[i], _sizes_remote[i],
428 &requests[i + _dest.size()]);
433 throw std::runtime_error(
"Scatter::type not recognized");
535 void scatter_rev(std::span<T> local_data, std::span<const T> remote_data,
540 auto pack_fn = [](
auto&& in,
auto&& idx,
auto&& out)
542 for (std::size_t i = 0; i < idx.size(); ++i)
545 auto unpack_fn = [](
auto&& in,
auto&& idx,
auto&& out,
auto op)
547 for (std::size_t i = 0; i < idx.size(); ++i)
548 out[idx[i]] = op(out[idx[i]], in[i]);
550 std::vector<MPI_Request> request(1, MPI_REQUEST_NULL);
552 std::span<T>(local_buffer), pack_fn,
553 std::span<MPI_Request>(request));
554 scatter_rev_end(std::span<const T>(local_buffer), local_data, unpack_fn, op,
555 std::span<MPI_Request>(request));