Update operator.hpp

config/config.hpp

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

fem/CMakeLists.txt

@@ -57,6 +57,8 @@ set(SRCS
   integ/lininteg_domain_grad.cpp
   integ/lininteg_domain_vectorfe.cpp
   integ/nonlininteg_vecconvection_pa.cpp
+  integ/nonlininteg_vecconvection_pa_diag.cpp
+  integ/nonlininteg_vecconvection_pa_grad.cpp
   integ/nonlininteg_vecconvection_mf.cpp
   coefficient.cpp
   complex_fem.cpp

fem/bilininteg.hpp

@@ -2651,14 +2651,22 @@ public:
    void AddMultMF(const Vector &x, Vector &y) const override;
    bool SupportsCeed() const override { return DeviceCanUseCeed(); }
 
+   // PA AddMultPA kernels
    using VectorMassAddMultPAType =
       void(*)(const int,  const int,
               const Array<real_t>&, const Vector&,
               const Vector&, Vector&, const int, const int);
-
    MFEM_REGISTER_KERNELS(VectorMassAddMultPA,
                          VectorMassAddMultPAType,
                          (int, int, int));
+
+   // PA DiagonalPA kernels
+   using VectorMassAssembleDiagonalPAType =
+      void(*)(const int, const int, const int,
+              const real_t*, const real_t*, real_t*);
+   MFEM_REGISTER_KERNELS(VectorMassAssembleDiagonalPA,
+                         VectorMassAssembleDiagonalPAType,
+                         (int /*dim*/, int /*q1d*/));
 };
 
 
@@ -3060,6 +3068,24 @@ public:
    void AddMultPA(const Vector &x, Vector &y) const override;
    void AddMultTransposePA(const Vector &x, Vector &y) const override;
 
+   using VectorDivergenceAddMultPAType =
+      void (*)(const int ne,
+               const Array<real_t> &b, const Array<real_t> &g, const Array<real_t> &bt,
+               const Vector &op, const Vector &x, Vector &y,
+               const int tr_d1d, const int te_d1d, const int q1d);
+   MFEM_REGISTER_KERNELS(VectorDivergenceAddMultPA,
+                         VectorDivergenceAddMultPAType,
+                         (int, int, int, int));
+
+   using VectorDivergenceAddMultTransposePAType =
+      void (*)(const int ne,
+               const Array<real_t> &bt, const Array<real_t> &gt, const Array<real_t> &b,
+               const Vector &q, const Vector &x, Vector &y,
+               const int tr_d1d, const int te_d1d, const int q1d);
+   MFEM_REGISTER_KERNELS(VectorDivergenceAddMultTransposePA,
+                         VectorDivergenceAddMultTransposePAType,
+                         (int, int, int, int));
+
    static const IntegrationRule &GetRule(const FiniteElement &trial_fe,
                                          const FiniteElement &test_fe,
                                          const ElementTransformation &Trans);

fem/integ/bilininteg_elasticity_kernels.cpp

