diff --git a/CHANGELOG b/CHANGELOG --- a/CHANGELOG +++ b/CHANGELOG @@ -1,5 +1,10 @@ # Changelog +- 0.3.1 (2026-10-31) + * Uses SSE2 on x86_64, along with AVX2 on CPUs that support it, + detected at runtime. With AVX2, 'cipher' is roughly 3.5x faster on + short inputs and over 10x faster on long ones. + - 0.3.0 (2026-09-06) * Adds a counter overflow error. diff --git a/bench/Main.hs b/bench/Main.hs --- a/bench/Main.hs +++ b/bench/Main.hs @@ -35,9 +35,17 @@ non = fromJust . B16.decode $ "000000090000004a00000000" +plain_1k :: BS.ByteString +plain_1k = BS.replicate 1024 0x42 + +plain_4k :: BS.ByteString +plain_4k = BS.replicate 4096 0x42 + suite :: Benchmark suite = bgroup "ppad-chacha" [ bench "cipher" $ nf (ChaCha20.cipher key 1 non) plain + , bench "cipher (1024B)" $ nf (ChaCha20.cipher key 1 non) plain_1k + , bench "cipher (4096B)" $ nf (ChaCha20.cipher key 1 non) plain_4k ] diff --git a/cbits/chacha20_arm.c b/cbits/chacha20_arm.c deleted file mode 100644 --- a/cbits/chacha20_arm.c +++ /dev/null @@ -1,199 +0,0 @@ -#include -#include -#include - -#if defined(__aarch64__) - -#include - -/* - * ChaCha20 NEON kernel using intra-block parallelism. The 16-word - * state matrix - * - * s00 s01 s02 s03 - * s04 s05 s06 s07 - * s08 s09 s10 s11 - * s12 s13 s14 s15 - * - * is held in four 128-bit NEON registers v0..v3, one per row. A - * column quarter-round on (s00, s04, s08, s12), (s01, s05, s09, s13), - * etc., becomes one set of element-wise vector operations on - * (v0, v1, v2, v3) — four quarter-rounds in parallel. Diagonal - * rounds are reached by left-rotating v1, v2, v3 by 1, 2, 3 lanes - * respectively with VEXT before the second quarter-round, then - * rotating back. - */ - -/* 32-bit left rotations. Rotate-by-16 reduces to REV32.u16; the - * others compile to a shift-shift-or pair (the compiler folds rotate- - * by-8 to a TBL with a constant shuffle on some targets). */ -#define ROTL32_16(x) \ - vreinterpretq_u32_u16(vrev32q_u16(vreinterpretq_u16_u32(x))) -#define ROTL32_12(x) \ - vorrq_u32(vshlq_n_u32((x), 12), vshrq_n_u32((x), 20)) -#define ROTL32_8(x) \ - vorrq_u32(vshlq_n_u32((x), 8), vshrq_n_u32((x), 24)) -#define ROTL32_7(x) \ - vorrq_u32(vshlq_n_u32((x), 7), vshrq_n_u32((x), 25)) - -#define QUARTER(v0, v1, v2, v3) \ - do { \ - v0 = vaddq_u32(v0, v1); \ - v3 = veorq_u32(v3, v0); v3 = ROTL32_16(v3); \ - v2 = vaddq_u32(v2, v3); \ - v1 = veorq_u32(v1, v2); v1 = ROTL32_12(v1); \ - v0 = vaddq_u32(v0, v1); \ - v3 = veorq_u32(v3, v0); v3 = ROTL32_8(v3); \ - v2 = vaddq_u32(v2, v3); \ - v1 = veorq_u32(v1, v2); v1 = ROTL32_7(v1); \ - } while (0) - -/* 20-round ChaCha20 core: 10 iterations of (column + diagonal). */ -static inline void chacha20_core(uint32x4_t *v0, uint32x4_t *v1, - uint32x4_t *v2, uint32x4_t *v3, - uint32x4_t s0, uint32x4_t s1, - uint32x4_t s2, uint32x4_t s3) { - uint32x4_t a = s0, b = s1, c = s2, d = s3; - for (int i = 0; i < 10; i++) { - QUARTER(a, b, c, d); - /* shift rows: row 1 left 1, row 2 left 2, row 3 left 3. */ - b = vextq_u32(b, b, 1); - c = vextq_u32(c, c, 2); - d = vextq_u32(d, d, 3); - QUARTER(a, b, c, d); - /* shift back. */ - b = vextq_u32(b, b, 3); - c = vextq_u32(c, c, 2); - d = vextq_u32(d, d, 1); - } - *v0 = vaddq_u32(a, s0); - *v1 = vaddq_u32(b, s1); - *v2 = vaddq_u32(c, s2); - *v3 = vaddq_u32(d, s3); -} - -static const uint32_t chacha_constants[4] = { - 0x61707865u, 0x3320646eu, 0x79622d32u, 0x6b206574u -}; - -/* Set up the constant rows of the state from key + nonce. s3 - * (counter + nonce) varies per block and is built inside the loop. */ -static inline void chacha20_setup(const uint8_t key[32], - const uint8_t nonce[12], - uint32x4_t *s0, uint32x4_t *s1, - uint32x4_t *s2, - uint32_t *n0, uint32_t *n1, - uint32_t *n2) { - *s0 = vld1q_u32(chacha_constants); - *s1 = vreinterpretq_u32_u8(vld1q_u8(key)); - *s2 = vreinterpretq_u32_u8(vld1q_u8(key + 16)); - memcpy(n0, nonce + 0, 4); - memcpy(n1, nonce + 4, 4); - memcpy(n2, nonce + 8, 4); -} - -/* - * Generate one 64-byte ChaCha20 keystream block at 'out'. - */ -void chacha20_block_arm(const uint8_t key[32], uint32_t counter, - const uint8_t nonce[12], uint8_t out[64]) { - uint32x4_t s0, s1, s2; - uint32_t n0, n1, n2; - chacha20_setup(key, nonce, &s0, &s1, &s2, &n0, &n1, &n2); - - uint32_t s3_in[4] = { counter, n0, n1, n2 }; - uint32x4_t s3 = vld1q_u32(s3_in); - uint32x4_t v0, v1, v2, v3; - chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); - - vst1q_u8(out + 0, vreinterpretq_u8_u32(v0)); - vst1q_u8(out + 16, vreinterpretq_u8_u32(v1)); - vst1q_u8(out + 32, vreinterpretq_u8_u32(v2)); - vst1q_u8(out + 48, vreinterpretq_u8_u32(v3)); -} - -/* - * Encrypt/decrypt 'inlen' bytes at 'in' into 'out' using ChaCha20 - * with the given key, starting counter, and nonce. Stream cipher, - * so the same routine decrypts. - */ -void chacha20_cipher_arm(const uint8_t key[32], uint32_t counter, - const uint8_t nonce[12], - const uint8_t *in, uint8_t *out, - size_t inlen) { - uint32x4_t s0, s1, s2; - uint32_t n0, n1, n2; - chacha20_setup(key, nonce, &s0, &s1, &s2, &n0, &n1, &n2); - - size_t pos = 0; - while (pos + 64 <= inlen) { - uint32_t s3_in[4] = { counter, n0, n1, n2 }; - uint32x4_t s3 = vld1q_u32(s3_in); - uint32x4_t v0, v1, v2, v3; - chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); - - uint8x16_t i0 = vld1q_u8(in + pos + 0); - uint8x16_t i1 = vld1q_u8(in + pos + 16); - uint8x16_t i2 = vld1q_u8(in + pos + 32); - uint8x16_t i3 = vld1q_u8(in + pos + 48); - - vst1q_u8(out + pos + 0, - veorq_u8(i0, vreinterpretq_u8_u32(v0))); - vst1q_u8(out + pos + 16, - veorq_u8(i1, vreinterpretq_u8_u32(v1))); - vst1q_u8(out + pos + 32, - veorq_u8(i2, vreinterpretq_u8_u32(v2))); - vst1q_u8(out + pos + 48, - veorq_u8(i3, vreinterpretq_u8_u32(v3))); - - pos += 64; - counter++; - } - - /* trailing partial block (< 64 bytes) */ - if (pos < inlen) { - uint32_t s3_in[4] = { counter, n0, n1, n2 }; - uint32x4_t s3 = vld1q_u32(s3_in); - uint32x4_t v0, v1, v2, v3; - chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); - - uint8_t block[64]; - vst1q_u8(block + 0, vreinterpretq_u8_u32(v0)); - vst1q_u8(block + 16, vreinterpretq_u8_u32(v1)); - vst1q_u8(block + 32, vreinterpretq_u8_u32(v2)); - vst1q_u8(block + 48, vreinterpretq_u8_u32(v3)); - - size_t remaining = inlen - pos; - for (size_t i = 0; i < remaining; i++) { - out[pos + i] = in[pos + i] ^ block[i]; - } - } -} - -int chacha20_arm_available(void) { - return 1; -} - -#else - -/* stubs for non-aarch64 builds; never reached because dispatch is - * gated on 'chacha20_arm_available' returning 0 */ - -void chacha20_block_arm(const uint8_t *key, uint32_t counter, - const uint8_t *nonce, uint8_t *out) { - (void)key; (void)counter; (void)nonce; (void)out; -} - -void chacha20_cipher_arm(const uint8_t *key, uint32_t counter, - const uint8_t *nonce, - const uint8_t *in, uint8_t *out, - size_t inlen) { - (void)key; (void)counter; (void)nonce; - (void)in; (void)out; (void)inlen; -} - -int chacha20_arm_available(void) { - return 0; -} - -#endif diff --git a/cbits/chacha20_native.c b/cbits/chacha20_native.c new file mode 100644 --- /dev/null +++ b/cbits/chacha20_native.c @@ -0,0 +1,593 @@ +#include +#include +#include + +#if defined(__aarch64__) + +#include + +#elif defined(__x86_64__) + +#include +#include + +#endif + +#if defined(__aarch64__) || defined(__x86_64__) + +static const uint32_t chacha_constants[4] = { + 0x61707865u, 0x3320646eu, 0x79622d32u, 0x6b206574u +}; + +#endif + +#if defined(__aarch64__) + +/* + * ChaCha20 NEON kernel using intra-block parallelism. The 16-word + * state matrix + * + * s00 s01 s02 s03 + * s04 s05 s06 s07 + * s08 s09 s10 s11 + * s12 s13 s14 s15 + * + * is held in four 128-bit NEON registers v0..v3, one per row. A + * column quarter-round on (s00, s04, s08, s12), (s01, s05, s09, s13), + * etc., becomes one set of element-wise vector operations on + * (v0, v1, v2, v3) — four quarter-rounds in parallel. Diagonal + * rounds are reached by left-rotating v1, v2, v3 by 1, 2, 3 lanes + * respectively with VEXT before the second quarter-round, then + * rotating back. + */ + +/* 32-bit left rotations. Rotate-by-16 reduces to REV32.u16; the + * others compile to a shift-shift-or pair (the compiler folds rotate- + * by-8 to a TBL with a constant shuffle on some targets). */ +#define ROTL32_16(x) \ + vreinterpretq_u32_u16(vrev32q_u16(vreinterpretq_u16_u32(x))) +#define ROTL32_12(x) \ + vorrq_u32(vshlq_n_u32((x), 12), vshrq_n_u32((x), 20)) +#define ROTL32_8(x) \ + vorrq_u32(vshlq_n_u32((x), 8), vshrq_n_u32((x), 24)) +#define ROTL32_7(x) \ + vorrq_u32(vshlq_n_u32((x), 7), vshrq_n_u32((x), 25)) + +#define QUARTER(v0, v1, v2, v3) \ + do { \ + v0 = vaddq_u32(v0, v1); \ + v3 = veorq_u32(v3, v0); v3 = ROTL32_16(v3); \ + v2 = vaddq_u32(v2, v3); \ + v1 = veorq_u32(v1, v2); v1 = ROTL32_12(v1); \ + v0 = vaddq_u32(v0, v1); \ + v3 = veorq_u32(v3, v0); v3 = ROTL32_8(v3); \ + v2 = vaddq_u32(v2, v3); \ + v1 = veorq_u32(v1, v2); v1 = ROTL32_7(v1); \ + } while (0) + +/* 20-round ChaCha20 core: 10 iterations of (column + diagonal). */ +static inline void chacha20_core(uint32x4_t *v0, uint32x4_t *v1, + uint32x4_t *v2, uint32x4_t *v3, + uint32x4_t s0, uint32x4_t s1, + uint32x4_t s2, uint32x4_t s3) { + uint32x4_t a = s0, b = s1, c = s2, d = s3; + for (int i = 0; i < 10; i++) { + QUARTER(a, b, c, d); + /* shift rows: row 1 left 1, row 2 left 2, row 3 left 3. */ + b = vextq_u32(b, b, 1); + c = vextq_u32(c, c, 2); + d = vextq_u32(d, d, 3); + QUARTER(a, b, c, d); + /* shift back. */ + b = vextq_u32(b, b, 3); + c = vextq_u32(c, c, 2); + d = vextq_u32(d, d, 1); + } + *v0 = vaddq_u32(a, s0); + *v1 = vaddq_u32(b, s1); + *v2 = vaddq_u32(c, s2); + *v3 = vaddq_u32(d, s3); +} + +/* Set up the constant rows of the state from key + nonce. s3 + * (counter + nonce) varies per block and is built inside the loop. */ +static inline void chacha20_setup(const uint8_t key[32], + const uint8_t nonce[12], + uint32x4_t *s0, uint32x4_t *s1, + uint32x4_t *s2, + uint32_t *n0, uint32_t *n1, + uint32_t *n2) { + *s0 = vld1q_u32(chacha_constants); + *s1 = vreinterpretq_u32_u8(vld1q_u8(key)); + *s2 = vreinterpretq_u32_u8(vld1q_u8(key + 16)); + memcpy(n0, nonce + 0, 4); + memcpy(n1, nonce + 4, 4); + memcpy(n2, nonce + 8, 4); +} + +/* + * Generate one 64-byte ChaCha20 keystream block at 'out'. + */ +void chacha20_block_native(const uint8_t key[32], uint32_t counter, + const uint8_t nonce[12], uint8_t out[64]) { + uint32x4_t s0, s1, s2; + uint32_t n0, n1, n2; + chacha20_setup(key, nonce, &s0, &s1, &s2, &n0, &n1, &n2); + + uint32_t s3_in[4] = { counter, n0, n1, n2 }; + uint32x4_t s3 = vld1q_u32(s3_in); + uint32x4_t v0, v1, v2, v3; + chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + + vst1q_u8(out + 0, vreinterpretq_u8_u32(v0)); + vst1q_u8(out + 16, vreinterpretq_u8_u32(v1)); + vst1q_u8(out + 32, vreinterpretq_u8_u32(v2)); + vst1q_u8(out + 48, vreinterpretq_u8_u32(v3)); +} + +/* + * Encrypt/decrypt 'inlen' bytes at 'in' into 'out' using ChaCha20 + * with the given key, starting counter, and nonce. Stream cipher, + * so the same routine decrypts. + */ +void chacha20_cipher_native(const uint8_t key[32], uint32_t counter, + const uint8_t nonce[12], + const uint8_t *in, uint8_t *out, + size_t inlen) { + uint32x4_t s0, s1, s2; + uint32_t n0, n1, n2; + chacha20_setup(key, nonce, &s0, &s1, &s2, &n0, &n1, &n2); + + size_t pos = 0; + while (pos + 64 <= inlen) { + uint32_t s3_in[4] = { counter, n0, n1, n2 }; + uint32x4_t s3 = vld1q_u32(s3_in); + uint32x4_t v0, v1, v2, v3; + chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + + uint8x16_t i0 = vld1q_u8(in + pos + 0); + uint8x16_t i1 = vld1q_u8(in + pos + 16); + uint8x16_t i2 = vld1q_u8(in + pos + 32); + uint8x16_t i3 = vld1q_u8(in + pos + 48); + + vst1q_u8(out + pos + 0, + veorq_u8(i0, vreinterpretq_u8_u32(v0))); + vst1q_u8(out + pos + 16, + veorq_u8(i1, vreinterpretq_u8_u32(v1))); + vst1q_u8(out + pos + 32, + veorq_u8(i2, vreinterpretq_u8_u32(v2))); + vst1q_u8(out + pos + 48, + veorq_u8(i3, vreinterpretq_u8_u32(v3))); + + pos += 64; + counter++; + } + + /* trailing partial block (< 64 bytes) */ + if (pos < inlen) { + uint32_t s3_in[4] = { counter, n0, n1, n2 }; + uint32x4_t s3 = vld1q_u32(s3_in); + uint32x4_t v0, v1, v2, v3; + chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + + uint8_t block[64]; + vst1q_u8(block + 0, vreinterpretq_u8_u32(v0)); + vst1q_u8(block + 16, vreinterpretq_u8_u32(v1)); + vst1q_u8(block + 32, vreinterpretq_u8_u32(v2)); + vst1q_u8(block + 48, vreinterpretq_u8_u32(v3)); + + size_t remaining = inlen - pos; + for (size_t i = 0; i < remaining; i++) { + out[pos + i] = in[pos + i] ^ block[i]; + } + } +} + +int chacha20_native_available(void) { + return 1; +} + +#elif defined(__x86_64__) + +/* + * ChaCha20 SSE2 kernel, using the same intra-block layout as the NEON + * kernel: the state matrix is held in four 128-bit registers, one per + * row, so that a column round is one set of element-wise vector + * operations. Diagonal rounds are reached by left-rotating rows 1, 2 + * and 3 by 1, 2 and 3 lanes respectively with PSHUFD, then rotating + * back. SSE2 is part of the x86-64 baseline, so no runtime check is + * needed. + */ + +/* 32-bit left rotations. Rotate-by-16 swaps the 16-bit halves of + * each lane with PSHUFLW/PSHUFHW; the others are a shift-shift-or + * pair. */ +#define ROTL32_16(x) \ + _mm_shufflehi_epi16(_mm_shufflelo_epi16((x), 0xb1), 0xb1) +#define ROTL32_N(x, n) \ + _mm_or_si128(_mm_slli_epi32((x), (n)), _mm_srli_epi32((x), 32 - (n))) + +#define QUARTER(v0, v1, v2, v3) \ + do { \ + v0 = _mm_add_epi32(v0, v1); \ + v3 = _mm_xor_si128(v3, v0); v3 = ROTL32_16(v3); \ + v2 = _mm_add_epi32(v2, v3); \ + v1 = _mm_xor_si128(v1, v2); v1 = ROTL32_N(v1, 12); \ + v0 = _mm_add_epi32(v0, v1); \ + v3 = _mm_xor_si128(v3, v0); v3 = ROTL32_N(v3, 8); \ + v2 = _mm_add_epi32(v2, v3); \ + v1 = _mm_xor_si128(v1, v2); v1 = ROTL32_N(v1, 7); \ + } while (0) + +/* 20-round ChaCha20 core: 10 iterations of (column + diagonal). */ +static inline void chacha20_core(__m128i *v0, __m128i *v1, + __m128i *v2, __m128i *v3, + __m128i s0, __m128i s1, + __m128i s2, __m128i s3) { + __m128i a = s0, b = s1, c = s2, d = s3; + for (int i = 0; i < 10; i++) { + QUARTER(a, b, c, d); + /* shift rows: row 1 left 1, row 2 left 2, row 3 left 3. */ + b = _mm_shuffle_epi32(b, _MM_SHUFFLE(0, 3, 2, 1)); + c = _mm_shuffle_epi32(c, _MM_SHUFFLE(1, 0, 3, 2)); + d = _mm_shuffle_epi32(d, _MM_SHUFFLE(2, 1, 0, 3)); + QUARTER(a, b, c, d); + /* shift back. */ + b = _mm_shuffle_epi32(b, _MM_SHUFFLE(2, 1, 0, 3)); + c = _mm_shuffle_epi32(c, _MM_SHUFFLE(1, 0, 3, 2)); + d = _mm_shuffle_epi32(d, _MM_SHUFFLE(0, 3, 2, 1)); + } + *v0 = _mm_add_epi32(a, s0); + *v1 = _mm_add_epi32(b, s1); + *v2 = _mm_add_epi32(c, s2); + *v3 = _mm_add_epi32(d, s3); +} + +/* Set up the constant rows of the state from key + nonce. s3 + * (counter + nonce) varies per block and is built inside the loop. */ +static inline void chacha20_setup(const uint8_t key[32], + const uint8_t nonce[12], + __m128i *s0, __m128i *s1, + __m128i *s2, + uint32_t *n0, uint32_t *n1, + uint32_t *n2) { + *s0 = _mm_loadu_si128((const __m128i *)chacha_constants); + *s1 = _mm_loadu_si128((const __m128i *)key); + *s2 = _mm_loadu_si128((const __m128i *)(key + 16)); + memcpy(n0, nonce + 0, 4); + memcpy(n1, nonce + 4, 4); + memcpy(n2, nonce + 8, 4); +} + +/* Build row 3 (counter + nonce) for one block. */ +static inline __m128i chacha20_row3(uint32_t counter, uint32_t n0, + uint32_t n1, uint32_t n2) { + return _mm_setr_epi32((int)counter, (int)n0, (int)n1, (int)n2); +} + +/* + * ChaCha20 AVX2 kernel: two blocks at once, in the same row layout. + * Each 256-bit register holds one row of the state for two + * consecutive blocks, the low 128-bit lane for block n and the high + * lane for block n + 1. VPSHUFD and VPSHUFB act within each lane, so + * the rounds are those of the SSE2 kernel on wider registers; only + * row 3, whose counters differ, needs per-lane setup. + * + * Each round is one long dependency chain, so long inputs are + * processed four blocks at a time, as two interleaved pairs, to keep + * more of the vector units busy. + * + * AVX2 is not part of the x86-64 baseline, so the kernel is compiled + * for it alone and selected at runtime by 'have_avx2'. + */ + +/* + * Return 1 if the CPU implements AVX2 and the OS saves and restores + * the ymm registers, 0 otherwise. + */ +static int avx2_detect(void) { + unsigned int eax, ebx, ecx, edx; + if (!__get_cpuid(1, &eax, &ebx, &ecx, &edx)) + return 0; + if (!(ecx & bit_OSXSAVE) || !(ecx & bit_AVX)) + return 0; + /* XCR0 bits 1 and 2: xmm and ymm state enabled by the OS */ + unsigned int xcr0_lo, xcr0_hi; + __asm__ volatile ("xgetbv" : "=a"(xcr0_lo), "=d"(xcr0_hi) : "c"(0)); + if ((xcr0_lo & 6) != 6) + return 0; + if (!__get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx)) + return 0; + return (ebx & bit_AVX2) != 0; +} + +/* 'avx2_detect', cached; CPUID can be slow, especially under a + * hypervisor. Concurrent first calls compute the same value. */ +static int have_avx2(void) { + static int cached = -1; + int v = __atomic_load_n(&cached, __ATOMIC_RELAXED); + if (v < 0) { + v = avx2_detect(); + __atomic_store_n(&cached, v, __ATOMIC_RELAXED); + } + return v; +} + +/* 32-bit left rotations. Rotate-by-16 and rotate-by-8 move whole + * bytes, so are single VPSHUFB byte permutations. */ +__attribute__((target("avx2"))) +static inline __m256i rotl32_16_avx2(__m256i x) { + const __m256i idx = _mm256_setr_epi8( + 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13, + 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13); + return _mm256_shuffle_epi8(x, idx); +} + +__attribute__((target("avx2"))) +static inline __m256i rotl32_8_avx2(__m256i x) { + const __m256i idx = _mm256_setr_epi8( + 3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14, + 3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14); + return _mm256_shuffle_epi8(x, idx); +} + +#define ROTL32_N_AVX2(x, n) \ + _mm256_or_si256(_mm256_slli_epi32((x), (n)), \ + _mm256_srli_epi32((x), 32 - (n))) + +#define QUARTER_AVX2(v0, v1, v2, v3) \ + do { \ + v0 = _mm256_add_epi32(v0, v1); \ + v3 = _mm256_xor_si256(v3, v0); v3 = rotl32_16_avx2(v3); \ + v2 = _mm256_add_epi32(v2, v3); \ + v1 = _mm256_xor_si256(v1, v2); v1 = ROTL32_N_AVX2(v1, 12); \ + v0 = _mm256_add_epi32(v0, v1); \ + v3 = _mm256_xor_si256(v3, v0); v3 = rotl32_8_avx2(v3); \ + v2 = _mm256_add_epi32(v2, v3); \ + v1 = _mm256_xor_si256(v1, v2); v1 = ROTL32_N_AVX2(v1, 7); \ + } while (0) + +/* 20-round ChaCha20 core over two blocks. */ +__attribute__((target("avx2"))) +static inline void chacha20_core_avx2(__m256i *v0, __m256i *v1, + __m256i *v2, __m256i *v3, + __m256i s0, __m256i s1, + __m256i s2, __m256i s3) { + __m256i a = s0, b = s1, c = s2, d = s3; + for (int i = 0; i < 10; i++) { + QUARTER_AVX2(a, b, c, d); + b = _mm256_shuffle_epi32(b, _MM_SHUFFLE(0, 3, 2, 1)); + c = _mm256_shuffle_epi32(c, _MM_SHUFFLE(1, 0, 3, 2)); + d = _mm256_shuffle_epi32(d, _MM_SHUFFLE(2, 1, 0, 3)); + QUARTER_AVX2(a, b, c, d); + b = _mm256_shuffle_epi32(b, _MM_SHUFFLE(2, 1, 0, 3)); + c = _mm256_shuffle_epi32(c, _MM_SHUFFLE(1, 0, 3, 2)); + d = _mm256_shuffle_epi32(d, _MM_SHUFFLE(0, 3, 2, 1)); + } + *v0 = _mm256_add_epi32(a, s0); + *v1 = _mm256_add_epi32(b, s1); + *v2 = _mm256_add_epi32(c, s2); + *v3 = _mm256_add_epi32(d, s3); +} + +/* 20-round ChaCha20 core over four blocks, as two interleaved + * pairs so that their independent dependency chains overlap. */ +__attribute__((target("avx2"))) +static inline void chacha20_core4_avx2(__m256i v[8], const __m256i s[8]) { + __m256i a = s[0], b = s[1], c = s[2], d = s[3]; + __m256i e = s[4], f = s[5], g = s[6], h = s[7]; + for (int i = 0; i < 10; i++) { + QUARTER_AVX2(a, b, c, d); + QUARTER_AVX2(e, f, g, h); + b = _mm256_shuffle_epi32(b, _MM_SHUFFLE(0, 3, 2, 1)); + c = _mm256_shuffle_epi32(c, _MM_SHUFFLE(1, 0, 3, 2)); + d = _mm256_shuffle_epi32(d, _MM_SHUFFLE(2, 1, 0, 3)); + f = _mm256_shuffle_epi32(f, _MM_SHUFFLE(0, 3, 2, 1)); + g = _mm256_shuffle_epi32(g, _MM_SHUFFLE(1, 0, 3, 2)); + h = _mm256_shuffle_epi32(h, _MM_SHUFFLE(2, 1, 0, 3)); + QUARTER_AVX2(a, b, c, d); + QUARTER_AVX2(e, f, g, h); + b = _mm256_shuffle_epi32(b, _MM_SHUFFLE(2, 1, 0, 3)); + c = _mm256_shuffle_epi32(c, _MM_SHUFFLE(1, 0, 3, 2)); + d = _mm256_shuffle_epi32(d, _MM_SHUFFLE(0, 3, 2, 1)); + f = _mm256_shuffle_epi32(f, _MM_SHUFFLE(2, 1, 0, 3)); + g = _mm256_shuffle_epi32(g, _MM_SHUFFLE(1, 0, 3, 2)); + h = _mm256_shuffle_epi32(h, _MM_SHUFFLE(0, 3, 2, 1)); + } + v[0] = _mm256_add_epi32(a, s[0]); + v[1] = _mm256_add_epi32(b, s[1]); + v[2] = _mm256_add_epi32(c, s[2]); + v[3] = _mm256_add_epi32(d, s[3]); + v[4] = _mm256_add_epi32(e, s[4]); + v[5] = _mm256_add_epi32(f, s[5]); + v[6] = _mm256_add_epi32(g, s[6]); + v[7] = _mm256_add_epi32(h, s[7]); +} + +/* XOR one block pair's keystream rows into 128 bytes of input. */ +__attribute__((target("avx2"))) +static inline void chacha20_xor2_avx2(const uint8_t *in, uint8_t *out, + __m256i v0, __m256i v1, + __m256i v2, __m256i v3) { + /* gather each block's rows: block n from the low lanes, + * block n + 1 from the high lanes */ + __m256i k0 = _mm256_permute2x128_si256(v0, v1, 0x20); + __m256i k1 = _mm256_permute2x128_si256(v2, v3, 0x20); + __m256i k2 = _mm256_permute2x128_si256(v0, v1, 0x31); + __m256i k3 = _mm256_permute2x128_si256(v2, v3, 0x31); + + __m256i i0 = _mm256_loadu_si256((const __m256i *)(in + 0)); + __m256i i1 = _mm256_loadu_si256((const __m256i *)(in + 32)); + __m256i i2 = _mm256_loadu_si256((const __m256i *)(in + 64)); + __m256i i3 = _mm256_loadu_si256((const __m256i *)(in + 96)); + + _mm256_storeu_si256((__m256i *)(out + 0), _mm256_xor_si256(i0, k0)); + _mm256_storeu_si256((__m256i *)(out + 32), _mm256_xor_si256(i1, k1)); + _mm256_storeu_si256((__m256i *)(out + 64), _mm256_xor_si256(i2, k2)); + _mm256_storeu_si256((__m256i *)(out + 96), _mm256_xor_si256(i3, k3)); +} + +/* + * Encrypt/decrypt the longest prefix of 'in' that is a whole number + * of 128-byte block pairs, starting at 'counter': four blocks at a + * time while at least 256 bytes remain, then two. Returns the number + * of bytes processed. + */ +__attribute__((target("avx2"))) +static size_t chacha20_cipher_avx2(const uint8_t key[32], + uint32_t counter, + const uint8_t nonce[12], + const uint8_t *in, uint8_t *out, + size_t inlen) { + __m256i s0 = _mm256_broadcastsi128_si256( + _mm_loadu_si128((const __m128i *)chacha_constants)); + __m256i s1 = _mm256_broadcastsi128_si256( + _mm_loadu_si128((const __m128i *)key)); + __m256i s2 = _mm256_broadcastsi128_si256( + _mm_loadu_si128((const __m128i *)(key + 16))); + uint32_t n0, n1, n2; + memcpy(&n0, nonce + 0, 4); + memcpy(&n1, nonce + 4, 4); + memcpy(&n2, nonce + 8, 4); + + size_t pos = 0; + while (pos + 256 <= inlen) { + uint32_t c1 = counter + 1, c2 = counter + 2, c3 = counter + 3; + __m256i s[8] = { + s0, s1, s2, + _mm256_setr_epi32((int)counter, (int)n0, (int)n1, (int)n2, + (int)c1, (int)n0, (int)n1, (int)n2), + s0, s1, s2, + _mm256_setr_epi32((int)c2, (int)n0, (int)n1, (int)n2, + (int)c3, (int)n0, (int)n1, (int)n2) + }; + __m256i v[8]; + chacha20_core4_avx2(v, s); + chacha20_xor2_avx2(in + pos, out + pos, v[0], v[1], v[2], v[3]); + chacha20_xor2_avx2(in + pos + 128, out + pos + 128, + v[4], v[5], v[6], v[7]); + pos += 256; + counter += 4; + } + while (pos + 128 <= inlen) { + uint32_t c1 = counter + 1; + __m256i s3 = _mm256_setr_epi32( + (int)counter, (int)n0, (int)n1, (int)n2, + (int)c1, (int)n0, (int)n1, (int)n2); + __m256i v0, v1, v2, v3; + chacha20_core_avx2(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + chacha20_xor2_avx2(in + pos, out + pos, v0, v1, v2, v3); + pos += 128; + counter += 2; + } + return pos; +} + +/* + * Generate one 64-byte ChaCha20 keystream block at 'out'. + */ +void chacha20_block_native(const uint8_t key[32], uint32_t counter, + const uint8_t nonce[12], uint8_t out[64]) { + __m128i s0, s1, s2; + uint32_t n0, n1, n2; + chacha20_setup(key, nonce, &s0, &s1, &s2, &n0, &n1, &n2); + + __m128i s3 = chacha20_row3(counter, n0, n1, n2); + __m128i v0, v1, v2, v3; + chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + + _mm_storeu_si128((__m128i *)(out + 0), v0); + _mm_storeu_si128((__m128i *)(out + 16), v1); + _mm_storeu_si128((__m128i *)(out + 32), v2); + _mm_storeu_si128((__m128i *)(out + 48), v3); +} + +/* + * Encrypt/decrypt 'inlen' bytes at 'in' into 'out' using ChaCha20 + * with the given key, starting counter, and nonce. Stream cipher, + * so the same routine decrypts. + */ +void chacha20_cipher_native(const uint8_t key[32], uint32_t counter, + const uint8_t nonce[12], + const uint8_t *in, uint8_t *out, + size_t inlen) { + __m128i s0, s1, s2; + uint32_t n0, n1, n2; + chacha20_setup(key, nonce, &s0, &s1, &s2, &n0, &n1, &n2); + + /* block pairs via AVX2 where available, the rest via SSE2 */ + size_t pos = 0; + if (inlen >= 128 && have_avx2()) { + pos = chacha20_cipher_avx2(key, counter, nonce, in, out, inlen); + counter += (uint32_t)(pos / 64); + } + + while (pos + 64 <= inlen) { + __m128i s3 = chacha20_row3(counter, n0, n1, n2); + __m128i v0, v1, v2, v3; + chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + + __m128i i0 = _mm_loadu_si128((const __m128i *)(in + pos + 0)); + __m128i i1 = _mm_loadu_si128((const __m128i *)(in + pos + 16)); + __m128i i2 = _mm_loadu_si128((const __m128i *)(in + pos + 32)); + __m128i i3 = _mm_loadu_si128((const __m128i *)(in + pos + 48)); + + _mm_storeu_si128((__m128i *)(out + pos + 0), + _mm_xor_si128(i0, v0)); + _mm_storeu_si128((__m128i *)(out + pos + 16), + _mm_xor_si128(i1, v1)); + _mm_storeu_si128((__m128i *)(out + pos + 32), + _mm_xor_si128(i2, v2)); + _mm_storeu_si128((__m128i *)(out + pos + 48), + _mm_xor_si128(i3, v3)); + + pos += 64; + counter++; + } + + /* trailing partial block (< 64 bytes) */ + if (pos < inlen) { + __m128i s3 = chacha20_row3(counter, n0, n1, n2); + __m128i v0, v1, v2, v3; + chacha20_core(&v0, &v1, &v2, &v3, s0, s1, s2, s3); + + uint8_t block[64]; + _mm_storeu_si128((__m128i *)(block + 0), v0); + _mm_storeu_si128((__m128i *)(block + 16), v1); + _mm_storeu_si128((__m128i *)(block + 32), v2); + _mm_storeu_si128((__m128i *)(block + 48), v3); + + size_t remaining = inlen - pos; + for (size_t i = 0; i < remaining; i++) { + out[pos + i] = in[pos + i] ^ block[i]; + } + } +} + +int chacha20_native_available(void) { + return 1; +} + +#else + +/* stubs for other architectures; never reached because dispatch is + * gated on 'chacha20_native_available' returning 0 */ + +void chacha20_block_native(const uint8_t *key, uint32_t counter, + const uint8_t *nonce, uint8_t *out) { + (void)key; (void)counter; (void)nonce; (void)out; +} + +void chacha20_cipher_native(const uint8_t *key, uint32_t counter, + const uint8_t *nonce, + const uint8_t *in, uint8_t *out, + size_t inlen) { + (void)key; (void)counter; (void)nonce; + (void)in; (void)out; (void)inlen; +} + +int chacha20_native_available(void) { + return 0; +} + +#endif diff --git a/lib/Crypto/Cipher/ChaCha20.hs b/lib/Crypto/Cipher/ChaCha20.hs --- a/lib/Crypto/Cipher/ChaCha20.hs +++ b/lib/Crypto/Cipher/ChaCha20.hs @@ -27,6 +27,7 @@ -- testing , ChaCha(..) , _chacha + , _cipher , _parse_key , _parse_nonce , _quarter @@ -35,7 +36,7 @@ ) where import Control.Monad.ST -import qualified Crypto.Cipher.ChaCha20.Arm as Arm +import qualified Crypto.Cipher.ChaCha20.Native as Native import qualified Data.Bits as B import Data.Bits ((.|.), (.<<.), (.^.)) import qualified Data.ByteString as BS @@ -292,8 +293,8 @@ block key@(BI.PS _ _ kl) counter nonce@(BI.PS _ _ nl) | kl /= 32 = Left InvalidKey | nl /= 12 = Left InvalidNonce - | Arm.chacha20_arm_available = - Right (Arm.block key counter nonce) + | Native.chacha20_native_available = + Right (Native.block key counter nonce) | otherwise = pure $ runST $ do let k = _parse_key key n = _parse_nonce nonce @@ -353,8 +354,8 @@ | kl /= 32 = Left InvalidKey | nl /= 12 = Left InvalidNonce | pl > room = Left CounterOverflow - | Arm.chacha20_arm_available = - Right (Arm.cipher raw_key counter raw_nonce plaintext) + | Native.chacha20_native_available = + Right (Native.cipher raw_key counter raw_nonce plaintext) | otherwise = pure $ runST $ do let key = _parse_key raw_key non = _parse_nonce raw_nonce diff --git a/lib/Crypto/Cipher/ChaCha20/Arm.hs b/lib/Crypto/Cipher/ChaCha20/Arm.hs deleted file mode 100644 --- a/lib/Crypto/Cipher/ChaCha20/Arm.hs +++ /dev/null @@ -1,82 +0,0 @@ -{-# OPTIONS_HADDOCK hide #-} -{-# LANGUAGE BangPatterns #-} - --- | --- Module: Crypto.Cipher.ChaCha20.Arm --- Copyright: (c) 2025 Jared Tobin --- License: MIT --- Maintainer: Jared Tobin --- --- ARM NEON support for the ChaCha20 stream cipher. - -module Crypto.Cipher.ChaCha20.Arm ( - chacha20_arm_available - , block - , cipher - ) where - -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BI -import Data.Word (Word8, Word32) -import Foreign.C.Types (CInt(..), CSize(..)) -import Foreign.ForeignPtr (withForeignPtr) -import Foreign.Ptr (Ptr, plusPtr) -import System.IO.Unsafe (unsafeDupablePerformIO) - --- ffi ------------------------------------------------------------------------ - -foreign import ccall unsafe "chacha20_block_arm" - c_chacha20_block - :: Ptr Word8 -> Word32 -> Ptr Word8 -> Ptr Word8 -> IO () - -foreign import ccall unsafe "chacha20_cipher_arm" - c_chacha20_cipher - :: Ptr Word8 -> Word32 -> Ptr Word8 - -> Ptr Word8 -> Ptr Word8 -> CSize -> IO () - -foreign import ccall unsafe "chacha20_arm_available" - c_chacha20_arm_available :: IO CInt - --- utilities ------------------------------------------------------------------ - -fi :: (Integral a, Num b) => a -> b -fi = fromIntegral -{-# INLINE fi #-} - --- api ------------------------------------------------------------------------ - --- | Are ARM NEON extensions available? -chacha20_arm_available :: Bool -chacha20_arm_available = - unsafeDupablePerformIO c_chacha20_arm_available /= 0 -{-# NOINLINE chacha20_arm_available #-} - --- | One 64-byte ChaCha20 keystream block for the given (already- --- validated) key, counter, and nonce. -block :: BS.ByteString -> Word32 -> BS.ByteString -> BS.ByteString -block (BI.PS kfp koff _) counter (BI.PS nfp noff _) = - BI.unsafeCreate 64 $ \dst -> - withForeignPtr kfp $ \kp0 -> - withForeignPtr nfp $ \np0 -> - c_chacha20_block (kp0 `plusPtr` koff) - counter - (np0 `plusPtr` noff) - dst - --- | XOR the plaintext with the ChaCha20 keystream derived from the --- given (already-validated) key, counter, and nonce. -cipher - :: BS.ByteString -> Word32 -> BS.ByteString -> BS.ByteString - -> BS.ByteString -cipher (BI.PS kfp koff _) counter (BI.PS nfp noff _) - (BI.PS pfp poff plen) = - BI.unsafeCreate plen $ \dst -> - withForeignPtr kfp $ \kp0 -> - withForeignPtr nfp $ \np0 -> - withForeignPtr pfp $ \pp0 -> - c_chacha20_cipher (kp0 `plusPtr` koff) - counter - (np0 `plusPtr` noff) - (pp0 `plusPtr` poff) - dst - (fi plen) diff --git a/lib/Crypto/Cipher/ChaCha20/Native.hs b/lib/Crypto/Cipher/ChaCha20/Native.hs new file mode 100644 --- /dev/null +++ b/lib/Crypto/Cipher/ChaCha20/Native.hs @@ -0,0 +1,82 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module: Crypto.Cipher.ChaCha20.Native +-- Copyright: (c) 2025 Jared Tobin +-- License: MIT +-- Maintainer: Jared Tobin +-- +-- Vectorised ChaCha20, via ARM NEON, or x86 SSE2 and AVX2. + +module Crypto.Cipher.ChaCha20.Native ( + chacha20_native_available + , block + , cipher + ) where + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BI +import Data.Word (Word8, Word32) +import Foreign.C.Types (CInt(..), CSize(..)) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Ptr (Ptr, plusPtr) +import System.IO.Unsafe (unsafeDupablePerformIO) + +-- ffi ------------------------------------------------------------------------ + +foreign import ccall unsafe "chacha20_block_native" + c_chacha20_block + :: Ptr Word8 -> Word32 -> Ptr Word8 -> Ptr Word8 -> IO () + +foreign import ccall unsafe "chacha20_cipher_native" + c_chacha20_cipher + :: Ptr Word8 -> Word32 -> Ptr Word8 + -> Ptr Word8 -> Ptr Word8 -> CSize -> IO () + +foreign import ccall unsafe "chacha20_native_available" + c_chacha20_native_available :: IO CInt + +-- utilities ------------------------------------------------------------------ + +fi :: (Integral a, Num b) => a -> b +fi = fromIntegral +{-# INLINE fi #-} + +-- api ------------------------------------------------------------------------ + +-- | Is a vectorised kernel available? +chacha20_native_available :: Bool +chacha20_native_available = + unsafeDupablePerformIO c_chacha20_native_available /= 0 +{-# NOINLINE chacha20_native_available #-} + +-- | One 64-byte ChaCha20 keystream block for the given (already- +-- validated) key, counter, and nonce. +block :: BS.ByteString -> Word32 -> BS.ByteString -> BS.ByteString +block (BI.PS kfp koff _) counter (BI.PS nfp noff _) = + BI.unsafeCreate 64 $ \dst -> + withForeignPtr kfp $ \kp0 -> + withForeignPtr nfp $ \np0 -> + c_chacha20_block (kp0 `plusPtr` koff) + counter + (np0 `plusPtr` noff) + dst + +-- | XOR the plaintext with the ChaCha20 keystream derived from the +-- given (already-validated) key, counter, and nonce. +cipher + :: BS.ByteString -> Word32 -> BS.ByteString -> BS.ByteString + -> BS.ByteString +cipher (BI.PS kfp koff _) counter (BI.PS nfp noff _) + (BI.PS pfp poff plen) = + BI.unsafeCreate plen $ \dst -> + withForeignPtr kfp $ \kp0 -> + withForeignPtr nfp $ \np0 -> + withForeignPtr pfp $ \pp0 -> + c_chacha20_cipher (kp0 `plusPtr` koff) + counter + (np0 `plusPtr` noff) + (pp0 `plusPtr` poff) + dst + (fi plen) diff --git a/ppad-chacha.cabal b/ppad-chacha.cabal --- a/ppad-chacha.cabal +++ b/ppad-chacha.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: ppad-chacha -version: 0.3.0 +version: 0.3.1 synopsis: A fast ChaCha20 stream cipher license: MIT license-file: LICENSE @@ -37,13 +37,13 @@ exposed-modules: Crypto.Cipher.ChaCha20 other-modules: - Crypto.Cipher.ChaCha20.Arm + Crypto.Cipher.ChaCha20.Native build-depends: base >= 4.9 && < 5 , bytestring >= 0.9 && < 0.13 , primitive >= 0.8 && < 0.10 c-sources: - cbits/chacha20_arm.c + cbits/chacha20_native.c if arch(aarch64) cc-options: -march=armv8-a if flag(sanitize) diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -6,6 +6,7 @@ module Main where +import Control.Monad.ST (runST) import qualified Crypto.Cipher.ChaCha20 as ChaCha import qualified Data.ByteString as BS import qualified Data.ByteString.Base16 as B16 @@ -28,6 +29,7 @@ , crypt3 , counter_boundary , counter_overflow + , dispatch ] quarter :: TestTree @@ -188,3 +190,24 @@ plain = BS.replicate 129 0 out = ChaCha.cipher overflow_key con overflow_non plain H.assertEqual mempty (Left ChaCha.CounterOverflow) out + +-- dispatching vs pure + +-- the dispatching 'cipher' agrees with the pure Haskell implementation +-- for every input length from 0 to 640 bytes, covering partial blocks, +-- short tails, and any multi-block paths. the second counter puts the +-- final block of a 640-byte input at the maximum counter value. +dispatch :: TestTree +dispatch = H.testCase "chacha20 encrypt (dispatch vs pure)" $ + for_ keys $ \key -> + for_ [1, maxBound - 9] $ \con -> + for_ [0..640] $ \l -> do + let msg = BS.take l plain + Right out = ChaCha.cipher key con block_non msg + ref = runST $ + ChaCha._cipher (ChaCha._parse_key key) con + (ChaCha._parse_nonce block_non) msg + H.assertEqual ("length " <> show l) ref out + where + keys = [block_key, BS.replicate 32 0xff] + plain = BS.pack [fromIntegral (i * 7 + 3) | i <- [0..639 :: Int]]