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/complex_fem.cpp

@@ -237,81 +237,6 @@ ComplexGridFunction::ProjectBdrCoefficientTangent(VectorCoefficient
    gfi->SyncAliasMemory(*this);
 }
 
-real_t
-ComplexGridFunction::ComputeLpError(const real_t p,
-                                    Coefficient &exsolr,
-                                    Coefficient &exsoli,
-                                    Coefficient *weight,
-                                    const IntegrationRule *irs[],
-                                    const Array<int> *elems) const
-{
-   real_t error = 0.0;
-   const FiniteElement *fe;
-   ElementTransformation *T;
-   Vector valsr;
-   Vector valsi;
-
-   const GridFunction& gf_r = real();
-   const GridFunction& gf_i = imag();
-
-   for (int i = 0; i < fes->GetNE(); i++)
-   {
-      if (elems != NULL && (*elems)[i] == 0) { continue; }
-      fe = fes->GetFE(i);
-      const IntegrationRule *ir;
-      if (irs)
-      {
-         ir = irs[fe->GetGeomType()];
-      }
-      else
-      {
-         int intorder = 2*fe->GetOrder() + 3;
-         ir = &(IntRules.Get(fe->GetGeomType(), intorder));
-      }
-      real_t elem_error = 0.0;
-      gf_r.GetValues(i, *ir, valsr);
-      gf_i.GetValues(i, *ir, valsi);
-      T = fes->GetElementTransformation(i);
-      for (int j = 0; j < ir->GetNPoints(); j++)
-      {
-         const IntegrationPoint &ip = ir->IntPoint(j);
-         T->SetIntPoint(&ip);
-         real_t diffr = valsr(j) - exsolr.Eval(*T, ip);
-         real_t diffi = valsi(j) - exsoli.Eval(*T, ip);
-         real_t diff = hypot(diffr, diffi);
-         if (p < infinity())
-         {
-            diff = pow(diff, p);
-            if (weight)
-            {
-               diff *= weight->Eval(*T, ip);
-            }
-            elem_error += ip.weight * T->Weight() * diff;
-         }
-         else
-         {
-            if (weight)
-            {
-               diff *= weight->Eval(*T, ip);
-            }
-            error = std::max(error, diff);
-         }
-      }
-      if (p < infinity())
-      {
-         // negative quadrature weights may cause the error to be negative
-         error += fabs(elem_error);
-      }
-   }
-
-   if (p < infinity())
-   {
-      error = pow(error, 1./p);
-   }
-
-   return error;
-}
-
 void ComplexGridFunction::Save(std::ostream &os) const
 {
    os << "ComplexGridFunction\n";

fem/complex_fem.hpp

@@ -166,75 +166,6 @@ public:
       return sqrt(err_r * err_r + err_i * err_i);
    }
 
-   /// @brief Returns Max|u_ex - u_h| error for complex-valued H1 or L2 elements
-   ///
-   /// Compute the $L_\infty$ error across the entire domain.
-   ///
-   /// @param[in] exsolr  Coefficient object reproducing the real part of the
-   ///                    anticipated values of the scalar field, Re(u_ex).
-   /// @param[in] exsoli  Coefficient object reproducing the imaginary part of
-   ///                    the anticipated values of the scalar field, Im(u_ex).
-   /// @param[in] irs        Optional pointer to an array of custom integration
-   ///                       rules e.g. higher order than the default rules. If
-   ///                       present the array will be indexed by
-   ///                       Geometry::Type.
-   ///
-   /// @note Uses ComputeLpError internally. See the ComputeLpError
-   ///       documentation for generalizations of this error computation.
-   ///
-   /// @note If an array of integration rules is provided through @a irs, be
-   ///       sure to include valid rules for each element type that may occur
-   ///       in the list of elements.
-   ///
-   virtual real_t ComputeMaxError(Coefficient &exsolr,
-                                  Coefficient &exsoli,
-                                  const IntegrationRule *irs[] = NULL) const
-   {
-      return ComputeLpError(infinity(), exsolr, exsoli, NULL, irs);
-   }
-
-   /// @brief Returns ||u_ex - u_h||_Lp for complex-valued H1 or L2 elements
-   ///
-   /// Computes:
-   ///    $$(\sum_{elems} \int_{elem} w \, |u_{ex} - u_h|^p)^{1/p}$$
-   /// Where:
-   ///    $$|u_{ex} - u_h| = \sqrt{Re(u_{ex} - u_h)^2 + Im(u_{ex} - u_h)^2}$$
-   ///
-   /// @param[in] p       Real value indicating the exponent of the $L^p$ norm.
-   ///                    To avoid domain errors p should have a positive value,
-   ///                    either finite or infinite.
-   /// @param[in] exsolr  Coefficient object reproducing the real part of the
-   ///                    anticipated values of the scalar field, Re(u_ex).
-   /// @param[in] exsoli  Coefficient object reproducing the imaginary part of
-   ///                    the anticipated values of the scalar field, Im(u_ex).
-   /// @param[in] weight  Optional pointer to a Coefficient object reproducing
-   ///                    a weighting function, w.
-   /// @param[in] irs     Optional pointer to an array of custom integration
-   ///                    rules e.g. higher order than the default rules. If
-   ///                    present the array will be indexed by Geometry::Type.
-   /// @param[in] elems   Optional pointer to a marker array, with a length
-   ///                    equal to the number of local elements, indicating
-   ///                    which elements to integrate over. Only those elements
-   ///                    corresponding to non-zero entries in @a elems will
-   ///                    contribute to the computed L2 error.
-   ///
-   /// @note If an array of integration rules is provided through @a irs, be
-   ///       sure to include valid rules for each element type that may occur
-   ///       in the list of elements.
-   ///
-   /// @note Quadratures with negative weights (as in some simplex integration
-   ///       rules in MFEM) can produce negative integrals even with
-   ///       non-negative integrands. To avoid returning negative errors this
-   ///       function uses the absolute values of the element-wise integrals.
-   ///       This may lead to results which are not entirely consistent with
-   ///       such integration rules.
-   virtual real_t ComputeLpError(const real_t p,
-                                 Coefficient &exsolr,
-                                 Coefficient &exsoli,
-                                 Coefficient *weight = NULL,
-                                 const IntegrationRule *irs[] = NULL,
-                                 const Array<int> *elems = NULL) const;
-
    /// Save the ComplexGridFunction to an output stream.
    virtual void Save(std::ostream &out) const;
 

fem/datacollection.cpp

@@ -1181,14 +1181,12 @@ void ParaViewDataCollection::SaveGFieldVTU(std::ostream &os, int ref_,
    DenseMatrix vval, pmat;
    std::vector<char> buf;
    int vec_dim = it->second->VectorDim();
-   int map_type = it->second->FESpace()->GetTypicalFE()->GetMapType();
    os << "<DataArray type=\"" << GetDataTypeString()
       << "\" Name=\"" << it->first
       << "\" NumberOfComponents=\"" << vec_dim << "\" "
       << VTKComponentLabels(vec_dim) << " "
       << "format=\"" << GetDataFormatString() << "\" >" << '\n';
-   if (vec_dim == 1 && (map_type == FiniteElement::VALUE ||
-                        map_type == FiniteElement::INTEGRAL))
+   if (vec_dim == 1)
    {
       for (int i = 0; i < mesh->GetNE(); i++)
       {

fem/datacollection.hpp

@@ -25,35 +25,21 @@
 namespace mfem
 {
 
-/// Lightweight adaptor over an std::map from type K to type to V
-template<typename K, typename V,
-         typename = typename std::enable_if<std::is_default_constructible<V>::value>::type>
-class GenericFieldMap
+/// Lightweight adaptor over an std::map from strings to pointer to T
+template<typename T>
+class NamedFieldsMap
 {
-private:
-   static constexpr bool ValueIsPointer = std::is_pointer<V>::value;
-
 public:
-   typedef std::map<K, V> MapType;
+   typedef std::map<std::string, T*> MapType;
    typedef typename MapType::iterator iterator;
    typedef typename MapType::const_iterator const_iterator;
 
-   /// Register field @a field with name @a key
-   /// Only enabled if the template parameter V is not a pointer
-   template<typename = std::enable_if<!ValueIsPointer, bool>>
-   void Register(const K& key, V field)
+   /// Register field @a field with name @a fname
+   /** Replace existing field associated with @a fname (and optionally
+       delete associated pointer if @a own_data is true) */
+   void Register(const std::string& fname, T* field, bool own_data)
    {
-      field_map[key] = field;
-   }
-
-   /// Register field @a field with name @a key
-   /** Replace existing field associated with @a key (and optionally
-       delete associated pointer if @a own_data is true).
-       Only enabled if the template parameter V is a pointer*/
-   template<typename = std::enable_if<ValueIsPointer, bool>>
-   void Register(const K& key, V field, bool own_data)
-   {
-      V& ref = field_map[key];
+      T*& ref = field_map[fname];
       if (own_data)
       {
          delete ref; // if newly allocated -> ref is null -> OK
@@ -61,40 +47,23 @@ public:
       ref = field;
    }
 
-   /// Unregister association between field @a field and name @a key
-   /// Only enabled if the template parameter V is not a pointer
-   template<typename = std::enable_if<!ValueIsPointer, bool>>
-   void Deregister(const K& key)
+   /// Unregister association between field @a field and name @a fname
+   /** Optionally delete associated pointer if @a own_data is true */
+   void Deregister(const std::string& fname, bool own_data)
    {
-      iterator it = field_map.find(key);
-      if ( it != field_map.end() )
-      {
-         field_map.erase(it);
-      }
-   }
-
-   /// Unregister association between field @a field and name @a key
-   /** Optionally delete associated pointer if @a own_data is true.
-       Only enabled if the template parameter V is a pointer */
-   template<typename = std::enable_if<ValueIsPointer, bool>>
-   void Deregister(const K& key, bool own_data)
-   {
-      iterator it = field_map.find(key);
+      iterator it = field_map.find(fname);
       if ( it != field_map.end() )
       {
          if (own_data)
          {
             delete it->second;
-            it->second = nullptr;
          }
          field_map.erase(it);
       }
    }
 
    /// Clear all associations between names and fields
-   /** Delete associated pointers when @a own_data is true.
-       Only enabled if the template parameter V is a pointer */
-   template<typename = std::enable_if<ValueIsPointer, bool>>
+   /** Delete associated pointers when @a own_data is true */
    void DeleteData(bool own_data)
    {
       for (iterator it = field_map.begin(); it != field_map.end(); ++it)
@@ -107,37 +76,22 @@ public:
       }
    }
 
-   /// Predicate to check if a field is associated with name @a key
-   bool Has(const K& key) const
+   /// Predicate to check if a field is associated with name @a fname
+   bool Has(const std::string& fname) const
    {
-      return field_map.find(key) != field_map.end();
+      return field_map.find(fname) != field_map.end();
    }
 
-   /// Get a pointer to the field associated with name @a key
-   /** @return Field associated with @a key or NULL,
-       if value is pointer and key not found */
-   V Get(const K& key) const
+   /// Get a pointer to the field associated with name @a fname
+   /** @return Pointer to field associated with @a fname or NULL */
+   T* Get(const std::string& fname) const
    {
-      const_iterator it = field_map.find(key);
-      if (it != field_map.end())
-      {
-         return it->second;
-      }
-      else
-      {
-         if constexpr (ValueIsPointer)
-         {
-            return nullptr;
-         }
-         else
-         {
-            return V(); // Return default-constructed value for non-pointer types
-         }
-      }
+      const_iterator it = field_map.find(fname);
+      return it != field_map.end() ? it->second : NULL;
    }
 
    /// Returns a const reference to the underlying map
-   const MapType &GetMap() const { return field_map; }
+   const MapType& GetMap() const { return field_map; }
 
    /// Returns the number of registered fields
    int NumFields() const { return field_map.size(); }
@@ -152,24 +106,21 @@ public:
    /// Returns an end const iterator to the registered fields
    const_iterator end() const { return field_map.end(); }
 
-   /// Returns an iterator to the field @a key
-   iterator find(const K& key)
-   { return field_map.find(key); }
+   /// Returns an iterator to the field @a fname
+   iterator find(const std::string& fname)
+   { return field_map.find(fname); }
 
-   /// Returns a const iterator to the field @a key
-   const_iterator find(const K& key) const
-   { return field_map.find(key); }
+   /// Returns a const iterator to the field @a fname
+   const_iterator find(const std::string& fname) const
+   { return field_map.find(fname); }
 
-   /// Clears the map of registered fields
+   /// Clears the map of registered fields without reclaiming memory
    void clear() { field_map.clear(); }
 
 protected:
    MapType field_map;
 };
 
-/// Lightweight adaptor over an std::map from strings to pointer to T
-template<typename T>
-using NamedFieldsMap = GenericFieldMap<std::string, T*>;
 
 /** A class for collecting finite element data that is part of the same
     simulation. Currently, this class groups together grid functions (fields),

fem/dfem/util.hpp

@@ -671,20 +671,6 @@ public:
                     MPI_COMM_WORLD);
    }
 
-   Operator& GetGradient(const Vector &x0) const override
-   {
-      x = x0;
-      f.UseDevice(x.UseDevice());
-      xpev.UseDevice(x.UseDevice());
-
-      op.Mult(x, f);
-      const real_t xnorm_local = x.Norml2();
-      MPI_Allreduce(&xnorm_local, &xnorm, 1, MPITypeMap<real_t>::mpi_type, MPI_SUM,
-                    MPI_COMM_WORLD);
-
-      return const_cast<FDJacobian&>(*this);
-   }
-
    void Mult(const Vector &v, Vector &y) const override
    {
       // See [1] for choice of eps.
@@ -739,11 +725,11 @@ public:
 
 private:
    const Operator &op;
-   mutable Vector x, f;
+   Vector x, f;
    mutable Vector xpev;
    real_t lambda = 1.0e-6;
    real_t fixed_eps;
-   mutable real_t xnorm;
+   real_t xnorm;
 };
 
 /// @brief Find the index of a field descriptor in a vector of field descriptors.

fem/gslib.cpp

@@ -76,9 +76,6 @@ namespace gslib
 #ifndef GSLIB_RELEASE_VERSION //gslib v1.0.7
 #define GSLIB_RELEASE_VERSION 10007
 #endif
-static_assert(std::is_same_v<uint,unsigned int>,
-              "GSLIB's integer-type, 'uint', defined in gslib.h, must be the same as 'unsigned int'!");
-
 extern "C" {
    struct hash_data_3
    {

fem/gslib.hpp

@@ -326,22 +326,22 @@ protected:
    void findptsedge_setup_2(DEV_STRUCT &devs,
                             const double *const elx[2],
                             const unsigned n,
-                            const unsigned int nel,
+                            const uint nel,
                             const unsigned m,
                             const double bbox_rel_size_inc,
-                            const unsigned int local_hash_size,
-                            const unsigned int global_hash_size,
+                            const uint local_hash_size,
+                            const uint global_hash_size,
                             const Vector *aabb_sz_inc);
 
    /// Preprocess 3D surface mesh needed for FindPoints.
    void findptssurf_setup_3(DEV_STRUCT &devs,
                             const double *const elx[3],
                             const unsigned n,
-                            const unsigned int nel,
+                            const uint nel,
                             const unsigned m,
                             const double bbox_rel_size_inc,
-                            const unsigned int local_hash_size,
-                            const unsigned int global_hash_size,
+                            const uint local_hash_size,
+                            const uint global_hash_size,
                             const int rD,
                             const Vector *aabb_sz_inc);
 

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.hpp

@@ -224,9 +224,6 @@ public:
    /** @see GetGradient(const Vector &) */
    Operator &GetGradient(const Vector &x, bool finalize) const;
 
-   /// Suppress a warning about hiding overloaded virtual function.
-   using Operator::GetGradient;
-
    /// Update the NonlinearForm to propagate updates of the associated FE space.
    /** After calling this method, the essential boundary conditions need to be
        set again. */

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;
 };
 
 }

linalg/CMakeLists.txt

@@ -27,7 +27,6 @@ list(APPEND SRCS
   handle.cpp
   matrix.cpp
   mma.cpp
-  multivector.cpp
   ode.cpp
   operator.cpp
   ordering.cpp
@@ -64,7 +63,6 @@ list(APPEND HDRS
   linalg.hpp
   matrix.hpp
   mma.hpp
-  multivector.hpp
   ode.hpp
   operator.hpp
   ordering.hpp

linalg/cudss.cpp

@@ -15,20 +15,10 @@
 
 #ifdef MFEM_USE_CUDSS
 
-#if CUDSS_VERSION >= 800
 #ifdef MFEM_USE_SINGLE
-#define CUDSS_REAL_T CUDSS_R_32F
+#define CUDA_REAL_T CUDA_R_32F
 #else
-#define CUDSS_REAL_T CUDSS_R_64F
-#endif
-#define CUDSS_INT_T CUDSS_R_32I
-#else
-#ifdef MFEM_USE_SINGLE
-#define CUDSS_REAL_T CUDA_R_32F
-#else
-#define CUDSS_REAL_T CUDA_R_64F
-#endif
-#define CUDSS_INT_T CUDA_R_32I
+#define CUDA_REAL_T CUDA_R_64F
 #endif
 
 // Define a cuDSS error check macro, MFEM_CUDSS_CHECK(x), where x returns/is of
@@ -75,13 +65,8 @@ CuDSSSolver::CuDSSSolver(MPI_Comm comm_) : mpi_comm(comm_)
 #endif
    MFEM_CUDSS_CHECK(cudssSetCommLayer(handle, comm_lib));
 
-#if CUDSS_VERSION >= 800
-   MFEM_CUDSS_CHECK(cudssDataSet(handle, solverData, CUDSS_DATA_COMM_HOST,
-                                 &mpi_comm, sizeof(MPI_Comm *)));
-#else
    MFEM_CUDSS_CHECK(cudssDataSet(handle, solverData, CUDSS_DATA_COMM,
                                  &mpi_comm, sizeof(MPI_Comm *)));
-#endif
 }
 #endif // MFEM_USE_MPI
 
@@ -272,19 +257,11 @@ void CuDSSSolver::SetMatrixCuDSS(int *csr_offsets, int *csr_columns,
          CuMemcpyDtoD(csr_offsets_d, csr_offsets, (n_loc + 1) * sizeof(int));
          CuMemcpyDtoD(csr_columns_d, csr_columns, nnz * sizeof(int));
 
-#if CUDSS_VERSION >= 800
          MFEM_CUDSS_CHECK(
             cudssMatrixCreateCsr(
                Ac.get(), n_global, n_global, nnz, csr_offsets_d, NULL,
-               csr_columns_d, csr_values_d, CUDSS_INT_T, CUDSS_INT_T, CUDSS_REAL_T,
-               mat_type, mview, CUDSS_BASE_ZERO));
-#else
-         MFEM_CUDSS_CHECK(
-            cudssMatrixCreateCsr(
-               Ac.get(), n_global, n_global, nnz, csr_offsets_d, NULL,
-               csr_columns_d, csr_values_d, CUDSS_INT_T, CUDSS_REAL_T,
-               mat_type, mview, CUDSS_BASE_ZERO));
-#endif
+               csr_columns_d, csr_values_d, CUDA_R_32I, CUDA_REAL_T, mat_type, mview,
+               CUDSS_BASE_ZERO));
       }
       else    // !reorder_reuse
       {
@@ -292,19 +269,11 @@ void CuDSSSolver::SetMatrixCuDSS(int *csr_offsets, int *csr_columns,
          {
             MFEM_CUDSS_CHECK(cudssMatrixDestroy(*Ac));
          }
-#if CUDSS_VERSION >= 800
          MFEM_CUDSS_CHECK(
             cudssMatrixCreateCsr(
-               Ac.get(), n_global, n_global, nnz, csr_offsets, NULL,
-               csr_columns, csr_values_d, CUDSS_INT_T, CUDSS_INT_T, CUDSS_REAL_T,
-               mat_type, mview, CUDSS_BASE_ZERO));
-#else
-         MFEM_CUDSS_CHECK(
-            cudssMatrixCreateCsr(
-               Ac.get(), n_global, n_global, nnz, csr_offsets, NULL,
-               csr_columns, csr_values_d, CUDSS_INT_T, CUDSS_REAL_T,
-               mat_type, mview, CUDSS_BASE_ZERO));
-#endif
+               Ac.get(), n_global, n_global, nnz, csr_offsets, NULL, csr_columns,
+               csr_values_d, CUDA_R_32I, CUDA_REAL_T, mat_type, mview,
+               CUDSS_BASE_ZERO));
       }
 #ifdef MFEM_USE_MPI
       if (Mpi::IsInitialized())
