/home/noah/src/trueno/src/backends/sse2/mod.rs
Line | Count | Source |
1 | | //! SSE2 backend implementation (x86_64 baseline SIMD) |
2 | | //! |
3 | | //! This backend uses SSE2 intrinsics for 128-bit SIMD operations. |
4 | | //! SSE2 is available on all x86_64 CPUs as a baseline requirement. |
5 | | //! |
6 | | //! # Performance |
7 | | //! |
8 | | //! Expected speedup: 4x for operations on aligned f32 vectors (4 elements per register) |
9 | | //! |
10 | | //! # Safety |
11 | | //! |
12 | | //! All SSE2 intrinsics are marked `unsafe` by Rust. This module carefully isolates |
13 | | //! all unsafe code and verifies correctness through comprehensive testing. |
14 | | |
15 | | mod ops; |
16 | | |
17 | | #[cfg(target_arch = "x86_64")] |
18 | | use std::arch::x86_64::*; |
19 | | |
20 | | use super::VectorBackend; |
21 | | |
22 | | /// SSE2 backend (128-bit SIMD for x86_64) |
23 | | pub struct Sse2Backend; |
24 | | |
25 | | impl VectorBackend for Sse2Backend { |
26 | | #[inline] |
27 | | #[target_feature(enable = "sse2")] |
28 | 0 | unsafe fn add(a: &[f32], b: &[f32], result: &mut [f32]) { |
29 | 0 | ops::arithmetic::add(a, b, result); |
30 | 0 | } |
31 | | |
32 | | #[inline] |
33 | | #[target_feature(enable = "sse2")] |
34 | 0 | unsafe fn sub(a: &[f32], b: &[f32], result: &mut [f32]) { |
35 | 0 | ops::arithmetic::sub(a, b, result); |
36 | 0 | } |
37 | | |
38 | | #[inline] |
39 | | #[target_feature(enable = "sse2")] |
40 | 0 | unsafe fn mul(a: &[f32], b: &[f32], result: &mut [f32]) { |
41 | 0 | ops::arithmetic::mul(a, b, result); |
42 | 0 | } |
43 | | |
44 | | #[inline] |
45 | | #[target_feature(enable = "sse2")] |
46 | 0 | unsafe fn div(a: &[f32], b: &[f32], result: &mut [f32]) { |
47 | 0 | ops::arithmetic::div(a, b, result); |
48 | 0 | } |
49 | | |
50 | | #[inline] |
51 | | #[target_feature(enable = "sse2")] |
52 | 0 | unsafe fn dot(a: &[f32], b: &[f32]) -> f32 { |
53 | 0 | ops::reductions::dot(a, b) |
54 | 0 | } |
55 | | |
56 | | #[inline] |
57 | | #[target_feature(enable = "sse2")] |
58 | 0 | unsafe fn sum(a: &[f32]) -> f32 { |
59 | 0 | ops::reductions::sum(a) |
60 | 0 | } |
61 | | |
62 | | #[inline] |
63 | | #[target_feature(enable = "sse2")] |
64 | 0 | unsafe fn max(a: &[f32]) -> f32 { |
65 | 0 | ops::reductions::max(a) |
66 | 0 | } |
67 | | |
68 | | #[inline] |
69 | | #[target_feature(enable = "sse2")] |
70 | 0 | unsafe fn min(a: &[f32]) -> f32 { |
71 | 0 | ops::reductions::min(a) |
72 | 0 | } |
73 | | |
74 | | #[inline] |
75 | | #[target_feature(enable = "sse2")] |
76 | 0 | unsafe fn argmax(a: &[f32]) -> usize { |
77 | 0 | ops::reductions::argmax(a) |
78 | 0 | } |
79 | | |
80 | | #[inline] |
81 | | #[target_feature(enable = "sse2")] |
82 | 0 | unsafe fn argmin(a: &[f32]) -> usize { |
83 | 0 | ops::reductions::argmin(a) |
84 | 0 | } |
85 | | |
86 | | #[inline] |
87 | | #[target_feature(enable = "sse2")] |
88 | 0 | unsafe fn sum_kahan(a: &[f32]) -> f32 { |
89 | 0 | ops::reductions::sum_kahan(a) |
90 | 0 | } |
91 | | |
92 | | #[inline] |
93 | | #[target_feature(enable = "sse2")] |
94 | 0 | unsafe fn norm_l2(a: &[f32]) -> f32 { |
95 | 0 | if a.is_empty() { return 0.0; } |
96 | 0 | Self::dot(a, a).sqrt() |
97 | 0 | } |
98 | | |
99 | | #[inline] |
100 | | #[target_feature(enable = "sse2")] |
101 | 0 | unsafe fn norm_l1(a: &[f32]) -> f32 { |
102 | 0 | if a.is_empty() { return 0.0; } |
103 | 0 | let len = a.len(); |
104 | 0 | let mut i = 0; |
105 | 0 | let mut acc = _mm_setzero_ps(); |
106 | 0 | let sign_mask = _mm_set1_ps(f32::from_bits(0x7FFF_FFFF)); |
107 | 0 | while i + 4 <= len { |
108 | 0 | acc = _mm_add_ps(acc, _mm_and_ps(_mm_loadu_ps(a.as_ptr().add(i)), sign_mask)); |
109 | 0 | i += 4; |
110 | 0 | } |
111 | 0 | let mut result = { |
112 | 0 | let temp = _mm_add_ps(acc, _mm_movehl_ps(acc, acc)); |
113 | 0 | let temp = _mm_add_ss(temp, _mm_shuffle_ps(temp, temp, 1)); |
114 | 0 | _mm_cvtss_f32(temp) |
115 | | }; |
116 | 0 | for &val in &a[i..] { result += val.abs(); } |
117 | 0 | result |
118 | 0 | } |
119 | | |
120 | | #[inline] |
121 | | #[target_feature(enable = "sse2")] |
122 | 0 | unsafe fn norm_linf(a: &[f32]) -> f32 { |
123 | 0 | if a.is_empty() { return 0.0; } |
124 | 0 | let len = a.len(); |
125 | 0 | let mut i = 0; |
126 | 0 | let mut max_vec = _mm_setzero_ps(); |
127 | 0 | let sign_mask = _mm_set1_ps(f32::from_bits(0x7FFF_FFFF)); |
128 | 0 | while i + 4 <= len { |
129 | 0 | let va = _mm_loadu_ps(a.as_ptr().add(i)); |
130 | 0 | max_vec = _mm_max_ps(max_vec, _mm_and_ps(va, sign_mask)); |
131 | 0 | i += 4; |
132 | 0 | } |
133 | 0 | let mut result = { |
134 | 0 | let temp = _mm_max_ps(max_vec, _mm_movehl_ps(max_vec, max_vec)); |
135 | 0 | let temp = _mm_max_ss(temp, _mm_shuffle_ps(temp, temp, 1)); |
136 | 0 | _mm_cvtss_f32(temp) |
137 | | }; |
138 | 0 | for &val in &a[i..] { let abs_val = val.abs(); if abs_val > result { result = abs_val; } } |
139 | 0 | result |
140 | 0 | } |
141 | | |
142 | | #[inline] |
143 | | #[target_feature(enable = "sse2")] |
144 | 0 | unsafe fn scale(a: &[f32], scalar: f32, result: &mut [f32]) { |
145 | 0 | let len = a.len(); |
146 | 0 | let mut i = 0; |
147 | 0 | let scalar_vec = _mm_set1_ps(scalar); |
148 | 0 | while i + 4 <= len { |
149 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_mul_ps(_mm_loadu_ps(a.as_ptr().add(i)), scalar_vec)); |
150 | 0 | i += 4; |
151 | 0 | } |
152 | 0 | for j in i..len { result[j] = a[j] * scalar; } |
153 | 0 | } |
154 | | |
155 | | #[inline] |
156 | | #[target_feature(enable = "sse2")] |
157 | 0 | unsafe fn abs(a: &[f32], result: &mut [f32]) { |
158 | 0 | let len = a.len(); |
159 | 0 | let mut i = 0; |
160 | 0 | let sign_mask = _mm_set1_ps(f32::from_bits(0x7FFF_FFFF)); |
161 | 0 | while i + 4 <= len { |
162 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_and_ps(_mm_loadu_ps(a.as_ptr().add(i)), sign_mask)); |
163 | 0 | i += 4; |
164 | 0 | } |
165 | 0 | for j in i..len { result[j] = a[j].abs(); } |
166 | 0 | } |
167 | | |
168 | | #[inline] |
169 | | #[target_feature(enable = "sse2")] |
170 | 0 | unsafe fn clamp(a: &[f32], min_val: f32, max_val: f32, result: &mut [f32]) { |
171 | 0 | let len = a.len(); |
172 | 0 | let mut i = 0; |
173 | 0 | let min_vec = _mm_set1_ps(min_val); |
174 | 0 | let max_vec = _mm_set1_ps(max_val); |
175 | 0 | while i + 4 <= len { |
176 | 0 | let va = _mm_loadu_ps(a.as_ptr().add(i)); |
177 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_min_ps(_mm_max_ps(va, min_vec), max_vec)); |
178 | 0 | i += 4; |
179 | 0 | } |
180 | 0 | for j in i..len { result[j] = a[j].max(min_val).min(max_val); } |
181 | 0 | } |
182 | | |
183 | | #[inline] |
184 | | #[target_feature(enable = "sse2")] |
185 | 0 | unsafe fn lerp(a: &[f32], b: &[f32], t: f32, result: &mut [f32]) { |
186 | 0 | let len = a.len(); |
187 | 0 | let mut i = 0; |
188 | 0 | let t_vec = _mm_set1_ps(t); |
189 | 0 | while i + 4 <= len { |
190 | 0 | let va = _mm_loadu_ps(a.as_ptr().add(i)); |
191 | 0 | let vb = _mm_loadu_ps(b.as_ptr().add(i)); |
192 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_add_ps(va, _mm_mul_ps(t_vec, _mm_sub_ps(vb, va)))); |
193 | 0 | i += 4; |
194 | 0 | } |
195 | 0 | for j in i..len { result[j] = a[j] + t * (b[j] - a[j]); } |
196 | 0 | } |
197 | | |
198 | | #[inline] |
199 | | #[target_feature(enable = "sse2")] |
200 | 0 | unsafe fn fma(a: &[f32], b: &[f32], c: &[f32], result: &mut [f32]) { |
201 | 0 | let len = a.len(); |
202 | 0 | let mut i = 0; |
203 | 0 | while i + 4 <= len { |
204 | 0 | let va = _mm_loadu_ps(a.as_ptr().add(i)); |
205 | 0 | let vb = _mm_loadu_ps(b.as_ptr().add(i)); |
206 | 0 | let vc = _mm_loadu_ps(c.as_ptr().add(i)); |
207 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_add_ps(_mm_mul_ps(va, vb), vc)); |
208 | 0 | i += 4; |
209 | 0 | } |
210 | 0 | for j in i..len { result[j] = a[j] * b[j] + c[j]; } |
211 | 0 | } |
212 | | |
213 | | #[inline] |
214 | | #[target_feature(enable = "sse2")] |
215 | 0 | unsafe fn relu(a: &[f32], result: &mut [f32]) { |
216 | 0 | let len = a.len(); |
217 | 0 | let mut i = 0; |
218 | 0 | let zero = _mm_setzero_ps(); |
219 | 0 | while i + 4 <= len { |
220 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_max_ps(_mm_loadu_ps(a.as_ptr().add(i)), zero)); |
221 | 0 | i += 4; |
222 | 0 | } |
223 | 0 | for j in i..len { result[j] = a[j].max(0.0); } |
224 | 0 | } |
225 | | |
226 | | #[inline] |
227 | | #[target_feature(enable = "sse2")] |
228 | 0 | unsafe fn exp(a: &[f32], result: &mut [f32]) { |
229 | | // Polynomial approximation for exp - range reduction + polynomial |
230 | 0 | let len = a.len(); |
231 | 0 | let mut i = 0; |
232 | 0 | let ln2 = _mm_set1_ps(std::f32::consts::LN_2); |
233 | 0 | let inv_ln2 = _mm_set1_ps(1.0 / std::f32::consts::LN_2); |
234 | 0 | let c1 = _mm_set1_ps(1.0); |
235 | 0 | let c2 = _mm_set1_ps(0.5); |
236 | 0 | let c3 = _mm_set1_ps(0.166_666_67); |
237 | 0 | let c4 = _mm_set1_ps(0.041_666_668); |
238 | 0 | let c5 = _mm_set1_ps(0.008_333_334); |
239 | 0 | while i + 4 <= len { |
240 | 0 | let x = _mm_loadu_ps(a.as_ptr().add(i)); |
241 | 0 | let k = _mm_cvtps_epi32(_mm_mul_ps(x, inv_ln2)); |
242 | 0 | let kf = _mm_cvtepi32_ps(k); |
243 | 0 | let r = _mm_sub_ps(x, _mm_mul_ps(kf, ln2)); |
244 | 0 | let mut poly = _mm_add_ps(c1, _mm_mul_ps(r, c5)); |
245 | 0 | poly = _mm_add_ps(c1, _mm_mul_ps(r, _mm_add_ps(c4, _mm_mul_ps(r, poly)))); |
246 | 0 | poly = _mm_add_ps(c1, _mm_mul_ps(r, _mm_add_ps(c3, _mm_mul_ps(r, poly)))); |
247 | 0 | poly = _mm_add_ps(c1, _mm_mul_ps(r, _mm_add_ps(c2, _mm_mul_ps(r, poly)))); |
248 | 0 | poly = _mm_add_ps(c1, _mm_mul_ps(r, poly)); |
249 | 0 | let exp_k = _mm_castsi128_ps(_mm_slli_epi32(_mm_add_epi32(k, _mm_set1_epi32(127)), 23)); |
250 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_mul_ps(poly, exp_k)); |
251 | 0 | i += 4; |
252 | 0 | } |
253 | 0 | for j in i..len { result[j] = a[j].exp(); } |
254 | 0 | } |
255 | | |
256 | | #[inline] |
257 | | #[target_feature(enable = "sse2")] |
258 | 0 | unsafe fn sigmoid(a: &[f32], result: &mut [f32]) { |
259 | | // sigmoid(x) = 1 / (1 + exp(-x)) |
260 | 0 | let len = a.len(); |
261 | 0 | let mut i = 0; |
262 | 0 | let one = _mm_set1_ps(1.0); |
263 | 0 | let neg_one = _mm_set1_ps(-1.0); |
264 | 0 | let ln2 = _mm_set1_ps(std::f32::consts::LN_2); |
265 | 0 | let inv_ln2 = _mm_set1_ps(1.0 / std::f32::consts::LN_2); |
266 | 0 | let c2 = _mm_set1_ps(0.5); |
267 | 0 | let c3 = _mm_set1_ps(0.166_666_67); |
268 | 0 | let c4 = _mm_set1_ps(0.041_666_668); |
269 | 0 | let c5 = _mm_set1_ps(0.008_333_334); |
270 | 0 | while i + 4 <= len { |
271 | 0 | let x = _mm_loadu_ps(a.as_ptr().add(i)); |
272 | 0 | let neg_x = _mm_mul_ps(x, neg_one); |
273 | 0 | let k = _mm_cvtps_epi32(_mm_mul_ps(neg_x, inv_ln2)); |
274 | 0 | let kf = _mm_cvtepi32_ps(k); |
275 | 0 | let r = _mm_sub_ps(neg_x, _mm_mul_ps(kf, ln2)); |
276 | 0 | let mut poly = _mm_add_ps(one, _mm_mul_ps(r, c5)); |
277 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c4, _mm_mul_ps(r, poly)))); |
278 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c3, _mm_mul_ps(r, poly)))); |
279 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c2, _mm_mul_ps(r, poly)))); |
280 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, poly)); |
281 | 0 | let exp_k = _mm_castsi128_ps(_mm_slli_epi32(_mm_add_epi32(k, _mm_set1_epi32(127)), 23)); |
282 | 0 | let exp_neg_x = _mm_mul_ps(poly, exp_k); |
283 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_div_ps(one, _mm_add_ps(one, exp_neg_x))); |
284 | 0 | i += 4; |
285 | 0 | } |
286 | 0 | for j in i..len { result[j] = 1.0 / (1.0 + (-a[j]).exp()); } |
287 | 0 | } |
288 | | |
289 | | #[inline] |
290 | | #[target_feature(enable = "sse2")] |
291 | 0 | unsafe fn gelu(a: &[f32], result: &mut [f32]) { |
292 | | // GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) |
293 | 0 | let len = a.len(); |
294 | 0 | let mut i = 0; |
295 | 0 | let half = _mm_set1_ps(0.5); |
296 | 0 | let one = _mm_set1_ps(1.0); |
297 | 0 | let sqrt_2_pi = _mm_set1_ps(0.797_884_56); |
298 | 0 | let coeff = _mm_set1_ps(0.044_715); |
299 | 0 | while i + 4 <= len { |
300 | 0 | let x = _mm_loadu_ps(a.as_ptr().add(i)); |
301 | 0 | let x3 = _mm_mul_ps(_mm_mul_ps(x, x), x); |
302 | 0 | let inner = _mm_mul_ps(sqrt_2_pi, _mm_add_ps(x, _mm_mul_ps(coeff, x3))); |
303 | 0 | // tanh approximation: (e^2x - 1) / (e^2x + 1) |
304 | 0 | let two_inner = _mm_add_ps(inner, inner); |
305 | 0 | let exp_2x = Self::exp_approx_sse2(two_inner); |
306 | 0 | let tanh_val = _mm_div_ps(_mm_sub_ps(exp_2x, one), _mm_add_ps(exp_2x, one)); |
307 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_mul_ps(half, _mm_mul_ps(x, _mm_add_ps(one, tanh_val)))); |
308 | 0 | i += 4; |
309 | 0 | } |
310 | 0 | for j in i..len { |
311 | 0 | let x = a[j]; |
312 | 0 | result[j] = 0.5 * x * (1.0 + ((0.797_884_56 * (x + 0.044_715 * x * x * x)) as f64).tanh() as f32); |
313 | 0 | } |
314 | 0 | } |
315 | | |
316 | | #[inline] |
317 | | #[target_feature(enable = "sse2")] |
318 | 0 | unsafe fn swish(a: &[f32], result: &mut [f32]) { |
319 | | // swish(x) = x * sigmoid(x) |
320 | 0 | let len = a.len(); |
321 | 0 | let mut i = 0; |
322 | 0 | let one = _mm_set1_ps(1.0); |
323 | 0 | let neg_one = _mm_set1_ps(-1.0); |
324 | 0 | let ln2 = _mm_set1_ps(std::f32::consts::LN_2); |
325 | 0 | let inv_ln2 = _mm_set1_ps(1.0 / std::f32::consts::LN_2); |
326 | 0 | let c2 = _mm_set1_ps(0.5); |
327 | 0 | let c3 = _mm_set1_ps(0.166_666_67); |
328 | 0 | let c4 = _mm_set1_ps(0.041_666_668); |
329 | 0 | let c5 = _mm_set1_ps(0.008_333_334); |
330 | 0 | while i + 4 <= len { |
331 | 0 | let x = _mm_loadu_ps(a.as_ptr().add(i)); |
332 | 0 | let neg_x = _mm_mul_ps(x, neg_one); |
333 | 0 | let k = _mm_cvtps_epi32(_mm_mul_ps(neg_x, inv_ln2)); |
334 | 0 | let kf = _mm_cvtepi32_ps(k); |
335 | 0 | let r = _mm_sub_ps(neg_x, _mm_mul_ps(kf, ln2)); |
336 | 0 | let mut poly = _mm_add_ps(one, _mm_mul_ps(r, c5)); |
337 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c4, _mm_mul_ps(r, poly)))); |
338 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c3, _mm_mul_ps(r, poly)))); |
339 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c2, _mm_mul_ps(r, poly)))); |
340 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, poly)); |
341 | 0 | let exp_k = _mm_castsi128_ps(_mm_slli_epi32(_mm_add_epi32(k, _mm_set1_epi32(127)), 23)); |
342 | 0 | let exp_neg_x = _mm_mul_ps(poly, exp_k); |
343 | 0 | let sigmoid = _mm_div_ps(one, _mm_add_ps(one, exp_neg_x)); |
344 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_mul_ps(x, sigmoid)); |
345 | 0 | i += 4; |
346 | 0 | } |
347 | 0 | for j in i..len { result[j] = a[j] / (1.0 + (-a[j]).exp()); } |
348 | 0 | } |
349 | | |
350 | | #[inline] |
351 | | #[target_feature(enable = "sse2")] |
352 | 0 | unsafe fn tanh(a: &[f32], result: &mut [f32]) { |
353 | | // tanh(x) = (e^2x - 1) / (e^2x + 1) |
354 | 0 | let len = a.len(); |
355 | 0 | let mut i = 0; |
356 | 0 | let one = _mm_set1_ps(1.0); |
357 | 0 | let two = _mm_set1_ps(2.0); |
358 | 0 | while i + 4 <= len { |
359 | 0 | let x = _mm_loadu_ps(a.as_ptr().add(i)); |
360 | 0 | let exp_2x = Self::exp_approx_sse2(_mm_mul_ps(two, x)); |
361 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_div_ps(_mm_sub_ps(exp_2x, one), _mm_add_ps(exp_2x, one))); |
362 | 0 | i += 4; |
363 | 0 | } |
364 | 0 | for j in i..len { |
365 | 0 | let exp_2x = (2.0 * a[j]).exp(); |
366 | 0 | result[j] = (exp_2x - 1.0) / (exp_2x + 1.0); |
367 | 0 | } |
368 | 0 | } |
369 | | |
370 | | #[inline] |
371 | | #[target_feature(enable = "sse2")] |
372 | 0 | unsafe fn sqrt(a: &[f32], result: &mut [f32]) { |
373 | 0 | let len = a.len(); |
374 | 0 | let mut i = 0; |
375 | 0 | while i + 4 <= len { |
376 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_sqrt_ps(_mm_loadu_ps(a.as_ptr().add(i)))); |
377 | 0 | i += 4; |
378 | 0 | } |
379 | 0 | for j in i..len { result[j] = a[j].sqrt(); } |
380 | 0 | } |
381 | | |
382 | | #[inline] |
383 | | #[target_feature(enable = "sse2")] |
384 | 0 | unsafe fn recip(a: &[f32], result: &mut [f32]) { |
385 | 0 | let len = a.len(); |
386 | 0 | let mut i = 0; |
387 | 0 | let one = _mm_set1_ps(1.0); |
388 | 0 | while i + 4 <= len { |
389 | 0 | _mm_storeu_ps(result.as_mut_ptr().add(i), _mm_div_ps(one, _mm_loadu_ps(a.as_ptr().add(i)))); |
390 | 0 | i += 4; |
391 | 0 | } |
392 | 0 | for j in i..len { result[j] = a[j].recip(); } |
393 | 0 | } |
394 | | |
395 | 0 | unsafe fn ln(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::ln(a, result); } |
396 | 0 | unsafe fn log2(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::log2(a, result); } |
397 | 0 | unsafe fn log10(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::log10(a, result); } |
398 | 0 | unsafe fn sin(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::sin(a, result); } |
399 | 0 | unsafe fn cos(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::cos(a, result); } |
400 | 0 | unsafe fn tan(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::tan(a, result); } |
401 | 0 | unsafe fn floor(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::floor(a, result); } |
402 | 0 | unsafe fn ceil(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::ceil(a, result); } |
403 | 0 | unsafe fn round(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::round(a, result); } |
404 | | } |
405 | | |
406 | | impl Sse2Backend { |
407 | | /// SSE2 exp approximation helper |
408 | | #[inline] |
409 | | #[target_feature(enable = "sse2")] |
410 | 0 | unsafe fn exp_approx_sse2(x: __m128) -> __m128 { |
411 | 0 | let ln2 = _mm_set1_ps(std::f32::consts::LN_2); |
412 | 0 | let inv_ln2 = _mm_set1_ps(1.0 / std::f32::consts::LN_2); |
413 | 0 | let one = _mm_set1_ps(1.0); |
414 | 0 | let c2 = _mm_set1_ps(0.5); |
415 | 0 | let c3 = _mm_set1_ps(0.166_666_67); |
416 | 0 | let c4 = _mm_set1_ps(0.041_666_668); |
417 | 0 | let c5 = _mm_set1_ps(0.008_333_334); |
418 | 0 | let k = _mm_cvtps_epi32(_mm_mul_ps(x, inv_ln2)); |
419 | 0 | let kf = _mm_cvtepi32_ps(k); |
420 | 0 | let r = _mm_sub_ps(x, _mm_mul_ps(kf, ln2)); |
421 | 0 | let mut poly = _mm_add_ps(one, _mm_mul_ps(r, c5)); |
422 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c4, _mm_mul_ps(r, poly)))); |
423 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c3, _mm_mul_ps(r, poly)))); |
424 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, _mm_add_ps(c2, _mm_mul_ps(r, poly)))); |
425 | 0 | poly = _mm_add_ps(one, _mm_mul_ps(r, poly)); |
426 | 0 | let exp_k = _mm_castsi128_ps(_mm_slli_epi32(_mm_add_epi32(k, _mm_set1_epi32(127)), 23)); |
427 | 0 | _mm_mul_ps(poly, exp_k) |
428 | 0 | } |
429 | | } |
430 | | |
431 | | #[cfg(test)] |
432 | | mod tests { |
433 | | use super::*; |
434 | | |
435 | | #[test] |
436 | | fn test_sse2_add() { |
437 | | let a = [1.0, 2.0, 3.0, 4.0, 5.0]; |
438 | | let b = [10.0, 20.0, 30.0, 40.0, 50.0]; |
439 | | let mut result = [0.0f32; 5]; |
440 | | unsafe { Sse2Backend::add(&a, &b, &mut result); } |
441 | | assert_eq!(result, [11.0, 22.0, 33.0, 44.0, 55.0]); |
442 | | } |
443 | | |
444 | | #[test] |
445 | | fn test_sse2_sub() { |
446 | | let a = [10.0, 20.0, 30.0, 40.0, 50.0]; |
447 | | let b = [1.0, 2.0, 3.0, 4.0, 5.0]; |
448 | | let mut result = [0.0f32; 5]; |
449 | | unsafe { Sse2Backend::sub(&a, &b, &mut result); } |
450 | | assert_eq!(result, [9.0, 18.0, 27.0, 36.0, 45.0]); |
451 | | } |
452 | | |
453 | | #[test] |
454 | | fn test_sse2_mul() { |
455 | | let a = [1.0, 2.0, 3.0, 4.0, 5.0]; |
456 | | let b = [2.0, 3.0, 4.0, 5.0, 6.0]; |
457 | | let mut result = [0.0f32; 5]; |
458 | | unsafe { Sse2Backend::mul(&a, &b, &mut result); } |
459 | | assert_eq!(result, [2.0, 6.0, 12.0, 20.0, 30.0]); |
460 | | } |
461 | | |
462 | | #[test] |
463 | | fn test_sse2_div() { |
464 | | let a = [10.0, 20.0, 30.0, 40.0, 50.0]; |
465 | | let b = [2.0, 4.0, 5.0, 8.0, 10.0]; |
466 | | let mut result = [0.0f32; 5]; |
467 | | unsafe { Sse2Backend::div(&a, &b, &mut result); } |
468 | | assert_eq!(result, [5.0, 5.0, 6.0, 5.0, 5.0]); |
469 | | } |
470 | | |
471 | | #[test] |
472 | | fn test_sse2_dot() { |
473 | | let a = [1.0, 2.0, 3.0, 4.0, 5.0]; |
474 | | let b = [2.0, 3.0, 4.0, 5.0, 6.0]; |
475 | | let result = unsafe { Sse2Backend::dot(&a, &b) }; |
476 | | assert!((result - 70.0).abs() < 1e-6); |
477 | | } |
478 | | |
479 | | #[test] |
480 | | fn test_sse2_sum() { |
481 | | let a = [1.0, 2.0, 3.0, 4.0, 5.0]; |
482 | | let result = unsafe { Sse2Backend::sum(&a) }; |
483 | | assert!((result - 15.0).abs() < 1e-6); |
484 | | } |
485 | | |
486 | | #[test] |
487 | | fn test_sse2_max() { |
488 | | let a = [1.0, 5.0, 3.0, 2.0, 4.0]; |
489 | | let result = unsafe { Sse2Backend::max(&a) }; |
490 | | assert!((result - 5.0).abs() < 1e-6); |
491 | | } |
492 | | |
493 | | #[test] |
494 | | fn test_sse2_min() { |
495 | | let a = [5.0, 1.0, 3.0, 2.0, 4.0]; |
496 | | let result = unsafe { Sse2Backend::min(&a) }; |
497 | | assert!((result - 1.0).abs() < 1e-6); |
498 | | } |
499 | | |
500 | | #[test] |
501 | | fn test_sse2_argmax() { |
502 | | let a = [1.0, 5.0, 3.0, 2.0, 4.0]; |
503 | | let result = unsafe { Sse2Backend::argmax(&a) }; |
504 | | assert_eq!(result, 1); |
505 | | } |
506 | | |
507 | | #[test] |
508 | | fn test_sse2_argmin() { |
509 | | let a = [5.0, 1.0, 3.0, 2.0, 4.0]; |
510 | | let result = unsafe { Sse2Backend::argmin(&a) }; |
511 | | assert_eq!(result, 1); |
512 | | } |
513 | | |
514 | | #[test] |
515 | | fn test_sse2_norm_linf() { |
516 | | let a = [-5.0, 1.0, 3.0, 2.0, -4.0]; |
517 | | let result = unsafe { Sse2Backend::norm_linf(&a) }; |
518 | | assert!((result - 5.0).abs() < 1e-6); |
519 | | } |
520 | | |
521 | | #[test] |
522 | | fn test_sse2_scale() { |
523 | | let a = [1.0, 2.0, 3.0, 4.0, 5.0]; |
524 | | let mut result = [0.0f32; 5]; |
525 | | unsafe { Sse2Backend::scale(&a, 2.0, &mut result); } |
526 | | assert_eq!(result, [2.0, 4.0, 6.0, 8.0, 10.0]); |
527 | | } |
528 | | |
529 | | #[test] |
530 | | fn test_sse2_abs() { |
531 | | let a = [-1.0, 2.0, -3.0, 4.0, -5.0]; |
532 | | let mut result = [0.0f32; 5]; |
533 | | unsafe { Sse2Backend::abs(&a, &mut result); } |
534 | | assert_eq!(result, [1.0, 2.0, 3.0, 4.0, 5.0]); |
535 | | } |
536 | | |
537 | | #[test] |
538 | | fn test_sse2_clamp() { |
539 | | let a = [-1.0, 0.5, 1.5, 2.0, 3.0]; |
540 | | let mut result = [0.0f32; 5]; |
541 | | unsafe { Sse2Backend::clamp(&a, 0.0, 1.0, &mut result); } |
542 | | assert_eq!(result, [0.0, 0.5, 1.0, 1.0, 1.0]); |
543 | | } |
544 | | |
545 | | #[test] |
546 | | fn test_sse2_relu() { |
547 | | let a = [-1.0, 0.0, 1.0, -2.0, 3.0]; |
548 | | let mut result = [0.0f32; 5]; |
549 | | unsafe { Sse2Backend::relu(&a, &mut result); } |
550 | | assert_eq!(result, [0.0, 0.0, 1.0, 0.0, 3.0]); |
551 | | } |
552 | | |
553 | | #[test] |
554 | | fn test_sse2_exp() { |
555 | | let a = [0.0, 1.0, -1.0, 2.0]; |
556 | | let mut result = [0.0f32; 4]; |
557 | | unsafe { Sse2Backend::exp(&a, &mut result); } |
558 | | assert!((result[0] - 1.0).abs() < 0.05); |
559 | | assert!((result[1] - std::f32::consts::E).abs() < 0.1); |
560 | | } |
561 | | |
562 | | #[test] |
563 | | fn test_sse2_sigmoid() { |
564 | | let a = [0.0, 1.0, -1.0, 10.0]; |
565 | | let mut result = [0.0f32; 4]; |
566 | | unsafe { Sse2Backend::sigmoid(&a, &mut result); } |
567 | | assert!((result[0] - 0.5).abs() < 0.01); |
568 | | assert!(result[1] > 0.5); |
569 | | assert!(result[2] < 0.5); |
570 | | } |
571 | | |
572 | | #[test] |
573 | | fn test_sse2_sqrt() { |
574 | | let a = [1.0, 4.0, 9.0, 16.0, 25.0]; |
575 | | let mut result = [0.0f32; 5]; |
576 | | unsafe { Sse2Backend::sqrt(&a, &mut result); } |
577 | | assert_eq!(result, [1.0, 2.0, 3.0, 4.0, 5.0]); |
578 | | } |
579 | | |
580 | | #[test] |
581 | | fn test_sse2_sum_kahan() { |
582 | | let a: Vec<f32> = (1..=16).map(|i| i as f32).collect(); |
583 | | let result = unsafe { Sse2Backend::sum_kahan(&a) }; |
584 | | assert!((result - 136.0).abs() < 1e-3); |
585 | | } |
586 | | |
587 | | #[test] |
588 | | fn test_sse2_norm_l2() { |
589 | | let a = vec![3.0, 4.0]; |
590 | | let result = unsafe { Sse2Backend::norm_l2(&a) }; |
591 | | assert!((result - 5.0).abs() < 1e-5); |
592 | | } |
593 | | |
594 | | #[test] |
595 | | fn test_sse2_norm_l1() { |
596 | | let a = vec![-1.0, 2.0, -3.0, 4.0]; |
597 | | let result = unsafe { Sse2Backend::norm_l1(&a) }; |
598 | | assert!((result - 10.0).abs() < 1e-5); |
599 | | } |
600 | | |
601 | | #[test] |
602 | | fn test_sse2_lerp() { |
603 | | let a = vec![0.0; 16]; |
604 | | let b = vec![10.0; 16]; |
605 | | let mut result = vec![0.0; 16]; |
606 | | unsafe { Sse2Backend::lerp(&a, &b, 0.5, &mut result); } |
607 | | assert!(result.iter().all(|&x| (x - 5.0).abs() < 1e-5)); |
608 | | } |
609 | | |
610 | | #[test] |
611 | | fn test_sse2_fma() { |
612 | | let a = vec![2.0; 16]; |
613 | | let b = vec![3.0; 16]; |
614 | | let c = vec![1.0; 16]; |
615 | | let mut result = vec![0.0; 16]; |
616 | | unsafe { Sse2Backend::fma(&a, &b, &c, &mut result); } |
617 | | assert!(result.iter().all(|&x| (x - 7.0).abs() < 1e-5)); |
618 | | } |
619 | | |
620 | | #[test] |
621 | | fn test_sse2_gelu() { |
622 | | let a = vec![0.0, 1.0]; |
623 | | let mut result = vec![0.0; 2]; |
624 | | unsafe { Sse2Backend::gelu(&a, &mut result); } |
625 | | assert!((result[0]).abs() < 1e-5); |
626 | | assert!((result[1] - 0.841_192).abs() < 1e-2); |
627 | | } |
628 | | |
629 | | #[test] |
630 | | fn test_sse2_swish() { |
631 | | let a = vec![0.0, 1.0]; |
632 | | let mut result = vec![0.0; 2]; |
633 | | unsafe { Sse2Backend::swish(&a, &mut result); } |
634 | | assert!((result[0]).abs() < 1e-5); |
635 | | assert!((result[1] - 0.731_059).abs() < 1e-2); |
636 | | } |
637 | | |
638 | | #[test] |
639 | | fn test_sse2_tanh() { |
640 | | let a = vec![0.0, 1.0]; |
641 | | let mut result = vec![0.0; 2]; |
642 | | unsafe { Sse2Backend::tanh(&a, &mut result); } |
643 | | assert!((result[0]).abs() < 1e-5); |
644 | | assert!((result[1] - 0.761_594_2).abs() < 1e-2); |
645 | | } |
646 | | |
647 | | #[test] |
648 | | fn test_sse2_recip() { |
649 | | let a = vec![2.0, 4.0, 5.0]; |
650 | | let mut result = vec![0.0; 3]; |
651 | | unsafe { Sse2Backend::recip(&a, &mut result); } |
652 | | assert!((result[0] - 0.5).abs() < 1e-5); |
653 | | assert!((result[1] - 0.25).abs() < 1e-5); |
654 | | assert!((result[2] - 0.2).abs() < 1e-5); |
655 | | } |
656 | | |
657 | | #[test] |
658 | | fn test_sse2_transcendental() { |
659 | | let a = vec![1.0, std::f32::consts::E, 10.0]; |
660 | | let mut ln_result = vec![0.0; 3]; |
661 | | let mut log10_result = vec![0.0; 3]; |
662 | | unsafe { |
663 | | Sse2Backend::ln(&a, &mut ln_result); |
664 | | Sse2Backend::log10(&a, &mut log10_result); |
665 | | } |
666 | | assert!((ln_result[0]).abs() < 1e-5); |
667 | | assert!((ln_result[1] - 1.0).abs() < 1e-4); |
668 | | assert!((log10_result[2] - 1.0).abs() < 1e-5); |
669 | | } |
670 | | |
671 | | #[test] |
672 | | fn test_sse2_trig() { |
673 | | let a = vec![0.0, std::f32::consts::FRAC_PI_2]; |
674 | | let mut sin_result = vec![0.0; 2]; |
675 | | let mut cos_result = vec![0.0; 2]; |
676 | | unsafe { |
677 | | Sse2Backend::sin(&a, &mut sin_result); |
678 | | Sse2Backend::cos(&a, &mut cos_result); |
679 | | } |
680 | | assert!((sin_result[0]).abs() < 1e-5); |
681 | | assert!((sin_result[1] - 1.0).abs() < 1e-5); |
682 | | assert!((cos_result[0] - 1.0).abs() < 1e-5); |
683 | | } |
684 | | |
685 | | #[test] |
686 | | fn test_sse2_rounding() { |
687 | | let a = vec![1.3, 1.5, 1.7, -1.3, -1.5, -1.7]; |
688 | | let mut floor_result = vec![0.0; 6]; |
689 | | let mut ceil_result = vec![0.0; 6]; |
690 | | unsafe { |
691 | | Sse2Backend::floor(&a, &mut floor_result); |
692 | | Sse2Backend::ceil(&a, &mut ceil_result); |
693 | | } |
694 | | assert_eq!(floor_result, vec![1.0, 1.0, 1.0, -2.0, -2.0, -2.0]); |
695 | | assert_eq!(ceil_result, vec![2.0, 2.0, 2.0, -1.0, -1.0, -1.0]); |
696 | | } |
697 | | } |