CS267 · 第 18 课

结构化网格

泊松方程是科学计算的"hello world"。本课系统比较求解结构化网格上线性方程组的经典算法——从 Jacobi、Gauss-Seidel、SOR、共轭梯度到多重网格——理解为何 Multigrid 能突破 O(N²) 的壁垒，达到 O(N) 的最优复杂度。

授课：James Demmel

日期：2025-03-20

预计用时：75 分钟

难度：困难

视频：课堂录像

课程大纲

1. 泊松方程回顾

1D 泊松方程

一维泊松方程写作 ∂²u/∂x² = f(x)，在均匀网格上用中心差分离散后，得到一个三对角线性系统 Tu = f。矩阵 T 的每一行只有三个非零元：对角线为 2，次对角线和超对角线为 -1。

离散化后：
  -u(i-1) + 2*u(i) - u(i+1) = h² · f(i)

矩阵 T (n×n):
    2  -1
   -1   2  -1
       -1   2  -1
            ... ...
                 -1   2  -1
                      -1   2

对应的 stencil：
         [ -1   2   -1 ]

Stencil 是结构化网格的核心抽象

所谓 stencil，就是描述每个网格点如何依赖其邻居的模板。1D 泊松方程的 stencil 是 [-1, 2, -1]，2D 是十字形 5 点格式，3D 是 7 点格式。Stencil 计算天然适合并行：每个点的更新相互独立（Jacobi），或近邻间有依赖（Gauss-Seidel）。

2D 泊松方程

二维泊松方程 ∂²u/∂x² + ∂²u/∂y² = f(x,y) 在 n×n 网格上离散化后，矩阵 T 变为块三对角形式。每个网格点 (i,j) 依赖于上下左右四个邻居，对角线元素为 4，邻居位置为 -1。对应的 5 点 stencil 为：

2D 5 点 stencil：
              [  -1  ]
         [ -1   4  -1 ]
              [  -1  ]

矩阵 T (N = n² 个未知量)：
    4  -1          -1
   -1   4  -1          -1
       -1   4  -1          -1
   -1          -1   4  -1          -1
       -1          -1   4  -1          -1
            ...     ...     ...

3D 情况完全类似，使用 7 点 stencil：中心为 6，六个面邻居各为 -1。未知量总数 N = n³。

2. 算法比较总览

James Demmel 在《Applied Numerical Linear Algebra》(1997) 中系统总结了求解 2D/3D 泊松方程的各类算法复杂度。N = n² (2D) 或 N = n³ (3D) 为未知量总数。

算法 串行复杂度 PRAM 复杂度 内存 处理器数
Dense LUN³NN²N²
Band LUN² (N⁷ᐟ³)NN³ᐟ² (N⁵ᐟ³)N (N⁴ᐟ³)
JacobiN² (N⁵ᐟ³)N (N²ᐟ³)NN
Explicit InverseN²log NN²N²
Conjugate GradientN³ᐟ² (N⁴ᐟ³)N¹ᐟ² (N¹ᐟ³)·log NNN
Red/Black SORN³ᐟ² (N⁴ᐟ³)N¹ᐟ² (N¹ᐟ³)NN
Sparse LUN³ᐟ² (N²)N¹ᐟ²N·log N (N⁴ᐟ³)N
FFTN·log Nlog NNN
MultigridNlog² NNN
Lower boundNlog N——

PRAM 是理想化模型

PRAM（Parallel Random Access Machine）假设通信零成本、处理器间无竞争。实际机器的通信延迟和带宽使得 PRAM 复杂度只是一个下界。但 Multigrid 即使在考虑通信后，仍是最优选择——它的串行复杂度 O(N) 已达下界。

上表中括号内为 3D 结果。可以看出，Multigrid 在串行复杂度上达到理论下界 O(N)，PRAM 复杂度 log² N 也接近下界 log N。FFT 串行复杂度为 O(N log N)，稍逊于 Multigrid，但 FFT 要求规则网格和均匀步长。

3. Jacobi 方法

公式推导

将泊松方程 Tu = b 改写，使 u(i,j) 单独出现在等式左侧：

u(i,j) = ( u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1) + b(i,j) ) / 4

设 u(i,j,m) 为第 m 步迭代后 u(i,j) 的近似值，Jacobi 迭代公式为：

u(i,j,m+1) = ( u(i-1,j,m) + u(i+1,j,m) + u(i,j-1,m) + u(i,j+1,m) + b(i,j) ) / 4

