【infra之路】AI 编译器学习笔记（八）：向量化——把小块变成 SIMD 指令

上一篇把迭代空间切成了小块。这一篇讲切出来的小块怎样变成 SIMD 指令：vector dialect 是什么、矩阵乘为什么要用"外积"的方式向量化、块大小和寄存器数量的关系，以及什么时候会退化成标量。以 x86 AVX2 为例，API 以 LLVM 20 为准，IR 为示意。

---

一、硬件：SIMD

普通加法指令一次算一个数，SIMD（单指令多数据）指令一次对一整排数做同样的操作。AVX2 中：

• 一个 ymm 寄存器 256 位，放 8 个 f32；
• 一条 vfmadd 同时完成 8 次乘加；
• 共 16 个 ymm 寄存器。

标量循环一次一个元素，向量化后一次 8 个，理论上快 8 倍。

GPU 不同：单个线程是标量的，并行靠大量线程。GPU 上的"向量化"主要指向量化访存（一次读写 128 位）和 Tensor Core 的矩阵指令，见第八节。

二、在流水线中的位置

linalg（大）
  → tiling：切成形状固定的小 linalg op
  → 向量化：小 linalg op → vector dialect      ← 本篇
  → vector 降级：多维 vector → 一维 vector
  → llvm dialect → LLVM 后端选出具体 SIMD 指令

向量化在 tiling 之后，作用在小块上，原因见第六节。

三、vector dialect

vector<4x8xf32> 表示"4×8 个 f32 组成的一整块，假定放在寄存器中"。和 tensor、memref 不同，它不在内存里、没有地址，是直接拿来计算的一组值。

一个 4×8 的 relu 小块向量化后：

%v = vector.transfer_read %x[%c0, %c0], %pad : tensor<4x8xf32>, vector<4x8xf32>
%z = arith.constant dense<0.0> : vector<4x8xf32>
%m = arith.maximumf %v, %z : vector<4x8xf32>
%w = vector.transfer_write %m, %init[%c0, %c0] : vector<4x8xf32>, tensor<4x8xf32>

• vector.transfer_read：从 tensor/memref 读一整块；%pad 是越界时的填充值。
• arith 的 op 对标量和向量都适用，作用在 vector 上就是逐元素计算。
• vector.transfer_write：写回。

本质上是把 linalg.generic region 里"一个点怎么算"换成"整块怎么算"。

四、矩阵乘：vector.contract 与外积

matmul 小块向量化后得到 vector.contract：

%r = vector.contract {
       indexing_maps = [
         affine_map<(i, j, k) -> (i, k)>,
         affine_map<(i, j, k) -> (k, j)>,
         affine_map<(i, j, k) -> (i, j)>],
       iterator_types = ["parallel", "parallel", "reduction"]
     } %a, %b, %acc : vector<4x1xf32>, vector<1x8xf32> into vector<4x8xf32>

它与 linalg.generic 的写法几乎相同：映射和维度类型原样保留。结构信息从高层一路传下来没有丢，这是 Linalg 设计的好处。

vector.contract 还不是硬件指令，接下来会降成外积形式，这是 CPU GEMM 微内核的经典做法。

设寄存器块为 4×8，累加器 acc 正好占 4 个 ymm （每行一个）。每走一步 k：

1. 读 B 的第 k 行连续 8 个数 → 一个 ymm（b0..b7）
2. A 第 k 列的 4 个数，每个广播成一个 ymm（8 份相同的值）
3. acc[0] += 广播(A[0][k]) * b     ← 一条 vfmadd，8 次乘加
   acc[1] += 广播(A[1][k]) * b
   acc[2] += 广播(A[2][k]) * b
   acc[3] += 广播(A[3][k]) * b

每步 4 条 FMA、32 次乘加。累加器全程留在寄存器，k 走完才写回一次。

五、为什么偏偏是外积：向量该沿哪个维度

外积方案中，8 个 lane 沿 j 方向：

• B[k][j..j+7] 沿 j 连续，一次读入；
• C[i][j..j+7] 也沿 j 连续，累加器按行直接存回；
• A[i][k] 在 8 个 lane 上相同，所以广播。

我一开始想反过来：对 A 按行取 8 个连续的数 A[i][k..k+7]，让向量沿 k。这样做有两个问题：

1. 与之相乘的是 B[k..k+7][j]，即 B 的一列，8 个不同的数，无法用广播代替；而行优先存储下一列是不连续的，要逐个读或 gather。（只有 B 恰好对称时，列才等于行，这只是特例。）
2. k 是 reduction 维度，8 个 lane 的乘积最后要水平求和成一个数，需要多条 shuffle 指令。

