From abd890b179794ef30e05fcd835fedf3f654f6346 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 5 Sep 2026 21:26:14 +0000 Subject: [PATCH 1/4] amx_matmul: make matmul_f32 exact; AMX f32 paths become named opt-ins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `matmul_f32` on an AMX host downcast both operands to BF16 once and ran TDPBF16PS. The one-rounding-then-f32-accumulate discipline was correct and intact; what it cannot do is recover the 16 significand bits that single rounding discards. Measured: ~1e-3 relative error (bf16 mantissa, eps 7.8e-3, f32 accumulate), uniform across aligned and ragged shapes, while the API returned Ok(()) under an f32 name. Routing burn's burn-ndarray through it failed 50 linalg tests (25 qr, 13 lu, 7 svd, 3 det, 1 attention) that pass at 1826/1826 on the exact path. The existing test could not see this: its inputs `(i+j)*0.5` and `(i*3+j)*0.25` are all exactly BF16-representable, so the downcast was lossless by construction, and its AMX tolerance was 1%. It now uses irrational inputs at 1e-5. Two candidate fixes were built and benchmarked against each other and against the CPU f32 paths (`gemm_paths_bench`, Xeon w/ AMX + AVX-512, square f32, release): size native gemm_f32 F32x16 sgemm_blocked AMX 3-pass split AMX 1-pass 256 0.42ms 2.9e-7 0.37ms 2.9e-7 2.78ms 1.4e-6 0.77ms 3.7e-4 512 3.51ms 1.2e-6 3.22ms 1.2e-6 20.9ms 1.4e-6 6.37ms 2.2e-4 1024 28.8ms 1.4e-6 27.3ms 1.4e-6 159ms 1.5e-6 45.9ms 1.6e-4 AMX loses on both axes at every size. The three-pass BF16 hi/lo split does reach f32 grade (~400x better than one pass) but is ~6x slower than the exact CPU kernel it was meant to replace, and even the lossy single pass is slower than plain f32. Packing, conversion and passes cost more than the tile unit saves for f32. So: - `matmul_f32` now delegates to `backend::native::gemm_f32` and is exact, on every host. It never touches AMX. - `matmul_f32_amx_split` (the three-pass split) and `matmul_f32_bf16_fast` (the old single pass) are kept as explicitly named opt-ins with the table above in their docs, so the measurement is reproducible and a future host or a real f32 tile op can be re-measured against them. Neither is recommended. - `three_pass_split_beats_one_bf16_pass` pins all three two-sided: the default at <=1e-5, the split at <=1e-5 and >=50x better than one pass, and the one pass at >1e-4 — the last so that if AMX ever gains a real f32 op the cost is re-justified rather than assumed. - `gemm_paths_bench` (`#[ignore]`, `--release`) is the bench above. The one-rounding technique was the right half of the answer; this change adds the other half, then measures that neither half beats not using AMX for f32 at all. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01EfrUJH3UNnv5NpDH4jDGHq --- src/hpc/amx_matmul.rs | 258 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 242 insertions(+), 16 deletions(-) diff --git a/src/hpc/amx_matmul.rs b/src/hpc/amx_matmul.rs index d8322455..a241f361 100644 --- a/src/hpc/amx_matmul.rs +++ b/src/hpc/amx_matmul.rs @@ -412,8 +412,14 @@ pub fn vnni_pack_i8(src: &[i8], dst: &mut [i8], k: usize, n: usize) { // ═══════════════════════════════════════════════════════════════════════════ // // Three entry points operating on `ArrayView2` / `ArrayViewMut2`: -// matmul_f32 — f32 × f32 → f32 (BF16 compute via AMX TDPBF16PS, -// f32 fallback on hosts without AMX) +// matmul_f32 — f32 × f32 → f32, EXACT: `backend::native::gemm_f32` +// (matrixmultiply). Never AMX — measured slower and +// less accurate than CPU f32 at every size (see +// `matmul_f32_amx_split`'s table). +// matmul_f32_amx_split — three-pass BF16 split on AMX, f32-grade but ~6x +// slower than `matmul_f32`. Kept for re-measurement. +// matmul_f32_bf16_fast — ONE BF16 pass on AMX (~1e-3 rel). Still slower +// than `matmul_f32`. Opt in by name only. // matmul_bf16_to_f32 — BF16 × BF16 → f32 (AMX TDPBF16PS or `bf16_gemm_f32`) // matmul_i8_to_i32 — i8 × i8 → i32 (AMX TDPBUSD or scalar `int8_gemm_i32`) // @@ -714,23 +720,137 @@ unsafe fn bf16_gemm_vdpbf16ps(a: &[BF16], b: &[BF16], c: &mut [f32], m: usize, n /// computation runs in pure f32 and is bit-stable. /// /// `out` must be row-contiguous; inputs may be strided. -pub fn matmul_f32( +/// Splits `v` into a BF16 head and a BF16 tail such that +/// `head.to_f32() + tail.to_f32()` reproduces `v` to ~17 significand bits. +/// +/// Both halves are exactly BF16-representable, so a later `from_f32_rounded` +/// on either is the identity — that is what keeps the three-pass product at +/// one rounding per operand rather than two. +fn split_bf16(v: &[f32]) -> (Vec, Vec) { + let mut hi = Vec::with_capacity(v.len()); + let mut lo = Vec::with_capacity(v.len()); + for &x in v { + let h = BF16::from_f32_rounded(x); + hi.push(h); + lo.push(BF16::from_f32_rounded(x - h.to_f32())); + } + (hi, lo) +} + +/// Single-BF16-pass `f32` matmul: one third of [`matmul_f32`]'s tile work, +/// at BF16 precision (~1e-3 relative error, measured). +/// +/// Use only where BF16-grade accuracy has been measured as acceptable for the +/// workload. [`matmul_f32`] is the default because its name promises f32 and +/// a silent 1e-3 is outside what f32 callers — linear algebra especially — +/// tolerate: routing burn's `burn-ndarray` through this path fails 50 of its +/// linalg tests (qr / lu / svd / det) that pass on the three-pass version. +pub fn matmul_f32_bf16_fast( lhs: ArrayView2<'_, f32>, rhs: ArrayView2<'_, f32>, mut out: ArrayViewMut2<'_, f32>, ) -> Result<(), MatmulError> { let (m, n, k) = check_shapes(&lhs, &rhs, &out)?; - let a_f32 = pack_contig(&lhs); let b_f32 = pack_contig(&rhs); let mut c = vec![0.0f32; m * n]; if amx_available() { - // AMX path: down-cast to BF16 (RNE, ~1 ULP at BF16 mantissa - // precision), then dispatch through the shared BF16 helper - // which picks `TDPBF16PS` tile kernel for 16/16/32-aligned - // shapes and the scalar `bf16_gemm_f32` reference otherwise. let a_bf16: Vec = a_f32.iter().map(|&v| BF16::from_f32_rounded(v)).collect(); let b_bf16: Vec = b_f32.iter().map(|&v| BF16::from_f32_rounded(v)).collect(); bf16_gemm_dispatch(&a_bf16, &b_bf16, &mut c, m, n, k); + } else { + for i in 0..m { + for p in 0..k { + let av = a_f32[i * k + p]; + for j in 0..n { + c[i * n + j] += av * b_f32[p * n + j]; + } + } + } + } + + write_contig(&mut out, &c); + Ok(()) +} + +/// Exact f32 matmul. Never routes through AMX — measured slower AND less +/// accurate than the CPU f32 GEMM at every size tried (see the table in +/// [`matmul_f32_amx_split`]). +/// +/// Delegates to `backend::native::gemm_f32` (matrixmultiply), which is +/// f32 end to end: no BF16 downcast, no tile emulation. +pub fn matmul_f32( + lhs: ArrayView2<'_, f32>, rhs: ArrayView2<'_, f32>, mut out: ArrayViewMut2<'_, f32>, +) -> Result<(), MatmulError> { + let (m, n, k) = check_shapes(&lhs, &rhs, &out)?; + let a_f32 = pack_contig(&lhs); + let b_f32 = pack_contig(&rhs); + let mut c = vec![0.0f32; m * n]; + crate::backend::native::gemm_f32(m, n, k, 1.0, &a_f32, k, &b_f32, n, 0.0, &mut c, n); + write_contig(&mut out, &c); + Ok(()) +} + +/// AMX three-pass BF16-split matmul: f32-grade accuracy on the tile unit. +/// +/// KEPT FOR THE RECORD, NOT RECOMMENDED. Benchmarked against the exact CPU +/// paths on a Xeon with AMX + AVX-512 (`gemm_paths_bench`), square f32: +/// +/// | size | native gemm_f32 | F32x16 sgemm_blocked | this (3-pass) | 1-pass BF16 | +/// |-------|-----------------|----------------------|---------------|-------------| +/// | 256 | 0.42ms 2.9e-7 | 0.37ms 2.9e-7 | 2.78ms 1.4e-6 | 0.77ms 3.7e-4 | +/// | 512 | 3.51ms 1.2e-6 | 3.22ms 1.2e-6 | 20.9ms 1.4e-6 | 6.37ms 2.2e-4 | +/// | 1024 | 28.8ms 1.4e-6 | 27.3ms 1.4e-6 | 159ms 1.5e-6 | 45.9ms 1.6e-4 | +/// +/// AMX loses on BOTH axes at every size — the packing, the BF16 conversion +/// and (here) three passes cost more than the tile unit saves, and even the +/// single pass is slower than plain f32 while being ~1000x less accurate. +/// Exposed so the measurement is reproducible and so a future host (or a +/// real f32 tile op) can be re-measured against it, not because a caller +/// should reach for it. +pub fn matmul_f32_amx_split( + lhs: ArrayView2<'_, f32>, rhs: ArrayView2<'_, f32>, mut out: ArrayViewMut2<'_, f32>, +) -> Result<(), MatmulError> { + let (m, n, k) = check_shapes(&lhs, &rhs, &out)?; + + let a_f32 = pack_contig(&lhs); + let b_f32 = pack_contig(&rhs); + let mut c = vec![0.0f32; m * n]; + + if amx_available() { + // AMX has no f32 tile op, so f32 operands must reach the tile + // kernel as BF16. A single RNE downcast discards 16 of f32's 24 + // significand bits — measured ~1e-3 relative error, which is + // BF16 grade, not f32 grade, and this function's name promises + // f32. So split each operand into a BF16 head plus a BF16 tail: + // + // a = a_hi + a_lo, b = b_hi + b_lo (all four exactly BF16) + // a·b = a_hi·b_hi + a_hi·b_lo + a_lo·b_hi + a_lo·b_lo + // + // and drop the last term: it is O(2^-18) relative, below the f32 + // significand, so three passes suffice. Each pass still rounds + // exactly once (its inputs are already BF16-representable, so the + // downcast inside the kernel is the identity) and accumulates in + // f32 — the one-rounding discipline is preserved, and the bits it + // used to discard are now carried by the tail. + // + // Measured on Xeon w/ AMX, vs an f32 reference: + // 16x32x16 1-pass 1.09e-3 -> 3-pass 2.6e-6 + // 32x64x32 1-pass 8.40e-4 -> 3-pass 2.1e-6 + // i.e. ~400x, at 3x the tile work. Callers that want the single + // BF16 pass ask for it by name: `matmul_f32_bf16_fast`. + let (a_hi, a_lo) = split_bf16(&a_f32); + let (b_hi, b_lo) = split_bf16(&b_f32); + + let mut scratch = vec![0.0f32; m * n]; + bf16_gemm_dispatch(&a_hi, &b_hi, &mut c, m, n, k); + bf16_gemm_dispatch(&a_hi, &b_lo, &mut scratch, m, n, k); + for (dst, add) in c.iter_mut().zip(scratch.iter()) { + *dst += *add; + } + bf16_gemm_dispatch(&a_lo, &b_hi, &mut scratch, m, n, k); + for (dst, add) in c.iter_mut().zip(scratch.iter()) { + *dst += *add; + } } else { // Pure f32 reference path. for i in 0..m { @@ -956,20 +1076,74 @@ mod tests { assert!(r < 0.01, "bf16 matmul exceeded 1% relative error: {}", r); } + /// Inputs whose significands do NOT fit in BF16's 8 bits. + /// + /// The previous version of this test used `(i + j) * 0.5` and + /// `(i * 3 + j) * 0.25` — every one of those values is exactly + /// BF16-representable, so the downcast was lossless and the test + /// reported success no matter how much precision the kernel threw + /// away. Paired with a 1% AMX tolerance, it could neither detect + /// nor fail on the ~1e-3 error the single-pass path actually had. + fn irrational_pair(m: usize, k: usize, n: usize) -> (Array2, Array2) { + let a = Array2::::from_shape_fn((m, k), |(i, j)| (((i * 7 + j * 3) % 97) as f32).sqrt() * 0.3137 - 1.0); + let b = Array2::::from_shape_fn((k, n), |(i, j)| (((i * 5 + j * 2) % 89) as f32).sqrt() * 0.2713 - 0.5); + (a, b) + } + #[test] fn matmul_f32_16x16() { - let m = 16; - let n = 16; - let k = 16; - let a = Array2::::from_shape_fn((m, k), |(i, j)| ((i + j) as f32) * 0.5); - let b = Array2::::from_shape_fn((k, n), |(i, j)| ((i * 3 + j) as f32) * 0.25); + let (m, k, n) = (16, 16, 16); + let (a, b) = irrational_pair(m, k, n); let mut out = Array2::::zeros((m, n)); matmul_f32(a.view(), b.view(), out.view_mut()).expect("f32 matmul"); let expect = ref_matmul_f32(&a, &b); - // Without AMX the path is exact; with AMX up to 1% bf16 error allowed. - let tol = if amx_available() { 0.01 } else { 1e-5 }; + // f32 grade on both paths: exact without AMX, three-pass BF16 split + // with it. 1e-5 is ~100x tighter than one BF16 pass can reach, so + // this fails if `matmul_f32` ever regresses to a single downcast. let r = rel_max(&out, &expect); - assert!(r <= tol, "f32 matmul exceeded {} tol: {}", tol, r); + assert!(r <= 1e-5, "f32 matmul exceeded 1e-5 tol: {}", r); + } + + /// Two-sided: the split must beat one pass by a wide margin AND the + /// single-pass path must still be measurably lossy. + /// + /// If the second half ever fails, AMX gained a real f32 op (or the + /// host has none and both paths are the exact reference) — either way + /// the three-pass cost should be re-justified rather than assumed. + #[test] + fn three_pass_split_beats_one_bf16_pass() { + if !amx_available() { + return; // both paths are the exact f32 reference; nothing to compare + } + for (m, k, n) in [(16usize, 32usize, 16usize), (32, 64, 32)] { + let (a, b) = irrational_pair(m, k, n); + let expect = ref_matmul_f32(&a, &b); + + let mut exact = Array2::::zeros((m, n)); + matmul_f32(a.view(), b.view(), exact.view_mut()).expect("exact"); + let r_exact = rel_max(&exact, &expect); + assert!(r_exact <= 1e-5, "{m}x{k}x{n}: default matmul_f32 {r_exact} is not f32 grade"); + + let mut split = Array2::::zeros((m, n)); + matmul_f32_amx_split(a.view(), b.view(), split.view_mut()).expect("split"); + let r_split = rel_max(&split, &expect); + + let mut fast = Array2::::zeros((m, n)); + matmul_f32_bf16_fast(a.view(), b.view(), fast.view_mut()).expect("fast"); + let r_fast = rel_max(&fast, &expect); + + assert!(r_split <= 1e-5, "{m}x{k}x{n}: split {r_split} is not f32 grade"); + assert!( + r_fast > 1e-4, + "{m}x{k}x{n}: single pass {r_fast} is unexpectedly accurate \ + — AMX may have a real f32 path now; re-justify the 3x cost" + ); + assert!( + r_split * 50.0 < r_fast, + "{m}x{k}x{n}: split {r_split} vs fast {r_fast} \ + — less than the ~400x measured; the tail term may not be reaching the kernel" + ); + } } #[test] @@ -1090,4 +1264,56 @@ mod tests { assert!((*v - 4.0).abs() < 1e-4); } } + /// Accuracy + throughput of every f32 GEMM path the fork offers, on + /// one host. `cargo test --release --lib gemm_paths_bench -- --ignored --nocapture`. + #[test] + #[ignore] + fn gemm_paths_bench() { + use std::time::Instant; + let time = |f: &mut dyn FnMut()| { + f(); // warm + let reps = 5; + let t = Instant::now(); + for _ in 0..reps { + f(); + } + t.elapsed().as_secs_f64() * 1e3 / reps as f64 + }; + let have_avx512 = std::is_x86_feature_detected!("avx512f"); + eprintln!("host amx={} avx512f={}", amx_available(), have_avx512); + if !have_avx512 { + eprintln!("no avx512f — skipping the F32x16 arm"); + return; + } + for &sz in &[256usize, 512, 1024] { + let (m, k, n) = (sz, sz, sz); + let (a, b) = irrational_pair(m, k, n); + let expect = ref_matmul_f32(&a, &b); + let av: Vec = a.iter().copied().collect(); + let bv: Vec = b.iter().copied().collect(); + let mut c = vec![0.0f32; m * n]; + + let mut mm_ms = + time(&mut || crate::backend::native::gemm_f32(m, n, k, 1.0, &av, k, &bv, n, 0.0, &mut c, n)); + let r_mm = rel_max(&Array2::from_shape_vec((m, n), c.clone()).unwrap(), &expect); + let mut blk_ms = time(&mut || { + c.iter_mut().for_each(|x| *x = 0.0); + // SAFETY: `#[target_feature(enable = "avx512f")]`; guarded by the + // runtime check below, and this bench only runs when it passes. + unsafe { crate::backend::kernels_avx512::sgemm_blocked(m, n, k, 1.0, &av, k, &bv, n, &mut c, n) } + }); + let r_blk = rel_max(&Array2::from_shape_vec((m, n), c.clone()).unwrap(), &expect); + + let mut out = Array2::::zeros((m, n)); + let mut fast_ms = time(&mut || matmul_f32_bf16_fast(a.view(), b.view(), out.view_mut()).unwrap()); + let r_fast = rel_max(&out, &expect); + let mut split_ms = time(&mut || matmul_f32_amx_split(a.view(), b.view(), out.view_mut()).unwrap()); + let r_split = rel_max(&out, &expect); + + for v in [&mut mm_ms, &mut blk_ms, &mut fast_ms, &mut split_ms] { + *v = (*v * 100.0).round() / 100.0; + } + eprintln!("{sz}^3 matrixmultiply {mm_ms:>8.2}ms rel={r_mm:.1e} | F32x16 sgemm_blocked {blk_ms:>8.2}ms rel={r_blk:.1e} | AMX 1-pass {fast_ms:>8.2}ms rel={r_fast:.1e} | AMX 3-pass {split_ms:>8.2}ms rel={r_split:.1e}"); + } + } } From 7b21f6cfa26ddca0c3adc59597011981aa1b93fe Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 5 Sep 2026 21:38:42 +0000 Subject: [PATCH 2/4] plans: gemm-ternlog-mask-consolidation-v1 (DRAFT) + blackboard entry for #303 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Source-first plan for the operator ask "consolidate gemm, ternlogq chaining and cached mask reuse". Inventory: 54 gemm/matmul entry points across 12 files, 4 unified. Three seams with one rule each; six frozen decisions; a measure-first Wave 0 of six probes; nothing in W1+ built. Cached-mask reuse is consumed from lance-graph-java's mask-risc-lowering-v1 (v4.2), not re-planned — this plan's one ask of it is a carving kind for a compacted GEMM row-index list (D-GTM-6). Blackboard records the #303 finding, the bench that decided it, and the loose ends (burn's amx-f32 gate, blas_level3's empty grep, the bf16_tile_gemm_16x16 name duplicate). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01EfrUJH3UNnv5NpDH4jDGHq --- .claude/blackboard.md | 31 ++- .../gemm-ternlog-mask-consolidation-v1.md | 180 ++++++++++++++++++ 2 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 .claude/plans/gemm-ternlog-mask-consolidation-v1.md diff --git a/.claude/blackboard.md b/.claude/blackboard.md index a73b019d..c1a5c3cc 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -3,7 +3,36 @@ > **Read this first.** The "Polyglot Notebook" architecture below is a > separate/older program, not the current epoch. -## 2026-09-04 (latest) — W1.5 signature primitives: PR #293/#294/#295 landed, no board entry until now +## 2026-09-05 — AMX f32 GEMM was silently bf16; `matmul_f32` made exact; consolidation plan filed + +**Finding (measured, PR #303):** `hpc::amx_matmul::matmul_f32` downcast both operands +to BF16 once and ran `TDPBF16PS` — ~1e-3 relative error under an f32 name, `Ok(())` +returned. Found via burn (burn#9): 50 linalg tests failed (25 qr / 13 lu / 7 svd / +3 det / 1 attention) vs 1826/1826 on the exact path. The one-rounding + f32-accumulate +discipline was correctly implemented and still insufficient: single rounding discards +16 of 24 significand bits. The existing test used `(i+j)*0.5` / `(i*3+j)*0.25` inputs — +all exactly BF16-representable — at 1% tolerance, so it could neither see nor fail on it. + +**Decision:** `matmul_f32` delegates to `backend::native::gemm_f32` (exact, every host). +Both AMX-f32 variants (`matmul_f32_amx_split` 3-pass hi/lo split, `matmul_f32_bf16_fast` +1-pass) kept as NAMED opt-ins carrying the bench table. Why not the split: measured +`gemm_paths_bench` — AMX loses on BOTH axes at every size (1024³: F32x16 `sgemm_blocked` +27.3 ms / 1.4e-6; matrixmultiply 28.8 ms; AMX 3-pass 159 ms / 1.5e-6; AMX 1-pass +45.9 ms / 1.6e-4). AMX is a BF16/INT8 unit; `matmul_bf16_to_f32` / `matmul_i8_to_i32` +untouched. + +**Plan filed:** `.claude/plans/gemm-ternlog-mask-consolidation-v1.md` (DRAFT v1) — +54 `gemm|matmul` entry points across 12 files, 4 unified; one facade per dtype +(D-GTM-F3), F32x16 `sgemm_blocked` as the exact default (D-GTM-F1), ternlog chaining +at T1 feeding a compacted-index GEMM prefilter (D-GTM-F4). Cached-mask reuse is +CONSUMED from lance-graph-java `mask-risc-lowering-v1` (v4.2), not re-planned. W0 is +six measurement probes; nothing in W1+ is built. + +**Loose ends:** burn's `amx-f32` feature becomes dead once #303 merges — delete it +(burn#9 follow-up). `hpc/blas_level3.rs` shows zero `pub fn` in the inventory grep +(D-GTM-0b). `simd_ops.rs:587` duplicates `hpc/bf16_tile_gemm.rs:45` by name (D-GTM-0a). + +## 2026-09-04 — W1.5 signature primitives: PR #293/#294/#295 landed, no board entry until now Three merged PRs closing W1.5 signature-kernel work items went unrecorded on this blackboard — corrected here. diff --git a/.claude/plans/gemm-ternlog-mask-consolidation-v1.md b/.claude/plans/gemm-ternlog-mask-consolidation-v1.md new file mode 100644 index 00000000..98d0c772 --- /dev/null +++ b/.claude/plans/gemm-ternlog-mask-consolidation-v1.md @@ -0,0 +1,180 @@ +# gemm-ternlog-mask-consolidation-v1 — one GEMM entry per dtype, masks as the prefilter, ternlog as the mask ALU + +> **Status:** DRAFT v1 (2026-09-05). Source-first: every "exists" row cites `file:line` +> at ndarray `claude/great-curie-d2ufyl` HEAD (PR #303 + this doc). Every "proposed" +> row carries a falsifier. Nothing in §5 is built. No kernel, no ABI symbol. +> +> **Operator ask (2026-09-05, verbatim in spirit):** *consolidate gemm, ternlogq chaining, +> and cached mask reuse; MKL and gemm and blasgraph are all reverse-engineered; MKL could +> benefit from ternlogq.* Plus the W0 evidence that started it: the fork's AMX f32 GEMM +> silently computed in bf16 (burn#9 → ndarray#303), and the benchmark that decided the +> fix showed **AMX loses to the hand-rolled F32x16 kernel on both axes at every size**. +> +> **Scope fence, read first.** *Cached mask reuse* already has an authoritative, +> council-corrected owner: lance-graph-java `.claude/plans/mask-risc-lowering-v1.md` +> (v4.2 — the memo, rail-group factoring, voxelmasking, "focus of awareness are cached +> masks"). This plan does NOT re-plan it. It consumes its decisions (§2.3) and adds the +> two things nobody owns: **(a)** the GEMM surface itself, and **(b)** the seam where a +> cached mask becomes a GEMM row/col prefilter — the only place "ternlogq chaining" and +> "GEMM" actually meet in code today (`hpc/prefilter.rs:189`). + +## §0 — The model, in one picture + +``` + T2 where / hop / plan_eval Mask × WideFieldMask → Mask (lgj-abi exports) + │ cached masks (mask-risc-lowering-v1: memo, rail groups) + ▼ + T1 ternlog chain m = ternlog(a, b, c) … one VPTERNLOGQ per 512 bits + │ (ndarray::simd::mask_ternlog_assign; lgj kernels::simd_mask_ternlog_assign) + ▼ + T1 mask → GEMM prefilter pruned_gemm_rows(mask, …) (hpc/prefilter.rs) + │ + ▼ + T0/T1 ONE gemm per dtype gemm_f32 / gemm_f64 / gemm_bf16 / gemm_i8 (backend::mod.rs) + ↓ dispatch: F32x16 sgemm_blocked | matrixmultiply | AMX(bf16,i8 only) +``` + +Three claims, each falsifiable in §6: +1. **There is one GEMM per dtype at the facade; everything else is a backend.** Today there are ~50 `pub fn *gemm*|*matmul*` across 12 files (§1.1). +2. **A mask is the prefilter, and it arrives already chained.** The T2 tier hands T1 a mask that ternlog has already conjoined; GEMM never recomputes a predicate. +3. **AMX is a BF16/INT8 unit, not an f32 unit.** Measured (§1.3). No f32 path routes through it by default, ever. + +## §1 — What EXISTS (documentation register — every row cites source) + +### §1.1 — The GEMM inventory (the consolidation target) + +`grep -rn 'pub fn [a-z_0-9]*\(gemm\|matmul\)' src/` at HEAD, grouped by what they are: + +| group | entry points | file:line | role today | +|---|---|---|---| +| **Facade, unified** | `cblas_sgemm`, `cblas_dgemm`, `gemm_i8`, `gemm_bf16` | `backend/mod.rs:149,158,186,211` | the ONLY dtype-unified surface; `cblas_gemm_s8s8s32`, `cblas_gemm_bf16bf16f32` at `:268,:274` are aliases | +| Native f32/f64 | `gemm_f32`, `gemm_f64`, `gemm_f32_tiled`, `sgemm_nr/mr`, `dgemm_nr/mr` | `backend/native.rs:210,260,425,50-77` | `gemm_f32` = `matrixmultiply::sgemm` (crates.io), NOT the hand-rolled kernel | +| **Hand-rolled F32x16** | `sgemm_blocked`, `dgemm_blocked` | `backend/kernels_avx512.rs:665,833` | packed panels MR=6/NR=16/KC=256/MC=72/NC=256 (`:538-542`); **fastest exact path measured**; `pub(crate)`, unreachable from outside | +| Hand-rolled AVX2 twin | `sgemm_blocked`, `dgemm_blocked` | `simd_avx2.rs:462,479` | second copy, portable backend | +| Reverse-engineered "MKL" | `gemm_f32`, `gemm_f64`, `sgemm`, `dgemm`, `sgemm_bf16`, `sgemm_int8` | `backend/mkl.rs:195,219,384,429,479,531` | CBLAS-shaped API names over pure Rust; feature `intel-mkl` | +| OpenBLAS-shaped | `gemm_f32`, `gemm_f64` | `backend/openblas.rs:90,116` | feature `openblas` | +| AMX (ArrayView API) | `matmul_bf16_to_f32`, `matmul_i8_to_i32`, `matmul_f32`, `matmul_f32_bf16_fast`, `matmul_f32_amx_split` | `hpc/amx_matmul.rs:531,887,781,748,810` | post-#303: `matmul_f32` is exact (delegates to native); the two AMX-f32 variants are named opt-ins with the bench table in their docs | +| AMX tiles | `bf16_tile_gemm_16x16`, `_packed`, `_tier` | `hpc/bf16_tile_gemm.rs:45,203,66` | the TDPBF16PS 16×16 tile | +| AMX tiles (dup) | `bf16_tile_gemm_16x16` | `simd_ops.rs:587` | **duplicate name of the above** — verify same body (W0) | +| INT8 | `int8_tile_gemm_16x16`, `int8_gemm_amx_tiled`, `int8_gemm_vnni`, `gemm_u8_i8` | `hpc/int8_tile_gemm.rs:47,358`, `hpc/vnni_gemm.rs:46`, `simd_int_ops.rs:253` | AMX → VNNI zmm → VNNI ymm → scalar | +| Quantized refs | `bf16_gemm_f32`, `mixed_precision_gemm`, `int8_gemm_i32`, `int8_gemm_f32`, `int8_gemm_per_channel_f32` | `hpc/quantized.rs:444,484,618,633,655` | scalar references + per-channel scaling | +| f64 tiled | `gemm_f64_tiled`, `gemm_f64_tiled_fma` | `simd_ops.rs:952,1003` | separate f64 blocking, not `dgemm_blocked` | +| Batched | `batched_gemm_f32`, `batched_gemm_4d_f32` | `hpc/linalg/batched.rs:60,124` | loops over 2-D GEMM | +| **Mask-prefiltered** | `pruned_gemm_rows` | `hpc/prefilter.rs:189` | the ONLY existing mask→GEMM bridge | +| Runtime re-exports | `matmul_f32`, `matmul_bf16_to_f32`, `matmul_i8_to_i32` (×2), `gemm_u8_i8` | `simd_runtime/matmul.rs:43,33,80,193,227` | feature `runtime-dispatch` mirror of the AMX API | +| Misc | `matmul_vec`, `matmul_i8_to_i32_wasm` | `hpc/models/layers.rs:174`, `simd_wasm.rs:1474` | model layer GEMV; wasm arm | + +**Counted, not estimated:** 54 entry points, 12 files, **4 unified.** `hpc/blas_level3.rs` (named in CLAUDE.md as "BLAS L3 gemm/syrk/trsm/symm") returned **zero** `pub fn` hits in this grep — W0 must resolve whether that module is empty, macro-generated, or misnamed. + +### §1.2 — ternlog: already a T1 primitive, already chained once + +- Facade: `U64x8::ternlog` / `U32x16::ternlog` — `simd_avx2.rs:3559,3618` (portable backend; W1a-#9), with the truth-table constants module `simd.rs:570` (`AND3 0x80, AND2_ANDNOT 0x40, AND_ANDNOT2 0x10, OR2_AND 0xA8, XOR3 0x96, MAJ3 0xE8, AND2 0xC0, OR3 0xFE`). `mask_ternlog` / `mask_ternlog_assign` re-exported at `simd.rs:733`. +- lgj-abi consumes it correctly by tier: `kernels.rs:100` `pub use ndarray::simd::ternlog;`, `kernels.rs:111` `simd_mask_ternlog_assign` → `ndarray::simd::mask_ternlog_assign::`. `exports.rs` names `kernels::ternlog::AND3`, never `ndarray::simd` (membrane-tiers.md:37). +- **It is already the hop conjunction:** lgj LATEST_STATE:10-13 — the two-AND conjunction became ONE `mask_ternlog_assign::` pass, one `VPTERNLOGQ` per 512 bits. +- **And it is NOT where the time goes:** LATEST_STATE:27-29 — "the ternlog wire is one of three mask passes and cannot account for 5×"; the bulk was `eq_u32_strided_to_mask` at `stride_bytes == 4`. This plan inherits that measurement: **ternlog chaining is a correctness/shape win first, a speed win only where a measurement says so.** +- jitson already whitelists `vpternlogd/vpternlogq` (`hpc/jitson/validator.rs:39,327`) — the JIT tier can emit it. + +### §1.3 — The W0 evidence this plan is built on (measured 2026-09-05, PR #303) + +`gemm_paths_bench` (`hpc/amx_matmul.rs`, `#[ignore]`, `--release`), Xeon w/ AMX + AVX-512, square f32, irrational inputs: + +| size | `native::gemm_f32` (matrixmultiply) | **F32x16 `sgemm_blocked`** | AMX 3-pass split | AMX 1-pass bf16 | +|---|---|---|---|---| +| 256³ | 0.42 ms · 2.9e-7 | **0.37 ms · 2.9e-7** | 2.78 ms · 1.4e-6 | 0.77 ms · 3.7e-4 | +| 512³ | 3.51 ms · 1.2e-6 | **3.22 ms · 1.2e-6** | 20.9 ms · 1.4e-6 | 6.37 ms · 2.2e-4 | +| 1024³ | 28.8 ms · 1.4e-6 | **27.3 ms · 1.4e-6** | 159 ms · 1.5e-6 | 45.9 ms · 1.6e-4 | + +Three consequences, all frozen in §4: the hand-rolled kernel is the exact default; AMX is BF16/INT8-only; the "one rounding + f32 accumulate" discipline was correctly implemented in the AMX path and was still insufficient for an f32 contract — the missing half was operand splitting, and even that half loses to not using AMX. + +### §1.4 — Mask caching: owned elsewhere, consumed here + +- lgj-abi `registry.rs:842-849` — `set_cached_carving` / `cached_carving` on a mask (per-generation, invalidated on close). +- Java `Mask.java:43,185-191` — the packed-bit word lane is resolved once and cached, re-validated per facade call (the `epoch` recheck; `ISS-LGJ-CACHED-DESCRIPTOR-CROSS-THREAD-WINDOW` still open). +- `NativePattern.java:54` — `scratchMask` reused as the destination of terminal ops. +- **The plan that owns all of this:** `lance-graph-java/.claude/plans/mask-risc-lowering-v1.md` v4.2 — §5 rail-group factoring (why the memo MAY hit), §14 voxelmasking (the vertical axis is enumerated, not cached), §3c reuse map (OSM native; weather splits). Its frozen decisions and gates D-MRL-* are **inputs** here, never re-decided. + +## §2 — Three seams, one rule each + +### §2.1 — GEMM facade rule: **one name per (dtype-in, dtype-out); backends are `pub(crate)`** + +`backend::mod.rs:149-215` is already that shape for four signatures. The rule makes it total: every row in §1.1 that is not one of `gemm_f32 / gemm_f64 / gemm_bf16 / gemm_i8` (+ the batched wrappers) either (a) becomes the backend one of those dispatches to, (b) is a named opt-in with a measured reason to exist, or (c) is deleted after W0 proves it a duplicate. The `simd_runtime/matmul.rs` mirror is (a). The MKL/OpenBLAS-shaped modules are (a) behind their features. The two AMX-f32 variants are (b), already documented. `simd_ops.rs:587` vs `hpc/bf16_tile_gemm.rs:45` is (c) pending W0. + +### §2.2 — ternlog chaining rule: **a predicate is composed at T2, lowered to ONE ternlog chain at T1, and never re-evaluated below** + +A 3-input `IMM` is the whole truth table (`simd.rs:561-563`), so any Boolean over three masks is one pass; an n-input predicate is ⌈(n−1)/2⌉ chained passes. The chain is *emitted* by the T2 planner (mask-risc-lowering-v1 Wave 1/2), *executed* by `mask_ternlog_assign`, and its RESULT is what reaches the GEMM prefilter. GEMM never sees a predicate, only a mask. This is the "chaining" the operator named; the *cache* of the intermediate masks is mask-risc-lowering-v1's memo (§1.4), not ours. + +**Where MKL benefits (the operator's hunch, scoped honestly):** in a GEMM, ternlog pays at the *edges* — the K-tail / N-tail lane masks, sign/abs prologues, and beta==0/1 select — not in the FMA core, which is arithmetic. In the reverse-engineered `backend/mkl.rs` the same tail handling appears per-kernel (`sgemm`, `sgemm_bf16`, `sgemm_int8`). Consolidating those tails onto `mask_ternlog` is a shape win (one tail idiom) and *possibly* a speed win; W0-3 measures before anyone claims it. + +### §2.3 — Mask→GEMM seam rule: **the mask is the row/col fetch list; the prefilter is bulk, not per-row** + +`hpc/prefilter.rs:189 pruned_gemm_rows` is the seed. The rule from mask-risc-lowering-v1 §14.7 ("the mask IS the fetch list") applies verbatim: a cached mask selects which rows of A (or cols of B) enter the packed panel; the pack loop consumes the mask by word (`popcount` + `pdep`-style expansion or a compacted index list built ONCE per mask generation), never by testing bits inside the micro-kernel. The mask's `cached_carving` (`registry.rs:842`) is the natural home for that compacted index list — **that is the one place this plan asks mask-risc-lowering-v1 for something**: a carving slot whose payload is a GEMM row-index list, keyed by mask generation. + +## §3 — Non-goals (each with why) + +- **Not** re-deciding anything in mask-risc-lowering-v1 (memo policy, rail groups, voxelmasking). Consumed as inputs. +- **Not** an f32 AMX path of any kind. Measured; §1.3. Re-open only with a new measurement on a host with a true f32 tile op. +- **Not** touching `blasgraph` (lance-graph) — that is HDR/16384-bit semiring algebra with its own sparse formats and (per its own plan doc) an *unwritten* tropical GEMM. It is a consumer of §2.1's facade once one exists, not part of the consolidation. +- **Not** a new SIMD backend, a new crate, or any `#[cfg(target_arch)]` above T1 (simd-savant / kernel-membrane-warden). +- **Not** `array_windows` for GEMM: sliding windows overlap; packed panels do not. Recorded so it is not re-proposed. + +## §4 — Frozen decisions (the council attacks these, not the prose) + +| id | decision | evidence | +|---|---|---| +| D-GTM-F1 | `gemm_f32` default = hand-rolled F32x16 `sgemm_blocked` on `avx512f` hosts, `matrixmultiply` otherwise. Exact on both. | §1.3: 0.37 vs 0.42 ms, identical error | +| D-GTM-F2 | AMX serves `gemm_bf16` and `gemm_i8` ONLY. Any f32 AMX path is a named opt-in carrying its measured table. | §1.3; ndarray#303 | +| D-GTM-F3 | The facade is `backend::{gemm_f32, gemm_f64, gemm_bf16, gemm_i8}` (+ batched). Every other GEMM symbol is `pub(crate)`, a documented opt-in, or deleted. | §1.1 count 54→4 | +| D-GTM-F4 | Mask→GEMM prefilter consumes a compacted row-index list built once per mask generation, stored as a mask carving. Never bit-tests inside the micro-kernel. | §2.3; mask-risc §14.7 | +| D-GTM-F5 | Every accuracy test in this surface uses inputs whose significands exceed 8 bits, and tolerances at f32 grade (1e-5) for f32 APIs. | the vacuous `(i+j)*0.5` test that hid the bf16 loss (#303) | +| D-GTM-F6 | Every new kernel lands with a two-sided pin: the fast path must beat the reference by a stated factor AND the reference must still be measurably slower — so a regression in either direction fails. | `three_pass_split_beats_one_bf16_pass` pattern | + +## §5 — What is PROPOSED (plan register), substrate-first + +### Wave 0 — measure before minting (no production code; all probes `#[ignore]` + `--release`) + +| id | probe | falsifier / gate | +|---|---|---| +| D-GTM-0a | Diff `simd_ops.rs:587 bf16_tile_gemm_16x16` vs `hpc/bf16_tile_gemm.rs:45`; diff `simd_avx2.rs:462 sgemm_blocked` vs `kernels_avx512.rs:665`. | byte-identical bodies → delete one; divergent → record which is canonical and why | +| D-GTM-0b | Resolve `hpc/blas_level3.rs`: what does it export, if anything? | zero `pub fn` → either delete the CLAUDE.md claim or find the macro | +| D-GTM-0c | Extend `gemm_paths_bench` to f64 (`dgemm_blocked` vs `gemm_f64_tiled_fma` vs matrixmultiply) and to non-square / K-tail shapes (e.g. 1000×1000×1000, 17×33×15 scaled). | if `sgemm_blocked` loses on any tail shape, D-GTM-F1 gains a shape guard, not a revert | +| D-GTM-0d | **The MKL-ternlog hunch:** in `backend/mkl.rs` `sgemm` (`:384`), time the tail-lane handling with the current idiom vs one `mask_ternlog` select, K∈{255,257,1023,1025}. | < 3% end-to-end → record as shape-only win, no speed claim; ≥ 3% → W1 item | +| D-GTM-0e | `pruned_gemm_rows` (`prefilter.rs:189`): measure per-row bit-test vs compacted-index pack at 10%/50%/90% mask density. | the crossover density decides D-GTM-F4's threshold, or proves compaction always wins | +| D-GTM-0f | Count real call sites of every §1.1 symbol across ndarray, lance-graph, lance-graph-java, burn (`grep -rn`). | symbols with zero external callers are `pub(crate)` candidates for W1 with no consumer wave | + +### Wave 1 — ndarray (the facade and the backends) + +- D-GTM-1: `backend::gemm_f32` → runtime `avx512f` check → `sgemm_blocked` (one `unsafe` call with a SAFETY comment; the `#[target_feature]` boundary stays inside `kernels_avx512.rs`), else matrixmultiply. `dgemm` per D-GTM-0c. Two-sided pin per D-GTM-F6. +- D-GTM-2: demote every non-facade GEMM symbol per D-GTM-F3, in one commit per file, each with the W0-0f caller count in its message. `simd_runtime/matmul.rs` becomes a thin re-export of the facade. +- D-GTM-3: `hpc/amx_matmul::matmul_f32` (the ArrayView API) delegates to `backend::gemm_f32` — already does post-#303; make the delegation explicit in the doc and drop `pack_contig` when the views are already contiguous. +- D-GTM-4 (gated on D-GTM-0d ≥ 3%): tail-lane selects in `mkl.rs` `sgemm`/`sgemm_bf16`/`sgemm_int8` become `mask_ternlog` idioms. + +### Wave 2 — the mask→GEMM seam (ndarray + lgj-abi, one PR each) + +- D-GTM-5 (ndarray): `pruned_gemm_rows` takes a `&[u32]` compacted index list, not a mask; a `mask_to_row_indices(&[u64]) -> Vec` T1 primitive builds it (popcount + expand, one pass). Density crossover per D-GTM-0e. +- D-GTM-6 (lgj-abi, **filed against mask-risc-lowering-v1, not built here**): a carving kind whose payload is that index list, keyed by mask generation, invalidated with the mask. This is the single ask of the other plan. + +### Wave 3 — consumers (last, by the STOP rule) + +- burn-ndarray `simd_dispatch::matmul_2d` → `backend::gemm_f32` once D-GTM-1 lands; delete the `amx-f32` feature (burn#9 follow-up). +- blasgraph / bgz-tensor / jc: no change unless W0-0f finds a caller of a demoted symbol. + +## §6 — Pre-registered gates (decided before any worker runs) + +- G1: after W1, `grep -rn 'pub fn [a-z_0-9]*\(gemm\|matmul\)' src/` outside `backend/mod.rs` + batched + documented opt-ins returns **0**. Falsifier: the grep. +- G2: `gemm_paths_bench` f32 default row is `sgemm_blocked`'s number ±5% on the same host; error ≤ 1e-5 rel at every shape. Falsifier: the bench. +- G3: burn tensor suite stays 1826/1826 against the W1 fork. Falsifier: `cargo test -p burn-backend-tests --features ndarray --test tensor`. +- G4: no new `#[cfg(target_arch)]`, `_mm*`, or `core::arch` above `simd_*.rs` / `kernels_avx512.rs`. Falsifier: `simd-savant` grep. +- G5: every new test passes D-GTM-F5's input rule — a reviewer can grep the fixture for `* 0.5` / `* 0.25` / `as f32` on small ints and find none in an accuracy test. + +## §7 — Worker allocation (by role, per the workspace model policy) + +- **Opus:** this plan; W0 result synthesis (D-GTM-0a..0f read together decide W1's shape); every disable-run adjudication; central `fmt`/`clippy -D warnings`/tests; all commits/pushes. +- **Sonnet:** each W1 demotion commit (one file, one shape: "make these symbols `pub(crate)`, cite the caller count"); the f64 bench extension; the `mask_to_row_indices` primitive against a written spec. Edit-only; no `cargo build/check` per `.claude/rules/agent-cargo-hygiene.md`; the orchestrator compiles once. +- **Haiku:** only the guarded-executor role — run a pre-written `#[ignore]` bench card and paste the tail. + +## §8 — Open questions (answered by W0, not by discussion) + +1. Is `sgemm_blocked`'s 12% lead over matrixmultiply stable across shapes and hosts, or an artifact of square 2ⁿ sizes? (D-GTM-0c) +2. Does the reverse-engineered MKL path have any caller today, or is it a second facade with zero consumers? (D-GTM-0f decides whether W1 demotes it or deletes it) +3. At what mask density does compacted-index packing beat per-row bit tests? (D-GTM-0e) +4. Is there a real shape where an f32 AMX path wins? Nothing measured says yes; the opt-ins exist so the question stays cheap to re-ask. From 2230c51b32e424db948c12dd3d0dafc7517bb64d Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 5 Sep 2026 21:44:27 +0000 Subject: [PATCH 3/4] amx_matmul: fall back to exact GEMM when BF16 cannot carry the input; doc examples MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit (#303) found a real hole in both AMX opt-ins: `f32::INFINITY` splits as `h = inf`, `x - h = NaN`, and the three-pass cross term then multiplies `inf` by zero. `f32::MAX` is the finite version — its RNE BF16 head rounds to `inf` (bit-trick `from_f32_rounded`: `0x7F7F` + round + sticky = `0x7F80`). Either way a NaN lands in the product under an `Ok(())`. `bf16_safe` checks every value is finite AND its BF16 head is finite; `matmul_f32_amx_split` and `matmul_f32_bf16_fast` take the tile path only when both operands pass, else the exact CPU reference they already carried. `matmul_f32` is untouched — it never went near BF16. Two-sided test: the premise (`f32::MAX` → BF16 inf) is asserted before anything else, then `inf` and `f32::MAX` poison one cell of an otherwise irrational fixture and both opt-ins must return no NaN and match the reference; ordinary large values (`-3.0e38`) and subnormals must NOT trip the guard. Clippy caught the first fixture: `3.5e38` is above `f32::MAX` — the "representable" example wasn't. Runnable `# Example` blocks on all three public fns (the repo rule); 3/3 doctests pass. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01EfrUJH3UNnv5NpDH4jDGHq --- src/hpc/amx_matmul.rs | 102 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 100 insertions(+), 2 deletions(-) diff --git a/src/hpc/amx_matmul.rs b/src/hpc/amx_matmul.rs index a241f361..6f41b182 100644 --- a/src/hpc/amx_matmul.rs +++ b/src/hpc/amx_matmul.rs @@ -720,6 +720,17 @@ unsafe fn bf16_gemm_vdpbf16ps(a: &[BF16], b: &[BF16], c: &mut [f32], m: usize, n /// computation runs in pure f32 and is bit-stable. /// /// `out` must be row-contiguous; inputs may be strided. +/// `true` iff every value survives a BF16 downcast finite. +/// +/// Two ways to fail: a non-finite input (`inf - inf = NaN` in the tail), and a +/// finite input above BF16's max finite (~3.39e38) whose RNE head rounds to +/// `inf` — `f32::MAX` does. Either poisons the tile product with a NaN cross +/// term, so both AMX opt-ins fall back to the exact CPU GEMM for such inputs. +fn bf16_safe(v: &[f32]) -> bool { + v.iter() + .all(|&x| x.is_finite() && BF16::from_f32_rounded(x).to_f32().is_finite()) +} + /// Splits `v` into a BF16 head and a BF16 tail such that /// `head.to_f32() + tail.to_f32()` reproduces `v` to ~17 significand bits. /// @@ -737,6 +748,24 @@ fn split_bf16(v: &[f32]) -> (Vec, Vec) { (hi, lo) } +/// # Example +/// +/// ``` +/// use ndarray::Array2; +/// use ndarray::hpc::amx_matmul::matmul_f32_bf16_fast; +/// +/// let a = Array2::::from_shape_fn((16, 32), |(i, j)| (i * 32 + j) as f32 * 0.01); +/// let b = Array2::::from_shape_fn((32, 16), |(i, j)| (i + j) as f32 * 0.02); +/// let mut c = Array2::::zeros((16, 16)); +/// matmul_f32_bf16_fast(a.view(), b.view(), c.view_mut()).unwrap(); +/// +/// // With AMX: one BF16 pass, expect ~1e-3 relative error. +/// // Without AMX (or for inputs that would overflow BF16): the exact CPU GEMM. +/// let exact = a.dot(&b); +/// let rel = (&c - &exact).iter().map(|x| x.abs()).fold(0.0f32, f32::max) +/// / exact.iter().map(|x| x.abs()).fold(0.0f32, f32::max); +/// assert!(rel < 5e-3); +/// ``` /// Single-BF16-pass `f32` matmul: one third of [`matmul_f32`]'s tile work, /// at BF16 precision (~1e-3 relative error, measured). /// @@ -753,7 +782,7 @@ pub fn matmul_f32_bf16_fast( let b_f32 = pack_contig(&rhs); let mut c = vec![0.0f32; m * n]; - if amx_available() { + if amx_available() && bf16_safe(&a_f32) && bf16_safe(&b_f32) { let a_bf16: Vec = a_f32.iter().map(|&v| BF16::from_f32_rounded(v)).collect(); let b_bf16: Vec = b_f32.iter().map(|&v| BF16::from_f32_rounded(v)).collect(); bf16_gemm_dispatch(&a_bf16, &b_bf16, &mut c, m, n, k); @@ -772,6 +801,23 @@ pub fn matmul_f32_bf16_fast( Ok(()) } +/// # Example +/// +/// ``` +/// use ndarray::Array2; +/// use ndarray::hpc::amx_matmul::matmul_f32; +/// +/// let a = Array2::::from_shape_fn((4, 8), |(i, j)| (i * 8 + j) as f32 * 0.37); +/// let b = Array2::::from_shape_fn((8, 4), |(i, j)| (i + j) as f32 * 0.11); +/// let mut c = Array2::::zeros((4, 4)); +/// matmul_f32(a.view(), b.view(), c.view_mut()).unwrap(); +/// +/// // Exact on every host — never routed through AMX. +/// let exact = a.dot(&b); +/// for (x, y) in c.iter().zip(exact.iter()) { +/// assert!((x - y).abs() <= 1e-5 * y.abs().max(1.0)); +/// } +/// ``` /// Exact f32 matmul. Never routes through AMX — measured slower AND less /// accurate than the CPU f32 GEMM at every size tried (see the table in /// [`matmul_f32_amx_split`]). @@ -790,6 +836,24 @@ pub fn matmul_f32( Ok(()) } +/// # Example +/// +/// ``` +/// use ndarray::Array2; +/// use ndarray::hpc::amx_matmul::{matmul_f32, matmul_f32_amx_split}; +/// +/// let a = Array2::::from_shape_fn((16, 32), |(i, j)| ((i * 7 + j) % 13) as f32 * 0.3137); +/// let b = Array2::::from_shape_fn((32, 16), |(i, j)| ((i + j * 5) % 11) as f32 * 0.2713); +/// let (mut split, mut exact) = (Array2::::zeros((16, 16)), Array2::::zeros((16, 16))); +/// matmul_f32_amx_split(a.view(), b.view(), split.view_mut()).unwrap(); +/// matmul_f32(a.view(), b.view(), exact.view_mut()).unwrap(); +/// +/// // f32-grade (~1e-6 relative) on AMX — but ~6x SLOWER than `matmul_f32`. +/// // Exists to keep the measurement reproducible, not as a path to reach for. +/// let rel = (&split - &exact).iter().map(|x| x.abs()).fold(0.0f32, f32::max) +/// / exact.iter().map(|x| x.abs()).fold(0.0f32, f32::max); +/// assert!(rel <= 1e-5); +/// ``` /// AMX three-pass BF16-split matmul: f32-grade accuracy on the tile unit. /// /// KEPT FOR THE RECORD, NOT RECOMMENDED. Benchmarked against the exact CPU @@ -816,7 +880,7 @@ pub fn matmul_f32_amx_split( let b_f32 = pack_contig(&rhs); let mut c = vec![0.0f32; m * n]; - if amx_available() { + if amx_available() && bf16_safe(&a_f32) && bf16_safe(&b_f32) { // AMX has no f32 tile op, so f32 operands must reach the tile // kernel as BF16. A single RNE downcast discards 16 of f32's 24 // significand bits — measured ~1e-3 relative error, which is @@ -1316,4 +1380,38 @@ mod tests { eprintln!("{sz}^3 matrixmultiply {mm_ms:>8.2}ms rel={r_mm:.1e} | F32x16 sgemm_blocked {blk_ms:>8.2}ms rel={r_blk:.1e} | AMX 1-pass {fast_ms:>8.2}ms rel={r_fast:.1e} | AMX 3-pass {split_ms:>8.2}ms rel={r_split:.1e}"); } } + /// The guard: non-finite inputs and inputs whose BF16 head overflows must + /// NOT reach the tile path. Two-sided — the first half proves the guard + /// fires (a NaN would otherwise appear in the split's cross term), the + /// second half proves it stays silent on ordinary inputs. + #[test] + fn amx_opt_ins_fall_back_for_inputs_bf16_cannot_carry() { + // Sanity on the premise the guard rests on, independent of AMX: + assert!(!BF16::from_f32_rounded(f32::MAX).to_f32().is_finite(), "f32::MAX must round to BF16 inf"); + assert!(!bf16_safe(&[f32::INFINITY])); + assert!(!bf16_safe(&[f32::MAX])); + assert!(!bf16_safe(&[f32::NAN])); + assert!(bf16_safe(&[1.0, -3.0e38, 1e-40]), "large-but-BF16-representable and subnormal are fine"); + + let (m, k, n) = (16, 32, 16); + for (name, poison) in [("inf", f32::INFINITY), ("f32::MAX", f32::MAX)] { + let (mut a, b) = irrational_pair(m, k, n); + a[[3, 5]] = poison; + let expect = ref_matmul_f32(&a, &b); + type Path = fn(ArrayView2<'_, f32>, ArrayView2<'_, f32>, ArrayViewMut2<'_, f32>) -> Result<(), MatmulError>; + let paths: [(&str, Path); 2] = [("split", matmul_f32_amx_split), ("fast", matmul_f32_bf16_fast)]; + for (which, f) in paths { + let mut out = Array2::::zeros((m, n)); + f(a.view(), b.view(), out.view_mut()).unwrap(); + assert!(out.iter().all(|x| !x.is_nan()), "{which}/{name}: NaN leaked into the result"); + // Fell back to the exact path: bit-for-bit the reference. + for (x, y) in out.iter().zip(expect.iter()) { + assert!( + x.to_bits() == y.to_bits() || (x - y).abs() <= 1e-5 * y.abs().max(1.0), + "{which}/{name}: {x} vs reference {y}" + ); + } + } + } + } } From 318c5f9775f602ec1b5e3b68a0a50abf72e89819 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 5 Sep 2026 21:50:12 +0000 Subject: [PATCH 4/4] gemm_paths_bench: 2048/4096 + deep-K/skinny/wide shapes; plan D-GTM-F1 gains a shape guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Operator objection to the bench: 256/512/1024 are undersized. Extended to 2048³, 4096³ (1.6 s per GEMM, well past cache-resident), 1024×4096×1024, 256×8192×256 and 64×2048×8192. AMX does not converge with size: 1-pass stays 1.75x slower than the F32x16 kernel at 4096³, 3-pass 5x. The conclusion about AMX for f32 stands at every size measured. Size did expose one real thing: on the two skinny rectangles matrixmultiply beats sgemm_blocked by 10-19% — its MC=72/NC=256 blocking is tuned for square-ish operands. D-GTM-F1 now carries a shape guard whose threshold D-GTM-0c measures, rather than a blanket default. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01EfrUJH3UNnv5NpDH4jDGHq --- .../plans/gemm-ternlog-mask-consolidation-v1.md | 16 ++++++++++++++-- src/hpc/amx_matmul.rs | 12 +++++++++--- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/.claude/plans/gemm-ternlog-mask-consolidation-v1.md b/.claude/plans/gemm-ternlog-mask-consolidation-v1.md index 98d0c772..abd7cb86 100644 --- a/.claude/plans/gemm-ternlog-mask-consolidation-v1.md +++ b/.claude/plans/gemm-ternlog-mask-consolidation-v1.md @@ -84,7 +84,19 @@ Three claims, each falsifiable in §6: | 512³ | 3.51 ms · 1.2e-6 | **3.22 ms · 1.2e-6** | 20.9 ms · 1.4e-6 | 6.37 ms · 2.2e-4 | | 1024³ | 28.8 ms · 1.4e-6 | **27.3 ms · 1.4e-6** | 159 ms · 1.5e-6 | 45.9 ms · 1.6e-4 | -Three consequences, all frozen in §4: the hand-rolled kernel is the exact default; AMX is BF16/INT8-only; the "one rounding + f32 accumulate" discipline was correctly implemented in the AMX path and was still insufficient for an f32 contract — the missing half was operand splitting, and even that half loses to not using AMX. +**Re-measured past the cache-resident regime (operator objection: "256/512/1024 are undersized"):** + +| shape | matrixmultiply | F32x16 `sgemm_blocked` | AMX 1-pass | AMX 3-pass | +|---|---|---|---|---| +| 2048³ | 218.7 ms | **216.7 ms** | 371.2 ms · 1.2e-4 | 1099 ms · 2.3e-6 | +| 4096³ | 1737 ms | **1616 ms** | 2826 ms · 1.1e-4 | 8106 ms · 2.9e-6 | +| 1024×4096×1024 (deep K) | 115.0 ms | **109.6 ms** | 206.6 ms | 580.7 ms | +| 256×8192×256 (skinny) | **14.4 ms** | 17.1 ms | 33.2 ms | 99.3 ms | +| 64×2048×8192 (wide N) | **35.6 ms** | 39.3 ms | 148.5 ms | 369.6 ms | + +AMX does not converge with size — 1-pass is 1.75× slower than F32x16 at 4096³, 3-pass 5×. "Undersized" was not hiding an AMX win. What size DID expose: on the two skinny rectangles `matrixmultiply` beats `sgemm_blocked` by 10–19% — the hand-rolled MC=72/NC=256 blocking is tuned for square-ish operands. **D-GTM-F1 therefore needs the shape guard D-GTM-0c anticipated**, not a blanket default. + +Three consequences, all frozen in §4: the hand-rolled kernel is the exact default *for square-ish shapes, with a measured rectangle guard*; AMX is BF16/INT8-only; the "one rounding + f32 accumulate" discipline was correctly implemented in the AMX path and was still insufficient for an f32 contract — the missing half was operand splitting, and even that half loses to not using AMX. ### §1.4 — Mask caching: owned elsewhere, consumed here @@ -121,7 +133,7 @@ A 3-input `IMM` is the whole truth table (`simd.rs:561-563`), so any Boolean ove | id | decision | evidence | |---|---|---| -| D-GTM-F1 | `gemm_f32` default = hand-rolled F32x16 `sgemm_blocked` on `avx512f` hosts, `matrixmultiply` otherwise. Exact on both. | §1.3: 0.37 vs 0.42 ms, identical error | +| D-GTM-F1 | `gemm_f32` default = hand-rolled F32x16 `sgemm_blocked` on `avx512f` hosts for square-ish shapes; `matrixmultiply` for skinny/wide rectangles (threshold from D-GTM-0c) and on other hosts. Exact on both. | §1.3: wins 256³–4096³ (up to 7%); LOSES 10–19% at 256×8192×256 and 64×2048×8192 | | D-GTM-F2 | AMX serves `gemm_bf16` and `gemm_i8` ONLY. Any f32 AMX path is a named opt-in carrying its measured table. | §1.3; ndarray#303 | | D-GTM-F3 | The facade is `backend::{gemm_f32, gemm_f64, gemm_bf16, gemm_i8}` (+ batched). Every other GEMM symbol is `pub(crate)`, a documented opt-in, or deleted. | §1.1 count 54→4 | | D-GTM-F4 | Mask→GEMM prefilter consumes a compacted row-index list built once per mask generation, stored as a mask carving. Never bit-tests inside the micro-kernel. | §2.3; mask-risc §14.7 | diff --git a/src/hpc/amx_matmul.rs b/src/hpc/amx_matmul.rs index 6f41b182..b483f2a3 100644 --- a/src/hpc/amx_matmul.rs +++ b/src/hpc/amx_matmul.rs @@ -1349,8 +1349,14 @@ mod tests { eprintln!("no avx512f — skipping the F32x16 arm"); return; } - for &sz in &[256usize, 512, 1024] { - let (m, k, n) = (sz, sz, sz); + // Square 256..4096 (crosses the L2/L3-resident boundary), plus deep-K + // and wide-N rectangles — small squares alone can flatter packing cost. + let shapes: Vec<(usize, usize, usize)> = [256usize, 512, 1024, 2048, 4096] + .iter() + .map(|&s| (s, s, s)) + .chain([(1024, 4096, 1024), (256, 8192, 256), (64, 2048, 8192)]) + .collect(); + for (m, k, n) in shapes { let (a, b) = irrational_pair(m, k, n); let expect = ref_matmul_f32(&a, &b); let av: Vec = a.iter().copied().collect(); @@ -1377,7 +1383,7 @@ mod tests { for v in [&mut mm_ms, &mut blk_ms, &mut fast_ms, &mut split_ms] { *v = (*v * 100.0).round() / 100.0; } - eprintln!("{sz}^3 matrixmultiply {mm_ms:>8.2}ms rel={r_mm:.1e} | F32x16 sgemm_blocked {blk_ms:>8.2}ms rel={r_blk:.1e} | AMX 1-pass {fast_ms:>8.2}ms rel={r_fast:.1e} | AMX 3-pass {split_ms:>8.2}ms rel={r_split:.1e}"); + eprintln!("{m}x{k}x{n} matrixmultiply {mm_ms:>8.2}ms rel={r_mm:.1e} | F32x16 sgemm_blocked {blk_ms:>8.2}ms rel={r_blk:.1e} | AMX 1-pass {fast_ms:>8.2}ms rel={r_fast:.1e} | AMX 3-pass {split_ms:>8.2}ms rel={r_split:.1e}"); } } /// The guard: non-finite inputs and inputs whose BF16 head overflows must