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.

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)
 {

fem/bilinearform.hpp

@@ -462,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

@@ -173,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

@@ -53,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
@@ -117,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

@@ -34,13 +34,20 @@ void BilinearFormIntegrator::AssembleDiagonalPA(Vector&) const
 
 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

@@ -56,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
@@ -1726,6 +1729,8 @@ public:
 
    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);
 };
@@ -1768,6 +1773,8 @@ public:
 
    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

@@ -118,7 +118,6 @@ static void PAMassAssembleDiagonal2D(const int NE,
    MFEM_VERIFY(D1D <= MAX_D1D, "");
    MFEM_VERIFY(Q1D <= MAX_Q1D, "");
    auto B = Reshape(_B.Read(), Q1D, D1D);
-   // auto Bt = Reshape(_Bt.Read(), D1D, Q1D); // ?? (TODO atb@llnl.gov)
    auto op = Reshape(_op.Read(), Q1D, Q1D, NE);
    auto y = Reshape(_diag.ReadWrite(), D1D, D1D, NE);
    MFEM_FORALL(e, NE,
@@ -170,7 +169,6 @@ static void PAMassAssembleDiagonal3D(const int NE,
    MFEM_VERIFY(D1D <= MAX_D1D, "");
    MFEM_VERIFY(Q1D <= MAX_Q1D, "");
    auto B = Reshape(_B.Read(), Q1D, D1D);
-   // auto Bt = Reshape(_Bt.Read(), D1D, Q1D); // ?? (TODO atb@llnl.gov)
    auto op = Reshape(_op.Read(), Q1D, Q1D, Q1D, NE);
    auto y = Reshape(_diag.ReadWrite(), D1D, D1D, D1D, NE);
    MFEM_FORALL(e, NE,
@@ -467,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>
@@ -492,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);
    });
 }
 
@@ -733,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)
 {
@@ -750,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);
+
    });
 }
 
@@ -966,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"

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