sha2: remove RISC-V hack for unaligned loads

sha2/README.md

@@ -63,8 +63,7 @@ SHA-256 and SHA-512 backends:
 - `soft`: portable software implementation
 - `loongarch64-asm`: `asm!`-based implementation for LoongArch64 targets
 - `riscv-zknh`: uses the RISC-V `Zknh` scalar crypto extension. Experimental,
-  requires Nightly compiler and to enable `Zknh` and `Zbkb` (or `Zbb`)
-  target features at compile time.
+  see the [section below](#about-the-risc-v-zknh-backend) for more information.
 - `wasm32-simd128`: uses the WASM `simd128` extension.
 
 SHA-256 only backends:
@@ -103,6 +102,16 @@ performance at the cost of a bigger resulting binary. You can disable unrolling
 by using `sha2_backend_soft = "compact"` and `sha2_backend_riscv_zknh = "compact"` configuration
 flags respectively.
 
+### About the RISC-V Zknh backend
+
+This is an experimental backend which requires a Nightly compiler and to enable the `Zknh`
+target feature at compile time. For much better code generation, it's recommended
+to enable the `Zbkb` (or `Zbb`) and `unaligned-scalar-mem` target features
+(see [this LLVM issue][Zicclsm] for more information). For example,
+you can do it by using `RUSTFLAGS="-C target-feature=+zicclsm,+zbkb,+unaligned-scalar-mem"`.
+
+[Zicclsm]: https://github.com/llvm/llvm-project/issues/110454
+
 ## License
 
 The crate is licensed under either of:

sha2/src/sha256.rs

@@ -5,11 +5,8 @@ cfg_if::cfg_if! {
     } else if #[cfg(any(sha2_backend = "riscv-zknh", sha2_256_backend = "riscv-zknh"))] {
         mod riscv_zknh;
 
-        #[cfg(not(all(
-            target_feature = "zknh",
-            any(target_feature = "zbb", target_feature = "zbkb")
-        )))]
-        compile_error!("riscv-zknh backend requires zknh and zbkb (or zbb) target features");
+        #[cfg(not(target_feature = "zknh"))]
+        compile_error!("riscv-zknh backend requires `zknh` target feature");
 
         fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
             // SAFETY: we checked above that the required target features are enabled
@@ -22,7 +19,7 @@ cfg_if::cfg_if! {
             target_feature = "sha",
             target_feature = "sse4.1",
         )))]
-        compile_error!("x86-sha backend requires sha and sse4.1 target features");
+        compile_error!("x86-sha backend requires `sha` and `sse4.1` target features");
 
         fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
             // SAFETY: we checked above that the required target features are enabled
@@ -32,7 +29,7 @@ cfg_if::cfg_if! {
         mod aarch64_sha2;
 
         #[cfg(not(target_feature = "sha2"))]
-        compile_error!("aarch64-sha2 backend requires sha2 target feature");
+        compile_error!("aarch64-sha2 backend requires `sha2` target feature");
 
         fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
             // SAFETY: we checked above that the required target features are enabled

sha2/src/sha256/riscv_zknh.rs

@@ -1,8 +1,6 @@
 #[cfg(not(any(target_arch = "riscv32", target_arch = "riscv64")))]
 compile_error!("riscv-zknh backend can be used only on riscv32 and riscv64 target arches");
 
-mod utils;
-
 #[cfg(target_arch = "riscv32")]
 use core::arch::riscv32::{sha256sig0, sha256sig1, sha256sum0, sha256sum1};
 #[cfg(target_arch = "riscv64")]
@@ -11,9 +9,20 @@ use core::arch::riscv64::{sha256sig0, sha256sig1, sha256sum0, sha256sum1};
 cfg_if::cfg_if! {
     if #[cfg(sha2_backend_riscv_zknh = "compact")] {
         mod compact;
-        pub(super) use compact::compress;
+        use compact::compress_block;
     } else {
         mod unroll;
-        pub(super) use unroll::compress;
+        use unroll::compress_block;
+    }
+}
+
+#[target_feature(enable = "zknh")]
+pub(super) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
+    for block in blocks {
+        let block: [u32; 16] = core::array::from_fn(|i| {
+            let chunk = block[4 * i..][..4].try_into().unwrap();
+            u32::from_be_bytes(chunk)
+        });
+        compress_block(state, block);
     }
 }

