删除 Rust 循环中的边界检查以尝试达到最佳编译器输出

Question

为了确定我是否 can/should 使用 Rust 而不是默认的 C/C++ 我正在研究各种极端情况，主要是考虑到这个问题：在 0.1%在确实重要的情况下，我能否始终获得与 gcc 一样好的编译器输出（具有适当的优化标志）？答案很可能是否定的，但让我们看看...

Reddit 上有一个相当特殊的示例，它研究了无分支排序算法的子例程的编译器输出。

这是基准 C 代码：

#include <stdint.h>
#include <stdlib.h>
int32_t* foo(int32_t* elements, int32_t* buffer, int32_t pivot)
{
    size_t buffer_index = 0;

    for (size_t i = 0; i < 64; ++i) {
        buffer[buffer_index] = (int32_t)i;
        buffer_index += (size_t)(elements[i] < pivot);
    }
}

这里是带有编译器输出的 godbolt link。

Rust 的第一次尝试是这样的：

pub fn foo0(elements: &Vec<i32>, mut buffer: [i32; 64], pivot: i32) -> () {
    let mut buffer_index: usize = 0;
    for i in 0..buffer.len() {
        buffer[buffer_index] = i as i32;
        buffer_index += (elements[i] < pivot) as usize; 
    }
}

正在进行大量边界检查，请参阅 godbolt。

下一次尝试消除第一次边界检查：

pub unsafe fn foo1(elements: &Vec<i32>, mut buffer: [i32; 64], pivot: i32) -> () {
    let mut buffer_index: usize = 0;
    for i in 0..buffer.len() {
        unsafe {
            buffer[buffer_index] = i as i32;
            buffer_index += (elements.get_unchecked(i) < &pivot) as usize; 
        }
    }
}

这样好一点（见同上神栓link）

最后，让我们尝试完全删除边界检查：

use std::ptr;

pub unsafe fn foo2(elements: &Vec<i32>, mut buffer: [i32; 64], pivot: i32) -> () {
    let mut buffer_index: usize = 0;
    unsafe {
        for i in 0..buffer.len() {
            ptr::replace(&mut buffer[buffer_index], i as i32);
            buffer_index += (elements.get_unchecked(i) < &pivot) as usize; 
        }
    }
}

这会产生与 foo1 相同的输出，因此 ptr::replace 仍会执行边界检查。对于那些 unsafe 操作，我肯定超出了我的理解范围。这就引出了我的两个问题：

如何消除边界检查？
像这样分析边缘情况是否有意义？或者，如果提供整个算法而不是其中的一小部分，Rust 编译器是否会看穿所有这些。

关于最后一点，我很好奇，总的来说，Rust 是否可以被屠杀到“字面意思”的程度，即接近金属，就像 C 一样。经验丰富的 Rust 程序员可能会对这一行调查感到畏缩，但这里是...

Answer 1

How can the bounds check be eliminated?

数组，通过它们对切片的解引用强制，也有一个 unchecked form of mutable get。

pub unsafe fn foo(elements: &Vec<i32>, mut buffer: [i32; 64], pivot: i32) {
    let mut buffer_index: usize = 0;
    for i in 0..buffer.len() {
        unsafe {
            *buffer.get_unchecked_mut(buffer_index) = i as i32;
            buffer_index += (elements.get_unchecked(i) < &pivot) as usize; 
        }
    }
}

这可能会产生与使用 Clang 编译等效 C 代码所获得的机器码相同的机器码。 https://godbolt.org/z/ddxP1P

Does it even make sense to analyze edge cases edge cases like this? Or would the Rust compiler see through all this if presented with the whole algorithm instead of only a small fraction thereof.

一如既往，benchmark 任何这些情况，即使您已经确定了那部分代码中的瓶颈。否则，这是一个过早的优化，有一天可能会后悔。特别是在 Rust 中，编写 unsafe 代码的决定不应掉以轻心。可以肯定地说，在许多情况下，仅消除边界检查的工作量和风险就超过了预期的性能优势。

Regarding the last point, I'm curious, in general, whether Rust can be butchered to the point where it is as "literal", i.e. close to the metal, as C is.

不，您不希望这样做有两个主要原因：

尽管 Rust 具有强大的抽象功能，但不为未使用的东西付费的原则仍然非常切题，这与 C++ 类似。参见 what makes an abstraction zero-cost。在边界检查的情况下，这仅仅是语言设计决策的结果，即当编译器无法确保此类访问是内存安全时始终执行空间检查。
C is not that low-level 无论如何。它可能看起来很真实并且接近金属，直到它真的不是。

