测量 Gauss-Seidel 循环携带依赖并通过循环展开修复它

# 让高斯-赛德尔法重焕生机！——让-克里斯托夫·卢瓦佐博士 来源：https://loiseaujc.github.io/posts/blog-title/make_gauss_seidel_great_again.html --- 上次，我们探讨了**高斯-赛德尔法** (https://loiseaujc.github.io/posts/blog-title/gauss_seidel_experiments.html)。我们得出了以下结论：理论上高斯-赛德尔法收敛所需的迭代次数仅为雅可比法的一半，数值测试也证实了这一点。然而，尽管所需的迭代次数减半，高斯-赛德尔求解器的计算时间却是雅可比法的4到5倍。仔细思考，这似乎令人费解：一个在纸面上表现更优的算法，为何在实际数值实现中反而表现更差？正如我们所见，原因在于，与雅可比法不同，高斯-赛德尔的更新规则存在**循环依赖性** (https://en.wikipedia.org/wiki/Loop_dependence_analysis)，这阻碍了编译器对代码进行向量化。因此，我们今天要探讨的问题是：**这是否意味着高斯-赛德尔法已走进死胡同？我们能否在不牺牲雅可比法硬件效率的前提下，恢复其收敛优势？** 答案是肯定的！我们有两条不同的路径：一条是相当通用的方法，需要你尽可能地帮助编译器（但你需要对编译器和CPU的工作原理有所了解）；另一条则非常特定于我们当前的二维泊松测试案例，让数学为你指引方向。两条路都不会一帆风顺，但我们将在探索中学到很多。最终，我们会对两者都进行探讨，但现在，我们先沿着第一条路，一起看看什么是*循环展开*。 ## 等等，之前的问题是什么来着？ 让我们回顾一下。上次，我们在相同的二维泊松问题上比较了两种更新规则：朴素的雅可比迭代，以及它略显复杂的“表亲”——高斯-赛德尔法。 ```fortran do j = 2, n-1 do i = 2, n-1 u(i, j) = 0.25_dp*(dx2*b(i,j) + v(i+1,j) + v(i-1,j) & + v(i,j+1) + v(i,j-1)) end do end do ``` 雅可比法总是从*上一次迭代*的值`v`中读取，并写入一个新的数组`u`。相反，高斯-赛德尔法是原地更新： ```fortran do j = 2, n-1 do i = 2, n-1 u(i, j) = 0.25_dp*(dx2*b(i,j) + u(i+1,j) + u(i-1,j) & + u(i,j+1) + u(i,j-1)) end do end do ``` 这仅仅一行的代码差异，就为你换来了一个收敛更快的方法，达到给定精度所需的扫描次数恰好减半。在纸面上，高斯-赛德尔法应该完胜。然而，当我们实际运行时，高斯-赛德尔法的**墙钟时间**（wall-clock time）**是雅可比法的4到5倍**，尽管它只需要一半的迭代次数。那一行代码的差异正在让我们付出巨大的代价，而这在收敛性理论中却毫无体现。 **具体的循环依赖** 正如我们所发现的，罪魁祸首在于高斯-赛德尔更新中的`u(i-1,j)`在你读取它时**已经被覆盖**了。它不是上一次扫描的旧值，而是本次扫描刚刚计算出的邻居值。这正是高斯-赛德尔法收敛更快的原因：你正在单次扫描内传播信息，而不是等待整个迭代完成。但这也意味着，在完成`u(i-1,j)`的写入之前，迭代`i`无法开始。雅可比法则没有这样的限制：雅可比扫描中的每一次读取都来自`v`，而该扫描中没有任何操作会触及`v`。 这是**循环依赖** (https://en.wikipedia.org/wiki/Loop_dependence_analysis) 的教科书定义，也恰恰是编译器在决定是否能够向量化循环、重排指令或在乱序核心上重叠迭代之前，必须进行推理分析的内容。雅可比法：没有此类依赖，编译器可以自由发挥。高斯-赛德尔法：存在硬依赖，编译器束手束脚。 **为“束缚”量化** 故事讲得不错，但在你能测量它之前，它仍然只是个故事。所以这一次，我们没有采信编译器的“一面之词”，而是使用了**OSACA** (https://github.com/RRZE-HPC/OSACA) ——开源架构代码分析器——直接分析两个内核编译后的汇编代码。 OSACA读取标记过的汇编循环，并假设一个理想的乱序核心，为每个循环体报告两个数字：**关键路径（CP）**，即通过单次迭代的最长依赖链长度；以及**循环依赖（LCD）**，即该链中*必须*从一次迭代传递到下一次迭代的部分。LCD才是这里的关键数字。它是每迭代周期数的硬性下限，任何巧妙的调度都无法突破它。 以下是我们使用`-O3 -march=native`选项、用普通`do`循环（无`do concurrent`等花哨写法）编译的内核分析结果： **雅可比和标准高斯-赛德尔内核的`osaca`分析总结**。两个内核均使用`gfortran 15.3`编译，选项为`-O3 -mtune=native -march=native`。具体数字可能因您使用的CPU和编译器而异。 | 指标 | 雅可比（向量化，每次迭代处理4个元素） | 高斯-赛德尔（标量，每次迭代1个元素） | 主导限制 | | :---------------- | :------------------------------------ | :---------------------------------- | :------- | | **关键路径（CP）** | 24 cycles | 24 cycles | | | **循环依赖（LCD）**| 1 cycle | 12 cycles | | | **端口压力下限** | 3 cycles | 2.5 cycles | | 有两个点立即凸显出来。首先，两个内核的**关键路径*完全相同***——24个周期。这很合理：在两种情况下，它本质上是相同的浮点操作链： - 三次加法：`tmp1 = v(i+1, j) + v(i-1, j) + v(i, j+1) + v(i, j-1)` - 一次乘法：`tmp2 = b(i, j) * dx2` - 一次融合乘加：`tmp3 = 0.25_dp * tmp1 + 0.25_dp*tmp2` 外加一次加载和一次存储。算术操作没有改变；改变的只是允许重叠的部分。 其次，也是整个问题的关键：雅可比法的循环依赖仅为可忽略的1个周期（只是循环计数器递增，与实际计算无关）。而高斯-赛德尔法的依赖是**12个周期**，OSACA甚至准确地指出了来源： ``` 46 | 12.0 | vaddsd %xmm1, %xmm0, %xmm0 | [46, 47, 48] ``` 回溯追踪，这来自西侧邻居项：`%xmm1`保存着本次扫描刚写入的`u(i,j)`，它在下一次循环迭代中立刻被作为`u(i-1,j)`读取。三条连续指令——西侧邻居加法、北侧邻居加法和最终乘法——每条都有4周期延迟，加起来正好是12个周期。这不是编译器偷懒；而是编译器正确地报告：*在`u(i-1,j)`存在之前，根本无法开始计算`u(i,j)`*，因为算法本身要求如此。 另一个明显的迹象，甚至不需要OSACA，直接从汇编指令助记符就能看出。雅可比的内循环完全由打包双精度指令构成——`vaddpd`、`vmulpd`、`vfmadd213pd`——每条指令在单个256位AVX寄存器中同时处理四个网格点。高斯-赛德尔的内循环使用的是相同指令的标量形式——`vaddsd`、`vmulsd`、`vfmadd213sd`——一次处理一个网格点。仅`pd`/`sd`后缀就能一目了然地告诉你，哪个循环被编译器向量化了，哪个没有。 **这些数字能对上吗？** 对于每个内核，CPU实际的持续每迭代开销同时受制于*两个*约束：吞吐量下限和循环依赖（LCD），因此实际主导的是两者中较大的那个。对于高斯-赛德尔法，主导的是LCD（12周期，而吞吐量下限仅为2.5周期）。对于雅可比法，主导的是吞吐量下限（3周期，而可忽略的1周期LCD无关紧要）。换句话说，高斯-赛德尔法是*延迟受限*的，而雅可比法是*吞吐量受限*的。这是两种完全不同的瓶颈，因此比较它们绝不能仅仅查看同一列数据。 将每个内核的主导约束归一化到每个网格点（雅可比法的除以4，因为它每个向量迭代处理四个元素），得到*预测的*每次扫描减慢比约为 \[\dfrac{12}{3 / 4} = 16\times\] 这比我们测量的4–5倍墙钟时间要高。但请记住，墙钟时间衡量的是*求解总时间*，而高斯-赛德尔法收敛所需的扫描次数仅为雅可比法的一半。一旦扣除这个因素，测量结果所隐含的**每次扫描实际减慢比**是 \[\left(\frac{T_{GS}}{T_J}\right) \Big/ \left(\frac{N_{GS}}{N_J}\right) \approx (4\text{–}5) \times 2 \approx 8\text{–}10\times\] 这与OSACA预测值的差距在两倍以内，对于一个*静态*模型（它完全无法感知内存系统）来说，这是相当不错的一致性。OSACA识别出的机制——高斯-赛德尔法在硬性12周期墙上受延迟限制，雅可比法轻松受吞吐量限制——完全正确，即使精确数字需要稍打折扣。 在任意合理的规模下，雅可比法真正的瓶颈几乎肯定是DRAM带宽，而不是OSACA报告的3周期计算下限，这将会增加其实际的每次扫描开销，从而缩小剩余的差距。 因此，诊断结果三重确认：理论、静态分析和测量结果完全一致。高斯-赛德尔法更快的收敛性，完全是以一个串行化的、每元素12周期的依赖链为代价换来的，而编译器自身根本无法打破它。这就引出了一个显而易见的问题：**这个依赖链是算法*固有的*，还是我们编写循环的方式造成的假象？** ## 那么，什么是循环展开？ 循环展开 (https://en.wikipedia.org/wiki/Loop_unrolling) 是优化编译器工具箱中最古老的技巧之一，其思想简单得令人动容。与其每次迭代执行一次循环体，并每次支付分支和计数器更新的开销，不如每次迭代将循环体重复写两遍、四遍或八遍，让计数器一次性前进相应的步数。 举一个简单的例子，与我们的泊松问题无关： ```fortran do i = 1, n a(i) = 2.0_dp * a(i) end do ``` 以2的因子展开后变成： ```fortran do i = 1, n, 2 a(i) = 2.0_dp * a(i) a(i+1) = 2.0_dp * a(i+1) end do ``` 这里我们暂时忽略`n`为奇数时的一些额外处理。分支更少，计数器递增更少，而且关键的是，编译器在*单次循环迭代内*能看到的独立工作量加倍了。最后这一点才是全部吸引力所在：两个语句并列摆放，它们之间没有依赖，编译器可以自由地交错执行它们、将它们打包成单条向量指令，或按硬件喜欢的方式调度。展开并没有让CPU做更少的工作；它通过一次提供更多可调度的任务，使*调度器*的工作更轻松。所以，很自然，我们应该问：同样的技巧能拯救高斯-赛德尔法吗？ **一个看似诱人却无效的捷径** 让我们尝试显而易见的做法，将高斯-赛德尔内循环展开2倍，其他不变： ```fortran do i = 2, n-2, 2 u(i, j) = 0.25_dp*(dx2*b(i,j) + u(i-1,j) + u(i+1,j) & + u(i,j-1) + u(i,j+1)) u(i+1, j) = 0.25_dp*(dx2*b(i+1,j) + u(i,j) + u(i+2,j) & + u(i+1,j-1) + u(i+1,j+1)) end do ``` 仔细看第二个语句：它读取`u(i,j)`，而这正是*第一个*语句两行之上刚刚计算完成的值。我们根本没有移除循环依赖。我们只是将它的一个副本移到了*循环体内部*，现在它变成了同次迭代内的依赖，而非跨迭代的依赖。编译器仍然无法独立处理这两个语句，仍然无法将它们一起向量化，仍然必须严格按照顺序执行它们。我们支付了展开的开销（代码变大，需要处理`n`为奇数的剩余项），却一无所获。 这一点值得我们稍作停留，因为它是一个非常常见的陷阱：展开*看起来*应该有助于解决依赖问题，因为它经常被用作向量化使能器，但它本身对于打破*真正的*依赖链毫无作用。它非常善于暴露已经存在的并行性（比如我们简单的`a(i) = 2*a(i)`例子），但对于需要*创造*的并行性却完全无能为力。高斯-赛德尔法就完全属于第二类：依赖不是编译器未能注意到的调度偶然性，而是深植于更新规则的数学之中。 所以，如果单独展开是死路一条，那么缺少的是什么？答案是，在展开能够发挥作用之前，依赖本身必须被*重写掉*。这意味着需要借助一点代数技巧，而不仅仅是循环重构。 ## 好吧，那我们就为高斯-赛德尔这样做！ 为了保持代数推导的清晰易读，我们先暂时抛开完整的二维网格，转而处理一维泊松方程 \(u'' = -f\)，在均匀网格上离散化为 \[\frac{u_{i+1} - 2 u_i + u_{i-1}}{\Delta x^2} = -f_i\] 解出 \(u_i\) 得到更新规则 \[u_i = \dfrac{1}{2}\left(u_{i-1} + u_{i+1} + \Delta x^2 f_i\right)\] 因此毫不奇怪，扮演之前 \(0.25\) 角色的常数现在是 \(c = \tfrac12\)。从左到右扫描并原地更新，高斯-赛德尔读作： ```fortran do i = 2, n-1 u(i) = c*(u(i-1) + u(i+1) + dx2*b(i)) end do ``` 这里恰好有一个邻居是问题所在：`u(i-1)`已被*上一次循环迭代*覆盖，而`u(i+1)`仍然是未被触及的旧值（扫描尚未到达）。依赖关系就此完全暴露，没有任何干扰。 **将其写成递归形式** 将所有*不*依赖于扫描自身进展的部分，分组到单个每列量中 \[t_i = u_{i+1} + \Delta x^2 f_i\] 该量在块开始前就完全已知（它是旧数据，未被当前计算的任何操作触及）。更新简化为一个简洁的、一项的仿射递归： \[u_i = c\left(u_{i-1} + t_i\right), \qquad u_{i-1} \equiv \texttt{um}\] 其中`um`是送入此块的上一个已更新值的简称。这正是让依赖链可见为*仅*是 \(u\) 上的递归的形式——仅此而已。 **提前两步求解递归** 我们不先计算 \(u_i\) 然后*再*从中计算 \(u_{i+1}\)，而是对两者都进行显式求解，完全用块开始*之前*就存在的量来表示： \[u_i = c\,\left(\texttt{um} + t_i\right)\] \[u_{i+1} = c\,\left(u_i + t_{i+1}\right) = c\Big(c\,\left(\texttt{um}+t_i\right) + t_{i+1}\Big) = c^2\,\texttt{um} + c^2\,t_i + c\,t_{i+1}\] 看看发生了什么：\(u_{i+1}\) **不再引用** \(u_i\)。两个输出现在都是`um`、\(t_i\) 和 \(t_{i+1}\) 的纯函数。所有量都存在于*这次*这对更新开始之前。单纯展开无法触及的对内依赖（上一节中的“仅展开”内核），已经通过代数方法消解了。 在代码中，定义 \(c_2 \equiv c^2\)，这给出了以下内核： ```fortran do i = 2, n-2, 2 um = u(i-1) t0 = u(i+1) + dx2*b(i) t1 = u(i+2) + dx2*b(i+1) u(i) = c *(um + t0) u(i+1) = c2*um + c2*t0 + c*t1 ! 不依赖于 u(i) ！ end do ``` 两个独立的表达式，其中一个输出的计算不读取另一个的输出。这正是编译器需要的，以便并行调度它们，而不是串行执行。 **回到二维：无需新东西** 这正是一维迂回策略的价值所在。在实际的二维模板中，北、南以及源项贡献对创建`i`依赖**没有任何作用**——它们要么是高斯-赛德尔扫描已经处理过的行中的值...