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

tract_linalg/x86_64/
mod.rs
