Merge branch 'nlconvpa' into FEBioMFEMiFSI-gpu

Fix header file kernels so users can instantiate their own specializations

CHANGELOG

@@ -70,6 +70,15 @@ Linear and nonlinear solvers
 
 GPU computing
 -------------
+- Added PA gradient and diagonal support for VectorConvectionNLFIntegrator
+ (AssembleGradPA, AddMultGradPA, AssembleGradDiagonalPA).
+
+- Improved partial assembly for VectorDivergenceIntegrator with shared-memory
+  kernels, kernel registration, and transpose support.
+
+- Improved partial-assembly diagonal kernels for VectorMassIntegrator (shared-
+  memory specializations) and ElasticityIntegrator (no scratch Q-vector).
+
 - Added device assembly support for 3D H(curl) VectorFEDomainLFIntegrator.
 
 - Added NVIDIA cuDSS library interface. Implementation examples have been

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
@@ -204,7 +206,11 @@ set(HDRS
   integ/bilininteg_mass_kernels.hpp
   integ/bilininteg_mass_pa_simplices.hpp
   integ/bilininteg_vecdiffusion_pa.hpp
+  integ/bilininteg_vecdiv_pa.hpp
   integ/bilininteg_vecmass_pa.hpp
+  integ/nonlininteg_vecconvection_pa.hpp
+  integ/nonlininteg_vecconvection_pa_diag.hpp
+  integ/nonlininteg_vecconvection_pa_grad.hpp
   coefficient.hpp
   complex_fem.hpp
   convergence.hpp

fem/bilininteg.hpp

@@ -2689,14 +2689,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*/));
 };
 
 
