Flatten dofmap storage

cpp/dolfinx/fem/DirichletBC.cpp

@@ -11,7 +11,6 @@
 #include <array>
 #include <dolfinx/common/IndexMap.h>
 #include <dolfinx/common/sort.h>
-#include <dolfinx/graph/AdjacencyList.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
 #include <dolfinx/mesh/cell_types.h>

cpp/dolfinx/fem/DofMap.cpp

@@ -19,6 +19,7 @@
 
 using namespace dolfinx;
 using namespace dolfinx::fem;
+namespace stdex = std::experimental;
 
 namespace
 {
@@ -42,7 +43,10 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   assert(cells);
 
   // Build set of dofs that are in the new dofmap (un-blocked)
-  std::vector<std::int32_t> dofs_view = dofmap_view.list().array();
+  auto dofs_view_md = dofmap_view.list();
+  std::vector<std::int32_t> dofs_view(dofs_view_md.data_handle(),
+                                      dofs_view_md.data_handle()
+                                          + dofs_view_md.size());
   dolfinx::radix_sort(std::span(dofs_view));
   dofs_view.erase(std::unique(dofs_view.begin(), dofs_view.end()),
                   dofs_view.end());
@@ -104,10 +108,11 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   }
 
   // Map dofs to new collapsed indices for new dofmap
-  const std::vector<std::int32_t>& dof_array_view = dofmap_view.list().array();
+  auto dof_array_view = dofmap_view.list();
   std::vector<std::int32_t> dofmap;
   dofmap.reserve(dof_array_view.size());
-  std::transform(dof_array_view.begin(), dof_array_view.end(),
+  std::transform(dof_array_view.data_handle(),
+                 dof_array_view.data_handle() + dof_array_view.size(),
                  std::back_inserter(dofmap),
                  [&old_to_new](auto idx_old) { return old_to_new[idx_old]; });
 
@@ -122,28 +127,30 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   ElementDofLayout element_dof_layout = dofmap_view.element_dof_layout().copy();
 
   // Create new dofmap and return
-  return DofMap(
-      std::move(element_dof_layout), index_map, 1,
-      graph::regular_adjacency_list(std::move(dofmap), cell_dimension), 1);
+  return DofMap(std::move(element_dof_layout), index_map, 1, std::move(dofmap),
+                1);
 }
 
 } // namespace
 
 //-----------------------------------------------------------------------------
-graph::AdjacencyList<std::int32_t>
-fem::transpose_dofmap(const graph::AdjacencyList<std::int32_t>& dofmap,
-                      std::int32_t num_cells)
+graph::AdjacencyList<std::int32_t> fem::transpose_dofmap(
+    std::experimental::mdspan<const std::int32_t,
+                              std::experimental::dextents<std::size_t, 2>>
+        dofmap,
+    std::int32_t num_cells)
 {
   // Count number of cell contributions to each global index
   const std::int32_t max_index = *std::max_element(
-      dofmap.array().begin(),
-      std::next(dofmap.array().begin(), dofmap.offsets()[num_cells]));
+      dofmap.data_handle(),
+      std::next(dofmap.data_handle(), num_cells * dofmap.extent(1)));
 
   std::vector<int> num_local_contributions(max_index + 1, 0);
-  for (int c = 0; c < num_cells; ++c)
+  for (std::int32_t c = 0; c < num_cells; ++c)
   {
-    for (auto dof : dofmap.links(c))
-      num_local_contributions[dof]++;
+    auto dofs = stdex::submdspan(dofmap, c, stdex::full_extent);
+    for (std::size_t d = 0; d < dofmap.extent(1); ++d)
+      num_local_contributions[dofs[d]]++;
   }
 
   // Compute offset for each global index
