Skip to main content

tract_linalg/generic/
silu.rs

1#![allow(clippy::excessive_precision)]
2use crate::frame::element_wise::ElementWiseKer;
3use tract_data::internal::*;
4
5#[derive(Clone, Debug)]
6pub struct SSiLU4;
7
8impl ElementWiseKer<f32> for SSiLU4 {
9    fn name() -> &'static str {
10        "generic"
11    }
12
13    fn alignment_bytes() -> usize {
14        16
15    }
16
17    fn alignment_items() -> usize {
18        4
19    }
20
21    fn nr() -> usize {
22        4
23    }
24
25    fn run(x: &mut [f32], _: ()) {
26        debug_assert!(x.len() % Self::nr() == 0);
27        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
28        x.iter_mut().for_each(|px| {
29            let sigmoid = 1.0 / (1.0 + (-*px).exp());
30            *px = *px * sigmoid;
31        });
32    }
33}
34
35#[derive(Clone, Debug)]
36pub struct HSiLU8;
37
38impl ElementWiseKer<f16> for HSiLU8 {
39    fn name() -> &'static str {
40        "generic"
41    }
42
43    fn alignment_bytes() -> usize {
44        16
45    }
46
47    fn alignment_items() -> usize {
48        4
49    }
50
51    fn nr() -> usize {
52        8
53    }
54
55    fn run(x: &mut [f16], _: ()) {
56        debug_assert!(x.len() % Self::nr() == 0);
57        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
58        x.iter_mut().for_each(|px| {
59            let x_f32 = px.to_f32();
60            let sigmoid = 1.0 / (1.0 + (-x_f32).exp());
61            *px = f16::from_f32(x_f32 * sigmoid);
62        });
63    }
64}
65
66#[cfg(test)]
67#[macro_use]
68pub mod s {
69    silu_frame_tests!(true, f32, crate::generic::silu::SSiLU4);
70}
71
72#[cfg(test)]
73#[macro_use]
74pub mod h {
75    silu_frame_tests!(true, tract_data::internal::f16, crate::generic::silu::HSiLU8);
76}