From 050f3089f145f203d1b4d83c4d3f9e9860e150e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Wed, 14 Oct 2026 20:39:40 +0300 Subject: [PATCH 1/7] sha2: tweak `riscv_zknh` backend --- sha2/src/sha256/riscv_zknh/compact.rs | 89 +++++++++++++-------------- 1 file changed, 43 insertions(+), 46 deletions(-) diff --git a/sha2/src/sha256/riscv_zknh/compact.rs b/sha2/src/sha256/riscv_zknh/compact.rs index caa19f91c..52695f3b6 100644 --- a/sha2/src/sha256/riscv_zknh/compact.rs +++ b/sha2/src/sha256/riscv_zknh/compact.rs @@ -4,52 +4,49 @@ use crate::consts::K32; #[inline] #[target_feature(enable = "zknh")] pub(super) fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) { - let mut s = *state; - - for r in 0..64 { - round(&mut s, &block, r); - if r < 48 { - schedule(&mut block, r) - } - } - - for i in 0..8 { - state[i] = state[i].wrapping_add(s[i]); + let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; + + for i in 0..64 { + let w = if i < 16 { + block[i] + } else { + let s0 = sha256sig0(block[(i - 15) % 16]); + let s1 = sha256sig1(block[(i - 2) % 16]); + let new_w = block[i % 16] + .wrapping_add(s1) + .wrapping_add(block[(i - 7) % 16]) + .wrapping_add(s0); + block[i % 16] = new_w; + new_w + }; + + let s1 = sha256sum1(e); // e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25); + let ch = (e & f) ^ ((!e) & g); + let t1 = s1 + .wrapping_add(ch) + .wrapping_add(K32[i]) + .wrapping_add(w) + .wrapping_add(h); + let s0 = sha256sum0(a); + let maj = (a & b) ^ (a & c) ^ (b & c); + let t2 = s0.wrapping_add(maj); + + h = g; + g = f; + f = e; + e = d.wrapping_add(t1); + d = c; + c = b; + b = a; + a = t1.wrapping_add(t2); } -} - -#[inline] -#[target_feature(enable = "zknh")] -fn schedule(block: &mut [u32; 16], r: usize) { - block[r % 16] = block[r % 16] - .wrapping_add(sha256sig1(block[(r + 14) % 16])) - .wrapping_add(block[(r + 9) % 16]) - .wrapping_add(sha256sig0(block[(r + 1) % 16])); -} - -#[inline] -#[target_feature(enable = "zknh")] -fn round(state: &mut [u32; 8], block: &[u32; 16], r: usize) { - let n = K32.len() - r; - let [a, b, c, d, e, f, g, h] = core::array::from_fn(|i| (n + i) % 8); - - state[h] = state[h] - .wrapping_add(sha256sum1(state[e])) - .wrapping_add(ch(state[e], state[f], state[g])) - .wrapping_add(K32[r]) - .wrapping_add(block[r % 16]); - state[d] = state[d].wrapping_add(state[h]); - state[h] = state[h] - .wrapping_add(sha256sum0(state[a])) - .wrapping_add(maj(state[a], state[b], state[c])) -} - -#[inline(always)] -fn ch(x: u32, y: u32, z: u32) -> u32 { - (x & y) ^ (!x & z) -} -#[inline(always)] -fn maj(x: u32, y: u32, z: u32) -> u32 { - (x & y) ^ (x & z) ^ (y & z) + state[0] = state[0].wrapping_add(a); + state[1] = state[1].wrapping_add(b); + state[2] = state[2].wrapping_add(c); + state[3] = state[3].wrapping_add(d); + state[4] = state[4].wrapping_add(e); + state[5] = state[5].wrapping_add(f); + state[6] = state[6].wrapping_add(g); + state[7] = state[7].wrapping_add(h); } From a854c373fd36bdfd589e010003c48b68dc3f62ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Wed, 14 Oct 2026 21:29:40 +0300 Subject: [PATCH 2/7] move block laoding into compress functions --- sha2/src/sha256/riscv_zknh.rs | 8 +--- sha2/src/sha256/riscv_zknh/compact.rs | 10 +++-- sha2/src/sha256/riscv_zknh/unroll.rs | 8 +++- sha2/src/sha512/riscv_zknh.rs | 8 +--- sha2/src/sha512/riscv_zknh/compact.rs | 58 +++++++++++++++++++++++++++ sha2/src/sha512/riscv_zknh/unroll.rs | 8 +++- 6 files changed, 81 insertions(+), 19 deletions(-) diff --git a/sha2/src/sha256/riscv_zknh.rs b/sha2/src/sha256/riscv_zknh.rs index bad22ecae..67161e311 100644 --- a/sha2/src/sha256/riscv_zknh.rs +++ b/sha2/src/sha256/riscv_zknh.rs @@ -18,13 +18,7 @@ cfg_if::cfg_if! { #[target_feature(enable = "zknh")] pub(super) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { - for block in blocks.iter().map(load_block) { + for block in blocks { compress_block(state, block); } } - -fn load_block(block: &[u8; 64]) -> [u32; 16] { - let (chunks, tail) = block.as_chunks(); - assert!(tail.is_empty()); - core::array::from_fn(|i| u32::from_be_bytes(chunks[i])) -} diff --git a/sha2/src/sha256/riscv_zknh/compact.rs b/sha2/src/sha256/riscv_zknh/compact.rs index 52695f3b6..9ab6c9bfe 100644 --- a/sha2/src/sha256/riscv_zknh/compact.rs +++ b/sha2/src/sha256/riscv_zknh/compact.rs @@ -3,8 +3,12 @@ use crate::consts::K32; #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) { - let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; +pub(super) fn compress_block(state: &mut [u32; 8], block: &[u8; 64]) { + let mut block: [u32; 16] = { + let (chunks, tail) = block.as_chunks(); + assert!(tail.is_empty()); + core::array::from_fn(|i| u32::from_be_bytes(chunks[i])) + }; for i in 0..64 { let w = if i < 16 { @@ -20,7 +24,7 @@ pub(super) fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) { new_w }; - let s1 = sha256sum1(e); // e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25); + let s1 = sha256sum1(e); let ch = (e & f) ^ ((!e) & g); let t1 = s1 .wrapping_add(ch) diff --git a/sha2/src/sha256/riscv_zknh/unroll.rs b/sha2/src/sha256/riscv_zknh/unroll.rs index dc3e9dc16..eb250c941 100644 --- a/sha2/src/sha256/riscv_zknh/unroll.rs +++ b/sha2/src/sha256/riscv_zknh/unroll.rs @@ -3,7 +3,13 @@ use crate::consts::K32; #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u32; 8], mut block: [u32; 16]) { +pub(super) fn compress_block(state: &mut [u32; 8], block: &[u8; 64]) { + let mut block: [u32; 16] = { + let (chunks, tail) = block.as_chunks(); + assert!(tail.is_empty()); + core::array::from_fn(|i| u32::from_be_bytes(chunks[i])) + }; + let s = &mut state.clone(); let b = &mut block; diff --git a/sha2/src/sha512/riscv_zknh.rs b/sha2/src/sha512/riscv_zknh.rs index ecb14fd5f..03f2ced76 100644 --- a/sha2/src/sha512/riscv_zknh.rs +++ b/sha2/src/sha512/riscv_zknh.rs @@ -13,17 +13,11 @@ cfg_if::cfg_if! { #[target_feature(enable = "zknh")] pub(super) fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) { - for block in blocks.iter().map(load_block) { + for block in blocks { compress_block(state, block); } } -fn load_block(block: &[u8; 128]) -> [u64; 16] { - let (chunks, tail) = block.as_chunks(); - assert!(tail.is_empty()); - core::array::from_fn(|i| u64::from_be_bytes(chunks[i])) -} - cfg_if::cfg_if! { if #[cfg(target_arch = "riscv64")] { use core::arch::riscv64::{sha512sig0, sha512sig1, sha512sum0, sha512sum1}; diff --git a/sha2/src/sha512/riscv_zknh/compact.rs b/sha2/src/sha512/riscv_zknh/compact.rs index 23b83ae1c..e7c902f37 100644 --- a/sha2/src/sha512/riscv_zknh/compact.rs +++ b/sha2/src/sha512/riscv_zknh/compact.rs @@ -1,6 +1,63 @@ use super::{sha512sig0, sha512sig1, sha512sum0, sha512sum1}; use crate::consts::K64; +#[inline] +#[target_feature(enable = "zknh")] +pub(super) fn compress_block(state: &mut [u64; 8], block: &[u8; 128]) { + let mut block: [u64; 16] = { + let (chunks, tail) = block.as_chunks(); + assert!(tail.is_empty()); + core::array::from_fn(|i| u64::from_be_bytes(chunks[i])) + }; + + let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; + + for i in 0..80 { + let w = if i < 16 { + block[i] + } else { + let s0 = sha512sig0(block[(i - 15) % 16]); + let s1 = sha512sig1(block[(i - 2) % 16]); + let new_w = block[i % 16] + .wrapping_add(s0) + .wrapping_add(block[(i - 7) % 16]) + .wrapping_add(s1); + block[i % 16] = new_w; + new_w + }; + + let s1 = sha512sum1(e); + let ch = (e & f) ^ ((!e) & g); + let t1 = s1 + .wrapping_add(ch) + .wrapping_add(K64[i]) + .wrapping_add(w) + .wrapping_add(h); + let s0 = sha512sum0(a); + let maj = (a & b) ^ (a & c) ^ (b & c); + let t2 = s0.wrapping_add(maj); + + h = g; + g = f; + f = e; + e = d.wrapping_add(t1); + d = c; + c = b; + b = a; + a = t1.wrapping_add(t2); + } + + state[0] = state[0].wrapping_add(a); + state[1] = state[1].wrapping_add(b); + state[2] = state[2].wrapping_add(c); + state[3] = state[3].wrapping_add(d); + state[4] = state[4].wrapping_add(e); + state[5] = state[5].wrapping_add(f); + state[6] = state[6].wrapping_add(g); + state[7] = state[7].wrapping_add(h); +} + +/* #[inline] #[target_feature(enable = "zknh")] pub(super) fn compress_block(state: &mut [u64; 8], mut block: [u64; 16]) { @@ -53,3 +110,4 @@ fn ch(x: u64, y: u64, z: u64) -> u64 { fn maj(x: u64, y: u64, z: u64) -> u64 { (x & y) ^ (x & z) ^ (y & z) } +*/ diff --git a/sha2/src/sha512/riscv_zknh/unroll.rs b/sha2/src/sha512/riscv_zknh/unroll.rs index 8aad2038a..e46ea4337 100644 --- a/sha2/src/sha512/riscv_zknh/unroll.rs +++ b/sha2/src/sha512/riscv_zknh/unroll.rs @@ -3,7 +3,13 @@ use crate::consts::K64; #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u64; 8], mut block: [u64; 16]) { +pub(super) fn compress_block(state: &mut [u64; 8], block: &[u8; 128]) { + let mut block: [u64; 16] = { + let (chunks, tail) = block.as_chunks(); + assert!(tail.is_empty()); + core::array::from_fn(|i| u64::from_be_bytes(chunks[i])) + }; + let s = &mut state.clone(); let b = &mut block; From ae49411f7f98324622f59a0bae6c39126d2f4313 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Wed, 14 Oct 2026 21:31:31 +0300 Subject: [PATCH 3/7] fix --- sha2/src/sha256/riscv_zknh/compact.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sha2/src/sha256/riscv_zknh/compact.rs b/sha2/src/sha256/riscv_zknh/compact.rs index 9ab6c9bfe..b0b263401 100644 --- a/sha2/src/sha256/riscv_zknh/compact.rs +++ b/sha2/src/sha256/riscv_zknh/compact.rs @@ -10,6 +10,8 @@ pub(super) fn compress_block(state: &mut [u32; 8], block: &[u8; 64]) { core::array::from_fn(|i| u32::from_be_bytes(chunks[i])) }; + let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; + for i in 0..64 { let w = if i < 16 { block[i] From 10c7ef56073d7d04781aa2a8be8220d78503bd3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Wed, 14 Oct 2026 21:40:14 +0300 Subject: [PATCH 4/7] remove odl code --- sha2/src/sha512/riscv_zknh/compact.rs | 55 --------------------------- 1 file changed, 55 deletions(-) diff --git a/sha2/src/sha512/riscv_zknh/compact.rs b/sha2/src/sha512/riscv_zknh/compact.rs index e7c902f37..436bdcbed 100644 --- a/sha2/src/sha512/riscv_zknh/compact.rs +++ b/sha2/src/sha512/riscv_zknh/compact.rs @@ -56,58 +56,3 @@ pub(super) fn compress_block(state: &mut [u64; 8], block: &[u8; 128]) { state[6] = state[6].wrapping_add(g); state[7] = state[7].wrapping_add(h); } - -/* -#[inline] -#[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u64; 8], mut block: [u64; 16]) { - let mut s = *state; - - for r in 0..80 { - round(&mut s, &block, r); - if r < 64 { - schedule(&mut block, r) - } - } - - for i in 0..8 { - state[i] = state[i].wrapping_add(s[i]); - } -} - -#[inline] -#[target_feature(enable = "zknh")] -fn schedule(block: &mut [u64; 16], r: usize) { - block[r % 16] = block[r % 16] - .wrapping_add(sha512sig1(block[(r + 14) % 16])) - .wrapping_add(block[(r + 9) % 16]) - .wrapping_add(sha512sig0(block[(r + 1) % 16])); -} - -#[inline] -#[target_feature(enable = "zknh")] -fn round(state: &mut [u64; 8], block: &[u64; 16], r: usize) { - let n = K64.len() - r; - let [a, b, c, d, e, f, g, h] = core::array::from_fn(|i| (n + i) % 8); - - state[h] = state[h] - .wrapping_add(sha512sum1(state[e])) - .wrapping_add(ch(state[e], state[f], state[g])) - .wrapping_add(K64[r]) - .wrapping_add(block[r % 16]); - state[d] = state[d].wrapping_add(state[h]); - state[h] = state[h] - .wrapping_add(sha512sum0(state[a])) - .wrapping_add(maj(state[a], state[b], state[c])) -} - -#[inline(always)] -fn ch(x: u64, y: u64, z: u64) -> u64 { - (x & y) ^ (!x & z) -} - -#[inline(always)] -fn maj(x: u64, y: u64, z: u64) -> u64 { - (x & y) ^ (x & z) ^ (y & z) -} -*/ From 4a819ff0bd8865f6795178908233a610657f9daa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Fri, 16 Oct 2026 16:16:03 +0300 Subject: [PATCH 5/7] tweak sha256 --- sha2/src/sha256/riscv_zknh.rs | 11 +---- sha2/src/sha256/riscv_zknh/compact.rs | 61 ++++++++++++++++++--------- sha2/src/sha256/riscv_zknh/unroll.rs | 9 +++- 3 files changed, 50 insertions(+), 31 deletions(-) diff --git a/sha2/src/sha256/riscv_zknh.rs b/sha2/src/sha256/riscv_zknh.rs index 67161e311..98677e862 100644 --- a/sha2/src/sha256/riscv_zknh.rs +++ b/sha2/src/sha256/riscv_zknh.rs @@ -9,16 +9,9 @@ use core::arch::riscv64::{sha256sig0, sha256sig1, sha256sum0, sha256sum1}; cfg_if::cfg_if! { if #[cfg(sha2_backend_riscv_zknh = "compact")] { mod compact; - use compact::compress_block; + pub(super) use compact::compress; } else { mod unroll; - use unroll::compress_block; - } -} - -#[target_feature(enable = "zknh")] -pub(super) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { - for block in blocks { - compress_block(state, block); + pub(super) use unroll::compress; } } diff --git a/sha2/src/sha256/riscv_zknh/compact.rs b/sha2/src/sha256/riscv_zknh/compact.rs index b0b263401..f0580c1dd 100644 --- a/sha2/src/sha256/riscv_zknh/compact.rs +++ b/sha2/src/sha256/riscv_zknh/compact.rs @@ -1,37 +1,56 @@ use super::{sha256sig0, sha256sig1, sha256sum0, sha256sum1}; use crate::consts::K32; +const ROUNDS: usize = K32.len(); + +#[target_feature(enable = "zknh")] +pub fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { + // Unforuntately, LLVM is unable to remove `memset` for + // `fn message_schedule(block: &[u8; 64]) -> [u32; ROUNDS] { ... }` + // See: https://rust.godbolt.org/z/s61sP3GKj + // So to prevent unnecessary `memset` call for each block we reuse the buffer. + let mut msg = [0u32; 64]; + + for block in blocks { + // Pre-computing the whole message schedule instead of doing it during each round + // with a ring buffer results in a smaller binary with less branching + // at the cost of using additional stack space + message_schedule(block, &mut msg); + rounds(state, &msg); + } +} + #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u32; 8], block: &[u8; 64]) { - let mut block: [u32; 16] = { - let (chunks, tail) = block.as_chunks(); - assert!(tail.is_empty()); - core::array::from_fn(|i| u32::from_be_bytes(chunks[i])) - }; +fn message_schedule(block: &[u8; 64], msg: &mut [u32; ROUNDS]) { + let (block_chunks, tail) = block.as_chunks(); + assert!(tail.is_empty()); - let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; + let n = block_chunks.len(); + for i in 0..n { + msg[i] = u32::from_be_bytes(block_chunks[i]); + } + + for i in n..ROUNDS { + msg[i] = msg[i - 16] + .wrapping_add(sha256sig1(msg[i - 2])) + .wrapping_add(msg[i - 7]) + .wrapping_add(sha256sig0(msg[i - 15])); + } +} - for i in 0..64 { - let w = if i < 16 { - block[i] - } else { - let s0 = sha256sig0(block[(i - 15) % 16]); - let s1 = sha256sig1(block[(i - 2) % 16]); - let new_w = block[i % 16] - .wrapping_add(s1) - .wrapping_add(block[(i - 7) % 16]) - .wrapping_add(s0); - block[i % 16] = new_w; - new_w - }; +#[inline] +#[target_feature(enable = "zknh")] +fn rounds(state: &mut [u32; 8], msg: &[u32; ROUNDS]) { + let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; + for i in 0..ROUNDS { let s1 = sha256sum1(e); let ch = (e & f) ^ ((!e) & g); let t1 = s1 .wrapping_add(ch) .wrapping_add(K32[i]) - .wrapping_add(w) + .wrapping_add(msg[i]) .wrapping_add(h); let s0 = sha256sum0(a); let maj = (a & b) ^ (a & c) ^ (b & c); diff --git a/sha2/src/sha256/riscv_zknh/unroll.rs b/sha2/src/sha256/riscv_zknh/unroll.rs index eb250c941..962f7a3da 100644 --- a/sha2/src/sha256/riscv_zknh/unroll.rs +++ b/sha2/src/sha256/riscv_zknh/unroll.rs @@ -1,9 +1,16 @@ use super::{sha256sig0, sha256sig1, sha256sum0, sha256sum1}; use crate::consts::K32; +#[target_feature(enable = "zknh")] +pub fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { + for block in blocks { + compress_block(state, block); + } +} + #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u32; 8], block: &[u8; 64]) { +fn compress_block(state: &mut [u32; 8], block: &[u8; 64]) { let mut block: [u32; 16] = { let (chunks, tail) = block.as_chunks(); assert!(tail.is_empty()); From 9fcc98d00283e46b0a16fc2c46aa225dc24f0b90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Fri, 16 Oct 2026 16:27:04 +0300 Subject: [PATCH 6/7] tweak sha512 --- sha2/src/sha256/riscv_zknh/compact.rs | 12 ++--- sha2/src/sha256/riscv_zknh/unroll.rs | 2 +- sha2/src/sha512/riscv_zknh.rs | 11 +---- sha2/src/sha512/riscv_zknh/compact.rs | 67 +++++++++++++++++---------- sha2/src/sha512/riscv_zknh/unroll.rs | 9 +++- 5 files changed, 60 insertions(+), 41 deletions(-) diff --git a/sha2/src/sha256/riscv_zknh/compact.rs b/sha2/src/sha256/riscv_zknh/compact.rs index f0580c1dd..e946c41e8 100644 --- a/sha2/src/sha256/riscv_zknh/compact.rs +++ b/sha2/src/sha256/riscv_zknh/compact.rs @@ -4,25 +4,25 @@ use crate::consts::K32; const ROUNDS: usize = K32.len(); #[target_feature(enable = "zknh")] -pub fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { +pub(crate) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { // Unforuntately, LLVM is unable to remove `memset` for // `fn message_schedule(block: &[u8; 64]) -> [u32; ROUNDS] { ... }` // See: https://rust.godbolt.org/z/s61sP3GKj // So to prevent unnecessary `memset` call for each block we reuse the buffer. - let mut msg = [0u32; 64]; + let mut msg = [0u32; ROUNDS]; for block in blocks { // Pre-computing the whole message schedule instead of doing it during each round // with a ring buffer results in a smaller binary with less branching // at the cost of using additional stack space - message_schedule(block, &mut msg); + schedule(block, &mut msg); rounds(state, &msg); } } #[inline] #[target_feature(enable = "zknh")] -fn message_schedule(block: &[u8; 64], msg: &mut [u32; ROUNDS]) { +fn schedule(block: &[u8; 64], msg: &mut [u32; ROUNDS]) { let (block_chunks, tail) = block.as_chunks(); assert!(tail.is_empty()); @@ -33,9 +33,9 @@ fn message_schedule(block: &[u8; 64], msg: &mut [u32; ROUNDS]) { for i in n..ROUNDS { msg[i] = msg[i - 16] - .wrapping_add(sha256sig1(msg[i - 2])) + .wrapping_add(sha256sig0(msg[i - 15])) .wrapping_add(msg[i - 7]) - .wrapping_add(sha256sig0(msg[i - 15])); + .wrapping_add(sha256sig1(msg[i - 2])); } } diff --git a/sha2/src/sha256/riscv_zknh/unroll.rs b/sha2/src/sha256/riscv_zknh/unroll.rs index 962f7a3da..e868c0514 100644 --- a/sha2/src/sha256/riscv_zknh/unroll.rs +++ b/sha2/src/sha256/riscv_zknh/unroll.rs @@ -2,7 +2,7 @@ use super::{sha256sig0, sha256sig1, sha256sum0, sha256sum1}; use crate::consts::K32; #[target_feature(enable = "zknh")] -pub fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { +pub(crate) fn compress(state: &mut [u32; 8], blocks: &[[u8; 64]]) { for block in blocks { compress_block(state, block); } diff --git a/sha2/src/sha512/riscv_zknh.rs b/sha2/src/sha512/riscv_zknh.rs index 03f2ced76..caa63ae4b 100644 --- a/sha2/src/sha512/riscv_zknh.rs +++ b/sha2/src/sha512/riscv_zknh.rs @@ -4,17 +4,10 @@ compile_error!("riscv-zknh backend can be used only on riscv32 and riscv64 targe cfg_if::cfg_if! { if #[cfg(sha2_backend_riscv_zknh = "compact")] { mod compact; - use compact::compress_block; + pub(super) use compact::compress; } else { mod unroll; - use unroll::compress_block; - } -} - -#[target_feature(enable = "zknh")] -pub(super) fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) { - for block in blocks { - compress_block(state, block); + pub(super) use unroll::compress; } } diff --git a/sha2/src/sha512/riscv_zknh/compact.rs b/sha2/src/sha512/riscv_zknh/compact.rs index 436bdcbed..86d436adc 100644 --- a/sha2/src/sha512/riscv_zknh/compact.rs +++ b/sha2/src/sha512/riscv_zknh/compact.rs @@ -1,39 +1,58 @@ use super::{sha512sig0, sha512sig1, sha512sum0, sha512sum1}; use crate::consts::K64; +const ROUNDS: usize = K64.len(); + +#[target_feature(enable = "zknh")] +pub(crate) fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) { + // Unforuntately, LLVM is unable to remove `memset` for + // `fn message_schedule(block: &[u8; 128]) -> [u64; ROUNDS] { ... }` + // See: https://rust.godbolt.org/z/ncxzYrMPh + // So to prevent unnecessary `memset` call for each block we reuse the buffer. + let mut msg = [0u64; ROUNDS]; + + for block in blocks { + // Pre-computing the whole message schedule instead of doing it during each round + // with a ring buffer results in a smaller binary with less branching + // at the cost of using additional stack space + schedule(block, &mut msg); + rounds(state, &msg); + } +} + #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u64; 8], block: &[u8; 128]) { - let mut block: [u64; 16] = { - let (chunks, tail) = block.as_chunks(); - assert!(tail.is_empty()); - core::array::from_fn(|i| u64::from_be_bytes(chunks[i])) - }; +fn schedule(block: &[u8; 128], msg: &mut [u64; ROUNDS]) { + let (block_chunks, tail) = block.as_chunks(); + assert!(tail.is_empty()); + let n = block_chunks.len(); + for i in 0..n { + msg[i] = u64::from_be_bytes(block_chunks[i]); + } + + for i in n..ROUNDS { + msg[i] = msg[i - 16] + .wrapping_add(sha512sig0(msg[i - 15])) + .wrapping_add(msg[i - 7]) + .wrapping_add(sha512sig1(msg[i - 2])); + } +} + +#[inline] +#[target_feature(enable = "zknh")] +fn rounds(state: &mut [u32; 8], msg: &[u32; ROUNDS]) { let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; - for i in 0..80 { - let w = if i < 16 { - block[i] - } else { - let s0 = sha512sig0(block[(i - 15) % 16]); - let s1 = sha512sig1(block[(i - 2) % 16]); - let new_w = block[i % 16] - .wrapping_add(s0) - .wrapping_add(block[(i - 7) % 16]) - .wrapping_add(s1); - block[i % 16] = new_w; - new_w - }; - - let s1 = sha512sum1(e); + for i in 0..ROUNDS { + let s1 = sha256sum1(e); let ch = (e & f) ^ ((!e) & g); let t1 = s1 .wrapping_add(ch) - .wrapping_add(K64[i]) - .wrapping_add(w) + .wrapping_add(K32[i]) + .wrapping_add(msg[i]) .wrapping_add(h); - let s0 = sha512sum0(a); + let s0 = sha256sum0(a); let maj = (a & b) ^ (a & c) ^ (b & c); let t2 = s0.wrapping_add(maj); diff --git a/sha2/src/sha512/riscv_zknh/unroll.rs b/sha2/src/sha512/riscv_zknh/unroll.rs index e46ea4337..56f40e52e 100644 --- a/sha2/src/sha512/riscv_zknh/unroll.rs +++ b/sha2/src/sha512/riscv_zknh/unroll.rs @@ -1,9 +1,16 @@ use super::{sha512sig0, sha512sig1, sha512sum0, sha512sum1}; use crate::consts::K64; +#[target_feature(enable = "zknh")] +pub(crate) fn compress(state: &mut [u64; 8], blocks: &[[u8; 128]]) { + for block in blocks { + compress_block(state, block); + } +} + #[inline] #[target_feature(enable = "zknh")] -pub(super) fn compress_block(state: &mut [u64; 8], block: &[u8; 128]) { +fn compress_block(state: &mut [u64; 8], block: &[u8; 128]) { let mut block: [u64; 16] = { let (chunks, tail) = block.as_chunks(); assert!(tail.is_empty()); From ad842411ce852499699fd1451eb93f6d090708fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= Date: Fri, 16 Oct 2026 16:30:51 +0300 Subject: [PATCH 7/7] fix --- sha2/src/sha512/riscv_zknh/compact.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sha2/src/sha512/riscv_zknh/compact.rs b/sha2/src/sha512/riscv_zknh/compact.rs index 86d436adc..a6b491f44 100644 --- a/sha2/src/sha512/riscv_zknh/compact.rs +++ b/sha2/src/sha512/riscv_zknh/compact.rs @@ -41,18 +41,18 @@ fn schedule(block: &[u8; 128], msg: &mut [u64; ROUNDS]) { #[inline] #[target_feature(enable = "zknh")] -fn rounds(state: &mut [u32; 8], msg: &[u32; ROUNDS]) { +fn rounds(state: &mut [u64; 8], msg: &[u64; ROUNDS]) { let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *state; for i in 0..ROUNDS { - let s1 = sha256sum1(e); + let s1 = sha512sum1(e); let ch = (e & f) ^ ((!e) & g); let t1 = s1 .wrapping_add(ch) - .wrapping_add(K32[i]) + .wrapping_add(K64[i]) .wrapping_add(msg[i]) .wrapping_add(h); - let s0 = sha256sum0(a); + let s0 = sha512sum0(a); let maj = (a & b) ^ (a & c) ^ (b & c); let t2 = s0.wrapping_add(maj);