65.9K
CodeProject 正在变化。 阅读更多。
Home

使用 TBB 容器和 OpenMP 构建有限元稀疏矩阵

Igor Gribanov
starIconstarIconstarIconstarIconstarIcon

5.00/5 (6投票s)

2021年10月6日

CPOL

9分钟阅读

viewsIcon

7832

一种用于组装压缩稀疏行 (CSR) 格式稀疏矩阵的高效算法。

引言

有限元法 (FEM) 最初是作为一种计算桁架中力的工程工具而构思的。它使工程师能够“看到”肉眼看不见的力。人们很快意识到，除了分析静态力分布之外，FEM 还可以应用于物理学的各个领域——流体力学、传热、电磁学。总的来说，FEM 是一种通过将微分方程转换为计算机“理解”的东西——浮点运算来“求解”微分方程的方法。

虽然 FEM 有各种类型，但其中许多归结为处理以下形式的线性方程组

$\begin{aligned} \mathbf{Ax}=\mathbf{b}, \end{aligned}$

其中 \(\mathbf{A}\) 是一个稀疏矩阵。

求解此类方程的库已经存在——MKL、LAPACK、cuSPARSE——举几个例子。FEM 代码准备矩阵 \(\mathbf{A}\) 和右侧 \(\mathbf{b}\)，然后让求解器发挥其魔力。但稀疏矩阵究竟是如何准备的？有没有办法优化这个过程？

在学术界工作的几年里，我编写了一种稀疏矩阵组装算法。这种数据结构与众不同。它是一种具有特定属性的非标准映射。大多数尝试使用“标准”哈希集和红黑树的方法都大错特错。

本文将讨论稀疏线性系统组装中可能出现的陷阱。它可能对开发自定义 FEM 实现的研究人员有所帮助。

一个简单示例

考虑一个质量-弹簧系统，其中每个弹簧都有其独特的刚度常数 \(k_n\) 和静止长度 \(d_n\)。连接节点 \(i\) 和 \(j\) 的弹簧 \(n\) 对节点 \(i\) 施加以下力。

$\begin{aligned} F_{n,i} = k_n(x_j-x_i-d_n).\end{aligned}$

对于作用在节点 \(i\) 上的每一个力，都有一个符号相反的力作用在入射节点 \(j\) 上

$\begin{aligned} F_{n,j}= -F_{n,i}.\end{aligned}$

我们正在寻找平衡时的节点位移——当所有弹簧的力相加为零时

$\begin{aligned} \sum_{n} {F_{n,k}} = 0.\end{aligned}$

一个线性方程组出现

$\begin{aligned} \mathbf{Ax}=\mathbf{b}. \end{aligned}$

其中 \(\mathbf{x}=[x_1,x_2,...,x_N]^T\) 是节点位置向量，\(\mathbf{A}\) 是节点相互作用矩阵，\(\mathbf{b}\) 是线性系统的右侧。后者不为零，因为初始弹簧长度 \(d_n\) 贡献常数项。

\(\mathbf{x}\) 的长度将表示为 N，矩阵 \(\mathbf{A}\) 的大小为 NxN。

对于连接节点 \(i\) 和 \(j\) 的每个弹簧，矩阵 \(\mathbf{A}\) 中都有非零条目 \(e_{ii},e_{ij},e_{ji},e_{jj}\)。除了这些元素，矩阵大部分都填充了零——它是稀疏的。

本文将讨论用于 FEM 应用的稀疏矩阵的准备和存储方法。

COO 列表 - 超级简单

大多数求解器库都接受所谓的坐标列表 (COO) 格式的稀疏矩阵，该格式将非零条目表示为 (row, column, value) 元组。Row 和 column 编码为整数，而 value 条目具有浮点类型。

要创建这样一个元组列表，需要遍历所有弹簧并将相应的条目附加到列表中。

// make a stretchy list
std::vector<std::tuple<int,int,double>> COO_list;

// append values from each finite element
for(const auto &spring : springs)
{
  COO_list.insert(A.end(),spring.Entries_A());
}

