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

Author: 0x0ece

src/ballet/chacha/fd_chacha.h

@@ -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_rng.c

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

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

src/ballet/wsample/fd_wsample.c

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

@@ -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/test_wsample.c

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

src/disco/shred/Local.mk

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