Skip to main content

tract_linalg/generic/
hardswish.rs

1#![allow(clippy::excessive_precision)]
2use tract_data::internal::*;
3use tract_num_traits::Zero;
4
5routine_ew_rust!(generic;
6    f32,
7    generic_hardswish_f32_4n,
8    4,
9    4,
10    fn run(x: &mut [f32], _: ()) {
11        debug_assert!(x.len() % Self::nr() == 0);
12        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
13        const INV6: f32 = 1.0 / 6.0;
14        x.iter_mut().for_each(|px| {
15            let relu6 = (*px + 3.0).clamp(0.0, 6.0);
16            *px = *px * relu6 * INV6;
17        });
18    },
19    func(Hardswish)
20);
21
22routine_ew_rust!(generic;
23    f16,
24    generic_hardswish_f16_8n,
25    8,
26    8,
27    fn run(x: &mut [f16], _: ()) {
28        debug_assert!(x.len() % Self::nr() == 0);
29        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
30        let three = f16::from_f32(3.0);
31        let six = f16::from_f32(6.0);
32        let inv6 = f16::from_f32(1.0 / 6.0);
33        x.iter_mut().for_each(|px| {
34            let relu6 = ((*px + three).min(six)).max(f16::zero());
35            *px = *px * relu6 * inv6;
36        });
37    },
38    func(Hardswish)
39);