CS267 · 第 19 课

稀疏矩阵-向量乘法与迭代求解器

稀疏矩阵是科学计算、图分析与机器学习的基本数据结构。SpMV 看似简单，却受限于内存带宽；理解它的存储格式、并行策略与自动调优，是掌握现代高性能计算的关键一步。

---

稀疏矩阵无处不在

当矩阵中绝大多数元素为零（通常 >90% 甚至 >99.9%）时，我们称之为稀疏矩阵。用稠密数据结构存储它们会浪费内存与计算。稀疏矩阵并非特例——它们是现代高性能计算中最常见的矩阵类型。

稀疏矩阵出现的领域极为广泛：

领域	具体应用	稀疏性来源
气候模拟	大气/海洋偏微分方程离散	局部差分格式
图像分析	图像分割（Image Contour Detection）	像素间的局部亲和矩阵
Web 搜索	Google PageRank	Web 图的邻接矩阵
推荐系统	Netflix 评分矩阵	用户仅评价极少数物品
图神经网络	GCN、GraphSAGE 的消息传播	图的邻接矩阵
文本分析	潜在语义索引（LSI）	词-文档 bag-of-words 矩阵
结构工程	有限元分析（NASA 结构问题）	局部单元连接性
线性规划	运筹优化	约束矩阵高度稀疏

以图像分割为例：我们希望识别图像中不同物体的边界。为此构造一个亲和矩阵（affinity matrix）\(A\)，其中 \(A_{ij}\) 表示像素 \(i\) 与像素 \(j\) 属于同一物体的"亲和力"。由于只有相邻且颜色相近的像素才有关联，这个矩阵极其稀疏。最终通过求解 \(Ax = \lambda x\) 的特征问题来找到分割边界。

---

稀疏矩阵存储格式

稀疏矩阵存储的核心矛盾是：既要节省空间，又要支持高效的矩阵-向量乘法。不同的稀疏模式需要不同的存储策略。

CSR（Compressed Sparse Row）：压缩稀疏行

CSR 是最常用的稀疏矩阵格式，由三个数组组成：

• val（长度 nnz）：所有非零元素的值，按行优先顺序排列
• ind（长度 nnz）：每个非零元素的列索引
• ptr（长度 m+1）：每行在 val/ind 中的起始位置，ptr[m] = nnz

CSR 与图的邻接表有着深刻的同构关系：ptr 就是每行的邻接表起始偏移，ind 就是邻居列表，val 就是边权重。一个有向图天然对应一个稀疏矩阵，无向图对应对称矩阵（只需存储一半）。这种同构使得许多图算法可以自然地表示为稀疏线性代数运算。

其他存储格式

格式	全称	特点	适用场景
CSC	Compressed Sparse Column	按列压缩，类似 CSR 的列视角	列操作频繁（如 MATLAB）
COO	Coordinate Format	(row, col, val) 三元组列表	构建/修改矩阵，GPU 友好
DIAG	Diagonal	按对角线存储，带偏移量	差分方程（有限差分/有限元）
Symmetric	对称存储	只存上/下三角，索引复杂	对称矩阵（省一半空间）
Blocked	分块格式	每个 r×c 块连续存储	矩阵具有天然块结构

---

串行 SpMV

SpMV（Sparse Matrix-Vector Multiplication）是稀疏线性代数中最核心的 kernel：\(y \leftarrow y + Ax\)。CSR 格式下的串行实现非常简洁：

def spmv_csr(ptr, ind, val, x, y, m):
    """y = y + A*x，A 以 CSR 格式存储，m 行"""
    for i in range(m):              # 遍历每一行
        for k in range(ptr[i], ptr[i+1]):  # 遍历该行所有非零元
            y[i] += val[k] * x[ind[k]]

分析这个 kernel 的内存访问模式：

• y[i]：每行只写一次，可完全放入寄存器——重用度最高
• val[k], ind[k]：顺序访问，空间局部性极佳（每对约 8 字节）
• x[ind[k]]：间接寻址，列索引不规则时完全是随机访问——这是性能瓶颈

矩阵 \(A\) 本身没有任何跨行的重用——每行读取一次就丢弃。唯一可能的重用来自 \(x\) 向量的缓存命中：如果多行访问 \(x\) 的同一位置且时间上足够接近，就能命中缓存。这就是为什么"近对角"或"良序"的矩阵 SpMV 性能更好——它们的列索引集中在较小范围内，\(x\) 的 working set 能放入缓存。

---

并行 SpMV

