ElementWiseSmoother: also test with MassIntegrator

There are lots of kernels left to work on, but this is kind of the
full minimal set; ie, for one device configuration I have the full set.

examples/ex1.cpp

@@ -194,9 +194,12 @@ int main(int argc, char *argv[])
       umf_solver.Mult(B, X);
 #endif
    }
-   else // No preconditioning for now in partial assembly mode.
+   else // Jacobi preconditioning in partial assembly mode
    {
-      CG(*A, B, X, 1, 2000, 1e-12, 0.0);
+      Vector diag_pa(fespace->GetTrueVSize());
+      a->AssembleDiagonal(diag_pa);
+      OperatorJacobiSmoother M(diag_pa, ess_tdof_list);
+      PCG(*A, M, B, X, 1, 2000, 1e-12, 0.0);
    }
 
    // 12. Recover the solution as a finite element grid function.

examples/ex1p.cpp

@@ -208,9 +208,18 @@ int main(int argc, char *argv[])
 
    // 13. Solve the linear system A X = B.
    //     * With full assembly, use the BoomerAMG preconditioner from hypre.
-   //     * With partial assembly, use no preconditioner, for now.
+   //     * With partial assembly, use jacobi smoothing, for now.
    Solver *prec = NULL;
-   if (!pa) { prec = new HypreBoomerAMG; }
+   if (pa)
+   {
+      Vector diag_pa(fespace->GetTrueVSize());
+      a->AssembleDiagonal(diag_pa);
+      prec = new OperatorJacobiSmoother(diag_pa, ess_tdof_list);
+   }
+   else
+   {
+      prec = new HypreBoomerAMG;
+   }
    CGSolver cg(MPI_COMM_WORLD);
    cg.SetRelTol(1e-12);
    cg.SetMaxIter(2000);

fem/CMakeLists.txt

@@ -32,6 +32,7 @@ set(SRCS
   nonlininteg.cpp
   staticcond.cpp
   tmop.cpp
+  elementwisesmoother.cpp
   )
 
 set(HDRS
@@ -64,6 +65,7 @@ set(HDRS
   tfespace.hpp
   tintrules.hpp
   tmop.hpp
+  elementwisesmoother.hpp
   )
 
 if (MFEM_USE_SIDRE)

fem/bilinearform.cpp

@@ -294,6 +294,7 @@ void BilinearForm::ComputeElementMatrix(int i, DenseMatrix &elmat)
    }
 }
 
