From ecce2c17e5d32dfbe39460b9e0a7bf7a609ea12c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 15:30:36 +0000 Subject: [PATCH 1/6] Add I32x16::dot4_u8_i8_add_masked and the 9-bit cost probe dot4_u8_i8_add_masked is dot4_u8_i8_add with the bytes of `a` whose bit in `keep` is clear replaced by 0. Two calls over the same magnitudes, one with the positive plane and one with its complement, give an exact 9-bit sign-magnitude dot product against i8 weights (the full i8 range, -128 included) with i32 accumulators and no widening of the operands. AVX-512 uses one masked move before vpdpbusd; AVX2, AVX, NEON and wasm expand the mask with one AND and one byte compare per register; nightly selects through core::simd::Mask. examples/u9_cost_probe.rs times the candidates per 64 lanes and measures the reconstruction error of each 9-bit layout on a unit Gaussian. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- .claude/hooks/facet-stride-exclude.txt | 1 + Cargo.toml | 4 + crates/neon-simd-parity/src/main.rs | 35 +- crates/wasm-simd-parity/src/lib.rs | 33 ++ examples/u9_cost_probe.rs | 521 +++++++++++++++++++++++++ src/simd.rs | 89 +++++ src/simd_avx.rs | 40 ++ src/simd_avx2.rs | 40 ++ src/simd_avx512.rs | 27 ++ src/simd_masking_ops.rs | 15 + src/simd_neon.rs | 40 ++ src/simd_nightly/i_word_types.rs | 28 ++ src/simd_scalar.rs | 27 ++ src/simd_wasm.rs | 40 ++ 14 files changed, 939 insertions(+), 1 deletion(-) create mode 100644 examples/u9_cost_probe.rs diff --git a/.claude/hooks/facet-stride-exclude.txt b/.claude/hooks/facet-stride-exclude.txt index 0c28dfd4..1d6dd571 100644 --- a/.claude/hooks/facet-stride-exclude.txt +++ b/.claude/hooks/facet-stride-exclude.txt @@ -4,3 +4,4 @@ src/hpc/perm.rs | R1 | Perm64 is a 64-lane permutation map; composing and materializing it is a genuine byte move, the operation itself (module doc: "composable maps"). examples/ternlog_popcnt_gap_probe.rs | R1 | Benchmark that measures the nibble-LUT popcount as its baseline against ternlog + popcnt; the shuffle is the thing being measured. vendor/chacha20/* | R1 | ChaCha20 rotate-left by 8 and 16 bits is a fixed byte rotation (pshufb / vqtbl1q with a constant mask): the cipher's defined operation, in vendored upstream code. +examples/u9_cost_probe.rs | R1 | Cost probe for 9-bit TurboVec scoring: the vpshufb / vpermb / vpermi2b table lookups are the candidates being timed against the lookup-free layout. diff --git a/Cargo.toml b/Cargo.toml index 3e33f058..31bc45af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,10 @@ required-features = ["std"] name = "sm9_presence_probe" required-features = ["std"] +[[example]] +name = "u9_cost_probe" +required-features = ["std"] + [[example]] name = "cpu_guard_probe" required-features = ["std"] diff --git a/crates/neon-simd-parity/src/main.rs b/crates/neon-simd-parity/src/main.rs index 1bf26724..5d58d2c6 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 + 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) + nibbles — bit-identical to scalar"); } #[cfg(not(target_arch = "aarch64"))] { @@ -530,6 +530,39 @@ mod checks { return Err(0x503); } } + // dot4_u8_i8_add_masked: each keep bit gates its own byte, and the + // pos - neg form is exact for every SM9 magnitude in both polarities. + for i in 0..64 { + let got = I32x16::splat(0) + .dot4_u8_i8_add_masked(U8x64::splat(7), !(1u64 << i), I8x64::splat(-3)) + .to_array(); + if got != core::array::from_fn(|j| if j == i / 4 { -63 } else { -84 }) { + return Err(0x504); + } + } + for w in [-128i8, 127, -1] { + for chunk in 0..4usize { + let mags: [u8; 64] = core::array::from_fn(|i| (64 * chunk + i) as u8); + let plane = 0x5A5A_C3C3_0F0F_9696u64.rotate_left(chunk as u32); + let (m, wv) = (U8x64::from_array(mags), I8x64::splat(w)); + let pos = I32x16::splat(0) + .dot4_u8_i8_add_masked(m, plane, wv) + .to_array(); + let neg = I32x16::splat(0) + .dot4_u8_i8_add_masked(m, !plane, wv) + .to_array(); + for j in 0..16 { + let mut want = 0i32; + for i in 4 * j..4 * j + 4 { + let v = i32::from(mags[i]); + want += if plane >> i & 1 == 1 { v } else { -v } * i32::from(w); + } + if pos[j] - neg[j] != want { + return Err(0x505); + } + } + } + } Ok(()) } } diff --git a/crates/wasm-simd-parity/src/lib.rs b/crates/wasm-simd-parity/src/lib.rs index 3fd2d58e..4c8690d8 100644 --- a/crates/wasm-simd-parity/src/lib.rs +++ b/crates/wasm-simd-parity/src/lib.rs @@ -596,5 +596,38 @@ fn check_dot4_and_nibbles() -> Result<(), u32> { return Err(0x503); } } + // dot4_u8_i8_add_masked: each keep bit gates its own byte, and the + // pos - neg form is exact for every SM9 magnitude in both polarities. + for i in 0..64 { + let got = I32x16::splat(0) + .dot4_u8_i8_add_masked(U8x64::splat(7), !(1u64 << i), I8x64::splat(-3)) + .to_array(); + if got != core::array::from_fn(|j| if j == i / 4 { -63 } else { -84 }) { + return Err(0x504); + } + } + for w in [-128i8, 127, -1] { + for chunk in 0..4usize { + let mags: [u8; 64] = core::array::from_fn(|i| (64 * chunk + i) as u8); + let plane = 0x5A5A_C3C3_0F0F_9696u64.rotate_left(chunk as u32); + let (m, wv) = (U8x64::from_array(mags), I8x64::splat(w)); + let pos = I32x16::splat(0) + .dot4_u8_i8_add_masked(m, plane, wv) + .to_array(); + let neg = I32x16::splat(0) + .dot4_u8_i8_add_masked(m, !plane, wv) + .to_array(); + for j in 0..16 { + let mut want = 0i32; + for i in 4 * j..4 * j + 4 { + let v = i32::from(mags[i]); + want += if plane >> i & 1 == 1 { v } else { -v } * i32::from(w); + } + if pos[j] - neg[j] != want { + return Err(0x505); + } + } + } + } Ok(()) } diff --git a/examples/u9_cost_probe.rs b/examples/u9_cost_probe.rs new file mode 100644 index 00000000..97a651fe --- /dev/null +++ b/examples/u9_cost_probe.rs @@ -0,0 +1,521 @@ +//! Cost probe for 9-bit sign-magnitude (SM9) arithmetic and 512-entry lookup. +//! +//! Measurement only: the raw intrinsics here are candidates being timed, not +//! a realization. Compares, per 64 lanes, a plain u8 x i8 four-byte dot +//! (`vpdpbusd`), the signed i8 x i8 dot, three exact ways to multiply a +//! sign-magnitude value in -255..=255 by an i8 weight, and the ways to look a +//! 9-bit centroid ID up in a 512-entry table. +//! +//! Build with the host's features (`.cargo/config.toml` is `target-cpu=native`): +//! `cargo run --release --example u9_cost_probe --features std`. + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512vnni", target_feature = "avx512vbmi"))] +mod probe { + use core::arch::x86_64::*; + use std::hint::black_box; + use std::time::Instant; + + /// Chunks of 64 lanes per pass: magnitudes + weights stay in L1. + const CHUNKS: usize = 256; + const PASSES: usize = 4000; + + struct Data { + mag: Vec<[u8; 64]>, + pos: Vec, + w: Vec<[i8; 64]>, + ids_lo: Vec<[u8; 64]>, + ids_hi: Vec, + } + + fn data() -> Data { + 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 >> 11 + }; + Data { + mag: (0..CHUNKS) + .map(|_| core::array::from_fn(|_| next() as u8)) + .collect(), + pos: (0..CHUNKS).map(|_| next() ^ (next() << 32)).collect(), + w: (0..CHUNKS) + .map(|_| core::array::from_fn(|_| i8::from_ne_bytes([next() as u8]))) + .collect(), + ids_lo: (0..CHUNKS) + .map(|_| core::array::from_fn(|_| next() as u8)) + .collect(), + ids_hi: (0..CHUNKS).map(|_| next() ^ (next() << 32)).collect(), + } + } + + /// Runs `body` over every chunk PASSES times; returns (ns, TSC ticks) per 64 lanes. + fn time(name: &str, mut body: impl FnMut() -> i64) -> (f64, f64) { + black_box(body()); + let t0 = Instant::now(); + // SAFETY: rdtsc has no preconditions. + let c0 = unsafe { _rdtsc() }; + let mut sink = 0i64; + for _ in 0..PASSES { + sink = sink.wrapping_add(body()); + } + // SAFETY: as above. + let c1 = unsafe { _rdtsc() }; + let ns = t0.elapsed().as_nanos() as f64; + black_box(sink); + let per = (PASSES * CHUNKS) as f64; + let (n, c) = (ns / per, (c1 - c0) as f64 / per); + println!("{name:<44} {n:7.3} ns/64 lanes {c:6.2} TSC/64 {:6.3} TSC/elem", c / 64.0); + (n, c) + } + + #[inline(always)] + unsafe fn ld(p: &[u8; 64]) -> __m512i { + unsafe { _mm512_loadu_si512(p.as_ptr() as *const __m512i) } + } + #[inline(always)] + unsafe fn ldi(p: &[i8; 64]) -> __m512i { + unsafe { _mm512_loadu_si512(p.as_ptr() as *const __m512i) } + } + #[inline(always)] + unsafe fn hsum(v: __m512i) -> i64 { + unsafe { _mm512_reduce_add_epi32(v) as i64 } + } + + /// Codes streamed from DRAM (64 MiB of code bytes), query weights L1-resident: + /// the shape of a scan over an index larger than cache. + fn stream(d: &Data, mag_t: &[u8; 512]) { + const BIG: usize = 1 << 20; // chunks of 64 code bytes + let codes: Vec<[u8; 64]> = (0..BIG).map(|i| d.mag[i % CHUNKS]).collect(); + let plane: Vec = (0..BIG).map(|i| d.pos[i % CHUNKS] ^ i as u64).collect(); + println!("== streamed from DRAM: 64 MiB of code bytes, weights in L1 =="); + let run = |name: &str, f: &dyn Fn() -> i64, dims: f64, bytes: f64| { + black_box(f()); + let t0 = Instant::now(); + let mut sink = 0i64; + for _ in 0..5 { + sink = sink.wrapping_add(f()); + } + let ns = t0.elapsed().as_nanos() as f64 / 5.0; + black_box(sink); + println!("{name:<44} {:7.4} ns/dim {:6.2} GB/s", ns / dims, bytes / ns); + }; + // SAFETY (both): avx512f/bw/vnni compiled in (module cfg); whole-array loads. + run( + "i4 permute-dot, 128 dims per 64 bytes", + &|| unsafe { + let lv = _mm512_broadcast_i32x4(_mm_loadu_si128(mag_t.as_ptr() as *const _)); + let mut acc = _mm512_setzero_si512(); + for (i, c) in codes.iter().enumerate() { + let codes = ld(c); + let w = ldi(&d.w[i % CHUNKS]); + let l = _mm512_shuffle_epi8(lv, _mm512_and_si512(codes, _mm512_set1_epi8(0x0F))); + let h = + _mm512_shuffle_epi8(lv, _mm512_and_si512(_mm512_srli_epi16(codes, 4), _mm512_set1_epi8(0x0F))); + acc = _mm512_dpbusd_epi32(_mm512_dpbusd_epi32(acc, l, w), h, w); + } + hsum(acc) + }, + (BIG * 128) as f64, + (BIG * 64) as f64, + ); + run( + "i9 stored SM9 C', 64 dims per 72 bytes", + &|| unsafe { + let (mut p, mut n) = (_mm512_setzero_si512(), _mm512_setzero_si512()); + for (i, (c, &pl)) in codes.iter().zip(&plane).enumerate() { + let (m, w) = (ld(c), ldi(&d.w[i % CHUNKS])); + p = _mm512_dpbusd_epi32(p, _mm512_maskz_mov_epi8(pl, m), w); + n = _mm512_dpbusd_epi32(n, _mm512_maskz_mov_epi8(!pl, m), w); + } + hsum(_mm512_sub_epi32(p, n)) + }, + (BIG * 64) as f64, + (BIG * 72) as f64, + ); + } + + pub fn run() { + let d = data(); + println!("== arithmetic (one accumulator chain per kernel, data in L1) =="); + // SAFETY (all kernels): the build has avx512f/bw/vnni/vbmi compiled in + // (module cfg); loads read whole 64-byte arrays. + let base = time("u8 x i8 dot4 (vpdpbusd, #386)", || unsafe { + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + acc = _mm512_dpbusd_epi32(acc, ld(&d.mag[c]), ldi(&d.w[c])); + } + hsum(acc) + }); + time("i8 x i8 dot4 (xor 0x80 bias, 1 dpbusd)", || unsafe { + let bias = _mm512_set1_epi8(-128); + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + acc = _mm512_dpbusd_epi32(acc, _mm512_xor_si512(ld(&d.mag[c]), bias), ldi(&d.w[c])); + } + hsum(acc) + }); + time("SM9 A: sign folded into w (w in -127..=127)", || unsafe { + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + let w = ldi(&d.w[c]); + let wn = _mm512_sub_epi8(_mm512_setzero_si512(), w); + let ws = _mm512_mask_blend_epi8(!d.pos[c], w, wn); + acc = _mm512_dpbusd_epi32(acc, ld(&d.mag[c]), ws); + } + hsum(acc) + }); + time("SM9 B: 2C limbs, dot(lo) - 256*dot(sign)", || unsafe { + let one = _mm512_set1_epi8(1); + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + let w = ldi(&d.w[c]); + let neg = !d.pos[c]; + // two's-complement low byte of -m is (256 - m) & 0xFF + let lo = _mm512_mask_sub_epi8(ld(&d.mag[c]), neg, _mm512_setzero_si512(), ld(&d.mag[c])); + let a = _mm512_dpbusd_epi32(acc, lo, w); + let t = _mm512_dpbusd_epi32(_mm512_setzero_si512(), _mm512_maskz_mov_epi8(neg, one), w); + acc = _mm512_sub_epi32(a, _mm512_slli_epi32(t, 8)); + } + hsum(acc) + }); + time("SM9 C: dot(m) - 2*dot(m & neg), one acc", || unsafe { + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + let (m, w) = (ld(&d.mag[c]), ldi(&d.w[c])); + let a = _mm512_dpbusd_epi32(acc, m, w); + let t = _mm512_dpbusd_epi32(_mm512_setzero_si512(), _mm512_maskz_mov_epi8(!d.pos[c], m), w); + acc = _mm512_sub_epi32(a, _mm512_add_epi32(t, t)); + } + hsum(acc) + }); + time("SM9 C': dot(m & pos) - dot(m & neg), two accs", || unsafe { + let (mut p, mut n) = (_mm512_setzero_si512(), _mm512_setzero_si512()); + for c in 0..CHUNKS { + let (m, w) = (ld(&d.mag[c]), ldi(&d.w[c])); + p = _mm512_dpbusd_epi32(p, _mm512_maskz_mov_epi8(d.pos[c], m), w); + n = _mm512_dpbusd_epi32(n, _mm512_maskz_mov_epi8(!d.pos[c], m), w); + } + hsum(_mm512_sub_epi32(p, n)) + }); + + println!("== 512-entry lookup: 9-bit ID (lo byte + hi plane) -> SM9 (magnitude + positive plane) =="); + let mag_t: [u8; 512] = core::array::from_fn(|i| (i as u32 * 97 % 256) as u8); + let pos_t: [u8; 512] = core::array::from_fn(|i| (i >= 256) as u8); + time("per-position reads (lut_u8 shape)", || { + let mut sink = 0i64; + for c in 0..CHUNKS { + let (lo, hi) = (&d.ids_lo[c], d.ids_hi[c]); + let mut out = [0u8; 64]; + let mut plane = 0u64; + for (i, o) in out.iter_mut().enumerate() { + let idx = usize::from(lo[i]) | ((((hi >> i) & 1) as usize) << 8); + *o = mag_t[idx]; + plane |= u64::from(pos_t[idx]) << i; + } + sink = sink.wrapping_add(i64::from(out[c & 63]) ^ plane as i64); + } + sink + }); + time("vpermi2b x4 + blends (mag) + x4 (sign)", || unsafe { + let t: [__m512i; 8] = core::array::from_fn(|k| _mm512_loadu_si512(mag_t.as_ptr().add(64 * k) as *const _)); + let s: [__m512i; 8] = core::array::from_fn(|k| _mm512_loadu_si512(pos_t.as_ptr().add(64 * k) as *const _)); + let mut sink = 0i64; + for c in 0..CHUNKS { + let lo = ld(&d.ids_lo[c]); + let b7 = _mm512_movepi8_mask(lo); + let hi = d.ids_hi[c]; + let pick = |tb: &[__m512i; 8]| { + let q0 = _mm512_permutex2var_epi8(tb[0], lo, tb[1]); + let q1 = _mm512_permutex2var_epi8(tb[2], lo, tb[3]); + let q2 = _mm512_permutex2var_epi8(tb[4], lo, tb[5]); + let q3 = _mm512_permutex2var_epi8(tb[6], lo, tb[7]); + let l = _mm512_mask_blend_epi8(b7, q0, q1); + let h = _mm512_mask_blend_epi8(b7, q2, q3); + _mm512_mask_blend_epi8(hi, l, h) + }; + let m = pick(&t); + let plane = _mm512_test_epi8_mask(pick(&s), pick(&s)); + sink = sink.wrapping_add(hsum(m) ^ plane as i64); + } + sink + }); + time("vpermi2b x2 + blend, antisymmetric codebook", || unsafe { + let t: [__m512i; 4] = + core::array::from_fn(|k| _mm512_loadu_si512(mag_t.as_ptr().add(256 + 64 * k) as *const _)); + let mut sink = 0i64; + for c in 0..CHUNKS { + let hi = d.ids_hi[c]; + // negative half reads entry 255 - lo of the positive table + let raw = ld(&d.ids_lo[c]); + let lo = _mm512_mask_blend_epi8(!hi, raw, _mm512_xor_si512(raw, _mm512_set1_epi8(-1))); + let b7 = _mm512_movepi8_mask(lo); + let q0 = _mm512_permutex2var_epi8(t[0], lo, t[1]); + let q1 = _mm512_permutex2var_epi8(t[2], lo, t[3]); + let m = _mm512_mask_blend_epi8(b7, q0, q1); + sink = sink.wrapping_add(hsum(m) ^ hi as i64); + } + sink + }); + + println!("== end to end per 64 lanes: lookup + SM9 dot vs i4 permute-dot =="); + time("i4 permute-dot (vpshufb 16-level + dpbusd)", || unsafe { + let lv = _mm512_broadcast_i32x4(_mm_loadu_si128(mag_t.as_ptr() as *const _)); + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + let codes = ld(&d.ids_lo[c]); + let l = _mm512_shuffle_epi8(lv, _mm512_and_si512(codes, _mm512_set1_epi8(0x0F))); + let h = _mm512_shuffle_epi8(lv, _mm512_and_si512(_mm512_srli_epi16(codes, 4), _mm512_set1_epi8(0x0F))); + acc = _mm512_dpbusd_epi32(acc, l, ldi(&d.w[c])); + acc = _mm512_dpbusd_epi32(acc, h, ldi(&d.w[c])); + } + hsum(acc) + }); + let e2e = time("i9 8+1: vpermi2b lookup + SM9 C'", || unsafe { + let t: [__m512i; 8] = core::array::from_fn(|k| _mm512_loadu_si512(mag_t.as_ptr().add(64 * k) as *const _)); + let s: [__m512i; 8] = core::array::from_fn(|k| _mm512_loadu_si512(pos_t.as_ptr().add(64 * k) as *const _)); + let (mut p, mut n) = (_mm512_setzero_si512(), _mm512_setzero_si512()); + for c in 0..CHUNKS { + let lo = ld(&d.ids_lo[c]); + let b7 = _mm512_movepi8_mask(lo); + let hi = d.ids_hi[c]; + let pick = |tb: &[__m512i; 8]| { + let q0 = _mm512_permutex2var_epi8(tb[0], lo, tb[1]); + let q1 = _mm512_permutex2var_epi8(tb[2], lo, tb[3]); + let q2 = _mm512_permutex2var_epi8(tb[4], lo, tb[5]); + let q3 = _mm512_permutex2var_epi8(tb[6], lo, tb[7]); + _mm512_mask_blend_epi8(hi, _mm512_mask_blend_epi8(b7, q0, q1), _mm512_mask_blend_epi8(b7, q2, q3)) + }; + let m = pick(&t); + let pos = _mm512_test_epi8_mask(pick(&s), pick(&s)); + let w = ldi(&d.w[c]); + p = _mm512_dpbusd_epi32(p, _mm512_maskz_mov_epi8(pos, m), w); + n = _mm512_dpbusd_epi32(n, _mm512_maskz_mov_epi8(!pos, m), w); + } + hsum(_mm512_sub_epi32(p, n)) + }); + time("4+5 additive: vpshufb C16 + vpermb R32, 2 dot4", || unsafe { + // value = C16[hi4] + R32[lo5], both stored +128-biased u8; the bias is + // removed once per query (2*128*sum(w)), outside this loop. + let c16 = _mm512_broadcast_i32x4(_mm_loadu_si128(mag_t.as_ptr() as *const _)); + let r32 = _mm512_loadu_si512(mag_t.as_ptr().add(64) as *const _); + let mut acc = _mm512_setzero_si512(); + for c in 0..CHUNKS { + let raw = ld(&d.ids_lo[c]); + let hi4 = _mm512_and_si512(_mm512_srli_epi16(raw, 4), _mm512_set1_epi8(0x0F)); + let lo5 = _mm512_and_si512(raw, _mm512_set1_epi8(0x1F)); + let w = ldi(&d.w[c]); + acc = _mm512_dpbusd_epi32(acc, _mm512_shuffle_epi8(c16, hi4), w); + acc = _mm512_dpbusd_epi32(acc, _mm512_permutexvar_epi8(lo5, r32), w); + } + hsum(acc) + }); + time("i9 values stored (no lookup) + SM9 C'", || unsafe { + let (mut p, mut n) = (_mm512_setzero_si512(), _mm512_setzero_si512()); + for c in 0..CHUNKS { + let (m, w) = (ld(&d.mag[c]), ldi(&d.w[c])); + p = _mm512_dpbusd_epi32(p, _mm512_maskz_mov_epi8(d.pos[c], m), w); + n = _mm512_dpbusd_epi32(n, _mm512_maskz_mov_epi8(!d.pos[c], m), w); + } + hsum(_mm512_sub_epi32(p, n)) + }); + stream(&d, &mag_t); + println!( + "note: i4 permute-dot scores 128 dims per 64 bytes, i9 scores 64 dims per 64 bytes + 8-byte plane; \ + per dimension the i9 end-to-end costs {:.2}x the u8 dot4 per 64 lanes", + e2e.0 / base.0 + ); + } +} + +/// Reconstruction error of the candidate 9-bit layouts against a unit Gaussian, +/// the coordinate distribution after TurboVec's rotation. The codebooks are +/// probe-local Lloyd-Max fits on a discretized density, not TurboVec's solver. +mod precision { + const BINS: usize = 1 << 20; + const LO: f64 = -7.0; + const HI: f64 = 7.0; + + fn grid() -> (Vec, Vec) { + let dx = (HI - LO) / BINS as f64; + let xs: Vec = (0..BINS).map(|i| LO + (i as f64 + 0.5) * dx).collect(); + let mut ws: Vec = xs.iter().map(|x| (-0.5 * x * x).exp()).collect(); + let z: f64 = ws.iter().sum(); + ws.iter_mut().for_each(|w| *w /= z); + (xs, ws) + } + + /// Lloyd-Max levels for weighted samples (xs sorted ascending). + fn lloyd(xs: &[f64], ws: &[f64], n: usize) -> Vec { + let (mut c, mut acc, mut lv) = (0.0, 0.0, vec![0.0; n]); + // Start at the high-resolution optimum: levels at the quantiles of the + // density raised to 1/3 (Panter-Dite), so few iterations are needed. + let w3: Vec = ws.iter().map(|w| w.cbrt()).collect(); + let z3: f64 = w3.iter().sum(); + let mut k = 0; + for (x, w) in xs.iter().zip(&w3) { + acc += w / z3; + while k < n && acc >= (k as f64 + 0.5) / n as f64 { + lv[k] = *x; + k += 1; + } + } + for _ in 0..1000 { + let (mut sw, mut sx) = (vec![0.0; n], vec![0.0; n]); + let mut j = 0; + for (x, w) in xs.iter().zip(ws) { + while j + 1 < n && *x > 0.5 * (lv[j] + lv[j + 1]) { + j += 1; + } + sw[j] += w; + sx[j] += w * x; + } + let mut moved = 0.0f64; + for i in 0..n { + if sw[i] > 0.0 { + let v = sx[i] / sw[i]; + moved = moved.max((v - lv[i]).abs()); + lv[i] = v; + } + } + c = moved; + if moved < 1e-12 { + break; + } + } + let _ = c; + lv + } + + fn mse(xs: &[f64], ws: &[f64], levels: &[f64]) -> f64 { + let mut l = levels.to_vec(); + l.sort_by(f64::total_cmp); + let mut j = 0; + let mut e = 0.0; + for (x, w) in xs.iter().zip(ws) { + while j + 1 < l.len() && *x > 0.5 * (l[j] + l[j + 1]) { + j += 1; + } + e += w * (x - l[j]) * (x - l[j]); + } + e + } + + fn snap(levels: &[f64], max_code: f64) -> (Vec, usize) { + let s = max_code / levels.iter().fold(0.0f64, |m, v| m.max(v.abs())); + let mut q: Vec = levels.iter().map(|v| (v * s).round() / s).collect(); + q.sort_by(f64::total_cmp); + q.dedup(); + let n = q.len(); + (q, n) + } + + pub fn run() { + let (xs, ws) = grid(); + let db = |m: f64| 10.0 * (1.0 / m).log10(); + println!("== precision on N(0,1): MSE and SNR (probe-local Lloyd-Max, not TurboVec's solver) =="); + let l16 = lloyd(&xs, &ws, 16); + let (q16, _) = snap(&l16, 127.0); + println!("i4 16 levels (f64) MSE {:.3e} {:5.2} dB", mse(&xs, &ws, &l16), db(mse(&xs, &ws, &l16))); + println!("i4 16 levels on i8 grid MSE {:.3e} {:5.2} dB", mse(&xs, &ws, &q16), db(mse(&xs, &ws, &q16))); + let l256 = lloyd(&xs, &ws, 256); + let (q256, n256) = snap(&l256, 127.0); + println!("i8 256 levels (f64) MSE {:.3e} {:5.2} dB", mse(&xs, &ws, &l256), db(mse(&xs, &ws, &l256))); + println!( + "i8 256 levels on i8 grid MSE {:.3e} {:5.2} dB ({n256} distinct)", + mse(&xs, &ws, &q256), + db(mse(&xs, &ws, &q256)) + ); + let l512 = lloyd(&xs, &ws, 512); + let (q512, n512) = snap(&l512, 255.0); + println!("i9 512 levels (f64, ideal) MSE {:.3e} {:5.2} dB", mse(&xs, &ws, &l512), db(mse(&xs, &ws, &l512))); + println!( + "i9 512 levels on SM9 grid MSE {:.3e} {:5.2} dB ({n512} distinct of 512)", + mse(&xs, &ws, &q512), + db(mse(&xs, &ws, &q512)) + ); + // 4+5 additive: coarse 16-level Lloyd, then one shared 32-level residual Lloyd. + let coarse = l16.clone(); + let mut rx = Vec::with_capacity(BINS); + { + let mut j = 0; + for x in &xs { + while j + 1 < coarse.len() && *x > 0.5 * (coarse[j] + coarse[j + 1]) { + j += 1; + } + rx.push(x - coarse[j]); + } + } + let mut order: Vec = (0..BINS).collect(); + order.sort_by(|a, b| rx[*a].total_cmp(&rx[*b])); + let (rxs, rws): (Vec, Vec) = order.iter().map(|&i| (rx[i], ws[i])).unzip(); + let r32 = lloyd(&rxs, &rws, 32); + let sum: Vec = coarse + .iter() + .flat_map(|c| r32.iter().map(move |r| c + r)) + .collect(); + // The stored-value layout is a uniform 9-bit grid, so the Lloyd IDs buy nothing + // there; what matters is the clip. Sweep the step for the best uniform SM9 grid. + let mut best = (f64::MAX, 0.0); + for k in 0..200 { + let clip = 2.5 + 0.01 * k as f64; + let grid: Vec = (-255..=255).map(|m| m as f64 * clip / 255.0).collect(); + let e = mse(&xs, &ws, &grid); + if e < best.0 { + best = (e, clip); + } + } + let mut best8 = (f64::MAX, 0.0); + for k in 0..200 { + let clip = 2.0 + 0.01 * k as f64; + let grid: Vec = (-127..=127).map(|m| m as f64 * clip / 127.0).collect(); + let e = mse(&xs, &ws, &grid); + if e < best8.0 { + best8 = (e, clip); + } + } + println!( + "i8 uniform grid, best clip MSE {:.3e} {:5.2} dB (clip {:.2} sigma, 255 values)", + best8.0, + db(best8.0), + best8.1 + ); + println!( + "SM9 uniform grid, best clip MSE {:.3e} {:5.2} dB (clip {:.2} sigma, 511 values)", + best.0, + db(best.0), + best.1 + ); + println!( + "4+5 C16+R32 additive (f64) MSE {:.3e} {:5.2} dB (residual codebook fit after the coarse split)", + mse(&xs, &ws, &sum), + db(mse(&xs, &ws, &sum)) + ); + // As executed: each table on its own i8 grid (two scales, two accumulators). + let q8 = |v: &[f64]| { + let sc = 127.0 / v.iter().fold(0.0f64, |m, x| m.max(x.abs())); + v.iter() + .map(|x| (x * sc).round() / sc) + .collect::>() + }; + let (cq, rq) = (q8(&coarse), q8(&r32)); + let sumq: Vec = cq + .iter() + .flat_map(|c| rq.iter().map(move |r| c + r)) + .collect(); + println!( + "4+5 C16+R32 on two i8 grids MSE {:.3e} {:5.2} dB", + mse(&xs, &ws, &sumq), + db(mse(&xs, &ws, &sumq)) + ); + } +} + +fn main() { + precision::run(); + #[cfg(all(target_arch = "x86_64", target_feature = "avx512vnni", target_feature = "avx512vbmi"))] + probe::run(); + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512vnni", target_feature = "avx512vbmi")))] + eprintln!("u9_cost_probe needs avx512vnni + avx512vbmi compiled in (target-cpu=native on such a host)"); +} diff --git a/src/simd.rs b/src/simd.rs index acafd0e4..a0c21619 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -2409,6 +2409,95 @@ mod dot4_nibble_facade_parity { } } + fn masked_reference(acc: [i32; 16], a: &[u8; 64], keep: u64, b: &[i8; 64]) -> [i32; 16] { + let kept: [u8; 64] = core::array::from_fn(|i| if keep >> i & 1 == 1 { a[i] } else { 0 }); + reference(acc, &kept, b) + } + + #[test] + fn dot4_masked_matches_reference() { + let (us, is) = (bytes(), signed()); + let mut s = 0x6A09_E667_F3BC_C908u64; + let mut checked = 0; + for a in &us { + for b in &is { + 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 acc: [i32; 16] = core::array::from_fn(|_| next(&mut s) as i32); + let got = I32x16::from_array(acc) + .dot4_u8_i8_add_masked(U8x64::from_array(*a), keep, I8x64::from_array(*b)) + .to_array(); + assert_eq!(got, masked_reference(acc, a, keep, b), "keep {keep:#x}\na {a:?}\nb {b:?}"); + checked += 1; + } + } + } + assert!(checked > 10_000, "corpus too small: {checked}"); + } + + /// Bit `i` of `keep` must gate byte `i` and nothing else. + #[test] + fn dot4_masked_gates_each_byte_by_its_own_bit() { + let a = U8x64::splat(7); + for i in 0..64 { + let got = I32x16::splat(0) + .dot4_u8_i8_add_masked(a, 1 << i, I8x64::splat(-3)) + .to_array(); + let want: [i32; 16] = core::array::from_fn(|j| if j == i / 4 { -21 } else { 0 }); + assert_eq!(got, want, "bit {i}"); + let all_but = I32x16::splat(0) + .dot4_u8_i8_add_masked(a, !(1 << i), I8x64::splat(-3)) + .to_array(); + let want: [i32; 16] = core::array::from_fn(|j| if j == i / 4 { -63 } else { -84 }); + assert_eq!(all_but, want, "all but bit {i}"); + } + } + + /// Every 9-bit sign-magnitude value (all 256 magnitudes in both polarities, + /// so 512 codewords including the never-touched `P=0, M=0` and the touched + /// `+0`) times every weight edge, through the two-accumulator form, against + /// an independent `i64` oracle. + #[test] + fn sm9_pos_minus_neg_is_exact_for_every_codeword() { + for w in [-128i8, -127, -1, 0, 1, 127] { + for positive in [false, true] { + for chunk in 0..4u32 { + let mags: [u8; 64] = core::array::from_fn(|i| (64 * chunk + i as u32) as u8); + let plane = if positive { u64::MAX } else { 0 }; + let (m, wv) = (U8x64::from_array(mags), I8x64::splat(w)); + let pos = I32x16::splat(0) + .dot4_u8_i8_add_masked(m, plane, wv) + .to_array(); + let neg = I32x16::splat(0) + .dot4_u8_i8_add_masked(m, !plane, wv) + .to_array(); + for j in 0..16 { + let want: i64 = (4 * j..4 * j + 4) + .map(|i| { + let v = i64::from(mags[i]); + (if positive { v } else { -v }) * i64::from(w) + }) + .sum(); + assert_eq!(i64::from(pos[j]) - i64::from(neg[j]), want, "w {w} positive {positive} lane {j}"); + } + } + } + } + } + + /// The doc promise: 16_448 worst-case calls fit an `i32` accumulator. + #[test] + fn dot4_masked_worst_case_holds_16448_calls() { + let (a, b) = (U8x64::splat(255), I8x64::splat(-128)); + let mut acc = I32x16::splat(0); + for _ in 0..16_448 { + acc = acc.dot4_u8_i8_add_masked(a, u64::MAX, b); + } + assert_eq!(acc.to_array(), [(-16_448i64 * 130_560) as i32; 16]); + assert!(-16_448i64 * 130_560 >= i64::from(i32::MIN)); + assert!(-16_449i64 * 130_560 < i64::from(i32::MIN)); + } + #[test] fn nibbles_match_reference() { for a in bytes() { diff --git a/src/simd_avx.rs b/src/simd_avx.rs index 5038558f..7511e6b7 100644 --- a/src/simd_avx.rs +++ b/src/simd_avx.rs @@ -357,6 +357,46 @@ impl I32x16 { } Self(out) } + + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let mut kept = a.to_array(); + // 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. + unsafe { + let bit = _mm_set1_epi64x(crate::simd_masking_ops::BIT_IN_BYTE as i64); + let (ps, pk) = (spread.as_ptr() as *const __m128i, kept.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); + _mm_storeu_si128(pk.add(q), _mm_and_si128(_mm_loadu_si128(pk.add(q)), m)); + } + } + self.dot4_u8_i8_add(U8x64::from_array(kept), b) + } } impl I32x16 { diff --git a/src/simd_avx2.rs b/src/simd_avx2.rs index 5239b943..42e5837a 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -3007,6 +3007,46 @@ impl I32x16 { } Self(out) } + + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let mut kept = a.to_array(); + // 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. 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. + unsafe { + let bit = _mm256_set1_epi64x(crate::simd_masking_ops::BIT_IN_BYTE as i64); + let (ps, pk) = (spread.as_ptr() as *const __m256i, kept.as_mut_ptr() as *mut __m256i); + for h in 0..2 { + let m = _mm256_cmpeq_epi8(_mm256_and_si256(_mm256_loadu_si256(ps.add(h)), bit), bit); + _mm256_storeu_si256(pk.add(h), _mm256_and_si256(_mm256_loadu_si256(pk.add(h)), m)); + } + } + self.dot4_u8_i8_add(U8x64::from_array(kept), b) + } } #[cfg(not(all(target_feature = "avx", not(target_feature = "avx2"))))] diff --git a/src/simd_avx512.rs b/src/simd_avx512.rs index b82aca21..6536849b 100644 --- a/src/simd_avx512.rs +++ b/src/simd_avx512.rs @@ -962,6 +962,33 @@ impl I32x16 { Self(acc) } + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> Self { + // SAFETY: avx512bw is compiled in on this arm; register-only. + self.dot4_u8_i8_add(U8x64(unsafe { _mm512_maskz_mov_epi8(keep, a.0) }), b) + } + pub const LANES: usize = 16; #[inline(always)] diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index aa5ee9cc..9bf0a1ca 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -5295,6 +5295,21 @@ fn live8(live: usize) -> u8 { } } +/// One bit per byte, bit `k` in byte `k`: ANDed with a word of +/// [`mask_byte_spread`] and compared for equality against itself, it turns +/// mask bits into 0xFF / 0x00 bytes. +pub(crate) const BIT_IN_BYTE: u64 = 0x8040_2010_0804_0201; + +/// The 64 mask bits of `keep` as eight little-endian words in which every +/// byte of word `j` holds mask byte `j`. Byte `i` of the result then carries +/// bit `i` of `keep` at bit position `i % 8`, so one AND with +/// [`BIT_IN_BYTE`] and one byte compare per register expand the mask without +/// moving a byte. +#[inline(always)] +pub(crate) fn mask_byte_spread(keep: u64) -> [u64; 8] { + core::array::from_fn(|j| (keep >> (8 * j) & 0xFF) * 0x0101_0101_0101_0101) +} + /// `live`-lane validity mask for a 64-lane group (one `u8x64` register is one /// whole mask word): all ones for a full group, the low `live` bits otherwise. #[inline(always)] diff --git a/src/simd_neon.rs b/src/simd_neon.rs index 31fb170d..b2144a4d 100644 --- a/src/simd_neon.rs +++ b/src/simd_neon.rs @@ -3332,6 +3332,46 @@ impl I32x16 { Self(out) } + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: crate::simd::U8x64, keep: u64, b: crate::simd::I8x64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let mut kept = a.to_array(); + for q in 0..4 { + // SAFETY: NEON baseline. Each load and store 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. + 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 p = kept.as_mut_ptr().add(16 * q); + vst1q_u8(p, vandq_u8(vld1q_u8(p), m)); + } + } + self.dot4_u8_i8_add(crate::simd::U8x64::from_array(kept), b) + } + /// 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 4122e72f..85bd9e41 100644 --- a/src/simd_nightly/i_word_types.rs +++ b/src/simd_nightly/i_word_types.rs @@ -389,6 +389,34 @@ impl I32x16 { Self(acc) } + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: super::U8x64, keep: u64, b: super::I8x64) -> Self { + use core::simd::{Mask, Simd}; + let kept = Mask::::from_bitmask(keep).select(a.0, Simd::splat(0)); + self.dot4_u8_i8_add(super::U8x64(kept), b) + } + pub const LANES: usize = 16; /// Per-lane `SaturatingRoundingDoublingHighMul` (gemmlowp / NEON `VQRDMULH`): diff --git a/src/simd_scalar.rs b/src/simd_scalar.rs index 1aa1f8b7..d90fef4b 100644 --- a/src/simd_scalar.rs +++ b/src/simd_scalar.rs @@ -1147,6 +1147,33 @@ impl I32x16 { Self::from_array(acc) } + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> Self { + let a = a.to_array(); + self.dot4_u8_i8_add(U8x64::from_array(core::array::from_fn(|i| if keep >> i & 1 == 1 { a[i] } else { 0 })), b) + } + /// Per-lane `SaturatingRoundingDoublingHighMul` (gemmlowp / NEON `VQRDMULH`): /// `round(2*a*b / 2^32)`, ties toward +inf, `MIN * MIN` saturating to `MAX`. /// diff --git a/src/simd_wasm.rs b/src/simd_wasm.rs index 0e941dcd..4317b7ff 100644 --- a/src/simd_wasm.rs +++ b/src/simd_wasm.rs @@ -2554,6 +2554,46 @@ pub mod wasm32_simd { Self(out) } + /// [`Self::dot4_u8_i8_add`] with the bytes of `a` whose bit in `keep` is 0 + /// replaced by 0: lane `j` gains the sum of `a[i]*b[i]` over the bytes + /// `i` in `4j..4j+4` with bit `i` of `keep` set. Exact for every `u8` and + /// every `i8`, `-128` included; only the final addition to the accumulator + /// wraps on `i32` overflow. + /// + /// Two calls over the same bytes give a 9-bit sign-magnitude dot product + /// (magnitude in `a`, positive plane in `keep`) without widening the + /// operands: `pos.dot4_u8_i8_add_masked(m, plane, w) - neg.dot4_u8_i8_add_masked(m, !plane, w)`. + /// Each accumulator moves by at most `4 * 255 * 128 = 130_560` per call, so + /// it holds 16_448 calls before it can wrap. + /// + /// ``` + /// use ndarray::simd::{I32x16, I8x64, U8x64}; + /// let a = U8x64::splat(255); + /// let b = I8x64::splat(-128); + /// // Lane 0 keeps bytes 0 and 3, every other lane keeps nothing. + /// let got = I32x16::splat(0).dot4_u8_i8_add_masked(a, 0b1001, b).to_array(); + /// assert_eq!(got[0], -2 * 255 * 128); + /// assert_eq!(got[1..], [0; 15]); + /// ``` + #[inline(always)] + pub fn dot4_u8_i8_add_masked(self, a: crate::simd::U8x64, keep: u64, b: crate::simd::I8x64) -> Self { + let spread = crate::simd_masking_ops::mask_byte_spread(keep); + let mut kept = a.to_array(); + let bit = u64x2_splat(crate::simd_masking_ops::BIT_IN_BYTE); + for q in 0..4 { + // SAFETY: simd128 is compiled in on this arm. Each load and store touches 16 bytes + // at offset 16*q of a 64-byte array (`spread` is `[u64; 8]`), 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. + unsafe { + let m = i8x16_eq(v128_and(v128_load((spread.as_ptr() as *const v128).add(q)), bit), bit); + let p = (kept.as_mut_ptr() as *mut v128).add(q); + v128_store(p, v128_and(v128_load(p), m)); + } + } + self.dot4_u8_i8_add(crate::simd::U8x64::from_array(kept), b) + } + /// Per-lane `SaturatingRoundingDoublingHighMul` (gemmlowp / NEON `VQRDMULH`): /// `round(2*a*b / 2^32)`, ties toward +inf, `MIN * MIN` saturating to `MAX`. /// From c498cf457663ebb7bc33fbdeb259eb6cfba53830 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 15:34:29 +0000 Subject: [PATCH 2/6] Import core::simd::Select for the nightly masked dot4 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- src/simd_nightly/i_word_types.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/simd_nightly/i_word_types.rs b/src/simd_nightly/i_word_types.rs index 85bd9e41..24a9b688 100644 --- a/src/simd_nightly/i_word_types.rs +++ b/src/simd_nightly/i_word_types.rs @@ -412,7 +412,7 @@ impl I32x16 { /// ``` #[inline(always)] pub fn dot4_u8_i8_add_masked(self, a: super::U8x64, keep: u64, b: super::I8x64) -> Self { - use core::simd::{Mask, Simd}; + use core::simd::{Mask, Select, Simd}; let kept = Mask::::from_bitmask(keep).select(a.0, Simd::splat(0)); self.dot4_u8_i8_add(super::U8x64(kept), b) } From 33a3e4db8712f215d84a3301ee3bb3a158a91c4c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 15:40:59 +0000 Subject: [PATCH 3/6] Fuse the AVX2 masked dot4 and time the facade on each tier The AVX2 arm applies the byte mask inside the dot4 loop instead of writing a masked copy first. The probe gains a facade section that times dot4, two plain accumulators and the two masked accumulators on whatever tier the build compiled. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- examples/u9_cost_probe.rs | 63 +++++++++++++++++++++++++++++++++++++++ src/simd_avx2.rs | 21 +++++++++---- 2 files changed, 78 insertions(+), 6 deletions(-) diff --git a/examples/u9_cost_probe.rs b/examples/u9_cost_probe.rs index 97a651fe..874a59de 100644 --- a/examples/u9_cost_probe.rs +++ b/examples/u9_cost_probe.rs @@ -512,8 +512,71 @@ mod precision { } } +/// The same two kernels through the `ndarray::simd` facade, on whatever tier +/// this build compiled: what a consumer actually calls. +mod facade { + use ndarray::simd::{I32x16, I8x64, U8x64}; + use std::hint::black_box; + use std::time::Instant; + + pub fn run() { + const CHUNKS: usize = 256; + const PASSES: usize = 4000; + 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 >> 11 + }; + let mag: Vec = (0..CHUNKS) + .map(|_| U8x64::from_array(core::array::from_fn(|_| next() as u8))) + .collect(); + let w: Vec = (0..CHUNKS) + .map(|_| I8x64::from_array(core::array::from_fn(|_| i8::from_ne_bytes([next() as u8])))) + .collect(); + let pos: Vec = (0..CHUNKS).map(|_| next() ^ (next() << 32)).collect(); + let time = |name: &str, f: &dyn Fn() -> i32| { + black_box(f()); + let t0 = Instant::now(); + let mut sink = 0i32; + for _ in 0..PASSES { + sink = sink.wrapping_add(f()); + } + let ns = t0.elapsed().as_nanos() as f64 / (PASSES * CHUNKS) as f64; + black_box(sink); + println!("{name:<44} {ns:7.3} ns/64 lanes"); + }; + println!("== facade (ndarray::simd, this build's tier) =="); + time("I32x16::dot4_u8_i8_add", &|| { + let mut acc = I32x16::splat(0); + for c in 0..CHUNKS { + acc = acc.dot4_u8_i8_add(mag[c], w[c]); + } + acc.to_array()[0] + }); + time("two dot4_u8_i8_add accumulators, no mask", &|| { + let (mut p, mut n) = (I32x16::splat(0), I32x16::splat(0)); + for c in 0..CHUNKS { + p = p.dot4_u8_i8_add(mag[c], w[c]); + n = n.dot4_u8_i8_add(mag[(c + 1) % CHUNKS], w[c]); + } + p.to_array()[0].wrapping_sub(n.to_array()[0]) + }); + time("SM9: two dot4_u8_i8_add_masked accumulators", &|| { + let (mut p, mut n) = (I32x16::splat(0), I32x16::splat(0)); + for c in 0..CHUNKS { + p = p.dot4_u8_i8_add_masked(mag[c], pos[c], w[c]); + n = n.dot4_u8_i8_add_masked(mag[c], !pos[c], w[c]); + } + p.to_array()[0].wrapping_sub(n.to_array()[0]) + }); + } +} + fn main() { precision::run(); + facade::run(); #[cfg(all(target_arch = "x86_64", target_feature = "avx512vnni", target_feature = "avx512vbmi"))] probe::run(); #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512vnni", target_feature = "avx512vbmi")))] diff --git a/src/simd_avx2.rs b/src/simd_avx2.rs index 42e5837a..7df23ac6 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -3032,20 +3032,29 @@ impl I32x16 { #[inline(always)] pub fn dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> Self { let spread = crate::simd_masking_ops::mask_byte_spread(keep); - let mut kept = a.to_array(); + 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. 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. + // `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; the rest is `dot4_u8_i8_add`. unsafe { + let low = _mm256_set1_epi16(0x00FF); let bit = _mm256_set1_epi64x(crate::simd_masking_ops::BIT_IN_BYTE as i64); - let (ps, pk) = (spread.as_ptr() as *const __m256i, kept.as_mut_ptr() as *mut __m256i); + let ps = spread.as_ptr() as *const __m256i; + 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 m = _mm256_cmpeq_epi8(_mm256_and_si256(_mm256_loadu_si256(ps.add(h)), bit), bit); - _mm256_storeu_si256(pk.add(h), _mm256_and_si256(_mm256_loadu_si256(pk.add(h)), m)); + let av = _mm256_and_si256(_mm256_loadu_si256(pa.add(h)), m); + let bv = _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.dot4_u8_i8_add(U8x64::from_array(kept), b) + Self(out) } } From 0c7531317e1b60b3939620ee24bbbdcf1fffd5ff Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 15:44:40 +0000 Subject: [PATCH 4/6] Record the 9-bit TurboVec9 measurements and handoff Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- .claude/blackboard.md | 20 +++ .claude/knowledge/turbovec9-u9-substrate.md | 150 ++++++++++++++++++++ 2 files changed, 170 insertions(+) create mode 100644 .claude/knowledge/turbovec9-u9-substrate.md diff --git a/.claude/blackboard.md b/.claude/blackboard.md index 44321e56..920e86a5 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3876,3 +3876,23 @@ Loose ends: AVX-VNNI (`avxvnni`, 256-bit `vpdpbusd`) is not used on the AVX2 arm is compiled in. Neither is measured. The signed form (`i8 x i8`, turbovec's `xor 0x80` + `zero` correction) is not a facade word yet. + +## dot4_u8_i8_add_masked + 9-bit TurboVec9 cost probe (2026-10-29, branch claude/u9-turbovec9-substrate) + +Follow-up to #386. Added on every backend: `I32x16::dot4_u8_i8_add_masked(a, keep, b)`; two calls with the positive +plane and its complement give an exact 9-bit sign-magnitude x i8 dot product (all i8 incl. -128, i32 accumulators, +no widening). Crate-private `simd_masking_ops::{BIT_IN_BYTE, mask_byte_spread}` expand a u64 mask to bytes with one +AND + one compare per register. Measurements, precision tables and the TurboVec handoff: +`.claude/knowledge/turbovec9-u9-substrate.md`; probe `examples/u9_cost_probe.rs` (R1-excluded: it times shuffle +lookups on purpose). + +Witnessed: 8/8 dot4 tests on native (VNNI), v4, v3, avx (qemu), v2 (qemu), aarch64 (qemu), nightly; neon/wasm parity +programs pass (0x504, 0x505). clippy -D warnings clean on native, v4, v3. Disable runs all red: unmasked AVX-512 move, +reversed BIT_IN_BYTE (AVX2), mask dropped in the fused AVX2 loop. + +Key numbers (AVX-512 VNNI, L1): masked pair 1.27-1.33 ns/64 lanes vs dot4 1.10-1.19; 512-entry lookup 6.1 ns +(vpermi2b) / 38-40 ns (per-position reads). Stored uniform SM9 46.0 dB vs i4 20.2 dB at the same per-dim L1 cost and +1.4x per dim streamed. AVX2: masked pair 4.5 ns vs dot4 1.5 ns. + +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. diff --git a/.claude/knowledge/turbovec9-u9-substrate.md b/.claude/knowledge/turbovec9-u9-substrate.md new file mode 100644 index 00000000..f1531264 --- /dev/null +++ b/.claude/knowledge/turbovec9-u9-substrate.md @@ -0,0 +1,150 @@ +# TurboVec9 on ndarray: 9-bit sign-magnitude scoring, measured + +READ BY: savant-architect, vector-synthesis, simd-savant, anyone wiring a +9-bit index (TurboVec9) or a DR512 dot product onto `ndarray::simd`. + +Status: probe + one facade word (`I32x16::dot4_u8_i8_add_masked`). Numbers +below come from `examples/u9_cost_probe.rs` on one AVX-512 VNNI + VBMI Xeon +(target-cpu=native), three runs, reported as ranges. They are not TurboVec +recall numbers; the codebooks in the probe are probe-local Lloyd-Max fits. + +## Two contracts, never mixed + +- **DR512 / Codec9** (`hpc::codec9`, `simd_stride::lut9`): u8 magnitude plus + one P bit per value, -255..=255, and `P=0, M=0` is the never-touched + sentinel. `lut9` belongs here and nowhere else. +- **TurboVec9**: an ID 0..=511 selects one of 512 centroids. Every ID is + valid. An ID is not a Codec9 value and is never read as sign-magnitude. + +The facade word below is pure arithmetic over a u8 magnitude and a u64 +plane. It does not know which contract produced them. For DR512 data both +`P=0, M=0` and `P=1, M=0` contribute 0 to a dot product; a caller that must +tell them apart reads the plane itself. + +## The facade word + +```rust +// Lane j gains the sum of a[i] * b[i] over bytes i in 4j..4j+4 with bit i of keep set. +I32x16::dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> I32x16 +``` + +9-bit sign-magnitude times i8, exact for every magnitude and every i8 +(-128 included), i32 accumulators, operands never widened: + +```rust +let (mut pos, mut neg) = (I32x16::splat(0), I32x16::splat(0)); +for blk in blocks { // 64 dims per block + pos = pos.dot4_u8_i8_add_masked(blk.magnitude, blk.positive, w); + neg = neg.dot4_u8_i8_add_masked(blk.magnitude, !blk.positive, w); +} +let score_lanes = pos - neg; // 16 lanes, each a 4-dim partial sum +``` + +Overflow: each accumulator moves by at most `4 * 255 * 128 = 130_560` per +call, so 16_448 calls (65_792 dims per lane) cannot wrap; the test +`dot4_masked_worst_case_holds_16448_calls` pins both sides of that bound. +No saturating instruction is involved (no `vpmaddubsw`, no `vpdpbusds`). + +Realizations: AVX-512 one `vmovdqu8{z}` before `vpdpbusd` (or the widen + +`vpmaddwd` path without VNNI); AVX2, AVX, NEON and wasm expand the mask +with one AND against `BIT_IN_BYTE` and one byte compare per register (no +shuffle, no per-lane loop); nightly selects through `core::simd::Mask`; +scalar is the reference loop. + +Witnessed: `simd::dot4_nibble_facade_parity` 8/8 on native (AVX-512 VNNI), +config-v4 (AVX-512 without VNNI), config-v3 (AVX2), config-avx (qemu +SandyBridge), config-v2 (scalar, qemu Nehalem), aarch64 (qemu) and +nightly-simd; `neon-parity.sh` and `wasm-parity.sh` pass (codes 0x504, +0x505). Disable runs, all red: unmasked AVX-512 move (3 tests), reversed +`BIT_IN_BYTE` on AVX2 (2 tests), mask dropped in the fused AVX2 loop (3). + +## Arithmetic: three exact forms compared (AVX-512 VNNI, L1, ns per 64 lanes) + +| kernel | ns / 64 lanes | +|---|---| +| u8 x i8 `vpdpbusd` (#386 `dot4_u8_i8_add`) | 1.10 - 1.19 | +| i8 x i8 via `xor 0x80` bias | 1.20 - 1.23 | +| A: sign folded into w (exact only for w in -127..=127) | 1.24 - 1.50 | +| B: two's-complement limbs, `dot(lo) - 256 * dot(sign)` | 1.88 - 1.89 | +| C: `dot(m) - 2 * dot(m & neg)`, one accumulator | 1.79 - 1.84 | +| **C': `dot(m & pos) - dot(m & neg)`, two accumulators** | **1.27 - 1.33** | + +The loop is one latency-bound chain. B and C put two dependent ops on the +accumulator; C' keeps two independent chains and lands within ~10% of the +plain u8 dot. A is not exact for w = -128 and was not added. + +Through the facade, same host: `dot4_u8_i8_add` 1.12, two masked +accumulators 1.29. On config-v3 (AVX2, no VNNI): `dot4_u8_i8_add` 1.51, +two plain accumulators 3.26, two masked accumulators 4.48 - 4.58. Without +VNNI every call widens to 16 bits, so a pair costs two widenings; a fused +SM9 op that widens once and negates in 16 bits would roughly halve that, +and is not built. + +## 512-entry lookup (WP2): keep it out of the scan + +ID = low byte + high-bit plane, decoded to magnitude + positive plane: + +| lookup | ns / 64 lanes | +|---|---| +| per-position reads (`lut_u8` shape, two 512-entry tables) | 38 - 40 | +| 4x `vpermi2b` + blends for magnitude, same for the sign | 6.07 - 6.16 | +| 2x `vpermi2b` on an antisymmetric codebook (`255 - lo` for negatives) | 6.79 - 6.94 | + +`lut9` and `lut_u8` keep their contracts; no 512-entry facade word was +added. The lookup alone costs ~5x the arithmetic it feeds. TurboVec's +levels do not depend on the query, so the decode can happen once at encode +time; then the scan stores values, not IDs (next section). + +## Precision vs cost: the layouts (N(0,1) coordinates) + +| layout | SNR | scan cost (AVX-512 VNNI) | +|---|---|---| +| i4, 16 Lloyd-Max levels (today) | 20.22 dB | 2.64 - 2.70 ns / 128 dims (L1) | +| i8 uniform grid, best clip | 40.54 dB | 1.10 - 1.19 ns / 64 dims | +| 512 Lloyd-Max IDs, executed by `vpermi2b` lookup | 49.86 dB | 5.39 - 5.53 ns / 64 dims | +| 512 Lloyd-Max levels snapped to the SM9 grid | 44.74 dB (329 distinct) | 1.29 - 1.31 ns / 64 dims | +| **SM9 uniform grid, best clip (4.21 sigma), stored** | **46.02 dB** | **1.29 - 1.31 ns / 64 dims** | +| 4+5 additive C16 + R32, two i8 grids | 43.25 dB | 2.58 - 2.65 ns / 64 dims | + +Streamed from DRAM (64 MiB of codes, weights in L1), per dimension: +i4 0.031 - 0.035 ns, stored SM9 0.043 - 0.049 ns (72 bytes per 64 dims). + +Reading the table: + +- Stored values make the 512 Lloyd-Max IDs pointless: snapping them to + the 9-bit grid keeps 329 distinct values and loses to a plain uniform + 9-bit grid with a tuned clip. The Lloyd-Max levels only pay off when + executed as IDs, and the lookup makes that path ~4x slower than the + stored one. +- The 4+5 additive competitor is worse than stored SM9 on both axes in this + probe. Its residual codebook is one shared Lloyd fit after a 16-level + split; a joint fit could do better, but it cannot beat the 512-point + optimum and it pays two lookups. + +## Answer to the acceptance question + +On AVX-512 VNNI: yes, for the stored-value layout. 64 dims cost one +masked-pair of `vpdpbusd`, ~1.3 ns in L1 (about the same per dimension as +today's i4 permute-dot) and 1.4x i4 per dimension when streamed from DRAM, +for +25.8 dB over i4. With IDs plus an in-scan 512 lookup: no, the lookup +dominates. On AVX2 the pair costs ~3x the plain dot; a fused SM9 op is the +next step there if AVX2 matters. + +## Handoff to the TurboVec owner (owns ABI, encode, pack, search, mutation, serialization, Python) + +Proposed scan layout, per 64 dimensions of one vector, vector-major as +today: 64 magnitude bytes plus one u64 positive plane (72 bytes). Scoring: +the two-accumulator loop above, then the existing `scale` / `zero` +correction (`zero` is 0 here: magnitudes are unbiased u8, weights stay i8). +Open decisions on the TurboVec side: grid clip per dimension or global, how +the encoder writes 0 (either plane bit scores 0), and recall measured on +real data with TurboVec's own codebook solver. None of this was changed +from this session. + +## Not done + +- No NEON or wasm timings (qemu timings mean nothing); NEON `SDOT` / + `USDOT` and AVX-VNNI (256-bit) are not used by the masked word. +- No fused single-widen SM9 op for non-VNNI tiers. +- No signed i8 x i8 facade word. +- Recall on a real index. From 98cb65ad8d4947eab577d4cd7fd748c64b6778b1 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 15:58:07 +0000 Subject: [PATCH 5/6] Make the probe kernels exact and consume every lane; build the AVX2 mask in registers Kernel B subtracted 256 * w for a negative zero (magnitude 0, plane bit clear): the sign limb now skips zero magnitudes. Every exact SM9 kernel, and the facade pair, must match an independent i64 reference before its time is printed. The facade kernels returned lane 0 only, which let the optimizer drop half of the AVX2 work; they now fold all 16 lanes. The AVX2 masked dot4 builds its byte mask from the u64 with a variable shift and a multiply instead of a stored spread. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- examples/u9_cost_probe.rs | 52 ++++++++++++++++++++++++++++++--------- src/simd_avx2.rs | 23 ++++++++++++----- 2 files changed, 57 insertions(+), 18 deletions(-) diff --git a/examples/u9_cost_probe.rs b/examples/u9_cost_probe.rs index 874a59de..680f8574 100644 --- a/examples/u9_cost_probe.rs +++ b/examples/u9_cost_probe.rs @@ -78,6 +78,12 @@ mod probe { unsafe fn ldi(p: &[i8; 64]) -> __m512i { unsafe { _mm512_loadu_si512(p.as_ptr() as *const __m512i) } } + /// Times `body` only after it matched the exact reference once. + fn checked(expected: i64, name: &str, mut body: impl FnMut() -> i64) -> (f64, f64) { + assert_eq!(body(), expected, "{name} does not compute the exact SM9 dot"); + time(name, body) + } + #[inline(always)] unsafe fn hsum(v: __m512i) -> i64 { unsafe { _mm512_reduce_add_epi32(v) as i64 } @@ -138,6 +144,15 @@ mod probe { pub fn run() { let d = data(); + // Independent i64 reference for the exact SM9 kernels: every kernel checked + // against it must match before its time is reported. + let sm9: i64 = (0..CHUNKS) + .flat_map(|c| (0..64).map(move |i| (c, i))) + .map(|(c, i)| { + let v = i64::from(d.mag[c][i]); + (if d.pos[c] >> i & 1 == 1 { v } else { -v }) * i64::from(d.w[c][i]) + }) + .sum(); println!("== arithmetic (one accumulator chain per kernel, data in L1) =="); // SAFETY (all kernels): the build has avx512f/bw/vnni/vbmi compiled in // (module cfg); loads read whole 64-byte arrays. @@ -166,21 +181,23 @@ mod probe { } hsum(acc) }); - time("SM9 B: 2C limbs, dot(lo) - 256*dot(sign)", || unsafe { + checked(sm9, "SM9 B: 2C limbs, dot(lo) - 256*dot(sign)", || unsafe { let one = _mm512_set1_epi8(1); let mut acc = _mm512_setzero_si512(); for c in 0..CHUNKS { let w = ldi(&d.w[c]); - let neg = !d.pos[c]; - // two's-complement low byte of -m is (256 - m) & 0xFF - let lo = _mm512_mask_sub_epi8(ld(&d.mag[c]), neg, _mm512_setzero_si512(), ld(&d.mag[c])); + let m = ld(&d.mag[c]); + // Two's-complement -m is 256 * (-1) + ((256 - m) & 0xFF) only for m != 0: + // a negative zero has no borrow, so the sign limb skips it. + let neg = !d.pos[c] & _mm512_test_epi8_mask(m, m); + let lo = _mm512_mask_sub_epi8(m, neg, _mm512_setzero_si512(), m); let a = _mm512_dpbusd_epi32(acc, lo, w); let t = _mm512_dpbusd_epi32(_mm512_setzero_si512(), _mm512_maskz_mov_epi8(neg, one), w); acc = _mm512_sub_epi32(a, _mm512_slli_epi32(t, 8)); } hsum(acc) }); - time("SM9 C: dot(m) - 2*dot(m & neg), one acc", || unsafe { + checked(sm9, "SM9 C: dot(m) - 2*dot(m & neg), one acc", || unsafe { let mut acc = _mm512_setzero_si512(); for c in 0..CHUNKS { let (m, w) = (ld(&d.mag[c]), ldi(&d.w[c])); @@ -190,7 +207,7 @@ mod probe { } hsum(acc) }); - time("SM9 C': dot(m & pos) - dot(m & neg), two accs", || unsafe { + checked(sm9, "SM9 C': dot(m & pos) - dot(m & neg), two accs", || unsafe { let (mut p, mut n) = (_mm512_setzero_si512(), _mm512_setzero_si512()); for c in 0..CHUNKS { let (m, w) = (ld(&d.mag[c]), ldi(&d.w[c])); @@ -311,7 +328,7 @@ mod probe { } hsum(acc) }); - time("i9 values stored (no lookup) + SM9 C'", || unsafe { + checked(sm9, "i9 values stored (no lookup) + SM9 C'", || unsafe { let (mut p, mut n) = (_mm512_setzero_si512(), _mm512_setzero_si512()); for c in 0..CHUNKS { let (m, w) = (ld(&d.mag[c]), ldi(&d.w[c])); @@ -547,13 +564,22 @@ mod facade { black_box(sink); println!("{name:<44} {ns:7.3} ns/64 lanes"); }; + // Every lane feeds the result, so no backend can drop part of the work. + let all = |v: I32x16| v.to_array().iter().fold(0i32, |s, x| s.wrapping_add(*x)); + let sm9: i64 = (0..CHUNKS) + .flat_map(|c| (0..64).map(move |i| (c, i))) + .map(|(c, i)| { + let v = i64::from(mag[c].to_array()[i]); + (if pos[c] >> i & 1 == 1 { v } else { -v }) * i64::from(w[c].to_array()[i]) + }) + .sum(); println!("== facade (ndarray::simd, this build's tier) =="); time("I32x16::dot4_u8_i8_add", &|| { let mut acc = I32x16::splat(0); for c in 0..CHUNKS { acc = acc.dot4_u8_i8_add(mag[c], w[c]); } - acc.to_array()[0] + all(acc) }); time("two dot4_u8_i8_add accumulators, no mask", &|| { let (mut p, mut n) = (I32x16::splat(0), I32x16::splat(0)); @@ -561,16 +587,18 @@ mod facade { p = p.dot4_u8_i8_add(mag[c], w[c]); n = n.dot4_u8_i8_add(mag[(c + 1) % CHUNKS], w[c]); } - p.to_array()[0].wrapping_sub(n.to_array()[0]) + all(p).wrapping_sub(all(n)) }); - time("SM9: two dot4_u8_i8_add_masked accumulators", &|| { + let masked = || { let (mut p, mut n) = (I32x16::splat(0), I32x16::splat(0)); for c in 0..CHUNKS { p = p.dot4_u8_i8_add_masked(mag[c], pos[c], w[c]); n = n.dot4_u8_i8_add_masked(mag[c], !pos[c], w[c]); } - p.to_array()[0].wrapping_sub(n.to_array()[0]) - }); + all(p).wrapping_sub(all(n)) + }; + assert_eq!(i64::from(masked()), sm9, "facade SM9 does not compute the exact dot"); + time("SM9: two dot4_u8_i8_add_masked accumulators", &masked); } } diff --git a/src/simd_avx2.rs b/src/simd_avx2.rs index 7df23ac6..f6f8acb1 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -3031,20 +3031,31 @@ impl I32x16 { /// ``` #[inline(always)] pub fn dot4_u8_i8_add_masked(self, a: U8x64, keep: u64, b: I8x64) -> Self { - let spread = crate::simd_masking_ops::mask_byte_spread(keep); 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. 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; the rest is `dot4_u8_i8_add`. + // load and store is a 32-byte unaligned access inside a 64-byte array; the mask is + // built in registers. Qword `k` of half `h` takes mask byte `4h + k` (variable shift), + // copies it into its four low bytes (`* 0x0101_0101`, in range because the byte is + // below 256) and its four high bytes (shift by 32); testing bit `i % 8` against + // `BIT_IN_BYTE` then gives 0xFF exactly where bit `i` of `keep` is set. The rest is + // `dot4_u8_i8_add`. unsafe { let low = _mm256_set1_epi16(0x00FF); let bit = _mm256_set1_epi64x(crate::simd_masking_ops::BIT_IN_BYTE as i64); - let ps = spread.as_ptr() as *const __m256i; + let keep_v = _mm256_set1_epi64x(keep as i64); + let byte = _mm256_set1_epi64x(0xFF); + let rep = _mm256_set1_epi64x(0x0101_0101); 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 m = _mm256_cmpeq_epi8(_mm256_and_si256(_mm256_loadu_si256(ps.add(h)), bit), bit); + 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 av = _mm256_and_si256(_mm256_loadu_si256(pa.add(h)), m); let bv = _mm256_loadu_si256(pb.add(h)); let b_even = _mm256_srai_epi16(_mm256_slli_epi16(bv, 8), 8); From 0a23a54f5b31a89614630ddccc3d28090ca4933f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 16:02:15 +0000 Subject: [PATCH 6/6] Correct the 9-bit measurements after the probe fixes Updates the handoff tables with the checked, all-lane measurements: the AVX2 masked pair costs 6.35-6.44 ns per 64 lanes, about 4x one dot4, not the 4.5 ns the lane-0 timing reported. Silences dead_code on the stored mask spread for the tiers that build the mask another way. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- .claude/blackboard.md | 4 +- .claude/knowledge/turbovec9-u9-substrate.md | 65 ++++++++++++--------- src/simd_masking_ops.rs | 2 + 3 files changed, 42 insertions(+), 29 deletions(-) diff --git a/.claude/blackboard.md b/.claude/blackboard.md index 920e86a5..4c67fc2f 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3892,7 +3892,9 @@ reversed BIT_IN_BYTE (AVX2), mask dropped in the fused AVX2 loop. Key numbers (AVX-512 VNNI, L1): masked pair 1.27-1.33 ns/64 lanes vs dot4 1.10-1.19; 512-entry lookup 6.1 ns (vpermi2b) / 38-40 ns (per-position reads). Stored uniform SM9 46.0 dB vs i4 20.2 dB at the same per-dim L1 cost and -1.4x per dim streamed. AVX2: masked pair 4.5 ns vs dot4 1.5 ns. +1.4x per dim streamed. AVX2: masked pair 6.35-6.44 ns vs dot4 1.59-1.79 ns (corrected after review: the first +facade timing returned lane 0 only and let the optimizer drop half the AVX2 work; probe kernel B also mishandled a +negative zero; both fixed, every exact kernel now checked against an i64 reference before it is timed). 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. diff --git a/.claude/knowledge/turbovec9-u9-substrate.md b/.claude/knowledge/turbovec9-u9-substrate.md index f1531264..b1a0a6b1 100644 --- a/.claude/knowledge/turbovec9-u9-substrate.md +++ b/.claude/knowledge/turbovec9-u9-substrate.md @@ -5,7 +5,9 @@ READ BY: savant-architect, vector-synthesis, simd-savant, anyone wiring a Status: probe + one facade word (`I32x16::dot4_u8_i8_add_masked`). Numbers below come from `examples/u9_cost_probe.rs` on one AVX-512 VNNI + VBMI Xeon -(target-cpu=native), three runs, reported as ranges. They are not TurboVec +(target-cpu=native), three runs, reported as ranges. Every exact SM9 kernel +must match an independent i64 reference before its time is printed, and +every facade kernel folds all 16 accumulator lanes into its result. They are not TurboVec recall numbers; the codebooks in the probe are probe-local Lloyd-Max fits. ## Two contracts, never mixed @@ -46,9 +48,11 @@ call, so 16_448 calls (65_792 dims per lane) cannot wrap; the test No saturating instruction is involved (no `vpmaddubsw`, no `vpdpbusds`). Realizations: AVX-512 one `vmovdqu8{z}` before `vpdpbusd` (or the widen + -`vpmaddwd` path without VNNI); AVX2, AVX, NEON and wasm expand the mask -with one AND against `BIT_IN_BYTE` and one byte compare per register (no -shuffle, no per-lane loop); nightly selects through `core::simd::Mask`; +`vpmaddwd` path without VNNI); AVX2 builds the byte mask in registers +(variable shift, multiply by `0x0101_0101`, shift by 32) and applies it +inside the dot loop; AVX, NEON and wasm expand a stored spread with one AND +against `BIT_IN_BYTE` and one byte compare per register (no shuffle, no +per-lane loop); nightly selects through `core::simd::Mask`; scalar is the reference loop. Witnessed: `simd::dot4_nibble_facade_parity` 8/8 on native (AVX-512 VNNI), @@ -56,29 +60,34 @@ config-v4 (AVX-512 without VNNI), config-v3 (AVX2), config-avx (qemu SandyBridge), config-v2 (scalar, qemu Nehalem), aarch64 (qemu) and nightly-simd; `neon-parity.sh` and `wasm-parity.sh` pass (codes 0x504, 0x505). Disable runs, all red: unmasked AVX-512 move (3 tests), reversed -`BIT_IN_BYTE` on AVX2 (2 tests), mask dropped in the fused AVX2 loop (3). +`BIT_IN_BYTE` on AVX2 (2 tests), mask dropped in the fused AVX2 loop (3), +wrong byte shifts for the upper AVX2 half (2). ## Arithmetic: three exact forms compared (AVX-512 VNNI, L1, ns per 64 lanes) | kernel | ns / 64 lanes | |---|---| -| u8 x i8 `vpdpbusd` (#386 `dot4_u8_i8_add`) | 1.10 - 1.19 | -| i8 x i8 via `xor 0x80` bias | 1.20 - 1.23 | -| A: sign folded into w (exact only for w in -127..=127) | 1.24 - 1.50 | -| B: two's-complement limbs, `dot(lo) - 256 * dot(sign)` | 1.88 - 1.89 | -| C: `dot(m) - 2 * dot(m & neg)`, one accumulator | 1.79 - 1.84 | -| **C': `dot(m & pos) - dot(m & neg)`, two accumulators** | **1.27 - 1.33** | +| u8 x i8 `vpdpbusd` (#386 `dot4_u8_i8_add`) | 1.12 - 1.14 | +| i8 x i8 via `xor 0x80` bias | 1.20 - 1.25 | +| A: sign folded into w (exact only for w in -127..=127) | 1.23 - 1.26 | +| B: two's-complement limbs, `dot(lo) - 256 * dot(sign)`, sign limb skips magnitude 0 | 1.89 - 1.92 | +| C: `dot(m) - 2 * dot(m & neg)`, one accumulator | 1.79 - 1.88 | +| **C': `dot(m & pos) - dot(m & neg)`, two accumulators** | **1.28 - 1.30** | The loop is one latency-bound chain. B and C put two dependent ops on the accumulator; C' keeps two independent chains and lands within ~10% of the plain u8 dot. A is not exact for w = -128 and was not added. -Through the facade, same host: `dot4_u8_i8_add` 1.12, two masked -accumulators 1.29. On config-v3 (AVX2, no VNNI): `dot4_u8_i8_add` 1.51, -two plain accumulators 3.26, two masked accumulators 4.48 - 4.58. Without -VNNI every call widens to 16 bits, so a pair costs two widenings; a fused -SM9 op that widens once and negates in 16 bits would roughly halve that, -and is not built. +Through the facade, same host: `dot4_u8_i8_add` 1.08 - 1.15, two masked +accumulators 1.34 - 1.38. On config-v3 (AVX2, no VNNI): `dot4_u8_i8_add` +1.59 - 1.79, two plain accumulators 2.75 - 2.91, two masked accumulators +6.35 - 6.44. Without VNNI each masked call adds the mask expansion (about +1.4 ns, measured as one masked accumulator against one plain one) on top of +its own 16-bit widening, so the pair costs about 4x one dot4. The loop is +instruction-bound, not memory-bound: building the mask in registers instead +of from a stored spread did not change the time. A fused SM9 op that widens +once and negates in 16 bits would avoid both the second widening and the +mask, and is not built. ## 512-entry lookup (WP2): keep it out of the scan @@ -86,9 +95,9 @@ ID = low byte + high-bit plane, decoded to magnitude + positive plane: | lookup | ns / 64 lanes | |---|---| -| per-position reads (`lut_u8` shape, two 512-entry tables) | 38 - 40 | -| 4x `vpermi2b` + blends for magnitude, same for the sign | 6.07 - 6.16 | -| 2x `vpermi2b` on an antisymmetric codebook (`255 - lo` for negatives) | 6.79 - 6.94 | +| per-position reads (`lut_u8` shape, two 512-entry tables) | 38 - 39 | +| 4x `vpermi2b` + blends for magnitude, same for the sign | 6.03 - 6.09 | +| 2x `vpermi2b` on an antisymmetric codebook (`255 - lo` for negatives) | 6.72 - 6.77 | `lut9` and `lut_u8` keep their contracts; no 512-entry facade word was added. The lookup alone costs ~5x the arithmetic it feeds. TurboVec's @@ -99,15 +108,15 @@ time; then the scan stores values, not IDs (next section). | layout | SNR | scan cost (AVX-512 VNNI) | |---|---|---| -| i4, 16 Lloyd-Max levels (today) | 20.22 dB | 2.64 - 2.70 ns / 128 dims (L1) | -| i8 uniform grid, best clip | 40.54 dB | 1.10 - 1.19 ns / 64 dims | -| 512 Lloyd-Max IDs, executed by `vpermi2b` lookup | 49.86 dB | 5.39 - 5.53 ns / 64 dims | -| 512 Lloyd-Max levels snapped to the SM9 grid | 44.74 dB (329 distinct) | 1.29 - 1.31 ns / 64 dims | -| **SM9 uniform grid, best clip (4.21 sigma), stored** | **46.02 dB** | **1.29 - 1.31 ns / 64 dims** | -| 4+5 additive C16 + R32, two i8 grids | 43.25 dB | 2.58 - 2.65 ns / 64 dims | +| i4, 16 Lloyd-Max levels (today) | 20.22 dB | 2.65 - 2.68 ns / 128 dims (L1) | +| i8 uniform grid, best clip | 40.54 dB | 1.12 - 1.14 ns / 64 dims | +| 512 Lloyd-Max IDs, executed by `vpermi2b` lookup | 49.86 dB | 5.37 - 5.70 ns / 64 dims | +| 512 Lloyd-Max levels snapped to the SM9 grid | 44.74 dB (329 distinct) | 1.29 - 1.32 ns / 64 dims | +| **SM9 uniform grid, best clip (4.21 sigma), stored** | **46.02 dB** | **1.29 - 1.32 ns / 64 dims** | +| 4+5 additive C16 + R32, two i8 grids | 43.25 dB | 2.61 - 2.65 ns / 64 dims | Streamed from DRAM (64 MiB of codes, weights in L1), per dimension: -i4 0.031 - 0.035 ns, stored SM9 0.043 - 0.049 ns (72 bytes per 64 dims). +i4 0.035 - 0.037 ns, stored SM9 0.048 - 0.052 ns (72 bytes per 64 dims). Reading the table: @@ -127,7 +136,7 @@ On AVX-512 VNNI: yes, for the stored-value layout. 64 dims cost one masked-pair of `vpdpbusd`, ~1.3 ns in L1 (about the same per dimension as today's i4 permute-dot) and 1.4x i4 per dimension when streamed from DRAM, for +25.8 dB over i4. With IDs plus an in-scan 512 lookup: no, the lookup -dominates. On AVX2 the pair costs ~3x the plain dot; a fused SM9 op is the +dominates. On AVX2 the pair costs ~4x one plain dot4; a fused SM9 op is the next step there if AVX2 matters. ## Handoff to the TurboVec owner (owns ABI, encode, pack, search, mutation, serialization, Python) diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index 9bf0a1ca..a880a06d 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -5305,6 +5305,8 @@ pub(crate) const BIT_IN_BYTE: u64 = 0x8040_2010_0804_0201; /// bit `i` of `keep` at bit position `i % 8`, so one AND with /// [`BIT_IN_BYTE`] and one byte compare per register expand the mask without /// moving a byte. +// Used by the AVX, NEON and wasm realizations only; other tiers build the mask differently. +#[allow(dead_code)] #[inline(always)] pub(crate) fn mask_byte_spread(keep: u64) -> [u64; 8] { core::array::from_fn(|j| (keep >> (8 * j) & 0xFF) * 0x0101_0101_0101_0101)