Fix merge with global API

Co-authored-by: John Camier <camierjs@gmail.com>

.gitignore

@@ -451,8 +451,8 @@ miniapps/plasma/pic/*.csv
 tests/unit/output_meshes
 tests/unit/unit_tests
 tests/unit/punit_tests
-tests/unit/gpu_unit_tests
-tests/unit/pgpu_unit_tests
+tests/unit/cunit_tests
+tests/unit/pcunit_tests
 tests/unit/sedov_tests_*
 tests/unit/psedov_tests_*
 tests/unit/tmop_pa_tests_*

config/config.hpp

@@ -157,4 +157,22 @@ constexpr real_t operator""_r(unsigned long long v)
 #endif
 #endif // MFEM_USE_MPI not defined
 
+#ifdef NVTX_DBG_HPP
+#include NVTX_DBG_HPP
+#else
+#define db1(...)
+#define dbg(...)
+#define dbl(...)
+#define dba(...)
+#define dbc(...)
+#define NVTX_MARK_FUNCTION
+#define NVTX_MARK_BEGIN(...)
+#define NVTX_INI(...)
+#define NVTX_END(...)
+#define NVTX_MARK_INI(...)
+#define NVTX_MARK_END(...)
+#define NVTX_MARK(...)
+#define NVTX(...)
+#endif
+
 #endif // MFEM_CONFIG_HPP

fem/bilininteg.hpp

@@ -2178,18 +2178,22 @@ class DiffusionIntegrator: public BilinearFormIntegrator
 {
 public:
 
-   using ApplyKernelType = void(*)(const int, const bool, const Array<real_t>&,
-                                   const Array<real_t>&, const Array<real_t>&,
-                                   const Array<real_t>&,
-                                   const Vector&, const Vector&,
-                                   Vector&, const int, const int);
+   using DiffusionApplyKernelType = void(*)(const int, const bool,
+                                            const Array<real_t>&,
+                                            const Array<real_t>&, const Array<real_t>&,
+                                            const Array<real_t>&,
+                                            const Vector&, const Vector&,
+                                            Vector&, const int, const int);
 
-   using DiagonalKernelType = void(*)(const int, const bool, const Array<real_t>&,
-                                      const Array<real_t>&, const Vector&, Vector&,
-                                      const int, const int);
+   using DiffusionDiagonalKernelType = void(*)(const int, const bool,
+                                               const Array<real_t>&,
+                                               const Array<real_t>&, const Vector&, Vector&,
+                                               const int, const int);
 
-   MFEM_REGISTER_KERNELS(ApplyPAKernels, ApplyKernelType, (int, int, int));
-   MFEM_REGISTER_KERNELS(DiagonalPAKernels, DiagonalKernelType, (int, int, int));
+   MFEM_REGISTER_KERNELS(DiffusionApplyPAKernel, DiffusionApplyKernelType,
+                         (int, int, int));
+   MFEM_REGISTER_KERNELS(DiffusionDiagonalPAKernel, DiffusionDiagonalKernelType,
+                         (int, int, int));
    struct Kernels { Kernels(); };
 
 protected:
@@ -2209,6 +2213,7 @@ private:
    const FiniteElementSpace *fespace;
    const DofToQuad *maps;         ///< Not owned
    const GeometricFactors *geom;  ///< Not owned
+public:
    int dim, ne, dofs1D, quad1D;
    Vector pa_data;
    bool symmetric = true; ///< False if using a nonsymmetric matrix coefficient
@@ -2350,8 +2355,8 @@ public:
    template <int DIM, int D1D, int Q1D>
    static void AddSpecialization()
    {
-      ApplyPAKernels::Specialization<DIM,D1D,Q1D>::Add();
-      DiagonalPAKernels::Specialization<DIM,D1D,Q1D>::Add();
+      DiffusionApplyPAKernel::Specialization<DIM,D1D,Q1D>::Add();
+      DiffusionDiagonalPAKernel::Specialization<DIM,D1D,Q1D>::Add();
    }
 protected:
    const IntegrationRule* GetDefaultIntegrationRule(

fem/dfem/action.hpp

@@ -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.
+#pragma once
+
+#include <cassert>
+#include <cstddef>
+
+// #include "fem/kernels.hpp"
+#include "fem/kernels3d.hpp"
+namespace ker = mfem::kernels::internal;
+namespace low = mfem::kernels::internal::low;
+#include "fem/kernel_dispatch.hpp"
+
+// #include "linalg/kernels.hpp"
+
+#include "util.hpp"
+
+#undef NVTX_COLOR
+#define NVTX_COLOR ::nvtx::kOrchid
+
+namespace mfem::future
+{
+
+/** @brief Zero-copy view of a contiguous block as a `tensor<T, n1>` */
+template<typename T, int n1>
+MFEM_HOST_DEVICE
+const tensor<T, n1>& as_tensor(const T* ptr)
+{
+   // std::launder makes this defined behavior under strict aliasing rules
+   return *std::launder(reinterpret_cast<const tensor<T, n1>*>(ptr));
+}
+
+// convenience overload if you prefer a mutable view
+template<typename T, int n1>
+MFEM_HOST_DEVICE
+tensor<T, n1>& as_tensor(T* ptr)
+{
+   return *std::launder(reinterpret_cast<tensor<T, n1>*>(ptr));
+}
+
+/** @brief Zero-copy view of a contiguous block as a `tensor<T, n1, n2>` */
+template<typename T, int n1, int n2>
+MFEM_HOST_DEVICE
+const tensor<T, n1, n2>& as_tensor(const T* ptr)
+{
+   // std::launder makes this defined behavior under strict aliasing rules
+   return *std::launder(reinterpret_cast<const tensor<T, n1, n2>*>(ptr));
+}
+
+// convenience overload if you prefer a mutable view
+template<typename T, int n1, int n2>
+MFEM_HOST_DEVICE
+tensor<T, n1, n2>& as_tensor(T* ptr)
+{
+   return *std::launder(reinterpret_cast<tensor<T, n1, n2>*>(ptr));
+}
+
+/** @brief Zero-copy view of a contiguous block as a `tensor<T, n1, n2, n3>` */
+template<typename T, int n1, int n2, int n3>
+MFEM_HOST_DEVICE
+const tensor<T, n1, n2, n3>& as_tensor(const T* ptr)
+{
+   // std::launder makes this defined behavior under strict aliasing rules
+   return *std::launder(reinterpret_cast<const tensor<T, n1, n2, n3>*>(ptr));
+}
+
+// convenience overload if you prefer a mutable view
+template<typename T, int n1, int n2, int n3>
+MFEM_HOST_DEVICE
+tensor<T, n1, n2, n3>& as_tensor(T* ptr)
+{
+   return *std::launder(reinterpret_cast<tensor<T, n1, n2, n3>*>(ptr));
+}
+
+/** @brief Zero-copy view of a contiguous block as a `tensor<T, n1, n2, n3, n4>` */
+template<typename T, int n1, int n2, int n3, int n4>
+MFEM_HOST_DEVICE
+const tensor<T, n1, n2, n3, n4>& as_tensor(const T* ptr)
+{
+   // std::launder makes this defined behavior under strict aliasing rules
+   return *std::launder(reinterpret_cast<const tensor<T, n1, n2, n3, n4>*>(ptr));
+}
+
+// convenience overload if you prefer a mutable view
+template<typename T, int n1, int n2, int n3, int n4>
+MFEM_HOST_DEVICE
+tensor<T, n1, n2, n3, n4>& as_tensor(T* ptr)
+{
+   return *std::launder(reinterpret_cast<tensor<T, n1, n2, n3, n4>*>(ptr));
+}
+
+
+template <std::size_t N>
+MFEM_HOST_DEVICE inline
+std::array<real_t*, N>
+load_field_e_ptr(const std::array<DeviceTensor<2>, N> &fields_e,
+                 const int e)
+{
+   std::array<real_t*, N> f;
+   for_constexpr<N>([&](auto i) { f[i] = &fields_e[i](0, e); });
+   return f;
+}
+
+namespace qf
+{
+
+template <int T_Q1D,
+          size_t num_args,
+          typename reg_t,
+          typename qfunc_t,
+          typename args_ts>
+MFEM_HOST_DEVICE inline
+void apply_kernel(reg_t &res /*output*/,
+                  reg_t &reg,
+                  const real_t *rd,
+                  const int qx, const int qy, const int qz,
+                  const qfunc_t &qfunc, args_ts &args)
+{
+   if constexpr (num_args == 2) // PAApply
+   {
+      // ∇u
+      tensor<real_t, 3> &arg_0 = get<0>(args);
+      arg_0[0] = reg[qz][qy][qx][0];
+      arg_0[1] = reg[qz][qy][qx][1];
+      arg_0[2] = reg[qz][qy][qx][2];
+
+      // D (PA data)
+      tensor<real_t, 3, 3> &arg_1 = get<1>(args);
+
+      if constexpr (T_Q1D > 0)
+      {
+         const auto *D = (const real_t (*)[T_Q1D][T_Q1D][3][3]) rd;
+         for (int k = 0; k < 3; k++)
+         {
+            for (int j = 0; j < 3; j++)
+            {
+               arg_1[k][j] = D[qx][qy][qz][k][j];
+            }
+         }
+      }
+      else
+      {
+         static_assert(false);
+         // const auto D = Reshape(r2, 3, 3, Q1D, Q1D, Q1D);
+         // for (int j = 0; j < 3; j++)
+         // {
+         //    for (int k = 0; k < 3; k++)
+         //    {
+         //       arg_1[k][j] = D(j, k, qz, qy, qx);
+         //    }
+         // }
+      }
+   }
+   else
+   {
+      // MFApply comes here
+      assert(false);
+      // MFEM_ABORT("Only two arguments (∇u and D) are supported in apply_kernel for now");
+   }
+
+   const auto r = get<0>(apply(qfunc, args));
+
+   if constexpr (decltype(r)::ndim == 1)
+   {
+      // process_qf_result_from_reg(r0, qx, qy, qz, r);
+      as_tensor<real_t, 3>(&res[qz][qy][qx][0]) = r;
+   }
+   else
+   {
+      static_assert(false);
+   }
+}
+
+} // namespace qf
+
+#define MFEM_D2Q_MAX_SIZE 4
+static MFEM_CONSTANT real_t Bi[MFEM_D2Q_MAX_SIZE][8*8], Bo[8*8];
+static MFEM_CONSTANT real_t Gi[MFEM_D2Q_MAX_SIZE][8*8], Go[8*8];
+
+template<size_t num_fields,
+         size_t num_inputs,
+         size_t num_outputs,
+         typename restriction_cb_t,
+         typename qfunc_t,
+         typename input_t,
+         typename output_fop_t>
+class NewActionCallback
+{
+   restriction_cb_t &restriction_cb;
+   qfunc_t &qfunc;
+   input_t &inputs;
+   const std::array<size_t, num_inputs> &input_to_field;
+   const std::array<DofToQuadMap, num_inputs> &input_dtq_maps;
+   const std::array<DofToQuadMap, num_outputs> &output_dtq_maps;
+   const int num_entities;
+   const int test_vdim;
+   const int num_test_dof;
+   const int dimension;
+   const ThreadBlocks &thread_blocks;
+   SharedMemoryInfo<num_fields, num_inputs, num_outputs> &shmem_info;
+   const Array<int> &attributes;
+   const output_fop_t &output_fop;
+   const Array<int> *elem_attributes;
+   // refs
+   std::vector<Vector> &fields_e;
+   Vector &residual_e;
+   std::function<void(Vector &, Vector &)> &output_restriction_transpose;
+   // args
+   std::vector<Vector> &solutions_l;
+   const std::vector<Vector> &parameters_l;
+   Vector &residual_l;
+
+public:
+   NewActionCallback() = delete;
+
+   NewActionCallback(const bool use_kernels_specialization,
+                     restriction_cb_t &restriction_cb,
+                     qfunc_t &qfunc,
+                     input_t &inputs,
+                     const std::array<size_t, num_inputs> &input_to_field,
+                     const std::array<DofToQuadMap, num_inputs> &input_dtq_maps,
+                     const std::array<DofToQuadMap, num_outputs> &output_dtq_maps,
+                     const int num_entities,
+                     const int test_vdim,
+                     const int num_test_dof,
+                     const int dimension,
+                     const ThreadBlocks &thread_blocks,
+                     SharedMemoryInfo<num_fields, num_inputs, num_outputs> &shmem_info,
+                     const Array<int> &attributes,
+                     const output_fop_t &output_fop,
+                     const Array<int> *elem_attributes,
+                     // refs
+                     std::vector<Vector> &fields_e,
+                     Vector &residual_e,
+                     std::function<void(Vector &, Vector &)> &output_restriction_transpose,
+                     // args
+                     std::vector<Vector> &solutions_l,
+                     const std::vector<Vector> &parameters_l,
+                     Vector &residual_l):
+      restriction_cb(restriction_cb),
+      qfunc(qfunc),
+      inputs(inputs),
+      input_to_field(input_to_field),
+      input_dtq_maps(input_dtq_maps),
+      output_dtq_maps(output_dtq_maps),
+      num_entities(num_entities),
+      test_vdim(test_vdim),
+      num_test_dof(num_test_dof),
+      dimension(dimension),
+      thread_blocks(thread_blocks),
+      shmem_info(shmem_info),
+      attributes(attributes),
+      output_fop(output_fop),
+      elem_attributes(elem_attributes),
+      fields_e(fields_e),
+      residual_e(residual_e),
+      output_restriction_transpose(output_restriction_transpose),
+      solutions_l(solutions_l),
+      parameters_l(parameters_l),
+      residual_l(residual_l)
+   {
+      if (!use_kernels_specialization) { return; }
+      NewActionCallbackKernels::template Specialization<3>::Add(); // 1
+      NewActionCallbackKernels::template Specialization<4>::Add(); // 2
+      NewActionCallbackKernels::template Specialization<5>::Add(); // 3
+      NewActionCallbackKernels::template Specialization<6>::Add(); // 4
+      NewActionCallbackKernels::template Specialization<7>::Add(); // 5
+      NewActionCallbackKernels::template Specialization<8>::Add(); // 6
+   }
+
+   template<int T_Q1D = 0>
+   static void action_callback_new(const int d1d,
+                                   restriction_cb_t &restriction_cb,
+                                   qfunc_t &qfunc,
+                                   [[maybe_unused]] input_t &inputs,
+                                   [[maybe_unused]] const std::array<size_t, num_inputs> &input_to_field,
+                                   const std::array<DofToQuadMap, num_inputs> &input_dtq_maps,
+                                   const std::array<DofToQuadMap, num_outputs> &output_dtq_maps,
+                                   [[maybe_unused]] const int dimension,
+                                   const int num_entities,
+                                   [[maybe_unused]] const int test_vdim,
+                                   [[maybe_unused]] const int num_test_dof,
+                                   const ThreadBlocks &thread_blocks,
+                                   [[maybe_unused]] SharedMemoryInfo<num_fields, num_inputs, num_outputs>
+                                   &shmem_info,
+                                   [[maybe_unused]] const Array<int> &attributes,
+                                   [[maybe_unused]] const output_fop_t &output_fop,
+                                   [[maybe_unused]] const Array<int> *elem_attributes,
+                                   // refs
+                                   std::vector<Vector> &fields_e,
+                                   Vector &residual_e,
+                                   std::function<void(Vector &, Vector &)> &output_restriction_transpose,
+                                   // args
+                                   std::vector<Vector> &solutions_l,
+                                   const std::vector<Vector> &parameters_l,
+                                   Vector &residual_l,
+                                   // fallback arguments
+                                   const int q1d)
+   {
+      NVTX_MARK_FUNCTION;
+      assert(dimension == 3);
+      static_assert(MFEM_D2Q_MAX_SIZE >= num_inputs, "MFEM_D2Q_MAX_SIZE error");
+
+      constexpr int DIM = 3;
+
+      [[maybe_unused]] static bool ini = (for_constexpr<num_inputs>([&](auto i)
+      {
+         const auto dtq = input_dtq_maps[i];
+         {
+            const auto [q, _, p] = dtq.B.GetShape();
+            const auto B = (const real_t*)input_dtq_maps[i].B;
+            dbg("Loading Bi[{}]: q={} p={}", i.value, q, p);
+            if (B) { Gpu(MemcpyToSymbol)(Bi[i], B, (p*q)*sizeof(real_t)); }
+         }
+         {
+            const auto [q, _, p] = dtq.G.GetShape();
+            const auto G = (const real_t*)input_dtq_maps[i].G;
+            if (G) { Gpu(MemcpyToSymbol)(Gi[i], G, (p*q)*sizeof(real_t)); }
+         }
+         if constexpr (i == 0) // output B
+         {
+            const auto dtq_o = output_dtq_maps[0];
+            const auto [q, _, p] = dtq_o.B.GetShape();
+            const auto B = (const real_t*)dtq_o.B;
+            if (B) { Gpu(MemcpyToSymbol)(Bo, B, (p*q)*sizeof(real_t)); }
+         }
+         if constexpr (i == 0) // output G
+         {
+            const auto dtq_o = output_dtq_maps[0];
+            const auto [q, _, p] = dtq_o.G.GetShape();
+            const auto G = (const real_t*)dtq_o.G;
+            if (G) { Gpu(MemcpyToSymbol)(Go, G, (p*q)*sizeof(real_t)); }
+            dbg("Loaded B and G to constant memory");
+         }
+      }), true);
+
+      // types
+      using qf_signature =
+         typename create_function_signature<decltype(&qfunc_t::operator())>::type;
+      using qf_param_ts = typename qf_signature::parameter_ts;
+
+      restriction_cb(solutions_l, parameters_l, fields_e);
+
+      NVTX_INI("res=0");
+      residual_e = 0.0;
+      NVTX_END("res=0");
+
+      // auto wrapped_fields_e =
+      // wrap_fields(fields_e, shmem_info.field_sizes, num_entities);
+
+      const bool has_attr = attributes.Size() > 0;
+      const auto d_attr = attributes.Read();
+      const auto d_elem_attr = elem_attributes->Read();
+
+      // const int vdim = input.vdim;
+      // const auto fields_e_ptr = load_field_e_ptr(wrapped_fields_e, e);
+      // const real_t *field_e_r = fields_e_ptr[input_to_field[i]];
+      // const auto fields_e_ptr = load_field_e_ptr(wrapped_fields_e, e);
+      const int NE = num_entities;
+      constexpr int VDIM = 1;
+
+      const auto XE = Reshape(fields_e[0].Read(), d1d, d1d, d1d, VDIM, NE);
+      const real_t *dx_ptr = fields_e[1].Read();
+
+      auto YE = Reshape(residual_e.ReadWrite(), d1d, d1d, d1d, VDIM, NE);
+
+      const auto B = (const real_t*)input_dtq_maps[0/*i*/].B;
+      const auto G = (const real_t*)input_dtq_maps[0/*i*/].G;
+
+      NVTX_INI("forall");
+      dfem::forall<T_Q1D*T_Q1D*T_Q1D>([=] MFEM_HOST_DEVICE (int e, void *)
+      {
+         if (has_attr && !d_attr[d_elem_attr[e] - 1]) { return; }
+
+         constexpr int MQ1 = T_Q1D > 0 ? T_Q1D : 8;
+
+         MFEM_SHARED real_t sm0[MQ1][MQ1][MQ1][3];
+         MFEM_SHARED real_t sm1[MQ1][MQ1][MQ1][3];
+         // real_t (&sm0_ptr)[MQ1][MQ1][MQ1][3] = sm0;
+         // real_t (&sm1_ptr)[MQ1][MQ1][MQ1][3] = sm1;
+
+         low::regs3d_t<DIM, MQ1> reg;
+         const real_t *rd = dx_ptr;
+
+         // const auto fields_e_ptr = load_field_e_ptr(wrapped_fields_e, e);
+
+         MFEM_SHARED real_t sB[MQ1][MQ1], sG[MQ1][MQ1];
+         // real_t (&sB_ptr)[MD1][MQ1] = sB;
+         // real_t (&sG_ptr)[MD1][MQ1] = sG;
+
+         // Interpolate
+         // for_constexpr<num_inputs>(
+         //    [ D1D, Q1D, MQ1, e,
+         //           &input_dtq_maps,
+         //           &sm0_ptr, &sm1_ptr,
+         //           &sB = sB_ptr, &sG = sG_ptr,
+         //           &inputs,
+         //           //  &fields_e_ptr,
+         //           &reg, &rd,
+         //           &input_to_field ] (auto i)
+         {
+            // const auto input = get<0/*i*/>(inputs);
+            // using field_operator_t = std::decay_t<decltype(input)>;
+
+            // if constexpr (is_gradient_fop<field_operator_t>::value) // Grad
+            {
+               // const int vdim = input.vdim;
+               // const real_t *field_e_r = fields_e_ptr[input_to_field[i]];
+               // const auto XE = Reshape(field_e_r, D1D, D1D, D1D, vdim);
+               // const auto sB = reinterpret_cast<const real_t (*)[MQ1]>(Bi[i]);
+               // const auto sG = reinterpret_cast<const real_t (*)[MQ1]>(Gi[i]);
+               low::LoadMatrix(d1d, q1d, B, sB);
+               low::LoadMatrix(d1d, q1d, G, sG);
+               // for (int c = 0; c < vdim; c++)
+               // constexpr int c = 0;
+               {
+                  low::LoadDofs3d(e, d1d, XE, sm0);
+                  low::Grad3d(d1d, q1d, sB, sG, sm0, sm1, reg);
+               }
+            }
+            // else if constexpr (is_identity_fop<field_operator_t>::value)   // Identity
+            {
+               // db1("Identity");
+               // rd = fields_e_ptr[input_to_field[i]];
+               // rd = dx_ptr;
+            }
+            // else if constexpr (is_weight_fop<field_operator_t>::value)   // Weight
+            // {
+            //    dbg("Weight");
+            //    rw = fields_e_ptr[input_to_field[i]]; // 🔥
+            // }
+            // else
+            {
+               // MFApply comes here
+               // assert(false);
+               // MFEM_ABORT("Only Grad and Identity field operators are supported");
+            }
+         }//); // for_constexpr<num_inputs>
+
+         MFEM_FOREACH_THREAD_DIRECT(qz,z,q1d)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy,y,q1d)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx,x,q1d)
+               {
+#if 0
+                  auto qf_args = decay_tuple<qf_param_ts> {};
+                  qf::apply_kernel<T_Q1D, num_inputs>
+                  (reg, reg, rd, qx, qy, qz, qfunc, qf_args);
+#elif 0
+                  real_t v[3], u[3] = { reg[qz][qy][qx][0],
+                                        reg[qz][qy][qx][1],
+                                        reg[qz][qy][qx][2]
+                                      };
+                  const auto *D = (real_t (*)[T_Q1D][T_Q1D][3][3]) rd;
+                  kernels::Mult(3, 3, &D[qx][qy][qz][0][0], u, v);
+                  reg[qz][qy][qx][0] = v[0];
+                  reg[qz][qy][qx][1] = v[1];
+                  reg[qz][qy][qx][2] = v[2];
+#elif 0
+                  const auto *D = (real_t (*)[T_Q1D][T_Q1D][3][3]) rd;
+                  const auto args = decay_tuple<qf_param_ts>
+                  {
+                     {{ reg[qz][qy][qx][0], reg[qz][qy][qx][1], reg[qz][qy][qx][2] }},
+                     {{
+                           {{ D[qx][qy][qz][0][0], D[qx][qy][qz][0][1], D[qx][qy][qz][0][2] }},
+                           {{ D[qx][qy][qz][1][0], D[qx][qy][qz][1][1], D[qx][qy][qz][1][2] }},
+                           {{ D[qx][qy][qz][2][0], D[qx][qy][qz][2][1], D[qx][qy][qz][2][2] }}
+                        }
+                     }
+                  };
+                  const auto r = get<0>(apply(qfunc, args));
+                  reg[qz][qy][qx][0] = r[0];
+                  reg[qz][qy][qx][1] = r[1];
+                  reg[qz][qy][qx][2] = r[2];
+#elif 0
+                  auto u = as_tensor<real_t, 3>(&reg[qz][qy][qx][0]);
+                  const auto *d = (real_t (*)[T_Q1D][T_Q1D][3][3]) rd;
+                  auto D = as_tensor<real_t, 3, 3>(&d[qx][qy][qz][0][0]);
+                  auto r = D * u;
+                  reg[qz][qy][qx][0] = r[0];
+                  reg[qz][qy][qx][1] = r[1];
+                  reg[qz][qy][qx][2] = r[2];
+#else
+                  auto args = decay_tuple<qf_param_ts> {};
+                  get<0>(args) = as_tensor<real_t, 3>(&reg[qz][qy][qx][0]);
+                  if constexpr (T_Q1D > 0)
+                  {
+                     get<1>(args) = as_tensor<real_t, 3, 3>(rd + 9*(qx*T_Q1D*T_Q1D + qy*T_Q1D + qz));
+                  }
+                  else
+                  {
+                     get<1>(args) = as_tensor<real_t, 3, 3>(rd + 9*(qx*q1d*q1d + qy*q1d + qz));
+                  }
+                  auto r = get<0>(apply(qfunc, args));
+                  if constexpr (decltype(r)::ndim == 1)
+                  {
+                     as_tensor<real_t, 3>(&reg[qz][qy][qx][0]) = r;
+                  }
+                  else { static_assert(false); }
+#endif
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         // Integrate
+         // if constexpr (is_gradient_fop<std::decay_t<output_fop_t>>::value) // Gradient
+         {
+            // const auto sB = reinterpret_cast<const real_t (*)[MQ1]>(Bo);
+            // const auto sG = reinterpret_cast<const real_t (*)[MQ1]>(Go);
+            low::GradTranspose3d(d1d, q1d, sB, sG, reg, sm1, sm0);
+            low::WriteDofs3d(d1d, 0, e, reg, YE);
+         }
+      },
+      num_entities, thread_blocks, 0, nullptr);
+      NVTX_END("forall");
+
+      NVTX_INI("out^T");
+      output_restriction_transpose(residual_e, residual_l);
+      NVTX_END("out^T");
+   }
+
+   using NewActionKernelType = decltype(&NewActionCallback::action_callback_new<>);
+   MFEM_REGISTER_KERNELS(NewActionCallbackKernels, NewActionKernelType, (int));
+
+   void Apply(const int d1d, const int q1d)
+   {
+      db1();
+      NewActionCallbackKernels::Run(q1d,
+                                    // args
+                                    d1d,
+                                    restriction_cb,
+                                    qfunc,
+                                    inputs,
+                                    input_to_field,
+                                    input_dtq_maps,
+                                    output_dtq_maps,
+                                    dimension,
+                                    num_entities,
+                                    test_vdim,
+                                    num_test_dof,
+                                    thread_blocks,
+                                    shmem_info,
+                                    attributes,
+                                    output_fop,
+                                    elem_attributes,
+                                    fields_e,
+                                    residual_e,
+                                    output_restriction_transpose,
+                                    solutions_l,
+                                    parameters_l,
+                                    residual_l,
+                                    // fallback arguments
+                                    q1d);
+   }
+};
+
+template<size_t num_fields, size_t num_inputs, size_t num_outputs,
+         typename restriction_cb_t, typename qfunc_t, typename input_t, typename output_fop_t>
+template<int T_Q1D>
+typename NewActionCallback<num_fields, num_inputs, num_outputs, restriction_cb_t, qfunc_t, input_t, output_fop_t>::NewActionKernelType
+NewActionCallback<num_fields, num_inputs, num_outputs, restriction_cb_t, qfunc_t, input_t, output_fop_t>::NewActionCallbackKernels::Kernel()
+{
+   return action_callback_new<T_Q1D>;
+}
+
+template<size_t num_fields, size_t num_inputs, size_t num_outputs,
+         typename restriction_cb_t, typename qfunc_t, typename input_t, typename output_fop_t>
+typename NewActionCallback<num_fields, num_inputs, num_outputs, restriction_cb_t, qfunc_t, input_t, output_fop_t>::NewActionKernelType
+NewActionCallback<num_fields, num_inputs, num_outputs, restriction_cb_t, qfunc_t, input_t, output_fop_t>::NewActionCallbackKernels::Fallback
+(int q1d)
+{
+   dbg("\x1b[33mFallback q1d:{}", q1d);
+   // MFEM_ABORT("No kernel for q1d=" << q1d);
+   // return nullptr;
+   return action_callback_new<>;
+}
+
+} // namespace mfem::future