+
 void BilinearForm::ComputeBdrElementMatrix(int i, DenseMatrix &elmat)
 {
    if (bbfi.Size())
@@ -321,6 +322,27 @@ void BilinearForm::AssembleElementMatrix(
    AssembleElementMatrix(i, elmat, vdofs, skip_zeros);
 }
 
+
+/// in is global, out is local to the element
+/// (this will set the size of out_element)
+void BilinearForm::ElementMatrixMult(int i, const Vector& in, Vector& out_element)
+{
+   Array<int> dofs;
+   FESpace()->GetElementDofs(i, dofs);
+   out_element.SetSize(dofs.Size());
+   out_element = 0.0;
+   if (ext)
+   {
+      ext->ElementMatrixMult(i, in, out_element);
+      return;
+   }
+   DenseMatrix elmat;
+   ComputeElementMatrix(i, elmat);
+   Vector subvec;
+   in.GetSubVector(dofs, subvec);
+   elmat.Mult(subvec, out_element);
+}
+
 void BilinearForm::AssembleElementMatrix(
    int i, const DenseMatrix &elmat, Array<int> &vdofs, int skip_zeros)
 {
@@ -608,6 +630,19 @@ void BilinearForm::ConformingAssemble()
    width = mat->Width();
 }
 
+
+void BilinearForm::AssembleDiagonal(Vector& diag) const
+{
+   if (ext)
+   {
+      ext->AssembleDiagonal(diag);
+   }
+   else
+   {
+      mfem_error("Not implemented, maybe assemble your bilinear form into a matrix an use SparseMatrix::GetDiag?");
+   }
+}
+
 void BilinearForm::FormLinearSystem(const Array<int> &ess_tdof_list, Vector &x,
                                     Vector &b, OperatorHandle &A, Vector &X,
                                     Vector &B, int copy_interior)

fem/bilinearform.hpp

@@ -319,6 +319,9 @@ public:
    /// Assembles the form i.e. sums over all domain/bdr integrators.
    void Assemble(int skip_zeros = 1);
 
+   /// Assemble diagonal of bilinear form into diag
+   void AssembleDiagonal(Vector& diag) const;
+
    /// Get the finite element space prolongation matrix
    virtual const Operator *GetProlongation() const
    { return fes->GetConformingProlongation(); }
@@ -459,6 +462,14 @@ public:
    void AssembleBdrElementMatrix(int i, const DenseMatrix &elmat,
                                  Array<int> &vdofs, int skip_zeros = 1);
 
+   /**
+      Multiply element-matrix by vector.
+
+      in is an l-vector (local to processor, but assembled across
+      elements), while out_element is local to the particular element.
+   */
+   void ElementMatrixMult(int i, const Vector& in, Vector& out_element);
+
    /// Eliminate essential boundary DOFs from the system.
    /** The array @a bdr_attr_is_ess marks boundary attributes that constitute
        the essential part of the boundary. By default, the diagonal at the

fem/bilinearform_ext.cpp

@@ -60,6 +60,31 @@ void PABilinearFormExtension::Assemble()
    }
 }
 
+void PABilinearFormExtension::AssembleDiagonal(Vector &y) const
+{
+   Array<BilinearFormIntegrator*> &integrators = *a->GetDBFI();
+
+   const int iSz = integrators.Size();
+   if (elem_restrict_lex)
+   {
+      localY = 0.0;
+      for (int i = 0; i < iSz; ++i)
+      {
+         integrators[i]->AssembleDiagonalPA(localY);
+      }
+      elem_restrict_lex->MultTranspose(localY, y);
+   }
+   else
+   {
+      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]->AssembleDiagonalPA(y);
+      }
+   }
+}
+
 void PABilinearFormExtension::Update()
 {
    FiniteElementSpace *fes = a->FESpace();
@@ -148,4 +173,46 @@ void PABilinearFormExtension::MultTranspose(const Vector &x, Vector &y) const
    }
 }
 
+/// TODO this needs some work
+void PABilinearFormExtension::ElementMatrixMult(int i, const Vector &x,
+                                                Vector &y_element)
+{
+   Array<BilinearFormIntegrator*> &integrators = *a->GetDBFI();
+   Vector yelem_preorder(y_element);
+
+   const int iSz = integrators.Size();
+   if (elem_restrict_lex)
+   {
+      elem_restrict_lex->Mult(x, localX);
+      // localY = 0.0;
+      for (int j = 0; j < iSz; ++j)
+      {
+         integrators[j]->AddMultElementPA(i, localX, yelem_preorder);
+      }
+
+      // reorder... (kinda like the transpose of elem_restrict_lex)
+      const FiniteElement *fe = trialFes->GetFE(i);
+      const TensorBasisElement* el =
+         dynamic_cast<const TensorBasisElement*>(fe);
+      MFEM_VERIFY(el, "Not a tensor element!");
+      const Array<int> &fe_dof_map = el->GetDofMap();
+      MFEM_VERIFY(fe_dof_map.Size() == yelem_preorder.Size(),
+                  "Sizes don't match!");
+      for (int j = 0; j < yelem_preorder.Size(); ++j)
+      {
+         // y_element[j] = yelem_preorder[fe_dof_map[j]];
+         y_element[fe_dof_map[j]] = yelem_preorder[j];
+      }
+
+   }
+   else
+   {
+      y_element = 0.0; // TODO: do I want this?
+      for (int i = 0; i < iSz; ++i)
+      {
+         integrators[i]->AddMultElementPA(i, x, y_element);
+      }
+   }
+}
+
 } // namespace mfem

fem/bilinearform_ext.hpp

@@ -42,6 +42,10 @@ public:
    virtual const Operator *GetRestriction() const;
 
    virtual void Assemble() = 0;
+   virtual void AssembleDiagonal(Vector& diag) const
+   {
+      mfem_error("Not implemented for this assembly level!");
+   }
    virtual void FormSystemMatrix(const Array<int> &ess_tdof_list,
                                  OperatorHandle &A) = 0;
    virtual void FormLinearSystem(const Array<int> &ess_tdof_list,
@@ -49,6 +53,13 @@ public:
                                  OperatorHandle &A, Vector &X, Vector &B,
                                  int copy_interior = 0) = 0;
    virtual void Update() = 0;
+
+   /// Here x is local-to-processor but contains all elements
+   /// (it's an l-vector), while y_element is striclty local to an element
+   virtual void ElementMatrixMult(int i, const Vector &x, Vector &y_element)
+   {
+      mfem_error("Not implemented!");
+   }
 };
 
 /// Data and methods for fully-assembled bilinear forms
@@ -103,6 +114,7 @@ public:
    PABilinearFormExtension(BilinearForm*);
 
    void Assemble();
+   void AssembleDiagonal(Vector& diag) const;
    void FormSystemMatrix(const Array<int> &ess_tdof_list, OperatorHandle &A);
    void FormLinearSystem(const Array<int> &ess_tdof_list,
                          Vector &x, Vector &b,
@@ -112,6 +124,7 @@ public:
    void Mult(const Vector &x, Vector &y) const;
    void MultTranspose(const Vector &x, Vector &y) const;
    void Update();
+   void ElementMatrixMult(int i, const Vector &x, Vector &y_element);
 };
 
 /// Data and methods for matrix-free bilinear forms

fem/bilininteg.cpp

@@ -26,15 +26,28 @@ void BilinearFormIntegrator::AssemblePA(const FiniteElementSpace&)
                "   is not implemented for this class.");
 }
 
+void BilinearFormIntegrator::AssembleDiagonalPA(Vector&) const
+{
+   mfem_error ("BilinearFormIntegrator::AssembleDiagonalPA (...)\n"
+               "   is not implemented for this class.");
+}
+
 void BilinearFormIntegrator::AddMultPA(const Vector &, Vector &) const
 {
-   mfem_error ("BilinearFormIntegrator::MultAssembled (...)\n"
+   mfem_error ("BilinearFormIntegrator::AddMultPA (...)\n"
+               "   is not implemented for this class.");
+}
+
+void BilinearFormIntegrator::AddMultElementPA(int element,
+                                              const Vector &, Vector &) const
+{
+   mfem_error ("BilinearFormIntegrator::AddMultElementPA (...)\n"
                "   is not implemented for this class.");
 }
 
 void BilinearFormIntegrator::AddMultTransposePA(const Vector &, Vector &) const
 {
-   mfem_error ("BilinearFormIntegrator::MultAssembledTranspose (...)\n"
+   mfem_error ("BilinearFormIntegrator::AddMultTransposePA (...)\n"
                "   is not implemented for this class.");
 }
 

fem/bilininteg.hpp

@@ -44,6 +44,9 @@ public:
        used later in the methods AddMultPA() and AddMultTransposePA(). */
    virtual void AssemblePA(const FiniteElementSpace &fes);
 
+   /// assemble diagonal into vector diag
+   virtual void AssembleDiagonalPA(Vector& diag) const;
+
    /// Method for partially assembled action.
    /** Perform the action of integrator on the input @a x and add the result to
        the output @a y. Both @a x and @a y are E-vectors, i.e. they represent
@@ -53,6 +56,9 @@ public:
        called. */
    virtual void AddMultPA(const Vector &x, Vector &y) const;
 
+   /// Method for partially assembled action of an element matrix.
+   virtual void AddMultElementPA(int element, const Vector &x, Vector &y) const;
+
    /// Method for partially assembled transposed action.
    /** Perform the transpose action of integrator on the input @a x and add the
        result to the output @a y. Both @a x and @a y are E-vectors, i.e. they
@@ -1719,8 +1725,12 @@ public:
 
    virtual void AssemblePA(const FiniteElementSpace&);
 
+   virtual void AssembleDiagonalPA(Vector& diag) const;
+
    virtual void AddMultPA(const Vector&, Vector&) const;
 
+   virtual void AddMultElementPA(int element, const Vector &x, Vector &y) const;
+
    static const IntegrationRule &GetRule(const FiniteElement &trial_fe,
                                          const FiniteElement &test_fe);
 };
@@ -1759,8 +1769,12 @@ public:
 
    virtual void AssemblePA(const FiniteElementSpace&);
 
+   virtual void AssembleDiagonalPA(Vector& diag) const;
+
    virtual void AddMultPA(const Vector&, Vector&) const;
 
+   virtual void AddMultElementPA(int element, const Vector &x, Vector &y) const;
+
    static const IntegrationRule &GetRule(const FiniteElement &trial_fe,
                                          const FiniteElement &test_fe,
                                          ElementTransformation &Trans);

fem/bilininteg_mass.cpp

@@ -102,6 +102,184 @@ void MassIntegrator::AssemblePA(const FiniteElementSpace &fes)
    }
 }
 
+template<const int T_D1D = 0,
+         const int T_Q1D = 0>
+static void PAMassAssembleDiagonal2D(const int NE,
+                                     const Array<double> &_B,
+                                     const Array<double> &_Bt,
+                                     const Vector &_op,
+                                     Vector &_diag,
+                                     const int d1d = 0,
+                                     const int q1d = 0)
+{
+   // see PAMassApply2D
+   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 op = Reshape(_op.Read(), Q1D, Q1D, NE);
+   auto y = Reshape(_diag.ReadWrite(), D1D, D1D, NE);
+   MFEM_FORALL(e, NE,
+   {
+      const int D1D = T_D1D ? T_D1D : d1d; // nvcc workaround
+      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      // the following variables are evaluated at compile time
+      constexpr int max_D1D = T_D1D ? T_D1D : MAX_D1D;
+      constexpr int max_Q1D = T_Q1D ? T_Q1D : MAX_Q1D;
+
+      double temp[max_Q1D][max_D1D];
+      for (int qx = 0; qx < Q1D; ++qx)
+      {
+         for (int dy = 0; dy < D1D; ++dy)
+         {
+            temp[qx][dy] = 0.0;
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               temp[qx][dy] += B(qy, dy) * B(qy, dy) * op(qx, qy, e);
+            }
+         }
+      }
+      for (int dy = 0; dy < D1D; ++dy)
+      {
+         for (int dx = 0; dx < D1D; ++dx)
+         {
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               // might need absolute values on next line
+               y(dx,dy,e) += B(qx, dx) * B(qx, dx) * temp[qx][dy];
+            }
+         }
+      }
+   });
+}
+
+template<const int T_D1D = 0,
+         const int T_Q1D = 0>
+static void PAMassAssembleDiagonal3D(const int NE,
+                                     const Array<double> &_B,
+                                     const Array<double> &_Bt,
+                                     const Vector &_op,
+                                     Vector &_diag,
+                                     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 op = Reshape(_op.Read(), Q1D, Q1D, Q1D, NE);
+   auto y = Reshape(_diag.ReadWrite(), D1D, D1D, D1D, NE);
+   MFEM_FORALL(e, NE,
+   {
+      const int D1D = T_D1D ? T_D1D : d1d; // nvcc workaround
+      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      // the following variables are evaluated at compile time
+      constexpr int max_D1D = T_D1D ? T_D1D : MAX_D1D;
+      constexpr int max_Q1D = T_Q1D ? T_Q1D : MAX_Q1D;
+
+      double temp[max_Q1D][max_Q1D][max_D1D];
+      for (int qx = 0; qx < Q1D; ++qx)
+      {
+         for (int qy = 0; qy < Q1D; ++qy)
+         {
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               temp[qx][qy][dz] = 0.0;
+               for (int qz = 0; qz < Q1D; ++qz)
+               {
+                  temp[qx][qy][dz] += B(qz, dz) * B(qz, dz) * op(qx, qy, qz, e);
+               }
+            }
+         }
+      }
+      double temp2[max_Q1D][max_D1D][max_D1D];
+      for (int qx = 0; qx < Q1D; ++qx)
+      {
+         for (int dz = 0; dz < D1D; ++dz)
+         {
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               temp2[qx][dy][dz] = 0.0;
+               for (int qy = 0; qy < Q1D; ++qy)
+               {
+                  temp2[qx][dy][dz] += B(qy, dy) * B(qy, dy) * temp[qx][qy][dz];
+               }
+            }
+         }
+      }
+      for (int dz = 0; dz < D1D; ++dz)
+      {
+         for (int dy = 0; dy < D1D; ++dy)
+         {
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               for (int qx = 0; qx < Q1D; ++qx)
+               {
+                  y(dx, dy, dz, e) += B(qx, dx) * B(qx, dx) * temp2[qx][dy][dz];
+               }
+            }
+         }
+      }
+   });
+}
+
+static void PAMassAssembleDiagonal(
+   const int dim, const int D1D,
+   const int Q1D, const int NE,
+   const Array<double> &B, const Array<double> &Bt,
+   const Vector &op, Vector &y)
+{
+#ifdef MFEM_USE_OCCA
+   if (DeviceCanUseOcca())
+   {
+      MFEM_ABORT("OCCA PA Mass Assemble Diagonal unknown kernel!");
+   }
+#endif // MFEM_USE_OCCA
+
+   if (Device::Allows(Backend::RAJA_CUDA))
+   {
+      if (dim == 2)
+      {
+         switch ((D1D << 4 ) | Q1D)
+         {
+            default:   return PAMassAssembleDiagonal2D(NE, B, Bt, op, y, D1D, Q1D);
+         }
+      }
+      if (dim == 3)
+      {
+         switch ((D1D << 4 ) | Q1D)
+         {
+            default:   return PAMassAssembleDiagonal3D(NE, B, Bt, op, y, D1D, Q1D);
+         }
+      }
+   }
+   else if (dim == 2)
+   {
+      switch ((D1D << 4 ) | Q1D)
+      {
+         // should look at smem routines in what follows?
+         default:   return PAMassAssembleDiagonal2D(NE, B, Bt, op, y, D1D, Q1D);
+      }
+   }
+   else if (dim == 3)
+   {
+      switch ((D1D << 4 ) | Q1D)
+      {
+         // should look at smem routines in what follows?
+         default:   return PAMassAssembleDiagonal3D(NE, B, Bt, op, y, D1D, Q1D);
+      }
+   }
+   MFEM_ABORT("Unknown kernel.");
+}
+
+void MassIntegrator::AssembleDiagonalPA(Vector& diag) const
+{
+   PAMassAssembleDiagonal(dim, dofs1D, quad1D, ne,
+                          maps->B, maps->Bt, pa_data, diag);
+}
+
 #ifdef MFEM_USE_OCCA
 // OCCA PA Mass Apply 2D kernel
 static void OccaPAMassApply2D(const int D1D,
@@ -287,6 +465,117 @@ static void PAMassApply2D(const int NE,
    });
 }
 
+template<const int T_D1D = 0,
+         const int T_Q1D = 0,
+         const int T_NBZ = 0>
+static void SmemPAMassApply2DElement(int e,
+                                     DeviceTensor<2,const double> b,
+                                     DeviceTensor<3,const double> op,
+                                     DeviceTensor<2,const double> x,
+                                     DeviceTensor<2,double> y,
+                                     const int d1d = 0,
+                                     const int q1d = 0)
+{
+   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;
+   constexpr int MDQ = (MQ1 > MD1) ? MQ1 : MD1;
+   MFEM_SHARED double BBt[MQ1*MD1];
+   double (*B)[MD1] = (double (*)[MD1]) BBt;
+   double (*Bt)[MQ1] = (double (*)[MQ1]) BBt;
+   MFEM_SHARED double sm0[NBZ][MDQ*MDQ];
+   MFEM_SHARED double sm1[NBZ][MDQ*MDQ];
+   double (*X)[MD1] = (double (*)[MD1]) (sm0 + tidz);
+   double (*DQ)[MQ1] = (double (*)[MQ1]) (sm1 + tidz);
+   double (*QQ)[MQ1] = (double (*)[MQ1]) (sm0 + tidz);
+   double (*QD)[MD1] = (double (*)[MD1]) (sm1 + tidz);
+   MFEM_FOREACH_THREAD(dy,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         // X[dy][dx] = x(dx,dy,e);
+         X[dy][dx] = x(dx,dy);
+      }
+   }
+   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(dy,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         double dq = 0.0;
+         for (int dx = 0; dx < D1D; ++dx)
+         {
+            dq += X[dy][dx] * B[qx][dx];
+         }
+         DQ[dy][qx] = dq;
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(qy,y,Q1D)
+   {
+      MFEM_FOREACH_THREAD(qx,x,Q1D)
+      {
+         double qq = 0.0;
+         for (int dy = 0; dy < D1D; ++dy)
+         {
+            qq += DQ[dy][qx] * B[qy][dy];
+         }
+         QQ[qy][qx] = qq * op(qx, qy, e);
+      }
+   }
+   MFEM_SYNC_THREAD;
+   if (tidz == 0)
+   {
+      MFEM_FOREACH_THREAD(d,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(q,x,Q1D)
+         {
+            Bt[d][q] = b(q,d);
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(qy,y,Q1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         double dq = 0.0;
+         for (int qx = 0; qx < Q1D; ++qx)
+         {
+            dq += QQ[qy][qx] * Bt[dx][qx];
+         }
+         QD[qy][dx] = dq;
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(dy,y,D1D)
+   {
+      MFEM_FOREACH_THREAD(dx,x,D1D)
+      {
+         double dd = 0.0;
+         for (int qy = 0; qy < Q1D; ++qy)
+         {
+            dd += (QD[qy][dx] * Bt[dy][qy]);
+         }
+         // y(dx, dy, e) += dd;
+         y(dx, dy) += dd;
+      }
+   }
+}
+
 template<const int T_D1D = 0,
          const int T_Q1D = 0,
          const int T_NBZ = 0>
@@ -312,102 +601,15 @@ static void SmemPAMassApply2D(const int 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;
-      constexpr int MDQ = (MQ1 > MD1) ? MQ1 : MD1;
-      MFEM_SHARED double BBt[MQ1*MD1];
-      double (*B)[MD1] = (double (*)[MD1]) BBt;
-      double (*Bt)[MQ1] = (double (*)[MQ1]) BBt;
-      MFEM_SHARED double sm0[NBZ][MDQ*MDQ];
-      MFEM_SHARED double sm1[NBZ][MDQ*MDQ];
-      double (*X)[MD1] = (double (*)[MD1]) (sm0 + tidz);
-      double (*DQ)[MQ1] = (double (*)[MQ1]) (sm1 + tidz);
-      double (*QQ)[MQ1] = (double (*)[MQ1]) (sm0 + tidz);
-      double (*QD)[MD1] = (double (*)[MD1]) (sm1 + tidz);
-      MFEM_FOREACH_THREAD(dy,y,D1D)
-      {
-         MFEM_FOREACH_THREAD(dx,x,D1D)
-         {
-            X[dy][dx] = x(dx,dy,e);
-         }
-      }
-      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(dy,y,D1D)
-      {
-         MFEM_FOREACH_THREAD(qx,x,Q1D)
-         {
-            double dq = 0.0;
-            for (int dx = 0; dx < D1D; ++dx)
-            {
-               dq += X[dy][dx] * B[qx][dx];
-            }
-            DQ[dy][qx] = dq;
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qy,y,Q1D)
-      {
-         MFEM_FOREACH_THREAD(qx,x,Q1D)
-         {
-            double qq = 0.0;
-            for (int dy = 0; dy < D1D; ++dy)
-            {
-               qq += DQ[dy][qx] * B[qy][dy];
-            }
-            QQ[qy][qx] = qq * op(qx, qy, e);
-         }
-      }
-      MFEM_SYNC_THREAD;
-      if (tidz == 0)
-      {
-         MFEM_FOREACH_THREAD(d,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(q,x,Q1D)
-            {
-               Bt[d][q] = b(q,d);
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qy,y,Q1D)
-      {
-         MFEM_FOREACH_THREAD(dx,x,D1D)
-         {
-            double dq = 0.0;
-            for (int qx = 0; qx < Q1D; ++qx)
-            {
-               dq += QQ[qy][qx] * Bt[dx][qx];
-            }
-            QD[qy][dx] = dq;
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(dy,y,D1D)
-      {
-         MFEM_FOREACH_THREAD(dx,x,D1D)
-         {
-            double dd = 0.0;
-            for (int qy = 0; qy < Q1D; ++qy)
-            {
-               dd += (QD[qy][dx] * Bt[dy][qy]);
-            }
-            y(dx, dy, e) += dd;
-         }
-      }
+      auto x_element = DeviceTensor<2, const double>(
+         (const double*) x + e * D1D * D1D,
+         D1D, D1D);
+      auto y_element = DeviceTensor<2, double>(
+         (double*) y + e * D1D * D1D,
+         D1D, D1D);
+
+      SmemPAMassApply2DElement<T_D1D, T_Q1D, T_NBZ>(e, b, op, x_element,
+                                                    y_element, d1d, q1d);
    });
 }
 
@@ -553,14 +755,174 @@ static void PAMassApply3D(const int NE,
    });
 }
 
+
+template<const int T_D1D = 0,
+         const int T_Q1D = 0>
+static void SmemPAMassApply3DElement(int e,
+                                     DeviceTensor<2,const double> b,
+                                     DeviceTensor<4,const double> op,
+                                     DeviceTensor<3,const double> x,
+                                     DeviceTensor<3,double> y,
+                                     const int d1d = 0,
+                                     const int q1d = 0)
+{
+   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 MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
+   constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
+   constexpr int MDQ = (MQ1 > MD1) ? MQ1 : MD1;
+   MFEM_SHARED double sDQ[MQ1*MD1];
+   double (*B)[MD1] = (double (*)[MD1]) sDQ;
+   double (*Bt)[MQ1] = (double (*)[MQ1]) sDQ;
+   MFEM_SHARED double sm0[MDQ*MDQ*MDQ];
+   MFEM_SHARED double sm1[MDQ*MDQ*MDQ];
+   double (*X)[MD1][MD1]   = (double (*)[MD1][MD1]) sm0;
+   double (*DDQ)[MD1][MQ1] = (double (*)[MD1][MQ1]) sm1;
+   double (*DQQ)[MQ1][MQ1] = (double (*)[MQ1][MQ1]) sm0;
+   double (*QQQ)[MQ1][MQ1] = (double (*)[MQ1][MQ1]) sm1;
+   double (*QQD)[MQ1][MD1] = (double (*)[MQ1][MD1]) sm0;
+   double (*QDD)[MD1][MD1] = (double (*)[MD1][MD1]) sm1;
+   MFEM_FOREACH_THREAD(dz,z,D1D)
+   {
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            // X[dz][dy][dx] = x(dx,dy,dz,e);
+            X[dz][dy][dx] = x(dx,dy,dz);
+         }
+      }
+   }
+   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(dz,z,D1D)
+   {
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            double u = 0.0;
+            for (int dx = 0; dx < D1D; ++dx)
+            {
+               u += X[dz][dy][dx] * B[qx][dx];
+            }
+            DDQ[dz][dy][qx] = u;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(dz,z,D1D)
+   {
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            double u = 0.0;
+            for (int dy = 0; dy < D1D; ++dy)
+            {
+               u += DDQ[dz][dy][qx] * B[qy][dy];
+            }
+            DQQ[dz][qy][qx] = u;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(qz,z,Q1D)
+   {
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(qx,x,Q1D)
+         {
+            double u = 0.0;
+            for (int dz = 0; dz < D1D; ++dz)
+            {
+               u += DQQ[dz][qy][qx] * B[qz][dz];
+            }
+            QQQ[qz][qy][qx] = u * op(qx,qy,qz,e);
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   if (tidz == 0)
+   {
+      MFEM_FOREACH_THREAD(d,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(q,x,Q1D)
+         {
+            Bt[d][q] = b(q,d);
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(qz,z,Q1D)
+   {
+      MFEM_FOREACH_THREAD(qy,y,Q1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u = 0.0;
+            for (int qx = 0; qx < Q1D; ++qx)
+            {
+               u += QQQ[qz][qy][qx] * Bt[dx][qx];
+            }
+            QQD[qz][qy][dx] = u;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(qz,z,Q1D)
+   {
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u = 0.0;
+            for (int qy = 0; qy < Q1D; ++qy)
+            {
+               u += QQD[qz][qy][dx] * Bt[dy][qy];
+            }
+            QDD[qz][dy][dx] = u;
+         }
+      }
+   }
+   MFEM_SYNC_THREAD;
+   MFEM_FOREACH_THREAD(dz,z,D1D)
+   {
+      MFEM_FOREACH_THREAD(dy,y,D1D)
+      {
+         MFEM_FOREACH_THREAD(dx,x,D1D)
+         {
+            double u = 0.0;
+            for (int qz = 0; qz < Q1D; ++qz)
+            {
+               u += QDD[qz][dy][dx] * Bt[dz][qz];
+            }
+            // y(dx,dy,dz,e) += u;
+            y(dx,dy,dz) += u;
+         }
+      }
+   }
+}
+
+
 template<const int T_D1D = 0,
          const int T_Q1D = 0>
 static void SmemPAMassApply3D(const int NE,
-                              const Array<double> &b_,
-                              const Array<double> &bt_,
-                              const Vector &op_,
-                              const Vector &x_,
-                              Vector &y_,
+                              const Array<double> &_b,
+                              const Array<double> &_bt,
+                              const Vector &_op,
+                              const Vector &_x,
+                              Vector &_y,
                               const int d1d = 0,
                               const int q1d = 0)
 {
@@ -570,156 +932,22 @@ static void SmemPAMassApply3D(const int NE,
    constexpr int M1D = T_D1D ? T_D1D : MAX_D1D;
    MFEM_VERIFY(D1D <= M1D, "");
    MFEM_VERIFY(Q1D <= M1Q, "");
-   auto b = Reshape(b_.Read(), Q1D, D1D);
-   auto op = Reshape(op_.Read(), Q1D, Q1D, Q1D, NE);
-   auto x = Reshape(x_.Read(), D1D, D1D, D1D, NE);
-   auto y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
+   auto b = Reshape(_b.Read(), Q1D, D1D);
+   auto op = Reshape(_op.Read(), Q1D, Q1D, Q1D, NE);
+   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, 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 MQ1 = T_Q1D ? T_Q1D : MAX_Q1D;
-      constexpr int MD1 = T_D1D ? T_D1D : MAX_D1D;
-      constexpr int MDQ = (MQ1 > MD1) ? MQ1 : MD1;
-      MFEM_SHARED double sDQ[MQ1*MD1];
-      double (*B)[MD1] = (double (*)[MD1]) sDQ;
-      double (*Bt)[MQ1] = (double (*)[MQ1]) sDQ;
-      MFEM_SHARED double sm0[MDQ*MDQ*MDQ];
-      MFEM_SHARED double sm1[MDQ*MDQ*MDQ];
-      double (*X)[MD1][MD1]   = (double (*)[MD1][MD1]) sm0;
-      double (*DDQ)[MD1][MQ1] = (double (*)[MD1][MQ1]) sm1;
-      double (*DQQ)[MQ1][MQ1] = (double (*)[MQ1][MQ1]) sm0;
-      double (*QQQ)[MQ1][MQ1] = (double (*)[MQ1][MQ1]) sm1;
-      double (*QQD)[MQ1][MD1] = (double (*)[MQ1][MD1]) sm0;
-      double (*QDD)[MD1][MD1] = (double (*)[MD1][MD1]) sm1;
-      MFEM_FOREACH_THREAD(dz,z,D1D)
-      {
-         MFEM_FOREACH_THREAD(dy,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(dx,x,D1D)
-            {
-               X[dz][dy][dx] = x(dx,dy,dz,e);
-            }
-         }
-      }
-      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(dz,z,D1D)
-      {
-         MFEM_FOREACH_THREAD(dy,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(qx,x,Q1D)
-            {
-               double u = 0.0;
-               for (int dx = 0; dx < D1D; ++dx)
-               {
-                  u += X[dz][dy][dx] * B[qx][dx];
-               }
-               DDQ[dz][dy][qx] = u;
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(dz,z,D1D)
-      {
-         MFEM_FOREACH_THREAD(qy,y,Q1D)
-         {
-            MFEM_FOREACH_THREAD(qx,x,Q1D)
-            {
-               double u = 0.0;
-               for (int dy = 0; dy < D1D; ++dy)
-               {
-                  u += DDQ[dz][dy][qx] * B[qy][dy];
-               }
-               DQQ[dz][qy][qx] = u;
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qz,z,Q1D)
-      {
-         MFEM_FOREACH_THREAD(qy,y,Q1D)
-         {
-            MFEM_FOREACH_THREAD(qx,x,Q1D)
-            {
-               double u = 0.0;
-               for (int dz = 0; dz < D1D; ++dz)
-               {
-                  u += DQQ[dz][qy][qx] * B[qz][dz];
-               }
-               QQQ[qz][qy][qx] = u * op(qx,qy,qz,e);
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      if (tidz == 0)
-      {
-         MFEM_FOREACH_THREAD(d,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(q,x,Q1D)
-            {
-               Bt[d][q] = b(q,d);
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qz,z,Q1D)
-      {
-         MFEM_FOREACH_THREAD(qy,y,Q1D)
-         {
-            MFEM_FOREACH_THREAD(dx,x,D1D)
-            {
-               double u = 0.0;
-               for (int qx = 0; qx < Q1D; ++qx)
-               {
-                  u += QQQ[qz][qy][qx] * Bt[dx][qx];
-               }
-               QQD[qz][qy][dx] = u;
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(qz,z,Q1D)
-      {
-         MFEM_FOREACH_THREAD(dy,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(dx,x,D1D)
-            {
-               double u = 0.0;
-               for (int qy = 0; qy < Q1D; ++qy)
-               {
-                  u += QQD[qz][qy][dx] * Bt[dy][qy];
-               }
-               QDD[qz][dy][dx] = u;
-            }
-         }
-      }
-      MFEM_SYNC_THREAD;
-      MFEM_FOREACH_THREAD(dz,z,D1D)
-      {
-         MFEM_FOREACH_THREAD(dy,y,D1D)
-         {
-            MFEM_FOREACH_THREAD(dx,x,D1D)
-            {
-               double u = 0.0;
-               for (int qz = 0; qz < Q1D; ++qz)
-               {
-                  u += QDD[qz][dy][dx] * Bt[dz][qz];
-               }
-               y(dx,dy,dz,e) += u;
-            }
-         }
-      }
+      auto x_element = DeviceTensor<3, const double>(
+         (const double*) x + e * D1D * D1D * D1D,
+         D1D, D1D, D1D);
+      auto y_element = DeviceTensor<3, double>(
+         (double*) y + e * D1D * D1D * D1D,
+         D1D, D1D, D1D);
+
+      SmemPAMassApply3DElement<T_D1D, T_Q1D>(e, b, op, x_element,
+                                             y_element, D1D, Q1D);
+
    });
 }
 
@@ -786,4 +1014,60 @@ void MassIntegrator::AddMultPA(const Vector &x, Vector &y) const
    PAMassApply(dim, dofs1D, quad1D, ne, maps->B, maps->Bt, pa_data, x, y);
 }
 
+/**
+   as written this wants a global e-vector for x and a local
+   element vector for y
+
+   not sure that's the right interface
+*/
+void MassIntegrator::AddMultElementPA(int element, const Vector &x,
+                                      Vector &y) const
+{
+   if (dim == 2)
+   {
+      auto B = Reshape(maps->B.Read(), quad1D, dofs1D);
+      auto op = Reshape(pa_data.Read(), quad1D, quad1D, ne);
+
+      auto _x = Reshape(x.Read(), dofs1D, dofs1D, ne);
+      auto _y = Reshape(y.ReadWrite(), dofs1D, dofs1D);
+
+      auto x_element = DeviceTensor<2, const double>(
+         (const double*) _x + element * dofs1D * dofs1D,
+         dofs1D, dofs1D);
+
+      SmemPAMassApply2DElement(element,
+                               B,
+                               op,
+                               x_element,
+                               _y,
+                               dofs1D,
+                               quad1D);
+   }
+   else if (dim == 3)
+   {
+      auto B = Reshape(maps->B.Read(), quad1D, dofs1D);
+      auto op = Reshape(pa_data.Read(), quad1D, quad1D, quad1D, ne);
+      auto _x = Reshape(x.Read(), dofs1D, dofs1D, dofs1D, ne);
+      auto _y = Reshape(y.ReadWrite(), dofs1D, dofs1D, dofs1D);
+
+      auto x_element = DeviceTensor<3, const double>(
+         (const double*) _x + element * dofs1D * dofs1D * dofs1D,
+         dofs1D, dofs1D, dofs1D);
+
+      SmemPAMassApply3DElement(element,
+                               B,
+                               op,
+                               x_element,
+                               _y,
+                               dofs1D,
+                               quad1D);
+   }
+   else
+   {
+      mfem_error("Not implemented!");
+   }
+
+}
+
+
 } // namespace mfem

