Add RNG support for avx512 (not benchmarked)

Author: caelunshun

chacha20/Cargo.toml

@@ -20,9 +20,7 @@ rand_core-compatible RNGs based on those ciphers.
 
 [dependencies]
 cfg-if = "1"
-cipher = { version = "0.5.0-rc.1", optional = true, features = [
-    "stream-wrapper",
-] }
+cipher = { version = "0.5.0-rc.1", optional = true, features = ["stream-wrapper"] }
 rand_core = { version = "0.10.0-rc.1", optional = true, default-features = false }
 
 # `zeroize` is an explicit dependency because this crate may be used without the `cipher` crate

chacha20/src/backends.rs

@@ -7,7 +7,9 @@ cfg_if! {
         pub(crate) mod soft;
     } else if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
         cfg_if! {
-            if #[cfg(chacha20_force_avx2)] {
+            if #[cfg(chacha20_force_avx512)] {
+                pub(crate) mod avx512;
+            } else if #[cfg(chacha20_force_avx2)] {
                 pub(crate) mod avx2;
             } else if #[cfg(chacha20_force_sse2)] {
                 pub(crate) mod sse2;

chacha20/src/backends/avx512.rs

@@ -85,36 +85,46 @@ where
 }
 
 #[inline]
-#[target_feature(enable = "avx512")]
+#[target_feature(enable = "avx512f")]
 #[cfg(feature = "rng")]
 pub(crate) unsafe fn rng_inner<R, V>(core: &mut ChaChaCore<R, V>, buffer: &mut [u32; 64])
 where
     R: Rounds,
     V: Variant,
 {
+    use core::slice;
+
+    use crate::rng::BLOCK_WORDS;
+
     let state_ptr = core.state.as_ptr() as *const __m128i;
     let v = [
-        _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(0))),
-        _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(1))),
-        _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(2))),
+        _mm512_broadcast_i32x4(_mm_loadu_si128(state_ptr.add(0))),
+        _mm512_broadcast_i32x4(_mm_loadu_si128(state_ptr.add(1))),
+        _mm512_broadcast_i32x4(_mm_loadu_si128(state_ptr.add(2))),
     ];
-    let mut c = _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(3)));
-    c = _mm256_add_epi64(c, _mm256_set_epi64x(0, 1, 0, 0));
-    let mut ctr = [c; N];
-    for i in 0..N {
+    let mut c = _mm512_broadcast_i32x4(_mm_loadu_si128(state_ptr.add(3)));
+    c = _mm512_add_epi64(c, _mm512_set_epi64(0, 3, 0, 2, 0, 1, 0, 0));
+    let mut ctr = [c; MAX_N];
+    for i in 0..MAX_N {
         ctr[i] = c;
-        c = _mm256_add_epi64(c, _mm256_set_epi64x(0, 2, 0, 2));
+        c = _mm512_add_epi64(c, _mm512_set_epi64(0, 4, 0, 4, 0, 4, 0, 4));
     }
     let mut backend = Backend::<R, V> {
         v,
         ctr,
         _pd: PhantomData,
     };
 
-    backend.rng_gen_par_ks_blocks(buffer);
+    let buffer = slice::from_raw_parts_mut(
+        buffer.as_mut_ptr().cast::<Block>(),
+        buffer.len() / BLOCK_WORDS as usize,
+    );
+    backend.gen_par_ks_blocks_inner::<4, { 4 / BLOCKS_PER_VECTOR }>(buffer.try_into().unwrap());
 
-    core.state[12] = _mm256_extract_epi32(backend.ctr[0], 0) as u32;
-    core.state[13] = _mm256_extract_epi32(backend.ctr[0], 1) as u32;
+    core.state[12] =
+        _mm256_extract_epi32::<0>(_mm512_extracti32x8_epi32::<0>(backend.ctr[0])) as u32;
+    core.state[13] =
+        _mm256_extract_epi32::<1>(_mm512_extracti32x8_epi32::<0>(backend.ctr[0])) as u32;
 }
 
 struct Backend<R: Rounds, V: Variant> {
@@ -233,31 +243,6 @@ impl<R: Rounds, V: Variant> StreamCipherBackend for Backend<R, V> {
     }
 }
 
