Skip to content

Commit

Permalink
Remove NBX call from sub-IndexMap creation (#3392)
Browse files Browse the repository at this point in the history 
* Remove nbx call

* Simplify (and verify)

* Start the complicated process of sending new dest ranks after determining ownership

* Communicate new dest ranks

* Fix code (with verification)

* Remove verification

* Remove iostream

* Revert loop we do not use

* Add back linebreak

* Remove duplicate sort

* Fix source and dest confusion in alltoall

* Remove nbx comparison

* Fix typo

* Update cpp/dolfinx/common/IndexMap.cpp

Co-authored-by: Joe Dean <jpd62@cam.ac.uk>

* Update cpp/dolfinx/common/IndexMap.cpp

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Update cpp/dolfinx/common/IndexMap.cpp

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Update cpp/dolfinx/common/IndexMap.cpp

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Simplifications

* Logic simplifications

* Small edits

---------

Co-authored-by: Garth N. Wells <gnw20@cam.ac.uk>
Co-authored-by: Joe Dean <jpd62@cam.ac.uk>
  • Loading branch information
@jorgensd @garth-wells @jpdean
3 people authored Sep 24, 2024
1 parent 6021fd5 commit 0da58d1
Showing 1 changed file with 116 additions and 33 deletions.
149 changes: 116 additions & 33 deletions cpp/dolfinx/common/IndexMap.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
// Copyright (C) 2015-2022 Chris Richardson, Garth N. Wells and Igor Baratta
// Copyright (C) 2015-2024 Chris Richardson, Garth N. Wells, Igor Baratta,
// Joseph P. Dean and Jørgen S. Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
Expand Down Expand Up @@ -98,8 +99,7 @@ communicate_ghosts_to_owners(MPI_Comm comm, std::span<const int> src,
{
auto it = std::ranges::lower_bound(src, owners[i]);
assert(it != src.end() and *it == owners[i]);
std::size_t r = std::distance(src.begin(), it);
if (include_ghost[i])
if (std::size_t r = std::distance(src.begin(), it); include_ghost[i])
{
send_data[r].push_back(ghosts[i]);
pos_to_ghost[r].push_back(i);
Expand Down Expand Up @@ -191,16 +191,20 @@ compute_submap_indices(const IndexMap& imap,
std::span(is_in_submap.cbegin() + imap.size_local(),
is_in_submap.cend()));

// --- Step 2 ---: Create a map from the indices in `recv_indices` (i.e.
// indices owned by this process that are in `indices` on other processes) to
// their owner in the submap. This is required since not all indices in
// `recv_indices` will necessarily be in `indices` on this process, and thus
// other processes must own them in the submap.

// --- Step 2 ---: Create a map from the indices in `recv_indices`
// (i.e. indices owned by this process that are in `indices` on other
// processes) to their owner in the submap. This is required since not
// all indices in `recv_indices` will necessarily be in `indices` on
// this process, and thus other processes must own them in the submap.
// If ownership of received index doesn't change, then this process
// has the receiving rank as a destination
std::vector<int> recv_owners(send_disp.back());
std::vector<int> submap_dest;
submap_dest.reserve(1);
const int rank = dolfinx::MPI::rank(imap.comm());
{
// Flag to track if the owner of any indices have changed in the submap
// Flag to track if the owner of any indices have changed in the
// submap
bool owners_changed = false;

// Create a map from (global) indices in `recv_indices` to a list of
Expand All @@ -221,11 +225,14 @@ compute_submap_indices(const IndexMap& imap,
assert(idx_local >= 0);
assert(idx_local < local_range[1]);

// Check if index is included in the submap on this process. If so,
// this process remains its owner in the submap. Otherwise,
// Check if index is included in the submap on this process. If
// so, this process remains its owner in the submap. Otherwise,
// add the process that sent it to the list of possible owners.
if (is_in_submap[idx_local])
{
global_idx_to_possible_owner.push_back({idx, rank});
submap_dest.push_back(dest[i]);
}
else
{
owners_changed = true;
Expand All @@ -239,18 +246,69 @@ compute_submap_indices(const IndexMap& imap,

std::ranges::sort(global_idx_to_possible_owner);

// Choose the submap owner for each index in `recv_indices`
// Choose the submap owner for each index in `recv_indices` and pack
// destination ranks for each process that has received new indices.
// During ownership determination, we know what other processes
// requires this index, and add them to the destination set.
std::vector<int> send_owners;
send_owners.reserve(recv_indices.size());
for (auto idx : recv_indices)
send_owners.reserve(1);
std::vector<int> new_owner_dest_ranks;
new_owner_dest_ranks.reserve(1);
std::vector<int> new_owner_dest_ranks_offsets(recv_sizes.size() + 1, 0);
std::vector<std::int32_t> new_owner_dest_ranks_sizes(recv_sizes.size());
new_owner_dest_ranks_sizes.reserve(1);
for (std::size_t i = 0; i < recv_sizes.size(); ++i)
{
// NOTE: Could choose new owner in a way that is is better for
// load balancing, though the impact is probably only very small
auto it = std::ranges::lower_bound(global_idx_to_possible_owner, idx,
std::ranges::less(),
[](auto e) { return e.first; });
assert(it != global_idx_to_possible_owner.end() and it->first == idx);
send_owners.push_back(it->second);
for (int j = recv_disp[i]; j < recv_disp[i + 1]; ++j)
{
std::int64_t idx = recv_indices[j];

// NOTE: Could choose new owner in a way that is is better for
// load balancing, though the impact is probably only very small
auto it = std::ranges::lower_bound(global_idx_to_possible_owner, idx,
std::ranges::less(),
[](auto e) { return e.first; });
assert(it != global_idx_to_possible_owner.end() and it->first == idx);
send_owners.push_back(it->second);

// If rank that sent this ghost is the submap owner, send all
// other ranks
if (it->second == dest[i])
{
// Find upper limit of recv index and pack all ranks from
// ownership determination (except new owner) as dest ranks
auto it_upper = std::ranges::upper_bound(
it, global_idx_to_possible_owner.end(), idx, std::ranges::less(),
[](auto e) { return e.first; });
std::transform(std::next(it), it_upper,
std::back_inserter(new_owner_dest_ranks),
[](auto e) { return e.second; });
}
}

// Remove duplicate new dest ranks from recv process. The new
// owning process can have taken ownership of multiple indices
// from the same rank.
if (auto dest_begin = std::next(new_owner_dest_ranks.begin(),
new_owner_dest_ranks_offsets[i]);
dest_begin != new_owner_dest_ranks.end())
{
std::ranges::sort(dest_begin, new_owner_dest_ranks.end());
auto [unique_end, range_end]
= std::ranges::unique(dest_begin, new_owner_dest_ranks.end());
new_owner_dest_ranks.erase(unique_end, range_end);

std::size_t num_unique_dest_ranks
= std::distance(dest_begin, unique_end);
new_owner_dest_ranks_sizes[i] = num_unique_dest_ranks;
new_owner_dest_ranks_offsets[i + 1]
= new_owner_dest_ranks_offsets[i] + num_unique_dest_ranks;
}
else
{
new_owner_dest_ranks_sizes[i] = 0;
new_owner_dest_ranks_offsets[i + 1] = new_owner_dest_ranks_offsets[i];
}
}

// Create neighbourhood comm (owner -> ghost)
Expand All @@ -267,6 +325,36 @@ compute_submap_indices(const IndexMap& imap,
comm1);
dolfinx::MPI::check_error(imap.comm(), ierr);

// Communicate number of received dest_ranks from each process
std::vector<int> recv_dest_ranks_sizes(imap.src().size());
recv_dest_ranks_sizes.reserve(1);
ierr = MPI_Neighbor_alltoall(new_owner_dest_ranks_sizes.data(), 1, MPI_INT,
recv_dest_ranks_sizes.data(), 1, MPI_INT,
comm1);
dolfinx::MPI::check_error(imap.comm(), ierr);

// Communicate new dest ranks
std::vector<int> recv_dest_rank_disp(imap.src().size() + 1, 0);
std::partial_sum(recv_dest_ranks_sizes.begin(), recv_dest_ranks_sizes.end(),
std::next(recv_dest_rank_disp.begin()));
std::vector<int> recv_dest_ranks(recv_dest_rank_disp.back());
recv_dest_ranks.reserve(1);
ierr = MPI_Neighbor_alltoallv(
new_owner_dest_ranks.data(), new_owner_dest_ranks_sizes.data(),
new_owner_dest_ranks_offsets.data(), MPI_INT, recv_dest_ranks.data(),
recv_dest_ranks_sizes.data(), recv_dest_rank_disp.data(), MPI_INT,
comm1);
dolfinx::MPI::check_error(imap.comm(), ierr);

// Append new submap dest ranks and remove duplicates
std::ranges::copy(recv_dest_ranks, std::back_inserter(submap_dest));
std::ranges::sort(submap_dest);
{
auto [unique_end, range_end] = std::ranges::unique(submap_dest);
submap_dest.erase(unique_end, range_end);
submap_dest.shrink_to_fit();
}

// Free the communicator
ierr = MPI_Comm_free(&comm1);
dolfinx::MPI::check_error(imap.comm(), ierr);
Expand All @@ -278,6 +366,7 @@ compute_submap_indices(const IndexMap& imap,
// Local indices (w.r.t. original map) owned by this process in the
// submap
std::vector<std::int32_t> submap_owned;
submap_owned.reserve(indices.size());

// Local indices (w.r.t. original map) ghosted by this process in the
// submap
Expand All @@ -289,9 +378,9 @@ compute_submap_indices(const IndexMap& imap,

{
// Add owned indices to submap_owned
for (std::int32_t v : indices)
if (v < imap.size_local())
submap_owned.push_back(v);
std::copy_if(
indices.begin(), indices.end(), std::back_inserter(submap_owned),
[local_size = imap.size_local()](auto i) { return i < local_size; });

// FIXME: Could just create when making send_indices
std::vector<std::int32_t> send_indices_local(send_indices.size());
Expand Down Expand Up @@ -348,12 +437,6 @@ compute_submap_indices(const IndexMap& imap,
submap_ghost = std::move(submap_ghost1);
}

// Compute submap destination ranks
// FIXME: Remove call to NBX
std::vector<int> submap_dest
= dolfinx::MPI::compute_graph_edges_nbx(imap.comm(), submap_src);
std::ranges::sort(submap_dest);

return {std::move(submap_owned), std::move(submap_ghost),
std::move(submap_ghost_owners), std::move(submap_src),
std::move(submap_dest)};
Expand Down Expand Up @@ -1078,8 +1161,8 @@ graph::AdjacencyList<int> IndexMap::index_to_dest_ranks() const
dolfinx::MPI::check_error(_comm.comm(), ierr);

// Prepare displacement vectors
std::vector<int> send_disp(dest.size() + 1, 0),
recv_disp(src.size() + 1, 0);
std::vector<int> send_disp(dest.size() + 1, 0);
std::vector<int> recv_disp(src.size() + 1, 0);
std::partial_sum(send_sizes.begin(), send_sizes.end(),
std::next(send_disp.begin()));
std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
Expand Down

0 comments on commit 0da58d1

Please sign in to comment.