From ce767ed5f41b0b915faaa26ad3fcf1945ed4083b Mon Sep 17 00:00:00 2001
From: Daniel Arndt <arndtd@ornl.gov>
Date: Fri, 2 Aug 2024 16:01:49 -0400
Subject: [PATCH] Portable MatrixFree: compute diagonal

---
 include/deal.II/matrix_free/tools.h           | 117 +++++++++++-
 .../matrix_free_kokkos/compute_diagonal_01.cc |  90 ++++++++++
 ...iagonal_01.with_p4est=true.mpirun=1.output |  21 +++
 ...iagonal_01.with_p4est=true.mpirun=3.output |  65 +++++++
 .../compute_diagonal_util.h                   | 167 ++++++++++++++++++
 5 files changed, 458 insertions(+), 2 deletions(-)
 create mode 100644 tests/matrix_free_kokkos/compute_diagonal_01.cc
 create mode 100644 tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=1.output
 create mode 100644 tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=3.output
 create mode 100644 tests/matrix_free_kokkos/compute_diagonal_util.h

diff --git a/include/deal.II/matrix_free/tools.h b/include/deal.II/matrix_free/tools.h
index e23acadee2..21531c3785 100644
--- a/include/deal.II/matrix_free/tools.h
+++ b/include/deal.II/matrix_free/tools.h
@@ -21,9 +21,10 @@
 
 #include <deal.II/matrix_free/fe_evaluation.h>
 #include <deal.II/matrix_free/matrix_free.h>
+#include <deal.II/matrix_free/portable_fe_evaluation.h>
+#include <deal.II/matrix_free/portable_matrix_free.h>
 #include <deal.II/matrix_free/vector_access_internal.h>
 
