diff --git a/cpp/dolfinx/fem/dofmapbuilder.cpp b/cpp/dolfinx/fem/dofmapbuilder.cpp index cd330b063ed..2dc73c2ada6 100644 --- a/cpp/dolfinx/fem/dofmapbuilder.cpp +++ b/cpp/dolfinx/fem/dofmapbuilder.cpp @@ -7,7 +7,6 @@ #include "dofmapbuilder.h" #include "ElementDofLayout.h" #include -#include #include #include #include @@ -27,14 +26,21 @@ using namespace dolfinx; namespace { -template -using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan< - T, MDSPAN_IMPL_STANDARD_NAMESPACE::dextents>; - +/// @brief dofmap as a flattened 2D array +/// The number of cells is array.size() / width +/// and therefore array.size() must be a multiple of width +struct dofmap_t +{ + std::int32_t width; + std::vector array; +}; //----------------------------------------------------------------------------- -/// Build a graph for owned dofs and apply graph reordering function -/// @param[in] dofmap The local dofmap (cell -> dofs) +/// Build a graph for owned dofs and apply graph reordering function with +/// multiple dofmaps. The dofmaps are 2D arrays, of fixed width, stored in +/// `dofmap_t` format. The dofmaps all refer to dof indices in the same range +/// [0:owned_size). +/// @param[in] dofmaps The local dofmaps (cell -> dofs) /// @param[in] owned_size Number of dofs owned by this process /// @param[in] original_to_contiguous Map from dof indices in @p dofmap /// to new indices that are ordered such that owned indices are [0, @@ -43,7 +49,7 @@ using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan< /// @return Map from original_to_contiguous[i] to new index after /// reordering std::vector -reorder_owned(mdspan2_t dofmap, std::int32_t owned_size, +reorder_owned(const std::vector& dofmaps, std::int32_t owned_size, const std::vector& original_to_contiguous, const std::function( const graph::AdjacencyList&)>& reorder_fn) @@ -52,22 +58,22 @@ reorder_owned(mdspan2_t dofmap, std::int32_t owned_size, // Compute maximum number of graph out edges edges per dof std::vector num_edges(owned_size); - for (std::size_t cell = 0; cell < dofmap.extent(0); ++cell) + for (const auto& dofmap : dofmaps) { - std::span nodes( - dofmap.data_handle() + cell * dofmap.extent(1), dofmap.extent(1)); - for (auto n0 : nodes) + std::size_t num_cells = dofmap.array.size() / dofmap.width; + std::vector node_temp; + for (std::size_t cell = 0; cell < num_cells; ++cell) { - const std::int32_t node_0 = original_to_contiguous[n0]; - - // Skip unowned node - if (node_0 >= owned_size) - continue; - for (auto n1 : nodes) + node_temp.clear(); + for (std::int32_t i = 0; i < dofmap.width; ++i) { - if (n0 != n1 and original_to_contiguous[n1] < owned_size) - ++num_edges[node_0]; + std::int32_t node + = original_to_contiguous[dofmap.array[cell * dofmap.width + i]]; + if (node < owned_size) + node_temp.push_back(node); } + for (std::int32_t node : node_temp) + num_edges[node] += node_temp.size() - 1; } } @@ -76,21 +82,30 @@ reorder_owned(mdspan2_t dofmap, std::int32_t owned_size, std::partial_sum(num_edges.begin(), num_edges.end(), std::next(offsets.begin(), 1)); std::vector edges(offsets.back()); - for (std::size_t cell = 0; cell < dofmap.extent(0); ++cell) + for (const auto& dofmap : dofmaps) { - std::span nodes( - dofmap.data_handle() + cell * dofmap.extent(1), dofmap.extent(1)); - for (auto n0 : nodes) + std::size_t num_cells = dofmap.array.size() / dofmap.width; + std::vector node_temp; + + for (std::size_t cell = 0; cell < num_cells; ++cell) { - const std::int32_t node_0 = original_to_contiguous[n0]; - if (node_0 >= owned_size) - continue; - for (auto n1 : nodes) + node_temp.clear(); + for (std::int32_t i = 0; i < dofmap.width; ++i) { - if (const std::int32_t node_1 = original_to_contiguous[n1]; - n0 != n1 and node_1 < owned_size) + std::int32_t node + = original_to_contiguous[dofmap.array[cell * dofmap.width + i]]; + if (node < owned_size) + node_temp.push_back(node); + } + + for (std::size_t i = 0; i < node_temp.size(); ++i) + { + std::int32_t node_0 = node_temp[i]; + for (std::size_t j = i + 1; j < node_temp.size(); ++j) { + std::int32_t node_1 = node_temp[j]; edges[offsets[node_0]++] = node_1; + edges[offsets[node_1]++] = node_0; } } } @@ -132,7 +147,7 @@ reorder_owned(mdspan2_t dofmap, std::int32_t owned_size, /// * local-to-global map for each local dof /// * local-to-entity map for each local dof /// Entities are represented as {dimension, mesh entity index}. -std::tuple, std::vector, +std::tuple, std::vector>> build_basic_dofmap(const mesh::Topology& topology, const fem::ElementDofLayout& element_dof_layout) @@ -175,20 +190,19 @@ build_basic_dofmap(const mesh::Topology& topology, // Allocate dofmap memory const std::int32_t num_cells = topology.index_map(D)->size_local() + topology.index_map(D)->num_ghosts(); - std::vector dofs(element_dof_layout.num_dofs() * num_cells); + dofmap_t dofs; + dofs.width = element_dof_layout.num_dofs(); + dofs.array.resize(dofs.width * num_cells); // Loop over cells and build dofmap const std::vector>> entity_dofs = element_dof_layout.entity_dofs_all(); assert(entity_dofs.size() == D + 1); - const int dofmap_width = element_dof_layout.num_dofs(); - // for (int c = group_offsets[i]; c < group_offsets[i + 1]; ++c) for (std::int32_t c = 0; c < num_cells; ++c) { // Wrap dofs for cell c - std::span dofs_c(dofs.data() + c * dofmap_width, - dofmap_width); + std::span dofs_c(dofs.array.data() + c * dofs.width, dofs.width); // Iterate over each topological dimension for this element std::int32_t offset_local = 0; @@ -232,7 +246,7 @@ build_basic_dofmap(const mesh::Topology& topology, std::int32_t offset_local = 0; std::int64_t offset_global = 0; - // Dof -> (dim, entity index) marker + // Dof -> (index_map, entity index) marker std::vector> dof_entity(local_size); // Storage for local-to-global map @@ -272,29 +286,31 @@ build_basic_dofmap(const mesh::Topology& topology, /// the end, i.e. in the positions [M, ..., N), where N is the total /// number of dofs on this process. /// -/// @param [in] dofmap The basic dofmap data -/// @param [in] dof_entity Map from dof index to (dim, entity_index), -/// where entity_index is the process-wise mesh entity index -/// @param [in] topology The mesh topology +/// @param [in] dofmaps The basic dofmap data in multiple dofmaps sharing the +/// same range +/// @param [in] dof_entity Map from dof index to (index_map, entity_index), +/// where entity_index is the local mesh entity index in the given index_map +/// @param [in] index_maps The set of IndexMaps, one for each topological +/// entity type used in the dofmap. The location in this array is referred to by +/// the first item in each entry of @p dof_entity /// @param [in] reorder_fn Graph reordering function that is applied for /// dof re-ordering /// @return The pair (old-to-new local index map, M), where M is the /// number of dofs owned by this process std::pair, std::int32_t> compute_reordering_map( - mdspan2_t dofmap, + const std::vector& dofmaps, const std::vector>& dof_entity, - const mesh::Topology& topology, + const std::vector>& index_maps, const std::function( const graph::AdjacencyList&)>& reorder_fn) { common::Timer t0("Compute dof reordering map"); // Get mesh entity ownership offset for each topological dimension - const int D = topology.dim(); - std::vector offset(D + 1, -1); + std::vector offset(index_maps.size(), -1); for (std::size_t d = 0; d < offset.size(); ++d) { - auto map = topology.index_map(d); + auto map = index_maps[d]; if (map) offset[d] = map->size_local(); } @@ -311,19 +327,16 @@ std::pair, std::int32_t> compute_reordering_map( // owned dofs. Set to -1 for unowned dofs. std::vector original_to_contiguous(dof_entity.size(), -1); std::int32_t counter_owned(0), counter_unowned(owned_size); - for (std::size_t cell = 0; cell < dofmap.extent(0); ++cell) + for (auto dofmap : dofmaps) { - auto dofs = MDSPAN_IMPL_STANDARD_NAMESPACE::MDSPAN_IMPL_PROPOSED_NAMESPACE:: - submdspan(dofmap, cell, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent); - for (std::size_t i = 0; i < dofs.size(); ++i) + for (std::int32_t dof : dofmap.array) { - if (original_to_contiguous[dofs[i]] == -1) + if (original_to_contiguous[dof] == -1) { - const std::pair& e = dof_entity[dofs[i]]; - if (e.second < offset[e.first]) - original_to_contiguous[dofs[i]] = counter_owned++; + if (auto e = dof_entity[dof]; e.second < offset[e.first]) + original_to_contiguous[dof] = counter_owned++; else - original_to_contiguous[dofs[i]] = counter_unowned++; + original_to_contiguous[dof] = counter_unowned++; } } } @@ -348,8 +361,8 @@ std::pair, std::int32_t> compute_reordering_map( // Re-order using graph ordering // Apply graph reordering to owned dofs - const std::vector node_remap - = reorder_owned(dofmap, owned_size, original_to_contiguous, reorder_fn); + const std::vector node_remap = reorder_owned( + dofmaps, owned_size, original_to_contiguous, reorder_fn); std::transform(original_to_contiguous.begin(), original_to_contiguous.end(), original_to_contiguous.begin(), [&node_remap, owned_size](auto index) @@ -361,35 +374,34 @@ std::pair, std::int32_t> compute_reordering_map( //----------------------------------------------------------------------------- /// Get global indices for unowned dofs -/// @param [in] topology The mesh topology +/// @param [in] index_maps Set of index maps corresponding to dofs in @p +/// dof_entity, below. /// @param [in] num_owned The number of nodes owned by this process /// @param [in] process_offset The node offset for this process, i.e. /// the global index of owned node i is i + process_offset /// @param [in] global_indices_old The old global index of the old local /// node i /// @param [in] old_to_new The old local index to new local index map -/// @param [in] dof_entity The ith entry gives (topological dim, local +/// @param [in] dof_entity The ith entry gives (index_map, local /// index) of the mesh entity to which node i (old local index) is -/// associated +/// associated. /// @returns The (0) global indices for unowned dofs, (1) owner rank of /// each unowned dof std::pair, std::vector> get_global_indices( - const mesh::Topology& topology, const std::int32_t num_owned, - const std::int64_t process_offset, + const std::vector>& index_maps, + std::int32_t num_owned, std::int64_t process_offset, const std::vector& global_indices_old, const std::vector& old_to_new, const std::vector>& dof_entity) { assert(dof_entity.size() == global_indices_old.size()); - const int D = topology.dim(); - // Build list of flags for owned mesh entities that are shared, i.e. // are a ghost on a neighbor - std::vector> shared_entity(D + 1); + std::vector> shared_entity(index_maps.size()); for (std::size_t d = 0; d < shared_entity.size(); ++d) { - auto map = topology.index_map(d); + auto map = index_maps[d]; if (map) { shared_entity[d] = std::vector(map->size_local(), false); @@ -402,9 +414,9 @@ std::pair, std::vector> get_global_indices( // Build list of (global old, global new) index pairs for dofs that // are ghosted on other processes - std::vector> global(D + 1); // Loop over all dofs + std::vector> global(index_maps.size()); for (std::size_t i = 0; i < dof_entity.size(); ++i) { // Topological dimension of mesh entity that dof is associated with @@ -420,14 +432,14 @@ std::pair, std::vector> get_global_indices( } std::vector requests_dim; - std::vector requests(D + 1); - std::vector comm(D + 1, MPI_COMM_NULL); - std::vector> all_dofs_received(D + 1); - std::vector> disp_recv(D + 1); - for (int d = 0; d <= D; ++d) + std::vector requests(index_maps.size()); + std::vector comm(index_maps.size(), MPI_COMM_NULL); + std::vector> all_dofs_received(index_maps.size()); + std::vector> disp_recv(index_maps.size()); + for (std::size_t d = 0; d < index_maps.size(); ++d) { // FIXME: This should check which dimension are needed by the dofmap - auto map = topology.index_map(d); + auto map = index_maps[d]; if (map) { std::span src = map->src(); @@ -463,7 +475,8 @@ std::pair, std::vector> get_global_indices( // Build [local_new - num_owned] -> global old array broken down by // dimension - std::vector> local_new_to_global_old(D + 1); + std::vector> local_new_to_global_old( + index_maps.size()); for (std::size_t i = 0; i < global_indices_old.size(); ++i) { const int d = dof_entity[i].first; @@ -483,7 +496,7 @@ std::pair, std::vector> get_global_indices( MPI_Waitany(requests_dim.size(), requests.data(), &idx, MPI_STATUS_IGNORE); d = requests_dim[idx]; - std::span src = topology.index_map(d)->src(); + std::span src = index_maps[d]->src(); // Build (global old, global new) map for dofs of dimension d std::vector>> @@ -516,10 +529,10 @@ std::pair, std::vector> get_global_indices( } } - for (std::size_t i = 0; i < comm.size(); ++i) + for (MPI_Comm c : comm) { - if (comm[i] != MPI_COMM_NULL) - MPI_Comm_free(&comm[i]); + if (c != MPI_COMM_NULL) + MPI_Comm_free(&c); } return {std::move(local_to_global_new), std::move(local_to_global_new_owner)}; @@ -546,29 +559,35 @@ fem::build_dofmap_data( const auto [node_graph0, local_to_global0, dof_entity0] = build_basic_dofmap(topology, element_dof_layout); - // Compute global dofmap offset - std::int64_t offset = 0; + std::vector> index_maps(D + 1); for (int d = 0; d <= D; ++d) { if (element_dof_layout.num_entity_dofs(d) > 0) { assert(topology.index_map(d)); - offset += topology.index_map(d)->local_range()[0] - * element_dof_layout.num_entity_dofs(d); + index_maps[d] = topology.index_map(d); } } // Build re-ordering map for data locality and get number of owned // nodes - mdspan2_t _node_graph0( - node_graph0.data(), node_graph0.size() / element_dof_layout.num_dofs(), - element_dof_layout.num_dofs()); - const auto [old_to_new, num_owned] - = compute_reordering_map(_node_graph0, dof_entity0, topology, reorder_fn); + const auto [old_to_new, num_owned] = compute_reordering_map( + {node_graph0}, dof_entity0, index_maps, reorder_fn); + + // Compute global dofmap offset + std::int64_t offset = 0; + for (int d = 0; d <= D; ++d) + { + if (element_dof_layout.num_entity_dofs(d) > 0) + { + offset += index_maps[d]->local_range()[0] + * element_dof_layout.num_entity_dofs(d); + } + } // Get global indices for unowned dofs const auto [local_to_global_unowned, local_to_global_owner] - = get_global_indices(topology, num_owned, offset, local_to_global0, + = get_global_indices(index_maps, num_owned, offset, local_to_global0, old_to_new, dof_entity0); assert(local_to_global_unowned.size() == local_to_global_owner.size()); @@ -577,21 +596,9 @@ fem::build_dofmap_data( local_to_global_owner); // Build re-ordered dofmap - std::vector dofmap(node_graph0.size()); - for (std::size_t cell = 0; cell < _node_graph0.extent(0); ++cell) - { - // Get dof order on this cell - auto old_nodes = MDSPAN_IMPL_STANDARD_NAMESPACE:: - MDSPAN_IMPL_PROPOSED_NAMESPACE::submdspan( - _node_graph0, cell, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent); - const std::int32_t local_dim0 = old_nodes.size(); - std::span dofs(dofmap.data() + cell * local_dim0, local_dim0); - for (std::int32_t j = 0; j < local_dim0; ++j) - { - std::int32_t old_node = old_nodes[j]; - dofs[j] = old_to_new[old_node]; - } - } + std::vector dofmap(node_graph0.array.size()); + for (std::size_t i = 0; i < dofmap.size(); ++i) + dofmap[i] = old_to_new[node_graph0.array[i]]; return {std::move(index_map), element_dof_layout.block_size(), std::move(dofmap)}; diff --git a/python/dolfinx/wrappers/fem.cpp b/python/dolfinx/wrappers/fem.cpp index e9b22e8849a..b9ba3807bcf 100644 --- a/python/dolfinx/wrappers/fem.cpp +++ b/python/dolfinx/wrappers/fem.cpp @@ -988,7 +988,7 @@ void fem(nb::module_& m) return std::tuple(std::move(map), bs, std::move(dofmap)); }, nb::arg("comm"), nb::arg("topology"), nb::arg("layout"), - "Build and dofmap on a mesh."); + "Build a dofmap on a mesh."); m.def( "transpose_dofmap", [](nb::ndarray, nb::c_contig> dofmap,