urng 1.0.0

Universal Random Number Generator
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use std::slice::from_raw_parts_mut;

use rayon::prelude::*;
use wrapn::{wrap, wu64};

use crate::_internal::{fill_chunk_auto, prefer_nt};
use crate::cbrng::b64::Philox64;
use crate::{i2f_bits, u2f_01};

/// Creates a new heap-allocated `Philox64` and returns a raw pointer to it.
/// The caller is responsible for freeing it with [`philox64_free`].
#[unsafe(no_mangle)]
pub extern "C" fn philox64_new(seed: u64) -> *mut Philox64 {
    Box::into_raw(Box::new(Philox64::new(seed)))
}
/// Frees a `Philox64` instance previously created by [`philox64_new`].
/// Does nothing if `ptr` is null.
#[unsafe(no_mangle)]
pub extern "C" fn philox64_free(ptr: *mut Philox64) {
    if !ptr.is_null() {
        unsafe { drop(Box::from_raw(ptr)) }
    }
}

/// Parallel batch size for Philox64 (elements per thread task).
const PHILOX64_PAR_CHUNK: usize = 0x20000;

/// Fills `buffer` in parallel from counter-mode Philox64 blocks, mapping
/// each raw `u64` through `map`. Four blocks (8 outputs = one 64-byte
/// cache line for 8-byte `T`) are batched per generator call so the
/// non-temporal path can stream whole lines.
#[inline(always)]
fn philox64_fill<T, M>(buffer: &mut [T], c0: [wu64; 2], k: [wu64; 2], map: M)
where
    T: Copy + Default + Send,
    M: Fn(u64) -> T + Sync,
{
    let nt = prefer_nt::<T>(buffer.len());
    buffer
        .par_chunks_mut(PHILOX64_PAR_CHUNK)
        .enumerate()
        .for_each(|(chunk_idx, chunk)| {
            let mut block = ((chunk_idx * PHILOX64_PAR_CHUNK) / 2) as u64;
            unsafe {
                fill_chunk_auto(chunk, nt, || {
                    let mut out = [T::default(); 8];
                    for blk in out.chunks_exact_mut(2) {
                        let mut c = c0;
                        let (new_c0, overflow) = c[0].overflowing_add(block);
                        c[0] = wrap!(new_c0);
                        if overflow {
                            c[1] += 1;
                        }
                        let r = Philox64::compute(c, k);
                        blk[0] = map(*r[0]);
                        blk[1] = map(*r[1]);
                        block = block.wrapping_add(1);
                    }
                    out
                });
            }
        });
}

/// Advances the 128-bit counter past everything `philox64_fill` consumed
/// (4 blocks per 8-output batch).
#[inline(always)]
fn philox64_advance(rng: &mut Philox64, count: usize) {
    let num_blocks = (count.div_ceil(8) * 4) as u64;
    let (new_c0, carry) = rng.c[0].overflowing_add(num_blocks);
    rng.c[0] = new_c0.into();
    if carry {
        rng.c[1] += 1;
    }
}

/// Fills `out[0..count]` with raw `u64` random values using parallel counter-based generation.
#[unsafe(no_mangle)]
pub extern "C" fn philox64_next_u64s(ptr: *mut Philox64, out: *mut u64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        philox64_fill(buffer, rng.c, rng.k, |x| x);
        philox64_advance(rng, count);
    }
}

/// Fills `out[0..count]` with `f64` values in `[0, 1)` using parallel counter-based generation.
#[unsafe(no_mangle)]
pub extern "C" fn philox64_next_f64s(ptr: *mut Philox64, out: *mut f64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        philox64_fill(buffer, rng.c, rng.k, |x| u2f_01!(f64, 64, x));
        philox64_advance(rng, count);
    }
}

/// Fills `out[0..count]` with `i64` values in `[min, max]` using parallel counter-based generation.
#[unsafe(no_mangle)]
pub extern "C" fn philox64_rand_i64s(
    ptr: *mut Philox64,
    out: *mut i64,
    count: usize,
    min: i64,
    max: i64,
) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        let range = (max as i128 - min as i128 + 1) as u128;
        philox64_fill(buffer, rng.c, rng.k, |x| {
            ((x as u128 * range) >> 64) as i64 + min
        });
        philox64_advance(rng, count);
    }
}

/// Fills `out[0..count]` with `f64` values in `[min, max)` using parallel counter-based generation.
#[unsafe(no_mangle)]
pub extern "C" fn philox64_rand_f64s(
    ptr: *mut Philox64,
    out: *mut f64,
    count: usize,
    min: f64,
    max: f64,
) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        let mult = max - min;
        philox64_fill(buffer, rng.c, rng.k, |x| u2f_01!(f64, 64, x) * mult + min);
        philox64_advance(rng, count);
    }
}