Fixed element access on shared faces with MPI. Added more coefficients to Parallel PA DG Diffusion test

Also add (failing) unit test for piecewise constant diffusion
coefficient

fem/coefficient.cpp

@@ -2195,6 +2195,8 @@ void CoefficientVector::SetConstant(const DenseSymmetricMatrix &constant)
 
 int CoefficientVector::GetVDim() const { return vdim; }
 
+void CoefficientVector::SetVDim(int vdim_) { vdim = vdim_; }
+
 CoefficientVector::~CoefficientVector()
 {
    delete qf;

fem/coefficient.hpp

@@ -2606,6 +2606,9 @@ public:
    /// Return the number of values per quadrature point.
    int GetVDim() const;
 
+   /// Override the vector dimension (number of values per quadrature point).
+   void SetVDim(int vdim_);
+
    ~CoefficientVector();
 };
 

fem/integ/bilininteg_dgdiffusion_kernels.hpp

@@ -43,7 +43,7 @@ static void PADGDiffusionApply2D(const int NF, const Array<real_t> &b,
    auto G_ = Reshape(g.Read(), Q1D, D1D);
 
    auto pa =
-      Reshape(pa_data.Read(), 6, Q1D, NF); // (q, 1/h, J00, J01, J10, J11)
+      Reshape(pa_data.Read(), 5, Q1D, NF); // (J00, J01, J10, J11, q/h)
 
    auto x = Reshape(x_.Read(), D1D, 2, NF);
    auto y = Reshape(y_.ReadWrite(), D1D, 2, NF);
@@ -109,8 +109,8 @@ static void PADGDiffusionApply2D(const int NF, const Array<real_t> &b,
 
          MFEM_FOREACH_THREAD(p, x, Q1D)
          {
-            const real_t Je_side[] = {pa(2 + 2 * side, p, f),
-                                      pa(2 + 2 * side + 1, p, f)
+            const real_t Je_side[] = {pa(2 * side + 0, p, f),
+                                      pa(2 * side + 1, p, f)
                                      };
 
             Bu[p] = 0.0;
@@ -133,11 +133,10 @@ static void PADGDiffusionApply2D(const int NF, const Array<real_t> &b,
       {
          MFEM_FOREACH_THREAD(p, x, Q1D)
          {
-            const real_t q = pa(0, p, f);
-            const real_t hi = pa(1, p, f);
+            const real_t q = pa(4, p, f);
             const real_t jump = Bu0[p] - Bu1[p];
             const real_t avg = Bdu0[p] + Bdu1[p]; // = {Q du/dn} * w * det(J)
-            r[p] = -avg + hi * q * jump;
+            r[p] = -avg + q * jump;
          }
       }
       MFEM_SYNC_THREAD;
@@ -173,8 +172,8 @@ static void PADGDiffusionApply2D(const int NF, const Array<real_t> &b,
          {
             for (int p = 0; p < Q1D; ++p)
             {
-               const real_t Je[] = {pa(2 + 2 * side, p, f),
-                                    pa(2 + 2 * side + 1, p, f)
+               const real_t Je[] = {pa(2 * side + 0, p, f),
+                                    pa(2 * side + 1, p, f)
                                    };
                const real_t jump = Bu0[p] - Bu1[p];
                const real_t r_p = Je[0] * jump; // normal

fem/integ/bilininteg_dgdiffusion_pa.cpp

@@ -14,6 +14,7 @@
 #include "../fe/face_map_utils.hpp"
 #include "../gridfunc.hpp"
 #include "../qfunction.hpp"
+#include "../../mesh/pmesh.hpp"
 
 #include "bilininteg_dgdiffusion_kernels.hpp"
 
@@ -42,25 +43,24 @@ static void PADGDiffusionSetup2D(const int Q1D, const int NE, const int NF,
    const auto n = Reshape(face_geom.normal.Read(), Q1D, 2, NF);
 
    const bool const_q = (q.Size() == coeff_dim);
-   const auto Q = const_q ? Reshape(q.Read(), coeff_dim, 1, 1)
-                  : Reshape(q.Read(), coeff_dim, Q1D, NF);
+   const auto Q = const_q ? Reshape(q.Read(), coeff_dim, 1, 1, 1)
+                  : Reshape(q.Read(), coeff_dim, Q1D, 2, NF);
 
    const auto W = w.Read();
 
    // (normal0, normal1, e0, e1, fid0, fid1)
    const auto face_info = Reshape(face_info_.Read(), 6, NF);
 
-   // (q, 1/h, J0_0, J0_1, J1_0, J1_1)
-   auto pa = Reshape(pa_data.Write(), 6, Q1D, NF);
-
-   auto get_coeff = [const_q] MFEM_HOST_DEVICE (const decltype(Q) &Q, int i,
-                                                int qx, int e)
-   {
-      return const_q ? Q(i,0,0) : Q(i,qx,e);
-   };
+   // (J0_0, J0_1, J1_0, J1_1, {Q/h})
+   auto pa = Reshape(pa_data.Write(), 5, Q1D, NF);
 
    mfem::forall(NF, [=] MFEM_HOST_DEVICE(int f) -> void
    {
+      auto get_coeff = [&] (int i, int qx, int side, int e)
+      {
+         return const_q ? Q(i, 0, 0, 0) : Q(i, qx, side, e);
+      };
+
       const int normal_dir[] = {face_info(0, f), face_info(1, f)};
       const int fid[] = {face_info(4, f), face_info(5, f)};
 
@@ -77,28 +77,28 @@ static void PADGDiffusionSetup2D(const int Q1D, const int NE, const int NF,
 
       for (int p = 0; p < Q1D; ++p)
       {
-         real_t qh = 0.0;
-         real_t hi = 0.0;
+         const real_t nvec[2] = {n(p, 0, f), n(p, 1, f)};
+         const real_t dJf = detJf(p, f);
 
-         real_t Qtn[2];
-         if (coeff_dim > 1)
-         {
-            // matrix coefficient
-            Qtn[0] = get_coeff(Q,0,p,f)*n(p,0,f) + get_coeff(Q,1,p,f)*n(p,1,f);
-            Qtn[1] = get_coeff(Q,2,p,f)*n(p,0,f) + get_coeff(Q,3,p,f)*n(p,1,f);
-            qh = Qtn[0]*n(p,0,f) + Qtn[1]*n(p,1,f);
-         }
-         else
-         {
-            qh = get_coeff(Q, 0, p, f);
-            Qtn[0] = qh*n(p,0,f);
-            Qtn[1] = qh*n(p,1,f);
-         }
-
-         pa(0, p, f) = kappa * qh * W[p] * detJf(p, f);
+         real_t qhi = 0.0;
 
          for (int side = 0; side < nsides; ++side)
          {
+            real_t Qtn[2];
+            real_t nqn = 0.0;
+            if (coeff_dim > 1)
+            {
+               Qtn[0] = get_coeff(0, p, side, f) * nvec[0] + get_coeff(1, p, side, f) * nvec[1];
+               Qtn[1] = get_coeff(2, p, side, f) * nvec[0] + get_coeff(3, p, side, f) * nvec[1];
+               nqn = Qtn[0] * nvec[0] + Qtn[1] * nvec[1];
+            }
+            else
+            {
+               nqn = get_coeff(0, p, side, f);
+               Qtn[0] = nqn * nvec[0];
+               Qtn[1] = nqn * nvec[1];
+            }
+
             int i, j;
             internal::FaceIdxToVolIdx2D(p, Q1D, fid[0], fid[1], side, i, j);
 
@@ -115,7 +115,6 @@ static void PADGDiffusionSetup2D(const int Q1D, const int NE, const int NF,
             nJi[1] = -Qtn[0]*J(i, j, 1, 0, e) + Qtn[1]*J(i, j, 0, 0, e);
 
             const real_t dJe = detJ(i, j, e);
-            const real_t dJf = detJf(p, f);
 
             const real_t w = factor * W[p] * dJf / dJe;
 
@@ -123,20 +122,20 @@ static void PADGDiffusionSetup2D(const int Q1D, const int NE, const int NF,
             const int ti = 1 - ni;
 
             // Normal
-            pa(2 + 2 * side + 0, p, f) = w * nJi[ni];
+            pa(2 * side + 0, p, f) = w * nJi[ni];
             // Tangential
-            pa(2 + 2 * side + 1, p, f) = sgn * w * nJi[ti];
+            pa(2 * side + 1, p, f) = sgn * w * nJi[ti];
 
-            hi += factor * dJf / dJe;
+            qhi += kappa * nqn * w * dJf;
          }
 
          if (nsides == 1)
          {
-            pa(4, p, f) = 0.0;
-            pa(5, p, f) = 0.0;
+            pa(2, p, f) = 0.0;
+            pa(3, p, f) = 0.0;
          }
 
-         pa(1, p, f) = hi;
+         pa(4, p, f) = qhi; // {Q/h}
       }
    });
 }
@@ -163,8 +162,8 @@ static void PADGDiffusionSetup3D(const int Q1D, const int NE, const int NF,
    const auto n = Reshape(face_geom.normal.Read(), Q1D, Q1D, 3, NF);
 
    const bool const_q = (q.Size() == coeff_dim);
-   const auto Q = const_q ? Reshape(q.Read(), coeff_dim, 1, 1, 1)
-                  : Reshape(q.Read(), coeff_dim, Q1D, Q1D, NF);
+   const auto Q = const_q ? Reshape(q.Read(), coeff_dim, 1, 1, 1, 1)
+                  : Reshape(q.Read(), coeff_dim, Q1D, Q1D, 2, NF);
 
    const auto W = Reshape(w.Read(), Q1D, Q1D);
 
@@ -174,17 +173,16 @@ static void PADGDiffusionSetup3D(const int Q1D, const int NE, const int NF,
    constexpr int _fid_ = 4; // offset in face_info for local face id
    constexpr int _or_ = 5;  // offset in face_info for orientation
 
-   // (J00, J01, J02, J10, J11, J12, q/h)
+   // (J00, J01, J02, J10, J11, J12, {Q/h})
    const auto pa = Reshape(pa_data.Write(), 7, Q1D, Q1D, NF);
 
-   auto get_coeff = [const_q] MFEM_HOST_DEVICE (const decltype(Q) &Q, int i,
-                                                int qx, int qy, int e)
-   {
-      return const_q ? Q(i,0,0,0) : Q(i,qx,qy,e);
-   };
-
    mfem::forall_2D(NF, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int f) -> void
    {
+      auto get_coeff = [&] (int i, int qx, int qy, int side, int e)
+      {
+         return const_q ? Q(i, 0, 0, 0, 0) : Q(i, qx, qy, side, e);
+      };
+
       MFEM_SHARED int perm[2][3];
       MFEM_SHARED int el[2];
       MFEM_SHARED bool shared[2];
@@ -219,33 +217,34 @@ static void PADGDiffusionSetup3D(const int Q1D, const int NE, const int NF,
          MFEM_FOREACH_THREAD(p2, y, Q1D)
          {
             const real_t dJf = detJf(p1, p2, f);
+            const real_t nvec[3] = {n(p1, p2, 0, f), n(p1, p2, 1, f), n(p1, p2, 2, f)};
 
-            real_t hi = 0.0;
-
-            real_t Qtn[3];
-            real_t qh = 0.0;
-
-            if (coeff_dim > 1)
-            {
-               // matrix coefficient
-               for (int d = 0; d < 3; ++d)
-               {
-                  Qtn[d] = get_coeff(Q,0+3*d,p1,p2,f)*n(p1,p2,0,f)
-                           + get_coeff(Q,1+3*d,p1,p2,f)*n(p1,p2,1,f)
-                           + get_coeff(Q,2+3*d,p1,p2,f)*n(p1,p2,2,f);
-                  qh += Qtn[d] * n(p1,p2,d,f);
-               }
-            }
-            else
-            {
-               qh = get_coeff(Q,0,p1,p2,f);
-               Qtn[0] = qh * n(p1,p2,0,f);
-               Qtn[1] = qh * n(p1,p2,1,f);
-               Qtn[2] = qh * n(p1,p2,2,f);
-            }
+            real_t qhi = 0.0;
 
             for (int side = 0; side < nsides; ++side)
             {
+               real_t Qtn[3]; // n' * Q
+               real_t nqn = 0.0; // n' * Q * n
+
+               if (coeff_dim > 1)
+               {
+                  // matrix coefficient
+                  for (int d = 0; d < 3; ++d)
+                  {
+                     Qtn[d] = get_coeff(0+3*d, p1, p2, side, f) * nvec[0]
+                              + get_coeff(1+3*d, p1, p2, side, f) * nvec[1]
+                              + get_coeff(2+3*d, p1, p2, side, f) * nvec[2];
+                     nqn += Qtn[d] * nvec[d];
+                  }
+               }
+               else
+               {
+                  nqn = get_coeff(0, p1, p2, side, f);
+                  for (int d = 0; d < 3; ++d)
+                     Qtn[d] = nqn * nvec[d];
+               }
+
+
                int i, j, k;
                internal::FaceIdxToVolIdx3D(p1 + Q1D * p2, Q1D, fid[0], fid[1],
                                            side, ortn[1], i, j, k);
@@ -279,16 +278,16 @@ static void PADGDiffusionSetup3D(const int Q1D, const int NE, const int NF,
                // *INDENT-ON*
 
                const real_t dJe = detJe(i, j, k, e);
-               const real_t val = factor * W(p1, p2) * dJf / dJe;
+               const real_t w = factor * W(p1, p2) * dJf / dJe;
 
                for (int d = 0; d < 3; ++d)
                {
                   const int idx = std::abs(perm[side][d]) - 1;
                   const int sgn = (perm[side][d] < 0) ? -1 : 1;
-                  pa(3 * side + d, p1, p2, f) = sgn * val * nJi[idx];
+                  pa(3 * side + d, p1, p2, f) = sgn * w * nJi[idx];
                }
 
-               hi += factor * dJf / dJe;
+               qhi += kappa * nqn * w * dJf;
             }
 
             if (nsides == 1)
@@ -298,7 +297,7 @@ static void PADGDiffusionSetup3D(const int Q1D, const int NE, const int NF,
                pa(5, p1, p2, f) = 0.0;
             }
 
-            pa(6, p1, p2, f) = kappa * hi * qh * W(p1, p2) * dJf;
+            pa(6, p1, p2, f) = qhi; // {Q/h}
          }
       }
    });
@@ -535,15 +534,94 @@ void DGDiffusionIntegrator::SetupPA(const FiniteElementSpace &fes,
    dofs1D = maps->ndof;
    quad1D = maps->nqpt;
 
-   const int pa_size = (dim == 2) ? (6 * q1d * nf) : (7 * q1d * q1d * nf);
+   const int pa_size = (dim == 2) ? (5 * q1d * nf) : (7 * q1d * q1d * nf);
    pa_data.SetSize(pa_size, Device::GetMemoryType());
 
    // Evaluate the coefficient at the face quadrature points.
    FaceQuadratureSpace fqs(mesh, ir, type);
    CoefficientVector q(fqs, CoefficientStorage::CONSTANTS);
-   if (Q) { q.Project(*Q); }
-   else if (MQ) { q.Project(*MQ); }
-   else { q.SetConstant(1.0); }
+
+   if (Q == nullptr && MQ == nullptr)
+   {
+      q.SetConstant(1.0);
+   }
+   else if (auto *const_q = dynamic_cast<ConstantCoefficient*>(Q))
+   {
+      q.SetConstant(const_q->constant);
+   }
+   else if (auto *const_mq = dynamic_cast<MatrixConstantCoefficient*>(MQ))
+   {
+      q.SetConstant(const_mq->GetMatrix());
+   }
+   else
+   {
+      const int q_dim = MQ ? dim*dim : 1;
+      nq = ir.GetNPoints();
+      q.SetSize(2 * nq * nf * q_dim);
+      q.SetVDim(q_dim);
+      auto C = Reshape(q.HostWrite(), q_dim, nq, 2, nf);
+
+      auto *pmesh = dynamic_cast<ParMesh*>(&mesh);
+
+      int f_ind = 0;
+      for (int f = 0; f < mesh.GetNumFacesWithGhost(); ++f)
+      {
+         Mesh::FaceInformation face = mesh.GetFaceInformation(f);
+         if (face.IsNonconformingCoarse() || !face.IsOfFaceType(type))
+         {
+            // We skip nonconforming coarse faces as they are treated
+            // by the corresponding nonconforming fine faces.
+            continue;
+         }
+         FaceElementTransformations &T = [&] () -> FaceElementTransformations&
+         {
+            if (face.IsShared() && pmesh)
+               return *pmesh->GetSharedFaceTransformationsByLocalIndex(f, true);
+            return *mesh.GetFaceElementTransformations(f);
+         }();
+
+
+         for (int q = 0; q < nq; ++q)
+         {
+            // Convert to lexicographic ordering
+            int iq =
+               ToLexOrdering(dim, face.element[0].local_face_id, quad1D, q);
+
+            T.SetAllIntPoints(&ir.IntPoint(q));
+            const IntegrationPoint &eip1 = T.GetElement1IntPoint();
+            const IntegrationPoint &eip2 = T.GetElement2IntPoint();
+
+            if (Q)
+            {
+               real_t q_val = Q->Eval(*T.Elem1, eip1);
+               C(0, iq, 0, f_ind) = q_val;
+               C(0, iq, 1, f_ind) = (face.IsInterior()) ? Q->Eval(*T.Elem2, eip2) : q_val;
+            }
+            else if (MQ)
+            {
+               DenseMatrix mq_val(dim, dim);
+               MQ->Eval(mq_val, *T.Elem1, eip1);
+
+               DenseMatrix mq_val_2(dim, dim);
+               if (face.IsInterior())
+                  MQ->Eval(mq_val_2, *T.Elem2, eip2);
+               else
+                  mq_val_2 = mq_val;
+
+               for (int j = 0; j < dim; ++j)
+               {
+                  for (int i = 0; i < dim; ++i)
+                  {
+                     C(i + dim*j, iq, 0, f_ind) = mq_val(i, j);
+                     C(i + dim*j, iq, 1, f_ind) = mq_val_2(i, j);
+                  }
+               }
+            }
+         }
+         f_ind++;
+      }
+      MFEM_VERIFY(f_ind == nf, "Incorrect number of faces.");
+   }
 
    const int coeff_dim = q.GetVDim();
 

tests/unit/fem/test_pa_kernels.cpp

@@ -841,6 +841,13 @@ MakeCoeff<SymmetricMatrixConstantCoefficient>(int dim)
    return std::make_unique<SymmetricMatrixConstantCoefficient>(A);
 }
 
+template <> std::unique_ptr<PWConstCoefficient>
+MakeCoeff<PWConstCoefficient>(int)
+{
+   Vector c({1.0, 2.0});
+   return std::make_unique<PWConstCoefficient>(c);
+}
+
 template <typename CoeffType = ConstantCoefficient,
           typename FES = FiniteElementSpace>
 void test_dg_diffusion(FES &fes)
@@ -916,12 +923,18 @@ TEST_CASE("PA DG Diffusion", "[PartialAssembly], [GPU]")
    Mesh mesh = Mesh::LoadFromFile(mesh_fname.c_str());
    const int dim = mesh.Dimension();
 
+   for (int i = 0; i < mesh.GetNE(); ++i)
+   {
+      mesh.SetAttribute(i, 1 + (i % 2));
+   }
+
    DG_FECollection fec(order, dim, BasisType::GaussLobatto);
    FiniteElementSpace fes(&mesh, &fec);
 
    test_dg_diffusion<ConstantCoefficient>(fes);
    test_dg_diffusion<MatrixConstantCoefficient>(fes);
    test_dg_diffusion<SymmetricMatrixConstantCoefficient>(fes);
+   test_dg_diffusion<PWConstCoefficient>(fes);
 }
 
 #ifdef MFEM_USE_MPI
@@ -933,6 +946,11 @@ TEST_CASE("Parallel PA DG Diffusion", "[PartialAssembly][Parallel][GPU]")
    CAPTURE(order, mesh_fname);
 
    Mesh serial_mesh = Mesh::LoadFromFile(mesh_fname.c_str());
+   for (int i = 0; i < serial_mesh.GetNE(); ++i)
+   {
+      serial_mesh.SetAttribute(i, 1 + (i % 2));
+   }
+
    ParMesh mesh(MPI_COMM_WORLD, serial_mesh);
    serial_mesh.Clear();
 
@@ -943,6 +961,8 @@ TEST_CASE("Parallel PA DG Diffusion", "[PartialAssembly][Parallel][GPU]")
 
    test_dg_diffusion<ConstantCoefficient>(fes);
    test_dg_diffusion<MatrixConstantCoefficient>(fes);
+   test_dg_diffusion<SymmetricMatrixConstantCoefficient>(fes);
+   test_dg_diffusion<PWConstCoefficient>(fes);
 }
 
 #endif