另请参阅：

Why does my code run slower when I remove bounds checks?

Answer 2

您可以使用老式指针算法来实现。

const N: usize = 64;
pub fn foo2(elements: &Vec<i32>, mut buffer: [i32; N], pivot: i32) -> () {
    assert!(elements.len() >= N);
    let elements = &elements[..N];
    let mut buff_ptr = buffer.as_mut_ptr();
    for (i, &elem) in elements.iter().enumerate(){
        unsafe{
            // SAFETY: We increase ptr strictly less or N times
            *buff_ptr = i as i32;
            if elem < pivot{
                buff_ptr = buff_ptr.add(1);
            }
        }
    }
}

这个版本编译成：

example::foo2:
        push    rax
        cmp     qword ptr [rdi + 16], 64
        jb      .LBB7_4
        mov     r9, qword ptr [rdi]
        lea     r8, [r9 + 256]
        xor     edi, edi

        // Loop goes here
.LBB7_2:
        mov     ecx, dword ptr [r9 + 4*rdi]
        mov     dword ptr [rsi], edi
        lea     rax, [rsi + 4]
        cmp     ecx, edx
        cmovge  rax, rsi
        mov     ecx, dword ptr [r9 + 4*rdi + 4]
        lea     esi, [rdi + 1]
        mov     dword ptr [rax], esi
        lea     rsi, [rax + 4]
        cmp     ecx, edx
        cmovge  rsi, rax
        mov     eax, dword ptr [r9 + 4*rdi + 8]
        lea     ecx, [rdi + 2]
        mov     dword ptr [rsi], ecx
        lea     rcx, [rsi + 4]
        cmp     eax, edx
        cmovge  rcx, rsi
        mov     r10d, dword ptr [r9 + 4*rdi + 12]
        lea     esi, [rdi + 3]
        lea     rax, [r9 + 4*rdi + 16]
        add     rdi, 4
        mov     dword ptr [rcx], esi
        lea     rsi, [rcx + 4]
        cmp     r10d, edx
        cmovge  rsi, rcx
        // Conditional branch to the loop beginning
        cmp     rax, r8
        jne     .LBB7_2
        pop     rax
        ret
.LBB7_4:
        call    std::panicking::begin_panic
        ud2

如你所见，循环展开，单分支为循环迭代跳转。

然而，令我惊讶的是，这个函数并没有被淘汰，因为它没有任何作用：它应该被编译成简单的noop。大概，内联后会变成这样吧。

此外，我要说的是，将参数更改为 &mut 不会更改代码：

example::foo2:
        push    rax
        cmp     qword ptr [rdi + 16], 64
        jb      .LBB7_4
        mov     r9, qword ptr [rdi]
        lea     r8, [r9 + 256]
        xor     edi, edi
.LBB7_2:
        mov     ecx, dword ptr [r9 + 4*rdi]
        mov     dword ptr [rsi], edi
        lea     rax, [rsi + 4]
        cmp     ecx, edx
        cmovge  rax, rsi
        mov     ecx, dword ptr [r9 + 4*rdi + 4]
        lea     esi, [rdi + 1]
        mov     dword ptr [rax], esi
        lea     rsi, [rax + 4]
        cmp     ecx, edx
        cmovge  rsi, rax
        mov     eax, dword ptr [r9 + 4*rdi + 8]
        lea     ecx, [rdi + 2]
        mov     dword ptr [rsi], ecx
        lea     rcx, [rsi + 4]
        cmp     eax, edx
        cmovge  rcx, rsi
        mov     r10d, dword ptr [r9 + 4*rdi + 12]
        lea     esi, [rdi + 3]
        lea     rax, [r9 + 4*rdi + 16]
        add     rdi, 4
        mov     dword ptr [rcx], esi
        lea     rsi, [rcx + 4]
        cmp     r10d, edx
        cmovge  rsi, rcx
        cmp     rax, r8
        jne     .LBB7_2
        pop     rax
        ret
.LBB7_4:
        call    std::panicking::begin_panic
        ud2

不幸的是，rustc 可能发出该函数接受缓冲区参数作为 LLVM IR 中的指针。

删除 Rust 循环中的边界检查以尝试达到最佳编译器输出

Remove bounds checking in Rust loop in attempt to reach optimal compiler output

loops

rust

bounds-check-elimination