Share common Scatterer between Vectors on same FunctionSpace

cpp/dolfinx/common/Scatterer.h

@@ -294,16 +294,16 @@ class Scatterer
     for (std::size_t i = 0; i < _src.size(); ++i)
     {
       int ierr = MPI_Irecv(recv_buffer + _displs_remote[i], _sizes_remote[i],
-                           dolfinx::MPI::mpi_t<T>, _src[i], MPI_ANY_TAG,
+                           dolfinx::MPI::mpi_t<T>, _src[i], p2p_tag_fwd,
                            _comm0.comm(), &requests[i]);
       dolfinx::MPI::check_error(_comm0.comm(), ierr);
     }
 
     for (std::size_t i = 0; i < _dest.size(); ++i)
     {
       int ierr = MPI_Isend(send_buffer + _displs_local[i], _sizes_local[i],
-                           dolfinx::MPI::mpi_t<T>, _dest[i], 0, _comm0.comm(),
-                           &requests[i + _src.size()]);
+                           dolfinx::MPI::mpi_t<T>, _dest[i], p2p_tag_fwd,
+                           _comm0.comm(), &requests[i + _src.size()]);
       dolfinx::MPI::check_error(_comm0.comm(), ierr);
     }
   }
@@ -380,7 +380,7 @@ class Scatterer
     for (std::size_t i = 0; i < _dest.size(); i++)
     {
       int ierr = MPI_Irecv(recv_buffer + _displs_local[i], _sizes_local[i],
-                           dolfinx::MPI::mpi_t<T>, _dest[i], MPI_ANY_TAG,
+                           dolfinx::MPI::mpi_t<T>, _dest[i], p2p_tag_rev,
                            _comm0.comm(), &requests[i]);
       dolfinx::MPI::check_error(_comm0.comm(), ierr);
     }
@@ -390,8 +390,8 @@ class Scatterer
     for (std::size_t i = 0; i < _src.size(); i++)
     {
       int ierr = MPI_Isend(send_buffer + _displs_remote[i], _sizes_remote[i],
-                           dolfinx::MPI::mpi_t<T>, _src[i], 0, _comm0.comm(),
-                           &requests[i + _dest.size()]);
+                           dolfinx::MPI::mpi_t<T>, _src[i], p2p_tag_rev,
+                           _comm0.comm(), &requests[i + _dest.size()]);
       dolfinx::MPI::check_error(_comm0.comm(), ierr);
     }
   }
@@ -495,6 +495,14 @@ class Scatterer
   }
 
 private:
+  // MPI tags for point-to-point forward/reverse scatter messages.
+  // Both directions share communicator _comm0, so a fixed, distinct
+  // tag per direction is required to stop a forward message being
+  // matched against a reverse receive (or vice versa) irrespective of
+  // the relative order in which the two directions are posted.
+  static constexpr int p2p_tag_fwd = 0;
+  static constexpr int p2p_tag_rev = 1;
+
   // Communicator where the source ranks own the indices in the callers
   // halo, and the destination ranks 'ghost' indices owned by the
   // caller. I.e.,

cpp/dolfinx/fem/Function.h

@@ -55,9 +55,7 @@ class Function
   /// @brief Create function on given function space.
   /// @param[in] V The function space
   explicit Function(std::shared_ptr<const FunctionSpace<geometry_type>> V)
