Introduce DofMapSpan

FEniCS · Feb 26, 2025 · d12d72d · d12d72d
1 parent fda506a
commit d12d72d
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 54 deletions.
diff --git a/cpp/dolfinx/common/types.h b/cpp/dolfinx/common/types.h
@@ -10,6 +10,8 @@
 #include <concepts>
 #include <type_traits>
 
+#include <basix/mdspan.hpp>
+
 namespace dolfinx
 {
 /// @private This concept is used to constrain the a template type to floating
@@ -36,4 +38,11 @@ struct scalar_value<T, std::void_t<typename T::value_type>>
 /// @private Convenience typedef
 template <scalar T>
 using scalar_value_t = typename scalar_value<T>::type;
+
+/// Namespace containing the `mdspan` implementation
+namespace md = MDSPAN_IMPL_STANDARD_NAMESPACE;
+
+/// @private DofMap span data layout
+using DofMapSpan = md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>;
+
 } // namespace dolfinx
diff --git a/cpp/dolfinx/fem/assemble_matrix_impl.h b/cpp/dolfinx/fem/assemble_matrix_impl.h
@@ -12,6 +12,7 @@
 #include "traits.h"
 #include "utils.h"
 #include <algorithm>
+#include <dolfinx/common/types.h>
 #include <dolfinx/la/utils.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