@@ -365,10 +334,10 @@ void CuDSSSolver::SetNumRHS(int nrhs_) const
       }
       // Create empty RHS and solution vectors
       MFEM_CUDSS_CHECK(cudssMatrixCreateDn(&xc, n_global, nrhs_, n_global, NULL,
-                                           CUDSS_REAL_T, CUDSS_LAYOUT_COL_MAJOR));
+                                           CUDA_REAL_T, CUDSS_LAYOUT_COL_MAJOR));
 
       MFEM_CUDSS_CHECK(cudssMatrixCreateDn(&yc, n_global, nrhs_, n_global, NULL,
-                                           CUDSS_REAL_T, CUDSS_LAYOUT_COL_MAJOR));
+                                           CUDA_REAL_T, CUDSS_LAYOUT_COL_MAJOR));
 
 #ifdef MFEM_USE_MPI
       MFEM_CUDSS_CHECK(cudssMatrixSetDistributionRow1d(xc, row_start, row_end));

linalg/linalg.hpp

@@ -15,7 +15,6 @@
 // Linear algebra header file
 
 #include "vector.hpp"
-#include "multivector.hpp"
 #include "operator.hpp"
 #include "matrix.hpp"
 #include "sparsemat.hpp"

linalg/multivector.cpp

@@ -1,60 +0,0 @@
-// 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 "multivector.hpp"
-
-namespace mfem
-{
-
-MultiVector::MultiVector(const Array<int> &vector_sizes)
-{
-   SetSizes(vector_sizes);
-}
-
-MultiVector::MultiVector(const Array<int> &vector_sizes, MemoryType mt)
-{
-   SetSizes(vector_sizes, mt);
-}
-
-MultiVector::MultiVector(Vector &base, const Array<int> &vector_sizes)
-{
-   MakeRef(base, vector_sizes);
-}
-
-void MultiVector::SetSizes(const Array<int> &vector_sizes)
-{
-   blocks.resize(vector_sizes.Size());
-   for (int i = 0; i < vector_sizes.Size(); i++)
-   {
-      operator[](i).SetSize(vector_sizes[i]);
-   }
-}
-
-void MultiVector::SetSizes(const Array<int> &vector_sizes, MemoryType mt)
-{
-   blocks.resize(vector_sizes.Size());
-   for (int i = 0; i < vector_sizes.Size(); i++)
-   {
-      operator[](i).SetSize(vector_sizes[i], mt);
-   }
-}
-
-void MultiVector::MakeRef(Vector &base, const Array<int> &vector_sizes)
-{
-   blocks.resize(vector_sizes.Size());
-   for (int offset = 0, i = 0; i < vector_sizes.Size(); i++)
-   {
-      blocks[i].emplace<0>(base, offset, vector_sizes[i]);
-      offset += vector_sizes[i];
-   }
-}
-
-} // namespace mfem

linalg/multivector.hpp

@@ -1,198 +0,0 @@
-// 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.
-
-#ifndef MFEM_MULTIVECTOR_HPP
-#define MFEM_MULTIVECTOR_HPP
-
-#include "../general/array.hpp"
-#include "vector.hpp"
-#include <vector>
-#include <array>
-#include <variant>
-
-namespace mfem
-{
-
-/// Class representing an array of Vectors with generally different sizes.
-/** This class is similar to BlockVector with the following two main
-    differences:
-    - the data for the individual Vector blocks does not need to be part of one
-      big contiguous memory allocation;
-    - this class does not inherit from class Vector (as a consequence of the
-      first bullet).
-
-    Internally, each Vector block is represented as either:
-    - (default) a Vector object constructed and owned by this class; this
-      object, in turn, as any Vector object, can own its Memory allocation or
-      refer to a sub-Memory of another Memory object; or
-    - a pointer to an externally allocated Vector or classes derived from
-      Vector. */
-class MultiVector
-{
-private:
-   std::vector<std::variant<Vector,Vector*>> blocks;
-
-public:
-   /// Create an empty MultiVector with zero blocks.
-   MultiVector() = default;
-
-   /** @brief Create a MultiVector with @a num_blocks blocks. The individual
-       Vector blocks are default initialized, i.e. they all have size zero. */
-   MultiVector(int num_blocks)
-      : blocks(num_blocks) { }
-
-   /** @brief Construct a MultiVector with number of blocks and individual block
-       Vector sizes given by @a vector_sizes.
-
-       @note The memory of the individual Vector blocks is NOT initialized. */
-   MultiVector(const Array<int> &vector_sizes);
-
-   /** @brief Construct a MultiVector with number of blocks and individual block
-       Vector sizes given by @a vector_sizes. All Vector blocks use the
-       MemoryType @a mt.
-
-       @note The memory of the individual Vector blocks is NOT initialized. */
-   MultiVector(const Array<int> &vector_sizes, MemoryType mt);
-
-   /** @brief Construct a MultiVector referencing data within a given monolithic
-       Vector @a base.
-
-       With this constructor, the Memory flags of @a base and of the individual
-       Vector blocks may need to be explicitly synchronized when data is moved
-       between host and device. */
-   MultiVector(Vector &base, const Array<int> &vector_sizes);
-
-   /** @brief Construct a MultiVector referencing multiple Vectors given as
-       arguments.
-
-       The VectorTypes reference arguments are expected to be static_cast-able
-       to (Vector &) which is the case if the types are derived from Vector,
-       e.g. HypreParVector, GridFunction, etc.
-
-       With this constructor, operations on individual Vector blocks are
-       performed directly on the objects @a vs. In particular, there is no need
-       to synchronize the Memory flags of @a vs and the ones of the individual
-       Vector blocks when data is moved between host and device. */
-   template <typename... VectorTypes,
-             std::enable_if_t<
-                std::conjunction_v<
-                   std::is_convertible<VectorTypes&,Vector&>...>, bool> = true>
-   MultiVector(VectorTypes &...vs) { MakeRef(vs...); }
-
-   /// Return the number of Vectors in the MultiVector.
-   int NumBlocks() const { return blocks.size(); }
-
-   /** @brief Set the number of Vectors in the MultiVector. Existing Vector
-       blocks will remain unmodified. New Vector blocks will be default
-       initialized, i.e. they all have size zero. */
-   void SetNumBlocks(int num_blocks) { blocks.resize(num_blocks); }
-
-   /// Read-write access to the i-th Vector.
-   inline Vector &operator[](int i);
-
-   /// Read-only access to the i-th Vector.
-   inline const Vector &operator[](int i) const;
-
-   /** @brief Update the MultiVector according to the given @a vector_sizes.
-
-       This method can be used to add or remove blocks. The individual Vector
-       sizes are updated using the method Vector::SetSize(int). */
-   void SetSizes(const Array<int> &vector_sizes);
-
-   /** @brief Update the MultiVector according to the given @a vector_sizes and
-       MemoryType @a mt.
-
-       This method can be used to add or remove blocks. The individual Vector
-       sizes and MemoryType are updated using the method
-       Vector::SetSize(int, MemoryType). */
-   void SetSizes(const Array<int> &vector_sizes, MemoryType mt);
-
-   /** @brief Update the MultiVector to reference data within a given monolithic
-       Vector @a base.
-
-       After calling this method, the Memory flags of @a base and of the
-       individual Vector blocks may need to be explicitly synchronized when data
-       is moved between host and device.*/
-   void MakeRef(Vector &base, const Array<int> &vector_sizes);
-
-   /** @brief Update the @a i-th MultiVector block to reference data within the
-       given monolithic Vector @a base at the given @a offset and with the given
-       @a size.
-
-       After calling this method, the Memory flags of @a base and of the @a i-th
-       Vector block may need to be explicitly synchronized when data is moved
-       between host and device.*/
-   inline void MakeRef(int i, Vector &base, int offset, int size)
-   {
-      blocks[i].emplace<0>(base, offset, size);
-   }
-
-   /** @brief Update the MultiVector to reference multiple Vectors given as
-       arguments.
-
-       The VectorTypes reference arguments are expected to be static_cast-able
-       to (Vector &) which is the case if the types are derived from Vector,
-       e.g. HypreParVector, GridFunction, etc.
-
-       After calling this method, operations on individual Vector blocks are
-       performed directly on the objects @a vs. In particular, there is no need
-       to synchronize the Memory flags of @a vs and the ones of the individual
-       Vector blocks when data is moved between host and device. */
-   template <typename... VectorTypes,
-             std::enable_if_t<
-                std::conjunction_v<
-                   std::is_convertible<VectorTypes&,Vector&>...>, bool> = true>
-   inline void MakeRef(VectorTypes &...vs);
-
-   /** @brief Update the @a i-th MultiVector block to reference the given
-       Vector @a v.
-
-       After calling this method, operations on the @a i-th Vector block are
-       performed directly on the Vector @a v. In particular, there is no need
-       to synchronize the Memory flags of @a v and the ones of the @a i-th
-       Vector blocks when data is moved between host and device. */
-   inline void MakeRef(int i, Vector &v) { blocks[i] = &v; }
-};
-
-// Inline and template methods
-
-inline Vector &MultiVector::operator[](int i)
-{
-   auto &bi = blocks[i];
-   return (bi.index() == 0) ? std::get<0>(bi) : *std::get<1>(bi);
-}
-
-inline const Vector &MultiVector::operator[](int i) const
-{
-   auto &bi = blocks[i];
-   return (bi.index() == 0) ? std::get<0>(bi) : *std::get<1>(bi);
-}
-
-template <typename... VectorTypes,
-          std::enable_if_t<
-             std::conjunction_v<
-                std::is_convertible<VectorTypes&,Vector&>...>, bool>>
-inline void MultiVector::MakeRef(VectorTypes &...vs)
-{
-   blocks.resize(sizeof...(vs));
-   if constexpr (sizeof...(vs) > 0)
-   {
-      const std::array vs_p{&static_cast<Vector&>(vs)...};
-      for (std::size_t i = 0; i < sizeof...(vs); i++)
-      {
-         blocks[i] = vs_p[i];
-      }
-   }
-}
-
-} // namespace mfem
-
-#endif // MFEM_MULTIVECTOR_HPP

linalg/operator.cpp

@@ -111,16 +111,6 @@ void Operator::ArrayAddMultTranspose(const Array<const Vector *> &X,
    }
 }
 
