Refactor FEEvaluationBase::read_cell_data() and add FEEvaluationBase::set_cell_data()

diff --git a/include/deal.II/matrix_free/fe_evaluation.h b/include/deal.II/matrix_free/fe_evaluation.h
index cbde241e1f8fc035ea31a8e7bba0941debd96baa..ee9afe610cc943e6b2eff59869b2ec92ef105a45 100644 (file)
--- a/include/deal.II/matrix_free/fe_evaluation.h
+++ b/include/deal.II/matrix_free/fe_evaluation.h
@@ -726,6 +726,16 @@ public:
   VectorizedArrayType
   read_cell_data(const AlignedVector<VectorizedArrayType> &array) const;
 
+  /**
+   * Provides a unified interface to set data in a vector of
+   * VectorizedArray fields of length MatrixFree::n_macro_cells() +
+   * MatrixFree::n_macro_ghost_cells() for both cells (plain read) and faces
+   * (indirect addressing).
+   */
+  void
+  set_cell_data(AlignedVector<VectorizedArrayType> &array,
+                const VectorizedArrayType &         value) const;
+
   /**
    * The same as above, just for std::array of length of VectorizedArrayType for
    * arbitrary data type.
@@ -735,6 +745,16 @@ public:
   read_cell_data(const AlignedVector<std::array<T, VectorizedArrayType::size()>>
                    &array) const;
 
+  /**
+   * The same as above, just for std::array of length of VectorizedArrayType for
+   * arbitrary data type.
+   */
+  template <typename T>
+  void
+  set_cell_data(
+    AlignedVector<std::array<T, VectorizedArrayType::size()>> &array,
+    const std::array<T, VectorizedArrayType::size()> &         value) const;
+
   /**
    * Return the id of the cells this FEEvaluation or FEFaceEvaluation is
    * associated with.
@@ -3821,6 +3841,45 @@ FEEvaluationBase<dim, n_components_, Number, is_face, VectorizedArrayType>::
 }
 
 
+namespace internal
+{
+  template <int dim,
+            int n_components_,
+            typename Number,
+            bool is_face,
+            typename VectorizedArrayType,
+            typename VectorizedArrayType2,
+            typename GlobalVectorType,
+            typename FU>
+  inline void
+  process_cell_data(
+    const FEEvaluationBase<dim,
+                           n_components_,
+                           Number,
+                           is_face,
+                           VectorizedArrayType> &       phi,
+    const MatrixFree<dim, Number, VectorizedArrayType> *matrix_info,
+    GlobalVectorType &                                  array,
+    VectorizedArrayType2 &                              out,
+    const FU &                                          fu)
+  {
+    Assert(matrix_info != nullptr, ExcNotImplemented());
+    AssertDimension(array.size(),
+                    matrix_info->get_task_info().cell_partition_data.back());
+
+    // 1) collect ids of cell
+    const auto cells = phi.get_cell_ids();
+
+    // 2) actually gather values
+    for (unsigned int i = 0; i < VectorizedArrayType::size(); ++i)
+      if (cells[i] != numbers::invalid_unsigned_int)
+        fu(out[i],
+           array[cells[i] / VectorizedArrayType::size()]
+                [cells[i] % VectorizedArrayType::size()]);
+  }
+} // namespace internal
+
+
 
 template <int dim,
           int n_components_,
@@ -3831,24 +3890,34 @@ inline VectorizedArrayType
 FEEvaluationBase<dim, n_components_, Number, is_face, VectorizedArrayType>::
   read_cell_data(const AlignedVector<VectorizedArrayType> &array) const
 {
-  Assert(matrix_info != nullptr, ExcNotImplemented());
-  AssertDimension(array.size(),
-                  matrix_info->get_task_info().cell_partition_data.back());
-
-  // 1) collect ids of cell
-  const auto cells = this->get_cell_ids();
-
-  // 2) actually gather values
   VectorizedArrayType out = Number(1.);
-  for (unsigned int i = 0; i < VectorizedArrayType::size(); ++i)
-    if (cells[i] != numbers::invalid_unsigned_int)
-      out[i] = array[cells[i] / VectorizedArrayType::size()]
-                    [cells[i] % VectorizedArrayType::size()];
+  internal::process_cell_data(
+    *this, this->matrix_info, array, out, [](auto &local, const auto &global) {
+      local = global;
+    });
   return out;
 }
 
 
 
+template <int dim,
+          int n_components_,
+          typename Number,
+          bool is_face,
+          typename VectorizedArrayType>
+inline void
+FEEvaluationBase<dim, n_components_, Number, is_face, VectorizedArrayType>::
+  set_cell_data(AlignedVector<VectorizedArrayType> &array,
+                const VectorizedArrayType &         in) const
+{
+  internal::process_cell_data(
+    *this, this->matrix_info, array, in, [](const auto &local, auto &global) {
+      global = local;
+    });
+}
+
+
+
 template <int dim,
           int n_components_,
           typename Number,
@@ -3860,24 +3929,36 @@ FEEvaluationBase<dim, n_components_, Number, is_face, VectorizedArrayType>::
   read_cell_data(const AlignedVector<std::array<T, VectorizedArrayType::size()>>
                    &array) const
 {
-  Assert(matrix_info != nullptr, ExcNotImplemented());
-  AssertDimension(array.size(),
-                  matrix_info->get_task_info().cell_partition_data.back());
-
-  // 1) collect ids of cell
-  const auto cells = this->get_cell_ids();
-
-  // 2) actually gather values
   std::array<T, VectorizedArrayType::size()> out;
-  for (unsigned int i = 0; i < VectorizedArrayType::size(); ++i)
-    if (cells[i] != numbers::invalid_unsigned_int)
-      out[i] = array[cells[i] / VectorizedArrayType::size()]
-                    [cells[i] % VectorizedArrayType::size()];
+  internal::process_cell_data(
+    *this, this->matrix_info, array, out, [](auto &local, const auto &global) {
+      local = global;
+    });
   return out;
 }
 
 
 
+template <int dim,
+          int n_components_,
+          typename Number,
+          bool is_face,
+          typename VectorizedArrayType>
+template <typename T>
+inline void
+FEEvaluationBase<dim, n_components_, Number, is_face, VectorizedArrayType>::
+  set_cell_data(
+    AlignedVector<std::array<T, VectorizedArrayType::size()>> &array,
+    const std::array<T, VectorizedArrayType::size()> &         in) const
+{
+  internal::process_cell_data(
+    *this, this->matrix_info, array, in, [](const auto &local, auto &global) {
+      global = local;
+    });
+}
+
+
+
 namespace internal
 {
   // allows to select between block vectors and non-block vectors, which

diff --git a/tests/matrix_free/read_cell_data_01.cc b/tests/matrix_free/read_cell_data_01.cc
index a59169cba3e7c0dc785837a824828474488d5e35..592b360d71a84079c962cdf9d32650c3762c59cb 100644 (file)
--- a/tests/matrix_free/read_cell_data_01.cc
+++ b/tests/matrix_free/read_cell_data_01.cc
@@ -108,10 +108,18 @@ private:
     unsigned int n_cells = data.n_cell_batches() + data.n_ghost_cell_batches();
     cell_ids.resize(n_cells);
 
+    FEEvaluation<dim, 1> fe_eval(data);
     for (unsigned int cell = 0; cell < n_cells; cell++)
-      for (auto lane = 0u; lane < data.n_active_entries_per_cell_batch(cell);
-           lane++)
-        cell_ids[cell][lane] = data.get_cell_iterator(cell, lane)->id();
+      {
+        fe_eval.reinit(cell);
+
+        std::array<CellId, VectorizedArrayType::size()> cell_ids_local;
+        for (auto lane = 0u; lane < data.n_active_entries_per_cell_batch(cell);
+             lane++)
+          cell_ids_local[lane] = data.get_cell_iterator(cell, lane)->id();
+
+        fe_eval.set_cell_data(cell_ids, cell_ids_local);
+      }
   }
 
   void