tract_linalg/x86_64/
mod.rs

1pub mod mmm;
2
3mod amd_avx512_linear;
4mod amd_fma_linear;
5mod intel_avx512_linear;
6mod intel_avx512_mmv_linear;
7mod intel_fma_linear;
8
9/// CPU vendor, the axis (with the AVX-512-vs-FMA tier, read from the instruction set)
10/// that selects a per-target `LinearCostModel`. `TRACT_X86_KIND=intel|amd|other`
11/// overrides the CPUID probe (for forcing a cohort under emulation or in CI).
12#[derive(PartialEq, Clone, Copy)]
13pub(crate) enum Vendor {
14    Intel,
15    Amd,
16    Other,
17}
18
19pub(crate) fn vendor() -> Vendor {
20    if let Ok(k) = std::env::var("TRACT_X86_KIND") {
21        return match k.as_str() {
22            "intel" => Vendor::Intel,
23            "amd" => Vendor::Amd,
24            _ => Vendor::Other,
25        };
26    }
27    cpuid_vendor()
28}
29
30#[cfg(target_arch = "x86_64")]
31fn cpuid_vendor() -> Vendor {
32    // `unsafe` is required on the MSRV (1.91); newer rustc deems it redundant.
33    #[allow(unused_unsafe)]
34    let id = unsafe { std::arch::x86_64::__cpuid(0) };
35    let mut s = [0u8; 12];
36    s[0..4].copy_from_slice(&id.ebx.to_le_bytes());
37    s[4..8].copy_from_slice(&id.edx.to_le_bytes());
38    s[8..12].copy_from_slice(&id.ecx.to_le_bytes());
39    match &s {
40        b"GenuineIntel" => Vendor::Intel,
41        b"AuthenticAMD" => Vendor::Amd,
42        _ => Vendor::Other,
43    }
44}
45
46#[cfg(not(target_arch = "x86_64"))]
47fn cpuid_vendor() -> Vendor {
48    Vendor::Other
49}
50
51pub mod act;
52pub mod act_f16;
53pub mod act_f16_fp16;
54
55// CPUID probes, tile permission syscalls and uarch burst measurements: host machinery
56// with no kernel to enumerate, and only ever consulted by `isa_set`.
57#[cfg(target_arch = "x86_64")]
58pub mod amx;
59#[cfg(target_arch = "x86_64")]
60pub mod amx_bf16;
61#[cfg(target_arch = "x86_64")]
62pub mod avxvnni;
63pub mod by_scalar;
64pub mod erf;
65pub mod exp;
66#[cfg(tract_avx512vnni)]
67pub mod fma_width;
68pub mod ln;
69pub mod max;
70pub mod min;
71pub mod panel_extract;
72pub mod rms_norm;
73pub mod softmax;
74
75/// A CPUID feature probe, answering false in a build that does not target x86_64: the
76/// kernels it gates are bail stubs there, so nothing may select them.
77macro_rules! cpu_feature {
78    ($id:ident = $feature:tt) => {
79        #[cfg(target_arch = "x86_64")]
80        const $id: fn() -> bool = || is_x86_feature_detected!($feature);
81        #[cfg(not(target_arch = "x86_64"))]
82        const $id: fn() -> bool = || false;
83    };
84}
85
86cpu_feature!(AVX = "avx");
87cpu_feature!(AVX2 = "avx2");
88cpu_feature!(FMA = "fma");
89cpu_feature!(AVX512F = "avx512f");
90cpu_feature!(AVX512FP16 = "avx512fp16");
91cpu_feature!(F16C = "f16c");
92
93#[cfg(tract_avx512vnni)]
94cpu_feature!(AVX512VNNI = "avx512vnni");
95
96routine_ew_extern!(x86_64; Tanh, f32, fma_tanh_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
97routine_ew_extern!(x86_64; Sigmoid, f32, fma_sigmoid_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
98routine_ew_extern!(x86_64; Silu, f32, fma_silu_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
99
100// AVX-without-FMA ports of the fma kernels above (each vfmadd132ps expanded
101// to an in-place vmulps+vaddps pair) for CPUs outside the fma tier.
102routine_ew_extern!(x86_64; Tanh, f32, avx_tanh_f32, 8, 8, isa(X86_64Avx));
103routine_ew_extern!(x86_64; Sigmoid, f32, avx_sigmoid_f32, 8, 8, isa(X86_64Avx));
104
105// AVX-512 (zmm, 16-wide) variants. The assembly lives in x86_64/avx512/; the
106// main loop handles 64 lanes (4 zmm) per iteration with a 16-lane tail, so
107// nr()=16 (any multiple of 16 is safe).
108routine_ew_extern!(x86_64; Tanh, f32, avx512_tanh_f32, 16, 16, isa(X86_64Avx512f));
109routine_ew_extern!(x86_64; Sigmoid, f32, avx512_sigmoid_f32, 16, 16, isa(X86_64Avx512f));
110routine_ew_extern!(x86_64; Silu, f32, avx512_silu_f32, 16, 16, isa(X86_64Avx512f));
111
112// Correct, and slower than the f32 round-trip above on every AVX-512_FP16 part measured, so it
113// is declared to keep its tests running and never preferred. A part where fp16 mul and max
114// saturate their ports would want the boost dropped.
115/// What CPUID says this core has, in the shared vocabulary.
116pub fn isa_set() -> crate::isa::IsaSet {
117    use crate::isa::{Isa, IsaSet};
118    let mut set = IsaSet::of_arch(crate::isa::Arch::X86_64);
119    for (isa, probe) in [
120        (Isa::X86_64Avx, AVX),
121        (Isa::X86_64Avx2, AVX2),
122        (Isa::X86_64Fma, FMA),
123        (Isa::X86_64F16c, F16C),
124        (Isa::X86_64Avx512f, AVX512F),
125        (Isa::X86_64Avx512Fp16, AVX512FP16),
126    ] {
127        if probe() {
128            set = set.with(isa);
129        }
130    }
131    #[cfg(tract_avx512vnni)]
132    if AVX512VNNI() {
133        set = set.with(Isa::X86_64Avx512Vnni);
134    }
135    #[cfg(tract_avxvnni)]
136    if avxvnni::has_avxvnni() {
137        set = set.with(Isa::X86_64AvxVnni);
138    }
139    #[cfg(tract_amx_int8)]
140    if amx::has_amx_int8() {
141        set = set.with(Isa::X86_64AmxInt8);
142    }
143    #[cfg(tract_amx_bf16)]
144    if amx_bf16::has_amx_bf16() {
145        set = set.with(Isa::X86_64AmxBf16);
146    }
147    set
148}
tract_linalg/x86_64/mod.rs

tract_linalg/x86_64/
mod.rs
