1use std::cmp::min;
17use std::collections::HashMap;
18use std::time::Instant;
19use torsh_core::sync::RwLockExt;
20
21use scirs2_core::parallel_ops::*;
23use torsh_core::{
24 dtype::FloatElement,
25 error::{Result, TorshError},
26};
27
28#[derive(Debug, Clone)]
33pub struct AlgorithmConfig {
34 pub enable_adaptive_selection: bool,
36 pub min_size_for_advanced: usize,
38 pub l1_cache_size: usize,
40 pub l2_cache_size: usize,
41 pub l3_cache_size: usize,
42 pub enable_operation_fusion: bool,
44 pub max_fusion_chain: usize,
46 pub enable_numerical_stability: bool,
48 pub scheduling_strategy: SchedulingStrategy,
50}
51
52impl Default for AlgorithmConfig {
53 fn default() -> Self {
54 Self {
55 enable_adaptive_selection: true,
56 min_size_for_advanced: 64,
57 l1_cache_size: 32 * 1024, l2_cache_size: 256 * 1024, l3_cache_size: 8 * 1024 * 1024, enable_operation_fusion: true,
61 max_fusion_chain: 8,
62 enable_numerical_stability: true,
63 scheduling_strategy: SchedulingStrategy::WorkStealing,
64 }
65 }
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum SchedulingStrategy {
71 Static,
73 WorkStealing,
75 Adaptive,
77 NumaAware,
79}
80
81pub struct AlgorithmicOptimizer {
83 config: AlgorithmConfig,
84 performance_history: std::sync::RwLock<HashMap<OperationSignature, PerformanceMetrics>>,
86}
87
88impl AlgorithmicOptimizer {
89 pub fn new() -> Self {
91 Self::with_config(AlgorithmConfig::default())
92 }
93
94 pub fn with_config(config: AlgorithmConfig) -> Self {
96 Self {
97 config,
98 performance_history: std::sync::RwLock::new(HashMap::new()),
99 }
100 }
101
102 pub fn optimized_matmul<T>(
104 &self,
105 a: &[T],
106 b: &[T],
107 c: &mut [T],
108 m: usize, k: usize, n: usize, ) -> Result<()>
112 where
113 T: FloatElement + Send + Sync + std::ops::AddAssign,
114 {
115 #[cfg(feature = "profiling")]
116 {
117 }
119 let signature = OperationSignature::MatMul { m, k, n };
120
121 let algorithm = self.select_matmul_algorithm(&signature);
123
124 let start_time = Instant::now();
125
126 match algorithm {
127 MatMulAlgorithm::Naive => self.naive_matmul(a, b, c, m, k, n)?,
128 MatMulAlgorithm::Blocked => self.blocked_matmul(a, b, c, m, k, n)?,
129 MatMulAlgorithm::Strassen => self.strassen_matmul(a, b, c, m, k, n)?,
130 MatMulAlgorithm::CacheOblivious => self.cache_oblivious_matmul(a, b, c, m, k, n)?,
131 MatMulAlgorithm::Parallel => self.parallel_matmul(a, b, c, m, k, n)?,
132 }
133
134 let duration = start_time.elapsed();
136 self.record_performance(signature, algorithm, duration);
137
138 Ok(())
139 }
140
141 fn select_matmul_algorithm(&self, signature: &OperationSignature) -> MatMulAlgorithm {
143 if !self.config.enable_adaptive_selection {
144 return MatMulAlgorithm::Blocked; }
146
147 if let Some(metrics) = self.performance_history.read_or_recover().get(signature) {
149 return metrics
150 .best_algorithm
151 .clone()
152 .unwrap_or(MatMulAlgorithm::Blocked);
153 }
154
155 match signature {
157 OperationSignature::MatMul { m, k, n } => {
158 let total_size = m * k * n;
159
160 if total_size < 1000 {
161 MatMulAlgorithm::Naive
162 } else if total_size < 10000 {
163 MatMulAlgorithm::Blocked
164 } else if *m >= 1024 && *k >= 1024 && *n >= 1024 {
165 MatMulAlgorithm::Strassen
166 } else if total_size > 100000 {
167 MatMulAlgorithm::Parallel
168 } else {
169 MatMulAlgorithm::CacheOblivious
170 }
171 }
172 }
173 }
174
175 fn naive_matmul<T>(
177 &self,
178 a: &[T],
179 b: &[T],
180 c: &mut [T],
181 m: usize,
182 k: usize,
183 n: usize,
184 ) -> Result<()>
185 where
186 T: FloatElement + std::ops::AddAssign,
187 {
188 for i in 0..m {
189 for j in 0..n {
190 let mut sum = <T as torsh_core::TensorElement>::zero();
191 for l in 0..k {
192 sum += a[i * k + l] * b[l * n + j];
193 }
194 c[i * n + j] = sum;
195 }
196 }
197 Ok(())
198 }
199
200 fn blocked_matmul<T>(
202 &self,
203 a: &[T],
204 b: &[T],
205 c: &mut [T],
206 m: usize,
207 k: usize,
208 n: usize,
209 ) -> Result<()>
210 where
211 T: FloatElement + std::ops::AddAssign,
212 {
213 let block_size = self.calculate_optimal_block_size(m, k, n);
215
216 for i_block in (0..m).step_by(block_size) {
217 for j_block in (0..n).step_by(block_size) {
218 for k_block in (0..k).step_by(block_size) {
219 let i_end = min(i_block + block_size, m);
220 let j_end = min(j_block + block_size, n);
221 let k_end = min(k_block + block_size, k);
222
223 for i in i_block..i_end {
225 for j in j_block..j_end {
226 let mut sum = if k_block == 0 {
227 <T as torsh_core::TensorElement>::zero()
228 } else {
229 c[i * n + j]
230 };
231 for l in k_block..k_end {
232 sum += a[i * k + l] * b[l * n + j];
233 }
234 c[i * n + j] = sum;
235 }
236 }
237 }
238 }
239 }
240 Ok(())
241 }
242
243 fn strassen_matmul<T>(
245 &self,
246 a: &[T],
247 b: &[T],
248 c: &mut [T],
249 m: usize,
250 k: usize,
251 n: usize,
252 ) -> Result<()>
253 where
254 T: FloatElement + Send + Sync + std::ops::AddAssign,
255 {
256 if m != k || k != n || m < 128 {
258 return self.blocked_matmul(a, b, c, m, k, n);
259 }
260
261 self.strassen_recursive(a, b, c, m, 0, 0, 0, 0, 0, 0)
262 }
263
264 fn strassen_recursive<T>(
266 &self,
267 a: &[T],
268 b: &[T],
269 c: &mut [T],
270 n: usize,
271 a_row: usize,
272 a_col: usize,
273 b_row: usize,
274 b_col: usize,
275 c_row: usize,
276 c_col: usize,
277 ) -> Result<()>
278 where
279 T: FloatElement + Send + Sync + std::ops::AddAssign,
280 {
281 if n <= 64 {
282 for i in 0..n {
284 for j in 0..n {
285 let mut sum = <T as torsh_core::TensorElement>::zero();
286 for k in 0..n {
287 let a_val = a[(a_row + i) * n + (a_col + k)];
288 let b_val = b[(b_row + k) * n + (b_col + j)];
289 sum += a_val * b_val;
290 }
291 c[(c_row + i) * n + (c_col + j)] = sum;
292 }
293 }
294 return Ok(());
295 }
296
297 let half = n / 2;
298
299 let temp_size = half * half;
301 let mut m1 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
302 let mut m2 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
303 let mut m3 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
304 let mut m4 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
305 let mut m5 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
306 let mut m6 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
307 let mut m7 = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
308
309 let mut temp_a = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
311 let mut temp_b = vec![<T as torsh_core::TensorElement>::zero(); temp_size];
312
313 let add_quadrants = |temp: &mut [T],
315 quad1_row: usize,
316 quad1_col: usize,
317 quad2_row: usize,
318 quad2_col: usize,
319 source: &[T]| {
320 for i in 0..half {
321 for j in 0..half {
322 let val1 = source[(quad1_row + i) * n + (quad1_col + j)];
323 let val2 = source[(quad2_row + i) * n + (quad2_col + j)];
324 temp[i * half + j] = val1 + val2;
325 }
326 }
327 };
328
329 let sub_quadrants = |temp: &mut [T],
331 quad1_row: usize,
332 quad1_col: usize,
333 quad2_row: usize,
334 quad2_col: usize,
335 source: &[T]| {
336 for i in 0..half {
337 for j in 0..half {
338 let val1 = source[(quad1_row + i) * n + (quad1_col + j)];
339 let val2 = source[(quad2_row + i) * n + (quad2_col + j)];
340 temp[i * half + j] = val1 - val2;
341 }
342 }
343 };
344
345 add_quadrants(&mut temp_a, a_row, a_col, a_row + half, a_col + half, a);
347 add_quadrants(&mut temp_b, b_row, b_col, b_row + half, b_col + half, b);
348 self.blocked_matmul(&temp_a, &temp_b, &mut m1, half, half, half)?;
349
350 add_quadrants(
352 &mut temp_a,
353 a_row + half,
354 a_col,
355 a_row + half,
356 a_col + half,
357 a,
358 );
359 for i in 0..half {
360 for j in 0..half {
361 temp_b[i * half + j] = b[(b_row + i) * n + (b_col + j)];
362 }
363 }
364 self.blocked_matmul(&temp_a, &temp_b, &mut m2, half, half, half)?;
365
366 for i in 0..half {
368 for j in 0..half {
369 temp_a[i * half + j] = a[(a_row + i) * n + (a_col + j)];
370 }
371 }
372 sub_quadrants(
373 &mut temp_b,
374 b_row,
375 b_col + half,
376 b_row + half,
377 b_col + half,
378 b,
379 );
380 self.blocked_matmul(&temp_a, &temp_b, &mut m3, half, half, half)?;
381
382 for i in 0..half {
384 for j in 0..half {
385 temp_a[i * half + j] = a[(a_row + half + i) * n + (a_col + half + j)];
386 }
387 }
388 sub_quadrants(&mut temp_b, b_row + half, b_col, b_row, b_col, b);
389 self.blocked_matmul(&temp_a, &temp_b, &mut m4, half, half, half)?;
390
391 add_quadrants(&mut temp_a, a_row, a_col, a_row, a_col + half, a);
393 for i in 0..half {
394 for j in 0..half {
395 temp_b[i * half + j] = b[(b_row + half + i) * n + (b_col + half + j)];
396 }
397 }
398 self.blocked_matmul(&temp_a, &temp_b, &mut m5, half, half, half)?;
399
400 sub_quadrants(&mut temp_a, a_row + half, a_col, a_row, a_col, a);
402 add_quadrants(&mut temp_b, b_row, b_col, b_row, b_col + half, b);
403 self.blocked_matmul(&temp_a, &temp_b, &mut m6, half, half, half)?;
404
405 sub_quadrants(
407 &mut temp_a,
408 a_row,
409 a_col + half,
410 a_row + half,
411 a_col + half,
412 a,
413 );
414 add_quadrants(
415 &mut temp_b,
416 b_row + half,
417 b_col,
418 b_row + half,
419 b_col + half,
420 b,
421 );
422 self.blocked_matmul(&temp_a, &temp_b, &mut m7, half, half, half)?;
423
424 for i in 0..half {
427 for j in 0..half {
428 c[(c_row + i) * n + (c_col + j)] =
429 m1[i * half + j] + m4[i * half + j] - m5[i * half + j] + m7[i * half + j];
430 }
431 }
432
433 for i in 0..half {
435 for j in 0..half {
436 c[(c_row + i) * n + (c_col + half + j)] = m3[i * half + j] + m5[i * half + j];
437 }
438 }
439
440 for i in 0..half {
442 for j in 0..half {
443 c[(c_row + half + i) * n + (c_col + j)] = m2[i * half + j] + m4[i * half + j];
444 }
445 }
446
447 for i in 0..half {
449 for j in 0..half {
450 c[(c_row + half + i) * n + (c_col + half + j)] =
451 m1[i * half + j] - m2[i * half + j] + m3[i * half + j] + m6[i * half + j];
452 }
453 }
454
455 Ok(())
456 }
457
458 fn cache_oblivious_matmul<T>(
460 &self,
461 a: &[T],
462 b: &[T],
463 c: &mut [T],
464 m: usize,
465 k: usize,
466 n: usize,
467 ) -> Result<()>
468 where
469 T: FloatElement + std::ops::AddAssign,
470 {
471 self.cache_oblivious_recursive(a, b, c, m, k, n, 0, 0, 0, 0, 0, 0)
472 }
473
474 fn cache_oblivious_recursive<T>(
476 &self,
477 a: &[T],
478 b: &[T],
479 c: &mut [T],
480 m: usize,
481 k: usize,
482 n: usize,
483 a_row: usize,
484 a_col: usize,
485 b_row: usize,
486 b_col: usize,
487 c_row: usize,
488 c_col: usize,
489 ) -> Result<()>
490 where
491 T: FloatElement + std::ops::AddAssign,
492 {
493 if m <= 32 || k <= 32 || n <= 32 {
495 return self
496 .naive_matmul_region(a, b, c, m, k, n, a_row, a_col, b_row, b_col, c_row, c_col);
497 }
498
499 if m >= k && m >= n {
501 let m1 = m / 2;
502 let m2 = m - m1;
503
504 self.cache_oblivious_recursive(
506 a, b, c, m1, k, n, a_row, a_col, b_row, b_col, c_row, c_col,
507 )?;
508
509 self.cache_oblivious_recursive(
511 a,
512 b,
513 c,
514 m2,
515 k,
516 n,
517 a_row + m1,
518 a_col,
519 b_row,
520 b_col,
521 c_row + m1,
522 c_col,
523 )?;
524 } else if k >= n {
525 let k1 = k / 2;
526 let k2 = k - k1;
527
528 self.cache_oblivious_recursive(
530 a, b, c, m, k1, n, a_row, a_col, b_row, b_col, c_row, c_col,
531 )?;
532
533 self.cache_oblivious_recursive(
534 a,
535 b,
536 c,
537 m,
538 k2,
539 n,
540 a_row,
541 a_col + k1,
542 b_row + k1,
543 b_col,
544 c_row,
545 c_col,
546 )?;
547 } else {
548 let n1 = n / 2;
549 let n2 = n - n1;
550
551 self.cache_oblivious_recursive(
553 a, b, c, m, k, n1, a_row, a_col, b_row, b_col, c_row, c_col,
554 )?;
555
556 self.cache_oblivious_recursive(
558 a,
559 b,
560 c,
561 m,
562 k,
563 n2,
564 a_row,
565 a_col,
566 b_row,
567 b_col + n1,
568 c_row,
569 c_col + n1,
570 )?;
571 }
572
573 Ok(())
574 }
575
576 fn naive_matmul_region<T>(
578 &self,
579 a: &[T],
580 b: &[T],
581 c: &mut [T],
582 m: usize,
583 k: usize,
584 n: usize,
585 a_row: usize,
586 a_col: usize,
587 b_row: usize,
588 b_col: usize,
589 c_row: usize,
590 c_col: usize,
591 ) -> Result<()>
592 where
593 T: FloatElement + std::ops::AddAssign,
594 {
595 for i in 0..m {
596 for j in 0..n {
597 let mut sum = <T as torsh_core::TensorElement>::zero();
598 for l in 0..k {
599 let a_idx = (a_row + i) * k + (a_col + l);
600 let b_idx = (b_row + l) * n + (b_col + j);
601 sum += a[a_idx] * b[b_idx];
602 }
603 let c_idx = (c_row + i) * n + (c_col + j);
604 c[c_idx] += sum; }
606 }
607 Ok(())
608 }
609
610 fn parallel_matmul<T>(
612 &self,
613 a: &[T],
614 b: &[T],
615 c: &mut [T],
616 m: usize,
617 k: usize,
618 n: usize,
619 ) -> Result<()>
620 where
621 T: FloatElement + Send + Sync + std::ops::AddAssign,
622 {
623 let num_cores = get_num_threads();
624 let block_size = self.calculate_optimal_block_size(m, k, n);
625
626 let total_operations = m * k * n;
628 let min_work_per_core = 100_000; let should_parallelize = num_cores > 1 && total_operations > min_work_per_core * num_cores;
630
631 if !should_parallelize {
632 return self.blocked_matmul(a, b, c, m, k, n);
634 }
635
636 let work_items: Vec<_> = (0..m)
638 .step_by(block_size)
639 .flat_map(|i| (0..n).step_by(block_size).map(move |j| (i, j)))
640 .collect();
641
642 let results: Result<Vec<_>> = parallel_map_result(&work_items, |&(i_block, j_block)| {
644 let i_end = min(i_block + block_size, m);
645 let j_end = min(j_block + block_size, n);
646
647 let mut block_results = Vec::new();
648 for i in i_block..i_end {
649 for j in j_block..j_end {
650 let mut sum = <T as torsh_core::TensorElement>::zero();
651 for l in 0..k {
652 sum += a[i * k + l] * b[l * n + j];
653 }
654 let idx = i * n + j;
655 block_results.push((idx, sum));
656 }
657 }
658 Ok(block_results)
659 });
660
661 for block_results in results? {
663 for (idx, value) in block_results {
664 c[idx] = value;
665 }
666 }
667
668 Ok(())
669 }
670
671 fn calculate_optimal_block_size(&self, m: usize, k: usize, n: usize) -> usize {
673 let element_size = std::mem::size_of::<f32>(); let l1_elements = self.config.l1_cache_size / element_size;
679
680 let cache_optimal = (l1_elements as f64 / 3.0).sqrt() as usize;
683
684 let dim_optimal = m.min(k).min(n);
686
687 let optimal_block = cache_optimal.min(dim_optimal);
689
690 let clamped = optimal_block.clamp(16, 256);
692
693 let log2 = (clamped as f64).log2().round() as u32;
695 2usize.pow(log2).min(256)
696 }
697
698 fn record_performance(
700 &self,
701 signature: OperationSignature,
702 algorithm: MatMulAlgorithm,
703 duration: std::time::Duration,
704 ) {
705 let mut history = self.performance_history.write_or_recover();
706 let metrics = history
707 .entry(signature)
708 .or_insert_with(PerformanceMetrics::default);
709
710 metrics.update_performance(algorithm, duration);
711 }
712
713 pub fn optimized_conv2d<T>(
715 &self,
716 input: &[T],
717 kernel: &[T],
718 output: &mut [T],
719 input_h: usize,
720 input_w: usize,
721 kernel_h: usize,
722 kernel_w: usize,
723 stride: usize,
724 padding: usize,
725 ) -> Result<()>
726 where
727 T: FloatElement + Send + Sync + std::ops::AddAssign,
728 {
729 #[cfg(feature = "profiling")]
730 {
731 }
733
734 let output_h = (input_h + 2 * padding - kernel_h) / stride + 1;
736 let output_w = (input_w + 2 * padding - kernel_w) / stride + 1;
737 let expected_output_size = output_h * output_w;
738
739 if output.len() < expected_output_size {
741 return Err(torsh_core::error::TorshError::InvalidShape(format!(
742 "Output buffer too small: expected at least {} ({}x{}) elements, got {}",
743 expected_output_size,
744 output_h,
745 output_w,
746 output.len()
747 )));
748 }
749
750 if kernel_h * kernel_w <= 9 && input_h * input_w > 10000 {
766 self.direct_conv2d(
768 input, kernel, output, input_h, input_w, kernel_h, kernel_w, stride, padding,
769 )
770 } else if kernel_h >= 7 && kernel_w >= 7 {
771 self.fft_conv2d(
773 input, kernel, output, input_h, input_w, kernel_h, kernel_w, stride, padding,
774 )
775 } else {
776 self.winograd_conv2d(
778 input, kernel, output, input_h, input_w, kernel_h, kernel_w, stride, padding,
779 )
780 }
781 }
782
783 fn direct_conv2d<T>(
785 &self,
786 input: &[T],
787 kernel: &[T],
788 output: &mut [T],
789 input_h: usize,
790 input_w: usize,
791 kernel_h: usize,
792 kernel_w: usize,
793 stride: usize,
794 padding: usize,
795 ) -> Result<()>
796 where
797 T: FloatElement + Send + Sync + std::ops::AddAssign,
798 {
799 let output_h = (input_h + 2 * padding - kernel_h) / stride + 1;
800 let output_w = (input_w + 2 * padding - kernel_w) / stride + 1;
801
802 let output_positions: Vec<_> = (0..output_h)
804 .flat_map(|out_y| (0..output_w).map(move |out_x| (out_y, out_x)))
805 .collect();
806
807 let results: Vec<_> = parallel_map_collect(output_positions, |(out_y, out_x)| {
808 let mut sum = <T as torsh_core::TensorElement>::zero();
809
810 for ky in 0..kernel_h {
811 for kx in 0..kernel_w {
812 let in_y = out_y * stride + ky;
813 let in_x = out_x * stride + kx;
814
815 if in_y >= padding
816 && in_y < input_h + padding
817 && in_x >= padding
818 && in_x < input_w + padding
819 {
820 let input_y = in_y - padding;
821 let input_x = in_x - padding;
822
823 if input_y < input_h && input_x < input_w {
824 sum += input[input_y * input_w + input_x] * kernel[ky * kernel_w + kx];
825 }
826 }
827 }
828 }
829
830 (out_y * output_w + out_x, sum)
831 });
832
833 for (idx, value) in results {
835 output[idx] = value;
836 }
837
838 Ok(())
839 }
840
841 fn fft_conv2d<T>(
843 &self,
844 input: &[T],
845 kernel: &[T],
846 output: &mut [T],
847 input_h: usize,
848 input_w: usize,
849 kernel_h: usize,
850 kernel_w: usize,
851 stride: usize,
852 padding: usize,
853 ) -> Result<()>
854 where
855 T: FloatElement + std::ops::AddAssign,
856 {
857 self.direct_conv2d(
860 input, kernel, output, input_h, input_w, kernel_h, kernel_w, stride, padding,
861 )
862 }
863
864 fn winograd_conv2d<T>(
866 &self,
867 input: &[T],
868 kernel: &[T],
869 output: &mut [T],
870 input_h: usize,
871 input_w: usize,
872 kernel_h: usize,
873 kernel_w: usize,
874 stride: usize,
875 padding: usize,
876 ) -> Result<()>
877 where
878 T: FloatElement + std::ops::AddAssign,
879 {
880 self.direct_conv2d(
883 input, kernel, output, input_h, input_w, kernel_h, kernel_w, stride, padding,
884 )
885 }
886
887 pub fn execute_fused_operations<T>(
889 &self,
890 operations: &[FusedOperation<T>],
891 inputs: &[&[T]],
892 outputs: &mut [&mut [T]],
893 ) -> Result<()>
894 where
895 T: FloatElement + Send + Sync + std::ops::AddAssign,
896 {
897 if !self.config.enable_operation_fusion {
898 return Err(TorshError::InvalidArgument(
899 "Operation fusion disabled".to_string(),
900 ));
901 }
902
903 #[cfg(feature = "profiling")]
904 {
905 }
907
908 let compiled = self.compile_fusion(operations)?;
910 compiled.execute(inputs, outputs)
911 }
912
913 fn compile_fusion<T>(&self, operations: &[FusedOperation<T>]) -> Result<CompiledFusion<T>>
915 where
916 T: FloatElement + std::ops::AddAssign,
917 {
918 let plan = ExecutionPlan {
920 operations: operations.to_vec(),
921 optimization_level: OptimizationLevel::Aggressive,
922 };
923
924 Ok(CompiledFusion {
925 plan,
926 estimated_flops: self.estimate_fusion_flops(operations),
927 })
928 }
929
930 fn estimate_fusion_flops<T>(&self, operations: &[FusedOperation<T>]) -> usize
932 where
933 T: FloatElement + std::ops::AddAssign,
934 {
935 operations.len() * 1000 }
938
939 pub fn get_performance_stats(&self) -> AlgorithmPerformanceStats {
941 let history = self.performance_history.read_or_recover();
942
943 let mut total_operations = 0;
944 let mut algorithm_counts = HashMap::new();
945
946 for metrics in history.values() {
947 total_operations += metrics.execution_count;
948 if let Some(ref algorithm) = metrics.best_algorithm {
949 *algorithm_counts.entry(algorithm.clone()).or_insert(0) += 1;
950 }
951 }
952
953 AlgorithmPerformanceStats {
954 total_operations,
955 unique_operation_signatures: history.len(),
956 algorithm_distribution: algorithm_counts,
957 average_speedup: self.calculate_average_speedup(&history),
958 }
959 }
960
961 fn calculate_average_speedup(
963 &self,
964 history: &HashMap<OperationSignature, PerformanceMetrics>,
965 ) -> f64 {
966 if history.is_empty() {
967 return 1.0;
968 }
969
970 let speedups: Vec<f64> = history
971 .values()
972 .filter_map(|metrics| metrics.best_speedup)
973 .collect();
974
975 if speedups.is_empty() {
976 1.0
977 } else {
978 speedups.iter().sum::<f64>() / speedups.len() as f64
979 }
980 }
981}
982
983impl Default for AlgorithmicOptimizer {
984 fn default() -> Self {
985 Self::new()
986 }
987}
988
989#[derive(Debug, Clone, Hash, PartialEq, Eq)]
991enum OperationSignature {
992 MatMul { m: usize, k: usize, n: usize },
993}
994
995#[derive(Debug, Clone, PartialEq, Eq, Hash)]
997pub enum MatMulAlgorithm {
998 Naive,
999 Blocked,
1000 Strassen,
1001 CacheOblivious,
1002 Parallel,
1003}
1004
1005#[derive(Debug, Clone)]
1007struct PerformanceMetrics {
1008 execution_count: usize,
1009 algorithm_timings: HashMap<MatMulAlgorithm, Vec<std::time::Duration>>,
1010 best_algorithm: Option<MatMulAlgorithm>,
1011 best_speedup: Option<f64>,
1012}
1013
1014impl Default for PerformanceMetrics {
1015 fn default() -> Self {
1016 Self {
1017 execution_count: 0,
1018 algorithm_timings: HashMap::new(),
1019 best_algorithm: None,
1020 best_speedup: None,
1021 }
1022 }
1023}
1024
1025impl PerformanceMetrics {
1026 fn update_performance(&mut self, algorithm: MatMulAlgorithm, duration: std::time::Duration) {
1027 self.execution_count += 1;
1028 self.algorithm_timings
1029 .entry(algorithm.clone())
1030 .or_insert_with(Vec::new)
1031 .push(duration);
1032
1033 let avg_duration = self.average_duration(&algorithm);
1035 let current_best_duration = self
1036 .best_algorithm
1037 .as_ref()
1038 .map(|alg| self.average_duration(alg))
1039 .unwrap_or(std::time::Duration::from_secs(u64::MAX));
1040
1041 if avg_duration < current_best_duration {
1042 let speedup = current_best_duration.as_secs_f64() / avg_duration.as_secs_f64();
1043 self.best_algorithm = Some(algorithm);
1044 self.best_speedup = Some(speedup);
1045 }
1046 }
1047
1048 fn average_duration(&self, algorithm: &MatMulAlgorithm) -> std::time::Duration {
1049 static EMPTY_VEC: Vec<std::time::Duration> = Vec::new();
1050 let timings = self.algorithm_timings.get(algorithm).unwrap_or(&EMPTY_VEC);
1051 if timings.is_empty() {
1052 return std::time::Duration::from_secs(u64::MAX);
1053 }
1054
1055 let total_nanos: u128 = timings.iter().map(|d| d.as_nanos()).sum();
1056 std::time::Duration::from_nanos((total_nanos / timings.len() as u128) as u64)
1057 }
1058}
1059
1060#[derive(Debug, Clone)]
1062pub enum FusedOperation<T> {
1063 ElementwiseAdd {
1064 alpha: T,
1065 },
1066 ElementwiseMul {
1067 scale: T,
1068 },
1069 ReLU,
1070 Sigmoid,
1071 MatMul {
1072 transpose_a: bool,
1073 transpose_b: bool,
1074 },
1075}
1076
1077#[allow(dead_code)]
1079#[derive(Debug, Clone, Hash, PartialEq, Eq)]
1080struct FusionSignature {
1081 operation_types: Vec<String>,
1082 tensor_shapes: Vec<Vec<usize>>,
1083}
1084
1085#[allow(dead_code)]
1086impl FusionSignature {
1087 fn from_operations<T>(operations: &[FusedOperation<T>]) -> Self
1088 where
1089 T: FloatElement + std::ops::AddAssign,
1090 {
1091 let operation_types = operations.iter().map(|op| format!("{:?}", op)).collect();
1092
1093 Self {
1094 operation_types,
1095 tensor_shapes: vec![], }
1097 }
1098}
1099
1100#[allow(dead_code)]
1102#[derive(Debug, Clone)]
1103struct CompiledFusion<T> {
1104 plan: ExecutionPlan<T>,
1105 estimated_flops: usize,
1106}
1107
1108impl<T> CompiledFusion<T> {
1109 fn execute(&self, inputs: &[&[T]], outputs: &mut [&mut [T]]) -> Result<()>
1110 where
1111 T: FloatElement + std::ops::AddAssign,
1112 {
1113 self.plan.execute(inputs, outputs)
1115 }
1116}
1117
1118#[allow(dead_code)]
1120#[derive(Debug, Clone)]
1121struct ExecutionPlan<T> {
1122 operations: Vec<FusedOperation<T>>,
1123 optimization_level: OptimizationLevel,
1124}
1125
1126impl<T> ExecutionPlan<T> {
1127 fn execute(&self, inputs: &[&[T]], outputs: &mut [&mut [T]]) -> Result<()>
1128 where
1129 T: FloatElement + std::ops::AddAssign,
1130 {
1131 if outputs.is_empty() || inputs.is_empty() {
1132 return Ok(());
1133 }
1134
1135 let output = outputs.get_mut(0).ok_or_else(|| {
1138 torsh_core::error::TorshError::InvalidShape("No output buffer".to_string())
1139 })?;
1140
1141 if let Some(first_input) = inputs.first() {
1143 if first_input.len() == output.len() {
1144 output.copy_from_slice(first_input);
1145 }
1146 }
1147
1148 for op in &self.operations {
1150 match op {
1151 FusedOperation::ElementwiseAdd { alpha } => {
1152 for val in output.iter_mut() {
1153 *val += *alpha;
1154 }
1155 }
1156 FusedOperation::ElementwiseMul { scale } => {
1157 for val in output.iter_mut() {
1158 *val = *val * *scale;
1159 }
1160 }
1161 FusedOperation::ReLU => {
1162 let zero = <T as torsh_core::dtype::TensorElement>::zero();
1163 for val in output.iter_mut() {
1164 if *val < zero {
1165 *val = zero;
1166 }
1167 }
1168 }
1169 FusedOperation::Sigmoid => {
1170 let one = <T as num_traits::One>::one();
1171 for val in output.iter_mut() {
1172 let exp_neg = (-*val).exp();
1174 *val = one / (one + exp_neg);
1175 }
1176 }
1177 FusedOperation::MatMul { .. } => {
1178 }
1181 }
1182 }
1183
1184 Ok(())
1185 }
1186}
1187
1188#[allow(dead_code)]
1190#[derive(Debug, Clone, Copy)]
1191enum OptimizationLevel {
1192 Conservative,
1193 Moderate,
1194 Aggressive,
1195}
1196
1197#[derive(Debug)]
1199pub struct AlgorithmPerformanceStats {
1200 pub total_operations: usize,
1201 pub unique_operation_signatures: usize,
1202 pub algorithm_distribution: HashMap<MatMulAlgorithm, usize>,
1203 pub average_speedup: f64,
1204}
1205
1206#[cfg(test)]
1207mod tests {
1208 use super::*;
1209
1210 #[test]
1211 fn test_algorithm_config_default() {
1212 let config = AlgorithmConfig::default();
1213 assert!(config.enable_adaptive_selection);
1214 assert!(config.enable_operation_fusion);
1215 assert!(config.enable_numerical_stability);
1216 }
1217
1218 #[test]
1219 fn test_algorithmic_optimizer_creation() {
1220 let optimizer = AlgorithmicOptimizer::new();
1221 let stats = optimizer.get_performance_stats();
1222
1223 assert_eq!(stats.total_operations, 0);
1224 assert_eq!(stats.unique_operation_signatures, 0);
1225 }
1226
1227 #[test]
1228 fn test_algorithm_selection() {
1229 let optimizer = AlgorithmicOptimizer::new();
1230 let signature = OperationSignature::MatMul {
1231 m: 100,
1232 k: 100,
1233 n: 100,
1234 };
1235
1236 let algorithm = optimizer.select_matmul_algorithm(&signature);
1237 assert!(matches!(algorithm, MatMulAlgorithm::Parallel));
1239 }
1240
1241 #[test]
1242 fn test_small_matrix_multiplication() {
1243 let optimizer = AlgorithmicOptimizer::new();
1244
1245 let a = vec![1.0f32, 2.0, 3.0, 4.0]; let b = vec![5.0f32, 6.0, 7.0, 8.0]; let mut c = vec![0.0f32; 4]; optimizer
1250 .optimized_matmul(&a, &b, &mut c, 2, 2, 2)
1251 .expect("optimized_matmul should succeed");
1252
1253 assert!((c[0] - 19.0).abs() < 1e-6);
1255 assert!((c[1] - 22.0).abs() < 1e-6);
1256 assert!((c[2] - 43.0).abs() < 1e-6);
1257 assert!((c[3] - 50.0).abs() < 1e-6);
1258 }
1259
1260 #[test]
1261 fn test_block_size_calculation() {
1262 let optimizer = AlgorithmicOptimizer::new();
1263 let block_size = optimizer.calculate_optimal_block_size(1000, 1000, 1000);
1264
1265 assert!(block_size >= 16);
1266 assert!(block_size <= 256);
1267 }
1268
1269 #[test]
1270 fn test_conv2d_basic() {
1271 let optimizer = AlgorithmicOptimizer::new();
1272
1273 let input = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
1275 let kernel = vec![1.0f32, 0.0, 0.0, 1.0];
1276 let mut output = vec![0.0f32; 4]; optimizer
1279 .optimized_conv2d(&input, &kernel, &mut output, 3, 3, 2, 2, 1, 0)
1280 .expect("operation should succeed");
1281
1282 assert!(output.iter().all(|&x| x >= 0.0));
1284 }
1285
1286 #[test]
1287 fn test_performance_metrics() {
1288 let mut metrics = PerformanceMetrics::default();
1289
1290 let duration = std::time::Duration::from_millis(100);
1291 metrics.update_performance(MatMulAlgorithm::Blocked, duration);
1292
1293 assert_eq!(metrics.execution_count, 1);
1294 assert!(metrics.best_algorithm.is_some());
1295 }
1296
1297 #[test]
1298 fn test_fusion_signature() {
1299 let operations = vec![
1300 FusedOperation::ElementwiseAdd { alpha: 1.0f32 },
1301 FusedOperation::ReLU,
1302 ];
1303
1304 let signature = FusionSignature::from_operations(&operations);
1305 assert_eq!(signature.operation_types.len(), 2);
1306 }
1307}