fem/elementwisesmoother.cpp

@@ -0,0 +1,131 @@
+// Copyright (c) 2010, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-443211. All Rights
+// reserved. See file COPYRIGHT for details.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability see http://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the GNU Lesser General Public License (as published by the Free
+// Software Foundation) version 2.1 dated February 1999.
+
+#include "fem.hpp"
+
+namespace mfem
+{
+
+ElementWiseSmoother::ElementWiseSmoother(const mfem::FiniteElementSpace& fespace)
+   :
+   mfem::Solver(fespace.GetVSize()),
+   fespace_(fespace)
+{
+}
+
+void ElementWiseSmoother::Mult(const Vector& b, Vector& x) const
+{
+   Vector local_residual;
+   Vector local_correction;
+   Vector b_local;
+   Array<int> local_dofs;
+
+   if (!iterative_mode)
+   {
+      x = 0.0;
+   }
+
+   for (int e = 0; e < fespace_.GetNE(); ++e)
+   {
+      fespace_.GetElementDofs(e, local_dofs);
+      b.GetSubVector(local_dofs, b_local);
+      local_correction.SetSize(b_local.Size());
+      ElementResidual(e, b_local, x, local_residual);
+
+      LocalSmoother(e, local_residual, local_correction);
+
+      x.AddElementVector(local_dofs, local_correction);
+   }
+}
+
+void ElementWiseSmoother::ElementResidual(int e, const Vector& b_local,
+                                          const Vector& x,
+                                          Vector& r_local) const
+{
+   r_local.SetSize(b_local.Size());
+   r_local = 0.0;
+
+   GetElementFromMatVec(e, x, r_local);
+   r_local -= b_local;
+   r_local *= -1.0;
+}
+
+ElementWiseJacobi::ElementWiseJacobi(const mfem::FiniteElementSpace& fespace,
+                                     mfem::BilinearForm& aform,
+                                     const mfem::Vector& global_diag,
+                                     double scale)
+   :
+   ElementWiseSmoother(fespace),
+   aform_(aform),
+   global_diag_(global_diag),
+   scale_(scale)
+{
+   const Table& el_dof = fespace_.GetElementToDofTable();
+   Table dof_el;
+   mfem::Transpose(el_dof, dof_el);
+   mfem::Mult(el_dof, dof_el, el_to_el_);
+}
+
+// x is *global*, y is local to the element
+void ElementWiseJacobi::GetElementFromMatVec(int element, const mfem::Vector& x,
+                                             mfem::Vector& y_element) const
+{
+   Array<int> neighbors;
+   el_to_el_.GetRow(element, neighbors);
+   Array<int> row_dofs;
+   Array<int> col_dofs;
+   fespace_.GetElementDofs(element, row_dofs);
+   Vector x_local;
+   Vector z_element;
+   DenseMatrix elmat;
+
+   y_element.SetSize(row_dofs.Size());
+   z_element.SetSize(row_dofs.Size());
+   y_element = 0.0;
+
+   for (int i = 0; i < neighbors.Size(); ++i)
+   {
+      int ne = neighbors[i];
+      fespace_.GetElementDofs(ne, col_dofs);
+
+      z_element = 0.0;
+      aform_.ElementMatrixMult(ne, x, z_element);
+
+      // okay, this next section seems very inefficient
+      // (could probably store and precompute some kind of map?)
+      for (int j = 0; j < row_dofs.Size(); ++j)
+      {
+         const int rd = row_dofs[j];
+         for (int k = 0; k < col_dofs.Size(); ++k)
+         {
+            const int cd = col_dofs[k];
+            if (rd == cd)
+            {
+               y_element[j] += z_element[k];
+               break;
+            }
+         }
+      }
+   }
+}
+
+void ElementWiseJacobi::LocalSmoother(int e, const Vector& in, Vector& out) const
+{
+   DenseMatrix elmat;
+   Array<int> local_dofs;
+   fespace_.GetElementDofs(e, local_dofs);
+   for (int i = 0; i < in.Size(); ++i)
+   {
+      out[i] = (scale_ / global_diag_(local_dofs[i])) * in[i];
+   }
+}
+
+}

