diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index 8874da94b2..b0be2a786d 100644 --- a/linalg/Cargo.toml +++ b/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 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); 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/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 4be0c756af..c10c3f1ba4 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -141,6 +141,32 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 ( 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(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(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 ( 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 f01e781d25..8399612a76 100644 --- a/linalg/src/x86_64/mod.rs +++ b/linalg/src/x86_64/mod.rs @@ -1,5 +1,6 @@ pub mod mmm; +mod ndarray_bf16_gemm; mod ndarray_gemm; mod amd_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..8de0c696b7 --- /dev/null +++ b/linalg/src/x86_64/ndarray_bf16_gemm.rs @@ -0,0 +1,357 @@ +#![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 an +//! additional candidate alongside the hand-tuned AVX-512 asm kernels and the f32-exact +//! `ndarray_avx512_mmm_f32_16x8` (`ndarray_gemm.rs`). +//! +//! 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 +//! 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. + +#[cfg(target_arch = "x86_64")] +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. +#[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 { + 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]; + } + } + } +} + +// `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 { + 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(all(test, target_arch = "x86_64"))] +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 bf16_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_avx512_bf16_mmm_f32_16x16"), + "the bf16 candidate must never appear in automatic dispatch (n={n:?})" + ); + } + } +} + +/// 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(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; + 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(), + ); + } + } + } +}