From 65bd54e9f74ed9c6dd23b7c453ff080fa737607c Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 12:35:31 +0000 Subject: [PATCH 1/5] Add an ndarray-bf16_tile_gemm-backed AVX-512/AMX f32 GEMM kernel as a second additional candidate PR #4's ndarray_avx512_mmm_f32_16x8 went through BlasLevel3::blas_gemm, which allocates and re-packs B on every tile call, so it lost badly to the hand-tuned asm kernel. This pilot tries a different ndarray entry point instead: hpc::bf16_tile_gemm (via ndarray::simd), whose tile primitive takes a pre-packed VNNI B and runs with zero allocation inside, dispatching at runtime to AMX TDPBF16PS, AVX-512 VDPBF16PS, or a decode+FMA polyfill. ndarray_avx512_bf16_mmm_f32_16x16 registers additively at the fixed 16x16 tile geometry the primitive requires, truncates its operands to bf16, and calls bf16_tile_gemm_16x16_packed once per AddMatMul step. That step is still called once per output tile (that's how MatMatMulKer invokes any kernel body), so the per-call pack is real work, not something hoisted above the tile loop -- the module doc spells this out rather than overclaiming a structural fix. The accumulate arithmetic itself is bit-exact across tiers for bf16-exact operands; the actual precision cost is the one-time f32->bf16 truncation of the inputs, which is real and stated plainly, not glossed as approximate math. Tested with a dedicated relative-tolerance test rather than the exact-bit test_mmm_kernel! macros, since those assume f32-exact output. Default dispatch is unchanged (dispatch_stays_default test), and no existing kernel, asm file, or preference is touched. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- Cargo.lock | 20 +- Cargo.toml | 7 + linalg/Cargo.toml | 11 + linalg/benches/ndarray_bf16_gemm.rs | 53 ++++ linalg/src/x86_64/mmm.rs | 10 + linalg/src/x86_64/mod.rs | 2 + linalg/src/x86_64/ndarray_bf16_gemm.rs | 352 +++++++++++++++++++++++++ 7 files changed, 453 insertions(+), 2 deletions(-) create mode 100644 linalg/benches/ndarray_bf16_gemm.rs create mode 100644 linalg/src/x86_64/ndarray_bf16_gemm.rs diff --git a/Cargo.lock b/Cargo.lock index d07531612e..a7cd8d1b5e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1517,6 +1517,13 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fractal" +version = "0.1.0" +dependencies = [ + "libm", +] + [[package]] name = "fs-err" version = "3.3.1" @@ -2499,13 +2506,14 @@ dependencies = [ [[package]] name = "ndarray" version = "0.17.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" dependencies = [ + "fractal", "matrixmultiply", "num-complex", "num-integer", "num-traits", + "p64", + "paste", "portable-atomic", "portable-atomic-util", "rawpointer", @@ -2946,6 +2954,13 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "p64" +version = "0.1.0" +dependencies = [ + "fractal", +] + [[package]] name = "page_size" version = "0.6.0" @@ -4987,6 +5002,7 @@ dependencies = [ "libc", "log", "minijinja", + "ndarray", "nu-ansi-term", "num-traits", "pastey 0.2.3", diff --git a/Cargo.toml b/Cargo.toml index 3cbe58e5c3..c100986bdb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -265,3 +265,10 @@ opt-level = 2 debug = false # strip = "debuginfo" does not work on android and ios incremental = false + +# Pilot v2 (see linalg/src/x86_64/ndarray_bf16_gemm.rs): path-patches the AdaWorldAPI +# ndarray fork so the bf16 tile GEMM candidate kernel can build locally. Not a permanent +# dependency change -- this only resolves in a checkout that has ../ndarray alongside +# tract, matching PR #4's pilot-v1 setup. +[patch.crates-io] +ndarray = { path = "../ndarray" } diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index ef4f3d92b3..dd83e77f16 100644 --- a/linalg/Cargo.toml +++ b/linalg/Cargo.toml @@ -33,6 +33,13 @@ tract-data.workspace = true [target.'cfg(target_arch = "riscv64")'.dependencies] libc.workspace = true +# Pilot v2: one x86_64 GEMM kernel candidate calls the AdaWorldAPI ndarray fork's +# `hpc::bf16_tile_gemm` tile primitives (AMX / AVX-512-VNNI-bf16 / FMA-polyfill tiers) +# through the canonical `ndarray::simd` consumer surface, alongside the existing +# hand-tuned AVX-512 asm kernels and pilot-v1's f32 `blas_gemm` candidate. +[target.'cfg(target_arch = "x86_64")'.dependencies] +ndarray.workspace = true + [build-dependencies] cc.workspace = true half.workspace = true @@ -81,6 +88,10 @@ harness = false name = "mat_vec" harness = false +[[bench]] +name = "ndarray_bf16_gemm" +harness = false + [[bench]] name = "mm_for_wavenet_hw" harness = false diff --git a/linalg/benches/ndarray_bf16_gemm.rs b/linalg/benches/ndarray_bf16_gemm.rs new file mode 100644 index 0000000000..408fa4bb76 --- /dev/null +++ b/linalg/benches/ndarray_bf16_gemm.rs @@ -0,0 +1,53 @@ +// Compares the hand-tuned AVX-512 asm f32 GEMM kernel against the additive bf16-tile-based +// candidate (see linalg/src/x86_64/ndarray_bf16_gemm.rs) on a full matrix multiply -- the +// kernel's own panel-walking machinery loops its tile many times over m/n/k, so this measures +// the whole GEMM each candidate produces, not one microkernel tile call. +use criterion::*; +use tract_data::internal::*; +use tract_linalg::mmm::{AsInputValue, FusedSpec}; + +fn gemm_f32(c: &mut Criterion) { + let mut group = c.benchmark_group("gemm_f32_bf16_pilot"); + for &(m, k, n) in &[(512usize, 512usize, 512usize), (1024, 1024, 1024)] { + group.throughput(Throughput::Elements((2 * m * k * n) as u64)); + for (label, mmm) in [ + ("asm_16x8", tract_linalg::x86_64::mmm::avx512_mmm_f32_16x8.mmm()), + ( + "ndarray_bf16_16x16", + tract_linalg::x86_64::mmm::ndarray_avx512_bf16_mmm_f32_16x16.mmm(), + ), + ] { + group.bench_with_input( + BenchmarkId::new(label, format!("{m}x{k}x{n}")), + &(m, k, n), + |be, &(m, k, n)| { + let packing = &mmm.packings()[0]; + 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: 0, + }, + FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())), + ], + ) + .unwrap() + }); + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, gemm_f32); +criterion_main!(benches); diff --git a/linalg/src/x86_64/mmm.rs b/linalg/src/x86_64/mmm.rs index 3818885d2f..e35a1eefe2 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -127,6 +127,16 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_128x1(128, 1)@(512,4) isa(X86_64Avx MMMExternKernel!(x86_64; avx512_mmm_f32_16x1 ( 16, 1)@(512,4) isa(X86_64Avx512f)); MMMExternKernel!(x86_64; avx512_mmm_f32_16x12( 16,12)@(512,4) isa(X86_64Avx512f)); MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 ( 16, 8)@(512,4) isa(X86_64Avx512f)); + +// Pilot v2: 16x16 tile geometry (matching ndarray's fixed bf16_tile_gemm_16x16 shape). The +// AddMatMul accumulation truncates operands to bf16 and calls the AdaWorldAPI ndarray fork's +// `simd::bf16_tile_gemm_16x16_packed` (AMX / AVX-512-VNNI-bf16 / FMA-polyfill tiers) instead of +// hand-written asm -- see ndarray_bf16_gemm.rs's module doc for the precision tradeoff and the +// honest read on how this compares structurally and numerically to pilot v1's blas_gemm +// candidate. Purely additive: it carries no boost, so retain_best ties it with the asm kernels +// on preference, and every x86_64 dispatch tier below (amd/intel_avx512_linear) still names its +// own asm kernels explicitly and never sees this one. +MMMRustKernel!(x86_64; ndarray_bf16_gemm::kernel => ndarray_avx512_bf16_mmm_f32_16x16(16, 16) isa(X86_64Avx512f)); 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 bd9ca27c8b..a90c2340b3 100644 --- a/linalg/src/x86_64/mod.rs +++ b/linalg/src/x86_64/mod.rs @@ -1,5 +1,7 @@ pub mod mmm; +mod ndarray_bf16_gemm; + mod amd_avx512_linear; mod amd_fma_linear; mod intel_avx512_linear; diff --git a/linalg/src/x86_64/ndarray_bf16_gemm.rs b/linalg/src/x86_64/ndarray_bf16_gemm.rs new file mode 100644 index 0000000000..3b5f02cd13 --- /dev/null +++ b/linalg/src/x86_64/ndarray_bf16_gemm.rs @@ -0,0 +1,352 @@ +#![allow(clippy::needless_range_loop)] +//! An f32 GEMM `MatMatMulKer` body whose `AddMatMul` step truncates its operands to bf16 and +//! calls into the AdaWorldAPI ndarray fork's `simd::bf16_tile_gemm_16x16_packed` tile primitive +//! (AMX `TDPBF16PS` → AVX-512 `VDPBF16PS` → decode+FMA polyfill, selected at runtime) as a +//! second additional candidate alongside the hand-tuned AVX-512 asm kernels and pilot v1's +//! f32-exact `ndarray_avx512_mmm_f32_16x8` (`ndarray_gemm.rs`). +//! +//! **Where this stands relative to pilot v1's structural problem:** `ndarray_gemm.rs`'s +//! `blas_gemm` call allocates a fresh `Array` and re-packs its B operand on *every tile call* +//! via a full BLAS-level3 entry point (allocation + a generic path selection on top of the +//! repack), which was measured 5-10x slower than the hand-tuned asm kernel for exactly that +//! reason. This kernel's `AddMatMul` step is *also* invoked once per output tile (that is how +//! `MatMatMulKer`'s fused-op interpreter calls into any kernel body — one call per (MR, NR) +//! tile, carrying that tile's full K depth), so it still allocates and VNNI-packs its A/B +//! operands **once per tile call**, not once per whole-matrix GEMM — hoisting the pack any +//! further up would mean restructuring the packed-panel format `MatMatMulKer` hands the kernel, +//! which is out of scope for this pilot. What *is* structurally different from pilot v1: the +//! per-call work here is a single `PackedBf16B::pack` (one VNNI interleave over a `k×16` +//! buffer) plus a bf16 truncation pass, calling directly into a tile primitive with no +//! allocation inside — not a generic BLAS-level3 entry point that re-derives packing/path +//! selection from scratch on every call. The measured benchmark numbers below are the honest +//! comparison; see them before assuming this claim translates into a win. +//! +//! **Precision, stated plainly:** the accumulate arithmetic (`C += A·B`) is bit-exact across +//! all three `bf16_tile_gemm` tiers for bf16-exact-integer operands with accumulation below +//! 2^24 (verified by `assert_eq!` parity tests in ndarray's own `hpc::bf16_tile_gemm`), and for +//! general float operands the tiers agree with each other exactly up to accumulation order — +//! this kernel introduces **no additional lossiness of its own**. The precision this kernel +//! trades away versus the native f32 asm kernel is entirely the one-time f32→bf16 truncation of +//! the input operands themselves before they reach any tile primitive: bf16 keeps a 7-bit +//! mantissa against f32's 23-bit, so every element of A and B loses precision at pack time, not +//! merely at accumulation time. This is a real, user-visible precision change for a general +//! inference engine and must not be read as "approximate GEMM" (the arithmetic is not +//! approximate) or as "safe to swap in for f32 workloads" (real model weights are not +//! bf16-exact integers, so the tier-parity bit-exactness above does not extend to them). +//! +//! Goes through `ndarray::simd::*` (`f32_to_bf16_batch_rne`, `PackedBf16B`, +//! `bf16_tile_gemm_16x16_packed`, `bf16_tile_gemm_tier`), the canonical consumer-facing +//! re-export, never `ndarray::hpc::bf16_tile_gemm::*` directly — see the ndarray fork's own +//! `CLAUDE.md` ("all SIMD from `ndarray::simd`"). +//! +//! Tile shape is fixed by the ndarray primitive at M=16, N=16, K a multiple of 32, so this +//! kernel is registered at the matching (16, 16) `MatMatMulKer` geometry rather than reusing +//! pilot v1's 16x8 -- packed A/B panels are padded up to the next multiple of 32 in K with +//! zero rows/columns, which contribute nothing to the accumulation. + +use ndarray::simd::{PackedBf16B, bf16_tile_gemm_16x16_packed, f32_to_bf16_batch_rne}; + +use crate::frame::mmm::FusedKerSpec; +use crate::frame::mmm::OutputStoreKer; + +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` is packed k-major, MR(=16) contiguous per k-step (`pa[ik * 16 + i]`); `pb` is packed +/// k-major NR(=16) contiguous per k-step (`pb[ik * 16 + j]`), which is already row-major +/// `B[K, 16]` -- no transpose needed. `A_panel` (`pa`) is transposed into a small contiguous +/// `(16, K)` buffer, same as pilot v1, because the bf16 tile primitive wants row-major `A[16, K]`. +/// +/// Both operands are truncated to bf16 with `f32_to_bf16_batch_rne` (round-to-nearest-even, +/// the hot-loop-safe path -- never the scalar RNE fn, which is test-only). K is padded up to +/// the next multiple of 32 with zero rows in A and zero rows in B: the padding columns/rows +/// contribute `0 * anything = 0` to every accumulated cell, so the padding is inert. +/// +/// B is packed into VNNI layout via `PackedBf16B::pack`, once per `AddMatMul` call -- i.e. once +/// per (MR, NR) output tile, since that is the granularity `MatMatMulKer` calls a kernel body +/// at. Unlike pilot v1 this pack is a single VNNI interleave straight into the tile primitive +/// (no BLAS-level3 entry point re-deriving packing/dispatch from scratch), but it is still +/// real per-tile allocation and work, not something hoisted above the tile loop. +unsafe fn add_mat_mul_bf16(pa: *const u8, pb: *const u8, k: usize, ab: &mut [[f32; TILE]; TILE]) { + unsafe { + if k == 0 { + return; + } + let a = pa as *const f32; + let b = pb as *const f32; + + let k_padded = k.next_multiple_of(32); + + let mut a_row_major = vec![0f32; TILE * k]; + for i in 0..TILE { + for ik in 0..k { + a_row_major[i * k + ik] = *a.add(ik * TILE + i); + } + } + let mut a_bf16 = vec![0u16; TILE * k_padded]; + for i in 0..TILE { + f32_to_bf16_batch_rne( + &a_row_major[i * k..i * k + k], + &mut a_bf16[i * k_padded..i * k_padded + k], + ); + } + + let b_row_major = std::slice::from_raw_parts(b, k * TILE); + let mut b_bf16 = vec![0u16; k_padded * TILE]; + f32_to_bf16_batch_rne(b_row_major, &mut b_bf16[..k * TILE]); + + let packed_b = PackedBf16B::pack(&b_bf16, k_padded); + + let mut c_tile = [0f32; TILE * TILE]; + bf16_tile_gemm_16x16_packed(&a_bf16, &packed_b, &mut c_tile); + + for i in 0..TILE { + for j in 0..TILE { + ab[i][j] += c_tile[i * TILE + j]; + } + } + } +} + +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, f32-only, one packing (index 0, plain f32×f32), fixed 16x16 +/// tile geometry (the shape `bf16_tile_gemm_16x16_packed` is built for). Same fused-op +/// interpreter shape as `crate::generic::mmm::kernel` and pilot v1's `ndarray_gemm::kernel`; +/// the `AddMatMul` arm is the only place this diverges from the generic reference. +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(..) => { + // Integer-quantization epilogue ops: this kernel only declares an f32 + // accumulator packing, so a caller never reaches these arms. + unreachable!("quantization ops are not reachable on the f32-only packing") + } + FusedKerSpec::AddMatMul { k, pa, pb, packing } => { + assert_eq!(packing, 0, "this kernel only declares packing 0 (f32 x f32)"); + add_mat_mul_bf16(pa, pb, k, &mut ab); + } + FusedKerSpec::Store(tile) => store(&tile, &ab), + }; + pnl = pnl.add(1); + } + } + 0 +} + +#[cfg(test)] +mod dispatch_stays_default { + use crate::frame::mmm::{MmmDispatch, Query}; + use tract_data::internal::DatumType; + + #[test] + fn adding_the_bf16_candidate_does_not_change_default_pick() { + let dispatch = MmmDispatch::native(); + let query = Query::plain(DatumType::F32, Some(64), Some(256), Some(32)); + let suitable = dispatch.suitable(&query); + assert!( + suitable.iter().any(|(mmm, _, _)| mmm.name() == "ndarray_avx512_bf16_mmm_f32_16x16"), + "the new candidate should be suitable wherever avx512f is native" + ); + if let Some((picked, _, _)) = dispatch.pick(&query) { + assert_ne!( + picked.name(), + "ndarray_avx512_bf16_mmm_f32_16x16", + "default dispatch must still prefer the hand-tuned asm kernel" + ); + assert_ne!( + picked.name(), + "ndarray_avx512_mmm_f32_16x8", + "default dispatch must still prefer the hand-tuned asm kernel" + ); + } + } +} + +/// Tolerance-based correctness test. This kernel is inherently bf16-precision (see the module +/// doc): the exact-bit `test_mmm_kernel!` macro family compares kernel output against an f32 +/// reference with `==`/ULP-tight bounds, which this kernel cannot pass by construction, so it +/// gets a dedicated relative-tolerance check instead, run directly against `MatMatMulKer` +/// through the same fused-op path the real dispatcher uses (`AddMatMul` + `Store`), against a +/// naive f32 reference GEMM over inputs deliberately chosen to be exactly bf16-representable +/// (so this test's own tolerance is measuring accumulation-order/tier drift, not re-measuring +/// the f32->bf16 truncation the module doc already documents and asserts is real). +#[cfg(test)] +mod bf16_tolerance { + use crate::frame::mmm::FusedSpec; + use crate::x86_64::mmm::ndarray_avx512_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_avx512_bf16_mmm_f32_16x16.mmm(); + if !mmm.built() || !mmm.runnable() { + eprintln!("skipping: ndarray_avx512_bf16_mmm_f32_16x16 not runnable on this host"); + return; + } + let packing = &mmm.packings()[0]; + 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: 0, + }, + 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 be9c14d1ec8de570cf1b3b28de903d6ed64c4d6d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 12:44:20 +0000 Subject: [PATCH 2/5] Gate the bf16 GEMM candidate's ndarray import to x86_64 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. The new kernel's unconditional `use ndarray::simd::*` broke aarch64-apple-darwin CI. Gate the ndarray-backed implementation and its test modules to target_arch = "x86_64", with a stub for other arches that is never reached at runtime since MMMRustKernel!(x86_64; ...) marks the real kernel unbuilt there. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/src/x86_64/ndarray_bf16_gemm.rs | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/linalg/src/x86_64/ndarray_bf16_gemm.rs b/linalg/src/x86_64/ndarray_bf16_gemm.rs index 3b5f02cd13..c643552d6b 100644 --- a/linalg/src/x86_64/ndarray_bf16_gemm.rs +++ b/linalg/src/x86_64/ndarray_bf16_gemm.rs @@ -44,6 +44,7 @@ //! pilot v1's 16x8 -- packed A/B panels are padded up to the next multiple of 32 in K with //! zero rows/columns, which contribute nothing to the accumulation. +#[cfg(target_arch = "x86_64")] use ndarray::simd::{PackedBf16B, bf16_tile_gemm_16x16_packed, f32_to_bf16_batch_rne}; use crate::frame::mmm::FusedKerSpec; @@ -96,6 +97,7 @@ const TILE: usize = 16; /// at. Unlike pilot v1 this pack is a single VNNI interleave straight into the tile primitive /// (no BLAS-level3 entry point re-deriving packing/dispatch from scratch), but it is still /// real per-tile allocation and work, not something hoisted above the tile loop. +#[cfg(target_arch = "x86_64")] unsafe fn add_mat_mul_bf16(pa: *const u8, pb: *const u8, k: usize, ab: &mut [[f32; TILE]; TILE]) { unsafe { if k == 0 { @@ -137,6 +139,22 @@ unsafe fn add_mat_mul_bf16(pa: *const u8, pb: *const u8, k: usize, ab: &mut [[f3 } } +// `linalg/src/lib.rs` compiles this 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 dependency on x86_64 (`linalg/Cargo.toml`). This +// stub keeps the crate compiling there; `MMMRustKernel!(x86_64; ...)` marks the real +// kernel `built(cfg!(target_arch = "x86_64"))`, so `MmmDispatch` never selects it and +// this arm never runs off x86_64. +#[cfg(not(target_arch = "x86_64"))] +unsafe fn add_mat_mul_bf16( + _pa: *const u8, + _pb: *const u8, + _k: usize, + _ab: &mut [[f32; TILE]; TILE], +) { + unreachable!("ndarray_bf16_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 { @@ -234,7 +252,7 @@ pub(super) unsafe fn kernel(mut pnl: *const FusedKerSpec) -> isize { 0 } -#[cfg(test)] +#[cfg(all(test, target_arch = "x86_64"))] mod dispatch_stays_default { use crate::frame::mmm::{MmmDispatch, Query}; use tract_data::internal::DatumType; @@ -271,7 +289,7 @@ mod dispatch_stays_default { /// naive f32 reference GEMM over inputs deliberately chosen to be exactly bf16-representable /// (so this test's own tolerance is measuring accumulation-order/tier drift, not re-measuring /// the f32->bf16 truncation the module doc already documents and asserts is real). -#[cfg(test)] +#[cfg(all(test, target_arch = "x86_64"))] mod bf16_tolerance { use crate::frame::mmm::FusedSpec; use crate::x86_64::mmm::ndarray_avx512_bf16_mmm_f32_16x16; From ac7b3326979bb96c1bf42c16f12d962a31e4b11b Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 12:53:47 +0000 Subject: [PATCH 3/5] Exclude the bf16 GEMM candidate from automatic dispatch The symbolic-N fallback in core::ops::einsum::kernel_selection:: strategize picks the largest-nr kernel per packing group, bypassing preferred/boost entirely. This kernel's nr=16 exceeds every existing f32 AVX-512 kernel's nr (max 12), so a real f32 model with a dynamic N dimension could have silently landed on this bf16-truncating kernel. Register it via MMMRustKernel!'s lower-level form, which skips the inventory::submit! that makes a kernel discoverable by MmmDispatch::native() -- the kernel stays directly constructible for this pilot's own bench/tests, but is never selected automatically. Rewrote dispatch_stays_default to assert non-discoverability for both a concrete and a symbolic N, and trimmed the module doc to the current contract instead of narrating pilot-v1 history and benchmark numbers. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/src/x86_64/mmm.rs | 18 ++++++-- linalg/src/x86_64/ndarray_bf16_gemm.rs | 57 ++++++++++---------------- 2 files changed, 36 insertions(+), 39 deletions(-) diff --git a/linalg/src/x86_64/mmm.rs b/linalg/src/x86_64/mmm.rs index e35a1eefe2..005207acd5 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -133,10 +133,20 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 ( 16, 8)@(512,4) isa(X86_64Avx // `simd::bf16_tile_gemm_16x16_packed` (AMX / AVX-512-VNNI-bf16 / FMA-polyfill tiers) instead of // hand-written asm -- see ndarray_bf16_gemm.rs's module doc for the precision tradeoff and the // honest read on how this compares structurally and numerically to pilot v1's blas_gemm -// candidate. Purely additive: it carries no boost, so retain_best ties it with the asm kernels -// on preference, and every x86_64 dispatch tier below (amd/intel_avx512_linear) still names its -// own asm kernels explicitly and never sees this one. -MMMRustKernel!(x86_64; ndarray_bf16_gemm::kernel => ndarray_avx512_bf16_mmm_f32_16x16(16, 16) isa(X86_64Avx512f)); +// candidate. +// +// Deliberately NOT registered through the `(x86_64; ...)` macro sugar, which also +// `inventory::submit!`s an `MmmRoutine` that `MmmDispatch::native()` (and so +// `core::ops::einsum::kernel_selection::strategize`) discovers automatically. This kernel's +// nr=16 is larger than every existing f32 AVX-512 kernel (max nr=12, `avx512_mmm_f32_16x12`), +// so the symbolic-N grouped fallback in `strategize` -- which picks the largest-`nr` kernel +// per packing group, bypassing `preferred`/boost entirely -- would silently select this +// bf16-truncating kernel for real f32 models with a dynamic N dimension. Calling the lower-level +// form directly skips that `inventory::submit!`, so the kernel stays reachable for direct +// construction (this pilot's own bench/tests) but invisible to automatic dispatch -- the +// concrete guarantee "purely additive, no behavior change" actually requires. +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)); 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/ndarray_bf16_gemm.rs b/linalg/src/x86_64/ndarray_bf16_gemm.rs index c643552d6b..8de0c696b7 100644 --- a/linalg/src/x86_64/ndarray_bf16_gemm.rs +++ b/linalg/src/x86_64/ndarray_bf16_gemm.rs @@ -1,25 +1,14 @@ #![allow(clippy::needless_range_loop)] //! An f32 GEMM `MatMatMulKer` body whose `AddMatMul` step truncates its operands to bf16 and //! calls into the AdaWorldAPI ndarray fork's `simd::bf16_tile_gemm_16x16_packed` tile primitive -//! (AMX `TDPBF16PS` → AVX-512 `VDPBF16PS` → decode+FMA polyfill, selected at runtime) as a -//! second additional candidate alongside the hand-tuned AVX-512 asm kernels and pilot v1's -//! f32-exact `ndarray_avx512_mmm_f32_16x8` (`ndarray_gemm.rs`). +//! (AMX `TDPBF16PS` → AVX-512 `VDPBF16PS` → decode+FMA polyfill, selected at runtime) as an +//! additional candidate alongside the hand-tuned AVX-512 asm kernels and the f32-exact +//! `ndarray_avx512_mmm_f32_16x8` (`ndarray_gemm.rs`). //! -//! **Where this stands relative to pilot v1's structural problem:** `ndarray_gemm.rs`'s -//! `blas_gemm` call allocates a fresh `Array` and re-packs its B operand on *every tile call* -//! via a full BLAS-level3 entry point (allocation + a generic path selection on top of the -//! repack), which was measured 5-10x slower than the hand-tuned asm kernel for exactly that -//! reason. This kernel's `AddMatMul` step is *also* invoked once per output tile (that is how -//! `MatMatMulKer`'s fused-op interpreter calls into any kernel body — one call per (MR, NR) -//! tile, carrying that tile's full K depth), so it still allocates and VNNI-packs its A/B -//! operands **once per tile call**, not once per whole-matrix GEMM — hoisting the pack any -//! further up would mean restructuring the packed-panel format `MatMatMulKer` hands the kernel, -//! which is out of scope for this pilot. What *is* structurally different from pilot v1: the -//! per-call work here is a single `PackedBf16B::pack` (one VNNI interleave over a `k×16` -//! buffer) plus a bf16 truncation pass, calling directly into a tile primitive with no -//! allocation inside — not a generic BLAS-level3 entry point that re-derives packing/path -//! selection from scratch on every call. The measured benchmark numbers below are the honest -//! comparison; see them before assuming this claim translates into a win. +//! Registered outside automatic dispatch (see `mmm.rs`'s registration comment for this kernel): +//! this kernel's `AddMatMul` step allocates and VNNI-packs its A/B operands once per output-tile +//! call (that is the granularity `MatMatMulKer`'s fused-op interpreter calls a kernel body at), +//! reachable only by direct construction, not through `MmmDispatch::native()`. //! //! **Precision, stated plainly:** the accumulate arithmetic (`C += A·B`) is bit-exact across //! all three `bf16_tile_gemm` tiers for bf16-exact-integer operands with accumulation below @@ -257,25 +246,23 @@ mod dispatch_stays_default { use crate::frame::mmm::{MmmDispatch, Query}; use tract_data::internal::DatumType; + /// This kernel is registered without `inventory::submit!` (see `mmm.rs`'s registration + /// comment) specifically so it never reaches `MmmDispatch::native()` -- for both a concrete + /// and a symbolic (`None`) N, since the symbolic-N fallback in + /// `core::ops::einsum::kernel_selection::strategize` picks the largest-`nr` kernel per + /// packing group, bypassing `preferred`/boost entirely, and this kernel's nr=16 exceeds + /// every existing f32 AVX-512 kernel's nr. #[test] - fn adding_the_bf16_candidate_does_not_change_default_pick() { + fn bf16_candidate_is_not_reachable_through_automatic_dispatch() { let dispatch = MmmDispatch::native(); - let query = Query::plain(DatumType::F32, Some(64), Some(256), Some(32)); - let suitable = dispatch.suitable(&query); - assert!( - suitable.iter().any(|(mmm, _, _)| mmm.name() == "ndarray_avx512_bf16_mmm_f32_16x16"), - "the new candidate should be suitable wherever avx512f is native" - ); - if let Some((picked, _, _)) = dispatch.pick(&query) { - assert_ne!( - picked.name(), - "ndarray_avx512_bf16_mmm_f32_16x16", - "default dispatch must still prefer the hand-tuned asm kernel" - ); - assert_ne!( - picked.name(), - "ndarray_avx512_mmm_f32_16x8", - "default dispatch must still prefer the hand-tuned asm kernel" + 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_avx512_bf16_mmm_f32_16x16"), + "the bf16 candidate must never appear in automatic dispatch (n={n:?})" ); } } From fe9a74ce35554b1d4a8d08260af6ebd6fc7dd87a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 13:12:31 +0000 Subject: [PATCH 4/5] Skip MMMKernel!'s auto-generated exact-bit tests for lossy kernels MMMKernel! unconditionally generates test_mmm_kernel!'s bit-exact suite for any registered kernel, regardless of whether it goes through the inventory-submitting sugar or the raw form this pilot's kernel uses -- the earlier dispatch-exclusion fix didn't touch test generation. This kernel's accumulate path truncates operands to bf16, so it cannot pass an exact-vs-f32-reference comparison by construction, and CI caught the resulting failures (x86_64::mmm::test_ndarray_avx512_bf16_mmm_f32_16x16::{frame,fuse} ::prop, fuse::packed_packed_bug_3) that a too-narrow local test filter had missed. Added an additive lossy_no_exact_tests flag to MMMKernel! (default behavior unchanged for every other kernel) and set it for this one; its own bf16_tolerance module remains its real correctness test. Re-ran the benchmark after the fix to confirm the registration/test change didn't touch the compute path: numbers are unchanged within noise from the previously reported run. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/src/frame/mmm/macros.rs | 6 +++++- linalg/src/frame/mmm/tests/mod.rs | 15 +++++++++++++++ linalg/src/x86_64/mmm.rs | 8 +++++++- 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/linalg/src/frame/mmm/macros.rs b/linalg/src/frame/mmm/macros.rs index 011d95fd30..6cc6ac3ca9 100644 --- a/linalg/src/frame/mmm/macros.rs +++ b/linalg/src/frame/mmm/macros.rs @@ -79,6 +79,7 @@ macro_rules! MMMRustKernel { $(boost($boost:expr))? $(store($($store:ty),*))? $(row_major_store($rms:expr))? + $(lossy_no_exact_tests($lossy_no_exact_tests:literal))? ) => { paste! { mod [] { @@ -101,6 +102,7 @@ macro_rules! MMMRustKernel { $(boost($boost))? $(store($($store),*))? $(row_major_store($rms))? + $(lossy_no_exact_tests($lossy_no_exact_tests))? ); } } @@ -120,6 +122,7 @@ macro_rules! MMMKernel { $(boost($boost:expr))? $(store($($store:ty),*))? $(row_major_store($rms:expr))? + $(lossy_no_exact_tests($lossy_no_exact_tests:literal))? ) => { paste! { lazy_static::lazy_static! { @@ -160,8 +163,9 @@ macro_rules! MMMKernel { #[cfg(test)] mod [] { + #[allow(unused_imports)] use super::$id; - test_mmm_kernel!($ti, &*super::$id); + maybe_test_mmm_kernel!($(lossy_no_exact_tests($lossy_no_exact_tests))? ; $ti, &*super::$id); $(mmm_packed_packed_tests!(&*super::$id, $pid : $pnum);)* $($(mmm_store_test!(&*super::$id, $store);)*)? } diff --git a/linalg/src/frame/mmm/tests/mod.rs b/linalg/src/frame/mmm/tests/mod.rs index beb4fb25d1..98b398b120 100644 --- a/linalg/src/frame/mmm/tests/mod.rs +++ b/linalg/src/frame/mmm/tests/mod.rs @@ -27,6 +27,21 @@ macro_rules! test_mmm_kernel { }; } +/// Gate for `MMMKernel!`'s `lossy_no_exact_tests` flag: a kernel whose accumulate arithmetic +/// isn't exact against its declared datum type (e.g. an internal bf16 truncation) can't pass +/// `test_mmm_kernel!`'s bit-exact suite by construction, and needs its own tolerance-based +/// tests instead of this one. +#[cfg(test)] +macro_rules! maybe_test_mmm_kernel { + (lossy_no_exact_tests(true) ; $ti:tt, $ker:expr) => {}; + (lossy_no_exact_tests(false) ; $ti:tt, $ker:expr) => { + test_mmm_kernel!($ti, $ker); + }; + (; $ti:tt, $ker:expr) => { + test_mmm_kernel!($ti, $ker); + }; +} + #[macro_export] macro_rules! test_mmm_kernel_f16 { ($ker: expr) => { diff --git a/linalg/src/x86_64/mmm.rs b/linalg/src/x86_64/mmm.rs index 005207acd5..edc0c1b7b7 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -145,8 +145,14 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 ( 16, 8)@(512,4) isa(X86_64Avx // form directly skips that `inventory::submit!`, so the kernel stays reachable for direct // construction (this pilot's own bench/tests) but invisible to automatic dispatch -- the // concrete guarantee "purely additive, no behavior change" actually requires. +// +// `lossy_no_exact_tests(true)`: this kernel's accumulate path truncates its f32 operands to +// bf16 before compute, so it cannot pass `MMMKernel!`'s auto-generated bit-exact test suite +// (`test_mmm_kernel!`) by construction -- that suite compares against an exact f32 reference. +// `ndarray_bf16_gemm.rs`'s own `bf16_tolerance` module is this kernel's real correctness test. 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)); + built(cfg!(target_arch = "x86_64")) arch(Some(crate::isa::Arch::X86_64)) isa(X86_64Avx512f) + 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)); From eed2f665214407b90efd1d612c538aeb9d395a81 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 13:22:58 +0000 Subject: [PATCH 5/5] Add standalone benchmarks decomposing the bf16-AMX GEMM kernel's gap against the asm baseline into raw tile-primitive throughput versus per-tile conversion/packing overhead. The new harness calls ndarray's bf16_tile_gemm_16x16_packed directly, outside MatMatMulKer, with operands pre-converted and pre-packed for one case and only the A operand converted at runtime for another, so the AMX arithmetic itself can be measured apart from PR #5's kernel body. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/Cargo.toml | 4 + linalg/benches/amx_bf16_gap_decomposition.rs | 207 +++++++++++++++++++ 2 files changed, 211 insertions(+) create mode 100644 linalg/benches/amx_bf16_gap_decomposition.rs diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index 16b9cef686..b0be2a786d 100644 --- a/linalg/Cargo.toml +++ b/linalg/Cargo.toml @@ -92,6 +92,10 @@ harness = false name = "ndarray_bf16_gemm" harness = false +[[bench]] +name = "amx_bf16_gap_decomposition" +harness = false + [[bench]] name = "ndarray_gemm" harness = false diff --git a/linalg/benches/amx_bf16_gap_decomposition.rs b/linalg/benches/amx_bf16_gap_decomposition.rs new file mode 100644 index 0000000000..e5d0a44e7f --- /dev/null +++ b/linalg/benches/amx_bf16_gap_decomposition.rs @@ -0,0 +1,207 @@ +// Diagnostic-only benchmark: decomposes the gap between the hand-tuned AVX-512 asm f32 GEMM +// kernel (`avx512_mmm_f32_16x8`) and the bf16-tile-based candidate added in PR #5 +// (`linalg/src/x86_64/ndarray_bf16_gemm.rs`) into its two possible causes -- the raw AMX/VNNI +// tile arithmetic itself, versus the per-tile f32->bf16 conversion + VNNI packing that PR #5's +// kernel body performs on every `AddMatMul` call (once per 16x16 output tile). +// +// Four cases at the same m=k=n shapes as `ndarray_bf16_gemm.rs`: +// B0 -- the existing asm kernel through the normal `MatMatMulKer` path (sanity baseline). +// B1 -- the raw `bf16_tile_gemm_16x16_packed` primitive called directly in a hand-rolled +// tiling loop over the whole m x k x n matmul, with A and B already fully converted to +// bf16 and B already VNNI-packed *outside* the timed closure. Zero conversion, zero +// allocation, zero packing inside the timed portion -- isolates whether the tile +// arithmetic itself (AMX TDPBF16PS on this host) is fast. +// B2 -- same tiling loop and same primitive, but A is converted from f32 to bf16 *inside* the +// timed closure every iteration (simulating a runtime activation matrix), while B stays +// pre-converted and pre-packed outside the loop (simulating a weight matrix packed once +// at model-load time and reused across many activations). +// B3 -- not implemented here: it is the existing `ndarray_bf16_16x16` case in +// `ndarray_bf16_gemm.rs`, included in the PR comment report by re-running that bench. +// +// 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. +use criterion::*; +use ndarray::simd::{PackedBf16B, bf16_tile_gemm_16x16_packed, f32_to_bf16_batch_rne}; +use std::hint::black_box; +use tract_data::internal::*; +use tract_linalg::mmm::{AsInputValue, FusedSpec}; + +const TILE: usize = 16; + +/// Walks the m x k x n output in 16x16 tiles (m, n assumed multiples of 16; k padded to a +/// multiple of 32 by the caller) and accumulates each tile via `bf16_tile_gemm_16x16_packed`. +/// `a_bf16` is row-major `[m, k_padded]`; `b_packed` is one `PackedBf16B` per 16-column tile of +/// B, indexed `b_tiles[j_tile]`. No allocation, no conversion -- pure tile-primitive calls. +fn tiled_matmul_packed( + m: usize, + n: usize, + k_padded: usize, + a_bf16: &[u16], + b_tiles: &[PackedBf16B], + c: &mut [f32], +) { + let m_tiles = m / TILE; + let n_tiles = n / TILE; + let mut tile_c = [0f32; TILE * TILE]; + for it in 0..m_tiles { + let a_row_tile = &a_bf16[it * TILE * k_padded..(it + 1) * TILE * k_padded]; + for jt in 0..n_tiles { + tile_c.fill(0.0); + bf16_tile_gemm_16x16_packed(a_row_tile, &b_tiles[jt], &mut tile_c); + for i in 0..TILE { + for j in 0..TILE { + c[(it * TILE + i) * n + jt * TILE + j] = tile_c[i * TILE + j]; + } + } + } + } +} + +fn gap_decomposition(c: &mut Criterion) { + let mut group = c.benchmark_group("amx_bf16_gap_decomposition"); + + for &(m, k, n) in &[(512usize, 512usize, 512usize), (1024, 1024, 1024)] { + group.throughput(Throughput::Elements((2 * m * k * n) as u64)); + let k_padded = k.next_multiple_of(32); + let m_tiles = m / TILE; + let n_tiles = n / TILE; + + // ---- B0: existing asm kernel through the normal MatMatMulKer path ---- + { + let mmm = tract_linalg::x86_64::mmm::avx512_mmm_f32_16x8.mmm(); + group.bench_with_input( + BenchmarkId::new("B0_asm_16x8", format!("{m}x{k}x{n}")), + &(m, k, n), + |be, &(m, k, n)| { + let packing = &mmm.packings()[0]; + 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: 0, + }, + FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())), + ], + ) + .unwrap() + }); + }, + ); + } + + // ---- B1: raw AMX tile primitive, A and B fully pre-converted + pre-packed outside + // the timed loop. Isolates the tile arithmetic itself. ---- + { + let a_f32 = vec![0f32; m * k_padded]; + let mut a_bf16 = vec![0u16; m * k_padded]; + f32_to_bf16_batch_rne(&a_f32, &mut a_bf16); + + let b_f32 = vec![0f32; k_padded * n]; + let mut b_bf16_rm = vec![0u16; k_padded * n]; + f32_to_bf16_batch_rne(&b_f32, &mut b_bf16_rm); + // One PackedBf16B per 16-column tile of B, row-major over k_padded. + let mut b_tiles: Vec = Vec::with_capacity(n_tiles); + for jt in 0..n_tiles { + let mut col_major = vec![0u16; k_padded * TILE]; + for kk in 0..k_padded { + for jj in 0..TILE { + col_major[kk * TILE + jj] = b_bf16_rm[kk * n + jt * TILE + jj]; + } + } + b_tiles.push(PackedBf16B::pack(&col_major, k_padded)); + } + let mut c_out = vec![0f32; m * n]; + + group.bench_function( + BenchmarkId::new("B1_raw_amx_prepacked", format!("{m}x{k}x{n}")), + |be| { + be.iter(|| { + tiled_matmul_packed(m, n, k_padded, &a_bf16, &b_tiles, &mut c_out); + black_box(&c_out); + }); + }, + ); + let _ = m_tiles; + } + + // ---- B2: A converted f32->bf16 inside the timed loop (runtime activation); + // B pre-converted + pre-packed outside (weight matrix packed once at load time). ---- + { + let a_f32 = vec![0f32; m * k_padded]; + + let b_f32 = vec![0f32; k_padded * n]; + let mut b_bf16_rm = vec![0u16; k_padded * n]; + f32_to_bf16_batch_rne(&b_f32, &mut b_bf16_rm); + let mut b_tiles: Vec = Vec::with_capacity(n_tiles); + for jt in 0..n_tiles { + let mut col_major = vec![0u16; k_padded * TILE]; + for kk in 0..k_padded { + for jj in 0..TILE { + col_major[kk * TILE + jj] = b_bf16_rm[kk * n + jt * TILE + jj]; + } + } + b_tiles.push(PackedBf16B::pack(&col_major, k_padded)); + } + let mut c_out = vec![0f32; m * n]; + let mut a_bf16_scratch = vec![0u16; m * k_padded]; + + group.bench_function( + BenchmarkId::new("B2_activation_runtime_convert", format!("{m}x{k}x{n}")), + |be| { + be.iter(|| { + f32_to_bf16_batch_rne(&a_f32, &mut a_bf16_scratch); + tiled_matmul_packed(m, n, k_padded, &a_bf16_scratch, &b_tiles, &mut c_out); + black_box(&c_out); + }); + }, + ); + } + + // ---- B3: existing PR #5 kernel through MatMatMulKer -- included for cross-reference; + // see `ndarray_bf16_gemm.rs` for the primary measurement of this case. ---- + { + let mmm = tract_linalg::x86_64::mmm::ndarray_avx512_bf16_mmm_f32_16x16.mmm(); + group.bench_with_input( + BenchmarkId::new("B3_pr5_kernel_as_is", format!("{m}x{k}x{n}")), + &(m, k, n), + |be, &(m, k, n)| { + let packing = &mmm.packings()[0]; + 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: 0, + }, + FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())), + ], + ) + .unwrap() + }); + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, gap_decomposition); +criterion_main!(benches);