-void Operator::Mult(const MultiVector &, MultiVector &)
-{
-   MFEM_ABORT("this method is not overriden for this class!");
-}
-
-Operator &Operator::GetGradient(const MultiVector &) const
-{
-   MFEM_ABORT("this method is not overriden for this class!");
-}
-
 void Operator::FormLinearSystem(const Array<int> &ess_tdof_list,
                                 Vector &x, Vector &b,
                                 Operator* &Aout, Vector &X, Vector &B,

linalg/operator.hpp

@@ -13,7 +13,6 @@
 #define MFEM_OPERATOR
 
 #include "vector.hpp"
-#include "multivector.hpp"
 
 namespace mfem
 {
@@ -130,16 +129,6 @@ public:
    virtual void ArrayAddMultTranspose(const Array<const Vector *> &X,
                                       Array<Vector *> &Y, const real_t a = 1.0) const;
 
-   /** @brief Operator application, y = A(x), where the input @a x and the
-       output @a y are MultiVector objects, i.e. they generally use
-       non-contiguous memory representation.
-
-       The typical use case for this method are block operators like
-       DifferentiableOperator.
-
-       The base class implementation for the method is to generate an error. */
-   virtual void Mult(const MultiVector &x, MultiVector &y);
-
    /** @brief Evaluate the gradient operator at the point @a x. The default
        behavior in class Operator is to generate an error. */
    virtual Operator &GetGradient(const Vector &x) const
@@ -148,16 +137,6 @@ public:
       return const_cast<Operator &>(*this);
    }
 
-   /** @brief Evaluate the gradient operator at the point @a x. The input @a x
-       is provided as a MultiVector, i.e. it generally uses non-contiguous
-       memory representation.
-
-       The typical use case for this method are block operators like
-       DifferentiableOperator.
-
-       The base class implementation for the method is to generate an error. */
-   virtual Operator &GetGradient(const MultiVector &x) const;
-
    /** @brief Computes the diagonal entries into @a diag. Typically, this
        operation only makes sense for linear Operator%s. In some cases, only an
        approximation of the diagonal is computed. */

miniapps/CMakeLists.txt

@@ -22,7 +22,6 @@ add_subdirectory(common)
 add_subdirectory(contact)
 add_subdirectory(dfem)
 add_subdirectory(diag-smoothers)
-add_subdirectory(multiapp)
 add_subdirectory(dpg)
 add_subdirectory(electromagnetics)
 add_subdirectory(fluids/navier)

miniapps/hdiv-linear-solver/change_basis.hpp

@@ -80,8 +80,6 @@ public:
    // limitations
    void MultRT_2D(const Vector &x, Vector &y, Mode mode) const;
    void MultRT_3D(const Vector &x, Vector &y, Mode mode) const;
-   // suppress warning about hiding overloaded virtual function:
-   using Operator::Mult;
 };
 
 } // namespace mfem

miniapps/multiapp/CMakeLists.txt

@@ -1,52 +0,0 @@
-# Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
-# at the Lawrence Livermore National Laboratory. All Rights reserved. See files
-# LICENSE and NOTICE for details. LLNL-CODE-806117.
-#
-# This file is part of the MFEM library. For more information and source code
-# availability visit https://mfem.org.
-#
-# MFEM is free software; you can redistribute it and/or modify it under the
-# terms of the BSD-3 license. We welcome feedback and contributions, see file
-# CONTRIBUTING.md for details.
-
-set(MESH_FILES
-  backward-facing-step.msh
-  channel-cylinder.msh
-)
-
-# Add a target to copy the mesh files from the source directory; used by sample
-# runs.
-set(SRC_MESH_FILES)
-foreach(MESH_FILE ${MESH_FILES})
-  list(APPEND SRC_MESH_FILES ${CMAKE_CURRENT_SOURCE_DIR}/${MESH_FILE})
-endforeach()
-add_custom_command(OUTPUT data_is_copied
-  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${SRC_MESH_FILES} .
-  COMMAND ${CMAKE_COMMAND} -E touch data_is_copied
-  COMMENT "Copying multiapp miniapps data files ...")
-add_custom_target(copy_miniapps_multiapp_data DEPENDS data_is_copied)
-
-
-list(APPEND MULTIAPP_COMMON_SOURCES
-	multiapp.cpp)
-
-list(APPEND MULTIAPP_COMMON_HEADERS
-	multiapp.hpp)
-
-set(MULTIAPP_COMMON_FILES
-	EXTRA_SOURCES ${MULTIAPP_COMMON_SOURCES}
-	EXTRA_HEADERS ${MULTIAPP_COMMON_HEADERS})
-
-
-# Parallel apps.
-if (MFEM_USE_MPI)
-  add_mfem_miniapp(coupled-diffusion
-    MAIN coupled-diffusion.cpp
-    ${MFEM_MINIAPPS_COMMON_HEADERS} ${MULTIAPP_COMMON_FILES}
-    LIBRARIES mfem-common)
-  add_dependencies(coupled-diffusion copy_miniapps_multiapp_data)
-
-# Add parallel tests.
-  # if (MFEM_ENABLE_TESTING)
-  # endif()
-endif()

miniapps/multiapp/coupled-diffusion.cpp

