Skip to main content

trustformers_optim/
simd_optimizations.rs

1//! SIMD Optimizations for Optimizers
2#![cfg_attr(test, allow(unused_variables, unused_mut))]
3//!
4//! This module provides SIMD-optimized implementations of optimizer operations
5//! for improved performance on x86_64, ARM, and other architectures.
6
7use anyhow::{anyhow, Result};
8#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
9use std::arch::x86_64::*;
10
11/// SIMD-optimized operations configuration
12#[derive(Debug, Clone)]
13pub struct SIMDConfig {
14    /// Enable AVX2 operations (x86_64)
15    pub enable_avx2: bool,
16    /// Enable AVX-512 operations (x86_64)
17    pub enable_avx512: bool,
18    /// Enable NEON operations (ARM)
19    pub enable_neon: bool,
20    /// Minimum vector size for SIMD operations
21    pub min_vector_size: usize,
22    /// Enable unrolled loops
23    pub enable_unrolling: bool,
24}
25
26impl Default for SIMDConfig {
27    fn default() -> Self {
28        Self {
29            enable_avx2: true,
30            enable_avx512: true,
31            enable_neon: true,
32            min_vector_size: 8,
33            enable_unrolling: true,
34        }
35    }
36}
37
38/// SIMD operations for optimizer kernels
39pub struct SIMDOptimizer {
40    config: SIMDConfig,
41}
42
43impl SIMDOptimizer {
44    /// Create a new SIMD optimizer with configuration
45    pub fn new(config: SIMDConfig) -> Self {
46        Self { config }
47    }
48
49    /// Detect available SIMD instruction sets
50    pub fn detect_capabilities() -> SIMDConfig {
51        SIMDConfig {
52            enable_avx2: {
53                #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
54                {
55                    is_x86_feature_detected!("avx2")
56                }
57                #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
58                {
59                    false
60                }
61            },
62            enable_avx512: {
63                #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
64                {
65                    is_x86_feature_detected!("avx512f")
66                }
67                #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
68                {
69                    false
70                }
71            },
72            enable_neon: cfg!(target_arch = "aarch64"),
73            min_vector_size: 8,
74            enable_unrolling: true,
75        }
76    }
77
78    /// SIMD-optimized Adam update with AVX2
79    ///
80    /// # Safety
81    ///
82    /// Caller must ensure the CPU supports AVX2 instructions. This is guaranteed
83    /// when the function is invoked via the `#[target_feature(enable = "avx2")]`
84    /// dispatch path only on hardware where the feature is available.
85    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
86    #[target_feature(enable = "avx2")]
87    pub unsafe fn adam_update_avx2(
88        &self,
89        params: &mut [f32],
90        gradients: &[f32],
91        momentum: &mut [f32],
92        velocity: &mut [f32],
93        lr: f32,
94        beta1: f32,
95        beta2: f32,
96        eps: f32,
97        step: i32,
98    ) -> Result<()> {
99        if params.len() != gradients.len()
100            || params.len() != momentum.len()
101            || params.len() != velocity.len()
102        {
103            return Err(anyhow!("All arrays must have the same length"));
104        }
105
106        let bias_correction1 = 1.0 - beta1.powi(step);
107        let bias_correction2 = 1.0 - beta2.powi(step);
108        let corrected_lr = lr * (bias_correction2.sqrt() / bias_correction1);
109
110        // SIMD constants
111        let beta1_vec = _mm256_set1_ps(beta1);
112        let beta2_vec = _mm256_set1_ps(beta2);
113        let one_minus_beta1 = _mm256_set1_ps(1.0 - beta1);
114        let one_minus_beta2 = _mm256_set1_ps(1.0 - beta2);
115        let eps_vec = _mm256_set1_ps(eps);
116        let lr_vec = _mm256_set1_ps(corrected_lr);
117
118        let len = params.len();
119        let chunks = len / 8;
120        let _remainder = len % 8;
121
122        // Process 8 elements at a time with AVX2
123        for i in 0..chunks {
124            let idx = i * 8;
125
126            // Load values
127            let p = _mm256_loadu_ps(params.as_ptr().add(idx));
128            let g = _mm256_loadu_ps(gradients.as_ptr().add(idx));
129            let m = _mm256_loadu_ps(momentum.as_ptr().add(idx));
130            let v = _mm256_loadu_ps(velocity.as_ptr().add(idx));
131
132            // Update momentum: m = β₁ * m + (1 - β₁) * g
133            let m_new = _mm256_fmadd_ps(beta1_vec, m, _mm256_mul_ps(one_minus_beta1, g));
134
135            // Update velocity: v = β₂ * v + (1 - β₂) * g²
136            let g_sq = _mm256_mul_ps(g, g);
137            let v_new = _mm256_fmadd_ps(beta2_vec, v, _mm256_mul_ps(one_minus_beta2, g_sq));
138
139            // Update parameters: p = p - α * m / (√v + ε)
140            let v_sqrt = _mm256_sqrt_ps(v_new);
141            let v_sqrt_eps = _mm256_add_ps(v_sqrt, eps_vec);
142            let update = _mm256_div_ps(m_new, v_sqrt_eps);
143            let p_new = _mm256_fnmadd_ps(lr_vec, update, p);
144
145            // Store results
146            _mm256_storeu_ps(params.as_mut_ptr().add(idx), p_new);
147            _mm256_storeu_ps(momentum.as_mut_ptr().add(idx), m_new);
148            _mm256_storeu_ps(velocity.as_mut_ptr().add(idx), v_new);
149        }
150
151        // Handle remaining elements
152        for i in (chunks * 8)..len {
153            let g = gradients[i];
154            let m = momentum[i];
155            let v = velocity[i];
156
157            let m_new = beta1 * m + (1.0 - beta1) * g;
158            let v_new = beta2 * v + (1.0 - beta2) * g * g;
159
160            momentum[i] = m_new;
161            velocity[i] = v_new;
162            params[i] -= corrected_lr * m_new / (v_new.sqrt() + eps);
163        }
164
165        Ok(())
166    }
167
168    /// SIMD-optimized AdamW update with AVX2 (decoupled weight decay)
169    ///
170    /// # Safety
171    ///
172    /// Caller must ensure the CPU supports AVX2 instructions. This is guaranteed
173    /// when the function is invoked via the `#[target_feature(enable = "avx2")]`
174    /// dispatch path only on hardware where the feature is available.
175    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
176    #[target_feature(enable = "avx2")]
177    pub unsafe fn adamw_update_avx2(
178        &self,
179        params: &mut [f32],
180        gradients: &[f32],
181        momentum: &mut [f32],
182        velocity: &mut [f32],
183        lr: f32,
184        beta1: f32,
185        beta2: f32,
186        eps: f32,
187        weight_decay: f32,
188        step: i32,
189    ) -> Result<()> {
190        if params.len() != gradients.len()
191            || params.len() != momentum.len()
192            || params.len() != velocity.len()
193        {
194            return Err(anyhow!("All arrays must have the same length"));
195        }
196
197        let bias_correction1 = 1.0 - beta1.powi(step);
198        let bias_correction2 = 1.0 - beta2.powi(step);
199        let corrected_lr = lr * (bias_correction2.sqrt() / bias_correction1);
200
201        // SIMD constants
202        let beta1_vec = _mm256_set1_ps(beta1);
203        let beta2_vec = _mm256_set1_ps(beta2);
204        let one_minus_beta1 = _mm256_set1_ps(1.0 - beta1);
205        let one_minus_beta2 = _mm256_set1_ps(1.0 - beta2);
206        let eps_vec = _mm256_set1_ps(eps);
207        let lr_vec = _mm256_set1_ps(corrected_lr);
208        let wd_vec = _mm256_set1_ps(1.0 - lr * weight_decay);
209
210        let len = params.len();
211        let chunks = len / 8;
212
213        for i in 0..chunks {
214            let idx = i * 8;
215
216            let p = _mm256_loadu_ps(params.as_ptr().add(idx));
217            let g = _mm256_loadu_ps(gradients.as_ptr().add(idx));
218            let m = _mm256_loadu_ps(momentum.as_ptr().add(idx));
219            let v = _mm256_loadu_ps(velocity.as_ptr().add(idx));
220
221            // Apply weight decay first: p = p * (1 - lr * wd)
222            let p_decayed = _mm256_mul_ps(p, wd_vec);
223
224            // Update momentum and velocity
225            let m_new = _mm256_fmadd_ps(beta1_vec, m, _mm256_mul_ps(one_minus_beta1, g));
226            let g_sq = _mm256_mul_ps(g, g);
227            let v_new = _mm256_fmadd_ps(beta2_vec, v, _mm256_mul_ps(one_minus_beta2, g_sq));
228
229            // Update parameters
230            let v_sqrt = _mm256_sqrt_ps(v_new);
231            let v_sqrt_eps = _mm256_add_ps(v_sqrt, eps_vec);
232            let update = _mm256_div_ps(m_new, v_sqrt_eps);
233            let p_new = _mm256_fnmadd_ps(lr_vec, update, p_decayed);
234
235            _mm256_storeu_ps(params.as_mut_ptr().add(idx), p_new);
236            _mm256_storeu_ps(momentum.as_mut_ptr().add(idx), m_new);
237            _mm256_storeu_ps(velocity.as_mut_ptr().add(idx), v_new);
238        }
239
240        // Handle remaining elements
241        for i in (chunks * 8)..len {
242            let p = params[i];
243            let g = gradients[i];
244            let m = momentum[i];
245            let v = velocity[i];
246
247            let p_decayed = p * (1.0 - lr * weight_decay);
248            let m_new = beta1 * m + (1.0 - beta1) * g;
249            let v_new = beta2 * v + (1.0 - beta2) * g * g;
250
251            momentum[i] = m_new;
252            velocity[i] = v_new;
253            params[i] = p_decayed - corrected_lr * m_new / (v_new.sqrt() + eps);
254        }
255
256        Ok(())
257    }
258
259    /// SIMD-optimized SGD with momentum update
260    ///
261    /// # Safety
262    ///
263    /// Caller must ensure the CPU supports AVX2 instructions. This is guaranteed
264    /// when the function is invoked via the `#[target_feature(enable = "avx2")]`
265    /// dispatch path only on hardware where the feature is available.
266    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
267    #[target_feature(enable = "avx2")]
268    pub unsafe fn sgd_momentum_update_avx2(
269        &self,
270        params: &mut [f32],
271        gradients: &[f32],
272        momentum: &mut [f32],
273        lr: f32,
274        momentum_factor: f32,
275        weight_decay: f32,
276        nesterov: bool,
277    ) -> Result<()> {
278        if params.len() != gradients.len() || params.len() != momentum.len() {
279            return Err(anyhow!("All arrays must have the same length"));
280        }
281
282        let lr_vec = _mm256_set1_ps(lr);
283        let momentum_vec = _mm256_set1_ps(momentum_factor);
284        let wd_vec = _mm256_set1_ps(weight_decay);
285
286        let len = params.len();
287        let chunks = len / 8;
288
289        for i in 0..chunks {
290            let idx = i * 8;
291
292            let p = _mm256_loadu_ps(params.as_ptr().add(idx));
293            let g = _mm256_loadu_ps(gradients.as_ptr().add(idx));
294            let m = _mm256_loadu_ps(momentum.as_ptr().add(idx));
295
296            // Apply weight decay to gradient: g = g + wd * p
297            let g_wd = _mm256_fmadd_ps(wd_vec, p, g);
298
299            // Update momentum: m = momentum * m + g
300            let m_new = _mm256_fmadd_ps(momentum_vec, m, g_wd);
301
302            // Update parameters
303            let update = if nesterov {
304                // Nesterov: p = p - lr * (momentum * m + g)
305                _mm256_fmadd_ps(momentum_vec, m_new, g_wd)
306            } else {
307                // Standard: p = p - lr * m
308                m_new
309            };
310
311            let p_new = _mm256_fnmadd_ps(lr_vec, update, p);
312
313            _mm256_storeu_ps(params.as_mut_ptr().add(idx), p_new);
314            _mm256_storeu_ps(momentum.as_mut_ptr().add(idx), m_new);
315        }
316
317        // Handle remaining elements
318        for i in (chunks * 8)..len {
319            let p = params[i];
320            let g = gradients[i] + weight_decay * p;
321            let m = momentum[i];
322
323            let m_new = momentum_factor * m + g;
324            momentum[i] = m_new;
325
326            if nesterov {
327                params[i] = p - lr * (momentum_factor * m_new + g);
328            } else {
329                params[i] = p - lr * m_new;
330            }
331        }
332
333        Ok(())
334    }
335
336    /// SIMD-optimized gradient clipping
337    ///
338    /// # Safety
339    ///
340    /// Caller must ensure the CPU supports AVX2 instructions. This is guaranteed
341    /// when the function is invoked via the `#[target_feature(enable = "avx2")]`
342    /// dispatch path only on hardware where the feature is available.
343    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
344    #[target_feature(enable = "avx2")]
345    pub unsafe fn clip_gradients_avx2(&self, gradients: &mut [f32], max_norm: f32) -> Result<f32> {
346        let len = gradients.len();
347        let chunks = len / 8;
348
349        // Compute global norm
350        let mut norm_sq_vec = _mm256_setzero_ps();
351
352        for i in 0..chunks {
353            let idx = i * 8;
354            let g = _mm256_loadu_ps(gradients.as_ptr().add(idx));
355            let g_sq = _mm256_mul_ps(g, g);
356            norm_sq_vec = _mm256_add_ps(norm_sq_vec, g_sq);
357        }
358
359        // Horizontal sum of norm_sq_vec
360        let mut norm_sq = 0.0f32;
361        let norm_sq_array: [f32; 8] = std::mem::transmute(norm_sq_vec);
362        for &val in &norm_sq_array {
363            norm_sq += val;
364        }
365
366        // Add remaining elements
367        for i in (chunks * 8)..len {
368            norm_sq += gradients[i] * gradients[i];
369        }
370
371        let global_norm = norm_sq.sqrt();
372
373        if global_norm > max_norm {
374            let scale = max_norm / global_norm;
375            let scale_vec = _mm256_set1_ps(scale);
376
377            // Scale gradients
378            for i in 0..chunks {
379                let idx = i * 8;
380                let g = _mm256_loadu_ps(gradients.as_ptr().add(idx));
381                let g_scaled = _mm256_mul_ps(g, scale_vec);
382                _mm256_storeu_ps(gradients.as_mut_ptr().add(idx), g_scaled);
383            }
384
385            // Scale remaining elements
386            for i in (chunks * 8)..len {
387                gradients[i] *= scale;
388            }
389        }
390
391        Ok(global_norm)
392    }
393
394    /// SIMD-optimized vector addition (for gradient accumulation)
395    ///
396    /// # Safety
397    ///
398    /// Caller must ensure the CPU supports AVX2 instructions. This is guaranteed
399    /// when the function is invoked via the `#[target_feature(enable = "avx2")]`
400    /// dispatch path only on hardware where the feature is available.
401    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
402    #[target_feature(enable = "avx2")]
403    pub unsafe fn vector_add_avx2(&self, a: &mut [f32], b: &[f32], scale: f32) -> Result<()> {
404        if a.len() != b.len() {
405            return Err(anyhow!("Vectors must have the same length"));
406        }
407
408        let scale_vec = _mm256_set1_ps(scale);
409        let len = a.len();
410        let chunks = len / 8;
411
412        for i in 0..chunks {
413            let idx = i * 8;
414            let a_vec = _mm256_loadu_ps(a.as_ptr().add(idx));
415            let b_vec = _mm256_loadu_ps(b.as_ptr().add(idx));
416            let result = _mm256_fmadd_ps(b_vec, scale_vec, a_vec);
417            _mm256_storeu_ps(a.as_mut_ptr().add(idx), result);
418        }
419
420        // Handle remaining elements
421        for i in (chunks * 8)..len {
422            a[i] += scale * b[i];
423        }
424
425        Ok(())
426    }
427
428    /// SIMD-optimized dot product
429    ///
430    /// # Safety
431    ///
432    /// Caller must ensure the CPU supports AVX2 instructions. This is guaranteed
433    /// when the function is invoked via the `#[target_feature(enable = "avx2")]`
434    /// dispatch path only on hardware where the feature is available.
435    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
436    #[target_feature(enable = "avx2")]
437    pub unsafe fn dot_product_avx2(&self, a: &[f32], b: &[f32]) -> Result<f32> {
438        if a.len() != b.len() {
439            return Err(anyhow!("Vectors must have the same length"));
440        }
441
442        let len = a.len();
443        let chunks = len / 8;
444        let mut result_vec = _mm256_setzero_ps();
445
446        for i in 0..chunks {
447            let idx = i * 8;
448            let a_vec = _mm256_loadu_ps(a.as_ptr().add(idx));
449            let b_vec = _mm256_loadu_ps(b.as_ptr().add(idx));
450            let prod = _mm256_mul_ps(a_vec, b_vec);
451            result_vec = _mm256_add_ps(result_vec, prod);
452        }
453
454        // Horizontal sum
455        let result_array: [f32; 8] = std::mem::transmute(result_vec);
456        let mut result = result_array.iter().sum::<f32>();
457
458        // Add remaining elements
459        for i in (chunks * 8)..len {
460            result += a[i] * b[i];
461        }
462
463        Ok(result)
464    }
465
466    /// Fallback implementations for non-x86 architectures
467    pub fn adam_update_fallback(
468        &self,
469        params: &mut [f32],
470        gradients: &[f32],
471        momentum: &mut [f32],
472        velocity: &mut [f32],
473        lr: f32,
474        beta1: f32,
475        beta2: f32,
476        eps: f32,
477        step: i32,
478    ) -> Result<()> {
479        if params.len() != gradients.len()
480            || params.len() != momentum.len()
481            || params.len() != velocity.len()
482        {
483            return Err(anyhow!("All arrays must have the same length"));
484        }
485
486        let bias_correction1 = 1.0 - beta1.powi(step);
487        let bias_correction2 = 1.0 - beta2.powi(step);
488        let corrected_lr = lr * (bias_correction2.sqrt() / bias_correction1);
489
490        for i in 0..params.len() {
491            let g = gradients[i];
492            let m = momentum[i];
493            let v = velocity[i];
494
495            let m_new = beta1 * m + (1.0 - beta1) * g;
496            let v_new = beta2 * v + (1.0 - beta2) * g * g;
497
498            momentum[i] = m_new;
499            velocity[i] = v_new;
500            params[i] -= corrected_lr * m_new / (v_new.sqrt() + eps);
501        }
502
503        Ok(())
504    }
505
506    /// Auto-dispatch to best available implementation
507    pub fn adam_update(
508        &self,
509        params: &mut [f32],
510        gradients: &[f32],
511        momentum: &mut [f32],
512        velocity: &mut [f32],
513        lr: f32,
514        beta1: f32,
515        beta2: f32,
516        eps: f32,
517        step: i32,
518    ) -> Result<()> {
519        if params.len() < self.config.min_vector_size {
520            return self.adam_update_fallback(
521                params, gradients, momentum, velocity, lr, beta1, beta2, eps, step,
522            );
523        }
524
525        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
526        {
527            if self.config.enable_avx2 && is_x86_feature_detected!("avx2") {
528                return unsafe {
529                    self.adam_update_avx2(
530                        params, gradients, momentum, velocity, lr, beta1, beta2, eps, step,
531                    )
532                };
533            }
534        }
535
536        self.adam_update_fallback(
537            params, gradients, momentum, velocity, lr, beta1, beta2, eps, step,
538        )
539    }
540
541    /// Auto-dispatch AdamW
542    pub fn adamw_update(
543        &self,
544        params: &mut [f32],
545        gradients: &[f32],
546        momentum: &mut [f32],
547        velocity: &mut [f32],
548        lr: f32,
549        beta1: f32,
550        beta2: f32,
551        eps: f32,
552        weight_decay: f32,
553        step: i32,
554    ) -> Result<()> {
555        if params.len() < self.config.min_vector_size {
556            // Fallback implementation
557            let bias_correction1 = 1.0 - beta1.powi(step);
558            let bias_correction2 = 1.0 - beta2.powi(step);
559            let corrected_lr = lr * (bias_correction2.sqrt() / bias_correction1);
560
561            for i in 0..params.len() {
562                let p = params[i];
563                let g = gradients[i];
564                let m = momentum[i];
565                let v = velocity[i];
566
567                let p_decayed = p * (1.0 - lr * weight_decay);
568                let m_new = beta1 * m + (1.0 - beta1) * g;
569                let v_new = beta2 * v + (1.0 - beta2) * g * g;
570
571                momentum[i] = m_new;
572                velocity[i] = v_new;
573                params[i] = p_decayed - corrected_lr * m_new / (v_new.sqrt() + eps);
574            }
575            return Ok(());
576        }
577
578        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
579        {
580            if self.config.enable_avx2 && is_x86_feature_detected!("avx2") {
581                return unsafe {
582                    self.adamw_update_avx2(
583                        params,
584                        gradients,
585                        momentum,
586                        velocity,
587                        lr,
588                        beta1,
589                        beta2,
590                        eps,
591                        weight_decay,
592                        step,
593                    )
594                };
595            }
596        }
597
598        // Fallback
599        self.adamw_update(
600            params,
601            gradients,
602            momentum,
603            velocity,
604            lr,
605            beta1,
606            beta2,
607            eps,
608            weight_decay,
609            step,
610        )
611    }
612
613    /// Get performance statistics
614    pub fn get_performance_info(&self) -> SIMDPerformanceInfo {
615        SIMDPerformanceInfo {
616            avx2_available: {
617                #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
618                {
619                    is_x86_feature_detected!("avx2")
620                }
621                #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
622                {
623                    false
624                }
625            },
626            avx512_available: {
627                #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
628                {
629                    is_x86_feature_detected!("avx512f")
630                }
631                #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
632                {
633                    false
634                }
635            },
636            neon_available: cfg!(target_arch = "aarch64"),
637            vector_width: {
638                #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
639                {
640                    if is_x86_feature_detected!("avx2") {
641                        8
642                    } else {
643                        1
644                    }
645                }
646                #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
647                {
648                    1
649                }
650            },
651            recommended_min_size: self.config.min_vector_size,
652        }
653    }
654}
655
656impl Default for SIMDOptimizer {
657    fn default() -> Self {
658        Self::new(SIMDOptimizer::detect_capabilities())
659    }
660}
661
662/// SIMD performance information
663#[derive(Debug, Clone)]
664pub struct SIMDPerformanceInfo {
665    pub avx2_available: bool,
666    pub avx512_available: bool,
667    pub neon_available: bool,
668    pub vector_width: usize,
669    pub recommended_min_size: usize,
670}
671
672#[cfg(test)]
673mod tests {
674    use super::*;
675
676    #[test]
677    fn test_simd_config_detection() {
678        let config = SIMDOptimizer::detect_capabilities();
679        // Test will pass regardless of actual hardware capabilities
680        assert!(config.min_vector_size > 0);
681    }
682
683    #[test]
684    fn test_adam_update_fallback() {
685        let optimizer = SIMDOptimizer::default();
686        let mut params = vec![1.0, 2.0, 3.0, 4.0];
687        let gradients = vec![0.1, 0.2, 0.3, 0.4];
688        let mut momentum = vec![0.0; 4];
689        let mut velocity = vec![0.0; 4];
690
691        optimizer
692            .adam_update_fallback(
693                &mut params,
694                &gradients,
695                &mut momentum,
696                &mut velocity,
697                0.001,
698                0.9,
699                0.999,
700                1e-8,
701                1,
702            )
703            .expect("Operation failed in test");
704
705        // Check that parameters were updated
706        assert!(params[0] < 1.0);
707        assert!(momentum[0] > 0.0);
708        assert!(velocity[0] > 0.0);
709    }
710
711    #[test]
712    fn test_auto_dispatch_adam() {
713        let optimizer = SIMDOptimizer::default();
714        let mut params = vec![1.0; 16];
715        let gradients = vec![0.1; 16];
716        let mut momentum = vec![0.0; 16];
717        let mut velocity = vec![0.0; 16];
718
719        optimizer
720            .adam_update(
721                &mut params,
722                &gradients,
723                &mut momentum,
724                &mut velocity,
725                0.001,
726                0.9,
727                0.999,
728                1e-8,
729                1,
730            )
731            .expect("Operation failed in test");
732
733        // Verify update occurred
734        assert!(params.iter().all(|&p| p < 1.0));
735        assert!(momentum.iter().all(|&m| m > 0.0));
736    }
737
738    #[test]
739    fn test_performance_info() {
740        let optimizer = SIMDOptimizer::default();
741        let info = optimizer.get_performance_info();
742
743        assert!(info.vector_width > 0);
744        assert!(info.recommended_min_size > 0);
745    }
746
747    #[test]
748    fn test_vector_operations() {
749        let optimizer = SIMDOptimizer::default();
750        let mut a = [1.0, 2.0, 3.0, 4.0];
751        let b = [0.5, 0.5, 0.5, 0.5];
752
753        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
754        {
755            if is_x86_feature_detected!("avx2") {
756                unsafe {
757                    optimizer.vector_add_avx2(&mut a, &b, 2.0).expect("Operation failed in test");
758                }
759                assert_eq!(a, [2.0f32, 3.0, 4.0, 5.0]);
760            }
761        }
762    }
763
764    #[test]
765    fn test_dot_product() {
766        let optimizer = SIMDOptimizer::default();
767        let a = [1.0, 2.0, 3.0, 4.0];
768        let b = [1.0, 1.0, 1.0, 1.0];
769
770        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
771        {
772            if is_x86_feature_detected!("avx2") {
773                unsafe {
774                    let result =
775                        optimizer.dot_product_avx2(&a, &b).expect("Operation failed in test");
776                    assert_eq!(result, 10.0);
777                }
778            }
779        }
780    }
781}