考虑下面的例子来计算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
}

当我try 对代码进行基准测试时,第一个示例得到了约23GB/s(这接近我的RAM速度的理论最大值).第二个例子是8GB/s.

在查看带Cargo 的部件时,第一个示例转化为展开的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内部函数会导致如此未优化的代码.

推荐答案

It appears you forgot to tell rustc it was allowed to use AVX2 instructions everywhere, so it couldn't inline those functions.相反,你会得到一个彻底的灾难,只有包装器函数被编译为AVX2使用函数,或类似的东西.

对我来说,-O -C target-cpu=skylake-avx512(https://godbolt.org/z/csY5or43T)运行良好,因此它甚至可以内联您使用的AVX512VL加载,_mm256_load_epi321,然后在紧循环中优化为vpaddd ymm0, ymm0, ymmword ptr [rdi + 4*rax](AVX2)的内存源操作数.

在GCC/clang中,在本例中会出现类似"调用always_inline foobar时内联失败"的错误,而不是正常工作但速度较慢的asm.(见this for details).这可能是Rust在准备好进入黄金时段之前应该解决的问题,要么像MSVC一样,实际使用内在函数将指令内联到函数中,要么拒绝像GCC/clang那样编译.

Footnote 1:

有了-O -C target-cpu=skylake(只有AVX2),它将其他所有内容都内联起来,包括vpaddd ymm,但仍然调用一个函数,该函数使用AVX vmovaps将32字节从内存复制到内存.它需要AVX512VL来内联本机,但后来在优化过程中,它意识到,如果没有屏蔽,它只需要256位加载,而不需要臃肿的AVX-512指令.有点愚蠢的是,英特尔甚至提供了_mm256_mask[z]_loadu_epi32的无屏蔽版本,需要AVX-512.或者哑巴认为GCC/CLAN/RUSTC认为它是AVX512内在的.

Rust相关问答推荐

关于Rust 中回归的逻辑

为什么这是&q;,而让&q;循环是无限循环?

当T不执行Copy时,如何返回Arc Mutex T后面的值?

对于已经被认为是未定义行为的相同数据,纯粹存在`&;[u32]`和`&;mut[u32]`吗?

获取已知数量的输入

这是什么:`impl Trait for T {}`?

全面的 Rust Ch.16.2 - 使用捕获和 const 表达式的 struct 模式匹配

pyO3 和 Panics

有什么方法可以通过使用生命周期来减轻嵌套生成器中的当生成器产生时borrow 可能仍在使用错误?

是否可以通过可变引用推进可变切片?

Rust 生命周期:这两种类型声明为不同的生命周期

使用 `clap` 在 Rust CLI 工具中设置布尔标志

仅在运行测试时生成调试输出

为什么在 macOS / iOS 上切换 WiFi 网络时 reqwest 响应会挂起?

你能告诉我如何在 Rust 中使用定时器吗?

如何异步记忆选项中的 struct 字段

如何为枚举中的单个或多个值返回迭代器

在 Rust 中枚举字符串的最佳方式? (字符()与 as_bytes())

如何制作具有关联类型的特征的类型擦除版本?

如果返回类型是通用的,我可以返回 &str 输入的一部分吗?