优化稀疏下三角线性系统的反向求解

Question

我有 nxn 下三角矩阵 A 的压缩稀疏列 (csc) 表示，主对角线上有零，并且想求解 b

(A + I)' * x = b

这是我计算这个的例程：

void backsolve(const int*__restrict__ Lp,
               const int*__restrict__ Li,
               const double*__restrict__ Lx,
               const int n,
               double*__restrict__ x) {
  for (int i=n-1; i>=0; --i) {
      for (int j=Lp[i]; j<Lp[i+1]; ++j) {
          x[i] -= Lx[j] * x[Li[j]];
      }
  }
}

因此，b通过参数传入x，并被解决方案覆盖。Lp, Li,Lx分别是稀疏矩阵标准 csc 表示中的行、索引和数据指针。这个函数是程序中的顶级热点，用行

x[i] -= Lx[j] * x[Li[j]];

花费的大部分时间。编译gcc-8.3 -O3 -mfma -mavx -mavx512f给出

backsolve(int const*, int const*, double const*, int, double*):
        lea     eax, [rcx-1]
        movsx   r11, eax
        lea     r9, [r8+r11*8]
        test    eax, eax
        js      .L9
.L5:
        movsx   rax, DWORD PTR [rdi+r11*4]
        mov     r10d, DWORD PTR [rdi+4+r11*4]
        cmp     eax, r10d
        jge     .L6
        vmovsd  xmm0, QWORD PTR [r9]
.L7:
        movsx   rcx, DWORD PTR [rsi+rax*4]
        vmovsd  xmm1, QWORD PTR [rdx+rax*8]
        add     rax, 1
        vfnmadd231sd    xmm0, xmm1, QWORD PTR [r8+rcx*8]
        vmovsd  QWORD PTR [r9], xmm0
        cmp     r10d, eax
        jg      .L7
.L6:
        sub     r11, 1
        sub     r9, 8
        test    r11d, r11d
        jns     .L5
        ret
.L9:
        ret

根据 vtune 的说法，

vmovsd  QWORD PTR [r9], xmm0

是最慢的部分。我几乎没有装配经验，不知道如何进一步诊断或优化此操作。我尝试使用不同的标志进行编译以启用/禁用 SSE、FMA 等，但没有任何效果。

处理器：至强 Skylake

问题我可以做些什么来优化这个功能？

Answer 1

这在很大程度上取决于矩阵的确切稀疏模式和所使用的平台。我在这个维度为 3006、有 19554 个非零条目的矩阵的下三角上使用gcc 8.3.0编译器标志-O3 -march=native（位于我的 CPU 上）测试了一些内容。希望这与您的设置有些接近，但无论如何我希望这些可以让您知道从哪里开始。-march=skylake

为了计时，我使用了google/benchmark和这个源文件。它定义了benchBacksolveBaseline问题中给出的实现的基准以及benchBacksolveOptimized建议的“优化”实现的基准。还有benchFillRhs一个单独的基准测试，用于为右侧生成一些不完全微不足道的值。要获得“纯”反求解的时间，benchFillRhs应减去所需的时间。

1. 严格向后迭代

实现中的外部循环向后迭代列，而内部循环向前迭代当前列。似乎向后迭代每一列也会更加一致：

for (int i=n-1; i>=0; --i) {
    for (int j=Lp[i+1]-1; j>=Lp[i]; --j) {
        x[i] -= Lx[j] * x[Li[j]];
    }
}

这几乎没有改变程序集（https://godbolt.org/z/CBZAT5），但基准计时显示出明显的改进：

------------------------------------------------------------------
Benchmark                        Time             CPU   Iterations
------------------------------------------------------------------
benchFillRhs                  2737 ns         2734 ns      5120000
benchBacksolveBaseline       17412 ns        17421 ns       829630
benchBacksolveOptimized      16046 ns        16040 ns       853333

我认为这是由某种更可预测的缓存访问引起的，但我没有进一步研究它。

2. 内循环中更少的加载/存储

由于 A 是下三角形，所以我们有i < Li[j]。因此我们知道不会因为内循环中 的x[Li[j]]变化而改变。x[i]我们可以通过使用临时变量将这些知识放入我们的实现中：

for (int i=n-1; i>=0; --i) {
    double xi_temp = x[i];
    for (int j=Lp[i+1]-1; j>=Lp[i]; --j) {
        xi_temp -= Lx[j] * x[Li[j]];
    }
    x[i] = xi_temp;
}

这使得将gcc 8.3.0存储从内部循环内部移动到内存中，直接在其结束之后（https://godbolt.org/z/vM4gPD）。我的系统上测试矩阵的基准显示了一个小改进：

------------------------------------------------------------------
Benchmark                        Time             CPU   Iterations
------------------------------------------------------------------
benchFillRhs                  2737 ns         2740 ns      5120000
benchBacksolveBaseline       17410 ns        17418 ns       814545
benchBacksolveOptimized      15155 ns        15147 ns       887129

3.展开循环

虽然clang在第一次建议的代码更改后已经开始展开循环，但gcc 8.3.0仍然没有。因此，让我们通过另外传递来尝试一下-funroll-loops。

------------------------------------------------------------------
Benchmark                        Time             CPU   Iterations
------------------------------------------------------------------
benchFillRhs                  2733 ns         2734 ns      5120000
benchBacksolveBaseline       15079 ns        15081 ns       953191
benchBacksolveOptimized      14392 ns        14385 ns       963441

请注意，基线也会得到改善，因为该实现中的循环也已展开。我们的优化版本也从循环展开中受益匪浅，但可能没有我们希望的那么多。查看生成的程序集 ( https://godbolt.org/z/_LJC5f )，似乎gcc8 次展开可能有点过头了。对于我的设置，实际上我只需一个简单的手动展开就可以做得更好一些。因此，再次放下标志-funroll-loops并使用如下所示的方式实现展开：

for (int i=n-1; i>=0; --i) {
    const int col_begin = Lp[i];
    const int col_end = Lp[i+1];
    const bool is_col_nnz_odd = (col_end - col_begin) & 1;
    double xi_temp = x[i];
    int j = col_end - 1;
    if (is_col_nnz_odd) {
        xi_temp -= Lx[j] * x[Li[j]];
        --j;
    }
    for (; j >= col_begin; j -= 2) {
        xi_temp -= Lx[j - 0] * x[Li[j - 0]] +
                   Lx[j - 1] * x[Li[j - 1]];
    }
    x[i] = xi_temp;
}

据此我测量：

------------------------------------------------------------------
Benchmark                        Time             CPU   Iterations
------------------------------------------------------------------
benchFillRhs                  2728 ns         2729 ns      5090909
benchBacksolveBaseline       17451 ns        17449 ns       822018
benchBacksolveOptimized      13440 ns        13443 ns      1018182

其他算法

所有这些版本仍然使用稀疏矩阵结构上向后求解的相同简单实现。本质上，对此类稀疏矩阵结构进行操作可能会在内存流量方面产生严重问题。至少对于矩阵分解，有更复杂的方法，可以对由稀疏结构组装而成的密集子矩阵进行操作。例如超节点方法和多额方法。我对此有点模糊，但我认为此类方法也会将这种想法应用于布局并使用密集矩阵运算进行下三角向后求解（例如 Cholesky 型分解）。因此，如果您不被迫坚持直接适用于稀疏结构的简单方法，那么研究这些方法可能是值得的。例如，参见戴维斯的这项调查。