Coverage Report

Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/home/noah/src/trueno/src/backends/avx512/mod.rs
Line
Count
Source
1
//! AVX-512 backend implementation (x86_64 advanced SIMD)
2
//!
3
//! This backend uses AVX-512 intrinsics for 512-bit SIMD operations.
4
//! AVX-512 is available on Intel Skylake-X/Sapphire Rapids (2017+) and AMD Zen 4 (2022+) CPUs.
5
//!
6
//! # Performance
7
//!
8
//! Expected speedup: 16x for operations on f32 vectors (16 elements per register)
9
//! This provides 2x improvement over AVX2 (8 elements) and ~16x over scalar.
10
//!
11
//! # Safety
12
//!
13
//! All AVX-512 intrinsics are marked `unsafe` by Rust. This module carefully isolates
14
//! all unsafe code and verifies correctness through comprehensive testing.
15
16
mod ops;
17
18
#[cfg(target_arch = "x86_64")]
19
use std::arch::x86_64::*;
20
21
use super::VectorBackend;
22
23
/// AVX-512 backend (512-bit SIMD for x86_64)
24
pub struct Avx512Backend;
25
26
impl VectorBackend for Avx512Backend {
27
    #[inline]
28
    #[target_feature(enable = "avx512f")]
29
0
    unsafe fn add(a: &[f32], b: &[f32], result: &mut [f32]) {
30
0
        ops::arithmetic::add(a, b, result);
31
0
    }
32
33
    #[inline]
34
    #[target_feature(enable = "avx512f")]
35
0
    unsafe fn sub(a: &[f32], b: &[f32], result: &mut [f32]) {
36
0
        ops::arithmetic::sub(a, b, result);
37
0
    }
38
39
    #[inline]
40
    #[target_feature(enable = "avx512f")]
41
0
    unsafe fn mul(a: &[f32], b: &[f32], result: &mut [f32]) {
42
0
        ops::arithmetic::mul(a, b, result);
43
0
    }
44
45
    #[inline]
46
    #[target_feature(enable = "avx512f")]
47
0
    unsafe fn div(a: &[f32], b: &[f32], result: &mut [f32]) {
48
0
        ops::arithmetic::div(a, b, result);
49
0
    }
50
51
    #[inline]
52
    #[target_feature(enable = "avx512f")]
53
0
    unsafe fn dot(a: &[f32], b: &[f32]) -> f32 {
54
0
        ops::reductions::dot(a, b)
55
0
    }
56
57
    #[inline]
58
    #[target_feature(enable = "avx512f")]
59
0
    unsafe fn sum(a: &[f32]) -> f32 {
60
0
        ops::reductions::sum(a)
61
0
    }
62
63
    #[inline]
64
    #[target_feature(enable = "avx512f")]
65
0
    unsafe fn max(a: &[f32]) -> f32 {
66
0
        ops::reductions::max(a)
67
0
    }
68
69
    #[inline]
70
    #[target_feature(enable = "avx512f")]
71
0
    unsafe fn min(a: &[f32]) -> f32 {
72
0
        ops::reductions::min(a)
73
0
    }
74
75
    #[inline]
76
    #[target_feature(enable = "avx512f")]
77
0
    unsafe fn argmax(a: &[f32]) -> usize {
78
0
        ops::reductions::argmax(a)
79
0
    }
80
81
    #[inline]
82
    #[target_feature(enable = "avx512f")]
83
0
    unsafe fn argmin(a: &[f32]) -> usize {
84
0
        ops::reductions::argmin(a)
85
0
    }
86
87
0
    unsafe fn sum_kahan(a: &[f32]) -> f32 {
88
0
        ops::reductions::sum_kahan(a)
89
0
    }
90
91
    #[inline]
92
    #[target_feature(enable = "avx512f")]
93
0
    unsafe fn norm_l2(a: &[f32]) -> f32 {
94
0
        if a.is_empty() { return 0.0; }
95
0
        let len = a.len();
96
0
        let mut i = 0;
97
0
        let mut acc = _mm512_setzero_ps();
98
0
        while i + 16 <= len {
99
0
            let va = _mm512_loadu_ps(a.as_ptr().add(i));
100
0
            acc = _mm512_add_ps(acc, _mm512_mul_ps(va, va));
101
0
            i += 16;
102
0
        }
103
0
        let mut sum_sq = _mm512_reduce_add_ps(acc);
104
0
        for &val in &a[i..] { sum_sq += val * val; }
105
0
        sum_sq.sqrt()
106
0
    }
107
108
    #[inline]
109
    #[target_feature(enable = "avx512f")]
110
0
    unsafe fn norm_l1(a: &[f32]) -> f32 {
111
0
        let len = a.len();
112
0
        let mut i = 0;
113
0
        let sign_mask = _mm512_set1_ps(f32::from_bits(0x7FFF_FFFF));
114
0
        let mut acc = _mm512_setzero_ps();
115
0
        while i + 16 <= len {
116
0
            acc = _mm512_add_ps(acc, _mm512_and_ps(_mm512_loadu_ps(a.as_ptr().add(i)), sign_mask));
117
0
            i += 16;
118
0
        }
119
0
        let mut result = _mm512_reduce_add_ps(acc);
120
0
        for &val in &a[i..] { result += val.abs(); }
121
0
        result
122
0
    }
123
124
    #[inline]
125
    #[target_feature(enable = "avx512f")]
126
0
    unsafe fn norm_linf(a: &[f32]) -> f32 {
127
0
        let len = a.len();
128
0
        let mut i = 0;
129
0
        let sign_mask = _mm512_set1_ps(f32::from_bits(0x7FFF_FFFF));
130
0
        let mut max_vec = _mm512_setzero_ps();
131
0
        while i + 16 <= len {
132
0
            max_vec = _mm512_max_ps(max_vec, _mm512_and_ps(_mm512_loadu_ps(a.as_ptr().add(i)), sign_mask));
133
0
            i += 16;
134
0
        }
135
0
        let mut result = _mm512_reduce_max_ps(max_vec);
136
0
        for &val in &a[i..] { let abs_val = val.abs(); if abs_val > result { result = abs_val; } }
137
0
        result
138
0
    }
139
140
    #[inline]
141
    #[target_feature(enable = "avx512f")]
142
0
    unsafe fn scale(a: &[f32], scalar: f32, result: &mut [f32]) {
143
0
        let len = a.len();
144
0
        let mut i = 0;
145
0
        let scalar_vec = _mm512_set1_ps(scalar);
146
0
        while i + 16 <= len {
147
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_mul_ps(_mm512_loadu_ps(a.as_ptr().add(i)), scalar_vec));
148
0
            i += 16;
149
0
        }
150
0
        for j in i..len { result[j] = a[j] * scalar; }
151
0
    }