-
 DEAL_II_NAMESPACE_OPEN
 
 /**
@@ -104,7 +105,27 @@ namespace MatrixFreeTools
     const unsigned int first_selected_component = 0,
     const unsigned int first_vector_component   = 0);
 
-
+  /**
+   * Same as above but for Portable::MatrixFree.
+   */
+  template <int dim,
+            int fe_degree,
+            int n_q_points_1d,
+            int n_components,
+            typename Number,
+            typename MemorySpace,
+            typename QuadOperation>
+  void
+  compute_diagonal(
+    const Portable::MatrixFree<dim, Number>                 &matrix_free,
+    LinearAlgebra::distributed::Vector<Number, MemorySpace> &diagonal_global,
+    const QuadOperation                                     &quad_operation,
+    EvaluationFlags::EvaluationFlags                         evaluation_flags,
+    EvaluationFlags::EvaluationFlags                         integration_flags,
+    const unsigned int                                       dof_no  = 0,
+    const unsigned int                                       quad_no = 0,
+    const unsigned int first_selected_component                      = 0,
+    const unsigned int first_vector_component                        = 0);
 
   /**
    * Same as above but with a class and a function pointer.
@@ -1340,6 +1361,98 @@ namespace MatrixFreeTools
 
   } // namespace internal
 
+
+  template <int dim, int fe_degree, typename Number, typename QuadOperation>
+  class CellAction
+  {
+  public:
+    CellAction(const QuadOperation                   &quad_operation,
+               const EvaluationFlags::EvaluationFlags evaluation_flags,
+               const EvaluationFlags::EvaluationFlags integration_flags)
+      : m_quad_operation(quad_operation)
+      , m_evaluation_flags(evaluation_flags)
+      , m_integration_flags(integration_flags)
+    {}
+
+    KOKKOS_FUNCTION void
+    operator()(const unsigned int /*cell*/,
+               const typename Portable::MatrixFree<dim, Number>::Data *gpu_data,
+               Portable::SharedData<dim, Number> *shared_data,
+               const Number *,
+               Number *dst) const
+    {
+      Portable::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, Number> fe_eval(
+        gpu_data, shared_data);
+      constexpr int dofs_per_cell = decltype(fe_eval)::tensor_dofs_per_cell;
+      double        diagonal[dofs_per_cell] = {};
+      for (unsigned int i = 0; i < dofs_per_cell; ++i)
+        {
+          Kokkos::parallel_for(
+            Kokkos::TeamThreadRange(shared_data->team_member, dofs_per_cell),
+            [&](int j) { fe_eval.submit_dof_value(i == j ? 1 : 0, j); });
+          fe_eval.evaluate(m_evaluation_flags);
+          fe_eval.apply_for_each_quad_point(m_quad_operation);
+          fe_eval.integrate(m_integration_flags);
+          Kokkos::single(Kokkos::PerTeam(shared_data->team_member),
+                         [&] { diagonal[i] = fe_eval.get_dof_value(i); });
+        }
+
+      Kokkos::single(Kokkos::PerTeam(shared_data->team_member), [&] {
+        for (unsigned int i = 0; i < dofs_per_cell; ++i)
+          fe_eval.submit_dof_value(diagonal[i], i);
+      });
+      fe_eval.distribute_local_to_global(dst);
+    };
+
+    static constexpr unsigned int n_local_dofs = QuadOperation::n_local_dofs;
+
+  private:
+    const QuadOperation                   &m_quad_operation;
+    const EvaluationFlags::EvaluationFlags m_evaluation_flags;
+    const EvaluationFlags::EvaluationFlags m_integration_flags;
+  };
+
+
+  template <int dim,
+            int fe_degree,
+            int n_q_points_1d,
+            int n_components,
+            typename Number,
+            typename MemorySpace,
+            typename QuadOperation>
+  void
+  compute_diagonal(
+    const Portable::MatrixFree<dim, Number>                 &matrix_free,
+    LinearAlgebra::distributed::Vector<Number, MemorySpace> &diagonal_global,
+    const QuadOperation                                     &quad_operation,
+    EvaluationFlags::EvaluationFlags                         evaluation_flags,
+    EvaluationFlags::EvaluationFlags                         integration_flags,
+    const unsigned int                                       dof_no,
+    const unsigned int                                       quad_no,
+    const unsigned int first_selected_component,
+    const unsigned int first_vector_component)
+  {
+    (void)dof_no;
+    (void)quad_no;
+    (void)first_selected_component;
+    (void)first_vector_component;
+    Assert(dof_no == 0, ExcNotImplemented());
+    Assert(quad_no == 0, ExcNotImplemented());
+    Assert(first_selected_component == 0, ExcNotImplemented());
+    Assert(first_vector_component == 0, ExcNotImplemented());
+
+    matrix_free.initialize_dof_vector(diagonal_global);
+
+
+    CellAction<dim, fe_degree, Number, QuadOperation> cell_action(
+      quad_operation, evaluation_flags, integration_flags);
+    // diagonal_global.zero_out_ghost_values();
+    LinearAlgebra::distributed::Vector<Number, MemorySpace> dummy;
+    matrix_free.cell_loop(cell_action, dummy, diagonal_global);
+
+    matrix_free.set_constrained_values(Number(1.), diagonal_global);
+  }
+
   template <int dim,
             int fe_degree,
             int n_q_points_1d,
