SIMD 执行允许一条硬件指令同时处理多个数据点。处理器不再是加载一个数字，将其加到另一个数字，然后存储结果，而是加载一组“矢量”数字（内存中连续的数据元素），一次性对它们全部执行操作，然后存储结果。对于主要由逐元素操作和矩阵乘法构成的机器学习任务，矢量化是一项必要的优化，可以带来与目标架构矢量宽度成比例的性能提升。

了解矢量寄存器

在硬件层面，矢量化依赖于专用的矢量寄存器。与可能只保存单个 32 位或 64 位值的标准通用寄存器不同，矢量寄存器要宽很多。常见的指令集架构（ISA）包含矢量扩展，例如 x86 CPU 上的 AVX-512（512 位）或 ARM 处理器上的 NEON（128 位）。

如果使用单精度浮点数（float32），一个 512 位的寄存器可以容纳 16 个值（ $512 / 32 = 16$ ）。这意味着一条指令在理论上可以完成 16 次加法或乘法，而标量处理器在相同时间内只能完成一次。

标量执行与 SIMD 执行处理四个数据元素的对比。标量方法需要四个不同的指令周期，而 SIMD 方法一步即可处理所有元素。

循环分段与重写

当 ML 编译器对循环进行矢量化时，它会从根本上改变迭代结构。编译器会进行一个常被称为“分段（stripmining）”的过程。它调整循环的步长，使其与硬件的矢量宽度相匹配。

假设将两个数组 $A$ 和 $B$ 进行简单的逐元素相加，存入结果数组 $C$ 。在标准的标量实现中，代码会遍历每个索引 $i$ ：

// 标量循环
for (int i = 0; i < 1024; i++) {
    C[i] = A[i] + B[i];
}

如果目标硬件的矢量宽度为 8（例如，使用 float32 的 AVX2），编译器会将此循环转换为每次迭代处理 8 个元素。生成的伪代码如下所示：

// 矢量化循环
for (int i = 0; i < 1024; i += 8) {
    // 将 8 个元素从内存载入矢量寄存器
    vec_float8 a_val = vector_load(A + i);
    vec_float8 b_val = vector_load(B + i);
    
    // 同时对所有 8 个通道执行加法
    vec_float8 c_val = vector_add(a_val, b_val);
    
    // 将结果矢量存回内存
    vector_store(C + i, c_val);
}

这种转换大幅减少了循环开销，因为循环条件检查和索引增量发生的频率降低了 8 倍。更重要的是，它使软件指令流与硬件的并行执行单元保持一致。

处理循环尾部和边界

矢量化面临的一个主要问题是，元素总数 $N$ 不总是能被矢量宽度 $W$ 整除。如果 $N = 100$ 且我们的矢量宽度为 8，那么主矢量化循环只能处理前 96 个元素（ $12 \times 8$ ）。剩下的 4 个元素构成“尾部”或“收尾部分”。

编译器通过生成两个不同部分的代码来处理这种情况：

矢量循环： 处理大部分数据，其中索引完美适配矢量寄存器。
标量清理循环： 每次迭代处理一个元素，以完成剩余工作。

逻辑结构如下：

$\text{矢量迭代次数} = \lfloor \frac{N}{W} \rfloor$ $\text{余数} = N \pmod W$

在编译器中间表示（IR）中，你常会看到循环被明确地拆分。如果大小 $N$ 在编译时已知（静态形状），编译器可能会进一步优化尾部处理。如果 $N$ 是动态的，编译器必须生成运行时检查，以确定矢量循环运行的次数以及标量循环是否必需。

矢量化操作的控制流生成过程。编译器通过对边界元素回退到标量执行来确保处理任意输入大小。

数据对齐和内存访问

当从“对齐”的内存地址载入数据时，矢量指令的效率最高。如果地址是矢量大小（以字节计）的倍数，则称其为对齐的。例如，一个 256 位（32 字节）的矢量载入操作，如果内存地址以 0x0、0x20、0x40 等结尾，则效果最佳。

如果数据未对齐，硬件可能需要执行多次内存请求并拼接数据，这会带来性能损失。在一些较旧的架构中，未对齐的矢量载入甚至会直接导致程序崩溃。现代 ML 编译器在内存规划上投入了大量精力，以确保为张量分配的缓冲区能对齐到 64 字节或 128 字节边界。

在优化内核时，你可能会遇到“掩码”载入或存储。如果编译器无法保证循环在矢量化阶段不会读取超出数组末尾，它可能会使用掩码，即一个布尔矢量，用于禁用特定通道。这使得即使处理尾部元素也能安全使用矢量指令，尽管与干净、对齐的矢量载入相比，这通常会带来轻微的性能开销。

编译器的矢量化原语

在 TVM 或通过 MLIR 优化的循环等标准深度学习编译器中，矢量化通常是明确的。与依赖启发式方法进行“自动矢量化”的 C++ 编译器（如 GCC 或 Clang）不同，ML 编译器允许开发人员或自动调优代理明确标记循环以进行矢量化。

例如，在调度语言中，你可能会看到如下命令：

schedule[output].vectorize(inner_axis)

该指令强制编译器在 inner_axis 处按硬件的矢量因子展开循环，并将标量操作替换为矢量内在函数（如 llvm.vector.fadd）。如果操作包含阻止并行执行的依赖关系，例如 A[i] 依赖于 A[i-1] 的递归关系，编译器将报错或无法应用此转换，因为矢量化严格要求矢量块内的数据独立性。

这部分内容有帮助吗？