152
153
    #[inline]
154
    #[target_feature(enable = "avx512f")]
155
0
    unsafe fn abs(a: &[f32], result: &mut [f32]) {
156
0
        let len = a.len();
157
0
        let mut i = 0;
158
0
        let sign_mask = _mm512_set1_ps(f32::from_bits(0x7FFF_FFFF));
159
0
        while i + 16 <= len {
160
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_and_ps(_mm512_loadu_ps(a.as_ptr().add(i)), sign_mask));
161
0
            i += 16;
162
0
        }
163
0
        for j in i..len { result[j] = a[j].abs(); }
164
0
    }
165
166
    #[inline]
167
    #[target_feature(enable = "avx512f")]
168
0
    unsafe fn clamp(a: &[f32], min_val: f32, max_val: f32, result: &mut [f32]) {
169
0
        let len = a.len();
170
0
        let mut i = 0;
171
0
        let min_vec = _mm512_set1_ps(min_val);
172
0
        let max_vec = _mm512_set1_ps(max_val);
173
0
        while i + 16 <= len {
174
0
            let va = _mm512_loadu_ps(a.as_ptr().add(i));
175
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_min_ps(_mm512_max_ps(va, min_vec), max_vec));
176
0
            i += 16;
177
0
        }
178
0
        for j in i..len { result[j] = a[j].max(min_val).min(max_val); }
179
0
    }
180
181
    #[inline]
182
    #[target_feature(enable = "avx512f")]