fem/dfem/backends/util.hpp

@@ -28,8 +28,8 @@ struct FieldBasis
    std::function<void(const Vector &, Vector &)> transpose;
 };
 
-FieldBasis FromQI(const QuadratureInterpolator *qi,
-                  QuadratureInterpolator::EvalFlags mode)
+inline FieldBasis FromQI(const QuadratureInterpolator *qi,
+                         QuadratureInterpolator::EvalFlags mode)
 {
    return
    {
@@ -62,7 +62,7 @@ FieldBasis FromQI(const QuadratureInterpolator *qi,
 }
 
 // QuadratureFunction identity copy
-FieldBasis FromQF()
+inline FieldBasis FromQF()
 {
    return
    {
@@ -72,7 +72,7 @@ FieldBasis FromQF()
 }
 
 // User-defined parameter space B
-FieldBasis FromPS(const Operator *B, const Operator *Bt)
+inline FieldBasis FromPS(const Operator *B, const Operator *Bt)
 {
    return
    {
@@ -81,7 +81,7 @@ FieldBasis FromPS(const Operator *B, const Operator *Bt)
    };
 }
 
-FieldBasis FieldBasisFromWeight(const IntegrationRule &ir)
+inline FieldBasis FieldBasisFromWeight(const IntegrationRule &ir)
 {
    return
    {
@@ -102,9 +102,9 @@ FieldBasis FieldBasisFromWeight(const IntegrationRule &ir)
    };
 }
 
-const FieldBasis GetFieldBasis(const FieldDescriptor &f,
-                               const IntegrationRule &ir,
-                               QuadratureInterpolator::EvalFlags mode)
+inline const FieldBasis GetFieldBasis(const FieldDescriptor &f,
+                                      const IntegrationRule &ir,
+                                      QuadratureInterpolator::EvalFlags mode)
 {
    return std::visit([&ir, &mode](auto && arg) -> FieldBasis
    {

fem/dfem/doperator.hpp

@@ -10,16 +10,24 @@
 // CONTRIBUTING.md for details.
 #pragma once
 
+#include <cassert>
+#include <type_traits>
+#include <utility>
+
 #include "../../config/config.hpp"
 
 #ifdef MFEM_USE_MPI
 #include "../fespace.hpp"
 #include "../linalg/multivector.hpp"
 
-#include "util.hpp"
+#include "action.hpp"
 #include "integrator_ctx.hpp"
 
 #include "backends/global_qf/prelude.hpp"
+#include "util.hpp"
+
+#undef NVTX_COLOR
+#define NVTX_COLOR ::nvtx::kTurquoise
 
 namespace mfem::future
 {
@@ -80,19 +88,20 @@ public:
       const std::vector<derivative_action_t> &derivative_actions,
       const std::vector<derivative_action_t> &derivative_actions_transpose,
       const FieldDescriptor &direction,
-      const vector_t &x,
+      [[maybe_unused]] const vector_t &x,
       const std::vector<FieldDescriptor> &infds,
       const std::vector<FieldDescriptor> &outfds) :
       Operator(height, width),
       derivative_actions(derivative_actions),
-      derivative_actions_transpose(derivative_actions_transpose),
-      direction(direction),
       infds(infds),
-      outfds(outfds)
+      outfds(outfds),
+      direction(direction),
+      derivative_actions_transpose(derivative_actions_transpose)
    {
       daction_l.resize(outfds.size());
       daction_e.resize(outfds.size());
       infields_e.resize(infds.size());
+      NVTX_MARK_FUNCTION;
       infields_l.resize(infds.size());
       for (size_t i = 0; i < infds.size(); i++)
       {
@@ -120,6 +129,7 @@ public:
    /// direction_t on T-dofs.
    void Mult(const Vector &x, Vector &y) const override
    {
+      NVTX_MARK_FUNCTION;
       MFEM_ASSERT(dynamic_cast<const BlockVector*>(&y),
                   "y needs to be a BlockVector");
       auto &by = static_cast<BlockVector &>(y);
@@ -376,10 +386,10 @@ public:
    {
       // TODO: Do we want those extensive checks here?
       static_assert(
-         std::is_same_v<x_t, MultiVector> && std::is_same_v<y_t, MultiVector> ||
-         std::is_same_v<x_t, BlockVector> && std::is_same_v<y_t, MultiVector> ||
-         std::is_same_v<x_t, MultiVector> && std::is_same_v<y_t, BlockVector> ||
-         std::is_same_v<x_t, BlockVector> && std::is_same_v<y_t, BlockVector>,
+         (std::is_same_v<x_t, MultiVector> && std::is_same_v<y_t, MultiVector>) ||
+         (std::is_same_v<x_t, BlockVector> && std::is_same_v<y_t, MultiVector>) ||
+         (std::is_same_v<x_t, MultiVector> && std::is_same_v<y_t, BlockVector>) ||
+         (std::is_same_v<x_t, BlockVector> && std::is_same_v<y_t, BlockVector>),
          "input and output vector types are incompatible");
 
       prolongation(infds, x, infields_l);
@@ -492,6 +502,10 @@ public:
    std::shared_ptr<DerivativeOperator> GetDerivative(
       size_t derivative_id, const MultiVector &x);
 
+   void UseNewKernels() { use_new_kernels = true; }
+
+   void UseKernelsSpecialization() { use_kernels_specialization = true; }
+
 private:
    const ParMesh &mesh;
 
@@ -532,6 +546,8 @@ private:
    std::map<size_t, size_t> assembled_vector_sizes;
 
    bool use_tensor_product_structure = true;
+   bool use_new_kernels = false;
+   bool use_kernels_specialization = false;
 
    size_t test_space_field_idx = SIZE_MAX;
 };
@@ -604,6 +620,7 @@ void DifferentiableOperator::AddIntegrator(
    static constexpr size_t num_outputs =
       tuple_size<decltype(outputs)>::value;
 
+   // constexpr int MQ1 = 8; // qfunc_t::MQ1;
    using qf_signature =
       typename get_function_signature<qfunc_t>::type;
    using qf_param_ts = typename qf_signature::parameter_ts;
@@ -686,12 +703,12 @@ void DifferentiableOperator::AddIntegrator(
       }
    }
 
-   ElementDofOrdering element_dof_ordering = ElementDofOrdering::NATIVE;
-   DofToQuad::Mode doftoquad_mode = DofToQuad::Mode::FULL;
+   // 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;
+      // element_dof_ordering = ElementDofOrdering::LEXICOGRAPHIC;
+      // doftoquad_mode = DofToQuad::Mode::TENSOR;
    }
 
    const int num_entities = GetNumEntities<entity_t>(mesh);

fem/dfem/interpolate.hpp

@@ -9,8 +9,23 @@
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
 #pragma once
+// #define NVTX_COLOR nvtx::kPeru
 
 #include "util.hpp"
+#include "fem/kernels.hpp"
+
+///////////////////////////////////////////////////////////////////////////////
+template <class T>
+inline std::enable_if_t<!std::numeric_limits<T>::is_integer, bool>
+AlmostEq(T x, T y, T tolerance = 15.0 * std::numeric_limits<T>::epsilon())
+{
+   const T neg = std::abs(x - y);
+   constexpr T min = std::numeric_limits<T>::min();
+   constexpr T eps = std::numeric_limits<T>::epsilon();
+   const T min_abs = std::min(std::abs(x), std::abs(y));
+   if (std::abs(min_abs) == 0.0) { return neg < eps; }
+   return (neg / (1.0 + std::max(min, min_abs))) < tolerance;
+}
 
 namespace mfem::future
 {
@@ -30,6 +45,7 @@ void map_field_to_quadrature_data_tensor_product_3d(
 
    if constexpr (is_value_fop<std::decay_t<field_operator_t>>::value)
    {
+      dbg("Value");
       auto [q1d, unused, d1d] = B.GetShape();
       const int vdim = input.vdim;
       const auto field = Reshape(&field_e[0], d1d, d1d, d1d, vdim);
@@ -94,10 +110,11 @@ void map_field_to_quadrature_data_tensor_product_3d(
    else if constexpr (
       is_gradient_fop<std::decay_t<field_operator_t>>::value)
    {
-      const auto [q1d, unused, d1d] = B.GetShape();
+      // dbg("Gradient");
+      const auto [q1d, B_dim, d1d] = B.GetShape();
       const int vdim = input.vdim;
       const int dim = input.dim;
-      const auto field = Reshape(&field_e[0], d1d, d1d, d1d, vdim);
+      const auto field = Reshape(&std::as_const(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);
@@ -106,7 +123,30 @@ void map_field_to_quadrature_data_tensor_product_3d(
       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++)
+      // constexpr int MQ1 = T_Q1D > 0 ? T_Q1D : 8;
+      // static constexpr int DIM = 3;
+      // MFEM_VERIFY(q1d <= MQ1, "q1d > MQ1");
+      // MFEM_SHARED real_t smem[MQ1][MQ1];
+
+      // kernels::internal::d_regs3d_t<DIM, MQ1> r0, r1;
+      // real_t sB[MQ1][MQ1], sG[MQ1][MQ1];
+
+      /*
+      {
+         assert(B_dim == 1 && "1D B required!");
+         kernels::internal::LoadMatrix(d1d, q1d, B, sB);
+         kernels::internal::LoadMatrix(d1d, q1d, G, sG);
+         for (int qx = 0; qx < q1d; qx++)
+         {
+            for (int dx = 0; dx < d1d; dx++)
+            {
+               assert(AlmostEq(B(qx, 0, dx), sB[dx][qx]));
+               assert(AlmostEq(G(qx, 0, dx), sG[dx][qx]));
+            }
+         }
+      }*/
+
+      for (int c = 0; c < vdim; c++)
       {
          MFEM_FOREACH_THREAD(dz, z, d1d)
          {
@@ -117,7 +157,7 @@ void map_field_to_quadrature_data_tensor_product_3d(
                   real_t uv[2] = {0.0, 0.0};
                   for (int dx = 0; dx < d1d; dx++)
                   {
-                     const real_t f = field(dx, dy, dz, vd);
+                     const real_t f = field(dx, dy, dz, c);
                      uv[0] += f * B(qx, 0, dx);
                      uv[1] += f * G(qx, 0, dx);
                   }
@@ -163,19 +203,59 @@ void map_field_to_quadrature_data_tensor_product_3d(
                      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];
+                  fqp(c, 0, qx, qy, qz) = uvw[0];
+                  fqp(c, 1, qx, qy, qz) = uvw[1];
+                  fqp(c, 2, qx, qy, qz) = uvw[2];
                }
             }
          }
          MFEM_SYNC_THREAD;
       }
+
+      /*
+      {
+         for (int c = 0; c < vdim; c++)
+         {
+            kernels::internal::LoadDofs3d(d1d, c, field, r0);
+            for (int d = 0; d < DIM; d++)
+            {
+               for (int dz = 0; dz < d1d; dz++)
+               {
+                  for (int dy = 0; dy < d1d; dy++)
+                  {
+                     for (int dx = 0; dx < d1d; dx++)
+                     {
+                        const real_t f = field(dx, dy, dz, c);
+                        assert(AlmostEq(f, r0[d][dz][dy][dx]));
+                     }
+                  }
+               }
+            }
+
+            kernels::internal::Grad3d(d1d, q1d, smem, sB, sG, r0, r1, c);
+            for (int qz = 0; qz < q1d; qz++)
+            {
+               for (int qy = 0; qy < q1d; qy++)
+               {
+                  for (int qx = 0; qx < q1d; qx++)
+                  {
+                     if (!AlmostEq(fqp(c, d, qx, qy, qz), r1[d][qz][qy][qx]))
+                     {
+                        dbg("\x1b[31m[{}:d] {} {}", c, fqp(c, d, qx, qy, qz), r1[d][qz][qy][qx]);
+                        dbg("❌❌❌"), std::exit(EXIT_FAILURE);
+                     }
+                  }
+               }
+            }
+         }
+         // dbg("✅✅✅✅✅✅✅✅✅✅✅✅✅✅✅");//, std::exit(EXIT_SUCCESS);
+      }*/
    }
    // TODO: Create separate function for clarity
    else if constexpr (
       std::is_same_v<std::decay_t<field_operator_t>, Weight>)
    {
+      // dbg("None");
       const int num_qp = integration_weights.GetShape()[0];
       // TODO: eeek
       const int q1d = (int)floor(std::pow(num_qp, 1.0/input.dim) + 0.5);
@@ -518,6 +598,9 @@ void map_fields_to_quadrature_data(
    const int &dimension,
    const bool &use_sum_factorization = false)
 {
+   // dbg();
+   assert(use_sum_factorization && "❌ use_sum_factorization required");
+
    // When the input_to_field map returns -1, this means the requested input
    // is the integration weight. Weights don't have a user defined field
    // attached to them and we create a dummy field which is not accessed
@@ -578,6 +661,7 @@ void map_field_to_quadrature_data_conditional(
    const int &dimension,
    const bool &use_sum_factorization = false)
 {
+   assert(false && "❌ condition not implemented");
    if (condition)
    {
       if (use_sum_factorization)
@@ -619,6 +703,7 @@ void map_fields_to_quadrature_data_conditional(
    const std::array<bool, num_inputs> &conditions,
    const bool &use_sum_factorization = false)
 {
+   assert(false && "❌ condition not implemented");
    for_constexpr<num_inputs>([&](auto i)
    {
       map_field_to_quadrature_data_conditional(
@@ -627,7 +712,7 @@ void map_fields_to_quadrature_data_conditional(
    });
 }
 
-template <size_t num_inputs, typename field_operator_ts>
+template <int T_Q1D, 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,
@@ -660,7 +745,7 @@ void map_direction_to_quadrature_data_conditional(
             }
             else if (dimension == 3)
             {
-               map_field_to_quadrature_data_tensor_product_3d(
+               map_field_to_quadrature_data_tensor_product_3d<T_Q1D>(
                   directions_qp[i], dtqmaps[i], direction_e, get<i>(fops),
                   integration_weights, scratch_mem);
             }

fem/dfem/parameterspace.hpp

@@ -125,6 +125,18 @@ public:
       return lsize;
    }
 
+   const Operator* GetB() const override
+   {
+      MFEM_ABORT("UniformParameterSpace does not support GetB");
+      return nullptr;
+   }
+
+   const Operator* GetBt() const override
+   {
+      MFEM_ABORT("UniformParameterSpace does not support GetBt");
+      return nullptr;
+   }
+
 private:
    /// T-vector size
    int tsize;

fem/dfem/qfunction_transform.hpp

@@ -243,6 +243,8 @@ void process_qf_arg(
    }
 }
 
+// const tensor<real_t, DIM>       ∇u
+// const tensor<real_t, DIM, DIM>  D (PA_DATA)
 template <typename arg_type>
 MFEM_HOST_DEVICE inline
 void process_qf_arg(const DeviceTensor<2> &u, arg_type &arg, int qp)

fem/dfem/qp.hpp

@@ -0,0 +1,76 @@
+// 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 "tuple.hpp"
+#include "../linalg/tensor.hpp"
+
+using namespace mfem::future;
+using mfem::future::tensor;
+
+// Helper to add dimension to tensor type
+template<typename T, int qp>
+struct AddQPDimension;
+
+// Specialization for tensor<real_t, dim>
+template<typename real_t, int dim, int qp>
+struct AddQPDimension<tensor<real_t, dim>, qp>
+{
+   using type = tensor<real_t, dim, qp>;
+};
+
+// Specialization for tensor<real_t, dim, dim>
+template<typename real_t, int dim, int qp>
+struct AddQPDimension<tensor<real_t, dim, dim>, qp>
+{
+   using type = tensor<real_t, dim, dim, qp>;
+};
+
+// Specialization for real_t (transforms to tensor<real_t, qp>)
+template<typename real_t, int qp>
+struct AddQPDimension
+{
+   using type = tensor<real_t, qp>;
+};
+
+// Helper to transform tuple
+template<typename Tuple, int qp>
+struct TransformTupleQP {};
+
+// Specialization for mfem::future::tuple
+template<int qp, typename... Types>
+struct TransformTupleQP<mfem::future::tuple<Types...>, qp>
+{
+   using type = mfem::future::tuple<typename AddQPDimension<Types, qp>::type...>;
+};
+
+template<int qp, typename... Types>
+struct TransformTupleQP<std::tuple<Types...>, qp>
+{
+   using type = std::tuple<typename AddQPDimension<Types, qp>::type...>;
+};
+
+// Function to transform tuple type with qp dimension
+template<int qp, typename qf_param_ts>
+struct add_qp_dimension
+{
+   using type = typename TransformTupleQP<qf_param_ts, qp>::type;
+};
+
+// Helper alias template for cleaner usage
+template<int qp, typename qf_param_ts>
+using add_qp_dimension_t = typename add_qp_dimension<qp, qf_param_ts>::type;
+
+// ...AddDomainIntegrator...
+// {
+//    constexpr int Q1D = 4;
+//    using qf_param_augmentd_ts = add_qp_dimension_t<Q1D, decay_tuple<qf_param_ts>>;
+// }

fem/dfem/util.hpp

@@ -39,6 +39,9 @@
 #include "parameterspace.hpp"
 #include "tuple.hpp"
 
+#undef NVTX_COLOR
+#define NVTX_COLOR ::nvtx::kLightBlue
+
 namespace mfem::future
 {
 
@@ -79,7 +82,7 @@ constexpr void for_constexpr(lambda&& f,
 }
 
 template <typename lambda>
-constexpr void for_constexpr(lambda&& f, std::integer_sequence<std::size_t>) {}
+constexpr void for_constexpr(lambda&&, std::integer_sequence<std::size_t>) {}
 
 template <int... n, typename lambda>
 constexpr void for_constexpr(lambda&& f)
@@ -88,7 +91,7 @@ constexpr void for_constexpr(lambda&& f)
 }
 
 template <typename lambda, typename arg_t>
-constexpr void for_constexpr_with_arg(lambda&& f, arg_t&& arg,
+constexpr void for_constexpr_with_arg(lambda&&, arg_t&&,
                                       std::integer_sequence<std::size_t>)
 {
    // Base case - do nothing for empty sequence
@@ -567,7 +570,7 @@ auto get_marked_entries(
 /// @param t the tuple to filter fields from.
 /// @returns a tuple containing only the fields with field IDs not equal to -1.
 template <typename... Ts>
-constexpr auto filter_fields(const std::tuple<Ts...>& t)
+constexpr auto filter_fields(const std::tuple<Ts...>&)
 {
    return std::tuple_cat(
              std::conditional_t<Ts::GetFieldId() != -1, std::tuple<Ts>, std::tuple<>> {}...);
@@ -602,7 +605,7 @@ struct ThreadBlocks
 
 #if defined(MFEM_USE_CUDA_OR_HIP)
 template <typename func_t>
-__global__ void forall_kernel_shmem(func_t f, int n)
+__global__ void forall_kernel_extern_shmem(func_t f, int n)
 {
    int i = blockIdx.x;
    extern __shared__ real_t shmem[];
@@ -611,23 +614,48 @@ __global__ void forall_kernel_shmem(func_t f, int n)
       f(i, shmem);
    }
 }
+template <typename func_t>
+__global__ void forall_kernel_static_smem(func_t f, int n)
+{
+   int i = blockIdx.x;
+   if (i >= n) { return; }
+   f(i, nullptr);
+}
+template <int MAX_THREADS_PER_BLOCK, typename func_t>
+__global__
+MFEM_LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK)
+static void forall_kernel_static_smem_launch_bounds(func_t f, int n)
+{
+   for (int k = blockIdx.x; k < n; k += gridDim.x) { f(k, nullptr); }
+}
 #endif
 
-template <typename func_t>
+template </*typename kernel_tag,*/ typename func_t>
 void forall(func_t f,
             const int &N,
-            const ThreadBlocks &blocks,
-            int num_shmem = 0,
+            [[maybe_unused]] const ThreadBlocks &blocks,
+            [[maybe_unused]] int num_shmem = 0,
             real_t *shmem = nullptr)
 {
+   db1();
    if (Device::Allows(Backend::CUDA_MASK) ||
        Device::Allows(Backend::HIP_MASK))
    {
 #if defined(MFEM_USE_CUDA_OR_HIP)
       // int gridsize = (N + Z - 1) / Z;
       int num_bytes = num_shmem * sizeof(decltype(shmem));
+      db1("num_bytes:{}", num_bytes);
+      db1("block: {}x{}x{}", blocks.x, blocks.y, blocks.z);
       dim3 block_size(blocks.x, blocks.y, blocks.z);
-      forall_kernel_shmem<<<N, block_size, num_bytes>>>(f, N);
+      // ForallKernel<kernel_tag>::run<<<N, block_size, num_bytes>>>(f, N);
+      if (num_bytes > 0)
+      {
+         forall_kernel_extern_shmem<<<N, block_size, num_bytes>>>(f, N);
+      }
+      else
+      {
+         forall_kernel_static_smem<<<N, block_size>>>(f, N);
+      }
 #if defined(MFEM_USE_CUDA)
       MFEM_GPU_CHECK(cudaGetLastError());
 #elif defined(MFEM_USE_HIP)
@@ -638,6 +666,7 @@ void forall(func_t f,
    }
    else if (Device::Allows(Backend::CPU_MASK))
    {
+      db1("CPU_MASK");
       MFEM_ASSERT(!((bool)num_shmem != (bool)shmem),
                   "Backend::CPU needs a pre-allocated shared memory block");
       for (int i = 0; i < N; i++)
@@ -651,6 +680,69 @@ void forall(func_t f,
    }
 }
 
+namespace dfem
+{
+
+template <int MAX_THREADS_PER_BLOCK = 0, typename func_t>
+void forall(func_t f,
+            const int &N,
+            [[maybe_unused]] const ThreadBlocks &blocks,
+            [[maybe_unused]] int num_shmem = 0,
+            real_t *shmem = nullptr)
+{
+   db1();
+   if (Device::Allows(Backend::CUDA_MASK) ||
+       Device::Allows(Backend::HIP_MASK))
+   {
+#if defined(MFEM_USE_CUDA_OR_HIP)
+      int num_bytes = num_shmem * sizeof(decltype(shmem));
+      db1("num_bytes:{}", num_bytes);
+      db1("block: {}x{}x{}", blocks.x, blocks.y, blocks.z);
+      db1("MAX_THREADS_PER_BLOCK:{}", MAX_THREADS_PER_BLOCK);
+      dim3 block_size(blocks.x, blocks.y, blocks.z);
+      if constexpr (MAX_THREADS_PER_BLOCK > 0)
+      {
+         assert(num_bytes == 0);
+         forall_kernel_static_smem_launch_bounds
+         <MAX_THREADS_PER_BLOCK><<<N, block_size>>> (f, N);
+      }
+      else
+      {
+         static_assert(MAX_THREADS_PER_BLOCK == 0);
+         if (num_bytes == 0)
+         {
+            forall_kernel_static_smem<<<N, block_size>>>(f, N);
+         }
+         else
+         {
+            forall_kernel_extern_shmem<<<N, block_size, num_bytes>>>(f, N);
+         }
+      }
+#if defined(MFEM_USE_CUDA)
+      MFEM_GPU_CHECK(cudaGetLastError());
+#elif defined(MFEM_USE_HIP)
+      MFEM_GPU_CHECK(hipGetLastError());
+#endif
+      // MFEM_DEVICE_SYNC; // ⚠️
+#endif
+   }
+   else if (Device::Allows(Backend::CPU_MASK))
+   {
+      db1("CPU_MASK");
+      MFEM_ASSERT(!((bool)num_shmem != (bool)shmem),
+                  "Backend::CPU needs a pre-allocated shared memory block");
+      for (int i = 0; i < N; i++)
+      {
+         f(i, shmem);
+      }
+   }
+   else
+   {
+      MFEM_ABORT("no compute backend available");
+   }
+}
+}
+
 /// @todo To be removed.
 class FDJacobian : public Operator
 {
@@ -982,6 +1074,7 @@ std::variant<const QuadratureInterpolator *, const Operator *>get_qinterp(
 inline
 const Operator *get_prolongation(const FieldDescriptor &f)
 {
+   NVTX("get P");
    return std::visit([](auto&& arg) -> const Operator*
    {
       using T = std::decay_t<decltype(arg)>;
@@ -1016,6 +1109,7 @@ inline
 const Operator *get_element_restriction(const FieldDescriptor &f,
                                         ElementDofOrdering o)
 {
+   NVTX("get ER");
    return std::visit([&o](auto&& arg) -> const Operator*
    {
       using T = std::decay_t<decltype(arg)>;
@@ -1055,6 +1149,7 @@ const Operator *get_face_restriction(const FieldDescriptor &f,
                                      FaceType ft,
                                      L2FaceValues m)
 {
+   NVTX("get FR");
    return std::visit([&o, &ft, &m](auto&& arg) -> const Operator*
    {
       using T = std::decay_t<decltype(arg)>;
@@ -1093,6 +1188,7 @@ inline
 const Operator *get_restriction(const FieldDescriptor &f,
                                 const ElementDofOrdering &o)
 {
+   NVTX("get R");
    if constexpr (std::is_same_v<entity_t, Entity::Element>)
    {
       return get_element_restriction(f, o);
@@ -1118,12 +1214,14 @@ inline std::tuple<std::function<void(const Vector&, Vector&)>, int>
 get_restriction_transpose(
    const FieldDescriptor &f,
    const ElementDofOrdering &o,
-   const fop_t &fop)
+   [[maybe_unused]] const fop_t &fop)
 {
+   NVTX("get R^T");
    if constexpr (is_sum_fop<fop_t>::value)
    {
       auto RT = [=](const Vector &v_e, Vector &v_l)
       {
+         NVTX("R^T sum");
          v_l += v_e;
       };
       return std::make_tuple(RT, 1);
@@ -1133,6 +1231,7 @@ get_restriction_transpose(
       const Operator *R = get_restriction<entity_t>(f, o);
       std::function<void(const Vector&, Vector&)> RT = [=](const Vector &x, Vector &y)
       {
+         NVTX("R^T+");
          R->AddMultTranspose(x, y);
       };
       return std::make_tuple(RT, R->Height());
@@ -1152,8 +1251,14 @@ get_restriction_transpose(
 inline
 void prolongation(const FieldDescriptor field, const Vector &x, Vector &field_l)
 {
+   NVTX("P");
    const auto P = get_prolongation(field);
+
+   NVTX_INI("SetSize");
    field_l.SetSize(P->Height());
+   NVTX_END("SetSize");
+
+   NVTX_INI("P->Mult");
    P->Mult(x, field_l);
 }
 
@@ -1182,6 +1287,7 @@ void prolongation(const std::array<FieldDescriptor, N> fields,
                   const Vector &x,
                   std::array<Vector, M> &fields_l)
 {
+   NVTX("P");
    int data_offset = 0;
    for (int i = 0; i < N; i++)
    {
@@ -1189,9 +1295,14 @@ void prolongation(const std::array<FieldDescriptor, N> fields,
       const int width = P->Width();
       // const Vector x_i(x.GetData() + data_offset, width);
       const Vector x_i(const_cast<Vector&>(x), data_offset, width);
-      fields_l[i].SetSize(P->Height());
 
+      NVTX_INI("SetSize");
+      fields_l[i].SetSize(P->Height());
+      NVTX_END("SetSize");
+
+      NVTX_INI("P->Mult");
       P->Mult(x_i, fields_l[i]);
+      NVTX_END("P->Mult");
       data_offset += width;
    }
 }
@@ -1466,6 +1577,7 @@ void get_lvectors(const std::vector<FieldDescriptor> fields,
                   const Vector &x,
                   std::vector<Vector> &fields_l)
 {
+   NVTX("get_lvectors");
    int data_offset = 0;
    for (std::size_t i = 0; i < fields.size(); i++)
    {
@@ -1492,13 +1604,15 @@ template <typename fop_t>
 inline
 std::function<void(const Vector&, Vector&)> get_prolongation_transpose(
    const FieldDescriptor &f,
-   const fop_t &fop,
+   [[maybe_unused]] const fop_t &fop,
    MPI_Comm mpi_comm)
 {
+   NVTX("get P^T");
    if constexpr (is_sum_fop<fop_t>::value)
    {
       auto PT = [=](const Vector &r_local, Vector &y)
       {
+         NVTX("P^T sum");
          MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
          real_t local_sum = r_local.Sum();
          MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM, mpi_comm);
@@ -1509,6 +1623,7 @@ std::function<void(const Vector&, Vector&)> get_prolongation_transpose(
    {
       auto PT = [=](const Vector &r_local, Vector &y)
       {
+         NVTX("P^T Identity");
          y = r_local;
       };
       return PT;
@@ -1516,6 +1631,7 @@ std::function<void(const Vector&, Vector&)> get_prolongation_transpose(
    const Operator *P = get_prolongation(f);
    auto PT = [=](const Vector &r_local, Vector &y)
    {
+      NVTX("P^T");
       P->MultTranspose(r_local, y);
    };
    return PT;
@@ -1534,12 +1650,19 @@ void restriction(const FieldDescriptor u,
                  Vector &field_e,
                  ElementDofOrdering ordering)
 {
+   NVTX("R");
    const auto R = get_restriction<entity_t>(u, ordering);
    MFEM_ASSERT(R->Width() == u_l.Size(),
                "restriction not applicable to given data size");
    const int height = R->Height();
+
+   NVTX_INI("SetSize");
    field_e.SetSize(height);
+   NVTX_END("SetSize");
+
+   NVTX_INI("R->Mult");
    R->Mult(u_l, field_e);
+   NVTX_END("R->Mult");
 }
 
 /// @brief Apply the restriction operator to a vector of fields.
@@ -1557,14 +1680,29 @@ void restriction(const std::vector<FieldDescriptor> u,
                  ElementDofOrdering ordering,
                  const int offset = 0)
 {
+   NVTX("R");
    for (std::size_t i = 0; i < u.size(); i++)
    {
       const auto R = get_restriction<entity_t>(u[i], ordering);
       MFEM_ASSERT(R->Width() == u_l[i].Size(),
                   "restriction not applicable to given data size");
       const int height = R->Height();
+
+      // NVTX_INI("SetSize");
       fields_e[i + offset].SetSize(height);
-      R->Mult(u_l[i], fields_e[i + offset]);
+      // NVTX_END("SetSize");
+
+      // NVTX_INI("R->Mult");
+      if (dynamic_cast<const IdentityOperator*>(R))
+      {
+         NVTX("Identity");
+         fields_e[i + offset].NewMemoryAndSize(u_l[i].GetMemory(), u_l[i].Size(), false);
+      }
+      else
+      {
+         R->Mult(u_l[i], fields_e[i + offset]);
+      }
+      // NVTX_END("R->Mult");
    }
 }
 
@@ -1576,14 +1714,21 @@ void element_restriction(const std::array<FieldDescriptor, N> u,
                          ElementDofOrdering ordering,
                          const int offset = 0)
 {
+   NVTX("ER");
    for (int i = 0; i < N; i++)
    {
       const auto R = get_element_restriction(u[i], ordering);
       MFEM_ASSERT(R->Width() == u_l[i].Size(),
                   "element restriction not applicable to given data size");
       const int height = R->Height();
+
+      NVTX_INI("SetSize");
       fields_e[i + offset].SetSize(height);
+      NVTX_END("SetSize");
+
+      NVTX_INI("R->Mult");
       R->Mult(u_l[i], fields_e[i + offset]);
+      NVTX_END("R->Mult");
    }
 }
 
@@ -1905,7 +2050,7 @@ get_shmem_info(
    const std::array<DofToQuadMap, num_outputs> &output_dtq_maps,
    const std::vector<FieldDescriptor> &fields,
    const int &num_entities,
-   const input_t &inputs,
+   [[maybe_unused]] const input_t &inputs,
    const int &num_qp,
    const std::vector<int> &input_size_on_qp,
    const int &residual_size_on_qp,

fem/integ/bilininteg_diffusion_kernels.hpp

@@ -1064,6 +1064,8 @@ inline void SmemPADiffusionApply3D(const int NE,
          }
       }
       MFEM_SYNC_THREAD;
+
+      // Grad X
       MFEM_FOREACH_THREAD_DIRECT(dz,z,D1D)
       {
          MFEM_FOREACH_THREAD_DIRECT(dy,y,D1D)
@@ -1084,6 +1086,8 @@ inline void SmemPADiffusionApply3D(const int NE,
          }
       }
       MFEM_SYNC_THREAD;
+
+      // Grad Y
       MFEM_FOREACH_THREAD_DIRECT(dz,z,D1D)
       {
          MFEM_FOREACH_THREAD_DIRECT(qy,y,Q1D)
@@ -1105,6 +1109,8 @@ inline void SmemPADiffusionApply3D(const int NE,
          }
       }
       MFEM_SYNC_THREAD;
+
+      // Grad Z + Q-function
       MFEM_FOREACH_THREAD_DIRECT(qz,z,Q1D)
       {
          MFEM_FOREACH_THREAD_DIRECT(qy,y,Q1D)
@@ -1217,20 +1223,23 @@ inline void SmemPADiffusionApply3D(const int NE,
 
 namespace
 {
-using ApplyKernelType = DiffusionIntegrator::ApplyKernelType;
-using DiagonalKernelType = DiffusionIntegrator::DiagonalKernelType;
+using DiffusionApplyKernelType =
+   DiffusionIntegrator::DiffusionApplyKernelType;
+
+using DiffusionDiagonalKernelType =
+   DiffusionIntegrator::DiffusionDiagonalKernelType;
 }
 
 template<int DIM, int T_D1D, int T_Q1D>
-ApplyKernelType DiffusionIntegrator::ApplyPAKernels::Kernel()
+DiffusionApplyKernelType DiffusionIntegrator::DiffusionApplyPAKernel::Kernel()
 {
    if constexpr (DIM == 2) { return internal::SmemPADiffusionApply2D<T_D1D,T_Q1D>; }
    else if constexpr (DIM == 3) { return internal::SmemPADiffusionApply3D<T_D1D, T_Q1D>; }
    MFEM_ABORT("");
 }
 
-inline
-ApplyKernelType DiffusionIntegrator::ApplyPAKernels::Fallback(int DIM, int, int)
+inline DiffusionApplyKernelType
+DiffusionIntegrator::DiffusionApplyPAKernel::Fallback(int DIM, int, int)
 {
    if (DIM == 2) { return internal::PADiffusionApply2D; }
    else if (DIM == 3) { return internal::PADiffusionApply3D; }
@@ -1238,15 +1247,16 @@ ApplyKernelType DiffusionIntegrator::ApplyPAKernels::Fallback(int DIM, int, int)
 }
 
 template<int DIM, int D1D, int Q1D>
-DiagonalKernelType DiffusionIntegrator::DiagonalPAKernels::Kernel()
+DiffusionDiagonalKernelType
+DiffusionIntegrator::DiffusionDiagonalPAKernel::Kernel()
 {
    if constexpr (DIM == 2) { return internal::SmemPADiffusionDiagonal2D<D1D,Q1D>; }
    else if constexpr (DIM == 3) { return internal::SmemPADiffusionDiagonal3D<D1D, Q1D>; }
    MFEM_ABORT("");
 }
 
-inline DiagonalKernelType
-DiffusionIntegrator::DiagonalPAKernels::Fallback(int DIM, int, int)
+inline DiffusionDiagonalKernelType
+DiffusionIntegrator::DiffusionDiagonalPAKernel::Fallback(int DIM, int, int)
 {
    if (DIM == 2) { return internal::PADiffusionDiagonal2D; }
    else if (DIM == 3) { return internal::PADiffusionDiagonal3D; }

fem/integ/bilininteg_diffusion_pa.cpp

@@ -31,8 +31,8 @@ void DiffusionIntegrator::AssembleDiagonalPA(Vector &diag)
       const Array<real_t> &B = maps->B;
       const Array<real_t> &G = maps->G;
       const Vector &Dv = pa_data;
-      DiagonalPAKernels::Run(dim, dofs1D, quad1D, ne, symmetric, B, G, Dv,
-                             diag, dofs1D, quad1D);
+      DiffusionDiagonalPAKernel::Run(dim, dofs1D, quad1D, ne, symmetric, B, G, Dv,
+                                     diag, dofs1D, quad1D);
    }
 }
 
@@ -68,8 +68,8 @@ void DiffusionIntegrator::AddMultPA(const Vector &x, Vector &y) const
       }
 #endif // MFEM_USE_OCCA
 
-      ApplyPAKernels::Run(dim, dofs1D, quad1D, ne, symmetric, B, G, Bt,
-                          Gt, Dv, x, y, dofs1D, quad1D);
+      DiffusionApplyPAKernel::Run(dim, dofs1D, quad1D, ne, symmetric, B, G, Bt,
+                                  Gt, Dv, x, y, dofs1D, quad1D);
    }
 }
 
@@ -174,9 +174,9 @@ void DiffusionIntegrator::AddAbsMultPA(const Vector &x, Vector &y) const
    abs_pa_data.Abs();
    auto abs_maps = maps->Abs();
 
-   ApplyPAKernels::Run(dim, dofs1D, quad1D, ne, symmetric,
-                       abs_maps.B, abs_maps.G, abs_maps.Bt, abs_maps.Gt,
-                       abs_pa_data, x, y, dofs1D, quad1D);
+   DiffusionApplyPAKernel::Run(dim, dofs1D, quad1D, ne, symmetric,
+                               abs_maps.B, abs_maps.G, abs_maps.Bt, abs_maps.Gt,
+                               abs_pa_data, x, y, dofs1D, quad1D);
 }
 
 void DiffusionIntegrator::AddAbsMultTransposePA(const Vector &x,

fem/kernels.hpp

@@ -207,6 +207,28 @@ inline MFEM_HOST_DEVICE void WriteDofs2d(const int e, const int d1d,
    MFEM_SYNC_THREAD;
 }
 
+/// Load 3D input DIM vector at element offset into given register tensor
+template <int VDIM, int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void LoadDofs3d(const int d1d, const int c,
+                                        const DeviceTensor<4, const real_t> &X,
+                                        vd_regs3d_t<VDIM, DIM, MQ1> &Y)
+{
+   for (int d = 0; d < DIM; d++)
+   {
+      for (int dz = 0; dz < d1d; ++dz)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(dx, x, d1d)
+            {
+               Y[c][d][dz][dy][dx] = X(dx, dy, dz, c);
+            }
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
 /// Load 3D input VDIM*DIM vector into given register tensor, specific component
 template <int VDIM, int DIM, int MQ1>
 inline MFEM_HOST_DEVICE void LoadDofs3d(const int e, const int d1d, const int c,
@@ -332,6 +354,28 @@ inline MFEM_HOST_DEVICE void WriteDofs3d(const int e, const int d1d,
    MFEM_SYNC_THREAD;
 }
 
+/// Write 3D DIM vector into given device tensor for specific component
+template <int VDIM, int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void WriteDofs3d(const int d1d, const int c,
+                                         vd_regs3d_t<VDIM, DIM, MQ1> &X,
+                                         DeviceTensor<4, real_t> &Y)
+{
+   for (int dz = 0; dz < d1d; ++dz)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dy, y, d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dx, x, d1d)
+         {
+            for (int d = 0; d < DIM; ++d)
+            {
+               Y(dx, dy, dz, c) += X(c, d, dz, dy, dx);
+            }
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
 /// 2D scalar contraction, X direction
 template <bool Transpose, int MQ1>
 inline MFEM_HOST_DEVICE void ContractX2d(const int d1d, const int q1d,

fem/kernels3d.hpp

@@ -0,0 +1,332 @@
+// 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 "../config/config.hpp"
+#include "../general/forall.hpp"
+#include "../linalg/dtensor.hpp"
+
+#include "kernels.hpp" // IWYU pragma: keep
+
+namespace mfem::kernels::internal::low
+{
+
+#if ((defined(MFEM_USE_CUDA) && defined(__CUDA_ARCH__)) ||       \
+     (defined(MFEM_USE_HIP) && defined(__HIP_DEVICE_COMPILE__)))
+template <int DIM, int N>
+// struct regs3d_device_wrapper: mfem::future::tensor<real_t, DIM, 0, 0, 0> {};
+struct regs3d_device_wrapper: mfem::future::tensor<real_t, 0, 0, 0, DIM> {};
+template <int DIM, int N>
+using regs3d_t = regs3d_device_wrapper<DIM, N>;
+#else
+template <int DIM, int N>
+using regs3d_t = mfem::future::tensor<real_t, N, N, N, DIM>;
+// using regs3d_t = mfem::future::tensor<real_t, DIM, N, N, N>;
+#endif
+
+///////////////////////////////////////////////////////////////////////////////
+/// Load 2D matrix into shared memory
+template <int MQ1>
+inline MFEM_HOST_DEVICE void LoadMatrix(const int d1d, const int q1d,
+                                        const real_t *M, real_t (*N)[MQ1])
+{
+   if (MFEM_THREAD_ID(z) == 0)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dy, y, d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, q1d)
+         {
+            N[dy][qx] = M[dy * q1d + qx];
+         }
+      }
+   }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+template <int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void LoadDofs3d(const int e, const int d1d,
+                                        const DeviceTensor<5, const real_t> &XE,
+                                        real_t (&sm0)[MQ1][MQ1][MQ1][DIM])
+{
+   MFEM_FOREACH_THREAD_DIRECT(dy,y,d1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dx,x,d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dz,z,d1d)
+         {
+            sm0[dz][dy][dx][0] = XE(dx, dy, dz, 0, e);
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient, 1/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradX(const int d1d, const int q1d,
+                                   const real_t (*B)[MQ1],
+                                   const real_t (*G)[MQ1],
+                                   const real_t (&sm0)[MQ1][MQ1][MQ1][DIM],
+                                   real_t (&sm1)[MQ1][MQ1][MQ1][DIM])
+{
+   MFEM_FOREACH_THREAD_DIRECT(dz,z,d1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dy,y,d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx,x,q1d)
+         {
+            real_t u = 0.0, v = 0.0;
+            MFEM_UNROLL(MQ1)
+            for (int dx = 0; dx < d1d; ++dx)
+            {
+               const auto x = sm0[dz][dy][dx][0];
+               u = std::fma(B[dx][qx], x, u);
+               v = std::fma(G[dx][qx], x, v);
+            }
+            sm1[dz][dy][qx][0] = u;
+            sm1[dz][dy][qx][1] = v;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient, 2/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradY(const int d1d, const int q1d,
+                                   const real_t (*B)[MQ1],
+                                   const real_t (*G)[MQ1],
+                                   const real_t (&sm1)[MQ1][MQ1][MQ1][DIM],
+                                   real_t (&sm0)[MQ1][MQ1][MQ1][DIM])
+{
+   MFEM_FOREACH_THREAD_DIRECT(dz,z,d1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(qy,y,q1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx,x,q1d)
+         {
+            real_t u = 0.0, v = 0.0, w = 0.0;
+            MFEM_UNROLL(MQ1)
+            for (int dy = 0; dy < d1d; ++dy)
+            {
+               u = std::fma(sm1[dz][dy][qx][1], B[dy][qy], u);
+               v = std::fma(sm1[dz][dy][qx][0], G[dy][qy], v);
+               w = std::fma(sm1[dz][dy][qx][0], B[dy][qy], w);
+            }
+            sm0[dz][qy][qx][0] = u;
+            sm0[dz][qy][qx][1] = v;
+            sm0[dz][qy][qx][2] = w;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient, 3/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradZ(const int d1d, const int q1d,
+                                   const real_t (*B)[MQ1],
+                                   const real_t (*G)[MQ1],
+                                   const real_t (&sm0)[MQ1][MQ1][MQ1][DIM],
+                                   regs3d_t<DIM,MQ1> &reg)
+{
+   MFEM_FOREACH_THREAD_DIRECT(qz,z,q1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(qy,y,q1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx,x,q1d)
+         {
+            real_t u[3] = {0.0, 0.0, 0.0};
+            MFEM_UNROLL(MQ1)
+            for (int dz = 0; dz < d1d; ++dz)
+            {
+               u[0] = std::fma(B[dz][qz], sm0[dz][qy][qx][0], u[0]);
+               u[1] = std::fma(B[dz][qz], sm0[dz][qy][qx][1], u[1]);
+               u[2] = std::fma(G[dz][qz], sm0[dz][qy][qx][2], u[2]);
+            }
+            reg[qz][qy][qx][0] = u[0];
+            reg[qz][qy][qx][1] = u[1];
+            reg[qz][qy][qx][2] = u[2];
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D scalar gradient
+template <int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void Grad3d(const int d1d, const int q1d,
+                                    const real_t (*B)[MQ1],
+                                    const real_t (*G)[MQ1],
+                                    real_t (&sm0)[MQ1][MQ1][MQ1][DIM],
+                                    real_t (&sm1)[MQ1][MQ1][MQ1][DIM],
+                                    regs3d_t<DIM,MQ1> &reg)
+{
+   GradX(d1d, q1d, B, G, sm0, sm1); // Grad X
+   GradY(d1d, q1d, B, G, sm1, sm0); // Grad Y
+   GradZ(d1d, q1d, B, G, sm0, reg); // Grad Z
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient Transposed, 1/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradTranspose3dX(const int d1d, const int q1d,
+                                              const real_t (*B)[MQ1],
+                                              const real_t (*G)[MQ1],
+                                              regs3d_t<DIM,MQ1> &reg,
+                                              real_t (&sm1)[MQ1][MQ1][MQ1][DIM],
+                                              real_t (&sm0)[MQ1][MQ1][MQ1][DIM])
+
+{
+   MFEM_FOREACH_THREAD_DIRECT(qz,z,q1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(qy,y,q1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx,x,q1d)
+         {
+            sm1[qz][qy][qx][0] = reg[qz][qy][qx][0];
+            sm1[qz][qy][qx][1] = reg[qz][qy][qx][1];
+            sm1[qz][qy][qx][2] = reg[qz][qy][qx][2];
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+
+   MFEM_FOREACH_THREAD_DIRECT(qz,z,q1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(qy,y,q1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dx,x,d1d)
+         {
+            real_t u = 0.0, v = 0.0, w = 0.0;
+            MFEM_UNROLL(MQ1)
+            for (int qx = 0; qx < q1d; ++qx)
+            {
+               u = std::fma(sm1[qz][qy][qx][0], G[dx][qx], u);
+               v = std::fma(sm1[qz][qy][qx][1], B[dx][qx], v);
+               w = std::fma(sm1[qz][qy][qx][2], B[dx][qx], w);
+            }
+            sm0[qz][qy][dx][0] = u;
+            sm0[qz][qy][dx][1] = v;
+            sm0[qz][qy][dx][2] = w;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient Transposed, 2/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradTranspose3dY(const int d1d, const int q1d,
+                                              const real_t (*B)[MQ1],
+                                              const real_t (*G)[MQ1],
+                                              real_t (&sm0)[MQ1][MQ1][MQ1][DIM],
+                                              real_t (&sm1)[MQ1][MQ1][MQ1][DIM])
+{
+   MFEM_FOREACH_THREAD_DIRECT(qz,z,q1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dy,y,d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dx,x,d1d)
+         {
+            real_t u = 0.0, v = 0.0, w = 0.0;
+            MFEM_UNROLL(MQ1)
+            for (int qy = 0; qy < q1d; ++qy)
+            {
+               u = std::fma(sm0[qz][qy][dx][0], B[dy][qy], u);
+               v = std::fma(sm0[qz][qy][dx][1], G[dy][qy], v);
+               w = std::fma(sm0[qz][qy][dx][2], B[dy][qy], w);
+            }
+            sm1[qz][dy][dx][0] = u;
+            sm1[qz][dy][dx][1] = v;
+            sm1[qz][dy][dx][2] = w;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient Transposed, 3/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradTranspose3dZ(const int d1d, const int q1d,
+                                              const real_t (*B)[MQ1],
+                                              const real_t (*G)[MQ1],
+                                              real_t (&sm1)[MQ1][MQ1][MQ1][DIM],
+                                              regs3d_t<DIM,MQ1> &reg)
+{
+   MFEM_FOREACH_THREAD_DIRECT(dz,z,d1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dy,y,d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dx,x,d1d)
+         {
+            real_t u = 0.0, v = 0.0, w = 0.0;
+            MFEM_UNROLL(MQ1)
+            for (int qz = 0; qz < q1d; ++qz)
+            {
+               u = std::fma(sm1[qz][dy][dx][0], B[dz][qz], u);
+               v = std::fma(sm1[qz][dy][dx][1], B[dz][qz], v);
+               w = std::fma(sm1[qz][dy][dx][2], G[dz][qz], w);
+            }
+            reg[dz][dy][dx][0] = u;
+            reg[dz][dy][dx][1] = v;
+            reg[dz][dy][dx][2] = w;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D scalar gradient transposed
+template <int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void GradTranspose3d(const int d1d, const int q1d,
+                                             const real_t (*B)[MQ1],
+                                             const real_t (*G)[MQ1],
+                                             regs3d_t<DIM,MQ1> &reg,
+                                             real_t (&sm1)[MQ1][MQ1][MQ1][DIM],
+                                             real_t (&sm0)[MQ1][MQ1][MQ1][DIM])
+{
+   GradTranspose3dX(d1d, q1d, B, G, reg, sm1, sm0); // Grad^T X
+   GradTranspose3dY(d1d, q1d, B, G, sm0, sm1); // Grad^T Y
+   GradTranspose3dZ(d1d, q1d, B, G, sm1, reg); // Grad^T Z
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// 3D Scalar Gradient Transposed, 3/3
+template<int DIM, int MQ1>
+inline MFEM_HOST_DEVICE void WriteDofs3d(const int d1d,
+                                         const int c, const int e,
+                                         regs3d_t<DIM,MQ1> &reg,
+                                         const DeviceTensor<5, real_t> &YE)
+{
+   MFEM_FOREACH_THREAD_DIRECT(dz,z,d1d)
+   {
+      MFEM_FOREACH_THREAD_DIRECT(dy,y,d1d)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dx,x,d1d)
+         {
+            const real_t u = reg[dz][dy][dx][0];
+            const real_t v = reg[dz][dy][dx][1];
+            const real_t w = reg[dz][dy][dx][2];
+            YE(dx, dy, dz, c, e) += (u + v + w);
+         }
+      }
+   }
+}
+
+} // namespace mfem::kernels::internal

general/backends.hpp

@@ -27,6 +27,14 @@
 #endif
 #include "hip.hpp"
 
+#if defined(MFEM_USE_CUDA)
+#define Gpu(...) Cu##__VA_ARGS__
+#elif defined(MFEM_USE_HIP)
+#define Gpu(...) Hip##__VA_ARGS__
+#else
+#define Gpu(...) __VA_ARGS__
+#endif
+
 #ifdef MFEM_USE_OCCA
 #include "occa.hpp"
 #endif
@@ -48,6 +56,7 @@ constexpr bool mfem_use_gpu = false;
 #define MFEM_DEVICE
 #define MFEM_HOST
 #define MFEM_LAMBDA
+#define MFEM_CONSTANT
 // #define MFEM_HOST_DEVICE // defined in config/config.hpp
 // MFEM_DEVICE_SYNC is made available for debugging purposes
 #define MFEM_DEVICE_SYNC
@@ -65,6 +74,13 @@ constexpr bool mfem_use_gpu = false;
 #define MFEM_THREAD_SIZE(k) 1
 #define MFEM_FOREACH_THREAD(i,k,N) for(int i=0; i<N; i++)
 #define MFEM_FOREACH_THREAD_DIRECT(i,k,N) MFEM_FOREACH_THREAD(i,k,N)
+
+inline const void* MemcpyToSymbol(const void *d_sym, const void *h_src,
+                                  size_t bytes)
+{
+   memcpy(const_cast<void *>(d_sym), h_src, bytes);
+   return d_sym;
+}
 #endif
 
 // 'double' and 'float' atomicAdd implementation for previous versions of CUDA

general/cuda.cpp

@@ -175,6 +175,17 @@ void* CuMemcpyDtoHAsync(void *dst, const void *src, size_t bytes)
    return dst;
 }
 
+const void* CuMemcpyToSymbol(const void *d_sym, const void *h_src,
+                             size_t bytes)
+{
+#ifdef MFEM_USE_CUDA
+   MFEM_GPU_CHECK(cudaMemcpyToSymbol(d_sym, h_src, bytes));
+   return d_sym;
+#endif
+   MFEM_ABORT("CUDA has no shadow host copy of device symbols");
+   return memcpy(const_cast<void*>(d_sym), h_src, bytes);
+}
+
 void CuCheckLastError()
 {
 #ifdef MFEM_USE_CUDA

general/cuda.hpp

@@ -25,6 +25,8 @@ constexpr bool mfem_use_gpu = true;
 #define MFEM_HOST __host__
 #define MFEM_LAMBDA __host__
 #define MFEM_LAUNCH_BOUNDS __launch_bounds__
+#define MFEM_DEVICE __device__
+#define MFEM_CONSTANT __constant__
 // #define MFEM_HOST_DEVICE __host__ __device__ // defined in config/config.hpp
 #define MFEM_DEVICE_SYNC MFEM_GPU_CHECK(cudaDeviceSynchronize())
 #define MFEM_STREAM_SYNC MFEM_GPU_CHECK(cudaStreamSynchronize(0))
@@ -94,6 +96,10 @@ void* CuMemcpyDtoH(void *h_dst, const void *d_src, size_t bytes);
 /// Copies memory from Device to Host
 void* CuMemcpyDtoHAsync(void *h_dst, const void *d_src, size_t bytes);
 
+/// Copies data to the given symbol on the device.
+const void* CuMemcpyToSymbol(const void *d_sym, const void *h_src,
+                             size_t bytes);
+
 /// Check the error code returned by cudaGetLastError(), aborting on error.
 void CuCheckLastError();
 

general/forall.hpp

@@ -1090,6 +1090,12 @@ inline void forall_2D_batch(int N, int X, int Y, int BZ, lambda &&body)
    ForallWrap<2>(true, N, body, X, Y, BZ);
 }
 
+template<int MAX_THREADS_PER_BLOCK, typename lambda>
+inline void forall_2D_batch(int N, int X, int Y, int BZ, lambda &&body)
+{
+   ForallWrap<2, MAX_THREADS_PER_BLOCK>(true, N, body, X, Y, BZ);
+}
+
 template<typename lambda>
 inline void forall_3D(int N, int X, int Y, int Z, lambda &&body)
 {

general/hip.cpp

@@ -175,6 +175,15 @@ void* HipMemcpyDtoHAsync(void *dst, const void *src, size_t bytes)
    return dst;
 }
 
+const void* HipMemcpyToSymbol(const void *d_sym, const void *h_src,
+                              size_t bytes)
+{
+#ifdef MFEM_USE_HIP
+   MFEM_GPU_CHECK(hipMemcpyToSymbol(d_sym, h_src, bytes));
+#endif
+   return memcpy(const_cast<void*>(d_sym), h_src, bytes);
+}
+
 void HipCheckLastError()
 {
 #ifdef MFEM_USE_HIP

general/hip.hpp

@@ -21,8 +21,9 @@
 #if defined(MFEM_USE_HIP) && defined(__HIP__)
 #define MFEM_USE_CUDA_OR_HIP
 constexpr bool mfem_use_gpu = true;
-#define MFEM_DEVICE __device__
 #define MFEM_HOST __host__
+#define MFEM_DEVICE __device__
+#define MFEM_CONSTANT __constant__
 #define MFEM_LAMBDA __host__ __device__
 #define MFEM_LAUNCH_BOUNDS __launch_bounds__
 // #define MFEM_HOST_DEVICE __host__ __device__ // defined in config/config.hpp
@@ -96,6 +97,10 @@ void* HipMemcpyDtoH(void *h_dst, const void *d_src, size_t bytes);
 /// Copies memory from Device to Host
 void* HipMemcpyDtoHAsync(void *h_dst, const void *d_src, size_t bytes);
 
+/// Copies data to the given symbol on the device.
+const void* HipMemcpyToSymbol(const void *d_sym, const void *h_src,
+                              size_t bytes);
+
 /// Check the error code returned by hipGetLastError(), aborting on error.
 void HipCheckLastError();
 

general/nvtx.hpp

@@ -0,0 +1 @@
+../../stash/debug/nvtx.hpp

linalg/tensor.hpp

@@ -165,6 +165,23 @@ struct tensor<T, n0, n1, n2>
    tensor < T, n1, n2 > values[n0];
 };
 
+template < typename T, int n1, int n2 >
+struct tensor<T, 0, n1, n2>
+{
+   using type = T;
+   static constexpr int ndim      = 3;
+   static constexpr int first_dim = 0;
+   MFEM_HOST_DEVICE tensor< T, n1, n2 >& operator[](int /*i*/) { return values; }
+   MFEM_HOST_DEVICE const tensor< T, n1, n2 >& operator[](int /*i*/) const { return values; }
+   MFEM_HOST_DEVICE tensor< T, n1, n2 >& operator()(int /*i*/) { return values; }
+   MFEM_HOST_DEVICE const tensor< T, n1, n2 >& operator()(int /*i*/) const { return values; }
+   MFEM_HOST_DEVICE tensor< T, n2 >& operator()(int /*i*/, int j) { return values[j]; }
+   MFEM_HOST_DEVICE const tensor< T, n2 >& operator()(int i, int j) const { return values[i][j]; }
+   MFEM_HOST_DEVICE T& operator()(int /*i*/, int j, int k) { return values[j][k]; }
+   MFEM_HOST_DEVICE const T& operator()(int /*i*/, int j, int k) const { return values[j][k]; }
+   tensor < T, n1, n2 > values;
+};
+
 template < typename T, int n0, int n1, int n2, int n3 >
 struct tensor<T, n0, n1, n2, n3>
 {
@@ -184,6 +201,26 @@ struct tensor<T, n0, n1, n2, n3>
    tensor < T, n1, n2, n3 > values[n0];
 };
 
+template < typename T, int n1, int n2, int n3 >
+struct tensor<T, 0, n1, n2, n3>
+{
+   using type = T;
+   static constexpr int ndim      = 4;
+   static constexpr int first_dim = 0;
+   MFEM_HOST_DEVICE tensor< T, n1, n2, n3 >& operator[](int /*i*/) { return values; }
+   MFEM_HOST_DEVICE const tensor< T, n1, n2, n3 >& operator[](int /*i*/) const { return values; }
+   MFEM_HOST_DEVICE tensor< T, n1, n2, n3 >& operator()(int /*i*/) { return values; }
+   MFEM_HOST_DEVICE const tensor< T, n1, n2, n3 >& operator()(int /*i*/) const { return values; }
+   MFEM_HOST_DEVICE tensor< T, n2, n3 >& operator()(int /*i*/, int j) { return values[j]; }
+   MFEM_HOST_DEVICE const tensor< T, n2, n3 >& operator()(int /*i*/, int j) const { return values[j]; }
+   MFEM_HOST_DEVICE tensor< T, n3 >& operator()(int /*i*/, int j, int k) { return values[j][k]; }
+   MFEM_HOST_DEVICE const tensor< T, n3 >& operator()(int /*i*/, int j,
+                                                      int k) const { return values[j][k]; }
+   MFEM_HOST_DEVICE T& operator()(int /*i*/, int j, int k, int l) { return values[j][k][l]; }
+   MFEM_HOST_DEVICE const T&  operator()(int /*i*/, int j, int k, int l) const { return values[j][k][l]; }
+   tensor < T, n1, n2, n3 > values;
+};
+
 template < typename T, int n0, int n1, int n2, int n3, int n4 >
 struct tensor<T, n0, n1, n2, n3, n4>
 {

makefile

@@ -807,7 +807,7 @@ FORMAT_EXCLUDE = general/tinyxml2.cpp tests/unit/catch.hpp
 FORMAT_LIST = $(filter-out $(FORMAT_EXCLUDE),$(wildcard $(FORMAT_FILES)))
 
 COUT_CERR_FILES = $(foreach dir,$(DIRS),$(dir)/*.[ch]pp)
-COUT_CERR_EXCLUDE = '^general/error\.cpp' '^general/globals\.[ch]pp'
+COUT_CERR_EXCLUDE = '^general/error\.cpp' '^general/globals\.[ch]pp' '^general/nvtx\.hpp'
 
 DEPRECATION_WARNING := \
 "This feature is planned for removal in the next release."\

tests/benchmarks/CMakeLists.txt

@@ -32,7 +32,11 @@ function(add_benchmark name)
     endif(MFEM_USE_CUDA)
 
     add_executable(bench_${name} ${${NAME}_BENCH_SRCS})
-    target_link_libraries(bench_${name} mfem pthread)
+    if (fmt_FOUND)
+        target_link_libraries(bench_${name} mfem pthread fmt::fmt)
+    else()
+        target_link_libraries(bench_${name} mfem pthread)
+    endif()
     add_dependencies(${MFEM_ALL_BENCHMARKS_TARGET_NAME} bench_${name})
 
     add_test(NAME bench_${name}_cpu
@@ -51,6 +55,7 @@ endfunction(add_benchmark)
 #-------------------------------------------------------------------------------
 add_benchmark(assembly_levels)
 add_benchmark(ceed)
+add_benchmark(dfem)
 add_benchmark(dg_amr)
 add_benchmark(elasticity)
 add_benchmark(tmop)

tests/benchmarks/bench_dfem.cpp

@@ -0,0 +1,845 @@
+// 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 "bench.hpp" // IWYU pragma: keep
+
+#ifdef MFEM_USE_BENCHMARK
+
+#include <cstdlib>
+#include <memory>
+
+#include <fem/qinterp/det.cpp>
+#include <fem/qinterp/grad.hpp> // IWYU pragma: keep
+#include "fem/integ/lininteg_domain_kernels.hpp" // IWYU pragma: keep
+
+#include "fem/dfem/doperator.hpp"
+#include <linalg/tensor.hpp>
+
+#include <fem/kernels3d.hpp>
+namespace ker = mfem::kernels::internal;
+namespace low = mfem::kernels::internal::low;
+
+#include "bench_dfem_mma.hpp"
+
+#undef NVTX_COLOR
+#define NVTX_COLOR ::nvtx::kNvidia
+
+using namespace mfem;
+
+using mfem::future::tuple;
+using mfem::future::tensor;
+
+using future::DifferentiableOperator;
+using future::UniformParameterSpace;
+using future::ParameterFunction;
+using future::FieldDescriptor;
+using future::make_tensor;
+using future::Gradient;
+using future::Weight;
+using future::Identity;
+
+/// info //////////////////////////////////////////////////////////////////////
+static void DumpVersionInfo()
+{
+   mfem::out << "\x1b[33m";
+   mfem::out << "version 0: PA std" << std::endl;
+   mfem::out << "version 1: PA reg" << std::endl; // can do high order
+   mfem::out << "version 2: PA low" << std::endl;
+   mfem::out << "version 3: PA mma" << std::endl;
+   // mfem::out << "version 4: PA ∂fem new, not specialized" << std::endl;
+   mfem::out << "version 5: PA ∂fem new, specialized" << std::endl;
+   // mfem::out << "version 6: PA ∂fem std" << std::endl; // ⚠️ max p=3
+   // mfem::out << "version 7: MF ∂fem std" << std::endl;
+   // mfem::out << "version 8: MF ∂fem new" << std::endl; // ⚠️ not supported
+   mfem::out << "\x1b[m" << std::endl;
+}
+
+// Custom benchmark arguments generator ///////////////////////////////////////
+static void CustomArguments(bm::Benchmark *b) noexcept
+{
+   constexpr int MAX_NDOFS = 8 * 1024 * (mfem_use_gpu ? 1024 : 8);
+
+   const auto versions = { 0, 1, 2, 3, /*4,*/ 5, /*6, 7, 8*/ };
+
+   const auto orders = { 6, 5, 4, 3, 2, 1 };
+
+   constexpr auto ndofs = [](int n) constexpr noexcept -> int
+   {
+      return (n + 1) * (n + 1) * (n + 1);
+   };
+
+   constexpr auto inc = [](int n) constexpr noexcept -> int
+   {
+      return n < 160 ?  4 : n < 240 ?  8 : n < 320 ? 16 : 32;
+   };
+
+   for (auto k : versions)
+   {
+      for (auto p : orders)
+      {
+         for (int n = 16; ndofs(n) <= MAX_NDOFS; n += inc(n))
+         {
+            b->Args({k, p, n});
+         }
+      }
+   }
+}
+
+/// Basic Kernels Specializations /////////////////////////////////////////////
+static void AddBasicKernelSpecializations()
+{
+   using Det = QuadratureInterpolator::DetKernels;
+   Det::Specialization<3, 3, 2, 2>::Add();
+   Det::Specialization<3, 3, 2, 3>::Add();
+   Det::Specialization<3, 3, 2, 5>::Add();
+   Det::Specialization<3, 3, 2, 6>::Add();
+   // Others might exceed memory limits
+
+   using Grad = QuadratureInterpolator::GradKernels;
+   Grad::Specialization<3, QVectorLayout::byVDIM,  false, 3, 2, 3>::Add();
+   Grad::Specialization<3, QVectorLayout::byVDIM,  false, 3, 2, 4>::Add();
+   Grad::Specialization<3, QVectorLayout::byVDIM,  false, 3, 2, 5>::Add();
+   Grad::Specialization<3, QVectorLayout::byVDIM,  false, 3, 2, 6>::Add();
+   Grad::Specialization<3, QVectorLayout::byVDIM,  false, 3, 2, 7>::Add();
+   Grad::Specialization<3, QVectorLayout::byVDIM,  false, 3, 2, 8>::Add();
+   Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 7>::Add();
+   Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 8>::Add();
+
+   using LIN = DomainLFIntegrator::AssembleKernels;
+   LIN::Specialization<3, 7, 7>::Add();
+   LIN::Specialization<3, 6, 6>::Add();
+   LIN::Specialization<3, 8, 8>::Add();
+}
+
+/// Globals ///////////////////////////////////////////////////////////////////
+Device *device_ptr = nullptr;
+static int gD1D = 0, gQ1D = 0;
+
+/// StiffnessIntegrator ///////////////////////////////////////////////////////
+struct StiffnessIntegrator : public BilinearFormIntegrator
+{
+   const FiniteElementSpace *fes;
+   const real_t *B, *G, *DX;
+   int ne, d1d, q1d;
+   Vector J0, dx;
+   Vector &qdata;
+
+public:
+   StiffnessIntegrator(Vector &qdata): qdata(qdata)
+   {
+      StiffnessKernels::Specialization<2,3>::Add();  // 1
+      StiffnessKernels::Specialization<3,4>::Add();  // 2
+      StiffnessKernels::Specialization<4,5>::Add();  // 3
+      StiffnessKernels::Specialization<5,6>::Add();  // 4
+      StiffnessKernels::Specialization<6,7>::Add();  // 5
+      StiffnessKernels::Specialization<7,8>::Add();  // 6
+      StiffnessKernels::Specialization<9,10>::Add();  // 8
+   }
+
+   void AssemblePA(const FiniteElementSpace &fespace) override
+   {
+      NVTX();
+      fes = &fespace;
+      auto *mesh = fes->GetMesh();
+      const int DIM = mesh->Dimension();
+      ne = mesh->GetNE();
+      const auto p = fes->GetFE(0)->GetOrder();
+      const auto q = 2 * p + mesh->GetElementTransformation(0)->OrderW();
+      const auto type = mesh->GetElementBaseGeometry(0);
+      const IntegrationRule &ir = IntRules.Get(type, q);
+      const int NQPT = ir.GetNPoints();
+      d1d = p + 1;
+      q1d = IntRules.Get(Geometry::SEGMENT, ir.GetOrder()).GetNPoints();
+      MFEM_VERIFY(d1d == gD1D, "D1D mismatch: " << d1d << " != " << gD1D);
+      MFEM_VERIFY(q1d == gQ1D, "Q1D mismatch: " << q1d << " != " << gQ1D);
+      MFEM_VERIFY(NQPT == q1d * q1d * q1d, "");
+      const DofToQuad *maps =
+         &fes->GetFE(0)->GetDofToQuad(ir, DofToQuad::TENSOR);
+      const GridFunction *nodes = (mesh->EnsureNodes(), mesh->GetNodes());
+      const FiniteElementSpace *nfes = nodes->FESpace();
+      const int nVDIM = nfes->GetVDim();
+      dx.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType());
+      J0.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType());
+      dx.UseDevice(true), J0.UseDevice(true);
+      B = maps->B.Read(), G = maps->G.Read(), DX = dx.Read();
+
+      const Operator *NR =
+         nfes->GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+      const QuadratureInterpolator *nqi = nfes->GetQuadratureInterpolator(ir);
+      nqi->SetOutputLayout(QVectorLayout::byVDIM);
+      const int nd = nfes->GetFE(0)->GetDof();
+      Vector xe(nVDIM * nd * ne, Device::GetDeviceMemoryType());
+      NR->Mult(*nodes, (xe.UseDevice(true), xe));
+      nqi->Derivatives(xe, J0);
+
+      const int Q1D = q1d;
+      const auto w_r = ir.GetWeights().Read();
+      const auto W = Reshape(w_r, q1d, q1d, q1d);
+      const auto J = Reshape(J0.Read(), 3, 3, q1d, q1d, q1d, ne);
+      auto DX_w = Reshape(dx.Write(), 3, 3, q1d, q1d, q1d, ne);
+
+      mfem::forall_3D(ne, Q1D, Q1D, Q1D,[=] MFEM_HOST_DEVICE(int e)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qz, z, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  const real_t w = W(qx, qy, qz);
+                  const real_t *Jtr = &J(0, 0, qx, qy, qz, e);
+                  const real_t detJ = kernels::Det<3>(Jtr);
+                  const real_t wd = w * detJ;
+                  const real_t D[9] = { wd, 0.0, 0.0,
+                                        0.0, wd, 0.0,
+                                        0.0, 0.0, wd
+                                      };
+                  real_t Jrt[9], A[9];
+                  kernels::CalcInverse<3>(Jtr, Jrt);
+                  kernels::MultABt(3, 3, 3, D, Jrt, A);
+                  kernels::Mult(3, 3, 3, A, Jrt, &DX_w(0, 0, qx, qy, qz, e));
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      });
+      qdata = dx;
+   }
+
+   //////////////////////////////////////////////////////////////////
+   template <int T_D1D = 0, int T_Q1D = 0>
+   static void StiffnessMult(const int NE, const real_t *b, const real_t *g,
+                             const real_t *dx, const real_t *xe, real_t *ye,
+                             const int d1d, const int q1d)
+   {
+      const int D1D = T_D1D ? T_D1D : d1d;
+      const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+      constexpr int DIM = 3, VDIM = 1;
+      const auto XE = Reshape(xe, D1D, D1D, D1D, VDIM, NE);
+      const auto DX = Reshape(dx, 3, 3, Q1D, Q1D, Q1D, NE);
+      auto YE = Reshape(ye, D1D, D1D, D1D, VDIM, NE);
+
+      mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+      {
+         constexpr int MD1 = T_D1D > 0 ? kernels::internal::SetMaxOf(T_D1D) : 8;
+         constexpr int MQ1 = T_Q1D > 0 ? kernels::internal::SetMaxOf(T_Q1D) : 8;
+
+         MFEM_SHARED real_t smem[MQ1][MQ1];
+         MFEM_SHARED real_t sB[MD1][MQ1], sG[MD1][MQ1];
+         ker::vd_regs3d_t<VDIM, DIM, MQ1> r0, r1;
+
+         ker::LoadMatrix(D1D, Q1D, b, sB);
+         ker::LoadMatrix(D1D, Q1D, g, sG);
+
+         ker::LoadDofs3d(e, D1D, XE, r0);
+         ker::Grad3d(D1D, Q1D, smem, sB, sG, r0, r1);
+
+         for (int qz = 0; qz < Q1D; qz++)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  real_t v[3], u[3] = { r1[0][0][qz][qy][qx],
+                                        r1[0][1][qz][qy][qx],
+                                        r1[0][2][qz][qy][qx]
+                                      };
+                  const real_t *dx = &DX(0, 0, qx, qy, qz, e);
+                  kernels::Mult(3, 3, dx, u, v);
+                  r0[0][0][qz][qy][qx] = v[0];
+                  r0[0][1][qz][qy][qx] = v[1];
+                  r0[0][2][qz][qy][qx] = v[2];
+               }
+            }
+         }
+         ker::GradTranspose3d(D1D, Q1D, smem, sB, sG, r0, r1);
+         ker::WriteDofs3d(e, D1D, r1, YE);
+      });
+   }
+
+   using StiffnessKernelType = decltype(&StiffnessMult<>);
+   MFEM_REGISTER_KERNELS(StiffnessKernels, StiffnessKernelType, (int, int));
+
+   void AddMultPA(const Vector &x, Vector &y) const override
+   {
+      db1("\x1b[32md1d:{} q1d:{}", d1d, q1d);
+      StiffnessKernels::Run(d1d, q1d,
+                            ne, B, G, DX, x.Read(), y.ReadWrite(),
+                            d1d, q1d);
+   }
+};
+
+template <int D1D, int Q1D>
+StiffnessIntegrator::StiffnessKernelType
+StiffnessIntegrator::StiffnessKernels::Kernel()
+{
+   db1("D1D:{} Q1D:{}", D1D, Q1D);
+   return StiffnessMult<D1D, Q1D>;
+}
+
+StiffnessIntegrator::StiffnessKernelType
+StiffnessIntegrator::StiffnessKernels::Fallback([[maybe_unused]] int d1d,
+                                                [[maybe_unused]] int q1d)
+{
+   dbg("\x1b[33mFallback d1d:{} q1d:{}", d1d, q1d);
+   // MFEM_ABORT("No kernel for d1d=" << d1d << " q1d=" << q1d);
+   // return nullptr;
+   return StiffnessMult<>;
+}
+
+/// PADiffLowIntegrator ///////////////////////////////////////////////////////
+struct PADiffLowIntegrator : public BilinearFormIntegrator
+{
+   const FiniteElementSpace *fes;
+   const real_t *B, *G, *DX;
+   int ne, d1d, q1d;
+   Vector J0, dx;
+
+public: // for nvcc
+   //////////////////////////////////////////////////////////////////
+   template <int T_Q1D = 0>
+   static void PADiffLowMult(const int ne, const int d1d,
+                             const real_t *b, const real_t *g,
+                             const real_t *dx, const real_t *xe,
+                             real_t *ye,
+                             const int q1d)
+   {
+      constexpr int DIM = 3, VDIM = 1;
+
+      const auto XE = Reshape(xe, d1d, d1d, d1d, VDIM, ne);
+      auto YE = Reshape(ye, d1d, d1d, d1d, VDIM, ne);
+
+      mfem::forall_3D<T_Q1D*T_Q1D*T_Q1D>(ne, q1d, q1d, q1d,
+                                         [=] MFEM_HOST_DEVICE(int e)
+      {
+         constexpr int MQ1 = T_Q1D;
+
+         MFEM_SHARED real_t sm0[MQ1][MQ1][MQ1][3];
+         MFEM_SHARED real_t sm1[MQ1][MQ1][MQ1][3];
+         MFEM_SHARED real_t sB[MQ1][MQ1];
+         MFEM_SHARED real_t sG[MQ1][MQ1];
+
+         low::regs3d_t<DIM, MQ1> reg;
+
+         low::LoadMatrix(d1d, q1d, b, sB);
+         low::LoadMatrix(d1d, q1d, g, sG);
+         low::LoadDofs3d(e, d1d, XE, sm0); // Load & sync
+
+         // Grad: sm0 -X-> sm1 -Y-> sm0 -Z-> reg
+         low::Grad3d(d1d, q1d, sB, sG, sm0, sm1, reg); // Grad 3D
+
+         // Q-function
+         MFEM_FOREACH_THREAD_DIRECT(qz,z,q1d)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy,y,q1d)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx,x,q1d)
+               {
+                  // pull
+                  real_t v[3], u[3] = { reg[qz][qy][qx][0],
+                                        reg[qz][qy][qx][1],
+                                        reg[qz][qy][qx][2]
+                                      };
+                  //  Q-function
+                  kernels::Mult(3, 3, dx + 9*(qx*q1d*q1d + qy*q1d + qz), u, v);
+                  // push
+                  reg[qz][qy][qx][0] = v[0];
+                  reg[qz][qy][qx][1] = v[1];
+                  reg[qz][qy][qx][2] = v[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         // Grad^T: reg -=-> sm1 -X^T-> sm0 -Y^T-> sm1 -Z^T-> reg -> YE
+         low::GradTranspose3d(d1d, q1d, sB, sG, reg, sm1, sm0); // Grad^T 3D
+         low::WriteDofs3d(d1d, 0, e, reg, YE); // Write YE
+      });
+   }
+
+   using PADiffLowKernelType = decltype(&PADiffLowMult<>);
+   MFEM_REGISTER_KERNELS(PADiffLowKernels, PADiffLowKernelType, (int));
+
+public:
+   PADiffLowIntegrator()
+   {
+      PADiffLowKernels::Specialization<3>::Add();  // 1
+      PADiffLowKernels::Specialization<4>::Add();  // 2
+      PADiffLowKernels::Specialization<5>::Add();  // 3
+      PADiffLowKernels::Specialization<6>::Add();  // 4
+      PADiffLowKernels::Specialization<7>::Add();  // 5
+      PADiffLowKernels::Specialization<8>::Add();  // 6
+   }
+
+   void AssemblePA(const FiniteElementSpace &fespace) override
+   {
+      NVTX();
+      fes = &fespace;
+      auto *mesh = fes->GetMesh();
+      const int DIM = mesh->Dimension();
+      ne = mesh->GetNE();
+      const auto p = fes->GetFE(0)->GetOrder();
+      const auto q = 2 * p + mesh->GetElementTransformation(0)->OrderW();
+      const auto type = mesh->GetElementBaseGeometry(0);
+      const IntegrationRule &ir = IntRules.Get(type, q);
+      const int NQPT = ir.GetNPoints();
+      d1d = p + 1;
+      q1d = IntRules.Get(Geometry::SEGMENT, ir.GetOrder()).GetNPoints();
+      MFEM_VERIFY(d1d == gD1D, "D1D mismatch: " << d1d << " != " << gD1D);
+      MFEM_VERIFY(q1d == gQ1D, "Q1D mismatch: " << q1d << " != " << gQ1D);
+      MFEM_VERIFY(NQPT == q1d * q1d * q1d, "");
+      const DofToQuad *maps =
+         &fes->GetFE(0)->GetDofToQuad(ir, DofToQuad::TENSOR);
+      const GridFunction *nodes = (mesh->EnsureNodes(), mesh->GetNodes());
+      const FiniteElementSpace *nfes = nodes->FESpace();
+      const int nVDIM = nfes->GetVDim();
+      dx.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType());
+      J0.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType());
+      dx.UseDevice(true), J0.UseDevice(true);
+      B = maps->B.Read(), G = maps->G.Read(), DX = dx.Read();
+
+      const Operator *NR =
+         nfes->GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+      const QuadratureInterpolator *nqi = nfes->GetQuadratureInterpolator(ir);
+      nqi->SetOutputLayout(QVectorLayout::byVDIM);
+      const int nd = nfes->GetFE(0)->GetDof();
+      Vector xe(nVDIM * nd * ne, Device::GetDeviceMemoryType());
+      NR->Mult(*nodes, (xe.UseDevice(true), xe));
+      nqi->Derivatives(xe, J0);
+
+      const int Q1D = q1d;
+      const auto w_r = ir.GetWeights().Read();
+      const auto W = Reshape(w_r, q1d, q1d, q1d);
+      const auto J = Reshape(J0.Read(), 3, 3, q1d, q1d, q1d, ne);
+      auto DX_w = Reshape(dx.Write(), 3, 3, q1d, q1d, q1d, ne);
+
+      mfem::forall_3D(ne, Q1D, Q1D, Q1D,[=] MFEM_HOST_DEVICE(int e)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qz, z, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  const real_t w = W(qx, qy, qz);
+                  const real_t *Jtr = &J(0, 0, qx, qy, qz, e);
+                  const real_t detJ = kernels::Det<3>(Jtr);
+                  const real_t wd = w * detJ;
+                  const real_t D[9] = { wd, 0.0, 0.0,
+                                        0.0, wd, 0.0,
+                                        0.0, 0.0, wd
+                                      };
+                  real_t Jrt[9], A[9];
+                  kernels::CalcInverse<3>(Jtr, Jrt);
+                  kernels::MultABt(3, 3, 3, D, Jrt, A);
+                  kernels::Mult(3, 3, 3, A, Jrt, &DX_w(0, 0, qz, qy, qx, e));
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      });
+   }
+
+   void AddMultPA(const Vector &x, Vector &y) const override
+   {
+      db1("\x1b[32md1d:{} q1d:{}", d1d, q1d);
+      PADiffLowKernels::Run(q1d,
+                            ne, d1d, B, G, DX, x.Read(), y.ReadWrite(),
+                            q1d);
+   }
+};
+template <int Q1D>
+PADiffLowIntegrator::PADiffLowKernelType
+PADiffLowIntegrator::PADiffLowKernels::Kernel()
+{
+   db1("Q1D:{}", Q1D);
+   return PADiffLowMult<Q1D>;
+}
+
+PADiffLowIntegrator::PADiffLowKernelType
+PADiffLowIntegrator::PADiffLowKernels::Fallback(int q1d)
+{
+   dbg("\x1b[33mFallback d1d:{} q1d:{}", q1d);
+   MFEM_ABORT("No kernel for q1d=" << q1d);
+   return nullptr;
+   // return StiffnessMult<>;
+}
+
+/// BakeOff ///////////////////////////////////////////////////////////////////
+template <int VDIM, bool GLL>
+struct BakeOff
+{
+   static constexpr int DIM = 3;
+   const int p, c, q, n, nx, ny, nz;
+   const bool check_x, check_y, check_z, checked;
+   Mesh smesh;
+   ParMesh pmesh;
+   H1_FECollection fec;
+   ParFiniteElementSpace pfes;
+   const Geometry::Type geom_type;
+   IntegrationRules irs;
+   const IntegrationRule *ir;
+   ConstantCoefficient one;
+   Vector uvec;
+   VectorConstantCoefficient unit_vec;
+   const int dofs;
+   ParGridFunction *nodes;
+   ParFiniteElementSpace& mfes;
+   ParGridFunction x, y;
+   ParBilinearForm a;
+   std::unique_ptr<DifferentiableOperator> dop;
+   const int elem_size, total_size, d1d, q1d;
+   UniformParameterSpace qd_ps;
+   ParameterFunction qdata;
+
+   double mdofs{};
+
+   BakeOff(int p, int side):
+      p(p), c(side), q(2 * p + (GLL ? -1 : 3)), n((assert(c >= p), c / p)),
+      nx(n + (p * (n + 1) * p * n * p * n < c * c * c ? 1 : 0)),
+      ny(n + (p * (n + 1) * p * (n + 1) * p * n < c * c * c ? 1 : 0)), nz(n),
+      check_x(p * nx * p * ny * p * nz <= c * c * c),
+      check_y(p * (nx + 1) * p * (ny + 1) * p * nz > c * c * c),
+      check_z(p * (nx + 1) * p * (ny + 1) * p * (nz + 1) > c * c * c),
+      checked((assert(check_x &&check_y &&check_z), true)),
+      smesh(Mesh::MakeCartesian3D(nx, ny, nz, Element::HEXAHEDRON)),
+      pmesh(MPI_COMM_WORLD, (smesh.EnsureNodes(), smesh)),
+      fec(p, DIM, BasisType::GaussLobatto),
+      pfes(&pmesh, &fec, VDIM),//, Ordering::byNODES),
+      geom_type(pmesh.GetTypicalElementGeometry()),
+      irs(0, GLL ? Quadrature1D::GaussLobatto : Quadrature1D::GaussLegendre),
+      ir(&irs.Get(geom_type, q)), one(1.0), uvec(DIM),
+      unit_vec((uvec = 1.0, uvec /= uvec.Norml2(), uvec)),
+      dofs(pfes.GetTrueVSize()),
+      nodes(static_cast<ParGridFunction*>(pmesh.GetNodes())),
+      mfes(*nodes->ParFESpace()),
+      x(&pfes),
+      y(&pfes),
+      a(&pfes),
+      elem_size(DIM * DIM * ir->GetNPoints()),
+      total_size(elem_size * pmesh.GetNE()),
+      d1d(p + 1),
+      q1d(IntRules.Get(Geometry::SEGMENT, ir->GetOrder()).GetNPoints()),
+      qd_ps(pmesh, *ir, DIM*DIM),
+      qdata(qd_ps)
+   {
+      NVTX_MARK_FUNCTION;
+      dbg("p:{} q:{}", p, q);
+      smesh.Clear();
+      x = 0.0;
+
+      gD1D = d1d, gQ1D = q1d;
+      dbg("D1D: {}, Q1D: {}", gD1D, gQ1D);
+      qdata.UseDevice(true);
+      qdata = 0.0;
+      MFEM_VERIFY(q1d*q1d*q1d == ir->GetNPoints(), "");
+   }
+
+   virtual void Benchmark() { MFEM_ABORT("Not implemented."); }
+
+   [[nodiscard]] double SumMdofs() const noexcept { return mdofs; }
+
+   [[nodiscard]] double MDofs() const noexcept { return 1e-6 * dofs; }
+};
+
+/// Q-Functions ///////////////////////////////////////////////////////////////
+template<int DIM>
+struct MFApply
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, DIM>& Gu,
+                   const tensor<real_t, DIM, DIM>& J,
+                   const real_t& w) const
+   {
+      auto invJ = inv(J);
+      return tuple{((Gu * invJ)) * transpose(invJ) * det(J) * w};
+   }
+};
+
+template<int DIM>
+struct PASetup
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()([[maybe_unused]] const real_t &u,
+                   const tensor<real_t, DIM, DIM> &J,
+                   const real_t &w)const
+   {
+      return tuple{inv(J) * transpose(inv(J)) * det(J) * w};
+   }
+};
+
+
+template<int DIM>
+struct PAApply
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, DIM> &Gu,
+                   const tensor<real_t, DIM, DIM> &q) const
+   {
+      return tuple{q * Gu};
+   };
+};
+
+/// Diffusion /////////////////////////////////////////////////////////////////
+template <int VDIM = 1, bool GLL = false>
+struct Diffusion : public BakeOff<VDIM, GLL>
+{
+   static constexpr int DIM = 3;
+   static constexpr int U = 0, Ξ = 1, Q = 2;
+
+   const real_t rtol = 0.0;
+   const int max_it = 32, print_lvl = -1;
+
+   Array<int> ess_tdof_list, ess_bdr, all_domain_attr;
+   ParLinearForm b;
+   FieldDescriptor u_fd, Ξ_fd, q_fd;
+   std::vector<FieldDescriptor> u_sol, q_param, Ξ_q_params;
+   OperatorPtr A;
+   Operator *A_ptr;
+   Vector B, X;
+   CGSolver cg;
+
+   using BakeOff<VDIM, GLL>::a;
+   using BakeOff<VDIM, GLL>::ir;
+   using BakeOff<VDIM, GLL>::one;
+   using BakeOff<VDIM, GLL>::pmesh;
+   using BakeOff<VDIM, GLL>::pfes;
+   using BakeOff<VDIM, GLL>::mfes;
+   using BakeOff<VDIM, GLL>::x;
+   using BakeOff<VDIM, GLL>::y;
+   using BakeOff<VDIM, GLL>::mdofs;
+   using BakeOff<VDIM, GLL>::dop;
+   using BakeOff<VDIM, GLL>::nodes;
+   using BakeOff<VDIM, GLL>::qdata;
+   using BakeOff<VDIM, GLL>::qd_ps;
+   using BakeOff<VDIM, GLL>::dofs;
+
+   Diffusion(int version, int order, int side):
+      BakeOff<VDIM, GLL>(order, side),
+      ess_bdr(pmesh.bdr_attributes.Max()),
+      all_domain_attr(pmesh.bdr_attributes.Max()),
+      b(&pfes),
+      u_fd{U, &pfes}, Ξ_fd{Ξ, &mfes}, q_fd{Q, &qd_ps},
+      u_sol{u_fd},
+      q_param {q_fd},
+      Ξ_q_params {Ξ_fd, q_fd},
+      cg(MPI_COMM_WORLD)
+   {
+      // dbg("pmesh.bdr_attributes.Max():{}",pmesh.bdr_attributes.Max());
+      static_assert(VDIM == 1 && GLL == false);
+
+      ess_bdr = 1;
+      all_domain_attr = 1;
+      pfes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+      b.AddDomainIntegrator(new DomainLFIntegrator(this->one));
+      b.UseFastAssembly(true);
+      b.Assemble();
+
+      // MF setup ///////////////////////////////////////////////////
+      const auto dMFOperatorSetup = [&] (bool use_new_kernels,
+                                         bool use_kernels_specialization)
+      {
+         dbg("MF ∂fem {} kernels", use_new_kernels ? "NEW" : "STD");
+         std::vector<FieldDescriptor> in_fds = {{U, &pfes}, {Ξ, &mfes}};
+         std::vector<FieldDescriptor> out_fds = {{U, &pfes}};
+         dop = std::make_unique<DifferentiableOperator>(in_fds, out_fds, pmesh);
+         // dop->SetParameters({nodes});
+         if (use_kernels_specialization) { dop->UseKernelsSpecialization(); }
+         if (use_new_kernels) { dop->UseNewKernels(); }
+         // MFApply<DIM> mf_apply;
+         // dop->AddDomainIntegrator(mf_apply,
+         //                          tuple{Gradient<U>{}, Gradient<Ξ>{}, Weight{}}, // local API 🔥
+         //                          tuple{Gradient<U>{}},
+         //                          *ir, ess_bdr);
+         dop->FormLinearSystem(ess_tdof_list, x, b, A_ptr, X, B);
+         A.Reset(A_ptr);
+      };
+
+      // PA setup ///////////////////////////////////////////////////
+      const auto dPAOperatorSetup = [&] (bool use_new_kernels,
+                                         bool use_kernels_specialization)
+      {
+#if 0
+         dbg("[PA ∂fem] Setup");
+         auto Iu = Identity<U> {};
+         auto GΞ = Gradient<Ξ> {};
+         auto W = Weight{};
+         tuple Iu_GΞ_W = {Iu, GΞ, W};
+         PASetup<DIM> pa_setup_qf;
+         DifferentiableOperator dSetup(u_sol, Ξ_q_params, pmesh);
+         if (use_kernels_specialization) { dSetup.UseKernelsSpecialization(); }
+         if (use_new_kernels) { dSetup.UseNewKernels(); }
+         dSetup.AddDomainIntegrator(pa_setup_qf, Iu_GΞ_W, tuple{Iq}, *ir, ess_bdr);
+         dSetup.SetParameters({nodes, &qdata});
+         X.SetSize(pfes.GetTrueVSize());
+         pfes.GetRestrictionMatrix()->Mult(x, X);
+         dSetup.Mult(X, qdata);
+#else
+         dbg("[PA ∂fem] Setup (borrowing PA setup)");
+         {
+            ParBilinearForm bf(&pfes);
+            bf.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+            bf.AddDomainIntegrator(new StiffnessIntegrator(qdata));
+            bf.Assemble();
+         }
+#endif
+         dbg("[PA ∂fem] Apply");
+         // auto Iq = Identity<Q> {};
+         // auto Gu = Gradient<U> {};
+         // tuple Gu_Iq = {Gu, Iq};
+         // PAApply<DIM> pa_apply_qf;
+         dop = std::make_unique<DifferentiableOperator>(u_sol, q_param, pmesh);
+         dop->SetMultLevel(DifferentiableOperator::MultLevel::LVECTOR);
+         if (use_kernels_specialization) { dop->UseKernelsSpecialization(); }
+         if (use_new_kernels) { dop->UseNewKernels(); }
+         else { dbg("[PA ∂fem] NOT using kernels specialization"); }
+         // dop->AddDomainIntegrator(pa_apply_qf, Gu_Iq, tuple{Gu}, *ir, ess_bdr); // local API 🔥
+         assert(qdata*qdata > 0.0);
+         // dop->SetParameters({ &qdata });
+         dop->FormLinearSystem(ess_tdof_list, x, b, A_ptr, X, B);
+         A.Reset(A_ptr);
+         dbg("[PA ∂fem] done");
+      };
+
+      if (version <= 3) // std, reg, low & mma
+      {
+         a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+         if (version == 0) { a.AddDomainIntegrator(new DiffusionIntegrator(ir)); }
+         if (version == 1) { a.AddDomainIntegrator(new StiffnessIntegrator(qdata)); }
+         if (version == 2) { a.AddDomainIntegrator(new PADiffLowIntegrator()); }
+         if (version == 3) { a.AddDomainIntegrator(new PADiffMmaIntegrator()); }
+         a.Assemble();
+         a.FormLinearSystem(ess_tdof_list, x, b, A, X, B);
+         if (version == 0)
+         {
+            BilinearFormIntegrator *bfi = a.GetDBFI()->operator[](0);
+            auto *di = dynamic_cast<DiffusionIntegrator*>(bfi);
+            assert(di);
+            const int d1d = di->dofs1D, q1d = di->quad1D;
+            // dbg("\x1b[33md1d:{} q1d:{}", d1d, q1d);
+            MFEM_VERIFY(d1d == gD1D, "D1D mismatch: " << d1d << " != " << gD1D);
+            MFEM_VERIFY(q1d == gQ1D, "Q1D mismatch: " << q1d << " != " << gQ1D);
+         }
+      }
+      else if (version == 4) // PA ∂fem new kernels, not specialized
+      {
+         dPAOperatorSetup(true, false);
+      }
+      else if (version == 5) // PA ∂fem new kernels, specialized
+      {
+         dPAOperatorSetup(true, true);
+      }
+      else if (version == 6) // PA ∂fem std
+      {
+         dPAOperatorSetup(false, false);
+      }
+      else if (version == 7) // MF ∂fem std
+      {
+         dMFOperatorSetup(false, false);
+      }
+      else if (version == 8) // MF ∂fem new kernels
+      {
+         MFEM_ABORT("MF ∂fem new kernels not implemented");
+         // dMFOperatorSetup(true, true);
+      }
+      else { MFEM_ABORT("Invalid version"); }
+
+      cg.SetOperator(*A);
+      cg.iterative_mode = false;
+      cg.SetAbsTol(0.0);
+      if (dofs < 128 * 1024) // check
+      {
+         cg.SetPrintLevel(3/*-1*/);
+         cg.SetMaxIter(2000);
+         cg.SetRelTol(1e-8);
+         cg.Mult(B, X);
+         MFEM_VERIFY(cg.GetConverged(), "❌ CG solver did not converge.");
+         // mfem::out << (cg.GetConverged() ? "✅" : "❌") << std::endl;
+         // mfem::out << "✅" << std::endl;
+      }
+      cg.SetPrintLevel(print_lvl);
+      cg.SetMaxIter(max_it);
+      cg.SetRelTol(rtol);
+      Benchmark();
+      mdofs = 0.0;
+   }
+
+   void Benchmark() override
+   {
+      NVTX_MARK_FUNCTION;
+      cg.Mult(B, X);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs() * cg.GetNumIterations();
+   }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+#define BakeOff_Problem(i, Problem)                                  \
+   static void BP##i(bm::State &state)                               \
+   {                                                                 \
+      const auto version = static_cast<int>(state.range(0));         \
+      const auto order = static_cast<int>(state.range(1));           \
+      const auto side = static_cast<int>(state.range(2));            \
+      Problem ker(version, order, side);                             \
+      while (state.KeepRunning()) { ker.Benchmark(); }               \
+      bm::Counter::Flags flags = bm::Counter::kIsRate;               \
+      state.counters["MDof/s"] = bm::Counter(ker.SumMdofs(), flags); \
+      state.counters["Dofs"] = bm::Counter(ker.dofs);                \
+      state.counters["p"] = bm::Counter(order);                      \
+      state.counters["version"] = bm::Counter(version);              \
+   }                                                                 \
+   BENCHMARK(BP##i)                                                  \
+      ->Apply(CustomArguments)                                       \
+      ->Unit(bm::kMillisecond)
+
+BakeOff_Problem(3, Diffusion);
+
+/// main //////////////////////////////////////////////////////////////////////
+int main(int argc, char *argv[])
+{
+   dbg();
+   DumpVersionInfo();
+   AddBasicKernelSpecializations();
+   static mfem::MPI_Session mpi(argc, argv);
+
+   bm::ConsoleReporter CR;
+   bm::Initialize(&argc, argv);
+
+   // Device setup, cpu by default
+   std::string device_context = "cpu",
+               kernels_context = "std",
+               kernels_specialization = "yes";
+   const auto global_context = bmi::GetGlobalContext();
+   if (global_context != nullptr)
+   {
+      const auto device = global_context->find("device");
+      if (device != global_context->end())
+      {
+         mfem::out << device->first << " : "
+                   << device->second << std::endl;
+         device_context = device->second;
+      }
+   }
+   dbg("device_config: {}", device_context);
+   Device device(device_context.c_str());
+   device_ptr = &device;
+   device.Print();
+
+   if (bm::ReportUnrecognizedArguments(argc, argv)) { return EXIT_FAILURE; }
+
+   bm::RunSpecifiedBenchmarks(&CR);
+
+   return EXIT_SUCCESS;
+}
+
+#endif // MFEM_USE_BENCHMARK

tests/benchmarks/bench_dfem_mma.hpp

@@ -0,0 +1,841 @@
+// 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 "fem/bilininteg.hpp"
+#include <fem/quadinterpolator.hpp>
+#include "general/forall.hpp"
+#include "linalg/dtensor.hpp"
+#include "linalg/kernels.hpp"
+
+using namespace mfem;
+
+/// MMA ///////////////////////////////////////////////////////////////////////
+namespace mma
+{
+
+MFEM_HOST_DEVICE inline int getThreadIdx()
+{
+#ifdef __CUDA_ARCH__
+   return threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
+#else
+   return 0;
+#endif
+}
+
+MFEM_HOST_DEVICE inline int getWarpId(int thread)
+{
+   return thread / 32;
+}
+
+MFEM_HOST_DEVICE inline int getLaneId(int thread)
+{
+   return thread % 32;
+}
+
+MFEM_HOST_DEVICE inline int getGroupId(int laneId)
+{
+   return laneId / 4;
+}
+
+MFEM_HOST_DEVICE inline int getThreadIdInGroup(int laneId)
+{
+   return laneId % 4;
+}
+
+/// Load B1d & G1d matrices into shared memory
+template<int MD1, int MQ1>
+MFEM_HOST_DEVICE inline void LoadBG(const int D1D, const int Q1D,
+                                    const ConstDeviceMatrix &b,
+                                    const ConstDeviceMatrix &g,
+                                    real_t (&sBG)[2][MQ1*MD1])
+{
+   DeviceMatrix B(sBG[0], D1D, Q1D);
+   DeviceMatrix G(sBG[1], D1D, Q1D);
+   int tid = getThreadIdx();
+   if (tid < D1D * Q1D)
+   {
+      int q = tid / D1D;
+      int d = tid % D1D;
+      B(d,q) = b(q,d);
+      G(d,q) = g(q,d);
+   }
+}
+
+/// Load Bt1d & Gt1d matrices into shared memory
+template<int MD1, int MQ1>
+MFEM_HOST_DEVICE inline void LoadBtGt(const int D1D, const int Q1D,
+                                      //   const ConstDeviceMatrix &bt,
+                                      //   const ConstDeviceMatrix &gt,
+                                      const ConstDeviceMatrix &b,
+                                      const ConstDeviceMatrix &g,
+                                      real_t (&sBG)[2][MQ1*MD1])
+{
+   DeviceMatrix Bt(sBG[0], Q1D, D1D);
+   DeviceMatrix Gt(sBG[1], Q1D, D1D);
+
+   int thread = getThreadIdx();
+   if (thread < D1D * Q1D)
+   {
+      int q = thread % Q1D;
+      int d = thread / Q1D;
+      //   Bt(q,d) = bt(d,q);
+      //   Gt(q,d) = gt(d,q);
+      Bt(q,d) = b(q,d);
+      Gt(q,d) = g(q,d);
+   }
+}
+
+/// Load 3D input vector into shared memory
+template<int MQ1>
+MFEM_HOST_DEVICE inline void LoadX(const int e, const int D1D,
+                                   const DeviceTensor<4, const real_t> &x,
+                                   real_t (&sm)[3][MQ1*MQ1*MQ1])
+{
+   const int DDD = D1D * D1D * D1D;
+   DeviceCube X(sm[0], D1D,D1D,D1D);
+   int tid = getThreadIdx();
+   if (tid < DDD)
+   {
+      int dx = tid % D1D;
+      int div = tid / D1D;
+      int dy = div % D1D;
+      int dz = div / D1D;
+      X(dx,dy,dz) = x(dx,dy,dz,e);
+   }
+}
+
+// using the m8n8k4 DMMA instriction
+constexpr int mmaM = 8;
+[[maybe_unused]] constexpr int mmaN = 8;
+constexpr int mmaK = 4;
+
+MFEM_HOST_DEVICE inline void dmmaSync([[maybe_unused]] double aReg[1],
+                                      [[maybe_unused]] double bReg[1],
+                                      [[maybe_unused]] double cReg[2])
+{
+#ifdef __CUDA_ARCH__
+   asm volatile("mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64 {%0,%1}, {%2}, {%3}, {%0,%1};"
+                : "+d"(cReg[0]), "+d"(cReg[1]) : "d"(aReg[0]), "d"(bReg[0]));
+#endif
+}
+
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void dmma_GradX(const int m, const int n, const int k,
+                                        const real_t (&BG)[2][MQ1*MD1],
+                                        const real_t (*A)[MDQ*MDQ*MDQ],
+                                        real_t (*C)[MDQ*MDQ*MDQ])
+{
+   ConstDeviceMatrix B(BG[0], k, n);
+   ConstDeviceMatrix G(BG[1], k, n);
+
+   int thread = getThreadIdx();
+   int warpId = getWarpId(thread);
+   int laneId = getLaneId(thread);
+   int groupId = getGroupId(laneId);
+   int threadIdInGroup = getThreadIdInGroup(laneId);
+
+   // using the m8n8k4 DMMA instriction
+
+   int mPass = (m + mmaM - 1) / mmaM;
+   if (warpId < mPass)   // Spread the warps.
+   {
+
+      int aRowInWarp = groupId;
+      int aColumnInWarp = threadIdInGroup;
+      int bRowInWarp = threadIdInGroup;
+      int bColumnInWarp = groupId;
+
+      constexpr int magicNumber =
+         0b100011111010110001101000; // jump table [0,5,1,6,2,7,3,4]
+      int mM = warpId;
+      double cReg[4] = {};
+      for (int mK = 0; mK < (k + mmaK - 1) / mmaK; mK++)
+      {
+         double bReg[1];
+         double gReg[1];
+         int bRow = bRowInWarp + mK * mmaK;
+         int bColumn = (magicNumber >> (3 * bColumnInWarp)) & 0b111;
+         if (bColumn < n && bRow < k)
+         {
+            bReg[0] = B(bRow, bColumn);
+            gReg[0] = G(bRow, bColumn);
+         }
+         else
+         {
+            bReg[0] = 0;
+            gReg[0] = 0;
+         }
+         double aReg[1];
+         int aRow = aRowInWarp * mPass +  mM;
+         int aColumn = aColumnInWarp + mK * mmaK;
+         if (aRow < m && aColumn < k)
+         {
+            ConstDeviceMatrix aA(A[0], k, m);
+            aReg[0] = aA(aColumn, aRow);
+         }
+         else
+         {
+            aReg[0] = 0;
+         }
+         dmmaSync(aReg, gReg, &cReg[0]);
+         dmmaSync(aReg, bReg, &cReg[2]);
+      }
+      for (int d = 0; d < 2; d++)
+      {
+#pragma unroll
+         for (int i = 0; i < 2; i++)
+         {
+            int cRow = groupId * mPass + mM;
+            int cColumn = (magicNumber >> (3 * (threadIdInGroup * 2 + i))) & 0b111;
+            if (cRow < m && cColumn < n)
+            {
+               DeviceMatrix cC(C[d], m, n);
+               cC(cRow, cColumn) = cReg[d * 2 + i];
+            }
+         }
+      }
+   }
+}
+
+/// 3D Gradient, 1/3
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void GradX(const int D1D, const int Q1D,
+                                   const real_t (&sBG)[2][MQ1*MD1],
+                                   const real_t (*sDDD)[MDQ*MDQ*MDQ],
+                                   real_t (*sDDQ)[MDQ*MDQ*MDQ])
+{
+   dmma_GradX<MD1, MQ1>(D1D * D1D, Q1D, D1D, sBG, sDDD, sDDQ);
+}
+
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void dmma_GradY(const int m, const int n,
+                                        const int k,
+                                        const real_t (&BG)[2][MQ1*MD1],
+                                        const real_t (*A)[MDQ*MDQ*MDQ],
+                                        real_t (*C)[MDQ*MDQ*MDQ])
+{
+   ConstDeviceMatrix B(BG[0], k, n);
+   ConstDeviceMatrix G(BG[1], k, n);
+
+   int thread = getThreadIdx();
+   int warpId = getWarpId(thread);
+   int laneId = getLaneId(thread);
+   int groupId = getGroupId(laneId);
+   int threadIdInGroup = getThreadIdInGroup(laneId);
+
+   // using the m8n8k4 DMMA instriction
+
+   int mPass = (m + mmaM - 1) / mmaM;
+   if (warpId < mPass)   // Spread the warps.
+   {
+
+      int aRowInWarp = groupId;
+      int aColumnInWarp = threadIdInGroup;
+      int bRowInWarp = threadIdInGroup;
+      int bColumnInWarp = groupId;
+
+      constexpr int magicNumber =
+         0b100011111010110001101000; // jump table [0,5,1,6,2,7,3,4]
+      int mM = warpId;
+      double cReg[6] = {};
+      for (int mK = 0; mK < (k + mmaK - 1) / mmaK; mK++)
+      {
+         double bReg[1];
+         double gReg[1];
+         int bRow = bRowInWarp + mK * mmaK;
+         int bColumn = (magicNumber >> (3 * bColumnInWarp)) & 0b111;
+         if (bColumn < n && bRow < k)
+         {
+            bReg[0] = B(bRow, bColumn);
+            gReg[0] = G(bRow, bColumn);
+         }
+         else
+         {
+            bReg[0] = 0;
+            gReg[0] = 0;
+         }
+         double agReg[1];
+         double abReg[1];
+         int aRow = aRowInWarp * mPass +  mM;
+         int aColumn = aColumnInWarp + mK * mmaK;
+         if (aRow < m && aColumn < k)
+         {
+            ConstDeviceMatrix gA(A[0], k, m);
+            ConstDeviceMatrix bA(A[1], k, m);
+            agReg[0] = gA(aColumn, aRow);
+            abReg[0] = bA(aColumn, aRow);
+         }
+         else
+         {
+            agReg[0] = 0;
+            abReg[0] = 0;
+         }
+         dmmaSync(agReg, bReg, &cReg[0]);
+         dmmaSync(abReg, gReg, &cReg[2]);
+         dmmaSync(abReg, bReg, &cReg[4]);
+      }
+      for (int d = 0; d < 3; d++)
+      {
+#pragma unroll
+         for (int i = 0; i < 2; i++)
+         {
+            int cRow = groupId * mPass + mM;
+            int cColumn = (magicNumber >> (3 * (threadIdInGroup * 2 + i))) & 0b111;
+            if (cRow < m && cColumn < n)
+            {
+               DeviceMatrix cC(C[d], m, n);
+               cC(cRow, cColumn) = cReg[d * 2 + i];
+            }
+         }
+      }
+   }
+}
+
+/// 3D Gradient, 2/3
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void GradY(const int D1D, const int Q1D,
+                                   const real_t (&sBG)[2][MQ1*MD1],
+                                   const real_t (*sDDQ)[MDQ*MDQ*MDQ],
+                                   real_t (*sDQQ)[MDQ*MDQ*MDQ])
+{
+   dmma_GradY<MD1, MQ1>(D1D * Q1D, Q1D, D1D, sBG, sDDQ, sDQQ);
+}
+
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void dmma_GradZ(const int m, const int n,
+                                        const int k,
+                                        const real_t (&BG)[2][MQ1*MD1],
+                                        const real_t (*A)[MDQ*MDQ*MDQ],
+                                        real_t (*C)[MDQ*MDQ*MDQ],
+                                        int gIdx)
+{
+   ConstDeviceMatrix B(BG[0], k, n);
+   ConstDeviceMatrix G(BG[1], k, n);
+
+   int thread = getThreadIdx();
+   int warpId = getWarpId(thread);
+   int laneId = getLaneId(thread);
+   int groupId = getGroupId(laneId);
+   int threadIdInGroup = getThreadIdInGroup(laneId);
+
+   // using the m8n8k4 DMMA instriction
+
+   int mPass = (m + mmaM - 1) / mmaM;
+   if (warpId < mPass)   // Spread the warps.
+   {
+
+      int aRowInWarp = groupId;
+      int aColumnInWarp = threadIdInGroup;
+      int bRowInWarp = threadIdInGroup;
+      int bColumnInWarp = groupId;
+
+      constexpr int magicNumber =
+         0b100011111010110001101000; // jump table [0,5,1,6,2,7,3,4]
+      int mM = warpId;
+      double cReg[6] = {};
+      for (int mK = 0; mK < (k + mmaK - 1) / mmaK; mK++)
+      {
+         double bReg[1];
+         double gReg[1];
+         int bRow = bRowInWarp + mK * mmaK;
+         int bColumn = (magicNumber >> (3 * bColumnInWarp)) & 0b111;
+         if (bColumn < n && bRow < k)
+         {
+            bReg[0] = B(bRow, bColumn);
+            gReg[0] = G(bRow, bColumn);
+         }
+         else
+         {
+            bReg[0] = 0;
+            gReg[0] = 0;
+         }
+         for (int d = 0; d < 3; d++)
+         {
+            double aReg[1];
+            int aRow = aRowInWarp * mPass +  mM;
+            int aColumn = aColumnInWarp + mK * mmaK;
+            if (aRow < m && aColumn < k)
+            {
+               ConstDeviceMatrix aA(A[d], k, m);
+               aReg[0] = aA(aColumn, aRow);
+            }
+            else
+            {
+               aReg[0] = 0;
+            }
+            dmmaSync(aReg, d == gIdx ? gReg : bReg, &cReg[d * 2]);
+         }
+      }
+      for (int d = 0; d < 3; d++)
+      {
+#pragma unroll
+         for (int i = 0; i < 2; i++)
+         {
+            int cRow = groupId * mPass + mM;
+            int cColumn = (magicNumber >> (3 * (threadIdInGroup * 2 + i))) & 0b111;
+            if (cRow < m && cColumn < n)
+            {
+               DeviceMatrix cC(C[d], m, n);
+               cC(cRow, cColumn) = cReg[d * 2 + i];
+            }
+         }
+      }
+   }
+}
+
+/// 3D Gradient, 3/3
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void GradZ(const int D1D, const int Q1D,
+                                   const real_t (&sBG)[2][MQ1*MD1],
+                                   const real_t (*sDQQ)[MDQ*MDQ*MDQ],
+                                   real_t (*sQQQ)[MDQ*MDQ*MDQ])
+{
+   dmma_GradZ<MD1, MQ1>(Q1D * Q1D, Q1D, D1D, sBG, sDQQ, sQQQ, 2);
+}
+
+/// 3D Transposed Gradient, 1/3
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void GradZt(const int D1D, const int Q1D,
+                                    const real_t (&sBG)[2][MQ1*MD1],
+                                    const real_t (*sQQQ)[MDQ*MDQ*MDQ],
+                                    real_t (*sDQQ)[MDQ*MDQ*MDQ])
+{
+   ConstDeviceMatrix Bt(sBG[0], Q1D, D1D);
+   ConstDeviceMatrix Gt(sBG[1], Q1D, D1D);
+   int thread = getThreadIdx();
+   int warpId = getWarpId(thread);
+   int laneId = getLaneId(thread);
+   int groupId = getGroupId(laneId);
+   int threadIdInGroup = getThreadIdInGroup(laneId);
+
+   // using the m8n8k4 DMMA instriction
+   // qy (Q1D), qz (Q1D) === M, dx (D1D) === N, qx (Q1D) === K
+
+   int mPass = (Q1D * Q1D + mmaM - 1) / mmaM;
+   if (warpId < mPass)   // Spread the warps to calculate the 3 directions.
+   {
+
+      int aRowInWarp = groupId;
+      int aColumnInWarp = threadIdInGroup;
+      int bRowInWarp = threadIdInGroup;
+      int bColumnInWarp = groupId;
+
+      constexpr int magicNumber =
+         0b100011111010110001101000; // jump table [0,5,1,6,2,7,3,4]
+      int mM = warpId;
+      double cReg[6] = {};
+      for (int mK = 0; mK < (Q1D + mmaK - 1) / mmaK; mK++)
+      {
+         double BtReg[1];
+         double GtReg[1];
+         int bRow = bRowInWarp + mK * mmaK;
+         int bColumn = (magicNumber >> (3 * bColumnInWarp)) & 0b111;
+         if (bColumn < D1D && bRow < Q1D)
+         {
+            BtReg[0] = Bt(bRow, bColumn);
+            GtReg[0] = Gt(bRow, bColumn);
+         }
+         else
+         {
+            BtReg[0] = 0;
+            GtReg[0] = 0;
+         }
+         for (int d = 0; d < 3; d++)
+         {
+            double aReg[1];
+            int aRow = aRowInWarp * mPass +  mM;
+            int aColumn = aColumnInWarp + mK * mmaK;
+            if (aRow < Q1D * Q1D && aColumn < Q1D)
+            {
+               ConstDeviceMatrix XxBBG(sQQQ[d], Q1D, Q1D * Q1D);
+               aReg[0] = XxBBG(aColumn, aRow);
+            }
+            else
+            {
+               aReg[0] = 0;
+            }
+
+            dmmaSync(aReg, d == 0 ? GtReg : BtReg, &cReg[d * 2]);
+         }
+      }
+      for (int d = 0; d < 3; d++)
+      {
+#pragma unroll
+         for (int i = 0; i < 2; i++)
+         {
+            int cRow = groupId * mPass + mM;
+            int cColumn = (magicNumber >> (3 * (threadIdInGroup * 2 + i))) & 0b111;
+            if (cRow < Q1D * Q1D && cColumn < D1D)
+            {
+               DeviceMatrix Xx(sDQQ[d], Q1D * Q1D, D1D); // qy, qz, dx
+               Xx(cRow, cColumn) = cReg[d * 2 + i];
+            }
+         }
+      }
+   }
+}
+
+/// 3D Transposed Gradient, 2/3
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void GradYt(const int D1D, const int Q1D,
+                                    const real_t (&sBG)[2][MQ1*MD1],
+                                    const real_t (*sDQQ)[MDQ*MDQ*MDQ],
+                                    real_t (*sDDQ)[MDQ*MDQ*MDQ])
+{
+   ConstDeviceMatrix Bt(sBG[0], Q1D, D1D);
+   ConstDeviceMatrix Gt(sBG[1], Q1D, D1D);
+   int thread = getThreadIdx();
+   int warpId = getWarpId(thread);
+   int laneId = getLaneId(thread);
+   int groupId = getGroupId(laneId);
+   int threadIdInGroup = getThreadIdInGroup(laneId);
+
+   // using the m8n8k4 DMMA instriction
+   // dx (D1D), qz (Q1D) === M, dy (D1D) === N, qy (Q1D) === K
+
+   int mPass = (D1D * Q1D + mmaM - 1) / mmaM;
+   if (warpId < mPass)   // Spread the warps.
+   {
+
+      int aRowInWarp = groupId;
+      int aColumnInWarp = threadIdInGroup;
+      int bRowInWarp = threadIdInGroup;
+      int bColumnInWarp = groupId;
+
+      constexpr int magicNumber =
+         0b100011111010110001101000; // jump table [0,5,1,6,2,7,3,4]
+      int mM = warpId;
+      double cReg[6] = {}; // initialized to zero
+      for (int mK = 0; mK < (Q1D + mmaK - 1) / mmaK; mK++)
+      {
+         double BtReg[1];
+         double GtReg[1];
+         int bRow = bRowInWarp + mK * mmaK;
+         int bColumn = (magicNumber >> (3 * bColumnInWarp)) & 0b111;
+         if (bColumn < D1D && bRow < Q1D)
+         {
+            BtReg[0] = Bt(bRow, bColumn);
+            GtReg[0] = Gt(bRow, bColumn);
+         }
+         else
+         {
+            BtReg[0] = 0;
+            GtReg[0] = 0;
+         }
+         for (int d = 0; d < 3; d++)
+         {
+            double aReg[1];
+
+            int aRow = aRowInWarp * mPass + mM;
+            int aColumn = aColumnInWarp + mK * mmaK;
+            if (aRow < D1D * Q1D && aColumn < Q1D)
+            {
+               ConstDeviceMatrix XxBB(sDQQ[d], Q1D, D1D * Q1D); // qy, qz, dx
+               aReg[0] = XxBB(aColumn, aRow);
+            }
+            else
+            {
+               aReg[0] = 0;
+            }
+
+            dmmaSync(aReg, d == 1 ? GtReg : BtReg, &cReg[d * 2]);
+         }
+      }
+      for (int d = 0; d < 3; d++)
+      {
+#pragma unroll
+         for (int i = 0; i < 2; i++)
+         {
+            int cRow = groupId * mPass + mM;
+            int cColumn = (magicNumber >> (3 * (threadIdInGroup * 2 + i))) & 0b111;
+            if (cRow < D1D * Q1D && cColumn < D1D)
+            {
+               DeviceMatrix Xx(sDDQ[d], D1D * Q1D, D1D); // qz, dx, dy
+               Xx(cRow, cColumn) = cReg[d * 2 + i];
+            }
+         }
+      }
+   }
+}
+
+/// 3D Transposed Gradient, 3/3
+template<int MD1, int MQ1, int MDQ = (MQ1 > MD1 ? MQ1 : MD1)>
+MFEM_HOST_DEVICE inline void GradXt(const int D1D, const int Q1D,
+                                    const real_t (&sBG)[2][MQ1*MD1],
+                                    const real_t (&sDDQ)[3][MDQ*MDQ*MDQ],
+                                    const DeviceTensor<4> &Y, // output
+                                    const int e)
+{
+   ConstDeviceMatrix Bt(sBG[0], Q1D, D1D);
+   ConstDeviceMatrix Gt(sBG[1], Q1D, D1D);
+   int thread = getThreadIdx();
+   int warpId = getWarpId(thread);
+   int laneId = getLaneId(thread);
+   int groupId = getGroupId(laneId);
+   int threadIdInGroup = getThreadIdInGroup(laneId);
+
+   // using the m8n8k4 DMMA instriction
+   // dx (D1D), dy (D1D) === M, dz (D1D) === N, qz (Q1D) === K
+
+   int mPass = (D1D * D1D + mmaM - 1) / mmaM;
+   if (warpId < mPass)   // Spread the warps to calculate the 3 directions.
+   {
+
+      int aRowInWarp = groupId;
+      int aColumnInWarp = threadIdInGroup;
+      int bRowInWarp = threadIdInGroup;
+      int bColumnInWarp = groupId;
+
+      constexpr int magicNumber =
+         0b100011111010110001101000; // jump table [0,5,1,6,2,7,3,4]
+      int mM = warpId;
+      {
+         double BtReg[1];
+         double GtReg[1];
+         double cReg[2] = {}; // initialized to zero
+
+         for (int mK = 0; mK < (Q1D + mmaK - 1) / mmaK; mK++)
+         {
+            int bRow = bRowInWarp + mK * mmaK;
+            int bColumn = (magicNumber >> (3 * bColumnInWarp)) & 0b111;
+            if (bColumn < D1D && bRow < Q1D)
+            {
+               BtReg[0] = Bt(bRow, bColumn);
+               GtReg[0] = Gt(bRow, bColumn);
+            }
+            else
+            {
+               BtReg[0] = 0;
+               GtReg[0] = 0;
+            }
+            for (int d = 0; d < 3; d++)
+            {
+               double aReg[1];
+               int aRow = aRowInWarp * mPass + mM;
+               int aColumn = aColumnInWarp + mK * mmaK;
+               if (aRow < D1D * D1D && aColumn < Q1D)
+               {
+                  ConstDeviceMatrix Xx(sDDQ[d], Q1D, D1D * D1D); // qz, dx, dy
+                  aReg[0] = Xx(aColumn, aRow);
+               }
+               else
+               {
+                  aReg[0] = 0;
+               }
+
+               dmmaSync(aReg, d == 2 ? GtReg : BtReg, cReg);
+            }
+         }
+#pragma unroll
+         for (int i = 0; i < 2; i++)
+         {
+            int cRow = groupId * mPass + mM;
+            int cColumn = (magicNumber >> (3 * (threadIdInGroup * 2 + i))) & 0b111;
+            if (cRow < D1D * D1D && cColumn < D1D)
+            {
+               int dx = cRow % D1D;
+               int dy = cRow / D1D;
+               int dz = cColumn;
+               Y(dx,dy,dz,e) += cReg[i];
+            }
+         }
+      }
+   }
+}
+
+} // namespace mma
+
+/// PADiffMmaIntegrator ///////////////////////////////////////////////////////
+struct PADiffMmaIntegrator : public BilinearFormIntegrator
+{
+   const FiniteElementSpace *fes;
+   const real_t *B, *G, *DX;
+   int ne, d1d, q1d;
+   Vector J0, dx;
+
+public: // for nvcc
+   //////////////////////////////////////////////////////////////////
+   template <int T_D1D = 0, int T_Q1D = 0>
+   static void PADiffMmaMult(const int ne,
+                             const real_t *b, const real_t *g,
+                             const real_t *dx, const real_t *xe,
+                             real_t *ye,
+                             const int, const int)
+   {
+      constexpr int Q1D = T_Q1D, D1D = T_D1D;
+
+      const auto B = Reshape(b, Q1D, D1D);
+      const auto G = Reshape(g, Q1D, D1D);
+
+      const auto XE = Reshape(xe, D1D, D1D, D1D, ne);
+      const auto DX = Reshape(dx, 3, 3, Q1D, Q1D, Q1D, ne);
+      auto YE = Reshape(ye, D1D, D1D, D1D, ne);
+
+      mfem::forall_3D(ne, ((Q1D * Q1D * Q1D + 31) / 32) * 32, 1, 1,
+                      [=] MFEM_HOST_DEVICE(int e)
+      {
+         constexpr int MQ1 = T_Q1D, MD1 = T_D1D;
+
+         MFEM_SHARED real_t sm0[3][MQ1*MQ1*MQ1];
+         MFEM_SHARED real_t sm1[3][MQ1*MQ1*MQ1];
+         MFEM_SHARED real_t BG[2][MD1*MQ1];
+
+         mma::LoadBG<MD1, MQ1>(D1D, Q1D, B, G, BG);
+         mma::LoadX<MQ1>(e, D1D, XE, sm0);
+         MFEM_SYNC_THREAD;
+
+         mma::GradX<MD1, MQ1>(D1D, Q1D, BG, sm0, sm1);
+         MFEM_SYNC_THREAD;
+         mma::GradY<MD1, MQ1>(D1D, Q1D, BG, sm1, sm0);
+         MFEM_SYNC_THREAD;
+         mma::GradZ<MD1, MQ1>(D1D, Q1D, BG, sm0, sm1);
+         MFEM_SYNC_THREAD;
+
+         int thread = mma::getThreadIdx();
+         if (thread < Q1D * Q1D * Q1D)
+         {
+            int qx = thread % Q1D;
+            int div = thread / Q1D;
+            int qy = div % Q1D;
+            int qz = div / Q1D;
+
+            {
+               // pull
+               real_t v[3], u[3] = { sm1[0][qz + qy*Q1D + qx*Q1D*Q1D],
+                                     sm1[1][qz + qy*Q1D + qx*Q1D*Q1D],
+                                     sm1[2][qz + qy*Q1D + qx*Q1D*Q1D]
+                                   };
+               //  Q-function
+               const real_t *dx = &DX(0, 0, qx, qy, qz, e);
+               kernels::Mult(3, 3, dx, u, v);
+               // push
+               sm0[0][qz + qy*Q1D + qx*Q1D*Q1D] = v[0];
+               sm0[1][qz + qy*Q1D + qx*Q1D*Q1D] = v[1];
+               sm0[2][qz + qy*Q1D + qx*Q1D*Q1D] = v[2];
+            }
+         }
+
+         mma::LoadBtGt<MD1,MQ1>(D1D, Q1D, B, G, BG);
+         MFEM_SYNC_THREAD;
+         mma::GradZt<MD1, MQ1>(D1D, Q1D, BG, sm0, sm1);
+         MFEM_SYNC_THREAD;
+         mma::GradYt<MD1, MQ1>(D1D, Q1D, BG, sm1, sm0);
+         MFEM_SYNC_THREAD;
+         mma::GradXt<MD1,MQ1>(D1D, Q1D, BG, sm0, YE, e);
+      });
+   }
+
+   using PADiffMmaKernelType = decltype(&PADiffMmaMult<>);
+   MFEM_REGISTER_KERNELS(PADiffMmaKernels, PADiffMmaKernelType, (int, int));
+
+public:
+   PADiffMmaIntegrator()
+   {
+      // PADiffMmaKernels::Specialization<2,3>::Add();  // 1 ❌
+      PADiffMmaKernels::Specialization<3,4>::Add();  // 2
+      PADiffMmaKernels::Specialization<4,5>::Add();  // 3
+      PADiffMmaKernels::Specialization<5,6>::Add();  // 4
+      PADiffMmaKernels::Specialization<6,7>::Add();  // 5
+      PADiffMmaKernels::Specialization<7,8>::Add();  // 6
+   }
+
+   void AssemblePA(const FiniteElementSpace &fespace) override
+   {
+      NVTX();
+      fes = &fespace;
+      auto *mesh = fes->GetMesh();
+      const int DIM = mesh->Dimension();
+      ne = mesh->GetNE();
+      const auto p = fes->GetFE(0)->GetOrder();
+      const auto q = 2 * p + mesh->GetElementTransformation(0)->OrderW();
+      const auto type = mesh->GetElementBaseGeometry(0);
+      const IntegrationRule &ir = IntRules.Get(type, q);
+      const int NQPT = ir.GetNPoints();
+      d1d = p + 1;
+      q1d = IntRules.Get(Geometry::SEGMENT, ir.GetOrder()).GetNPoints();
+      MFEM_VERIFY(NQPT == q1d * q1d * q1d, "");
+      const DofToQuad *maps =
+         &fes->GetFE(0)->GetDofToQuad(ir, DofToQuad::TENSOR);
+      const GridFunction *nodes = (mesh->EnsureNodes(), mesh->GetNodes());
+      const FiniteElementSpace *nfes = nodes->FESpace();
+      const int nVDIM = nfes->GetVDim();
+      dx.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType());
+      J0.SetSize(nVDIM * DIM * NQPT * ne, Device::GetDeviceMemoryType());
+      dx.UseDevice(true), J0.UseDevice(true);
+      B = maps->B.Read(), G = maps->G.Read(), DX = dx.Read();
+
+      const Operator *NR =
+         nfes->GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+      const QuadratureInterpolator *nqi = nfes->GetQuadratureInterpolator(ir);
+      nqi->SetOutputLayout(QVectorLayout::byVDIM);
+      const int nd = nfes->GetFE(0)->GetDof();
+      Vector xe(nVDIM * nd * ne, Device::GetDeviceMemoryType());
+      NR->Mult(*nodes, (xe.UseDevice(true), xe));
+      nqi->Derivatives(xe, J0);
+
+      const int Q1D = q1d;
+      const auto w_r = ir.GetWeights().Read();
+      const auto W = Reshape(w_r, q1d, q1d, q1d);
+      const auto J = Reshape(J0.Read(), 3, 3, q1d, q1d, q1d, ne);
+      auto DX_w = Reshape(dx.Write(), 3, 3, q1d, q1d, q1d, ne);
+
+      mfem::forall_3D(ne, Q1D, Q1D, Q1D,[=] MFEM_HOST_DEVICE(int e)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qz, z, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  const real_t w = W(qx, qy, qz);
+                  const real_t *Jtr = &J(0, 0, qx, qy, qz, e);
+                  const real_t detJ = kernels::Det<3>(Jtr);
+                  const real_t wd = w * detJ;
+                  const real_t D[9] = { wd, 0.0, 0.0,
+                                        0.0, wd, 0.0,
+                                        0.0, 0.0, wd
+                                      };
+                  real_t Jrt[9], A[9];
+                  kernels::CalcInverse<3>(Jtr, Jrt);
+                  kernels::MultABt(3, 3, 3, D, Jrt, A);
+                  kernels::Mult(3, 3, 3, A, Jrt, &DX_w(0, 0, qz, qy, qx, e));
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      });
+   }
+
+   void AddMultPA(const Vector &x, Vector &y) const override
+   {
+      db1("\x1b[32md1d:{} q1d:{}", d1d, q1d);
+      PADiffMmaKernels::Run(d1d, q1d,
+                            ne, B, G, DX, x.Read(), y.ReadWrite(),
+                            d1d, q1d);
+   }
+};
+template <int D1D, int Q1D>
+PADiffMmaIntegrator::PADiffMmaKernelType
+PADiffMmaIntegrator::PADiffMmaKernels::Kernel()
+{
+   db1("D1D:{} Q1D:{}", D1D, Q1D);
+   return PADiffMmaMult<D1D, Q1D>;
+}
+
+PADiffMmaIntegrator::PADiffMmaKernelType
+PADiffMmaIntegrator::PADiffMmaKernels::Fallback(int d1d, int q1d)
+{
+   dbg("\x1b[33mFallback d1d:{} q1d:{}", d1d, q1d);
+   MFEM_ABORT("No kernel for q1d=" << q1d);
+   return nullptr;
+   //    return PADiffMmaMult;
+}

tests/benchmarks/makefile

@@ -20,8 +20,8 @@ CONFIG_MK = $(or $(wildcard $(MFEM_BUILD_DIR)/config/config.mk),\
 MFEM_LIB_FILE = mfem_is_not_built
 -include $(CONFIG_MK)
 
-SEQ_TESTS = bench_assembly_levels bench_ceed bench_dg_amr bench_elasticity \
-            bench_tmop bench_vector bench_virtuals
+SEQ_TESTS = bench_assembly_levels bench_ceed bench_dfem bench_dg_amr \
+			bench_elasticity bench_tmop bench_vector bench_virtuals
 PAR_TESTS = 
 ifeq ($(MFEM_USE_MPI),NO)
    TESTS = $(SEQ_TESTS)

tests/unit/CMakeLists.txt

@@ -36,6 +36,7 @@ include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR})
 #    for d in dfem general linalg mesh fem enzyme; do ls -1 $d/*.cpp; done
 set(UNIT_TESTS_SRCS
   # dfem/test_diffusion.cpp
+  # dfem/test_diffusion_q1d.cpp
   # dfem/test_divergence.cpp
   # dfem/test_lvector_interface.cpp
   # dfem/test_mass.cpp

tests/unit/dfem/test_diffusion.cpp

@@ -302,7 +302,7 @@ void diffusion(const char *filename, int p)
 
 TEST_CASE("dFEM Diffusion", "[Parallel][dFEM][GPU]")
 {
-   const bool all_tests = launch_all_non_regression_tests;
+   // const bool all_tests = launch_all_non_regression_tests;
 
    const auto p = !all_tests ? 2 : GENERATE(1, 2, 3);
 

tests/unit/dfem/test_diffusion_q1d.cpp

@@ -0,0 +1,338 @@
+// 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.
+#define NVTX_COLOR nvtx::kGold
+
+#include "unit_tests.hpp"
+#include "mfem.hpp"
+// #include <type_traits>
+#include "fem/dfem/doperator.hpp"
+#include "fem/dfem/util.hpp"
+
+#include <fem/integ/bilininteg_diffusion_kernels.hpp>
+
+#ifdef MFEM_USE_MPI
+
+using namespace mfem;
+using namespace mfem::future;
+using mfem::future::tensor;
+using mfem::future::dual;
+
+using DOperator = future::DifferentiableOperator;
+
+enum class MQ1Settings : int { kRuntime,
+                               kCompileTime,
+                               kDefault
+                             };
+
+namespace dfem_pa_kernels
+{
+
+///////////////////////////////////////////////////////////////////////////////
+template <typename T, int DIM, int T_MQ1 = 0> struct Diffusion
+{
+   using dvecd_t = tensor<T, DIM>;
+   using matd_t = tensor<real_t, DIM, DIM>;
+
+   struct MFApply
+   {
+      static constexpr int MQ1 = T_MQ1;
+      MFEM_HOST_DEVICE inline auto operator()(const dvecd_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::future::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::future::tuple{ inv(J) * transpose(inv(J)) * det(J) * w * rho };
+      }
+   };
+
+   struct PAApply
+   {
+      MFEM_HOST_DEVICE inline auto operator()(const dvecd_t &dudxi,
+                                              const matd_t &q) const
+      {
+         return mfem::future::tuple{ q * dudxi };
+      };
+   };
+};
+
+///////////////////////////////////////////////////////////////////////////////
+template <typename T, int DIM, std::size_t... MQ1s>
+struct MFDiffusionFactory
+{
+   static auto All()
+   {
+      // could also use a map instead of a tuple
+      return mfem::future::make_tuple(typename Diffusion<T, DIM, MQ1s>::MFApply{}...);
+   }
+};
+
+template <typename T, int DIM>
+using MFDiffusionFactory_1_4 = MFDiffusionFactory<T, DIM, 1, 2, 3, 4>;
+
+template <typename T, int DIM>
+class MFDiffusionQFs
+{
+   using MFApplyTuple = decltype(MFDiffusionFactory_1_4<T, DIM>::All());
+   MFApplyTuple mf_qfs;
+
+public:
+   MFDiffusionQFs(): mf_qfs(MFDiffusionFactory_1_4<T, DIM>::All()) {}
+
+   template <typename F>
+   void run(int i, F&& f)
+   {
+      MFEM_VERIFY(i >= 1, "Index must be >= 1");
+      const auto I = static_cast<size_t>(i - 1);
+      runtime_get_impl(I, std::forward<F>(f),
+                       std::make_index_sequence<mfem::future::tuple_size<MFApplyTuple>::value>());
+   }
+
+private:
+   template <typename F, size_t... I>
+   void runtime_get_impl(size_t index, F&& f, std::index_sequence<I...>)
+   {
+      using fun_ptr = std::function<void(F&&)>;
+      fun_ptr table[] = { [&](F&& f) { f(mfem::future::get<I>(mf_qfs)); } ... };
+      if (index < mfem::future::tuple_size<MFApplyTuple>::value)
+      {
+         table[index](std::forward<F>(f));
+      }
+      else
+      {
+         throw std::out_of_range("Index out of bounds");
+      }
+   }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+template <int DIM>
+void DFemDiffusion(const char *filename, int p, const int r,
+                   const MQ1Settings mq1_setting)
+{
+   dbg("DIM:{}", DIM);
+   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.attributes.Size() > 0)
+   {
+      all_domain_attr.SetSize(pmesh.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");
+
+   ParGridFunction x(&pfes), y(&pfes), z(&pfes);
+   Vector X(pfes.GetTrueVSize()), Y(pfes.GetTrueVSize()), Z(pfes.GetTrueVSize());
+
+   X.Randomize(1);
+   x.SetFromTrueDofs(X);
+
+   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();
+
+   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, "");
+
+   UniformParameterSpace rho_ps(pmesh, *ir, 1);
+
+   static constexpr int U = 0, Coords = 1, Rho = 3;
+   const auto sol = std::vector{ FieldDescriptor{ U, &pfes } };
+
+   SECTION("DFEM Matrix free")
+   {
+      // fields = {solutions, parameters}
+      dbg("fields = {{solutions, parameters}} = {{{{U}}, {{Rho, Coords}}}}");
+      DOperator dop_mf(sol, {{Rho, &rho_ps}, {Coords, mfes}}, pmesh);
+
+      dbg("AddDomainIntegrator: {{∇U, Rho, ∇Coords, Weight}} -> {{∇U}}");
+      if (mq1_setting == MQ1Settings::kRuntime)
+      {
+         dbg("MQ1Settings::kRuntime");
+         MFEM_VERIFY(q1d == (int)floor(std::pow(ir->GetNPoints(), 1.0/DIM) + 0.5),
+                     "q1d and ir->GetNPoints() have to match");
+         auto add_domain_integrator = [&](auto &qf)
+         {
+            dbg("q1d:{} MQ1:{}", q1d, qf.MQ1);
+            MFEM_VERIFY(q1d == qf.MQ1, "q1d and qf.MQ1 have to match");
+            dop_mf.AddDomainIntegrator(qf,
+                                       tuple{ Gradient<U>{}, Identity<Rho>{},
+                                              Gradient<Coords>{}, Weight{} },
+                                       tuple{ Gradient<U>{} }, *ir,
+                                       all_domain_attr);
+         };
+         // select the right qf from the factory
+         MFDiffusionQFs<real_t, DIM> {}.run(q1d, add_domain_integrator);
+      }
+      else if (mq1_setting == MQ1Settings::kCompileTime) // hardcoded, MQ1 = 2,3,4,5
+      {
+         dbg("MQ1Settings::kCompileTime");
+         dbg("q1d:{}", q1d);
+         if (q1d == 2)
+         {
+            typename Diffusion<real_t, DIM, 2>::MFApply mf_apply_qf;
+            MFEM_VERIFY(q1d == 2, "q1d and 2 have to match");
+            dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                       tuple{ Gradient<U>{}, Identity<Rho>{},
+                                              Gradient<Coords>{}, Weight{} },
+                                       tuple{ Gradient<U>{} }, *ir,
+                                       all_domain_attr);
+         }
+         else if (q1d == 3)
+         {
+            typename Diffusion<real_t, DIM, 3>::MFApply mf_apply_qf;
+            MFEM_VERIFY(q1d == 3, "q1d and 3 have to match");
+            dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                       tuple{ Gradient<U>{}, Identity<Rho>{},
+                                              Gradient<Coords>{}, Weight{} },
+                                       tuple{ Gradient<U>{} }, *ir,
+                                       all_domain_attr);
+         }
+         else if (q1d == 4)
+         {
+            typename Diffusion<real_t, DIM, 4>::MFApply mf_apply_qf;
+            MFEM_VERIFY(q1d == 4, "q1d and 4 have to match");
+            dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                       tuple{ Gradient<U>{}, Identity<Rho>{},
+                                              Gradient<Coords>{}, Weight{} },
+                                       tuple{ Gradient<U>{} }, *ir,
+                                       all_domain_attr);
+         }
+         else if (q1d == 5)
+         {
+            typename Diffusion<real_t, DIM, 5>::MFApply mf_apply_qf;
+            MFEM_VERIFY(q1d == 5, "q1d and 5 have to match");
+            dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                       tuple{ Gradient<U>{}, Identity<Rho>{},
+                                              Gradient<Coords>{}, Weight{} },
+                                       tuple{ Gradient<U>{} }, *ir,
+                                       all_domain_attr);
+         }
+         else { MFEM_ABORT("Not supported q1d:" << q1d); }
+      }
+      else // MQ1Settings::kDefault, MQ1 = 0
+      {
+         dbg("MQ1Settings::kDefault");
+         typename Diffusion<real_t, DIM>::MFApply mf_apply_qf;
+         dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                    tuple{ Gradient<U>{}, Identity<Rho>{},
+                                           Gradient<Coords>{}, Weight{} },
+                                    tuple{ Gradient<U>{} }, *ir,
+                                    all_domain_attr);
+      }
+      dop_mf.SetParameters({ &rho_coeff_cv, nodes });
+
+      pfes.GetRestrictionMatrix()->Mult(x, X);
+      dop_mf.Mult(X, Z);
+
+      blf_fa.Mult(x, y);
+      pfes.GetProlongationMatrix()->MultTranspose(y, Y);
+      Y -= Z;
+
+      real_t norm_global = 0.0;
+      real_t norm_local = Y.Normlinf();
+      MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
+                    pmesh.GetComm());
+
+      REQUIRE(norm_global == MFEM_Approx(0.0));
+      MPI_Barrier(MPI_COMM_WORLD);
+   }
+
+}
+
+TEST_CASE("DFEM Diffusion Q1D", "[Parallel][DFEM][MQ1]")
+{
+   // const bool all_tests = launch_all_non_regression_tests;
+
+   // const auto p = !all_tests ? 2 : GENERATE(1, 2, 3);
+   // const auto r = !all_tests ? 1 : GENERATE(0, 1, 2, 3);
+   const int p = 2, r = 1;
+   dbg("p:{} r:{}", p, r);
+
+   const auto mq1_setting = MQ1Settings::kCompileTime;
+   /*const auto mq1_setting = GENERATE(MQ1Settings::kRuntime,
+                                     MQ1Settings::kCompileTime,
+                                     MQ1Settings::kDefault);*/
+
+   DiffusionIntegrator::AddSpecialization<3,3,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))
+   {
+#if 0
+      const auto filename =
+         GENERATE("../../data/fichera.mesh",
+                  "../../data/fichera-q3.mesh",
+                  "../../data/inline-hex.mesh",
+                  "../../data/toroid-hex.mesh",
+                  "../../data/periodic-cube.mesh");
+#else
+      const auto filename = "../../data/fichera.mesh";
+#endif
+      dbg("DFemDiffusion");
+      DFemDiffusion<3>(filename, p, r, mq1_setting);
+   }
+}
+
+} // namespace dfem_pa_kernels
+
+#endif

tests/unit/dfem/test_divergence.cpp

@@ -11,6 +11,7 @@
 
 #include "../unit_tests.hpp"
 #include "mfem.hpp"
+#include <fem/dfem/doperator.hpp>
 
 #ifdef MFEM_USE_MPI
 

tests/unit/dfem/test_mass.cpp

@@ -12,6 +12,7 @@
 #include "../unit_tests.hpp"
 #include "../linalg/test_same_matrices.hpp"
 #include "mfem.hpp"
+#include <fem/dfem/doperator.hpp>
 
 #ifdef MFEM_USE_MPI
 

tests/unit/fem/test_pa_kernels.cpp

@@ -962,9 +962,9 @@ TEST_CASE("Dispatch Map Specializations")
 
    DiffusionIntegrator{};
    REQUIRE_FALSE(
-      DiffusionIntegrator::ApplyPAKernels::GetDispatchTable().empty());
+      DiffusionIntegrator::DiffusionApplyPAKernel::GetDispatchTable().empty());
    REQUIRE_FALSE(
-      DiffusionIntegrator::DiagonalPAKernels::GetDispatchTable().empty());
+      DiffusionIntegrator::DiffusionDiagonalPAKernel::GetDispatchTable().empty());
 
    Mesh mesh = Mesh::MakeCartesian2D(2, 2, Element::QUADRILATERAL);
    H1_FECollection fec(1, mesh.Dimension());