Cleanup

Serial DFEM diffusion tests

CMakeLists.txt

@@ -526,9 +526,11 @@ if (MFEM_USE_TRIBOL)
   endif()
 endif()
 
-# Enzyme
 if (MFEM_USE_ENZYME)
-   find_package(ENZYME REQUIRED)
+    find_package(Enzyme REQUIRED HINTS ${ENZYME_DIR})
+    message(STATUS "Enzyme found in ${ENZYME_DIR}.")
+    set(ENZYME_INCLUDE_DIRS ${ENZYME_DIR}/include)
+    set(ENZYME_FOUND 1)
 endif()
 
 # MFEM_TIMER_TYPE
@@ -680,6 +682,9 @@ if (MFEM_USE_MPI)
     target_link_libraries(mfem PUBLIC ${MPI_CXX_LINK_FLAGS})
   endif()
 endif()
+if (MFEM_USE_ENZYME)
+  target_link_libraries(mfem PUBLIC ClangEnzymeFlags)
+endif()
 
 set_target_properties(mfem PROPERTIES VERSION "${mfem_VERSION}")
 set_target_properties(mfem PROPERTIES SOVERSION "${mfem_VERSION}")

examples/CMakeLists.txt

@@ -249,3 +249,5 @@ endif()
 if(MFEM_USE_MOONOLITH)
   add_subdirectory(moonolith)
 endif()
+
+add_subdirectory(dfem)

examples/dfem/CMakeLists.txt

@@ -0,0 +1,116 @@
+# Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+# at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+# LICENSE and NOTICE for details. LLNL-CODE-806117.
+#
+# This file is part of the MFEM library. For more information and source code
+# availability visit https://mfem.org.
+#
+# MFEM is free software; you can redistribute it and/or modify it under the
+# terms of the BSD-3 license. We welcome feedback and contributions, see file
+# CONTRIBUTING.md for details.
+
+set(DFEM_EXAMPLES_SRCS)
+
+if (MFEM_USE_MPI)
+    list(APPEND DFEM_EXAMPLES_SRCS
+        plasticity.cpp
+        laghos.cpp
+    )
+endif()
+
+# Include the source directory where mfem.hpp and mfem-performance.hpp are.
+include_directories(BEFORE ${PROJECT_BINARY_DIR})
+
+# Add "test_dfem" target, see below.
+add_custom_target(test_dfem
+  ${CMAKE_CTEST_COMMAND} -R dfem USES_TERMINAL)
+
+# Add one executable per cpp file, adding "dfem_" as prefix so the CMake
+# target is unique from those in the non-dFEM examples. Also sets
+# "test_dfem" as a target that depends on the given dFEM examples.
+set(PFX dfem_)
+add_mfem_examples(DFEM_EXAMPLES_SRCS ${PFX} "" test_dfem)
+
+# Remove "dfem_" prefix from exectuable name for consistency with GNU build
+# system.
+foreach(SRC_FILE ${DFEM_EXAMPLES_SRCS})
+    get_filename_component(SRC_FILENAME ${SRC_FILE} NAME)
+    string(REPLACE ".cpp" "" TARGET_NAME "${PFX}${SRC_FILENAME}")
+    string(REPLACE ${PFX} "" EXE_NAME ${TARGET_NAME})
+    set_target_properties(${TARGET_NAME} PROPERTIES OUTPUT_NAME ${EXE_NAME})
+endforeach()
+
+# Testing.
+# The dFEM tests can be run separately using the target "test_dfem"
+# which builds the examples and runs:
+#   ctest -R dfem
+
+if (MFEM_ENABLE_TESTING)
+    # Command line options for the tests.
+    # Example 9: test CVODE with CV_ADAMS (non-stiff implicit) time stepping
+    # set(EX9_COMMON_OPTS -m ../../data/periodic-hexagon.mesh -p 0 -s 7)
+    # set(EX9_TEST_OPTS ${EX9_COMMON_OPTS}  -r  2 -dt 0.0018 -vs 25)
+    # set(EX9P_TEST_OPTS ${EX9_COMMON_OPTS} -rp 1 -dt 0.0009 -vs 50)
+    # Example 10: test CVODE with CV_BDF (stiff implicit) time stepping
+    # set(EX10_COMMON_OPTS -m ../../data/beam-quad.mesh -o 2 -s 5 -dt 0.15 -tf 6 -vs 10)
+    # set(EX10_TEST_OPTS ${EX10_COMMON_OPTS}  -r  2)
+    # set(EX10P_TEST_OPTS ${EX10_COMMON_OPTS} -rp 1)
+    # Example 16: test ARKODE with implicit time stepping using mass form
+    # set(EX16_COMMON_OPTS -s 15)
+    # set(EX16_TEST_OPTS ${EX16_COMMON_OPTS})
+    # set(EX16P_TEST_OPTS ${EX16_COMMON_OPTS})
+
+    # Add the tests: one test per source file.
+    foreach(SRC_FILE ${DFEM_EXAMPLES_SRCS})
+        get_filename_component(SRC_FILENAME ${SRC_FILE} NAME)
+        string(REPLACE ".cpp" "" TEST_NAME ${SRC_FILENAME})
+        string(TOUPPER ${TEST_NAME} UP_TEST_NAME)
+        set(TEST_NAME ${PFX}${TEST_NAME})
+
+        set(THIS_TEST_OPTIONS "-no-vis")
+        list(APPEND THIS_TEST_OPTIONS ${${UP_TEST_NAME}_TEST_OPTS})
+        # message(STATUS "Test ${TEST_NAME} options: ${THIS_TEST_OPTIONS}")
+
+        if (NOT (${TEST_NAME} MATCHES ".*p$"))
+            add_test(NAME ${TEST_NAME}_ser
+        COMMAND ${TEST_NAME} ${THIS_TEST_OPTIONS})
+        else()
+            add_test(NAME ${TEST_NAME}_np=${MFEM_MPI_NP}
+        COMMAND ${MPIEXEC} ${MPIEXEC_NUMPROC_FLAG} ${MFEM_MPI_NP}
+        ${MPIEXEC_PREFLAGS}
+        $<TARGET_FILE:${TEST_NAME}> ${THIS_TEST_OPTIONS}
+        ${MPIEXEC_POSTFLAGS})
+        endif()
+    endforeach()
+
+    # Add CUDA/HIP tests.
+    set(DEVICE_EXAMPLES
+      # parallel examples with device support:
+      # ex9p
+    )
+    set(MFEM_TEST_DEVICE)
+    if (MFEM_USE_CUDA)
+        set(MFEM_TEST_DEVICE "cuda")
+    elseif (MFEM_USE_HIP)
+        set(MFEM_TEST_DEVICE "hip")
+    endif()
+    if (MFEM_TEST_DEVICE)
+        foreach(TEST_NAME ${DEVICE_EXAMPLES})
+            string(TOUPPER ${TEST_NAME} UP_TEST_NAME)
+
+            set(THIS_TEST_OPTIONS "-no-vis" "-d" "${MFEM_TEST_DEVICE}")
+            list(APPEND THIS_TEST_OPTIONS ${${UP_TEST_NAME}_TEST_OPTS})
+
+            if (NOT (${TEST_NAME} MATCHES ".*p$"))
+                add_test(NAME ${PFX}${TEST_NAME}_${MFEM_TEST_DEVICE}_ser
+          COMMAND ${PFX}${TEST_NAME} ${THIS_TEST_OPTIONS})
+            else()
+                add_test(NAME ${PFX}${TEST_NAME}_${MFEM_TEST_DEVICE}_np=${MFEM_MPI_NP}
+          COMMAND ${MPIEXEC} ${MPIEXEC_NUMPROC_FLAG} ${MFEM_MPI_NP}
+          ${MPIEXEC_PREFLAGS}
+          $<TARGET_FILE:${PFX}${TEST_NAME}> ${THIS_TEST_OPTIONS}
+          ${MPIEXEC_POSTFLAGS})
+            endif()
+        endforeach()
+    endif(MFEM_TEST_DEVICE)
+endif(MFEM_ENABLE_TESTING)

examples/dfem/plasticity.cpp

@@ -0,0 +1,587 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#include <mfem.hpp>
+
+// TODO: Do we want this to be included from mfem.hpp automatically now?
+#include <fem/dfem/doperator.hpp>
+#include <linalg/tensor.hpp>
+
+#include <fstream>
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+template <typename T, int dim>
+MFEM_HOST_DEVICE inline
+tensor<T, 3, 3> tensor_to_3D(const tensor<T, dim, dim>& A)
+{
+   tensor<T, 3, 3> A3D{};
+   for (int i = 0; i < dim; i++)
+   {
+      for (int j = 0; j < dim; j++)
+      {
+         A3D[i][j] = A[i][j];
+      }
+   }
+   return A3D;
+}
+
+template <typename Material, int dim = DIMENSION>
+struct InternalStateQFunction
+{
+   InternalStateQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const tensor<real_t, 10> &internal_state,
+      const double &w) const
+   {
+      auto invJ = inv(J);
+      auto dudX = dudxi * invJ;
+      auto dudX3D = tensor_to_3D(dudX);
+      //auto internal_state_new = get<1>(material(dudX3D, internal_state));
+      auto [stress, internal_state_new] = material(dudX3D, internal_state);
+      // real_t vm = sqrt(1.5)*norm(dev(stress));
+      // out << vm << " " << internal_state_new[9] << std::endl;
+      return mfem::tuple{internal_state_new};
+   }
+
+   Material material;
+};
+
+template <typename Material, int dim = DIMENSION>
+struct MomentumRefStateQFunction
+{
+   MomentumRefStateQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const tensor<real_t, 10> &internal_state,
+      const double &w) const
+   {
+      auto invJ = inv(J);
+      auto dudX = dudxi * invJ;
+      auto dudX3D = tensor_to_3D(dudX);
+      auto [P3D, Qnew] = material(dudX3D, internal_state);
+      auto P = mfem::internal::make_tensor<dim, dim>([&P3D](int i, int j) { return P3D[i][j]; });
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{P * JxW};
+   }
+
+   Material material;
+};
+
+
+struct J2SmallStrain
+{
+   static constexpr int dim = 3;         ///< spatial dimension
+   static constexpr int n_internal_states = 10;
+   static constexpr double tol =
+      1e-10;  ///< relative tolerance on residual mag to judge convergence of return map
+
+   real_t E;                 ///< Young's modulus
+   real_t nu;                ///< Poisson's ratio
+   real_t sigma_y;           ///< Yield strength
+   real_t Hi;                ///< Isotropic hardening modulus
+   real_t density;           ///< Mass density
+
+   /// @brief variables required to characterize the hysteresis response
+   struct InternalState
+   {
+      tensor<double, dim, dim> plastic_strain;  ///< plastic strain
+      double accumulated_plastic_strain;        ///< uniaxial equivalent plastic strain
+   };
+
+   MFEM_HOST_DEVICE inline
+   InternalState unpack_internal_state(const tensor<real_t, n_internal_states> &
+                                       packed_state) const
+   {
+      // we could use type punning here to avoid copies
+      auto plastic_strain = mfem::internal::make_tensor<dim, dim>(
+      [&packed_state](int i, int j) { return packed_state[dim*i + j]; });
+      real_t accumulated_plastic_strain = packed_state[n_internal_states - 1];
+      return {plastic_strain, accumulated_plastic_strain};
+   }
+
+   MFEM_HOST_DEVICE inline
+   tensor<real_t, n_internal_states> pack_internal_state(const
+                                                         tensor<real_t, dim, dim> & plastic_strain,
+                                                         real_t accumulated_plastic_strain) const
+   {
+      tensor<real_t, n_internal_states> packed_state{};
+      for (int i = 0, ij = 0; i < dim; i++)
+      {
+         for (int j = 0; j < dim; j++, ij++)
+         {
+            packed_state[ij] = plastic_strain[i][j];
+         }
+      }
+      packed_state[n_internal_states - 1] = accumulated_plastic_strain;
+      return packed_state;
+   }
+
+   MFEM_HOST_DEVICE inline
+   tuple<tensor<real_t, dim, dim>, tensor<real_t, n_internal_states>>
+                                                                   operator()(const tensor<real_t, dim, dim> & dudX,
+                                                                              const tensor<real_t, n_internal_states> & internal_state) const
+   {
+      auto I = mfem::internal::Identity<dim>();
+      const real_t K = E / (3.0 * (1.0 - 2.0 * nu));
+      const real_t G = 0.5 * E / (1.0 + nu);
+
+      auto [plastic_strain, accumulated_plastic_strain] = unpack_internal_state(
+                                                             internal_state);
+
+      // (i) elastic predictor
+      auto el_strain = sym(dudX) - plastic_strain;
+      auto p = K * tr(el_strain);
+      auto s = 2.0 * G * dev(el_strain);
+      auto q = sqrt(1.5) * norm(s);
+      [[maybe_unused]] real_t delta_eqps = 0.0;
+
+      [[maybe_unused]] auto flow_strength = [this](real_t eqps) { return this->sigma_y + this->Hi*eqps; };
+
+      // (ii) admissibility
+      if (q - (sigma_y + Hi*accumulated_plastic_strain) > tol*sigma_y)
+      {
+         // (iii) return mapping
+         real_t delta_eqps = (q - sigma_y - Hi*accumulated_plastic_strain)/(3*G + Hi);
+         auto Np = 1.5 * s / q;
+         s -= 2.0 * G * delta_eqps * Np;
+         plastic_strain += delta_eqps * Np;
+         accumulated_plastic_strain += delta_eqps;
+      }
+      auto stress = s + p * I;
+      auto internal_state_new = pack_internal_state(plastic_strain,
+                                                    accumulated_plastic_strain);
+      return {stress, internal_state_new};
+   }
+};
+
+class ElasticityOperator : public Operator
+{
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+   static constexpr int InternalState = 2;
+
+public:
+   class ElasticityJacobianOperator : public Operator
+   {
+   public:
+      ElasticityJacobianOperator(const ElasticityOperator *elasticity,
+                                 const Vector &x) :
+         Operator(elasticity->Height()),
+         elasticity(elasticity),
+         z(elasticity->Height())
+      {
+         ParGridFunction u(&elasticity->displacement_fes);
+         u.SetFromTrueDofs(x);
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (elasticity->displacement_fes.GetParMesh()->GetNodes());
+         momentum_du = elasticity->momentum->GetDerivative(Displacement, {&u}, {mesh_nodes, &elasticity->internal_state});
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         z = x;
+         z.SetSubVector(elasticity->displacement_ess_tdof, 0.0);
+
+         momentum_du->Mult(z, y);
+
+         for (int i = 0; i < elasticity->displacement_ess_tdof.Size(); i++)
+         {
+            y[elasticity->displacement_ess_tdof[i]] =
+               x[elasticity->displacement_ess_tdof[i]];
+         }
+      }
+
+      const ElasticityOperator *elasticity;
+      std::shared_ptr<DerivativeOperator> momentum_du;
+      mutable Vector z;
+   };
+
+   template <typename Material>
+   ElasticityOperator(ParFiniteElementSpace &displacement_fes,
+                      Array<int> &vel_ess_tdofs,
+                      const IntegrationRule &displacement_ir,
+                      ParametricFunction &internal_state,
+                      Material material) :
+      Operator(displacement_fes.GetTrueVSize()),
+      density(1.0e3),
+      body_force(displacement_fes.GetTrueVSize()),
+      displacement_ess_tdof(vel_ess_tdofs),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      internal_state(internal_state)
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Displacement, &displacement_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Coordinates, &mesh_fes},
+            FieldDescriptor{InternalState, &internal_state.space}
+         };
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+         momentum->DisableTensorProductStructure();
+
+         mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, None<InternalState>{}, Weight{}};
+         mfem::tuple outputs{Gradient<Displacement>{}};
+
+         auto momentum_qf = MomentumRefStateQFunction<Material, DIMENSION> {.material = material};
+         auto derivatives = std::integer_sequence<size_t, Displacement> {};
+         Array<int> solid_domain_attr(mesh->attributes.Max());
+         solid_domain_attr[0] = 1;
+         momentum->AddDomainIntegrator(
+            momentum_qf, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+      }
+
+      {
+         Vector g(DIMENSION);
+         g = 0.0;
+
+         ParLinearForm body_force_lf(&displacement_fes);
+         body_force_coef = new VectorConstantCoefficient(g);
+         auto integ = new VectorDomainLFIntegrator(*body_force_coef);
+         integ->SetIntRule(&displacement_ir);
+         body_force_lf.AddDomainIntegrator(integ);
+         body_force_lf.Assemble();
+         body_force_lf.ParallelAssemble(body_force);
+      }
+   }
+
+   void Mult(const Vector &displacement, Vector &r) const override
+   {
+      momentum->SetParameters({mesh_nodes, &internal_state});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.SetSubVector(displacement_ess_tdof, 0.0);
+   }
+
+   void Reaction(const Vector &displacement, Vector &r) const
+   {
+      momentum->SetParameters({mesh_nodes, &internal_state});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.Neg();
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      jacobian_operator = std::make_shared<ElasticityJacobianOperator>(this, x);
+      return *jacobian_operator;
+
+      // fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      // return *fd_jacobian;
+   }
+
+   real_t density;
+   std::shared_ptr<DifferentiableOperator> momentum;
+   mutable std::shared_ptr<HypreParMatrix> A;
+   VectorConstantCoefficient *body_force_coef = nullptr;
+   Vector body_force;
+
+   ParGridFunction *mesh_nodes;
+
+   const Array<int> displacement_ess_tdof;
+
+   ParFiniteElementSpace &displacement_fes;
+   IntegrationRule displacement_ir;
+
+   ParametricFunction& internal_state;
+
+   mutable std::shared_ptr<ElasticityJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+};
+
+
+class InternalStateUpdater : public Operator
+{
+public:
+
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+   static constexpr int InternalState = 2;
+
+   template <typename Material>
+   InternalStateUpdater(ParFiniteElementSpace &displacement_fes,
+                        const IntegrationRule &displacement_ir,
+                        ParametricFunction &internal_state,
+                        Material material) :
+      Operator(displacement_fes.GetTrueVSize()),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      internal_state(internal_state)
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      auto solutions = std::vector
+      {
+         FieldDescriptor{Displacement, &displacement_fes}
+      };
+
+      auto parameters = std::vector
+      {
+         FieldDescriptor{Coordinates, &mesh_fes},
+         FieldDescriptor{InternalState, &internal_state.space}
+      };
+
+      op = std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+      op->DisableTensorProductStructure();
+
+      mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, None<InternalState>{}, Weight{}};
+      mfem::tuple outputs{None<InternalState>{}};
+
+      auto qfunction = InternalStateQFunction<Material, DIMENSION> {.material = material};
+      // just a placeholder for now. We want vjps wrt both displacement and old internal state eventually
+      auto derivatives = std::integer_sequence<size_t, Displacement> {};
+      Array<int> solid_domain_attr(mesh->attributes.Max());
+      solid_domain_attr[0] = 1;
+      op->AddDomainIntegrator(
+         qfunction, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+   }
+
+   void Mult(const Vector &displacement, Vector& internal_state_new) const override
+   {
+      op->SetParameters({mesh_nodes, &internal_state});
+      op->Mult(displacement, internal_state_new);
+   }
+
+   void VjpDisplacement(ParGridFunction &u, Vector& internal_state_old,
+                        Vector& internal_state_new_bar, Vector& displacement_bar) const
+   {
+      // u, internal_state_old, internal_state_new_bar should be const
+      out << "Sizes " << "u " << u.Size() << ", qold " << internal_state_old.Size() <<
+          ", qbar " << internal_state_new_bar.Size() << ", ubar " <<
+          displacement_bar.Size() << std::endl;
+      auto grad_op = op->GetDerivative(Displacement, {&u}, {mesh_nodes, &internal_state_old});
+      out << "grad_op " << grad_op->Height() << " " << grad_op->Width() << std::endl;
+      out << "grad_op^T " << grad_op->Width() << " " << grad_op->Height() <<
+          std::endl;
+      grad_op->MultTranspose(internal_state_new_bar, displacement_bar);
+   }
+
+   ParGridFunction *mesh_nodes;
+   ParFiniteElementSpace &displacement_fes;
+   std::shared_ptr<DifferentiableOperator> op;
+   IntegrationRule displacement_ir;
+   ParametricFunction& internal_state;
+};
+
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   int polynomial_order = 1;
+   int ir_order = 2;
+   int refinements = 0;
+   int nonlinear_solver_type = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--refinements", "");
+   args.AddOption(&ir_order, "-iro", "--integration-rule-order", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&nonlinear_solver_type, "-nls", "--nonlinear-solver", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL,
+                                            false, 1.0, 0.1);
+   mesh_serial.EnsureNodes();
+   auto mesh_beam = ParMesh(MPI_COMM_WORLD, mesh_serial);
+
+   out << "#el: " << mesh_beam.GetNE() << "\n";
+
+   H1_FECollection displacement_fec(polynomial_order, dim);
+   ParFiniteElementSpace displacement_fes(&mesh_beam, &displacement_fec, dim);
+
+   HYPRE_BigInt global_size = displacement_fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of unknowns: " << global_size << "\n";
+   }
+
+   const IntegrationRule &displacement_ir =
+      IntRules.Get(displacement_fes.GetFE(0)->GetGeomType(),
+                   2 * ir_order + displacement_fes.GetFE(0)->GetOrder());
+
+   constexpr int n_internal_state_variables = 10;
+   ParametricSpace internal_state_space(dim, n_internal_state_variables,
+                                        displacement_ir.GetNPoints(),
+                                        n_internal_state_variables*displacement_ir.GetNPoints()*mesh_beam.GetNE());
+
+   ParametricFunction internal_state(internal_state_space);
+   internal_state = 0.0;
+   ParametricFunction internal_state_old(internal_state_space);
+   internal_state_old = 0.0;
+
+   Array<int> bdr_attr_is_ess(mesh_beam.bdr_attributes.Max());
+   Array<int> displacement_ess_tdof;
+   Array<int> bc_tdof;
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[0] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 1);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[3] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[1] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   ParGridFunction u(&displacement_fes);
+   u = 0.0;
+
+   using Material = J2SmallStrain; // StVenantKirchhoff
+   Material material{.E = 1000.0, .nu = 0.25, .sigma_y = 0.53333, .Hi = 40.0, .density = 1.0};
+   // Material material{.mu = 0.5e6, .nu = 0.4};
+
+   ElasticityOperator elasticity(displacement_fes, displacement_ess_tdof,
+                                 displacement_ir, internal_state, material);
+
+   CGSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-10);
+   solver.SetMaxIter(1000);
+   solver.SetPrintLevel(2);
+
+   std::shared_ptr<NewtonSolver> nonlinear_solver;
+   if (nonlinear_solver_type == 0)
+   {
+      nonlinear_solver = std::make_shared<NewtonSolver>(MPI_COMM_WORLD);
+   }
+   // else if (nonlinear_solver_type == 1)
+   // {
+   //    nonlinear_solver = std::make_shared<KINSolver>(MPI_COMM_WORLD, KIN_LINESEARCH);
+   // }
+   else
+   {
+      MFEM_ABORT("invalid nonlinear solver type");
+   }
+   nonlinear_solver->SetOperator(elasticity);
+   nonlinear_solver->SetRelTol(1e-9);
+   nonlinear_solver->SetMaxIter(25);
+   nonlinear_solver->SetSolver(solver);
+   nonlinear_solver->SetPrintLevel(1);
+
+   // variables for output
+   QuadratureSpace output_internal_state_space(mesh_beam, displacement_ir);
+   QuadratureFunction output_internal_state(&output_internal_state_space,
+                                            internal_state.GetData(), material.n_internal_states);
+   Vector r(displacement_fes.GetTrueVSize());
+   ParGridFunction reaction(&displacement_fes);
+   Vector end_forces_x(bc_tdof.Size());
+
+   ParaViewDataCollection dc("dfem_plasticity", &mesh_beam);
+   dc.SetHighOrderOutput(true);
+   dc.SetLevelsOfDetail(1);
+   dc.RegisterField("displacement", &u);
+   dc.RegisterField("reaction", &reaction);
+   dc.RegisterQField("internal_state", &output_internal_state);
+   dc.SetCycle(0);
+   dc.Save();
+
+   InternalStateUpdater internal_state_update(displacement_fes, displacement_ir,
+                                              internal_state, material);
+   //Vector q(internal_state_space.GetTotalSize());
+
+   auto applied_displacement = [](double t) { return 1.2e-2*t; };
+
+   real_t time = 0.0;
+   std::ofstream history_file("history_output.csv");
+   history_file << applied_displacement(time) << " " << 0.0 << std::endl;
+
+   Vector zero, x(displacement_fes.GetTrueVSize());
+
+   constexpr int max_cycles = 3;
+   const real_t dt = 1.0/(max_cycles - 1);
+   for (int cycle = 1; cycle < max_cycles; cycle++)
+   {
+      time += dt;
+      out << "-------------------------------------------" << std::endl;
+      out << "TIME STEP " << cycle << std::endl;
+      out << "t = " << time << std::endl;
+
+      real_t ubc = applied_displacement(time);
+      u.SetSubVector(bc_tdof, ubc);
+
+      u.GetTrueDofs(x);
+      nonlinear_solver->Mult(zero, x);
+      u.SetFromTrueDofs(x);
+
+      // update internal variables
+      internal_state_old.Set(1.0, internal_state);
+      internal_state_update.Mult(u, internal_state);
+
+      // Compute reactions
+      elasticity.Reaction(x, r);
+      reaction.SetFromTrueDofs(r);
+      reaction.GetSubVector(bc_tdof, end_forces_x);
+      real_t force = -end_forces_x.Sum();
+      out << "u = " << applied_displacement(time) << ", Force = " << force <<
+          std::endl;
+      history_file << applied_displacement(time) << " " << force << std::endl;
+
+      output_internal_state = internal_state;
+
+      dc.SetCycle(cycle);
+      dc.SetTime(time);
+      dc.Save();
+   }
+
+   // try to use the derivative to see if it works
+   ParametricFunction internal_state_bar(internal_state_space);
+   internal_state_bar = 1.0;
+   //ParGridFunction u_bar(displacement_fes);
+   Vector u_bar(displacement_fes.GetTrueVSize());
+   internal_state_update.VjpDisplacement(u, internal_state_old, internal_state_bar,
+                                         u_bar);
+
+   pretty_print(u_bar);
+
+   history_file.close();
+   return 0;
+}

