portable matrix-free: fix GPU crash in compute_diagonal()

fixes dealii#18210

The Functor passed to Kokkos::parallel_for() is placed into constant memory, which is read-only. This means the HelmholtzOperatorQuad can not be modified by setting a member variable to the current cell index for example. Instead, add access functions for the necessary functions to FEEvaluation and query the information as needed.

Author: tjhei

doc/news/changes/minor/20250319tjhei

@@ -0,0 +1,5 @@
+New: Portable::FEEvaluation::get_current_cell_index() and
+Portable::FEEvaluation::get_matrix_free_data() allow querying
+the current Portable::MatrixFree::Data object and the current cell index respectively.
+<br>
+(Daniel Arndt, Timo Heister, 2025/03/19)

examples/step-64/step-64.cc

@@ -122,39 +122,23 @@ namespace Step64
   class HelmholtzOperatorQuad
   {
   public:
-    DEAL_II_HOST_DEVICE HelmholtzOperatorQuad(
-      const typename Portable::MatrixFree<dim, double>::Data *gpu_data,
-      double                                                 *coef,
-      int                                                     cell)
-      : gpu_data(gpu_data)
-      , coef(coef)
-      , cell(cell)
+    DEAL_II_HOST_DEVICE HelmholtzOperatorQuad(double *coef)
+      : coef(coef)
     {}
 
     DEAL_II_HOST_DEVICE void operator()(
       Portable::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, double> *fe_eval,
       const int q_point) const;
 
-    DEAL_II_HOST_DEVICE void set_matrix_free_data(
-      const typename Portable::MatrixFree<dim, double>::Data &data)
-    {
-      gpu_data = &data;
-    }
 
-    DEAL_II_HOST_DEVICE void set_cell(int new_cell)
-    {
-      cell = new_cell;
-    }
 
     static const unsigned int n_q_points =
       dealii::Utilities::pow(fe_degree + 1, dim);
 
     static const unsigned int n_local_dofs = n_q_points;
 
   private:
-    const typename Portable::MatrixFree<dim, double>::Data *gpu_data;
-    double                                                 *coef;
-    int                                                     cell;
+    double *coef;
   };
 
 
@@ -170,10 +154,17 @@ namespace Step64
     Portable::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, double> *fe_eval,
     const int q_point) const
   {
-    const unsigned int pos =
-      gpu_data->local_q_point_id(cell, n_q_points, q_point);
+    const int cell_index = fe_eval->get_current_cell_index();
+    const typename Portable::MatrixFree<dim, double>::Data *gpu_data =
+      fe_eval->get_matrix_free_data();
 
-    fe_eval->submit_value(coef[pos] * fe_eval->get_value(q_point), q_point);
+    const unsigned int position =
+      gpu_data->local_q_point_id(cell_index, n_q_points, q_point);
+    auto coeff = coef[position];
+
+    auto value = fe_eval->get_value(q_point);
+
+    fe_eval->submit_value(coeff * value, q_point);
     fe_eval->submit_gradient(fe_eval->get_gradient(q_point), q_point);
   }
 
@@ -218,7 +209,7 @@ namespace Step64
   // vector.
   template <int dim, int fe_degree>
   DEAL_II_HOST_DEVICE void LocalHelmholtzOperator<dim, fe_degree>::operator()(
-    const unsigned int                                      cell,
+    const unsigned int /*cell*/,
     const typename Portable::MatrixFree<dim, double>::Data *gpu_data,
     Portable::SharedData<dim, double>                      *shared_data,
     const double                                           *src,
@@ -229,7 +220,7 @@ namespace Step64
     fe_eval.read_dof_values(src);
     fe_eval.evaluate(EvaluationFlags::values | EvaluationFlags::gradients);
     fe_eval.apply_for_each_quad_point(
-      HelmholtzOperatorQuad<dim, fe_degree>(gpu_data, coef, cell));
+      HelmholtzOperatorQuad<dim, fe_degree>(coef));
     fe_eval.integrate(EvaluationFlags::values | EvaluationFlags::gradients);
     fe_eval.distribute_local_to_global(dst);
   }
