Line data    Source code

       1             : #ifndef HEADER_fd_src_ballet_chacha_fd_chacha_rng_h
       2             : #define HEADER_fd_src_ballet_chacha_fd_chacha_rng_h
       3             : 
       4             : /* fd_chacha_rng provides APIs for ChaCha-based RNG, as used in the
       5             :    Solana protocol.  This API should only be used where necessary.
       6             :    fd_rng is a better choice in all other cases. */
       7             : 
       8             : #include "fd_chacha.h"
       9             : #include "../../util/bits/fd_uwide.h"
      10             : 
      11             : /* FD_CHACHA_RNG_DEBUG controls debug logging.  0 is off; 1 is on. */
      12             : 
      13             : #ifndef FD_CHACHA_RNG_DEBUG
      14             : #define FD_CHACHA_RNG_DEBUG 0
      15             : #endif
      16             : 
      17             : /* Solana uses different mechanisms of mapping a ulong to an unbiased
      18             :    integer in [0, n) in different parts of the code.  In particular,
      19             :    leader schedule generation uses MODE_MOD and Turbine uses MODE_SHIFT.
      20             :    See the note in fd_chacha_rng_ulong_roll for more details. */
      21  1030743717 : #define FD_CHACHA_RNG_MODE_MOD   1
      22      284361 : #define FD_CHACHA_RNG_MODE_SHIFT 2
      23             : 
      24             : /* Solana uses different ChaCha algorithms. Leader schedule generation
      25             :    uses ChaCha20 and Turbine switched from ChaCha20 to ChaCha8. */
      26      310887 : #define FD_CHACHA_RNG_ALGO_CHACHA20   1
      27    20859340 : #define FD_CHACHA_RNG_ALGO_CHACHA8    2
      28             : 
      29             : /* FD_CHACHA_RNG_BUFSZ is the internal buffer size of pre-generated
      30             :    ChaCha20 blocks.  Multiple of block size (64 bytes) and a power of 2. */
      31             : 
      32             : #if FD_HAS_AVX512
      33     4406819 : #define FD_CHACHA_RNG_BUFSZ (16*FD_CHACHA_BLOCK_SZ)
      34             : #elif FD_HAS_AVX
      35     8813590 : #define FD_CHACHA_RNG_BUFSZ (8*FD_CHACHA_BLOCK_SZ)
      36             : #elif FD_HAS_SVE2
      37             : #define FD_CHACHA_RNG_BUFSZ (4*FD_CHACHA_BLOCK_SZ)
      38             : #else
      39             : #define FD_CHACHA_RNG_BUFSZ (256UL)
      40             : #endif
      41             : 
      42             : struct __attribute__((aligned(32UL))) fd_chacha_rng_private {
      43             :   /* ChaCha20 encryption key */
      44             :   uchar key[ 32UL ] __attribute__((aligned(32UL)));
      45             : 
      46             :   /* Ring buffer of pre-generated ChaCha20 RNG data.
      47             :      Note: We currently assume all reads are 8 byte.  This means the
      48             :            cursor is always aligned by 8 and strictly increases in
      49             :            increments of 8.  Thus, we really only have to refill the
      50             :            buffer if buf_off==buf_fill.  */
      51             :   uchar buf[ FD_CHACHA_RNG_BUFSZ ] __attribute__((aligned(FD_CHACHA_BLOCK_SZ)));
      52             :   ulong buf_off;   /* Total number of bytes consumed */
      53             :   ulong buf_fill;  /* Total number of bytes produced
      54             :                       Always aligned by FD_CHACHA_BLOCK_SZ */
      55             : 
      56             :   int mode;
      57             :   int algo;
      58             : };
      59             : typedef struct fd_chacha_rng_private fd_chacha_rng_t;
      60             : 
      61             : FD_PROTOTYPES_BEGIN
      62             : 
      63             : /* fd_chacha_rng_{align,footprint} give the needed alignment and
      64             :    footprint of a memory region suitable to hold a ChaCha20-based RNG.
      65             : 
      66             :    fd_chacha_rng_new formats a memory region with suitable alignment
      67             :    and footprint for holding a chacha20_rng object.  Assumes shmem
      68             :    points on the caller to the first byte of the memory region owned by
      69             :    the caller to use.  `mode` must be one of the FD_CHACHA_RNG_MODE_*
      70             :    constants defined above and dictates what mode this object will use
      71             :    to generate random numbers. Returns shmem on success and NULL on
      72             :    failure (logs details).  The memory region will be owned by the
      73             :    object on successful return.  The caller is not joined on return.
      74             : 
      75             :    fd_chacha_rng_join joins the caller to a chacha20_rng object.
      76             :    Assumes shrng points to the first byte of the memory region holding
      77             :    the object.  Returns a local handle to the join on success (this is
      78             :    not necessarily a simple cast of the address) and NULL on failure
      79             :    (logs details).
      80             : 
      81             :    fd_chacha_rng_leave leaves the caller's current local join to a
      82             :    ChaCha20 RNG object.  Returns a pointer to the memory region holding
      83             :    the object on success this is not necessarily a simple cast of the
      84             :    address) and NULL on failure (logs details).  The caller is not
      85             :    joined on successful return.
      86             : 
      87             :    fd_chacha_rng_delete unformats a memory region that holds a ChaCha20
      88             :    RNG object.  Assumes shrng points on the caller to the first byte of
      89             :    the memory region holding the state and that nobody is joined.
      90             :    Returns a pointer to the memory region on success and NULL on failure
      91             :    (logs details).  The caller has ownership of the memory region on
      92             :    successful return. */
      93             : 
      94             : FD_FN_CONST ulong
      95             : fd_chacha_rng_align( void );
      96             : 
      97             : FD_FN_CONST ulong
      98             : fd_chacha_rng_footprint( void );
      99             : 
     100             : void *
     101             : fd_chacha_rng_new( void * shmem, int mode );
     102             : 
     103             : fd_chacha_rng_t *
     104             : fd_chacha_rng_join( void * shrng );
     105             : 
     106             : void *
     107             : fd_chacha_rng_leave( fd_chacha_rng_t * );
     108             : 
     109             : void *
     110             : fd_chacha_rng_delete( void * shrng );
     111             : 
     112             : /* fd_chacha_rng_init starts a ChaCha{8,20} RNG stream.  rng is
     113             :    assumed to be a current local join to a chacha_rng object with no
     114             :    other concurrent operation that would modify the state while this is
     115             :    executing.  seed points to the first byte of the RNG seed byte vector
     116             :    with 32 byte size.  Any preexisting state for an in-progress or
     117             :    recently completed calculation will be discarded.  Returns rng (on
     118             :    return, rng will have the state of a new in-progress calculation).
     119             : 
     120             :    The param algo is expected to be FD_CHACHA_RNG_ALGO_CHACHA20 or
     121             :    FD_CHACHA_RNG_ALGO_CHACH8. If invalid, it defaults to chacha20.
     122             : 
     123             :    Compatible with Rust fn rand_chacha::ChaCha20Rng::from_seed
     124             :    https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html#method.from_seed
     125             :    (and ChaCha8Rng) */
     126             : 
     127             : fd_chacha_rng_t *
     128             : fd_chacha_rng_init( fd_chacha_rng_t * rng,
     129             :                     void const *      key,
     130             :                     int               algo );
     131             : 
     132             : /* The refill function.  Not part of the public API. */
     133             : 
     134             : #if FD_HAS_AVX512
     135             : void fd_chacha8_rng_refill_avx512 ( fd_chacha_rng_t * rng );
     136             : void fd_chacha20_rng_refill_avx512( fd_chacha_rng_t * rng );
     137             : #endif
     138             : 
     139             : #if FD_HAS_AVX
     140             : void fd_chacha8_rng_refill_avx ( fd_chacha_rng_t * rng );
     141             : void fd_chacha20_rng_refill_avx( fd_chacha_rng_t * rng );
     142             : #endif
     143             : 
     144             : #if FD_HAS_SVE2
     145             : void fd_chacha8_rng_refill_sve2 ( fd_chacha_rng_t * rng );
     146             : void fd_chacha20_rng_refill_sve2( fd_chacha_rng_t * rng );
     147             : #endif
     148             : 
     149             : void fd_chacha8_rng_refill_seq ( fd_chacha_rng_t * rng );
     150             : void fd_chacha20_rng_refill_seq( fd_chacha_rng_t * rng );
     151             : 
     152             : #if FD_HAS_AVX512
     153     1446797 : #define fd_chacha8_rng_private_refill  fd_chacha8_rng_refill_avx512
     154     2570823 : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_avx512
     155             : #elif FD_HAS_AVX
     156     4788712 : #define fd_chacha8_rng_private_refill  fd_chacha8_rng_refill_avx
     157    10275754 : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_avx
     158             : #elif FD_HAS_SVE2
     159             : #define fd_chacha8_rng_private_refill  fd_chacha8_rng_refill_sve2
     160             : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_sve2
     161             : #else
     162             : #define fd_chacha8_rng_private_refill  fd_chacha8_rng_refill_seq
     163             : #define fd_chacha20_rng_private_refill fd_chacha20_rng_refill_seq
     164             : #endif
     165             : 
     166             : /* fd_chacha_rng_avail returns the number of buffered bytes. */
     167             : 
     168             : FD_FN_PURE static inline ulong
     169  1376708593 : fd_chacha_rng_avail( fd_chacha_rng_t const * rng ) {
     170  1376708593 :   return rng->buf_fill - rng->buf_off;
     171  1376708593 : }
     172             : 
     173             : /* fd_chacha_rng_ulong read a 64-bit integer in [0,2^64) from the
     174             :    RNG stream. */
     175             : 
     176             : static inline ulong
     177  1376688210 : fd_chacha_rng_ulong( fd_chacha_rng_t * rng ) {
     178  1376688210 :   if( FD_UNLIKELY( fd_chacha_rng_avail( rng ) < sizeof(ulong) ) ) {
     179    17289271 :     if( rng->algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
     180     4454321 :       fd_chacha8_rng_private_refill( rng );
     181    12834950 :     } else {
     182    12834950 :       fd_chacha20_rng_private_refill( rng );
     183    12834950 :     }
     184    17289271 :   }
     185  1376688210 :   ulong x = FD_LOAD( ulong, rng->buf + (rng->buf_off % FD_CHACHA_RNG_BUFSZ) );
     186  1376688210 :   rng->buf_off += 8U;
     187  1376688210 :   return x;
     188  1376688210 : }
     189             : 
     190             : /* fd_chacha_rng_read32 reads 32 bytes from the RNG stream into buf.
     191             :    buf may be unaligned. */
     192             : 
     193             : static inline void
     194             : fd_chacha_rng_read32( fd_chacha_rng_t * rng,
     195       20383 :                       void *            buf ) {
     196       20383 :   rng->buf_off = fd_ulong_align_up( rng->buf_off, 32UL );
     197       20383 :   if( FD_UNLIKELY( fd_chacha_rng_avail( rng )<32UL ) ) {
     198          39 :     if( rng->algo==FD_CHACHA_RNG_ALGO_CHACHA8 ) {
     199          39 :       fd_chacha8_rng_private_refill( rng );
     200          39 :     } else {
     201           0 :       fd_chacha20_rng_private_refill( rng );
     202           0 :     }
     203          39 :   }
     204             : 
     205       20383 :   fd_memcpy( buf, rng->buf + (rng->buf_off % FD_CHACHA_RNG_BUFSZ), 32UL );
     206       20383 :   rng->buf_off += 32UL;
     207       20383 : }
     208             : 
     209             : /* fd_chacha_rng_ulong_roll returns an uniform IID rand in [0,n)
     210             :    analogous to fd_rng_ulong_roll.  Rejection method based using
     211             :    fd_chacha_rng_ulong.
     212             : 
     213             :    Compatible with Rust type
     214             :    <rand_chacha::ChaCha20Rng as rand::Rng>::gen<rand::distributions::Uniform<u64>>()
     215             :    as of version 0.7.0 of the crate
     216             :    https://docs.rs/rand/latest/rand/distributions/struct.Uniform.html */
     217             : 
     218             : static inline ulong
     219             : fd_chacha_rng_ulong_roll( fd_chacha_rng_t * rng,
     220  1030732890 :                           ulong             n ) {
     221             :   /* We use a pretty standard rejection-sampling based approach here,
     222             :      but for future reference, here's an explanation:
     223             : 
     224             :      We know that v can take 2^64 values, and so any method that maps
     225             :      each of the 2^64 values to the range directly [0, n) will not be
     226             :      uniform distribution when 2^64 is not divisible by n.  This
     227             :      motivates using rejection sampling.
     228             : 
     229             :      The most basic approach is to map v from [0, n*floor(2^64/n) ) to
     230             :      [0, n) using v%n, but that puts a modulus on the critical path.  To
     231             :      avoid that, the Rust rand crate uses a different approach: compute
     232             :      v*n/2^64, which is also in [0, n).
     233             : 
     234             :      Now the question to answer is which values to throw out.  We pick a
     235             :      large integer k such that k*n<=2^64 and map [0, k*n) -> 0, [2^64,
     236             :      2^64+k*n) -> 1, etc.  Since k*n might be 2^64 and then not fit in a
     237             :      long, we define zone=k*n-1 <= ULONG_MAX, and make the intervals
     238             :      closed instead of half-open.
     239             : 
     240             :      Here's where the mode comes in.  Depending on what method you call
     241             :      and what datatype you use, the Rust crate uses different values of
     242             :      k.  When MODE_MOD is set, we use largest possible value of k,
     243             :      namely floor(2^64/n).  You can compute zone directly as follows:
     244             :                zone  = k*n-1
     245             :                      = floor(2^64/n)*n - 1
     246             :                      = 2^64 - (2^64%n) - 1
     247             :                      = 2^64-1 - (2^64-n)%n, since n<2^64
     248             :                      = 2^64-1 - ((2^64-1)-n+1)%n
     249             :      Which is back to having a mod... But at least if n is a
     250             :      compile-time constant then the whole zone computation becomes a
     251             :      compile-time constant.
     252             : 
     253             :      When MODE_SHIFT is set, we use uses almost the largest possible
     254             :      power of two for k.  Precisely, it uses the smallest power of two
     255             :      such that k*n >= 2^63, which is the largest power of two such that
     256             :      k*n<=2^64 unless n is a power of two.  This approach eliminates the
     257             :      mod calculation but increases the expected number of samples
     258             :      required. */
     259  1030732890 :   ulong const zone = fd_ulong_if( rng->mode==FD_CHACHA_RNG_MODE_MOD,
     260  1030732890 :                                   ULONG_MAX - (ULONG_MAX-n+1UL)%n,
     261  1030732890 :                                   (n << (63 - fd_ulong_find_msb( n ) )) - 1UL );
     262             : 
     263  1310363745 :   for( int i=0; 1; i++ ) {
     264  1310363745 :     ulong v = fd_chacha_rng_ulong( rng );
     265  1310363745 :     ulong hi, lo;
     266  1310363745 :     fd_uwide_mul( &hi, &lo, v, n );
     267             : 
     268             : #   if FD_CHACHA_RNG_DEBUG
     269             :     FD_LOG_DEBUG(( "roll (attempt %d): n=%016lx zone: %016lx v=%016lx lo=%016lx hi=%016lx", i, n, zone, v, lo, hi ));
     270             : #   else
     271  1310363745 :     (void)i;
     272  1310363745 : #   endif /* FD_CHACHA_RNG_DEBUG */
     273             : 
     274  1310363745 :     if( FD_LIKELY( lo<=zone ) ) return hi;
     275  1310363745 :   }
     276  1030732890 : }
     277             : 
     278             : FD_PROTOTYPES_END
     279             : 
     280             : #endif /* HEADER_fd_src_ballet_chacha20_fd_chacha20_rng_h */