fem/elementwisesmoother.hpp

@@ -0,0 +1,102 @@
+// Copyright (c) 2010, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-443211. All Rights
+// reserved. See file COPYRIGHT for details.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability see http://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the GNU Lesser General Public License (as published by the Free
+// Software Foundation) version 2.1 dated February 1999.
+
+#ifndef MFEM_ELEMENTWISESMOOTHER
+#define MFEM_ELEMENTWISESMOOTHER
+
+#include "../config/config.hpp"
+#include "fespace.hpp"
+
+#ifdef MFEM_USE_MPI
+#include <mpi.h>
+#endif
+
+namespace mfem
+{
+
+/**
+   Applies a smoother element-by-element.
+
+   The interface basically requires you to implement GetElementFromMatVec()
+   and LocalSmoother()
+*/
+class ElementWiseSmoother : public mfem::Solver
+{
+public:
+   ElementWiseSmoother(const mfem::FiniteElementSpace& fespace);
+   virtual ~ElementWiseSmoother() {}
+
+   virtual void SetOperator(const mfem::Operator &op) {}
+
+   virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const;
+
+   /**
+      b and r are local to the element, but x has to be global.
+   */
+   virtual void ElementResidual(int element, const mfem::Vector& b,
+                                const mfem::Vector& x,
+                                mfem::Vector& r) const;
+
+   /**
+      This produces the same result in y_element as if you did a
+      global matvec with (global) x, and then extracted the
+      dofs corresponding to a particular element.
+
+      Since elements have a bounded number of neighbors, this can
+      be done more efficiently than a global matvec.
+   */
+   virtual void GetElementFromMatVec(int element, const mfem::Vector& x,
+                                     mfem::Vector& y_element) const = 0;
+
+   /**
+      Interface is that `in` should be a residual, and then `out` is a correction
+      to the solution. (Ie, in is not a right-hand side, and out is not a current
+      iterate; out gets totally overwritten, not updated.)
+
+      I am not sure this is the *correct* interface for this use, but it is the
+      *current* interface.
+   */
+   virtual void LocalSmoother(int element, const mfem::Vector& in,
+                              mfem::Vector& out) const = 0;
+
+protected:
+   const mfem::FiniteElementSpace& fespace_;
+};
+
+/**
+   Our first implementation is Jacobi within elements and Gauss-Seidel
+   between them.
+*/
+class ElementWiseJacobi : public ElementWiseSmoother
+{
+public:
+   ElementWiseJacobi(const mfem::FiniteElementSpace& fespace,
+                     mfem::BilinearForm& aform,
+                     const mfem::Vector& global_diag,
+                     double scale=1.0);
+
+   virtual void GetElementFromMatVec(int element, const mfem::Vector& x,
+                                     mfem::Vector& y) const;
+
+   virtual void LocalSmoother(int element, const mfem::Vector& in,
+                              mfem::Vector& out) const;
+
+protected:
+   mfem::BilinearForm& aform_;
+   const mfem::Vector& global_diag_;
+   double scale_;
+
+   mfem::Table el_to_el_;
+};
+
+}
+
+#endif // MFEM_ELEMENTWISESMOOTHER

