SIMD 性能缓慢 - 无内联

Question

考虑以下计算 i32 数组总和的示例：

示例 1：简单的 for 循环

pub fn vec_sum_for_loop_i32(src: &[i32]) -> i32 {
    let mut sum = 0;
    for c in src {
        sum += *c;
    }

    sum
}

示例 2：显式 SIMD 总和：

use std::arch::x86_64::*;
// #[inline]
pub fn vec_sum_simd_direct_loop(src: &[i32]) -> i32 {
    #[cfg(debug_assertions)]
    assert!(src.as_ptr() as u64 % 64 == 0);
    #[cfg(debug_assertions)]
    assert!(src.len() % (std::mem::size_of::<__m256i>() / std::mem::size_of::<i32>()) == 0);

    let p_src = src.as_ptr();
    let batch_size = std::mem::size_of::<__m256i>() / std::mem::size_of::<i32>();

    #[cfg(debug_assertions)]
    assert!(src.len() % batch_size == 0);

    let result: i32;
    unsafe {
        let mut offset: isize = 0;
        let total: isize = src.len() as isize;
        let mut curr_sum = _mm256_setzero_si256();

        while offset < total {
            let curr = _mm256_load_epi32(p_src.offset(offset));
            curr_sum = _mm256_add_epi32(curr_sum, curr);
            offset += 8;
        }

        // this can be reduced with hadd.
        let a0 = _mm256_extract_epi32::<0>(curr_sum);
        let a1 = _mm256_extract_epi32::<1>(curr_sum);
        let a2 = _mm256_extract_epi32::<2>(curr_sum);
        let a3 = _mm256_extract_epi32::<3>(curr_sum);
        let a4 = _mm256_extract_epi32::<4>(curr_sum);
        let a5 = _mm256_extract_epi32::<5>(curr_sum);
        let a6 = _mm256_extract_epi32::<6>(curr_sum);
        let a7 = _mm256_extract_epi32::<7>(curr_sum);

        result = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7;
    }

    result
}

当我尝试对代码进行基准测试时，第一个示例得到了 ~23GB/s（接近我的 RAM 速度的理论最大值）。第二个例子得到了 8GB/s。

在查看带有 cargo asm 的程序集时，第一个示例转换为展开的 SIMD 优化循环：

.LBB11_7:
 sum += *c;
 movdqu  xmm2, xmmword, ptr, [rcx, +, 4*rax]
 paddd   xmm2, xmm0
 movdqu  xmm0, xmmword, ptr, [rcx, +, 4*rax, +, 16]
 paddd   xmm0, xmm1
 movdqu  xmm1, xmmword, ptr, [rcx, +, 4*rax, +, 32]
 movdqu  xmm3, xmmword, ptr, [rcx, +, 4*rax, +, 48]
 movdqu  xmm4, xmmword, ptr, [rcx, +, 4*rax, +, 64]
 paddd   xmm4, xmm1
 paddd   xmm4, xmm2
 movdqu  xmm2, xmmword, ptr, [rcx, +, 4*rax, +, 80]
 paddd   xmm2, xmm3
 paddd   xmm2, xmm0
 movdqu  xmm0, xmmword, ptr, [rcx, +, 4*rax, +, 96]
 paddd   xmm0, xmm4
 movdqu  xmm1, xmmword, ptr, [rcx, +, 4*rax, +, 112]
 paddd   xmm1, xmm2
 add     rax, 32
 add     r11, -4
 jne     .LBB11_7
.LBB11_8:
 test    r10, r10
 je      .LBB11_11
 lea     r11, [rcx, +, 4*rax]
 add     r11, 16
 shl     r10, 5
 xor     eax, eax

第二个例子没有展开任何循环，甚至没有内联代码 _mm256_add_epi32:

...
movaps  xmmword, ptr, [rbp, +, 320], xmm7
 movaps  xmmword, ptr, [rbp, +, 304], xmm6
 and     rsp, -32
 mov     r12, rdx
 mov     rdi, rcx
 lea     rcx, [rsp, +, 32]
 let mut curr_sum = _mm256_setzero_si256();
 call    core::core_arch::x86::avx::_mm256_setzero_si256
 movaps  xmm6, xmmword, ptr, [rsp, +, 32]
 movaps  xmm7, xmmword, ptr, [rsp, +, 48]
 while offset < total {
 test    r12, r12
 jle     .LBB13_3
 xor     esi, esi
 lea     rbx, [rsp, +, 384]
 lea     r14, [rsp, +, 64]
 lea     r15, [rsp, +, 96]
.LBB13_2:
 let curr = _mm256_load_epi32(p_src.offset(offset));
 mov     rcx, rbx
 mov     rdx, rdi
 call    core::core_arch::x86::avx512f::_mm256_load_epi32
 curr_sum = _mm256_add_epi32(curr_sum, curr);
 movaps  xmmword, ptr, [rsp, +, 112], xmm7
 movaps  xmmword, ptr, [rsp, +, 96], xmm6
 mov     rcx, r14
 mov     rdx, r15
 mov     r8, rbx
 call    core::core_arch::x86::avx2::_mm256_add_epi32
 movaps  xmm6, xmmword, ptr, [rsp, +, 64]
 movaps  xmm7, xmmword, ptr, [rsp, +, 80]
 offset += 8;
 add     rsi, 8
 while offset < total {
 add     rdi, 32
 cmp     rsi, r12
...

这当然是非常简单的示例，我不打算使用手工制作的 SIMD 进行简单求和。但它仍然让我感到困惑，为什么显式 SIMD 如此缓慢以及为什么使用 SIMD 内在函数会导致这种未优化的代码。

Answer 1

看来你忘了告诉 rustc 它被允许在任何地方使用 AVX2 指令，所以它不能内联这些函数。相反，你得到了一个彻底的灾难，只有包装函数被编译为使用 AVX2 的函数，或类似的东西。

我用 -O -C target-cpu=skylake-avx512 (https://godbolt.org/z/csY5or43T) 效果很好，所以它甚至可以内联您使用的 AVX512VL 负载，_mm256_load_epi32¹，然后在紧密循环内将其优化为 vpaddd ymm0, ymm0, ymmword ptr [rdi + 4*rax] (AVX2) 的内存源操作数。

在 GCC / clang 中，在这种情况下，您会收到类似“在调用 always_inline foobar 时内联失败”这样的错误，而不是工作但速度慢的 asm。（参见）。这是 Rust 可能应该在黄金时间准备就绪之前解决的问题，要么像 MSVC 那样使用内在函数将指令实际内联到函数中，要么拒绝像 GCC/clang.

那样编译

脚注 1： 如果您不想使用 AVX512，请参阅。

使用 -O -C target-cpu=skylake（仅 AVX2），它内联所有其他内容，包括 vpaddd ymm，但仍然调用一个函数，使用 AVX vmovaps 将 32 个字节从内存复制到内存.它需要 AVX512VL 来内联内部函数，但在优化过程的后期，它意识到在没有屏蔽的情况下，它只是一个 256 位加载，它应该在没有臃肿的 AVX-512 指令的情况下完成。英特尔甚至提供了需要 AVX-512 的 no-masking 版本的 _mm256_mask[z]_loadu_epi32，这有点愚蠢。或愚蠢 gcc/clang/rustc 认为它是 AVX512 内在的。

SIMD 性能缓慢 - 无内联

Slow SIMD performance - no inlining

sse

simd

rust

avx2