]> https://gitweb.dealii.org/ - dealii.git/commitdiff
Use a functor 18124/head
authorDaniel Arndt <arndtd@ornl.gov>
Wed, 5 Mar 2025 03:51:20 +0000 (03:51 +0000)
committerDaniel Arndt <arndtd@ornl.gov>
Wed, 5 Mar 2025 03:58:20 +0000 (03:58 +0000)
tests/matrix_free_kokkos/evaluate_2d_shape.cc

index e82a4fea4f4f3d821e7398fd606a48de8d95ba0b..93b39787194835a5a38eb6901a090bc01ee8e508 100644 (file)
 using TeamHandle = Kokkos::TeamPolicy<
   MemorySpace::Default::kokkos_space::execution_space>::member_type;
 
+// FIXME Inlining the functor creates "invalid device code" errors
 template <int M, int N, int type, bool add, bool dof_to_quad>
-DEAL_II_HOST_DEVICE void
-evaluate_tensor_product(
-  const TeamHandle                                          &team_member,
-  Kokkos::View<double *, MemorySpace::Default::kokkos_space> shape_values,
-  Kokkos::View<double *, MemorySpace::Default::kokkos_space> shape_gradients,
-  Kokkos::View<double *, MemorySpace::Default::kokkos_space> co_shape_gradients,
-  Kokkos::View<double *, MemorySpace::Default::kokkos_space> dst,
-  Kokkos::View<double *, MemorySpace::Default::kokkos_space> src,
-  Kokkos::View<double *, MemorySpace::Default::kokkos_space> tmp)
+struct EvaluateTensorProduct
 {
-  Kokkos::View<
-    double *,
-    MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
-    Kokkos::MemoryTraits<Kokkos::Unmanaged>>
-    dummy_scratch(team_member.team_shmem(), 0);
-
-  Portable::internal::EvaluatorTensorProduct<
-    Portable::internal::evaluate_general,
-    2,
-    M,
-    N,
-    double>
-    evaluator(team_member,
-              shape_values,
-              shape_gradients,
-              co_shape_gradients,
-              dummy_scratch);
-
-  constexpr int d0 = dof_to_quad ? 0 : 1;
-  constexpr int d1 = dof_to_quad ? 1 : 0;
-
-  if (type == 0)
-    {
-      evaluator.template values<d0, dof_to_quad, false, false>(src, tmp);
-      team_member.team_barrier();
-      evaluator.template values<d1, dof_to_quad, add, false>(tmp, dst);
-    }
-  if (type == 1)
-    {
-      evaluator.template gradients<d0, dof_to_quad, false, false>(src, tmp);
-      team_member.team_barrier();
-      evaluator.template gradients<d1, dof_to_quad, add, false>(tmp, dst);
-    }
-}
+  DEAL_II_HOST_DEVICE void
+  operator()(const TeamHandle &team_member) const
+  {
+    Kokkos::View<
+      double *,
+      MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
+      Kokkos::MemoryTraits<Kokkos::Unmanaged>>
+      dummy_scratch(team_member.team_shmem(), 0);
+
+    Portable::internal::EvaluatorTensorProduct<
+      Portable::internal::evaluate_general,
+      2,
+      M,
+      N,
+      double>
+      evaluator(team_member,
+                shape_values,
+                shape_gradients,
+                co_shape_gradients,
+                dummy_scratch);
+
+    constexpr int d0 = dof_to_quad ? 0 : 1;
+    constexpr int d1 = dof_to_quad ? 1 : 0;
+
+    if (type == 0)
+      {
+        evaluator.template values<d0, dof_to_quad, false, false>(src, tmp);
+        team_member.team_barrier();
+        evaluator.template values<d1, dof_to_quad, add, false>(tmp, dst);
+      }
+    if (type == 1)
+      {
+        evaluator.template gradients<d0, dof_to_quad, false, false>(src, tmp);
+        team_member.team_barrier();
+        evaluator.template gradients<d1, dof_to_quad, add, false>(tmp, dst);
+      }
+  }
+
+  Kokkos::View<double *, MemorySpace::Default::kokkos_space> shape_values;
+  Kokkos::View<double *, MemorySpace::Default::kokkos_space> shape_gradients;
+  Kokkos::View<double *, MemorySpace::Default::kokkos_space> co_shape_gradients;
+  Kokkos::View<double *, MemorySpace::Default::kokkos_space> dst;
+  Kokkos::View<double *, MemorySpace::Default::kokkos_space> src;
+  Kokkos::View<double *, MemorySpace::Default::kokkos_space> tmp;
+};
 
 template <int M, int N, int type, bool add>
 void
@@ -170,16 +174,9 @@ test()
   MemorySpace::Default::kokkos_space::execution_space exec;
   Kokkos::TeamPolicy<MemorySpace::Default::kokkos_space::execution_space>
     team_policy(exec, 1, Kokkos::AUTO);
-  Kokkos::parallel_for(
-    team_policy, KOKKOS_LAMBDA(const TeamHandle &team_member) {
-      evaluate_tensor_product<M, N, type, add, false>(team_member,
-                                                      shape_values,
-                                                      shape_gradients,
-                                                      co_shape_gradients,
-                                                      y_dev,
-                                                      x_dev,
-                                                      tmp_dev);
-    });
+  EvaluateTensorProduct<M, N, type, add, false> functor_to_dof{
+    shape_values, shape_gradients, co_shape_gradients, y_dev, x_dev, tmp_dev};
+  Kokkos::parallel_for(team_policy, functor_to_dof);
 
   // Check the results on the host
   Kokkos::deep_copy(y_host, y_dev);
@@ -208,16 +205,9 @@ test()
   Kokkos::deep_copy(x_dev, x_host);
 
   // Launch the kernel
-  Kokkos::parallel_for(
-    team_policy, KOKKOS_LAMBDA(const TeamHandle &team_member) {
-      evaluate_tensor_product<M, N, type, add, true>(team_member,
-                                                     shape_values,
-                                                     shape_gradients,
-                                                     co_shape_gradients,
-                                                     x_dev,
-                                                     y_dev,
-                                                     tmp_dev);
-    });
+  EvaluateTensorProduct<M, N, type, add, true> functor_to_quad{
+    shape_values, shape_gradients, co_shape_gradients, x_dev, y_dev, tmp_dev};
+  Kokkos::parallel_for(team_policy, functor_to_quad);
 
   // Check the results on the host
   Kokkos::deep_copy(x_host, x_dev);

In the beginning the Universe was created. This has made a lot of people very angry and has been widely regarded as a bad move.

Douglas Adams


Typeset in Trocchi and Trocchi Bold Sans Serif.