@@ -91,15 +91,15 @@ void ElasticityAddMultPA(const int dim, const int nDofs,
 void ElasticityAssembleDiagonalPA(const int dim, const int nDofs,
                                   const CoefficientVector &lambda,
                                   const CoefficientVector &mu, const GeometricFactors &geom,
-                                  const DofToQuad &maps, QuadratureFunction &QVec, Vector &diag)
+                                  const DofToQuad &maps, const IntegrationRule &ir, Vector &diag)
 {
    switch (dim)
    {
       case 2:
-         ElasticityAssembleDiagonalPA_<2>(nDofs, lambda, mu, geom, maps, QVec, diag);
+         ElasticityAssembleDiagonalPA_<2>(nDofs, lambda, mu, geom, maps, ir, diag);
          break;
       case 3:
-         ElasticityAssembleDiagonalPA_<3>(nDofs, lambda, mu, geom, maps, QVec, diag);
+         ElasticityAssembleDiagonalPA_<3>(nDofs, lambda, mu, geom, maps, ir, diag);
          break;
       default:
          MFEM_ABORT("Only dimensions 2 and 3 supported.");

fem/integ/bilininteg_elasticity_kernels.hpp

@@ -38,7 +38,6 @@
 #include "../../linalg/vector.hpp"
 #include "../../linalg/tensor.hpp"
 #include "../quadinterpolator.hpp"
-#include "../bilininteg.hpp"
 #include "../coefficient.hpp"
 #include "../qfunction.hpp"
 
@@ -133,12 +132,12 @@ void ElasticityAssembleEA(const int dim, const int i_block, const int j_block,
 /// @param[in] mu Quadrature function for second Lame param.
 /// @param[in] geom Geometric factors corresponding to fespace.
 /// @param[in] maps DofToQuad maps for one element (assume elements all same).
-/// @param QVec Scratch Q-Vector. nQuad x dim x dim x dim x dim x numEls.
+/// @param[in] ir Integration rule.
 /// @param[out] diag diagonal of A. nDofs x dim x numEls.
 void ElasticityAssembleDiagonalPA(const int dim, const int nDofs,
                                   const CoefficientVector &lambda,
                                   const CoefficientVector &mu, const GeometricFactors &geom,
-                                  const DofToQuad &maps, QuadratureFunction &QVec, Vector &diag);
+                                  const DofToQuad &maps, const IntegrationRule &ir, Vector &diag);
 
 /// Templated implementation of ElasticityAddMultPA.
 template<int dim, int i_block = -1, int j_block = -1>
@@ -280,77 +279,62 @@ void ElasticityAddMultPA_(const int nDofs, const FiniteElementSpace &fespace,
 template<int dim>
 void ElasticityAssembleDiagonalPA_(const int nDofs,
                                    const CoefficientVector &lambda,
-                                   const CoefficientVector &mu, const GeometricFactors &geom,
-                                   const DofToQuad &maps, QuadratureFunction &QVec, Vector &diag)
+                                   const CoefficientVector &mu,
+                                   const GeometricFactors &geom,
+                                   const DofToQuad &maps,
+                                   const IntegrationRule &ir,
+                                   Vector &diag)
 {
-   using future::tensor;
-   using future::make_tensor;
    using future::det;
    using future::inv;
+   using future::make_tensor;
+   using future::tensor;
 
    // Assuming all elements are the same
-   const auto &ir = QVec.GetIntRule(0);
    static constexpr int d = dim;
    const int numPoints = ir.GetNPoints();
-   const int numEls = lambda.Size()/numPoints;
+   const int numEls = lambda.Size() / numPoints;
+
    const auto lamDev = Reshape(lambda.Read(), numPoints, numEls);
    const auto muDev = Reshape(mu.Read(), numPoints, numEls);
    const auto J = Reshape(geom.J.Read(), numPoints, d, d, numEls);
-   auto Q = Reshape(QVec.ReadWrite(), numPoints, d,d, d, numEls);
    const real_t *ipWeights = ir.GetWeights().Read();
-   mfem::forall_2D(numEls, numPoints,1, [=] MFEM_HOST_DEVICE (int e)
-   {
-      MFEM_FOREACH_THREAD(p, x,numPoints)
-      {
-         auto invJ = inv(make_tensor<d, d>(
-         [&](int i, int j) { return J(p, i, j, e); }));
-         const real_t w = ipWeights[p] /det(invJ);
-         for (int n = 0; n < d; n++)
-         {
-            for (int m = 0; m < d; m++)
-            {
-               for (int q = 0; q < d; q++)
-               {
-                  // compute contraction of 4*sym(grad(u))sym(grad(v)) term.
-                  // this contraction could be made slightly cheaper using Voigt
-                  // notation, but repeated entries are summed for simplicity.
-                  real_t contraction = 0.;
-                  for (int a = 0; a < d; a++)
-                  {
-                     for (int b = 0; b < d; b++)
-                     {
-                        contraction += ((a == q)*invJ(m,b) + (b==q)*invJ(m,a))*((a == q)
-                                                                                *invJ(n, b) + (b==q)*invJ(n,a));
-                     }
-                  }
-                  // lambda*div(u)*div(v) + 2*mu*sym(grad(u))*sym(grad(v))
-                  // contraction = 4*sym(grad(u))sym(grad(v))
-                  Q(p,m,n,q,e) = w*(lamDev(p, e)*invJ(m,q)*invJ(n,q)
-                                    + 0.5*muDev(p, e)*contraction);
-               }
-            }
-         }
-      }
-   });
-
-   // Reduce quadrature function to an E-Vector
-   const auto QRead = Reshape(QVec.Read(), numPoints, d, d, d, numEls);
-   auto diagDev = Reshape(diag.Write(), nDofs, d, numEls);
    const auto G = Reshape(maps.G.Read(), numPoints, d, nDofs);
+   auto diagDev = Reshape(diag.Write(), nDofs, d, numEls);
+
    mfem::forall_2D(numEls, d, nDofs, [=] MFEM_HOST_DEVICE (int e)
    {
-      MFEM_FOREACH_THREAD(i, y, nDofs)
+      MFEM_FOREACH_THREAD_DIRECT(i, y, nDofs)
       {
-         MFEM_FOREACH_THREAD(q, x, d)
+         MFEM_FOREACH_THREAD_DIRECT(q, x, d)
          {
-            real_t sum = 0.;
-            for (int n = 0; n < d; n++)
+            real_t sum = 0.0;
+            for (int p = 0; p < numPoints; p++)
             {
-               for (int m = 0; m < d; m++)
+               auto invJ = inv(make_tensor<d, d>([&](int r, int c) { return J(p, r, c, e); }));
+               const real_t w = ipWeights[p] / det(invJ);
+
+               for (int n = 0; n < d; n++)
                {
-                  for (int p = 0; p < numPoints; p++ )
+                  for (int m = 0; m < d; m++)
                   {
-                     sum += QRead(p,m,n,q,e)*G(p,m,i)*G(p,n,i);
+                     // compute contraction of 4*sym(grad(u))sym(grad(v)) term.
+                     // this contraction could be made slightly cheaper using Voigt
+                     // notation, but repeated entries are summed for simplicity.
+                     real_t contraction = 0.;
+                     for (int a = 0; a < d; a++)
+                     {
+                        for (int b = 0; b < d; b++)
+                        {
+                           contraction += ((a == q) * invJ(m, b) + (b == q) * invJ(m, a)) *
+                                          ((a == q) * invJ(n, b) + (b == q) * invJ(n, a));
+                        }
+                     }
+                     // lambda*div(u)*div(v) + 2*mu*sym(grad(u))*sym(grad(v))
+                     // contraction = 4*sym(grad(u))sym(grad(v))
+                     const real_t Q = w *
+                                      (lamDev(p, e) * invJ(m, q) * invJ(n, q) + 0.5 * muDev(p, e) * contraction);
+                     sum += Q * G(p, m, i) * G(p, n, i);
                   }
                }
             }

fem/integ/bilininteg_elasticity_pa.cpp

@@ -10,7 +10,6 @@
 // CONTRIBUTING.md for details.
 
 #include "../bilininteg.hpp"
-#include "../gridfunc.hpp"
 #include "../qfunction.hpp"
 #include "bilininteg_elasticity_kernels.hpp"
 
@@ -59,9 +58,8 @@ void ElasticityIntegrator::AssemblePA(const FiniteElementSpace &fes)
 
 void ElasticityIntegrator::AssembleDiagonalPA(Vector &diag)
 {
-   q_vec->SetVDim(vdim*vdim*vdim*vdim);
    internal::ElasticityAssembleDiagonalPA(vdim, ndofs, *lambda_quad, *mu_quad,
-                                          *geom, *maps, *q_vec, diag);
+                                          *geom, *maps, *IntRule, diag);
 }
 
 void ElasticityIntegrator::AddMultPA(const Vector &x, Vector &y) const

fem/integ/bilininteg_vecmass_pa.cpp

@@ -205,157 +205,40 @@ void VectorMassIntegrator::AddMultPA(const Vector &x, Vector &y) const
 
 }
 
-template <const int T_D1D = 0, const int T_Q1D = 0>
-static void PAVectorMassAssembleDiagonal2D(const int NE,
-                                           const Array<real_t> &b,
-                                           const Vector &pa_data, Vector &diag,
-                                           const int d1d = 0, const int q1d = 0)
-{
-   constexpr int VDIM = 2;
-   const int D1D = T_D1D ? T_D1D : d1d;
-   const int Q1D = T_Q1D ? T_Q1D : q1d;
-   MFEM_VERIFY(D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
-   MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
-   const auto B = Reshape(b.Read(), Q1D, D1D);
-   const auto D = Reshape(pa_data.Read(), Q1D, Q1D, NE);
-   auto Y = Reshape(diag.ReadWrite(), D1D, D1D, VDIM, NE);
-
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE(int e)
-   {
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
-      constexpr int max_D1D = T_D1D ? T_D1D : DofQuadLimits::MAX_D1D;
-      constexpr int max_Q1D = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
-
-      real_t temp[max_Q1D][max_D1D];
-      for (int qx = 0; qx < Q1D; ++qx)
-      {
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            temp[qx][dy] = 0.0;
-            for (int qy = 0; qy < Q1D; ++qy)
-            {
-               temp[qx][dy] += B(qy, dy) * B(qy, dy) * D(qx, qy, e);
-            }
-         }
-      }
-      for (int dy = 0; dy < D1D; ++dy)
-      {
-         for (int dx = 0; dx < D1D; ++dx)
-         {
-            real_t temp1 = 0.0;
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               temp1 += B(qx, dx) * B(qx, dx) * temp[qx][dy];
-            }
-            Y(dx, dy, 0, e) = temp1;
-            Y(dx, dy, 1, e) = temp1;
-         }
-      }
-   });
-}
-
-template <const int T_D1D = 0, const int T_Q1D = 0>
-static void PAVectorMassAssembleDiagonal3D(const int NE,
-                                           const Array<real_t> &B_,
-                                           const Vector &pa_data, Vector &diag,
-                                           const int d1d = 0, const int q1d = 0)
-{
-   constexpr int VDIM = 3;
-   const int D1D = T_D1D ? T_D1D : d1d;
-   const int Q1D = T_Q1D ? T_Q1D : q1d;
-   MFEM_VERIFY(D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
-   MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
-   const auto B = Reshape(B_.Read(), Q1D, D1D);
-   MFEM_VERIFY(pa_data.Size() == Q1D * Q1D * Q1D * NE, "pa_data size error");
-   const auto D = Reshape(pa_data.Read(), Q1D, Q1D, Q1D, NE);
-   auto Y = Reshape(diag.ReadWrite(), D1D, D1D, D1D, VDIM, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE(int e)
-   {
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
-      // the following variables are evaluated at compile time
-      constexpr int max_D1D = T_D1D ? T_D1D : DofQuadLimits::MAX_D1D;
-      constexpr int max_Q1D = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
-
-      real_t temp[max_Q1D][max_Q1D][max_D1D];
-      for (int qx = 0; qx < Q1D; ++qx)
-      {
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            for (int dz = 0; dz < D1D; ++dz)
-            {
-               temp[qx][qy][dz] = 0.0;
-               for (int qz = 0; qz < Q1D; ++qz)
-               {
-                  temp[qx][qy][dz] +=
-                     B(qz, dz) * B(qz, dz) * D(qx, qy, qz, e);
-               }
-            }
-         }
-      }
-      real_t temp2[max_Q1D][max_D1D][max_D1D];
-      for (int qx = 0; qx < Q1D; ++qx)
-      {
-         for (int dz = 0; dz < D1D; ++dz)
-         {
-            for (int dy = 0; dy < D1D; ++dy)
-            {
-               temp2[qx][dy][dz] = 0.0;
-               for (int qy = 0; qy < Q1D; ++qy)
-               {
-                  temp2[qx][dy][dz] +=
-                     B(qy, dy) * B(qy, dy) * temp[qx][qy][dz];
-               }
-            }
-         }
-      }
-      for (int dz = 0; dz < D1D; ++dz)
-      {
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               real_t temp3 = 0.0;
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  temp3 += B(qx, dx) * B(qx, dx) * temp2[qx][dy][dz];
-               }
-               Y(dx, dy, dz, 0, e) = temp3;
-               Y(dx, dy, dz, 1, e) = temp3;
-               Y(dx, dy, dz, 2, e) = temp3;
-            }
-         }
-      }
-   });
-}
-
-static void PAVectorMassAssembleDiagonal(const int dim, const int D1D,
-                                         const int Q1D, const int NE,
-                                         const Array<real_t> &B,
-                                         const Vector &pa_data,
-                                         Vector &diag)
-{
-   if (dim == 2)
-   {
-      return PAVectorMassAssembleDiagonal2D(NE, B, pa_data, diag, D1D, Q1D);
-   }
-   else if (dim == 3)
-   {
-      return PAVectorMassAssembleDiagonal3D(NE, B, pa_data, diag, D1D, Q1D);
-   }
-   MFEM_ABORT("Dimension not implemented.");
-}
-
 void VectorMassIntegrator::AssembleDiagonalPA(Vector &diag)
 {
-   if (DeviceCanUseCeed()) { ceedOp->GetDiagonal(diag); }
-   else
-   {
-      MFEM_VERIFY(coeff_vdim == 1, "coeff_vdim != 1");
-      MFEM_VERIFY(!VQ && !MQ, "VQ and MQ not supported");
-      PAVectorMassAssembleDiagonal(dim, dofs1D, quad1D, ne, maps->B, pa_data, diag);
-   }
+   if (DeviceCanUseCeed()) { return ceedOp->GetDiagonal(diag); }
+
+   MFEM_VERIFY(coeff_vdim == 1, "coeff_vdim != 1");
+   MFEM_VERIFY(!VQ && !MQ, "VQ and MQ not supported");
+
+   // Add the VectorMassAssembleDiagonalPA specializations
+   static const auto vector_mass_assemble_diagonal_kernel_specializations =
+      ( // 2D
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 2>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 3>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 4>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 5>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 6>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 7>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<2, 8>::Add(),
+         // 3D
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 2>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 3>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 4>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 5>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 6>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 7>::Add(),
+         VectorMassIntegrator::VectorMassAssembleDiagonalPA::Specialization<3, 8>::Add(),
+         true);
+   MFEM_CONTRACT_VAR(vector_mass_assemble_diagonal_kernel_specializations);
+
+   VectorMassAssembleDiagonalPA::Run(dim, quad1D, // templated arguments
+                                     ne, dofs1D, quad1D,
+                                     maps->B.Read(),
+                                     pa_data.Read(),
+                                     diag.ReadWrite());
+
 }
 
 } // namespace mfem