fem/fem.hpp

@@ -31,6 +31,7 @@
 #include "estimators.hpp"
 #include "staticcond.hpp"
 #include "tmop.hpp"
+#include "elementwisesmoother.hpp"
 
 #ifdef MFEM_USE_MPI
 #include "pfespace.hpp"

fem/pbilinearform.cpp

@@ -246,6 +246,13 @@ void ParBilinearForm::Assemble(int skip_zeros)
    }
 }
 
+void ParBilinearForm::AssembleDiagonal(Vector& diag) const
+{
+   Vector local(pfes->GetVSize());
+   BilinearForm::AssembleDiagonal(local);
+   pfes->GetProlongationMatrix()->MultTranspose(local, diag);
+}
+
 void ParBilinearForm
 ::ParallelEliminateEssentialBC(const Array<int> &bdr_attr_is_ess,
                                HypreParMatrix &A, const HypreParVector &X,

fem/pbilinearform.hpp

@@ -89,6 +89,9 @@ public:
    /// Assemble the local matrix
    void Assemble(int skip_zeros = 1);
 
+   /// Assemble local diagonal (on truedofs)
+   void AssembleDiagonal(Vector& diag) const;
+
    /// Returns the matrix assembled on the true dofs, i.e. P^t A P.
    /** The returned matrix has to be deleted by the caller. */
    HypreParMatrix *ParallelAssemble() { return ParallelAssemble(mat); }

linalg/solvers.cpp

@@ -10,6 +10,7 @@
 // Software Foundation) version 2.1 dated February 1999.
 
 #include "linalg.hpp"
+#include "../general/forall.hpp"
 #include "../general/globals.hpp"
 #include <iostream>
 #include <iomanip>
@@ -103,6 +104,49 @@ void IterativeSolver::SetOperator(const Operator &op)
 }
 
 