-      : _function_space(V), _x(std::make_shared<la::Vector<value_type>>(
-                                V->dofmaps().front()->index_map,
-                                V->dofmaps().front()->index_map_bs()))
+      : _function_space(V), _x(std::make_shared<la::Vector<value_type>>(*V))
   {
     if (!V->component().empty())
     {
@@ -122,8 +120,7 @@ class Function
     auto [V, map] = _function_space->collapse();
 
     // Create new vector
-    auto x = std::make_shared<la::Vector<value_type>>(
-        V.dofmap()->index_map, V.dofmap()->index_map_bs());
+    auto x = std::make_shared<la::Vector<value_type>>(V);
 
     // Copy values into new vector
     std::span<const value_type> x_old = _x->array();

cpp/dolfinx/fem/FunctionSpace.h

@@ -15,6 +15,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <dolfinx/common/IndexMap.h>
+#include <dolfinx/common/Scatterer.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
@@ -29,7 +30,10 @@ namespace dolfinx::fem
 /// degrees-of-freedom.
 /// @tparam T The floating point (real) type of the mesh geometry and
 /// the finite element basis.
-template <std::floating_point T>
+/// @tparam ScatterT Scatterer type used to scatter/gather dof data
+/// between the local (owned + ghost) and global dof layouts.
+template <std::floating_point T,
+          typename ScatterT = common::Scatterer<std::vector<std::int32_t>>>
 class FunctionSpace
 {
 public:
@@ -45,6 +49,8 @@ class FunctionSpace
                 std::shared_ptr<const FiniteElement<geometry_type>> element,
                 std::shared_ptr<const DofMap> dofmap)
       : _mesh(mesh), _elements{element}, _dofmaps{std::move(dofmap)},
+        _scatterer(std::make_shared<ScatterT>(
+            *_dofmaps.front()->index_map, _dofmaps.front()->index_map_bs())),
         _id(boost::uuids::random_generator()()), _root_space_id(_id)
   {
     // Do nothing
@@ -63,6 +69,8 @@ class FunctionSpace
       std::vector<std::shared_ptr<const FiniteElement<geometry_type>>> elements,
       std::vector<std::shared_ptr<const DofMap>> dofmaps)
       : _mesh(mesh), _elements(elements), _dofmaps(std::move(dofmaps)),
+        _scatterer(std::make_shared<ScatterT>(
+            *_dofmaps.front()->index_map, _dofmaps.front()->index_map_bs())),
         _id(boost::uuids::random_generator()()), _root_space_id(_id)
   {
     std::vector<mesh::CellType> cell_types = mesh->topology()->cell_types();
@@ -400,6 +408,12 @@ class FunctionSpace
     return _dofmaps;
   }
 
+  /// @brief The scatterer associated with the function space.
+  ///
+  /// Used to scatter data between the local (owned + ghost) dof layout
+  /// and the global layout.
+  std::shared_ptr<const ScatterT> scatterer() const { return _scatterer; }
+
 private:
   // The mesh
   std::shared_ptr<const mesh::Mesh<geometry_type>> _mesh;
@@ -410,6 +424,10 @@ class FunctionSpace
   // The dofmap
   std::vector<std::shared_ptr<const DofMap>> _dofmaps;
 
+  // An associated scatterer for the function space. This is used to scatter
+  // data from the local dof layout to the global dof layout and vice versa.
+  std::shared_ptr<ScatterT> _scatterer;
+
   // The component w.r.t. to root space
   std::vector<int> _component;
 

cpp/dolfinx/fem/assembler.h

@@ -37,8 +37,6 @@ template <dolfinx::scalar T, std::floating_point U>
 class Expression;
 template <dolfinx::scalar T, std::floating_point U>
 class Form;
-template <std::floating_point T>
-class FunctionSpace;
 
 /// @brief Evaluate an Expression on cells or facets.
 ///

cpp/dolfinx/io/vtk_utils.h

@@ -26,12 +26,6 @@
 
 namespace dolfinx
 {
-namespace fem
-{
-template <std::floating_point T>
-class FunctionSpace;
-}
-
 namespace mesh
 {
 enum class CellType : std::int8_t;

cpp/dolfinx/la/Vector.h

@@ -10,6 +10,7 @@
 #include <algorithm>
 #include <cmath>
 #include <complex>
+#include <concepts>
 #include <dolfinx/common/IndexMap.h>
 #include <dolfinx/common/Scatterer.h>
 #include <dolfinx/common/types.h>
@@ -20,6 +21,12 @@
 #include <type_traits>
 #include <vector>
 
+namespace dolfinx::fem
+{
+template <std::floating_point T, typename ScatterT>
+class FunctionSpace;
+} // namespace dolfinx::fem
+
 namespace dolfinx::la
 {
 /// @brief la::Vector scatter pack/unpack function concept.
@@ -137,6 +144,22 @@ class Vector
   {
   }
 
+  /// @brief Create a distributed vector from a FunctionSpace, inheriting a
+  /// shared Scatterer.
+  /// @tparam U Floating point (geometry) type of the FunctionSpace.
+  /// @param V FunctionSpace defining the block size and dofmap layout of the
+  /// vector, and sharing a common Scatterer.
+  template <std::floating_point U>
+  explicit Vector(
+      const fem::FunctionSpace<U, common::Scatterer<ScatterContainer>>& V)
+      : _map(V.dofmap()->index_map), _bs(V.dofmap()->index_map_bs()),
+        _x(_bs * (_map->size_local() + _map->num_ghosts())),
+        _scatterer(V.scatterer()),
+        _buffer_local(_scatterer->local_indices().size()),
+        _buffer_remote(_scatterer->remote_indices().size())
+  {
+  }
+
   /// Copy constructor
   Vector(const Vector& x) = default;