CS267 · 第 15 课

并行矩阵乘法

从 BLAS 层次结构到通信下界，再到现代并行矩阵乘法算法（SUMMA、2.5D、Strassen），理解如何在最小化通信的前提下实现高效并行稠密矩阵乘法。

历史背景与并行 motif

矩阵乘法是科学计算中最基础的 kernel 之一。从 1970 年代至今，稠密线性代数（Dense Linear Algebra, DLA）的发展史就是一部与存储层次结构（memory hierarchy）和通信开销博弈的历史。

LAPACK 支持 28 种矩阵类型，涵盖对称、三角、带状、三对角等结构。理解这些结构对算法选择至关重要。

BLAS 层次结构

基本线性代数子程序（BLAS）按数据重用能力分为三个层次。核心思想：Level k 的算法将 O(n^k) 操作作用于 O(n^k) 数据，从而在 cache 中实现高效重用。

为什么要最小化通信？

在并行和异构系统中，通信（数据移动）而非算术计算成为瓶颈。执行时间模型为：

T = (#flops) · γ + (#words) · β + (#messages) · α

其中：

• γ = 每 flop 时间（≈ 1/clock）
• β = 每 word 传输时间（≈ 1/bandwidth）
• α = 每次消息启动延迟（latency）

通信还消耗能量：移动 1 bit 的能量约是执行 1 flop 的 100–1000 倍。因此最小化通信同时优化能耗。

分块矩阵乘法回顾

经典的 cache-oblivious 分块 GEMM 算法：将 A、B、C 分块为大小为 b×b 的子矩阵，使得 3b² 数据能放入 cache。

// Blocked GEMM: C = A * B + C
// A: m x k, B: k x n, C: m x n
for (ii = 0; ii < m; ii += b)
  for (jj = 0; jj < n; jj += b)
    for (pp = 0; pp < k; pp += b)
      // Micro-kernel: C[ii:ii+b][jj:jj+b] += A[ii:ii+b][pp:pp+b] * B[pp:pp+b][jj:jj+b]
      for (i = ii; i < min(ii+b, m); i++)
        for (j = jj; j < min(jj+b, n); j++)
          for (p = pp; p < min(pp+b, k); p++)
            C[i][j] += A[i][p] * B[p][j];

通信下界

在并行矩阵乘法中，我们关心两个下界：最少需要移动多少数据（words），以及最少需要发送多少消息（messages）。

#words moved ≥ Ω(n³ / √M)

#messages ≥ Ω(n³ / M^(3/2))

#words = Ω(F / √M)
#messages = Ω(F / M^(3/2))

达到下界的算法

存在一系列算法能够渐近达到上述下界：

并行数据布局

在将矩阵分布到 P 个处理器时，数据布局决定了通信模式。

1D 布局

• Blocked：每个处理器持有连续的行/列块。通信仅在边界处发生。
• Cyclic：行/列循环分配（row i → processor i mod P）。更好的负载均衡。
• Block-Cyclic：每 b 行/列为一组循环分配。结合 blocked 的通信效率与 cyclic 的均衡性。

2D 布局

处理器排列为 P_r × P_c 网格。矩阵 A 按行分布、B 按列分布、C 按 (i,j) 分布。这是 ScaLAPACK 等库的标准选择。

并行矩阵-向量乘法

矩阵-向量乘法 y = Ax 是迭代方法（CG、GMRES）的核心。在 2D 布局下，每个处理器持有 A 的一个子块和 x 的对应部分：

1. 在拥有 x 的处理器列内做 broadcast（或 reduce）
2. 每个处理器计算本地 A_ij · x_j
3. 在行内做 reduce 得到 y_i

1D 矩阵乘法：总线与环形网络

最简单的并行 matmul：每个处理器持有一行（或一列）的数据。

Naïve 总线算法

// 每个处理器 i 持有 A 的第 i 行，需要 B 的所有列
for each processor i (in parallel):
    broadcast A[i][:] to all processors
    compute C[i][:] = A[i][:] * B

问题：每次 broadcast 都使用整个总线带宽，无法流水线化。

Ring 旋转算法

// 环形拓扑：处理器 i 将数据传给 (i+1) mod P
for step = 0 to P-1:
    each processor i:
        send A_local to processor (i+1) mod P
        C_local += A_local * B_panel
        receive new A_local from (i-1) mod P

SUMMA 算法

Scalable Universal Matrix Multiplication Algorithm (SUMMA) 是 2D 并行 GEMM 的经典算法，基于 outer product 形式：

// SUMMA: C = A * B, A: m x k, B: k x n
// 处理器网格: P_r x P_c
for (p = 0; p < k; p++) {
    // 拥有 A[:, p] 的处理器列 broadcast 到本列所有处理器
    if (I_own_column_p) broadcast A_col to my row;
    // 拥有 B[p, :] 的处理器行 broadcast 到本行所有处理器
    if (I_own_row_p) broadcast B_row to my col;
    // 每个处理器更新本地 C 块
    C_local += A_col * B_row;  // outer product
}

2.5D 矩阵乘法

2.5D 算法通过复制数据到 c 层来减少通信，在内存和通信之间取得平衡：

// 2.5D GEMM: 使用 c 份数据副本
// 将 P 个处理器组织为 c 个 P^(1/3) x P^(1/3) 的层
for each layer l in [0, c):
    // 每层执行独立的 SUMMA 子问题
    C_l += A_l * B_l;
// 跨层 reduce 结果
Reduce(C_0, C_1, ..., C_{c-1}) -> C;

#words = O(n³ / (P·√M)) · (c^(1/2))
#messages = O(n³ / M^(3/2)) · (1/c)

Strassen 算法

Strassen（1969）发现矩阵乘法可以用 7 次乘法而非 8 次完成：

// 将 A, B, C 分为 2x2 块
// 计算 7 个中间矩阵：
M1 = (A11 + A22) * (B11 + B22)
M2 = (A21 + A22) * B11
M3 = A11 * (B12 - B22)
M4 = A22 * (B21 - B11)
M5 = (A11 + A12) * B22
M6 = (A21 - A11) * (B11 + B12)
M7 = (A12 - A22) * (B21 + B22)

// 组合结果：
C11 = M1 + M4 - M5 + M7
C12 = M3 + M5
C21 = M2 + M4
C22 = M1 - M2 + M3 + M6

T(n) = 7·T(n/2) + O(n²) ⇒ T(n) = O(n^log₂7) ≈ O(n^2.81)

并行 Strassen 性能

并行 Strassen 的通信下界也降低：#words = Ω(n^log₂7 / √M)。这意味着 Strassen 不仅减少 flop，还减少通信量。然而递归开销和内存需求使其在实际中通常只在大规模问题（n > 10000）时才有优势。

ScaLAPACK PDGEMM

ScaLAPACK 的 PDGEMM 实现了基于 SUMMA 的分布式 GEMM：

• 使用 2D block-cyclic 分布
• 基于 BLACS（BLAS Communication）子程序
• 调用本地 BLAS3 DGEMM 作为 micro-kernel
• 支持任意 m, k, n 和 α, β 参数

// ScaLAPACK PDGEMM 调用示例
pdgemm_("N", "N", &m, &n, &k, &alpha,
        A, &ia, &ja, descA,
        B, &ib, &jb, descB,
        &beta,
        C, &ic, &jc, descC);

本课要点

知识检查