tract_linalg/
x86_64_fma.rs

1use crate::Ops;
2use crate::frame::element_wise::ElementWiseKer;
3use crate::frame::reduce::{MapReduceKer, ReduceKer};
4use crate::x86_64_fma::softmax::x86_64_avx512_softmax2_fastcompact_f16_64n;
5use crate::x86_64_fma::softmax::x86_64_fma_softmax2_fastcompact_f32_32n;
6
7pub mod mmm;
8
9pub mod act;
10pub mod act_f16;
11pub mod act_f16_fp16;
12
13pub mod amx;
14pub mod amx_bf16;
15pub mod avxvnni;
16pub mod by_scalar;
17pub mod erf;
18#[cfg(tract_avx512vnni)]
19pub mod fma_width;
20mod intel;
21pub mod max;
22pub mod min;
23pub mod panel_extract;
24pub mod rms_norm;
25pub mod softmax;
26
27const AVX2: fn() -> bool = || is_x86_feature_detected!("avx2");
28const FMA: fn() -> bool = || is_x86_feature_detected!("fma");
29const AVX512F: fn() -> bool = || is_x86_feature_detected!("avx512f");
30#[cfg(tract_avx512vnni)]
31const AVX512VNNI: fn() -> bool = || is_x86_feature_detected!("avx512vnni");
32
33tanh_impl!(f32, fma_tanh_f32, 8, 8, is_x86_feature_detected!("fma"));
34sigmoid_impl!(f32, fma_sigmoid_f32, 8, 8, is_x86_feature_detected!("fma"));
35
36// AVX-512 (zmm, 16-wide) variants. The assembly lives in x86_64/avx512/; the
37// main loop handles 64 lanes (4 zmm) per iteration with a 16-lane tail, so
38// nr()=16 (any multiple of 16 is safe).
39tanh_impl!(f32, avx512_tanh_f32, 16, 16, is_x86_feature_detected!("avx512f"));
40sigmoid_impl!(f32, avx512_sigmoid_f32, 16, 16, is_x86_feature_detected!("avx512f"));
41
42fn plug_avx2(_ops: &mut Ops) {}
43
44fn plug_fma(ops: &mut Ops) {
45    panel_extract::plug(ops);
46
47    ops.sigmoid_f32 = Box::new(|| fma_sigmoid_f32::ew());
48    ops.tanh_f32 = Box::new(|| fma_tanh_f32::ew());
49
50    ops.mul_by_scalar_f32 = Box::new(|| by_scalar::x86_64_avx_f32_mul_by_scalar_32n::ew());
51    ops.max_f32 = Box::new(|| max::x86_64_fma_max_f32_32n::red());
52    ops.min_f32 = Box::new(|| min::x86_64_fma_min_f32_32n::red());
53    ops.softmax2_fastcompact_f32 = Box::new(|| x86_64_fma_softmax2_fastcompact_f32_32n::red());
54
55    log::info!("sigmoid_f32, tanh_f32: x86_64/fma activated");
56}
57
58/// On hosts that also support AVX-512_FP16 (Sapphire Rapids / Granite Rapids /
59/// later, and recent Xeon-D / consumer parts), upgrade the f16 element-wise
60/// kernels from the f32-roundtrip implementations in `act_f16.rs` to the
61/// native f16 implementations in `act_f16_fp16.rs` where the native path is
62/// actually faster on this uarch. We benched each op against its f32-roundtrip
63/// equivalent on Sapphire Rapids and only plug in the ones that win:
64///
65///   hardswish_f16:  8.71 → 31.6 Gelem/s  (3.62× native) — plug in
66///   leaky_relu_f16: 9.44 →  5.85 Gelem/s (0.62× native — regression) — keep
67///                   the f32-roundtrip version from act_f16.rs. The native
68///                   kernel exists in act_f16_fp16.rs for future revisits but
69///                   is not wired here.
70fn plug_avx512fp16(ops: &mut Ops) {
71    ops.hardswish_f16 = Box::new(|| act_f16_fp16::x86_64_avx512fp16_hardswish_f16_128n::ew());
72
73    log::info!("hardswish_f16: x86_64/avx512fp16 native activated");
74}
75
76fn plug_avx512f(ops: &mut Ops) {
77    ops.sigmoid_f32 = Box::new(|| avx512_sigmoid_f32::ew());
78    ops.tanh_f32 = Box::new(|| avx512_tanh_f32::ew());
79    ops.hardswish_f32 = Box::new(|| act::x86_64_avx512_hardswish_f32_64n::ew());
80    ops.leaky_relu_f32 = Box::new(|| act::x86_64_avx512_leaky_relu_f32_64n::ew());
81    ops.silu_f32 = Box::new(|| act::x86_64_avx512_silu_f32_16n::ew());
82    ops.gelu_f32 = Box::new(|| act::x86_64_avx512_gelu_f32_16n::ew());
83
84    ops.sigmoid_f16 = Box::new(|| act_f16::x86_64_avx512_sigmoid_f16_16n::ew());
85    ops.tanh_f16 = Box::new(|| act_f16::x86_64_avx512_tanh_f16_16n::ew());
86    ops.hardswish_f16 = Box::new(|| act_f16::x86_64_avx512_hardswish_f16_64n::ew());
87    ops.leaky_relu_f16 = Box::new(|| act_f16::x86_64_avx512_leaky_relu_f16_64n::ew());
88    ops.silu_f16 = Box::new(|| act_f16::x86_64_avx512_silu_f16_16n::ew());
89    ops.gelu_f16 = Box::new(|| act_f16::x86_64_avx512_gelu_f16_16n::ew());
90
91    ops.max_f32 = Box::new(|| max::x86_64_avx512_max_f32_64n::red());
92    ops.softmax2_fastcompact_f32 =
93        Box::new(|| softmax::x86_64_avx512_softmax2_fastcompact_f32_64n::red());
94    ops.softmax2_fastcompact_f16 = Box::new(|| x86_64_avx512_softmax2_fastcompact_f16_64n::red());
95
96    ops.erf_f32 = Box::new(|| erf::x86_64_avx512_erf_f32_64n::ew());
97
98    ops.rms_norm_f32 = Box::new(rms_norm::rms_norm_f32);
99
100    log::info!(
101        "sigmoid_f32, tanh_f32, hardswish_f32, leaky_relu_f32, \
102         silu_f32, gelu_f32, \
103         sigmoid_f16, tanh_f16, hardswish_f16, leaky_relu_f16, \
104         silu_f16, gelu_f16, \
105         max_f32, softmax2_fastcompact_f32, softmax2_fastcompact_f16, erf_f32, \
106         rms_norm_f32: x86_64/avx512f activated"
107    );
108}
109
110pub fn plug(ops: &mut Ops) {
111    mmm::plug(ops);
112    if is_x86_feature_detected!("avx2") {
113        plug_avx2(ops);
114        if is_x86_feature_detected!("fma") {
115            plug_fma(ops);
116            if is_x86_feature_detected!("avx512f") {
117                plug_avx512f(ops);
118                if is_x86_feature_detected!("avx512fp16") {
119                    plug_avx512fp16(ops);
120                }
121            }
122        }
123    }
124}
tract_linalg/x86_64_fma.rs

tract_linalg/
x86_64_fma.rs