sha2/src/sha256/riscv_zknh/compact.rs

@@ -2,14 +2,7 @@ use super::{sha256sig0, sha256sig1, sha256sum0, sha256sum1};
 use crate::consts::K32;
 
 #[target_feature(enable = "zknh")]
-pub(in super::super) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
-    for block in blocks.iter().map(super::utils::load_block) {
-        compress_block(state, block);
-    }
-}
-
-#[target_feature(enable = "zknh")]
-fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) {
+pub(super) fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) {
     let mut s = *state;
 
     for r in 0..64 {

sha2/src/sha256/riscv_zknh/unroll.rs

@@ -2,14 +2,7 @@ use super::{sha256sig0, sha256sig1, sha256sum0, sha256sum1};
 use crate::consts::K32;
 
 #[target_feature(enable = "zknh")]
-pub(in super::super) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
-    for block in blocks.iter().map(super::utils::load_block) {
-        compress_block(state, block);
-    }
-}
-
-#[target_feature(enable = "zknh")]
-fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) {
+pub(super) fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) {
     let s = &mut state.clone();
     let b = &mut block;
 
@@ -82,7 +75,7 @@ fn round<const R: usize>(state: &mut [u32; 8], block: &[u32; 16], k: &[u32]) {
     state[h] = state[h]
         .wrapping_add(sha256sum1(state[e]))
         .wrapping_add(ch(state[e], state[f], state[g]))
-        .wrapping_add(super::utils::opaque_load::<R>(k))
+        .wrapping_add(opaque_load::<R>(k))
         .wrapping_add(block[R]);
     state[d] = state[d].wrapping_add(state[h]);
     state[h] = state[h]
@@ -99,3 +92,33 @@ fn ch(x: u32, y: u32, z: u32) -> u32 {
 fn maj(x: u32, y: u32, z: u32) -> u32 {
     (x & y) ^ (x & z) ^ (y & z)
 }
+
+/// This function returns `k[R]`, but prevents the compiler from inlining the indexed value
+fn opaque_load<const R: usize>(k: &[u32]) -> u32 {
+    assert!(R < k.len());
+    let dst;
+
+    #[cfg(target_arch = "riscv64")]
+    unsafe {
+        core::arch::asm!(
+            "lwu {dst}, 4*{R}({k})",
+            R = const R,
+            k = in(reg) k.as_ptr(),
+            dst = out(reg) dst,
+            options(pure, readonly, nostack, preserves_flags),
+        );
+    }
+
+    #[cfg(target_arch = "riscv32")]
+    unsafe {
+        core::arch::asm!(
+            "lw {dst}, 4*{R}({k})",
+            R = const R,
+            k = in(reg) k.as_ptr(),
+            dst = out(reg) dst,
+            options(pure, readonly, nostack, preserves_flags),
+        );
+    }
+
+    dst
+}

sha2/src/sha512.rs

@@ -5,11 +5,8 @@ cfg_if::cfg_if! {
     } else if #[cfg(any(sha2_backend = "riscv-zknh", sha2_256_backend = "riscv-zknh"))] {
         mod riscv_zknh;
 
-        #[cfg(not(all(
-            target_feature = "zknh",
-            any(target_feature = "zbb", target_feature = "zbkb")
-        )))]
-        compile_error!("riscv-zknh backend requires zknh and zbkb (or zbb) target features");
+        #[cfg(not(target_feature = "zknh"))]
+        compile_error!("riscv-zknh backend requires `zknh` target features");
 
         fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
             // SAFETY: we checked above that the required target features are enabled