@@ -3098,6 +3106,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,67 @@ 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++)
+               const auto invJ = inv(make_tensor<d, d>([&](int r, int c)
                {
-                  for (int p = 0; p < numPoints; p++ )
+                  return J(p, r, c, e);
+               }));
+               const real_t w = ipWeights[p] / det(invJ);
+
+               for (int n = 0; n < d; n++)
+               {
+                  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.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_vecdiv_pa.hpp

@@ -0,0 +1,365 @@
+// 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/array.hpp"
+#include "../../general/forall.hpp"
+#include "../../linalg/dtensor.hpp"
+#include "../../linalg/vector.hpp"
+#include "../bilininteg.hpp"
+#include "../kernels.hpp"
+
+namespace mfem
+{
+
+/// \cond DO_NOT_DOCUMENT
+
+namespace internal
+{
+
+// Shared memory PA Divergence Apply 2D kernel
+template<int T_TR_D1D = 0, int T_TE_D1D = 0, int T_Q1D = 0>
+inline void SmemPADivergenceApply2D(const int NE,
+                                    const Array<real_t> &b_,
+                                    const Array<real_t> &g_,
+                                    const Array<real_t> &bt_,
+                                    const Vector &q_,
+                                    const Vector &x_,
+                                    Vector &y_,
+                                    const int tr_d1d = 0,
+                                    const int te_d1d = 0,
+                                    const int q1d = 0)
+{
+   const int TR_D1D = T_TR_D1D ? T_TR_D1D : tr_d1d;
+   const int TE_D1D = T_TE_D1D ? T_TE_D1D : te_d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   MFEM_VERIFY(TR_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(TE_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+
+   const auto B = b_.Read(), G = g_.Read(), Bt = bt_.Read();
+   const auto Q = Reshape(q_.Read(), Q1D, Q1D, 2, 2, NE);
+   const auto X = Reshape(x_.Read(), TR_D1D, TR_D1D, 2, NE);
+   auto Y = Reshape(y_.ReadWrite(), TE_D1D, TE_D1D, 1, NE);
+
+   mfem::forall_2D<T_Q1D * T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   {
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1];
+      MFEM_SHARED real_t sB[MQ1][MQ1], sG[MQ1][MQ1];
+
+      kernels::internal::vd_regs2d_t<2, 2, MQ1> g0, g1;
+      kernels::internal::v_regs2d_t<1, MQ1> r0, r1;
+
+      kernels::internal::LoadMatrix(TR_D1D, Q1D, B, sB);
+      kernels::internal::LoadMatrix(TR_D1D, Q1D, G, sG);
+
+      kernels::internal::LoadDofs2d(e, TR_D1D, X, g0);
+      kernels::internal::Grad2d(TR_D1D, Q1D, smem, sB, sG, g0, g1);
+
+      MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            r0[0][qy][qx] =
+               g1[0][0][qy][qx] * Q(qx, qy, 0, 0, e) +
+               g1[0][1][qy][qx] * Q(qx, qy, 1, 0, e) +
+               g1[1][0][qy][qx] * Q(qx, qy, 0, 1, e) +
+               g1[1][1][qy][qx] * Q(qx, qy, 1, 1, e);
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      kernels::internal::LoadMatrix<MQ1,true>(TE_D1D, Q1D, Bt, sB);
+      kernels::internal::EvalTranspose2d(TE_D1D, Q1D, smem, sB, r0, r1);
+      kernels::internal::WriteDofs2d(e, TE_D1D, r1, Y);
+   });
+}
+
+// Shared memory PA Divergence Apply 2D kernel transpose
+template<int T_TR_D1D = 0, int T_TE_D1D = 0, int T_Q1D = 0>
+inline void SmemPADivergenceApplyTranspose2D(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 = 0,
+                                             const int te_d1d = 0,
+                                             const int q1d = 0)
+{
+   const int TR_D1D = T_TR_D1D ? T_TR_D1D : tr_d1d;
+   const int TE_D1D = T_TE_D1D ? T_TE_D1D : te_d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   MFEM_VERIFY(TR_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(TE_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+
+   const auto Bt = bt.Read(), Gt = gt.Read(), B = b.Read();
+   const auto Q = Reshape(q_.Read(), Q1D, Q1D, 2, 2, NE);
+   const auto X  = Reshape(x_.Read(), TE_D1D, TE_D1D, 1, NE);
+   auto Y  = Reshape(y_.ReadWrite(), TR_D1D, TR_D1D, 2, NE);
+
+   mfem::forall_2D<T_Q1D * T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   {
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1];
+      MFEM_SHARED real_t sB[MQ1][MQ1], sG[MQ1][MQ1];
+
+      kernels::internal::v_regs2d_t<1, MQ1> r0, r1;
+      kernels::internal::vd_regs2d_t<2, 2, MQ1> g0, g1;
+
+      kernels::internal::LoadMatrix(TE_D1D, Q1D, B, sB);
+      kernels::internal::LoadDofs2d(e, TE_D1D, X, r0);
+      kernels::internal::Eval2d(TE_D1D, Q1D, smem, sB, r0, r1);
+
+      MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            g0[0][0][qy][qx] = r1[0][qy][qx] * Q(qx, qy, 0, 0, e);
+            g0[0][1][qy][qx] = r1[0][qy][qx] * Q(qx, qy, 1, 0, e);
+            g0[1][0][qy][qx] = r1[0][qy][qx] * Q(qx, qy, 0, 1, e);
+            g0[1][1][qy][qx] = r1[0][qy][qx] * Q(qx, qy, 1, 1, e);
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      kernels::internal::LoadMatrix<MQ1,true>(TR_D1D, Q1D, Bt, sB);
+      kernels::internal::LoadMatrix<MQ1,true>(TR_D1D, Q1D, Gt, sG);
+      kernels::internal::GradTranspose2d(TR_D1D, Q1D, smem, sB, sG, g0, g1);
+      kernels::internal::WriteDofs2d(e, TR_D1D, g1, Y);
+   });
+}
+
+// Shared memory PA Divergence Apply 3D kernel transpose
+template<int T_TR_D1D = 0, int T_TE_D1D = 0, int T_Q1D = 0>
+inline void SmemPADivergenceApplyTranspose3D(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_,
+                                             int tr_d1d = 0,
+                                             int te_d1d = 0,
+                                             int q1d = 0)
+{
+   const int TR_D1D = T_TR_D1D ? T_TR_D1D : tr_d1d;
+   const int TE_D1D = T_TE_D1D ? T_TE_D1D : te_d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   MFEM_VERIFY(TR_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(TE_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+
+   const auto Bt = bt.Read(), Gt = gt.Read(), B = b.Read();
+   const auto Q = Reshape(q_.Read(), Q1D, Q1D, Q1D, 3, 3, NE);
+   const auto X = Reshape(x_.Read(), TE_D1D, TE_D1D, TE_D1D, 1, NE);
+   auto Y = Reshape(y_.ReadWrite(), TR_D1D, TR_D1D, TR_D1D, 3, NE);
+
+   mfem::forall_2D<T_Q1D * T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   {
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1];
+      MFEM_SHARED real_t sB[MQ1][MQ1], sG[MQ1][MQ1];
+
+      kernels::internal::v_regs3d_t<1, MQ1> r0, r1;
+      kernels::internal::vd_regs3d_t<3, 3, MQ1> g0, g1;
+
+      kernels::internal::LoadMatrix(TE_D1D, Q1D, B, sB);
+      kernels::internal::LoadDofs3d(e, TE_D1D, X, r0);
+      kernels::internal::Eval3d(TE_D1D, Q1D, smem, sB, r0, r1);
+
+      for (int qz = 0; qz < Q1D; qz++)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+            {
+               const auto r = r1[0][qz][qy][qx];
+               g0[0][0][qz][qy][qx] = r * Q(qx, qy, qz, 0, 0, e);
+               g0[0][1][qz][qy][qx] = r * Q(qx, qy, qz, 1, 0, e);
+               g0[0][2][qz][qy][qx] = r * Q(qx, qy, qz, 2, 0, e);
+
+               g0[1][0][qz][qy][qx] = r * Q(qx, qy, qz, 0, 1, e);
+               g0[1][1][qz][qy][qx] = r * Q(qx, qy, qz, 1, 1, e);
+               g0[1][2][qz][qy][qx] = r * Q(qx, qy, qz, 2, 1, e);
+
+               g0[2][0][qz][qy][qx] = r * Q(qx, qy, qz, 0, 2, e);
+               g0[2][1][qz][qy][qx] = r * Q(qx, qy, qz, 1, 2, e);
+               g0[2][2][qz][qy][qx] = r * Q(qx, qy, qz, 2, 2, e);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      kernels::internal::LoadMatrix<MQ1,true>(TR_D1D, Q1D, Bt, sB);
+      kernels::internal::LoadMatrix<MQ1,true>(TR_D1D, Q1D, Gt, sG);
+      kernels::internal::GradTranspose3d(TR_D1D, Q1D, smem, sB, sG, g0, g1);
+      kernels::internal::WriteDofs3d(e, TR_D1D, g1, Y);
+   });
+}
+
+// Shared memory PA Divergence Apply 3D kernel
+template<int T_TR_D1D = 0, int T_TE_D1D = 0, int T_Q1D = 0>
+inline void SmemPADivergenceApply3D(const int NE,
+                                    const Array<real_t> &b_,
+                                    const Array<real_t> &g_,
+                                    const Array<real_t> &bt_,
+                                    const Vector &q_,
+                                    const Vector &x_,
+                                    Vector &y_,
+                                    const int tr_d1d = 0,
+                                    const int te_d1d = 0,
+                                    const int q1d = 0)
+{
+   const int TR_D1D = T_TR_D1D ? T_TR_D1D : tr_d1d;
+   const int TE_D1D = T_TE_D1D ? T_TE_D1D : te_d1d;
+   const int Q1D = T_Q1D ? T_Q1D : q1d;
+
+   MFEM_VERIFY(TR_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(TE_D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+
+   const auto B = b_.Read(), G = g_.Read(), Bt = bt_.Read();
+   const auto Q = Reshape(q_.Read(), Q1D, Q1D, Q1D, 3,3, NE);
+   const auto X = Reshape(x_.Read(), TR_D1D, TR_D1D, TR_D1D, 3, NE);
+   auto Y = Reshape(y_.ReadWrite(), TE_D1D, TE_D1D, TE_D1D, 1, NE);
+
+   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   {
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1];
+      MFEM_SHARED real_t sB[MQ1][MQ1], sG[MQ1][MQ1];
+
+      kernels::internal::vd_regs3d_t<3, 3, MQ1> g0, g1;
+      kernels::internal::v_regs3d_t<1, MQ1> r0, r1;
+
+      kernels::internal::LoadMatrix(TR_D1D, Q1D, B, sB);
+      kernels::internal::LoadMatrix(TR_D1D, Q1D, G, sG);
+
+      kernels::internal::LoadDofs3d(e, TR_D1D, X, g0);
+      kernels::internal::Grad3d(TR_D1D, Q1D, smem, sB, sG, g0, g1);
+
+      for (int qz = 0; qz < Q1D; qz++)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+            {
+               r0[0][qz][qy][qx] =
+                  // c = 0
+                  g1[0][0][qz][qy][qx] * Q(qx, qy, qz, 0, 0, e) +
+                  g1[0][1][qz][qy][qx] * Q(qx, qy, qz, 1, 0, e) +
+                  g1[0][2][qz][qy][qx] * Q(qx, qy, qz, 2, 0, e) +
+                  // c = 1
+                  g1[1][0][qz][qy][qx] * Q(qx, qy, qz, 0, 1, e) +
+                  g1[1][1][qz][qy][qx] * Q(qx, qy, qz, 1, 1, e) +
+                  g1[1][2][qz][qy][qx] * Q(qx, qy, qz, 2, 1, e) +
+                  // c = 2
+                  g1[2][0][qz][qy][qx] * Q(qx, qy, qz, 0, 2, e) +
+                  g1[2][1][qz][qy][qx] * Q(qx, qy, qz, 1, 2, e) +
+                  g1[2][2][qz][qy][qx] * Q(qx, qy, qz, 2, 2, e);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      kernels::internal::LoadMatrix<MQ1, true>(TE_D1D, Q1D, Bt, sB);
+      kernels::internal::EvalTranspose3d(TE_D1D, Q1D, smem, sB, r0, r1);
+      kernels::internal::WriteDofs3d(e, TE_D1D, r1, Y);
+   });
+}
+
+} // namespace internal
+
+template<int DIM, int T_TR_D1D, int T_TE_D1D, int T_Q1D>
+VectorDivergenceIntegrator::VectorDivergenceAddMultPAType
+VectorDivergenceIntegrator::VectorDivergenceAddMultPA::Kernel()
+{
+   static_assert(T_TR_D1D <= T_Q1D && T_TE_D1D <= T_Q1D);
+   if constexpr (DIM == 2)
+   {
+      return internal::SmemPADivergenceApply2D<T_TR_D1D, T_TE_D1D, T_Q1D>;
+   }
+   else if constexpr (DIM == 3)
+   {
+      return internal::SmemPADivergenceApply3D<T_TR_D1D, T_TE_D1D, T_Q1D>;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+inline VectorDivergenceIntegrator::VectorDivergenceAddMultPAType
+VectorDivergenceIntegrator::VectorDivergenceAddMultPA::Fallback
+(int dim, int tr_d1d, int te_d1d, int q1d)
+{
+   MFEM_VERIFY(tr_d1d <= q1d && te_d1d <= q1d, "");
+   MFEM_VERIFY(tr_d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(te_d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   if (dim == 2)
+   {
+      return internal::SmemPADivergenceApply2D;
+   }
+   else if (dim == 3)
+   {
+      return internal::SmemPADivergenceApply3D;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+template<int DIM, int T_TR_D1D, int T_TE_D1D, int T_Q1D>
+VectorDivergenceIntegrator::VectorDivergenceAddMultTransposePAType
+VectorDivergenceIntegrator::VectorDivergenceAddMultTransposePA::Kernel()
+{
+   static_assert(T_TR_D1D <= T_Q1D && T_TE_D1D <= T_Q1D);
+   if constexpr (DIM == 2)
+   {
+      return internal::SmemPADivergenceApplyTranspose2D<T_TR_D1D, T_TE_D1D, T_Q1D>;
+   }
+   else if constexpr (DIM == 3)
+   {
+      return internal::SmemPADivergenceApplyTranspose3D<T_TR_D1D, T_TE_D1D, T_Q1D>;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+inline VectorDivergenceIntegrator::VectorDivergenceAddMultTransposePAType
+VectorDivergenceIntegrator::VectorDivergenceAddMultTransposePA::Fallback
+(int dim, int tr_d1d, int te_d1d, int q1d)
+{
+   MFEM_VERIFY(tr_d1d <= q1d && te_d1d <= q1d, "");
+   MFEM_VERIFY(tr_d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(te_d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   if (dim == 2)
+   {
+      return internal::SmemPADivergenceApplyTranspose2D;
+   }
+   else if (dim == 3)
+   {
+      return internal::SmemPADivergenceApplyTranspose3D;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+/// \endcond DO_NOT_DOCUMENT
+
+} // namespace mfem

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,146 @@ 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;
+         }
+      }
+   });
+}
+
+// T_MDQ <= 10 so the Q1D^3 thread block stays within the 1024/block GPU limit
+template <int T_Q1D = 0, int T_MDQ = 10>
+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 +328,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 +345,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.cpp

@@ -9,21 +9,51 @@
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
 
-#include "../../general/forall.hpp"
-#include "../nonlininteg.hpp"
 #include "../ceed/integrators/nlconvection/nlconvection.hpp"
+#include "./nonlininteg_vecconvection_pa.hpp" // IWYU pragma: keep
+#include "./nonlininteg_vecconvection_pa_grad.hpp" // IWYU pragma: keep
+#include "./nonlininteg_vecconvection_pa_diag.hpp" // IWYU pragma: keep
 
 namespace mfem
 {
 
+VectorConvectionNLFIntegrator::Kernels::Kernels()
+{
+   // 2D
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 2, 2>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 2, 3>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 3, 4>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 3, 5>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 4, 5>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 4, 6>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 5, 7>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 5, 8>();
+   VectorConvectionNLFIntegrator::AddSpecialization<2, 6, 8>();
+   // 3D
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 2, 3>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 2, 4>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 2, 5>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 3, 4>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 3, 5>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 3, 6>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 4, 5>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 4, 6>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 4, 7>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 4, 8>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 5, 6>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 5, 7>();
+   VectorConvectionNLFIntegrator::AddSpecialization<3, 5, 8>();
+}
+
 void VectorConvectionNLFIntegrator::AssemblePA(const FiniteElementSpace &fes)
 {
    MFEM_ASSERT(fes.GetOrdering() == Ordering::byNODES,
                "PA Only supports Ordering::byNODES!");
    Mesh *mesh = fes.GetMesh();
    const FiniteElement &el = *fes.GetTypicalFE();
-   ElementTransformation &T = *mesh->GetTypicalElementTransformation();
-   const IntegrationRule *ir = IntRule ? IntRule : &GetRule(el, T);
+   ElementTransformation &Tr = *mesh->GetTypicalElementTransformation();
+   const IntegrationRule *ir = IntRule ? IntRule : &GetRule(el, Tr);
+
    if (DeviceCanUseCeed())
    {
       delete ceedOp;
@@ -39,769 +69,124 @@ void VectorConvectionNLFIntegrator::AssemblePA(const FiniteElementSpace &fes)
       }
       return;
    }
-   dim = mesh->Dimension();
-   ne = fes.GetMesh()->GetNE();
+
+   ne = mesh->GetNE();
    nq = ir->GetNPoints();
-   geom = mesh->GetGeometricFactors(*ir, GeometricFactors::JACOBIANS);
+   dim = mesh->Dimension();
+   MFEM_VERIFY(dim == 2 || dim == 3, "Dimension not supported");
+
+   const MemoryType mt = pa_mt == MemoryType::DEFAULT
+                         ? Device::GetDeviceMemoryType()
+                         : pa_mt;
+   pa_adj.SetSize(ne * nq * dim * dim, mt);
+   geom = mesh->GetGeometricFactors(*ir, GeometricFactors::JACOBIANS, mt);
    maps = &el.GetDofToQuad(*ir, DofToQuad::TENSOR);
-   pa_data.SetSize(ne * nq * dim * dim, Device::GetMemoryType());
-   real_t COEFF = 1.0;
-   if (Q)
-   {
-      ConstantCoefficient *cQ = dynamic_cast<ConstantCoefficient *>(Q);
-      MFEM_VERIFY(cQ != NULL, "only ConstantCoefficient is supported!");
-      COEFF = cQ->constant;
-   }
-   const int NE = ne;
-   const int NQ = nq;
-   auto W = ir->GetWeights().Read();
-   if (dim == 1)
-   {
-      MFEM_ABORT("dim==1 not supported!");
-   }
+   d1d = maps->ndof;
+   q1d = maps->nqpt;
+
+   QuadratureSpace qs(*mesh, *ir);
+   CoefficientVector coeff(Q, qs, CoefficientStorage::COMPRESSED);
+
+   const int nq1d = q1d * q1d * (dim==3 ? q1d : 1);
+   MFEM_VERIFY(coeff.Size() == 1 || coeff.Size() == nq1d*ne, "Invalid coeff");
+   MFEM_VERIFY(ir->GetWeights().Size() == nq1d, "Invalid weights size");
+
+   const auto w_r = ir->GetWeights().Read();
+   const bool const_coeff = coeff.Size() == 1;
+
    if (dim == 2)
    {
-      auto J = Reshape(geom->J.Read(), NQ, 2, 2, NE);
-      auto G = Reshape(pa_data.Write(), NQ, 2, 2, NE);
-      mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+      const int Q1D = q1d;
+      constexpr int VDIM = 2, DIM = 2;
+      const auto W = Reshape(w_r, Q1D, Q1D);
+      const auto C = const_coeff ?
+                     Reshape(coeff.Read(), 1, 1, 1) :
+                     Reshape(coeff.Read(), Q1D, Q1D, ne);
+      const auto J = Reshape(geom->J.Read(), Q1D, Q1D, VDIM, DIM, ne);
+      auto A = Reshape(pa_adj.Write(), VDIM, DIM, Q1D, Q1D, ne);
+
+      mfem::forall_2D(ne, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
       {
-         for (int q = 0; q < NQ; ++q)
+         MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
          {
-            const real_t J11 = J(q, 0, 0, e);
-            const real_t J12 = J(q, 0, 1, e);
-            const real_t J21 = J(q, 1, 0, e);
-            const real_t J22 = J(q, 1, 1, e);
-            // Store wq * Q * adj(J)
-            G(q, 0, 0, e) = W[q] * COEFF * J22;  // 1,1
-            G(q, 0, 1, e) = W[q] * COEFF * -J12; // 1,2
-            G(q, 1, 0, e) = W[q] * COEFF * -J21; // 2,1
-            G(q, 1, 1, e) = W[q] * COEFF * J11;  // 2,2
+            MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+            {
+               const real_t J11 = J(qx, qy, 0, 0, e), J12 = J(qx, qy, 0, 1, e);
+               const real_t J21 = J(qx, qy, 1, 0, e), J22 = J(qx, qy, 1, 1, e);
+               // adj(J)
+               const real_t A11 = +J22, A12 = -J12;
+               const real_t A21 = -J21, A22 = +J11;
+               // Store w * coeff * adj(J)
+               const real_t w = W(qx, qy);
+               const real_t c = const_coeff ? C(0, 0, 0) : C(qx, qy, e);
+               A(0, 0, qx, qy, e) = w * c * A11;
+               A(1, 0, qx, qy, e) = w * c * A12;
+               A(0, 1, qx, qy, e) = w * c * A21;
+               A(1, 1, qx, qy, e) = w * c * A22;
+            }
          }
       });
    }
-   if (dim == 3)
+   else if (dim == 3)
    {
-      auto J = Reshape(geom->J.Read(), NQ, 3, 3, NE);
-      auto G = Reshape(pa_data.Write(), NQ, 3, 3, NE);
-      mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+      const int Q1D = q1d;
+      constexpr int VDIM = 3, DIM = 3;
+      const auto W = Reshape(w_r, Q1D, Q1D, Q1D);
+      const auto C = const_coeff ?
+                     Reshape(coeff.Read(), 1, 1, 1, 1) :
+                     Reshape(coeff.Read(), Q1D, Q1D, Q1D, ne);
+      const auto J = Reshape(geom->J.Read(), Q1D, Q1D, Q1D, VDIM, DIM, ne);
+      auto A = Reshape(pa_adj.Write(), VDIM, DIM, Q1D, Q1D, Q1D, ne);
+
+      mfem::forall_3D(ne, Q1D, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
       {
-         for (int q = 0; q < NQ; ++q)
+         MFEM_FOREACH_THREAD_DIRECT(qz, z, Q1D)
          {
-            const real_t J11 = J(q, 0, 0, e);
-            const real_t J21 = J(q, 1, 0, e);
-            const real_t J31 = J(q, 2, 0, e);
-            const real_t J12 = J(q, 0, 1, e);
-            const real_t J22 = J(q, 1, 1, e);
-            const real_t J32 = J(q, 2, 1, e);
-            const real_t J13 = J(q, 0, 2, e);
-            const real_t J23 = J(q, 1, 2, e);
-            const real_t J33 = J(q, 2, 2, e);
-            const real_t cw = W[q] * COEFF;
-            // adj(J)
-            const real_t A11 = (J22 * J33) - (J23 * J32);
-            const real_t A12 = (J32 * J13) - (J12 * J33);
-            const real_t A13 = (J12 * J23) - (J22 * J13);
-            const real_t A21 = (J31 * J23) - (J21 * J33);
-            const real_t A22 = (J11 * J33) - (J13 * J31);
-            const real_t A23 = (J21 * J13) - (J11 * J23);
-            const real_t A31 = (J21 * J32) - (J31 * J22);
-            const real_t A32 = (J31 * J12) - (J11 * J32);
-            const real_t A33 = (J11 * J22) - (J12 * J21);
-            // Store wq * Q * adj(J)
-            G(q, 0, 0, e) = cw * A11; // 1,1
-            G(q, 0, 1, e) = cw * A12; // 1,2
-            G(q, 0, 2, e) = cw * A13; // 1,3
-            G(q, 1, 0, e) = cw * A21; // 2,1
-            G(q, 1, 1, e) = cw * A22; // 2,2
-            G(q, 1, 2, e) = cw * A23; // 2,3
-            G(q, 2, 0, e) = cw * A31; // 3,1
-            G(q, 2, 1, e) = cw * A32; // 3,2
-            G(q, 2, 2, e) = cw * A33; // 3,3
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+               {
+                  const real_t J11 = J(qx, qy, qz, 0, 0, e),
+                               J12 = J(qx, qy, qz, 0, 1, e),
+                               J13 = J(qx, qy, qz, 0, 2, e);
+                  const real_t J21 = J(qx, qy, qz, 1, 0, e),
+                               J22 = J(qx, qy, qz, 1, 1, e),
+                               J23 = J(qx, qy, qz, 1, 2, e);
+                  const real_t J31 = J(qx, qy, qz, 2, 0, e),
+                               J32 = J(qx, qy, qz, 2, 1, e),
+                               J33 = J(qx, qy, qz, 2, 2, e);
+                  const real_t c =
+                     const_coeff ? C(0, 0, 0, 0) : C(qx, qy, qz, e);
+                  const real_t cw = W(qx, qy, qz) * c;
+                  // adj(J)
+                  const real_t A11 = (J22 * J33) - (J23 * J32);
+                  const real_t A12 = (J32 * J13) - (J12 * J33);
+                  const real_t A13 = (J12 * J23) - (J22 * J13);
+                  const real_t A21 = (J31 * J23) - (J21 * J33);
+                  const real_t A22 = (J11 * J33) - (J13 * J31);
+                  const real_t A23 = (J21 * J13) - (J11 * J23);
+                  const real_t A31 = (J21 * J32) - (J31 * J22);
+                  const real_t A32 = (J31 * J12) - (J11 * J32);
+                  const real_t A33 = (J11 * J22) - (J12 * J21);
+                  // Store wq * coeff * adj(J)
+                  A(0, 0, qx, qy, qz, e) = cw * A11;
+                  A(1, 0, qx, qy, qz, e) = cw * A12;
+                  A(2, 0, qx, qy, qz, e) = cw * A13;
+                  A(0, 1, qx, qy, qz, e) = cw * A21;
+                  A(1, 1, qx, qy, qz, e) = cw * A22;
+                  A(2, 1, qx, qy, qz, e) = cw * A23;
+                  A(0, 2, qx, qy, qz, e) = cw * A31;
+                  A(1, 2, qx, qy, qz, e) = cw * A32;
+                  A(2, 2, qx, qy, qz, e) = cw * A33;
+               }
+            }
          }
       });
    }
-}
-
-// PA Convection NL 2D kernel
-template<int T_D1D = 0, int T_Q1D = 0>
-static void PAConvectionNLApply2D(const int NE,
-                                  const Array<real_t> &b,
-                                  const Array<real_t> &g,
-                                  const Array<real_t> &bt,
-                                  const Vector &q_,
-                                  const Vector &x_,
-                                  Vector &y_,
-                                  const int d1d = 0,
-                                  const int q1d = 0)
-{
-   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, "");
-   auto B = Reshape(b.Read(), Q1D, D1D);
-   auto G = Reshape(g.Read(), Q1D, D1D);
-   auto Bt = Reshape(bt.Read(), D1D, Q1D);
-   auto Q = Reshape(q_.Read(), Q1D * Q1D, 2, 2, NE);
-   auto x = Reshape(x_.Read(), D1D, D1D, 2, NE);
-   auto y = Reshape(y_.ReadWrite(), D1D, D1D, 2, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   else
    {
-      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 data[max_Q1D][max_Q1D][2];
-      real_t grad0[max_Q1D][max_Q1D][2];
-      real_t grad1[max_Q1D][max_Q1D][2];
-      real_t Z[max_Q1D][max_Q1D][2];
-      for (int qy = 0; qy < Q1D; ++qy)
-      {
-         for (int qx = 0; qx < Q1D; ++qx)
-         {
-            data[qy][qx][0] = 0.0;
-            data[qy][qx][1] = 0.0;
-            grad0[qy][qx][0] = 0.0;
-            grad0[qy][qx][1] = 0.0;
-            grad1[qy][qx][0] = 0.0;
-            grad1[qy][qx][1] = 0.0;
-         }
-      }
-      for (int dy = 0; dy < D1D; ++dy)
-      {
-         real_t dataX[max_Q1D][2];
-         real_t gradX0[max_Q1D][2];
-         real_t gradX1[max_Q1D][2];
-         for (int qx = 0; qx < Q1D; ++qx)
-         {
-            dataX[qx][0] = 0.0;
-            dataX[qx][1] = 0.0;
-            gradX0[qx][0] = 0.0;
-            gradX0[qx][1] = 0.0;
-            gradX1[qx][0] = 0.0;
-            gradX1[qx][1] = 0.0;
-         }
-         for (int dx = 0; dx < D1D; ++dx)
-         {
-            const real_t s0 = x(dx, dy, 0, e);
-            const real_t s1 = x(dx, dy, 1, e);
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               const real_t Bx = B(qx, dx);
-               const real_t Gx = G(qx, dx);
-               dataX[qx][0] += s0 * Bx;
-               dataX[qx][1] += s1 * Bx;
-               gradX0[qx][0] += s0 * Gx;
-               gradX0[qx][1] += s0 * Bx;
-               gradX1[qx][0] += s1 * Gx;
-               gradX1[qx][1] += s1 * Bx;
-            }
-         }
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            const real_t By = B(qy, dy);
-            const real_t Gy = G(qy, dy);
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               data[qy][qx][0] += dataX[qx][0] * By;
-               data[qy][qx][1] += dataX[qx][1] * By;
-               grad0[qy][qx][0] += gradX0[qx][0] * By;
-               grad0[qy][qx][1] += gradX0[qx][1] * Gy;
-               grad1[qy][qx][0] += gradX1[qx][0] * By;
-               grad1[qy][qx][1] += gradX1[qx][1] * Gy;
-            }
-         }
-      }
-      for (int qy = 0; qy < Q1D; ++qy)
-      {
-         for (int qx = 0; qx < Q1D; ++qx)
-         {
-            const int q = qx + qy * Q1D;
-            const real_t u1 = data[qy][qx][0];
-            const real_t u2 = data[qy][qx][1];
-            const real_t grad00 = grad0[qy][qx][0];
-            const real_t grad01 = grad0[qy][qx][1];
-            const real_t grad10 = grad1[qy][qx][0];
-            const real_t grad11 = grad1[qy][qx][1];
-            const real_t Dxu1 = grad00 * Q(q, 0, 0, e) + grad01 * Q(q, 1, 0, e);
-            const real_t Dyu1 = grad00 * Q(q, 0, 1, e) + grad01 * Q(q, 1, 1, e);
-            const real_t Dxu2 = grad10 * Q(q, 0, 0, e) + grad11 * Q(q, 1, 0, e);
-            const real_t Dyu2 = grad10 * Q(q, 0, 1, e) + grad11 * Q(q, 1, 1, e);
-            Z[qy][qx][0] = u1 * Dxu1 + u2 * Dyu1;
-            Z[qy][qx][1] = u1 * Dxu2 + u2 * Dyu2;
-         }
-      }
-      for (int qy = 0; qy < Q1D; ++qy)
-      {
-         real_t Y[max_D1D][2];
-         for (int dx = 0; dx < D1D; ++dx)
-         {
-            Y[dx][0] = 0.0;
-            Y[dx][1] = 0.0;
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               const real_t Btx = Bt(dx, qx);
-               Y[dx][0] += Btx * Z[qy][qx][0];
-               Y[dx][1] += Btx * Z[qy][qx][1];
-            }
-         }
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               const real_t Bty = Bt(dy, qy);
-               y(dx, dy, 0, e) += Bty * Y[dx][0];
-               y(dx, dy, 1, e) += Bty * Y[dx][1];
-            }
-         }
-      }
-   });
-}
-
-// PA Convection NL 3D kernel
-template<int T_D1D = 0, int T_Q1D = 0>
-static void PAConvectionNLApply3D(const int NE,
-                                  const Array<real_t> &b,
-                                  const Array<real_t> &g,
-                                  const Array<real_t> &bt,
-                                  const Vector &q_,
-                                  const Vector &x_,
-                                  Vector &y_,
-                                  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, "");
-
-   auto B = Reshape(b.Read(), Q1D, D1D);
-   auto G = Reshape(g.Read(), Q1D, D1D);
-   auto Bt = Reshape(bt.Read(), D1D, Q1D);
-   auto Q = Reshape(q_.Read(), Q1D * Q1D * Q1D, VDIM, VDIM, NE);
-   auto x = Reshape(x_.Read(), D1D, D1D, D1D, VDIM, NE);
-   auto y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, VDIM, NE);
-
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
-   {
-      constexpr int VDIM = 3;
-      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 data[max_Q1D][max_Q1D][max_Q1D][VDIM];
-      real_t grad0[max_Q1D][max_Q1D][max_Q1D][VDIM];
-      real_t grad1[max_Q1D][max_Q1D][max_Q1D][VDIM];
-      real_t grad2[max_Q1D][max_Q1D][max_Q1D][VDIM];
-      real_t Z[max_Q1D][max_Q1D][max_Q1D][VDIM];
-      for (int qz = 0; qz < Q1D; ++qz)
-      {
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               data[qz][qy][qx][0] = 0.0;
-               data[qz][qy][qx][1] = 0.0;
-               data[qz][qy][qx][2] = 0.0;
-
-               grad0[qz][qy][qx][0] = 0.0;
-               grad0[qz][qy][qx][1] = 0.0;
-               grad0[qz][qy][qx][2] = 0.0;
-
-               grad1[qz][qy][qx][0] = 0.0;
-               grad1[qz][qy][qx][1] = 0.0;
-               grad1[qz][qy][qx][2] = 0.0;
-
-               grad2[qz][qy][qx][0] = 0.0;
-               grad2[qz][qy][qx][1] = 0.0;
-               grad2[qz][qy][qx][2] = 0.0;
-            }
-         }
-      }
-      for (int dz = 0; dz < D1D; ++dz)
-      {
-         real_t dataXY[max_Q1D][max_Q1D][VDIM];
-         real_t gradXY0[max_Q1D][max_Q1D][VDIM];
-         real_t gradXY1[max_Q1D][max_Q1D][VDIM];
-         real_t gradXY2[max_Q1D][max_Q1D][VDIM];
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               dataXY[qy][qx][0] = 0.0;
-               dataXY[qy][qx][1] = 0.0;
-               dataXY[qy][qx][2] = 0.0;
-
-               gradXY0[qy][qx][0] = 0.0;
-               gradXY0[qy][qx][1] = 0.0;
-               gradXY0[qy][qx][2] = 0.0;
-
-               gradXY1[qy][qx][0] = 0.0;
-               gradXY1[qy][qx][1] = 0.0;
-               gradXY1[qy][qx][2] = 0.0;
-
-               gradXY2[qy][qx][0] = 0.0;
-               gradXY2[qy][qx][1] = 0.0;
-               gradXY2[qy][qx][2] = 0.0;
-            }
-         }
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            real_t dataX[max_Q1D][VDIM];
-            real_t gradX0[max_Q1D][VDIM];
-            real_t gradX1[max_Q1D][VDIM];
-            real_t gradX2[max_Q1D][VDIM];
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               dataX[qx][0] = 0.0;
-               dataX[qx][1] = 0.0;
-               dataX[qx][2] = 0.0;
-
-               gradX0[qx][0] = 0.0;
-               gradX0[qx][1] = 0.0;
-               gradX0[qx][2] = 0.0;
-
-               gradX1[qx][0] = 0.0;
-               gradX1[qx][1] = 0.0;
-               gradX1[qx][2] = 0.0;
-
-               gradX2[qx][0] = 0.0;
-               gradX2[qx][1] = 0.0;
-               gradX2[qx][2] = 0.0;
-            }
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               const real_t s0 = x(dx, dy, dz, 0, e);
-               const real_t s1 = x(dx, dy, dz, 1, e);
-               const real_t s2 = x(dx, dy, dz, 2, e);
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  const real_t Bx = B(qx, dx);
-                  const real_t Gx = G(qx, dx);
-
-                  dataX[qx][0] += s0 * Bx;
-                  dataX[qx][1] += s1 * Bx;
-                  dataX[qx][2] += s2 * Bx;
-
-                  gradX0[qx][0] += s0 * Gx;
-                  gradX0[qx][1] += s0 * Bx;
-                  gradX0[qx][2] += s0 * Bx;
-
-                  gradX1[qx][0] += s1 * Gx;
-                  gradX1[qx][1] += s1 * Bx;
-                  gradX1[qx][2] += s1 * Bx;
-
-                  gradX2[qx][0] += s2 * Gx;
-                  gradX2[qx][1] += s2 * Bx;
-                  gradX2[qx][2] += s2 * Bx;
-               }
-            }
-            for (int qy = 0; qy < Q1D; ++qy)
-            {
-               const real_t By = B(qy, dy);
-               const real_t Gy = G(qy, dy);
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  dataXY[qy][qx][0] += dataX[qx][0] * By;
-                  dataXY[qy][qx][1] += dataX[qx][1] * By;
-                  dataXY[qy][qx][2] += dataX[qx][2] * By;
-
-                  gradXY0[qy][qx][0] += gradX0[qx][0] * By;
-                  gradXY0[qy][qx][1] += gradX0[qx][1] * Gy;
-                  gradXY0[qy][qx][2] += gradX0[qx][2] * By;
-
-                  gradXY1[qy][qx][0] += gradX1[qx][0] * By;
-                  gradXY1[qy][qx][1] += gradX1[qx][1] * Gy;
-                  gradXY1[qy][qx][2] += gradX1[qx][2] * By;
-
-                  gradXY2[qy][qx][0] += gradX2[qx][0] * By;
-                  gradXY2[qy][qx][1] += gradX2[qx][1] * Gy;
-                  gradXY2[qy][qx][2] += gradX2[qx][2] * By;
-               }
-            }
-         }
-         for (int qz = 0; qz < Q1D; ++qz)
-         {
-            const real_t Bz = B(qz, dz);
-            const real_t Gz = G(qz, dz);
-            for (int qy = 0; qy < Q1D; ++qy)
-            {
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  data[qz][qy][qx][0] += dataXY[qy][qx][0] * Bz;
-                  data[qz][qy][qx][1] += dataXY[qy][qx][1] * Bz;
-                  data[qz][qy][qx][2] += dataXY[qy][qx][2] * Bz;
-
-                  grad0[qz][qy][qx][0] += gradXY0[qy][qx][0] * Bz;
-                  grad0[qz][qy][qx][1] += gradXY0[qy][qx][1] * Bz;
-                  grad0[qz][qy][qx][2] += gradXY0[qy][qx][2] * Gz;
-
-                  grad1[qz][qy][qx][0] += gradXY1[qy][qx][0] * Bz;
-                  grad1[qz][qy][qx][1] += gradXY1[qy][qx][1] * Bz;
-                  grad1[qz][qy][qx][2] += gradXY1[qy][qx][2] * Gz;
-
-                  grad2[qz][qy][qx][0] += gradXY2[qy][qx][0] * Bz;
-                  grad2[qz][qy][qx][1] += gradXY2[qy][qx][1] * Bz;
-                  grad2[qz][qy][qx][2] += gradXY2[qy][qx][2] * Gz;
-               }
-            }
-         }
-      }
-      for (int qz = 0; qz < Q1D; ++qz)
-      {
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               const int q = qx + Q1D * (qy + qz * Q1D);
-
-               const real_t u1 = data[qz][qy][qx][0];
-               const real_t u2 = data[qz][qy][qx][1];
-               const real_t u3 = data[qz][qy][qx][2];
-
-               const real_t grad00 = grad0[qz][qy][qx][0];
-               const real_t grad01 = grad0[qz][qy][qx][1];
-               const real_t grad02 = grad0[qz][qy][qx][2];
-
-               const real_t grad10 = grad1[qz][qy][qx][0];
-               const real_t grad11 = grad1[qz][qy][qx][1];
-               const real_t grad12 = grad1[qz][qy][qx][2];
-
-               const real_t grad20 = grad2[qz][qy][qx][0];
-               const real_t grad21 = grad2[qz][qy][qx][1];
-               const real_t grad22 = grad2[qz][qy][qx][2];
-
-               const real_t Dxu1 = grad00 * Q(q, 0, 0, e)
-                                   + grad01 * Q(q, 1, 0, e)
-                                   + grad02 * Q(q, 2, 0, e);
-               const real_t Dyu1 = grad00 * Q(q, 0, 1, e)
-                                   + grad01 * Q(q, 1, 1, e)
-                                   + grad02 * Q(q, 2, 1, e);
-               const real_t Dzu1 = grad00 * Q(q, 0, 2, e)
-                                   + grad01 * Q(q, 1, 2, e)
-                                   + grad02 * Q(q, 2, 2, e);
-
-               const real_t Dxu2 = grad10 * Q(q, 0, 0, e)
-                                   + grad11 * Q(q, 1, 0, e)
-                                   + grad12 * Q(q, 2, 0, e);
-               const real_t Dyu2 = grad10 * Q(q, 0, 1, e)
-                                   + grad11 * Q(q, 1, 1, e)
-                                   + grad12 * Q(q, 2, 1, e);
-               const real_t Dzu2 = grad10 * Q(q, 0, 2, e)
-                                   + grad11 * Q(q, 1, 2, e)
-                                   + grad12 * Q(q, 2, 2, e);
-
-               const real_t Dxu3 = grad20 * Q(q, 0, 0, e)
-                                   + grad21 * Q(q, 1, 0, e)
-                                   + grad22 * Q(q, 2, 0, e);
-               const real_t Dyu3 = grad20 * Q(q, 0, 1, e)
-                                   + grad21 * Q(q, 1, 1, e)
-                                   + grad22 * Q(q, 2, 1, e);
-               const real_t Dzu3 = grad20 * Q(q, 0, 2, e)
-                                   + grad21 * Q(q, 1, 2, e)
-                                   + grad22 * Q(q, 2, 2, e);
-
-               Z[qz][qy][qx][0] = u1 * Dxu1 + u2 * Dyu1 + u3 * Dzu1;
-               Z[qz][qy][qx][1] = u1 * Dxu2 + u2 * Dyu2 + u3 * Dzu2;
-               Z[qz][qy][qx][2] = u1 * Dxu3 + u2 * Dyu3 + u3 * Dzu3;
-            }
-         }
-      }
-      for (int qz = 0; qz < Q1D; ++qz)
-      {
-         real_t opXY[max_D1D][max_D1D][VDIM];
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               opXY[dy][dx][0] = 0.0;
-               opXY[dy][dx][1] = 0.0;
-               opXY[dy][dx][2] = 0.0;
-            }
-         }
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            real_t opX[max_D1D][VDIM];
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               opX[dx][0] = 0.0;
-               opX[dx][1] = 0.0;
-               opX[dx][2] = 0.0;
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  const real_t Btx = Bt(dx, qx);
-                  opX[dx][0] += Btx * Z[qz][qy][qx][0];
-                  opX[dx][1] += Btx * Z[qz][qy][qx][1];
-                  opX[dx][2] += Btx * Z[qz][qy][qx][2];
-               }
-            }
-            for (int dy = 0; dy < D1D; ++dy)
-            {
-               for (int dx = 0; dx < D1D; ++dx)
-               {
-                  const real_t Bty = Bt(dy, qy);
-                  opXY[dy][dx][0] += Bty * opX[dx][0];
-                  opXY[dy][dx][1] += Bty * opX[dx][1];
-                  opXY[dy][dx][2] += Bty * opX[dx][2];
-               }
-            }
-         }
-         for (int dz = 0; dz < D1D; ++dz)
-         {
-            for (int dy = 0; dy < D1D; ++dy)
-            {
-               for (int dx = 0; dx < D1D; ++dx)
-               {
-                  const real_t Btz = Bt(dz, qz);
-                  y(dx, dy, dz, 0, e) += Btz * opXY[dy][dx][0];
-                  y(dx, dy, dz, 1, e) += Btz * opXY[dy][dx][1];
-                  y(dx, dy, dz, 2, e) += Btz * opXY[dy][dx][2];
-               }
-            }
-         }
-      }
-   });
-}
-
-template<int T_D1D = 0, int T_Q1D = 0, int T_MAX_D1D = 0, int T_MAX_Q1D = 0>
-static void SmemPAConvectionNLApply3D(const int NE,
-                                      const Array<real_t> &b_,
-                                      const Array<real_t> &g_,
-                                      const Vector &d_,
-                                      const Vector &x_,
-                                      Vector &y_,
-                                      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;
-   constexpr int MD1 = T_D1D ? T_D1D : T_MAX_D1D;
-   constexpr int MQ1 = T_Q1D ? T_Q1D : T_MAX_Q1D;
-   MFEM_VERIFY(D1D <= MD1, "");
-   MFEM_VERIFY(Q1D <= MQ1, "");
-
-   auto b = Reshape(b_.Read(), Q1D, D1D);
-   auto g = Reshape(g_.Read(), Q1D, D1D);
-   auto D = Reshape(d_.Read(), Q1D * Q1D * Q1D, VDIM, VDIM, NE);
-   auto x = Reshape(x_.Read(), D1D, D1D, D1D, VDIM, NE);
-   auto Y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, VDIM, NE);
-
-   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
-   {
-      const int tidz = MFEM_THREAD_ID(z);
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
-      constexpr int MD1 = T_D1D ? T_D1D : T_MAX_D1D;
-      constexpr int MQ1 = T_Q1D ? T_Q1D : T_MAX_Q1D;
-      MFEM_SHARED real_t BG[2][MQ1 * MD1];
-      real_t(*B)[MD1] = (real_t(*)[MD1])(BG + 0);
-      real_t(*G)[MD1] = (real_t(*)[MD1])(BG + 1);
-      real_t(*Bt)[MQ1] = (real_t(*)[MQ1])(BG + 0);
-      MFEM_SHARED real_t U[2][MQ1][MQ1][MQ1];
-      MFEM_SHARED real_t sm0[3][MQ1 * MQ1 * MQ1];
-      MFEM_SHARED real_t sm1[3][MQ1 * MQ1 * MQ1];
-      real_t(*DDQ0)[MD1][MQ1] = (real_t(*)[MD1][MQ1])(sm0 + 0);
-      real_t(*DDQ1)[MD1][MQ1] = (real_t(*)[MD1][MQ1])(sm0 + 1);
-      real_t(*X)[MD1][MD1] = (real_t(*)[MD1][MD1])(sm0 + 2);
-      real_t(*DQQ0)[MQ1][MQ1] = (real_t(*)[MQ1][MQ1])(sm1 + 0);
-      real_t(*DQQ1)[MQ1][MQ1] = (real_t(*)[MQ1][MQ1])(sm1 + 1);
-      real_t(*DQQ2)[MQ1][MQ1] = (real_t(*)[MQ1][MQ1])(sm1 + 2);
-      real_t(*QQQ0)[MQ1][MQ1] = (real_t(*)[MQ1][MQ1])(sm0 + 0);
-      real_t(*QQQ1)[MQ1][MQ1] = (real_t(*)[MQ1][MQ1])(sm0 + 1);
-      real_t(*QQQ2)[MQ1][MQ1] = (real_t(*)[MQ1][MQ1])(sm0 + 2);
-      real_t(*QQD0)[MQ1][MD1] = (real_t(*)[MQ1][MD1])(sm1 + 0);
-      real_t(*QDD0)[MD1][MD1] = (real_t(*)[MD1][MD1])(sm0 + 0);
-      MFEM_SHARED real_t Z[MQ1][MQ1][MQ1];
-
-      for (int cy = 0; cy < VDIM; ++cy)
-      {
-         if (tidz == 0)
-         {
-            MFEM_FOREACH_THREAD(q, x, Q1D)
-            {
-               MFEM_FOREACH_THREAD(d, y, D1D)
-               {
-                  B[q][d] = b(q, d);
-                  G[q][d] = g(q, d);
-               }
-            }
-         }
-         MFEM_FOREACH_THREAD(qz, z, Q1D)
-         {
-            MFEM_FOREACH_THREAD(qy, y, Q1D)
-            {
-               MFEM_FOREACH_THREAD(qx, x, Q1D) { Z[qz][qy][qx] = 0.0; }
-            }
-         }
-         MFEM_SYNC_THREAD;
-         for (int c = 0; c < VDIM; ++c)
-         {
-            MFEM_FOREACH_THREAD(dz, z, D1D)
-            {
-               MFEM_FOREACH_THREAD(dy, y, D1D)
-               {
-                  MFEM_FOREACH_THREAD(dx, x, D1D)
-                  {
-                     X[dz][dy][dx] = x(dx, dy, dz, cy, e);
-                     U[0][dz][dy][dx] = x(dx, dy, dz, c, e);
-                  }
-               }
-            }
-            MFEM_SYNC_THREAD;
-            MFEM_FOREACH_THREAD(dz, z, D1D)
-            {
-               MFEM_FOREACH_THREAD(dy, y, D1D)
-               {
-                  MFEM_FOREACH_THREAD(qx, x, Q1D)
-                  {
-                     real_t u = 0.0;
-                     real_t v = 0.0;
-                     real_t z = 0.0;
-                     for (int dx = 0; dx < D1D; ++dx)
-                     {
-                        const real_t coord = X[dz][dy][dx];
-                        const real_t value = U[0][dz][dy][dx];
-                        u += coord * B[qx][dx];
-                        v += coord * G[qx][dx];
-                        z += value * B[qx][dx];
-                     }
-                     DDQ0[dz][dy][qx] = u;
-                     DDQ1[dz][dy][qx] = v;
-                     U[1][dz][dy][qx] = z;
-                  }
-               }
-            }
-            MFEM_SYNC_THREAD;
-            MFEM_FOREACH_THREAD(dz, z, D1D)
-            {
-               MFEM_FOREACH_THREAD(qy, y, Q1D)
-               {
-                  MFEM_FOREACH_THREAD(qx, x, Q1D)
-                  {
-                     real_t u = 0.0;
-                     real_t v = 0.0;
-                     real_t w = 0.0;
-                     real_t z = 0.0;
-                     for (int dy = 0; dy < D1D; ++dy)
-                     {
-                        u += DDQ1[dz][dy][qx] * B[qy][dy];
-                        v += DDQ0[dz][dy][qx] * G[qy][dy];
-                        w += DDQ0[dz][dy][qx] * B[qy][dy];
-                        z += U[1][dz][dy][qx] * B[qy][dy];
-                     }
-                     DQQ0[dz][qy][qx] = u;
-                     DQQ1[dz][qy][qx] = v;
-                     DQQ2[dz][qy][qx] = w;
-                     U[0][dz][qy][qx] = z;
-                  }
-               }
-            }
-            MFEM_SYNC_THREAD;
-            MFEM_FOREACH_THREAD(qz, z, Q1D)
-            {
-               MFEM_FOREACH_THREAD(qy, y, Q1D)
-               {
-                  MFEM_FOREACH_THREAD(qx, x, Q1D)
-                  {
-                     real_t u = 0.0;
-                     real_t v = 0.0;
-                     real_t w = 0.0;
-                     real_t z = 0.0;
-                     for (int dz = 0; dz < D1D; ++dz)
-                     {
-                        u += DQQ0[dz][qy][qx] * B[qz][dz];
-                        v += DQQ1[dz][qy][qx] * B[qz][dz];
-                        w += DQQ2[dz][qy][qx] * G[qz][dz];
-                        z += U[0][dz][qy][qx] * B[qz][dz];
-                     }
-                     QQQ0[qz][qy][qx] = u;
-                     QQQ1[qz][qy][qx] = v;
-                     QQQ2[qz][qy][qx] = w;
-                     U[1][qz][qy][qx] = z;
-                  }
-               }
-            }
-            MFEM_SYNC_THREAD;
-            MFEM_FOREACH_THREAD(qz, z, Q1D)
-            {
-               MFEM_FOREACH_THREAD(qy, y, Q1D)
-               {
-                  MFEM_FOREACH_THREAD(qx, x, Q1D)
-                  {
-                     const int q = qx + (qy + qz * Q1D) * Q1D;
-                     const real_t z = U[1][qz][qy][qx];
-                     const real_t gX = QQQ0[qz][qy][qx];
-                     const real_t gY = QQQ1[qz][qy][qx];
-                     const real_t gZ = QQQ2[qz][qy][qx];
-                     const real_t d = gX * D(q, 0, c, e) + gY * D(q, 1, c, e)
-                                      + gZ * D(q, 2, c, e);
-                     Z[qz][qy][qx] += z * d;
-                  }
-               }
-            }
-            MFEM_SYNC_THREAD;
-         } // for each conv component
-         if (tidz == 0)
-         {
-            MFEM_FOREACH_THREAD(d, y, D1D)
-            {
-               MFEM_FOREACH_THREAD(q, x, Q1D) { Bt[d][q] = b(q, d); }
-            }
-         }
-         MFEM_SYNC_THREAD;
-         MFEM_FOREACH_THREAD(qz, z, Q1D)
-         {
-            MFEM_FOREACH_THREAD(qy, y, Q1D)
-            {
-               MFEM_FOREACH_THREAD(dx, x, D1D)
-               {
-                  real_t u = 0.0;
-                  for (int qx = 0; qx < Q1D; ++qx)
-                  {
-                     u += Z[qz][qy][qx] * Bt[dx][qx];
-                  }
-                  QQD0[qz][qy][dx] = u;
-               }
-            }
-         }
-         MFEM_SYNC_THREAD;
-         MFEM_FOREACH_THREAD(qz, z, Q1D)
-         {
-            MFEM_FOREACH_THREAD(dy, y, D1D)
-            {
-               MFEM_FOREACH_THREAD(dx, x, D1D)
-               {
-                  real_t u = 0.0;
-                  for (int qy = 0; qy < Q1D; ++qy)
-                  {
-                     u += QQD0[qz][qy][dx] * Bt[dy][qy];
-                  }
-                  QDD0[qz][dy][dx] = u;
-               }
-            }
-         }
-         MFEM_SYNC_THREAD;
-         MFEM_FOREACH_THREAD(dz, z, D1D)
-         {
-            MFEM_FOREACH_THREAD(dy, y, D1D)
-            {
-               MFEM_FOREACH_THREAD(dx, x, D1D)
-               {
-                  real_t u = 0.0;
-                  for (int qz = 0; qz < Q1D; ++qz)
-                  {
-                     u += QDD0[qz][dy][dx] * Bt[dz][qz];
-                  }
-                  Y(dx, dy, dz, cy, e) += u;
-               }
-            }
-         }
-         MFEM_SYNC_THREAD;
-      }
-   });
+      MFEM_ABORT("dim " << dim << " not supported!");
+   }
 }
 
 void VectorConvectionNLFIntegrator::AddMultPA(const Vector &x, Vector &y) const
