Skip to main content

tape_sha256/
lib.rs

1//! SHA-256 in bulk: many independent messages at once, or iterated chains
2//!
3//! One hash cannot be vectorised, since each of the 64 rounds depends on the
4//! last. Independent messages can, though: run them in lockstep, one per SIMD
5//! lane, and every round becomes an elementwise vector operation.
6//!
7//! Worth reaching for when you have many messages and need none of them early.
8//! The canonical case is a Merkle tree, where the whole bottom level can be
9//! hashed in one pass.
10//!
11//! ```
12//! use tape_sha256::hash_many;
13//!
14//! let msgs: Vec<&[u8]> = vec![b"one", b"two", b"three"];
15//! let mut out = vec![[0u8; 32]; msgs.len()];
16//! hash_many(&msgs, &mut out);
17//! ```
18//!
19//! When every leaf shares a domain-separation prefix, hash_many_prefixed
20//! avoids materialising `prefix || body` per message:
21//!
22//! ```
23//! use tape_sha256::hash_many_prefixed;
24//!
25//! const LEAF: &[u8] = b"\x00SOLANA_MERKLE_SHREDS_LEAF";
26//! let leaves: Vec<&[u8]> = vec![&[1u8; 64], &[2u8; 64]];
27//! let mut out = vec![[0u8; 32]; leaves.len()];
28//! hash_many_prefixed(LEAF, &leaves, &mut out);
29//! ```
30//!
31//! # Correctness
32//!
33//! Output is bit-identical to any conforming SHA-256. Every backend is gated
34//! against the independent `sha2` crate over lengths covering all block and
35//! padding edge cases, and each SIMD backend is cross-checked against the
36//! portable one.
37//!
38//! # Backend selection
39//!
40//! By default the best kernel for the running CPU is picked and falls back
41//! gracefully. The `scalar`, `avx2`, `avx512`, and `neon` features pin one
42//! kernel at build time instead, skipping detection; a pinned kernel the CPU
43//! lacks will fault, so pin only what the whole fleet supports.
44//!
45//! # Using both threads of a core
46//!
47//! These entry points are stateless, so two threads may call them concurrently
48//! on disjoint halves of a batch. Doing that on the two SMT siblings of one
49//! physical core measured 1.5x the single-thread rate on Zen 5 (9.83 vs 14.84
50//! us for 64 Merkle leaves), because a single thread at ~2 instructions per
51//! cycle is limited by its own dependency chains rather than by issue width,
52//! and the sibling fills the gaps.
53//!
54//! This crate never spawns or pins a thread. Thread topology is the caller's
55//! policy, not a hashing library's. Split the batch, hand each half to a thread
56//! you have pinned, and join:
57//!
58//! ```no_run
59//! use tape_sha256::hash_many;
60//!
61//! # let msgs: Vec<&[u8]> = vec![];
62//! # let mut out = vec![[0u8; 32]; 0];
63//! let (m0, m1) = msgs.split_at(msgs.len() / 2);
64//! let (o0, o1) = out.split_at_mut(msgs.len() / 2);
65//! std::thread::scope(|s| {
66//!     s.spawn(|| hash_many(m0, o0));
67//!     s.spawn(|| hash_many(m1, o1));
68//! });
69//! ```
70//!
71//! How much this is worth depends heavily on the microarchitecture.
72//!
73//! Three further caveats decide whether this is worth anything:
74//!
75//! The two threads must be siblings of the *same* physical core. Pinned to
76//! different cores this is ordinary multicore parallelism, which spends a core
77//! to get it; the 1.5x claim is specifically about using a sibling that would
78//! otherwise idle. Sibling pairs are listed in
79//! `/sys/devices/system/cpu/cpu<N>/topology/thread_siblings_list`. Getting the
80//! pairing wrong degrades silently, so it is worth asserting.
81//!
82//! The sibling has to actually be idle. Under a fully loaded process there is
83//! no spare thread to claim and the gain goes away. Use an equal split of the
84//! same kernel: an uneven or mixed-kernel split leaves a straggler and measured
85//! worse than a single thread.
86//!
87//! # Hash chains
88//!
89//! Everything above is one pass over messages that already exist. A chain
90//! instead feeds each digest back in as the next message, which is what
91//! Solana's proof of history is.
92//!
93//! [`hash_chain`] is the exception to the whole premise of this crate: one
94//! chain has no independent work to lane up, so it ignores the multi-buffer
95//! kernels and drives the CPU's SHA-256 unit at its latency floor.
96//!
97//! [`hash_chains`] puts the premise back. Independent chains, one per
98//! verified entry, say run in lockstep exactly like independent messages
99//! do, which trades that latency floor for a throughput one. See both.
100
101#[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
102mod avx2;
103#[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
104mod avx512;
105#[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
106mod avx512x2;
107mod batch;
108mod chain;
109mod core;
110mod lanes;
111#[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
112mod neon;
113#[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
114mod neon_sha2;
115#[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
116mod shani;
117#[cfg(all(
118    target_arch = "wasm32",
119    target_feature = "simd128",
120    not(feature = "scalar")
121))]
122mod simd128;
123
124pub use batch::Message;
125
126use {
127    batch::{drive, drive_pairs, drive_slices},
128    lanes::Scalar,
129};
130
131/// Number of messages the active backend hashes per pass
132///
133/// Batches at least this large amortise the transpose fully; smaller ones
134/// still work but leave lanes idle.
135pub fn lane_width() -> usize {
136    dispatch::lane_width()
137}
138
139/// Name of the active backend, for logging and benchmarks
140pub fn backend() -> &'static str {
141    dispatch::backend()
142}
143
144/// Hashes every message in `msgs`, writing one digest per message to `out`
145///
146/// # Panics
147///
148/// Panics if `msgs.len() != out.len()`.
149pub fn hash_many(msgs: &[&[u8]], out: &mut [[u8; 32]]) {
150    dispatch::hash_slices(&[], msgs, out);
151}
152
153/// Hashes `prefix || body` for each body in `bodies`
154///
155/// Equivalent to concatenating and calling hash_many, but without building the
156/// concatenation.
157///
158/// # Panics
159///
160/// Panics if `bodies.len() != out.len()`.
161pub fn hash_many_prefixed(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
162    dispatch::hash_slices(prefix, bodies, out);
163}
164
165/// Hashes `prefix || left || right` for each pair, without joining them first
166///
167/// For Merkle interior nodes, where a parent hashes a domain prefix and two
168/// child hashes held in separate buffers. Saves the caller a staging buffer
169/// and the join.
170///
171/// # Panics
172///
173/// Panics unless `left.len() == right.len() == out.len()`.
174pub fn hash_pairs(prefix: &[u8], left: &[&[u8]], right: &[&[u8]], out: &mut [[u8; 32]]) {
175    dispatch::hash_pairs(prefix, left, right, out);
176}
177
178/// Hashes pre-built messages, for callers that want per-message prefixes
179///
180/// # Panics
181///
182/// Panics if `msgs.len() != out.len()`.
183pub fn hash_messages(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
184    dispatch::hash(msgs, out);
185}
186
187/// Hashes `seed`, then that digest, and so on `n` times
188///
189/// The message is always 32 bytes, which makes the padded block constant above
190/// word 8 and lets the state feed back without the byte swap a general
191/// `sha256(digest)` would pay twice per link.
192///
193/// Returns `seed` unchanged when `n` is 0.
194///
195/// ```
196/// use tape_sha256::hash_chain;
197///
198/// let seed = [0u8; 32];
199/// assert_eq!(hash_chain(&seed, 3), hash_chain(&hash_chain(&seed, 1), 2));
200/// ```
201pub fn hash_chain(seed: &[u8; 32], n: u64) -> [u8; 32] {
202    chain::hash_chain(seed, n)
203}
204
205/// Hashes many independent chains at once, one per lane
206///
207/// `out[i]` is `seeds[i]` hashed back into itself `lens[i]` times, the same
208/// answer `hash_chain` gives, but the chains run in lockstep so the SHA-256
209/// unit is throughput-bound instead of latency-bound.
210///
211/// A single chain cannot go faster than one block compression's latency, and
212/// `hash_chain` already sits there. Independent chains are a different
213/// problem, and Solana replay has them: every `Entry` publishes its ending
214/// hash, so entry `i`'s segment starts from entry `i - 1`'s known hash and all
215/// segments can be verified at once. Ragged lengths are scheduled, not paid
216/// for: chains run longest-first and a lane that finishes picks up the next
217/// chain, so a batch costs about `lens.iter().sum() / width` links, however
218/// the lengths fall.
219///
220/// Uses `backend`'s kernel, not `chain_backend`'s.
221///
222/// # Panics
223///
224/// Panics unless `seeds.len() == lens.len() == out.len()`.
225///
226/// ```
227/// use tape_sha256::{hash_chain, hash_chains};
228///
229/// let seeds = [[1u8; 32], [2u8; 32]];
230/// let lens = [7u64, 9];
231/// let mut out = [[0u8; 32]; 2];
232/// hash_chains(&seeds, &lens, &mut out);
233/// assert_eq!(out[0], hash_chain(&seeds[0], 7));
234/// assert_eq!(out[1], hash_chain(&seeds[1], 9));
235/// ```
236pub fn hash_chains(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
237    assert_eq!(seeds.len(), lens.len());
238    assert_eq!(seeds.len(), out.len());
239    chain::hash_chains(seeds, lens, out);
240}
241
242/// Name of the kernel `hash_chain` uses
243///
244/// Not the same choice as `backend`: the chain cares only whether a SHA-256
245/// unit exists, never how wide the vector unit is.
246pub fn chain_backend() -> &'static str {
247    chain::backend()
248}
249
250/// Backends exposed for differential testing and benchmarking
251///
252/// Not stable API; callers should use hash_many and let dispatch pick. Each
253/// dispatchable backend has two entry points: one taking pre-built messages,
254/// and a `_slices` form that builds them on the stack so the wrappers need no
255/// allocation. Both `serial` and `portable8` are dispatchable — which one the
256/// no-SIMD fallback picks is per-architecture, see `dispatch::PORTABLE`.
257#[doc(hidden)]
258pub mod backends {
259    use super::*;
260
261    // Concrete scalar group functions; see GroupFn for why kernels are
262    // instantiated once here and called through pointers.
263    pub(crate) fn scalar1(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
264        batch::hash_lanes::<Scalar<1>, 1>(msgs, out)
265    }
266    pub(crate) fn scalar8(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
267        batch::hash_lanes::<Scalar<8>, 8>(msgs, out)
268    }
269    /// Portable 8-lane kernel
270    ///
271    /// The reference every SIMD backend is checked against, and the no-SIMD
272    /// fallback on register-rich targets.
273    pub fn portable8(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
274        unsafe { drive(8, scalar8, msgs, out) }
275    }
276
277    pub fn portable8_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
278        unsafe { drive_slices(8, scalar8, prefix, bodies, out) }
279    }
280
281    /// Single-lane kernel: ordinary serial SHA-256 through the same core
282    ///
283    /// Shows the lane machinery is not what makes output correct, and is the
284    /// baseline speedups are quoted against.
285    pub fn serial(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
286        unsafe { drive(1, scalar1, msgs, out) }
287    }
288
289    pub fn serial_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
290        unsafe { drive_slices(1, scalar1, prefix, bodies, out) }
291    }
292
293    /// 4-lane AArch64 NEON kernel
294    ///
295    /// NEON is baseline on AArch64, so this is always safe to call there.
296    #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
297    pub fn neon4(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
298        unsafe { drive(4, crate::neon::group, msgs, out) }
299    }
300
301    #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
302    pub fn neon4_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
303        unsafe { drive_slices(4, crate::neon::group, prefix, bodies, out) }
304    }
305
306    /// 4-lane wasm simd128 kernel
307    ///
308    /// Exists only in builds with `-C target-feature=+simd128`; an engine
309    /// that lacks simd128 rejects the module at instantiation, so a call
310    /// that runs at all is safe.
311    #[cfg(all(
312        target_arch = "wasm32",
313        target_feature = "simd128",
314        not(feature = "scalar")
315    ))]
316    pub fn simd128_4(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
317        unsafe { drive(4, crate::simd128::group, msgs, out) }
318    }
319
320    #[cfg(all(
321        target_arch = "wasm32",
322        target_feature = "simd128",
323        not(feature = "scalar")
324    ))]
325    pub fn simd128_4_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
326        unsafe { drive_slices(4, crate::simd128::group, prefix, bodies, out) }
327    }
328
329    /// Two-wave 8-lane wasm simd128 kernel, for benchmarking against the
330    /// single wave; never dispatched to
331    #[cfg(all(
332        target_arch = "wasm32",
333        target_feature = "simd128",
334        not(feature = "scalar")
335    ))]
336    pub fn simd128_8_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
337        unsafe { drive_slices(8, crate::simd128::group8, prefix, bodies, out) }
338    }
339
340    /// Four-wave 16-lane wasm simd128 kernel; never dispatched to
341    #[cfg(all(
342        target_arch = "wasm32",
343        target_feature = "simd128",
344        not(feature = "scalar")
345    ))]
346    pub fn simd128_16_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
347        unsafe { drive_slices(16, crate::simd128::group16, prefix, bodies, out) }
348    }
349
350    /// Multi-stream kernel on the ARMv8 SHA-256 crypto extension
351    ///
352    /// # Safety
353    ///
354    /// The running CPU must support the `sha2` extension.
355    #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
356    pub unsafe fn neon_sha2x4(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
357        drive(
358            crate::neon_sha2::STREAMS,
359            crate::neon_sha2::group,
360            msgs,
361            out,
362        )
363    }
364
365    /// # Safety
366    ///
367    /// The running CPU must support the `sha2` extension.
368    #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
369    pub unsafe fn neon_sha2x4_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
370        drive_slices(
371            crate::neon_sha2::STREAMS,
372            crate::neon_sha2::group,
373            prefix,
374            bodies,
375            out,
376        )
377    }
378
379    /// Multi-stream kernel on the x86 SHA-NI extension
380    ///
381    /// # Safety
382    ///
383    /// The running CPU must support SHA-NI, SSSE3, and SSE4.1.
384    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
385    pub unsafe fn shani_x4(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
386        drive(crate::shani::STREAMS, crate::shani::group, msgs, out)
387    }
388
389    /// # Safety
390    ///
391    /// The running CPU must support SHA-NI, SSSE3, and SSE4.1.
392    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
393    pub unsafe fn shani_x4_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
394        drive_slices(
395            crate::shani::STREAMS,
396            crate::shani::group,
397            prefix,
398            bodies,
399            out,
400        )
401    }
402
403    /// 8-lane AVX2 kernel
404    ///
405    /// # Safety
406    ///
407    /// The running CPU must support AVX2.
408    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
409    pub unsafe fn avx2_8(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
410        drive(8, crate::avx2::group, msgs, out)
411    }
412
413    /// # Safety
414    ///
415    /// The running CPU must support AVX2.
416    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
417    pub unsafe fn avx2_8_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
418        drive_slices(8, crate::avx2::group, prefix, bodies, out)
419    }
420
421    /// 16-lane AVX-512 kernel
422    ///
423    /// # Safety
424    ///
425    /// The running CPU must support AVX-512F and AVX-512BW.
426    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
427    pub unsafe fn avx512_16(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
428        drive(16, crate::avx512::group, msgs, out)
429    }
430
431    /// # Safety
432    ///
433    /// The running CPU must support AVX-512F and AVX-512BW.
434    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
435    pub unsafe fn avx512_16_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
436        drive_slices(16, crate::avx512::group, prefix, bodies, out)
437    }
438
439    /// 32-message AVX-512 interlace: two 16-lane compressions with their
440    /// rounds interlaced in one thread
441    ///
442    /// # Safety
443    ///
444    /// The running CPU must support AVX-512F and AVX-512BW.
445    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
446    pub unsafe fn avx512_16x2(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
447        drive(crate::avx512x2::WIDTH, crate::avx512x2::group2, msgs, out)
448    }
449
450    /// # Safety
451    ///
452    /// The running CPU must support AVX-512F and AVX-512BW.
453    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
454    pub unsafe fn avx512_16x2_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
455        drive_slices(
456            crate::avx512x2::WIDTH,
457            crate::avx512x2::group2,
458            prefix,
459            bodies,
460            out,
461        )
462    }
463
464    /// Fallback where no SHA-256 unit exists
465    pub fn chain_portable(seed: &[u8; 32], n: u64) -> [u8; 32] {
466        crate::chain::portable(seed, n)
467    }
468
469    /// Chain on the x86 SHA-NI extension
470    ///
471    /// # Safety
472    ///
473    /// The running CPU must support SHA-NI, SSSE3, and SSE4.1.
474    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
475    pub unsafe fn chain_shani(seed: &[u8; 32], n: u64) -> [u8; 32] {
476        crate::shani::chain(seed, n)
477    }
478
479    /// Chain on the ARMv8 SHA-256 crypto extension
480    ///
481    /// # Safety
482    ///
483    /// The running CPU must support the `sha2` extension.
484    #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
485    pub unsafe fn chain_neon_sha2(seed: &[u8; 32], n: u64) -> [u8; 32] {
486        crate::neon_sha2::chain(seed, n)
487    }
488
489    /// Lockstep chain kernels, one entry per dispatchable backend
490    ///
491    /// Each pins one step kernel under the shared scheduler, so any number of
492    /// chains streams through that kernel's width. `hash_chains` picks between
493    /// them and the serial chain; these exist so the choice can be measured
494    /// rather than asserted.
495    pub mod chains {
496        use crate::chain::run_scheduled;
497
498        pub fn portable1(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
499            // SAFETY: the scalar steps need no CPU feature.
500            unsafe { run_scheduled(1, crate::chain::steps_scalar1, seeds, lens, out) }
501        }
502        pub fn portable8(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
503            // SAFETY: as above.
504            unsafe { run_scheduled(8, crate::chain::steps_scalar8, seeds, lens, out) }
505        }
506
507        /// # Safety
508        ///
509        /// The running CPU must support the `sha2` extension.
510        #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
511        pub unsafe fn neon_sha2x4(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
512            run_scheduled(4, crate::neon_sha2::steps4, seeds, lens, out)
513        }
514
515        /// # Safety
516        ///
517        /// The running CPU must support the `sha2` extension.
518        #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
519        pub unsafe fn neon_sha2x8(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
520            run_scheduled(8, crate::neon_sha2::steps8, seeds, lens, out)
521        }
522
523        /// # Safety
524        ///
525        /// NEON is baseline on AArch64.
526        #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
527        pub unsafe fn neon4(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
528            run_scheduled(4, crate::neon::steps, seeds, lens, out)
529        }
530
531        /// # Safety
532        ///
533        /// The running CPU must support SHA-NI, SSSE3, and SSE4.1.
534        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
535        pub unsafe fn shani_x4(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
536            run_scheduled(4, crate::shani::steps4, seeds, lens, out)
537        }
538
539        /// # Safety
540        ///
541        /// The running CPU must support AVX2.
542        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
543        pub unsafe fn avx2_8(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
544            run_scheduled(8, crate::avx2::steps, seeds, lens, out)
545        }
546
547        /// # Safety
548        ///
549        /// The running CPU must support AVX-512F and AVX-512BW.
550        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
551        pub unsafe fn avx512_16(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
552            run_scheduled(16, crate::avx512::steps, seeds, lens, out)
553        }
554
555        /// # Safety
556        ///
557        /// The running CPU must support AVX-512F and AVX-512BW.
558        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
559        pub unsafe fn avx512_16x2(seeds: &[[u8; 32]], lens: &[u64], out: &mut [[u8; 32]]) {
560            run_scheduled(32, crate::avx512x2::steps2, seeds, lens, out)
561        }
562    }
563}
564
565mod dispatch {
566    use super::*;
567
568    /// The kernel to fall back on where no SIMD exists
569    ///
570    /// Multi-buffer costs a real general-purpose register per lane in scalar
571    /// code rather than riding along free in a vector lane, so the width that
572    /// wins is set by the register file. x86-64 has 16 GPRs and cannot absorb
573    /// even two lanes: eight measured 1.9x slower than one on Zen 4 (418.78 vs
574    /// 218.30 us). aarch64 has 31, and eight lanes measure 13% faster than one
575    /// (93.7 vs 106.5).
576    ///
577    /// One lane is the default because it is the safe side of that trade: it
578    /// gives up ~13% where the registers exist, while the wide kernel loses
579    /// 1.9x where they do not, and most targets are nearer x86-64 than aarch64
580    /// (32-bit x86 has 8 GPRs; wasm is JIT-ed onto whatever the host has).
581    /// aarch64 is listed because it was measured, not because it is 64-bit.
582    ///
583    /// Unreachable on an aarch64 build without `scalar`, since NEON is
584    /// baseline there and nothing falls through to it.
585    #[allow(dead_code)]
586    const PORTABLE: Kernel = if cfg!(target_arch = "aarch64") {
587        Kernel::Portable8
588    } else {
589        Kernel::Portable1
590    };
591
592    /// Which kernel this build and CPU resolve to
593    ///
594    /// Entry points, reported name, and lane width all read from a single
595    /// select, so they cannot drift. An earlier version kept two detection
596    /// ladders and derived the width by string-matching the backend name,
597    /// where a rename would silently have changed the reported width.
598    // Which variants exist depends on target and features; the rest are only
599    // match arms, which dead_code would otherwise flag.
600    #[allow(dead_code)]
601    #[derive(Clone, Copy, PartialEq, Eq)]
602    pub(crate) enum Kernel {
603        Portable1,
604        Portable8,
605        #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
606        Neon4,
607        #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
608        NeonSha2x4,
609        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
610        ShaNiX4,
611        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
612        Avx2_8,
613        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
614        Avx512_16,
615        #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
616        Avx512_16x2,
617        #[cfg(all(
618            target_arch = "wasm32",
619            target_feature = "simd128",
620            not(feature = "scalar")
621        ))]
622        Simd128_4,
623    }
624
625    impl Kernel {
626        pub(crate) fn name(self) -> &'static str {
627            match self {
628                Kernel::Portable1 => "portable-1",
629                Kernel::Portable8 => "portable-8",
630                #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
631                Kernel::Neon4 => "neon-4",
632                #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
633                Kernel::NeonSha2x4 => "neon-sha2-x4",
634                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
635                Kernel::ShaNiX4 => "shani-x4",
636                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
637                Kernel::Avx2_8 => "avx2-8",
638                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
639                Kernel::Avx512_16 => "avx512-16",
640                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
641                Kernel::Avx512_16x2 => "avx512-2x16",
642                #[cfg(all(
643                    target_arch = "wasm32",
644                    target_feature = "simd128",
645                    not(feature = "scalar")
646                ))]
647                Kernel::Simd128_4 => "simd128-4",
648            }
649        }
650
651        pub(crate) fn width(self) -> usize {
652            match self {
653                Kernel::Portable1 => 1,
654                Kernel::Portable8 => 8,
655                #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
656                Kernel::Neon4 => 4,
657                #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
658                Kernel::NeonSha2x4 => crate::neon_sha2::STREAMS,
659                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
660                Kernel::ShaNiX4 => crate::shani::STREAMS,
661                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
662                Kernel::Avx2_8 => 8,
663                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
664                Kernel::Avx512_16 => 16,
665                #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
666                Kernel::Avx512_16x2 => crate::avx512x2::WIDTH,
667                #[cfg(all(
668                    target_arch = "wasm32",
669                    target_feature = "simd128",
670                    not(feature = "scalar")
671                ))]
672                Kernel::Simd128_4 => 4,
673            }
674        }
675    }
676
677    /// Effective AMD CPU family, or 0 on anything else
678    ///
679    /// 0x19 is Zen 3 / Zen 4, 0x1a is Zen 5; the AVX-512 dispatch below is
680    /// per-family because the kernels rank differently on each. Cached:
681    /// `select` runs on every public call, and two `cpuid`s per call would
682    /// be measurable against a single small message.
683    #[cfg(all(
684        target_arch = "x86_64",
685        not(any(feature = "scalar", feature = "avx2", feature = "avx512"))
686    ))]
687    fn amd_family() -> u32 {
688        use std::sync::OnceLock;
689        static FAMILY: OnceLock<u32> = OnceLock::new();
690        *FAMILY.get_or_init(|| {
691            use std::arch::x86_64::__cpuid;
692            // SAFETY: leaves 0 and 1 exist on every x86_64 CPU. The blocks
693            // are redundant from 1.95, where `__cpuid` became safe, but the
694            // MSRV toolchain still requires them.
695            #[allow(unused_unsafe)]
696            let v = unsafe { __cpuid(0) };
697            // "AuthenticAMD" in the ebx/edx/ecx registers.
698            if (v.ebx, v.edx, v.ecx) != (0x6874_7541, 0x6974_6e65, 0x444d_4163) {
699                return 0;
700            }
701            #[allow(unused_unsafe)]
702            let eax = unsafe { __cpuid(1) }.eax;
703            // AMD reports base family 0xf and puts the rest in the extended
704            // field, so the effective family is their sum.
705            ((eax >> 8) & 0xf) + ((eax >> 20) & 0xff)
706        })
707    }
708
709    /// Selects the best kernel for this target and CPU
710    #[inline]
711    #[allow(clippy::needless_return)]
712    pub(crate) fn select() -> Kernel {
713        #[cfg(feature = "scalar")]
714        return PORTABLE;
715
716        #[cfg(all(target_arch = "x86_64", feature = "avx512", not(feature = "scalar")))]
717        return Kernel::Avx512_16;
718        #[cfg(all(target_arch = "x86_64", feature = "avx2", not(feature = "scalar")))]
719        return Kernel::Avx2_8;
720        #[cfg(all(target_arch = "aarch64", feature = "neon", not(feature = "scalar")))]
721        return Kernel::Neon4;
722
723        #[cfg(not(any(
724            feature = "scalar",
725            all(target_arch = "x86_64", any(feature = "avx2", feature = "avx512")),
726            all(target_arch = "aarch64", feature = "neon"),
727        )))]
728        {
729            #[cfg(target_arch = "x86_64")]
730            {
731                use std::sync::OnceLock;
732                // Each feature probe is itself a cached atomic load, but the
733                // ladder runs on every public call and stacks half a dozen
734                // of them; resolve the kernel once instead.
735                static KERNEL: OnceLock<Kernel> = OnceLock::new();
736                return *KERNEL.get_or_init(|| {
737                    let sha = have_shani();
738                    if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
739                        // AVX-512 capability alone does not decide; the family
740                        // does, and each arm here is a measurement. Zen 5 runs
741                        // the 2x16 interlace (12.7 vs 14.6 us single-wave on
742                        // EPYC 9B45): its native 512-bit datapath leaves
743                        // dependency stalls the second wave fills.
744                        //
745                        // Zen 4 double-pumps 512-bit ops, so the interlace measured
746                        // exactly neutral there (23.69 vs 23.72 on EPYC 9B14)
747                        // and the SHA-NI streams beat every 16-lane kernel
748                        // instead (21.05).
749                        //
750                        // Intel keeps the single wave for a different reason:
751                        // the interlace wins the cycles (998 vs 1053 per block
752                        // step on Emerald Rapids, 1004 vs 1048 on Granite Rapids)
753                        // and loses the clock, since doubling 512-bit density
754                        // costs 7 to 11% of frequency.
755                        let fam = amd_family();
756                        if fam == 0x1a {
757                            return Kernel::Avx512_16x2;
758                        }
759                        if !(sha && fam == 0x19) {
760                            return Kernel::Avx512_16;
761                        }
762                    }
763                    // Without AVX-512 (or on Zen 4), the dedicated SHA unit beats
764                    // 8-lane integer multi-buffer by a wide margin.
765                    // Measured on Zen 3, AVX2 loses even to serial SHA-NI.
766                    if sha {
767                        return Kernel::ShaNiX4;
768                    }
769                    if is_x86_feature_detected!("avx2") {
770                        return Kernel::Avx2_8;
771                    }
772                    PORTABLE
773                });
774            }
775            #[cfg(target_arch = "aarch64")]
776            {
777                // The dedicated SHA-256 unit beats integer multi-buffer, so
778                // prefer it and fall back to NEON lanes.
779                if std::arch::is_aarch64_feature_detected!("sha2") {
780                    return Kernel::NeonSha2x4;
781                }
782                Kernel::Neon4
783            }
784            #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
785            {
786                // Wasm has no runtime probing; simd128 is a build-time fact.
787                #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
788                {
789                    Kernel::Simd128_4
790                }
791                #[cfg(not(all(target_arch = "wasm32", target_feature = "simd128")))]
792                {
793                    PORTABLE
794                }
795            }
796        }
797    }
798
799    pub(crate) fn hash(msgs: &[Message<'_>], out: &mut [[u8; 32]]) {
800        match select() {
801            Kernel::Portable1 => backends::serial(msgs, out),
802            Kernel::Portable8 => backends::portable8(msgs, out),
803            #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
804            Kernel::Neon4 => backends::neon4(msgs, out),
805            #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
806            Kernel::NeonSha2x4 => unsafe { backends::neon_sha2x4(msgs, out) },
807            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
808            Kernel::ShaNiX4 => unsafe { backends::shani_x4(msgs, out) },
809            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
810            Kernel::Avx2_8 => unsafe { backends::avx2_8(msgs, out) },
811            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
812            k @ (Kernel::Avx512_16 | Kernel::Avx512_16x2) => unsafe {
813                let (w, group) = avx512_parts(k);
814                let cut = avx512_tail_cut(msgs.len());
815                let (head, tail) = msgs.split_at(cut);
816                let (out_head, out_tail) = out.split_at_mut(cut);
817                drive(w, group, head, out_head);
818                if !tail.is_empty() {
819                    backends::shani_x4(tail, out_tail);
820                }
821            },
822            #[cfg(all(
823                target_arch = "wasm32",
824                target_feature = "simd128",
825                not(feature = "scalar")
826            ))]
827            Kernel::Simd128_4 => backends::simd128_4(msgs, out),
828        }
829    }
830
831    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
832    #[inline]
833    pub(crate) fn have_shani() -> bool {
834        is_x86_feature_detected!("sha")
835            && is_x86_feature_detected!("ssse3")
836            && is_x86_feature_detected!("sse4.1")
837    }
838
839    /// Where to split an avx512 batch so its remainder runs on SHA-NI
840    ///
841    /// A remainder under 13 messages costs a flat 16-lane group mostly
842    /// running empty lanes, while SHA-NI streams pay only for the messages
843    /// present; at 13 or more the flat group is cheaper again. Routed when
844    /// the CPU has SHA-NI. Returns the batch length when no split should
845    /// happen.
846    ///
847    /// The cut is at wave (16-lane) granularity for both AVX-512 kernels:
848    /// the interlace processes 32 per chunk but degrades internally to
849    /// 16-lane groups for a partial chunk, so its routable straggler is
850    /// also `len % 16` -- cutting at `len % 32` would strand a 17..=28
851    /// remainder as 16 plus the exact mostly-empty group this exists to
852    /// avoid.
853    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
854    #[inline]
855    fn avx512_tail_cut(len: usize) -> usize {
856        let rem = len % 16;
857        if (1..=12).contains(&rem) && have_shani() {
858            len - rem
859        } else {
860            len
861        }
862    }
863
864    /// Chunk width and group function for an AVX-512 kernel choice
865    #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
866    fn avx512_parts(k: Kernel) -> (usize, batch::GroupFn) {
867        match k {
868            Kernel::Avx512_16x2 => (
869                crate::avx512x2::WIDTH,
870                crate::avx512x2::group2 as batch::GroupFn,
871            ),
872            _ => (16, crate::avx512::group as batch::GroupFn),
873        }
874    }
875
876    pub(crate) fn hash_slices(prefix: &[u8], bodies: &[&[u8]], out: &mut [[u8; 32]]) {
877        match select() {
878            Kernel::Portable1 => backends::serial_slices(prefix, bodies, out),
879            Kernel::Portable8 => backends::portable8_slices(prefix, bodies, out),
880            #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
881            Kernel::Neon4 => backends::neon4_slices(prefix, bodies, out),
882            #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
883            Kernel::NeonSha2x4 => unsafe { backends::neon_sha2x4_slices(prefix, bodies, out) },
884            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
885            Kernel::ShaNiX4 => unsafe { backends::shani_x4_slices(prefix, bodies, out) },
886            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
887            Kernel::Avx2_8 => unsafe { backends::avx2_8_slices(prefix, bodies, out) },
888            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
889            k @ (Kernel::Avx512_16 | Kernel::Avx512_16x2) => unsafe {
890                let (w, group) = avx512_parts(k);
891                let cut = avx512_tail_cut(bodies.len());
892                let (head, tail) = bodies.split_at(cut);
893                let (out_head, out_tail) = out.split_at_mut(cut);
894                drive_slices(w, group, prefix, head, out_head);
895                if !tail.is_empty() {
896                    backends::shani_x4_slices(prefix, tail, out_tail);
897                }
898            },
899            #[cfg(all(
900                target_arch = "wasm32",
901                target_feature = "simd128",
902                not(feature = "scalar")
903            ))]
904            Kernel::Simd128_4 => backends::simd128_4_slices(prefix, bodies, out),
905        }
906    }
907
908    pub(crate) fn hash_pairs(prefix: &[u8], left: &[&[u8]], right: &[&[u8]], out: &mut [[u8; 32]]) {
909        let k = select();
910        let width = k.width();
911        match k {
912            Kernel::Portable1 => unsafe {
913                drive_pairs(width, backends::scalar1, prefix, left, right, out)
914            },
915            Kernel::Portable8 => unsafe {
916                drive_pairs(width, backends::scalar8, prefix, left, right, out)
917            },
918            #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
919            Kernel::Neon4 => unsafe {
920                drive_pairs(width, crate::neon::group, prefix, left, right, out)
921            },
922            #[cfg(all(target_arch = "aarch64", not(feature = "scalar")))]
923            Kernel::NeonSha2x4 => unsafe {
924                drive_pairs(width, crate::neon_sha2::group, prefix, left, right, out)
925            },
926            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
927            Kernel::ShaNiX4 => unsafe {
928                drive_pairs(width, crate::shani::group, prefix, left, right, out)
929            },
930            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
931            Kernel::Avx2_8 => unsafe {
932                drive_pairs(width, crate::avx2::group, prefix, left, right, out)
933            },
934            #[cfg(all(target_arch = "x86_64", not(feature = "scalar")))]
935            k @ (Kernel::Avx512_16 | Kernel::Avx512_16x2) => unsafe {
936                let (w, group) = avx512_parts(k);
937                let cut = avx512_tail_cut(left.len());
938                drive_pairs(
939                    w,
940                    group,
941                    prefix,
942                    &left[..cut],
943                    &right[..cut],
944                    &mut out[..cut],
945                );
946                if cut < left.len() {
947                    drive_pairs(
948                        crate::shani::STREAMS,
949                        crate::shani::group,
950                        prefix,
951                        &left[cut..],
952                        &right[cut..],
953                        &mut out[cut..],
954                    );
955                }
956            },
957            #[cfg(all(
958                target_arch = "wasm32",
959                target_feature = "simd128",
960                not(feature = "scalar")
961            ))]
962            Kernel::Simd128_4 => unsafe {
963                drive_pairs(width, crate::simd128::group, prefix, left, right, out)
964            },
965        }
966    }
967
968    pub(crate) fn lane_width() -> usize {
969        select().width()
970    }
971
972    pub(crate) fn backend() -> &'static str {
973        select().name()
974    }
975}
976
977#[cfg(test)]
978mod tests {
979    use super::*;
980
981    #[test]
982    fn empty_message() {
983        let mut out = [[0u8; 32]; 1];
984        hash_many(&[b""], &mut out);
985        // FIPS 180-4 known answer for the empty string.
986        let expect = hex("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
987        assert_eq!(out[0], expect);
988    }
989
990    #[test]
991    fn abc() {
992        let mut out = [[0u8; 32]; 1];
993        hash_many(&[b"abc"], &mut out);
994        let expect = hex("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
995        assert_eq!(out[0], expect);
996    }
997
998    #[test]
999    fn prefix_matches_concatenation() {
1000        let prefix = b"\x00SOLANA_MERKLE_SHREDS_LEAF";
1001        let bodies: Vec<Vec<u8>> = (0..8usize).map(|i| vec![i as u8; 100 + i * 37]).collect();
1002        let body_refs: Vec<&[u8]> = bodies.iter().map(|b| b.as_slice()).collect();
1003
1004        let mut via_prefix = vec![[0u8; 32]; bodies.len()];
1005        hash_many_prefixed(prefix, &body_refs, &mut via_prefix);
1006
1007        let joined: Vec<Vec<u8>> = bodies
1008            .iter()
1009            .map(|b| [prefix.as_slice(), b].concat())
1010            .collect();
1011        let joined_refs: Vec<&[u8]> = joined.iter().map(|b| b.as_slice()).collect();
1012        let mut via_concat = vec![[0u8; 32]; bodies.len()];
1013        hash_many(&joined_refs, &mut via_concat);
1014
1015        assert_eq!(via_prefix, via_concat);
1016    }
1017
1018    fn hex(s: &str) -> [u8; 32] {
1019        let mut out = [0u8; 32];
1020        for (i, b) in out.iter_mut().enumerate() {
1021            *b = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).unwrap();
1022        }
1023        out
1024    }
1025}