14 #include <dolfinx/common/Timer.h>
15 #include <dolfinx/common/log.h>
16 #include <dolfinx/graph/AdjacencyList.h>
21 #include <type_traits>
25 #define MPICH_IGNORE_CXX_SEEK 1
45 explicit Comm(MPI_Comm
comm,
bool duplicate =
true);
63 MPI_Comm
comm()
const noexcept;
71 int rank(MPI_Comm comm);
75 int size(MPI_Comm comm);
92 const std::int64_t n = N /
size;
93 const std::int64_t r = N %
size;
97 return {
rank * (n + 1),
rank * (n + 1) + n + 1};
99 return {
rank * n + r,
rank * n + r + n};
113 const std::size_t n = N /
size;
114 const std::size_t r = N %
size;
116 if (index < r * (n + 1))
119 return index / (n + 1);
124 return r + (index - r * (n + 1)) / n;
152 const std::span<const int>& edges);
179 const std::span<const int>& edges);
200 template <
typename T>
201 std::pair<std::vector<std::int32_t>, std::vector<T>>
203 std::array<std::int64_t, 2> shape,
204 std::int64_t rank_offset);
227 template <
typename T>
229 MPI_Comm comm,
const std::span<const std::int64_t>& indices,
230 const std::span<const T>& x, std::array<std::int64_t, 2> shape,
231 std::int64_t rank_offset);
256 template <
typename T>
258 const std::span<const std::int64_t>& indices,
259 const std::span<const T>& x,
int shape1);
261 template <
typename T>
267 template <
typename T>
270 if constexpr (std::is_same_v<T, float>)
272 else if constexpr (std::is_same_v<T, double>)
274 else if constexpr (std::is_same_v<T, std::complex<double>>)
275 return MPI_C_DOUBLE_COMPLEX;
276 else if constexpr (std::is_same_v<T, std::complex<float>>)
277 return MPI_C_FLOAT_COMPLEX;
278 else if constexpr (std::is_same_v<T, short int>)
280 else if constexpr (std::is_same_v<T, int>)
282 else if constexpr (std::is_same_v<T, unsigned int>)
284 else if constexpr (std::is_same_v<T, long int>)
286 else if constexpr (std::is_same_v<T, unsigned long>)
287 return MPI_UNSIGNED_LONG;
288 else if constexpr (std::is_same_v<T, long long>)
289 return MPI_LONG_LONG;
290 else if constexpr (std::is_same_v<T, unsigned long long>)
291 return MPI_UNSIGNED_LONG_LONG;
292 else if constexpr (std::is_same_v<T, bool>)
294 else if constexpr (std::is_same_v<T, std::int8_t>)
298 static_assert(!std::is_same_v<T, T>);
302 template <
typename T>
303 std::pair<std::vector<std::int32_t>, std::vector<T>>
305 std::array<std::int64_t, 2> shape,
306 std::int64_t rank_offset)
310 assert(x.size() % shape[1] == 0);
311 const std::int32_t shape0_local = x.size() / shape[1];
313 LOG(2) <<
"Sending data to post offices (distribute_to_postoffice)";
316 std::vector<int> row_to_dest(shape0_local);
317 for (std::int32_t i = 0; i < shape0_local; ++i)
320 row_to_dest[i] = dest;
325 std::vector<std::array<std::int32_t, 2>> dest_to_index;
326 dest_to_index.reserve(shape0_local);
327 for (std::int32_t i = 0; i < shape0_local; ++i)
329 std::size_t idx = i + rank_offset;
331 dest_to_index.push_back({dest, i});
333 std::sort(dest_to_index.begin(), dest_to_index.end());
337 std::vector<int> dest;
338 std::vector<std::int32_t> num_items_per_dest,
339 pos_to_neigh_rank(shape0_local, -1);
341 auto it = dest_to_index.begin();
342 while (it != dest_to_index.end())
344 const int neigh_rank = dest.size();
347 dest.push_back((*it)[0]);
351 = std::find_if(it, dest_to_index.end(),
352 [r = dest.back()](
auto& idx) { return idx[0] != r; });
355 num_items_per_dest.push_back(std::distance(it, it1));
358 for (
auto e = it; e != it1; ++e)
359 pos_to_neigh_rank[(*e)[1]] = neigh_rank;
369 <<
"Number of neighbourhood source ranks in distribute_to_postoffice: "
374 MPI_Dist_graph_create_adjacent(comm, src.size(), src.data(), MPI_UNWEIGHTED,
375 dest.size(), dest.data(), MPI_UNWEIGHTED,
376 MPI_INFO_NULL,
false, &neigh_comm);
379 std::vector<std::int32_t> send_disp = {0};
380 std::partial_sum(num_items_per_dest.begin(), num_items_per_dest.end(),
381 std::back_inserter(send_disp));
384 std::vector<T> send_buffer_data(shape[1] * send_disp.back());
385 std::vector<std::int64_t> send_buffer_index(send_disp.back());
387 std::vector<std::int32_t> send_offsets = send_disp;
388 for (std::int32_t i = 0; i < shape0_local; ++i)
390 if (
int neigh_dest = pos_to_neigh_rank[i]; neigh_dest != -1)
392 std::size_t pos = send_offsets[neigh_dest];
393 send_buffer_index[pos] = i + rank_offset;
394 std::copy_n(std::next(x.begin(), i * shape[1]), shape[1],
395 std::next(send_buffer_data.begin(), shape[1] * pos));
396 ++send_offsets[neigh_dest];
403 std::vector<int> num_items_recv(src.size());
404 num_items_per_dest.reserve(1);
405 num_items_recv.reserve(1);
406 MPI_Neighbor_alltoall(num_items_per_dest.data(), 1, MPI_INT,
407 num_items_recv.data(), 1, MPI_INT, neigh_comm);
410 std::vector<std::int32_t> recv_disp(num_items_recv.size() + 1, 0);
411 std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
412 std::next(recv_disp.begin()));
415 std::vector<std::int64_t> recv_buffer_index(recv_disp.back());
416 MPI_Neighbor_alltoallv(send_buffer_index.data(), num_items_per_dest.data(),
417 send_disp.data(), MPI_INT64_T,
418 recv_buffer_index.data(), num_items_recv.data(),
419 recv_disp.data(), MPI_INT64_T, neigh_comm);
422 MPI_Datatype compound_type;
423 MPI_Type_contiguous(shape[1], dolfinx::MPI::mpi_type<T>(), &compound_type);
424 MPI_Type_commit(&compound_type);
425 std::vector<T> recv_buffer_data(shape[1] * recv_disp.back());
426 MPI_Neighbor_alltoallv(send_buffer_data.data(), num_items_per_dest.data(),
427 send_disp.data(), compound_type,
428 recv_buffer_data.data(), num_items_recv.data(),
429 recv_disp.data(), compound_type, neigh_comm);
430 MPI_Type_free(&compound_type);
431 MPI_Comm_free(&neigh_comm);
433 LOG(2) <<
"Completed send data to post offices.";
437 std::vector<std::int32_t> index_local(recv_buffer_index.size());
438 std::transform(recv_buffer_index.cbegin(), recv_buffer_index.cend(),
439 index_local.begin(), [r0](
auto idx) { return idx - r0; });
441 return {index_local, recv_buffer_data};
444 template <
typename T>
446 MPI_Comm comm,
const std::span<const std::int64_t>& indices,
447 const std::span<const T>& x, std::array<std::int64_t, 2> shape,
448 std::int64_t rank_offset)
451 assert(shape[1] > 0);
455 assert(x.size() % shape[1] == 0);
456 const std::int64_t shape0_local = x.size() / shape[1];
463 comm, x, {shape[0], shape[1]}, rank_offset);
464 assert(post_indices.size() == post_x.size() / shape[1]);
470 std::vector<std::tuple<int, std::int64_t, std::int32_t>> src_to_index;
471 for (std::size_t i = 0; i < indices.size(); ++i)
473 std::size_t idx = indices[i];
475 src_to_index.push_back({src, idx, i});
477 std::sort(src_to_index.begin(), src_to_index.end());
481 std::vector<std::int32_t> num_items_per_src;
482 std::vector<int> src;
484 auto it = src_to_index.begin();
485 while (it != src_to_index.end())
487 src.push_back(std::get<0>(*it));
488 auto it1 = std::find_if(it, src_to_index.end(),
489 [r = src.back()](
auto& idx)
490 { return std::get<0>(idx) != r; });
491 num_items_per_src.push_back(std::distance(it, it1));
498 const std::vector<int> dest
500 LOG(INFO) <<
"Neighbourhood destination ranks from post office in "
501 "distribute_data (rank, num dests, num dests/mpi_size): "
502 <<
rank <<
", " << dest.size() <<
", "
503 <<
static_cast<double>(dest.size()) /
size;
507 MPI_Comm neigh_comm0;
508 MPI_Dist_graph_create_adjacent(comm, dest.size(), dest.data(), MPI_UNWEIGHTED,
509 src.size(), src.data(), MPI_UNWEIGHTED,
510 MPI_INFO_NULL,
false, &neigh_comm0);
513 std::vector<int> num_items_recv(dest.size());
514 num_items_per_src.reserve(1);
515 num_items_recv.reserve(1);
516 MPI_Neighbor_alltoall(num_items_per_src.data(), 1, MPI_INT,
517 num_items_recv.data(), 1, MPI_INT, neigh_comm0);
520 std::vector<std::int32_t> send_disp = {0};
521 std::partial_sum(num_items_per_src.begin(), num_items_per_src.end(),
522 std::back_inserter(send_disp));
523 std::vector<std::int32_t> recv_disp = {0};
524 std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
525 std::back_inserter(recv_disp));
529 assert(send_disp.back() == (
int)src_to_index.size());
530 std::vector<std::int64_t> send_buffer_index(src_to_index.size());
531 std::transform(src_to_index.cbegin(), src_to_index.cend(),
532 send_buffer_index.begin(),
533 [](
auto& x) { return std::get<1>(x); });
536 std::vector<std::int64_t> recv_buffer_index(recv_disp.back());
537 MPI_Neighbor_alltoallv(send_buffer_index.data(), num_items_per_src.data(),
538 send_disp.data(), MPI_INT64_T,
539 recv_buffer_index.data(), num_items_recv.data(),
540 recv_disp.data(), MPI_INT64_T, neigh_comm0);
542 MPI_Comm_free(&neigh_comm0);
550 const std::array<std::int64_t, 2> postoffice_range
552 std::vector<std::int32_t> post_indices_map(
553 postoffice_range[1] - postoffice_range[0], -1);
554 for (std::size_t i = 0; i < post_indices.size(); ++i)
556 assert(post_indices[i] < (
int)post_indices_map.size());
557 post_indices_map[post_indices[i]] = i;
561 std::vector<T> send_buffer_data(shape[1] * recv_disp.back());
562 for (std::size_t p = 0; p < recv_disp.size() - 1; ++p)
564 int offset = recv_disp[p];
565 for (std::int32_t i = recv_disp[p]; i < recv_disp[p + 1]; ++i)
567 std::int64_t index = recv_buffer_index[i];
568 if (index >= rank_offset and index < (rank_offset + shape0_local))
571 std::int32_t local_index = index - rank_offset;
572 std::copy_n(std::next(x.begin(), shape[1] * local_index), shape[1],
573 std::next(send_buffer_data.begin(), shape[1] * offset));
578 auto local_index = index - postoffice_range[0];
579 std::int32_t pos = post_indices_map[local_index];
581 std::copy_n(std::next(post_x.begin(), shape[1] * pos), shape[1],
582 std::next(send_buffer_data.begin(), shape[1] * offset));
589 MPI_Dist_graph_create_adjacent(comm, src.size(), src.data(), MPI_UNWEIGHTED,
590 dest.size(), dest.data(), MPI_UNWEIGHTED,
591 MPI_INFO_NULL,
false, &neigh_comm0);
593 MPI_Datatype compound_type0;
594 MPI_Type_contiguous(shape[1], dolfinx::MPI::mpi_type<T>(), &compound_type0);
595 MPI_Type_commit(&compound_type0);
597 std::vector<T> recv_buffer_data(shape[1] * send_disp.back());
598 MPI_Neighbor_alltoallv(send_buffer_data.data(), num_items_recv.data(),
599 recv_disp.data(), compound_type0,
600 recv_buffer_data.data(), num_items_per_src.data(),
601 send_disp.data(), compound_type0, neigh_comm0);
603 MPI_Type_free(&compound_type0);
604 MPI_Comm_free(&neigh_comm0);
606 std::vector<std::int32_t> index_pos_to_buffer(indices.size(), -1);
607 for (std::size_t i = 0; i < src_to_index.size(); ++i)
608 index_pos_to_buffer[std::get<2>(src_to_index[i])] = i;
611 std::vector<T> x_new(shape[1] * indices.size());
612 for (std::size_t i = 0; i < indices.size(); ++i)
614 const std::int64_t index = indices[i];
615 if (index >= rank_offset and index < (rank_offset + shape0_local))
618 auto local_index = index - rank_offset;
619 std::copy_n(std::next(x.begin(), shape[1] * local_index), shape[1],
620 std::next(x_new.begin(), shape[1] * i));
627 auto local_index = index - postoffice_range[0];
628 std::int32_t pos = post_indices_map[local_index];
630 std::copy_n(std::next(post_x.begin(), shape[1] * pos), shape[1],
631 std::next(x_new.begin(), shape[1] * i));
636 std::int32_t pos = index_pos_to_buffer[i];
638 std::copy_n(std::next(recv_buffer_data.begin(), shape[1] * pos),
639 shape[1], std::next(x_new.begin(), shape[1] * i));
647 template <
typename T>
649 const std::span<const std::int64_t>& indices,
650 const std::span<const T>& x,
int shape1)
653 assert(x.size() % shape1 == 0);
654 const std::int64_t shape0_local = x.size() / shape1;
656 std::int64_t shape0(0), rank_offset(0);
657 MPI_Allreduce(&shape0_local, &shape0, 1, MPI_INT64_T, MPI_SUM, comm);
658 MPI_Exscan(&shape0_local, &rank_offset, 1, MPI_INT64_T, MPI_SUM, comm);
A duplicate MPI communicator and manage lifetime of the communicator.
Definition: MPI.h:42
Comm(MPI_Comm comm, bool duplicate=true)
Duplicate communicator and wrap duplicate.
Definition: MPI.cpp:12
~Comm()
Destructor (frees wrapped communicator)
Definition: MPI.cpp:39
MPI_Comm comm() const noexcept
Return the underlying MPI_Comm object.
Definition: MPI.cpp:73
A timer can be used for timing tasks. The basic usage is.
Definition: Timer.h:31
MPI support functionality.
Definition: MPI.h:30
std::pair< std::vector< std::int32_t >, std::vector< T > > distribute_to_postoffice(MPI_Comm comm, const std::span< const T > &x, std::array< std::int64_t, 2 > shape, std::int64_t rank_offset)
Distribute row data to 'post office' ranks.
Definition: MPI.h:304
std::vector< T > distribute_from_postoffice(MPI_Comm comm, const std::span< const std::int64_t > &indices, const std::span< const T > &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:445
std::vector< int > compute_graph_edges_nbx(MPI_Comm comm, const std::span< const int > &edges)
Determine incoming graph edges using the NBX consensus algorithm.
Definition: MPI.cpp:151
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:108
std::vector< int > compute_graph_edges_pcx(MPI_Comm comm, const std::span< const int > &edges)
Determine incoming graph edges using the PCX consensus algorithm.
Definition: MPI.cpp:91
std::vector< T > distribute_data(MPI_Comm comm, const std::span< const std::int64_t > &indices, const std::span< const T > &x, int shape1)
Distribute rows of a rectangular data array to ranks where they are required (scalable version).
Definition: MPI.h:648
int size(MPI_Comm comm)
Return size of the group (number of processes) associated with the communicator.
Definition: MPI.cpp:83
int rank(MPI_Comm comm)
Return process rank for the communicator.
Definition: MPI.cpp:75
constexpr MPI_Datatype mpi_type()
MPI Type.
Definition: MPI.h:268
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:84
tag
MPI communication tags.
Definition: MPI.h:34