diff --git a/.claude/blackboard.md b/.claude/blackboard.md index e9625644..b194b08e 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -1,3 +1,111 @@ +## 2026-09-16 (13) — the `VPTERNLOGQ` tail is a DESCENT, not a pad (5–8×); a 64×2 re-apply on a full-width mask is NOT (0.5–0.7×); the GEMM block-stop tail is INERT (0.99–1.02×) + +Three probes, one question in three places (operator: *"instead of padding the +tail you could simply split the N×64×8 + N×64×2 tail"*, then *"would the gather +vs re-apply be faster with 64×2 instead of 64×8"*, then *"would the MKL GEMM +stop logic profit from a tail optimization"*). All AVX-512 (`.cargo/config-v4.toml`, +`avx512f=true` printed by each program), release, `black_box` on inputs AND +outputs, every arm gated bit-identical before timing. Branch +`claude/c64-6502-falsifier-shztkk`, PR #311. + +### 1. `examples/ternlogq_tail_descent_probe.rs` — the TAIL. YES. + +`mask_ternlog` chunks over `U64x8` (512 rows) and ends on three `pad_tail`s +(zero-fill three `[u64; 8]`, one zmm op, prefix copy). The alternative: descend +zmm → ymm → xmm (`4 + 2 + 1`, every lane live, all in vector registers). + +Crux — for a remainder of `t` words, tail only, ns/call, 3 runs: + +| t | shape | P padded zmm | **G greedy** | X all-xmm | winner | +|---:|---|---:|---:|---:|---| +| 1 | `1` | 13.1–13.8 | 1.78–1.90 | 1.49–1.60 | G = X (identical code) | +| 2 | `2` | 13.0–13.7 | **1.74–1.94** | 2.57–2.78 | G | +| 3 | `2+1` | 17.1–18.6 | **2.13–2.36** | 2.82–3.23 | G | +| 4 | `4` vs `2+2` | 12.8–13.6 | **1.59–1.76** | 2.92–3.00 | G | +| 5 | `4+1` vs `2+2+1` | 17.3–18.3 | **2.08–2.13** | 3.23–3.49 | G | +| **6** | **`4+2`** vs `2+2+2` | 17.0–18.3 | **2.26–2.74** | 3.29–3.68 | **G** | +| 7 | `4+2+1` vs `2+2+2+1` | 16.8–18.5 | **2.38–2.57** | 3.68–3.95 | G | + +Greedy widest-first wins every `t ≥ 2`: 1.3–1.6× over all-xmm, 5–8× over +padding. The operator's crux (`t = 6`): one ymm + one xmm beats three xmm. + +End to end through the real chunk loop: + +| words | rows | tail | padded ns | descend ns | ratio | +|---:|---:|---:|---:|---:|---:| +| 3 (`ogar-r2il` `CallMask`) | 192 | 3 | 18.1 | 2.3 | **7.7–8.0×** | +| 1–7 | 64–448 | 1–7 | 13–23 | 2.0–2.6 | 5.9–8.0× | +| 9 | 576 | 1 | 14.2 | 3.7 | 4.2× | +| 11 | 704 | 3 | 21.9 | 3.6 | 5.4–5.6× | +| 31 | 1 984 | 7 | 23.6 | 6.8 | 2.5–3.3× | +| 194 | 12 416 | 2 | 41.4 | 26.5 | 1.3–1.4× | +| 8 / 16 / 24 / 64 | — | none | — | — | 1.0–1.4× (loop shape, not tail) | + +asm: 33 zmm + 3 ymm + 6 xmm `vpternlogq`, folded memory operands, **zero** GPR +and/or/xor on lane data — a descent is not the scalar peel +`scripts/codegen-witness.sh` caps at `SLICE_GPR_CAP=6`. Follow-up named, not +built: `U64x4::ternlog` / `U64x2::ternlog` on the facade + rewire +`mask_ternlog`'s tail; an un-gated `pack` sibling of `pack_under` to +retire the 12 hand-rolled `if !tail.is_empty()` sites. + +### 2. `examples/ternlogq_sparse_reapply_probe.rs` — FULL-WIDTH sparse frontier. NO. + +1 024 words (65 536 rows, the MQ / `lgj_hop` population), `dst = src ∧ gate ∧ +elig`, arms: F8 full zmm pass (shipped), S8/S4/S2 chunk-skip at zmm/ymm/xmm +(`vptestmq` → kortest), S1 GPR floor, W per-bit gather walk. ns/call, 3 runs: + +| frontier | shape | live | F8 | S8 | S4 | S2 | S1 | W | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| 0.01 % | uniform | 7 | 143–164 | **121–141** | 148–168 | 239–277 | 192–233 | 404–435 | +| 0.01 % | clustered | 7 | 148–163 | **125–136** | 144–145 | 240–242 | 223 | 358–368 | +| 0.1 % | clustered | 66 | 127–161 | **103–132** | 137–145 | 227–240 | 198–223 | 394–421 | +| 1 % | uniform | 655 | **153–162** | 190–197 | 216–222 | 277–289 | 198–224 | 918–1 013 | +| 1 % | clustered | 655 | 152–162 | **123–132** | 129–146 | 190–240 | 223–227 | 933–981 | +| 10 % | uniform | 6 554 | **153–162** | 188–200 | 213–222 | 314–334 | 216–224 | 9 987–10 849 | +| 10 % | clustered | 6 554 | 152–161 | **135–141** | 144–157 | 239–258 | 193–222 | 6 187–6 926 | +| 100 % | either | 65 536 | **145–161** | 172–205 | 194–222 | 313–339 | 201–229 | 52–58 k | + +S2 is 1.5–2.1× SLOWER than the full pass everywhere; S4 never beats S8; the +skip is worth ≤ 1.24× on clustered frontiers and LOSES on uniform ≥ 1 % +(branch mispredicts). 150 ns for 32 KiB of traffic is L1 bandwidth; narrower +chunks are more iterations, not less work. The gather walks all 1 024 words +before knowing they are empty (≥ 360 ns at 7 bits, ~0.8 ns/bit after). The +64×2 rung is a tail instrument only. + +### 3. `kernels_avx512.rs::block_stop_probe` (ignored test) — the GEMM stop. INERT. + +`sgemm_blocked`'s M-stop pads the last `MR=6` panel and the ukernel computes all +six accumulators regardless of `mr_eff`; every power-of-two `m` has such a tail +(128 = 21·6+2, 256 = 42·6+4, 512 = 85·6+2, 1024 = 170·6+4). A test-local `R×16` +tail ukernel (R ∈ {2, 4}) on the tail tile only, bit-identical, best of 9: + +| m×n×k | tail | shipped ms | desc ms | ratio | FMA waste | +|---|---:|---:|---:|---:|---:| +| 126×256×256 | 0 | 0.183 | 0.185 | 0.987× | 0.0 % | +| 128×256×256 | 2 | 0.190 | 0.192 | 0.989× | 3.1 % | +| 130×256×256 | 4 | 0.190 | 0.191 | 0.992× | 1.5 % | +| 132×256×256 | 0 | 0.193 | 0.190 | 1.012× | 0.0 % | +| 128³ | 2 | 0.056 | 0.055 | 1.015× | 3.1 % | +| 256³ | 4 | 0.347 | 0.350 | 0.992× | 0.8 % | +| 512³ | 2 | 3.110 | 3.054 | 1.018× | 0.8 % | +| 1024³ | 4 | 32.57 | 32.04 | 1.016× | 0.2 % | + +Noise. ~66 GFLOP/s at 1024³ (half of one core's FMA peak): packing and memory +traffic hide the padded rows. K-stop has no waste; N-stop padding is on lanes +the FMA unit processes anyway — no lane-width descent applies to GEMM. The BF16 +`vdpbf16ps` path's stop problem is ONE accumulator chain per row +(`amx_matmul.rs:686-697`, latency-bound), not its tails. + +### The rule the three share + +The tail descent paid where the pad cost **loads and a copy** (3 zero-fills + +prefix copy per call). It buys nothing where the padding is **ALU on lanes the +unit processes anyway** (GEMM N-stop, sparse re-apply, GEMM M-stop hidden +behind bandwidth). Sibling finding the same day, lance-graph #1241: the facet's +per-axis LCP was gathering + re-folding per call; reading the single `u128` +register masked to the axis bytes (the `-f` done ONCE at mint) took it +12.5 → 5.8 ns for both axes. + ## 2026-09-16 (12) — ⊘ the G1 ratio was INFLATED by dead-store elimination; corrected 6.75× → 5.84×, and the conclusion survives codex P2 on PR #309, and it was load-bearing. Entry (11)'s numbers are diff --git a/Cargo.toml b/Cargo.toml index cd3f8bbd..c525bab3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,6 +63,14 @@ required-features = ["std"] name = "r2il_column_scan_probe" required-features = ["std"] +[[example]] +name = "ternlogq_tail_descent_probe" +required-features = ["std"] + +[[example]] +name = "ternlogq_sparse_reapply_probe" +required-features = ["std"] + [[example]] name = "hex_tenant_mq_probe" required-features = ["std"] diff --git a/examples/ternlogq_sparse_reapply_probe.rs b/examples/ternlogq_sparse_reapply_probe.rs new file mode 100644 index 00000000..31606892 --- /dev/null +++ b/examples/ternlogq_sparse_reapply_probe.rs @@ -0,0 +1,351 @@ +//! Sparse frontier: re-apply the `VPTERNLOGQ` chain per chunk, or walk the bits? +//! +//! Follow-up to `ternlogq_tail_descent_probe` (the TAIL) for the other place a +//! 64×2 rung could matter: a **sparse frontier over a full-width mask**. The +//! MQ cost model (`hex_tenant_mq_probe`, plan §14) is `step = x·ternlogq + n`: +//! ternlogq = one full pass over every word (291 ns at 1 024 words), `n` = the +//! per-active-bit walk (17.3 µs). `lgj_hop` has the same two halves — three +//! full mask passes, then a `trailing_zeros` scatter over the selected set. +//! +//! The question (operator, 2026-09-16): *"would the gather vs re-apply be +//! faster with 64×2 instead of 64×8?"* At 1 024 words there is no tail, so the +//! descent can only pay if the re-apply is **gated per chunk on a non-empty +//! source** — then the chunk width sets how many dead rows each live bit drags +//! through the ALU: 512 (zmm), 256 (ymm), 128 (xmm), 64 (a GPR word). Against +//! that stands the per-bit gather, O(frontier) and blind to the mask geometry. +//! +//! Arms, all computing `dst = src ∧ gate ∧ elig` (AND3) and gated bit-identical: +//! +//! | arm | shape | +//! |---|---| +//! | `F8` | the shipped full pass, `mask_ternlog` over `U64x8` chunks | +//! | `S8` | zmm chunks, skip when the 8-word src chunk is all zero (`vptestmq` → kortest) | +//! | `S4` | ymm chunks, same skip at 4 words | +//! | `S2` | xmm chunks, same skip at 2 words | +//! | `S1` | one word: `if s != 0 { s & g & e }` — GPR, the shape `codegen-witness.sh` caps; here a FLOOR, not a candidate | +//! | `W` | the gather: `trailing_zeros` walk of src, per bit one 8-byte payload read + one dst bit — O(frontier) | +//! +//! Two frontier shapes at each density, because a BFS frontier on a Morton +//! lattice is CLUSTERED, not uniform: `uniform` scatters the live bits over +//! all 1 024 words; `clustered` packs them into a contiguous run of words. +//! +//! # Measured 2026-09-16 — Xeon @ 2.10 GHz, `avx512f=true avx512vl=true`, release, 3 runs +//! +//! ns per call, 1 024 words (65 536 rows), operands ~62 % dense, best of 3 × 30 ms. +//! +//! | frontier | shape | live | F8 | S8 | S4 | S2 | S1 (gpr) | W (gather) | +//! |---|---|---:|---:|---:|---:|---:|---:|---:| +//! | 0.01 % | uniform | 7 | 143–164 | **121–141** | 148–168 | 239–277 | 192–233 | 404–435 | +//! | 0.01 % | clustered | 7 | 148–163 | **125–136** | 144–145 | 240–242 | 223 | 358–368 | +//! | 0.1 % | clustered | 66 | 127–161 | **103–132** | 137–145 | 227–240 | 198–223 | 394–421 | +//! | 1 % | uniform | 655 | **153–162** | 190–197 | 216–222 | 277–289 | 198–224 | 918–1013 | +//! | 1 % | clustered | 655 | 152–162 | **123–132** | 129–146 | 190–240 | 223–227 | 933–981 | +//! | 10 % | uniform | 6 554 | **153–162** | 188–200 | 213–222 | 314–334 | 216–224 | 9 987–10 849 | +//! | 10 % | clustered | 6 554 | 152–161 | **135–141** | 144–157 | 239–258 | 193–222 | 6 187–6 926 | +//! | 100 % | either | 65 536 | **145–161** | 172–205 | 194–222 | 313–339 | 201–229 | 52 351–58 329 | +//! +//! **Answer: no — 64×2 is not the faster re-apply at any density.** `S2` is +//! 1.5–2.1× SLOWER than the full zmm pass everywhere, `S4` never beats `S8`, +//! and the skip itself is worth at most **1.24×** (clustered frontiers; on a +//! uniform frontier at ≥ 1 % the branch mispredicts and `S8` LOSES to `F8`). +//! The full pass is 150 ns for 24 KiB of reads + 8 KiB of writes — L1-resident, +//! already at the bandwidth the chunk shape cannot improve on. Chunk width sets +//! wasted rows per live bit only in principle; at this population the pass is +//! bandwidth, not ALU, so narrower chunks just mean more loop iterations. +//! +//! **The gather never wins here, and the reason is its own loop.** `W` walks +//! all 1 024 source words before it can know they are empty, so it costs ≥ 360 ns +//! even at 7 live bits (2.6–3.5× the mask pass) and grows ~0.8 ns per live bit +//! (payload L1/L2-resident; lgj's 512-byte-strided rows would be ~25 ns/bit, +//! which is the number the operator quoted). A gather that walked an INDEX LIST +//! would be O(bits) and would win below ~0.2 % — and that list is exactly the +//! serialization of the population the mask doctrine forbids (lgj R1); it is not +//! a candidate, and the mask pass is within 1.5× of what it would cost at the +//! sparsest arm anyway. +//! +//! So for a full-width frontier the 64×2 rung is a TAIL instrument only +//! (`ternlogq_tail_descent_probe`: 5–8× on 1..7-word masks); on a 1 024-word +//! mask it is the wrong tool, and the one lever that shows is chunk-skip on a +//! clustered frontier, worth ≤ 1.24×. +//! +//! AVX-512 only (v4 config), gated on `avx512f` AND `avx512vl` AND `avx512dq` +//! (Codex P2 on #311): `S4`/`S2` are VL encodings and `vptestmq` is DQ; a +//! `#[target_feature]` attribute is a caller precondition, not a CPU check, so +//! an F-only target takes the no-op `main`. +//! +//! ```text +//! env -u RUSTFLAGS cargo --config .cargo/config-v4.toml run --release \ +//! --example ternlogq_sparse_reapply_probe +//! ``` + +#[cfg(all( + target_arch = "x86_64", + target_feature = "avx512f", + target_feature = "avx512vl", + target_feature = "avx512dq" +))] +mod probe { + use ndarray::simd::mask_ternlog; + use ndarray::simd::ternlog::AND3; + use std::arch::x86_64::*; + use std::hint::black_box; + use std::time::Instant; + + pub const ROWS: usize = 65_536; + pub const WORDS: usize = ROWS / 64; + + fn splitmix(s: &mut u64) -> u64 { + *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = *s; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + /// `F8` — what ships. + fn f8(src: &[u64], gate: &[u64], elig: &[u64], dst: &mut [u64]) { + mask_ternlog::(src, gate, elig, dst); + } + + /// `S8` — zmm chunks, skipped when the src chunk is all zero. + #[target_feature(enable = "avx512f,avx512vl")] + unsafe fn s8(src: &[u64], gate: &[u64], elig: &[u64], dst: &mut [u64]) { + let n = src.len(); + let mut i = 0; + // SAFETY: caller guarantees avx512f+avx512vl+avx512dq and equal-length + // slices; the loop runs only while `i + 8 <= n`, so every zmm + // load/store of 8 words at `i` is inside `src`/`gate`/`elig`/`dst`. + unsafe { + while i + 8 <= n { + let s = _mm512_loadu_si512(src.as_ptr().add(i).cast()); + if _mm512_test_epi64_mask(s, s) == 0 { + _mm512_storeu_si512(dst.as_mut_ptr().add(i).cast(), _mm512_setzero_si512()); + } else { + let v = _mm512_ternarylogic_epi64::( + s, + _mm512_loadu_si512(gate.as_ptr().add(i).cast()), + _mm512_loadu_si512(elig.as_ptr().add(i).cast()), + ); + _mm512_storeu_si512(dst.as_mut_ptr().add(i).cast(), v); + } + i += 8; + } + } + debug_assert_eq!(i, n, "probe sizes are multiples of 8 words"); + } + + /// `S4` — ymm chunks, skipped when the 4-word src chunk is all zero. + #[target_feature(enable = "avx512f,avx512vl")] + unsafe fn s4(src: &[u64], gate: &[u64], elig: &[u64], dst: &mut [u64]) { + let n = src.len(); + let mut i = 0; + // SAFETY: caller guarantees avx512f+avx512vl+avx512dq and equal-length + // slices; the loop runs only while `i + 4 <= n`, so every ymm + // load/store of 4 words at `i` is inside `src`/`gate`/`elig`/`dst`. + unsafe { + while i + 4 <= n { + let s = _mm256_loadu_si256(src.as_ptr().add(i).cast()); + if _mm256_test_epi64_mask(s, s) == 0 { + _mm256_storeu_si256(dst.as_mut_ptr().add(i).cast(), _mm256_setzero_si256()); + } else { + let v = _mm256_ternarylogic_epi64::( + s, + _mm256_loadu_si256(gate.as_ptr().add(i).cast()), + _mm256_loadu_si256(elig.as_ptr().add(i).cast()), + ); + _mm256_storeu_si256(dst.as_mut_ptr().add(i).cast(), v); + } + i += 4; + } + } + } + + /// `S2` — xmm chunks, skipped when the 2-word src chunk is all zero. + #[target_feature(enable = "avx512f,avx512vl")] + unsafe fn s2(src: &[u64], gate: &[u64], elig: &[u64], dst: &mut [u64]) { + let n = src.len(); + let mut i = 0; + // SAFETY: caller guarantees avx512f+avx512vl+avx512dq and equal-length + // slices; the loop runs only while `i + 2 <= n`, so every xmm + // load/store of 2 words at `i` is inside `src`/`gate`/`elig`/`dst`. + unsafe { + while i + 2 <= n { + let s = _mm_loadu_si128(src.as_ptr().add(i).cast()); + if _mm_test_epi64_mask(s, s) == 0 { + _mm_storeu_si128(dst.as_mut_ptr().add(i).cast(), _mm_setzero_si128()); + } else { + let v = _mm_ternarylogic_epi64::( + s, + _mm_loadu_si128(gate.as_ptr().add(i).cast()), + _mm_loadu_si128(elig.as_ptr().add(i).cast()), + ); + _mm_storeu_si128(dst.as_mut_ptr().add(i).cast(), v); + } + i += 2; + } + } + } + + /// `S1` — the GPR floor. Not a candidate (fails the codegen witness); it + /// bounds what "skip at the finest granularity" can buy. + fn s1(src: &[u64], gate: &[u64], elig: &[u64], dst: &mut [u64]) { + for (((&s, &g), &e), d) in src.iter().zip(gate).zip(elig).zip(dst.iter_mut()) { + *d = if s == 0 { 0 } else { s & g & e }; + } + } + + /// `W` — the gather. Walk src's set bits; per bit read one 8-byte payload + /// (the hop's decode read, from a 65 536 × u64 column) and decide from it. + /// `payload[row] & 1` is prearranged to equal `gate ∧ elig` at that row, so + /// the walk computes the SAME predicate and gates bit-identical. + fn w(src: &[u64], payload: &[u64], dst: &mut [u64]) { + for (wi, (&sw, d)) in src.iter().zip(dst.iter_mut()).enumerate() { + let mut bits = sw; + let mut out = 0u64; + while bits != 0 { + let b = bits.trailing_zeros(); + bits &= bits - 1; + let row = wi * 64 + b as usize; + out |= (payload[row] & 1) << b; + } + *d = out; + } + } + + /// Time `f` to a 30 ms floor, best of 3 rounds; ns per call. + fn time(mut f: impl FnMut()) -> f64 { + let mut best = f64::MAX; + for _ in 0..3 { + let mut reps = 0u64; + let t0 = Instant::now(); + while t0.elapsed().as_millis() < 30 { + f(); + reps += 1; + } + best = best.min(t0.elapsed().as_nanos() as f64 / reps as f64); + } + best + } + + fn frontier(seed: &mut u64, live_bits: usize, clustered: bool) -> Vec { + let mut m = vec![0u64; WORDS]; + let mut placed = 0; + // clustered: live bits confined to the first `span` words, where span is + // the smallest run that fits them at ~50% fill; uniform: anywhere. + let span = if clustered { + ((live_bits * 2) / 64).clamp(1, WORDS) + } else { + WORDS + }; + while placed < live_bits { + let r = splitmix(seed) as usize; + let wi = r % span; + let b = (r >> 20) % 64; + if m[wi] >> b & 1 == 0 { + m[wi] |= 1 << b; + placed += 1; + } + } + m + } + + pub fn main() { + println!( + "realization: avx512f={} avx512vl={} rows={ROWS} words={WORDS}", + cfg!(target_feature = "avx512f"), + cfg!(target_feature = "avx512vl") + ); + let mut seed = 0x1234_5678_9ABC_DEF0u64; + // Operands at ~62% density each (the MQ fixture's permeability). + let dense = |seed: &mut u64| -> Vec { + (0..WORDS) + .map(|_| splitmix(seed) | splitmix(seed) & splitmix(seed)) + .collect() + }; + let gate = dense(&mut seed); + let elig = dense(&mut seed); + // payload column: bit 0 = gate ∧ elig at that row, upper bits noise. + let payload: Vec = (0..ROWS) + .map(|r| (splitmix(&mut seed) & !1) | ((gate[r / 64] & elig[r / 64]) >> (r % 64) & 1)) + .collect(); + + println!( + "\n{:<9} {:<10} {:>8} | {:>8} {:>8} {:>8} {:>8} {:>8} {:>9} | {:<12}", + "frontier", "shape", "live", "F8", "S8", "S4", "S2", "S1(gpr)", "W(gather)", "best vector" + ); + for &pct in &[0.01f64, 0.1, 1.0, 10.0, 100.0] { + let live = ((ROWS as f64) * pct / 100.0).round().max(1.0) as usize; + for clustered in [false, true] { + let src = frontier(&mut seed, live, clustered); + let mut d_ref = vec![0u64; WORDS]; + f8(&src, &gate, &elig, &mut d_ref); + let mut d = vec![0u64; WORDS]; + // equivalence gate, every arm, before any timing + // SAFETY (this and the timed calls below): the active `main` is + // cfg-gated on avx512f+avx512vl+avx512dq, and all four buffers + // are `WORDS` long. + unsafe { s8(&src, &gate, &elig, &mut d) }; + assert_eq!(d, d_ref, "S8"); + unsafe { s4(&src, &gate, &elig, &mut d) }; + assert_eq!(d, d_ref, "S4"); + unsafe { s2(&src, &gate, &elig, &mut d) }; + assert_eq!(d, d_ref, "S2"); + s1(&src, &gate, &elig, &mut d); + assert_eq!(d, d_ref, "S1"); + w(&src, &payload, &mut d); + assert_eq!(d, d_ref, "W"); + + let t_f8 = time(|| f8(black_box(&src), black_box(&gate), black_box(&elig), black_box(&mut d))); + let t_s8 = + time(|| unsafe { s8(black_box(&src), black_box(&gate), black_box(&elig), black_box(&mut d)) }); + let t_s4 = + time(|| unsafe { s4(black_box(&src), black_box(&gate), black_box(&elig), black_box(&mut d)) }); + let t_s2 = + time(|| unsafe { s2(black_box(&src), black_box(&gate), black_box(&elig), black_box(&mut d)) }); + let t_s1 = time(|| s1(black_box(&src), black_box(&gate), black_box(&elig), black_box(&mut d))); + let t_w = time(|| w(black_box(&src), black_box(&payload), black_box(&mut d))); + let best = [("S8", t_s8), ("S4", t_s4), ("S2", t_s2)] + .iter() + .fold(("F8", t_f8), |acc, &(n, t)| if t < acc.1 { (n, t) } else { acc }); + println!( + "{:<9} {:<10} {:>8} | {:>8.0} {:>8.0} {:>8.0} {:>8.0} {:>8.0} {:>9.0} | {} ({:.2}x F8, {:.2}x W)", + format!("{pct}%"), + if clustered { "clustered" } else { "uniform" }, + live, + t_f8, + t_s8, + t_s4, + t_s2, + t_s1, + t_w, + best.0, + t_f8 / best.1, + t_w / best.1 + ); + } + } + println!("\nns per call; equivalence asserted for every arm before timing."); + } +} + +#[cfg(all( + target_arch = "x86_64", + target_feature = "avx512f", + target_feature = "avx512vl", + target_feature = "avx512dq" +))] +fn main() { + probe::main(); +} + +#[cfg(not(all( + target_arch = "x86_64", + target_feature = "avx512f", + target_feature = "avx512vl", + target_feature = "avx512dq" +)))] +fn main() { + println!( + "ternlogq_sparse_reapply_probe: needs x86_64 + avx512f + avx512vl + avx512dq — run under .cargo/config-v4.toml" + ); +} diff --git a/examples/ternlogq_tail_descent_probe.rs b/examples/ternlogq_tail_descent_probe.rs new file mode 100644 index 00000000..c5eaabf4 --- /dev/null +++ b/examples/ternlogq_tail_descent_probe.rs @@ -0,0 +1,395 @@ +//! `VPTERNLOGQ` tail: pad to zmm, or descend zmm → ymm → xmm? +//! +//! `mask_ternlog` chunks over `U64x8::LANES` = **8 words = 512 rows**, then +//! ends on three `pad_tail`s — zero-fill three 8-word stack arrays, run the +//! full 512-bit op, copy the live prefix back out. Anything under 512 rows has +//! no full chunk at all, so the whole operation IS that padded tail +//! (`ogar-r2il`'s `CallMask` is `[u64; 3]`). +//! +//! The alternative (operator, 2026-09-16): **do not pad — descend.** `VPTERNLOGQ` +//! has ymm and xmm encodings under AVX512VL, so a 1..7-word remainder splits as +//! 4 + 2 + 1 with every lane live, no zero-fill, no prefix copy — and, unlike a +//! scalar peel, **it stays entirely in vector registers**. That distinction is +//! load-bearing: `scripts/codegen-witness.sh` caps GPR logic on lane data at +//! `SLICE_GPR_CAP=6` and records that an exact-length scalar tail measured 16 GPR +//! ops on aarch64 (LLVM unrolled it to 7×(and, orr)), which is what `pad_tail` +//! removed. A vector descent is not that peel and does not reintroduce it. +//! +//! The facade carries `ternlog` on `U64x8` and `U32x16` only — there is no +//! `U64x4`/`U64x2` wrapper — so the split arm calls the intrinsics such a wrapper +//! would hold. Measuring whether they are worth adding is the point of the probe. +//! +//! AVX-512 only, deliberately: v3 is the GitHub/distribution baseline, not a +//! deployment target. +//! Gated on `avx512f` AND `avx512vl` (Codex P2 on #311): the ymm/xmm rungs are +//! VL encodings, and a `#[target_feature]` attribute is a caller precondition, +//! not a CPU check — an F-only target (Knights Landing) takes the no-op `main`. +//! +//! # Measured 2026-09-16 — Xeon @ 2.10 GHz, `avx512f=true`, release, 3 runs +//! +//! ## The crux: for a remainder of `t` words, which decomposition? +//! +//! Tail only, no body loop. ns per call. +//! +//! | t | shape | **P** padded zmm | **G** greedy | **X** all-xmm | winner | +//! |---:|---|---:|---:|---:|---| +//! | 1 | `1` | 13.1-13.8 | 1.78-1.90 | 1.49-1.60 | G = X, same instructions | +//! | 2 | `2` | 13.0-13.7 | **1.74-1.94** | 2.57-2.78 | **G** | +//! | 3 | `2+1` | 17.1-18.6 | **2.13-2.36** | 2.82-3.23 | **G** | +//! | 4 | `4` vs `2+2` | 12.8-13.6 | **1.59-1.76** | 2.92-3.00 | **G** | +//! | 5 | `4+1` vs `2+2+1` | 17.3-18.3 | **2.08-2.13** | 3.23-3.49 | **G** | +//! | **6** | **`4+2`** vs `2+2+2` | 17.0-18.3 | **2.26-2.74** | 3.29-3.68 | **G** | +//! | 7 | `4+2+1` vs `2+2+2+1` | 16.8-18.5 | **2.38-2.57** | 3.68-3.95 | **G** | +//! +//! **Greedy widest-first wins at every `t >= 2`** — 1.3-1.6x over all-xmm and +//! **5-8x over padding**. The `t = 6` case: `4+2` at 2.26-2.74 ns against +//! `2+2+2` at 3.29-3.68 ns — one ymm plus one xmm beats three xmm. Fewer wider +//! ops win; three 128-bit ops do not pay for avoiding one 256-bit one. +//! +//! At `t = 1` G and X compile to the SAME instructions (neither the ymm nor the +//! pair rung is reachable), so the 0.3 ns gap is layout noise. Recorded as a tie +//! rather than a ranking — a "winner" there is an artifact of the comparison. +//! +//! ## End to end, through the real chunk loop +//! +//! | words | rows | tail | padded | descend | pad/descend | +//! |---:|---:|---:|---:|---:|---:| +//! | 3 | 192 | 3 | 18.1 | 2.3 | **7.7-8.0x** | +//! | 1-7 | 64-448 | 1-7 | 13-23 | 2.0-2.6 | **5.9-8.0x** | +//! | 9 | 576 | 1 | 14.2 | 3.7 | **4.2x** | +//! | 11 | 704 | 3 | 21.9 | 3.6 | **5.4-5.6x** | +//! | 31 | 1984 | 7 | 23.6 | 6.8 | **2.5-3.3x** | +//! | 194 | 12416 | 2 | 41.4 | 26.5 | 1.3-1.4x | +//! | 8 / 16 / 24 / 64 | - | **none** | - | - | 1.0-1.4x | +//! +//! `ogar-r2il`'s `CallMask` is `[u64; 3]` — three words, zero full chunks — so +//! it sits at the 7.7-8.0x row: 18 ns to AND three words, because three 8-word +//! stack arrays are zeroed and partly filled to compute 192 live rows inside a +//! 512-row register. +//! +//! **What is NOT the tail:** the no-tail rows still show 1.0-1.4x. That is loop +//! shape (`as_chunks` + zip + `from_array`/`to_array` round-trips vs direct +//! `loadu`/`storeu`), so the tail-attributable factor is ~5-7x at small word +//! counts, not the full 8x. Subtract it before quoting these. +//! +//! ## The emitted code, checked rather than claimed +//! +//! `--emit=asm` on this probe: **33 zmm + 3 ymm + 6 xmm `vpternlogq`**, several +//! with folded memory operands (`vpternlogq $128, (%r13,%r10,8), %ymm0, %ymm1`), +//! and **zero** GPR `and`/`or`/`xor` on lane data. That is what separates a +//! descent from a scalar peel: `scripts/codegen-witness.sh` caps GPR logic at +//! `SLICE_GPR_CAP=6` and records an exact-length scalar tail measuring 16 GPR ops +//! on aarch64 — the thing `pad_tail` was introduced to remove. A ymm/xmm descent +//! is not that peel and does not reintroduce it. +//! +//! ```text +//! env -u RUSTFLAGS cargo --config .cargo/config-v4.toml run --release \ +//! --example ternlogq_tail_descent_probe +//! ``` + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +use std::hint::black_box; +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +use std::time::Instant; + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +use ndarray::simd::ternlog::AND3; +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +use ndarray::simd::U64x8; + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +use std::arch::x86_64::*; + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +#[target_feature(enable = "avx512f,avx512vl")] +/// **P** — pad the remainder into one zmm (today's shape). +unsafe fn tail_padded(a: &[u64], b: &[u64], c: &[u64], dst: &mut [u64]) { + let t = dst.len(); + let (mut pa, mut pb, mut pc) = ([0u64; 8], [0u64; 8], [0u64; 8]); + pa[..t].copy_from_slice(a); + pb[..t].copy_from_slice(b); + pc[..t].copy_from_slice(c); + let v = U64x8::from_array(pa) + .ternlog::(U64x8::from_array(pb), U64x8::from_array(pc)) + .to_array(); + dst.copy_from_slice(&v[..t]); +} + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +#[target_feature(enable = "avx512f,avx512vl")] +/// **G** — greedy widest-first: ymm, then xmm, then xmm-low. +unsafe fn tail_greedy(a: &[u64], b: &[u64], c: &[u64], dst: &mut [u64]) { + let (t, mut i) = (dst.len(), 0usize); + // SAFETY: caller guarantees avx512f+avx512vl (the `#[target_feature]` + // precondition) and equal-length slices; each rung reads/writes exactly + // `4`, `2` or `1` words at `i` only when `t - i` covers it, so every + // unaligned load/store stays inside `a`/`b`/`c`/`dst`. + unsafe { + if t - i >= 4 { + let v = _mm256_ternarylogic_epi64::( + _mm256_loadu_si256(a.as_ptr().add(i).cast()), + _mm256_loadu_si256(b.as_ptr().add(i).cast()), + _mm256_loadu_si256(c.as_ptr().add(i).cast()), + ); + _mm256_storeu_si256(dst.as_mut_ptr().add(i).cast(), v); + i += 4; + } + while t - i >= 2 { + let v = _mm_ternarylogic_epi64::( + _mm_loadu_si128(a.as_ptr().add(i).cast()), + _mm_loadu_si128(b.as_ptr().add(i).cast()), + _mm_loadu_si128(c.as_ptr().add(i).cast()), + ); + _mm_storeu_si128(dst.as_mut_ptr().add(i).cast(), v); + i += 2; + } + if t - i == 1 { + let v = _mm_ternarylogic_epi64::( + _mm_loadl_epi64(a.as_ptr().add(i).cast()), + _mm_loadl_epi64(b.as_ptr().add(i).cast()), + _mm_loadl_epi64(c.as_ptr().add(i).cast()), + ); + _mm_storel_epi64(dst.as_mut_ptr().add(i).cast(), v); + } + } +} + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +#[target_feature(enable = "avx512f,avx512vl")] +/// **X** — never widen past xmm: pairs, then the odd one. +unsafe fn tail_all_xmm(a: &[u64], b: &[u64], c: &[u64], dst: &mut [u64]) { + let (t, mut i) = (dst.len(), 0usize); + // SAFETY: same contract as `tail_greedy` — avx512f+avx512vl from the + // caller, equal-length slices, and each xmm/xmm-low access is guarded by + // `t - i >= 2` / `== 1`, so no access leaves the slices. + unsafe { + while t - i >= 2 { + let v = _mm_ternarylogic_epi64::( + _mm_loadu_si128(a.as_ptr().add(i).cast()), + _mm_loadu_si128(b.as_ptr().add(i).cast()), + _mm_loadu_si128(c.as_ptr().add(i).cast()), + ); + _mm_storeu_si128(dst.as_mut_ptr().add(i).cast(), v); + i += 2; + } + if t - i == 1 { + let v = _mm_ternarylogic_epi64::( + _mm_loadl_epi64(a.as_ptr().add(i).cast()), + _mm_loadl_epi64(b.as_ptr().add(i).cast()), + _mm_loadl_epi64(c.as_ptr().add(i).cast()), + ); + _mm_storel_epi64(dst.as_mut_ptr().add(i).cast(), v); + } + } +} + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +/// Today's shape: three zero-filled 8-word arrays, one zmm op, prefix copy out. +fn tern_padded(a: &[u64], b: &[u64], c: &[u64], dst: &mut [u64]) { + const L: usize = U64x8::LANES; + let (ca, ta) = a.as_chunks::(); + let (cb, tb) = b.as_chunks::(); + let (cc, tc) = c.as_chunks::(); + let (cd, td) = dst.as_chunks_mut::(); + for (((x, y), z), d) in ca.iter().zip(cb).zip(cc).zip(cd.iter_mut()) { + *d = U64x8::from_array(*x) + .ternlog::(U64x8::from_array(*y), U64x8::from_array(*z)) + .to_array(); + } + if !ta.is_empty() { + let (mut pa, mut pb, mut pc) = ([0u64; L], [0u64; L], [0u64; L]); + pa[..ta.len()].copy_from_slice(ta); + pb[..tb.len()].copy_from_slice(tb); + pc[..tc.len()].copy_from_slice(tc); + let v = U64x8::from_array(pa) + .ternlog::(U64x8::from_array(pb), U64x8::from_array(pc)) + .to_array(); + td.copy_from_slice(&v[..td.len()]); + } +} + +/// zmm for the body, then ymm → xmm → xmm-low for the remainder. Every lane +/// live; no zero-fill, no prefix copy, and no GPR logic on lane data. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +#[target_feature(enable = "avx512f,avx512vl")] +unsafe fn tern_descend(a: &[u64], b: &[u64], c: &[u64], dst: &mut [u64]) { + const L: usize = U64x8::LANES; + let n = dst.len(); + let full = n / L; + for g in 0..full { + let (i, p) = (g * L, dst.as_mut_ptr()); + // SAFETY: caller guarantees avx512f+avx512vl and equal-length slices; + // `g < n / L` keeps the 8-word window `[i, i + L)` inside every slice. + unsafe { + let va = _mm512_loadu_si512(a.as_ptr().add(i).cast()); + let vb = _mm512_loadu_si512(b.as_ptr().add(i).cast()); + let vc = _mm512_loadu_si512(c.as_ptr().add(i).cast()); + _mm512_storeu_si512(p.add(i).cast(), _mm512_ternarylogic_epi64::(va, vb, vc)); + } + } + let mut i = full * L; + // SAFETY: same contract; the remainder `n - i < L` is descended by rungs + // each guarded by `n - i >= 4` / `>= 2` / `== 1`, so every ymm / xmm / + // 64-bit access stays inside the slices. + unsafe { + // ymm rung — 4 live words. + if n - i >= 4 { + let va = _mm256_loadu_si256(a.as_ptr().add(i).cast()); + let vb = _mm256_loadu_si256(b.as_ptr().add(i).cast()); + let vc = _mm256_loadu_si256(c.as_ptr().add(i).cast()); + _mm256_storeu_si256(dst.as_mut_ptr().add(i).cast(), _mm256_ternarylogic_epi64::(va, vb, vc)); + i += 4; + } + // xmm rung — 2 live words. + if n - i >= 2 { + let va = _mm_loadu_si128(a.as_ptr().add(i).cast()); + let vb = _mm_loadu_si128(b.as_ptr().add(i).cast()); + let vc = _mm_loadu_si128(c.as_ptr().add(i).cast()); + _mm_storeu_si128(dst.as_mut_ptr().add(i).cast(), _mm_ternarylogic_epi64::(va, vb, vc)); + i += 2; + } + // xmm-low rung — the final single word, still vector (64-bit load/store, + // so it never reads past the slice). + if n - i == 1 { + let va = _mm_loadl_epi64(a.as_ptr().add(i).cast()); + let vb = _mm_loadl_epi64(b.as_ptr().add(i).cast()); + let vc = _mm_loadl_epi64(c.as_ptr().add(i).cast()); + _mm_storel_epi64(dst.as_mut_ptr().add(i).cast(), _mm_ternarylogic_epi64::(va, vb, vc)); + } + } +} + +#[cfg(not(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl")))] +fn main() { + // The descent rungs are `VPTERNLOGQ` ymm/xmm under AVX512VL, so this probe + // has nothing to measure off v4 x86_64. It still has to BUILD on the + // matrix's aarch64 / wasm32 / v3 rows, hence a running no-op rather than a + // `compile_error!` or a constant assert (which clippy rejects anyway). + println!("skipped: needs x86_64 + avx512f + avx512vl — run under .cargo/config-v4.toml"); +} + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "avx512vl"))] +fn main() { + println!("realization: avx512f=true\n"); + + for words in 1..=40usize { + let a: Vec = (0..words) + .map(|i| 0xF0F0_F0F0_0000_1111u64 ^ (i as u64 * 31)) + .collect(); + let b: Vec = (0..words) + .map(|i| 0xFFFF_0000_FFFF_0000u64 ^ (i as u64 * 17)) + .collect(); + let c: Vec = (0..words) + .map(|i| 0x0F0F_0F0F_1111_0000u64 ^ (i as u64 * 7)) + .collect(); + let (mut d1, mut d2) = (vec![0u64; words], vec![0u64; words]); + tern_padded(&a, &b, &c, &mut d1); + // SAFETY: avx512f+avx512vl asserted above; every load/store is within + // the slices (the final rung uses a 64-bit load, not a 128-bit one). + unsafe { tern_descend(&a, &b, &c, &mut d2) }; + assert_eq!(d1, d2, "arms disagree at words={words}"); + } + println!("equivalence: 1..=40 words, bit-identical\n"); + + macro_rules! bench { + ($body:expr, $reps:expr) => {{ + $body; + let t0 = Instant::now(); + for _ in 0..$reps { + $body; + } + t0.elapsed().as_secs_f64() / ($reps as f64) * 1e9 + }}; + } + + // ---- the crux: for a tail of t words, which decomposition wins? ---- + // t = 6 is the interesting one: 1 ymm + 1 xmm, or 3 xmm, or one padded zmm? + println!("== TAIL-ONLY, per remainder length: which split? (ns per call) =="); + println!("{:>5} {:>10} {:>10} {:>10} {}", "t", "P padded", "G greedy", "X all-xmm", "winner"); + for t in 1..=7usize { + let a: Vec = (0..t).map(|i| 0xF0F0_1111u64 ^ i as u64).collect(); + let b: Vec = (0..t).map(|i| 0xFF00_2222u64 ^ i as u64).collect(); + let c: Vec = (0..t).map(|i| 0x0F0F_4444u64 ^ i as u64).collect(); + let (mut dp, mut dg, mut dx) = (vec![0u64; t], vec![0u64; t], vec![0u64; t]); + // SAFETY: avx512f+avx512vl hold on this arm; every access is in-slice. + unsafe { + tail_padded(&a, &b, &c, &mut dp); + tail_greedy(&a, &b, &c, &mut dg); + tail_all_xmm(&a, &b, &c, &mut dx); + } + assert_eq!(dp, dg, "greedy disagrees at t={t}"); + assert_eq!(dp, dx, "all-xmm disagrees at t={t}"); + let reps = 3_000_000; + // SAFETY: as above. + let p = bench!( + { + unsafe { tail_padded(black_box(&a), black_box(&b), black_box(&c), &mut dp) }; + black_box(&dp); + }, + reps + ); + let g = bench!( + { + unsafe { tail_greedy(black_box(&a), black_box(&b), black_box(&c), &mut dg) }; + black_box(&dg); + }, + reps + ); + let x = bench!( + { + unsafe { tail_all_xmm(black_box(&a), black_box(&b), black_box(&c), &mut dx) }; + black_box(&dx); + }, + reps + ); + // At t == 1 the G and X bodies are the SAME instructions (neither the + // ymm nor the pair rung is reachable, both fall to xmm-low), so a gap + // there is layout noise and a "winner" would be an artifact of this + // comparison rather than a strategy difference. Say so, do not rank it. + let best = if t == 1 { + "G=X (identical code)" + } else if p <= g && p <= x { + "P" + } else if g <= x { + "G" + } else { + "X" + }; + let shape = match t { + 1 => "1", + 2 => "2", + 3 => "2+1", + 4 => "4 | 2+2", + 5 => "4+1 | 2+2+1", + 6 => "4+2 | 2+2+2", + _ => "4+2+1 | 2+2+2+1", + }; + println!("{:>5} {:>10.2} {:>10.2} {:>10.2} {} [{}]", t, p, g, x, best, shape); + } + println!(); + + println!("{:>6} {:>7} {:>5} {:>11} {:>11} {:>10}", "words", "rows", "tail", "padded", "descend", "pad/desc"); + for &words in &[1usize, 2, 3, 4, 5, 6, 7, 8, 9, 11, 16, 24, 31, 64, 194] { + let a: Vec = (0..words).map(|i| 0xF0F0u64 ^ i as u64).collect(); + let b: Vec = (0..words).map(|i| 0xFF00u64 ^ i as u64).collect(); + let c: Vec = (0..words).map(|i| 0x0F0Fu64 ^ i as u64).collect(); + let (mut d1, mut d2) = (vec![0u64; words], vec![0u64; words]); + let reps = (3_000_000 / words.max(1)).max(2000); + let p = bench!( + { + tern_padded(black_box(&a), black_box(&b), black_box(&c), &mut d1); + black_box(&d1); + }, + reps + ); + let s = bench!( + { + // SAFETY: as above. + unsafe { tern_descend(black_box(&a), black_box(&b), black_box(&c), &mut d2) }; + black_box(&d2); + }, + reps + ); + println!("{:>6} {:>7} {:>5} {:>11.2} {:>11.2} {:>10.2}", words, words * 64, words % 8, p, s, p / s); + } +} diff --git a/src/backend/kernels_avx512.rs b/src/backend/kernels_avx512.rs index ce80a7d8..5e3963f5 100644 --- a/src/backend/kernels_avx512.rs +++ b/src/backend/kernels_avx512.rs @@ -965,3 +965,211 @@ pub fn hamming_batch(query: &[u8], database: &[u8], num_rows: usize, row_bytes: }) .collect() } + +/// Block-stop probe (operator, 2026-09-16: *"would the MKL GEMM stop logic +/// profit from a tail optimization?"*). `sgemm_blocked`'s M-stop pads the last +/// `MR = 6` panel to six rows and the ukernel computes all six accumulators +/// regardless of `mr_eff`; every power-of-two `m` has such a tail (128 = 21·6+2, +/// 256 = 42·6+4, 512 = 85·6+2, 1024 = 170·6+4). This test-local variant runs an +/// `R × 16` ukernel on the tail tile only (R ∈ {2, 4}) and times both against +/// each other, equivalence-gated first. `cargo test --release --lib +/// block_stop_probe -- --ignored --nocapture` under `.cargo/config-v4.toml`. +/// +/// Measured 2026-09-16 (Xeon @ 2.10 GHz, `avx512f=true`, release, best of 9): +/// +/// | m×n×k | tail | shipped ms | desc ms | ratio | FMA waste | +/// |---|---:|---:|---:|---:|---:| +/// | 126×256×256 | 0 | 0.183 | 0.185 | 0.987× | 0.0 % | +/// | 128×256×256 | 2 | 0.190 | 0.192 | 0.989× | 3.1 % | +/// | 130×256×256 | 4 | 0.190 | 0.191 | 0.992× | 1.5 % | +/// | 132×256×256 | 0 | 0.193 | 0.190 | 1.012× | 0.0 % | +/// | 128³ | 2 | 0.056 | 0.055 | 1.015× | 3.1 % | +/// | 256³ | 4 | 0.347 | 0.350 | 0.992× | 0.8 % | +/// | 512³ | 2 | 3.110 | 3.054 | 1.018× | 0.8 % | +/// | 1024³ | 4 | 32.57 | 32.04 | 1.016× | 0.2 % | +/// +/// **Inert.** Bit-identical output, 0.99–1.02× everywhere — noise. The padded +/// accumulator rows never reach the critical path: at 1024³ the kernel runs at +/// ~66 GFLOP/s (about half of one core's FMA peak), so packing and memory +/// traffic hide a 0.2–3 % FMA surplus entirely. The K-stop has no waste and the +/// N-stop's padding is on lanes the FMA unit processes anyway, so no lane-width +/// descent applies to GEMM at all. Recorded so the tail question is not +/// re-opened for the F32x16 path; the BF16 `vdpbf16ps` path's stop problem is a +/// single accumulator chain per row (`amx_matmul.rs`), not its tails. +#[cfg(all(test, target_arch = "x86_64", target_feature = "avx512f"))] +mod block_stop_probe { + use super::*; + use std::hint::black_box; + use std::time::Instant; + + #[target_feature(enable = "avx512f")] + unsafe fn ukernel_rows( + kc: usize, alpha: f32, a_packed: &[f32], b_packed: &[f32], c: &mut [f32], ldc: usize, nr_eff: usize, + ) { + let mut acc = [_mm512_setzero_ps(); R]; + for p in 0..kc { + let bv = _mm512_loadu_ps(b_packed[p * SGEMM_NR..].as_ptr()); + let a_off = p * SGEMM_MR; + for r in 0..R { + acc[r] = _mm512_fmadd_ps(_mm512_set1_ps(a_packed[a_off + r]), bv, acc[r]); + } + } + let alpha_v = _mm512_set1_ps(alpha); + for (r, v) in acc.iter().enumerate() { + let row_ptr = c[r * ldc..].as_mut_ptr(); + let v = _mm512_mul_ps(*v, alpha_v); + if nr_eff == SGEMM_NR { + _mm512_storeu_ps(row_ptr, _mm512_add_ps(_mm512_loadu_ps(row_ptr), v)); + } else { + let mask: u16 = (1u32 << nr_eff) as u16 - 1; + _mm512_mask_storeu_ps(row_ptr, mask, _mm512_add_ps(_mm512_maskz_loadu_ps(mask, row_ptr), v)); + } + } + } + + /// `sgemm_blocked` with the M-tail tile dispatched to an `R×16` ukernel. + #[target_feature(enable = "avx512f")] + fn sgemm_blocked_desc( + m: usize, n: usize, k: usize, alpha: f32, a: &[f32], lda: usize, b: &[f32], ldb: usize, c: &mut [f32], + ldc: usize, + ) { + let mut a_packed = vec![0.0f32; SGEMM_MC * SGEMM_KC]; + let mut b_packed = vec![0.0f32; SGEMM_KC * SGEMM_NC]; + let mut kk = 0; + while kk < k { + let kc = SGEMM_KC.min(k - kk); + let mut jj = 0; + while jj < n { + let nc = SGEMM_NC.min(n - jj); + pack_b_f32(b, ldb, kc, nc, kk, jj, &mut b_packed); + let mut ii = 0; + while ii < m { + let mc = SGEMM_MC.min(m - ii); + pack_a_f32(a, lda, mc, kc, ii, kk, &mut a_packed); + let mut ir = 0; + while ir < mc { + let mr_eff = SGEMM_MR.min(mc - ir); + let mut jr = 0; + while jr < nc { + let nr_eff = SGEMM_NR.min(nc - jr); + let a_off = (ir / SGEMM_MR) * (SGEMM_MR * kc); + let b_off = (jr / SGEMM_NR) * (SGEMM_NR * kc); + let cc = &mut c[(ii + ir) * ldc + (jj + jr)..]; + // SAFETY: this fn is `#[target_feature(enable = "avx512f")]`, + // so the callee's feature precondition holds; `a_off`/`b_off` + // index whole packed panels of `mr_eff`/`nr_eff` live rows + // inside `a_packed`/`b_packed`, and `cc` starts at `(ii+ir, + // jj+jr)` with `ldc` so every masked store lands in `c`. + unsafe { + match mr_eff { + 2 => ukernel_rows::<2>( + kc, + alpha, + &a_packed[a_off..], + &b_packed[b_off..], + cc, + ldc, + nr_eff, + ), + 4 => ukernel_rows::<4>( + kc, + alpha, + &a_packed[a_off..], + &b_packed[b_off..], + cc, + ldc, + nr_eff, + ), + _ => sgemm_ukernel_6x16( + kc, + alpha, + &a_packed[a_off..], + &b_packed[b_off..], + cc, + ldc, + mr_eff, + nr_eff, + ), + } + } + jr += SGEMM_NR; + } + ir += SGEMM_MR; + } + ii += mc; + } + jj += nc; + } + kk += kc; + } + } + + fn fill(n: usize, seed: u64) -> Vec { + let mut s = seed; + (0..n) + .map(|_| { + s = s + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((s >> 40) as f32 / (1u64 << 24) as f32) - 0.5 + }) + .collect() + } + + #[test] + #[ignore] + fn block_stop_probe() { + assert!(std::is_x86_feature_detected!("avx512f")); + eprintln!("realization: avx512f=true MR={SGEMM_MR} NR={SGEMM_NR} MC={SGEMM_MC}"); + eprintln!( + "{:>5} {:>5} {:>5} {:>4} | {:>10} {:>10} | {:>7} {:>8}", + "m", "n", "k", "tail", "shipped_ms", "desc_ms", "ratio", "fma_waste" + ); + let shapes: Vec<(usize, usize, usize)> = [(126, 256, 256), (128, 256, 256), (130, 256, 256), (132, 256, 256)] + .into_iter() + .chain([128usize, 256, 512, 1024].into_iter().map(|s| (s, s, s))) + .collect(); + for (m, n, k) in shapes { + let a = fill(m * k, 1); + let b = fill(k * n, 2); + let mut c1 = vec![0.0f32; m * n]; + let mut c2 = vec![0.0f32; m * n]; + // SAFETY: `is_x86_feature_detected!("avx512f")` asserted above, so + // both `#[target_feature]` callees may run; buffers are `m*k`, + // `k*n`, `m*n` with the matching leading dimensions. + unsafe { + sgemm_blocked(m, n, k, 1.0, &a, k, &b, n, &mut c1, n); + sgemm_blocked_desc(m, n, k, 1.0, &a, k, &b, n, &mut c2, n); + } + assert!(c1.iter().zip(&c2).all(|(x, y)| x.to_bits() == y.to_bits()), "{m}x{n}x{k}: desc diverges"); + let time = |f: &mut dyn FnMut()| { + f(); + let reps = if m >= 1024 { 3 } else { 9 }; + let mut best = f64::MAX; + for _ in 0..reps { + let t = Instant::now(); + f(); + best = best.min(t.elapsed().as_secs_f64() * 1e3); + } + best + }; + let t1 = time(&mut || { + c1.fill(0.0); + // SAFETY: as the equivalence call above — avx512f asserted, same buffers. + unsafe { sgemm_blocked(m, n, k, 1.0, black_box(&a), k, black_box(&b), n, black_box(&mut c1), n) } + }); + let t2 = time(&mut || { + c2.fill(0.0); + // SAFETY: as the equivalence call above — avx512f asserted, same buffers. + unsafe { sgemm_blocked_desc(m, n, k, 1.0, black_box(&a), k, black_box(&b), n, black_box(&mut c2), n) } + }); + let tail = m % SGEMM_MR; + let waste = if tail == 0 { + 0.0 + } else { + (SGEMM_MR - tail) as f64 / m as f64 * 100.0 + }; + eprintln!("{m:>5} {n:>5} {k:>5} {tail:>4} | {t1:>10.3} {t2:>10.3} | {:>6.3}x {waste:>7.1}%", t1 / t2); + } + } +} diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index 6004e275..3ecc13c6 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -5034,7 +5034,7 @@ mod tests { #[test] fn mask_set_range_matches_scalar_reference_randomized() { - let mut seed = 0xB000_0000_5E7A_11u64; + let mut seed = 0x00B0_0000_005E_7A11_u64; for &words in &[1usize, 2, 3, 4, 7, 8, 16] { let capacity = words * 64; for _ in 0..40 {