AMX-native bf16 packed operands: pack once at prepare time, not per tile

linalg/benches/amx_bf16_gap_decomposition.rs

@@ -21,6 +21,19 @@
 // This file adds no new production kernel and touches no packing/plan code -- it is a
 // standalone harness calling `ndarray::simd::*` directly, bypassing `MatMatMulKer` entirely for
 // B1/B2.
+//
+// Two further cases measure the AMX-native packed kernel (`ndarray_bf16_native_gemm.rs` +
+// `ndarray_amx_native_pack.rs`) through the real `MatMatMulKer::run` path, at two different
+// operand-preparation lifetimes -- these are NOT the same measurement and are compared against
+// different B-cases above:
+//   P0 -- A and B both prepared (`prepare_one`) OUTSIDE the timed loop, matching B1's lifetime.
+//         Isolates tract's packed-execution abstraction tax over the raw AMX ceiling, with
+//         neither operand's preparation cost in the timed region. Compare against B1.
+//   P1 -- B prepared ONCE outside the timed loop and reused every iteration (a persistent
+//         weight, as in real inference); A's source stays f32 and is prepared (`prepare_one`)
+//         fresh INSIDE every timed iteration, once per whole matrix (not per tile) -- the same
+//         "runtime activation" lifetime B2 uses. This is the realistic inference-shaped
+//         acceptance metric. Compare against B2.
 use criterion::*;
 use ndarray::simd::{PackedBf16B, bf16_tile_gemm_16x16_packed, f32_to_bf16_batch_rne};
 use std::hint::black_box;
@@ -199,6 +212,75 @@ fn gap_decomposition(c: &mut Criterion) {
                 },
             );
         }
+
+        // ---- P0: AMX-native packed kernel through MatMatMulKer, A and B both prepared
+        // outside the timed loop -- compare against B1 (raw AMX ceiling, same lifetime). ----
+        {
+            let mmm = tract_linalg::x86_64::mmm::ndarray_amx_native_bf16_mmm_f32_16x16.mmm();
+            group.bench_with_input(
+                BenchmarkId::new("P0_amx_native_prepacked", format!("{m}x{k}x{n}")),
+                &(m, k, n),
+                |be, &(m, k, n)| {
+                    let packing = &mmm.packings()[1];
+                    let a = Tensor::zero::<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: 1,
+                                },
+                                FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())),
+                            ],
+                        )
+                        .unwrap()
+                    });
+                },
+            );
+        }
+
+        // ---- P1: AMX-native packed kernel through MatMatMulKer, B prepared once outside the
+        // timed loop and reused (persistent weight); A re-prepared once per iteration inside
+        // the timed loop, once for the whole matrix (runtime activation) -- compare against B2
+        // (same lifetime split). This is the realistic inference-shaped acceptance metric. ----
+        {
+            let mmm = tract_linalg::x86_64::mmm::ndarray_amx_native_bf16_mmm_f32_16x16.mmm();
+            group.bench_with_input(
+                BenchmarkId::new("P1_amx_native_runtime_a", format!("{m}x{k}x{n}")),
+                &(m, k, n),
+                |be, &(m, k, n)| {
+                    let packing = &mmm.packings()[1];
+                    let a = Tensor::zero::<f32>(&[m, k]).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 {
+                        let pa = packing.0.prepare_one(&a, 1, 0).unwrap();
+                        mmm.run(
+                            m,
+                            n,
+                            &[
+                                FusedSpec::AddMatMul {
+                                    a: AsInputValue::Borrowed(&*pa),
+                                    b: AsInputValue::Borrowed(&*pb),
+                                    packing: 1,
+                                },
+                                FusedSpec::Store(mmm.c_view(Some(1), Some(0)).wrap(&cc.view_mut())),
+                            ],
+                        )
+                        .unwrap();
+                        black_box(&cc);
+                    });
+                },
+            );
+        }
     }
     group.finish();
 }

linalg/src/x86_64/mmm.rs

@@ -167,6 +167,32 @@ MMMRustKernel!(ndarray_gemm::kernel::<16, 8> => ndarray_avx512_mmm_f32_16x8<f32>
 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));
+
+// AMX-native sibling of the kernel above: `AddMatMul` consumes operands already converted to
+// bf16 and VNNI-packed at `prepare_one` time (`ndarray_amx_native_pack.rs`'s
+// `NdarrayAmxBf16A`/`NdarrayAmxBf16B`, packing index 1) instead of redoing that work on every
+// tile call -- see `ndarray_bf16_native_gemm.rs`'s module doc. Same non-`inventory::submit!`
+// registration as the kernel above, for the same reason: reachable only by direct
+// construction, never through `MmmDispatch::native()`.
+//
+// `lossy_no_exact_tests(true)` for the same reason as above; packing 1's own
+// `mmm_packed_packed_tests!` (added by the `packing[1]` clause) is this kernel's auto-generated
+// correctness coverage, on top of `ndarray_bf16_native_gemm.rs`'s own `bf16_tolerance` module.
+//
+// `ndarray_amx_native_pack.rs` and `ndarray_bf16_native_gemm.rs` are both
+// `#![cfg(target_arch = "x86_64")]` at the module level (unlike the other pilot kernels in this
+// file, which stay compiled everywhere `feature = "foreign-inventory"` reaches with an internal
+// stub) -- this registration must match that gate or a foreign-inventory build on another arch
+// fails to resolve the now-nonexistent symbols.
+#[cfg(target_arch = "x86_64")]
+MMMRustKernel!(ndarray_bf16_native_gemm::kernel => ndarray_amx_native_bf16_mmm_f32_16x16<f32>(16, 16)
+    built(cfg!(target_arch = "x86_64")) arch(Some(crate::isa::Arch::X86_64)) isa(X86_64Avx512f)
+    packing[1] = amx_bf16_native => |k| k.with_packing(
+        crate::x86_64::ndarray_amx_native_pack::NdarrayAmxBf16A::new(16),
+        crate::x86_64::ndarray_amx_native_pack::NdarrayAmxBf16B::new(16),
+    );
+    lossy_no_exact_tests(true));
+
 MMMExternKernel!(x86_64; avx512_mmm_f32_32x6 <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,6 +1,8 @@
 pub mod mmm;
 
+mod ndarray_amx_native_pack;
 mod ndarray_bf16_gemm;
+mod ndarray_bf16_native_gemm;
 mod ndarray_gemm;
 
 mod amd_avx512_linear;