fem/CMakeLists.txt

@@ -175,6 +175,13 @@ set(HDRS
   dgmassinv.hpp
   dgmassinv_kernels.hpp
   doftrans.hpp
+  dfem/doperator.hpp
+  dfem/fieldoperator.hpp
+  dfem/integrate.hpp
+  dfem/parametricspace.hpp
+  dfem/qfunction.hpp
+  dfem/tuple.hpp
+  dfem/util.hpp
   eltrans.hpp
   estimators.hpp
   fe.hpp

fem/dfem/doperator.hpp

@@ -0,0 +1,776 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+#include <type_traits>
+#include <utility>
+
+#include "util.hpp"
+#include "interpolate.hpp"
+#include "qfunction.hpp"
+#include "integrate.hpp"
+
+#undef NVTX_COLOR
+#define NVTX_COLOR nvtx::kOrchid
+#include "general/nvtx.hpp"
+
+namespace mfem
+{
+
+using action_t =
+   std::function<void(std::vector<Vector> &, const std::vector<Vector> &, Vector &)>;
+
+using derivative_action_t =
+   std::function<void(std::vector<Vector> &, const Vector &, Vector &)>;
+
+using restriction_callback_t =
+   std::function<void(std::vector<Vector> &,
+                      const std::vector<Vector> &,
+                      std::vector<Vector> &)>;
+
+class DerivativeOperator : public Operator
+{
+public:
+   DerivativeOperator(
+      const int &height,
+      const int &width,
+      const std::vector<derivative_action_t> &derivative_actions,
+      const FieldDescriptor &direction,
+      const int &daction_l_size,
+      const std::vector<derivative_action_t> &derivative_actions_transpose,
+      const FieldDescriptor &transpose_direction,
+      const int &daction_transpose_l_size,
+      const std::vector<Vector *> &solutions_l,
+      const std::vector<Vector *> &parameters_l,
+      const restriction_callback_t &restriction_callback,
+      const std::function<void(Vector &, Vector &)> &prolongation_transpose) :
+      Operator(height, width),
+      derivative_actions(derivative_actions),
+      direction(direction),
+      daction_l(daction_l_size),
+      derivative_actions_transpose(derivative_actions_transpose),
+      transpose_direction(transpose_direction),
+      daction_transpose_l(daction_transpose_l_size),
+      prolongation_transpose(prolongation_transpose)
+   {
+      std::vector<Vector> s_l(solutions_l.size());
+      for (size_t i = 0; i < s_l.size(); i++)
+      {
+         s_l[i] = *solutions_l[i];
+      }
+
+      std::vector<Vector> p_l(parameters_l.size());
+      for (size_t i = 0; i < p_l.size(); i++)
+      {
+         p_l[i] = *parameters_l[i];
+      }
+
+      fields_e.resize(solutions_l.size() + parameters_l.size());
+      restriction_callback(s_l, p_l, fields_e);
+   }
+
+   void Mult(const Vector &direction_t, Vector &y) const override
+   {
+      // daction_l.SetSize(height);
+      daction_l = 0.0;
+
+      prolongation(direction, direction_t, direction_l);
+      for (size_t i = 0; i < derivative_actions.size(); i++)
+      {
+         derivative_actions[i](fields_e, direction_l, daction_l);
+      }
+      prolongation_transpose(daction_l, y);
+   };
+
+   void MultTranspose(const Vector &direction_t, Vector &y) const override
+   {
+      // daction_l.SetSize(width);
+      daction_l = 0.0;
+
+      prolongation(transpose_direction, direction_t, direction_l);
+      for (size_t i = 0; i < derivative_actions_transpose.size(); i++)
+      {
+         derivative_actions_transpose[i](fields_e, direction_l, daction_l);
+      }
+      prolongation_transpose(daction_l, y);
+   };
+
+private:
+   std::vector<derivative_action_t> derivative_actions;
+   FieldDescriptor direction;
+   mutable Vector daction_l;
+
+   std::vector<derivative_action_t> derivative_actions_transpose;
+   FieldDescriptor transpose_direction;
+   mutable Vector daction_transpose_l;
+
+   mutable std::vector<Vector> fields_e;
+
+   mutable Vector direction_l;
+
+   std::function<void(Vector &, Vector &)> prolongation_transpose;
+};
+
+class DifferentiableOperator : public Operator
+{
+public:
+   DifferentiableOperator(
+      const std::vector<FieldDescriptor> &solutions,
+      const std::vector<FieldDescriptor> &parameters,
+      const ParMesh &mesh);
+
+   void Mult(const Vector &solutions_t, Vector &y) const override
+   {
+      MFEM_ASSERT(!action_callbacks.empty(), "no integrators have been set");
+      prolongation(solutions, solutions_t, solutions_l);
+      for (auto &action : action_callbacks)
+      {
+         action(solutions_l, parameters_l, residual_l);
+      }
+      prolongation_transpose(residual_l, y);
+   }
+
+   void Mult(ParGridFunction &x, ParGridFunction &y) const
+   {
+      MFEM_ASSERT(!action_callbacks.empty(), "no integrators have been set");
+      MFEM_VERIFY(y.Size() == residual_l.Size(), "output size mismatch");
+      prolongation(solutions, x.GetTrueVector(), solutions_l);
+      for (auto &action : action_callbacks)
+      {
+         action(solutions_l, parameters_l, residual_l);
+      }
+      y = residual_l;
+   }
+
+   template <
+      typename func_t,
+      typename... input_ts,
+      typename... output_ts,
+      typename derivative_indices_t>
+   void AddDomainIntegrator(
+      func_t &qfunc,
+      mfem::tuple<input_ts...> inputs,
+      mfem::tuple<output_ts...> outputs,
+      const IntegrationRule &integration_rule,
+      const Array<int> domain_attributes,
+      const derivative_indices_t derivative_indices = {});
+
+   void SetParameters(std::vector<Vector *> p) const;
+
+   void DisableTensorProductStructure(bool disable = true)
+   {
+      use_tensor_product_structure = !disable;
+   }
+
+   std::shared_ptr<DerivativeOperator> GetDerivative(
+      size_t derivative_id,
+      std::vector<Vector *> solutions_l,
+      std::vector<Vector *> parameters_l)
+   {
+      MFEM_ASSERT(derivative_action_callbacks.find(derivative_id) !=
+                  derivative_action_callbacks.end(),
+                  "no derivative action has been found for ID " << derivative_id);
+
+      MFEM_ASSERT(solutions_l.size() == solutions.size(),
+                  "wrong number of solutions");
+
+      MFEM_ASSERT(parameters_l.size() == parameters.size(),
+                  "wrong number of parameters");
+
+      const size_t derivative_idx = FindIdx(derivative_id, fields);
+
+      return std::make_shared<DerivativeOperator>(
+                height,
+                GetTrueVSize(fields[derivative_idx]),
+                derivative_action_callbacks[derivative_id],
+                fields[derivative_idx],
+                residual_l.Size(),
+                daction_transpose_callbacks[derivative_id],
+                fields[test_space_field_idx],
+                GetTrueVSize(fields[test_space_field_idx]),
+                solutions_l,
+                parameters_l,
+                restriction_callback,
+                prolongation_transpose);
+   }
+
+private:
+   const ParMesh &mesh;
+
+   std::vector<action_t> action_callbacks;
+   std::map<size_t,
+       std::vector<derivative_action_t>> derivative_action_callbacks;
+   std::map<size_t,
+       std::vector<derivative_action_t>> daction_transpose_callbacks;
+
+   std::vector<FieldDescriptor> solutions;
+   std::vector<FieldDescriptor> parameters;
+   // solutions and parameters
+   std::vector<FieldDescriptor> fields;
+
+   mutable std::vector<Vector> solutions_l;
+   mutable std::vector<Vector> parameters_l;
+   mutable Vector residual_l;
+
+   mutable std::vector<Vector> fields_e;
+   mutable Vector residual_e;
+
+   std::function<void(Vector &, Vector &)> prolongation_transpose;
+   std::function<void(Vector &, Vector &)> output_restriction_transpose;
+   restriction_callback_t restriction_callback;
+
+   std::map<size_t, size_t> assembled_vector_sizes;
+
+   bool use_tensor_product_structure = true;
+
+   size_t test_space_field_idx = SIZE_MAX;
+};
+
+void DifferentiableOperator::SetParameters(std::vector<Vector *> p) const
+{
+   MFEM_ASSERT(parameters.size() == p.size(),
+               "number of parameters doesn't match descriptors");
+   for (size_t i = 0; i < parameters.size(); i++)
+   {
+      p[i]->Read();
+      parameters_l[i] = *p[i];
+   }
+}
+
+DifferentiableOperator::DifferentiableOperator(
+   const std::vector<FieldDescriptor> &solutions,
+   const std::vector<FieldDescriptor> &parameters,
+   const ParMesh &mesh) :
+   mesh(mesh),
+   solutions(solutions),
+   parameters(parameters)
+{
+   fields.resize(solutions.size() + parameters.size());
+   fields_e.resize(fields.size());
+   solutions_l.resize(solutions.size());
+   parameters_l.resize(parameters.size());
+
+   for (size_t i = 0; i < solutions.size(); i++)
+   {
+      fields[i] = solutions[i];
+   }
+
+   for (size_t i = 0; i < parameters.size(); i++)
+   {
+      fields[i + solutions.size()] = parameters[i];
+   }
+}
+
+template <
+   typename qfunc_t,
+   typename... input_ts,
+   typename... output_ts,
+   typename derivative_ids_t = std::make_index_sequence<0>>
+void DifferentiableOperator::AddDomainIntegrator(
+   qfunc_t &qfunc,
+   mfem::tuple<input_ts...> inputs,
+   mfem::tuple<output_ts...> outputs,
+   const IntegrationRule &integration_rule,
+   const Array<int> domain_attributes,
+   derivative_ids_t derivative_ids)
+{
+   using entity_t = Entity::Element;
+
+   static constexpr size_t num_inputs =
+      mfem::tuple_size<decltype(inputs)>::value;
+
+   static constexpr size_t num_outputs =
+      mfem::tuple_size<decltype(outputs)>::value;
+
+   using qf_signature =
+      typename create_function_signature<decltype(&qfunc_t::operator())>::type;
+   using qf_param_ts = typename qf_signature::parameter_ts;
+   using qf_output_t = typename qf_signature::return_t;
+
+   // Consistency checks
+   if constexpr (num_outputs > 1)
+   {
+      static_assert(always_false<qfunc_t>,
+                    "more than one output per quadrature functions is not supported right now");
+   }
+
+   if constexpr (std::is_same_v<qf_output_t, void>)
+   {
+      static_assert(always_false<qfunc_t>, "quadrature function has no return value");
+   }
+
+   constexpr size_t num_qfinputs = mfem::tuple_size<qf_param_ts>::value;
+   static_assert(num_qfinputs == num_inputs,
+                 "quadrature function inputs and descriptor inputs have to match");
+
+   constexpr size_t num_qf_outputs = mfem::tuple_size<qf_output_t>::value;
+   static_assert(num_qf_outputs == num_outputs,
+                 "quadrature function outputs and descriptor outputs have to match");
+
+   constexpr auto inout_tuple = std::tuple_cat(std::tuple<input_ts...> {},
+                                               std::tuple<output_ts...> {});
+   constexpr auto filtered_inout_tuple = filter_fields(inout_tuple);
+   constexpr size_t num_fields = count_unique_field_ids(filtered_inout_tuple);
+
+   MFEM_ASSERT(num_fields == solutions.size() + parameters.size(),
+               "Total number of fields doesn't match sum of solutions and parameters."
+               " This indicates that some fields are not used in the integrator,"
+               " which currently is not supported.");
+
+   auto dependency_map = make_dependency_map(mfem::tuple<input_ts...> {});
+
+   // pretty_print(dependency_map);
+
+   auto input_to_field =
+      create_descriptors_to_fields_map<entity_t>(fields, inputs);
+   auto output_to_field =
+      create_descriptors_to_fields_map<entity_t>(fields, outputs);
+
+   // TODO: factor out
+   std::vector<int> inputs_vdim(num_inputs);
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      inputs_vdim[i] = mfem::get<i>(inputs).vdim;
+   });
+
+   const int NE = mesh.GetNE();
+   if (NE == 0)
+   {
+      // use of GetElement(0), GetFE(0) in GetDofToQuad assume that NE > 0
+      MFEM_ABORT("Mesh with no elements is not yet supported!");
+   }
+
+   Array<int> elem_attributes;
+   if (NE > 0)
+   {
+      elem_attributes.SetSize(NE);
+      for (int i = 0; i < NE; ++i)
+      {
+         elem_attributes[i] = mesh.GetAttribute(i);
+      }
+   }
+
+   const auto output_fop = mfem::get<0>(outputs);
+   test_space_field_idx = FindIdx(output_fop.GetFieldId(), fields);
+
+   bool use_sum_factorization = false;
+   auto entity_element_type =  mesh.GetElement(0)->GetType();
+   if ((entity_element_type == Element::QUADRILATERAL ||
+        entity_element_type == Element::HEXAHEDRON) &&
+       use_tensor_product_structure == true)
+   {
+      use_sum_factorization = true;
+   }
+
+   ElementDofOrdering element_dof_ordering = ElementDofOrdering::NATIVE;
+   DofToQuad::Mode doftoquad_mode = DofToQuad::Mode::FULL;
+   if (use_sum_factorization)
+   {
+      element_dof_ordering = ElementDofOrdering::LEXICOGRAPHIC;
+      doftoquad_mode = DofToQuad::Mode::TENSOR;
+   }
+
+   auto [output_rt,
+         output_e_sz] = get_restriction_transpose<entity_t>
+                        (fields[test_space_field_idx],
+                         element_dof_ordering, output_fop);
+   auto &output_e_size = output_e_sz;
+
+   output_restriction_transpose = output_rt;
+   residual_e.SetSize(output_e_size);
+
+   // The explicit captures are necessary to avoid dependency on
+   // the specific instance of this class (this pointer).
+   restriction_callback =
+      [=, solutions = this->solutions, parameters = this->parameters]
+      (std::vector<Vector> &solutions_l,
+       const std::vector<Vector> &parameters_l,
+       std::vector<Vector> &fields_e)
+   {
+      restriction<entity_t>(solutions, solutions_l, fields_e,
+                            element_dof_ordering);
+      restriction<entity_t>(parameters, parameters_l, fields_e,
+                            element_dof_ordering,
+                            solutions.size());
+   };
+
+   prolongation_transpose = get_prolongation_transpose(
+                               fields[test_space_field_idx], output_fop, mesh.GetComm());
+
+   const int dimension = mesh.Dimension();
+   [[maybe_unused]] const int num_elements = GetNumEntities<Entity::Element>(mesh);
+   const int num_entities = GetNumEntities<entity_t>(mesh);
+   const int num_qp = integration_rule.GetNPoints();
+
+   if constexpr (is_one_fop<decltype(output_fop)>::value)
+   {
+      residual_l.SetSize(1);
+      height = 1;
+   }
+   else
+   {
+      const int residual_lsize = GetVSize(fields[test_space_field_idx]);
+      residual_l.SetSize(residual_lsize);
+      height = GetTrueVSize(fields[test_space_field_idx]);
+   }
+
+   // TODO: Is this a hack?
+   width = GetTrueVSize(fields[0]);
+
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(
+                          field,
+                          integration_rule,
+                          doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/dimension) + 0.5);
+
+   const int residual_size_on_qp =
+      GetSizeOnQP<entity_t>(output_fop,
+                            fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(inputs, dtq, input_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(outputs, dtq, output_to_field);
+
+   const int test_vdim = output_fop.vdim;
+   const int test_op_dim = output_fop.size_on_qp / output_fop.vdim;
+   MFEM_VERIFY(num_entities > 0,
+               "The number of entities must be greater than zero");
+   const int num_test_dof = output_e_size / output_fop.vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp =
+      get_input_size_on_qp(inputs, std::make_index_sequence<num_inputs> {});
+
+   auto action_shmem_info =
+      get_shmem_info<entity_t, num_fields, num_inputs, num_outputs>
+      (input_dtq_maps, output_dtq_maps, fields, num_entities, inputs, num_qp,
+       input_size_on_qp, residual_size_on_qp, element_dof_ordering);
+
+   Vector shmem_cache(action_shmem_info.total_size);
+
+   // print_shared_memory_info(action_shmem_info);
+
+   ThreadBlocks thread_blocks;
+   if (dimension == 3)
+   {
+      if (use_sum_factorization)
+      {
+         thread_blocks.x = q1d;
+         thread_blocks.y = q1d;
+         thread_blocks.z = q1d;
+      }
+   }
+   else if (dimension == 2)
+   {
+      if (use_sum_factorization)
+      {
+         thread_blocks.x = q1d;
+         thread_blocks.y = q1d;
+         thread_blocks.z = 1;
+      }
+   }
+
+   action_callbacks.push_back(
+      [=, restriction_callback = this->restriction_callback]
+      (std::vector<Vector> &solutions_l,
+       const std::vector<Vector> &parameters_l,
+       Vector &residual_l) mutable
+   {
+      restriction_callback(solutions_l, parameters_l, fields_e);
+
+      residual_e = 0.0;
+      auto ye = Reshape(residual_e.ReadWrite(), test_vdim, num_test_dof, num_entities);
+
+      auto wrapped_fields_e = wrap_fields(fields_e,
+                                          action_shmem_info.field_sizes,
+                                          num_entities);
+
+      const bool has_attr = domain_attributes.Size() > 0;
+      const auto d_domain_attr = domain_attributes.Read();
+      const auto d_elem_attr = elem_attributes.Read();
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         if (has_attr && !d_domain_attr[d_elem_attr[e] - 1]) { return; }
+
+         auto [input_dtq_shmem, output_dtq_shmem, fields_shmem, input_shmem,
+                                residual_shmem, scratch_shmem] =
+                  unpack_shmem(shmem, action_shmem_info, input_dtq_maps, output_dtq_maps,
+                               wrapped_fields_e, num_qp, e);
+
+         map_fields_to_quadrature_data(
+            input_shmem, fields_shmem, input_dtq_shmem, input_to_field, inputs, ir_weights,
+            scratch_shmem, dimension, use_sum_factorization);
+
+         call_qfunction<qf_param_ts>(
+            qfunc, input_shmem, residual_shmem,
+            residual_size_on_qp, num_qp, q1d, dimension, use_sum_factorization);
+
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields(
+            y, fhat, output_fop, output_dtq_shmem[0],
+            scratch_shmem, dimension, use_sum_factorization);
+      }, num_entities, thread_blocks, action_shmem_info.total_size, shmem_cache.ReadWrite());
+      output_restriction_transpose(residual_e, residual_l);
+   });
+
+   // Create the action of the derivatives
+   for_constexpr([&](auto derivative_id)
+   {
+      const size_t d_field_idx = FindIdx(derivative_id, fields);
+      const auto direction = fields[d_field_idx];
+      const int da_size_on_qp = GetSizeOnQP<entity_t>(output_fop,
+                                                      fields[test_space_field_idx]);
+
+      auto shmem_info =
+         get_shmem_info<entity_t, num_fields, num_inputs, num_outputs>
+         (input_dtq_maps, output_dtq_maps, fields, num_entities, inputs, num_qp,
+          input_size_on_qp, residual_size_on_qp, element_dof_ordering, d_field_idx);
+
+      Vector shmem_cache(shmem_info.total_size);
+
+      // print_shared_memory_info(shmem_info);
+
+      Vector direction_e;
+      Vector derivative_action_e(output_e_size);
+      derivative_action_e = 0.0;
+
+      const auto input_is_dependent = dependency_map[derivative_id];
+
+      derivative_action_callbacks[derivative_id].push_back(
+         [=, output_restriction_transpose = this->output_restriction_transpose](
+            std::vector<Vector> &fields_e, const Vector &direction_l,
+            Vector &derivative_action_l) mutable
+      {
+         restriction<entity_t>(direction, direction_l, direction_e, element_dof_ordering);
+         auto ye = Reshape(derivative_action_e.ReadWrite(), num_test_dof, test_vdim, num_entities);
+         auto wrapped_fields_e = wrap_fields(fields_e, shmem_info.field_sizes, num_entities);
+         auto wrapped_direction_e = Reshape(direction_e.ReadWrite(), shmem_info.direction_size, num_entities);
+
+         derivative_action_e = 0.0;
+         forall([=] MFEM_HOST_DEVICE (int e, double *shmem)
+         {
+            auto [input_dtq_shmem, output_dtq_shmem, fields_shmem, direction_shmem,
+                                   input_shmem, shadow_shmem_, residual_shmem, scratch_shmem] =
+            unpack_shmem(shmem, shmem_info, input_dtq_maps,
+                         output_dtq_maps, wrapped_fields_e, wrapped_direction_e, num_qp, e);
+            auto &shadow_shmem = shadow_shmem_;
+
+            map_fields_to_quadrature_data(
+               input_shmem, fields_shmem, input_dtq_shmem, input_to_field, inputs, ir_weights,
+               scratch_shmem, dimension, use_sum_factorization);
+
+            // TODO: Probably redundant
+            set_zero(shadow_shmem);
+
+            map_direction_to_quadrature_data_conditional(
+               shadow_shmem, direction_shmem, input_dtq_shmem, inputs, ir_weights,
+               scratch_shmem, input_is_dependent, dimension, use_sum_factorization);
+
+            call_qfunction_derivative_action<qf_param_ts>(
+               qfunc, input_shmem, shadow_shmem, residual_shmem,
+               da_size_on_qp, num_qp, q1d, dimension, use_sum_factorization);
+
+            auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+            auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+            map_quadrature_data_to_fields(
+               y, fhat, output_fop, output_dtq_shmem[0],
+               scratch_shmem, dimension, use_sum_factorization);
+         }, num_entities, thread_blocks, shmem_info.total_size, shmem_cache.ReadWrite());
+         output_restriction_transpose(derivative_action_e, derivative_action_l);
+      });
+   }, derivative_ids);
+
+   // Create the transpose action of the derivatives
+   if (!use_sum_factorization)
+   {
+      for_constexpr([&](auto derivative_id)
+      {
+         const size_t d_field_idx = FindIdx(derivative_id, fields);
+         const auto direction = fields[test_space_field_idx];
+         const int da_size_on_qp = GetSizeOnQP<entity_t>(output_fop,
+                                                         fields[test_space_field_idx]);
+
+         auto shmem_info =
+            get_shmem_info<entity_t, num_fields, num_inputs, num_outputs>
+            (input_dtq_maps, output_dtq_maps, fields, num_entities, inputs, num_qp,
+             input_size_on_qp, residual_size_on_qp, element_dof_ordering,
+             test_space_field_idx);
+
+         Vector shmem_cache(shmem_info.total_size);
+
+         // print_shared_memory_info(shmem_info);
+
+         auto [RT, e_size] = get_restriction_transpose<entity_t>(
+                                fields[d_field_idx],
+                                element_dof_ordering,
+                                mfem::get<0>(inputs)); // TODO
+
+         Vector direction_e;
+         Vector daction_transpose_e(e_size);
+         daction_transpose_e = 0.0;
+
+         const auto input_is_dependent = dependency_map[derivative_id];
+
+         const int trial_vdim = GetVDim(fields[0]);
+
+         int total_trial_op_dim = 0;
+         for_constexpr<num_inputs>([&](auto s)
+         {
+            if (!input_is_dependent[s])
+            {
+               return;
+            }
+            auto B = is_value_fop<decltype(mfem::get<s>(inputs))>::value ?
+                     input_dtq_maps[s].B : input_dtq_maps[s].G;
+            total_trial_op_dim += B.GetShape()[DofToQuadMap::Index::DIM];
+         });
+
+         daction_transpose_callbacks[derivative_id].push_back(
+            [=, restriction_transpose = RT](
+               std::vector<Vector> &fields_e, const Vector &direction_l,
+               Vector &daction_l) mutable
+         {
+            auto shmem = shmem_cache.ReadWrite();
+
+            restriction<entity_t>(direction, direction_l, direction_e, element_dof_ordering);
+            auto ye = Reshape(daction_transpose_e.ReadWrite(), num_test_dof, trial_vdim, num_entities);
+            auto wrapped_fields_e = wrap_fields(fields_e, shmem_info.field_sizes, num_entities);
+            auto wrapped_direction_e = Reshape(direction_e.ReadWrite(), shmem_info.direction_size, num_entities);
+
+            Vector a_qp_mem(test_vdim * test_op_dim * trial_vdim * total_trial_op_dim);
+            auto a_qp = Reshape(a_qp_mem.ReadWrite(), test_vdim, test_op_dim,
+                                trial_vdim, total_trial_op_dim);
+
+            Vector dir_mem(shmem_info.shadow_sizes[test_space_field_idx]);
+            auto dir = Reshape(dir_mem.ReadWrite(), input_size_on_qp[test_space_field_idx], num_qp);
+
+            daction_transpose_e = 0.0;
+            for (int e = 0; e < num_entities; e++)
+            {
+               auto [input_dtq_shmem, output_dtq_shmem, fields_shmem, direction_shmem,
+                                      input_shmem_, shadow_shmem_, residual_shmem_, scratch_shmem] =
+               unpack_shmem(shmem, shmem_info, input_dtq_maps,
+                            output_dtq_maps, wrapped_fields_e, wrapped_direction_e, num_qp, e);
+               // avoid captured structured bindings
+               auto &input_shmem = input_shmem_;
+               auto &shadow_shmem = shadow_shmem_;
+               auto &residual_shmem = residual_shmem_;
+
+               map_fields_to_quadrature_data(
+                  input_shmem, fields_shmem, input_dtq_shmem, input_to_field, inputs, ir_weights,
+                  scratch_shmem, dimension, use_sum_factorization);
+
+               set_zero(shadow_shmem);
+               std::array<bool, num_inputs> direction_is_dependent{false};
+               direction_is_dependent[test_space_field_idx] = true;
+
+               map_direction_to_quadrature_data_conditional(
+                  shadow_shmem, direction_shmem, input_dtq_shmem, inputs, ir_weights,
+                  scratch_shmem, direction_is_dependent, use_sum_factorization);
+
+               copy(shadow_shmem[test_space_field_idx], dir);
+               set_zero(shadow_shmem);
+
+               // pretty_print(dir_mem);
+
+               for (int q = 0; q < num_qp; q++)
+               {
+                  for (int j = 0; j < trial_vdim; j++)
+                  {
+                     size_t m_offset = 0;
+
+                     for_constexpr<num_inputs>([&](auto s)
+                     {
+                        if (!input_is_dependent[s])
+                        {
+                           return;
+                        }
+
+                        auto B = is_value_fop<std::decay_t<decltype(mfem::get<s>(inputs))>>::value ?
+                                 input_dtq_maps[s].B : input_dtq_maps[s].G;
+                        auto trial_op_dim = B.GetShape()[DofToQuadMap::Index::DIM];
+                        auto d_qp = Reshape(&(shadow_shmem[s])[0], trial_vdim, trial_op_dim, num_qp);
+                        for (int m = 0; m < trial_op_dim; m++)
+                        {
+                           d_qp(j, m, q) = 1.0;
+
+                           auto r = Reshape(&residual_shmem(0, q), da_size_on_qp);
+                           auto qf_args = decay_tuple<qf_param_ts> {};
+#ifdef MFEM_USE_ENZYME
+                           auto qf_shadow_args = decay_tuple<qf_param_ts> {};
+                           apply_kernel_fwddiff_enzyme(r, qfunc, qf_args, qf_shadow_args, input_shmem,
+                                                       shadow_shmem, q);
+#else
+                           MFEM_ABORT("Native dual support is not enabled!");
+                           // apply_kernel_native_dual(r, qfunc, qf_args, input_shmem, shadow_shmem, q);
+#endif
+                           d_qp(j, m, q) = 0.0;
+
+                           auto f = Reshape(&r(0), test_vdim, test_op_dim);
+                           for (int i = 0; i < test_vdim; i++)
+                           {
+                              for (int k = 0; k < test_op_dim; k++)
+                              {
+                                 a_qp(i, k, j, m + m_offset) = f(i, k);
+                              }
+                           }
+                        }
+                        m_offset += trial_op_dim;
+                     });
+                  }
+
+                  // pretty_print(a_qp_mem);
+
+                  // Multiply transpose of a_qp with direction
+                  // auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+                  // auto dir_qp = Reshape(&dir[0], trial_vdim, total_trial_op_dim, num_qp);
+                  auto fhat = Reshape(&residual_shmem(0, 0), trial_vdim, total_trial_op_dim,
+                                      num_qp);
+                  auto dir_qp = Reshape(&dir(0, 0), test_vdim, test_op_dim, num_qp);
+                  for (int i = 0; i < trial_vdim; i++)
+                  {
+                     for (int k = 0; k < total_trial_op_dim; k++)
+                     {
+                        fhat(i, k, q) = 0.0;
+                        for (int j = 0; j < test_vdim; j++)
+                        {
+                           for (int m = 0; m < test_op_dim; m++)
+                           {
+                              fhat(i, k, q) += a_qp(j, m, i, k) * dir_qp(j, m, q);
+                           }
+                        }
+                     }
+                  }
+               }
+
+               auto fhat = Reshape(&residual_shmem(0, 0), trial_vdim, total_trial_op_dim,
+                                   num_qp);
+               int num_trial_dof = input_dtq_shmem[0].B.GetShape()[DofToQuadMap::Index::DOF];
+               auto y = Reshape(&ye(0, 0, e), num_trial_dof, trial_vdim);
+               map_quadrature_data_to_fields(
+                  y, fhat, mfem::get<0>(inputs), input_dtq_shmem[0],
+                  scratch_shmem, dimension, use_sum_factorization);
+            }
+            restriction_transpose(daction_transpose_e, daction_l);
+         });
+      }, derivative_ids);
+   }
+}
+
+
+} // namespace mfem

