tract-linalg 0.23.7

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

pub mod mmm;

mod amd_avx512_linear;
mod amd_fma_linear;
mod intel_avx512_linear;
mod intel_avx512_mmv_linear;
mod intel_fma_linear;

/// CPU vendor, the axis (with the AVX-512-vs-FMA tier, read from the instruction set)
/// that selects a per-target `LinearCostModel`. `TRACT_X86_KIND=intel|amd|other`
/// overrides the CPUID probe (for forcing a cohort under emulation or in CI).
#[derive(PartialEq, Clone, Copy)]
pub(crate) enum Vendor {
    Intel,
    Amd,
    Other,
}

pub(crate) fn vendor() -> Vendor {
    if let Ok(k) = std::env::var("TRACT_X86_KIND") {
        return match k.as_str() {
            "intel" => Vendor::Intel,
            "amd" => Vendor::Amd,
            _ => Vendor::Other,
        };
    }
    cpuid_vendor()
}

#[cfg(target_arch = "x86_64")]
fn cpuid_vendor() -> Vendor {
    // `unsafe` is required on the MSRV (1.91); newer rustc deems it redundant.
    #[allow(unused_unsafe)]
    let id = unsafe { std::arch::x86_64::__cpuid(0) };
    let mut s = [0u8; 12];
    s[0..4].copy_from_slice(&id.ebx.to_le_bytes());
    s[4..8].copy_from_slice(&id.edx.to_le_bytes());
    s[8..12].copy_from_slice(&id.ecx.to_le_bytes());
    match &s {
        b"GenuineIntel" => Vendor::Intel,
        b"AuthenticAMD" => Vendor::Amd,
        _ => Vendor::Other,
    }
}

#[cfg(not(target_arch = "x86_64"))]
fn cpuid_vendor() -> Vendor {
    Vendor::Other
}

pub mod act;
pub mod act_f16;
pub mod act_f16_fp16;

// CPUID probes, tile permission syscalls and uarch burst measurements: host machinery
// with no kernel to enumerate, and only ever consulted by `isa_set`.
#[cfg(target_arch = "x86_64")]
pub mod amx;
#[cfg(target_arch = "x86_64")]
pub mod amx_bf16;
#[cfg(target_arch = "x86_64")]
pub mod avxvnni;
pub mod by_scalar;
pub mod erf;
pub mod exp;
#[cfg(tract_avx512vnni)]
pub mod fma_width;
pub mod ln;
pub mod max;
pub mod min;
pub mod panel_extract;
pub mod rms_norm;
pub mod softmax;

/// A CPUID feature probe, answering false in a build that does not target x86_64: the
/// kernels it gates are bail stubs there, so nothing may select them.
macro_rules! cpu_feature {
    ($id:ident = $feature:tt) => {
        #[cfg(target_arch = "x86_64")]
        const $id: fn() -> bool = || is_x86_feature_detected!($feature);
        #[cfg(not(target_arch = "x86_64"))]
        const $id: fn() -> bool = || false;
    };
}

cpu_feature!(AVX = "avx");
cpu_feature!(AVX2 = "avx2");
cpu_feature!(FMA = "fma");
cpu_feature!(AVX512F = "avx512f");
cpu_feature!(AVX512FP16 = "avx512fp16");
cpu_feature!(F16C = "f16c");

#[cfg(tract_avx512vnni)]
cpu_feature!(AVX512VNNI = "avx512vnni");

routine_ew_extern!(x86_64; Tanh, f32, fma_tanh_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
routine_ew_extern!(x86_64; Sigmoid, f32, fma_sigmoid_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));
routine_ew_extern!(x86_64; Silu, f32, fma_silu_f32, 8, 8, isa(X86_64Avx2, X86_64Fma));

// AVX-without-FMA ports of the fma kernels above (each vfmadd132ps expanded
// to an in-place vmulps+vaddps pair) for CPUs outside the fma tier.
routine_ew_extern!(x86_64; Tanh, f32, avx_tanh_f32, 8, 8, isa(X86_64Avx));
routine_ew_extern!(x86_64; Sigmoid, f32, avx_sigmoid_f32, 8, 8, isa(X86_64Avx));

// AVX-512 (zmm, 16-wide) variants. The assembly lives in x86_64/avx512/; the
// main loop handles 64 lanes (4 zmm) per iteration with a 16-lane tail, so
// nr()=16 (any multiple of 16 is safe).
routine_ew_extern!(x86_64; Tanh, f32, avx512_tanh_f32, 16, 16, isa(X86_64Avx512f));
routine_ew_extern!(x86_64; Sigmoid, f32, avx512_sigmoid_f32, 16, 16, isa(X86_64Avx512f));
routine_ew_extern!(x86_64; Silu, f32, avx512_silu_f32, 16, 16, isa(X86_64Avx512f));

// Correct, and slower than the f32 round-trip above on every AVX-512_FP16 part measured, so it
// is declared to keep its tests running and never preferred. A part where fp16 mul and max
// saturate their ports would want the boost dropped.
/// What CPUID says this core has, in the shared vocabulary.
pub fn isa_set() -> crate::isa::IsaSet {
    use crate::isa::{Isa, IsaSet};
    let mut set = IsaSet::of_arch(crate::isa::Arch::X86_64);
    for (isa, probe) in [
        (Isa::X86_64Avx, AVX),
        (Isa::X86_64Avx2, AVX2),
        (Isa::X86_64Fma, FMA),
        (Isa::X86_64F16c, F16C),
        (Isa::X86_64Avx512f, AVX512F),
        (Isa::X86_64Avx512Fp16, AVX512FP16),
    ] {
        if probe() {
            set = set.with(isa);
        }
    }
    #[cfg(tract_avx512vnni)]
    if AVX512VNNI() {
        set = set.with(Isa::X86_64Avx512Vnni);
    }
    #[cfg(tract_avxvnni)]
    if avxvnni::has_avxvnni() {
        set = set.with(Isa::X86_64AvxVnni);
    }
    #[cfg(tract_amx_int8)]
    if amx::has_amx_int8() {
        set = set.with(Isa::X86_64AmxInt8);
    }
    #[cfg(tract_amx_bf16)]
    if amx_bf16::has_amx_bf16() {
        set = set.with(Isa::X86_64AmxBf16);
    }
    set
}