即 u(i,j,m+1) 是其四个邻居和右端项 b(i,j) 的加权平均。动机：u(i,j,m+1) 的选取使得方程在 (i,j) 点被精确满足。

收敛性分析

Jacobi 方法是最近邻计算：每步迭代中，信息只能传播一个网格点。要在 n×n 网格上传播信息到对角，至少需要 O(n) = O(N¹ᐟ²) 步。因此：

收敛速度 = 信息传播速度

Jacobi 的慢收敛本质是信息瓶颈：每步只能与直接邻居交换信息。这是所有基于稀疏矩阵-向量乘（SpMV）的迭代法的共同局限——信息传播速度被限制在每步一个网格点。Multigrid 的突破正在于打破这一限制。

并行化

Jacobi 方法可归结为稀疏矩阵-向量乘法：

U(m+1) = (T/4 - I) · U(m) + B/4

每个 U(m+1) 的分量可独立更新，只需维护两份副本（第 m 步和第 m+1 步）。在分布式内存系统中，每个处理器拥有 n²/p 个待更新元素，每步只需与邻居通信边界上的 n/p¹ᐟ² 个值。当 n >> p 时，通信量相对较小，计算与通信比良好。

4. 通信避免的 Jacobi (Communication-Avoiding Jacobi)

核心思想

理想情况：我们希望以 1 次迭代的通信代价，完成 k >> 1 次迭代的计算。关键洞察：在 Jacobi 中我们不需要存储矩阵 A 本身（stencil 是隐式的），也不一定要保留所有中间结果，只需要最终输出 Aᵏx。

Matrix Powers Kernel [Ax, A²x, ..., Aᵏx]

用计算 [Ax, A²x, ..., Aᵏx] 替代 k 次独立的 y = Ax 迭代。以三对角矩阵 A、n=32、k=3 为例：

串行算法（假设 1/4 数据可放入缓存）：
  将数据分块，逐块流水线计算 Ax, A²x, A³x
  只需在块边界移动 O(n) 字，而非 O(kn) 字

并行算法（允许少量冗余计算）：
  Proc 1  |  Proc 2  |  Proc 3  |  Proc 4
  --------+----------+----------+--------
  每个处理器多算一些重叠区域（三角形区域）
  只需发送 O(1) 条消息，而非 O(k) 条

冗余计算换取通信减少

通信避免 Jacobi 的核心权衡是用少量冗余计算换取通信次数从 O(k) 降到 O(1)。这在通信代价远高于计算代价的现代并行机上尤其有效。同样的思想（Matrix Powers Kernel）也适用于 CG 等 Krylov 子空间方法。

对于 2D Laplacian，计算 [x, Ax, A²x, A³x] 时每个网格点依赖的区域随幂次扩大。并行实现中，处理器在重叠区域（图中三角形）进行冗余计算，但只需一次消息交换。

数值稳定性

Matrix Powers Kernel 虽然减少了通信，但多次矩阵幂可能放大舍入误差。实际应用中需要权衡通信节省与数值稳定性。对于条件数适中的泊松问题，这一方法效果良好。

5. Gauss-Seidel 与红黑 SOR

Gauss-Seidel：逐点更新

Jacobi 的改进方向之一是"使用最新可用的值"。按从左到右、从上到下的顺序更新时，u(i-1,j) 和 u(i,j-1) 已经是第 m+1 步的值：

for i = 1 to n:
  for j = 1 to n:
    u(i,j,m+1) = ( u(i-1,j,m+1) + u(i+1,j,m)
                 + u(i,j-1,m+1) + u(i,j+1,m) + b(i,j) ) / 4

但 Gauss-Seidel 无法并行化——更新 u(i,j) 依赖于同一次迭代中已更新的邻居，形成数据依赖链。

红黑排序 (Red-Black Ordering)

为了恢复并行性，采用红黑排序：将网格点按棋盘模式染成红色和黑色。更新策略为：

forall 黑点 u(i,j):
    u(i,j,m+1) = 邻居红点的平均 + b(i,j)/4   ← 所有红点都是第 m 步的值

forall 红点 u(i,j):
    u(i,j,m+1) = 邻居黑点的平均 + b(i,j)/4   ← 所有黑点已更新为第 m+1 步

图着色的视角

红黑排序本质上是图着色：泊松方程的依赖图 T 是二分图，因此 2-可着色。对一般图，可使用重复的最大独立集（Maximal Independent Sets）算法进行着色。同色节点之间无依赖，可同时更新。

逐次超松弛 (SOR)

将基本迭代步改写为"当前值 + 修正量"的形式：

u(i,j,m+1) = u(i,j,m) + correction(i,j,m)