fem/dfem/fieldoperator.hpp

@@ -0,0 +1,139 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+namespace mfem
+{
+
+template <int FIELD_ID = -1>
+class FieldOperator
+{
+public:
+   constexpr FieldOperator(int size_on_qp = 0) :
+      size_on_qp(size_on_qp) {};
+
+   static constexpr int GetFieldId() { return FIELD_ID; }
+
+   int size_on_qp = -1;
+
+   int dim = -1;
+
+   int vdim = -1;
+};
+
+template <int FIELD_ID = -1>
+class None : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr None() : FieldOperator<FIELD_ID>() {}
+};
+
+template< typename T >
+struct is_none_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_none_fop<None<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <typename T>
+struct DisableAD
+{
+   T& operator()() const { return fop; }
+   T fop;
+};
+
+class Weight : public FieldOperator<-1>
+{
+public:
+   constexpr Weight() : FieldOperator<-1>() {};
+};
+
+template< typename T >
+struct is_weight_fop
+{
+   static const bool value = false;
+};
+
+template <>
+struct is_weight_fop<Weight>
+{
+   static const bool value = true;
+};
+
+template <int FIELD_ID = -1>
+class Value : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr Value() : FieldOperator<FIELD_ID>() {};
+};
+
+template< typename T >
+struct is_value_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_value_fop<Value<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <typename T>
+struct is_value_fop<DisableAD<T>>
+{
+   static const bool value = is_value_fop<T>::value;
+};
+
+template <int FIELD_ID = -1>
+class Gradient : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr Gradient() : FieldOperator<FIELD_ID>() {};
+};
+
+template< typename T >
+struct is_gradient_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_gradient_fop<Gradient<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <int FIELD_ID = -1>
+class One : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr One() : FieldOperator<FIELD_ID>() {};
+};
+
+template< typename T >
+struct is_one_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_one_fop<One<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+} // namespace mfem

fem/dfem/integrate.hpp