这可以与 OpenMP 和 TBB 容器并行完成

// a container that will survive simultaneous insertions
tbb::concurrent_vector<std::tuple<int,int,double>> COO_list;

// go parallel!
#pragma omp parallel for
for(unsigned i=0;i<springs.size();i++)
{
  const Spring &spring = springs[i];
  A.push_back(spring.Aii());
  A.push_back(spring.Aij());
  A.push_back(spring.Aji());
  A.push_back(spring.Ajj());
}

这种方法存在一些问题。

  1. 由于所有线程同时将条目附加到同一个容器中，因此频繁发生竞争条件。当竞争条件不频繁时，并发容器效果最佳。
  2. 一些 (row,column) 对可能会出现多次。求解器随后会将它们加在一起，但它们占用的内存超出所需。
  3. 对于求解器来说，COO 不是一种方便的格式，因此它们会花费时间将其转换为更合适的表示形式。

尽管有明显的局限性，这种组装方法可能是代码复杂性和性能之间的最佳权衡。它绝对是一个不错的选择，特别是对于不突破 RAM/CPU 限制的研究代码。

比 COO 更糟糕的有缺陷选项

在拥有超过一百万个变量的系统中，组装步骤可能会花费大量时间。与一维弹簧模型（每个弹簧创建四个条目）不同，一些 FEM 元素可能会创建更大的条目。例如，一个内聚区元素可能连接两个四边形，即影响八个节点。此外，每个节点可能有三个自由度 (xyz)，使每个子条目的大小为 24x24，总共产生 576 个双精度值。将所有条目存储在一个列表中可能会影响性能并消耗超出所需的内存。

由于矩阵的每一行至少包含一个非零条目，因此矩阵在其行中是密集的。可以尝试将其存储为对向量的向量

std::vector<std::vector<std::pair<int,double>>> crazy_vector;

一些稀疏矩阵库使用这种方法，因为它允许通过动态调整数组大小来插入新的非零条目。这种存储格式还方便了对行条目的迭代，从而简化了矩阵乘法算法。

问题是这些好处都不能加快组装过程。向量的重新分配会导致内存碎片和整体性能损失，而向量上的插入操作速度极慢。

解决问题

稀疏矩阵库提供即时调整大小功能，但对于大多数 FEM 应用程序来说，这并不需要。由于我们知道哪些元素连接到哪些节点，因此我们可以在计算它们的值之前预先推断出非零条目的位置。

这使我们想到了压缩稀疏行 (CSR) 格式，它是许多求解器中的本机内部表示。将此格式提交给求解器不会触发转换步骤。与其他表示不同，它不适合动态调整大小，但对于存储和访问而言效率很高。

在 CSR 中，非零条目按顺序存储在单个分配块中。在该块分配之前，会创建“稀疏结构”，即应已知非零条目的列表。对于质量-弹簧示例，代码将如下所示

// STRUCTURE CREATION STAGE

// will store (i,j)-indices of non-zero entries
std::vector<std::pair<int,int>> s;

// iterate over the list of springs, and "mark" which matrix entries are used
for(const auto &spring : springs)
{
  s.emplace_back(spring.i, spring.j);
  s.emplace_back(spring.j, spring.i);
}

// add diagonal entries
for(int i=0;i<N;i++) structure.emplace_back(i,i);

// sort the resulting list -
// the consecutive elements of a row will reside next to each other
std::sort(s.begin(),s.end());

// remove duplicates
s.resize(std::distance(s.begin(),std::unique(s.begin(), s.end())));

// NNZ is the number of non-zero entries
unsigned NNZ = s.size();

// prepare a block of memory for non-zero values
std::vector<double> values;
values.resize(NNZ);

// READY TO COMPUTE THE MATRIX!

结果结构向量 s 包含稀疏矩阵的每个非零元素的 (i,j) 条目。s 的大小是非零元素的数量，表示为 NNZ。

这种方法看起来与使用 COO 列表的第一次尝试非常相似。我们只是创建 COO 列表而没有实际值。该算法标记哪些 (i,j) 单元格正在使用中，删除重复项并为值准备一个内存块。

