commit ffebd11fe9272f4ec3cbd7d331d0923076f5851c
parent c36cd59073c97a3cde999102ffd5dca7ccd7c02f
Author: Jared Tobin <jared@jtobin.io>
Date: Sat, 31 Oct 2026 14:35:13 +0400
cbits: add an x86 AVX2 kernel
On x86-64 CPUs with AVX2, chacha20_cipher_native now encrypts block
pairs with 256-bit registers, keeping the row layout of the SSE2
kernel: each register holds one row for two consecutive blocks, one
per 128-bit lane. VPSHUFD and VPSHUFB act within lanes, so the rounds
carry over unchanged, with rotate-by-16 and rotate-by-8 as single
byte shuffles. Inputs of 256 bytes or more go four blocks at a time,
as two interleaved pairs, since each round is otherwise a single
dependency chain. Whatever is left after the last whole pair, and the
block function, stay on SSE2.
AVX2 support (CPUID, plus XGETBV to confirm the OS saves the ymm
registers) is detected on first use and cached, and the kernel is
compiled for AVX2 alone via a target attribute, so CPUs without it
keep the SSE2 path. Counters are uint32_t, so per-lane counters wrap
as they do in the scalar path.
Diffstat:
2 files changed, 227 insertions(+), 2 deletions(-)
diff --git a/cbits/chacha20_native.c b/cbits/chacha20_native.c
@@ -8,7 +8,8 @@
#elif defined(__x86_64__)
-#include <emmintrin.h>
+#include <cpuid.h>
+#include <immintrin.h>
#endif
@@ -265,6 +266,224 @@ static inline __m128i chacha20_row3(uint32_t counter, uint32_t n0,
}
/*
+ * 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,
@@ -296,7 +515,13 @@ void chacha20_cipher_native(const uint8_t key[32], uint32_t counter,
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;
diff --git a/lib/Crypto/Cipher/ChaCha20/Native.hs b/lib/Crypto/Cipher/ChaCha20/Native.hs
@@ -7,7 +7,7 @@
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
--- Vectorised ChaCha20, via ARM NEON or x86 SSE2.
+-- Vectorised ChaCha20, via ARM NEON, or x86 SSE2 and AVX2.
module Crypto.Cipher.ChaCha20.Native (
chacha20_native_available