fem/integ/bilininteg_vecmass_pa.hpp

@@ -176,8 +176,145 @@ void SmemPAVectorMassApply3D(const int NE,
    });
 }
 
+template <int T_Q1D = 0, int T_MDQ = 16>
+static void SmemPAVectorMassAssembleDiagonal2D(const int ne,
+                                               const int d1d,
+                                               const int q1d,
+                                               const real_t *b_r,
+                                               const real_t *d_r,
+                                               real_t *y_rw)
+{
+   constexpr int VDIM = 2;
+
+   const int D1D = d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   MFEM_VERIFY(Q1D <= T_MDQ && D1D <= Q1D, "");
+
+   const auto B = Reshape(b_r, Q1D, D1D);
+   const auto D = Reshape(d_r, Q1D, Q1D, ne);
+   auto Y = Reshape(y_rw, D1D, D1D, VDIM, ne);
+
+   mfem::forall_2D<T_Q1D*T_Q1D>(
+      ne, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   {
+      constexpr int MQ1 = T_Q1D ? T_Q1D : T_MDQ;
+
+      MFEM_SHARED real_t sm[MQ1][MQ1];
+
+      MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+         {
+            real_t u = 0.0;
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               u += B(qy, dy) * B(qy, dy) * D(qx, qy, e);
+            }
+            sm[qx][dy] = u;
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(dx, x, D1D)
+         {
+            real_t u = 0.0;
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               u += B(qx, dx) * B(qx, dx) * sm[qx][dy];
+            }
+            Y(dx, dy, 0, e) = u;
+            Y(dx, dy, 1, e) = u;
+         }
+      }
+   });
+}
+
+template <int T_Q1D = 0, int T_MDQ = 12>
+static void SmemPAVectorMassAssembleDiagonal3D(const int ne,
+                                               const int d1d,
+                                               const int q1d,
+                                               const real_t *b_r,
+                                               const real_t *d_r,
+                                               real_t *y_rw)
+{
+   constexpr int VDIM = 3;
+
+   const int D1D = d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   MFEM_VERIFY(Q1D <= T_MDQ && D1D <= Q1D, "");
+
+   const auto B = Reshape(b_r, Q1D, D1D);
+   const auto D = Reshape(d_r, Q1D, Q1D, Q1D, ne);
+   auto Y = Reshape(y_rw, 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 ? T_Q1D : T_MDQ;
+
+      MFEM_SHARED real_t sm[2][MQ1][MQ1][MQ1];
+
+      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;
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  u += B(qz, dz) * B(qz, dz) * D(qx, qy, qz, e);
+               }
+               sm[0][dz][qy][qx] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD_DIRECT(dz, z, D1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+            {
+               real_t u = 0.0;
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  u += B(qy, dy) * B(qy, dy) * sm[0][dz][qy][qx];
+               }
+               sm[1][dz][dy][qx] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      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;
+               for (int qx = 0; qx < Q1D; ++qx)
+               {
+                  u += B(qx, dx) * B(qx, dx) * sm[1][dz][dy][qx];
+               }
+               Y(dx, dy, dz, 0, e) = u;
+               Y(dx, dy, dz, 1, e) = u;
+               Y(dx, dy, dz, 2, e) = u;
+            }
+         }
+      }
+   });
+}
+
 } // namespace internal
 
+// AddMultPA kernels
 template<int DIM, int T_D1D, int T_Q1D>
 VectorMassIntegrator::VectorMassAddMultPAType
 VectorMassIntegrator::VectorMassAddMultPA::Kernel()
@@ -190,11 +327,11 @@ VectorMassIntegrator::VectorMassAddMultPA::Kernel()
    {
       return internal::SmemPAVectorMassApply3D<T_D1D, T_Q1D>;
    }
-   MFEM_ABORT("Unsupported kernel");
+   else { MFEM_ABORT("Unsupported kernel"); }
 }
 
 inline VectorMassIntegrator::VectorMassAddMultPAType