@@ -1,895 +0,0 @@
-#include "mfem.hpp"
-#include "multiapp.hpp"
-#include <fstream>
-using namespace mfem;
-using namespace std;
-
-struct CaseContext
-{
-   int ser_ref = 1;         // Serial mesh refinement
-   int order = 3;           // Finite element order
-   bool visualization = true;// Visualization on/off
-   int grad_mode = 1;       // Gradient mode for the coupled operator - 0: finite difference,
-                            //                                          1: back/forward propagation
-   bool coupled = true;     // Coupled (true) vs. uncoupled (false) solves
-   int nl_iter = 50;        // Maximum number of nonlinear iterations
-   int lin_iter = 2000;     // Maximum number of linear iterations
-
-#if defined(MFEM_USE_DOUBLE)
-   real_t tol_nsolve = 1e-4;
-   real_t tol_lsolve = 1e-6;
-#elif defined(MFEM_USE_SINGLE)
-   real_t tol_nsolve = 1e-3;
-   real_t tol_lsolve = 1e-3;
-#else
-#error "Only single and double precision are supported!"
-   real_t tol_nsolve = 0;
-   real_t tol_lsolve = 0;
-#endif
-} ctx;
-
-void SetSolverParameters(IterativeSolver *solver, real_t rtol, real_t atol , int max_it,
-                         int print_level, bool iterative_mode);
-
-
-/// A functional diffusion coefficient (i.e., k(T))
-class FunctionalCoefficient : public Coefficient
-{
-public:
-   enum Mode { FUNC = 0, GRAD = 1};
-
-protected:
-   ParGridFunction *T_gf = nullptr;
-   real_t kref = 1.0;
-   real_t a0 = 0.0, a1 = 0.0, a2 = 0.0;
-   int findex = 0;
-   Mode mode = Mode::FUNC; // otherwise, grad
-
-public:
-   FunctionalCoefficient(ParGridFunction *T_gf, real_t kref): 
-                         T_gf(T_gf), kref(kref) { }
-
-   FunctionalCoefficient(ParGridFunction *T_gf, real_t kref, real_t a0):
-                         T_gf(T_gf), kref(kref), a0(a0) { findex = 1; }
-
-   FunctionalCoefficient(ParGridFunction *T_gf, real_t kref,
-                         real_t a0, real_t a1, real_t a2): T_gf(T_gf),
-                         kref(kref), a0(a0), a1(a1), a2(a2) { findex = 2; }
-
-    real_t Exponential(real_t x, bool eval_f) const
-    { 
-      real_t f = kref*exp(a0*x);
-      return (eval_f ? f : a0*f);
-    }
-    real_t Polynomial(real_t x, bool eval_f) const
-    { 
-      return (eval_f ? kref*(a0 + a1*x + a2*x*x) : kref*(a1 + 2*a2*x));
-    }
-
-    void SetMode(Mode mode) { this->mode = mode; }
-    Mode GetMode() const { return mode; }
-    
-    void UpdateGridFunction(ParGridFunction *gf) { T_gf = gf; }
-
-    real_t Eval(real_t x, bool eval_f) const
-    {
-        switch (findex)
-        {
-        case 1:
-         return Exponential(x, eval_f);
-        case 2:
-         return Polynomial(x, eval_f);
-        default:
-         return kref;
-        }
-    }
-
-    real_t Eval(ElementTransformation &Tr,
-              const IntegrationPoint &ip) override
-    {
-        real_t T = T_gf ? T_gf->GetValue(Tr, ip) : 0.0;
-        bool eval_f = (mode == Mode::FUNC);
-        return Eval(T, eval_f);
-    }
-};
-
-
-/// A coefficient defined by the product of grid functions, e.g. k(T) = prod_i x_i
-class GridFunctionProductCoefficient : public Coefficient
-{
-protected:
-   std::vector<ParGridFunction*> &x;
-
-public:
-   GridFunctionProductCoefficient(std::vector<ParGridFunction*> &x) : x(x) { }
-
-   real_t Eval(ElementTransformation &Tr, const IntegrationPoint &ip) override
-   {
-      real_t prod = 1.0;
-      for(size_t i = 0; i < x.size(); i++)
-      {
-         real_t val = x[i]->GetValue(Tr, ip);
-         prod *= val;
-      }
-      return prod;
-   }
-};
-
-class CoefficientIntegrator : public NonlinearFormIntegrator
-{
-protected:
-   FunctionalCoefficient *func = nullptr;
-   Vector shape;
-
-public:
-   CoefficientIntegrator(FunctionalCoefficient *func) : func(func) { }
-
-
-   void SetCoefficient(FunctionalCoefficient *f) { func = f; }
-
-   void AssembleElementVector(const FiniteElement &el,
-                              ElementTransformation &Tr,
-                              const Vector &elfun, Vector &elvect)
-   {
-      int dof = el.GetDof();
-      shape.SetSize(dof);
-      elvect.SetSize(dof);
-      elvect = 0.0;
-
-      const IntegrationRule *ir = &el.GetNodes();
-
-      for (int i = 0; i < ir->GetNPoints(); i++)
-      {
-         const IntegrationPoint &ip = ir->IntPoint(i);
-         el.CalcShape(ip, shape);
-         Tr.SetIntPoint(&ip);
-         real_t x = elfun * shape; // Evaluate the function at the integration point
-         real_t fval = func->Eval(x, true);
-         for (int j = 0; j < dof; j++)
-         {
-            elvect(j) += fval * shape(j);
-         }
-      }
-   }
-
-   void AssembleElementGrad(const FiniteElement &el, ElementTransformation &Tr,
-                            const Vector &elfun, DenseMatrix &elmat)
-   {
-      int dof = el.GetDof();
-      shape.SetSize(dof);
-      elmat.SetSize(dof);
-      elmat = 0.0;
-
-      const IntegrationRule *ir = &el.GetNodes();
-
-      for (int i = 0; i < ir->GetNPoints(); i++)
-      {
-         const IntegrationPoint &ip = ir->IntPoint(i);
-         el.CalcShape(ip, shape);
-         Tr.SetIntPoint(&ip);
-         real_t x = elfun * shape; // Evaluate the function at the integration point
-         real_t dfdx = func->Eval(x, false); // Evaluate the derivative of the function at the integration point
-         for (int j = 0; j < dof; j++)
-         {
-            elmat(j,j) += dfdx * shape(j); // Diagonal contribution to the Jacobian
-         }
-      }
-   }
-};
-
-
-class NonlinearDiffusionIntegrator : public NonlinearFormIntegrator
-{
-protected:
-   Coefficient *k;
-   Coefficient *dk;
-
-   Vector u, vec, shape;
-   DenseMatrix dshape, dshapedxt, adjJ;
-public:
-   NonlinearDiffusionIntegrator(Coefficient *kappa, Coefficient *dkappa) :
-                                k(kappa), dk(dkappa) { }
-
-   virtual void AssembleElementVector(const FiniteElement &el,
-                                      ElementTransformation &Tr,
-                                      const Vector &elfun, Vector &elvec)
-   {
-      int dim = el.GetDim();
-      int dof = el.GetDof();
-      real_t w;
-
-      elvec.SetSize(dof);
-      elvec = 0.0;
-
-      const IntegrationRule *ir = IntRule ? IntRule : &IntRules.Get(el.GetGeomType(), 2*el.GetOrder());
-      u.SetSize(dim);
-      vec.SetSize(dim);
-      dshape.SetSize(dof, dim);
-      adjJ.SetSize(dim, dim);
-
-      for (int i = 0; i < ir->GetNPoints(); i++)
-      {
-         const IntegrationPoint &ip = ir->IntPoint(i);
-         el.CalcDShape(ip, dshape);
-
-         Tr.SetIntPoint(&ip);
-         CalcAdjugate(Tr.Jacobian(), adjJ);
-         w = ip.weight / Tr.Weight();
-
-         dshape.MultTranspose(elfun, u);
-         adjJ.MultTranspose(u, vec);
-         if(k)
-         {
-            w *= k->Eval(Tr, ip);
-         }
-
-         vec *= w;
-         adjJ.Mult(vec, u);
-         dshape.AddMult(u, elvec);
-      }
-   }
-
-   void AssembleElementGrad(const FiniteElement &el, ElementTransformation &Tr,
-                            const Vector &elfun, DenseMatrix &elmat)
-   {
-      int dim = el.GetDim();
-      int dof = el.GetDof();
-      real_t w, k0 = 0.0, dk0 = 0.0;
-
-      elmat.SetSize(dof);
-      elmat = 0.0;
-
-      const IntegrationRule *ir = IntRule ? IntRule : &IntRules.Get(el.GetGeomType(), 2*el.GetOrder());
-      u.SetSize(dim);
-      shape.SetSize(dof);
-      vec.SetSize(dof);
-      dshape.SetSize(dof, dim);
-      dshapedxt.SetSize(dof, dim);
-
-      // f = grad(psi) * k(u) * grad(T)
-      // df/dT = grad(psi) ( k(u0) * grad(T) + k'(u0) * grad(u0) * T )
-      for (int i = 0; i < ir->GetNPoints(); i++)
-      {
-         const IntegrationPoint &ip = ir->IntPoint(i);
-         el.CalcShape(ip, shape);
-         el.CalcDShape(ip, dshape);
-
-         Tr.SetIntPoint(&ip);
-         w = ip.weight / Tr.Weight();
-
-         Mult(dshape, Tr.AdjugateJacobian(), dshapedxt);
-
-         k0  = k ? k->Eval(Tr, ip) : 0.0;
-         dk0 = dk ? dk->Eval(Tr, ip) : 0.0;
-
-         if(k0 != 0.0) // grad(psi) * k(u0) * grad(T)
-         {
-            real_t kdT = w*k0;
-            AddMult_a_AAt(kdT, dshapedxt, elmat);
-         }
-
-         if(dk0 != 0.0) // grad(psi) * (k'(T0) * grad(T0)) * T
-         {
-            dk0 = w*dk->Eval(Tr, ip);
-            dshapedxt.MultTranspose(elfun, u); // grad(T0) in physical space
-            u *= dk0; // k'(T0) * grad(T0)
-            dshapedxt.Mult(u, vec); // grad(psi) * k'(T0) * grad(T0)
-            AddMultVWt(vec, shape, elmat); // grad(psi) * k'(T0) * grad(T0) * T
-         }
-      }
-   }
-};
-
-/// An application that takes an input field T, and computes an output field k(T)
-// represented by the FunctionalCoefficient class.
-class DiffusionCoefficient : public GraphNode
-{
-public:
-   using Mode = FunctionalCoefficient::Mode;
-
-protected:
-   ParFiniteElementSpace &fes;
-   mutable ParGridFunction T, k;
-   mutable FunctionalCoefficient *kc;
-   // mutable Vector tdof, kdof, dk_dof, dT_dof;
-   mutable Mode mode = Mode::FUNC;
-
-   mutable ParNonlinearForm Nform;
-   mutable Operator *J = nullptr; // Jacobian for the nonlinear form
-
-   CoefficientIntegrator *coeff_integrator = nullptr;
-
-public:
-   DiffusionCoefficient(ParFiniteElementSpace &fes) :
-                        GraphNode(fes.GetTrueVSize()), fes(fes), T(&fes), k(&fes),
-                        kc(new FunctionalCoefficient(&T, 1.0, 5.0e-2)),
-                        Nform(&fes),
-                        coeff_integrator(new CoefficientIntegrator(kc))
-   {
-      k = 0.0;
-      T = 0.0;
-      k.ProjectCoefficient(*kc);
-
-      // Testing with the nonlinear form framework to compute k(T) and dk/dT
-      Nform.AddDomainIntegrator(coeff_integrator); // Transfer ownership
-      Nform.SetGradientType(Operator::Type::Hypre_ParCSR);
-      Nform.Setup();
-
-      SetInputOffsets(Array<int>({0, fes.GetTrueVSize()}));
-      SetOutputOffsets(Array<int>({0, fes.GetTrueVSize()}));
-   }
-
-   void SetMode(Mode mode) { this->mode = mode; }
-
-   FunctionalCoefficient* GetCoefficient() { return kc; }
-
-   void SetCoefficient(FunctionalCoefficient *fc)
-   {
-      if(kc) delete kc;
-      kc = fc;
-      kc->SetMode(mode);
-      kc->UpdateGridFunction(&T);
-      coeff_integrator->SetCoefficient(kc);
-   }
-
-   void Mult(const Vector &x, Vector &y) const override
-   {
-      BlockVector xb(x.GetData(), InputOffsets());
-      BlockVector yb(y.GetData(), OutputOffsets());
-
-      MultiVector xmv(1), ymv(1);
-      xmv.MakeRef(0, xb.GetBlock(0));
-      ymv.MakeRef(0, yb.GetBlock(0));
-
-      const_cast<DiffusionCoefficient*>(this)->Mult(xmv, ymv);
-   }
-
-   void Mult(const MultiVector &x, MultiVector &y) override
-   {
-      const Vector &tdof = x[0];
-      Vector &kdof = y[0];
-
-      Nform.Mult(tdof, kdof);
-      if(exec_mode == GraphNode::GRADIENT_MODE)
-      {
-         J = &Nform.GetGradient(tdof); // Store jacobian for JVP
-      }
-      else
-      {
-         J = nullptr; // Clear the Jacobian if not in gradient mode
-      }
-   }
-
-   // TODO: Possibly delete and only support MultiVector version of GradientMult
-   void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const override
-   {
-      BlockVector xb(x.GetData(), InputOffsets());
-      BlockVector dxb(dx.GetData(), InputOffsets());
-      BlockVector dyb(dy.GetData(), OutputOffsets());
-
-      MultiVector xmv(1), dxmv(1), dymv(1);
-      xmv.MakeRef(0, xb.GetBlock(0));
-      dxmv.MakeRef(0, dxb.GetBlock(0));
-      dymv.MakeRef(0, dyb.GetBlock(0));
-      const_cast<DiffusionCoefficient*>(this)->GradientMult(xmv, dxmv, dymv);
-   }
-
-   void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const override
-   {
-      const Vector &tdof = x[0];
-      const Vector &xadj = dx[0];
-      Vector &yadj = dy[0];
-
-      if(J)
-      {
-         J->Mult(xadj, yadj);
-      }
-      else
-      {
-         J = &Nform.GetGradient(tdof); // Store jacobian for JVP
-         J->Mult(xadj, yadj);
-      }
-   }
-
-   ~DiffusionCoefficient() override
-   {
-      if(kc) delete kc;
-   }
-};
-
-/// An application that takes n input fields x_i, and computes an output 
-/// field prod(x) := y = prod_i x_i.
-/// Also provides the derivative dy/dx_i = prod_{j!=i} x_j * dx_i/dx for i = 0,...,n-1.
-class ProductGridFunctions : public GraphNode
-{
-protected:
-
-   ParFiniteElementSpace &fes;
-   mutable std::vector<ParGridFunction*> x_gf;
-   mutable Vector dfdx;
-   mutable ParGridFunction y_gf;
-   mutable GridFunctionProductCoefficient prod_coeff;
-
-public:
-   ProductGridFunctions(ParFiniteElementSpace &fes, int n) :
-                     //   GraphNode(fes.GetTrueVSize()),
-                       GraphNode(fes.GetTrueVSize(), fes.GetTrueVSize() * n),
-                       fes(fes), x_gf(n),
-                       y_gf(&fes), prod_coeff(x_gf)
-   {
-      Array<int> offsets(n+1);
-      offsets[0] = 0;
-      for (int i = 0; i < n; i++)
-      {
-         x_gf[i]  = new ParGridFunction(&fes);
-         *x_gf[i] = 0.0;
-         offsets[i+1] = offsets[i] + fes.GetTrueVSize();
-      }
-      y_gf = 0.0;
-      y_gf.ProjectCoefficient(prod_coeff);
-
-      SetInputOffsets(offsets);
-      SetOutputOffsets(Array<int>({0, fes.GetTrueVSize()}));
-   }
-
-   void Mult(const Vector &x, Vector &y) const override
-   {
-      BlockVector xb(x.GetData(), InputOffsets());
-      BlockVector yb(y.GetData(), OutputOffsets());
-
-      MultiVector xmv(x_gf.size()), ymv(1);
-      for (size_t i = 0; i < x_gf.size(); i++)
-      {
-         xmv.MakeRef(i, xb.GetBlock(i));
-      }
-      ymv.MakeRef(0, yb.GetBlock(0));
-      const_cast<ProductGridFunctions*>(this)->Mult(xmv, ymv);
-   }
-
-   void Mult(const MultiVector &x, MultiVector &y) override
-   {
-      for (size_t i = 0; i < x_gf.size(); i++)
-      {
-         const Vector &x_dof = x[i];
-         x_gf[i]->SetFromTrueDofs(x_dof);
-      }
-
-      Field *out_field = OutputField(0);
-      Vector &y_dof = y[0];
-      y_gf.ProjectCoefficient(prod_coeff);
-      y_gf.GetTrueDofs(y_dof);
-   }
-
-   // TODO: Possibly delete and only support MultiVector version of GradientMult
-   void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const override
-   {
-      BlockVector xb(x.GetData(), InputOffsets());
-      BlockVector dxb(dx.GetData(), InputOffsets());
-      BlockVector dyb(dy.GetData(), OutputOffsets());
-
-      MultiVector xmv(x_gf.size()), dxmv(x_gf.size()), dymv(1);
-      for (size_t i = 0; i < x_gf.size(); i++)
-      {
-         xmv.MakeRef(i, xb.GetBlock(i));
-         dxmv.MakeRef(i, dxb.GetBlock(i));
-      }
-      dymv.MakeRef(0, dyb.GetBlock(0));
-      const_cast<ProductGridFunctions*>(this)->GradientMult(xmv, dxmv, dymv);
-   }
-
-   void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const override
-   {
-      // Jacobian vector product for y = prod_i x_i is:
-      // dy/dx = sum_i (prod_{j!=i} x_j * dx_i/dx)
-      for (size_t i = 0; i < x_gf.size(); i++)
-      {
-         const Vector &x_dof = x[i];
-         x_gf[i]->SetFromTrueDofs(x_dof); // Set all x_i
-      }
-
-      Vector &jvp = dy[0];
-      jvp = 0.0;
-      for (size_t i = 0; i < x_gf.size(); i++)
-      {
-         const Vector &x_dof = x[i];
-         const Vector &dx_dof = dx[i]; // Get dx_i/dx
-
-         x_gf[i]->SetFromTrueDofs(dx_dof); // Set x_i = dx_i/dx for i-th term in the product
-         y_gf.ProjectCoefficient(prod_coeff); // Recompute product with x_i replaced by dx_i/dx
-         y_gf.GetTrueDofs(dfdx); // Get prod_{j!=i} x_j * dx_i/dx for i-th term
-         jvp += dfdx; // Accumulate contribution from i-th term
-         x_gf[i]->SetFromTrueDofs(x_dof); // reset to original value for next iteration
-      }
-   }
-
-   ~ProductGridFunctions() override
-   {
-      for (size_t i = 0; i < x_gf.size(); i++)
-      {
-         if(x_gf[i]) delete x_gf[i];
-      }
-   }
-};
-
-
-/// An application that represents the nonlinear diffusion operator: f(T) = -Div(k(u) grad(T)) 
-/// with input field T and k, and output field f(T).
-class DiffusionOperator : public GraphNode
-{
-public:
-
-   // Mesh and finite element space
-   ParMesh &mesh;
-   ParFiniteElementSpace &fes;
-
-   /// Essential dof array.
-   Array<int> ess_tdofs;
-
-   /// Grid functions for the temperature and heat flux
-   mutable ParGridFunction T, k, dk;
-   mutable GridFunctionCoefficient k_gfc, dk_gfc;
-   mutable ParNonlinearForm Nform;
-   mutable ParLinearForm bform;
-   mutable Vector b;
-
-   ConstantCoefficient zero_coeff, one_coeff;
-
-   mutable FunctionalCoefficient *kc = nullptr;
-   mutable HypreParMatrix *dfdk_mat = nullptr, *dfdT_mat = nullptr;
-
-public:
-
-   DiffusionOperator(ParFiniteElementSpace &fes_) :
-                     // GraphNode(fes_.GetTrueVSize()),
-                     GraphNode(fes_.GetTrueVSize(),2*fes_.GetTrueVSize()),
-                     mesh(*fes_.GetParMesh()), fes(fes_),
-                     T(&fes), k(&fes), dk(&fes),
-                     k_gfc(&k), dk_gfc(&dk),
-                     Nform(&fes), bform(&fes),
-                     zero_coeff(0.0), one_coeff(1.0)
-   {
-      fes.GetBoundaryTrueDofs(ess_tdofs);
-      T = 0.0;
-      k = 0.0;
-      dk = 0.0;
-
-      bform.AddDomainIntegrator(new DomainLFIntegrator(one_coeff));
-      Nform.AddDomainIntegrator(new NonlinearDiffusionIntegrator(&k_gfc, &dk_gfc));
-      Nform.SetGradientType(Operator::Type::Hypre_ParCSR);
-
-      b.SetSize(fes.GetTrueVSize()); b = 0.0;
-      Assemble();
-
-      SetInputOffsets(Array<int>({0, fes.GetTrueVSize(), 2*fes.GetTrueVSize()}));
-      SetOutputOffsets(Array<int>({0, fes.GetTrueVSize()}));
-   }
-
-   void SetCoefficient(FunctionalCoefficient *fc) { kc = fc; }
-
-   void Assemble()
-   {
-      AssembleLinearForms();
-      AssembleBilinearForms();
-      AssembleNonlinearForms();
-   }
-
-   void AssembleBilinearForms()
-   {}
-
-   void AssembleNonlinearForms()
-   {
-      Nform.SetEssentialTrueDofs(ess_tdofs);
-      Nform.Setup();
-   }
-
-   void AssembleLinearForms()
-   {
-      bform.Assemble();
-      bform.ParallelAssemble(b);
-   }
-
-   void Mult(const Vector &x, Vector &y) const override
-   {
-      BlockVector xb(x.GetData(), InputOffsets());
-      BlockVector yb(y.GetData(), OutputOffsets());
-
-      MultiVector xmv(2), ymv(1);
-      xmv.MakeRef(0, xb.GetBlock(0));
-      xmv.MakeRef(1, xb.GetBlock(1));
-      ymv.MakeRef(0, yb.GetBlock(0));
-
-      const_cast<DiffusionOperator*>(this)->Mult(xmv, ymv);
-   }
-
-   void Mult(const MultiVector &x, MultiVector &y) override
-   {
-      const Vector &tdofs = x[0];
-      const Vector &kdofs = x[1];
-      Vector &fdofs = y[0];
-
-      k.SetFromTrueDofs(kdofs); // update for use in k_gfc
-
-      if(exec_mode == GraphNode::GRADIENT_MODE)
-      {
-         if(dfdT_mat) delete dfdT_mat;
-         if(dfdk_mat) delete dfdk_mat;
-
-         dk = 0.0;
-         k.SetFromTrueDofs(kdofs);
-         Operator* grad = &Nform.GetGradient(tdofs);
-         dfdT_mat = new HypreParMatrix(dynamic_cast<const HypreParMatrix&>(*grad)); // deep copy
-
-         dk = 1.0;
-         k  = 0.0;
-         grad = &Nform.GetGradient(tdofs);
-         dfdk_mat = new HypreParMatrix(dynamic_cast<const HypreParMatrix&>(*grad)); // deep copy
-      }
-      else
-      {
-         if(dfdT_mat) { delete dfdT_mat; dfdT_mat = nullptr; }
-         if(dfdk_mat) { delete dfdk_mat; dfdk_mat = nullptr; }
-      }
-
-      Nform.Mult(tdofs, fdofs);
-      fdofs.SetSubVector(ess_tdofs, 0.0);
-   }
-
-   // Exact block jacobian [df/dT, df/dk]
-   Operator& GetGradient(const Vector &x) const override
-   {
-      MFEM_ABORT("GetGradient not implemented for DiffusionOperator");
-   }
-
-   // TODO: Possibly delete and only support MultiVector version of GradientMult
-   void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const override
-   {
-      BlockVector xb(x.GetData(), InputOffsets());
-      BlockVector dxb(dx.GetData(), InputOffsets());
-      BlockVector dyb(dy.GetData(), OutputOffsets());
-
-      Vector &Tadj = dxb.GetBlock(0);
-      Vector &kadj = dxb.GetBlock(1);
-      Vector &yadj = dyb.GetBlock(0);
-
-      Vector &tdofs = xb.GetBlock(0);
-      Vector &kdofs = xb.GetBlock(1);
-
-      dfdT_mat->Mult(Tadj, yadj);
-      dfdk_mat->AddMult(kadj, yadj);
-   }
-
-   void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const override
-   {
-      const Vector &Tadj = dx[0];
-      const Vector &kadj = dx[1];
-      Vector &yadj = dy[0];
-
-      const Vector &tdofs = x[0];
-      const Vector &kdofs = x[1];
-
-      dfdT_mat->Mult(Tadj, yadj);
-      dfdk_mat->AddMult(kadj, yadj);
-   }
-
-   /// @brief Destroy the DiffusionOperator object
-   ~DiffusionOperator() override
-   {
-      if(dfdT_mat) delete dfdT_mat;
-      if(dfdk_mat) delete dfdk_mat;
-   }
-};
-
-
-int main(int argc, char *argv[])
-{
-   Mpi::Init();
-   Hypre::Init();
-
-   using GradMode = DAGraph::GradMode;
-
-   OptionsParser args(argc, argv);
-   args.AddOption(&ctx.order, "-o", "--order",
-                  "Finite element order (polynomial degree).");
-   args.AddOption(&ctx.visualization, "-vis", "--visualization", "-no-vis",
-                  "--no-visualization",
-                  "Enable or disable GLVis visualization.");
-   args.AddOption(&ctx.ser_ref, "-rs", "--serial-refine",
-                  "Number of times to refine the mesh in serial.");
-
-   args.AddOption(&ctx.grad_mode, "-gm", "--grad-mode",
-                  "Gradient mode for the coupled operator (0: exact, 1: finite difference, 2: algorithmic differentiation)");
-   args.AddOption(&ctx.coupled, "-cp", "--coupled", "-ucp", "--uncoupled",
-                  "Coupled (true) vs. uncoupled (false) solves.");
-   args.ParseCheck();
-
-
-   int order = ctx.order;
-   std::string mesh_file = "../../data/star.mesh";
-   Mesh *serial_mesh = new Mesh(mesh_file);
-   int dim = serial_mesh->Dimension();
-
-   for (int i = 0; i < ctx.ser_ref; ++i) { serial_mesh->UniformRefinement(); }
-   serial_mesh->SetCurvature(order, false, dim, Ordering::byNODES);
-
-   ParMesh pmesh(MPI_COMM_WORLD, *serial_mesh);
-   delete serial_mesh;
-   pmesh.UniformRefinement();
-
-   // Finite element spaces
-   H1_FECollection fec(order, dim);
-   ParFiniteElementSpace fes(&pmesh, &fec);
-
-   // Build all operator nodes
-   DiffusionCoefficient diff_coeff_1(fes);
-   diff_coeff_1.SetName("k(T1)");
-   diff_coeff_1.SetCoefficient(new FunctionalCoefficient(nullptr, 1.0, 3.5e-2));
-   // diff_coeff_1.SetCoefficient(new FunctionalCoefficient(nullptr, 1.0, 1.0, 0.1, 0.0));
-
-   DiffusionCoefficient diff_coeff_2(fes);
-   diff_coeff_2.SetName("k(T2)");
-   diff_coeff_2.SetCoefficient(new FunctionalCoefficient(nullptr, 1.0, 1.0, 2.0, 0.0));
-   // diff_coeff_2.SetCoefficient(new FunctionalCoefficient(nullptr, 1.5, 2.5e-2));
-
-   ProductGridFunctions prod_coeff(fes, 2);
-   prod_coeff.SetName("k(T1,T2)");
-
-   DiffusionOperator diff_op1(fes);
-   diff_op1.SetName("Div(k(T1,T2) grad(T1))");
-   diff_op1.SetCoefficient(diff_coeff_1.GetCoefficient());
-
-   DiffusionOperator diff_op2(fes);
-   diff_op2.SetName("Div(k(T1,T2) grad(T2))");
-   diff_op2.SetCoefficient(diff_coeff_2.GetCoefficient());
-
-
-   // Build the DAG in any order, and then sort it to ensure the correct execution order
-   DAGraph dag(5);
-   dag.AddOperator(&diff_coeff_1);
-   dag.AddOperator(&diff_op1, fes.GetTrueVSize());
-   dag.AddOperator(&diff_op2, fes.GetTrueVSize());
-   dag.AddOperator(&diff_coeff_2);
-   dag.AddOperator(&prod_coeff);
-
-   Vector k1vec(fes.GetTrueVSize()); k1vec = 0.0;
-   Vector k2vec(fes.GetTrueVSize()); k2vec = 0.0;
-   Vector kpvec(fes.GetTrueVSize()); kpvec = 0.0;
-
-   Vector k1adj(fes.GetTrueVSize()); k1adj = 0.0;
-   Vector k2adj(fes.GetTrueVSize()); k2adj = 0.0;
-   // Vector kpadj(fes.GetTrueVSize()); kpadj = 0.0;
-
-   // Input fields get data from 'x' in DAGraph::Mult(x, y)
-   Field T1_field(nullptr, nullptr);
-   Field T2_field(nullptr, nullptr);
-
-   // Write space for data and adjoint only needed
-   // for the intermediate fields k1, k2, and k_prod
-   Field k1_field(&k1vec, &k1adj);
-   Field k2_field(&k2vec, &k2adj);
-   Field kp_field(&kpvec, &kpvec); // can use same space for data & adjoint
-
-   // Output fields get data from 'y' in DAGraph::Mult(x, y)
-   Field f1_field(nullptr, nullptr);
-   Field f2_field(nullptr, nullptr);
-
-
-   // Add input and output to the DAG
-   int sz = fes.GetTrueVSize();
-   dag.AddInput(&T1_field, sz);
-   dag.AddInput(&T2_field, sz);
-   dag.AddOutput(&f1_field, sz);
-   dag.AddOutput(&f2_field, sz);
-
-   // Form connections between the nodes in the DAG
-   diff_coeff_1.AddInput(&T1_field);
-   diff_coeff_1.AddOutput(&k1_field);
-
-   diff_coeff_2.AddInput(&T2_field);
-   diff_coeff_2.AddOutput(&k2_field);
-
-   prod_coeff.AddInputs(&k1_field, &k2_field);
-   prod_coeff.AddOutput(&kp_field);
-
-   diff_op1.AddInput(&T1_field);
-   diff_op1.AddOutput(&f1_field);
-
-   diff_op2.AddInput(&T2_field);
-   diff_op2.AddOutput(&f2_field);
-
-   if(ctx.coupled)
-   {
-      diff_op1.AddInput(&kp_field); // kp_field
-      diff_op2.AddInput(prod_coeff.OutputField(0)); // Can also use kp_field directly
-   }
-   else
-   {
-      diff_op1.AddInput(&k1_field); // Can also use diff_coeff_1.OutputField(0)
-      diff_op2.AddInput(&k2_field); // Can also use diff_coeff_2.OutputField(0)
-   }
-
-   // Assemble DAG: topological sort, validate nodes, etc.
-   dag.Assemble();
-
-   std::string output_prefix = ctx.coupled ? "Coupled_Diffusion" : "Uncoupled_Diffusion";
-
-   if(Mpi::Root())
-   {
-      std::ofstream fout(output_prefix+"-dag.txt");
-      fout << "{\n";
-      dag.Save(fout);
-      fout << "}\n";
-      fout << std::flush;
-      fout.close();
-   }
-
-   // Set initial guess and boundary conditions for T1 and T2
-   Array<int> ess_tdofs;
-   fes.GetBoundaryTrueDofs(ess_tdofs);
-
-   int T1_idx = 0;
-   int T2_idx = 1;
-
-   BlockVector xb(dag.InputOffsets());
-   BlockVector yb(dag.OutputOffsets());
-
-   xb.GetBlock(T1_idx).Randomize();
-   xb.GetBlock(T2_idx).Randomize();
-   xb.GetBlock(T1_idx).SetSubVector(ess_tdofs, 0.0);
-   xb.GetBlock(T2_idx).SetSubVector(ess_tdofs, 0.0);
-
-   // Build the nonlinear solver and linear solver for the DAG
-   NewtonSolver newton_solver(pmesh.GetComm());
-   GMRESSolver linear_solver(pmesh.GetComm());
-   linear_solver.SetKDim(500);
-   SetSolverParameters(&newton_solver, ctx.tol_nsolve, 0.0, ctx.nl_iter, 1, true);
-   SetSolverParameters(&linear_solver, ctx.tol_lsolve, 0.0, ctx.lin_iter, 1, false);
-
-   newton_solver.SetPreconditioner(linear_solver);
-   linear_solver.SetPrintLevel(1);
-
-   // Set the gradient mode for the DAG and solve the coupled system
-   GradMode gm = static_cast<GradMode>(ctx.grad_mode);
-   dag.SetGradientMode(gm);
-   newton_solver.SetOperator(dag);
-   newton_solver.Mult(xb, yb);
-
-   ParaViewDataCollection *pv = nullptr;
-   if (ctx.visualization)
-   {
-      std::string pv_prefix;
-      switch (ctx.grad_mode)
-      {
-         case 0: pv_prefix = "FD"; break;
-         case 1: pv_prefix = "MF"; break;
-         default: pv_prefix = "Unknown"; break;
-      }
-
-      pv = new ParaViewDataCollection(output_prefix+"-"+pv_prefix, &pmesh);
-      pv->SetLevelsOfDetail(order);
-      pv->SetDataFormat(VTKFormat::BINARY);
-      pv->SetHighOrderOutput(true);
-
-      ParGridFunction T1_gf(&fes);
-      ParGridFunction T2_gf(&fes);
-      T1_gf.SetFromTrueDofs(yb.GetBlock(T1_idx));
-      T2_gf.SetFromTrueDofs(yb.GetBlock(T2_idx));
-
-      pv->RegisterField("T1", &T1_gf);
-      pv->RegisterField("T2", &T2_gf);
-      pv->Save();
-      delete pv;
-   }
-
-   std::cout << "Finished solving the coupled diffusion problem." << std::endl;
-   return 0;
-}
-
-void SetSolverParameters(IterativeSolver *solver, real_t rtol, real_t atol,
-                         int max_it, int print_level, bool iterative_mode)
-{
-    solver->SetRelTol(rtol);
-    solver->SetAbsTol(atol);
-    solver->SetMaxIter(max_it);
-    solver->SetPrintLevel(print_level);
-    solver->iterative_mode = iterative_mode;
-}

