From 688f7ec72d8e6053fe83210367a8631425227c26 Mon Sep 17 00:00:00 2001
From: Martin Kronbichler <kronbichler@lnm.mw.tum.de>
Date: Tue, 11 May 2021 13:17:17 +0200
Subject: [PATCH] Come up with a more generic optimized path

---
 .../multigrid/mg_transfer_global_coarsening.h |  7 +-
 .../mg_transfer_global_coarsening.templates.h | 79 ++++++++++---------
 2 files changed, 45 insertions(+), 41 deletions(-)

diff --git a/include/deal.II/multigrid/mg_transfer_global_coarsening.h b/include/deal.II/multigrid/mg_transfer_global_coarsening.h
index f543fd8877..c975908a82 100644
--- a/include/deal.II/multigrid/mg_transfer_global_coarsening.h
+++ b/include/deal.II/multigrid/mg_transfer_global_coarsening.h
@@ -299,8 +299,11 @@ private:
   std::shared_ptr<const Utilities::MPI::Partitioner> partitioner_coarse;
 
   /**
-   * Internal vector needed for collecting all degrees of freedom of the
-   * fine cells.
+   * Internal vector needed for collecting all degrees of freedom of the fine
+   * cells. It is only initialized if the fine-level DoF indices touch DoFs
+   * other than the locally active ones (which we always assume can be
+   * accessed by the given vectors in the prolongate/restrict functions),
+   * otherwise it is left at size zero.
    */
   mutable LinearAlgebra::distributed::Vector<Number> vec_fine;
 
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 8f01b5462e..3540d9fcbe 100644
--- a/include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h
+++ b/include/deal.II/multigrid/mg_transfer_global_coarsening.templates.h
@@ -1238,12 +1238,6 @@ namespace internal
         transfer.schemes[1].fine_element_is_continuous =
           fe_fine.dofs_per_vertex > 0;
 
-      Assert(
-        transfer.schemes[0].fine_element_is_continuous ||
-          constraint_fine.n_constraints() == 0,
-        ExcMessage(
-          "For discontinuous elements, no constraints are supported at the moment."));
-
       // count coarse cells for each scheme (0, 1)
       {
         transfer.schemes[0].n_coarse_cells = 0; // reset
@@ -1665,12 +1659,6 @@ namespace internal
           transfer.schemes[fe_index_pair.second].fine_element_is_continuous =
             dof_handler_fine.get_fe(fe_index_pair.first.second)
               .dofs_per_vertex > 0;
-
-          Assert(
-            transfer.schemes[fe_index_pair.second].fine_element_is_continuous ||
-              constraint_fine.n_constraints() == 0,
-            ExcMessage(
-              "For discontinuous elements, no constraints are supported at the moment."));
         }
 
       precompute_constraints(dof_handler_coarse, constraint_coarse, transfer);
@@ -1786,11 +1774,16 @@ namespace internal
         for (unsigned int i = 0; i < fe_index_pairs.size(); i++)
           {
             level_dof_indices_coarse_[i] =
-              &transfer.schemes[i].level_dof_indices_coarse[0];
+              transfer.schemes[i].level_dof_indices_coarse.data();
             level_dof_indices_fine_[i] =
-              &transfer.schemes[i].level_dof_indices_fine[0];
+              transfer.schemes[i].level_dof_indices_fine.data();
           }
 
+        bool     fine_indices_touch_non_active_dofs = false;
+        IndexSet locally_active_dofs;
+        DoFTools::extract_locally_active_dofs(dof_handler_coarse,
+                                              locally_active_dofs);
+
         process_cells([&](const auto &cell_coarse, const auto &cell_fine) {
           const auto fe_pair_no =
             fe_index_pairs[std::pair<unsigned int, unsigned int>(
@@ -1805,22 +1798,34 @@ namespace internal
                 local_dof_indices_coarse
                   [fe_pair_no][lexicographic_numbering_coarse[fe_pair_no][i]]);
 
-
           cell_fine->get_dof_indices(local_dof_indices_fine[fe_pair_no]);
           for (unsigned int i = 0;
                i < transfer.schemes[fe_pair_no].dofs_per_cell_fine;
                i++)
-            level_dof_indices_fine_[fe_pair_no][i] =
-              transfer.partitioner_fine->global_to_local(
-                local_dof_indices_fine
-                  [fe_pair_no][lexicographic_numbering_fine[fe_pair_no][i]]);
-
+            {
+              level_dof_indices_fine_[fe_pair_no][i] =
+                transfer.partitioner_fine->global_to_local(
+                  local_dof_indices_fine
+                    [fe_pair_no][lexicographic_numbering_fine[fe_pair_no][i]]);
+              if (!locally_active_dofs.is_element(
+                    local_dof_indices_fine
+                      [fe_pair_no]
+                      [lexicographic_numbering_fine[fe_pair_no][i]]))
+                fine_indices_touch_non_active_dofs = true;
+            }
 
           level_dof_indices_coarse_[fe_pair_no] +=
             transfer.schemes[fe_pair_no].dofs_per_cell_coarse;
           level_dof_indices_fine_[fe_pair_no] +=
             transfer.schemes[fe_pair_no].dofs_per_cell_fine;
         });
+
+        // if all access goes to the locally active dofs on all ranks, we do
+        // not need the vec_fine vector
+        if (Utilities::MPI::max(static_cast<unsigned int>(
+                                  fine_indices_touch_non_active_dofs),
+                                comm) == 0)
+          transfer.vec_fine.reinit(0);
       }
 
       // ------------------------- prolongation matrix -------------------------
