Simplifying lifting

cpp/dolfinx/fem/assemble_matrix_impl.h

@@ -30,6 +30,8 @@
 /// @brief Execute kernel over cells and accumulate result in a matrix.
 ///
 /// @tparam T Matrix/form scalar type.
+/// @tparam LiftingMode If set true, only execute mat_set on cells with BCs in
+/// column space.
 /// @param mat_set Function that accumulates computed entries into a
 /// matrix.
 /// @param[in] x_dofmap Degree-of-freedom map for the mesh geometry.
@@ -60,7 +62,7 @@
 /// function mesh.
 /// @param cell_info1 Cell permutation information for the trial
 /// function mesh.
-template <dolfinx::scalar T>
+template <dolfinx::scalar T, bool LiftingMode = false>
 void assemble_cells_matrix(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
     md::mdspan<const scalar_value_t<T>,
@@ -101,6 +103,29 @@
     std::int32_t cell0 = cells0[c];
     std::int32_t cell1 = cells1[c];
 
+    std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
+    std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
+
+    // In "LiftingMode" only execute kernel if there are BCs on column space
+    if constexpr (LiftingMode)
+    {
+      auto has_bc = [&]()
+      {
+        for (std::int32_t dof : dofs1)
+        {
+          for (int k = 0; k < bs1; ++k)
+          {
+            if (bc1[bs1 * dof + k])
+              return true;
+          }
+        }
+        return false;
+      };
+
+      if (!has_bc())
+        continue;
+    }
+
     // Get cell coordinates/geometry
     auto x_dofs = md::submdspan(x_dofmap, cell, md::full_extent);
     for (std::size_t i = 0; i < x_dofs.size(); ++i)
@@ -115,38 +140,38 @@
     P0(Ae, cell_info0, cell0, ndim1);  // B = P0 \tilde{A}
     P1T(Ae, cell_info1, cell1, ndim0); // A =  B P1_T
 
-    // Zero rows/columns for essential bcs
-    std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
-    std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
-
-    if (!bc0.empty())
+    // Do not clear BC rows/cols in "LiftingMode"
+    if constexpr (!LiftingMode)
     {
-      for (int i = 0; i < num_dofs0; ++i)
+      // Zero rows and columns for BCs
+      if (!bc0.empty())
       {
-        for (int k = 0; k < bs0; ++k)
+        for (int i = 0; i < num_dofs0; ++i)
         {
-          if (bc0[bs0 * dofs0[i] + k])
+          for (int k = 0; k < bs0; ++k)
           {
-            // Zero row bs0 * i + k
-            const int row = bs0 * i + k;
-            std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
+            if (bc0[bs0 * dofs0[i] + k])
+            {
+              // Zero row bs0 * i + k
+              const int row = bs0 * i + k;
+              std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
+            }
           }
         }
       }
-    }
-
-    if (!bc1.empty())
-    {
-      for (int j = 0; j < num_dofs1; ++j)
+      if (!bc1.empty())
       {
-        for (int k = 0; k < bs1; ++k)
+        for (int j = 0; j < num_dofs1; ++j)
         {
-          if (bc1[bs1 * dofs1[j] + k])
+          for (int k = 0; k < bs1; ++k)
           {
-            // Zero column bs1 * j + k
-            const int col = bs1 * j + k;
-            for (int row = 0; row < ndim0; ++row)
-              Ae[row * ndim1 + col] = 0;
+            if (bc1[bs1 * dofs1[j] + k])
+            {
+              // Zero col bs1 * j + k
+              const int col = bs1 * j + k;
+              for (int row = 0; row < ndim0; ++row)
+                Ae[row * ndim1 + col] = 0;
+            }
           }
         }
       }
@@ -168,6 +193,8 @@
 /// from cell used to define the entity.
 ///
 /// @tparam T Matrix/form scalar type.
+/// @tparam LiftingMode If set true, only execute mat_set on cells with BCs in
+/// column space.
 /// @param[in] mat_set Function that accumulates computed entries into a
 /// matrix.
 /// @param[in] x_dofmap Dofmap for the mesh geometry.
@@ -198,7 +225,7 @@
 /// function mesh.
 /// @param[in] perms Entity permutation integer. Empty if entity
 /// permutations are not required.
-template <dolfinx::scalar T>
+template <dolfinx::scalar T, bool LiftingMode = false>
 void assemble_entities(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
     md::mdspan<const scalar_value_t<T>,
@@ -236,6 +263,7 @@
   const int ndim0 = bs0 * num_dofs0;
   const int ndim1 = bs1 * num_dofs1;
   std::vector<T> Ae(ndim0 * ndim1);
+
   assert(entities0.size() == entities.size());
   assert(entities1.size() == entities.size());
   for (std::size_t f = 0; f < entities.extent(0); ++f)
@@ -248,6 +276,28 @@
     std::int32_t cell0 = entities0(f, 0);
     std::int32_t cell1 = entities1(f, 0);
 
+    std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
+    std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
+
+    // Check for BCs on column space
+    if constexpr (LiftingMode)
+    {
+      auto has_bc = [&]()
+      {
+        for (std::int32_t dof : dofs1)
+        {
+          for (int k = 0; k < bs1; ++k)
+          {
+            if (bc1[bs1 * dof + k])
+              return true;
+          }
+        }
+        return false;
+      };
+      if (!has_bc())
+        continue;
+    }
+
     // Get cell coordinates/geometry
     auto x_dofs = md::submdspan(x_dofmap, cell, md::full_extent);
     for (std::size_t i = 0; i < x_dofs.size(); ++i)
@@ -263,36 +313,38 @@
     P0(Ae, cell_info0, cell0, ndim1);
     P1T(Ae, cell_info1, cell1, ndim0);
 
-    // Zero rows/columns for essential bcs
-    std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
-    std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
-    if (!bc0.empty())
+    // Don't clear rows/cols in LiftingMode
+    if constexpr (!LiftingMode)
     {
-      for (int i = 0; i < num_dofs0; ++i)
+      // Zero rows and columns for BCs
+      if (!bc0.empty())
       {
-        for (int k = 0; k < bs0; ++k)
+        for (int i = 0; i < num_dofs0; ++i)
         {
-          if (bc0[bs0 * dofs0[i] + k])
+          for (int k = 0; k < bs0; ++k)
           {
-            // Zero row bs0 * i + k
-            const int row = bs0 * i + k;
-            std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
+            if (bc0[bs0 * dofs0[i] + k])
+            {
+              // Zero row bs0 * i + k
+              const int row = bs0 * i + k;
+              std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
+            }
           }
         }
       }
-    }
-    if (!bc1.empty())
-    {
-      for (int j = 0; j < num_dofs1; ++j)
+      if (!bc1.empty())
       {
-        for (int k = 0; k < bs1; ++k)
+        for (int j = 0; j < num_dofs1; ++j)
         {
-          if (bc1[bs1 * dofs1[j] + k])
+          for (int k = 0; k < bs1; ++k)
           {
-            // Zero column bs1 * j + k
-            const int col = bs1 * j + k;
-            for (int row = 0; row < ndim0; ++row)
-              Ae[row * ndim1 + col] = 0;
+            if (bc1[bs1 * dofs1[j] + k])
+            {
+              // Zero column bs1 * j + k
+              const int col = bs1 * j + k;
+              for (int row = 0; row < ndim0; ++row)
+                Ae[row * ndim1 + col] = 0;
+            }
           }
         }
       }
@@ -306,6 +358,8 @@
 /// a matrix.
 ///
 /// @tparam T Matrix/form scalar type.
+/// @tparam LiftingMode If set true, only execute mat_set on cells with BCs in
+/// column space.
 /// @param mat_set Function that accumulates computed entries into a
 /// matrix.
 /// @param[in] x_dofmap Dofmap for the mesh geometry.
@@ -338,7 +392,7 @@
 /// function mesh.
 /// @param[in] perms Facet permutation integer. Empty if facet
 /// permutations are not required.
-template <dolfinx::scalar T>
+template <dolfinx::scalar T, bool LiftingMode = false>
 void assemble_interior_facets(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
     md::mdspan<const scalar_value_t<T>,
@@ -384,7 +438,7 @@
   const int num_cols = bs1 * 2 * dmap1_size;
 
   // Temporaries for joint dofmaps
-  std::vector<T> Ae(num_rows * num_cols), be(num_rows);
+  std::vector<T> Ae(num_rows * num_cols);
   std::vector<std::int32_t> dmapjoint0(2 * dmap0_size);
   std::vector<std::int32_t> dmapjoint1(2 * dmap1_size);
   assert(facets0.size() == facets.size());
@@ -433,6 +487,26 @@
     std::ranges::copy(dmap1_cell0, dmapjoint1.begin());
     std::ranges::copy(dmap1_cell1, std::next(dmapjoint1.begin(), dmap1_size));
 
+    // Check for BCs on column space
+    if constexpr (LiftingMode)
+    {
+      auto has_bc = [&]()
+      {
+        for (std::int32_t dof : dmapjoint1)
+        {
+          for (int k = 0; k < bs1; ++k)
+          {
+            if (bc1[bs1 * dof + k])
+              return true;
+          }
+        }
+        return false;
+      };
+
+      if (!has_bc())
+        continue;
+    }
+
     // Tabulate tensor
     std::ranges::fill(Ae, 0);
     std::array perm = perms.empty()
@@ -474,33 +548,37 @@
       }
     }
 
-    // Zero rows/columns for essential bcs
-    if (!bc0.empty())
+    // Don't clear rows/cols in LiftingMode
+    if constexpr (!LiftingMode)
     {
-      for (std::size_t i = 0; i < dmapjoint0.size(); ++i)
+      // Zero rows and columns for BCs
+      if (!bc0.empty())
       {
-        for (int k = 0; k < bs0; ++k)
+        for (std::size_t i = 0; i < dmapjoint0.size(); ++i)
         {
-          if (bc0[bs0 * dmapjoint0[i] + k])
+          for (int k = 0; k < bs0; ++k)
           {
-            // Zero row bs0 * i + k
-            std::fill_n(std::next(Ae.begin(), num_cols * (bs0 * i + k)),
-                        num_cols, 0);
+            if (bc0[bs0 * dmapjoint0[i] + k])
+            {
+              // Zero row bs0 * i + k
+              int row = bs0 * i + k;
+              std::fill_n(std::next(Ae.begin(), num_cols * row), num_cols, 0);
+            }
           }
         }
       }
-    }
-    if (!bc1.empty())
-    {
-      for (std::size_t j = 0; j < dmapjoint1.size(); ++j)
+      if (!bc1.empty())
       {
-        for (int k = 0; k < bs1; ++k)
+        for (std::size_t j = 0; j < dmapjoint1.size(); ++j)
         {
-          if (bc1[bs1 * dmapjoint1[j] + k])
+          for (int k = 0; k < bs1; ++k)
           {
-            // Zero column bs1 * j + k
-            for (int m = 0; m < num_rows; ++m)
-              Ae[m * num_cols + bs1 * j + k] = 0;
+            if (bc1[bs1 * dmapjoint1[j] + k])
+            {
+              // Zero column bs1 * j + k
+              for (int m = 0; m < num_rows; ++m)
+                Ae[m * num_cols + bs1 * j + k] = 0;
+            }
           }
         }
       }
@@ -518,6 +596,8 @@
 ///
 /// @tparam T Scalar type.
 /// @tparam U Geometry type.
+/// @tparam LiftingMode. Set to true to call the kernel only on cells with BCs
+/// in bc1.
 /// @param[in] mat_set Function that accumulates computed entries into a
 /// matrix.
 /// @param[in] a Bilinear form to assemble.
@@ -528,7 +608,7 @@
 /// applied.
 /// @param bc1 Marker for columns with Dirichlet boundary conditions
 /// applied.
-template <dolfinx::scalar T, std::floating_point U>
+template <dolfinx::scalar T, std::floating_point U, bool LiftingMode = false>
 void assemble_matrix(
     la::MatSet<T> auto mat_set, const Form<T, U>& a,
     md::mdspan<const scalar_value_t<T>,
@@ -606,7 +686,7 @@
       std::span cells1 = a.domain_arg(IntegralType::cell, 1, i, cell_type_idx);
       auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
       assert(cells.size() * cstride == coeffs.size());
-      impl::assemble_cells_matrix(
+      impl::assemble_cells_matrix<T, LiftingMode>(
           mat_set, x_dofmap, x, cells, {dofs0, bs0, cells0}, P0,
           {dofs1, bs1, cells1}, P1T, bc0, bc1, fn,
           md::mdspan(coeffs.data(), cells.size(), cstride), constants,
@@ -651,7 +731,7 @@
       std::span facets0 = a.domain_arg(IntegralType::interior_facet, 0, i, 0);
       std::span facets1 = a.domain_arg(IntegralType::interior_facet, 1, i, 0);
       assert((facets.size() / 4) * 2 * cstride == coeffs.size());
-      impl::assemble_interior_facets(
+      impl::assemble_interior_facets<T, LiftingMode>(
           mat_set, x_dofmap, x,
           mdspanx22_t(facets.data(), facets.size() / 4, 2, 2),
           {*dofmap0, bs0,
@@ -696,7 +776,7 @@
         std::span e1 = a.domain_arg(itg_type, 1, i, 0);
         mdspanx2_t entities1(e1.data(), e1.size() / 2, 2);
         assert((entities.size() / 2) * cstride == coeffs.size());
-        impl::assemble_entities(
+        impl::assemble_entities<T, LiftingMode>(
             mat_set, x_dofmap, x, entities, {dofs0, bs0, entities0}, P0,
             {dofs1, bs1, entities1}, P1T, bc0, bc1, fn,
             md::mdspan(coeffs.data(), entities.extent(0), cstride), constants,