Optimize away bounds check in loop indexing into slice, given an assertion

[joshtriplett]

I wrote a simple loop indexing into a slice, to test rustc's ability to optimize away bounds checks if it knows an index is in bounds. Even with this very simple test case, I can't seem to get rust to omit the bounds checks no matter what assert! I add. (I know that I could trivially write this code using iterators instead, but I'm trying to figure out rust's ability to optimize here.)

Test case (edited since original posting to augment the assert! further):

#[no_mangle]
fn f(slice: &[u64], start: usize, end: usize) -> u64 {
    let mut total = 0;
    assert!(start < end && start < slice.len() && end <= slice.len());
    for i in start..end {
        total += slice[i];
    }
    total
}

I put that into the compiler explorer, with -O, and the resulting assembly looks like this:

f:
        push    rax
        cmp     rdx, rcx
        jae     .LBB5_8
        cmp     rsi, rdx
        jbe     .LBB5_8
        cmp     rsi, rcx
        jb      .LBB5_8
        xor     eax, eax
.LBB5_4:
        cmp     rdx, rsi
        jae     .LBB5_7
        add     rax, qword ptr [rdi + 8*rdx]
        add     rdx, 1
        cmp     rcx, rdx
        jne     .LBB5_4
        pop     rcx
        ret
.LBB5_7:
        lea     rax, [rip + .L__unnamed_5]
        mov     rdi, rdx
        mov     rdx, rax
        call    qword ptr [rip + core::panicking::panic_bounds_check@GOTPCREL]
        ud2
.LBB5_8:
        call    std::panicking::begin_panic
        ud2

Based on the x86 calling convention, rdi contains the slice base address, rsi contains the slice length, rdx contains start, and rcx contains end.

So, the first three comparisons verify the assertion and jump to .LBB5_8 if it fails, to panic.

Then inside the loop, there's still another comparison of rdx to rsi, and a jump to .LBB5_7 to panic if out of bounds.

As far as I can tell, that's exactly the same comparison. Shouldn't rustc be able to optimize away that bounds check?

Things I've tested:

• I tried replacing the assert! with an if and unreachable!, or an if and unsafe { std::hint::unreachable_unchecked() }, but in both cases the loop still checked if the index was in bounds on each iteration.
• I tried using -Zmutable-noalias=yes, which didn't help.
• I tried various forms of the assertion condition.

Ideally, rustc should be able to optimize away the bounds check in the loop, based on the assertion. Even better would be if rustc could hoist the bounds check out of the loop even without the assertion, but that seems like a harder problem.

[ecstatic-morse]

Did you mean end < slice.len()?

[joshtriplett]

@ecstatic-morse I tried including that as well, among other permutations of the assert condition, and it didn't change the in-loop bounds check.

[the8472]

Adding std::intrinsics::assume(i < slice.len()) inside the loop eliminates the check and vectorizes it.

[RalfJung]

Looks like LLVM is unable to propagate what it learned on an earlier conditional, into the later loop?

[ecstatic-morse]

@joshtriplett Nevertheless, you should update your example. Without checking end < slice.len(), f(&[0], 0, 1337) will trigger an index-out-of-bounds panic without hitting the assertion.

[joshtriplett]

@ecstatic-morse Done; edited the code and provided the new corresponding assembly. Doesn't affect the code of the loop, which still includes the bounds check.

[the8472]

This works and considering that an exclusive range is used it presumably also is the intended use of that function

#[no_mangle]
fn f(slice: &[u64], start: usize, end: usize) -> u64 {
    let mut total = 0;
    assert!(start < end && end <= slice.len());
    let end = std::cmp::min(slice.len(), end);
    for i in start..end {
        total += slice[i];
    }
    total
}

[joshtriplett]

@the8472 I can confirm that that code eliminates the bounds check. And interestingly, if I reverse the two arguments to min, that does not eliminate the bounds check.

This works (no bounds check, vectorized):

#[no_mangle]
fn f(slice: &[u64], start: usize, end: usize) -> u64 {
    let mut total = 0;
    assert!(start < end && end <= slice.len());
    let end = std::cmp::min(slice.len(), end);
    for i in start..end {
        total += slice[i];
    }
    total
}

This doesn't work (includes the bounds check in the loop):

#[no_mangle]
fn f(slice: &[u64], start: usize, end: usize) -> u64 {
    let mut total = 0;
    assert!(start < end && end <= slice.len());
    let end = std::cmp::min(end, slice.len());
    for i in start..end {
        total += slice[i];
    }
    total
}

[joshtriplett]

The only differences between the two argument orders to min are the direction of the comparison and which argument gets returned if they compare equal. And sure enough, I can confirm that open-coding the equivalent only works if we use slice.len() in place of end in the case where end == slice.len().

This works (no bounds check in the loop):

#[no_mangle]
fn f(slice: &[u64], start: usize, mut end: usize) -> u64 {
    let mut total = 0;
    assert!(start < end && start < slice.len() && end <= slice.len());
    if end >= slice.len() { end = slice.len(); }
    for i in start..end {
        total += slice[i];
    }
    total
}

This doesn't work (bounds check in the loop):

#[no_mangle]
fn f(slice: &[u64], start: usize, mut end: usize) -> u64 {
    let mut total = 0;
    assert!(start < end && start < slice.len() && end <= slice.len());
    if end > slice.len() { end = slice.len(); }
    for i in start..end {
        total += slice[i];
    }
    total
}

[joshtriplett]

It looks like there are two separate bugs here:

1. rustc should optimize away the bounds check given just the assert!, without needing the redundant call to min or the equivalent open-coded change to end.
2. The optimization shouldn't require assigning end = slice.len() in the case where end == slice.len().

[joshtriplett]

I produced a corresponding set of (naively translated) C++ test cases, and confirmed the same behavior from clang trunk.

The following C++ code does not optimize away the bounds check:

#include <cassert>
#include <cstdint>
#include <vector>
using namespace std;

uint64_t f(vector<uint64_t> slice, size_t start, size_t end)
{
    uint64_t total = 0;
    assert(start < end && start < slice.size() && end <= slice.size());
    for (size_t i = start; i < end; i++) {
        total += slice.at(i);
    }
    return total;
}

Nor does this:

#include <cassert>
#include <cstdint>
#include <vector>
using namespace std;

uint64_t f(vector<uint64_t> slice, size_t start, size_t end)
{
    uint64_t total = 0;
    assert(start < end && start < slice.size() && end <= slice.size());
    if (end > slice.size())
        end = slice.size();
    for (size_t i = start; i < end; i++) {
        total += slice.at(i);
    }
    return total;
}

But this does (note the >= in the if):

#include <cassert>
#include <cstdint>
#include <vector>
using namespace std;

uint64_t f(vector<uint64_t> slice, size_t start, size_t end)
{
    uint64_t total = 0;
    assert(start < end && start < slice.size() && end <= slice.size());
    if (end >= slice.size())
        end = slice.size();
    for (size_t i = start; i < end; i++) {
        total += slice.at(i);
    }
    return total;
}