tract_linalg/x86_64/
act_f16.rs

1//! AVX-512 f16 element-wise activations for cores without native f16 arithmetic.
2//!
3//! Each kernel round-trips through the matching f32 AVX-512 kernel via
4//! `ew_kernel_via_f32!`: convert an f16 chunk into a 64-byte-aligned f32 scratch
5//! (the f32 kernels assume 64-byte-aligned input), run the f32 kernel, convert
6//! back.
7//! Conversion is driven through `std::arch` intrinsics directly (see the
8//! helpers below) because rustc + LLVM do not autovectorize the scalar
9//! `f16::to_f32` / `f16::from_f32` loops.
10
11use tract_data::internal::f16;
12
13const CHUNK: usize = 256;
14
15/// SiLU's f32 scratch length, wider than `CHUNK`.
16///
17/// Every call into an f32 kernel pays a fixed cost that does not shrink with the
18/// length passed to it: MXCSR is saved, overwritten and restored, resynchronising
19/// the FP pipeline at both ends. `avx512_silu_f32` needs a longer call than
20/// `CHUNK` to amortise that and to get its four 64-lane groups in flight across
21/// the divide; at `CHUNK` the ymm `fma_silu_f32`, whose groups are 16 lanes wide,
22/// wins instead. Must stay a multiple of `nr`.
23const SILU_CHUNK: usize = 1024;
24
25// Vectorized f16 <-> f32 helpers using vcvtph2ps / vcvtps2ph. Rustc + LLVM
26// do NOT autovectorize the scalar `.to_f32()` loop (the half crate's method
27// has branches / function-call overhead), so we drive the conversion with
28// intrinsics directly. Both helpers process 16 lanes per iteration; the tail
29// (which only fires for the 1-15 leftover lanes inside a CHUNK-sized batch)
30// falls back to scalar.
31#[cfg(target_arch = "x86_64")]
32#[target_feature(enable = "avx512f")]
33unsafe fn cvt_f16_to_f32(src: &[f16], dst: &mut [f32]) {
34    use core::arch::x86_64::*;
35    let n = src.len();
36    debug_assert!(dst.len() >= n);
37    let chunks = n / 16;
38    unsafe {
39        for k in 0..chunks {
40            let m = _mm256_loadu_si256(src.as_ptr().add(k * 16) as *const __m256i);
41            let z = _mm512_cvtph_ps(m);
42            _mm512_storeu_ps(dst.as_mut_ptr().add(k * 16), z);
43        }
44        for k in (chunks * 16)..n {
45            *dst.get_unchecked_mut(k) = src.get_unchecked(k).to_f32();
46        }
47    }
48}
49
50bail_stub!(x86_64; unsafe fn cvt_f16_to_f32(&[f16], &mut [f32]));
51
52#[cfg(target_arch = "x86_64")]
53#[target_feature(enable = "avx512f")]
54unsafe fn cvt_f32_to_f16(src: &[f32], dst: &mut [f16]) {
55    use core::arch::x86_64::*;
56    let n = src.len();
57    debug_assert!(dst.len() >= n);
58    let chunks = n / 16;
59    unsafe {
60        for k in 0..chunks {
61            let z = _mm512_loadu_ps(src.as_ptr().add(k * 16));
62            // _MM_FROUND_TO_NEAREST_INT == 0 (round-to-nearest-even, matches f16::from_f32)
63            let m = _mm512_cvtps_ph::<0>(z);
64            _mm256_storeu_si256(dst.as_mut_ptr().add(k * 16) as *mut __m256i, m);
65        }
66        for k in (chunks * 16)..n {
67            *dst.get_unchecked_mut(k) = f16::from_f32(*src.get_unchecked(k));
68        }
69    }
70}
71
72bail_stub!(x86_64; unsafe fn cvt_f32_to_f16(&[f32], &mut [f16]));
73
74// hardswish_f16
75routine_ew_via_f32!(x86_64;
76    x86_64_avx512_hardswish_f16_64n,
77    64,
78    32,
79    CHUNK,
80    64,
81    cvt_f16_to_f32,
82    cvt_f32_to_f16,
83    super::act::x86_64_avx512_hardswish_f32_64n,
84    func(Hardswish),
85    isa(X86_64Avx512f)
86);
87
88routine_ew_via_f32!(x86_64;
89    x86_64_avx512_leaky_relu_f16_64n,
90    64,
91    32,
92    CHUNK,
93    64,
94    cvt_f16_to_f32,
95    cvt_f32_to_f16,
96    super::act::x86_64_avx512_leaky_relu_f32_64n,
97    func(LeakyRelu),
98    param(alpha => alpha.to_f32()),
99    isa(X86_64Avx512f)
100);
101
102routine_ew_via_f32!(x86_64;
103    x86_64_avx512_sigmoid_f16_16n,
104    16,
105    16,
106    CHUNK,
107    64,
108    cvt_f16_to_f32,
109    cvt_f32_to_f16,
110    super::avx512_sigmoid_f32,
111    func(Sigmoid),
112    isa(X86_64Avx512f)
113);
114
115routine_ew_via_f32!(x86_64;
116    x86_64_avx512_tanh_f16_16n,
117    16,
118    16,
119    CHUNK,
120    64,
121    cvt_f16_to_f32,
122    cvt_f32_to_f16,
123    super::avx512_tanh_f32,
124    func(Tanh),
125    isa(X86_64Avx512f)
126);
127
128routine_ew_via_f32!(x86_64;
129    x86_64_avx512_silu_f16_16n,
130    16,
131    16,
132    SILU_CHUNK,
133    64,
134    cvt_f16_to_f32,
135    cvt_f32_to_f16,
136    super::avx512_silu_f32,
137    func(Silu),
138    isa(X86_64Avx512f)
139);
140
141routine_ew_via_f32!(x86_64;
142    x86_64_avx512_gelu_f16_16n,
143    16,
144    16,
145    CHUNK,
146    64,
147    cvt_f16_to_f32,
148    cvt_f32_to_f16,
149    super::act::x86_64_avx512_gelu_f32_16n,
150    func(Gelu),
151    isa(X86_64Avx512f)
152);
tract_linalg/x86_64/act_f16.rs

tract_linalg/x86_64/
act_f16.rs
