diff --git a/.claude/blackboard.md b/.claude/blackboard.md index 88a610ab..a9b1b83d 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3949,3 +3949,11 @@ tiers; 6 disable runs red. Review fix (Codex, PR #391): PL's ordinary ripple now marks an all-ones result FULL (as L); new test + per-add FULL invariant in check_all, Y7 red; 8 tests on 5 tiers. Re-measured on 3 tiers: PL ~2.2 ns per FULL column; verdicts unchanged. +## 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 the abacus probe (#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`. 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;