183
0
    unsafe fn lerp(a: &[f32], b: &[f32], t: f32, result: &mut [f32]) {
184
0
        let len = a.len();
185
0
        let mut i = 0;
186
0
        let t_vec = _mm512_set1_ps(t);
187
0
        while i + 16 <= len {
188
0
            let va = _mm512_loadu_ps(a.as_ptr().add(i));
189
0
            let vb = _mm512_loadu_ps(b.as_ptr().add(i));
190
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_fmadd_ps(t_vec, _mm512_sub_ps(vb, va), va));
191
0
            i += 16;
192
0
        }
193
0
        for j in i..len { result[j] = a[j] + t * (b[j] - a[j]); }
194
0
    }
195
196
    #[inline]
197
    #[target_feature(enable = "avx512f")]
198
0
    unsafe fn fma(a: &[f32], b: &[f32], c: &[f32], result: &mut [f32]) {
199
0
        let len = a.len();
200
0
        let mut i = 0;
201
0
        while i + 16 <= len {
202
0
            let va = _mm512_loadu_ps(a.as_ptr().add(i));
203
0
            let vb = _mm512_loadu_ps(b.as_ptr().add(i));
204
0
            let vc = _mm512_loadu_ps(c.as_ptr().add(i));
205
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_fmadd_ps(va, vb, vc));
206
0
            i += 16;
207
0
        }
208
0
        for j in i..len { result[j] = a[j] * b[j] + c[j]; }
209
0
    }
210
211
    #[inline]
212
    #[target_feature(enable = "avx512f")]
213
0
    unsafe fn relu(a: &[f32], result: &mut [f32]) {
214
0
        let len = a.len();
215
0
        let mut i = 0;
216
0
        let zero = _mm512_setzero_ps();
217
0
        while i + 16 <= len {
218
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_max_ps(_mm512_loadu_ps(a.as_ptr().add(i)), zero));
219
0
            i += 16;
220
0
        }
221
0
        for j in i..len { result[j] = a[j].max(0.0); }
222
0
    }
223
224
    #[inline]
225
    #[target_feature(enable = "avx512f")]
226
0
    unsafe fn exp(a: &[f32], result: &mut [f32]) {
227
0
        let len = a.len();
228
0
        let mut i = 0;
229
0
        let ln2 = _mm512_set1_ps(std::f32::consts::LN_2);
230
0
        let inv_ln2 = _mm512_set1_ps(1.0 / std::f32::consts::LN_2);
231
0
        let one = _mm512_set1_ps(1.0);
232
0
        let c2 = _mm512_set1_ps(0.5);
233
0
        let c3 = _mm512_set1_ps(0.166_666_67);
234
0
        let c4 = _mm512_set1_ps(0.041_666_668);
235
0
        let c5 = _mm512_set1_ps(0.008_333_334);
236
0
        while i + 16 <= len {
237
0
            let x = _mm512_loadu_ps(a.as_ptr().add(i));
238
0
            let k = _mm512_cvtps_epi32(_mm512_mul_ps(x, inv_ln2));
239
0
            let kf = _mm512_cvtepi32_ps(k);
240
0
            let r = _mm512_sub_ps(x, _mm512_mul_ps(kf, ln2));
241
0
            let mut poly = _mm512_fmadd_ps(r, c5, one);
242
0
            poly = _mm512_fmadd_ps(r, _mm512_fmadd_ps(r, poly, c4), one);
243
0
            poly = _mm512_fmadd_ps(r, _mm512_fmadd_ps(r, poly, c3), one);
244
0
            poly = _mm512_fmadd_ps(r, _mm512_fmadd_ps(r, poly, c2), one);
245
0
            poly = _mm512_fmadd_ps(r, poly, one);
246
0
            let exp_k = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_add_epi32(k, _mm512_set1_epi32(127)), 23));
247
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_mul_ps(poly, exp_k));
248
0
            i += 16;
249
0
        }
250
0
        for j in i..len { result[j] = a[j].exp(); }
251
0
    }
252
253
    #[inline]
254
    #[target_feature(enable = "avx512f")]
255
0
    unsafe fn sigmoid(a: &[f32], result: &mut [f32]) {
256
        // sigmoid(x) = 1 / (1 + exp(-x))
257
0
        let len = a.len();
258
0
        for j in 0..len { result[j] = 1.0 / (1.0 + (-a[j]).exp()); }
259
0
    }
260
261
    #[inline]
262
    #[target_feature(enable = "avx512f")]
