chacha20: AVX2 detection

Automatically detects the availability of AVX2 on x86(_64) architectures
based on CPUID, and falls back to the SSE2" backend if unavailable.

Author: tarcieri

.github/workflows/chacha20.yml

@@ -31,14 +31,16 @@ jobs:
       - uses: actions/checkout@v1
       - uses: actions-rs/toolchain@v1
         with:
-          profile: minimal
           toolchain: ${{ matrix.rust }}
           target: ${{ matrix.target }}
           override: true
-      - run: cargo build --no-default-features --release --target ${{ matrix.target }}
-        env:
-          CARGO_INCREMENTAL: 0
-          RUSTFLAGS: -D warnings
+          profile: minimal
+      - run: cargo build --target ${{ matrix.target }} --release --no-default-features --features cipher
+      - run: cargo build --target ${{ matrix.target }} --release --no-default-features --features force-soft
+      - run: cargo build --target ${{ matrix.target }} --release --no-default-features --features legacy
+      - run: cargo build --target ${{ matrix.target }} --release --no-default-features --features rng
+      - run: cargo build --target ${{ matrix.target }} --release --no-default-features --features xchacha20
+      - run: cargo build --target ${{ matrix.target }} --release --no-default-features --features cipher,force-soft,legacy,rng,xchacha20,zeroize
 
   test:
     runs-on: ubuntu-latest

Cargo.lock

@@ -17,19 +17,24 @@ readme = "README.md"
 edition = "2018"
 
 [dependencies]
+cfg-if = "1"
 cipher = { version = "=0.3.0-pre", optional = true }
 rand_core = { version = "0.5", optional = true, default-features = false }
 zeroize = { version = "1", optional = true, default-features = false }
 
+[target.'cfg(any(target_arch = "x86_64", target_arch = "x86"))'.dependencies]
+cpuid-bool = "0.2"
+
 [dev-dependencies]
 cipher = { version = "=0.3.0-pre", features = ["dev"] }
 hex-literal = "0.2"
 
 [features]
 default = ["xchacha20"]
+force-soft = []
 legacy = ["cipher"]
-xchacha20 = ["cipher"]
 rng = ["rand_core"]
+xchacha20 = ["cipher"]
 
 [package.metadata.docs.rs]
 all-features = true

chacha20/Cargo.toml

@@ -2,50 +2,26 @@
 //!
 //! <https://tools.ietf.org/html/rfc8439#section-2.3>
 
-// TODO(tarcieri): figure out what circumstances these occur in
-#![allow(unused_imports)]
-
-pub(crate) mod soft;
-
-use crate::rounds::Rounds;
-
-#[cfg(all(
-    any(target_arch = "x86", target_arch = "x86_64"),
-    target_feature = "sse2",
-    not(target_feature = "avx2")
-))]
-mod sse2;
-
-#[cfg(all(
-    any(target_arch = "x86", target_arch = "x86_64"),
-    target_feature = "avx2"
-))]
-mod avx2;
-
-#[cfg(not(all(
-    any(target_arch = "x86", target_arch = "x86_64"),
-    any(target_feature = "sse2", target_feature = "avx2")
-)))]
-pub(crate) use self::soft::{State, BUFFER_SIZE};
-
-#[cfg(all(
-    any(target_arch = "x86", target_arch = "x86_64"),
-    target_feature = "sse2",
-    not(target_feature = "avx2")
-))]
-pub(crate) use self::sse2::{State, BUFFER_SIZE};
-
-#[cfg(all(
-    any(target_arch = "x86", target_arch = "x86_64"),
-    target_feature = "avx2"
-))]
-pub(crate) use self::avx2::{State, BUFFER_SIZE};
-
-use core::fmt::{self, Debug};
-
-/// Common debug impl for all blocks
-impl<R: Rounds> Debug for State<R> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
-        f.write_str("State {{  .. }}")
+use cfg_if::cfg_if;
+
+cfg_if! {
+    if #[cfg(all(
+        any(target_arch = "x86", target_arch = "x86_64"),
+        not(feature = "force-soft")
+    ))] {
+        pub(crate) mod autodetect;
+        pub(crate) mod avx2;
+        pub(crate) mod sse2;
+
+        #[cfg(feature = "cipher")]
+        pub(crate) use self::autodetect::{State, BUFFER_SIZE};
+
+        #[cfg(feature = "xchacha20")]
+        pub(crate) mod soft;
+    } else {
+        pub(crate) mod soft;
+
+        #[cfg(feature = "cipher")]
+        pub(crate) use self::soft::{State, BUFFER_SIZE};
     }
 }