-VectorMassIntegrator::VectorMassAddMultPA::Fallback(int dim, int d1d, int q1d)
+VectorMassIntegrator::VectorMassAddMultPA::Fallback(int dim, int, int)
 {
    if (dim == 2)
    {
@@ -207,6 +344,36 @@ VectorMassIntegrator::VectorMassAddMultPA::Fallback(int dim, int d1d, int q1d)
    else { MFEM_ABORT("Unsupported kernel"); }
 }
 
+// DiagonalPA kernels
+template<int DIM, int T_Q1D>
+VectorMassIntegrator::VectorMassAssembleDiagonalPAType
+VectorMassIntegrator::VectorMassAssembleDiagonalPA::Kernel()
+{
+   if constexpr (DIM == 2)
+   {
+      return internal::SmemPAVectorMassAssembleDiagonal2D<T_Q1D>;
+   }
+   else if constexpr (DIM == 3)
+   {
+      return internal::SmemPAVectorMassAssembleDiagonal3D<T_Q1D>;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+inline VectorMassIntegrator::VectorMassAssembleDiagonalPAType
+VectorMassIntegrator::VectorMassAssembleDiagonalPA::Fallback(int dim, int)
+{
+   if (dim == 2)
+   {
+      return internal::SmemPAVectorMassAssembleDiagonal2D;
+   }
+   else if (dim == 3)
+   {
+      return internal::SmemPAVectorMassAssembleDiagonal3D;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
 /// \endcond DO_NOT_DOCUMENT
 
 } // namespace mfem

fem/integ/nonlininteg_vecconvection_pa_diag.cpp

@@ -0,0 +1,351 @@
+// 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 <utility>
+
+#include "../kernels.hpp"
+#include "../nonlininteg.hpp"
+#include "../../general/forall.hpp"
+
+namespace mfem
+{
+
+template<int T_D1D = 0, int T_Q1D = 0, int T_MDQ = 16>
+static void SmemPAConvectionNLGradDiagonal2D(const int NE,
+                                             const real_t *b,
+                                             const real_t *g,
+                                             const real_t *a,
+                                             const real_t *u,
+                                             real_t *de,
+                                             const int d1d,
+                                             const int q1d)
+{
+   static constexpr int VDIM = 2, DIM = 2;
+   const int D1D = T_D1D ? T_D1D : d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   const auto A = Reshape(a, VDIM, DIM, Q1D, Q1D, NE);
+   const auto U = Reshape(u, D1D, D1D, VDIM, NE);
+   auto D = Reshape(de, D1D, D1D, VDIM, NE);
+
+   mfem::forall_2D<T_Q1D * T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   {
+      constexpr int MD1 = T_D1D ? T_D1D : T_MDQ;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : T_MDQ;
+
+      MFEM_SHARED real_t sM[3][MQ1][MQ1], sQ[3][MQ1][MQ1];
+      MFEM_SHARED real_t sB[MD1][MQ1], sG[MD1][MQ1];
+
+      kernels::internal::v_regs2d_t<VDIM, MQ1> r0, r1;
+      kernels::internal::vd_regs2d_t<VDIM, DIM, MQ1> g0, g1;
+
+      kernels::internal::LoadMatrix(D1D, Q1D, b, sB);
+      kernels::internal::LoadMatrix(D1D, Q1D, g, sG);
+
+      kernels::internal::LoadDofs2d(e, D1D, U, r0);
+      kernels::internal::Eval2d(D1D, Q1D, sM[0], sB, r0, r1);
+
+      kernels::internal::LoadDofs2d(e, D1D, U, g0);
+      kernels::internal::Grad2d(D1D, Q1D, sM[0], sB, sG, g0, g1);
+
+      for (int v = 0; v < VDIM; ++v)
+      {
+         future::tensor<real_t, VDIM> e_v = {};
+         e_v[v] = real_t(1);
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               const future::tensor<real_t, VDIM> u_val =
+               {
+                  r1[0][qy][qx], r1[1][qy][qx]
+               };
+               const future::tensor<real_t, VDIM, DIM> Q_adj =
+               {
+                  {  { A(0, 0, qx, qy, e), A(1, 0, qx, qy, e) },
+                     { A(0, 1, qx, qy, e), A(1, 1, qx, qy, e) }
+                  }
+               };
+               const future::tensor<real_t, VDIM, DIM> grad_U =
+               {
+                  {  { g1[0][0][qy][qx], g1[1][0][qy][qx] },
+                     { g1[0][1][qy][qx], g1[1][1][qy][qx] }
+                  }
+               };
+               const auto one = Q_adj * u_val;
+               const auto two = transpose(grad_U) * (Q_adj * e_v);
+               sQ[0][qx][qy] = one[0];
+               sQ[1][qx][qy] = one[1];
+               sQ[2][qx][qy] = two[v];
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+            {
+               real_t s[3] = {};
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  const real_t By = sB[dy][qy], Gy = sG[dy][qy];
+                  s[0] += By * By * sQ[0][qx][qy];
+                  s[1] += Gy * By * sQ[1][qx][qy];
+                  s[2] += By * By * sQ[2][qx][qy];
+               }
+               sM[0][qx][dy] = s[0];
+               sM[1][qx][dy] = s[1];
+               sM[2][qx][dy] = s[2];
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(dx, x, D1D)
+            {
+               real_t d = 0.0;
+               for (int qx = 0; qx < Q1D; ++qx)
+               {
+                  const real_t Bx = sB[dx][qx], Gx = sG[dx][qx];
+                  d += Gx * Bx * sM[0][qx][dy] +
+                       Bx * Bx * sM[1][qx][dy] +
+                       Bx * Bx * sM[2][qx][dy];
+               }
+               D(dx, dy, v, e) += d;
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   });
+}
+
+template<int T_D1D = 0, int T_Q1D = 0, int T_MDQ = 16>
+static void SmemPAConvectionNLGradDiagonal3D(const int NE,
+                                             const real_t *b,
+                                             const real_t *g,
+                                             const real_t *a,
+                                             const real_t *u,
+                                             real_t *de,
+                                             const int d1d,
+                                             const int q1d)
+{
+   static constexpr int VDIM = 3, DIM = 3;
+   const int D1D = T_D1D ? T_D1D : d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   const auto A = Reshape(a, VDIM, DIM, Q1D, Q1D, Q1D, NE);
+   const auto U = Reshape(u, D1D, D1D, D1D, VDIM, NE);
+   auto D = Reshape(de, 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 ? T_D1D : T_MDQ;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : T_MDQ;
+
+      MFEM_SHARED real_t sM[4][MQ1][MQ1], sQ[4][MQ1][MQ1];
+      MFEM_SHARED real_t sB[MD1][MQ1], sG[MD1][MQ1];
+
+      kernels::internal::v_regs3d_t<VDIM, MQ1> r0, r1;
+      kernels::internal::vd_regs3d_t<VDIM, DIM, MQ1> g0, g1;
+
+      kernels::internal::LoadMatrix(D1D, Q1D, b, sB);
+      kernels::internal::LoadMatrix(D1D, Q1D, g, sG);
+
+      kernels::internal::LoadDofs3d(e, D1D, U, r0);
+      kernels::internal::Eval3d(D1D, Q1D, sM[0], sB, r0, r1);
+
+      kernels::internal::LoadDofs3d(e, D1D, U, g0);
+      kernels::internal::Grad3d(D1D, Q1D, sM[0], sB, sG, g0, g1);
+
+      for (int v = 0; v < VDIM; ++v)
+      {
+         future::tensor<real_t, VDIM> e_v = {};
+         e_v[v] = real_t(1);
+         for (int dz = 0; dz < D1D; ++dz)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  real_t s[4] = {};
+                  for (int qz = 0; qz < Q1D; ++qz)
+                  {
+                     const future::tensor<real_t, VDIM> u_val =
+                     {
+                        r1[0][qz][qy][qx], r1[1][qz][qy][qx], r1[2][qz][qy][qx]
+                     };
+                     const future::tensor<real_t, VDIM, DIM> Q_adj = {{
+                           {A(0,0,qx,qy,qz,e), A(1,0,qx,qy,qz,e), A(2,0,qx,qy,qz,e)},
+                           {A(0,1,qx,qy,qz,e), A(1,1,qx,qy,qz,e), A(2,1,qx,qy,qz,e)},
+                           {A(0,2,qx,qy,qz,e), A(1,2,qx,qy,qz,e), A(2,2,qx,qy,qz,e)}
+                        }
+                     };
+                     const future::tensor<real_t, VDIM, DIM> grad_U = {{
+                           {g1[0][0][qz][qy][qx], g1[1][0][qz][qy][qx], g1[2][0][qz][qy][qx]},
+                           {g1[0][1][qz][qy][qx], g1[1][1][qz][qy][qx], g1[2][1][qz][qy][qx]},
+                           {g1[0][2][qz][qy][qx], g1[1][2][qz][qy][qx], g1[2][2][qz][qy][qx]}
+                        }
+                     };
+                     const auto one = Q_adj * u_val;
+                     const auto two = transpose(grad_U) * (Q_adj * e_v);
+
+                     const real_t Bz = sB[dz][qz], Gz = sG[dz][qz];
+                     s[0] += one[0] * Bz * Bz;
+                     s[1] += one[1] * Bz * Bz;
+                     s[2] += one[2] * Bz * Gz;
+                     s[3] += two[v] * Bz * Bz;
+                  }
+                  sQ[0][qx][qy] = s[0];
+                  sQ[1][qx][qy] = s[1];
+                  sQ[2][qx][qy] = s[2];
+                  sQ[3][qx][qy] = s[3];
+               }
+            }
+            MFEM_SYNC_THREAD;
+
+            MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  real_t s[4] = {};
+                  for (int qy = 0; qy < Q1D; ++qy)
+                  {
+                     const real_t By = sB[dy][qy], Gy = sG[dy][qy];
+                     s[0] += By * By * sQ[0][qx][qy];
+                     s[1] += Gy * By * sQ[1][qx][qy];
+                     s[2] += By * By * sQ[2][qx][qy];
+                     s[3] += By * By * sQ[3][qx][qy];
+                  }
+                  sM[0][dy][qx] = s[0];
+                  sM[1][dy][qx] = s[1];
+                  sM[2][dy][qx] = s[2];
+                  sM[3][dy][qx] = s[3];
+               }
+            }
+            MFEM_SYNC_THREAD;
+
+            MFEM_FOREACH_THREAD_DIRECT(dy, y, D1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(dx, x, D1D)
+               {
+                  real_t d = 0.0;
+                  for (int qx = 0; qx < Q1D; ++qx)
+                  {
+                     const real_t Bx = sB[dx][qx], Gx = sG[dx][qx];
+                     d += Gx * Bx * sM[0][dy][qx];
+                     d += Bx * Bx * sM[1][dy][qx];
+                     d += Bx * Bx * sM[2][dy][qx];
+                     d += Bx * Bx * sM[3][dy][qx];
+                  }
+                  D(dx, dy, dz, v, e) += d;
+               }
+            }
+            MFEM_SYNC_THREAD;
+         }
+      }
+   });
+}
+
+void VectorConvectionNLFIntegrator::AssembleGradDiagonalPA(Vector &de) const
+{
+   if (dim == 2)
+   {
+      if (static auto ini = false; !std::exchange(ini, true))
+      {
+         VectorConvectionNLFGradDiagPA2D::Specialization<2, 2>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<2, 3>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<3, 4>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<3, 5>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<4, 5>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<4, 6>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<5, 7>::Add();
+         VectorConvectionNLFGradDiagPA2D::Specialization<5, 8>::Add();
+      }
+      VectorConvectionNLFGradDiagPA2D::Run(d1d, q1d, ne,
+                                           maps->B.Read(),
+                                           maps->G.Read(),
+                                           pa_adj.Read(),
+                                           pa_u.Read(),
+                                           de.ReadWrite(),
+                                           d1d, q1d);
+   }
+   else if (dim == 3)
+   {
+      if (static auto ini = false; !std::exchange(ini, true))
+      {
+         VectorConvectionNLFGradDiagPA3D::Specialization<2, 3>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<2, 4>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<2, 5>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<3, 4>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<3, 5>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<3, 6>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<4, 6>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<4, 7>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<4, 8>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<5, 7>::Add();
+         VectorConvectionNLFGradDiagPA3D::Specialization<5, 8>::Add();
+      }
+      VectorConvectionNLFGradDiagPA3D::Run(d1d, q1d, ne,
+                                           maps->B.Read(),
+                                           maps->G.Read(),
+                                           pa_adj.Read(),
+                                           pa_u.Read(),
+                                           de.ReadWrite(),
+                                           d1d, q1d);
+   }
+   else
+   {
+      MFEM_ABORT("Unsupported dimension");
+   }
+}
+
+/// \cond DO_NOT_DOCUMENT
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPA2D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return SmemPAConvectionNLGradDiagonal2D<T_D1D, T_Q1D>;
+}
+
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPA2D::Fallback
+(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= 16, "d1d > 16 is not supported");
+   MFEM_VERIFY(q1d <= 16, "q1d > 16 is not supported");
+   return SmemPAConvectionNLGradDiagonal2D<>;
+}
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPA3D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return SmemPAConvectionNLGradDiagonal3D<T_D1D, T_Q1D>;
+}
+
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFGradDiagPA3D::Fallback
+(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= 16, "d1d > 16 is not supported");
+   MFEM_VERIFY(q1d <= 16, "q1d > 16 is not supported");
+   return SmemPAConvectionNLGradDiagonal3D<>;
+}
+
+/// \endcond DO_NOT_DOCUMENT
+
+} // namespace mfem

