Line data Source code
1 : #ifndef HEADER_fd_src_ballet_chacha_fd_chacha_rng_h
2 : #define HEADER_fd_src_ballet_chacha_fd_chacha_rng_h
3 :
4 : /* fd_chacha_rng provides APIs for ChaCha-based RNG, as used in the
5 : Solana protocol. This API should only be used where necessary.
6 : fd_rng is a better choice in all other cases. */
7 :
8 : #include "fd_chacha.h"
9 : #include "../../util/bits/fd_uwide.h"
10 :
11 : /* FD_CHACHA_RNG_DEBUG controls debug logging. 0 is off; 1 is on. */
12 :
13 : #ifndef FD_CHACHA_RNG_DEBUG
14 : #define FD_CHACHA_RNG_DEBUG 0
15 : #endif
16 :
17 : /* Solana uses different mechanisms of mapping a ulong to an unbiased
18 : integer in [0, n) in different parts of the code. In particular,
19 : leader schedule generation uses MODE_MOD and Turbine uses MODE_SHIFT.
20 : See the note in fd_chacha_rng_ulong_roll for more details. */
21 20705151 : #define FD_CHACHA_RNG_MODE_MOD 1
22 0 : #define FD_CHACHA_RNG_MODE_SHIFT 2
23 :
24 : /* Solana uses different ChaCha algorithms. Leader schedule generation
25 : uses ChaCha20 and Turbine switched from ChaCha20 to ChaCha8. */
26 1294 : #define FD_CHACHA_RNG_ALGO_CHACHA20 1
27 277835 : #define FD_CHACHA_RNG_ALGO_CHACHA8 2
28 :
29 : /* FD_CHACHA_RNG_BUFSZ is the internal buffer size of pre-generated
30 : ChaCha20 blocks. Multiple of block size (64 bytes) and a power of 2. */
31 :
32 : #if FD_HAS_AVX512
33 0 : #define FD_CHACHA_RNG_BUFSZ (16*FD_CHACHA_BLOCK_SZ)
34 : #elif FD_HAS_AVX
35 : #define FD_CHACHA_RNG_BUFSZ (8*FD_CHACHA_BLOCK_SZ)
36 : #elif FD_HAS_SVE2
37 : #define FD_CHACHA_RNG_BUFSZ (4*FD_CHACHA_BLOCK_SZ)
38 : #else
39 : #define FD_CHACHA_RNG_BUFSZ (256UL)
40 : #endif
41 :
42 : struct __attribute__((aligned(32UL))) fd_chacha_rng_private {
43 : /* ChaCha20 encryption key */
44 : uchar key[ 32UL ] __attribute__((aligned(32UL)));
45 :
46 : /* Ring buffer of pre-generated ChaCha20 RNG data.
47 : Note: We currently assume all reads are 8 byte. This means the
48 : cursor is always aligned by 8 and strictly increases in
49 : increments of 8. Thus, we really only have to refill the
50 : buffer if buf_off==buf_fill. pre_buf is used for one
51 : optimization where it's helpful to stage data before buf. */
52 : uchar pre_buf[ 128 ] __attribute__((aligned(FD_CHACHA_BLOCK_SZ)));
53 : uchar buf[ FD_CHACHA_RNG_BUFSZ ] __attribute__((aligned(FD_CHACHA_BLOCK_SZ)));
54 : ulong buf_off; /* Total number of bytes consumed */
55 : ulong buf_fill; /* Total number of bytes produced
56 : Always aligned by FD_CHACHA_BLOCK_SZ */
57 :
58 : int mode;
59 : int algo;
60 : };
61 : typedef struct fd_chacha_rng_private fd_chacha_rng_t;
62 :
63 : FD_PROTOTYPES_BEGIN
64 :
65 : /* fd_chacha_rng_{align,footprint} give the needed alignment and
66 : footprint of a memory region suitable to hold a ChaCha20-based RNG.
67 :
68 : fd_chacha_rng_new formats a memory region with suitable alignment
69 : and footprint for holding a chacha20_rng object. Assumes shmem
70 : points on the caller to the first byte of the memory region owned by
71 : the caller to use. `mode` must be one of the FD_CHACHA_RNG_MODE_*
72 : constants defined above and dictates what mode this object will use
73 : to generate random numbers. Returns shmem on success and NULL on
74 : failure (logs details). The memory region will be owned by the
75 : object on successful return. The caller is not joined on return.
76 :
77 : fd_chacha_rng_join joins the caller to a chacha20_rng object.
78 : Assumes shrng points to the first byte of the memory region holding
79 : the object. Returns a local handle to the join on success (this is
80 : not necessarily a simple cast of the address) and NULL on failure
81 : (logs details).
82 :
83 : fd_chacha_rng_leave leaves the caller's current local join to a
84 : ChaCha20 RNG object. Returns a pointer to the memory region holding
85 : the object on success this is not necessarily a simple cast of the
86 : address) and NULL on failure (logs details). The caller is not
87 : joined on successful return.
88 :
89 : fd_chacha_rng_delete unformats a memory region that holds a ChaCha20
90 : RNG object. Assumes shrng points on the caller to the first byte of
91 : the memory region holding the state and that nobody is joined.
92 : Returns a pointer to the memory region on success and NULL on failure
93 : (logs details). The caller has ownership of the memory region on
94 : successful return. */
95 :
96 : FD_FN_CONST ulong
97 : fd_chacha_rng_align( void );
98 :
99 : FD_FN_CONST ulong
100 : fd_chacha_rng_footprint( void );
101 :
102 : void *
103 : fd_chacha_rng_new( void * shmem, int mode );
104 :
105 : fd_chacha_rng_t *
106 : fd_chacha_rng_join( void * shrng );
107 :
108 : void *
109 : fd_chacha_rng_leave( fd_chacha_rng_t * );
110 :
111 : void *
112 : fd_chacha_rng_delete( void * shrng );
113 :
114 : /* fd_chacha_rng_init starts a ChaCha{8,20} RNG stream. rng is
115 : assumed to be a current local join to a chacha_rng object with no
116 : other concurrent operation that would modify the state while this is
117 : executing. seed points to the first byte of the RNG seed byte vector
118 : with 32 byte size. Any preexisting state for an in-progress or
119 : recently completed calculation will be discarded. Returns rng (on
120 : return, rng will have the state of a new in-progress calculation).
121 :
122 : The param algo is expected to be FD_CHACHA_RNG_ALGO_CHACHA20 or
123 : FD_CHACHA_RNG_ALGO_CHACH8. If invalid, it defaults to chacha20.
124 :
125 : Compatible with Rust fn rand_chacha::ChaCha20Rng::from_seed
126 : https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html#method.from_seed
127 : (and ChaCha8Rng) */
128 :
129 : fd_chacha_rng_t *
130 : fd_chacha_rng_init( fd_chacha_rng_t * rng,
131 : void const * key,
132 : int algo );
133 :
134 : /* The refill function. Not part of the public API. */
135 :
136 : #if FD_HAS_AVX512
137 : void fd_chacha8_rng_refill_avx512 ( fd_chacha_rng_t * rng );
138 : void fd_chacha20_rng_refill_avx512( fd_chacha_rng_t * rng );
139 : #endif
140 :
141 : #if FD_HAS_AVX
142 : void fd_chacha8_rng_refill_avx ( fd_chacha_rng_t * rng );
143 : void fd_chacha20_rng_refill_avx( fd_chacha_rng_t * rng );
144 : #endif
145 :
146 : #if FD_HAS_SVE2
147 : void fd_chacha8_rng_refill_sve2 ( fd_chacha_rng_t * rng );
148 : void fd_chacha20_rng_refill_sve2( fd_chacha_rng_t * rng );
149 : #endif
150 :
151 : void fd_chacha8_rng_refill_seq ( fd_chacha_rng_t * rng );
152 : void fd_chacha20_rng_refill_seq( fd_chacha_rng_t * rng );
153 :
154 : #if FD_HAS_AVX512
155 0 : #define fd_chacha8_rng_private_refill fd_chacha8_rng_refill_avx512
156 277835 : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_avx512
157 : #elif FD_HAS_AVX
158 : #define fd_chacha8_rng_private_refill fd_chacha8_rng_refill_avx
159 : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_avx
160 : #elif FD_HAS_SVE2
161 : #define fd_chacha8_rng_private_refill fd_chacha8_rng_refill_sve2
162 : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_sve2
163 : #else
164 : #define fd_chacha8_rng_private_refill fd_chacha8_rng_refill_seq
165 : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_seq
166 : #endif
167 :
168 : /* fd_chacha_rng_avail returns the number of buffered bytes. */
169 :
170 : FD_FN_PURE static inline ulong
171 19054228 : fd_chacha_rng_avail( fd_chacha_rng_t const * rng ) {
172 19054228 : return rng->buf_fill - rng->buf_off;
173 19054228 : }
174 :
175 : /* fd_chacha_rng_ulong read a 64-bit integer in [0,2^64) from the
176 : RNG stream. */
177 :
178 : static inline ulong
179 18449674 : fd_chacha_rng_ulong( fd_chacha_rng_t * rng ) {
180 18449674 : if( FD_UNLIKELY( fd_chacha_rng_avail( rng ) < sizeof(ulong) ) ) {
181 277188 : if( rng->algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
182 0 : fd_chacha8_rng_private_refill( rng );
183 277188 : } else {
184 277188 : fd_chacha20_rng_private_refill( rng );
185 277188 : }
186 277188 : }
187 18449674 : ulong x = FD_LOAD( ulong, rng->buf + (rng->buf_off % FD_CHACHA_RNG_BUFSZ) );
188 18449674 : rng->buf_off += 8U;
189 18449674 : return x;
190 18449674 : }
191 :
192 : /* fd_chacha_rng_read32 reads 32 bytes from the RNG stream into buf.
193 : buf may be unaligned. */
194 :
195 : static inline void
196 : fd_chacha_rng_read32( fd_chacha_rng_t * rng,
197 0 : void * buf ) {
198 0 : rng->buf_off = fd_ulong_align_up( rng->buf_off, 32UL );
199 0 : if( FD_UNLIKELY( fd_chacha_rng_avail( rng )<32UL ) ) {
200 0 : if( rng->algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
201 0 : fd_chacha8_rng_private_refill( rng );
202 0 : } else {
203 0 : fd_chacha20_rng_private_refill( rng );
204 0 : }
205 0 : }
206 :
207 0 : fd_memcpy( buf, rng->buf + (rng->buf_off % FD_CHACHA_RNG_BUFSZ), 32UL );
208 0 : rng->buf_off += 32UL;
209 0 : }
210 :
211 : /* fd_chacha_rng_ulong_roll returns an uniform IID rand in [0,n)
212 : analogous to fd_rng_ulong_roll. Rejection method based using
213 : fd_chacha_rng_ulong.
214 :
215 : Compatible with Rust type
216 : <rand_chacha::ChaCha20Rng as rand::Rng>::gen<rand::distributions::Uniform<u64>>()
217 : as of version 0.7.0 of the crate
218 : https://docs.rs/rand/latest/rand/distributions/struct.Uniform.html */
219 :
220 : static inline ulong
221 : fd_chacha_rng_ulong_roll( fd_chacha_rng_t * rng,
222 20704504 : ulong n ) {
223 : /* We use a pretty standard rejection-sampling based approach here,
224 : but for future reference, here's an explanation:
225 :
226 : We know that v can take 2^64 values, and so any method that maps
227 : each of the 2^64 values to the range directly [0, n) will not be
228 : uniform distribution when 2^64 is not divisible by n. This
229 : motivates using rejection sampling.
230 :
231 : The most basic approach is to map v from [0, n*floor(2^64/n) ) to
232 : [0, n) using v%n, but that puts a modulus on the critical path. To
233 : avoid that, the Rust rand crate uses a different approach: compute
234 : v*n/2^64, which is also in [0, n).
235 :
236 : Now the question to answer is which values to throw out. We pick a
237 : large integer k such that k*n<=2^64 and map [0, k*n) -> 0, [2^64,
238 : 2^64+k*n) -> 1, etc. Since k*n might be 2^64 and then not fit in a
239 : long, we define zone=k*n-1 <= ULONG_MAX, and make the intervals
240 : closed instead of half-open.
241 :
242 : Here's where the mode comes in. Depending on what method you call
243 : and what datatype you use, the Rust crate uses different values of
244 : k. When MODE_MOD is set, we use largest possible value of k,
245 : namely floor(2^64/n). You can compute zone directly as follows:
246 : zone = k*n-1
247 : = floor(2^64/n)*n - 1
248 : = 2^64 - (2^64%n) - 1
249 : = 2^64-1 - (2^64-n)%n, since n<2^64
250 : = 2^64-1 - ((2^64-1)-n+1)%n
251 : Which is back to having a mod... But at least if n is a
252 : compile-time constant then the whole zone computation becomes a
253 : compile-time constant.
254 :
255 : When MODE_SHIFT is set, we use uses almost the largest possible
256 : power of two for k. Precisely, it uses the smallest power of two
257 : such that k*n >= 2^63, which is the largest power of two such that
258 : k*n<=2^64 unless n is a power of two. This approach eliminates the
259 : mod calculation but increases the expected number of samples
260 : required. */
261 20704504 : ulong const zone = fd_ulong_if( rng->mode==FD_CHACHA_RNG_MODE_MOD,
262 20704504 : ULONG_MAX - (ULONG_MAX-n+1UL)%n,
263 20704504 : (n << (63 - fd_ulong_find_msb( n ) )) - 1UL );
264 :
265 >1844*10^16 : for( int i=0; 1; i++ ) {
266 17978610 : ulong v = fd_chacha_rng_ulong( rng );
267 17978610 : ulong hi, lo;
268 17978610 : fd_uwide_mul( &hi, &lo, v, n );
269 :
270 : # if FD_CHACHA_RNG_DEBUG
271 : FD_LOG_DEBUG(( "roll (attempt %d): n=%016lx zone: %016lx v=%016lx lo=%016lx hi=%016lx", i, n, zone, v, lo, hi ));
272 : # else
273 17978610 : (void)i;
274 17978610 : # endif /* FD_CHACHA_RNG_DEBUG */
275 :
276 19227274 : if( FD_LIKELY( lo<=zone ) ) return hi;
277 17978610 : }
278 20704504 : }
279 :
280 : /* fd_chacha_rng_ulong_roll_decreasing is an optimized version of
281 : for( ulong i=0UL; i<cnt; i++ )
282 : out[i] = fd_chacha_rng_ulong_roll( rng, n_initial-i );
283 : cnt must be in [0, n_initial). rng==NULL is not allowed, even if
284 : cnt==0. */
285 : void
286 : fd_chacha_rng_ulong_roll_decreasing( fd_chacha_rng_t * rng,
287 : ulong n_initial,
288 : ulong * out,
289 : ulong cnt );
290 :
291 : FD_PROTOTYPES_END
292 :
293 : #endif /* HEADER_fd_src_ballet_chacha20_fd_chacha20_rng_h */
|
