589 const U& x, std::array<std::int64_t, 2> shape,
590 std::int64_t rank_offset)
592 assert(rank_offset >= 0 or x.empty());
593 using T = std::ranges::range_value_t<U>;
596 assert(shape[1] > 0);
600 assert(x.size() % shape[1] == 0);
601 const std::int64_t shape0_local = x.size() / shape[1];
609 auto [send_dest, num_items_per_send_dest, pos_to_neigh_rank]
610 = impl::postoffice_plan(
size,
rank, shape0_local, shape[0], rank_offset);
615 std::vector<std::tuple<int, std::int64_t, std::int32_t>> src_to_index;
616 for (std::size_t i = 0; i < indices.size(); ++i)
618 std::int64_t idx = indices[i];
619 if (idx >= rank_offset and idx < rank_offset + shape0_local)
622 src_to_index.push_back({src, idx, i});
628 std::vector<std::int32_t> perm(src_to_index.size());
629 std::iota(perm.begin(), perm.end(), 0);
631 {
return std::get<0>(src_to_index[i]); });
632 std::vector<std::tuple<int, std::int64_t, std::int32_t>> sorted(
633 src_to_index.size());
634 for (std::size_t i = 0; i < perm.size(); ++i)
635 sorted[i] = src_to_index[perm[i]];
636 src_to_index = std::move(sorted);
641 std::vector<std::int32_t> num_items_per_src;
642 std::vector<int> src;
644 auto it = src_to_index.begin();
645 while (it != src_to_index.end())
647 src.push_back(std::get<0>(*it));
648 auto it1 = std::ranges::find_if(it, src_to_index.end(),
649 [r = src.back()](
auto& idx)
650 { return std::get<0>(idx) != r; });
651 num_items_per_src.push_back(std::ranges::distance(it, it1));
661 comm, send_dest,
static_cast<int>(tag::consensus_nbx), src,
662 static_cast<int>(tag::consensus_nbx) + 1);
664 "Neighbourhood destination ranks from post office in "
665 "distribute_data (rank, num dests, num dests/mpi_size): {}, {}, {}",
666 rank, dest.size(),
static_cast<double>(dest.size()) /
size);
670 auto [post_indices, post_x] = impl::postoffice_exchange(
671 comm, x, {shape[0], shape[1]}, rank_offset, send_dest,
672 num_items_per_send_dest, pos_to_neigh_rank, post_src);
673 assert(post_indices.size() == post_x.size() / shape[1]);
677 MPI_Comm neigh_comm0;
678 int err = MPI_Dist_graph_create_adjacent(
679 comm, dest.size(), dest.data(), MPI_UNWEIGHTED, src.size(), src.data(),
680 MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &neigh_comm0);
684 std::vector<int> num_items_recv(dest.size());
685 num_items_per_src.reserve(1);
686 num_items_recv.reserve(1);
687 err = MPI_Neighbor_alltoall(num_items_per_src.data(), 1, MPI_INT,
688 num_items_recv.data(), 1, MPI_INT, neigh_comm0);
692 std::vector<std::int32_t> send_disp{0};
693 std::partial_sum(num_items_per_src.begin(), num_items_per_src.end(),
694 std::back_inserter(send_disp));
695 std::vector<std::int32_t> recv_disp = {0};
696 std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
697 std::back_inserter(recv_disp));
701 assert(send_disp.back() ==
static_cast<int>(src_to_index.size()));
702 std::vector<std::int64_t> send_buffer_index(src_to_index.size());
703 std::ranges::transform(src_to_index, send_buffer_index.begin(),
704 [](
auto x) { return std::get<1>(x); });
707 std::vector<std::int64_t> recv_buffer_index(recv_disp.back());
708 err = MPI_Neighbor_alltoallv(
709 send_buffer_index.data(), num_items_per_src.data(), send_disp.data(),
710 MPI_INT64_T, recv_buffer_index.data(), num_items_recv.data(),
711 recv_disp.data(), MPI_INT64_T, neigh_comm0);
714 err = MPI_Comm_free(&neigh_comm0);
723 const std::array<std::int64_t, 2> postoffice_range
725 std::vector<std::int32_t> post_indices_map(
726 postoffice_range[1] - postoffice_range[0], -1);
727 for (std::size_t i = 0; i < post_indices.size(); ++i)
729 assert(post_indices[i] <
static_cast<int>(post_indices_map.size()));
730 post_indices_map[post_indices[i]] = i;
734 std::vector<T> send_buffer_data(shape[1] * recv_disp.back());
735 for (std::int32_t i = 0; i < recv_disp.back(); ++i)
737 std::int64_t index = recv_buffer_index[i];
738 if (index >= rank_offset and index < (rank_offset + shape0_local))
741 std::int32_t local_index = index - rank_offset;
742 std::copy_n(std::next(x.begin(), shape[1] * local_index), shape[1],
743 std::next(send_buffer_data.begin(), shape[1] * i));
748 std::int64_t local_index = index - postoffice_range[0];
749 std::int32_t pos = post_indices_map[local_index];
751 std::copy_n(std::next(post_x.begin(), shape[1] * pos), shape[1],
752 std::next(send_buffer_data.begin(), shape[1] * i));
756 err = MPI_Dist_graph_create_adjacent(
757 comm, src.size(), src.data(), MPI_UNWEIGHTED, dest.size(), dest.data(),
758 MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &neigh_comm0);
761 MPI_Datatype compound_type0;
763 MPI_Type_commit(&compound_type0);
765 std::vector<T> recv_buffer_data(shape[1] * send_disp.back());
766 err = MPI_Neighbor_alltoallv(
767 send_buffer_data.data(), num_items_recv.data(), recv_disp.data(),
768 compound_type0, recv_buffer_data.data(), num_items_per_src.data(),
769 send_disp.data(), compound_type0, neigh_comm0);
772 err = MPI_Type_free(&compound_type0);
774 err = MPI_Comm_free(&neigh_comm0);
777 std::vector<std::int32_t> index_pos_to_buffer(indices.size(), -1);
778 for (std::size_t i = 0; i < src_to_index.size(); ++i)
779 index_pos_to_buffer[std::get<2>(src_to_index[i])] = i;
782 std::vector<T> x_new(shape[1] * indices.size());
783 for (std::size_t i = 0; i < indices.size(); ++i)
785 const std::int64_t index = indices[i];
786 if (index >= rank_offset and index < (rank_offset + shape0_local))
789 std::int64_t local_index = index - rank_offset;
790 std::copy_n(std::next(x.begin(), shape[1] * local_index), shape[1],
791 std::next(x_new.begin(), shape[1] * i));
793 else if (std::int32_t pos = index_pos_to_buffer[i]; pos != -1)
798 std::copy_n(std::next(recv_buffer_data.begin(), shape[1] * pos), shape[1],
799 std::next(x_new.begin(), shape[1] * i));
804 std::int64_t local_index = index - postoffice_range[0];
805 std::int32_t bag_pos = post_indices_map[local_index];
806 assert(bag_pos != -1);
807 std::copy_n(std::next(post_x.begin(), shape[1] * bag_pos), shape[1],
808 std::next(x_new.begin(), shape[1] * i));