From 00f504f2c5daa31ea918d3bb0a5f1d397d30e0cf Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Wed, 29 Nov 2023 17:32:06 +0000 Subject: [PATCH] Updated Python module, Numpy module and randquik CLI for new core. --- meson.build | 7 ++++++ randquik/cha.py | 30 ++++++++++++++-------- randquik/nprand.py | 32 ------------------------ setup.py | 13 +++++++++- src/cha1block.h | 23 +++++++++-------- src/cha4block.h | 15 +++++------ src/cha8block.h | 21 ++++++++-------- src/chacha20.h | 62 ++++++++++++++++++++++------------------------ src/npbitgen.h | 24 ++++++++---------- src/nprand.pyx | 11 ++++---- 10 files changed, 117 insertions(+), 121 deletions(-) delete mode 100644 randquik/nprand.py diff --git a/meson.build b/meson.build index e5e8662..3ab737e 100644 --- a/meson.build +++ b/meson.build @@ -6,3 +6,10 @@ executable( install: true, ) dependency('threads') + +library( + 'randquik-chacha20', + 'src/chacha20.c', + build_by_default: true, + c_args: ['-Wall', '-O3', '-march=native'], +) diff --git a/randquik/cha.py b/randquik/cha.py index 64d45d6..6b60722 100644 --- a/randquik/cha.py +++ b/randquik/cha.py @@ -12,15 +12,20 @@ if not src.is_dir(): ffi = cffi.FFI() ffi.cdef( """ + typedef uint64_t (*genfunc)( + uint8_t* out, size_t outsize, uint32_t state[16], unsigned rounds + ); typedef struct cha_ctx { uint32_t input[16]; - uint8_t unconsumed[64]; - uint8_t uncount; + uint8_t unconsumed[512]; + uint32_t offset, end; + unsigned rounds; + genfunc gen; } cha_ctx; - int cha_generate(uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16]); + int cha_generate(uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16], unsigned rounds); - void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv); + void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds); void cha_wipe(cha_ctx* ctx); int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen); """ @@ -65,11 +70,11 @@ def _processBuffer(out): class Cha: - def __init__(self, key: bytes | Any, iv: bytes | Any): + def __init__(self, key: bytes | Any, iv: bytes | Any, *, rounds=8): """Construct a generator that holds its internal state, moving forward on each call.""" key, iv = _processKeys(key, iv) self.ctx = ffi.new("cha_ctx*") - lib.cha_init(self.ctx, key, iv) + lib.cha_init(self.ctx, key, iv, rounds) def __del__(self): lib.cha_wipe(self.ctx) @@ -82,15 +87,20 @@ class Cha: def generate_into( - out: bytearray | memoryview | Any, key: bytes | Any, iv: bytes | Any = bytes(16) + out: bytearray | memoryview | Any, + key: bytes | Any, + iv: bytes | Any = bytes(16), + *, + rounds=8, ): """Fill in random bytes into an existing array (buffer interface)""" key, iv = _processKeys(key, iv) outbuf, outlen = _processBuffer(out) - lib.cha_generate(outbuf, outlen, key, iv) + lib.cha_generate(outbuf, outlen, key, iv, rounds) return out -def generate(outlen: int, key: bytes | Any, iv: bytes | Any = bytes(16)): +def generate(outlen: int, key: bytes | Any, iv: bytes | Any = bytes(16), *, rounds=8): """Return a bytearray of random bytes""" - return generate_into(bytearray(outlen), key, iv) + assert outlen >= 0 + return generate_into(bytearray(outlen), key, iv, rounds=rounds) diff --git a/randquik/nprand.py b/randquik/nprand.py deleted file mode 100644 index 2f14a29..0000000 --- a/randquik/nprand.py +++ /dev/null @@ -1,32 +0,0 @@ -import secrets -import sys - -import cha -import numpy as np - - -class ChaRandom(np.random.BitGenerator): - def __init__(self, seed=None): - super().__init__(seed) - sys.stderr.write("Construct\n") - if seed is None: - key = secrets.token_bytes(32) - else: - key = ( - np.random.SeedSequence(seed, pool_size=8) - .generate_state(4, dtype=np.uint64) - .tobytes() - ) - self._generator = cha.Cha(key, bytes(8) + b"NumpyGen") - - def random_raw(self, size=None): - sys.stderr.write(f"Random raw {size=}\n") - if size is None: - return int.from_bytes(self._generator(bytearray(8)), "little") - ret = np.empty(size, np.uint64) - self._generator(ret.data) - return ret - - def spawn(self, n): - sys.stderr.write(f"Spawn {n=}\n") - raise NotImplementedError diff --git a/setup.py b/setup.py index 6706d91..3b8febe 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,17 @@ +import os from distutils.core import setup import numpy from Cython.Build import cythonize +from setuptools import Extension -setup(ext_modules=cythonize("src/nprand.pyx"), include_dirs=[numpy.get_include()]) +os.environ["CFLAGS"] = "-O3 -march=native -Wall -Wextra" +extensions = [ + Extension( + "nprand", + ["src/nprand.pyx"], + include_dirs=[numpy.get_include()], + define_macros=[("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")], + ) +] +setup(ext_modules=cythonize(extensions)) diff --git a/src/cha1block.h b/src/cha1block.h index 581ce62..6ef78d9 100644 --- a/src/cha1block.h +++ b/src/cha1block.h @@ -7,19 +7,22 @@ QUARTERSTEP(a, b, d, 16); QUARTERSTEP(c, d, b, 12); \ QUARTERSTEP(a, b, d, 8); QUARTERSTEP(c, d, b, 7); } -static inline uint64_t _cha_block(uint8_t out[CHA_BLOCK_SIZE], uint32_t state[16], unsigned rounds) { +static inline uint64_t _cha_block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) { + unsigned blocks = bufsize / 64; + uint32_t* out = (uint32_t*)buf; uint32_t x[16]; - for (unsigned i = 0; i < 16; ++i) x[i] = state[i]; // Faster than memcpy - for (unsigned i = rounds / 2; i-->0;) { - // Mix columns, then diagonals - for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + j], x[8 + j], x[12 + j]); - for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + (j+1)%4], x[8 + (j+2)%4], x[12 + (j+3)%4]); + for (unsigned b = blocks; b-->0;) { + for (unsigned i = 0; i < 16; ++i) x[i] = state[i]; // Faster than memcpy + for (unsigned i = rounds / 2; i-->0;) { + // Mix columns, then diagonals + for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + j], x[8 + j], x[12 + j]); + for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + (j+1)%4], x[8 + (j+2)%4], x[12 + (j+3)%4]); + } + for (unsigned i = 0; i < 16; ++i) *out++ = x[i] + state[i]; + ++*(uint64_t*)(state + 12); // Increment counter } - uint32_t* buf = (uint32_t*)out; - for (unsigned i = 0; i < 16; ++i) buf[i] = x[i] + state[i]; memset(x, 0, sizeof x); - ++*(uint64_t*)(state + 12); // Increment counter - return CHA_BLOCK_SIZE; + return blocks * CHA_BLOCK_SIZE; } #undef QUARTERROUND diff --git a/src/cha4block.h b/src/cha4block.h index c7ff692..48d1641 100644 --- a/src/cha4block.h +++ b/src/cha4block.h @@ -1,5 +1,6 @@ #if defined(__x86_64__) #include // SSE2 +#include // SSSE3 #elif defined(__aarch64__) #include "sse2neon.h" #endif @@ -21,7 +22,7 @@ x[C] = _mm_add_epi32(x[C], x[D]); \ x[B] = VEC4_ROT(_mm_xor_si128(x[B], x[C]), 7) -#define ONEQUAD(A, B, C, D, CT) \ +#define ONEQUAD(A, B, C, D, OUT) \ { \ /* Add original block */ \ x[A] = _mm_add_epi32(x[A], orig[A]); \ @@ -37,11 +38,11 @@ x[B] = _mm_unpackhi_epi64(abl, cdl); /* a1 b1 c1 d1 */ \ x[C] = _mm_unpacklo_epi64(abh, cdh); /* a2 b2 c2 d2 */ \ x[D] = _mm_unpackhi_epi64(abh, cdh); /* a3 b3 c3 d3 */ \ - \ - _mm_storeu_si128((__m128i*)(CT), x[A]); \ - _mm_storeu_si128((__m128i*)(CT + 64), x[B]); \ - _mm_storeu_si128((__m128i*)(CT + 128), x[C]); \ - _mm_storeu_si128((__m128i*)(CT + 192), x[D]); \ + /* Write out 1/4 of each block */ \ + _mm_storeu_si128((__m128i*)(OUT), x[A]); \ + _mm_storeu_si128((__m128i*)(OUT + 64), x[B]); \ + _mm_storeu_si128((__m128i*)(OUT + 128), x[C]); \ + _mm_storeu_si128((__m128i*)(OUT + 192), x[D]); \ } #define COUNTER_INCREMENT(a, b, c, d) \ @@ -87,7 +88,7 @@ _cha_4block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) { COUNTER_INCREMENT(4, 4, 4, 4); buf += 256; } - // Update counter + // Store counter state[12] = _mm_cvtsi128_si32(orig[12]); state[13] = _mm_cvtsi128_si32(orig[13]); return batches * 256; diff --git a/src/cha8block.h b/src/cha8block.h index 795fe0a..c9c6270 100644 --- a/src/cha8block.h +++ b/src/cha8block.h @@ -1,5 +1,4 @@ #include // AVX2 -#include // SSSE3 // clang-format off @@ -53,16 +52,16 @@ #define ONEOCTO(A, B, C, D, A2, B2, C2, D2, c) \ { \ - TRANSPOSE(A, B, C, D); \ - TRANSPOSE(A2, B2, C2, D2); \ - _mm256_storeu_si256((__m256i*)(c), _mm256_permute2x128_si256(x[A], x[A2], 0x20)); \ - _mm256_storeu_si256((__m256i*)(c + 64), _mm256_permute2x128_si256(x[B], x[B2], 0x20)); \ - _mm256_storeu_si256((__m256i*)(c + 128), _mm256_permute2x128_si256(x[C], x[C2], 0x20)); \ - _mm256_storeu_si256((__m256i*)(c + 192), _mm256_permute2x128_si256(x[D], x[D2], 0x20)); \ - _mm256_storeu_si256((__m256i*)(c + 256), _mm256_permute2x128_si256(x[A], x[A2], 0x31)); \ - _mm256_storeu_si256((__m256i*)(c + 320), _mm256_permute2x128_si256(x[B], x[B2], 0x31)); \ - _mm256_storeu_si256((__m256i*)(c + 384), _mm256_permute2x128_si256(x[C], x[C2], 0x31)); \ - _mm256_storeu_si256((__m256i*)(c + 448), _mm256_permute2x128_si256(x[D], x[D2], 0x31)); \ + TRANSPOSE(A, B, C, D); \ + TRANSPOSE(A2, B2, C2, D2); \ + _mm256_storeu_si256((__m256i*)(c), _mm256_permute2x128_si256(x[A], x[A2], 0x20)); \ + _mm256_storeu_si256((__m256i*)(c + 64), _mm256_permute2x128_si256(x[B], x[B2], 0x20)); \ + _mm256_storeu_si256((__m256i*)(c + 128), _mm256_permute2x128_si256(x[C], x[C2], 0x20)); \ + _mm256_storeu_si256((__m256i*)(c + 192), _mm256_permute2x128_si256(x[D], x[D2], 0x20)); \ + _mm256_storeu_si256((__m256i*)(c + 256), _mm256_permute2x128_si256(x[A], x[A2], 0x31)); \ + _mm256_storeu_si256((__m256i*)(c + 320), _mm256_permute2x128_si256(x[B], x[B2], 0x31)); \ + _mm256_storeu_si256((__m256i*)(c + 384), _mm256_permute2x128_si256(x[C], x[C2], 0x31)); \ + _mm256_storeu_si256((__m256i*)(c + 448), _mm256_permute2x128_si256(x[D], x[D2], 0x31)); \ } #define COUNTER_INCREMENT(addv) \ diff --git a/src/chacha20.h b/src/chacha20.h index 3a6af50..08b74d4 100644 --- a/src/chacha20.h +++ b/src/chacha20.h @@ -11,12 +11,13 @@ #pragma GCC target("ssse3") #pragma GCC target("avx2") #endif +#include "cha4block.h" #include "cha8block.h" -#else +#elif defined(__aarch64__) +#include "cha4block.h" #endif #include "cha1block.h" -#include "cha4block.h" #include #include @@ -28,11 +29,15 @@ #include #include +typedef uint64_t (*genfunc)( + uint8_t* out, size_t outsize, uint32_t state[16], unsigned rounds +); typedef struct cha_ctx { uint32_t state[16]; uint8_t unconsumed[BATCH_SIZE]; uint32_t offset, end; unsigned rounds; + genfunc gen; } cha_ctx; /// @brief Initialize cha_ctx @@ -52,6 +57,18 @@ void cha_init( memset(ctx->unconsumed, 0, sizeof ctx->unconsumed); ctx->offset = ctx->end = 0; ctx->rounds = rounds; +#if defined(__x86_64__) + if (__builtin_cpu_supports("avx2")) + ctx->gen = _cha_8block; + else if (__builtin_cpu_supports("ssse3")) + ctx->gen = _cha_4block; + else + ctx->gen = _cha_block; +#elif defined(__aarch64__) + ctx->gen = _cha_4block; +#else + ctx->gen = _cha_block; +#endif } /// Dispose of sensitive data within the context @@ -65,39 +82,18 @@ void cha_seek_blocks(cha_ctx* ctx, int64_t offset) { ctx->offset = ctx->end = 0; } -uint64_t cha_generate_batch( - uint8_t* out, size_t outsize, uint32_t* state, unsigned rounds -) { -#if defined(__x86_64__) - if (__builtin_cpu_supports("ssse3")) { - if (__builtin_cpu_supports("avx2")) { - return _cha_8block(out, outsize, state, rounds); - } - return _cha_4block(out, outsize, state, rounds); - } -#elif defined(__aarch64__) - return _cha_4block(out, outsize, state); -#endif - unsigned count = 0; - unsigned n = _cha_block(out, state, rounds); - count += n; - out += n; - return count; -} - /// @brief Incremental generation, keeps state between calls /// @param ctx ChaCha context /// @param out output buffer /// @param outlen output buffer length -void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) { +static inline void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) { // The included header will mess with these variables uint8_t* end = out + outlen; if (ctx->offset) { // Need to generate stored buffer? if (ctx->end == 0) - ctx->end = cha_generate_batch( - ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds - ); + ctx->end = + ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds); // Deliver stored bytes first uint64_t N = ctx->end - ctx->offset; if (N > outlen) @@ -105,15 +101,16 @@ void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) { memcpy(out, ctx->unconsumed + ctx->offset, N); ctx->offset += N; out += N; + if (ctx->offset == ctx->end) + ctx->offset = ctx->end = 0; if (out == end) return; } - out += cha_generate_batch(out, end - out, ctx->state, ctx->rounds); + out += ctx->gen(out, end - out, ctx->state, ctx->rounds); const uint32_t N = end - out; if (N) { - ctx->end = cha_generate_batch( - ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds - ); + ctx->end = + ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds); memcpy(out, ctx->unconsumed, N); ctx->offset = N; } @@ -125,10 +122,11 @@ void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) { /// @param key 32 byte key /// @param iv 16 bytes, where normally initial 4-8 bytes are 0 (counter) void cha_generate( - uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16] + uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16], + unsigned rounds ) { cha_ctx ctx; - cha_init(&ctx, key, iv, 20); + cha_init(&ctx, key, iv, rounds); cha_update(&ctx, out, outlen); cha_wipe(&ctx); } diff --git a/src/npbitgen.h b/src/npbitgen.h index 51a7eb5..a682e8b 100644 --- a/src/npbitgen.h +++ b/src/npbitgen.h @@ -1,19 +1,17 @@ #include "chacha20.h" - -static uint64_t cha_uint64(void *st) { - cha_ctx *ctx = (cha_ctx *)st; - uint64_t ret; - cha_update(ctx, (uint8_t *)&ret, sizeof ret); +static uint64_t cha_uint64(void* st) { + cha_ctx* ctx = (cha_ctx*)st; + if (ctx->offset == ctx->end) { + ctx->offset = 0; + ctx->end = + ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds); + } + register uint64_t ret = *(uint64_t*)(ctx->unconsumed + ctx->offset); + ctx->offset += sizeof(uint64_t); return ret; } -static uint32_t cha_uint32(void *st) { - cha_ctx *ctx = (cha_ctx *)st; - uint32_t ret; - cha_update(ctx, (uint8_t *)&ret, sizeof ret); - return ret; - -} -static double cha_double(void *st) { +static uint32_t cha_uint32(void* st) { return cha_uint64(st); } +static double cha_double(void* st) { return cha_uint64(st) / (UINT64_MAX + 1.0); } diff --git a/src/nprand.pyx b/src/nprand.pyx index 5fc1ae0..aad27e9 100644 --- a/src/nprand.pyx +++ b/src/nprand.pyx @@ -15,10 +15,11 @@ np.import_array() cdef extern from "npbitgen.h": struct cha_ctx: uint32_t state[16] - uint8_t unconsumed[64] - uint8_t uncount + uint8_t unconsumed[512] + uint32_t offset, end; + unsigned rounds; - void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv) nogil + void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds) nogil uint64_t cha_uint64(void *state) nogil uint32_t cha_uint32(void *state) nogil double cha_double(void *state) nogil @@ -27,7 +28,7 @@ cdef class Cha(BitGenerator): cdef cha_ctx rng_state - def __init__(self, seed=None): + def __init__(self, seed=None, rounds=8): BitGenerator.__init__(self, seed) self._bitgen.state = &self.rng_state self._bitgen.next_uint64 = &cha_uint64 @@ -36,4 +37,4 @@ cdef class Cha(BitGenerator): self._bitgen.next_raw = &cha_uint64 # Generated state is ChaCha20 key key = self._seed_seq.generate_state(4, np.uint64) - cha_init(&self.rng_state, np.PyArray_DATA(key), b"NumpRand") + cha_init(&self.rng_state, np.PyArray_DATA(key), b"NumpRand", rounds)