Skip to content

Commit 7aa2637

Browse files
authored
Add SIMD benches (#1553)
Add benchmarks for uniform distribution / single-sample variates for `Simd` types: `u8x8, u8x16, u8x32, u8x64, i16x8, i16x16, i16x32`. Also, change the name of the non-SIMD benchmarks to "x1" e.g. `sample_i16x1/SmallRng/distr`.
2 parents d3dd415 + 5ec6390 commit 7aa2637

File tree

2 files changed

+78
-25
lines changed

2 files changed

+78
-25
lines changed

benches/Cargo.toml

+5
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@ version = "0.1.0"
44
edition = "2021"
55
publish = false
66

7+
[features]
8+
# Option (requires nightly Rust): experimental SIMD support
9+
simd_support = ["rand/simd_support"]
10+
11+
712
[dependencies]
813

914
[dev-dependencies]

benches/benches/uniform.rs

+73-25
Original file line numberDiff line numberDiff line change
@@ -8,66 +8,114 @@
88

99
//! Implement benchmarks for uniform distributions over integer types
1010
11+
#![cfg_attr(feature = "simd_support", feature(portable_simd))]
12+
1113
use core::time::Duration;
1214
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
1315
use rand::distr::uniform::{SampleRange, Uniform};
1416
use rand::prelude::*;
1517
use rand_chacha::ChaCha8Rng;
1618
use rand_pcg::{Pcg32, Pcg64};
19+
#[cfg(feature = "simd_support")]
20+
use std::simd::{num::SimdUint, Simd};
1721

1822
const WARM_UP_TIME: Duration = Duration::from_millis(1000);
1923
const MEASUREMENT_TIME: Duration = Duration::from_secs(3);
2024
const SAMPLE_SIZE: usize = 100_000;
2125
const N_RESAMPLES: usize = 10_000;
2226

2327
macro_rules! sample {
24-
($R:ty, $T:ty, $U:ty, $g:expr) => {
28+
(@range $T:ty, $U:ty, 1, $rng:ident) => {{
29+
assert_eq!(<$T>::BITS, <$U>::BITS);
30+
let bits = (<$T>::BITS / 2);
31+
let mask = (1 as $U).wrapping_neg() >> bits;
32+
let x = $rng.random::<$U>();
33+
((x >> bits) * (x & mask)) as $T
34+
}};
35+
36+
(@range $T:ty, $U:ty, $len:tt, $rng:ident) => {{
37+
let bits = (<$T>::BITS / 2);
38+
let mask = Simd::splat((1 as $U).wrapping_neg() >> bits);
39+
let bits = Simd::splat(bits as $U);
40+
let x = $rng.random::<Simd<$U, $len>>();
41+
((x >> bits) * (x & mask)).cast()
42+
}};
43+
44+
(@MIN $T:ty, 1) => {
45+
<$T>::MIN
46+
};
47+
48+
(@MIN $T:ty, $len:tt) => {
49+
Simd::<$T, $len>::splat(<$T>::MIN)
50+
};
51+
52+
(@wrapping_add $lhs:expr, $rhs:expr, 1) => {
53+
$lhs.wrapping_add($rhs)
54+
};
55+
56+
(@wrapping_add $lhs:expr, $rhs:expr, $len:tt) => {
57+
($lhs + $rhs)
58+
};
59+
60+
($R:ty, $T:ty, $U:ty, $len:tt, $g:expr) => {
2561
$g.bench_function(BenchmarkId::new(stringify!($R), "single"), |b| {
2662
let mut rng = <$R>::from_rng(&mut rand::rng());
27-
let x = rng.random::<$U>();
28-
let bits = (<$T>::BITS / 2);
29-
let mask = (1 as $U).wrapping_neg() >> bits;
30-
let range = (x >> bits) * (x & mask);
31-
let low = <$T>::MIN;
32-
let high = low.wrapping_add(range as $T);
63+
let range = sample!(@range $T, $U, $len, rng);
64+
let low = sample!(@MIN $T, $len);
65+
let high = sample!(@wrapping_add low, range, $len);
3366

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

3770
$g.bench_function(BenchmarkId::new(stringify!($R), "distr"), |b| {
3871
let mut rng = <$R>::from_rng(&mut rand::rng());
39-
let x = rng.random::<$U>();
40-
let bits = (<$T>::BITS / 2);
41-
let mask = (1 as $U).wrapping_neg() >> bits;
42-
let range = (x >> bits) * (x & mask);
43-
let low = <$T>::MIN;
44-
let high = low.wrapping_add(range as $T);
45-
let dist = Uniform::<$T>::new_inclusive(<$T>::MIN, high).unwrap();
72+
let range = sample!(@range $T, $U, $len, rng);
73+
let low = sample!(@MIN $T, $len);
74+
let high = sample!(@wrapping_add low, range, $len);
75+
let dist = Uniform::new_inclusive(low, high).unwrap();
4676

4777
b.iter(|| dist.sample(&mut rng));
4878
});
4979
};
5080

51-
($c:expr, $T:ty, $U:ty) => {{
52-
let mut g = $c.benchmark_group(concat!("sample", stringify!($T)));
81+
// Entrypoint:
82+
// $T is the output type (integer)
83+
// $U is the unsigned version of the output type
84+
// $len is the width for SIMD or 1 for non-SIMD
85+
($c:expr, $T:ty, $U:ty, $len:tt) => {{
86+
let mut g = $c.benchmark_group(concat!("sample_", stringify!($T), "x", stringify!($len)));
5387
g.sample_size(SAMPLE_SIZE);
5488
g.warm_up_time(WARM_UP_TIME);
5589
g.measurement_time(MEASUREMENT_TIME);
5690
g.nresamples(N_RESAMPLES);
57-
sample!(SmallRng, $T, $U, g);
58-
sample!(ChaCha8Rng, $T, $U, g);
59-
sample!(Pcg32, $T, $U, g);
60-
sample!(Pcg64, $T, $U, g);
91+
sample!(SmallRng, $T, $U, $len, g);
92+
sample!(ChaCha8Rng, $T, $U, $len, g);
93+
sample!(Pcg32, $T, $U, $len, g);
94+
sample!(Pcg64, $T, $U, $len, g);
6195
g.finish();
6296
}};
6397
}
6498

6599
fn sample(c: &mut Criterion) {
66-
sample!(c, i8, u8);
67-
sample!(c, i16, u16);
68-
sample!(c, i32, u32);
69-
sample!(c, i64, u64);
70-
sample!(c, i128, u128);
100+
sample!(c, i8, u8, 1);
101+
sample!(c, i16, u16, 1);
102+
sample!(c, i32, u32, 1);
103+
sample!(c, i64, u64, 1);
104+
sample!(c, i128, u128, 1);
105+
#[cfg(feature = "simd_support")]
106+
sample!(c, u8, u8, 8);
107+
#[cfg(feature = "simd_support")]
108+
sample!(c, u8, u8, 16);
109+
#[cfg(feature = "simd_support")]
110+
sample!(c, u8, u8, 32);
111+
#[cfg(feature = "simd_support")]
112+
sample!(c, u8, u8, 64);
113+
#[cfg(feature = "simd_support")]
114+
sample!(c, i16, u16, 8);
115+
#[cfg(feature = "simd_support")]
116+
sample!(c, i16, u16, 16);
117+
#[cfg(feature = "simd_support")]
118+
sample!(c, i16, u16, 32);
71119
}
72120

73121
criterion_group! {

0 commit comments

Comments
 (0)