为什么 Rust 编译器不能自动向量化这个 FP 点积实现？

Question

让我们考虑一个简单的简化，例如点积：

pub fn add(a:&[f32], b:&[f32]) -> f32 {
    a.iter().zip(b.iter()).fold(0.0, |c,(x,y)| c+x*y))
}

使用 rustc 1.68 与-C opt-level=3 -C target-feature=+avx2,+fma 我得到

.LBB0_5:
        vmovss  xmm1, dword ptr [rdi + 4*rsi]
        vmulss  xmm1, xmm1, dword ptr [rdx + 4*rsi]
        vmovss  xmm2, dword ptr [rdi + 4*rsi + 4]
        vaddss  xmm0, xmm0, xmm1
        vmulss  xmm1, xmm2, dword ptr [rdx + 4*rsi + 4]
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 8]
        vmulss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 8]
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 12]
        vmulss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 12]
        lea     rax, [rsi + 4]
        vaddss  xmm0, xmm0, xmm1
        mov     rsi, rax
        cmp     rcx, rax
        jne     .LBB0_5

这是一个带有循环展开的标量实现，甚至没有将 mul+add 收缩到 FMA 中。从这段代码到simd代码应该很容易，为什么rustc不对此进行优化？

如果我替换f32为，i32我会得到所需的自动矢量化：

.LBB0_5:
        vmovdqu ymm4, ymmword ptr [rdx + 4*rax]
        vmovdqu ymm5, ymmword ptr [rdx + 4*rax + 32]
        vmovdqu ymm6, ymmword ptr [rdx + 4*rax + 64]
        vmovdqu ymm7, ymmword ptr [rdx + 4*rax + 96]
        vpmulld ymm4, ymm4, ymmword ptr [rdi + 4*rax]
        vpaddd  ymm0, ymm4, ymm0
        vpmulld ymm4, ymm5, ymmword ptr [rdi + 4*rax + 32]
        vpaddd  ymm1, ymm4, ymm1
        vpmulld ymm4, ymm6, ymmword ptr [rdi + 4*rax + 64]
        vpmulld ymm5, ymm7, ymmword ptr [rdi + 4*rax + 96]
        vpaddd  ymm2, ymm4, ymm2
        vpaddd  ymm3, ymm5, ymm3
        add     rax, 32
        cmp     r8, rax
        jne     .LBB0_5

Answer 1

这是因为浮点不具有关联性，即一般意义a+(b+c) != (a+b)+c。因此，对浮点数求和就变成了串行任务，因为编译器不会将其重新排序((a+b)+c)+d为(a+b)+(c+d). 最后一个可以矢量化，第一个不能。

在大多数情况下，程序员并不关心求和顺序的差异。

gcc 和 clang 提供了-fassociative-math允许编译器重新排序浮点运算以提高性能的标志。

rustc 不提供此功能，据我所知 llvm 也不接受会改变此行为的标志。

在夜间 Rust 中，您可以使用#![feature(core_intrinsics)]来获得优化：

#![feature(core_intrinsics)]
pub fn add(a:&[f32], b:&[f32]) -> f32 {
    unsafe {
        a.iter().zip(b.iter()).fold(0.0, |c,(x,y)| std::intrinsics::fadd_fast(c,x*y))
    }
}

这不使用 fma。所以对于 fma 你必须使用：

#![feature(core_intrinsics)]
pub fn add(a:&[f32], b:&[f32]) -> f32 {
    unsafe {
        a.iter().zip(b.iter()).fold(0.0, |c,(&x,&y)| std::intrinsics::fadd_fast(c,std::intrinsics::fmul_fast(x,y)))
    }
}

我不知道有一个稳定的 Rust 解决方案，它不涉及显式的 simd 内在函数。