poly1305: Expose 4-block processing on the AVX2 backend

If the Poly1305 state is unaligned (a previous update left the AVX2
backend with 1-3 cached blocks), this falls back to the single-block
API that incurs an additional copy per block.

Author: str4d

poly1305/src/backend/avx2.rs

@@ -16,7 +16,7 @@
 // length to be known, which is incompatible with the streaming API of UniversalHash.
 
 use universal_hash::{
-    consts::{U1, U16},
+    consts::{U16, U4},
     crypto_common::{BlockSizeUser, ParBlocksSizeUser},
     generic_array::GenericArray,
     UhfBackend,
@@ -27,6 +27,9 @@ use crate::{Block, Key, Tag};
 mod helpers;
 use self::helpers::*;
 
+/// Four Poly1305 blocks (64-bytes)
+type ParBlocks = universal_hash::ParBlocks<State>;
+
 #[derive(Copy, Clone)]
 struct Initialized {
     p: Aligned4x130,
@@ -65,6 +68,15 @@ impl State {
         }
     }
 
+    /// Process four Poly1305 blocks at once.
+    #[target_feature(enable = "avx2")]
+    pub(crate) unsafe fn compute_par_blocks(&mut self, blocks: &ParBlocks) {
+        assert!(self.partial_block.is_none());
+        assert_eq!(self.num_cached_blocks, 0);
+
+        self.process_blocks(Aligned4x130::from_par_blocks(blocks));
+    }
+
     /// Compute a Poly1305 block
     #[target_feature(enable = "avx2")]
     pub(crate) unsafe fn compute_block(&mut self, block: &Block, partial: bool) {
@@ -83,13 +95,18 @@ impl State {
             self.num_cached_blocks = 0;
         }
 
+        self.process_blocks(Aligned4x130::from_blocks(&self.cached_blocks));
+    }
+
+    /// Compute a Poly1305 block
+    #[target_feature(enable = "avx2")]
+    unsafe fn process_blocks(&mut self, blocks: Aligned4x130) {
         if let Some(inner) = &mut self.initialized {
             // P <-- R^4 * P + blocks
-            inner.p =
-                (&inner.p * inner.r4).reduce() + Aligned4x130::from_blocks(&self.cached_blocks);
+            inner.p = (&inner.p * inner.r4).reduce() + blocks;
         } else {
             // Initialize the polynomial.
-            let p = Aligned4x130::from_blocks(&self.cached_blocks);
+            let p = blocks;
 
             // Initialize the multiplier (used to merge down the polynomial during
             // finalization).
@@ -160,11 +177,23 @@ impl BlockSizeUser for State {
 }
 
 impl ParBlocksSizeUser for State {
-    type ParBlocksSize = U1;
+    type ParBlocksSize = U4;
 }
 
 impl UhfBackend for State {
     fn proc_block(&mut self, block: &Block) {
         unsafe { self.compute_block(block, false) };
     }
+
+    fn proc_par_blocks(&mut self, blocks: &ParBlocks) {
+        if self.num_cached_blocks == 0 {
+            // Fast path.
+            unsafe { self.compute_par_blocks(blocks) };
+        } else {
+            // We are unaligned; use the slow fallback.
+            for block in blocks {
+                self.proc_block(block);
+            }
+        }
+    }
 }

poly1305/src/backend/avx2/helpers.rs

@@ -8,6 +8,7 @@ use core::arch::x86::*;
 #[cfg(target_arch = "x86_64")]
 use core::arch::x86_64::*;
 
+use super::ParBlocks;
 use crate::{Block, Key};
 
 const fn set02(x3: u8, x2: u8, x1: u8, x0: u8) -> i32 {
@@ -964,25 +965,44 @@ impl Aligned4x130 {
     /// Panics if `src.len() < 64`.
     #[target_feature(enable = "avx2")]
     pub(super) unsafe fn from_blocks(src: &[Block; 4]) -> Self {
+        let (lo, hi) = src.split_at(2);
+        let blocks_23 = _mm256_loadu_si256(hi.as_ptr() as *const _);
+        let blocks_01 = _mm256_loadu_si256(lo.as_ptr() as *const _);
+
+        Self::from_loaded_blocks(blocks_01, blocks_23)
+    }
+
+    /// Aligns four 16-byte Poly1305 blocks at 26-bit boundaries within 32-bit words, and
+    /// sets the high bit for each block.
+    #[target_feature(enable = "avx2")]
+    pub(super) unsafe fn from_par_blocks(src: &ParBlocks) -> Self {
+        let (lo, hi) = src.split_at(2);
+        let blocks_23 = _mm256_loadu_si256(hi.as_ptr() as *const _);
+        let blocks_01 = _mm256_loadu_si256(lo.as_ptr() as *const _);
+
+        Self::from_loaded_blocks(blocks_01, blocks_23)
+    }
+
+    /// Aligns four 16-byte Poly1305 blocks at 26-bit boundaries within 32-bit words, and
+    /// sets the high bit for each block.
+    ///
+    /// The four blocks must be in the following 32-bit word layout:
+    ///      [b33, b32, b31, b30, b23, b22, b21, b20]
+    ///      [b13, b12, b11, b10, b03, b02, b01, b00]
+    #[target_feature(enable = "avx2")]
+    unsafe fn from_loaded_blocks(blocks_01: __m256i, blocks_23: __m256i) -> Self {
         // 26-bit mask on each 32-bit word.
         let mask_26 = _mm256_set1_epi32(0x3ffffff);
         // Sets bit 24 of each 32-bit word.
         let set_hibit = _mm256_set1_epi32(1 << 24);
 
-        // - Load the four blocks into the following 32-bit word layout:
-        //      [b33, b32, b31, b30, b23, b22, b21, b20]
-        //      [b13, b12, b11, b10, b03, b02, b01, b00]
-        //
         // - Unpack the upper and lower 64 bits:
         //      [b33, b32, b13, b12, b23, b22, b03, b02]
         //      [b31, b30, b11, b10, b21, b20, b01, b00]
         //
         // - Swap the middle two 64-bit words:
         // a0 = [b33, b32, b23, b22, b13, b12, b03, b02]
         // a1 = [b31, b30, b21, b20, b11, b10, b01, b00]
-        let (lo, hi) = src.split_at(2);
-        let blocks_23 = _mm256_loadu_si256(hi.as_ptr() as *const _);
-        let blocks_01 = _mm256_loadu_si256(lo.as_ptr() as *const _);
         let a0 = _mm256_permute4x64_epi64(
             _mm256_unpackhi_epi64(blocks_01, blocks_23),
             set02(3, 1, 2, 0),