chacha20/src/backend.rs

@@ -0,0 +1,81 @@
+//! Autodetection support for AVX2 CPU intrinsics on x86 CPUs, with fallback
+//! to the SSE2 backend when it's unavailable (the `sse2` target feature is
+//! enabled-by-default on all x86(_64) CPUs)
+
+use crate::{rounds::Rounds, IV_SIZE, KEY_SIZE, BLOCK_SIZE};
+use super::{avx2, sse2};
+
+/// Size of buffers passed to `generate` and `apply_keystream` for this
+/// backend, which operates on two blocks in parallel for optimal performance.
+pub(crate) const BUFFER_SIZE: usize = BLOCK_SIZE * 2;
+
+cpuid_bool::new!(avx2_cpuid, "avx2");
+
+pub struct State<R: Rounds> {
+    inner: Inner<R>,
+    token: avx2_cpuid::InitToken,
+}
+
+union Inner<R: Rounds> {
+    avx2: avx2::State<R>,
+    sse2: sse2::State<R>,
+}
+
+impl<R: Rounds> State<R> {
+    /// Initialize ChaCha block function with the given key size, IV, and
+    /// number of rounds.
+    #[inline]
+    pub(crate) fn new(key: &[u8; KEY_SIZE], iv: [u8; IV_SIZE]) -> Self {
+        let (token, avx2_present) = avx2_cpuid::init_get();
+
+        let inner = if avx2_present {
+            Inner {
+                avx2: avx2::State::new(key, iv),
+            }
+        } else {
+            Inner {
+                sse2: sse2::State::new(key, iv),
+            }
+        };
+
+        Self { inner, token }
+    }
+
+    #[inline]
+    pub(crate) fn generate(&self, counter: u64, output: &mut [u8]) {
+        if self.token.get() {
+            unsafe { self.inner.avx2.generate(counter, output) }
+        } else {
+            unsafe { self.inner.sse2.generate(counter, output) }
+        }
+    }
+
+    #[inline]
+    #[cfg(feature = "cipher")]
+    pub(crate) fn apply_keystream(&self, counter: u64, output: &mut [u8]) {
+        if self.token.get() {
+            unsafe { self.inner.avx2.apply_keystream(counter, output) }
+        } else {
+            unsafe { self.inner.sse2.apply_keystream(counter, output) }
+        }
+    }
+}
+
+impl<R: Rounds> Clone for State<R> {
+    fn clone(&self) -> Self {
+        let inner = if self.token.get() {
+            Inner {
+                avx2: unsafe { self.inner.avx2.clone() },
+            }
+        } else {
+            Inner {
+                sse2: unsafe { self.inner.sse2.clone() },
+            }
+        };
+
+        Self {
+            inner,
+            token: self.token,
+        }
+    }
+}

chacha20/src/backend/autodetect.rs

@@ -10,19 +10,16 @@
 
 use crate::{rounds::Rounds, BLOCK_SIZE, CONSTANTS, IV_SIZE, KEY_SIZE};
 use core::{convert::TryInto, marker::PhantomData};
+use super::autodetect::BUFFER_SIZE;
 
 #[cfg(target_arch = "x86")]
 use core::arch::x86::*;
 #[cfg(target_arch = "x86_64")]
 use core::arch::x86_64::*;
 
