Skip to content

Commit 2f27839

Browse files
dancixxdancixx
committed
feat: make distr gen more faster
1 parent b5c033e commit 2f27839

5 files changed

Lines changed: 109 additions & 28 deletions

File tree

‎src/distributions.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -292,11 +292,18 @@ impl SimdFloatExt for f32 {
292292
v.round_float()
293293
}
294294

295+
const PREFERS_F32_WN: bool = true;
296+
295297
#[inline(always)]
296298
fn from_f64_fast(v: f64) -> f32 {
297299
v as f32
298300
}
299301

302+
#[inline(always)]
303+
fn from_f32_fast(v: f32) -> f32 {
304+
v
305+
}
306+
300307
fn pi() -> f32 {
301308
std::f32::consts::PI
302309
}
@@ -385,6 +392,8 @@ impl SimdFloatExt for f64 {
385392
f64x8::from_i32x8(v)
386393
}
387394

395+
const PREFERS_F32_WN: bool = false;
396+
388397
#[inline(always)]
389398
fn from_f64_fast(v: f64) -> f64 {
390399
v

‎src/distributions/exp.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,15 +78,15 @@ fn efix<T: SimdFloatExt>(hz: i32, iz: usize, tables: &ExpZigTables, rng: &mut Si
7878

7979
loop {
8080
if iz == 0 {
81-
return T::from_f64_fast(ZIG_EXP_R - (1.0f64 - rng.random::<f64>()).ln());
81+
return T::from_f64_fast(ZIG_EXP_R - (1.0f64 - rng.next_f64()).ln());
8282
}
8383

8484
let x = (hz.unsigned_abs() as f64) * tables.we[iz];
85-
if tables.fe[iz] + rng.random::<f64>() * (tables.fe[iz - 1] - tables.fe[iz]) < (-x).exp() {
85+
if tables.fe[iz] + rng.next_f64() * (tables.fe[iz - 1] - tables.fe[iz]) < (-x).exp() {
8686
return T::from_f64_fast(x);
8787
}
8888

89-
hz = rng.random::<i32>();
89+
hz = rng.next_i32();
9090
iz = (hz & 0xFF) as usize;
9191
let abs_hz = hz.unsigned_abs() as i64;
9292
if abs_hz < tables.ke[iz] as i64 {

‎src/distributions/normal.rs‎

Lines changed: 62 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use crate::simd_rng::SimdRng;
1818
struct ZigTables {
1919
kn: [i32; 128],
2020
wn: [f64; 128],
21+
wn_f32: [f32; 128],
2122
fn_tab: [f64; 128],
2223
}
2324

@@ -27,6 +28,7 @@ fn zig_tables() -> &'static ZigTables {
2728
ZIG_TABLES.get_or_init(|| {
2829
let mut kn = [0i32; 128];
2930
let mut wn = [0.0f64; 128];
31+
let mut wn_f32 = [0.0f32; 128];
3032
let mut fn_tab = [0.0f64; 128];
3133

3234
let mut dn = 3.442619855899f64;
@@ -63,7 +65,16 @@ fn zig_tables() -> &'static ZigTables {
6365
wn[i] = dn / m1;
6466
}
6567

66-
ZigTables { kn, wn, fn_tab }
68+
for i in 0..128 {
69+
wn_f32[i] = wn[i] as f32;
70+
}
71+
72+
ZigTables {
73+
kn,
74+
wn,
75+
wn_f32,
76+
fn_tab,
77+
}
6778
})
6879
}
6980

@@ -79,8 +90,8 @@ fn nfix<T: SimdFloatExt>(hz: i32, iz: usize, tables: &ZigTables, rng: &mut SimdR
7990

8091
if iz == 0 {
8192
loop {
82-
let u1: f64 = rng.random();
83-
let u2: f64 = rng.random();
93+
let u1: f64 = rng.next_f64();
94+
let u2: f64 = rng.next_f64();
8495
let x_tail = -0.2904764 * (-u1.ln());
8596
let y = -u2.ln();
8697
if y + y >= x_tail * x_tail {
@@ -94,13 +105,13 @@ fn nfix<T: SimdFloatExt>(hz: i32, iz: usize, tables: &ZigTables, rng: &mut SimdR
94105
}
95106
}
96107

97-
if tables.fn_tab[iz] + rng.random::<f64>() * (tables.fn_tab[iz - 1] - tables.fn_tab[iz])
108+
if tables.fn_tab[iz] + rng.next_f64() * (tables.fn_tab[iz - 1] - tables.fn_tab[iz])
98109
< (-0.5 * x * x).exp()
99110
{
100111
return T::from_f64_fast(x);
101112
}
102113

103-
hz = rng.random::<i32>();
114+
hz = rng.next_i32();
104115
iz = (hz & 127) as usize;
105116
if (hz.unsigned_abs() as i64) < tables.kn[iz] as i64 {
106117
return T::from_f64_fast(hz as f64 * tables.wn[iz]);
@@ -191,16 +202,29 @@ impl<T: SimdFloatExt, const N: usize> SimdNormal<T, N> {
191202
let abs_hz = hz.abs();
192203
let accept = abs_hz.simd_lt(kn_vals);
193204

194-
let wn_arr: [T; 8] = [
195-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[0] as usize)),
196-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[1] as usize)),
197-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[2] as usize)),
198-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[3] as usize)),
199-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[4] as usize)),
200-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[5] as usize)),
201-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[6] as usize)),
202-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[7] as usize)),
203-
];
205+
let wn_arr: [T; 8] = if T::PREFERS_F32_WN {
206+
[
207+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[0] as usize)),
208+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[1] as usize)),
209+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[2] as usize)),
210+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[3] as usize)),
211+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[4] as usize)),
212+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[5] as usize)),
213+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[6] as usize)),
214+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[7] as usize)),
215+
]
216+
} else {
217+
[
218+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[0] as usize)),
219+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[1] as usize)),
220+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[2] as usize)),
221+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[3] as usize)),
222+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[4] as usize)),
223+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[5] as usize)),
224+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[6] as usize)),
225+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[7] as usize)),
226+
]
227+
};
204228
let hz_float = T::simd_from_i32x8(hz);
205229
let wn_simd = T::simd_from_array(wn_arr);
206230
let result = hz_float * wn_simd;
@@ -308,16 +332,29 @@ impl<T: SimdFloatExt, const N: usize> SimdNormal<T, N> {
308332
let abs_hz = hz.abs();
309333
let accept = abs_hz.simd_lt(kn_vals);
310334

311-
let wn_arr: [T; 8] = [
312-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[0] as usize)),
313-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[1] as usize)),
314-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[2] as usize)),
315-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[3] as usize)),
316-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[4] as usize)),
317-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[5] as usize)),
318-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[6] as usize)),
319-
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[7] as usize)),
320-
];
335+
let wn_arr: [T; 8] = if T::PREFERS_F32_WN {
336+
[
337+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[0] as usize)),
338+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[1] as usize)),
339+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[2] as usize)),
340+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[3] as usize)),
341+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[4] as usize)),
342+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[5] as usize)),
343+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[6] as usize)),
344+
T::from_f32_fast(*tables.wn_f32.get_unchecked(iz_arr[7] as usize)),
345+
]
346+
} else {
347+
[
348+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[0] as usize)),
349+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[1] as usize)),
350+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[2] as usize)),
351+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[3] as usize)),
352+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[4] as usize)),
353+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[5] as usize)),
354+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[6] as usize)),
355+
T::from_f64_fast(*tables.wn.get_unchecked(iz_arr[7] as usize)),
356+
]
357+
};
321358
let hz_float = T::simd_from_i32x8(hz);
322359
let wn_simd = T::simd_from_array(wn_arr);
323360
let result = hz_float * wn_simd;

