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;
65#[cfg(tract_avx512vnni)]
66pub mod fma_width;
67pub mod max;
68pub mod min;
69pub mod panel_extract;
70pub mod rms_norm;
71pub mod softmax;
72
73/// A CPUID feature probe, answering false in a build that does not target x86_64: the
74/// kernels it gates are bail stubs there, so nothing may select them.
75macro_rules! cpu_feature {
76    ($id:ident = $feature:tt) => {
77        #[cfg(target_arch = "x86_64")]
78        const $id: fn() -> bool = || is_x86_feature_detected!($feature);
79        #[cfg(not(target_arch = "x86_64"))]
80        const $id: fn() -> bool = || false;
81    };
82}
83
84cpu_feature!(AVX = "avx");
85cpu_feature!(AVX2 = "avx2");
86cpu_feature!(FMA = "fma");
87cpu_feature!(AVX512F = "avx512f");
88cpu_feature!(AVX512FP16 = "avx512fp16");
89cpu_feature!(F16C = "f16c");
90
91#[cfg(tract_avx512vnni)]
92cpu_feature!(AVX512VNNI = "avx512vnni");
93
94routine_ew_extern!(x86_64; Tanh, f32, fma_tanh_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
95routine_ew_extern!(x86_64; Sigmoid, f32, fma_sigmoid_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
96routine_ew_extern!(x86_64; Silu, f32, fma_silu_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
97
98// AVX-without-FMA ports of the fma kernels above (each vfmadd132ps expanded
99// to an in-place vmulps+vaddps pair) for CPUs outside the fma tier.
100routine_ew_extern!(x86_64; Tanh, f32, avx_tanh_f32, 8, 8, isa(X86_64Avx));
101routine_ew_extern!(x86_64; Sigmoid, f32, avx_sigmoid_f32, 8, 8, isa(X86_64Avx));
102
103// AVX-512 (zmm, 16-wide) variants. The assembly lives in x86_64/avx512/; the
104// main loop handles 64 lanes (4 zmm) per iteration with a 16-lane tail, so
105// nr()=16 (any multiple of 16 is safe).
106routine_ew_extern!(x86_64; Tanh, f32, avx512_tanh_f32, 16, 16, isa(X86_64Avx512f));
107routine_ew_extern!(x86_64; Sigmoid, f32, avx512_sigmoid_f32, 16, 16, isa(X86_64Avx512f));
108routine_ew_extern!(x86_64; Silu, f32, avx512_silu_f32, 16, 16, isa(X86_64Avx512f));
109
110// Correct, and slower than the f32 round-trip above on every AVX-512_FP16 part measured, so it
111// is declared to keep its tests running and never preferred. A part where fp16 mul and max
112// saturate their ports would want the boost dropped.
113/// What CPUID says this core has, in the shared vocabulary.
114pub fn isa_set() -> crate::isa::IsaSet {
115    use crate::isa::{Isa, IsaSet};
116    let mut set = IsaSet::of_arch(crate::isa::Arch::X86_64);
117    for (isa, probe) in [
118        (Isa::X86_64Avx, AVX),
119        (Isa::X86_64Avx2, AVX2),
120        (Isa::X86_64Fma, FMA),
121        (Isa::X86_64F16c, F16C),
122        (Isa::X86_64Avx512f, AVX512F),
123        (Isa::X86_64Avx512Fp16, AVX512FP16),
124    ] {
125        if probe() {
126            set = set.with(isa);
127        }
128    }
129    #[cfg(tract_avx512vnni)]
130    if AVX512VNNI() {
131        set = set.with(Isa::X86_64Avx512Vnni);
132    }
133    #[cfg(tract_avxvnni)]
134    if avxvnni::has_avxvnni() {
135        set = set.with(Isa::X86_64AvxVnni);
136    }
137    #[cfg(tract_amx_int8)]
138    if amx::has_amx_int8() {
139        set = set.with(Isa::X86_64AmxInt8);
140    }
141    #[cfg(tract_amx_bf16)]
142    if amx_bf16::has_amx_bf16() {
143        set = set.with(Isa::X86_64AmxBf16);
144    }
145    set
146}
tract_linalg/x86_64/mod.rs

tract_linalg/x86_64/
mod.rs