@@ -812,26 +197,13 @@ void VectorConvectionNLFIntegrator::AddMultPA(const Vector &x, Vector &y) const
    }
    else
    {
-      const int NE = ne;
-      const int D1D = maps->ndof;
-      const int Q1D = maps->nqpt;
-      const Vector &QV = pa_data;
-      const Array<real_t> &B = maps->B;
-      const Array<real_t> &G = maps->G;
-      const Array<real_t> &Bt = maps->Bt;
-      if (dim == 2)
-      {
-         return PAConvectionNLApply2D(NE, B, G, Bt, QV, x, y, D1D, Q1D);
-      }
-      if (dim == 3)
-      {
-         constexpr int T_MAX_D1D = 8;
-         constexpr int T_MAX_Q1D = 8;
-         MFEM_VERIFY(D1D <= T_MAX_D1D && Q1D <= T_MAX_Q1D, "Not yet implemented!");
-         return SmemPAConvectionNLApply3D<0, 0, T_MAX_D1D, T_MAX_Q1D>
-                (NE, B, G, QV, x, y, D1D, Q1D);
-      }
-      MFEM_ABORT("Not yet implemented!");
+      AddMultPAKernels::Run(dim, d1d, q1d, ne,
+                            maps->B.Read(),
+                            maps->G.Read(),
+                            pa_adj.Read(),
+                            x.Read(),
+                            y.ReadWrite(),
+                            d1d, q1d);
    }
 }
 

