chacha: add ChaCha8 backend and bench

Author: riptl

src/ballet/chacha/Local.mk

@@ -0,0 +1,22 @@
+# Base ChaCha support
+$(call add-hdrs,fd_chacha.h)
+ifdef FD_HAS_SSE
+$(call add-objs,fd_chacha_sse,fd_ballet)
+else
+$(call add-objs,fd_chacha,fd_ballet)
+endif
+$(call make-unit-test,test_chacha,test_chacha,fd_ballet fd_util)
+$(call run-unit-test,test_chacha)
+
+# ChaCha-RNG support (Rust rand_chacha compatible)
+$(call add-hdrs,fd_chacha_rng.h)
+$(call add-objs,fd_chacha_rng,fd_ballet)
+ifdef FD_HAS_AVX512
+$(call add-objs,fd_chacha_rng_avx512,fd_ballet)
+endif
+ifdef FD_HAS_AVX
+$(call add-objs,fd_chacha_rng_avx,fd_ballet)
+endif
+$(call make-unit-test,test_chacha_rng,test_chacha_rng,fd_ballet fd_util)
+$(call make-unit-test,test_chacha_rng_roll,test_chacha_rng_roll,fd_ballet fd_util)
+$(call run-unit-test,test_chacha_rng)

src/ballet/chacha20/fd_chacha20.c renamed to src/ballet/chacha/fd_chacha.c

@@ -1,25 +1,26 @@
-#include "fd_chacha20.h"
+#include "fd_chacha.h"
 
 /* Reference implementation of the ChaCha20 block function.
 
    FIXME Not optimized for high performance.  Trivially parallelizable
          via SSE or AVX if required. */
 
 static inline void
