Refactor extended lambda and enable RAJA usage inside domain integral…

… kernels

src/serac/numerics/functional/domain_integral_kernels.hpp

@@ -10,7 +10,10 @@
 #include "serac/numerics/functional/function_signature.hpp"
 #include "serac/numerics/functional/differentiate_wrt.hpp"
 
+#include <RAJA/index/RangeSegment.hpp>
+#include <RAJA/RAJA.hpp>
 #include <array>
+#include <cstdint>
 
 namespace serac {
 
@@ -159,8 +162,14 @@ void evaluation_kernel_impl(FunctionSignature<test(trials...)>, const std::vecto
   [[maybe_unused]] tuple u = {
       reinterpret_cast<const typename decltype(type<indices>(trial_elements))::dof_type*>(inputs[indices])...};
 
+#if defined(RAJA_ENABLE_OPENMP)
+  using policy = RAJA::omp_parallel_for_exec;
+#else
+  using policy = RAJA::simd_exec;
+#endif
+
   // for each element in the domain
-  for (uint32_t e = 0; e < num_elements; e++) {
+  RAJA::forall<policy>(RAJA::TypedRangeSegment<uint32_t>(0, num_elements), [&](uint32_t e) {
     // load the jacobians and positions for each quadrature point in this element
     auto J_e = J[e];
     auto x_e = x[e];
@@ -201,7 +210,9 @@ void evaluation_kernel_impl(FunctionSignature<test(trials...)>, const std::vecto
 
     // (batch) integrate the material response against the test-space basis functions
     test_element::integrate(get_value(qf_outputs), rule, &r[e]);
-  }
+  });
+
+  return;
 }
 
 //clang-format off

src/serac/physics/heat_transfer.hpp

@@ -63,7 +63,7 @@ const TimesteppingOptions default_static_options = {TimestepMethod::QuasiStatic}
 /**
  * @brief An object containing the solver for a thermal conduction PDE
  *
- * This is a generic linear thermal diffusion oeprator of the form
+ * This is a generic linear thermal diffusion operator of the form
  *
  * \f[
  * \mathbf{M} \frac{\partial \mathbf{u}}{\partial t} = -\kappa \mathbf{K} \mathbf{u} + \mathbf{f}
@@ -276,6 +276,41 @@ class HeatTransfer<order, dim, Parameters<parameter_space...>, std::integer_sequ
     cycle_ += 1;
   }
 
+  template <typename MaterialType>
+  struct ThermalMaterialIntegrand {
+    ThermalMaterialIntegrand(MaterialType material) : material_(material) {}
+
+    template <typename X, typename T, typename dT_dt, typename Shape, typename... Params>
+    auto operator()(X x, T temperature, dT_dt dtemp_dt, Shape shape, Params... params)
+    {
+      // Get the value and the gradient from the input tuple
+      auto [u, du_dX] = temperature;
+      auto [p, dp_dX] = shape;
+      auto du_dt      = get<VALUE>(dtemp_dt);
+
+      auto I_plus_dp_dX     = I + dp_dX;
+      auto inv_I_plus_dp_dX = inv(I_plus_dp_dX);
+      auto det_I_plus_dp_dX = det(I_plus_dp_dX);
+
+      // Note that the current configuration x = X + p, where X is the original reference
+      // configuration and p is the shape displacement. We need the gradient with
+      // respect to the perturbed reference configuration x = X + p for the material model. Therefore, we calculate
+      // du/dx = du/dX * dX/dx = du/dX * (dx/dX)^-1 = du/dX * (I + dp/dX)^-1
+
+      auto du_dx = dot(du_dX, inv_I_plus_dp_dX);
+
+      auto [heat_capacity, heat_flux] = material_(x + p, u, du_dx, params...);
+
+      // Note that the return is integrated in the perturbed reference
+      // configuration, hence the det(I + dp_dx) = det(dx/dX)
+      return serac::tuple{heat_capacity * du_dt * det_I_plus_dp_dX,
+                          -1.0 * dot(inv_I_plus_dp_dX, heat_flux) * det_I_plus_dp_dX};
+    }
+
+  private:
+    MaterialType material_;
+  };
+
   /**
    * @brief Set the thermal material model for the physics solver
    *
@@ -300,33 +335,9 @@ class HeatTransfer<order, dim, Parameters<parameter_space...>, std::integer_sequ
   template <int... active_parameters, typename MaterialType>
   void setMaterial(DependsOn<active_parameters...>, MaterialType material)
   {
-    residual_->AddDomainIntegral(
-        Dimension<dim>{}, DependsOn<0, 1, 2, active_parameters + NUM_STATE_VARS...>{},
-        [material](auto x, auto temperature, auto dtemp_dt, auto shape, auto... params) {
-          // Get the value and the gradient from the input tuple
-          auto [u, du_dX] = temperature;
-          auto [p, dp_dX] = shape;
-          auto du_dt      = get<VALUE>(dtemp_dt);
-
-          auto I_plus_dp_dX     = I + dp_dX;
-          auto inv_I_plus_dp_dX = inv(I_plus_dp_dX);
-          auto det_I_plus_dp_dX = det(I_plus_dp_dX);
-
-          // Note that the current configuration x = X + p, where X is the original reference
-          // configuration and p is the shape displacement. We need the gradient with
-          // respect to the perturbed reference configuration x = X + p for the material model. Therefore, we calculate
-          // du/dx = du/dX * dX/dx = du/dX * (dx/dX)^-1 = du/dX * (I + dp/dX)^-1
-
-          auto du_dx = dot(du_dX, inv_I_plus_dp_dX);
-
-          auto [heat_capacity, heat_flux] = material(x + p, u, du_dx, params...);
-
-          // Note that the return is integrated in the perturbed reference
-          // configuration, hence the det(I + dp_dx) = det(dx/dX)
-          return serac::tuple{heat_capacity * du_dt * det_I_plus_dp_dX,
-                              -1.0 * dot(inv_I_plus_dp_dX, heat_flux) * det_I_plus_dp_dX};
-        },
-        mesh_);
+    ThermalMaterialIntegrand<MaterialType> integrand(material);
+    residual_->AddDomainIntegral(Dimension<dim>{}, DependsOn<0, 1, 2, active_parameters + NUM_STATE_VARS...>{},
+                                 integrand, mesh_);
   }
 
   /// @overload