如何让Rust自动向量化生成VMULPS而非VMULSS？

我在支持AVX512的CPU上实现了16x16矩阵乘法，使用--release -C target-cpu=native编译。手动展开其中一个轴后，理论上优化器应该识别连续浮点数内存，将数据打包进ymm/zmm寄存器并使用vmulps做向量乘法，但实际生成的是标量乘法指令vmulss，性能显著降低；而看似未优化的朴素三重循环却生成了vmulps、vaddps等向量指令，性能快4倍。以下是具体分析：

两段代码的优化差异根源 #

手动展开的代码为何未被向量化 #

手动展开的代码片段：

for h in 0..16 {
    let aval = a[(v * n) + h];
    let bvals = &b[(h * n)..(h * n + 16)];
    g0 += aval * bvals[0];
    g1 += aval * bvals[1];
    g2 += aval * bvals[2];
    // ... 直到g15
}

这段代码的问题在于：

• 累加目标是分散的标量变量：g0到g15是16个独立的标量，优化器需要将这些分散的操作重新组合成向量操作，但这种变量分散的结构会增加优化器的分析难度，使其无法识别出可以将这些变量打包成向量寄存器的元素。
• 内存访问模式的间接性：通过切片bvals访问内存，虽然逻辑上是连续的，但切片的边界检查和引用语义可能让优化器需要额外确认内存的连续性与安全性，干扰了向量化判断。
• 标量操作的独立性：每个gX += aval * bvals[X]都是独立的标量更新，优化器难以将这些操作合并为批量的向量乘法与累加。

朴素三重循环为何能被高效向量化 #

朴素循环代码：

pub fn matmul(a: &[f32], b: &[f32], r: &mut [f32], m: usize, k: usize, n: usize) {
    let a = &a[0..(m * k)];
    let b = &b[0..(k * n)];
    for t in 0..k {
        for y in 0..m {
            for x in 0..n {
                r[(y * n) + x] += a[(y * n) + t] * b[(t * n) + x];
            }
        }
    }
}

这段代码的结构更贴合优化器的向量化逻辑：

• 内层循环的连续内存访问：固定t和y后，内层x循环中，b[(t*n)+x]是连续的内存块，r[(y*n)+x]也是连续的内存块。优化器能直接识别这种连续访问模式，将连续的f32元素加载到向量寄存器。
• 单一重复的操作模式：内层循环是重复执行的乘法累加操作，优化器可以轻松将这个循环替换为向量指令：广播a[(y*n)+t]到向量寄存器，与b的向量块做vmulps，再用vaddps累加到r的向量块上。
• 清晰的索引计算：直接通过数组索引访问内存，没有切片的间接性，优化器能更准确地分析内存访问的连续性和安全性。

不用unsafe/intrinsics实现高效矩阵乘法的建议 #

• 保持循环结构简洁：避免手动展开导致的变量分散问题，让优化器更容易识别向量化模式。优先使用朴素的三重循环，并选择对缓存和向量化友好的循环顺序（比如上述代码的t->y->x顺序，即KMN顺序）。
• 用数组存储累加值：如果需要手动展开，不要使用16个独立的标量变量，而是用数组let mut g = [0.0f32; 16];存储累加值，优化器更容易将数组识别为向量寄存器的载体，进而生成向量指令。
• 避免不必要的内存间接引用：尽量直接使用数组索引访问内存，减少切片、引用等可能干扰优化器分析的结构。
• 利用编译优化提示：可以尝试添加#[inline(always)]属性，或使用#[target_feature(enable = "avx512f")]（需配合target-cpu=native）来提示优化器，但注意不要过度干预，以免适得其反。

内容的提问来源于stack exchange，提问作者user22608671