Skip to content

simd-0332: reduce chacha rounds for turbine from 20 to 8 #7597

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
0x0ece wants to merge 1 commit into
base: main
Choose a base branch
from
+1,004 −175
Open

simd-0332: reduce chacha rounds for turbine from 20 to 8 #7597

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/ballet/chacha/fd_chacha.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@

#define FD_CHACHA_BLOCK_SZ (64UL)

/* FD_CHACHA20_KEY_SZ is the size of the ChaCha20 encryption key */
/* FD_CHACHA_KEY_SZ is the size of the ChaCha20 encryption key */

#define FD_CHACHA20_KEY_SZ (32UL)
#define FD_CHACHA_KEY_SZ (32UL)

FD_PROTOTYPES_BEGIN

Expand Down
19 changes: 14 additions & 5 deletions src/ballet/chacha/fd_chacha_rng.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ fd_chacha_rng_new( void * shmem, int mode ) {
return NULL;
}
((fd_chacha_rng_t *)shmem)->mode = mode;

((fd_chacha_rng_t *)shmem)->algo = FD_CHACHA_RNG_ALGO_CHACHA20;
return shmem;
}

Expand Down Expand Up @@ -59,12 +59,21 @@ fd_chacha_rng_delete( void * shrng ) {
}

fd_chacha_rng_t *
fd_chacha20_rng_init( fd_chacha_rng_t * rng,
void const * key ) {
memcpy( rng->key, key, FD_CHACHA20_KEY_SZ );
fd_chacha_rng_init( fd_chacha_rng_t * rng,
void const * key,
int algo ) {
memcpy( rng->key, key, FD_CHACHA_KEY_SZ );
rng->buf_off = 0UL;
rng->buf_fill = 0UL;
fd_chacha20_rng_private_refill( rng );

/* invalid algo defaults to chacha20 */
rng->algo = algo;
if( algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
fd_chacha8_rng_private_refill( rng );
} else {
fd_chacha20_rng_private_refill( rng );
}

return rng;
}

Expand Down
68 changes: 31 additions & 37 deletions src/ballet/chacha/fd_chacha_rng.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,7 @@
fd_rng is a better choice in all other cases. */

#include "fd_chacha.h"
#if !FD_HAS_INT128
#include "../../util/bits/fd_uwide.h"
#endif

/* FD_CHACHA_RNG_DEBUG controls debug logging. 0 is off; 1 is on. */

Expand All @@ -19,10 +17,15 @@
/* Solana uses different mechanisms of mapping a ulong to an unbiased
integer in [0, n) in different parts of the code. In particular,
leader schedule generation uses MODE_MOD and Turbine uses MODE_SHIFT.
See the note in fd_chacha20_rng_ulong_roll for more details. */
See the note in fd_chacha_rng_ulong_roll for more details. */
#define FD_CHACHA_RNG_MODE_MOD 1
#define FD_CHACHA_RNG_MODE_SHIFT 2

/* Solana uses different ChaCha algorithms. Leader schedule generation
uses ChaCha20 and Turbine switched from ChaCha20 to ChaCha8. */
#define FD_CHACHA_RNG_ALGO_CHACHA20 1
#define FD_CHACHA_RNG_ALGO_CHACHA8 2

/* FD_CHACHA_RNG_BUFSZ is the internal buffer size of pre-generated
ChaCha20 blocks. Multiple of block size (64 bytes) and a power of 2. */

Expand All @@ -49,6 +52,7 @@ struct __attribute__((aligned(32UL))) fd_chacha_rng_private {
Always aligned by FD_CHACHA_BLOCK_SZ */

int mode;
int algo;
};
typedef struct fd_chacha_rng_private fd_chacha_rng_t;

Expand Down Expand Up @@ -103,24 +107,25 @@ fd_chacha_rng_leave( fd_chacha_rng_t * );
void *
fd_chacha_rng_delete( void * shrng );

/* fd_chacha_rng{8,20}_init starts a ChaCha{8,20} RNG stream. rng is
assumed to be a current local join to a chacha20_rng object with no
/* fd_chacha_rng_init starts a ChaCha{8,20} RNG stream. rng is
assumed to be a current local join to a chacha_rng object with no
other concurrent operation that would modify the state while this is
executing. seed points to the first byte of the RNG seed byte vector
with 32 byte size. Any preexisting state for an in-progress or
recently completed calculation will be discarded. Returns rng (on
return, rng will have the state of a new in-progress calculation).

Compatible with Rust fn rand_chacha::ChaCha20Rng::from_seed
https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html#method.from_seed */
The param algo is expected to be FD_CHACHA_RNG_ALGO_CHACHA20 or
FD_CHACHA_RNG_ALGO_CHACH8. If invalid, it defaults to chacha20.

fd_chacha_rng_t *
fd_chacha8_rng_init( fd_chacha_rng_t * rng,
void const * key );
Compatible with Rust fn rand_chacha::ChaCha20Rng::from_seed
https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html#method.from_seed
(and ChaCha8Rng) */

fd_chacha_rng_t *
fd_chacha20_rng_init( fd_chacha_rng_t * rng,
void const * key );
fd_chacha_rng_init( fd_chacha_rng_t * rng,
void const * key,
int algo );

/* The refill function. Not part of the public API. */

Expand Down Expand Up @@ -155,39 +160,35 @@ fd_chacha_rng_avail( fd_chacha_rng_t const * rng ) {
return rng->buf_fill - rng->buf_off;
}

/* fd_chacha{8,20}_rng_ulong read a 64-bit integer in [0,2^64) from the
/* fd_chacha_rng_ulong read a 64-bit integer in [0,2^64) from the
RNG stream. */

static inline ulong
fd_chacha8_rng_ulong( fd_chacha_rng_t * rng ) {
if( FD_UNLIKELY( fd_chacha_rng_avail( rng ) < sizeof(ulong) ) )
fd_chacha8_rng_private_refill( rng );
ulong x = FD_LOAD( ulong, rng->buf + (rng->buf_off % FD_CHACHA_RNG_BUFSZ) );
rng->buf_off += 8U;
return x;
}

static inline ulong
fd_chacha20_rng_ulong( fd_chacha_rng_t * rng ) {
if( FD_UNLIKELY( fd_chacha_rng_avail( rng ) < sizeof(ulong) ) )
fd_chacha20_rng_private_refill( rng );
fd_chacha_rng_ulong( fd_chacha_rng_t * rng ) {
if( FD_UNLIKELY( fd_chacha_rng_avail( rng ) < sizeof(ulong) ) ) {
if( rng->algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
fd_chacha8_rng_private_refill( rng );
} else {
fd_chacha20_rng_private_refill( rng );
}
}
ulong x = FD_LOAD( ulong, rng->buf + (rng->buf_off % FD_CHACHA_RNG_BUFSZ) );
rng->buf_off += 8U;
return x;
}

/* fd_chacha20_rng_ulong_roll returns an uniform IID rand in [0,n)
/* fd_chacha_rng_ulong_roll returns an uniform IID rand in [0,n)
analogous to fd_rng_ulong_roll. Rejection method based using
fd_chacha20_rng_ulong.
fd_chacha_rng_ulong.

Compatible with Rust type
<rand_chacha::ChaCha20Rng as rand::Rng>::gen<rand::distributions::Uniform<u64>>()
as of version 0.7.0 of the crate
https://docs.rs/rand/latest/rand/distributions/struct.Uniform.html */

static inline ulong
fd_chacha20_rng_ulong_roll( fd_chacha_rng_t * rng,
ulong n ) {
fd_chacha_rng_ulong_roll( fd_chacha_rng_t * rng,
ulong n ) {
/* We use a pretty standard rejection-sampling based approach here,
but for future reference, here's an explanation:

Expand Down Expand Up @@ -231,16 +232,9 @@ fd_chacha20_rng_ulong_roll( fd_chacha_rng_t * rng,
(n << (63 - fd_ulong_find_msb( n ) )) - 1UL );

for( int i=0; 1; i++ ) {
ulong v = fd_chacha20_rng_ulong( rng );
#if FD_HAS_INT128
/* Compiles to one mulx instruction */
uint128 res = (uint128)v * (uint128)n;
ulong hi = (ulong)(res>>64);
ulong lo = (ulong) res;
#else
ulong v = fd_chacha_rng_ulong( rng );
ulong hi, lo;
fd_uwide_mul( &hi, &lo, v, n );
#endif

# if FD_CHACHA_RNG_DEBUG
FD_LOG_DEBUG(( "roll (attempt %d): n=%016lx zone: %016lx v=%016lx lo=%016lx hi=%016lx", i, n, zone, v, lo, hi ));
Expand Down
30 changes: 15 additions & 15 deletions src/ballet/chacha/test_chacha_rng.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,20 +31,20 @@ main( int argc,
0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f
};
FD_TEST( fd_chacha20_rng_init( rng, key ) );
FD_TEST( fd_chacha_rng_init( rng, key, FD_CHACHA_RNG_ALGO_CHACHA20 ) );

/* Test output */

FD_TEST( fd_chacha20_rng_ulong( rng )==0x6a19c5d97d2bfd39UL );
FD_TEST( fd_chacha_rng_ulong( rng )==0x6a19c5d97d2bfd39UL );
ulong x = 0UL;
for( ulong i=0UL; i<100000UL; i++ ) x ^= fd_chacha20_rng_ulong( rng );
for( ulong i=0UL; i<100000UL; i++ ) x ^= fd_chacha_rng_ulong( rng );
FD_TEST( x==0xb425be48c89d4f75UL );

#define RNG_TEST( name ) \
#define RNG_TEST( name, algo ) \
do { \
FD_LOG_NOTICE(( "Benchmarking " #name )); \
key[ 0 ]++; \
FD_TEST( fd_chacha20_rng_init( rng, key ) ); \
FD_TEST( fd_chacha_rng_init( rng, key, algo ) ); \
\
/* warmup */ \
for( ulong rem=1000000UL; rem; rem-- ) name( rng ); \
Expand All @@ -62,14 +62,14 @@ main( int argc,
FD_LOG_NOTICE(( " ~%6.3f ns / ulong", ns )); \
} while(0);

RNG_TEST( fd_chacha20_rng_ulong );
RNG_TEST( fd_chacha8_rng_ulong );
RNG_TEST( fd_chacha_rng_ulong, FD_CHACHA_RNG_ALGO_CHACHA8 );
RNG_TEST( fd_chacha_rng_ulong, FD_CHACHA_RNG_ALGO_CHACHA20 );

#define REFILL_TEST( name, stride ) \
#define REFILL_TEST( name, stride, algo ) \
do { \
FD_LOG_NOTICE(( "Benchmarking " #name )); \
key[ 0 ]++; \
FD_TEST( fd_chacha20_rng_init( rng, key ) ); \
FD_TEST( fd_chacha_rng_init( rng, key, algo ) ); \
\
/* warmup */ \
for( ulong rem=100000UL; rem; rem-- ) { \
Expand All @@ -90,15 +90,15 @@ main( int argc,
} while(0);

# if FD_HAS_AVX512
REFILL_TEST( fd_chacha8_rng_refill_avx512, 16*FD_CHACHA_BLOCK_SZ );
REFILL_TEST( fd_chacha20_rng_refill_avx512, 16*FD_CHACHA_BLOCK_SZ );
REFILL_TEST( fd_chacha8_rng_refill_avx512, 16*FD_CHACHA_BLOCK_SZ, FD_CHACHA_RNG_ALGO_CHACHA8 );
REFILL_TEST( fd_chacha20_rng_refill_avx512, 16*FD_CHACHA_BLOCK_SZ, FD_CHACHA_RNG_ALGO_CHACHA20 );
# endif
# if FD_HAS_AVX
REFILL_TEST( fd_chacha8_rng_refill_avx, 8*FD_CHACHA_BLOCK_SZ );
REFILL_TEST( fd_chacha20_rng_refill_avx, 8*FD_CHACHA_BLOCK_SZ );
REFILL_TEST( fd_chacha8_rng_refill_avx, 8*FD_CHACHA_BLOCK_SZ, FD_CHACHA_RNG_ALGO_CHACHA8 );
REFILL_TEST( fd_chacha20_rng_refill_avx, 8*FD_CHACHA_BLOCK_SZ, FD_CHACHA_RNG_ALGO_CHACHA20 );
# endif
REFILL_TEST( fd_chacha8_rng_refill_seq, 1*FD_CHACHA_BLOCK_SZ );
REFILL_TEST( fd_chacha20_rng_refill_seq, 1*FD_CHACHA_BLOCK_SZ );
REFILL_TEST( fd_chacha8_rng_refill_seq, 1*FD_CHACHA_BLOCK_SZ, FD_CHACHA_RNG_ALGO_CHACHA8 );
REFILL_TEST( fd_chacha20_rng_refill_seq, 1*FD_CHACHA_BLOCK_SZ, FD_CHACHA_RNG_ALGO_CHACHA20 );

/* Test leave/delete */

Expand Down
Loading
Loading