+OperatorJacobiSmoother::OperatorJacobiSmoother(const Vector &d,
+                                               const Array<int>& ess_tdofs,
+                                               const double dmpng)
+   :
+   Solver(d.Size()),
+   N(d.Size()),
+   dinv(N),
+   diag(d),
+   damping(dmpng),
+   ess_tdof_list(ess_tdofs),
+   residual(N) { Setup(); }
+
+void OperatorJacobiSmoother::Setup()
+{
+   residual.UseDevice(true);
+   const double delta = damping;
+   auto D = diag.Read();
+   auto X = dinv.Write();
+   MFEM_FORALL(i, N, X[i] = delta / D[i]; );
+   auto I = ess_tdof_list.Read();
+   MFEM_FORALL(i, ess_tdof_list.Size(), X[I[i]] = delta; );
+}
+
+void OperatorJacobiSmoother::Mult(const Vector& x, Vector &y) const
+{
+   if (iterative_mode && oper)
+   {
+      oper->Mult(y, residual);  // r = A x
+      subtract(x, residual, residual); // r = b - A x
+   }
+   else
+   {
+      residual = x;
+      y.UseDevice(true);
+      y = 0.0;
+   }
+   auto X = dinv.Read();
+   auto R = residual.Read();
+   auto Y = y.ReadWrite();
+   MFEM_FORALL(i, N, Y[i] += X[i] * R[i]; );
+}
+
+
 void SLISolver::UpdateVectors()
 {
    r.SetSize(width);

linalg/solvers.hpp

@@ -75,6 +75,44 @@ public:
 };
 
 
