gemmkit 0.1.2

//! Parses the `GEMMKIT_REQUIRE_ISA` pin and provides the memoized ISA-selection plumbing. Every
//! dispatch family's `select_*` ladder builds from `x86_isa_detected!` and `memoized_select!`

/// x86 feature probe shared by every `select_*` ladder. Under `std`, it runs a real runtime
/// CPUID query through `is_x86_feature_detected!`. Without `std`, it falls back to a
/// compile-time `cfg!(target_feature = ...)` check, because `is_x86_feature_detected!` itself
/// needs `std`. A `no_std` build therefore only ever sees whatever target-features it was
/// compiled with
#[cfg(all(feature = "std", any(target_arch = "x86", target_arch = "x86_64")))]
macro_rules! x86_isa_detected {
    ($feat:tt) => {
        is_x86_feature_detected!($feat)
    };
}
#[cfg(all(not(feature = "std"), any(target_arch = "x86", target_arch = "x86_64")))]
macro_rules! x86_isa_detected {
    ($feat:tt) => {
        cfg!(target_feature = $feat)
    };
}

/// An explicitly requested kernel, parsed from `GEMMKIT_REQUIRE_ISA`
///
/// Only the `std` build of [`forced_isa`] ever constructs a non-`Auto` variant, because it
/// parses the env var. The `no_std` build always returns `Auto`. Every `select_*` ladder still
/// matches on the full enum regardless of `std`, so every variant has to stay in the type. The
/// `dead_code` allowance below silences the resulting unused-variant warning on a `no_std` build,
/// instead of `#[cfg]`-ing the variants out
#[cfg_attr(not(feature = "std"), allow(dead_code))]
#[derive(Copy, Clone, PartialEq, Eq)]
pub(super) enum ForcedIsa {
    /// No pin: auto-select the best available ISA (the default)
    Auto,
    /// Scalar fallback, used when no other ISA applies
    Scalar,
    /// AVX2 + FMA widen kernel
    Fma,
    /// AVX-512 Foundation (`avx512f`) widen kernel
    Avx512F,
    /// AVX-512 VNNI: the `i8` `vpdpbusd` dot kernel
    Avx512Vnni,
    /// AVX-512 BF16: the `bf16` `vdpbf16ps` dot kernel
    Avx512Bf16,
    /// NEON, the baseline (and only) SIMD ISA on aarch64
    Neon,
    /// WebAssembly `simd128`. Unlike `Neon`, this is not automatically part of a wasm32 build.
    /// It needs the compile-time `-C target-feature=+simd128` flag, which is easy to forget.
    /// Pinning it makes the build assert the SIMD path is live, and panic if the flag is
    /// absent, instead of silently falling back to scalar
    Simd128,
}

/// Parse `GEMMKIT_REQUIRE_ISA` into a [`ForcedIsa`]. Unset or empty maps to
/// [`ForcedIsa::Auto`]. An unrecognized value panics rather than silently falling back to
/// auto-selection, so a typo in a CI config fails loudly. Each dispatch type calls this at
/// most once, because the result it feeds into `select_*` is memoized in that type's `OnceLock`
#[cfg(feature = "std")]
pub(super) fn forced_isa() -> ForcedIsa {
    match std::env::var("GEMMKIT_REQUIRE_ISA") {
        Err(_) => ForcedIsa::Auto,
        Ok(s) => {
            let t = s.trim();
            if t.is_empty() || t.eq_ignore_ascii_case("auto") {
                ForcedIsa::Auto
            } else if t.eq_ignore_ascii_case("scalar") {
                ForcedIsa::Scalar
            } else if t.eq_ignore_ascii_case("fma") || t.eq_ignore_ascii_case("avx2") {
                ForcedIsa::Fma
            } else if t.eq_ignore_ascii_case("avx512f") {
                ForcedIsa::Avx512F
            } else if t.eq_ignore_ascii_case("avx512vnni") || t.eq_ignore_ascii_case("vnni") {
                ForcedIsa::Avx512Vnni
            } else if t.eq_ignore_ascii_case("avx512bf16") || t.eq_ignore_ascii_case("bf16") {
                ForcedIsa::Avx512Bf16
            } else if t.eq_ignore_ascii_case("neon") {
                ForcedIsa::Neon
            } else if t.eq_ignore_ascii_case("simd128") || t.eq_ignore_ascii_case("wasm") {
                ForcedIsa::Simd128
            } else {
                panic!(
                    "GEMMKIT_REQUIRE_ISA: unknown value `{t}` (expected scalar|fma|avx512f|avx512vnni|avx512bf16|neon|simd128|auto)"
                )
            }
        }
    }
}
// No env var access without std: every select_* ladder always runs its auto-detect path
#[cfg(not(feature = "std"))]
pub(super) fn forced_isa() -> ForcedIsa {
    ForcedIsa::Auto
}

/// Emits the memoized dispatch accessor for one element type: a `#[cfg(feature = "std")]`
/// `OnceLock<$ty>`, plus a `fn $accessor() -> $ty`. Under `std`, the accessor runs `$select`
/// once via `get_or_init`. Without `std`, it runs `$select` on every call, because there is no
/// `OnceLock` to cache into
///
/// The 6-arg form adds a trailing `$feat` literal that also gates the accessor and its static
/// on that crate feature. The static still always needs `std` too. Every dispatch family's
/// `dispatched_*` slot builds through one of these 2 arms
macro_rules! memoized_select {
    ($static:ident, $accessor:ident, $ty:ty, $select:ident, $doc:literal) => {
        #[cfg(feature = "std")]
        static $static: OnceLock<$ty> = OnceLock::new();
        #[doc = $doc]
        #[inline]
        fn $accessor() -> $ty {
            #[cfg(feature = "std")]
            {
                *$static.get_or_init($select)
            }
            #[cfg(not(feature = "std"))]
            {
                $select()
            }
        }
    };
    ($static:ident, $accessor:ident, $ty:ty, $select:ident, $doc:literal, $feat:literal) => {
        #[cfg(all(feature = "std", feature = $feat))]
        static $static: OnceLock<$ty> = OnceLock::new();
        #[doc = $doc]
        #[cfg(feature = $feat)]
        #[inline]
        fn $accessor() -> $ty {
            #[cfg(feature = "std")]
            {
                *$static.get_or_init($select)
            }
            #[cfg(not(feature = "std"))]
            {
                $select()
            }
        }
    };
}