diff --git a/.claude/blackboard.md b/.claude/blackboard.md index 44321e56..4c67fc2f 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3876,3 +3876,25 @@ 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 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/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/.claude/knowledge/turbovec9-u9-substrate.md b/.claude/knowledge/turbovec9-u9-substrate.md new file mode 100644 index 00000000..b1a0a6b1 --- /dev/null +++ b/.claude/knowledge/turbovec9-u9-substrate.md @@ -0,0 +1,159 @@ +# 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. 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 + +- **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 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), +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), +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.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.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 + +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 - 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 +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.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.035 - 0.037 ns, stored SM9 0.048 - 0.052 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 ~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) + +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. 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..680f8574 --- /dev/null +++ b/examples/u9_cost_probe.rs @@ -0,0 +1,612 @@ +//! 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) } + } + /// 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 } + } + + /// 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(); + // 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. + 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) + }); + 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 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) + }); + 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])); + 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) + }); + 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])); + 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) + }); + 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])); + 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)) + ); + } +} + +/// 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"); + }; + // 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]); + } + all(acc) + }); + 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]); + } + all(p).wrapping_sub(all(n)) + }); + 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]); + } + 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); + } +} + +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")))] + 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..f6f8acb1 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -3007,6 +3007,66 @@ 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 mut out = self.0; + // SAFETY: AVX2 is a run-time precondition of this arm (see `U64x8::avx2_halves`); every + // load and store is a 32-byte unaligned access inside a 64-byte array; the mask is + // built in registers. 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 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 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); + 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"))))] 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..a880a06d 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -5295,6 +5295,23 @@ 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. +// 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) +} + /// `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..24a9b688 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, Select, 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`. ///