Add an ndarray-blas_gemm-backed AVX-512 f32 GEMM kernel candidate

linalg/Cargo.toml

@@ -33,9 +33,10 @@ tract-data.workspace = true
 [target.'cfg(target_arch = "riscv64")'.dependencies]
 libc.workspace = true
 
-# Pilot: one x86_64 element-wise kernel (leaky_relu) calls the AdaWorldAPI ndarray
-# fork's `simd::F32x16` polyfill instead of hand-rolled intrinsics/asm, as an
-# additional candidate alongside the existing hand-tuned AVX-512 kernel.
+# 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.
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 ndarray.workspace = true
 
@@ -87,6 +88,10 @@ harness = false
 name = "mat_vec"
 harness = false
 
+[[bench]]
+name = "ndarray_gemm"
+harness = false
+
 [[bench]]
 name = "mm_for_wavenet_hw"
 harness = false

linalg/benches/ndarray_gemm.rs

@@ -0,0 +1,50 @@
+// Compares the hand-tuned AVX-512 asm 16x8 GEMM kernel against the additive
+// ndarray-blas_gemm-backed candidate of the same tile geometry, on a full matrix multiply
+// (not just one microkernel tile call): the kernel's own panel-walking machinery loops the
+// tile many times over m/n/k, so this measures the whole GEMM each candidate produces.
+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_16x8");
+    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", tract_linalg::x86_64::mmm::avx512_mmm_f32_16x8.mmm()),
+            ("ndarray", tract_linalg::x86_64::mmm::ndarray_avx512_mmm_f32_16x8.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/x86_64/mmm.rs

@@ -127,6 +127,20 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_128x1<f32>(128, 1)@(512,4) isa(X86_64Avx
 MMMExternKernel!(x86_64; avx512_mmm_f32_16x1 <f32>( 16, 1)@(512,4) isa(X86_64Avx512f));
 MMMExternKernel!(x86_64; avx512_mmm_f32_16x12<f32>( 16,12)@(512,4) isa(X86_64Avx512f));
 MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 <f32>( 16, 8)@(512,4) isa(X86_64Avx512f));
+
+// Pilot: same 16x8 tile geometry as avx512_mmm_f32_16x8 above, so the two are directly
+// comparable, but the AddMatMul accumulation calls into the AdaWorldAPI ndarray fork's
+// `BlasLevel3::blas_gemm` instead of hand-written asm.
+//
+// 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`'s symbolic-N fallback, which picks the
+// largest-`nr` kernel per packing group and bypasses `preferred`/boost entirely) would
+// discover automatically. 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.
+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));
 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,7 @@
 pub mod mmm;
 
+mod ndarray_gemm;
+
 mod amd_avx512_linear;
 mod amd_fma_linear;
 mod intel_avx512_linear;

linalg/src/x86_64/ndarray_gemm.rs

