4. 稀疏矩阵向量乘的向量化与SIMD优化

稀疏矩阵向量乘（SpMV）计算 y = A·x，其中 A 为 m×n 稀疏矩阵，x 为稠密向量。它是共轭梯度、PageRank、图神经网络等应用的核心内核。SpMV 的算术强度极低（每读取一个非零元只做一次乘加），性能瓶颈几乎完全落在内存带宽和非规则访存上。SIMD 向量化的价值不在于减少浮点运算次数，而在于用一条指令处理多个非零元或连续输出，从而摊薄地址计算与循环控制开销，并提高有效访存吞吐。

需要先明确一个前提：SpMV 的向量化难度与稀疏结构强相关。随机稀疏矩阵的列索引不可预测，直接对 CSR 的 values 做 SIMD 收集（gather）往往得不偿失；而对角线、带状等结构化稀疏，索引具有规律性，可以设计出连续访存、无冲突的向量化内核。因此本节按“结构决定重排策略，重排策略决定内核形态”的思路展开。

三类稀疏结构的访存特征与重排思路

对角线结构：非零元集中在若干条对角线上，可用“对角线存储”（DIA）格式，每条对角线用一个长度为 n 的数组保存，配合偏移量 offset。此时 y[i] += diag[d][i - offset] * x[i - offset]，内层循环对 i 连续，天然适合 SIMD。重排的关键是把对角线按 offset 排序，并把长度不足 n 的对角线用零填充到统一长度，使每条对角线都能按向量宽度整块处理。

带状结构：非零元位于主对角线附近宽度为 b 的带内，可视为对角线的推广。若带宽较小（如 b ≤ 8），可直接用“带存储 + 分块”策略：把矩阵按行切成 R 行的块，每块内非零元按列偏移落入固定槽位，形成规整的 R×b 小矩阵，便于用向量寄存器一次装载一列。

随机结构：列索引无规律，CSR 是最自然的格式。此时不应强求对非零元做 gather，而应转向“按行分组 + 输出向量化”：把若干行打包成一个行组，使组内非零元总数接近向量宽度的整数倍，再对行组做展开。另一种有效手段是 ELL 格式重排：将每行非零元填充到相同长度 max_nnz，列索引矩阵按列主序存储，使同一列位置上的多个行可以连续读取。

向量化内核的实现要点

以 AVX2（8 个 float32）为例，对角线内核的核心循环可以写成：

// DIA 格式：diag[d] 长度 n，offset[d] 为对角偏移
for (int d = 0; d < ndiag; ++d) {
    const float* dv = diag[d];
    int off = offset[d];
    int lo = off > 0 ? off : 0;
    int hi = off < 0 ? n + off : n;
    for (int i = lo; i < hi; i += 8) {
        __m256 yv = _mm256_loadu_ps(&y[i]);
        __m256 xv = _mm256_loadu_ps(&x[i - off]);
        __m256 av = _mm256_loadu_ps(&dv[i]);
        yv = _mm256_fmadd_ps(av, xv, yv);
        _mm256_storeu_ps(&y[i], yv);
    }
}

这里有三处值得注意。第一，lo 与 hi 的边界处理保证 x[i-off] 不越界，同时避免对零填充区做无效计算。第二，_mm256_fmadd_ps 把乘加融合为一条指令，减少中间结果写回。第三，若 n 不是 8 的倍数，尾部需用标量循环收尾，或把对角线数组填充到 8 的倍数，使主循环完全向量化。

对 ELL 格式，内核形态不同：列索引矩阵 col_idx 按列主序存储，同一列位置上的 8 行可以连续读取，但 x 的访问仍需按索引收集。此时更稳妥的做法是“行组展开”：把 8 行打包，对组内每个非零元做标量乘加，但把 8 个输出累加器放在向量寄存器中，最后一次性写回 y。这样避免了 gather，代价是组内非零元数量不齐时需要掩码或填充。

标量与向量化实现的性能对比

下面给出一组教学示例参数（非实测，用于说明量级关系）：矩阵规模 20000×20000，每行平均 6 个非零元，共约 12 万非零元；x、y 均为 float32；平台为支持 AVX2 的 4 核 CPU，内存带宽约 40 GB/s。标量 CSR 内核每处理一个非零元需一次索引读取、一次 x 读取、一次乘加和一次 y 累加，且 y 的累加存在写后读依赖。向量化 DIA 内核则把 8 个连续元素打包处理，索引计算被摊薄，y 的更新变为整块读写。

判断向量化是否真正生效，不能只看总耗时。建议同时观察三个指标：一是有效非零元吞吐（非零元数除以耗时），若向量化后吞吐提升但总时间未降，说明填充或收集开销吃掉了收益；二是缓存未命中率，DIA 与带存储应显著低于 CSR；三是尾部标量循环占比，若占比过高，应调整分块大小或填充策略。对随机结构，若 ELL 的填充率超过 30%，通常意味着行长度分布过于不均，此时应改用分块 CSR 或按行组分桶，而不是继续加大填充。

最后需要强调，SIMD 优化必须与数据布局协同设计。先根据稀疏结构选定 DIA、带存储或 ELL，再针对该布局写出连续访存的内核，最后用有效吞吐和缓存行为验证收益；反过来先写内核再迁就布局，往往只能得到 gather 密集、收益有限的实现。