CS267 · 第 6b 课 · James Demmel

通信避免矩阵乘法（及更多）

通信（数据移动）是科学计算中最昂贵的操作。本课探讨如何证明并构造通信最优的算法，涵盖矩阵乘法、线性代数、CNNs 等。

授课教师：James Demmel

日期：2025-02-06

难度：高级

视频：Lecture 6b Video

通信 = 数据移动

通信是科学计算中最昂贵的操作（时间或能量）：

主内存 ↔ 缓存
处理器 ↔ 处理器（通过网络）

目标：为形似"嵌套循环访问数组"的算法，证明并构造通信最优算法。

包括：矩阵乘法、线性代数（稠密和稀疏）、N-body、CNNs 等。

N-body 的通信下界

简化模型

for i = 0:n
  for j = 0:n
    access A(i), B(j)

若缓存大小为 M，只能访问 M 个 A 和 B 的条目，最多能执行多少循环迭代？

直觉分析

若缓存中有 M/2 个 A 条目和 M/2 个 B 条目：

可执行 \((M/2) \times (M/2) = M^2/4\) 次迭代
总共需要 \(n^2\) 次迭代
需要刷新缓存 \(n^2 / (M^2/4) = 4(n/M)^2\) 次
每次刷新需要 M 次读
总读次数：\(4(n/M)^2 \cdot M = 4n^2 / M\)

通信下界： \(\Omega(n^2 / M) = \Omega(\text{\#iterations} / M)\)

最优分块

// 读 M/5 个 A 的条目: A([0,1,2,...,M/2-1])
// 读 M/2 个 B 的条目: B([0,1,2,...,M/2-1])
// 执行 (M/2)² = M²/4 次迭代

矩阵乘法的通信下界

几何证明（Loomis-Whitney 不等式）

将矩阵乘法的迭代表示为 3D 空间中的点 \((i,j,k)\)：

for i = 1:n, for j = 1:n, for k = 1:n
  C(i,j) += A(i,k) * B(k,j)

每个迭代访问三个"面"上的元素：

A face: \((i,k)\)
B face: \((j,k)\)
C face: \((i,j)\)

Loomis-Whitney 不等式 (1949)：

\[\text{\#cubes in 3D set} \leq (\text{area(A shadow)} \cdot \text{area(B shadow)} \cdot \text{area(C shadow)})^{1/2}\]

若有 M 个数据可用：

每个面最多 M 个方块
最多 \(M^{3/2}\) 次迭代
需要刷新缓存 \(n^3 / M^{3/2}\) 次
通信下界：\(\Omega(n^3 / M^{1/2}) = \Omega(\text{\#iterations} / M^{1/2})\)

并行情况

对于 P 个处理器：

每处理器迭代数：\(n^3/P\)
每处理器数据：\(M = 3n^2/P\)
通信下界：\(\Omega(n^2 / P^{1/2})\)

最优算法

6 层嵌套循环的分块矩阵乘法达到了这个下界：

for i = 1 to N step b
  for j = 1 to N step b
    for k = 1 to N step b
      // b×b 的块矩阵乘法，b = O(M^{1/2})
      C[i:i+b, j:j+b] += A[i:i+b, k:k+b] * B[k:k+b, j:j+b]

一般化：嵌套循环的通信下界

一般形式

for i1 = 1:n, for i2 = i1:m, ..., for ik = i3:i4
  C(i1+2*i3-i7) = func(A(i2+3*i4, i1, i2, ...), B(...), ...)

表示为：

迭代空间：\((i_1, i_2, ..., i_k) \in S \subset \mathbb{Z}^k\)
数组访问由投影 \(\phi_j\) 定义

一般通信下界

定理：

\[\text{words\_moved} = \Omega(\#iterations / M^{s_{HBL}-1})\]

其中 \(s_{HBL}\) 是一个线性规划的最优值：

\[\text{minimize } s_{HBL} = \sum_j e_j \text{ subject to } \text{rank}(H) \leq \sum_j e_j \cdot \text{rank}(\phi_j(H))\]

证明依赖于纯数学的最新成果：

Christ/Tao/Carbery/Bennett
Hölder-Brascamp-Lieb (HBL) 不等式在 Abel 群上的推广
HBL 推广了 Cauchy-Schwartz、Loomis-Whitney 等

是否可达？

定理：总能构造一个最优分块算法达到下界。

假设/限制：

达到 \(\Omega(\#iterations / M^{s_{HBL}-1})\)（在 O() 意义上）
取决于循环依赖（仅 reductions 如 matmul 可任意重排）
假设循环边界足够大以容纳分块

注意：对于 matrix-vector multiply，同样的下界形式适用，但不可达（因为循环边界不够大）。

CNNs 的通信下界

CNN 计算

for k = 1:K, for h = 1:H, for w = 1:W, for r = 1:R,
    for s = 1:S, for c = 1:C, for b = 1:B
  Out(k, h, w, b) += Image(r+w, s+h, c, b) * Filter(k, r, s, c)

总迭代数 \(N = K \cdot H \cdot W \cdot R \cdot S \cdot C \cdot B\)

通信下界

\[\text{words\_moved} = \Omega(\max(\]

\(BKHW\) — 输出大小
\(\sigma_H \sigma_W BCWH\) — 图像大小
\(CKRS\) — 滤波器大小
\(N/M\) — 与 n-body 相同的下界
\(N / (M^{1/2} (RS/(\sigma_H \sigma_W))^{1/2})\) — 新下界

\[)\]

新下界的优势：

比 matmul 下界好 \((RS/(\sigma_H \sigma_W))^{1/2}\) 倍
适用于数据不 fit in cache 但单个 R×S 滤波器 fit 的情况
许多情况下可达（用 Mathematica 求解）

自动化工具

iocomplexity.corse.inria.fr/iolb — 自动计算并构造通信最优的分块。

一般化分块示例

"倾斜" N-body

for i = 0:n
  for j = 0:n
    access A(i), B(i+j)

最优分块：读 A([0,1,2,3,4]) 和 B([4,5,6,7,8])，执行 25 次迭代。

"扭曲" N-body

for i = 0:n
  for j = 0:n
    access A(3*i-j), B(i-2*j)

最优分块：读 A([0,5,10,15,20]) 和 B([0,-5,-10,-15,-20])，执行 25 次迭代。

测验

Q1：N-body 的通信下界是什么？

Q2：矩阵乘法的通信下界是什么？

Q3：证明矩阵乘法通信下界使用了什么不等式？

Q4：为什么 matrix-vector multiply 的通信下界不可达？

Q5：CNNs 的新通信下界比 matmul 好多少倍？

要点总结

通信下界是理论极限

通过 Loomis-Whitney 不等式和 HBL 不等式，我们可以证明各种算法的通信下界。

分块可以达到最优

对于 matmul 等算法，分块矩阵乘法达到了通信下界。

一般化框架

s_HBL 线性规划框架可以处理任意"嵌套循环访问数组"的算法。

CNNs 有更好的下界

CNNs 的新通信下界比 matmul 好，因为可以利用滤波器的局部性。