From 9c333ffeab1171d167cb6a59bd7c3fb174bac116 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 13:40:37 +0000 Subject: [PATCH 1/3] Add AMX-native packed operand formats so the bf16 GEMM kernel packs once, not per tile Adds NdarrayAmxBf16A/NdarrayAmxBf16B, two MMMInputFormat implementations that convert to bf16 (and, for B, VNNI-pack) once at prepare_one time and hand a new kernel (ndarray_avx512_bf16_native_mmm_f32_16x16, packing index 1) owned Rust panels it consumes directly -- no allocation, conversion, or packing inside AddMatMul, unlike the existing kernel from PR #5 which redoes all three on every 16x16 tile call. The new benchmark cases (P0, P1) measure this through the real MatMatMulKer path at two operand-preparation lifetimes and compare against the earlier B1/B2 raw-AMX baselines. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/benches/amx_bf16_gap_decomposition.rs | 82 ++++ linalg/src/x86_64/mmm.rs | 19 + linalg/src/x86_64/mod.rs | 2 + linalg/src/x86_64/ndarray_amx_native_pack.rs | 430 ++++++++++++++++++ linalg/src/x86_64/ndarray_bf16_native_gemm.rs | 301 ++++++++++++ 5 files changed, 834 insertions(+) create mode 100644 linalg/src/x86_64/ndarray_amx_native_pack.rs create mode 100644 linalg/src/x86_64/ndarray_bf16_native_gemm.rs diff --git a/linalg/benches/amx_bf16_gap_decomposition.rs b/linalg/benches/amx_bf16_gap_decomposition.rs index e5d0a44e7f..7bbeee2aae 100644 --- a/linalg/benches/amx_bf16_gap_decomposition.rs +++ b/linalg/benches/amx_bf16_gap_decomposition.rs @@ -21,6 +21,19 @@ // This file adds no new production kernel and touches no packing/plan code -- it is a // standalone harness calling `ndarray::simd::*` directly, bypassing `MatMatMulKer` entirely for // B1/B2. +// +// Two further cases measure the AMX-native packed kernel (`ndarray_bf16_native_gemm.rs` + +// `ndarray_amx_native_pack.rs`) through the real `MatMatMulKer::run` path, at two different +// operand-preparation lifetimes -- these are NOT the same measurement and are compared against +// different B-cases above: +// P0 -- A and B both prepared (`prepare_one`) OUTSIDE the timed loop, matching B1's lifetime. +// Isolates tract's packed-execution abstraction tax over the raw AMX ceiling, with +// neither operand's preparation cost in the timed region. Compare against B1. +// P1 -- B prepared ONCE outside the timed loop and reused every iteration (a persistent +// weight, as in real inference); A's source stays f32 and is prepared (`prepare_one`) +// fresh INSIDE every timed iteration, once per whole matrix (not per tile) -- the same +// "runtime activation" lifetime B2 uses. This is the realistic inference-shaped +// acceptance metric. Compare against B2. use criterion::*; use ndarray::simd::{PackedBf16B, bf16_tile_gemm_16x16_packed, f32_to_bf16_batch_rne}; use std::hint::black_box; @@ -199,6 +212,75 @@ fn gap_decomposition(c: &mut Criterion) { }, ); } + + // ---- P0: AMX-native packed kernel through MatMatMulKer, A and B both prepared + // outside the timed loop -- compare against B1 (raw AMX ceiling, same lifetime). ---- + { + let mmm = tract_linalg::x86_64::mmm::ndarray_amx_native_bf16_mmm_f32_16x16.mmm(); + group.bench_with_input( + BenchmarkId::new("P0_amx_native_prepacked", format!("{m}x{k}x{n}")), + &(m, k, n), + |be, &(m, k, n)| { + let packing = &mmm.packings()[1]; + let a = Tensor::zero::(&[m, k]).unwrap(); + let pa = packing.0.prepare_one(&a, 1, 0).unwrap(); + let b = Tensor::zero::(&[k, n]).unwrap(); + let pb = packing.1.prepare_one(&b, 0, 1).unwrap(); + let mut cc = Tensor::zero::(&[n, m]).unwrap(); + be.iter(|| unsafe { + mmm.run( + m, + n, + &[ + FusedSpec::AddMatMul { + a: AsInputValue::Borrowed(&*pa), + b: AsInputValue::Borrowed(&*pb), + packing: 1, + }, + FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())), + ], + ) + .unwrap() + }); + }, + ); + } + + // ---- P1: AMX-native packed kernel through MatMatMulKer, B prepared once outside the + // timed loop and reused (persistent weight); A re-prepared once per iteration inside + // the timed loop, once for the whole matrix (runtime activation) -- compare against B2 + // (same lifetime split). This is the realistic inference-shaped acceptance metric. ---- + { + let mmm = tract_linalg::x86_64::mmm::ndarray_amx_native_bf16_mmm_f32_16x16.mmm(); + group.bench_with_input( + BenchmarkId::new("P1_amx_native_runtime_a", format!("{m}x{k}x{n}")), + &(m, k, n), + |be, &(m, k, n)| { + let packing = &mmm.packings()[1]; + let a = Tensor::zero::(&[m, k]).unwrap(); + let b = Tensor::zero::(&[k, n]).unwrap(); + let pb = packing.1.prepare_one(&b, 0, 1).unwrap(); + let mut cc = Tensor::zero::(&[n, m]).unwrap(); + be.iter(|| unsafe { + let pa = packing.0.prepare_one(&a, 1, 0).unwrap(); + mmm.run( + m, + n, + &[ + FusedSpec::AddMatMul { + a: AsInputValue::Borrowed(&*pa), + b: AsInputValue::Borrowed(&*pb), + packing: 1, + }, + FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())), + ], + ) + .unwrap(); + black_box(&cc); + }); + }, + ); + } } group.finish(); } diff --git a/linalg/src/x86_64/mmm.rs b/linalg/src/x86_64/mmm.rs index c10c3f1ba4..24a6cdee98 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -167,6 +167,25 @@ MMMRustKernel!(ndarray_gemm::kernel::<16, 8> => ndarray_avx512_mmm_f32_16x8 MMMRustKernel!(ndarray_bf16_gemm::kernel => ndarray_avx512_bf16_mmm_f32_16x16(16, 16) built(cfg!(target_arch = "x86_64")) arch(Some(crate::isa::Arch::X86_64)) isa(X86_64Avx512f) lossy_no_exact_tests(true)); + +// AMX-native sibling of the kernel above: `AddMatMul` consumes operands already converted to +// bf16 and VNNI-packed at `prepare_one` time (`ndarray_amx_native_pack.rs`'s +// `NdarrayAmxBf16A`/`NdarrayAmxBf16B`, packing index 1) instead of redoing that work on every +// tile call -- see `ndarray_bf16_native_gemm.rs`'s module doc. Same non-`inventory::submit!` +// registration as the kernel above, for the same reason: reachable only by direct +// construction, never through `MmmDispatch::native()`. +// +// `lossy_no_exact_tests(true)` for the same reason as above; packing 1's own +// `mmm_packed_packed_tests!` (added by the `packing[1]` clause) is this kernel's auto-generated +// correctness coverage, on top of `ndarray_bf16_native_gemm.rs`'s own `bf16_tolerance` module. +MMMRustKernel!(ndarray_bf16_native_gemm::kernel => ndarray_amx_native_bf16_mmm_f32_16x16(16, 16) + built(cfg!(target_arch = "x86_64")) arch(Some(crate::isa::Arch::X86_64)) isa(X86_64Avx512f) + packing[1] = amx_bf16_native => |k| k.with_packing( + crate::x86_64::ndarray_amx_native_pack::NdarrayAmxBf16A::new(16), + crate::x86_64::ndarray_amx_native_pack::NdarrayAmxBf16B::new(16), + ); + lossy_no_exact_tests(true)); + MMMExternKernel!(x86_64; avx512_mmm_f32_32x6 ( 32, 6)@(512,4) isa(X86_64Avx512f)); MMMExternKernel!(x86_64; avx512_mmm_f32_32x5 ( 32, 5)@(512,4) isa(X86_64Avx512f)); MMMExternKernel!(x86_64; avx512_mmm_f32_48x4 ( 48, 4)@(512,4) isa(X86_64Avx512f)); diff --git a/linalg/src/x86_64/mod.rs b/linalg/src/x86_64/mod.rs index 8399612a76..40edebef6b 100644 --- a/linalg/src/x86_64/mod.rs +++ b/linalg/src/x86_64/mod.rs @@ -1,6 +1,8 @@ pub mod mmm; +mod ndarray_amx_native_pack; mod ndarray_bf16_gemm; +mod ndarray_bf16_native_gemm; mod ndarray_gemm; mod amd_avx512_linear; diff --git a/linalg/src/x86_64/ndarray_amx_native_pack.rs b/linalg/src/x86_64/ndarray_amx_native_pack.rs new file mode 100644 index 0000000000..eb592f31ec --- /dev/null +++ b/linalg/src/x86_64/ndarray_amx_native_pack.rs @@ -0,0 +1,430 @@ +#![allow(clippy::needless_range_loop)] +//! Packing formats that hold the ndarray AMX-native bf16 tile representation +//! built once at `prepare_one`/`prepare_one_view` time, so the kernel body in +//! `ndarray_bf16_native_gemm.rs` never converts or VNNI-packs anything inside +//! `AddMatMul`. +//! +//! Panel geometry mirrors `amx_bf16.rs`'s `PackedAmxBf16A`/`PackedBf16K2` -- +//! r=16, K padded to a multiple of 32 -- but panels are stored as owned Rust +//! values (`Vec` row-major bf16 for A, `ndarray::simd::PackedBf16B` for +//! B) rather than raw bytes in a `Blob`: the consuming kernel calls +//! `ndarray::simd::bf16_tile_gemm_16x16_packed`'s typed API directly, so +//! there is nothing to gain from a byte-blob indirection and every extra +//! layer would be an extra copy. +//! +//! `AMX_BF16_A` is meant for activation-lifetime operands (converted once +//! per matrix, reused across every tile of that matmul); `AMX_BF16_B` is +//! meant for constant-weight-lifetime operands (converted and VNNI-packed +//! once, reused across every matmul that shares the weight). + +use std::alloc::Layout; +use std::fmt::Display; +use std::hash::{Hash, Hasher}; + +use tract_data::internal::*; + +use ndarray::simd::{PackedBf16B, f32_to_bf16_batch_rne}; + +use crate::WeightType; +use crate::frame::mmm::{ + EagerPackedInput, MMMInputFormat, MMMInputValue, PackedExoticFact, PackedMatrixStorage, +}; + +const R: usize = 16; + +fn k_padded(k: usize) -> usize { + k.next_multiple_of(32) +} + +/// Round every f32 element of `tensor` through bf16 (round-to-nearest-even), +/// matching what the packers in this module do at pack time. Non-f32 tensors +/// pass through unchanged. Lets a reference f32 matmul reproduce the +/// kernel's bf16 rounding. +fn simulate_bf16_precision_loss(mut tensor: Tensor) -> TractResult { + if tensor.datum_type() == f32::datum_type() { + let mut plain = tensor.try_as_plain_mut()?; + let slice = plain.as_slice_mut::()?; + let mut bits = vec![0u16; slice.len()]; + f32_to_bf16_batch_rne(slice, &mut bits); + for (v, b) in slice.iter_mut().zip(bits.iter()) { + *v = f32::from_bits((*b as u32) << 16); + } + } + Ok(tensor) +} + +// ───────────────────────────── A: activation lifetime ────────────────────── + +/// Packing format for the AMX-native bf16 kernel's A operand: bf16-converted, +/// panel-native (row-major `[16, k_padded]` per panel) data, built once per +/// `prepare_one`/`prepare_one_view` call. Never carries partially-converted +/// or per-tile state -- the whole matrix is converted in one +/// `f32_to_bf16_batch_rne` pass per panel row at pack time. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub struct NdarrayAmxBf16A { + r: usize, +} + +impl NdarrayAmxBf16A { + pub fn new(r: usize) -> Self { + assert_eq!(r, R, "ndarray's bf16 tile primitive is fixed at r=16"); + NdarrayAmxBf16A { r } + } +} + +impl Display for NdarrayAmxBf16A { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "NdarrayAmxBf16A[{}]", self.r) + } +} + +impl MMMInputFormat for NdarrayAmxBf16A { + fn prepare_tensor(&self, t: &Tensor, k_axis: usize, mn_axis: usize) -> TractResult { + Ok(PackedMatrixStorage::new(self.prepare_one(t, k_axis, mn_axis)?) + .into_tensor(t.datum_type())) + } + + fn prepare_one_view( + &self, + t: &TensorView, + k_axis: usize, + mn_axis: usize, + ) -> TractResult> { + let k = t.shape()[k_axis]; + let mn = t.shape()[mn_axis]; + let kp = k_padded(k); + let panels_count = mn.div_ceil(self.r); + let st = t.strides(); + let (ks, ms) = (st[k_axis], st[mn_axis]); + + let mut panels: Vec> = Vec::with_capacity(panels_count); + let mut row_f32 = vec![0f32; k]; + unsafe { + let src = t.as_ptr_unchecked::(); + for p in 0..panels_count { + let pw = self.r.min(mn - p * self.r); + let mn0 = (p * self.r) as isize; + let mut panel = vec![0u16; self.r * kp]; + for lm in 0..pw { + let srow_base = src.offset((mn0 + lm as isize) * ms); + for kk in 0..k { + row_f32[kk] = *srow_base.offset(kk as isize * ks); + } + f32_to_bf16_batch_rne(&row_f32, &mut panel[lm * kp..lm * kp + k]); + } + panels.push(panel); + } + } + + Ok(Box::new(NdarrayAmxBf16AValue { + fact: PackedExoticFact { format: Box::new(self.clone()), mn: mn.to_dim(), k }, + format: self.clone(), + panels, + k, + mn, + })) + } + + fn k_alignment(&self) -> usize { + 32 + } + + fn r(&self) -> usize { + self.r + } + + fn precursor(&self) -> WeightType { + WeightType::Plain(f32::datum_type()) + } + + fn simulate_precision_loss(&self, tensor: Tensor) -> TractResult { + simulate_bf16_precision_loss(tensor) + } + + fn merge_with<'o, 'a: 'o, 'b: 'o>( + &'a self, + o: &'b dyn MMMInputFormat, + ) -> Option<&'o dyn MMMInputFormat> { + o.downcast_ref::().filter(|x| x.r == self.r).map(|_| self as _) + } + + fn mem_size(&self, k: TDim, mn: TDim) -> TDim { + mn.divceil(self.r) * (self.r * k_padded(k.to_usize().unwrap_or(0)) * 2) + } + + fn extract_at_mn_f16(&self, _: &EagerPackedInput, _: usize, _: &mut [f16]) -> TractResult<()> { + bail!("no f16 extract") + } + + fn extract_at_mn_f32(&self, _: &EagerPackedInput, _: usize, _: &mut [f32]) -> TractResult<()> { + bail!("no f32 extract") + } +} + +/// One prepared A operand: `panels[p]` is row-major bf16 `[16, k_padded]`, +/// ready to feed `ndarray::simd::bf16_tile_gemm_16x16_packed` as `a_bf16` +/// with no further conversion. +#[derive(Clone, Debug)] +pub struct NdarrayAmxBf16AValue { + fact: PackedExoticFact, + format: NdarrayAmxBf16A, + panels: Vec>, + k: usize, + mn: usize, +} + +impl Hash for NdarrayAmxBf16AValue { + fn hash(&self, state: &mut H) { + self.format.hash(state); + self.k.hash(state); + self.mn.hash(state); + self.panels.hash(state); + } +} + +impl PartialEq for NdarrayAmxBf16AValue { + fn eq(&self, other: &Self) -> bool { + self.format == other.format + && self.k == other.k + && self.mn == other.mn + && self.panels == other.panels + } +} +impl Eq for NdarrayAmxBf16AValue {} + +impl Display for NdarrayAmxBf16AValue { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{} value (mn={} k={})", self.format, self.mn, self.k) + } +} + +impl MMMInputValue for NdarrayAmxBf16AValue { + fn format(&self) -> &dyn MMMInputFormat { + &self.format + } + fn scratch_panel_buffer_layout(&self) -> Option { + None + } + fn panel_bytes(&self, i: usize, _buffer: Option<*mut u8>) -> TractResult<*const u8> { + Ok(self.panels[i].as_ptr() as *const u8) + } + fn mn(&self) -> usize { + self.mn + } + fn k(&self) -> usize { + self.k + } + fn exotic_fact(&self) -> &dyn ExoticFact { + &self.fact + } + fn extract_at_mn_f16(&self, _mn: usize, _slice: &mut [f16]) -> TractResult<()> { + bail!("no f16 extract") + } + fn extract_at_mn_f32(&self, _mn: usize, _slice: &mut [f32]) -> TractResult<()> { + bail!("no f32 extract") + } +} + +// ───────────────────────────── B: constant-weight lifetime ───────────────── + +/// Packing format for the AMX-native bf16 kernel's B operand: bf16-converted +/// AND VNNI-packed (`ndarray::simd::PackedBf16B`) once per +/// `prepare_one`/`prepare_one_view` call, reusable across every matmul that +/// shares the packed weight -- the shape a constant weight tensor would be +/// packed into once at model-load time. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub struct NdarrayAmxBf16B { + r: usize, +} + +impl NdarrayAmxBf16B { + pub fn new(r: usize) -> Self { + assert_eq!(r, R, "ndarray's bf16 tile primitive is fixed at r=16"); + NdarrayAmxBf16B { r } + } +} + +impl Display for NdarrayAmxBf16B { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "NdarrayAmxBf16B[{}]", self.r) + } +} + +impl MMMInputFormat for NdarrayAmxBf16B { + fn prepare_tensor(&self, t: &Tensor, k_axis: usize, mn_axis: usize) -> TractResult { + Ok(PackedMatrixStorage::new(self.prepare_one(t, k_axis, mn_axis)?) + .into_tensor(t.datum_type())) + } + + fn prepare_one_view( + &self, + t: &TensorView, + k_axis: usize, + mn_axis: usize, + ) -> TractResult> { + let k = t.shape()[k_axis]; + let mn = t.shape()[mn_axis]; + let kp = k_padded(k); + let panels_count = mn.div_ceil(self.r); + let st = t.strides(); + let (ks, ms) = (st[k_axis], st[mn_axis]); + + let mut panels: Vec = Vec::with_capacity(panels_count); + let mut row_major_bf16 = vec![0u16; kp * self.r]; + let mut row_f32 = vec![0f32; self.r]; + unsafe { + let src = t.as_ptr_unchecked::(); + for p in 0..panels_count { + let pw = self.r.min(mn - p * self.r); + let mn0 = (p * self.r) as isize; + row_major_bf16.fill(0); + for kk in 0..k { + let srow_base = src.offset(kk as isize * ks + mn0 * ms); + for lm in 0..pw { + row_f32[lm] = *srow_base.offset(lm as isize * ms); + } + let mut row_bf16 = [0u16; R]; + f32_to_bf16_batch_rne(&row_f32[..pw], &mut row_bf16[..pw]); + row_major_bf16[kk * self.r..kk * self.r + pw].copy_from_slice(&row_bf16[..pw]); + } + panels.push(PackedBf16B::pack(&row_major_bf16, kp)); + } + } + + Ok(Box::new(NdarrayAmxBf16BValue { + fact: PackedExoticFact { format: Box::new(self.clone()), mn: mn.to_dim(), k }, + format: self.clone(), + panels, + k, + mn, + })) + } + + fn k_alignment(&self) -> usize { + 32 + } + + fn r(&self) -> usize { + self.r + } + + fn precursor(&self) -> WeightType { + WeightType::Plain(f32::datum_type()) + } + + fn simulate_precision_loss(&self, tensor: Tensor) -> TractResult { + simulate_bf16_precision_loss(tensor) + } + + fn merge_with<'o, 'a: 'o, 'b: 'o>( + &'a self, + o: &'b dyn MMMInputFormat, + ) -> Option<&'o dyn MMMInputFormat> { + o.downcast_ref::().filter(|x| x.r == self.r).map(|_| self as _) + } + + fn mem_size(&self, k: TDim, mn: TDim) -> TDim { + mn.divceil(self.r) * (k_padded(k.to_usize().unwrap_or(0)) * self.r * 2) + } + + fn extract_at_mn_f16(&self, _: &EagerPackedInput, _: usize, _: &mut [f16]) -> TractResult<()> { + bail!("no f16 extract") + } + + fn extract_at_mn_f32(&self, _: &EagerPackedInput, _: usize, _: &mut [f32]) -> TractResult<()> { + bail!("no f32 extract") + } +} + +/// One prepared B operand: `panels[p]` is an `ndarray::simd::PackedBf16B` +/// ready to feed `bf16_tile_gemm_16x16_packed` directly, with no further +/// conversion or VNNI packing. +pub struct NdarrayAmxBf16BValue { + fact: PackedExoticFact, + format: NdarrayAmxBf16B, + panels: Vec, + k: usize, + mn: usize, +} + +impl Clone for NdarrayAmxBf16BValue { + fn clone(&self) -> Self { + NdarrayAmxBf16BValue { + fact: self.fact.clone(), + format: self.format.clone(), + panels: self + .panels + .iter() + .map(|p| PackedBf16B::from_le_bytes(p.as_le_bytes(), p.k())) + .collect(), + k: self.k, + mn: self.mn, + } + } +} + +impl std::fmt::Debug for NdarrayAmxBf16BValue { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{} ({} panels)", self.format, self.panels.len()) + } +} + +impl Hash for NdarrayAmxBf16BValue { + fn hash(&self, state: &mut H) { + self.format.hash(state); + self.k.hash(state); + self.mn.hash(state); + for p in &self.panels { + p.data().hash(state); + p.k().hash(state); + } + } +} + +impl PartialEq for NdarrayAmxBf16BValue { + fn eq(&self, other: &Self) -> bool { + self.format == other.format + && self.k == other.k + && self.mn == other.mn + && self.panels.len() == other.panels.len() + && self + .panels + .iter() + .zip(other.panels.iter()) + .all(|(a, b)| a.k() == b.k() && a.data() == b.data()) + } +} +impl Eq for NdarrayAmxBf16BValue {} + +impl Display for NdarrayAmxBf16BValue { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{} value (mn={} k={})", self.format, self.mn, self.k) + } +} + +impl MMMInputValue for NdarrayAmxBf16BValue { + fn format(&self) -> &dyn MMMInputFormat { + &self.format + } + fn scratch_panel_buffer_layout(&self) -> Option { + None + } + fn panel_bytes(&self, i: usize, _buffer: Option<*mut u8>) -> TractResult<*const u8> { + Ok(&self.panels[i] as *const PackedBf16B as *const u8) + } + fn mn(&self) -> usize { + self.mn + } + fn k(&self) -> usize { + self.k + } + fn exotic_fact(&self) -> &dyn ExoticFact { + &self.fact + } + fn extract_at_mn_f16(&self, _mn: usize, _slice: &mut [f16]) -> TractResult<()> { + bail!("no f16 extract") + } + fn extract_at_mn_f32(&self, _mn: usize, _slice: &mut [f32]) -> TractResult<()> { + bail!("no f32 extract") + } +} diff --git a/linalg/src/x86_64/ndarray_bf16_native_gemm.rs b/linalg/src/x86_64/ndarray_bf16_native_gemm.rs new file mode 100644 index 0000000000..fb79d8e217 --- /dev/null +++ b/linalg/src/x86_64/ndarray_bf16_native_gemm.rs @@ -0,0 +1,301 @@ +#![allow(clippy::needless_range_loop)] +//! An f32 GEMM `MatMatMulKer` body whose `AddMatMul` step consumes operands already packed +//! into `ndarray::simd::PackedBf16B` / panel-native bf16 form (`ndarray_amx_native_pack.rs`'s +//! `NdarrayAmxBf16A`/`NdarrayAmxBf16B` packing) and calls +//! `ndarray::simd::bf16_tile_gemm_16x16_packed` directly -- no allocation, no f32->bf16 +//! conversion, and no VNNI packing inside `AddMatMul` itself, unlike `ndarray_bf16_gemm.rs`'s +//! kernel, which redoes all three on every tile call. +//! +//! Registered outside automatic dispatch (see `mmm.rs`'s registration comment for this kernel): +//! reachable only by direct construction, not through `MmmDispatch::native()`. +//! +//! Precision tradeoff is identical to `ndarray_bf16_gemm.rs`'s kernel (see its module doc): +//! operands are truncated to bf16 at pack time, one-time and lossy versus f32; the tile +//! arithmetic itself introduces no further lossiness beyond bf16-precision accumulation order. + +use crate::frame::mmm::FusedKerSpec; +use crate::frame::mmm::OutputStoreKer; + +#[cfg(target_arch = "x86_64")] +use ndarray::simd::{PackedBf16B, bf16_tile_gemm_16x16_packed}; + +macro_rules! scalar { + ($ab: expr, $m: expr, $f: expr) => { + for i in 0..$ab.len() { + for j in 0..$ab[0].len() { + $ab[i][j] = $f($m, $ab[i][j]) + } + } + }; +} + +macro_rules! per_row { + ($ab: expr, $m: expr, $f: expr) => { + for i in 0..$ab.len() { + for j in 0..$ab[0].len() { + $ab[i][j] = $f(*$m.add(i), $ab[i][j]) + } + } + }; +} + +macro_rules! per_col { + ($ab: expr, $m: expr, $f: expr) => { + for i in 0..$ab.len() { + for j in 0..$ab[0].len() { + $ab[i][j] = $f(*$m.add(j), $ab[i][j]) + } + } + }; +} + +const TILE: usize = 16; + +/// `pa` points at an `NdarrayAmxBf16AValue` panel: row-major bf16 `[16, k_padded]`, ready to +/// feed `bf16_tile_gemm_16x16_packed` as `a_bf16` with no further work. `pb` points at one +/// `ndarray::simd::PackedBf16B` value (not a byte blob -- `NdarrayAmxBf16BValue::panel_bytes` +/// hands back a pointer to the `PackedBf16B` itself, cast to `*const u8`), reinterpreted back +/// in place. Both were built once at `prepare_one`/`prepare_one_view` time by +/// `ndarray_amx_native_pack.rs`; this function performs no allocation, no bf16 conversion, and +/// no VNNI packing. +#[cfg(target_arch = "x86_64")] +unsafe fn add_mat_mul_amx_native( + pa: *const u8, + pb: *const u8, + k: usize, + ab: &mut [[f32; TILE]; TILE], +) { + unsafe { + if k == 0 { + return; + } + let k_padded = k.next_multiple_of(32); + let a_bf16 = std::slice::from_raw_parts(pa as *const u16, TILE * k_padded); + let b = &*(pb as *const PackedBf16B); + debug_assert_eq!(b.k(), k_padded); + + let mut c_tile = [0f32; TILE * TILE]; + bf16_tile_gemm_16x16_packed(a_bf16, b, &mut c_tile); + + for i in 0..TILE { + for j in 0..TILE { + ab[i][j] += c_tile[i * TILE + j]; + } + } + } +} + +// See `ndarray_bf16_gemm.rs`'s identical stub for why this exists: `linalg/src/lib.rs` +// compiles this module tree under `feature = "foreign-inventory"` on any host arch, but +// `ndarray` is only a dependency on x86_64. +#[cfg(not(target_arch = "x86_64"))] +unsafe fn add_mat_mul_amx_native( + _pa: *const u8, + _pb: *const u8, + _k: usize, + _ab: &mut [[f32; TILE]; TILE], +) { + unreachable!("ndarray_bf16_native_gemm's kernel is x86_64-only and unbuilt elsewhere") +} + +unsafe fn add_unicast(ab: &mut [[f32; TILE]; TILE], other: &OutputStoreKer) { + unsafe { + for i in 0..TILE { + for j in 0..TILE { + let value: *const f32 = other + .ptr + .offset(other.row_byte_stride * i as isize + other.col_byte_stride * j as isize) + as _; + ab[i][j] += *value; + } + } + } +} + +unsafe fn store(tile: &OutputStoreKer, ab: &[[f32; TILE]; TILE]) { + unsafe { + for i in 0..TILE { + for j in 0..TILE { + let loc: *mut f32 = tile + .ptr + .offset(tile.row_byte_stride * i as isize + tile.col_byte_stride * j as isize) + as _; + *loc = ab[i][j]; + } + } + } +} + +/// The `MatMatMulKer` inner loop. `AddMatMul` only accepts packing index 1 (this kernel's +/// `NdarrayAmxBf16A`/`NdarrayAmxBf16B` packing) -- packing 0 (the framework's default f32 +/// packing) is never a valid call here, since this kernel is only ever reached by direct +/// construction with the native packing explicitly selected. +pub(super) unsafe fn kernel(mut pnl: *const FusedKerSpec) -> isize { + unsafe { + let mut ab = [[0f32; TILE]; TILE]; + loop { + if pnl.is_null() { + break; + } + match *pnl { + FusedKerSpec::Done => break, + FusedKerSpec::Clear => ab = [[0f32; TILE]; TILE], + FusedKerSpec::LoadTile(col_major, _row_major) => { + for row in 0..TILE { + for col in 0..TILE { + ab[row][col] = *col_major.add(col * TILE + row); + } + } + } + FusedKerSpec::ScalarAdd(a) => scalar!(ab, a, |a, b| a + b), + FusedKerSpec::ScalarMul(a) => scalar!(ab, a, |a, b| a * b), + FusedKerSpec::ScalarMin(m) => scalar!(ab, m, |a: f32, b: f32| a.min(b)), + FusedKerSpec::ScalarMax(m) => scalar!(ab, m, |a: f32, b: f32| a.max(b)), + FusedKerSpec::ScalarSub(m) => scalar!(ab, m, |a, b| a - b), + FusedKerSpec::ScalarSubF(m) => scalar!(ab, m, |a, b| b - a), + FusedKerSpec::LeakyRelu(m) => { + scalar!(ab, m, |a, b| if b > 0.0 { b } else { a * b }) + } + FusedKerSpec::PerRowMin(m) => per_row!(ab, m, |a: f32, b: f32| a.min(b)), + FusedKerSpec::PerRowMax(m) => per_row!(ab, m, |a: f32, b: f32| a.max(b)), + FusedKerSpec::PerRowAdd(m) => per_row!(ab, m, |a, b| a + b), + FusedKerSpec::PerRowMul(m) => per_row!(ab, m, |a, b| a * b), + FusedKerSpec::PerRowSub(m) => per_row!(ab, m, |a, b| a - b), + FusedKerSpec::PerRowSubF(m) => per_row!(ab, m, |a, b| b - a), + FusedKerSpec::PerColMin(m) => per_col!(ab, m, |a: f32, b: f32| a.min(b)), + FusedKerSpec::PerColMax(m) => per_col!(ab, m, |a: f32, b: f32| a.max(b)), + FusedKerSpec::PerColAdd(m) => per_col!(ab, m, |a, b| a + b), + FusedKerSpec::PerColMul(m) => per_col!(ab, m, |a, b| a * b), + FusedKerSpec::PerColSub(m) => per_col!(ab, m, |a, b| a - b), + FusedKerSpec::PerColSubF(m) => per_col!(ab, m, |a, b| b - a), + FusedKerSpec::AddRowColProducts(rows, cols) => { + for i in 0..TILE { + for j in 0..TILE { + ab[i][j] += *rows.add(i) * *cols.add(j); + } + } + } + FusedKerSpec::AddUnicast(other) => add_unicast(&mut ab, &other), + FusedKerSpec::ShiftLeft(_) + | FusedKerSpec::RoundingShiftRight(..) + | FusedKerSpec::QScale(..) => { + unreachable!("quantization ops are not reachable on the f32-only packing") + } + FusedKerSpec::AddMatMul { k, pa, pb, packing } => { + assert_eq!(packing, 1, "this kernel only declares packing 1 (AMX-native bf16)"); + add_mat_mul_amx_native(pa, pb, k, &mut ab); + } + FusedKerSpec::Store(tile) => store(&tile, &ab), + }; + pnl = pnl.add(1); + } + } + 0 +} + +#[cfg(all(test, target_arch = "x86_64"))] +mod dispatch_stays_default { + use crate::frame::mmm::{MmmDispatch, Query}; + use tract_data::internal::DatumType; + + /// Same guard as `ndarray_bf16_gemm.rs`'s: this kernel is registered without + /// `inventory::submit!`, so it must never surface through `MmmDispatch::native()`, for + /// both a concrete and a symbolic (`None`) N. + #[test] + fn amx_native_candidate_is_not_reachable_through_automatic_dispatch() { + let dispatch = MmmDispatch::native(); + for n in [Some(32), None] { + let query = Query::plain(DatumType::F32, Some(64), Some(256), n); + let suitable = dispatch.suitable(&query); + assert!( + suitable + .iter() + .all(|(mmm, _, _)| mmm.name() != "ndarray_amx_native_bf16_mmm_f32_16x16"), + "the AMX-native candidate must never appear in automatic dispatch (n={n:?})" + ); + } + } +} + +/// Tolerance-based correctness test, same shape as `ndarray_bf16_gemm.rs`'s `bf16_tolerance` +/// module: this kernel is bf16-precision by construction, run directly against `MatMatMulKer` +/// through the real fused-op path (`AddMatMul` + `Store`) with packing 1 explicitly selected. +#[cfg(all(test, target_arch = "x86_64"))] +mod bf16_tolerance { + use crate::frame::mmm::FusedSpec; + use crate::x86_64::mmm::ndarray_amx_native_bf16_mmm_f32_16x16; + use ndarray::simd::f32_to_bf16_batch_rne; + use tract_data::internal::*; + + fn bf16_exact_value(x: f32) -> f32 { + let mut bits = [0u16; 1]; + f32_to_bf16_batch_rne(&[x], &mut bits); + f32::from_bits((bits[0] as u32) << 16) + } + + #[test] + fn matches_naive_f32_reference_within_bf16_tolerance() { + let (m, k, n) = (32usize, 64usize, 32usize); + let mut a = vec![0f32; m * k]; + let mut b = vec![0f32; k * n]; + for (i, v) in a.iter_mut().enumerate() { + *v = bf16_exact_value(((i % 13) as f32 - 6.0) * 0.5); + } + for (i, v) in b.iter_mut().enumerate() { + *v = bf16_exact_value(((i % 11) as f32 - 5.0) * 0.5); + } + + let mut expected = vec![0f32; m * n]; + for i in 0..m { + for j in 0..n { + let mut acc = 0f32; + for kk in 0..k { + acc += a[i * k + kk] * b[kk * n + j]; + } + expected[i * n + j] = acc; + } + } + + let mmm = ndarray_amx_native_bf16_mmm_f32_16x16.mmm(); + if !mmm.built() || !mmm.runnable() { + eprintln!("skipping: ndarray_amx_native_bf16_mmm_f32_16x16 not runnable on this host"); + return; + } + let packing = &mmm.packings()[1]; + let a_tensor = Tensor::from_shape(&[m, k], &a).unwrap(); + let pa = packing.0.prepare_one(&a_tensor, 1, 0).unwrap(); + let b_tensor = Tensor::from_shape(&[k, n], &b).unwrap(); + let pb = packing.1.prepare_one(&b_tensor, 0, 1).unwrap(); + let mut c = Tensor::zero::(&[n, m]).unwrap(); + + unsafe { + mmm.run( + m, + n, + &[ + FusedSpec::AddMatMul { + a: crate::mmm::AsInputValue::Borrowed(&*pa), + b: crate::mmm::AsInputValue::Borrowed(&*pb), + packing: 1, + }, + FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&c.view_mut())), + ], + ) + .unwrap(); + } + + let got = unsafe { c.as_slice_unchecked::() }; + for i in 0..m { + for j in 0..n { + let e = expected[i * n + j]; + let g = got[j * m + i]; + let tol = 1e-2 * e.abs().max(1.0); + assert!( + (e - g).abs() <= tol, + "mismatch at ({i},{j}): expected {e}, got {g} (tol {tol}), tier={}", + ndarray::simd::bf16_tile_gemm_tier(), + ); + } + } + } +} From 3116f222010713ac3e7593bcb9458e7325bb36ed Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 13:46:58 +0000 Subject: [PATCH 2/3] Gate the AMX-native packing files to x86_64 Same class of bug as PR #5's earlier fix: linalg/src/lib.rs compiles the x86_64 module tree under feature = "foreign-inventory" on any host arch, to enumerate x86_64 kernel names as metadata for cross-compiled builds, but ndarray is only a Cargo dependency on x86_64. ndarray_amx_native_pack.rs and ndarray_bf16_native_gemm.rs use ndarray types throughout rather than in one or two functions, so rather than per-item stubs (this file's other pilot kernels' pattern) both get a whole-module #![cfg(target_arch = "x86_64")] gate, and the mmm.rs registration that references their symbols is gated to match. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/src/x86_64/mmm.rs | 7 +++++++ linalg/src/x86_64/ndarray_amx_native_pack.rs | 1 + linalg/src/x86_64/ndarray_bf16_native_gemm.rs | 1 + 3 files changed, 9 insertions(+) diff --git a/linalg/src/x86_64/mmm.rs b/linalg/src/x86_64/mmm.rs index 24a6cdee98..5a9b1a1e41 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -178,6 +178,13 @@ MMMRustKernel!(ndarray_bf16_gemm::kernel => ndarray_avx512_bf16_mmm_f32_16x16 ndarray_amx_native_bf16_mmm_f32_16x16(16, 16) built(cfg!(target_arch = "x86_64")) arch(Some(crate::isa::Arch::X86_64)) isa(X86_64Avx512f) packing[1] = amx_bf16_native => |k| k.with_packing( diff --git a/linalg/src/x86_64/ndarray_amx_native_pack.rs b/linalg/src/x86_64/ndarray_amx_native_pack.rs index eb592f31ec..3a042b39e2 100644 --- a/linalg/src/x86_64/ndarray_amx_native_pack.rs +++ b/linalg/src/x86_64/ndarray_amx_native_pack.rs @@ -1,3 +1,4 @@ +#![cfg(target_arch = "x86_64")] #![allow(clippy::needless_range_loop)] //! Packing formats that hold the ndarray AMX-native bf16 tile representation //! built once at `prepare_one`/`prepare_one_view` time, so the kernel body in diff --git a/linalg/src/x86_64/ndarray_bf16_native_gemm.rs b/linalg/src/x86_64/ndarray_bf16_native_gemm.rs index fb79d8e217..c3664762cf 100644 --- a/linalg/src/x86_64/ndarray_bf16_native_gemm.rs +++ b/linalg/src/x86_64/ndarray_bf16_native_gemm.rs @@ -1,3 +1,4 @@ +#![cfg(target_arch = "x86_64")] #![allow(clippy::needless_range_loop)] //! An f32 GEMM `MatMatMulKer` body whose `AddMatMul` step consumes operands already packed //! into `ndarray::simd::PackedBf16B` / panel-native bf16 form (`ndarray_amx_native_pack.rs`'s From 2aae5c2ef15af5f1dde4ffcd824be13ee4e728d7 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 14:12:28 +0000 Subject: [PATCH 3/3] Remove section-banner comments Codex review flagged the A/B section dividers as banner comments against the repo's inline-comment style; the doc comments directly above each section already carry the structure. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/src/x86_64/ndarray_amx_native_pack.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/linalg/src/x86_64/ndarray_amx_native_pack.rs b/linalg/src/x86_64/ndarray_amx_native_pack.rs index 3a042b39e2..a3dcfa5c45 100644 --- a/linalg/src/x86_64/ndarray_amx_native_pack.rs +++ b/linalg/src/x86_64/ndarray_amx_native_pack.rs @@ -54,8 +54,6 @@ fn simulate_bf16_precision_loss(mut tensor: Tensor) -> TractResult { Ok(tensor) } -// ───────────────────────────── A: activation lifetime ────────────────────── - /// Packing format for the AMX-native bf16 kernel's A operand: bf16-converted, /// panel-native (row-major `[16, k_padded]` per panel) data, built once per /// `prepare_one`/`prepare_one_view` call. Never carries partially-converted @@ -226,8 +224,6 @@ impl MMMInputValue for NdarrayAmxBf16AValue { } } -// ───────────────────────────── B: constant-weight lifetime ───────────────── - /// Packing format for the AMX-native bf16 kernel's B operand: bf16-converted /// AND VNNI-packed (`ndarray::simd::PackedBf16B`) once per /// `prepare_one`/`prepare_one_view` call, reusable across every matmul that