From 9449b94775c279a2d8580677c4156128294efd86 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 07:33:23 +0000 Subject: [PATCH 1/5] Patch ndarray to the AdaWorldAPI fork so tract can reach its SIMD polyfill Adds a [patch.crates-io] entry pointing ndarray at the AdaWorldAPI fork (git, master branch, so it resolves in any checkout including CI). The fork is API-compatible (same 0.17.2) but carries an additional simd/hpc module (crate::simd::F32x16, AVX-512/AVX2/NEON/scalar tier dispatch) that tract's own linalg kernels can be wired to later. This patch alone changes no behavior; tract's element-wise/reduce kernels still use their existing intrinsics path until a follow-up adds a parallel, non-overlapping kernel registration that calls into it. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- Cargo.lock | 22 ++++++++++++++++++++-- Cargo.toml | 11 +++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d07531612e..0b85cbeb3c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1517,6 +1517,14 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fractal" +version = "0.1.0" +source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#db3a7dde568eb6d7f8806faff58bc0aa17d2a733" +dependencies = [ + "libm", +] + [[package]] name = "fs-err" version = "3.3.1" @@ -2499,13 +2507,15 @@ dependencies = [ [[package]] name = "ndarray" version = "0.17.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#db3a7dde568eb6d7f8806faff58bc0aa17d2a733" dependencies = [ + "fractal", "matrixmultiply", "num-complex", "num-integer", "num-traits", + "p64", + "paste", "portable-atomic", "portable-atomic-util", "rawpointer", @@ -2946,6 +2956,14 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "p64" +version = "0.1.0" +source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#db3a7dde568eb6d7f8806faff58bc0aa17d2a733" +dependencies = [ + "fractal", +] + [[package]] name = "page_size" version = "0.6.0" diff --git a/Cargo.toml b/Cargo.toml index 3cbe58e5c3..e501110142 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -265,3 +265,14 @@ opt-level = 2 debug = false # strip = "debuginfo" does not work on android and ios incremental = false + +# AdaWorldAPI/ndarray fork: same 0.17.2 API surface as upstream, plus a +# `simd`/`hpc` module (crate::simd::F32x16 etc.) with a LazyLock-cached +# AVX-512 -> AVX2 -> NEON -> scalar dispatch tier (src/simd.rs, +# src/simd_avx512.rs). Patching it in is transparent for every existing +# ndarray:: call site; it does not by itself accelerate any tract kernel +# until tract's own linalg ops are rewritten to call crate::simd::* instead +# of their current SIMD path. Pinned via git (not a sibling path) so this +# resolves in any checkout, CI included. +[patch.crates-io] +ndarray = { git = "https://github.com/AdaWorldAPI/ndarray", branch = "master" } From 08554d6e620a65afb7f747255b53e12e1ef9965b Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 07:48:50 +0000 Subject: [PATCH 2/5] Bump MSRV to 1.97 to match the ndarray fork's pinned toolchain CI was failing across the whole 1.91 matrix because the ndarray fork patched in by the previous commit requires rustc 1.97 (its rust-toolchain.toml is permanently pinned there for AdaWorldAPI-stack alignment). CI derives its tested toolchain straight from Cargo.toml's rust-version, so bumping that one field re-targets every job at once; updated the README rustc badge to match per the comment above that field. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- Cargo.toml | 2 +- README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e501110142..7bee0b95be 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -124,7 +124,7 @@ default-members = [ [workspace.package] # MSRV source of truth: CI derives the tested toolchain from this. Keep the # README rustc badge in sync when bumping. -rust-version = "1.91" +rust-version = "1.97" [workspace.dependencies] anstyle = "1.0.2" diff --git a/README.md b/README.md index 0755a39815..655556c04c 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ ![tract-logo](assets/tract-logo/PNG/tract-horizontal-blue.png) ![Rust](https://img.shields.io/badge/rust-%23000000.svg?style=for-the-badge&logo=rust&logoColor=white) -![rustc >= 1.91.0](https://img.shields.io/badge/rustc-%3E%3D1.91.0-brightgreen) +![rustc >= 1.97.0](https://img.shields.io/badge/rustc-%3E%3D1.97.0-brightgreen) ![MIT/Apache 2](https://img.shields.io/crates/l/tract) [![Native Linux test status](https://github.com/sonos/tract/workflows/Native%20Linux/badge.svg)](https://github.com/sonos/tract/actions) [![Embedded targets status](https://github.com/sonos/tract/workflows/Embedded%20targets/badge.svg)](https://github.com/sonos/tract/actions) From 259288ada9bbaf59e794e9f76bfd68d7fba67edd Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 07:50:17 +0000 Subject: [PATCH 3/5] Address codex review: allow the ndarray fork's git source in cargo-deny Both api/rs/deny.toml and cli/deny.toml gate git dependencies with an allow-git list that only permitted rustformers/llm; the ndarray patch resolves ndarray/fractal/p64 from AdaWorldAPI/ndarray, which the cargo-deny CI job would otherwise reject. Also trimmed the patch.crates-io comment to the current constraint (why the git pin exists) instead of narrating a follow-up change that hasn't happened. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- Cargo.toml | 5 +---- api/rs/deny.toml | 1 + cli/deny.toml | 1 + 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 7bee0b95be..8264ba49b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -269,10 +269,7 @@ incremental = false # AdaWorldAPI/ndarray fork: same 0.17.2 API surface as upstream, plus a # `simd`/`hpc` module (crate::simd::F32x16 etc.) with a LazyLock-cached # AVX-512 -> AVX2 -> NEON -> scalar dispatch tier (src/simd.rs, -# src/simd_avx512.rs). Patching it in is transparent for every existing -# ndarray:: call site; it does not by itself accelerate any tract kernel -# until tract's own linalg ops are rewritten to call crate::simd::* instead -# of their current SIMD path. Pinned via git (not a sibling path) so this +# src/simd_avx512.rs). Pinned via git (not a sibling path) so this # resolves in any checkout, CI included. [patch.crates-io] ndarray = { git = "https://github.com/AdaWorldAPI/ndarray", branch = "master" } diff --git a/api/rs/deny.toml b/api/rs/deny.toml index 0ec5a9aa9c..00db08df14 100644 --- a/api/rs/deny.toml +++ b/api/rs/deny.toml @@ -49,6 +49,7 @@ skip = [ # trusted git sources. allow-git = [ "https://github.com/rustformers/llm.git", + "https://github.com/AdaWorldAPI/ndarray", ] [licenses] diff --git a/cli/deny.toml b/cli/deny.toml index bee6142bb0..dab870aa0e 100644 --- a/cli/deny.toml +++ b/cli/deny.toml @@ -35,6 +35,7 @@ deny = [ # trusted git sources. allow-git = [ "https://github.com/rustformers/llm.git", + "https://github.com/AdaWorldAPI/ndarray", ] [licenses] From 5e8250fe4de6aed9df47eabb2a0cc7739b3bfce7 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 12:00:33 +0000 Subject: [PATCH 4/5] Add an ndarray-blas_gemm-backed AVX-512 f32 GEMM kernel as an additional 16x8 candidate avx512_mmm_f32_16x8 was the only f32 GEMM tile with no ndarray-backed sibling, so this adds ndarray_avx512_mmm_f32_16x8: same tile geometry, same fused-op interpreter as the generic reference kernel, but the AddMatMul accumulation calls ndarray::simd::BlasLevel3::blas_gemm (the AdaWorldAPI fork's canonical consumer-facing re-export, not hpc::blas_level3 directly) instead of a hand-written inner-product loop. Registered purely additively through the existing MMMRustKernel! machinery; no existing kernel, tier, or preference is touched, and a new test pins that default dispatch still picks the asm kernel. Requires a companion ndarray-fork change (adding the blas_level3 re-export to simd.rs, made in the local checkout) that is not yet pushed upstream, so this currently only builds against a local path-patched ndarray, not the git-pinned fork. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- Cargo.lock | 1 + linalg/Cargo.toml | 11 ++ linalg/benches/ndarray_gemm.rs | 50 +++++++ linalg/src/x86_64/mmm.rs | 8 ++ linalg/src/x86_64/mod.rs | 2 + linalg/src/x86_64/ndarray_gemm.rs | 209 ++++++++++++++++++++++++++++++ 6 files changed, 281 insertions(+) create mode 100644 linalg/benches/ndarray_gemm.rs create mode 100644 linalg/src/x86_64/ndarray_gemm.rs diff --git a/Cargo.lock b/Cargo.lock index 0b85cbeb3c..abde11c3ee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5005,6 +5005,7 @@ dependencies = [ "libc", "log", "minijinja", + "ndarray", "nu-ansi-term", "num-traits", "pastey 0.2.3", diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index ef4f3d92b3..1827e4a0ad 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: one x86_64 GEMM kernel calls the AdaWorldAPI ndarray fork's +# `hpc::blas_level3::BlasLevel3::blas_gemm` (AVX-512-capable native backend) +# instead of hand-rolled asm, as an additional candidate alongside the +# existing hand-tuned AVX-512 asm kernels. +[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_gemm" +harness = false + [[bench]] name = "mm_for_wavenet_hw" harness = false diff --git a/linalg/benches/ndarray_gemm.rs b/linalg/benches/ndarray_gemm.rs new file mode 100644 index 0000000000..30053d7591 --- /dev/null +++ b/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::(&[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..dce00b958a 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -127,6 +127,14 @@ 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: 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. Purely additive: it carries no +// boost, so retain_best ties it with the asm kernel on preference, and every x86_64 +// dispatch tier below (amd/intel_avx512_linear) still names avx512_mmm_f32_16x8 +// explicitly and never sees this one. +MMMRustKernel!(x86_64; ndarray_gemm::kernel::<16, 8> => ndarray_avx512_mmm_f32_16x8(16, 8) 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..f01e781d25 100644 --- a/linalg/src/x86_64/mod.rs +++ b/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; diff --git a/linalg/src/x86_64/ndarray_gemm.rs b/linalg/src/x86_64/ndarray_gemm.rs new file mode 100644 index 0000000000..c4f0a354a8 --- /dev/null +++ b/linalg/src/x86_64/ndarray_gemm.rs @@ -0,0 +1,209 @@ +#![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`"). + +use ndarray::ArrayView2; +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. +unsafe fn add_mat_mul_ndarray( + 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]]; + } + } + } +} + +unsafe fn add_unicast( + 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(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( + mut pnl: *const FusedKerSpec, +) -> 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::(&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::(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_ndarray_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_mmm_f32_16x8"), + "the new candidate should be suitable wherever avx512f is native" + ); + if let Some((picked, _, _)) = dispatch.pick(&query) { + assert_ne!( + picked.name(), + "ndarray_avx512_mmm_f32_16x8", + "default dispatch must still prefer the hand-tuned asm kernel" + ); + } + } +} From f83a442ff0d928e60c34620dc65c7cbb18f88b78 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 12:55:30 +0000 Subject: [PATCH 5/5] Exclude the ndarray GEMM candidate from automatic dispatch Same defensive fix as pilot v2 (PR #5): 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=8 currently ties rather than exceeds the existing max (avx512_mmm_f32_16x12's nr=12), but relying on that ordering to hold is fragile and inconsistent with the "purely additive, no behavior change" guarantee this pilot claims. 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. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht --- linalg/src/x86_64/mmm.rs | 16 +++++++++++----- linalg/src/x86_64/ndarray_gemm.rs | 24 ++++++++++++------------ 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/linalg/src/x86_64/mmm.rs b/linalg/src/x86_64/mmm.rs index dce00b958a..4be0c756af 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -130,11 +130,17 @@ MMMExternKernel!(x86_64; avx512_mmm_f32_16x8 ( 16, 8)@(512,4) isa(X86_64Avx // 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. Purely additive: it carries no -// boost, so retain_best ties it with the asm kernel on preference, and every x86_64 -// dispatch tier below (amd/intel_avx512_linear) still names avx512_mmm_f32_16x8 -// explicitly and never sees this one. -MMMRustKernel!(x86_64; ndarray_gemm::kernel::<16, 8> => ndarray_avx512_mmm_f32_16x8(16, 8) isa(X86_64Avx512f)); +// `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(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 ( 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_gemm.rs b/linalg/src/x86_64/ndarray_gemm.rs index 71d6ff8319..97e54040a7 100644 --- a/linalg/src/x86_64/ndarray_gemm.rs +++ b/linalg/src/x86_64/ndarray_gemm.rs @@ -208,20 +208,20 @@ 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 adding_the_ndarray_candidate_does_not_change_default_pick() { + fn ndarray_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_mmm_f32_16x8"), - "the new candidate should be suitable wherever avx512f is native" - ); - if let Some((picked, _, _)) = dispatch.pick(&query) { - 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_mmm_f32_16x8"), + "the ndarray candidate must never appear in automatic dispatch (n={n:?})" ); } }