@@ -0,0 +1,448 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+#include "util.hpp"
+
+namespace mfem
+{
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields_impl(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+   // assuming the quadrature point residual has to "play nice with
+   // the test function"
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [num_qp, cdim, num_dof] = B.GetShape();
+      const int vdim = output.vdim > 0 ? output.vdim : cdim ;
+      for (int dof = 0; dof < num_dof; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double acc = 0.0;
+            for (int qp = 0; qp < num_qp; qp++)
+            {
+               acc += B(qp, 0, dof) * f(vd, 0, qp);
+            }
+            y(dof, vd) += acc;
+         }
+      }
+   }
+   else if constexpr (
+      is_gradient_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [num_qp, dim, num_dof] = G.GetShape();
+      const int vdim = output.vdim;
+      for (int dof = 0; dof < num_dof; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double acc = 0.0;
+            for (int d = 0; d < dim; d++)
+            {
+               for (int qp = 0; qp < num_qp; qp++)
+               {
+                  acc += G(qp, d, dof) * f(vd, d, qp);
+               }
+            }
+            y(dof, vd) += acc;
+         }
+      }
+   }
+   else if constexpr (is_one_fop<std::decay_t<output_t>>::value)
+   {
+      // This is the "integral over all quadrature points type" applying
+      // B = 1 s.t. B^T * C \in R^1.
+      const auto [num_qp, unused, unused1] = B.GetShape();
+      auto cc = Reshape(&f(0, 0, 0), num_qp);
+      for (int i = 0; i < num_qp; i++)
+      {
+         y(0, 0) += cc(i);
+      }
+   }
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [num_qp, unused, num_dof] = B.GetShape();
+      const auto vdim = output.vdim;
+      auto cc = Reshape(&f(0, 0, 0), num_qp * vdim);
+      auto yy = Reshape(&y(0, 0), num_qp * vdim);
+      for (int i = 0; i < num_qp * vdim; i++)
+      {
+         yy(i) = cc(i);
+      }
+   }
+   else
+   {
+      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+                 " this field descriptor");
+   }
+}
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields_tensor_impl_2d(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq,
+   std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               double acc = 0.0;
+               for (int qx = 0; qx < q1d; qx++)
+               {
+                  acc += fqp(vd, 0, qx, qy) * B(qx, 0, dx);
+               }
+               s0(qy, dx) = acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               double acc = 0.0;
+               for (int qy = 0; qy < q1d; qy++)
+               {
+                  acc += s0(qy, dx) * B(qy, 0, dy);
+               }
+               yd(dx, dy, vd) += acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_gradient_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = G.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, d1d);
+      auto s1 = Reshape(&scratch_mem[1](0), q1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int qx = 0; qx < q1d; qx++)
+               {
+                  uv[0] += fqp(vd, 0, qx, qy) * G(qx, 0, dx);
+                  uv[1] += fqp(vd, 1, qx, qy) * B(qx, 0, dx);
+               }
+               s0(qy, dx) = uv[0];
+               s1(qy, dx) = uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int qy = 0; qy < q1d; qy++)
+               {
+                  uv[0] += s0(qy, dx) * B(qy, 0, dy);
+                  uv[1] += s1(qy, dx) * G(qy, 0, dy);
+               }
+               yd(dx, dy, vd) += uv[0] + uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+
+      // TODO: Check if this is the right fix for all cases
+      auto fqp = Reshape(&f(0, 0, 0), output.size_on_qp, q1d);
+      auto yqp = Reshape(&y(0, 0), output.size_on_qp, q1d);
+      for (int sq = 0; sq < output.size_on_qp; sq++)
+      {
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            yqp(sq, qx) = fqp(sq, qx);
+         }
+         MFEM_SYNC_THREAD;
+      }
+
+      // auto fqp = Reshape(&f(0, 0, 0), output.size_on_qp, q1d, q1d);
+      // auto yqp = Reshape(&y(0, 0), output.size_on_qp, q1d, q1d);
+
+      // for (int sq = 0; sq < output.size_on_qp; sq++)
+      // {
+      //    MFEM_FOREACH_THREAD(qx, x, q1d)
+      //    {
+      //       MFEM_FOREACH_THREAD(qy, y, q1d)
+      //       {
+      //          yqp(sq, qx, qy) = fqp(sq, qx, qy);
+      //       }
+      //    }
+      //    MFEM_SYNC_THREAD;
+      // }
+   }
+   else
+   {
+      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+                 " this field descriptor with sum factorization on tensor product elements");
+   }
+}
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields_tensor_impl_3d(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq,
+   std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, q1d, d1d);
+      auto s1 = Reshape(&scratch_mem[1](0), q1d, d1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  double acc = 0.0;
+                  for (int qx = 0; qx < q1d; qx++)
+                  {
+                     acc += fqp(vd, 0, qx, qy, qz) * B(qx, 0, dx);
+                  }
+                  s0(qz, qy, dx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  double acc = 0.0;
+                  for (int qy = 0; qy < q1d; qy++)
+                  {
+                     acc += s0(qz, qy, dx) * B(qy, 0, dy);
+                  }
+                  s1(qz, dy, dx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               MFEM_FOREACH_THREAD(dz, z, d1d)
+               {
+                  double acc = 0.0;
+                  for (int qz = 0; qz < q1d; qz++)
+                  {
+                     acc += s1(qz, dy, dx) * B(qz, 0, dz);
+                  }
+                  yd(dx, dy, dz, vd) += acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_gradient_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = G.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, q1d, d1d);
+      auto s1 = Reshape(&scratch_mem[1](0), q1d, q1d, d1d);
+      auto s2 = Reshape(&scratch_mem[2](0), q1d, q1d, d1d);
+      auto s3 = Reshape(&scratch_mem[3](0), q1d, d1d, d1d);
+      auto s4 = Reshape(&scratch_mem[4](0), q1d, d1d, d1d);
+      auto s5 = Reshape(&scratch_mem[5](0), q1d, d1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(dx, x, d1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int qx = 0; qx < q1d; qx++)
+                  {
+                     uvw[0] += fqp(vd, 0, qx, qy, qz) * G(qx, 0, dx);
+                     uvw[1] += fqp(vd, 1, qx, qy, qz) * B(qx, 0, dx);
+                     uvw[2] += fqp(vd, 2, qx, qy, qz) * B(qx, 0, dx);
+                  }
+                  s0(qz, qy, dx) = uvw[0];
+                  s1(qz, qy, dx) = uvw[1];
+                  s2(qz, qy, dx) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(dx, x, d1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int qy = 0; qy < q1d; qy++)
+                  {
+                     uvw[0] += s0(qz, qy, dx) * B(qy, 0, dy);
+                     uvw[1] += s1(qz, qy, dx) * G(qy, 0, dy);
+                     uvw[2] += s2(qz, qy, dx) * B(qy, 0, dy);
+                  }
+                  s3(qz, dy, dx) = uvw[0];
+                  s4(qz, dy, dx) = uvw[1];
+                  s5(qz, dy, dx) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(dx, x, d1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int qz = 0; qz < q1d; qz++)
+                  {
+                     uvw[0] += s3(qz, dy, dx) * B(qz, 0, dz);
+                     uvw[1] += s4(qz, dy, dx) * B(qz, 0, dz);
+                     uvw[2] += s5(qz, dy, dx) * G(qz, 0, dz);
+                  }
+                  yd(dx, dy, dz, vd) += uvw[0] + uvw[1] + uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      auto fqp = Reshape(&f(0, 0, 0), output.size_on_qp, q1d, q1d, q1d);
+      auto yqp = Reshape(&y(0, 0), output.size_on_qp, q1d, q1d, q1d);
+
+      for (int sq = 0; sq < output.size_on_qp; sq++)
+      {
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  yqp(sq, qx, qy, qz) = fqp(sq, qx, qy, qz);
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else
+   {
+      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+                 " this field descriptor with sum factorization on tensor product elements");
+   }
+}
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq,
+   std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const int &dimension,
+   const bool &use_sum_factorization)
+{
+   if (use_sum_factorization)
+   {
+      if (dimension == 2)
+      {
+         map_quadrature_data_to_fields_tensor_impl_2d(y, f, output, dtq, scratch_mem);
+      }
+      else if (dimension == 3)
+      {
+         map_quadrature_data_to_fields_tensor_impl_3d(y, f, output, dtq, scratch_mem);
+      }
+   }
+   else
+   {
+      map_quadrature_data_to_fields_impl(y, f, output, dtq);
+   }
+}
+
+}

fem/dfem/interpolate.hpp

@@ -0,0 +1,579 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+#include "util.hpp"
+
+namespace mfem
+{
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE inline
+void map_field_to_quadrature_data_tensor_product_3d(
+   DeviceTensor<2> &field_qp,
+   const DofToQuadMap &dtq,
+   const DeviceTensor<1> &field_e,
+   const field_operator_t &input,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<field_operator_t>>::value)
+   {
+      auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, q1d, q1d, q1d);
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, d1d, q1d);
+      auto s1 = Reshape(&scratch_mem[1](0), d1d, q1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  double acc = 0.0;
+                  for (int dx = 0; dx < d1d; dx++)
+                  {
+                     acc += B(qx, 0, dx) * field(dx, dy, dz, vd);
+                  }
+                  s0(dz, dy, qx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               MFEM_FOREACH_THREAD(qy, y, q1d)
+               {
+                  double acc = 0.0;
+                  for (int dy = 0; dy < d1d; dy++)
+                  {
+                     acc += s0(dz, dy, qx) * B(qy, 0, dy);
+                  }
+                  s1(dz, qy, qx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  double acc = 0.0;
+                  for (int dz = 0; dz < d1d; dz++)
+                  {
+                     acc += s1(dz, qy, qx) * B(qz, 0, dz);
+                  }
+                  fqp(vd, qx, qy, qz) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (
+      is_gradient_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const int dim = input.dim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, dim, q1d, q1d, q1d);
+
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, d1d, q1d);
+      auto s1 = Reshape(&scratch_mem[1](0), d1d, d1d, q1d);
+      auto s2 = Reshape(&scratch_mem[2](0), d1d, q1d, q1d);
+      auto s3 = Reshape(&scratch_mem[3](0), d1d, q1d, q1d);
+      auto s4 = Reshape(&scratch_mem[4](0), d1d, q1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  real_t uv[2] = {0.0, 0.0};
+                  for (int dx = 0; dx < d1d; dx++)
+                  {
+                     const real_t f = field(dx, dy, dz, vd);
+                     uv[0] += f * B(qx, 0, dx);
+                     uv[1] += f * G(qx, 0, dx);
+                  }
+                  s0(dz, dy, qx) = uv[0];
+                  s1(dz, dy, qx) = uv[1];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int dy = 0; dy < d1d; dy++)
+                  {
+                     const real_t s0i = s0(dz, dy, qx);
+                     uvw[0] += s1(dz, dy, qx) * B(qy, 0, dy);
+                     uvw[1] += s0i * G(qy, 0, dy);
+                     uvw[2] += s0i * B(qy, 0, dy);
+                  }
+                  s2(dz, qy, qx) = uvw[0];
+                  s3(dz, qy, qx) = uvw[1];
+                  s4(dz, qy, qx) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int dz = 0; dz < d1d; dz++)
+                  {
+                     uvw[0] += s2(dz, qy, qx) * B(qz, 0, dz);
+                     uvw[1] += s3(dz, qy, qx) * B(qz, 0, dz);
+                     uvw[2] += s4(dz, qy, qx) * G(qz, 0, dz);
+                  }
+                  fqp(vd, 0, qx, qy, qz) = uvw[0];
+                  fqp(vd, 1, qx, qy, qz) = uvw[1];
+                  fqp(vd, 2, qx, qy, qz) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   // TODO: Create separate function for clarity
+   else if constexpr (
+      std::is_same_v<std::decay_t<field_operator_t>, Weight>)
+   {
+      const int num_qp = integration_weights.GetShape()[0];
+      // TODO: eeek
+      const int q1d = (int)floor(pow(num_qp, 1.0/input.dim) + 0.5);
+      auto w = Reshape(&integration_weights[0], q1d, q1d, q1d);
+      auto f = Reshape(&field_qp[0], q1d, q1d, q1d);
+      MFEM_FOREACH_THREAD(qx, x, q1d)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(qz, z, q1d)
+            {
+               f(qx, qy, qz) = w(qx, qy, qz);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   }
+   else if constexpr (is_none_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const int q1d = B.GetShape()[0];
+      auto field = Reshape(&field_e[0], input.size_on_qp, q1d * q1d * q1d);
+      field_qp = field;
+   }
+   else
+   {
+      static_assert(always_false<std::decay_t<field_operator_t>>,
+                    "can't map field to quadrature data");
+   }
+}
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE inline
+void map_field_to_quadrature_data_tensor_product_2d(
+   DeviceTensor<2> &field_qp,
+   const DofToQuadMap &dtq,
+   const DeviceTensor<1> &field_e,
+   const field_operator_t &input,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<field_operator_t>>::value)
+   {
+      auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, q1d, q1d);
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               double acc = 0.0;
+               for (int dx = 0; dx < d1d; dx++)
+               {
+                  acc += B(qx, 0, dx) * field(dx, dy, vd);
+               }
+               s0(dy, qx) = acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               double acc = 0.0;
+               for (int dy = 0; dy < d1d; dy++)
+               {
+                  acc += s0(dy, qx) * B(qy, 0, dy);
+               }
+               fqp(vd, qx, qy) = acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (
+      is_gradient_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const int dim = input.dim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, dim, q1d, q1d);
+
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, q1d);
+      auto s1 = Reshape(&scratch_mem[1](0), d1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int dx = 0; dx < d1d; dx++)
+               {
+                  const real_t f = field(dx, dy, vd);
+                  uv[0] += f * B(qx, 0, dx);
+                  uv[1] += f * G(qx, 0, dx);
+               }
+               s0(dy, qx) = uv[0];
+               s1(dy, qx) = uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int dy = 0; dy < d1d; dy++)
+               {
+                  const real_t s0i = s0(dy, qx);
+                  uv[0] += s1(dy, qx) * B(qy, 0, dy);
+                  uv[1] += s0i * G(qy, 0, dy);
+               }
+               fqp(vd, 0, qx, qy) = uv[0];
+               fqp(vd, 1, qx, qy) = uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   // TODO: Create separate function for clarity
+   else if constexpr (
+      std::is_same_v<std::decay_t<field_operator_t>, Weight>)
+   {
+      const int num_qp = integration_weights.GetShape()[0];
+      // TODO: eeek
+      const int q1d = (int)floor(pow(num_qp, 1.0/input.dim) + 0.5);
+      auto w = Reshape(&integration_weights[0], q1d, q1d);
+      auto f = Reshape(&field_qp[0], q1d, q1d);
+      MFEM_FOREACH_THREAD(qx, x, q1d)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            f(qx, qy) = w(qx, qy);
+         }
+      }
+      MFEM_SYNC_THREAD;
+   }
+   else if constexpr (is_none_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const int q1d = B.GetShape()[0];
+      auto field = Reshape(&field_e[0], input.size_on_qp, q1d * q1d);
+      field_qp = field;
+   }
+   else
+   {
+      static_assert(always_false<std::decay_t<field_operator_t>>,
+                    "can't map field to quadrature data");
+   }
+}
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE
+void map_field_to_quadrature_data(
+   DeviceTensor<2> field_qp,
+   const DofToQuadMap &dtq,
+   const DeviceTensor<1> &field_e,
+   const field_operator_t &input,
+   const DeviceTensor<1, const double> &integration_weights)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+   if constexpr (is_value_fop<field_operator_t>::value)
+   {
+      auto [num_qp, dim, num_dof] = B.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e(0), num_dof, vdim);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         for (int qp = 0; qp < num_qp; qp++)
+         {
+            double acc = 0.0;
+            for (int dof = 0; dof < num_dof; dof++)
+            {
+               acc += B(qp, 0, dof) * field(dof, vd);
+            }
+            field_qp(vd, qp) = acc;
+         }
+      }
+   }
+   else if constexpr (is_gradient_fop<field_operator_t>::value)
+   {
+      const auto [num_qp, dim, num_dof] = G.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e(0), num_dof, vdim);
+
+      auto f = Reshape(&field_qp[0], vdim, dim, num_qp);
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         for (int qp = 0; qp < num_qp; qp++)
+         {
+            for (int d = 0; d < dim; d++)
+            {
+               double acc = 0.0;
+               for (int dof = 0; dof < num_dof; dof++)
+               {
+                  acc += G(qp, d, dof) * field(dof, vd);
+               }
+               f(vd, d, qp) = acc;
+            }
+         }
+      }
+   }
+   // else if constexpr (std::is_same_v<field_operator_t, FaceNormal>)
+   // {
+   //    auto normal = geometric_factors.normal;
+   //    auto [num_qp, dim, num_entities] = normal.GetShape();
+   //    auto f = Reshape(&field_qp[0], dim, num_qp);
+   //    for (int qp = 0; qp < num_qp; qp++)
+   //    {
+   //       for (int d = 0; d < dim; d++)
+   //       {
+   //          f(d, qp) = normal(qp, d, entity_idx);
+   //       }
+   //    }
+   // }
+   // TODO: Create separate function for clarity
+   else if constexpr (std::is_same_v<field_operator_t, Weight>)
+   {
+      const int num_qp = integration_weights.GetShape()[0];
+      auto f = Reshape(&field_qp[0], num_qp);
+      for (int qp = 0; qp < num_qp; qp++)
+      {
+         f(qp) = integration_weights(qp);
+      }
+   }
+   else if constexpr (is_none_fop<field_operator_t>::value)
+   {
+      auto [num_qp, unused, num_dof] = B.GetShape();
+      const int size_on_qp = input.size_on_qp;
+      const auto field = Reshape(&field_e[0], size_on_qp * num_qp);
+      auto f = Reshape(&field_qp[0], size_on_qp * num_qp);
+      for (int i = 0; i < size_on_qp * num_qp; i++)
+      {
+         f(i) = field(i);
+      }
+   }
+   else
+   {
+      static_assert(always_false<field_operator_t>,
+                    "can't map field to quadrature data");
+
+   }
+}
+
+template <typename field_operator_ts, size_t num_inputs, size_t num_fields>
+MFEM_HOST_DEVICE inline
+void map_fields_to_quadrature_data(
+   std::array<DeviceTensor<2>, num_inputs> &fields_qp,
+   const std::array<DeviceTensor<1>, num_fields> &fields_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   const std::array<int, num_inputs> &input_to_field,
+   const field_operator_ts &fops,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const int &dimension,
+   const bool &use_sum_factorization = false)
+{
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      if (use_sum_factorization)
+      {
+         if (dimension == 2)
+         {
+            map_field_to_quadrature_data_tensor_product_2d(
+               fields_qp[i], dtqmaps[i], fields_e[input_to_field[i]], mfem::get<i>(fops),
+               integration_weights, scratch_mem);
+         }
+         else if (dimension == 3)
+         {
+            map_field_to_quadrature_data_tensor_product_3d(
+               fields_qp[i], dtqmaps[i], fields_e[input_to_field[i]], mfem::get<i>(fops),
+               integration_weights, scratch_mem);
+         }
+         else
+         {
+#if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
+            MFEM_ABORT("unsupported dimension");
+#endif
+         }
+      }
+      else
+      {
+         map_field_to_quadrature_data(
+            fields_qp[i], dtqmaps[i], fields_e[input_to_field[i]], mfem::get<i>(fops),
+            integration_weights);
+      }
+   });
+}
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE
+void map_field_to_quadrature_data_conditional(
+   DeviceTensor<2> &field_qp,
+   const DeviceTensor<1> &field_e,
+   const DofToQuadMap &dtqmap,
+   field_operator_t &fop,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const bool &condition,
+   const int &dimension,
+   const bool &use_sum_factorization = false)
+{
+   if (condition)
+   {
+      if (use_sum_factorization)
+      {
+         if (dimension == 2)
+         {
+            map_field_to_quadrature_data_tensor_product_3d(
+               field_qp, dtqmap, field_e, fop, integration_weights, scratch_mem);
+         }
+         else if (dimension == 3)
+         {
+            map_field_to_quadrature_data_tensor_product_2d(
+               field_qp, dtqmap, field_e, fop, integration_weights, scratch_mem);
+         }
+      }
+      else
+      {
+         map_field_to_quadrature_data(
+            field_qp, dtqmap, field_e, fop, integration_weights);
+      }
+   }
+}
+
+template <size_t num_fields, size_t num_inputs, typename field_operator_ts>
+MFEM_HOST_DEVICE
+void map_fields_to_quadrature_data_conditional(
+   std::array<DeviceTensor<2>, num_inputs> &fields_qp,
+   const std::array<DeviceTensor<1, const double>, num_fields> &fields_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   field_operator_ts fops,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const std::array<bool, num_inputs> &conditions,
+   const bool &use_sum_factorization = false)
+{
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      map_field_to_quadrature_data_conditional(
+         fields_qp[i], fields_e[i], dtqmaps[i], mfem::get<i>(fops), integration_weights,
+         scratch_mem, conditions[i], use_sum_factorization);
+   });
+}
+
+template <size_t num_inputs, typename field_operator_ts>
+MFEM_HOST_DEVICE
+void map_direction_to_quadrature_data_conditional(
+   std::array<DeviceTensor<2>, num_inputs> &directions_qp,
+   const DeviceTensor<1> &direction_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   field_operator_ts fops,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const std::array<bool, num_inputs> &conditions,
+   const int &dimension,
+   const bool &use_sum_factorization = false)
+{
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      if (conditions[i])
+      {
+         if (use_sum_factorization)
+         {
+            if (dimension == 2)
+            {
+               map_field_to_quadrature_data_tensor_product_2d(
+                  directions_qp[i], dtqmaps[i], direction_e, mfem::get<i>(fops),
+                  integration_weights, scratch_mem);
+            }
+            else if (dimension == 3)
+            {
+               map_field_to_quadrature_data_tensor_product_3d(
+                  directions_qp[i], dtqmaps[i], direction_e, mfem::get<i>(fops),
+                  integration_weights, scratch_mem);
+            }
+         }
+         else
+         {
+            map_field_to_quadrature_data(
+               directions_qp[i], dtqmaps[i], direction_e, mfem::get<i>(fops),
+               integration_weights);
+         }
+      }
+   });
+}
+
+}