共享内存并行是最简单的 SpMV 并行化场景，因为每行的计算天然独立。

OpenMP 并行

在 CSR 的串行 SpMV 上直接加 OpenMP #pragma omp for 是最自然的并行方式：

#pragma omp parallel num_threads(thread_num)
{
#pragma omp for private(j, i, tmp) schedule(static)
   for (int i=0; i<m; i++)  {
       for (j = ptr[i]; j < ptr[i+1]; j++)  {
           tmp = ind[j];
           y[i] += val[j] * x[tmp];
       }
    }
}

调度策略的选择取决于非零元的分布：

• static：每行非零元数量相同时最优（如规则网格上的 PDE），零开销
• dynamic：行与行之间非零元数量差异巨大时避免负载不均，但有调度开销

SIMD 向量化（AVX2 Gather）

现代 CPU 提供 gather 指令（如 AVX2 的 vpgatherdd），可以从不规则地址加载向量。一个 8-wide 的 AVX2 SpMV 实现：

void avx2_csr_spmv(float *A, int32_t *nIdx, int32_t **indices,
                   float *x, int32_t m, float *y) {
  int32_t A_offset = 0;
  for(int32_t i = 0; i < m; i++) {
    int32_t nElem = nIdx[i]; float t = 0.0f;
    __m256 vT = _mm256_setzero_ps();
    int32_t smLen = nElem - (nElem & 7);  // 8 的倍数
    for(int32_t j = 0; j < smLen; j+=8) {
      __m256i vIdx = _mm256_loadu_si256((__m256i*)&(indices[i][j]));
      __m256 vX = _mm256_i32gather_ps((float const*)x, vIdx, 4);  // gather x
      __m256 vA = _mm256_loadu_ps(&A[A_offset + j]);              // 连续加载 A
      vT = _mm256_add_ps(vT, _mm256_mul_ps(vX, vA));              // FMA
    }
    t += sum8(vT);  // 水平归约
    for(int32_t j = smLen; j < nElem; j++)  // 尾部标量处理
      t += x[indices[i][j]] * A[A_offset + j];
    y[i] = t;
    A_offset += nElem;
  }
}

注意 gather 指令虽然强大，但性能仍受限于内存随机访问的延迟。当 \(x\) 的 working set 超过缓存容量时，gather 的吞吐会急剧下降。

CUDA 实现

GPU 上的朴素 CSR SpMV 让每行映射到一个线程：

// Parallel SpMV using CSR format
__global__ void spmv(int *ptr, int *ind, float *val,
                     int m, float *x, float *y) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < m; i += blockDim.x * gridDim.x) {
        float yi = 0;
        for (int j = ptr[i]; j < ptr[i+1]; j++) {
            yi += val[j] * x[ind[j]];
        }
        y[i] = yi;
    }
}

Segmented Suffix Scan（分段后缀扫描）

这是一种向量机时代的经典 SpMV 算法，思想在今天仍有启发性。核心是把 CSR 的"行边界"编码为 FLAG 数组，然后用分段扫描（segmented scan）完成行内归约：

1. FLAG：在 ptr 标记的行起始位置为 1，其余为 0
2. PROD = VAL * X(IND)：逐元素乘（gather x）
3. SUMS = SUM_SCAN(PROD, segment=FLAG)：分段前缀和，FLAG=1 处重置
4. Y = SUMS(PTR)：在 ptr 位置取值即为每行结果

这种表达把"行边界"信息从循环控制中完全消除，转化为数据并行的扫描操作，天然适合向量处理器和 GPU 的 SIMT 执行模型。

Diagonal 格式与 OpenMP

对角格式特别适合具有带状结构的矩阵（如有限差分生成的三对角矩阵）。每条对角线存储为一个一维数组，附带一个列偏移量：

for each diagonal k do
  #pragma omp parallel for
  for each row i do
    column = i + columnoffset[k]
    if (column >= 0 && column < n)
      y[i] = y[i] + val[k][column] * x[column]

这种格式的访问模式更规则（对角线上是步长为 1 的顺序访问），且天然适合 GPU 和向量机。进阶变体包括"锯齿对角"（jagged diagonal）以避免过度填充，以及按对角线长度排序以改善负载均衡。

---

分布式内存 SpMV

当矩阵大到单机内存放不下时，必须跨节点分布。SpMV 的分布式并行策略与稠密矩阵既有相似之处，也有本质差异。

行并行（1D Row Blocking）

