Loading...
Searching...
No Matches
utils.h
// Copyright (C) 2012-2024 Chris N. Richardson, Garth N. Wells, Jørgen S. Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once
#include <array>
#include <basix/mdspan.hpp>
#include <dolfinx/common/types.h>
#include <dolfinx/fem/ElementDofLayout.h>
#include <dolfinx/mesh/Topology.h>
#include <mpi.h>
#include <span>
#include <utility>
#include <vector>
namespace dolfinx
{
namespace io
{
template <typename T>
std::pair<std::vector<std::int32_t>, std::vector<T>> distribute_entity_data(
const mesh::Topology& topology, std::span<const std::int64_t> nodes_g,
std::int64_t num_nodes_g, const fem::ElementDofLayout& cmap_dof_layout,
md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> xdofmap,
int entity_dim,
md::mdspan<const std::int64_t, md::dextents<std::size_t, 2>> entities,
std::span<const T> data)
{
assert(entities.extent(0) == data.size());
spdlog::info("XDMF distribute entity data");
mesh::CellType cell_type = topology.cell_type();
// Get layout of dofs on 0th cell entity of dimension entity_dim
std::vector<int> cell_vertex_dofs;
for (int i = 0; i < mesh::cell_num_entities(cell_type, 0); ++i)
{
const std::vector<int>& local_index = cmap_dof_layout.entity_dofs(0, i);
assert(local_index.size() == 1);
cell_vertex_dofs.push_back(local_index[0]);
}
// -- A. Convert from list of entities by 'nodes' to list of entities
// by 'vertex nodes'
auto to_vertex_entities
= [](const fem::ElementDofLayout& cmap_dof_layout, int entity_dim,
std::span<const int> cell_vertex_dofs, mesh::CellType cell_type,
auto entities)
{
// Use ElementDofLayout of the cell to get vertex dof indices (local
// to a cell), i.e. build a map from local vertex index to associated
// local dof index
const std::vector<int>& entity_layout
= cmap_dof_layout.entity_closure_dofs(entity_dim, 0);
std::vector<int> entity_vertex_dofs;
for (std::size_t i = 0; i < cell_vertex_dofs.size(); ++i)
{
auto it = std::find(entity_layout.begin(), entity_layout.end(),
cell_vertex_dofs[i]);
if (it != entity_layout.end())
entity_vertex_dofs.push_back(std::distance(entity_layout.begin(), it));
}
const std::size_t num_vert_per_e = mesh::cell_num_entities(
mesh::cell_entity_type(cell_type, entity_dim, 0), 0);
assert(entities.extent(1) == entity_layout.size());
std::vector<std::int64_t> entities_v(entities.extent(0) * num_vert_per_e);
for (std::size_t e = 0; e < entities.extent(0); ++e)
{
std::span entity(entities_v.data() + e * num_vert_per_e, num_vert_per_e);
for (std::size_t i = 0; i < num_vert_per_e; ++i)
entity[i] = entities(e, entity_vertex_dofs[i]);
std::ranges::sort(entity);
}
std::array shape{entities.extent(0), num_vert_per_e};
return std::pair(std::move(entities_v), shape);
};
const auto [entities_v_b, shapev] = to_vertex_entities(
cmap_dof_layout, entity_dim, cell_vertex_dofs, cell_type, entities);
md::mdspan<const std::int64_t, md::dextents<std::size_t, 2>> entities_v(
entities_v_b.data(), shapev);
MPI_Comm comm = topology.comm();
MPI_Datatype compound_type;
MPI_Type_contiguous(entities_v.extent(1), MPI_INT64_T, &compound_type);
MPI_Type_commit(&compound_type);
// -- B. Send entities and entity data to postmaster
auto send_entities_to_postmaster
= [](MPI_Comm comm, MPI_Datatype compound_type, std::int64_t num_nodes_g,
auto entities, std::span<const T> data)
{
const int size = dolfinx::MPI::size(comm);
// Determine destination by index of first vertex
std::vector<int> dest0;
dest0.reserve(entities.extent(0));
for (std::size_t e = 0; e < entities.extent(0); ++e)
{
dest0.push_back(
dolfinx::MPI::index_owner(size, entities(e, 0), num_nodes_g));
}
std::vector<int> perm(dest0.size());
std::iota(perm.begin(), perm.end(), 0);
std::ranges::sort(perm, [&dest0](auto x0, auto x1)
{ return dest0[x0] < dest0[x1]; });
// Note: dest[perm[i]] is ordered with increasing i
// Build list of neighbour dest ranks and count number of entities to
// send to each post office
std::vector<int> dest;
std::vector<std::int32_t> num_items_send;
{
auto it = perm.begin();
while (it != perm.end())
{
dest.push_back(dest0[*it]);
auto it1
= std::find_if(it, perm.end(), [&dest0, r = dest.back()](auto idx)
{ return dest0[idx] != r; });
num_items_send.push_back(std::distance(it, it1));
it = it1;
}
}
// Compute send displacements
std::vector<int> send_disp(num_items_send.size() + 1, 0);
std::partial_sum(num_items_send.begin(), num_items_send.end(),
std::next(send_disp.begin()));
// Determine src ranks. Sort ranks so that ownership determination is
// deterministic for a given number of ranks.
std::vector<int> src = dolfinx::MPI::compute_graph_edges_nbx(comm, dest);
std::ranges::sort(src);
// Create neighbourhood communicator for sending data to post
// offices
MPI_Comm comm0;
int err = MPI_Dist_graph_create_adjacent(
comm, src.size(), src.data(), MPI_UNWEIGHTED, dest.size(), dest.data(),
MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm0);
dolfinx::MPI::check_error(comm, err);
// Send number of items to post offices (destinations)
std::vector<int> num_items_recv(src.size());
num_items_send.reserve(1);
num_items_recv.reserve(1);
MPI_Neighbor_alltoall(num_items_send.data(), 1, MPI_INT,
num_items_recv.data(), 1, MPI_INT, comm0);
dolfinx::MPI::check_error(comm, err);
// Compute receive displacements
std::vector<int> recv_disp(num_items_recv.size() + 1, 0);
std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
std::next(recv_disp.begin()));
// Prepare send buffer
std::vector<std::int64_t> send_buffer;
std::vector<T> send_values_buffer;
send_buffer.reserve(entities.size());
send_values_buffer.reserve(data.size());
for (std::size_t e = 0; e < entities.extent(0); ++e)
{
auto idx = perm[e];
auto it = std::next(entities.data_handle(), idx * entities.extent(1));
send_buffer.insert(send_buffer.end(), it, it + entities.extent(1));
send_values_buffer.push_back(data[idx]);
}
std::vector<std::int64_t> recv_buffer(recv_disp.back()
* entities.extent(1));
err = MPI_Neighbor_alltoallv(send_buffer.data(), num_items_send.data(),
send_disp.data(), compound_type,
recv_buffer.data(), num_items_recv.data(),
recv_disp.data(), compound_type, comm0);
dolfinx::MPI::check_error(comm, err);
std::vector<T> recv_values_buffer(recv_disp.back());
err = MPI_Neighbor_alltoallv(
send_values_buffer.data(), num_items_send.data(), send_disp.data(),
dolfinx::MPI::mpi_t<T>, recv_values_buffer.data(),
num_items_recv.data(), recv_disp.data(), dolfinx::MPI::mpi_t<T>, comm0);
dolfinx::MPI::check_error(comm, err);
err = MPI_Comm_free(&comm0);
dolfinx::MPI::check_error(comm, err);
std::array shape{recv_buffer.size() / (entities.extent(1)),
(entities.extent(1))};
return std::tuple<std::vector<std::int64_t>, std::vector<T>,
std::array<std::size_t, 2>>(
std::move(recv_buffer), std::move(recv_values_buffer), shape);
};
const auto [entitiesp_b, entitiesp_v, shapep] = send_entities_to_postmaster(
comm, compound_type, num_nodes_g, entities_v, data);
md::mdspan<const std::int64_t, md::dextents<std::size_t, 2>> entitiesp(
entitiesp_b.data(), shapep);
// -- C. Send mesh global indices to postmaster
auto indices_to_postoffice = [](MPI_Comm comm, std::int64_t num_nodes,
std::span<const std::int64_t> indices)
{
int size = dolfinx::MPI::size(comm);
std::vector<std::pair<int, std::int64_t>> dest_to_index;
std::ranges::transform(
indices, std::back_inserter(dest_to_index),
[size, num_nodes](auto n)
{
return std::pair(dolfinx::MPI::index_owner(size, n, num_nodes), n);
});
std::ranges::sort(dest_to_index);
// Build list of neighbour dest ranks and count number of indices to
// send to each post office
std::vector<int> dest;
std::vector<std::int32_t> num_items_send;
{
auto it = dest_to_index.begin();
while (it != dest_to_index.end())
{
dest.push_back(it->first);
auto it1
= std::find_if(it, dest_to_index.end(), [r = dest.back()](auto idx)
{ return idx.first != r; });
num_items_send.push_back(std::distance(it, it1));
it = it1;
}
}
// Compute send displacements
std::vector<int> send_disp(num_items_send.size() + 1, 0);
std::partial_sum(num_items_send.begin(), num_items_send.end(),
std::next(send_disp.begin()));
// Determine src ranks. Sort ranks so that ownership determination is
// deterministic for a given number of ranks.
std::vector<int> src = dolfinx::MPI::compute_graph_edges_nbx(comm, dest);
std::ranges::sort(src);
// Create neighbourhood communicator for sending data to post offices
MPI_Comm comm0;
int err = MPI_Dist_graph_create_adjacent(
comm, src.size(), src.data(), MPI_UNWEIGHTED, dest.size(), dest.data(),
MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm0);
dolfinx::MPI::check_error(comm, err);
// Send number of items to post offices (destination) that I will be
// sending
std::vector<int> num_items_recv(src.size());
num_items_send.reserve(1);
num_items_recv.reserve(1);
MPI_Neighbor_alltoall(num_items_send.data(), 1, MPI_INT,
num_items_recv.data(), 1, MPI_INT, comm0);
dolfinx::MPI::check_error(comm, err);
// Compute receive displacements
std::vector<int> recv_disp(num_items_recv.size() + 1, 0);
std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
std::next(recv_disp.begin()));
// Prepare send buffer
std::vector<std::int64_t> send_buffer;
send_buffer.reserve(indices.size());
std::ranges::transform(dest_to_index, std::back_inserter(send_buffer),
[](auto x) { return x.second; });
std::vector<std::int64_t> recv_buffer(recv_disp.back());
err = MPI_Neighbor_alltoallv(send_buffer.data(), num_items_send.data(),
send_disp.data(), MPI_INT64_T,
recv_buffer.data(), num_items_recv.data(),
recv_disp.data(), MPI_INT64_T, comm0);
dolfinx::MPI::check_error(comm, err);
err = MPI_Comm_free(&comm0);
dolfinx::MPI::check_error(comm, err);
return std::tuple(std::move(recv_buffer), std::move(recv_disp),
std::move(src), std::move(dest));
};
const auto [nodes_g_p, recv_disp, src, dest]
= indices_to_postoffice(comm, num_nodes_g, nodes_g);
// D. Send entities to possible owners, based on first entity index
auto candidate_ranks
= [](MPI_Comm comm, MPI_Datatype compound_type,
std::span<const std::int64_t> indices_recv,
std::span<const int> indices_recv_disp, std::span<const int> src,
std::span<const int> dest, auto entities, std::span<const T> data)
{
// Build map from received global node indices to neighbourhood
// ranks that have the node
std::multimap<std::int64_t, int> node_to_rank;
for (std::size_t i = 0; i < indices_recv_disp.size() - 1; ++i)
for (int j = indices_recv_disp[i]; j < indices_recv_disp[i + 1]; ++j)
node_to_rank.insert({indices_recv[j], i});
std::vector<std::vector<std::int64_t>> send_data(dest.size());
std::vector<std::vector<T>> send_values(dest.size());
for (std::size_t e = 0; e < entities.extent(0); ++e)
{
std::span e_recv(entities.data_handle() + e * entities.extent(1),
entities.extent(1));
auto [it0, it1] = node_to_rank.equal_range(entities(e, 0));
for (auto it = it0; it != it1; ++it)
{
int p = it->second;
send_data[p].insert(send_data[p].end(), e_recv.begin(), e_recv.end());
send_values[p].push_back(data[e]);
}
}
MPI_Comm comm0;
int err = MPI_Dist_graph_create_adjacent(
comm, src.size(), src.data(), MPI_UNWEIGHTED, dest.size(), dest.data(),
MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm0);
dolfinx::MPI::check_error(comm, err);
std::vector<int> num_items_send;
num_items_send.reserve(send_data.size());
for (auto& x : send_data)
num_items_send.push_back(x.size() / entities.extent(1));
std::vector<int> num_items_recv(src.size());
num_items_send.reserve(1);
num_items_recv.reserve(1);
err = MPI_Neighbor_alltoall(num_items_send.data(), 1, MPI_INT,
num_items_recv.data(), 1, MPI_INT, comm0);
dolfinx::MPI::check_error(comm, err);
// Compute send displacements
std::vector<std::int32_t> send_disp(num_items_send.size() + 1, 0);
std::partial_sum(num_items_send.begin(), num_items_send.end(),
std::next(send_disp.begin()));
// Compute receive displacements
std::vector<std::int32_t> recv_disp(num_items_recv.size() + 1, 0);
std::partial_sum(num_items_recv.begin(), num_items_recv.end(),
std::next(recv_disp.begin()));
// Prepare send buffers
std::vector<std::int64_t> send_buffer;
std::vector<T> send_values_buffer;
for (auto& x : send_data)
send_buffer.insert(send_buffer.end(), x.begin(), x.end());
for (auto& v : send_values)
send_values_buffer.insert(send_values_buffer.end(), v.begin(), v.end());
std::vector<std::int64_t> recv_buffer(entities.extent(1)
* recv_disp.back());
err = MPI_Neighbor_alltoallv(send_buffer.data(), num_items_send.data(),
send_disp.data(), compound_type,
recv_buffer.data(), num_items_recv.data(),
recv_disp.data(), compound_type, comm0);
dolfinx::MPI::check_error(comm, err);
std::vector<T> recv_values_buffer(recv_disp.back());
err = MPI_Neighbor_alltoallv(
send_values_buffer.data(), num_items_send.data(), send_disp.data(),
dolfinx::MPI::mpi_t<T>, recv_values_buffer.data(),
num_items_recv.data(), recv_disp.data(), dolfinx::MPI::mpi_t<T>, comm0);
dolfinx::MPI::check_error(comm, err);
err = MPI_Comm_free(&comm0);
dolfinx::MPI::check_error(comm, err);
std::array shape{recv_buffer.size() / entities.extent(1),
entities.extent(1)};
return std::tuple<std::vector<std::int64_t>, std::vector<T>,
std::array<std::size_t, 2>>(
std::move(recv_buffer), std::move(recv_values_buffer), shape);
};
// NOTE: src and dest are transposed here because we're reversing the
// direction of communication
const auto [entities_data_b, entities_values, shape_eb]
= candidate_ranks(comm, compound_type, nodes_g_p, recv_disp, dest, src,
entitiesp, std::span(entitiesp_v));
md::mdspan<const std::int64_t, md::dextents<std::size_t, 2>> entities_data(
entities_data_b.data(), shape_eb);
// -- E. From the received (key, value) data, determine which keys
// (entities) are on this process.
//
// TODO: We have already received possibly tagged entities from other
// ranks, so we could use the received data to avoid creating
// the std::map for *all* entities and just for candidate
// entities.
auto select_entities
= [](const mesh::Topology& topology, auto xdofmap,
std::span<const std::int64_t> nodes_g,
std::span<const int> cell_vertex_dofs, auto entities_data,
std::span<const T> entities_values)
{
spdlog::info("XDMF build map");
auto c_to_v = topology.connectivity(topology.dim(), 0);
if (!c_to_v)
throw std::runtime_error("Missing cell-vertex connectivity.");
std::map<std::int64_t, std::int32_t> input_idx_to_vertex;
for (int c = 0; c < c_to_v->num_nodes(); ++c)
{
auto vertices = c_to_v->links(c);
std::span xdofs(xdofmap.data_handle() + c * xdofmap.extent(1),
xdofmap.extent(1));
for (std::size_t v = 0; v < vertices.size(); ++v)
input_idx_to_vertex[nodes_g[xdofs[cell_vertex_dofs[v]]]] = vertices[v];
}
std::vector<std::int32_t> entities;
std::vector<T> data;
std::vector<std::int32_t> entity(entities_data.extent(1));
for (std::size_t e = 0; e < entities_data.extent(0); ++e)
{
bool entity_found = true;
for (std::size_t i = 0; i < entities_data.extent(1); ++i)
{
if (auto it = input_idx_to_vertex.find(entities_data(e, i));
it == input_idx_to_vertex.end())
{
// As soon as this received index is not in locally owned
// input global indices skip the entire entity
entity_found = false;
break;
}
else
entity[i] = it->second;
}
if (entity_found)
{
entities.insert(entities.end(), entity.begin(), entity.end());
data.push_back(entities_values[e]);
}
}
return std::pair(std::move(entities), std::move(data));
};
MPI_Type_free(&compound_type);
return select_entities(topology, xdofmap, nodes_g, cell_vertex_dofs,
entities_data, std::span(entities_values));
}
//-----------------------------------------------------------------------------}
} // namespace io
} // namespace dolfinx
dolfinx::fem::ElementDofLayout
Definition ElementDofLayout.h:31
dolfinx::fem::ElementDofLayout::entity_closure_dofs
const std::vector< int > & entity_closure_dofs(int dim, int entity_index) const
Definition ElementDofLayout.cpp:80
dolfinx::fem::ElementDofLayout::entity_dofs
const std::vector< int > & entity_dofs(int dim, int entity_index) const
Definition ElementDofLayout.cpp:73
dolfinx::mesh::Topology
Topology stores the topology of a mesh, consisting of mesh entities and connectivity (incidence relat...
Definition Topology.h:41
dolfinx::mesh::Topology::connectivity
std::shared_ptr< const graph::AdjacencyList< std::int32_t > > connectivity(std::array< int, 2 > d0, std::array< int, 2 > d1) const
Get the connectivity from entities of topological dimension d0 to dimension d1.
Definition Topology.cpp:840
dolfinx::mesh::Topology::dim
int dim() const noexcept
Topological dimension of the mesh.
Definition Topology.cpp:786
dolfinx::mesh::Topology::comm
MPI_Comm comm() const
Mesh MPI communicator.
Definition Topology.cpp:1018
dolfinx::mesh::Topology::cell_type
CellType cell_type() const
Cell type.
Definition Topology.cpp:796
dolfinx::MPI::mpi_t
MPI_Datatype mpi_t
Retrieves the MPI data type associated to the provided type.
Definition MPI.h:280
dolfinx::MPI::compute_graph_edges_nbx
std::vector< int > compute_graph_edges_nbx(MPI_Comm comm, std::span< const int > edges, int tag=static_cast< int >(tag::consensus_nbx))
Determine incoming graph edges using the NBX consensus algorithm.
Definition MPI.cpp:162
dolfinx::MPI::index_owner
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:113
dolfinx::MPI::check_error
void check_error(MPI_Comm comm, int code)
Check MPI error code. If the error code is not equal to MPI_SUCCESS, then std::abort is called.
Definition MPI.cpp:80
dolfinx::MPI::size
int size(MPI_Comm comm)
Definition MPI.cpp:72
dolfinx::io
Support for file IO.
Definition cells.h:119
dolfinx::io::distribute_entity_data
std::pair< std::vector< std::int32_t >, std::vector< T > > distribute_entity_data(const mesh::Topology &topology, std::span< const std::int64_t > nodes_g, std::int64_t num_nodes_g, const fem::ElementDofLayout &cmap_dof_layout, md::mdspan< const std::int32_t, md::dextents< std::size_t, 2 > > xdofmap, int entity_dim, md::mdspan< const std::int64_t, md::dextents< std::size_t, 2 > > entities, std::span< const T > data)
Get owned entities and associated data from input entities defined by global 'node' indices.
Definition utils.h:60
dolfinx::mesh::cell_entity_type
CellType cell_entity_type(CellType type, int d, int index)
Return type of cell for entity of dimension d at given entity index.
Definition cell_types.cpp:64
dolfinx::mesh::CellType
CellType
Cell type identifier.
Definition cell_types.h:22
dolfinx::mesh::cell_num_entities
int cell_num_entities(CellType type, int dim)
Definition cell_types.cpp:139
dolfinx
Top-level namespace.
Definition defines.h:12