Pilot v2: bf16 tile GEMM (AMX/AVX-512) as an additional ndarray-backed candidate

linalg/Cargo.toml

@@ -34,9 +34,9 @@ tract-data.workspace = true
 libc.workspace = true
 
 # Pilot: additional x86_64 kernel candidates call the AdaWorldAPI ndarray fork's
-# `simd::*` surface (F32x16 polyfill for leaky_relu; `BlasLevel3::blas_gemm` for
-# GEMM) instead of hand-rolled intrinsics/asm, alongside the existing hand-tuned
-# AVX-512 kernels.
+# `simd::*` surface (F32x16 polyfill for leaky_relu; `BlasLevel3::blas_gemm` and
+# `hpc::bf16_tile_gemm` tile primitives for GEMM) instead of hand-rolled
+# intrinsics/asm, alongside the existing hand-tuned AVX-512 kernels.
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 ndarray.workspace = true
 
@@ -88,6 +88,14 @@ harness = false
 name = "mat_vec"
 harness = false
 
+[[bench]]
+name = "ndarray_bf16_gemm"
+harness = false
+
+[[bench]]
+name = "amx_bf16_gap_decomposition"
+harness = false
+
 [[bench]]
 name = "ndarray_gemm"
 harness = false

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::<f32>(&[m, k]).unwrap();
+                    let pa = packing.0.prepare_one(&a, 1, 0).unwrap();
+                    let b = Tensor::zero::<f32>(&[k, n]).unwrap();
+                    let pb = packing.1.prepare_one(&b, 0, 1).unwrap();
+                    let mut cc = Tensor::zero::<f32>(&[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<PackedBf16B> = 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<PackedBf16B> = 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::<f32>(&[m, k]).unwrap();
+                    let pa = packing.0.prepare_one(&a, 1, 0).unwrap();
+                    let b = Tensor::zero::<f32>(&[k, n]).unwrap();
+                    let pb = packing.1.prepare_one(&b, 0, 1).unwrap();
+                    let mut cc = Tensor::zero::<f32>(&[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);

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::<f32>(&[m, k]).unwrap();
+                    let pa = packing.0.prepare_one(&a, 1, 0).unwrap();
+                    let b = Tensor::zero::<f32>(&[k, n]).unwrap();
+                    let pb = packing.1.prepare_one(&b, 0, 1).unwrap();
+                    let mut cc = Tensor::zero::<f32>(&[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);

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 [<sys_ $id>] {
@@ -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 [<test_$id>] {
+                #[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);)*)?
             }

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) => {

linalg/src/x86_64/mmm.rs

@@ -141,6 +141,32 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 <f32>( 16, 8)@(512,4) isa(X86_64Avx
 // own bench/tests) but invisible to automatic dispatch.
 MMMRustKernel!(ndarray_gemm::kernel::<16, 8> => ndarray_avx512_mmm_f32_16x8<f32>(16, 8)
     built(cfg!(target_arch = "x86_64")) arch(Some(crate::isa::Arch::X86_64)) 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.
+//
+// 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.
+//
+// `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<f32>(16, 16)
+    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 <f32>( 32, 6)@(512,4) isa(X86_64Avx512f));
 MMMExternKernel!(x86_64; avx512_mmm_f32_32x5 <f32>( 32, 5)@(512,4) isa(X86_64Avx512f));
 MMMExternKernel!(x86_64; avx512_mmm_f32_48x4 <f32>( 48, 4)@(512,4) isa(X86_64Avx512f));

linalg/src/x86_64/mod.rs

@@ -1,5 +1,6 @@
 pub mod mmm;
 
+mod ndarray_bf16_gemm;
 mod ndarray_gemm;
 
 mod amd_avx512_linear;