将矩阵 \(A\) 和结果 \(y\) 按行划分给各处理器。每个处理器 \(p\) 持有 \(A_p\) 和 \(y_p\)，计算 \(y_p \leftarrow A_p x\)。关键问题：\(A_p\) 中某些列索引可能指向 \(x\) 中由其他处理器持有的分量，因此需要通信获取这些 \(x\) 的元素。

• 简单方案：完全复制 \(x\) 到所有处理器（当 \(x\) 不太大时可行）
• 高效方案：每个处理器只向"需要它拥有的 \(x\) 分量"的邻居发送数据（point-to-point 通信）

列并行（1D Column Blocking）

将矩阵 \(A\) 和向量 \(x\) 按列划分。每个处理器计算一个局部的 \(temp_y \leftarrow A_j x_j\)，然后通过稀疏归约（sparse sum-reduce）将所有 \(temp_y\) 累加到最终的 \(y\) 上。

2D 并行

将处理器组织成 \(p_1 \times p_2\) 的网格。矩阵 \(A\) 也按 2D 块划分，每个处理器只与同行/同列的处理器通信。这是 NAS CG benchmark 采用的策略，在 \(p\) 很大且非零元分布均匀时效果最好。

矩阵重排序：图划分的力量

重排序不会改变数学结果，但会彻底改变通信模式。这是分布式 SpMV 性能调优中最强大的技术之一。

HPCG Benchmark

HPCG（High Performance Conjugate Gradients）是 TOP500 的另一个 benchmark，专门测试稀疏线性系统的求解性能。它基于一个合成的 3D 热传导 PDE（27 点 stencil，每行 27 个非零元，对称正定）。

系统	HPCG [Pflop/s]	Rmax [Pflop/s]	HPCG/Peak	HPCG/HPL
Fugaku (A64FX)	16	442	3%	82%
Frontier (AMD EPYC)	14	1353	.7%	66%
Aurora (Intel Xeon Max)	5.6	1012	.3%	—
LUMI (AMD EPYC)	4.6	380	.9%	—

注意 HPCG/Typical 比值仅 0.3%-3%——这直观地说明了稀疏计算在当代超算上远远无法达到峰值。HPCG/HPL 比值（如 Fugaku 的 82%）则反映不同架构对稀疏 vs 稠密负载的适配度差异。

---

Register Blocking 与自动调优

SpMV 的瓶颈是从内存中取矩阵数据。每个非零元只产生 2 个 flop，却需要加载列索引（4 字节）和值（4 字节）——共 8 字节。如果能减少每个非零元的索引开销，就能提升有效带宽利用率。

r×c 分块 SpMV

核心思想：不要为每个非零元存储一个列索引，而是把矩阵分成 \(r \times c\) 的小块，每个块只存一个列索引。块内的 \(r \times c\) 个元素被当作稠密小矩阵处理（无需索引）。

• 当 \(r \times c\) 增大时，存储量下降（最多减少 ~2×，假设所有数据都是 32 位）
• 从内存取矩阵的时间减少 → 性能提升
• 需要同时修改数据结构和算法

以一个 NASA 结构分析矩阵（n=21200, nnz=1.5M）为例：\(r=c=8\) 时存储最小，看起来是最佳选择……但真的如此吗？

不同架构的最优块大小差异巨大：

架构	最佳块大小	峰值百分比
Power3	—	13%
Power4	—	14%
Itanium 2	4×2	31%
Itanium 1	—	7%
Ultra 2i	—	9%
Pentium III	—	19%

Fill Estimation（填充估计）

大多数矩阵没有整齐的块结构。以 FEM 流体力学矩阵（N=16614, NNZ=1.1M）为例：3×3 的块看起来很自然，但许多块并不完整——为了形成完整的 r×c 块，需要填充显式零（explicit zeros）。

填充的代价：

• 增加存储量（1.5× 的"填充开销"很常见）
• 增加无效计算（对零做乘法）
• 但循环可以展开，向量化更高效

在 Pentium III 上，1.5× 填充开销的 3×3 分块反而比无填充的实现快 1.5×——因为向量化收益超过了填充代价。这种 trade-off 强烈依赖于架构。

OSKI：Optimized Sparse Kernel Interface

OSKI（以及其多核版本 pOSKI）是 Berkeley 的 BeBOP 小组开发的自动调优稀疏 kernel 库。它的核心思想是：将矩阵格式的选择与优化延迟到运行时，通过搜索找到最佳配置。

自动选择块大小的三步流程

