Skip to content

Commit

Permalink
Add assumes to avoid bounds checks
Browse files Browse the repository at this point in the history
Surprisingly the compiler removed these bounds checks automatically when
targeting x86_64, but not on aarch64.
  • Loading branch information
jhorstmann committed May 9, 2024
1 parent 5038a49 commit 2a52a46
Showing 1 changed file with 14 additions and 2 deletions.
16 changes: 14 additions & 2 deletions library/alloc/src/str.rs
Original file line number Diff line number Diff line change
Expand Up @@ -623,9 +623,15 @@ fn convert_while_ascii(s: &str, convert: fn(&u8) -> u8) -> (String, &str) {
let mut i = 0_usize;

// process the input in chunks to enable auto-vectorization
let mut is_ascii = [false; N];
while slice.len() >= N {
// Safety: out_slice was allocated with same lengths as input slice and gets updated with
// the same offsets
unsafe {
core::intrinsics::assume(slice.len() == out_slice.len());
}

let chunk = &slice[..N];
let mut is_ascii = [false; N];

for j in 0..N {
is_ascii[j] = chunk[j] <= 127;
Expand All @@ -634,7 +640,7 @@ fn convert_while_ascii(s: &str, convert: fn(&u8) -> u8) -> (String, &str) {
// auto-vectorization for this check is a bit fragile,
// sum and comparing against the chunk size gives the best result,
// specifically a pmovmsk instruction on x86.
if is_ascii.into_iter().map(|x| x as u8).sum::<u8>() as usize != N {
if is_ascii.iter().map(|x| *x as u8).sum::<u8>() as usize != N {
break;
}

Expand All @@ -649,6 +655,12 @@ fn convert_while_ascii(s: &str, convert: fn(&u8) -> u8) -> (String, &str) {

// handle the remainder as individual bytes
while !slice.is_empty() {
// Safety: out_slice was allocated with same lengths as input slice and gets updated with
// the same offsets
unsafe {
core::intrinsics::assume(slice.len() == out_slice.len());
}

let byte = slice[0];
if byte > 127 {
break;
Expand Down

0 comments on commit 2a52a46

Please sign in to comment.