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
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
12 changes: 11 additions & 1 deletion src/ballet/chacha/fd_chacha_rng.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,10 +58,20 @@ fd_chacha_rng_delete( void * shrng ) {
return shrng;
}

fd_chacha_rng_t *
fd_chacha8_rng_init( fd_chacha_rng_t * rng,
void const * key ) {
memcpy( rng->key, key, FD_CHACHA_KEY_SZ );
rng->buf_off = 0UL;
rng->buf_fill = 0UL;
fd_chacha8_rng_private_refill( rng );
return rng;
}

fd_chacha_rng_t *
fd_chacha20_rng_init( fd_chacha_rng_t * rng,
void const * key ) {
memcpy( rng->key, key, FD_CHACHA20_KEY_SZ );
memcpy( rng->key, key, FD_CHACHA_KEY_SZ );
rng->buf_off = 0UL;
rng->buf_fill = 0UL;
fd_chacha20_rng_private_refill( rng );
Expand Down
33 changes: 23 additions & 10 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 Down Expand Up @@ -231,16 +229,31 @@ 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_chacha20_rng_ulong( rng );
ulong hi, lo;
fd_uwide_mul( &hi, &lo, v, n );

# 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 ));
# else
(void)i;
# endif /* FD_CHACHA_RNG_DEBUG */

if( FD_LIKELY( lo<=zone ) ) return hi;
}
}