@@ -0,0 +1,228 @@
+#![allow(clippy::needless_range_loop)]
+//! An f32 GEMM `MatMatMulKer` body whose `AddMatMul` step calls into the AdaWorldAPI
+//! ndarray fork's `simd::BlasLevel3::blas_gemm` instead of a hand-written inner-product
+//! loop, as an additional candidate alongside the hand-tuned AVX-512 asm kernels. Every
+//! other fused op (bias, min/max, per-row/per-col, store) is the same scalar Rust the
+//! generic reference kernel uses, so only the matmul accumulation itself is delegated.
+//!
+//! Goes through `ndarray::simd::BlasLevel3`, the canonical consumer-facing re-export,
+//! never `ndarray::hpc::blas_level3` directly — see the ndarray fork's own `CLAUDE.md`
+//! ("all SIMD from `ndarray::simd`").
+
+#[cfg(target_arch = "x86_64")]
+use ndarray::ArrayView2;
+#[cfg(target_arch = "x86_64")]
+use ndarray::simd::BlasLevel3;
+
+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])
+            }
+        }
+    };
+}
+
+/// `pa` is packed k-major, MR contiguous per k-step (`pa[ik * MR + i]`); `pb` likewise for
+/// NR. That makes the panel-pair product `ab[i][j] += sum_ik pa[ik*MR+i] * pb[ik*NR+j]` the
+/// matrix product `A_panel^T . B_panel` where `A_panel` is `(k, MR)` row-major and `B_panel`
+/// is `(k, NR)` row-major. `A_panel` is transposed into a small contiguous `(MR, k)` buffer
+/// (a copy, not a stride trick) so both operands reach `blas_gemm` as contiguous slices and
+/// take the real backend path instead of ndarray's non-contiguous fallback loop.
+#[cfg(target_arch = "x86_64")]
+unsafe fn add_mat_mul_ndarray<const MR: usize, const NR: usize>(
+    pa: *const u8,
+    pb: *const u8,
+    k: usize,
+    ab: &mut [[f32; NR]; MR],
+) {
+    unsafe {
+        if k == 0 {
+            return;
+        }
+        let a = pa as *const f32;
+        let b = pb as *const f32;
+
+        let mut a_t = vec![0f32; MR * k];
+        for i in 0..MR {
+            for ik in 0..k {
+                a_t[i * k + ik] = *a.add(ik * MR + i);
+            }
+        }
+        let a_view = ArrayView2::from_shape((MR, k), &a_t).unwrap();
+        let b_slice = std::slice::from_raw_parts(b, k * NR);
+        let b_view = ArrayView2::from_shape((k, NR), b_slice).unwrap();
+
+        let prod = a_view.blas_gemm(1.0f32, &b_view, 0.0f32);
+        for i in 0..MR {
+            for j in 0..NR {
+                ab[i][j] += prod[[i, 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_ndarray<const MR: usize, const NR: usize>(
+    _pa: *const u8,
+    _pb: *const u8,
+    _k: usize,
+    _ab: &mut [[f32; NR]; MR],
+) {
+    unreachable!("ndarray_gemm's kernel is x86_64-only and unbuilt elsewhere")
+}
+
+unsafe fn add_unicast<const MR: usize, const NR: usize>(
+    ab: &mut [[f32; NR]; MR],
+    other: &OutputStoreKer,
+) {
+    unsafe {
+        for i in 0..MR {
+            for j in 0..NR {
+                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<const MR: usize, const NR: usize>(tile: &OutputStoreKer, ab: &[[f32; NR]; MR]) {
+    unsafe {
+        for i in 0..MR {
+            for j in 0..NR {
+                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). Same
+/// fused-op interpreter shape as `crate::generic::mmm::kernel`; the `AddMatMul` arm is the
+/// only place this diverges from it.
+pub(super) unsafe fn kernel<const MR: usize, const NR: usize>(
+    mut pnl: *const FusedKerSpec<f32>,
+) -> isize {
+    unsafe {
+        let mut ab = [[0f32; NR]; MR];
+        loop {
+            if pnl.is_null() {
+                break;
+            }
+            match *pnl {
+                FusedKerSpec::Done => break,
+                FusedKerSpec::Clear => ab = [[0f32; NR]; MR],
+                FusedKerSpec::LoadTile(col_major, _row_major) => {
+                    for row in 0..MR {
+                        for col in 0..NR {
+                            ab[row][col] = *col_major.add(col * MR + 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..MR {
+                        for j in 0..NR {
+                            ab[i][j] += *rows.add(i) * *cols.add(j);
+                        }
+                    }
+                }
+                FusedKerSpec::AddUnicast(other) => add_unicast::<MR, NR>(&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_ndarray::<MR, NR>(pa, pb, k, &mut ab);
+                }
+                FusedKerSpec::Store(tile) => store::<MR, NR>(&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) 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.
+    #[test]
+    fn ndarray_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_mmm_f32_16x8"),
+                "the ndarray candidate must never appear in automatic dispatch (n={n:?})"
+            );
+        }
+    }
+}