在Eigen中实现稀疏矩阵的多线程乘法

在做李导数的实验时，我曾绞尽脑汁给出了一个稍微优于Eigen默认实现的稀疏矩阵的并行乘法，前几天复现LSFM时对这个方法进一步地优化，性能有了进一步的提升，整理后记录在这里。这个方法本质上也适用于任何列压缩稀疏矩阵乘以行压缩稀疏矩阵的情形。

本文给出了一种基于Eigen的稀疏系数矩阵多线程乘法，本质上也适用于任何列压缩稀疏矩阵*行压缩稀疏矩阵的情形。在Eigen3中，目前原生支持的与稀疏矩阵乘法有关的多线程计算并不多，从前面的记录可以看出，只有

行优先稀疏矩阵 * 稠密矩阵(或稠密向量)

才能直接被Eigen多线程计算。而在做参数化实现的时候，有不少算法都涉及了稀疏矩阵的乘法，例如：

Least squares conformal maps for automatic texture atlas generation
Authalic Parameterization of General Surfaces Using Lie Advection
……

都涉及了以 $A^TA\mathbf{x}=A^T\mathbf{b}$ 来曲线求解线性方程组 $A\mathbf{x}=\mathbf{b}$ 的操作，其中A为稀疏矩阵。在矩阵规模愈发庞大时，稀疏矩阵的乘法将耗费大量的时间（而且整个程序只有这一部分是单核在跑），这让人很不愉悦。

在做李导数的实验时，我曾绞尽脑汁给出了一个稍微优于Eigen默认实现的稀疏矩阵的并行乘法，前几天复现LSFM时对这个方法进一步地优化，性能有了进一步的提升，整理后记录在这里。

Naive的方法

矩阵乘法 $C=AB$ ，无非就是A中每i行与B中每j列做向量内积，得到的结果作为Cij的值，即 $C_{ij} = A_i \cdot B_j \\$ 。对于稀疏矩阵也是如此。

一种比较Naive的思路就是用A的每一行与B的每一列之间做一个稀疏向量之间的内积，这个方法要求A是行压缩而B是列压缩。但是，稀疏矩阵的乘积C往往还是稀疏矩阵，既然是稀疏矩阵那么就意味着我们要做的大量的点积的结果其实是0，如下面的例子示范的。这就导致了这种方法效率非常低下，即便是多线程的依然比不过Eigen的单线程乘法。

$\begin{matrix} A_i=[1 0 0 0 0 1 0 0] \\ B_j=[0 0 0 1 0 0 0 0]^T \end{matrix}$

一个不错的方法

Naive方法效率低下的原因便在于它做了大量无效的点积运算，为了克服这个缺点，我们希望简单粗暴地直接定位非0元素。显然，前面稀疏行向量Ai和稀疏列向量Bj的内积非0的必要条件是，Ai中的某k列和Bj中的某k行同时非0——只有这个条件满足时，Ai和Bj的内积才可能非0，Cij处才可能有非0元素。但是，常规地去判断是否存在这样一个k却又和前面直接做点积没有什么差异。

既然如此，我们可以考虑一个和Naive方法逆向的思维：

假定A为列压缩而B为行压缩，矩阵乘法要求A的列数等于B的行数；
然后对于A的每k列（对应B的每k行，反之亦然），我们可以直接遍历得到A在第k列的非零元素Aik，以及B在第k行的非零元素Bkj；
对Aik和Bkj相乘，得到结果三元组(i, j, Aik*Bjk)，累加到Cij处。

这种方法相对于Naive方法的优点在于，每次计算基本上都是有效的，不会踯躅于大量的零元素中。它的一个特点是，每一次Aik和Bkj相乘，我们得到的只是Cij一个部分结果，只有走完k的循环后，才能累加得到完整的结果。虽然我没有翻看过Eigen自身做稀疏矩阵乘法的实现，但从效率上反推应该也是类似的方法。

并行化

前面方法在并行化上一个最大的痛点在于数据写相关上。例如，在并行化循环k时，完全可能存在这样的k1和k2，使得Aik1、Aik2、Bk1j、Bk2j均为非零元素，这样就会产生对Cij的同时写入操作。另一方面，正如在前面的文章里提到过的，稀疏矩阵在插入元素时，要么预分配空间，要么使用三元组数组进行初始化。这些都是并行化时需要解决的问题。