@@ -19,7 +16,7 @@ cfg_if::cfg_if! {
         mod x86_avx2;
 
         #[cfg(not(target_feature = "avx2"))]
-        compile_error!("x86-avx2 backend requires avx2 target feature");
+        compile_error!("x86-avx2 backend requires `avx2` target feature");
 
         fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
             // SAFETY: we checked above that the required target features are enabled
@@ -29,7 +26,7 @@ cfg_if::cfg_if! {
         mod aarch64_sha3;
 
         #[cfg(not(target_feature = "sha3"))]
-        compile_error!("aarch64-sha3 backend requires sha3 target feature");
+        compile_error!("aarch64-sha3 backend requires `sha3` target feature");
 
         fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
             // SAFETY: we checked above that the required target features are enabled

sha2/src/sha512/riscv_zknh.rs

@@ -1,52 +1,59 @@
 #[cfg(not(any(target_arch = "riscv32", target_arch = "riscv64")))]
 compile_error!("riscv-zknh backend can be used only on riscv32 and riscv64 target arches");
 
-mod utils;
-
 cfg_if::cfg_if! {
     if #[cfg(sha2_backend_riscv_zknh = "compact")] {
         mod compact;
-        pub(super) use compact::compress;
+        use compact::compress_block;
     } else {
         mod unroll;
-        pub(super) use unroll::compress;
+        use unroll::compress_block;
     }
 }
 
-#[cfg(target_arch = "riscv64")]
-use core::arch::riscv64::{sha512sig0, sha512sig1, sha512sum0, sha512sum1};
-
-#[cfg(target_arch = "riscv32")]
-use core::arch::riscv32::*;
-
-#[cfg(target_arch = "riscv32")]
-#[target_feature(enable = "zknh")]
-fn sha512sum0(x: u64) -> u64 {
-    let a = sha512sum0r((x >> 32) as u32, x as u32);
-    let b = sha512sum0r(x as u32, (x >> 32) as u32);
-    ((a as u64) << 32) | (b as u64)
-}
-
-#[cfg(target_arch = "riscv32")]
-#[target_feature(enable = "zknh")]
-fn sha512sum1(x: u64) -> u64 {
-    let a = sha512sum1r((x >> 32) as u32, x as u32);
-    let b = sha512sum1r(x as u32, (x >> 32) as u32);
-    ((a as u64) << 32) | (b as u64)
-}
-
-#[cfg(target_arch = "riscv32")]
 #[target_feature(enable = "zknh")]
-fn sha512sig0(x: u64) -> u64 {
-    let a = sha512sig0h((x >> 32) as u32, x as u32);
-    let b = sha512sig0l(x as u32, (x >> 32) as u32);
-    ((a as u64) << 32) | (b as u64)
+pub(super) fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
+    for block in blocks {
+        let block: [u64; 16] = core::array::from_fn(|i| {
+            let chunk = block[8 * i..][..8].try_into().unwrap();
+            u64::from_be_bytes(chunk)
+        });
+        compress_block(state, block);
+    }
 }
 
-#[cfg(target_arch = "riscv32")]
-#[target_feature(enable = "zknh")]
-fn sha512sig1(x: u64) -> u64 {
-    let a = sha512sig1h((x >> 32) as u32, x as u32);
-    let b = sha512sig1l(x as u32, (x >> 32) as u32);
-    ((a as u64) << 32) | (b as u64)
+cfg_if::cfg_if! {
+    if #[cfg(target_arch = "riscv64")] {
+        use core::arch::riscv64::{sha512sig0, sha512sig1, sha512sum0, sha512sum1};
+    } else {
+        use core::arch::riscv32::*;
+
+        #[target_feature(enable = "zknh")]
+        fn sha512sum0(x: u64) -> u64 {
+            let a = sha512sum0r((x >> 32) as u32, x as u32);
+            let b = sha512sum0r(x as u32, (x >> 32) as u32);
+            ((a as u64) << 32) | (b as u64)
+        }
+
+        #[target_feature(enable = "zknh")]
+        fn sha512sum1(x: u64) -> u64 {
+            let a = sha512sum1r((x >> 32) as u32, x as u32);
+            let b = sha512sum1r(x as u32, (x >> 32) as u32);
+            ((a as u64) << 32) | (b as u64)
+        }
+
+        #[target_feature(enable = "zknh")]
+        fn sha512sig0(x: u64) -> u64 {
+            let a = sha512sig0h((x >> 32) as u32, x as u32);
+            let b = sha512sig0l(x as u32, (x >> 32) as u32);
+            ((a as u64) << 32) | (b as u64)
+        }
+
+        #[target_feature(enable = "zknh")]
+        fn sha512sig1(x: u64) -> u64 {
+            let a = sha512sig1h((x >> 32) as u32, x as u32);
+            let b = sha512sig1l(x as u32, (x >> 32) as u32);
+            ((a as u64) << 32) | (b as u64)
+        }
+    }
 }