但是现在值在一个单独的列表中！

为什么两个列表比一个好？嗯，组装阶段现在可以真正并行了。当两个弹簧写入同一个 (i,j) 单元格时，会发生竞争条件。这种情况虽然并非不可能，但可能性很小。因此，组装阶段得到了改进。典型的实现将如下所示

// ASSEMBLY STAGE

// reset the matrix to zero
std::fill(values.begin(), values.end(), 0);

// additively populate the entries of the matrix from each spring
#pragma omp parallel for
for(unsigned i=0;i<springs.size();i++)
{
  const Spring &spring = springs[i];
#pragma omp atomic
  values[get_offset(spring.i, spring.i)] += spring.Aii();
#pragma omp atomic
  values[get_offset(spring.i, spring.j)] += spring.Aij();
#pragma omp atomic
  values[get_offset(spring.j, spring.i)] += spring.Aji();
#pragma omp atomic
  values[get_offset(spring.j, spring.j)] += spring.Ajj();
}

映射，映射……

我们现在需要一个映射函数 get_offset(int i, int j)，它指向每个 (i,j) 条目存储的位置。一个简单的解决方案是使用 std::map 或 std::unordered_map，但让我们考虑一下缺点。

  1. 这些数据结构在每次准备时都会动态分配内存。
  2. 插入复杂度为 O(log(NNZ))，因此对于数百万个 NNZ，将会出现明显的性能下降。
  3. 对 (i,j) 列表进行排序（或保持排序顺序）计算成本很高。

最终解决方案

让我们利用稀疏矩阵在行方面是密集的这一事实。我们知道每行至少包含一个非零条目，并且预先已知行数 N。因此，(i,j) 对可以表示为向量的向量（再次）。

// expected number of non-zero entries per row
constexpr unsigned nCon = 8;

// each row will store indices of the connected neighbors
std::vector<std::unique_ptr<tbb::concurrent_vector<unsigned>>> rows;

// preallocate upfront
// avoid copying the whole "vector of vectors" by using pointers
if(rows.size()<N)
{
  rows.reserve(N);
  while(rows.size()<N)
    rows.push_back(std::make_unique<tbb::concurrent_vector<unsigned>>(nCon*2));
}

我们回到了起点，但又不完全是。这次情况不同。我们没有动态分配整个矩阵，而只分配结构数组。事实上，我们完全避免了不必要的重新分配。

每行的最大条目数或多或少是预先知道的——它取决于特定的 FEM 算法和网格类型。尽管如此，它通常在给定应用程序中是常数。例如，在二维网格中，使用三角形元素，每个节点通常被 6-7 个三角形包围。将此上限存储在 nCon 常量中可确保行向量不会超出其初始容量。此外，整个内存结构可以在后续计算步骤中重用，并且只发生一次分配。

最后，由于使用了 TBB 并发向量，结构可以并行填充。我没有比较顺序算法和并行算法，但并行算法对于超大型网格可能更快，而顺序算法对于较小的几何形状是最佳的。对于顺序算法，可以使用Boost 的小向量 boost::container::small_vector<unsigned,nCon*2>。

让我们看看非零元素是如何标记的。

// STRUCTURE CREATION STAGE - each spring "marks" its non-zero matrix entries
#pragma omp parallel for
for(unsigned k=0;k<springs.size();k++)
{
  const Spring &spring = springs[k];
  rows[spring.i]->push_back(j);
  rows[spring.j]->push_back(i);
}

我们现在并行排序许多小数组，而不是排序一个巨大的 (i,j) 索引数组！重复的 (i,j) 条目像以前一样被删除。最后，我们计算非零元素的总数 NNZ，这将在下一步中用到。

// SORT EACH ROW SEPARATELY AND REMOVE THE DUPLICATES; COUNT NNZ

// counter for the total number of non-zero entries
unsigned NNZ = 0;