fem/integ/nonlininteg_vecconvection_pa.hpp

@@ -0,0 +1,209 @@
+// 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"
+#include "../nonlininteg.hpp"
+
+namespace mfem
+{
+
+/// \cond DO_NOT_DOCUMENT
+
+namespace internal
+{
+
+// PA Convection NL 2D kernel
+template<int T_D1D = 0, int T_Q1D = 0>
+inline void SmemPAConvectionNLApply2D(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 = 0,
+                                      const int q1d = 0)
+{
+   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 B = Reshape(b, Q1D, D1D);
+   const auto G = Reshape(g, Q1D, D1D);
+   const auto A = Reshape(a, VDIM, DIM, Q1D, Q1D, NE);
+   const auto X = Reshape(x, 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 : DofQuadLimits::MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1], sB[MD1][MQ1], sG[MD1][MQ1];
+
+      kernels::internal::vd_regs2d_t<VDIM, DIM, MQ1> g0, g1;
+      kernels::internal::v_regs2d_t<VDIM, MQ1> r0, r1;
+      kernels::internal::v_regs2d_t<VDIM, MQ1> s0, s1;
+
+      kernels::internal::LoadMatrix(D1D, Q1D, B, sB);
+      kernels::internal::LoadMatrix(D1D, Q1D, G, sG);
+
+      kernels::internal::LoadDofs2d(e, D1D, X, r0);
+      kernels::internal::Eval2d(D1D, Q1D, smem, sB, r0, r1); // u vector-value
+      kernels::internal::LoadDofs2d(e, D1D, X, g0);
+      kernels::internal::Grad2d(D1D, Q1D, smem, sB, sG, g0, g1); // u vector-gradient
+
+      MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+         {
+            const future::tensor<real_t, 2> U =
+            {
+               r1[0][qy][qx], r1[1][qy][qx]
+            };
+            const future::tensor<real_t, 2,2> gradU = {{
+                  {g1[0][0][qy][qx], g1[1][0][qy][qx]},
+                  {g1[0][1][qy][qx], g1[1][1][qy][qx]},
+               }
+            };
+            const future::tensor<real_t, 2,2> Q = {{
+                  {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, 2> conv = transpose(gradU) * (Q * U);
+            s0[0][qy][qx] = conv[0];
+            s0[1][qy][qx] = conv[1];
+         }
+      }
+      MFEM_SYNC_THREAD;
+      kernels::internal::EvalTranspose2d(D1D, Q1D, smem, sB, s0, s1);
+      kernels::internal::WriteDofs2d(e, D1D, s1, Y);
+   });
+}
+
+// PA Convection NL 3D kernel
+template<int T_D1D = 0, int T_Q1D = 0>
+inline void SmemPAConvectionNLApply3D(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 = 0,
+                                      const int q1d = 0)
+{
+   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 B = Reshape(b, Q1D, D1D);
+   const auto G = Reshape(g, Q1D, D1D);
+   const auto A = Reshape(a, VDIM, DIM, Q1D, Q1D, Q1D, NE);
+   const auto X = Reshape(x, 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 : DofQuadLimits::MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      MFEM_SHARED real_t smem[MQ1][MQ1], sB[MD1][MQ1], sG[MD1][MQ1];
+
+      kernels::internal::vd_regs3d_t<VDIM, DIM, MQ1> g0, g1;
+      kernels::internal::v_regs3d_t<VDIM, MQ1> r0, r1;
+      kernels::internal::v_regs3d_t<VDIM, MQ1> s0, s1;
+
+      kernels::internal::LoadMatrix(D1D, Q1D, B, sB);
+      kernels::internal::LoadMatrix(D1D, Q1D, G, sG);
+
+      kernels::internal::LoadDofs3d(e, D1D, X, r0);
+      kernels::internal::Eval3d(D1D, Q1D, smem, sB, r0, r1); // u vector-value
+      kernels::internal::LoadDofs3d(e, D1D, X, g0);
+      kernels::internal::Grad3d(D1D, Q1D, smem, sB, sG, g0, g1); // u vector-gradient
+
+      for (int qz = 0; qz < Q1D; qz++)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qy, y, Q1D)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qx, x, Q1D)
+            {
+               const future::tensor<real_t, 3> U =
+               {
+                  r1[0][qz][qy][qx], r1[1][qz][qy][qx], r1[2][qz][qy][qx]
+               };
+               const future::tensor<real_t, 3,3> gradU = {{
+                     {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 future::tensor<real_t, 3,3> Q = {{
+                     {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, 3> conv = transpose(gradU) * (Q * U);
+               s0[0][qz][qy][qx] = conv[0];
+               s0[1][qz][qy][qx] = conv[1];
+               s0[2][qz][qy][qx] = conv[2];
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      kernels::internal::EvalTranspose3d(D1D, Q1D, smem, sB, s0, s1);
+      kernels::internal::WriteDofs3d(e, D1D, s1, Y);
+   });
+}
+
+} // namespace internal
+
+template<int DIM, int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::AddMultPAType
+VectorConvectionNLFIntegrator::AddMultPAKernels::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   if constexpr (DIM == 2)
+   {
+      return internal::SmemPAConvectionNLApply2D<T_D1D, T_Q1D>;
+   }
+   else if constexpr (DIM == 3)
+   {
+      return internal::SmemPAConvectionNLApply3D<T_D1D, T_Q1D>;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+inline VectorConvectionNLFIntegrator::AddMultPAType
+VectorConvectionNLFIntegrator::AddMultPAKernels::Fallback
+(int dim, int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   if (dim == 2)
+   {
+      return internal::SmemPAConvectionNLApply2D<>;
+   }
+   else if (dim == 3)
+   {
+      return internal::SmemPAConvectionNLApply3D<>;
+   }
+   else { MFEM_ABORT("Unsupported kernel"); }
+}
+
+/// \endcond DO_NOT_DOCUMENT
+
+} // namespace mfem

fem/integ/nonlininteg_vecconvection_pa_diag.cpp

@@ -0,0 +1,50 @@
+// 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 "../ceed/interface/util.hpp"
+#include "./nonlininteg_vecconvection_pa_diag.hpp" // IWYU pragma: keep
+
+namespace mfem
+{
+
+void VectorConvectionNLFIntegrator::AssembleGradDiagonalPA(Vector &de) const
+{
+   MFEM_VERIFY(!DeviceCanUseCeed(),
+               "VectorConvectionNLFIntegrator PA gradients are not supported "
+               "with the libCEED backend");
+
+   if (dim == 2)
+   {
+      GradDiagPA2D::Run(d1d, q1d, ne,
+                        maps->B.Read(),
+                        maps->G.Read(),
+                        pa_adj.Read(),
+                        pa_u.Read(),
+                        de.ReadWrite(),
+                        d1d, q1d);
+   }
+   else if (dim == 3)
+   {
+      GradDiagPA3D::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");
+   }
+}
+
+} // namespace mfem

fem/integ/nonlininteg_vecconvection_pa_diag.hpp

@@ -0,0 +1,302 @@
+// 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"
+#include "../nonlininteg.hpp"
+
+namespace mfem
+{
+
+/// \cond DO_NOT_DOCUMENT
+
+namespace internal
+{
+
+template<int T_D1D = 0, int T_Q1D = 0>
+inline 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 : DofQuadLimits::MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      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>
+inline 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 : DofQuadLimits::MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      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;
+         }
+      }
+   });
+}
+
+} // namespace internal
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::GradDiagPAType
+VectorConvectionNLFIntegrator::GradDiagPA2D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return internal::SmemPAConvectionNLGradDiagonal2D<T_D1D, T_Q1D>;
+}
+
+inline VectorConvectionNLFIntegrator::GradDiagPAType
+VectorConvectionNLFIntegrator::GradDiagPA2D::Fallback(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   return internal::SmemPAConvectionNLGradDiagonal2D<>;
+}
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::GradDiagPAType
+VectorConvectionNLFIntegrator::GradDiagPA3D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return internal::SmemPAConvectionNLGradDiagonal3D<T_D1D, T_Q1D>;
+}
+
+inline VectorConvectionNLFIntegrator::GradDiagPAType
+VectorConvectionNLFIntegrator::GradDiagPA3D::Fallback(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   return internal::SmemPAConvectionNLGradDiagonal3D<>;
+}
+
+/// \endcond DO_NOT_DOCUMENT
+
+} // namespace mfem