‎src/simd_rng.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,10 @@ pub struct SimdRng {
169169
f32_engine: Xoshiro128PP8,
170170
u64_buf: [u64; 4],
171171
u64_idx: usize,
172+
f64_scalar_buf: [f64; 8],
173+
f64_scalar_idx: usize,
174+
i32_scalar_buf: [i32; 8],
175+
i32_scalar_idx: usize,
172176
}
173177

174178
impl SimdRng {
@@ -182,6 +186,10 @@ impl SimdRng {
182186
f32_engine: Xoshiro128PP8::new_from_u64(seed32),
183187
u64_buf: [0; 4],
184188
u64_idx: 4,
189+
f64_scalar_buf: [0.0; 8],
190+
f64_scalar_idx: 8,
191+
i32_scalar_buf: [0; 8],
192+
i32_scalar_idx: 8,
185193
}
186194
}
187195

@@ -226,6 +234,28 @@ impl SimdRng {
226234
(a[7] >> 8) as f32 * F32_SCALE,
227235
]
228236
}
237+
238+
#[inline(always)]
239+
pub fn next_f64(&mut self) -> f64 {
240+
if self.f64_scalar_idx >= 8 {
241+
self.f64_scalar_buf = self.next_f64_array();
242+
self.f64_scalar_idx = 0;
243+
}
244+
let v = self.f64_scalar_buf[self.f64_scalar_idx];
245+
self.f64_scalar_idx += 1;
246+
v
247+
}
248+
249+
#[inline(always)]
250+
pub fn next_i32(&mut self) -> i32 {
251+
if self.i32_scalar_idx >= 8 {
252+
self.i32_scalar_buf = self.next_i32x8().to_array();
253+
self.i32_scalar_idx = 0;
254+
}
255+
let v = self.i32_scalar_buf[self.i32_scalar_idx];
256+
self.i32_scalar_idx += 1;
257+
v
258+
}
229259
}
230260

231261
impl Default for SimdRng {

‎src/traits.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,12 @@ pub trait SimdFloatExt: num_traits::Float + Default + Send + Sync + 'static {
145145
fn fill_uniform_simd(rng: &mut crate::simd_rng::SimdRng, out: &mut [Self]);
146146
fn sample_uniform<R: Rng + ?Sized>(rng: &mut R) -> Self;
147147
fn simd_from_i32x8(v: wide::i32x8) -> Self::Simd;
148+
const PREFERS_F32_WN: bool = false;
148149
fn from_f64_fast(v: f64) -> Self;
150+
#[inline(always)]
151+
fn from_f32_fast(v: f32) -> Self {
152+
Self::from_f64_fast(v as f64)
153+
}
149154
fn pi() -> Self;
150155
fn two_pi() -> Self;
151156
fn min_positive_val() -> Self;

0 commit comments

Comments
 (0)