ORE v2 (4/n): NEON + AVX2 SIMD backends

packages/ore-rs/src/primitives.rs

@@ -1,6 +1,7 @@
 pub mod hash;
 pub mod prf;
 pub mod prp;
+pub mod simd;
 
 use aes::cipher::{consts::U16, generic_array::GenericArray};
 use aes::Block;
@@ -10,36 +11,6 @@ pub type PrfKey = GenericArray<u8, U16>;
 pub type HashKey = GenericArray<u8, U16>;
 pub const NONCE_SIZE: usize = 16;
 
-/// Pack one bit per element of `src` into `out`, LSB-first, eight elements per
-/// byte: the bit derived from `src[j]` lands in `out[j / 8]` at position
-/// `j % 8` — the bit order used by `RightBlock32::set_bit`. `bit` extracts the
-/// 0/1 value of each element; `merge` folds each packed byte into the existing
-/// `out` byte (assign for a fresh buffer, XOR to overlay). This is the single
-/// home of the right-ciphertext bit-packing convention, shared by
-/// [`Hash::hash_all_into`] and [`Prp::indicator_mask_xor`].
-///
-/// Panics unless `out.len() * 8 == src.len()`, which also guarantees
-/// `chunks_exact(8)` consumes `src` with no dropped remainder.
-pub(crate) fn pack_bits_lsb_first<T>(
-    out: &mut [u8],
-    src: &[T],
-    bit: impl Fn(&T) -> u8,
-    merge: impl Fn(&mut u8, u8),
-) {
-    assert_eq!(
-        out.len() * 8,
-        src.len(),
-        "pack_bits_lsb_first: out.len()*8 must equal src.len()"
-    );
-    for (slot, chunk) in out.iter_mut().zip(src.chunks_exact(8)) {
-        let mut byte = 0u8;
-        for (i, elem) in chunk.iter().enumerate() {
-            byte |= (bit(elem) & 1) << i;
-        }
-        merge(slot, byte);
-    }
-}
-
 pub trait Prf {
     fn new(key: &PrfKey) -> Self;
     fn encrypt_all(&self, data: &mut [AesBlock]);
@@ -69,9 +40,9 @@ pub trait Prp<T>: Sized {
     fn invert(&self, data: T) -> PrpResult<T>;
 
     /// XOR the indicator mask for `data` into `out`: bit `j` of the mask is
-    /// `1` iff `invert(j) > data`. Bit order matches `RightBlock32::set_bit`
-    /// (LSB-first within each byte). `out.len() * 8` must equal the
-    /// permutation domain.
+    /// `1` iff `invert(j) > data`. Bit order matches
+    /// `RightBitVec::set_bit` (LSB-first within each byte). `out.len() * 8`
+    /// must equal the permutation domain.
     ///
     /// This is the bulk form of the per-`j` `invert`-and-compare loop the
     /// right-ciphertext encoder needs; implementations walk their inverse

packages/ore-rs/src/primitives/hash.rs

@@ -32,15 +32,18 @@ impl Hash for Aes128Z2Hash {
     }
 
     fn hash_all_into(&self, data: &mut [AesBlock], out: &mut [u8]) {
+        debug_assert_eq!(out.len() * 8, data.len());
         self.cipher.encrypt_blocks(data);
-        // Pack the Z2 (1-bit) outputs LSB-first, eight blocks per byte (assign
-        // into a fresh block); see `pack_bits_lsb_first` for the convention.
-        crate::primitives::pack_bits_lsb_first(
-            out,
-            data,
-            |block| block[0],
-            |slot, byte| *slot = byte,
-        );
+
+        // Pack the Z2 (1-bit) outputs LSB-first, eight blocks per byte —
+        // the same bit order as `RightBitVec::set_bit`. The 256-block case
+        // (Bit8's per-block RO output) has a vectorised gather; other sizes
+        // use the scalar pack.
+        if data.len() == 256 {
+            crate::primitives::simd::lsb_mask_256(data, out);
+        } else {
+            crate::primitives::simd::scalar::lsb_mask(data, out);
+        }
     }
 }
 

packages/ore-rs/src/primitives/prp.rs

@@ -82,17 +82,12 @@ impl Prp<u8> for KnuthShufflePRP<u8, 256> {
     }
 
     fn indicator_mask_xor(&self, data: u8, out: &mut [u8]) {
+        debug_assert_eq!(out.len() * 8, 256);
+
         // `invert(j)` is `self.permutation[j]` (see `invert` above), so the
-        // mask is one linear pass over the table: a bytewise `> data` compare,
-        // XORed over the hash bits. Branch-free with a fixed trip count — the
-        // scalar form of a SIMD compare-and-movemask (v2 plan §3). Bit order
-        // is shared with `hash_all_into` via `pack_bits_lsb_first`.
-        crate::primitives::pack_bits_lsb_first(
-            out,
-            &self.permutation,
-            |&p| u8::from(p > data),
-            |slot, byte| *slot ^= byte,
-        );
+        // mask is one pass over the table: a bytewise `> data` compare
+        // packed to bits — vectorised where the target supports it.
+        crate::primitives::simd::gt_mask_xor_256(&self.permutation, data, out);
     }
 }
 

