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.

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

Add SIMD benches #1553

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 5 additions & 0 deletions benches/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,11 @@ version = "0.1.0"
edition = "2021"
publish = false

[features]
# Option (requires nightly Rust): experimental SIMD support
simd_support = ["rand/simd_support"]


[dependencies]

[dev-dependencies]
Expand Down
98 changes: 73 additions & 25 deletions benches/benches/uniform.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,66 +8,114 @@

//! Implement benchmarks for uniform distributions over integer types

#![cfg_attr(feature = "simd_support", feature(portable_simd))]

use core::time::Duration;
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use rand::distr::uniform::{SampleRange, Uniform};
use rand::prelude::*;
use rand_chacha::ChaCha8Rng;
use rand_pcg::{Pcg32, Pcg64};
#[cfg(feature = "simd_support")]
use std::simd::{num::SimdUint, Simd};

const WARM_UP_TIME: Duration = Duration::from_millis(1000);
const MEASUREMENT_TIME: Duration = Duration::from_secs(3);
const SAMPLE_SIZE: usize = 100_000;
const N_RESAMPLES: usize = 10_000;

macro_rules! sample {
($R:ty, $T:ty, $U:ty, $g:expr) => {
(@range $T:ty, $U:ty, 1, $rng:ident) => {{
assert_eq!(<$T>::BITS, <$U>::BITS);
let bits = (<$T>::BITS / 2);
let mask = (1 as $U).wrapping_neg() >> bits;
let x = $rng.random::<$U>();
((x >> bits) * (x & mask)) as $T
}};

(@range $T:ty, $U:ty, $len:tt, $rng:ident) => {{
let bits = (<$T>::BITS / 2);
let mask = Simd::splat((1 as $U).wrapping_neg() >> bits);
let bits = Simd::splat(bits as $U);
let x = $rng.random::<Simd<$U, $len>>();
((x >> bits) * (x & mask)).cast()
}};

(@MIN $T:ty, 1) => {
<$T>::MIN
};

(@MIN $T:ty, $len:tt) => {
Simd::<$T, $len>::splat(<$T>::MIN)
};

(@wrapping_add $lhs:expr, $rhs:expr, 1) => {
$lhs.wrapping_add($rhs)
};

(@wrapping_add $lhs:expr, $rhs:expr, $len:tt) => {
($lhs + $rhs)
};

($R:ty, $T:ty, $U:ty, $len:tt, $g:expr) => {
$g.bench_function(BenchmarkId::new(stringify!($R), "single"), |b| {
let mut rng = <$R>::from_rng(&mut rand::rng());
let x = rng.random::<$U>();
let bits = (<$T>::BITS / 2);
let mask = (1 as $U).wrapping_neg() >> bits;
let range = (x >> bits) * (x & mask);
let low = <$T>::MIN;
let high = low.wrapping_add(range as $T);
let range = sample!(@range $T, $U, $len, rng);
let low = sample!(@MIN $T, $len);
let high = sample!(@wrapping_add low, range, $len);

b.iter(|| (low..=high).sample_single(&mut rng));
});

$g.bench_function(BenchmarkId::new(stringify!($R), "distr"), |b| {
let mut rng = <$R>::from_rng(&mut rand::rng());
let x = rng.random::<$U>();
let bits = (<$T>::BITS / 2);
let mask = (1 as $U).wrapping_neg() >> bits;
let range = (x >> bits) * (x & mask);
let low = <$T>::MIN;
let high = low.wrapping_add(range as $T);
let dist = Uniform::<$T>::new_inclusive(<$T>::MIN, high).unwrap();
let range = sample!(@range $T, $U, $len, rng);
let low = sample!(@MIN $T, $len);
let high = sample!(@wrapping_add low, range, $len);
let dist = Uniform::new_inclusive(low, high).unwrap();

b.iter(|| dist.sample(&mut rng));
});
};

($c:expr, $T:ty, $U:ty) => {{
let mut g = $c.benchmark_group(concat!("sample", stringify!($T)));
// Entrypoint:
// $T is the output type (integer)
// $U is the unsigned version of the output type
// $len is the width for SIMD or 1 for non-SIMD
($c:expr, $T:ty, $U:ty, $len:tt) => {{
let mut g = $c.benchmark_group(concat!("sample_", stringify!($T), "x", stringify!($len)));
g.sample_size(SAMPLE_SIZE);
g.warm_up_time(WARM_UP_TIME);
g.measurement_time(MEASUREMENT_TIME);
g.nresamples(N_RESAMPLES);
sample!(SmallRng, $T, $U, g);
sample!(ChaCha8Rng, $T, $U, g);
sample!(Pcg32, $T, $U, g);
sample!(Pcg64, $T, $U, g);
sample!(SmallRng, $T, $U, $len, g);
sample!(ChaCha8Rng, $T, $U, $len, g);
sample!(Pcg32, $T, $U, $len, g);
sample!(Pcg64, $T, $U, $len, g);
g.finish();
}};
}

fn sample(c: &mut Criterion) {
sample!(c, i8, u8);
sample!(c, i16, u16);
sample!(c, i32, u32);
sample!(c, i64, u64);
sample!(c, i128, u128);
sample!(c, i8, u8, 1);
sample!(c, i16, u16, 1);
sample!(c, i32, u32, 1);
sample!(c, i64, u64, 1);
sample!(c, i128, u128, 1);
#[cfg(feature = "simd_support")]
sample!(c, u8, u8, 8);
#[cfg(feature = "simd_support")]
sample!(c, u8, u8, 16);
#[cfg(feature = "simd_support")]
sample!(c, u8, u8, 32);
#[cfg(feature = "simd_support")]
sample!(c, u8, u8, 64);
#[cfg(feature = "simd_support")]
sample!(c, i16, u16, 8);
#[cfg(feature = "simd_support")]
sample!(c, i16, u16, 16);
#[cfg(feature = "simd_support")]
sample!(c, i16, u16, 32);
}

criterion_group! {
Expand Down
Loading