tract-linalg 0.23.8

//! The wasm kernel tree, behind `simd128` and — for the activations — `relaxed-simd`.
//!
//! Running its tests needs `wasmtime` and the wasi target:
//!
//! ```text
//! RUSTFLAGS='-C target-feature=+simd128,+relaxed-simd' \
//! CARGO_TARGET_WASM32_WASIP1_RUNNER=wasmtime \
//! cargo test --target wasm32-wasip1 -p tract-linalg
//! ```
use crate::DatumType;

#[macro_use]
mod madd;
#[macro_use]
mod fuse;

mod act;
mod act_f32;
#[cfg(all(test, target_arch = "wasm32", target_feature = "simd128"))]
mod dispatch_tests;
mod exp;
mod ln;
mod mmm_f32_gemm;
mod mmm_f32_gemv;
mod mmm_i32;
mod reduce;

pub use act::*;
pub use act_f32::*;
pub use exp::*;
pub use ln::*;
pub use mmm_f32_gemm::*;
pub use mmm_f32_gemv::*;
pub use mmm_i32::*;

/// Every kernel this tier names must be one written for wasm: its answer is held to the suitable
/// list, and a generic kernel counts as no opinion, leaving the N>1 rule to pick max(nr*mr) among
/// the surviving GEMV kernels — i.e. wasm_f32_32x1, a matrix-vector kernel, for every GEMM.
fn preferred(
    _isa: &crate::isa::IsaSet,
    dt: DatumType,
    query: &crate::mmm::Query,
    _suitable: &[crate::mmm::Suitable],
) -> Option<&'static str> {
    match (dt, query.n) {
        // int8 -> i32 matmul: SIMD kernel (the generic scalar is the tier below).
        (DatumType::I32, Some(1)) => None,
        (DatumType::I32, _) => Some(wasm_i32_4x4.name.as_str()),
        // GEMV routes by M-band to the kernel whose MR fits. Bands derived from
        // benches/wasm.rs: at each edge, using the next-larger kernel beats halving outer
        // iterations of the smaller one (1 outer with ILP-absorbed padding > 2 outer with the
        // kernel preamble doubled). M=4/8/16 are exact tile fits at the lower edges; M=17/9/5
        // are the first values where the next-larger kernel wins.
        (DatumType::F32, Some(1)) => Some(match query.m.unwrap_or(0) {
            0..=4 => &wasm_f32_4x1.name,
            5..=8 => &wasm_f32_8x1.name,
            9..=16 => &wasm_f32_16x1.name,
            _ => &wasm_f32_32x1.name,
        }),
        // GEMM by N: a k-step of 4x16 splats one A value into four multiply-adds where 8x8
        // gets two, so it is 1.09-1.11x on the 64-wide shapes and 1.21x aggregate. It only
        // pays when N is a multiple of 16: at N=8 a 16-wide tile wastes half its columns
        // (0.58x) and N=24 pays that on its second tile, and gating on N>=16 instead cost
        // fastenhancer tiny and base 4% each. An N unknown at optimisation time stays on 8x8.
        // Rank the two on a full tiled GEMM: over a single tile the wider tile's doubled
        // B-panel re-reads do not show.
        (DatumType::F32, Some(n)) if n % 16 == 0 => Some(wasm_f32_4x16.name.as_str()),
        (DatumType::F32, _) => Some(wasm_f32_8x8.name.as_str()),
        _ => None,
    }
}

inventory::submit! {
    crate::mmm_tiers::MmmTier {
        arch: Some(crate::isa::Arch::Wasm32Simd128),
        precedence: 1,
        name: "wasm-simd128",
        applies: |_| true,
        preferred,
    }
}

/// What this build offers, in the shared vocabulary. Unlike the other trees this is a build
/// question rather than a probe: wasm features are enabled at compile time and a module cannot
/// ask the engine what it got.
pub fn isa_set() -> crate::isa::IsaSet {
    use crate::isa::{Isa, IsaSet};
    let mut set = IsaSet::of_arch(crate::isa::Arch::Wasm32Simd128).with(Isa::Wasm32Simd128);
    if cfg!(target_feature = "relaxed-simd") {
        set = set.with(Isa::Wasm32RelaxedSimd);
    }
    set
}