Optimize PreconditionRelaxation

diff --git a/include/deal.II/lac/precondition.h b/include/deal.II/lac/precondition.h
index cc428ae8dbdd87d43d24b807dc3b6591bb38db60..c75ede1e546a18bdbd1e89d5d9be7a76fdc5984c 100644 (file)
--- a/include/deal.II/lac/precondition.h
+++ b/include/deal.II/lac/precondition.h
@@ -525,6 +525,31 @@ protected:
 
 namespace internal
 {
+  // a helper type-trait that leverage SFINAE to figure out if MatrixType has
+  // ... MatrixType::vmult(VectorType &, const VectorType&,
+  // std::function<...>, std::function<...>) const
+  template <typename MatrixType, typename VectorType>
+  using vmult_functions_t = decltype(std::declval<MatrixType const>().vmult(
+    std::declval<VectorType &>(),
+    std::declval<const VectorType &>(),
+    std::declval<
+      const std::function<void(const unsigned int, const unsigned int)> &>(),
+    std::declval<
+      const std::function<void(const unsigned int, const unsigned int)> &>()));
+
+  template <typename MatrixType,
+            typename VectorType,
+            typename PreconditionerType>
+  constexpr bool has_vmult_with_std_functions =
+    is_supported_operation<vmult_functions_t, MatrixType, VectorType> &&
+      std::is_same<PreconditionerType, DiagonalMatrix<VectorType>>::value &&
+    (std::is_same<VectorType,
+                  dealii::Vector<typename VectorType::value_type>>::value ||
+     std::is_same<
+       VectorType,
+       LinearAlgebra::distributed::Vector<typename VectorType::value_type,
+                                          MemorySpace::Host>>::value);
+
   namespace PreconditionRelaxation
   {
     template <typename T, typename VectorType>
@@ -1051,45 +1076,157 @@ namespace internal
 
     template <typename MatrixType,
               typename VectorType,
-              std::enable_if_t<!IsBlockVector<VectorType>::value, VectorType>
-                * = nullptr>
+              std::enable_if_t<
+                !IsBlockVector<VectorType>::value &&
+                  !has_vmult_with_std_functions<MatrixType,
+                                                VectorType,
+                                                DiagonalMatrix<VectorType>>,
+                VectorType> * = nullptr>
     void
     step_operations(const MatrixType &                A,
                     const DiagonalMatrix<VectorType> &preconditioner,
                     VectorType &                      dst,
                     const VectorType &                src,
                     const double                      relaxation,
-                    VectorType &                      residual,
+                    VectorType &                      tmp,
                     VectorType &,
                     const unsigned int i,
                     const bool         transposed)
     {
+      using Number = typename VectorType::value_type;
+
       if (i == 0)
         {
-          const auto dst_ptr  = dst.begin();
-          const auto src_ptr  = src.begin();
-          const auto diag_ptr = preconditioner.get_vector().begin();
+          Number *      dst_ptr  = dst.begin();
+          const Number *src_ptr  = src.begin();
+          const Number *diag_ptr = preconditioner.get_vector().begin();
 
-          for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
-            dst_ptr[i] = relaxation * src_ptr[i] * diag_ptr[i];
+          if (relaxation == 1.0)
+            {
+              DEAL_II_OPENMP_SIMD_PRAGMA
+              for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
+                dst_ptr[i] = src_ptr[i] * diag_ptr[i];
+            }
+          else
+            {
+              DEAL_II_OPENMP_SIMD_PRAGMA
+              for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
+                dst_ptr[i] = relaxation * src_ptr[i] * diag_ptr[i];
+            }
         }
       else
         {
-          residual.reinit(src, true);
+          tmp.reinit(src, true);
 
-          const auto dst_ptr      = dst.begin();
-          const auto src_ptr      = src.begin();
-          const auto residual_ptr = residual.begin();
-          const auto diag_ptr     = preconditioner.get_vector().begin();
+          Number *      dst_ptr  = dst.begin();
+          const Number *src_ptr  = src.begin();
+          const Number *tmp_ptr  = tmp.begin();
+          const Number *diag_ptr = preconditioner.get_vector().begin();
 
           if (transposed)
-            A.Tvmult(residual, dst);
+            Tvmult(A, tmp, dst);
           else
-            A.vmult(residual, dst);
+            A.vmult(tmp, dst);
 
-          for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
-            dst_ptr[i] +=
-              relaxation * (src_ptr[i] - residual_ptr[i]) * diag_ptr[i];
+          if (relaxation == 1.0)
+            {
+              DEAL_II_OPENMP_SIMD_PRAGMA
+              for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
+                dst_ptr[i] += (src_ptr[i] - tmp_ptr[i]) * diag_ptr[i];
+            }
+          else
+            {
+              DEAL_II_OPENMP_SIMD_PRAGMA
+              for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
+                dst_ptr[i] +=
+                  relaxation * (src_ptr[i] - tmp_ptr[i]) * diag_ptr[i];
+            }
+        }
+    }
+
+    template <typename MatrixType,
+              typename VectorType,
+              std::enable_if_t<
+                !IsBlockVector<VectorType>::value &&
+                  has_vmult_with_std_functions<MatrixType,
+                                               VectorType,
+                                               DiagonalMatrix<VectorType>>,
+                VectorType> * = nullptr>
+    void
+    step_operations(const MatrixType &                A,
+                    const DiagonalMatrix<VectorType> &preconditioner,
+                    VectorType &                      dst,
+                    const VectorType &                src,
+                    const double                      relaxation,
+                    VectorType &                      tmp,
+                    VectorType &,
+                    const unsigned int i,
+                    const bool         transposed)
+    {
+      using Number = typename VectorType::value_type;
+
+      if (i == 0)
+        {
+          Number *      dst_ptr  = dst.begin();
+          const Number *src_ptr  = src.begin();
+          const Number *diag_ptr = preconditioner.get_vector().begin();
+
+          if (relaxation == 1.0)
+            {
+              DEAL_II_OPENMP_SIMD_PRAGMA
+              for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
+                dst_ptr[i] = src_ptr[i] * diag_ptr[i];
+            }
+          else
+            {
+              DEAL_II_OPENMP_SIMD_PRAGMA
+              for (unsigned int i = 0; i < dst.locally_owned_size(); ++i)
+                dst_ptr[i] = relaxation * src_ptr[i] * diag_ptr[i];
+            }
+        }
+      else
+        {
+          tmp.reinit(src, true);
+
+          Assert(transposed == false, ExcNotImplemented());
+
+          A.vmult(
+            tmp,
+            dst,
+            [&](const unsigned int start_range, const unsigned int end_range) {
+              // zero 'tmp' before running the vmult
+              // operation
+              if (end_range > start_range)
+                std::memset(tmp.begin() + start_range,
+                            0,
+                            sizeof(Number) * (end_range - start_range));
+            },
+            [&](const unsigned int begin, const unsigned int end) {
+              const Number *dst_ptr  = dst.begin();
+              const Number *src_ptr  = src.begin();
+              Number *      tmp_ptr  = tmp.begin();
+              const Number *diag_ptr = preconditioner.get_vector().begin();
+
+              // for efficiency reason, write back to temp_vector that is
+              // already read (avoid read-for-ownership)
+              if (relaxation == 1.0)
+                {
+                  DEAL_II_OPENMP_SIMD_PRAGMA
+                  for (std::size_t i = begin; i < end; ++i)
+                    tmp_ptr[i] =
+                      dst_ptr[i] + (src_ptr[i] - tmp_ptr[i]) * diag_ptr[i];
+                }
+              else
+                {
+                  DEAL_II_OPENMP_SIMD_PRAGMA
+                  for (std::size_t i = begin; i < end; ++i)
+                    tmp_ptr[i] = dst_ptr[i] + relaxation *
+                                                (src_ptr[i] - tmp_ptr[i]) *
+                                                diag_ptr[i];
+                }
+            });
+
+          tmp.swap(dst);
         }
     }
 
@@ -2536,29 +2673,6 @@ namespace internal
         solution.swap(solution_old);
     }
 
-    // a helper type-trait that leverage SFINAE to figure out if MatrixType has
-    // ... MatrixType::vmult(VectorType &, const VectorType&,
-    // std::function<...>, std::function<...>) const
-    template <typename MatrixType, typename VectorType>
-    using vmult_functions_t = decltype(std::declval<MatrixType const>().vmult(
-      std::declval<VectorType &>(),
-      std::declval<const VectorType &>(),
-      std::declval<
-        const std::function<void(const unsigned int, const unsigned int)> &>(),
-      std::declval<const std::function<void(const unsigned int,
-                                            const unsigned int)> &>()));
-
-    template <typename MatrixType,
-              typename VectorType,
-              typename PreconditionerType>
-    constexpr bool has_vmult_with_std_functions =
-      is_supported_operation<vmult_functions_t, MatrixType, VectorType>
-        &&  std::is_same<PreconditionerType, DiagonalMatrix<VectorType>>::value
-          &&std::is_same<
-            VectorType,
-            LinearAlgebra::distributed::Vector<typename VectorType::value_type,
-                                               MemorySpace::Host>>::value;
-
     // We need to have a separate declaration for static const members
 
     template <

diff --git a/tests/lac/precondition_relaxation_01.cc b/tests/lac/precondition_relaxation_01.cc
index 2127481b66a308dce17d9d5a546ca98851df5449..884a4779a2401a1123f6ba387421f09064411a8b 100644 (file)
--- a/tests/lac/precondition_relaxation_01.cc
+++ b/tests/lac/precondition_relaxation_01.cc
@@ -63,25 +63,64 @@ private:
   DiagonalMatrix<VectorType> diagonal_matrix;
 };
 