规则：向量沿着既连续、又是 parallel 的维度 ，行优先的矩阵乘里就是 j。

补充两点：

• 如果事先把 B 转置或**打包（packing）**存放，沿 k 的"点积"方案也能连续读，水平求和可以拖到 k 走完再做一次。GEMV 常用这种写法。真实的 BLAS 库也会对 A、B 做 packing，让任意方向的访问都连续。
• GEMM 仍偏爱外积，原因是复用：读进来的一行 B 要和 4 个广播值各做一次 FMA，读一次用 4 次；点积方案每读一对向量只用一次。每次访存换来多少计算（计算访存比），决定了能否接近硬件峰值。

六、块大小与寄存器：spill

块大小的约束：

• 列方向取 8 的倍数，正好装满 ymm；
• 累加器、B 的一行、A 的广播值加起来，要在 16 个寄存器以内。

以 8×16 的寄存器块为例：累加器 128 个 f32，每行占 2 个 ymm，8 行共 16 个 ymm；再加 B 的一行（2 个）和至少 1 个广播值，至少需要 19 个，超过 16 个。

这时编译器会 spill：把一部分值临时存到栈上，用时再读回来，每次多出一对内存存取，性能明显下降。

要和另一件事区分开：一行 16 个数拆成两条 8 宽指令是正常的展开（unroll），不损失性能；spill 才是寄存器不够的代价。

换成 AVX-512（32 个 zmm，每个 16 个 f32），同样的 8×16 累加器只占 8 个寄存器，就放得下了。合适的块大小取决于目标硬件，这就是"寄存器 tiling"的含义。

这也解释了为什么必须先 tiling 再向量化：直接对 128×128 的 matmul 向量化会得到 vector<128x128xf32>，展开后是两千多个寄存器值，只会疯狂 spill。

七、多维 vector 的降级，以及何时退化成标量

硬件寄存器是一维的。vector<4x8xf32> 降级时被展开成 4 个 vector<8xf32>，各对应一个 ymm。形状对不齐时（如 vector<4x12xf32>，每行拆成 8 + 4）会浪费半个寄存器，所以 tiling 时最内层维度应是硬件宽度的整数倍。

退化成标量的典型原因：

1. 动态形状 ：块大小不是编译期常量，向量化器不知道生成多大的 vector。解决办法：先 tiling 成固定大小、余数单独处理（剥离或填充）；或者向量化时给出 vector_sizes，用 mask 让越界的 lane 不读不写。什么都不做时向量化器不会报错，只是跳过这个 op，最终降成标量循环------性能异常时要先检查关键 op 是否被向量化。
2. 不连续访存：跨步访问需要 gather 或逐个读，收益大减。
3. op 没有向量版本：只能逐元素调用。
4. 别名不确定 （LLVM 自动向量化中）：不确定两个指针是否重叠就不敢整排读写，C 的 restrict 就是为此存在。

八、GPU 上的对应物

• 向量化访存 ：每个线程一次读 4 个连续 f32（PTX 的 ld.global.v4.f32）。
• Tensor Core ：一条 mma 指令由一个 warp 合作完成一个小矩阵乘（如 16×8×16），相当于把外积微内核直接做进硬件。MLIR 中 vector.contract 可降成 nvgpu.mma.sync，再到 NVVM intrinsic。

Triton 的 tl.dot 最终落在 Tensor Core 指令上；在 TTGIR 层，编译器要决定块中每个元素归哪个线程、在哪个寄存器，即 Triton 的 layout 机制。

九、动手试

transform dialect 中的 transform.structured.vectorize（单个 op，可指定 vector_sizes）和 transform.structured.vectorize_children_and_apply_patterns（函数内所有可向量化的 op）。接着上一篇的实验：先 tiling，再向量化，用 --mlir-print-ir-after-all 观察 vector.contract 如何一步步变成 vector.fma。语法以 LLVM 源码 mlir/test/Dialect/Linalg/ 下 vectorize 相关测试为准。

---

小结

• vector dialect 表示寄存器中的一整块值；transfer_read/write 负责与内存交换。
• matmul 向量化为 vector.contract，再降为外积：读一行 B、广播 A、逐行 FMA，累加器常驻寄存器。
• 向量沿连续且 parallel 的维度；外积胜在复用。
• 块大小受寄存器数量约束，超出就 spill；展开不等于 spill。
• 先 tiling 再向量化；动态形状、不连续访存等会导致退化成标量。

下一篇：Bufferization------tensor 怎样变成 memref，insert_slice 为什么最后没有真的复制。