Line data Source code
1 : #include "fd_chacha_rng.h"
2 : #if FD_HAS_AVX512
3 : #include "../../util/simd/fd_avx512.h"
4 : #include "../../util/simd/fd_sse.h"
5 : #endif
6 :
7 : FD_FN_CONST ulong
8 0 : fd_chacha_rng_align( void ) {
9 0 : return alignof(fd_chacha_rng_t);
10 0 : }
11 :
12 : FD_FN_CONST ulong
13 0 : fd_chacha_rng_footprint( void ) {
14 0 : return sizeof(fd_chacha_rng_t);
15 0 : }
16 :
17 : void *
18 647 : fd_chacha_rng_new( void * shmem, int mode ) {
19 647 : if( FD_UNLIKELY( !shmem ) ) {
20 0 : FD_LOG_WARNING(( "NULL shmem" ));
21 0 : return NULL;
22 0 : }
23 647 : if( FD_UNLIKELY( !fd_ulong_is_aligned( (ulong)shmem, alignof(fd_chacha_rng_t) ) ) ) {
24 0 : FD_LOG_WARNING(( "misaligned shmem" ));
25 0 : return NULL;
26 0 : }
27 647 : memset( shmem, 0, sizeof(fd_chacha_rng_t) );
28 647 : if( FD_UNLIKELY( (mode!=FD_CHACHA_RNG_MODE_MOD) & (mode!=FD_CHACHA_RNG_MODE_SHIFT) ) ) {
29 0 : FD_LOG_WARNING(( "invalid mode" ));
30 0 : return NULL;
31 0 : }
32 647 : ((fd_chacha_rng_t *)shmem)->mode = mode;
33 647 : ((fd_chacha_rng_t *)shmem)->algo = FD_CHACHA_RNG_ALGO_CHACHA20;
34 647 : return shmem;
35 647 : }
36 :
37 : fd_chacha_rng_t *
38 647 : fd_chacha_rng_join( void * shrng ) {
39 647 : if( FD_UNLIKELY( !shrng ) ) {
40 0 : FD_LOG_WARNING(( "NULL shrng" ));
41 0 : return NULL;
42 0 : }
43 647 : return (fd_chacha_rng_t *)shrng;
44 647 : }
45 :
46 : void *
47 647 : fd_chacha_rng_leave( fd_chacha_rng_t * rng ) {
48 647 : if( FD_UNLIKELY( !rng ) ) {
49 0 : FD_LOG_WARNING(( "NULL rng" ));
50 0 : return NULL;
51 0 : }
52 647 : return (void *)rng;
53 647 : }
54 :
55 : void *
56 647 : fd_chacha_rng_delete( void * shrng ) {
57 647 : if( FD_UNLIKELY( !shrng ) ) {
58 0 : FD_LOG_WARNING(( "NULL shrng" ));
59 0 : return NULL;
60 0 : }
61 647 : memset( shrng, 0, sizeof(fd_chacha_rng_t) );
62 647 : return shrng;
63 647 : }
64 :
65 : fd_chacha_rng_t *
66 : fd_chacha_rng_init( fd_chacha_rng_t * rng,
67 : void const * key,
68 647 : int algo ) {
69 647 : memcpy( rng->key, key, FD_CHACHA_KEY_SZ );
70 647 : rng->buf_off = 0UL;
71 647 : rng->buf_fill = 0UL;
72 :
73 : /* invalid algo defaults to chacha20 */
74 647 : rng->algo = algo;
75 647 : if( algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
76 0 : fd_chacha8_rng_private_refill( rng );
77 647 : } else {
78 647 : fd_chacha20_rng_private_refill( rng );
79 647 : }
80 :
81 647 : return rng;
82 647 : }
83 :
84 : static void
85 : fd_chacha_rng_refill_seq( fd_chacha_rng_t * rng,
86 0 : void * (* block_fn)( void *, void const *, void const * ) ) {
87 0 : ulong fill_target = FD_CHACHA_RNG_BUFSZ - FD_CHACHA_BLOCK_SZ;
88 :
89 0 : ulong buf_avail;
90 0 : while( (buf_avail=(rng->buf_fill - rng->buf_off))<fill_target ) {
91 0 : ulong idx = rng->buf_fill >> 6;
92 0 : uint idx_nonce[4] __attribute__((aligned(16))) =
93 0 : { (uint)idx, (uint)(idx>>32), 0U, 0U };
94 0 : block_fn( rng->buf + (rng->buf_fill % FD_CHACHA_RNG_BUFSZ),
95 0 : rng->key,
96 0 : idx_nonce );
97 0 : rng->buf_fill += (uint)FD_CHACHA_BLOCK_SZ;
98 0 : }
99 0 : }
100 :
101 : void
102 0 : fd_chacha8_rng_refill_seq( fd_chacha_rng_t * rng ) {
103 0 : fd_chacha_rng_refill_seq( rng, fd_chacha8_block );
104 0 : }
105 :
106 : void
107 0 : fd_chacha20_rng_refill_seq( fd_chacha_rng_t * rng ) {
108 0 : fd_chacha_rng_refill_seq( rng, fd_chacha20_block );
109 0 : }
110 :
111 : void
112 : fd_chacha_rng_ulong_roll_decreasing( fd_chacha_rng_t * rng,
113 : ulong n_initial,
114 : ulong * out,
115 0 : ulong cnt ) {
116 : #if !FD_HAS_AVX512
117 : int use_fast_path = 0;
118 : #else
119 0 : int use_fast_path = (rng->mode==FD_CHACHA_RNG_MODE_SHIFT) & (cnt>=16UL) & (n_initial<(1UL<<40));
120 0 : #endif
121 :
122 0 : if( use_fast_path ) {
123 0 : #if FD_HAS_AVX512
124 : /* Using MODE_SHIFT, the expected probability of rejection is about
125 : 25%, which is sadly pretty high. We amortize that by guessing many
126 : values with AVX multiplication. wwv_mul is very slow on Intel CPUs
127 : but decent on recent AMD CPUs. */
128 :
129 0 : ulong buf_avail = fd_chacha_rng_avail( rng );
130 0 : while( cnt>=14UL ) {
131 : /* In this loop, we'll produce 0-8 samples and consume 8-15 ulongs
132 : from the stream.
133 :
134 : The loop invariant is that buf_avail is in [0,
135 : BUFSZ+6*sizeof(ulong)] at the start and end of each loop
136 : iteration. We need 15 ulongs of contiguous bytes though, and
137 : the only way to get more random data is from a refill that
138 : produces BUFSZ more data. That means buf_avail can peak at
139 : BUFSZ+14*sizeof(ulong), but then we always consume at least 8
140 : ulongs from the stream, so it ends at no more than
141 : BUFSZ+6*sizeof(ulong).
142 :
143 : We can't return from this function until buf_avail is <= BUFSZ,
144 : which means we need to exit this loop early enough to consume
145 : at least 6 more ulongs in the worst case. Since each iteration
146 : of the loop might produce 8 random samples, cnt>=6 at the end
147 : means cnt>=14 at the start. */
148 0 : if( FD_UNLIKELY( buf_avail<15UL*sizeof(ulong) ) ) {
149 0 : memcpy( rng->pre_buf, rng->buf+FD_CHACHA_RNG_BUFSZ-128UL, 128UL );
150 0 : rng->buf_off = rng->buf_fill;
151 0 : if( rng->algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
152 0 : fd_chacha8_rng_private_refill( rng );
153 0 : } else {
154 0 : fd_chacha20_rng_private_refill( rng );
155 0 : }
156 0 : buf_avail += FD_CHACHA_RNG_BUFSZ;
157 0 : }
158 :
159 0 : ulong const * load_base = (ulong const *)(rng->buf+FD_CHACHA_RNG_BUFSZ-buf_avail);
160 0 : wwv_t r0 = wwv_ldu( load_base+0UL );
161 0 : wwv_t r1 = wwv_ldu( load_base+1UL );
162 0 : wwv_t r2 = wwv_ldu( load_base+2UL );
163 0 : wwv_t r3 = wwv_ldu( load_base+3UL );
164 0 : wwv_t r4 = wwv_ldu( load_base+4UL );
165 0 : wwv_t r5 = wwv_ldu( load_base+5UL );
166 0 : wwv_t r6 = wwv_ldu( load_base+6UL );
167 0 : wwv_t r7 = wwv_ldu( load_base+7UL );
168 :
169 0 : wwv_t const ii = wwv( 0UL, 1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL );
170 :
171 0 : wwv_t n_initial_v = wwv_bcast( n_initial );
172 0 : wwv_t n = wwv_sub( n_initial_v, ii );
173 :
174 : /* Recomputing in parallel seems to be faster than broadcasting
175 : from n... */
176 0 : wwv_t n_initial_0 = n_initial_v;
177 0 : wwv_t n_initial_1 = wwv_sub( n_initial_v, wwv_bcast( 1UL ) );
178 0 : wwv_t n_initial_2 = wwv_sub( n_initial_v, wwv_bcast( 2UL ) );
179 0 : wwv_t n_initial_3 = wwv_sub( n_initial_v, wwv_bcast( 3UL ) );
180 0 : wwv_t n_initial_4 = wwv_sub( n_initial_v, wwv_bcast( 4UL ) );
181 0 : wwv_t n_initial_5 = wwv_sub( n_initial_v, wwv_bcast( 5UL ) );
182 0 : wwv_t n_initial_6 = wwv_sub( n_initial_v, wwv_bcast( 6UL ) );
183 0 : wwv_t n_initial_7 = wwv_sub( n_initial_v, wwv_bcast( 7UL ) );
184 :
185 : /* Recall that we need to check lo<=zone, where zone is
186 : n<<lzcnt(n) - 1. That's equivalent to lo<(n<<lzcnt(n)), but
187 : n<<lzcnt(n) could be 2^64. However, that's equivalent to
188 : (lo>>lzcnt(n)) < n, which again is equivalent to
189 : (lo>>lzcnt(n)) - n < 0. Given that n<2^40, we don't need to
190 : worry about overflow. */
191 :
192 0 : wwv_t prod0 = wwv_mul( r0, n_initial_0 );
193 0 : wwv_t prod1 = wwv_mul( r1, n_initial_1 );
194 0 : wwv_t prod2 = wwv_mul( r2, n_initial_2 );
195 0 : wwv_t prod3 = wwv_mul( r3, n_initial_3 );
196 0 : wwv_t prod4 = wwv_mul( r4, n_initial_4 );
197 0 : wwv_t prod5 = wwv_mul( r5, n_initial_5 );
198 0 : wwv_t prod6 = wwv_mul( r6, n_initial_6 );
199 0 : wwv_t prod7 = wwv_mul( r7, n_initial_7 );
200 :
201 : /* Only the high bit of each ulong is meaningful */
202 0 : wwv_t le0 = wwv_sub( wwv_shr_vector( prod0, _mm512_lzcnt_epi64( n_initial_0 ) ), n_initial_0 );
203 0 : wwv_t le1 = wwv_sub( wwv_shr_vector( prod1, _mm512_lzcnt_epi64( n_initial_1 ) ), n_initial_1 );
204 0 : wwv_t le2 = wwv_sub( wwv_shr_vector( prod2, _mm512_lzcnt_epi64( n_initial_2 ) ), n_initial_2 );
205 0 : wwv_t le3 = wwv_sub( wwv_shr_vector( prod3, _mm512_lzcnt_epi64( n_initial_3 ) ), n_initial_3 );
206 0 : wwv_t le4 = wwv_sub( wwv_shr_vector( prod4, _mm512_lzcnt_epi64( n_initial_4 ) ), n_initial_4 );
207 0 : wwv_t le5 = wwv_sub( wwv_shr_vector( prod5, _mm512_lzcnt_epi64( n_initial_5 ) ), n_initial_5 );
208 0 : wwv_t le6 = wwv_sub( wwv_shr_vector( prod6, _mm512_lzcnt_epi64( n_initial_6 ) ), n_initial_6 );
209 0 : wwv_t le7 = wwv_sub( wwv_shr_vector( prod7, _mm512_lzcnt_epi64( n_initial_7 ) ), n_initial_7 );
210 :
211 : /* We want to compute something like "if I had to consume k ulongs
212 : before getting to produce sample j, how many ulongs do I have
213 : to consume to produce sample j?" If le_i is
214 : [ a 0 0 0... , b 0, c 0, d 0, e 0, f 0, g 0, h 0 ]
215 : Where the 0s are the rest of the ulong, then we want to produce
216 : [ abcdefgh, 0bcdefgh, 00cdefgh, ... ]
217 : which we can do with one permutexvar. Then we lzcnt it, which
218 : gives us 8*(ignoring all lanes prior to the current one, what
219 : is the first non-zero lane in le). That's only 8 pieces of
220 : information, each <=6 bits, so we compact it. */
221 : #if 0
222 : static wwb_t const perm = wwb( 63, 55, 47, 39, 31, 23, 15, 7,
223 : 63, 55, 47, 39, 31, 23, 15, 0,
224 : 63, 55, 47, 39, 31, 23, 0, 0,
225 : 63, 55, 47, 39, 31, 0, 0, 0,
226 : 63, 55, 47, 39, 0, 0, 0, 0,
227 : 63, 55, 47, 0, 0, 0, 0, 0,
228 : 63, 55, 0, 0, 0, 0, 0, 0,
229 : 63, 0, 0, 0, 0, 0, 0, 0 );
230 : #else
231 0 : static ulong const perm_q[ 8 ] __attribute__((aligned(64))) = {
232 0 : 0x070F171F272F373FUL, /* lane 0 */
233 0 : 0x000F171F272F373FUL, /* lane 1 */
234 0 : 0x0000171F272F373FUL, /* lane 2 */
235 0 : 0x0000001F272F373FUL, /* lane 3 */
236 0 : 0x00000000272F373FUL, /* lane 4 */
237 0 : 0x00000000002F373FUL, /* lane 5 */
238 0 : 0x000000000000373FUL, /* lane 6 */
239 0 : 0x000000000000003FUL, /* lane 7 */
240 0 : };
241 0 : wwv_t perm = wwv_ld( perm_q );
242 0 : #endif
243 0 : const ulong DIAG_MASK = 0x0103070F1F3F7FFFUL;
244 0 : wwv_t w_table0 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le0 ) ), 3 );
245 0 : wwv_t w_table1 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le1 ) ), 3 );
246 0 : wwv_t w_table2 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le2 ) ), 3 );
247 0 : wwv_t w_table3 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le3 ) ), 3 );
248 0 : wwv_t w_table4 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le4 ) ), 3 );
249 0 : wwv_t w_table5 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le5 ) ), 3 );
250 0 : wwv_t w_table6 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le6 ) ), 3 );
251 0 : wwv_t w_table7 = wwv_shr( _mm512_lzcnt_epi64( _mm512_maskz_permutexvar_epi8( DIAG_MASK, perm, le7 ) ), 3 );
252 :
253 : /* Compact to the low 8 bytes, filling the rest with 8 */
254 0 : wwb_t const compact = _mm512_castsi128_si512( vb( 0, 8, 16, 24, 32, 40, 48, 56,
255 0 : 64,64, 64, 64, 64, 64, 64, 64 ) );
256 0 : wwb_t const eight = wwb_bcast( 8 );
257 0 : vb_t table0 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table0, compact, eight ) );
258 0 : vb_t table1 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table1, compact, eight ) );
259 0 : vb_t table2 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table2, compact, eight ) );
260 0 : vb_t table3 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table3, compact, eight ) );
261 0 : vb_t table4 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table4, compact, eight ) );
262 0 : vb_t table5 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table5, compact, eight ) );
263 0 : vb_t table6 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table6, compact, eight ) );
264 0 : vb_t table7 = _mm512_castsi512_si128( _mm512_permutex2var_epi8( w_table7, compact, eight ) );
265 :
266 : /* Up until this point, everything has been parallelizable, but
267 : now we need to combine all the information. We've formatted it
268 : in such a way that we can do that with just shifting. For
269 : example, if (suppressing all but the high bit)
270 : le0 = [ 0x80.., 0, 0x80.., 0x80.., 0 ... ]
271 : le1 = [ 0x80.., 0x80.., 0, 0, 0 ... ]
272 : le2 = [ 0, 0x80.., 0, 0x80.., 0 ... ]
273 : le3 = [ 0x80.., 0x80.., 0, 0x80.., 0 ... ]
274 : le4 = [ 0x80.., 0, 0, 0, 0x80 ... ]
275 :
276 : Then table are:
277 : table0 = [ 0, 2, 2, 3, ??, ... ]
278 : table1 = [ 0, 1, ??, ??, ?? ... ]
279 : table2 = [ 1, 1, 3, 3, ?? ... ]
280 : table3 = [ 0, 1, 3, 3, ?? ... ]
281 : table4 = [ 0, 4, 4, 4, 4 ... ]
282 :
283 : We can interpret these as a function, and applying them can be
284 : done with vpshufb, one of my favorite instructions. Composing
285 : them can also be done with vpshufb, and we can apply the
286 : composition of the whole thing with lower latency using a tree. */
287 :
288 0 : vb_t table01 = _mm_shuffle_epi8( table1, table0 );
289 0 : vb_t table23 = _mm_shuffle_epi8( table3, table2 );
290 0 : vb_t table45 = _mm_shuffle_epi8( table5, table4 );
291 0 : vb_t table67 = _mm_shuffle_epi8( table7, table6 );
292 0 : vb_t table0123 = _mm_shuffle_epi8( table23, table01 );
293 0 : vb_t table4567 = _mm_shuffle_epi8( table67, table45 );
294 0 : vb_t table_all0 = _mm_maskz_shuffle_epi8( 1, table4567, table0123 ); /* Only need byte 1, shaves a cycle off critical path */
295 :
296 : /* Extract byte0 from table_all. Once we do this, the next
297 : iteration can probably begin. */
298 0 : ulong final = (ulong)_mm_cvtsi128_si64( table_all0 );
299 :
300 : /* We need to produce the rest of the table values that aren't on
301 : the critical path. */
302 0 : vb_t table012 = _mm_shuffle_epi8( table2, table01 );
303 0 : vb_t table01234 = _mm_shuffle_epi8( table4, table0123 );
304 0 : vb_t table012345 = _mm_shuffle_epi8( table45, table0123 );
305 0 : vb_t table0123456 = _mm_shuffle_epi8( table6, table012345 );
306 :
307 : /* Now combine all of the byte 0s from all these vectors. */
308 0 : vb_t counter01 = _mm_unpacklo_epi8 ( table0, table01 );
309 0 : vb_t counter23 = _mm_unpacklo_epi8 ( table012, table0123 );
310 0 : vb_t counter45 = _mm_unpacklo_epi8 ( table01234, table012345 );
311 0 : vb_t counter67 = _mm_unpacklo_epi8 ( table0123456, table_all0 );
312 0 : vb_t counter0123 = _mm_unpacklo_epi16( counter01, counter23 );
313 0 : vb_t counter4567 = _mm_unpacklo_epi16( counter45, counter67 );
314 0 : vb_t counter_all = _mm_unpacklo_epi32( counter0123, counter4567 );
315 :
316 0 : ulong produced;
317 0 : ulong consumed;
318 0 : if( FD_LIKELY( final<8 ) ) {
319 : /* About 95% at p=0.7 */
320 0 : produced = 8UL;
321 0 : consumed = final+8UL;
322 0 : } else {
323 : /* Find the first index where it's 8. */
324 0 : produced = (ulong)fd_ulong_find_lsb( (ulong)_mm_cvtsi128_si64( counter_all ) & 0x0808080808080808UL ) >> 3;
325 : /* If we didn't produce all 8, then we don't need to retry a
326 : bunch of bad random data on the last value that we did
327 : produce. */
328 0 : consumed = produced + 8UL;
329 0 : }
330 : /* The next iteration can definitely begin now while we compute
331 : the final values. */
332 :
333 : /* We want to construct the final r vector with
334 : _mm512_permutex2var_epi64, with the indices coming from the
335 : counter values we produced above. Remember that lane j in
336 : vector i, is just lane i+j in r0, where values >=8 come from
337 : load_base+8. */
338 0 : wwv_t counter_wide = wwv_add( ii, _mm512_cvtepu8_epi64( counter_all ) );
339 0 : wwv_t r = _mm512_permutex2var_epi64( r0, counter_wide, wwv_ldu( load_base+8UL ) );
340 :
341 :
342 : /* Finally... we want the high part of r * n. There's no mulhi
343 : instruction though, but it's not too bad.
344 : Let r = r_hi * 2^52 + r_lo, with r_lo<2^52. n<2^52 by
345 : hypothesis.
346 :
347 : r*n = r_hi * n * 2^52 + r_lo*n
348 : r_lo*n = m52hi( r_lo, n )*2^52 + m52lo( r_lo, n ), where m52lo(-,-) < 2^52
349 :
350 : Since r_hi < 2^(64-52) and n<2^40, r_hi * n < 2^52, so
351 : r_hi * n = m52lo( r_hi, n ). So then
352 : r*n = 2^52 * (m52hi( r_lo, n ) + m52lo( r_hi, n )) +
353 : m52lo( r_lo, n ),
354 : Since we want to discard the low 64 bits and we know the last
355 : term doesn't carry into the higher 52 bits, we can just omit
356 : it.
357 : r*n>>64 = (m52hi( r_lo, n ) + m52lo( r_hi, n ))>>12 */
358 0 : wwv_t high_product = wwv_madd52hi( wwv_madd52lo( wwv_zero(), wwv_shr( r, 52 ), n ),
359 0 : r, n );
360 :
361 0 : wwv_t samples = wwv_shr( high_product, 12 );
362 0 : wwv_stu( out, samples );
363 :
364 0 : n_initial -= produced;
365 0 : out += produced;
366 0 : cnt -= produced;
367 0 : buf_avail -= consumed*sizeof(ulong);
368 0 : }
369 0 : while( (cnt>0) && (buf_avail>FD_CHACHA_RNG_BUFSZ) ) {
370 0 : ulong zone = (n_initial << (63 - fd_ulong_find_msb( n_initial) )) - 1UL;
371 0 : ulong const * load_base = (ulong const *)(rng->buf+FD_CHACHA_RNG_BUFSZ-buf_avail);
372 0 : ulong v = *load_base;
373 0 : ulong hi, lo;
374 0 : fd_uwide_mul( &hi, &lo, v, n_initial );
375 0 : *out = hi;
376 :
377 0 : ulong consumed = 1UL;
378 0 : ulong produced = (lo<=zone);
379 :
380 0 : n_initial -= produced;
381 0 : out += produced;
382 0 : cnt -= produced;
383 0 : buf_avail -= consumed*sizeof(ulong);
384 0 : }
385 0 : rng->buf_off = rng->buf_fill - buf_avail;
386 0 : #endif
387 0 : }
388 0 : for( ulong i=0UL; i<cnt; i++ ) out[i] = fd_chacha_rng_ulong_roll( rng, n_initial-i );
389 0 : }
|
