From 6d15c1b7b72fa6a056bf4bf8923197606d2a5db5 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 20:26:38 +0000 Subject: [PATCH 1/2] Add U64x8::sum8_u8_add_masked: masked byte sums per 8-byte lane Lane j gains the sum of the bytes 8j..8j+8 of a U8x64 whose keep bit is set, accumulated per lane with no horizontal reduction (vpsadbw on the kept bytes on x86, pairwise widening adds on NEON and wasm). Realized in every backend: AVX-512, AVX2, AVX, NEON, wasm, scalar and nightly. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- crates/neon-simd-parity/src/main.rs | 36 +++++++++++++++- crates/wasm-simd-parity/src/lib.rs | 34 +++++++++++++++ src/simd.rs | 67 +++++++++++++++++++++++++++++ src/simd_avx.rs | 42 ++++++++++++++++++ src/simd_avx2.rs | 55 +++++++++++++++++++++++ src/simd_avx512.rs | 28 ++++++++++++ src/simd_neon.rs | 42 ++++++++++++++++++ src/simd_nightly/u_word_types.rs | 32 ++++++++++++++ src/simd_scalar.rs | 32 ++++++++++++++ src/simd_wasm.rs | 44 +++++++++++++++++++ 10 files changed, 411 insertions(+), 1 deletion(-) diff --git a/crates/neon-simd-parity/src/main.rs b/crates/neon-simd-parity/src/main.rs index 5d58d2c6..46fa2a53 100644 --- a/crates/neon-simd-parity/src/main.rs +++ b/crates/neon-simd-parity/src/main.rs @@ -18,7 +18,7 @@ fn main() { eprintln!("neon-simd-parity FAILED: lane/op code = {rc}"); std::process::exit(rc as i32); } - println!("neon-simd-parity OK: U32x16 / F32x16 / I8x16 lanes, ternlog (256 tables, U32x16 + U64x8), U64x8 algebra, I32x16 compares, dot4 (+ masked) + nibbles — bit-identical to scalar"); + println!("neon-simd-parity OK: U32x16 / F32x16 / I8x16 lanes, ternlog (256 tables, U32x16 + U64x8), U64x8 algebra, I32x16 compares, dot4 (+ masked), masked sum8 + nibbles — bit-identical to scalar"); } #[cfg(not(target_arch = "aarch64"))] { @@ -563,6 +563,40 @@ mod checks { } } } + // sum8_u8_add_masked: each keep bit gates its own byte into its own + // 8-byte lane, and random bytes / masks / accumulators (u64::MAX + // included, so the wrap is exercised) match the reference. + for i in 0..64 { + let got = ndarray::simd::U64x8::splat(0) + .sum8_u8_add_masked(U8x64::splat(7), !(1u64 << i)) + .to_array(); + if got != core::array::from_fn(|j| if j == i / 8 { 49 } else { 56 }) { + return Err(0x506); + } + } + let mut t = 0x243F_6A88_85A3_08D3u64; + for round in 0..40 { + t ^= t << 13; + t ^= t >> 7; + t ^= t << 17; + let a: [u8; 64] = core::array::from_fn(|i| (t.rotate_left(i as u32) >> 3) as u8 ^ (i as u8).wrapping_mul(37)); + let keep = match round { + 0 => 0, + 1 => u64::MAX, + 2 => 0x8000_0000_0000_0001, + _ => t.rotate_left(round), + }; + let acc: [u64; 8] = core::array::from_fn(|j| if round % 3 == 0 { u64::MAX - j as u64 } else { t.wrapping_mul(j as u64 + 1) }); + let got = ndarray::simd::U64x8::from_array(acc) + .sum8_u8_add_masked(U8x64::from_array(a), keep) + .to_array(); + let want: [u64; 8] = core::array::from_fn(|j| { + (8 * j..8 * j + 8).fold(acc[j], |s, i| if keep >> i & 1 == 1 { s.wrapping_add(u64::from(a[i])) } else { s }) + }); + if got != want { + return Err(0x507); + } + } Ok(()) } } diff --git a/crates/wasm-simd-parity/src/lib.rs b/crates/wasm-simd-parity/src/lib.rs index 4c8690d8..1efa6c79 100644 --- a/crates/wasm-simd-parity/src/lib.rs +++ b/crates/wasm-simd-parity/src/lib.rs @@ -629,5 +629,39 @@ fn check_dot4_and_nibbles() -> Result<(), u32> { } } } + // sum8_u8_add_masked: each keep bit gates its own byte into its own + // 8-byte lane, and random bytes / masks / accumulators (u64::MAX + // included, so the wrap is exercised) match the reference. + for i in 0..64 { + let got = ndarray::simd::U64x8::splat(0) + .sum8_u8_add_masked(U8x64::splat(7), !(1u64 << i)) + .to_array(); + if got != core::array::from_fn(|j| if j == i / 8 { 49 } else { 56 }) { + return Err(0x506); + } + } + let mut t = 0x243F_6A88_85A3_08D3u64; + for round in 0..40 { + t ^= t << 13; + t ^= t >> 7; + t ^= t << 17; + let a: [u8; 64] = core::array::from_fn(|i| (t.rotate_left(i as u32) >> 3) as u8 ^ (i as u8).wrapping_mul(37)); + let keep = match round { + 0 => 0, + 1 => u64::MAX, + 2 => 0x8000_0000_0000_0001, + _ => t.rotate_left(round), + }; + let acc: [u64; 8] = core::array::from_fn(|j| if round % 3 == 0 { u64::MAX - j as u64 } else { t.wrapping_mul(j as u64 + 1) }); + let got = ndarray::simd::U64x8::from_array(acc) + .sum8_u8_add_masked(U8x64::from_array(a), keep) + .to_array(); + let want: [u64; 8] = core::array::from_fn(|j| { + (8 * j..8 * j + 8).fold(acc[j], |s, i| if keep >> i & 1 == 1 { s.wrapping_add(u64::from(a[i])) } else { s }) + }); + if got != want { + return Err(0x507); + } + } Ok(()) } diff --git a/src/simd.rs b/src/simd.rs index a0c21619..6b6d018b 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -2485,6 +2485,73 @@ mod dot4_nibble_facade_parity { } } + fn sum8_reference(acc: [u64; 8], a: &[u8; 64], keep: u64) -> [u64; 8] { + core::array::from_fn(|j| { + (8 * j..8 * j + 8).fold(acc[j], |s, i| { + if keep >> i & 1 == 1 { + s.wrapping_add(u64::from(a[i])) + } else { + s + } + }) + }) + } + + #[test] + fn sum8_masked_matches_reference() { + use super::U64x8; + let mut s = 0xBB67_AE85_84CA_A73Bu64; + let mut checked = 0; + for a in &bytes() { + for _ in 0..8 { + let rand = next(&mut s) << 32 | next(&mut s) << 1 | (next(&mut s) & 1); + for keep in [0, u64::MAX, 0xAAAA_AAAA_AAAA_AAAA, 0x8000_0000_0000_0001, rand] { + let rand_acc: [u64; 8] = core::array::from_fn(|_| next(&mut s) << 31 ^ next(&mut s)); + for acc in [[0u64; 8], [u64::MAX; 8], rand_acc] { + let got = U64x8::from_array(acc) + .sum8_u8_add_masked(U8x64::from_array(*a), keep) + .to_array(); + assert_eq!(got, sum8_reference(acc, a, keep), "keep {keep:#x} acc {acc:?}\na {a:?}"); + checked += 1; + } + } + } + } + assert!(checked > 5_000, "corpus too small: {checked}"); + } + + /// Bit `i` of `keep` must gate byte `i` into lane `i / 8` and nothing else. + #[test] + fn sum8_masked_gates_each_byte_by_its_own_bit() { + use super::U64x8; + let a = U8x64::splat(7); + for i in 0..64 { + let one = U64x8::splat(0).sum8_u8_add_masked(a, 1 << i).to_array(); + let want: [u64; 8] = core::array::from_fn(|j| if j == i / 8 { 7 } else { 0 }); + assert_eq!(one, want, "bit {i}"); + let all_but = U64x8::splat(0).sum8_u8_add_masked(a, !(1 << i)).to_array(); + let want: [u64; 8] = core::array::from_fn(|j| if j == i / 8 { 49 } else { 56 }); + assert_eq!(all_but, want, "all but bit {i}"); + } + } + + /// With every byte kept, the eight lanes add up to `U8x64::sum_bytes_u64`, + /// and repeated calls accumulate per lane without a reduction. + #[test] + fn sum8_masked_all_kept_agrees_with_sum_bytes() { + use super::U64x8; + for a in &bytes() { + let v = U8x64::from_array(*a); + let lanes = U64x8::splat(0).sum8_u8_add_masked(v, u64::MAX).to_array(); + assert_eq!(lanes.iter().sum::(), v.sum_bytes_u64()); + } + let (a, mut acc) = (U8x64::splat(255), U64x8::splat(0)); + for _ in 0..1000 { + acc = acc.sum8_u8_add_masked(a, u64::MAX); + } + assert_eq!(acc.to_array(), [1000 * 8 * 255; 8]); + } + /// The doc promise: 16_448 worst-case calls fit an `i32` accumulator. #[test] fn dot4_masked_worst_case_holds_16448_calls() { diff --git a/src/simd_avx.rs b/src/simd_avx.rs index 7511e6b7..c4585f99 100644 --- a/src/simd_avx.rs +++ b/src/simd_avx.rs @@ -166,6 +166,48 @@ pub fn dot_i8(a: &[u8], b: &[u8]) -> i64 { // U64x8 (array-backed; only these methods need AVX2 on the AVX2 arm) // --------------------------------------------------------------------------- +impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: U8x64, keep: u64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let mut out = self.0; + // SAFETY: SSE2 is compiled in (module docs); every load and store is a 16-byte + // unaligned access inside a 64-byte array. Byte `i` of `spread` repeats mask byte + // `i / 8`, so testing bit `i % 8` against `BIT_IN_BYTE` gives 0xFF exactly where bit + // `i` of `keep` is set. Quarter `q` covers bytes `16q..16q+16`, i.e. lanes `2q, 2q+1`, + // which is what `sad_epu8` produces. + unsafe { + let bit = _mm_set1_epi64x(crate::simd_masking_ops::BIT_IN_BYTE as i64); + let zero = _mm_setzero_si128(); + let (ps, pa) = (spread.as_ptr() as *const __m128i, a.0.as_ptr() as *const __m128i); + let po = out.as_mut_ptr() as *mut __m128i; + for q in 0..4 { + let m = _mm_cmpeq_epi8(_mm_and_si128(_mm_loadu_si128(ps.add(q)), bit), bit); + let kept = _mm_and_si128(_mm_loadu_si128(pa.add(q)), m); + _mm_storeu_si128(po.add(q), _mm_add_epi64(_mm_loadu_si128(po.add(q)), _mm_sad_epu8(kept, zero))); + } + } + Self(out) + } +} + impl U64x8 { /// `(x << n) | (x >> (64 - n))` per 64-bit lane of one 256-bit half, with /// `1 <= n <= 63` guaranteed by the callers (see the AVX2 realization). diff --git a/src/simd_avx2.rs b/src/simd_avx2.rs index f6f8acb1..52724f33 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -2515,6 +2515,61 @@ impl U64x4 { } } +#[cfg(not(all(target_feature = "avx", not(target_feature = "avx2"))))] +impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: U8x64, keep: u64) -> Self { + let mut out = self.0; + // SAFETY: AVX2 is a run-time precondition of this arm (see `U64x8::avx2_halves`); every + // load and store is a 32-byte unaligned access inside a 64-byte array; the mask is + // built in registers as in `I32x16::dot4_u8_i8_add_masked`: qword `k` of half `h` + // takes mask byte `4h + k` (variable shift), copies it into its eight bytes + // (`* 0x0101_0101`, then shift by 32), and testing bit `i % 8` against `BIT_IN_BYTE` + // gives 0xFF exactly where bit `i` of `keep` is set. Half `h` covers bytes + // `32h..32h+32`, i.e. lanes `4h..4h+4`, which is what `sad_epu8` produces. + unsafe { + let bit = _mm256_set1_epi64x(crate::simd_masking_ops::BIT_IN_BYTE as i64); + let keep_v = _mm256_set1_epi64x(keep as i64); + let byte = _mm256_set1_epi64x(0xFF); + let rep = _mm256_set1_epi64x(0x0101_0101); + let zero = _mm256_setzero_si256(); + let pa = a.0.as_ptr() as *const __m256i; + let po = out.as_mut_ptr() as *mut __m256i; + for h in 0..2 { + let shifts = if h == 0 { + _mm256_setr_epi64x(0, 8, 16, 24) + } else { + _mm256_setr_epi64x(32, 40, 48, 56) + }; + let mb = _mm256_mul_epu32(_mm256_and_si256(_mm256_srlv_epi64(keep_v, shifts), byte), rep); + let spread = _mm256_or_si256(mb, _mm256_slli_epi64(mb, 32)); + let m = _mm256_cmpeq_epi8(_mm256_and_si256(spread, bit), bit); + let kept = _mm256_and_si256(_mm256_loadu_si256(pa.add(h)), m); + let acc = _mm256_add_epi64(_mm256_loadu_si256(po.add(h)), _mm256_sad_epu8(kept, zero)); + _mm256_storeu_si256(po.add(h), acc); + } + } + Self(out) + } +} + // ── W1a-#5: U64x8::popcnt / xor_popcount (AVX2 scalar polyfill) ───────────── // The avx2_int_type! macro generated U64x8 as a scalar polyfill in this file. // We add popcnt + xor_popcount to match the API surface of the AVX-512 backend. diff --git a/src/simd_avx512.rs b/src/simd_avx512.rs index 6536849b..226ecde7 100644 --- a/src/simd_avx512.rs +++ b/src/simd_avx512.rs @@ -2164,6 +2164,34 @@ impl PartialEq for U32x16 { #[repr(transparent)] pub struct U64x8(pub __m512i); +impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: U8x64, keep: u64) -> Self { + // SAFETY: avx512bw is compiled in on this arm; register-only. + Self(unsafe { + let kept = _mm512_maskz_mov_epi8(keep, a.0); + _mm512_add_epi64(self.0, _mm512_sad_epu8(kept, _mm512_setzero_si512())) + }) + } +} + impl U64x8 { pub const LANES: usize = 8; diff --git a/src/simd_neon.rs b/src/simd_neon.rs index b2144a4d..f89b58fa 100644 --- a/src/simd_neon.rs +++ b/src/simd_neon.rs @@ -2874,6 +2874,48 @@ impl Default for U64x8 { } } +#[cfg(target_arch = "aarch64")] +impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: crate::simd::U8x64, keep: u64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let a = a.to_array(); + let mut out = self.0; + for (q, lane) in out.iter_mut().enumerate() { + // SAFETY: NEON baseline. Each load touches 16 bytes at offset 16*q of a 64-byte + // array (`spread` is `[u64; 8]`, 64 bytes), q < 4. Byte `i` of `spread` repeats + // mask byte `i / 8`; `vtst` against `BIT_IN_BYTE` gives 0xFF exactly where bit `i` + // of `keep` is set. The three pairwise widening adds sum bytes 0..8 and 8..16 of + // the quarter into the two `u64` lanes `2q`, `2q + 1`. + unsafe { + let bit = vreinterpretq_u8_u64(vdupq_n_u64(crate::simd_masking_ops::BIT_IN_BYTE)); + let m = vtstq_u8(vld1q_u8((spread.as_ptr() as *const u8).add(16 * q)), bit); + let kept = vandq_u8(vld1q_u8(a.as_ptr().add(16 * q)), m); + let sums = vpaddlq_u32(vpaddlq_u16(vpaddlq_u8(kept))); + *lane = U64x2(vaddq_u64(lane.0, sums)); + } + } + Self(out) + } +} + #[cfg(target_arch = "aarch64")] impl U64x8 { pub const LANES: usize = 8; diff --git a/src/simd_nightly/u_word_types.rs b/src/simd_nightly/u_word_types.rs index f58f40f9..f03d26e7 100644 --- a/src/simd_nightly/u_word_types.rs +++ b/src/simd_nightly/u_word_types.rs @@ -17,6 +17,38 @@ use core::simd::{u16x16, u16x32, u32x16, u32x8, u64x4, u64x8}; #[repr(transparent)] pub struct U64x8(pub u64x8); +impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: super::U8x64, keep: u64) -> Self { + use core::simd::{Mask, Select, Simd, ToBytes}; + let kept = Mask::::from_bitmask(keep).select(a.0, Simd::splat(0)); + // Each u64 lane holds its eight kept bytes; fold them pairwise in place. + let x = u64x8::from_ne_bytes(kept); + let (m8, m16) = (Simd::splat(0x00FF_00FF_00FF_00FFu64), Simd::splat(0x0000_FFFF_0000_FFFFu64)); + let s = (x & m8) + ((x >> 8) & m8); + let s = (s & m16) + ((s >> 16) & m16); + let s = (s & Simd::splat(0xFFFF_FFFF)) + (s >> 32); + Self(self.0 + s) + } +} + impl U64x8 { pub const LANES: usize = 8; diff --git a/src/simd_scalar.rs b/src/simd_scalar.rs index d90fef4b..4d0ca167 100644 --- a/src/simd_scalar.rs +++ b/src/simd_scalar.rs @@ -528,6 +528,38 @@ impl_int_type!(U16x32, u16, 32, 0u16); impl_int_type!(U32x16, u32, 16, 0u32); impl_int_type!(U64x8, u64, 8, 0u64); +impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: U8x64, keep: u64) -> Self { + let (a, mut o) = (a.to_array(), self.to_array()); + for (j, lane) in o.iter_mut().enumerate() { + for (i, &x) in a.iter().enumerate().skip(8 * j).take(8) { + if keep >> i & 1 == 1 { + *lane = lane.wrapping_add(u64::from(x)); + } + } + } + Self::from_array(o) + } +} + /// u64 ARX rotate — the BLAKE2b / argon2 lane. /// /// Scalar per-lane loops, and **measured not to vectorize**: the codegen diff --git a/src/simd_wasm.rs b/src/simd_wasm.rs index 4317b7ff..cd345a19 100644 --- a/src/simd_wasm.rs +++ b/src/simd_wasm.rs @@ -2043,6 +2043,50 @@ pub mod wasm32_simd { } } + impl U64x8 { + /// Lane `j` gains the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit + /// in `keep` is set: `vpsadbw` against zero on the kept bytes, accumulated + /// per lane with no horizontal reduction. Exact; each call adds at most + /// `8 * 255 = 2_040` to a lane, so the wrapping `u64` lanes hold more than + /// 9 * 10^15 calls. Reduce once at the end (`to_array`, `reduce_sum`). + /// + /// A masked byte sum is what a power-of-two anchor needs: with + /// `|b| = low7 + 128 * bit7`, `a * |b| = a * low7 + 128 * (sum of a where bit 7 + /// of b is set)`, so the anchor term is this word with `keep = sign & bit7`. + /// + /// ``` + /// use ndarray::simd::{U64x8, U8x64}; + /// let a = U8x64::splat(255); + /// // Lane 0 keeps bytes 0 and 7, lane 1 keeps byte 8, the rest keep nothing. + /// let got = U64x8::splat(0).sum8_u8_add_masked(a, 0b1_1000_0001).to_array(); + /// assert_eq!(got, [510, 255, 0, 0, 0, 0, 0, 0]); + /// ``` + #[inline(always)] + pub fn sum8_u8_add_masked(self, a: crate::simd::U8x64, keep: u64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let a = a.to_array(); + let bit = u64x2_splat(crate::simd_masking_ops::BIT_IN_BYTE); + let low32 = u64x2_splat(0xFFFF_FFFF); + let mut out = self.0; + for (q, lane) in out.iter_mut().enumerate() { + // SAFETY: simd128 is compiled in on this arm. Each load touches 16 bytes at + // offset 16*q of a 64-byte array, q < 4; `v128_load` needs no alignment. Byte `i` + // of `spread` repeats mask byte `i / 8`, so testing bit `i % 8` against + // `BIT_IN_BYTE` gives 0xFF exactly where bit `i` of `keep` is set. + let kept = unsafe { + let m = i8x16_eq(v128_and(v128_load((spread.as_ptr() as *const v128).add(q)), bit), bit); + v128_and(v128_load((a.as_ptr() as *const v128).add(q)), m) + }; + // Bytes to u16 pairs to u32 quads; each u64 lane then holds two u32 sums, + // added by masking and shifting (no shuffle). + let s32 = u32x4_extadd_pairwise_u16x8(u16x8_extadd_pairwise_u8x16(kept)); + let s64 = i64x2_add(v128_and(s32, low32), u64x2_shr(s32, 32)); + *lane = U64x2(i64x2_add(lane.0, s64)); + } + Self(out) + } + } + impl U64x8 { pub const LANES: usize = 8; From 3dfa3b09eda5e29c0d94e5c9ae20bbcd490da5e9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 21:01:33 +0000 Subject: [PATCH 2/2] Document U64x8::sum8_u8_add_masked Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- .claude/blackboard.md | 9 +++++++++ .claude/knowledge/turbovec9-u9-substrate.md | 20 ++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/.claude/blackboard.md b/.claude/blackboard.md index 4c67fc2f..afd83ef0 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3898,3 +3898,12 @@ negative zero; both fixed, every exact kernel now checked against an i64 referen Loose ends: fused single-widen SM9 op for non-VNNI tiers; NEON SDOT/USDOT and AVX-VNNI unused; no NEON/wasm timings; recall on a real TurboVec index is the TurboVec owner's measurement. + +## U64x8::sum8_u8_add_masked: vertical masked byte sum (2026-10-29, branch claude/sum8-u8-masked) + +Facade word closing the gap #389 recorded for candidate E's anchor term: lane j += sum of kept bytes 8j..8j+8, no +horizontal reduction. One realization per backend (AVX-512 maskz + vpsadbw; AVX2 register mask + vpsadbw; AVX psadbw per +quarter; NEON uaddlp/uadalp; wasm extadd_pairwise; scalar; nightly SWAR). Parity in simd.rs + wasm/neon harnesses +(0x506/0x507); 7 tiers green; 7 disable runs red. 0.687 ns/64 bytes on AVX-512 (281-285x the reference loop), 1.99 ns on +AVX2 (98x). Knowledge: `turbovec9-u9-substrate.md` § second facade word. Not yet wired into #389. + diff --git a/.claude/knowledge/turbovec9-u9-substrate.md b/.claude/knowledge/turbovec9-u9-substrate.md index b1a0a6b1..f9fbbfa2 100644 --- a/.claude/knowledge/turbovec9-u9-substrate.md +++ b/.claude/knowledge/turbovec9-u9-substrate.md @@ -157,3 +157,23 @@ from this session. - No fused single-widen SM9 op for non-VNNI tiers. - No signed i8 x i8 facade word. - Recall on a real index. + +## Second facade word: `U64x8::sum8_u8_add_masked` (2026-10-29) + +`U64x8::sum8_u8_add_masked(self, a: U8x64, keep: u64) -> U64x8` adds to lane `j` the sum of the bytes `a[i]`, `i` in `8j..8j+8`, whose bit `i` in `keep` is set. It is exact. Each call adds at most 8 * 255 = 2,040 per lane, so the u64 lanes wrap only after more than 9 * 10^15 calls. There is no horizontal reduction; reduce once at the end. + +It was added for the power-of-two anchor term of candidate E in PR #389. With `|b| = low7 + 128*bit7`, `a*|b| = a*low7 + 128*(sum of a where bit 7 of b is set)`. The facade only had `U8x64::sum_bytes_u64`, which reduces horizontally on every call. + +| Tier | Lowering | +|---|---| +| AVX-512 | zero-masked `vmovdqu8`, `vpsadbw` against zero, `vpaddq` | +| AVX2 | mask built in registers as in `dot4_u8_i8_add_masked`, AND, `vpsadbw`, `vpaddq`, two halves (LLVM emits 2 `vpshufb`; none in the source) | +| AVX without AVX2 | `mask_byte_spread`, `cmpeq`, AND, `psadbw` per 16-byte quarter | +| NEON | `vtstq_u8` mask, AND, `uaddlp` x2, `uadalp`; no scalar loop | +| wasm simd128 | `i8x16_eq` mask, AND, two `extadd_pairwise`, fold of 32-bit halves, `i64x2_add`; no shuffle | +| scalar | per-lane loop, the correctness anchor | +| nightly-simd | `Mask::from_bitmask(keep).select`, SWAR fold | + +Measured on an AVX-512 VNNI Xeon (4,096 random vectors and masks, medians of 1,001 runs): 0.687 ns per 64 bytes at target-cpu=native against 192.8 - 196.0 ns for the plain reference loop (281 - 285x); 1.987 - 1.991 ns at x86-64-v3 against 194.1 - 195.4 ns (98x). NEON and wasm were not timed. + +Tests: `sum8_masked_matches_reference`, `sum8_masked_gates_each_byte_by_its_own_bit`, `sum8_masked_all_kept_agrees_with_sum_bytes` in `dot4_nibble_facade_parity`, plus the doc example (run as a doctest on the native and v3 builds) and parity codes 0x506 and 0x507 in `wasm-simd-parity` and `neon-simd-parity`.