fem/dfem/parametricspace.hpp

@@ -0,0 +1,126 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+#include "../fe/fe_base.hpp"
+
+namespace mfem
+{
+
+class ParametricSpace
+{
+
+public:
+   /// spatial_dim is the dimension of the spatial domain (e.g. 2 for 2D)
+   /// local_size is the size of the data on a single quadrature point
+   /// element_size is the size of the data on an element divided by vdim
+   /// total_size is the size of the data for all elements
+   ParametricSpace(int spatial_dim, int local_size, int element_size,
+                   int total_size) :
+      spatial_dim(spatial_dim),
+      local_size(local_size),
+      element_size(element_size),
+      total_size(total_size),
+      identity(total_size)
+   {
+      // dtq.ndof = (int)floor(pow(element_size, 1.0/spatial_dim) + 0.5);
+      dtq.ndof = element_size;
+      dtq.nqpt = dtq.ndof;
+   }
+
+   ParametricSpace(int local_size) :
+      local_size(local_size),
+      element_size(local_size),
+      total_size(local_size),
+      identity(local_size)
+   {
+      dtq.ndof = (int)floor(pow(element_size, 1.0/spatial_dim) + 0.5);
+      dtq.nqpt = dtq.ndof;
+   }
+
+   ParametricSpace(int spatial_dim, int local_size, int element_size,
+                   int total_size, int d1d, int q1d) :
+      spatial_dim(spatial_dim),
+      local_size(local_size),
+      element_size(element_size),
+      total_size(total_size),
+      identity(total_size)
+   {
+      dtq.ndof = d1d;
+      dtq.nqpt = q1d;
+   }
+
+   int Dimension() const
+   {
+      return spatial_dim;
+   }
+
+   int GetLocalSize() const
+   {
+      return local_size;
+   }
+
+   int GetElementSize() const
+   {
+      return element_size;
+   }
+
+   int GetTotalSize() const
+   {
+      return total_size;
+   }
+
+   const DofToQuad &GetDofToQuad() const
+   {
+      return dtq;
+   }
+
+   const Operator *GetProlongation() const
+   {
+      return &identity;
+   }
+
+   const Operator *GetRestriction() const
+   {
+      return &identity;
+   }
+
+private:
+   int spatial_dim;
+
+   // Hint for the local dimension. E.g. the size on the quadrature point or vdim.
+   int local_size;
+
+   // Size of the data on an element
+   int element_size;
+
+   int total_size;
+
+   IdentityOperator identity;
+
+   DofToQuad dtq;
+};
+
+class ParametricFunction : public Vector
+{
+public:
+   ParametricFunction(ParametricSpace &space) :
+      Vector(space.GetTotalSize()),
+      space(space)
+   {}
+
+   ParametricSpace &space;
+
+   using Vector::operator=;
+
+};
+
+}

fem/dfem/qfunction.hpp

@@ -0,0 +1,272 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+#include "util.hpp"
+#include "../linalg/tensor.hpp"
+
+namespace mfem
+{
+
+template <typename func_t, typename... arg_ts>
+MFEM_HOST_DEVICE inline
+auto qfunction_wrapper(const func_t &f, arg_ts &&...args)
+{
+   return f(args...);
+}
+
+template <typename T0, typename T1>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(const T0 &, T1 &)
+{
+   static_assert(always_false<T0, T1>,
+                 "process_kf_arg not implemented for arg type");
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1, T> &u,
+   T &arg)
+{
+   arg = u(0);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1, T> &u,
+   internal::tensor<T> &arg)
+{
+   arg(0) = u(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::tensor<T, n> &arg)
+{
+   for (int i = 0; i < n; i++)
+   {
+      arg(i) = u(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::tensor<T, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i) = u((i * m) + j);
+      }
+   }
+}
+
+template <typename arg_type>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(const DeviceTensor<2> &u, arg_type &arg, int qp)
+{
+   const auto u_qp = Reshape(&u(0, qp), u.GetShape()[0]);
+   process_kf_arg(u_qp, arg);
+}
+
+template <size_t num_fields, typename kf_args>
+MFEM_HOST_DEVICE inline
+void process_kf_args(
+   const std::array<DeviceTensor<2>, num_fields> &u,
+   kf_args &args,
+   const int &qp)
+{
+   for_constexpr<mfem::tuple_size<kf_args>::value>([&](auto i)
+   {
+      process_kf_arg(u[i], mfem::get<i>(args), qp);
+      // out << mfem::get<i>(args) << ", ";
+   });
+}
+
+template <typename T0, typename T1>
+MFEM_HOST_DEVICE inline
+Vector process_kf_result(T0, T1)
+{
+   static_assert(always_false<T0, T1>,
+                 "process_kf_result not implemented for result type");
+   return Vector{};
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const double &x)
+{
+   r(0) = x;
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<T> &x)
+{
+   r(0) = x(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<T, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      r(i) = x(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<T, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         r(i + n * j) = x(i, j);
+      }
+   }
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   double &arg)
+{
+   arg = u(0);
+}
+
+template <int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<double, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i) = u((i * m) + j);
+      }
+   }
+}
+
+template <typename qfunc_t, typename args_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel(
+   DeviceTensor<1, double> &f_qp,
+   const qfunc_t &qfunc,
+   args_ts &args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   int qp)
+{
+   process_kf_args(u, args, qp);
+   process_kf_result(f_qp, mfem::get<0>(mfem::apply(qfunc, args)));
+}
+
+#ifdef MFEM_USE_ENZYME
+// Version for active function arguments only
+//
+// This is an Enzyme regression and can be removed in later versions.
+template <typename qfunc_t, typename arg_ts, std::size_t... Is,
+          typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme_indexed(qfunc_t &qfunc, arg_ts &&args,
+                                  arg_ts &&shadow_args,
+                                  std::index_sequence<Is...>,
+                                  inactive_arg_ts &&inactive_args,
+                                  std::index_sequence<>)
+{
+   using qf_return_t = typename create_function_signature<
+                       decltype(&qfunc_t::operator())>::type::return_t;
+   return __enzyme_fwddiff<qf_return_t>(
+             qfunction_wrapper<qfunc_t, decltype(mfem::get<Is>(args))...>, enzyme_const,
+             (void *)&qfunc, enzyme_dup, &mfem::get<Is>(args)..., enzyme_interleave,
+             &mfem::get<Is>(shadow_args)...);
+}
+
+// Interleave function arguments for enzyme
+template <typename qfunc_t, typename arg_ts, std::size_t... Is,
+          typename inactive_arg_ts, std::size_t... Js>
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme_indexed(qfunc_t &qfunc, arg_ts &&args,
+                                  arg_ts &&shadow_args,
+                                  std::index_sequence<Is...>,
+                                  inactive_arg_ts &&inactive_args,
+                                  std::index_sequence<Js...>)
+{
+   using qf_return_t = typename create_function_signature<
+                       decltype(&qfunc_t::operator())>::type::return_t;
+   return __enzyme_fwddiff<qf_return_t>(
+             qfunction_wrapper<qfunc_t, decltype(mfem::get<Is>(args))...,
+             decltype(mfem::get<Js>(inactive_args))...>,
+             enzyme_const, (void *)&qfunc, enzyme_dup, &mfem::get<Is>(args)...,
+             enzyme_const, &mfem::get<Js>(inactive_args)..., enzyme_interleave,
+             &mfem::get<Is>(shadow_args)...);
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme(qfunc_t &qfunc, arg_ts &&args,
+                          arg_ts &&shadow_args,
+                          inactive_arg_ts &&inactive_args)
+{
+   auto arg_indices = std::make_index_sequence<
+                      mfem::tuple_size<std::remove_reference_t<arg_ts>>::value> {};
+
+   auto inactive_arg_indices = std::make_index_sequence<
+                               mfem::tuple_size<std::remove_reference_t<inactive_arg_ts>>::value> {};
+
+   return fwddiff_apply_enzyme_indexed(qfunc, args, shadow_args, arg_indices,
+                                       inactive_args, inactive_arg_indices);
+}
+
+template <typename qfunc_t, typename arg_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel_fwddiff_enzyme(
+   DeviceTensor<1, double> &f_qp,
+   qfunc_t &qfunc,
+   arg_ts &args,
+   arg_ts &shadow_args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &v,
+   int qp_idx)
+{
+   // out << "\nargs: ";
+   process_kf_args(u, args, qp_idx);
+   // out << "\nshadow args: ";
+   process_kf_args(v, shadow_args, qp_idx);
+   // out << "\n";
+   process_kf_result(f_qp,
+                     mfem::get<0>(fwddiff_apply_enzyme(qfunc, args, shadow_args, mfem::tuple<> {})));
+}
+#endif // MFEM_USE_ENZYME
+
+} // namespace mfem

fem/dfem/tuo_rocm_dfem_build_instructions.md

@@ -0,0 +1,49 @@
+export LC_USER=andrej1
+module load rocmcc/6.3.1-cce-19.0.0-magic cmake/3.29.2
+
+export MPICH_CC=amdclang
+export MPICH_CXX=amdclang++
+export ROCM_PATH=/opt/rocm-6.3.1
+export LLVM_DIR=$ROCM_PATH/lib/llvm
+export MPI_DIR=/usr/tce/packages/cray-mpich/cray-mpich-8.1.32-rocmcc-6.3.1-cce-19.0.0-magic
+
+export CMAKE_PREFIX_PATH=$CMAKE_PREFIX_PATH:$ROCM_PATH/lib/cmake/hip:$ROCM_PATH/lib/cmake/hipblas:$ROCM_PATH/lib/cmake/hipblas-common:$ROCM_PATH/lib/cmake/hipsparse:$ROCM_PATH/lib/cmake/rocsparse:$ROCM_PATH/lib/cmake/rocrand
+
+export BASE_DIR=/usr/workspace/$LC_USER/dfem-tuo-magic
+export LOCAL_DIR=/usr/workspace/$LC_USER/dfem-tuo-magic/local
+mkdir -p $LOCAL_DIR
+export PATH=$LOCAL_DIR/bin:$PATH
+cd $BASE_DIR
+
+## Enzyme
+git clone --depth 1 https://github.com/EnzymeAD/Enzyme.git
+pushd Enzyme/enzyme
+CC=amdclang CXX=amdclang++ cmake -B build -DLLVM_DIR=$LLVM_DIR -DCMAKE_INSTALL_PREFIX=$LOCAL_DIR
+cmake --build build -j && cmake --install build
+popd
+
+## hypre
+curl https://github.com/hypre-space/hypre/archive/refs/tags/v2.32.0.tar.gz -o hypre-v2.32.0.tar.gz -L
+tar xzf hypre-v2.32.0.tar.gz
+pushd hypre-2.32.0/src
+CC=mpicc CXX=mpicxx CXXFLAGS="std=c++17 -fPIC" CFLAGS="-fPIC" ROCM_PATH=$ROCM_PATH ./configure --disable-fortran --prefix=$LOCAL_DIR --with-MPI-libs="mpi mpich" --with-MPI-lib-dirs=$MPI_DIR/lib --with-MPI-include=$MPI_DIR/include --enable-shared --with-hip
+make -j install
+popd
+
+## metis
+curl -OL https://github.com/mfem/tpls/raw/gh-pages/parmetis-4.0.3.tar.gz
+tar xzf parmetis-4.0.3.tar.gz
+pushd parmetis-4.0.3
+cmake -B build -DCMAKE_CXX_FLAGS="-fPIC" -DCMAKE_C_FLAGS="-fPIC" -DGKLIB_PATH=$BASE_DIR/parmetis-4.0.3/metis/GKlib -DMETIS_PATH=$BASE_DIR/parmetis-4.0.3/metis -DCMAKE_INSTALL_PREFIX=$LOCAL_DIR -DSHARED=1 -DCMAKE_C_COMPILER=mpicc -DCMAKE_CXX_COMPILER=mpicxx
+cmake --build build -j && cmake --install build
+popd
+pushd parmetis-4.0.3/metis
+cmake -B build -DCMAKE_CXX_FLAGS="-fPIC" -DCMAKE_C_FLAGS="-fPIC" -DGKLIB_PATH=$BASE_DIR/parmetis-4.0.3/metis/GKlib -DCMAKE_INSTALL_PREFIX=$LOCAL_DIR -DSHARED=1 -DCMAKE_C_COMPILER=mpicc -DCMAKE_CXX_COMPILER=mpicxx
+cmake --build build -j && cmake --install build
+popd
+
+git clone https://github.com/mfem/mfem.git
+git switch dfem-phase1-dev
+pushd mfem
+CXX=mpicxx cmake -B build-opt -DCMAKE_BUILD_TYPE=Release -DMFEM_USE_HIP=ON -DCMAKE_HIP_ARCHITECTURES="gfx942" -DCMAKE_HIP_PLATFORM="amd"
+cmake --build build-opt -j

fem/dfem/tuo_rocm_dfem_user.cmake

@@ -0,0 +1,31 @@
+if (NOT CMAKE_BUILD_TYPE)
+  set(CMAKE_BUILD_TYPE "Release" CACHE STRING
+      "Build type: Debug, Release, RelWithDebInfo, or MinSizeRel." FORCE)
+endif()
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+set(CMAKE_CXX_STANDARD 17)
+# set(CMAKE_CXX_FLAGS "--save-temps -Rpass-analysis=kernel-resource-usage -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false")
+
+set(MFEM_PRECISION "double" CACHE STRING
+    "Floating-point precision to use: single, or double")
+
+option(BUILD_SHARED_LIBS "Enable shared library build of MFEM" ON)
+option(MFEM_USE_MPI "Enable MPI parallel build" ON)
+option(MFEM_USE_METIS "Enable METIS usage" ${MFEM_USE_MPI})
+option(MFEM_USE_ENZYME "Enable Enzyme" ON)
+option(MFEM_USE_HIP "Enable HIP" ON)
+
+set(MFEM_MPI_NP 4 CACHE STRING "Number of processes used for MPI tests")
+
+option(MFEM_ENABLE_TESTING ON)
+
+set(HIP_ARCH "gfx942" CACHE STRING "Target HIP architecture.")
+
+# Make sure all dirs are absolute
+set(ENZYME_DIR "/usr/workspace/andrej1/dfem-tuo-magic/local/cmake/Enzyme" CACHE PATH "Path to the Enzyme library.")
+set(HYPRE_DIR "/usr/workspace/andrej1/dfem-tuo-magic/local" CACHE PATH "Path to the hypre library.")
+set(METIS_DIR "/usr/workspace/andrej1/dfem-tuo-magic/local" CACHE PATH "Path to the METIS library.")
+
+set(CMAKE_SKIP_PREPROCESSED_SOURCE_RULES ON) # Skip *.i rules
+set(CMAKE_SKIP_ASSEMBLY_SOURCE_RULES ON) # Skip *.s rules

fem/dfem/tuple.hpp