1. 离线基准测试（Offline Benchmark）：对每种 r×c 块大小，用稠密矩阵模拟稀疏格式跑 SpMV，测得 Mflops(r,c)。只需在每台机器上做一次。
2. 运行时搜索（Run-time Search）：对真实矩阵采样，估计每种 r×c 的 Fill(r,c)。通过控制采样比例 \(s \in [0,1]\) 控制开销。
3. 启发式模型（Heuristic Model）：选择使 \(\text{time} \sim \text{Fill}(r,c) / \text{Mflops}(r,c)\) 最小的 r, c。

自适应填充估计

通过计算统计置信区间来自动控制采样比例 \(s\)：监测采样方差，当方差足够小时停止采样。这避免了过度采样或采样不足。

机器学习辅助自动调优

后续研究将统计学习和贝叶斯优化引入自动调优：

• Vuduc, Demmel, Bilmes：基于统计模型的实证搜索性能调优
• Ganapathi：用统计机器学习预测系统利用率与性能
• Bergstra 等：用 Boosted Regression Trees 进行预测性自动调优
• Snoek 等：贝叶斯优化
• OpenTuner（Amarasinghe 组）：可扩展的程序自动调优框架
• GPTune（Sid-Lakhdar 等）：高斯过程自动调优

---

Cache Blocking

对于 CSR 上的 SpMV，唯一的数据重用机会在 \(x\) 向量中。如果矩阵的列索引集中在较小范围内（如近对角矩阵），\(x\) 的 working set 能放入缓存，自然就有良好的重用。

Cache blocking 对"短宽"矩阵（行数少、列数多）有效。以 LSI（潜在语义索引）矩阵为例：

• 矩阵大小：10k × 255k，3.7M 非零元
• 基线性能（无 cache blocking）：44 Mflop/s
• 最佳分块（16k × 16k）：210 Mflop/s（约 4.8× 加速）

---

Roofline 模型分析 SpMV

Roofline 模型为 SpMV 提供了直观的性能分析框架：

优化层次	关键效果
Out-of-the-box parallel	2 个单位步长流，无 ILP/DLP，FP 占 12-25%，compulsory flop:byte < 0.166
NUMA + SW prefetch	利用所有内存通道，减少 DRAM 延迟
Matrix compression (Register blocking)	提升 ILP、DLP、flop:byte 比和 FP 占比
SIMD	每周期更多 flop
ILP（指令级并行）	隐藏延迟

Register blocking 是 SpMV 上最有效的单优化：它同时改善了计算强度（减少索引开销）、ILP（块内循环展开）、DLP（向量化）和 FP 占比（减少整数地址计算）。实测中 register blocking 可带来最高 4× 的加速。

其他串行优化技术的典型收益：

优化技术	最高加速比
Register blocking	4× over CSR
Variable block splitting	2.1× over CSR, 1.8× over RB
Diagonals	2× over CSR
Reordering + splitting	2× over CSR
Symmetry	2.8× over CSR, 2.6× over RB
Cache blocking	2.8× over CSR
Multiple vectors (SpMM)	7× over CSR

---

Krylov 子空间方法

在实际应用中，我们很少只做一次 SpMV。迭代求解器（Iterative Solvers）的核心是一系列 SpMV 调用，用于求解 \(Ax = b\) 或 \(Ax = \lambda x\)。

Krylov 子空间方法通过构造 Krylov 子空间 \(\mathcal{K}_k(A, x) = \text{span}\{x, Ax, A^2x, \ldots, A^{k-1}x\}\)，在其中寻找最佳近似解。

经典 Krylov 方法概览：

方法	适用问题	矩阵要求	核心操作
CG (Conjugate Gradient)	\(Ax=b\)	对称正定 (SPD)	SpMV + 向量内积
GMRES	\(Ax=b\)	一般非对称	SpMV + Modified Gram-Schmidt
Lanczos	\(Ax=\lambda x\)	对称	SpMV + 三项递推
Arnoldi	\(Ax=\lambda x\)	一般非对称	SpMV + MGS

---

避免通信的迭代求解器

既然 SpMV 受限于内存/通信，而 Krylov 方法要做 \(k\) 次 SpMV，能否重新组织算法，使得总通信量不依赖于 \(k\)？这就是 Communication-Avoiding (CA) 迭代求解器的核心思想。

Matrix Powers Kernel：\([Ax, A^2x, \ldots, A^k x]\)

传统 Krylov 方法每次迭代做一次 SpMV，共需 \(k\) 次通信。Matrix Powers Kernel 一次性计算 \(x, Ax, A^2x, \ldots, A^k x\) 这一整组向量，通过冗余计算换取通信减少：

