rust - 缓慢的 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 中，在这种情况下，您会收到类似“inlining failed in call to always_inline foobar”的错误，而不是工作但 asm 速度慢。 (见 this for details)。这可能是 Rust 在准备好黄金时间之前应该解决的问题，要么像 MSVC 那样使用内在函数将指令实际内联到函数中，要么拒绝像 GCC/clang 那样编译。

脚注 1:见 How to emulate _mm256_loadu_epi32 with gcc or clang?如果您不是要使用 AVX512。

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