static inline ulong
fd_chacha8_rng_ulong_roll( fd_chacha_rng_t * rng,
ulong n ) {
ulong const zone = fd_ulong_if( rng->mode==FD_CHACHA_RNG_MODE_MOD,
ULONG_MAX - (ULONG_MAX-n+1UL)%n,
(n << (63 - fd_ulong_find_msb( n ) )) - 1UL );

for( int i=0; 1; i++ ) {
ulong v = fd_chacha8_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
37 changes: 23 additions & 14 deletions src/ballet/wsample/fd_wsample.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,8 @@ struct __attribute__((aligned(64UL))) fd_wsample_private {
uint height;
char restore_enabled;
char poisoned_mode;
/* Two bytes of padding here */
uchar rng_algo;
/* One byte of padding here */

fd_chacha_rng_t * rng;

Expand Down Expand Up @@ -302,6 +303,7 @@ fd_wsample_new_init( void * shmem,
sampler->height = (uint)height;
sampler->restore_enabled = (char)!!restore_enabled;
sampler->poisoned_mode = 0;
sampler->rng_algo = 0;
sampler->rng = rng;

fd_memset( sampler->tree, (char)0, internal_cnt*sizeof(tree_ele_t) );
Expand Down Expand Up @@ -389,18 +391,25 @@ fd_wsample_delete( void * shmem ) {
return shmem;
}



fd_chacha_rng_t * fd_wsample_get_rng( fd_wsample_t * sampler ) { return sampler->rng; }


/* TODO: Should this function exist at all? */
void
fd_wsample_seed_rng( fd_chacha_rng_t * rng,
uchar seed[static 32] ) {
fd_chacha20_rng_init( rng, seed );
fd_wsample_seed_rng( fd_wsample_t * sampler,
uchar seed[ 32 ],
int use_chacha8 ) {
sampler->rng_algo = fd_uchar_if( use_chacha8, FD_WSAMPLE_RNG_CHACHA8, FD_WSAMPLE_RNG_CHACHA20 );
if( FD_UNLIKELY( sampler->rng_algo==FD_WSAMPLE_RNG_CHACHA8 ) ) {
fd_chacha8_rng_init( sampler->rng, seed );
} else {
fd_chacha20_rng_init( sampler->rng, seed );
}
}

ulong
fd_wsample_rng_ulong_roll( fd_wsample_t * sampler, ulong n ) {
if( FD_UNLIKELY( sampler->rng_algo==FD_WSAMPLE_RNG_CHACHA8 ) ) {
return fd_chacha8_rng_ulong_roll( sampler->rng, n );
}
return fd_chacha20_rng_ulong_roll( sampler->rng, n );
}

fd_wsample_t *
fd_wsample_restore_all( fd_wsample_t * sampler ) {
Expand Down Expand Up @@ -597,7 +606,7 @@ fd_wsample_remove_idx( fd_wsample_t * sampler,
but operations with mask registers are frustratingly slow. Instead,
we first define x'=x+1, so then v0<=x is equivalent to v0-x'<0, or
whether v0-x' has the high bit set. Because of
fd_chacha20_rng_ulong_roll's contract, we know x<ULONG_MAX, so forming
fd_wsample_rng_ulong_roll's contract, we know x<ULONG_MAX, so forming
x' is safe. We then use a _mm512_permutexvar_epi8 to essentially
broadcast the high byte from each of the ulongs (really we just need
the high bit) to the whole vector. A popcnt gives us our base
Expand Down Expand Up @@ -734,7 +743,7 @@ fd_wsample_sample_and_remove_many( fd_wsample_t * sampler,
for( ulong i=0UL; i<cnt; i++ ) {
if( FD_UNLIKELY( !sampler->unremoved_weight ) ) { idxs[ i ] = FD_WSAMPLE_EMPTY; continue; }
if( FD_UNLIKELY( sampler->poisoned_mode ) ) { idxs[ i ] = FD_WSAMPLE_INDETERMINATE; continue; }
ulong unif = fd_chacha20_rng_ulong_roll( sampler->rng, sampler->unremoved_weight+sampler->poisoned_weight );
ulong unif = fd_wsample_rng_ulong_roll( sampler, sampler->unremoved_weight+sampler->poisoned_weight );
if( FD_UNLIKELY( unif>=sampler->unremoved_weight ) ) {
idxs[ i ] = FD_WSAMPLE_INDETERMINATE;
sampler->poisoned_mode = 1;
Expand Down Expand Up @@ -771,7 +780,7 @@ ulong
fd_wsample_sample( fd_wsample_t * sampler ) {
if( FD_UNLIKELY( !sampler->unremoved_weight ) ) return FD_WSAMPLE_EMPTY;
if( FD_UNLIKELY( sampler->poisoned_mode ) ) return FD_WSAMPLE_INDETERMINATE;
ulong unif = fd_chacha20_rng_ulong_roll( sampler->rng, sampler->unremoved_weight+sampler->poisoned_weight );
ulong unif = fd_wsample_rng_ulong_roll( sampler, sampler->unremoved_weight+sampler->poisoned_weight );
if( FD_UNLIKELY( unif>=sampler->unremoved_weight ) ) return FD_WSAMPLE_INDETERMINATE;
return (ulong)fd_wsample_map_sample( sampler, unif );
}
Expand All @@ -780,7 +789,7 @@ ulong
fd_wsample_sample_and_remove( fd_wsample_t * sampler ) {
if( FD_UNLIKELY( !sampler->unremoved_weight ) ) return FD_WSAMPLE_EMPTY;
if( FD_UNLIKELY( sampler->poisoned_mode ) ) return FD_WSAMPLE_INDETERMINATE;
ulong unif = fd_chacha20_rng_ulong_roll( sampler->rng, sampler->unremoved_weight+sampler->poisoned_weight );
ulong unif = fd_wsample_rng_ulong_roll( sampler, sampler->unremoved_weight+sampler->poisoned_weight );
if( FD_UNLIKELY( unif>=sampler->unremoved_weight ) ) {
sampler->poisoned_mode = 1;
return FD_WSAMPLE_INDETERMINATE;
Expand Down
12 changes: 9 additions & 3 deletions src/ballet/wsample/fd_wsample.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@
struct fd_wsample_private;
typedef struct fd_wsample_private fd_wsample_t;

#define FD_WSAMPLE_RNG_CHACHA20 (0U)
#define FD_WSAMPLE_RNG_CHACHA8 (1U)

#define FD_WSAMPLE_ALIGN (64UL)
/* fd_leaders really wants a compile time-compatible footprint... The
internal count is 1/8 * (9^ceil(log_9(ele_cnt)) - 1) */
Expand Down Expand Up @@ -135,10 +138,10 @@ void * fd_wsample_new_fini( void * shmem, ulong poisoned_weight );
/* fd_wsample_get_rng returns the value provided for rng in new. */
fd_chacha_rng_t * fd_wsample_get_rng( fd_wsample_t * sampler );

/* fd_wsample_seed_rng seeds the ChaCha20 rng with the provided seed in
/* fd_wsample_seed_rng seeds the ChaCha rng with the provided seed in
preparation for sampling. This function is compatible with Solana's
ChaChaRng::from_seed. */
void fd_wsample_seed_rng( fd_chacha_rng_t * rng, uchar seed[static 32] );
void fd_wsample_seed_rng( fd_wsample_t * sampler, uchar seed[ 32 ], int use_chacha8 );

/* fd_wsample_sample{_and_remove}{,_many} produces one or cnt (in the
_many case) weighted random samples from the sampler. If the
Expand Down Expand Up @@ -191,6 +194,9 @@ void fd_wsample_remove_idx( fd_wsample_t * sampler, ulong idx );
in which case no elements are restored. */
fd_wsample_t * fd_wsample_restore_all( fd_wsample_t * sampler );


/* fd_wsample_rng_ulong_roll returns an uniform IID rand in [0,n)
analogous to fd_rng_ulong_roll. Internally it uses chacha8 or
chacha20, based on how the wsample was initialized. */
ulong fd_wsample_rng_ulong_roll( fd_wsample_t * sampler, ulong n );

#endif /* HEADER_fd_src_ballet_wsample_fd_wsample_h */
68 changes: 63 additions & 5 deletions src/ballet/wsample/test_wsample.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ test_matches_solana( void ) {

void * partial = fd_wsample_new_init( _shmem, rng, 2UL, 0, FD_WSAMPLE_HINT_FLAT );
fd_wsample_t * tree = fd_wsample_join( fd_wsample_new_fini( fd_wsample_new_add( fd_wsample_new_add( partial, 2UL ), 1UL ), 0UL ) );
fd_wsample_seed_rng( fd_wsample_get_rng( tree ), zero_seed );
fd_wsample_seed_rng( tree, zero_seed, 0 /* use_chacha8 */ );

FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
Expand All @@ -214,13 +214,10 @@ test_matches_solana( void ) {
ulong weights2[18] = { 78, 70, 38, 27, 21, 82, 42, 21, 77, 77, 17, 4, 50, 96, 83, 33, 16, 72 };

memset( zero_seed, 48, 32UL );
fd_chacha20_rng_init( rng, zero_seed );

partial = fd_wsample_new_init( _shmem, rng, 18UL, 0, FD_WSAMPLE_HINT_FLAT );
for( ulong i=0UL; i<18UL; i++ ) partial = fd_wsample_new_add( partial, weights2[i] );
tree = fd_wsample_join( fd_wsample_new_fini( partial, 0UL ) );
fd_wsample_seed_rng( fd_wsample_get_rng( tree ), zero_seed );

fd_wsample_seed_rng( tree, zero_seed, 0 /* use_chacha8 */ );

FD_TEST( fd_wsample_sample_and_remove( tree ) == 9UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 3UL );
Expand All @@ -245,6 +242,66 @@ test_matches_solana( void ) {
fd_chacha_rng_delete( fd_chacha_rng_leave( rng ) );
}

static void
test_matches_solana_chacha8( void ) {
/* Adopted from test_repeated_leader_schedule_specific: */
fd_chacha_rng_t _rng[1];
fd_chacha_rng_t * rng = fd_chacha_rng_join( fd_chacha_rng_new( _rng, FD_CHACHA_RNG_MODE_MOD ) );
uchar zero_seed[32] = {0};

void * partial = fd_wsample_new_init( _shmem, rng, 2UL, 0, FD_WSAMPLE_HINT_FLAT );
fd_wsample_t * tree = fd_wsample_join( fd_wsample_new_fini( fd_wsample_new_add( fd_wsample_new_add( partial, 2UL ), 1UL ), 0UL ) );
fd_wsample_seed_rng( tree, zero_seed, 1 /* use_chacha8 */ );

FD_TEST( fd_wsample_sample( tree ) == 1UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );
FD_TEST( fd_wsample_sample( tree ) == 0UL );

fd_wsample_delete( fd_wsample_leave( tree ) );
fd_chacha_rng_delete( fd_chacha_rng_leave( rng ) );

rng = fd_chacha_rng_join( fd_chacha_rng_new( _rng, FD_CHACHA_RNG_MODE_SHIFT ) );

/* Adopted from test_weighted_shuffle_hard_coded, except they handle
the special case for 0 weights inside their WeightedShuffle object,
and the test case initially used i32 as weights, which made their
Chacha20 object generate i32s instead of u64s. */
ulong weights2[18] = { 78, 70, 38, 27, 21, 82, 42, 21, 77, 77, 17, 4, 50, 96, 83, 33, 16, 72 };

memset( zero_seed, 48, 32UL );
partial = fd_wsample_new_init( _shmem, rng, 18UL, 0, FD_WSAMPLE_HINT_FLAT );
for( ulong i=0UL; i<18UL; i++ ) partial = fd_wsample_new_add( partial, weights2[i] );
tree = fd_wsample_join( fd_wsample_new_fini( partial, 0UL ) );
fd_wsample_seed_rng( tree, zero_seed, 1 /* use_chacha8 */ );

FD_TEST( fd_wsample_sample_and_remove( tree ) == 13UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 8UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 6UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 14UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 0UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 17UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 1UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 12UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 3UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 16UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 5UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 15UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 9UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 2UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 4UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 7UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 10UL );
FD_TEST( fd_wsample_sample_and_remove( tree ) == 11UL );

fd_wsample_delete( fd_wsample_leave( tree ) );
fd_chacha_rng_delete( fd_chacha_rng_leave( rng ) );
}

static void
test_sharing( void ) {
fd_chacha_rng_t _rng[1];
Expand Down Expand Up @@ -472,6 +529,7 @@ main( int argc,
FD_TEST( fd_wsample_footprint( MAX, 1 )<MAX_FOOTPRINT );

test_matches_solana();
test_matches_solana_chacha8();
test_map();
test_sharing();
test_restore_disabled();
Expand Down
2 changes: 2 additions & 0 deletions src/disco/shred/Local.mk
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,11 @@ ifdef FD_HAS_ALLOCA
$(call add-objs,fd_shred_tile,fd_disco)
endif
$(call make-unit-test,test_shred_dest,test_shred_dest,fd_disco fd_flamenco fd_ballet fd_util)
$(call make-unit-test,test_shred_dest_conformance,test_shred_dest_conformance,fd_disco fd_flamenco fd_ballet fd_util)
$(call make-unit-test,test_fec_resolver,test_fec_resolver,fd_flamenco fd_disco fd_ballet fd_util fd_tango fd_reedsol)
$(call make-unit-test,test_stake_ci,test_stake_ci,fd_disco fd_flamenco fd_ballet fd_util fd_tango fd_reedsol)
$(call run-unit-test,test_shred_dest,)
$(call run-unit-test,test_shred_dest_conformance,)
$(call run-unit-test,test_fec_resolver,)
$(call run-unit-test,test_stake_ci,)
ifdef FD_HAS_HOSTED
Expand Down
Loading
Loading