Second Attempt of VAES

.gitignore

@@ -1,3 +1,4 @@
 .idea/
+.vscode/
 Cargo.lock
 target

Cargo.toml

@@ -47,6 +47,9 @@ atomic-polyfill = [ "dep:atomic-polyfill", "once_cell/atomic-polyfill"]
 # Nightly-only support for AES intrinsics on 32-bit ARM
 nightly-arm-aes = []
 
+# Nightly-only support for VAES intrinsics with 256 SIMD registers
+vaes = []
+
 [[bench]]
 name = "ahash"
 path = "tests/bench.rs"

smhasher/ahash-cbindings/Cargo.toml

@@ -17,4 +17,7 @@ lto = 'fat'
 debug-assertions = false
 
 [dependencies]
-ahash = { path = "../../", default-features = false }
+ahash = { path = "../../", default-features = false }
+
+[features]
+vaes = ["ahash/vaes"]

smhasher/ahash-cbindings/install.sh

@@ -1 +1,9 @@
-RUSTFLAGS="-C opt-level=3 -C target-cpu=native -C codegen-units=1" cargo build --release && sudo cp target/release/libahash_c.a /usr/local/lib/
+
+# check if args contains vaes
+if [[ $* == *vaes* ]]; then
+    export CARGO_OPTS="--features=vaes"
+else
+    export CARGO_OPTS=""
+fi
+
+RUSTFLAGS="-C opt-level=3 -C target-cpu=native -C codegen-units=1" cargo build ${CARGO_OPTS} --release && sudo cp target/release/libahash_c.a /usr/local/lib/

smhasher/ahash-cbindings/src/lib.rs

@@ -1,6 +1,5 @@
 use ahash::*;
 use core::slice;