fem/integ/nonlininteg_vecconvection_pa_grad.cpp

@@ -0,0 +1,319 @@
+// 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 <utility>
+
+#include "../kernels.hpp"
+#include "../nonlininteg.hpp"
+#include "../../general/forall.hpp"
+
+namespace mfem
+{
+
+void VectorConvectionNLFIntegrator::AssembleGradPA(
+   const Vector &u, const FiniteElementSpace &fes)
+{
+   this->pa_u = u;
+   AssemblePA(fes);
+
+   if (static auto done = false; !std::exchange(done, true))
+   {
+      // 2D
+      VectorConvectionNLFAddMultGradPA2D::Specialization<2, 2>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<2, 3>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<3, 4>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<3, 5>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<4, 5>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<4, 6>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<5, 7>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<5, 8>::Add();
+      VectorConvectionNLFAddMultGradPA2D::Specialization<6, 8>::Add();
+      // 3D
+      VectorConvectionNLFAddMultGradPA3D::Specialization<2, 3>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<2, 4>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<2, 5>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<3, 4>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<3, 5>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<3, 6>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<4, 5>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<4, 6>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<5, 6>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<4, 7>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<4, 8>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<5, 7>::Add();
+      VectorConvectionNLFAddMultGradPA3D::Specialization<5, 8>::Add();
+   }
+}
+
+template<int T_D1D = 0, int T_Q1D = 0, int T_MDQ = 16>
+static void SmemPAConvectionNLGradApply2D(const int ne,
+                                          const real_t *b,
+                                          const real_t *g,
+                                          const real_t *a,
+                                          const real_t *u,
+                                          const real_t *du,
+                                          real_t *y,
+                                          const int d1d,
+                                          const int q1d)
+{
+   static constexpr int VDIM = 2, DIM = 2;
+   const int D1D = T_D1D ? T_D1D : d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   const auto A = Reshape(a, VDIM, DIM, Q1D, Q1D, ne);
+   const auto U = Reshape(u, D1D, D1D, VDIM, ne);
+   const auto dU = Reshape(du, D1D, D1D, VDIM, ne);
+   auto Y = Reshape(y, D1D, D1D, VDIM, ne);
+
+   mfem::forall_2D<T_Q1D * T_Q1D>(ne, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   {
+      constexpr int MD1 = T_D1D ? T_D1D : T_MDQ;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : T_MDQ;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1];
+      MFEM_SHARED real_t sB[MD1][MQ1], sG[MD1][MQ1];
+
+      kernels::internal::vd_regs2d_t<VDIM, DIM, MQ1> g0, g1, g2;
+      kernels::internal::v_regs2d_t<DIM, MQ1> r0, r1, r2;
+
+      kernels::internal::LoadMatrix(D1D, Q1D, b, sB);
+      kernels::internal::LoadMatrix(D1D, Q1D, g, sG);
+
+      kernels::internal::LoadDofs2d(e, D1D, dU, g0);
+      kernels::internal::Grad2d(D1D, Q1D, smem, sB, sG, g0, g1); // δu gradient
+
+      kernels::internal::LoadDofs2d(e, D1D, U, r0);
+      kernels::internal::Eval2d(D1D, Q1D, smem, sB, r0, r2); // u value
+
+      kernels::internal::LoadDofs2d(e, D1D, dU, r0);
+      kernels::internal::Eval2d(D1D, Q1D, smem, sB, r0, r1); // δu value
+
+      kernels::internal::LoadDofs2d(e, D1D, U, g0);
+      kernels::internal::Grad2d(D1D, Q1D, smem, sB, sG, g0, g2); // u gradient
+
+      MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            // First part of the Jacobian: u·∇δu
+            const future::tensor<real_t, DIM> u_val =
+            {
+               r2[0][qy][qx], r2[1][qy][qx]
+            };
+            const future::tensor<real_t, VDIM, DIM> Q_adj =
+            {
+               {  { A(0, 0, qx, qy, e), A(1, 0, qx, qy, e) },
+                  { A(0, 1, qx, qy, e), A(1, 1, qx, qy, e) }
+               }
+            };
+            const future::tensor<real_t, VDIM, DIM> grad_dU =
+            {
+               {  { g1[0][0][qy][qx], g1[1][0][qy][qx] },
+                  { g1[0][1][qy][qx], g1[1][1][qy][qx] }
+               }
+            };
+            const auto one = transpose(grad_dU) * (Q_adj * u_val);
+
+            // Second part of the Jacobian: δu·∇u
+            const future::tensor<real_t, DIM> du_val =
+            {
+               r1[0][qy][qx], r1[1][qy][qx]
+            };
+            const future::tensor<real_t, VDIM, DIM> grad_U =
+            {
+               {  { g2[0][0][qy][qx], g2[1][0][qy][qx] },
+                  { g2[0][1][qy][qx], g2[1][1][qy][qx] }
+               }
+            };
+            const auto two = transpose(grad_U) * (Q_adj * du_val);
+
+            // u⋅∇δu + δu⋅∇u
+            r0[0][qy][qx] = one[0] + two[0];
+            r0[1][qy][qx] = one[1] + two[1];
+         }
+      }
+      MFEM_SYNC_THREAD;
+      kernels::internal::EvalTranspose2d(D1D, Q1D, smem, sB, r0, r1);
+      kernels::internal::WriteDofs2d(e, D1D, r1, Y);
+   });
+}
+
+template<int T_D1D = 0, int T_Q1D = 0, int T_MDQ = 16>
+static void SmemPAConvectionNLGradApply3D(const int ne,
+                                          const real_t *b,
+                                          const real_t *g,
+                                          const real_t *a,
+                                          const real_t *u,
+                                          const real_t *du,
+                                          real_t *y,
+                                          const int d1d,
+                                          const int q1d)
+{
+   static constexpr int VDIM = 3, DIM = 3;
+   const int D1D = T_D1D ? T_D1D : d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   const auto A = Reshape(a, VDIM, DIM, Q1D, Q1D, Q1D, ne);
+   const auto U = Reshape(u, D1D, D1D, D1D, VDIM, ne);
+   const auto dU = Reshape(du, D1D, D1D, D1D, VDIM, ne);
+   auto Y = Reshape(y, 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 ? T_D1D : T_MDQ;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : T_MDQ;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1];
+      MFEM_SHARED real_t sB[MD1][MQ1], sG[MD1][MQ1];
+
+      kernels::internal::v_regs3d_t<VDIM, MQ1> r0, r1, r2;
+      kernels::internal::vd_regs3d_t<VDIM, DIM, MQ1> g0, g1, g2;
+
+      kernels::internal::LoadMatrix(D1D, Q1D, b, sB);
+      kernels::internal::LoadMatrix(D1D, Q1D, g, sG);
+
+      kernels::internal::LoadDofs3d(e, D1D, dU, g0);
+      kernels::internal::Grad3d(D1D, Q1D, smem, sB, sG, g0, g1); // δu gradient
+
+      kernels::internal::LoadDofs3d(e, D1D, U, r0);
+      kernels::internal::Eval3d(D1D, Q1D, smem, sB, r0, r2); // u value
+
+      kernels::internal::LoadDofs3d(e, D1D, dU, r0);
+      kernels::internal::Eval3d(D1D, Q1D, smem, sB, r0, r1); // δu value
+
+      kernels::internal::LoadDofs3d(e, D1D, U, g0);
+      kernels::internal::Grad3d(D1D, Q1D, smem, sB, sG, g0, g2); // u gradient
+
+      for (int qz = 0; qz < Q1D; qz++)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+            {
+               // First part of the Jacobian: u·∇δu
+               const future::tensor<real_t, DIM> u_val =
+               {
+                  r2[0][qz][qy][qx],
+                  r2[1][qz][qy][qx],
+                  r2[2][qz][qy][qx]
+               };
+               const future::tensor<real_t, VDIM, DIM> Q_adj = {{
+                     {A(0,0,qx,qy,qz,e), A(1,0,qx,qy,qz,e), A(2,0,qx,qy,qz,e)},
+                     {A(0,1,qx,qy,qz,e), A(1,1,qx,qy,qz,e), A(2,1,qx,qy,qz,e)},
+                     {A(0,2,qx,qy,qz,e), A(1,2,qx,qy,qz,e), A(2,2,qx,qy,qz,e)}
+                  }
+               };
+               const future::tensor<real_t, DIM, DIM> grad_dU = {{
+                     {g1[0][0][qz][qy][qx], g1[1][0][qz][qy][qx], g1[2][0][qz][qy][qx]},
+                     {g1[0][1][qz][qy][qx], g1[1][1][qz][qy][qx], g1[2][1][qz][qy][qx]},
+                     {g1[0][2][qz][qy][qx], g1[1][2][qz][qy][qx], g1[2][2][qz][qy][qx]}
+                  }
+               };
+               const auto one = transpose(grad_dU) * (Q_adj * u_val);
+
+               // Second part of the Jacobian: δu·∇u
+               const future::tensor<real_t, DIM> du_val =
+               {
+                  r1[0][qz][qy][qx], r1[1][qz][qy][qx], r1[2][qz][qy][qx]
+               };
+               const future::tensor<real_t, VDIM, DIM> grad_U = {{
+                     {g2[0][0][qz][qy][qx], g2[1][0][qz][qy][qx], g2[2][0][qz][qy][qx]},
+                     {g2[0][1][qz][qy][qx], g2[1][1][qz][qy][qx], g2[2][1][qz][qy][qx]},
+                     {g2[0][2][qz][qy][qx], g2[1][2][qz][qy][qx], g2[2][2][qz][qy][qx]}
+                  }
+               };
+               const auto two = transpose(grad_U) * (Q_adj * du_val);
+
+               // u⋅∇δu + δu⋅∇u
+               r0[0][qz][qy][qx] = one[0] + two[0];
+               r0[1][qz][qy][qx] = one[1] + two[1];
+               r0[2][qz][qy][qx] = one[2] + two[2];
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      kernels::internal::EvalTranspose3d(D1D, Q1D, smem, sB, r0, r1);
+      kernels::internal::WriteDofs3d(e, D1D, r1, Y);
+   });
+}
+
+void VectorConvectionNLFIntegrator::AddMultGradPA(const Vector &x,
+                                                  Vector &y) const
+{
+   if (dim == 2)
+   {
+      VectorConvectionNLFAddMultGradPA2D::Run(d1d, q1d, ne,
+                                              maps->B.Read(),
+                                              maps->G.Read(),
+                                              pa_adj.Read(),
+                                              pa_u.Read(),
+                                              x.Read(),
+                                              y.ReadWrite(),
+                                              d1d, q1d);
+   }
+   else if (dim == 3)
+   {
+      VectorConvectionNLFAddMultGradPA3D::Run(d1d, q1d, ne,
+                                              maps->B.Read(),
+                                              maps->G.Read(),
+                                              pa_adj.Read(),
+                                              pa_u.Read(),
+                                              x.Read(),
+                                              y.ReadWrite(),
+                                              d1d, q1d);
+   }
+   else
+   {
+      MFEM_ABORT("Unsupported dimension");
+   }
+}
+
+/// \cond DO_NOT_DOCUMENT
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPA2D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return SmemPAConvectionNLGradApply2D<T_D1D, T_Q1D>;
+}
+
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPA2D::Fallback
+(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= 16, "d1d > 16 is not supported");
+   MFEM_VERIFY(q1d <= 16, "q1d > 16 is not supported");
+   return SmemPAConvectionNLGradApply2D<>;
+}
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPA3D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return SmemPAConvectionNLGradApply3D<T_D1D, T_Q1D>;
+}
+
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPAType
+VectorConvectionNLFIntegrator::VectorConvectionNLFAddMultGradPA3D::Fallback
+(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= 16, "d1d > 16 is not supported");
+   MFEM_VERIFY(q1d <= 16, "q1d > 16 is not supported");
+   return SmemPAConvectionNLGradApply3D<>;
+}
+
+/// \endcond DO_NOT_DOCUMENT
+
+} // namespace mfem

