luma_tensor/device/cpu/kernels/
element.rs

1//! Per-element numeric traits. Generic CPU kernels are written over these, then
2//! the storage enums dispatch to the concrete `f32`/`f64`/`i32`/… instantiation.
3//!
4//! This is the CPU analogue of luma-core's `NumDType`/`FloatDType`/`IntDType`,
5//! trimmed to exactly what the kernels need.
6
7use std::iter::{Product, Sum};
8
9/// Numeric element shared by float and int kinds.
10pub trait CpuNum:
11    Copy
12    + PartialOrd
13    + Send
14    + Sync
15    + 'static
16    + std::ops::Add<Output = Self>
17    + std::ops::Sub<Output = Self>
18    + std::ops::Mul<Output = Self>
19    + std::ops::Div<Output = Self>
20    + Sum
21    + Product
22{
23    const ZERO: Self;
24    const ONE: Self;
25
26    fn from_f64(v: f64) -> Self;
27    fn to_f64(self) -> f64;
28    fn from_usize(v: usize) -> Self;
29    fn to_usize(self) -> usize;
30
31    fn minimum(a: Self, b: Self) -> Self {
32        if a > b { b } else { a }
33    }
34    fn maximum(a: Self, b: Self) -> Self {
35        if a < b { b } else { a }
36    }
37    fn abs(self) -> Self;
38    fn signum(self) -> Self;
39}
40
41/// Float-only element: transcendental math and activations.
42pub trait CpuFloat: CpuNum + std::ops::Neg<Output = Self> {
43    fn exp(self) -> Self;
44    fn ln(self) -> Self;
45    fn sin(self) -> Self;
46    fn cos(self) -> Self;
47    fn tanh(self) -> Self;
48    fn sqrt(self) -> Self;
49    fn floor(self) -> Self;
50    fn ceil(self) -> Self;
51    fn round(self) -> Self;
52    fn powf(self, e: Self) -> Self;
53    fn recip(self) -> Self;
54    fn erf(self) -> Self;
55
56    fn sqr(self) -> Self {
57        self * self
58    }
59    fn relu(self) -> Self {
60        Self::maximum(self, Self::ZERO)
61    }
62    fn leaky_relu(self, negative_slope: Self) -> Self {
63        if self > Self::ZERO { self } else { self * negative_slope }
64    }
65    fn sigmoid(self) -> Self {
66        Self::ONE / (Self::ONE + (-self).exp())
67    }
68    fn silu(self) -> Self {
69        self / (Self::ONE + (-self).exp())
70    }
71    /// 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 x^3)))
72    fn gelu(self) -> Self {
73        let sqrt_2_over_pi = Self::from_f64(0.797_884_560_802_865_4);
74        let coef = Self::from_f64(0.044715);
75        let half = Self::from_f64(0.5);
76        let x3 = self * self * self;
77        let inner = sqrt_2_over_pi * (self + coef * x3);
78        half * self * (Self::ONE + inner.tanh())
79    }
80    /// 0.5 * x * (1 + erf(x / sqrt(2)))
81    fn gelu_erf(self) -> Self {
82        let frac_1_sqrt_2 = Self::from_f64(std::f64::consts::FRAC_1_SQRT_2);
83        let half = Self::from_f64(0.5);
84        half * self * (Self::ONE + (self * frac_1_sqrt_2).erf())
85    }
86}
87
88/// Int-only element. `MAX` doubles as the padding sentinel used by indexing ops
89/// (matching luma-core's `I::max_value()` convention).
90pub trait CpuInt: CpuNum + Ord + Eq {
91    const MAX: Self;
92}
93
94// ---- f32 / f64 ----
95macro_rules! impl_cpu_float {
96    ($t:ty, $erf:path) => {
97        impl CpuNum for $t {
98            const ZERO: Self = 0.0;
99            const ONE: Self = 1.0;
100            fn from_f64(v: f64) -> Self {
101                v as $t
102            }
103            fn to_f64(self) -> f64 {
104                self as f64
105            }
106            fn from_usize(v: usize) -> Self {
107                v as $t
108            }
109            fn to_usize(self) -> usize {
110                self as usize
111            }
112            fn abs(self) -> Self {
113                <$t>::abs(self)
114            }
115            fn signum(self) -> Self {
116                <$t>::signum(self)
117            }
118        }
119        impl CpuFloat for $t {
120            fn exp(self) -> Self {
121                <$t>::exp(self)
122            }
123            fn ln(self) -> Self {
124                <$t>::ln(self)
125            }
126            fn sin(self) -> Self {
127                <$t>::sin(self)
128            }
129            fn cos(self) -> Self {
130                <$t>::cos(self)
131            }
132            fn tanh(self) -> Self {
133                <$t>::tanh(self)
134            }
135            fn sqrt(self) -> Self {
136                <$t>::sqrt(self)
137            }
138            fn floor(self) -> Self {
139                <$t>::floor(self)
140            }
141            fn ceil(self) -> Self {
142                <$t>::ceil(self)
143            }
144            fn round(self) -> Self {
145                <$t>::round(self)
146            }
147            fn powf(self, e: Self) -> Self {
148                <$t>::powf(self, e)
149            }
150            fn recip(self) -> Self {
151                <$t>::recip(self)
152            }
153            fn erf(self) -> Self {
154                $erf(self)
155            }
156        }
157    };
158}
159
160impl_cpu_float!(f32, libm::erff);
161impl_cpu_float!(f64, libm::erf);
162
163// ---- i32 / u32 / u8 ----
164macro_rules! impl_cpu_int {
165    ($t:ty) => {
166        impl CpuNum for $t {
167            const ZERO: Self = 0;
168            const ONE: Self = 1;
169            fn from_f64(v: f64) -> Self {
170                v as $t
171            }
172            fn to_f64(self) -> f64 {
173                self as f64
174            }
175            fn from_usize(v: usize) -> Self {
176                v as $t
177            }
178            fn to_usize(self) -> usize {
179                self as usize
180            }
181            fn abs(self) -> Self {
182                // unsigned types have no abs; wrap the signed case only.
183                #[allow(unused_comparisons)]
184                if self < 0 { Self::ZERO.wrapping_sub(self) } else { self }
185            }
186            fn signum(self) -> Self {
187                #[allow(unused_comparisons)]
188                if self > 0 {
189                    Self::ONE
190                } else if self < 0 {
191                    // unreachable for unsigned; wrapping_sub avoids const-overflow.
192                    Self::ZERO.wrapping_sub(Self::ONE)
193                } else {
194                    Self::ZERO
195                }
196            }
197        }
198        impl CpuInt for $t {
199            const MAX: Self = <$t>::MAX;
200        }
201    };
202}
203
204impl_cpu_int!(i32);
205impl_cpu_int!(u32);
206impl_cpu_int!(u8);
luma_tensor/device/cpu/kernels/element.rs

luma_tensor/device/cpu/kernels/
element.rs
