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

Question

我有 n x n 下三角矩阵 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

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

处理器:至强 Skylake

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

Answer 1

这应该在很大程度上取决于矩阵的精确稀疏模式和所使用的平台。我用 gcc 8.3.0 测试了一些东西和编译器标志 -O3 -march=native (在我的 CPU 上是 -march=skylake)在 this matrix 的下三角形上维度为 3006，具有 19554 个非零条目。希望这有点接近您的设置，但无论如何我希望这些可以让您知道从哪里开始。

我使用的时间安排 google/benchmark与 this source file .它定义了 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 )，它看起来像 gcc 8 次展开可能走得有点远。对于我的设置，我实际上只需一个简单的手动展开就可以做得更好。所以放下旗帜 -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 型分解)。因此，如果您没有被迫坚持直接适用于稀疏结构的简单方法，那么研究这些方法可能是值得的。参见例如 this survey通过戴维斯。