Merge branch 'main' into main

Author: uvilla

cpp/demo/custom_kernel/main.cpp

@@ -168,7 +168,7 @@ double assemble_vector1(const mesh::Geometry<T>& g, const fem::DofMap& dofmap,
   common::Timer timer("Assembler1 lambda (vector)");
   fem::impl::assemble_cells<T, 1>([](auto, auto, auto, auto) {},
                                   b.mutable_array(), g.dofmap(), g.x(), cells,
-                                  dofmap.map(), 1, kernel, {}, {}, 0, {});
+                                  {dofmap.map(), 1, cells}, kernel, {}, {}, 0, {});
   b.scatter_rev(std::plus<T>());
   return la::squared_norm(b);
 }

cpp/dolfinx/common/IndexMap.cpp

@@ -164,6 +164,8 @@ communicate_ghosts_to_owners(MPI_Comm comm, std::span<const int> src,
 /// @param[in] imap An index map.
 /// @param[in] indices List of entity indices (indices local to the
 /// process).
+/// @param[in] order Control the order in which ghost indices appear in
+/// the new map.
 /// @param[in] allow_owner_change Allows indices that are not included
 /// by their owning process but included on sharing processes to be
 /// included in the submap. These indices will be owned by one of the
@@ -176,7 +178,7 @@ std::tuple<std::vector<std::int32_t>, std::vector<std::int32_t>,
            std::vector<int>, std::vector<int>, std::vector<int>>
 compute_submap_indices(const IndexMap& imap,
                        std::span<const std::int32_t> indices,
-                       bool allow_owner_change)
+                       IndexMapOrder order, bool allow_owner_change)
 {
   std::span<const int> src = imap.src();
   std::span<const int> dest = imap.dest();
@@ -203,7 +205,7 @@ compute_submap_indices(const IndexMap& imap,
   // `recv_indices` will necessarily be in `indices` on this process, and thus
   // other processes must own them in the submap.
 
-  std::vector<int> recv_owners(send_disp.back()), submap_dest;
+  std::vector<int> recv_owners(send_disp.back());
   const int rank = dolfinx::MPI::rank(imap.comm());
   {
     // Flag to track if the owner of any indices have changed in the submap
@@ -231,11 +233,7 @@ compute_submap_indices(const IndexMap& imap,
           // 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});
-            // Add the sending process as a destination rank in the submap
-            submap_dest.push_back(dest[i]);
-          }
           else
           {
             owners_changed = true;
@@ -250,10 +248,6 @@ compute_submap_indices(const IndexMap& imap,
 
     std::sort(global_idx_to_possible_owner.begin(),
               global_idx_to_possible_owner.end());
-    std::sort(submap_dest.begin(), submap_dest.end());
-    submap_dest.erase(std::unique(submap_dest.begin(), submap_dest.end()),
-                      submap_dest.end());
-    submap_dest.shrink_to_fit();
 
     // Choose the submap owner for each index in `recv_indices`
     std::vector<int> send_owners;
@@ -350,6 +344,37 @@ compute_submap_indices(const IndexMap& imap,
                    submap_src.end());
   submap_src.shrink_to_fit();
 
+  // If required, preserve the order of the ghost indices
+  if (order == IndexMapOrder::preserve)
+  {
+    // Build (old postion, new position) list for ghosts and sort
+    std::vector<std::pair<std::int32_t, std::int32_t>> pos;
+    pos.reserve(submap_ghost.size());
+    for (std::int32_t idx : submap_ghost)
+      pos.emplace_back(idx, pos.size());
+    std::sort(pos.begin(), pos.end());
+
+    // Order ghosts in the sub-map by their position in the parent map
+    std::vector<int> submap_ghost_owners1;
+    submap_ghost_owners1.reserve(submap_ghost_owners.size());
+    std::vector<std::int32_t> submap_ghost1;
+    submap_ghost1.reserve(submap_ghost.size());
+    for (auto [_, idx] : pos)
+    {
+      submap_ghost_owners1.push_back(submap_ghost_owners[idx]);
+      submap_ghost1.push_back(submap_ghost[idx]);
+    }
+
+    submap_ghost_owners = std::move(submap_ghost_owners1);
+    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::sort(submap_dest.begin(), submap_dest.end());
+
   return {std::move(submap_owned), std::move(submap_ghost),
           std::move(submap_ghost_owners), std::move(submap_src),
           std::move(submap_dest)};
@@ -725,14 +750,14 @@ common::stack_index_maps(
 std::pair<IndexMap, std::vector<std::int32_t>>
 common::create_sub_index_map(const IndexMap& imap,
                              std::span<const std::int32_t> indices,
-                             bool allow_owner_change)
+                             IndexMapOrder order, bool allow_owner_change)
 {
   // Compute the owned, ghost, and ghost owners of submap indices.
   // NOTE: All indices are local and numbered w.r.t. the original (imap)
   // index map
-  const auto [submap_owned, submap_ghost, submap_ghost_owners, submap_src,
-              submap_dest]
-      = compute_submap_indices(imap, indices, allow_owner_change);
+  auto [submap_owned, submap_ghost, submap_ghost_owners, submap_src,
+        submap_dest]
+      = compute_submap_indices(imap, indices, order, allow_owner_change);
 
   // Compute submap offset for this rank
   std::int64_t submap_local_size = submap_owned.size();

cpp/dolfinx/common/IndexMap.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2015-2022 Chris Richardson, Garth N. Wells and Igor Baratta
+// Copyright (C) 2015-2024 Chris Richardson, Garth N. Wells and Igor Baratta
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -19,6 +19,14 @@ namespace dolfinx::common
 // Forward declaration
 class IndexMap;
 
+/// Enum to control preservation of ghost index ordering in
+/// sub-IndexMaps.
+enum class IndexMapOrder : bool
+{
+  preserve = true, ///< Preserve the ordering of ghost indices
+  any = false      ///< Allow arbitrary ordering of ghost indices in sub-maps
+};
+
 /// @brief Given a sorted vector of indices (local numbering, owned or
 /// ghost) and an index map, this function returns the indices owned by
 /// this process, including indices that might have been in the list of
@@ -60,6 +68,8 @@ stack_index_maps(
 /// @param[in] imap Parent map to create a new sub-map from.
 /// @param[in] indices Local indices in `imap` (owned and ghost) to
 /// include in the new index map.
+/// @param[in] order Control the order in which ghost indices appear in
+/// the new map.
 /// @param[in] allow_owner_change If `true`, indices that are not
 /// included in `indices` by their owning process can be included in
 /// `indices` by processes that ghost the indices to be included in the
@@ -70,10 +80,9 @@ stack_index_maps(
 /// @return The (i) new index map and (ii) a map from local indices in
 /// the submap to local indices in the original (this) map.
 /// @pre `indices` must be sorted and must not contain duplicates.
-std::pair<IndexMap, std::vector<std::int32_t>>
-create_sub_index_map(const IndexMap& imap,
-                     std::span<const std::int32_t> indices,
-                     bool allow_owner_change = false);
+std::pair<IndexMap, std::vector<std::int32_t>> create_sub_index_map(
+    const IndexMap& imap, std::span<const std::int32_t> indices,
+    IndexMapOrder order = IndexMapOrder::any, bool allow_owner_change = false);
 
 /// This class represents the distribution index arrays across
 /// processes. An index array is a contiguous collection of `N+1`
@@ -227,18 +236,19 @@ class IndexMap
 
   /// @brief Returns the imbalance of the current IndexMap.
   ///
-  /// The imbalance is a measure of load balancing across all processes, defined
-  /// as the maximum number of indices on any process divided by the average
-  /// number of indices per process. This function calculates the imbalance
-  /// separately for owned indices and ghost indices and returns them as a
-  /// std::array<double, 2>. If the total number of owned or ghost indices is
-  /// zero, the respective entry in the array is set to -1.
+  /// The imbalance is a measure of load balancing across all processes,
+  /// defined as the maximum number of indices on any process divided by
+  /// the average number of indices per process. This function
+  /// calculates the imbalance separately for owned indices and ghost
+  /// indices and returns them as a std::array<double, 2>. If the total
+  /// number of owned or ghost indices is zero, the respective entry in
+  /// the array is set to -1.
   ///
-  /// @note This is a collective operation and must be called by all processes
-  /// in the communicator associated with the IndexMap.
+  /// @note This is a collective operation and must be called by all
+  /// processes in the communicator associated with the IndexMap.
   ///
-  /// @return An array containing the imbalance in owned indices
-  /// (first element) and the imbalance in ghost indices (second element).
+  /// @return An array containing the imbalance in owned indices (first
+  /// element) and the imbalance in ghost indices (second element).
   std::array<double, 2> imbalance() const;
 
 private:

cpp/dolfinx/common/MPI.cpp

@@ -116,10 +116,10 @@ dolfinx::MPI::compute_graph_edges_pcx(MPI_Comm comm, std::span<const int> edges)
   dolfinx::MPI::check_error(comm, err);
 
   std::vector<MPI_Request> send_requests(edges.size());
-  std::byte send_buffer;
+  std::vector<std::byte> send_buffer(edges.size());
   for (std::size_t e = 0; e < edges.size(); ++e)
   {
-    int err = MPI_Isend(&send_buffer, 1, MPI_BYTE, edges[e],
+    int err = MPI_Isend(send_buffer.data() + e, 1, MPI_BYTE, edges[e],
                         static_cast<int>(tag::consensus_pcx), comm,
                         &send_requests[e]);
     dolfinx::MPI::check_error(comm, err);
@@ -168,10 +168,10 @@ dolfinx::MPI::compute_graph_edges_nbx(MPI_Comm comm, std::span<const int> edges)
 
   // Start non-blocking synchronised send
   std::vector<MPI_Request> send_requests(edges.size());
-  std::byte send_buffer;
+  std::vector<std::byte> send_buffer(edges.size());
   for (std::size_t e = 0; e < edges.size(); ++e)
   {
-    int err = MPI_Issend(&send_buffer, 1, MPI_BYTE, edges[e],
+    int err = MPI_Issend(send_buffer.data() + e, 1, MPI_BYTE, edges[e],
                          static_cast<int>(tag::consensus_pex), comm,
                          &send_requests[e]);
     dolfinx::MPI::check_error(comm, err);

cpp/dolfinx/fem/DofMap.cpp

@@ -60,9 +60,8 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   std::vector<std::int32_t> sub_imap_to_imap;
   if (bs_view == 1)
   {
-    auto [_index_map, _sub_imap_to_imap]
-        = dolfinx::common::create_sub_index_map(*dofmap_view.index_map,
-                                                dofs_view);
+    auto [_index_map, _sub_imap_to_imap] = common::create_sub_index_map(
+        *dofmap_view.index_map, dofs_view, common::IndexMapOrder::preserve);
     index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
     sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }
@@ -73,9 +72,8 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
     std::transform(dofs_view.begin(), dofs_view.end(),
                    std::back_inserter(indices),
                    [bs_view](auto idx) { return idx / bs_view; });
-    auto [_index_map, _sub_imap_to_imap]
-        = dolfinx::common::create_sub_index_map(*dofmap_view.index_map,
-                                                indices);
+    auto [_index_map, _sub_imap_to_imap] = common::create_sub_index_map(
+        *dofmap_view.index_map, indices, common::IndexMapOrder::preserve);
     index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
     sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }

cpp/dolfinx/fem/assemble_matrix_impl.h

@@ -40,7 +40,7 @@ using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
 /// @param dofmap0 Test function (row) degree-of-freedom data holding
 /// the (0) dofmap, (1) dofmap block size and (2) dofmap cell indices.
 /// @param P0 Function that applies transformation P_0 A in-place to
-/// transform trial degrees-of-freedom.
+/// transform test degrees-of-freedom.
 /// @param dofmap1 Trial function (column) degree-of-freedom data
 /// holding the (0) dofmap, (1) dofmap block size and (2) dofmap cell
 /// indices.
@@ -168,7 +168,7 @@ void assemble_cells(
 /// @param dofmap0 Test function (row) degree-of-freedom data holding
 /// the (0) dofmap, (1) dofmap block size and (2) dofmap cell indices.
 /// @param P0 Function that applies transformation P0.A in-place to
-/// transform trial degrees-of-freedom.
+/// transform test degrees-of-freedom.
 /// @param dofmap1 Trial function (column) degree-of-freedom data
 /// holding the (0) dofmap, (1) dofmap block size and (2) dofmap cell
 /// indices.
@@ -218,13 +218,13 @@ void assemble_exterior_facets(
   assert(facets1.size() == facets.size());
   for (std::size_t index = 0; index < facets.size(); index += 2)
   {
-    // Cell in the integration domain
+    // Cell in the integration domain, local facet index relative to the
+    // integration domain cell, and cells in the test and trial function
+    // meshes
     std::int32_t cell = facets[index];
-    // Cell in the test function mesh
+    std::int32_t local_facet = facets[index + 1];
     std::int32_t cell0 = facets0[index];
-    // Cell in the trial function mesh
     std::int32_t cell1 = facets1[index];
-    std::int32_t local_facet = facets[index + 1];
 
     // Get cell coordinates/geometry
     auto x_dofs = MDSPAN_IMPL_STANDARD_NAMESPACE::submdspan(
@@ -295,7 +295,7 @@ void assemble_exterior_facets(
 /// @param dofmap0 Test function (row) degree-of-freedom data holding
 /// the (0) dofmap, (1) dofmap block size and (2) dofmap cell indices.
 /// @param P0 Function that applies transformation P0.A in-place to
-/// transform trial degrees-of-freedom.
+/// transform test degrees-of-freedom.
 /// @param dofmap1 Trial function (column) degree-of-freedom data
 /// holding the (0) dofmap, (1) dofmap block size and (2) dofmap cell
 /// indices.