From abd99df33cb6f011f088acea7d5b4482beb03243 Mon Sep 17 00:00:00 2001
From: Martin Kronbichler <martin.kronbichler@rub.de>
Date: Sat, 20 Apr 2024 07:56:20 +0200
Subject: [PATCH] FEPointEvaluation: Allow evaluating more points together

---
 .../deal.II/matrix_free/fe_point_evaluation.h | 40 +++++++++++++-----
 .../tensor_product_point_kernels.h            | 41 ++++++++++++++++++-
 2 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/include/deal.II/matrix_free/fe_point_evaluation.h b/include/deal.II/matrix_free/fe_point_evaluation.h
index d66926e889..9ee92802d6 100644
--- a/include/deal.II/matrix_free/fe_point_evaluation.h
+++ b/include/deal.II/matrix_free/fe_point_evaluation.h
@@ -1842,8 +1842,6 @@ FEPointEvaluationBase<n_components_, dim, spacedim, Number>::
   , current_face_number(other.current_face_number)
   , fast_path(other.fast_path)
   , is_reinitialized(false)
-  , shapes(other.shapes)
-  , shapes_faces(other.shapes_faces)
   , is_interior(other.is_interior)
 {
   connection_is_reinitialized = mapping_info->connect_is_reinitialized(
@@ -1879,8 +1877,6 @@ FEPointEvaluationBase<n_components_, dim, spacedim, Number>::
   , current_face_number(other.current_face_number)
   , fast_path(other.fast_path)
   , is_reinitialized(false)
-  , shapes(other.shapes)
-  , shapes_faces(other.shapes_faces)
   , is_interior(other.is_interior)
 {
   connection_is_reinitialized = mapping_info->connect_is_reinitialized(
@@ -2064,13 +2060,16 @@ FEPointEvaluationBase<n_components_, dim, spacedim, Number>::do_reinit()
   if (!is_linear && fast_path)
     {
       const std::size_t n_shapes = poly.size();
+      if (is_face)
+        shapes_faces.resize_fast(n_q_batches * n_shapes);
+      else
+        shapes.resize_fast(n_q_batches * n_shapes);
 
       for (unsigned int qb = 0; qb < n_q_batches; ++qb)
         if (is_face)
           {
             if (dim > 1)
               {
-                shapes_faces.resize_fast(n_q_batches * n_shapes);
                 internal::compute_values_of_array(
                   shapes_faces.data() + qb * n_shapes,
                   poly,
@@ -2080,12 +2079,31 @@ FEPointEvaluationBase<n_components_, dim, spacedim, Number>::do_reinit()
           }
         else
           {
-            shapes.resize_fast(n_q_batches * n_shapes);
-            internal::compute_values_of_array(
-              shapes.data() + qb * n_shapes,
-              poly,
-              unit_point_ptr[qb],
-              update_flags & UpdateFlags::update_gradients ? 1 : 0);
+            if (update_flags & UpdateFlags::update_gradients)
+              {
+                internal::compute_values_of_array(shapes.data() + qb * n_shapes,
+                                                  poly,
+                                                  unit_point_ptr[qb],
+                                                  1);
+              }
+            else if (qb + 1 < n_q_batches)
+              {
+                // Use function with reduced overhead to compute for two
+                // points at once
+                internal::compute_values_of_array_in_pairs(
+                  shapes.data() + qb * n_shapes,
+                  poly,
+                  unit_point_ptr[qb],
+                  unit_point_ptr[qb + 1]);
+                ++qb;
+              }
+            else
+              {
+                internal::compute_values_of_array(shapes.data() + qb * n_shapes,
+                                                  poly,
+                                                  unit_point_ptr[qb],
+                                                  0);
+              }
           }
     }
 
diff --git a/include/deal.II/matrix_free/tensor_product_point_kernels.h b/include/deal.II/matrix_free/tensor_product_point_kernels.h
index 0fed110392..9b7bb7a2a9 100644
--- a/include/deal.II/matrix_free/tensor_product_point_kernels.h
+++ b/include/deal.II/matrix_free/tensor_product_point_kernels.h
@@ -52,8 +52,8 @@ namespace internal
 
 
   /**
-   * Computes the values and derivatives of the 1d polynomials @p poly at the
-   * specified point @p p and stores it in @p shapes.
+   * Compute the values and derivatives of the 1d polynomials @p poly at the
+   * specified point @p p and store them in @p shapes.
    */
   template <int dim, typename Number>
   inline void
@@ -90,6 +90,43 @@ namespace internal
 
 
 
+  /**
+   * Evaluate the 1d polynomials @p poly at the two specified points @p p0 and
+   * @p p1 and store them in @p shapes. This function can be used as a more
+   * efficient alternative to the compute_values_of_array() function, because
+   * of reduced overhead when querying the polynomials (which usually have
+   * loop bounds that are not known at compile time).
+   */
+  template <int dim, typename Number>
+  inline void
+  compute_values_of_array_in_pairs(
+    dealii::ndarray<Number, 2, dim>                    *shapes,
+    const std::vector<Polynomials::Polynomial<double>> &poly,
+    const Point<dim, Number>                           &p0,
+    const Point<dim, Number>                           &p1)
+  {
+    // Use 'int' variable here to let the compiler apply additional
+    // optimizations, in particular regarding multiplications and additions in
+    // loop increments that are known not to overflow/wrap around (as is the
+    // case for unsigned int).
+    const int n_shapes = poly.size();
+
+    std::array<Number, 2 * dim> point, result;
+    for (unsigned int d = 0; d < dim; ++d)
+      point[d] = p0[d];
+    for (unsigned int d = 0; d < dim; ++d)
+      point[dim + d] = p1[d];
+    for (int i = 0; i < n_shapes; ++i)
+      {
+        poly[i].values_of_array(point, 0, &result);
+        for (unsigned int j = 0; j < 2; ++j)
+          for (unsigned int d = 0; d < dim; ++d)
+            shapes[j * n_shapes + i][0][d] = result[j * dim + d];
+      }
+  }
+
+
+
   /**
    * Interpolate inner dimensions of tensor product shape functions.
    */
-- 
2.39.5