Add some comments

Big speedup, but can we do this in a portable way?

.gitignore

@@ -15,6 +15,19 @@
 CMakeCache.txt
 CMakeFiles/
 
+# NSYS files
+*.nsys-rep
+
+# Benchmark files
+*.tex
+*.org
+*.aux
+*.log
+*.gnuplot
+*.table
+*.pdf
+*.png
+
 # Backup files
 *~
 
@@ -316,6 +329,9 @@ tests/unit/debug_device_tests
 
 # Benchmark binaries
 tests/benchmarks/bench_ceed
+tests/benchmarks/bench_solvers
+tests/benchmarks/bench_linext
+tests/benchmarks/bench_lor_assembly
 tests/benchmarks/bench_tmop
 tests/benchmarks/bench_vector
 tests/benchmarks/bench_virtuals

CHANGELOG

@@ -81,6 +81,14 @@ Version 4.3.1 (development)
 - Added ParMesh Adjaceny Set (adjset) creation support to the Conduit Mesh
   Blueprint MFEM wrapper functions in ConduitDataCollection.
 
+- Added full assembly / device support for the LinearForm integrators:
+  * DomainLF: L(v) := (f, v)
+  * VectorDomainLF: L(v) := ((f1,...,fn), (v1,...,vn))
+  * DomainLFGrad: L(v) := (f,grad v)
+  * VectorDomainLFGrad:  L(v) := ((f1x,f1y,f1z,...,fnx,fny,fnz),grad(v1,...,vn))
+  It can be enabled by setting the linear assembly level to
+  'LinearAssemblyLevel::FULL' with the 'LinearForm::SetAssemblyLevel' method.
+
 - `HypreParVector` and `Vector` now support move semantics, and the copy
   constructor for `HypreParVector` now copies the local vector data.
 

CONTRIBUTING.md

@@ -141,6 +141,7 @@ The MFEM source code has the following structure:
   │   ├── tools
   │   └── toys
   └── tests
+      ├── benchmarks
       ├── convergence
       ├── gitlab
       ├── mem_manager

config/defaults.mk

@@ -428,7 +428,7 @@ GSLIB_LIB = -L$(GSLIB_DIR)/lib -lgs
 
 # CUDA library configuration
 CUDA_OPT =
-CUDA_LIB = -lcusparse
+CUDA_LIB = -lcusparse -lnvToolsExt
 
 # HIP library configuration (currently not needed)
 HIP_OPT =

examples/ex1.cpp

@@ -120,8 +120,8 @@ int main(int argc, char *argv[])
    //    largest number that gives a final mesh with no more than 50,000
    //    elements.
    {
-      int ref_levels =
-         (int)floor(log(50000./mesh.GetNE())/log(2.)/dim);
+      const int NE = Device::IsEnabled() ? 1e5 : 50000.;
+      const int ref_levels = (int)floor(log(NE/mesh.GetNE())/log(2.)/dim);
       for (int l = 0; l < ref_levels; l++)
       {
          mesh.UniformRefinement();
@@ -171,6 +171,7 @@ int main(int argc, char *argv[])
    LinearForm b(&fespace);
    ConstantCoefficient one(1.0);
    b.AddDomainIntegrator(new DomainLFIntegrator(one));
+   if (Device::IsEnabled()) { b.SetAssemblyLevel(LinearAssemblyLevel::FULL); }
    b.Assemble();
 
    // 8. Define the solution vector x as a finite element grid function
@@ -184,7 +185,10 @@ int main(int argc, char *argv[])
    //    domain integrator.
    BilinearForm a(&fespace);
    if (pa) { a.SetAssemblyLevel(AssemblyLevel::PARTIAL); }
-   a.AddDomainIntegrator(new DiffusionIntegrator(one));
+   // 'Default' E2E kernel
+   //a.AddDomainIntegrator(new DiffusionIntegrator(one));
+   // Can be E2E, L2L, libCeed, etc.
+   a.AddDomainIntegrator(new MassIntegrator(one));
 
    // 10. Assemble the bilinear form and the corresponding linear system,
    //     applying any necessary transformations such as: eliminating boundary
@@ -197,7 +201,7 @@ int main(int argc, char *argv[])
    Vector B, X;
    a.FormLinearSystem(ess_tdof_list, x, b, A, X, B);
 
-   cout << "Size of linear system: " << A->Height() << endl;
+   //cout << "Size of linear system: " << A->Height() << endl;
 
    // 11. Solve the linear system A X = B.
    if (!pa)
@@ -223,10 +227,70 @@ int main(int argc, char *argv[])
             ceed::AlgebraicSolver M(a, ess_tdof_list);
             PCG(*A, M, B, X, 1, 400, 1e-12, 0.0);
          }
-         else
+         /*else
          {
             OperatorJacobiSmoother M(a, ess_tdof_list);
             PCG(*A, M, B, X, 1, 400, 1e-12, 0.0);
+         }*/
+         else
+         {
+            const int myid = 0;
+            const int max_it = 50;
+            const int print_lvl = -1;
+            const double rtol = 1e-12;
+
+            CGSolver cg;
+            cg.SetRelTol(rtol);
+            cg.SetOperator(*A);
+
+            // Warm-up CG solve (in case of JIT to avoid timing it)
+            {
+               Vector Y(X);
+               cg.SetMaxIter(2);
+               cg.SetPrintLevel(-1);
+               cg.Mult(B, Y);
+               MFEM_DEVICE_SYNC;
+            }
+
+            // benchmark this problem
+            {
+               tic_toc.Clear();
+               cg.SetMaxIter(max_it);
+               cg.SetPrintLevel(print_lvl);
+               {
+                  tic_toc.Start();
+                  cg.Mult(B, X);
+                  MFEM_DEVICE_SYNC;
+                  tic_toc.Stop();
+               }
+            }
+            // Final norm check
+            const double final_norm = cg.GetFinalNorm();
+            MFEM_VERIFY(final_norm < sqrt(rtol),
+                        "FinalNorm (" << final_norm << ") Error!");
+            // Number of iteration check
+            const int num_iter = cg.GetNumIterations();
+            MFEM_VERIFY(num_iter <= max_it,
+                        "NumIterations (" << num_iter << ") Error!");
+
+            const double rt = tic_toc.RealTime();
+            const double rt_min = rt, rt_max = rt;
+            const int dofs = fespace.GetVSize();
+            const int cg_iter = cg.GetNumIterations();
+            const double mdofs_max = ((1e-6 * dofs) * cg_iter) / rt_max;
+            const double mdofs_min = ((1e-6 * dofs) * cg_iter) / rt_min;
+            if (myid == 0)
+            {
+               mfem::out << "Total CG time:    " << rt_max << " (" << rt_min << ") sec."
+                         << std::endl;
+               mfem::out << "Time per CG step: "
+                         << rt_max / cg_iter << " ("
+                         << rt_min / cg_iter << ") sec." << std::endl;
+               mfem::out << "\033[32m";
+               mfem::out << "\"DOFs/sec\" in CG: " << mdofs_max << " ("
+                         << mdofs_min << ") million.";
+               mfem::out << "\033[m" << std::endl;
+            }
          }
       }
       else

examples/ex1p.cpp

@@ -128,7 +128,7 @@ int main(int argc, char *argv[])
    //    more than 10,000 elements.
    {
       int ref_levels =
-         (int)floor(log(10000./mesh.GetNE())/log(2.)/dim);
+         (int)floor(log(1000./mesh.GetNE())/log(2.)/dim);
       for (int l = 0; l < ref_levels; l++)
       {
          mesh.UniformRefinement();
@@ -197,6 +197,7 @@ int main(int argc, char *argv[])
    ParLinearForm b(&fespace);
    ConstantCoefficient one(1.0);
    b.AddDomainIntegrator(new DomainLFIntegrator(one));
+   if (Device::IsEnabled()) { b.SetAssemblyLevel(LinearAssemblyLevel::FULL); }
    b.Assemble();
 
    // 10. Define the solution vector x as a parallel finite element grid
@@ -210,7 +211,8 @@ int main(int argc, char *argv[])
    //     Diffusion domain integrator.
    ParBilinearForm a(&fespace);
    if (pa) { a.SetAssemblyLevel(AssemblyLevel::PARTIAL); }
-   a.AddDomainIntegrator(new DiffusionIntegrator(one));
+   //a.AddDomainIntegrator(new DiffusionIntegrator(one));
+   a.AddDomainIntegrator(new MassIntegrator(one));
 
    // 12. Assemble the parallel bilinear form and the corresponding linear
    //     system, applying any necessary transformations such as: parallel
@@ -245,13 +247,67 @@ int main(int argc, char *argv[])
    {
       prec = new HypreBoomerAMG;
    }
+
+   const int max_it = 50;
+   const int print_lvl = -1;
+   const double rtol = 1e-12;
+
    CGSolver cg(MPI_COMM_WORLD);
-   cg.SetRelTol(1e-12);
-   cg.SetMaxIter(2000);
-   cg.SetPrintLevel(1);
+   cg.SetRelTol(rtol);
+   //cg.SetMaxIter(max_it);
+   //cg.SetPrintLevel(print_lvl);
    if (prec) { cg.SetPreconditioner(*prec); }
    cg.SetOperator(*A);
-   cg.Mult(B, X);
+
+   // Warm-up CG solve (in case of JIT to avoid timing it)
+   {
+      Vector Y(X);
+      cg.SetMaxIter(2);
+      cg.SetPrintLevel(-1);
+      cg.Mult(B, Y);
+      MFEM_DEVICE_SYNC;
+   }
+
+   // benchmark this problem
+   {
+      tic_toc.Clear();
+      cg.SetMaxIter(max_it);
+      cg.SetPrintLevel(print_lvl);
+      {
+         tic_toc.Start();
+         cg.Mult(B, X);
+         MFEM_DEVICE_SYNC;
+         tic_toc.Stop();
+      }
+   }
+   // Final norm check
+   const double final_norm = cg.GetFinalNorm();
+   MFEM_VERIFY(final_norm < sqrt(rtol),
+               "FinalNorm (" << final_norm << ") Error!");
+   // Number of iteration check
+   const int num_iter = cg.GetNumIterations();
+   MFEM_VERIFY(num_iter <= max_it,
+               "NumIterations (" << num_iter << ") Error!");
+
+   const double rt = tic_toc.RealTime();
+   const double rt_min = rt, rt_max = rt;
+   HYPRE_BigInt dofs = fespace.GlobalTrueVSize();
+   const int cg_iter = cg.GetNumIterations();
+   const double mdofs_max = ((1e-6 * dofs) * cg_iter) / rt_max;
+   const double mdofs_min = ((1e-6 * dofs) * cg_iter) / rt_min;
+   if (myid == 0)
+   {
+      mfem::out << "Total CG time:    " << rt_max << " (" << rt_min << ") sec."
+                << std::endl;
+      mfem::out << "Time per CG step: "
+                << rt_max / cg_iter << " ("
+                << rt_min / cg_iter << ") sec." << std::endl;
+      mfem::out << "\033[32m";
+      mfem::out << "\"DOFs/sec\" in CG: " << mdofs_max << " ("
+                << mdofs_min << ") million.";
+      mfem::out << "\033[m" << std::endl;
+   }
+   //cg.Mult(B, X);
    delete prec;
 
    // 14. Recover the parallel grid function corresponding to X. This is the

examples/ex6.cpp

@@ -125,6 +125,7 @@ int main(int argc, char *argv[])
    BilinearFormIntegrator *integ = new DiffusionIntegrator(one);
    a.AddDomainIntegrator(integ);
    b.AddDomainIntegrator(new DomainLFIntegrator(one));
+   if (Device::IsEnabled()) { b.SetAssemblyLevel(LinearAssemblyLevel::FULL); }
 
    // 7. The solution vector x and the associated finite element grid function
    //    will be maintained over the AMR iterations. We initialize it to zero.

examples/ex6p.cpp

@@ -202,6 +202,7 @@ int main(int argc, char *argv[])
    BilinearFormIntegrator *integ = new DiffusionIntegrator(one);
    a.AddDomainIntegrator(integ);
    b.AddDomainIntegrator(new DomainLFIntegrator(one));
+   if (Device::IsEnabled()) { b.SetAssemblyLevel(LinearAssemblyLevel::FULL); }
 
    // 12. The solution vector x and the associated finite element grid function
    //     will be maintained over the AMR iterations. We initialize it to zero.

fem/CMakeLists.txt

@@ -21,6 +21,8 @@ set(SRCS
   bilininteg_dgtrace_ea.cpp
   bilininteg_diffusion_mf.cpp
   bilininteg_diffusion_pa.cpp
+  bilininteg_diffusion_pa_diag_fast.cpp
+  bilininteg_diffusion_pa_fast.cpp
   bilininteg_diffusion_ea.cpp
   bilininteg_divergence.cpp
   bilininteg_hcurl.cpp
@@ -29,6 +31,11 @@ set(SRCS
   bilininteg_gradient.cpp
   bilininteg_mass_mf.cpp
   bilininteg_mass_pa.cpp
+  bilininteg_mass_pa_diag.cpp
+  bilininteg_mass_pa_fast.cpp
+  bilininteg_mass_pa_fast_amd.cpp
+  bilininteg_mass_pa_fast_hip.cpp
+  bilininteg_mass_pa_fast_diag.cpp
   bilininteg_mass_ea.cpp
   bilininteg_transpose_ea.cpp
   bilininteg_vecdiffusion.cpp
@@ -68,7 +75,16 @@ set(SRCS
   ceed/operator.cpp
   ceed/util.cpp
   linearform.cpp
+  linearform_ext.cpp
   lininteg.cpp
+  lininteg_domain.cpp
+  lininteg_domain_grad.cpp
+  lininteg_vector_domain.cpp
+  lininteg_vector_domain_grad.cpp
+  lor.cpp
+  lor_assembly.cpp
+  lor_assembly_ker.cpp
+  lor_assembly_ni3.cpp
   multigrid.cpp
   nonlinearform.cpp
   nonlinearform_ext.cpp
@@ -161,7 +177,12 @@ set(HDRS
   ceed/operator.hpp
   ceed/util.hpp
   linearform.hpp
+  linearform_ext.hpp
   lininteg.hpp
+  lininteg_domain.hpp
+  lininteg_domain_grad.hpp
+  lor.hpp
+  lor_assembly.hpp
   multigrid.hpp
   nonlinearform.hpp
   nonlinearform_ext.hpp

fem/bilinearform_ext.cpp

@@ -318,14 +318,31 @@ void PABilinearFormExtension::Assemble()
 void PABilinearFormExtension::AssembleDiagonal(Vector &y) const
 {
    Array<BilinearFormIntegrator*> &integrators = *a->GetDBFI();
+   const int integrators_size = integrators.Size();
 
-   const int iSz = integrators.Size();
-   if (elem_restrict && !DeviceCanUseCeed())
+   // Scan the different action types we need to launch
+   bool E2E_actions = false,
+        L2L_actions = false;
+   for (int i = 0; i < integrators_size; ++i)
+   {
+      const ActionType action_type = integrators[i]->GetActionType();
+      E2E_actions |= action_type == ActionType::E2E;
+      L2L_actions |= action_type == ActionType::L2L;
+   }
+
+   // typically this is a large vector, so store on device
+   y.UseDevice(true);
+
+   // If E2E kernels are present, do their computation on local vectors
+   if (E2E_actions && elem_restrict)
    {
       localY = 0.0;
-      for (int i = 0; i < iSz; ++i)
+      for (int i = 0; i < integrators_size; ++i)
       {
-         integrators[i]->AssembleDiagonalPA(localY);
+         if (integrators[i]->GetActionType() == ActionType::E2E)
+         {
+            integrators[i]->AssembleDiagonalPA(localY);
+         }
       }
       const ElementRestriction* H1elem_restrict =
          dynamic_cast<const ElementRestriction*>(elem_restrict);
@@ -340,11 +357,20 @@ void PABilinearFormExtension::AssembleDiagonal(Vector &y) const
    }
    else
    {
-      y.UseDevice(true); // typically this is a large vector, so store on device
+      // otherwise, initialize the y output
       y = 0.0;
-      for (int i = 0; i < iSz; ++i)
+   }
+
+   // Continue with the computation of the L2L integrators
+   if (L2L_actions || !elem_restrict)
+   {
+      for (int i = 0; i < integrators_size; ++i)
       {
-         integrators[i]->AssembleDiagonalPA(y);
+         const ActionType action = integrators[i]->GetActionType();
+         if (action == ActionType::L2L || !elem_restrict)
+         {
+            integrators[i]->AssembleDiagonalPA(y);
+         }
       }
    }
 }
@@ -383,27 +409,52 @@ void PABilinearFormExtension::FormLinearSystem(const Array<int> &ess_tdof_list,
 void PABilinearFormExtension::Mult(const Vector &x, Vector &y) const
 {
    Array<BilinearFormIntegrator*> &integrators = *a->GetDBFI();
+   const int integrators_size = integrators.Size();
 
-   const int iSz = integrators.Size();
-   if (DeviceCanUseCeed() || !elem_restrict)
+   // Scan the different action types we need to launch
+   bool E2E_actions = false,
+        L2L_actions = false;
+   for (int i = 0; i < integrators_size; ++i)
    {
-      y.UseDevice(true); // typically this is a large vector, so store on device
-      y = 0.0;
-      for (int i = 0; i < iSz; ++i)
-      {
-         integrators[i]->AddMultPA(x, y);
-      }
+      const ActionType action_type = integrators[i]->GetActionType();
+      E2E_actions |= action_type == ActionType::E2E;
+      L2L_actions |= action_type == ActionType::L2L;
    }
-   else
+
+   // typically this is a large vector, so store on device
+   y.UseDevice(true);
+
+   // If E2E kernels are present, do their computation on local vectors
+   if (E2E_actions && elem_restrict)
    {
       elem_restrict->Mult(x, localX);
       localY = 0.0;
-      for (int i = 0; i < iSz; ++i)
+      for (int i = 0; i < integrators_size; ++i)
       {
-         integrators[i]->AddMultPA(localX, localY);
+         if (integrators[i]->GetActionType() == ActionType::E2E)
+         {
+            integrators[i]->AddMultPA(localX, localY);
+         }
       }
       elem_restrict->MultTranspose(localY, y);
    }
+   else
+   {
+      // otherwise, initialize the y output
+      y = 0.0;
+   }
+
+   // Continue with the computation of the L2L integrators
+   if (L2L_actions || !elem_restrict)
+   {
+      for (int i = 0; i < integrators_size; ++i)
+      {
+         if (integrators[i]->GetActionType() == ActionType::L2L || !elem_restrict)
+         {
+            integrators[i]->AddMultPA(x, y);
+         }
+      }
+   }
 
    Array<BilinearFormIntegrator*> &intFaceIntegrators = *a->GetFBFI();
    const int iFISz = intFaceIntegrators.Size();

fem/bilininteg.hpp

@@ -244,6 +244,8 @@ public:
                                     Vector &flux, Vector *d_energy = NULL)
    { return 0.0; }
 
+   virtual bool SupportsBatchedLOR() const { return false; }
+
    virtual ~BilinearFormIntegrator() { }
 };
 
@@ -1977,7 +1979,7 @@ private:
    const DofToQuad *maps;         ///< Not owned
    const GeometricFactors *geom;  ///< Not owned
    int dim, ne, dofs1D, quad1D;
-   Vector pa_data;
+   Vector pa_data, CoG;
    bool symmetric = true; ///< False if using a nonsymmetric matrix coefficient
 
 public:
@@ -2058,6 +2060,8 @@ public:
                                          const FiniteElement &test_fe);
 
    bool SupportsCeed() const { return DeviceCanUseCeed(); }
+
+   bool SupportsBatchedLOR() const { return true; }
 };
 
 /** Class for local mass matrix assembling a(u,v) := (Q u, v) */

fem/bilininteg_convection_pa.cpp

@@ -768,8 +768,8 @@ void SmemPAConvectionApply3D(const int ne,
 
 void ConvectionIntegrator::AssemblePA(const FiniteElementSpace &fes)
 {
-   const MemoryType mt = (pa_mt == MemoryType::DEFAULT) ?
-                         Device::GetDeviceMemoryType() : pa_mt;
+   const MemoryType mt = (memory_type == MemoryType::DEFAULT) ?
+                         Device::GetDeviceMemoryType() : memory_type;
    // Assumes tensor-product elements
    Mesh *mesh = fes.GetMesh();
    const FiniteElement &el = *fes.GetFE(0);

fem/bilininteg_diffusion_pa.cpp

@@ -9,18 +9,45 @@
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
 
+#define MFEM_DEBUG_COLOR 87
+#include "../general/debug.hpp"
 #include "../general/forall.hpp"
 #include "bilininteg.hpp"
 #include "gridfunc.hpp"
 #include "ceed/diffusion.hpp"
+#include <string>
+
+#define MFEM_NVTX_COLOR Olive
+#include "../general/nvtx.hpp"
 
 using namespace std;
 
 namespace mfem
 {
 
+void NDK_PADiffusionAssembleDiagonal(const int dim,
+                                     const int D1D,
+                                     const int Q1D,
+                                     const int NE,
+                                     const bool symm,
+                                     const FiniteElementSpace *fes,
+                                     const DofToQuad *maps,
+                                     const Vector &D,
+                                     Vector &Y);
+
 // PA Diffusion Integrator
 
+void NDK_PADiffusionApply(const int dim,
+                          const int D1D,
+                          const int Q1D,
+                          const int NE,
+                          const Vector &CoG,
+                          const FiniteElementSpace *fes,
+                          const DofToQuad *maps,
+                          const Vector &D,
+                          const Vector &X,
+                          Vector &Y);
+
 // OCCA 2D Assemble kernel
 #ifdef MFEM_USE_OCCA
 static void OccaPADiffusionSetup2D(const int D1D,
@@ -349,10 +376,115 @@ static void PADiffusionSetup(const int dim,
    }
 }
 
+////////////////////////////////////////////////////////////////////////////////
+template<typename T> MFEM_HOST_DEVICE inline
+void HouseholderReflect(T *A, const T *v,
+                        const T b, const int m, const int n,
+                        const int row, const int col)
+{
+   for (int j = 0; j < n; j++)
+   {
+      T w = A[0*row + j*col];
+      for (int i = 1; i < m; i++) { w += v[i] * A[i*row + j*col]; }
+      A[0*row + j*col] -= b * w;
+      for (int i = 1; i < m; i++) { A[i*row + j*col] -= b * w * v[i]; }
+   }
+}
+
+template<int Q1D, typename T> MFEM_HOST_DEVICE inline
+void HouseholderApplyQ(T *A, const T *Q, const T *tau,
+                       const int k, const int row, const int col)
+{
+   T v[Q1D];
+   for (int ii=0; ii<k; ii++)
+   {
+      const int i = k-1-ii;
+      for (int j = i+1; j < Q1D; j++) { v[j] = Q[j*k+i]; }
+      // Apply Householder reflector (I - tau v v^T) coG^T
+      HouseholderReflect(&A[i*row], &v[i], tau[i], Q1D-i, Q1D, row, col);
+   }
+}
+
+template<int D1D, int Q1D, typename T> MFEM_HOST_DEVICE inline
+void QRFactorization(T *mat, T *tau)
+{
+   T v[Q1D];
+   DeviceMatrix B(mat, D1D, Q1D);
+   for (int i = 0; i < D1D; i++)
+   {
+      // Calculate Householder vector, magnitude
+      T sigma = 0.0;
+      v[i] = B(i,i);
+      for (int j = i + 1; j < Q1D; j++)
+      {
+         v[j] = B(i,j);
+         sigma += v[j] * v[j];
+      }
+      T norm = std::sqrt(v[i]*v[i] + sigma); // norm of v[i:m]
+      T Rii = -copysign(norm, v[i]);
+      v[i] -= Rii;
+      // norm of v[i:m] after modification above and scaling below
+      //   norm = sqrt(v[i]*v[i] + sigma) / v[i];
+      //   tau = 2 / (norm*norm)
+      tau[i] = 2 * v[i]*v[i] / (v[i]*v[i] + sigma);
+      for (int j=i+1; j<Q1D; j++) { v[j] /= v[i]; }
+      // Apply Householder reflector to lower right panel
+      HouseholderReflect(&mat[i*D1D+i+1], &v[i], tau[i],
+                         Q1D-i, D1D-i-1, D1D, 1);
+      // Save v
+      B(i,i) = Rii;
+      for (int j=i+1; j<Q1D; j++) { B(i,j) = v[j]; }
+   }
+}
+
+template<int D1D, int Q1D>
+void GetCollocatedGrad(const ConstDeviceMatrix &b,
+                       const ConstDeviceMatrix &g,
+                       const DeviceMatrix &CoG)
+{
+   double tau[Q1D];
+   double B1d[Q1D*D1D];
+   double G1d[Q1D*D1D];
+   DeviceMatrix B(B1d, D1D, Q1D);
+   DeviceMatrix G(G1d, D1D, Q1D);
+
+   for (int d = 0; d < D1D; d++)
+   {
+      for (int q = 0; q < Q1D; q++)
+      {
+         B(d,q) = b(q,d);
+         G(d,q) = g(q,d);
+      }
+   }
+   QRFactorization<D1D,Q1D>(B1d, tau);
+   // Apply Rinv, colograd1d = grad1d Rinv
+   for (int i = 0; i < Q1D; i++)
+   {
+      CoG(0,i) = G(0,i)/B(0,0);
+      for (int j = 1; j < D1D; j++)
+      {
+         CoG(j,i) = G(j,i);
+         for (int k = 0; k < j; k++) { CoG(j,i) -= B(j,k)*CoG(k,i); }
+         CoG(j,i) /= B(j,j);
+      }
+      for (int j = D1D; j < Q1D; j++) { CoG(j,i) = 0.0; }
+   }
+   // Apply Qtranspose, colograd = colograd Qtranspose
+   HouseholderApplyQ<Q1D>((double*)CoG, B1d, tau, D1D, 1, Q1D);
+}
+
+////////////////////////////////////////////////////////////////////////////////
 void DiffusionIntegrator::AssemblePA(const FiniteElementSpace &fes)
 {
-   const MemoryType mt = (pa_mt == MemoryType::DEFAULT) ?
-                         Device::GetDeviceMemoryType() : pa_mt;
+   const MemoryType mt = (memory_type == MemoryType::DEFAULT) ?
+                         Device::GetDeviceMemoryType() : memory_type;
+
+   // If device options allow fast kernels, set the action type to L2L
+   action_type =
+      (Device::FastKernelsEnabled())?
+      ActionType::L2L: // all fast kernel are L2L
+      ActionType::E2E; // default is E2E
+
    // Assuming the same element type
    fespace = &fes;
    Mesh *mesh = fes.GetMesh();
@@ -365,6 +497,7 @@ void DiffusionIntegrator::AssemblePA(const FiniteElementSpace &fes)
       MFEM_VERIFY(!VQ && !MQ && !SMQ,
                   "Only scalar coefficient supported for DiffusionIntegrator"
                   " with libCEED");
+      action_type = ActionType::L2L;
       ceedOp = new ceed::PADiffusionIntegrator(fes, *ir, Q);
       return;
    }
@@ -490,6 +623,40 @@ void DiffusionIntegrator::AssemblePA(const FiniteElementSpace &fes)
          }
       }
    }
+
+   if (Device::FastKernelsEnabled())
+   {
+      NVTX("CoG");
+      const int D1D = dofs1D;
+      const int Q1D = quad1D;
+      const int id = (D1D << 4) | Q1D;
+
+      CoG.SetSize(Q1D*Q1D);
+      CoG.UseDevice(true);
+      assert(CoG.UseDevice());
+
+      void (*KoG)(const ConstDeviceMatrix &b,
+                  const ConstDeviceMatrix &g,
+                  const DeviceMatrix &CoG) = nullptr;
+
+      switch (id) // orders 1~8
+      {
+         case 0x23: KoG=GetCollocatedGrad<2,3>; break; // 1
+         case 0x34: KoG=GetCollocatedGrad<3,4>; break; // 2
+         case 0x45: KoG=GetCollocatedGrad<4,5>; break; // 3
+         case 0x56: KoG=GetCollocatedGrad<5,6>; break; // 4
+         case 0x67: KoG=GetCollocatedGrad<6,7>; break; // 5
+         case 0x78: KoG=GetCollocatedGrad<7,8>; break; // 6
+         //case 0x89: KoG=GetCollocatedGrad<8,9>; break; // 7
+         //case 0x9A: KoG=GetCollocatedGrad<9,10>; break; // 8
+         default: MFEM_ABORT("Unknown kernel 0x" << std::hex << id << std::dec);
+      }
+
+      KoG(ConstDeviceMatrix(maps->B.HostRead(),Q1D,D1D),
+          ConstDeviceMatrix(maps->G.HostRead(),Q1D,D1D),
+          DeviceMatrix(CoG.HostReadWrite(),Q1D,Q1D));
+   }
+
    pa_data.SetSize((symmetric ? symmDims : MQfullDim) * nq * ne, mt);
    PADiffusionSetup(dim, sdim, dofs1D, quad1D, coeffDim, ne, ir->GetWeights(),
                     geom->J, coeff, pa_data);
@@ -869,6 +1036,7 @@ static void SmemPADiffusionDiagonal3D(const int NE,
                   }
                }
             }
+            MFEM_SYNC_THREAD;
          }
       }
    });
@@ -925,6 +1093,11 @@ void DiffusionIntegrator::AssembleDiagonalPA(Vector &diag)
    {
       ceedOp->GetDiagonal(diag);
    }
+   else if (Device::FastKernelsEnabled())
+   {
+      NDK_PADiffusionAssembleDiagonal(dim, dofs1D, quad1D, ne, symmetric,
+                                      fespace, maps, pa_data, diag);
+   }
    else
    {
       if (pa_data.Size()==0) { AssemblePA(*fespace); }
@@ -1180,6 +1353,7 @@ static void SmemPADiffusionApply2D(const int NE,
    auto D = Reshape(d_.Read(), Q1D*Q1D, symmetric ? 3 : 4, NE);
    auto x = Reshape(x_.Read(), D1D, D1D, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
+
    MFEM_FORALL_2D(e, NE, Q1D, Q1D, NBZ,
    {
       const int tidz = MFEM_THREAD_ID(z);
@@ -1220,6 +1394,7 @@ static void SmemPADiffusionApply2D(const int NE,
          }
       }
       MFEM_SYNC_THREAD;
+
       MFEM_FOREACH_THREAD(dy,y,D1D)
       {
          MFEM_FOREACH_THREAD(qx,x,Q1D)
@@ -1545,6 +1720,7 @@ static void SmemPADiffusionApply3D(const int NE,
                                    const int d1d = 0,
                                    const int q1d = 0)
 {
+   MFEM_NVTX;
    const int D1D = T_D1D ? T_D1D : d1d;
    const int Q1D = T_Q1D ? T_Q1D : q1d;
    constexpr int M1Q = T_Q1D ? T_Q1D : MAX_Q1D;
@@ -1595,6 +1771,7 @@ static void SmemPADiffusionApply3D(const int NE,
             }
          }
       }
+      MFEM_SYNC_THREAD;
       if (MFEM_THREAD_ID(z) == 0)
       {
          MFEM_FOREACH_THREAD(dy,y,D1D)
@@ -1791,6 +1968,7 @@ static void SmemPADiffusionApply3D(const int NE,
             }
          }
       }
+      MFEM_SYNC_THREAD;
    });
 }
 
@@ -1830,8 +2008,11 @@ static void PADiffusionApply(const int dim,
       switch (ID)
       {
          case 0x22: return SmemPADiffusionApply2D<2,2,16>(NE,symm,B,G,D,X,Y);
+         case 0x23: return SmemPADiffusionApply2D<2,3,16>(NE,symm,B,G,D,X,Y);
          case 0x33: return SmemPADiffusionApply2D<3,3,16>(NE,symm,B,G,D,X,Y);
+         case 0x34: return SmemPADiffusionApply2D<3,4,16>(NE,symm,B,G,D,X,Y);
          case 0x44: return SmemPADiffusionApply2D<4,4,8>(NE,symm,B,G,D,X,Y);
+         case 0x45: return SmemPADiffusionApply2D<4,5,8>(NE,symm,B,G,D,X,Y);
          case 0x55: return SmemPADiffusionApply2D<5,5,8>(NE,symm,B,G,D,X,Y);
          case 0x66: return SmemPADiffusionApply2D<6,6,4>(NE,symm,B,G,D,X,Y);
          case 0x77: return SmemPADiffusionApply2D<7,7,4>(NE,symm,B,G,D,X,Y);
@@ -1868,6 +2049,12 @@ void DiffusionIntegrator::AddMultPA(const Vector &x, Vector &y) const
    {
       ceedOp->AddMult(x, y);
    }
+   else if (Device::FastKernelsEnabled())
+   {
+      NDK_PADiffusionApply(dim, dofs1D, quad1D, ne, CoG,
+                           fespace, maps,
+                           pa_data, x, y);
+   }
    else
    {
       PADiffusionApply(dim, dofs1D, quad1D, ne, symmetric,

fem/bilininteg_diffusion_pa_diag_fast.cpp

@@ -0,0 +1,186 @@
+// Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#define MFEM_NVTX_COLOR Pink
+#include "../general/nvtx.hpp"
+
+#define MFEM_DEBUG_COLOR 206
+#include "../general/debug.hpp"
+#include "../general/forall.hpp"
+#include "bilininteg.hpp"
+#include "gridfunc.hpp"
+#include "ceed/diffusion.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+// Shared memory PA Diffusion Diagonal 3D kernel
+template<int D1D, int Q1D>
+static void NDK_SmemPADiffusionDiag3D(const int ndofs,
+                                      const int NE,
+                                      const bool symmetric,
+                                      const int *map_,
+                                      const double *b_,
+                                      const double *g_,
+                                      const double *d_,
+                                      double *y_)
+{
+   MFEM_NVTX;
+   constexpr int DIM = 3;
+   const auto MAP = Reshape(map_, D1D,D1D,D1D, NE);
+   auto b = Reshape(b_, Q1D, D1D);
+   auto g = Reshape(g_, Q1D, D1D);
+   auto D = Reshape(d_, Q1D*Q1D*Q1D, symmetric ? 6 : 9, NE);
+   auto Y = Reshape(y_, D1D, D1D, D1D, NE);
+
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, Q1D,
+   {
+      const int tidz = MFEM_THREAD_ID(z);
+      MFEM_SHARED double BG[2][Q1D*D1D];
+      double (*B)[D1D] = (double (*)[D1D]) (BG+0);
+      double (*G)[D1D] = (double (*)[D1D]) (BG+1);
+      MFEM_SHARED double QQD[Q1D][Q1D][D1D];
+      MFEM_SHARED double QDD[Q1D][D1D][D1D];
+      if (tidz == 0)
+      {
+         MFEM_FOREACH_THREAD(d,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(q,x,Q1D)
+            {
+               B[q][d] = b(q,d);
+               G[q][d] = g(q,d);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      for (int i = 0; i < DIM; ++i)
+      {
+         for (int j = 0; j < DIM; ++j)
+         {
+            // first tensor contraction, along z direction
+            MFEM_FOREACH_THREAD(qx,x,Q1D)
+            {
+               MFEM_FOREACH_THREAD(qy,y,Q1D)
+               {
+                  MFEM_FOREACH_THREAD(dz,z,D1D)
+                  {
+                     QQD[qx][qy][dz] = 0.0;
+                     for (int qz = 0; qz < Q1D; ++qz)
+                     {
+                        const int q = qx + (qy + qz * Q1D) * Q1D;
+                        const int ksym = j >= i ?
+                                         3 - (3-i)*(2-i)/2 + j:
+                                         3 - (3-j)*(2-j)/2 + i;
+                        const int k = symmetric ? ksym : (i*DIM) + j;
+                        const double O = D(q,k,e);
+                        const double Bz = B[qz][dz];
+                        const double Gz = G[qz][dz];
+                        const double L = i==2 ? Gz : Bz;
+                        const double R = j==2 ? Gz : Bz;
+                        QQD[qx][qy][dz] += L * O * R;
+                     }
+                  }
+               }
+            }
+            MFEM_SYNC_THREAD;
+            // second tensor contraction, along y direction
+            MFEM_FOREACH_THREAD(qx,x,Q1D)
+            {
+               MFEM_FOREACH_THREAD(dz,z,D1D)
+               {
+                  MFEM_FOREACH_THREAD(dy,y,D1D)
+                  {
+                     QDD[qx][dy][dz] = 0.0;
+                     for (int qy = 0; qy < Q1D; ++qy)
+                     {
+                        const double By = B[qy][dy];
+                        const double Gy = G[qy][dy];
+                        const double L = i==1 ? Gy : By;
+                        const double R = j==1 ? Gy : By;
+                        QDD[qx][dy][dz] += L * QQD[qx][qy][dz] * R;
+                     }
+                  }
+               }
+            }
+            MFEM_SYNC_THREAD;
+            // third tensor contraction, along x direction
+            MFEM_FOREACH_THREAD(dz,z,D1D)
+            {
+               MFEM_FOREACH_THREAD(dy,y,D1D)
+               {
+                  MFEM_FOREACH_THREAD(dx,x,D1D)
+                  {
+                     for (int qx = 0; qx < Q1D; ++qx)
+                     {
+                        const double Bx = B[qx][dx];
+                        const double Gx = G[qx][dx];
+                        const double L = i==0 ? Gx : Bx;
+                        const double R = j==0 ? Gx : Bx;
+                        const double lvr = L * QDD[qx][dy][dz] * R;
+                        const int gid = MAP(dx,dy,dz,e);
+                        const int idx = gid >= 0 ? gid : -1 - gid;
+                        AtomicAdd(Y[idx], lvr);
+                     }
+                  }
+               }
+            }
+            MFEM_SYNC_THREAD;
+         }
+      }
+   });
+}
+
+void NDK_PADiffusionAssembleDiagonal(const int dim,
+                                     const int D1D,
+                                     const int Q1D,
+                                     const int NE,
+                                     const bool symm,
+                                     const FiniteElementSpace *fes,
+                                     const DofToQuad *maps,
+                                     const Vector &D,
+                                     Vector &Y)
+{
+   MFEM_NVTX;
+   const int ND = fes->GetNDofs();
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes->GetElementRestriction(ordering);
+   const ElementRestriction* ER = dynamic_cast<const ElementRestriction*>(ERop);
+   assert(ER);
+   const int *m = ER->GatherMap().Read();
+   const double *b = maps->B.Read();
+   const double *g = maps->G.Read();
+   const double *d = D.Read();
+   double *y = Y.ReadWrite();
+
+   if (dim == 2) {  assert(false); }
+   else if (dim == 3)
+   {
+      switch ((D1D << 4 ) | Q1D)
+      {
+         case 0x22: return NDK_SmemPADiffusionDiag3D<2,2>(ND,NE,symm,m,b,g,d,y);
+         case 0x23: return NDK_SmemPADiffusionDiag3D<2,3>(ND,NE,symm,m,b,g,d,y);
+         case 0x34: return NDK_SmemPADiffusionDiag3D<3,4>(ND,NE,symm,m,b,g,d,y);
+         case 0x45: return NDK_SmemPADiffusionDiag3D<4,5>(ND,NE,symm,m,b,g,d,y);
+         case 0x46: return NDK_SmemPADiffusionDiag3D<4,6>(ND,NE,symm,m,b,g,d,y);
+         case 0x56: return NDK_SmemPADiffusionDiag3D<5,6>(ND,NE,symm,m,b,g,d,y);
+         case 0x67: return NDK_SmemPADiffusionDiag3D<6,7>(ND,NE,symm,m,b,g,d,y);
+         case 0x78: return NDK_SmemPADiffusionDiag3D<7,8>(ND,NE,symm,m,b,g,d,y);
+         //case 0x89: return NDK_SmemPADiffusionDiag3D<8,9>(ND,NE,symm,m,b,g,d,y);
+         //case 0x9A: return NDK_SmemPADiffusionDiag3D<9,10>(ND,NE,symm,m,b,g,d,y);
+         default: MFEM_ABORT("Not implemented!");
+      }
+   }
+   MFEM_ABORT("Unknown kernel.");
+}
+
+} // namespace mfem

fem/bilininteg_diffusion_pa_fast.cpp

@@ -0,0 +1,332 @@
+// Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#define MFEM_NVTX_COLOR Pink
+#include "../general/nvtx.hpp"
+
+#define MFEM_DEBUG_COLOR 206
+#include "../general/debug.hpp"
+#include "../general/forall.hpp"
+#include "bilininteg.hpp"
+#include "gridfunc.hpp"
+#include "ceed/diffusion.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+////////////////////////////////////////////////////////////////////////////////
+template<int D1D, int Q1D, int NBZ, int NBK> static
+MFEM_GLOBAL MFEM_LAUNCH_BOUNDS(Q1D*Q1D*NBZ,NBK)
+void NDK_PADiffApply(const int NE,
+                     const int *MAP,
+                     const double *B,
+                     const double *G,
+                     const double *D,
+                     const double *X,
+                     double *Y)
+{
+   double r_qt, r_q[Q1D];
+   MFEM_SHARED double s_B[D1D][Q1D];
+   MFEM_SHARED double s_G[Q1D][Q1D];
+   MFEM_SHARED double s_Iq[NBZ][Q1D][Q1D][Q1D];
+   MFEM_SHARED double s_Gqr[NBZ][Q1D][Q1D];
+   MFEM_SHARED double s_Gqs[NBZ][Q1D][Q1D];
+
+   for (int be = MFEM_BLOCK_ID(x); be < (NE+NBZ-1)/NBZ; be += MFEM_GRID_DIM(x))
+   {
+      const int tz = MFEM_THREAD_ID(z);
+      const int e = be * NBZ + tz;
+      if (e>=NE) { return; }
+
+      // Scatter X
+      MFEM_FOREACH_THREAD(j,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            if (tz == 0) { s_G[j][i] = G[i+Q1D*j]; } // ok with init
+            if (tz == 0 && j<D1D) { s_B[j][i] = B[i+Q1D*j]; } // ok
+            if (j<D1D && i<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int k = 0; k < D1D; k++)
+               {
+                  const int gid = MAP[e*D1D*D1D*D1D + k*D1D*D1D + j*D1D + i];
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  r_q[k] = X[idx];
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Grad1X
+      MFEM_FOREACH_THREAD(b,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,D1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int k=0; k<Q1D; ++k)
+            {
+               double u = 0.0;
+               MFEM_UNROLL(D1D)
+               for (int c=0; c<D1D; ++c) { u += s_B[c][k] * r_q[c]; }
+               s_Iq[tz][k][b][a] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Grad1Y
+      MFEM_FOREACH_THREAD(k,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int b=0; b<D1D; ++b) { r_q[b] = s_Iq[tz][k][b][a]; }
+            MFEM_UNROLL(Q1D)
+            for (int j=0; j<Q1D; ++j)
+            {
+               double u = 0.0;
+               MFEM_UNROLL(D1D)
+               for (int b=0; b<D1D; ++b) { u += s_B[b][j] * r_q[b]; }
+               s_Iq[tz][k][j][a] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Grad1Z
+      MFEM_FOREACH_THREAD(k,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(j,x,Q1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int a=0; a<D1D; ++a) { r_q[a] = s_Iq[tz][k][j][a]; }
+            MFEM_UNROLL(Q1D)
+            for (int i=0; i<Q1D; ++i)
+            {
+               double u = 0.0;
+               MFEM_UNROLL(D1D)
+               for (int a=0; a<D1D; ++a) { u += s_B[a][i] * r_q[a]; }
+               s_Iq[tz][k][j][i] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Flush
+      MFEM_FOREACH_THREAD(j,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int k = 0; k < Q1D; ++k) { r_q[k] = 0.0; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Q-Function
+      MFEM_UNROLL(Q1D)
+      for (int k = 0; k < Q1D; ++k)
+      {
+         MFEM_SYNC_THREAD;
+         MFEM_FOREACH_THREAD(j,y,Q1D)
+         {
+            MFEM_FOREACH_THREAD(i,x,Q1D)
+            {
+               double qr = 0.0, qs = 0.0, qt = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int m = 0; m < Q1D; ++m)
+               {
+                  const double Dim = s_G[i][m];
+                  const double Djm = s_G[j][m];
+                  const double Dkm = s_G[k][m];
+                  qr += Dim * s_Iq[tz][k][j][m];
+                  qs += Djm * s_Iq[tz][k][m][i];
+                  qt += Dkm * s_Iq[tz][m][j][i];
+               }
+               //(d, Q1D,Q1D,Q1D, 6, NE);
+               const int ebase = e * 6*Q1D*Q1D*Q1D;
+               const int gbase = i + j*Q1D + k*Q1D*Q1D;
+               const double D00 = D[0*Q1D*Q1D*Q1D + gbase + ebase];
+               const double D01 = D[1*Q1D*Q1D*Q1D + gbase + ebase];
+               const double D02 = D[2*Q1D*Q1D*Q1D + gbase + ebase];
+               const double D11 = D[3*Q1D*Q1D*Q1D + gbase + ebase];
+               const double D12 = D[4*Q1D*Q1D*Q1D + gbase + ebase];
+               const double D22 = D[5*Q1D*Q1D*Q1D + gbase + ebase];
+
+               s_Gqr[tz][j][i] = D00*qr + D01*qs + D02*qt;
+               s_Gqs[tz][j][i] = D01*qr + D11*qs + D12*qt;
+               r_qt = D02*qr + D12*qs + D22*qt;
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(j,y,Q1D)
+         {
+            MFEM_FOREACH_THREAD(i,x,Q1D)
+            {
+               double Aqtmp = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int m = 0; m < Q1D; ++m)
+               {
+                  const double Dmi = s_G[m][i];
+                  const double Dmj = s_G[m][j];
+                  const double Dkm = s_G[k][m];
+                  Aqtmp += Dmi * s_Gqr[tz][j][m];
+                  Aqtmp += Dmj * s_Gqs[tz][m][i];
+                  r_q[m] += Dkm * r_qt;
+               }
+               r_q[k] += Aqtmp;
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+
+      // GradZT
+      MFEM_FOREACH_THREAD(j,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               double u = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int k=0; k<Q1D; ++k) { u += s_B[c][k] * r_q[k]; }
+               s_Iq[tz][c][j][i] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // GradYT
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int j=0; j<Q1D; ++j) { r_q[j] = s_Iq[tz][c][j][i]; }
+            MFEM_UNROLL(D1D)
+            for (int b=0; b<D1D; ++b)
+            {
+               double u = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int j=0; j<Q1D; ++j) { u += s_B[b][j] * r_q[j]; }
+               s_Iq[tz][c][b][i] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // GradXT
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(b,x,D1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int i=0; i<Q1D; ++i) { r_q[i] = s_Iq[tz][c][b][i]; }
+            MFEM_UNROLL(D1D)
+            for (int a=0; a<D1D; ++a)
+            {
+               double u = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int i=0; i<Q1D; ++i) { u += s_B[a][i] * r_q[i]; }
+               s_Iq[tz][c][b][a] = u;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Gather
+      MFEM_FOREACH_THREAD(j,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int k = 0; k < D1D; k++)
+            {
+               const int gid = MAP[e*D1D*D1D*D1D + k*D1D*D1D + j*D1D + i];
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               const double output = s_Iq[tz][k][j][i];
+               AtomicAdd(Y[idx], output);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   }
+}
+
+void NDK_PADiffusionApply(const int dim,
+                          const int D1D,
+                          const int Q1D,
+                          const int NE,
+                          const Vector &CoG,
+                          const FiniteElementSpace *fes,
+                          const DofToQuad *maps,
+                          const Vector &D,
+                          const Vector &X,
+                          Vector &Y)
+{
+   //dbg();
+   MFEM_NVTX;
+   assert(dim == 3);
+
+   const double *b = maps->B.Read();
+   const double *d = D.Read();
+   const double *x = X.Read();
+   double *y = Y.ReadWrite();
+
+   const int id = (D1D << 4) | Q1D;
+
+   const int ND = fes->GetNDofs();
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes->GetElementRestriction(ordering);
+   const ElementRestriction *ER = dynamic_cast<const ElementRestriction*>(ERop);
+   assert(ER);
+   const int *map = ER->GatherMap().Read();
+
+   const auto dM = Reshape(map, D1D,D1D,D1D, NE);
+   const auto dB = Reshape(b, Q1D,D1D);
+   const auto dG = Reshape(CoG.Read(), Q1D,Q1D);
+   const auto dD = Reshape(d, Q1D,Q1D,Q1D, 6, NE);
+   const auto dX = Reshape(x, ND);
+   auto dY = Reshape(y, ND);
+
+   void (*Ker)(const int NE,
+               const int *MAP,
+               const double *B,
+               const double *G,
+               const double *D,
+               const double *X,
+               double *Y) = nullptr;
+
+   int NBZ = 1;
+
+   switch (id) // orders 1~8
+   {
+      case 0x23: Ker=NDK_PADiffApply<2,3,16,5>; NBZ=16; break; // 1
+      case 0x34: Ker=NDK_PADiffApply<3,4,8,6>; NBZ=8; break; // 2
+      case 0x45: Ker=NDK_PADiffApply<4,5,4,4>; NBZ=4; break; // 3
+      case 0x56: Ker=NDK_PADiffApply<5,6,4,3>; NBZ=4; break; // 4
+      case 0x67: Ker=NDK_PADiffApply<6,7,1,3>; break; // 5
+      case 0x78: Ker=NDK_PADiffApply<7,8,1,0>; break; // 6
+      //case 0x89: Ker=NDK_PADiffApply<8,9,1,0>; break; // 7
+      //case 0x9A: Ker=NDK_PADiffApply<9,10,1,1>; break; // 8
+      default: MFEM_ABORT("Unknown kernel 0x" << std::hex << id << std::dec);
+   }
+   MFEM_CONTRACT_VAR(NBZ);
+   MFEM_LAUNCH_KERNEL(Ker,(NE+NBZ-1)/NBZ,dim3(Q1D,Q1D,NBZ),0,NE,dM,dB,dG,dD,dX,dY);
+}
+
+} // namespace mfem

fem/bilininteg_mass_pa.cpp

@@ -19,14 +19,49 @@ using namespace std;
 namespace mfem
 {
 
+// Forward declaration of non-deterministic 'fast' kernels
+void NDK_PAMassApply(const int dim,
+                     const int D1D,
+                     const int Q1D,
+                     const int NE,
+                     const FiniteElementSpace *fes,
+                     const DofToQuad *maps,
+                     const Vector &D,
+                     const Vector &X,
+                     Vector &Y);
+
+void NDK_AMD_PAMassApply(const int dim,
+                         const int D1D,
+                         const int Q1D,
+                         const int NE,
+                         const FiniteElementSpace *fes,
+                         const DofToQuad *maps,
+                         const Vector &D,
+                         const Vector &X,
+                         Vector &Y);
+
+void NDK_HIP_PAMassApply(const int dim,
+                         const int D1D,
+                         const int Q1D,
+                         const int NE,
+                         const FiniteElementSpace *fes,
+                         const DofToQuad *maps,
+                         const Vector &D,
+                         const Vector &X,
+                         Vector &Y);
+
 // PA Mass Integrator
 
 // PA Mass Assemble kernel
 
 void MassIntegrator::AssemblePA(const FiniteElementSpace &fes)
 {
-   const MemoryType mt = (pa_mt == MemoryType::DEFAULT) ?
-                         Device::GetDeviceMemoryType() : pa_mt;
+   const MemoryType mt = (memory_type == MemoryType::DEFAULT) ?
+                         Device::GetDeviceMemoryType() : memory_type;
+
+   // If device options allow fast kernels, set the action type to L2L
+   action_type =
+      Device::FastKernelsEnabled() ? ActionType::L2L : ActionType::E2E;
 
    // Assuming the same element type
    fespace = &fes;
@@ -38,7 +73,7 @@ void MassIntegrator::AssemblePA(const FiniteElementSpace &fes)
    if (DeviceCanUseCeed())
    {
       delete ceedOp;
-      ceedOp = new ceed::PAMassIntegrator(fes, *ir, Q);
+      ceedOp = new ceed::PAMassIntegrator(fes, *ir, Q, action_type);
       return;
    }
    dim = mesh->Dimension();
@@ -154,328 +189,6 @@ void MassIntegrator::AssemblePA(const FiniteElementSpace &fes)
    }
 }
 
-template<int T_D1D = 0, int T_Q1D = 0>
-static void PAMassAssembleDiagonal2D(const int NE,
-                                     const Array<double> &b,
-                                     const Vector &d,
-                                     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 <= MAX_D1D, "");
-   MFEM_VERIFY(Q1D <= MAX_Q1D, "");
-   auto B = Reshape(b.Read(), Q1D, D1D);
-   auto D = Reshape(d.Read(), Q1D, Q1D, NE);
-   auto Y = Reshape(y.ReadWrite(), D1D, D1D, NE);
-   MFEM_FORALL(e, NE,
-   {
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
-      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
-      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-      double QD[MQ1][MD1];
-      for (int qx = 0; qx < Q1D; ++qx)
-      {
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            QD[qx][dy] = 0.0;
-            for (int qy = 0; qy < Q1D; ++qy)
-            {
-               QD[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)
-         {
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               Y(dx,dy,e) += B(qx, dx) * B(qx, dx) * QD[qx][dy];
-            }
-         }
-      }
-   });
-}
-
-template<int T_D1D = 0, int T_Q1D = 0, int T_NBZ = 0>
-static void SmemPAMassAssembleDiagonal2D(const int NE,
-                                         const Array<double> &b_,
-                                         const Vector &d_,
-                                         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;
-   constexpr int NBZ = T_NBZ ? T_NBZ : 1;
-   constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-   constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
-   MFEM_VERIFY(D1D <= MD1, "");
-   MFEM_VERIFY(Q1D <= MQ1, "");
-   auto b = Reshape(b_.Read(), Q1D, D1D);
-   auto D = Reshape(d_.Read(), Q1D, Q1D, NE);
-   auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
-   MFEM_FORALL_2D(e, NE, Q1D, Q1D, NBZ,
-   {
-      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 NBZ = T_NBZ ? T_NBZ : 1;
-      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
-      MFEM_SHARED double B[MQ1][MD1];
-      MFEM_SHARED double QDZ[NBZ][MQ1][MD1];
-      double (*QD)[MD1] = (double (*)[MD1])(QDZ + tidz);
-      if (tidz == 0)
-      {
-         MFEM_FOREACH_THREAD(d,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(q,x,Q1D)
-            {
-               B[q][d] = b(q,d);
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qx,x,Q1D)
-      {
-         MFEM_FOREACH_THREAD(dy,y,D1D)
-         {
-            QD[qx][dy] = 0.0;
-            for (int qy = 0; qy < Q1D; ++qy)
-            {
-               QD[qx][dy] += B[qy][dy] * B[qy][dy] * D(qx, qy, e);
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(dy,y,D1D)
-      {
-         MFEM_FOREACH_THREAD(dx,x,D1D)
-         {
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               // might need absolute values on next line
-               Y(dx,dy,e) += B[qx][dx] * B[qx][dx] * QD[qx][dy];
-            }
-         }
-      }
-   });
-}
-
-template<int T_D1D = 0, int T_Q1D = 0>
-static void PAMassAssembleDiagonal3D(const int NE,
-                                     const Array<double> &b,
-                                     const Vector &d,
-                                     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 <= MAX_D1D, "");
-   MFEM_VERIFY(Q1D <= MAX_Q1D, "");
-   auto B = Reshape(b.Read(), Q1D, D1D);
-   auto D = Reshape(d.Read(), Q1D, Q1D, Q1D, NE);
-   auto Y = Reshape(y.ReadWrite(), D1D, D1D, D1D, NE);
-   MFEM_FORALL(e, NE,
-   {
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
-      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
-      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-      double QQD[MQ1][MQ1][MD1];
-      double QDD[MQ1][MD1][MD1];
-      for (int qx = 0; qx < Q1D; ++qx)
-      {
-         for (int qy = 0; qy < Q1D; ++qy)
-         {
-            for (int dz = 0; dz < D1D; ++dz)
-            {
-               QQD[qx][qy][dz] = 0.0;
-               for (int qz = 0; qz < Q1D; ++qz)
-               {
-                  QQD[qx][qy][dz] += B(qz, dz) * B(qz, dz) * D(qx, qy, qz, e);
-               }
-            }
-         }
-      }
-      for (int qx = 0; qx < Q1D; ++qx)
-      {
-         for (int dz = 0; dz < D1D; ++dz)
-         {
-            for (int dy = 0; dy < D1D; ++dy)
-            {
-               QDD[qx][dy][dz] = 0.0;
-               for (int qy = 0; qy < Q1D; ++qy)
-               {
-                  QDD[qx][dy][dz] += B(qy, dy) * B(qy, dy) * QQD[qx][qy][dz];
-               }
-            }
-         }
-      }
-      for (int dz = 0; dz < D1D; ++dz)
-      {
-         for (int dy = 0; dy < D1D; ++dy)
-         {
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               double t = 0.0;
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  t += B(qx, dx) * B(qx, dx) * QDD[qx][dy][dz];
-               }
-               Y(dx, dy, dz, e) += t;
-            }
-         }
-      }
-   });
-}
-
-template<int T_D1D = 0, int T_Q1D = 0>
-static void SmemPAMassAssembleDiagonal3D(const int NE,
-                                         const Array<double> &b_,
-                                         const Vector &d_,
-                                         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;
-   constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-   constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
-   MFEM_VERIFY(D1D <= MD1, "");
-   MFEM_VERIFY(Q1D <= MQ1, "");
-   auto b = Reshape(b_.Read(), Q1D, D1D);
-   auto D = Reshape(d_.Read(), Q1D, Q1D, Q1D, NE);
-   auto Y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
-   MFEM_FORALL_3D(e, NE, Q1D, Q1D, Q1D,
-   {
-      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 : MAX_D1D;
-      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-      MFEM_SHARED double B[MQ1][MD1];
-      MFEM_SHARED double QQD[MQ1][MQ1][MD1];
-      MFEM_SHARED double QDD[MQ1][MD1][MD1];
-      if (tidz == 0)
-      {
-         MFEM_FOREACH_THREAD(d,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(q,x,Q1D)
-            {
-               B[q][d] = b(q,d);
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qx,x,Q1D)
-      {
-         MFEM_FOREACH_THREAD(qy,y,Q1D)
-         {
-            MFEM_FOREACH_THREAD(dz,z,D1D)
-            {
-               QQD[qx][qy][dz] = 0.0;
-               for (int qz = 0; qz < Q1D; ++qz)
-               {
-                  QQD[qx][qy][dz] += B[qz][dz] * B[qz][dz] * D(qx, qy, qz, e);
-               }
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qx,x,Q1D)
-      {
-         MFEM_FOREACH_THREAD(dz,z,D1D)
-         {
-            MFEM_FOREACH_THREAD(dy,y,D1D)
-            {
-               QDD[qx][dy][dz] = 0.0;
-               for (int qy = 0; qy < Q1D; ++qy)
-               {
-                  QDD[qx][dy][dz] += B[qy][dy] * B[qy][dy] * QQD[qx][qy][dz];
-               }
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(dz,z,D1D)
-      {
-         MFEM_FOREACH_THREAD(dy,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(dx,x,D1D)
-            {
-               double t = 0.0;
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  t += B[qx][dx] * B[qx][dx] * QDD[qx][dy][dz];
-               }
-               Y(dx, dy, dz, e) += t;
-            }
-         }
-      }
-   });
-}
-
-static void PAMassAssembleDiagonal(const int dim, const int D1D,
-                                   const int Q1D, const int NE,
-                                   const Array<double> &B,
-                                   const Vector &D,
-                                   Vector &Y)
-{
-   if (dim == 2)
-   {
-      switch ((D1D << 4 ) | Q1D)
-      {
-         case 0x22: return SmemPAMassAssembleDiagonal2D<2,2,16>(NE,B,D,Y);
-         case 0x33: return SmemPAMassAssembleDiagonal2D<3,3,16>(NE,B,D,Y);
-         case 0x44: return SmemPAMassAssembleDiagonal2D<4,4,8>(NE,B,D,Y);
-         case 0x55: return SmemPAMassAssembleDiagonal2D<5,5,8>(NE,B,D,Y);
-         case 0x66: return SmemPAMassAssembleDiagonal2D<6,6,4>(NE,B,D,Y);
-         case 0x77: return SmemPAMassAssembleDiagonal2D<7,7,4>(NE,B,D,Y);
-         case 0x88: return SmemPAMassAssembleDiagonal2D<8,8,2>(NE,B,D,Y);
-         case 0x99: return SmemPAMassAssembleDiagonal2D<9,9,2>(NE,B,D,Y);
-         default:   return PAMassAssembleDiagonal2D(NE,B,D,Y,D1D,Q1D);
-      }
-   }
-   else if (dim == 3)
-   {
-      switch ((D1D << 4 ) | Q1D)
-      {
-         case 0x23: return SmemPAMassAssembleDiagonal3D<2,3>(NE,B,D,Y);
-         case 0x24: return SmemPAMassAssembleDiagonal3D<2,4>(NE,B,D,Y);
-         case 0x26: return SmemPAMassAssembleDiagonal3D<2,6>(NE,B,D,Y);
-         case 0x34: return SmemPAMassAssembleDiagonal3D<3,4>(NE,B,D,Y);
-         case 0x35: return SmemPAMassAssembleDiagonal3D<3,5>(NE,B,D,Y);
-         case 0x45: return SmemPAMassAssembleDiagonal3D<4,5>(NE,B,D,Y);
-         case 0x48: return SmemPAMassAssembleDiagonal3D<4,8>(NE,B,D,Y);
-         case 0x56: return SmemPAMassAssembleDiagonal3D<5,6>(NE,B,D,Y);
-         case 0x67: return SmemPAMassAssembleDiagonal3D<6,7>(NE,B,D,Y);
-         case 0x78: return SmemPAMassAssembleDiagonal3D<7,8>(NE,B,D,Y);
-         case 0x89: return SmemPAMassAssembleDiagonal3D<8,9>(NE,B,D,Y);
-         default:   return PAMassAssembleDiagonal3D(NE,B,D,Y,D1D,Q1D);
-      }
-   }
-   MFEM_ABORT("Unknown kernel.");
-}
-
-void MassIntegrator::AssembleDiagonalPA(Vector &diag)
-{
-   if (DeviceCanUseCeed())
-   {
-      ceedOp->GetDiagonal(diag);
-   }
-   else
-   {
-      PAMassAssembleDiagonal(dim, dofs1D, quad1D, ne, maps->B, pa_data, diag);
-   }
-}
-
-
 #ifdef MFEM_USE_OCCA
 // OCCA PA Mass Apply 2D kernel
 static void OccaPAMassApply2D(const int D1D,
@@ -1151,6 +864,159 @@ static void SmemPAMassApply3D(const int NE,
    });
 }
 
+template<int D1D, int Q1D>
+void SmemPAMassApply3D_v1(const int NE,
+                          const Array<double> &b_,
+                          const Array<double> &bt_,
+                          const Vector &d_,
+                          const Vector &x_,
+                          Vector &y_,
+                          const int d1d = 0,
+                          const int q1d = 0)
+{
+   MFEM_CONTRACT_VAR(bt_);
+   MFEM_CONTRACT_VAR(d1d);
+   MFEM_CONTRACT_VAR(q1d);
+   const auto B = Reshape(b_.Read(), Q1D,D1D);
+   const auto D = Reshape(d_.Read(), Q1D,Q1D,Q1D, NE);
+   const auto X = Reshape(x_.Read(), D1D,D1D,D1D, NE);
+   auto Y = Reshape(y_.ReadWrite(), D1D,D1D,D1D, NE);
+
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, 1,
+   {
+      double u[Q1D];
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[Q1D][Q1D][Q1D];
+
+      // Load input, B & X interpolation
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            s_B[qx][dy] = B(qx,dy);
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double Bx = B(qx,dx);
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz)
+               {
+                  u[dz] += X(dx,dy,dz,e) * Bx;
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[dz][dy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Y interpolation
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { u[qy] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               const double zyX = s_q[dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int qy = 0; qy < Q1D; ++qy) { u[qy] += zyX * s_B[qy][dy]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { s_q[dz][qy][qx] = u[qy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Z interpolation, Q-function & Zt projection
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            // Z interpolation
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz) { u[qz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               const double zYX = s_q[dz][qy][qx];
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; ++qz) { u[qz] += zYX * s_B[qz][dz]; }
+            }
+
+            // Q-function
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               s_q[qz][qy][qx] = u[qz] * D(qx,qy,qz,e);
+            }
+
+            // Zt projection
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const double ZYX = s_q[qz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz) { u[dz] += ZYX * s_B[qz][dz]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[dz][qy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Yt projection
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { u[dy] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               const double zYX = s_q[dz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dy = 0; dy < D1D; ++dy) { u[dy] += zYX * s_B[qy][dy]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { s_q[dz][dy][qx] = u[dy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Xt projection & save output
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dy,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx) { u[dx] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               const double zyX = s_q[dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dx = 0; dx < D1D; ++dx) { u[dx] += zyX * s_B[qx][dx]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double output = u[dx];
+               Y(dx,dy,dz,e) += output;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
 static void PAMassApply(const int dim,
                         const int D1D,
                         const int Q1D,
@@ -1201,30 +1067,47 @@ static void PAMassApply(const int dim,
    }
    else if (dim == 3)
    {
+      const int ver = Device::KernelsVersion();
+      const int id = (ver << 8) | (D1D << 4) | Q1D;
+
+      static int ini = 0;
+      if (!ini++) { printf("\033[33mkernel #0x%x\033[m\n",id); }
+
       switch (id)
       {
-         case 0x22: return SmemPAMassApply3D<2,2>(NE,B,Bt,D,X,Y);
-         case 0x23: return SmemPAMassApply3D<2,3>(NE,B,Bt,D,X,Y);
-         case 0x24: return SmemPAMassApply3D<2,4>(NE,B,Bt,D,X,Y);
-         case 0x26: return SmemPAMassApply3D<2,6>(NE,B,Bt,D,X,Y);
-         case 0x34: return SmemPAMassApply3D<3,4>(NE,B,Bt,D,X,Y);
-         case 0x35: return SmemPAMassApply3D<3,5>(NE,B,Bt,D,X,Y);
-         case 0x36: return SmemPAMassApply3D<3,6>(NE,B,Bt,D,X,Y);
-         case 0x37: return SmemPAMassApply3D<3,7>(NE,B,Bt,D,X,Y);
-         case 0x45: return SmemPAMassApply3D<4,5>(NE,B,Bt,D,X,Y);
-         case 0x46: return SmemPAMassApply3D<4,6>(NE,B,Bt,D,X,Y);
-         case 0x48: return SmemPAMassApply3D<4,8>(NE,B,Bt,D,X,Y);
-         case 0x56: return SmemPAMassApply3D<5,6>(NE,B,Bt,D,X,Y);
-         case 0x58: return SmemPAMassApply3D<5,8>(NE,B,Bt,D,X,Y);
-         case 0x67: return SmemPAMassApply3D<6,7>(NE,B,Bt,D,X,Y);
-         case 0x78: return SmemPAMassApply3D<7,8>(NE,B,Bt,D,X,Y);
-         case 0x89: return SmemPAMassApply3D<8,9>(NE,B,Bt,D,X,Y);
-         case 0x9A: return SmemPAMassApply3D<9,10>(NE,B,Bt,D,X,Y);
-         default:   return PAMassApply3D(NE,B,Bt,D,X,Y,D1D,Q1D);
+         case 0x123: return SmemPAMassApply3D_v1<2,3>(NE,B,Bt,D,X,Y);
+         case 0x124: return SmemPAMassApply3D_v1<2,4>(NE,B,Bt,D,X,Y);
+         case 0x134: return SmemPAMassApply3D_v1<3,4>(NE,B,Bt,D,X,Y);
+         case 0x136: return SmemPAMassApply3D_v1<3,6>(NE,B,Bt,D,X,Y);
+         case 0x145: return SmemPAMassApply3D_v1<4,5>(NE,B,Bt,D,X,Y);
+         case 0x148: return SmemPAMassApply3D_v1<4,8>(NE,B,Bt,D,X,Y);
+         case 0x156: return SmemPAMassApply3D_v1<5,6>(NE,B,Bt,D,X,Y);
+         case 0x158: return SmemPAMassApply3D_v1<5,8>(NE,B,Bt,D,X,Y);
+         case 0x167: return SmemPAMassApply3D_v1<6,7>(NE,B,Bt,D,X,Y);
+         case 0x178: return SmemPAMassApply3D_v1<7,8>(NE,B,Bt,D,X,Y);
+
+         case 0x022: return SmemPAMassApply3D<2,2>(NE,B,Bt,D,X,Y);
+         case 0x023: return SmemPAMassApply3D<2,3>(NE,B,Bt,D,X,Y);
+         case 0x024: return SmemPAMassApply3D<2,4>(NE,B,Bt,D,X,Y);
+         case 0x026: return SmemPAMassApply3D<2,6>(NE,B,Bt,D,X,Y);
+         case 0x034: return SmemPAMassApply3D<3,4>(NE,B,Bt,D,X,Y);
+         case 0x035: return SmemPAMassApply3D<3,5>(NE,B,Bt,D,X,Y);
+         case 0x036: return SmemPAMassApply3D<3,6>(NE,B,Bt,D,X,Y);
+         case 0x037: return SmemPAMassApply3D<3,7>(NE,B,Bt,D,X,Y);
+         case 0x045: return SmemPAMassApply3D<4,5>(NE,B,Bt,D,X,Y);
+         case 0x046: return SmemPAMassApply3D<4,6>(NE,B,Bt,D,X,Y);
+         case 0x048: return SmemPAMassApply3D<4,8>(NE,B,Bt,D,X,Y);
+         case 0x056: return SmemPAMassApply3D<5,6>(NE,B,Bt,D,X,Y);
+         case 0x058: return SmemPAMassApply3D<5,8>(NE,B,Bt,D,X,Y);
+         case 0x067: return SmemPAMassApply3D<6,7>(NE,B,Bt,D,X,Y);
+         case 0x078: return SmemPAMassApply3D<7,8>(NE,B,Bt,D,X,Y);
+         case 0x089: return SmemPAMassApply3D<8,9>(NE,B,Bt,D,X,Y);
+         case 0x09A: return SmemPAMassApply3D<9,10>(NE,B,Bt,D,X,Y);
+
+         default:   break; //return PAMassApply3D(NE,B,Bt,D,X,Y,D1D,Q1D);
       }
    }
-   mfem::out << "Unknown kernel 0x" << std::hex << id << std::endl;
-   MFEM_ABORT("Unknown kernel.");
+   MFEM_ABORT("Unknown kernel 0x" << std::hex << id);
 }
 
 void MassIntegrator::AddMultPA(const Vector &x, Vector &y) const
@@ -1233,6 +1116,35 @@ void MassIntegrator::AddMultPA(const Vector &x, Vector &y) const
    {
       ceedOp->AddMult(x, y);
    }
+   else if (Device::FastKernelsEnabled())
+   {
+      const int version = Device::KernelsVersion();
+      MFEM_VERIFY(version < 4 || version==7, "Unsupported version!");
+      if (version == 3) // AMD
+      {
+         NDK_AMD_PAMassApply(dim, dofs1D, quad1D, ne,
+                             fespace, maps,
+                             pa_data, x, y);
+      }
+      // 4 E-vector
+      // 5 fused
+      // 6 MMA
+      else if (version == 7) // HIP
+      {
+         NDK_HIP_PAMassApply(dim, dofs1D, quad1D, ne,
+                             fespace, maps,
+                             pa_data, x, y);
+      }
+      // 0 legacy
+      // 1 fast
+      // 2 libP
+      else
+      {
+         NDK_PAMassApply(dim, dofs1D, quad1D, ne,
+                         fespace, maps,
+                         pa_data, x, y);
+      }
+   }
    else
    {
       PAMassApply(dim, dofs1D, quad1D, ne, maps->B, maps->Bt, pa_data, x, y);

fem/bilininteg_mass_pa_diag.cpp

@@ -0,0 +1,358 @@
+// Copyright (c) 2010-2021, 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 "../general/forall.hpp"
+#include "bilininteg.hpp"
+#include "gridfunc.hpp"
+#include "ceed/mass.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+void NDK_PAMassAssembleDiagonal(const int dim,
+                                const int D1D,
+                                const int Q1D,
+                                const int NE,
+                                const FiniteElementSpace *fes,
+                                const DofToQuad *maps,
+                                const Vector &D,
+                                Vector &Y);
+
+template<int T_D1D = 0, int T_Q1D = 0>
+static void PAMassAssembleDiagonal2D(const int NE,
+                                     const Array<double> &b,
+                                     const Vector &d,
+                                     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 <= MAX_D1D, "");
+   MFEM_VERIFY(Q1D <= MAX_Q1D, "");
+   auto B = Reshape(b.Read(), Q1D, D1D);
+   auto D = Reshape(d.Read(), Q1D, Q1D, NE);
+   auto Y = Reshape(y.ReadWrite(), D1D, D1D, NE);
+   MFEM_FORALL(e, NE,
+   {
+      const int D1D = T_D1D ? T_D1D : d1d;
+      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+      double QD[MQ1][MD1];
+      for (int qx = 0; qx < Q1D; ++qx)
+      {
+         for (int dy = 0; dy < D1D; ++dy)
+         {
+            QD[qx][dy] = 0.0;
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               QD[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)
+         {
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               Y(dx,dy,e) += B(qx, dx) * B(qx, dx) * QD[qx][dy];
+            }
+         }
+      }
+   });
+}
+
+template<int T_D1D = 0, int T_Q1D = 0, int T_NBZ = 0>
+static void SmemPAMassAssembleDiagonal2D(const int NE,
+                                         const Array<double> &b_,
+                                         const Vector &d_,
+                                         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;
+   constexpr int NBZ = T_NBZ ? T_NBZ : 1;
+   constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+   constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
+   MFEM_VERIFY(D1D <= MD1, "");
+   MFEM_VERIFY(Q1D <= MQ1, "");
+   auto b = Reshape(b_.Read(), Q1D, D1D);
+   auto D = Reshape(d_.Read(), Q1D, Q1D, NE);
+   auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
+   MFEM_FORALL_2D(e, NE, Q1D, Q1D, NBZ,
+   {
+      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 NBZ = T_NBZ ? T_NBZ : 1;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
+      MFEM_SHARED double B[MQ1][MD1];
+      MFEM_SHARED double QDZ[NBZ][MQ1][MD1];
+      double (*QD)[MD1] = (double (*)[MD1])(QDZ + tidz);
+      if (tidz == 0)
+      {
+         MFEM_FOREACH_THREAD(d,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(q,x,Q1D)
+            {
+               B[q][d] = b(q,d);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         MFEM_FOREACH_THREAD(dy,y,D1D)
+         {
+            QD[qx][dy] = 0.0;
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               QD[qx][dy] += B[qy][dy] * B[qy][dy] * D(qx, qy, e);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               // might need absolute values on next line
+               Y(dx,dy,e) += B[qx][dx] * B[qx][dx] * QD[qx][dy];
+            }
+         }
+      }
+   });
+}
+
+template<int T_D1D = 0, int T_Q1D = 0>
+static void PAMassAssembleDiagonal3D(const int NE,
+                                     const Array<double> &b,
+                                     const Vector &d,
+                                     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 <= MAX_D1D, "");
+   MFEM_VERIFY(Q1D <= MAX_Q1D, "");
+   auto B = Reshape(b.Read(), Q1D, D1D);
+   auto D = Reshape(d.Read(), Q1D, Q1D, Q1D, NE);
+   auto Y = Reshape(y.ReadWrite(), D1D, D1D, D1D, NE);
+   MFEM_FORALL(e, NE,
+   {
+      const int D1D = T_D1D ? T_D1D : d1d;
+      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+      double QQD[MQ1][MQ1][MD1];
+      double QDD[MQ1][MD1][MD1];
+      for (int qx = 0; qx < Q1D; ++qx)
+      {
+         for (int qy = 0; qy < Q1D; ++qy)
+         {
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               QQD[qx][qy][dz] = 0.0;
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  QQD[qx][qy][dz] += B(qz, dz) * B(qz, dz) * D(qx, qy, qz, e);
+               }
+            }
+         }
+      }
+      for (int qx = 0; qx < Q1D; ++qx)
+      {
+         for (int dz = 0; dz < D1D; ++dz)
+         {
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               QDD[qx][dy][dz] = 0.0;
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  QDD[qx][dy][dz] += B(qy, dy) * B(qy, dy) * QQD[qx][qy][dz];
+               }
+            }
+         }
+      }
+      for (int dz = 0; dz < D1D; ++dz)
+      {
+         for (int dy = 0; dy < D1D; ++dy)
+         {
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               double t = 0.0;
+               for (int qx = 0; qx < Q1D; ++qx)
+               {
+                  t += B(qx, dx) * B(qx, dx) * QDD[qx][dy][dz];
+               }
+               Y(dx, dy, dz, e) += t;
+            }
+         }
+      }
+   });
+}
+
+template<int T_D1D = 0, int T_Q1D = 0>
+static void SmemPAMassAssembleDiagonal3D(const int NE,
+                                         const Array<double> &b_,
+                                         const Vector &d_,
+                                         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;
+   constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+   constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
+   MFEM_VERIFY(D1D <= MD1, "");
+   MFEM_VERIFY(Q1D <= MQ1, "");
+   const auto b = Reshape(b_.Read(), Q1D, D1D);
+   const auto D = Reshape(d_.Read(), Q1D, Q1D, Q1D, NE);
+   auto Y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, Q1D,
+   {
+      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 : MAX_D1D;
+      constexpr int MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+      MFEM_SHARED double B[MQ1][MD1];
+      MFEM_SHARED double QQD[MQ1][MQ1][MD1];
+      MFEM_SHARED double QDD[MQ1][MD1][MD1];
+      if (tidz == 0)
+      {
+         MFEM_FOREACH_THREAD(d,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(q,x,Q1D)
+            {
+               B[q][d] = b(q,d);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qy,y,Q1D)
+         {
+            MFEM_FOREACH_THREAD(dz,z,D1D)
+            {
+               QQD[qx][qy][dz] = 0.0;
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  QQD[qx][qy][dz] += B[qz][dz] * B[qz][dz] * D(qx, qy, qz, e);
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         MFEM_FOREACH_THREAD(dz,z,D1D)
+         {
+            MFEM_FOREACH_THREAD(dy,y,D1D)
+            {
+               QDD[qx][dy][dz] = 0.0;
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  QDD[qx][dy][dz] += B[qy][dy] * B[qy][dy] * QQD[qx][qy][dz];
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(dz,z,D1D)
+      {
+         MFEM_FOREACH_THREAD(dy,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(dx,x,D1D)
+            {
+               double t = 0.0;
+               for (int qx = 0; qx < Q1D; ++qx)
+               {
+                  t += B[qx][dx] * B[qx][dx] * QDD[qx][dy][dz];
+               }
+               Y(dx, dy, dz, e) += t;
+            }
+         }
+      }
+   });
+}
+
+static void PAMassAssembleDiagonal(const int dim, const int D1D,
+                                   const int Q1D, const int NE,
+                                   const Array<double> &B,
+                                   const Vector &D,
+                                   Vector &Y)
+{
+   if (dim == 2)
+   {
+      switch ((D1D << 4 ) | Q1D)
+      {
+         case 0x22: return SmemPAMassAssembleDiagonal2D<2,2,16>(NE,B,D,Y);
+         case 0x33: return SmemPAMassAssembleDiagonal2D<3,3,16>(NE,B,D,Y);
+         case 0x44: return SmemPAMassAssembleDiagonal2D<4,4,8>(NE,B,D,Y);
+         case 0x55: return SmemPAMassAssembleDiagonal2D<5,5,8>(NE,B,D,Y);
+         case 0x66: return SmemPAMassAssembleDiagonal2D<6,6,4>(NE,B,D,Y);
+         case 0x77: return SmemPAMassAssembleDiagonal2D<7,7,4>(NE,B,D,Y);
+         case 0x88: return SmemPAMassAssembleDiagonal2D<8,8,2>(NE,B,D,Y);
+         case 0x99: return SmemPAMassAssembleDiagonal2D<9,9,2>(NE,B,D,Y);
+         default:   return PAMassAssembleDiagonal2D(NE,B,D,Y,D1D,Q1D);
+      }
+   }
+   else if (dim == 3)
+   {
+      switch ((D1D << 4 ) | Q1D)
+      {
+         case 0x23: return SmemPAMassAssembleDiagonal3D<2,3>(NE,B,D,Y);
+         case 0x24: return SmemPAMassAssembleDiagonal3D<2,4>(NE,B,D,Y);
+         case 0x26: return SmemPAMassAssembleDiagonal3D<2,6>(NE,B,D,Y);
+         case 0x34: return SmemPAMassAssembleDiagonal3D<3,4>(NE,B,D,Y);
+         case 0x35: return SmemPAMassAssembleDiagonal3D<3,5>(NE,B,D,Y);
+         case 0x45: return SmemPAMassAssembleDiagonal3D<4,5>(NE,B,D,Y);
+         case 0x48: return SmemPAMassAssembleDiagonal3D<4,8>(NE,B,D,Y);
+         case 0x56: return SmemPAMassAssembleDiagonal3D<5,6>(NE,B,D,Y);
+         case 0x67: return SmemPAMassAssembleDiagonal3D<6,7>(NE,B,D,Y);
+         case 0x78: return SmemPAMassAssembleDiagonal3D<7,8>(NE,B,D,Y);
+         case 0x89: return SmemPAMassAssembleDiagonal3D<8,9>(NE,B,D,Y);
+         default:   return PAMassAssembleDiagonal3D(NE,B,D,Y,D1D,Q1D);
+      }
+   }
+   MFEM_ABORT("Unknown kernel.");
+}
+
+void MassIntegrator::AssembleDiagonalPA(Vector &diag)
+{
+   if (DeviceCanUseCeed())
+   {
+      ceedOp->GetDiagonal(diag);
+   }
+   else if (Device::FastKernelsEnabled())
+   {
+      NDK_PAMassAssembleDiagonal(dim, dofs1D, quad1D, ne,
+                                 fespace, maps,
+                                 pa_data, diag);
+   }
+   else
+   {
+      PAMassAssembleDiagonal(dim, dofs1D, quad1D, ne, maps->B, pa_data, diag);
+   }
+}
+
+} // namespace mfem

fem/bilininteg_mass_pa_fast.cpp

@@ -0,0 +1,684 @@
+// Copyright (c) 2010-2021, 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 "../general/forall.hpp"
+#include "gridfunc.hpp"
+#include "restriction.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+// Fast '0' non-deterministic 3D mass kernel
+template<int D1D, int Q1D>
+void NDK_SmemPAMassApply3D(const int ndofs,
+                           const int NE,
+                           const int *map,
+                           const double *b_,
+                           const double *d_,
+                           const double *x_,
+                           double *y_)
+{
+   const auto MAP = Reshape(map, D1D,D1D,D1D, NE);
+   const auto b = Reshape(b_, Q1D, D1D);
+   const auto D = Reshape(d_, Q1D, Q1D, Q1D, NE);
+   const auto X = Reshape(x_, ndofs);
+   const auto X1 = Reshape(x_, D1D,D1D,D1D, NE);
+   auto Y = Reshape(y_, ndofs);
+   auto Y1 = Reshape(y_, D1D,D1D,D1D, NE);
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, 1,
+   {
+      MFEM_SHARED double sDQ[Q1D*Q1D];
+      double (*B)[D1D] = (double (*)[D1D]) sDQ;
+      double (*Bt)[Q1D] = (double (*)[Q1D]) sDQ;
+      MFEM_SHARED double sm0[Q1D*Q1D*Q1D];
+      MFEM_SHARED double sm1[Q1D*Q1D*Q1D];
+      double (*DDQ)[D1D][Q1D] = (double (*)[D1D][Q1D]) sm1;
+      double (*DQQ)[Q1D][Q1D] = (double (*)[Q1D][Q1D]) sm0;
+      double (*QQQ)[Q1D][Q1D] = (double (*)[Q1D][Q1D]) sm1;
+      double (*QQD)[Q1D][D1D] = (double (*)[Q1D][D1D]) sm0;
+      double (*QDD)[D1D][D1D] = (double (*)[D1D][D1D]) sm1;
+
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            B[qx][dy] = b(qx,dy);
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            double u[D1D];
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; dz++)
+            {
+               u[dz] = 0;
+            }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz)
+               {
+                  if (map)
+                  {
+                     const int gid = MAP(dx, dy, dz, e);
+                     const int idx = gid >= 0 ? gid : -1 - gid;
+                     u[dz] += X(idx) * B[qx][dx];
+                  }
+                  else
+                  {
+                     u[dz] += X1(dx,dy,dz,e) * B[qx][dx];
+                  }
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               DDQ[dz][dy][qx] = u[dz];
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            double u[D1D];
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; dz++)
+            {
+               u[dz] = 0;
+            }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; dz++)
+               {
+                  u[dz] += DDQ[dz][dy][qx] * B[qy][dy];
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; dz++)
+            {
+               DQQ[dz][qy][qx] = u[dz];
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            double u[Q1D];
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; qz++)
+            {
+               u[qz] = 0;
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; qz++)
+               {
+                  u[qz] += DQQ[dz][qy][qx] * B[qz][dz];
+               }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; qz++)
+            {
+               QQQ[qz][qy][qx] = u[qz] * D(qx,qy,qz,e);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(d,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(q,x,Q1D)
+         {
+            Bt[d][q] = b(q,d);
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u[Q1D];
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               u[qz] = 0;
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  u[qz] += QQQ[qz][qy][qx] * Bt[dx][qx];
+               }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               QQD[qz][qy][dx] = u[qz];
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u[Q1D];
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               u[qz] = 0;
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  u[qz] += QQD[qz][qy][dx] * Bt[dy][qy];
+               }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               QDD[qz][dy][dx] = u[qz];
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u[D1D];
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               u[dz] = 0;
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz)
+               {
+                  u[dz] += QDD[qz][dy][dx] * Bt[dz][qz];
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               if (map)
+               {
+                  const int gid = MAP(dx, dy, dz, e);
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  AtomicAdd(Y(idx), u[dz]);
+               }
+               else
+               {
+                  Y1(dx,dy,dz,e) += u[dz];
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
+// Fast '1' non-deterministic 3D mass kernel
+// Smem version melded toward registers + BZ-batch
+template<int D1D, int Q1D, int NBZ>
+void NDK_SmRgPAMassApply3D(const int ndofs,
+                           const int NE,
+                           const int *map,
+                           const double *b_,
+                           const double *d_,
+                           const double *x_,
+                           double *y_)
+{
+   const auto MAP = Reshape(map, D1D,D1D,D1D, NE);
+   const auto B = Reshape(b_, Q1D, D1D);
+   const auto D = Reshape(d_, Q1D, Q1D, Q1D, NE);
+   const auto X = Reshape(x_, ndofs);
+   const auto X1 = Reshape(x_, D1D,D1D,D1D, NE);
+   auto Y = Reshape(y_, ndofs);
+   auto Y1 = Reshape(y_, D1D,D1D,D1D, NE);
+
+   MFEM_FORALL_3D(be, (NE+NBZ-1)/NBZ, Q1D, Q1D, NBZ,
+   {
+      double u[Q1D];
+      const int tz = MFEM_THREAD_ID(z);
+      const int e = be * MFEM_THREAD_SIZE(z) + tz;
+
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[NBZ][Q1D][Q1D][Q1D];
+
+      // Load input, B & X interpolation
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            if (tz == 0) { s_B[qx][dy] = B(qx,dy); }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double Bx = B(qx,dx);
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz)
+               {
+                  const int gid = map ? MAP(dx,dy,dz,e) : 0;
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  u[dz] += (map ? X(idx) : X1(dx,dy,dz,e)) * Bx;
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[tz][dz][dy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Y interpolation
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { u[qy] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               const double zyX = s_q[tz][dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int qy = 0; qy < Q1D; ++qy) { u[qy] += zyX * s_B[qy][dy]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { s_q[tz][dz][qy][qx] = u[qy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Z interpolation, Q-function & Zt projection
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            // Z interpolation
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz) { u[qz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               const double zYX = s_q[tz][dz][qy][qx];
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; ++qz) { u[qz] += zYX * s_B[qz][dz]; }
+            }
+
+            // Q-function
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               s_q[tz][qz][qy][qx] = u[qz] * D(qx,qy,qz,e);
+            }
+
+            // Zt projection
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const double ZYX = s_q[tz][qz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz) { u[dz] += ZYX * s_B[qz][dz]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[tz][dz][qy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Yt projection
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { u[dy] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               const double zYX = s_q[tz][dz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dy = 0; dy < D1D; ++dy) { u[dy] += zYX * s_B[qy][dy]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { s_q[tz][dz][dy][qx] = u[dy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Xt projection & save output
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dy,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx) { u[dx] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               const double zyX = s_q[tz][dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dx = 0; dx < D1D; ++dx) { u[dx] += zyX * s_B[qx][dx]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double output = u[dx];
+               if (map)
+               {
+                  const int gid = MAP(dx,dy,dz,e);
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  AtomicAdd(Y(idx), output);
+               }
+               else
+               {
+                  Y1(dx,dy,dz,e) += output;
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
+// Fast '2' non-deterministic 3D mass kernel
+// Smem version with registers
+template<int D1D, int Q1D>
+void NDK_RegsPAMassApply3D(const int ndofs,
+                           const int NE,
+                           const int *map,
+                           const double *b_,
+                           const double *d_,
+                           const double *x_,
+                           double *y_)
+{
+   const auto MAP = Reshape(map, D1D,D1D,D1D, NE);
+   const auto B = Reshape(b_, Q1D,D1D);
+   const auto D = Reshape(d_, Q1D,Q1D,Q1D, NE);
+   const auto X = Reshape(x_, ndofs);
+   const auto X1 = Reshape(x_, D1D,D1D,D1D, NE);
+   auto Y = Reshape(y_, ndofs);
+   auto Y1 = Reshape(y_, D1D,D1D,D1D, NE);
+
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, 1,
+   {
+      double r_wk[Q1D];
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[Q1D][Q1D][Q1D];
+
+      // Load s_B, load X in shared memory
+      MFEM_FOREACH_THREAD(b,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,Q1D)
+         {
+            if (a<D1D) { s_B[b][a] = B(b,a); }
+
+            MFEM_UNROLL(Q1D)
+            for (int i=0; i<Q1D; ++i) { r_wk[i] = 0.0; }
+
+            if (a<D1D && b<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int c=0; c<D1D; ++c)
+               {
+                  const int gid = map ? MAP(a,b,c,e) : 0;
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  s_q[c][b][a] = map ? X(idx) : X1(a,b,c,e);
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Interpolate in X direction
+      MFEM_FOREACH_THREAD(c,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(b,x,Q1D)
+         {
+            if (b<D1D && c<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int a=0; a<D1D; ++a)
+               {
+                  const double q_cba = s_q[c][b][a];
+                  MFEM_UNROLL(Q1D)
+                  for (int i=0; i<Q1D; ++i) { r_wk[i] += s_B[i][a]*q_cba; }
+               }
+               // reg => s_mem
+               MFEM_UNROLL(Q1D)
+               for (int i=0; i<Q1D; ++i) { s_q[c][b][i] = r_wk[i]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int j=0; j<Q1D; ++j) { r_wk[j] = 0.0; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Interpolate in Y direction
+      MFEM_FOREACH_THREAD(c,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            if (c<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int b=0; b<D1D; ++b)
+               {
+                  const double q_cbi = s_q[c][b][i];
+                  MFEM_UNROLL(Q1D)
+                  for (int j=0; j<Q1D; ++j) { r_wk[j] += s_B[j][b]*q_cbi; }
+               }
+               MFEM_UNROLL(Q1D)
+               for (int j=0; j<Q1D; ++j) { s_q[c][j][i] = r_wk[j]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int k=0; k<Q1D; ++k) { r_wk[k] = 0.0; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(j,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            // Interpolate in Z direction
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               const double q_cji = s_q[c][j][i];
+               MFEM_UNROLL(Q1D)
+               for (int k=0; k<Q1D; ++k) { r_wk[k] += s_B[k][c]*q_cji; }
+            }
+
+            // Scale by Jacobian and integration weights
+            MFEM_UNROLL(Q1D)
+            for (int k=0; k<Q1D; ++k) { r_wk[k] *= D(i,j,k,e); }
+
+            // Project back in Z direction
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               double q_cji = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int k=0; k<Q1D; ++k) { q_cji += s_B[k][c] * r_wk[k]; }
+               s_q[c][j][i] = q_cji;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Project back in Y direction
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int j=0; j<Q1D; ++j) { r_wk[j] = s_q[c][j][i]; }
+            MFEM_UNROLL(D1D)
+            for (int b=0; b<D1D; ++b)
+            {
+               double q_cbi = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int j=0; j<Q1D; ++j) { q_cbi += s_B[j][b] * r_wk[j]; }
+               s_q[c][b][i] = q_cbi;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Project back in X direction
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(b,x,D1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int i=0; i<Q1D; ++i) { r_wk[i] = s_q[c][b][i]; }
+            MFEM_UNROLL(D1D)
+            for (int a=0; a<D1D; ++a)
+            {
+               double q_cba = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int i=0; i<Q1D; ++i) { q_cba += s_B[i][a] * r_wk[i]; }
+               s_q[c][b][a] = q_cba;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Save back to memory
+      MFEM_FOREACH_THREAD(b,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               const double q_cba = s_q[c][b][a];
+               const int gid = map ? MAP(a,b,c,e) : 0;
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               AtomicAdd(map?Y(idx):Y1(a,b,c,e), q_cba);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
+void NDK_PAMassApply(const int dim,
+                     const int D1D,
+                     const int Q1D,
+                     const int NE,
+                     const FiniteElementSpace *fes,
+                     const DofToQuad *maps,
+                     const Vector &D,
+                     const Vector &X,
+                     Vector &Y)
+{
+   const int ND = fes->GetNDofs();
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes->GetElementRestriction(ordering);
+   const ElementRestriction* ER = dynamic_cast<const ElementRestriction*>(ERop);
+   const int *map = ER ? ER->GatherMap().Read() : nullptr;
+   const double *b = maps->B.Read();
+   const double *d = D.Read();
+   const double *x = X.Read();
+   double *y = Y.ReadWrite();
+
+   assert(dim == 3);
+   const int ver = Device::KernelsVersion();
+   const int id = (ver << 8) | (D1D << 4) | Q1D;
+
+   //printf("\033[32mkernel #0x%x\033[m\n",id); fflush(0);
+
+   switch (id) // orders 1~6
+   {
+      // Fast '0': Legacy non-deterministic 3D mass kernel
+      case 0x023: return NDK_SmemPAMassApply3D<2,3>(ND,NE,map,b,d,x,y);
+      case 0x024: return NDK_SmemPAMassApply3D<2,4>(ND,NE,map,b,d,x,y);
+      case 0x034: return NDK_SmemPAMassApply3D<3,4>(ND,NE,map,b,d,x,y);
+      case 0x036: return NDK_SmemPAMassApply3D<3,6>(ND,NE,map,b,d,x,y);
+      case 0x045: return NDK_SmemPAMassApply3D<4,5>(ND,NE,map,b,d,x,y);
+      case 0x046: return NDK_SmemPAMassApply3D<4,6>(ND,NE,map,b,d,x,y);
+      case 0x048: return NDK_SmemPAMassApply3D<4,8>(ND,NE,map,b,d,x,y);
+      case 0x056: return NDK_SmemPAMassApply3D<5,6>(ND,NE,map,b,d,x,y);
+      case 0x058: return NDK_SmemPAMassApply3D<5,8>(ND,NE,map,b,d,x,y);
+      case 0x067: return NDK_SmemPAMassApply3D<6,7>(ND,NE,map,b,d,x,y);
+      case 0x078: return NDK_SmemPAMassApply3D<7,8>(ND,NE,map,b,d,x,y);
+
+      // Fast '1': Legacy & half smem non-deterministic 3D mass kernel + Z-batch
+      case 0x123: return NDK_SmRgPAMassApply3D<2,3,32>(ND,NE,map,b,d,x,y);//1
+      case 0x124: return NDK_SmRgPAMassApply3D<2,4,16>(ND,NE,map,b,d,x,y);
+      case 0x134: return NDK_SmRgPAMassApply3D<3,4,16>(ND,NE,map,b,d,x,y);//2
+      case 0x136: return NDK_SmRgPAMassApply3D<3,6,8>(ND,NE,map,b,d,x,y);
+      case 0x145: return NDK_SmRgPAMassApply3D<4,5,4>(ND,NE,map,b,d,x,y);//3
+      case 0x146: return NDK_SmRgPAMassApply3D<4,6,4>(ND,NE,map,b,d,x,y);
+      case 0x148: return NDK_SmRgPAMassApply3D<4,8,4>(ND,NE,map,b,d,x,y);
+      case 0x156: return NDK_SmRgPAMassApply3D<5,6,4>(ND,NE,map,b,d,x,y);//4
+      case 0x158: return NDK_SmRgPAMassApply3D<5,8,1>(ND,NE,map,b,d,x,y);
+      case 0x167: return NDK_SmRgPAMassApply3D<6,7,1>(ND,NE,map,b,d,x,y);//5
+      case 0x178: return NDK_SmRgPAMassApply3D<7,8,1>(ND,NE,map,b,d,x,y);//6
+
+      // Fast '2': libP non-deterministic 3D mass kernel
+      case 0x223: return NDK_RegsPAMassApply3D<2,3>(ND,NE,map,b,d,x,y);
+      case 0x224: return NDK_RegsPAMassApply3D<2,4>(ND,NE,map,b,d,x,y);
+      case 0x234: return NDK_RegsPAMassApply3D<3,4>(ND,NE,map,b,d,x,y);
+      case 0x236: return NDK_RegsPAMassApply3D<3,6>(ND,NE,map,b,d,x,y);
+      case 0x245: return NDK_RegsPAMassApply3D<4,5>(ND,NE,map,b,d,x,y);
+      case 0x246: return NDK_RegsPAMassApply3D<4,6>(ND,NE,map,b,d,x,y);
+      case 0x248: return NDK_RegsPAMassApply3D<4,8>(ND,NE,map,b,d,x,y);
+      case 0x256: return NDK_RegsPAMassApply3D<5,6>(ND,NE,map,b,d,x,y);
+      case 0x258: return NDK_RegsPAMassApply3D<5,8>(ND,NE,map,b,d,x,y);
+      case 0x267: return NDK_RegsPAMassApply3D<6,7>(ND,NE,map,b,d,x,y);
+      case 0x278: return NDK_RegsPAMassApply3D<7,8>(ND,NE,map,b,d,x,y);
+
+      // Fast '3': libP + AMD specific non-deterministic 3D mass kernel
+      // in bilininteg_mass_pa_fast_amd.cpp
+
+      default: break;
+   }
+
+   MFEM_ABORT("Unknown kernel 0x" << std::hex << id);
+}
+
+} // namespace mfem

fem/bilininteg_mass_pa_fast_amd.cpp

@@ -0,0 +1,250 @@
+// Copyright (c) 2010-2021, 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 "../general/forall.hpp"
+#include "gridfunc.hpp"
+#include "restriction.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+// Fast '3' non-deterministic 3D mass kernel
+template<int D1D, int Q1D>
+void NDK_AMD_PAMassApply3D(const int ndofs,
+                           const int NE,
+                           const int *map,
+                           const double *b_,
+                           const double *d_,
+                           const double *x_,
+                           double *y_)
+{
+   const auto MAP = Reshape(map, D1D,D1D,D1D, NE);
+   const auto B = Reshape(b_, Q1D,D1D);
+   const auto D = Reshape(d_, Q1D,Q1D,Q1D, NE);
+   const auto X = Reshape(x_, ndofs);
+   const auto X1 = Reshape(x_, D1D,D1D,D1D, NE);
+   auto Y = Reshape(y_, ndofs);
+   auto Y1 = Reshape(y_, D1D,D1D,D1D, NE);
+
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, 1,
+   {
+      double r_wk[Q1D];
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[Q1D][Q1D][Q1D];
+
+      // Load s_B, load X in shared memory
+      MFEM_FOREACH_THREAD(b,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,Q1D)
+         {
+            if (a<D1D) { s_B[b][a] = B(b,a); }
+
+            MFEM_UNROLL(Q1D)
+            for (int i=0; i<Q1D; ++i) { r_wk[i] = 0.0; }
+
+            if (a<D1D && b<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int c=0; c<D1D; ++c)
+               {
+                  const int gid = map ? MAP(a,b,c,e) : 0;
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  s_q[c][b][a] = map ? X(idx) : X1(a,b,c,e);
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Interpolate in X direction
+      MFEM_FOREACH_THREAD(c,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(b,x,Q1D)
+         {
+            if (b<D1D && c<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int a=0; a<D1D; ++a)
+               {
+                  const double q_cba = s_q[c][b][a];
+                  MFEM_UNROLL(Q1D)
+                  for (int i=0; i<Q1D; ++i) { r_wk[i] += s_B[i][a]*q_cba; }
+               }
+               // reg => s_mem
+               MFEM_UNROLL(Q1D)
+               for (int i=0; i<Q1D; ++i) { s_q[c][b][i] = r_wk[i]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int j=0; j<Q1D; ++j) { r_wk[j] = 0.0; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Interpolate in Y direction
+      MFEM_FOREACH_THREAD(c,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            if (c<D1D)
+            {
+               MFEM_UNROLL(D1D)
+               for (int b=0; b<D1D; ++b)
+               {
+                  const double q_cbi = s_q[c][b][i];
+                  MFEM_UNROLL(Q1D)
+                  for (int j=0; j<Q1D; ++j) { r_wk[j] += s_B[j][b]*q_cbi; }
+               }
+               MFEM_UNROLL(Q1D)
+               for (int j=0; j<Q1D; ++j) { s_q[c][j][i] = r_wk[j]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int k=0; k<Q1D; ++k) { r_wk[k] = 0.0; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(j,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            // Interpolate in Z direction
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               const double q_cji = s_q[c][j][i];
+               MFEM_UNROLL(Q1D)
+               for (int k=0; k<Q1D; ++k) { r_wk[k] += s_B[k][c]*q_cji; }
+            }
+
+            // Scale by Jacobian and integration weights
+            MFEM_UNROLL(Q1D)
+            for (int k=0; k<Q1D; ++k) { r_wk[k] *= D(i,j,k,e); }
+
+            // Project back in Z direction
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               double q_cji = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int k=0; k<Q1D; ++k) { q_cji += s_B[k][c] * r_wk[k]; }
+               s_q[c][j][i] = q_cji;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Project back in Y direction
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int j=0; j<Q1D; ++j) { r_wk[j] = s_q[c][j][i]; }
+            MFEM_UNROLL(D1D)
+            for (int b=0; b<D1D; ++b)
+            {
+               double q_cbi = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int j=0; j<Q1D; ++j) { q_cbi += s_B[j][b] * r_wk[j]; }
+               s_q[c][b][i] = q_cbi;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Project back in X direction
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(b,x,D1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int i=0; i<Q1D; ++i) { r_wk[i] = s_q[c][b][i]; }
+            MFEM_UNROLL(D1D)
+            for (int a=0; a<D1D; ++a)
+            {
+               double q_cba = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int i=0; i<Q1D; ++i) { q_cba += s_B[i][a] * r_wk[i]; }
+               s_q[c][b][a] = q_cba;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Save back to memory
+      MFEM_FOREACH_THREAD(b,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               const double q_cba = s_q[c][b][a];
+               const int gid = map ? MAP(a,b,c,e) : 0;
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               AtomicAdd(map?Y(idx):Y1(a,b,c,e), q_cba);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
+void NDK_AMD_PAMassApply(const int dim,
+                         const int D1D,
+                         const int Q1D,
+                         const int NE,
+                         const FiniteElementSpace *fes,
+                         const DofToQuad *maps,
+                         const Vector &D,
+                         const Vector &X,
+                         Vector &Y)
+{
+   const int ND = fes->GetNDofs();
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes->GetElementRestriction(ordering);
+   const ElementRestriction* ER = dynamic_cast<const ElementRestriction*>(ERop);
+   const int *map = ER ? ER->GatherMap().Read() : nullptr;
+   const double *b = maps->B.Read();
+   const double *d = D.Read();
+   const double *x = X.Read();
+   double *y = Y.ReadWrite();
+
+   assert(dim == 3);
+   const int ver = Device::KernelsVersion();
+   const int id = (ver << 8) | (D1D << 4) | Q1D;
+
+   //printf("\033[32mkernel #0x%x\033[m\n",id); fflush(0);
+
+   switch (id) // orders 1~6
+   {
+      // Fast '3': libP + AMD specific non-deterministic 3D mass kernel
+      case 0x323: return NDK_AMD_PAMassApply3D<2,3>(ND,NE,map,b,d,x,y);
+      case 0x324: return NDK_AMD_PAMassApply3D<2,4>(ND,NE,map,b,d,x,y);
+      case 0x334: return NDK_AMD_PAMassApply3D<3,4>(ND,NE,map,b,d,x,y);
+      case 0x336: return NDK_AMD_PAMassApply3D<3,6>(ND,NE,map,b,d,x,y);
+      case 0x345: return NDK_AMD_PAMassApply3D<4,5>(ND,NE,map,b,d,x,y);
+      case 0x346: return NDK_AMD_PAMassApply3D<4,6>(ND,NE,map,b,d,x,y);
+      case 0x348: return NDK_AMD_PAMassApply3D<4,8>(ND,NE,map,b,d,x,y);
+      case 0x356: return NDK_AMD_PAMassApply3D<5,6>(ND,NE,map,b,d,x,y);
+      case 0x358: return NDK_AMD_PAMassApply3D<5,8>(ND,NE,map,b,d,x,y);
+      case 0x367: return NDK_AMD_PAMassApply3D<6,7>(ND,NE,map,b,d,x,y);
+      case 0x378: return NDK_AMD_PAMassApply3D<7,8>(ND,NE,map,b,d,x,y);
+
+      default: break;
+   }
+
+   MFEM_ABORT("Unknown kernel 0x" << std::hex << id);
+}
+
+} // namespace mfem

fem/bilininteg_mass_pa_fast_diag.cpp

@@ -0,0 +1,269 @@
+// Copyright (c) 2010-2021, 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 "../general/forall.hpp"
+#include "gridfunc.hpp"
+#include "restriction.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+// Fast '0' non-deterministic 3D mass diag kernel
+template<int D1D, int Q1D>
+static void NDK_SmemPAMassDiag3D(const int ndofs,
+                                 const int NE,
+                                 const int *map_,
+                                 const double *b_,
+                                 const double *d_,
+                                 double *y_)
+{
+   const auto MAP = Reshape(map_, D1D,D1D,D1D, NE);
+   const auto b = Reshape(b_, Q1D, D1D);
+   const auto D = Reshape(d_, Q1D, Q1D, Q1D, NE);
+   auto Y = Reshape(y_, ndofs);
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, Q1D,
+   {
+      const int tidz = MFEM_THREAD_ID(z);
+      MFEM_SHARED double B[Q1D][D1D];
+      MFEM_SHARED double QQD[Q1D][Q1D][D1D];
+      MFEM_SHARED double QDD[Q1D][D1D][D1D];
+      if (tidz == 0)
+      {
+         MFEM_FOREACH_THREAD(d,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(q,x,Q1D)
+            {
+               B[q][d] = b(q,d);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qy,y,Q1D)
+         {
+            MFEM_FOREACH_THREAD(dz,z,D1D)
+            {
+               QQD[qx][qy][dz] = 0.0;
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  QQD[qx][qy][dz] += B[qz][dz] * B[qz][dz] * D(qx, qy, qz, e);
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         MFEM_FOREACH_THREAD(dz,z,D1D)
+         {
+            MFEM_FOREACH_THREAD(dy,y,D1D)
+            {
+               QDD[qx][dy][dz] = 0.0;
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  QDD[qx][dy][dz] += B[qy][dy] * B[qy][dy] * QQD[qx][qy][dz];
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+      MFEM_FOREACH_THREAD(dz,z,D1D)
+      {
+         MFEM_FOREACH_THREAD(dy,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(dx,x,D1D)
+            {
+               double t = 0.0;
+               for (int qx = 0; qx < Q1D; ++qx)
+               {
+                  t += B[qx][dx] * B[qx][dx] * QDD[qx][dy][dz];
+               }
+               const int gid = MAP(dx, dy, dz, e);
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               AtomicAdd(Y(idx), t);
+            }
+         }
+      }
+   });
+}
+
+// Fast non-deterministic 3D mass diag kernel
+// Less smem version with registers
+template<int D1D, int Q1D>
+static void NDK_RegsPAMassDiag3D(const int ndofs,
+                                 const int NE,
+                                 const int *map_,
+                                 const double *b_,
+                                 const double *d_,
+                                 double *y_)
+{
+   const auto MAP = Reshape(map_, D1D,D1D,D1D, NE);
+   const auto B = Reshape(b_, Q1D, D1D);
+   const auto D = Reshape(d_, Q1D, Q1D, Q1D, NE);
+   auto Y = Reshape(y_, ndofs);
+
+   MFEM_FORALL_3D(e, NE, Q1D, Q1D, 1,
+   {
+      double r_wk[Q1D];
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[Q1D][Q1D][Q1D];
+
+
+      MFEM_FOREACH_THREAD(d,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(q,x,Q1D)
+         {
+            if (d<D1D) { s_B[q][d] = B(q,d); }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(j,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int k=0; k<Q1D; ++k) { r_wk[k] = D(i,j,k,e); }
+            for (int c=0; c<D1D; ++c)
+            {
+               double q_cji = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int k=0; k<Q1D; ++k)
+               {
+                  const double Bkc = s_B[k][c];
+                  q_cji += Bkc * Bkc * r_wk[k];
+               }
+               s_q[c][j][i] = q_cji;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(i,x,Q1D)
+         {
+            for (int j=0; j<Q1D; ++j) { r_wk[j] = s_q[c][j][i]; }
+            MFEM_UNROLL(D1D)
+            for (int b=0; b<D1D; ++b)
+            {
+               double q_cbi = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int j=0; j<Q1D; ++j)
+               {
+                  const double Bjb = s_B[j][b];
+                  q_cbi += Bjb * Bjb * r_wk[j];
+               }
+               s_q[c][b][i] = q_cbi;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(c,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(b,x,D1D)
+         {
+            for (int i=0; i<Q1D; ++i) { r_wk[i] = s_q[c][b][i]; }
+            MFEM_UNROLL(D1D)
+            for (int a=0; a<D1D; ++a)
+            {
+               double q_cba = 0.0;
+               MFEM_UNROLL(Q1D)
+               for (int i=0; i<Q1D; ++i)
+               {
+                  const double Bia = s_B[i][a];
+                  q_cba += Bia * Bia * r_wk[i];
+               }
+               s_q[c][b][a] = q_cba;
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      MFEM_FOREACH_THREAD(b,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(a,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int c=0; c<D1D; ++c)
+            {
+               const double q_cba = s_q[c][b][a];
+               const int gid = MAP(a,b,c,e);
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               AtomicAdd(Y(idx), q_cba);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
+void NDK_PAMassAssembleDiagonal(const int dim,
+                                const int D1D,
+                                const int Q1D,
+                                const int NE,
+                                const FiniteElementSpace *fes,
+                                const DofToQuad *maps,
+                                const Vector &D,
+                                Vector &Y)
+{
+   const int ND = fes->GetNDofs();
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes->GetElementRestriction(ordering);
+   const ElementRestriction* ER = dynamic_cast<const ElementRestriction*>(ERop);
+   assert(ER);
+   const int *map = ER->GatherMap().Read();
+   const double *b = maps->B.Read();
+   const double *d = D.Read();
+   double *y = Y.ReadWrite();
+
+   assert(dim == 3);
+   const int ver = Device::KernelsVersion();
+   const int vid = (ver << 8) | (D1D << 4) | Q1D;
+
+   // Fast '0' non-deterministic 3D mass diag kernel
+   switch (vid) // orders 1~6
+   {
+      case 0x023: return NDK_SmemPAMassDiag3D<2,3>(ND,NE,map,b,d,y);
+      case 0x024: return NDK_SmemPAMassDiag3D<2,4>(ND,NE,map,b,d,y);
+      case 0x034: return NDK_SmemPAMassDiag3D<3,4>(ND,NE,map,b,d,y);
+      case 0x045: return NDK_SmemPAMassDiag3D<4,5>(ND,NE,map,b,d,y);
+      case 0x046: return NDK_SmemPAMassDiag3D<4,6>(ND,NE,map,b,d,y);
+      case 0x056: return NDK_SmemPAMassDiag3D<5,6>(ND,NE,map,b,d,y);
+      case 0x058: return NDK_SmemPAMassDiag3D<5,8>(ND,NE,map,b,d,y);
+      case 0x067: return NDK_SmemPAMassDiag3D<6,7>(ND,NE,map,b,d,y);
+      case 0x078: return NDK_SmemPAMassDiag3D<7,8>(ND,NE,map,b,d,y);
+      default: break;
+   }
+
+   // all other versions should fold back through these kernels
+   switch ((D1D << 4) | Q1D)
+   {
+      case 0x23: return NDK_RegsPAMassDiag3D<2,3>(ND,NE,map,b,d,y);
+      case 0x24: return NDK_RegsPAMassDiag3D<2,4>(ND,NE,map,b,d,y);
+      case 0x34: return NDK_RegsPAMassDiag3D<3,4>(ND,NE,map,b,d,y);
+      case 0x45: return NDK_RegsPAMassDiag3D<4,5>(ND,NE,map,b,d,y);
+      case 0x46: return NDK_RegsPAMassDiag3D<4,6>(ND,NE,map,b,d,y);
+      case 0x56: return NDK_RegsPAMassDiag3D<5,6>(ND,NE,map,b,d,y);
+      case 0x58: return NDK_RegsPAMassDiag3D<5,8>(ND,NE,map,b,d,y);
+      case 0x67: return NDK_RegsPAMassDiag3D<6,7>(ND,NE,map,b,d,y);
+      case 0x78: return NDK_RegsPAMassDiag3D<7,8>(ND,NE,map,b,d,y);
+      default: break;
+   }
+
+   MFEM_ABORT("Unknown kernel 0x" << std::hex << vid);
+}
+
+} // namespace mfem

fem/bilininteg_mass_pa_fast_hip.cpp

@@ -0,0 +1,224 @@
+// Copyright (c) 2010-2021, 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 "../general/forall.hpp"
+#include "gridfunc.hpp"
+#include "restriction.hpp"
+
+using namespace std;
+
+namespace mfem
+{
+
+// Fast '7' HIP
+template<int D1D, int Q1D, int NBZ=1, int NBK=1> MFEM_GLOBAL static
+//MFEM_LAUNCH_BOUNDS(Q1D*Q1D*NBZ,NBK)
+void HIP_PAMassApply(const int NE,
+                     const int* MAP,
+                     const double* B,
+                     const double* D,
+                     const double* X,
+                     double* Y)
+{
+   double u[Q1D];
+   const int tz = MFEM_THREAD_ID(z);
+   MFEM_SHARED double s_B[Q1D][D1D];
+   MFEM_SHARED double s_q[NBZ][Q1D][Q1D][Q1D];
+
+   for (int be = MFEM_BLOCK_ID(x); be < (NE+NBZ-1)/NBZ; be += MFEM_GRID_DIM(x))
+   {
+      const int e = be * NBZ + tz;
+      if (e>=NE) { return; }
+
+      // Load input, B & X interpolation
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            if (tz == 0) { s_B[qx][dy] = B[qx+Q1D*dy]; }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double Bx = B[qx+Q1D*dx];
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz)
+               {
+                  const int gid = MAP[e*D1D*D1D*D1D+dz*D1D*D1D+dy*D1D+dx];
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  u[dz] = fma(X[idx], Bx, u[dz]);
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[tz][dz][dy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Y interpolation
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { u[qy] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               const double zyX = s_q[tz][dz][dy][qx];
+               MFEM_UNROLL(Q1D)
+               for (int qy = 0; qy < Q1D; ++qy) { u[qy] = fma(zyX,s_B[qy][dy],u[qy]); }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { s_q[tz][dz][qy][qx] = u[qy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Z interpolation, Q-function & Zt projection
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            // Z interpolation
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz) { u[qz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               const double zYX = s_q[tz][dz][qy][qx];
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; ++qz) { u[qz] = fma(zYX,s_B[qz][dz],u[qz]); }
+            }
+
+            // Q-function
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const int idx = e*Q1D*Q1D*Q1D + qx + qy*Q1D + qz*Q1D*Q1D;
+               s_q[tz][qz][qy][qx] = u[qz] * D[idx];
+            }
+
+            // Zt projection
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const double ZYX = s_q[tz][qz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz) { u[dz] = fma(ZYX,s_B[qz][dz],u[dz]); }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[tz][dz][qy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Yt projection
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { u[dy] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               const double zYX = s_q[tz][dz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dy = 0; dy < D1D; ++dy) { u[dy] = fma(zYX,s_B[qy][dy],u[dy]); }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { s_q[tz][dz][dy][qx] = u[dy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Xt projection & save output
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dy,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx) { u[dx] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               const double zyX = s_q[tz][dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dx = 0; dx < D1D; ++dx) { u[dx] = fma(zyX,s_B[qx][dx],u[dx]); }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double output = u[dx];
+               const int gid = MAP[e*D1D*D1D*D1D+dz*D1D*D1D+dy*D1D+dx];
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               AtomicAdd(Y[idx], output);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   }
+}
+
+
+void NDK_HIP_PAMassApply(const int dim,
+                         const int D1D,
+                         const int Q1D,
+                         const int NE,
+                         const FiniteElementSpace *fes,
+                         const DofToQuad *maps,
+                         const Vector &d,
+                         const Vector &x,
+                         Vector &y)
+{
+   assert(dim == 3);
+
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes->GetElementRestriction(ordering);
+   const ElementRestriction *ER = dynamic_cast<const ElementRestriction*>(ERop);
+   assert(ER);
+
+   const auto M = ER->GatherMap().Read();
+   const double *B = maps->B.Read();
+   const double *D = d.Read();
+   const double *X = x.Read();
+   double *Y = y.ReadWrite();
+
+   void (*Ker)(const int NE,
+               const int *M,
+               const double *B,
+               const double *D,
+               const double *X,
+               double *Y) = nullptr;
+
+   const int ver = Device::KernelsVersion();
+   const int id = (ver << 8) | (D1D << 4) | Q1D;
+
+   switch (id) // orders 1~8
+   {
+      case 0x723: Ker=HIP_PAMassApply<2,3>; break; // 1
+      case 0x734: Ker=HIP_PAMassApply<3,4>; break; // 2
+      case 0x745: Ker=HIP_PAMassApply<4,5>; break; // 3
+      case 0x756: Ker=HIP_PAMassApply<5,6>; break; // 4
+      case 0x767: Ker=HIP_PAMassApply<6,7>; break; // 5
+      case 0x778: Ker=HIP_PAMassApply<7,8>; break; // 6
+      case 0x789: Ker=HIP_PAMassApply<8,9>; break; // 7
+      case 0x79A: Ker=HIP_PAMassApply<9,10>; break; // 8
+      default: MFEM_ABORT("Unknown kernel 0x" << std::hex << id << std::dec);
+   }
+   MFEM_LAUNCH_KERNEL(Ker,NE,dim3(Q1D,Q1D,1),0,NE,M,B,D,X,Y);
+}
+
+} // namespace mfem

fem/bilininteg_vecmass.cpp

@@ -34,7 +34,7 @@ void VectorMassIntegrator::AssemblePA(const FiniteElementSpace &fes)
    if (DeviceCanUseCeed())
    {
       delete ceedOp;
-      ceedOp = new ceed::PAMassIntegrator(fes, *ir, Q);
+      ceedOp = new ceed::PAMassIntegrator(fes, *ir, Q, action_type);
       return;
    }
    dim = mesh->Dimension();

fem/ceed/mass.cpp

@@ -48,7 +48,8 @@ struct MassOperatorInfo : public OperatorInfo
 
 PAMassIntegrator::PAMassIntegrator(const mfem::FiniteElementSpace &fes,
                                    const mfem::IntegrationRule &irm,
-                                   mfem::Coefficient *Q)
+                                   mfem::Coefficient *Q,
+                                   mfem::ActionType &action_type)
    : PAIntegrator()
 {
 #ifdef MFEM_USE_CEED
@@ -57,6 +58,7 @@ PAMassIntegrator::PAMassIntegrator(const mfem::FiniteElementSpace &fes,
 #else
    MFEM_ABORT("MFEM must be built with MFEM_USE_CEED=YES to use libCEED.");
 #endif
+   action_type = mfem::ActionType::L2L;
 }
 
 MFMassIntegrator::MFMassIntegrator(const mfem::FiniteElementSpace &fes,

fem/ceed/mass.hpp

@@ -27,7 +27,8 @@ class PAMassIntegrator : public PAIntegrator
 public:
    PAMassIntegrator(const mfem::FiniteElementSpace &fes,
                     const mfem::IntegrationRule &irm,
-                    mfem::Coefficient *Q);
+                    mfem::Coefficient *Q,
+                    mfem::ActionType &action_type);
 };
 
 /// Represent a MassIntegrator with AssemblyLevel::None using libCEED.

fem/ceed/util.cpp

@@ -65,6 +65,8 @@ void RemoveBasisAndRestriction(const mfem::FiniteElementSpace *fes)
          itr++;
       }
    }
+#else
+   MFEM_CONTRACT_VAR(fes);
 #endif
 }
 

fem/fe/fe_base.hpp

@@ -749,11 +749,12 @@ public:
                            DenseMatrix &div) const;
 
    /** @brief Get an Array<int> that maps lexicographically ordered indices to
-       the indices of the respective nodes/dofs/basis functions. Lexicographic
-       ordering of nodes is defined in terms of reference-space coordinates
-       (x,y,z). Lexicographically ordered nodes are listed first in order of
-       increasing x-coordinate, and then in order of increasing y-coordinate,
-       and finally in order of increasing z-coordinate.
+       the indices of the respective nodes/dofs/basis functions.
+
+       Lexicographic ordering of nodes is defined in terms of reference-space
+       coordinates (x,y,z). Lexicographically ordered nodes are listed first in
+       order of increasing x-coordinate, and then in order of increasing
+       y-coordinate, and finally in order of increasing z-coordinate.
 
        For example, the six nodes of a quadratic triangle are lexicographically
        ordered as follows:

fem/kernels.hpp

@@ -48,6 +48,26 @@ MFEM_HOST_DEVICE inline void LoadB(const int D1D, const int Q1D,
    MFEM_SYNC_THREAD;
 }
 
+/// Load B1d matrice into shared memory
+MFEM_HOST_DEVICE inline void LoadB(const int D1D, const int Q1D,
+                                   const ConstDeviceMatrix &b,
+                                   const DeviceMatrix &B)
+{
+   const int tidz = MFEM_THREAD_ID(z);
+
+   if (tidz == 0)
+   {
+      MFEM_FOREACH_THREAD(d,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(q,x,Q1D)
+         {
+            B(q,d) = b(q,d);
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
 /// Load Bt1d matrices into shared memory
 template<int MD1, int MQ1>
 MFEM_HOST_DEVICE inline void LoadBt(const int D1D, const int Q1D,
@@ -120,6 +140,29 @@ MFEM_HOST_DEVICE inline void LoadBGt(const int D1D, const int Q1D,
    MFEM_SYNC_THREAD;
 }
 
+/// Load Bt1d & Gt1d matrices into shared memory w/o MAX
+MFEM_HOST_DEVICE inline void LoadBGt(const int D1D, const int Q1D,
+                                     const ConstDeviceMatrix &b,
+                                     const ConstDeviceMatrix &g,
+                                     const DeviceMatrix &Bt,
+                                     const DeviceMatrix &Gt)
+{
+   const int tidz = MFEM_THREAD_ID(z);
+
+   if (tidz == 0)
+   {
+      MFEM_FOREACH_THREAD(d,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(q,x,Q1D)
+         {
+            Bt(d,q) = b(q,d);
+            Gt(d,q) = g(q,d);
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
 /// Load 2D input scalar into shared memory
 template<int MD1, int NBZ>
 MFEM_HOST_DEVICE inline void LoadX(const int e, const int D1D,
@@ -628,6 +671,141 @@ MFEM_HOST_DEVICE inline void GradXt(const int D1D, const int Q1D,
    MFEM_SYNC_THREAD;
 }
 
+/// Atomic 2D Transposed Evaluation, 1/2
+MFEM_HOST_DEVICE inline void Atomic2DEvalYt(const int D1D, const int Q1D,
+                                            const DeviceMatrix &B,
+                                            const DeviceMatrix &QQ,
+                                            const DeviceMatrix &QD)
+{
+   MFEM_FOREACH_THREAD(qy,y,Q1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         double u = 0.0;
+         for (int qx = 0; qx < Q1D; ++qx) { u += QQ(qy,qx) * B(qx,dx); }
+         QD(qy,dx) = u;
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+/// Atomic 2D Transposed Evaluation, 2/2
+MFEM_HOST_DEVICE inline void Atomic2DEvalXt(const int D1D, const int Q1D,
+                                            const DeviceMatrix &B,
+                                            const DeviceMatrix &Q,
+                                            const DeviceTensor<3,const int> &I,
+                                            const DeviceMatrix &Y,
+                                            const int c,
+                                            const int e,
+                                            const bool byVDIM)
+{
+   MFEM_FOREACH_THREAD(dy,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         double u = 0.0;
+         for (int qy = 0; qy < Q1D; ++qy) { u += Q(qy,dx) * B(qy,dy); }
+         const int gid = I(dx,dy,e);
+         const int idx = gid >= 0 ? gid : -1 - gid;
+         if (byVDIM) { AtomicAdd(Y(c,idx), u); }
+         else { AtomicAdd(Y(idx,c), u); }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+MFEM_HOST_DEVICE inline
+void Atomic2DEvalTranspose(const int D1D,
+                           const int Q1D,
+                           const DeviceMatrix &B,
+                           const DeviceMatrix &QQ,
+                           const DeviceMatrix &QD,
+                           const DeviceTensor<3,const int> &I,
+                           const DeviceMatrix &Y,
+                           const int c,
+                           const int e,
+                           const bool byVDIM)
+{
+   kernels::internal::Atomic2DEvalYt(D1D,Q1D,B,QQ,QD);
+   kernels::internal::Atomic2DEvalXt(D1D,Q1D,B,QD,I,Y,c,e,byVDIM);
+}
+
+/// Atomic 2D Transposed Gradient, 1/2
+MFEM_HOST_DEVICE inline void AtomicGradYt(const int D1D, const int Q1D,
+                                          const DeviceMatrix &Bt,
+                                          const DeviceMatrix &Gt,
+                                          const DeviceMatrix &QQ0,
+                                          const DeviceMatrix &QQ1,
+                                          const DeviceMatrix &DQ0,
+                                          const DeviceMatrix &DQ1)
+{
+   MFEM_FOREACH_THREAD(qy,y,Q1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         double u = 0.0, v = 0.0;
+         for (int qx = 0; qx < Q1D; ++qx)
+         {
+            u += Gt(dx,qx) * QQ0(qy,qx);
+            v += Bt(dx,qx) * QQ1(qy,qx);
+         }
+         DQ0(dx,qy) = u;
+         DQ1(dx,qy) = v;
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+/// Atomic 2D Transposed Gradient, 2/2
+MFEM_HOST_DEVICE inline void AtomicGradXt(const int D1D, const int Q1D,
+                                          const DeviceMatrix &Bt,
+                                          const DeviceMatrix &Gt,
+                                          const DeviceMatrix &DQ0,
+                                          const DeviceMatrix &DQ1,
+                                          const DeviceTensor<3,const int> &I,
+                                          const DeviceMatrix &Y,
+                                          const int c,
+                                          const int e,
+                                          const bool byVDIM)
+{
+   MFEM_FOREACH_THREAD(dy,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         double u = 0.0, v = 0.0;
+         for (int qy = 0; qy < Q1D; ++qy)
+         {
+            u += DQ0(dx,qy) * Bt(dy,qy);
+            v += DQ1(dx,qy) * Gt(dy,qy);
+         }
+         const double sum = u + v;
+         const int gid = I(dx,dy,e);
+         const int idx = gid >= 0 ? gid : -1-gid;
+         if (byVDIM) { AtomicAdd(Y(c,idx), sum); }
+         else { AtomicAdd(Y(idx,c), sum); }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+MFEM_HOST_DEVICE inline
+void Atomic2DGradTranspose(const int D1D, const int Q1D,
+                           const DeviceMatrix &Bt,
+                           const DeviceMatrix &Gt,
+                           const DeviceMatrix &QQ0,
+                           const DeviceMatrix &QQ1,
+                           const DeviceMatrix &DQ0,
+                           const DeviceMatrix &DQ1,
+                           const DeviceTensor<3,const int> &I,
+                           const DeviceMatrix &Y,
+                           const int c,
+                           const int e,
+                           const bool byVDIM)
+{
+   kernels::internal::AtomicGradYt(D1D,Q1D,Bt,Gt,QQ0,QQ1,DQ0,DQ1);
+   kernels::internal::AtomicGradXt(D1D,Q1D,Bt,Gt,DQ0,DQ1,I,Y,c,e,byVDIM);
+}
+
 /// Load 3D scalar input vector into shared memory
 MFEM_HOST_DEVICE inline void LoadX(const int e, const int D1D,
                                    const DeviceTensor<4, const double> &x,
@@ -1543,6 +1721,269 @@ MFEM_HOST_DEVICE inline void GradXt(const int D1D, const int Q1D,
    }
 }
 
+/// Atomic 3D Transposed Evaluation, 1/3
+MFEM_HOST_DEVICE inline void Atomic3DEvalZt(const int D1D, const int Q1D,
+                                            double *u,
+                                            const DeviceMatrix &B,
+                                            const DeviceCube &Q)
+{
+   MFEM_FOREACH_THREAD(qx,x,Q1D)
+   {
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+         for (int qz = 0; qz < Q1D; ++qz)
+         {
+            const double ZYX = Q(qz,qy,qx);
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] += ZYX * B(qz,dz); }
+         }
+         for (int dz = 0; dz < D1D; ++dz) { Q(dz,qy,qx) = u[dz]; }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+/// Atomic 3D Transposed Evaluation, 2/3
+MFEM_HOST_DEVICE inline void Atomic3DEvalYt(const int D1D, const int Q1D,
+                                            double *u,
+                                            const DeviceMatrix &B,
+                                            const DeviceCube &Q)
+{
+   MFEM_FOREACH_THREAD(dz,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         for (int dy = 0; dy < D1D; ++dy) { u[dy] = 0.0; }
+         for (int qy = 0; qy < Q1D; ++qy)
+         {
+            const double zYX = Q(dz,qy,qx);
+            for (int dy = 0; dy < D1D; ++dy) { u[dy] += zYX * B(qy,dy); }
+         }
+         for (int dy = 0; dy < D1D; ++dy) { Q(dz,dy,qx) = u[dy]; }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+/// Atomic 3D Transposed Evaluation, 3/3
+MFEM_HOST_DEVICE inline void Atomic3DEvalXt(const int D1D, const int Q1D,
+                                            double *u,
+                                            const DeviceMatrix &B,
+                                            const DeviceCube &Q,
+                                            const DeviceTensor<4,const int> &I,
+                                            const DeviceMatrix &Y,
+                                            const int c,
+                                            const int e,
+                                            const bool byVDIM)
+{
+   MFEM_FOREACH_THREAD(dz,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(dy,x,D1D)
+      {
+         for (int dx = 0; dx < D1D; ++dx) { u[dx] = 0.0; }
+         for (int qx = 0; qx < Q1D; ++qx)
+         {
+            const double zyX = Q(dz,dy,qx);
+            for (int dx = 0; dx < D1D; ++dx) { u[dx] += zyX * B(qx,dx); }
+         }
+         for (int dx = 0; dx < D1D; ++dx)
+         {
+            const double val = u[dx];
+            const int gid = I(dx,dy,dz,e);
+            const int idx = gid >= 0 ? gid : -1 - gid;
+            if (byVDIM) { AtomicAdd(Y(c,idx), val); }
+            else { AtomicAdd(Y(idx,c), val); }
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+MFEM_HOST_DEVICE inline
+void Atomic3DEvalTranspose(const int D1D,
+                           const int Q1D,
+                           double *u,
+                           const DeviceMatrix &B,
+                           const DeviceCube &Q,
+                           const DeviceTensor<4,const int> &I,
+                           const DeviceMatrix &Y,
+                           const int c,
+                           const int e,
+                           const bool byVDIM)
+{
+   kernels::internal::Atomic3DEvalZt(D1D,Q1D,u,B,Q);
+   kernels::internal::Atomic3DEvalYt(D1D,Q1D,u,B,Q);
+   kernels::internal::Atomic3DEvalXt(D1D,Q1D,u,B,Q,I,Y,c,e,byVDIM);
+}
+
+// Half of B and G are stored in shared to get B, Bt, G and Gt.
+// Indices computation for SmemPADiffusionApply3D.
+static MFEM_HOST_DEVICE inline int qi(const int q, const int d, const int Q)
+{
+   return (q<=d) ? q : Q-1-q;
+}
+
+static MFEM_HOST_DEVICE inline int dj(const int q, const int d, const int D)
+{
+   return (q<=d) ? d : D-1-d;
+}
+
+static MFEM_HOST_DEVICE inline int qk(const int q, const int d, const int Q)
+{
+   return (q<=d) ? Q-1-q : q;
+}
+
+static MFEM_HOST_DEVICE inline int dl(const int q, const int d, const int D)
+{
+   return (q<=d) ? D-1-d : d;
+}
+
+static MFEM_HOST_DEVICE inline double sign(const int q, const int d)
+{
+   return (q<=d) ? -1.0 : 1.0;
+}
+
+/// Atomic 3D Transposed Gradient, 1/3
+MFEM_HOST_DEVICE inline void AtomicGradZt(const int D1D, const int Q1D,
+                                          const DeviceMatrix &Bt,
+                                          const DeviceMatrix &Gt,
+                                          const DeviceCube &QQQ0,
+                                          const DeviceCube &QQQ1,
+                                          const DeviceCube &QQQ2,
+                                          const DeviceCube &QQD0,
+                                          const DeviceCube &QQD1,
+                                          const DeviceCube &QQD2)
+{
+   MFEM_FOREACH_THREAD(qz,z,Q1D)
+   {
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u = 0.0, v = 0.0, w = 0.0;
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               const int i = qi(qx,dx,Q1D);
+               const int j = dj(qx,dx,D1D);
+               const int k = qk(qx,dx,Q1D);
+               const int l = dl(qx,dx,D1D);
+               const double s = sign(qx,dx);
+               u += QQQ0(qz,qy,qx) * Gt(l,k) * s;
+               v += QQQ1(qz,qy,qx) * Bt(j,i);
+               w += QQQ2(qz,qy,qx) * Bt(j,i);
+            }
+            QQD0(qz,qy,dx) = u;
+            QQD1(qz,qy,dx) = v;
+            QQD2(qz,qy,dx) = w;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+/// Atomic 3D Transposed Gradient, 2/3
+MFEM_HOST_DEVICE inline void AtomicGradYt(const int D1D, const int Q1D,
+                                          const DeviceMatrix &Bt,
+                                          const DeviceMatrix &Gt,
+                                          const DeviceCube &QQD0,
+                                          const DeviceCube &QQD1,
+                                          const DeviceCube &QQD2,
+                                          const DeviceCube &QDD0,
+                                          const DeviceCube &QDD1,
+                                          const DeviceCube &QDD2)
+{
+   MFEM_FOREACH_THREAD(qz,z,Q1D)
+   {
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u = 0.0, v = 0.0, w = 0.0;
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               const int i = qi(qy,dy,Q1D);
+               const int j = dj(qy,dy,D1D);
+               const int k = qk(qy,dy,Q1D);
+               const int l = dl(qy,dy,D1D);
+               const double s = sign(qy,dy);
+               u += QQD0(qz,qy,dx) * Bt(j,i);
+               v += QQD1(qz,qy,dx) * Gt(l,k) * s;
+               w += QQD2(qz,qy,dx) * Bt(j,i);
+            }
+            QDD0(qz,dy,dx) = u;
+            QDD1(qz,dy,dx) = v;
+            QDD2(qz,dy,dx) = w;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+/// Atomic 3D Transposed Gradient, 3/3
+MFEM_HOST_DEVICE inline void AtomicGradXt(const int D1D, const int Q1D,
+                                          const DeviceMatrix &Bt,
+                                          const DeviceMatrix &Gt,
+                                          const DeviceCube &QDD0,
+                                          const DeviceCube &QDD1,
+                                          const DeviceCube &QDD2,
+                                          const DeviceTensor<4,const int> &I,
+                                          const DeviceMatrix &Y,
+                                          const int c,
+                                          const int e,
+                                          const bool byVDIM)
+{
+   MFEM_FOREACH_THREAD(dz,z,D1D)
+   {
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u = 0.0, v = 0.0, w = 0.0;
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const int i = qi(qz,dz,Q1D);
+               const int j = dj(qz,dz,D1D);
+               const int k = qk(qz,dz,Q1D);
+               const int l = dl(qz,dz,D1D);
+               const double s = sign(qz,dz);
+               u += QDD0(qz,dy,dx) * Bt(j,i);
+               v += QDD1(qz,dy,dx) * Bt(j,i);
+               w += QDD2(qz,dy,dx) * Gt(l,k) * s;
+            }
+            const double sum = u + v + w;
+            const int gid = I(dx,dy,dz,e);
+            const int idx = gid >= 0 ? gid : -1-gid;
+            if (byVDIM) { AtomicAdd(Y(c,idx), sum); }
+            else { AtomicAdd(Y(idx,c), sum); }
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
+
+MFEM_HOST_DEVICE inline void Atomic3DGrad(const int D1D, const int Q1D,
+                                          const DeviceMatrix &Bt,
+                                          const DeviceMatrix &Gt,
+                                          const DeviceCube &QQ0,
+                                          const DeviceCube &QQ1,
+                                          const DeviceCube &QQ2,
+                                          const DeviceCube &QD0,
+                                          const DeviceCube &QD1,
+                                          const DeviceCube &QD2,
+                                          const DeviceCube &DD0,
+                                          const DeviceCube &DD1,
+                                          const DeviceCube &DD2,
+                                          const DeviceTensor<4,const int> &I,
+                                          const DeviceMatrix &Y,
+                                          const int c,
+                                          const int e,
+                                          const bool byVDIM)
+{
+   kernels::internal::AtomicGradZt(D1D,Q1D,Bt,Gt,QQ0,QQ1,QQ2,QD0,QD1,QD2);
+   kernels::internal::AtomicGradYt(D1D,Q1D,Bt,Gt,QD0,QD1,QD2,DD0,DD1,DD2);
+   kernels::internal::AtomicGradXt(D1D,Q1D,Bt,Gt,DD0,DD1,DD2,I,Y,c,e,byVDIM);
+}
+
 } // namespace kernels::internal
 
 } // namespace kernels

fem/linearform.cpp

@@ -23,6 +23,7 @@ LinearForm::LinearForm(FiniteElementSpace *f, LinearForm *lf)
    UseDevice(true);
 
    fes = f;
+   ext = nullptr;
    extern_lfs = 1;
 
    // Copy the pointers to the integrators
@@ -69,14 +70,14 @@ void LinearForm::AddDomainIntegrator(LinearFormIntegrator *lfi,
 
 void LinearForm::AddBoundaryIntegrator (LinearFormIntegrator * lfi)
 {
-   boundary_integs.Append (lfi);
+   boundary_integs.Append(lfi);
    boundary_integs_marker.Append(NULL); // NULL -> all attributes are active
 }
 
 void LinearForm::AddBoundaryIntegrator (LinearFormIntegrator * lfi,
                                         Array<int> &bdr_attr_marker)
 {
-   boundary_integs.Append (lfi);
+   boundary_integs.Append(lfi);
    boundary_integs_marker.Append(&bdr_attr_marker);
 }
 
@@ -99,8 +100,35 @@ void LinearForm::AddInteriorFaceIntegrator(LinearFormIntegrator *lfi)
    interior_face_integs.Append(lfi);
 }
 
+void LinearForm::SetAssemblyLevel(LinearAssemblyLevel assembly_level)
+{
+   if (ext)
+   {
+      MFEM_ABORT("the assembly level has already been set!");
+   }
+   assembly = assembly_level;
+   switch (assembly)
+   {
+      case LinearAssemblyLevel::LEGACY:
+         break;
+      case LinearAssemblyLevel::FULL:
+         ext = new FullLinearFormExtension(this);
+         break;
+      default:
+         mfem_error("Unknown assembly level");
+   }
+}
+
 void LinearForm::Assemble()
 {
+   Vector::operator=(0.0);
+
+   // The above operation is executed on device because of UseDevice().
+   // The first use of AddElementVector() below will move it back to host
+   // because both 'vdofs' and 'elemvect' are on host.
+
+   if (ext) { return ext->Assemble(); }
+
    Array<int> vdofs;
    ElementTransformation *eltrans;
    DofTransformation *doftrans;
@@ -108,12 +136,6 @@ void LinearForm::Assemble()
 
    int i;
 
-   Vector::operator=(0.0);
-
-   // The above operation is executed on device because of UseDevice().
-   // The first use of AddElementVector() below will move it back to host
-   // because both 'vdofs' and 'elemvect' are on host.
-
    if (domain_integs.Size())
    {
       for (int k = 0; k < domain_integs.Size(); k++)
@@ -143,7 +165,7 @@ void LinearForm::Assemble()
                {
                   doftrans->TransformDual(elemvect);
                }
-               AddElementVector (vdofs, elemvect);
+               AddElementVector(vdofs, elemvect);
             }
          }
       }
@@ -274,6 +296,18 @@ void LinearForm::Assemble()
    }
 }
 
+void LinearForm::Update()
+{
+   SetSize(fes->GetVSize()); ResetDeltaLocations();
+   if (ext) { ext->Update(); }
+}
+
+void LinearForm::Update(FiniteElementSpace *f)
+{
+   fes = f;
+   Update();
+}
+
 void LinearForm::Update(FiniteElementSpace *f, Vector &v, int v_offset)
 {
    MFEM_ASSERT(v.Size() >= v_offset + f->GetVSize(), "");
@@ -281,6 +315,7 @@ void LinearForm::Update(FiniteElementSpace *f, Vector &v, int v_offset)
    v.UseDevice(true);
    this->Vector::MakeRef(v, v_offset, fes->GetVSize());
    ResetDeltaLocations();
+   if (ext) { ext->Update(); }
 }
 
 void LinearForm::MakeRef(FiniteElementSpace *f, Vector &v, int v_offset)
@@ -292,6 +327,8 @@ void LinearForm::AssembleDelta()
 {
    if (domain_delta_integs.Size() == 0) { return; }
 
+   if (ext) { return ext->AssembleDelta(); }
+
    if (!HaveDeltaLocations())
    {
       int sdim = fes->GetMesh()->SpaceDimension();
@@ -356,6 +393,8 @@ LinearForm::~LinearForm()
       for (k=0; k < interior_face_integs.Size(); k++)
       { delete interior_face_integs[k]; }
    }
+
+   delete ext;
 }
 
 }

fem/linearform.hpp

@@ -14,18 +14,39 @@
 
 #include "../config/config.hpp"
 #include "lininteg.hpp"
+#include "linearform_ext.hpp"
 #include "gridfunc.hpp"
 
 namespace mfem
 {
 
+/** @brief Enumeration defining the assembly level for linear form classes
+ *  derived from Vector. */
+enum class LinearAssemblyLevel
+{
+   /// Legacy assembled form, i.e. a global vector in MFEM format.
+   /// This assembly level is ALWAYS performed on the host.
+   LEGACY = 0,
+   /// Fully assembled form, i.e. a global vector in MFEM format.
+   /// This assembly is compatible with device execution.
+   FULL
+};
+
 /// Vector with associated FE space and LinearFormIntegrators.
 class LinearForm : public Vector
 {
+   friend LinearFormExtension;
+
 protected:
    /// FE space on which the LinearForm lives. Not owned.
    FiniteElementSpace *fes;
 
+   /// The assembly level of the form (legacy or full)
+   LinearAssemblyLevel assembly;
+
+   /** @brief Extension for supporting different assembly levels. */
+   LinearFormExtension *ext;
+
    /** @brief Indicates the LinearFormIntegrator%s stored in #domain_integs,
        #domain_delta_integs, #boundary_integs, and #boundary_face_integs are
        owned by another LinearForm. */
@@ -75,7 +96,7 @@ public:
    /// Creates linear form associated with FE space @a *f.
    /** The pointer @a f is not owned by the newly constructed object. */
    LinearForm(FiniteElementSpace *f) : Vector(f->GetVSize())
-   { fes = f; extern_lfs = 0; UseDevice(true); }
+   { fes = f; ext = nullptr; extern_lfs = 0; UseDevice(true); }
 
    /** @brief Create a LinearForm on the FiniteElementSpace @a f, using the
        same integrators as the LinearForm @a lf.
@@ -90,7 +111,7 @@ public:
    /** The associated FiniteElementSpace can be set later using one of the
        methods: Update(FiniteElementSpace *) or
        Update(FiniteElementSpace *, Vector &, int). */
-   LinearForm() { fes = NULL; extern_lfs = 0; UseDevice(true); }
+   LinearForm() { fes = NULL; ext = nullptr; extern_lfs = 0; UseDevice(true); }
 
    /// Construct a LinearForm using previously allocated array @a data.
    /** The LinearForm does not assume ownership of @a data which is assumed to
@@ -98,7 +119,7 @@ public:
        for externally allocated array, the pointer @a data can be NULL. The data
        array can be replaced later using the method SetData(). */
    LinearForm(FiniteElementSpace *f, double *data) : Vector(data, f->GetVSize())
-   { fes = f; extern_lfs = 0; }
+   { fes = f; ext = nullptr; extern_lfs = 0; }
 
    /// Copy assignment. Only the data of the base class Vector is copied.
    /** It is assumed that this object and @a rhs use FiniteElementSpace%s that
@@ -155,6 +176,9 @@ public:
        coefficients. */
    Array<LinearFormIntegrator*> *GetDLFI() { return &domain_integs; }
 
+   /// Access the element attribute marker (of length mesh->attributes)
+   Array<Array<int>*> *GetDLFIM() { return &domain_integs_marker; }
+
    /** @brief Access all integrators added with AddDomainIntegrator() which are
        DeltaLFIntegrator%s with delta coefficients. */
    Array<DeltaLFIntegrator*> *GetDLFI_Delta() { return &domain_delta_integs; }
@@ -173,6 +197,10 @@ public:
        corresponding pointer (to Array<int>) will be NULL. */
    Array<Array<int>*> *GetFLFI_Marker() { return &boundary_face_integs_marker; }
 
+   /// Set the desired assembly level, default is LinearAssemblyLevel::LEGACY.
+   /** This method must be called before assembly. */
+   void SetAssemblyLevel(LinearAssemblyLevel);
+
    /// Assembles the linear form i.e. sums over all domain/bdr integrators.
    void Assemble();
 
@@ -184,11 +212,10 @@ public:
        updated, e.g. after its associated Mesh object has been refined.
 
        @note This method does not perform assembly. */
-   void Update() { SetSize(fes->GetVSize()); ResetDeltaLocations(); }
+   void Update();
 
    /// Associate a new FE space, @a *f, with this object and Update() it. */
-   void Update(FiniteElementSpace *f)
-   { fes = f; SetSize(f->GetVSize()); ResetDeltaLocations(); }
+   void Update(FiniteElementSpace *f);
 
    /** @brief Associate a new FE space, @a *f, with this object and use the data
        of @a v, offset by @a v_offset, to initialize this object's Vector::data.

fem/linearform_ext.cpp

@@ -0,0 +1,104 @@
+// Copyright (c) 2010-2021, 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 "linearform.hpp"
+#include "../general/forall.hpp"
+
+// Implementations of FullLinearFormExtension.
+
+namespace mfem
+{
+
+FullLinearFormExtension::FullLinearFormExtension(LinearForm *lf):
+   LinearFormExtension(lf)
+{
+   Update();
+}
+
+void FullLinearFormExtension::Assemble()
+{
+   MFEM_VERIFY(lf->Size() == lf->FESpace()->GetVSize(), "");
+
+   // Filter out the unsupported integrators
+   MFEM_VERIFY(lf->GetBLFI()->Size() == 0,
+               "Integrators added with AddBoundaryIntegrator() "
+               "are not supported!");
+
+   MFEM_VERIFY(lf->GetDLFI_Delta()->Size() == 0, ""
+               "Integrators added with AddDomainIntegrator() which are "
+               "DeltaLFIntegrators with delta coefficients "
+               "are not supported!");
+
+   MFEM_VERIFY(lf->GetIFLFI()->Size() == 0,
+               "Integrators added with AddInteriorFaceIntegrator() "
+               "are not supported!");
+
+   MFEM_VERIFY(lf->GetFLFI()->Size() == 0,
+               "Integrators added with AddBdrFaceIntegrator() "
+               " are not supported!");
+
+   const FiniteElementSpace &fes = *lf->FESpace();
+   const Array<Array<int>*> &domain_integs_marker = *lf->GetDLFIM();
+   const int mesh_attributes_size = fes.GetMesh()->attributes.Size();
+   const Array<LinearFormIntegrator*> &domain_integs = *lf->GetDLFI();
+
+   for (int k = 0; k < domain_integs.Size(); ++k)
+   {
+      // Get the markers for this integrator
+      const Array<int> *domain_integs_marker_k = domain_integs_marker[k];
+
+      // check if there are markers for this integrator
+      const bool has_markers_k = domain_integs_marker_k != nullptr;
+
+      if (has_markers_k)
+      {
+         // Element attribute marker should be of length mesh->attributes
+         MFEM_VERIFY(mesh_attributes_size == domain_integs_marker_k->Size(),
+                     "invalid element marker for domain linear form "
+                     "integrator #" << k << ", counting from zero");
+      }
+
+      const int NE = fes.GetNE();
+      auto markers_w = markers.Write();
+
+      // if there are no markers, just use the whole linear form (1)
+      if (!has_markers_k)
+      {
+         // done this way as operator= for array are still done on the host
+         MFEM_FORALL(e, NE, markers_w[e] = 1;);
+      }
+      else
+      {
+         // otherwise, scan the attributes to set the markers to 0 or 1
+         const auto attr = attributes.Read();
+         const auto dimk = domain_integs_marker_k->Read();
+         MFEM_FORALL(e, NE, markers_w[e] = dimk[attr[e]-1] == 1;);
+      }
+
+      domain_integs[k]->AssembleFull(fes, markers, *lf);
+   }
+}
+
+void FullLinearFormExtension::Update()
+{
+   MFEM_VERIFY(lf->Size() == lf->FESpace()->GetVSize(), "");
+
+   const int ne = lf->FESpace()->GetNE();
+   const Mesh &mesh = *lf->FESpace()->GetMesh();
+
+   markers.SetSize(ne);
+
+   // Gather the attributes on the host from all the elements
+   attributes.SetSize(ne);
+   for (int i = 0; i < ne; ++i) { attributes[i] = mesh.GetAttribute(i); }
+}
+
+} // namespace mfem

fem/linearform_ext.hpp

@@ -0,0 +1,71 @@
+// Copyright (c) 2010-2021, 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_LINEARFORM_EXT
+#define MFEM_LINEARFORM_EXT
+
+#include "../general/array.hpp"
+
+namespace mfem
+{
+
+class LinearForm;
+class FullLinearFormExtension;
+
+/// Class extending the LinearForm class to support different AssemblyLevels.
+class LinearFormExtension
+{
+protected:
+   /// Linear form from which this extension depends. Not owned.
+   LinearForm *lf;
+
+public:
+   LinearFormExtension(LinearForm *lf): lf(lf) { }
+
+   virtual ~LinearFormExtension() { }
+
+   /// Updates the linear form extension
+   virtual void Update() = 0;
+
+   /// Assemble at the level given for the linear form extension
+   virtual void Assemble() = 0;
+
+   /// Assembles delta functions of the linear form extension
+   virtual void AssembleDelta() = 0;
+};
+
+/// Data and methods for fully-assembled linear forms
+class FullLinearFormExtension : public LinearFormExtension
+{
+private:
+   /// Attributes of all mesh elements.
+   Array<int> attributes;
+
+   /// Temporary markers for device kernels.
+   Array<int> markers;
+
+public:
+   FullLinearFormExtension(LinearForm *lf);
+
+   /// Fully assembles the linear form, compatible with device execution.
+   /// Only integrators added with AddDomainIntegrator are supported.
+   void Assemble() override;
+
+   /// Fully assembles the delta functions of the linear form.
+   /// Not yet supported.
+   void AssembleDelta() override { MFEM_ABORT("Not yet supported!"); }
+
+   void Update() override;
+};
+
+} // namespace mfem
+
+#endif // MFEM_LINEARFORM_EXT

fem/lininteg.cpp

@@ -9,13 +9,21 @@
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
 
-
 #include "fem.hpp"
 #include <cmath>
 
 namespace mfem
 {
 
+void LinearFormIntegrator::AssembleFull(const FiniteElementSpace&,
+                                        const Array<int> &markers,
+                                        Vector &b)
+{
+   MFEM_CONTRACT_VAR(b);
+   MFEM_CONTRACT_VAR(markers);
+   mfem_error("LinearFormIntegrator::AssembleFull(...)");
+}
+
 void LinearFormIntegrator::AssembleRHSElementVect(
    const FiniteElement &el, FaceElementTransformations &Tr, Vector &elvect)
 {
@@ -280,12 +288,11 @@ void VectorDomainLFIntegrator::AssembleRHSElementVect(
       val = Tr.Weight();
 
       el.CalcShape(ip, shape);
-      Q.Eval (Qvec, Tr, ip);
+      Q.Eval(Qvec, Tr, ip);
 
       for (int k = 0; k < vdim; k++)
       {
          cf = val * Qvec(k);
-
          for (int s = 0; s < dof; s++)
          {
             elvect(dof*k+s) += ip.weight * cf * shape(s);
@@ -311,6 +318,54 @@ void VectorDomainLFIntegrator::AssembleDeltaElementVect(
    MultVWt(shape, Qvec, elvec_as_mat);
 }
 
+void VectorDomainLFGradIntegrator::AssembleRHSElementVect(
+   const FiniteElement &el, ElementTransformation &Tr, Vector &elvect)
+{
+   const int dim = el.GetDim();
+   const int dof = el.GetDof();
+   const int vdim = Q.GetVDim();
+   const int sdim = Tr.GetSpaceDim();
+
+   dshape.SetSize(dof,sdim);
+
+   elvect.SetSize(dof*(vdim/sdim));
+   elvect = 0.0;
+
+   const IntegrationRule *ir = IntRule;
+   if (ir == NULL)
+   {
+      int intorder = 2 * el.GetOrder();
+      ir = &IntRules.Get(el.GetGeomType(), intorder);
+   }
+
+   Vector pelvect(dof);
+   Vector part_x(dim);
+
+   for (int q = 0; q < ir->GetNPoints(); q++)
+   {
+      const IntegrationPoint &ip = ir->IntPoint(q);
+
+      Tr.SetIntPoint(&ip);
+      el.CalcPhysDShape(Tr, dshape);
+
+      Q.Eval(Qvec, Tr, ip);
+      Qvec *= ip.weight * Tr.Weight();
+
+      for (int k = 0; k < vdim/sdim; k++)
+      {
+         for (int d=0; d < sdim; ++d) { part_x(d) = Qvec(k*sdim+d); }
+         dshape.Mult(part_x, pelvect);
+         for (int s = 0; s < dof; ++s) { elvect(s+k*dof) += pelvect(s); }
+      }
+   }
+}
+
+void VectorDomainLFGradIntegrator::AssembleDeltaElementVect(
+   const FiniteElement&, ElementTransformation&, Vector&)
+{
+   MFEM_ABORT("Not implemented!");
+}
+
 void VectorBoundaryLFIntegrator::AssembleRHSElementVect(
    const FiniteElement &el, ElementTransformation &Tr, Vector &elvect)
 {

fem/lininteg.hpp

@@ -27,6 +27,11 @@ protected:
    LinearFormIntegrator(const IntegrationRule *ir = NULL) { IntRule = ir; }
 
 public:
+   /// Method defining full assembly
+   virtual void AssembleFull(const FiniteElementSpace &fes,
+                             const Array<int> &markers,
+                             Vector &b);
+
    /** Given a particular Finite Element and a transformation (Tr)
        computes the element vector, elvect. */
    virtual void AssembleRHSElementVect(const FiniteElement &el,
@@ -110,6 +115,10 @@ public:
    DomainLFIntegrator(Coefficient &QF, const IntegrationRule *ir)
       : DeltaLFIntegrator(QF, ir), Q(QF), oa(1), ob(1) { }
 
+   virtual void AssembleFull(const FiniteElementSpace &fes,
+                             const Array<int> &markers,
+                             Vector &b);
+
    /** Given a particular Finite Element and a transformation (Tr)
        computes the element right hand side element vector, elvect. */
    virtual void AssembleRHSElementVect(const FiniteElement &el,
@@ -136,6 +145,10 @@ public:
    DomainLFGradIntegrator(VectorCoefficient &QF)
       : DeltaLFIntegrator(QF), Q(QF) { }
 
+   virtual void AssembleFull(const FiniteElementSpace &fes,
+                             const Array<int> &markers,
+                             Vector &b);
+
    /** Given a particular Finite Element and a transformation (Tr)
        computes the element right hand side element vector, elvect. */
    virtual void AssembleRHSElementVect(const FiniteElement &el,
@@ -223,6 +236,10 @@ public:
    VectorDomainLFIntegrator(VectorCoefficient &QF)
       : DeltaLFIntegrator(QF), Q(QF) { }
 
+   virtual void AssembleFull(const FiniteElementSpace &fes,
+                             const Array<int> &markers,
+                             Vector &b);
+
    /** Given a particular Finite Element and a transformation (Tr)
        computes the element right hand side element vector, elvect. */
    virtual void AssembleRHSElementVect(const FiniteElement &el,
@@ -236,6 +253,37 @@ public:
    using LinearFormIntegrator::AssembleRHSElementVect;
 };
 
+/** Class for domain integrator L(v) := (f, grad v), where
+    f=(f1x,f1y,f1z,...,fnx,fny,fnz) and v=(v1,...,vn). */
+class VectorDomainLFGradIntegrator : public DeltaLFIntegrator
+{
+private:
+   Vector shape, Qvec;
+   VectorCoefficient &Q;
+   DenseMatrix dshape;
+
+public:
+   /// Constructs the domain integrator (Q, grad v)
+   VectorDomainLFGradIntegrator(VectorCoefficient &QF)
+      : DeltaLFIntegrator(QF), Q(QF) { }
+
+   virtual void AssembleFull(const FiniteElementSpace &fes,
+                             const Array<int> &markers,
+                             Vector &b) override;
+
+   /** Given a particular Finite Element and a transformation (Tr)
+       computes the element right hand side element vector, elvect. */
+   virtual void AssembleRHSElementVect(const FiniteElement &el,
+                                       ElementTransformation &Tr,
+                                       Vector &elvect) override;
+
+   virtual void AssembleDeltaElementVect(const FiniteElement &fe,
+                                         ElementTransformation &Trans,
+                                         Vector &elvect) override;
+
+   using LinearFormIntegrator::AssembleRHSElementVect;
+};
+
 /** Class for boundary integration of L(v) := (g, v), where
     f=(f1,...,fn) and v=(v1,...,vn). */
 class VectorBoundaryLFIntegrator : public LinearFormIntegrator

fem/lininteg_domain.cpp

@@ -0,0 +1,104 @@
+// Copyright (c) 2010-2021, 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 "fem.hpp"
+#include "lininteg_domain.hpp"
+
+namespace mfem
+{
+
+using namespace internal::linearform_extension;
+
+void DomainLFIntegrator::AssembleFull(const FiniteElementSpace &fes,
+                                      const Array<int> &markers,
+                                      Vector &y)
+{
+   const int vdim = fes.GetVDim();
+   MFEM_VERIFY(vdim == 1, "vdim should be equal to 1!");
+   GetOrder_f qorder_fct = [&](int el_order) { return oa * el_order + ob; };
+   const IntegrationRule *ir = GetIntRuleFromOrder(fes, IntRule, qorder_fct);
+
+   Vector coeff;
+   const int NQ = ir->GetNPoints();
+   const int NE = fes.GetMesh()->GetNE();
+
+   if (ConstantCoefficient *cQ =
+          dynamic_cast<ConstantCoefficient*>(&Q))
+   {
+      coeff.SetSize(1);
+      coeff(0) = cQ->constant;
+   }
+   else if (QuadratureFunctionCoefficient *cQ =
+               dynamic_cast<QuadratureFunctionCoefficient*>(&Q))
+   {
+      const QuadratureFunction &qfun = cQ->GetQuadFunction();
+      MFEM_VERIFY(qfun.Size() == vdim*NE*NQ,
+                  "Incompatible QuadratureFunction dimension \n");
+      MFEM_VERIFY(ir == &qfun.GetSpace()->GetElementIntRule(0),
+                  "IntegrationRule used within integrator and in"
+                  " QuadratureFunction appear to be different.\n");
+      qfun.Read();
+      coeff.MakeRef(const_cast<QuadratureFunction&>(qfun),0);
+   }
+   else
+   {
+      coeff.SetSize(NQ * NE);
+      auto C = Reshape(coeff.HostWrite(), NQ, NE);
+      for (int e = 0; e < NE; ++e)
+      {
+         ElementTransformation& T = *fes.GetElementTransformation(e);
+         for (int q = 0; q < NQ; ++q)
+         {
+            C(q,e) = Q.Eval(T, ir->IntPoint(q));
+         }
+      }
+   }
+
+   const int id = GetKernelId(fes,ir);
+   const int dim = fes.GetMesh()->Dimension();
+
+   LinearFormExtensionKernel_f ker = nullptr;
+   if (dim==2) { ker=VectorDomainLFIntegratorAssemble2D; }
+   if (dim==3) { ker=VectorDomainLFIntegratorAssemble3D; }
+
+   switch (id)
+   {
+      // 2D kernels, q=p+1
+      case 0x222: ker=VectorDomainLFIntegratorAssemble2D<2,2>; break;
+      case 0x233: ker=VectorDomainLFIntegratorAssemble2D<3,3>; break;
+      case 0x244: ker=VectorDomainLFIntegratorAssemble2D<4,4>; break;
+      case 0x255: ker=VectorDomainLFIntegratorAssemble2D<5,5>; break;
+
+      // 2D kernels, q=p+2
+      case 0x223: ker=VectorDomainLFIntegratorAssemble2D<2,3>; break;
+      case 0x234: ker=VectorDomainLFIntegratorAssemble2D<3,4>; break;
+      case 0x245: ker=VectorDomainLFIntegratorAssemble2D<4,5>; break;
+      case 0x256: ker=VectorDomainLFIntegratorAssemble2D<5,6>; break;
+
+      // 3D kernels, q=p+1, BENCH_FULL_DomainLF_3D_GLL
+      case 0x322: ker=VectorDomainLFIntegratorAssemble3D<2,2>; break;
+      case 0x333: ker=VectorDomainLFIntegratorAssemble3D<3,3>; break;
+      case 0x344: ker=VectorDomainLFIntegratorAssemble3D<4,4>; break;
+      case 0x355: ker=VectorDomainLFIntegratorAssemble3D<5,5>; break;
+      case 0x366: ker=VectorDomainLFIntegratorAssemble3D<6,6>; break;
+
+      // 3D kernels, q=p+2, BENCH_FULL_DomainLF_3D_GL
+      case 0x323: ker=VectorDomainLFIntegratorAssemble3D<2,3>; break;
+      case 0x334: ker=VectorDomainLFIntegratorAssemble3D<3,4>; break;
+      case 0x345: ker=VectorDomainLFIntegratorAssemble3D<4,5>; break;
+      case 0x356: ker=VectorDomainLFIntegratorAssemble3D<5,6>; break;
+      case 0x367: ker=VectorDomainLFIntegratorAssemble3D<6,7>; break;
+   }
+   MFEM_VERIFY(ker, "Unexpected kernel error!");
+   Launch(ker,fes,ir,coeff,markers,y);
+}
+
+} // namespace mfem

fem/lininteg_domain.hpp

@@ -0,0 +1,263 @@
+// Copyright (c) 2010-2021, 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.
+
+// Internal header, included only by .cpp files
+
+#include "fem.hpp"
+
+#include "../fem/kernels.hpp"
+#include "../general/forall.hpp"
+#include "../linalg/kernels.hpp"
+
+namespace mfem
+{
+
+namespace internal
+{
+
+namespace linearform_extension
+{
+
+/// Signature of the kernels used for linear form extension
+using LinearFormExtensionKernel_f = void (*)(const int vdim,
+                                             const bool byVDIM,
+                                             const int ND,
+                                             const int NE,
+                                             const int d,
+                                             const int q,
+                                             const int *markers,
+                                             const double *b,
+                                             const double *g,
+                                             const int *idx,
+                                             const double *J,
+                                             const double *detJ,
+                                             const double *weights,
+                                             const Vector &coeff,
+                                             double *output);
+
+/// Signature of the function used to compute the quadrature order
+using GetOrder_f = std::function<int(int)>;
+
+/// Internal helper function to get the integration rule
+inline const IntegrationRule *GetIntRuleFromOrder(const FiniteElementSpace &fes,
+                                                  const IntegrationRule *IntRule,
+                                                  const GetOrder_f &qorder_fct)
+{
+   const FiniteElement &fe = *fes.GetFE(0);
+   const int qorder = qorder_fct(fe.GetOrder());
+   const Geometry::Type geom_type = fe.GetGeomType();
+   return IntRule ? IntRule : &IntRules.Get(geom_type, qorder);
+}
+
+/// Internal helper function to encode the ID of a LinearFormExtension kernel
+/// The ID is equal to: (dim << 8) |
+///                     (1D number of degrees of freedom << 4) |
+///                     (1D number of quadrature points)
+inline int GetKernelId(const FiniteElementSpace &fes,
+                       const IntegrationRule *ir)
+{
+   Mesh *mesh = fes.GetMesh();
+   const int dim = mesh->Dimension();
+   const FiniteElement &el = *fes.GetFE(0);
+   const DofToQuad &maps = el.GetDofToQuad(*ir, DofToQuad::TENSOR);
+   const int d = maps.ndof;
+   const int q = maps.nqpt;
+   return (dim << 8) | (d << 4) | q;
+}
+
+/// Internal helper function to launch the LinearFormExtension kernel
+inline void Launch(const LinearFormExtensionKernel_f &kernel,
+                   const FiniteElementSpace &fes,
+                   const IntegrationRule *ir,
+                   const Vector &coeff,
+                   const Array<int> &markers,
+                   Vector &y)
+{
+   Mesh *mesh = fes.GetMesh();
+   const int vdim = fes.GetVDim();
+   const bool byVDIM = fes.GetOrdering() == Ordering::byVDIM;
+
+   const FiniteElement &el = *fes.GetFE(0);
+   constexpr int flags = GeometricFactors::JACOBIANS |
+                         GeometricFactors::DETERMINANTS;
+   const MemoryType mt = Device::GetDeviceMemoryType();
+   const GeometricFactors *geom = mesh->GetGeometricFactors(*ir, flags, mt);
+   const DofToQuad &maps = el.GetDofToQuad(*ir, DofToQuad::TENSOR);
+   constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   const Operator *ERop = fes.GetElementRestriction(ordering);
+   const ElementRestriction* ER = dynamic_cast<const ElementRestriction*>(ERop);
+   MFEM_ASSERT(ER, "Not supported!");
+
+   const int *M = markers.Read();
+   const double *B = maps.B.Read();
+   const double *G = maps.G.Read();
+   const double *J = geom->J.Read();
+   const double *detJ = geom->detJ.Read();
+   const int *I = ER->GatherMap().Read();
+   const double *W = ir->GetWeights().Read();
+   double *Y = y.ReadWrite();
+
+   const int ND = fes.GetNDofs();
+   const int NE = fes.GetMesh()->GetNE();
+
+   const int d = maps.ndof;
+   const int q = maps.nqpt;
+
+   kernel(vdim, byVDIM, ND, NE, d, q, M, B, G, I, J, detJ, W, coeff, Y);
+}
+
+/// Internal assembly kernel for the 2D (Vector)DomainLFIntegrator
+template<int D=0, int Q=0> static
+void VectorDomainLFIntegratorAssemble2D(const int vdim,
+                                        const bool byVDIM,
+                                        const int ND,
+                                        const int NE,
+                                        const int d,
+                                        const int q,
+                                        const int *markers,
+                                        const double *b,
+                                        const double *, // g
+                                        const int *idx,
+                                        const double *, // jacobians
+                                        const double *detJ,
+                                        const double *weights,
+                                        const Vector &coeff,
+                                        double *y)
+{
+   constexpr bool USE_SMEM = D > 0 && Q > 0;
+
+   const bool cst_coeff = coeff.Size() == vdim;
+
+   const auto F = coeff.Read();
+   const auto M = Reshape(markers, NE);
+   const auto B = Reshape(b, q,d);
+   const auto DetJ = Reshape(detJ, q,q, NE);
+   const auto W = Reshape(weights, q,q);
+   const auto I = Reshape(idx, d,d, NE);
+   const auto C = cst_coeff ? Reshape(F,vdim,1,1,1) : Reshape(F,vdim,q,q,NE);
+
+   auto Y = Reshape(y, byVDIM ? vdim : ND, byVDIM ? ND : vdim);
+
+   const int sm_size = 2*q*(d+q);
+   constexpr int GRID = USE_SMEM ? 0 : 128;
+   double *gmem = ScratchMem<GRID>(sm_size);
+
+   MFEM_FORALL_3D_GRID(e, NE, q,q,1, GRID,
+   {
+      if (M(e) == 0) { /* ignore */ return; }
+
+      const int bid = MFEM_BLOCK_ID(x);
+      constexpr int SM_SIZE = 2*Q*(D+Q);
+      constexpr bool USE_SMEM = D > 0 && Q > 0;
+      MFEM_SHARED double SMEM[USE_SMEM ? SM_SIZE : 1];
+      double *sm = USE_SMEM ? SMEM : (gmem + sm_size*bid);
+      const DeviceMatrix Bt(DeviceMemAlloc(sm,q*d), q,d);
+      const DeviceMatrix QQ(DeviceMemAlloc(sm,q*q), q,q);
+      const DeviceMatrix QD(DeviceMemAlloc(sm,q*d), q,d);
+
+      kernels::internal::LoadB(d,q,B,Bt);
+
+      for (int c = 0; c < vdim; ++c)
+      {
+         const double cst_val = C(c,0,0,0);
+         MFEM_FOREACH_THREAD(x,x,q)
+         {
+            MFEM_FOREACH_THREAD(y,y,q)
+            {
+               const double detJ = DetJ(x,y,e);
+               const double coeff_val = cst_coeff ? cst_val : C(c,x,y,e);
+               QQ(y,x) = W(x,y) * coeff_val * detJ;
+            }
+         }
+         MFEM_SYNC_THREAD;
+         kernels::internal::Atomic2DEvalTranspose(d,q,Bt,QQ,QD,I,Y,c,e,byVDIM);
+      }
+   });
+}
+
+/// Internal assembly kernel for the 2D (Vector)DomainLFIntegrator
+template<int D=0, int Q=0> static
+void VectorDomainLFIntegratorAssemble3D(const int vdim,
+                                        const bool byVDIM,
+                                        const int ND,
+                                        const int NE,
+                                        const int d,
+                                        const int q,
+                                        const int *markers,
+                                        const double *b,
+                                        const double *, // g
+                                        const int *idx,
+                                        const double *, // jacobians
+                                        const double *detJ,
+                                        const double *weights,
+                                        const Vector &coeff,
+                                        double *y)
+{
+   constexpr bool USE_SMEM = D > 0 && Q > 0;
+
+   const bool cst_coeff = coeff.Size() == vdim;
+
+   const auto F = coeff.Read();
+   const auto M = Reshape(markers, NE);
+   const auto B = Reshape(b, q,d);
+   const auto DetJ = Reshape(detJ, q,q,q, NE);
+   const auto W = Reshape(weights, q,q,q);
+   const auto I = Reshape(idx, d,d,d, NE);
+   const auto C = cst_coeff ? Reshape(F,vdim,1,1,1,1):Reshape(F,vdim,q,q,q,NE);
+
+   auto Y = Reshape(y, byVDIM ? vdim : ND, byVDIM ? ND : vdim);
+
+   const int sm_size = q*d + q*q*q;
+   const int GRID = USE_SMEM ? 0 : 128;
+   double *gmem = ScratchMem<GRID>(sm_size);
+   MFEM_VERIFY(q < 32, "Unsupported quadrature order!");
+
+   MFEM_FORALL_3D_GRID(e, NE, q,q,1, GRID,
+   {
+      if (M(e) == 0) { /* ignore */ return; }
+
+      double u[Q>0?Q:32];
+
+      const int bid = MFEM_BLOCK_ID(x);
+      constexpr int SM_SIZE = Q*D + Q*Q*Q;
+      constexpr bool USE_SMEM = D > 0 && Q > 0;
+      MFEM_SHARED double SMEM[USE_SMEM ? SM_SIZE : 1];
+      double *sm = USE_SMEM ? SMEM : (gmem + sm_size*bid);
+      const DeviceCube QQQ(DeviceMemAlloc(sm,q*q*q), q,q,q);
+      const DeviceMatrix Bt(DeviceMemAlloc(sm,q*d), q,d);
+      kernels::internal::LoadB(d,q,B,Bt);
+
+      for (int c = 0; c < vdim; ++c)
+      {
+         const double cst_val = C(c,0,0,0,0);
+         MFEM_FOREACH_THREAD(x,x,q)
+         {
+            MFEM_FOREACH_THREAD(y,y,q)
+            {
+               for (int z = 0; z < q; ++z)
+               {
+                  const double detJ = DetJ(x,y,z,e);
+                  const double coeff_val = cst_coeff ? cst_val : C(c,x,y,z,e);
+                  QQQ(z,y,x) = W(x,y,z) * coeff_val * detJ;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         kernels::internal::Atomic3DEvalTranspose(d,q,u,Bt,QQQ,I,Y,c,e,byVDIM);
+      }
+   });
+}
+
+} // namespace linearform_extension
+
+} // namespace internal
+
+} // namespace mfem

fem/lininteg_domain_grad.cpp

@@ -0,0 +1,107 @@
+// Copyright (c) 2010-2021, 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 "fem.hpp"
+#include "lininteg_domain.hpp"
+#include "lininteg_domain_grad.hpp"
+
+namespace mfem
+{
+
+using namespace internal::linearform_extension;
+
+void DomainLFGradIntegrator::AssembleFull(const FiniteElementSpace &fes,
+                                          const Array<int> &markers,
+                                          Vector &y)
+{
+   MFEM_VERIFY(fes.GetVDim()==1, "vdim != 1");
+   GetOrder_f gof = [](const int el_order) { return 2.0 * el_order; };
+   const IntegrationRule *ir = GetIntRuleFromOrder(fes, IntRule, gof);
+
+   Vector coeff;
+   const int NQ = ir->GetNPoints();
+   const int NE = fes.GetMesh()->GetNE();
+
+   if (VectorConstantCoefficient *vcQ =
+          dynamic_cast<VectorConstantCoefficient*>(&Q))
+   {
+      coeff = vcQ->GetVec();
+   }
+   else if (VectorQuadratureFunctionCoefficient *vqfQ =
+               dynamic_cast<VectorQuadratureFunctionCoefficient*>(&Q))
+   {
+      const QuadratureFunction &qfun = vqfQ->GetQuadFunction();
+      MFEM_VERIFY(qfun.Size() == NE*NQ,
+                  "Incompatible QuadratureFunction dimension \n");
+      MFEM_VERIFY(ir == &qfun.GetSpace()->GetElementIntRule(0),
+                  "IntegrationRule used within integrator and in"
+                  " QuadratureFunction appear to be different.\n");
+      qfun.Read();
+      coeff.MakeRef(const_cast<QuadratureFunction&>(qfun),0);
+   }
+   else
+   {
+      const int qvdim = Q.GetVDim();
+      Vector Qvec(qvdim);
+      coeff.SetSize(qvdim * NQ * NE);
+      auto C = Reshape(coeff.HostWrite(), qvdim, NQ, NE);
+      for (int e = 0; e < NE; ++e)
+      {
+         ElementTransformation& T = *fes.GetElementTransformation(e);
+         for (int q = 0; q < NQ; ++q)
+         {
+            Q.Eval(Qvec, T, ir->IntPoint(q));
+            for (int c=0; c<qvdim; ++c)
+            {
+               C(c,q,e) = Qvec[c];
+            }
+         }
+      }
+   }
+
+   const int id = GetKernelId(fes,ir);
+   const int dim = fes.GetMesh()->Dimension();
+
+   LinearFormExtensionKernel_f ker = nullptr;
+   if (dim==2) { ker=VectorDomainLFGradIntegratorAssemble2D; }
+   if (dim==3) { ker=VectorDomainLFGradIntegratorAssemble3D; }
+
+   switch (id)
+   {
+      // 2D kernels, q=p+1
+      case 0x222: ker=VectorDomainLFGradIntegratorAssemble2D<2,2>; break;
+      case 0x233: ker=VectorDomainLFGradIntegratorAssemble2D<3,3>; break;
+      case 0x244: ker=VectorDomainLFGradIntegratorAssemble2D<4,4>; break;
+      case 0x255: ker=VectorDomainLFGradIntegratorAssemble2D<5,5>; break;
+
+      // 2D kernels, q=p+2
+      case 0x223: ker=VectorDomainLFGradIntegratorAssemble2D<2,3>; break;
+      case 0x234: ker=VectorDomainLFGradIntegratorAssemble2D<3,4>; break;
+      case 0x245: ker=VectorDomainLFGradIntegratorAssemble2D<4,5>; break;
+      case 0x256: ker=VectorDomainLFGradIntegratorAssemble2D<5,6>; break;
+
+      // 3D kernels, q=p+1
+      case 0x322: ker=VectorDomainLFGradIntegratorAssemble3D<2,2>; break;
+      case 0x333: ker=VectorDomainLFGradIntegratorAssemble3D<3,3>; break;
+      case 0x344: ker=VectorDomainLFGradIntegratorAssemble3D<4,4>; break;
+      case 0x355: ker=VectorDomainLFGradIntegratorAssemble3D<5,5>; break;
+
+      // 3D kernels, q=p+2
+      case 0x323: ker=VectorDomainLFGradIntegratorAssemble3D<2,3>; break;
+      case 0x334: ker=VectorDomainLFGradIntegratorAssemble3D<3,4>; break;
+      case 0x345: ker=VectorDomainLFGradIntegratorAssemble3D<4,5>; break;
+      case 0x356: ker=VectorDomainLFGradIntegratorAssemble3D<5,6>; break;
+   }
+   MFEM_VERIFY(ker, "Unexpected kernel error!");
+   Launch(ker,fes,ir,coeff,markers,y);
+}
+
+} // namespace mfem

fem/lininteg_domain_grad.hpp

@@ -0,0 +1,239 @@
+// Copyright (c) 2010-2021, 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.
+
+// Internal header, included only by .cpp files
+
+#include "fem.hpp"
+
+#include "../general/forall.hpp"
+#include "../fem/kernels.hpp"
+#include "../linalg/kernels.hpp"
+
+namespace mfem
+{
+
+namespace internal
+{
+
+namespace linearform_extension
+{
+
+////////////////////////////////////////////////////////////////////////////////
+template<int D=0, int Q=0> static
+void VectorDomainLFGradIntegratorAssemble2D(const int vdim,
+                                            const bool byVDIM,
+                                            const int ND,
+                                            const int NE,
+                                            const int d,
+                                            const int q,
+                                            const int *markers,
+                                            const double *b,
+                                            const double *g,
+                                            const int *idx,
+                                            const double *jacobians,
+                                            const double *detJ,
+                                            const double *weights,
+                                            const Vector &coeff,
+                                            double *y)
+{
+   constexpr int DIM = 2;
+   constexpr bool USE_SMEM = D > 0 && Q > 0;
+
+   const bool cst_coeff = coeff.Size() == vdim*DIM;
+
+   const auto F = coeff.Read();
+   const auto M = Reshape(markers, NE);
+   const auto B = Reshape(b, q,d);
+   const auto G = Reshape(g, q,d);
+   const auto J = Reshape(jacobians, q,q, DIM,DIM, NE);
+   const auto DetJ = Reshape(detJ, q,q, NE);
+   const auto W = Reshape(weights, q,q);
+   const auto I = Reshape(idx, d,d, NE);
+   const auto C = cst_coeff ?
+                  Reshape(F,DIM,vdim,1,1,1):
+                  Reshape(F,DIM,vdim,q,q,NE);
+
+   auto Y = Reshape(y, byVDIM ? vdim : ND, byVDIM ? ND : vdim);
+
+   const int sm_size = 2*d*q + 4*q*q;
+   const int GRID = USE_SMEM ? 0 : 128;
+   double *gmem = ScratchMem<GRID>(sm_size);
+
+   MFEM_FORALL_3D_GRID(e, NE, q,q,1, GRID,
+   {
+      if (M(e) == 0) { /* ignore */ return; }
+
+      const int bid = MFEM_BLOCK_ID(x);
+      constexpr int SM_SIZE = 2*D*Q + 4*Q*Q;
+      constexpr bool USE_SMEM = D > 0 && Q > 0;
+      MFEM_SHARED double SMEM[USE_SMEM ? SM_SIZE : 1];
+      double *sm = USE_SMEM ? SMEM : (gmem + sm_size*bid);
+
+      const DeviceMatrix Bt(DeviceMemAlloc(sm,q*d), d,q);
+      const DeviceMatrix Gt(DeviceMemAlloc(sm,q*d), d,q);
+      kernels::internal::LoadBGt(d,q,B,G,Bt,Gt);
+
+      const DeviceMatrix QQ0(DeviceMemAlloc(sm,q*q), q,q);
+      const DeviceMatrix QQ1(DeviceMemAlloc(sm,q*q), q,q);
+
+      const DeviceMatrix DQ0(DeviceMemAlloc(sm,d*q), d,q);
+      const DeviceMatrix DQ1(DeviceMemAlloc(sm,d*q), d,q);
+
+      for (int c = 0; c < vdim; ++c)
+      {
+         const double cst_val0 = C(0,c,0,0,0);
+         const double cst_val1 = C(1,c,0,0,0);
+
+         MFEM_FOREACH_THREAD(x,x,q)
+         {
+            MFEM_FOREACH_THREAD(y,y,q)
+            {
+               double Jloc[4], Jinv[4];
+               Jloc[0] = J(x,y,0,0,e);
+               Jloc[1] = J(x,y,1,0,e);
+               Jloc[2] = J(x,y,0,1,e);
+               Jloc[3] = J(x,y,1,1,e);
+               const double detJ = DetJ(x,y,e);
+               kernels::CalcInverse<2>(Jloc, Jinv);
+               const double weight = W(x,y);
+               const double u = cst_coeff ? cst_val0 : C(0,c,x,y,e);
+               const double v = cst_coeff ? cst_val1 : C(1,c,x,y,e);
+               QQ0(y,x) = Jinv[0]*u + Jinv[2]*v;
+               QQ1(y,x) = Jinv[1]*u + Jinv[3]*v;
+               QQ0(y,x) *= weight * detJ;
+               QQ1(y,x) *= weight * detJ;
+            }
+         }
+         MFEM_SYNC_THREAD;
+         kernels::internal::Atomic2DGradTranspose(d,q,Bt,Gt,
+                                                  QQ0,QQ1,DQ0,DQ1,
+                                                  I,Y,c,e,byVDIM);
+      }
+   });
+}
+
+template<int D=0, int Q=0> static
+void VectorDomainLFGradIntegratorAssemble3D(const int vdim,
+                                            const bool byVDIM,
+                                            const int ND,
+                                            const int NE,
+                                            const int d,
+                                            const int q,
+                                            const int *markers,
+                                            const double *b,
+                                            const double *g,
+                                            const int *idx,
+                                            const double *jacobians,
+                                            const double *detJ,
+                                            const double *weights,
+                                            const Vector &coeff,
+                                            double *y)
+{
+   constexpr int DIM = 3;
+   constexpr bool USE_SMEM = D > 0 && Q > 0;
+
+   const bool cst_coeff = coeff.Size() == vdim*DIM;
+
+   const auto F = coeff.Read();
+   const auto M = Reshape(markers, NE);
+   const auto B = Reshape(b, q,d);
+   const auto G = Reshape(g, q,d);
+   const auto J = Reshape(jacobians, q,q,q, DIM,DIM, NE);
+   const auto DetJ = Reshape(detJ, q,q,q, NE);
+   const auto W = Reshape(weights, q,q,q);
+   const auto I = Reshape(idx, d,d,d, NE);
+   const auto C = cst_coeff ?
+                  Reshape(F,DIM,vdim,1,1,1,1):
+                  Reshape(F,DIM,vdim,q,q,q,NE);
+
+   auto Y = Reshape(y, byVDIM ? vdim : ND, byVDIM ? ND : vdim);
+
+   const int sm_size = 2*q*d + 6*q*q*q;
+
+   const int GRID = USE_SMEM ? 0 : 128;
+   double *gmem = ScratchMem<GRID>(sm_size);
+
+   MFEM_FORALL_3D_GRID(e, NE, q,q,1, GRID,
+   {
+      if (M(e) == 0) { /* ignore */ return; }
+
+      const int bid = MFEM_BLOCK_ID(x);
+      constexpr int SM_SIZE = 2*Q*D + 6*Q*Q*Q;
+      constexpr bool USE_SMEM = D > 0 && Q > 0;
+      MFEM_SHARED double SMEM[USE_SMEM ? SM_SIZE : 1];
+      double *sm = USE_SMEM ? SMEM : (gmem + sm_size*bid);
+
+      const DeviceMatrix Bt(DeviceMemAlloc(sm,q*d), d,q);
+      const DeviceMatrix Gt(DeviceMemAlloc(sm,q*d), d,q);
+      kernels::internal::LoadBGt(d,q,B,G,Bt,Gt);
+
+      const DeviceCube QQ0(DeviceMemAlloc(sm,q*q*q), q,q,q);
+      const DeviceCube QQ1(DeviceMemAlloc(sm,q*q*q), q,q,q);
+      const DeviceCube QQ2(DeviceMemAlloc(sm,q*q*q), q,q,q);
+
+      const DeviceCube QD0(DeviceMemAlloc(sm,q*q*q), q,q,d);
+      const DeviceCube QD1(DeviceMemAlloc(sm,q*q*q), q,q,d);
+      const DeviceCube QD2(DeviceMemAlloc(sm,q*q*q), q,q,d);
+
+      const DeviceCube DD0(QQ0,q,d,d);
+      const DeviceCube DD1(QQ1,q,d,d);
+      const DeviceCube DD2(QQ2,q,d,d);
+
+      for (int c = 0; c < vdim; ++c)
+      {
+         const double cst_val_0 = C(0,c,0,0,0,0);
+         const double cst_val_1 = C(1,c,0,0,0,0);
+         const double cst_val_2 = C(2,c,0,0,0,0);
+
+         MFEM_FOREACH_THREAD(x,x,q)
+         {
+            MFEM_FOREACH_THREAD(y,y,q)
+            {
+               for (int z = 0; z < q; ++z)
+               {
+                  double Jloc[9], Jinv[9];
+                  for (int j = 0; j < 3; j++)
+                  {
+                     for (int i = 0; i < 3; i++)
+                     {
+                        Jloc[i+3*j] = J(x,y,z,i,j,e);
+                     }
+                  }
+                  const double detJ = DetJ(x,y,z,e);
+                  kernels::CalcInverse<3>(Jloc, Jinv);
+                  const double weight = W(x,y,z);
+                  const double u = cst_coeff ? cst_val_0 : C(0,c,x,y,z,e);
+                  const double v = cst_coeff ? cst_val_1 : C(1,c,x,y,z,e);
+                  const double w = cst_coeff ? cst_val_2 : C(2,c,x,y,z,e);
+                  QQ0(z,y,x) = Jinv[0]*u + Jinv[3]*v + Jinv[6]*w;
+                  QQ1(z,y,x) = Jinv[1]*u + Jinv[4]*v + Jinv[7]*w;
+                  QQ2(z,y,x) = Jinv[2]*u + Jinv[5]*v + Jinv[8]*w;
+                  QQ0(z,y,x) *= weight * detJ;
+                  QQ1(z,y,x) *= weight * detJ;
+                  QQ2(z,y,x) *= weight * detJ;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         kernels::internal::Atomic3DGrad(d,q,Bt,Gt,
+                                         QQ0,QQ1,QQ2,
+                                         QD0,QD1,QD2,
+                                         DD0,DD1,DD2,
+                                         I,Y,c,e,byVDIM);
+      }
+   });
+}
+
+} // namespace linearform_extension
+
+} // namespace internal
+
+} // namespace mfem

fem/lininteg_vector_domain.cpp

@@ -0,0 +1,102 @@
+// Copyright (c) 2010-2021, 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 "fem.hpp"
+#include "lininteg_domain.hpp"
+
+namespace mfem
+{
+
+using namespace internal::linearform_extension;
+
+void VectorDomainLFIntegrator::AssembleFull(const FiniteElementSpace &fes,
+                                            const Array<int> &markers,
+                                            Vector &y)
+{
+   const int vdim = fes.GetVDim();
+   GetOrder_f gof = [](const int el_order) { return 2.0 * el_order; };
+   const IntegrationRule *ir = GetIntRuleFromOrder(fes, IntRule, gof);
+
+   Vector coeff;
+   const int NQ = ir->GetNPoints();
+   const int NE = fes.GetMesh()->GetNE();
+
+   if (VectorConstantCoefficient *vcQ =
+          dynamic_cast<VectorConstantCoefficient*>(&Q))
+   {
+      coeff = vcQ->GetVec();
+   }
+   else if (VectorQuadratureFunctionCoefficient *vQ =
+               dynamic_cast<VectorQuadratureFunctionCoefficient*>(&Q))
+   {
+      const QuadratureFunction &qfun = vQ->GetQuadFunction();
+      MFEM_VERIFY(qfun.Size() == vdim*NE*NQ,
+                  "Incompatible QuadratureFunction dimension \n");
+      MFEM_VERIFY(ir == &qfun.GetSpace()->GetElementIntRule(0),
+                  "IntegrationRule used within integrator and in"
+                  " QuadratureFunction appear to be different.\n");
+      qfun.Read();
+      coeff.MakeRef(const_cast<QuadratureFunction&>(qfun),0);
+   }
+   else
+   {
+      Vector Qvec(vdim);
+      coeff.SetSize(vdim * NQ * NE);
+      auto C = Reshape(coeff.HostWrite(), vdim, NQ, NE);
+      for (int e = 0; e < NE; ++e)
+      {
+         ElementTransformation& T = *fes.GetElementTransformation(e);
+         for (int q = 0; q < NQ; ++q)
+         {
+            Q.Eval(Qvec, T, ir->IntPoint(q));
+            for (int c=0; c<vdim; ++c) { C(c,q,e) = Qvec[c]; }
+         }
+      }
+   }
+
+   LinearFormExtensionKernel_f ker = nullptr;
+   const int id = GetKernelId(fes,ir);
+   const int dim = fes.GetMesh()->Dimension();
+
+   if (dim==2) { ker = VectorDomainLFIntegratorAssemble2D<>; }
+   if (dim==3) { ker = VectorDomainLFIntegratorAssemble3D<>; }
+
+   switch (id)
+   {
+      // 2D kernels, q=p+1
+      case 0x222: ker=VectorDomainLFIntegratorAssemble2D<2,2>; break;
+      case 0x233: ker=VectorDomainLFIntegratorAssemble2D<3,3>; break;
+      case 0x244: ker=VectorDomainLFIntegratorAssemble2D<4,4>; break;
+      case 0x255: ker=VectorDomainLFIntegratorAssemble2D<5,5>; break;
+
+      // 2D kernels, q=p+2
+      case 0x223: ker=VectorDomainLFIntegratorAssemble2D<2,3>; break;
+      case 0x234: ker=VectorDomainLFIntegratorAssemble2D<3,4>; break;
+      case 0x245: ker=VectorDomainLFIntegratorAssemble2D<4,5>; break;
+      case 0x256: ker=VectorDomainLFIntegratorAssemble2D<5,6>; break;
+
+      // 3D kernels, q=p+1
+      case 0x322: ker=VectorDomainLFIntegratorAssemble3D<2,2>; break;
+      case 0x333: ker=VectorDomainLFIntegratorAssemble3D<3,3>; break;
+      case 0x344: ker=VectorDomainLFIntegratorAssemble3D<4,4>; break;
+      case 0x355: ker=VectorDomainLFIntegratorAssemble3D<5,5>; break;
+
+      // 3D kernels, q=p+2
+      case 0x323: ker=VectorDomainLFIntegratorAssemble3D<2,3>; break;
+      case 0x334: ker=VectorDomainLFIntegratorAssemble3D<3,4>; break;
+      case 0x345: ker=VectorDomainLFIntegratorAssemble3D<4,5>; break;
+      case 0x356: ker=VectorDomainLFIntegratorAssemble3D<5,6>; break;
+   }
+   MFEM_VERIFY(ker, "Unexpected kernel error!");
+   Launch(ker,fes,ir,coeff,markers,y);
+}
+
+} // namespace mfem

fem/lininteg_vector_domain_grad.cpp

@@ -0,0 +1,116 @@
+// Copyright (c) 2010-2021, 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 "fem.hpp"
+#include "lininteg_domain.hpp"
+#include "lininteg_domain_grad.hpp"
+
+namespace mfem
+{
+
+using namespace internal::linearform_extension;
+
+void VectorDomainLFGradIntegrator::AssembleFull(const FiniteElementSpace &fes,
+                                                const Array<int> &markers,
+                                                Vector &y)
+{
+   const int vdim = fes.GetVDim();
+   GetOrder_f gof = [](const int el_order) { return 2.0 * el_order; };
+   const IntegrationRule *ir = GetIntRuleFromOrder(fes, IntRule, gof);
+
+   Vector coeff;
+   const int NQ = ir->GetNPoints();
+   const int NE = fes.GetMesh()->GetNE();
+   const int NS = fes.GetMesh()->SpaceDimension();
+
+   if (VectorConstantCoefficient *vcQ =
+          dynamic_cast<VectorConstantCoefficient*>(&Q))
+   {
+      coeff = vcQ->GetVec();
+   }
+   else if (QuadratureFunctionCoefficient *qfQ =
+               dynamic_cast<QuadratureFunctionCoefficient*>(&Q))
+   {
+      const QuadratureFunction &qfun = qfQ->GetQuadFunction();
+      MFEM_VERIFY(qfun.Size() == NE*NQ,
+                  "Incompatible QuadratureFunction dimension \n");
+      MFEM_VERIFY(ir == &qfun.GetSpace()->GetElementIntRule(0),
+                  "IntegrationRule used within integrator and in"
+                  " QuadratureFunction appear to be different.\n");
+      qfun.Read();
+      coeff.MakeRef(const_cast<QuadratureFunction&>(qfun),0);
+   }
+   else if (VectorQuadratureFunctionCoefficient* vqfQ =
+               dynamic_cast<VectorQuadratureFunctionCoefficient*>(&Q))
+   {
+      const QuadratureFunction &qFun = vqfQ->GetQuadFunction();
+      MFEM_VERIFY(qFun.Size() == vdim * NS * NQ * NE,
+                  "Incompatible QuadratureFunction dimension \n");
+      MFEM_VERIFY(ir == &qFun.GetSpace()->GetElementIntRule(0),
+                  "IntegrationRule used within integrator and in"
+                  " QuadratureFunction appear to be different");
+      qFun.Read();
+      coeff.MakeRef(const_cast<QuadratureFunction &>(qFun),0);
+   }
+   else
+   {
+      Vector Qvec(vdim);
+      coeff.SetSize(vdim * NQ * NE);
+      auto C = Reshape(coeff.HostWrite(), vdim, NQ, NE);
+      for (int e = 0; e < NE; ++e)
+      {
+         ElementTransformation &Tr = *fes.GetElementTransformation(e);
+         for (int q = 0; q < NQ; ++q)
+         {
+            Q.Eval(Qvec, Tr, ir->IntPoint(q));
+            for (int c = 0; c<vdim; ++c) { C(c,q,e) = Qvec[c]; }
+         }
+      }
+   }
+
+   LinearFormExtensionKernel_f ker = nullptr;
+   const int id = GetKernelId(fes,ir);
+   const int dim = fes.GetMesh()->Dimension();
+
+   if (dim==2) { ker=VectorDomainLFGradIntegratorAssemble2D; }
+   if (dim==3) { ker=VectorDomainLFGradIntegratorAssemble3D; }
+
+   switch (id)
+   {
+      // 2D kernels, q=p+1
+      case 0x222: ker=VectorDomainLFGradIntegratorAssemble2D<2,2>; break;
+      case 0x233: ker=VectorDomainLFGradIntegratorAssemble2D<3,3>; break;
+      case 0x244: ker=VectorDomainLFGradIntegratorAssemble2D<4,4>; break;
+      case 0x255: ker=VectorDomainLFGradIntegratorAssemble2D<5,5>; break;
+
+      // 2D kernels, q=p+2
+      case 0x223: ker=VectorDomainLFGradIntegratorAssemble2D<2,3>; break;
+      case 0x234: ker=VectorDomainLFGradIntegratorAssemble2D<3,4>; break;
+      case 0x245: ker=VectorDomainLFGradIntegratorAssemble2D<4,5>; break;
+      case 0x256: ker=VectorDomainLFGradIntegratorAssemble2D<5,6>; break;
+
+      // 3D kernels, q=p+1
+      case 0x322: ker=VectorDomainLFGradIntegratorAssemble3D<2,2>; break;
+      case 0x333: ker=VectorDomainLFGradIntegratorAssemble3D<3,3>; break;
+      case 0x344: ker=VectorDomainLFGradIntegratorAssemble3D<4,4>; break;
+      case 0x355: ker=VectorDomainLFGradIntegratorAssemble3D<5,5>; break;
+
+      // 3D kernels, q=p+2
+      case 0x323: ker=VectorDomainLFGradIntegratorAssemble3D<2,3>; break;
+      case 0x334: ker=VectorDomainLFGradIntegratorAssemble3D<3,4>; break;
+      case 0x345: ker=VectorDomainLFGradIntegratorAssemble3D<4,5>; break;
+      case 0x356: ker=VectorDomainLFGradIntegratorAssemble3D<5,6>; break;
+   }
+   MFEM_VERIFY(ker, "Unexpected kernel error!");
+   Launch(ker,fes,ir,coeff,markers,y);
+}
+
+} // namespace mfem

fem/lor.cpp

@@ -10,8 +10,21 @@
 // CONTRIBUTING.md for details.
 
 #include "lor.hpp"
+#include "lor_assembly.hpp"
+#include "restriction.hpp"
 #include "pbilinearform.hpp"
 
+#include "../mfem-performance.hpp"
+
+#include "../general/forall.hpp"
+
+#define MFEM_DEBUG_COLOR 220
+#include "../general/debug.hpp"
+
+#define MFEM_NVTX_COLOR Turquoise
+#include "../general/nvtx.hpp"
+#include <limits.h>
+
 namespace mfem
 {
 
@@ -21,12 +34,15 @@ void LORBase::AddIntegrators(BilinearForm &a_from,
                              AddIntegratorFn add_integrator,
                              const IntegrationRule *ir)
 {
+   MFEM_NVTX;
    Array<BilinearFormIntegrator*> *integrators = (a_from.*get_integrators)();
    for (int i=0; i<integrators->Size(); ++i)
    {
-      (a_to.*add_integrator)((*integrators)[i]);
-      ir_map[(*integrators)[i]] = ((*integrators)[i])->GetIntegrationRule();
-      if (ir) { ((*integrators)[i])->SetIntegrationRule(*ir); }
+      BilinearFormIntegrator *integrator = (*integrators)[i];
+      if (!integrator->SupportsBatchedLOR()) { supports_batched_assembly = false; }
+      (a_to.*add_integrator)(integrator);
+      ir_map[integrator] = integrator->GetIntegrationRule();
+      if (ir) { integrator->SetIntegrationRule(*ir); }
    }
 }
 
@@ -38,21 +54,24 @@ void LORBase::AddIntegratorsAndMarkers(BilinearForm &a_from,
                                        AddIntegratorFn add_integrator,
                                        const IntegrationRule *ir)
 {
+   MFEM_NVTX;
    Array<BilinearFormIntegrator*> *integrators = (a_from.*get_integrators)();
    Array<Array<int>*> *markers = (a_from.*get_markers)();
 
    for (int i=0; i<integrators->Size(); ++i)
    {
+      BilinearFormIntegrator *integrator = (*integrators)[i];
       if (*markers[i])
       {
-         (a_to.*add_integrator_marker)((*integrators)[i], *(*markers[i]));
+         (a_to.*add_integrator_marker)(integrator, *(*markers[i]));
       }
       else
       {
-         (a_to.*add_integrator)((*integrators)[i]);
+         (a_to.*add_integrator)(integrator);
       }
-      ir_map[(*integrators)[i]] = ((*integrators)[i])->GetIntegrationRule();
-      if (ir) { ((*integrators)[i])->SetIntegrationRule(*ir); }
+      if (!integrator->SupportsBatchedLOR()) { supports_batched_assembly = false; }
+      ir_map[integrator] = integrator->GetIntegrationRule();
+      if (ir) { integrator->SetIntegrationRule(*ir); }
    }
 }
 
@@ -84,6 +103,7 @@ int LORBase::GetLOROrder() const
 
 void LORBase::ConstructLocalDofPermutation(Array<int> &perm_) const
 {
+   MFEM_NVTX;
    FESpaceType type = GetFESpaceType();
    MFEM_VERIFY(type != H1 && type != L2, "");
 
@@ -203,6 +223,7 @@ void LORBase::ConstructLocalDofPermutation(Array<int> &perm_) const
 
 void LORBase::ConstructDofPermutation() const
 {
+   MFEM_NVTX;
    FESpaceType type = GetFESpaceType();
    if (type == H1 || type == L2)
    {
@@ -261,8 +282,24 @@ const OperatorHandle &LORBase::GetAssembledSystem() const
    return A;
 }
 
+const LORRestriction *LORBase::GetLORRestriction() const
+{
+   if (R_lor == NULL)
+   {
+      R_lor = new LORRestriction(fes_ho);
+   }
+   return R_lor;
+}
+
 void LORBase::AssembleSystem_(BilinearForm &a_ho, const Array<int> &ess_dofs)
 {
+   dbg();
+   MFEM_NVTX;
+   // By default, we want to use "batched assembly", however this is only
+   // supported for certain integrators. We set it to true here, and then when
+   // we loop through the integrators, if we encounter unsupported integrators,
+   // we set it to false.
+   supports_batched_assembly = true;
    a->UseExternalIntegrators();
    AddIntegrators(a_ho, *a, &BilinearForm::GetDBFI,
                   &BilinearForm::AddDomainIntegrator, ir_el);
@@ -276,8 +313,35 @@ void LORBase::AssembleSystem_(BilinearForm &a_ho, const Array<int> &ess_dofs)
                             &BilinearForm::GetBFBFI_Marker,
                             &BilinearForm::AddBdrFaceIntegrator,
                             &BilinearForm::AddBdrFaceIntegrator, ir_face);
-   a->Assemble();
-   a->FormSystemMatrix(ess_dofs, A);
+
+   if (supports_batched_assembly)
+   {
+      dbg("supports_batched_assembly");
+      fes_ho.GetMesh()->EnsureNodes();
+#ifdef MFEM_USE_MPI
+      ParFiniteElementSpace *pfes_ho =
+         dynamic_cast<ParFiniteElementSpace*>(&fes_ho);
+      if (pfes_ho)
+      {
+         dbg("=> PARALLEL AssembleBatchedLOR");
+         ParAssembleBatchedLOR(*this, *a, fes_ho, ess_dofs, A);
+      }
+      else
+      {
+         dbg("=> SEQUENTIAL AssembleBatchedLOR");
+         AssembleBatchedLOR(*this, *a, fes_ho, ess_dofs, A);
+      }
+#else
+      AssembleBatchedLOR(*this, *a, fes_ho, ess_dofs, A);
+#endif
+   }
+   else
+   {
+      dbg("NOT supports_batched_assembly");
+      a->Assemble();
+      a->FormSystemMatrix(ess_dofs, A);
+   }
+
    ResetIntegrationRules(&BilinearForm::GetDBFI);
    ResetIntegrationRules(&BilinearForm::GetFBFI);
    ResetIntegrationRules(&BilinearForm::GetBBFI);
@@ -286,6 +350,7 @@ void LORBase::AssembleSystem_(BilinearForm &a_ho, const Array<int> &ess_dofs)
 
 void LORBase::SetupProlongationAndRestriction()
 {
+   MFEM_NVTX;
    if (!HasSameDofNumbering())
    {
       Array<int> p;
@@ -348,6 +413,7 @@ void CheckBasisType(const FiniteElementSpace &fes)
 LORBase::LORBase(FiniteElementSpace &fes_ho_)
    : irs(0, Quadrature1D::GaussLobatto), fes_ho(fes_ho_)
 {
+   MFEM_NVTX;
    Mesh &mesh_ = *fes_ho_.GetMesh();
    int dim = mesh_.Dimension();
    Array<Geometry::Type> geoms;
@@ -363,6 +429,8 @@ LORBase::LORBase(FiniteElementSpace &fes_ho_)
       ir_face = NULL;
    }
    a = NULL;
+   R_lor = NULL;
+   supports_batched_assembly = true;
 }
 
 LORBase::~LORBase()
@@ -378,14 +446,17 @@ LORDiscretization::LORDiscretization(BilinearForm &a_ho_,
                                      int ref_type)
    : LORDiscretization(*a_ho_.FESpace(), ref_type)
 {
+   MFEM_NVTX;
    AssembleSystem(a_ho_, ess_tdof_list);
 }
 
 LORDiscretization::LORDiscretization(FiniteElementSpace &fes_ho,
                                      int ref_type) : LORBase(fes_ho)
 {
+   MFEM_NVTX;
    CheckBasisType(fes_ho);
 
+#if 0
    Mesh &mesh_ho = *fes_ho.GetMesh();
    // For H1, ND and RT spaces, use refinement = element order, for DG spaces,
    // use refinement = element order + 1 (since LOR is p = 0 in this case).
@@ -400,21 +471,29 @@ LORDiscretization::LORDiscretization(FiniteElementSpace &fes_ho,
    fec = fes_ho.FEColl()->Clone(GetLOROrder());
    fes = new FiniteElementSpace(mesh, fec);
    SetupProlongationAndRestriction();
-
+#else
+   mesh = NULL;
+   fec = NULL;
+   fes = NULL;
+#endif
    A.SetType(Operator::MFEM_SPARSEMAT);
 }
 
 void LORDiscretization::AssembleSystem(BilinearForm &a_ho,
                                        const Array<int> &ess_dofs)
 {
+   dbg();
+   MFEM_NVTX;
    delete a;
-   a = new BilinearForm(&GetFESpace());
+   // a = new BilinearForm(&GetFESpace());
+   a = new BilinearForm(&fes_ho);
    AssembleSystem_(a_ho, ess_dofs);
 }
 
 SparseMatrix &LORDiscretization::GetAssembledMatrix() const
 {
-   MFEM_VERIFY(a != NULL && A.Ptr() != NULL, "No LOR system assembled");
+   dbg();
+   MFEM_VERIFY(a != nullptr && A.Ptr() != nullptr, "No LOR system assembled");
    return *A.As<SparseMatrix>();
 }
 
@@ -425,20 +504,26 @@ ParLORDiscretization::ParLORDiscretization(ParBilinearForm &a_ho_,
                                            int ref_type)
    : ParLORDiscretization(*a_ho_.ParFESpace(), ref_type)
 {
+   dbg();
+   MFEM_NVTX;
    AssembleSystem(a_ho_, ess_tdof_list);
 }
 
 ParLORDiscretization::ParLORDiscretization(ParFiniteElementSpace &fes_ho,
                                            int ref_type) : LORBase(fes_ho)
 {
+   dbg();
+   MFEM_NVTX;
    if (fes_ho.GetMyRank() == 0) { CheckBasisType(fes_ho); }
    // TODO: support variable-order spaces in parallel
    MFEM_VERIFY(!fes_ho.IsVariableOrder(),
                "Cannot construct LOR operators on variable-order spaces");
 
+#if 0
    int order = fes_ho.GetMaxElementOrder();
    if (GetFESpaceType() == L2) { ++order; }
 
+   NVTX("ParMesh");
    ParMesh &mesh_ho = *fes_ho.GetParMesh();
    ParMesh *pmesh = new ParMesh(ParMesh::MakeRefined(mesh_ho, order, ref_type));
    mesh = pmesh;
@@ -447,6 +532,11 @@ ParLORDiscretization::ParLORDiscretization(ParFiniteElementSpace &fes_ho,
    ParFiniteElementSpace *pfes = new ParFiniteElementSpace(pmesh, fec);
    fes = pfes;
    SetupProlongationAndRestriction();
+#else
+   mesh = NULL;
+   fec = NULL;
+   fes = NULL;
+#endif
 
    A.SetType(Operator::Hypre_ParCSR);
 }
@@ -454,14 +544,19 @@ ParLORDiscretization::ParLORDiscretization(ParFiniteElementSpace &fes_ho,
 void ParLORDiscretization::AssembleSystem(ParBilinearForm &a_ho,
                                           const Array<int> &ess_dofs)
 {
+   dbg();
+   MFEM_NVTX;
    delete a;
-   a = new ParBilinearForm(&GetParFESpace());
+   // a = new ParBilinearForm(&GetParFESpace());
+   a = new ParBilinearForm(&dynamic_cast<ParFiniteElementSpace&>(fes_ho));
    AssembleSystem_(a_ho, ess_dofs);
 }
 
 HypreParMatrix &ParLORDiscretization::GetAssembledMatrix() const
 {
-   MFEM_VERIFY(a != NULL && A.Ptr() != NULL, "No LOR system assembled");
+   dbg();
+   MFEM_NVTX;
+   MFEM_VERIFY(a != nullptr && A.Ptr() != nullptr, "No LOR system assembled");
    return *A.As<HypreParMatrix>();
 }
 
@@ -470,6 +565,368 @@ ParFiniteElementSpace &ParLORDiscretization::GetParFESpace() const
    return static_cast<ParFiniteElementSpace&>(*fes);
 }
 
-#endif
+#endif // MFEM_USE_MPI
+
+int LORRestriction::GetNRefinedElements(const FiniteElementSpace &fes)
+{
+   int ref = fes.GetMaxElementOrder();
+   int dim = fes.GetMesh()->Dimension();
+   return pow(ref, dim);
+}
+
+FiniteElementCollection *LORRestriction::GetLowOrderFEC(
+   const FiniteElementSpace &fes)
+{
+   return fes.FEColl()->Clone(1);
+}
+
+LORRestriction::LORRestriction(const FiniteElementSpace &fes_ho)
+   : fes_ho(fes_ho),
+     fec_lo(GetLowOrderFEC(fes_ho)),
+     geom(fes_ho.GetMesh()->GetElementGeometry(0)),
+     ne_ref(GetNRefinedElements(fes_ho)),
+     ne(fes_ho.GetNE()*ne_ref),
+     vdim(fes_ho.GetVDim()),
+     byvdim(fes_ho.GetOrdering() == Ordering::byVDIM),
+     ndofs(fes_ho.GetNDofs()),
+     dof(fec_lo->GetFE(geom, 1)->GetDof()),
+
+     offsets(ndofs+1),
+     indices(ne*dof),
+     gatherMap(ne*dof),
+
+     dof_glob2loc(),
+     dof_glob2loc_offsets(),
+     el_dof_lex()
+{
+   MFEM_NVTX;
+   SetupLocalToElement();
+   SetupGlobalToLocal();
+
+   NVTX("EnsureNodes");
+}
+
+void LORRestriction::SetupLocalToElement()
+{
+   MFEM_NVTX;
+   MFEM_VERIFY(ne>0, "ne==0 not supported");
+   const FiniteElement *fe = fec_lo->GetFE(geom, 1);
+   const TensorBasisElement* el =
+      dynamic_cast<const TensorBasisElement*>(fe);
+   MFEM_VERIFY(el, "!TensorBasisElement");
+
+   const Array<int> &fe_dof_map = el->GetDofMap();
+   MFEM_VERIFY(fe_dof_map.Size() > 0, "invalid dof map");
+
+   const FiniteElement *fe_ho = fes_ho.GetFE(0);
+   const TensorBasisElement* tel_ho =
+      dynamic_cast<const TensorBasisElement*>(fe_ho);
+   MFEM_VERIFY(tel_ho, "!TensorBasisElement");
+   const Array<int> &fe_dof_map_ho = tel_ho->GetDofMap();
+
+   int order = fes_ho.GetMaxElementOrder();
+   RefinedGeometry &RG = *GlobGeometryRefiner.Refine(geom, order);
+   Array<int> local_dof_map(dof*ne_ref);
+   for (int ie_lo = 0; ie_lo < ne_ref; ++ie_lo)
+   {
+      for (int i = 0; i < dof; ++i)
+      {
+         int cart_idx = RG.RefGeoms[i + dof*ie_lo]; // local Cartesian index
+         local_dof_map[i + dof*ie_lo] = fe_dof_map_ho[cart_idx];
+      }
+   }
+
+   const Table& e2dTable_ho = fes_ho.GetElementToDofTable();
+
+   auto d_offsets = offsets.Write();
+   const int NDOFS = ndofs;
+   dbg();
+   MFEM_FORALL(i, NDOFS+1, d_offsets[i] = 0;);
+
+   const Memory<int> &J = e2dTable_ho.GetJMemory();
+   const MemoryClass mc = Device::GetDeviceMemoryClass();
+   const int *d_elementMap = J.Read(mc, J.Capacity());
+   const int *d_local_dof_map = local_dof_map.Read();
+   const int DOF = dof;
+   const int DOF_ho = fe_ho->GetDof();
+   const int NE = ne;
+   const int NR_REF = ne_ref;
+
+   dbg();
+   MFEM_FORALL(e, NE,
+   {
+      const int e_ho = e/NR_REF;
+      const int i_ref = e%NR_REF;
+      for (int d = 0; d < DOF; ++d)
+      {
+         const int d_ho = d_local_dof_map[d + i_ref*DOF];
+         const int sgid = d_elementMap[DOF_ho*e_ho + d_ho];  // signed
+         const int gid = (sgid >= 0) ? sgid : -1 - sgid;
+         AtomicAdd(d_offsets[gid+1], 1);
+      }
+   });
+
+   // Aggregate to find offsets for each global dof
+   offsets.HostReadWrite();
+   for (int i = 1; i <= ndofs; ++i) { offsets[i] += offsets[i - 1]; }
+
+   // For each global dof, fill in all local nodes that point to it
+   auto d_gather = gatherMap.Write();
+   auto d_indices = indices.Write();
+   auto drw_offsets = offsets.ReadWrite();
+   const auto dof_map_mem = fe_dof_map.GetMemory();
+   const auto d_dof_map = fe_dof_map.GetMemory().Read(mc,dof_map_mem.Capacity());
+
+   dbg();
+   MFEM_FORALL(e, NE,
+   {
+      const int e_ho = e/NR_REF;
+      const int i_ref = e%NR_REF;
+      for (int d = 0; d < DOF; ++d)
+      {
+         int d_ho = d_local_dof_map[d + i_ref*DOF];
+         const int sdid = d_dof_map[d];  // signed
+         // const int did = d;
+         const int sgid = d_elementMap[DOF_ho*e_ho + d_ho];  // signed
+         const int gid = (sgid >= 0) ? sgid : -1-sgid;
+         const int lid = DOF*e + d;
+         const bool plus = (sgid >= 0 && sdid >= 0) || (sgid < 0 && sdid < 0);
+         d_gather[lid] = plus ? gid : -1-gid;
+         d_indices[AtomicAdd(drw_offsets[gid], 1)] = plus ? lid : -1-lid;
+      }
+   });
+
+   dbg();
+   offsets.HostReadWrite();
+   for (int i = ndofs; i > 0; --i) { offsets[i] = offsets[i - 1]; }
+   offsets[0] = 0;
+}
+
+void LORRestriction::SetupGlobalToLocal()
+{
+   MFEM_NVTX;
+   const int ndof = fes_ho.GetVSize();
+   const int nel_ho = fes_ho.GetMesh()->GetNE();
+   const int order = fes_ho.GetMaxElementOrder();
+   const int dim = fes_ho.GetMesh()->Dimension();
+   MFEM_VERIFY(dim==3, "Not supported");
+   const int nd1d = order + 1;
+   const int ndof_per_el = nd1d*nd1d*nd1d;
+
+   dof_glob2loc.SetSize(2*ndof_per_el*nel_ho);
+   dof_glob2loc_offsets.SetSize(ndof+1);
+   el_dof_lex.SetSize(ndof_per_el*nel_ho);
+
+   Array<int> dofs;
+
+   const Array<int> &lex_map =
+      dynamic_cast<const NodalFiniteElement&>
+      (*fes_ho.GetFE(0)).GetLexicographicOrdering();
+
+   dof_glob2loc_offsets = 0;
+   const Memory<int> &I = fes_ho.GetElementToDofTable().GetIMemory();
+   const Memory<int> &J = fes_ho.GetElementToDofTable().GetJMemory();
+   I.Read(MemoryClass::HOST, I.Capacity());
+   J.Read(MemoryClass::HOST, J.Capacity());
+
+   for (int iel_ho=0; iel_ho<nel_ho; ++iel_ho)
+   {
+      fes_ho.GetElementDofs(iel_ho, dofs);
+      for (int i=0; i<ndof_per_el; ++i)
+      {
+         const int dof = dofs[lex_map[i]];
+         el_dof_lex[i + iel_ho*ndof_per_el] = dof;
+         dof_glob2loc_offsets[dof+1] += 2;
+      }
+   }
+
+   dof_glob2loc_offsets.PartialSum();
+
+   // Sanity check
+   MFEM_VERIFY(dof_glob2loc_offsets[ndof] == dof_glob2loc.Size(), "");
+
+   Array<int> dof_ptr(ndof);
+
+   for (int i=0; i<ndof; ++i) { dof_ptr[i] = dof_glob2loc_offsets[i]; }
+
+   for (int iel_ho=0; iel_ho<nel_ho; ++iel_ho)
+   {
+      fes_ho.GetElementDofs(iel_ho, dofs);
+      for (int i=0; i<ndof_per_el; ++i)
+      {
+         const int dof = dofs[lex_map[i]];
+         dof_glob2loc[dof_ptr[dof]++] = iel_ho;
+         dof_glob2loc[dof_ptr[dof]++] = i;
+      }
+   }
+}
+
+static MFEM_HOST_DEVICE int GetMinElt(const int *my_elts, const int nbElts,
+                                      const int *nbr_elts, const int nbrNbElts)
+{
+   // Find the minimal element index found in both my_elts[] and nbr_elts[]
+   int min_el = INT_MAX;
+   for (int i = 0; i < nbElts; i++)
+   {
+      const int e_i = my_elts[i];
+      if (e_i >= min_el) { continue; }
+      for (int j = 0; j < nbrNbElts; j++)
+      {
+         if (e_i==nbr_elts[j])
+         {
+            min_el = e_i; // we already know e_i < min_el
+            break;
+         }
+      }
+   }
+   return min_el;
+}
+
+int LORRestriction::FillI(SparseMatrix &mat) const
+{
+   MFEM_NVTX;
+   static constexpr int Max = 16;
+   const int all_dofs = ndofs;
+   const int vd = vdim;
+   const int elt_dofs = dof;
+   auto I = mat.ReadWriteI();
+   auto d_offsets = offsets.Read();
+   auto d_indices = indices.Read();
+   auto d_gatherMap = gatherMap.Read();
+   MFEM_FORALL(i_L, vd*all_dofs+1, { I[i_L] = 0; });
+   MFEM_FORALL(e, ne,
+   {
+      for (int i = 0; i < elt_dofs; i++)
+      {
+         int i_elts[Max];
+         const int i_E = e*elt_dofs + i;
+         const int i_L = d_gatherMap[i_E];
+         const int i_offset = d_offsets[i_L];
+         const int i_nextOffset = d_offsets[i_L+1];
+         const int i_nbElts = i_nextOffset - i_offset;
+         for (int e_i = 0; e_i < i_nbElts; ++e_i)
+         {
+            const int i_E = d_indices[i_offset+e_i];
+            i_elts[e_i] = i_E/elt_dofs;
+         }
+         for (int j = 0; j < elt_dofs; j++)
+         {
+            const int j_E = e*elt_dofs + j;
+            const int j_L = d_gatherMap[j_E];
+            const int j_offset = d_offsets[j_L];
+            const int j_nextOffset = d_offsets[j_L+1];
+            const int j_nbElts = j_nextOffset - j_offset;
+            if (i_nbElts == 1 || j_nbElts == 1) // no assembly required
+            {
+               AtomicAdd(I[i_L],1);
+            }
+            else // assembly required
+            {
+               int j_elts[Max];
+               for (int e_j = 0; e_j < j_nbElts; ++e_j)
+               {
+                  const int j_E = d_indices[j_offset+e_j];
+                  const int elt = j_E/elt_dofs;
+                  j_elts[e_j] = elt;
+               }
+               const int min_e = GetMinElt(i_elts, i_nbElts, j_elts, j_nbElts);
+               if (e == min_e) // add the nnz only once
+               {
+                  AtomicAdd(I[i_L],1);
+               }
+            }
+         }
+      }
+   });
+   // We need to sum the entries of I, we do it on CPU as it is very sequential.
+   auto h_I = mat.HostReadWriteI();
+   const int nTdofs = vd*all_dofs;
+   int sum = 0;
+   for (int i = 0; i < nTdofs; i++)
+   {
+      const int nnz = h_I[i];
+      h_I[i] = sum;
+      sum+=nnz;
+   }
+   h_I[nTdofs] = sum;
+   // We return the number of nnz
+   return h_I[nTdofs];
+}
+
+void LORRestriction::FillJAndZeroData(SparseMatrix &mat) const
+{
+   MFEM_NVTX;
+   static constexpr int Max = 8;
+   const int all_dofs = ndofs;
+   const int vd = vdim;
+   const int elt_dofs = dof;
+   auto I = mat.ReadWriteI();
+   auto J = mat.WriteJ();
+   auto Data = mat.WriteData();
+   const int NE = ne;
+   auto d_offsets = offsets.Read();
+   auto d_indices = indices.Read();
+   auto d_gatherMap = gatherMap.Read();
+
+   MFEM_FORALL(e, NE,
+   {
+      for (int i = 0; i < elt_dofs; i++)
+      {
+         int i_elts[Max];
+         const int i_E = e*elt_dofs + i;
+         const int i_L = d_gatherMap[i_E];
+         const int i_offset = d_offsets[i_L];
+         const int i_nextOffset = d_offsets[i_L+1];
+         const int i_nbElts = i_nextOffset - i_offset;
+         for (int e_i = 0; e_i < i_nbElts; ++e_i)
+         {
+            const int i_E = d_indices[i_offset+e_i];
+            i_elts[e_i] = i_E/elt_dofs;
+         }
+         for (int j = 0; j < elt_dofs; j++)
+         {
+            const int j_E = e*elt_dofs + j;
+            const int j_L = d_gatherMap[j_E];
+            const int j_offset = d_offsets[j_L];
+            const int j_nextOffset = d_offsets[j_L+1];
+            const int j_nbElts = j_nextOffset - j_offset;
+            if (i_nbElts == 1 || j_nbElts == 1) // no assembly required
+            {
+               const int nnz = AtomicAdd(I[i_L],1);
+               J[nnz] = j_L;
+               Data[nnz] = 0.0;
+            }
+            else // assembly required
+            {
+               int j_elts[Max];
+               for (int e_j = 0; e_j < j_nbElts; ++e_j)
+               {
+                  const int j_E = d_indices[j_offset+e_j];
+                  const int elt = j_E/elt_dofs;
+                  j_elts[e_j] = elt;
+               }
+               const int min_e = GetMinElt(i_elts, i_nbElts, j_elts, j_nbElts);
+               if (e == min_e) // add the nnz only once
+               {
+                  const int nnz = AtomicAdd(I[i_L],1);
+                  J[nnz] = j_L;
+                  Data[nnz] = 0.0;
+               }
+            }
+         }
+      }
+   });
+   // We need to shift again the entries of I, we do it on CPU as it is very
+   // sequential.
+   auto h_I = mat.HostReadWriteI();
+   const int size = vd*all_dofs;
+   for (int i = 0; i < size; i++) { h_I[size-i] = h_I[size-(i+1)]; }
+   h_I[0] = 0;
+}
+
+LORRestriction::~LORRestriction()
+{
+   delete fec_lo;
+}
 
 } // namespace mfem

fem/lor.hpp

@@ -69,6 +69,10 @@ protected:
    BilinearForm *a;
    OperatorHandle A;
    mutable Array<int> perm;
+   bool supports_batched_assembly;
+
+   /// The LOR element restriction operator.
+   mutable class LORRestriction *R_lor;
 
    /// Constructs the local DOF (ldof) permutation. In parallel this is used as
    /// an intermediate step in computing the DOF permutation (see
@@ -120,6 +124,8 @@ public:
    /// Returns the low-order refined finite element space.
    FiniteElementSpace &GetFESpace() const { return *fes; }
 
+   /// Returns the low-order restriction.
+   const LORRestriction *GetLORRestriction() const;
    ~LORBase();
 };
 
@@ -259,6 +265,51 @@ public:
    ~LORSolver() { if (own_lor) { delete lor; } }
 };
 
+
+/// Create a low-order refined version of a Restriction.
+/// Only used here for the FillI and FillJAndZeroData methods.
+class LORRestriction
+{
+   const FiniteElementSpace &fes_ho;
+   FiniteElementCollection *fec_lo;
+   const Geometry::Type geom;
+   const int ne_ref;
+   const int ne;
+   const int vdim;
+   const bool byvdim;
+   const int ndofs;
+   const int dof;
+
+   Array<int> offsets;
+   Array<int> indices;
+   Array<int> gatherMap;
+
+   Array<int> dof_glob2loc;
+   Array<int> dof_glob2loc_offsets;
+   Array<int> el_dof_lex;
+
+protected:
+   static int GetNRefinedElements(const FiniteElementSpace &fes);
+   static FiniteElementCollection *GetLowOrderFEC(const FiniteElementSpace &fes);
+
+public:
+   LORRestriction(const FiniteElementSpace &fes_ho);
+
+   int FillI(SparseMatrix &mat) const;
+   void FillJAndZeroData(SparseMatrix &mat) const;
+
+   const Array<int> &GatherMap() const { return el_dof_lex; }
+   const Array<int> &Indices() const { return dof_glob2loc; }
+   const Array<int> &Offsets() const { return dof_glob2loc_offsets; }
+
+   ~LORRestriction();
+
+   // Device lambda cannot have private or protected access
+public:
+   void SetupLocalToElement();
+   void SetupGlobalToLocal();
+};
+
 } // namespace mfem
 
 #endif

fem/lor_assembly.cpp

@@ -0,0 +1,412 @@
+// Copyright (c) 2010-2021, 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 "lor.hpp"
+#include "lor_assembly.hpp"
+#include "../linalg/dtensor.hpp"
+#include "../general/forall.hpp"
+
+#define MFEM_DEBUG_COLOR 226
+#include "../general/debug.hpp"
+
+#define MFEM_NVTX_COLOR DarkOrchid
+#include "../general/nvtx.hpp"
+
+#include <thrust/execution_policy.h>
+#include <thrust/scan.h>
+
+namespace mfem
+{
+
+// Defined in lor_assembly_ker.cpp
+template <int order, bool use_smem = true>
+void Assemble3DBatchedLOR(const Array<int> &dof_glob2loc,
+                          const Array<int> &dof_glob2loc_offsets,
+                          const Array<int> &el_dof_lex,
+                          Mesh &mesh_ho,
+                          SparseMatrix &A_mat);
+
+static void AssembleBatchedLORWithoutBC(LORBase &lor_disc,
+                                        BilinearForm &form_lor,
+                                        FiniteElementSpace &fes_ho,
+                                        OperatorHandle &Ah)
+{
+   MFEM_NVTX;
+   Mesh &mesh_ho = *fes_ho.GetMesh();
+   const int dim = mesh_ho.Dimension();
+   const int order = fes_ho.GetMaxElementOrder();
+   dbg("dim:%d order:%d",dim,order);
+
+   const bool has_to_init = Ah.Ptr() == nullptr;
+   dbg("has_to_init: %s", has_to_init?"yes":"no");
+   SparseMatrix *A = has_to_init ? nullptr : Ah.As<SparseMatrix>();
+
+   dbg("GetLORRestriction");
+   const LORRestriction *R = lor_disc.GetLORRestriction();
+   MFEM_VERIFY(R,"LOR Restriction error!");
+
+   if (has_to_init)
+   {
+      MFEM_VERIFY(UsesTensorBasis(fes_ho),
+                  "Batched LOR assembly requires tensor basis");
+      if (Device::IsEnabled()||true)
+      {
+         dbg("Device::IsEnabled()");
+#ifdef MFEM_USE_MPI
+         ParFiniteElementSpace *pfes_ho =
+            dynamic_cast<ParFiniteElementSpace*>(&fes_ho);
+         if (pfes_ho && pfes_ho->GetNRanks() > 1)
+         {
+            dbg("Device::IsEnabled() and multiple ranks!");
+            const int width = pfes_ho->GetVSize();
+            const int height = pfes_ho->GetVSize();
+            dbg("HxW: %dx%d",height,width);
+            A = new SparseMatrix(height, width, 0);
+         }
+         else
+#endif
+         {
+            dbg("Device::IsEnabled() but one rank!");
+            const int width = fes_ho.GetVSize();
+            const int height = fes_ho.GetVSize();
+            dbg("HxW: %dx%d",height,width);
+            A = new SparseMatrix(height, width, 0);
+         }
+         A->GetMemoryI().New(A->Height()+1, A->GetMemoryI().GetMemoryType());
+         const int nnz = R->FillI(*A);
+         A->GetMemoryJ().New(nnz, A->GetMemoryJ().GetMemoryType());
+         A->GetMemoryData().New(nnz, A->GetMemoryData().GetMemoryType());
+         R->FillJAndZeroData(*A); // J, A = 0.0
+      }
+      else
+      {
+         MFEM_ABORT("");
+         dbg("NOT Device::IsEnabled()");
+         // the sparsity pattern is defined from the map: element->dof
+         const int ndofs = fes_ho.GetVSize();
+         dbg("ndofs:%d",ndofs);
+         const Table &elem_dof = form_lor.FESpace()->GetElementToDofTable();
+         Table dof_dof, dof_elem;
+         Transpose(elem_dof, dof_elem, ndofs);
+         mfem::Mult(dof_elem, elem_dof, dof_dof);
+         dof_dof.SortRows();
+         int *I = dof_dof.GetI();
+         int *J = dof_dof.GetJ();
+         double *data = Memory<double>(I[ndofs]);
+         A = new SparseMatrix(I,J,data,ndofs,ndofs,true,true,true);
+         dof_dof.LoseData();
+         *A = 0.0;
+         dbg("done");
+      }
+   }
+
+   void (*Kernel)(const Array<int> &dof_glob2loc,
+                  const Array<int> &dof_glob2loc_offsets,
+                  const Array<int> &el_dof_lex,
+                  Mesh &mesh_ho,
+                  SparseMatrix &A_mat) = nullptr;
+
+   if (dim == 2) { MFEM_ABORT("Unsuported!"); }
+   else if (dim == 3)
+   {
+      switch (order)
+      {
+         case 1: Kernel = Assemble3DBatchedLOR<1>; break;
+         case 2: Kernel = Assemble3DBatchedLOR<2>; break;
+         case 3: Kernel = Assemble3DBatchedLOR<3>; break;
+         case 4: Kernel = Assemble3DBatchedLOR<4>; break;
+         case 5: Kernel = Assemble3DBatchedLOR<5>; break;
+         case 6: Kernel = Assemble3DBatchedLOR<6,false>; break;/*
+         case 7: Kernel = Assemble3DBatchedLOR<7,false>; break;
+         case 8: Kernel = Assemble3DBatchedLOR<8,false>; break;
+         case 9: Kernel = Assemble3DBatchedLOR<9,false>; break;
+         case 10: Kernel = Assemble3DBatchedLOR<10,false>; break;
+         case 11: Kernel = Assemble3DBatchedLOR<11,false>; break;
+         case 12: Kernel = Assemble3DBatchedLOR<12,false>; break;
+         case 13: Kernel = Assemble3DBatchedLOR<13,false>; break;
+         case 14: Kernel = Assemble3DBatchedLOR<14,false>; break;
+         case 15: Kernel = Assemble3DBatchedLOR<15,false>; break;
+         case 16: Kernel = Assemble3DBatchedLOR<16,false>; break;*/
+         default: MFEM_ABORT("Kernel not ready!");
+      }
+   }
+
+   Kernel(R->Indices(),
+          R->Offsets(),
+          R->GatherMap(),
+          mesh_ho, *A);
+
+   A->Finalize();
+
+   if (has_to_init) { Ah.Reset(A); } // A now owns A_mat
+}
+
+
+void AssembleBatchedLOR(LORBase &lor_disc,
+                        BilinearForm &form_lor,
+                        FiniteElementSpace &fes_ho,
+                        const Array<int> &ess_dofs,
+                        OperatorHandle &Ah)
+{
+   MFEM_NVTX;
+   AssembleBatchedLORWithoutBC(lor_disc, form_lor, fes_ho, Ah);
+
+   // Set essential dofs to 0.0
+   const int n_ess_dofs = ess_dofs.Size();
+   const auto ess_dofs_d = ess_dofs.Read();
+
+   const auto I = Ah.As<SparseMatrix>()->ReadI();
+   const auto J = Ah.As<SparseMatrix>()->ReadJ();
+   auto dA = Ah.As<SparseMatrix>()->ReadWriteData();
+
+   MFEM_FORALL(i, n_ess_dofs,
+   {
+      const int idof = ess_dofs_d[i];
+      for (int j=I[idof]; j<I[idof+1]; ++j)
+      {
+         const int jdof = J[j];
+         if (jdof != idof)
+         {
+            dA[j] = 0.0;
+            for (int k=I[jdof]; k<I[jdof+1]; ++k)
+            {
+               if (J[k] == idof)
+               {
+                  dA[k] = 0.0;
+                  break;
+               }
+            }
+         }
+      }
+   });
+}
+
+
+#ifdef MFEM_USE_MPI
+
+void ParAssembleBatchedLOR(LORBase &lor_disc,
+                           BilinearForm &form_lor,
+                           FiniteElementSpace &fes_ho,
+                           const Array<int> &ess_dofs,
+                           OperatorHandle &Ah)
+{
+   dbg();
+   MFEM_NVTX;
+   ParFiniteElementSpace *pfes_ho =
+      dynamic_cast<ParFiniteElementSpace*>(&fes_ho);
+   assert(pfes_ho);
+
+   OperatorHandle A_local(Operator::MFEM_SPARSEMAT);
+   AssembleBatchedLORWithoutBC(lor_disc, form_lor, fes_ho, A_local);
+   MFEM_VERIFY(A_local.As<SparseMatrix>()->Finalized(),
+               "the local matrix must be finalized");
+
+   NVTX("Parallel");
+   OperatorHandle dA(Operator::Hypre_ParCSR),
+                  Ph(Operator::Hypre_ParCSR);
+   {
+      NVTX("MakeSquareBlockDiag");
+      dA.MakeSquareBlockDiag(pfes_ho->GetComm(),
+                             pfes_ho->GlobalVSize(),
+                             pfes_ho->GetDofOffsets(),
+                             A_local.As<SparseMatrix>());
+   }
+   Ph.ConvertFrom(pfes_ho->Dof_TrueDof_Matrix());
+
+   {
+      NVTX("MakePtAP");
+      Ah.MakePtAP(dA, Ph);
+   }
+
+   {
+      dbg("EliminateRowsCols");
+      NVTX("EliminateRowsCols");
+      HypreParMatrix *A_mat = Ah.As<HypreParMatrix>();
+      hypre_ParCSRMatrix *A = *A_mat;
+      A_mat->HypreReadWrite();
+
+      hypre_CSRMatrix *diag = hypre_ParCSRMatrixDiag(A);
+      hypre_CSRMatrix *offd = hypre_ParCSRMatrixOffd(A);
+
+      HYPRE_Int diag_nrows = hypre_CSRMatrixNumRows(diag);
+      HYPRE_Int offd_ncols = hypre_CSRMatrixNumCols(offd);
+      dbg("diag_nrows:%d offd_ncols:%d", diag_nrows, offd_ncols);
+
+      const int n_ess_dofs = ess_dofs.Size();
+      const auto ess_dofs_d = ess_dofs.Read();
+      dbg("n_ess_dofs:%d", n_ess_dofs);
+
+      // Start communication to figure out which columns need to be eliminated in
+      // the off-diagonal block
+      hypre_ParCSRCommHandle *comm_handle;
+      HYPRE_Int *int_buf_data, *eliminate_row, *eliminate_col;
+      {
+         eliminate_row = hypre_CTAlloc(HYPRE_Int, diag_nrows, HYPRE_MEMORY_DEVICE);
+         eliminate_col = hypre_CTAlloc(HYPRE_Int, offd_ncols, HYPRE_MEMORY_DEVICE);
+
+         // Get the communication package for A, creating it if it does not
+         // already exist.
+         hypre_ParCSRCommPkg *comm_pkg = hypre_ParCSRMatrixCommPkg(A);
+         if (!comm_pkg)
+         {
+            hypre_MatvecCommPkgCreate(A);
+            comm_pkg = hypre_ParCSRMatrixCommPkg(A);
+         }
+
+         MFEM_FORALL(i, diag_nrows,
+         {
+            eliminate_row[i] = 0;
+         });
+         MFEM_FORALL(i, n_ess_dofs,
+         {
+            eliminate_row[ess_dofs_d[i]] = 1;
+         });
+
+         // Use a matvec communication pattern to find (in eliminate_col) which of
+         // the local offd columns are to be eliminated
+         HYPRE_Int num_sends = hypre_ParCSRCommPkgNumSends(comm_pkg);
+         HYPRE_Int int_buf_sz = hypre_ParCSRCommPkgSendMapStart(comm_pkg, num_sends);
+         dbg("num_sends:%d", num_sends);
+         int_buf_data = hypre_CTAlloc(HYPRE_Int, int_buf_sz, HYPRE_MEMORY_DEVICE);
+
+         hypre_ParCSRCommPkgCopySendMapElmtsToDevice(comm_pkg);
+         HYPRE_Int *send_map_elmts = hypre_ParCSRCommPkgDeviceSendMapElmts(comm_pkg);
+         MFEM_FORALL(i, int_buf_sz,
+         {
+            int k = send_map_elmts[i];
+            int_buf_data[i] = eliminate_row[k];
+         });
+
+         // Try to use device-aware MPI for the communication
+         comm_handle = hypre_ParCSRCommHandleCreate_v2(
+            11, comm_pkg, HYPRE_MEMORY_DEVICE, int_buf_data,
+            HYPRE_MEMORY_DEVICE, eliminate_col);
+      }
+
+      // Eliminate rows and columns in the diagonal block
+      {
+         dbg("Eliminate rows and columns in the diagonal block");
+         const auto I = diag->i;
+         const auto J = diag->j;
+         auto data = diag->data;
+
+         MFEM_FORALL(i, n_ess_dofs,
+         {
+            const int idof = ess_dofs_d[i];
+            for (int j=I[idof]; j<I[idof+1]; ++j)
+            {
+               const int jdof = J[j];
+               if (jdof != idof)
+               {
+                  data[j] = 0.0;
+                  for (int k=I[jdof]; k<I[jdof+1]; ++k)
+                  {
+                     if (J[k] == idof)
+                     {
+                        data[k] = 0.0;
+                        break;
+                     }
+                  }
+               }
+            }
+         });
+      }
+
+      // Eliminate rows in the off-diagonal block
+      {
+         dbg("Eliminate rows in the off-diagonal block");
+         const auto I = offd->i;
+         auto data = offd->data;
+         MFEM_FORALL(i, n_ess_dofs,
+         {
+            const int idof = ess_dofs_d[i];
+            for (int j=I[idof]; j<I[idof+1]; ++j)
+            {
+               data[j] = 0.0;
+            }
+         });
+      }
+
+      // Wait for MPI communication to finish
+      Array<HYPRE_Int> cols_to_eliminate;
+      {
+         dbg("Wait for MPI communication to finish");
+         hypre_ParCSRCommHandleDestroy(comm_handle);
+
+         // Convert a "column marker" array to a "column index" array.
+
+         // First we need to know how many columns are marked, and do a partial
+         // sum to figure out which index they correspond to.
+         HYPRE_Int *col_idx = hypre_CTAlloc(HYPRE_Int, offd_ncols, HYPRE_MEMORY_DEVICE);
+         thrust::exclusive_scan(
+            thrust::device, eliminate_col, eliminate_col + offd_ncols, col_idx
+         );
+
+         // Because we do an exclusive scan above, the last entry is missing
+         // from the partial sum. To compute the total sum, get last entry of
+         // both arrays and add them.
+         HYPRE_Int partial_sum, last_element;
+         hypre_TMemcpy(&partial_sum, col_idx + offd_ncols - 1, HYPRE_Int, 1, HYPRE_MEMORY_HOST, HYPRE_MEMORY_DEVICE);
+         hypre_TMemcpy(&last_element, eliminate_col + offd_ncols - 1, HYPRE_Int, 1, HYPRE_MEMORY_HOST, HYPRE_MEMORY_DEVICE);
+         HYPRE_Int ncols_to_eliminate = partial_sum + last_element;
+
+         // Perform the conversion to "column index" array.
+         cols_to_eliminate.SetSize(ncols_to_eliminate);
+         HYPRE_Int *cols = cols_to_eliminate.Write();
+         MFEM_FORALL(i, offd_ncols,
+         {
+            if (eliminate_col[i])
+            {
+               cols[col_idx[i]] = i;
+            }
+         });
+
+         hypre_TFree(int_buf_data, HYPRE_MEMORY_DEVICE);
+         hypre_TFree(eliminate_row, HYPRE_MEMORY_DEVICE);
+         hypre_TFree(eliminate_col, HYPRE_MEMORY_DEVICE);
+         hypre_TFree(col_idx, HYPRE_MEMORY_DEVICE);
+      }
+
+      // Eliminate columns in the off-diagonal block
+      {
+         dbg("Eliminate columns in the off-diagonal block");
+         const int ncols_to_eliminate = cols_to_eliminate.Size();
+         const int nrows_offd = hypre_CSRMatrixNumRows(offd);
+         const auto cols = cols_to_eliminate.Read();
+         const auto I = offd->i;
+         const auto J = offd->j;
+         auto data = offd->data;
+         dbg("ncols_to_eliminate:%d nrows_offd:%d", ncols_to_eliminate, nrows_offd);
+         // Note: could also try a different strategy, looping over nnz in the
+         // matrix and then doing a binary search in ncols_to_eliminate to see if
+         // the column should be eliminated.
+         MFEM_FORALL(idx, ncols_to_eliminate,
+         {
+            const int j = cols[idx];
+            for (int i=0; i<nrows_offd; ++i)
+            {
+               for (int jj=I[i]; jj<I[i+1]; ++jj)
+               {
+                  if (J[jj] == j)
+                  {
+                     data[jj] = 0.0;
+                     break;
+                  }
+               }
+            }
+         });
+      }
+   }
+}
+
+#endif // MFEM_USE_MPI
+
+} // namespace mfem

fem/lor_assembly.hpp

@@ -0,0 +1,40 @@
+// Copyright (c) 2010-2021, 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_LOR_ASSEMBLY
+#define MFEM_LOR_ASSEMBLY
+
+#include "lor.hpp"
+#include "bilinearform.hpp"
+#include "pbilinearform.hpp"
+
+namespace mfem
+{
+
+void AssembleBatchedLOR(LORBase &lor_disc,
+                        BilinearForm &form_lo,
+                        FiniteElementSpace &fes_ho,
+                        const Array<int> &ess_dofs,
+                        OperatorHandle &A);
+
+#ifdef MFEM_USE_MPI
+
+void ParAssembleBatchedLOR(LORBase &lor_disc,
+                           BilinearForm &form_lo,
+                           FiniteElementSpace &fes_ho,
+                           const Array<int> &ess_dofs,
+                           OperatorHandle &A);
+
+#endif // MFEM_USE_MPI
+
+} // namespace mfem
+
+#endif

fem/lor_assembly_ker.cpp

@@ -0,0 +1,533 @@
+// Copyright (c) 2010-2021, 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 "fem.hpp"
+#include "../general/forall.hpp"
+
+#define MFEM_DEBUG_COLOR 187
+#include "../general/debug.hpp"
+
+#define MFEM_NVTX_COLOR SlateBlue
+#include "../general/nvtx.hpp"
+
+namespace mfem
+{
+
+template<int D1D, int Q1D>
+void NodalInterpolation3D(const int NE,
+                          const Vector& localL, Vector& localH,
+                          const Array<double>& B);
+
+template <int order, bool USE_SMEM = true>
+void Assemble3DBatchedLOR(const Array<int> &dof_glob2loc_,
+                          const Array<int> &dof_glob2loc_offsets_,
+                          const Array<int> &el_dof_lex_,
+                          Mesh &mesh_ho,
+                          SparseMatrix &A_mat)
+{
+   const int nel_ho = mesh_ho.GetNE();
+
+   static constexpr int nv = 8;
+   static constexpr int dim = 3;
+   static constexpr int ddm2 = (dim*(dim+1))/2;
+   static constexpr int nd1d = order + 1;
+   static constexpr int ndof_per_el = nd1d*nd1d*nd1d;
+   static constexpr int nnz_per_row = 27;
+   static constexpr int nnz_per_el = nnz_per_row * ndof_per_el;
+   static constexpr int sz_grad_A = 3*3*2*2*2*2;
+   static constexpr int sz_grad_B = sz_grad_A*2;
+   static constexpr int sz_local_mat = 8*8;
+
+   static constexpr int GRID = USE_SMEM ? 0 : 128;
+   double *GM = nullptr;
+   static Vector *d_buffer = nullptr;
+   if (!USE_SMEM)
+   {
+      if (!d_buffer)
+      {
+         d_buffer = new Vector();
+         d_buffer->UseDevice(true);
+      }
+      d_buffer->SetSize(nnz_per_el*GRID);
+      GM = d_buffer->Write();
+   }
+
+   const auto el_dof_lex = Reshape(el_dof_lex_.Read(), ndof_per_el, nel_ho);
+   const auto dof_glob2loc = dof_glob2loc_.Read();
+   const auto K = dof_glob2loc_offsets_.Read();
+
+   const auto I = A_mat.ReadI();
+   const auto J = A_mat.ReadJ();
+   auto A = A_mat.ReadWriteData();
+
+   const GridFunction *nodal_gf = mesh_ho.GetNodes();
+   const FiniteElementSpace *nodal_fes = nodal_gf->FESpace();
+   const Operator *nodal_restriction = nodal_fes->GetElementRestriction(
+                                          ElementDofOrdering::LEXICOGRAPHIC);
+   const int nodal_nd1d = nodal_fes->GetMaxElementOrder() + 1;
+
+   IntegrationRules irs(0, Quadrature1D::GaussLobatto);
+   const IntegrationRule &ir = irs.Get(Geometry::Type::CUBE, 2*nd1d - 3);
+   MFEM_VERIFY(ir.Size() == ndof_per_el, "");
+
+   // Get the map from mesh nodes to LOR vertices
+   const DofToQuad& maps =
+      nodal_fes->GetFE(0)->GetDofToQuad(ir, DofToQuad::TENSOR);
+
+   // Map from nodal E-vector to L-vector
+   Vector nodes_loc(nodal_restriction->Height());
+   nodes_loc.UseDevice(true);
+   nodal_restriction->Mult(*nodal_gf, nodes_loc);
+
+   // Get nodal points at the LOR vertices
+   Vector X_loc(dim*ndof_per_el*nel_ho);
+   X_loc.UseDevice(true);
+
+   // Get the LOR vertex coordinates
+   MFEM_VERIFY(nd1d==order+1, "nd1d!=order+1");
+   switch (nodal_nd1d)
+   {
+      case 2:
+      {
+         NodalInterpolation3D<2,nd1d>(nel_ho, nodes_loc, X_loc, maps.B);
+         break;
+      }
+      case 4:
+      {
+         NodalInterpolation3D<4,nd1d>(nel_ho, nodes_loc, X_loc, maps.B);
+         break;
+      }
+      case 6:
+      {
+         NodalInterpolation3D<6,nd1d>(nel_ho, nodes_loc, X_loc, maps.B);
+         break;
+      }
+      default: MFEM_ABORT("Unsuported mesh order!");
+   }
+   auto X = X_loc.Read();
+
+   // Last GRID dimension is lowered to avoid too many resources
+   MFEM_FORALL_3D_GRID(iel_ho, nel_ho, order, order, USE_SMEM?order:1, GRID,
+   {
+      const int bid = MFEM_BLOCK_ID(x);
+      MFEM_SHARED double smem[USE_SMEM ? nnz_per_el : 1];
+      double *V_ = USE_SMEM ? smem : GM + nnz_per_el*bid;
+      DeviceTensor<4> V(V_, nnz_per_row, nd1d, nd1d, nd1d);
+
+      // Assemble a sparse matrix over the macro-element by looping over each
+      // subelement.
+      // V(j,i) stores the jth nonzero in the ith row of the sparse matrix.
+      MFEM_FOREACH_THREAD(iz,z,nd1d)
+      {
+         MFEM_FOREACH_THREAD(iy,y,nd1d)
+         {
+            MFEM_FOREACH_THREAD(ix,x,nd1d)
+            {
+               //MFEM_UNROLL(27)
+               for (int j=0; j<nnz_per_row; ++j)
+               {
+                  V(j,ix,iy,iz) = 0.0;
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Compute geometric factors at quadrature points
+      MFEM_FOREACH_THREAD(kz,z,order)
+      {
+         MFEM_FOREACH_THREAD(ky,y,order)
+         {
+            MFEM_FOREACH_THREAD(kx,x,order)
+            {
+               double Q_[ddm2*nv];
+               double grad_A_[sz_grad_A];
+               double grad_B_[sz_grad_B];
+               double local_mat_[sz_local_mat];
+
+               DeviceTensor<4> Q(Q_, ddm2,2,2,2);
+               DeviceTensor<2> local_mat(local_mat_, 8, 8);
+               DeviceTensor<6> grad_A(grad_A_, 3, 3, 2, 2, 2, 2);
+               DeviceTensor<7> grad_B(grad_B_, 3, 3, 2, 2, 2, 2, 2);
+
+               // local_mat is the local (dense) stiffness matrix
+               for (int i=0; i<sz_local_mat; ++i) { local_mat[i] = 0.0; }
+
+               // Intermediate quantities
+               // (see e.g. Mora and Demkowicz for notation).
+               for (int i=0; i<sz_grad_A; ++i) { grad_A[i] = 0.0; }
+               for (int i=0; i<sz_grad_B; ++i) { grad_B[i] = 0.0; }
+
+               const int v0 = kx + nd1d*(ky + nd1d*kz);
+               const int v1 = kx + 1 + nd1d*(ky + nd1d*kz);
+               const int v2 = kx + 1 + nd1d*(ky + 1 + nd1d*kz);
+               const int v3 = kx + nd1d*(ky + 1 + nd1d*kz);
+               const int v4 = kx + nd1d*(ky + nd1d*(kz + 1));
+               const int v5 = kx + 1 + nd1d*(ky + nd1d*(kz + 1));
+               const int v6 = kx + 1 + nd1d*(ky + 1 + nd1d*(kz + 1));
+               const int v7 = kx + nd1d*(ky + 1 + nd1d*(kz + 1));
+
+               const int e0 = dim*(v0 + ndof_per_el*iel_ho);
+               const int e1 = dim*(v1 + ndof_per_el*iel_ho);
+               const int e2 = dim*(v2 + ndof_per_el*iel_ho);
+               const int e3 = dim*(v3 + ndof_per_el*iel_ho);
+               const int e4 = dim*(v4 + ndof_per_el*iel_ho);
+               const int e5 = dim*(v5 + ndof_per_el*iel_ho);
+               const int e6 = dim*(v6 + ndof_per_el*iel_ho);
+               const int e7 = dim*(v7 + ndof_per_el*iel_ho);
+
+               const double v0x = X[e0 + 0];
+               const double v0y = X[e0 + 1];
+               const double v0z = X[e0 + 2];
+
+               const double v1x = X[e1 + 0];
+               const double v1y = X[e1 + 1];
+               const double v1z = X[e1 + 2];
+
+               const double v2x = X[e2 + 0];
+               const double v2y = X[e2 + 1];
+               const double v2z = X[e2 + 2];
+
+               const double v3x = X[e3 + 0];
+               const double v3y = X[e3 + 1];
+               const double v3z = X[e3 + 2];
+
+               const double v4x = X[e4 + 0];
+               const double v4y = X[e4 + 1];
+               const double v4z = X[e4 + 2];
+
+               const double v5x = X[e5 + 0];
+               const double v5y = X[e5 + 1];
+               const double v5z = X[e5 + 2];
+
+               const double v6x = X[e6 + 0];
+               const double v6y = X[e6 + 1];
+               const double v6z = X[e6 + 2];
+
+               const double v7x = X[e7 + 0];
+               const double v7y = X[e7 + 1];
+               const double v7z = X[e7 + 2];
+
+               //MFEM_UNROLL(2)
+               for (int iqz=0; iqz<2; ++iqz)
+               {
+                  //MFEM_UNROLL(2)
+                  for (int iqy=0; iqy<2; ++iqy)
+                  {
+                     //MFEM_UNROLL(2)
+                     for (int iqx=0; iqx<2; ++iqx)
+                     {
+
+                        const double x = iqx;
+                        const double y = iqy;
+                        const double z = iqz;
+                        const double w = 1.0/8.0;
+
+                        // c: (1-x)(1-y)(1-z)v0[c] + x (1-y)(1-z)v1[c] + x y (1-z)v2[c] + (1-x) y (1-z)v3[c]
+                        //  + (1-x)(1-y) z   v4[c] + x (1-y) z   v5[c] + x y z    v6[c] + (1-x) y z    v7[c]
+                        const double J11 = -(1-y)*(1-z)*v0x
+                                           + (1-y)*(1-z)*v1x + y*(1-z)*v2x - y*(1-z)*v3x
+                                           - (1-y)*z*v4x + (1-y)*z*v5x + y*z*v6x - y*z*v7x;
+
+                        const double J12 = -(1-x)*(1-z)*v0x
+                                           - x*(1-z)*v1x + x*(1-z)*v2x + (1-x)*(1-z)*v3x
+                                           - (1-x)*z*v4x - x*z*v5x + x*z*v6x + (1-x)*z*v7x;
+
+                        const double J13 = -(1-x)*(1-y)*v0x - x*(1-y)*v1x
+                                           - x*y*v2x - (1-x)*y*v3x + (1-x)*(1-y)*v4x
+                                           + x*(1-y)*v5x + x*y*v6x + (1-x)*y*v7x;
+
+                        const double J21 = -(1-y)*(1-z)*v0y + (1-y)*(1-z)*v1y
+                                           + y*(1-z)*v2y - y*(1-z)*v3y - (1-y)*z*v4y
+                                           + (1-y)*z*v5y + y*z*v6y - y*z*v7y;
+
+                        const double J22 = -(1-x)*(1-z)*v0y - x*(1-z)*v1y
+                                           + x*(1-z)*v2y + (1-x)*(1-z)*v3y- (1-x)*z*v4y -
+                                           x*z*v5y + x*z*v6y + (1-x)*z*v7y;
+
+                        const double J23 = -(1-x)*(1-y)*v0y - x*(1-y)*v1y
+                                           - x*y*v2y - (1-x)*y*v3y + (1-x)*(1-y)*v4y
+                                           + x*(1-y)*v5y + x*y*v6y + (1-x)*y*v7y;
+
+                        const double J31 = -(1-y)*(1-z)*v0z + (1-y)*(1-z)*v1z
+                                           + y*(1-z)*v2z - y*(1-z)*v3z- (1-y)*z*v4z +
+                                           (1-y)*z*v5z + y*z*v6z - y*z*v7z;
+
+                        const double J32 = -(1-x)*(1-z)*v0z - x*(1-z)*v1z
+                                           + x*(1-z)*v2z + (1-x)*(1-z)*v3z - (1-x)*z*v4z
+                                           - x*z*v5z + x*z*v6z + (1-x)*z*v7z;
+
+                        const double J33 = -(1-x)*(1-y)*v0z - x*(1-y)*v1z
+                                           - x*y*v2z - (1-x)*y*v3z + (1-x)*(1-y)*v4z
+                                           + x*(1-y)*v5z + x*y*v6z + (1-x)*y*v7z;
+
+                        const double detJ = J11 * (J22 * J33 - J32 * J23) -
+                                            J21 * (J12 * J33 - J32 * J13) +
+                                            J31 * (J12 * J23 - J22 * J13);
+                        const double w_detJ = w/detJ;
+
+                        // adj(J)
+                        const double A11 = (J22 * J33) - (J23 * J32);
+                        const double A12 = (J32 * J13) - (J12 * J33);
+                        const double A13 = (J12 * J23) - (J22 * J13);
+                        const double A21 = (J31 * J23) - (J21 * J33);
+                        const double A22 = (J11 * J33) - (J13 * J31);
+                        const double A23 = (J21 * J13) - (J11 * J23);
+                        const double A31 = (J21 * J32) - (J31 * J22);
+                        const double A32 = (J31 * J12) - (J11 * J32);
+                        const double A33 = (J11 * J22) - (J12 * J21);
+
+                        Q(0,iqz,iqy,iqx) = w_detJ * (A11*A11 + A12*A12 + A13*A13); // 1,1
+                        Q(1,iqz,iqy,iqx) = w_detJ * (A11*A21 + A12*A22 + A13*A23); // 2,1
+                        Q(2,iqz,iqy,iqx) = w_detJ * (A11*A31 + A12*A32 + A13*A33); // 3,1
+                        Q(3,iqz,iqy,iqx) = w_detJ * (A21*A21 + A22*A22 + A23*A23); // 2,2
+                        Q(4,iqz,iqy,iqx) = w_detJ * (A21*A31 + A22*A32 + A23*A33); // 3,2
+                        Q(5,iqz,iqy,iqx) = w_detJ * (A31*A31 + A32*A32 + A33*A33); // 3,3
+                     }
+                  }
+               }
+
+               //MFEM_UNROLL(2)
+               for (int iqx=0; iqx<2; ++iqx)
+               {
+                  //MFEM_UNROLL(2)
+                  for (int jz=0; jz<2; ++jz)
+                  {
+                     // Note loop starts at iz=jz here, taking advantage of
+                     // symmetries.
+                     //MFEM_UNROLL(2)
+                     for (int iz=jz; iz<2; ++iz)
+                     {
+                        //MFEM_UNROLL(2)
+                        for (int iqy=0; iqy<2; ++iqy)
+                        {
+                           //MFEM_UNROLL(2)
+                           for (int iqz=0; iqz<2; ++iqz)
+                           {
+                              const double biz = (iz == iqz) ? 1.0 : 0.0;
+                              const double giz = (iz == 0) ? -1.0 : 1.0;
+
+                              const double bjz = (jz == iqz) ? 1.0 : 0.0;
+                              const double gjz = (jz == 0) ? -1.0 : 1.0;
+
+                              const double J11 = Q(0,iqz,iqy,iqx);
+                              const double J21 = Q(1,iqz,iqy,iqx);
+                              const double J31 = Q(2,iqz,iqy,iqx);
+                              const double J12 = J21;
+                              const double J22 = Q(3,iqz,iqy,iqx);
+                              const double J32 = Q(4,iqz,iqy,iqx);
+                              const double J13 = J31;
+                              const double J23 = J32;
+                              const double J33 = Q(5,iqz,iqy,iqx);
+
+                              grad_A(0,0,iqy,iz,jz,iqx) += J11*biz*bjz;
+                              grad_A(1,0,iqy,iz,jz,iqx) += J21*biz*bjz;
+                              grad_A(2,0,iqy,iz,jz,iqx) += J31*giz*bjz;
+                              grad_A(0,1,iqy,iz,jz,iqx) += J12*biz*bjz;
+                              grad_A(1,1,iqy,iz,jz,iqx) += J22*biz*bjz;
+                              grad_A(2,1,iqy,iz,jz,iqx) += J32*giz*bjz;
+                              grad_A(0,2,iqy,iz,jz,iqx) += J13*biz*gjz;
+                              grad_A(1,2,iqy,iz,jz,iqx) += J23*biz*gjz;
+                              grad_A(2,2,iqy,iz,jz,iqx) += J33*giz*gjz;
+                           }
+                           //MFEM_UNROLL(2)
+                           for (int jy=0; jy<2; ++jy)
+                           {
+                              //MFEM_UNROLL(2)
+                              for (int iy=0; iy<2; ++iy)
+                              {
+                                 const double biy = (iy == iqy) ? 1.0 : 0.0;
+                                 const double giy = (iy == 0) ? -1.0 : 1.0;
+
+                                 const double bjy = (jy == iqy) ? 1.0 : 0.0;
+                                 const double gjy = (jy == 0) ? -1.0 : 1.0;
+
+                                 grad_B(0,0,iy,jy,iz,jz,iqx) += biy*bjy*grad_A(0,0,iqy,iz,jz,iqx);
+                                 grad_B(1,0,iy,jy,iz,jz,iqx) += giy*bjy*grad_A(1,0,iqy,iz,jz,iqx);
+                                 grad_B(2,0,iy,jy,iz,jz,iqx) += biy*bjy*grad_A(2,0,iqy,iz,jz,iqx);
+                                 grad_B(0,1,iy,jy,iz,jz,iqx) += biy*gjy*grad_A(0,1,iqy,iz,jz,iqx);
+                                 grad_B(1,1,iy,jy,iz,jz,iqx) += giy*gjy*grad_A(1,1,iqy,iz,jz,iqx);
+                                 grad_B(2,1,iy,jy,iz,jz,iqx) += biy*gjy*grad_A(2,1,iqy,iz,jz,iqx);
+                                 grad_B(0,2,iy,jy,iz,jz,iqx) += biy*bjy*grad_A(0,2,iqy,iz,jz,iqx);
+                                 grad_B(1,2,iy,jy,iz,jz,iqx) += giy*bjy*grad_A(1,2,iqy,iz,jz,iqx);
+                                 grad_B(2,2,iy,jy,iz,jz,iqx) += biy*bjy*grad_A(2,2,iqy,iz,jz,iqx);
+                              }
+                           }
+                        }
+                        //MFEM_UNROLL(2)
+                        for (int jy=0; jy<2; ++jy)
+                        {
+                           //MFEM_UNROLL(2)
+                           for (int jx=0; jx<2; ++jx)
+                           {
+                              //MFEM_UNROLL(2)
+                              for (int iy=0; iy<2; ++iy)
+                              {
+                                 //MFEM_UNROLL(2)
+                                 for (int ix=0; ix<2; ++ix)
+                                 {
+                                    const double bix = (ix == iqx) ? 1.0 : 0.0;
+                                    const double gix = (ix == 0) ? -1.0 : 1.0;
+
+                                    const double bjx = (jx == iqx) ? 1.0 : 0.0;
+                                    const double gjx = (jx == 0) ? -1.0 : 1.0;
+
+                                    int ii_loc = ix + 2*iy + 4*iz;
+                                    int jj_loc = jx + 2*jy + 4*jz;
+
+                                    // Only store the lower-triangular part of
+                                    // the matrix (by symmetry).
+                                    if (jj_loc > ii_loc) { continue; }
+
+                                    double val = 0.0;
+                                    val += gix*gjx*grad_B(0,0,iy,jy,iz,jz,iqx);
+                                    val += bix*gjx*grad_B(1,0,iy,jy,iz,jz,iqx);
+                                    val += bix*gjx*grad_B(2,0,iy,jy,iz,jz,iqx);
+                                    val += gix*bjx*grad_B(0,1,iy,jy,iz,jz,iqx);
+                                    val += bix*bjx*grad_B(1,1,iy,jy,iz,jz,iqx);
+                                    val += bix*bjx*grad_B(2,1,iy,jy,iz,jz,iqx);
+                                    val += gix*bjx*grad_B(0,2,iy,jy,iz,jz,iqx);
+                                    val += bix*bjx*grad_B(2,2,iy,jy,iz,jz,iqx);
+                                    val += bix*bjx*grad_B(1,2,iy,jy,iz,jz,iqx);
+
+                                    local_mat(ii_loc, jj_loc) += val;
+                                 }
+                              }
+                           }
+                        }
+                     }
+                  }
+               }
+               // Assemble the local matrix into the macro-element sparse matrix
+               // in a format similar to coordinate format. The (I,J) arrays
+               // are implicit (not stored explicitly).
+               //MFEM_UNROLL(8)
+               for (int ii_loc=0; ii_loc<8; ++ii_loc)
+               {
+                  const int ix = ii_loc%2;
+                  const int iy = (ii_loc/2)%2;
+                  const int iz = ii_loc/2/2;
+
+                  for (int jj_loc=0; jj_loc<8; ++jj_loc)
+                  {
+                     const int jx = jj_loc%2;
+                     const int jy = (jj_loc/2)%2;
+                     const int jz = jj_loc/2/2;
+                     const int jj_off = (jx-ix+1) + 3*(jy-iy+1) + 9*(jz-iz+1);
+
+                     if (jj_loc <= ii_loc)
+                     {
+                        AtomicAdd(V(jj_off, ix+kx, iy+ky, iz+kz), local_mat(ii_loc, jj_loc));
+                     }
+                     else
+                     {
+                        AtomicAdd(V(jj_off, ix+kx, iy+ky, iz+kz), local_mat(jj_loc, ii_loc));
+                     }
+                  }
+               }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Place the macro-element sparse matrix into the global sparse matrix.
+      MFEM_FOREACH_THREAD(iz,z,nd1d)
+      {
+         MFEM_FOREACH_THREAD(iy,y,nd1d)
+         {
+            MFEM_FOREACH_THREAD(ix,x,nd1d)
+            {
+               double col_ptr[nnz_per_row]; // 27
+
+               const int ii_el = ix + nd1d*(iy + nd1d*iz);
+               const int ii = el_dof_lex(ii_el, iel_ho);
+
+               // Set column pointer to avoid searching in the row
+               for (int j = I[ii], end = I[ii+1]; j < end; j++)
+               {
+                  const int jj = J[j];
+                  int jj_el = -1;
+                  for (int k = K[jj], k_end = K[jj+1]; k < k_end; k += 2)
+                  {
+                     if (dof_glob2loc[k] == iel_ho)
+                     {
+                        jj_el = dof_glob2loc[k+1];
+                        break;
+                     }
+                  }
+                  if (jj_el < 0) { continue; }
+                  const int jx = jj_el%nd1d;
+                  const int jy = (jj_el/nd1d)%nd1d;
+                  const int jz = jj_el/nd1d/nd1d;
+                  const int jj_off = (jx-ix+1) + 3*(jy-iy+1) + 9*(jz-iz+1);
+                  col_ptr[jj_off] = j;
+               }
+
+               const int jx_begin = (ix > 0) ? ix - 1 : 0;
+               const int jx_end = (ix < order) ? ix + 1 : order;
+
+               const int jy_begin = (iy > 0) ? iy - 1 : 0;
+               const int jy_end = (iy < order) ? iy + 1 : order;
+
+               const int jz_begin = (iz > 0) ? iz - 1 : 0;
+               const int jz_end = (iz < order) ? iz + 1 : order;
+
+               for (int jz=jz_begin; jz<=jz_end; ++jz)
+               {
+                  for (int jy=jy_begin; jy<=jy_end; ++jy)
+                  {
+                     for (int jx=jx_begin; jx<=jx_end; ++jx)
+                     {
+                        const int jj_off = (jx-ix+1) + 3*(jy-iy+1) + 9*(jz-iz+1);
+                        const double Vji = V(jj_off, ix, iy, iz);
+                        const int col_ptr_jj = col_ptr[jj_off];
+                        if ((ix == 0 && jx == 0) || (ix == order && jx == order) ||
+                            (iy == 0 && jy == 0) || (iy == order && jy == order) ||
+                            (iz == 0 && jz == 0) || (iz == order && jz == order))
+                        {
+                           AtomicAdd(A[col_ptr_jj], Vji);
+                        }
+                        else
+                        {
+                           A[col_ptr_jj] += Vji;
+                        }
+                     }
+                  }
+               }
+            }
+         }
+      }
+   });
+}
+
+#define LOR_KERNEL_INSTANCE(order,use_smem) \
+template void Assemble3DBatchedLOR<order,use_smem>\
+    (const Array<int> &,const Array<int> &, const Array<int> &,\
+     Mesh &, SparseMatrix &)
+
+LOR_KERNEL_INSTANCE(1,true);
+LOR_KERNEL_INSTANCE(2,true);
+LOR_KERNEL_INSTANCE(3,true);
+LOR_KERNEL_INSTANCE(4,true);
+LOR_KERNEL_INSTANCE(5,true);
+LOR_KERNEL_INSTANCE(6,false);/*
+LOR_KERNEL_INSTANCE(7,false);
+LOR_KERNEL_INSTANCE(8,false);
+LOR_KERNEL_INSTANCE(9,false);
+LOR_KERNEL_INSTANCE(10,false);
+LOR_KERNEL_INSTANCE(11,false);
+LOR_KERNEL_INSTANCE(12,false);
+LOR_KERNEL_INSTANCE(13,false);
+LOR_KERNEL_INSTANCE(14,false);
+LOR_KERNEL_INSTANCE(15,false);
+LOR_KERNEL_INSTANCE(16,false);*/
+
+} // namespace mfem

fem/lor_assembly_ni3.cpp

@@ -0,0 +1,128 @@
+// Copyright (c) 2010-2021, 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 "fem.hpp"
+#include "../general/forall.hpp"
+
+#define MFEM_DEBUG_COLOR 227
+#include "../general/debug.hpp"
+
+#define MFEM_NVTX_COLOR MediumVioletRed
+#include "../general/nvtx.hpp"
+
+namespace mfem
+{
+
+#define M1D 8
+
+template<int D1D, int Q1D>
+void NodalInterpolation3D(const int NE,
+                          const Vector& localL, Vector& localH,
+                          const Array<double> &B)
+{
+   MFEM_NVTX;
+   dbg("D1D:%d Q1D:%d", D1D, Q1D);
+
+   static constexpr int VDIM = 3;
+
+   const auto x_ = Reshape(localL.Read(), D1D, D1D, D1D, VDIM, NE);
+   const auto B_ = Reshape(B.Read(), Q1D, D1D);
+
+   auto y_ = Reshape(localH.Write(), VDIM, Q1D, Q1D, Q1D, NE);
+
+   {
+      NVTX("localH = 0.0");
+      localH = 0.0;
+   }
+
+   {
+      NVTX("NodalInterpolation3D Kernel");
+      MFEM_FORALL(e, NE,
+      {
+         for (int vd = 0; vd < VDIM; ++vd)
+         {
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               double sol_xy[M1D][M1D];
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  for (int qx = 0; qx < Q1D; ++qx)
+                  {
+                     sol_xy[qy][qx] = 0.0;
+                  }
+               }
+               for (int dy = 0; dy < D1D; ++dy)
+               {
+                  double sol_x[M1D];
+                  for (int qx = 0; qx < Q1D; ++qx)
+                  {
+                     sol_x[qx] = 0;
+                  }
+                  for (int dx = 0; dx < D1D; ++dx)
+                  {
+                     const double s = x_(dx, dy, dz, vd, e);
+                     for (int qx = 0; qx < Q1D; ++qx)
+                     {
+                        sol_x[qx] += B_(qx, dx) * s;
+                     }
+                  }
+                  for (int qy = 0; qy < Q1D; ++qy)
+                  {
+                     const double wy = B_(qy, dy);
+                     for (int qx = 0; qx < Q1D; ++qx)
+                     {
+                        sol_xy[qy][qx] += wy * sol_x[qx];
+                     }
+                  }
+               }
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  const double wz = B_(qz, dz);
+                  for (int qy = 0; qy < Q1D; ++qy)
+                  {
+                     for (int qx = 0; qx < Q1D; ++qx)
+                     {
+                        y_(vd, qx, qy, qz, e) += wz * sol_xy[qy][qx];
+                     }
+                  }
+               }
+            }
+         }
+      });
+   }
+}
+
+#define NODAL_INTERP_3D_INSTANCE(D1D,Q1D) \
+template void NodalInterpolation3D<D1D,Q1D>\
+    (const int, const Vector&, Vector&,const Array<double>&)
+
+NODAL_INTERP_3D_INSTANCE(2,2);
+NODAL_INTERP_3D_INSTANCE(2,3);
+NODAL_INTERP_3D_INSTANCE(2,4);
+NODAL_INTERP_3D_INSTANCE(2,5);
+NODAL_INTERP_3D_INSTANCE(2,6);
+NODAL_INTERP_3D_INSTANCE(2,7);
+
+NODAL_INTERP_3D_INSTANCE(4,2);
+NODAL_INTERP_3D_INSTANCE(4,3);
+NODAL_INTERP_3D_INSTANCE(4,4);
+NODAL_INTERP_3D_INSTANCE(4,5);
+NODAL_INTERP_3D_INSTANCE(4,6);
+NODAL_INTERP_3D_INSTANCE(4,7);
+
+NODAL_INTERP_3D_INSTANCE(6,2);
+NODAL_INTERP_3D_INSTANCE(6,3);
+NODAL_INTERP_3D_INSTANCE(6,4);
+NODAL_INTERP_3D_INSTANCE(6,5);
+NODAL_INTERP_3D_INSTANCE(6,6);
+NODAL_INTERP_3D_INSTANCE(6,7);
+
+} // namespace mfem

fem/nonlininteg.hpp

@@ -21,6 +21,14 @@
 namespace mfem
 {
 
+// E or L I/O vectors
+enum class ActionType
+{
+   E2E,   ///< E => E kernels
+   L2E,   ///< L => E kernels
+   L2L    ///< L => L kernels
+};
+
 /** @brief This class is used to express the local action of a general nonlinear
     finite element operator. In addition it may provide the capability to
     assemble the local gradient operator and to compute the local energy. */
@@ -32,7 +40,9 @@ protected:
    // CEED extension
    ceed::Operator* ceedOp;
 
-   MemoryType pa_mt = MemoryType::DEFAULT;
+   MemoryType memory_type = MemoryType::DEFAULT;
+
+   ActionType action_type = ActionType::E2E;
 
    NonlinearFormIntegrator(const IntegrationRule *ir = NULL)
       : IntRule(ir), ceedOp(NULL) { }
@@ -46,8 +56,11 @@ public:
    void SetIntegrationRule(const IntegrationRule &ir) { SetIntRule(&ir); }
 
    /// Set the memory type used for GeometricFactors and other large allocations
-   /// in PA extensions.
-   void SetPAMemoryType(MemoryType mt) { pa_mt = mt; }
+   void SetMemoryType(MemoryType mt) { memory_type = mt; }
+   MFEM_DEPRECATED void SetPAMemoryType(MemoryType mt) { memory_type = mt; }
+
+   /// Set the kernel type used in PA extensions.
+   ActionType GetActionType() const { return action_type; }
 
    /// Get the integration rule of the integrator (possibly NULL).
    const IntegrationRule *GetIntegrationRule() const { return IntRule; }

fem/restriction.hpp

@@ -74,6 +74,23 @@ public:
    /** Fill the J and Data arrays of SparseMatrix corresponding to the sparsity
        pattern given by this ElementRestriction, and the values of ea_data. */
    void FillJAndData(const Vector &ea_data, SparseMatrix &mat) const;
+
+   /**
+    * @brief GatherMap
+    * @return the mapping from L dofs to E dofs.
+    */
+   const Array<int> &GatherMap() const { return gatherMap; }
+
+   /**
+    * @brief Indices
+    * @return the local to global indices.
+    */
+   const Array<int> &Indices() const { return indices; }
+   /**
+    * @brief Offsets
+    * @return the offsets for each global dof.
+    */
+   const Array<int> &Offsets() const { return offsets; }
 };
 
 /// Operator that converts L2 FiniteElementSpace L-vectors to E-vectors.

fem/tbilininteg.hpp

@@ -384,9 +384,9 @@ struct TDiffusionKernel<2,2,complex_t>
          const complex_t w_det_J = Q.get(q,i,k) / (J11 * J22 - J21 * J12);
          internal::MatrixOps<2,2>::Symm<Symm>::Set(
             A.layout.ind1(i), A,
-            + w_det_J * (J12*J12 + J22*J22), // (1,1)
+            w_det_J * (J12*J12 + J22*J22), // (1,1)
             - w_det_J * (J11*J12 + J21*J22), // (2,1)
-            + w_det_J * (J11*J11 + J21*J21)  // (2,2)
+            w_det_J * (J11*J11 + J21*J21)  // (2,2)
          );
       }
    }

fem/tintrules.hpp

@@ -19,6 +19,8 @@
 namespace mfem
 {
 
+static IntegrationRules LobattoIntRules(0, Quadrature1D::GaussLobatto);
+
 // Templated integration rules, cf. intrules.?pp
 
 template <Geometry::Type G, int Q, int Order, typename real_t>
@@ -52,7 +54,7 @@ public:
 
    static const IntegrationRule &GetIntRule()
    {
-      return IntRules.Get(geom, order);
+      return LobattoIntRules.Get(geom, order);
    }
 
    // Multi-component weight assignment. qpt_layout_t must be (qpts x n1 x ...)
@@ -201,10 +203,10 @@ public:
 
 template <int Dim, int Q, typename real_t>
 class GaussIntegrationRule
-   : public TProductIntegrationRule<Dim, Q, 2*Q-1, real_t>
+   : public TProductIntegrationRule<Dim, Q, 2*Q-3, real_t>
 {
 public:
-   typedef TProductIntegrationRule<Dim,Q,2*Q-1,real_t> base_class;
+   typedef TProductIntegrationRule<Dim,Q,2*Q-3,real_t> base_class;
 
    using base_class::geom;
    using base_class::order;
@@ -226,11 +228,11 @@ public:
 
    static const IntegrationRule &Get1DIntRule()
    {
-      return IntRules.Get(Geometry::SEGMENT, order);
+      return LobattoIntRules.Get(Geometry::SEGMENT, order);
    }
    static const IntegrationRule &GetIntRule()
    {
-      return IntRules.Get(geom, order);
+      return LobattoIntRules.Get(geom, order);
    }
 };
 
@@ -239,15 +241,15 @@ class TIntegrationRule;
 
 template <int Order, typename real_t>
 class TIntegrationRule<Geometry::SEGMENT, Order, real_t>
-   : public GaussIntegrationRule<1, Order/2+1, real_t> { };
+   : public GaussIntegrationRule<1, Order/2+2, real_t> { };
 
 template <int Order, typename real_t>
 class TIntegrationRule<Geometry::SQUARE, Order, real_t>
-   : public GaussIntegrationRule<2, Order/2+1, real_t> { };
+   : public GaussIntegrationRule<2, Order/2+2, real_t> { };
 
 template <int Order, typename real_t>
 class TIntegrationRule<Geometry::CUBE, Order, real_t>
-   : public GaussIntegrationRule<3, Order/2+1, real_t> { };
+   : public GaussIntegrationRule<3, Order/2+2, real_t> { };
 
 // Triangle integration rules (based on intrules.cpp)
 // These specializations define the number of quadrature points for each rule as

fem/tmop/tmop_pa.cpp

@@ -50,8 +50,8 @@ void TMOP_Integrator::AssembleGradPA(const Vector &xe,
 
 void TMOP_Integrator::AssemblePA_Limiting()
 {
-   const MemoryType mt = (pa_mt == MemoryType::DEFAULT) ?
-                         Device::GetDeviceMemoryType() : pa_mt;
+   const MemoryType mt = (memory_type == MemoryType::DEFAULT) ?
+                         Device::GetDeviceMemoryType() : memory_type;
    // Return immediately if limiting is not enabled
    if (lim_coeff == nullptr) { return; }
    MFEM_VERIFY(lim_nodes0, "internal error");
@@ -177,8 +177,8 @@ void TMOP_Integrator::ComputeAllElementTargets(const Vector &xe) const
 
 void TMOP_Integrator::AssemblePA(const FiniteElementSpace &fes)
 {
-   const MemoryType mt = (pa_mt == MemoryType::DEFAULT) ?
-                         Device::GetDeviceMemoryType() : pa_mt;
+   const MemoryType mt = (memory_type == MemoryType::DEFAULT) ?
+                         Device::GetDeviceMemoryType() : memory_type;
    PA.enabled = true;
    PA.fes = &fes;
    Mesh *mesh = fes.GetMesh();

fem/tmop_tools.cpp

@@ -240,7 +240,7 @@ ParAdvectorCGOper::ParAdvectorCGOper(const Vector &x_start,
    ConvectionIntegrator *Kinteg = new ConvectionIntegrator(u_coeff);
    if (al == AssemblyLevel::PARTIAL)
    {
-      Kinteg->SetPAMemoryType(mt);
+      Kinteg->SetMemoryType(mt);
    }
    K.AddDomainIntegrator(Kinteg);
    K.SetAssemblyLevel(al);
@@ -250,7 +250,7 @@ ParAdvectorCGOper::ParAdvectorCGOper(const Vector &x_start,
    MassIntegrator *Minteg = new MassIntegrator;
    if (al == AssemblyLevel::PARTIAL)
    {
-      Minteg->SetPAMemoryType(mt);
+      Minteg->SetMemoryType(mt);
    }
    M.AddDomainIntegrator(Minteg);
    M.SetAssemblyLevel(al);

fem/transfer.cpp

@@ -1192,6 +1192,114 @@ void Prolongation2D(const int NE, const int D1D, const int Q1D,
    });
 }
 
+template <int D1D, int Q1D, int NBZ = 1>
+void FastProlongation3D(const int lND,
+                        const int hNQ,
+                        const int NE,
+                        const int *lmap,
+                        const Vector& x,
+                        const int *hmap,
+                        Vector& y,
+                        const Array<double>& B,
+                        const Vector& mask)
+{
+   const auto B_ = Reshape(B.Read(), Q1D, D1D);
+   const auto lMAP = Reshape(lmap, D1D,D1D,D1D, NE);
+   const auto hMAP = Reshape(hmap, Q1D,Q1D,Q1D, NE);
+   const auto X = Reshape(x.Read(), lND);
+   const auto M = Reshape(mask.Read(), Q1D, Q1D, Q1D, NE);
+   auto Y = Reshape(y.ReadWrite(), hNQ);
+
+   y = 0.0;
+
+   assert(Q1D<=8);
+
+   MFEM_FORALL_3D(be, (NE+NBZ-1)/NBZ, Q1D, Q1D, NBZ,
+   {
+      double u[Q1D];
+      const int tz = MFEM_THREAD_ID(z);
+      const int e = be * MFEM_THREAD_SIZE(z) + tz;
+
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[NBZ][Q1D][Q1D][Q1D];
+
+      // Load input, B & X interpolation
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            if (tz == 0) { s_B[qx][dy] = B_(qx,dy); }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const double Bx = B_(qx,dx);
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz)
+               {
+                  const int gid = lMAP(dx,dy,dz,e);
+                  const int idx = gid >= 0 ? gid : -1 - gid;
+                  u[dz] += X(idx)* Bx;
+               }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[tz][dz][dy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Y interpolation
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { u[qy] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               const double zyX = s_q[tz][dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int qy = 0; qy < Q1D; ++qy) { u[qy] += zyX * s_B[qy][dy]; }
+            }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy) { s_q[tz][dz][qy][qx] = u[qy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Z interpolation & Q-function
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            // Z interpolation
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz) { u[qz] = 0.0; }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               const double zYX = s_q[tz][dz][qy][qx];
+               MFEM_UNROLL(Q1D)
+               for (int qz = 0; qz < Q1D; ++qz) { u[qz] += zYX * s_B[qz][dz]; }
+            }
+
+            // Q-function
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const int gid = hMAP(qx,qy,qz,e);
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               const double value = M(qx, qy, qz, e) * u[qz];
+               AtomicAdd(Y[idx], value);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
 void Prolongation3D(const int NE, const int D1D, const int Q1D,
                     const Vector& localL, Vector& localH,
                     const Array<double>& B, const Vector& mask)
@@ -1303,6 +1411,7 @@ void Restriction2D(const int NE, const int D1D, const int Q1D,
       }
    });
 }
+
 void Restriction3D(const int NE, const int D1D, const int Q1D,
                    const Vector& localH, Vector& localL,
                    const Array<double>& Bt, const Vector& mask)
@@ -1364,6 +1473,129 @@ void Restriction3D(const int NE, const int D1D, const int Q1D,
       }
    });
 }
+
+template<int D1D, int Q1D, int NBZ = 1>
+void FastRestriction3D(const int lND,
+                       const int hNQ,
+                       const int NE,
+                       const int *hmap,
+                       const Vector& x,
+                       const int *lmap,
+                       Vector& y,
+                       const Array<double>& Bt,
+                       const Vector& mask)
+{
+   const auto Bt_ = Reshape(Bt.Read(), D1D, Q1D);
+   const auto lMAP = Reshape(lmap, D1D,D1D,D1D, NE);
+   const auto hMAP = Reshape(hmap, Q1D,Q1D,Q1D, NE);
+   const auto M = Reshape(mask.Read(), Q1D, Q1D, Q1D, NE);
+   const auto X = Reshape(x.Read(), hNQ);
+   auto Y = Reshape(y.ReadWrite(), lND);
+
+   y = 0.0;
+
+   assert(Q1D<=8);
+
+   MFEM_FORALL_3D(be, (NE+NBZ-1)/NBZ, Q1D, Q1D, NBZ,
+   {
+      double u[Q1D];
+      const int tz = MFEM_THREAD_ID(z);
+      const int e = be * MFEM_THREAD_SIZE(z) + tz;
+
+      MFEM_SHARED double s_B[Q1D][D1D];
+      MFEM_SHARED double s_q[NBZ][Q1D][Q1D][Q1D];
+
+      // Load B
+      if (tz == 0)
+      {
+         MFEM_FOREACH_THREAD(dy,y,D1D)
+         {
+            MFEM_FOREACH_THREAD(qx,x,Q1D)
+            {
+               { s_B[qx][dy] = Bt_(dy,qx); }
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Q-function & Zt projection
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            // Q-function
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const int gid = hMAP(qx,qy,qz,e);
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               s_q[tz][qz][qy][qx] = M(qx, qy, qz, e) * X(idx);
+            }
+
+            // Zt projection
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { u[dz] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               const double ZYX = s_q[tz][qz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dz = 0; dz < D1D; ++dz) { u[dz] += ZYX * s_B[qz][dz]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dz = 0; dz < D1D; ++dz) { s_q[tz][dz][qy][qx] = u[dz]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Yt projection
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { u[dy] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               const double zYX = s_q[tz][dz][qy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dy = 0; dy < D1D; ++dy) { u[dy] += zYX * s_B[qy][dy]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dy = 0; dy < D1D; ++dy) { s_q[tz][dz][dy][qx] = u[dy]; }
+         }
+      }
+      MFEM_SYNC_THREAD;
+
+      // Xt projection & save output
+      MFEM_FOREACH_THREAD(dz,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dy,x,D1D)
+         {
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx) { u[dx] = 0.0; }
+            MFEM_UNROLL(Q1D)
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               const double zyX = s_q[tz][dz][dy][qx];
+               MFEM_UNROLL(D1D)
+               for (int dx = 0; dx < D1D; ++dx) { u[dx] += zyX * s_B[qx][dx]; }
+            }
+            MFEM_UNROLL(D1D)
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               const int gid = lMAP(dx,dy,dz,e);
+               const int idx = gid >= 0 ? gid : -1 - gid;
+               const double value = u[dx];
+               AtomicAdd(Y[idx], value);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   });
+}
+
 } // namespace TransferKernels
 
 
@@ -1380,22 +1612,69 @@ void TensorProductPRefinementTransferOperator::Mult(const Vector& x,
       return;
    }
 
-   elem_restrict_lex_l->Mult(x, localL);
-   if (dim == 2)
+   if (Device::FastKernelsEnabled())
    {
-      TransferKernels::Prolongation2D(NE, D1D, Q1D, localL, localH, B, mask);
-   }
-   else if (dim == 3)
-   {
-      TransferKernels::Prolongation3D(NE, D1D, Q1D, localL, localH, B, mask);
+      assert(dim==3);
+      constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+      // lFESpace map idx
+      const int lND = lFESpace.GetNDofs();
+      const Operator *lERop = lFESpace.GetElementRestriction(ordering);
+      const ElementRestriction* lER = dynamic_cast<const ElementRestriction*>(lERop);
+      assert(lER);
+      const int *lmap = lER->GatherMap().Read();
+      assert(lER->GatherMap().Size() == D1D*D1D*D1D*NE);
+      // hFESpace map idx
+      const int hNQ = hFESpace.GetNDofs();
+      const Operator *hERop = elem_restrict_lex_h;
+      const ElementRestriction* hER = dynamic_cast<const ElementRestriction*>(hERop);
+      assert(hER);
+      const int *hmap = hER->GatherMap().Read();
+      assert(hER->GatherMap().Size() == Q1D*Q1D*Q1D*NE);
+
+      void (*Ker)(const int lND,
+                  const int hNQ,
+                  const int NE,
+                  const int *lmap,
+                  const Vector& x,
+                  const int *hmap,
+                  Vector& y,
+                  const Array<double>& B,
+                  const Vector& mask) = nullptr;
+
+      const int id = (D1D << 4) | Q1D;
+      switch (id) // orders 1~8
+      {
+         case 0x23: Ker=TransferKernels::FastProlongation3D<2,3>; break; // 1
+         case 0x34: Ker=TransferKernels::FastProlongation3D<3,4>; break; // 2
+         case 0x35: Ker=TransferKernels::FastProlongation3D<3,5>; break;
+         case 0x45: Ker=TransferKernels::FastProlongation3D<4,5>; break; // 3
+         case 0x56: Ker=TransferKernels::FastProlongation3D<5,6>; break; // 4
+         case 0x57: Ker=TransferKernels::FastProlongation3D<5,7>; break;
+         case 0x67: Ker=TransferKernels::FastProlongation3D<6,7>; break; // 5
+         case 0x78: Ker=TransferKernels::FastProlongation3D<7,8>; break; // 6
+         default: MFEM_ABORT("Unknown kernel 0x" << std::hex << id << std::dec);
+      }
+      Ker(lND, hNQ, NE, lmap, x, hmap, y, B, mask);
    }
    else
    {
-      MFEM_ABORT("TensorProductPRefinementTransferOperator::Mult not "
-                 "implemented for dim = "
-                 << dim);
+      elem_restrict_lex_l->Mult(x, localL);
+      if (dim == 2)
+      {
+         TransferKernels::Prolongation2D(NE, D1D, Q1D, localL, localH, B, mask);
+      }
+      else if (dim == 3)
+      {
+         TransferKernels::Prolongation3D(NE, D1D, Q1D, localL, localH, B, mask);
+      }
+      else
+      {
+         MFEM_ABORT("TensorProductPRefinementTransferOperator::Mult not "
+                    "implemented for dim = "
+                    << dim);
+      }
+      elem_restrict_lex_h->MultTranspose(localH, y);
    }
-   elem_restrict_lex_h->MultTranspose(localH, y);
 }
 
 void TensorProductPRefinementTransferOperator::MultTranspose(const Vector& x,
@@ -1406,22 +1685,68 @@ void TensorProductPRefinementTransferOperator::MultTranspose(const Vector& x,
       return;
    }
 
-   elem_restrict_lex_h->Mult(x, localH);
-   if (dim == 2)
+   if (Device::FastKernelsEnabled())
    {
-      TransferKernels::Restriction2D(NE, D1D, Q1D, localH, localL, Bt, mask);
-   }
-   else if (dim == 3)
-   {
-      TransferKernels::Restriction3D(NE, D1D, Q1D, localH, localL, Bt, mask);
+      assert(dim==3);
+      // lFESpace map idx
+      const int lND = lFESpace.GetNDofs();
+      constexpr ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
+      const Operator *lERop = lFESpace.GetElementRestriction(ordering);
+      const ElementRestriction* lER = dynamic_cast<const ElementRestriction*>(lERop);
+      assert(lER);
+      const int *lmap = lER->GatherMap().Read();
+      // hFESpace map idx
+      const int hNQ = hFESpace.GetNDofs();
+      const Operator *hERop = elem_restrict_lex_h;
+      const ElementRestriction* hER = dynamic_cast<const ElementRestriction*>(hERop);
+      assert(hER);
+      const int *hmap = hER->GatherMap().Read();
+      assert(hER->GatherMap().Size() == Q1D*Q1D*Q1D*NE);
+
+      void (*Ker)(const int lND,
+                  const int hNQ,
+                  const int NE,
+                  const int *hmap,
+                  const Vector& x,
+                  const int *lmap,
+                  Vector& y,
+                  const Array<double>& Bt,
+                  const Vector& mask) = nullptr;
+
+      const int id = (D1D << 4) | Q1D;
+      switch (id) // orders 1~8
+      {
+         case 0x23: Ker=TransferKernels::FastRestriction3D<2,3>; break; // 1
+         case 0x34: Ker=TransferKernels::FastRestriction3D<3,4>; break; // 2
+         case 0x35: Ker=TransferKernels::FastRestriction3D<3,5>; break;
+         case 0x45: Ker=TransferKernels::FastRestriction3D<4,5>; break; // 3
+         case 0x56: Ker=TransferKernels::FastRestriction3D<5,6>; break; // 4
+         case 0x57: Ker=TransferKernels::FastRestriction3D<5,7>; break;
+         case 0x67: Ker=TransferKernels::FastRestriction3D<6,7>; break; // 5
+         case 0x78: Ker=TransferKernels::FastRestriction3D<7,8>; break; // 6
+         default: MFEM_ABORT("Unknown kernel 0x" << std::hex << id << std::dec);
+      }
+      Ker(lND, hNQ, NE, hmap, x, lmap, y, Bt, mask);
    }
    else
    {
-      MFEM_ABORT("TensorProductPRefinementTransferOperator::MultTranspose not "
-                 "implemented for dim = "
-                 << dim);
+      elem_restrict_lex_h->Mult(x, localH);
+      if (dim == 2)
+      {
+         TransferKernels::Restriction2D(NE, D1D, Q1D, localH, localL, Bt, mask);
+      }
+      else if (dim == 3)
+      {
+         TransferKernels::Restriction3D(NE, D1D, Q1D, localH, localL, Bt, mask);
+      }
+      else
+      {
+         MFEM_ABORT("TensorProductPRefinementTransferOperator::MultTranspose not "
+                    "implemented for dim = "
+                    << dim);
+      }
+      elem_restrict_lex_l->MultTranspose(localL, y);
    }
-   elem_restrict_lex_l->MultTranspose(localL, y);
 }
 
 #ifdef MFEM_USE_MPI

general/backends.hpp

@@ -18,6 +18,7 @@
 #include <cusparse.h>
 #include <library_types.h>
 #include <cuda_runtime.h>
+#include <curand.h>
 #include <cuda.h>
 #endif
 #include "cuda.hpp"
@@ -39,20 +40,26 @@
 #endif
 
 #if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
+#define MFEM_GLOBAL
 #define MFEM_DEVICE
 #define MFEM_LAMBDA
+#define MFEM_GLOBAL
 #define MFEM_HOST_DEVICE
 // MFEM_DEVICE_SYNC is made available for debugging purposes
 #define MFEM_DEVICE_SYNC
 // MFEM_STREAM_SYNC is used for UVM and MPI GPU-Aware kernels
 #define MFEM_STREAM_SYNC
+#define MFEM_LAUNCH_BOUNDS(...)
+#define MFEM_LAUNCH_KERNEL(Kernel,Grid,Block,Smem,...) Kernel(__VA_ARGS__)
 #endif
 
 #if !((defined(MFEM_USE_CUDA) && defined(__CUDA_ARCH__)) || \
       (defined(MFEM_USE_HIP)  && defined(__HIP_DEVICE_COMPILE__)))
 #define MFEM_SHARED
 #define MFEM_SYNC_THREAD
+#define MFEM_GRID_DIM(k) 1
 #define MFEM_BLOCK_ID(k) 0
+#define MFEM_BLOCK_DIM(k) 1
 #define MFEM_THREAD_ID(k) 0
 #define MFEM_THREAD_SIZE(k) 1
 #define MFEM_FOREACH_THREAD(i,k,N) for(int i=0; i<N; i++)
@@ -91,4 +98,12 @@ MFEM_HOST_DEVICE T AtomicAdd(T &add, const T val)
 #endif
 }
 
+/// Helper function to return and increment a given pointer with a given size
+template<typename T> MFEM_HOST_DEVICE static
+inline T *DeviceMemAlloc(T* &mem, size_t size) noexcept
+{
+   T* base = mem;
+   return (mem += size, base);
+}
+
 #endif // MFEM_BACKENDS_HPP

general/cuda.hpp

@@ -19,9 +19,13 @@
 #define MFEM_CUDA_BLOCKS 256
 
 #ifdef MFEM_USE_CUDA
+#define MFEM_GLOBAL __global__
 #define MFEM_DEVICE __device__
 #define MFEM_LAMBDA __host__
+#define MFEM_GLOBAL __global__
 #define MFEM_HOST_DEVICE __host__ __device__
+#define MFEM_LAUNCH_BOUNDS(NTH,NBK) __launch_bounds__(NTH,NBK)
+#define MFEM_DEVICE_CHECK MFEM_GPU_CHECK((cudaDeviceSynchronize(),cudaGetLastError()))
 #define MFEM_DEVICE_SYNC MFEM_GPU_CHECK(cudaDeviceSynchronize())
 #define MFEM_STREAM_SYNC MFEM_GPU_CHECK(cudaStreamSynchronize(0))
 // Define a CUDA error check macro, MFEM_GPU_CHECK(x), where x returns/is of
@@ -37,12 +41,15 @@
       } \
    } \
    while (0)
+#define MFEM_LAUNCH_KERNEL(Kernel,Grid,Block,Smem,...) \
+    Kernel<<<Grid,Block,sizeof(double)*(Smem),0>>>(__VA_ARGS__)
 #endif // MFEM_USE_CUDA
 
 // Define the MFEM inner threading macros
 #if defined(MFEM_USE_CUDA) && defined(__CUDA_ARCH__)
 #define MFEM_SHARED __shared__
 #define MFEM_SYNC_THREAD __syncthreads()
+#define MFEM_GRID_DIM(k) gridDim.k
 #define MFEM_BLOCK_ID(k) blockIdx.k
 #define MFEM_THREAD_ID(k) threadIdx.k
 #define MFEM_THREAD_SIZE(k) blockDim.k

general/debug.hpp

@@ -0,0 +1,178 @@
+// Copyright (c) 2010-2021, 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_DEBUG_HPP
+#define MFEM_DEBUG_HPP
+
+#include <string>
+#include <cstring>
+#include <iomanip>
+#include <iostream>
+
+#include "globals.hpp"
+
+#ifdef MFEM_USE_MPI
+#include <mpi.h>
+#endif
+
+namespace mfem
+{
+
+class Debug
+{
+   const bool debug = false;
+public:
+   inline Debug() {}
+
+   inline Debug(const int mpi_rank,
+                const char *FILE, const int LINE,
+                const char *FUNC, int COLOR): debug(true)
+   {
+      if (!debug) { return; }
+      const char *base = Strrnchr(FILE,'/', 2);
+      const char *file = base ? base + 1 : FILE;
+      const uint8_t color = COLOR ? COLOR : 20 + Checksum8(FILE) % 210;
+      mfem::out << "\033[38;5;" << std::to_string(color) << "m";
+      mfem::out << mpi_rank << std::setw(30) << file << ":";
+      mfem::out << "\033[2m" << std::setw(4) << LINE << "\033[22m: ";
+      if (FUNC) { mfem::out << "[" << FUNC << "] "; }
+      mfem::out << "\033[1m";
+   }
+
+   ~Debug()
+   {
+      if (!debug) { return; }
+      mfem::out << "\033[m";
+      mfem::out << std::endl;
+   }
+
+   template <typename T>
+   inline void operator<<(const T &arg) const noexcept { mfem::out << arg; }
+
+   template<typename T, typename... Args>
+   inline void operator()(const char *fmt, const T &arg,
+                          Args... args) const noexcept
+   {
+      if (!debug) { return; }
+      for (; *fmt != '\0'; fmt++ )
+      {
+         if (*fmt == '%')
+         {
+            fmt++;
+            const char c = *fmt;
+            if (c == 'p') { operator<<(arg); }
+            if (c == 's' || c == 'd' || c == 'f') { operator<<(arg); }
+            if (c == 'x' || c == 'X')
+            {
+               mfem::out << std::hex;
+               if (c == 'X') { mfem::out << std::uppercase; }
+               operator<<(arg);
+               mfem::out << std::nouppercase << std::dec;
+            }
+            if (c == '.')
+            {
+               fmt++;
+               const char c = *fmt;
+               char num[8] = { 0 };
+               for (int k = 0; *fmt != '\0'; fmt++, k++)
+               {
+                  if (*fmt == 'e' || *fmt == 'f') { break; }
+                  if (*fmt < 0x30 || *fmt > 0x39) { break; }
+                  num[k] = *fmt;
+               }
+               const int fx = std::atoi(num);
+               if (c == 'e') { mfem::out << std::scientific; }
+               if (c == 'f') { mfem::out << std::fixed; }
+               mfem::out << std::setprecision(fx);
+               operator<<(arg);
+               mfem::out << std::setprecision(6);
+            }
+            return operator()(fmt + 1, args...);
+         }
+         operator<<(*fmt);
+      }
+   }
+
+   template<typename T>
+   inline void operator()(const T &arg) const noexcept
+   {
+      if (!debug) { return; }
+      operator<<(arg);
+   }
+
+   inline void operator()() const noexcept { }
+
+public:
+   static const Debug Set(const char *FILE, const int LINE, const char *FUNC,
+                          int COLOR = 0)
+   {
+      static int mpi_dbg = 0, mpi_rank = 0;
+      static bool env_mpi = false, env_dbg = false;
+      static bool ini_dbg = false;
+      if (!ini_dbg)
+      {
+         const char *DBG = getenv("MFEM_DEBUG");
+         const char *MPI = getenv("MFEM_DEBUG_MPI");
+         env_dbg = DBG != nullptr;
+         env_mpi = MPI != nullptr;
+#ifdef MFEM_USE_MPI
+         int mpi_ini = false;
+         MPI_Initialized(&mpi_ini);
+         if (mpi_ini) { MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank); }
+         mpi_dbg = atoi(env_mpi ? MPI : "0");
+#endif
+         ini_dbg = true;
+      }
+      const bool debug = (env_dbg && (!env_mpi || mpi_rank == mpi_dbg));
+      return debug ? Debug(mpi_rank, FILE, LINE, FUNC, COLOR) : Debug();
+   }
+
+private:
+   inline uint8_t Checksum8(const char *bfr)
+   {
+      unsigned int chk = 0;
+      size_t len = strlen(bfr);
+      for (; len; len--,bfr++) { chk += static_cast<unsigned int>(*bfr); }
+      return (uint8_t) chk;
+   }
+
+   inline const char *Strrnchr(const char *s, const unsigned char c, int n)
+   {
+      size_t len = strlen(s);
+      char *p = const_cast<char*>(s) + len - 1;
+      for (; n; n--,p--,len--)
+      {
+         for (; len; p--,len--)
+            if (*p == c) { break; }
+         if (!len) { return nullptr; }
+         if (n == 1) { return p; }
+      }
+      return nullptr;
+   }
+
+};
+
+#ifndef MFEM_DEBUG_COLOR
+#define MFEM_DEBUG_COLOR 0
+#endif
+
+#define dbg(...) \
+    mfem::Debug::Set(__FILE__,__LINE__,__FUNCTION__,MFEM_DEBUG_COLOR).\
+    operator()(__VA_ARGS__)
+
+} // mfem namespace
+
+#define DBG(...) { printf("\033[33m");  \
+                   printf(__VA_ARGS__); \
+                   printf(" \n\033[m"); \
+                   fflush(0); }
+
+#endif // MFEM_DEBUG_HPP

general/device.cpp

@@ -363,6 +363,19 @@ void Device::UpdateMemoryTypeAndClass()
       device_mem_type = MemoryType::MANAGED;
    }
 
+   // Enable the non deterministic kernels shortcut when requested
+   if (device_option && !strncmp(device_option, "fast", 4))
+   {
+      deterministic_kernels = DETERMINISTIC_KERNELS::OFF;
+      kernels_version = atoi(device_option + 4);
+   }
+
+   // Enable the version shortcut when requested
+   if (device_option && !strncmp(device_option, "ver", 3))
+   {
+      kernels_version = atoi(device_option + 3);
+   }
+
    // Enable the DEBUG mode when requested
    if (debug)
    {

general/device.hpp

@@ -123,11 +123,14 @@ class Device
 private:
    friend class MemoryManager;
    enum MODES {SEQUENTIAL, ACCELERATED};
+   enum DETERMINISTIC_KERNELS {ON, OFF};
 
    static bool device_env, mem_host_env, mem_device_env, mem_types_set;
    static Device device_singleton;
 
    MODES mode = Device::SEQUENTIAL;
+   /// Device kernels are set to be deterministic by default.
+   DETERMINISTIC_KERNELS deterministic_kernels = DETERMINISTIC_KERNELS::ON;
    int dev = 0;   ///< Device ID of the configured device.
    int ngpu = -1; ///< Number of detected devices; -1: not initialized.
    /// Bitwise-OR of all configured backends.
@@ -145,6 +148,8 @@ private:
    MemoryClass device_mem_class = MemoryClass::HOST;
 
    char *device_option = NULL;
+   /// Kernels version number set from the device shortcut
+   int kernels_version = 0;
    Device(Device const&);
    void operator=(Device const&);
    static Device& Get() { return device_singleton; }
@@ -248,6 +253,18 @@ public:
    /// The opposite of IsEnabled().
    static inline bool IsDisabled() { return !IsEnabled(); }
 
+   /// Return true if the fast kernels are allowed.
+   static inline bool FastKernelsEnabled()
+   {
+      return Get().deterministic_kernels == OFF;
+   }
+
+   /// Return the kernel version to use, set from the device shortcut.
+   static inline int KernelsVersion()
+   {
+      return Get().kernels_version;
+   }
+
    /// Get the device id of the configured device.
    static inline int GetId() { return Get().dev; }
 

general/forall.hpp

@@ -32,12 +32,31 @@ const int MAX_D1D = 14;
 const int MAX_Q1D = 14;
 #endif
 
+/** @brief Create a scratch memory on the device. */
+template<int GRID, typename T = double>
+static T *ScratchMem(const int sm_size)
+{
+   if (GRID==0) { return nullptr; }
+   static Memory<T> data;
+   if (sm_size*GRID > data.Capacity())
+   {
+      data.Delete();
+      data.New(sm_size*GRID, Device::GetDeviceMemoryType());
+      data.UseDevice(true);
+   }
+   return data.Write(Device::GetDeviceMemoryClass(), data.Capacity());
+}
+
 // MFEM pragma macros that can be used inside MFEM_FORALL macros.
 #define MFEM_PRAGMA(X) _Pragma(#X)
 
 // MFEM_UNROLL pragma macro that can be used inside MFEM_FORALL macros.
 #if defined(MFEM_USE_CUDA)
 #define MFEM_UNROLL(N) MFEM_PRAGMA(unroll(N))
+#elif defined(__clang__)
+#define MFEM_UNROLL(N) MFEM_PRAGMA(unroll N)
+#elif __GNUC__ >= 8
+#define MFEM_UNROLL(N) MFEM_PRAGMA(GCC unroll N)
 #else
 #define MFEM_UNROLL(N)
 #endif

general/hip.hpp

@@ -19,6 +19,7 @@
 #define MFEM_HIP_BLOCKS 256
 
 #ifdef MFEM_USE_HIP
+#define MFEM_GLOBAL __global__
 #define MFEM_DEVICE __device__
 #define MFEM_LAMBDA __host__ __device__
 #define MFEM_HOST_DEVICE __host__ __device__
@@ -37,6 +38,8 @@
       } \
    } \
    while (0)
+#define MFEM_LAUNCH_KERNEL(Kernel,Grid,Block,Smem,...) \
+    hipLaunchKernelGGL(Kernel,Grid,Block,sizeof(double)*(Smem),0,__VA_ARGS__)
 #endif // MFEM_USE_HIP
 
 // Define the MFEM inner threading macros
@@ -44,6 +47,8 @@
 #define MFEM_SHARED __shared__
 #define MFEM_SYNC_THREAD __syncthreads()
 #define MFEM_BLOCK_ID(k) hipBlockIdx_ ##k
+#define MFEM_BLOCK_DIM(k) hipBlockDim_ ##k
+#define MFEM_GRID_DIM(k) hipGridDim_ ##k
 #define MFEM_THREAD_ID(k) hipThreadIdx_ ##k
 #define MFEM_THREAD_SIZE(k) hipBlockDim_ ##k
 #define MFEM_FOREACH_THREAD(i,k,N) \

general/nvtx.hpp

@@ -0,0 +1,314 @@
+// Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
+// reserved. See files LICENSE and NOTICE for details.
+//
+// This file is part of CEED, a collection of benchmarks, miniapps, software
+// libraries and APIs for efficient high-order finite element and spectral
+// element discretizations for exascale applications. For more information and
+// source code availability see http://github.com/ceed.
+//
+// The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
+// a collaborative effort of two U.S. Department of Energy organizations (Office
+// of Science and the National Nuclear Security Administration) responsible for
+// the planning and preparation of a capable exascale ecosystem, including
+// software, applications, hardware, advanced system engineering and early
+// testbed platforms, in support of the nation's exascale computing imperative.
+
+#ifndef MFEM_NVVP_HPP
+#define MFEM_NVVP_HPP
+
+#include <string>
+#include <cstring>
+#include <cassert>
+
+#include "globals.hpp"
+#include "backends.hpp"
+
+#if defined(MFEM_USE_CUDA) //||1
+#include <cuda.h>
+#include <nvToolsExt.h>
+#include <cudaProfiler.h>
+#include <cuda_runtime.h>
+#else // MFEM_USE_CUDA
+typedef struct
+{
+   uint16_t version;
+   uint16_t size;
+   uint32_t category;
+   int32_t colorType;
+   uint32_t color;
+   int32_t payloadType;
+   int32_t reserved0;
+   int32_t messageType;
+   struct { const char* ascii; } message;
+} nvtxEventAttributes_t;
+#define NVTX_VERSION 2
+#define NVTX_COLOR_ARGB 1
+#define NVTX_MESSAGE_TYPE_ASCII 1
+#define NVTX_EVENT_ATTRIB_STRUCT_SIZE \
+    ((uint16_t)(sizeof(nvtxEventAttributes_t)))
+inline int nvtxRangePushEx(const nvtxEventAttributes_t*) { return 0; }
+inline int nvtxRangePop(void) { return 0; }
+#endif // MFEM_USE_CUDA
+
+namespace mfem
+{
+
+// en.wikipedia.org/wiki/Web_colors#Hex_triplet
+typedef enum
+{
+   Black, NavyBlue, DarkBlue, MediumBlue, Blue, DarkGreen, WebGreen, Teal,
+   DarkCyan, DeepSkyBlue, DarkTurquoise, MediumSpringGreen, Green, Lime,
+   SpringGreen, Aqua, Cyan, MidnightBlue, DodgerBlue, LightSeaGreen,
+   ForestGreen, SeaGreen, DarkSlateGray, LimeGreen, MediumSeaGreen,
+   Turquoise, RoyalBlue, SteelBlue, DarkSlateBlue, MediumTurquoise, Indigo,
+   DarkOliveGreen, CadetBlue, Cornflower, RebeccaPurple, MediumAquamarine,
+   DimGray, SlateBlue, OliveDrab, SlateGray, LightSlateGray,
+   MediumSlateBlue, LawnGreen, WebMaroon, WebPurple, Chartreuse,
+   Aquamarine, Olive, WebGray, SkyBlue, LightSkyBlue, BlueViolet, DarkRed,
+   DarkMagenta, SaddleBrown, DarkSeaGreen, LightGreen, MediumPurple,
+   DarkViolet, PaleGreen, DarkOrchid, YellowGreen, Purple, Sienna, Brown,
+   DarkGray, LightBlue, GreenYellow, PaleTurquoise, Maroon,
+   LightSteelBlue, PowderBlue, Firebrick, DarkGoldenrod, MediumOrchid,
+   RosyBrown, DarkKhaki, Gray, Silver, MediumVioletRed, IndianRed, Peru,
+   Chocolate, Tan, LightGray, Thistle, Orchid, Goldenrod, PaleVioletRed,
+   Crimson, Gainsboro, Plum, Burlywood, LightCyan, Lavender, DarkSalmon,
+   Violet, PaleGoldenrod, LightCoral, Khaki, AliceBlue, Honeydew, Azure,
+   SandyBrown, Wheat, Beige, WhiteSmoke, MintCream, GhostWhite, Salmon,
+   AntiqueWhite, Linen, LightGoldenrod, OldLace, Red, Fuchsia, Magenta,
+   DeepPink, OrangeRed, Tomato, HotPink, Coral, DarkOrange, LightSalmon,
+   Orange, LightPink, Pink, Gold, PeachPuff, NavajoWhite, Moccasin,
+   Bisque, MistyRose, BlanchedAlmond, PapayaWhip, LavenderBlush, Seashell,
+   Cornsilk, LemonChiffon, FloralWhite, Snow, Yellow, LightYellow, Ivory,
+   White
+} COLOR_NAMES;
+
+static constexpr uint32_t HEX_COLORS[] =
+{
+   0x000000, 0x000080, 0x00008B, 0x0000CD, 0x0000FF, 0x006400, 0x008000,
+   0x008080, 0x008B8B, 0x00BFFF, 0x00CED1, 0x00FA9A, 0x00FF00, 0x00FF00,
+   0x00FF7F, 0x00FFFF, 0x00FFFF, 0x191970, 0x1E90FF, 0x20B2AA, 0x228B22,
+   0x2E8B57, 0x2F4F4F, 0x32CD32, 0x3CB371, 0x40E0D0, 0x4169E1, 0x4682B4,
+   0x483D8B, 0x48D1CC, 0x4B0082, 0x556B2F, 0x5F9EA0, 0x6495ED, 0x663399,
+   0x66CDAA, 0x696969, 0x6A5ACD, 0x6B8E23, 0x708090, 0x778899, 0x7B68EE,
+   0x7CFC00, 0x7F0000, 0x7F007F, 0x7FFF00, 0x7FFFD4, 0x808000, 0x808080,
+   0x87CEEB, 0x87CEFA, 0x8A2BE2, 0x8B0000, 0x8B008B, 0x8B4513, 0x8FBC8F,
+   0x90EE90, 0x9370DB, 0x9400D3, 0x98FB98, 0x9932CC, 0x9ACD32, 0xA020F0,
+   0xA0522D, 0xA52A2A, 0xA9A9A9, 0xADD8E6, 0xADFF2F, 0xAFEEEE, 0xB03060,
+   0xB0C4DE, 0xB0E0E6, 0xB22222, 0xB8860B, 0xBA55D3, 0xBC8F8F, 0xBDB76B,
+   0xBEBEBE, 0xC0C0C0, 0xC71585, 0xCD5C5C, 0xCD853F, 0xD2691E, 0xD2B48C,
+   0xD3D3D3, 0xD8BFD8, 0xDA70D6, 0xDAA520, 0xDB7093, 0xDC143C, 0xDCDCDC,
+   0xDDA0DD, 0xDEB887, 0xE0FFFF, 0xE6E6FA, 0xE9967A, 0xEE82EE, 0xEEE8AA,
+   0xF08080, 0xF0E68C, 0xF0F8FF, 0xF0FFF0, 0xF0FFFF, 0xF4A460, 0xF5DEB3,
+   0xF5F5DC, 0xF5F5F5, 0xF5FFFA, 0xF8F8FF, 0xFA8072, 0xFAEBD7, 0xFAF0E6,
+   0xFAFAD2, 0xFDF5E6, 0xFF0000, 0xFF00FF, 0xFF00FF, 0xFF1493, 0xFF4500,
+   0xFF6347, 0xFF69B4, 0xFF7F50, 0xFF8C00, 0xFFA07A, 0xFFA500, 0xFFB6C1,
+   0xFFC0CB, 0xFFD700, 0xFFDAB9, 0xFFDEAD, 0xFFE4B5, 0xFFE4C4, 0xFFE4E1,
+   0xFFEBCD, 0xFFEFD5, 0xFFF0F5, 0xFFF5EE, 0xFFF8DC, 0xFFFACD, 0xFFFAF0,
+   0xFFFAFA, 0xFFFF00, 0xFFFFE0, 0xFFFFF0, 0xFFFFFF
+};
+
+static const int NUM_HEX_COLORS = sizeof(HEX_COLORS)/sizeof(uint32_t);
+
+class Nvtx
+{
+   const bool nvtx = false;
+   const bool enforce_kernel_sync = false;
+   const char *base, *file;
+   const uint32_t color = Black;
+   mutable std::string ascii;
+   mutable nvtxEventAttributes_t event;
+
+public:
+   Nvtx() { }
+
+   Nvtx(bool enforce_kernel_sync,
+        const char *FILE, const int LINE, const char *FUNC, uint32_t COLOR):
+      nvtx(true),
+      enforce_kernel_sync(enforce_kernel_sync),
+      base(Strrnchr(FILE,'/', 2)),
+      file(base ? base + 1 : FILE),
+      color(COLOR),
+      ascii(file),
+      event({})
+   {
+      event.version = NVTX_VERSION;
+      event.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+      event.colorType = NVTX_COLOR_ARGB;
+      event.color = HEX_COLORS[color % NUM_HEX_COLORS];
+      event.messageType = NVTX_MESSAGE_TYPE_ASCII;
+
+      ascii += ":";
+      ascii += std::to_string(LINE);
+      ascii += ":[";
+      ascii += FUNC;
+      ascii += "] ";
+   }
+
+   Nvtx(const char *title, int color = Wheat, bool enforce_kernel_sync = true):
+      nvtx(true),
+      enforce_kernel_sync(enforce_kernel_sync),
+      color(color),
+      ascii(title),
+      event({})
+   {
+      event.version = NVTX_VERSION;
+      event.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+      event.colorType = NVTX_COLOR_ARGB;
+      event.color = HEX_COLORS[color % NUM_HEX_COLORS];
+      event.messageType = NVTX_MESSAGE_TYPE_ASCII;
+      event.message.ascii = ascii.c_str();
+      nvtxRangePushEx(&event); // push
+   }
+
+   ~Nvtx()
+   {
+      if (!nvtx) { return; }
+      if (enforce_kernel_sync)
+      {
+         nvtxEventAttributes_t eks = {};
+         eks.version = NVTX_VERSION;
+         eks.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+         eks.category = 0; // user value
+         eks.colorType = NVTX_COLOR_ARGB;
+         eks.messageType = NVTX_MESSAGE_TYPE_ASCII;
+         eks.message.ascii = "Sync";
+         eks.color = HEX_COLORS[Yellow];
+         nvtxRangePushEx(&eks);
+         MFEM_STREAM_SYNC;
+         nvtxRangePop();
+      }
+      nvtxRangePop(); // pop
+   }
+
+   // used through MFEM_NVTX
+   inline void operator()() const noexcept
+   {
+      if (!nvtx) { return; }
+      event.message.ascii = ascii.c_str();
+      nvtxRangePushEx(&event); // push
+   }
+
+   template<typename T>
+   inline void operator()(const T &arg) const noexcept
+   {
+      if (!nvtx) { return; }
+      operator<<(arg);
+      event.message.ascii = ascii.c_str();
+      nvtxRangePushEx(&event); // push
+   }
+
+   template<typename T, typename... Args>
+   inline void operator()(const char *fmt, const T &arg,
+                          Args... args) const noexcept
+   {
+      if (!nvtx) { return; }
+
+      for (; *fmt != '\0'; fmt++ )
+      {
+         if (*fmt == '%')
+         {
+            fmt++;
+            const char c = *fmt;
+            if (c == 'p') { operator<<(arg); }
+            if (c == 's' || c == 'd' || c == 'f') { operator<<(arg); }
+            if (c == 'x' || c == 'X')
+            {
+               mfem::out << std::hex;
+               if (c == 'X') { mfem::out << std::uppercase; }
+               operator<<(arg);
+               mfem::out << std::nouppercase << std::dec;
+            }
+            if (c == '.')
+            {
+               fmt++;
+               const char c = *fmt;
+               char num[8] = { 0 };
+               for (int k = 0; *fmt != '\0'; fmt++, k++)
+               {
+                  if (*fmt == 'e' || *fmt == 'f') { break; }
+                  if (*fmt < 0x30 || *fmt > 0x39) { break; }
+                  num[k] = *fmt;
+               }
+               const int fx = std::atoi(num);
+               if (c == 'e') { mfem::out << std::scientific; }
+               if (c == 'f') { mfem::out << std::fixed; }
+               mfem::out << std::setprecision(fx);
+               operator<<(arg);
+               mfem::out << std::setprecision(6);
+            }
+            return operator()(fmt + 1, args...);
+         }
+         operator<<(*fmt);
+      }
+      // should never be here
+      assert(false);
+   }
+
+   template <typename T>
+   inline void operator<<(const T &arg) const noexcept
+   {
+      if (!nvtx) { return; }
+      ascii += arg;
+   }
+
+   inline void operator<<(const int &arg) const noexcept
+   {
+      if (!nvtx) { return; }
+      ascii += std::to_string(arg);
+   }
+
+
+public:
+   static const Nvtx Set(const char *FILE, const int LINE, const char *FUNC,
+                         uint32_t COLOR)
+   {
+      static bool env_nvtx = false;
+      static bool env_eks = false;
+      static bool ini_nvtx = false;
+      if (!ini_nvtx)
+      {
+         env_nvtx = getenv("MFEM_NVTX") != nullptr;
+         env_eks = getenv("MFEM_EKS") != nullptr;
+         ini_nvtx = true;
+      }
+      return env_nvtx ? Nvtx(env_eks, FILE, LINE, FUNC, COLOR) : Nvtx();
+   }
+
+private:
+   inline const char *Strrnchr(const char *s, const unsigned char c, int n)
+   {
+      size_t len = strlen(s);
+      char *p = const_cast<char*>(s) + len - 1;
+      for (; n; n--,p--,len--)
+      {
+         for (; len; p--,len--)
+            if (*p == c) { break; }
+         if (!len) { return nullptr; }
+         if (n == 1) { return p; }
+      }
+      return nullptr;
+   }
+};
+
+#ifndef MFEM_NVTX_COLOR
+#define MFEM_NVTX_COLOR SeaGreen
+#endif
+
+// Helpers for generating unique variable names
+#define NVTX_PRIVATE_NAME(name) NVTX_PRIVATE_CONCAT(name, __LINE__)
+#define NVTX_PRIVATE_CONCAT(a, b) NVTX_PRIVATE_CONCAT2(a, b)
+#define NVTX_PRIVATE_CONCAT2(a, b) a##b
+
+// temporary object which is only alive in the expression
+// __PRETTY_FUNCTION__
+#define NVTX(...) \
+    mfem::Nvtx NVTX_PRIVATE_NAME(nvtx) = \
+    mfem::Nvtx::Set(__FILE__,__LINE__,__FUNCTION__,MFEM_NVTX_COLOR);\
+    NVTX_PRIVATE_NAME(nvtx).operator()(__VA_ARGS__)
+
+#define MFEM_NVTX NVTX()
+
+} // namespace mfem
+
+#endif // MFEM_NVVP_HPP

linalg/hypre.cpp

@@ -758,6 +758,7 @@ HypreParMatrix::HypreParMatrix(MPI_Comm comm, HYPRE_BigInt glob_size,
 
    // FIXME:
 #ifdef HYPRE_BIGINT
+   diag->HostReadWriteJ();
    CopyCSR_J(A->diag, diag->GetJ());
 #endif
 

linalg/hypre.hpp

@@ -989,7 +989,8 @@ public:
       ABORT_HYPRE_ERRORS   ///< Abort on hypre errors (default in base class)
    };
 
-protected:
+   //protected:
+public:
    /// The linear system matrix
    const HypreParMatrix *A;
 

linalg/sparsemat.cpp

@@ -181,6 +181,7 @@ SparseMatrix::SparseMatrix(const SparseMatrix &mat, bool copy_graph,
 {
    if (mat.Finalized())
    {
+      mat.HostReadI();
       const int nnz = mat.I[height];
       if (copy_graph)
       {
@@ -1432,6 +1433,7 @@ int SparseMatrix::NumNonZeroElems() const
 {
    if (A != NULL)  // matrix is finalized
    {
+      HostReadI();
       return I[height];
    }
    else

linalg/vector.cpp

@@ -15,6 +15,9 @@
 #include "vector.hpp"
 #include "../general/forall.hpp"
 
+#define MFEM_NVTX_COLOR SlateBlue
+#include "../general/nvtx.hpp"
+
 #if defined(MFEM_USE_SUNDIALS)
 #include "sundials.hpp"
 #if defined(MFEM_USE_MPI)
@@ -771,14 +774,13 @@ void Vector::PrintHash(std::ostream &out) const
 
 void Vector::Randomize(int seed)
 {
+   MFEM_NVTX;
+   if (seed == 0) {  seed = (int)time(0); }
+
+#ifndef MFEM_USE_CUDA
    // static unsigned int seed = time(0);
    const double max = (double)(RAND_MAX) + 1.;
 
-   if (seed == 0)
-   {
-      seed = (int)time(0);
-   }
-
    // srand(seed++);
    srand((unsigned)seed);
 
@@ -787,6 +789,15 @@ void Vector::Randomize(int seed)
    {
       data[i] = std::abs(rand()/max);
    }
+#else
+   if (curng == nullptr)
+   {
+      curng = new curandGenerator_t();
+      curandCreateGenerator(curng, CURAND_RNG_PSEUDO_DEFAULT);
+      curandSetPseudoRandomGeneratorSeed(*curng, (unsigned)seed);
+   }
+   curandGenerateUniformDouble(*curng, Write(), size);
+#endif
 }
 
 double Vector::Norml2() const

linalg/vector.hpp

@@ -17,6 +17,7 @@
 #include "../general/adios2stream.hpp"
 #endif
 #include "../general/globals.hpp"
+#include "../general/backends.hpp"
 #include "../general/mem_manager.hpp"
 #include "../general/device.hpp"
 #ifdef MFEM_USE_SUNDIALS
@@ -63,6 +64,9 @@ protected:
 
    Memory<double> data;
    int size;
+#ifdef MFEM_USE_CUDA
+   curandGenerator_t *curng = nullptr;
+#endif
 
 public:
 
@@ -635,6 +639,13 @@ template<> inline void Swap<Vector>(Vector &a, Vector &b)
 inline Vector::~Vector()
 {
    data.Delete();
+#ifdef MFEM_USE_CUDA
+   if (curng)
+   {
+      curandDestroyGenerator(*curng);
+      delete curng;
+   }
+#endif
 }
 
 inline double DistanceSquared(const double *x, const double *y, const int n)

linalg/wamg.cpp

@@ -0,0 +1,394 @@
+// Copyright (c) 2010-2021, 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 <cassert>
+#include <iomanip>
+#include <iostream>
+using namespace std;
+
+#include "wamg.hpp"
+#include "wavelets.hpp"
+
+#define MFEM_DEBUG_COLOR 82
+#include "../general/debug.hpp"
+
+#define MFEM_NVTX_COLOR Cornflower
+#include "../general/nvtx.hpp"
+
+#include "../general/forall.hpp"
+#include "../general/socketstream.hpp"
+
+namespace mfem
+{
+
+////////////////////////////////////////////////////////////////////////////////
+WaveletRecursiveLevel::WaveletRecursiveLevel(Wavelet::Type &wavelet,
+                                             const bool &lowpass,
+                                             const Operator &A):
+   W(Wavelet::New(wavelet, A.Height(), lowpass)),
+   Wt(new TransposeOperator(W)),
+   WAWt(new RAPOperator(*Wt,A,*Wt))
+{
+   MFEM_NVTX;
+   MFEM_VERIFY((wavelet == Wavelet::HAAR && lowpass) ||
+               (wavelet == Wavelet::DAUBECHIES && lowpass),
+               "Wavelet spec error!");
+   dbg("A: %dx%d", A.Height(), A.Width());
+   dbg("W: %dx%d", W->Height(), W->Width());
+   dbg("WAWt: %dx%d", WAWt->Height(), WAWt->Width());
+   MFEM_VERIFY(A.Height() == A.Width(), "Operator should be square!");
+   MFEM_VERIFY(WAWt->Height() == WAWt->Width(), "WAWt operator should be square!");
+
+   MFEM_VERIFY(W->Width() == A.Width(),
+               "Dimensions error: " << WAWt->Height() <<"x"<< WAWt->Width());
+}
+
+WaveletRecursiveLevel::~WaveletRecursiveLevel()
+{
+   delete W;
+   delete Wt;
+   delete WAWt;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+WaveletRecursiveLevelFA::WaveletRecursiveLevelFA(FiniteElementSpace &fes,
+                                                 Wavelet::Type &wavelet,
+                                                 const OperatorHandle &Ah)
+{
+   MFEM_NVTX;
+   MFEM_VERIFY((wavelet == Wavelet::HAAR) ||
+               (wavelet == Wavelet::DAUBECHIES), "Wavelet argument error!");
+
+   switch (wavelet)
+   {
+      case Wavelet::HAAR:
+         W = new HaarWavelet(Ah->Height(), true); break;
+      case Wavelet::DAUBECHIES:
+         W = new DaubechiesWavelet(Ah->Height(), true); break;
+      default: assert(false);
+   }
+   Wt = new TransposeOperator(W); // the prolongator
+
+   // Get the forward and backward matrices
+   M = mfem::Transpose(*W->GetMatrix()); // Prepare for the Rt of mfem::RAP
+   tM = W->GetTransposedMatrix();
+
+#ifdef MFEM_USE_MPI
+   ParFiniteElementSpace *pfes = dynamic_cast<ParFiniteElementSpace*>(&fes);
+   assert(pfes);
+
+   OperatorHandle H(Operator::Hypre_ParCSR),
+                  tH(Operator::Hypre_ParCSR);
+
+   // Compute the row and colum offsets
+   HYPRE_BigInt locals[2];
+   Array<HYPRE_BigInt> row_offsets;
+   Array<HYPRE_BigInt> col_offsets;
+   int glob_num_rows, glob_num_cols;
+   assert(HYPRE_AssumedPartitionCheck());
+   Array<HYPRE_BigInt> *offsets[2] = { &row_offsets, &col_offsets };
+
+   auto CreateRectangularHypreMatrix = [&] (OperatorHandle &H,
+                                            SparseMatrix *A)
+   {
+      NVTX("CreateRectangularHypreMatrix");
+      locals[0] = A->Height();
+      locals[1] = A->Width();
+      pfes->GetParMesh()->GenerateOffsets(2, locals, offsets);
+      glob_num_rows = row_offsets[row_offsets.Size()-1];
+      glob_num_cols = col_offsets[col_offsets.Size()-1];
+      dbg("glob_num_rows:%d glob_num_cols:%d", glob_num_rows, glob_num_cols);
+      H.MakeRectangularBlockDiag(pfes->GetComm(),
+                                 glob_num_rows, glob_num_cols,
+                                 row_offsets, col_offsets, A);
+   };
+
+   CreateRectangularHypreMatrix(H, M);
+   H.SetOperatorOwner(false);
+   assert(H.Ptr());
+   assert(H.Is<HypreParMatrix>());
+
+   CreateRectangularHypreMatrix(tH, tM);
+   tH.SetOperatorOwner(false);
+   assert(tH.Ptr());
+   assert(tH.Is<HypreParMatrix>());
+
+   const HypreParMatrix *A = Ah.As<HypreParMatrix>();
+   {
+      NVTX("mfem::RAP");
+      assert(A);
+      assert(*H.As<HypreParMatrix>());
+      assert(*tH.As<HypreParMatrix>());
+      MAMt.Reset(mfem::RAP(H.As<HypreParMatrix>(),
+                           A,
+                           tH.As<HypreParMatrix>()),
+                 false);
+   }
+#else
+   assert(false);
+   MFEM_CONTRACT_VAR(fes);
+   const SparseMatrix *A = Ah.As<SparseMatrix>();
+   MAMt.Reset(new RAPOperator(*Wt,*A,*Wt), false);
+#endif // MFEM_USE_MPI
+
+   MFEM_VERIFY(A->Height() == A->Width(), "A should be square!");
+   MFEM_VERIFY(MAMt->Height() == MAMt->Width(), "MAMt should be square!");
+   MFEM_VERIFY(W->Width() == A->Width(),
+               "Dimensions error: " << MAMt->Height() <<"x"<< MAMt->Width());
+}
+
+OperatorHandle WaveletRecursiveLevelFA::OpHandle()
+{
+   return MAMt;
+   //return OperatorHandle(MAMt, false);
+}
+
+Operator *WaveletRecursiveLevelFA::Prolongator() { return Wt; }
+
+WaveletRecursiveLevelFA::~WaveletRecursiveLevelFA()
+{
+   delete W;
+   delete Wt;
+   //delete MAMt;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// \brief The IdentitySolver class
+class IdentitySolver: public Solver
+{
+public:
+   IdentitySolver(int s): Solver(s) { }
+   void SetOperator(const Operator&) { }
+   virtual void Mult(const Vector &x, Vector &y) const { y = x; }
+   virtual void MultTranspose(const Vector &x, Vector &y) const  { y = x; }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// \brief WAMGR solver with DIAGONAL
+WAMGRSolver::WAMGRSolver(FiniteElementSpace &fes,
+                         Wavelet::Type wavelet,
+                         const bool lowpass,
+                         wargs_t args,
+                         const bool to_full) :
+   Multigrid()
+{
+   MFEM_NVTX;
+   const int n = args.op_h->Height();
+   const int max_depth = args.mg_depth;
+   const int max_ndofs = args.mg_ndofs;
+   dbg("COARSE WAMGRSolver n:%d, %s wavelets max_levels:%d",
+       n, Wavelet::GetType(wavelet).c_str(),  max_depth);
+
+   Array<Solver*> Smoothers;
+   Array<Operator*> Operators;
+   Array<Operator*> Prolongators;
+
+   const bool round_up = true;
+   int m = (round_up ? (n+(n%2&1)) : (n-(n%2&1))) >> 1;
+   dbg("n:%d m:%d",n,m);
+
+   OperatorHandle Op_h;
+   Op_h.Reset(args.op_h.Ptr(),false);
+
+   for (int depth = 1; true; depth+=1)
+   {
+      NVTX("DEPTH %d",depth);
+      dbg("\033[31mDEPTH:%d",depth);
+      WaveletLevel *L = nullptr;
+      if (!to_full)
+      {
+         L = new WaveletRecursiveLevel(wavelet, lowpass, *Op_h);
+      }
+      else
+      {
+         assert(lowpass);
+         L = new WaveletRecursiveLevelFA(fes, wavelet, Op_h);
+      }
+      assert(L);
+      assert(L->OpHandle().Ptr());
+      assert(L->OpHandle()->Height() == m && L->OpHandle()->Width() == m);
+
+      dbg("Smoother: (%dx%d)", Op_h->Height(), Op_h->Width());
+      Solver* smoother = nullptr;
+      if (depth==1)
+      {
+         NVTX("Smoother @ 1");
+         if (!args.smoother_order)
+         {
+            dbg("Jacobi smoother");
+            smoother = new OperatorJacobiSmoother(args.diag,
+                                                  args.ess_tdof_list);
+         }
+         else
+         {
+            dbg("Chebyshev smoother");
+            OperatorChebyshevSmoother *chebyshev_smoother =
+               new OperatorChebyshevSmoother(*Op_h.Ptr(),
+                                             args.diag,
+                                             args.ess_tdof_list,
+                                             args.smoother_order
+#ifdef MFEM_USE_MPI
+                                             ,MPI_COMM_WORLD
+#endif // MFEM_USE_MPI
+                                            );
+            smoother = chebyshev_smoother;
+         }
+      }
+      else
+      {
+         NVTX("Smoother @ %d",depth);
+         Vector diag(Op_h->Width());
+         Op_h->AssembleDiagonal(diag);
+         if (!args.smoother_order)
+         {
+            dbg("Jacobi smoother");
+            smoother = new OperatorJacobiSmoother(diag, Array<int>());
+         }
+         else
+         {
+            dbg("Chebyshev smoother");
+            OperatorChebyshevSmoother *chebyshev_smoother =
+               new OperatorChebyshevSmoother(*Op_h.Ptr(),
+                                             diag,
+                                             Array<int>(),
+                                             args.smoother_order
+#ifdef MFEM_USE_MPI
+                                             ,MPI_COMM_WORLD
+#endif // MFEM_USE_MPI
+                                            );
+            smoother = chebyshev_smoother;
+         }
+      }
+      Smoothers.Append(smoother);
+      Operators.Append(Op_h.Ptr());
+
+      Operator *P = L->Prolongator();
+      Prolongators.Append(P);
+
+      OperatorHandle LOp_h;
+      LOp_h.Reset(L->OpHandle().Ptr(),false);
+      dbg("max_ndofs set to %d", max_ndofs);
+      const bool coarse_enough = m < max_ndofs; // stop when too small
+      const bool reached_max_levels = depth >= max_depth;
+      const bool local_coarse_enough_OR_reached_max_levels =
+         coarse_enough || reached_max_levels;
+      bool coarse_enough_OR_reached_max_levels;
+
+#ifdef MFEM_USE_MPI
+      {
+         NVTX("MPI_Allreduce");
+         MPI_Allreduce(&local_coarse_enough_OR_reached_max_levels,
+                       &coarse_enough_OR_reached_max_levels,
+                       1, MPI_CXX_BOOL, MPI_LOR, MPI_COMM_WORLD);
+      }
+#else
+      coarse_enough_OR_reached_max_levels =
+         local_coarse_enough_OR_reached_max_levels;
+#endif // MFEM_USE_MPI
+
+      if (coarse_enough_OR_reached_max_levels)
+      {
+         NVTX("coarse_enough_OR_reached_max_levels");
+         if (to_full)
+         {
+            dbg("Sparse & Full level 0: %dx%d", LOp_h->Height(), LOp_h->Width());
+            CGSolver* wcg =  new CGSolver(
+#ifdef MFEM_USE_MPI
+               MPI_COMM_WORLD
+#endif
+            );
+            wcg->SetMaxIter(args.max_iter);
+            wcg->SetRelTol(1e-8);
+            wcg->SetAbsTol(1e-8);
+            wcg->SetPrintLevel(args.print_level);
+            wcg->SetOperator(*LOp_h);
+            wcg->iterative_mode = false;
+
+            Smoothers.Append(wcg);
+            Operators.Append(LOp_h.Ptr());
+         }
+         else
+         {
+            const int depth = Prolongators.Size();
+            dbg("Coarse solver level 0 (%dx%d), depth:%d",
+                LOp_h->Height(), LOp_h->Width(), depth);
+            CGSolver *wcg = new CGSolver(
+#ifdef MFEM_USE_MPI
+               MPI_COMM_WORLD
+#endif // MFEM_USE_MPI
+            );
+            wcg->SetMaxIter(args.max_iter);
+            wcg->SetRelTol(1e-8);
+            wcg->SetAbsTol(1e-8);
+            wcg->SetPrintLevel(args.print_level);
+            wcg->SetOperator(*LOp_h);
+            wcg->iterative_mode = false;
+
+            Smoothers.Append(wcg);
+            Operators.Append(LOp_h.Ptr());
+         }
+         break;
+      }
+
+      Op_h.Reset(LOp_h.Ptr(), false);
+      m = (round_up ? (m+(m%2&1)) : (m-(m%2&1))) >> 1;
+   }
+
+   const int depth = Prolongators.Size();
+
+   // Coarse solver
+   NVTX("AddLevels");
+   AddLevel(Operators[depth], Smoothers[depth], true, true);
+
+   for (int level = 1; level <= depth; level+=1)
+   {
+      NVTX("level %d", level);
+      const int idx = depth-level;
+      assert(idx>=0);
+      const int Hop = Operators[idx]->Height();
+      const int Wop = Operators[idx]->Width();
+      assert(Smoothers[idx]->Height() == Hop);
+      assert(Smoothers[idx]->Width() == Wop);
+      dbg("@%d %dx%d",level,Hop,Wop);
+      AddLevel(Operators[idx], Smoothers[idx], false, true);
+      // Add the prolongation operator associated with this level
+      prolongations.Append(Prolongators[idx]);
+      ownedProlongations.Append(false);
+   }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// \brief WAMG solver
+WAMG::WAMG(FiniteElementSpace &fes,
+           Wavelet::Type wavelet, wargs_t args):
+   wavelet_solver(fes, wavelet, lowpass, args, to_full)
+{
+   dbg();
+}
+
+void WAMG::Mult(const Vector &x, Vector &y) const
+{
+   wavelet_solver.Mult(x,y);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// \brief faWAMG solver
+faWAMG::faWAMG(FiniteElementSpace &fes,
+               Wavelet::Type wavelet,
+               wargs_t args):
+   wavelet_solver(fes, wavelet, lowpass, args, to_full)
+{
+   dbg();
+}
+
+void faWAMG::Mult(const Vector &x, Vector &y) const {wavelet_solver.Mult(x,y);}
+
+} // namespace mfem
+

linalg/wamg.hpp

@@ -0,0 +1,128 @@
+// Copyright (c) 2010-2021, 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
+
+#ifndef MFEM_WAMG_HPP
+#define MFEM_WAMG_HPP
+
+#include "../fem/fem.hpp"
+#include "wavelets.hpp"
+
+namespace mfem
+{
+
+////////////////////////////////////////////////////////////////////////////////
+struct wargs_t
+{
+   const OperatorHandle &op_h;
+   Vector &diag;
+   Array<int> &ess_tdof_list;
+   int smoother_order, max_iter, print_level;
+   const int mg_depth, mg_ndofs;
+   BilinearForm *a;
+
+   wargs_t(const OperatorHandle &op,
+           Vector &diag,
+           Array<int> &ess_tdof_list,
+           int smoother_order,
+           int max_iter,
+           int print_level,
+           const int mg_depth,
+           const int mg_ndofs,
+           BilinearForm *a = nullptr):
+      op_h(op),
+      diag(diag),
+      ess_tdof_list(ess_tdof_list),
+      smoother_order(smoother_order),
+      max_iter(max_iter),
+      print_level(print_level),
+      mg_depth(mg_depth),
+      mg_ndofs(mg_ndofs),
+      a(a) {}
+};
+
+////////////////////////////////////////////////////////////////////////////////
+struct WaveletLevel
+{
+   virtual OperatorHandle OpHandle() = 0;
+   virtual Operator *Prolongator() = 0;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// Wavelet Recursive level associated with an PA operator
+struct WaveletRecursiveLevel: public WaveletLevel
+{
+   Operator *W, *Wt, *WAWt;
+   WaveletRecursiveLevel(Wavelet::Type &wavelet, const bool &lowpass,
+                         const Operator &A);
+   ~WaveletRecursiveLevel();
+   OperatorHandle OpHandle() override { return OperatorHandle(WAWt, false); }
+   Operator *Prolongator() override { return Wt; }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// Wavelet Recursive level associated with a FULL operator
+struct WaveletRecursiveLevelFA: public WaveletLevel
+{
+   Wavelet *W;
+   Operator *Wt;
+   SparseMatrix *M, *tM;
+   OperatorHandle MAMt;
+   WaveletRecursiveLevelFA(FiniteElementSpace&,
+                           Wavelet::Type &wavelet,
+                           const OperatorHandle &Op_h);
+   ~WaveletRecursiveLevelFA();
+   OperatorHandle OpHandle() override;
+   Operator *Prolongator() override;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+class WAMGRSolver : public Multigrid
+{
+public:
+   WAMGRSolver(FiniteElementSpace &fes,
+               Wavelet::Type wavelet,
+               const bool lowpass,
+               wargs_t args,
+               const bool to_full);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// @brief WAMG solver
+/// wavelet == HAAR | DAUBECHIES & lowpass
+class WAMG : public Solver
+{
+   const bool lowpass = true;
+   const bool to_full = false;
+   WAMGRSolver wavelet_solver;
+public:
+   WAMG(FiniteElementSpace &fes, Wavelet::Type wavelet, wargs_t args);
+   void Mult(const Vector&, Vector&) const override;
+   void SetOperator(const Operator&) override { assert(false); }
+   void AssembleDiagonal(Vector&) const override { assert(false); }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// @brief faWAMG solver
+/// wavelet == HAAR | DAUBECHIES & lowpass
+class faWAMG : public Solver
+{
+   const bool lowpass = true;
+   const bool to_full = true;
+   WAMGRSolver wavelet_solver;
+public:
+   faWAMG(FiniteElementSpace &fes, Wavelet::Type wavelet, wargs_t args);
+   void Mult(const Vector&, Vector&) const override;
+   void SetOperator(const Operator&) override { assert(false); }
+   void AssembleDiagonal(Vector&) const override { assert(false); }
+};
+
+} // mfem namespace
+
+#endif // MFEM_WAMG_HPP

linalg/wavelets.cpp

@@ -0,0 +1,910 @@
+// Copyright (c) 2010-2021, 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 "wavelets.hpp"
+#include "../general/forall.hpp"
+
+#define MFEM_NVTX_COLOR Orange
+#include "../general/nvtx.hpp"
+
+using namespace mfem;
+
+namespace mfem
+{
+
+////////////////////////////////////////////////////////////////////////////////
+MFEM_HOST_DEVICE inline int RotateLeft(int i, int M) { return (i+1) % M; }
+MFEM_HOST_DEVICE inline int RotateLeft2(int i, int M) { return (i+2) % M; }
+MFEM_HOST_DEVICE inline int RotateRight(int i, int M) { return (i+M-1) % M; }
+MFEM_HOST_DEVICE inline int RotateRight2(int i, int M) { return (i+M-2) % M; }
+
+////////////////////////////////////////////////////////////////////////////////
+Operator *Wavelet::New(const Wavelet::Type &wavelet, int n, bool lowpass)
+{
+   switch (wavelet)
+   {
+      case Wavelet::HAAR: return new HaarWavelet(n,lowpass);
+      case Wavelet::DAUBECHIES: return new DaubechiesWavelet(n,lowpass);
+      // Lowpass to add to the folowing wavelets
+      case Wavelet::CDF53: return new CDF53Wavelet(n);
+      case Wavelet::CDF97: return new CDF97Wavelet(n);
+      case Wavelet::LEGENDRE2: return new Legendre2Wavelet(n);
+      case Wavelet::LEGENDRE3: return new Legendre3Wavelet(n);
+      default: MFEM_ABORT("No Wavelet Operator found!");
+   }
+   return nullptr;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+std::string Wavelet::GetType(const Wavelet::Type &wavelet)
+{
+   switch (wavelet)
+   {
+      case Wavelet::HAAR: return std::string("HAAR");
+      case Wavelet::DAUBECHIES: return std::string("DAUBECHIES");
+      case Wavelet::CDF53: return std::string("CDF53");
+      case Wavelet::CDF97: return std::string("CDF97");
+      case Wavelet::LEGENDRE2: return std::string("LEGENDRE2");
+      case Wavelet::LEGENDRE3: return std::string("LEGENDRE3");
+      default: MFEM_ABORT("Not a valid wavelet type!");
+   }
+   return std::string("???");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+SparseMatrix *Wavelet::GetMatrix()
+{
+   MFEM_ABORT("Not implemented");
+   return nullptr;
+}
+
+SparseMatrix *Wavelet::GetTransposedMatrix()
+{
+   MFEM_ABORT("Not implemented");
+   return nullptr;
+}
+
+double Wavelet::GetOddRowValue(const int, const int, const bool)
+{
+   MFEM_ABORT("Not implemented");
+   return 0.0;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+SparseMatrix *Wavelet::GetEvenMatrix(const double *coeffs, const int rowsize)
+{
+   MFEM_NVTX;
+   const int ncols = width;
+   const int nrows = height;
+   const int nnz = nrows * rowsize;
+
+   SparseMatrix *A = new SparseMatrix(nrows, ncols, rowsize);
+
+   int *J = A->GetJ();
+   int *I = A->GetI();
+   int row_offsets = 0;
+   double *data = A->GetData();
+
+   assert(m*2 == width);
+   assert(((width-rowsize)%2&1) == 0);
+   const int shifts = 1 + (width-rowsize)/2; // 0 + possible shifts
+   const int warped = m - shifts;
+
+   // Lowpass
+   for (int i=0, j=0; i<shifts; i++, j+=2)
+   {
+      I[i] = row_offsets;
+      for (int nz=0; nz<rowsize; nz++)
+      {
+         J[I[i]+nz] = j + nz;
+         *data++ = coeffs[nz];
+      }
+      row_offsets += rowsize;
+   }
+
+   // Warped lines
+   for (int i=shifts, nz=0; i<shifts+warped; i++)
+   {
+      // Lowpass-warped
+      for (int j=0; nz<rowsize-2; nz++,j++)
+      {
+         J[I[i]+nz] = j;
+         *data++ = coeffs[nz+2];
+      }
+      for (int j=width-2; nz<rowsize; nz++, j++)
+      {
+         J[I[i]+nz] = j;
+         *data++ = coeffs[nz-2];
+      }
+      row_offsets += rowsize;
+   }
+
+   if (!lowpass)
+   {
+      for (int i=shifts+warped, nz=0; i<shifts+2*warped; i++)
+      {
+         // Highpass-warped
+         I[i] = row_offsets;
+         double sign = -1.0;
+         for (int j=0; nz<rowsize-2; nz++,j++)
+         {
+            J[I[i]+nz] = j;
+            *data++ = sign*coeffs[rowsize-1-nz-2];
+            sign *= -1.0;
+         }
+         for (int j=width-2; nz<rowsize; nz++,j++)
+         {
+            J[I[i]+nz] = j;
+            *data++ = sign*coeffs[rowsize-1-nz+2];
+            sign *= -1.0;
+         }
+         row_offsets += rowsize;
+      }
+
+      // Highpass
+      for (int i=shifts+2*warped, j=0; i<height; i++, j+=2)
+      {
+         I[i] = row_offsets;
+         double sign = -1.0;
+         for (int nz=0; nz<rowsize; nz++)
+         {
+            J[I[i]+nz] = j+nz;
+            *data++ = sign*coeffs[rowsize-1-nz];
+            sign *= -1.0;
+         }
+         row_offsets += rowsize;
+      }
+   }
+
+   I[height] = row_offsets;
+   MFEM_VERIFY(I[0] == 0, "Error");
+   MFEM_VERIFY(I[height] == nnz, "Error");
+   A->Finalize();
+   return A;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+SparseMatrix *Wavelet::GetOddMatrix(const double *coeffs, const int rowsize,
+                                    const bool for_transpose)
+{
+   MFEM_NVTX;
+   const int nrows = height;
+   const int ncols = width;
+   // sparse matrix with flexible sparsity structure
+   SparseMatrix *A = new SparseMatrix(nrows, ncols);
+
+   Vector x(width), y(height);
+
+   const int shifts = 1 + (width-rowsize)/2;
+   const int warped = m - shifts - 1;
+
+   // Lowpass
+   for (int i=0, j=0; i<shifts; i++, j+=2)
+   {
+      for (int nz=0; nz<rowsize; nz++)
+      {
+         A->Set(i, j+nz, coeffs[nz]);
+      }
+   }
+
+   // Warped lines
+   if (!for_transpose)
+   {
+      // lowpass, warped
+      for (int i=shifts; i<shifts+warped; i++)
+      {
+         // DAUBECHIES only
+         assert(type != HAAR);
+         assert(rowsize == 4);
+         for (int nz=0; nz<rowsize-1; nz++)
+         {
+            const int j[3] = {width-3, width-2, width-1};
+            const double val = GetOddRowValue(i, j[nz], false);
+            A->Set(i, j[nz], val);
+         }
+      }
+
+      // middle line
+      {
+         const int i = m - 1;
+         for (int nz=0; nz<rowsize-1; nz++)
+         {
+            const int j = type == HAAR ? width-1 :
+                          // DAUBECHIES
+                          nz==0 ? 0 :
+                          nz==1 ? 1 :
+                          width-1;
+            const double val = GetOddRowValue(i, j, false);
+            A->Set(i, j, val);
+         }
+      }
+
+      if (!lowpass)
+      {
+         // highpass, warped
+         for (int i=shifts+warped+1; i<shifts+2*warped+1; i++)
+         {
+            assert(type != HAAR);
+            assert(rowsize == 4);
+            for (int nz = 0; nz < rowsize-1; nz++)
+            {
+               const int j[3] = {0, 1, width-1};
+               const double val = GetOddRowValue(i, j[nz], false);
+               A->Set(i, j[nz], val);
+            }
+         }
+      }
+   }
+   else // transpose
+   {
+      if (lowpass)
+      {
+         const bool root = type == HAAR ? true : false;
+
+         for (int i=shifts; i<shifts+warped; i++)
+         {
+            // DAUBECHIES only
+            assert(type != HAAR);
+            assert(rowsize == 4);
+            for (int nz=0; nz<rowsize-1; nz++)
+            {
+               const int j[3] = {width-3, width-2, width-1};
+               const double val = GetOddRowValue(i, j[nz], root);
+               A->Set(i, j[nz], val);
+            }
+         }
+
+         // middle line
+         {
+            const int i = m-1;
+            for (int nz=0; nz<rowsize-1; nz++)
+            {
+               const int j = type == HAAR ? width-1 :
+                             // DAUBECHIES
+                             nz==0 ? 0 :
+                             nz==1 ? 1 :
+                             nz==2 ? width-1 :
+                             -1;
+               const double val = GetOddRowValue(i, j, root);
+               A->Set(i, j, val);
+            }
+         }
+      }
+      else // !lowpass
+      {
+         for (int i=shifts; i<shifts+warped; i++)
+         {
+            // DAUBECHIES only
+            assert(type != HAAR);
+            assert(rowsize == 4);
+            for (int nz = 0; nz < rowsize-2; nz++)
+            {
+               const int j[2] = {width-3, width-2};
+               const double val = GetOddRowValue(i, j[nz], true);
+               A->Set(i,j[nz], val);
+            }
+         }
+
+         // middle line
+         {
+            const int i = m-1;
+            for (int nz=0; nz<rowsize+1; nz++)
+            {
+               const int j = type == HAAR ? width-1 :
+                             // DAUBECHIES
+                             nz==0 ? 0 :
+                             nz==1 ? 1 :
+                             nz==2 ? width-3 :
+                             nz==3 ? width-2 :
+                             width-1;
+               const double val = GetOddRowValue(i, j, true);
+               A->Set(i,j, val);
+            }
+         }
+
+         for (int i=shifts+warped+1; i<shifts+2*warped+1; i++)
+         {
+            // DAUBECHIES only
+            assert(type != HAAR);
+            assert(rowsize == 4);
+            for (int nz=0; nz<rowsize+1; nz++)
+            {
+               const int j[5] = {0, 1, width-3, width-2, width-1};
+               const double val = GetOddRowValue(i, j[nz], true);
+               A->Set(i, j[nz], val);
+            }
+         }
+      }
+   }
+
+   // Highpass
+   if (!lowpass)
+   {
+      for (int i=shifts+2*warped+1, j=0; i<height; i++, j+=2)
+      {
+         double sign = -1.0;
+         for (int nz = 0; nz < rowsize; nz++)
+         {
+            A->Set(i, j+nz, sign*coeffs[rowsize-1-nz]);
+            sign *= -1.0;
+         }
+      }
+   }
+
+   A->Finalize();
+   return A;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+SparseMatrix *Wavelet::GetMatrix(const double *coeffs, const int rowsize)
+{
+   MFEM_NVTX;
+   return odd ?
+          GetOddMatrix(coeffs, rowsize, false) :
+          GetEvenMatrix(coeffs, rowsize);
+}
+
+SparseMatrix *Wavelet::GetTransposedMatrix(const double *coeffs,
+                                           const int rowsize)
+{
+   MFEM_NVTX;
+   return odd ?
+          Transpose(*GetOddMatrix(coeffs, rowsize, true)) :
+          Transpose(*GetEvenMatrix(coeffs, rowsize));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+HaarWavelet::HaarWavelet(int n, bool lowpass): Wavelet(HAAR, n,lowpass) { }
+
+void HaarWavelet::Mult(const Vector &x, Vector &y) const
+{
+   MFEM_NVTX;
+   const int M = m;
+   const int O = odd;
+   const bool LOWPASS = lowpass;
+
+   const double isq2 = 1.0 / sqrt(2.0);
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const double Si = X[2*i];
+      const double Di = (O && i==M-1) ? 0.0 : X[2*i+1];
+      const double ds = (Di - Si) * isq2;
+      Y[i] = 2.0 * Si * isq2 + ds;
+      if (LOWPASS) { return; }
+      if (O && i==M-1) { return; }
+      Y[M+i] = ds;
+   });
+}
+
+double HaarWavelet::GetOddRowValue(const int I, const int J,
+                                   const bool root)
+{
+   MFEM_NVTX;
+   assert(odd);
+   assert(I == m-1);
+   assert(I == J/2);
+   assert(J == width-1);
+   const double sq2 = sqrt(2.0);
+   const double isq2 = 1.0 / sq2;
+   const double Di = root ? -1.0 : 0.0;
+   const double Si = (1.0 - Di) * isq2;
+   return Si; // root ? sq2 : isq2;
+}
+
+void HaarWavelet::MultTranspose(const Vector &x, Vector &y) const
+{
+   MFEM_NVTX;
+   const int M = m;
+   const int O = odd;
+   const bool LOWPASS = lowpass;
+
+   const double sq2 = sqrt(2.0);
+   const double isq2 = 1.0 / sq2;
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const double Di = (O && i==M-1) ? -X[M-1] : LOWPASS ? 0.0 : X[M+i];
+      const double Si = (X[i] - Di) * isq2;
+      Y[2*i] = Si;
+      if (O && i==M-1) { return; }
+      Y[2*i+1] = Di * sq2 + Si;
+   });
+}
+
+SparseMatrix *HaarWavelet::GetMatrix()
+{
+   MFEM_NVTX;
+   constexpr int rowsize = 2;
+   const double isq2 = 1.0 / sqrt(2.0);
+   const double coeffs[rowsize] = { isq2, isq2 };
+   return Wavelet::GetMatrix(coeffs, rowsize);
+}
+
+SparseMatrix *HaarWavelet::GetTransposedMatrix()
+{
+   MFEM_NVTX;
+   constexpr int rowsize = 2;
+   const double isq2 = 1.0 / sqrt(2.0);
+   const double coeffs[rowsize] = { isq2, isq2 };
+   return Wavelet::GetTransposedMatrix(coeffs, rowsize);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+DaubechiesWavelet::DaubechiesWavelet(int n, bool lowpass):
+   Wavelet(DAUBECHIES, n, lowpass) {}
+
+void DaubechiesWavelet::Mult(const Vector &x, Vector &y) const
+{
+   const int M = m;
+   const int O = odd;
+   const bool LOWPASS = lowpass;
+
+   const double sq2 = sqrt(2.0);
+   const double sq3 = sqrt(3.0);
+   const double sq34mh = sq3 / 4.0 - 0.5;
+   const double sq2i3m1 = sq2 / (sq3 - 1.0);
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int r = RotateRight(i,M);
+      double Si = X[2*i];
+      double Sl = X[2*l];
+      double Sr = X[2*r];
+      double Di = (O && i==M-1) ? 0.0 : X[2*i+1];
+      double Dl = (O && l==M-1) ? 0.0 : X[2*l+1];
+      double Dr = (O && r==M-1) ? 0.0 : X[2*r+1];
+      Di -= sq3 * Si;
+      Dl -= sq3 * Sl;
+      Dr -= sq3 * Sr;
+      Si += sq3 * Di / 4.0 + sq34mh * Dl;
+      Sr += sq3 * Dr / 4.0 + sq34mh * Di;
+      Di += Sr;
+      Si *= sq2i3m1;
+      Di /= sq2i3m1;
+      Y[i] = Si;
+      if (LOWPASS) { return; }
+      if (O && i==M-1) { return; }
+      Y[M+i] = Di;
+   });
+}
+
+double DaubechiesWavelet::GetOddRowValue(const int I,
+                                         const int J,
+                                         const bool root)
+{
+   assert(odd);
+
+   const auto X = [&](int i) { return I==i ? 1.0 : 0.0; };
+
+   const double sq3 = sqrt(3.0);
+   const double msq3q = -sq3 / 4.0;
+   const double hmsq34 = 0.5 + msq3q;
+   const double sq2i3m1 = sqrt(2.0) / (sq3 - 1.0);
+
+   const int i = J/2;
+   const int l = RotateLeft(i,m);
+   const int r = RotateRight(i,m);
+   double Si = X(i);
+   double Sr = X(r);
+   const bool oi = i==m-1, ol = l==m-1;
+   const bool ox = oi || ol;
+   const double Xi =
+      ox && root ? (3-2*sq3)*X(m)-sq3*X(m-1)+(2-sq3)*X(2*(m-1)%m) : 0.0;
+   double Di = oi ? Xi : X(m+i);
+   double Dl = ol ? Xi : X(m+l);
+   Si /= sq2i3m1;
+   Sr /= sq2i3m1;
+   Di *= sq2i3m1;
+   Dl *= sq2i3m1;
+   Di -= Sr;
+   Dl -= Si;
+   Si += msq3q * Di + hmsq34 * Dl;
+   Di += sq3 * Si;
+   return J%2 ? Di : Si;
+}
+
+void DaubechiesWavelet::MultTranspose(const Vector &x, Vector &y) const
+{
+   const int M = m;
+   const int O = odd;
+   const bool LOWPASS = lowpass;
+
+   const double sq3 = sqrt(3.0);
+   const double msq3q = -sq3 / 4.0;
+   const double hmsq34 = 0.5 + msq3q;
+   const double sq2i3m1 = sqrt(2.0) / (sq3 - 1.0);
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int r = RotateRight(i,M);
+      double Si = X[i];
+      double Sr = X[r];
+      const bool oi = O && i==M-1;
+      const bool ol = O && l==M-1;
+      const bool ox = oi || ol;
+      const double Xi =
+      ox ? LOWPASS ? 0.0 : (3-2*sq3)*X[M]-sq3*X[M-1]+(2-sq3)*X[2*(M-1)%M] : 0.0;
+      double Di = oi ? Xi : LOWPASS ? 0.0 : X[M+i];
+      double Dl = ol ? Xi : LOWPASS ? 0.0 : X[M+l];
+      Si /= sq2i3m1;
+      Sr /= sq2i3m1;
+      Di *= sq2i3m1;
+      Dl *= sq2i3m1;
+      Di -= Sr;
+      Dl -= Si;
+      Si += msq3q * Di + hmsq34 * Dl;
+      Di += sq3 * Si;
+      Y[2*i] = Si;
+      if (O && i==M-1) { return; }
+      Y[2*i+1] = Di;
+   });
+}
+
+SparseMatrix *DaubechiesWavelet::GetMatrix()
+{
+   constexpr int rowsize = 4;
+   const double sqrt2 = sqrt(2);
+   const double coeffs[rowsize] = { sqrt2*(1+sqrt(3))/8.,
+                                    sqrt2*(3+sqrt(3))/8.,
+                                    sqrt2*(3-sqrt(3))/8.,
+                                    sqrt2*(1-sqrt(3))/8.
+                                  };
+   return Wavelet::GetMatrix(coeffs, rowsize);
+}
+
+SparseMatrix *DaubechiesWavelet::GetTransposedMatrix()
+{
+   constexpr int rowsize = 4;
+   const double sqrt2 = sqrt(2);
+   const double coeffs[rowsize] = { sqrt2*(1+sqrt(3))/8.,
+                                    sqrt2*(3+sqrt(3))/8.,
+                                    sqrt2*(3-sqrt(3))/8.,
+                                    sqrt2*(1-sqrt(3))/8.
+                                  };
+   return Wavelet::GetTransposedMatrix(coeffs, rowsize);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+CDF53Wavelet::CDF53Wavelet(int n): Wavelet(CDF53, n) { }
+
+void CDF53Wavelet::Mult(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const double isq2 = 1.0 / sqrt(2.0);
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int r = RotateRight(i,M);
+      double Si = X[2*i];
+      double Sl = X[2*l];
+      double Di = (odd && i==M-1) ? 0.0 : X[2*i+1];
+      double Dl = (odd && l==M-1) ? 0.0 : X[2*l+1];
+      double Dr = (odd && r==M-1) ? 0.0 : X[2*r+1];
+      Si += Di/2.0 + Dr/2.0;
+      Sl += Dl/2.0 + Di/2.0;
+      Di -= Si/4.0 + Sl/4.0;
+      Si *= isq2;
+      Di /= isq2;
+      Y[i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[M+i] = Di;
+   });
+}
+
+void CDF53Wavelet::MultTranspose(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const double isq2 = 1.0 / sqrt(2.0);
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int r = RotateRight(i,M);
+      double Si = X[i];
+      double Sl = X[l];
+      double Sr = X[r];
+      double Di = (odd && i==M-1) ? -(X[0] + X[M-1]) / 2.0: X[M+i];
+      double Dr = (odd && r==M-1) ? -(X[0] + X[M-1]) / 2.0: X[M+r];
+      Si /= isq2;
+      Sl /= isq2;
+      Sr /= isq2;
+      Di *= isq2;
+      Dr *= isq2;
+      Di += Si / 4.0 + Sl / 4.0;
+      Dr += Sr / 4.0 + Si / 4.0;
+      Si -= Di / 2.0 + Dr / 2.0;
+      Y[2*i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[2*i+1] = Di;
+   });
+}
+
+////////////////////////////////////////////////////////////////////////////////
+CDF97Wavelet::CDF97Wavelet(int n): Wavelet(CDF97, n) { }
+
+void CDF97Wavelet::Mult(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int r = RotateRight(i,M);
+      const int k = RotateLeft2(i,M);
+      const int s = RotateRight2(i,M);
+      double Si = X[2*i];
+      double Sk = X[2*k];
+      double Sl = X[2*l];
+      double Sr = X[2*r];
+      double Ss = X[2*s];
+      double Di = (odd && i==M-1) ? 0.0 : X[2*i+1];
+      double Dl = (odd && l==M-1) ? 0.0 : X[2*l+1];
+      double Dr = (odd && r==M-1) ? 0.0 : X[2*r+1];
+      double Ds = (odd && s==M-1) ? 0.0 : X[2*s+1];
+      Di -= alpha * (Si + Sl);
+      Dl -= alpha * (Sl + Sk);
+      Dr -= alpha * (Sr + Si);
+      Ds -= alpha * (Ss + Sr);
+      Si -=  beta * (Di + Dr);
+      Sl -=  beta * (Dl + Di);
+      Sr -=  beta * (Dr + Ds);
+      Di += gamma * (Si + Sl);
+      Dr += gamma * (Sr + Si);
+      Si += delta * (Di + Dr);
+      Si *= kappa;
+      Di /= kappa;
+      Y[i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[M+i] = Di;
+   });
+}
+
+void CDF97Wavelet::MultTranspose(const Vector &x, Vector &y) const
+{
+   const int n = height;
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   constexpr double a = 0.490316548523234605288;
+   constexpr double b = 0.075687794783398174545;
+   constexpr double c = 0.129734462057497172811;
+   constexpr double d = 0.044363215797333603554;
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int r = RotateRight(i,M);
+      const int k = RotateLeft2(i,M);
+      const int s = RotateRight2(i,M);
+      const bool io = odd && i==M-1;
+      const bool lo = odd && l==M-1;
+      const bool ro = odd && r==M-1;
+      const bool so = odd && s==M-1;
+      const bool ko = odd && k==M-1;
+      const bool xo = io || lo || ro || so || ko;
+      const double Xi = xo ?
+      - a * (X[0] + X[M-1]) + b * (X[1] + X[2*(M-1)%M])
+      + c * (X[M] + X[n-1]) - d * (X[M+1] + X[n-2]) : 0.0;
+      double Si = X[i];
+      double Sl = X[l];
+      double Sk = X[k];
+      double Sr = X[r];
+      double Di = io ? Xi : X[M+i];
+      double Dl = lo ? Xi : X[M+l];
+      double Dk = ko ? Xi : X[M+k];
+      double Dr = ro ? Xi : X[M+r];
+      double Ds = so ? Xi : X[M+s];
+      Si /= kappa;
+      Sl /= kappa;
+      Sk /= kappa;
+      Sr /= kappa;
+      Di *= kappa;
+      Dk *= kappa;
+      Dl *= kappa;
+      Dr *= kappa;
+      Ds *= kappa;
+      Si -= delta * (Di + Dr);
+      Sr -= delta * (Dr + Ds);
+      Sl -= delta * (Dl + Di);
+      Sk -= delta * (Dk + Dl);
+      Di -= gamma * (Si + Sl);
+      Dr -= gamma * (Sr + Si);
+      Dl -= gamma * (Sl + Sk);
+      Si +=  beta * (Di + Dr);
+      Sl +=  beta * (Dl + Di);
+      Di += alpha * (Si + Sl);
+      Y[2*i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[2*i+1] = Di;
+   });
+}
+
+////////////////////////////////////////////////////////////////////////////////
+Legendre2Wavelet::Legendre2Wavelet(int n): Wavelet(LEGENDRE2, n) { }
+
+void Legendre2Wavelet::Mult(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const double k2 = sqrt(2.0) / 5.0;
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      double Si = X[2*i];
+      double Sl = X[2*l];
+      double Di = (odd && i==M-1) ? 0.0 : X[2*i+1];
+      double Dl = (odd && l==M-1) ? 0.0 : X[2*l+1];
+      Di += Si * 3.0 / 5.0;
+      Dl += Sl * 3.0 / 5.0;
+      Si += (15.0 * Di + 25.0 * Dl) / 16.0;
+      Si *= k2;
+      Di /= k2;
+      Y[i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[M+i] = Di;
+   });
+}
+
+void Legendre2Wavelet::MultTranspose(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const double k2 = sqrt(2.0) / 5.0;
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      double Si = X[i];
+      const bool oi = odd && i==M-1;
+      const bool ol = odd && l==M-1;
+      const bool ox = oi || ol;
+      const double Xi = ox ? (3.0 * (8.0 * X[M-1] - X[M])) / 5.0 : 0.0;
+      double Di = oi ? Xi : +X[M+i];
+      double Dl = ol ? Xi : +X[M+l];
+      Si /= k2;
+      Di *= k2;
+      Dl *= k2;
+      Si -= (15.0 * Di + 25.0 * Dl) / 16.0;
+      Di -= 3.0 * Si / 5.0;
+      Y[2*i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[2*i+1] = Di;
+   });
+}
+
+////////////////////////////////////////////////////////////////////////////////
+Legendre3Wavelet::Legendre3Wavelet(int n): Wavelet(LEGENDRE3, n) { }
+
+void Legendre3Wavelet::Mult(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const double k3 = 153.0 / (343.0 * sqrt(2.0));
+   const double i12 = 1215.0 / 2744.0;
+   const double i59 = 5.0 / 9.0;
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int r = RotateRight(i,M);
+      const int l = RotateLeft(i,M);
+      const int k = RotateLeft2(i,M);
+      double Si = X[2*i];
+      double Sr = X[2*r];
+      double Sl = X[2*l];
+      double Sk = X[2*k];
+      double Di = (odd && i==M-1) ? 0.0 : X[2*i+1];
+      double Dr = (odd && r==M-1) ? 0.0 : X[2*r+1];
+      double Dl = (odd && l==M-1) ? 0.0 : X[2*l+1];
+      double Dk = (odd && k==M-1) ? 0.0 : X[2*k+1];
+      Si += Di * i59;
+      Sr += Dr * i59;
+      Sl += Dl * i59;
+      Sk += Dk * i59;
+      Di += 45.0 * Si / 56.0 + i12 * Sr;
+      Dl += 45.0 * Sl / 56.0 + i12 * Si;
+      Dk += 45.0 * Sk / 56.0 + i12 * Sl;
+      Si += 1715.0 * Dl / 7344.0 + 16807.0 * Dk / 22032.0;
+      Si *= k3;
+      Di /= k3;
+      Y[i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[M+i] = Di;
+   });
+}
+
+void Legendre3Wavelet::MultTranspose(const Vector &x, Vector &y) const
+{
+   const int M = (height+1)>>1;
+   const int odd = height%2 &1;
+
+   const double k3 = 153.0 / (343.0 * sqrt(2.0));
+
+   const auto X = x.Read();
+   auto Y = y.Write();
+
+   MFEM_FORALL(i, M,
+   {
+      const int l = RotateLeft(i,M);
+      const int k = RotateLeft2(i,M);
+      const int r = RotateRight(i,M);
+      double Si = X[i];
+      double Sr = X[r];
+      const bool io = odd && i==M-1;
+      const bool il = odd && l==M-1;
+      const bool ik = odd && k==M-1;
+      const bool xo = io || il|| ik;
+      const double Xi = xo ?
+      (5.*(864.*X[2*(M-1)%M]+1568.*X[M-1]-17.*(6.*X[M]+7.*X[M+1])))/1071. : 0.0;
+      double Di = io ? Xi : X[M+i];
+      double Dl = il ? Xi : X[M+l];
+      double Dk = ik ? Xi : X[M+k];
+      Si /= k3;
+      Sr /= k3;
+      Di *= k3;
+      Dl *= k3;
+      Dk *= k3;
+      Si -= 1715.0 * Dl / 7344.0 + 16807.0 * Dk / 22032.0;
+      Sr -= 1715.0 * Di / 7344.0 + 16807.0 * Dl / 22032.0;
+      Di -= 45.0 * Si / 56.0 + 1215.0 * Sr / 2744.0;
+      Si -= Di * 5.0 / 9.0;
+      Y[2*i] = Si;
+      if (odd && i==M-1) { return; }
+      Y[2*i+1] = Di;
+   });
+}
+
+} // namespace mfem

linalg/wavelets.hpp

@@ -0,0 +1,126 @@
+// Copyright (c) 2010-2021, 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_WAVELETS_HPP
+#define MFEM_WAVELETS_HPP
+
+#include "linalg.hpp"
+
+namespace mfem
+{
+
+////////////////////////////////////////////////////////////////////////////////
+struct Wavelet : public Operator
+{
+   enum Type { HAAR = 0,
+               DAUBECHIES,
+               CDF53,
+               CDF97,
+               LEGENDRE2,
+               LEGENDRE3
+             };
+
+   const Type type;
+   const int m, odd;
+   const bool lowpass;
+
+   Wavelet(Type type, int s, bool lowpass = false):
+      Operator(lowpass?(s+1)>>1:s,s),
+      type(type),
+      m(lowpass?height:(height+1)>>1),
+      odd(width%2 &1),
+      lowpass(lowpass) {}
+
+   Wavelet(Type type, int h, int w):
+      Operator(h,w),
+      type(type),
+      m(height),
+      odd(width%2 &1),
+      lowpass(true) { }
+
+   static Operator* New(const Wavelet::Type&, int n, bool lowpass = false);
+   static std::string GetType(const Wavelet::Type&);
+
+   SparseMatrix *GetEvenMatrix(const double *coeffs, const int rowsize);
+   SparseMatrix *GetOddMatrix(const double *coeffs, const int rowsize,
+                              bool for_transpose);
+   SparseMatrix *GetMatrix(const double *coeffs, const int rowsize);
+   SparseMatrix *GetTransposedMatrix(const double *coeffs, const int rowsize);
+
+   virtual SparseMatrix *GetMatrix();
+   virtual SparseMatrix *GetTransposedMatrix();
+   virtual double GetOddRowValue(const int i, const int j, const bool root);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+struct HaarWavelet : public Wavelet
+{
+   HaarWavelet(int n, bool lowpass = false);
+   void Mult(const Vector &x, Vector &y) const;
+   void MultTranspose(const Vector &x, Vector &y) const;
+   SparseMatrix *GetMatrix();
+   SparseMatrix *GetTransposedMatrix();
+   double GetOddRowValue(const int i, const int j, const bool root);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// \brief The DaubechiesWavelet of order 2
+struct DaubechiesWavelet : public Wavelet
+{
+   DaubechiesWavelet(int n, bool lowpass = false);
+   void Mult(const Vector &x, Vector &y) const;
+   void MultTranspose(const Vector &x, Vector &y) const;
+   SparseMatrix *GetMatrix();
+   SparseMatrix *GetTransposedMatrix();
+   double GetOddRowValue(const int i, const int j, const bool root);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+struct CDF53Wavelet : public Wavelet
+{
+   CDF53Wavelet(int n);
+   void Mult(const Vector &x, Vector &y) const;
+   void MultTranspose(const Vector &x, Vector &y) const;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+class CDF97Wavelet : public Wavelet
+{
+   static constexpr double alpha = 1.58613434205992355842832;
+   static constexpr double  beta = 0.05298011857296141462412;
+   static constexpr double gamma = 0.88291107553093329591979;
+   static constexpr double delta = 0.44350685204397115211560;
+   static constexpr double kappa = 1.14960439886024115979508;
+public:
+   CDF97Wavelet(int n);
+   void Mult(const Vector &x, Vector &y) const;
+   void MultTranspose(const Vector &x, Vector &y) const;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+struct Legendre2Wavelet : public Wavelet
+{
+   Legendre2Wavelet(int n);
+   void Mult(const Vector &x, Vector &y) const;
+   void MultTranspose(const Vector &x, Vector &y) const;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+struct Legendre3Wavelet : public Wavelet
+{
+   Legendre3Wavelet(int n);
+   void Mult(const Vector &x, Vector &y) const;
+   void MultTranspose(const Vector &x, Vector &y) const;
+};
+
+} // namespace mfem
+
+#endif // MFEM_WAVELETS_HPP

makefile

@@ -711,6 +711,7 @@ ASTYLE = $(ASTYLE_BIN) --options=$(SRC)config/mfem.astylerc
 ASTYLE_VER = "Artistic Style Version 3.1"
 FORMAT_FILES = $(foreach dir,$(DIRS) $(EM_DIRS) config,$(dir)/*.?pp)
 FORMAT_FILES += tests/unit/*.cpp
+FORMAT_FILES += tests/benchmarks/*.?pp
 UNIT_TESTS_SUBDIRS = general linalg mesh fem miniapps ceed
 FORMAT_FILES += $(foreach dir,$(UNIT_TESTS_SUBDIRS),tests/unit/$(dir)/*.?pp)
 FORMAT_LIST = $(filter-out general/tinyxml2.cpp,$(wildcard $(FORMAT_FILES)))

miniapps/solvers/lor_solvers.cpp

@@ -76,6 +76,7 @@ int main(int argc, char *argv[])
    int ref_levels = 1;
    int order = 3;
    const char *fe = "h";
+   const char *device_config = "cpu";
    bool visualization = true;
 
    OptionsParser args(argc, argv);
@@ -88,14 +89,20 @@ int main(int argc, char *argv[])
    args.AddOption(&visualization, "-vis", "--visualization", "-no-vis",
                   "--no-visualization",
                   "Enable or disable GLVis visualization.");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
    args.ParseCheck();
 
+   Device device(device_config);
+   device.Print();
+
    bool H1 = false, ND = false, RT = false, L2 = false;
    if (string(fe) == "h") { H1 = true; }
    else if (string(fe) == "n") { ND = true; }
    else if (string(fe) == "r") { RT = true; }
    else if (string(fe) == "l") { L2 = true; }
    else { MFEM_ABORT("Bad FE type. Must be 'h', 'n', 'r', or 'l'."); }
+   assert(H1);
 
    if (RT) { grad_div_problem = true; }
    double kappa = (order+1)*(order+1); // Penalty used for DG discretizations
@@ -125,7 +132,9 @@ int main(int argc, char *argv[])
    BilinearForm a(&fes);
    if (H1 || L2)
    {
-      a.AddDomainIntegrator(new MassIntegrator);
+      // #warning no MassIntegrator
+      // all warnings being treated as errors
+      //a.AddDomainIntegrator(new MassIntegrator);
       a.AddDomainIntegrator(new DiffusionIntegrator);
    }
    else
@@ -201,5 +210,14 @@ int main(int argc, char *argv[])
       dc.Save();
    }
 
+   if (visualization)
+   {
+      char vishost[] = "localhost";
+      int  visport   = 19916;
+      socketstream sol_sock(vishost, visport);
+      sol_sock.precision(8);
+      sol_sock << "solution\n" << mesh << x << flush;
+   }
+
    return 0;
 }

miniapps/solvers/plor_solvers.cpp

@@ -71,13 +71,17 @@ bool grad_div_problem = false;
 int main(int argc, char *argv[])
 {
    MPI_Session mpi;
+   const int num_procs = mpi.WorldSize();
+   const int myid = mpi.WorldRank();
 
    const char *mesh_file = "../../data/star.mesh";
    int ser_ref_levels = 1, par_ref_levels = 1;
    int order = 3;
    const char *fe = "h";
    const char *device_config = "cpu";
-   bool visualization = true;
+   bool visualization = false;
+   bool compute_L2_error = false;
+   int config_dev_modulo = 4;
 
    OptionsParser args(argc, argv);
    args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
@@ -91,11 +95,18 @@ int main(int argc, char *argv[])
    args.AddOption(&visualization, "-vis", "--visualization", "-no-vis",
                   "--no-visualization",
                   "Enable or disable GLVis visualization.");
+   args.AddOption(&compute_L2_error, "-l2", "--compute-L2-error", "-no-l2",
+                  "--no-compute-L2-error",
+                  "Enable or disable GLVis visualization.");
    args.AddOption(&device_config, "-d", "--device",
                   "Device configuration string, see Device::Configure().");
+   args.AddOption(&config_dev_modulo, "-dm", "--device-modulo",
+                  "Number of devices available on the node.");
    args.ParseCheck();
 
-   Device device(device_config);
+   const int dev = myid % config_dev_modulo;
+
+   Device device(device_config, dev);
    device.Print();
 
    bool H1 = false, ND = false, RT = false, L2 = false;
@@ -104,6 +115,7 @@ int main(int argc, char *argv[])
    else if (string(fe) == "r") { RT = true; }
    else if (string(fe) == "l") { L2 = true; }
    else { MFEM_ABORT("Bad FE type. Must be 'h', 'n', 'r', or 'l'."); }
+   assert(H1);
 
    if (RT) { grad_div_problem = true; }
    double kappa = (order+1)*(order+1); // Penalty used for DG discretizations
@@ -140,7 +152,9 @@ int main(int argc, char *argv[])
    ParBilinearForm a(&fes);
    if (H1 || L2)
    {
-      a.AddDomainIntegrator(new MassIntegrator);
+      // #warning no MassIntegrator
+      // all warnings being treated as errors
+      //a.AddDomainIntegrator(new MassIntegrator);
       a.AddDomainIntegrator(new DiffusionIntegrator);
    }
    else
@@ -187,10 +201,12 @@ int main(int argc, char *argv[])
    }
    else if (RT && dim == 3)
    {
+      assert(false);
       solv_lor.reset(new LORSolver<HypreADS>(lor, &fes_lor));
    }
    else
    {
+      assert(false);
       solv_lor.reset(new LORSolver<HypreAMS>(lor, &fes_lor));
    }
 
@@ -201,13 +217,18 @@ int main(int argc, char *argv[])
    cg.SetPrintLevel(1);
    cg.SetOperator(*A);
    cg.SetPreconditioner(*solv_lor);
-   cg.Mult(B, X);
+   {
+      cg.Mult(B, X);
+   }
 
    a.RecoverFEMSolution(X, b, x);
 
-   double er =
-      (H1 || L2) ? x.ComputeL2Error(u_coeff) : x.ComputeL2Error(u_vec_coeff);
-   if (mpi.Root()) { cout << "L2 error: " << er << endl; }
+   if (compute_L2_error)
+   {
+      double er =
+         (H1 || L2) ? x.ComputeL2Error(u_coeff) : x.ComputeL2Error(u_vec_coeff);
+      if (mpi.Root()) { cout << "L2 error: " << er << endl; }
+   }
 
    if (visualization)
    {
@@ -227,5 +248,15 @@ int main(int argc, char *argv[])
       dc.Save();
    }
 
+   if (visualization)
+   {
+      char vishost[] = "localhost";
+      int  visport   = 19916;
+      socketstream sol_sock(vishost, visport);
+      sol_sock << "parallel " << num_procs << " " << myid << "\n";
+      sol_sock.precision(8);
+      sol_sock << "solution\n" << mesh << x << flush;
+   }
+
    return 0;
 }

tests/benchmarks/CMakeLists.txt

@@ -47,6 +47,9 @@ endfunction(add_benchmark)
 #-------------------------------------------------------------------------------
 if (MFEM_USE_BENCHMARK)
     add_benchmark(ceed)
+    add_benchmark(solvers)
+    add_benchmark(linext)
+    add_benchmark(lor_assembly)
     add_benchmark(tmop)
     add_benchmark(vector)
     add_benchmark(virtuals)

tests/benchmarks/bench.hpp

@@ -13,12 +13,11 @@
 #define MFEM_TESTS_BENCH_HPP
 
 #include "mfem.hpp"
+using namespace mfem;
 
 #ifdef MFEM_USE_BENCHMARK
 
 #include "benchmark/benchmark.h"
-
-using namespace mfem;
 namespace bm = benchmark;
 namespace bmi = benchmark::internal;
 
@@ -26,15 +25,59 @@ namespace benchmark
 {
 namespace internal
 {
+
 extern std::map<std::string, std::string> *global_context;
-}
-}
+
+template<typename T>
+void FindInContext(const char *context, T &config)
+{
+   const auto found = bmi::global_context->find(context);
+   if (found != bmi::global_context->end()) { config = found->second; }
+};
+
+void FindInContext(const char *context, const char* &config)
+{
+   const auto found = bmi::global_context->find(context);
+   if (found != bmi::global_context->end()) { config = found->second.c_str(); }
+};
+
+void FindInContext(const char *context, bool &config)
+{
+   const auto found = bmi::global_context->find(context);
+   if (found != bmi::global_context->end())
+   { config = !strncmp(found->second.c_str(),"true",4); }
+};
+
+void FindInContext(const char *context, int &config)
+{
+   const auto found = bmi::global_context->find(context);
+   if (found != bmi::global_context->end())
+   { config = std::stoi(found->second.c_str()); }
+};
+
+} // namespace internal
+} // namespace benchmark
+
+#endif // MFEM_USE_BENCHMARK
 
 namespace mfem
 {
 
+template<class T>
+typename std::enable_if<!std::numeric_limits<T>::is_integer, bool>::type
+almost_equal(T x, T y, T tolerance = 1e-14)
+{
+   const T neg = std::abs(x - y);
+   constexpr T min = std::numeric_limits<T>::min();
+   constexpr T eps = std::numeric_limits<T>::epsilon();
+   const T min_abs = std::min(std::abs(x), std::abs(y));
+   if (std::abs(min_abs)==0.0) { return neg < eps; }
+   return (neg/std::max(min, min_abs)) < tolerance;
+}
+
 constexpr std::size_t KB = (1<<10);
 
+#ifdef MFEM_USE_BENCHMARK
 // Specific MFEM Reporter
 class Reporter : public benchmark::BenchmarkReporter
 {
@@ -66,9 +109,17 @@ public:
       }
    }
 };
+#endif // MFEM_USE_BENCHMARK
+
+struct NoReporter : public ::benchmark::BenchmarkReporter
+{
+   explicit NoReporter() {}
+   bool ReportContext(const Context &) { return true; }
+   void ReportRuns(const std::vector<Run> &) { }
+   operator NoReporter*() { return this; }
+   //void Finalize() {}
+};
 
 } // namespace mfem
 
-#endif // MFEM_USE_BENCHMARK
-
 #endif // MFEM_TESTS_BENCH_HPP

tests/benchmarks/bench_ceed.cpp

@@ -13,6 +13,9 @@
 
 #ifdef MFEM_USE_BENCHMARK
 
+namespace mfem
+{
+
 /*
   This benchmark contains the implementation of the CEED's bake-off problems:
   high-order kernels/benchmarks designed to test and compare the performance
@@ -23,7 +26,8 @@
 
 struct BakeOff
 {
-   const int N, p, q, dim = 3;
+   const int p, c, q, n, nx, ny, nz, dim = 3;
+   const bool check_x, check_y, check_z, checked;
    Mesh mesh;
    H1_FECollection fec;
    FiniteElementSpace fes;
@@ -37,11 +41,19 @@ struct BakeOff
    BilinearForm a;
    double mdofs;
 
-   BakeOff(int p, int vdim, bool GLL):
-      N(Device::IsEnabled()?32:8),
+   BakeOff(int p, int side, int vdim, bool gll):
       p(p),
-      q(2*p + (GLL?-1:3)),
-      mesh(Mesh::MakeCartesian3D(N,N,N,Element::HEXAHEDRON)),
+      c(side),
+      q(2*p + (gll?-1:3)),
+      n((assert(c>=p),c/p)),
+      nx(n + (p*(n+1)*p*n*p*n < c*c*c ?1:0)),
+      ny(n + (p*(n+1)*p*(n+1)*p*n < c*c*c ?1:0)),
+      nz(n),
+      check_x(p*nx * p*ny * p*nz <= c*c*c),
+      check_y(p*(nx+1) * p*(ny+1) * p*nz > c*c*c),
+      check_z(p*(nx+1) * p*(ny+1) * p*(nz+1) > c*c*c),
+      checked((assert(check_x && check_y && check_z), true)),
+      mesh(Mesh::MakeCartesian3D(nx,ny,nz,Element::HEXAHEDRON)),
       fec(p, dim, BasisType::GaussLobatto),
       fes(&mesh, &fec, vdim),
       geom_type(fes.GetFE(0)->GetGeomType()),
@@ -53,7 +65,7 @@ struct BakeOff
       x(&fes),
       y(&fes),
       a(&fes),
-      mdofs(0.0) {}
+      mdofs(0.0) { }
 
    virtual void benchmark() = 0;
 
@@ -77,8 +89,8 @@ struct Problem: public BakeOff
    Vector B, X;
    CGSolver cg;
 
-   Problem(int order):
-      BakeOff(order,VDIM,GLL),
+   Problem(int order, int side):
+      BakeOff(order,side,VDIM,GLL),
       ess_bdr(mesh.bdr_attributes.Max()),
       b(&fes)
    {
@@ -110,13 +122,34 @@ struct Problem: public BakeOff
    }
 };
 
+static void OrderSideArgs(bmi::Benchmark *b)
+{
+   const auto est = [](int c) { return (c+1)*(c+1)*(c+1); };
+   for (int p = 1; p <= 6; ++p)
+   {
+      for (int c = p; est(c) <= 2*1024*1024; c += 1)
+      {
+         if (c<10) { continue; }
+         b->Args({p, c});
+      }
+   }
+}
+
 /// Bake-off Problems (BPs)
-#define BakeOff_Problem(i,Kernel,VDIM,p_eq_q)\
+#define BakeOff_Problem(i,Kernel,VDIM,GLL)\
 static void BP##i(bm::State &state){\
-   Problem<Kernel##Integrator,VDIM,p_eq_q> ker(state.range(0));\
+   const int p = state.range(0);\
+   const int side = state.range(1);\
+   Problem<Kernel##Integrator,VDIM,GLL> ker(p,side);\
    while (state.KeepRunning()) { ker.benchmark(); }\
-   state.counters["MDof/s"] = bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);}\
-BENCHMARK(BP##i)->DenseRange(1,6)->Unit(bm::kMillisecond);
+   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(p);\
+}\
+BENCHMARK(BP##i)\
+       -> Apply(OrderSideArgs)\
+       -> Unit(bm::kMillisecond);
 
 /// BP1: scalar PCG with mass matrix, q=p+2
 BakeOff_Problem(1,Mass,1,false)
@@ -143,7 +176,7 @@ struct Kernel: public BakeOff
 {
    GridFunction y;
 
-   Kernel(int order): BakeOff(order,VDIM,GLL), y(&fes)
+   Kernel(int order, int side): BakeOff(order, side, VDIM, GLL), y(&fes)
    {
       x.Randomize(1);
       a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
@@ -164,10 +197,17 @@ struct Kernel: public BakeOff
 /// Generic CEED BKi
 #define BakeOff_Kernel(i,KER,VDIM,GLL)\
 static void BK##i(bm::State &state){\
-   Kernel<KER##Integrator,VDIM,GLL> ker(state.range(0));\
+   const int p = state.range(0);\
+   const int side = state.range(1);\
+   Kernel<KER##Integrator,VDIM,GLL> ker(p,side);\
    while (state.KeepRunning()) { ker.benchmark(); }\
-   state.counters["MDof/s"] = bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);}\
-BENCHMARK(BK##i)->DenseRange(1,6)->Unit(bm::kMillisecond);
+   state.counters["MDof/s"] = bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);\
+   state.counters["Dofs"] = bm::Counter(ker.dofs);\
+   state.counters["p"] = bm::Counter(p);\
+ }\
+ BENCHMARK(BK##i)\
+        -> Apply(OrderSideArgs)\
+        -> Unit(bm::kMillisecond);
 
 /// BK1: scalar E-vector-to-E-vector evaluation of mass matrix, q=p+2
 BakeOff_Kernel(1,Mass,1,false)
@@ -187,10 +227,15 @@ BakeOff_Kernel(5,Diffusion,1,true)
 /// BK6: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
 BakeOff_Kernel(6,VectorDiffusion,3,true)
 
+} // namespace mfem
+
 /**
  * @brief main entry point
- * --benchmark_filter=BK1/6
+ * --benchmark_filter=BP1/6
  * --benchmark_context=device=cpu
+ * --benchmark_out=bp1_cpu_fast.org
+ * --benchmark_out_format=csv
+ * --benchmark_min_time=2
  */
 int main(int argc, char *argv[])
 {
@@ -216,4 +261,8 @@ int main(int argc, char *argv[])
    return 0;
 }
 
+#else // MFEM_USE_BENCHMARK
+
+int main(int, char *[]) { return 0; }
+
 #endif // MFEM_USE_BENCHMARK

tests/benchmarks/bench_ceed.tex

@@ -0,0 +1,82 @@
+\documentclass[tikz]{standalone}
+%\documentclass[preview, border=4mm, convert={density=600,outext=.png}]{standalone}
+\usepackage{pgfplots}
+\pgfplotsset{compat=newest, table/col sep=comma}
+\usepackage{tikzscale}
+\usetikzlibrary{backgrounds}
+\usetikzlibrary{calc}
+
+\newcommand{\DATA}{bp1.org}
+\newcommand{\TITLE}{BP1}
+
+\newcommand{\XMIN}{1e3}
+\newcommand{\XMAX}{4e6}
+
+\newcommand{\YMIN}{0}
+\newcommand{\YMAX}{4}
+
+\usepackage{times}
+\usepackage{bm}
+\usepackage{fontspec}
+\setmainfont{Calibri}
+
+\definecolor{RYB2}{RGB}{31,120,180}   % Dark blue
+\definecolor{RYB4}{RGB}{51,160,44}    % Dark green
+\definecolor{RYB6}{RGB}{227,26,28}    % Dark red
+\definecolor{RYB8}{RGB}{255,127,0}    % Dark orange
+\definecolor{RYB10}{RGB}{160,60,140}  % Dark purple (modified)
+\definecolor{RYB12}{RGB}{177,89,40}   % Dark brown
+
+\definecolor{firebrick}{HTML}{B22222}
+\definecolor{steelblue4}{HTML}{3B77AF}
+
+\pgfplotsset{tick scale binop=\times}
+
+\pgfplotsset{every axis/.append style={
+axis line style={->},
+label style={font=\Large},
+tick label style={font=\large},
+title style={font=\Large\bfseries, text=steelblue4}
+}}
+
+\begin{document}
+\pgfplotscreateplotcyclelist{CycleList}{%
+RYB2!50!black,every mark/.append style={fill=RYB2},mark=*,smooth\\%
+RYB4!50!black,every mark/.append style={fill=RYB4},mark=*,smooth\\%
+RYB6!50!black,every mark/.append style={fill=RYB6},mark=*,smooth\\%
+RYB8!50!black,every mark/.append style={fill=RYB8},mark=*,smooth\\%
+RYB10!50!black,every mark/.append style={fill=RYB10},mark=*,smooth\\%
+RYB12!50!black,every mark/.append style={fill=RYB12},mark=*,smooth\\%
+}
+
+\begin{tikzpicture}[background rectangle/.style={fill=white},
+show background rectangle, scale=4.0]
+\begin{semilogxaxis}[
+    grid=both,
+    major grid style={line width=.1pt,draw=gray!50},
+    minor grid style={line width=.1pt,draw=gray!50},
+    width=4in,
+    xlabel={Degrees of Freedom (DOF)},
+    xlabel style={at={(axis description cs:0.5,-0.05)}},
+    ylabel={GDOF/s},
+    ylabel style={at={(axis description cs:-0.025,0.5)}},
+    cycle list name=CycleList,
+    %legend cell align=left,
+    legend pos=north west,
+    legend entries={$p=1$,$p=2$,$p=3$,$p=4$,$p=5$,$p=6$},
+    title=\TITLE,
+    ymin=\YMIN,
+    ymax=\YMAX,
+    xmin=\XMIN,
+    xmax=\XMAX
+  ]
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($11):($13==1?$12/1e3:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($11):($13==2?$12/1e3:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($11):($13==3?$12/1e3:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($11):($13==4?$12/1e3:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($11):($13==5?$12/1e3:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($11):($13==6?$12/1e3:NaN)};
+
+\end{semilogxaxis}
+\end{tikzpicture}
+\end{document}

tests/benchmarks/bench_linext.cpp

@@ -0,0 +1,217 @@
+// Copyright (c) 2010-2021, 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"
+
+#ifdef MFEM_USE_BENCHMARK
+
+#include "tests/unit/fem/test_linearform_ext.hpp"
+
+namespace mfem
+{
+
+namespace linearform_ext_tests
+{
+
+void LinearFormExtTest::Run() { MFEM_ABORT("Should use the virtuals!"); }
+void LinearFormExtTest::Description() { MFEM_ABORT("Should use the virtuals!"); }
+
+////////////////////////////////////////////////////////////////////////////////
+/// TEST for LinearFormExtension
+struct Test: public LinearFormExtTest
+{
+   Test(int N, int dim, int vdim, int ordering, bool gll, int problem, int order):
+      LinearFormExtTest(N, dim, vdim, ordering, gll,
+                        problem, order,
+                        true) { }
+
+   void Description() override { /* */ }
+
+   void Run() override
+   {
+      AssembleBoth();
+      MFEM_DEVICE_SYNC;
+      const double tolerance = 1e-13;
+      const double dtd = lf_full*lf_full;
+      const double rtr = lf_legacy*lf_legacy;
+      const bool almost_eq = almost_equal(dtd, rtr, tolerance);
+      MFEM_VERIFY(almost_eq, "almost_equal test error!");
+   }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+constexpr int _2D = 2;
+constexpr int _3D = 3;
+constexpr int _GL = false; // Gauss-Legendre, q=p+2
+constexpr int _GLL = true; // Gauss-Legendre-Lobatto, q=p+1
+constexpr int VDIM = 24;
+
+/// Scalar Linear Form Extension Tests
+#define LinExtTest(Problem,dim,vdim,gll)\
+static void TEST_##Problem##dim##gll(bm::State &state){\
+   const int p = state.range(0);\
+   Test ker(4, dim,vdim,Ordering::byVDIM,gll,LinearFormExtTest::Problem,p);\
+   while(state.KeepRunning()) { ker.Run(); }\
+   state.counters["MDof/s"] = bm::Counter(ker.SumMdofs(),bm::Counter::kIsRate);}\
+BENCHMARK(TEST_##Problem##dim##gll)->DenseRange(1,6)->Unit(bm::kMillisecond);
+
+/// Scalar linear form tests & Gauss-Legendre-Lobatto, q=p+1
+LinExtTest(DomainLF,_2D,1,_GLL)
+LinExtTest(DomainLF,_3D,1,_GLL)
+
+/// Vector linear form tests & Gauss-Legendre-Lobatto, q=p+1
+LinExtTest(VectorDomainLF,_2D,VDIM,_GLL)
+LinExtTest(VectorDomainLF,_3D,VDIM,_GLL)
+
+/// Grad linear form tests & Gauss-Legendre-Lobatto, q=p+1
+LinExtTest(DomainLFGrad,_2D,1,_GLL)
+LinExtTest(DomainLFGrad,_3D,1,_GLL)
+
+/// Vector Grad linear form tests & Gauss-Legendre-Lobatto, q=p+1
+LinExtTest(VectorDomainLFGrad,_2D,VDIM,_GLL)
+LinExtTest(VectorDomainLFGrad,_3D,VDIM,_GLL)
+
+/// Scalar linear form tests & Gauss-Legendre, q=p+2
+LinExtTest(DomainLF,_2D,1,_GL)
+LinExtTest(DomainLF,_3D,1,_GL)
+
+/// Vector linear form tests & Gauss-Legendre, q=p+2
+LinExtTest(VectorDomainLF,_2D,VDIM,_GL)
+LinExtTest(VectorDomainLF,_3D,VDIM,_GL)
+
+/// Grad linear form tests & Gauss-Legendre, q=p+2
+LinExtTest(DomainLFGrad,_2D,1,_GL)
+LinExtTest(DomainLFGrad,_3D,1,_GL)
+
+/// Vector Grad linear form tests & Gauss-Legendre, q=p+2
+LinExtTest(VectorDomainLFGrad,_2D,VDIM,_GL)
+LinExtTest(VectorDomainLFGrad,_3D,VDIM,_GL)
+
+
+////////////////////////////////////////////////////////////////////////////////
+/// BENCH for LinearFormExtension
+template<enum LinearAssemblyLevel LAL>
+struct Bench: public LinearFormExtTest
+{
+   Bench(int dim, int vdim, int ordering, bool gll, int problem, int p):
+      LinearFormExtTest(Device::IsEnabled()?24:4,
+                        dim, vdim, ordering, gll,
+                        problem, p,
+                        false)
+   { }
+
+   void Description() override { /* */ }
+
+   void Run() override
+   {
+      MFEM_DEVICE_SYNC;
+      this->mdofs += this->MDofs();
+      if (LAL==LinearAssemblyLevel::FULL) { lf_full.Assemble(); }
+      if (LAL==LinearAssemblyLevel::LEGACY) { lf_legacy.Assemble(); }
+   }
+};
+
+/// Linear Form Extension Scalar Benchs
+#define LinExtBench(Problem,lal,dim,vdim,gll)\
+static void BENCH_##lal##_##Problem##dim##gll(bm::State &state){\
+   const int p = state.range(0);\
+   Bench<LinearAssemblyLevel::lal> ker(dim,vdim,Ordering::byVDIM,gll,LinearFormExtTest::Problem, p);\
+   while (state.KeepRunning()) { ker.Run(); }\
+   state.counters["MDof/s"] = bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);}\
+BENCHMARK(BENCH_##lal##_##Problem##dim##gll)->DenseRange(1,6)->Unit(bm::kMicrosecond);
+
+/// Scalar linear form bench & Gauss-Legendre-Lobatto, q=p+1
+LinExtBench(DomainLF,LEGACY,_2D,1,_GLL)
+LinExtBench(DomainLF,  FULL,_2D,1,_GLL)
+LinExtBench(DomainLF,LEGACY,_3D,1,_GLL)
+LinExtBench(DomainLF,  FULL,_3D,1,_GLL)
+
+/// Vector linear form bench & Gauss-Legendre-Lobatto, q=p+1
+LinExtBench(VectorDomainLF,LEGACY,_2D,VDIM,_GLL)
+LinExtBench(VectorDomainLF,  FULL,_2D,VDIM,_GLL)
+LinExtBench(VectorDomainLF,LEGACY,_3D,VDIM,_GLL)
+LinExtBench(VectorDomainLF,  FULL,_3D,VDIM,_GLL)
+
+/// Grad Scalar linear form bench & Gauss-Legendre-Lobatto, q=p+1
+LinExtBench(DomainLFGrad,LEGACY,_2D,1,_GLL)
+LinExtBench(DomainLFGrad,  FULL,_2D,1,_GLL)
+LinExtBench(DomainLFGrad,LEGACY,_3D,1,_GLL)
+LinExtBench(DomainLFGrad,  FULL,_3D,1,_GLL)
+
+/// Vector Grad linear form bench & Gauss-Legendre-Lobatto, q=p+1
+LinExtBench(VectorDomainLFGrad,LEGACY,_2D,VDIM,_GLL)
+LinExtBench(VectorDomainLFGrad,  FULL,_2D,VDIM,_GLL)
+LinExtBench(VectorDomainLFGrad,LEGACY,_3D,VDIM,_GLL)
+LinExtBench(VectorDomainLFGrad,  FULL,_3D,VDIM,_GLL)
+
+/// Scalar linear form bench & Gauss-Legendre, q=p+2
+LinExtBench(DomainLF,LEGACY,_2D,1,_GL)
+LinExtBench(DomainLF,  FULL,_2D,1,_GL)
+LinExtBench(DomainLF,LEGACY,_3D,1,_GL)
+LinExtBench(DomainLF,  FULL,_3D,1,_GL)
+
+/// Vector linear form bench & Gauss-Legendre, q=p+2
+LinExtBench(VectorDomainLF,LEGACY,_2D,VDIM,_GL)
+LinExtBench(VectorDomainLF,  FULL,_2D,VDIM,_GL)
+LinExtBench(VectorDomainLF,LEGACY,_3D,VDIM,_GL)
+LinExtBench(VectorDomainLF,  FULL,_3D,VDIM,_GL)
+
+/// Grad Scalar linear form bench & Gauss-Legendre, q=p+2
+LinExtBench(DomainLFGrad,LEGACY,_2D,1,_GL)
+LinExtBench(DomainLFGrad,  FULL,_2D,1,_GL)
+LinExtBench(DomainLFGrad,LEGACY,_3D,1,_GL)
+LinExtBench(DomainLFGrad,  FULL,_3D,1,_GL)
+
+/// Vector Grad linear form bench & Gauss-Legendre, q=p+2
+LinExtBench(VectorDomainLFGrad,LEGACY,_2D,VDIM,_GL)
+LinExtBench(VectorDomainLFGrad,  FULL,_2D,VDIM,_GL)
+LinExtBench(VectorDomainLFGrad,LEGACY,_3D,VDIM,_GL)
+LinExtBench(VectorDomainLFGrad,  FULL,_3D,VDIM,_GL)
+
+} // namespace linearform_ext_tests
+
+} // namespace mfem
+
+/** ****************************************************************************
+ * @brief main entry point, some options are for example:
+ * --benchmark_filter=TEST --benchmark_min_time=0.01
+ * --benchmark_filter=BENCH_FULL --benchmark_min_time=0.1
+ * --benchmark_context=device=cuda
+ **************************************************************************** */
+int main(int argc, char *argv[])
+{
+   bm::ConsoleReporter CR;
+   bm::Initialize(&argc, argv);
+
+   // Device setup, CPU by default
+   std::string device_config = "cpu";
+   if (bmi::global_context != nullptr)
+   {
+      const auto device = bmi::global_context->find("device");
+      if (device != bmi::global_context->end())
+      {
+         mfem::out << device->first << " : " << device->second << std::endl;
+         device_config = device->second;
+      }
+   }
+   Device device(device_config.c_str());
+   device.Print();
+
+   if (bm::ReportUnrecognizedArguments(argc, argv)) { return 1; }
+   bm::RunSpecifiedBenchmarks(&CR);
+   return 0;
+}
+
+#else // MFEM_USE_BENCHMARK
+
+int main(int, char *[]) { return 0; }
+
+#endif // MFEM_USE_BENCHMARK

tests/benchmarks/bench_lor_assembly.cpp

@@ -0,0 +1,368 @@
+// Copyright (c) 2010-2021, 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"
+
+#ifdef MFEM_USE_BENCHMARK
+
+#include "fem/lor.hpp"
+#include "fem/lor_assembly.hpp"
+
+#define MFEM_DEBUG_COLOR 119
+#include "general/debug.hpp"
+
+#include <cassert>
+#include <cmath>
+
+constexpr int SEED = 0x100001b3;
+
+struct LORBench
+{
+   const int p, c, q, n, nx, ny, nz, dim = 3;
+   const bool check_x, check_y, check_z, checked;
+   Mesh mesh;
+   H1_FECollection fec;
+   FiniteElementSpace mfes, fes_ho;
+   Array<int> ess_bdr_ho, ess_dofs_ho;
+   LORDiscretization lor_disc;
+   IntegrationRules irs;
+   const IntegrationRule &ir_el;
+   FiniteElementSpace &fes_lo;
+   Array<int> ess_bdr_lo, ess_dofs_lo;
+   BilinearForm a_legacy, a_full;
+   OperatorHandle A_batched, A_deviced;
+   SparseMatrix *A_full;
+   GridFunction x;
+   const int dofs;
+   double mdof;
+
+   LORBench(int p, int side):
+      p(p),
+      c(side),
+      q(2*p + 2),
+      n((assert(c>=p),c/p)),
+      nx(n + (p*(n+1)*p*n*p*n < c*c*c ?1:0)),
+      ny(n + (p*(n+1)*p*(n+1)*p*n < c*c*c ?1:0)),
+      nz(n),
+      check_x(p*nx * p*ny * p*nz <= c*c*c),
+      check_y(p*(nx+1) * p*(ny+1) * p*nz > c*c*c),
+      check_z(p*(nx+1) * p*(ny+1) * p*(nz+1) > c*c*c),
+      checked((assert(check_x && check_y && check_z), true)),
+      mesh(Mesh::MakeCartesian3D(nx,ny,nz, Element::HEXAHEDRON)),
+      fec(p, dim, BasisType::GaussLobatto),
+      mfes(&mesh, &fec, dim),
+      fes_ho(&mesh, &fec),
+      ess_bdr_ho(mesh.bdr_attributes.Max()),
+      lor_disc(fes_ho, BasisType::GaussLobatto),
+      irs(0, Quadrature1D::GaussLobatto),
+      ir_el(irs.Get(Geometry::Type::CUBE, 1)),
+      fes_lo(lor_disc.GetFESpace()),
+      ess_bdr_lo(fes_lo.GetMesh()->bdr_attributes.Max()),
+      a_legacy(&fes_lo),
+      a_full(&fes_lo),
+      A_deviced(),
+      A_full(nullptr),
+      x(&mfes),
+      dofs(fes_ho.GetVSize()),
+      mdof(0.0)
+   {
+      dbg("p:%d side:%d dofs:%d/%d",p,side,dofs, fes_lo.GetVSize());
+      a_legacy.AddDomainIntegrator(new DiffusionIntegrator(&ir_el));
+      a_legacy.SetAssemblyLevel(AssemblyLevel::LEGACY);
+
+      a_full.AddDomainIntegrator(new DiffusionIntegrator(&ir_el));
+      a_full.SetAssemblyLevel(AssemblyLevel::FULL);
+
+      SetupRandomMesh();
+      // Make sure that SetCurvature is called on the LOR mesh
+      fes_lo.GetMesh()->EnsureNodes();
+
+      ess_bdr_ho = 1;
+      fes_ho.GetEssentialTrueDofs(ess_bdr_ho, ess_dofs_ho);
+
+      ess_bdr_lo = 1;
+      fes_lo.GetEssentialTrueDofs(ess_bdr_lo, ess_dofs_lo);
+
+      tic_toc.Clear();
+   }
+
+   void SetupRandomMesh() noexcept
+   {
+      mesh.SetNodalFESpace(&mfes);
+      mesh.SetNodalGridFunction(&x);
+      const double jitter = 1./(M_PI*M_PI);
+      const double h0 = mesh.GetElementSize(0);
+      GridFunction rdm(&mfes);
+      rdm.Randomize(SEED);
+      rdm -= 0.5; // Shift to random values in [-0.5,0.5]
+      rdm *= jitter * h0; // Scale the random values to be of same order
+      x -= rdm;
+   }
+
+   void SanityChecks()
+   {
+      dbg();
+      constexpr double EPS = 1e-15;
+
+      Vector x(dofs), y(dofs);
+      x.Randomize(SEED);
+      y.Randomize(SEED);
+
+      OperatorHandle A_legacy, A_full, A_deviced;
+
+      BilinearForm a_legacy(&fes_lo);
+      a_legacy.AddDomainIntegrator(new DiffusionIntegrator(&ir_el));
+      a_legacy.SetAssemblyLevel(AssemblyLevel::LEGACY);
+      MFEM_DEVICE_SYNC;
+      tic();
+      a_legacy.Assemble();
+      MFEM_DEVICE_SYNC;
+      dbg(" Legacy time = %f",toc());
+      a_legacy.FormSystemMatrix(ess_dofs_lo, A_legacy);
+      a_legacy.Finalize();
+      A_legacy.As<SparseMatrix>()->HostReadWriteI();
+      A_legacy.As<SparseMatrix>()->HostReadWriteJ();
+      A_legacy.As<SparseMatrix>()->HostReadWriteData();
+      const double dot_legacy = A_legacy.As<SparseMatrix>()->InnerProduct(x,y);
+
+      MFEM_DEVICE_SYNC;
+      tic();
+      a_full.Assemble();
+      MFEM_DEVICE_SYNC;
+      dbg("   Full time = %f",toc());
+      a_full.FormSystemMatrix(ess_dofs_lo, A_full); /// BC NOT DONE !!!
+      constexpr bool still_have_to_remove_the_bc = true;
+      if (still_have_to_remove_the_bc)
+      {
+         a_full.EliminateVDofs(ess_dofs_lo, Operator::DIAG_KEEP);
+         a_full.Finalize();
+      }
+      a_full.SpMat().HostReadWriteI();
+      a_full.SpMat().HostReadWriteJ();
+      a_full.SpMat().HostReadWriteData();
+      const double dot_full = a_full.SpMat().InnerProduct(x,y);
+      MFEM_VERIFY(almost_equal(dot_legacy, dot_full), "dot_full error!");
+      a_full.SpMat().Add(-1.0, *A_legacy.As<SparseMatrix>());
+      const double max_norm_full = a_full.SpMat().MaxNorm();
+      MFEM_VERIFY(max_norm_full < EPS, "max_norm_full error!");
+
+      MFEM_DEVICE_SYNC;
+      tic();
+      AssembleBatchedLOR(lor_disc, a_legacy, fes_ho, ess_dofs_lo, A_deviced);
+      MFEM_DEVICE_SYNC;
+      dbg("Deviced time = %f",toc());
+      A_deviced.As<SparseMatrix>()->HostReadWriteI();
+      A_deviced.As<SparseMatrix>()->HostReadWriteJ();
+      A_deviced.As<SparseMatrix>()->HostReadWriteData();
+      const double dot_device = A_deviced.As<SparseMatrix>()->InnerProduct(x,y);
+      MFEM_VERIFY(almost_equal(dot_legacy, dot_device), "dot_device error!");
+      A_deviced.As<SparseMatrix>()->Add(-1.0, *A_legacy.As<SparseMatrix>());
+      const double max_norm_deviced = A_deviced.As<SparseMatrix>()->MaxNorm();
+      MFEM_VERIFY(max_norm_deviced < EPS, "max_norm_deviced");
+   }
+
+   void GLVis(Mesh &mesh, GridFunction *x = nullptr)
+   {
+      char vishost[] = "localhost";
+      int  visport   = 19916;
+      socketstream sol_sock(vishost, visport);
+      sol_sock.precision(8);
+      if (!x) { sol_sock << "mesh\n" << mesh; }
+      else { sol_sock << "solution\n" << mesh << *x ;}
+      sol_sock << std::flush;
+   }
+
+   void Test()
+   {
+      MFEM_DEVICE_SYNC;
+      tic();
+      AssembleBatchedLOR(lor_disc, a_legacy, fes_ho, ess_dofs_lo, A_deviced);
+      MFEM_DEVICE_SYNC;
+      dbg(" Deviced time = %f",toc());
+      A_deviced.Clear(); // forcing initialization phase
+      dbg("Exiting!");
+      std::exit(0);
+   }
+
+   void Dump()
+   {
+      OperatorHandle A_legacy;
+
+      MFEM_DEVICE_SYNC;
+      a_legacy.Assemble();
+      MFEM_DEVICE_SYNC;
+      a_legacy.FormSystemMatrix(ess_dofs_lo, A_legacy);
+      a_legacy.Finalize();
+      A_legacy.As<SparseMatrix>()->HostReadWriteI();
+      A_legacy.As<SparseMatrix>()->HostReadWriteJ();
+      A_legacy.As<SparseMatrix>()->HostReadWriteData();
+
+      dbg("Saving 'A.mtx' file");
+      {
+         std::ofstream mtx_file("A.mtx");
+         A_legacy.As<SparseMatrix>()->PrintMM(mtx_file);
+      }
+
+      dbg("fes_lo.GetVSize: %d", fes_lo.GetVSize());
+      Mesh *mesh_lo =fes_lo.GetMesh();
+      GridFunction ids_lo(&fes_lo);
+      Array<int> dofs(fes_lo.GetVSize());
+      fes_lo.GetVDofs(0,dofs);
+      assert(fes_lo.GetVDim()==1);
+      for (int i=0; i<fes_lo.GetVSize(); i++) { ids_lo(i) = dofs[i]; }
+      GLVis(*mesh_lo, &ids_lo);
+      dbg("LO GLVis done!");
+
+      dbg("fes_ho.GetVSize: %d", fes_ho.GetVSize());
+      GridFunction ids(&fes_ho);
+      Array<int> vdofs(fes_ho.GetVSize());
+      fes_ho.GetVDofs(0, vdofs);
+      assert(fes_ho.GetVDim()==1);
+      for (int i=0; i<fes_ho.GetVSize(); i++) { ids(i) = vdofs[i]; }
+      GLVis(mesh, &ids);
+      dbg("HO GLVis done!");
+
+      dbg("Exiting!");
+      std::exit(0);
+   }
+
+   void KerLegacy()
+   {
+      MFEM_DEVICE_SYNC;
+      tic_toc.Start();
+      a_legacy.Assemble();
+      MFEM_DEVICE_SYNC;
+      tic_toc.Stop();
+      mdof += 1e-6 * dofs;
+   }
+
+   void KerFull()
+   {
+      MFEM_DEVICE_SYNC;
+      tic_toc.Start();
+      a_full.Assemble();
+      MFEM_DEVICE_SYNC;
+      tic_toc.Stop();
+      mdof += 1e-6 * dofs;
+   }
+
+   void KerDeviced()
+   {
+      MFEM_DEVICE_SYNC;
+      tic_toc.Start();
+      AssembleBatchedLOR(lor_disc, a_legacy, fes_ho, ess_dofs_lo, A_deviced);
+      MFEM_DEVICE_SYNC;
+      tic_toc.Stop();
+      mdof += 1e-6 * dofs;
+   }
+
+   void AllFull()
+   {
+      LORDiscretization lor_disc(fes_ho, BasisType::GaussLobatto);
+      FiniteElementSpace fes_lo(lor_disc.GetFESpace());
+      BilinearForm bf_full(&fes_lo);
+      bf_full.AddDomainIntegrator(new DiffusionIntegrator(&ir_el));
+      bf_full.SetAssemblyLevel(AssemblyLevel::FULL);
+      MFEM_DEVICE_SYNC;
+      tic_toc.Start();
+      bf_full.Assemble();
+      MFEM_DEVICE_SYNC;
+      tic_toc.Stop();
+      mdof += 1e-6 * dofs;
+   }
+
+   void AllDeviced()
+   {
+      LORDiscretization lor_disc(fes_ho, BasisType::GaussLobatto);
+      FiniteElementSpace fes_lo(lor_disc.GetFESpace());
+      BilinearForm a_legacy(&fes_lo);
+      a_legacy.AddDomainIntegrator(new DiffusionIntegrator(&ir_el));
+      a_legacy.SetAssemblyLevel(AssemblyLevel::LEGACY);
+      MFEM_DEVICE_SYNC;
+      tic_toc.Start();
+      OperatorHandle A_deviced;
+      AssembleBatchedLOR(lor_disc, a_legacy, fes_ho, ess_dofs_lo, A_deviced);
+      MFEM_DEVICE_SYNC;
+      tic_toc.Stop();
+      mdof += 1e-6 * dofs;
+   }
+
+   double Mdofs() const { return mdof / tic_toc.RealTime(); }
+};
+
+// The different orders the tests can run
+#define P_ORDERS bm::CreateDenseRange(1,8,1)
+
+// The different sides of the mesh
+#define N_SIDES bm::CreateDenseRange(2,20,1)
+#define MAX_NDOFS 2*1024*1024
+
+/// Kernels definitions and registrations
+#define Benchmark(Name)\
+static void Name(bm::State &state){\
+   const int p = state.range(0);\
+   const int side = state.range(1);\
+   LORBench lor(p, side);\
+   if (lor.dofs > MAX_NDOFS) { state.SkipWithError("MAX_NDOFS"); }\
+   while (state.KeepRunning()) { lor.Name(); }\
+   bm::Counter::Flags flags = bm::Counter::kIsIterationInvariantRate;\
+   state.counters["Ker_(Dofs/s)"] = bm::Counter(lor.dofs, flags);\
+   state.counters["All_(MDof/s)"] = bm::Counter(lor.Mdofs());\
+   state.counters["dofs"] = bm::Counter(lor.dofs);\
+   state.counters["p"] = bm::Counter(p);\
+}\
+BENCHMARK(Name)\
+            -> ArgsProduct({P_ORDERS,N_SIDES})\
+            -> Unit(bm::kMillisecond);
+
+Benchmark(SanityChecks)
+
+Benchmark(KerLegacy)
+Benchmark(KerFull)
+Benchmark(KerDeviced)
+
+Benchmark(AllFull)
+Benchmark(AllDeviced)
+
+Benchmark(Dump)
+Benchmark(Test)
+
+/**
+ * @brief main entry point
+ * --benchmark_filter=Batched/4/16
+ * --benchmark_filter=\(Batched\|Deviced\|Full\)/4/16
+ * --benchmark_context=device=cuda
+ */
+int main(int argc, char *argv[])
+{
+   bm::ConsoleReporter CR;
+   bm::Initialize(&argc, argv);
+
+   // Device setup, cpu by default
+   std::string device_config = "cpu";
+   if (bmi::global_context != nullptr)
+   {
+      const auto device = bmi::global_context->find("device");
+      if (device != bmi::global_context->end())
+      {
+         mfem::out << device->first << " : " << device->second << std::endl;
+         device_config = device->second;
+      }
+   }
+   Device device(device_config.c_str());
+   device.Print();
+
+   if (bm::ReportUnrecognizedArguments(argc, argv)) { return 1; }
+   bm::RunSpecifiedBenchmarks(&CR);
+   return 0;
+}
+
+#endif // MFEM_USE_BENCHMARK

tests/benchmarks/bench_solvers.tex

@@ -0,0 +1,81 @@
+\documentclass[tikz]{standalone}
+%\documentclass[preview, border=4mm, convert={density=600,outext=.png}]{standalone}
+\usepackage{pgfplots}
+\pgfplotsset{compat=newest, table/col sep=comma}
+\usepackage{tikzscale}
+\usetikzlibrary{backgrounds}
+\usetikzlibrary{calc}
+
+\newcommand{\DATA}{bps3_sawmill.org}
+\newcommand{\TITLE}{BPS3-MFEM results}
+
+\newcommand{\XMIN}{1e3}
+\newcommand{\XMAX}{4e6}
+
+\newcommand{\YMIN}{0}
+\newcommand{\YMAX}{200}
+
+\usepackage{times}
+\usepackage{bm}
+\usepackage{fontspec}
+\setmainfont{Calibri}
+
+\definecolor{RYB2}{RGB}{31,120,180}   % Dark blue
+\definecolor{RYB4}{RGB}{51,160,44}    % Dark green
+\definecolor{RYB6}{RGB}{227,26,28}    % Dark red
+\definecolor{RYB8}{RGB}{255,127,0}    % Dark orange
+\definecolor{RYB10}{RGB}{160,60,140}  % Dark purple (modified)
+\definecolor{RYB12}{RGB}{177,89,40}   % Dark brown
+
+\definecolor{firebrick}{HTML}{B22222}
+\definecolor{steelblue4}{HTML}{3B77AF}
+
+\pgfplotsset{tick scale binop=\times}
+
+\pgfplotsset{every axis/.append style={
+axis line style={->},
+label style={font=\Large},
+tick label style={font=\large},
+title style={font=\Large\bfseries, text=steelblue4}
+}}
+
+\begin{document}
+\pgfplotscreateplotcyclelist{CycleList}{%
+RYB2!50!black,every mark/.append style={fill=RYB2},mark=*,smooth\\%
+RYB4!50!black,every mark/.append style={fill=RYB4},mark=*,smooth\\%
+RYB6!50!black,every mark/.append style={fill=RYB6},mark=*,smooth\\%
+RYB8!50!black,every mark/.append style={fill=RYB8},mark=*,smooth\\%
+RYB10!50!black,every mark/.append style={fill=RYB10},mark=*,smooth\\%
+RYB12!50!black,every mark/.append style={fill=RYB12},mark=*,smooth\\%
+}
+
+\begin{tikzpicture}[background rectangle/.style={fill=white},
+show background rectangle, scale=4.0]
+\begin{semilogxaxis}[
+    grid=both,
+    major grid style={line width=.1pt,draw=gray!50},
+    minor grid style={line width=.1pt,draw=gray!50},
+    width=4in,
+    xlabel={\#DOFs},
+    xlabel style={at={(axis description cs:0.5,-0.1)}},
+    ylabel={Iterations},
+    ylabel style={at={(axis description cs:-0.1,0.5)}},
+    cycle list name=CycleList,
+    %legend cell align=left,
+    legend pos=north west,
+    legend entries={$p=1$,$p=2$,$p=3$,$p=4$,$p=5$,$p=6$},
+    title=\TITLE,
+    ymin=\YMIN,
+    ymax=\YMAX,
+    xmin=\XMIN,
+    xmax=\XMAX
+  ]
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($13):($15==1?$14:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($13):($15==2?$14:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($13):($15==3?$14:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($13):($15==4?$14:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($13):($15==5?$14:NaN)};
+  \addplot +[raw gnuplot] gnuplot { set datafile separator ','; plot "\DATA" using ($13):($15==6?$14:NaN)};
+\end{semilogxaxis}
+\end{tikzpicture}
+\end{document}

tests/benchmarks/bench_tmop.cpp

@@ -18,6 +18,9 @@
 #include <memory>
 #include <cmath>
 
+namespace mfem
+{
+
 struct TMOP
 {
    const int N, p, q, dim = 3;
@@ -164,6 +167,8 @@ BENCHMARK_TMOP(AddMultGradPA)
 BENCHMARK_TMOP(GetLocalStateEnergyPA)
 BENCHMARK_TMOP(AssembleGradDiagonalPA)
 
+} // namespace mfem
+
 /**
  * @brief main entry point
  * --benchmark_filter=AddMultPA/4
@@ -193,4 +198,8 @@ int main(int argc, char *argv[])
    return 0;
 }
 
+#else // MFEM_USE_BENCHMARK
+
+int main(int, char *[]) { return 0; }
+
 #endif // MFEM_USE_BENCHMARK

tests/benchmarks/bench_vector.cpp

@@ -222,4 +222,8 @@ int main(int argc, char *argv[])
    return 0;
 }
 
+#else // MFEM_USE_BENCHMARK
+
+int main(int, char *[]) { return 0; }
+
 #endif // MFEM_USE_BENCHMARK

tests/benchmarks/bench_virtuals.cpp

@@ -86,4 +86,8 @@ int main(int argc, char *argv[])
    return 0;
 }
 
+#else // MFEM_USE_BENCHMARK
+
+int main(int, char *[]) { return 0; }
+
 #endif // MFEM_USE_BENCHMARK

tests/benchmarks/makefile

@@ -12,13 +12,13 @@
 # Use the MFEM build directory
 MFEM_DIR ?= ../..
 MFEM_BUILD_DIR ?= ../..
-SRC = $(if $(MFEM_DIR:../../..=),$(MFEM_DIR)/tests/benchmarks/,)
+SRC = $(if $(MFEM_DIR:../..=),$(MFEM_DIR)/tests/benchmarks/,)
 CONFIG_MK = $(MFEM_BUILD_DIR)/config/config.mk
 
 MFEM_LIB_FILE = mfem_is_not_built
 -include $(CONFIG_MK)
 
-SEQ_TESTS = bench_ceed bench_tmop bench_vector bench_virtuals
+SEQ_TESTS = bench_ceed bench_linext bench_solvers bench_tmop bench_vector bench_virtuals
 PAR_TESTS = 
 ifeq ($(MFEM_USE_MPI),NO)
    TESTS = $(SEQ_TESTS)
@@ -39,7 +39,7 @@ all: $(TESTS)
 
 # Rules for building the TESTS
 
-%: $(SRC)%.cpp $(MFEM_LIB_FILE) $(CONFIG_MK)
+%: $(SRC)%.cpp bench.hpp $(MFEM_LIB_FILE) $(CONFIG_MK)
 	$(MFEM_CXX) $(MFEM_FLAGS) $< -o $@ $(MFEM_LIBS)
 
 # Rules for compiling miniapp dependencies

tests/unit/fem/test_linearform_ext.cpp

@@ -0,0 +1,87 @@
+// Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#include "unit_tests.hpp"
+
+#include "fem/test_linearform_ext.hpp"
+
+using namespace mfem;
+using namespace linearform_ext_tests;
+
+namespace mfem
+{
+
+namespace linearform_ext_tests
+{
+
+void LinearFormExtTest::Description()
+{
+   const bool scalar = problem == LinearFormExtTest::DomainLF ||
+                       problem == LinearFormExtTest::DomainLFGrad;
+   if (scalar) { MFEM_VERIFY(vdim == 1, "VDIM should be 1"); }
+   const bool grad = problem == LinearFormExtTest::DomainLFGrad ||
+                     problem == LinearFormExtTest::VectorDomainLFGrad;
+
+   mfem::out << "[LinearFormExt]"
+             << " p=" << p
+             << " q=" << q
+             << (ordering==Ordering::byNODES ? " byNODES" : " byVDIM ")
+             << " "<< dim << "D"
+             << " "<< vdim << "-"
+             << (scalar ? "Scalar" : "Vector")
+             << (grad ? "Grad" : "")
+             << std::endl;
+}
+
+void LinearFormExtTest::Run()
+{
+   Description();
+
+   AssembleBoth();
+
+   // Test the difference to verify the orderings
+   Vector difference = lf_legacy;
+   difference -= lf_full;
+   REQUIRE(0.0 == MFEM_Approx(difference * difference));
+
+   REQUIRE(lf_full * lf_full == MFEM_Approx(lf_legacy * lf_legacy));
+}
+
+} // namespace linearform_ext_tests
+
+} // namespace mfem
+
+TEST_CASE("Linear Form Extension", "[LinearformExt], [CUDA]")
+{
+   const auto N = GENERATE(3,4);
+   const auto p = GENERATE(1,3,6); // limitations: 2D:11, 3D:6
+   const auto dim = GENERATE(2,3);
+   const auto gll = GENERATE(false,true); // q=p+2, q=p+1
+
+   SECTION("Scalar")
+   {
+      const auto vdim = 1;
+      const auto ordering = Ordering::byNODES;
+      const auto problem = GENERATE(LinearFormExtTest::DomainLF,
+                                    LinearFormExtTest::DomainLFGrad);
+      LinearFormExtTest(N, dim, vdim, ordering, gll, problem, p, true).Run();
+   }
+
+   SECTION("Vector")
+   {
+      const auto vdim = GENERATE(1,5);
+      const auto ordering = GENERATE(Ordering::byVDIM, Ordering::byNODES);
+      const auto problem = GENERATE(LinearFormExtTest::VectorDomainLF,
+                                    LinearFormExtTest::VectorDomainLFGrad);
+      LinearFormExtTest(N, dim, vdim, ordering, gll, problem, p, true).Run();
+   }
+} // test case
+

tests/unit/fem/test_linearform_ext.hpp

@@ -0,0 +1,229 @@
+// Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#include "mfem.hpp"
+#include <functional>
+
+namespace mfem
+{
+
+namespace linearform_ext_tests
+{
+
+constexpr int SEED = 0x100001b3;
+
+struct LinearFormExtTest
+{
+   enum
+   {
+      DomainLF = 1,
+      DomainLFGrad = 2,
+      VectorDomainLF = 3,
+      VectorDomainLFGrad = 4
+   };
+
+   const int dim, vdim, ordering;
+   const bool gll, test;
+   const int problem, N, p, q;
+   const Element::Type type;
+   Mesh mesh;
+   H1_FECollection fec;
+   FiniteElementSpace vfes, mfes;
+   GridFunction x;
+   const Geometry::Type geom_type;
+   IntegrationRules IntRulesGLL;
+   const IntegrationRule *irGLL, *ir;
+
+   Array<int> elem_marker;
+   Vector one_vec, dim_vec, vdim_vec, vdim_dim_vec;
+   ConstantCoefficient constant_coeff;
+   VectorConstantCoefficient dim_constant_coeff;
+   VectorConstantCoefficient vdim_constant_coeff;
+   VectorConstantCoefficient vdim_dim_constant_coeff;
+   std::function<void(const Vector&, Vector&)>
+   vdim_vector_function = [&](const Vector&, Vector &y)
+   {
+      y.SetSize(vdim);
+      y.Randomize(SEED);
+   };
+   std::function<void(const Vector&, Vector&)> vector_f;
+   VectorFunctionCoefficient vdim_function_coeff;
+
+   LinearForm lf_full, lf_legacy;
+
+   const int dofs;
+   double mdofs;
+
+   LinearFormExtTest(int N, int dim, int vdim, int ordering,
+                     bool gll,
+                     int problem, int order,
+                     bool test):
+      dim(dim),
+      vdim(vdim),
+      ordering(ordering),
+      gll(gll),
+      test(test),
+      problem(problem),
+      N(N),
+      p(order),
+      q(2*p + (gll?-1:3)),
+      type(dim==3 ?
+           Element::HEXAHEDRON :
+           Element::QUADRILATERAL),
+      mesh(dim==2 ?
+           Mesh::MakeCartesian2D(N,N,type):
+           Mesh::MakeCartesian3D(N,N,N,type)),
+      fec(p, dim),
+      vfes(&mesh, &fec, vdim, ordering),
+      mfes(&mesh, &fec, dim),
+      x(&mfes),
+      geom_type(vfes.GetFE(0)->GetGeomType()),
+      IntRulesGLL(0, Quadrature1D::GaussLobatto),
+      irGLL(&IntRulesGLL.Get(geom_type, q)),
+      ir(&IntRules.Get(geom_type, q)),
+      elem_marker(),
+      one_vec(1),
+      dim_vec(dim),
+      vdim_vec(vdim),
+      vdim_dim_vec(vdim*dim),
+      constant_coeff(M_PI),
+      dim_constant_coeff((dim_vec.Randomize(SEED), dim_vec)),
+      vdim_constant_coeff((vdim_vec.Randomize(SEED), vdim_vec)),
+      vdim_dim_constant_coeff((vdim_dim_vec.Randomize(SEED), vdim_dim_vec)),
+      vector_f(vdim_vector_function),
+      vdim_function_coeff(vdim, vector_f),
+      lf_full(&vfes),
+      lf_legacy(&vfes),
+      dofs(vfes.GetVSize()),
+      mdofs(0.0)
+   {
+      SetupMarkers();
+      SetupRandomMesh();
+
+      LinearFormIntegrator *integ_full = nullptr;
+      LinearFormIntegrator *integ_legacy = nullptr;
+
+      switch (problem)
+      {
+         case DomainLF:
+         {
+            integ_full = new DomainLFIntegrator(constant_coeff);
+            integ_legacy = new DomainLFIntegrator(constant_coeff);
+            break;
+         }
+         case DomainLFGrad:
+         {
+            integ_full = new DomainLFGradIntegrator(dim_constant_coeff);
+            integ_legacy = new DomainLFGradIntegrator(dim_constant_coeff);
+            break;
+         }
+         case VectorDomainLF:
+         {
+            if (test)
+            {
+               integ_full = new VectorDomainLFIntegrator(vdim_function_coeff);
+               integ_legacy = new VectorDomainLFIntegrator(vdim_function_coeff);
+            }
+            else // !test => bench, we don't want to spend time building coeff
+            {
+               integ_full = new VectorDomainLFIntegrator(vdim_constant_coeff);
+               integ_legacy = new VectorDomainLFIntegrator(vdim_constant_coeff);
+            }
+            break;
+         }
+         case VectorDomainLFGrad:
+         {
+            integ_full = new VectorDomainLFGradIntegrator(vdim_dim_constant_coeff);
+            integ_legacy = new VectorDomainLFGradIntegrator(vdim_dim_constant_coeff);
+            break;
+         }
+         default: { MFEM_ABORT("Unknown Problem!"); }
+      }
+
+      integ_full->SetIntRule(gll ? irGLL : ir);
+      integ_legacy->SetIntRule(gll ? irGLL : ir);
+
+      MFEM_VERIFY(mesh.attributes.Size() == elem_marker.Size(),
+                  "Markers attributes error!");
+      lf_full.AddDomainIntegrator(integ_full, elem_marker);
+      lf_legacy.AddDomainIntegrator(integ_legacy, elem_marker);
+
+      lf_full.SetAssemblyLevel(LinearAssemblyLevel::FULL);
+      lf_legacy.SetAssemblyLevel(LinearAssemblyLevel::LEGACY);
+   }
+
+   void AssembleBoth()
+   {
+      lf_full.Assemble();
+      lf_legacy.Assemble();
+   }
+
+   virtual void Run();
+
+   virtual void Description();
+
+   void SetupRandomMesh()
+   {
+      mesh.SetNodalFESpace(&mfes);
+      mesh.SetNodalGridFunction(&x);
+      const double jitter = 1./(M_PI*M_PI);
+      const double h0 = mesh.GetElementSize(0);
+      GridFunction rdm(&mfes);
+      rdm.Randomize(SEED);
+      rdm -= 0.5; // Shift to random values in [-0.5,0.5]
+      rdm *= jitter * h0; // Scale the random values to be of same order
+      x -= rdm;
+   }
+
+   void SetupMarkers()
+   {
+      // Initial mesh from MakeCartesian has attributes size of 1
+      MFEM_VERIFY(mesh.attributes.Size() == 1, "Initial attributes error!");
+
+      // Add attributes for interior/exterior domain
+      const double radius = sqrt(dim)/2.0;
+      Vector center(dim), diff(dim);
+      center = 0.5;
+      Array<int> vertices;
+      for (int e = 0; e < mesh.GetNE(); e++)
+      {
+         Element *el = mesh.GetElement(e);
+         el->GetVertices(vertices);
+         bool interior = true;
+         for (int j = 0; j < vertices.Size(); j++)
+         {
+            const Vector coord(mesh.GetVertex(vertices[j]), dim);
+            subtract(coord, center, diff);
+            const double norm = diff.Norml2();
+            MFEM_VERIFY(norm >= 0.0 && norm <= radius,"");
+            if (norm > 0.5) { interior = false; break; }
+         }
+         mesh.SetAttribute(e, interior ? 2 : 1);
+      }
+      mesh.SetAttributes();
+
+      // Marked mesh should now have two attributes
+      MFEM_VERIFY(mesh.attributes.Size() == 2, "Marked attributes error!");
+
+      elem_marker.SetSize(2);
+      elem_marker[0] = 0; // ignore exterior
+      elem_marker[1] = 1; // include interior
+   }
+
+   double SumMdofs() const { return mdofs; }
+
+   double MDofs() const { return 1e-6 * dofs; }
+};
+
+} // namespace linearform_ext_tests
+
+} // namespace mfem
+