diff --git a/.travis/cross.sh b/.travis/cross.sh index 00423f5eba..e897a559a0 100755 --- a/.travis/cross.sh +++ b/.travis/cross.sh @@ -254,6 +254,11 @@ case "$PLATFORM" in cargo check --target $PLATFORM --features getrandom-js -p tract-onnx -p tract-tensorflow RUSTFLAGS='-C target-feature=+simd128' CARGO_TARGET_WASM32_WASIP1_RUNNER=wasmtime \ cargo test --target=$PLATFORM -p tract-linalg -p tract-core -p test-unit-core + # The wasm backend picks its multiply-add form, its int8 packing and its + # sigmoid/tanh kernels at compile time on +relaxed-simd, so the run above + # leaves that half of it untested. + RUSTFLAGS='-C target-feature=+simd128,+relaxed-simd' CARGO_TARGET_WASM32_WASIP1_RUNNER=wasmtime \ + cargo test --target=$PLATFORM -p tract-linalg -p tract-core -p test-unit-core ;; wasm32-*) rustup target add $PLATFORM diff --git a/linalg/benches/wasm.rs b/linalg/benches/wasm.rs index a87cd3aa41..029d546ce5 100644 --- a/linalg/benches/wasm.rs +++ b/linalg/benches/wasm.rs @@ -40,6 +40,10 @@ fn main() { eprintln!(); eprintln!("=== int8 relaxed-dot prototype: relaxed_dot vs widening (4x4 tile) ==="); bench_relaxed_dot::run(); + + eprintln!(); + eprintln!("=== sigmoid/tanh: wasm relaxed-SIMD vs generic scalar ==="); + bench_activations::run(); } #[cfg(not(target_feature = "relaxed-simd"))] eprintln!("\n(int8 relaxed-dot prototype skipped — rebuild with +relaxed-simd)"); @@ -110,19 +114,8 @@ mod bench_8x8 { elapsed.as_secs_f64() / iters as f64 * 1e9 } - fn pick(name: &str) -> Box { - let mut ops = tract_linalg::generic(); - tract_linalg::wasm::plug(&mut ops); - for impl_ in ops.mmm_impls() { - if impl_.name() == name { - return impl_.clone(); - } - } - panic!("kernel {name} not registered") - } - fn bench_shape(label: &str, m: usize, k: usize, n: usize, iters: usize) { - let k88 = pick("wasm_f32_8x8"); + let k88 = crate::util::pick("wasm_f32_8x8"); let ns = run_one(&*k88, m, k, n, iters); let m_tiles = m.div_ceil(8); let n_tiles = n.div_ceil(8); @@ -212,19 +205,8 @@ mod bench_32x1 { elapsed.as_secs_f64() / iters as f64 * 1e9 } - fn pick(name: &str) -> Box { - let mut ops = tract_linalg::generic(); - tract_linalg::wasm::plug(&mut ops); - for impl_ in ops.mmm_impls() { - if impl_.name() == name { - return impl_.clone(); - } - } - panic!("kernel {name} not registered") - } - fn bench_min_of_n(label: &str, m: usize, k: usize, iters: usize, repetitions: usize) { - let kernel = pick("wasm_f32_32x1"); + let kernel = crate::util::pick("wasm_f32_32x1"); let mut samples: Vec = Vec::with_capacity(repetitions); for _ in 0..repetitions { samples.push(run_one(&*kernel, m, k, iters)); @@ -314,19 +296,8 @@ mod bench_16x1 { elapsed.as_secs_f64() / iters as f64 * 1e9 } - fn pick(name: &str) -> Box { - let mut ops = tract_linalg::generic(); - tract_linalg::wasm::plug(&mut ops); - for impl_ in ops.mmm_impls() { - if impl_.name() == name { - return impl_.clone(); - } - } - panic!("kernel {name} not registered") - } - fn bench_min_of_n(label: &str, m: usize, k: usize, iters: usize, repetitions: usize) { - let kernel = pick("wasm_f32_16x1"); + let kernel = crate::util::pick("wasm_f32_16x1"); let mut samples: Vec = Vec::with_capacity(repetitions); for _ in 0..repetitions { samples.push(run_one(&*kernel, m, k, iters)); @@ -420,17 +391,6 @@ mod bench_i8_4x4 { elapsed.as_secs_f64() / iters as f64 * 1e9 } - fn pick(name: &str) -> Box { - let mut ops = tract_linalg::generic(); - tract_linalg::wasm::plug(&mut ops); - for impl_ in ops.mmm_impls() { - if impl_.name() == name { - return impl_.clone(); - } - } - panic!("kernel {name} not registered") - } - fn min_of_n( kernel: &dyn MatMatMul, m: usize, @@ -445,8 +405,8 @@ mod bench_i8_4x4 { } fn bench(label: &str, m: usize, k: usize, n: usize, iters: usize, reps: usize) { - let wasm = pick("wasm_i32_4x4"); - let generic = pick("generic_i32_4x4"); + let wasm = crate::util::pick("wasm_i32_4x4"); + let generic = crate::util::pick("generic_i32_4x4"); let w = min_of_n(&*wasm, m, k, n, iters, reps); let g = min_of_n(&*generic, m, k, n, iters, reps); let tiles = m.div_ceil(4) * n.div_ceil(4); @@ -682,3 +642,71 @@ mod bench_relaxed_dot { bench(1536, 8_000, 8); } } + +#[cfg(all(target_arch = "wasm32", target_feature = "relaxed-simd"))] +mod bench_activations { + //! Microbench: WASM SIMD sigmoid/tanh vs the generic scalar fallback. + //! Sizes mirror typical RNN/transformer hidden dims (256, 512, 1024). + //! Requires `+relaxed-simd`; without it the slots hold the scalar + //! polynomial and the comparison is vacuous. + + use std::time::Instant; + use tract_linalg::element_wise::ElementWiseKer; + + fn ns_per_call>(buf: &mut [f32], iters: usize) -> f64 { + // Warmup + for _ in 0..50 { + K::run(buf, ()); + } + let t0 = Instant::now(); + for _ in 0..iters { + K::run(buf, ()); + } + let elapsed = t0.elapsed(); + elapsed.as_secs_f64() / iters as f64 * 1e9 + } + + fn bench(label: &str, n: usize, iters: usize) { + // Same input for both kernels — rebuild between to avoid post-clamp + // saturation mucking up the measurement. + let make = || (0..n).map(|i| ((i % 37) as f32 - 18.0) * 0.5).collect::>(); + + let mut buf = make(); + let scalar_sig = ns_per_call::(&mut buf, iters); + let mut buf = make(); + let simd_sig = ns_per_call::(&mut buf, iters); + let mut buf = make(); + let scalar_tanh = ns_per_call::(&mut buf, iters); + let mut buf = make(); + let simd_tanh = ns_per_call::(&mut buf, iters); + + eprintln!( + "{label} n={n} iters={iters}: \ + sigmoid scalar={scalar_sig:.0} ns simd={simd_sig:.0} ns ({:.2}x); \ + tanh scalar={scalar_tanh:.0} ns simd={simd_tanh:.0} ns ({:.2}x)", + scalar_sig / simd_sig, + scalar_tanh / simd_tanh, + ); + } + pub fn run() { + bench("hidden=256", 256, 5_000); + bench("hidden=512", 512, 3_000); + bench("hidden=1024", 1024, 2_000); + } +} + +#[cfg(target_arch = "wasm32")] +mod util { + use tract_linalg::mmm::MatMatMul; + + pub fn pick(name: &str) -> Box { + let mut ops = tract_linalg::generic(); + tract_linalg::wasm::plug(&mut ops); + for impl_ in ops.mmm_impls() { + if impl_.name() == name { + return impl_.clone(); + } + } + panic!("kernel {name} not registered") + } +} diff --git a/linalg/src/wasm.rs b/linalg/src/wasm.rs index 6feaab47f6..85f6c4d93a 100644 --- a/linalg/src/wasm.rs +++ b/linalg/src/wasm.rs @@ -56,7 +56,7 @@ pub fn plug(ops: &mut Ops) { // max(mr)=32x1 for every M, leaving up to ~37% on the table for // small-M GEMV. ops.mmm_f32 = Box::new(|_m, _k, _n| wasm_f32_8x8.mmm()); - // Bands derived from microbench_dispatch_gemv. At each band edge, using + // Bands derived from benches/wasm.rs. At each band edge, using // the next-larger kernel beats halving outer iterations of the smaller // one (1 outer with ILP-absorbed padding > 2 outer with kernel preamble // doubled). M=4/8/16 are exact tile fits at the lower edges; M=17/9/5 @@ -76,336 +76,3 @@ pub fn plug(ops: &mut Ops) { ops.tanh_f32 = Box::new(|| WasmTanh4Relaxed::ew()); } } - -#[cfg(test)] -mod microbench_32x1 { - //! Quick microbench: time per-call cost for the kernel kit's GEMV path - //! on DFN3-shaped inputs. Compares 16x1 vs 32x1 head-to-head by - //! dispatching the named kernel directly. - //! - //! Run with: - //! RUSTFLAGS='-C target-feature=+simd128' \ - //! CARGO_TARGET_WASM32_WASIP1_RUNNER='wasmtime --env RUST_TEST_NOCAPTURE=1 --' \ - //! cargo test --release --target=wasm32-wasip1 -p tract-linalg \ - //! wasm::microbench_32x1::microbench -- --nocapture --ignored - - use crate::mmm::{AsInputValue, FusedSpec}; - use std::time::Instant; - use tract_data::internal::*; - use tract_data::prelude::*; - - fn run_one(kernel: &dyn crate::mmm::MatMatMul, m: usize, k: usize, iters: usize) -> f64 { - // Pack A (m,k) and B (k,1) - let packing = &kernel.packings()[0]; - let a = Tensor::zero::(&[m, k]).unwrap(); - let pa = packing.0.prepare_one(&a, 1, 0).unwrap(); - let b = Tensor::zero::(&[k, 1]).unwrap(); - let pb = packing.1.prepare_one(&b, 0, 1).unwrap(); - let mut c = Tensor::zero::(&[m, 1]).unwrap(); - - // Warmup - for _ in 0..50 { - unsafe { - kernel - .run( - m, - 1, - &[ - FusedSpec::AddMatMul { - a: AsInputValue::Borrowed(&*pa), - b: AsInputValue::Borrowed(&*pb), - packing: 0, - }, - FusedSpec::Store(kernel.c_view(Some(0), Some(0)).wrap(&c.view_mut())), - ], - ) - .unwrap(); - } - } - - // Timed - let t0 = Instant::now(); - for _ in 0..iters { - unsafe { - kernel - .run( - m, - 1, - &[ - FusedSpec::AddMatMul { - a: AsInputValue::Borrowed(&*pa), - b: AsInputValue::Borrowed(&*pb), - packing: 0, - }, - FusedSpec::Store(kernel.c_view(Some(0), Some(0)).wrap(&c.view_mut())), - ], - ) - .unwrap(); - } - } - let elapsed = t0.elapsed(); - elapsed.as_secs_f64() / iters as f64 * 1e9 // ns/call - } - - fn pick(name: &str) -> Box { - let mut ops = crate::generic(); - super::plug(&mut ops); - for impl_ in ops.mmm_impls() { - if impl_.name() == name { - return impl_.clone(); - } - } - panic!("kernel {name} not registered") - } - - fn bench_shape(label: &str, m: usize, k: usize, iters: usize) { - let k16 = pick("wasm_f32_16x1"); - let k32 = pick("wasm_f32_32x1"); - let ns16 = run_one(&*k16, m, k, iters); - let ns32 = run_one(&*k32, m, k, iters); - let calls16 = m.div_ceil(16); - let calls32 = m.div_ceil(32); - let delta = (ns32 - ns16) / ns16 * 100.0; - eprintln!( - "{label} (m={m}, k={k}, iters={iters}): 16x1={ns16:.1} ns/call ({calls16} kernel calls); 32x1={ns32:.1} ns/call ({calls32} kernel calls); Δ={delta:+.2}% ; per-frame call ns: 16x1={n16:.1} 32x1={n32:.1} pf-Δ={dpf:+.2}%", - n16 = ns16 * calls16 as f64, - n32 = ns32 * calls32 as f64, - dpf = (ns32 * calls32 as f64 - ns16 * calls16 as f64) / (ns16 * calls16 as f64) * 100.0, - ); - } - - #[test] - #[ignore] - fn microbench() { - eprintln!("=== DFN3 GEMV microbench: 16x1 vs 32x1 ==="); - // DFN3 GRU gates (highest call count) - bench_shape("GRU m=256 k=256", 256, 256, 5_000); - // post-RNN - bench_shape("post-rnn m=256 k=512", 256, 512, 3_000); - // frame encoder - bench_shape("frame-encoder m=64 k=96", 64, 96, 20_000); - // perfect tile - bench_shape("perfect-tile m=32 k=256", 32, 256, 20_000); - } - - /// Numerical-equivalence sanity check between 16x1 and 32x1 kernels on a - /// real-shape matmul with non-trivial inputs. - /// - /// Under `+simd128` (no relaxed-simd): both kernels emit - /// `f32x4_add(f32x4_mul(...))` via `madd_f32x4!`, so the K-loop order is - /// identical and outputs are bit-identical. - /// - /// Under `+simd128,+relaxed-simd`: 32x1 uses `f32x4.relaxed_madd` (fused - /// FMA) via `madd_f32x4!`, while 16x1 uses separate `mul+add` via - /// `madd_f32x4_nofma!` to avoid the destructive-accumulator recurrence - /// that throttles ≤4-accumulator kernels (see header comment on - /// `madd_f32x4_nofma`). Outputs drift by ≤1 ulp per K-step from the - /// rounding difference between fused and separate ops. We accept that - /// drift with a generous relative tolerance. - #[test] - fn numerical_consistency_16x1_vs_32x1() { - let m = 256usize; - let k = 256usize; - let mut a_data = vec![0f32; m * k]; - for (i, x) in a_data.iter_mut().enumerate() { - *x = ((i % 13) as f32 - 6.0) * 0.1 + ((i / 17) % 11) as f32 * 0.07; - } - let mut b_data = vec![0f32; k]; - for (i, x) in b_data.iter_mut().enumerate() { - *x = (i as f32).sin() * 0.5; - } - let a = Tensor::from_shape(&[m, k], &a_data).unwrap(); - let b = Tensor::from_shape(&[k, 1], &b_data).unwrap(); - - let run = |name: &str| -> Vec { - let kernel = pick(name); - let packing = &kernel.packings()[0]; - let pa = packing.0.prepare_one(&a, 1, 0).unwrap(); - let pb = packing.1.prepare_one(&b, 0, 1).unwrap(); - let mut c = Tensor::zero::(&[m, 1]).unwrap(); - unsafe { - kernel - .run( - m, - 1, - &[ - FusedSpec::AddMatMul { - a: AsInputValue::Borrowed(&*pa), - b: AsInputValue::Borrowed(&*pb), - packing: 0, - }, - FusedSpec::Store(kernel.c_view(Some(0), Some(0)).wrap(&c.view_mut())), - ], - ) - .unwrap(); - } - c.try_as_plain().unwrap().as_slice::().unwrap().to_vec() - }; - - let c16 = run("wasm_f32_16x1"); - let c32 = run("wasm_f32_32x1"); - - #[cfg(not(target_feature = "relaxed-simd"))] - { - for (i, (x16, x32)) in c16.iter().zip(c32.iter()).enumerate() { - assert!( - x16.to_bits() == x32.to_bits(), - "row {i}: 16x1={x16} (bits 0x{:x}) != 32x1={x32} (bits 0x{:x})", - x16.to_bits(), - x32.to_bits() - ); - } - eprintln!("bit-identity OK over m={m} k={k} ({} rows)", m); - } - - #[cfg(target_feature = "relaxed-simd")] - { - // K=256 accumulator drift on fp32 between FMA and separate mul+add - // can grow up to roughly K × 0.5 ulp ≈ 128 ulp in the accumulator. - // For small-magnitude outputs that translates to ~1e-4 relative. - // We use 1e-4 as the tolerance — tight enough to catch real bugs - // (typically 1e-2+ drift) but generous for legitimate FMA drift. - let mut max_abs = 0.0f32; - let mut max_rel = 0.0f32; - for (i, (x16, x32)) in c16.iter().zip(c32.iter()).enumerate() { - let abs = (x16 - x32).abs(); - let scale = x16.abs().max(x32.abs()).max(1.0e-9); - let rel = abs / scale; - assert!( - rel < 1.0e-4, - "row {i}: relative drift {rel:e} too large; 16x1={x16} 32x1={x32}" - ); - if abs > max_abs { - max_abs = abs; - } - if rel > max_rel { - max_rel = rel; - } - } - eprintln!( - "relaxed-simd consistency OK over m={m} k={k}: max abs={max_abs:.3e}, max rel={max_rel:.3e}" - ); - } - } -} - -#[cfg(test)] -mod microbench_dispatch_gemv { - //! Microbench: 4x1 vs 8x1 vs 16x1 vs 32x1 GEMV kernels across the M - //! range. Drives the dispatch-fix decision — the M-band callback in - //! plug() routes small-M to smaller kernels, but only takes effect - //! once the kernels are tagged ManuallyOptimized (otherwise - //! kernel_selection::strategize bypasses the callback and always - //! picks max(mr) = 32x1). - //! - //! Run with: - //! RUSTFLAGS='-C target-feature=+simd128' \ - //! CARGO_TARGET_WASM32_WASIP1_RUNNER='wasmtime --env RUST_TEST_NOCAPTURE=1 --' \ - //! cargo test --release --target=wasm32-wasip1 -p tract-linalg \ - //! wasm::microbench_dispatch_gemv::microbench -- --nocapture --ignored - - use crate::mmm::{AsInputValue, FusedSpec}; - use std::time::Instant; - use tract_data::internal::*; - use tract_data::prelude::*; - - fn run_one(kernel: &dyn crate::mmm::MatMatMul, m: usize, k: usize, iters: usize) -> f64 { - let packing = &kernel.packings()[0]; - let a = Tensor::zero::(&[m, k]).unwrap(); - let pa = packing.0.prepare_one(&a, 1, 0).unwrap(); - let b = Tensor::zero::(&[k, 1]).unwrap(); - let pb = packing.1.prepare_one(&b, 0, 1).unwrap(); - let mut c = Tensor::zero::(&[m, 1]).unwrap(); - - for _ in 0..50 { - unsafe { - kernel - .run( - m, - 1, - &[ - FusedSpec::AddMatMul { - a: AsInputValue::Borrowed(&*pa), - b: AsInputValue::Borrowed(&*pb), - packing: 0, - }, - FusedSpec::Store(kernel.c_view(Some(0), Some(0)).wrap(&c.view_mut())), - ], - ) - .unwrap(); - } - } - - let t0 = Instant::now(); - for _ in 0..iters { - unsafe { - kernel - .run( - m, - 1, - &[ - FusedSpec::AddMatMul { - a: AsInputValue::Borrowed(&*pa), - b: AsInputValue::Borrowed(&*pb), - packing: 0, - }, - FusedSpec::Store(kernel.c_view(Some(0), Some(0)).wrap(&c.view_mut())), - ], - ) - .unwrap(); - } - } - let elapsed = t0.elapsed(); - elapsed.as_secs_f64() / iters as f64 * 1e9 - } - - fn pick(name: &str) -> Box { - let mut ops = crate::generic(); - super::plug(&mut ops); - for impl_ in ops.mmm_impls() { - if impl_.name() == name { - return impl_.clone(); - } - } - panic!("kernel {name} not registered") - } - - fn bench_shape(label: &str, m: usize, k: usize, iters: usize) { - let k4 = pick("wasm_f32_4x1"); - let k8 = pick("wasm_f32_8x1"); - let k16 = pick("wasm_f32_16x1"); - let k32 = pick("wasm_f32_32x1"); - let n4 = run_one(&*k4, m, k, iters); - let n8 = run_one(&*k8, m, k, iters); - let n16 = run_one(&*k16, m, k, iters); - let n32 = run_one(&*k32, m, k, iters); - let entries = [("4x1", n4), ("8x1", n8), ("16x1", n16), ("32x1", n32)]; - let winner = entries.iter().min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()).unwrap(); - let delta_vs_32 = (winner.1 - n32) / n32 * 100.0; - eprintln!( - "{label} (m={m} k={k}): 4x1={n4:.0} 8x1={n8:.0} 16x1={n16:.0} 32x1={n32:.0} ns; \ - winner={} ({:.0} ns, Δ vs 32x1: {delta_vs_32:+.1}%)", - winner.0, winner.1 - ); - } - - #[test] - #[ignore] - fn microbench() { - eprintln!("=== WASM GEMV dispatch microbench: 4x1 vs 8x1 vs 16x1 vs 32x1 ==="); - // M ≤ 16 — small-M region; the M-band callback's choices win clearly. - bench_shape("M=1 k=512", 1, 512, 50_000); - bench_shape("M=8 k=64 ", 8, 64, 50_000); - bench_shape("M=8 k=512", 8, 512, 20_000); - bench_shape("M=12 k=256", 12, 256, 50_000); - bench_shape("M=16 k=96 ", 16, 96, 50_000); - bench_shape("M=16 k=256", 16, 256, 30_000); - // M ≥ 17 — 32x1 wins (16x1 needs 2 outer iters, 32x1 single iter - // with ILP absorbing the row padding). - bench_shape("M=24 k=256", 24, 256, 30_000); - bench_shape("M=32 k=256", 32, 256, 20_000); - bench_shape("M=64 k=96 ", 64, 96, 20_000); - bench_shape("M=100 k=256", 100, 256, 10_000); - bench_shape("M=256 k=256", 256, 256, 5_000); - } -} diff --git a/linalg/src/wasm/act.rs b/linalg/src/wasm/act.rs index d1c73ecf76..baadd3fe81 100644 --- a/linalg/src/wasm/act.rs +++ b/linalg/src/wasm/act.rs @@ -215,62 +215,3 @@ mod test_wasm_sigmoid_relaxed { mod test_wasm_tanh_relaxed { tanh_frame_tests!(true, f32, crate::wasm::WasmTanh4Relaxed); } - -#[cfg(all(test, target_feature = "relaxed-simd"))] -mod microbench_activations { - //! Microbench: WASM SIMD sigmoid/tanh vs the generic scalar fallback. - //! Sizes mirror typical RNN/transformer hidden dims (256, 512, 1024). - //! - //! Run with: - //! RUSTFLAGS='-C target-feature=+simd128' \ - //! CARGO_TARGET_WASM32_WASIP1_RUNNER='wasmtime --env RUST_TEST_NOCAPTURE=1 --' \ - //! cargo test --release --target=wasm32-wasip1 -p tract-linalg \ - //! wasm::microbench_activations::microbench -- --nocapture --ignored - use crate::frame::element_wise::ElementWiseKer; - use std::time::Instant; - - fn ns_per_call>(buf: &mut [f32], iters: usize) -> f64 { - // Warmup - for _ in 0..50 { - K::run(buf, ()); - } - let t0 = Instant::now(); - for _ in 0..iters { - K::run(buf, ()); - } - let elapsed = t0.elapsed(); - elapsed.as_secs_f64() / iters as f64 * 1e9 - } - - fn bench(label: &str, n: usize, iters: usize) { - // Same input for both kernels — rebuild between to avoid post-clamp - // saturation mucking up the measurement. - let make = || (0..n).map(|i| ((i % 37) as f32 - 18.0) * 0.5).collect::>(); - - let mut buf = make(); - let scalar_sig = ns_per_call::(&mut buf, iters); - let mut buf = make(); - let simd_sig = ns_per_call::(&mut buf, iters); - let mut buf = make(); - let scalar_tanh = ns_per_call::(&mut buf, iters); - let mut buf = make(); - let simd_tanh = ns_per_call::(&mut buf, iters); - - eprintln!( - "{label} n={n} iters={iters}: \ - sigmoid scalar={scalar_sig:.0} ns simd={simd_sig:.0} ns ({:.2}x); \ - tanh scalar={scalar_tanh:.0} ns simd={simd_tanh:.0} ns ({:.2}x)", - scalar_sig / simd_sig, - scalar_tanh / simd_tanh, - ); - } - - #[test] - #[ignore] - fn microbench() { - eprintln!("=== WASM SIMD activations: scalar vs simd ==="); - bench("hidden=256", 256, 5_000); - bench("hidden=512", 512, 3_000); - bench("hidden=1024", 1024, 2_000); - } -} diff --git a/linalg/src/wasm/dispatch_tests.rs b/linalg/src/wasm/dispatch_tests.rs index e8cacb44ab..b17fec51fa 100644 --- a/linalg/src/wasm/dispatch_tests.rs +++ b/linalg/src/wasm/dispatch_tests.rs @@ -83,3 +83,229 @@ mod dispatch_trace { } } } +use crate::mmm::{AsInputValue, FusedSpec}; +use tract_data::internal::*; + +fn pick(name: &str) -> Box { + let mut ops = crate::generic(); + crate::wasm::plug(&mut ops); + for impl_ in ops.mmm_impls() { + if impl_.name() == name { + return impl_.clone(); + } + } + panic!("kernel {name} not registered") +} + +/// Numerical-equivalence sanity check between 16x1 and 32x1 kernels on a +/// real-shape matmul with non-trivial inputs. +/// +/// Under `+simd128` (no relaxed-simd): both kernels emit +/// `f32x4_add(f32x4_mul(...))` via `madd_f32x4!`, so the K-loop order is +/// identical and outputs are bit-identical. +/// +/// Under `+simd128,+relaxed-simd`: 32x1 uses `f32x4.relaxed_madd` (fused +/// FMA) via `madd_f32x4!`, while 16x1 uses separate `mul+add` via +/// `madd_f32x4_nofma!` to avoid the destructive-accumulator recurrence +/// that throttles ≤4-accumulator kernels (see header comment on +/// `madd_f32x4_nofma`). Outputs drift by ≤1 ulp per K-step from the +/// rounding difference between fused and separate ops. We accept that +/// drift with a generous relative tolerance. +#[test] +fn numerical_consistency_16x1_vs_32x1() { + let m = 256usize; + let k = 256usize; + let mut a_data = vec![0f32; m * k]; + for (i, x) in a_data.iter_mut().enumerate() { + *x = ((i % 13) as f32 - 6.0) * 0.1 + ((i / 17) % 11) as f32 * 0.07; + } + let mut b_data = vec![0f32; k]; + for (i, x) in b_data.iter_mut().enumerate() { + *x = (i as f32).sin() * 0.5; + } + let a = Tensor::from_shape(&[m, k], &a_data).unwrap(); + let b = Tensor::from_shape(&[k, 1], &b_data).unwrap(); + + let run = |name: &str| -> Vec { + let kernel = pick(name); + let packing = &kernel.packings()[0]; + let pa = packing.0.prepare_one(&a, 1, 0).unwrap(); + let pb = packing.1.prepare_one(&b, 0, 1).unwrap(); + let mut c = Tensor::zero::(&[m, 1]).unwrap(); + unsafe { + kernel + .run( + m, + 1, + &[ + FusedSpec::AddMatMul { + a: AsInputValue::Borrowed(&*pa), + b: AsInputValue::Borrowed(&*pb), + packing: 0, + }, + FusedSpec::Store(kernel.c_view(Some(0), Some(0)).wrap(&c.view_mut())), + ], + ) + .unwrap(); + } + c.try_as_plain().unwrap().as_slice::().unwrap().to_vec() + }; + + let c16 = run("wasm_f32_16x1"); + let c32 = run("wasm_f32_32x1"); + + #[cfg(not(target_feature = "relaxed-simd"))] + { + for (i, (x16, x32)) in c16.iter().zip(c32.iter()).enumerate() { + assert!( + x16.to_bits() == x32.to_bits(), + "row {i}: 16x1={x16} (bits 0x{:x}) != 32x1={x32} (bits 0x{:x})", + x16.to_bits(), + x32.to_bits() + ); + } + eprintln!("bit-identity OK over m={m} k={k} ({} rows)", m); + } + + #[cfg(target_feature = "relaxed-simd")] + { + // K=256 accumulator drift on fp32 between FMA and separate mul+add + // can grow up to roughly K × 0.5 ulp ≈ 128 ulp in the accumulator. + // For small-magnitude outputs that translates to ~1e-4 relative. + // We use 1e-4 as the tolerance — tight enough to catch real bugs + // (typically 1e-2+ drift) but generous for legitimate FMA drift. + let mut max_abs = 0.0f32; + let mut max_rel = 0.0f32; + for (i, (x16, x32)) in c16.iter().zip(c32.iter()).enumerate() { + let abs = (x16 - x32).abs(); + let scale = x16.abs().max(x32.abs()).max(1.0e-9); + let rel = abs / scale; + assert!( + rel < 1.0e-4, + "row {i}: relative drift {rel:e} too large; 16x1={x16} 32x1={x32}" + ); + if abs > max_abs { + max_abs = abs; + } + if rel > max_rel { + max_rel = rel; + } + } + eprintln!( + "relaxed-simd consistency OK over m={m} k={k}: max abs={max_abs:.3e}, max rel={max_rel:.3e}" + ); + } +} + +/// `AddRowColProducts` and `AddMatMul` with `k = 1` both compute +/// `c[i][j] += a[i] * b[j]`, so a kernel has to use the same multiply-add form +/// for both arms. Under `+relaxed-simd` the fused form keeps the full product +/// before adding while the separate form rounds it first, so a kernel that +/// fuses one arm and not the other returns two different answers for the same +/// arithmetic. +/// +/// The operands make that difference observable: with `a = b = 1 + 2^-12` the +/// product `1 + 2^-11 + 2^-24` needs 25 significand bits and rounds to +/// `1 + 2^-11`, so against `c = -1` the fused form yields `2^-11 + 2^-24` and +/// the separate form `2^-11`. Without `+relaxed-simd` both arms are `mul` then +/// `add` and the two agree trivially. +#[cfg(test)] +fn check_madd_pairing>(ker: &K) { + use crate::mmm::{FusedKerSpec, OutputStoreKer}; + + if !ker.is_supported_here() { + return; + } + let (mr, nr) = (ker.mr(), ker.nr()); + let v = 1f32 + 2f32.powi(-12); + + let (pack_a, pack_b) = &ker.packings()[0]; + let k = pack_a.k_alignment().max(pack_b.k_alignment()); + let mut a_data = vec![0f32; mr * k]; + let mut b_data = vec![0f32; k * nr]; + for i in 0..mr { + a_data[i * k] = v; + } + b_data[..nr].copy_from_slice(&vec![v; nr]); + let a = Tensor::from_shape(&[mr, k], &a_data).unwrap(); + let b = Tensor::from_shape(&[k, nr], &b_data).unwrap(); + let pa = pack_a.prepare_one(&a, 1, 0).unwrap(); + let pb = pack_b.prepare_one(&b, 0, 1).unwrap(); + + let rows = vec![v; mr]; + let cols = vec![v; nr]; + + let run = |op: FusedKerSpec| -> Vec { + let out = vec![0f32; mr * nr]; + let item = std::mem::size_of::(); + let store = OutputStoreKer { + ptr: out.as_ptr() as *mut u8, + row_byte_stride: (item * nr) as isize, + col_byte_stride: item as isize, + item_size: item, + }; + let ops = [ + FusedKerSpec::Clear, + FusedKerSpec::ScalarAdd(-1.0), + op, + FusedKerSpec::Store(store), + FusedKerSpec::Done, + ]; + assert_eq!(ker.kernel(&ops), 0); + out + }; + + let from_row_col = run(FusedKerSpec::AddRowColProducts(rows.as_ptr(), cols.as_ptr())); + let from_mat_mul = run(FusedKerSpec::AddMatMul { + k, + pa: pa.panel_bytes(0, None).unwrap(), + pb: pb.panel_bytes(0, None).unwrap(), + packing: 0, + }); + + for (i, (rc, mm)) in from_row_col.iter().zip(from_mat_mul.iter()).enumerate() { + assert_eq!( + rc.to_bits(), + mm.to_bits(), + "{}: cell {i} is {rc:e} from AddRowColProducts but {mm:e} from AddMatMul — \ + the two arms disagree on whether the multiply-add is fused", + ker.name() + ); + } +} + +#[test] +fn add_row_col_products_and_add_mat_mul_agree_on_fusion() { + check_madd_pairing(&*crate::wasm::wasm_f32_4x4); + check_madd_pairing(&*crate::wasm::wasm_f32_4x1); + check_madd_pairing(&*crate::wasm::wasm_f32_8x1); + check_madd_pairing(&*crate::wasm::wasm_f32_16x1); + check_madd_pairing(&*crate::wasm::wasm_f32_32x1); + check_madd_pairing(&*crate::wasm::wasm_f32_8x8); +} + +/// `wasm_f32_4x4` is registered at `TargetOptimized` while every other kernel +/// is `ManuallyOptimized`, so `strategize`'s `retain()` drops it before +/// selection and neither `mmm_f32` nor `mmv_f32` ever names it. Promoting it +/// without also giving it a dispatch band would silently put a 4-wide tile in +/// front of shapes the 8x8 and the GEMV kernels currently own. +#[test] +fn dispatch_never_returns_wasm_f32_4x4() { + let mut ops = crate::generic(); + crate::wasm::plug(&mut ops); + for m in [1usize, 3, 4, 5, 8, 9, 16, 17, 32, 64, 256, 1024] { + for n in [1usize, 2, 4, 8, 10, 64, 256] { + for k in [1usize, 64, 576] { + let mmm = ops + .mmm(tract_data::prelude::DatumType::F32, Some(m), Some(k), Some(n)) + .unwrap(); + assert_ne!( + mmm.name(), + "wasm_f32_4x4", + "m={m} k={k} n={n} dispatched to wasm_f32_4x4, which is registered \ + TargetOptimized and has no dispatch band" + ); + } + } + } +} diff --git a/linalg/src/wasm/mmm_f32_gemm.rs b/linalg/src/wasm/mmm_f32_gemm.rs index 97fd7a11f3..f147f1e9aa 100644 --- a/linalg/src/wasm/mmm_f32_gemm.rs +++ b/linalg/src/wasm/mmm_f32_gemm.rs @@ -284,6 +284,11 @@ unsafe fn kernel_f32_4x4(mut pnl: *const FusedKerSpec) -> isize { } } +// Reachable only by name, never through dispatch: it is the one kernel left at +// TargetOptimized, and strategize's retain() keeps only the top quality tier. +// Kept because it is the only f32 kernel besides 8x8 whose C tile is +// two-dimensional, so the generated store and packing tests cover that layout +// on a second shape. `dispatch_never_returns_wasm_f32_4x4` holds this in place. MMMRustKernel!(kernel_f32_4x4 => wasm_f32_4x4(4,4)@(4,4) quality(ImplementationQuality::TargetOptimized)); /// WASM SIMD f32 8x8 kernel — wide MM tile (8 rows × 8 cols, 16 v128 accumulators).