-fd_chacha20_quarter_round( uint * a,
-                           uint * b,
-                           uint * c,
-                           uint * d ) {
+fd_chacha_quarter_round( uint * a,
+                         uint * b,
+                         uint * c,
+                         uint * d ) {
   *a += *b; *d ^= *a; *d = fd_uint_rotate_left(*d, 16);
   *c += *d; *b ^= *c; *b = fd_uint_rotate_left(*b, 12);
   *a += *b; *d ^= *a; *d = fd_uint_rotate_left(*d,  8);
   *c += *d; *b ^= *c; *b = fd_uint_rotate_left(*b,  7);
 }
 
-void *
-fd_chacha20_block( void *       _block,
-                   void const * _key,
-                   void const * _idx_nonce ) {
+__attribute__((always_inline)) static inline void *
+fd_chacha_block( void *       _block,
+                 void const * _key,
+                 void const * _idx_nonce,
+                 ulong        rnd2_cnt ) {
 
   uint *       block     = __builtin_assume_aligned( _block,     64UL );
   uint const * key       = __builtin_assume_aligned( _key,       32UL );
@@ -55,15 +56,15 @@ fd_chacha20_block( void *       _block,
   /* Run the ChaCha round function 20 times.
      (Each iteration does a column round and a diagonal round.) */
 
-  for( ulong i=0UL; i<10UL; i++ ) {
-    fd_chacha20_quarter_round( &block[ 0 ], &block[ 4 ], &block[  8 ], &block[ 12 ] );
-    fd_chacha20_quarter_round( &block[ 1 ], &block[ 5 ], &block[  9 ], &block[ 13 ] );
-    fd_chacha20_quarter_round( &block[ 2 ], &block[ 6 ], &block[ 10 ], &block[ 14 ] );
-    fd_chacha20_quarter_round( &block[ 3 ], &block[ 7 ], &block[ 11 ], &block[ 15 ] );
-    fd_chacha20_quarter_round( &block[ 0 ], &block[ 5 ], &block[ 10 ], &block[ 15 ] );
-    fd_chacha20_quarter_round( &block[ 1 ], &block[ 6 ], &block[ 11 ], &block[ 12 ] );
-    fd_chacha20_quarter_round( &block[ 2 ], &block[ 7 ], &block[  8 ], &block[ 13 ] );
-    fd_chacha20_quarter_round( &block[ 3 ], &block[ 4 ], &block[  9 ], &block[ 14 ] );
+  for( ulong i=0UL; i<rnd2_cnt; i++ ) {
+    fd_chacha_quarter_round( &block[ 0 ], &block[ 4 ], &block[  8 ], &block[ 12 ] );
+    fd_chacha_quarter_round( &block[ 1 ], &block[ 5 ], &block[  9 ], &block[ 13 ] );
+    fd_chacha_quarter_round( &block[ 2 ], &block[ 6 ], &block[ 10 ], &block[ 14 ] );
+    fd_chacha_quarter_round( &block[ 3 ], &block[ 7 ], &block[ 11 ], &block[ 15 ] );
+    fd_chacha_quarter_round( &block[ 0 ], &block[ 5 ], &block[ 10 ], &block[ 15 ] );
+    fd_chacha_quarter_round( &block[ 1 ], &block[ 6 ], &block[ 11 ], &block[ 12 ] );
+    fd_chacha_quarter_round( &block[ 2 ], &block[ 7 ], &block[  8 ], &block[ 13 ] );
+    fd_chacha_quarter_round( &block[ 3 ], &block[ 4 ], &block[  9 ], &block[ 14 ] );
   }
 
   /* Complete the block by adding the input state */
@@ -74,3 +75,16 @@ fd_chacha20_block( void *       _block,
   return (void *)block;
 }
 
+void *
+fd_chacha8_block( void *       _block,
+                  void const * _key,
+                  void const * _idx_nonce ) {
+  return fd_chacha_block( _block, _key, _idx_nonce, 4UL );
+}
+
+void *
+fd_chacha20_block( void *       _block,
+                   void const * _key,
+                   void const * _idx_nonce ) {
+  return fd_chacha_block( _block, _key, _idx_nonce, 10UL );
+}

src/ballet/chacha20/fd_chacha20.h renamed to src/ballet/chacha/fd_chacha.h

@@ -3,9 +3,9 @@
 
 #include "../fd_ballet_base.h"
 
-/* FD_CHACHA20_BLOCK_SZ is the output size of the ChaCha20 block function. */
+/* FD_CHACHA_BLOCK_SZ is the output size of the ChaCha20 block function. */
 
-#define FD_CHACHA20_BLOCK_SZ (64UL)
+#define FD_CHACHA_BLOCK_SZ (64UL)
 
 /* FD_CHACHA20_KEY_SZ is the size of the ChaCha20 encryption key */
 
@@ -23,6 +23,11 @@ FD_PROTOTYPES_BEGIN
 
    FIXME this should probably do multiple blocks */
 
+void *
+fd_chacha8_block( void *       block,
+                  void const * key,
+                  void const * idx_nonce );
+
 void *
 fd_chacha20_block( void *       block,
                    void const * key,

src/ballet/chacha/fd_chacha_rng.c

@@ -0,0 +1,96 @@
+#include "fd_chacha_rng.h"
+
+FD_FN_CONST ulong
+fd_chacha_rng_align( void ) {
+  return alignof(fd_chacha_rng_t);
+}
+
+FD_FN_CONST ulong
+fd_chacha_rng_footprint( void ) {
+  return sizeof(fd_chacha_rng_t);
+}
+
+void *
+fd_chacha_rng_new( void * shmem, int mode ) {
+  if( FD_UNLIKELY( !shmem ) ) {
+    FD_LOG_WARNING(( "NULL shmem" ));
+    return NULL;
+  }
+  if( FD_UNLIKELY( !fd_ulong_is_aligned( (ulong)shmem, alignof(fd_chacha_rng_t) ) ) ) {
+    FD_LOG_WARNING(( "misaligned shmem" ));
+    return NULL;
+  }
+  memset( shmem, 0, sizeof(fd_chacha_rng_t) );
+  if( FD_UNLIKELY( (mode!=FD_CHACHA_RNG_MODE_MOD) & (mode!=FD_CHACHA_RNG_MODE_SHIFT) ) ) {
+    FD_LOG_WARNING(( "invalid mode" ));
+    return NULL;
+  }
+  ((fd_chacha_rng_t *)shmem)->mode = mode;
+
+  return shmem;
+}
+
+fd_chacha_rng_t *
+fd_chacha_rng_join( void * shrng ) {
+  if( FD_UNLIKELY( !shrng ) ) {
+    FD_LOG_WARNING(( "NULL shrng" ));
+    return NULL;
+  }
+  return (fd_chacha_rng_t *)shrng;
+}
+
+void *
+fd_chacha_rng_leave( fd_chacha_rng_t * rng ) {
+  if( FD_UNLIKELY( !rng ) ) {
+    FD_LOG_WARNING(( "NULL rng" ));
+    return NULL;
+  }
+  return (void *)rng;
+}
+
+void *
+fd_chacha_rng_delete( void * shrng ) {
+  if( FD_UNLIKELY( !shrng ) ) {
+    FD_LOG_WARNING(( "NULL shrng" ));
+    return NULL;
+  }
+  memset( shrng, 0, sizeof(fd_chacha_rng_t) );
+  return 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 );
+  rng->buf_off  = 0UL;
+  rng->buf_fill = 0UL;
+  fd_chacha20_rng_private_refill( rng );
+  return rng;
+}
+
+static void
+fd_chacha_rng_refill_seq( fd_chacha_rng_t * rng,
+                          void * (* block_fn)( void *, void const *, void const * ) ) {
+  ulong fill_target = FD_CHACHA_RNG_BUFSZ - FD_CHACHA_BLOCK_SZ;
+
+  ulong buf_avail;
+  while( (buf_avail=(rng->buf_fill - rng->buf_off))<fill_target ) {
+    ulong idx = rng->buf_fill >> 6;
+    uint idx_nonce[4] __attribute__((aligned(16))) =
+      { (uint)idx, 0U, 0U, 0U };
+    block_fn( rng->buf + (rng->buf_fill % FD_CHACHA_RNG_BUFSZ),
+              rng->key,
+              idx_nonce );
+    rng->buf_fill += (uint)FD_CHACHA_BLOCK_SZ;
+  }
+}
+
+void
+fd_chacha8_rng_refill_seq( fd_chacha_rng_t * rng ) {
+  fd_chacha_rng_refill_seq( rng, fd_chacha8_block );
+}
+
+void
+fd_chacha20_rng_refill_seq( fd_chacha_rng_t * rng ) {
+  fd_chacha_rng_refill_seq( rng, fd_chacha20_block );
+}