Portable::MatrixFree: add dof_handler index

Author: tjhei

include/deal.II/matrix_free/portable_fe_evaluation.h

@@ -121,10 +121,16 @@ namespace Portable
       tensor_dofs_per_component * n_components;
 
     /**
-     * Constructor.
+     * Constructor. You will need to provide a pointer to the
+     * Portable::MatrixFree::Data object, which is typically provided to the
+     * functor inside the
+     * Portable::MatrixFree::cell_loop() and the index @p dof_index of the DoFHandler
+     * if more than one was provided when the Portable::MatrixFree object was
+     * initialized.
      */
     DEAL_II_HOST_DEVICE
-    explicit FEEvaluation(const data_type *data);
+    explicit FEEvaluation(const data_type   *data,
+                          const unsigned int dof_index = 0);
 
     /**
      * Return the index of the current cell.
@@ -255,12 +261,15 @@ namespace Portable
             typename Number>
   DEAL_II_HOST_DEVICE
   FEEvaluation<dim, fe_degree, n_q_points_1d, n_components_, Number>::
-    FEEvaluation(const data_type *data)
+    FEEvaluation(const data_type *data, const unsigned int dof_index)
     : data(data)
     , precomputed_data(data->precomputed_data)
     , shared_data(data->shared_data)
     , cell_id(data->team_member.league_rank())
-  {}
+  {
+    (void)dof_index;
+    AssertIndexRange(dof_index, data->n_dofhandler);
+  }
 
 
 

include/deal.II/matrix_free/portable_matrix_free.h

@@ -375,6 +375,7 @@ namespace Portable
        */
       TeamHandle team_member;
 
+      const unsigned int       n_dofhandler;
       const int                cell_index;
       const PrecomputedData   *precomputed_data;
       SharedData<dim, Number> *shared_data;

include/deal.II/matrix_free/portable_matrix_free.templates.h

@@ -418,8 +418,11 @@ namespace Portable
 
         SharedData<dim, Number> shared_data(values, gradients, scratch_pad);
 
+        const unsigned int cell_index = team_member.league_rank();
+
         typename MatrixFree<dim, Number>::Data data{team_member,
-                                                    team_member.league_rank(),
+                                                    /* n_dofhandler */ 1,
+                                                    cell_index,
                                                     &gpu_data,
                                                     &shared_data};
 
@@ -675,14 +678,16 @@ namespace Portable
               Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member,
                                                            n_q_points),
                                    [&](const int q_point) {
-                                     const int cell = team_member.league_rank();
+                                     const int cell_index =
+                                       team_member.league_rank();
 
                                      Data data{team_member,
-                                               cell,
+                                               /* n_dofhandler */ 1,
+                                               cell_index,
                                                &color_data,
-                                               /*shared_data*/ nullptr};
+                                               /* shared_data */ nullptr};
 
-                                     func(&data, cell, q_point);
+                                     func(&data, cell_index, q_point);
                                    });
             });
         }