@@ -154,10 +161,11 @@ fem::transpose_dofmap(const graph::AdjacencyList<std::int32_t>& dofmap,
   std::vector<std::int32_t> data(index_offsets.back());
   std::vector<int> pos = index_offsets;
   int cell_offset = 0;
-  for (int c = 0; c < num_cells; ++c)
+  for (std::int32_t c = 0; c < num_cells; ++c)
   {
-    for (auto dof : dofmap.links(c))
-      data[pos[dof]++] = cell_offset++;
+    auto dofs = stdex::submdspan(dofmap, c, stdex::full_extent);
+    for (std::size_t d = 0; d < dofmap.extent(1); ++d)
+      data[pos[dofs[d]]++] = cell_offset++;
   }
 
   // Sort the source indices for each global index
@@ -172,7 +180,6 @@ fem::transpose_dofmap(const graph::AdjacencyList<std::int32_t>& dofmap,
   return graph::AdjacencyList(std::move(data), std::move(index_offsets));
 }
 //-----------------------------------------------------------------------------
-/// Equality operator
 bool DofMap::operator==(const DofMap& map) const
 {
   return this->_index_map_bs == map._index_map_bs
@@ -190,17 +197,18 @@ DofMap DofMap::extract_sub_dofmap(std::span<const int> component) const
       = this->element_dof_layout().sub_view(component);
 
   // Build dofmap by extracting from parent
-  const int num_cells = this->_dofmap.num_nodes();
+  const std::int32_t num_cells = this->_dofmap.size() / this->_shape1;
+
   // FIXME X: how does sub_element_map_view hand block sizes?
   const std::int32_t dofs_per_cell = sub_element_map_view.size();
   std::vector<std::int32_t> dofmap(num_cells * dofs_per_cell);
   const int bs_parent = this->bs();
-  for (int c = 0; c < num_cells; ++c)
+  for (std::int32_t c = 0; c < num_cells; ++c)
   {
-    auto cell_dmap_parent = this->_dofmap.links(c);
+    auto cell_dmap_parent = this->cell_dofs(c);
     for (std::int32_t i = 0; i < dofs_per_cell; ++i)
     {
-      const std::div_t pos = std::div(sub_element_map_view[i], bs_parent);
+      std::div_t pos = std::div(sub_element_map_view[i], bs_parent);
       dofmap[c * dofs_per_cell + i]
           = bs_parent * cell_dmap_parent[pos.quot] + pos.rem;
     }
@@ -209,11 +217,9 @@ DofMap DofMap::extract_sub_dofmap(std::span<const int> component) const
   // FIXME X
 
   // Set element dof layout and cell dimension
-  ElementDofLayout sub_element_dof_layout
-      = _element_dof_layout.sub_layout(component);
-  return DofMap(
-      std::move(sub_element_dof_layout), this->index_map, this->index_map_bs(),
-      graph::regular_adjacency_list(std::move(dofmap), dofs_per_cell), 1);
+  ElementDofLayout sub_dof_layout = _element_dof_layout.sub_layout(component);
+  return DofMap(std::move(sub_dof_layout), this->index_map,
+                this->index_map_bs(), std::move(dofmap), 1);
 }
 //-----------------------------------------------------------------------------
 std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
@@ -260,7 +266,7 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
   auto cells = topology.connectivity(tdim, 0);
   assert(cells);
   const int bs = dofmap_new.bs();
-  for (int c = 0; c < cells->num_nodes(); ++c)
+  for (std::int32_t c = 0; c < cells->num_nodes(); ++c)
   {
     std::span<const std::int32_t> cell_dofs_view = this->cell_dofs(c);
     std::span<const std::int32_t> cell_dofs = dofmap_new.cell_dofs(c);
@@ -278,9 +284,13 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
   return {std::move(dofmap_new), std::move(collapsed_map)};
 }
 //-----------------------------------------------------------------------------
-const graph::AdjacencyList<std::int32_t>& DofMap::list() const
+std::experimental::mdspan<const std::int32_t,
+                          std::experimental::dextents<std::size_t, 2>>
+DofMap::list() const
 {
-  return _dofmap;
+  return std::experimental::mdspan<const std::int32_t,
+                                   std::experimental::dextents<std::size_t, 2>>(
+      _dofmap.data(), _dofmap.size() / _shape1, _shape1);
 }
 //-----------------------------------------------------------------------------
 int DofMap::index_map_bs() const { return _index_map_bs; }

cpp/dolfinx/fem/DofMap.h

@@ -10,6 +10,7 @@
 #pragma once
 
 #include "ElementDofLayout.h"
+#include <basix/mdspan.hpp>
 #include <concepts>
 #include <cstdlib>
 #include <dolfinx/common/MPI.h>
@@ -58,9 +59,11 @@ namespace dolfinx::fem
 /// typically used to exclude ghost cell contributions.
 /// @return Map from global (process-wise) index to positions in an
 /// unaassembled array. The links for each node are sorted.
-graph::AdjacencyList<std::int32_t>
-transpose_dofmap(const graph::AdjacencyList<std::int32_t>& dofmap,
-                 std::int32_t num_cells);
+graph::AdjacencyList<std::int32_t> transpose_dofmap(
+    std::experimental::mdspan<const std::int32_t,
+                              std::experimental::dextents<std::size_t, 2>>
+        dofmap,
+    std::int32_t num_cells);
 
 /// @brief Degree-of-freedom map.
 ///
@@ -85,12 +88,14 @@ class DofMap
   /// each cell.
   /// @param[in] bs The block size of the `dofmap`.
   template <std::convertible_to<fem::ElementDofLayout> E,
-            std::convertible_to<graph::AdjacencyList<std::int32_t>> U>
+            std::convertible_to<std::vector<std::int32_t>> U>
   DofMap(E&& element, std::shared_ptr<const common::IndexMap> index_map,
          int index_map_bs, U&& dofmap, int bs)
       : index_map(index_map), _index_map_bs(index_map_bs),
         _element_dof_layout(std::forward<E>(element)),
-        _dofmap(std::forward<U>(dofmap)), _bs(bs)
+        _dofmap(std::forward<U>(dofmap)), _bs(bs),
+        _shape1(_element_dof_layout.num_dofs()
+                * _element_dof_layout.block_size() / _bs)
   {
     // Do nothing
   }
@@ -115,12 +120,12 @@ class DofMap
   bool operator==(const DofMap& map) const;
 
   /// @brief Local-to-global mapping of dofs on a cell
-  /// @param[in] cell The cell index
+  /// @param[in] c The cell index
   /// @return Local-global dof map for the cell (using process-local
   /// indices)
-  std::span<const std::int32_t> cell_dofs(int cell) const
+  std::span<const std::int32_t> cell_dofs(std::int32_t c) const
   {
-    return _dofmap.links(cell);
+    return std::span<const std::int32_t>(_dofmap.data() + _shape1 * c, _shape1);
   }
 
   /// @brief Return the block size for the dofmap
@@ -133,10 +138,9 @@ class DofMap
 
   /// @brief Create a "collapsed" dofmap (collapses a sub-dofmap)
   /// @param[in] comm MPI Communicator
-  /// @param[in] topology The mesh topology that the dofmap is defined
-  /// on
-  /// @param[in] reorder_fn The graph re-ordering function to apply to
-  /// the dof data
+  /// @param[in] topology Mesh topology that the dofmap is defined on
+  /// @param[in] reorder_fn Graph re-ordering function to apply to the
+  /// dof data
   /// @return The collapsed dofmap
   std::pair<DofMap, std::vector<std::int32_t>> collapse(
       MPI_Comm comm, const mesh::Topology& topology,
@@ -147,7 +151,9 @@ class DofMap
 
   /// @brief Get dofmap data
   /// @return The adjacency list with dof indices for each cell
-  const graph::AdjacencyList<std::int32_t>& list() const;
+  std::experimental::mdspan<const std::int32_t,
+                            std::experimental::dextents<std::size_t, 2>>
+  list() const;
 
   /// Layout of dofs on an element
   const ElementDofLayout& element_dof_layout() const
@@ -170,9 +176,12 @@ class DofMap
   ElementDofLayout _element_dof_layout;
 
   // Cell-local-to-dof map (dofs for cell dofmap[cell])
-  graph::AdjacencyList<std::int32_t> _dofmap;
+  std::vector<std::int32_t> _dofmap;
 
   // Block size for the dofmap
   int _bs = -1;
+
+  // Number of columns in _dofmap
+  int _shape1 = -1;
 };
 } // namespace dolfinx::fem

cpp/dolfinx/fem/Expression.h

@@ -134,15 +134,16 @@ class Expression
   void eval(std::span<const std::int32_t> cells, std::span<T> values,
             std::array<std::size_t, 2> vshape) const
   {
+    namespace stdex = std::experimental;
+
     // Prepare coefficients and constants
     const auto [coeffs, cstride] = pack_coefficients(*this, cells);
     const std::vector<T> constant_data = pack_constants(*this);
     const auto& fn = this->get_tabulate_expression();
 
     // Prepare cell geometry
     assert(_mesh);
-    const graph::AdjacencyList<std::int32_t>& x_dofmap
-        = _mesh->geometry().dofmap();
+    auto x_dofmap = _mesh->geometry().dofmap();
 
     // Get geometry data
     auto cmaps = _mesh->geometry().cmaps();
@@ -189,7 +190,7 @@ class Expression
     for (std::size_t c = 0; c < cells.size(); ++c)
     {
       const std::int32_t cell = cells[c];
-      auto x_dofs = x_dofmap.links(cell);
+      auto x_dofs = stdex::submdspan(x_dofmap, cell, stdex::full_extent);
       for (std::size_t i = 0; i < x_dofs.size(); ++i)
       {
         std::copy_n(std::next(x_g.begin(), 3 * x_dofs[i]), 3,

cpp/dolfinx/fem/Function.h

@@ -407,8 +407,7 @@ class Function
     const CoordinateElement& cmap = mesh->geometry().cmaps()[0];
 
     // Get geometry data
-    const graph::AdjacencyList<std::int32_t>& x_dofmap
-        = mesh->geometry().dofmap();
+    auto x_dofmap = mesh->geometry().dofmap();
     const std::size_t num_dofs_g = cmap.dim();
     std::span<const U> x_g = mesh->geometry().x();
 
@@ -501,7 +500,7 @@ class Function
         continue;
 
       // Get cell geometry (coordinate dofs)
-      auto x_dofs = x_dofmap.links(cell_index);
+      auto x_dofs = stdex::submdspan(x_dofmap, cell_index, stdex::full_extent);
       assert(x_dofs.size() == num_dofs_g);
       for (std::size_t i = 0; i < num_dofs_g; ++i)
       {

cpp/dolfinx/fem/FunctionSpace.h

@@ -15,7 +15,6 @@
 #include <cstddef>
 #include <cstdint>
 #include <dolfinx/common/IndexMap.h>
-#include <dolfinx/graph/AdjacencyList.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
@@ -217,14 +216,11 @@ class FunctionSpace
 
     // Get coordinate map
     if (_mesh->geometry().cmaps().size() > 1)
-    {
       throw std::runtime_error("Mixed topology not supported");
-    }
     const CoordinateElement& cmap = _mesh->geometry().cmaps()[0];
 
     // Prepare cell geometry
-    const graph::AdjacencyList<std::int32_t>& x_dofmap
-        = _mesh->geometry().dofmap();
+    auto x_dofmap = _mesh->geometry().dofmap();
     const std::size_t num_dofs_g = cmap.dim();
     std::span<const T> x_g = _mesh->geometry().x();
 
@@ -270,7 +266,7 @@ class FunctionSpace
     for (int c = 0; c < num_cells; ++c)
     {
       // Extract cell geometry
-      auto x_dofs = x_dofmap.links(c);
+      auto x_dofs = stdex::submdspan(x_dofmap, c, stdex::full_extent);
       for (std::size_t i = 0; i < x_dofs.size(); ++i)
         for (std::size_t j = 0; j < gdim; ++j)
           coordinate_dofs(i, j) = x_g[3 * x_dofs[i] + j];