// process rows in parallel
#pragma omp parallel for reduction(+:NNZ)
for(unsigned i=0;i<N;i++)
{
  tbb::concurrent_vector<unsigned> &rn = *rows[i];

  // add the diagonal entry
  rn.push_back(i);

  // sort the column indices (j-entries)
  std::sort(rn.begin(),rn.end());

  // remove the duplicates
  rn.resize(std::distance(rn.begin(),std::unique(rn.begin(),rn.end())));

  // count the non-zero entries
  NNZ += rn.size();
}

我们现在继续创建结构数组 COL_INDEX 和 ROW_INDEX，它们是 CSR 表示的一部分并传递给求解器。它们还允许在 (i,j)->(offset) 映射中进行 O(1) 查找。COL_INDEX 的大小为 NNZ，ROW_INDEX 的大小为 N+1。我们将这两个数组都存储为 std::vector<int> 并根据需要调整大小。为了填充它们，我们遍历排序后的 (i,j) 对并顺序枚举条目。

// PREPARE CSR STRUCTURE ARRAYS
std::vector<int> COL_INDEX, ROW_INDEX;

// resize as needed
COL_INDEX.resize(NNZ);
ROW_INDEX.resize(N+1);
ROW_INDEX[N] = NNZ;

// enumerate the sorted indices to populate COL_INDEX and ROW_INDEX
unsigned count = 0;
for(unsigned row=0;row<N;row++)
{
  ROW_INDEX[row] = count;
  tbb::concurrent_vector<unsigned> &rn = *rows[row];
  for(unsigned const &local_column : rn)
  {
    COL_INDEX[count] = local_column;
    count++;
  }
}

拼图的最后一块

在组装阶段，矩阵条目以加法方式写入 values 数组。对于每个 (i,j) 条目，都有一个到内存位置的 offset。高效地推断这些偏移量很重要，因为此操作会重复多次。在早期版本中，我尝试了各种方法，从简单的 std::map<std::pair<int,int>,int> 开始，然后发展到每行映射，甚至独立的向量向量。但有没有更好的方法？

有一种优雅的方法可以使用 COL_INDEX 和 ROW_INDEX。后者已经存储了每行第一个非零元素的偏移量，ROW_INDEX[i]。我们只需要找到与 j 匹配的列索引。

我们遍历 COL_INDEX 在 &COL_INDEX[ROW_INDEX[i]] 和 &COL_INDEX[ROW_INDEX[i+1]] 之间的部分，并找到与 j 匹配的值。

// for a given (i,j)-entry, return an offset in the array of values
unsigned get_offset(int row, int column)
{
  // where the row begins
  int col_offset_begin = ROW_INDEX[row];

  // where the row ends
  int col_offset_end = ROW_INDEX[row+1];

  // subrange in COL_INDEX that stores j-indices for a given row
  int *start_pt = &COL_INDEX[col_offset_begin];
  int *end_pt = &COL_INDEX[col_offset_end];

  // run a binary search on COL_INDEX to find our "column"
  auto it = std::lower_bound(start_pt, end_pt, column);
  return std::distance(start_pt,it)+col_offset_begin;
}

我不太确定在这种情况下二分查找 std::lower_bound 还是常规查找 std::find 更快。根据特定的 FEM 设置，行可能只有很少的条目，因此直接线性查找可能会表现更好。无论哪种方式，这都是一个优雅的解决方案，具有常数时间查找。

结论

有限元法有多种多样。学生和研究人员不断发明新的修改。并非所有 FEM 代码都使用稀疏矩阵——有些可以不用它们。本文不声称提供最优算法——它是一种可行算法，并且仍有优化的空间。

在我的博士和博士后研究中，我编写了各种 FEM 代码——包括内聚区材料、壳体和板元素、2D Neo-Hookean 材料等。有些使用线性求解器，有些使用二次优化器。我的重点更多是 FEM 公式，而不是矩阵数学的实现。确实，一些竞技程序员会很快想出最优的组装过程，但对我来说，这个代码多年来慢慢演变成现在的样子。使用此算法的有效项目可以在这里找到。请不要犹豫留下评论，特别是提出改进建议。此算法可能对其他 FEM 研究人员具有实际价值，这也是本文的动机！

© . All rights reserved.