+template <typename SparseMatrixType>
+class MySparseMatrix : public Subscriptor
+{
+public:
+  MySparseMatrix(const SparseMatrixType &sparse_matrix)
+    : sparse_matrix(sparse_matrix)
+  {}
+
+  template <typename VectorType>
+  void
+  vmult(VectorType &      dst,
+        const VectorType &src,
+        const std::function<void(const unsigned int, const unsigned int)>
+          &operation_before_matrix_vector_product,
+        const std::function<void(const unsigned int, const unsigned int)>
+          &operation_after_matrix_vector_product) const
+  {
+    operation_before_matrix_vector_product(0, src.size());
+
+    sparse_matrix.vmult(dst, src);
+
+    operation_after_matrix_vector_product(0, src.size());
+  }
+
+private:
+  const SparseMatrixType &sparse_matrix;
+};
 
 
 template <typename PreconditionerType, typename VectorType>
 std::tuple<double, double, double, double>
-test(const PreconditionerType &preconditioner, const VectorType &src)
+test(const PreconditionerType &preconditioner,
+     const VectorType &        src,
+     const bool                test_transposed = true)
 {
   VectorType dst;
   dst.reinit(src);
 
+  dst = 1.0;
   preconditioner.vmult(dst, src);
   const double norm_0 = dst.l2_norm();
 
+  if (test_transposed)
+    {
+      dst = 1.0;
+      preconditioner.Tvmult(dst, src);
+    }
+  const double norm_2 = dst.l2_norm();
+
+  dst = 1.0;
   preconditioner.step(dst, src);
   const double norm_1 = dst.l2_norm();
 
-  preconditioner.Tvmult(dst, src);
-  const double norm_2 = dst.l2_norm();
-
-  preconditioner.Tstep(dst, src);
+  if (test_transposed)
+    {
+      dst = 1.0;
+      preconditioner.Tstep(dst, src);
+    }
   const double norm_3 = dst.l2_norm();
 
   return std::tuple<double, double, double, double>{norm_0,
@@ -159,6 +198,30 @@ main()
           results.emplace_back(test(preconditioner, src));
         }
 
+        {
+          // Test PreconditionRelaxation + DiagonalMatrix: optimized path with
+          // lambdas
+          using PreconditionerType = DiagonalMatrix<VectorType>;
+
+          using MyMatrixType = MySparseMatrix<MatrixType>;
+
+          MyMatrixType my_system_matrix(system_matrix);
+
+          PreconditionRelaxation<MyMatrixType, PreconditionerType>
+            preconditioner;
+
+          PreconditionRelaxation<MyMatrixType,
+                                 PreconditionerType>::AdditionalData ad;
+          ad.relaxation     = relaxation;
+          ad.n_iterations   = n_iterations;
+          ad.preconditioner = std::make_shared<PreconditionerType>();
+          ad.preconditioner->get_vector() = diagonal;
+
+          preconditioner.initialize(my_system_matrix, ad);
+
+          results.emplace_back(test(preconditioner, src, false));
+        }
+
         {
           // Test PreconditionRelaxation + DiagonalMatrix: optimized path
           using PreconditionerType = DiagonalMatrix<VectorType>;