diff --git a/.github/workflows/sntrup-kem.yml b/.github/workflows/sntrup-kem.yml index a52d579..d813f8e 100644 --- a/.github/workflows/sntrup-kem.yml +++ b/.github/workflows/sntrup-kem.yml @@ -53,6 +53,8 @@ jobs: - run: cargo test --no-default-features - run: cargo test - run: cargo test --all-features + # Stable exercises SIMD; the pre-1.95 x86 MSRV uses its scalar fallback. + - run: cargo test --features kem,serde,std - run: cargo test --features serde,force-scalar cross: diff --git a/Cargo.lock b/Cargo.lock index bc0835d..74e532b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1353,8 +1353,11 @@ dependencies = [ "criterion", "getrandom", "hex", + "hybrid-array", + "kem", "rand", "rand_chacha", + "rand_core", "serde", "serde_json", "serdect", diff --git a/sntrup-kem/Cargo.toml b/sntrup-kem/Cargo.toml index bbd4bc2..178414e 100644 --- a/sntrup-kem/Cargo.toml +++ b/sntrup-kem/Cargo.toml @@ -5,11 +5,16 @@ authors = ["Michael Lodder "] license = "MIT OR Apache-2.0" keywords = ["sntrup", "kem", "post-quantum", "cryptography", "NTRU"] description = "Pure Rust implementation of the Streamlined NTRU Prime KEM for all parameter sizes" +documentation = "https://docs.rs/sntrup-kem" homepage = "https://github.com/RustCrypto/KEMs/tree/master/sntrup-kem" repository = "https://github.com/RustCrypto/KEMs" categories = ["algorithms", "cryptography"] readme = "README.md" edition = "2024" +rust-version = "1.85" + +[package.metadata.docs.rs] +features = ["kem", "serde", "std", "alloc"] [features] default = ["kgen", "ecap", "dcap"] @@ -21,6 +26,8 @@ force-scalar = [] std = [] serde = ["dep:serdect", "dep:serde"] js = ["getrandom/wasm_js"] +# Implementations of the `kem` crate traits, covering all three operations at once. +kem = ["dep:kem", "dep:hybrid-array", "dep:rand_core", "kgen", "ecap", "dcap"] [dependencies] hex = "0.4" @@ -28,6 +35,9 @@ rand = "0.10.0" rand_chacha = "0.10.0" subtle = "2" getrandom = { version = "0.4", optional = true } +hybrid-array = { version = "0.4.14", features = ["extra-sizes"], optional = true } +kem = { version = "0.3", optional = true } +rand_core = { version = "0.10", optional = true } serde = { version = "1", optional = true, default-features = false } serdect = { version = "0.4", optional = true } # sha2 0.11 dropped the `asm` feature; hardware SHA acceleration is now selected @@ -40,9 +50,15 @@ zeroize = { version = "1", features = ["derive"] } criterion = "0.7" serde_json = "1" +[[example]] +name = "kem_traits" +path = "examples/kem_traits.rs" +required-features = ["kem"] + [[bench]] name = "mod" harness = false +required-features = ["kgen", "ecap", "dcap"] [lints.rust] missing_docs = "deny" diff --git a/sntrup-kem/README.md b/sntrup-kem/README.md index 50b3bd4..c12b547 100644 --- a/sntrup-kem/README.md +++ b/sntrup-kem/README.md @@ -35,7 +35,7 @@ All key and ciphertext sizes are in bytes. Sizes are fixed per parameter set usi - All six parameter sizes: sntrup653, sntrup761, sntrup857, sntrup953, sntrup1013, sntrup1277 - IND-CCA2 secure with implicit rejection - Constant-time operations throughout (branchless sort, constant-time comparison and selection) -- SIMD acceleration (AVX2 on x86_64, NEON on aarch64) with automatic detection +- SIMD acceleration with automatic run-time detection: AVX-512 and AVX2 (plus AVX-VNNI where present) on x86_64, NEON on aarch64 - Optional `serde` support via the `serde` feature - Deterministic key generation from a 32-byte seed @@ -48,10 +48,17 @@ The KEM API is split into three default features so downstream crates can pull i | `kgen` | **yes** | Key generation: `SntrupKem::generate_key`, `SntrupKem::generate_key_deterministic` | | `ecap` | **yes** | Encapsulation: `EncapsulationKey::encapsulate` | | `dcap` | **yes** | Decapsulation: `DecapsulationKey::decapsulate` | -| `force-scalar` | no | Disable SIMD (AVX2/NEON) and use pure-Rust scalar code | +| `alloc` | no | Allocator-dependent APIs | +| `std` | no | Standard-library integration; implies `alloc` | +| `force-scalar` | no | Compile out every SIMD kernel and use the portable scalar code paths only | +| `kem` | no | Implements the [`kem`](https://docs.rs/kem) crate's traits (`Encapsulate`, `Decapsulate`, `Kem`, ...) so this crate can be used generically alongside other KEMs. See [`sntrup_kem::kem`](src/kem.rs) and `examples/kem_traits.rs`. | | `serde` | no | Enables `Serialize`/`Deserialize` for all key and ciphertext types (via `serdect` for constant-time hex encoding) | | `js` | no | Enables WebAssembly support for `wasm32-unknown-unknown` by configuring `getrandom` to use JavaScript's `crypto.getRandomValues()` | +The synchronized x86_64 SIMD implementation requires Rust 1.95 or newer. Builds +with the Rust 1.85 MSRV automatically use the portable scalar paths; AArch64 +builds retain NEON acceleration on Rust 1.85. + To use only a subset of the KEM API, disable defaults and pick the features you need: ```toml @@ -175,6 +182,30 @@ let ek2 = EncapsulationKey::::try_from(ek_bytes).unwrap(); assert_eq!(ek, ek2); ``` +### `kem` crate integration + +With the `kem` feature enabled, the [`kem`](https://docs.rs/kem) module implements that crate's +traits for every parameter set, so Streamlined NTRU Prime can be used in generic code alongside +other KEMs. The traits and the parameter-set marker types are re-exported there, so no direct +dependency on the `kem` crate is needed: + +```rust +# #[cfg(feature = "kem")] { +use sntrup_kem::kem::{Decapsulate, Encapsulate, Kem, Sntrup761Params}; +use rand::SeedableRng; +use rand::rngs::{StdRng, SysRng}; + +let mut rng = StdRng::try_from_rng(&mut SysRng).expect("OS randomness"); + +let (dk, ek) = Sntrup761Params::generate_keypair_from_rng(&mut rng); +let (ct, sent) = ek.encapsulate_with_rng(&mut rng); +assert_eq!(dk.decapsulate(&ct), sent); +# } +``` + +Run `cargo run --release --example kem_traits --features kem` for KEM-generic code and key +export. + ## WebAssembly To compile for `wasm32-unknown-unknown`, enable the `js` feature so that `getrandom` uses JavaScript's `crypto.getRandomValues()` for randomness: @@ -206,10 +237,56 @@ For `wasm32-wasi` (or `wasm32-wasip1`), the `js` feature is **not** needed since This implementation has not undergone any security auditing and while care has been taken no guarantees can be made for either correctness or the constant time running of the underlying functions. **Please use at your own risk.** +Secret-derived heap temporaries (multiply scratch, Euclidean-inversion state, sampling +randomness, hash intermediates) are wiped with the [`zeroize`](https://docs.rs/zeroize) crate +before being freed. One documented exception: `generate_key_deterministic`'s ChaCha20 RNG state +cannot be wiped because `rand_chacha` offers no zeroization support. + #### Algorithm Streamlined NTRU Prime was first published in 2016. The algorithm still requires careful security review. Please see [here](https://ntruprime.cr.yp.to/warnings.html) for further warnings from the authors regarding NTRU Prime and lattice-based encryption schemes. +## Performance + +`cargo bench` runs this crate's Criterion suite (`benches/mod.rs`) across all six parameter +sets. The synchronized standalone implementation also has a +[comparison harness](https://github.com/mikelodder7/sntrup/tree/main/benches/comparison) for +sntrup761 — the parameter set with independent PQClean and liboqs implementations. + +This crate is faster than both C references on every operation, on both +architectures, while also zeroizing every secret-derived scratch buffer — which neither C +reference does. + +On x86_64 (AMD Ryzen AI 9 HX 370, Zen 5), sntrup761, against liboqs's AVX2 build: + +| Operation | sntrup | liboqs | PQClean | +|-----------|-------:|-------:|--------:| +| keypair | 106.6 µs | 107.9 µs (0.99x) | 4545.7 µs (42.7x) | +| encapsulate | 10.5 µs | 11.5 µs (0.91x) | 239.1 µs (22.9x) | +| decapsulate | 8.4 µs | 8.4 µs (1.00x) | 607.0 µs (72.6x) | + +On aarch64 (Apple M2 Max), sntrup761, against their portable C builds: keypair 684 µs +(2.7x), encapsulate 37.0 µs (1.40x), decapsulate 77.1 µs (1.19x). + +Two things drive the x86_64 numbers. Key generation runs the Bernstein–Yang divstep inversion +through **AVX-512**, 32 coefficients per step — neither PQClean nor liboqs has a 512-bit path +for this KEM. Encapsulation and decapsulation run sntrup761's polynomial multiply as a +number-theoretic transform (Good's 3x512 decomposition over the primes 7681 and 10753, +recombined by CRT). Every other parameter set, and all of aarch64, uses a schoolbook kernel +that computes each output coefficient as a contiguous dot product spread across eight +independent widening multiply-accumulate chains (`smlal`-family on NEON, `pmaddwd`/`vpdpwssd` +on x86_64) — a shape taken from disassembling what clang's autovectorizer produces for +PQClean's reference C and then out-tuning it. + +See the standalone implementation's +[benchmark results](https://github.com/mikelodder7/sntrup/blob/main/benches/comparison/RESULTS.md) +for the full investigation narrative, including machine and build details. + +**A SIMD-testing gotcha every contributor should read:** `--all-features` enables +`force-scalar`, which silently compiles the SIMD kernels out of the test binary. The permanent +kernel-vs-scalar differential tests in `src/rq.rs` and `src/r3.rs` only exercise SIMD when +built with a feature set that leaves `force-scalar` off, e.g. `--features kem,serde,std`. + # License Licensed under either of diff --git a/sntrup-kem/benches/mod.rs b/sntrup-kem/benches/mod.rs index 4b15433..852505e 100644 --- a/sntrup-kem/benches/mod.rs +++ b/sntrup-kem/benches/mod.rs @@ -1,4 +1,4 @@ -#![allow(missing_docs)] +#![allow(missing_docs, clippy::mod_module_files)] use criterion::{Criterion, criterion_group, criterion_main}; use sntrup_kem::*; diff --git a/sntrup-kem/build.rs b/sntrup-kem/build.rs new file mode 100644 index 0000000..21d1297 --- /dev/null +++ b/sntrup-kem/build.rs @@ -0,0 +1,36 @@ +//! Compiler-version compatibility configuration for x86 SIMD kernels. + +use std::{env, process::Command}; + +fn rustc_version() -> Option<(u32, u32)> { + let rustc = env::var_os("RUSTC")?; + let output = Command::new(rustc).arg("--version").output().ok()?; + let stdout = String::from_utf8(output.stdout).ok()?; + let version = stdout.split_whitespace().nth(1)?; + let mut components = version.split('.'); + let major = components.next()?.parse().ok()?; + let minor = components.next()?.parse().ok()?; + Some((major, minor)) +} + +fn main() { + println!("cargo:rerun-if-env-changed=RUSTC"); + println!("cargo:rerun-if-env-changed=TARGET"); + + if env::var("CARGO_CFG_TARGET_ARCH").as_deref() != Ok("x86_64") { + return; + } + + // AVX-512/AVX-VNNI intrinsics and the safe `target_feature` calling rules + // used by the synchronized implementation require Rust 1.95. Older x86 + // compilers use the same constant-time scalar paths as `force-scalar`. + let supports_simd = match rustc_version() { + Some((1, minor)) => minor >= 95, + Some((major, _)) => major > 1, + None => false, + }; + + if !supports_simd { + println!("cargo:rustc-cfg=feature=\"force-scalar\""); + } +} diff --git a/sntrup-kem/examples/kem_traits.rs b/sntrup-kem/examples/kem_traits.rs new file mode 100644 index 0000000..6dc3b0c --- /dev/null +++ b/sntrup-kem/examples/kem_traits.rs @@ -0,0 +1,63 @@ +/* + Copyright Michael Lodder. All Rights Reserved. + SPDX-License-Identifier: MIT OR Apache-2.0 +*/ +//! Streamlined NTRU Prime through the [`kem`](https://docs.rs/kem) crate traits, in code that +//! is generic over the parameter set and would work just as well over any other KEM. +//! +//! Run with: +//! +//! ```sh +//! cargo run --release --example kem_traits --features kem +//! ``` + +use rand::SeedableRng; +use rand::rngs::{StdRng, SysRng}; +use rand_core::CryptoRng; +use sntrup_kem::kem::{ + Decapsulate, DecapsulationKey, Decapsulator, Encapsulate, EncapsulationKey, Generate, Kem, + KemSizes, KeyExport, Sntrup653Params, Sntrup761Params, Sntrup1277Params, TryKeyInit, +}; + +/// Establish a shared secret and hand back the sizes involved. +/// +/// Nothing here names Streamlined NTRU Prime: the same function compiles against any KEM +/// whose key types implement the `kem` traits. +fn round_trip(mut rng: impl CryptoRng) -> (usize, usize) +where + K: KemSizes + + Kem, DecapsulationKey = DecapsulationKey>, +{ + let (dk, ek) = K::generate_keypair_from_rng(&mut rng); + let (ct, sent) = ek.encapsulate_with_rng(&mut rng); + let received = dk.decapsulate(&ct); + assert_eq!(sent, received); + (ct.len(), received.len()) +} + +fn main() { + // A cryptographically secure generator, seeded once from the operating system. + let mut rng = StdRng::try_from_rng(&mut SysRng).expect("the OS RNG is available"); + + for (name, (ct, ss)) in [ + ("Sntrup653Params", round_trip::(&mut rng)), + ("Sntrup761Params", round_trip::(&mut rng)), + ("Sntrup1277Params", round_trip::(&mut rng)), + ] { + println!("{name}: ciphertext {ct} bytes, shared secret {ss} bytes"); + } + + // Exporting or importing a key moves the whole key by value. Streamlined NTRU Prime keys + // are small (at most a few kilobytes), so unlike KEMs with megabyte-scale keys this needs + // no special thread-stack handling. + let dk = DecapsulationKey::::generate_from_rng(&mut rng); + let ek = dk.encapsulation_key(); + let exported = ek.to_bytes(); + let imported = + EncapsulationKey::::new(&exported).expect("exported key round-trips"); + assert_eq!(&imported, dk.encapsulation_key()); + println!( + "Sntrup761Params: exported and reimported {} bytes", + exported.len() + ); +} diff --git a/sntrup-kem/src/cpu.rs b/sntrup-kem/src/cpu.rs new file mode 100644 index 0000000..2776e0a --- /dev/null +++ b/sntrup-kem/src/cpu.rs @@ -0,0 +1,92 @@ +//! Runtime detection of AVX2 support, cached after the first check. +//! +//! The AVX2 kernels throughout this crate are marked `#[target_feature(enable = "avx2")]`, +//! which lets them compile regardless of the crate's ambient compilation flags. Whether to +//! *call* them is decided here, at runtime, so a default `cargo build --release` uses AVX2 on +//! any capable x86_64 host instead of silently falling back to scalar unless the caller passes +//! `RUSTFLAGS="-C target-feature=+avx2"` (or `target-cpu=native`). +//! +//! aarch64 needs no equivalent: NEON is a baseline guarantee of the architecture. + +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +use core::sync::atomic::{AtomicU8, Ordering}; + +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +static AVX2_STATE: AtomicU8 = AtomicU8::new(0); + +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +static AVXVNNI_STATE: AtomicU8 = AtomicU8::new(0); + +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +static AVX512_STATE: AtomicU8 = AtomicU8::new(0); + +/// Returns `true` if the host CPU supports AVX2. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[inline] +pub(crate) fn has_avx2() -> bool { + match AVX2_STATE.load(Ordering::Relaxed) { + 1 => true, + 2 => false, + _ => { + let detected = std::is_x86_feature_detected!("avx2"); + AVX2_STATE.store(if detected { 1 } else { 2 }, Ordering::Relaxed); + detected + } + } +} + +/// Returns `true` if the host CPU supports AVX-VNNI (the VEX-encoded `vpdpwssd` +/// family — Zen 5, Alder Lake and later). +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[inline] +pub(crate) fn has_avxvnni() -> bool { + match AVXVNNI_STATE.load(Ordering::Relaxed) { + 1 => true, + 2 => false, + _ => { + let detected = std::is_x86_feature_detected!("avxvnni"); + AVXVNNI_STATE.store(if detected { 1 } else { 2 }, Ordering::Relaxed); + detected + } + } +} + +/// Returns `true` if the host CPU supports the AVX-512 subsets the 512-bit +/// kernels use: `F` for the base instruction set, `BW` for 16-bit lane +/// arithmetic and `VL` so 256-bit forms remain available alongside. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[inline] +pub(crate) fn has_avx512() -> bool { + match AVX512_STATE.load(Ordering::Relaxed) { + 1 => true, + 2 => false, + _ => { + let detected = std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vl"); + AVX512_STATE.store(if detected { 1 } else { 2 }, Ordering::Relaxed); + detected + } + } +} + +#[cfg(all(test, target_arch = "x86_64", not(feature = "force-scalar")))] +mod tests { + use super::*; + + #[test] + fn detection_is_cached_and_consistent() { + let first = has_avx2(); + for _ in 0..8 { + assert_eq!(has_avx2(), first); + } + let first = has_avxvnni(); + for _ in 0..8 { + assert_eq!(has_avxvnni(), first); + } + let first = has_avx512(); + for _ in 0..8 { + assert_eq!(has_avx512(), first); + } + } +} diff --git a/sntrup-kem/src/kem.rs b/sntrup-kem/src/kem.rs index da02d07..43e0081 100644 --- a/sntrup-kem/src/kem.rs +++ b/sntrup-kem/src/kem.rs @@ -1,79 +1,299 @@ -//! Internal KEM operations for Streamlined NTRU Prime. +/* + Copyright Michael Lodder. All Rights Reserved. + SPDX-License-Identifier: MIT OR Apache-2.0 +*/ +//! Implementations of the traits from the [`kem`] crate. //! -//! Top-level keygen/encaps/decaps functions that delegate to `utils` for -//! the core cryptographic operations. +//! These wrap the [`crate::EncapsulationKey`] / [`crate::DecapsulationKey`] pair in the +//! fixed-size array types the `kem` traits use, so Streamlined NTRU Prime can be dropped into +//! generic code alongside other KEMs. The traits are re-exported here, so no direct dependency +//! on the `kem` crate is needed. +//! +//! # Example +//! +//! ``` +//! use sntrup_kem::kem::{Decapsulate, Encapsulate, Kem, Sntrup761Params}; +//! use rand::SeedableRng; +//! use rand::rngs::{StdRng, SysRng}; +//! +//! // A cryptographically secure generator, seeded once from the operating system. +//! let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap(); +//! +//! let (dk, ek) = Sntrup761Params::generate_keypair_from_rng(&mut rng); +//! let (ct, sent) = ek.encapsulate_with_rng(&mut rng); +//! let received = dk.decapsulate(&ct); +//! +//! assert_eq!(sent, received); +//! ``` + +use hybrid_array::{Array, ArraySize}; +/// The [`kem`] crate's traits, re-exported so that callers need not depend on a +/// version-matched copy of that crate themselves. +pub use kem::{ + Ciphertext, Decapsulate, Decapsulator, Encapsulate, Generate, InvalidKey, Kem, Key, KeyExport, + KeySizeUser, SharedKey, TryDecapsulate, TryKeyInit, +}; +use rand_core::{CryptoRng, TryCryptoRng}; +use zeroize::Zeroizing; -use crate::params::SntrupParameters; -use crate::{r3, utils, zx}; -use rand::CryptoRng; -use zeroize::Zeroize; +/// The parameter-set marker types, re-exported so callers can name them from this module. +pub use crate::{ + Sntrup653Params, Sntrup761Params, Sntrup857Params, Sntrup953Params, Sntrup1013Params, + Sntrup1277Params, +}; -/// Generate a Streamlined NTRU Prime key pair. +fn array_from_slice(bytes: &[u8]) -> Array { + let mut array = Array::default(); + array.copy_from_slice(bytes); + array +} + +/// Compile-time sizes used by the [`kem`] trait implementations. /// -/// Returns `(pk_bytes, sk_bytes)` as `Vec`. -#[cfg(feature = "kgen")] -pub(crate) fn keygen(params: &SntrupParameters, rng: &mut impl CryptoRng) -> (Vec, Vec) { - let p = params.p; - - // Generate g and its reciprocal in R3 - let mut g = vec![0i8; p]; - let mut gr = loop { - zx::random::random_small(&mut g, rng); - let (mask, mut gr) = r3::reciprocal(&g, p); - if mask == 0 { - break gr; - } - // Rejected reciprocal is still derived from the secret g — wipe it. - gr.zeroize(); - }; +/// Generic code over these traits names key types as `EncapsulationKey` and +/// `DecapsulationKey`, both of which require `K` to implement this trait, so it is part +/// of the public vocabulary even though every implementation lives in this crate. +pub trait KemSizes: + crate::SntrupParams + Copy + Clone + core::fmt::Debug + Eq + Ord + Send + Sync + 'static +{ + /// Encapsulation key size. + type EncapsulationKeySize: ArraySize; + /// Decapsulation key size. + type DecapsulationKeySize: ArraySize; + /// Ciphertext size. + type CiphertextSize: ArraySize; + /// Shared key size. + type SharedKeySize: ArraySize; +} + +/// A Streamlined NTRU Prime encapsulation key for use with the [`kem`] traits. +#[derive(Clone)] +pub struct EncapsulationKey(crate::EncapsulationKey

); + +/// A Streamlined NTRU Prime decapsulation key for use with the [`kem`] traits. +/// +/// The encapsulation key is kept alongside the private key because [`Decapsulator`] exposes +/// it. Unlike some KEMs, recovering it from the private key alone is cheap here — it is a +/// slice out of the embedded public key, not a recomputation — but caching it avoids an +/// allocation on every call. +pub struct DecapsulationKey { + key: crate::DecapsulationKey

, + encapsulation_key: EncapsulationKey

, +} - // Generate f with Hamming weight w - let mut f = vec![0i8; p]; - zx::random::random_tsmall(&mut f, p, params.w, rng); +impl core::fmt::Debug for EncapsulationKey

{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + self.0.fmt(f) + } +} - // Generate random rho for implicit rejection (raw random bytes, per PQClean) - let mut rho = vec![0u8; params.small_encode_size]; - rng.fill_bytes(&mut rho); +impl PartialEq for EncapsulationKey

{ + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} - let result = utils::derive_key(&f, &g, &gr, &rho, params); +impl Eq for EncapsulationKey

{} - // Zeroize secret intermediates - f.zeroize(); - g.zeroize(); - gr.zeroize(); - rho.zeroize(); +impl core::fmt::Debug for DecapsulationKey

{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("DecapsulationKey") + .field("algorithm", &P::NAME) + .finish_non_exhaustive() + } +} - result +impl KeySizeUser for EncapsulationKey

{ + type KeySize = P::EncapsulationKeySize; } -/// Encapsulate with a public key. -/// -/// Returns `(ciphertext_bytes, shared_secret_bytes)`. -#[cfg(feature = "ecap")] -pub(crate) fn encaps( - pk: &[u8], - params: &SntrupParameters, - rng: &mut impl CryptoRng, -) -> (Vec, Vec) { - let p = params.p; +impl TryKeyInit for EncapsulationKey

{ + fn new(key: &Key) -> Result { + crate::EncapsulationKey::

::try_from(key.as_slice()) + .map(Self) + .map_err(|_| InvalidKey) + } +} - // Generate random r with Hamming weight w - let mut r = vec![0i8; p]; - zx::random::random_tsmall(&mut r, p, params.w, rng); +impl KeyExport for EncapsulationKey

{ + fn to_bytes(&self) -> Key { + array_from_slice(self.0.as_ref()) + } +} - let (ct, ss) = utils::create_cipher(&r, pk, params); +impl KeySizeUser for DecapsulationKey

{ + type KeySize = P::DecapsulationKeySize; +} - // Zeroize secret intermediate - r.zeroize(); +impl TryKeyInit for DecapsulationKey

{ + /// Import a private key. Unlike KEMs that must rerun key generation to recover the + /// matching public key, the public key here is embedded in the private key's byte layout, + /// so recovering it is a slice, not a recomputation. + fn new(key: &Key) -> Result { + let key = crate::DecapsulationKey::

::try_from(key.as_slice()).map_err(|_| InvalidKey)?; + let encapsulation_key = EncapsulationKey(key.encapsulation_key()); + Ok(Self { + key, + encapsulation_key, + }) + } +} - (ct, ss.to_vec()) +impl KeyExport for DecapsulationKey

{ + fn to_bytes(&self) -> Key { + array_from_slice(self.key.as_ref()) + } } -/// Decapsulate with a secret key. -/// -/// Returns shared secret bytes. -#[cfg(feature = "dcap")] -pub(crate) fn decaps(sk: &[u8], ct: &[u8], params: &SntrupParameters) -> Vec { - let ss = utils::decapsulate_inner(ct, sk, params); - ss.to_vec() +impl Generate for DecapsulationKey

{ + fn try_generate_from_rng(rng: &mut R) -> Result { + let mut seed = Zeroizing::new([0u8; 32]); + rng.try_fill_bytes(seed.as_mut())?; + let (ek, dk) = crate::SntrupKem::

::generate_key_deterministic(&seed); + Ok(Self { + key: dk, + encapsulation_key: EncapsulationKey(ek), + }) + } +} + +impl

Decapsulator for DecapsulationKey

+where + P: KemSizes + Kem>, +{ + type Kem = P; + + fn encapsulation_key(&self) -> &EncapsulationKey

{ + &self.encapsulation_key + } +} + +impl

Decapsulate for DecapsulationKey

+where + P: KemSizes + Kem>, +{ + fn decapsulate(&self, ct: &Ciphertext

) -> SharedKey

{ + let Ok(ct) = crate::Ciphertext::

::try_from(ct.as_slice()) else { + return SharedKey::

::default(); + }; + array_from_slice(self.key.decapsulate(&ct).as_ref()) + } +} + +impl

Encapsulate for EncapsulationKey

+where + P: KemSizes + Kem, +{ + type Kem = P; + + fn encapsulate_with_rng(&self, mut rng: &mut R) -> (Ciphertext

, SharedKey

) + where + R: CryptoRng + ?Sized, + { + // The reborrow lets an unsized `R` (permitted by this trait) satisfy the concrete + // `impl rand::CryptoRng` bound on the underlying `encapsulate` method: `&mut R` + // implements `CryptoRng` via rand_core's blanket impl regardless of whether `R` itself + // is `Sized`, and a reference is always `Sized`. + let (ct, ss) = self.0.encapsulate(&mut rng); + (array_from_slice(ct.as_ref()), array_from_slice(ss.as_ref())) + } +} + +macro_rules! impl_kem { + ($($params:ident, $ek:ident, $dk:ident, $ct:ident, $ss:ident;)+) => { + $( + impl KemSizes for $params { + type EncapsulationKeySize = hybrid_array::sizes::$ek; + type DecapsulationKeySize = hybrid_array::sizes::$dk; + type CiphertextSize = hybrid_array::sizes::$ct; + type SharedKeySize = hybrid_array::sizes::$ss; + } + + impl Kem for $params { + type DecapsulationKey = DecapsulationKey; + type EncapsulationKey = EncapsulationKey; + type SharedKeySize = hybrid_array::sizes::$ss; + type CiphertextSize = hybrid_array::sizes::$ct; + } + )+ + }; +} + +impl_kem! { + Sntrup653Params, U994, U1518, U897, U32; + Sntrup761Params, U1158, U1763, U1039, U32; + Sntrup857Params, U1322, U1999, U1184, U32; + Sntrup953Params, U1505, U2254, U1349, U32; + Sntrup1013Params, U1623, U2417, U1455, U32; + Sntrup1277Params, U2067, U3059, U1847, U32; +} + +#[cfg(test)] +#[allow(clippy::unwrap_used)] +mod tests { + use super::*; + use rand_core::SeedableRng; + + fn traits_round_trip(seed: u8) + where + K: KemSizes + + Kem, DecapsulationKey = DecapsulationKey>, + { + let mut rng = rand_chacha::ChaCha8Rng::from_seed([seed; 32]); + let (dk, ek) = K::generate_keypair_from_rng(&mut rng); + + // The private key must not print its contents. + assert!(!format!("{dk:?}").contains(&hex::encode(dk.key.as_ref()))); + + let (ct, sent) = ek.encapsulate_with_rng(&mut rng); + let received = dk.decapsulate(&ct); + assert_eq!(sent, received); + + // Round tripping through the byte-array representation preserves both keys. + let ek_bytes = ek.to_bytes(); + let imported = EncapsulationKey::::new(&ek_bytes).unwrap(); + assert_eq!(imported, ek); + assert_eq!(dk.encapsulation_key(), &ek); + + let dk_bytes = dk.to_bytes(); + let imported = DecapsulationKey::::new(&dk_bytes).unwrap(); + assert_eq!(imported.decapsulate(&ct), sent); + assert_eq!(imported.encapsulation_key(), &ek); + } + + macro_rules! kem_trait_tests { + ($($name:ident, $params:ident, $seed:expr;)+) => { + $( + #[test] + fn $name() { + traits_round_trip::<$params>($seed); + } + )+ + }; + } + + // One per parameter set. + kem_trait_tests! { + round_trip_653, Sntrup653Params, 0x30; + round_trip_761, Sntrup761Params, 0x31; + round_trip_857, Sntrup857Params, 0x32; + round_trip_953, Sntrup953Params, 0x33; + round_trip_1013, Sntrup1013Params, 0x34; + round_trip_1277, Sntrup1277Params, 0x35; + } + + /// A ciphertext whose bytes are corrupted must not decapsulate to the original secret, + /// and must still return *some* shared key rather than propagate an error observably — + /// the `kem` crate's `Decapsulate` trait cannot report failure. + #[test] + fn corrupted_ciphertext_yields_a_different_key_without_panicking() { + let mut rng = rand_chacha::ChaCha8Rng::from_seed([0x41; 32]); + let (dk, ek) = Sntrup761Params::generate_keypair_from_rng(&mut rng); + let (ct, sent) = ek.encapsulate_with_rng(&mut rng); + + let mut damaged = ct; + let last = damaged.len() - 1; + damaged[last] ^= 0xFF; + assert_ne!(dk.decapsulate(&damaged), sent); + } } diff --git a/sntrup-kem/src/lib.rs b/sntrup-kem/src/lib.rs index 14bb092..f1d73f1 100644 --- a/sntrup-kem/src/lib.rs +++ b/sntrup-kem/src/lib.rs @@ -42,30 +42,64 @@ //! //! # Features //! -//! - `kgen`: Key generation (default) -//! - `ecap`: Encapsulation (default) -//! - `dcap`: Decapsulation (default) -//! - `serde`: Serde serialization support via `serdect` +//! - `kgen`: key generation (default) +//! - `ecap`: encapsulation (default) +//! - `dcap`: decapsulation (default) +//! - `kem`: implementations of the [`kem`](https://docs.rs/kem) crate's traits, +//! covering all three operations at once +//! - `serde`: `Serialize`/`Deserialize` for every key and ciphertext type, via +//! `serdect` for constant-time hex +//! - `alloc`: allocator-dependent APIs +//! - `std`: standard-library integration; implies `alloc` +//! - `js`: WebAssembly randomness for `wasm32-unknown-unknown` +//! - `force-scalar`: compile out every SIMD kernel and use the portable scalar +//! code paths only +//! +//! # SIMD dispatch +//! +//! Kernel selection happens at run time, from cached CPU-feature probes, so a +//! default `cargo build --release` uses the widest available instruction set +//! without needing `RUSTFLAGS` or `target-cpu=native`. Every SIMD kernel has a +//! scalar counterpart that produces bit-identical output, and the +//! `force-scalar` feature compiles the SIMD paths out entirely. +//! +//! | Target | Selected when | Used for | +//! |--------|---------------|----------| +//! | AVX-512 (F/BW/VL) | probed at run time | divstep inversion, 32 coefficients per step | +//! | AVX2 | probed at run time | everything else on x86_64 | +//! | AVX-VNNI | probed at run time | the schoolbook multiply, where available | +//! | NEON | baseline on `aarch64` | all vector kernels | +//! | scalar | no SIMD, or `force-scalar` | everything | +//! +//! The polynomial multiply takes two different shapes. For sntrup761 on x86_64 +//! it is a number-theoretic transform (Good's 3x512 decomposition over the +//! primes 7681 and 10753, recombined by CRT). Every other parameter set, and +//! all of `aarch64`, uses a schoolbook kernel that computes each output +//! coefficient as a contiguous dot product spread across eight independent +//! widening multiply-accumulate chains. -// The `kgen`/`ecap`/`dcap` features select which KEM operations are compiled. -// Building with a subset (or none) of them leaves some shared internal helpers -// (`ct`, `r3`, `rq`, `zx`, `utils`, and their imports) without a caller — that is -// expected, not a defect. Dead-code/unused-import enforcement is therefore scoped -// to the full-feature build (default + `--all-features`); partial builds tolerate -// the uncalled helpers so the crate stays warning-clean under `-D warnings`. +// Partial feature builds intentionally leave shared implementation helpers +// without callers. Keep those configurations warning-clean for workspace CI. #![cfg_attr( not(all(feature = "kgen", feature = "ecap", feature = "dcap")), allow(dead_code, unused_imports) )] +mod cpu; mod ct; mod error; -mod kem; +#[cfg(feature = "kem")] +pub mod kem; +mod ops; mod params; mod r3; mod rq; +mod scratch; +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +mod simd; mod types; mod utils; +mod wipe; mod zx; pub use error::Error; diff --git a/sntrup-kem/src/ops.rs b/sntrup-kem/src/ops.rs new file mode 100644 index 0000000..c0a9247 --- /dev/null +++ b/sntrup-kem/src/ops.rs @@ -0,0 +1,84 @@ +//! Internal KEM operations for Streamlined NTRU Prime. +//! +//! Top-level keygen/encaps/decaps functions that delegate to `utils` for +//! the core cryptographic operations. + +use crate::params::SntrupParameters; +use crate::{r3, utils, zx}; +use rand::CryptoRng; +use zeroize::Zeroize; + +/// Generate a Streamlined NTRU Prime key pair. +/// +/// Returns `(pk_bytes, sk_bytes)` as `Vec`. +#[cfg(feature = "kgen")] +pub(crate) fn keygen(params: &SntrupParameters, rng: &mut impl CryptoRng) -> (Vec, Vec) { + let p = params.p; + + // Generate g and its reciprocal in R3 + let mut g = vec![0i8; p]; + let mut gr = loop { + zx::random::random_small(&mut g, rng); + let (mask, mut gr) = r3::reciprocal(&g, p); + if mask == 0 { + break gr; + } + // Rejected reciprocal is still derived from the secret g — wipe it. + gr.zeroize(); + }; + + // Generate f with Hamming weight w + let mut f = vec![0i8; p]; + zx::random::random_tsmall(&mut f, p, params.w, rng); + + // Generate random rho for implicit rejection (raw random bytes, per PQClean) + let mut rho = vec![0u8; params.small_encode_size]; + rng.fill_bytes(&mut rho); + + let result = utils::derive_key(&f, &g, &gr, &rho, params); + + // Zeroize secret intermediates + f.zeroize(); + g.zeroize(); + gr.zeroize(); + rho.zeroize(); + + result +} + +/// Encapsulate with a public key, supplied pre-decoded. +/// +/// `h` is the decoded public-key polynomial and `pk_hash` is Hash4(pk); both are +/// per-key constants the caller caches so repeated encapsulations skip re-deriving +/// them. +/// +/// Returns `(ciphertext_bytes, shared_secret_bytes)`. +#[cfg(feature = "ecap")] +pub(crate) fn encaps( + h: &[i16], + pk_hash: &[u8; 32], + params: &SntrupParameters, + rng: &mut impl CryptoRng, +) -> (Vec, Vec) { + let p = params.p; + + // Generate random r with Hamming weight w + let mut r = vec![0i8; p]; + zx::random::random_tsmall(&mut r, p, params.w, rng); + + let (ct, ss) = utils::create_cipher(&r, h, pk_hash, params); + + // Zeroize secret intermediate + r.zeroize(); + + (ct, ss.to_vec()) +} + +/// Decapsulate with a secret key. +/// +/// Returns shared secret bytes. +#[cfg(feature = "dcap")] +pub(crate) fn decaps(sk: &[u8], h: &[i16], ct: &[u8], params: &SntrupParameters) -> Vec { + let ss = utils::decapsulate_inner(ct, sk, h, params); + ss.to_vec() +} diff --git a/sntrup-kem/src/params.rs b/sntrup-kem/src/params.rs index 900630c..28e40d3 100644 --- a/sntrup-kem/src/params.rs +++ b/sntrup-kem/src/params.rs @@ -3,6 +3,10 @@ /// Shared secret size in bytes. pub(crate) const SS_BYTES: usize = 32; +/// Largest `p` across the supported parameter sets — bounds stack scratch in +/// the allocation-free codec and KEM paths. +pub(crate) const MAX_P: usize = 1277; + /// Internal runtime parameter set for Streamlined NTRU Prime. #[doc(hidden)] #[derive(Debug, Clone, Copy)] diff --git a/sntrup-kem/src/r3.rs b/sntrup-kem/src/r3.rs index 3b1e206..a6b4bc0 100644 --- a/sntrup-kem/src/r3.rs +++ b/sntrup-kem/src/r3.rs @@ -1,52 +1,110 @@ +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +mod bitsliced; pub mod mod3; mod vector; use crate::ct::{smaller_mask, swap_int}; +use crate::wipe::wipe; -#[allow(clippy::cast_possible_wrap)] +/// Reciprocal in R/3, dispatched: bitsliced divstep on x86_64/AVX2, the +/// elimination form elsewhere. Same `(mask, r)` contract on both paths. +#[allow(unsafe_code)] pub fn reciprocal(s: &[i8], p: usize) -> (isize, Vec) { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return bitsliced::reciprocal_divstep(s, p); + } + } + #[allow(unreachable_code)] + reciprocal_eliminate(s, p) +} + +/// Top-coefficient elimination form: the non-x86/force-scalar path and the +/// differential oracle for the bitsliced divstep port. +#[allow(clippy::cast_possible_wrap)] +fn reciprocal_eliminate(s: &[i8], p: usize) -> (isize, Vec) { let loops = 2 * p + 1; + + // Buffers are padded to a multiple of the widest SIMD block (32 i8 lanes) so the + // vector kernels never fall into their scalar tail loops inside the hot iteration. + // Padding lanes start at zero; in u/v they provably stay zero (their source lanes + // are zero too), and in f/g anything written above index p only ever propagates + // upward — index p and below, the only entries ever read, are untouched by it. + let pad = |len: usize| (len + 31) & !31; + let mut r = vec![0i8; p]; - let mut f = vec![0i8; p + 1]; + let mut f = vec![0i8; pad(p + 1)]; f[0] = -1; f[1] = -1; f[p] = 1; - let mut g = vec![0i8; p + 1]; + let mut g = vec![0i8; pad(p + 1)]; g[..p].copy_from_slice(&s[..p]); + let fg_len = f.len(); let mut d = p as isize; let mut e = p as isize; - let mut u = vec![0i8; loops + 1]; - let mut v = vec![0i8; loops + 1]; + let mut u = vec![0i8; pad(loops + 1)]; + let mut v = vec![0i8; pad(loops + 1)]; + let uv_cap = u.len(); v[0] = 1; - for _ in 0..loops { + for i in 0..loops { let c = mod3::quotient(g[p], f[p]); - vector::minus_product_shift(&mut g, p + 1, &f, c); - vector::minus_product_shift(&mut v, loops + 1, &u, c); + // The swap mask needs the *post-shift* leading coefficient, so compute that + // single element scalar-first: new g[p] = freeze(g[p-1] - f[p-1]·c). This + // lets the shift and the conditional swap run as one fused memory pass. + let new_gp = mod3::minus_product(g[p - 1], f[p - 1], c); e -= 1; - let m = smaller_mask(e, d) & mod3::mask_set(g[p]); + let m = smaller_mask(e, d) & mod3::mask_set(new_gp); let (e_tmp, d_tmp) = swap_int(e, d, m); e = e_tmp; d = d_tmp; - vector::swap(&mut f, &mut g, p + 1, m); - vector::swap(&mut u, &mut v, loops + 1, m); + // After iteration i, the support of u and v is confined to indices 0..=i+1 + // (v starts as {0}, each shift grows it by one, and swaps only exchange the + // two): entries past that window are zero and stay zero, so processing + // `i + 2` elements computes exactly what the full-length pass would. The + // bound depends only on the public loop counter, never on secret data, so + // the constant-time property is unchanged. + let uv_len = pad(i + 2).min(uv_cap); + vector::minus_product_shift_cswap(&mut g, &mut f, fg_len, c, m); + vector::minus_product_shift_cswap(&mut v, &mut u, uv_len, c, m); } vector::product(&mut r, p, &u[p..], mod3::reciprocal(f[p])); + // The Euclidean state is derived from the secret input — wipe it before returning. + wipe(&mut f); + wipe(&mut g); + wipe(&mut u); + wipe(&mut v); (smaller_mask(0, d), r) } #[allow(unsafe_code)] pub fn mult(h: &mut [i8], f: &[i8], g: &[i8], p: usize) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return mult_avx2(h, f, g, p); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + { + // Same NTT machine as rq::mult, single-prime (product coefficients are + // bounded by p, far inside 7681/2). p = 761 only — see rq::ntt. + if p == 761 && crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return crate::rq::ntt::mult3_761(h, f, g); + } + } + if crate::cpu::has_avxvnni() { + // SAFETY: AVX2 + AVX-VNNI support confirmed by has_avxvnni() + unsafe { + return mult_avxvnni(h, f, g, p); + } + } + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return mult_avx2(h, f, g, p); + } + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -78,165 +136,371 @@ fn mult_scalar(h: &mut [i8], f: &[i8], g: &[i8], p: usize) { fg[i - p + 1] = mod3::freeze(fg[i - p + 1] as i32 + fg[i] as i32); } h[..p].copy_from_slice(&fg[..p]); + // At least one operand is secret at every call site — wipe the product scratch. + wipe(&mut fg); +} + +/// Row-major schoolbook multiplication for R3 polynomials on x86_64, expanded once per +/// instruction level by the macro below (`mult_avx2` via `crate::simd::mac_madd`, +/// `mult_avxvnni` via the fused `crate::simd::mac_vnni`). +/// +/// Same structure as `rq::mult`'s AVX2 kernel (see its doc comment): contiguous dot products +/// over widened copies of `f` and a reversed `g`, four independent `_mm256_madd_epi16` +/// accumulators (16 multiply-accumulates per instruction), one write per output coefficient, +/// and a single mod-3 freeze per output during the `x^p ≡ x + 1` fold (|folded sum| ≤ 3p, +/// well inside `mod3::freeze`'s i32 domain). +macro_rules! r3_mult_x86_kernel { + ($name:ident, $features:literal, $mac:path) => { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + #[target_feature(enable = $features)] + #[allow( + unsafe_code, + clippy::cast_possible_truncation, + clippy::needless_range_loop + )] + unsafe fn $name(h: &mut [i8], f: &[i8], g: &[i8], p: usize) { + unsafe { + use core::arch::x86_64::*; + + let fg_len = p * 2 - 1; + + let mut f16 = vec![0i16; p]; + let mut g_rev = vec![0i16; p]; + for i in 0..p { + f16[i] = f[i] as i16; + g_rev[i] = g[p - 1 - i] as i16; + } + + // Raw i16 row sums (|sum| ≤ p ≤ 1277), padded with one zero so the fold below may + // read `fg[k + p]` unconditionally at `k = p - 1`. + let mut fg = vec![0i16; fg_len + 1]; + for (i, out) in fg[..fg_len].iter_mut().enumerate() { + let jlo = i.saturating_sub(p - 1); + let len = i.min(p - 1) - jlo + 1; + let fp = f16.as_ptr().add(jlo); + // `p - 1 + jlo` never drops below `i` (jlo = max(0, i-p+1)); the naive + // `p - 1 - i + jlo` ordering underflows in debug builds when i ≥ p. + let gp = g_rev.as_ptr().add(p - 1 + jlo - i); + + let mut acc0 = _mm256_setzero_si256(); + let mut acc1 = _mm256_setzero_si256(); + let mut acc2 = _mm256_setzero_si256(); + let mut acc3 = _mm256_setzero_si256(); + let mut k = 0usize; + while k + 64 <= len { + acc0 = $mac( + acc0, + _mm256_loadu_si256(fp.add(k) as *const __m256i), + _mm256_loadu_si256(gp.add(k) as *const __m256i), + ); + acc1 = $mac( + acc1, + _mm256_loadu_si256(fp.add(k + 16) as *const __m256i), + _mm256_loadu_si256(gp.add(k + 16) as *const __m256i), + ); + acc2 = $mac( + acc2, + _mm256_loadu_si256(fp.add(k + 32) as *const __m256i), + _mm256_loadu_si256(gp.add(k + 32) as *const __m256i), + ); + acc3 = $mac( + acc3, + _mm256_loadu_si256(fp.add(k + 48) as *const __m256i), + _mm256_loadu_si256(gp.add(k + 48) as *const __m256i), + ); + k += 64; + } + while k + 16 <= len { + acc0 = $mac( + acc0, + _mm256_loadu_si256(fp.add(k) as *const __m256i), + _mm256_loadu_si256(gp.add(k) as *const __m256i), + ); + k += 16; + } + let s = _mm256_add_epi32( + _mm256_add_epi32(acc0, acc1), + _mm256_add_epi32(acc2, acc3), + ); + let s4 = + _mm_add_epi32(_mm256_castsi256_si128(s), _mm256_extracti128_si256(s, 1)); + let s2 = _mm_add_epi32(s4, _mm_shuffle_epi32(s4, 0b0000_1110)); + let s1 = _mm_add_epi32(s2, _mm_shuffle_epi32(s2, 0b0000_0001)); + let mut sum = _mm_cvtsi128_si32(s1); + while k < len { + sum += *fp.add(k) as i32 * *gp.add(k) as i32; + k += 1; + } + *out = sum as i16; + } + + // Fold x^p ≡ x + 1 with a single mod-3 freeze per output coefficient: fg[i] (i ≥ p) + // contributes to outputs i-p and i-p+1, and no fold target is itself ≥ p, so every + // output is independent. + h[0] = mod3::freeze(i32::from(fg[0]) + i32::from(fg[p])); + for k in 1..p { + h[k] = mod3::freeze( + i32::from(fg[k]) + i32::from(fg[k + p]) + i32::from(fg[k + p - 1]), + ); + } + } + } + }; } -/// Column-major schoolbook multiplication with AVX2 for R3 polynomials. -/// Uses _mm256_sign_epi16 for {-1,0,1} multiplication and i16 accumulators. -/// Processes 16 coefficients per SIMD instruction. -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] -#[target_feature(enable = "avx2")] +r3_mult_x86_kernel!(mult_avx2, "avx2", crate::simd::mac_madd); +r3_mult_x86_kernel!(mult_avxvnni, "avx2,avxvnni", crate::simd::mac_vnni); + +/// Row-major schoolbook multiplication with NEON for R3 polynomials. +/// +/// Same structure as `rq::mult`'s NEON kernel (see its doc comment for the reasoning): each +/// output coefficient's convolution sum is a contiguous dot product over `f` and a pre-reversed +/// `g`, held across EIGHT independent widening accumulators so the multiply-accumulate +/// latency is hidden, and written to memory once. Here the operands are ternary i8, so +/// `vmlal_s8`/`vmlal_high_s8` (i8×i8→i16, 8 lanes per instruction) process 64 elements per +/// unrolled iteration, and the i16 accumulator lanes stay far from overflow (each lane absorbs +/// at most `p/8` unit products; the final cross-lane sum is bounded by `p ≤ 1277`). +/// +/// The `x^p ≡ x + 1` fold then adds three raw row sums (|sum| ≤ 3p = 3831, well inside +/// `mod3::freeze`'s i32 domain) with a single freeze per output coefficient, instead of the +/// reference's three. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] #[allow( unsafe_code, clippy::cast_possible_truncation, clippy::needless_range_loop )] -unsafe fn mult_avx2(h: &mut [i8], f: &[i8], g: &[i8], p: usize) { +unsafe fn mult_neon(h: &mut [i8], f: &[i8], g: &[i8], p: usize) { unsafe { - use core::arch::x86_64::*; + use core::arch::aarch64::*; - let g_pad_len = (p + 15) & !15; // multiple of 16 - let fg_pad_len = p + g_pad_len; // >= 2p-1 let fg_len = p * 2 - 1; - // Sign-extend g to i16, padded - let mut g_pad = vec![0i16; g_pad_len]; + let mut g_rev = vec![0i8; p]; for i in 0..p { - g_pad[i] = g[i] as i16; + g_rev[i] = g[p - 1 - i]; } - // i16 accumulators (max value: ±p, fits in i16 for p <= 1277) - let mut fg = vec![0i16; fg_pad_len]; + // Raw i16 row sums, padded with one zero so the fold below may read `fg[k + p]` + // unconditionally at `k = p - 1`. + let mut fg = vec![0i16; fg_len + 1]; + for (i, out) in fg[..fg_len].iter_mut().enumerate() { + let jlo = i.saturating_sub(p - 1); + let len = i.min(p - 1) - jlo + 1; + let fp = f.as_ptr().add(jlo); + // `p - 1 + jlo` never drops below `i` (jlo = max(0, i-p+1)); the naive + // `p - 1 - i + jlo` ordering underflows in debug builds when i ≥ p. + let gp = g_rev.as_ptr().add(p - 1 + jlo - i); - // Column-major accumulation: fg[j+k] += f[j] * g[k] - for j in 0..p { - let fj = _mm256_set1_epi16(f[j] as i16); + let mut acc0 = vdupq_n_s16(0); + let mut acc1 = vdupq_n_s16(0); + let mut acc2 = vdupq_n_s16(0); + let mut acc3 = vdupq_n_s16(0); + let mut acc4 = vdupq_n_s16(0); + let mut acc5 = vdupq_n_s16(0); + let mut acc6 = vdupq_n_s16(0); + let mut acc7 = vdupq_n_s16(0); let mut k = 0usize; - while k + 16 <= g_pad_len { - let gk = _mm256_loadu_si256(g_pad.as_ptr().add(k) as *const __m256i); - // sign_epi16: if fj>0 → gk, if fj==0 → 0, if fj<0 → -gk - let prod = _mm256_sign_epi16(gk, fj); - let acc = _mm256_loadu_si256(fg.as_ptr().add(j + k) as *const __m256i); - _mm256_storeu_si256( - fg.as_mut_ptr().add(j + k) as *mut __m256i, - _mm256_add_epi16(acc, prod), - ); + while k + 64 <= len { + let f0 = vld1q_s8(fp.add(k)); + let f1 = vld1q_s8(fp.add(k + 16)); + let f2 = vld1q_s8(fp.add(k + 32)); + let f3 = vld1q_s8(fp.add(k + 48)); + let g0 = vld1q_s8(gp.add(k)); + let g1 = vld1q_s8(gp.add(k + 16)); + let g2 = vld1q_s8(gp.add(k + 32)); + let g3 = vld1q_s8(gp.add(k + 48)); + acc0 = vmlal_s8(acc0, vget_low_s8(f0), vget_low_s8(g0)); + acc1 = vmlal_high_s8(acc1, f0, g0); + acc2 = vmlal_s8(acc2, vget_low_s8(f1), vget_low_s8(g1)); + acc3 = vmlal_high_s8(acc3, f1, g1); + acc4 = vmlal_s8(acc4, vget_low_s8(f2), vget_low_s8(g2)); + acc5 = vmlal_high_s8(acc5, f2, g2); + acc6 = vmlal_s8(acc6, vget_low_s8(f3), vget_low_s8(g3)); + acc7 = vmlal_high_s8(acc7, f3, g3); + k += 64; + } + while k + 16 <= len { + let f0 = vld1q_s8(fp.add(k)); + let g0 = vld1q_s8(gp.add(k)); + acc0 = vmlal_s8(acc0, vget_low_s8(f0), vget_low_s8(g0)); + acc1 = vmlal_high_s8(acc1, f0, g0); k += 16; } + let total = vaddq_s16( + vaddq_s16(vaddq_s16(acc0, acc1), vaddq_s16(acc2, acc3)), + vaddq_s16(vaddq_s16(acc4, acc5), vaddq_s16(acc6, acc7)), + ); + let mut sum = i32::from(vaddvq_s16(total)); + while k < len { + sum += i32::from(*fp.add(k)) * i32::from(*gp.add(k)); + k += 1; + } + *out = sum as i16; } - // Vectorized mod-3 freeze: mulhrs(a, 10923) gives floor((a*10923+16384)/32768) - // which is the correct quotient for |a| <= 1277. - // Result: a - 3*q is in {-1, 0, 1}. - let k10923 = _mm256_set1_epi16(10923); - let three16 = _mm256_set1_epi16(3); - - let mut fg8 = vec![0i8; fg_len]; - let mut i = 0usize; - while i + 32 <= fg_len { - // Process 32 values: two batches of 16 i16 → 32 i8 - let a0 = _mm256_loadu_si256(fg.as_ptr().add(i) as *const __m256i); - let q0 = _mm256_mulhrs_epi16(a0, k10923); - let r0 = _mm256_sub_epi16(a0, _mm256_mullo_epi16(q0, three16)); + // Fold x^p ≡ x + 1 with a single mod-3 freeze per output coefficient: fg[i] (i ≥ p) + // contributes to outputs i-p and i-p+1, and no fold target is itself ≥ p, so every + // output is independent. + h[0] = mod3::freeze(i32::from(fg[0]) + i32::from(fg[p])); + for k in 1..p { + h[k] = mod3::freeze(i32::from(fg[k]) + i32::from(fg[k + p]) + i32::from(fg[k + p - 1])); + } + } +} - let a1 = _mm256_loadu_si256(fg.as_ptr().add(i + 16) as *const __m256i); - let q1 = _mm256_mulhrs_epi16(a1, k10923); - let r1 = _mm256_sub_epi16(a1, _mm256_mullo_epi16(q1, three16)); +#[cfg(test)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +mod tests { - // Pack 16+16 i16 → 32 i8, fix AVX2 lane ordering - let packed = _mm256_permute4x64_epi64(_mm256_packs_epi16(r0, r1), 0xD8); - _mm256_storeu_si256(fg8.as_mut_ptr().add(i) as *mut __m256i, packed); - i += 32; + /// The NTT mod-3 multiply must agree with the schoolbook kernel exactly. + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + #[test] + fn ntt_mult3_matches_scalar() { + if !crate::cpu::has_avx2() { + return; } - while i < fg_len { - fg8[i] = mod3::freeze(fg[i] as i32); - i += 1; + let p = 761usize; + let mut state = 0xabcd_ef01_2345_6789u64 | 1; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + }; + for trial in 0..8 { + let f: Vec = (0..p).map(|_| ((next() % 3) as i8) - 1).collect(); + let g: Vec = (0..p).map(|_| ((next() % 3) as i8) - 1).collect(); + let mut want = vec![0i8; p]; + mult_scalar(&mut want, &f, &g, p); + let mut got = vec![0i8; p]; + // SAFETY: AVX2 confirmed above. + unsafe { crate::rq::ntt::mult3_761(&mut got, &f, &g) }; + assert_eq!(got, want, "ntt mult3 vs scalar: trial={trial}"); } - - // Reduction: x^p ≡ x + 1 (mod x^p - x - 1) - for i in (p..(p * 2) - 1).rev() { - fg8[i - p] = mod3::freeze(fg8[i - p] as i32 + fg8[i] as i32); - fg8[i - p + 1] = mod3::freeze(fg8[i - p + 1] as i32 + fg8[i] as i32); + // Extremes: all +1 and all -1. + for &(fv, gv) in &[(1i8, 1i8), (-1, 1), (1, -1), (-1, -1)] { + let f = vec![fv; p]; + let g = vec![gv; p]; + let mut want = vec![0i8; p]; + mult_scalar(&mut want, &f, &g, p); + let mut got = vec![0i8; p]; + unsafe { crate::rq::ntt::mult3_761(&mut got, &f, &g) }; + assert_eq!(got, want, "ntt mult3 extremes f={fv} g={gv}"); } - h[..p].copy_from_slice(&fg8[..p]); } -} - -/// Column-major schoolbook multiplication with NEON for R3 polynomials. -/// Uses vmulq_s16 for {-1,0,1} multiplication and i16 accumulators. -/// Processes 8 coefficients per SIMD instruction. -#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] -#[allow( - unsafe_code, - clippy::cast_possible_truncation, - clippy::needless_range_loop -)] -unsafe fn mult_neon(h: &mut [i8], f: &[i8], g: &[i8], p: usize) { - unsafe { - use core::arch::aarch64::*; - let g_pad_len = (p + 7) & !7; // multiple of 8 - let fg_pad_len = p + g_pad_len; // >= 2p-1 - let fg_len = p * 2 - 1; - - // Sign-extend g to i16, padded - let mut g_pad = vec![0i16; g_pad_len]; - for i in 0..p { - g_pad[i] = g[i] as i16; + /// The bitsliced divstep port must agree with the elimination oracle on + /// both the invertibility mask and (when invertible) the reciprocal vector, + /// for every parameter size — including non-invertible inputs. + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + #[test] + fn bitsliced_reciprocal_matches_eliminate() { + if !crate::cpu::has_avx2() { + return; } - - // i16 accumulators (max value: ±p, fits in i16 for p <= 1277) - let mut fg = vec![0i16; fg_pad_len]; - - // Column-major accumulation: fg[j+k] += f[j] * g[k] - // vmulq_s16(gk, fj): for fj in {-1,0,1} this produces correct signed product - for j in 0..p { - let fj = vdupq_n_s16(f[j] as i16); - let mut k = 0usize; - while k + 8 <= g_pad_len { - let gk = vld1q_s16(g_pad.as_ptr().add(k)); - let prod = vmulq_s16(gk, fj); - let acc = vld1q_s16(fg.as_ptr().add(j + k)); - vst1q_s16(fg.as_mut_ptr().add(j + k), vaddq_s16(acc, prod)); - k += 8; + let mut state = 0x5eed_5eed_5eed_5eedu64 | 1; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + }; + for &p in &[653usize, 761, 857, 953, 1013, 1277] { + let mut invertible_seen = false; + let mut singular_seen = false; + for trial in 0..12 { + let g: Vec = (0..p).map(|_| ((next() % 3) as i8) - 1).collect(); + let (want_mask, want) = reciprocal_eliminate(&g, p); + // SAFETY: AVX2 confirmed above. + let (got_mask, got) = unsafe { bitsliced::reciprocal_divstep(&g, p) }; + assert_eq!(got_mask, want_mask, "mask p={p} trial={trial}"); + if want_mask == 0 { + invertible_seen = true; + assert_eq!(got, want, "value p={p} trial={trial}"); + } else { + singular_seen = true; + } } + // g = 0 is always singular. + let zero = vec![0i8; p]; + let (want_mask, _) = reciprocal_eliminate(&zero, p); + let (got_mask, _) = unsafe { bitsliced::reciprocal_divstep(&zero, p) }; + assert_ne!(want_mask, 0, "zero must be singular p={p}"); + assert_eq!(got_mask, want_mask, "zero mask p={p}"); + assert!(invertible_seen, "no invertible sample hit for p={p}"); + let _ = singular_seen; } + } + use super::*; - // Vectorized mod-3 freeze: vqrdmulhq_s16(a, 10923) gives correct quotient - // for |a| <= 1277. Result: a - 3*q is in {-1, 0, 1}. - let k10923 = vdupq_n_s16(10923); - let three16 = vdupq_n_s16(3); + /// Deterministic xorshift64* so the test needs no RNG crates or features. + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + } - let mut fg8 = vec![0i8; fg_len]; - let mut i = 0usize; - while i + 16 <= fg_len { - // Process 16 values: two batches of 8 i16 → 16 i8 - let a0 = vld1q_s16(fg.as_ptr().add(i)); - let q0 = vqrdmulhq_s16(a0, k10923); - let r0 = vsubq_s16(a0, vmulq_s16(q0, three16)); + fn random_ternary(p: usize, seed: u64) -> Vec { + let mut s = seed | 1; + (0..p).map(|_| ((next(&mut s) % 3) as i8) - 1).collect() + } - let a1 = vld1q_s16(fg.as_ptr().add(i + 8)); - let q1 = vqrdmulhq_s16(a1, k10923); - let r1 = vsubq_s16(a1, vmulq_s16(q1, three16)); + /// Compare every compiled-in SIMD kernel against the scalar reference. Catches the class + /// of bug the KAT/roundtrip suite can miss when run with `--all-features`, which enables + /// `force-scalar` and silently compiles the SIMD kernels out of the test entirely. + fn check_case(p: usize, f: &[i8], g: &[i8], label: &str) { + let mut want = vec![0i8; p]; + mult_scalar(&mut want, f, g, p); - // Pack 8+8 i16 → 16 i8 (naturally ordered, no permute needed) - let packed = vcombine_s8(vqmovn_s16(r0), vqmovn_s16(r1)); - vst1q_s8(fg8.as_mut_ptr().add(i), packed); - i += 16; + #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] + { + let mut got = vec![0i8; p]; + // SAFETY: NEON is baseline on aarch64 + unsafe { mult_neon(&mut got, f, g, p) }; + assert_eq!(got, want, "r3 mult_neon vs scalar: {label} p={p}"); + } + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + let mut got = vec![0i8; p]; + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { mult_avx2(&mut got, f, g, p) }; + assert_eq!(got, want, "r3 mult_avx2 vs scalar: {label} p={p}"); } - while i < fg_len { - fg8[i] = mod3::freeze(fg[i] as i32); - i += 1; + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avxvnni() { + let mut got = vec![0i8; p]; + // SAFETY: AVX2 + AVX-VNNI support confirmed by has_avxvnni() + unsafe { mult_avxvnni(&mut got, f, g, p) }; + assert_eq!(got, want, "r3 mult_avxvnni vs scalar: {label} p={p}"); + } + + let mut got = vec![0i8; p]; + mult(&mut got, f, g, p); + assert_eq!(got, want, "dispatched r3 mult vs scalar: {label} p={p}"); + } + + #[test] + fn simd_mult_matches_scalar_random() { + for p in [653usize, 761, 857, 953, 1013, 1277] { + for seed in 1..=8u64 { + let f = random_ternary(p, seed.wrapping_mul(0x9E37_79B9_7F4A_7C15)); + let g = random_ternary(p, seed.wrapping_mul(0xD1B5_4A32_D192_ED03)); + check_case(p, &f, &g, "random"); + } } + } - // Reduction: x^p ≡ x + 1 (mod x^p - x - 1) - for i in (p..(p * 2) - 1).rev() { - fg8[i - p] = mod3::freeze(fg8[i - p] as i32 + fg8[i] as i32); - fg8[i - p + 1] = mod3::freeze(fg8[i - p + 1] as i32 + fg8[i] as i32); + /// All-ones operands maximize accumulator magnitude, probing the i16 headroom the NEON + /// widening-accumulate staging depends on. + #[test] + fn simd_mult_extremes_match_scalar() { + for p in [653usize, 761, 857, 953, 1013, 1277] { + let ones = vec![1i8; p]; + let neg = vec![-1i8; p]; + check_case(p, &ones, &ones, "all +1"); + check_case(p, &ones, &neg, "+1 × -1"); } - h[..p].copy_from_slice(&fg8[..p]); } } diff --git a/sntrup-kem/src/r3/bitsliced.rs b/sntrup-kem/src/r3/bitsliced.rs new file mode 100644 index 0000000..17e5dc0 --- /dev/null +++ b/sntrup-kem/src/r3/bitsliced.rs @@ -0,0 +1,359 @@ +//! Bitsliced constant-time R/3 inversion (Bernstein–Yang divstep), ported from +//! the SUPERCOP AVX2 `crypto_core_inv3sntrup761` and generalized over all six +//! parameter sets. +//! +//! Ternary coefficients are stored as two bitplanes — plane 0 is the nonzero +//! bit, plane 1 the negative bit: `0 → (0,0)`, `1 → (1,0)`, `-1 → (1,1)` — with +//! 256 coefficients per `__m256i`, so a whole polynomial is `numvec = +//! ceil((p+1)/256)` registers per plane. Within a register, coefficient `j` +//! lives at bit `(j % 256) / 4` of 64-bit word `j % 4`: that interleaving makes +//! the shift-by-one-coefficient (`divx`/`timesx`) a 64-bit-word rotation plus a +//! single scalar shift, instead of a cross-register bit shift. +//! +//! The divstep loop eliminates the constant term (inputs are reversed) with +//! pure boolean operations — no multiplies anywhere. The reference hand-unrolls +//! five phases with fixed register counts; here the two register-width +//! schedules are computed per iteration from public data only, which is what +//! the phases encode: the V/R side grows one coefficient per iteration +//! (`min(numvec, k/256 + 1)`) and the F/G side shrinks with the remaining +//! iteration count (`ceil((2p - 1 - k)/256)`, the divstep degree-sum +//! invariant). Both formulas reproduce the reference's phase boundaries for +//! p = 761 exactly. +#![allow( + unsafe_code, + clippy::cast_possible_truncation, + clippy::cast_possible_wrap +)] + +use crate::wipe::wipe; +use core::arch::x86_64::*; + +const NUMVEC_MAX: usize = 5; // ceil((1277 + 1) / 256) + +#[inline] +fn numvec(p: usize) -> usize { + (p + 1).div_ceil(256) +} + +/// Bit-transpose 256 bytes (values 0/1) into one register, in the interleaved +/// coefficient order described in the module docs. +#[target_feature(enable = "avx2")] +fn frombits(b: &[i8]) -> __m256i { + unsafe { + const TRANSPOSE: [i8; 32] = [ + 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15, 16, 20, 24, 28, 17, 21, 25, 29, + 18, 22, 26, 30, 19, 23, 27, 31, + ]; + + let ld = |k: usize| _mm256_loadu_si256(b.as_ptr().add(32 * k) as *const __m256i); + let (b0, b1, b2, b3) = (ld(0), ld(1), ld(2), ld(3)); + let (b4, b5, b6, b7) = (ld(4), ld(5), ld(6), ld(7)); + + let c0 = _mm256_unpacklo_epi32(b0, b1); + let c1 = _mm256_unpackhi_epi32(b0, b1); + let c2 = _mm256_unpacklo_epi32(b2, b3); + let c3 = _mm256_unpackhi_epi32(b2, b3); + let c4 = _mm256_unpacklo_epi32(b4, b5); + let c5 = _mm256_unpackhi_epi32(b4, b5); + let c6 = _mm256_unpacklo_epi32(b6, b7); + let c7 = _mm256_unpackhi_epi32(b6, b7); + + let d0 = _mm256_or_si256(c0, _mm256_slli_epi32(c1, 2)); + let d2 = _mm256_or_si256(c2, _mm256_slli_epi32(c3, 2)); + let d4 = _mm256_or_si256(c4, _mm256_slli_epi32(c5, 2)); + let d6 = _mm256_or_si256(c6, _mm256_slli_epi32(c7, 2)); + + let e0 = _mm256_unpacklo_epi64(d0, d2); + let e2 = _mm256_unpackhi_epi64(d0, d2); + let e4 = _mm256_unpacklo_epi64(d4, d6); + let e6 = _mm256_unpackhi_epi64(d4, d6); + + let f0 = _mm256_or_si256(e0, _mm256_slli_epi32(e2, 1)); + let f4 = _mm256_or_si256(e4, _mm256_slli_epi32(e6, 1)); + + let g0 = _mm256_permute2x128_si256::<0x20>(f0, f4); + let g4 = _mm256_permute2x128_si256::<0x31>(f0, f4); + + let h = _mm256_or_si256(g0, _mm256_slli_epi32(g4, 4)); + let h = _mm256_shuffle_epi8(h, _mm256_loadu_si256(TRANSPOSE.as_ptr() as *const __m256i)); + let h = _mm256_permute4x64_epi64::<0xd8>(h); + _mm256_shuffle_epi32::<0xd8>(h) + } +} + +/// Inverse of [`frombits`]: one register back to 256 bytes of 0/1. +#[target_feature(enable = "avx2")] +fn tobits(h: __m256i, b: &mut [i8]) { + unsafe { + const TRANSPOSE: [i8; 32] = [ + 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15, 16, 20, 24, 28, 17, 21, 25, 29, + 18, 22, 26, 30, 19, 23, 27, 31, + ]; + + let h = _mm256_shuffle_epi32::<0xd8>(h); + let h = _mm256_permute4x64_epi64::<0xd8>(h); + let h = _mm256_shuffle_epi8(h, _mm256_loadu_si256(TRANSPOSE.as_ptr() as *const __m256i)); + + let g0 = _mm256_and_si256(h, _mm256_set1_epi8(15)); + let g4 = _mm256_and_si256(_mm256_srli_epi32(h, 4), _mm256_set1_epi8(15)); + + let f0 = _mm256_permute2x128_si256::<0x20>(g0, g4); + let f4 = _mm256_permute2x128_si256::<0x31>(g0, g4); + + let e0 = _mm256_and_si256(f0, _mm256_set1_epi8(5)); + let e2 = _mm256_and_si256(_mm256_srli_epi32(f0, 1), _mm256_set1_epi8(5)); + let e4 = _mm256_and_si256(f4, _mm256_set1_epi8(5)); + let e6 = _mm256_and_si256(_mm256_srli_epi32(f4, 1), _mm256_set1_epi8(5)); + + let d0 = _mm256_unpacklo_epi32(e0, e2); + let d2 = _mm256_unpackhi_epi32(e0, e2); + let d4 = _mm256_unpacklo_epi32(e4, e6); + let d6 = _mm256_unpackhi_epi32(e4, e6); + + let one = _mm256_set1_epi8(1); + let c0 = _mm256_and_si256(d0, one); + let c1 = _mm256_and_si256(_mm256_srli_epi32(d0, 2), one); + let c2 = _mm256_and_si256(d2, one); + let c3 = _mm256_and_si256(_mm256_srli_epi32(d2, 2), one); + let c4 = _mm256_and_si256(d4, one); + let c5 = _mm256_and_si256(_mm256_srli_epi32(d4, 2), one); + let c6 = _mm256_and_si256(d6, one); + let c7 = _mm256_and_si256(_mm256_srli_epi32(d6, 2), one); + + let st = |k: usize, v: __m256i, b: &mut [i8]| { + _mm256_storeu_si256(b.as_mut_ptr().add(32 * k) as *mut __m256i, v); + }; + st(0, _mm256_unpacklo_epi64(c0, c1), b); + st(1, _mm256_unpackhi_epi64(c0, c1), b); + st(2, _mm256_unpacklo_epi64(c2, c3), b); + st(3, _mm256_unpackhi_epi64(c2, c3), b); + st(4, _mm256_unpacklo_epi64(c4, c5), b); + st(5, _mm256_unpackhi_epi64(c4, c5), b); + st(6, _mm256_unpacklo_epi64(c6, c7), b); + st(7, _mm256_unpackhi_epi64(c6, c7), b); + } +} + +/// Bitplane-encode a ternary polynomial (given as bytes of 0/1 per plane). +#[target_feature(enable = "avx2")] +fn planes_from_small(dst0: &mut [__m256i], dst1: &mut [__m256i], s: &[i8], n: usize) { + let mut b0 = [0i8; NUMVEC_MAX * 256]; + let mut b1 = [0i8; NUMVEC_MAX * 256]; + for (i, &si) in s.iter().enumerate() { + b0[i] = si & 1; + b1[i] = (si >> 1) & b0[i]; + } + for i in 0..n { + dst0[i] = frombits(&b0[256 * i..]); + dst1[i] = frombits(&b1[256 * i..]); + } +} + +#[inline] +fn negative_mask(x: i32) -> i32 { + x >> 31 +} + +#[target_feature(enable = "avx2")] +fn swap(f: &mut [__m256i], g: &mut [__m256i], len: usize, mask: __m256i) { + for i in 0..len { + let flip = _mm256_and_si256(mask, _mm256_xor_si256(f[i], g[i])); + f[i] = _mm256_xor_si256(f[i], flip); + g[i] = _mm256_xor_si256(g[i], flip); + } +} + +/// `g -= c·f` in GF(3) on the bitplane encoding (then the caller divides by x). +#[target_feature(enable = "avx2")] +fn eliminate( + f0: &[__m256i], + f1: &[__m256i], + g0: &mut [__m256i], + g1: &mut [__m256i], + len: usize, + c0: __m256i, + c1: __m256i, +) { + for i in 0..len { + let f0i = _mm256_and_si256(f0[i], c0); + let f1i = _mm256_and_si256(_mm256_xor_si256(f1[i], c1), f0i); + let g0i = g0[i]; + let g1i = g1[i]; + + let t = _mm256_xor_si256(g0i, f0i); + g0[i] = _mm256_or_si256(t, _mm256_xor_si256(g1i, f1i)); + g1[i] = _mm256_and_si256(_mm256_xor_si256(g1i, f0i), _mm256_xor_si256(f1i, t)); + } +} + +/// Multiply V by the unit `c` (bitplane-broadcast masks), for the final scale. +#[target_feature(enable = "avx2")] +fn scale(f0: &mut [__m256i], f1: &mut [__m256i], n: usize, c0: __m256i, c1: __m256i) { + for i in 0..n { + let f0i = _mm256_and_si256(f0[i], c0); + f0[i] = f0i; + f1[i] = _mm256_and_si256(_mm256_xor_si256(f1[i], c1), f0i); + } +} + +/// Coefficient 0 (bit 0 of word 0 of register 0) as an all-ones/zero i32 mask. +#[target_feature(enable = "avx2")] +fn bit0mask(f: &[__m256i]) -> i32 { + -(_mm_cvtsi128_si32(_mm256_castsi256_si128(f[0])) & 1) +} + +/// Shift down by one coefficient: rotate each register's 64-bit words and fix +/// word 0 with a scalar shift chained from the next register. +#[target_feature(enable = "avx2")] +fn divx(f: &mut [__m256i], len: usize) { + let mut lows = [0u64; NUMVEC_MAX]; + for i in 0..len { + lows[i] = u64::from_ne_bytes(_mm_cvtsi128_si64(_mm256_castsi256_si128(f[i])).to_ne_bytes()); + } + for i in 0..len { + let next = if i + 1 < len { lows[i + 1] } else { 0 }; + let low = (lows[i] >> 1) | (next << 63); + let v = _mm256_blend_epi32::<0x3>( + f[i], + _mm256_set_epi64x(0, 0, 0, i64::from_ne_bytes(low.to_ne_bytes())), + ); + f[i] = _mm256_permute4x64_epi64::<0x39>(v); + } +} + +/// Shift up by one coefficient: inverse rotation, carries flow upward. +#[target_feature(enable = "avx2")] +fn timesx(f: &mut [__m256i], len: usize) { + let mut rot = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut lows = [0u64; NUMVEC_MAX]; + for i in 0..len { + rot[i] = _mm256_permute4x64_epi64::<0x93>(f[i]); + lows[i] = + u64::from_ne_bytes(_mm_cvtsi128_si64(_mm256_castsi256_si128(rot[i])).to_ne_bytes()); + } + for i in (0..len).rev() { + let prev = if i > 0 { lows[i - 1] } else { 0 }; + let low = (lows[i] << 1) | (prev >> 63); + f[i] = _mm256_blend_epi32::<0x3>( + rot[i], + _mm256_set_epi64x(0, 0, 0, i64::from_ne_bytes(low.to_ne_bytes())), + ); + } +} + +/// Constant-time reciprocal in R/3 via bitsliced divstep. +/// +/// Returns `(mask, r)` with the same contract as the elimination form: `mask == +/// 0` iff `s` is invertible, and `r` the reciprocal (valid when invertible). +#[target_feature(enable = "avx2")] +pub fn reciprocal_divstep(s: &[i8], p: usize) -> (isize, Vec) { + let n = numvec(p); + let total = 2 * p - 1; + + let mut f0 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut f1 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut g0 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut g1 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut v0 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut v1 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut r0 = [_mm256_setzero_si256(); NUMVEC_MAX]; + let mut r1 = [_mm256_setzero_si256(); NUMVEC_MAX]; + + // f = reversal of x^p - x - 1: coefficients 1 at 0, -1 at p-1, -1 at p. + let mut fs = [0i8; NUMVEC_MAX * 256]; + fs[0] = 1; + fs[p - 1] = -1; + fs[p] = -1; + planes_from_small(&mut f0, &mut f1, &fs[..n * 256], n); + + // g = reversal of s. + let mut gs = [0i8; NUMVEC_MAX * 256]; + for i in 0..p { + gs[i] = s[p - 1 - i]; + } + planes_from_small(&mut g0, &mut g1, &gs[..n * 256], n); + + // r = 1; v = 0. + r0[0] = _mm256_set_epi32(0, 0, 0, 0, 0, 0, 0, 1); + + let mut minusdelta: i32 = -1; + + for k in 0..total { + // Public width schedules (see module docs): V grows one coefficient + // per iteration; G's bit-length is bounded by the remaining + // iteration count (divstep degree-sum invariant). + let vw = n.min(k / 256 + 1); + let fw = n.min((total - k).div_ceil(256)); + + timesx(&mut v0, vw); + timesx(&mut v1, vw); + + let swapmask = negative_mask(minusdelta) & bit0mask(&g0); + let c0 = bit0mask(&f0) & bit0mask(&g0); + let c1 = (bit0mask(&f1) ^ bit0mask(&g1)) & c0; + + minusdelta ^= swapmask & (minusdelta ^ -minusdelta); + minusdelta -= 1; + + let swapvec = _mm256_set1_epi32(swapmask); + swap(&mut f0, &mut g0, fw, swapvec); + swap(&mut f1, &mut g1, fw, swapvec); + + let c0v = _mm256_set1_epi32(c0); + let c1v = _mm256_set1_epi32(c1); + eliminate(&f0, &f1, &mut g0, &mut g1, fw, c0v, c1v); + divx(&mut g0, fw); + divx(&mut g1, fw); + + swap(&mut v0, &mut r0, vw, swapvec); + swap(&mut v1, &mut r1, vw, swapvec); + eliminate(&v0, &v1, &mut r0, &mut r1, vw, c0v, c1v); + } + + // Scale V by the unit f0 and unpack, reversing back. + let c0v = _mm256_set1_epi32(bit0mask(&f0)); + let c1v = _mm256_set1_epi32(bit0mask(&f1)); + scale(&mut v0, &mut v1, n, c0v, c1v); + + let mut b0 = [0i8; NUMVEC_MAX * 256]; + let mut b1 = [0i8; NUMVEC_MAX * 256]; + for i in 0..n { + tobits(v0[i], &mut b0[256 * i..]); + tobits(v1[i], &mut b1[256 * i..]); + } + let mut out = vec![0i8; p]; + for (i, o) in out.iter_mut().enumerate() { + let (x0, x1) = (b0[p - 1 - i], b1[p - 1 - i]); + *o = x0 + 2 * x1 - 4 * (x0 & x1); + } + + // Everything above is derived from the secret input. The unpacked byte planes + // and the input scratch zeroize directly; the bitplane registers are wiped + // through their raw bytes, since `__m256i` has no `Zeroize` impl. + wipe(&mut b0); + wipe(&mut b1); + wipe(&mut fs); + wipe(&mut gs); + for regs in [ + f0.as_mut_ptr(), + f1.as_mut_ptr(), + g0.as_mut_ptr(), + g1.as_mut_ptr(), + v0.as_mut_ptr(), + v1.as_mut_ptr(), + r0.as_mut_ptr(), + r1.as_mut_ptr(), + ] { + // SAFETY: each array is NUMVEC_MAX `__m256i` values on this frame, so + // reinterpreting it as u64 words is in-bounds and correctly aligned + // (`__m256i` is 32-byte aligned, and its size is a multiple of 8). + unsafe { + wipe(core::slice::from_raw_parts_mut( + regs.cast::(), + NUMVEC_MAX * size_of::<__m256i>() / 8, + )); + } + } + + (negative_mask(minusdelta) as isize, out) +} diff --git a/sntrup-kem/src/r3/vector.rs b/sntrup-kem/src/r3/vector.rs index 6747cd1..38eed59 100644 --- a/sntrup-kem/src/r3/vector.rs +++ b/sntrup-kem/src/r3/vector.rs @@ -6,18 +6,18 @@ )] use super::mod3; +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +use crate::cpu::has_avx2; #[inline(always)] #[allow(clippy::cast_possible_truncation)] pub fn swap(x: &mut [i8], y: &mut [i8], n: usize, mask: isize) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return swap_avx2(x, y, n, mask); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return swap_avx2(x, y, n, mask); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -39,11 +39,7 @@ fn swap_scalar(x: &mut [i8], y: &mut [i8], n: usize, mask: isize) { } /// 32 i8 elements per SIMD iteration. -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] unsafe fn swap_avx2(x: &mut [i8], y: &mut [i8], n: usize, mask: isize) { unsafe { @@ -101,14 +97,12 @@ unsafe fn swap_neon(x: &mut [i8], y: &mut [i8], n: usize, mask: isize) { #[inline(always)] pub fn product(z: &mut [i8], n: usize, x: &[i8], c: i8) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return product_avx2(z, n, x, c); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return product_avx2(z, n, x, c); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -127,11 +121,7 @@ fn product_scalar(z: &mut [i8], n: usize, x: &[i8], c: i8) { /// For c in {-1, 0, 1}: _mm256_sign_epi8(x, c) computes x * c. /// Processes 32 elements per iteration. -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] unsafe fn product_avx2(z: &mut [i8], n: usize, x: &[i8], c: i8) { unsafe { @@ -196,18 +186,16 @@ unsafe fn product_neon(z: &mut [i8], n: usize, x: &[i8], c: i8) { } } -/// Fused minus_product and shift: `z[i+1] = freeze(z[i] - y[i]*c)`, `z[0] = 0`. +/// Fused minus_product and shift: `z[i+1] = freeze(z[i] - y[i]*c), z[0] = 0`. /// Processes backward to avoid overwrite conflicts, eliminating a separate memmove. #[inline(always)] pub fn minus_product_shift(z: &mut [i8], n: usize, y: &[i8], c: i8) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return minus_product_shift_avx2(z, n, y, c); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return minus_product_shift_avx2(z, n, y, c); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -225,11 +213,104 @@ fn minus_product_shift_scalar(z: &mut [i8], n: usize, y: &[i8], c: i8) { z[0] = 0; } -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +/// Fused `minus_product_shift` + conditional swap — one memory pass instead of two. +/// +/// Same contract as the rq version: exactly `minus_product_shift(z, n, y, c)` then +/// `swap(z, y, n, mask)`; non-AVX2 targets (including aarch64/NEON, untouched) +/// take the two-pass fallback. +#[inline(always)] +pub fn minus_product_shift_cswap(z: &mut [i8], y: &mut [i8], n: usize, c: i8, mask: isize) { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return minus_product_shift_cswap_avx2(z, y, n, c, mask); + } + } + minus_product_shift(z, n, y, c); + swap(z, y, n, mask); +} + +/// AVX2 fused kernel. The mod-3 fixup is a `vpshufb` register LUT: `r ∈ [-2,2]` +/// biased to `[0,4]` indexes the table `[1, -1, 0, 1, -1]` (freeze of `r`), +/// replacing the 6-op compare/mask chain with add + shuffle. vpshufb is an +/// in-lane register permute with data-independent timing — constant-time safe. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx2")] +unsafe fn minus_product_shift_cswap_avx2(z: &mut [i8], y: &mut [i8], n: usize, c: i8, mask: isize) { + unsafe { + use core::arch::x86_64::*; + let cv = _mm256_set1_epi8(c); + let two = _mm256_set1_epi8(2); + #[rustfmt::skip] + let table = _mm256_setr_epi8( + 1, -1, 0, 1, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 1, -1, 0, 1, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + ); + let mv = _mm256_set1_epi8(mask as i8); + + let mut j = (n - 2) as isize; + + while j >= 31 { + let start = (j - 31) as usize; + let zv = _mm256_loadu_si256(z.as_ptr().add(start) as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr().add(start) as *const __m256i); + let r = _mm256_sub_epi8(zv, _mm256_sign_epi8(yv, cv)); + let w = _mm256_shuffle_epi8(table, _mm256_add_epi8(r, two)); + + let y1 = _mm256_loadu_si256(y.as_ptr().add(start + 1) as *const __m256i); + let new_z = _mm256_blendv_epi8(w, y1, mv); + let new_y = _mm256_blendv_epi8(y1, w, mv); + _mm256_storeu_si256(z.as_mut_ptr().add(start + 1) as *mut __m256i, new_z); + _mm256_storeu_si256(y.as_mut_ptr().add(start + 1) as *mut __m256i, new_y); + j -= 32; + } + + // Bottom overlapped block; the +1 loads' topmost byte (index 32) is + // post-swap and is preserved via the keep mask rather than recomputed. + if j >= 0 && n >= 33 && n & 31 == 0 { + let zv = _mm256_loadu_si256(z.as_ptr() as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr() as *const __m256i); + let r = _mm256_sub_epi8(zv, _mm256_sign_epi8(yv, cv)); + let w = _mm256_shuffle_epi8(table, _mm256_add_epi8(r, two)); + + let z1 = _mm256_loadu_si256(z.as_ptr().add(1) as *const __m256i); + let y1 = _mm256_loadu_si256(y.as_ptr().add(1) as *const __m256i); + let new_z = _mm256_blendv_epi8(w, y1, mv); + let new_y = _mm256_blendv_epi8(y1, w, mv); + #[rustfmt::skip] + let keep = _mm256_setr_epi8( + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, + ); + _mm256_storeu_si256( + z.as_mut_ptr().add(1) as *mut __m256i, + _mm256_blendv_epi8(new_z, z1, keep), + ); + _mm256_storeu_si256( + y.as_mut_ptr().add(1) as *mut __m256i, + _mm256_blendv_epi8(new_y, y1, keep), + ); + j = -1; + } + + // Scalar remainder (only when n < 33). + let mi = mask as i8; + while j >= 0 { + let k = (j + 1) as usize; + let w = mod3::minus_product(z[k - 1], y[k - 1], c); + let yk = y[k]; + z[k] = (mi & yk) | (!mi & w); + y[k] = (mi & w) | (!mi & yk); + j -= 1; + } + let y0 = y[0]; + z[0] = mi & y0; + y[0] = !mi & y0; + } +} + +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] unsafe fn minus_product_shift_avx2(z: &mut [i8], n: usize, y: &[i8], c: i8) { unsafe { @@ -257,7 +338,25 @@ unsafe fn minus_product_shift_avx2(z: &mut [i8], n: usize, y: &[i8], c: i8) { j -= 32; } - // Scalar remainder + // The backward loop strands `(n - 2) % 32` bottom elements. When n is a + // multiple of 32 and the body ran at least once, a final full-width block at + // start = 0 covers them: + // z[0..32] is still original (higher blocks only wrote z[32..] and beyond), + // and the overlap element it rewrites (z[32]) gets the identical value the + // previous block computed from the same inputs. + if j >= 0 && n >= 33 && n & 31 == 0 { + let zv = _mm256_loadu_si256(z.as_ptr() as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr() as *const __m256i); + let yc = _mm256_sign_epi8(yv, cv); + let r = _mm256_sub_epi8(zv, yc); + let add = _mm256_and_si256(three, _mm256_cmpeq_epi8(r, neg2)); + let sub = _mm256_and_si256(three, _mm256_cmpeq_epi8(r, pos2)); + let r = _mm256_add_epi8(_mm256_sub_epi8(r, sub), add); + _mm256_storeu_si256(z.as_mut_ptr().add(1) as *mut __m256i, r); + j = -1; + } + + // Scalar remainder (only when n < 33) while j >= 0 { z[(j + 1) as usize] = mod3::minus_product(z[j as usize], y[j as usize], c); j -= 1; @@ -304,3 +403,43 @@ unsafe fn minus_product_shift_neon(z: &mut [i8], n: usize, y: &[i8], c: i8) { z[0] = 0; } } + +#[cfg(test)] +mod tests { + use super::*; + + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + /// The fused kernel must match scalar minus_product_shift + scalar swap + /// exactly, for both mask values, at lengths exercising the vector loop, + /// the overlapped bottom block, and the scalar path. + #[test] + fn fused_cswap_matches_two_pass_reference() { + let mut s = 0xfeed_face_cafe_beefu64; + for &n in &[2usize, 9, 32, 33, 65, 762, 768, 1536] { + for &mask in &[0isize, -1] { + for &c in &[-1i8, 0, 1] { + let z0: Vec = (0..n).map(|_| ((next(&mut s) % 3) as i8) - 1).collect(); + let y0: Vec = (0..n).map(|_| ((next(&mut s) % 3) as i8) - 1).collect(); + + let mut z_ref = z0.clone(); + let mut y_ref = y0.clone(); + minus_product_shift_scalar(&mut z_ref, n, &y_ref, c); + swap_scalar(&mut z_ref, &mut y_ref, n, mask); + + let mut z_got = z0.clone(); + let mut y_got = y0.clone(); + minus_product_shift_cswap(&mut z_got, &mut y_got, n, c, mask); + + assert_eq!(z_got, z_ref, "z mismatch n={n} mask={mask} c={c}"); + assert_eq!(y_got, y_ref, "y mismatch n={n} mask={mask} c={c}"); + } + } + } + } +} diff --git a/sntrup-kem/src/rq.rs b/sntrup-kem/src/rq.rs index 90cc218..fbd6ddd 100644 --- a/sntrup-kem/src/rq.rs +++ b/sntrup-kem/src/rq.rs @@ -1,44 +1,188 @@ +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +pub mod codec761; pub mod encoding; pub mod modq; +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +pub(crate) mod ntt; mod vector; use crate::ct::{smaller_mask, swap_int}; use crate::params::SntrupParameters; +use crate::wipe::wipe; -#[allow(clippy::cast_possible_wrap)] +/// Reciprocal of `3·s` in R/q, dispatched: divstep on x86_64/AVX2, the +/// top-coefficient-elimination form elsewhere. Both produce identical canonical +/// output (differentially tested). pub fn reciprocal3(s: &[i8], params: &SntrupParameters) -> Vec { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + return reciprocal3_divstep(s, params); + } + #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] + return reciprocal3_divstep(s, params); + #[allow(unreachable_code)] + reciprocal3_eliminate(s, params) +} + +/// Constant-time divstep inversion (Bernstein–Yang, ported from the SUPERCOP +/// AVX2 `crypto_core_invsntrup761`, generalized over all six parameter sets). +/// +/// The polynomials are processed *reversed*, so eliminating the leading +/// coefficient of the original becomes eliminating the constant term, and the +/// divstep `/x` is a shift-down. Each iteration eliminates via the cross-multiply +/// `g' = (f0·g − g0·f)/x` — two Montgomery products per element and **no +/// division anywhere in the loop** (the old form pays a ~20-freeze Fermat +/// quotient chain per iteration). The second loop's f/g window shrinks by one +/// per iteration on a public schedule. Montgomery 2⁻¹⁶ factors cancel in the +/// output because each divstep scales a complete row of the transition matrix, +/// and the result is the ratio `v/f0` of same-row entries (times the 1/3 folded +/// into r's initialization). +/// +/// The swap decision uses `delta` and the strictly-frozen `g0`: masks only, +/// applied through blends — no secret-dependent branch, index, or bound. +#[cfg(all( + any(target_arch = "x86_64", target_arch = "aarch64"), + not(feature = "force-scalar") +))] +#[allow(clippy::cast_possible_truncation)] +fn reciprocal3_divstep(s: &[i8], params: &SntrupParameters) -> Vec { + let p = params.p; + let q = params.q; + let b1 = params.barrett1; + let b2 = params.barrett2; + // The widest pass is 32 lanes, and `x[1..]` may be touched up to + // `x[1 ..= ceil32(len)]`. Padding to the widest kernel costs the narrower + // ones nothing: the extra lanes stay zero, and zero is a fixed point of the + // elimination step, so they neither affect nor are affected by the result. + let ppad = 1 + p.next_multiple_of(32); + + /// -1 if x != 0 (x canonical), else 0. + fn nonzero_mask(x: i16) -> i32 { + let v = u32::from(u16::from_ne_bytes(x.to_ne_bytes())); + -i32::from_ne_bytes(((!v).wrapping_add(1) >> 31).to_ne_bytes()) + } + /// -1 if x < 0, else 0. + fn negative_mask(x: i32) -> i32 { + x >> 31 + } + + let mut f = vec![0i16; ppad]; + f[0] = 1; + f[p - 1] = -1; + f[p] = -1; + // g = reversal of s (the reversal makes the divstep shift go downward). + let mut g = vec![0i16; ppad]; + for i in 0..p { + g[i] = i16::from(s[p - 1 - i]); + } + let mut v = vec![0i16; ppad]; + let mut r = vec![0i16; ppad]; + // Folds the "3" of 1/(3s) into the Bezout side. + r[0] = modq::reciprocal(3, q, b1, b2); + + let mut delta: i32 = 1; + + let step = |f: &mut Vec, + g: &mut Vec, + v: &mut Vec, + r: &mut Vec, + delta: &mut i32, + fg_len: usize, + vr_len: usize| { + let g0 = modq::freeze(i32::from(g[0]), q, b1, b2); + let f0 = modq::freeze(i32::from(f[0]), q, b1, b2); + + let swap = negative_mask(-*delta) & nonzero_mask(g0); + *delta ^= swap & (*delta ^ -*delta); + *delta += 1; + + let flip = (swap as i16) & (f0 ^ g0); + let f0 = f0 ^ flip; + let g0 = g0 ^ flip; + f[0] = f0; + + let mask = swap as isize; + // Buffers provide the capacity the widest kernel requires; the + // dispatchers pick the AVX-512, AVX2 or NEON implementation. + vector::swapeliminate(f, g, fg_len, f0, g0, mask, q); + vector::xswapeliminate(v, r, vr_len, f0, g0, mask, q); + }; + + for loop_i in 0..p { + step(&mut f, &mut g, &mut v, &mut r, &mut delta, p, loop_i + 1); + } + for loop_i in (1..p).rev() { + step(&mut f, &mut g, &mut v, &mut r, &mut delta, loop_i, p); + } + + let scale = modq::reciprocal(modq::freeze(i32::from(f[0]), q, b1, b2), q, b1, b2); + let mut out = vec![0i16; p]; + for (i, o) in out.iter_mut().enumerate() { + let vi = modq::freeze(i32::from(v[p - i]), q, b1, b2); + *o = modq::product(scale, vi, q, b1, b2); + } + // The divstep state is derived from the secret input — wipe it before returning. + wipe(&mut f); + wipe(&mut g); + wipe(&mut v); + wipe(&mut r); + out +} + +/// Top-coefficient elimination form (the pre-divstep algorithm): the non-x86 +/// and force-scalar path, and the differential oracle for the divstep port. +#[allow(clippy::cast_possible_wrap)] +fn reciprocal3_eliminate(s: &[i8], params: &SntrupParameters) -> Vec { let p = params.p; let q = params.q; let b1 = params.barrett1; let b2 = params.barrett2; let loops = 2 * p + 1; + // Buffers are padded to a multiple of this path's SIMD block (16 i16 lanes) so the + // vector kernels never fall into their scalar tail loops inside the hot iteration. + // Padding lanes start at zero; in u/v they provably stay zero (their source lanes + // are zero too), and in f/g anything written above index p only ever propagates + // upward — index p and below, the only entries ever read, are untouched by it. + let pad = |len: usize| (len + 15) & !15; + let mut r = vec![0i16; p]; - let mut f = vec![0i16; p + 1]; + let mut f = vec![0i16; pad(p + 1)]; f[0] = -1; f[1] = -1; f[p] = 1; - let mut g = vec![0i16; p + 1]; + let mut g = vec![0i16; pad(p + 1)]; for i in 0..p { g[i] = (3 * s[i]) as i16; } + let fg_len = f.len(); let mut d = p as isize; let mut e = p as isize; - let mut u = vec![0i16; loops + 1]; - let mut v = vec![0i16; loops + 1]; + let mut u = vec![0i16; pad(loops + 1)]; + let mut v = vec![0i16; pad(loops + 1)]; + let uv_cap = u.len(); v[0] = 1; - for _ in 0..loops { + for i in 0..loops { let c = modq::quotient(g[p], f[p], q, b1, b2); - vector::minus_product_shift(&mut g, p + 1, &f, c, q, b1, b2); - vector::minus_product_shift(&mut v, loops + 1, &u, c, q, b1, b2); + // The swap mask needs the *post-shift* leading coefficient, so compute that + // single element scalar-first: new g[p] = freeze(g[p-1] - f[p-1]·c). This + // lets the shift and the conditional swap run as one fused memory pass. + let new_gp = modq::minus_product(g[p - 1], f[p - 1], c, q, b1, b2); e -= 1; - let m = smaller_mask(e, d) & modq::mask_set(g[p]); + let m = smaller_mask(e, d) & modq::mask_set(new_gp); let (e_tmp, d_tmp) = swap_int(e, d, m); e = e_tmp; d = d_tmp; - vector::swap(&mut f, &mut g, p + 1, m); - vector::swap(&mut u, &mut v, loops + 1, m); + // After iteration i, the support of u and v is confined to indices 0..=i+1 + // (v starts as {0}, each shift grows it by one, and swaps only exchange the + // two): entries past that window are zero and stay zero, so processing + // `i + 2` elements computes exactly what the full-length pass would. The + // bound depends only on the public loop counter, never on secret data, so + // the constant-time property is unchanged. + let uv_len = pad(i + 2).min(uv_cap); + vector::minus_product_shift_cswap(&mut g, &mut f, fg_len, c, m, params); + vector::minus_product_shift_cswap(&mut v, &mut u, uv_len, c, m, params); } vector::product( &mut r, @@ -49,6 +193,11 @@ pub fn reciprocal3(s: &[i8], params: &SntrupParameters) -> Vec { b1, b2, ); + // The Euclidean state is derived from the secret input — wipe it before returning. + wipe(&mut f); + wipe(&mut g); + wipe(&mut u); + wipe(&mut v); // Note: unlike r3::reciprocal, no invertibility check is returned here. // For these parameter sets q is prime and x^p - x - 1 is irreducible mod q, // so R/q is a field and the weight-w secret f is always invertible — the @@ -67,14 +216,29 @@ pub fn round3(h: &mut [i16], params: &SntrupParameters) { #[allow(unsafe_code)] pub fn mult(h: &mut [i16], f: &[i16], g: &[i8], params: &SntrupParameters) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 availability verified by cfg target_feature - unsafe { - return mult_avx2(h, f, g, params); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + { + // The NTT machine (Good 3x512, primes 7681/10753) covers products up to + // 1536 coefficients, so it serves p = 761; larger sets need the factor-5 + // variant and keep the schoolbook kernels for now. + if params.p == 761 && crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return ntt::mult761(h, f, g); + } + } + if crate::cpu::has_avxvnni() { + // SAFETY: AVX2 + AVX-VNNI support confirmed by has_avxvnni() + unsafe { + return mult_avxvnni(h, f, g, params); + } + } + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return mult_avx2(h, f, g, params); + } + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -111,180 +275,633 @@ fn mult_scalar(h: &mut [i16], f: &[i16], g: &[i8], params: &SntrupParameters) { fg[i - p + 1] = modq::freeze(fg[i - p + 1] as i32 + fg[i] as i32, q, b1, b2); } h[..p].copy_from_slice(&fg[..p]); + // At least one operand is secret at every call site — wipe the product scratch. + wipe(&mut fg); } -/// Column-major schoolbook multiplication with AVX2. -/// Processes 8 i32 multiply-accumulates per SIMD instruction. -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] -#[target_feature(enable = "avx2")] +/// Row-major schoolbook multiplication for x86_64, expanded once per instruction level +/// by the macro below: `mult_avx2` accumulates with `vpmaddwd` + `vpaddd` via +/// `crate::simd::mac_madd`, and `mult_avxvnni` uses the fused `vpdpwssd` via +/// `crate::simd::mac_vnni` (one instruction and one dependency fewer per 16 MACs). +/// +/// Same structure as the NEON kernel below (see its doc comment): contiguous dot products over +/// `f` and a pre-reversed `g`, four independent accumulators, one store per output +/// coefficient, and the `x^p ≡ x + 1` fold applied to raw sums with vectorized freezes. +/// +/// The widening multiply-accumulate here is `_mm256_madd_epi16` (pmaddwd): i16×i16 products +/// pairwise-summed into i32 lanes — 16 multiply-accumulates per instruction, double the old +/// column-major kernel's 8-lane `_mm256_mullo_epi32` density. A dot product only needs the +/// total, so pmaddwd's pairwise fold loses nothing. (An earlier note in this crate claimed +/// AVX2 has no i16-widening MAC — wrong: pmaddwd is exactly that for dot-product shapes.) +macro_rules! rq_mult_x86_kernel { + ($name:ident, $features:literal, $mac:path) => { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + #[target_feature(enable = $features)] + #[allow( + unsafe_code, + clippy::cast_possible_truncation, + clippy::cast_possible_wrap, + clippy::needless_range_loop + )] + unsafe fn $name(h: &mut [i16], f: &[i16], g: &[i8], params: &SntrupParameters) { + unsafe { + use core::arch::x86_64::*; + + let p = params.p; + let q = params.q; + let b1 = params.barrett1; + let b2 = params.barrett2; + let fg_len = p * 2 - 1; + + let mut g_rev = vec![0i16; p]; + for i in 0..p { + g_rev[i] = g[p - 1 - i] as i16; + } + + // Raw i32 convolution sums, padded with one zero so the fold below may read + // `fg32[k + p]` unconditionally at `k = p - 1`. + let mut fg32 = vec![0i32; fg_len + 1]; + for (i, out) in fg32[..fg_len].iter_mut().enumerate() { + let jlo = i.saturating_sub(p - 1); + let len = i.min(p - 1) - jlo + 1; + let fp = f.as_ptr().add(jlo); + // `p - 1 + jlo` never drops below `i` (jlo = max(0, i-p+1)); the naive + // `p - 1 - i + jlo` ordering underflows in debug builds when i ≥ p. + let gp = g_rev.as_ptr().add(p - 1 + jlo - i); + + let mut acc0 = _mm256_setzero_si256(); + let mut acc1 = _mm256_setzero_si256(); + let mut acc2 = _mm256_setzero_si256(); + let mut acc3 = _mm256_setzero_si256(); + let mut k = 0usize; + while k + 64 <= len { + acc0 = $mac( + acc0, + _mm256_loadu_si256(fp.add(k) as *const __m256i), + _mm256_loadu_si256(gp.add(k) as *const __m256i), + ); + acc1 = $mac( + acc1, + _mm256_loadu_si256(fp.add(k + 16) as *const __m256i), + _mm256_loadu_si256(gp.add(k + 16) as *const __m256i), + ); + acc2 = $mac( + acc2, + _mm256_loadu_si256(fp.add(k + 32) as *const __m256i), + _mm256_loadu_si256(gp.add(k + 32) as *const __m256i), + ); + acc3 = $mac( + acc3, + _mm256_loadu_si256(fp.add(k + 48) as *const __m256i), + _mm256_loadu_si256(gp.add(k + 48) as *const __m256i), + ); + k += 64; + } + while k + 16 <= len { + acc0 = $mac( + acc0, + _mm256_loadu_si256(fp.add(k) as *const __m256i), + _mm256_loadu_si256(gp.add(k) as *const __m256i), + ); + k += 16; + } + let s = _mm256_add_epi32( + _mm256_add_epi32(acc0, acc1), + _mm256_add_epi32(acc2, acc3), + ); + let s4 = + _mm_add_epi32(_mm256_castsi256_si128(s), _mm256_extracti128_si256(s, 1)); + let s2 = _mm_add_epi32(s4, _mm_shuffle_epi32(s4, 0b0000_1110)); + let s1 = _mm_add_epi32(s2, _mm_shuffle_epi32(s2, 0b0000_0001)); + let mut sum = _mm_cvtsi128_si32(s1); + while k < len { + sum += *fp.add(k) as i32 * *gp.add(k) as i32; + k += 1; + } + *out = sum; + } + + let qv = _mm256_set1_epi32(q); + let kb1 = _mm256_set1_epi32(b1); + let kb2 = _mm256_set1_epi32(b2); + let k134m = _mm256_set1_epi32(134_217_728); + // Strict-canonical correction bound (see modq::freeze). + let hqv = _mm256_set1_epi32((q - 1) >> 1); + let nhqv = _mm256_set1_epi32(-((q - 1) >> 1)); + + // Vectorized in-place Barrett freeze of the raw sums (the pad entry stays zero). + let mut i = 0usize; + while i + 8 <= fg_len { + let a = _mm256_loadu_si256(fg32.as_ptr().add(i) as *const __m256i); + let t = _mm256_srai_epi32(_mm256_mullo_epi32(a, kb1), 20); + let b = _mm256_sub_epi32(a, _mm256_mullo_epi32(t, qv)); + let t = + _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b, kb2), k134m), 28); + let r = _mm256_sub_epi32(b, _mm256_mullo_epi32(t, qv)); + let r = _mm256_sub_epi32(r, _mm256_and_si256(_mm256_cmpgt_epi32(r, hqv), qv)); + let r = _mm256_add_epi32(r, _mm256_and_si256(_mm256_cmpgt_epi32(nhqv, r), qv)); + _mm256_storeu_si256(fg32.as_mut_ptr().add(i) as *mut __m256i, r); + i += 8; + } + while i < fg_len { + fg32[i] = modq::freeze(fg32[i], q, b1, b2) as i32; + i += 1; + } + + // Fold x^p ≡ x + 1 and freeze once more: fg32[i] (i ≥ p) contributes to outputs i-p + // and i-p+1, so h[k] = freeze(fg32[k] + fg32[k+p] + fg32[k+p-1]) for k ≥ 1, and + // h[0] = freeze(fg32[0] + fg32[p]). + h[0] = modq::freeze(fg32[0] + fg32[p], q, b1, b2); + let mut k = 1usize; + while k + 8 <= p { + let a = _mm256_add_epi32( + _mm256_add_epi32( + _mm256_loadu_si256(fg32.as_ptr().add(k) as *const __m256i), + _mm256_loadu_si256(fg32.as_ptr().add(k + p) as *const __m256i), + ), + _mm256_loadu_si256(fg32.as_ptr().add(k + p - 1) as *const __m256i), + ); + let t = _mm256_srai_epi32(_mm256_mullo_epi32(a, kb1), 20); + let b = _mm256_sub_epi32(a, _mm256_mullo_epi32(t, qv)); + let t = + _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b, kb2), k134m), 28); + let r = _mm256_sub_epi32(b, _mm256_mullo_epi32(t, qv)); + let r = _mm256_sub_epi32(r, _mm256_and_si256(_mm256_cmpgt_epi32(r, hqv), qv)); + let r = _mm256_add_epi32(r, _mm256_and_si256(_mm256_cmpgt_epi32(nhqv, r), qv)); + let packed = + _mm_packs_epi32(_mm256_castsi256_si128(r), _mm256_extracti128_si256(r, 1)); + _mm_storeu_si128(h.as_mut_ptr().add(k) as *mut __m128i, packed); + k += 8; + } + while k < p { + h[k] = modq::freeze(fg32[k] + fg32[k + p] + fg32[k + p - 1], q, b1, b2); + k += 1; + } + + // At least one operand is secret at every call site — wipe the + // reversed copy and the product scratch. + wipe(&mut g_rev); + wipe(&mut fg32); + } + } + }; +} + +rq_mult_x86_kernel!(mult_avx2, "avx2", crate::simd::mac_madd); +rq_mult_x86_kernel!(mult_avxvnni, "avx2,avxvnni", crate::simd::mac_vnni); + +/// Row-major schoolbook multiplication with NEON. +/// +/// Each output coefficient's convolution sum is a contiguous dot product held in registers and +/// written to memory once. `g` is pre-reversed (`g_rev[k] = g[p-1-k]`) so row `i`'s terms +/// `f[j]·g[i-j]` become `f[jlo+t]·g_rev[g0+t]` — both operands advance together. +/// +/// The inner loop uses `vmlal_s16`/`vmlal_high_s16` (i16×i16→i32 widening multiply-accumulate, +/// 8 lanes per instruction pair) across EIGHT independent accumulators, mirroring the structure +/// clang's autovectorizer emits for PQClean's reference C (`smlal`/`smlal2` over `v0`–`v7`): +/// with multi-cycle MAC latency and multiple SIMD pipes, ~8 independent chains are needed to +/// keep the pipes full. An earlier 2-chain attempt (each chain serially feeding its own low and +/// high halves) was latency-bound and only tied the previous column-major kernel; 8 chains is +/// what makes row-major win. Row-major also eliminates the column-major kernel's +/// store-to-load-forwarding hazard, where each `j` iteration reloaded accumulator vectors +/// partially overlapping the previous iteration's stores at a one-element offset. +/// +/// Raw i32 row sums are kept unreduced (bounded by `p·(q-1)/2 < 5.1M`, inside the Barrett +/// window), frozen once in a vectorized pass, then the `x^p ≡ x + 1` fold adds three +/// already-frozen values (|sum| ≤ 3·(q-1)/2) with one final vectorized freeze — two freezes +/// per output coefficient where the reference implementation spends three, and no serial +/// dependency anywhere: `fg[i]` for `i ≥ p` is never itself a fold target, so every fold lane +/// is independent. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] #[allow( unsafe_code, clippy::cast_possible_truncation, clippy::cast_possible_wrap, clippy::needless_range_loop )] -unsafe fn mult_avx2(h: &mut [i16], f: &[i16], g: &[i8], params: &SntrupParameters) { +unsafe fn mult_neon(h: &mut [i16], f: &[i16], g: &[i8], params: &SntrupParameters) { unsafe { - use core::arch::x86_64::*; + use core::arch::aarch64::*; let p = params.p; let q = params.q; let b1 = params.barrett1; let b2 = params.barrett2; - - // Pad to multiples of 8 so SIMD loops need no remainder handling - let g_pad_len = (p + 7) & !7; - let fg_pad_len = p + g_pad_len; let fg_len = p * 2 - 1; - let mut g_pad = vec![0i8; g_pad_len]; - g_pad[..p].copy_from_slice(&g[..p]); - let mut fg = vec![0i32; fg_pad_len]; + let mut g_rev = vec![0i16; p]; + for i in 0..p { + g_rev[i] = g[p - 1 - i] as i16; + } + + // Raw i32 convolution sums, padded with one zero so the fold below may read + // `fg32[k + p]` unconditionally at `k = p - 1`. + let mut fg32 = vec![0i32; fg_len + 1]; + for (i, out) in fg32[..fg_len].iter_mut().enumerate() { + let jlo = i.saturating_sub(p - 1); + let len = i.min(p - 1) - jlo + 1; + let fp = f.as_ptr().add(jlo); + // `p - 1 + jlo` never drops below `i` (jlo = max(0, i-p+1)); the naive + // `p - 1 - i + jlo` ordering underflows in debug builds when i ≥ p. + let gp = g_rev.as_ptr().add(p - 1 + jlo - i); - // Accumulate f[j]*g[k] into fg[j+k] - for j in 0..p { - let fj = _mm256_set1_epi32(f[j] as i32); + let mut acc0 = vdupq_n_s32(0); + let mut acc1 = vdupq_n_s32(0); + let mut acc2 = vdupq_n_s32(0); + let mut acc3 = vdupq_n_s32(0); + let mut acc4 = vdupq_n_s32(0); + let mut acc5 = vdupq_n_s32(0); + let mut acc6 = vdupq_n_s32(0); + let mut acc7 = vdupq_n_s32(0); let mut k = 0usize; - while k < g_pad_len { - let gb = _mm_loadl_epi64(g_pad.as_ptr().add(k) as *const __m128i); - let gk = _mm256_cvtepi8_epi32(gb); - let prod = _mm256_mullo_epi32(fj, gk); - let acc = _mm256_loadu_si256(fg.as_ptr().add(j + k) as *const __m256i); - _mm256_storeu_si256( - fg.as_mut_ptr().add(j + k) as *mut __m256i, - _mm256_add_epi32(acc, prod), - ); + while k + 32 <= len { + let f0 = vld1q_s16(fp.add(k)); + let f1 = vld1q_s16(fp.add(k + 8)); + let f2 = vld1q_s16(fp.add(k + 16)); + let f3 = vld1q_s16(fp.add(k + 24)); + let g0 = vld1q_s16(gp.add(k)); + let g1 = vld1q_s16(gp.add(k + 8)); + let g2 = vld1q_s16(gp.add(k + 16)); + let g3 = vld1q_s16(gp.add(k + 24)); + acc0 = vmlal_s16(acc0, vget_low_s16(f0), vget_low_s16(g0)); + acc1 = vmlal_high_s16(acc1, f0, g0); + acc2 = vmlal_s16(acc2, vget_low_s16(f1), vget_low_s16(g1)); + acc3 = vmlal_high_s16(acc3, f1, g1); + acc4 = vmlal_s16(acc4, vget_low_s16(f2), vget_low_s16(g2)); + acc5 = vmlal_high_s16(acc5, f2, g2); + acc6 = vmlal_s16(acc6, vget_low_s16(f3), vget_low_s16(g3)); + acc7 = vmlal_high_s16(acc7, f3, g3); + k += 32; + } + while k + 8 <= len { + let f0 = vld1q_s16(fp.add(k)); + let g0 = vld1q_s16(gp.add(k)); + acc0 = vmlal_s16(acc0, vget_low_s16(f0), vget_low_s16(g0)); + acc1 = vmlal_high_s16(acc1, f0, g0); k += 8; } + let total = vaddq_s32( + vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3)), + vaddq_s32(vaddq_s32(acc4, acc5), vaddq_s32(acc6, acc7)), + ); + let mut sum = vaddvq_s32(total); + while k < len { + sum += *fp.add(k) as i32 * *gp.add(k) as i32; + k += 1; + } + *out = sum; } - // Vectorized Barrett freeze: i32 -> i16 - let qv = _mm256_set1_epi32(q); - let kb1 = _mm256_set1_epi32(b1); - let kb2 = _mm256_set1_epi32(b2); - let k134m = _mm256_set1_epi32(134_217_728); + let qv = vdupq_n_s32(q); + let kb1 = vdupq_n_s32(b1); + let kb2 = vdupq_n_s32(b2); + let k134m = vdupq_n_s32(134_217_728); + // Strict-canonical correction bound (see modq::freeze). + let hqv = vdupq_n_s32((q - 1) >> 1); + let nhqv = vdupq_n_s32(-((q - 1) >> 1)); - let mut fg16 = vec![0i16; fg_len]; + // Vectorized in-place Barrett freeze of the raw sums (the pad entry stays zero). let mut i = 0usize; - while i + 16 <= fg_len { - let a0 = _mm256_loadu_si256(fg.as_ptr().add(i) as *const __m256i); - let a1 = _mm256_loadu_si256(fg.as_ptr().add(i + 8) as *const __m256i); - - // freeze(a) = a - Q*((b1*a)>>20) then b - Q*((b2*b+134M)>>28) - let t = _mm256_srai_epi32(_mm256_mullo_epi32(a0, kb1), 20); - let b0 = _mm256_sub_epi32(a0, _mm256_mullo_epi32(t, qv)); - let t = _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b0, kb2), k134m), 28); - let r0 = _mm256_sub_epi32(b0, _mm256_mullo_epi32(t, qv)); - - let t = _mm256_srai_epi32(_mm256_mullo_epi32(a1, kb1), 20); - let b1v = _mm256_sub_epi32(a1, _mm256_mullo_epi32(t, qv)); - let t = _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b1v, kb2), k134m), 28); - let r1 = _mm256_sub_epi32(b1v, _mm256_mullo_epi32(t, qv)); - - // Pack 8+8 i32 -> 16 i16 and fix AVX2 lane ordering - let packed = _mm256_permute4x64_epi64(_mm256_packs_epi32(r0, r1), 0xD8); - _mm256_storeu_si256(fg16.as_mut_ptr().add(i) as *mut __m256i, packed); - i += 16; + while i + 4 <= fg_len { + let a = vld1q_s32(fg32.as_ptr().add(i)); + let t = vshrq_n_s32(vmulq_s32(a, kb1), 20); + let b = vsubq_s32(a, vmulq_s32(t, qv)); + let t = vshrq_n_s32(vaddq_s32(vmulq_s32(b, kb2), k134m), 28); + let r = vsubq_s32(b, vmulq_s32(t, qv)); + let r = vsubq_s32(r, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(r, hqv)), qv)); + let r = vaddq_s32(r, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(nhqv, r)), qv)); + vst1q_s32(fg32.as_mut_ptr().add(i), r); + i += 4; } while i < fg_len { - fg16[i] = modq::freeze(fg[i], q, b1, b2); + fg32[i] = modq::freeze(fg32[i], q, b1, b2) as i32; i += 1; } - // Reduction (scalar -- sequential dependencies prevent vectorization) - for i in (p..(p * 2) - 1).rev() { - fg16[i - p] = modq::freeze(fg16[i - p] as i32 + fg16[i] as i32, q, b1, b2); - fg16[i - p + 1] = modq::freeze(fg16[i - p + 1] as i32 + fg16[i] as i32, q, b1, b2); + // Fold x^p ≡ x + 1 and freeze once more: fg32[i] (i ≥ p) contributes to outputs i-p + // and i-p+1, so h[k] = freeze(fg32[k] + fg32[k+p] + fg32[k+p-1]) for k ≥ 1, and + // h[0] = freeze(fg32[0] + fg32[p]). + h[0] = modq::freeze(fg32[0] + fg32[p], q, b1, b2); + let mut k = 1usize; + while k + 4 <= p { + let a = vaddq_s32( + vaddq_s32( + vld1q_s32(fg32.as_ptr().add(k)), + vld1q_s32(fg32.as_ptr().add(k + p)), + ), + vld1q_s32(fg32.as_ptr().add(k + p - 1)), + ); + let t = vshrq_n_s32(vmulq_s32(a, kb1), 20); + let b = vsubq_s32(a, vmulq_s32(t, qv)); + let t = vshrq_n_s32(vaddq_s32(vmulq_s32(b, kb2), k134m), 28); + let r = vsubq_s32(b, vmulq_s32(t, qv)); + let r = vsubq_s32(r, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(r, hqv)), qv)); + let r = vaddq_s32(r, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(nhqv, r)), qv)); + vst1_s16(h.as_mut_ptr().add(k), vmovn_s32(r)); + k += 4; } - h[..p].copy_from_slice(&fg16[..p]); + while k < p { + h[k] = modq::freeze(fg32[k] + fg32[k + p] + fg32[k + p - 1], q, b1, b2); + k += 1; + } + + // At least one operand is secret at every call site — wipe the reversed copy + // and the product scratch. + wipe(&mut g_rev); + wipe(&mut fg32); } } -/// Column-major schoolbook multiplication with NEON. -/// Processes 4 i32 multiply-accumulates per SIMD instruction. -#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] +/// `out[i] = R3(freeze_q(3 · cf[i]))` — the scale-by-3 and lift-to-R3 step of +/// decapsulation (liboqs splits this into `crypto_core_scale3` and +/// `crypto_encode_pxfreeze3`). +/// +/// `cf` holds canonical mod-q coefficients, so `3·cf` stays within ±3(q−1)/2 +/// and both reductions run comfortably inside i32 lanes. +#[allow(unsafe_code)] +pub fn scale3_freeze3(out: &mut [i8], cf: &[i16], params: &SntrupParameters) { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return scale3_freeze3_avx2(out, cf, params); + } + } + #[allow(unreachable_code)] + scale3_freeze3_scalar(out, cf, params); +} + +fn scale3_freeze3_scalar(out: &mut [i8], cf: &[i16], params: &SntrupParameters) { + for (o, &c) in out.iter_mut().zip(cf.iter()) { + let scaled = modq::freeze(3 * i32::from(c), params.q, params.barrett1, params.barrett2); + *o = crate::r3::mod3::freeze(i32::from(scaled)); + } +} + +/// AVX2 form: 32 coefficients per iteration, via a threshold rather than the +/// literal composition the scalar path computes. +/// +/// Writing `s = 3c - kq` for the `k` that lands `s` in the centered range, +/// every parameter set has `q == 1 (mod 3)`, so `s == -k (mod 3)` and the +/// ternary result depends only on `k` — the value of `c mod 3` cannot reach the +/// output at all. And `|3c| <= 1.5(q - 1)` bounds `k` to `{-1, 0, 1}`, so `k` +/// is just the sign of `c` outside a dead zone. The whole two-Barrett, +/// two-freeze composition is therefore a pair of comparisons, which +/// `scale3_collapses_to_a_threshold_on_c` verifies exhaustively over every +/// representable `c` for all six parameter sets. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx2")] +#[allow(unsafe_code, clippy::cast_possible_truncation)] +unsafe fn scale3_freeze3_avx2(out: &mut [i8], cf: &[i16], params: &SntrupParameters) { + unsafe { + use core::arch::x86_64::*; + + let q = params.q; + // Smallest `c` with `k == 1`: `3c` must reach `q/2`, so `t = ceil((q + 1) / 6)`. + let t = (q + 6) / 6; + // `cmpgt` only tests strict `>`, so compare against the neighbours. + let hi = _mm256_set1_epi16((t - 1) as i16); + let lo = _mm256_set1_epi16((1 - t) as i16); + + let n = out.len().min(cf.len()); + let mut i = 0usize; + while i + 32 <= n { + let c0 = _mm256_loadu_si256(cf.as_ptr().add(i) as *const __m256i); + let c1 = _mm256_loadu_si256(cf.as_ptr().add(i + 16) as *const __m256i); + // Each mask is all-ones (-1) when set, so the difference is -1 above + // the band, +1 below it and 0 inside. + let v0 = _mm256_sub_epi16(_mm256_cmpgt_epi16(c0, hi), _mm256_cmpgt_epi16(lo, c0)); + let v1 = _mm256_sub_epi16(_mm256_cmpgt_epi16(c1, hi), _mm256_cmpgt_epi16(lo, c1)); + // `packs` interleaves the two 128-bit halves; the permute undoes it. + let packed = _mm256_permute4x64_epi64::<0b1101_1000>(_mm256_packs_epi16(v0, v1)); + _mm256_storeu_si256(out.as_mut_ptr().add(i) as *mut __m256i, packed); + i += 32; + } + while i < n { + let c = i32::from(cf[i]); + out[i] = i8::from(c <= -t) - i8::from(c >= t); + i += 1; + } + } +} + +#[cfg(test)] #[allow( - unsafe_code, clippy::cast_possible_truncation, clippy::cast_possible_wrap, - clippy::needless_range_loop + clippy::cast_sign_loss )] -unsafe fn mult_neon(h: &mut [i16], f: &[i16], g: &[i8], params: &SntrupParameters) { - unsafe { - use core::arch::aarch64::*; +mod tests { + use super::*; + use crate::params::SntrupParams; - let p = params.p; - let q = params.q; - let b1 = params.barrett1; - let b2 = params.barrett2; + /// Deterministic xorshift64* so the test needs no RNG crates or features. + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + } - // Pad to multiples of 4 so SIMD loops need no remainder handling - let g_pad_len = (p + 3) & !3; - let fg_pad_len = p + g_pad_len; - let fg_len = p * 2 - 1; + fn random_case(params: &SntrupParameters, seed: u64) -> (Vec, Vec) { + let mut s = seed | 1; + let q12 = params.q12; + let f = (0..params.p) + .map(|_| ((next(&mut s) % (2 * q12 as u64 + 1)) as i32 - q12) as i16) + .collect(); + let g = (0..params.p) + .map(|_| ((next(&mut s) % 3) as i8) - 1) + .collect(); + (f, g) + } - let mut g_pad = vec![0i8; g_pad_len]; - g_pad[..p].copy_from_slice(&g[..p]); - let mut fg = vec![0i32; fg_pad_len]; + /// The divstep port must produce byte-identical output to the + /// top-coefficient-elimination oracle for every parameter set: same + /// canonical reciprocal of 3·s, including the ternary edge patterns. + #[cfg(all( + any(target_arch = "x86_64", target_arch = "aarch64"), + not(feature = "force-scalar") + ))] + #[test] + fn divstep_reciprocal_matches_eliminate() { + #[cfg(target_arch = "x86_64")] + if !crate::cpu::has_avx2() { + return; + } + for params in all_params() { + let p = params.p; + for seed in 0..8u64 { + let (_, g) = random_case(params, seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let want = reciprocal3_eliminate(&g, params); + let got = reciprocal3_divstep(&g, params); + assert_eq!(got, want, "divstep vs eliminate: p={p} seed={seed}"); + } + // Deterministic patterns: monomial and alternating ternary. + let mut mono = vec![0i8; p]; + mono[0] = 1; + assert_eq!( + reciprocal3_divstep(&mono, params), + reciprocal3_eliminate(&mono, params), + "monomial p={p}" + ); + let alt: Vec = (0..p).map(|i| [1i8, -1, 0][i % 3]).collect(); + assert_eq!( + reciprocal3_divstep(&alt, params), + reciprocal3_eliminate(&alt, params), + "alternating p={p}" + ); + } + } - // Accumulate f[j]*g[k] into fg[j+k] - for j in 0..p { - let fj = vdupq_n_s32(f[j] as i32); - let mut k = 0usize; - while k + 4 <= g_pad_len { - // Sign-extend 4 i8 -> i16 -> i32 - let gb = vld1_s8(g_pad.as_ptr().add(k)); - let g16 = vmovl_s8(gb); - let gk = vmovl_s16(vget_low_s16(g16)); - let prod = vmulq_s32(fj, gk); - let acc = vld1q_s32(fg.as_ptr().add(j + k)); - vst1q_s32(fg.as_mut_ptr().add(j + k), vaddq_s32(acc, prod)); - k += 4; + /// The NTT multiply must agree with the schoolbook kernel exactly on + /// p = 761 — random operands plus extreme (±(q−1)/2 × ±1) inputs. + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + #[test] + fn ntt_mult_matches_scalar() { + if !crate::cpu::has_avx2() { + return; + } + let params = crate::params::Sntrup761Params::params(); + let p = params.p; + for seed in 0..8u64 { + let (f, g) = random_case(params, seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let mut want = vec![0i16; p]; + mult_scalar(&mut want, &f, &g, params); + let mut got = vec![0i16; p]; + // SAFETY: AVX2 confirmed above. + unsafe { ntt::mult761(&mut got, &f, &g) }; + assert_eq!(got, want, "ntt vs scalar: random seed={seed}"); + } + let hq = params.q12 as i16; + for &(fv, gv) in &[(hq, 1i8), (-hq, 1), (hq, -1), (-hq, -1)] { + let f = vec![fv; p]; + let g = vec![gv; p]; + let mut want = vec![0i16; p]; + mult_scalar(&mut want, &f, &g, params); + let mut got = vec![0i16; p]; + unsafe { ntt::mult761(&mut got, &f, &g) }; + assert_eq!(got, want, "ntt vs scalar: extreme f={fv} g={gv}"); + } + } + + /// The vectorized scale-by-3/lift-to-R3 step must match the scalar form + /// exactly on every parameter set, including the canonical extremes. + #[test] + fn scale3_freeze3_matches_scalar() { + for params in all_params() { + let p = params.p; + let hq = params.q12 as i16; + for seed in 0..6u64 { + let (f, _) = random_case(params, seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let mut want = vec![0i8; p]; + scale3_freeze3_scalar(&mut want, &f, params); + let mut got = vec![0i8; p]; + scale3_freeze3(&mut got, &f, params); + assert_eq!(got, want, "scale3 random p={p} seed={seed}"); + } + for &v in &[0i16, 1, -1, hq, -hq] { + let f = vec![v; p]; + let mut want = vec![0i8; p]; + scale3_freeze3_scalar(&mut want, &f, params); + let mut got = vec![0i8; p]; + scale3_freeze3(&mut got, &f, params); + assert_eq!(got, want, "scale3 const {v} p={p}"); } } + } - // Vectorized Barrett freeze: i32 -> i16 - let qv = vdupq_n_s32(q); - let kb1 = vdupq_n_s32(b1); - let kb2 = vdupq_n_s32(b2); - let k134m = vdupq_n_s32(134_217_728); + #[test] + fn scale3_collapses_to_a_threshold_on_c() { + for params in all_params() { + let q = params.q; + let half = (q - 1) / 2; + let t = (q + 6) / 6; + for c in -half..=half { + let mut got = [0i8; 1]; + scale3_freeze3_scalar(&mut got, &[c as i16], params); + let want = if c >= t { + -1i8 + } else if c <= -t { + 1 + } else { + 0 + }; + assert_eq!(got[0], want, "q={q} c={c}"); + } + } + } - let mut fg16 = vec![0i16; fg_len]; - let mut i = 0usize; - while i + 8 <= fg_len { - // Process 8 values: two batches of 4 i32 -> 8 i16 - let a0 = vld1q_s32(fg.as_ptr().add(i)); - let a1 = vld1q_s32(fg.as_ptr().add(i + 4)); - - let t = vshrq_n_s32(vmulq_s32(a0, kb1), 20); - let b0 = vsubq_s32(a0, vmulq_s32(t, qv)); - let t = vshrq_n_s32(vaddq_s32(vmulq_s32(b0, kb2), k134m), 28); - let r0 = vsubq_s32(b0, vmulq_s32(t, qv)); - - let t = vshrq_n_s32(vmulq_s32(a1, kb1), 20); - let b1v = vsubq_s32(a1, vmulq_s32(t, qv)); - let t = vshrq_n_s32(vaddq_s32(vmulq_s32(b1v, kb2), k134m), 28); - let r1 = vsubq_s32(b1v, vmulq_s32(t, qv)); - - // Pack 4+4 i32 -> 8 i16 (naturally ordered, no permute needed) - let packed = vcombine_s16(vmovn_s32(r0), vmovn_s32(r1)); - vst1q_s16(fg16.as_mut_ptr().add(i), packed); - i += 8; + fn all_params() -> [&'static SntrupParameters; 6] { + [ + crate::params::Sntrup653Params::params(), + crate::params::Sntrup761Params::params(), + crate::params::Sntrup857Params::params(), + crate::params::Sntrup953Params::params(), + crate::params::Sntrup1013Params::params(), + crate::params::Sntrup1277Params::params(), + ] + } + + /// Compare every compiled-in SIMD kernel against the scalar reference. Catches the class + /// of bug the KAT/roundtrip suite can miss when run with `--all-features`, which enables + /// `force-scalar` and silently compiles the SIMD kernels out of the test entirely. + fn check_case(params: &SntrupParameters, f: &[i16], g: &[i8], label: &str) { + let p = params.p; + let mut want = vec![0i16; p]; + mult_scalar(&mut want, f, g, params); + + #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] + { + let mut got = vec![0i16; p]; + // SAFETY: NEON is baseline on aarch64 + unsafe { mult_neon(&mut got, f, g, params) }; + assert_eq!(got, want, "mult_neon vs scalar: {label} p={p}"); } - while i < fg_len { - fg16[i] = modq::freeze(fg[i], q, b1, b2); - i += 1; + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + let mut got = vec![0i16; p]; + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { mult_avx2(&mut got, f, g, params) }; + assert_eq!(got, want, "mult_avx2 vs scalar: {label} p={p}"); + } + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avxvnni() { + let mut got = vec![0i16; p]; + // SAFETY: AVX2 + AVX-VNNI support confirmed by has_avxvnni() + unsafe { mult_avxvnni(&mut got, f, g, params) }; + assert_eq!(got, want, "mult_avxvnni vs scalar: {label} p={p}"); } - // Reduction (scalar -- sequential dependencies prevent vectorization) - for i in (p..(p * 2) - 1).rev() { - fg16[i - p] = modq::freeze(fg16[i - p] as i32 + fg16[i] as i32, q, b1, b2); - fg16[i - p + 1] = modq::freeze(fg16[i - p + 1] as i32 + fg16[i] as i32, q, b1, b2); + let mut got = vec![0i16; p]; + mult(&mut got, f, g, params); + assert_eq!(got, want, "dispatched mult vs scalar: {label} p={p}"); + } + + #[test] + fn simd_mult_matches_scalar_random() { + for params in all_params() { + for seed in 1..=8u64 { + let (f, g) = random_case(params, seed.wrapping_mul(0x9E37_79B9_7F4A_7C15)); + check_case(params, &f, &g, "random"); + } + } + } + + /// Extremes: `f` saturated at ±(q-1)/2 and `g` all ±1 maximize accumulator magnitude, + /// probing the overflow headroom the SIMD kernels' Barrett-freeze staging depends on. + #[test] + fn simd_mult_extremes_match_scalar() { + for params in all_params() { + let q12 = params.q12 as i16; + let f_max: Vec = (0..params.p).map(|_| q12).collect(); + let f_alt: Vec = (0..params.p) + .map(|i| if i % 2 == 0 { q12 } else { -q12 }) + .collect(); + let g_ones = vec![1i8; params.p]; + let g_neg = vec![-1i8; params.p]; + check_case(params, &f_max, &g_ones, "f=+max g=+1"); + check_case(params, &f_max, &g_neg, "f=+max g=-1"); + check_case(params, &f_alt, &g_ones, "f=alt g=+1"); } - h[..p].copy_from_slice(&fg16[..p]); } } diff --git a/sntrup-kem/src/rq/codec761.rs b/sntrup-kem/src/rq/codec761.rs new file mode 100644 index 0000000..84a6db6 --- /dev/null +++ b/sntrup-kem/src/rq/codec761.rs @@ -0,0 +1,854 @@ +//! Generated variable-radix codecs for p = 761, ported from the reference's +//! auto-generated `crypto_decode_761x1531` and friends. +//! +//! Profiling liboqs directly (2026-08-05) showed its codecs run 9-16x faster +//! than our generic implementation — `crypto_decode_761x1531` in 0.13 us +//! against our 2.14 us — because they are vectorized at *every* radix level, +//! while ours only vectorizes levels whose moduli happen to be uniform. +//! +//! These are mechanical translations of generated C. Like `rq::ntt`, they are +//! p = 761 specific (the radix chain depends on p and q); the generic codec in +//! the parent module remains the path for every other parameter set, and is the +//! differential oracle in tests. +// Machine-translated from generated C: the upstream naming (R0, A0, S1) and +// parenthesisation are preserved deliberately so the port can be diffed against +// its source. Regenerate rather than hand-edit. +// The sign-losing and truncating casts below are not accidents: they reproduce +// C's `int16`/`uint16` wrapping semantics, which the radix decomposition relies +// on. Silencing them module-wide is deliberate for this generated translation +// and must not be copied into hand-written code. +#![allow( + unsafe_code, + unused_parens, + unused_assignments, + non_snake_case, + clippy::all, + clippy::cast_sign_loss, + clippy::cast_possible_truncation, + clippy::cast_possible_wrap +)] + +use core::arch::x86_64::*; + +/// C `int16` truncation: assignments to an `int16` wrap. +#[inline(always)] +fn trunc(x: i32) -> i32 { + x as i16 as i32 +} + +#[inline(always)] +fn mullo(x: i32, y: i32) -> i32 { + (x as i16).wrapping_mul(y as i16) as i32 +} + +#[inline(always)] +fn mulhi(x: i32, y: i32) -> i32 { + (((x as i16 as i32) * (y as i16 as i32)) >> 16) as i16 as i32 +} + +#[inline] +#[target_feature(enable = "avx2")] +fn add(x: __m256i, y: __m256i) -> __m256i { + _mm256_add_epi16(x, y) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn sub(x: __m256i, y: __m256i) -> __m256i { + _mm256_sub_epi16(x, y) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn mulloconst(x: __m256i, y: i16) -> __m256i { + _mm256_mullo_epi16(x, _mm256_set1_epi16(y)) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn mulhiconst(x: __m256i, y: i16) -> __m256i { + _mm256_mulhi_epi16(x, _mm256_set1_epi16(y)) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn shiftleftconst(x: __m256i) -> __m256i { + _mm256_slli_epi16::(x) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn subconst(x: __m256i, y: i16) -> __m256i { + sub(x, _mm256_set1_epi16(y)) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn ifgesubconst(x: __m256i, y: i16) -> __m256i { + let y16 = _mm256_set1_epi16(y); + let top16 = _mm256_set1_epi16(y.wrapping_sub(1)); + sub(x, _mm256_and_si256(_mm256_cmpgt_epi16(x, top16), y16)) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn ifnegaddconst(x: __m256i, y: i16) -> __m256i { + add( + x, + _mm256_and_si256(_mm256_srai_epi16::<15>(x), _mm256_set1_epi16(y)), + ) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn ldu(s: &[u8], at: isize) -> __m256i { + unsafe { _mm256_loadu_si256(s.as_ptr().offset(at) as *const __m256i) } +} + +#[inline] +#[target_feature(enable = "avx2")] +fn ldu8(s: &[u8], at: isize) -> __m256i { + unsafe { _mm256_cvtepu8_epi16(_mm_loadu_si128(s.as_ptr().offset(at) as *const __m128i)) } +} + +#[inline] +#[target_feature(enable = "avx2")] +fn ldr(r: &[i16], at: isize) -> __m256i { + unsafe { _mm256_loadu_si256(r.as_ptr().offset(at) as *const __m256i) } +} + +#[inline] +#[target_feature(enable = "avx2")] +fn str_(r: &mut [i16], at: isize, v: __m256i) { + unsafe { _mm256_storeu_si256(r.as_mut_ptr().offset(at) as *mut __m256i, v) } +} + +/// Decode the rounded (761 x 1531) ciphertext representation into `R0[..761]`. +#[target_feature(enable = "avx2")] +pub fn decode_761x1531(r0: &mut [i16], s: &[u8]) { + let mut R0 = &mut *r0; + let mut R1 = [0i16; 381]; + let mut R2 = [0i16; 191]; + let mut R3 = [0i16; 96]; + let mut R4 = [0i16; 48]; + let mut R5 = [0i16; 24]; + let mut R6 = [0i16; 12]; + let mut R7 = [0i16; 6]; + let mut R8 = [0i16; 3]; + let mut R9 = [0i16; 2]; + let mut R10 = [0i16; 1]; + let (mut A0, mut A1, mut A2): (__m256i, __m256i, __m256i); + let (mut S0, mut S1): (__m256i, __m256i); + let (mut B0, mut B1, mut C0, mut C1): (__m256i, __m256i, __m256i, __m256i); + let mut i: isize; + let (mut a0, mut a1, mut a2): (i32, i32, i32) = (0, 0, 0); + + let mut si: isize = 1007; + a1 = 0; + a1 += { + si -= 1; + s[si as usize] as i32 + }; + a1 = mulhi(a1, -84) - mulhi(mullo(a1, -4828), 3475); + a1 += { + si -= 1; + s[si as usize] as i32 + }; + a1 += (a1 >> 15) & 3475; + R10[0] = a1 as i16; + + // R10 ------> R9: reconstruct mod 1*[593]+[1500] + + i = 0; + si -= 1; + a0 = R10[0] as i32; + a2 = a0; + a0 = mulhi(a0, 60) - mulhi(mullo(a0, -28292), 593); + a0 += (s[(si + (1 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 593; + a1 = trunc((a2 << 8) + (s[(si + (i)) as usize] as i32) - a0); + a1 = mullo(a1, -31055); + + // invalid inputs might need reduction mod 1500 + a1 -= 1500; + a1 += (a1 >> 15) & 1500; + + R9[0] = a0 as i16; + R9[1] = a1 as i16; + si -= 0; + + // R9 ------> R8: reconstruct mod 2*[6232]+[1500] + + R8[2] = R9[1]; + si -= 2; + i = 0; + while i >= 0 { + a0 = R9[(i) as usize] as i32; + a2 = a0; + a0 = mulhi(a0, 672) - mulhi(mullo(a0, -2692), 6232); + a0 += (s[(si + (2 * i + 1)) as usize] as i32); + a0 = mulhi(a0, 672) - mulhi(mullo(a0, -2692), 6232); + a0 += (s[(si + (2 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 6232; + a1 = trunc( + (a2 << 13) + + ((s[(si + (2 * i + 1)) as usize] as i32) << 5) + + (((s[(si + (2 * i)) as usize] as i32) - a0) >> 3), + ); + a1 = mullo(a1, 12451); + + // invalid inputs might need reduction mod 6232 + a1 -= 6232; + a1 += (a1 >> 15) & 6232; + + R8[(2 * i) as usize] = a0 as i16; + R8[(2 * i + 1) as usize] = a1 as i16; + i -= 1; + } + + // R8 ------> R7: reconstruct mod 5*[1263]+[304] + + i = 0; + si -= 1; + a0 = R8[2] as i32; + a2 = a0; + a0 = mulhi(a0, -476) - mulhi(mullo(a0, -13284), 1263); + a0 += (s[(si + (1 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 1263; + a1 = trunc((a2 << 8) + (s[(si + (i)) as usize] as i32) - a0); + a1 = mullo(a1, -22001); + + // invalid inputs might need reduction mod 304 + a1 -= 304; + a1 += (a1 >> 15) & 304; + + R7[4] = a0 as i16; + R7[5] = a1 as i16; + si -= 2; + i = 1; + while i >= 0 { + a0 = R8[(i) as usize] as i32; + a2 = a0; + a0 = mulhi(a0, -476) - mulhi(mullo(a0, -13284), 1263); + a0 += (s[(si + (1 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 1263; + a1 = trunc((a2 << 8) + (s[(si + (i)) as usize] as i32) - a0); + a1 = mullo(a1, -22001); + + // invalid inputs might need reduction mod 1263 + a1 -= 1263; + a1 += (a1 >> 15) & 1263; + + R7[(2 * i) as usize] = a0 as i16; + R7[(2 * i + 1) as usize] = a1 as i16; + i -= 1; + } + + // R7 ------> R6: reconstruct mod 11*[9097]+[2188] + + i = 0; + si -= 2; + a0 = R7[5] as i32; + a0 = mulhi(a0, 2348) - mulhi(mullo(a0, -1844), 9097); + a0 += (s[(si + (2 * i + 1)) as usize] as i32); + a0 = mulhi(a0, 2348) - mulhi(mullo(a0, -1844), 9097); + a0 += (s[(si + (2 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 9097; + a1 = trunc( + ((s[(si + (2 * i + 1)) as usize] as i32) << 8) + (s[(si + (2 * i)) as usize] as i32) - a0, + ); + a1 = mullo(a1, 17081); + + // invalid inputs might need reduction mod 2188 + a1 -= 2188; + a1 += (a1 >> 15) & 2188; + + R6[10] = a0 as i16; + R6[11] = a1 as i16; + si -= 10; + i = 4; + while i >= 0 { + a0 = R7[(i) as usize] as i32; + a0 = mulhi(a0, 2348) - mulhi(mullo(a0, -1844), 9097); + a0 += (s[(si + (2 * i + 1)) as usize] as i32); + a0 = mulhi(a0, 2348) - mulhi(mullo(a0, -1844), 9097); + a0 += (s[(si + (2 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 9097; + a1 = trunc( + ((s[(si + (2 * i + 1)) as usize] as i32) << 8) + (s[(si + (2 * i)) as usize] as i32) + - a0, + ); + a1 = mullo(a1, 17081); + + // invalid inputs might need reduction mod 9097 + a1 -= 9097; + a1 += (a1 >> 15) & 9097; + + R6[(2 * i) as usize] = a0 as i16; + R6[(2 * i + 1) as usize] = a1 as i16; + i -= 1; + } + + // R6 ------> R5: reconstruct mod 23*[1526]+[367] + + i = 0; + si -= 1; + a0 = R6[11] as i32; + a2 = a0; + a0 = mulhi(a0, 372) - mulhi(mullo(a0, -10994), 1526); + a0 += (s[(si + (1 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 1526; + a1 = trunc((a2 << 7) + (((s[(si + (i)) as usize] as i32) - a0) >> 1)); + a1 = mullo(a1, -18381); + + // invalid inputs might need reduction mod 367 + a1 -= 367; + a1 += (a1 >> 15) & 367; + + R5[22] = a0 as i16; + R5[23] = a1 as i16; + si -= 11; + i = 10; + while i >= 0 { + a0 = R6[(i) as usize] as i32; + a2 = a0; + a0 = mulhi(a0, 372) - mulhi(mullo(a0, -10994), 1526); + a0 += (s[(si + (1 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 1526; + a1 = trunc((a2 << 7) + (((s[(si + (i)) as usize] as i32) - a0) >> 1)); + a1 = mullo(a1, -18381); + + // invalid inputs might need reduction mod 1526 + a1 -= 1526; + a1 += (a1 >> 15) & 1526; + + R5[(2 * i) as usize] = a0 as i16; + R5[(2 * i + 1) as usize] = a1 as i16; + i -= 1; + } + + // R5 ------> R4: reconstruct mod 47*[625]+[150] + + i = 0; + si -= 1; + a0 = R5[23] as i32; + a2 = a0; + a0 = mulhi(a0, -284) - mulhi(mullo(a0, -26844), 625); + a0 += (s[(si + (1 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 625; + a1 = trunc((a2 << 8) + (s[(si + (i)) as usize] as i32) - a0); + a1 = mullo(a1, 32401); + + // invalid inputs might need reduction mod 150 + a1 -= 150; + a1 += (a1 >> 15) & 150; + + R4[46] = a0 as i16; + R4[47] = a1 as i16; + si -= 23; + i = 7; + loop { + A0 = ldr(&R5, (i)); + A2 = A0; + S0 = ldu8(s, si + (i)); + A0 = sub( + mulhiconst(A0, -284), + mulhiconst(mulloconst(A0, -26844), 625), + ); + A0 = add(A0, S0); + A0 = ifnegaddconst(A0, 625); + A1 = add(shiftleftconst::<8>(A2), sub(S0, A0)); + A1 = mulloconst(A1, 32401); + + // invalid inputs might need reduction mod 625 + A1 = ifgesubconst(A1, 625); + + // A0: r0r2r4r6r8r10r12r14 r16r18r20r22r24r26r28r30 + // A1: r1r3r5r7r9r11r13r15 r17r19r21r23r25r27r29r31 + B0 = _mm256_unpacklo_epi16(A0, A1); + B1 = _mm256_unpackhi_epi16(A0, A1); + // B0: r0r1r2r3r4r5r6r7 r16r17r18r19r20r21r22r23 + // B1: r8r9r10r11r12r13r14r15 r24r25r26r27r28r29r30r31 + C0 = _mm256_permute2x128_si256(B0, B1, 0x20); + C1 = _mm256_permute2x128_si256(B0, B1, 0x31); + // C0: r0r1r2r3r4r5r6r7 r8r9r10r11r12r13r14r15 + // C1: r16r17r18r19r20r21r22r23 r24r25r26r27r28r29r30r31 + str_(&mut R4, (2 * i), C0); + str_(&mut R4, (2 * i) + 16, C1); + if i == 0 { + break; + } + i = -16 - ((!15) & -i); + } + + // R4 ------> R3: reconstruct mod 95*[6400]+[1531] + + i = 0; + si -= 2; + a0 = R4[47] as i32; + a2 = a0; + a0 = mulhi(a0, 2816) - mulhi(mullo(a0, -2621), 6400); + a0 += (s[(si + (2 * i + 1)) as usize] as i32); + a0 = mulhi(a0, 2816) - mulhi(mullo(a0, -2621), 6400); + a0 += (s[(si + (2 * i + 0)) as usize] as i32); + a0 += (a0 >> 15) & 6400; + a1 = trunc( + (a2 << 8) + + (s[(si + (2 * i + 1)) as usize] as i32) + + (((s[(si + (2 * i)) as usize] as i32) - a0) >> 8), + ); + a1 = mullo(a1, 23593); + + // invalid inputs might need reduction mod 1531 + a1 -= 1531; + a1 += (a1 >> 15) & 1531; + + R3[94] = a0 as i16; + R3[95] = a1 as i16; + si -= 94; + i = 31; + loop { + A0 = ldr(&R4, (i)); + A2 = A0; + S0 = ldu(s, si + (2 * i)); + S1 = _mm256_srli_epi16::<8>(S0); + S0 = _mm256_and_si256(S0, _mm256_set1_epi16(255)); + A0 = sub( + mulhiconst(A0, 2816), + mulhiconst(mulloconst(A0, -2621), 6400), + ); + A0 = add(A0, S1); + A0 = sub( + mulhiconst(A0, 2816), + mulhiconst(mulloconst(A0, -2621), 6400), + ); + A0 = add(A0, S0); + A0 = ifnegaddconst(A0, 6400); + A1 = add( + add(shiftleftconst::<8>(A2), S1), + _mm256_srai_epi16::<8>(sub(S0, A0)), + ); + A1 = mulloconst(A1, 23593); + + // invalid inputs might need reduction mod 6400 + A1 = ifgesubconst(A1, 6400); + + // A0: r0r2r4r6r8r10r12r14 r16r18r20r22r24r26r28r30 + // A1: r1r3r5r7r9r11r13r15 r17r19r21r23r25r27r29r31 + B0 = _mm256_unpacklo_epi16(A0, A1); + B1 = _mm256_unpackhi_epi16(A0, A1); + // B0: r0r1r2r3r4r5r6r7 r16r17r18r19r20r21r22r23 + // B1: r8r9r10r11r12r13r14r15 r24r25r26r27r28r29r30r31 + C0 = _mm256_permute2x128_si256(B0, B1, 0x20); + C1 = _mm256_permute2x128_si256(B0, B1, 0x31); + // C0: r0r1r2r3r4r5r6r7 r8r9r10r11r12r13r14r15 + // C1: r16r17r18r19r20r21r22r23 r24r25r26r27r28r29r30r31 + str_(&mut R3, (2 * i), C0); + str_(&mut R3, (2 * i) + 16, C1); + if i == 0 { + break; + } + i = -16 - ((!15) & -i); + } + + // R3 ------> R2: reconstruct mod 190*[1280]+[1531] + + R2[190] = R3[95]; + si -= 95; + i = 79; + loop { + A0 = ldr(&R3, (i)); + A2 = A0; + S0 = ldu8(s, si + (i)); + A0 = sub( + mulhiconst(A0, 256), + mulhiconst(mulloconst(A0, -13107), 1280), + ); + A0 = add(A0, S0); + A0 = ifnegaddconst(A0, 1280); + A1 = add(A2, _mm256_srai_epi16::<8>(sub(S0, A0))); + A1 = mulloconst(A1, -13107); + + // invalid inputs might need reduction mod 1280 + A1 = ifgesubconst(A1, 1280); + + // A0: r0r2r4r6r8r10r12r14 r16r18r20r22r24r26r28r30 + // A1: r1r3r5r7r9r11r13r15 r17r19r21r23r25r27r29r31 + B0 = _mm256_unpacklo_epi16(A0, A1); + B1 = _mm256_unpackhi_epi16(A0, A1); + // B0: r0r1r2r3r4r5r6r7 r16r17r18r19r20r21r22r23 + // B1: r8r9r10r11r12r13r14r15 r24r25r26r27r28r29r30r31 + C0 = _mm256_permute2x128_si256(B0, B1, 0x20); + C1 = _mm256_permute2x128_si256(B0, B1, 0x31); + // C0: r0r1r2r3r4r5r6r7 r8r9r10r11r12r13r14r15 + // C1: r16r17r18r19r20r21r22r23 r24r25r26r27r28r29r30r31 + str_(&mut R2, (2 * i), C0); + str_(&mut R2, (2 * i) + 16, C1); + if i == 0 { + break; + } + i = -16 - ((!15) & -i); + } + + // R2 ------> R1: reconstruct mod 380*[9157]+[1531] + + R1[380] = R2[190]; + si -= 380; + i = 174; + loop { + A0 = ldr(&R2, (i)); + S0 = ldu(s, si + (2 * i)); + S1 = _mm256_srli_epi16::<8>(S0); + S0 = _mm256_and_si256(S0, _mm256_set1_epi16(255)); + A0 = sub( + mulhiconst(A0, 1592), + mulhiconst(mulloconst(A0, -1832), 9157), + ); + A0 = add(A0, S1); + A0 = sub( + mulhiconst(A0, 1592), + mulhiconst(mulloconst(A0, -1832), 9157), + ); + A0 = add(A0, S0); + A0 = ifnegaddconst(A0, 9157); + A1 = add(shiftleftconst::<8>(S1), sub(S0, A0)); + A1 = mulloconst(A1, 25357); + + // invalid inputs might need reduction mod 9157 + A1 = ifgesubconst(A1, 9157); + + // A0: r0r2r4r6r8r10r12r14 r16r18r20r22r24r26r28r30 + // A1: r1r3r5r7r9r11r13r15 r17r19r21r23r25r27r29r31 + B0 = _mm256_unpacklo_epi16(A0, A1); + B1 = _mm256_unpackhi_epi16(A0, A1); + // B0: r0r1r2r3r4r5r6r7 r16r17r18r19r20r21r22r23 + // B1: r8r9r10r11r12r13r14r15 r24r25r26r27r28r29r30r31 + C0 = _mm256_permute2x128_si256(B0, B1, 0x20); + C1 = _mm256_permute2x128_si256(B0, B1, 0x31); + // C0: r0r1r2r3r4r5r6r7 r8r9r10r11r12r13r14r15 + // C1: r16r17r18r19r20r21r22r23 r24r25r26r27r28r29r30r31 + str_(&mut R1, (2 * i), C0); + str_(&mut R1, (2 * i) + 16, C1); + if i == 0 { + break; + } + i = -16 - ((!15) & -i); + } + + // R1 ------> R0: reconstruct mod 761*[1531] + + R0[760] = trunc(3 * R1[380] as i32 - 2295) as i16; + si -= 380; + i = 364; + loop { + A0 = ldr(&R1, (i)); + A2 = A0; + S0 = ldu8(s, si + (i)); + A0 = sub( + mulhiconst(A0, 518), + mulhiconst(mulloconst(A0, -10958), 1531), + ); + A0 = add(A0, S0); + A0 = ifnegaddconst(A0, 1531); + A1 = add(shiftleftconst::<8>(A2), sub(S0, A0)); + A1 = mulloconst(A1, 15667); + + // invalid inputs might need reduction mod 1531 + A1 = ifgesubconst(A1, 1531); + + A0 = mulloconst(A0, 3); + A1 = mulloconst(A1, 3); + A0 = subconst(A0, 2295); + A1 = subconst(A1, 2295); + // A0: r0r2r4r6r8r10r12r14 r16r18r20r22r24r26r28r30 + // A1: r1r3r5r7r9r11r13r15 r17r19r21r23r25r27r29r31 + B0 = _mm256_unpacklo_epi16(A0, A1); + B1 = _mm256_unpackhi_epi16(A0, A1); + // B0: r0r1r2r3r4r5r6r7 r16r17r18r19r20r21r22r23 + // B1: r8r9r10r11r12r13r14r15 r24r25r26r27r28r29r30r31 + C0 = _mm256_permute2x128_si256(B0, B1, 0x20); + C1 = _mm256_permute2x128_si256(B0, B1, 0x31); + // C0: r0r1r2r3r4r5r6r7 r8r9r10r11r12r13r14r15 + // C1: r16r17r18r19r20r21r22r23 r24r25r26r27r28r29r30r31 + str_(&mut R0, (2 * i), C0); + str_(&mut R0, (2 * i) + 16, C1); + if i == 0 { + break; + } + i = -16 - ((!15) & -i); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::params::SNTRUP761; + + /// The ported generated decoder must agree exactly with the generic + /// variable-radix implementation on every input, including the ciphertexts + /// the KEM actually produces and adversarial random byte strings. + #[test] + fn decode_761x1531_matches_generic() { + if !crate::cpu::has_avx2() { + return; + } + let params = &SNTRUP761; + let len = params.rounded_encode_size; + let mut state = 0x2468_ace0_1357_9bdfu64 | 1; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + }; + for trial in 0..64 { + let s: Vec = (0..len) + .map(|_| match trial % 4 { + 0 => (next() & 0xff) as u8, + 1 => 0, + 2 => 0xff, + _ => (next() % 3) as u8, + }) + .collect(); + let mut want = vec![0i16; params.p]; + crate::rq::encoding::rounded_decode_into(&s, &mut want, params); + let mut got = vec![0i16; params.p]; + // SAFETY: AVX2 confirmed above. + unsafe { decode_761x1531(&mut got, &s) }; + assert_eq!(got, want, "decode_761x1531 mismatch, trial {trial}"); + } + } +} + +// --------------------------------------------------------------------------- +// crypto_encode_761x1531round +// --------------------------------------------------------------------------- +// +// The reference emits this as six near-identical loops over two shapes, walking +// `reading`/`writing`/`out` pointers with a mid-loop back-off on the final +// iteration (the same overlapping-tail trick the decoder uses). Both shapes are +// factored into helpers here and instantiated per radix level; the back-off is +// passed in rather than duplicated. + +/// Shuffle that splits each 32-bit lane into "two bytes out, two bytes carried". +#[inline] +#[target_feature(enable = "avx2")] +fn shuf_carry2() -> __m256i { + _mm256_set_epi8( + 15, 14, 11, 10, 7, 6, 3, 2, 13, 12, 9, 8, 5, 4, 1, 0, 15, 14, 11, 10, 7, 6, 3, 2, 13, 12, + 9, 8, 5, 4, 1, 0, + ) +} + +/// Shuffle that splits each 32-bit lane into "one byte out, two bytes carried". +#[inline] +#[target_feature(enable = "avx2")] +fn shuf_carry1() -> __m256i { + _mm256_set_epi8( + 12, 8, 4, 0, 12, 8, 4, 0, 14, 13, 10, 9, 6, 5, 2, 1, 12, 8, 4, 0, 12, 8, 4, 0, 14, 13, 10, + 9, 6, 5, 2, 1, + ) +} + +/// Shape A: 16 u16 in -> 8 u16 carried + 8 bytes emitted, per iteration. +#[target_feature(enable = "avx2")] +unsafe fn enc_pass_a( + rd_base: *const u16, + wr_base: *mut u16, + out: *mut u8, + iters: usize, + m: i32, + back: (usize, usize, usize), + round: bool, +) -> usize { + unsafe { + let (mut rd, mut wr, mut o) = (0isize, 0isize, 0isize); + let mut i = iters; + while i > 0 { + i -= 1; + if i == 0 { + rd -= back.0 as isize; + wr -= back.1 as isize; + o -= back.2 as isize; + } + let mut x = _mm256_loadu_si256(rd_base.offset(rd) as *const __m256i); + if round { + // round-to-nearest-multiple-of-3, recentre, then divide by 3 + x = _mm256_mulhrs_epi16(x, _mm256_set1_epi16(10923)); + x = _mm256_add_epi16(x, _mm256_add_epi16(x, x)); + x = _mm256_add_epi16(x, _mm256_set1_epi16(2295)); + x = _mm256_and_si256(x, _mm256_set1_epi16(16383)); + x = _mm256_mulhi_epi16(x, _mm256_set1_epi16(21846)); + } + let y = _mm256_and_si256(x, _mm256_set1_epi32(65535)); + let hi = _mm256_mullo_epi32(_mm256_srli_epi32::<16>(x), _mm256_set1_epi32(m)); + let x = _mm256_permute4x64_epi64::<0xd8>(_mm256_shuffle_epi8( + _mm256_add_epi32(y, hi), + shuf_carry1(), + )); + _mm_storeu_si128( + wr_base.offset(wr) as *mut __m128i, + _mm256_castsi256_si128(x), + ); + let mut s0 = _mm256_extract_epi32::<4>(x) as u32; + for k in 0..4 { + *out.offset(o + k) = s0 as u8; + s0 >>= 8; + } + let mut s0 = _mm256_extract_epi32::<6>(x) as u32; + for k in 0..4 { + *out.offset(o + 4 + k) = s0 as u8; + s0 >>= 8; + } + rd += 16; + wr += 8; + o += 8; + } + o as usize + } +} + +/// Shape B: 32 u16 in -> 16 u16 carried + 32 bytes emitted, per iteration. +#[target_feature(enable = "avx2")] +unsafe fn enc_pass_b( + rd_base: *const u16, + wr_base: *mut u16, + out: *mut u8, + iters: usize, + m: i32, + back: (usize, usize, usize), +) -> usize { + unsafe { + let (mut rd, mut wr, mut o) = (0isize, 0isize, 0isize); + let mut i = iters; + while i > 0 { + i -= 1; + if i == 0 { + rd -= back.0 as isize; + wr -= back.1 as isize; + o -= back.2 as isize; + } + let comb = |v: __m256i| -> __m256i { + let y = _mm256_and_si256(v, _mm256_set1_epi32(65535)); + let hi = _mm256_mullo_epi32(_mm256_srli_epi32::<16>(v), _mm256_set1_epi32(m)); + _mm256_permute4x64_epi64::<0xd8>(_mm256_shuffle_epi8( + _mm256_add_epi32(y, hi), + shuf_carry2(), + )) + }; + let x = comb(_mm256_loadu_si256(rd_base.offset(rd) as *const __m256i)); + let x2 = comb(_mm256_loadu_si256(rd_base.offset(rd + 16) as *const __m256i)); + _mm256_storeu_si256( + wr_base.offset(wr) as *mut __m256i, + _mm256_permute2f128_si256::<0x31>(x, x2), + ); + _mm256_storeu_si256( + out.offset(o) as *mut __m256i, + _mm256_permute2f128_si256::<0x20>(x, x2), + ); + rd += 32; + wr += 16; + o += 32; + } + o as usize + } +} + +/// Round each coefficient to a multiple of 3 and encode the 761 x 1531 +/// representation into `out[..1007]`. +#[target_feature(enable = "avx2")] +pub fn encode_761x1531round(out: &mut [u8], r0: &[i16]) { + unsafe { + let mut r = [0u16; 381]; + // Level 0 reads the caller's coefficients and writes carries into `r`; + // every later level reads and writes `r` in place. + let rp = r.as_mut_ptr(); + let op = out.as_mut_ptr(); + + // Level 0 reads the caller's coefficients; every later level reads and + // writes `r` in place, exactly as the reference aliases its buffers. + let mut o = enc_pass_a(r0.as_ptr().cast::(), rp, op, 48, 1531, (8, 4, 4), true); + r[380] = ((((3 * ((10923 * i32::from(r0[760]) + 16384) >> 15) + 2295) & 16383) * 10923) + >> 15) as u16; + + o += enc_pass_b(rp, rp, op.add(o), 12, 9157, (4, 2, 4)); + r[190] = r[380]; + o += enc_pass_a(rp, rp, op.add(o), 12, 1280, (2, 1, 1), false); + r[95] = r[190]; + o += enc_pass_b(rp, rp, op.add(o), 3, 6400, (0, 0, 0)); + o += enc_pass_a(rp, rp, op.add(o), 3, 625, (0, 0, 0), false); + o += enc_pass_a(rp, rp, op.add(o), 2, 1526, (8, 4, 4), false); + + // Scalar tail: the last four radix levels are too small to vectorize. + for i in 0..6 { + let r2 = u32::from(r[2 * i]) + u32::from(r[2 * i + 1]) * 9097; + *op.add(o) = r2 as u8; + *op.add(o + 1) = (r2 >> 8) as u8; + o += 2; + r[i] = (r2 >> 16) as u16; + } + for i in 0..3 { + let r2 = u32::from(r[2 * i]) + u32::from(r[2 * i + 1]) * 1263; + *op.add(o) = r2 as u8; + o += 1; + r[i] = (r2 >> 8) as u16; + } + let r2 = u32::from(r[0]) + u32::from(r[1]) * 6232; + *op.add(o) = r2 as u8; + *op.add(o + 1) = (r2 >> 8) as u8; + o += 2; + r[0] = (r2 >> 16) as u16; + r[1] = r[2]; + let r2 = u32::from(r[0]) + u32::from(r[1]) * 593; + *op.add(o) = r2 as u8; + o += 1; + r[0] = (r2 >> 8) as u16; + *op.add(o) = r[0] as u8; + *op.add(o + 1) = (r[0] >> 8) as u8; + } +} + +#[cfg(test)] +mod encode_tests { + use super::*; + use crate::params::SNTRUP761; + + /// The ported encoder must agree byte-for-byte with the generic + /// implementation on the coefficient ranges decapsulation actually produces. + #[test] + fn encode_761x1531round_matches_generic() { + if !crate::cpu::has_avx2() { + return; + } + let params = &SNTRUP761; + let hq = params.q12 as i16; + let mut state = 0x1357_9bdf_2468_ace0u64 | 1; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + }; + for trial in 0..32 { + let f: Vec = (0..params.p) + .map(|_| match trial % 4 { + 0 => ((next() % (2 * hq as u64 + 1)) as i32 - i32::from(hq)) as i16, + 1 => hq, + 2 => -hq, + _ => 0, + }) + .collect(); + // The reference's `..._round` rounds internally; our generic encoder + // expects input that `round3` has already processed. Compare like + // for like by rounding only on the generic side. + let mut rounded = f.clone(); + crate::rq::round3(&mut rounded, params); + let mut want = vec![0u8; params.rounded_encode_size]; + crate::rq::encoding::rounded_encode_into(&rounded, &mut want, params); + let mut got = vec![0u8; params.rounded_encode_size]; + // SAFETY: AVX2 confirmed above. + unsafe { encode_761x1531round(&mut got, &f) }; + assert_eq!(got, want, "encode mismatch, trial {trial}"); + } + } +} diff --git a/sntrup-kem/src/rq/encoding.rs b/sntrup-kem/src/rq/encoding.rs index d1e239a..6536b63 100644 --- a/sntrup-kem/src/rq/encoding.rs +++ b/sntrup-kem/src/rq/encoding.rs @@ -112,20 +112,151 @@ fn encode_single(out: &mut [u8], mut val: u32, mut modulus: u32) -> usize { pos } -/// Iterative variable-radix decoding. Forward pass computes moduli and byte -/// offsets at each level; backward pass expands decoded values from base case. -#[allow(clippy::cast_possible_truncation)] -fn decode(out: &mut [u16], s: &[u8], m_in: &[u16], n_start: usize) { - if n_start == 0 { - return; - } - if n_start == 1 { - decode_single(out, s, m_in[0]); - return; +/// AVX2 expansion of one decode level whose pairs all share a modulus `m` and +/// bottom-byte count `bb` — the shape levels 0-2 always have (see RESULTS.md). +/// +/// Handles pairs `[lo, hi)` backward in blocks of eight 32-bit lanes: +/// `combined = out[i]·256^bb + LE(s[start + i·bb ..][..bb])`, then +/// `out[2i] = combined mod m`, `out[2i+1] = (combined / m) mod m`, replicating +/// `uint32_divmod_uint14`'s two Barrett steps plus speculative correction +/// lanewise. Callers guarantee `lo >= 8`, so a block's writes (`out[2i..2i+16]`) +/// never overlap its reads (`out[i..i+8]`). +/// +/// `bb` is 1 or 2 for every level this is invoked on; other values fall back to +/// the scalar path. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx2")] +#[allow(unsafe_code)] +unsafe fn decode_level_avx2( + out: &mut [u16], + s: &[u8], + m: u16, + bb: usize, + start: usize, + lo: usize, + hi: usize, +) { + unsafe { + use core::arch::x86_64::*; + + let m32 = u32::from(m); + #[allow(clippy::cast_possible_truncation)] + let v = (0x8000_0000u64 / u64::from(m32)) as u32; + let mv = _mm256_set1_epi32(i32::from_ne_bytes(m32.to_ne_bytes())); + let vv = _mm256_set1_epi64x(i64::from(v)); + + // 8-lane u32 multiply-high-by-v, keeping bit 31 upward (the scalar + // code's `(x as u64 * v) >> 31`). AVX2 only multiplies even 32-bit + // lanes, so run the even and odd halves separately and reblend. + let qhat = |x: __m256i| -> __m256i { + let pe = _mm256_srli_epi64::<31>(_mm256_mul_epu32(x, vv)); + let po = _mm256_srli_epi64::<31>(_mm256_mul_epu32(_mm256_srli_epi64::<32>(x), vv)); + _mm256_blend_epi32::<0b1010_1010>(pe, _mm256_slli_epi64::<32>(po)) + }; + + // One Barrett step: returns (quotient_part, remainder). + let step = |x: __m256i| -> (__m256i, __m256i) { + let q = qhat(x); + (q, _mm256_sub_epi32(x, _mm256_mullo_epi32(q, mv))) + }; + + let mut i = hi; + while i >= lo + 8 { + i -= 8; + + let ov = _mm256_cvtepu16_epi32(_mm_loadu_si128(out.as_ptr().add(i) as *const __m128i)); + let bytes = if bb == 1 { + _mm256_cvtepu8_epi32(_mm_loadl_epi64(s.as_ptr().add(start + i) as *const __m128i)) + } else { + _mm256_cvtepu16_epi32(_mm_loadu_si128( + s.as_ptr().add(start + 2 * i) as *const __m128i + )) + }; + let shift = if bb == 1 { 8 } else { 16 }; + let combined = _mm256_add_epi32(_mm256_sllv_epi32(ov, _mm256_set1_epi32(shift)), bytes); + + // divmod: two Barrett steps then the speculative correction. + let (q0, r0) = step(combined); + let (q1, r1) = step(r0); + let q = _mm256_add_epi32(q0, q1); + let rsub = _mm256_sub_epi32(r1, mv); + let mask = _mm256_srai_epi32::<31>(rsub); + let rem = _mm256_add_epi32(rsub, _mm256_and_si256(mask, mv)); + let quo = _mm256_add_epi32(_mm256_add_epi32(q, _mm256_set1_epi32(1)), mask); + + // out[2i+1] = quo mod m (one more reduction; quo < m^2 / ... but a + // single Barrett step plus correction suffices, matching + // uint32_mod_uint14). + let (_, hr0) = step(quo); + let (_, hr1) = step(hr0); + let hsub = _mm256_sub_epi32(hr1, mv); + let hmask = _mm256_srai_epi32::<31>(hsub); + let hi_val = _mm256_add_epi32(hsub, _mm256_and_si256(hmask, mv)); + + // Interleave (rem, hi_val) into out[2i .. 2i+16] as u16. + let lo16 = _mm256_packus_epi32(rem, hi_val); // lanes: r0..3 h0..3 r4..7 h4..7 + let fixed = _mm256_permute4x64_epi64::<0b11_01_10_00>(lo16); // r0..3 r4..7 h0..3 h4..7 + let r16 = _mm256_castsi256_si128(fixed); + let h16 = _mm256_extracti128_si256::<1>(fixed); + _mm_storeu_si128( + out.as_mut_ptr().add(2 * i) as *mut __m128i, + _mm_unpacklo_epi16(r16, h16), + ); + _mm_storeu_si128( + out.as_mut_ptr().add(2 * i + 8) as *mut __m128i, + _mm_unpackhi_epi16(r16, h16), + ); + } } +} - // --- Forward pass: compute level sizes, moduli, and bottom-byte totals --- +/// Per-(modulus, length) decode plan: the moduli tree, per-level bottom-byte +/// counts, and level offsets. +/// +/// These depend only on the starting modulus and `p`, both public and fixed per +/// parameter set, yet recomputing them cost ~0.75 µs on every decode — 1.5 µs of +/// each decapsulation, which runs two. There are exactly twelve live +/// combinations (six parameter sets x {Rq, rounded}), so they are computed once +/// on first use and reused thereafter. +struct DecodePlan { + ns: [usize; MAX_LEVELS], + num_levels: usize, + all_m: [u16; MAX_M_STORAGE], + level_m_offset: [usize; MAX_LEVELS + 1], + level_bottom_total: [usize; MAX_LEVELS], + level_bottom_start: [usize; MAX_LEVELS], + all_bb: [u8; MAX_M_STORAGE], + /// Total bottom bytes across all levels — where the base case is read from. + cum_bottom: usize, +} + +/// Twelve slots: `(parameter set, codec kind)`. Keyed by the pair actually +/// requested, so a miss simply builds and stores its own plan. +static PLANS: [std::sync::OnceLock<(u16, usize, Box)>; 12] = + [const { std::sync::OnceLock::new() }; 12]; + +fn plan_for(m0: u16, n_start: usize) -> &'static DecodePlan { + for slot in &PLANS { + // An occupied slot for a different key is skipped; an empty one is + // claimed for this key. Twelve slots cover every live combination. + if let Some((km, kn, plan)) = slot.get() { + if *km == m0 && *kn == n_start { + return plan; + } + continue; + } + let built = slot.get_or_init(|| (m0, n_start, Box::new(build_plan(m0, n_start)))); + if built.0 == m0 && built.1 == n_start { + return &built.2; + } + } + // Slots exhausted (cannot happen for the supported parameter sets): fall + // back to leaking one plan rather than failing. + Box::leak(Box::new(build_plan(m0, n_start))) +} +#[allow(clippy::cast_possible_truncation)] +fn build_plan(m0: u16, n_start: usize) -> DecodePlan { let mut ns = [0usize; MAX_LEVELS]; let mut num_levels = 0; { @@ -137,13 +268,13 @@ fn decode(out: &mut [u16], s: &[u8], m_in: &[u16], n_start: usize) { } } - // Flat storage for moduli at every level (including paired output for base case) let mut all_m = [0u16; MAX_M_STORAGE]; + let mut all_bb = [0u8; MAX_M_STORAGE]; let mut level_m_offset = [0usize; MAX_LEVELS + 1]; let mut level_bottom_total = [0usize; MAX_LEVELS]; + let mut level_bottom_start = [0usize; MAX_LEVELS]; - // Level 0 input moduli - all_m[..n_start].copy_from_slice(&m_in[..n_start]); + all_m[..n_start].fill(m0); level_m_offset[0] = 0; let mut m_pos = n_start; @@ -153,34 +284,70 @@ fn decode(out: &mut [u16], s: &[u8], m_in: &[u16], n_start: usize) { let m_off = level_m_offset[level]; level_m_offset[level + 1] = m_pos; let mut total_bottom = 0usize; - for i in 0..n2 { if 2 * i + 1 < n { - let mut cm = (all_m[m_off + 2 * i] as u32) * (all_m[m_off + 2 * i + 1] as u32); + let mut cm = u32::from(all_m[m_off + 2 * i]) * u32::from(all_m[m_off + 2 * i + 1]); let mut bb = 0usize; while cm >= 16384 { bb += 1; cm = (cm + 255) >> 8; } total_bottom += bb; + all_bb[m_pos] = bb as u8; all_m[m_pos] = cm as u16; } else { all_m[m_pos] = all_m[m_off + 2 * i]; } m_pos += 1; } - level_bottom_total[level] = total_bottom; } - // Cumulative bottom-byte start positions - let mut level_bottom_start = [0usize; MAX_LEVELS]; - let mut cum_bottom = 0usize; + // Cumulative bottom-byte start positions, level 0 upward. + let mut cum = 0usize; for level in 0..num_levels { - level_bottom_start[level] = cum_bottom; - cum_bottom += level_bottom_total[level]; + level_bottom_start[level] = cum; + cum += level_bottom_total[level]; + } + + DecodePlan { + ns, + num_levels, + all_m, + level_m_offset, + level_bottom_total, + level_bottom_start, + all_bb, + cum_bottom: cum, + } +} + +/// Iterative variable-radix decoding: the moduli tree and byte offsets come +/// from a cached [`DecodePlan`]; the backward pass expands decoded values from +/// the base case. +#[allow(clippy::cast_possible_truncation)] +fn decode(out: &mut [u16], s: &[u8], m_in: &[u16], n_start: usize) { + if n_start == 0 { + return; + } + if n_start == 1 { + decode_single(out, s, m_in[0]); + return; } + // The moduli tree, per-level byte counts and offsets depend only on the + // (public, fixed) starting modulus and length, so they are built once and + // cached rather than recomputed on every call. + let plan = plan_for(m_in[0], n_start); + let ns = &plan.ns; + let num_levels = plan.num_levels; + let all_m = &plan.all_m; + let all_bb = &plan.all_bb; + let level_m_offset = &plan.level_m_offset; + let level_bottom_total = &plan.level_bottom_total; + let level_bottom_start = &plan.level_bottom_start; + let cum_bottom = plan.cum_bottom; + // --- Decode base case (n = 1) --- let base_m_off = level_m_offset[num_levels]; decode_single(out, &s[cum_bottom..], all_m[base_m_off]); @@ -195,16 +362,57 @@ fn decode(out: &mut [u16], s: &[u8], m_in: &[u16], n_start: usize) { // Process backwards: reads from out[i], writes to out[2*i] / out[2*i+1]. let mut bpos = level_bottom_start[level] + level_bottom_total[level]; - for i in (0..n2).rev() { - if 2 * i + 1 < n { - // Recompute bottom-byte count for this pair - let mut cm = (all_m[m_off + 2 * i] as u32) * (all_m[m_off + 2 * i + 1] as u32); - let mut bb = 0usize; - while cm >= 16384 { - bb += 1; - cm = (cm + 255) >> 8; + // Uniform-modulus fast path. The kernel covers the half-open pair range + // `[simd_lo, simd_hi)`; the scalar loop below walks every pair anyway to + // thread `bpos` (whose step varies with each pair's byte count) and + // simply skips the work the kernel already did. + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + let mut simd_lo = n2; + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + { + let n_full = n / 2; + if n_full >= 16 && crate::cpu::has_avx2() { + let m0 = all_m[m_off]; + let bb = all_bb[level_m_offset[level + 1]] as usize; + // Every full pair on this level must share the modulus. A + // prefix-only relaxation is NOT sound here: the kernel's stores + // reach out[2*n_uniform + 14], so any higher pair still needing + // its input would have to run first, and the byte-cursor + // threading then has to be split to match. See RESULTS.md. + let n_uniform = if (0..n_full) + .all(|k| all_m[m_off + 2 * k] == m0 && all_m[m_off + 2 * k + 1] == m0) + { + n_full + } else { + 0 + }; + if n_uniform >= 16 && (bb == 1 || bb == 2) { + let start = level_bottom_start[level]; + // The unpaired tail element is the highest index the scalar + // loop would visit, so it must be copied BEFORE the kernel + // runs — the kernel's stores reach out[2*n_full - 1] and + // would otherwise clobber out[n2 - 1] before it is read. + if n % 2 == 1 { + out[2 * (n2 - 1)] = out[n2 - 1]; + } + // SAFETY: AVX2 confirmed by has_avx2(). `lo = 0` is sound: + // within a block every load precedes every store, and + // blocks descend, so a block's writes (out[2i..2i+16]) + // always sit above any lower block's reads (out[i..i+8]). + unsafe { + decode_level_avx2(out, s, m0, bb, start, 8, n_uniform); + } + simd_lo = 8 + (n_uniform - 8) % 8; + bpos = start + simd_lo * bb; } + } + } + #[cfg(not(all(target_arch = "x86_64", not(feature = "force-scalar"))))] + let simd_lo = n2; + for i in (0..simd_lo).rev() { + if 2 * i + 1 < n { + let bb = all_bb[level_m_offset[level + 1] + i] as usize; bpos -= bb; let mut combined = out[i] as u32; for j in (0..bb).rev() { @@ -248,8 +456,10 @@ pub fn rq_encode(f: &[i16], params: &SntrupParameters) -> Vec { out } +/// Allocation-free Rq decoder: writes into `out[..p]`, using stack +/// scratch bounded by [`crate::params::MAX_P`]. #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] -pub fn rq_decode(c: &[u8], params: &SntrupParameters) -> Vec { +pub fn rq_decode_into(c: &[u8], out: &mut [i16], params: &SntrupParameters) { let p = params.p; let q12 = params.q12; let q_u16 = params.q as u16; @@ -257,8 +467,11 @@ pub fn rq_decode(c: &[u8], params: &SntrupParameters) -> Vec { let b1 = params.barrett1; let b2 = params.barrett2; - let m = vec![q_u16; p]; - let mut r = vec![0u16; p]; + let mut m = [0u16; crate::params::MAX_P]; + m[..p].fill(q_u16); + let m = &m[..p]; + let mut r_buf = [0u16; crate::params::MAX_P]; + let r = &mut r_buf[..p]; // Callers pass exactly `pk_size` bytes, so borrow directly on the hot path. // Only allocate-and-pad if the input is short (defensive; never happens via // the public API, where `EncapsulationKey::try_from` enforces the size). @@ -270,32 +483,63 @@ pub fn rq_decode(c: &[u8], params: &SntrupParameters) -> Vec { padded[..c.len()].copy_from_slice(c); &padded }; - decode(&mut r, s, &m, p); - let mut f = vec![0i16; p]; - for (fi, &ri) in f.iter_mut().zip(r.iter()) { + decode(r, s, m, p); + for (fi, &ri) in out[..p].iter_mut().zip(r.iter()) { *fi = modq::freeze(ri as i32 - q12, q, b1, b2); } - f } +/// Round to multiples of 3 and encode in one step. +/// +/// The reference fuses these (`crypto_encode_761x1531round`), and the ported +/// p = 761 kernel does the rounding internally — so it takes the *un-rounded* +/// coefficients. Other parameter sets round in place first, as before. +pub fn round_and_encode_into(f: &mut [i16], out: &mut [u8], params: &SntrupParameters) { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if params.p == 761 && crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2(). + unsafe { + return super::codec761::encode_761x1531round(out, f); + } + } + super::round3(f, params); + rounded_encode_into(f, out, params); +} + +/// Allocation-free rounded encoder: writes into +/// `out[..rounded_encode_size]`. #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] -pub fn rounded_encode(f: &[i16], params: &SntrupParameters) -> Vec { +pub fn rounded_encode_into(f: &[i16], out: &mut [u8], params: &SntrupParameters) { let p = params.p; let q12 = params.q12; let q_rounded = (params.q as u16).div_ceil(3); - let mut r = vec![0u16; p]; + let mut r_buf = [0u16; crate::params::MAX_P]; + let r = &mut r_buf[..p]; for (ri, &fi) in r.iter_mut().zip(f.iter()) { *ri = (((fi as i32 + q12) * 10923) >> 15) as u16; } - let mut m = vec![q_rounded; p]; - let mut out = vec![0u8; params.rounded_encode_size]; - encode(&mut out, &mut r, &mut m, p); - out + let mut m = [0u16; crate::params::MAX_P]; + m[..p].fill(q_rounded); + encode(&mut out[..params.rounded_encode_size], r, &mut m[..p], p); + // On the decapsulation path `f` is the re-encrypted candidate, secret until (and unless) + // the constant-time ciphertext comparison succeeds — wipe the working representation, + // which `encode` mutates in place across pairing levels. `m` holds only public moduli. + crate::wipe::wipe(&mut r_buf); } +/// Allocation-free rounded decoder: writes into `out[..p]`. #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] -pub fn rounded_decode(c: &[u8], params: &SntrupParameters) -> Vec { +pub fn rounded_decode_into(c: &[u8], out: &mut [i16], params: &SntrupParameters) { + // p = 761 has a ported copy of the reference's generated codec, which is + // vectorized at every radix level rather than only the uniform ones. + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if params.p == 761 && c.len() >= params.rounded_encode_size && crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2(). + unsafe { + return super::codec761::decode_761x1531(out, c); + } + } let p = params.p; let q12 = params.q12; let q_rounded = (params.q as u16).div_ceil(3); @@ -303,8 +547,11 @@ pub fn rounded_decode(c: &[u8], params: &SntrupParameters) -> Vec { let b1 = params.barrett1; let b2 = params.barrett2; - let m = vec![q_rounded; p]; - let mut r = vec![0u16; p]; + let mut m = [0u16; crate::params::MAX_P]; + m[..p].fill(q_rounded); + let m = &m[..p]; + let mut r_buf = [0u16; crate::params::MAX_P]; + let r = &mut r_buf[..p]; // Callers pass exactly `rounded_encode_size` bytes, so borrow directly on the // hot path. Only allocate-and-pad if the input is short (defensive; never // happens via the public API, where `Ciphertext::try_from` enforces the size). @@ -316,10 +563,8 @@ pub fn rounded_decode(c: &[u8], params: &SntrupParameters) -> Vec { padded[..c.len()].copy_from_slice(c); &padded }; - decode(&mut r, s, &m, p); - let mut f = vec![0i16; p]; - for (fi, &ri) in f.iter_mut().zip(r.iter()) { + decode(r, s, m, p); + for (fi, &ri) in out[..p].iter_mut().zip(r.iter()) { *fi = modq::freeze(ri as i32 * 3 - q12, q, b1, b2); } - f } diff --git a/sntrup-kem/src/rq/modq.rs b/sntrup-kem/src/rq/modq.rs index 36c2ea2..c3bf018 100644 --- a/sntrup-kem/src/rq/modq.rs +++ b/sntrup-kem/src/rq/modq.rs @@ -1,12 +1,24 @@ -/// Barrett reduction: freezes `a` into the range (-q/2, q/2). +/// Barrett reduction: freezes `a` into the canonical range [-(q-1)/2, (q-1)/2]. /// /// `barrett1` = floor(2^20 / q), `barrett2` = floor(2^28 / q). +/// +/// The two Barrett steps alone can land up to ±3 outside the canonical range for +/// a few thousand inputs per parameter set (exhaustively scanned over the live +/// |a| ≤ (q-1)/2 + ((q-1)/2)^2 window). A non-canonical coefficient wraps +/// negative in `rq_encode`'s `+ q12` bias and corrupts the variable-radix +/// encoding, so a final branchless correction makes the output strictly +/// canonical — matching the reference implementation's exact freeze. Every SIMD +/// freeze implementation applies the same correction so all paths produce +/// byte-identical results. #[inline(always)] #[allow(clippy::cast_possible_truncation)] pub fn freeze(a: i32, q: i32, barrett1: i32, barrett2: i32) -> i16 { let mut b = a; b -= q * ((barrett1 * b) >> 20); b -= q * ((barrett2 * b + 134_217_728) >> 28); + let hq = (q - 1) >> 1; + b -= q & ((hq - b) >> 31); // b > hq → subtract q + b += q & ((b + hq) >> 31); // b < -hq → add q b as i16 } @@ -15,10 +27,30 @@ pub fn product(a: i16, b: i16, q: i32, b1: i32, b2: i32) -> i16 { freeze(a as i32 * b as i32, q, b1, b2) } +/// The two Barrett steps without the strict-canonical correction: output may land +/// up to ±3 outside ±(q-1)/2, but is always the correct residue and maps a zero +/// residue to literal 0 (exhaustively scanned per parameter set). Used only for +/// intermediate chain values inside [`reciprocal`], where the serial dependency +/// makes the correction's latency expensive and a later strict [`freeze`] +/// canonicalizes anything that escapes. +#[inline(always)] +#[allow(clippy::cast_possible_truncation)] +fn freeze_loose(a: i32, q: i32, barrett1: i32, barrett2: i32) -> i16 { + let mut b = a; + b -= q * ((barrett1 * b) >> 20); + b -= q * ((barrett2 * b + 134_217_728) >> 28); + b as i16 +} + #[inline(always)] -pub fn square(a: i16, q: i32, b1: i32, b2: i32) -> i16 { +fn product_loose(a: i16, b: i16, q: i32, b1: i32, b2: i32) -> i16 { + freeze_loose(a as i32 * b as i32, q, b1, b2) +} + +#[inline(always)] +fn square_loose(a: i16, q: i32, b1: i32, b2: i32) -> i16 { let a32 = a as i32; - freeze(a32 * a32, q, b1, b2) + freeze_loose(a32 * a32, q, b1, b2) } /// Compute `a1^(q-2) mod q` via Fermat's little theorem using binary @@ -31,12 +63,15 @@ pub fn reciprocal(a1: i16, q: i32, b1: i32, b2: i32) -> i16 { // Find the highest set bit position let bits = 32 - exp.leading_zeros(); // number of significant bits - // Square-and-multiply from the second-highest bit down + // Square-and-multiply from the second-highest bit down. Chain values stay in + // the loose (residue-correct, near-canonical) domain; callers that need a + // canonical result apply a strict freeze afterwards (quotient's final + // product does). let mut result = a1; for i in (0..(bits - 1)).rev() { - result = square(result, q, b1, b2); + result = square_loose(result, q, b1, b2); if (exp >> i) & 1 == 1 { - result = product(result, a1, q, b1, b2); + result = product_loose(result, a1, q, b1, b2); } } result @@ -61,3 +96,53 @@ pub fn mask_set(x: i16) -> isize { r >>= 31; r as isize } + +#[cfg(test)] +mod tests { + use super::*; + + /// Strict canonical-range property over the live input window, all six + /// parameter sets. The pre-fix two-step Barrett freeze violated this for a + /// few thousand inputs per set (up to ±3 outside canonical), which wraps + /// negative in `rq_encode`'s bias and corrupts the variable-radix encoding — + /// this test fails against that implementation and pins the fix. + #[test] + fn freeze_is_strictly_canonical_and_residue_correct() { + for &(q, b1, b2) in &[ + (4621i32, 226i32, 58084i32), + (4591, 228, 58464), + (5167, 202, 51948), + (6343, 165, 42324), + (7177, 146, 37410), + (7879, 133, 34073), + ] { + let hq = (q - 1) / 2; + let lim = hq + hq * hq; + // Exhaustive canonical-range check over the full live window — the + // pre-fix freeze's violations are scattered through the interior, so + // sampling would miss them. + let mut a = -lim; + while a <= lim { + let r = i32::from(freeze(a, q, b1, b2)); + assert!( + (-hq..=hq).contains(&r), + "freeze({a}) = {r} outside canonical ±{hq} for q={q}" + ); + a += 1; + } + // Residue correctness on a strided sample (division is the slow part). + let mut a = -lim; + while a <= lim { + let r = i32::from(freeze(a, q, b1, b2)); + assert_eq!((r - a).rem_euclid(q), 0, "wrong residue for q={q} a={a}"); + a += 997; + } + // Every zero-residue input must freeze to literal 0 (mask_set relies on it). + let mut a = -(lim / q) * q; + while a <= lim { + assert_eq!(freeze(a, q, b1, b2), 0, "freeze({a}) != 0 for q={q}"); + a += q; + } + } + } +} diff --git a/sntrup-kem/src/rq/ntt.rs b/sntrup-kem/src/rq/ntt.rs new file mode 100644 index 0000000..05ce565 --- /dev/null +++ b/sntrup-kem/src/rq/ntt.rs @@ -0,0 +1,1525 @@ +//! NTT-based multiplication in R/q (ported from the SUPERCOP AVX2 +//! `crypto_core_multsntrup761` and its generated `_ntt` kernels) — the algorithm +//! liboqs dispatches to on x86_64. +//! +//! Good's trick maps the 768-coefficient operands into 3 interleaved tracks of +//! 512 (3 and 512 are coprime, and 3·512 = 1536 ≥ 2p−1 = 1521), then two +//! 512-point NTTs run over the NTT-friendly primes 7681 and 10753 — 4591 itself +//! is not NTT-friendly, and both primes are ≡ 1 mod 2^10 so 512-point +//! transforms exist. A Karatsuba-shaped 3×3 pointwise stage, inverse NTTs, and +//! CRT recombination bring the result back mod 4591. All arithmetic is 16-bit +//! lanes: signed Montgomery products and `mulhrs` squeezes. +//! +//! **Scope: p = 761 only** (see `super::mult`'s dispatcher). The 3×512 machine +//! holds products up to 1536 coefficients, so p ≤ 768 fits; p ≥ 857 needs the +//! Good factor 5 variant (5·512 = 2560) with a 5×5 pointwise stage, which the +//! reference generates as a separate parameter set. The twiddle tables are +//! prime-specific, not p-specific, so that extension reuses them. +//! +//! `ntt512`/`invntt512` are mechanical translations of the reference's +//! auto-generated kernels, and the twiddle tables are extracted verbatim; both +//! are validated against the schoolbook implementation by differential tests. +#![allow( + unsafe_code, + clippy::cast_possible_truncation, + clippy::too_many_lines, + clippy::needless_range_loop +)] + +use crate::wipe::wipe; +use core::arch::x86_64::*; + +const Q: i16 = 4591; + +#[inline] +#[target_feature(enable = "avx2")] +fn ld(qdata: &[i16; 1696], off: usize) -> __m256i { + unsafe { _mm256_loadu_si256(qdata.as_ptr().add(off) as *const __m256i) } +} + +#[inline] +#[target_feature(enable = "avx2")] +fn add16(a: __m256i, b: __m256i) -> __m256i { + _mm256_add_epi16(a, b) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn sub16(a: __m256i, b: __m256i) -> __m256i { + _mm256_sub_epi16(a, b) +} + +/// `x·y·2^-16 mod± q` with a pre-scaled table pair (y, y·q^-1 mod 2^16). +#[inline] +#[target_feature(enable = "avx2")] +fn mulmod_scaled(x: __m256i, y: __m256i, yqinv: __m256i, qv: __m256i) -> __m256i { + let b = _mm256_mulhi_epi16(x, y); + let d = _mm256_mullo_epi16(x, yqinv); + let e = _mm256_mulhi_epi16(d, qv); + sub16(b, e) +} + +/// The reference's `reduce_x16`: mulhrs-based partial reduction. +#[inline] +#[target_feature(enable = "avx2")] +fn reduce(x: __m256i, qv: __m256i, qrv: __m256i) -> __m256i { + let y = _mm256_mulhrs_epi16(x, qrv); + let y = _mm256_mullo_epi16(y, qv); + sub16(x, y) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn perm_lo(a: __m256i, b: __m256i) -> __m256i { + _mm256_permute2x128_si256::<0x20>(a, b) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn perm_hi(a: __m256i, b: __m256i) -> __m256i { + _mm256_permute2x128_si256::<0x31>(a, b) +} + +/// Squeeze toward `mod± q` via `mulhrs` with `c = round(2^15/q)`-ish constant. +#[inline] +#[target_feature(enable = "avx2")] +fn squeeze(x: __m256i, c: i16, q: i16) -> __m256i { + sub16( + x, + _mm256_mullo_epi16( + _mm256_mulhrs_epi16(x, _mm256_set1_epi16(c)), + _mm256_set1_epi16(q), + ), + ) +} + +/// `x·y·2^-16 mod± q`, deriving `y·q^-1` on the fly (broadcast operands). +#[inline] +#[target_feature(enable = "avx2")] +fn mulmod(x: __m256i, y: __m256i, qinv: i16, q: i16) -> __m256i { + let yqinv = _mm256_mullo_epi16(y, _mm256_set1_epi16(qinv)); + let b = _mm256_mulhi_epi16(x, y); + let d = _mm256_mullo_epi16(x, yqinv); + let e = _mm256_mulhi_epi16(d, _mm256_set1_epi16(q)); + sub16(b, e) +} + +#[inline] +#[target_feature(enable = "avx2")] +fn squeeze_4591(x: __m256i) -> __m256i { + squeeze(x, 7, 4591) +} +#[inline] +#[target_feature(enable = "avx2")] +fn squeeze_7681(x: __m256i) -> __m256i { + squeeze(x, 4, 7681) +} +#[inline] +#[target_feature(enable = "avx2")] +fn squeeze_10753(x: __m256i) -> __m256i { + squeeze(x, 3, 10753) +} +#[inline] +#[target_feature(enable = "avx2")] +fn mulmod_4591(x: __m256i, y: __m256i) -> __m256i { + mulmod(x, y, 15631, 4591) +} +#[inline] +#[target_feature(enable = "avx2")] +fn mulmod_7681(x: __m256i, y: __m256i) -> __m256i { + mulmod(x, y, -7679, 7681) +} +#[inline] +#[target_feature(enable = "avx2")] +fn mulmod_10753(x: __m256i, y: __m256i) -> __m256i { + mulmod(x, y, -10751, 10753) +} + +/// Fully reduce to `[-(q-1)/2, (q-1)/2]` (the reference's `freeze_4591_x16`). +#[inline] +#[target_feature(enable = "avx2")] +fn freeze_4591(x: __m256i) -> __m256i { + let q = _mm256_set1_epi16(Q); + let x = add16(x, _mm256_and_si256(q, _mm256_srai_epi16::<15>(x))); + let m = _mm256_srai_epi16::<15>(sub16(x, _mm256_set1_epi16((Q + 1) / 2))); + _mm256_blendv_epi8(sub16(x, q), x, m) +} + +const MASKS: [[i16; 16]; 3] = [ + [-1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1], + [0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0], + [0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0], +]; + +#[inline] +#[target_feature(enable = "avx2")] +fn mask(t: usize) -> __m256i { + unsafe { _mm256_loadu_si256(MASKS[t].as_ptr() as *const __m256i) } +} + +/// Good's permutation: 768 coefficients into 3 tracks of 512. +/// +/// Coefficient `i` belongs to track `i mod 3`; because a 16-lane block spans +/// indices `16b .. 16b+15` and `16 ≡ 1 (mod 3)`, the three lane masks rotate by +/// one per block: track `t` takes `mask[(t + 3 − b mod 3) mod 3]` of the block +/// at block-index `b`. That single formula reproduces all six unrolled cases in +/// the reference. +#[target_feature(enable = "avx2")] +fn good(fpad: &mut [i16], f: &[i16; 768]) { + unsafe { + // Reindexing the track loop by `u = (t - b) mod 3` makes the mask index + // a constant and moves the block-dependent rotation into the store + // address, which is scalar arithmetic. The three masks then stay in + // registers instead of being reloaded six times per block. The second + // operand's block index is always `b0 + 2 (mod 3)`, so its mask is + // `m[(u + 1) % 3]`. + let m = [mask(0), mask(1), mask(2)]; + let mut j = 0usize; + while j < 512 { + let b0 = (j / 16) % 3; + let f0 = _mm256_loadu_si256(f.as_ptr().add(j) as *const __m256i); + if j < 256 { + let f1 = _mm256_loadu_si256(f.as_ptr().add(512 + j) as *const __m256i); + for u in 0..3 { + let v = _mm256_or_si256( + _mm256_and_si256(f0, m[u]), + _mm256_and_si256(f1, m[(u + 1) % 3]), + ); + let track = (u + b0) % 3; + _mm256_storeu_si256(fpad.as_mut_ptr().add(512 * track + j) as *mut __m256i, v); + } + } else { + for u in 0..3 { + let v = _mm256_and_si256(f0, m[u]); + let track = (u + b0) % 3; + _mm256_storeu_si256(fpad.as_mut_ptr().add(512 * track + j) as *mut __m256i, v); + } + } + j += 16; + } + } +} + +/// Inverse of [`good`]: the 3 tracks back to 1536 coefficients. Output block at +/// block-index `B` gathers `mask[(t + 3 − B mod 3) mod 3]` from track `t` — the +/// transpose of the forward formula. +#[target_feature(enable = "avx2")] +fn ungood(h: &mut [i16; 1536], fpad: &[i16]) { + unsafe { + // Same reindexing as `good`, in the opposite direction: with `u` as the + // loop variable the masks are constants held in registers and the + // block-dependent rotation lands on the load address instead. + let m = [mask(0), mask(1), mask(2)]; + let mut j = 0usize; + while j < 512 { + let jb = j / 16; + for k in 0..3 { + // Block index of output block `k`: `(k * 32 + jb) % 3`, and + // `32 == 2 (mod 3)`. + let b = (2 * k + jb) % 3; + let mut g = _mm256_setzero_si256(); + for u in 0..3 { + let track = _mm256_loadu_si256( + fpad.as_ptr().add(512 * ((u + b) % 3) + j) as *const __m256i + ); + g = _mm256_or_si256(g, _mm256_and_si256(track, m[u])); + } + _mm256_storeu_si256(h.as_mut_ptr().add(k * 512 + j) as *mut __m256i, g); + } + j += 16; + } + } +} + +/// One prime's forward-NTT / pointwise / inverse-NTT pass, shared by the +/// mod-q (dual-prime) and mod-3 (single-prime) multiplies. The pointwise stage +/// is the Karatsuba-shaped 3x3 product over Good's three tracks. +macro_rules! prime_pass { + ($f:expr, $g:expr, $out:expr, $sq:ident, $mm:ident, $qdata:expr) => {{ + // One contiguous 6x512 buffer: Good's three f-tracks then three g-tracks, + // exactly the layout `ntt512` batches over. Writing the tracks straight + // into it avoids copying 6 KB in and out per prime. + // SAFETY: `good` stores to every 16-coefficient block of all three of + // its output tracks unconditionally, so the two calls below fill all + // 6 x 512 elements before `ntt512` reads any. + let mut fg_slot = core::mem::MaybeUninit::<[i16; 6 * 512]>::uninit(); + let fg = crate::scratch::uninit(&mut fg_slot); + good(&mut fg[..3 * 512], $f); + good(&mut fg[3 * 512..], $g); + ntt512(fg, 6, $qdata); + + // SAFETY: the pointwise loop below stores to `i`, `512 + i` and + // `1024 + i` for every `i` in `(0..512).step_by(16)`, covering all + // 3 x 512 elements before `invntt512` reads them. + let mut hpad_slot = core::mem::MaybeUninit::<[i16; 3 * 512]>::uninit(); + let hpad = crate::scratch::uninit(&mut hpad_slot); + let mut i = 0usize; + while i < 512 { + let f0 = $sq(_mm256_loadu_si256(fg.as_ptr().add(i) as *const __m256i)); + let f1 = $sq(_mm256_loadu_si256( + fg.as_ptr().add(512 + i) as *const __m256i + )); + let f2 = $sq(_mm256_loadu_si256( + fg.as_ptr().add(1024 + i) as *const __m256i + )); + let g0 = $sq(_mm256_loadu_si256( + fg.as_ptr().add(1536 + i) as *const __m256i + )); + let g1 = $sq(_mm256_loadu_si256( + fg.as_ptr().add(2048 + i) as *const __m256i + )); + let g2 = $sq(_mm256_loadu_si256( + fg.as_ptr().add(2560 + i) as *const __m256i + )); + let dsum = add16(add16($mm(f0, g0), $mm(f1, g1)), $mm(f2, g2)); + let r0 = add16(dsum, $mm(sub16(f2, f1), sub16(g1, g2))); + let r1 = add16(dsum, $mm(sub16(f1, f0), sub16(g0, g1))); + let r2 = add16(dsum, $mm(sub16(f0, f2), sub16(g2, g0))); + _mm256_storeu_si256(hpad.as_mut_ptr().add(i) as *mut __m256i, $sq(r0)); + _mm256_storeu_si256(hpad.as_mut_ptr().add(512 + i) as *mut __m256i, $sq(r1)); + _mm256_storeu_si256(hpad.as_mut_ptr().add(1024 + i) as *mut __m256i, $sq(r2)); + i += 16; + } + + invntt512(hpad, 3, $qdata); + ungood($out, &hpad[..]); + }}; +} + +/// 1536-coefficient product mod q, via both primes plus CRT (the reference's +/// `mult768`). +#[target_feature(enable = "avx2")] +fn mult768(h: &mut [i16; 1536], f: &[i16; 768], g: &[i16; 768]) { + unsafe { + // SAFETY: each is filled by `ungood`, whose store covers `k * 512 + j` + // for all `k < 3` and every `j` in `(0..512).step_by(16)` — all 1536 + // elements — before the CRT loop reads them. + let mut h7681_slot = core::mem::MaybeUninit::<[i16; 1536]>::uninit(); + let h7681 = crate::scratch::uninit(&mut h7681_slot); + let mut h10753_slot = core::mem::MaybeUninit::<[i16; 1536]>::uninit(); + let h10753 = crate::scratch::uninit(&mut h10753_slot); + prime_pass!(f, g, h7681, squeeze_7681, mulmod_7681, &QDATA_7681); + prime_pass!(f, g, h10753, squeeze_10753, mulmod_10753, &QDATA_10753); + + // CRT the two residues back to mod 4591. + let mut i = 0usize; + while i < 1536 { + let u1 = mulmod_10753( + _mm256_loadu_si256(h10753.as_ptr().add(i) as *const __m256i), + _mm256_set1_epi16(1268), + ); + let u2 = mulmod_7681( + _mm256_loadu_si256(h7681.as_ptr().add(i) as *const __m256i), + _mm256_set1_epi16(956), + ); + let t = mulmod_7681(sub16(u2, u1), _mm256_set1_epi16(-2539)); + let t = add16(u1, mulmod_4591(t, _mm256_set1_epi16(-710))); + _mm256_storeu_si256(h.as_mut_ptr().add(i) as *mut __m256i, t); + i += 16; + } + } +} + +/// 1536-coefficient product mod 3. Product coefficients are bounded by p = 761, +/// far inside 7681/2, so a single prime suffices — no CRT. +#[target_feature(enable = "avx2")] +fn mult768_3(h: &mut [i16; 1536], f: &[i16; 768], g: &[i16; 768]) { + unsafe { + // SAFETY: filled in full by `ungood`, as in `mult768`. + let mut h7681_slot = core::mem::MaybeUninit::<[i16; 1536]>::uninit(); + let h7681 = crate::scratch::uninit(&mut h7681_slot); + prime_pass!(f, g, h7681, squeeze_7681, mulmod_7681, &QDATA_7681); + let mut i = 0usize; + while i < 1536 { + let u = mulmod_7681( + _mm256_loadu_si256(h7681.as_ptr().add(i) as *const __m256i), + _mm256_set1_epi16(956), + ); + _mm256_storeu_si256(h.as_mut_ptr().add(i) as *mut __m256i, u); + i += 16; + } + } +} + +#[inline] +#[target_feature(enable = "avx2")] +fn squeeze_3(x: __m256i) -> __m256i { + squeeze(x, 10923, 3) +} + +/// Fully reduce to {-1, 0, 1}. +#[inline] +#[target_feature(enable = "avx2")] +fn freeze_3(x: __m256i) -> __m256i { + let three = _mm256_set1_epi16(3); + let x = add16(x, _mm256_and_si256(three, _mm256_srai_epi16::<15>(x))); + let m = _mm256_srai_epi16::<15>(sub16(x, _mm256_set1_epi16(2))); + _mm256_blendv_epi8(sub16(x, three), x, m) +} + +/// `h = f · g` in R/3 for p = 761, via the single-prime NTT. +#[target_feature(enable = "avx2")] +pub fn mult3_761(h: &mut [i8], f: &[i8], g: &[i8]) { + unsafe { + const P: usize = 761; + // SAFETY: the loop writes `0..P` and the tail clear covers `P..768`, + // so both are complete before `mult768_3` reads them. The 768-length + // padding is what makes the transform's zero-extension work, so it has + // to be written explicitly now that the buffer is not pre-zeroed. + let mut fp_slot = core::mem::MaybeUninit::<[i16; 768]>::uninit(); + let fp = crate::scratch::uninit(&mut fp_slot); + let mut gp_slot = core::mem::MaybeUninit::<[i16; 768]>::uninit(); + let gp = crate::scratch::uninit(&mut gp_slot); + for k in 0..P { + fp[k] = i16::from(f[k]); + gp[k] = i16::from(g[k]); + } + for k in P..768 { + fp[k] = 0; + gp[k] = 0; + } + + // SAFETY: `mult768_3` writes all 1536 coefficients of its output. + let mut fg_slot = core::mem::MaybeUninit::<[i16; 1536]>::uninit(); + let fg = crate::scratch::uninit(&mut fg_slot); + mult768_3(fg, fp, gp); + + fg[0] -= fg[P - 1]; + // SAFETY: the loop below stores every 16-element block of `0..768`. + let mut out_slot = core::mem::MaybeUninit::<[i16; 768]>::uninit(); + let out = crate::scratch::uninit(&mut out_slot); + let mut i = 0usize; + while i < 768 { + let a = _mm256_loadu_si256(fg.as_ptr().add(i) as *const __m256i); + let b = _mm256_loadu_si256(fg.as_ptr().add(i + P) as *const __m256i); + let c = _mm256_loadu_si256(fg.as_ptr().add(i + P - 1) as *const __m256i); + let x = freeze_3(squeeze_3(add16(a, add16(b, c)))); + _mm256_storeu_si256(out.as_mut_ptr().add(i) as *mut __m256i, x); + i += 16; + } + for k in 0..P { + h[k] = out[k] as i8; + } + // Both operands are secret at every call site — wipe the working buffers. + wipe(fp); + wipe(gp); + wipe(fg); + wipe(out); + } +} + +#[target_feature(enable = "avx2")] +pub fn mult761(h: &mut [i16], f: &[i16], g: &[i8]) { + unsafe { + const P: usize = 761; + // SAFETY: the loop below stores every 16-element block of `0..768`. + let mut fp_slot = core::mem::MaybeUninit::<[i16; 768]>::uninit(); + let fp = crate::scratch::uninit(&mut fp_slot); + // SAFETY: written over `0..P` by the copy loop and `P..768` by the tail + // clear, both before `mult768` reads it. + let mut gp_slot = core::mem::MaybeUninit::<[i16; 768]>::uninit(); + let gp = crate::scratch::uninit(&mut gp_slot); + let mut i = 0usize; + while i < 768 { + let x = if i < P { + _mm256_loadu_si256(f.as_ptr().add(i) as *const __m256i) + } else { + _mm256_setzero_si256() + }; + _mm256_storeu_si256( + fp.as_mut_ptr().add(i) as *mut __m256i, + freeze_4591(squeeze_4591(x)), + ); + i += 16; + } + // The last block overruns p; clear the tail explicitly. + for k in P..768 { + fp[k] = 0; + } + for k in 0..P { + gp[k] = i16::from(g[k]); + } + for k in P..768 { + gp[k] = 0; + } + + // SAFETY: `mult768` writes all 1536 coefficients of its output. + let mut fg_slot = core::mem::MaybeUninit::<[i16; 1536]>::uninit(); + let fg = crate::scratch::uninit(&mut fg_slot); + mult768(fg, fp, gp); + + fg[0] -= fg[P - 1]; + let mut i = 0usize; + while i < 768 { + let a = _mm256_loadu_si256(fg.as_ptr().add(i) as *const __m256i); + let b = _mm256_loadu_si256(fg.as_ptr().add(i + P) as *const __m256i); + let c = _mm256_loadu_si256(fg.as_ptr().add(i + P - 1) as *const __m256i); + let x = freeze_4591(squeeze_4591(add16(a, add16(b, c)))); + _mm256_storeu_si256(fp.as_mut_ptr().add(i) as *mut __m256i, x); + i += 16; + } + h[..P].copy_from_slice(&fp[..P]); + // `g` is secret at every call site — wipe the operand copies and the + // product scratch. + wipe(fp); + wipe(gp); + wipe(fg); + } +} + +#[target_feature(enable = "avx2")] +fn ntt512(f: &mut [i16], reps: usize, qdata: &[i16; 1696]) { + unsafe { + let qv = ld(qdata, 768); + let qround32v = ld(qdata, 1616); + let mut reps = reps; + let base = 0usize; + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a0 = _mm256_loadu_si256(fp.add(0) as *const __m256i); + let a16 = _mm256_loadu_si256(fp.add(256) as *const __m256i); + let b0 = add16(a0, a16); + let b16 = sub16(a0, a16); + let a8 = _mm256_loadu_si256(fp.add(128) as *const __m256i); + let a24 = _mm256_loadu_si256(fp.add(384) as *const __m256i); + let b8 = add16(a8, a24); + let mut b24 = sub16(a8, a24); + let a4 = _mm256_loadu_si256(fp.add(64) as *const __m256i); + let a20 = _mm256_loadu_si256(fp.add(320) as *const __m256i); + let b4 = add16(a4, a20); + let b20 = sub16(a4, a20); + let a12 = _mm256_loadu_si256(fp.add(192) as *const __m256i); + let a28 = _mm256_loadu_si256(fp.add(448) as *const __m256i); + let b12 = add16(a12, a28); + let mut b28 = sub16(a12, a28); + let c0 = add16(b0, b8); + let c8 = sub16(b0, b8); + let mut c4 = add16(b4, b12); + let mut c12 = sub16(b4, b12); + b24 = mulmod_scaled(b24, ld(qdata, 1632), ld(qdata, 1552), qv); + let c16 = add16(b16, b24); + let c24 = sub16(b16, b24); + b28 = mulmod_scaled(b28, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut c20 = add16(b20, b28); + let mut c28 = sub16(b20, b28); + c4 = reduce(c4, qv, qround32v); + let mut d0 = add16(c0, c4); + let mut d4 = sub16(c0, c4); + c12 = mulmod_scaled(c12, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut d8 = add16(c8, c12); + let mut d12 = sub16(c8, c12); + c20 = mulmod_scaled(c20, ld(qdata, 1664), ld(qdata, 1584), qv); + let mut d16 = add16(c16, c20); + let mut d20 = sub16(c16, c20); + c28 = mulmod_scaled(c28, ld(qdata, 1680), ld(qdata, 1600), qv); + let mut d24 = add16(c24, c28); + let mut d28 = sub16(c24, c28); + d0 = reduce(d0, qv, qround32v); + d4 = mulmod_scaled(d4, ld(qdata, 256), ld(qdata, 1040), qv); + d8 = mulmod_scaled(d8, ld(qdata, 384), ld(qdata, 1168), qv); + d12 = mulmod_scaled(d12, ld(qdata, 448), ld(qdata, 1232), qv); + d16 = mulmod_scaled(d16, ld(qdata, 512), ld(qdata, 1296), qv); + d20 = mulmod_scaled(d20, ld(qdata, 576), ld(qdata, 1360), qv); + d24 = mulmod_scaled(d24, ld(qdata, 704), ld(qdata, 1488), qv); + d28 = mulmod_scaled(d28, ld(qdata, 640), ld(qdata, 1424), qv); + let e0 = perm_lo(d0, d4); + let e4 = perm_hi(d0, d4); + let e8 = perm_lo(d8, d12); + let e12 = perm_hi(d8, d12); + let e16 = perm_lo(d16, d20); + let e20 = perm_hi(d16, d20); + let e24 = perm_lo(d24, d28); + let e28 = perm_hi(d24, d28); + _mm256_storeu_si256(fp.add(0) as *mut __m256i, e0); + _mm256_storeu_si256(fp.add(64) as *mut __m256i, e4); + _mm256_storeu_si256(fp.add(128) as *mut __m256i, e8); + _mm256_storeu_si256(fp.add(192) as *mut __m256i, e12); + _mm256_storeu_si256(fp.add(256) as *mut __m256i, e16); + _mm256_storeu_si256(fp.add(320) as *mut __m256i, e20); + _mm256_storeu_si256(fp.add(384) as *mut __m256i, e24); + _mm256_storeu_si256(fp.add(448) as *mut __m256i, e28); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a1 = _mm256_loadu_si256(fp.add(16) as *const __m256i); + let a17 = _mm256_loadu_si256(fp.add(272) as *const __m256i); + let b1 = add16(a1, a17); + let b17 = sub16(a1, a17); + let a9 = _mm256_loadu_si256(fp.add(144) as *const __m256i); + let a25 = _mm256_loadu_si256(fp.add(400) as *const __m256i); + let b9 = add16(a9, a25); + let mut b25 = sub16(a9, a25); + let a5 = _mm256_loadu_si256(fp.add(80) as *const __m256i); + let a21 = _mm256_loadu_si256(fp.add(336) as *const __m256i); + let b5 = add16(a5, a21); + let b21 = sub16(a5, a21); + let a13 = _mm256_loadu_si256(fp.add(208) as *const __m256i); + let a29 = _mm256_loadu_si256(fp.add(464) as *const __m256i); + let b13 = add16(a13, a29); + let mut b29 = sub16(a13, a29); + let c1 = add16(b1, b9); + let c9 = sub16(b1, b9); + let mut c5 = add16(b5, b13); + let mut c13 = sub16(b5, b13); + b25 = mulmod_scaled(b25, ld(qdata, 1632), ld(qdata, 1552), qv); + let c17 = add16(b17, b25); + let c25 = sub16(b17, b25); + b29 = mulmod_scaled(b29, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut c21 = add16(b21, b29); + let mut c29 = sub16(b21, b29); + c5 = reduce(c5, qv, qround32v); + let mut d1 = add16(c1, c5); + let mut d5 = sub16(c1, c5); + c13 = mulmod_scaled(c13, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut d9 = add16(c9, c13); + let mut d13 = sub16(c9, c13); + c21 = mulmod_scaled(c21, ld(qdata, 1664), ld(qdata, 1584), qv); + let mut d17 = add16(c17, c21); + let mut d21 = sub16(c17, c21); + c29 = mulmod_scaled(c29, ld(qdata, 1680), ld(qdata, 1600), qv); + let mut d25 = add16(c25, c29); + let mut d29 = sub16(c25, c29); + d1 = reduce(d1, qv, qround32v); + d5 = mulmod_scaled(d5, ld(qdata, 272), ld(qdata, 1056), qv); + d9 = mulmod_scaled(d9, ld(qdata, 400), ld(qdata, 1184), qv); + d13 = mulmod_scaled(d13, ld(qdata, 464), ld(qdata, 1248), qv); + d17 = mulmod_scaled(d17, ld(qdata, 528), ld(qdata, 1312), qv); + d21 = mulmod_scaled(d21, ld(qdata, 592), ld(qdata, 1376), qv); + d25 = mulmod_scaled(d25, ld(qdata, 720), ld(qdata, 1504), qv); + d29 = mulmod_scaled(d29, ld(qdata, 656), ld(qdata, 1440), qv); + let e1 = perm_lo(d1, d5); + let e5 = perm_hi(d1, d5); + let e9 = perm_lo(d9, d13); + let e13 = perm_hi(d9, d13); + let e17 = perm_lo(d17, d21); + let e21 = perm_hi(d17, d21); + let e25 = perm_lo(d25, d29); + let e29 = perm_hi(d25, d29); + _mm256_storeu_si256(fp.add(16) as *mut __m256i, e1); + _mm256_storeu_si256(fp.add(80) as *mut __m256i, e5); + _mm256_storeu_si256(fp.add(144) as *mut __m256i, e9); + _mm256_storeu_si256(fp.add(208) as *mut __m256i, e13); + _mm256_storeu_si256(fp.add(272) as *mut __m256i, e17); + _mm256_storeu_si256(fp.add(336) as *mut __m256i, e21); + _mm256_storeu_si256(fp.add(400) as *mut __m256i, e25); + _mm256_storeu_si256(fp.add(464) as *mut __m256i, e29); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a2 = _mm256_loadu_si256(fp.add(32) as *const __m256i); + let a18 = _mm256_loadu_si256(fp.add(288) as *const __m256i); + let b2 = add16(a2, a18); + let b18 = sub16(a2, a18); + let a10 = _mm256_loadu_si256(fp.add(160) as *const __m256i); + let a26 = _mm256_loadu_si256(fp.add(416) as *const __m256i); + let b10 = add16(a10, a26); + let mut b26 = sub16(a10, a26); + let a6 = _mm256_loadu_si256(fp.add(96) as *const __m256i); + let a22 = _mm256_loadu_si256(fp.add(352) as *const __m256i); + let b6 = add16(a6, a22); + let b22 = sub16(a6, a22); + let a14 = _mm256_loadu_si256(fp.add(224) as *const __m256i); + let a30 = _mm256_loadu_si256(fp.add(480) as *const __m256i); + let b14 = add16(a14, a30); + let mut b30 = sub16(a14, a30); + let c2 = add16(b2, b10); + let c10 = sub16(b2, b10); + let mut c6 = add16(b6, b14); + let mut c14 = sub16(b6, b14); + b26 = mulmod_scaled(b26, ld(qdata, 1632), ld(qdata, 1552), qv); + let c18 = add16(b18, b26); + let c26 = sub16(b18, b26); + b30 = mulmod_scaled(b30, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut c22 = add16(b22, b30); + let mut c30 = sub16(b22, b30); + c6 = reduce(c6, qv, qround32v); + let mut d2 = add16(c2, c6); + let mut d6 = sub16(c2, c6); + c14 = mulmod_scaled(c14, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut d10 = add16(c10, c14); + let mut d14 = sub16(c10, c14); + c22 = mulmod_scaled(c22, ld(qdata, 1664), ld(qdata, 1584), qv); + let mut d18 = add16(c18, c22); + let mut d22 = sub16(c18, c22); + c30 = mulmod_scaled(c30, ld(qdata, 1680), ld(qdata, 1600), qv); + let mut d26 = add16(c26, c30); + let mut d30 = sub16(c26, c30); + d2 = reduce(d2, qv, qround32v); + d6 = mulmod_scaled(d6, ld(qdata, 288), ld(qdata, 1072), qv); + d10 = mulmod_scaled(d10, ld(qdata, 416), ld(qdata, 1200), qv); + d14 = mulmod_scaled(d14, ld(qdata, 480), ld(qdata, 1264), qv); + d18 = mulmod_scaled(d18, ld(qdata, 544), ld(qdata, 1328), qv); + d22 = mulmod_scaled(d22, ld(qdata, 608), ld(qdata, 1392), qv); + d26 = mulmod_scaled(d26, ld(qdata, 736), ld(qdata, 1520), qv); + d30 = mulmod_scaled(d30, ld(qdata, 672), ld(qdata, 1456), qv); + let e2 = perm_lo(d2, d6); + let e6 = perm_hi(d2, d6); + let e10 = perm_lo(d10, d14); + let e14 = perm_hi(d10, d14); + let e18 = perm_lo(d18, d22); + let e22 = perm_hi(d18, d22); + let e26 = perm_lo(d26, d30); + let e30 = perm_hi(d26, d30); + _mm256_storeu_si256(fp.add(32) as *mut __m256i, e2); + _mm256_storeu_si256(fp.add(96) as *mut __m256i, e6); + _mm256_storeu_si256(fp.add(160) as *mut __m256i, e10); + _mm256_storeu_si256(fp.add(224) as *mut __m256i, e14); + _mm256_storeu_si256(fp.add(288) as *mut __m256i, e18); + _mm256_storeu_si256(fp.add(352) as *mut __m256i, e22); + _mm256_storeu_si256(fp.add(416) as *mut __m256i, e26); + _mm256_storeu_si256(fp.add(480) as *mut __m256i, e30); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a3 = _mm256_loadu_si256(fp.add(48) as *const __m256i); + let a19 = _mm256_loadu_si256(fp.add(304) as *const __m256i); + let b3 = add16(a3, a19); + let b19 = sub16(a3, a19); + let a11 = _mm256_loadu_si256(fp.add(176) as *const __m256i); + let a27 = _mm256_loadu_si256(fp.add(432) as *const __m256i); + let b11 = add16(a11, a27); + let mut b27 = sub16(a11, a27); + let a7 = _mm256_loadu_si256(fp.add(112) as *const __m256i); + let a23 = _mm256_loadu_si256(fp.add(368) as *const __m256i); + let b7 = add16(a7, a23); + let b23 = sub16(a7, a23); + let a15 = _mm256_loadu_si256(fp.add(240) as *const __m256i); + let a31 = _mm256_loadu_si256(fp.add(496) as *const __m256i); + let b15 = add16(a15, a31); + let mut b31 = sub16(a15, a31); + let c3 = add16(b3, b11); + let c11 = sub16(b3, b11); + let mut c7 = add16(b7, b15); + let mut c15 = sub16(b7, b15); + b27 = mulmod_scaled(b27, ld(qdata, 1632), ld(qdata, 1552), qv); + let c19 = add16(b19, b27); + let c27 = sub16(b19, b27); + b31 = mulmod_scaled(b31, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut c23 = add16(b23, b31); + let mut c31 = sub16(b23, b31); + c7 = reduce(c7, qv, qround32v); + let mut d3 = add16(c3, c7); + let mut d7 = sub16(c3, c7); + c15 = mulmod_scaled(c15, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut d11 = add16(c11, c15); + let mut d15 = sub16(c11, c15); + c23 = mulmod_scaled(c23, ld(qdata, 1664), ld(qdata, 1584), qv); + let mut d19 = add16(c19, c23); + let mut d23 = sub16(c19, c23); + c31 = mulmod_scaled(c31, ld(qdata, 1680), ld(qdata, 1600), qv); + let mut d27 = add16(c27, c31); + let mut d31 = sub16(c27, c31); + d3 = reduce(d3, qv, qround32v); + d7 = mulmod_scaled(d7, ld(qdata, 304), ld(qdata, 1088), qv); + d11 = mulmod_scaled(d11, ld(qdata, 432), ld(qdata, 1216), qv); + d15 = mulmod_scaled(d15, ld(qdata, 496), ld(qdata, 1280), qv); + d19 = mulmod_scaled(d19, ld(qdata, 560), ld(qdata, 1344), qv); + d23 = mulmod_scaled(d23, ld(qdata, 624), ld(qdata, 1408), qv); + d27 = mulmod_scaled(d27, ld(qdata, 752), ld(qdata, 1536), qv); + d31 = mulmod_scaled(d31, ld(qdata, 688), ld(qdata, 1472), qv); + let e3 = perm_lo(d3, d7); + let e7 = perm_hi(d3, d7); + let e11 = perm_lo(d11, d15); + let e15 = perm_hi(d11, d15); + let e19 = perm_lo(d19, d23); + let e23 = perm_hi(d19, d23); + let e27 = perm_lo(d27, d31); + let e31 = perm_hi(d27, d31); + _mm256_storeu_si256(fp.add(48) as *mut __m256i, e3); + _mm256_storeu_si256(fp.add(112) as *mut __m256i, e7); + _mm256_storeu_si256(fp.add(176) as *mut __m256i, e11); + _mm256_storeu_si256(fp.add(240) as *mut __m256i, e15); + _mm256_storeu_si256(fp.add(304) as *mut __m256i, e19); + _mm256_storeu_si256(fp.add(368) as *mut __m256i, e23); + _mm256_storeu_si256(fp.add(432) as *mut __m256i, e27); + _mm256_storeu_si256(fp.add(496) as *mut __m256i, e31); + } + reps *= 2; + reps *= 2; + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 128 * r); + let a0 = _mm256_loadu_si256(fp.add(0) as *const __m256i); + let a2 = _mm256_loadu_si256(fp.add(32) as *const __m256i); + let b0 = add16(a0, a2); + let b2 = sub16(a0, a2); + let a4 = _mm256_loadu_si256(fp.add(64) as *const __m256i); + let a6 = _mm256_loadu_si256(fp.add(96) as *const __m256i); + let b4 = add16(a4, a6); + let b6 = sub16(a4, a6); + let a1 = _mm256_loadu_si256(fp.add(16) as *const __m256i); + let a3 = _mm256_loadu_si256(fp.add(48) as *const __m256i); + let b1 = add16(a1, a3); + let mut b3 = sub16(a1, a3); + let a5 = _mm256_loadu_si256(fp.add(80) as *const __m256i); + let a7 = _mm256_loadu_si256(fp.add(112) as *const __m256i); + let b5 = add16(a5, a7); + let mut b7 = sub16(a5, a7); + let mut c0 = add16(b0, b1); + let mut c1 = sub16(b0, b1); + let mut c4 = add16(b4, b5); + let mut c5 = sub16(b4, b5); + b3 = mulmod_scaled(b3, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut c2 = add16(b2, b3); + let mut c3 = sub16(b2, b3); + b7 = mulmod_scaled(b7, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut c6 = add16(b6, b7); + let mut c7 = sub16(b6, b7); + c0 = reduce(c0, qv, qround32v); + c4 = reduce(c4, qv, qround32v); + c1 = mulmod_scaled(c1, ld(qdata, 128), ld(qdata, 912), qv); + c5 = mulmod_scaled(c5, ld(qdata, 144), ld(qdata, 928), qv); + c2 = mulmod_scaled(c2, ld(qdata, 192), ld(qdata, 976), qv); + c6 = mulmod_scaled(c6, ld(qdata, 208), ld(qdata, 992), qv); + c3 = mulmod_scaled(c3, ld(qdata, 224), ld(qdata, 1008), qv); + c7 = mulmod_scaled(c7, ld(qdata, 240), ld(qdata, 1024), qv); + let d0 = _mm256_unpacklo_epi16(c0, c2); + let d2 = _mm256_unpackhi_epi16(c0, c2); + let d1 = _mm256_unpacklo_epi16(c1, c3); + let d3 = _mm256_unpackhi_epi16(c1, c3); + let d4 = _mm256_unpacklo_epi16(c4, c6); + let d6 = _mm256_unpackhi_epi16(c4, c6); + let d5 = _mm256_unpacklo_epi16(c5, c7); + let d7 = _mm256_unpackhi_epi16(c5, c7); + let e0 = add16(d0, d4); + let e4 = sub16(d0, d4); + let e2 = add16(d2, d6); + let e6 = sub16(d2, d6); + let e1 = add16(d1, d5); + let e5 = sub16(d1, d5); + let e3 = add16(d3, d7); + let e7 = sub16(d3, d7); + let f0 = _mm256_unpacklo_epi32(e0, e1); + let f1 = _mm256_unpackhi_epi32(e0, e1); + let f2 = _mm256_unpacklo_epi32(e2, e3); + let f3 = _mm256_unpackhi_epi32(e2, e3); + let f4 = _mm256_unpacklo_epi32(e4, e5); + let f5 = _mm256_unpackhi_epi32(e4, e5); + let f6 = _mm256_unpacklo_epi32(e6, e7); + let f7 = _mm256_unpackhi_epi32(e6, e7); + _mm256_storeu_si256(fp.add(0) as *mut __m256i, f0); + _mm256_storeu_si256(fp.add(16) as *mut __m256i, f1); + _mm256_storeu_si256(fp.add(32) as *mut __m256i, f2); + _mm256_storeu_si256(fp.add(48) as *mut __m256i, f3); + _mm256_storeu_si256(fp.add(64) as *mut __m256i, f4); + _mm256_storeu_si256(fp.add(80) as *mut __m256i, f5); + _mm256_storeu_si256(fp.add(96) as *mut __m256i, f6); + _mm256_storeu_si256(fp.add(112) as *mut __m256i, f7); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 128 * r); + let a0 = _mm256_loadu_si256(fp.add(0) as *const __m256i); + let a2 = _mm256_loadu_si256(fp.add(32) as *const __m256i); + let mut b0 = add16(a0, a2); + let mut b2 = sub16(a0, a2); + let a1 = _mm256_loadu_si256(fp.add(16) as *const __m256i); + let a3 = _mm256_loadu_si256(fp.add(48) as *const __m256i); + let mut b1 = add16(a1, a3); + let mut b3 = sub16(a1, a3); + let a4 = _mm256_loadu_si256(fp.add(64) as *const __m256i); + let mut a6 = _mm256_loadu_si256(fp.add(96) as *const __m256i); + a6 = mulmod_scaled(a6, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut b4 = add16(a4, a6); + let mut b6 = sub16(a4, a6); + let a5 = _mm256_loadu_si256(fp.add(80) as *const __m256i); + let mut a7 = _mm256_loadu_si256(fp.add(112) as *const __m256i); + a7 = mulmod_scaled(a7, ld(qdata, 1632), ld(qdata, 1552), qv); + let mut b5 = add16(a5, a7); + let mut b7 = sub16(a5, a7); + b0 = reduce(b0, qv, qround32v); + b1 = reduce(b1, qv, qround32v); + b2 = mulmod_scaled(b2, ld(qdata, 0), ld(qdata, 784), qv); + b3 = mulmod_scaled(b3, ld(qdata, 16), ld(qdata, 800), qv); + b4 = mulmod_scaled(b4, ld(qdata, 64), ld(qdata, 848), qv); + b5 = mulmod_scaled(b5, ld(qdata, 80), ld(qdata, 864), qv); + b6 = mulmod_scaled(b6, ld(qdata, 96), ld(qdata, 880), qv); + b7 = mulmod_scaled(b7, ld(qdata, 112), ld(qdata, 896), qv); + let c0 = _mm256_unpacklo_epi64(b0, b4); + let c4 = _mm256_unpackhi_epi64(b0, b4); + let c1 = _mm256_unpacklo_epi64(b1, b5); + let c5 = _mm256_unpackhi_epi64(b1, b5); + let c2 = _mm256_unpacklo_epi64(b2, b6); + let c6 = _mm256_unpackhi_epi64(b2, b6); + let c3 = _mm256_unpacklo_epi64(b3, b7); + let c7 = _mm256_unpackhi_epi64(b3, b7); + let d0 = add16(c0, c1); + let d1 = sub16(c0, c1); + let d4 = add16(c4, c5); + let mut d5 = sub16(c4, c5); + let d2 = add16(c2, c3); + let d3 = sub16(c2, c3); + let d6 = add16(c6, c7); + let mut d7 = sub16(c6, c7); + let e0 = add16(d0, d4); + let e4 = sub16(d0, d4); + let e2 = add16(d2, d6); + let e6 = sub16(d2, d6); + d5 = mulmod_scaled(d5, ld(qdata, 1632), ld(qdata, 1552), qv); + let e1 = add16(d1, d5); + let e5 = sub16(d1, d5); + d7 = mulmod_scaled(d7, ld(qdata, 1632), ld(qdata, 1552), qv); + let e3 = add16(d3, d7); + let e7 = sub16(d3, d7); + _mm256_storeu_si256(fp.add(0) as *mut __m256i, e0); + _mm256_storeu_si256(fp.add(16) as *mut __m256i, e1); + _mm256_storeu_si256(fp.add(32) as *mut __m256i, e2); + _mm256_storeu_si256(fp.add(48) as *mut __m256i, e3); + _mm256_storeu_si256(fp.add(64) as *mut __m256i, e4); + _mm256_storeu_si256(fp.add(80) as *mut __m256i, e5); + _mm256_storeu_si256(fp.add(96) as *mut __m256i, e6); + _mm256_storeu_si256(fp.add(112) as *mut __m256i, e7); + } + } +} + +#[target_feature(enable = "avx2")] +fn invntt512(f: &mut [i16], reps: usize, qdata: &[i16; 1696]) { + unsafe { + let qv = ld(qdata, 768); + let qround32v = ld(qdata, 1616); + let mut reps = reps; + let base = 0usize; + reps *= 4; + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 128 * r); + let a3 = _mm256_loadu_si256(fp.add(48) as *const __m256i); + let a7 = _mm256_loadu_si256(fp.add(112) as *const __m256i); + let b3 = add16(a3, a7); + let mut b7 = sub16(a3, a7); + b7 = mulmod_scaled(b7, ld(qdata, 1648), ld(qdata, 1568), qv); + let a1 = _mm256_loadu_si256(fp.add(16) as *const __m256i); + let a5 = _mm256_loadu_si256(fp.add(80) as *const __m256i); + let b1 = add16(a1, a5); + let mut b5 = sub16(a1, a5); + b5 = mulmod_scaled(b5, ld(qdata, 1648), ld(qdata, 1568), qv); + let a2 = _mm256_loadu_si256(fp.add(32) as *const __m256i); + let a6 = _mm256_loadu_si256(fp.add(96) as *const __m256i); + let b2 = add16(a2, a6); + let b6 = sub16(a2, a6); + let a0 = _mm256_loadu_si256(fp.add(0) as *const __m256i); + let a4 = _mm256_loadu_si256(fp.add(64) as *const __m256i); + let b0 = add16(a0, a4); + let b4 = sub16(a0, a4); + let c6 = add16(b6, b7); + let c7 = sub16(b6, b7); + let c2 = add16(b2, b3); + let c3 = sub16(b2, b3); + let c4 = add16(b4, b5); + let c5 = sub16(b4, b5); + let c0 = add16(b0, b1); + let c1 = sub16(b0, b1); + let mut d3 = _mm256_unpacklo_epi64(c3, c7); + let mut d7 = _mm256_unpackhi_epi64(c3, c7); + let mut d2 = _mm256_unpacklo_epi64(c2, c6); + let mut d6 = _mm256_unpackhi_epi64(c2, c6); + let mut d1 = _mm256_unpacklo_epi64(c1, c5); + let mut d5 = _mm256_unpackhi_epi64(c1, c5); + let mut d0 = _mm256_unpacklo_epi64(c0, c4); + let mut d4 = _mm256_unpackhi_epi64(c0, c4); + d7 = mulmod_scaled(d7, ld(qdata, 80), ld(qdata, 864), qv); + d6 = mulmod_scaled(d6, ld(qdata, 64), ld(qdata, 848), qv); + d5 = mulmod_scaled(d5, ld(qdata, 112), ld(qdata, 896), qv); + d4 = mulmod_scaled(d4, ld(qdata, 96), ld(qdata, 880), qv); + d3 = mulmod_scaled(d3, ld(qdata, 48), ld(qdata, 832), qv); + d2 = mulmod_scaled(d2, ld(qdata, 32), ld(qdata, 816), qv); + d1 = reduce(d1, qv, qround32v); + d0 = reduce(d0, qv, qround32v); + let e5 = add16(d5, d7); + let mut e7 = sub16(d5, d7); + e7 = mulmod_scaled(e7, ld(qdata, 1648), ld(qdata, 1568), qv); + let e4 = add16(d4, d6); + let mut e6 = sub16(d4, d6); + e6 = mulmod_scaled(e6, ld(qdata, 1648), ld(qdata, 1568), qv); + let e1 = add16(d1, d3); + let e3 = sub16(d1, d3); + let e0 = add16(d0, d2); + let e2 = sub16(d0, d2); + _mm256_storeu_si256(fp.add(0) as *mut __m256i, e0); + _mm256_storeu_si256(fp.add(16) as *mut __m256i, e1); + _mm256_storeu_si256(fp.add(32) as *mut __m256i, e2); + _mm256_storeu_si256(fp.add(48) as *mut __m256i, e3); + _mm256_storeu_si256(fp.add(64) as *mut __m256i, e4); + _mm256_storeu_si256(fp.add(80) as *mut __m256i, e5); + _mm256_storeu_si256(fp.add(96) as *mut __m256i, e6); + _mm256_storeu_si256(fp.add(112) as *mut __m256i, e7); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 128 * r); + let a6 = _mm256_loadu_si256(fp.add(96) as *const __m256i); + let a7 = _mm256_loadu_si256(fp.add(112) as *const __m256i); + let b6 = _mm256_unpacklo_epi32(a6, a7); + let b7 = _mm256_unpackhi_epi32(a6, a7); + let c6 = _mm256_unpacklo_epi32(b6, b7); + let c7 = _mm256_unpackhi_epi32(b6, b7); + let a4 = _mm256_loadu_si256(fp.add(64) as *const __m256i); + let a5 = _mm256_loadu_si256(fp.add(80) as *const __m256i); + let b4 = _mm256_unpacklo_epi32(a4, a5); + let b5 = _mm256_unpackhi_epi32(a4, a5); + let c4 = _mm256_unpacklo_epi32(b4, b5); + let c5 = _mm256_unpackhi_epi32(b4, b5); + let a2 = _mm256_loadu_si256(fp.add(32) as *const __m256i); + let a3 = _mm256_loadu_si256(fp.add(48) as *const __m256i); + let b2 = _mm256_unpacklo_epi32(a2, a3); + let b3 = _mm256_unpackhi_epi32(a2, a3); + let c2 = _mm256_unpacklo_epi32(b2, b3); + let c3 = _mm256_unpackhi_epi32(b2, b3); + let a0 = _mm256_loadu_si256(fp.add(0) as *const __m256i); + let a1 = _mm256_loadu_si256(fp.add(16) as *const __m256i); + let b0 = _mm256_unpacklo_epi32(a0, a1); + let b1 = _mm256_unpackhi_epi32(a0, a1); + let c0 = _mm256_unpacklo_epi32(b0, b1); + let c1 = _mm256_unpackhi_epi32(b0, b1); + let d3 = add16(c3, c7); + let d7 = sub16(c3, c7); + let d1 = add16(c1, c5); + let d5 = sub16(c1, c5); + let d2 = add16(c2, c6); + let d6 = sub16(c2, c6); + let d0 = add16(c0, c4); + let d4 = sub16(c0, c4); + let e5 = _mm256_unpacklo_epi16(d5, d7); + let e7 = _mm256_unpackhi_epi16(d5, d7); + let f5 = _mm256_unpacklo_epi16(e5, e7); + let f7 = _mm256_unpackhi_epi16(e5, e7); + let mut g5 = _mm256_unpacklo_epi16(f5, f7); + let mut g7 = _mm256_unpackhi_epi16(f5, f7); + let e4 = _mm256_unpacklo_epi16(d4, d6); + let e6 = _mm256_unpackhi_epi16(d4, d6); + let f4 = _mm256_unpacklo_epi16(e4, e6); + let f6 = _mm256_unpackhi_epi16(e4, e6); + let mut g4 = _mm256_unpacklo_epi16(f4, f6); + let mut g6 = _mm256_unpackhi_epi16(f4, f6); + let e1 = _mm256_unpacklo_epi16(d1, d3); + let e3 = _mm256_unpackhi_epi16(d1, d3); + let f1 = _mm256_unpacklo_epi16(e1, e3); + let f3 = _mm256_unpackhi_epi16(e1, e3); + let mut g1 = _mm256_unpacklo_epi16(f1, f3); + let mut g3 = _mm256_unpackhi_epi16(f1, f3); + let e0 = _mm256_unpacklo_epi16(d0, d2); + let e2 = _mm256_unpackhi_epi16(d0, d2); + let f0 = _mm256_unpacklo_epi16(e0, e2); + let f2 = _mm256_unpackhi_epi16(e0, e2); + let mut g0 = _mm256_unpacklo_epi16(f0, f2); + let mut g2 = _mm256_unpackhi_epi16(f0, f2); + g7 = mulmod_scaled(g7, ld(qdata, 208), ld(qdata, 992), qv); + g3 = mulmod_scaled(g3, ld(qdata, 192), ld(qdata, 976), qv); + g6 = mulmod_scaled(g6, ld(qdata, 240), ld(qdata, 1024), qv); + g2 = mulmod_scaled(g2, ld(qdata, 224), ld(qdata, 1008), qv); + g5 = mulmod_scaled(g5, ld(qdata, 176), ld(qdata, 960), qv); + g1 = mulmod_scaled(g1, ld(qdata, 160), ld(qdata, 944), qv); + g4 = reduce(g4, qv, qround32v); + g0 = reduce(g0, qv, qround32v); + let h6 = add16(g6, g7); + let mut h7 = sub16(g6, g7); + h7 = mulmod_scaled(h7, ld(qdata, 1648), ld(qdata, 1568), qv); + let h2 = add16(g2, g3); + let mut h3 = sub16(g2, g3); + h3 = mulmod_scaled(h3, ld(qdata, 1648), ld(qdata, 1568), qv); + let h4 = add16(g4, g5); + let h5 = sub16(g4, g5); + let h0 = add16(g0, g1); + let h1 = sub16(g0, g1); + let i5 = add16(h5, h7); + let i7 = sub16(h5, h7); + let i1 = add16(h1, h3); + let i3 = sub16(h1, h3); + let i4 = add16(h4, h6); + let i6 = sub16(h4, h6); + let i0 = add16(h0, h2); + let i2 = sub16(h0, h2); + _mm256_storeu_si256(fp.add(0) as *mut __m256i, i0); + _mm256_storeu_si256(fp.add(16) as *mut __m256i, i1); + _mm256_storeu_si256(fp.add(32) as *mut __m256i, i2); + _mm256_storeu_si256(fp.add(48) as *mut __m256i, i3); + _mm256_storeu_si256(fp.add(64) as *mut __m256i, i4); + _mm256_storeu_si256(fp.add(80) as *mut __m256i, i5); + _mm256_storeu_si256(fp.add(96) as *mut __m256i, i6); + _mm256_storeu_si256(fp.add(112) as *mut __m256i, i7); + } + reps /= 2; + reps /= 2; + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a27 = _mm256_loadu_si256(fp.add(432) as *const __m256i); + let a31 = _mm256_loadu_si256(fp.add(496) as *const __m256i); + let mut b27 = perm_lo(a27, a31); + let mut b31 = perm_hi(a27, a31); + let a19 = _mm256_loadu_si256(fp.add(304) as *const __m256i); + let a23 = _mm256_loadu_si256(fp.add(368) as *const __m256i); + let mut b19 = perm_lo(a19, a23); + let mut b23 = perm_hi(a19, a23); + let a11 = _mm256_loadu_si256(fp.add(176) as *const __m256i); + let a15 = _mm256_loadu_si256(fp.add(240) as *const __m256i); + let mut b11 = perm_lo(a11, a15); + let mut b15 = perm_hi(a11, a15); + let a3 = _mm256_loadu_si256(fp.add(48) as *const __m256i); + let a7 = _mm256_loadu_si256(fp.add(112) as *const __m256i); + let mut b3 = perm_lo(a3, a7); + let mut b7 = perm_hi(a3, a7); + b31 = mulmod_scaled(b31, ld(qdata, 624), ld(qdata, 1408), qv); + b27 = mulmod_scaled(b27, ld(qdata, 560), ld(qdata, 1344), qv); + b23 = mulmod_scaled(b23, ld(qdata, 688), ld(qdata, 1472), qv); + b19 = mulmod_scaled(b19, ld(qdata, 752), ld(qdata, 1536), qv); + b15 = mulmod_scaled(b15, ld(qdata, 432), ld(qdata, 1216), qv); + b11 = mulmod_scaled(b11, ld(qdata, 496), ld(qdata, 1280), qv); + b7 = mulmod_scaled(b7, ld(qdata, 368), ld(qdata, 1152), qv); + b3 = reduce(b3, qv, qround32v); + let c27 = add16(b27, b31); + let mut c31 = sub16(b27, b31); + c31 = mulmod_scaled(c31, ld(qdata, 1664), ld(qdata, 1584), qv); + let c19 = add16(b19, b23); + let mut c23 = sub16(b19, b23); + c23 = mulmod_scaled(c23, ld(qdata, 1680), ld(qdata, 1600), qv); + let c11 = add16(b11, b15); + let mut c15 = sub16(b11, b15); + c15 = mulmod_scaled(c15, ld(qdata, 1648), ld(qdata, 1568), qv); + let c3 = add16(b3, b7); + let c7 = sub16(b3, b7); + let d23 = add16(c23, c31); + let mut d31 = sub16(c23, c31); + d31 = mulmod_scaled(d31, ld(qdata, 1648), ld(qdata, 1568), qv); + let mut d19 = add16(c19, c27); + let mut d27 = sub16(c19, c27); + d27 = mulmod_scaled(d27, ld(qdata, 1648), ld(qdata, 1568), qv); + let d7 = add16(c7, c15); + let d15 = sub16(c7, c15); + let mut d3 = add16(c3, c11); + let d11 = sub16(c3, c11); + d19 = reduce(d19, qv, qround32v); + d3 = reduce(d3, qv, qround32v); + let e15 = add16(d15, d31); + let e31 = sub16(d15, d31); + let e7 = add16(d7, d23); + let e23 = sub16(d7, d23); + let e11 = add16(d11, d27); + let e27 = sub16(d11, d27); + let e3 = add16(d3, d19); + let e19 = sub16(d3, d19); + _mm256_storeu_si256(fp.add(48) as *mut __m256i, e3); + _mm256_storeu_si256(fp.add(112) as *mut __m256i, e7); + _mm256_storeu_si256(fp.add(176) as *mut __m256i, e11); + _mm256_storeu_si256(fp.add(240) as *mut __m256i, e15); + _mm256_storeu_si256(fp.add(304) as *mut __m256i, e19); + _mm256_storeu_si256(fp.add(368) as *mut __m256i, e23); + _mm256_storeu_si256(fp.add(432) as *mut __m256i, e27); + _mm256_storeu_si256(fp.add(496) as *mut __m256i, e31); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a26 = _mm256_loadu_si256(fp.add(416) as *const __m256i); + let a30 = _mm256_loadu_si256(fp.add(480) as *const __m256i); + let mut b26 = perm_lo(a26, a30); + let mut b30 = perm_hi(a26, a30); + let a18 = _mm256_loadu_si256(fp.add(288) as *const __m256i); + let a22 = _mm256_loadu_si256(fp.add(352) as *const __m256i); + let mut b18 = perm_lo(a18, a22); + let mut b22 = perm_hi(a18, a22); + let a10 = _mm256_loadu_si256(fp.add(160) as *const __m256i); + let a14 = _mm256_loadu_si256(fp.add(224) as *const __m256i); + let mut b10 = perm_lo(a10, a14); + let mut b14 = perm_hi(a10, a14); + let a2 = _mm256_loadu_si256(fp.add(32) as *const __m256i); + let a6 = _mm256_loadu_si256(fp.add(96) as *const __m256i); + let mut b2 = perm_lo(a2, a6); + let mut b6 = perm_hi(a2, a6); + b30 = mulmod_scaled(b30, ld(qdata, 608), ld(qdata, 1392), qv); + b26 = mulmod_scaled(b26, ld(qdata, 544), ld(qdata, 1328), qv); + b22 = mulmod_scaled(b22, ld(qdata, 672), ld(qdata, 1456), qv); + b18 = mulmod_scaled(b18, ld(qdata, 736), ld(qdata, 1520), qv); + b14 = mulmod_scaled(b14, ld(qdata, 416), ld(qdata, 1200), qv); + b10 = mulmod_scaled(b10, ld(qdata, 480), ld(qdata, 1264), qv); + b6 = mulmod_scaled(b6, ld(qdata, 352), ld(qdata, 1136), qv); + b2 = reduce(b2, qv, qround32v); + let c26 = add16(b26, b30); + let mut c30 = sub16(b26, b30); + c30 = mulmod_scaled(c30, ld(qdata, 1664), ld(qdata, 1584), qv); + let c18 = add16(b18, b22); + let mut c22 = sub16(b18, b22); + c22 = mulmod_scaled(c22, ld(qdata, 1680), ld(qdata, 1600), qv); + let c10 = add16(b10, b14); + let mut c14 = sub16(b10, b14); + c14 = mulmod_scaled(c14, ld(qdata, 1648), ld(qdata, 1568), qv); + let c2 = add16(b2, b6); + let c6 = sub16(b2, b6); + let d22 = add16(c22, c30); + let mut d30 = sub16(c22, c30); + d30 = mulmod_scaled(d30, ld(qdata, 1648), ld(qdata, 1568), qv); + let mut d18 = add16(c18, c26); + let mut d26 = sub16(c18, c26); + d26 = mulmod_scaled(d26, ld(qdata, 1648), ld(qdata, 1568), qv); + let d6 = add16(c6, c14); + let d14 = sub16(c6, c14); + let mut d2 = add16(c2, c10); + let d10 = sub16(c2, c10); + d18 = reduce(d18, qv, qround32v); + d2 = reduce(d2, qv, qround32v); + let e14 = add16(d14, d30); + let e30 = sub16(d14, d30); + let e6 = add16(d6, d22); + let e22 = sub16(d6, d22); + let e10 = add16(d10, d26); + let e26 = sub16(d10, d26); + let e2 = add16(d2, d18); + let e18 = sub16(d2, d18); + _mm256_storeu_si256(fp.add(32) as *mut __m256i, e2); + _mm256_storeu_si256(fp.add(96) as *mut __m256i, e6); + _mm256_storeu_si256(fp.add(160) as *mut __m256i, e10); + _mm256_storeu_si256(fp.add(224) as *mut __m256i, e14); + _mm256_storeu_si256(fp.add(288) as *mut __m256i, e18); + _mm256_storeu_si256(fp.add(352) as *mut __m256i, e22); + _mm256_storeu_si256(fp.add(416) as *mut __m256i, e26); + _mm256_storeu_si256(fp.add(480) as *mut __m256i, e30); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a25 = _mm256_loadu_si256(fp.add(400) as *const __m256i); + let a29 = _mm256_loadu_si256(fp.add(464) as *const __m256i); + let mut b25 = perm_lo(a25, a29); + let mut b29 = perm_hi(a25, a29); + let a17 = _mm256_loadu_si256(fp.add(272) as *const __m256i); + let a21 = _mm256_loadu_si256(fp.add(336) as *const __m256i); + let mut b17 = perm_lo(a17, a21); + let mut b21 = perm_hi(a17, a21); + let a9 = _mm256_loadu_si256(fp.add(144) as *const __m256i); + let a13 = _mm256_loadu_si256(fp.add(208) as *const __m256i); + let mut b9 = perm_lo(a9, a13); + let mut b13 = perm_hi(a9, a13); + let a1 = _mm256_loadu_si256(fp.add(16) as *const __m256i); + let a5 = _mm256_loadu_si256(fp.add(80) as *const __m256i); + let mut b1 = perm_lo(a1, a5); + let mut b5 = perm_hi(a1, a5); + b29 = mulmod_scaled(b29, ld(qdata, 592), ld(qdata, 1376), qv); + b25 = mulmod_scaled(b25, ld(qdata, 528), ld(qdata, 1312), qv); + b21 = mulmod_scaled(b21, ld(qdata, 656), ld(qdata, 1440), qv); + b17 = mulmod_scaled(b17, ld(qdata, 720), ld(qdata, 1504), qv); + b13 = mulmod_scaled(b13, ld(qdata, 400), ld(qdata, 1184), qv); + b9 = mulmod_scaled(b9, ld(qdata, 464), ld(qdata, 1248), qv); + b5 = mulmod_scaled(b5, ld(qdata, 336), ld(qdata, 1120), qv); + b1 = reduce(b1, qv, qround32v); + let c25 = add16(b25, b29); + let mut c29 = sub16(b25, b29); + c29 = mulmod_scaled(c29, ld(qdata, 1664), ld(qdata, 1584), qv); + let c17 = add16(b17, b21); + let mut c21 = sub16(b17, b21); + c21 = mulmod_scaled(c21, ld(qdata, 1680), ld(qdata, 1600), qv); + let c9 = add16(b9, b13); + let mut c13 = sub16(b9, b13); + c13 = mulmod_scaled(c13, ld(qdata, 1648), ld(qdata, 1568), qv); + let c1 = add16(b1, b5); + let c5 = sub16(b1, b5); + let d21 = add16(c21, c29); + let mut d29 = sub16(c21, c29); + d29 = mulmod_scaled(d29, ld(qdata, 1648), ld(qdata, 1568), qv); + let mut d17 = add16(c17, c25); + let mut d25 = sub16(c17, c25); + d25 = mulmod_scaled(d25, ld(qdata, 1648), ld(qdata, 1568), qv); + let d5 = add16(c5, c13); + let d13 = sub16(c5, c13); + let mut d1 = add16(c1, c9); + let d9 = sub16(c1, c9); + d17 = reduce(d17, qv, qround32v); + d1 = reduce(d1, qv, qround32v); + let e13 = add16(d13, d29); + let e29 = sub16(d13, d29); + let e5 = add16(d5, d21); + let e21 = sub16(d5, d21); + let e9 = add16(d9, d25); + let e25 = sub16(d9, d25); + let e1 = add16(d1, d17); + let e17 = sub16(d1, d17); + _mm256_storeu_si256(fp.add(16) as *mut __m256i, e1); + _mm256_storeu_si256(fp.add(80) as *mut __m256i, e5); + _mm256_storeu_si256(fp.add(144) as *mut __m256i, e9); + _mm256_storeu_si256(fp.add(208) as *mut __m256i, e13); + _mm256_storeu_si256(fp.add(272) as *mut __m256i, e17); + _mm256_storeu_si256(fp.add(336) as *mut __m256i, e21); + _mm256_storeu_si256(fp.add(400) as *mut __m256i, e25); + _mm256_storeu_si256(fp.add(464) as *mut __m256i, e29); + } + for r in 0..reps { + let fp = f.as_mut_ptr().add(base + 512 * r); + let a24 = _mm256_loadu_si256(fp.add(384) as *const __m256i); + let a28 = _mm256_loadu_si256(fp.add(448) as *const __m256i); + let mut b24 = perm_lo(a24, a28); + let mut b28 = perm_hi(a24, a28); + let a16 = _mm256_loadu_si256(fp.add(256) as *const __m256i); + let a20 = _mm256_loadu_si256(fp.add(320) as *const __m256i); + let mut b16 = perm_lo(a16, a20); + let mut b20 = perm_hi(a16, a20); + let a8 = _mm256_loadu_si256(fp.add(128) as *const __m256i); + let a12 = _mm256_loadu_si256(fp.add(192) as *const __m256i); + let mut b8 = perm_lo(a8, a12); + let mut b12 = perm_hi(a8, a12); + let a0 = _mm256_loadu_si256(fp.add(0) as *const __m256i); + let a4 = _mm256_loadu_si256(fp.add(64) as *const __m256i); + let mut b0 = perm_lo(a0, a4); + let mut b4 = perm_hi(a0, a4); + b28 = mulmod_scaled(b28, ld(qdata, 576), ld(qdata, 1360), qv); + b24 = mulmod_scaled(b24, ld(qdata, 512), ld(qdata, 1296), qv); + b20 = mulmod_scaled(b20, ld(qdata, 640), ld(qdata, 1424), qv); + b16 = mulmod_scaled(b16, ld(qdata, 704), ld(qdata, 1488), qv); + b12 = mulmod_scaled(b12, ld(qdata, 384), ld(qdata, 1168), qv); + b8 = mulmod_scaled(b8, ld(qdata, 448), ld(qdata, 1232), qv); + b4 = mulmod_scaled(b4, ld(qdata, 320), ld(qdata, 1104), qv); + b0 = reduce(b0, qv, qround32v); + let c24 = add16(b24, b28); + let mut c28 = sub16(b24, b28); + c28 = mulmod_scaled(c28, ld(qdata, 1664), ld(qdata, 1584), qv); + let c16 = add16(b16, b20); + let mut c20 = sub16(b16, b20); + c20 = mulmod_scaled(c20, ld(qdata, 1680), ld(qdata, 1600), qv); + let c8 = add16(b8, b12); + let mut c12 = sub16(b8, b12); + c12 = mulmod_scaled(c12, ld(qdata, 1648), ld(qdata, 1568), qv); + let c0 = add16(b0, b4); + let c4 = sub16(b0, b4); + let d20 = add16(c20, c28); + let mut d28 = sub16(c20, c28); + d28 = mulmod_scaled(d28, ld(qdata, 1648), ld(qdata, 1568), qv); + let mut d16 = add16(c16, c24); + let mut d24 = sub16(c16, c24); + d24 = mulmod_scaled(d24, ld(qdata, 1648), ld(qdata, 1568), qv); + let d4 = add16(c4, c12); + let d12 = sub16(c4, c12); + let mut d0 = add16(c0, c8); + let d8 = sub16(c0, c8); + d16 = reduce(d16, qv, qround32v); + d0 = reduce(d0, qv, qround32v); + let e12 = add16(d12, d28); + let e28 = sub16(d12, d28); + let e4 = add16(d4, d20); + let e20 = sub16(d4, d20); + let e8 = add16(d8, d24); + let e24 = sub16(d8, d24); + let e0 = add16(d0, d16); + let e16 = sub16(d0, d16); + _mm256_storeu_si256(fp.add(0) as *mut __m256i, e0); + _mm256_storeu_si256(fp.add(64) as *mut __m256i, e4); + _mm256_storeu_si256(fp.add(128) as *mut __m256i, e8); + _mm256_storeu_si256(fp.add(192) as *mut __m256i, e12); + _mm256_storeu_si256(fp.add(256) as *mut __m256i, e16); + _mm256_storeu_si256(fp.add(320) as *mut __m256i, e20); + _mm256_storeu_si256(fp.add(384) as *mut __m256i, e24); + _mm256_storeu_si256(fp.add(448) as *mut __m256i, e28); + } + } +} + +#[rustfmt::skip] +static QDATA_7681: [i16; 1696] = [ + -3593, -3593, -3593, -3593, -3625, -3625, -3625, -3625, -3593, -3593, -3593, -3593, -3625, -3625, -3625, -3625, + -3777, -3777, -3777, -3777, 3182, 3182, 3182, 3182, -3777, -3777, -3777, -3777, 3182, 3182, 3182, 3182, + -3593, -3593, -3593, -3593, -3182, -3182, -3182, -3182, -3593, -3593, -3593, -3593, -3182, -3182, -3182, -3182, + 3777, 3777, 3777, 3777, 3625, 3625, 3625, 3625, 3777, 3777, 3777, 3777, 3625, 3625, 3625, 3625, + -3593, -3593, -3593, -3593, 2194, 2194, 2194, 2194, -3593, -3593, -3593, -3593, 2194, 2194, 2194, 2194, + -3625, -3625, -3625, -3625, -1100, -1100, -1100, -1100, -3625, -3625, -3625, -3625, -1100, -1100, -1100, -1100, + -3593, -3593, -3593, -3593, 3696, 3696, 3696, 3696, -3593, -3593, -3593, -3593, 3696, 3696, 3696, 3696, + -3182, -3182, -3182, -3182, -2456, -2456, -2456, -2456, -3182, -3182, -3182, -3182, -2456, -2456, -2456, -2456, + -3593, 1701, 2194, 834, -3625, 2319, -1100, 121, -3593, 1701, 2194, 834, -3625, 2319, -1100, 121, + -3777, 1414, 2456, 2495, 3182, 2876, -3696, 2250, -3777, 1414, 2456, 2495, 3182, 2876, -3696, 2250, + -3593, -2250, 3696, -2876, -3182, -2495, -2456, -1414, -3593, -2250, 3696, -2876, -3182, -2495, -2456, -1414, + 3777, -121, 1100, -2319, 3625, -834, -2194, -1701, 3777, -121, 1100, -2319, 3625, -834, -2194, -1701, + -3593, 3364, 1701, -1599, 2194, 2557, 834, -2816, -3593, 3364, 1701, -1599, 2194, 2557, 834, -2816, + -3625, 617, 2319, 2006, -1100, -1296, 121, 1986, -3625, 617, 2319, 2006, -1100, -1296, 121, 1986, + -3593, 2237, -2250, -1483, 3696, 3706, -2876, 1921, -3593, 2237, -2250, -1483, 3696, 3706, -2876, 1921, + -3182, 2088, -2495, -1525, -2456, 1993, -1414, 2830, -3182, 2088, -2495, -1525, -2456, 1993, -1414, 2830, + -3593, 514, 3364, 438, 1701, 2555, -1599, -1738, 2194, 103, 2557, 1881, 834, -549, -2816, 638, + -3625, -1399, 617, -1760, 2319, 2535, 2006, 3266, -1100, -1431, -1296, 3174, 121, 3153, 1986, -810, + -3777, 2956, -2830, -679, 1414, 2440, -1993, -3689, 2456, 2804, 1525, 3555, 2495, 1535, -2088, -7, + 3182, -1321, -1921, -1305, 2876, -3772, -3706, 3600, -3696, -2043, 1483, -396, 2250, -2310, -2237, 1887, + -3593, -1887, 2237, 2310, -2250, 396, -1483, 2043, 3696, -3600, 3706, 3772, -2876, 1305, 1921, 1321, + -3182, 7, 2088, -1535, -2495, -3555, -1525, -2804, -2456, 3689, 1993, -2440, -1414, 679, 2830, -2956, + 3777, 810, -1986, -3153, -121, -3174, 1296, 1431, 1100, -3266, -2006, -2535, -2319, 1760, -617, 1399, + 3625, -638, 2816, 549, -834, -1881, -2557, -103, -2194, 1738, 1599, -2555, -1701, -438, -3364, -514, + -3593, -1532, 514, -373, 3364, -3816, 438, -3456, 1701, 783, 2555, 2883, -1599, 727, -1738, -2385, + 2194, -2160, 103, -2391, 2557, 2762, 1881, -2426, 834, 3310, -549, -1350, -2816, 1386, 638, -194, + -3625, 404, -1399, -3692, 617, -2764, -1760, -1054, 2319, 1799, 2535, -3588, 2006, 1533, 3266, 2113, + -1100, -2579, -1431, -1756, -1296, 1598, 3174, -2, 121, -3480, 3153, -2572, 1986, 2743, -810, 2919, + -3593, 2789, -1887, -921, 2237, -1497, 2310, -2133, -2250, -915, 396, 1390, -1483, 3135, 2043, -859, + 3696, 2732, -3600, -1464, 3706, 2224, 3772, -2665, -2876, 1698, 1305, 2835, 1921, 730, 1321, 486, + -3182, 3417, 7, -3428, 2088, -3145, -1535, 1168, -2495, -3831, -3555, -3750, -1525, 660, -2804, 2649, + -2456, 3405, 3689, -1521, 1993, 1681, -2440, 1056, -1414, 1166, 679, -2233, 2830, 2175, -2956, -1919, + -3593, -1404, -1532, 451, 514, -402, -373, 1278, 3364, -509, -3816, -3770, 438, -2345, -3456, -226, + 1701, -1689, 783, -1509, 2555, 2963, 2883, 1242, -1599, 1669, 727, 2719, -1738, 642, -2385, -436, + 2194, 3335, -2160, 1779, 103, 3745, -2391, 17, 2557, 2812, 2762, -1144, 1881, 83, -2426, -1181, + 834, -1519, 3310, 3568, -549, -796, -1350, 2072, -2816, -2460, 1386, 2891, 638, -2083, -194, -715, + -3593, -402, -3816, -226, 2555, 1669, -2385, 1779, 2557, 83, 3310, 2072, 638, 1012, -3692, 1295, + 2319, -3208, 1533, -2071, -1431, -2005, -2, 1586, 1986, -293, 1919, -929, -679, 777, -1681, -3461, + 2456, 3366, 3750, -1203, 1535, -3657, -3417, -1712, -1921, 2515, 2665, -1070, 3600, 2532, -3135, -2589, + 2250, -2258, 921, -658, -514, 509, 3456, 1509, 1599, -642, 2160, -17, -1881, 1519, 1350, -2891, + -3593, -3434, -1497, 893, 396, -2422, -859, 2965, 3706, -2339, 1698, -2937, 1321, -670, -3428, -3163, + -2495, -1072, 660, 1084, 3689, -179, 1056, -1338, 2830, 2786, -2919, -3677, -3153, -151, -1598, 3334, + 1100, -3314, 3588, 2262, 1760, -2230, -404, 2083, 2816, -3568, 2426, -2812, -103, 436, -727, -2963, + -1701, 3770, 373, 1404, 1887, -1649, 2133, -826, 1483, 434, -2732, 3287, -3772, -2378, -2835, 3723, + -3593, 658, 2789, 370, -1887, -3434, -921, -3752, 2237, 1649, -1497, 2258, 2310, 3581, -2133, 893, + -2250, 3794, -915, 826, 396, 2589, 1390, 592, -1483, -2422, 3135, 3214, 2043, -434, -859, -2532, + 3696, 1121, 2732, 2965, -3600, 2998, -1464, -3287, 3706, 1070, 2224, -589, 3772, -2339, -2665, 2070, + -2876, 2378, 1698, -2515, 1305, -2815, 2835, -2937, 1921, -1348, 730, -3723, 1321, 1712, 486, 2130, + 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, 7681, + -9, -9, -9, -9, -16425, -16425, -16425, -16425, -9, -9, -9, -9, -16425, -16425, -16425, -16425, + -28865, -28865, -28865, -28865, 10350, 10350, 10350, 10350, -28865, -28865, -28865, -28865, 10350, 10350, 10350, 10350, + -9, -9, -9, -9, -10350, -10350, -10350, -10350, -9, -9, -9, -9, -10350, -10350, -10350, -10350, + 28865, 28865, 28865, 28865, 16425, 16425, 16425, 16425, 28865, 28865, 28865, 28865, 16425, 16425, 16425, 16425, + -9, -9, -9, -9, -4974, -4974, -4974, -4974, -9, -9, -9, -9, -4974, -4974, -4974, -4974, + -16425, -16425, -16425, -16425, -7244, -7244, -7244, -7244, -16425, -16425, -16425, -16425, -7244, -7244, -7244, -7244, + -9, -9, -9, -9, -4496, -4496, -4496, -4496, -9, -9, -9, -9, -4496, -4496, -4496, -4496, + -10350, -10350, -10350, -10350, -14744, -14744, -14744, -14744, -10350, -10350, -10350, -10350, -14744, -14744, -14744, -14744, + -9, -20315, -4974, 18242, -16425, 18191, -7244, -11655, -9, -20315, -4974, 18242, -16425, 18191, -7244, -11655, + -28865, 20870, 14744, -22593, 10350, 828, 4496, 23754, -28865, 20870, 14744, -22593, 10350, 828, 4496, 23754, + -9, -23754, -4496, -828, -10350, 22593, -14744, -20870, -9, -23754, -4496, -828, -10350, 22593, -14744, -20870, + 28865, 11655, 7244, -18191, 16425, -18242, 4974, 20315, 28865, 11655, 7244, -18191, 16425, -18242, 4974, 20315, + -9, -10972, -20315, 23489, -4974, 25597, 18242, -2816, -9, -10972, -20315, 23489, -4974, 25597, 18242, -2816, + -16425, -19351, 18191, -3114, -7244, -9488, -11655, 19394, -16425, -19351, 18191, -3114, -7244, -9488, -11655, 19394, + -9, -7491, -23754, -15307, -4496, -15750, -828, -5759, -9, -7491, -23754, -15307, -4496, -15750, -828, -5759, + -10350, 22568, 22593, -20469, -14744, 31177, -20870, 26382, -10350, 22568, 22593, -20469, -14744, 31177, -20870, 26382, + -9, -14846, -10972, -21066, -20315, -24581, 23489, -23242, -4974, -4505, 25597, -26279, 18242, 21467, -2816, 15998, + -16425, -4983, -19351, 14624, 18191, -2073, -3114, 20674, -7244, -21399, -9488, 6246, -11655, -29103, 19394, -5930, + -28865, -23668, -26382, -28839, 20870, 6536, -31177, 16279, 14744, 29428, 20469, 29667, -22593, 9215, -22568, -11783, + 10350, -14121, 5759, -5913, 828, -1724, 15750, 11792, 4496, 25093, 15307, 26228, 23754, -21766, 7491, -6817, + -9, 6817, -7491, 21766, -23754, -26228, -15307, -25093, -4496, -11792, -15750, 1724, -828, 5913, -5759, 14121, + -10350, 11783, 22568, -9215, 22593, -29667, -20469, -29428, -14744, -16279, 31177, -6536, -20870, 28839, 26382, 23668, + 28865, 5930, -19394, 29103, 11655, -6246, 9488, 21399, 7244, -20674, 3114, 2073, -18191, -14624, 19351, 4983, + 16425, -15998, 2816, -21467, -18242, 26279, -25597, 4505, 4974, 23242, -23489, 24581, 20315, 21066, 10972, 14846, + -9, -32252, -14846, -19317, -10972, 8472, -21066, -3456, -20315, 16655, -24581, 12611, 23489, -12073, -23242, 29871, + -4974, 6032, -4505, 10409, 25597, 24266, -26279, 17030, 18242, 10478, 21467, 11962, -2816, -26262, 15998, -17602, + -16425, -22124, -4983, -26220, -19351, -8908, 14624, 32738, 18191, 13575, -2073, 27132, -3114, 24573, 20674, 27201, + -7244, 12269, -21399, -16092, -9488, -15810, 6246, 15358, -11655, -15768, -29103, 24052, 19394, -26441, -5930, -1689, + -9, 13541, 6817, -5529, -7491, 26663, 21766, -4693, -23754, 13933, -26228, 8558, -15307, -21953, -25093, -22875, + -4496, -7508, -11792, -30136, -15750, 26800, 1724, 17303, -828, 2722, 5913, -12013, -5759, 30426, 14121, 3558, + -10350, -24743, 11783, -21860, 22568, -32329, -9215, 9360, 22593, -7415, -29667, 25946, -20469, -21868, -29428, -25511, + -14744, 1869, -16279, 14351, 31177, 2193, -6536, 17440, -20870, 24718, 28839, -23225, 26382, 9855, 23668, -9599, + -9, -32124, -32252, 10179, -14846, 6766, -19317, 16638, -10972, -23549, 8472, -17082, -21066, -15145, -3456, 31518, + -20315, -6297, 16655, -12261, -24581, -11885, 12611, 30938, 23489, 28805, -12073, 26783, -23242, -14718, 29871, 5708, + -4974, 15111, 6032, -29453, -4505, 12449, 10409, 529, 25597, -32004, 24266, 2952, -26279, 18003, 17030, 24931, + 18242, -1007, 10478, -4624, 21467, 17636, 11962, 14360, -2816, 15972, -26262, 16715, 15998, 4573, -17602, -14539, + -9, 6766, 8472, 31518, -24581, 28805, 29871, -29453, 25597, 18003, 10478, 14360, 15998, 27636, -26220, 17167, + 18191, -7304, 24573, -22039, -21399, -4565, 15358, 10802, 19394, 21723, 9599, -9633, -28839, -2807, -2193, -30597, + 14744, -26330, -25946, -2739, 9215, 32695, 24743, -26288, 5759, 20435, -17303, 24530, 11792, 20964, 21953, 23523, + 23754, -27858, 5529, 6510, 14846, 23549, 3456, 12261, -23489, 14718, -6032, -529, 26279, 1007, -11962, -16715, + -9, 24214, 26663, 23933, -26228, -13686, -22875, -27243, -15750, 4317, 2722, 8839, 14121, -32414, -21860, -25179, + 22593, -25648, -21868, -964, -16279, -1715, 17440, -14650, 26382, -28958, 1689, -10333, 29103, -20119, 15810, 22790, + 7244, 20238, -27132, -2858, -14624, 19274, 22124, -4573, 2816, 4624, -17030, 32004, 4505, -5708, 12073, 11885, + 20315, 17082, 19317, 32124, -6817, 14223, 4693, -14138, 15307, 9650, 7508, -9513, -1724, -23882, 12013, -15221, + -9, -6510, 13541, -23182, 6817, 24214, -5529, -24232, -7491, -14223, 26663, 27858, 21766, 26621, -4693, 23933, + -23754, 29394, 13933, 14138, -26228, -23523, 8558, -23984, -15307, -13686, -21953, 26766, -25093, -9650, -22875, -20964, + -4496, -22943, -7508, -27243, -11792, -18506, -30136, 9513, -15750, -24530, 26800, 947, 1724, 4317, 17303, 29718, + -828, 23882, 2722, -20435, 5913, -10495, -12013, 8839, -5759, -3396, 30426, 15221, 14121, 26288, 3558, 27730, + -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, -28865, + 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, 28865, + -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, -16425, + -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, -10350, + 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, + -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, -3777, + 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, 3777, + -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, -3625, + -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, -3182, +]; + +#[rustfmt::skip] +static QDATA_10753: [i16; 1696] = [ + 1018, 1018, 1018, 1018, 3688, 3688, 3688, 3688, 1018, 1018, 1018, 1018, 3688, 3688, 3688, 3688, + -223, -223, -223, -223, -4188, -4188, -4188, -4188, -223, -223, -223, -223, -4188, -4188, -4188, -4188, + 1018, 1018, 1018, 1018, 4188, 4188, 4188, 4188, 1018, 1018, 1018, 1018, 4188, 4188, 4188, 4188, + 223, 223, 223, 223, -3688, -3688, -3688, -3688, 223, 223, 223, 223, -3688, -3688, -3688, -3688, + 1018, 1018, 1018, 1018, -376, -376, -376, -376, 1018, 1018, 1018, 1018, -376, -376, -376, -376, + 3688, 3688, 3688, 3688, -3686, -3686, -3686, -3686, 3688, 3688, 3688, 3688, -3686, -3686, -3686, -3686, + 1018, 1018, 1018, 1018, -2413, -2413, -2413, -2413, 1018, 1018, 1018, 1018, -2413, -2413, -2413, -2413, + 4188, 4188, 4188, 4188, -357, -357, -357, -357, 4188, 4188, 4188, 4188, -357, -357, -357, -357, + 1018, -3364, -376, 4855, 3688, 425, -3686, 2695, 1018, -3364, -376, 4855, 3688, 425, -3686, 2695, + -223, -3784, 357, -2236, -4188, 4544, 2413, 730, -223, -3784, 357, -2236, -4188, 4544, 2413, 730, + 1018, -730, -2413, -4544, 4188, 2236, -357, 3784, 1018, -730, -2413, -4544, 4188, 2236, -357, 3784, + 223, -2695, 3686, -425, -3688, -4855, 376, 3364, 223, -2695, 3686, -425, -3688, -4855, 376, 3364, + 1018, -5175, -3364, 2503, -376, 1341, 4855, -4875, 1018, -5175, -3364, 2503, -376, 1341, 4855, -4875, + 3688, -2629, 425, -4347, -3686, 3823, 2695, -4035, 3688, -2629, 425, -4347, -3686, 3823, 2695, -4035, + 1018, 5063, -730, 341, -2413, -3012, -4544, -5213, 1018, 5063, -730, 341, -2413, -3012, -4544, -5213, + 4188, 1520, 2236, 1931, -357, 918, 3784, 4095, 4188, 1520, 2236, 1931, -357, 918, 3784, 4095, + 1018, 3085, -5175, 2982, -3364, -4744, 2503, -4129, -376, -2576, 1341, -193, 4855, 3062, -4875, 4, + 3688, 2388, -2629, -4513, 425, 4742, -4347, 2935, -3686, -544, 3823, -2178, 2695, 847, -4035, 268, + -223, -1299, -4095, -1287, -3784, -4876, -918, 3091, 357, -4189, -1931, 4616, -2236, 2984, -1520, -3550, + -4188, -1009, 5213, -205, 4544, -4102, 3012, 2790, 2413, -1085, -341, -2565, 730, -4379, -5063, -1284, + 1018, 1284, 5063, 4379, -730, 2565, 341, 1085, -2413, -2790, -3012, 4102, -4544, 205, -5213, 1009, + 4188, 3550, 1520, -2984, 2236, -4616, 1931, 4189, -357, -3091, 918, 4876, 3784, 1287, 4095, 1299, + 223, -268, 4035, -847, -2695, 2178, -3823, 544, 3686, -2935, 4347, -4742, -425, 4513, 2629, -2388, + -3688, -4, 4875, -3062, -4855, 193, -1341, 2576, 376, 4129, -2503, 4744, 3364, -2982, 5175, -3085, + 1018, 5116, 3085, -3615, -5175, 400, 2982, 3198, -3364, 2234, -4744, -4828, 2503, 326, -4129, -512, + -376, 1068, -2576, -4580, 1341, 3169, -193, -2998, 4855, -635, 3062, -4808, -4875, -2740, 4, 675, + 3688, -1324, 2388, 5114, -2629, 5294, -4513, -794, 425, -864, 4742, -886, -4347, 336, 2935, -2045, + -3686, -3715, -544, 4977, 3823, -2737, -2178, 3441, 2695, 467, 847, 454, -4035, -779, 268, 2213, + 1018, 1615, 1284, 2206, 5063, 5064, 4379, 472, -730, -5341, 2565, -4286, 341, 2981, 1085, -1268, + -2413, -3057, -2790, -2884, -3012, -1356, 4102, -3337, -4544, 5023, 205, -636, -5213, 909, 1009, -2973, + 4188, 2271, 3550, -1572, 1520, 1841, -2984, 970, 2236, -4734, -4616, 578, 1931, -116, 4189, 1586, + -357, -2774, -3091, -1006, 918, -5156, 4876, 4123, 3784, -567, 1287, 151, 4095, 1458, 1299, 2684, + 1018, -3260, 5116, -1722, 3085, 5120, -3615, 3760, -5175, 73, 400, 4254, 2982, 2788, 3198, -2657, + -3364, 569, 2234, 1930, -4744, -2279, -4828, 5215, 2503, -4403, 326, 1639, -4129, 5068, -512, -5015, + -376, -4859, 1068, -40, -2576, 4003, -4580, -4621, 1341, 2487, 3169, -2374, -193, 2625, -2998, 4784, + 4855, 825, -635, 2118, 3062, -2813, -4808, -4250, -4875, -2113, -2740, -4408, 4, -1893, 675, 458, + 1018, 5120, 400, -2657, -4744, -4403, -512, -40, 1341, 2625, -635, -4250, 4, -3360, 5114, -5313, + 425, -2151, 336, -2662, -544, 5334, 3441, 2117, -4035, 2205, -2684, -3570, -1287, -4973, 5156, 2419, + 357, 1204, -578, 1635, 2984, -1111, -2271, 4359, 5213, -2449, 3337, 3453, 2790, 554, -2981, -1409, + 730, -279, -2206, 3524, -3085, -73, -3198, -1930, -2503, -5068, -1068, 4621, 193, -825, 4808, 4408, + 1018, 4428, 5064, -4000, 2565, 573, -1268, 3125, -3012, -4144, 5023, 1927, 1009, -2139, -1572, 3535, + 2236, 663, -116, 4967, -3091, -854, 4123, 1160, 4095, -1349, -2213, 1782, -847, 2062, 2737, 624, + 3686, -2283, 886, 4889, 4513, -4601, 1324, 1893, 4875, -2118, 2998, -2487, 2576, 5015, -326, 2279, + 3364, -4254, 3615, 3260, -1284, -1381, -472, -3891, -341, 2087, 3057, 4720, -4102, 3410, 636, 1689, + 1018, -3524, 1615, 5268, 1284, 4428, 2206, -834, 5063, 1381, 5064, 279, 4379, 2439, 472, -4000, + -730, -2015, -5341, 3891, 2565, 1409, -4286, 2605, 341, 573, 2981, 5356, 1085, -2087, -1268, -554, + -2413, 3135, -3057, 3125, -2790, -778, -2884, -4720, -3012, -3453, -1356, -355, 4102, -4144, -3337, -152, + -4544, -3410, 5023, 2449, 205, -97, -636, 1927, -5213, 2624, 909, -1689, 1009, -4359, -2973, -3419, + 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, 10753, + -6, -6, -6, -6, -408, -408, -408, -408, -6, -6, -6, -6, -408, -408, -408, -408, + -27359, -27359, -27359, -27359, 1956, 1956, 1956, 1956, -27359, -27359, -27359, -27359, 1956, 1956, 1956, 1956, + -6, -6, -6, -6, -1956, -1956, -1956, -1956, -6, -6, -6, -6, -1956, -1956, -1956, -1956, + 27359, 27359, 27359, 27359, 408, 408, 408, 408, 27359, 27359, 27359, 27359, 408, 408, 408, 408, + -6, -6, -6, -6, -20856, -20856, -20856, -20856, -6, -6, -6, -6, -20856, -20856, -20856, -20856, + -408, -408, -408, -408, -21094, -21094, -21094, -21094, -408, -408, -408, -408, -21094, -21094, -21094, -21094, + -6, -6, -6, -6, -10093, -10093, -10093, -10093, -6, -6, -6, -6, -10093, -10093, -10093, -10093, + -1956, -1956, -1956, -1956, -28517, -28517, -28517, -28517, -1956, -1956, -1956, -1956, -28517, -28517, -28517, -28517, + -6, -9508, -20856, -29449, -408, 18345, -21094, -7033, -6, -9508, -20856, -29449, -408, 18345, -21094, -7033, + -27359, -16072, 28517, -12476, 1956, -28224, 10093, 16090, -27359, -16072, 28517, -12476, 1956, -28224, 10093, 16090, + -6, -16090, -10093, 28224, -1956, 12476, -28517, 16072, -6, -16090, -10093, 28224, -1956, 12476, -28517, 16072, + 27359, 7033, 21094, -18345, 408, 29449, 20856, 9508, 27359, 7033, 21094, -18345, 408, 29449, 20856, 9508, + -6, -3639, -9508, 25543, -20856, 829, -29449, -17675, -6, -3639, -9508, 25543, -20856, 829, -29449, -17675, + -408, 18363, 18345, 7429, -21094, -10001, -7033, -4547, -408, 18363, 18345, 7429, -21094, -10001, -7033, -4547, + -6, 28103, -16090, 3925, -10093, 7228, 28224, 11683, -6, 28103, -16090, 3925, -10093, 7228, 28224, 11683, + -1956, -23056, 12476, 14731, -28517, 26518, 16072, 14847, -1956, -23056, 12476, 14731, -28517, 26518, 16072, 14847, + -6, -5619, -3639, -12378, -9508, 15736, 25543, 23007, -20856, -27152, 829, -22209, -29449, -20490, -17675, 22532, + -408, 16724, 18363, 22623, 18345, 5766, 7429, -31369, -21094, 15840, -10001, 19326, -7033, 3407, -4547, 2316, + -27359, 6381, -14847, 8441, -16072, -6924, -26518, -4589, 28517, 12707, -14731, -15864, -12476, 31656, 23056, 24098, + 1956, -31217, -11683, -24269, -28224, -5126, -7228, 20198, 10093, -573, -3925, -14341, 16090, 23781, -28103, -23812, + -6, 23812, 28103, -23781, -16090, 14341, 3925, 573, -10093, -20198, 7228, 5126, 28224, 24269, 11683, 31217, + -1956, -24098, -23056, -31656, 12476, 15864, 14731, -12707, -28517, 4589, 26518, 6924, 16072, -8441, 14847, -6381, + 27359, -2316, 4547, -3407, 7033, -19326, 10001, -15840, 21094, 31369, -7429, -5766, -18345, -22623, -18363, -16724, + 408, -22532, 17675, 20490, 29449, 22209, -829, 27152, 20856, -23007, -25543, -15736, 9508, 12378, 3639, 5619, + -6, -17412, -5619, 2017, -3639, 24976, -12378, 24702, -9508, -31558, 15736, 1316, 25543, -31418, 23007, -512, + -20856, -13268, -27152, 22044, 829, 8801, -22209, -12214, -29449, 11141, -20490, -17096, -17675, 32076, 22532, 17571, + -408, 13012, 16724, 4090, 18363, -30546, 22623, 16614, 18345, -17248, 5766, 22666, 7429, -7856, -31369, 31235, + -21094, 28541, 15840, -30351, -10001, -177, 19326, -31887, -7033, 25555, 3407, -31290, -4547, -13579, 2316, -2395, + -6, 4175, 23812, 7326, 28103, 17352, -23781, -28200, -16090, 11555, 14341, 6978, 3925, -1627, 573, 780, + -10093, 32271, -20198, 7356, 7228, 29364, 5126, 27895, 28224, -609, 24269, 21892, 11683, -7795, 31217, -18845, + -1956, 29407, -24098, -7716, -23056, -719, -31656, -8246, 12476, -26238, 15864, 11842, 14731, 1932, -12707, -11726, + -28517, 4394, 4589, 2066, 26518, -11300, 6924, -24037, 16072, 969, -8441, 14999, 14847, -11854, -6381, -19844, + -6, -13500, -17412, 32070, -5619, 5120, 2017, 11952, -3639, 1609, 24976, 9374, -12378, -23836, 24702, -8289, + -9508, -22471, -31558, 25482, 15736, -8935, 1316, 32351, 25543, 19661, -31418, 8295, 23007, -25652, -512, -19863, + -20856, 6917, -13268, -28712, -27152, 20899, 22044, 4083, 829, 951, 8801, 29370, -22209, 24641, -12214, 12976, + -29449, -22215, 11141, -29626, -20490, 30467, -17096, 13158, -17675, -24129, 32076, 7880, 22532, -30053, 17571, -8758, + -6, 5120, 24976, -8289, 15736, 19661, -512, -28712, 829, 24641, 11141, 13158, 22532, 13024, 4090, -27329, + 18345, -8807, -7856, -20070, 15840, -1834, -31887, -18875, -4547, 18077, 19844, -23026, 8441, -12653, 11300, 11123, + 28517, 31924, -11842, -14237, 31656, 16809, -29407, -5369, -11683, -16273, -27895, -29827, 20198, 7722, 1627, 9343, + 16090, -15127, -7326, -6716, 5619, -1609, -24702, -25482, -25543, 25652, 13268, -4083, 22209, 22215, 17096, -7880, + -6, -26292, 17352, 12384, 14341, 61, 780, 23093, 7228, -12336, -609, -7801, 31217, -6747, -7716, 6095, + 12476, 15511, 1932, 11623, 4589, 6314, -24037, -19320, 14847, 19643, 2395, -21770, -3407, -17394, 177, -23952, + 21094, -31467, -22666, -1767, -22623, -14329, -13012, 30053, 17675, 29626, 12214, -951, 27152, 19863, 31418, 8935, + 9508, -9374, -2017, 13500, -23812, -29541, 28200, 20173, -3925, -24025, -32271, -19856, -5126, -26286, -21892, -4967, + -6, 6716, 4175, -13164, 23812, -26292, 7326, -12098, 28103, 29541, 17352, 15127, -23781, -7289, -28200, 12384, + -16090, -29151, 11555, -20173, 14341, -9343, 6978, -22483, 3925, 61, -1627, 23788, 573, 24025, 780, -7722, + -10093, -18881, 32271, 23093, -20198, -24330, 7356, 19856, 7228, 29827, 29364, 15517, 5126, -12336, 27895, -4248, + 28224, 26286, -609, 16273, 24269, -5729, 21892, -7801, 11683, -30144, -7795, 4967, 31217, 5369, -18845, -8027, + -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, -27359, + 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, 27359, + -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, -408, + -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, -1956, + 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, + -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, -223, + 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, 223, + 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, 3688, + 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, 4188, +]; diff --git a/sntrup-kem/src/rq/vector.rs b/sntrup-kem/src/rq/vector.rs index d63ff7a..199d403 100644 --- a/sntrup-kem/src/rq/vector.rs +++ b/sntrup-kem/src/rq/vector.rs @@ -10,14 +10,12 @@ use crate::rq::modq; #[inline(always)] #[allow(clippy::cast_possible_truncation)] pub fn swap(x: &mut [i16], y: &mut [i16], n: usize, mask: isize) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return swap_avx2(x, y, n, mask); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return swap_avx2(x, y, n, mask); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -38,11 +36,7 @@ fn swap_scalar(x: &mut [i16], y: &mut [i16], n: usize, mask: isize) { } } -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] unsafe fn swap_avx2(x: &mut [i16], y: &mut [i16], n: usize, mask: isize) { unsafe { @@ -104,18 +98,16 @@ pub fn product(z: &mut [i16], n: usize, x: &[i16], c: i16, q: i32, b1: i32, b2: } } -/// Fused minus_product and shift: `z[i+1] = freeze(z[i] - y[i]*c)`, `z[0] = 0`. +/// Fused minus_product and shift: `z[i+1] = freeze(z[i] - y[i]*c), z[0] = 0`. /// Processes backward to avoid overwrite conflicts, eliminating a separate memmove. #[inline(always)] pub fn minus_product_shift(z: &mut [i16], n: usize, y: &[i16], c: i16, q: i32, b1: i32, b2: i32) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return minus_product_shift_avx2(z, n, y, c, q, b1, b2); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return minus_product_shift_avx2(z, n, y, c, q, b1, b2); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -141,11 +133,177 @@ fn minus_product_shift_scalar( z[0] = 0; } -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +/// Fused `minus_product_shift` + conditional swap: one memory pass instead of two. +/// +/// Semantics are exactly `minus_product_shift(z, n, y, c, ..)` followed by +/// `swap(z, y, n, mask)` — the caller derives `mask` from the scalar-computed +/// post-shift leading coefficient before invoking this. On x86_64 with AVX2 the +/// fused kernel performs 3 loads + 2 stores per block where the two-pass form +/// performs 4 loads + 3 stores; other targets (including aarch64/NEON, whose +/// kernels are intentionally untouched) fall back to the two-pass form. +#[inline(always)] +pub fn minus_product_shift_cswap( + z: &mut [i16], + y: &mut [i16], + n: usize, + c: i16, + mask: isize, + params: &crate::params::SntrupParameters, +) { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return minus_product_shift_cswap_avx2(z, y, n, c, mask, params); + } + } + minus_product_shift(z, n, y, c, params.q, params.barrett1, params.barrett2); + swap(z, y, n, mask); +} + +/// AVX2 fused kernel: the signed-Montgomery shift (see `minus_product_shift_avx2`) +/// with the conditional swap applied in-register via `blendv` before storing. +/// +/// The `+1`-shifted store means each block also loads `y[start+1..start+17]` for +/// the swap half. In the bottom overlapped block, the topmost of those lanes +/// (index 16) was already written by the previous block, so that lane is +/// preserved through a keep-mask blend against current memory instead of being +/// recomputed from a stale input. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx2")] +unsafe fn minus_product_shift_cswap_avx2( + z: &mut [i16], + y: &mut [i16], + n: usize, + c: i16, + mask: isize, + params: &crate::params::SntrupParameters, +) { + unsafe { + use core::arch::x86_64::*; + + let q = params.q; + let b1 = params.barrett1; + let b2 = params.barrett2; + + // Montgomery constants — see minus_product_shift_avx2 for the derivation. + let qw = q as u16; + let mut qinv = qw; + for _ in 0..3 { + qinv = qinv.wrapping_mul(2u16.wrapping_sub(qw.wrapping_mul(qinv))); + } + let cp = modq::freeze( + (modq::freeze((c as i32) << 8, q, b1, b2) as i32) << 8, + q, + b1, + b2, + ); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let cpqinv = (cp as u16).wrapping_mul(qinv) as i16; + + let cpv = _mm256_set1_epi16(cp); + let cpqv = _mm256_set1_epi16(cpqinv); + #[allow(clippy::cast_possible_truncation)] + let qv = _mm256_set1_epi16(q as i16); + #[allow(clippy::cast_possible_truncation)] + let hqv = _mm256_set1_epi16(((q - 1) / 2) as i16); + #[allow(clippy::cast_possible_truncation)] + let nhqv = _mm256_set1_epi16((-(q - 1) / 2) as i16); + let mv = _mm256_set1_epi16(mask as i16); + + let mut j = (n - 2) as isize; + + while j >= 15 { + let start = (j - 15) as usize; + let zv = _mm256_loadu_si256(z.as_ptr().add(start) as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr().add(start) as *const __m256i); + + let m = _mm256_mullo_epi16(yv, cpqv); + let t = _mm256_sub_epi16(_mm256_mulhi_epi16(yv, cpv), _mm256_mulhi_epi16(m, qv)); + let w = _mm256_sub_epi16(zv, t); + let gt = _mm256_cmpgt_epi16(w, hqv); + let lt = _mm256_cmpgt_epi16(nhqv, w); + let w = _mm256_sub_epi16(w, _mm256_and_si256(gt, qv)); + let w = _mm256_add_epi16(w, _mm256_and_si256(lt, qv)); + + // Conditional swap against y at the shifted (+1) position. + let y1 = _mm256_loadu_si256(y.as_ptr().add(start + 1) as *const __m256i); + let new_z = _mm256_blendv_epi8(w, y1, mv); + let new_y = _mm256_blendv_epi8(y1, w, mv); + _mm256_storeu_si256(z.as_mut_ptr().add(start + 1) as *mut __m256i, new_z); + _mm256_storeu_si256(y.as_mut_ptr().add(start + 1) as *mut __m256i, new_y); + j -= 16; + } + + // Bottom overlapped block (see minus_product_shift_avx2 for the coverage + // argument). Inputs at [0..16) are still original; the +1 loads' topmost + // lane (index 16) is post-swap, so it is preserved, not recomputed. + if j >= 0 && n >= 17 && n & 15 == 0 { + let zv = _mm256_loadu_si256(z.as_ptr() as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr() as *const __m256i); + let m = _mm256_mullo_epi16(yv, cpqv); + let t = _mm256_sub_epi16(_mm256_mulhi_epi16(yv, cpv), _mm256_mulhi_epi16(m, qv)); + let w = _mm256_sub_epi16(zv, t); + let gt = _mm256_cmpgt_epi16(w, hqv); + let lt = _mm256_cmpgt_epi16(nhqv, w); + let w = _mm256_sub_epi16(w, _mm256_and_si256(gt, qv)); + let w = _mm256_add_epi16(w, _mm256_and_si256(lt, qv)); + + let z1 = _mm256_loadu_si256(z.as_ptr().add(1) as *const __m256i); + let y1 = _mm256_loadu_si256(y.as_ptr().add(1) as *const __m256i); + let new_z = _mm256_blendv_epi8(w, y1, mv); + let new_y = _mm256_blendv_epi8(y1, w, mv); + // Keep the topmost i16 lane (bytes 30-31) from current memory. + let keep = _mm256_setr_epi8( + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, -1, -1, + ); + _mm256_storeu_si256( + z.as_mut_ptr().add(1) as *mut __m256i, + _mm256_blendv_epi8(new_z, z1, keep), + ); + _mm256_storeu_si256( + y.as_mut_ptr().add(1) as *mut __m256i, + _mm256_blendv_epi8(new_y, y1, keep), + ); + j = -1; + } + + // Scalar remainder (only when n < 17). + #[allow(clippy::cast_possible_truncation)] + let mi = mask as i16; + while j >= 0 { + let k = (j + 1) as usize; + let w = modq::minus_product(z[k - 1], y[k - 1], c, q, b1, b2); + let yk = y[k]; + z[k] = (mi & yk) | (!mi & w); + y[k] = (mi & w) | (!mi & yk); + j -= 1; + } + let y0 = y[0]; + z[0] = mi & y0; + y[0] = !mi & y0; + } +} + +/// AVX2 kernel in the 16-bit domain via signed Montgomery multiplication +/// (Seiler, "Faster Kyber" — the same shape Kyber's AVX2 code uses). +/// +/// Instead of widening to i32 lanes and running the two-step Barrett chain +/// (five `vpmulld` per 8 elements), each 16-lane block computes +/// `z - y·c mod± q` directly in i16 lanes: +/// +/// - `c' = c·2^16 mod± q` (scalar, two Barrett freezes so the i32 window holds), +/// so the Montgomery product `y ⊗ c' = y·c'·2^-16 ≡ y·c (mod q)` is exact. +/// - `m = mullo(y, c'·q^-1 mod 2^16)`, `t = mulhi(y, c') − mulhi(m, q)`: +/// `t ≡ y·c (mod± q)` with `|t| < q` — three 16-bit multiplies per 16 lanes. +/// - `w = z − t` lies in `(−3q/2, 3q/2)`, so a single branchless +/// compare-and-correct lands it in the canonical `[−(q−1)/2, (q−1)/2]`. +/// +/// All lane operations are branchless; the scalar precomputation uses the same +/// constant-time `freeze` and wrapping arithmetic, so the constant-time +/// property is unchanged. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] unsafe fn minus_product_shift_avx2( z: &mut [i16], @@ -158,74 +316,80 @@ unsafe fn minus_product_shift_avx2( ) { unsafe { use core::arch::x86_64::*; - let qv = _mm256_set1_epi32(q); - let kb1 = _mm256_set1_epi32(b1); - let kb2 = _mm256_set1_epi32(b2); - let k134m = _mm256_set1_epi32(134_217_728); - let cv = _mm256_set1_epi32(c as i32); + + // q^-1 mod 2^16 by Newton iteration (q is odd and public): each step + // doubles the number of correct low bits, 3 -> 6 -> 12 -> 24 >= 16. + let qw = q as u16; + let mut qinv = qw; + for _ in 0..3 { + qinv = qinv.wrapping_mul(2u16.wrapping_sub(qw.wrapping_mul(qinv))); + } + + // c' = c·2^16 mod± q, in two freezes of c·2^8 so each input stays well + // inside the Barrett window (|c·2^8| <= ~1M << 2^31 / barrett1). + let cp = modq::freeze( + (modq::freeze((c as i32) << 8, q, b1, b2) as i32) << 8, + q, + b1, + b2, + ); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let cpqinv = (cp as u16).wrapping_mul(qinv) as i16; + + let cpv = _mm256_set1_epi16(cp); + let cpqv = _mm256_set1_epi16(cpqinv); + #[allow(clippy::cast_possible_truncation)] + let qv = _mm256_set1_epi16(q as i16); + #[allow(clippy::cast_possible_truncation)] + let hqv = _mm256_set1_epi16(((q - 1) / 2) as i16); + #[allow(clippy::cast_possible_truncation)] + let nhqv = _mm256_set1_epi16((-(q - 1) / 2) as i16); let mut j = (n - 2) as isize; - // Process 16 at a time (two 8-wide batches for ILP), backward + // Process 16 i16 elements per iteration, backward. while j >= 15 { let start = (j - 15) as usize; + let zv = _mm256_loadu_si256(z.as_ptr().add(start) as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr().add(start) as *const __m256i); + + // t = y·c'·2^-16 mod± q = y·c mod± q, |t| < q. + let m = _mm256_mullo_epi16(yv, cpqv); + let t = _mm256_sub_epi16(_mm256_mulhi_epi16(yv, cpv), _mm256_mulhi_epi16(m, qv)); + + // w = z - t, then one branchless correction into canonical range. + let w = _mm256_sub_epi16(zv, t); + let gt = _mm256_cmpgt_epi16(w, hqv); + let lt = _mm256_cmpgt_epi16(nhqv, w); + let w = _mm256_sub_epi16(w, _mm256_and_si256(gt, qv)); + let w = _mm256_add_epi16(w, _mm256_and_si256(lt, qv)); - // Batch 0: elements start..start+8 - let zv0 = - _mm256_cvtepi16_epi32(_mm_loadu_si128(z.as_ptr().add(start) as *const __m128i)); - let yv0 = - _mm256_cvtepi16_epi32(_mm_loadu_si128(y.as_ptr().add(start) as *const __m128i)); - let a0 = _mm256_sub_epi32(zv0, _mm256_mullo_epi32(yv0, cv)); - - // Batch 1: elements start+8..start+16 - let zv1 = - _mm256_cvtepi16_epi32(_mm_loadu_si128(z.as_ptr().add(start + 8) as *const __m128i)); - let yv1 = - _mm256_cvtepi16_epi32(_mm_loadu_si128(y.as_ptr().add(start + 8) as *const __m128i)); - let a1 = _mm256_sub_epi32(zv1, _mm256_mullo_epi32(yv1, cv)); - - // Barrett freeze batch 0 - let t0 = _mm256_srai_epi32(_mm256_mullo_epi32(a0, kb1), 20); - let b0 = _mm256_sub_epi32(a0, _mm256_mullo_epi32(t0, qv)); - let t0 = _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b0, kb2), k134m), 28); - let r0 = _mm256_sub_epi32(b0, _mm256_mullo_epi32(t0, qv)); - - // Barrett freeze batch 1 - let t1 = _mm256_srai_epi32(_mm256_mullo_epi32(a1, kb1), 20); - let b1 = _mm256_sub_epi32(a1, _mm256_mullo_epi32(t1, qv)); - let t1 = _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b1, kb2), k134m), 28); - let r1 = _mm256_sub_epi32(b1, _mm256_mullo_epi32(t1, qv)); - - // Pack 8+8 i32 -> 16 i16 and store at offset +1 (the shift) - let packed = _mm256_permute4x64_epi64(_mm256_packs_epi32(r0, r1), 0xD8); - _mm256_storeu_si256(z.as_mut_ptr().add(start + 1) as *mut __m256i, packed); + // Store at offset +1 (the shift). + _mm256_storeu_si256(z.as_mut_ptr().add(start + 1) as *mut __m256i, w); j -= 16; } - // Process remaining 8 at a time - while j >= 7 { - let start = (j - 7) as usize; - let zv = - _mm256_cvtepi16_epi32(_mm_loadu_si128(z.as_ptr().add(start) as *const __m128i)); - let yv = - _mm256_cvtepi16_epi32(_mm_loadu_si128(y.as_ptr().add(start) as *const __m128i)); - let a = _mm256_sub_epi32(zv, _mm256_mullo_epi32(yv, cv)); - - let t = _mm256_srai_epi32(_mm256_mullo_epi32(a, kb1), 20); - let b = _mm256_sub_epi32(a, _mm256_mullo_epi32(t, qv)); - let t = _mm256_srai_epi32(_mm256_add_epi32(_mm256_mullo_epi32(b, kb2), k134m), 28); - let r = _mm256_sub_epi32(b, _mm256_mullo_epi32(t, qv)); - - let lo = _mm256_castsi256_si128(r); - let hi = _mm256_extracti128_si256(r, 1); - _mm_storeu_si128( - z.as_mut_ptr().add(start + 1) as *mut __m128i, - _mm_packs_epi32(lo, hi), - ); - j -= 8; + // The backward loop strands `(n - 2) % 16` bottom elements. When n is a + // multiple of 16 and the body ran at least once, a final full-width block at + // start = 0 covers them: + // z[0..16] is still original (higher blocks only wrote z[16..] and beyond), + // and the overlap element it rewrites (z[16]) gets the identical value the + // previous block computed from the same inputs. + if j >= 0 && n >= 17 && n & 15 == 0 { + let zv = _mm256_loadu_si256(z.as_ptr() as *const __m256i); + let yv = _mm256_loadu_si256(y.as_ptr() as *const __m256i); + let m = _mm256_mullo_epi16(yv, cpqv); + let t = _mm256_sub_epi16(_mm256_mulhi_epi16(yv, cpv), _mm256_mulhi_epi16(m, qv)); + let w = _mm256_sub_epi16(zv, t); + let gt = _mm256_cmpgt_epi16(w, hqv); + let lt = _mm256_cmpgt_epi16(nhqv, w); + let w = _mm256_sub_epi16(w, _mm256_and_si256(gt, qv)); + let w = _mm256_add_epi16(w, _mm256_and_si256(lt, qv)); + _mm256_storeu_si256(z.as_mut_ptr().add(1) as *mut __m256i, w); + j = -1; } - // Scalar remainder + // Scalar remainder (only when n < 17) while j >= 0 { z[(j + 1) as usize] = modq::minus_product(z[j as usize], y[j as usize], c, q, b1, b2); j -= 1; @@ -252,6 +416,11 @@ unsafe fn minus_product_shift_neon( let kb2 = vdupq_n_s32(b2); let k134m = vdupq_n_s32(134_217_728); let cv = vdupq_n_s32(c as i32); + // Strict-canonical correction bounds — see `modq::freeze`: the two Barrett steps + // alone can land a few counts outside ±(q-1)/2, and every freeze path must match + // the scalar reference byte-for-byte. + let hqv = vdupq_n_s32((q - 1) >> 1); + let nhqv = vdupq_n_s32(-((q - 1) >> 1)); let mut j = (n - 2) as isize; @@ -269,17 +438,27 @@ unsafe fn minus_product_shift_neon( let yv1 = vmovl_s16(vld1_s16(y.as_ptr().add(start + 4))); let a1 = vsubq_s32(zv1, vmulq_s32(yv1, cv)); - // Barrett freeze batch 0 + // Barrett freeze batch 0, with strict-canonical correction let t0 = vshrq_n_s32(vmulq_s32(a0, kb1), 20); let b0 = vsubq_s32(a0, vmulq_s32(t0, qv)); let t0 = vshrq_n_s32(vaddq_s32(vmulq_s32(b0, kb2), k134m), 28); let r0 = vsubq_s32(b0, vmulq_s32(t0, qv)); + let r0 = vsubq_s32(r0, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(r0, hqv)), qv)); + let r0 = vaddq_s32( + r0, + vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(nhqv, r0)), qv), + ); - // Barrett freeze batch 1 + // Barrett freeze batch 1, with strict-canonical correction let t1 = vshrq_n_s32(vmulq_s32(a1, kb1), 20); let b1 = vsubq_s32(a1, vmulq_s32(t1, qv)); let t1 = vshrq_n_s32(vaddq_s32(vmulq_s32(b1, kb2), k134m), 28); let r1 = vsubq_s32(b1, vmulq_s32(t1, qv)); + let r1 = vsubq_s32(r1, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(r1, hqv)), qv)); + let r1 = vaddq_s32( + r1, + vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(nhqv, r1)), qv), + ); // Pack 4+4 i32 -> 8 i16 (naturally ordered, no permute needed) let packed = vcombine_s16(vmovn_s32(r0), vmovn_s32(r1)); @@ -298,6 +477,8 @@ unsafe fn minus_product_shift_neon( let b = vsubq_s32(a, vmulq_s32(t, qv)); let t = vshrq_n_s32(vaddq_s32(vmulq_s32(b, kb2), k134m), 28); let r = vsubq_s32(b, vmulq_s32(t, qv)); + let r = vsubq_s32(r, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(r, hqv)), qv)); + let r = vaddq_s32(r, vandq_s32(vreinterpretq_s32_u32(vcgtq_s32(nhqv, r)), qv)); vst1_s16(z.as_mut_ptr().add(start + 1), vmovn_s32(r)); j -= 4; @@ -311,3 +492,535 @@ unsafe fn minus_product_shift_neon( z[0] = 0; } } + +/// Divstep elimination pass over `f[1..]`/`g[1..]` (SUPERCOP `vectormodq_swapeliminate`). +/// +/// Per 16-lane block: conditional swap of f/g by `mask`, then +/// `g_new = (f0·g − g0·f)·2⁻¹⁶ mod± q` via two signed-Montgomery products, +/// stored one position *down* (the divstep `/x` — inputs are reversed, so the +/// eliminated "leading" coefficient is the constant term). Processes +/// `len.next_multiple_of(16)` elements: every operation is lanewise (plus the +/// fixed −1 store shift), so overrun lanes never contaminate lanes inside the +/// true window; callers provide `1 + len.next_multiple_of(16)` capacity. +/// +/// `mask` is applied as a broadcast blend; no branch, index, or bound depends +/// on secret data. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx2")] +unsafe fn swapeliminate_avx2( + f: &mut [i16], + g: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + unsafe { + use core::arch::x86_64::*; + + let qw = q as u16; + let mut qinv = qw; + for _ in 0..3 { + qinv = qinv.wrapping_mul(2u16.wrapping_sub(qw.wrapping_mul(qinv))); + } + #[allow(clippy::cast_possible_truncation)] + let qv = _mm256_set1_epi16(q as i16); + let f0v = _mm256_set1_epi16(f0); + let g0v = _mm256_set1_epi16(g0); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let f0qinv = _mm256_set1_epi16((f0 as u16).wrapping_mul(qinv) as i16); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let g0qinv = _mm256_set1_epi16((g0 as u16).wrapping_mul(qinv) as i16); + let mv = _mm256_set1_epi16(mask as i16); + + let mut k = 0usize; + let blocks = len.div_ceil(16); + while k < blocks * 16 { + let fi = _mm256_loadu_si256(f.as_ptr().add(1 + k) as *const __m256i); + let gi = _mm256_loadu_si256(g.as_ptr().add(1 + k) as *const __m256i); + let fnew = _mm256_blendv_epi8(fi, gi, mv); + let gnew = _mm256_blendv_epi8(gi, fi, mv); + // (f0·g_new − g0·f_new)·2⁻¹⁶, each product |.| < q so the difference + // stays well inside i16. + let a = _mm256_sub_epi16( + _mm256_mulhi_epi16(gnew, f0v), + _mm256_mulhi_epi16(_mm256_mullo_epi16(gnew, f0qinv), qv), + ); + let b = _mm256_sub_epi16( + _mm256_mulhi_epi16(fnew, g0v), + _mm256_mulhi_epi16(_mm256_mullo_epi16(fnew, g0qinv), qv), + ); + let gout = _mm256_sub_epi16(a, b); + _mm256_storeu_si256(f.as_mut_ptr().add(1 + k) as *mut __m256i, fnew); + _mm256_storeu_si256(g.as_mut_ptr().add(k) as *mut __m256i, gout); + k += 16; + } + } +} + +/// Divstep Bezout-side pass over `v`/`r` (SUPERCOP `vectormodq_xswapeliminate`). +/// +/// Same elimination as [`swapeliminate`], but `v` is stored one position *up* +/// (multiply by x) and `r` in place, iterating backward so the +1-shifted +/// stores never clobber unread inputs. Capacity contract as above. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx2")] +unsafe fn xswapeliminate_avx2( + v: &mut [i16], + r: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + unsafe { + use core::arch::x86_64::*; + + let qw = q as u16; + let mut qinv = qw; + for _ in 0..3 { + qinv = qinv.wrapping_mul(2u16.wrapping_sub(qw.wrapping_mul(qinv))); + } + #[allow(clippy::cast_possible_truncation)] + let qv = _mm256_set1_epi16(q as i16); + let f0v = _mm256_set1_epi16(f0); + let g0v = _mm256_set1_epi16(g0); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let f0qinv = _mm256_set1_epi16((f0 as u16).wrapping_mul(qinv) as i16); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let g0qinv = _mm256_set1_epi16((g0 as u16).wrapping_mul(qinv) as i16); + let mv = _mm256_set1_epi16(mask as i16); + + // Descending traversal is required: the `v` store is shifted up by one, + // so an ascending pass would clobber the next block's first input. A + // forward variant that carries that input in a register measured 19% + // slower on key generation when this was the dispatched kernel — the + // extra per-iteration branch and loop-carried register dependency cost + // more than any prefetch gain. + let mut k = (len.div_ceil(16) * 16) as isize; + while k > 0 { + k -= 16; + let ku = k as usize; + let vi = _mm256_loadu_si256(v.as_ptr().add(ku) as *const __m256i); + let ri = _mm256_loadu_si256(r.as_ptr().add(ku) as *const __m256i); + let vnew = _mm256_blendv_epi8(vi, ri, mv); + let rnew = _mm256_blendv_epi8(ri, vi, mv); + let a = _mm256_sub_epi16( + _mm256_mulhi_epi16(rnew, f0v), + _mm256_mulhi_epi16(_mm256_mullo_epi16(rnew, f0qinv), qv), + ); + let b = _mm256_sub_epi16( + _mm256_mulhi_epi16(vnew, g0v), + _mm256_mulhi_epi16(_mm256_mullo_epi16(vnew, g0qinv), qv), + ); + let rout = _mm256_sub_epi16(a, b); + _mm256_storeu_si256(v.as_mut_ptr().add(ku + 1) as *mut __m256i, vnew); + _mm256_storeu_si256(r.as_mut_ptr().add(ku) as *mut __m256i, rout); + } + } +} + +/// AVX-512 form of [`swapeliminate_avx2`]: 32 coefficients per iteration. +/// +/// Neither PQClean nor liboqs has a 512-bit path for this KEM, and the divstep +/// elimination passes dominate key generation, so doubling the lane count was +/// worth roughly 12% of the whole operation. `mask` is all-ones or all-zero, so it maps +/// directly onto a `__mmask32` and the blend needs no vector constant at all. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx512f,avx512bw,avx512vl")] +#[allow(unsafe_code)] +unsafe fn swapeliminate_avx512( + f: &mut [i16], + g: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + unsafe { + use core::arch::x86_64::*; + + let (qv, f0v, g0v, f0qinv, g0qinv) = avx512_operands(f0, g0, q); + // -1 selects the swapped operand in every lane, 0 selects neither. + // `__mmask32` is `u32`; the all-ones/all-zero mask maps straight onto it. + #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] + let km: __mmask32 = mask as u32; + + let mut k = 0usize; + let blocks = len.div_ceil(32); + while k < blocks * 32 { + let fi = _mm512_loadu_si512(f.as_ptr().add(1 + k).cast()); + let gi = _mm512_loadu_si512(g.as_ptr().add(1 + k).cast()); + let fnew = _mm512_mask_blend_epi16(km, fi, gi); + let gnew = _mm512_mask_blend_epi16(km, gi, fi); + let a = _mm512_sub_epi16( + _mm512_mulhi_epi16(gnew, f0v), + _mm512_mulhi_epi16(_mm512_mullo_epi16(gnew, f0qinv), qv), + ); + let b = _mm512_sub_epi16( + _mm512_mulhi_epi16(fnew, g0v), + _mm512_mulhi_epi16(_mm512_mullo_epi16(fnew, g0qinv), qv), + ); + let gout = _mm512_sub_epi16(a, b); + _mm512_storeu_si512(f.as_mut_ptr().add(1 + k).cast(), fnew); + _mm512_storeu_si512(g.as_mut_ptr().add(k).cast(), gout); + k += 32; + } + } +} + +/// AVX-512 form of [`xswapeliminate_avx2`]: 32 coefficients per iteration, +/// descending for the same reason the 256-bit pass is. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx512f,avx512bw,avx512vl")] +#[allow(unsafe_code)] +unsafe fn xswapeliminate_avx512( + v: &mut [i16], + r: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + unsafe { + use core::arch::x86_64::*; + + let (qv, f0v, g0v, f0qinv, g0qinv) = avx512_operands(f0, g0, q); + // `__mmask32` is `u32`; the all-ones/all-zero mask maps straight onto it. + #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] + let km: __mmask32 = mask as u32; + + let mut k = (len.div_ceil(32) * 32) as isize; + while k > 0 { + k -= 32; + let ku = k as usize; + let vi = _mm512_loadu_si512(v.as_ptr().add(ku).cast()); + let ri = _mm512_loadu_si512(r.as_ptr().add(ku).cast()); + let vnew = _mm512_mask_blend_epi16(km, vi, ri); + let rnew = _mm512_mask_blend_epi16(km, ri, vi); + let a = _mm512_sub_epi16( + _mm512_mulhi_epi16(rnew, f0v), + _mm512_mulhi_epi16(_mm512_mullo_epi16(rnew, f0qinv), qv), + ); + let b = _mm512_sub_epi16( + _mm512_mulhi_epi16(vnew, g0v), + _mm512_mulhi_epi16(_mm512_mullo_epi16(vnew, g0qinv), qv), + ); + let rout = _mm512_sub_epi16(a, b); + _mm512_storeu_si512(v.as_mut_ptr().add(ku + 1).cast(), vnew); + _mm512_storeu_si512(r.as_mut_ptr().add(ku).cast(), rout); + } + } +} + +/// Broadcast Montgomery operands shared by the two 512-bit divstep kernels: +/// `q^-1 mod 2^16` by Newton iteration, then `q`, `f0`, `g0` and the two +/// pre-multiplied inverses. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx512f,avx512bw,avx512vl")] +#[allow(unsafe_code, clippy::cast_possible_truncation, clippy::cast_sign_loss)] +fn avx512_operands( + f0: i16, + g0: i16, + q: i32, +) -> ( + core::arch::x86_64::__m512i, + core::arch::x86_64::__m512i, + core::arch::x86_64::__m512i, + core::arch::x86_64::__m512i, + core::arch::x86_64::__m512i, +) { + use core::arch::x86_64::*; + let qw = q as u16; + let mut qinv = qw; + for _ in 0..3 { + qinv = qinv.wrapping_mul(2u16.wrapping_sub(qw.wrapping_mul(qinv))); + } + ( + _mm512_set1_epi16(q as i16), + _mm512_set1_epi16(f0), + _mm512_set1_epi16(g0), + _mm512_set1_epi16((f0 as u16).wrapping_mul(qinv) as i16), + _mm512_set1_epi16((g0 as u16).wrapping_mul(qinv) as i16), + ) +} + +/// Dispatching entry point for the divstep elimination pass over `f`/`g`. +/// +/// See [`swapeliminate_avx2`] for the semantics. The three kernels differ only +/// in width: AVX-512 takes 32 coefficients per step, AVX2 sixteen and NEON +/// eight. All produce identical output; `reciprocal3` keeps the pre-divstep +/// elimination algorithm as both the fallback and the differential oracle. +#[cfg(all( + any(target_arch = "x86_64", target_arch = "aarch64"), + not(feature = "force-scalar") +))] +#[inline(always)] +pub fn swapeliminate( + f: &mut [i16], + g: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + #[cfg(target_arch = "x86_64")] + { + if crate::cpu::has_avx512() { + // SAFETY: AVX-512 F/BW/VL confirmed present at runtime. + unsafe { swapeliminate_avx512(f, g, len, f0, g0, mask, q) } + } else { + // SAFETY: callers reach this only when has_avx2() is true. + unsafe { swapeliminate_avx2(f, g, len, f0, g0, mask, q) } + } + } + #[cfg(target_arch = "aarch64")] + // SAFETY: NEON is baseline on aarch64. + unsafe { + swapeliminate_neon(f, g, len, f0, g0, mask, q); + } +} + +/// Dispatching entry point for the Bezout-side divstep pass over `v`/`r`, with +/// the same width dispatch as [`swapeliminate`]. +#[cfg(all( + any(target_arch = "x86_64", target_arch = "aarch64"), + not(feature = "force-scalar") +))] +#[inline(always)] +pub fn xswapeliminate( + v: &mut [i16], + r: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + #[cfg(target_arch = "x86_64")] + { + if crate::cpu::has_avx512() { + // SAFETY: AVX-512 F/BW/VL confirmed present at runtime. + unsafe { xswapeliminate_avx512(v, r, len, f0, g0, mask, q) } + } else { + // SAFETY: callers reach this only when has_avx2() is true. + unsafe { xswapeliminate_avx2(v, r, len, f0, g0, mask, q) } + } + } + #[cfg(target_arch = "aarch64")] + // SAFETY: NEON is baseline on aarch64. + unsafe { + xswapeliminate_neon(v, r, len, f0, g0, mask, q); + } +} + +/// Montgomery constants shared by the NEON divstep kernels: `q^-1 mod 2^16` by +/// Newton iteration, plus the broadcast operands. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] +#[inline(always)] +fn neon_qinv(q: i32) -> u16 { + let qw = q as u16; + let mut qinv = qw; + for _ in 0..3 { + qinv = qinv.wrapping_mul(2u16.wrapping_sub(qw.wrapping_mul(qinv))); + } + qinv +} + +/// Signed high-half product of two i16 vectors — NEON has no single `mulhi`, +/// so widen with `vmull`/`vmull_high` and take the odd (high) halves with +/// `vuzp2q`. Three instructions, exact. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] +#[inline(always)] +unsafe fn mulhi_neon( + a: core::arch::aarch64::int16x8_t, + b: core::arch::aarch64::int16x8_t, +) -> core::arch::aarch64::int16x8_t { + unsafe { + use core::arch::aarch64::*; + let lo = vmull_s16(vget_low_s16(a), vget_low_s16(b)); + let hi = vmull_high_s16(a, b); + vuzp2q_s16(vreinterpretq_s16_s32(lo), vreinterpretq_s16_s32(hi)) + } +} + +/// `x·y·2^-16 mod± q` (signed Montgomery), with `yqinv = y·q^-1 mod 2^16`. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] +#[inline(always)] +unsafe fn montproduct_neon( + x: core::arch::aarch64::int16x8_t, + y: core::arch::aarch64::int16x8_t, + yqinv: core::arch::aarch64::int16x8_t, + qv: core::arch::aarch64::int16x8_t, +) -> core::arch::aarch64::int16x8_t { + unsafe { + use core::arch::aarch64::*; + let b = mulhi_neon(x, y); + let d = vmulq_s16(x, yqinv); + let e = mulhi_neon(d, qv); + vsubq_s16(b, e) + } +} + +/// NEON divstep elimination over `f[1..]`/`g[1..]`, eight lanes per block. +/// Mirrors [`swapeliminate_avx2`]: conditional swap by `mask`, then +/// `g_new = (f0·g − g0·f)·2^-16 mod± q` stored one position down. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] +unsafe fn swapeliminate_neon( + f: &mut [i16], + g: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + unsafe { + use core::arch::aarch64::*; + + let qinv = neon_qinv(q); + #[allow(clippy::cast_possible_truncation)] + let qv = vdupq_n_s16(q as i16); + let f0v = vdupq_n_s16(f0); + let g0v = vdupq_n_s16(g0); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let f0qinv = vdupq_n_s16((f0 as u16).wrapping_mul(qinv) as i16); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let g0qinv = vdupq_n_s16((g0 as u16).wrapping_mul(qinv) as i16); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let mv = vreinterpretq_u16_s16(vdupq_n_s16(mask as i16)); + + let mut k = 0usize; + let blocks = len.div_ceil(8); + while k < blocks * 8 { + let fi = vld1q_s16(f.as_ptr().add(1 + k)); + let gi = vld1q_s16(g.as_ptr().add(1 + k)); + let fnew = vbslq_s16(mv, gi, fi); + let gnew = vbslq_s16(mv, fi, gi); + let a = montproduct_neon(gnew, f0v, f0qinv, qv); + let b = montproduct_neon(fnew, g0v, g0qinv, qv); + let gout = vsubq_s16(a, b); + vst1q_s16(f.as_mut_ptr().add(1 + k), fnew); + vst1q_s16(g.as_mut_ptr().add(k), gout); + k += 8; + } + } +} + +/// NEON Bezout-side divstep pass. Mirrors [`xswapeliminate_avx2`]: `v` stores +/// one position up, `r` in place, iterating backward so the shifted stores +/// never clobber unread inputs. +#[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] +unsafe fn xswapeliminate_neon( + v: &mut [i16], + r: &mut [i16], + len: usize, + f0: i16, + g0: i16, + mask: isize, + q: i32, +) { + unsafe { + use core::arch::aarch64::*; + + let qinv = neon_qinv(q); + #[allow(clippy::cast_possible_truncation)] + let qv = vdupq_n_s16(q as i16); + let f0v = vdupq_n_s16(f0); + let g0v = vdupq_n_s16(g0); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let f0qinv = vdupq_n_s16((f0 as u16).wrapping_mul(qinv) as i16); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let g0qinv = vdupq_n_s16((g0 as u16).wrapping_mul(qinv) as i16); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let mv = vreinterpretq_u16_s16(vdupq_n_s16(mask as i16)); + + #[allow(clippy::cast_possible_wrap)] + let mut k = (len.div_ceil(8) * 8) as isize; + while k > 0 { + k -= 8; + #[allow(clippy::cast_sign_loss)] + let ku = k as usize; + let vi = vld1q_s16(v.as_ptr().add(ku)); + let ri = vld1q_s16(r.as_ptr().add(ku)); + let vnew = vbslq_s16(mv, ri, vi); + let rnew = vbslq_s16(mv, vi, ri); + let a = montproduct_neon(rnew, f0v, f0qinv, qv); + let b = montproduct_neon(vnew, g0v, g0qinv, qv); + let rout = vsubq_s16(a, b); + vst1q_s16(v.as_mut_ptr().add(ku + 1), vnew); + vst1q_s16(r.as_mut_ptr().add(ku), rout); + } + } +} + +#[cfg(test)] +#[allow(clippy::cast_possible_truncation)] +mod tests { + use super::*; + + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + /// The fused kernel must match scalar minus_product_shift + scalar swap + /// exactly, for both mask values, at every parameter set's q and at + /// lengths that exercise the vector loop, the overlapped bottom block, + /// and the scalar path. + #[test] + fn fused_cswap_matches_two_pass_reference() { + let mut s = 0x1234_5678_9abc_def1u64; + for &(q, b1, b2) in &[ + (4621i32, 226i32, 58084i32), + (4591, 228, 58464), + (5167, 202, 51948), + (6343, 165, 42324), + (7177, 146, 37410), + (7879, 133, 34073), + ] { + let params = crate::params::SntrupParameters { + p: 0, + q, + w: 0, + q12: (q - 1) / 2, + small_encode_size: 0, + rounded_encode_size: 0, + pk_size: 0, + sk_size: 0, + ct_size: 0, + barrett1: b1, + barrett2: b2, + }; + let hq = ((q - 1) / 2) as i16; + for &n in &[2usize, 5, 16, 17, 33, 762, 768, 1524] { + for &mask in &[0isize, -1] { + let sample = + |s: &mut u64| ((next(s) % (2 * hq as u64 + 1)) as i32 - hq as i32) as i16; + let z0: Vec = (0..n).map(|_| sample(&mut s)).collect(); + let y0: Vec = (0..n).map(|_| sample(&mut s)).collect(); + let c: i16 = sample(&mut s); + + let mut z_ref = z0.clone(); + let mut y_ref = y0.clone(); + minus_product_shift_scalar(&mut z_ref, n, &y_ref, c, q, b1, b2); + swap_scalar(&mut z_ref, &mut y_ref, n, mask); + + let mut z_got = z0.clone(); + let mut y_got = y0.clone(); + minus_product_shift_cswap(&mut z_got, &mut y_got, n, c, mask, ¶ms); + + assert_eq!(z_got, z_ref, "z mismatch q={q} n={n} mask={mask}"); + assert_eq!(y_got, y_ref, "y mismatch q={q} n={n} mask={mask}"); + } + } + } + } +} diff --git a/sntrup-kem/src/scratch.rs b/sntrup-kem/src/scratch.rs new file mode 100644 index 0000000..024bffb --- /dev/null +++ b/sntrup-kem/src/scratch.rs @@ -0,0 +1,108 @@ +//! Stack scratch buffers that skip Rust's implicit zero-fill. +//! +//! The SIMD kernels work out of multi-kilobyte fixed-size stack arrays that are +//! written in full before anything reads them. Declaring those as `[0i16; N]` +//! makes the compiler emit a `memset` it cannot then eliminate — the producers +//! are opaque `#[target_feature]` calls, so LLVM cannot prove the overwrite. +//! Measured at 5.9% of a decapsulation (roughly 61 KB zeroed per call), against +//! a reference implementation that pays none of it. +//! +//! The obligation this trades for is real: every element must be written before +//! it is read. Rather than leave that as a comment, debug builds fill the buffer +//! with [`POISON`] instead of leaving it indeterminate, so a producer that skips +//! an element produces visibly wrong output in the differential and KAT suites +//! instead of silent nondeterminism. Release builds skip the fill entirely, +//! which is the point. + +use core::mem::MaybeUninit; + +/// Debug-build fill byte for scratch declared through [`uninit`]. +/// +/// `0x5A5A` as an `i16` is 23130 — outside every coefficient range in the +/// crate by a wide margin, and not a plausible near-miss for zero. +const POISON: u8 = 0x5A; + +/// Views owned, uninitialized stack storage as the array it will become. +/// +/// # Safety +/// `T` must be a plain integer type, for which every bit pattern is a valid +/// value. The caller must write every element of the returned buffer before +/// reading any element of it. +#[inline(always)] +pub(crate) unsafe fn uninit(slot: &mut MaybeUninit<[T; N]>) -> &mut [T; N] { + if cfg!(debug_assertions) { + // SAFETY: `slot` is owned, properly aligned storage of exactly this + // size, and `T` admits every bit pattern, so filling it with bytes + // produces valid (if deliberately absurd) values. + unsafe { + core::ptr::write_bytes(slot.as_mut_ptr().cast::(), POISON, size_of::<[T; N]>()); + } + } + // SAFETY: debug builds initialized the storage above. Release builds rely on + // the caller's documented contract that every element is written before any + // read; `T` has no invalid bit patterns, so forming the reference itself is + // well-defined regardless. + unsafe { slot.assume_init_mut() } +} + +/// Declares a stack scratch buffer without the implicit zero-fill. +/// +/// Expands to `let $name: &mut [$t; $n]`. Each use site must justify, in a +/// `SAFETY:` comment, that every element is written before it is read — see +/// the module documentation for how debug builds check that claim. +/// +/// This form introduces its own `unsafe` block, so it is for callers in safe +/// code. Code already inside an `unsafe` block — the SIMD kernels — calls +/// [`uninit`] directly instead, to avoid a redundant nested block. +macro_rules! uninit_scratch { + ($name:ident: [$t:ty; $n:expr]) => { + let mut $name = core::mem::MaybeUninit::<[$t; $n]>::uninit(); + let $name = unsafe { $crate::scratch::uninit(&mut $name) }; + }; +} + +pub(crate) use uninit_scratch; + +#[cfg(test)] +mod tests { + use super::POISON; + + #[test] + fn debug_builds_poison_and_release_builds_do_not_fill() { + uninit_scratch!(buf: [i16; 64]); + if cfg!(debug_assertions) { + assert!( + buf.iter().all(|&x| x == 0x5A5A), + "debug builds must poison so a missed write is visible" + ); + } + // The contract is write-before-read; writing in full is always sound. + buf.fill(7); + assert!(buf.iter().all(|&x| x == 7)); + } + + #[test] + fn poison_is_outside_every_coefficient_range() { + // A poison value that could pass for a real coefficient would let a + // missed write survive the differential suites, so tie the check to the + // actual parameter sets rather than to a hardcoded bound. + use crate::params::SntrupParams; + let widest = [ + crate::params::Sntrup653Params::params(), + crate::params::Sntrup761Params::params(), + crate::params::Sntrup857Params::params(), + crate::params::Sntrup953Params::params(), + crate::params::Sntrup1013Params::params(), + crate::params::Sntrup1277Params::params(), + ] + .iter() + .map(|p| p.q12) + .max() + .unwrap_or(i32::MAX); + let poison = i16::from_ne_bytes([POISON, POISON]); + assert!( + i32::from(poison) > widest, + "poison {poison} is inside the +/-{widest} coefficient range" + ); + } +} diff --git a/sntrup-kem/src/simd.rs b/sntrup-kem/src/simd.rs new file mode 100644 index 0000000..65bd1f8 --- /dev/null +++ b/sntrup-kem/src/simd.rs @@ -0,0 +1,27 @@ +//! Shared x86_64 SIMD helpers for the multiply kernels in `rq` and `r3`. +//! +//! The two row-major schoolbook kernels differ only in how they accumulate +//! `i16×i16` dot-product terms into i32 lanes. These helpers isolate that step so +//! one kernel body (expanded per feature level by a macro in each module) serves +//! both instruction sets: +//! +//! - [`mac_madd`]: plain AVX2 — `vpmaddwd` then a separate `vpaddd`. +//! - [`mac_vnni`]: AVX-VNNI — a single fused `vpdpwssd` (Zen 5, Alder Lake+), +//! removing one instruction and one dependency per 16 multiply-accumulates. +#![allow(unsafe_code)] + +use core::arch::x86_64::{__m256i, _mm256_add_epi32, _mm256_madd_epi16}; + +/// `acc + Σ_pairs(a·b)` via `vpmaddwd` + `vpaddd`. +#[inline] +#[target_feature(enable = "avx2")] +pub(crate) fn mac_madd(acc: __m256i, a: __m256i, b: __m256i) -> __m256i { + _mm256_add_epi32(acc, _mm256_madd_epi16(a, b)) +} + +/// `acc + Σ_pairs(a·b)` via the fused `vpdpwssd`. +#[inline] +#[target_feature(enable = "avx2,avxvnni")] +pub(crate) fn mac_vnni(acc: __m256i, a: __m256i, b: __m256i) -> __m256i { + core::arch::x86_64::_mm256_dpwssd_avx_epi32(acc, a, b) +} diff --git a/sntrup-kem/src/types.rs b/sntrup-kem/src/types.rs index 62fe0bf..8fb0af3 100644 --- a/sntrup-kem/src/types.rs +++ b/sntrup-kem/src/types.rs @@ -10,6 +10,12 @@ use zeroize::Zeroize; #[derive(Clone)] pub struct EncapsulationKey { bytes: Vec, + /// Decoded public-key polynomial and Hash4(pk), cached on first encapsulation. + /// + /// Encapsulation re-derives both on every call otherwise — repeated work that is + /// identical per key (~10% of the operation). Both are public values, so this + /// needs no zeroization. Mirrors `DecapsulationKey::h_cache`. + pk_cache: std::sync::OnceLock<(Vec, [u8; 32])>, _marker: PhantomData

, } @@ -17,6 +23,13 @@ pub struct EncapsulationKey { #[derive(Clone)] pub struct DecapsulationKey { bytes: Vec, + /// Decoded public-key polynomial, cached on first decapsulation. + /// + /// Decapsulation re-encrypts against the public key embedded in this secret + /// key, and decoding it is ~13% of the operation — pure repeated work, since + /// it is identical on every call. The public key is not secret, so this + /// needs no zeroization. + h_cache: std::sync::OnceLock>, _marker: PhantomData

, } @@ -62,16 +75,57 @@ macro_rules! impl_from_vec { }; } -impl_from_vec!(EncapsulationKey); -impl_from_vec!(DecapsulationKey); impl_from_vec!(Ciphertext); impl_from_vec!(SharedSecret); +impl EncapsulationKey

{ + pub(crate) fn from_vec(bytes: Vec) -> Self { + Self { + bytes, + pk_cache: std::sync::OnceLock::new(), + _marker: PhantomData, + } + } + + /// The decoded public-key polynomial and Hash4(pk), computed once and reused. + #[cfg(feature = "ecap")] + fn cached_pk(&self) -> &(Vec, [u8; 32]) { + self.pk_cache.get_or_init(|| { + let params = P::params(); + let mut h = vec![0i16; params.p]; + crate::rq::encoding::rq_decode_into(&self.bytes, &mut h, params); + let mut pk_hash = [0u8; 32]; + crate::utils::hash_prefix(&mut pk_hash, 4, &self.bytes); + (h, pk_hash) + }) + } +} + // --------------------------------------------------------------------------- // DecapsulationKey: extract encapsulation key // --------------------------------------------------------------------------- impl DecapsulationKey

{ + pub(crate) fn from_vec(bytes: Vec) -> Self { + Self { + bytes, + h_cache: std::sync::OnceLock::new(), + _marker: PhantomData, + } + } + + /// The decoded public-key polynomial, computed once and reused. + fn cached_h(&self) -> &[i16] { + self.h_cache.get_or_init(|| { + let params = P::params(); + let ses = params.small_encode_size; + let pk = &self.bytes[2 * ses..2 * ses + params.pk_size]; + let mut h = vec![0i16; params.p]; + crate::rq::encoding::rq_decode_into(pk, &mut h, params); + h + }) + } + /// Get the encapsulation (public) key embedded in this decapsulation key. /// /// SK layout: f(small_enc) || ginv(small_enc) || pk(pk_size) || rho(small_enc) || hash4(32) @@ -160,10 +214,7 @@ macro_rules! impl_try_from { actual: bytes.len(), }); } - Ok(Self { - bytes: bytes.to_vec(), - _marker: PhantomData, - }) + Ok(Self::from_vec(bytes.to_vec())) } } @@ -191,7 +242,6 @@ macro_rules! impl_try_from { } impl_try_from!(EncapsulationKey, PK_BYTES); -impl_try_from!(DecapsulationKey, SK_BYTES); impl_try_from!(Ciphertext, CT_BYTES); // --------------------------------------------------------------------------- @@ -218,6 +268,26 @@ impl Eq for Ciphertext

{} // ConstantTimeEq / PartialEq / Eq (DecapsulationKey) // --------------------------------------------------------------------------- +impl TryFrom<&[u8]> for DecapsulationKey

{ + type Error = Error; + fn try_from(bytes: &[u8]) -> Result { + if bytes.len() != P::SK_BYTES { + return Err(Error::InvalidSize { + expected: P::SK_BYTES, + actual: bytes.len(), + }); + } + Ok(Self::from_vec(bytes.to_vec())) + } +} + +impl TryFrom> for DecapsulationKey

{ + type Error = Error; + fn try_from(bytes: Vec) -> Result { + Self::try_from(bytes.as_slice()) + } +} + impl ConstantTimeEq for DecapsulationKey

{ fn ct_eq(&self, other: &Self) -> subtle::Choice { self.bytes.as_slice().ct_eq(other.bytes.as_slice()) @@ -288,7 +358,7 @@ impl SntrupKem

{ pub fn generate_key( rng: &mut impl rand::CryptoRng, ) -> (EncapsulationKey

, DecapsulationKey

) { - let (pk, sk) = crate::kem::keygen(P::params(), rng); + let (pk, sk) = crate::ops::keygen(P::params(), rng); ( EncapsulationKey::from_vec(pk), DecapsulationKey::from_vec(sk), @@ -299,6 +369,11 @@ impl SntrupKem

{ /// /// The seed is expanded via ChaCha20Rng to derive the full key pair. /// Identical seeds always produce identical key pairs. + /// + /// Note: `rand_chacha` offers no zeroization support, so the RNG's internal state (which + /// contains the seed) is dropped without being wiped when this returns. Callers with + /// strict key-erasure requirements should treat the seed's residency in freed stack + /// memory as a known limitation of this function. pub fn generate_key_deterministic( seed: &[u8; 32], ) -> (EncapsulationKey

, DecapsulationKey

) { @@ -312,7 +387,8 @@ impl SntrupKem

{ impl EncapsulationKey

{ /// Encapsulate: produce a ciphertext and shared secret. pub fn encapsulate(&self, rng: &mut impl rand::CryptoRng) -> (Ciphertext

, SharedSecret

) { - let (ct, ss) = crate::kem::encaps(&self.bytes, P::params(), rng); + let (h, pk_hash) = self.cached_pk(); + let (ct, ss) = crate::ops::encaps(h, pk_hash, P::params(), rng); (Ciphertext::from_vec(ct), SharedSecret::from_vec(ss)) } } @@ -325,7 +401,7 @@ impl DecapsulationKey

{ /// On failure, returns a pseudorandom key derived from rho, /// indistinguishable from a valid key to an attacker. pub fn decapsulate(&self, ct: &Ciphertext

) -> SharedSecret

{ - let ss = crate::kem::decaps(&self.bytes, &ct.bytes, P::params()); + let ss = crate::ops::decaps(&self.bytes, self.cached_h(), &ct.bytes, P::params()); SharedSecret::from_vec(ss) } } @@ -364,10 +440,7 @@ mod serde_impl { ), )); } - Ok(Self { - bytes: buf, - _marker: PhantomData, - }) + Ok(Self::from_vec(buf)) } } }; diff --git a/sntrup-kem/src/utils.rs b/sntrup-kem/src/utils.rs index f2be3da..0b2fa6b 100644 --- a/sntrup-kem/src/utils.rs +++ b/sntrup-kem/src/utils.rs @@ -2,6 +2,7 @@ use sha2::{Digest, Sha512}; use zeroize::Zeroize; use crate::params::SntrupParameters; +use crate::scratch::uninit_scratch; use crate::{r3, rq, zx}; /// Hash prefix helper: SHA-512(prefix || input), truncated to 32 bytes. @@ -9,8 +10,10 @@ pub(crate) fn hash_prefix(out: &mut [u8; 32], prefix: u8, input: &[u8]) { let mut hasher = Sha512::new(); hasher.update([prefix]); hasher.update(input); - let digest = hasher.finalize(); + let mut digest = hasher.finalize(); out.copy_from_slice(&digest[..32]); + // The discarded upper half is still derived from (possibly secret) input — wipe it. + digest.zeroize(); } /// hash_confirm: Hash(2 || Hash(3 || r_enc) || cache) @@ -23,8 +26,10 @@ pub(crate) fn hash_confirm(out: &mut [u8; 32], r_enc: &[u8], cache: &[u8; 32]) { hasher.update([2u8]); hasher.update(inner); hasher.update(&cache[..]); - let digest = hasher.finalize(); + let mut digest = hasher.finalize(); out.copy_from_slice(&digest[..32]); + inner.zeroize(); + digest.zeroize(); } /// hash_session: Hash(b || Hash(3 || y) || z) @@ -36,8 +41,10 @@ pub(crate) fn hash_session(out: &mut [u8; 32], b: u8, y: &[u8], z: &[u8]) { hasher.update([b]); hasher.update(inner); hasher.update(z); - let digest = hasher.finalize(); + let mut digest = hasher.finalize(); out.copy_from_slice(&digest[..32]); + inner.zeroize(); + digest.zeroize(); } /// Constant-time: returns 0 if x == 0, -1 (0xFFFFFFFF) otherwise. @@ -57,14 +64,12 @@ fn int16_nonzero_mask(x: i16) -> i32 { clippy::cast_possible_wrap )] pub(crate) fn weightw_mask(r: &[i8], p: usize, w: usize) -> i32 { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return weightw_mask_avx2(r, p, w); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return weightw_mask_avx2(r, p, w); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -85,11 +90,7 @@ fn weightw_mask_scalar(r: &[i8], _p: usize, w: usize) -> i32 { } /// Count non-zero elements 32 at a time using AVX2. -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] #[allow( unsafe_code, @@ -163,14 +164,12 @@ unsafe fn weightw_mask_neon(r: &[i8], p: usize, w: usize) -> i32 { /// Returns 0 if equal, -1 otherwise. #[allow(unsafe_code, clippy::cast_possible_wrap)] fn ciphertexts_diff_mask(a: &[u8], b: &[u8]) -> i32 { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return ciphertexts_diff_mask_avx2(a, b); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return ciphertexts_diff_mask_avx2(a, b); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -192,11 +191,7 @@ fn ciphertexts_diff_mask_scalar(a: &[u8], b: &[u8]) -> i32 { } /// XOR-accumulate 32 bytes at a time, then horizontal OR. -#[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") -))] +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] #[allow(unsafe_code, clippy::cast_possible_wrap, clippy::cast_sign_loss)] unsafe fn ciphertexts_diff_mask_avx2(a: &[u8], b: &[u8]) -> i32 { @@ -311,37 +306,45 @@ pub(crate) fn derive_key( clippy::cast_sign_loss, clippy::cast_possible_wrap )] -pub(crate) fn create_cipher(r: &[i8], pk: &[u8], params: &SntrupParameters) -> (Vec, [u8; 32]) { +pub(crate) fn create_cipher( + r: &[i8], + h: &[i16], + pk_hash: &[u8; 32], + params: &SntrupParameters, +) -> (Vec, [u8; 32]) { let p = params.p; - let h = rq::encoding::rq_decode(pk, params); - let mut c = vec![0i16; p]; - rq::mult(&mut c, &h, r, params); - rq::round3(&mut c, params); + use crate::params::MAX_P; - let r_enc = zx::encoding::encode(r, p, params.small_encode_size); + // SAFETY: `rq::mult` writes all `p` product coefficients. + uninit_scratch!(c_buf: [i16; MAX_P]); + let c = &mut c_buf[..p]; + rq::mult(c, h, r, params); - // Compute confirm hash: Hash(2 || Hash(3 || r_enc) || Hash4(pk)) - let mut cache = [0u8; 32]; - hash_prefix(&mut cache, 4, pk); + const MAX_SES: usize = MAX_P.div_ceil(4) + 1; + let ses = params.small_encode_size; + // SAFETY: `encode_into` writes all `ses` bytes. + uninit_scratch!(r_enc_buf: [u8; MAX_SES]); + let r_enc = &mut r_enc_buf[..ses]; + zx::encoding::encode_into(r, r_enc, p, ses); + + // Compute confirm hash: Hash(2 || Hash(3 || r_enc) || Hash4(pk)); Hash4(pk) is + // the caller-cached `pk_hash`. let mut confirm = [0u8; 32]; - hash_confirm(&mut confirm, &r_enc, &cache); + hash_confirm(&mut confirm, r_enc, pk_hash); // Ciphertext layout: rounded(rounded_encode_size) || confirm_hash(32) let mut cstr = vec![0u8; params.ct_size]; - cstr[..params.rounded_encode_size].copy_from_slice(&rq::encoding::rounded_encode(&c, params)); + rq::encoding::round_and_encode_into(c, &mut cstr[..params.rounded_encode_size], params); cstr[params.rounded_encode_size..].copy_from_slice(&confirm); // Shared key: hash_session(1, r_enc, cstr) let mut k = [0u8; 32]; - hash_session(&mut k, 1, &r_enc, &cstr); + hash_session(&mut k, 1, r_enc, &cstr); - // Zeroize secret intermediates - // r_enc, cache, confirm are on the stack / local Vecs and will be dropped, - // but we zeroize explicitly for defense in depth. - let mut r_enc = r_enc; - r_enc.zeroize(); - cache.zeroize(); + // Zeroize secret intermediates (whole frames, padding included). `pk_hash` is + // public and caller-owned; nothing to wipe for it. + crate::wipe::wipe(r_enc_buf); confirm.zeroize(); (cstr, k) @@ -356,14 +359,29 @@ pub(crate) fn create_cipher(r: &[i8], pk: &[u8], params: &SntrupParameters) -> ( clippy::cast_sign_loss, clippy::cast_possible_wrap )] -pub(crate) fn decapsulate_inner(cstr: &[u8], sk: &[u8], params: &SntrupParameters) -> [u8; 32] { +pub(crate) fn decapsulate_inner( + cstr: &[u8], + sk: &[u8], + h: &[i16], + params: &SntrupParameters, +) -> [u8; 32] { let p = params.p; let w = params.w; let ses = params.small_encode_size; + use crate::params::MAX_P; + // Parse SK: f(ses) || ginv(ses) || pk(pk_size) || rho(ses) || cache(32) - let mut f = zx::encoding::decode(&sk[..ses], p); - let mut ginv = zx::encoding::decode(&sk[ses..(2 * ses)], p); + // All working buffers live on this frame, bounded by MAX_P — decapsulation + // performs no heap allocation. + // SAFETY: `decode_into` writes all `p` coefficients. + uninit_scratch!(f_buf: [i8; MAX_P]); + let f = &mut f_buf[..p]; + zx::encoding::decode_into(&sk[..ses], f, p); + // SAFETY: `decode_into` writes all `p` coefficients. + uninit_scratch!(ginv_buf: [i8; MAX_P]); + let ginv = &mut ginv_buf[..p]; + zx::encoding::decode_into(&sk[ses..(2 * ses)], ginv, p); let pk_start = 2 * ses; let pk_end = pk_start + params.pk_size; let rho_start = pk_end; @@ -374,24 +392,26 @@ pub(crate) fn decapsulate_inner(cstr: &[u8], sk: &[u8], params: &SntrupParameter cache.copy_from_slice(&sk[cache_start..cache_start + 32]); // Decrypt: Rounded_decode, multiply by f, Rq_mult3, R3_fromRq, R3_mult by ginv - let c = rq::encoding::rounded_decode(&cstr[..params.rounded_encode_size], params); - let mut cf = vec![0i16; p]; - rq::mult(&mut cf, &c, &f, params); - let mut t3 = vec![0i8; p]; - for i in 0..p { - t3[i] = r3::mod3::freeze(rq::modq::freeze( - 3 * cf[i] as i32, - params.q, - params.barrett1, - params.barrett2, - ) as i32); - } - let mut r = vec![0i8; p]; - r3::mult(&mut r, &t3, &ginv, p); + // SAFETY: `rounded_decode_into` writes all `p` coefficients. + uninit_scratch!(c_buf: [i16; MAX_P]); + let c = &mut c_buf[..p]; + rq::encoding::rounded_decode_into(&cstr[..params.rounded_encode_size], c, params); + // SAFETY: `rq::mult` writes all `p` product coefficients. + uninit_scratch!(cf_buf: [i16; MAX_P]); + let cf = &mut cf_buf[..p]; + rq::mult(cf, c, f, params); + // SAFETY: `scale3_freeze3` writes one output per input coefficient. + uninit_scratch!(t3_buf: [i8; MAX_P]); + let t3 = &mut t3_buf[..p]; + rq::scale3_freeze3(t3, cf, params); + // SAFETY: `r3::mult` writes all `p` product coefficients. + uninit_scratch!(r_buf: [i8; MAX_P]); + let r = &mut r_buf[..p]; + r3::mult(r, t3, ginv, p); // Weight mask: on failure, set r to default weight-W vector // (W ones followed by P-W zeros), matching PQClean's Decrypt - let w_mask = weightw_mask(&r, p, w); + let w_mask = weightw_mask(r, p, w); let not_mask = (!w_mask) as i8; for val in r[..w].iter_mut() { *val = ((*val ^ 1) & not_mask) ^ 1; @@ -401,24 +421,36 @@ pub(crate) fn decapsulate_inner(cstr: &[u8], sk: &[u8], params: &SntrupParameter } // Hide: encode r, re-encrypt with pk, compute confirm hash - let mut r_enc = zx::encoding::encode(&r, p, ses); - let h = rq::encoding::rq_decode(&sk[pk_start..pk_end], params); - let mut hr = vec![0i16; p]; - rq::mult(&mut hr, &h, &r, params); - rq::round3(&mut hr, params); - let mut cnew = vec![0u8; params.ct_size]; - cnew[..params.rounded_encode_size].copy_from_slice(&rq::encoding::rounded_encode(&hr, params)); + const MAX_SES: usize = MAX_P.div_ceil(4) + 1; + // SAFETY: `encode_into` writes all `ses` bytes. + uninit_scratch!(r_enc_buf: [u8; MAX_SES]); + let r_enc = &mut r_enc_buf[..ses]; + zx::encoding::encode_into(r, r_enc, p, ses); + // SAFETY: `rq::mult` writes all `p` product coefficients. + uninit_scratch!(hr_buf: [i16; MAX_P]); + let hr = &mut hr_buf[..p]; + rq::mult(hr, h, r, params); + + // ct_size = rounded_encode_size + 32; bound by the largest set. + const MAX_CT: usize = 1847 + 32; + // SAFETY: `round_and_encode_into` fills the rounded prefix and the confirm + // hash is copied over the remainder, together covering all of `..ct_size`. + uninit_scratch!(cnew_buf: [u8; MAX_CT]); + let cnew = &mut cnew_buf[..params.ct_size]; + rq::encoding::round_and_encode_into(hr, &mut cnew[..params.rounded_encode_size], params); let mut confirm = [0u8; 32]; - hash_confirm(&mut confirm, &r_enc, &cache); + hash_confirm(&mut confirm, r_enc, &cache); cnew[params.rounded_encode_size..].copy_from_slice(&confirm); // Compare full ciphertexts (rounded + confirm hash) - let mask = ciphertexts_diff_mask(cstr, &cnew); + let mask = ciphertexts_diff_mask(cstr, cnew); // Constant-time select: r_enc on success (mask=0), rho on failure (mask=-1) let rho = &sk[rho_start..rho_end]; - let mut selected = vec![0u8; ses]; - selected.copy_from_slice(&r_enc); + // SAFETY: `copy_from_slice` below writes all `ses` bytes. + uninit_scratch!(selected_buf: [u8; MAX_SES]); + let selected = &mut selected_buf[..ses]; + selected.copy_from_slice(r_enc); let mask_byte = mask as u8; for i in 0..ses { selected[i] ^= mask_byte & (selected[i] ^ rho[i]); @@ -427,20 +459,20 @@ pub(crate) fn decapsulate_inner(cstr: &[u8], sk: &[u8], params: &SntrupParameter // Hash session: prefix=1 on success (mask=0), prefix=0 on failure (mask=-1) let prefix = (1 + mask) as u8; let mut k = [0u8; 32]; - hash_session(&mut k, prefix, &selected, cstr); + hash_session(&mut k, prefix, selected, cstr); - // Zeroize secret intermediates - f.zeroize(); - ginv.zeroize(); + // Zeroize secret intermediates (the whole stack frames, padding included). + crate::wipe::wipe(f_buf); + crate::wipe::wipe(ginv_buf); cache.zeroize(); - cf.zeroize(); - t3.zeroize(); - r.zeroize(); - r_enc.zeroize(); - hr.zeroize(); - cnew.zeroize(); + crate::wipe::wipe(cf_buf); + crate::wipe::wipe(t3_buf); + crate::wipe::wipe(r_buf); + crate::wipe::wipe(r_enc_buf); + crate::wipe::wipe(hr_buf); + crate::wipe::wipe(cnew_buf); confirm.zeroize(); - selected.zeroize(); + crate::wipe::wipe(selected_buf); k } diff --git a/sntrup-kem/src/wipe.rs b/sntrup-kem/src/wipe.rs new file mode 100644 index 0000000..780b018 --- /dev/null +++ b/sntrup-kem/src/wipe.rs @@ -0,0 +1,102 @@ +//! Fast volatile wiping of plain-integer scratch buffers. + +use zeroize::{DefaultIsZeroes, Zeroize}; + +/// Wipe a plain-integer buffer with volatile stores, as wide as the platform +/// allows. +/// +/// [`Zeroize`] on a slice issues one volatile store *per element*, which the +/// compiler is not permitted to merge or vectorize. Across the multi-kilobyte +/// scratch buffers the SIMD kernels use, that granularity — not the wiping +/// itself — dominated the cost: measured at roughly 18% of decapsulation for the +/// NTT multiply alone. +/// +/// Re-viewing the buffer at a wider integer width keeps the volatile guarantee +/// exactly — any unaligned head and tail are still wiped, just at their own +/// width — while issuing proportionally fewer stores. `u64` is the portable +/// floor; on x86_64 a 32-byte store cuts it by a further factor of four, which +/// matters because the reference implementations wipe nothing at all and this +/// is cost they simply do not pay. +#[inline] +pub(crate) fn wipe(buf: &mut [T]) { + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 confirmed present at runtime, which implies AVX. + unsafe { + wipe_avx(buf); + } + return; + } + wipe_u64(buf); +} + +/// 32-byte volatile stores over the aligned interior. +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +#[target_feature(enable = "avx")] +unsafe fn wipe_avx(buf: &mut [T]) { + use core::arch::x86_64::__m256i; + // SAFETY: `T` is a plain integer with no padding and no invalid bit + // patterns, so viewing the buffer's bytes as `__m256i` is valid. + // `align_to_mut` guarantees the middle is correctly aligned for a volatile + // store of that width and hands back whatever head and tail are not. + unsafe { + let (head, mid, tail) = buf.align_to_mut::<__m256i>(); + head.zeroize(); + let zero: __m256i = core::mem::zeroed(); + for slot in mid { + core::ptr::write_volatile(slot, zero); + } + tail.zeroize(); + } +} + +/// Portable eight-byte-at-a-time form. +fn wipe_u64(buf: &mut [T]) { + // SAFETY: as above, at `u64` width. + unsafe { + let (head, mid, tail) = buf.align_to_mut::(); + head.zeroize(); + mid.zeroize(); + tail.zeroize(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Every byte must be cleared regardless of where the wide interior starts, + /// so exercise every offset and length across an alignment period. + #[test] + fn clears_every_byte_at_every_alignment_and_length() { + let mut backing = [0i16; 96]; + for off in 0..16usize { + for len in 0..64usize { + backing.fill(-1); + wipe(&mut backing[off..off + len]); + assert!( + backing[off..off + len].iter().all(|&x| x == 0), + "not cleared at off={off} len={len}" + ); + assert!( + backing[..off].iter().all(|&x| x == -1) + && backing[off + len..].iter().all(|&x| x == -1), + "wiped outside the slice at off={off} len={len}" + ); + } + } + } + + #[test] + fn both_widths_agree_on_byte_data() { + let mut a = [0u8; 200]; + let mut b = [0u8; 200]; + for off in 0..8usize { + a.fill(0xAA); + b.fill(0xAA); + wipe(&mut a[off..]); + wipe_u64(&mut b[off..]); + assert_eq!(a, b, "widths disagree at off={off}"); + } + } +} diff --git a/sntrup-kem/src/zx.rs b/sntrup-kem/src/zx.rs index 58d40b6..bfc4bf1 100644 --- a/sntrup-kem/src/zx.rs +++ b/sntrup-kem/src/zx.rs @@ -1,3 +1,8 @@ +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +mod codec3; +#[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] +mod djbsort; + /// Small-element (ternary) encoding and decoding. pub mod encoding { /// Encode a small polynomial `f` of length `p` into `small_encode_size` bytes. @@ -7,7 +12,23 @@ pub mod encoding { #[allow(clippy::cast_sign_loss)] pub fn encode(f: &[i8], p: usize, small_encode_size: usize) -> Vec { let mut c = vec![0u8; small_encode_size]; - for (byte, chunk) in c[..small_encode_size - 1].iter_mut().zip(f.chunks(4)) { + encode_into(f, &mut c, p, small_encode_size); + c + } + + /// Allocation-free form of [`encode`]: writes into `c[..small_encode_size]`. + #[allow(clippy::cast_sign_loss)] + pub fn encode_into(f: &[i8], c: &mut [u8], p: usize, small_encode_size: usize) { + let n = small_encode_size - 1; + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 confirmed present at runtime. `p - 1 == 4 * n` holds + // for every parameter set, which is the kernel's length contract. + unsafe { super::codec3::encode_avx2(&f[..4 * n], &mut c[..n]) }; + c[n] = (f[p - 1] + 1) as u8; + return; + } + for (byte, chunk) in c[..n].iter_mut().zip(f.chunks(4)) { let mut c0 = chunk[0] + 1; c0 += (chunk[1] + 1) << 2; c0 += (chunk[2] + 1) << 4; @@ -15,17 +36,21 @@ pub mod encoding { *byte = c0 as u8; } c[small_encode_size - 1] = (f[p - 1] + 1) as u8; - c } - /// Decode `small_encode_size` bytes into a small polynomial of length `p`. - /// - /// Inverse of [`encode`]: unpacks 4 trits per byte, last element from last byte. + /// Allocation-free decoder: writes into `f[..p]`. #[allow(clippy::cast_possible_wrap)] - pub fn decode(c: &[u8], p: usize) -> Vec { + pub fn decode_into(c: &[u8], f: &mut [i8], p: usize) { let small_encode_size = c.len(); - let mut f = vec![0i8; p]; - for (byte, chunk) in c[..small_encode_size - 1].iter().zip(f.chunks_mut(4)) { + let n = small_encode_size - 1; + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 confirmed present at runtime; `p - 1 == 4 * n`. + unsafe { super::codec3::decode_avx2(&c[..n], &mut f[..4 * n]) }; + f[p - 1] = ((c[n] & 3) as i8) - 1; + return; + } + for (byte, chunk) in c[..n].iter().zip(f.chunks_mut(4)) { let mut c0 = *byte; chunk[0] = ((c0 & 3) as i8) - 1; c0 >>= 2; @@ -36,7 +61,6 @@ pub mod encoding { chunk[3] = ((c0 & 3) as i8) - 1; } f[p - 1] = ((c[small_encode_size - 1] & 3) as i8) - 1; - f } } @@ -62,17 +86,28 @@ pub mod random { x[j] ^= c; } - /// Batcher bitonic sort on `n` elements of `x`, dispatching to SIMD when available. + /// Constant-time sort of `n` elements of `x`, dispatching to the best + /// available implementation. + /// + /// On x86_64 with AVX2 this is the port of djb's `crypto_sort_int32`, whose + /// register-blocked merges keep every lane live; the Batcher network below + /// runs at roughly half lane utilisation and is ~7.6x slower at p = 761. + /// Both are differentially tested against each other. #[allow(unsafe_code)] pub fn sort(x: &mut [i32], n: usize) { - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] - // SAFETY: AVX2 verified by cfg - unsafe { - return sort_avx2(x, n); + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if crate::cpu::has_avx2() { + // SAFETY: AVX2 support confirmed by has_avx2() + unsafe { + return super::djbsort::sort(x, n); + } + } + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] + if false { + // SAFETY: unreachable; retained as the differential oracle. + unsafe { + return sort_avx2(x, n); + } } #[cfg(all(target_arch = "aarch64", not(feature = "force-scalar")))] // SAFETY: NEON is baseline on aarch64 @@ -113,11 +148,7 @@ pub mod random { /// AVX2-accelerated Batcher bitonic sort. /// Uses _mm256_min/max_epi32 for 8 parallel comparators when stride >= 8. - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] #[allow(unsafe_code)] unsafe fn sort_avx2(x: &mut [i32], n: usize) { @@ -145,13 +176,9 @@ pub mod random { } } - /// Process one pass of comparators: `minmax(x[i+off0], x[i+off1])` + /// Process one pass of comparators: minmax(x[i+off0], x[i+off1]) /// for all i in 0..(n-off1) where i & p_mask == 0. - #[cfg(all( - target_arch = "x86_64", - target_feature = "avx2", - not(feature = "force-scalar") - ))] + #[cfg(all(target_arch = "x86_64", not(feature = "force-scalar")))] #[target_feature(enable = "avx2")] #[allow(unsafe_code)] unsafe fn minmax_pass_avx2(x: &mut [i32], n: usize, p_mask: usize, off0: usize, off1: usize) { @@ -190,8 +217,73 @@ pub mod random { i += 1; } } + } else if off0 == 0 && off1 == p_mask { + // Register-local pass at stride p ∈ {1,2,4}: within one 8-lane + // block, lane l pairs with lane l ^ p. One load, one permute, + // min/max, one const-immediate blend, one store. + let mut i0 = 0usize; + macro_rules! local_pass { + ($swap:expr, $imm:literal) => { + while i0 + 8 <= end { + let v = _mm256_loadu_si256(x.as_ptr().add(i0) as *const __m256i); + let w = $swap(v); + let mn = _mm256_min_epi32(v, w); + let mx = _mm256_max_epi32(v, w); + _mm256_storeu_si256( + x.as_mut_ptr().add(i0) as *mut __m256i, + _mm256_blend_epi32::<$imm>(mn, mx), + ); + i0 += 8; + } + }; + } + match p_mask { + 4 => local_pass!(|v| _mm256_permute4x64_epi64::<0x4E>(v), 0b1111_0000), + 2 => local_pass!(|v| _mm256_shuffle_epi32::<0x4E>(v), 0b1100_1100), + _ => local_pass!(|v| _mm256_shuffle_epi32::<0xB1>(v), 0b1010_1010), + } + for i in i0..end { + if i & p_mask == 0 { + int32_minmax(x, i + off0, i + off1); + } + } + } else if off1 >= 8 { + // Sub-pass with small selection stride p ∈ {1,2,4} but distant + // partner (off1 ≥ 8): two loads at the two offsets, min/max, and a + // const-immediate blend keeps inactive lanes (l & p ≠ 0) unchanged. + let mut i0 = 0usize; + macro_rules! masked_pass { + ($imm:literal) => { + while i0 + 8 <= end { + let a = _mm256_loadu_si256(x.as_ptr().add(i0 + off0) as *const __m256i); + let b = _mm256_loadu_si256(x.as_ptr().add(i0 + off1) as *const __m256i); + let mn = _mm256_min_epi32(a, b); + let mx = _mm256_max_epi32(a, b); + _mm256_storeu_si256( + x.as_mut_ptr().add(i0 + off0) as *mut __m256i, + _mm256_blend_epi32::<$imm>(a, mn), + ); + _mm256_storeu_si256( + x.as_mut_ptr().add(i0 + off1) as *mut __m256i, + _mm256_blend_epi32::<$imm>(b, mx), + ); + i0 += 8; + } + }; + } + match p_mask { + 4 => masked_pass!(0b0000_1111), + 2 => masked_pass!(0b0011_0011), + _ => masked_pass!(0b0101_0101), + } + for i in i0..end { + if i & p_mask == 0 { + int32_minmax(x, i + off0, i + off1); + } + } } else { - // Small strides: scalar + // Small p with nearby partner (off1 < 8): overlapping-store hazard, + // scalar. Only the (2,2,4), (1,1,2), (1,1,4) shapes land here. for i in 0..end { if i & p_mask == 0 { int32_minmax(x, i + off0, i + off1); @@ -239,11 +331,31 @@ pub mod random { let end = n.saturating_sub(off1); if p_mask >= 4 { + // Contiguous blocks of p_mask elements; four vectors per iteration for ILP. let mut i = 0; while i < end { if i & p_mask == 0 { let block_end = (i + p_mask).min(end); let mut j = i; + while j + 16 <= block_end { + let a0 = vld1q_s32(x.as_ptr().add(j + off0)); + let a1 = vld1q_s32(x.as_ptr().add(j + off0 + 4)); + let a2 = vld1q_s32(x.as_ptr().add(j + off0 + 8)); + let a3 = vld1q_s32(x.as_ptr().add(j + off0 + 12)); + let b0 = vld1q_s32(x.as_ptr().add(j + off1)); + let b1 = vld1q_s32(x.as_ptr().add(j + off1 + 4)); + let b2 = vld1q_s32(x.as_ptr().add(j + off1 + 8)); + let b3 = vld1q_s32(x.as_ptr().add(j + off1 + 12)); + vst1q_s32(x.as_mut_ptr().add(j + off0), vminq_s32(a0, b0)); + vst1q_s32(x.as_mut_ptr().add(j + off0 + 4), vminq_s32(a1, b1)); + vst1q_s32(x.as_mut_ptr().add(j + off0 + 8), vminq_s32(a2, b2)); + vst1q_s32(x.as_mut_ptr().add(j + off0 + 12), vminq_s32(a3, b3)); + vst1q_s32(x.as_mut_ptr().add(j + off1), vmaxq_s32(a0, b0)); + vst1q_s32(x.as_mut_ptr().add(j + off1 + 4), vmaxq_s32(a1, b1)); + vst1q_s32(x.as_mut_ptr().add(j + off1 + 8), vmaxq_s32(a2, b2)); + vst1q_s32(x.as_mut_ptr().add(j + off1 + 12), vmaxq_s32(a3, b3)); + j += 16; + } while j + 4 <= block_end { let a = vld1q_s32(x.as_ptr().add(j + off0)); let b = vld1q_s32(x.as_ptr().add(j + off1)); @@ -261,8 +373,77 @@ pub mod random { i += 1; } } + } else if off0 == 0 { + // Register-local first pass at stride p ∈ {1, 2}: within one 4-lane + // vector, lane l pairs with lane l ^ p. One load, one in-register + // partner shuffle, min/max, one constant-mask blend, one store. + // The blend keeps min in the low lane of each pair and max in the + // high lane, exactly the scalar comparator's writeback. + let mut i0 = 0usize; + if p_mask == 2 { + // Partner = lanes rotated by 2 (swap 64-bit halves). + let take_max = vcombine_u32(vdup_n_u32(0), vdup_n_u32(u32::MAX)); + while i0 + 4 <= end { + let v = vld1q_s32(x.as_ptr().add(i0)); + let w = vextq_s32::<2>(v, v); + let mn = vminq_s32(v, w); + let mx = vmaxq_s32(v, w); + vst1q_s32(x.as_mut_ptr().add(i0), vbslq_s32(take_max, mx, mn)); + i0 += 4; + } + } else { + // p = 1: partner = lanes swapped within each 64-bit pair. + let take_max = vreinterpretq_u32_u64(vdupq_n_u64(0xFFFF_FFFF_0000_0000)); + while i0 + 4 <= end { + let v = vld1q_s32(x.as_ptr().add(i0)); + let w = vrev64q_s32(v); + let mn = vminq_s32(v, w); + let mx = vmaxq_s32(v, w); + vst1q_s32(x.as_mut_ptr().add(i0), vbslq_s32(take_max, mx, mn)); + i0 += 4; + } + } + for i in i0..end { + if i & p_mask == 0 { + int32_minmax(x, i + off0, i + off1); + } + } + } else if off1 >= 4 && !(off0 == 1 && off1 == 2) { + // Sub-pass with small selection stride p ∈ {1, 2} (off0 == p) and a + // partner at off1 ≥ 4: two loads, min/max, constant-mask blends keep + // inactive lanes (l & p ≠ 0) at their loaded values. + // + // For off1 = 4 the two 4-lane windows overlap by off1 − off0 ∈ {2, 3} + // trailing lanes of the low window. Those overlapping low-window lanes + // are always inactive (their element index has bit `p` set), so the + // low store writes them back unchanged and the high store — issued + // after it — supplies their comparator results. The one shape where an + // overlapping low-window lane is *active*, (p, off0, off1) = (1, 1, 2), + // is excluded above and stays scalar: either store order would clobber + // a comparator result there. + let take_lo = if p_mask == 2 { + vcombine_u32(vdup_n_u32(u32::MAX), vdup_n_u32(0)) + } else { + vreinterpretq_u32_u64(vdupq_n_u64(0x0000_0000_FFFF_FFFF)) + }; + let mut i0 = 0usize; + while i0 + 4 <= end { + let a = vld1q_s32(x.as_ptr().add(i0 + off0)); + let b = vld1q_s32(x.as_ptr().add(i0 + off1)); + let mn = vminq_s32(a, b); + let mx = vmaxq_s32(a, b); + vst1q_s32(x.as_mut_ptr().add(i0 + off0), vbslq_s32(take_lo, mn, a)); + vst1q_s32(x.as_mut_ptr().add(i0 + off1), vbslq_s32(take_lo, mx, b)); + i0 += 4; + } + for i in i0..end { + if i & p_mask == 0 { + int32_minmax(x, i + off0, i + off1); + } + } } else { - // Small strides: scalar + // (1, 1, 2): overlapping windows with an active lane in the overlap — + // scalar is the only correct order. for i in 0..end { if i & p_mask == 0 { int32_minmax(x, i + off0, i + off1); @@ -300,9 +481,25 @@ pub mod random { /// then a constant-time sort shuffles them. #[allow(clippy::cast_possible_wrap)] pub fn random_tsmall(f: &mut [i8], p: usize, w: usize, rng: &mut impl Rng) { - let mut r = vec![0i32; p]; - for val in r.iter_mut() { - *val = rng.random(); + use crate::params::MAX_P; + use crate::scratch::uninit_scratch; + + // One bulk RNG call instead of `p` per-element calls. For any `rand_core` + // generator, `next_u32` is defined as the next four stream bytes little-endian, + // so a byte fill reinterpreted LE is value-identical to the per-element + // `rng.random::()` loop this replaces — the deterministic-keygen KATs + // pin that equivalence. + // SAFETY: `fill_bytes` writes all `4 * p` bytes before they are read. + uninit_scratch!(bytes_buf: [u8; 4 * MAX_P]); + let bytes = &mut bytes_buf[..4 * p]; + rng.fill_bytes(bytes); + + // SAFETY: every element of `r` is written from `bytes` before being read. + uninit_scratch!(r_buf: [i32; MAX_P]); + let r = &mut r_buf[..p]; + for (val, chunk) in r.iter_mut().zip(bytes.chunks_exact(4)) { + // SAFETY (index): chunks_exact(4) yields exactly 4-byte chunks. + *val = i32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); } for val in r[..w].iter_mut() { *val &= -2; @@ -310,9 +507,53 @@ pub mod random { for val in r[w..p].iter_mut() { *val = (*val & -3) | 1 } - sort_uint32(&mut r, p); + sort_uint32(r, p); for (fv, &rv) in f.iter_mut().zip(r.iter()) { *fv = ((rv & 3) as i8) - 1; } + // The tagged randomness fully determines the secret polynomial — wipe both + // frames (padding included), at wide-store granularity. + crate::wipe::wipe(bytes_buf); + crate::wipe::wipe(r_buf); + } +} + +#[cfg(test)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +mod sort_tests { + use super::random::sort; + + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + /// The dispatched sort must produce fully sorted output at every length that + /// exercises the vectorized large-stride, register-local, masked sub-pass, + /// and scalar paths — including the six parameter sizes. + #[test] + fn sort_orders_correctly_at_all_path_lengths() { + let mut s = 0x0dd_ba11u64 | 1; + for &n in &[ + 0usize, 1, 2, 3, 7, 8, 9, 15, 16, 17, 31, 64, 100, 653, 761, 857, 953, 1013, 1277, + ] { + for pattern in 0..7 { + let mut x: Vec = (0..n) + .map(|i| match pattern { + 0..=2 => next(&mut s) as i32, // random + 3 => 42, // all equal + 4 => i as i32, // sorted + 5 => (n - i) as i32, // reverse sorted + _ => (next(&mut s) % 4) as i32, // heavy duplicates + }) + .collect(); + let mut want = x.clone(); + want.sort_unstable(); + sort(&mut x, n); + assert_eq!(x, want, "sort mismatch at n={n} pattern={pattern}"); + } + } } } diff --git a/sntrup-kem/src/zx/codec3.rs b/sntrup-kem/src/zx/codec3.rs new file mode 100644 index 0000000..0084745 --- /dev/null +++ b/sntrup-kem/src/zx/codec3.rs @@ -0,0 +1,168 @@ +//! AVX2 kernels for the packed-ternary small-polynomial codec (`x3`). +//! +//! A small polynomial has coefficients in {-1, 0, 1}; the wire format shifts +//! each to {0, 1, 2} and packs four per byte, low trit first. Every parameter +//! set has `p - 1` divisible by 4, so the bulk is a clean 4:1 repack and only +//! the final coefficient needs separate handling — which the callers in +//! [`crate::zx::encoding`] do. +//! +//! The scalar forms these replace cost roughly 27x what liboqs spends on the +//! same work, which made them the largest non-algorithmic item in a +//! decapsulation profile despite touching only ~190 bytes. + +use core::arch::x86_64::*; + +/// Packs the `4 * out.len()` trits of `f` into `out`, four per byte. +/// +/// # Panics +/// Debug-only: `f.len()` must be exactly `4 * out.len()`. +#[target_feature(enable = "avx2")] +pub(crate) fn encode_avx2(f: &[i8], out: &mut [u8]) { + debug_assert_eq!(f.len(), 4 * out.len()); + unsafe { + // maddubs pairs adjacent trits as t0 + 4*t1 (max 10, no overflow of the + // signed 16-bit accumulator); madd then folds pairs of those as + // x0 + 16*x1, which is exactly the four-trit byte value. + let w_lo = _mm256_set1_epi16(0x0401); + let w_hi = _mm256_set1_epi32(0x0010_0001); + // Byte 0 of each 32-bit lane holds one packed output byte. + let gather = _mm256_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 0, 4, 8, 12, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, + ); + let one = _mm256_set1_epi8(1); + + let blocks = out.len() / 8; + for b in 0..blocks { + let t = _mm256_loadu_si256(f.as_ptr().add(32 * b) as *const __m256i); + let t = _mm256_add_epi8(t, one); + let p16 = _mm256_maddubs_epi16(t, w_lo); + let p32 = _mm256_madd_epi16(p16, w_hi); + let packed = _mm256_shuffle_epi8(p32, gather); + let lo = _mm256_castsi256_si128(packed); + let hi = _mm256_extracti128_si256(packed, 1); + _mm_storel_epi64( + out.as_mut_ptr().add(8 * b) as *mut __m128i, + _mm_unpacklo_epi32(lo, hi), + ); + } + + for i in (8 * blocks)..out.len() { + let q = &f[4 * i..4 * i + 4]; + let t = |x: i8| u8::from_ne_bytes((x + 1).to_ne_bytes()); + out[i] = t(q[0]) | (t(q[1]) << 2) | (t(q[2]) << 4) | (t(q[3]) << 6); + } + } +} + +/// Unpacks the `4 * src.len()` trits packed in `src` into `f`. +/// +/// # Panics +/// Debug-only: `f.len()` must be exactly `4 * src.len()`. +#[target_feature(enable = "avx2")] +pub(crate) fn decode_avx2(src: &[u8], f: &mut [i8]) { + debug_assert_eq!(f.len(), 4 * src.len()); + unsafe { + // Each input byte becomes four output bytes, so a 32-bit output lane is + // one input byte replicated; the four trits then live at shifts 0/2/4/6 + // of that lane. Shifting the whole lane by each amount and selecting the + // one byte that lands correctly is cheaper than any per-byte shift. + let spread = _mm256_setr_epi8( + 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 7, + 7, 7, 7, + ); + let m1 = _mm256_set1_epi32(0x0000_FF00); + let m2 = _mm256_set1_epi32(0x00FF_0000); + // 0xFF00_0000 does not fit a positive `i32`; build it by shifting. + let m3 = _mm256_slli_epi32::<24>(_mm256_set1_epi32(0xFF)); + let three = _mm256_set1_epi8(3); + let one = _mm256_set1_epi8(1); + + let blocks = src.len() / 8; + for b in 0..blocks { + let raw = _mm_loadl_epi64(src.as_ptr().add(8 * b) as *const __m128i); + let r = _mm256_shuffle_epi8(_mm256_broadcastsi128_si256(raw), spread); + let v = _mm256_blendv_epi8(r, _mm256_srli_epi32(r, 2), m1); + let v = _mm256_blendv_epi8(v, _mm256_srli_epi32(r, 4), m2); + let v = _mm256_blendv_epi8(v, _mm256_srli_epi32(r, 6), m3); + let v = _mm256_sub_epi8(_mm256_and_si256(v, three), one); + _mm256_storeu_si256(f.as_mut_ptr().add(32 * b) as *mut __m256i, v); + } + + for i in (8 * blocks)..src.len() { + let byte = src[i]; + for j in 0..4 { + f[4 * i + j] = i8::from_ne_bytes([((byte >> (2 * j)) & 3)]) - 1; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn scalar_encode(f: &[i8], out: &mut [u8]) { + for (byte, q) in out.iter_mut().zip(f.chunks(4)) { + let t = |x: i8| u8::from_ne_bytes((x + 1).to_ne_bytes()); + *byte = t(q[0]) | (t(q[1]) << 2) | (t(q[2]) << 4) | (t(q[3]) << 6); + } + } + + fn scalar_decode(src: &[u8], f: &mut [i8]) { + for (byte, q) in src.iter().zip(f.chunks_mut(4)) { + for (j, slot) in q.iter_mut().enumerate() { + *slot = i8::from_ne_bytes([((byte >> (2 * j)) & 3)]) - 1; + } + } + } + + /// Deterministic ternary stream: covers every trit in every lane position. + fn trits(n: usize, seed: u64) -> Vec { + let mut s = seed | 1; + (0..n) + .map(|_| { + s = s.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1); + i8::try_from((s >> 33) % 3).unwrap_or(0) - 1 + }) + .collect() + } + + #[test] + fn kernels_match_scalar_over_every_parameter_length() { + if !crate::cpu::has_avx2() { + return; + } + // Every real `small_encode_size - 1`, plus lengths around the 8-byte + // block boundary so the scalar tail is exercised in all residues. + let lens: Vec = (0usize..24).chain([163, 190, 214, 238, 253, 319]).collect(); + for n in lens { + for seed in 0..4u64 { + let f = trits(4 * n, seed * 7 + 1); + let (mut a, mut b) = (vec![0u8; n], vec![0u8; n]); + scalar_encode(&f, &mut a); + unsafe { encode_avx2(&f, &mut b) }; + assert_eq!(a, b, "encode mismatch at n={n} seed={seed}"); + + let (mut x, mut y) = (vec![0i8; 4 * n], vec![0i8; 4 * n]); + scalar_decode(&a, &mut x); + unsafe { decode_avx2(&a, &mut y) }; + assert_eq!(x, y, "decode mismatch at n={n} seed={seed}"); + assert_eq!(x, f, "decode is not the inverse of encode at n={n}"); + } + } + } + + #[test] + fn every_byte_value_round_trips() { + if !crate::cpu::has_avx2() { + return; + } + let src: Vec = (0..=255u8).collect(); + let mut f = vec![0i8; 4 * src.len()]; + unsafe { decode_avx2(&src, &mut f) }; + let mut back = vec![0u8; src.len()]; + unsafe { encode_avx2(&f, &mut back) }; + assert_eq!(src, back); + } +} diff --git a/sntrup-kem/src/zx/djbsort.rs b/sntrup-kem/src/zx/djbsort.rs new file mode 100644 index 0000000..d54ad6d --- /dev/null +++ b/sntrup-kem/src/zx/djbsort.rs @@ -0,0 +1,1381 @@ +//! Port of djb's AVX2 constant-time `crypto_sort_int32` (in progress). +//! +//! Motivation, from the round-8 comparator census: our Batcher network spends +//! 52% of its comparators in passes whose selection stride is below the vector +//! width, and those run at **half lane utilisation** — only the lanes with +//! `l & p == 0` do useful work. No per-pass fix reaches that; the structure has +//! to change. This implementation does, by keeping several merge stages resident +//! in registers (`merge16_finish` runs four stages with a single load and store; +//! `threestages` runs three across eight vectors) so every lane is live. +//! +//! Complete: base cases, the general power-of-two driver, and the +//! non-power-of-two dispatch. [`crate::zx::random::sort`] routes here on +//! x86_64/AVX2; the Batcher network it replaces is retained as the fallback for +//! other targets and as the differential oracle. +#![allow(unsafe_code)] + +use core::arch::x86_64::*; + +/// Branchless scalar min/max, from the reference's `int32_MINMAX`. +/// +/// The widened subtraction is deliberately truncated back to i32: only the sign +/// bit is used, and the wrap is what makes the comparison total across the full +/// i32 range without branching. +#[inline(always)] +#[allow(clippy::cast_possible_truncation)] +fn minmax1(a: &mut i32, b: &mut i32) { + let ab = *b ^ *a; + let mut c = (i64::from(*b) - i64::from(*a)) as i32; + c ^= ab & (c ^ *b); + c >>= 31; + c &= ab; + *a ^= c; + *b ^= c; +} + +/// Lanewise min/max of two vectors: `a` takes the minima, `b` the maxima. +#[inline] +#[target_feature(enable = "avx2")] +fn minmax8(a: &mut __m256i, b: &mut __m256i) { + let c = _mm256_min_epi32(*a, *b); + *b = _mm256_max_epi32(*a, *b); + *a = c; +} + +#[inline] +#[target_feature(enable = "avx2")] +fn ld(x: &[i32], i: usize) -> __m256i { + unsafe { _mm256_loadu_si256(x.as_ptr().add(i) as *const __m256i) } +} + +#[inline] +#[target_feature(enable = "avx2")] +fn st(x: &mut [i32], i: usize, v: __m256i) { + unsafe { _mm256_storeu_si256(x.as_mut_ptr().add(i) as *mut __m256i, v) } +} + +/// Stages 64 and 32 of a bitonic merge over 128-element blocks: four vectors +/// resident, four MINMAX ops, one store round trip. +#[target_feature(enable = "avx2")] +fn twostages_32(x: &mut [i32], n: usize) { + let mut base = 0usize; + let mut left = n; + while left > 0 { + let mut i = 0usize; + while i < 32 { + let mut x0 = ld(x, base + i); + let mut x1 = ld(x, base + i + 32); + let mut x2 = ld(x, base + i + 64); + let mut x3 = ld(x, base + i + 96); + + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + + st(x, base + i, x0); + st(x, base + i + 32, x1); + st(x, base + i + 64, x2); + st(x, base + i + 96, x3); + i += 8; + } + base += 128; + left -= 128; + } +} + +/// Stages 4q, 2q and q of a bitonic merge: eight vectors resident, twelve +/// MINMAX ops covering three stages, one load and store round trip. This is the +/// register blocking our per-pass network cannot express. +#[target_feature(enable = "avx2")] +fn threestages(x: &mut [i32], n: usize, q: usize) -> usize { + let mut k = 0usize; + while k + 8 * q <= n { + let mut i = k; + while i < k + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + let mut x4 = ld(x, i + 4 * q); + let mut x5 = ld(x, i + 5 * q); + let mut x6 = ld(x, i + 6 * q); + let mut x7 = ld(x, i + 7 * q); + + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + st(x, i + 4 * q, x4); + st(x, i + 5 * q, x5); + st(x, i + 6 * q, x6); + st(x, i + 7 * q, x7); + i += 8; + } + k += 8 * q; + } + k +} + +/// Stages 8, 4, 2 and 1 of a size-16 bitonic merge, held entirely in registers. +/// +/// This is the shape our per-pass sort cannot express: four comparator stages +/// with one load and one store, at full lane utilisation. The shuffles between +/// `minmax8` calls are what bring each stage's partners into aligned lanes. +#[target_feature(enable = "avx2")] +fn merge16_finish(x: &mut [i32], at: usize, mut x0: __m256i, mut x1: __m256i, flagdown: bool) { + minmax8(&mut x0, &mut x1); + + let mut b0 = _mm256_permute2x128_si256::<0x20>(x0, x1); // A0123 B0123 + let mut b1 = _mm256_permute2x128_si256::<0x31>(x0, x1); // A4567 B4567 + minmax8(&mut b0, &mut b1); + + let mut c0 = _mm256_unpacklo_epi64(b0, b1); // A0145 B0145 + let mut c1 = _mm256_unpackhi_epi64(b0, b1); // A2367 B2367 + minmax8(&mut c0, &mut c1); + + b0 = _mm256_unpacklo_epi32(c0, c1); // A0213 B0213 + b1 = _mm256_unpackhi_epi32(c0, c1); // A4657 B4657 + + c0 = _mm256_unpacklo_epi64(b0, b1); // A0246 B0246 + c1 = _mm256_unpackhi_epi64(b0, b1); // A1357 B1357 + minmax8(&mut c0, &mut c1); + + b0 = _mm256_unpacklo_epi32(c0, c1); // A0123 B0123 + b1 = _mm256_unpackhi_epi32(c0, c1); // A4567 B4567 + + x0 = _mm256_permute2x128_si256::<0x20>(b0, b1); + x1 = _mm256_permute2x128_si256::<0x31>(b0, b1); + + if flagdown { + let mask = _mm256_set1_epi32(-1); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + } + + st(x, at, x0); + st(x, at + 8, x1); +} + +/// Sort `x[at..at + n]` for a power-of-two `n`, ascending when `flagdown` is +/// false. Base cases only so far; larger `n` is not yet implemented. +#[target_feature(enable = "avx2")] +fn sort_2power(x: &mut [i32], at: usize, n: usize, flagdown: bool) { + if n == 8 { + // Odd-even sorting network on eight scalars. + let mut v = [0i32; 8]; + v.copy_from_slice(&x[at..at + 8]); + let [ + mut x0, + mut x1, + mut x2, + mut x3, + mut x4, + mut x5, + mut x6, + mut x7, + ] = v; + + minmax1(&mut x1, &mut x0); + minmax1(&mut x3, &mut x2); + minmax1(&mut x2, &mut x0); + minmax1(&mut x3, &mut x1); + minmax1(&mut x2, &mut x1); + + minmax1(&mut x5, &mut x4); + minmax1(&mut x7, &mut x6); + minmax1(&mut x6, &mut x4); + minmax1(&mut x7, &mut x5); + minmax1(&mut x6, &mut x5); + + minmax1(&mut x4, &mut x0); + minmax1(&mut x6, &mut x2); + minmax1(&mut x4, &mut x2); + + minmax1(&mut x5, &mut x1); + minmax1(&mut x7, &mut x3); + minmax1(&mut x5, &mut x3); + + minmax1(&mut x2, &mut x1); + minmax1(&mut x4, &mut x3); + minmax1(&mut x6, &mut x5); + + x[at..at + 8].copy_from_slice(&[x0, x1, x2, x3, x4, x5, x6, x7]); + return; + } + + if n == 16 { + let mut x0 = ld(x, at); + let mut x1 = ld(x, at + 8); + + let mut mask = _mm256_set_epi32(0, 0, -1, -1, 0, 0, -1, -1); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + + let mut b0 = _mm256_unpacklo_epi32(x0, x1); + let mut b1 = _mm256_unpackhi_epi32(x0, x1); + + let mut c0 = _mm256_unpacklo_epi64(b0, b1); + let mut c1 = _mm256_unpackhi_epi64(b0, b1); + minmax8(&mut c0, &mut c1); + + mask = _mm256_set_epi32(0, 0, -1, -1, -1, -1, 0, 0); + c0 = _mm256_xor_si256(c0, mask); + c1 = _mm256_xor_si256(c1, mask); + + b0 = _mm256_unpacklo_epi32(c0, c1); + b1 = _mm256_unpackhi_epi32(c0, c1); + minmax8(&mut b0, &mut b1); + + x0 = _mm256_unpacklo_epi64(b0, b1); + x1 = _mm256_unpackhi_epi64(b0, b1); + + b0 = _mm256_unpacklo_epi32(x0, x1); + b1 = _mm256_unpackhi_epi32(x0, x1); + + c0 = _mm256_unpacklo_epi64(b0, b1); + c1 = _mm256_unpackhi_epi64(b0, b1); + minmax8(&mut c0, &mut c1); + + b0 = _mm256_unpacklo_epi32(c0, c1); + b1 = _mm256_unpackhi_epi32(c0, c1); + + b0 = _mm256_xor_si256(b0, mask); + b1 = _mm256_xor_si256(b1, mask); + + c0 = _mm256_permute2x128_si256::<0x20>(b0, b1); + c1 = _mm256_permute2x128_si256::<0x31>(b0, b1); + minmax8(&mut c0, &mut c1); + + b0 = _mm256_permute2x128_si256::<0x20>(c0, c1); + b1 = _mm256_permute2x128_si256::<0x31>(c0, c1); + minmax8(&mut b0, &mut b1); + + x0 = _mm256_unpacklo_epi64(b0, b1); + x1 = _mm256_unpackhi_epi64(b0, b1); + + b0 = _mm256_unpacklo_epi32(x0, x1); + b1 = _mm256_unpackhi_epi32(x0, x1); + + c0 = _mm256_unpacklo_epi64(b0, b1); + c1 = _mm256_unpackhi_epi64(b0, b1); + minmax8(&mut c0, &mut c1); + + b0 = _mm256_unpacklo_epi32(c0, c1); + b1 = _mm256_unpackhi_epi32(c0, c1); + + x0 = _mm256_unpacklo_epi64(b0, b1); + x1 = _mm256_unpackhi_epi64(b0, b1); + + mask = _mm256_set1_epi32(-1); + if flagdown { + x1 = _mm256_xor_si256(x1, mask); + } else { + x0 = _mm256_xor_si256(x0, mask); + } + + merge16_finish(x, at, x0, x1, flagdown); + return; + } + + if n == 32 { + sort_2power(x, at, 16, true); + sort_2power(x, at + 16, 16, false); + + let mut x0 = ld(x, at); + let mut x1 = ld(x, at + 8); + let mut x2 = ld(x, at + 16); + let mut x3 = ld(x, at + 24); + + if flagdown { + let mask = _mm256_set1_epi32(-1); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + x2 = _mm256_xor_si256(x2, mask); + x3 = _mm256_xor_si256(x3, mask); + } + + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + + merge16_finish(x, at, x0, x1, flagdown); + merge16_finish(x, at + 16, x2, x3, flagdown); + return; + } + + // ---- general power-of-two path (ported from the reference) ---- + let mut mask; + let mut p: usize; + let mut q: usize; + p = n >> 3; + let mut i = 0; + while i < p { + let mut x0 = ld(x, i); + let mut x2 = ld(x, i + 2 * p); + let mut x4 = ld(x, i + 4 * p); + let mut x6 = ld(x, i + 6 * p); + + // odd-even stage instead of bitonic stage + + minmax8(&mut x4, &mut x0); + minmax8(&mut x6, &mut x2); + minmax8(&mut x2, &mut x0); + minmax8(&mut x6, &mut x4); + minmax8(&mut x2, &mut x4); + + st(x, i, x0); + st(x, i + 2 * p, x2); + st(x, i + 4 * p, x4); + st(x, i + 6 * p, x6); + + let mut x1 = ld(x, i + p); + let mut x3 = ld(x, i + 3 * p); + let mut x5 = ld(x, i + 5 * p); + let mut x7 = ld(x, i + 7 * p); + + minmax8(&mut x1, &mut x5); + minmax8(&mut x3, &mut x7); + minmax8(&mut x1, &mut x3); + minmax8(&mut x5, &mut x7); + minmax8(&mut x5, &mut x3); + + st(x, i + p, x1); + st(x, i + 3 * p, x3); + st(x, i + 5 * p, x5); + st(x, i + 7 * p, x7); + i += 8; + } + + if n >= 128 { + let mut flip; + let mut flipflip; + + mask = _mm256_set1_epi32(-1); + + let mut j = 0; + while j < n { + let mut x0 = ld(x, j); + let mut x1 = ld(x, j + 16); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + st(x, j, x0); + st(x, j + 16, x1); + j += 32; + } + + p = 8; + loop { + q = p >> 1; + while q >= 128 { + threestages(x, n, q >> 2); + q >>= 3; + } + if q == 64 { + twostages_32(x, n); + q = 16; + } + if q == 32 { + q = 8; + let mut k = 0; + while k < n { + let mut i = k; + while i < k + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + let mut x4 = ld(x, i + 4 * q); + let mut x5 = ld(x, i + 5 * q); + let mut x6 = ld(x, i + 6 * q); + let mut x7 = ld(x, i + 7 * q); + + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + st(x, i + 4 * q, x4); + st(x, i + 5 * q, x5); + st(x, i + 6 * q, x6); + st(x, i + 7 * q, x7); + i += 8; + } + k += 8 * q; + } + q = 4; + } + if q == 16 { + q = 8; + let mut k = 0; + while k < n { + let mut i = k; + while i < k + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + i += 8; + } + k += 4 * q; + } + q = 4; + } + if q == 8 { + let mut k = 0; + while k < n { + let mut x0 = ld(x, k); + let mut x1 = ld(x, k + q); + + minmax8(&mut x0, &mut x1); + + st(x, k, x0); + st(x, k + q, x1); + k += q + q; + } + } + + q = n >> 3; + flip = 0; + if p << 1 == q { + flip = 1; + } + flipflip = 1 - flip; + let mut j = 0; + while j < q { + let mut k = j; + while k < j + p + p { + let mut i = k; + while i < k + p { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + let mut x4 = ld(x, i + 4 * q); + let mut x5 = ld(x, i + 5 * q); + let mut x6 = ld(x, i + 6 * q); + let mut x7 = ld(x, i + 7 * q); + + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + + if flip != 0 { + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + x2 = _mm256_xor_si256(x2, mask); + x3 = _mm256_xor_si256(x3, mask); + x4 = _mm256_xor_si256(x4, mask); + x5 = _mm256_xor_si256(x5, mask); + x6 = _mm256_xor_si256(x6, mask); + x7 = _mm256_xor_si256(x7, mask); + } + + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + st(x, i + 4 * q, x4); + st(x, i + 5 * q, x5); + st(x, i + 6 * q, x6); + st(x, i + 7 * q, x7); + i += 8; + } + flip ^= 1; + k += p; + } + flip ^= flipflip; + j += p + p; + } + + if p << 4 == n { + break; + } + p <<= 1; + } + } + + let mut p = 4; + while p >= 1 { + let mut zi = 0usize; + if p == 4 { + mask = _mm256_set_epi32(0, 0, 0, 0, -1, -1, -1, -1); + while zi != n { + let mut x0 = ld(x, zi); + let mut x1 = ld(x, zi + 8); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + st(x, zi, x0); + st(x, zi + 8, x1); + zi += 16; + } + } else if p == 2 { + mask = _mm256_set_epi32(0, 0, -1, -1, -1, -1, 0, 0); + while zi != n { + let mut x0 = ld(x, zi); + let mut x1 = ld(x, zi + 8); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + let mut b0 = _mm256_permute2x128_si256(x0, x1, 0x20); + let mut b1 = _mm256_permute2x128_si256(x0, x1, 0x31); + minmax8(&mut b0, &mut b1); + let c0 = _mm256_permute2x128_si256(b0, b1, 0x20); + let c1 = _mm256_permute2x128_si256(b0, b1, 0x31); + st(x, zi, c0); + st(x, zi + 8, c1); + zi += 16; + } + } else { + mask = _mm256_set_epi32(0, -1, -1, 0, 0, -1, -1, 0); + while zi != n { + let mut x0 = ld(x, zi); + let mut x1 = ld(x, zi + 8); + x0 = _mm256_xor_si256(x0, mask); + x1 = _mm256_xor_si256(x1, mask); + let b0 = _mm256_permute2x128_si256(x0, x1, 0x20); + let b1 = _mm256_permute2x128_si256(x0, x1, 0x31); + let mut c0 = _mm256_unpacklo_epi64(b0, b1); + let mut c1 = _mm256_unpackhi_epi64(b0, b1); + minmax8(&mut c0, &mut c1); + let mut d0 = _mm256_unpacklo_epi64(c0, c1); + let mut d1 = _mm256_unpackhi_epi64(c0, c1); + minmax8(&mut d0, &mut d1); + let e0 = _mm256_permute2x128_si256(d0, d1, 0x20); + let e1 = _mm256_permute2x128_si256(d0, d1, 0x31); + st(x, zi, e0); + st(x, zi + 8, e1); + zi += 16; + } + } + + q = n >> 4; + while q >= 128 || q == 32 { + threestages(x, n, q >> 2); + q >>= 3; + } + while q >= 16 { + q >>= 1; + let mut j = 0; + while j < n { + let mut k = j; + while k < j + q { + let mut x0 = ld(x, k); + let mut x1 = ld(x, k + q); + let mut x2 = ld(x, k + 2 * q); + let mut x3 = ld(x, k + 3 * q); + + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + + st(x, k, x0); + st(x, k + q, x1); + st(x, k + 2 * q, x2); + st(x, k + 3 * q, x3); + k += 8; + } + j += 4 * q; + } + q >>= 1; + } + if q == 8 { + let mut j = 0; + while j < n { + let mut x0 = ld(x, j); + let mut x1 = ld(x, j + q); + + minmax8(&mut x0, &mut x1); + + st(x, j, x0); + st(x, j + q, x1); + j += 2 * q; + } + } + + q = n >> 3; + let mut k = 0; + while k < q { + let mut x0 = ld(x, k); + let mut x1 = ld(x, k + q); + let mut x2 = ld(x, k + 2 * q); + let mut x3 = ld(x, k + 3 * q); + let mut x4 = ld(x, k + 4 * q); + let mut x5 = ld(x, k + 5 * q); + let mut x6 = ld(x, k + 6 * q); + let mut x7 = ld(x, k + 7 * q); + + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + + st(x, k, x0); + st(x, k + q, x1); + st(x, k + 2 * q, x2); + st(x, k + 3 * q, x3); + st(x, k + 4 * q, x4); + st(x, k + 5 * q, x5); + st(x, k + 6 * q, x6); + st(x, k + 7 * q, x7); + k += 8; + } + p >>= 1; + } + + // everything is still masked with _mm256_set_epi32(0,-1,0,-1,0,-1,0,-1); + mask = _mm256_set1_epi32(-1); + + let mut i = 0; + while i < n { + let a0 = ld(x, i); + let a1 = ld(x, i + 8); + let a2 = ld(x, i + 16); + let a3 = ld(x, i + 24); + let a4 = ld(x, i + 32); + let a5 = ld(x, i + 40); + let a6 = ld(x, i + 48); + let a7 = ld(x, i + 56); + + let b0 = _mm256_unpacklo_epi32(a0, a1); + let b1 = _mm256_unpackhi_epi32(a0, a1); + let b2 = _mm256_unpacklo_epi32(a2, a3); + let b3 = _mm256_unpackhi_epi32(a2, a3); + let b4 = _mm256_unpacklo_epi32(a4, a5); + let b5 = _mm256_unpackhi_epi32(a4, a5); + let b6 = _mm256_unpacklo_epi32(a6, a7); + let b7 = _mm256_unpackhi_epi32(a6, a7); + + let mut c0 = _mm256_unpacklo_epi64(b0, b2); + let mut c1 = _mm256_unpacklo_epi64(b1, b3); + let mut c2 = _mm256_unpackhi_epi64(b0, b2); + let mut c3 = _mm256_unpackhi_epi64(b1, b3); + let mut c4 = _mm256_unpacklo_epi64(b4, b6); + let mut c5 = _mm256_unpacklo_epi64(b5, b7); + let mut c6 = _mm256_unpackhi_epi64(b4, b6); + let mut c7 = _mm256_unpackhi_epi64(b5, b7); + + if flagdown { + c2 = _mm256_xor_si256(c2, mask); + c3 = _mm256_xor_si256(c3, mask); + c6 = _mm256_xor_si256(c6, mask); + c7 = _mm256_xor_si256(c7, mask); + } else { + c0 = _mm256_xor_si256(c0, mask); + c1 = _mm256_xor_si256(c1, mask); + c4 = _mm256_xor_si256(c4, mask); + c5 = _mm256_xor_si256(c5, mask); + } + + let mut d0 = _mm256_permute2x128_si256(c0, c4, 0x20); + let mut d1 = _mm256_permute2x128_si256(c2, c6, 0x20); + let mut d2 = _mm256_permute2x128_si256(c1, c5, 0x20); + let mut d3 = _mm256_permute2x128_si256(c3, c7, 0x20); + let mut d4 = _mm256_permute2x128_si256(c0, c4, 0x31); + let mut d5 = _mm256_permute2x128_si256(c2, c6, 0x31); + let mut d6 = _mm256_permute2x128_si256(c1, c5, 0x31); + let mut d7 = _mm256_permute2x128_si256(c3, c7, 0x31); + + minmax8(&mut d0, &mut d1); + minmax8(&mut d2, &mut d3); + minmax8(&mut d4, &mut d5); + minmax8(&mut d6, &mut d7); + minmax8(&mut d0, &mut d2); + minmax8(&mut d1, &mut d3); + minmax8(&mut d4, &mut d6); + minmax8(&mut d5, &mut d7); + minmax8(&mut d0, &mut d4); + minmax8(&mut d1, &mut d5); + minmax8(&mut d2, &mut d6); + minmax8(&mut d3, &mut d7); + + let e0 = _mm256_unpacklo_epi32(d0, d1); + let e1 = _mm256_unpackhi_epi32(d0, d1); + let e2 = _mm256_unpacklo_epi32(d2, d3); + let e3 = _mm256_unpackhi_epi32(d2, d3); + let e4 = _mm256_unpacklo_epi32(d4, d5); + let e5 = _mm256_unpackhi_epi32(d4, d5); + let e6 = _mm256_unpacklo_epi32(d6, d7); + let e7 = _mm256_unpackhi_epi32(d6, d7); + + let f0 = _mm256_unpacklo_epi64(e0, e2); + let f1 = _mm256_unpacklo_epi64(e1, e3); + let f2 = _mm256_unpackhi_epi64(e0, e2); + let f3 = _mm256_unpackhi_epi64(e1, e3); + let f4 = _mm256_unpacklo_epi64(e4, e6); + let f5 = _mm256_unpacklo_epi64(e5, e7); + let f6 = _mm256_unpackhi_epi64(e4, e6); + let f7 = _mm256_unpackhi_epi64(e5, e7); + + let g0 = _mm256_permute2x128_si256(f0, f4, 0x20); + let g1 = _mm256_permute2x128_si256(f2, f6, 0x20); + let g2 = _mm256_permute2x128_si256(f1, f5, 0x20); + let g3 = _mm256_permute2x128_si256(f3, f7, 0x20); + let g4 = _mm256_permute2x128_si256(f0, f4, 0x31); + let g5 = _mm256_permute2x128_si256(f2, f6, 0x31); + let g6 = _mm256_permute2x128_si256(f1, f5, 0x31); + let g7 = _mm256_permute2x128_si256(f3, f7, 0x31); + + st(x, i, g0); + st(x, i + 8, g1); + st(x, i + 16, g2); + st(x, i + 24, g3); + st(x, i + 32, g4); + st(x, i + 40, g5); + st(x, i + 48, g6); + st(x, i + 56, g7); + i += 64; + } + + q = n >> 4; + while q >= 128 || q == 32 { + q >>= 2; + let mut j = 0; + while j < n { + let mut i = j; + while i < j + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + let mut x4 = ld(x, i + 4 * q); + let mut x5 = ld(x, i + 5 * q); + let mut x6 = ld(x, i + 6 * q); + let mut x7 = ld(x, i + 7 * q); + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + st(x, i + 4 * q, x4); + st(x, i + 5 * q, x5); + st(x, i + 6 * q, x6); + st(x, i + 7 * q, x7); + i += 8; + } + j += 8 * q; + } + q >>= 1; + } + while q >= 16 { + q >>= 1; + let mut j = 0; + while j < n { + let mut i = j; + while i < j + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + i += 8; + } + j += 4 * q; + } + q >>= 1; + } + if q == 8 { + let mut j = 0; + while j < n { + let mut x0 = ld(x, j); + let mut x1 = ld(x, j + q); + minmax8(&mut x0, &mut x1); + st(x, j, x0); + st(x, j + q, x1); + j += q + q; + } + } + + q = n >> 3; + let mut i = 0; + while i < q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + let mut x4 = ld(x, i + 4 * q); + let mut x5 = ld(x, i + 5 * q); + let mut x6 = ld(x, i + 6 * q); + let mut x7 = ld(x, i + 7 * q); + + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + + let b0 = _mm256_unpacklo_epi32(x0, x4); + let b1 = _mm256_unpackhi_epi32(x0, x4); + let b2 = _mm256_unpacklo_epi32(x1, x5); + let b3 = _mm256_unpackhi_epi32(x1, x5); + let b4 = _mm256_unpacklo_epi32(x2, x6); + let b5 = _mm256_unpackhi_epi32(x2, x6); + let b6 = _mm256_unpacklo_epi32(x3, x7); + let b7 = _mm256_unpackhi_epi32(x3, x7); + + let c0 = _mm256_unpacklo_epi64(b0, b4); + let c1 = _mm256_unpacklo_epi64(b1, b5); + let c2 = _mm256_unpackhi_epi64(b0, b4); + let c3 = _mm256_unpackhi_epi64(b1, b5); + let c4 = _mm256_unpacklo_epi64(b2, b6); + let c5 = _mm256_unpacklo_epi64(b3, b7); + let c6 = _mm256_unpackhi_epi64(b2, b6); + let c7 = _mm256_unpackhi_epi64(b3, b7); + + let mut d0 = _mm256_permute2x128_si256(c0, c4, 0x20); + let mut d1 = _mm256_permute2x128_si256(c1, c5, 0x20); + let mut d2 = _mm256_permute2x128_si256(c2, c6, 0x20); + let mut d3 = _mm256_permute2x128_si256(c3, c7, 0x20); + let mut d4 = _mm256_permute2x128_si256(c0, c4, 0x31); + let mut d5 = _mm256_permute2x128_si256(c1, c5, 0x31); + let mut d6 = _mm256_permute2x128_si256(c2, c6, 0x31); + let mut d7 = _mm256_permute2x128_si256(c3, c7, 0x31); + + if flagdown { + d0 = _mm256_xor_si256(d0, mask); + d1 = _mm256_xor_si256(d1, mask); + d2 = _mm256_xor_si256(d2, mask); + d3 = _mm256_xor_si256(d3, mask); + d4 = _mm256_xor_si256(d4, mask); + d5 = _mm256_xor_si256(d5, mask); + d6 = _mm256_xor_si256(d6, mask); + d7 = _mm256_xor_si256(d7, mask); + } + + st(x, i, d0); + st(x, i + q, d4); + st(x, i + 2 * q, d1); + st(x, i + 3 * q, d5); + st(x, i + 4 * q, d2); + st(x, i + 5 * q, d6); + st(x, i + 6 * q, d3); + st(x, i + 7 * q, d7); + i += 8; + } +} + +/// Scalar min/max on two slice positions. +#[inline] +fn minmax_at(x: &mut [i32], i: usize, j: usize) { + let (mut a, mut b) = (x[i], x[j]); + minmax1(&mut a, &mut b); + x[i] = a; + x[j] = b; +} + +/// Lanewise min/max between two windows `x[a..]` and `x[b..]` of length `len`, +/// handling a ragged tail by overlapping the final vector. +#[target_feature(enable = "avx2")] +fn minmax_vector(x: &mut [i32], a: usize, b: usize, len: usize) { + let mut n = len; + if n < 8 { + for t in 0..n { + minmax_at(x, a + t, b + t); + } + return; + } + if n & 7 != 0 { + let mut x0 = ld(x, a + n - 8); + let mut y0 = ld(x, b + n - 8); + minmax8(&mut x0, &mut y0); + st(x, a + n - 8, x0); + st(x, b + n - 8, y0); + n &= !7; + } + let mut o = 0usize; + while o < n { + let mut x0 = ld(x, a + o); + let mut y0 = ld(x, b + o); + minmax8(&mut x0, &mut y0); + st(x, a + o, x0); + st(x, b + o, y0); + o += 8; + } +} + +/// Sort `x[..n]` ascending, for any `n` (the reference's `int32_sort`). +#[target_feature(enable = "avx2")] +pub fn sort(x: &mut [i32], n: usize) { + if n <= 8 { + if n == 8 { + minmax_at(x, 0, 1); + minmax_at(x, 1, 2); + minmax_at(x, 2, 3); + minmax_at(x, 3, 4); + minmax_at(x, 4, 5); + minmax_at(x, 5, 6); + minmax_at(x, 6, 7); + } + if n >= 7 { + minmax_at(x, 0, 1); + minmax_at(x, 1, 2); + minmax_at(x, 2, 3); + minmax_at(x, 3, 4); + minmax_at(x, 4, 5); + minmax_at(x, 5, 6); + } + if n >= 6 { + minmax_at(x, 0, 1); + minmax_at(x, 1, 2); + minmax_at(x, 2, 3); + minmax_at(x, 3, 4); + minmax_at(x, 4, 5); + } + if n >= 5 { + minmax_at(x, 0, 1); + minmax_at(x, 1, 2); + minmax_at(x, 2, 3); + minmax_at(x, 3, 4); + } + if n >= 4 { + minmax_at(x, 0, 1); + minmax_at(x, 1, 2); + minmax_at(x, 2, 3); + } + if n >= 3 { + minmax_at(x, 0, 1); + minmax_at(x, 1, 2); + } + if n >= 2 { + minmax_at(x, 0, 1); + } + return; + } + + if n & (n - 1) == 0 { + sort_2power(x, 0, n, false); + return; + } + + let mut q: usize = 8; + let mut j: usize = 0; + while q < n - q { + q += q; + } + // n > q >= 8 + + if q <= 128 { + // n <= 256: pad to the next power of two with sentinels that sort to the + // end, sort that, and copy back. The reference type-puns an + // `int32x8[32]` for alignment; a plain array suffices here. + let mut y = [0x7fff_ffffi32; 256]; + y[..n].copy_from_slice(&x[..n]); + sort_2power(&mut y, 0, 2 * q, false); + x[..n].copy_from_slice(&y[..n]); + return; + } + + sort_2power(x, 0, q, true); + sort(&mut x[q..], n - q); + + while q >= 64 { + q >>= 2; + j = threestages(x, n, q); + minmax_vector(x, j, j + 4 * q, n.saturating_sub(4 * q).saturating_sub(j)); + if j + 4 * q <= n { + let mut i = j; + while i < j + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + let mut x2 = ld(x, i + 2 * q); + let mut x3 = ld(x, i + 3 * q); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + st(x, i, x0); + st(x, i + q, x1); + st(x, i + 2 * q, x2); + st(x, i + 3 * q, x3); + i += 8; + } + j += 4 * q; + } + minmax_vector(x, j, j + 2 * q, n.saturating_sub(2 * q).saturating_sub(j)); + if j + 2 * q <= n { + let mut i = j; + while i < j + q { + let mut x0 = ld(x, i); + let mut x1 = ld(x, i + q); + minmax8(&mut x0, &mut x1); + st(x, i, x0); + st(x, i + q, x1); + i += 8; + } + j += 2 * q; + } + minmax_vector(x, j, j + q, n.saturating_sub(q).saturating_sub(j)); + q >>= 1; + } + if q == 32 { + j = 0; + while j + 64 <= n { + let mut x0 = ld(x, j); + let mut x1 = ld(x, j + 8); + let mut x2 = ld(x, j + 16); + let mut x3 = ld(x, j + 24); + let mut x4 = ld(x, j + 32); + let mut x5 = ld(x, j + 40); + let mut x6 = ld(x, j + 48); + let mut x7 = ld(x, j + 56); + minmax8(&mut x0, &mut x4); + minmax8(&mut x1, &mut x5); + minmax8(&mut x2, &mut x6); + minmax8(&mut x3, &mut x7); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x4, &mut x6); + minmax8(&mut x5, &mut x7); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + minmax8(&mut x4, &mut x5); + minmax8(&mut x6, &mut x7); + let mut a0 = _mm256_permute2x128_si256(x0, x1, 0x20); + let mut a1 = _mm256_permute2x128_si256(x0, x1, 0x31); + let mut a2 = _mm256_permute2x128_si256(x2, x3, 0x20); + let mut a3 = _mm256_permute2x128_si256(x2, x3, 0x31); + let mut a4 = _mm256_permute2x128_si256(x4, x5, 0x20); + let mut a5 = _mm256_permute2x128_si256(x4, x5, 0x31); + let mut a6 = _mm256_permute2x128_si256(x6, x7, 0x20); + let mut a7 = _mm256_permute2x128_si256(x6, x7, 0x31); + minmax8(&mut a0, &mut a1); + minmax8(&mut a2, &mut a3); + minmax8(&mut a4, &mut a5); + minmax8(&mut a6, &mut a7); + let b0 = _mm256_permute2x128_si256(a0, a1, 0x20); + let b1 = _mm256_permute2x128_si256(a0, a1, 0x31); + let b2 = _mm256_permute2x128_si256(a2, a3, 0x20); + let b3 = _mm256_permute2x128_si256(a2, a3, 0x31); + let b4 = _mm256_permute2x128_si256(a4, a5, 0x20); + let b5 = _mm256_permute2x128_si256(a4, a5, 0x31); + let b6 = _mm256_permute2x128_si256(a6, a7, 0x20); + let b7 = _mm256_permute2x128_si256(a6, a7, 0x31); + let mut c0 = _mm256_unpacklo_epi64(b0, b1); + let mut c1 = _mm256_unpackhi_epi64(b0, b1); + let mut c2 = _mm256_unpacklo_epi64(b2, b3); + let mut c3 = _mm256_unpackhi_epi64(b2, b3); + let mut c4 = _mm256_unpacklo_epi64(b4, b5); + let mut c5 = _mm256_unpackhi_epi64(b4, b5); + let mut c6 = _mm256_unpacklo_epi64(b6, b7); + let mut c7 = _mm256_unpackhi_epi64(b6, b7); + minmax8(&mut c0, &mut c1); + minmax8(&mut c2, &mut c3); + minmax8(&mut c4, &mut c5); + minmax8(&mut c6, &mut c7); + let d0 = _mm256_unpacklo_epi32(c0, c1); + let d1 = _mm256_unpackhi_epi32(c0, c1); + let d2 = _mm256_unpacklo_epi32(c2, c3); + let d3 = _mm256_unpackhi_epi32(c2, c3); + let d4 = _mm256_unpacklo_epi32(c4, c5); + let d5 = _mm256_unpackhi_epi32(c4, c5); + let d6 = _mm256_unpacklo_epi32(c6, c7); + let d7 = _mm256_unpackhi_epi32(c6, c7); + let mut e0 = _mm256_unpacklo_epi64(d0, d1); + let mut e1 = _mm256_unpackhi_epi64(d0, d1); + let mut e2 = _mm256_unpacklo_epi64(d2, d3); + let mut e3 = _mm256_unpackhi_epi64(d2, d3); + let mut e4 = _mm256_unpacklo_epi64(d4, d5); + let mut e5 = _mm256_unpackhi_epi64(d4, d5); + let mut e6 = _mm256_unpacklo_epi64(d6, d7); + let mut e7 = _mm256_unpackhi_epi64(d6, d7); + minmax8(&mut e0, &mut e1); + minmax8(&mut e2, &mut e3); + minmax8(&mut e4, &mut e5); + minmax8(&mut e6, &mut e7); + let f0 = _mm256_unpacklo_epi32(e0, e1); + let f1 = _mm256_unpackhi_epi32(e0, e1); + let f2 = _mm256_unpacklo_epi32(e2, e3); + let f3 = _mm256_unpackhi_epi32(e2, e3); + let f4 = _mm256_unpacklo_epi32(e4, e5); + let f5 = _mm256_unpackhi_epi32(e4, e5); + let f6 = _mm256_unpacklo_epi32(e6, e7); + let f7 = _mm256_unpackhi_epi32(e6, e7); + st(x, j, f0); + st(x, j + 8, f1); + st(x, j + 16, f2); + st(x, j + 24, f3); + st(x, j + 32, f4); + st(x, j + 40, f5); + st(x, j + 48, f6); + st(x, j + 56, f7); + j += 64; + } + minmax_vector(x, j, j + 32, n.saturating_sub(32).saturating_sub(j)); + // goto continue16; (restructured as fallthrough) + } + // `goto continue16` from the q == 32 branch lands inside this body, past the + // `j = 0`, carrying j forward — hence the guarded reset. + if q >= 16 { + if q == 16 { + j = 0; + } + while j + 32 <= n { + let mut x0 = ld(x, j); + let mut x1 = ld(x, j + 8); + let mut x2 = ld(x, j + 16); + let mut x3 = ld(x, j + 24); + minmax8(&mut x0, &mut x2); + minmax8(&mut x1, &mut x3); + minmax8(&mut x0, &mut x1); + minmax8(&mut x2, &mut x3); + let mut a0 = _mm256_permute2x128_si256(x0, x1, 0x20); + let mut a1 = _mm256_permute2x128_si256(x0, x1, 0x31); + let mut a2 = _mm256_permute2x128_si256(x2, x3, 0x20); + let mut a3 = _mm256_permute2x128_si256(x2, x3, 0x31); + minmax8(&mut a0, &mut a1); + minmax8(&mut a2, &mut a3); + let b0 = _mm256_permute2x128_si256(a0, a1, 0x20); + let b1 = _mm256_permute2x128_si256(a0, a1, 0x31); + let b2 = _mm256_permute2x128_si256(a2, a3, 0x20); + let b3 = _mm256_permute2x128_si256(a2, a3, 0x31); + let mut c0 = _mm256_unpacklo_epi64(b0, b1); + let mut c1 = _mm256_unpackhi_epi64(b0, b1); + let mut c2 = _mm256_unpacklo_epi64(b2, b3); + let mut c3 = _mm256_unpackhi_epi64(b2, b3); + minmax8(&mut c0, &mut c1); + minmax8(&mut c2, &mut c3); + let d0 = _mm256_unpacklo_epi32(c0, c1); + let d1 = _mm256_unpackhi_epi32(c0, c1); + let d2 = _mm256_unpacklo_epi32(c2, c3); + let d3 = _mm256_unpackhi_epi32(c2, c3); + let mut e0 = _mm256_unpacklo_epi64(d0, d1); + let mut e1 = _mm256_unpackhi_epi64(d0, d1); + let mut e2 = _mm256_unpacklo_epi64(d2, d3); + let mut e3 = _mm256_unpackhi_epi64(d2, d3); + minmax8(&mut e0, &mut e1); + minmax8(&mut e2, &mut e3); + let f0 = _mm256_unpacklo_epi32(e0, e1); + let f1 = _mm256_unpackhi_epi32(e0, e1); + let f2 = _mm256_unpacklo_epi32(e2, e3); + let f3 = _mm256_unpackhi_epi32(e2, e3); + st(x, j, f0); + st(x, j + 8, f1); + st(x, j + 16, f2); + st(x, j + 24, f3); + j += 32; + } + minmax_vector(x, j, j + 16, n.saturating_sub(16).saturating_sub(j)); + // goto continue8; (restructured as fallthrough) + } + // q == 8; `goto continue8` from the block above lands here past the reset. + if q == 8 { + j = 0; + } + while j + 16 <= n { + let mut x0 = ld(x, j); + let mut x1 = ld(x, j + 8); + minmax8(&mut x0, &mut x1); + st(x, j, x0); + st(x, j + 8, x1); + let mut a0 = _mm256_permute2x128_si256(x0, x1, 0x20); + let mut a1 = _mm256_permute2x128_si256(x0, x1, 0x31); + minmax8(&mut a0, &mut a1); + let b0 = _mm256_permute2x128_si256(a0, a1, 0x20); + let b1 = _mm256_permute2x128_si256(a0, a1, 0x31); + let mut c0 = _mm256_unpacklo_epi64(b0, b1); + let mut c1 = _mm256_unpackhi_epi64(b0, b1); + minmax8(&mut c0, &mut c1); + let d0 = _mm256_unpacklo_epi32(c0, c1); + let d1 = _mm256_unpackhi_epi32(c0, c1); + let mut e0 = _mm256_unpacklo_epi64(d0, d1); + let mut e1 = _mm256_unpackhi_epi64(d0, d1); + minmax8(&mut e0, &mut e1); + let f0 = _mm256_unpacklo_epi32(e0, e1); + let f1 = _mm256_unpackhi_epi32(e0, e1); + st(x, j, f0); + st(x, j + 8, f1); + j += 16; + } + minmax_vector(x, j, j + 8, n.saturating_sub(8).saturating_sub(j)); + if j + 8 <= n { + minmax_at(x, j, j + 4); + minmax_at(x, j + 1, j + 5); + minmax_at(x, j + 2, j + 6); + minmax_at(x, j + 3, j + 7); + minmax_at(x, j, j + 2); + minmax_at(x, j + 1, j + 3); + minmax_at(x, j, j + 1); + minmax_at(x, j + 2, j + 3); + minmax_at(x, j + 4, j + 6); + minmax_at(x, j + 5, j + 7); + minmax_at(x, j + 4, j + 5); + minmax_at(x, j + 6, j + 7); + j += 8; + } + minmax_vector(x, j, j + 4, n.saturating_sub(4).saturating_sub(j)); + if j + 4 <= n { + minmax_at(x, j, j + 2); + minmax_at(x, j + 1, j + 3); + minmax_at(x, j, j + 1); + minmax_at(x, j + 2, j + 3); + j += 4; + } + if j + 3 <= n { + minmax_at(x, j, j + 2); + } + if j + 2 <= n { + minmax_at(x, j, j + 1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + state.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + /// The full `sort` entry point must sort every length correctly, including + /// the six parameter sizes and the awkward boundaries around each internal + /// dispatch (power-of-two, the <=256 padding path, and the recursive path). + #[test] + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + fn full_sort_matches_reference_ordering() { + if !crate::cpu::has_avx2() { + return; + } + let mut s = 0xfeed_face_dead_beefu64; + let lens = [ + 0usize, 1, 2, 3, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 100, 128, 200, 255, 256, + 257, 300, 511, 512, 513, 653, 761, 857, 953, 1013, 1024, 1277, 2048, + ]; + for &n in &lens { + for pattern in 0..7 { + let mut x: Vec = (0..n) + .map(|i| match pattern { + 0..=2 => next(&mut s) as i32, + 3 => 42, + 4 => i as i32, + 5 => (n - i) as i32, + _ => (next(&mut s) % 4) as i32, + }) + .collect(); + let mut want = x.clone(); + want.sort_unstable(); + // SAFETY: AVX2 confirmed above. + unsafe { sort(&mut x, n) }; + assert_eq!(x, want, "full sort n={n} pattern={pattern}"); + } + } + } + + /// The ported base cases must reproduce the reference's ordering exactly, + /// on random data and on the degenerate patterns that break naive networks. + /// + /// `flagdown` selects direction: `false` ascending, `true` descending. The + /// `n == 8` network ignores it and always descends — that is the reference's + /// documented precondition ("if n == 8 then flagdown"), not an omission. + #[test] + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + fn base_cases_sort_correctly() { + if !crate::cpu::has_avx2() { + return; + } + let mut s = 0x1234_5678_9abc_def1u64; + for &n in &[8usize, 16, 32, 64, 128, 256, 512, 1024] { + // n == 8 is descending-only by contract. + let dirs: &[bool] = if n == 8 { &[true] } else { &[false, true] }; + for &flagdown in dirs { + for pattern in 0..7 { + let mut x: Vec = (0..n) + .map(|i| match pattern { + 0..=2 => next(&mut s) as i32, + 3 => 42, + 4 => i as i32, + 5 => (n - i) as i32, + _ => (next(&mut s) % 4) as i32, + }) + .collect(); + let mut want = x.clone(); + want.sort_unstable(); + if flagdown { + want.reverse(); + } + // SAFETY: AVX2 confirmed above. + unsafe { sort_2power(&mut x, 0, n, flagdown) }; + assert_eq!(x, want, "djbsort n={n} down={flagdown} pattern={pattern}"); + } + } + } + } +}