@@ -0,0 +1,853 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+#pragma once
+
+// This is serac's tuple implementation
+
+#include <utility>
+#include <mfem.hpp>
+
+namespace mfem
+{
+
+/**
+ * @tparam T the types stored in the tuple
+ * @brief This is a class that mimics most of std::tuple's interface,
+ * except that it is usable in CUDA kernels and admits some arithmetic operator overloads.
+ *
+ * see https://en.cppreference.com/w/cpp/utility/tuple for more information about std::tuple
+ */
+template <typename... T>
+struct tuple
+{
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ */
+template <typename T0>
+struct tuple<T0>
+{
+   T0 v0;  ///< The first member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ */
+template <typename T0, typename T1>
+struct tuple<T0, T1>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2>
+struct tuple<T0, T1, T2>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3>
+struct tuple<T0, T1, T2, T3>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4>
+struct tuple<T0, T1, T2, T3, T4>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ * @tparam T5 The sixth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5>
+struct tuple<T0, T1, T2, T3, T4, T5>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ * @tparam T5 The sixth type stored in the tuple
+ * @tparam T6 The seventh type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
+struct tuple<T0, T1, T2, T3, T4, T5, T6>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+   T6 v6;  ///< The seventh member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ * @tparam T5 The sixth type stored in the tuple
+ * @tparam T6 The seventh type stored in the tuple
+ * @tparam T7 The eighth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7>
+struct tuple<T0, T1, T2, T3, T4, T5, T6, T7>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+   T6 v6;  ///< The seventh member of the tuple
+   T7 v7;  ///< The eighth member of the tuple
+};
+
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8>
+struct tuple<T0, T1, T2, T3, T4, T5, T6, T7, T8>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+   T6 v6;  ///< The seventh member of the tuple
+   T7 v7;  ///< The eighth member of the tuple
+   T8 v8;
+};
+
+/**
+ * @brief Class template argument deduction rule for tuples
+ * @tparam T The variadic template parameter for tuple types
+ */
+template <typename... T>
+MFEM_HOST_DEVICE
+tuple(T...) -> tuple<T...>;
+
+/**
+ * @brief helper function for combining a list of values into a tuple
+ * @tparam T types of the values to be tuple-d
+ * @param args the actual values to be put into a tuple
+ */
+template <typename... T>
+MFEM_HOST_DEVICE tuple<T...> make_tuple(const T&... args)
+{
+   return tuple<T...> {args...};
+}
+
+template <class... Types>
+struct tuple_size
+{
+};
+
+template <class... Types>
+struct tuple_size<mfem::tuple<Types...>> :
+                                      std::integral_constant<std::size_t, sizeof...(Types)>
+{
+};
+
+/**
+ * @tparam i the tuple index to access
+ * @tparam T the types stored in the tuple
+ * @brief return a reference to the ith tuple entry
+ */
+template <int i, typename... T>
+MFEM_HOST_DEVICE constexpr auto& get(tuple<T...>& values)
+{
+   static_assert(i < sizeof...(T), "");
+   if constexpr (i == 0)
+   {
+      return values.v0;
+   }
+   if constexpr (i == 1)
+   {
+      return values.v1;
+   }
+   if constexpr (i == 2)
+   {
+      return values.v2;
+   }
+   if constexpr (i == 3)
+   {
+      return values.v3;
+   }
+   if constexpr (i == 4)
+   {
+      return values.v4;
+   }
+   if constexpr (i == 5)
+   {
+      return values.v5;
+   }
+   if constexpr (i == 6)
+   {
+      return values.v6;
+   }
+   if constexpr (i == 7)
+   {
+      return values.v7;
+   }
+   if constexpr (i == 8)
+   {
+      return values.v8;
+   }
+}
+
+/**
+ * @tparam i the tuple index to access
+ * @tparam T the types stored in the tuple
+ * @brief return a copy of the ith tuple entry
+ */
+template <int i, typename... T>
+MFEM_HOST_DEVICE constexpr const auto& get(const tuple<T...>& values)
+{
+   static_assert(i < sizeof...(T), "");
+   if constexpr (i == 0)
+   {
+      return values.v0;
+   }
+   if constexpr (i == 1)
+   {
+      return values.v1;
+   }
+   if constexpr (i == 2)
+   {
+      return values.v2;
+   }
+   if constexpr (i == 3)
+   {
+      return values.v3;
+   }
+   if constexpr (i == 4)
+   {
+      return values.v4;
+   }
+   if constexpr (i == 5)
+   {
+      return values.v5;
+   }
+   if constexpr (i == 6)
+   {
+      return values.v6;
+   }
+   if constexpr (i == 7)
+   {
+      return values.v7;
+   }
+   if constexpr (i == 8)
+   {
+      return values.v8;
+   }
+}
+
+/**
+ * @brief a function intended to be used for extracting the ith type from a tuple.
+ *
+ * @note type<i>(my_tuple) returns a value, whereas get<i>(my_tuple) returns a reference
+ *
+ * @tparam i the index of the tuple to query
+ * @tparam T the types stored in the tuple
+ * @param values the tuple of values
+ * @return a copy of the ith entry of the input
+ */
+template <int i, typename... T>
+MFEM_HOST_DEVICE constexpr auto type(const tuple<T...>& values)
+{
+   static_assert(i < sizeof...(T), "");
+   if constexpr (i == 0)
+   {
+      return values.v0;
+   }
+   if constexpr (i == 1)
+   {
+      return values.v1;
+   }
+   if constexpr (i == 2)
+   {
+      return values.v2;
+   }
+   if constexpr (i == 3)
+   {
+      return values.v3;
+   }
+   if constexpr (i == 4)
+   {
+      return values.v4;
+   }
+   if constexpr (i == 5)
+   {
+      return values.v5;
+   }
+   if constexpr (i == 6)
+   {
+      return values.v6;
+   }
+   if constexpr (i == 7)
+   {
+      return values.v7;
+   }
+   if constexpr (i == 8)
+   {
+      return values.v8;
+   }
+}
+
+/**
+ * @brief A helper function for the + operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple sum
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto plus_helper(const tuple<S...>& x,
+                                            const tuple<T...>& y,
+                                            std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) + get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise sum of x and y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator+(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return plus_helper(x, y,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the += operator of tuples
+ *
+ * @tparam T the types stored in the tuples x and y
+ * @tparam i integer sequence used to index the tuples
+ * @param x tuple of values to be incremented
+ * @param y tuple of increment values
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr void plus_equals_helper(tuple<T...>& x,
+                                                   const tuple<T...>& y,
+                                                   std::integer_sequence<int, i...>)
+{
+   ((get<i>(x) += get<i>(y)), ...);
+}
+
+/**
+ * @tparam T the types stored in the tuples x and y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief add values contained in y, to the tuple x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator+=(tuple<T...>& x,
+                                           const tuple<T...>& y)
+{
+   return plus_equals_helper(x, y,
+                             std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the -= operator of tuples
+ *
+ * @tparam T the types stored in the tuples x and y
+ * @tparam i integer sequence used to index the tuples
+ * @param x tuple of values to be subracted from
+ * @param y tuple of values to subtract from x
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr void minus_equals_helper(tuple<T...>& x,
+                                                    const tuple<T...>& y,
+                                                    std::integer_sequence<int, i...>)
+{
+   ((get<i>(x) -= get<i>(y)), ...);
+}
+
+/**
+ * @tparam T the types stored in the tuples x and y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief add values contained in y, to the tuple x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator-=(tuple<T...>& x,
+                                           const tuple<T...>& y)
+{
+   return minus_equals_helper(x, y,
+                              std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the - operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple difference
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto minus_helper(const tuple<S...>& x,
+                                             const tuple<T...>& y,
+                                             std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) - get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise difference of x and y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator-(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return minus_helper(x, y,
+                       std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the - operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @return the returned tuple difference
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto unary_minus_helper(const tuple<T...>& x,
+                                                   std::integer_sequence<int, i...>)
+{
+   return tuple{-get<i>(x)...};
+}
+
+/**
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @brief return a tuple of values defined by applying the unary minus operator to each element of x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator-(const tuple<T...>& x)
+{
+   return unary_minus_helper(x,
+                             std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the / operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple ratio
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto div_helper(const tuple<S...>& x,
+                                           const tuple<T...>& y,
+                                           std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) / get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise division of x by y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator/(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return div_helper(x, y,
+                     std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the / operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a the constant numerator
+ * @return the returned tuple ratio
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto div_helper(const double a,
+                                           const tuple<T...>& x, std::integer_sequence<int, i...>)
+{
+   return tuple{a / get<i>(x)...};
+}
+
+/**
+ * @brief A helper function for the / operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a the constant denomenator
+ * @return the returned tuple ratio
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto div_helper(const tuple<T...>& x,
+                                           const double a, std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) / a...};
+}
+
+/**
+ * @tparam T the types stored in the tuple x
+ * @param a the numerator
+ * @param x a tuple of denominator values
+ * @brief return a tuple of values defined by division of a by the elements of x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator/(const double a, const tuple<T...>& x)
+{
+   return div_helper(a, x,
+                     std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of numerator values
+ * @param a a denominator
+ * @brief return a tuple of values defined by elementwise division of x by a
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator/(const tuple<T...>& x, const double a)
+{
+   return div_helper(x, a,
+                     std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the * operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple product
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto mult_helper(const tuple<S...>& x,
+                                            const tuple<T...>& y,
+                                            std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) * get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise multiplication of x and y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator*(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return mult_helper(x, y,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the * operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a a constant multiplier
+ * @return the returned tuple product
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto mult_helper(const double a,
+                                            const tuple<T...>& x, std::integer_sequence<int, i...>)
+{
+   return tuple{a * get<i>(x)...};
+}
+
+/**
+ * @brief A helper function for the * operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a a constant multiplier
+ * @return the returned tuple product
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto mult_helper(const tuple<T...>& x,
+                                            const double a, std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) * a...};
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @param a a scaling factor
+ * @param x the tuple object
+ * @brief multiply each component of x by the value a on the left
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator*(const double a, const tuple<T...>& x)
+{
+   return mult_helper(a, x,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @param x the tuple object
+ * @param a a scaling factor
+ * @brief multiply each component of x by the value a on the right
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator*(const tuple<T...>& x, const double a)
+{
+   return mult_helper(x, a,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @tparam i a list of indices used to acces each element of the tuple
+ * @param out the ostream to write the output to
+ * @param A the tuple of values
+ * @brief helper used to implement printing a tuple of values
+ */
+template <typename... T, std::size_t... i>
+auto& print_helper(std::ostream& out, const mfem::tuple<T...>& A,
+                   std::integer_sequence<size_t, i...>)
+{
+   out << "tuple{";
+   (..., (out << (i == 0 ? "" : ", ") << mfem::get<i>(A)));
+   out << "}";
+   return out;
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @param out the ostream to write the output to
+ * @param A the tuple of values
+ * @brief print a tuple of values
+ */
+template <typename... T>
+auto& operator<<(std::ostream& out, const mfem::tuple<T...>& A)
+{
+   return print_helper(out, A, std::make_integer_sequence<size_t, sizeof...(T)>());
+}
+
+/**
+ * @brief A helper to apply a lambda to a tuple
+ *
+ * @tparam lambda The functor type
+ * @tparam T The tuple types
+ * @tparam i The integer sequence to i
+ * @param f The functor to apply to the tuple
+ * @param args The input tuple
+ * @return The functor output
+ */
+template <typename lambda, typename... T, int... i>
+MFEM_HOST_DEVICE auto apply_helper(lambda f, tuple<T...>& args,
+                                   std::integer_sequence<int, i...>)
+{
+   return f(get<i>(args)...);
+}
+
+/**
+ * @tparam lambda a callable type
+ * @tparam T the types of arguments to be passed in to f
+ * @param f the callable object
+ * @param args a tuple of arguments
+ * @brief a way of passing an n-tuple to a function that expects n separate arguments
+ *
+ *   e.g. foo(bar, baz) is equivalent to apply(foo, mfem::tuple(bar,baz));
+ */
+template <typename lambda, typename... T>
+MFEM_HOST_DEVICE auto apply(lambda f, tuple<T...>& args)
+{
+   return apply_helper(f, std::move(args),
+                       std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @overload
+ */
+template <typename lambda, typename... T, int... i>
+MFEM_HOST_DEVICE auto apply_helper(lambda f, const tuple<T...>& args,
+                                   std::integer_sequence<int, i...>)
+{
+   return f(get<i>(args)...);
+}
+
+/**
+ * @tparam lambda a callable type
+ * @tparam T the types of arguments to be passed in to f
+ * @param f the callable object
+ * @param args a tuple of arguments
+ * @brief a way of passing an n-tuple to a function that expects n separate arguments
+ *
+ *   e.g. foo(bar, baz) is equivalent to apply(foo, mfem::tuple(bar,baz));
+ */
+template <typename lambda, typename... T>
+MFEM_HOST_DEVICE auto apply(lambda f, const tuple<T...>& args)
+{
+   return apply_helper(f, std::move(args),
+                       std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief a struct used to determine the type at index I of a tuple
+ *
+ * @note see: https://en.cppreference.com/w/cpp/utility/tuple/tuple_element
+ *
+ * @tparam I the index of the desired type
+ * @tparam T a tuple of different types
+ */
+template <size_t I, class T>
+struct tuple_element;
+
+// recursive case
+/// @overload
+template <size_t I, class Head, class... Tail>
+struct tuple_element<I, tuple<Head, Tail...>> : tuple_element<I - 1,
+                                              tuple<Tail...>>
+{
+};
+
+// base case
+/// @overload
+template <class Head, class... Tail>
+struct tuple_element<0, tuple<Head, Tail...>>
+{
+   using type = Head;  ///< the type at the specified index
+};
+
+/**
+ * @brief Trait for checking if a type is a @p mfem::tuple
+ */
+template <typename T>
+struct is_tuple : std::false_type
+{
+};
+
+/// @overload
+template <typename... T>
+struct is_tuple<mfem::tuple<T...>> : std::true_type
+{
+};
+
+/**
+ * @brief Trait for checking if a type if a @p mfem::tuple containing only @p mfem::tuple
+ */
+template <typename T>
+struct is_tuple_of_tuples : std::false_type
+{
+};
+
+/**
+ * @brief Trait for checking if a type if a @p mfem::tuple containing only @p mfem::tuple
+ */
+template <typename... T>
+struct is_tuple_of_tuples<mfem::tuple<T...>>
+{
+   static constexpr bool value = (is_tuple<T>::value &&
+                                  ...);  ///< true/false result of type check
+};
+
+}  // namespace mfem

fem/pfespace.cpp

@@ -23,6 +23,10 @@
 #include <limits>
 #include <list>
 
+#undef NVTX_COLOR
+#define NVTX_COLOR nvtx::kLavender
+#include "general/nvtx.hpp"
+
 namespace mfem
 {
 
@@ -1215,6 +1219,7 @@ const Operator *ParFiniteElementSpace::GetProlongationMatrix() const
 
       if (nd_strias) { return Dof_TrueDof_Matrix(); }
 
+      dbg();
       if (NRanks == 1)
       {
          Pconf = new IdentityOperator(GetTrueVSize());
@@ -1234,6 +1239,7 @@ const Operator *ParFiniteElementSpace::GetProlongationMatrix() const
    }
    else
    {
+      assert(false);
       return Dof_TrueDof_Matrix();
    }
 }
@@ -3646,6 +3652,7 @@ ConformingProlongationOperator::ConformingProlongationOperator(
 
 void ConformingProlongationOperator::Mult(const Vector &x, Vector &y) const
 {
+   dbg();
    MFEM_ASSERT(x.Size() == Width(), "");
    MFEM_ASSERT(y.Size() == Height(), "");
 

general/enzyme.hpp

@@ -9,6 +9,7 @@
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
 
+#include "../config/config.hpp"
 
 #ifndef MFEM_ENZYME_HPP
 #define MFEM_ENZYME_HPP
@@ -25,11 +26,27 @@ extern int enzyme_dup;
 extern int enzyme_dupnoneed;
 extern int enzyme_out;
 extern int enzyme_const;
+extern int enzyme_interleave;
+
+#if defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP)
+#define MFEM_DEVICE_EXTERN_STMT(name) extern __device__ int name;
+#else
+#define MFEM_DEVICE_EXTERN_STMT(name)
+#endif
+
+MFEM_DEVICE_EXTERN_STMT(enzyme_dup)
+MFEM_DEVICE_EXTERN_STMT(enzyme_dupnoneed)
+MFEM_DEVICE_EXTERN_STMT(enzyme_out)
+MFEM_DEVICE_EXTERN_STMT(enzyme_const)
+MFEM_DEVICE_EXTERN_STMT(enzyme_interleave)
 
 template <typename return_type, typename... Args>
+MFEM_HOST_DEVICE inline
 return_type __enzyme_autodiff(Args...);
 
+// warning: if inlined, triggers function '__enzyme_fwddiff' is not defined
 template <typename return_type, typename... Args>
+MFEM_HOST_DEVICE
 return_type __enzyme_fwddiff(Args...);
 
 #define MFEM_ENZYME_INACTIVENOFREE   __attribute__((enzyme_inactive, enzyme_nofree))

general/nvtx.hpp

@@ -0,0 +1,471 @@
+#pragma once
+
+#include <fmt/format.h>
+
+#include <array>
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <iomanip>
+#include <iostream>
+#include <memory>
+#include <mutex>
+#include <stack>
+#include <string>
+
+#include "../config/config.hpp"
+
+#ifdef MFEM_USE_MPI
+#include <mpi.h>
+#endif
+
+#ifdef MFEM_USE_CALIPER
+#include <caliper/cali.h>
+#endif
+
+#ifdef MFEM_USE_CUDA
+#include <cudaProfiler.h>
+#include <cuda_runtime_api.h>
+#include <nvToolsExt.h>
+#else
+struct nvtxEventAttributes_t
+{
+   int version;
+   int size;
+   int category;
+   int colorType;
+   uint32_t color;
+   int payloadType;
+   uint64_t payload;
+   int messageType;
+   struct
+   {
+      std::string ascii;
+   } message;
+};
+#define NVTX_VERSION 1
+#define NVTX_EVENT_ATTRIB_STRUCT_SIZE 256
+#define NVTX_COLOR_ARGB 0
+#define NVTX_MESSAGE_TYPE_ASCII 0
+#define nvtxRangePushEx(...)
+#define nvtxRangePop(...)
+#define cudaStreamSynchronize(...)
+#endif
+
+namespace nvtx
+{
+
+///////////////////////////////////////////////////////////////////////////////
+// https://en.wikipedia.org/wiki/Web_colors#Extended_colors
+// clang-format off
+enum color_names
+{
+   kBlack = 0, kNavyBlue, kDarkBlue, kMediumBlue, kBlue, kDarkGreen, kWebGreen, kTeal,
+   kDarkCyan, kDeepSkyBlue, kDarkTurquoise, kMediumSpringGreen, kGreen, kLime,
+   kSpringGreen, kAqua, kCyan, kMidnightBlue, kDodgerBlue, kLightSeaGreen, kForestGreen,
+   kSeaGreen, kDarkSlateGray, kLimeGreen, kMediumSeaGreen, kTurquoise, kRoyalBlue,
+   kSteelBlue, kDarkSlateBlue, kMediumTurquoise, kIndigo, kDarkOliveGreen, kCadetBlue,
+   kCornflower, kRebeccaPurple, kMediumAquamarine, kDimGray, kSlateBlue, kOliveDrab,
+   kSlateGray, kLightSlateGray, kMediumSlateBlue, kLawnGreen, kWebMaroon, kWebPurple,
+   kChartreuse, kAquamarine, kOlive, kWebGray, kSkyBlue, kLightSkyBlue, kBlueViolet,
+   kDarkRed, kDarkMagenta, kSaddleBrown, kDarkSeaGreen, kLightGreen, kMediumPurple,
+   kDarkViolet, kPaleGreen, kDarkOrchid, kYellowGreen, kPurple, kSienna, kBrown,
+   kDarkGray, kLightBlue, kGreenYellow, kPaleTurquoise, kMaroon, kLightSteelBlue,
+   kPowderBlue, kFirebrick, kDarkGoldenrod, kMediumOrchid, kRosyBrown, kDarkKhaki,
+   kGray, kSilver, kMediumVioletRed, kIndianRed, kPeru, kChocolate, kTan, kLightGray,
+   kThistle, kOrchid, kGoldenrod, kPaleVioletRed, kCrimson, kGainsboro, kPlum, kBurlywood,
+   kLightCyan, kLavender, kDarkSalmon, kViolet, kPaleGoldenrod, kLightCoral, kKhaki,
+   kAliceBlue, kHoneydew, kAzure, kSandyBrown, kWheat, kBeige, kWhiteSmoke, kMintCream,
+   kGhostWhite, kSalmon, kAntiqueWhite, kLinen, kLightGoldenrod, kOldLace, kRed,
+   kFuchsia, kMagenta, kDeepPink, kOrangeRed, kTomato, kHotPink, kCoral, kDarkOrange,
+   kLightSalmon, kOrange, kLightPink, kPink, kGold, kPeachPuff, kNavajoWhite, kMoccasin,
+   kBisque, kMistyRose, kBlanchedAlmond, kPapayaWhip, kLavenderBlush, kSeashell,
+   kCornsilk, kLemonChiffon, kFloralWhite, kSnow, kYellow, kLightYellow, kIvory, kWhite,
+   kNvidia
+};
+// clang-format on
+
+static constexpr int kNumHexColors = 146;
+static constexpr std::array<uint32_t, kNumHexColors> kHexColors =
+{
+   {
+      0x000000, 0x000080, 0x00008B, 0x0000CD, 0x0000FF, 0x006400, 0x008000,
+      0x008080, 0x008B8B, 0x00BFFF, 0x00CED1, 0x00FA9A, 0x00FF00, 0x00FF00,
+      0x00FF7F, 0x00FFFF, 0x00FFFF, 0x191970, 0x1E90FF, 0x20B2AA, 0x228B22,
+      0x2E8B57, 0x2F4F4F, 0x32CD32, 0x3CB371, 0x40E0D0, 0x4169E1, 0x4682B4,
+      0x483D8B, 0x48D1CC, 0x4B0082, 0x556B2F, 0x5F9EA0, 0x6495ED, 0x663399,
+      0x66CDAA, 0x696969, 0x6A5ACD, 0x6B8E23, 0x708090, 0x778899, 0x7B68EE,
+      0x7CFC00, 0x7F0000, 0x7F007F, 0x7FFF00, 0x7FFFD4, 0x808000, 0x808080,
+      0x87CEEB, 0x87CEFA, 0x8A2BE2, 0x8B0000, 0x8B008B, 0x8B4513, 0x8FBC8F,
+      0x90EE90, 0x9370DB, 0x9400D3, 0x98FB98, 0x9932CC, 0x9ACD32, 0xA020F0,
+      0xA0522D, 0xA52A2A, 0xA9A9A9, 0xADD8E6, 0xADFF2F, 0xAFEEEE, 0xB03060,
+      0xB0C4DE, 0xB0E0E6, 0xB22222, 0xB8860B, 0xBA55D3, 0xBC8F8F, 0xBDB76B,
+      0xBEBEBE, 0xC0C0C0, 0xC71585, 0xCD5C5C, 0xCD853F, 0xD2691E, 0xD2B48C,
+      0xD3D3D3, 0xD8BFD8, 0xDA70D6, 0xDAA520, 0xDB7093, 0xDC143C, 0xDCDCDC,
+      0xDDA0DD, 0xDEB887, 0xE0FFFF, 0xE6E6FA, 0xE9967A, 0xEE82EE, 0xEEE8AA,
+      0xF08080, 0xF0E68C, 0xF0F8FF, 0xF0FFF0, 0xF0FFFF, 0xF4A460, 0xF5DEB3,
+      0xF5F5DC, 0xF5F5F5, 0xF5FFFA, 0xF8F8FF, 0xFA8072, 0xFAEBD7, 0xFAF0E6,
+      0xFAFAD2, 0xFDF5E6, 0xFF0000, 0xFF00FF, 0xFF00FF, 0xFF1493, 0xFF4500,
+      0xFF6347, 0xFF69B4, 0xFF7F50, 0xFF8C00, 0xFFA07A, 0xFFA500, 0xFFB6C1,
+      0xFFC0CB, 0xFFD700, 0xFFDAB9, 0xFFDEAD, 0xFFE4B5, 0xFFE4C4, 0xFFE4E1,
+      0xFFEBCD, 0xFFEFD5, 0xFFF0F5, 0xFFF5EE, 0xFFF8DC, 0xFFFACD, 0xFFFAF0,
+      0xFFFAFA, 0xFFFF00, 0xFFFFE0, 0xFFFFF0, 0xFFFFFF, 0x76B900
+   }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+inline size_t static_strlen(const char *str)
+{
+   return *str == '\0' ? 0 : static_strlen(str + 1) + 1;
+}
+
+inline uint8_t static_checksum8(const char *bfr)
+{
+   unsigned int chk = 0;
+   size_t len = static_strlen(bfr);
+   for (; len; len--, bfr++) { chk += static_cast<unsigned int>(*bfr); }
+   return static_cast<uint8_t>(chk);
+}
+
+inline char *static_strrnchr(const char *str, const char c, int n)
+{
+   size_t len = static_strlen(str);
+   char *p = const_cast<char *>(str) + len - 1;
+   for (; n; n--, p--, len--)
+   {
+      for (; len; p--, len--)
+      {
+         if (*p == c) { break; }
+      }
+      if (!len) { return nullptr; }
+      if (n == 1) { return p; }
+   }
+   return nullptr;
+}
+
+inline uint32_t static_color(const uint8_t COLOR, const int RANK,
+                             const char *FILE)
+{
+   constexpr auto kMpiColorShift = 1;
+   const auto rank_shift = kMpiColorShift * RANK;
+   if (COLOR > 0) { return kHexColors[COLOR + rank_shift]; }
+   const auto file_color = static_checksum8(FILE);
+   return kHexColors[(file_color + rank_shift) % kNumHexColors];
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Helpers to generate unique variable names
+#define NVTX_FLF __FILE__, __LINE__, __FUNCTION__
+#define NVTX_PRIVATE_NAME(prefix) NVTX_PRIVATE_CONCAT(prefix, __LINE__)
+#define NVTX_PRIVATE_CONCAT(a, b) NVTX_PRIVATE_CONCAT2(a, b)
+#define NVTX_PRIVATE_CONCAT2(a, b) a##b
+
+#ifndef NVTX_COLOR
+#define NVTX_COLOR ::nvtx::kBlack
+#endif
+
+///////////////////////////////////////////////////////////////////////////////
+struct Debug
+{
+   const bool debug = false, end = true;
+
+   inline Debug() = default;
+
+   inline Debug(const int RANK, const char *FILE, const int LINE,
+                const char *FUNC, uint8_t COLOR, bool ini = true,
+                bool END = true): debug(true), end(END)
+   {
+      const char *base = static_strrnchr(FILE, '/', 2);
+      const char *file = base ? base + 1 : FILE;
+      const uint32_t rgb = static_color(COLOR, RANK, FILE);
+      const uint8_t r = (rgb >> 16) & 0xFF, g = (rgb >> 8) & 0xFF,
+                    b = rgb & 0xFF;
+      std::cout << "\033[38;2;";
+      std::cout << std::to_string(r) << ";";
+      std::cout << std::to_string(g) << ";";
+      std::cout << std::to_string(b) << "m";
+      if (ini)
+      {
+         std::cout << RANK << std::setw(64) << file << ":";
+         std::cout << "\033[2m" << std::setw(4) << std::left << LINE
+                   << "\033[22m: ";
+         if (FUNC) { std::cout << "[" << FUNC << "] "; }
+      }
+      std::cout << std::right << "\033[1m";
+   }
+
+   inline ~Debug()
+   {
+      if (debug) { std::cout << "\033[m" << (end ? "\n" : "") << std::flush; }
+   }
+
+   template <typename T>
+   inline void operator<<(const T &arg) const noexcept
+   {
+      if (debug) { std::cout << arg; }
+   }
+
+   template <typename T>
+   inline void operator()(const T &arg) const noexcept
+   {
+      if (debug) { this->operator<<(arg); }
+   }
+
+   template <typename... Args>
+   inline void operator()(const char *fmt, Args &&...args) const noexcept
+   {
+      // if (debug) { std::cout << fmt::format(fmt, std::forward<Args>(args)...); }
+      if (debug) { std::cout << fmt::format(fmt::runtime(fmt), std::forward<Args>(args)...); }
+   }
+
+   inline void operator()() const noexcept {}
+
+   static Debug Set(const char *FILE, const int LINE, const char *FUNC,
+                    uint8_t COLOR, bool INI = true, bool END = true)
+   {
+      static int mpi_rank = 0, dbg_mpi_rank = 0;
+      static bool env_mpi = false, env_dbg = false;
+      static bool ini = false;
+      if (!ini)
+      {
+         env_dbg = (::getenv("MFEM_DEBUG") != nullptr);
+         env_mpi = ::getenv("MFEM_DEBUG_MPI") != nullptr;
+         int mpi_flag = 0;
+         MPI_Initialized(&mpi_flag);
+         if (mpi_flag) { MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank); }
+         dbg_mpi_rank = atoi(env_mpi ? ::getenv("MFEM_DEBUG_MPI") : "0");
+         ini = true;
+      }
+      const bool debug = (env_dbg && (!env_mpi || (dbg_mpi_rank == mpi_rank)));
+      return debug ? Debug(mpi_rank, FILE, LINE, FUNC, COLOR, INI, END)
+             : Debug();
+   }
+};
+
+// Debug console traces, unnamed
+#define NVTX_DEBUG(...)                                                  \
+   ::nvtx::Debug::Set(NVTX_FLF, NVTX_COLOR).operator()(__VA_ARGS__)
+
+#define NVTX_DEBUG_NO_INI(...)                                \
+   ::nvtx::Debug::Set(NVTX_FLF, NVTX_COLOR, false, true) \
+      .operator()(__VA_ARGS__)
+
+#define NVTX_DEBUG_APPEND(...)                                 \
+   ::nvtx::Debug::Set(NVTX_FLF, NVTX_COLOR, false, false) \
+      .operator()(__VA_ARGS__)
+
+#define NVTX_DEBUG_NO_END(...)                                \
+   ::nvtx::Debug::Set(NVTX_FLF, NVTX_COLOR, true, false) \
+      .operator()(__VA_ARGS__)
+
+///////////////////////////////////////////////////////////////////////////////
+struct Nvtx
+{
+   const bool nvtx = false, enforce_kernel_sync = false;
+   const char *base, *file;
+   const uint32_t color = kBlack;
+   mutable std::string ascii;
+   mutable nvtxEventAttributes_t event;
+   mutable bool pushed = false;
+
+   inline Nvtx() = default;
+
+   Nvtx(bool enforce_kernel_sync, const char *FILE, const int LINE,
+        const char *FUNC, uint8_t COLOR):
+      nvtx(true), enforce_kernel_sync(enforce_kernel_sync),
+      base(static_strrnchr(FILE, '/', 2)), file(base ? base + 1 : FILE),
+      color(COLOR), ascii(file), event({})
+   {
+      event.version = NVTX_VERSION;
+      event.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+      event.colorType = NVTX_COLOR_ARGB;
+      event.color = static_color(COLOR, 0, FILE);
+      event.messageType = NVTX_MESSAGE_TYPE_ASCII;
+
+      ascii += ":";
+      ascii += std::to_string(LINE);
+      ascii += ":[";
+      ascii += FUNC;
+      ascii += "] ";
+
+      pushed = false;
+   }
+
+   explicit Nvtx(const char *title, uint8_t color = kWheat,
+                 bool enforce_kernel_sync = true):
+      nvtx(true), enforce_kernel_sync(enforce_kernel_sync), color(color),
+      ascii(title), event({})
+   {
+      event.version = NVTX_VERSION;
+      event.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+      event.colorType = NVTX_COLOR_ARGB;
+      event.color = static_color(color, 0, "");
+      event.messageType = NVTX_MESSAGE_TYPE_ASCII;
+      event.message.ascii = ascii.c_str();
+      nvtxRangePushEx(&event);
+      pushed = true;
+   }
+
+   inline void operator()() const
+   {
+      if (!nvtx) { return; }
+      event.message.ascii = ascii.c_str();
+      assert(!pushed);
+      nvtxRangePushEx(&event);
+      pushed = true;
+   }
+
+   template <typename T>
+   inline void operator()(const T &arg) const
+   {
+      if (!nvtx) { return; }
+      this->operator<<(arg);
+      event.message.ascii = ascii.c_str();
+      assert(!pushed);
+      nvtxRangePushEx(&event);
+      pushed = true;
+   }
+
+   template <typename... Args>
+   inline void operator()(fmt::format_string<Args...> fmt,
+                          Args &&...args) const
+   {
+      if (!nvtx) { return; }
+      ascii += fmt::format(fmt, std::forward<Args>(args)...);
+      event.message.ascii = ascii.c_str();
+      assert(!pushed);
+      nvtxRangePushEx(&event);
+      pushed = true;
+   }
+
+   template <typename T>
+   inline void operator<<(const T &arg) const
+   {
+      if (nvtx) { ascii += arg; }
+   }
+
+   inline ~Nvtx()
+   {
+      if (!nvtx) { return; }
+      if (enforce_kernel_sync)
+      {
+         nvtxEventAttributes_t eks = {};
+         eks.version = NVTX_VERSION;
+         eks.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+         eks.category = 0; // user value
+         eks.colorType = NVTX_COLOR_ARGB;
+         eks.messageType = NVTX_MESSAGE_TYPE_ASCII;
+         eks.message.ascii = "!"; // enforce kernel synchronization
+         eks.color = kHexColors[kYellow];
+         nvtxRangePushEx(&eks);
+         cudaStreamSynchronize(nullptr);
+         nvtxRangePop(/*eks*/);
+      }
+      assert(pushed);
+      nvtxRangePop(/*event*/);
+   }
+
+   using nvtx_ptr = std::unique_ptr<Nvtx>;
+   using nvtx_stack_t = std::stack<nvtx_ptr>;
+
+   static nvtx_ptr Set(const char *FILE, const int LINE, const char *FUNC,
+                       uint8_t COLOR)
+   {
+      static bool nvtx = false, eks = false;
+      static bool ini = false;
+      if (!ini)
+      {
+         eks = ::getenv("MFEM_EKS") != nullptr;
+         nvtx = ::getenv("MFEM_NVTX") != nullptr;
+         Nvtx force_first_eks("Init EKS", kYellow, true);
+         ini = true;
+      }
+      return nvtx_ptr(nvtx ? new Nvtx(eks, FILE, LINE, FUNC, COLOR)
+                      : new Nvtx());
+   }
+
+   static nvtx_stack_t &Stack()
+   {
+      auto nvtx_events = []() -> nvtx_stack_t &
+      {
+         static nvtx_stack_t events;
+         return events;
+      };
+      static std::once_flag ready;
+      // one touch to guarantee the object is ready
+      std::call_once(ready, [&] { nvtx_events(); });
+      return nvtx_events();
+   }
+};
+
+// Temporary object only alive for the current statement
+#define NVTX_(COLOR, ...)                                       \
+   NVTX_DEBUG(__VA_ARGS__);                                     \
+   std::unique_ptr<::nvtx::Nvtx> NVTX_PRIVATE_NAME(nvtx) = \
+      ::nvtx::Nvtx::Set(NVTX_FLF, COLOR);                  \
+   NVTX_PRIVATE_NAME(nvtx)->operator()(__VA_ARGS__)
+
+// Temporary object only alive for the current statement
+#define NVTX(...) NVTX_(NVTX_COLOR, __VA_ARGS__)
+
+// Begin(with color)/End NVTX event traces
+#define NVTX_BEGIN_(COLOR, ...)                                              \
+   NVTX_DEBUG(__VA_ARGS__);                                                  \
+   ::nvtx::Nvtx::Stack().push(::nvtx::Nvtx::Set(NVTX_FLF, COLOR)); \
+   ::nvtx::Nvtx::Stack().top()->operator()(__VA_ARGS__)
+
+// Begin/End NVTX event traces
+#define NVTX_BEGIN(...) NVTX_BEGIN_(NVTX_COLOR, __VA_ARGS__);
+
+#define NVTX_END(...)                        \
+   ::nvtx::Nvtx::Stack().top().reset(); \
+   ::nvtx::Nvtx::Stack().pop()
+
+#ifdef USE_CALIPER
+// CALIPER & NVTX marks
+#define NVTX_MARK_FUNCTION                                            \
+      NVTX();                                                            \
+      std::unique_ptr<cali::Function> __cali_ann##__func__;              \
+      __cali_ann##__func__ = std::make_unique<cali::Function>(__func__);
+
+#define NVTX_MARK(...)                                                   \
+      NVTX(__VA_ARGS__);                                                    \
+      std::unique_ptr<cali::Function> __cali_ann##__func__;                 \
+      __cali_ann##__func__ = std::make_unique<cali::Function>(__VA_ARGS__);
+
+#define NVTX_MARK_FUNCTION_NAME(STR_NAME)                                \
+      NVTX(STR_NAME);                                                       \
+      std::unique_ptr<cali::Function> __cali_ann##__func__;                 \
+      if (g_caliper) {                                                      \
+         __cali_ann##__func__ = std::make_unique<cali::Function>(STR_NAME); \
+      }
+
+#define NVTX_MARK_BEGIN(...)     \
+      CALI_MARK_BEGIN(__VA_ARGS__); \
+      NVTX_BEGIN(__VA_ARGS__);
+
+#define NVTX_MARK_END(...)     \
+      NVTX_END(__VA_ARGS__);      \
+      CALI_MARK_END(__VA_ARGS__);
+#else
+#define NVTX_MARK_FUNCTION NVTX()
+#define NVTX_MARK(...) NVTX(__VA_ARGS__)
+#define NVTX_MARK_FUNCTION_NAME(...) NVTX(__VA_ARGS__)
+#define NVTX_MARK_BEGIN(...) NVTX_BEGIN(__VA_ARGS__)
+#define NVTX_MARK_END(...) NVTX_END(__VA_ARGS__)
+#endif
+
+} // namespace nvtx
+
+// Debug console traces, unnamed
+#if 1
+#define dbg(...) NVTX_DEBUG(__VA_ARGS__)
+#define dbl(...) NVTX_DEBUG_NO_END(__VA_ARGS__)
+#define dba(...) NVTX_DEBUG_APPEND(__VA_ARGS__)
+#define dbc(...) NVTX_DEBUG_NO_INI(__VA_ARGS__)
+#else
+#define dbg(...)
+#define dbl(...) (void)0
+#define dba(...)
+#define dbc(...)
+#endif

linalg/dtensor.hpp

@@ -87,7 +87,9 @@ protected:
 
 public:
    /// Default constructor
-   DeviceTensor() = delete;
+   // DeviceTensor() = delete;
+   MFEM_HOST_DEVICE
+   DeviceTensor() {}
 
    /// Constructor to initialize a tensor from the Scalar array data_
    template <typename... Args> MFEM_HOST_DEVICE
@@ -122,6 +124,17 @@ public:
    {
       return data[i];
    }
+
+   /// Returns the shape of the tensor.
+   MFEM_HOST_DEVICE inline std::array<int, Dim> GetShape() const
+   {
+      std::array<int, Dim> s;
+      for (int i = 0; i < Dim; i++)
+      {
+         s[i] = sizes[i];
+      }
+      return s;
+   }
 };
 
 

linalg/tensor.hpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// Copyright (c) 2010-2024, Lawrence Livermore National Security, LLC. Produced
 // at the Lawrence Livermore National Laboratory. All Rights reserved. See files
 // LICENSE and NOTICE for details. LLNL-CODE-806117.
 //
@@ -19,6 +19,8 @@
 #define MFEM_INTERNAL_TENSOR_HPP
 
 #include "dual.hpp"
+#include "general/backends.hpp"
+#include <limits>
 #include <type_traits> // for std::false_type
 
 namespace mfem
@@ -436,6 +438,23 @@ tensor<decltype(f(n1, n2, n3, n4)), n1, n2, n3, n4>
    return A;
 }
 
+// needs to be generalized
+template <typename T, int m, int n> MFEM_HOST_DEVICE
+tensor<T, n> get_col(tensor<T, m, n> A, int j)
+{
+   tensor<T, n> c{};
+   c(0) = A[0][j];
+   c(1) = A[1][j];
+   return c;
+}
+
+/// @overload
+template <typename T> MFEM_HOST_DEVICE
+tensor<T, 1> get_col(tensor<T, 1, 1> A, int j)
+{
+   return tensor<T, 1> {A[0][0]};
+}
+
 /**
  * @brief return the sum of two tensors
  * @tparam S the underlying type of the lefthand argument
@@ -697,6 +716,20 @@ auto outer(S A, T B) -> decltype(A * B)
    return A * B;
 }
 
+template <typename T, int n, int m> MFEM_HOST_DEVICE
+tensor<T, n + m> flatten(tensor<T, n, m> A)
+{
+   tensor<T, n + m> B{};
+   for (int i = 0; i < n; i++)
+   {
+      for (int j = 0; j < m; j++)
+      {
+         B(i + j * m) = A(i, j);
+      }
+   }
+   return B;
+}
+
 /**
  * @overload
  * @note this overload implements the case where the left argument is a scalar, and the right argument is a tensor
@@ -1051,6 +1084,18 @@ decltype(S {} * T{})
    return AB;
 }
 
+template <typename T, int m> MFEM_HOST_DEVICE
+auto dot(const tensor<T, m>& A, const tensor<T, m>& B) ->
+decltype(T {})
+{
+   decltype(T{}) AB{};
+   for (int i = 0; i < m; i++)
+   {
+      AB += A[i] * B[i];
+   }
+   return AB;
+}
+
 template <typename S, typename T, int m, int... n> MFEM_HOST_DEVICE
 auto dot(const tensor<S, m>& A, const tensor<T, m, n...>& B) ->
 tensor<decltype(S {} * T{}), n...>
@@ -1321,6 +1366,12 @@ tensor<T, n, m> transpose(const tensor<T, m, n>& A)
  * @param[in] A The matrix to obtain the determinant of
  */
 template <typename T> MFEM_HOST_DEVICE
+T det(const tensor<T, 1, 1>& A)
+{
+   return A[0][0];
+}
+/// @overload
+template <typename T> MFEM_HOST_DEVICE
 T det(const tensor<T, 2, 2>& A)
 {
    return A[0][0] * A[1][1] - A[0][1] * A[1][0];
@@ -1335,6 +1386,145 @@ T det(const tensor<T, 3, 3>& A)
           A[2][0];
 }
 
+template <typename T> MFEM_HOST_DEVICE
+std::tuple<tensor<T, 1>, tensor<T, 1, 1>> eig(tensor<T, 1, 1> &A)
+{
+   return {tensor<T, 1>{A[0][0]}, tensor<T, 1, 1>{{{1.0}}}};
+}
+
+template <typename T> MFEM_HOST_DEVICE
+std::tuple<tensor<T, 2>, tensor<T, 2, 2>> eig(tensor<T, 2, 2> &A)
+{
+   tensor<T, 2> e;
+   tensor<T, 2, 2> v;
+
+   double d0 = A(0, 0);
+   double d2 = A(0, 1);
+   double d3 = A(1, 1);
+   double c, s;
+
+   if (d2 == 0.0)
+   {
+      c = 1.0;
+      s = 0.0;
+   }
+   else
+   {
+      double t;
+      const double zeta = (d3 - d0) / (2.0 * d2);
+      const double azeta = fabs(zeta);
+      if (azeta < std::sqrt(1.0/std::numeric_limits<T>::epsilon()))
+      {
+         t = copysign(1./(azeta + std::sqrt(1. + zeta*zeta)), zeta);
+      }
+      else
+      {
+         t = copysign(0.5/azeta, zeta);
+      }
+      c = std::sqrt(1./(1. + t*t));
+      s = c*t;
+      t *= d2;
+      d0 -= t;
+      d3 += t;
+   }
+
+   if (d0 <= d3)
+   {
+      e(0) = d0;
+      e(1) = d3;
+      v(0, 0) =  c;
+      v(1, 0) = -s;
+      v(0, 1) =  s;
+      v(1, 1) =  c;
+   }
+   else
+   {
+      e(0) = d3;
+      e(1) = d0;
+      v(0, 0) =  s;
+      v(1, 0) =  c;
+      v(0, 1) =  c;
+      v(1, 1) = -s;
+   }
+
+   return {e, v};
+}
+
+template <typename T> MFEM_HOST_DEVICE
+void GetScalingFactor(const T &d_max, T &mult)
+{
+   int d_exp;
+   if (d_max > 0.)
+   {
+      mult = frexp(d_max, &d_exp);
+      if (d_exp == std::numeric_limits<T>::max_exponent)
+      {
+         mult *= std::numeric_limits<T>::radix;
+      }
+      mult = d_max/mult;
+   }
+   else
+   {
+      mult = 1.;
+   }
+}
+
+template <typename T> MFEM_HOST_DEVICE
+T calcsv(const tensor<T, 1, 1> A, const int i)
+{
+   return A[0][0];
+}
+
+/**
+ * @brief Compute the i-th singular value of a 2x2 matrix A
+ */
+template <typename T> MFEM_HOST_DEVICE
+T calcsv(const tensor<T, 2, 2> A, const int i)
+{
+   double mult;
+   double d0, d1, d2, d3;
+   d0 = A(0, 0);
+   d1 = A(1, 0);
+   d2 = A(0, 1);
+   d3 = A(1, 1);
+
+   double d_max = fabs(d0);
+   if (d_max < fabs(d1)) { d_max = fabs(d1); }
+   if (d_max < fabs(d2)) { d_max = fabs(d2); }
+   if (d_max < fabs(d3)) { d_max = fabs(d3); }
+
+   GetScalingFactor(d_max, mult);
+
+   d0 /= mult;
+   d1 /= mult;
+   d2 /= mult;
+   d3 /= mult;
+
+   double t = 0.5*((d0+d2)*(d0-d2)+(d1-d3)*(d1+d3));
+   double s = d0*d2 + d1*d3;
+   s = std::sqrt(0.5*(d0*d0 + d1*d1 + d2*d2 + d3*d3) + std::sqrt(t*t + s*s));
+
+   if (s == 0.0)
+   {
+      return 0.0;
+   }
+   t = fabs(d0*d3 - d1*d2) / s;
+   if (t > s)
+   {
+      if (i == 0)
+      {
+         return t*mult;
+      }
+      return s*mult;
+   }
+   if (i == 0)
+   {
+      return s*mult;
+   }
+   return t*mult;
+}
+
+
 /**
  * @brief Return whether a square rank 2 tensor is symmetric
  *
@@ -1474,13 +1664,20 @@ tensor<T, n> linear_solve(tensor<T, n, n> A, const tensor<T, n> b)
 /**
  * @brief Inverts a matrix
  * @param[in] A The matrix to invert
- * @note Uses a shortcut for inverting a 2-by-2 matrix
+ * @note Uses a shortcut for inverting a 1x1, 2x2 and 3x3 matrix
  */
-inline MFEM_HOST_DEVICE tensor<real_t, 2, 2> inv(const tensor<real_t, 2, 2>& A)
+template <typename T>
+inline MFEM_HOST_DEVICE tensor<T, 1, 1> inv(const tensor<T, 1, 1>& A)
 {
-   real_t inv_detA(1.0 / det(A));
+   return tensor<T, 1, 1> {{{T{1.0} / A[0][0]}}};
+}
 
-   tensor<real_t, 2, 2> invA{};
+template <typename T>
+inline MFEM_HOST_DEVICE tensor<T, 2, 2> inv(const tensor<T, 2, 2>& A)
+{
+   T inv_detA(1.0 / det(A));
+
+   tensor<T, 2, 2> invA{};
 
    invA[0][0] = A[1][1] * inv_detA;
    invA[0][1] = -A[0][1] * inv_detA;
@@ -1494,11 +1691,12 @@ inline MFEM_HOST_DEVICE tensor<real_t, 2, 2> inv(const tensor<real_t, 2, 2>& A)
  * @overload
  * @note Uses a shortcut for inverting a 3-by-3 matrix
  */
-inline MFEM_HOST_DEVICE tensor<real_t, 3, 3> inv(const tensor<real_t, 3, 3>& A)
+template <typename T>
+inline MFEM_HOST_DEVICE tensor<T, 3, 3> inv(const tensor<T, 3, 3>& A)
 {
-   real_t inv_detA(1.0 / det(A));
+   T inv_detA(1.0 / det(A));
 
-   tensor<real_t, 3, 3> invA{};
+   tensor<T, 3, 3> invA{};
 
    invA[0][0] = (A[1][1] * A[2][2] - A[1][2] * A[2][1]) * inv_detA;
    invA[0][1] = (A[0][2] * A[2][1] - A[0][1] * A[2][2]) * inv_detA;
@@ -1520,7 +1718,7 @@ inline MFEM_HOST_DEVICE tensor<real_t, 3, 3> inv(const tensor<real_t, 3, 3>& A)
 template <typename T, int n> MFEM_HOST_DEVICE
 tensor<T, n, n> inv(const tensor<T, n, n>& A)
 {
-   auto abs  = [](real_t x) { return (x < 0) ? -x : x; };
+   auto abs  = [](T x) { return (x < 0) ? -x : x; };
    auto swap = [](tensor<T, n>& x, tensor<T, n>& y)
    {
       auto tmp = x;
@@ -1528,12 +1726,12 @@ tensor<T, n, n> inv(const tensor<T, n, n>& A)
       y        = tmp;
    };
 
-   tensor<real_t, n, n> B = Identity<n>();
+   tensor<T, n, n> B = Identity<n>();
 
    for (int i = 0; i < n; i++)
    {
       // Search for maximum in this column
-      real_t max_val = abs(A[i][i]);
+      T max_val = abs(A[i][i]);
 
       int max_row = i;
       for (int j = i + 1; j < n; j++)
@@ -1553,7 +1751,7 @@ tensor<T, n, n> inv(const tensor<T, n, n>& A)
       {
          if (A[j][i] != 0.0)
          {
-            real_t c = -A[j][i] / A[i][i];
+            T c = -A[j][i] / A[i][i];
             A[j] += c * A[i];
             B[j] += c * B[i];
             A[j][i] = 0;

tests/unit/CMakeLists.txt

@@ -17,6 +17,7 @@ include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR})
 # The following list can be updated using (in bash):
 #    for d in general linalg mesh fem enzyme; do ls -1 $d/*.cpp; done
 set(UNIT_TESTS_SRCS
+  dfem/test_diffusion.cpp
   general/test_array.cpp
   general/test_arrays_by_name.cpp
   general/test_error.cpp

tests/unit/dfem/test_diffusion.cpp

@@ -0,0 +1,240 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#include "unit_tests.hpp"
+
+#include "mfem.hpp"
+#include "fem/dfem/doperator.hpp"
+#include "linalg/tensor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+using DOperator = DifferentiableOperator;
+
+#undef NVTX_COLOR
+#define NVTX_COLOR nvtx::kAquamarine
+#include "general/nvtx.hpp"
+
+namespace dfem_pa_kernels
+{
+
+///////////////////////////////////////////////////////////////////////////////
+template <int DIM> struct Diffusion
+{
+   using vecd_t = tensor<real_t, DIM>;
+   using matd_t = tensor<real_t, DIM, DIM>;
+
+   struct MFApply
+   {
+      MFEM_HOST_DEVICE inline auto operator()(const vecd_t &dudxi,
+                                              const real_t &rho,
+                                              const matd_t &J,
+                                              const real_t &w) const
+      {
+         const auto invJ = inv(J), TinJ = transpose(invJ);
+         return mfem::tuple{ (dudxi * invJ) * TinJ * det(J) * w * rho };
+      }
+   };
+
+   struct PASetup
+   {
+      MFEM_HOST_DEVICE inline auto operator()(const real_t &u,
+                                              const real_t &rho,
+                                              const matd_t &J,
+                                              const real_t &w) const
+      {
+         return mfem::tuple{ inv(J) * transpose(inv(J)) * det(J) * w * rho };
+      }
+   };
+
+   struct PAApply
+   {
+      MFEM_HOST_DEVICE inline auto operator()(const vecd_t &dudxi,
+                                              const matd_t &q) const
+      {
+         return mfem::tuple{ q * dudxi };
+      };
+   };
+};
+
+///////////////////////////////////////////////////////////////////////////////
+template <int DIM>
+void DFemDiffusion(const char *filename, int p, const int r)
+{
+   CAPTURE(filename, DIM, p, r);
+
+   Mesh smesh(filename);
+   ParMesh pmesh(MPI_COMM_WORLD, smesh);
+   MFEM_VERIFY(pmesh.Dimension() == DIM, "Mesh dimension mismatch");
+
+   pmesh.EnsureNodes();
+   auto *nodes = static_cast<ParGridFunction *>(pmesh.GetNodes());
+   p = std::max(p, pmesh.GetNodalFESpace()->GetMaxElementOrder());
+   smesh.Clear();
+
+   Array<int> all_domain_attr;
+   if (pmesh.bdr_attributes.Size() > 0)
+   {
+      all_domain_attr.SetSize(pmesh.bdr_attributes.Max());
+      all_domain_attr = 1;
+   }
+
+   H1_FECollection fec(p, DIM);
+   ParFiniteElementSpace pfes(&pmesh, &fec);
+   ParFiniteElementSpace *mfes = nodes->ParFESpace();
+
+   const int NE = pfes.GetNE(), d1d(p + 1), q = 2 * p + r;
+   const auto *ir = &IntRules.Get(pmesh.GetTypicalElementGeometry(), q);
+   const int q1d(IntRules.Get(Geometry::SEGMENT, ir->GetOrder()).GetNPoints());
+   MFEM_VERIFY(d1d <= q1d, "q1d should be >= d1d");
+   MFEM_VERIFY(NE > 0, "Mesh with no elements is not yet supported!");
+
+   ParGridFunction x(&pfes), y(&pfes), z(&pfes);
+
+   x.Randomize(1);
+   x.SetTrueVector();
+   x.SetFromTrueVector();
+
+   auto rho = [](const Vector &xyz)
+   {
+      const real_t x = xyz(0), y = xyz(1), z = DIM == 3 ? xyz(2) : 0.0;
+      real_t r = M_PI * pow(x, 2);
+      if (DIM >= 2) { r += pow(y, 3); }
+      if (DIM >= 3) { r += pow(z, 4); }
+      return r;
+   };
+   FunctionCoefficient rho_coeff(rho);
+
+   ParBilinearForm blf_fa(&pfes);
+   blf_fa.AddDomainIntegrator(new DiffusionIntegrator(rho_coeff, ir));
+   blf_fa.Assemble();
+   blf_fa.Finalize();
+
+   SECTION("Partial assembly")
+   {
+      dbg("Partial assembly");
+      ParBilinearForm blf_pa(&pfes);
+      blf_pa.AddDomainIntegrator(new DiffusionIntegrator(rho_coeff, ir));
+      blf_pa.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+      blf_pa.Assemble();
+      blf_pa.Mult(x, z);
+
+      blf_fa.Mult(x, y);
+      y -= z;
+      REQUIRE(y.Normlinf() == MFEM_Approx(0.0));
+      MPI_Barrier(MPI_COMM_WORLD);
+   }
+
+   QuadratureSpace qs(pmesh, *ir);
+   CoefficientVector rho_coeff_cv(rho_coeff, qs);
+   MFEM_VERIFY(rho_coeff_cv.GetVDim() == 1, "Coefficient should be scalar");
+   MFEM_VERIFY(rho_coeff_cv.Size() == q1d * q1d * (DIM == 3 ? q1d : 1) * NE, "");
+
+   const int rho_local_size = 1;
+   const int rho_elem_size(rho_local_size * ir->GetNPoints());
+   const int rho_total_size(rho_elem_size * NE);
+   ParametricSpace rho_ps(DIM, rho_local_size, rho_elem_size, rho_total_size,
+                          DIM == 3 ? d1d : d1d * d1d, // 🔥 2D workaround
+                          DIM == 3 ? q1d : q1d * q1d);
+
+   static constexpr int U = 0, Coords = 1, Rho = 3;
+   const auto sol = std::vector{ FieldDescriptor{ U, &pfes } };
+
+   SECTION("DFEM Matrix free")
+   {
+      DOperator dop_mf(sol, {{Rho, &rho_ps}, {Coords, mfes}}, pmesh);
+      typename Diffusion<DIM>::MFApply mf_apply_qf;
+      dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                 mfem::tuple{ Gradient<U>{}, None<Rho>{},
+                                              Gradient<Coords>{}, Weight{} },
+                                 mfem::tuple{ Gradient<U>{} }, *ir,
+                                 all_domain_attr);
+      dop_mf.SetParameters({ &rho_coeff_cv, nodes });
+      dop_mf.Mult(x, z);
+      z.SetTrueVector(), z.SetFromTrueVector();
+      blf_fa.Mult(x, y);
+      y.SetTrueVector(), y.SetFromTrueVector();
+      y -= z;
+      REQUIRE(y.Normlinf() == MFEM_Approx(0.0));
+      MPI_Barrier(MPI_COMM_WORLD);
+   }
+
+   SECTION("DFEM Partial assembly")
+   {
+      static constexpr int QData = 2;
+      const int qd_local_size = DIM * DIM;
+      const int qd_elem_size(qd_local_size * ir->GetNPoints());
+      const int qd_total_size(qd_elem_size * NE);
+      ParametricSpace qd_ps(DIM, qd_local_size, qd_elem_size, qd_total_size,
+                            DIM == 3 ? d1d : d1d * d1d, // 🔥 2D workaround
+                            DIM == 3 ? q1d : q1d * q1d);
+      ParametricFunction qdata(qd_ps);
+      qdata.UseDevice(true);
+
+      DOperator dSetup(sol, {{Rho, &rho_ps}, {Coords, mfes}, {QData, &qd_ps}}, pmesh);
+      typename Diffusion<DIM>::PASetup pa_setup_qf;
+      dSetup.AddDomainIntegrator(
+         pa_setup_qf,
+         mfem::tuple{ None<U>{}, None<Rho>{}, Gradient<Coords>{}, Weight{} },
+         mfem::tuple{ None<QData>{} }, *ir, all_domain_attr);
+      dSetup.SetParameters({ &rho_coeff_cv, nodes, &qdata });
+      pfes.GetRestrictionMatrix()->Mult(x, x.GetTrueVector());
+      dSetup.Mult(x.GetTrueVector(), qdata);
+
+      DOperator dop_pa(sol, { { QData, &qd_ps } }, pmesh);
+      typename Diffusion<DIM>::PAApply pa_apply_qf;
+      dop_pa.AddDomainIntegrator(pa_apply_qf,
+                                 mfem::tuple{ Gradient<U>{}, None<QData>{} },
+                                 mfem::tuple{ Gradient<U>{} },
+                                 *ir, all_domain_attr);
+      dop_pa.SetParameters({ &qdata });
+      dop_pa.Mult(x, z);
+      z.SetTrueVector(), z.SetFromTrueVector();
+      blf_fa.Mult(x, y);
+      y.SetTrueVector(), y.SetFromTrueVector();
+      y -= z;
+      REQUIRE(y.Normlinf() == MFEM_Approx(0.0));
+      MPI_Barrier(MPI_COMM_WORLD);
+   }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+TEST_CASE("DFEM Diffusion", "[Parallel][DFEM]")
+{
+   const bool all_tests = launch_all_non_regression_tests;
+
+   const auto p = !all_tests ? 1 : GENERATE(1, 2, 3);
+   const auto r = !all_tests ? 0 : GENERATE(0, 1, 2, 3);
+
+   SECTION("2D p=" + std::to_string(p) + " r=" + std::to_string(r))
+   {
+      const auto filename =
+         GENERATE("../../data/star.mesh",
+                  "../../data/star-q3.mesh",
+                  "../../data/rt-2d-q3.mesh",
+                  "../../data/inline-quad.mesh",
+                  "../../data/periodic-square.mesh");
+      DFemDiffusion<2>(filename, p, r);
+   }
+
+   SECTION("3D p=" + std::to_string(p) + " r=" + std::to_string(r))
+   {
+      const auto filename =
+         GENERATE("../../data/fichera.mesh",
+                  "../../data/fichera-q3.mesh",
+                  "../../data/inline-hex.mesh",
+                  "../../data/toroid-hex.mesh",
+                  "../../data/periodic-cube.mesh");
+      DFemDiffusion<3>(filename, p, r);
+   }
+}
+
+} // namespace dfem_pa_kernels

tests/unit/punit_test_main.cpp

@@ -18,6 +18,10 @@
 #error "This test should be disabled without MFEM_USE_MPI!"
 #endif
 
+#undef NVTX_COLOR
+#define NVTX_COLOR nvtx::kOrange
+#include "general/nvtx.hpp"
+
 int main(int argc, char *argv[])
 {
 #ifdef MFEM_USE_SINGLE
@@ -32,6 +36,8 @@ int main(int argc, char *argv[])
 #endif
    mfem::Device device("cpu"); // make sure hypre runs on CPU, if possible
 
+   dbg();
+
    // Only run tests that are labeled with Parallel.
    return RunCatchSession(argc, argv, {"[Parallel]"}, Root());
 }