diff --git a/Cargo.lock b/Cargo.lock index abde11c3ee..b63f9742e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1520,7 +1520,7 @@ dependencies = [ [[package]] name = "fractal" version = "0.1.0" -source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#db3a7dde568eb6d7f8806faff58bc0aa17d2a733" +source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#e55af0e3d10fb0065cbad4de2df48f8567503309" dependencies = [ "libm", ] @@ -2507,7 +2507,7 @@ dependencies = [ [[package]] name = "ndarray" version = "0.17.2" -source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#db3a7dde568eb6d7f8806faff58bc0aa17d2a733" +source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#e55af0e3d10fb0065cbad4de2df48f8567503309" dependencies = [ "fractal", "matrixmultiply", @@ -2959,7 +2959,7 @@ checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" [[package]] name = "p64" version = "0.1.0" -source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#db3a7dde568eb6d7f8806faff58bc0aa17d2a733" +source = "git+https://github.com/AdaWorldAPI/ndarray?branch=master#e55af0e3d10fb0065cbad4de2df48f8567503309" dependencies = [ "fractal", ] diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index 810dc9343d..8874da94b2 100644 --- a/linalg/Cargo.toml +++ b/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 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..4be0c756af 100644 --- a/linalg/src/x86_64/mmm.rs +++ b/linalg/src/x86_64/mmm.rs @@ -127,6 +127,20 @@ 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. +// +// 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/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..97e54040a7 --- /dev/null +++ b/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( + 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( + _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( + 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(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:?})" + ); + } + } +}