263
0
    unsafe fn gelu(a: &[f32], result: &mut [f32]) {
264
0
        for j in 0..a.len() {
265
0
            let x = a[j];
266
0
            let inner = 0.797_884_56 * (x + 0.044_715 * x * x * x);
267
0
            result[j] = 0.5 * x * (1.0 + inner.tanh());
268
0
        }
269
0
    }
270
271
    #[inline]
272
    #[target_feature(enable = "avx512f")]
273
0
    unsafe fn swish(a: &[f32], result: &mut [f32]) {
274
0
        for j in 0..a.len() { result[j] = a[j] / (1.0 + (-a[j]).exp()); }
275
0
    }
276
277
    #[inline]
278
    #[target_feature(enable = "avx512f")]
279
0
    unsafe fn tanh(a: &[f32], result: &mut [f32]) {
280
0
        for j in 0..a.len() { result[j] = a[j].tanh(); }
281
0
    }
282
283
    #[inline]
284
    #[target_feature(enable = "avx512f")]
285
0
    unsafe fn sqrt(a: &[f32], result: &mut [f32]) {
286
0
        let len = a.len();
287
0
        let mut i = 0;
288
0
        while i + 16 <= len {
289
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_sqrt_ps(_mm512_loadu_ps(a.as_ptr().add(i))));
290
0
            i += 16;
291
0
        }
292
0
        for j in i..len { result[j] = a[j].sqrt(); }
293
0
    }
294
295
    #[inline]
296
    #[target_feature(enable = "avx512f")]
297
0
    unsafe fn recip(a: &[f32], result: &mut [f32]) {
298
0
        let len = a.len();
299
0
        let mut i = 0;
300
0
        let one = _mm512_set1_ps(1.0);
301
0
        while i + 16 <= len {
302
0
            _mm512_storeu_ps(result.as_mut_ptr().add(i), _mm512_div_ps(one, _mm512_loadu_ps(a.as_ptr().add(i))));
303
0
            i += 16;
304
0
        }
305
0
        for j in i..len { result[j] = a[j].recip(); }
306
0
    }
307
308
0
    unsafe fn ln(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::ln(a, result); }
309
0
    unsafe fn log2(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::log2(a, result); }
310
0
    unsafe fn log10(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::log10(a, result); }
311
0
    unsafe fn sin(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::sin(a, result); }
312
0
    unsafe fn cos(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::cos(a, result); }
313
0
    unsafe fn tan(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::tan(a, result); }
314
315
0
    unsafe fn floor(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::floor(a, result); }
316
0
    unsafe fn ceil(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::ceil(a, result); }
317
0
    unsafe fn round(a: &[f32], result: &mut [f32]) { super::scalar::ScalarBackend::round(a, result); }