@@ -2136,12 +2141,9 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::prolongate(
 
   const unsigned int n_lanes = VectorizedArrayType::size();
 
-  const bool fine_vectors_are_compatible =
-    this->vec_fine.get_partitioner()->is_globally_compatible(
-      *dst.get_partitioner());
+  const bool use_dst_inplace = this->vec_fine.size() == 0;
 
-  const auto vec_fine_ptr =
-    fine_vectors_are_compatible ? &dst : &this->vec_fine;
+  const auto vec_fine_ptr = use_dst_inplace ? &dst : &this->vec_fine;
 
   // the following code is equivalent to:
   // this->constraint_coarse.distribute(this->vec_coarse);
@@ -2301,7 +2303,7 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::prolongate(
   if (schemes.size() > 0 && schemes.front().fine_element_is_continuous)
     vec_fine_ptr->compress(VectorOperation::add);
 
-  if (fine_vectors_are_compatible == false)
+  if (use_dst_inplace == false)
     dst.copy_locally_owned_data_from(this->vec_fine);
 }
 
@@ -2317,18 +2319,14 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
 
   const unsigned int n_lanes = VectorizedArrayType::size();
 
-  const bool fine_vectors_are_compatible =
-    this->vec_fine.get_partitioner()->is_globally_compatible(
-      *src.get_partitioner());
+  const bool use_src_inplace = this->vec_fine.size() == 0;
+  const auto vec_fine_ptr    = use_src_inplace ? &src : &this->vec_fine;
 
-  const auto vec_fine_ptr =
-    fine_vectors_are_compatible ? &src : &this->vec_fine;
-
-  if (fine_vectors_are_compatible == false)
+  if (use_src_inplace == false)
     this->vec_fine.copy_locally_owned_data_from(src);
 
-  if (schemes.size() > 0 && schemes.front().fine_element_is_continuous ||
-      fine_vectors_are_compatible == false)
+  if ((schemes.size() > 0 && schemes.front().fine_element_is_continuous) ||
+      use_src_inplace == false)
     vec_fine_ptr->update_ghost_values();
 
   this->vec_coarse.copy_locally_owned_data_from(dst);
@@ -2491,8 +2489,11 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
 
   const unsigned int n_lanes = VectorizedArrayType::size();
 
-  this->vec_fine.copy_locally_owned_data_from(src);
-  this->vec_fine.update_ghost_values();
+  const bool use_src_inplace = this->vec_fine.size() == 0;
+  const auto vec_fine_ptr    = use_src_inplace ? &src : &this->vec_fine;
+  if ((schemes.size() > 0 && schemes.front().fine_element_is_continuous) ||
+      use_src_inplace == false)
+    vec_fine_ptr->update_ghost_values();
 
   this->vec_coarse = 0.0;
 
@@ -2513,7 +2514,7 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
             {
               for (unsigned int i = 0; i < scheme.dofs_per_cell_fine; ++i)
                 this->vec_coarse.local_element(indices_coarse[i]) =
-                  this->vec_fine.local_element(indices_fine[i]);
+                  vec_fine_ptr->local_element(indices_fine[i]);
 
               indices_fine += scheme.dofs_per_cell_fine;
               indices_coarse += scheme.dofs_per_cell_coarse;
@@ -2550,7 +2551,7 @@ MGTwoLevelTransfer<dim, LinearAlgebra::distributed::Vector<Number>>::
               {
                 for (unsigned int i = 0; i < scheme.dofs_per_cell_fine; ++i)
                   evaluation_data_fine[i][v] =
-                    this->vec_fine.local_element(indices[i]);
+                    vec_fine_ptr->local_element(indices[i]);
 
                 indices += scheme.dofs_per_cell_fine;
               }
-- 
2.39.5