• 串行场景：传统方法需 \(O(k)\) 次数据从慢内存移到快内存；Matrix Powers Kernel 只需 \(O(1)\) 次——渐近最优
• 并行场景：传统方法需 \(O(k \log p)\) 条消息；Matrix Powers Kernel 只需 \(O(\log p)\) 条——渐近最优

关键洞察：如果矩阵"良划分"（well-partitioned，即 surface-to-volume 比小），每个处理器可以持有足够多的" ghost zone "（halo），使得连续 \(k\) 次 SpMV 中的前若干次不需要通信。

CA-GMRES：避免通信的 GMRES

标准 GMRES 的通信成本：

• Sequential：\(O(k \cdot \nnz)\) 来自 SpMV + \(O(k^2 \cdot n)\) 来自 MGS
• Parallel：\(O(k)\) 条消息来自 SpMV + \(O(k^2 \cdot \log p)\) 来自 MGS

CA-GMRES 用 Matrix Powers Kernel + TSQR（Tall Skinny QR）重构算法：

# 标准 GMRES（k 步）
for i in range(k):
    w = A @ v[i-1]              # SpMV → 1 次通信
    MGS(w, v[0], ..., v[i-1])   # 正交化 → log(p) 次通信
    update v[i], H

# CA-GMRES（k 步）
W = [v, Av, A^2v, ..., A^kv]   # Matrix Powers Kernel → O(1) 通信
[Q, R] = TSQR(W)               # Tall Skinny QR → O(log p) 通信
Build H from R, solve LSQ

通信成本对比：

实现	串行数据移动	并行消息数
标准 GMRES	\(O(k \cdot \nnz) + O(k^2 \cdot n)\)	\(O(k) + O(k^2 \cdot \log p)\)
CA-GMRES	\(O(\nnz) + O(k \cdot n)\)	\(O(1) + O(\log p)\)

在 8 核 Intel Clovertown 上的实测显示：CA-GMRES 在良划分矩阵上可以获得显著加速，但随着 \(k\) 增大，数值稳定性成为瓶颈。

相关工作

• Mark Hoemmen 的博士论文：CA Krylov 方法的理论与实现
• Erin Carson 的博士论文：避免通信的数值线性代数
• BeBOP 小组的系列工作：bebop.cs.berkeley.edu
• TSQR：Tall Skinny QR factorization（通信最优的 QR 分解）

---

SpGEMM：稀疏矩阵乘法

稀疏矩阵-稀疏矩阵乘法（SpGEMM）\(C \leftarrow A \times B\) 是许多高阶算法的基础（如代数多重网格、图算法中的三角形计数、图神经网络中的邻接聚合）。与 SpMV 相比，SpGEMM 的挑战在于：

• 输出矩阵 \(C\) 的规模在计算前不可预测（符号乘法 vs 数值乘法）
• 需要动态数据结构来累加行内的部分积
• 负载均衡更加困难

SpGEMM 的常见实现策略包括：

• 行内归约（row-wise reduction）：对 \(A\) 的每一行，遍历其非零元 \(A_{ik}\)，将 \(A_{ik} \cdot B_{k,:}\) 累加到 \(C_{i,:}\)
• Heap/Merge 方法：当一行中有多个非零元需要合并时
• Hash 方法：用哈希表动态管理 \(C\) 的列索引

SpGEMM 是 GPU 上最具挑战性的稀疏 kernel 之一，也是近年来的研究热点。后续课程会更详细地讨论图算法中的 SpGEMM 应用。

---

知识检查

---

要点总结

延伸阅读

• Jim Demmel, Mark Hoemmen, Marghoob Mohiyuddin, Kathy Yelick — Communication-Avoiding Krylov Methods
• Rich Vuduc 博士论文：Automatic Performance Tuning of Sparse Matrix Kernels (PDF)
• Eun-Jin Im 博士论文：Optimizing the Performance of Sparse Matrix-Vector Multiply (PDF)
• Guy E. Blelloch, Michael A. Heroux, Marco Zagha — Segmented Operations for Sparse Matrix Computation on Vector Multiprocessors (CMU-CS-93-173)
• Nishtala, et al. (2007) — When Cache Blocking of Sparse Matrix Vector Multiply Works and Why
• Sam Williams 等 — Optimization of Sparse Matrix-Vector Multiplication on Emerging Multicore Platforms
• BeBOP (Berkeley Benchmarking and Optimization Group)
• SuiteSparse Matrix Collection（原 UF Sparse Matrix Collection）

---