miniapps/multiapp/multiapp.cpp

@@ -1,873 +0,0 @@
-// 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 "multiapp.hpp"
-
-namespace mfem
-{
-
-DAGraph::~DAGraph()
-{
-    for(int i=0; i < nnodes; i++)
-    {
-        if(node_owned[i] && nodes[i]) delete nodes[i];
-    }
-    if(grad) delete grad;
-}
-
-void DAGraph::Assemble()
-{
-    // Sort graph nodes topologically to ensure correct execution order
-    // Ordering is not unique, hence, id->index maps are needed
-    TopologicalSort();
-
-    // Collect all fields from the nodes into the field map
-    CollectFieldMaps();
-
-    // Compute depth of the graph nodes
-    ComputeDepth();
-
-    // Validate each node
-    for (auto &node : nodes)
-    {
-        ValidateNode(*node);
-    }
-
-    // Update width and height of the DAG from offsets
-    // Check that the input and output offsets are consistent
-    ValidateOffsets();
-    width  = input_offsets.Last();
-    height = output_offsets.Last();
-
-    // Delete any existing gradient operator as node ordering may have changed
-    if (grad) delete grad;
-
-    assembled = true;
-}
-
-void DAGraph::ValidateOffsets()
-{
-    // Check that the input and output offsets are consistent
-    // with the number of inputs and outputs
-    if(InputFields().Size() > 1)
-    {
-        MFEM_ASSERT(input_offsets.Size() == InputFields().Size() + 1,
-                    "Input offsets size inconsistent with number of input fields");
-    }
-    else
-    {
-        input_offsets = Array<int>({0, nodes[0]->Width()});
-    }
-
-    if(OutputFields().Size() > 1)
-    {
-        MFEM_ASSERT(output_offsets.Size() == OutputFields().Size() + 1,
-                    "Output offsets size inconsistent with number of output fields");
-    }
-    else
-    {
-        output_offsets = Array<int>({0, nodes.Last()->Height()});
-    }
-}
-
-void DAGraph::ValidateNode(GraphNode &node)
-{
-    // Validate that the node's input and output fields are consistent with the graph's field map
-    auto inputs = node.InputFields();
-    auto outputs = node.OutputFields();
-
-    // Check offsets match width and height of the node
-    MFEM_ASSERT(node.InputOffsets().Last() == node.Width(),
-                "Node ID: " << node.ID() << " input offsets do not match node width.");
-    MFEM_ASSERT(node.OutputOffsets().Last() == node.Height(),
-                "Node ID: " << node.ID() << " output offsets do not match node height.");
-
-    // Check number of input and output fields match the offsets
-    MFEM_ASSERT(node.InputOffsets().Size() == inputs.Size() + 1,
-                "Node input offsets size inconsistent with number of input fields");
-    MFEM_ASSERT(node.OutputOffsets().Size() == outputs.Size() + 1,
-                "Node output offsets size inconsistent with number of output fields");
-
-    // Check that all input and output fields are registered in the graph's field map
-    for(auto input_field : inputs)
-    {
-        MFEM_ASSERT(fid_to_index.Has(input_field->ID()),
-                    "Input field ID " << input_field->ID() << " not found in graph's field map");
-    }
-    for(auto output_field : outputs)
-    {
-        MFEM_ASSERT(fid_to_index.Has(output_field->ID()),
-                    "Output field ID " << output_field->ID() << " not found in graph's field map");
-    }
-}
-
-void DAGraph::TopologicalSort()
-{
-    Array<int> sorted_indices;
-    sorted_indices.Reserve(nnodes);
-
-    Array<bool> visited(nnodes);
-    visited = false; // Initialize all nodes as unvisited
-
-    // Perform a depth-first search to sort the nodes topologically
-    std::function<void(int)> DepthFirstSearch = [&](int node_index)
-    {
-        if(visited[node_index]) return;
-        visited[node_index] = true;
-        auto node = nodes[node_index];
-        // Visit all nodes that this node depends on
-        for(auto input_field : node->InputFields())
-        {
-            for(int j=0; j < nnodes; j++)
-            {
-                auto other_node = nodes[j];
-                if(other_node == node) continue;
-                for(auto output_field : other_node->OutputFields())
-                {
-                    if(input_field->ID() == output_field->ID()) // Compare by unique ID
-                    {
-                        DepthFirstSearch(j);
-                    }
-                }
-            }
-        }
-        sorted_indices.push_back(node_index);
-    };
-
-    for(int i=0; i < nnodes; i++)
-    {
-        DepthFirstSearch(i);
-    }
-
-    nodes.Permute(sorted_indices);
-    node_owned.Permute(sorted_indices);
-
-    // Update the node indices after sorting
-    for(int i=0; i < nnodes; i++)
-    {
-        nodes[i]->SetNodeIndex(i);
-    }
-
-    sorted = true;
-}
-
-void DAGraph::ComputeDepth()
-{
-    // Compute depth of ordered nodes
-    node_depth.SetSize(nnodes);
-    node_depth = 0;
-    for(int i=0; i < nnodes; i++)
-    {
-        int max_depth = 0;
-        auto node = nodes[i];
-        for(auto input_field : node->InputFields())
-        {
-            for(int j=0; j < i; j++)
-            {
-                auto other_node = nodes[j];
-                if(other_node == node) continue;
-                for(auto output_field : other_node->OutputFields())
-                {
-                    if(input_field->ID() == output_field->ID()) // Compare by unique ID
-                    {
-                        max_depth = std::max(max_depth, node_depth[j] + 1);
-                    }
-                }
-            }
-        }
-        node_depth[i] = max_depth;
-    }
-}
-
-void DAGraph::CollectFieldMaps()
-{
-    MFEM_ASSERT(sorted, "DAGraph must be topologically sorted before collecting fields");
-
-    fid_to_index.clear();
-    fid_to_field.clear();
-
-    int nfields = 0;
-    for (auto f : InputFields())
-    {
-        fid_to_index.Register(f->ID(), nfields++);
-        fid_to_field.Register(f->ID(), f);
-    }
-
-    for (auto &node : nodes)
-    {
-        for (auto f : node->OutputFields())
-        {
-            if (!fid_to_index.Has(f->ID()))
-            {
-                fid_to_index.Register(f->ID(), nfields++);
-            }
-            if (!fid_to_field.Has(f->ID()))
-            {
-                fid_to_field.Register(f->ID(), f);
-            }
-        }
-    }
-    // TODO: Possibly add all intermediate fields from nodes to the graph's FieldCollection
-}
-
-void DAGraph::Mult(const Vector &x, Vector &y) const
-{
-    MFEM_ASSERT(width == x.Size(), "Input vector size (" << x.Size()
-                << ") must match matrix width (" << width << ")");
-
-    MFEM_ASSERT(height == y.Size(), "Output vector size (" << y.Size()
-                << ") must match matrix height (" << height << ")");
-
-    auto inputs  = InputFields();
-    auto outputs = OutputFields();
-
-    BlockVector xb(x.GetData(), input_offsets);
-    BlockVector yb(y.GetData(), output_offsets);
-    MultiVector xmv(inputs.Size()), ymv(outputs.Size());
-
-    // Set the data pointers of the input and output fields
-    // of the graph to point to the corresponding blocks of
-    // the input and output vectors
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        xmv.MakeRef(i, xb.GetBlock(i));
-    }
-
-    for(int i=0; i < outputs.Size(); i++)
-    {
-        ymv.MakeRef(i, yb.GetBlock(i));
-    }
-
-    const_cast<DAGraph*>(this)->Mult(xmv, ymv);
-}
-
-void DAGraph::Mult(const MultiVector &x, MultiVector &y)
-{
-    auto inputs  = InputFields();
-    auto outputs = OutputFields();
-
-    MFEM_ASSERT(inputs.Size() == x.NumBlocks(), "Number of input blocks (" << x.NumBlocks()
-                << ") must match number of input fields (" << inputs.Size() << ")");
-
-    MFEM_ASSERT(outputs.Size() == y.NumBlocks(), "Number of output blocks (" << y.NumBlocks()
-                << ") must match number of output fields (" << outputs.Size() << ")");
-
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        inputs[i]->SetData(const_cast<Vector*>(&x[i]));
-    }
-    for (int i=0; i < outputs.Size(); i++)
-    {
-        outputs[i]->SetData(&y[i]);
-    }
-
-    auto index_map = GetFieldIdToIndexMap();
-    auto fld_map = GetFieldIdToFieldMap();
-    int nfields = index_map.NumFields();
-    MultiVector ymv(nfields); // TODO: Should this be a member function?
-
-    // Assemble the multivector from the individual fields based on their IDs
-    // This multivector contains all input, output, and intermediate fields in the graph
-    for (auto const& [id, idx] : index_map)
-    {
-        if (fld_map.Has(id))
-        {
-            auto field = fld_map.Get(id);
-            ymv.MakeRef(idx, *field->Data());
-        }
-        else
-        {
-            MFEM_ABORT("Field ID " << id << " not found in field map");
-        }
-    }
-    
-    Execute(x, ymv);
-
-    for(auto &f : inputs)
-    {
-        f->SetData(nullptr);
-    }
-    for(auto &f : outputs)
-    {
-        f->SetData(nullptr);
-    }
-}
-
-void DAGraph::Execute(const MultiVector &x, MultiVector &y) const
-{
-    MFEM_ASSERT(assembled, "DAGraph must be assembled before calling Execute()");
-
-    MFEM_ASSERT(x.NumBlocks() == InputFields().Size(),
-                "Number of input blocks (" << x.NumBlocks()
-                << ") must match number of input fields (" << InputFields().Size() << ")");
-    
-    auto index_map = GetFieldIdToIndexMap();
-    
-    MFEM_ASSERT(y.NumBlocks() == index_map.NumFields(),
-                "Number of output blocks (" << y.NumBlocks()
-                << ") must match number of fields (" << index_map.NumFields() << ")");
-
-    auto inputs  = InputFields();
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        int idx = index_map.Get(inputs[i]->ID());
-        if(&y[idx] != &x[i]) // copy data, if address is different
-        {
-            y[idx] = x[i];
-        }
-    }
-
-    if(input_type == InputType::VECTOR)
-    {
-        x_node.SetSize(MaxWidth());
-        y_node.SetSize(MaxHeight());
-
-        for (auto node : nodes)
-        {
-            x_node.SetSize(node->Width());
-            y_node.SetSize(node->Height());
-
-            // Assemble input fields into a single vector for the node
-            auto node_inputs = node->InputFields();
-            auto ioffsets = node->InputOffsets();
-            for (int i=0; i < node_inputs.Size(); i++)
-            {
-                auto in_field = node_inputs[i];
-                int idx = index_map.Get(in_field->ID());
-                x_node.SetVector(y[idx],ioffsets[i]);
-            }
-
-            node->Mult(x_node, y_node);
-
-            // Disassemble output vector back
-            auto node_outputs = node->OutputFields();
-            BlockVector ynb(y_node.GetData(), node->OutputOffsets());
-            for (int i=0; i < node_outputs.Size(); i++)
-            {
-                auto out_field = node_outputs[i];
-                int idx = index_map.Get(out_field->ID());
-                y[idx] = ynb.GetBlock(i);
-            }
-        }
-    }
-    else if(input_type == InputType::MULTIVECTOR)
-    {
-        for (auto node : nodes)
-        {
-            auto node_inputs = node->InputFields();
-            auto node_outputs = node->OutputFields();
-            xmv_node.SetNumBlocks(node_inputs.Size());
-            ymv_node.SetNumBlocks(node_outputs.Size());
-
-            for (int i=0; i < node_inputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_inputs[i]->ID());
-                xmv_node.MakeRef(i, y[idx]);
-            }
-            for (int i=0; i < node_outputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_outputs[i]->ID());
-                ymv_node.MakeRef(i, y[idx]);
-            }
-            node->Mult(xmv_node, ymv_node);
-        }
-    }
-    else if(input_type == InputType::NONE)
-    {
-        Vector x_unused, y_unused;
-        for (auto node : nodes)
-        {
-            node->Mult(x_unused, y_unused);
-        }
-    }
-    else
-    {
-        MFEM_ABORT("DAGraph::Execute() not implemented for input type: "
-                    << static_cast<int>(input_type));
-    }
-}
-
-Operator& DAGraph::GetGradient(const Vector &x) const
-{
-    // TODO: Should/could be removed
-    if(grad_mode == GradMode::FINITE_DIFF)
-    {
-        if(!grad)
-        {
-            grad = new future::FDJacobian(*this, x, 1e-6);
-        }
-        else
-        {
-            grad->GetGradient(x); // Update the FDJacobian with new point x
-        }
-        return *grad;
-    }
-
-    MFEM_ASSERT(static_cast<int>(grad_mode) < static_cast<int>(GradMode::NONE),
-                "DAGraph::GetGradient() called with invalid grad_mode: "
-                << static_cast<int>(grad_mode));
-
-    if(!grad)
-    {
-        grad = new GraphGradient(const_cast<DAGraph&>(*this));
-    }
-
-    if(grad_mode == GradMode::ASSEMBLED)
-    {
-        return grad->GetGradient(x); // Assemble the Jacobian matrix
-    }
-    else // GradMode::MATRIX_FREE
-    {
-        dynamic_cast<GraphGradient*>(grad)->Update(x); // Update the GraphGradient with new point x
-    }
-
-    return *grad;
-}
-
-GraphGradient::GraphGradient(DAGraph &dag) : Operator(dag.Height(), dag.Width()),
-                                             graph(&dag)
-{
-    MFEM_ASSERT(graph->IsAssembled(), "GraphGradient requires an assembled DAGraph.");
-    MFEM_ASSERT(graph->IsSorted(), "GraphGradient requires a topologically sorted DAGraph.");
-
-    auto index_map = graph->GetFieldIdToIndexMap();
-    auto field_map = graph->GetFieldIdToFieldMap();
-
-    MFEM_ASSERT(index_map.NumFields() == field_map.NumFields(),
-                "Mismatch in number of fields between index_map and field_map");
-
-    x_work.DeleteAll(); // Clear any existing pointers
-    x_work.SetSize(index_map.NumFields());
-    x_work = nullptr; // Initialize all pointers to nullptr
-    xlin.SetNumBlocks(index_map.NumFields());
-
-    for (auto const& [id, idx] : index_map)
-    {
-        MFEM_ASSERT(idx >= 0 && idx < x_work.Size(), "Index out of bounds for field ID: " << id);
-        MFEM_ASSERT(field_map.Has(id), "Field ID not found in field_map: " << id);
-
-        if(x_work[idx] == nullptr)
-        {
-            x_work[idx] = new Vector(); // Allocate a new Vector for this field
-        }
-        xlin.MakeRef(idx, *x_work[idx]); // Make xlin refer to the allocated Vector
-    }
-
-}
-
-void GraphGradient::Update(const Vector &x)
-{
-    MFEM_ASSERT(graph != nullptr, "GraphGradient operator requires a non-null DAGraph pointer.");
-
-    auto set_exec_mode = [&](DAGraph::ExecutionMode mode)
-    {
-        for (auto &node : graph->Nodes())
-        {
-            node->SetExecutionMode(mode);
-        }
-    };
-
-    auto inputs = graph->InputFields();
-    BlockVector xb(x.GetData(), graph->InputOffsets());
-    MultiVector xmv(inputs.Size());
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        xmv.MakeRef(i, xb.GetBlock(i));
-    }
-
-    set_exec_mode(DAGraph::ExecutionMode::GRADIENT_MODE);
-    graph->Execute(xmv, xlin); // Forward pass to populate fields for gradient computations
-    set_exec_mode(DAGraph::ExecutionMode::DEFAULT_MODE); // Reset execution mode for forward pass
-}
-
-void GraphGradient::Mult(const Vector &x, Vector &y) const
-{
-    MFEM_ASSERT(x.Size() == graph->Width(), "Input vector size (" << x.Size()
-                << ") must match graph width (" << graph->Width() << ")");
-
-    MFEM_ASSERT(y.Size() == graph->Height(), "Output vector size (" << y.Size()
-                << ") must match graph height (" << graph->Height() << ")");
-
-    auto in_offsets  = graph->InputOffsets();
-    auto out_offsets = graph->OutputOffsets();
-
-    auto inputs  = graph->InputFields();
-    auto outputs = graph->OutputFields();
-
-    BlockVector xb(x.GetData(), in_offsets);
-    BlockVector yb(y.GetData(), out_offsets);
-    MultiVector xmv(inputs.Size()), ymv(outputs.Size());
-
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        xmv.MakeRef(i, xb.GetBlock(i));
-    }
-
-    for(int i=0; i < outputs.Size(); i++)
-    {
-        ymv.MakeRef(i, yb.GetBlock(i));
-    }
-
-    const_cast<GraphGradient*>(this)->Mult(xmv, ymv); // Forward mode: compute JVP, y = J(z) * x
-}
-
-void GraphGradient::Mult(const MultiVector &x, MultiVector &y)
-{
-    auto inputs  = graph->InputFields();
-    auto outputs = graph->OutputFields();
-
-    MFEM_ASSERT(inputs.Size() == x.NumBlocks(), "Number of input blocks (" << x.NumBlocks()
-                << ") must match number of input fields (" << inputs.Size() << ")");
-
-    MFEM_ASSERT(outputs.Size() == y.NumBlocks(), "Number of output blocks (" << y.NumBlocks()
-                << ") must match number of output fields (" << outputs.Size() << ")");
-
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        inputs[i]->SetAdjoint(const_cast<Vector*>(&x[i]));
-    }
-    for (int i=0; i < outputs.Size(); i++)
-    {
-        outputs[i]->SetAdjoint(&y[i]);
-    }
-
-    auto index_map = graph->GetFieldIdToIndexMap();
-    auto fld_map = graph->GetFieldIdToFieldMap();
-    int nfields = index_map.NumFields();
-    MultiVector ymv(nfields); // TODO: Should this be a member function?
-
-    // Assemble the multivector from the individual fields based on their IDs
-    // This multivector contains all input, output, and intermediate fields in the graph
-    for (auto const& [id, idx] : index_map)
-    {
-        if (fld_map.Has(id))
-        {
-            auto field = fld_map.Get(id);
-            ymv.MakeRef(idx, *field->Adjoint());
-        }
-        else
-        {
-            MFEM_ABORT("Field ID " << id << " not found in field map");
-        }
-    }
-
-    Forward(x, ymv); // Forward mode: compute JVP, y = J(z) * x
-
-    for (auto &f : inputs)
-    {
-        f->SetAdjoint(nullptr);
-    }
-    for (auto &f : outputs)
-    {
-        f->SetAdjoint(nullptr);
-    }
-}
-
-void GraphGradient::MultTranspose(const Vector &x, Vector &y) const
-{
-    MFEM_ASSERT(x.Size() == graph->Height(), "Input vector size (" << x.Size()
-                << ") must match graph height (" << graph->Height() << ")");
-    MFEM_ASSERT(y.Size() == graph->Width(), "Output vector size (" << y.Size()
-                << ") must match graph width (" << graph->Width() << ")");
-
-    auto in_offsets  = graph->InputOffsets();
-    auto out_offsets = graph->OutputOffsets();
-
-    auto inputs  = graph->InputFields();
-    auto outputs = graph->OutputFields();
-
-    BlockVector xb(x.GetData(), out_offsets);
-    BlockVector yb(y.GetData(), in_offsets);
-    MultiVector xmv(outputs.Size()), ymv(inputs.Size());
-
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        xmv.MakeRef(i, xb.GetBlock(i));
-    }
-
-    for(int i=0; i < outputs.Size(); i++)
-    {
-        ymv.MakeRef(i, yb.GetBlock(i));
-    }
-    const_cast<GraphGradient*>(this)->MultTranspose(xmv, ymv); // Reverse mode: compute VJP, y = J(z)^T * x
-}
-
-void GraphGradient::MultTranspose(const MultiVector &x, MultiVector &y)
-{
-    auto inputs  = graph->InputFields();
-    auto outputs = graph->OutputFields();
-
-    MFEM_ASSERT(outputs.Size() == x.NumBlocks(), "Number of input blocks (" << x.NumBlocks()
-                << ") must match number of output fields (" << outputs.Size() << ")");
-
-    MFEM_ASSERT(inputs.Size() == y.NumBlocks(), "Number of output blocks (" << y.NumBlocks()
-                << ") must match number of input fields (" << inputs.Size() << ")");
-
-    for(int i=0; i < outputs.Size(); i++)
-    {
-        outputs[i]->SetAdjoint(const_cast<Vector*>(&x[i]));
-    }
-    for (int i=0; i < inputs.Size(); i++)
-    {
-        inputs[i]->SetAdjoint(&y[i]);
-    }
-
-    auto index_map = graph->GetFieldIdToIndexMap();
-    auto fld_map = graph->GetFieldIdToFieldMap();
-    int nfields = index_map.NumFields();
-    MultiVector ymv(nfields); // TODO: Should this be a member function?
-
-    for(auto const& [id, idx] : index_map)
-    {
-        if (fld_map.Has(id))
-        {
-            auto field = fld_map.Get(id);
-            ymv.MakeRef(idx, *field->Adjoint());
-        }
-        else
-        {
-            MFEM_ABORT("Field ID " << id << " not found in field map");
-        }
-    }
-
-    Reverse(x, ymv); // Reverse mode: compute VJP, y = J(z)^T * x
-
-    for (auto &f : outputs)
-    {
-        f->SetAdjoint(nullptr);
-    }
-    for (auto &f : inputs)
-    {
-        f->SetAdjoint(nullptr);
-    }
-}
-
-void GraphGradient::Forward(const MultiVector &x, MultiVector &y) const
-{
-    MFEM_ASSERT(x.NumBlocks() == graph->InputFields().Size(),
-                "Number of input blocks (" << x.NumBlocks()
-                << ") must match number of input fields (" << graph->InputFields().Size() << ")");
-    
-    auto in_type = graph->GetInputType();
-    auto index_map  = graph->GetFieldIdToIndexMap();
-    auto field_map = graph->GetFieldIdToFieldMap();
-
-    MFEM_ASSERT(y.NumBlocks() == index_map.NumFields(),
-                "Number of output blocks (" << y.NumBlocks()
-                << ") must match number of fields (" << index_map.NumFields() << ")");
-
-    auto inputs  = graph->InputFields();
-    for(int i=0; i < inputs.Size(); i++)
-    {
-        int idx = index_map.Get(inputs[i]->ID());
-        if(&y[idx] != &x[i]) // copy data, if address is different
-        {
-            y[idx] = x[i];
-        }
-    }
-
-    if(in_type == InputType::VECTOR)
-    {
-        x0.SetSize(graph->MaxWidth());
-        dx.SetSize(graph->MaxWidth());
-        dy.SetSize(graph->MaxHeight());
-
-        auto nodes = graph->Nodes();
-        for (auto node : nodes)
-        {
-            x0.SetSize(node->Width());
-            dx.SetSize(node->Width());
-            dy.SetSize(node->Height());
-
-            // Assemble input fields into a single vector for the node
-            auto node_inputs = node->InputFields();
-            auto ioffsets = node->InputOffsets();
-            for(int i=0; i < node_inputs.Size(); i++)
-            {
-                auto in_field = node_inputs[i];
-                MFEM_ASSERT(index_map.Has(in_field->ID()), "Input field ID not found in index_map");
-                int idx = index_map.Get(in_field->ID());
-                x0.SetVector(xlin[idx], ioffsets[i]);
-                dx.SetVector(y[idx], ioffsets[i]);
-            }
-
-            node->GradientMult(x0, dx, dy); // Compute JVP for the node
-
-            // Disassemble output vector back
-            auto node_outputs = node->OutputFields();
-            BlockVector ynb(dy.GetData(), node->OutputOffsets());
-            for(int i=0; i < node_outputs.Size(); i++)
-            {
-                auto out_field = node_outputs[i];
-                MFEM_ASSERT(index_map.Has(out_field->ID()), "Output field ID not found in index_map");
-                int idx = index_map.Get(out_field->ID());
-                y[idx] = ynb.GetBlock(i);
-            }
-        }
-    }
-    else if(in_type == InputType::MULTIVECTOR)
-    {
-        auto nodes = graph->Nodes();
-        for (auto node : nodes)
-        {
-            auto node_inputs = node->InputFields();
-            auto node_outputs = node->OutputFields();
-            x0_mv.SetNumBlocks(node_inputs.Size());
-            dx_mv.SetNumBlocks(node_inputs.Size());
-            dy_mv.SetNumBlocks(node_outputs.Size());
-
-            for(int i=0; i < node_inputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_inputs[i]->ID());
-                x0_mv.MakeRef(i, xlin[idx]);
-                dx_mv.MakeRef(i, y[idx]);
-            }
-            for(int i=0; i < node_outputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_outputs[i]->ID());
-                dy_mv.MakeRef(i, y[idx]);
-            }
-            node->GradientMult(x0_mv, dx_mv, dy_mv); // Compute JVP for the node
-        }
-    }
-    else if(in_type == InputType::NONE)
-    {
-        Vector x_unused, dx_unused, dy_unused;
-        auto nodes = graph->Nodes();
-        for (auto node : nodes)
-        {
-            node->GradientMult(x_unused, dx_unused, dy_unused);
-        }
-    }
-    else
-    {
-        MFEM_ABORT("GraphGradient::Forward() not implemented for input type: "
-                    << static_cast<int>(in_type));
-    }
-}
-
-void GraphGradient::Reverse(const MultiVector &x, MultiVector &y) const
-{
-    MFEM_ASSERT(x.NumBlocks() == graph->OutputFields().Size(),
-                "Number of input blocks (" << x.NumBlocks()
-                << ") must match number of output fields (" << graph->OutputFields().Size() << ")");
-
-    auto in_type = graph->GetInputType();
-    auto index_map  = graph->GetFieldIdToIndexMap();
-    auto field_map = graph->GetFieldIdToFieldMap();
-    int nnodes  = graph->Size();
-
-    MFEM_ASSERT(y.NumBlocks() == index_map.NumFields(),
-                "Number of output blocks (" << y.NumBlocks()
-                << ") must match number of fields (" << index_map.NumFields() << ")");
-
-    auto outputs  = graph->OutputFields();
-    for(int i=0; i < outputs.Size(); i++)
-    {
-        int idx = index_map.Get(outputs[i]->ID());
-        if(&y[idx] != &x[i]) // copy data, if address is different
-        {
-            y[idx] = x[i];
-        }
-    }
-
-    if(in_type == InputType::VECTOR)
-    {
-        x0.SetSize(graph->MaxWidth());
-        dx.SetSize(graph->MaxHeight());
-        dy.SetSize(graph->MaxWidth());
-
-        for (int i=nnodes-1; i >= 0; i--)
-        {
-            auto node = graph->GetNode(i);
-            x0.SetSize(node->Width());
-            dx.SetSize(node->Height());
-            dy.SetSize(node->Width());
-
-            auto node_inputs = node->InputFields();
-            auto ioffsets = node->InputOffsets();
-            for(int i=0; i < node_inputs.Size(); i++)
-            {
-                auto in_field = node_inputs[i];
-                MFEM_ASSERT(index_map.Has(in_field->ID()), "Input field ID not found in index_map");
-                int idx = index_map.Get(in_field->ID());
-                x0.SetVector(xlin[idx], ioffsets[i]);
-            }
-
-            auto node_outputs = node->OutputFields();
-            auto ooffsets = node->OutputOffsets();
-            for(int i=0; i < node_outputs.Size(); i++)
-            {
-                auto out_field = node_outputs[i];
-                MFEM_ASSERT(index_map.Has(out_field->ID()), "Output field ID not found in index_map");
-                int idx = index_map.Get(out_field->ID());
-                dx.SetVector(y[idx], ooffsets[i]);
-            }
-
-            node->GradientMultTranspose(x0, dx, dy); // Compute JVP for the node
-
-            BlockVector dynb(dy.GetData(), node->InputOffsets());
-            for(int i=0; i < node_inputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_inputs[i]->ID());
-                y[idx] = dynb.GetBlock(i);
-            }
-        }
-    }
-    else if(in_type == InputType::MULTIVECTOR)
-    {
-        for (int i=nnodes-1; i >= 0; i--)
-        {
-            auto node = graph->GetNode(i);
-            auto node_inputs = node->InputFields();
-            auto node_outputs = node->OutputFields();
-            x0_mv.SetNumBlocks(node_inputs.Size());
-            dx_mv.SetNumBlocks(node_outputs.Size());
-            dy_mv.SetNumBlocks(node_inputs.Size());
-
-            for(int i=0; i < node_inputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_inputs[i]->ID());
-                x0_mv.MakeRef(i, xlin[idx]);
-                dy_mv.MakeRef(i, y[idx]);
-            }
-            for(int i=0; i < node_outputs.Size(); i++)
-            {
-                int idx = index_map.Get(node_outputs[i]->ID());
-                dx_mv.MakeRef(i, y[idx]);
-            }
-            node->GradientMultTranspose(x0_mv, dx_mv, dy_mv); // Compute JVP for the node
-        }
-    }
-    else if(in_type == InputType::NONE)
-    {
-        Vector x_unused, dx_unused, dy_unused;
-        for (int i=nnodes-1; i >= 0; i--)
-        {
-            auto node = graph->GetNode(i);
-            node->GradientMultTranspose(x_unused, dx_unused, dy_unused); // Compute VJP for the node
-        }
-    }
-    else
-    {
-        MFEM_ABORT("GraphGradient::Reverse() not implemented for input type: "
-                    << static_cast<int>(in_type));
-    }
-}
-
-Operator& GraphGradient::GetGradient(const Vector &x) const
-{
-    // Used to build Jacobian matrix
-    MFEM_ABORT("GraphGradient::GetGradient() not implemented");
-}
-
-
-} // namespace mfem

