7

我编写了这个非常简单的 Rust 函数:

fn iterate(nums: &Box<[i32]>) -> i32 {
    let mut total = 0;
    let len = nums.len();
    for i in 0..len {
        if nums[i] > 0 {
            total += nums[i];
        } else {
            total -= nums[i];
        }
    }

    total
}

我编写了一个基本基准,它使用有序数组和随机数组调用该方法:

fn criterion_benchmark(c: &mut Criterion) {
    const SIZE: i32 = 1024 * 1024;

    let mut group = c.benchmark_group("Branch Prediction");

    // setup benchmarking for an ordered array
    let mut ordered_nums: Vec<i32> = vec![];
    for i in 0..SIZE {
        ordered_nums.push(i - SIZE/2);
    }
    let ordered_nums = ordered_nums.into_boxed_slice();
    group.bench_function("ordered", |b| b.iter(|| iterate(&ordered_nums)));

    // setup benchmarking for a shuffled array
    let mut shuffled_nums: Vec<i32> = vec![];
    for i in 0..SIZE {
        shuffled_nums.push(i - SIZE/2);
    }
    let mut rng = thread_rng();
    let mut shuffled_nums = shuffled_nums.into_boxed_slice();
    shuffled_nums.shuffle(&mut rng);
    group.bench_function("shuffled", |b| b.iter(|| iterate(&shuffled_nums)));

    group.finish();
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);

令我惊讶的是,这两个基准测试具有几乎完全相同的运行时间,而 Java 中的类似基准测试显示两者之间存在明显差异,这可能是由于在 shuffled 情况下分支预测失败所致。

我见过提到条件移动指令,但如果我otool -tv是可执行文件(我在 Mac 上运行),我在iterate方法输出中看不到任何内容。

谁能阐明为什么 Rust 中的有序和无序案例之间没有明显的性能差异?

4

1 回答 1

11

摘要:LLVM 能够通过使用cmov指令或 SIMD 指令的非常巧妙的组合来删除/隐藏分支。


我使用 Godbolt查看完整的程序集(带有-C opt-level=3)。我将在下面解释组件的重要部分。

它是这样开始的:

        mov     r9, qword ptr [rdi + 8]         ; r9 = nums.len()
        test    r9, r9                          ; if len == 0
        je      .LBB0_1                         ;     goto LBB0_1
        mov     rdx, qword ptr [rdi]            ; rdx = base pointer (first element)
        cmp     r9, 7                           ; if len > 7
        ja      .LBB0_5                         ;     goto LBB0_5
        xor     eax, eax                        ; eax = 0
        xor     esi, esi                        ; esi = 0
        jmp     .LBB0_4                         ; goto LBB0_4

.LBB0_1:
        xor     eax, eax                        ; return 0
        ret

在这里,函数区分 3 种不同的“状态”:

  • 切片为空→立即返回 0
  • 切片长度 ≤ 7 → 使用标准顺序算法 ( LBB0_4)
  • 切片长度 > 7 → 使用 SIMD 算法 ( LBB0_5)

那么让我们来看看这两种不同的算法吧!


标准顺序算法

请记住,rsi( esi) 和rax( eax) 设置为 0,这rdx是指向数据的基指针。

.LBB0_4:
        mov     ecx, dword ptr [rdx + 4*rsi]    ; ecx = nums[rsi]
        add     rsi, 1                          ; rsi += 1
        mov     edi, ecx                        ; edi = ecx
        neg     edi                             ; edi = -edi
        cmovl   edi, ecx                        ; if ecx >= 0 { edi = ecx }
        add     eax, edi                        ; eax += edi
        cmp     r9, rsi                         ; if rsi != len
        jne     .LBB0_4                         ;     goto LBB0_4
        ret                                     ; return eax

这是一个简单的循环遍历所有元素num。不过,在循环体中有一个小技巧:从原始元素ecx中,一个否定值存储在edi. 如果原始值为正,则使用cmovl,edi被原始值覆盖。这意味着它总是会变成正数(即包含原始元素的绝对值)。然后将其添加到(最后返回)。edieax

所以你的if分支隐藏在cmov指令中。正如您在此基准测试中所见,执行cmov指令所需的时间与条件的概率无关。这是一个非常了不起的指令!


SIMD算法

SIMD 版本包含很多指令,我不会在这里完全粘贴。主循环一次处理 16 个整数!

        movdqu  xmm5, xmmword ptr [rdx + 4*rdi]
        movdqu  xmm3, xmmword ptr [rdx + 4*rdi + 16]
        movdqu  xmm0, xmmword ptr [rdx + 4*rdi + 32]
        movdqu  xmm1, xmmword ptr [rdx + 4*rdi + 48]

它们从内存加载到寄存器xmm0xmm1和。这些寄存器中的每一个都包含四个 32 位值,但为了更容易理解,想象一下每个寄存器只包含一个值。以下所有指令分别对这些 SIMD 寄存器的每个值进行操作,因此心智模型很好!我在下面的解释听起来好像寄存器只包含一个值。xmm3xmm5xmm

主要技巧现在在以下说明中(其中句柄xmm5):

        movdqa  xmm6, xmm5      ; xmm6 = xmm5 (make a copy)
        psrad   xmm6, 31        ; logical right shift 31 bits (see below)
        paddd   xmm5, xmm6      ; xmm5 += xmm6
        pxor    xmm5, xmm6      ; xmm5 ^= xmm6

逻辑右移用符号位的值填充“空的高位”(左侧“移入”的那些)。通过移位 31,我们最终在每个位置都只有符号位!所以任何正数都会变成 32 个零,任何负数都会变成 32 个一。现在xmm6也是000...000(如果xmm5是肯定的)或111...111(如果xmm5是否定的)。

接下来这个人工xmm6被添加到xmm5。如果xmm5为正,xmm6则为 0,因此添加它不会改变xmm5。但是,如果xmm5是负数,我们将111...111其相加相当于减 1。最后,我们xmm5与 进行异或xmm6。同样,如果xmm5一开始是积极的,我们异或与000...000which 没有效果。如果xmm5一开始是负数,我们与 异111...111或,这意味着我们翻转所有位。所以对于这两种情况:

  • 如果元素是正数,我们什么都不做(addxor没有任何效果)
  • 如果元素是负数,我们减 1 并翻转所有位。这是一个二进制补码否定!

因此,通过这 4 条指令,我们计算了xmm5! 再次,由于这种摆弄技巧,没有分支。请记住,它xmm5实际上包含 4 个整数,所以它非常快!

该绝对值现在被添加到累加器中,并且对xmm包含来自切片的值的其他三个寄存器执行相同操作。(我们不会详细讨论剩余的代码。)


带 AVX2 的 SIMD

如果我们允许 LLVM 发出 AVX2 指令(通过-C target-feature=+avx2),它甚至可以使用该pabsd指令而不是四个“hacky”指令:

vpabsd  ymm2, ymmword ptr [rdx + 4*rdi]

它直接从内存中加载值,计算绝对值并将其存储ymm2在一条指令中!请记住,ymm寄存器是寄存器的两倍xmm(适合 8 个 32 位值)!

于 2020-01-04T10:39:46.077 回答