如果修正方向是好的，应当更激进地沿该方向移动（w > 1）：

u(i,j,m+1) = u(i,j,m) + w · correction(i,j,m)

这称为逐次超松弛 (SOR, Successive OverRelaxation)。对于泊松方程，最优松弛因子为：

w_opt = 2 / (1 + sin(π/(n+1)))

红黑 SOR 的收敛速度是 Jacobi 的两倍（步数从 O(n²) 降到 O(n)），但由于每步有两轮（红+黑），并行复杂度仍为 O(n)。串行复杂度为 O(n³) = O(N³ᐟ²)（2D），比 Jacobi 的 O(N²) 有明显改善。

6. 共轭梯度法 (Conjugate Gradient)

算法步骤

共轭梯度法 (CG) 是求解对称正定线性方程组 Ax = b 的 Krylov 子空间方法。对于泊松方程（A = T），CG 比 SOR 更通用（不依赖最优松弛因子的先验知识）：

# Conjugate Gradient Algorithm for Ax = b
x = initial_guess
r = b - A @ x          # 残差
rho = r @ r
p = r                  # 搜索方向
j = 1

while rho > tolerance:
    q = A @ p                       # 稀疏矩阵-向量乘（stencil）
    alpha = rho / (p @ q)
    x = x + alpha * p               # 更新解
    r = r - alpha * q               # 更新残差
    old_rho = rho
    rho = r @ r
    beta = rho / old_rho
    p = r + beta * p                # 更新搜索方向
    j += 1

收敛性

CG 比 SOR 更通用

SOR 需要知道最优松弛因子 w_opt（仅对特定问题如泊松方程可推导）。CG 不需要任何参数，适用于一般对称正定矩阵。CG 也可通过 Matrix Powers Kernel 实现通信避免版本（见 Carson 2015、Hoemmen 2010 论文）。

7. 多重网格 (Multigrid)

核心动机

回顾 Jacobi、SOR、CG 的共同局限：它们都基于 SpMV，信息每步只能传播一个网格点。要在 n×n 网格上传播信息至少需要 O(n) 步。Multigrid 的核心突破：粗网格上"远处问题看起来更简单"，信息可以更快传播。

统一主题：远处问题看起来更简单

这一思想贯穿多个领域：

多重网格算子

对于问题序列 P(m), P(m-1), ..., P(1)（从最细到最粗），定义三个核心算子：

算子符号功能
限制算子 (Restriction)R(i)将细网格问题映射到粗网格：b(i-1) = R(i)(b(i))
插值算子 (Interpolation)In(i-1)将粗网格解映射到细网格：x(i) = In(i-1)(x(i-1))
光滑算子 (Smoother)S(i)在单层网格上用加权 Jacobi/SOR 消除高频误差

V-Cycle 算法

# Multigrid V-Cycle: 求解 T(i)*x(i) = b(i)
def MGV(b, i, x):
    if i == 1:
        return exact_solution(b[1])      # 最粗网格直接求解
    else:
        x = S(b[i], x)                    # 前光滑：消除高频误差
        r = T[i] @ x - b[i]              # 计算残差
        d = In(i-1)(MGV(R(i)(r), i-1, 0))  # 递归求解粗网格修正
        x = x - d                         # 修正细网格解
        x = S(b[i], x)                    # 后光滑
        return x

之所以称为"V-Cycle"，是因为调用图呈 V 形：从最细网格逐层下降到最粗，再逐层上升回最细。

复杂度分析

串行复杂度：第 i 层有 O(4ⁱ) 个未知量（2D）。V-Cycle 总代价：