-/// Size of buffers passed to `generate` and `apply_keystream` for this
-/// backend, which operates on two blocks in parallel for optimal performance.
-pub(crate) const BUFFER_SIZE: usize = BLOCK_SIZE * 2;
-
 /// The ChaCha20 block function (AVX2 accelerated implementation for x86/x86_64)
 // TODO(tarcieri): zeroize?
-#[derive(Clone)]
+#[derive(Copy, Clone)]
 pub(crate) struct State<R: Rounds> {
     v0: __m256i,
     v1: __m256i,

chacha20/src/backend/avx2.rs

@@ -9,12 +9,13 @@ use crate::{rounds::Rounds, BLOCK_SIZE, CONSTANTS, IV_SIZE, KEY_SIZE, STATE_WORD
 use core::{convert::TryInto, marker::PhantomData};
 
 /// Size of buffers passed to `generate` and `apply_keystream` for this backend
+#[allow(dead_code)]
 pub(crate) const BUFFER_SIZE: usize = BLOCK_SIZE;
 
 /// The ChaCha20 block function (portable software implementation)
 // TODO(tarcieri): zeroize?
-#[allow(dead_code)]
 #[derive(Clone)]
+#[allow(dead_code)]
 pub(crate) struct State<R: Rounds> {
     /// Internal state of the block function
     state: [u32; STATE_WORDS],

chacha20/src/backend/soft.rs

@@ -4,20 +4,18 @@
 //!
 //! SSE2-optimized implementation for x86/x86-64 CPUs.
 
-use crate::{rounds::Rounds, BLOCK_SIZE, CONSTANTS, IV_SIZE, KEY_SIZE};
+use crate::{rounds::Rounds, CONSTANTS, IV_SIZE, BLOCK_SIZE, KEY_SIZE};
 use core::{convert::TryInto, marker::PhantomData};
+use super::autodetect::BUFFER_SIZE;
 
 #[cfg(target_arch = "x86")]
 use core::arch::x86::*;
 #[cfg(target_arch = "x86_64")]
 use core::arch::x86_64::*;
 
-/// Size of buffers passed to `generate` and `apply_keystream` for this backend
-pub(crate) const BUFFER_SIZE: usize = BLOCK_SIZE;
-
 /// The ChaCha20 block function (SSE2 accelerated implementation for x86/x86_64)
 // TODO(tarcieri): zeroize?
-#[derive(Clone)]
+#[derive(Copy, Clone)]
 pub(crate) struct State<R: Rounds> {
     v0: __m128i,
     v1: __m128i,
@@ -49,11 +47,13 @@ impl<R: Rounds> State<R> {
     pub(crate) fn generate(&self, counter: u64, output: &mut [u8]) {
         debug_assert_eq!(output.len(), BUFFER_SIZE);
 
-        unsafe {
-            let (mut v0, mut v1, mut v2) = (self.v0, self.v1, self.v2);
-            let mut v3 = iv_setup(self.iv, counter);
-            self.rounds(&mut v0, &mut v1, &mut v2, &mut v3);
-            store(v0, v1, v2, v3, output)
+        for (i, chunk) in output.chunks_exact_mut(BLOCK_SIZE).enumerate() {
+            unsafe {
+                let (mut v0, mut v1, mut v2) = (self.v0, self.v1, self.v2);
+                let mut v3 = iv_setup(self.iv, counter.checked_add(i as u64).unwrap());
+                self.rounds(&mut v0, &mut v1, &mut v2, &mut v3);
+                store(v0, v1, v2, v3, chunk)
+            }
         }
     }
 
@@ -63,15 +63,17 @@ impl<R: Rounds> State<R> {
     pub(crate) fn apply_keystream(&self, counter: u64, output: &mut [u8]) {
         debug_assert_eq!(output.len(), BUFFER_SIZE);
 
-        unsafe {
-            let (mut v0, mut v1, mut v2) = (self.v0, self.v1, self.v2);
-            let mut v3 = iv_setup(self.iv, counter);
-            self.rounds(&mut v0, &mut v1, &mut v2, &mut v3);
-
-            for (chunk, a) in output.chunks_mut(0x10).zip(&[v0, v1, v2, v3]) {
-                let b = _mm_loadu_si128(chunk.as_ptr() as *const __m128i);
-                let out = _mm_xor_si128(*a, b);
-                _mm_storeu_si128(chunk.as_mut_ptr() as *mut __m128i, out);
+        for (i, chunk) in output.chunks_exact_mut(BLOCK_SIZE).enumerate() {
+            unsafe {
+                let (mut v0, mut v1, mut v2) = (self.v0, self.v1, self.v2);
+                let mut v3 = iv_setup(self.iv, counter.checked_add(i as u64).unwrap());
+                self.rounds(&mut v0, &mut v1, &mut v2, &mut v3);
+
+                for (ch, a) in chunk.chunks_exact_mut(0x10).zip(&[v0, v1, v2, v3]) {
+                    let b = _mm_loadu_si128(ch.as_ptr() as *const __m128i);
+                    let out = _mm_xor_si128(*a, b);
+                    _mm_storeu_si128(ch.as_mut_ptr() as *mut __m128i, out);
+                }
             }
         }
     }
@@ -263,12 +265,12 @@ mod tests {
 
     #[test]
     fn generate_vs_scalar_impl() {
-        let mut soft_result = [0u8; BLOCK_SIZE];
+        let mut soft_result = [0u8; soft::BUFFER_SIZE];
         soft::State::<R20>::new(&R_KEY, R_IV).generate(R_CNT, &mut soft_result);
 
-        let mut simd_result = [0u8; BLOCK_SIZE];
+        let mut simd_result = [0u8; BUFFER_SIZE];
         State::<R20>::new(&R_KEY, R_IV).generate(R_CNT, &mut simd_result);
 
-        assert_eq!(&soft_result[..], &simd_result[..])
+        assert_eq!(&soft_result[..], &simd_result[..soft::BUFFER_SIZE])
     }
 }

chacha20/src/backend/sse2.rs

@@ -12,7 +12,7 @@ use crate::{
 macro_rules! impl_chacha_rng {
     ($name:ident, $core:ident, $rounds:ident, $doc:expr) => {
         #[doc = $doc]
-        #[derive(Clone, Debug)]
+        #[derive(Clone)]
         #[cfg_attr(docsrs, doc(cfg(feature = "rng")))]
         pub struct $name(BlockRng<$core>);
 
@@ -51,10 +51,10 @@ macro_rules! impl_chacha_rng {
         impl CryptoRng for $name {}
 
         #[doc = "Core random number generator, for use with [`rand_core::block::BlockRng`]"]
-        #[derive(Clone, Debug)]
+        #[derive(Clone)]
         #[cfg_attr(docsrs, doc(cfg(feature = "rng")))]
         pub struct $core {
-            block: Block<$rounds>,
+            block: State<$rounds>,
             counter: u64,
         }
 
@@ -63,7 +63,7 @@ macro_rules! impl_chacha_rng {
 
             #[inline]
             fn from_seed(seed: Self::Seed) -> Self {
-                let block = Block::new(&seed, Default::default());
+                let block = State::new(&seed, Default::default());
                 Self { block, counter: 0 }
             }
         }
@@ -75,7 +75,7 @@ macro_rules! impl_chacha_rng {
             fn generate(&mut self, results: &mut Self::Results) {
                 assert!(self.counter <= MAX_BLOCKS as u64, "maximum number of allowed ChaCha blocks exceeded");
 
-                // TODO(tarcieri): eliminate unsafety (replace w\ [u8; BLOCK_SIZE)
+                // TODO(tarcieri): eliminate unsafety (replace w\ `[u8; BLOCK_SIZE]`)
                 self.block.generate(self.counter, unsafe {
                     &mut *(results.as_mut_ptr() as *mut [u8; BUFFER_SIZE])
                 });