fem/kernels.hpp

@@ -83,7 +83,7 @@ constexpr int SetMaxOf(int n) { return NextMultipleOf<4>(n); }
 #endif // CUDA/HIP && DEVICE_COMPILE
 
 /// Load 2D matrix into shared memory
-template <int MQ1>
+template <int MQ1, bool TRANSPOSE = false>
 inline MFEM_HOST_DEVICE void LoadMatrix(const int d1d, const int q1d,
                                         const real_t *M, real_t (*N)[MQ1])
 {
@@ -91,7 +91,14 @@ inline MFEM_HOST_DEVICE void LoadMatrix(const int d1d, const int q1d,
    {
       MFEM_FOREACH_THREAD_DIRECT(qx, x, q1d)
       {
-         N[dy][qx] = M[dy * q1d + qx];
+         if constexpr (TRANSPOSE)
+         {
+            N[dy][qx] = M[qx * d1d + dy];
+         }
+         else
+         {
+            N[dy][qx] = M[dy * q1d + qx];
+         }
       }
    }
    MFEM_SYNC_THREAD;

fem/nonlininteg.hpp

@@ -18,6 +18,7 @@
 #include "fespace.hpp"
 #include "ceed/interface/operator.hpp"
 #include "integrator.hpp"
+#include "kernel_dispatch.hpp"
 
 namespace mfem
 {
@@ -384,10 +385,10 @@ private:
    DenseMatrix dshape, dshapex, EF, gradEF, ELV, elmat_comp;
    Vector shape;
    // PA extension
-   Vector pa_data;
+   int dim, ne, nq, d1d, q1d;
+   Vector pa_adj, pa_u;
    const DofToQuad *maps;         ///< Not owned
    const GeometricFactors *geom;  ///< Not owned
-   int dim, ne, nq;
 
 public:
    VectorConvectionNLFIntegrator(Coefficient &q): Q(&q) { }
@@ -411,12 +412,51 @@ public:
 
    void AssemblePA(const FiniteElementSpace &fes) override;
 
-   void AssembleMF(const FiniteElementSpace &fes) override;
+   void AssembleGradPA(const Vector &x, const FiniteElementSpace &fes) override;
 
    void AddMultPA(const Vector &x, Vector &y) const override;
 
-   void AddMultMF(const Vector &x, Vector &y) const override;
+   using VectorConvectionNLFAddMultPAType =
+      void(*)(const int ne, const real_t *B, const real_t *G, const real_t *A,
+              const real_t *x, real_t *y,
+              const int d1d, const int q1d);
+   MFEM_REGISTER_KERNELS(VectorConvectionNLFAddMultPA,
+                         VectorConvectionNLFAddMultPAType,
+                         (int, int, int));
 
+   void AddMultGradPA(const Vector &x, Vector &y) const override;
+
+   using VectorConvectionNLFAddMultGradPAType =
+      void(*)(const int ne, const real_t *B, const real_t *G, const real_t *A,
+              const real_t *u, const real_t *x, real_t *y,
+              const int d1d, const int q1d);
+
+   MFEM_REGISTER_KERNELS(VectorConvectionNLFAddMultGradPA2D,
+                         VectorConvectionNLFAddMultGradPAType,
+                         (int, int));
+
+   MFEM_REGISTER_KERNELS(VectorConvectionNLFAddMultGradPA3D,
+                         VectorConvectionNLFAddMultGradPAType,
+                         (int, int));
+
+   void AssembleGradDiagonalPA(Vector &) const override;
+
+   using VectorConvectionNLFGradDiagPAType =
+      void (*)(const int ne, const real_t *B, const real_t *G, const real_t *A,
+               const real_t *u, real_t *y,
+               const int d1d, const int q1d);
+
+   MFEM_REGISTER_KERNELS(VectorConvectionNLFGradDiagPA2D,
+                         VectorConvectionNLFGradDiagPAType,
+                         (int, int));
+
+   MFEM_REGISTER_KERNELS(VectorConvectionNLFGradDiagPA3D,
+                         VectorConvectionNLFGradDiagPAType,
+                         (int, int));
+
+   void AssembleMF(const FiniteElementSpace &fes) override;
+
+   void AddMultMF(const Vector &x, Vector &y) const override;
 
 protected:
    const IntegrationRule* GetDefaultIntegrationRule(

linalg/operator.cpp

@@ -738,6 +738,13 @@ void ConstrainedOperator::AddMult(const Vector &x, Vector &y,
    y.Add(a, w);
 }
 
+void ConstrainedOperator::AddMultTranspose(const Vector &x, Vector &y,
+                                           const real_t a) const
+{
+   MultTranspose(x, w);
+   y.Add(a, w);
+}
+
 RectangularConstrainedOperator::RectangularConstrainedOperator(
    Operator *A,
    const Array<int> &trial_list,

linalg/operator.hpp

@@ -1107,6 +1107,9 @@ public:
 
    void MultTranspose(const Vector &x, Vector &y) const override;
 
+   void AddMultTranspose(const Vector &x, Vector &y,
+                         const real_t a = 1.0) const override;
+
    void AbsMultTranspose(const Vector &x, Vector &y) const override;
 
    /** @brief Implementation of Mult or MultTranspose.

tests/benchmarks/CMakeLists.txt

@@ -56,3 +56,4 @@ add_benchmark(elasticity)
 add_benchmark(tmop)
 add_benchmark(vector)
 add_benchmark(virtuals)
+add_benchmark(nlvc)

tests/benchmarks/bench_nlvc.cpp

@@ -0,0 +1,244 @@
+// 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 <cassert>
+#include <cstdlib>
+#include <functional>
+
+#include "fem/qinterp/grad.hpp"
+
+using namespace mfem;
+
+// Custom benchmark arguments generator ///////////////////////////////////////
+static void CustomArguments(bm::Benchmark *b) noexcept
+{
+   constexpr int MAX_NDOFS = 8 * 1024 * (mfem_use_gpu ? 1024 : 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 p : orders)
+   {
+      for (int n = (mfem_use_gpu ? 16 : 8); ndofs(n) <= MAX_NDOFS; n += inc(n))
+      {
+         b->Args({p, n});
+      }
+   }
+}
+
+/// Basic Kernels Specializations /////////////////////////////////////////////
+static void AddBasicKernelSpecializations()
+{
+   using Grad = QuadratureInterpolator::GradKernels;
+   // 2D
+   Grad::Specialization<2, QVectorLayout::byNODES, false, 2,2,7>::Add();
+   Grad::Specialization<2, QVectorLayout::byNODES, false, 2,2,8>::Add();
+   Grad::Specialization<2, QVectorLayout::byNODES, false, 2,2,10>::Add();
+   // 3D
+   Grad::Specialization<3, QVectorLayout::byNODES, false, 3,2,7>::Add();
+   Grad::Specialization<3, QVectorLayout::byNODES, false, 3,2,9>::Add();
+   Grad::Specialization<3, QVectorLayout::byNODES, false, 3,2,10>::Add();
+}
+
+/// VectorConvectionNLFBenchmark //////////////////////////////////////////////
+template <int DIM>
+struct VectorConvectionNLFBenchmark
+{
+   const int p, c, q, n, nx, ny, nz;
+   const std::function<Mesh()> MakeCartesianMesh = [&]()
+   {
+      if constexpr (DIM == 2)
+      {
+         return Mesh::MakeCartesian2D(nx, ny, Element::QUADRILATERAL);
+      }
+      else
+      {
+         return Mesh::MakeCartesian3D(nx, ny, nz, Element::HEXAHEDRON);
+      }
+   };
+   Mesh mesh;
+   H1_FECollection fec;
+   FiniteElementSpace fes;
+   const Geometry::Type geom_type;
+   IntegrationRules irs;
+   const IntegrationRule *ir;
+   ConstantCoefficient const_coeff { M_2_SQRTPI };
+   NonlinearFormIntegrator *nlfi;
+   NonlinearForm nlf;
+   Operator *grad;
+   GridFunction x, dx, y_pa;
+   Vector xe, dxe, ye;
+   const int dofs;
+   const int q1d;
+   double mdofs{};
+
+   VectorConvectionNLFBenchmark(int p, int side):
+      p(p), c(side), q(2 * p + 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),
+      mesh(MakeCartesianMesh()),
+      fec(p, DIM),
+      fes(&mesh, &fec, DIM),
+      geom_type(mesh.GetTypicalElementGeometry()),
+      irs(0, Quadrature1D::GaussLegendre),
+      ir(&irs.Get(geom_type, q)),
+      nlfi(new VectorConvectionNLFIntegrator(const_coeff)),
+      nlf(&fes),
+      x(&fes),
+      dx(&fes),
+      y_pa(&fes),
+      dofs(fes.GetTrueVSize()),
+      q1d(IntRules.Get(Geometry::SEGMENT, ir->GetOrder()).GetNPoints())
+   {
+      MFEM_VERIFY(q1d*q1d*(DIM == 3 ? q1d : 1) == ir->GetNPoints(), "");
+
+      nlf.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+      nlf.AddDomainIntegrator(nlfi);
+      nlf.Setup();
+
+      dx.Randomize(0x9e3779b9), x.Randomize(0x100001b3);
+
+      grad = &nlf.GetGradient(x);
+
+      const Table &el2dof = fes.GetElementToDofTable();
+      const int e_size = el2dof.Size_of_connections()*fes.GetVDim();
+      const auto R = fes.GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+      MFEM_VERIFY(e_size == R->Height(), "Input/Output E-vector size mismatch!");
+      xe.SetSize(R->Height()), dxe.SetSize(R->Height()), ye.SetSize(R->Height());
+      xe.UseDevice(true), dxe.UseDevice(true), ye.UseDevice(true);
+      xe.Randomize(0x100001b3), dxe.Randomize(0x9e3779b9), ye = 0.0;
+
+      mdofs = 0.0;
+   }
+
+   void Setup()
+   {
+      nlfi->AssembleGradPA(xe, fes);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs();
+   }
+
+   void AddMult()
+   {
+      nlf.AddMult(x, y_pa);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs();
+   }
+
+   void AddMultPA()
+   {
+      nlfi->AddMultPA(xe, ye);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs();
+   }
+
+   void AddMultGrad()
+   {
+      grad->Mult(dx, y_pa);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs();
+   }
+
+   void AddMultGradPA()
+   {
+      nlfi->AddMultGradPA(dxe, ye);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs();
+   }
+
+   void AssembleGradDiagonal()
+   {
+      grad->AssembleDiagonal(ye);
+      MFEM_DEVICE_SYNC;
+      mdofs += this->MDofs();
+   }
+
+   [[nodiscard]] double SumMdofs() const noexcept { return mdofs; }
+
+   [[nodiscard]] double MDofs() const noexcept { return 1e-6 * dofs; }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+#define RegisterVectorConvectionNLFBenchmark(Benchmark, DIM)         \
+   static void Benchmark##DIM##d(bm::State &state)                   \
+   {                                                                 \
+      const auto order = static_cast<int>(state.range(0));           \
+      const auto side = static_cast<int>(state.range(1));            \
+      VectorConvectionNLFBenchmark<DIM> ker(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);                      \
+   }                                                                 \
+   BENCHMARK(Benchmark##DIM##d)                                      \
+      ->Apply(CustomArguments)                                       \
+      ->Unit(bm::kMillisecond)
+
+RegisterVectorConvectionNLFBenchmark(Setup,3);
+RegisterVectorConvectionNLFBenchmark(AddMult,3);
+RegisterVectorConvectionNLFBenchmark(AddMultPA,3);
+RegisterVectorConvectionNLFBenchmark(AddMultGrad,3);
+RegisterVectorConvectionNLFBenchmark(AddMultGradPA,3);
+RegisterVectorConvectionNLFBenchmark(AssembleGradDiagonal,3);
+
+RegisterVectorConvectionNLFBenchmark(Setup,2);
+RegisterVectorConvectionNLFBenchmark(AddMult,2);
+RegisterVectorConvectionNLFBenchmark(AddMultPA,2);
+RegisterVectorConvectionNLFBenchmark(AddMultGrad,2);
+RegisterVectorConvectionNLFBenchmark(AddMultGradPA,2);
+RegisterVectorConvectionNLFBenchmark(AssembleGradDiagonal,2);
+
+/// main //////////////////////////////////////////////////////////////////////
+int main(int argc, char *argv[])
+{
+   AddBasicKernelSpecializations();
+
+   bm::ConsoleReporter CR;
+   bm::Initialize(&argc, argv);
+
+   // Device setup, cpu by default
+   std::string device_context = "cpu";
+   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;
+      }
+   }
+   Device device(device_context.c_str());
+   device.Print();
+
+   if (bm::ReportUnrecognizedArguments(argc, argv)) { return EXIT_FAILURE; }
+
+   bm::RunSpecifiedBenchmarks(&CR);
+
+   return EXIT_SUCCESS;
+}
+
+#endif // MFEM_USE_BENCHMARK

tests/unit/CMakeLists.txt

@@ -147,6 +147,8 @@ set(UNIT_TESTS_SRCS
   fem/test_pa_grad.cpp
   fem/test_pa_idinterp.cpp
   fem/test_pa_kernels.cpp
+  fem/test_pa_nlvc.cpp
+  fem/test_pa_vecdiv.cpp
   fem/test_particleset.cpp
   fem/test_pgridfunc_save_serial.cpp
   fem/test_poly1d.cpp

tests/unit/fem/test_pa_nlvc.cpp

@@ -0,0 +1,132 @@
+// 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.
+
+#ifdef _WIN32
+#define _USE_MATH_DEFINES
+#include <cmath>
+#endif
+
+#include "unit_tests.hpp"
+#include "mfem.hpp"
+
+#include "fem/qinterp/grad.hpp" // IWYU pragma: keep
+
+using namespace mfem;
+
+namespace pa_kernels
+{
+
+template<int DIM>
+void test_nl_convection_pa_grad(const char *filename, int p)
+{
+   CAPTURE(filename, DIM, p);
+
+   Mesh mesh(filename);
+   MFEM_VERIFY(mesh.Dimension() == DIM, "Mesh dimension mismatch");
+
+   H1_FECollection fec(p, DIM);
+   FiniteElementSpace fes(&mesh, &fec, DIM);
+
+   GridFunction x(&fes), dx(&fes), y_fa(&fes), y_pa(&fes);
+   x.Randomize(0x100001b3);
+   dx.Randomize(0x9e3779b9);
+
+   ConstantCoefficient const_coeff(M_2_SQRTPI);
+   FunctionCoefficient funct_coeff([](const Vector &x)
+   { return M_1_PI + x[0] * x[0]; });
+
+   NonlinearForm nlf_fa(&fes);
+   nlf_fa.AddDomainIntegrator(new VectorConvectionNLFIntegrator);
+   nlf_fa.AddDomainIntegrator(new VectorConvectionNLFIntegrator(const_coeff));
+   nlf_fa.AddDomainIntegrator(new VectorConvectionNLFIntegrator(funct_coeff));
+
+   NonlinearForm nlf_pa(&fes);
+   nlf_pa.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   nlf_pa.AddDomainIntegrator(new VectorConvectionNLFIntegrator);
+   nlf_pa.AddDomainIntegrator(new VectorConvectionNLFIntegrator(const_coeff));
+   nlf_pa.AddDomainIntegrator(new VectorConvectionNLFIntegrator(funct_coeff));
+   nlf_pa.Setup();
+
+   SECTION("Action")
+   {
+      nlf_fa.Mult(x, y_fa), nlf_pa.Mult(x, y_pa);
+      y_fa -= y_pa;
+      REQUIRE(y_fa.Norml2() == MFEM_Approx(0.0));
+   }
+
+   SECTION("Gradient")
+   {
+      Operator &nlf_fa_grad = nlf_fa.GetGradient(x);
+      Operator &nlf_pa_grad = nlf_pa.GetGradient(x);
+      nlf_pa_grad.Mult(dx, y_pa);
+      nlf_fa_grad.Mult(dx, y_fa);
+      y_fa -= y_pa;
+      REQUIRE(y_fa.Norml2() == MFEM_Approx(0.0));
+   }
+
+   SECTION("Diagonal")
+   {
+      Vector diag_fa(fes.GetVSize()), diag_pa(fes.GetVSize());
+      dynamic_cast<SparseMatrix &>(nlf_fa.GetGradient(x)).GetDiag(diag_fa);
+      nlf_pa.GetGradient(x).AssembleDiagonal(diag_pa);
+      diag_fa -= diag_pa;
+      REQUIRE(diag_fa.Norml2() == MFEM_Approx(0.0));
+   }
+}
+
+TEST_CASE("NL Convection PA Gradient",
+          "[PartialAssembly][NonlinearPA][GPU][NLConv]")
+{
+   const bool all_tests = launch_all_non_regression_tests;
+   const auto p = !all_tests ? GENERATE(1, 2) : GENERATE(1, 2, 3, 4);
+
+   if (static auto done = false; !std::exchange(done, true))
+   {
+      using Grad = QuadratureInterpolator::GradKernels;
+      Grad::Specialization<2, QVectorLayout::byNODES, false, 2, 2, 7>::Add();
+      Grad::Specialization<2, QVectorLayout::byNODES, false, 2, 3, 7>::Add();
+      Grad::Specialization<2, QVectorLayout::byNODES, false, 2, 4, 8>::Add();
+      Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 7>::Add();
+      Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 3, 7>::Add();
+      Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 3, 8>::Add();
+      Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 4, 9>::Add();
+   }
+
+   SECTION("2D")
+   {
+      const auto filename2d = all_tests
+                              ? GENERATE("../../data/star-q2.mesh",
+                                         "../../data/star-q3.mesh",
+                                         "../../data/rt-2d-q3.mesh",
+                                         "../../data/inline-quad.mesh",
+                                         "../../data/periodic-square.mesh")
+                              : GENERATE("../../data/inline-quad.mesh",
+                                         "../../data/periodic-square.mesh");
+      test_nl_convection_pa_grad<2>(filename2d, p);
+   }
+
+   SECTION("3D")
+   {
+      const auto filename3d = all_tests
+                              ? GENERATE("../../data/beam-hex.mesh",
+                                         "../../data/fichera.mesh",
+                                         "../../data/fichera-q2.mesh",
+                                         "../../data/fichera-q3.mesh",
+                                         "../../data/inline-hex.mesh",
+                                         "../../data/periodic-cube.mesh",
+                                         "../../data/toroid-hex.mesh")
+                              : GENERATE("../../data/inline-hex.mesh",
+                                         "../../data/periodic-cube.mesh");
+      test_nl_convection_pa_grad<3>(filename3d, p);
+   }
+}
+
+} // namespace pa_kernels

tests/unit/fem/test_pa_vecdiv.cpp

@@ -0,0 +1,133 @@
+// 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.
+
+#ifdef _WIN32
+#define _USE_MATH_DEFINES
+#include <cmath>
+#endif
+
+#include "unit_tests.hpp"
+#include "mfem.hpp"
+
+#include "fem/qinterp/grad.hpp" // IWYU pragma: keep
+
+using namespace mfem;
+
+namespace pa_kernels
+{
+
+///////////////////////////////////////////////////////////////////////////////
+template <typename INTEGRATOR, bool TRANSPOSE>
+void pa_mixed_test(FiniteElementSpace &fes1,
+                   FiniteElementSpace &fes2)
+{
+   MixedBilinearForm bform_pa(&fes1, &fes2);
+   if constexpr (TRANSPOSE)
+   {
+      bform_pa.AddDomainIntegrator(new TransposeIntegrator(new INTEGRATOR));
+   }
+   else
+   {
+      bform_pa.AddDomainIntegrator(new INTEGRATOR);
+   }
+   bform_pa.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   bform_pa.Assemble();
+
+   MixedBilinearForm bform_fa(&fes1, &fes2);
+   if constexpr (TRANSPOSE)
+   {
+      bform_fa.AddDomainIntegrator(new TransposeIntegrator(new INTEGRATOR));
+   }
+   else
+   {
+      bform_fa.AddDomainIntegrator(new INTEGRATOR);
+   }
+   bform_fa.Assemble();
+   bform_fa.Finalize();
+
+   GridFunction x(&fes1), y_pa(&fes2), y_fa(&fes2);
+   x.Randomize(0x100001b3);
+
+   bform_pa.Mult(x, y_pa);
+   bform_fa.Mult(x, y_fa);
+
+   y_pa -= y_fa;
+   REQUIRE(y_pa.Normlinf() == MFEM_Approx(0.0));
+}
+
+///////////////////////////////////////////////////////////////////////////////
+template<int DIM>
+void test_pa_divergence(const char *filename, int p)
+{
+   CAPTURE(filename, DIM, p);
+   dbg("{} {} {}", filename, DIM, p);
+
+   Mesh mesh(filename);
+   MFEM_VERIFY(mesh.Dimension() == DIM, "Mesh dimension mismatch");
+
+   // Vector
+   H1_FECollection vfec(p, DIM);
+   FiniteElementSpace vfes(&mesh, &vfec, DIM);
+
+   // Scalar
+   H1_FECollection sfec(p, DIM);
+   FiniteElementSpace sfes(&mesh, &sfec);
+
+   pa_mixed_test<VectorDivergenceIntegrator, false>(vfes, sfes);
+   pa_mixed_test<VectorDivergenceIntegrator, true>(sfes, vfes);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+TEST_CASE("VecDivPA", "[PartialAssembly][VecDivPA][GPU]")
+{
+   const bool all_tests = launch_all_non_regression_tests;
+   const auto p = !all_tests ? GENERATE(1, 2) : GENERATE(1, 2, 3, 4);
+
+   if (static auto done = false; !std::exchange(done, true))
+   {
+      using Grad = QuadratureInterpolator::GradKernels;
+      Grad::Specialization<2, QVectorLayout::byNODES, false, 2, 3, 5>::Add();
+      Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 3, 7>::Add();
+      Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 4, 9>::Add();
+   }
+
+   SECTION("2D")
+   {
+      const auto filename2d =
+         all_tests
+         ? GENERATE("../../data/star-q2.mesh",
+                    "../../data/star-q3.mesh",
+                    "../../data/rt-2d-q3.mesh",
+                    "../../data/inline-quad.mesh",
+                    "../../data/periodic-square.mesh")
+         : GENERATE("../../data/inline-quad.mesh",
+                    "../../data/periodic-square.mesh");
+      test_pa_divergence<2>(filename2d, p);
+   }
+
+   SECTION("3D")
+   {
+      const auto filename3d =
+         all_tests
+         ? GENERATE("../../data/beam-hex.mesh",
+                    "../../data/fichera.mesh",
+                    "../../data/fichera-q2.mesh",
+                    "../../data/fichera-q3.mesh",
+                    "../../data/inline-hex.mesh",
+                    "../../data/periodic-cube.mesh",
+                    "../../data/toroid-hex.mesh")
+         : GENERATE("../../data/inline-hex.mesh",
+                    "../../data/periodic-cube.mesh");
+      test_pa_divergence<3>(filename3d, p);
+   }
+}
+
+} // namespace pa_kernels