-#[cfg(feature = "rng")]
-impl<R: Rounds, V: Variant> Backend<R, V> {
-    #[inline(always)]
-    fn rng_gen_par_ks_blocks(&mut self, blocks: &mut [u32; 64]) {
-        unsafe {
-            let vs = rounds::<R>(&self.v, &self.ctr);
-
-            let pb = PAR_BLOCKS as i32;
-            for c in self.ctr.iter_mut() {
-                *c = _mm256_add_epi64(*c, _mm256_set_epi64x(0, pb as i64, 0, pb as i64));
-            }
-
-            let mut block_ptr = blocks.as_mut_ptr() as *mut __m128i;
-            for v in vs {
-                let t: [__m128i; 8] = core::mem::transmute(v);
-                for i in 0..4 {
-                    _mm_storeu_si128(block_ptr.add(i), t[2 * i]);
-                    _mm_storeu_si128(block_ptr.add(4 + i), t[2 * i + 1]);
-                }
-                block_ptr = block_ptr.add(8);
-            }
-        }
-    }
-}
-
 #[inline]
 #[target_feature(enable = "avx512f")]
 unsafe fn rounds<const N: usize, R: Rounds>(

chacha20/src/lib.rs

@@ -185,7 +185,12 @@ cfg_if! {
         type Tokens = ();
     } else if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
         cfg_if! {
-            if #[cfg(chacha20_force_avx2)] {
+            if #[cfg(chacha20_force_avx512)] {
+                #[cfg(not(target_feature = "avx512f"))]
+                compile_error!("You must enable `avx512f` target feature with \
+                    `chacha20_force_avx512` configuration option");
+                type Tokens = ();
+            } else if #[cfg(chacha20_force_avx2)] {
                 #[cfg(not(target_feature = "avx2"))]
                 compile_error!("You must enable `avx2` target feature with \
                     `chacha20_force_avx2` configuration option");
@@ -248,7 +253,9 @@ impl<R: Rounds, V: Variant> ChaChaCore<R, V> {
                 let tokens = ();
             } else if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
                 cfg_if! {
-                    if #[cfg(chacha20_force_avx2)] {
+                    if #[cfg(chacha20_force_avx512)] {
+                        let tokens = ();
+                    } else if #[cfg(chacha20_force_avx2)] {
                         let tokens = ();
                     } else if #[cfg(chacha20_force_sse2)] {
                         let tokens = ();
@@ -299,7 +306,11 @@ impl<R: Rounds, V: Variant> StreamCipherCore for ChaChaCore<R, V> {
                 f.call(&mut backends::soft::Backend(self));
             } else if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
                 cfg_if! {
-                    if #[cfg(chacha20_force_avx2)] {
+                    if #[cfg(chacha20_force_avx512)] {
+                        unsafe {
+                            backends::avx512::inner::<R, _, V>(&mut self.state, f);
+                        }
+                    } else if #[cfg(chacha20_force_avx2)] {
                         unsafe {
                             backends::avx2::inner::<R, _, V>(&mut self.state, f);
                         }

chacha20/src/rng.rs

@@ -189,7 +189,11 @@ impl<R: Rounds, V: Variant> ChaChaCore<R, V> {
                 backends::soft::Backend(self).gen_ks_blocks(buffer);
             } else if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
                 cfg_if! {
-                    if #[cfg(chacha20_force_avx2)] {
+                    if #[cfg(chacha20_force_avx512)] {
+                        unsafe {
+                            backends::avx512::rng_inner::<R, V>(self, buffer);
+                        }
+                    } else if #[cfg(chacha20_force_avx2)] {
                         unsafe {
                             backends::avx2::rng_inner::<R, V>(self, buffer);
                         }
@@ -198,8 +202,12 @@ impl<R: Rounds, V: Variant> ChaChaCore<R, V> {
                             backends::sse2::rng_inner::<R, V>(self, buffer);
                         }
                     } else {
-                        let (avx2_token, sse2_token) = self.tokens;
-                        if avx2_token.get() {
+                        let (avx512_token, avx2_token, sse2_token) = self.tokens;
+                        if avx512_token.get() {
+                            unsafe {
+                                backends::avx512::rng_inner::<R, V>(self, buffer);
+                            }
+                        } else if avx2_token.get() {
                             unsafe {
                                 backends::avx2::rng_inner::<R, V>(self, buffer);
                             }