|
8 | 8 |
|
9 | 9 | //! Implement benchmarks for uniform distributions over integer types
|
10 | 10 |
|
| 11 | +#![cfg_attr(feature = "simd_support", feature(portable_simd))] |
| 12 | + |
11 | 13 | use core::time::Duration;
|
12 | 14 | use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
|
13 | 15 | use rand::distr::uniform::{SampleRange, Uniform};
|
14 | 16 | use rand::prelude::*;
|
15 | 17 | use rand_chacha::ChaCha8Rng;
|
16 | 18 | use rand_pcg::{Pcg32, Pcg64};
|
| 19 | +#[cfg(feature = "simd_support")] |
| 20 | +use std::simd::{num::SimdUint, Simd}; |
17 | 21 |
|
18 | 22 | const WARM_UP_TIME: Duration = Duration::from_millis(1000);
|
19 | 23 | const MEASUREMENT_TIME: Duration = Duration::from_secs(3);
|
20 | 24 | const SAMPLE_SIZE: usize = 100_000;
|
21 | 25 | const N_RESAMPLES: usize = 10_000;
|
22 | 26 |
|
23 | 27 | 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) => { |
25 | 61 | $g.bench_function(BenchmarkId::new(stringify!($R), "single"), |b| {
|
26 | 62 | 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); |
33 | 66 |
|
34 | 67 | b.iter(|| (low..=high).sample_single(&mut rng));
|
35 | 68 | });
|
36 | 69 |
|
37 | 70 | $g.bench_function(BenchmarkId::new(stringify!($R), "distr"), |b| {
|
38 | 71 | 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(); |
46 | 76 |
|
47 | 77 | b.iter(|| dist.sample(&mut rng));
|
48 | 78 | });
|
49 | 79 | };
|
50 | 80 |
|
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))); |
53 | 87 | g.sample_size(SAMPLE_SIZE);
|
54 | 88 | g.warm_up_time(WARM_UP_TIME);
|
55 | 89 | g.measurement_time(MEASUREMENT_TIME);
|
56 | 90 | 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); |
61 | 95 | g.finish();
|
62 | 96 | }};
|
63 | 97 | }
|
64 | 98 |
|
65 | 99 | 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); |
71 | 119 | }
|
72 | 120 |
|
73 | 121 | criterion_group! {
|
|
0 commit comments