@@ -24,11 +25,6 @@
 
 namespace dolfinx::fem::impl
 {
-/// @brief Typedef
-using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
-    const std::int32_t,
-    MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 2>>;
-
 /// @brief Execute kernel over cells and accumulate result in matrix.
 /// @tparam T Matrix/form scalar type.
 /// @param mat_set Function that accumulates computed entries into a
@@ -59,12 +55,11 @@ using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
 /// mesh
 template <dolfinx::scalar T>
 void assemble_cells(
-    la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
-    std::span<const scalar_value_t<T>> x,
-    std::span<const std::int32_t> cells,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
+    la::MatSet<T> auto mat_set, DofMapSpan x_dofmap,
+    std::span<const scalar_value_t<T>> x, std::span<const std::int32_t> cells,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap0,
     fem::DofTransformKernel<T> auto P0,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap1,
     fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
     std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
     std::span<const T> coeffs, int cstride, std::span<const T> constants,
@@ -193,12 +188,12 @@ void assemble_cells(
 /// permutations are not required.
 template <dolfinx::scalar T>
 void assemble_exterior_facets(
-    la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
+    la::MatSet<T> auto mat_set, DofMapSpan x_dofmap,
     std::span<const scalar_value_t<T>> x, int num_facets_per_cell,
     std::span<const std::int32_t> facets,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap0,
     fem::DofTransformKernel<T> auto P0,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap1,
     fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
     std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
     std::span<const T> coeffs, int cstride, std::span<const T> constants,
@@ -332,7 +327,7 @@ void assemble_exterior_facets(
 /// permutations are not required.
 template <dolfinx::scalar T>
 void assemble_interior_facets(
-    la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
+    la::MatSet<T> auto mat_set, DofMapSpan x_dofmap,
     std::span<const scalar_value_t<T>> x, int num_facets_per_cell,
     std::span<const std::int32_t> facets,
     std::tuple<const DofMap&, int, std::span<const std::int32_t>> dofmap0,
@@ -520,7 +515,7 @@ void assemble_matrix(
   for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
   {
     // Geometry dofmap and data
-    mdspan2_t x_dofmap = mesh->geometry().dofmap(cell_type_idx);
+    DofMapSpan x_dofmap = mesh->geometry().dofmap(cell_type_idx);
 
     // Get dofmap data
     std::shared_ptr<const fem::DofMap> dofmap0

diff --git a/cpp/dolfinx/fem/assemble_scalar_impl.h b/cpp/dolfinx/fem/assemble_scalar_impl.h
@@ -12,6 +12,7 @@
 #include "utils.h"
 #include <algorithm>
 #include <dolfinx/common/IndexMap.h>
+#include <dolfinx/common/types.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
@@ -22,7 +23,7 @@ namespace dolfinx::fem::impl
 {
 /// Assemble functional over cells
 template <dolfinx::scalar T>
-T assemble_cells(mdspan2_t x_dofmap, std::span<const scalar_value_t<T>> x,
+T assemble_cells(DofMapSpan x_dofmap, std::span<const scalar_value_t<T>> x,
                  std::span<const std::int32_t> cells, FEkernel<T> auto fn,
                  std::span<const T> constants, std::span<const T> coeffs,
                  int cstride)
@@ -58,7 +59,7 @@ T assemble_cells(mdspan2_t x_dofmap, std::span<const scalar_value_t<T>> x,
 
 /// Execute kernel over exterior facets and accumulate result
 template <dolfinx::scalar T>
-T assemble_exterior_facets(mdspan2_t x_dofmap,
+T assemble_exterior_facets(DofMapSpan x_dofmap,
                            std::span<const scalar_value_t<T>> x,
                            int num_facets_per_cell,
                            std::span<const std::int32_t> facets,
@@ -102,7 +103,7 @@ T assemble_exterior_facets(mdspan2_t x_dofmap,
 
 /// Assemble functional over interior facets
 template <dolfinx::scalar T>
-T assemble_interior_facets(mdspan2_t x_dofmap,
+T assemble_interior_facets(DofMapSpan x_dofmap,
                            std::span<const scalar_value_t<T>> x,
                            int num_facets_per_cell,
                            std::span<const std::int32_t> facets,
@@ -165,7 +166,7 @@ T assemble_interior_facets(mdspan2_t x_dofmap,
 /// Assemble functional into an scalar with provided mesh geometry.
 template <dolfinx::scalar T, std::floating_point U>
 T assemble_scalar(
-    const fem::Form<T, U>& M, mdspan2_t x_dofmap,
+    const fem::Form<T, U>& M, DofMapSpan x_dofmap,
     std::span<const scalar_value_t<T>> x, std::span<const T> constants,
     const std::map<std::pair<IntegralType, int>,
                    std::pair<std::span<const T>, int>>& coefficients)

diff --git a/cpp/dolfinx/fem/assemble_vector_impl.h b/cpp/dolfinx/fem/assemble_vector_impl.h
@@ -16,6 +16,7 @@
 #include <basix/mdspan.hpp>
 #include <cstdint>
 #include <dolfinx/common/IndexMap.h>
+#include <dolfinx/common/types.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
@@ -33,12 +34,6 @@ class DirichletBC;
 }
 namespace dolfinx::fem::impl
 {
-/// @cond
-using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
-    const std::int32_t,
-    MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 2>>;
-/// @endcond
-
 /// @brief Apply boundary condition lifting for cell integrals.
 /// @tparam T The scalar type.
 /// @tparam _bs0 The block size of the form test function dof map. If
@@ -78,12 +73,11 @@ using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
 /// @param[in] alpha Scaling to apply.
 template <dolfinx::scalar T, int _bs0 = -1, int _bs1 = -1>
 void _lift_bc_cells(
-    std::span<T> b, mdspan2_t x_dofmap,
-    std::span<const scalar_value_t<T>> x, FEkernel<T> auto kernel,
-    std::span<const std::int32_t> cells,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
+    std::span<T> b, DofMapSpan x_dofmap, std::span<const scalar_value_t<T>> x,
+    FEkernel<T> auto kernel, std::span<const std::int32_t> cells,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap0,
     fem::DofTransformKernel<T> auto P0,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap1,
     fem::DofTransformKernel<T> auto P1T, std::span<const T> constants,
     std::span<const T> coeffs, int cstride,
     std::span<const std::uint32_t> cell_info0,
@@ -265,12 +259,12 @@ void _lift_bc_cells(
 /// permutations are not required.
 template <dolfinx::scalar T, int _bs = -1>
 void _lift_bc_exterior_facets(
-    std::span<T> b, mdspan2_t x_dofmap,
-    std::span<const scalar_value_t<T>> x, int num_facets_per_cell,
-    FEkernel<T> auto kernel, std::span<const std::int32_t> facets,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
+    std::span<T> b, DofMapSpan x_dofmap, std::span<const scalar_value_t<T>> x,
+    int num_facets_per_cell, FEkernel<T> auto kernel,
+    std::span<const std::int32_t> facets,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap0,
     fem::DofTransformKernel<T> auto P0,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap1,
     fem::DofTransformKernel<T> auto P1T, std::span<const T> constants,
     std::span<const T> coeffs, int cstride,
     std::span<const std::uint32_t> cell_info0,
@@ -414,12 +408,12 @@ void _lift_bc_exterior_facets(
 /// @param[in] alpha The scaling to apply
 template <dolfinx::scalar T, int _bs = -1>
 void _lift_bc_interior_facets(
-    std::span<T> b, mdspan2_t x_dofmap,
-    std::span<const scalar_value_t<T>> x, int num_facets_per_cell,
-    FEkernel<T> auto kernel, std::span<const std::int32_t> facets,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
+    std::span<T> b, DofMapSpan x_dofmap, std::span<const scalar_value_t<T>> x,
+    int num_facets_per_cell, FEkernel<T> auto kernel,
+    std::span<const std::int32_t> facets,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap0,
     fem::DofTransformKernel<T> auto P0,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap1,
     fem::DofTransformKernel<T> auto P1T, std::span<const T> constants,
     std::span<const T> coeffs, int cstride,
     std::span<const std::uint32_t> cell_info0,
@@ -636,10 +630,9 @@ void _lift_bc_interior_facets(
 /// mesh
 template <dolfinx::scalar T, int _bs = -1>
 void assemble_cells(
-    fem::DofTransformKernel<T> auto P0, std::span<T> b, mdspan2_t x_dofmap,
-    std::span<const scalar_value_t<T>> x,
-    std::span<const std::int32_t> cells,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap,
+    fem::DofTransformKernel<T> auto P0, std::span<T> b, DofMapSpan x_dofmap,
+    std::span<const scalar_value_t<T>> x, std::span<const std::int32_t> cells,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap,
     FEkernel<T> auto kernel, std::span<const T> constants,
     std::span<const T> coeffs, int cstride,
     std::span<const std::uint32_t> cell_info0)
@@ -722,10 +715,10 @@ void assemble_cells(
 /// permutations are not required.
 template <dolfinx::scalar T, int _bs = -1>
 void assemble_exterior_facets(
-    fem::DofTransformKernel<T> auto P0, std::span<T> b, mdspan2_t x_dofmap,
+    fem::DofTransformKernel<T> auto P0, std::span<T> b, DofMapSpan x_dofmap,
     std::span<const scalar_value_t<T>> x, int num_facets_per_cell,
     std::span<const std::int32_t> facets,
-    std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap,
+    std::tuple<DofMapSpan, int, std::span<const std::int32_t>> dofmap,
     FEkernel<T> auto fn, std::span<const T> constants,
     std::span<const T> coeffs, int cstride,
     std::span<const std::uint32_t> cell_info0,
@@ -818,7 +811,7 @@ void assemble_exterior_facets(
 /// permutations are not required.
 template <dolfinx::scalar T, int _bs = -1>
 void assemble_interior_facets(
-    fem::DofTransformKernel<T> auto P0, std::span<T> b, mdspan2_t x_dofmap,
+    fem::DofTransformKernel<T> auto P0, std::span<T> b, DofMapSpan x_dofmap,
     std::span<const scalar_value_t<T>> x, int num_facets_per_cell,
     std::span<const std::int32_t> facets,
     std::tuple<const DofMap&, int, std::span<const std::int32_t>> dofmap,
@@ -930,9 +923,8 @@ void assemble_interior_facets(
 /// solution' in a Newton method
 /// @param[in] alpha Scaling to apply
 template <dolfinx::scalar T, std::floating_point U>
-void lift_bc(std::span<T> b, const Form<T, U>& a, mdspan2_t x_dofmap,
-             std::span<const scalar_value_t<T>> x,
-             std::span<const T> constants,
+void lift_bc(std::span<T> b, const Form<T, U>& a, DofMapSpan x_dofmap,
+             std::span<const scalar_value_t<T>> x, std::span<const T> constants,
              const std::map<std::pair<IntegralType, int>,
                             std::pair<std::span<const T>, int>>& coefficients,
              std::span<const T> bc_values1,
@@ -1105,7 +1097,7 @@ void apply_lifting(
       std::shared_ptr<const mesh::Mesh<U>> mesh = a[j]->get().mesh();
       if (!mesh)
         throw std::runtime_error("Unable to extract a mesh.");
-      mdspan2_t x_dofmap = mesh->geometry().dofmap();
+      DofMapSpan x_dofmap = mesh->geometry().dofmap();
       auto x = mesh->geometry().x();
 
       assert(a[j]->get().function_spaces().at(0));
@@ -1146,8 +1138,8 @@ void apply_lifting(
 /// @param[in] coefficients Packed coefficients that appear in `L`.
 template <dolfinx::scalar T, std::floating_point U>
 void assemble_vector(
-    std::span<T> b, const Form<T, U>& L,
-    std::span<const scalar_value_t<T>> x, std::span<const T> constants,
+    std::span<T> b, const Form<T, U>& L, std::span<const scalar_value_t<T>> x,
+    std::span<const T> constants,
     const std::map<std::pair<IntegralType, int>,
                    std::pair<std::span<const T>, int>>& coefficients)
 {
@@ -1163,7 +1155,7 @@ void assemble_vector(
   for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
   {
     // Geometry dofmap and data
-    mdspan2_t x_dofmap = mesh->geometry().dofmap(cell_type_idx);
+    DofMapSpan x_dofmap = mesh->geometry().dofmap(cell_type_idx);
 
     // Get dofmap data
     assert(L.function_spaces().at(0));