Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Support vec zero-alloc optimization for tuples and byte arrays #97581

Merged
merged 1 commit into from
Jul 25, 2022
Merged

Support vec zero-alloc optimization for tuples and byte arrays #97581

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 32 additions & 4 deletions library/alloc/src/vec/is_zero.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@ macro_rules! impl_is_zero {
};
}

impl_is_zero!(i8, |x| x == 0); // It is needed to impl for arrays and tuples of i8.
impl_is_zero!(i16, |x| x == 0);
impl_is_zero!(i32, |x| x == 0);
impl_is_zero!(i64, |x| x == 0);
impl_is_zero!(i128, |x| x == 0);
impl_is_zero!(isize, |x| x == 0);

impl_is_zero!(u8, |x| x == 0); // It is needed to impl for arrays and tuples of u8.
impl_is_zero!(u16, |x| x == 0);
impl_is_zero!(u32, |x| x == 0);
impl_is_zero!(u64, |x| x == 0);
Expand Down Expand Up @@ -54,15 +56,41 @@ unsafe impl<T: IsZero, const N: usize> IsZero for [T; N] {
fn is_zero(&self) -> bool {
// Because this is generated as a runtime check, it's not obvious that
// it's worth doing if the array is really long. The threshold here
// is largely arbitrary, but was picked because as of 2022-05-01 LLVM
// can const-fold the check in `vec![[0; 32]; n]` but not in
// `vec![[0; 64]; n]`: https://godbolt.org/z/WTzjzfs5b
// is largely arbitrary, but was picked because as of 2022-07-01 LLVM
// fails to const-fold the check in `vec![[1; 32]; n]`
// See https://github.com/rust-lang/rust/pull/97581#issuecomment-1166628022
// Feel free to tweak if you have better evidence.

N <= 32 && self.iter().all(IsZero::is_zero)
N <= 16 && self.iter().all(IsZero::is_zero)
}
}

// This is recursive macro.
macro_rules! impl_for_tuples {
// Stopper
() => {
// No use for implementing for empty tuple because it is ZST.
};
($first_arg:ident $(,$rest:ident)*) => {
unsafe impl <$first_arg: IsZero, $($rest: IsZero,)*> IsZero for ($first_arg, $($rest,)*){
#[inline]
fn is_zero(&self) -> bool{
// Destructure tuple to N references
// Rust allows to hide generic params by local variable names.
#[allow(non_snake_case)]
let ($first_arg, $($rest,)*) = self;

$first_arg.is_zero()
$( && $rest.is_zero() )*
}
}

impl_for_tuples!($($rest),*);
}
}

impl_for_tuples!(A, B, C, D, E, F, G, H);

// `Option<&T>` and `Option<Box<T>>` are guaranteed to represent `None` as null.
// For fat pointers, the bytes that would be the pointer metadata in the `Some`
// variant are padding in the `None` variant, so ignoring them and
Expand Down
27 changes: 14 additions & 13 deletions library/alloc/src/vec/spec_from_elem.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use core::ptr;
Mark-Simulacrum marked this conversation as resolved.
Show resolved Hide resolved

use crate::alloc::Allocator;
use crate::raw_vec::RawVec;
use core::ptr::{self};

use super::{ExtendElement, IsZero, Vec};

Expand All @@ -17,6 +18,18 @@ impl<T: Clone> SpecFromElem for T {
}
}

impl<T: Clone + IsZero> SpecFromElem for T {
#[inline]
default fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
Mark-Simulacrum marked this conversation as resolved.
Show resolved Hide resolved
if elem.is_zero() {
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
}
let mut v = Vec::with_capacity_in(n, alloc);
v.extend_with(n, ExtendElement(elem));
v
}
}

impl SpecFromElem for i8 {
Mark-Simulacrum marked this conversation as resolved.
Show resolved Hide resolved
#[inline]
fn from_elem<A: Allocator>(elem: i8, n: usize, alloc: A) -> Vec<i8, A> {
Expand Down Expand Up @@ -46,15 +59,3 @@ impl SpecFromElem for u8 {
}
}
}

impl<T: Clone + IsZero> SpecFromElem for T {
#[inline]
fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
if elem.is_zero() {
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
}
let mut v = Vec::with_capacity_in(n, alloc);
v.extend_with(n, ExtendElement(elem));
v
}
}
136 changes: 124 additions & 12 deletions src/test/codegen/vec-calloc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,29 +4,141 @@

#![crate_type = "lib"]

// CHECK-LABEL: @vec_zero_bytes
#[no_mangle]
pub fn vec_zero_bytes(n: usize) -> Vec<u8> {
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(
// CHECK-NOT: call {{.*}}llvm.memset

// CHECK: call {{.*}}__rust_alloc_zeroed(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(
// CHECK-NOT: call {{.*}}llvm.memset

// CHECK: ret void
vec![0; n]
}

// CHECK-LABEL: @vec_one_bytes
#[no_mangle]
pub fn vec_one_bytes(n: usize) -> Vec<u8> {
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(

// CHECK: call {{.*}}__rust_alloc(
// CHECK: call {{.*}}llvm.memset

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(

// CHECK: ret void
vec![1; n]
}

// CHECK-LABEL: @vec_zero_scalar
#[no_mangle]
pub fn vec_zero_scalar(n: usize) -> Vec<i32> {
// CHECK-NOT: __rust_alloc(
// CHECK: __rust_alloc_zeroed(
// CHECK-NOT: __rust_alloc(
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: call {{.*}}__rust_alloc_zeroed(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: ret void
vec![0; n]
}

// CHECK-LABEL: @vec_one_scalar
#[no_mangle]
pub fn vec_one_scalar(n: usize) -> Vec<i32> {
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(

// CHECK: call {{.*}}__rust_alloc(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(

// CHECK: ret void
vec![1; n]
}

// CHECK-LABEL: @vec_zero_rgb48
#[no_mangle]
pub fn vec_zero_rgb48(n: usize) -> Vec<[u16; 3]> {
// CHECK-NOT: __rust_alloc(
// CHECK: __rust_alloc_zeroed(
// CHECK-NOT: __rust_alloc(
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: call {{.*}}__rust_alloc_zeroed(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: ret void
vec![[0, 0, 0]; n]
}

// CHECK-LABEL: @vec_zero_array_32
// CHECK-LABEL: @vec_zero_array_16
#[no_mangle]
pub fn vec_zero_array_16(n: usize) -> Vec<[i64; 16]> {
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: call {{.*}}__rust_alloc_zeroed(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: ret void
vec![[0_i64; 16]; n]
}

// CHECK-LABEL: @vec_zero_tuple
#[no_mangle]
pub fn vec_zero_tuple(n: usize) -> Vec<(i16, u8, char)> {
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: call {{.*}}__rust_alloc_zeroed(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc(

// CHECK: ret void
vec![(0, 0, '\0'); n]
}

// CHECK-LABEL: @vec_non_zero_tuple
#[no_mangle]
pub fn vec_zero_array_32(n: usize) -> Vec<[i64; 32]> {
// CHECK-NOT: __rust_alloc(
// CHECK: __rust_alloc_zeroed(
// CHECK-NOT: __rust_alloc(
vec![[0_i64; 32]; n]
pub fn vec_non_zero_tuple(n: usize) -> Vec<(i16, u8, char)> {
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(

// CHECK: call {{.*}}__rust_alloc(

// CHECK-NOT: call {{.*}}alloc::vec::from_elem
// CHECK-NOT: call {{.*}}reserve
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(

// CHECK: ret void
vec![(0, 0, 'A'); n]
}