cubek_random/
normal.rs

1use cubecl::{prelude::*, std::tensor::View};
2use std::f32::consts::PI;
3
4use super::{PrngArgs, PrngRuntime, random};
5
6use crate::{RandomFamily, lcg_step, taus_step_0, taus_step_1, taus_step_2, to_unit_interval_open};
7
8#[derive(CubeLaunch, CubeType)]
9pub(crate) struct Normal {
10    mean: f32,
11    std: f32,
12}
13
14#[derive(Debug)]
15struct NormalFamily;
16
17impl RandomFamily for NormalFamily {
18    type Runtime = Normal;
19}
20
21#[cube]
22impl PrngRuntime for Normal {
23    fn inner_loop<E: Numeric, N: Size>(
24        args: Normal,
25        write_index_base: usize,
26        n_invocations: u32,
27        #[comptime] n_values_per_thread: usize,
28        state_0: &mut u32,
29        state_1: &mut u32,
30        state_2: &mut u32,
31        state_3: &mut u32,
32        output: &mut View<Vector<E, N>, usize, ReadWrite>,
33    ) {
34        let mean = f32::cast_from(args.mean);
35        let std = f32::cast_from(args.std);
36
37        let mut output_vector_0 = Vector::empty();
38        let mut output_vector_1 = Vector::empty();
39
40        let num_iterations = n_values_per_thread / N::value() / 2;
41        #[unroll(num_iterations <= 8)]
42        for vector_index in 0..num_iterations {
43            // vectorization
44            #[unroll]
45            for i in 0..N::value() {
46                // First random uniform integer
47                *state_0 = taus_step_0(*state_0);
48                *state_1 = taus_step_1(*state_1);
49                *state_2 = taus_step_2(*state_2);
50                *state_3 = lcg_step(*state_3);
51
52                let int_random = *state_0 ^ *state_1 ^ *state_2 ^ *state_3;
53                let unit_0 = to_unit_interval_open(int_random);
54
55                // Second random uniform integer
56                *state_0 = taus_step_0(*state_0);
57                *state_1 = taus_step_1(*state_1);
58                *state_2 = taus_step_2(*state_2);
59                *state_3 = lcg_step(*state_3);
60
61                let int_random = *state_0 ^ *state_1 ^ *state_2 ^ *state_3;
62                let unit_1 = to_unit_interval_open(int_random);
63
64                // Box-Muller transform
65                let coeff = unit_0.ln() * -2.0;
66                let coeff = coeff.sqrt() * std;
67                let trigo_arg = 2.0 * PI * unit_1;
68
69                let normal_0 = f32::cos(trigo_arg) * coeff + mean;
70                let normal_1 = f32::sin(trigo_arg) * coeff + mean;
71
72                output_vector_0[i] = E::cast_from(normal_0);
73                output_vector_1[i] = E::cast_from(normal_1);
74            }
75
76            let iteration_offset = vector_index * n_invocations as usize * 2;
77            let write_index_0 = write_index_base + iteration_offset;
78            let write_index_1 = write_index_0 + n_invocations as usize;
79
80            output[write_index_0] = output_vector_0;
81            output[write_index_1] = output_vector_1;
82        }
83    }
84}
85
86impl PrngArgs for Normal {
87    type Args = Self;
88
89    fn args<R: Runtime>(self) -> NormalLaunch<R> {
90        NormalLaunch::new(self.mean, self.std)
91    }
92}
93
94/// Pseudo-random generator with uniform distribution
95pub fn random_normal<R: Runtime>(
96    client: &ComputeClient<R>,
97    mean: f32,
98    std: f32,
99    out: TensorBinding<R>,
100    dtype: StorageType,
101) -> Result<(), LaunchError> {
102    random::<NormalFamily, R>(client, Normal { mean, std }, out, dtype)
103}
cubek_random/normal.rs

cubek_random/
normal.rs
