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

Question

我有压缩稀疏列(csc)表示的n下三角矩阵A，在主对角线上有零点，并且想要解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等，但是没有什么效果。

处理器: Xeon Skylake

问题，我能做些什么来优化这个函数？

Answer 1

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

在计时方面，我使用了google/基准和这个源文件。它定义了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

请注意，基线也有所改进，因为该实现中的循环也是展开的。我们的优化版本也从循环展开中得到了一些好处，但可能没有我们喜欢的那么多。查看生成的程序集(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类型的分解)。因此，如果不被迫坚持直接在稀疏结构上工作的简单方法，那么这些方法可能是值得研究的。例如，见戴维斯的这次调查。

Answer 2

对于索引类型，您可以使用unsigned而不是int来刮掉几个周期，它无论如何都必须是>= 0：

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

使用戈德波特的编译器探险家编译的内部循环代码略有不同，可能会更好地利用CPU管道。我不能测试，但你可以试试。

下面是为内部循环生成的代码：

.L8:
        mov     rax, rcx
.L5:
        mov     ecx, DWORD PTR [r10+rax*4]
        vmovsd  xmm1, QWORD PTR [r11+rax*8]
        vfnmadd231sd    xmm0, xmm1, QWORD PTR [r8+rcx*8]
        lea     rcx, [rax+1]
        vmovsd  QWORD PTR [r9], xmm0
        cmp     rdi, rax
        jne     .L8