diff --git a/.claude/blackboard.md b/.claude/blackboard.md index f8f21b17..44321e56 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3849,3 +3849,30 @@ clippy errors are all pre-existing (same set on master). Disable run: AVX2 cmpgt Trap recorded: `cargo --config clippy` DROPS the config (clippy is an external subcommand); the v3 lint pass reported clean in 0.05s while linting the native build. Put `--config` after the subcommand: `cargo clippy --config .cargo/config-v3.toml`. `cargo --config test` does apply it. + +## dot4_u8_i8_add + U8x64 nibble split on every backend (2026-10-29, branch claude/dot4-nibble-facade) + +Trigger: turbovec's VNNI "permute-dot" kernels are raw `vpdpbusd` on x86 and `SDOT` on aarch64 (the same +operation written twice), plus a GFNI-only macro variant for the high nibble. Adding the operation to the facade +lets those kernels move onto `ndarray::simd`. + +Added on AVX-512, AVX2, AVX, scalar, NEON, wasm and nightly: `I32x16::dot4_u8_i8_add(a: U8x64, b: I8x64)` +(`vpdpbusd` when avx512vnni is compiled in, otherwise widen-then-`madd`; see the contract doc's P0 section on why +`vpmaddubsw` is wrong) and `U8x64::low_nibbles` / `high_nibbles` (GFNI `gf2p8affine` when compiled in). + +Fixed on the way: on aarch64 and wasm, `U16x32` comes from the scalar file but `F32x16` is native, so the scalar +`to_f32x16_lo/hi` (PR #385) returned a type that is not `ndarray::simd::F32x16`, and the aarch64 lib tests did not +compile at all. They now return `crate::simd::F32x16`. CI never builds aarch64/wasm lib tests, only the two +parity programs, which is why it slipped through; both programs now check dot4 and the nibbles (codes 0x500..). + +Witnessed: `simd::dot4_nibble_facade_parity` 4/4 on native (AVX-512 + VNNI + GFNI), config-v4 (AVX-512 emulation), +config-v3 (AVX2), config-avx (qemu SandyBridge), config-v2 (scalar, qemu Nehalem), aarch64 (qemu, plus +`u16x32_facade_parity`) and nightly-simd; `neon-parity.sh` and `wasm-parity.sh` pass. Lib suite 2657 passed +(native). clippy -D warnings clean on native/v4/v3/v2 (`--config` AFTER the subcommand) and aarch64; the AVX arm's +8 errors are pre-existing, none on added lines. Disable runs, all red as expected: `vpmaddubsw` AVX2 realization +(saturation test + reference), wrong GFNI matrix (nibble test), `vpdpbusds` for `vpdpbusd` (reference, i32::MAX acc). + +Loose ends: AVX-VNNI (`avxvnni`, 256-bit `vpdpbusd`) is not used on the AVX2 arm; NEON could use `USDOT` when i8mm +is compiled in. Neither is measured. The signed form (`i8 x i8`, turbovec's `xor 0x80` + `zero` correction) is not +a facade word yet. + diff --git a/.claude/knowledge/vertical-simd-consumer-contract.md b/.claude/knowledge/vertical-simd-consumer-contract.md index be66b6da..1fc78c5b 100644 --- a/.claude/knowledge/vertical-simd-consumer-contract.md +++ b/.claude/knowledge/vertical-simd-consumer-contract.md @@ -80,6 +80,37 @@ Any saturating-abs primitive in ndarray that does NOT produce `i8::MAX` for `i8: --- +## Four-byte dot product: `vpmaddubsw` is NOT an emulation of `vpdpbusd` (P0, 2026-10-29) + +`I32x16::dot4_u8_i8_add(acc, a: U8x64, b: I8x64)` adds `a[4j..4j+4] · b[4j..4j+4]` +(u8 times signed i8) to lane `j`. AVX-512 VNNI does it in one `vpdpbusd`; every +other realization must give the same answer, and the obvious AVX2 emulation does not: + +- `vpmaddubsw` sums each pair of `u8 * i8` products into an **i16 with signed + saturation**. Two products reach `2 * 255 * 127 = 64_770`, past `i16::MAX`, so the + pair clamps to 32_767 and the dot product is wrong. `vpmaddwd` with ones afterwards + cannot recover it. +- The correct realization widens first: even and odd bytes into separate 16-bit lanes + (`a` zero-extended with `& 0x00FF` / logical `>> 8`, `b` sign-extended with + `<< 8` then arithmetic `>> 8` / arithmetic `>> 8`), then `vpmaddwd` on each, then + add. The even lanes give bytes 4j and 4j+2, the odd lanes 4j+1 and 4j+3; every + product and partial sum is exact in i32. The same shape is used on AVX-512 without + VNNI, AVX2, AVX (SSE2 quarters) and wasm (`i32x4_dot_i16x8`); NEON widens with + `vmovl` and sums with `vmull_s16` + `vpaddq_s32`; nightly takes byte k of each + 32-bit word by shifts. +- Only the final addition to the accumulator wraps, as `vpdpbusd` does. `vpdpbusds` + (the saturating form) is a different operation. + +**Mandatory tests:** `simd::dot4_nibble_facade_parity::dot4_does_not_saturate_pair_sums` +(all-255 times all-127 must give `4 * 255 * 127` per lane, times all -128 `-4 * 255 * 128`) and the +`i32::MAX` accumulator cases in `dot4_matches_reference`. Disable runs on 2026-10-29: +a `vpmaddubsw` AVX2 realization fails both; `vpdpbusds` in place of `vpdpbusd` fails +the second. + +`U8x64::low_nibbles` / `high_nibbles` are the nibble split every 4-bit LUT kernel +starts with. `high_nibbles` is one `gf2p8affine` (matrix `0x1020_4080_0000_0000`) when +GFNI is compiled in, otherwise a 16-bit shift by 4 and a mask. + ## W1a queue — 5 primitives ndarray must ship Each is a tight-scope PR. Recommended: one branch per primitive, parallel review. diff --git a/crates/neon-simd-parity/src/main.rs b/crates/neon-simd-parity/src/main.rs index 9af5eca8..1bf26724 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 — bit-identical to scalar"); + println!("neon-simd-parity OK: U32x16 / F32x16 / I8x16 lanes, ternlog (256 tables, U32x16 + U64x8), U64x8 algebra, I32x16 compares, dot4 + nibbles — bit-identical to scalar"); } #[cfg(not(target_arch = "aarch64"))] { @@ -50,6 +50,9 @@ mod checks { if let Err(code) = check_i32x16_compare() { return code; } + if let Err(code) = check_dot4_and_nibbles() { + return code; + } 0 } @@ -460,4 +463,73 @@ mod checks { } Ok(()) } + + /// `I32x16::dot4_u8_i8_add` and the `U8x64` nibble split against inline + /// scalar arithmetic. Bytes cover `0`/`255` and `-128`/`127` (a pair of + /// `255 * 127` products overflows `i16`, so a saturating emulation fails + /// here), accumulators cover `i32::MIN`/`i32::MAX` (only the final addition + /// may wrap), and a single nonzero byte must move exactly its own lane. + /// Return codes `0x500 + op`. + fn check_dot4_and_nibbles() -> Result<(), u32> { + use ndarray::simd::{I32x16, I8x64, U8x64}; + let mut s = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move || { + s = s + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (s >> 33) as u32 + }; + let edges_u = [0u8, 255, 128, 127]; + let edges_i = [-128i8, 127, -1, 0]; + for round in 0..40 { + let a: [u8; 64] = core::array::from_fn(|i| { + if round < 4 { + edges_u[(i + round) % 4] + } else { + next() as u8 + } + }); + let b: [i8; 64] = core::array::from_fn(|i| { + if round < 4 { + edges_i[(i + round) % 4] + } else { + i8::from_ne_bytes([next() as u8]) + } + }); + let rand: [i32; 16] = core::array::from_fn(|_| next() as i32); + for acc in [[0i32; 16], [i32::MAX; 16], [i32::MIN; 16], rand] { + let got = I32x16::from_array(acc) + .dot4_u8_i8_add(U8x64::from_array(a), I8x64::from_array(b)) + .to_array(); + let want: [i32; 16] = core::array::from_fn(|j| { + let mut sum = 0i32; + for i in 4 * j..4 * j + 4 { + sum += i32::from(a[i]) * i32::from(b[i]); + } + acc[j].wrapping_add(sum) + }); + if got != want { + return Err(0x500); + } + } + let x = U8x64::from_array(a); + if x.low_nibbles().to_array() != a.map(|v| v & 0x0F) { + return Err(0x501); + } + if x.high_nibbles().to_array() != a.map(|v| v >> 4) { + return Err(0x502); + } + } + for i in 0..64 { + let mut a = [0u8; 64]; + a[i] = 3; + let got = I32x16::splat(0) + .dot4_u8_i8_add(U8x64::from_array(a), I8x64::splat(-5)) + .to_array(); + if got != core::array::from_fn(|j| if j == i / 4 { -15 } else { 0 }) { + return Err(0x503); + } + } + Ok(()) + } } diff --git a/crates/wasm-simd-parity/src/lib.rs b/crates/wasm-simd-parity/src/lib.rs index 1dc5cd2a..3fd2d58e 100644 --- a/crates/wasm-simd-parity/src/lib.rs +++ b/crates/wasm-simd-parity/src/lib.rs @@ -35,6 +35,9 @@ pub extern "C" fn selfcheck() -> u32 { if let Err(code) = check_i32x16_compare() { return code; } + if let Err(code) = check_dot4_and_nibbles() { + return code; + } if let Err(code) = check_exp_and_round() { return code; } @@ -526,3 +529,72 @@ fn check_i32x16_compare() -> Result<(), u32> { } Ok(()) } + +/// `I32x16::dot4_u8_i8_add` and the `U8x64` nibble split against inline +/// scalar arithmetic. Bytes cover `0`/`255` and `-128`/`127` (a pair of +/// `255 * 127` products overflows `i16`, so a saturating emulation fails +/// here), accumulators cover `i32::MIN`/`i32::MAX` (only the final addition +/// may wrap), and a single nonzero byte must move exactly its own lane. +/// Return codes `0x500 + op`. +fn check_dot4_and_nibbles() -> Result<(), u32> { + use ndarray::simd::{I32x16, I8x64, U8x64}; + let mut s = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move || { + s = s + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (s >> 33) as u32 + }; + let edges_u = [0u8, 255, 128, 127]; + let edges_i = [-128i8, 127, -1, 0]; + for round in 0..40 { + let a: [u8; 64] = core::array::from_fn(|i| { + if round < 4 { + edges_u[(i + round) % 4] + } else { + next() as u8 + } + }); + let b: [i8; 64] = core::array::from_fn(|i| { + if round < 4 { + edges_i[(i + round) % 4] + } else { + i8::from_ne_bytes([next() as u8]) + } + }); + let rand: [i32; 16] = core::array::from_fn(|_| next() as i32); + for acc in [[0i32; 16], [i32::MAX; 16], [i32::MIN; 16], rand] { + let got = I32x16::from_array(acc) + .dot4_u8_i8_add(U8x64::from_array(a), I8x64::from_array(b)) + .to_array(); + let want: [i32; 16] = core::array::from_fn(|j| { + let mut sum = 0i32; + for i in 4 * j..4 * j + 4 { + sum += i32::from(a[i]) * i32::from(b[i]); + } + acc[j].wrapping_add(sum) + }); + if got != want { + return Err(0x500); + } + } + let x = U8x64::from_array(a); + if x.low_nibbles().to_array() != a.map(|v| v & 0x0F) { + return Err(0x501); + } + if x.high_nibbles().to_array() != a.map(|v| v >> 4) { + return Err(0x502); + } + } + for i in 0..64 { + let mut a = [0u8; 64]; + a[i] = 3; + let got = I32x16::splat(0) + .dot4_u8_i8_add(U8x64::from_array(a), I8x64::splat(-5)) + .to_array(); + if got != core::array::from_fn(|j| if j == i / 4 { -15 } else { 0 }) { + return Err(0x503); + } + } + Ok(()) +} diff --git a/src/simd.rs b/src/simd.rs index 9b9c7023..acafd0e4 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -2313,3 +2313,114 @@ mod u16x32_facade_parity { assert_eq!(U16x32::default().to_array(), [0u16; 32]); } } + +/// Parity of `I32x16::dot4_u8_i8_add` and the `U8x64` nibble split against a +/// plain-array reference, on whichever backend this build selected. The +/// corpus mixes a fixed-seed generator with the byte extremes (`0`, `255`, +/// `-128`, `127`) and accumulators at `i32::MIN` / `i32::MAX`, where only the +/// final addition may wrap. +#[cfg(test)] +mod dot4_nibble_facade_parity { + use super::{I32x16, I8x64, U8x64}; + + fn next(state: &mut u64) -> u64 { + *state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + *state >> 33 + } + + fn bytes() -> Vec<[u8; 64]> { + let mut s = 0x9E37_79B9_7F4A_7C15u64; + let mut out: Vec<[u8; 64]> = (0..48) + .map(|_| core::array::from_fn(|_| next(&mut s) as u8)) + .collect(); + for e in [0u8, 1, 0x0F, 0x10, 0x7F, 0x80, 0xF0, 0xFF] { + out.push([e; 64]); + } + out.push(core::array::from_fn(|i| [0u8, 255, 128, 127][i % 4])); + out + } + + fn signed() -> Vec<[i8; 64]> { + let mut s = 0xD1B5_4A32_D192_ED03u64; + let mut out: Vec<[i8; 64]> = (0..48) + .map(|_| core::array::from_fn(|_| i8::from_ne_bytes([next(&mut s) as u8]))) + .collect(); + for e in [0i8, 1, -1, 127, -128] { + out.push([e; 64]); + } + out.push(core::array::from_fn(|i| [-128i8, 127, -1, 0][i % 4])); + out + } + + fn reference(acc: [i32; 16], a: &[u8; 64], b: &[i8; 64]) -> [i32; 16] { + core::array::from_fn(|j| { + let sum: i32 = (4 * j..4 * j + 4) + .map(|i| i32::from(a[i]) * i32::from(b[i])) + .sum(); + acc[j].wrapping_add(sum) + }) + } + + #[test] + fn dot4_matches_reference() { + let (us, is) = (bytes(), signed()); + let mut s = 0x2545_F491_4F6C_DD1Du64; + let mut checked = 0; + for a in &us { + for b in &is { + let rand: [i32; 16] = core::array::from_fn(|_| next(&mut s) as i32); + for acc in [[0i32; 16], [i32::MAX; 16], [i32::MIN; 16], rand] { + let got = I32x16::from_array(acc) + .dot4_u8_i8_add(U8x64::from_array(*a), I8x64::from_array(*b)) + .to_array(); + assert_eq!(got, reference(acc, a, b), "acc {acc:?}\na {a:?}\nb {b:?}"); + checked += 1; + } + } + } + assert!(checked > 10_000, "corpus too small: {checked}"); + } + + /// A pair of `255 * 127` products is 64_770, which a saturating 16-bit + /// pair sum (`vpmaddubsw`) would clamp to 32_767. These lanes would expose + /// such an emulation in either sign. + #[test] + fn dot4_does_not_saturate_pair_sums() { + let acc = I32x16::splat(0); + let a = U8x64::splat(255); + assert_eq!(acc.dot4_u8_i8_add(a, I8x64::splat(127)).to_array(), [4 * 255 * 127; 16]); + assert_eq!(acc.dot4_u8_i8_add(a, I8x64::splat(-128)).to_array(), [-4 * 255 * 128; 16]); + } + + /// Lane `j` must read bytes `4j..4j+4` and nothing else: a single nonzero + /// byte moves exactly one lane. + #[test] + fn dot4_keeps_bytes_in_their_lane() { + for i in 0..64 { + let mut a = [0u8; 64]; + a[i] = 3; + let got = I32x16::splat(0) + .dot4_u8_i8_add(U8x64::from_array(a), I8x64::splat(-5)) + .to_array(); + let want: [i32; 16] = core::array::from_fn(|j| if j == i / 4 { -15 } else { 0 }); + assert_eq!(got, want, "byte {i}"); + } + } + + #[test] + fn nibbles_match_reference() { + for a in bytes() { + let x = U8x64::from_array(a); + assert_eq!(x.low_nibbles().to_array(), a.map(|v| v & 0x0F), "low_nibbles {a:?}"); + assert_eq!(x.high_nibbles().to_array(), a.map(|v| v >> 4), "high_nibbles {a:?}"); + } + let every: [u8; 64] = core::array::from_fn(|i| (i as u8).wrapping_mul(4).wrapping_add(3)); + for v in [every, every.map(|x| !x)] { + let x = U8x64::from_array(v); + assert_eq!(x.high_nibbles().to_array(), v.map(|b| b >> 4)); + assert_eq!(x.low_nibbles().to_array(), v.map(|b| b & 0x0F)); + } + } +} diff --git a/src/simd_avx.rs b/src/simd_avx.rs index a9013156..5038558f 100644 --- a/src/simd_avx.rs +++ b/src/simd_avx.rs @@ -28,7 +28,7 @@ use core::ops::{ Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Mul, Not, Shl, Shr, Sub, SubAssign, }; -use crate::simd_avx2::{F32x16, I32x16, U16x16, U16x32, U64x8, U8x32}; +use crate::simd_avx2::{F32x16, I32x16, I8x64, U16x16, U16x32, U64x8, U8x32, U8x64}; use crate::simd_avx512::{F32x8, I16x16, I8x32}; // --------------------------------------------------------------------------- @@ -266,10 +266,99 @@ impl Shr for U64x8 { } } +// --------------------------------------------------------------------------- +// U8x64 (array-backed) +// --------------------------------------------------------------------------- + +impl U8x64 { + /// The low four bits of every byte, `x & 0x0F`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).low_nibbles().to_array(), [0x07; 64]); + /// ``` + #[inline(always)] + pub fn low_nibbles(self) -> Self { + let mut out = [0u8; 64]; + // SAFETY: SSE2 is compiled in (module docs); four 16-byte unaligned loads and stores + // inside the 64-byte arrays. + unsafe { + let m = _mm_set1_epi8(0x0F); + let (src, dst) = (self.0.as_ptr() as *const __m128i, out.as_mut_ptr() as *mut __m128i); + for q in 0..4 { + _mm_storeu_si128(dst.add(q), _mm_and_si128(_mm_loadu_si128(src.add(q)), m)); + } + } + Self(out) + } + + /// The high four bits of every byte, moved down: `x >> 4`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).high_nibbles().to_array(), [0x0A; 64]); + /// ``` + #[inline(always)] + pub fn high_nibbles(self) -> Self { + let mut out = [0u8; 64]; + // SAFETY: as in `low_nibbles`. The mask clears the bits the 16-bit shift pulls in. + unsafe { + let m = _mm_set1_epi8(0x0F); + let (src, dst) = (self.0.as_ptr() as *const __m128i, out.as_mut_ptr() as *mut __m128i); + for q in 0..4 { + _mm_storeu_si128(dst.add(q), _mm_and_si128(_mm_srli_epi16(_mm_loadu_si128(src.add(q)), 4), m)); + } + } + Self(out) + } +} + // --------------------------------------------------------------------------- // I32x16 (array-backed) // --------------------------------------------------------------------------- +impl I32x16 { + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: U8x64, b: I8x64) -> Self { + 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. Even and odd bytes are widened into + // separate 16-bit lanes (zero-extended for `a`, sign-extended for `b`), so the 16-bit + // multiply-add sums exact products into i32 lanes: the even lanes give bytes 4j and + // 4j+2, the odd lanes 4j+1 and 4j+3. + unsafe { + let low = _mm_set1_epi16(0x00FF); + let (pa, pb) = (a.0.as_ptr() as *const __m128i, b.0.as_ptr() as *const __m128i); + let po = out.as_mut_ptr() as *mut __m128i; + for q in 0..4 { + let (av, bv) = (_mm_loadu_si128(pa.add(q)), _mm_loadu_si128(pb.add(q))); + let even = _mm_madd_epi16(_mm_and_si128(av, low), _mm_srai_epi16(_mm_slli_epi16(bv, 8), 8)); + let odd = _mm_madd_epi16(_mm_srli_epi16(av, 8), _mm_srai_epi16(bv, 8)); + let acc = _mm_add_epi32(_mm_loadu_si128(po.add(q)), _mm_add_epi32(even, odd)); + _mm_storeu_si128(po.add(q), acc); + } + } + Self(out) + } +} + impl I32x16 { /// `SaturatingRoundingDoublingHighMul` per lane (TFLM `fixedpoint.h`). /// AVX without AVX2 has no 256-bit integer multiply, so this arm goes diff --git a/src/simd_avx2.rs b/src/simd_avx2.rs index 939453f2..5239b943 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -2920,6 +2920,95 @@ impl U16x32 { } } +#[cfg(not(all(target_feature = "avx", not(target_feature = "avx2"))))] +impl U8x64 { + /// The low four bits of every byte, `x & 0x0F`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).low_nibbles().to_array(), [0x07; 64]); + /// ``` + #[inline(always)] + pub fn low_nibbles(self) -> Self { + let mut out = [0u8; 64]; + // SAFETY: AVX2 is a run-time precondition of this arm (see `U64x8::avx2_halves`); two + // 32-byte unaligned loads and stores inside the 64-byte arrays. + unsafe { + let m = _mm256_set1_epi8(0x0F); + let (src, dst) = (self.0.as_ptr() as *const __m256i, out.as_mut_ptr() as *mut __m256i); + for h in 0..2 { + _mm256_storeu_si256(dst.add(h), _mm256_and_si256(_mm256_loadu_si256(src.add(h)), m)); + } + } + Self(out) + } + + /// The high four bits of every byte, moved down: `x >> 4`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).high_nibbles().to_array(), [0x0A; 64]); + /// ``` + #[inline(always)] + pub fn high_nibbles(self) -> Self { + let mut out = [0u8; 64]; + // SAFETY: as in `low_nibbles`. The mask clears the bits the 16-bit shift pulls in. + unsafe { + let m = _mm256_set1_epi8(0x0F); + let (src, dst) = (self.0.as_ptr() as *const __m256i, out.as_mut_ptr() as *mut __m256i); + for h in 0..2 { + let v = _mm256_srli_epi16(_mm256_loadu_si256(src.add(h)), 4); + _mm256_storeu_si256(dst.add(h), _mm256_and_si256(v, m)); + } + } + Self(out) + } +} + +#[cfg(not(all(target_feature = "avx", not(target_feature = "avx2"))))] +impl I32x16 { + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: U8x64, b: I8x64) -> 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. Even and odd bytes are widened into + // separate 16-bit lanes (zero-extended for `a`, sign-extended for `b`), so the 16-bit + // multiply-add sums exact products into i32 lanes: the even lanes give bytes 4j and + // 4j+2, the odd lanes 4j+1 and 4j+3. + unsafe { + let low = _mm256_set1_epi16(0x00FF); + let (pa, pb) = (a.0.as_ptr() as *const __m256i, b.0.as_ptr() as *const __m256i); + let po = out.as_mut_ptr() as *mut __m256i; + for h in 0..2 { + let (av, bv) = (_mm256_loadu_si256(pa.add(h)), _mm256_loadu_si256(pb.add(h))); + let b_even = _mm256_srai_epi16(_mm256_slli_epi16(bv, 8), 8); + let even = _mm256_madd_epi16(_mm256_and_si256(av, low), b_even); + let odd = _mm256_madd_epi16(_mm256_srli_epi16(av, 8), _mm256_srai_epi16(bv, 8)); + let acc = _mm256_add_epi32(_mm256_loadu_si256(po.add(h)), _mm256_add_epi32(even, odd)); + _mm256_storeu_si256(po.add(h), acc); + } + } + Self(out) + } +} + #[cfg(not(all(target_feature = "avx", not(target_feature = "avx2"))))] impl I32x16 { /// Store two 256-bit halves back into a fresh array (see `avx2_halves`). diff --git a/src/simd_avx512.rs b/src/simd_avx512.rs index fba9cde4..b82aca21 100644 --- a/src/simd_avx512.rs +++ b/src/simd_avx512.rs @@ -527,6 +527,38 @@ impl F64Mask8 { pub struct U8x64(pub __m512i); impl U8x64 { + /// The low four bits of every byte, `x & 0x0F`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).low_nibbles().to_array(), [0x07; 64]); + /// ``` + #[inline(always)] + pub fn low_nibbles(self) -> Self { + // SAFETY: avx512f is compiled in on this arm; register-only. + Self(unsafe { _mm512_and_si512(self.0, _mm512_set1_epi8(0x0F)) }) + } + + /// The high four bits of every byte, moved down: `x >> 4`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).high_nibbles().to_array(), [0x0A; 64]); + /// ``` + #[inline(always)] + pub fn high_nibbles(self) -> Self { + // GF(2) affine map whose matrix rows send input bits 4..7 to output bits 0..3 and + // clear the rest: one instruction instead of a 16-bit shift and a mask. + #[cfg(target_feature = "gfni")] + // SAFETY: gfni and avx512f are compiled in on this arm; register-only. + let v = unsafe { _mm512_gf2p8affine_epi64_epi8(self.0, _mm512_set1_epi64(0x1020_4080_0000_0000), 0) }; + #[cfg(not(target_feature = "gfni"))] + // SAFETY: avx512f + avx512bw are compiled in on this arm (the file relies on BW for its + // byte/word ops); register-only. The mask clears the bits the 16-bit shift pulls in. + let v = unsafe { _mm512_and_si512(_mm512_srli_epi16(self.0, 4), _mm512_set1_epi8(0x0F)) }; + Self(v) + } + pub const LANES: usize = 64; #[inline(always)] @@ -891,6 +923,45 @@ impl PartialEq for U8x64 { pub struct I32x16(pub __m512i); impl I32x16 { + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: U8x64, b: I8x64) -> Self { + #[cfg(target_feature = "avx512vnni")] + // SAFETY: avx512vnni is compiled in on this arm; register-only. + let acc = unsafe { _mm512_dpbusd_epi32(self.0, a.0, b.0) }; + #[cfg(not(target_feature = "avx512vnni"))] + // SAFETY: avx512f + avx512bw are compiled in on this arm (the file relies on BW for its + // byte/word ops); register-only. Even and odd bytes are widened into + // separate 16-bit lanes (zero-extended for `a`, sign-extended for `b`), so the 16-bit + // multiply-add sums exact products into i32 lanes: the even lanes give bytes 4j and + // 4j+2, the odd lanes 4j+1 and 4j+3. + let acc = unsafe { + let a_even = _mm512_and_si512(a.0, _mm512_set1_epi16(0x00FF)); + let a_odd = _mm512_srli_epi16(a.0, 8); + let b_even = _mm512_srai_epi16(_mm512_slli_epi16(b.0, 8), 8); + let b_odd = _mm512_srai_epi16(b.0, 8); + let sum = _mm512_add_epi32(_mm512_madd_epi16(a_even, b_even), _mm512_madd_epi16(a_odd, b_odd)); + _mm512_add_epi32(self.0, sum) + }; + Self(acc) + } + pub const LANES: usize = 16; #[inline(always)] diff --git a/src/simd_neon.rs b/src/simd_neon.rs index 62b25a0c..31fb170d 100644 --- a/src/simd_neon.rs +++ b/src/simd_neon.rs @@ -3292,6 +3292,46 @@ fn rdpot_quad(x: int32x4_t, e: int32x4_t) -> int32x4_t { impl I32x16 { pub const LANES: usize = 16; + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: crate::simd::U8x64, b: crate::simd::I8x64) -> Self { + let (a, b) = (a.to_array(), b.to_array()); + let mut out = self.0; + for (q, quad) in out.iter_mut().enumerate() { + // SAFETY: NEON baseline. The two loads read 16 bytes at offset 16*q of 64-byte + // arrays, q < 4; the rest is register-only. Bytes are widened to 16 bits + // (zero-extended for `a`, sign-extended for `b`), multiplied exactly into i32, + // and pairwise-added twice so lane j holds the sum of bytes 4j..4j+4. + unsafe { + let av = vld1q_u8(a.as_ptr().add(16 * q)); + let bv = vld1q_s8(b.as_ptr().add(16 * q)); + let a_lo = vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(av))); + let a_hi = vreinterpretq_s16_u16(vmovl_high_u8(av)); + let (b_lo, b_hi) = (vmovl_s8(vget_low_s8(bv)), vmovl_high_s8(bv)); + let p01 = vpaddq_s32(vmull_s16(vget_low_s16(a_lo), vget_low_s16(b_lo)), vmull_high_s16(a_lo, b_lo)); + let p23 = vpaddq_s32(vmull_s16(vget_low_s16(a_hi), vget_low_s16(b_hi)), vmull_high_s16(a_hi, b_hi)); + quad.0 = vaddq_s32(quad.0, vpaddq_s32(p01, p23)); + } + } + Self(out) + } + /// Saturating rounding doubling high multiply (gemmlowp / TFLM /// `SaturatingRoundingDoublingHighMul`): per lane, `round(a * b / 2^31)` /// with the half-way case rounded toward `+inf`, i.e. the C++ definition diff --git a/src/simd_nightly/i_word_types.rs b/src/simd_nightly/i_word_types.rs index ed5f495c..4122e72f 100644 --- a/src/simd_nightly/i_word_types.rs +++ b/src/simd_nightly/i_word_types.rs @@ -353,6 +353,42 @@ impl core::fmt::Display for I16x32 { pub struct I32x16(pub i32x16); impl I32x16 { + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: super::U8x64, b: super::I8x64) -> Self { + use core::simd::num::SimdUint; + use core::simd::{u32x16, ToBytes}; + // View the 64 bytes of each operand as 16 words, so word j holds bytes 4j..4j+4 in + // some order; the order does not matter because all four products are summed. Byte k + // of a word is taken by shifts: logical (zero-extending) for `a`, arithmetic + // (sign-extending) for `b`. All four products and their sum are exact in i32. + let aw = u32x16::from_ne_bytes(a.0); + let bw = i32x16::from_ne_bytes(b.0.to_ne_bytes()); + let mut acc = self.0; + for k in 0..4u32 { + let ak = ((aw >> u32x16::splat(8 * k)) & u32x16::splat(0xFF)).cast::(); + let bk = (bw << i32x16::splat(24 - 8 * k as i32)) >> i32x16::splat(24); + acc += ak * bk; + } + Self(acc) + } + pub const LANES: usize = 16; /// Per-lane `SaturatingRoundingDoublingHighMul` (gemmlowp / NEON `VQRDMULH`): diff --git a/src/simd_nightly/u8_types.rs b/src/simd_nightly/u8_types.rs index 01d6c5cf..b2bb917f 100644 --- a/src/simd_nightly/u8_types.rs +++ b/src/simd_nightly/u8_types.rs @@ -29,6 +29,28 @@ use core::simd::{u8x32 as core_u8x32, u8x64 as core_u8x64, Simd}; pub struct U8x64(pub core_u8x64); impl U8x64 { + /// The low four bits of every byte, `x & 0x0F`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).low_nibbles().to_array(), [0x07; 64]); + /// ``` + #[inline(always)] + pub fn low_nibbles(self) -> Self { + Self(self.0 & Simd::splat(0x0F)) + } + + /// The high four bits of every byte, moved down: `x >> 4`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).high_nibbles().to_array(), [0x0A; 64]); + /// ``` + #[inline(always)] + pub fn high_nibbles(self) -> Self { + Self(self.0 >> Simd::splat(4)) + } + /// Number of `u8` lanes. pub const LANES: usize = 64; diff --git a/src/simd_scalar.rs b/src/simd_scalar.rs index 09d2733c..1aa1f8b7 100644 --- a/src/simd_scalar.rs +++ b/src/simd_scalar.rs @@ -1095,9 +1095,11 @@ impl U16x32 { /// let a = U16x32::from_array(core::array::from_fn(|i| i as u16 * 1000)); /// assert_eq!(a.to_f32x16_lo().to_array()[15], 15_000.0); /// ``` + // Returns the facade's `F32x16`, not this file's: on aarch64 and wasm this + // `U16x32` is the facade type but `F32x16` is the native NEON/wasm one. #[inline(always)] - pub fn to_f32x16_lo(self) -> F32x16 { - F32x16(core::array::from_fn(|i| self.0[i] as f32)) + pub fn to_f32x16_lo(self) -> crate::simd::F32x16 { + crate::simd::F32x16::from_array(core::array::from_fn(|i| self.0[i] as f32)) } /// Lanes 16..32 converted to `f32`. Every `u16` is exact in `f32`. @@ -1108,13 +1110,43 @@ impl U16x32 { /// assert_eq!(a.to_f32x16_hi().to_array()[0], 16_000.0); /// ``` #[inline(always)] - pub fn to_f32x16_hi(self) -> F32x16 { - F32x16(core::array::from_fn(|i| self.0[16 + i] as f32)) + pub fn to_f32x16_hi(self) -> crate::simd::F32x16 { + crate::simd::F32x16::from_array(core::array::from_fn(|i| self.0[16 + i] as f32)) } } // Extra methods for I32x16 that float types have via the macro impl I32x16 { + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: U8x64, b: I8x64) -> Self { + let (a, b) = (a.to_array(), b.to_array()); + let mut acc = self.to_array(); + for (j, lane) in acc.iter_mut().enumerate() { + let sum: i32 = (4 * j..4 * j + 4) + .map(|i| i32::from(a[i]) * i32::from(b[i])) + .sum(); + *lane = lane.wrapping_add(sum); + } + Self::from_array(acc) + } + /// Per-lane `SaturatingRoundingDoublingHighMul` (gemmlowp / NEON `VQRDMULH`): /// `round(2*a*b / 2^32)`, ties toward +inf, `MIN * MIN` saturating to `MAX`. /// @@ -1742,6 +1774,28 @@ impl MulAssign for U8x64 { // U8x64 extra methods — byte-level operations for palette codec, nibble, byte scan impl U8x64 { + /// The low four bits of every byte, `x & 0x0F`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).low_nibbles().to_array(), [0x07; 64]); + /// ``` + #[inline(always)] + pub fn low_nibbles(self) -> Self { + Self(self.0.map(|x| x & 0x0F)) + } + + /// The high four bits of every byte, moved down: `x >> 4`. + /// + /// ``` + /// use ndarray::simd::U8x64; + /// assert_eq!(U8x64::splat(0xA7).high_nibbles().to_array(), [0x0A; 64]); + /// ``` + #[inline(always)] + pub fn high_nibbles(self) -> Self { + Self(self.0.map(|x| x >> 4)) + } + #[inline(always)] pub fn reduce_min(self) -> u8 { *self.0.iter().min().unwrap_or(&0) diff --git a/src/simd_wasm.rs b/src/simd_wasm.rs index 3bd33f83..0e941dcd 100644 --- a/src/simd_wasm.rs +++ b/src/simd_wasm.rs @@ -2517,6 +2517,43 @@ pub mod wasm32_simd { impl I32x16 { pub const LANES: usize = 16; + /// Adds four-byte dot products to each 32-bit lane: lane `j` gains + /// `a[4j]*b[4j] + a[4j+1]*b[4j+1] + a[4j+2]*b[4j+2] + a[4j+3]*b[4j+3]`, + /// with each `u8` of `a` multiplied by the signed `i8` of `b` at the same + /// position. The four products are exact; only the final addition to the + /// accumulator wraps on `i32` overflow, as with AVX-512 VNNI `vpdpbusd`. + /// + /// No intermediate saturates. Two products can reach `2 * 255 * 127`, + /// which does not fit `i16`, so an emulation through the saturating + /// `vpmaddubsw` would be wrong; every realization widens to 16 bits first. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(127); + /// // 4 * 255 * 127 = 129_540 per lane, past what a 16-bit pair sum can hold. + /// assert_eq!(I32x16::splat(1).dot4_u8_i8_add(a, b).to_array(), [129_541; 16]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add(self, a: crate::simd::U8x64, b: crate::simd::I8x64) -> Self { + let (a, b) = (a.to_array(), b.to_array()); + let mut out = self.0; + for (q, quad) in out.iter_mut().enumerate() { + // SAFETY: simd128 is compiled in on this arm. The two loads read 16 bytes at + // offset 16*q of 64-byte arrays, q < 4; `v128_load` needs no alignment. Even + // and odd bytes are widened into separate 16-bit lanes (zero-extended for + // `a`, sign-extended for `b`), so `i32x4_dot_i16x8` sums exact products: the + // even lanes give bytes 4j and 4j+2, the odd lanes 4j+1 and 4j+3. + let (av, bv) = unsafe { + (v128_load(a.as_ptr().add(16 * q) as *const v128), v128_load(b.as_ptr().add(16 * q) as *const v128)) + }; + let even = i32x4_dot_i16x8(v128_and(av, u16x8_splat(0x00FF)), i16x8_shr(i16x8_shl(bv, 8), 8)); + let odd = i32x4_dot_i16x8(u16x8_shr(av, 8), i16x8_shr(bv, 8)); + quad.0 = i32x4_add(quad.0, i32x4_add(even, odd)); + } + Self(out) + } + /// Per-lane `SaturatingRoundingDoublingHighMul` (gemmlowp / NEON `VQRDMULH`): /// `round(2*a*b / 2^32)`, ties toward +inf, `MIN * MIN` saturating to `MAX`. ///