Implement much faster sha256 and sha512.

sha2/Cargo.toml

@@ -10,5 +10,15 @@ keywords = ["crypto", "sha2", "asm"]
 categories = ["cryptography", "no-std"]
 edition = "2018"
 
+[dependencies]
+cpufeatures = "0.2.1"
+
 [build-dependencies]
 cc = "1.0"
+
+[dev-dependencies]
+criterion = {version= "0.3.5",features=["html_reports"] }
+
+[[bench]]
+name = "sha2"
+harness = false

sha2/benches/sha2.rs

@@ -0,0 +1,17 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+pub fn criterion_benchmark(c: &mut Criterion) {
+    let mut state = Default::default();
+    let data = [[0u8; 64]];
+    c.bench_function("sha256", |b| {
+        b.iter(|| sha2_asm::compress256(black_box(&mut state), black_box(&data)))
+    });
+    let mut state = Default::default();
+    let data = [[0u8; 128]];
+    c.bench_function("sha512", |b| {
+        b.iter(|| sha2_asm::compress512(black_box(&mut state), black_box(&data)))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

sha2/build.rs

@@ -6,24 +6,35 @@ fn main() {
 
     let mut build256 = cc::Build::new();
     let (sha256_path, sha512_path) = if target_arch == "x86" {
-        ("src/sha256_x86.S", "src/sha512_x86.S")
+        (["src/sha256_x86.S"].to_vec(), ["src/sha512_x86.S"].to_vec())
     } else if target_arch == "x86_64" {
-        ("src/sha256_x64.S", "src/sha512_x64.S")
+        let sha512 = ["src/sha512_x64_avx2.S", "src/sha512_x64.S"].to_vec();
+        // Prioritizing sha-ni, cause it's fastest
+        let sha256 = [
+            "src/sha256_x64_ni.S",
+            "src/sha256_x64_avx2.S",
+            "src/sha256_x64.S",
+        ]
+        .to_vec();
+        (sha256, sha512)
     } else if target_arch == "aarch64" && target_vendor == "apple" {
         build256.flag("-march=armv8-a+crypto");
-        ("src/sha256_aarch64_apple.S", "")
+        (["src/sha256_aarch64_apple.S"].to_vec(), [""].to_vec())
     } else if target_arch == "aarch64" {
         build256.flag("-march=armv8-a+crypto");
-        ("src/sha256_aarch64.S", "")
+        (["src/sha256_aarch64.S"].to_vec(), [""].to_vec())
     } else {
         panic!("Unsupported target architecture");
     };
 
     if target_arch != "aarch64" {
         cc::Build::new()
             .flag("-c")
-            .file(sha512_path)
+            .files(sha512_path)
             .compile("libsha512.a");
     }
-    build256.flag("-c").file(sha256_path).compile("libsha256.a");
+    build256
+        .flag("-c")
+        .files(sha256_path)
+        .compile("libsha256.a");
 }

sha2/src/lib.rs

@@ -13,23 +13,47 @@
 #[cfg(not(any(target_arch = "x86_64", target_arch = "x86", target_arch = "aarch64")))]
 compile_error!("crate can only be used on x86, x86-64 and aarch64 architectures");
 
+cpufeatures::new!(cpuid_avx2, "avx2");
+cpufeatures::new!(cpuid_sha, "sha");
+
 #[link(name = "sha256", kind = "static")]
+#[allow(dead_code)]
 extern "C" {
     fn sha256_compress(state: &mut [u32; 8], block: &[u8; 64]);
+    // setting usize is safe, cause we are on 64 bit platform
+    fn sha256_transform_rorx(state: &mut [u32; 8], block: *const [u8; 64], num_blocks: usize);
+    fn sha256_ni_transform(digest: &mut [u32; 8], data: *const [u8; 64], nblk: usize);
 }
 
 /// Safe wrapper around assembly implementation of SHA256 compression function
+///
 #[inline]
 pub fn compress256(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
-    for block in blocks {
-        unsafe { sha256_compress(state, block) }
+    let token_sha = cpuid_sha::init();
+    if token_sha.get() {
+        if !blocks.is_empty() {
+            unsafe { sha256_ni_transform(state, blocks.as_ptr(), blocks.len()) }
+        }
+        return;
+    }
+    let token: cpuid_avx2::InitToken = cpuid_avx2::init();
+
+    if token.get() {
+        if !blocks.is_empty() {
+            unsafe { sha256_transform_rorx(state, blocks.as_ptr(), blocks.len()) }
+        }
+    } else {
+        for block in blocks {
+            unsafe { sha256_compress(state, block) }
+        }
     }
 }
 
 #[cfg(not(target_arch = "aarch64"))]
 #[link(name = "sha512", kind = "static")]
 extern "C" {
     fn sha512_compress(state: &mut [u64; 8], block: &[u8; 128]);
+    fn sha512_transform_rorx(state: &mut [u64; 8], block: *const [u8; 128], num_blocks: usize);
 }
 
 /// Safe wrapper around assembly implementation of SHA512 compression function
@@ -38,7 +62,14 @@ extern "C" {
 #[cfg(not(target_arch = "aarch64"))]
 #[inline]
 pub fn compress512(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
-    for block in blocks {
-        unsafe { sha512_compress(state, block) }
+    let token: cpuid_avx2::InitToken = cpuid_avx2::init();
+    if token.get() {
+        if !blocks.is_empty() {
+            unsafe { sha512_transform_rorx(state, blocks.as_ptr(), blocks.len()) }
+        }
+    } else {
+        for block in blocks {
+            unsafe { sha512_compress(state, block) }
+        }
     }
 }