-use std::hash::{BuildHasher};
 
 #[no_mangle]
 pub extern "C" fn ahash64(buf: *const (), len: usize, seed: u64) -> u64 {

src/aes_hash.rs

@@ -22,6 +22,55 @@ pub struct AHasher {
     key: u128,
 }
 
+#[cfg(any(
+    all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes", not(miri)),
+    all(target_arch = "aarch64", target_feature = "aes", not(miri)),
+    all(feature = "nightly-arm-aes", target_arch = "arm", target_feature = "aes", not(miri)),
+))]
+fn hash_batch_128b(data: &mut &[u8], hasher: &mut AHasher) {
+    let tail = read_last4_vec256(data);
+    let mut current = [
+        aesenc_vec256(convert_u128_to_vec256(hasher.key), tail[0]),
+        aesdec_vec256(convert_u128_to_vec256(hasher.key), tail[1]),
+        aesenc_vec256(convert_u128_to_vec256(hasher.key), tail[2]),
+        aesdec_vec256(convert_u128_to_vec256(hasher.key), tail[3]),
+    ];
+    let mut sum: [Vector256; 2] = [convert_u128_to_vec256(hasher.key), convert_u128_to_vec256(!hasher.key)];
+    sum[0] = add_by_64s_vec256(sum[0], tail[0]);
+    sum[1] = add_by_64s_vec256(sum[1], tail[1]);
+    sum[0] = shuffle_and_add_vec256(sum[0], tail[2]);
+    sum[1] = shuffle_and_add_vec256(sum[1], tail[3]);
+    while data.len() > 128 {
+        let (blocks, rest) = read4_vec256(data);
+        current[0] = aesdec_vec256(current[0], blocks[0]);
+        current[1] = aesdec_vec256(current[1], blocks[1]);
+        current[2] = aesdec_vec256(current[2], blocks[2]);
+        current[3] = aesdec_vec256(current[3], blocks[3]);
+        sum[0] = shuffle_and_add_vec256(sum[0], blocks[0]);
+        sum[1] = shuffle_and_add_vec256(sum[1], blocks[1]);
+        sum[0] = shuffle_and_add_vec256(sum[0], blocks[2]);
+        sum[1] = shuffle_and_add_vec256(sum[1], blocks[3]);
+        *data = rest;
+    }
+    let current = [
+        convert_vec256_to_u128(current[0]),
+        convert_vec256_to_u128(current[1]),
+        convert_vec256_to_u128(current[2]),
+        convert_vec256_to_u128(current[3]),
+    ];
+    let sum = [convert_vec256_to_u128(sum[0]), convert_vec256_to_u128(sum[1])];
+    hasher.hash_in_2(
+        aesenc(current[0][0], current[0][1]),
+        aesenc(current[1][0], current[1][1]),
+    );
+    hasher.hash_in(add_by_64s(sum[0][0].convert(), sum[0][1].convert()).convert());
+    hasher.hash_in_2(
+        aesenc(current[2][0], current[2][1]),
+        aesenc(current[3][0], current[3][1]),
+    );
+    hasher.hash_in(add_by_64s(sum[1][0].convert(), sum[1][1].convert()).convert());
+}
+
 impl AHasher {
     /// Creates a new hasher keyed to the provided keys.
     ///
@@ -47,6 +96,7 @@ impl AHasher {
     ///
     /// println!("Hash is {:x}!", hasher.finish());
     /// ```
+    #[allow(unused)]
     #[inline]
     pub(crate) fn new_with_keys(key1: u128, key2: u128) -> Self {
         let pi: [u128; 2] = PI.convert();
@@ -160,6 +210,9 @@ impl Hasher for AHasher {
             self.hash_in(value.convert());
         } else {
             if data.len() > 32 {
+                if data.len() > 128 {
+                    return hash_batch_128b(&mut data, self);
+                }
                 if data.len() > 64 {
                     let tail = data.read_last_u128x4();
                     let mut current: [u128; 4] = [self.key; 4];

src/convert.rs

@@ -19,6 +19,7 @@ macro_rules! convert {
     };
 }
 
+convert!([u128; 8], [u8; 128]);
 convert!([u128; 4], [u64; 8]);
 convert!([u128; 4], [u32; 16]);
 convert!([u128; 4], [u16; 32]);
@@ -79,12 +80,14 @@ pub(crate) trait ReadFromSlice {
     fn read_u128(&self) -> (u128, &[u8]);
     fn read_u128x2(&self) -> ([u128; 2], &[u8]);
     fn read_u128x4(&self) -> ([u128; 4], &[u8]);
+    fn read_u128x8(&self) -> ([u128; 8], &[u8]);
     fn read_last_u16(&self) -> u16;
     fn read_last_u32(&self) -> u32;
     fn read_last_u64(&self) -> u64;
     fn read_last_u128(&self) -> u128;
     fn read_last_u128x2(&self) -> [u128; 2];
     fn read_last_u128x4(&self) -> [u128; 4];
+    fn read_last_u128x8(&self) -> [u128; 8];
 }
 
 impl ReadFromSlice for [u8] {
@@ -124,6 +127,12 @@ impl ReadFromSlice for [u8] {
         (as_array!(value, 64).convert(), rest)
     }
 
+    #[inline(always)]
+    fn read_u128x8(&self) -> ([u128; 8], &[u8]) {
+        let (value, rest) = self.split_at(128);
+        (as_array!(value, 128).convert(), rest)
+    }
+
     #[inline(always)]
     fn read_last_u16(&self) -> u16 {
         let (_, value) = self.split_at(self.len() - 2);
@@ -159,4 +168,9 @@ impl ReadFromSlice for [u8] {
         let (_, value) = self.split_at(self.len() - 64);
         as_array!(value, 64).convert()
     }
+
+    fn read_last_u128x8(&self) -> [u128; 8] {
+        let (_, value) = self.split_at(self.len() - 128);
+        as_array!(value, 128).convert()
+    }
 }

src/lib.rs

@@ -99,6 +99,7 @@ Note the import of [HashMapExt]. This is needed for the constructor.
 #![cfg_attr(all(not(test), not(feature = "std")), no_std)]
 #![cfg_attr(feature = "specialize", feature(min_specialization))]
 #![cfg_attr(feature = "nightly-arm-aes", feature(stdarch_arm_neon_intrinsics))]
+#![cfg_attr(feature = "vaes", feature(stdsimd))]
 
 #[macro_use]
 mod convert;