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

Author: 0x0ece

agave

@@ -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
 

src/ballet/chacha/fd_chacha.h

@@ -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 );

src/ballet/chacha/fd_chacha_rng.c

@@ -252,6 +252,35 @@ fd_chacha20_rng_ulong_roll( fd_chacha_rng_t * rng,
   }
 }
 
+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 );
+#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 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 ));
+#   else
+    (void)i;
+#   endif /* FD_CHACHA_RNG_DEBUG */
+
+    if( FD_LIKELY( lo<=zone ) ) return hi;
+  }
+}
+
 FD_PROTOTYPES_END
 
 #endif /* HEADER_fd_src_ballet_chacha20_fd_chacha20_rng_h */

src/ballet/chacha/fd_chacha_rng.h

@@ -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;
 
@@ -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) );
@@ -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 ) {
@@ -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
@@ -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;
@@ -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 );
 }
@@ -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;

src/ballet/wsample/fd_wsample.c

@@ -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) */
@@ -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
@@ -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 */

src/ballet/wsample/fd_wsample.h

@@ -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 );
@@ -219,7 +219,7 @@ test_matches_solana( void ) {
   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 );

src/ballet/wsample/test_wsample.c

@@ -21,7 +21,7 @@ static const fd_pubkey_t null_pubkey = {{ 0 }};
 #include "../../util/tmpl/fd_map_dynamic.c"
 
 