packages/ore-rs/src/primitives/simd.rs

@@ -0,0 +1,254 @@
+//! Vectorised kernels for the data-parallel inner loops of right-ciphertext
+//! encoding, with portable scalar fallbacks (v2 plan §3).
+//!
+//! Two operations are covered:
+//!
+//! - **`gt_mask_xor_256`** — XOR the indicator mask `bit j = (table[j] > x)`
+//!   into a 32-byte bitvector. This is the bulk form of the per-block
+//!   unary encoding (a bytewise compare against a broadcast value, packed
+//!   to a bitmask — `cmgt`+pack on NEON, `vpcmpgtb`+`vpmovmskb` on AVX2).
+//! - **`lsb_mask_256`** — pack `bit j = blocks[j][0] & 1` (a stride-16
+//!   gather over AES outputs) into a 32-byte bitvector.
+//!
+//! Bit order everywhere is LSB-first within each byte, matching
+//! `RightBitVec::set_bit`.
+//!
+//! # Dispatch
+//!
+//! - **aarch64:** NEON is baseline — compile-time dispatch, no detection.
+//! - **x86_64:** AVX2 via `is_x86_feature_detected!` (cached by `std`) for
+//!   the indicator mask; the LSB gather stays scalar (the stride-16 gather
+//!   doesn't pay on AVX2 without AVX-512 VBMI, and the scalar form is
+//!   already branch-free and fast).
+//! - **everything else:** scalar.
+//!
+//! # Constant-time discipline
+//!
+//! Every path in this module — scalar and vector — has a fixed trip count,
+//! no data-dependent branches, and no data-dependent memory addressing
+//! (the NEON `tbl` indices are compile-time constants). Equivalence
+//! between scalar and vector paths is pinned by tests below; equivalence
+//! against the per-bit reference is pinned in `prp.rs`.
+
+use super::AesBlock;
+
+/// XOR the indicator mask for `x` over `table` into `out`:
+/// `out[j/8] ^= ((table[j] > x) as u8) << (j%8)`.
+#[inline]
+pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) {
+    debug_assert_eq!(out.len(), 32);
+
+    #[cfg(target_arch = "aarch64")]
+    // SAFETY: NEON is baseline on aarch64; `out` length asserted above.
+    unsafe {
+        neon::gt_mask_xor_256(table, x, out);
+    }
+
+    #[cfg(target_arch = "x86_64")]
+    if is_x86_feature_detected!("avx2") {
+        // SAFETY: AVX2 presence just checked; `out` length asserted above.
+        unsafe {
+            avx2::gt_mask_xor_256(table, x, out);
+            return;
+        }
+    }
+
+    // Scalar fallback; excluded on aarch64 (the NEON path above always handles
+    // it) so the dead tail needs no blanket `#[allow(unreachable_code)]`.
+    #[cfg(not(target_arch = "aarch64"))]
+    scalar::gt_mask_xor_256(table, x, out);
+}
+
+/// Pack the LSB of byte 0 of each of 256 AES blocks into `out`:
+/// `out[j/8] |= (blocks[j][0] & 1) << (j%8)` over zeroed positions
+/// (callers pass `out` slots they own; bits are assigned, not accumulated).
+#[inline]
+pub(crate) fn lsb_mask_256(blocks: &[AesBlock], out: &mut [u8]) {
+    // Real asserts (not debug_assert): the NEON path gathers `blocks` via raw
+    // pointers assuming exactly 256 blocks, so a shorter slice would read out
+    // of bounds (UB) in a release build without this check.
+    assert_eq!(blocks.len(), 256);
+    assert_eq!(out.len(), 32);
+
+    #[cfg(target_arch = "aarch64")]
+    // SAFETY: NEON is baseline on aarch64; lengths asserted above.
+    unsafe {
+        neon::lsb_mask_256(blocks, out);
+    }
+
+    // Scalar fallback; excluded on aarch64 (the NEON path above always handles
+    // it) so the dead tail needs no blanket `#[allow(unreachable_code)]`.
+    #[cfg(not(target_arch = "aarch64"))]
+    scalar::lsb_mask(blocks, out);
+}
+
+pub(crate) mod scalar {
+    use super::AesBlock;
+
+    // Used by the non-aarch64 `gt_mask_xor_256` dispatcher and as the test
+    // oracle; on aarch64 the lib always reaches NEON, so outside `cfg(test)`
+    // this fn is unreferenced there. Compiling it exactly where it's used keeps
+    // it from being dead code without an `allow`. (`lsb_mask` below stays
+    // always compiled — `hash.rs` calls it directly for non-256 inputs.)
+    #[cfg(any(not(target_arch = "aarch64"), test))]
+    pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) {
+        for (slot, chunk) in out.iter_mut().zip(table.chunks_exact(8)) {
+            let mut byte = 0u8;
+            for (bit, &p) in chunk.iter().enumerate() {
+                byte |= u8::from(p > x) << bit;
+            }
+            *slot ^= byte;
+        }
+    }
+
+    /// Length-generic scalar LSB pack: `blocks.len()` must be a multiple
+    /// of 8 and equal to `out.len() * 8`.
+    pub(crate) fn lsb_mask(blocks: &[AesBlock], out: &mut [u8]) {
+        for (slot, chunk) in out.iter_mut().zip(blocks.chunks_exact(8)) {
+            let mut byte = 0u8;
+            for (bit, block) in chunk.iter().enumerate() {
+                byte |= (block[0] & 1u8) << bit;
+            }
+            *slot = byte;
+        }
+    }
+}
+
+#[cfg(target_arch = "aarch64")]
+mod neon {
+    use super::AesBlock;
+    use core::arch::aarch64::*;
+
+    /// Lane k holds `1 << (k % 8)`: the movemask bit-position vector.
+    const BITPOS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
+
+    /// Pack a 16-lane 0x00/0xFF compare mask into 2 bytes, LSB-first. Each
+    /// lane contributes a distinct power of two within its 8-lane half, so
+    /// the horizontal add is exactly the movemask byte.
+    #[inline(always)]
+    unsafe fn pack_addv(m: uint8x16_t) -> (u8, u8) {
+        let mb = vandq_u8(m, vld1q_u8(BITPOS.as_ptr()));
+        (vaddv_u8(vget_low_u8(mb)), vaddv_u8(vget_high_u8(mb)))
+    }
+
+    #[target_feature(enable = "neon")]
+    pub(super) unsafe fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) {
+        let xv = vdupq_n_u8(x);
+        for i in 0..16 {
+            let t = vld1q_u8(table.as_ptr().add(i * 16));
+            let (lo, hi) = pack_addv(vcgtq_u8(t, xv));
+            out[2 * i] ^= lo;
+            out[2 * i + 1] ^= hi;
+        }
+    }
+
+    /// `IDX[g]` places the 4 block-start bytes of 64-byte group `g`
+    /// (offsets 0, 16, 32, 48) at lanes 4g..4g+3; all other lanes are 0xFF
+    /// (out-of-range, yielding 0 in `tbl`).
+    const TBL_IDX: [[u8; 16]; 4] = {
+        let mut idx = [[0xFFu8; 16]; 4];
+        let mut g = 0;
+        while g < 4 {
+            let mut k = 0;
+            while k < 4 {
+                idx[g][4 * g + k] = (16 * k) as u8;
+                k += 1;
+            }
+            g += 1;
+        }
+        idx
+    };
+
+    /// Gather byte 0 of 16 consecutive 16-byte blocks (256 contiguous
+    /// bytes) into one vector via constant-index 4-register table lookups.
+    #[inline(always)]
+    unsafe fn gather_block0_x16(base: *const u8) -> uint8x16_t {
+        let r0 = vqtbl4q_u8(vld1q_u8_x4(base), vld1q_u8(TBL_IDX[0].as_ptr()));
+        let r1 = vqtbl4q_u8(vld1q_u8_x4(base.add(64)), vld1q_u8(TBL_IDX[1].as_ptr()));
+        let r2 = vqtbl4q_u8(vld1q_u8_x4(base.add(128)), vld1q_u8(TBL_IDX[2].as_ptr()));
+        let r3 = vqtbl4q_u8(vld1q_u8_x4(base.add(192)), vld1q_u8(TBL_IDX[3].as_ptr()));
+        vorrq_u8(vorrq_u8(r0, r1), vorrq_u8(r2, r3))
+    }
+
+    #[target_feature(enable = "neon")]
+    pub(super) unsafe fn lsb_mask_256(blocks: &[AesBlock], out: &mut [u8]) {
+        let p = blocks.as_ptr() as *const u8;
+        let one = vdupq_n_u8(1);
+        for i in 0..16 {
+            let m = vtstq_u8(gather_block0_x16(p.add(i * 256)), one);
+            let (lo, hi) = pack_addv(m);
+            out[2 * i] = lo;
+            out[2 * i + 1] = hi;
+        }
+    }
+}
+
+#[cfg(target_arch = "x86_64")]
+mod avx2 {
+    use core::arch::x86_64::*;
+
+    /// AVX2 has only *signed* byte compares; XOR-ing 0x80 into both sides
+    /// maps unsigned `>` onto signed `>`.
+    #[target_feature(enable = "avx2")]
+    pub(super) unsafe fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) {
+        let bias = _mm256_set1_epi8(i8::MIN);
+        let xv = _mm256_xor_si256(_mm256_set1_epi8(x as i8), bias);
+        for i in 0..8 {
+            let t = _mm256_loadu_si256(table.as_ptr().add(i * 32) as *const __m256i);
+            let gt = _mm256_cmpgt_epi8(_mm256_xor_si256(t, bias), xv);
+            let mask = _mm256_movemask_epi8(gt) as u32;
+            let m = mask.to_le_bytes();
+            out[4 * i] ^= m[0];
+            out[4 * i + 1] ^= m[1];
+            out[4 * i + 2] ^= m[2];
+            out[4 * i + 3] ^= m[3];
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::{Rng, SeedableRng};
+    use rand_chacha::ChaCha20Rng;
+
+    fn rand_table(rng: &mut ChaCha20Rng) -> [u8; 256] {
+        let mut t = [0u8; 256];
+        rng.fill(&mut t[..]);
+        t
+    }
+
+    #[test]
+    fn gt_mask_dispatched_matches_scalar() {
+        let mut rng = ChaCha20Rng::seed_from_u64(7);
+        for _ in 0..5_000 {
+            let table = rand_table(&mut rng);
+            let x: u8 = rng.gen();
+            // Random starting contents so the XOR semantics are exercised.
+            let mut a = [0u8; 32];
+            rng.fill(&mut a[..]);
+            let mut b = a;
+
+            gt_mask_xor_256(&table, x, &mut a);
+            scalar::gt_mask_xor_256(&table, x, &mut b);
+            assert_eq!(a, b);
+        }
+    }
+
+    #[test]
+    fn lsb_mask_dispatched_matches_scalar() {
+        let mut rng = ChaCha20Rng::seed_from_u64(11);
+        for _ in 0..5_000 {
+            let mut blocks = [AesBlock::default(); 256];
+            for block in blocks.iter_mut() {
+                rng.fill(block.as_mut_slice());
+            }
+            let mut a = [0u8; 32];
+            let mut b = [0u8; 32];
+
+            lsb_mask_256(&blocks, &mut a);
+            scalar::lsb_mask(&blocks, &mut b);
+            assert_eq!(a, b);
+        }
+    }
+}

packages/ore-rs/src/scheme/width.rs

@@ -50,16 +50,23 @@ impl AesBlockBuf for [AesBlock; 256] {
 /// Per-block bitvector operations on a Right ciphertext block, one bit per
 /// value in the block domain.
 pub trait RightBitVec {
+    /// Set bit `bit` to `value` (`0` or `1`). (Bulk encoding writes via
+    /// [`Self::as_mut_bytes`]; this remains for width-generic callers.)
+    #[allow(dead_code)]
+    fn set_bit(&mut self, bit: usize, value: u8);
     /// Read bit `bit`. (The width-generic comparator lands with the Bit6
     /// scheme; the legacy comparator calls the inherent method.)
     #[allow(dead_code)]
     fn get_bit(&self, bit: usize) -> u8;
-    /// The raw bitvector bytes, LSB-first within each byte (bit `j` lives in
-    /// byte `j / 8` at position `j % 8`), for bulk mask construction.
+    /// The raw bitvector bytes, LSB-first within each byte (the same bit
+    /// order as [`Self::set_bit`]), for bulk mask construction.
     fn as_mut_bytes(&mut self) -> &mut [u8];
 }
 
 impl RightBitVec for RightBlock32 {
+    fn set_bit(&mut self, bit: usize, value: u8) {
+        RightBlock32::set_bit(self, bit, value)
+    }
     fn get_bit(&self, bit: usize) -> u8 {
         RightBlock32::get_bit(self, bit)
     }