1#![cfg_attr(test, allow(unused_variables, unused_mut))]
3use anyhow::{anyhow, Result};
8#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
9use std::arch::x86_64::*;
10
11#[derive(Debug, Clone)]
13pub struct SIMDConfig {
14 pub enable_avx2: bool,
16 pub enable_avx512: bool,
18 pub enable_neon: bool,
20 pub min_vector_size: usize,
22 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
38pub struct SIMDOptimizer {
40 config: SIMDConfig,
41}
42
43impl SIMDOptimizer {
44 pub fn new(config: SIMDConfig) -> Self {
46 Self { config }
47 }
48
49 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 #[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 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 for i in 0..chunks {
124 let idx = i * 8;
125
126 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 let m_new = _mm256_fmadd_ps(beta1_vec, m, _mm256_mul_ps(one_minus_beta1, g));
134
135 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 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 _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 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 #[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 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 let p_decayed = _mm256_mul_ps(p, wd_vec);
223
224 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 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 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 #[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 let g_wd = _mm256_fmadd_ps(wd_vec, p, g);
298
299 let m_new = _mm256_fmadd_ps(momentum_vec, m, g_wd);
301
302 let update = if nesterov {
304 _mm256_fmadd_ps(momentum_vec, m_new, g_wd)
306 } else {
307 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 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 #[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 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 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 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 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 for i in (chunks * 8)..len {
387 gradients[i] *= scale;
388 }
389 }
390
391 Ok(global_norm)
392 }
393
394 #[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 for i in (chunks * 8)..len {
422 a[i] += scale * b[i];
423 }
424
425 Ok(())
426 }
427
428 #[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 let result_array: [f32; 8] = std::mem::transmute(result_vec);
456 let mut result = result_array.iter().sum::<f32>();
457
458 for i in (chunks * 8)..len {
460 result += a[i] * b[i];
461 }
462
463 Ok(result)
464 }
465
466 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 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 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 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 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 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#[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 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 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 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}