Total = Σ(i=1 to m) O(4ⁱ) = O(4ᵐ) = O(# unknowns) = O(N)

PRAM 复杂度：每层常数时间（处理器数与未知量数匹配），共 m 层，总时间为 O(m) = O(log N)。

收敛定理

Multigrid 的最关键性质

定理：一次 Multigrid V-Cycle 后，误差满足 error_after ≤ 0.5 · error_before，与网格大小无关！

推论：可以在与网格大小无关的固定步数内将误差降到任意容差。这是 Multigrid 区别于所有其他方法的核心——Jacobi/SOR/CG 的收敛速度随网格变细而下降，Multigrid 不会。

算子细节

加权 Jacobi 光滑算子 S(i)：在 1D 中，加权 Jacobi 使用 1/3 而非 1/2 的权重：

x(j) := (x(j-1) + x(j) + x(j+1) + b(j)) / 3

加权 Jacobi 的选择是为了有效消除高频误差分量。经过几步加权 Jacobi 后，误差变得"光滑"（高频分量被阻尼），此时粗网格修正能高效处理低频分量。

限制算子 R(i)：1D 情况下对邻居加权平均：

x_coarse(i) = 1/4 · x_fine(i-1) + 1/2 · x_fine(i) + 1/4 · x_fine(i+1)

2D 情况下使用全部 8 个邻居（N, S, E, W, NE, NW, SE, SW）加权平均。

插值算子 In(i-1)：1D 情况下线性插值：

若 i 也是粗网格点：x_fine(i) = x_coarse(i)
否则：            x_fine(i) = 1/2 · x_coarse(left) + 1/2 · x_coarse(right)

2D 情况下取 4 个最近邻（NW, SW, NE, SE）的加权平均。

Multigrid 作为分治算法

Multigrid 的每一次 V-Cycle 在不同频率域上降低误差：细网格负责高频误差，粗网格负责低频误差。这种频率分解的视角使 Multigrid 成为一种分治算法——将全频谱误差分解到不同尺度上分别处理。

8. 并行多重网格

2D 实现

在 n = 2ᵐ⁺¹ × 2ᵐ⁺¹ 网格上使用 s × s 处理器网格（p = s²）：

性能模型

设 p = 4ᵏ 个处理器排成 2ᵏ × 2ᵏ 网格。对 V-Cycle 从第 m 层开始：

层级 j计算 (Flops)消息数 (latency)通信量 (bandwidth)
j ≥ kO(4ʲ⁻ᵏ)O(1)·αO(2ʲ⁻ᵏ)·β
j < kO(1)O(1)·αO(1)·β

方法比较

方法# Flops# Messages# Words sent
MultigridN/p + log p · log N(log N)²(N/p)¹ᐟ² + log p · log N
FFTN·log N / pp¹ᐟ²N/p
SORN³ᐟ² / pN¹ᐟ²N/p

SOR 在所有指标上都更差。MG 的总通信量（bandwidth）更优，但消息数（latency）取决于 α 和 β 的相对大小。当 α >> β（延迟主导）时，FFT 可能更优。

粗网格的并行效率

实际实现中不会粗到 P(1)。在并行机上，粗网格通信代价不可忽视：

现代框架如 Chombo 和 AMREx（LBL）仍需处理粗网格通信问题。

9. 非结构化网格上的多重网格

挑战

实际工程问题（如骨骼力学分析、流体力学）常使用非结构化有限元网格。在非结构化网格上应用 Multigrid 面临根本性挑战：

应用实例：骨骼力学分析

M. Adams 等（2004 Gordon Bell Prize）用 Multigrid 进行非线性弹性骨骼分析：

10. 参考实现与工具

Stencil 优化工具

工具来源功能
HalideRagan-Kelley et al.图像处理 DSL，支持 stencil 优化
PochoirLeiserson @ MITStencil 编译器，自动分块和并行化
SEJITSFox et al. @ UCBPython 领域专用 JIT 并行化
BebopDatta @ UCB面向多核的 stencil 自动调优

多重网格框架

框架来源功能
ChomboLBLAMR（自适应网格加密）+ Multigrid
AMRExLBL下一代 AMR 框架（GPU 支持）
BoxLibLBL块结构 AMR 基础库

测验

Q1：Jacobi 方法在 2D n×n 网格上的串行复杂度为何是 O(N²)？

Q2：Communication-Avoiding Jacobi 的核心思想是什么？

Q3：为什么 Multigrid 的收敛速度与网格大小无关？

Q4：红黑 SOR 中，红点和黑点的更新顺序是？

Q5：在非结构化网格上实现 Multigrid 的主要挑战不包括？

要点总结

信息传播速度决定收敛速度

Jacobi/SOR/CG 的瓶颈在于信息每步只能传播一个网格点。Multigrid 通过粗网格打破这一限制，实现 O(N) 最优复杂度。

频率分解是 Multigrid 的核心

细网格消除高频误差，粗网格处理低频误差。这种分治策略使每轮 V-Cycle 误差减半，与网格大小无关。

通信避免是并行算法的关键

Matrix Powers Kernel 用冗余计算换取通信减少，适用于 Jacobi、CG 等迭代法。在现代通信瓶颈显著的并行机上效果显著。

非结构化网格仍需突破

实际工程问题需要非结构化网格上的 Multigrid，但粗化、算子定义、并行效率仍是活跃的研究领域。

延伸阅读

← 第 17 课：随机化算法 第 19 课：稀疏矩阵-向量乘法 →