diff --git a/tests/matrix_free_kokkos/compute_diagonal_01.cc b/tests/matrix_free_kokkos/compute_diagonal_01.cc
new file mode 100644
index 0000000000..2865a5069a
--- /dev/null
+++ b/tests/matrix_free_kokkos/compute_diagonal_01.cc
@@ -0,0 +1,90 @@
+// ------------------------------------------------------------------------
+//
+// SPDX-License-Identifier: LGPL-2.1-or-later
+// Copyright (C) 2020 - 2023 by the deal.II authors
+//
+// This file is part of the deal.II library.
+//
+// Part of the source code is dual licensed under Apache-2.0 WITH
+// LLVM-exception OR LGPL-2.1-or-later. Detailed license information
+// governing the source code and code contributions can be found in
+// LICENSE.md and CONTRIBUTING.md at the top level directory of deal.II.
+//
+// ------------------------------------------------------------------------
+
+
+// Test MatrixFreeTools::compute_diagonal() for a Laplace operator
+
+#include "compute_diagonal_util.h"
+
+template <int dim,
+          int fe_degree,
+          int n_points     = fe_degree + 1,
+          int n_components = 1,
+          typename Number  = double>
+void
+test()
+{
+  parallel::distributed::Triangulation<dim> tria(MPI_COMM_WORLD);
+
+  GridGenerator::hyper_cube(tria, -numbers::PI / 2, numbers::PI / 2);
+
+  tria.refine_global(2);
+  // FIXME
+  // for (auto &cell : tria.active_cell_iterators())
+  //  if (cell->is_active() && cell->is_locally_owned() &&
+  //      cell->center()[0] < 0.0)
+  // tria.begin_active()->set_refine_flag();
+  // tria.execute_coarsening_and_refinement();
+
+  const FE_Q<dim>     fe_q(fe_degree);
+  const FESystem<dim> fe(fe_q, n_components);
+
+  // setup dof-handlers
+  DoFHandler<dim> dof_handler(tria);
+  dof_handler.distribute_dofs(fe);
+
+  AffineConstraints<Number> constraint;
+  DoFTools::make_hanging_node_constraints(dof_handler, constraint);
+
+  VectorTools::interpolate_boundary_values(dof_handler,
+                                           0,
+                                           Functions::ZeroFunction<dim, Number>(
+                                             n_components),
+                                           constraint);
+  constraint.close();
+
+  // constraint.print(std::cout);
+
+  typename Portable::MatrixFree<dim, Number>::AdditionalData additional_data;
+  additional_data.mapping_update_flags = update_values | update_gradients;
+
+  MappingQ<dim> mapping(1);
+  QGauss<1>     quad(fe_degree + 1);
+
+  Portable::MatrixFree<dim, Number> matrix_free;
+  matrix_free.reinit(mapping, dof_handler, constraint, quad, additional_data);
+
+
+  Test<dim, fe_degree, n_points, n_components, Number> test(matrix_free,
+                                                            constraint);
+
+  deallog << "dim: " << dim << ", fe_degree: " << fe_degree
+          << ", n_points: " << n_points << ", n_components: " << n_components
+          << ", Number: "
+          << (std::is_same_v<Number, double> ? "double" : "float") << std::endl;
+
+  test.do_test();
+}
+
+int
+main(int argc, char **argv)
+{
+  Utilities::MPI::MPI_InitFinalize mpi_initialization(argc, argv, 1);
+  MPILogInitAll                    all;
+
+  test<2, 1, 2, 1>();        // scalar
+  test<2, 1, 2, 1, float>(); // scalar
+  test<3, 1, 2, 1>();        // scalar
+  test<3, 1, 2, 1, float>(); // scalar
+}
diff --git a/tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=1.output b/tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=1.output
new file mode 100644
index 0000000000..bdb17d9d19
--- /dev/null
+++ b/tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=1.output
@@ -0,0 +1,21 @@
+
+DEAL:0::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #0
+Local range: [0, 25), global size: 25
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 2.667e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 2.667e+00 1.000e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:0::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #0
+Local range: [0, 25), global size: 25
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 2.667e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 2.667e+00 1.000e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:0::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #0
+Local range: [0, 125), global size: 125
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:0::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #0
+Local range: [0, 125), global size: 125
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 
diff --git a/tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=3.output b/tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=3.output
new file mode 100644
index 0000000000..e5ef82bad5
--- /dev/null
+++ b/tests/matrix_free_kokkos/compute_diagonal_01.with_p4est=true.mpirun=3.output
@@ -0,0 +1,65 @@
+
+DEAL:0::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #0
+Local range: [0, 9), global size: 25
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 2.667e+00 
+DEAL:0::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #0
+Local range: [0, 9), global size: 25
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 1.000e+00 2.667e+00 2.667e+00 
+DEAL:0::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #0
+Local range: [0, 63), global size: 125
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:0::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #0
+Local range: [0, 63), global size: 125
+Vector data:
+1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 
+
+DEAL:1::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #1
+Local range: [9, 21), global size: 25
+Vector data:
+1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 2.667e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:1::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #1
+Local range: [9, 21), global size: 25
+Vector data:
+1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 1.000e+00 1.000e+00 2.667e+00 2.667e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:1::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #1
+Local range: [63, 93), global size: 125
+Vector data:
+1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:1::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #1
+Local range: [63, 93), global size: 125
+Vector data:
+1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 
+
+
+DEAL:2::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #2
+Local range: [21, 25), global size: 25
+Vector data:
+2.667e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:2::dim: 2, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #2
+Local range: [21, 25), global size: 25
+Vector data:
+2.667e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:2::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: double
+Process #2
+Local range: [93, 125), global size: 125
+Vector data:
+1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 
+DEAL:2::dim: 3, fe_degree: 1, n_points: 2, n_components: 1, Number: float
+Process #2
+Local range: [93, 125), global size: 125
+Vector data:
+1.000e+00 2.094e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 2.094e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 
+
diff --git a/tests/matrix_free_kokkos/compute_diagonal_util.h b/tests/matrix_free_kokkos/compute_diagonal_util.h
new file mode 100644
index 0000000000..ae71b3f1c4
--- /dev/null
+++ b/tests/matrix_free_kokkos/compute_diagonal_util.h
@@ -0,0 +1,167 @@
+// ------------------------------------------------------------------------
+//
+// SPDX-License-Identifier: LGPL-2.1-or-later
+// Copyright (C) 2020 - 2024 by the deal.II authors
+//
+// This file is part of the deal.II library.
+//
+// Part of the source code is dual licensed under Apache-2.0 WITH
+// LLVM-exception OR LGPL-2.1-or-later. Detailed license information
+// governing the source code and code contributions can be found in
+// LICENSE.md and CONTRIBUTING.md at the top level directory of deal.II.
+//
+// ------------------------------------------------------------------------
+
+#include <deal.II/base/logstream.h>
+#include <deal.II/base/mpi.h>
+#include <deal.II/base/quadrature_lib.h>
+
+#include <deal.II/distributed/tria.h>
+
+#include <deal.II/dofs/dof_handler.h>
+#include <deal.II/dofs/dof_tools.h>
+
+#include <deal.II/fe/fe_dgq.h>
+#include <deal.II/fe/fe_q.h>
+#include <deal.II/fe/fe_system.h>
+#include <deal.II/fe/fe_tools.h>
+#include <deal.II/fe/mapping_q.h>
+
+#include <deal.II/grid/grid_generator.h>
+#include <deal.II/grid/grid_out.h>
+
+#include <deal.II/lac/dynamic_sparsity_pattern.h>
+#include <deal.II/lac/sparse_matrix.h>
+
+#include <deal.II/matrix_free/portable_fe_evaluation.h>
+#include <deal.II/matrix_free/portable_matrix_free.h>
+#include <deal.II/matrix_free/tools.h>
+#include <deal.II/matrix_free/vector_access_internal.h>
+
+#include <deal.II/multigrid/mg_constrained_dofs.h>
+
+#include <deal.II/numerics/matrix_creator.h>
+#include <deal.II/numerics/vector_tools.h>
+
+#include "../tests.h"
+
+
+template <int dim, int fe_degree, typename Number>
+class LaplaceOperatorQuad
+{
+public:
+  DEAL_II_HOST_DEVICE void
+  operator()(
+    Portable::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, Number> *fe_eval,
+    const int q_point) const
+  {
+    fe_eval->submit_gradient(fe_eval->get_gradient(q_point), q_point);
+  }
+
+  static const unsigned int n_q_points =
+    dealii::Utilities::pow(fe_degree + 1, dim);
+
+  static const unsigned int n_local_dofs =
+    dealii::Utilities::pow(fe_degree + 1, dim);
+};
+
+
+template <int dim,
+          int fe_degree,
+          int n_points,
+          int n_components,
+          typename Number>
+class Test
+{
+  using VectorType =
+    LinearAlgebra::distributed::Vector<Number, MemorySpace::Default>;
+
+public:
+  Test(const Portable::MatrixFree<dim, Number> &matrix_free,
+       const AffineConstraints<Number>         &constraints)
+    : matrix_free(matrix_free)
+    , constraints(constraints)
+  {}
+
+  void
+  do_test()
+  {
+    // compute diagonal with the new function
+    LinearAlgebra::distributed::Vector<Number, MemorySpace::Default>
+      diagonal_global;
+    LinearAlgebra::distributed::Vector<Number, MemorySpace::Host>
+      diagonal_global_host;
+    LinearAlgebra::distributed::Vector<Number, MemorySpace::Host>
+      diagonal_global_reference;
+
+    SparseMatrix<Number> A_ref;
+    SparsityPattern      sparsity_pattern;
+
+    {
+      matrix_free.initialize_dof_vector(diagonal_global);
+      LaplaceOperatorQuad<dim, fe_degree, Number> laplace_operator_quad;
+      MatrixFreeTools::
+        compute_diagonal<dim, fe_degree, n_points, n_components, Number>(
+          matrix_free,
+          diagonal_global,
+          laplace_operator_quad,
+          EvaluationFlags::gradients,
+          EvaluationFlags::gradients);
+
+      matrix_free.initialize_dof_vector(diagonal_global_host);
+      LinearAlgebra::ReadWriteVector<Number> rw_vector(
+        diagonal_global.get_partitioner()->locally_owned_range());
+      rw_vector.import_elements(diagonal_global, VectorOperation::insert);
+      diagonal_global_host.import_elements(rw_vector, VectorOperation::insert);
+      for (unsigned int i = 0; i < diagonal_global.size(); ++i)
+        {
+          if (diagonal_global.get_partitioner()->in_local_range(i))
+            {
+              Assert(diagonal_global_host[i] > 0,
+                     ExcMessage("Diagonal non-positive at position " +
+                                std::to_string(i)));
+              // std::cout << i << ": " << diagonal_global_host[i] << '\n';
+            }
+        }
+
+      diagonal_global_host.print(deallog.get_file_stream());
+    }
+
+    const bool test_matrix =
+      (Utilities::MPI::job_supports_mpi() == false) ||
+      (Utilities::MPI::n_mpi_processes(MPI_COMM_WORLD) == 1);
+
+    if (test_matrix)
+      {
+        DynamicSparsityPattern dsp(matrix_free.get_dof_handler().n_dofs());
+        DoFTools::make_sparsity_pattern(matrix_free.get_dof_handler(),
+                                        dsp,
+                                        constraints);
+        sparsity_pattern.copy_from(dsp);
+        A_ref.reinit(sparsity_pattern);
+
+        Function<dim, Number> *scaling = nullptr;
+        MatrixCreator::create_laplace_matrix(matrix_free.get_dof_handler(),
+                                             QGauss<dim>(fe_degree + 1),
+                                             A_ref,
+                                             scaling,
+                                             constraints);
+
+        for (unsigned int i = 0; i < diagonal_global.size(); ++i)
+          {
+            if (!constraints.is_constrained(i))
+              {
+                if (std::abs(A_ref(i, i) - diagonal_global_host(i)) > 1e-6)
+                  // std::cout << i << ": " << A_ref(i,i) << " should be " <<
+                  // diagonal_global_host(i) << '\n';
+                  Assert(std::abs(A_ref(i, i) - diagonal_global_host(i)) < 1e-6,
+                         ExcMessage("Wrong diagonal entry at position " +
+                                    std::to_string(i)));
+              }
+          }
+      }
+  }
+
+  const Portable::MatrixFree<dim, Number> &matrix_free;
+  const AffineConstraints<Number>         &constraints;
+};
-- 
2.39.5