Permit only required data to be passed to IndexMap scatters

cpp/dolfinx/common/IndexMap.cpp

@@ -180,7 +180,7 @@ common::stack_index_maps(
   {
     const int bs = maps[f].second;
     const std::vector<std::int32_t>& forward_indices
-        = maps[f].first.get().shared_indices().array();
+        = maps[f].first.get().scatter_fwd_indices().array();
     const std::int64_t offset = bs * maps[f].first.get().local_range()[0];
     for (std::int32_t local_index : forward_indices)
     {
@@ -343,8 +343,9 @@ IndexMap::IndexMap(MPI_Comm mpi_comm, std::int32_t local_size,
   // Create communicators with directed edges: (0) owner -> ghost,
   // (1) ghost -> owner
   {
-    // Dummy weight to work around an MPICH bug. Would prefer to pass
-    // MPI_UNWEIGHTED.
+    // Dummy weight to work around a bug in older MPICH versions. Issue
+    // appears in MPICH v3.3.2 (Ubuntu 20.04) and not in v3.4 (Ubuntu
+    // 20.10). Would prefer to pass MPI_UNWEIGHTED.
     std::vector<int> weights(std::max(halo_src_ranks.size(), dest_ranks.size()),
                              1);
 
@@ -525,12 +526,18 @@ std::vector<std::int64_t> IndexMap::global_indices() const
 }
 //-----------------------------------------------------------------------------
 const graph::AdjacencyList<std::int32_t>&
-IndexMap::shared_indices() const noexcept
+IndexMap::scatter_fwd_indices() const noexcept
 {
   assert(_shared_indices);
   return *_shared_indices;
 }
 //-----------------------------------------------------------------------------
+const std::vector<std::int32_t>&
+IndexMap::scatter_fwd_ghost_positions() const noexcept
+{
+  return _ghost_pos_recv_fwd;
+}
+//-----------------------------------------------------------------------------
 std::vector<int> IndexMap::ghost_owner_rank() const
 {
   int indegree(-1), outdegree(-2), weighted(-1);

cpp/dolfinx/fem/DofMap.cpp

@@ -93,7 +93,6 @@ fem::DofMap build_collapsed_dofmap(MPI_Comm comm, const DofMap& dofmap_view,
 
   // Send new global indices for owned dofs to non-owning process, and
   // receive new global indices from owner
-
   std::vector<std::int64_t> global_index_remote(
       dofmap_view.index_map->num_ghosts());
   dofmap_view.index_map->scatter_fwd(

cpp/dolfinx/fem/Form.h

@@ -315,9 +315,9 @@ class Form
         std::set<std::int32_t> fwd_shared;
         if (topology.index_map(tdim)->num_ghosts() == 0)
         {
-          fwd_shared.insert(
-              topology.index_map(tdim - 1)->shared_indices().array().begin(),
-              topology.index_map(tdim - 1)->shared_indices().array().end());
+          const std::vector<std::int32_t>& fwd_indices
+              = topology.index_map(tdim - 1)->scatter_fwd_indices().array();
+          fwd_shared.insert(fwd_indices.begin(), fwd_indices.end());
         }
 
         for (auto f = tagged_entities.begin(); f != entity_end; ++f)
@@ -401,9 +401,9 @@ class Form
         // Only need to consider shared facets when there are no ghost cells
         if (topology.index_map(tdim)->num_ghosts() == 0)
         {
-          fwd_shared_facets.insert(
-              topology.index_map(tdim - 1)->shared_indices().array().begin(),
-              topology.index_map(tdim - 1)->shared_indices().array().end());
+          const std::vector<std::int32_t>& fwd_indices
+              = topology.index_map(tdim - 1)->scatter_fwd_indices().array();
+          fwd_shared_facets.insert(fwd_indices.begin(), fwd_indices.end());
         }
 
         const int num_facets = topology.index_map(tdim - 1)->size_local();

cpp/dolfinx/fem/dofmapbuilder.cpp

@@ -392,7 +392,7 @@ std::pair<std::vector<std::int64_t>, std::vector<int>> get_global_indices(
     {
       shared_entity[d] = std::vector<bool>(map->size_local(), false);
       const std::vector<std::int32_t>& forward_indices
-          = map->shared_indices().array();
+          = map->scatter_fwd_indices().array();
       for (auto entity : forward_indices)
         shared_entity[d][entity] = true;
     }

cpp/dolfinx/la/Vector.h

@@ -76,10 +76,20 @@ class Vector
   void scatter_fwd_begin()
   {
     assert(_map);
-    const std::int32_t local_size = _bs * _map->size_local();
-    xtl::span<const T> xlocal(_x.data(), local_size);
-    _map->scatter_fwd_begin(xlocal, _datatype, _request, _buffer_send_fwd,
-                            _buffer_recv_fwd);
+
+    // Pack send buffer
+    const std::vector<std::int32_t>& indices
+        = _map->scatter_fwd_indices().array();
+    _buffer_send_fwd.resize(_bs * indices.size());
+    for (std::size_t i = 0; i < indices.size(); ++i)
+    {
+      std::copy_n(std::next(_x.cbegin(), _bs * indices[i]), _bs,
+                  std::next(_buffer_send_fwd.begin(), _bs * i));
+    }
+
+    _buffer_recv_fwd.resize(_bs * _map->num_ghosts());
+    _map->scatter_fwd_begin(xtl::span<const T>(_buffer_send_fwd), _datatype,
+                            _request, xtl::span<T>(_buffer_recv_fwd));
   }
 
   /// End scatter of local data from owner to ghosts on other ranks
@@ -89,8 +99,17 @@ class Vector
     assert(_map);
     const std::int32_t local_size = _bs * _map->size_local();
     xtl::span xremote(_x.data() + local_size, _map->num_ghosts() * _bs);
-    _map->scatter_fwd_end(xremote, _request,
-                          xtl::span<const T>(_buffer_recv_fwd));
+    _map->scatter_fwd_end(_request);
+
+    // Copy received data into ghost positions
+    const std::vector<std::int32_t>& scatter_fwd_ghost_pos
+        = _map->scatter_fwd_ghost_positions();
+    for (std::size_t i = 0; i < _map->num_ghosts(); ++i)
+    {
+      const int pos = scatter_fwd_ghost_pos[i];
+      std::copy_n(std::next(_buffer_recv_fwd.cbegin(), _bs * pos), _bs,
+                  std::next(xremote.begin(), _bs * i));
+    }
   }
 
   /// Scatter local data to ghost positions on other ranks
@@ -105,25 +124,55 @@ class Vector
   /// @note Collective MPI operation
   void scatter_rev_begin()
   {
+    // Pack send buffer
     const std::int32_t local_size = _bs * _map->size_local();
     xtl::span<const T> xremote(_x.data() + local_size,
                                _map->num_ghosts() * _bs);
-    _map->scatter_rev_begin(xremote, _datatype, _request, _buffer_recv_fwd,
-                            _buffer_send_fwd);
+    const std::vector<std::int32_t>& scatter_fwd_ghost_pos
+        = _map->scatter_fwd_ghost_positions();
+    _buffer_recv_fwd.resize(_bs * scatter_fwd_ghost_pos.size());
+    for (std::size_t i = 0; i < scatter_fwd_ghost_pos.size(); ++i)
+    {
+      const int pos = scatter_fwd_ghost_pos[i];
+      std::copy_n(std::next(xremote.cbegin(), _bs * i), _bs,
+                  std::next(_buffer_recv_fwd.begin(), _bs * pos));
+    }
+
+    // Resize receive buffer and begin scatter
+    _buffer_send_fwd.resize(_bs * _map->scatter_fwd_indices().array().size());
+    _map->scatter_rev_begin(xtl::span<const T>(_buffer_recv_fwd), _datatype,
+                            _request, xtl::span<T>(_buffer_send_fwd));
   }
 
-  /// End scatter of ghost data to owner. This process may receive data from
-  /// more than one process, and the received data can be summed or
+  /// End scatter of ghost data to owner. This process may receive data
+  /// from more than one process, and the received data can be summed or
   /// inserted into the local portion of the vector.
   /// @param op The operation to perform when adding/setting received
   /// values (add or insert)
   /// @note Collective MPI operation
-  void scatter_rev_end(dolfinx::common::IndexMap::Mode op)
+  void scatter_rev_end(common::IndexMap::Mode op)
   {
-    const std::int32_t local_size = _bs * _map->size_local();
-    xtl::span xlocal(_x.data(), local_size);
-    _map->scatter_rev_end(xlocal, _request,
-                          xtl::span<const T>(_buffer_send_fwd), op);
+    // Complete scatter
+    _map->scatter_rev_end(_request);
+
+    // Copy/accumulate into owned part of the vector
+    const std::vector<std::int32_t>& shared_indices
+        = _map->scatter_fwd_indices().array();
+    switch (op)
+    {
+    case common::IndexMap::Mode::insert:
+      for (std::size_t i = 0; i < shared_indices.size(); ++i)
+      {
+        std::copy_n(std::next(_buffer_send_fwd.cbegin(), _bs * i), _bs,
+                    std::next(_x.begin(), _bs * shared_indices[i]));
+      }
+      break;
+    case common::IndexMap::Mode::add:
+      for (std::size_t i = 0; i < shared_indices.size(); ++i)
+        for (int j = 0; j < _bs; ++j)
+          _x[shared_indices[i] * _bs + j] += _buffer_send_fwd[i * _bs + j];
+      break;
+    }
   }
 
   /// Scatter ghost data to owner. This process may receive data from
@@ -227,7 +276,9 @@ T inner_product(const Vector<T, Allocator>& a, const Vector<T, Allocator>& b)
 
   const T local = std::transform_reduce(
       x_a.begin(), x_a.begin() + local_size, x_b.begin(), static_cast<T>(0),
-      std::plus<T>(), [](T a, T b) -> T {
+      std::plus<T>(),
+      [](T a, T b) -> T
+      {
         if constexpr (std::is_same<T, std::complex<double>>::value
                       or std::is_same<T, std::complex<float>>::value)
           return std::conj(a) * b;

cpp/dolfinx/mesh/Topology.cpp

@@ -131,8 +131,8 @@ std::vector<bool> mesh::compute_boundary_facets(const Topology& topology)
   if (facets->num_ghosts() == 0)
   {
     fwd_shared_facets
-        = std::set<std::int32_t>(facets->shared_indices().array().begin(),
-                                 facets->shared_indices().array().end());
+        = std::set<std::int32_t>(facets->scatter_fwd_indices().array().begin(),
+                                 facets->scatter_fwd_indices().array().end());
   }
 
   std::shared_ptr<const graph::AdjacencyList<std::int32_t>> fc
@@ -439,7 +439,8 @@ mesh::create_topology(MPI_Comm comm,
          == node_remap.end());
   std::for_each(global_to_local_vertices.begin(),
                 global_to_local_vertices.end(),
-                [&remap = std::as_const(node_remap)](auto& v) {
+                [&remap = std::as_const(node_remap)](auto& v)
+                {
                   if (v.second >= 0)
                     v.second = remap[v.second];
                 });

cpp/dolfinx/mesh/utils.cpp

@@ -468,14 +468,14 @@ std::vector<std::int32_t> mesh::exterior_facet_indices(const Mesh& mesh)
   mesh.topology_mutable().create_connectivity(tdim - 1, tdim);
   auto f_to_c = topology.connectivity(tdim - 1, tdim);
   assert(topology.index_map(tdim - 1));
-  std::set<std::int32_t> fwd_shared_facets;
 
   // Only need to consider shared facets when there are no ghost cells
+  std::set<std::int32_t> fwd_shared_facets;
   if (topology.index_map(tdim)->num_ghosts() == 0)
   {
     fwd_shared_facets.insert(
-        topology.index_map(tdim - 1)->shared_indices().array().begin(),
-        topology.index_map(tdim - 1)->shared_indices().array().end());
+        topology.index_map(tdim - 1)->scatter_fwd_indices().array().begin(),
+        topology.index_map(tdim - 1)->scatter_fwd_indices().array().end());
   }
 
   // Find all owned facets (not ghost) with only one attached cell, which are