318
}
319
320
#[cfg(all(test, target_arch = "x86_64"))]
321
mod tests {
322
    use super::*;
323
    use crate::backends::scalar::ScalarBackend;
324
325
    fn avx512_test<F>(test_fn: F) where F: FnOnce() {
326
        if is_x86_feature_detected!("avx512f") { test_fn(); }
327
        else { println!("Skipping AVX-512 test (CPU does not support avx512f)"); }
328
    }
329
330
    #[test]
331
    fn test_avx512_add() {
332
        avx512_test(|| {
333
            let a = vec![1.0; 32];
334
            let b = vec![2.0; 32];
335
            let mut result = vec![0.0; 32];
336
            unsafe { Avx512Backend::add(&a, &b, &mut result); }
337
            assert!(result.iter().all(|&x| (x - 3.0).abs() < 1e-6));
338
        });
339
    }
340
341
    #[test]
342
    fn test_avx512_sub() {
343
        avx512_test(|| {
344
            let a = vec![5.0; 32];
345
            let b = vec![2.0; 32];
346
            let mut result = vec![0.0; 32];
347
            unsafe { Avx512Backend::sub(&a, &b, &mut result); }
348
            assert!(result.iter().all(|&x| (x - 3.0).abs() < 1e-6));
349
        });
350
    }
351
352
    #[test]
353
    fn test_avx512_mul() {
354
        avx512_test(|| {
355
            let a = vec![2.0; 32];
356
            let b = vec![3.0; 32];
357
            let mut result = vec![0.0; 32];
358
            unsafe { Avx512Backend::mul(&a, &b, &mut result); }
359
            assert!(result.iter().all(|&x| (x - 6.0).abs() < 1e-6));
360
        });
361
    }
362
363
    #[test]
364
    fn test_avx512_div() {
365
        avx512_test(|| {
366
            let a = vec![6.0; 32];
367
            let b = vec![2.0; 32];
368
            let mut result = vec![0.0; 32];
369
            unsafe { Avx512Backend::div(&a, &b, &mut result); }
370
            assert!(result.iter().all(|&x| (x - 3.0).abs() < 1e-6));
371
        });
372
    }
373
374
    #[test]
375
    fn test_avx512_dot() {
376
        avx512_test(|| {
377
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
378
            let b: Vec<f32> = (1..=32).map(|i| i as f32).collect();
379
            let result = unsafe { Avx512Backend::dot(&a, &b) };
380
            let expected: f32 = (1..=32).map(|i| (i * i) as f32).sum();
381
            assert!((result - expected).abs() < 1e-3);
382
        });
383
    }
384
385
    #[test]
386
    fn test_avx512_sum() {
387
        avx512_test(|| {
388
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
389
            let result = unsafe { Avx512Backend::sum(&a) };
390
            assert!((result - 528.0).abs() < 1e-3);
391
        });
392
    }
393
394
    #[test]
395
    fn test_avx512_max() {
396
        avx512_test(|| {
397
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
398
            let result = unsafe { Avx512Backend::max(&a) };
399
            assert!((result - 32.0).abs() < 1e-6);
400
        });
401
    }
402
403
    #[test]
404
    fn test_avx512_min() {
405
        avx512_test(|| {
406
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
407
            let result = unsafe { Avx512Backend::min(&a) };
408
            assert!((result - 1.0).abs() < 1e-6);
409
        });
410
    }
411
412
    #[test]
413
    fn test_avx512_argmax() {
414
        avx512_test(|| {
415
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
416
            let result = unsafe { Avx512Backend::argmax(&a) };
417
            assert_eq!(result, 31);
418
        });
419
    }
420
421
    #[test]
422
    fn test_avx512_argmin() {
423
        avx512_test(|| {
424
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
425
            let result = unsafe { Avx512Backend::argmin(&a) };
426
            assert_eq!(result, 0);
427
        });
428
    }
429
430
    #[test]
431
    fn test_avx512_sum_kahan() {
432
        avx512_test(|| {
433
            let a: Vec<f32> = (1..=32).map(|i| i as f32).collect();
434
            let result = unsafe { Avx512Backend::sum_kahan(&a) };
435
            assert!((result - 528.0).abs() < 1e-3);
436
        });
437
    }
438
439
    #[test]
440
    fn test_avx512_norm_l2() {
441
        avx512_test(|| {
442
            let a = vec![3.0, 4.0];
443
            let result = unsafe { Avx512Backend::norm_l2(&a) };
444
            assert!((result - 5.0).abs() < 1e-5);
445
        });
446
    }
447
448
    #[test]
449
    fn test_avx512_norm_l1() {
450
        avx512_test(|| {
451
            let a = vec![-1.0, 2.0, -3.0, 4.0];
452
            let result = unsafe { Avx512Backend::norm_l1(&a) };
453
            assert!((result - 10.0).abs() < 1e-5);
454
        });
455
    }
456
457
    #[test]
458
    fn test_avx512_norm_linf() {
459
        avx512_test(|| {
460
            let a = vec![-5.0, 2.0, -3.0, 4.0];
461
            let result = unsafe { Avx512Backend::norm_linf(&a) };
462
            assert!((result - 5.0).abs() < 1e-5);
463
        });
464
    }
465
466
    #[test]
467
    fn test_avx512_scale() {
468
        avx512_test(|| {
469
            let a = vec![1.0; 32];
470
            let mut result = vec![0.0; 32];
471
            unsafe { Avx512Backend::scale(&a, 3.0, &mut result); }
472
            assert!(result.iter().all(|&x| (x - 3.0).abs() < 1e-6));
473
        });
474
    }
475
476
    #[test]
477
    fn test_avx512_abs() {
478
        avx512_test(|| {
479
            let a = vec![-1.0, 2.0, -3.0, 4.0];
480
            let mut result = vec![0.0; 4];
481
            unsafe { Avx512Backend::abs(&a, &mut result); }
482
            assert_eq!(result, vec![1.0, 2.0, 3.0, 4.0]);
483
        });
484
    }
485
486
    #[test]
487
    fn test_avx512_clamp() {
488
        avx512_test(|| {
489
            let a = vec![0.0, 5.0, 10.0, 15.0];
490
            let mut result = vec![0.0; 4];
491
            unsafe { Avx512Backend::clamp(&a, 2.0, 12.0, &mut result); }
492
            assert_eq!(result, vec![2.0, 5.0, 10.0, 12.0]);
493
        });
494
    }
495
496
    #[test]
497
    fn test_avx512_lerp() {
498
        avx512_test(|| {
499
            let a = vec![0.0; 32];
500
            let b = vec![10.0; 32];
501
            let mut result = vec![0.0; 32];
502
            unsafe { Avx512Backend::lerp(&a, &b, 0.5, &mut result); }
503
            assert!(result.iter().all(|&x| (x - 5.0).abs() < 1e-5));
504
        });
505
    }
506
507
    #[test]
508
    fn test_avx512_fma() {
509
        avx512_test(|| {
510
            let a = vec![2.0; 32];
511
            let b = vec![3.0; 32];
512
            let c = vec![1.0; 32];
513
            let mut result = vec![0.0; 32];
514
            unsafe { Avx512Backend::fma(&a, &b, &c, &mut result); }
515
            assert!(result.iter().all(|&x| (x - 7.0).abs() < 1e-5));
516
        });
517
    }
518
519
    #[test]
520
    fn test_avx512_relu() {
521
        avx512_test(|| {
522
            let a = vec![-1.0, 0.0, 1.0, 2.0];
523
            let mut result = vec![0.0; 4];
524
            unsafe { Avx512Backend::relu(&a, &mut result); }
525
            assert_eq!(result, vec![0.0, 0.0, 1.0, 2.0]);
526
        });
527
    }
528
529
    #[test]
530
    fn test_avx512_exp() {
531
        avx512_test(|| {
532
            let a = vec![0.0, 1.0];
533
            let mut result = vec![0.0; 2];
534
            unsafe { Avx512Backend::exp(&a, &mut result); }
535
            assert!((result[0] - 1.0).abs() < 1e-4);
536
            assert!((result[1] - std::f32::consts::E).abs() < 1e-3);
537
        });
538
    }
539
540
    #[test]
541
    fn test_avx512_sigmoid() {
542
        avx512_test(|| {
543
            let a = vec![0.0];
544
            let mut result = vec![0.0; 1];
545
            unsafe { Avx512Backend::sigmoid(&a, &mut result); }
546
            assert!((result[0] - 0.5).abs() < 1e-5);
547
        });
548
    }
549
550
    #[test]
551
    fn test_avx512_gelu() {
552
        avx512_test(|| {
553
            let a = vec![0.0, 1.0];
554
            let mut result = vec![0.0; 2];
555
            unsafe { Avx512Backend::gelu(&a, &mut result); }
556
            assert!((result[0]).abs() < 1e-5);
557
            assert!((result[1] - 0.841_192).abs() < 1e-3);
558
        });
559
    }
560
561
    #[test]
562
    fn test_avx512_swish() {
563
        avx512_test(|| {
564
            let a = vec![0.0, 1.0];
565
            let mut result = vec![0.0; 2];
566
            unsafe { Avx512Backend::swish(&a, &mut result); }
567
            assert!((result[0]).abs() < 1e-5);
568
            assert!((result[1] - 0.731_059).abs() < 1e-3);
569
        });
570
    }
571
572
    #[test]
573
    fn test_avx512_tanh() {
574
        avx512_test(|| {
575
            let a = vec![0.0, 1.0];
576
            let mut result = vec![0.0; 2];
577
            unsafe { Avx512Backend::tanh(&a, &mut result); }
578
            assert!((result[0]).abs() < 1e-5);
579
            assert!((result[1] - 0.761_594_2).abs() < 1e-3);
580
        });
581
    }
582
583
    #[test]
584
    fn test_avx512_sqrt() {
585
        avx512_test(|| {
586
            let a = vec![4.0, 9.0, 16.0];
587
            let mut result = vec![0.0; 3];
588
            unsafe { Avx512Backend::sqrt(&a, &mut result); }
589
            assert!((result[0] - 2.0).abs() < 1e-5);
590
            assert!((result[1] - 3.0).abs() < 1e-5);
591
            assert!((result[2] - 4.0).abs() < 1e-5);
592
        });
593
    }
594
595
    #[test]
596
    fn test_avx512_recip() {
597
        avx512_test(|| {
598
            let a = vec![2.0, 4.0, 5.0];
599
            let mut result = vec![0.0; 3];
600
            unsafe { Avx512Backend::recip(&a, &mut result); }
601
            assert!((result[0] - 0.5).abs() < 1e-5);
602
            assert!((result[1] - 0.25).abs() < 1e-5);
603
            assert!((result[2] - 0.2).abs() < 1e-5);
604
        });
605
    }
606
607
    #[test]
608
    fn test_avx512_transcendental() {
609
        avx512_test(|| {
610
            let a = vec![1.0, std::f32::consts::E, 10.0];
611
            let mut ln_result = vec![0.0; 3];
612
            let mut log2_result = vec![0.0; 3];
613
            let mut log10_result = vec![0.0; 3];
614
            unsafe {
615
                Avx512Backend::ln(&a, &mut ln_result);
616
                Avx512Backend::log2(&a, &mut log2_result);
617
                Avx512Backend::log10(&a, &mut log10_result);
618
            }
619
            assert!((ln_result[0]).abs() < 1e-5);
620
            assert!((ln_result[1] - 1.0).abs() < 1e-4);
621
            assert!((log10_result[2] - 1.0).abs() < 1e-5);
622
        });
623
    }
624
625
    #[test]
626
    fn test_avx512_trig() {
627
        avx512_test(|| {
628
            let a = vec![0.0, std::f32::consts::FRAC_PI_2];
629
            let mut sin_result = vec![0.0; 2];
630
            let mut cos_result = vec![0.0; 2];
631
            let mut tan_result = vec![0.0; 2];
632
            unsafe {
633
                Avx512Backend::sin(&a, &mut sin_result);
634
                Avx512Backend::cos(&a, &mut cos_result);
635
                Avx512Backend::tan(&a, &mut tan_result);
636
            }
637
            assert!((sin_result[0]).abs() < 1e-5);
638
            assert!((sin_result[1] - 1.0).abs() < 1e-5);
639
            assert!((cos_result[0] - 1.0).abs() < 1e-5);
640
        });
641
    }
642
643
    #[test]
644
    fn test_avx512_rounding() {
645
        avx512_test(|| {
646
            let a = vec![1.3, 1.5, 1.7, -1.3, -1.5, -1.7];
647
            let mut floor_result = vec![0.0; 6];
648
            let mut ceil_result = vec![0.0; 6];
649
            let mut round_result = vec![0.0; 6];
650
            unsafe {
651
                Avx512Backend::floor(&a, &mut floor_result);
652
                Avx512Backend::ceil(&a, &mut ceil_result);
653
                Avx512Backend::round(&a, &mut round_result);
654
            }
655
            assert_eq!(floor_result, vec![1.0, 1.0, 1.0, -2.0, -2.0, -2.0]);
656
            assert_eq!(ceil_result, vec![2.0, 2.0, 2.0, -1.0, -1.0, -1.0]);
657
        });
658
    }
659
660
    #[test]
661
    fn test_avx512_backend_equivalence() {
662
        avx512_test(|| {
663
            let a: Vec<f32> = (0..100).map(|i| (i as f32) * 0.1).collect();
664
            let b: Vec<f32> = (0..100).map(|i| (100 - i) as f32 * 0.1).collect();
665
            let mut avx512_add = vec![0.0; 100];
666
            let mut scalar_add = vec![0.0; 100];
667
            unsafe {
668
                Avx512Backend::add(&a, &b, &mut avx512_add);
669
                ScalarBackend::add(&a, &b, &mut scalar_add);
670
            }
671
            for i in 0..100 {
672
                assert!((avx512_add[i] - scalar_add[i]).abs() < 1e-5, "add mismatch at {}", i);
673
            }
674
        });
675
    }
676
}