From 21f41d48883011cce6c216683f6e40dc41121cf2 Mon Sep 17 00:00:00 2001
From: Peter Munch <peterrmuench@gmail.com>
Date: Sun, 6 Mar 2022 08:50:42 +0100
Subject: [PATCH] GC: allow to pass in external partitioner

Conflicts:
	include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h
---
 .../multigrid/mg_transfer_global_coarsening.h | 116 +++++++++++++++++-
 .../mg_transfer_global_coarsening.templates.h |  23 ++++
 2 files changed, 133 insertions(+), 6 deletions(-)

diff --git a/include/deal.II/multigrid/mg_transfer_global_coarsening.h b/include/deal.II/multigrid/mg_transfer_global_coarsening.h
index 1a2379b94d..339fcad875 100644
--- a/include/deal.II/multigrid/mg_transfer_global_coarsening.h
+++ b/include/deal.II/multigrid/mg_transfer_global_coarsening.h
@@ -180,6 +180,16 @@ public:
   void
   interpolate(VectorType &dst, const VectorType &src) const;
 
+  /**
+   * Enable inplace vector operations if external and internal vectors
+   * are compatible.
+   */
+  void
+  enable_inplace_operations_if_possible(
+    const std::shared_ptr<const Utilities::MPI::Partitioner>
+      &partitioner_coarse,
+    const std::shared_ptr<const Utilities::MPI::Partitioner> &partitioner_fine);
+
   /**
    * Return the memory consumption of the allocated memory in this class.
    */
@@ -293,6 +303,16 @@ public:
   interpolate(LinearAlgebra::distributed::Vector<Number> &      dst,
               const LinearAlgebra::distributed::Vector<Number> &src) const;
 
+  /**
+   * Enable inplace vector operations if external and internal vectors
+   * are compatible.
+   */
+  void
+  enable_inplace_operations_if_possible(
+    const std::shared_ptr<const Utilities::MPI::Partitioner>
+      &partitioner_coarse,
+    const std::shared_ptr<const Utilities::MPI::Partitioner> &partitioner_fine);
+
   /**
    * Return the memory consumption of the allocated memory in this class.
    */
@@ -472,6 +492,25 @@ public:
     const std::function<void(const unsigned int, VectorType &)>
       &initialize_dof_vector = {});
 
+  /**
+   * Similar function to MGTransferMatrixFree::build() with the difference that
+   * the information for the prolongation for each level has already been built.
+   * So this function only tries to optimize the data structues of the two-level
+   * transfer operators, e.g., by enabling inplace vector operations, by
+   * checking if @p external_partitioners and the internal ones are compatible.
+   */
+  void
+  build(const std::vector<std::shared_ptr<const Utilities::MPI::Partitioner>>
+          &external_partitioners = {});
+
+  /**
+   * Same as above but taking a lambda for initializing vector instead of
+   * partitioners.
+   */
+  void
+  build(const std::function<void(const unsigned int, VectorType &)>
+          &initialize_dof_vector);
+
   /**
    * Perform prolongation.
    */
@@ -561,13 +600,12 @@ private:
   /**
    * Collection of the two-level transfer operators.
    */
-  const MGLevelObject<MGTwoLevelTransfer<dim, VectorType>> &transfer;
+  MGLevelObject<MGTwoLevelTransfer<dim, VectorType>> transfer;
 
   /**
-   * %Function to initialize internal level vectors.
+   * Function to initialize internal level vectors.
    */
-  const std::function<void(const unsigned int, VectorType &)>
-    initialize_dof_vector;
+  std::function<void(const unsigned int, VectorType &)> initialize_dof_vector;
 };
 
 
@@ -584,8 +622,74 @@ MGTransferGlobalCoarsening<dim, VectorType>::MGTransferGlobalCoarsening(
   const std::function<void(const unsigned int, VectorType &)>
     &initialize_dof_vector)
   : transfer(transfer)
-  , initialize_dof_vector(initialize_dof_vector)
-{}
+{
+  build(initialize_dof_vector);
+}
+
+
+
+template <int dim, typename VectorType>
+void
+MGTransferGlobalCoarsening<dim, VectorType>::build(
+  const std::vector<std::shared_ptr<const Utilities::MPI::Partitioner>>
+    &external_partitioners)
+{
+  if (external_partitioners.size() > 0)
+    {
+      Assert(
+        this->initialize_dof_vector == nullptr,
+        ExcMessage(
+          "A initialize_dof_vector function has already been registered in the constructor!"));
+
+      const unsigned int min_level = transfer.min_level();
+      const unsigned int max_level = transfer.max_level();
+
+      AssertDimension(external_partitioners.size(), max_level - min_level + 1);
+
+      this->initialize_dof_vector =
+        [external_partitioners, min_level](
+          const unsigned int level,
+          LinearAlgebra::distributed::Vector<typename VectorType::value_type>
+            &vec) { vec.reinit(external_partitioners[level - min_level]); };
+
+      for (unsigned int l = min_level + 1; l <= max_level; ++l)
+        transfer[l].enable_inplace_operations_if_possible(
+          external_partitioners[l - 1 - min_level],
+          external_partitioners[l - min_level]);
+    }
+}
+
+
+
+template <int dim, typename VectorType>
+void
+MGTransferGlobalCoarsening<dim, VectorType>::build(
+  const std::function<void(const unsigned int, VectorType &)>
+    &initialize_dof_vector)
+{
+  if (initialize_dof_vector)
+    {
+      const unsigned int n_levels =
+        transfer.max_level() - transfer.min_level() + 1;
+
+      std::vector<std::shared_ptr<const Utilities::MPI::Partitioner>>
+        external_partitioners(n_levels);
+
+      for (unsigned int level = 0; level < n_levels; ++level)
+        {
+          LinearAlgebra::distributed::Vector<typename VectorType::value_type>
+            vector;
+          initialize_dof_vector(level, vector);
+          external_partitioners[level] = vector.get_partitioner();
+        }
+
+      build(external_partitioners);
+    }
+  else
+    {
+      this->build();
+    }
+}
 
 
 
diff --git a/include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h b/include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h
index 05fcb5caae..94c7ca8eef 100644
--- a/include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h
+++ b/include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h
@@ -2900,6 +2900,29 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
 
 
 
+template <int dim, typename Number>
+void
+MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
+  enable_inplace_operations_if_possible(
+    const std::shared_ptr<const Utilities::MPI::Partitioner>
+      &partitioner_coarse,
+    const std::shared_ptr<const Utilities::MPI::Partitioner> &partitioner_fine)
+{
+  if (this->partitioner_coarse->is_globally_compatible(*partitioner_coarse))
+    {
+      this->partitioner_coarse = partitioner_coarse;
+      this->vec_coarse.reinit(0);
+    }
+
+  if (this->partitioner_fine->is_globally_compatible(*partitioner_fine))
+    {
+      this->partitioner_fine = partitioner_fine;
+      this->vec_fine.reinit(0);
+    }
+}
+
+
+
 template <int dim, typename Number>
 void
 MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
-- 
2.39.5