-/* This 45 byte struct gets hashed to compute the seed for Chacha20 to
+/* This 45 byte struct gets hashed to compute the seed for ChaCha to
    compute the shred destinations. */
 struct __attribute__((packed)) shred_dest_input {
   ulong slot;
@@ -150,14 +150,19 @@ void * fd_shred_dest_delete( void * mem ) {
   return mem;
 }
 
+static inline ulong
+fd_sdest_rng_ulong_roll( fd_shred_dest_t * sdest, ulong n ) {
+  return fd_wsample_rng_ulong_roll( sdest->staked, n );
+}
+
 /* sample_unstaked, sample_unstaked_noprepare, and
    prepare_unstaked_sampling are used to perform the specific form of
    unweighted random sampling that Solana uses for unstaked validators.
    In essence, you:
     1. construct a list of all the unstaked validators,
     2. delete the leader (if present)
     then repeatedly:
-    3. choose the chacha20rng_roll( |unstaked| )th element.
+    3. choose the chacha_rng_roll( |unstaked| )th element.
     4. swap the last element in unstaked with the chosen element
     5. return and remove the chosen element (which is now in the last
     position, so remove is O(1)).
@@ -191,7 +196,7 @@ sample_unstaked_noprepare( fd_shred_dest_t  * sdest,
   ulong unstaked_cnt = sdest->unstaked_cnt - (ulong)remove_in_interval;
   if( FD_UNLIKELY( unstaked_cnt==0UL ) ) return FD_WSAMPLE_EMPTY;
 
-  ulong sample = sdest->staked_cnt + fd_chacha20_rng_ulong_roll( sdest->rng, unstaked_cnt );
+  ulong sample = sdest->staked_cnt + fd_sdest_rng_ulong_roll( sdest, unstaked_cnt );
   return fd_ulong_if( (!remove_in_interval) | (sample<remove_idx), sample, sample+1UL );
 }
 
@@ -219,7 +224,7 @@ static inline ulong
 sample_unstaked( fd_shred_dest_t * sdest ) {
   if( FD_UNLIKELY( sdest->unstaked_unremoved_cnt==0UL ) ) return FD_WSAMPLE_EMPTY;
 
-  ulong sample = fd_chacha20_rng_ulong_roll( sdest->rng, sdest->unstaked_unremoved_cnt );
+  ulong sample = fd_sdest_rng_ulong_roll( sdest, sdest->unstaked_unremoved_cnt );
   ulong to_return = sdest->unstaked[sample];
   sdest->unstaked[sample] = sdest->unstaked[--sdest->unstaked_unremoved_cnt];
   return to_return;
@@ -261,7 +266,8 @@ fd_shred_dest_idx_t *
 fd_shred_dest_compute_first( fd_shred_dest_t          * sdest,
                              fd_shred_t const * const * input_shreds,
                              ulong                      shred_cnt,
-                             fd_shred_dest_idx_t      * out ) {
+                             fd_shred_dest_idx_t      * out,
+                             int                        use_chacha8 ) {
 
   if( FD_UNLIKELY( shred_cnt==0UL ) ) return out;
 
@@ -290,7 +296,7 @@ fd_shred_dest_compute_first( fd_shred_dest_t          * sdest,
 
   int any_staked_candidates = sdest->staked_cnt > (ulong)source_validator_is_staked;
   for( ulong i=0UL; i<shred_cnt; i++ ) {
-    fd_wsample_seed_rng( fd_wsample_get_rng( sdest->staked ), dest_hash_outputs[ i ] );
+    fd_wsample_seed_rng( sdest->staked, dest_hash_outputs[ i ], use_chacha8 );
     /* Map FD_WSAMPLE_INDETERMINATE to FD_SHRED_DEST_NO_DEST */
     if( FD_LIKELY( any_staked_candidates ) ) out[i] = (fd_shred_dest_idx_t)fd_ulong_min( fd_wsample_sample( sdest->staked ), FD_SHRED_DEST_NO_DEST );
     else                                     out[i] = (fd_shred_dest_idx_t)sample_unstaked_noprepare( sdest, sdest->source_validator_orig_idx );
@@ -308,7 +314,8 @@ fd_shred_dest_compute_children( fd_shred_dest_t          * sdest,
                                 ulong                      out_stride,
                                 ulong                      fanout,
                                 ulong                      dest_cnt,
-                                ulong                    * opt_max_dest_cnt ) {
+                                ulong                    * opt_max_dest_cnt,
+                                int                        use_chacha8 ) {
 
   /* The logic here is a little tricky since we are keeping track of
      staked and unstaked separately and only logically concatenating
@@ -358,7 +365,7 @@ fd_shred_dest_compute_children( fd_shred_dest_t          * sdest,
     if( FD_LIKELY( query && leader_is_staked ) ) fd_wsample_remove_idx( sdest->staked, leader_idx );
 
     ulong my_idx         = 0UL;
-    fd_wsample_seed_rng( fd_wsample_get_rng( sdest->staked ), dest_hash_outputs[ i ] ); /* Seeds both samplers since the rng is shared */
+    fd_wsample_seed_rng( sdest->staked, dest_hash_outputs[ i ], use_chacha8 ); /* Seeds both samplers since the rng is shared */
 
     if( FD_UNLIKELY( !i_am_staked ) ) {
       /* If there's excluded stake, we don't know about any unstaked

src/disco/shred/fd_shred_dest.c

@@ -161,7 +161,8 @@ fd_shred_dest_idx_t *
 fd_shred_dest_compute_first( fd_shred_dest_t          * sdest,
                              fd_shred_t const * const * input_shreds,
                              ulong                      shred_cnt,
-                             fd_shred_dest_idx_t      * out );
+                             fd_shred_dest_idx_t      * out,
+                             int                        use_chacha8 );
 
 /* fd_shred_dest_compute_children computes the source validator's
    children in the Turbine tree for each of the provided shreds.
@@ -205,7 +206,8 @@ fd_shred_dest_compute_children( fd_shred_dest_t          * sdest,
                                 ulong                      out_stride,
                                 ulong                      fanout,
                                 ulong                      dest_cnt,
-                                ulong                    * opt_max_dest_cnt );
+                                ulong                    * opt_max_dest_cnt,
+                                int                        use_chacha8 );
 
 /* fd_shred_dest_idx_to_dest maps a destination index (as produced by
    fd_shred_dest_compute_children or fd_shred_dest_compute_first) to an