fem/integ/nonlininteg_vecconvection_pa_grad.cpp

@@ -0,0 +1,64 @@
+// 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 "../ceed/interface/util.hpp"
+#include "./nonlininteg_vecconvection_pa_grad.hpp" // IWYU pragma: keep
+
+namespace mfem
+{
+
+void VectorConvectionNLFIntegrator::AssembleGradPA(
+   const Vector &u, const FiniteElementSpace &fes)
+{
+   MFEM_VERIFY(!DeviceCanUseCeed(),
+               "VectorConvectionNLFIntegrator PA gradients are not supported "
+               "with the libCEED backend");
+
+   this->pa_u = u;
+   AssemblePA(fes);
+}
+
+void VectorConvectionNLFIntegrator::AddMultGradPA(const Vector &x,
+                                                  Vector &y) const
+{
+   MFEM_VERIFY(!DeviceCanUseCeed(),
+               "VectorConvectionNLFIntegrator PA gradients are not supported "
+               "with the libCEED backend");
+
+   if (dim == 2)
+   {
+      AddMultGradPA2D::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)
+   {
+      AddMultGradPA3D::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");
+   }
+}
+
+} // namespace mfem

fem/integ/nonlininteg_vecconvection_pa_grad.hpp

@@ -0,0 +1,257 @@
+// 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"
+#include "../nonlininteg.hpp"
+
+namespace mfem
+{
+
+/// \cond DO_NOT_DOCUMENT
+
+namespace internal
+{
+
+template<int T_D1D = 0, int T_Q1D = 0>
+inline 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 : DofQuadLimits::MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      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>
+inline 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 : DofQuadLimits::MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
+
+      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);
+   });
+}
+
+} // namespace internal
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::AddMultGradPAType
+VectorConvectionNLFIntegrator::AddMultGradPA2D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return internal::SmemPAConvectionNLGradApply2D<T_D1D, T_Q1D>;
+}
+
+inline VectorConvectionNLFIntegrator::AddMultGradPAType
+VectorConvectionNLFIntegrator::AddMultGradPA2D::Fallback(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   return internal::SmemPAConvectionNLGradApply2D<>;
+}
+
+template<int T_D1D, int T_Q1D>
+VectorConvectionNLFIntegrator::AddMultGradPAType
+VectorConvectionNLFIntegrator::AddMultGradPA3D::Kernel()
+{
+   static_assert(T_D1D <= T_Q1D, "d1d > q1d is not supported");
+   return internal::SmemPAConvectionNLGradApply3D<T_D1D, T_Q1D>;
+}
+
+inline VectorConvectionNLFIntegrator::AddMultGradPAType
+VectorConvectionNLFIntegrator::AddMultGradPA3D::Fallback(int d1d, int q1d)
+{
+   MFEM_VERIFY(d1d <= q1d, "d1d > q1d is not supported");
+   MFEM_VERIFY(d1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
+   MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+   return internal::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/nonlinearform_ext.cpp

@@ -100,6 +100,17 @@ PANonlinearFormExtension::Gradient::Gradient(const PANonlinearFormExtension &e):
 
 void PANonlinearFormExtension::Gradient::AssembleGrad(const Vector &g)
 {
+   if (DeviceCanUseCeed())
+   {
+      for (int i = 0; i < ext.dnfi.Size(); ++i)
+      {
+         MFEM_VERIFY(dynamic_cast<VectorConvectionNLFIntegrator *>
+                     (ext.dnfi[i]) == nullptr,
+                     "VectorConvectionNLFIntegrator PA gradients are not supported "
+                     "with the libCEED backend");
+      }
+   }
+
    ext.elemR->Mult(g, ext.xe);
    for (int i = 0; i < ext.dnfi.Size(); ++i)
    {

fem/nonlininteg.cpp

@@ -954,4 +954,74 @@ void SkewSymmetricVectorConvectionNLFIntegrator::AssembleElementGrad(
    }
 }
 
+void ConvectiveVectorConvectionNLFIntegrator::AssemblePA(
+   const FiniteElementSpace &)
+{
+   MFEM_ABORT("ConvectiveVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void ConvectiveVectorConvectionNLFIntegrator::AssembleGradPA(
+   const Vector &, const FiniteElementSpace &)
+{
+   MFEM_ABORT("ConvectiveVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void ConvectiveVectorConvectionNLFIntegrator::AddMultPA(
+   const Vector &, Vector &) const
+{
+   MFEM_ABORT("ConvectiveVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void ConvectiveVectorConvectionNLFIntegrator::AddMultGradPA(
+   const Vector &, Vector &) const
+{
+   MFEM_ABORT("ConvectiveVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void ConvectiveVectorConvectionNLFIntegrator::AssembleGradDiagonalPA(
+   Vector &) const
+{
+   MFEM_ABORT("ConvectiveVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void SkewSymmetricVectorConvectionNLFIntegrator::AssemblePA(
+   const FiniteElementSpace &)
+{
+   MFEM_ABORT("SkewSymmetricVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void SkewSymmetricVectorConvectionNLFIntegrator::AssembleGradPA(
+   const Vector &, const FiniteElementSpace &)
+{
+   MFEM_ABORT("SkewSymmetricVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void SkewSymmetricVectorConvectionNLFIntegrator::AddMultPA(
+   const Vector &, Vector &) const
+{
+   MFEM_ABORT("SkewSymmetricVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void SkewSymmetricVectorConvectionNLFIntegrator::AddMultGradPA(
+   const Vector &, Vector &) const
+{
+   MFEM_ABORT("SkewSymmetricVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
+void SkewSymmetricVectorConvectionNLFIntegrator::AssembleGradDiagonalPA(
+   Vector &) const
+{
+   MFEM_ABORT("SkewSymmetricVectorConvectionNLFIntegrator does not support "
+              "partial assembly; use VectorConvectionNLFIntegrator");
+}
+
 }

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,15 +385,17 @@ 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) { }
+   struct Kernels { Kernels(); };
 
-   VectorConvectionNLFIntegrator() = default;
+   VectorConvectionNLFIntegrator(Coefficient &q): Q(&q) { static Kernels kernels; }
+
+   VectorConvectionNLFIntegrator() { static Kernels kernels; }
 
    static const IntegrationRule &GetRule(const FiniteElement &fe,
                                          const ElementTransformation &T);
@@ -411,12 +414,55 @@ 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 AddMultPAType =
+      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(AddMultPAKernels, AddMultPAType, (int, int, int));
 
+   void AddMultGradPA(const Vector &x, Vector &y) const override;
+
+   using AddMultGradPAType =
+      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(AddMultGradPA2D, AddMultGradPAType, (int, int));
+   MFEM_REGISTER_KERNELS(AddMultGradPA3D, AddMultGradPAType, (int, int));
+
+   void AssembleGradDiagonalPA(Vector &) const override;
+
+   using GradDiagPAType =
+      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(GradDiagPA2D, GradDiagPAType, (int, int));
+   MFEM_REGISTER_KERNELS(GradDiagPA3D, GradDiagPAType, (int, int));
+
+   template <int DIM, int D1D, int Q1D>
+   static void AddSpecialization()
+   {
+      AddMultPAKernels::Specialization<DIM, D1D, Q1D>::Add();
+      if constexpr (DIM == 2)
+      {
+         AddMultGradPA2D::Specialization<D1D, Q1D>::Add();
+         GradDiagPA2D::Specialization<D1D, Q1D>::Add();
+      }
+      else if constexpr (DIM == 3)
+      {
+         AddMultGradPA3D::Specialization<D1D, Q1D>::Add();
+         GradDiagPA3D::Specialization<D1D, Q1D>::Add();
+      }
+   }
+
+   void AssembleMF(const FiniteElementSpace &fes) override;
+
+   void AddMultMF(const Vector &x, Vector &y) const override;
 
 protected:
    const IntegrationRule* GetDefaultIntegrationRule(
@@ -430,7 +476,8 @@ protected:
 
 
 /** This class is used to assemble the convective form of the nonlinear term
-    arising in the Navier-Stokes equations $(u \cdot \nabla v, w )$ */
+    arising in the Navier-Stokes equations $(u \cdot \nabla v, w )$.
+    Partial assembly is not supported; use VectorConvectionNLFIntegrator. */
 class ConvectiveVectorConvectionNLFIntegrator :
    public VectorConvectionNLFIntegrator
 {
@@ -448,12 +495,20 @@ public:
                             ElementTransformation &trans,
                             const Vector &elfun,
                             DenseMatrix &elmat) override;
+
+   using NonlinearFormIntegrator::AssemblePA;
+   void AssemblePA(const FiniteElementSpace &fes) override;
+   void AssembleGradPA(const Vector &x, const FiniteElementSpace &fes) override;
+   void AddMultPA(const Vector &x, Vector &y) const override;
+   void AddMultGradPA(const Vector &x, Vector &y) const override;
+   void AssembleGradDiagonalPA(Vector &diag) const override;
 };
 
 
 /** This class is used to assemble the skew-symmetric form of the nonlinear term
     arising in the Navier-Stokes equations
-    $.5*(u \cdot \nabla v, w ) - .5*(u \cdot \nabla w, v )$ */
+    $.5*(u \cdot \nabla v, w ) - .5*(u \cdot \nabla w, v )$.
+    Partial assembly is not supported; use VectorConvectionNLFIntegrator. */
 class SkewSymmetricVectorConvectionNLFIntegrator :
    public VectorConvectionNLFIntegrator
 {
@@ -471,6 +526,13 @@ public:
                             ElementTransformation &trans,
                             const Vector &elfun,
                             DenseMatrix &elmat) override;
+
+   using NonlinearFormIntegrator::AssemblePA;
+   void AssemblePA(const FiniteElementSpace &fes) override;
+   void AssembleGradPA(const Vector &x, const FiniteElementSpace &fes) override;
+   void AddMultPA(const Vector &x, Vector &y) const override;
+   void AddMultGradPA(const Vector &x, Vector &y) const override;
+   void AssembleGradDiagonalPA(Vector &diag) const override;
 };
 
 }

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/benchmarks/makefile

@@ -21,7 +21,7 @@ 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
+            bench_nlvc bench_tmop bench_vector bench_virtuals
 PAR_TESTS = 
 ifeq ($(MFEM_USE_MPI),NO)
    TESTS = $(SEQ_TESTS)

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_vecdiv.cpp
+  fem/test_pa_nlvc.cpp
   fem/test_pa_simplices.cpp
   fem/test_particleset.cpp
   fem/test_pgridfunc_save_serial.cpp

tests/unit/fem/test_pa_diagonal.cpp

@@ -320,7 +320,8 @@ double test_vdiag_pa(int dim, int order)
 }
 
 TEST_CASE("Vector Mass Diagonal PA",
-          "[AssembleDiagonal][PartialAssembly][VectorPA][VectorDiagonalPA][VectorMassPA][CUDA]")
+          "[AssembleDiagonal][PartialAssembly]"
+          "[VectorPA][VectorDiagonalPA][VectorMassPA][GPU]")
 {
    const auto DIM = GENERATE(2, 3);
    const auto P = GENERATE(1, 2, 3);
@@ -328,8 +329,48 @@ TEST_CASE("Vector Mass Diagonal PA",
    REQUIRE(test_vdiag_pa<VectorMassIntegrator>(DIM,P) == MFEM_Approx(0.0));
 }
 
+TEST_CASE("Vector Mass Diagonal PA accumulate",
+          "[AssembleDiagonal][PartialAssembly]"
+          "[VectorPA][VectorDiagonalPA][VectorMassPA][GPU]")
+{
+   const auto dim = GENERATE(2, 3);
+   const auto order = GENERATE(1, 2, 3);
+   CAPTURE(dim, order);
+
+   Mesh mesh;
+   if (dim == 2)
+   {
+      mesh = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL);
+   }
+   else
+   {
+      mesh = Mesh::MakeCartesian3D(1, 1, 1, Element::HEXAHEDRON);
+   }
+
+   H1_FECollection fec(order, dim);
+   FiniteElementSpace fes(&mesh, &fec, dim);
+
+   VectorMassIntegrator integ;
+   integ.AssemblePA(fes);
+
+   const int n = fes.GetVSize();
+   Vector base(n), from_zero(n), from_nonzero(n);
+   base.Randomize(1);
+
+   from_zero = 0.0;
+   integ.AssembleDiagonalPA(from_zero);
+
+   from_nonzero = base;
+   integ.AssembleDiagonalPA(from_nonzero);
+
+   from_nonzero -= base;
+   from_nonzero -= from_zero;
+   REQUIRE(from_nonzero.Normlinf() == MFEM_Approx(0.0));
+}
+
 TEST_CASE("Vector Diffusion Diagonal PA",
-          "[AssembleDiagonal][PartialAssembly][VectorPA][VectorDiagonalPA][VectorDiffusionPA][CUDA]")
+          "[AssembleDiagonal][PartialAssembly]"
+          "[VectorPA][VectorDiagonalPA][VectorDiffusionPA][GPU]")
 {
    const auto DIM = GENERATE(2, 3);
    const auto P = GENERATE(1, 2, 3);
@@ -337,6 +378,46 @@ TEST_CASE("Vector Diffusion Diagonal PA",
    REQUIRE(test_vdiag_pa<VectorDiffusionIntegrator>(DIM,P) == MFEM_Approx(0.0));
 }
 
+TEST_CASE("Elasticity Diagonal PA",
+          "[AssembleDiagonal][PartialAssembly][ElasticityPA][GPU]")
+{
+   const auto dim = GENERATE(2, 3);
+   const auto order = GENERATE(1, 2);
+   CAPTURE(dim, order);
+
+   Mesh mesh;
+   if (dim == 2)
+   {
+      mesh = Mesh::MakeCartesian2D(2, 2, Element::QUADRILATERAL, 0, 1.0, 1.0);
+   }
+   else
+   {
+      mesh = Mesh::MakeCartesian3D(2, 2, 2, Element::HEXAHEDRON, 1.0, 1.0, 1.0);
+   }
+
+   H1_FECollection fec(order, dim);
+   FiniteElementSpace fes(&mesh, &fec, dim, Ordering::byNODES);
+
+   ConstantCoefficient lambda(1.0), mu(1.0);
+
+   BilinearForm form_pa(&fes);
+   form_pa.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   form_pa.AddDomainIntegrator(new ElasticityIntegrator(lambda, mu));
+   form_pa.Assemble();
+
+   BilinearForm form_fa(&fes);
+   form_fa.AddDomainIntegrator(new ElasticityIntegrator(lambda, mu));
+   form_fa.Assemble();
+   form_fa.Finalize();
+
+   Vector diag_pa(fes.GetVSize()), diag_fa(fes.GetVSize());
+   form_pa.AssembleDiagonal(diag_pa);
+   form_fa.SpMat().GetDiag(diag_fa);
+
+   diag_fa -= diag_pa;
+   REQUIRE(diag_fa.Normlinf() == MFEM_Approx(0.0));
+}
+
 TEST_CASE("Hcurl/Hdiv diagonal PA",
           "[GPU][PartialAssembly][AssembleDiagonal]")
 {

tests/unit/fem/test_pa_kernels.cpp

@@ -17,6 +17,8 @@
 #include "unit_tests.hpp"
 #include "mfem.hpp"
 
+#include "fem/integ/bilininteg_vecmass_pa.hpp" // IWYU pragma: keep
+
 using namespace mfem;
 
 namespace pa_kernels
@@ -443,13 +445,76 @@ real_t test_pa_vector_integrator(int dim, int sdim)
    return y_fa.Norml2();
 }
 
+void test_pa_vector_mass_kernels(int dim, int p, int q_order)
+{
+   CAPTURE(dim, p, q_order);
+
+   Mesh mesh = MakeCartesianNonaligned(dim, 2);
+   mesh.SetCurvature(p, false, dim);
+   const IntegrationRule &ir =
+      IntRules.Get(mesh.GetTypicalElementGeometry(), q_order);
+
+   H1_FECollection fec(p, dim);
+   FiniteElementSpace fes(&mesh, &fec, dim);
+
+   GridFunction x(&fes), y_fa(&fes), y_pa(&fes);
+   x.Randomize(1);
+
+   BilinearForm blf_fa(&fes);
+   blf_fa.SetAssemblyLevel(AssemblyLevel::LEGACY);
+   // blf_fa will take ownership of integ_fa
+   auto *integ_fa = new VectorMassIntegrator;
+   integ_fa->SetIntRule(&ir);
+   blf_fa.AddDomainIntegrator(integ_fa);
+   blf_fa.Assemble();
+   blf_fa.Finalize();
+   blf_fa.Mult(x, y_fa);
+
+   Vector diag_fa(fes.GetVSize());
+   blf_fa.SpMat().GetDiag(diag_fa);
+
+   BilinearForm blf_pa(&fes);
+   blf_pa.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   // blf_pa will take ownership of integ_pa
+   auto *integ_pa = new VectorMassIntegrator;
+   integ_pa->SetIntRule(&ir);
+   blf_pa.AddDomainIntegrator(integ_pa);
+   blf_pa.Assemble();
+   blf_pa.Mult(x, y_pa);
+
+   Vector diag_pa(fes.GetVSize());
+   blf_pa.AssembleDiagonal(diag_pa);
+
+   y_fa -= y_pa;
+   diag_fa -= diag_pa;
+
+   REQUIRE(y_fa.Norml2() == MFEM_Approx(0.0));
+   REQUIRE(diag_fa.Normlinf() == MFEM_Approx(0.0));
+}
+
 TEST_CASE("PA Vector Mass",
           "[PartialAssembly][VectorPA][VectorMassPA][GPU]")
 {
-   const auto DIM = GENERATE(2, 3);
-   CAPTURE(DIM);
-   REQUIRE(test_pa_vector_integrator<VectorMassIntegrator>(DIM, DIM)
-           == MFEM_Approx(0.0));
+   using Mass = VectorMassIntegrator;
+
+   SECTION("built-in specializations")
+   {
+      const auto DIM = GENERATE(2, 3);
+      CAPTURE(DIM);
+      REQUIRE(test_pa_vector_integrator<Mass>(DIM, DIM) == MFEM_Approx(0.0));
+   }
+
+   SECTION("user specializations")
+   {
+      constexpr int user_dim = 2, user_d1d = 2, user_q1d = 9;
+      using AddMult = VectorMassIntegrator::VectorMassAddMultPA;
+      using Diag = VectorMassIntegrator::VectorMassAssembleDiagonalPA;
+      AddMult::Specialization<user_dim, user_d1d, user_q1d>::Add();
+      Diag::Specialization<user_dim, user_q1d>::Add();
+
+      constexpr int p = 1, q_order = 2*user_q1d - 1;
+      test_pa_vector_mass_kernels(user_dim, p, q_order);
+   }
 }
 
 TEST_CASE("PA Vector Diffusion",
@@ -462,7 +527,7 @@ TEST_CASE("PA Vector Diffusion",
 }
 
 TEST_CASE("PA Vector Diffusion 2D/3D",
-          "[PartialAssembly][VectorPA][VectorDiffusionPA][CUDA]")
+          "[PartialAssembly][VectorPA][VectorDiffusionPA][GPU]")
 {
    const int DIM = 2, SDIM = 3;
    CAPTURE(DIM, SDIM);

tests/unit/fem/test_pa_nlvc.cpp

@@ -0,0 +1,135 @@
+// 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/integ/nonlininteg_vecconvection_pa.hpp" // IWYU pragma: keep
+#include "fem/integ/nonlininteg_vecconvection_pa_diag.hpp" // IWYU pragma: keep
+#include "fem/integ/nonlininteg_vecconvection_pa_grad.hpp" // IWYU pragma: keep
+#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 auto p_base = {1, 2}, p_extra = {3, 4};
+   const auto p = MFEM_GENERATE_RANGES(p_base, p_extra);
+
+   using Grad = QuadratureInterpolator::GradKernels;
+   using NLVC = VectorConvectionNLFIntegrator;
+   static const auto specializations =
+      (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, 5>::Add(),
+       Grad::Specialization<3, QVectorLayout::byNODES, false, 3, 4, 9>::Add(),
+       // User specialization example for p=4 (D1D=5, Q1D=9) in 3D
+       NLVC::AddSpecialization<3, 5, 9>(),
+       true);
+   MFEM_CONTRACT_VAR(specializations);
+
+   SECTION("2D")
+   {
+      const auto meshs = { "../../data/inline-quad.mesh" };
+      const auto extra = { "../../data/star-q2.mesh",
+                           "../../data/star-q3.mesh",
+                           "../../data/rt-2d-q3.mesh",
+                           "../../data/periodic-square.mesh"
+                         };
+      test_nl_convection_pa_grad<2>(MFEM_GENERATE_RANGES(meshs, extra), p);
+   }
+
+   SECTION("3D")
+   {
+      const auto meshs = { "../../data/inline-hex.mesh" };
+      const auto extra = { "../../data/fichera.mesh",
+                           "../../data/beam-hex.mesh",
+                           "../../data/toroid-hex.mesh",
+                           "../../data/fichera-q2.mesh",
+                           "../../data/fichera-q3.mesh",
+                           "../../data/periodic-cube.mesh"
+                         };
+      test_nl_convection_pa_grad<3>(MFEM_GENERATE_RANGES(meshs, extra), p);
+   }
+}
+
+} // namespace pa_kernels

tests/unit/fem/test_pa_vecdiv.cpp

@@ -0,0 +1,153 @@
+// 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/integ/bilininteg_vecdiv_pa.hpp" // IWYU pragma: keep
+#include "fem/qinterp/grad.hpp" // IWYU pragma: keep
+
+#include <algorithm>
+#include <utility>
+
+using namespace mfem;
+
+namespace pa_kernels
+{
+
+template <typename INTEGRATOR, bool TRANSPOSE>
+void pa_mixed_test(FiniteElementSpace &fes1,
+                   FiniteElementSpace &fes2,
+                   const IntegrationRule &ir)
+{
+   MixedBilinearForm bform_pa(&fes1, &fes2);
+   if constexpr (TRANSPOSE)
+   {
+      auto *integ = new TransposeIntegrator(new INTEGRATOR);
+      integ->SetIntRule(&ir);
+      bform_pa.AddDomainIntegrator(integ);
+   }
+   else
+   {
+      auto *integ = new INTEGRATOR;
+      integ->SetIntRule(&ir);
+      bform_pa.AddDomainIntegrator(integ);
+   }
+   bform_pa.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   bform_pa.Assemble();
+
+   MixedBilinearForm bform_fa(&fes1, &fes2);
+   if constexpr (TRANSPOSE)
+   {
+      auto *integ = new TransposeIntegrator(new INTEGRATOR);
+      integ->SetIntRule(&ir);
+      bform_fa.AddDomainIntegrator(integ);
+   }
+   else
+   {
+      auto *integ = new INTEGRATOR;
+      integ->SetIntRule(&ir);
+      bform_fa.AddDomainIntegrator(integ);
+   }
+   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 vp, int sp)
+{
+   CAPTURE(filename, DIM, vp, sp);
+
+   Mesh mesh(filename);
+   MFEM_VERIFY(mesh.Dimension() == DIM, "Mesh dimension mismatch");
+
+   // Vector
+   H1_FECollection vfec(vp, DIM);
+   FiniteElementSpace vfes(&mesh, &vfec, DIM);
+
+   // Scalar
+   H1_FECollection sfec(sp, DIM);
+   FiniteElementSpace sfes(&mesh, &sfec);
+
+   // Shared-memory PA kernels require q1d >= max(trial_d1d, test_d1d)
+   const auto &trial_fe = *vfes.GetTypicalFE();
+   const auto &test_fe = *sfes.GetTypicalFE();
+   const auto &Trans = *mesh.GetTypicalElementTransformation();
+   int order = Trans.OrderGrad(&trial_fe) + test_fe.GetOrder() + Trans.OrderJ();
+   const int min_q1d = std::max(vp, sp) + 1;
+   order = std::max(order, 2 * min_q1d - 1);
+   const IntegrationRule &ir = IntRules.Get(trial_fe.GetGeomType(), order);
+
+   pa_mixed_test<VectorDivergenceIntegrator, false>(vfes, sfes, ir);
+   pa_mixed_test<VectorDivergenceIntegrator, true>(sfes, vfes, ir);
+}
+
+TEST_CASE("VecDivPA", "[PartialAssembly][VecDivPA][GPU]")
+{
+   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();
+
+      using VDiv = VectorDivergenceIntegrator::VectorDivergenceAddMultPA;
+      using VDivT = VectorDivergenceIntegrator::VectorDivergenceAddMultTransposePA;
+      VDiv::Specialization<2, 2, 3, 3>::Add();
+      VDivT::Specialization<2, 2, 3, 3>::Add();
+   }
+
+   // Vector (vp) and scalar (sp) space orders
+   const auto vp_base = {1, 2}, vp_extra = {3, 4};
+   const auto sp_base = {1, 2}, sp_extra = {3, 4};
+   const auto vp = MFEM_GENERATE_RANGES(vp_base, vp_extra);
+   const auto sp = MFEM_GENERATE_RANGES(sp_base, sp_extra);
+
+   SECTION("2D")
+   {
+      const auto meshs = { "../../data/inline-quad.mesh" };
+      const auto extra = { "../../data/star-q2.mesh",
+                           "../../data/star-q3.mesh",
+                           "../../data/rt-2d-q3.mesh",
+                           "../../data/periodic-square.mesh"
+                         };
+      test_pa_divergence<2>(MFEM_GENERATE_RANGES(meshs, extra), vp, sp);
+   }
+
+   SECTION("3D")
+   {
+      const auto meshs = { "../../data/inline-hex.mesh" };
+      const auto extra = { "../../data/fichera.mesh",
+                           "../../data/beam-hex.mesh",
+                           "../../data/toroid-hex.mesh",
+                           "../../data/fichera-q2.mesh",
+                           "../../data/fichera-q3.mesh",
+                           "../../data/periodic-cube.mesh"
+                         };
+      test_pa_divergence<3>(MFEM_GENERATE_RANGES(meshs, extra), vp, sp);
+   }
+}
+
+} // namespace pa_kernels

tests/unit/unit_tests.hpp

@@ -36,4 +36,11 @@ inline Approx MFEM_Approx(double val,
    return Approx(val).margin(abs_tol).epsilon(rel_tol);
 }
 
+/** @brief Generate values from @a base, and also from @a extra if the
+    command line '--all' option is provided. */
+#define MFEM_GENERATE_RANGES(base, extra) \
+   (!launch_all_non_regression_tests \
+    ? GENERATE_COPY(from_range(base)) \
+    : GENERATE_COPY(from_range(base), from_range(extra)))
+
 #endif