miniapps/multiapp/multiapp.hpp

@@ -1,838 +0,0 @@
-// 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.
-
-
-#ifndef MFEM_MULTIAPP_HPP
-#define MFEM_MULTIAPP_HPP
-
-#include "mfem.hpp"
-
-namespace mfem
-{
-
-/// Forward declarations needed below
-class Field;
-class FieldCollection;
-class GraphNode;
-class DAGraph;
-class GraphGradient;
-
-
-/// @brief Base class for storing data (Vector) and distinguishing
-/// fields variables
-class Field
-{
-public:
-    enum Type ///< Not used for now, but could be used to distinguish between input/output fields
-    {
-        INPUT , ///< Input field
-        OUTPUT, ///< Output field
-        DEFAULT ///< Any field
-    };
-
-    friend class GraphNode;
-
-private:
-    Type type = Type::DEFAULT;
-    inline static int next_id = 0;
-
-protected:
-    Vector *data = nullptr;
-    Vector *adjoint = nullptr; // For storing derivative info
-    int id = -1; // initialized to invalid id
-
-    std::string name; // Optional name for the field
-    Operator *oper  = nullptr; // Operator that outputs this field
-
-    int GetValidID(int id_, int lb=0, int ub = std::numeric_limits<int>::max())
-    {
-        return (id_ >= lb && id_ <= ub) ? id_ : next_id++;
-    }
-
-public:
-
-    ///@brief Constructor for a Field of type Type with optional ID
-    Field(Vector *field, Vector *adjoint, Type type, int id_ = -1) :
-          type(type), data(field), adjoint(adjoint), id(GetValidID(id_)),
-          name("Field_" + std::to_string(id)) { }
-
-    ///@brief Constructor for a Field of Default type with optional ID
-    Field(Vector *field, Vector *adjoint, int id_ = -1) : 
-          Field(field, adjoint, Type::DEFAULT, id_) { }
-
-    ///@brief Constructor for an input field
-    Field(Vector *field, int id_ = -1) :
-          Field(field, nullptr, Type::DEFAULT, id_) { }
-
-    ///@brief Constructor for a Field of type Type
-    Field(Vector *field, Type type, int id_ = -1) :
-          Field(field, nullptr, type, id_) { }
-
-    ///@brief Get the stored internally stored data pointer
-    Vector* Data() const { return data; }
-    Vector* Adjoint() const { return adjoint; }
-    Operator* GetOperator() const { return oper; }
-
-    ///@brief Set the internally stored data pointer
-    virtual void SetData(Vector *field) { data = field; }
-    virtual void SetAdjoint(Vector *adj) { adjoint = adj; }
-    virtual void SetOperator(Operator *op) { oper = op; }
-
-    virtual void GetData(Vector &field) const { field = *data; }
-    virtual void GetAdjoint(Vector &adj) const { adj = *adjoint; }
-
-    std::string Name() const { return name; }
-    void SetName(const std::string &n) { name = n; }
-    int ID() const { return id; }
-
-    void SetID(int i)
-    {
-        MFEM_ASSERT(i >= 0, "ID must be non-negative.");
-        id = i;
-    }
-
-    bool IsInput() const {return (type == Type::INPUT);}
-    bool IsOutput() const {return (type == Type::OUTPUT);}
-    bool IsDefault() const {return (type == Type::DEFAULT);}
-
-    virtual ~Field() = default;
-
-protected:
-
-    ///@brief Set the type of the field (prevents changing type of input/output fields)
-    void SetType(Type t)
-    {
-        type = t;
-    }
-};
-
-/// @brief A collection of Fields, each identified by a name
-class FieldCollection
-{
-public:
-    using FieldMap = GenericFieldMap<std::string, Field*>;
-    using IndexMap = GenericFieldMap<std::string, int>;
-
-private:
-    std::string name; /// Name of the collection
-    Operator *oper = nullptr; /// Operator associated with this collection (not owned)
-    FieldMap fields; /// Map from field name to Field pointer
-    IndexMap index_map; /// Map from field name to index in input/output vectors
-
-    Array<Field*> input_fields;  // Input fields for this node
-    Array<Field*> output_fields; // Output fields for this node
-
-public:
-
-    FieldCollection() = default;
-
-    /// @brief Constructor with collection name and optional associated operator
-    FieldCollection(std::string collection_name, Operator *op = nullptr):
-                    name(collection_name), oper(op) {}
-
-    /// @brief Constructor with associated operator and default collection name
-    FieldCollection(Operator *op) : name("FieldCollection"), oper(op) {}
-
-    /// @brief Get the number of fields in the collection
-    int Size() const { return fields.NumFields(); }
-
-    /// @brief Set the name of the collection
-    void SetName(const std::string &collection_name) { name = collection_name;}
-
-    /// @brief Get the name of the collection
-    std::string Name() const { return name; }
-
-    /// @brief Set the operator associated with this collection
-    void SetOperator(Operator *op){ oper = op; }
-
-    /// @brief Get the operator associated with this collection
-    const Operator* GetOperator() const { return oper; }
-
-    /// @brief Get the field associated with the given name, or nullptr if not found
-    Field* GetField(const std::string &field_name) const
-    {
-        return fields.Get(field_name);
-    }
-
-    /// @brief Add a field to the collection with a given name and ownership flag
-    void AddField(const std::string &field_name, Field *field, bool own = false)
-    {
-        if(fields.Has(field_name))
-        {
-            MFEM_WARNING("FieldCollection::AddField: Field with name "
-                         << field_name << " already exists. Replacing existing field.");
-        }
-        fields.Register(field_name, field, own);
-    }
-
-    void AddInput(const std::string &field_name,
-                  Field *field, bool own = false)
-    {
-        bool has_field = fields.Has(field_name);
-        bool has_index = index_map.Has(field_name);
-        if(has_field && has_index)
-        {
-            int i = index_map.Get(field_name);
-            input_fields[i] = field;
-        }
-        else
-        {
-            input_fields.push_back(field);
-            index_map.Register(field_name, input_fields.Size() - 1);
-        }
-        AddField(field_name, field, own);
-    }
-
-    void AddOutput(const std::string &field_name,
-                   Field *field, bool own = false)
-    {
-        bool has_field = fields.Has(field_name);
-        bool has_index = index_map.Has(field_name);
-        if(has_field && has_index)
-        {
-            int i = index_map.Get(field_name);
-            output_fields[i] = field;
-        }
-        else
-        {
-            output_fields.push_back(field);
-            index_map.Register(field_name, output_fields.Size() - 1);
-        }
-
-        AddField(field_name, field, own);
-        if(field->GetOperator() == nullptr)
-        {
-            field->SetOperator(oper);
-        }
-    }
-
-    Array<Field*>& InputFields() { return input_fields; }
-    Array<Field*>& OutputFields() { return output_fields; }
-
-    Field* InputField(int i) const { return input_fields[i]; }
-    Field *InputField(const std::string &field_name) const
-    {
-        bool has_index = index_map.Has(field_name);
-        if(!has_index)
-        {
-            MFEM_WARNING("FieldCollection::InputField: Field with name "
-                         << field_name << " does not exist in the collection.");
-            return nullptr;
-        }
-
-        int index = index_map.Get(field_name);
-        MFEM_VERIFY(index >= 0 && index < input_fields.Size(),
-                    "FieldCollection::InputField: Invalid index for field name: "
-                    << field_name << ".");
-        return input_fields[index];
-    }
-
-    Field* OutputField(int i) const { return output_fields[i]; }
-    Field *OutputField(const std::string &field_name) const
-    {
-        bool has_index = index_map.Has(field_name);
-        if(!has_index)
-        {
-            MFEM_WARNING("FieldCollection::OutputField: Field with name "
-                         << field_name << " does not exist in the collection.");
-            return nullptr;
-        }
-
-        int index = index_map.Get(field_name);
-        MFEM_VERIFY(index >= 0 && index < output_fields.Size(),
-                    "FieldCollection::OutputField: Invalid index for field name: "
-                    << field_name << ".");
-        return output_fields[index];
-    }
-
-    FieldMap &Fields() { return fields; }
-    FieldMap Fields() const { return fields; }
-
-    virtual void Save (std::ostream &out) const
-    {
-        out << "\"Fields\":\n";
-        out << "{\n";
-        for (auto f = fields.begin(); f != fields.end(); ++f)
-        {
-            std::string f_name = f->first;
-            Field *f_obj = f->second;
-            // out << "  " << f_name << ": ID " << f_obj->ID() << ",\n";
-            // out << f_obj->ID() << ": " << f_name << ",\n";
-            out << '\"' << f_obj->ID() << "\": \"" << f_name << "\"";
-            if(f != std::prev(fields.end())) out << ",";
-            out << "\n";
-        }
-        out << "},\n";
-
-        out << "\"Inputs\":\n";
-        out << "{\n";
-        for (int i = 0; i < input_fields.Size(); ++i)
-        {
-            Field *f_obj = input_fields[i];
-            out << '\"' << f_obj->ID() << "\": \"" << f_obj->Name() << "\"";
-            if(i != input_fields.Size() - 1) out << ",";
-            out << "\n";
-        }
-        out << "},\n";
-
-        out << "\"Outputs\":\n";
-        out << "{\n";
-        for (int i = 0; i < output_fields.Size(); ++i)
-        {
-            Field *f_obj = output_fields[i];
-            out << '\"' << f_obj->ID() << "\": \"" << f_obj->Name() << "\"";
-            if(i != output_fields.Size() - 1) out << ",";
-            out << "\n";
-        }
-        out << "}\n";
-    }
-
-    Field* HasField(const Field &field) const
-    {
-        for (auto f = fields.begin(); f != fields.end(); ++f)
-        {
-            if(f->second == &field)
-            {
-                return f->second;
-            }
-        }
-        return nullptr;
-    }
-
-    Field* HasField(const std::string &field_name) const
-    {
-        return fields.Get(field_name);
-    }
-
-    Field* HasField(const int id) const
-    {
-        for (auto f = fields.begin(); f != fields.end(); ++f)
-        {
-            if(f->second->ID() == id)
-            {
-                return f->second;
-            }
-        }
-        return nullptr;
-    }
-
-    ~FieldCollection(){}
-
-};
-
-
-class GraphNode : public Operator
-{
-public:
-    enum ExecutionMode
-    {
-        GRADIENT_MODE, ///< Node is being executed as part of a gradient evaluation
-        DEFAULT_MODE   ///< Node is being executed as default, operator evaluation
-    };
-
-private:
-    inline static int next_id = 0;
-
-protected:
-    int id = -1;
-    int node_index = -1;
-    mutable ExecutionMode exec_mode = DEFAULT_MODE;
-
-    std::string name;
-    mutable FieldCollection field_collection; ///< Collection of fields associated with this node
-
-    // Offsets to be used for operation on BlockVector
-    Array<int> input_offsets;  ///< Offsets for input fields
-    Array<int> output_offsets; ///< Offsets for output fields
-
-    int GetValidID(int id_, int lb=0, int ub = std::numeric_limits<int>::max())
-    {
-        return (id_ >= lb && id_ <= ub) ? id_ : next_id++;
-    }
-
-public:
-
-    GraphNode(int h, int w) : Operator(h,w), id(GetValidID(-1)),
-                              name("Node_" + std::to_string(id)),
-                              field_collection(this) { }
-
-    GraphNode(int s = 0) : GraphNode(s, s) { }
-
-    void SetNodeIndex(int index){ node_index = index; }
-    int GetNodeIndex() const { return node_index; }
-
-    void SetExecutionMode(ExecutionMode mode) { exec_mode = mode; }
-    ExecutionMode GetExecutionMode() const { return exec_mode; }
-
-    void SetName(const std::string &name_) { name = name_; }
-    std::string Name() const { return name; }
-
-    void SetID(int id_) { id = id_; }
-    int ID() const { return id; }
-
-    FieldCollection::FieldMap& Fields() { return field_collection.Fields(); }
-    Field* Fields(const std::string &f) { return field_collection.GetField(f); }
-
-    FieldCollection::FieldMap Fields() const { return field_collection.Fields(); }
-    Field* Fields(const std::string &f) const { return field_collection.GetField(f); }
-
-    Array<Field*>& InputFields() const { return field_collection.InputFields(); }
-    Array<Field*>& OutputFields() const { return field_collection.OutputFields(); }
-    Field* InputField(int i) const { return field_collection.InputField(i); }
-    Field* OutputField(int i) const { return field_collection.OutputField(i); }
-
-
-    virtual void AddInput(const std::string &field_name,
-                          Field *field, bool own = false)
-    { field_collection.AddInput(field_name, field, own); }
-
-    virtual void AddInput(Field *field, bool own = false)
-    { AddInput(field->Name(), field, own); }
-
-    template<bool OwnInputs = false,
-             typename... Args,
-             bool AreFields = std::conjunction<std::is_base_of<Field, std::remove_pointer_t<Args>> ...>::value,
-             typename std::enable_if<AreFields, bool>::type = true >
-    void AddInputs(Args... args)
-    {
-        ((AddInput(std::forward<Args>(args), OwnInputs)), ...);
-    }
-
-    virtual void AddOutput(const std::string &field_name,
-                           Field *field, bool own = false)
-    { field_collection.AddOutput(field_name, field, own); }
-
-    virtual void AddOutput(Field *field, bool own = false)
-    { AddOutput(field->Name(), field, own); }
-
-    template<bool OwnOutputs = false,
-             typename... Args,
-             bool AreFields = std::conjunction<std::is_base_of<Field, std::remove_pointer_t<Args>> ...>::value,
-             typename std::enable_if<AreFields, bool>::type = true >
-    void AddOutputs(Args... args)
-    {
-        ((AddOutput(std::forward<Args>(args), OwnOutputs)), ...);
-    }
-
-    virtual void Save (std::ostream &out) const
-    {
-        out << "\"Node-" << id << "\" : " << std::endl;
-        out << "{\n";
-        out << "\"Name\": \"" << name << "\",\n";
-        field_collection.Save(out);
-        out << "}";
-    }
-
-    virtual void Mult(const Vector &x, Vector &y) const override
-    {
-        MFEM_ABORT("GraphNode::Mult() not implemented");
-    }
-
-    virtual void Mult(const MultiVector &x, MultiVector &y) override
-    {
-        MFEM_ABORT("GraphNode::Mult(MultiVector) not implemented");
-    }
-
-    using Operator::GetGradient;
-
-    // TODO: Possibly remove this and only support MultiVector version of GradientMult
-    virtual void GradientMult(const Vector &x, const Vector &dx, Vector &dy) const
-    {
-        MFEM_ABORT("GraphNode::GradientMult() not implemented");
-        GetGradient(x).Mult(dx, dy);
-    }
-
-    virtual void GradientMult(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const
-    {
-        MFEM_ABORT("GraphNode::GradientMult() not implemented");
-        GetGradient(x).Mult(dx, dy);
-    }
-
-    // TODO: Possibly remove this and only support MultiVector version of GradientMultTranspose
-    virtual void GradientMultTranspose(const Vector &x, const Vector &dx, Vector &dy) const
-    {
-        MFEM_ABORT("GraphNode::GradientMultTranspose() not implemented");
-        GetGradient(x).MultTranspose(dx, dy);
-    }
-
-    virtual void GradientMultTranspose(const MultiVector &x, const MultiVector &dx, MultiVector &dy) const
-    {
-        MFEM_ABORT("GraphNode::GradientMultTranspose() not implemented");
-        // GetGradient(x).MultTranspose(dx, dy); // Not yet implemented
-    }
-
-    /// @brief Return the input offsets for block starts.
-    Array<int>& InputOffsets() { return input_offsets; }
-
-    /// @brief Read only access to the input offsets for block starts.
-    const Array<int>& InputOffsets() const { return input_offsets; }
-
-    void SetInputOffsets(const Array<int> &offsets) { input_offsets = offsets; }
-
-    /// @brief Return the output offsets for block starts.
-    Array<int>& OutputOffsets() { return output_offsets; }
-
-    /// @brief Read only access to the output offsets for block starts.
-    const Array<int>& OutputOffsets() const { return output_offsets; }
-
-    void SetOutputOffsets(const Array<int> &offsets) { output_offsets = offsets; }
-
-    virtual ~GraphNode() = default;
-};
-
-
-/**
-   @brief An abstract, type-erased class to define the interface for 
-   operators, not inherited from @a GraphNode. It performs SFINAE 
-   checks for stored operator's member functions and override the Mult
-   to call the stored object's functions.
- */
-template <typename OpType>
-class AbstractOperator : public GraphNode
-{
-protected:
-    /// Define a template class 'check' to test for the existence of member functions
-    template <typename C>
-    class CheckMember{
-        private:
-
-        /// @brief A type trait to check if the erased class has the function Mult
-        /// with the needed signatures.
-        template<class T>
-        using Mult = decltype(std::declval<T&>().Mult(std::declval<const Vector&>(),
-                                                      std::declval<Vector&>()));
-
-        template<class T>
-        using MultPtr = decltype(std::declval<T&>().Mult(std::declval<const int>(),
-                                                         std::declval<const real_t*>(),
-                                                         std::declval<const int>(),
-                                                         std::declval<real_t*>()));
-        // ---------------------------------------------------------------------
-        
-        template <typename T, template<typename> typename Func, typename R>
-        static constexpr auto Check(T*) -> typename std::is_same< Func<T>, R>::type;
-
-        template <typename, template<typename> typename, typename >
-        static constexpr std::false_type Check(...);
-
-        // --- Check for the existence of the member functions
-        typedef decltype(Check<C,Mult,void>(0)) Has_Mult;
-        typedef decltype(Check<C,MultPtr,void>(0)) Has_MultPtr;
-    public:
-        static constexpr bool HasMult  = Has_Mult::value;
-        static constexpr bool HasMultPtr  = Has_MultPtr::value;
-    };
-
-    OpType *op;  ///< Pointer to the operator
-
-public:
-
-    constexpr bool HasExecute(){return CheckMember<OpType>::HasStep;}
-    constexpr bool HasMult(){return CheckMember<OpType>::HasMult;}
-
-
-    /// @brief Constructor for the type-erased AbstractOperator class
-    AbstractOperator(OpType *op_, int h, int w) : GraphNode(h,w), op(op_)
-    { }
-
-    /// @brief Constructor for the type-erased AbstractOperator class.
-    AbstractOperator(OpType *op_, int s = 0) : AbstractOperator(op_,s,s) {}
-
-    /**
-       @brief Perform Mult operation with the stored operator, if it exists.
-     */
-    void Mult(const Vector &x, Vector &y) const override
-    {
-        if constexpr (CheckMember<OpType>::HasMult)
-        {
-            op->Mult(x,y);
-        }
-        else if constexpr (CheckMember<OpType>::HasMultPtr)
-        {
-            op->Mult(x.Size(), x.GetData(), y.Size(), y.GetData());
-        }
-        else
-        {
-            MFEM_ABORT("The AbstractOperator does not have the function, "
-                       "Mult(const Vector&, Vector&) or "
-                       "Mult(int, double*, int, double*).");
-        }
-    }
-};
-
-
-/**
-   @brief A class to store and coupled multiple operators together.
- */
-class DAGraph : public GraphNode
-{
-public:
-
-    using IntToIntMap = GenericFieldMap<int, int>;
-    using IntToFieldMap = GenericFieldMap<int, Field*>;
-
-    enum class GradMode
-    {
-        FINITE_DIFF = 0,  ///< Finite difference Jacobian
-        MATRIX_FREE = 1,  ///< Matrix-free Jacobian
-        ASSEMBLED = 2,    ///< Assembled Jacobian
-        NONE = 3          ///< Not implemented
-    };
-
-    enum InputType
-    {
-        VECTOR,      ///< Asemble the input blockvector from individual fields
-        MULTIVECTOR, ///< Asemble the multivector from individual fields
-        NONE         ///< No input
-    };
-
-protected:
-    Array<GraphNode*> nodes; ///< Vector of individual operators
-    Array<bool> node_owned; ///< Whether the operators are owned
-    Array<int> node_depth; ///< Depth of each operator in the graph
-
-    int max_width  = 0;     ///< Largest operator width
-    int max_height = 0;     ///< Largest operator height
-    int nnodes     = 0;     ///< The number of nodes
-    bool sorted    = false; ///< True if the nodes are topologically sorted
-    bool assembled = false; ///< True if the graph is assembled
-
-    GradMode grad_mode = GradMode::MATRIX_FREE; ///< Gradient mode for the graph
-    mutable Operator *grad = nullptr; ///< Gradient operator
-
-    InputType input_type = InputType::MULTIVECTOR; ///< Input type for the graph
-    mutable Vector x_node, y_node; ///< Temporary vectors for evaluating nodes
-    mutable MultiVector xmv_node, ymv_node; ///< Temporary multivectors for evaluating nodes
-
-    IntToFieldMap fid_to_field; ///< Map from Field ID to Field pointer
-    IntToIntMap fid_to_index; ///< Map from ID to index in an array; needed since ordering is not unique
-
-    friend class GraphGradient;
-
-public:
-    /**
-       @brief Construct a new CoupledOperator object.
-       @param nop Total number of operators to couple
-     */
-    DAGraph(const int nop) : GraphNode()
-    {
-        nodes.Reserve(nop);
-        node_owned.Reserve(nop);
-    }
-
-    /**
-       @brief Construct a new CoupledOperator object for an 
-       abstract non/mfem operator.
-     */
-    template <class OpType>
-    DAGraph(const OpType &op) : DAGraph(1)
-    {
-        AddOperator(op);
-    }
-
-    /**
-       @brief Add an operator to the list of coupled operator and
-       return pointer to it. Not owned unless it's not derived from GraphNode.
-     */
-    template <class OpType>
-    GraphNode* AddOperator(OpType *op_, int h, int w)
-    {
-        // Add operator to list of operators
-        if constexpr(std::is_base_of<GraphNode, OpType>::value)
-        {
-            nodes.push_back(op_);
-            node_owned.Append(false);
-        } 
-        else
-        {
-            nodes.push_back(new AbstractOperator<OpType>(op_,h,w));
-            node_owned.Append(true);
-        }
-        nnodes++;
-
-        // Update size of the coupled operator and the block offsets
-        GraphNode* op = nodes.Last();
-        op->SetNodeIndex(nnodes-1); // Set the index of the operator
-
-        int ht = op->Height();
-        int wt = op->Width();
-
-        max_width = std::max(max_width, wt);
-        max_height = std::max(max_height, ht);
-        sorted = false;
-
-        return op;
-    }
-
-    /// @brief Add an operator to the list of coupled operator and return pointer to it.
-    template <class OpType>
-    GraphNode* AddOperator(OpType *op_, int s = 0) { return AddOperator(op_,s,s);}
-
-    /// @brief Get the number of coupled operators
-    int Size(){return nnodes;}
-
-    /// @brief Get the size of the largest operator
-    int MaxWidth() const {return max_width;}
-    int MaxHeight() const {return max_height;}
-
-    IntToIntMap &GetFieldIdToIndexMap() { return fid_to_index; }
-    IntToIntMap GetFieldIdToIndexMap() const { return fid_to_index; }
-
-    IntToFieldMap &GetFieldIdToFieldMap() { return fid_to_field; }
-    IntToFieldMap GetFieldIdToFieldMap() const { return fid_to_field; }
-
-    /// @brief Get the operator at index @a i
-    GraphNode* GetNode(const int i)
-    {
-        MFEM_ASSERT(i >= 0 && i < nnodes,
-               "index [" << i << "] is out of range [0," << nnodes << ")");
-        return nodes[i];
-    }
-
-    Array<GraphNode*>& Nodes() { return nodes; }
-
-    /// @brief Specify whether the operator at index @a i is owned.
-    void OwnNode(const int i, bool own = true)
-    {
-        MFEM_ASSERT(i >= 0 && i < nnodes,
-               "index [" << i << "] is out of range [0," << nnodes << ")");
-        node_owned[i] = own;
-    }
-
-    void Assemble();
-    bool IsAssembled() const { return assembled; }
-
-    void TopologicalSort();
-    bool IsSorted() const { return sorted; }
-
-    void ComputeDepth();
-
-    void ValidateOffsets();
-
-    void ValidateNode(GraphNode &node);
-
-    void CollectFieldMaps();
-
-    using GraphNode::AddInput;
-    void AddInput(Field *field, int sz, bool own = false)
-    {
-        if(input_offsets.Size() == 0)
-        {   // First entry
-            input_offsets.Append(0);
-        }
-        input_offsets.Append(input_offsets.Last() + sz);
-        AddInput(field, own);
-    }
-
-    using GraphNode::AddOutput;
-    void AddOutput(Field *field, int sz, bool own = false)
-    {
-        if(output_offsets.Size() == 0)
-        {   // First entry
-            output_offsets.Append(0);
-        }
-        output_offsets.Append(output_offsets.Last() + sz);
-        AddOutput(field, own);
-    }
-
-    /// @brief Set the gradient mode for the coupled operator
-    void SetGradientMode(GradMode mode)
-    {
-        if(mode != grad_mode)
-        {
-            if(grad) { delete grad; grad = nullptr; }
-            grad_mode = mode;
-        }
-    }
-
-    void SetInputType(InputType type) { input_type = type; }
-    InputType GetInputType() const { return input_type; }
-
-    /**
-       @brief Apply the operator to the vector @a x 
-       and return the result in @a y.
-     */
-    virtual void Mult(const Vector &x, Vector &y) const override;
-
-    virtual void Mult(const MultiVector &x, MultiVector &y) override;
-
-    virtual void Execute(const MultiVector &x, MultiVector &y) const;
-
-    virtual void Save (std::ostream &out) const
-    {
-        out << "\"DAGraph\":\n";
-        out << "{\n";
-        // out << "\"nodes\" : " << nnodes << ",\n";
-        out << "\"Nodes\":\n";
-        out << "{\n";
-        for (int i = 0; i < nodes.Size(); i++)
-        {
-            nodes[i]->Save(out);
-            if(i != nodes.Size()-1) out << ",";
-            out << "\n";
-        }
-        out << "},\n"; // End of Nodes
-        field_collection.Save(out);
-        out << "}\n";
-    }
-
-    Operator& GetGradient(const Vector &x) const override;
-
-    /// @brief Destroy the Coupled Application object
-    ~DAGraph();
-};
-
-
-
-class GraphGradient : public Operator
-{
-public:
-    using InputType = DAGraph::InputType;
-
-protected:
-    mutable DAGraph *graph = nullptr; ///< Pointer to the DAGraph for which this is the gradient operator
-    Array<Vector*> x_work; ///< Array to store linearization point (intermediate fields)
-    mutable MultiVector xlin;
-    mutable Vector x0, dx, dy;
-    mutable MultiVector x0_mv, dx_mv, dy_mv;
-
-public:
-    GraphGradient(DAGraph &dag);
-
-    void Update(const Vector &x);
-
-    void Mult(const Vector &x, Vector &y) const override;
-
-    void Mult(const MultiVector &x, MultiVector &y) override;
-
-    void MultTranspose(const Vector &x, Vector &y) const override;
-
-    void MultTranspose(const MultiVector &x, MultiVector &y);
-
-    Operator &GetGradient(const Vector &x) const override;
-
-    void Forward(const MultiVector &x, MultiVector &y) const;
-
-    void Reverse(const MultiVector &x, MultiVector &y) const;
-
-    ~GraphGradient()
-    {
-        for (auto &v : x_work)
-        {
-            if(v) { delete v; v = nullptr; }
-        }
-        x_work.DeleteAll();
-    }
-};
-
-} //mfem namespace
-
-#endif

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