在这里，我采用的方法是先将结果以三元组的形式存储出来。为了解决写冲突的问题，我将三元组按列（或行，在这里行或列是可以互换的，后同）分配到n+1个桶（Bucket）中，每个桶负责一段宽度（Stride）列号的三元组存储，并且满足Stride*(n+1)>=列数，再为每个Bucket配置一个互斥锁。可以直观地看图。

考虑到在遍历k时，需要写入的行号在一个比较均匀的稀疏矩阵下是基本随机的，因此这种操作可以在尽可能少地使用锁的情况下，尽可能少地阻塞线程。这样，再计算完成后，我们可以：

直接用这个桶的迭代器，通过Eigen自带的方法从三元组生成结果稀疏矩阵C，请跳过下一节。
或者，自行预分配空间，然后自行生成结果稀疏矩阵C（当然我就是这样做的），请看下一节。

预分配空间

本文的方法虽然对A、B的行列压缩有要求，但对结果C的却没有。只是当自行预分配空间时，我们需要根据C的行或列压缩，来决定每个Bucket是负责一段行还是一段列的三元组数据。

简而言之：

当C矩阵为列压缩时，每个Bucket存储宽度为Stride的列的三元组；
当C矩阵为行压缩时，每个Bucket存储宽度为Stride的行的三元组。

因为列压缩矩阵要求通过每列的非零元素个数来预分配（行压缩矩阵相反），我们Bucket中的元素的分配方式与其对应的话，可以有效地避免预分配空间时的数据相关性。例如，对于列压缩矩阵，我们可以并行地考察第n个Bucket中的三元组，来确定[Stride * n, Stride * (n+1))这个区间中所有列的非零元素个数。当然，还要注意的一个问题就是计数的时候要判断重复。

在预分配好空间后，正如在前面的文章里提到过的，我们可以在桶间并行地将元素累加到结果C中地相应位置。

桶容量的估计

估计每个桶的容量大小是很重要的，虽然使用STL的vector可以自动resize，但是这会显著地降低向桶中插入元素地效率，尤其是插入元素还是在临界区中地操作。

假设稀疏矩阵的元素是比较平均的，我们可以得到A中每列的非零元素平均为 $nnz_A/cols_A$ ，B中每行的非零元素平均为 $nnz_B/rows_B$ ，那么总计插入到同种的三元组个数即为

$n_{Triplet} = \frac{nnz_A \cdot nnz_B} {cols_A \cdot rows_B} = \frac{nnz_A \cdot nnz_B} {cols_A^2}$

那么平均每个桶中的元素及是

$n_{TripletPerBucket} = \frac{n_{Triplet} * cols_A}{n_{Bucket}}= \frac{nnz_A \cdot nnz_B} {cols_A \cdot n_{Bucket}}$

一般而言，对于比较均匀的稀疏矩阵，可以取2倍于估计的容量值。

代码实现

SparseProduct.h

在代码里，Bucket的总量默认取8192，这个大小可以根据实际情况进行调整，但建议至少在机器线程数的1个数量级以上。

性能

这里给出2个测试案例，分别对应开头给出的两篇参数化方法涉及的稀疏矩阵乘法。测试所用的三维模型为Buddha，470K面，具体情况为：

测试案例1：LSFM参数化涉及的矩阵乘法，矩阵A为471538*941014，矩阵B为941014*471538，稀疏度均为1.27e-5。
测试案例2：Lie Advection参数化涉及的矩阵乘法，矩阵A为235771*235771，矩阵B为235771，稀疏度均为2.96e-5。

测试机器为Intel Core-i7 6700K @ 4.38GHz，16GB 单通道DDR4 2400MHz，Windows 10.0.17134.407。我们对每个方法用每个案例做2组实验，每一组进行连续4次矩阵乘，并记录下时间，每一次的时间取2组实验中的最小值记录。最后的结果如下表所示。

方法	案例	第1次	第2次	第3次	第4次
Eigen单线程	1	1141ms	1176ms	1162ms	1162ms
LCC的多线程	1	345ms	298ms	313ms	298ms
Eigen单线程	2	471ms	479ms	476ms	475ms
LCC的多线程	2	140ms	115ms	121ms	120ms

因为某种不明的机制（可能是内存分配），我们的方法在第1次计算时总会稍慢，但接下来对同样规模的矩阵进行计算就会快一些。总的来说，案例1的加速比在3.30~3.94，案例2的加速比在3.36~4.16。在4核CPU上可以认为取得了非常不错的效率。