+/// Jacobi smoothing with given vector, no matrix necessary
+/** Potentially useful with tensorized operators, for example.
+    This is just a very basic Jacobi iteration, if you want
+    tolerances, iteration control, etc. wrap this with SLISolver. */
+class OperatorJacobiSmoother : public Solver
+{
+public:
+   /** Application is by *inverse* of the given vector.
+       It is assumed the underlying operator acts as the identity
+       on entries in ess_tdof_list, corresponding to (assembled) DIAG_ONE
+       policy or ConstratinedOperator in the matrix-free setting. */
+   OperatorJacobiSmoother(const Vector &d,
+                          const Array<int>& ess_tdof_list,
+                          const double damping=1.0);
+   ~OperatorJacobiSmoother() {}
+
+   void Mult(const Vector&x, Vector &y) const;
+
+   void SetOperator(const Operator &op_)
+   {
+      oper = &op_;
+   }
+
+   void Setup();
+
+private:
+   const int N;
+   Vector dinv;
+   const Vector &diag;
+   const double damping;
+   const Array<int>& ess_tdof_list;
+   mutable Vector residual;
+   /// could use IterativeSolver as base class to have this
+   /// but don't want tolerances, preconditioner, etc.
+   const Operator* oper;
+};
+
+
 /// Stationary linear iteration: x <- x + B (b - A x)
 class SLISolver : public IterativeSolver
 {

tests/unit/fem/test_assemblediagonalpa.cpp

@@ -0,0 +1,117 @@
+// Copyright (c) 2010, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-443211. All Rights
+// reserved. See file COPYRIGHT for details.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability see http://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the GNU Lesser General Public License (as published by the Free
+// Software Foundation) version 2.1 dated February 1999.
+
+#include "mfem.hpp"
+#include "catch.hpp"
+
+using namespace mfem;
+
+namespace assemblediagonalpa
+{
+
+TEST_CASE("massdiag")
+{
+   for (int dimension = 2; dimension < 4; ++dimension)
+   {
+      for (int ne = 1; ne < 3; ++ne)
+      {
+         std::cout << "Testing " << dimension << "D partial assembly mass diagonal: "
+                   << std::pow(ne, dimension) << " elements." << std::endl;
+         for (int order = 1; order < 5; ++order)
+         {
+            Mesh * mesh;
+            if (dimension == 2)
+            {
+               mesh = new Mesh(ne, ne, Element::QUADRILATERAL, 1, 1.0, 1.0);
+            }
+            else
+            {
+               mesh = new Mesh(ne, ne, ne, Element::HEXAHEDRON, 1, 1.0, 1.0, 1.0);
+            }
+            FiniteElementCollection *h1_fec = new H1_FECollection(order, dimension);
+            FiniteElementSpace h1_fespace(mesh, h1_fec);
+            BilinearForm paform(&h1_fespace);
+            ConstantCoefficient one(1.0);
+            paform.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+            paform.AddDomainIntegrator(new MassIntegrator(one));
+            paform.Assemble();
+            Vector pa_diag(h1_fespace.GetVSize());
+            paform.AssembleDiagonal(pa_diag);
+
+            BilinearForm assemblyform(&h1_fespace);
+            assemblyform.AddDomainIntegrator(new MassIntegrator(one));
+            assemblyform.Assemble();
+            assemblyform.Finalize();
+            Vector assembly_diag(h1_fespace.GetVSize());
+            assemblyform.SpMat().GetDiag(assembly_diag);
+
+            assembly_diag -= pa_diag;
+            double error = assembly_diag.Norml2();
+            std::cout << "    order: " << order << ", error norm: " << error << std::endl;
+            REQUIRE(assembly_diag.Norml2() < 1.e-12);
+
+            delete mesh;
+            delete h1_fec;
+         }
+      }
+   }
+}
+
+TEST_CASE("diffusiondiag")
+{
+   for (int dimension = 2; dimension < 4; ++dimension)
+   {
+      for (int ne = 1; ne < 3; ++ne)
+      {
+         std::cout << "Testing " << dimension <<
+                   "D partial assembly diffusion diagonal: "
+                   << std::pow(ne, dimension) << " elements." << std::endl;
+         for (int order = 1; order < 5; ++order)
+         {
+            Mesh * mesh;
+            if (dimension == 2)
+            {
+               mesh = new Mesh(ne, ne, Element::QUADRILATERAL, 1, 1.0, 1.0);
+            }
+            else
+            {
+               mesh = new Mesh(ne, ne, ne, Element::HEXAHEDRON, 1, 1.0, 1.0, 1.0);
+            }
+            FiniteElementCollection *h1_fec = new H1_FECollection(order, dimension);
+            FiniteElementSpace h1_fespace(mesh, h1_fec);
+            BilinearForm paform(&h1_fespace);
+            ConstantCoefficient one(1.0);
+            paform.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+            paform.AddDomainIntegrator(new DiffusionIntegrator(one));
+            paform.Assemble();
+            Vector pa_diag(h1_fespace.GetVSize());
+            paform.AssembleDiagonal(pa_diag);
+
+            BilinearForm assemblyform(&h1_fespace);
+            assemblyform.AddDomainIntegrator(new DiffusionIntegrator(one));
+            assemblyform.Assemble();
+            assemblyform.Finalize();
+            Vector assembly_diag(h1_fespace.GetVSize());
+            assemblyform.SpMat().GetDiag(assembly_diag);
+
+            assembly_diag -= pa_diag;
+            double error = assembly_diag.Norml2();
+            std::cout << "    order: " << order << ", error norm: " << error << std::endl;
+            REQUIRE(assembly_diag.Norml2() < 1.e-12);
+
+            delete mesh;
+            delete h1_fec;
+         }
+      }
+   }
+}
+
+} // namespace assemblediagonalpa

tests/unit/fem/test_elementmatmult.cpp

@@ -0,0 +1,105 @@
+// Copyright (c) 2010, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-443211. All Rights
+// reserved. See file COPYRIGHT for details.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability see http://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the GNU Lesser General Public License (as published by the Free
+// Software Foundation) version 2.1 dated February 1999.
+
+#include "mfem.hpp"
+#include "catch.hpp"
+
+using namespace mfem;
+
+namespace elementmatmult
+{
+
+TEST_CASE("elementmatmult")
+{
+   for (int dimension = 2; dimension < 4; ++dimension)
+   {
+      for (int integrator = 0; integrator < 3; ++integrator)
+      {
+         for (int ne = 1; ne < 3; ++ne)
+         {
+            std::cout << "Testing " << dimension << "D partial assembly element mat mult: "
+                      << std::pow(ne, dimension) << " elements,"
+                      << " integrator " << integrator << std::endl;
+            for (int order = 1; order < 5; ++order)
+            {
+               Mesh * mesh;
+               if (dimension == 2)
+               {
+                  mesh = new Mesh(ne, ne, Element::QUADRILATERAL, 1, 1.0, 1.0);
+               }
+               else
+               {
+                  mesh = new Mesh(ne, ne, ne, Element::HEXAHEDRON, 1, 1.0, 1.0, 1.0);
+               }
+               FiniteElementCollection *h1_fec = new H1_FECollection(order, dimension);
+               FiniteElementSpace h1_fespace(mesh, h1_fec);
+               Array<int> ess_tdof_list;
+               // Array<int> ess_bdr(mesh->bdr_attributes.Max());
+               // ess_bdr = 1;
+               // h1_fespace.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+               BilinearForm paform(&h1_fespace);
+               ConstantCoefficient one(1.0);
+               paform.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+               if (integrator < 2)
+               {
+                  paform.AddDomainIntegrator(new DiffusionIntegrator(one));
+               }
+               if (integrator > 0)
+               {
+                  paform.AddDomainIntegrator(new MassIntegrator(one));
+               }
+               paform.Assemble();
+
+               GridFunction x(&h1_fespace);
+               x = 0.0;
+               GridFunction b(&h1_fespace);
+               b = 1.0;
+               BilinearForm assemblyform(&h1_fespace);
+               if (integrator < 2)
+               {
+                  assemblyform.AddDomainIntegrator(new DiffusionIntegrator(one));
+               }
+               if (integrator > 0)
+               {
+                  assemblyform.AddDomainIntegrator(new MassIntegrator(one));
+               }
+               // assemblyform.SetDiagonalPolicy(Matrix::DIAG_ONE); // todo?
+               assemblyform.Assemble();
+               assemblyform.Finalize();
+               OperatorPtr A_assembly;
+               Vector B, X;
+               assemblyform.FormLinearSystem(ess_tdof_list, x, b, A_assembly, X, B);
+
+               for (int elno = 0; elno < std::min(8, mesh->GetNE()); ++elno)
+               {
+                  Vector xin(h1_fespace.GetTrueVSize());
+                  xin.Randomize();
+                  Vector y_assembly, y_pa;
+                  assemblyform.ElementMatrixMult(elno, xin, y_assembly);
+                  paform.ElementMatrixMult(elno, xin, y_pa);
+
+                  y_assembly -= y_pa;
+                  double error = y_assembly.Norml2();
+                  std::cout << "    order: " << order << ", elno " << elno
+                            << ", error norm: " << error << std::endl;
+                  REQUIRE(y_assembly.Norml2() < 1.e-14);
+               }
+
+               delete mesh;
+               delete h1_fec;
+            }
+         }
+      }
+   }
+}
+
+} // namespace elementmatmult

tests/unit/fem/test_elementwisesmoother.cpp

@@ -0,0 +1,109 @@
+// Copyright (c) 2010, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-443211. All Rights
+// reserved. See file COPYRIGHT for details.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability see http://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the GNU Lesser General Public License (as published by the Free
+// Software Foundation) version 2.1 dated February 1999.
+
+#include "mfem.hpp"
+#include "catch.hpp"
+
+using namespace mfem;
+
+namespace elementwisesmoother
+{
+
+TEST_CASE("elementwisesmoother")
+{
+   for (int dimension = 2; dimension < 4; ++dimension)
+   {
+      for (int integrator = 0; integrator < 2; ++integrator)
+      {
+         const int ne = 1;
+         std::cout << "Testing " << dimension << "D element-wise smoother with "
+                   << "integrator " << integrator << std::endl;
+         for (int order = 1; order < 5; ++order)
+         {
+            Mesh * mesh;
+            if (dimension == 2)
+            {
+               mesh = new Mesh(ne, ne, Element::QUADRILATERAL, 1, 1.0, 1.0);
+            }
+            else
+            {
+               mesh = new Mesh(ne, ne, ne, Element::HEXAHEDRON, 1, 1.0, 1.0, 1.0);
+            }
+            FiniteElementCollection *h1_fec = new H1_FECollection(order, dimension);
+            FiniteElementSpace h1_fespace(mesh, h1_fec);
+            Array<int> ess_tdof_list;
+ 
+            BilinearForm paform(&h1_fespace);
+            ConstantCoefficient one(1.0); 
+            paform.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+            if (integrator < 2)
+            {
+               paform.AddDomainIntegrator(new DiffusionIntegrator(one));
+            }
+            if (integrator > 0)
+            {
+               paform.AddDomainIntegrator(new MassIntegrator(one));
+            }
+            paform.Assemble();
+
+            BilinearForm assemblyform(&h1_fespace);
+            if (integrator < 2)
+            {
+               assemblyform.AddDomainIntegrator(new DiffusionIntegrator(one));
+            }
+            if (integrator > 0)
+            {
+               assemblyform.AddDomainIntegrator(new MassIntegrator(one));
+            }
+            assemblyform.SetDiagonalPolicy(Matrix::DIAG_ONE);
+            assemblyform.Assemble();
+            assemblyform.Finalize();
+            const SparseMatrix& A_explicit = assemblyform.SpMat();
+
+            Vector pa_diag(h1_fespace.GetVSize());
+            paform.AssembleDiagonal(pa_diag);
+            ElementWiseJacobi e_pa_jacobi(h1_fespace, paform, pa_diag);
+            ElementWiseJacobi e_assembly_jacobi(h1_fespace, assemblyform, pa_diag);
+            DSmoother mat_jacobi(A_explicit);
+
+            Vector xin(h1_fespace.GetTrueVSize());
+            xin.Randomize();
+            Vector y_mat(xin);
+            y_mat = 0.0;
+            Vector y_assembly(xin);
+            y_assembly = 0.0;
+            Vector y_pa(xin);
+            y_pa = 0.0;
+
+            e_pa_jacobi.Mult(xin, y_pa);
+            e_assembly_jacobi.Mult(xin, y_assembly);
+            mat_jacobi.Mult(xin, y_mat);
+
+            y_pa -= y_mat;
+            double pa_error = y_pa.Norml2();
+            std::cout << "  order: " << order << ", pa error norm: " << pa_error
+                      << std::endl;
+            REQUIRE(pa_error < 1.e-12);
+
+            y_assembly -= y_mat;
+            double assembly_error = y_assembly.Norml2();
+            std::cout << "  order: " << order << ", assembly error norm: "
+                      << assembly_error << std::endl;
+            REQUIRE(assembly_error < 1.e-12);
+
+            delete mesh;
+            delete h1_fec;
+         }
+      }
+   }
+}
+
+} // namespace elementwisesmoother

tests/unit/fem/test_operatorjacobismoother.cpp

@@ -0,0 +1,90 @@
+// Copyright (c) 2010, Lawrence Livermore National Security, LLC. Produced at
+// the Lawrence Livermore National Laboratory. LLNL-CODE-443211. All Rights
+// reserved. See file COPYRIGHT for details.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability see http://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the GNU Lesser General Public License (as published by the Free
+// Software Foundation) version 2.1 dated February 1999.
+
+#include "mfem.hpp"
+#include "catch.hpp"
+
+using namespace mfem;
+
+namespace operatorjacobismoother
+{
+
+TEST_CASE("operatorjacobismoother")
+{
+   for (int dimension = 2; dimension < 4; ++dimension)
+   {
+      for (int ne = 1; ne < 3; ++ne)
+      {
+         std::cout << "Testing " << dimension << "D partial assembly smoother: "
+                   << std::pow(ne, dimension) << " elements." << std::endl;
+         for (int order = 1; order < 5; ++order)
+         {
+            Mesh * mesh;
+            if (dimension == 2)
+            {
+               mesh = new Mesh(ne, ne, Element::QUADRILATERAL, 1, 1.0, 1.0);
+            }
+            else
+            {
+               mesh = new Mesh(ne, ne, ne, Element::HEXAHEDRON, 1, 1.0, 1.0, 1.0);
+            }
+            FiniteElementCollection *h1_fec = new H1_FECollection(order, dimension);
+            FiniteElementSpace h1_fespace(mesh, h1_fec);
+            Array<int> ess_tdof_list;
+            Array<int> ess_bdr(mesh->bdr_attributes.Max());
+            ess_bdr = 1;
+            h1_fespace.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+            BilinearForm paform(&h1_fespace);
+            ConstantCoefficient one(1.0);
+            paform.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+            paform.AddDomainIntegrator(new DiffusionIntegrator(one));
+            paform.Assemble();
+            Vector pa_diag(h1_fespace.GetVSize());
+            paform.AssembleDiagonal(pa_diag);
+            OperatorJacobiSmoother pa_smoother(pa_diag, ess_tdof_list);
+
+            GridFunction x(&h1_fespace);
+            x = 0.0;
+            GridFunction b(&h1_fespace);
+            b = 1.0;
+            BilinearForm assemblyform(&h1_fespace);
+            assemblyform.AddDomainIntegrator(new DiffusionIntegrator(one));
+            assemblyform.SetDiagonalPolicy(Matrix::DIAG_ONE);
+            assemblyform.Assemble();
+            assemblyform.Finalize();
+            OperatorPtr A_assembly;
+            Vector B, X;
+            assemblyform.FormLinearSystem(ess_tdof_list, x, b, A_assembly, X, B);
+            DSmoother assembly_smoother((SparseMatrix&)(*A_assembly));
+
+            Vector xin(h1_fespace.GetTrueVSize());
+            xin.Randomize();
+            Vector y_assembly(xin);
+            y_assembly = 0.0;
+            Vector y_pa(xin);
+            y_pa = 0.0;
+            assembly_smoother.Mult(xin, y_assembly);
+            pa_smoother.Mult(xin, y_pa);
+
+            y_assembly -= y_pa;
+            double error = y_assembly.Norml2();
+            std::cout << "    order: " << order << ", error norm: " << error << std::endl;
+            REQUIRE(y_assembly.Norml2() < 1.e-12);
+
+            delete mesh;
+            delete h1_fec;
+         }
+      }
+   }
+}
+
+} // namespace operatorjacobismoother