@@ -363,7 +354,7 @@ namespace Step64
     initialize_dof_vector(inverse_diagonal);
 
     HelmholtzOperatorQuad<dim, fe_degree> helmholtz_operator_quad(
-      nullptr, coef.get_values(), -1);
+      coef.get_values());
 
     MatrixFreeTools::compute_diagonal<dim, fe_degree, fe_degree + 1, 1, double>(
       mf_data,

include/deal.II/matrix_free/portable_fe_evaluation.h

@@ -126,6 +126,22 @@ namespace Portable
     DEAL_II_HOST_DEVICE
     FEEvaluation(const data_type *data, SharedData<dim, Number> *shdata);
 
+    /**
+     * Return the index of the current cell.
+     */
+    DEAL_II_HOST_DEVICE
+    int
+    get_current_cell_index();
+
+    /**
+     * Return a pointer to the MatrixFree<dim, Number>::Data object on device
+     * that contains necessary constraint, dof index, and shape function
+     * information for evaluation used in the matrix-free kernels.
+     */
+    DEAL_II_HOST_DEVICE
+    const data_type *
+    get_matrix_free_data();
+
     /**
      * For the vector @p src, read out the values on the degrees of freedom of
      * the current cell, and store them internally. Similar functionality as
@@ -270,6 +286,38 @@ namespace Portable
 
 
 
+  template <int dim,
+            int fe_degree,
+            int n_q_points_1d,
+            int n_components_,
+            typename Number>
+  DEAL_II_HOST_DEVICE int
+  FEEvaluation<dim, fe_degree, n_q_points_1d, n_components_, Number>::
+    get_current_cell_index()
+  {
+    return cell_id;
+  }
+
+
+
+  template <int dim,
+            int fe_degree,
+            int n_q_points_1d,
+            int n_components_,
+            typename Number>
+  DEAL_II_HOST_DEVICE const typename FEEvaluation<dim,
+                                                  fe_degree,
+                                                  n_q_points_1d,
+                                                  n_components_,
+                                                  Number>::data_type *
+  FEEvaluation<dim, fe_degree, n_q_points_1d, n_components_, Number>::
+    get_matrix_free_data()
+  {
+    return data;
+  }
+
+
+
   template <int dim,
             int fe_degree,
             int n_q_points_1d,

include/deal.II/matrix_free/tools.h

@@ -1407,8 +1407,7 @@ namespace MatrixFreeTools
         Portable::
           FEEvaluation<dim, fe_degree, n_q_points_1d, n_components, Number>
             fe_eval(gpu_data, shared_data);
-        m_quad_operation.set_matrix_free_data(*gpu_data);
-        m_quad_operation.set_cell(cell);
+
         constexpr int dofs_per_cell = decltype(fe_eval)::tensor_dofs_per_cell;
         typename decltype(fe_eval)::value_type
           diagonal[dofs_per_cell / n_components] = {};
@@ -1503,7 +1502,7 @@ namespace MatrixFreeTools
         dealii::Utilities::pow(n_q_points_1d, dim);
 
     private:
-      mutable QuadOperation                  m_quad_operation;
+      const QuadOperation                    m_quad_operation;
       const EvaluationFlags::EvaluationFlags m_evaluation_flags;
       const EvaluationFlags::EvaluationFlags m_integration_flags;
     };

tests/matrix_free_kokkos/compute_diagonal_util.h

@@ -59,15 +59,6 @@ class LaplaceOperatorQuad
     fe_eval->submit_gradient(fe_eval->get_gradient(q_point), q_point);
   }
 
-  DEAL_II_HOST_DEVICE
-  void
-  set_cell(int)
-  {}
-
-  DEAL_II_HOST_DEVICE void
-  set_matrix_free_data(const typename Portable::MatrixFree<dim, Number>::Data &)
-  {}
-
   static const unsigned int n_q_points =
     dealii::Utilities::pow(fe_degree + 1, dim);