LeoVasanko/RandQuik
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Updated Python module, Numpy module and randquik CLI for new core.

Browse Source
This commit is contained in:
LeoVasanko
2023-11-29 17:32:06 +00:00
parent 236f7eedf6
commit 00f504f2c5
10 changed files with 117 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
meson.build
+7
View File
@@ -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'],
)
randquik/cha.py
+20 -10
View File
@@ -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)
randquik/nprand.py
-32
View File
@@ -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
setup.py
+12 -1
View File
@@ -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))
src/cha1block.h
+13 -10
View File
@@ -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
src/cha4block.h
+8 -7
View File
@@ -1,5 +1,6 @@
#if defined(__x86_64__)
#include <emmintrin.h> // SSE2
#include <tmmintrin.h> // 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;
src/cha8block.h
+10 -11
View File
@@ -1,5 +1,4 @@
#include <immintrin.h> // AVX2
#include <tmmintrin.h> // 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) \
src/chacha20.h
+30 -32
View File
@@ -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 <assert.h>
#include <stdbool.h>
@@ -28,11 +29,15 @@
#include <time.h>
#include <unistd.h>
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);
}
src/npbitgen.h
+11 -13
View File
@@ -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);
}
src/nprand.pyx
+6 -5
View File
@@ -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 = <void *>&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, <uint8_t *>np.PyArray_DATA(key), b"NumpRand")
cha_init(&self.rng_state, <uint8_t *>np.PyArray_DATA(key), b"NumpRand", rounds)