1#![allow(dead_code)]
22
23use anyhow::Result;
24use serde::{Deserialize, Serialize};
25use std::collections::HashMap;
26use tokio::time::{Duration, Instant};
27use tracing::{debug, info};
28
29#[derive(Debug)]
34pub struct AICodeAnalyzer {
35 config: AIAnalysisConfig,
36 analysis_cache: HashMap<String, CachedAnalysis>,
37 pattern_database: ModelPatternDatabase,
38 performance_monitor: AnalysisPerformanceMonitor,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct AIAnalysisConfig {
44 pub enable_deep_analysis: bool,
46 pub enable_pattern_recognition: bool,
48 pub enable_optimization_suggestions: bool,
50 pub enable_vulnerability_detection: bool,
52 pub enable_performance_prediction: bool,
54 pub max_analysis_time_secs: u64,
56 pub confidence_threshold: f64,
58 pub enable_caching: bool,
60 pub cache_expiration_hours: u64,
62}
63
64impl Default for AIAnalysisConfig {
65 fn default() -> Self {
66 Self {
67 enable_deep_analysis: true,
68 enable_pattern_recognition: true,
69 enable_optimization_suggestions: true,
70 enable_vulnerability_detection: true,
71 enable_performance_prediction: true,
72 max_analysis_time_secs: 30,
73 confidence_threshold: 0.75,
74 enable_caching: true,
75 cache_expiration_hours: 24,
76 }
77 }
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize)]
82struct CachedAnalysis {
83 result: CodeAnalysisResult,
84 timestamp: std::time::SystemTime,
85 code_hash: String,
86}
87
88#[derive(Debug)]
90struct AnalysisPerformanceMonitor {
91 analysis_count: u64,
92 total_analysis_time: Duration,
93 cache_hits: u64,
94 cache_misses: u64,
95}
96
97impl AnalysisPerformanceMonitor {
98 fn new() -> Self {
99 Self {
100 analysis_count: 0,
101 total_analysis_time: Duration::from_secs(0),
102 cache_hits: 0,
103 cache_misses: 0,
104 }
105 }
106
107 fn record_analysis(&mut self, duration: Duration, cache_hit: bool) {
108 self.analysis_count += 1;
109 self.total_analysis_time += duration;
110 if cache_hit {
111 self.cache_hits += 1;
112 } else {
113 self.cache_misses += 1;
114 }
115 }
116
117 fn average_analysis_time(&self) -> Duration {
118 if self.analysis_count > 0 {
119 self.total_analysis_time / self.analysis_count as u32
120 } else {
121 Duration::from_secs(0)
122 }
123 }
124
125 fn cache_hit_rate(&self) -> f64 {
126 let total = self.cache_hits + self.cache_misses;
127 if total > 0 {
128 self.cache_hits as f64 / total as f64
129 } else {
130 0.0
131 }
132 }
133}
134
135impl AICodeAnalyzer {
136 pub fn new(config: AIAnalysisConfig) -> Self {
138 Self {
139 config,
140 analysis_cache: HashMap::new(),
141 pattern_database: ModelPatternDatabase::new(),
142 performance_monitor: AnalysisPerformanceMonitor::new(),
143 }
144 }
145
146 pub async fn analyze_model_code(
148 &mut self,
149 code: &str,
150 context: ModelContext,
151 ) -> Result<CodeAnalysisResult> {
152 let start_time = Instant::now();
153 let code_hash = self.compute_code_hash(code);
154
155 if self.config.enable_caching {
157 if let Some(cached) = self.get_cached_analysis(&code_hash) {
158 let result = cached.result.clone();
159 self.performance_monitor.record_analysis(start_time.elapsed(), true);
160 return Ok(result);
161 }
162 }
163
164 info!(
165 "Starting AI code analysis for {} lines of code",
166 code.lines().count()
167 );
168
169 let mut result = CodeAnalysisResult::new();
170
171 if self.config.enable_pattern_recognition {
173 let patterns = self.detect_code_patterns(code, &context).await?;
174 result.detected_patterns = patterns;
175 }
176
177 if self.config.enable_deep_analysis {
179 let issues = self.perform_deep_analysis(code, &context).await?;
180 result.identified_issues = issues;
181 }
182
183 if self.config.enable_optimization_suggestions {
185 let optimizations = self.generate_optimization_suggestions(code, &context).await?;
186 result.optimization_suggestions = optimizations;
187 }
188
189 if self.config.enable_vulnerability_detection {
191 let vulnerabilities = self.detect_vulnerabilities(code, &context).await?;
192 result.security_issues = vulnerabilities;
193 }
194
195 if self.config.enable_performance_prediction {
197 let predictions = self.predict_performance_characteristics(code, &context).await?;
198 result.performance_predictions = predictions;
199 }
200
201 result.quality_score = self.calculate_quality_score(&result);
203 result.analysis_metadata = AnalysisMetadata {
204 analysis_duration: start_time.elapsed(),
205 confidence_score: self.calculate_confidence_score(&result),
206 analyzer_version: "1.0.0".to_string(),
207 timestamp: std::time::SystemTime::now(),
208 };
209
210 if self.config.enable_caching {
212 self.cache_analysis(code_hash, &result);
213 }
214
215 self.performance_monitor.record_analysis(start_time.elapsed(), false);
216
217 info!(
218 "AI code analysis completed in {:?} with quality score: {:.2}",
219 start_time.elapsed(),
220 result.quality_score
221 );
222
223 Ok(result)
224 }
225
226 pub async fn analyze_tensor_operations(
228 &self,
229 operations: &[TensorOperation],
230 ) -> Result<TensorOptimizationReport> {
231 debug!("Analyzing {} tensor operations", operations.len());
232
233 let mut report = TensorOptimizationReport::new();
234
235 report.fusion_opportunities = self.detect_fusion_opportunities(operations).await?;
237 report.memory_optimizations = self.detect_memory_optimizations(operations).await?;
238 report.parallelization_opportunities =
239 self.detect_parallelization_opportunities(operations).await?;
240 report.redundant_operations = self.detect_redundant_operations(operations).await?;
241
242 report.estimated_speedup = self.estimate_optimization_speedup(&report);
244 report.estimated_memory_savings = self.estimate_memory_savings(&report);
245
246 Ok(report)
247 }
248
249 pub async fn automated_debugging_assistance(
251 &self,
252 error_context: &ErrorContext,
253 ) -> Result<DebuggingAssistance> {
254 info!(
255 "Providing automated debugging assistance for error: {}",
256 error_context.error_type
257 );
258
259 let mut assistance = DebuggingAssistance::new();
260
261 assistance.probable_causes = self.analyze_error_patterns(error_context).await?;
263 assistance.suggested_fixes = self.generate_suggested_fixes(error_context).await?;
264 assistance.debugging_steps = self.generate_debugging_steps(error_context).await?;
265 assistance.related_documentation = self.find_related_documentation(error_context).await?;
266
267 assistance.confidence_score = self.calculate_debugging_confidence(&assistance);
269
270 Ok(assistance)
271 }
272
273 pub fn get_performance_metrics(&self) -> AnalysisPerformanceMetrics {
275 AnalysisPerformanceMetrics {
276 total_analyses: self.performance_monitor.analysis_count,
277 average_analysis_time: self.performance_monitor.average_analysis_time(),
278 cache_hit_rate: self.performance_monitor.cache_hit_rate(),
279 cached_results: self.analysis_cache.len(),
280 }
281 }
282
283 async fn detect_code_patterns(
286 &self,
287 code: &str,
288 context: &ModelContext,
289 ) -> Result<Vec<DetectedPattern>> {
290 debug!("Detecting code patterns");
291
292 let mut patterns = Vec::new();
293
294 if code.contains("torch.cuda.empty_cache()") && context.model_type == ModelType::Production
296 {
297 patterns.push(DetectedPattern {
298 pattern_type: PatternType::AntiPattern,
299 name: "Frequent CUDA Cache Clearing".to_string(),
300 description: "Frequent CUDA cache clearing can hurt performance".to_string(),
301 severity: Severity::Medium,
302 confidence: 0.85,
303 recommendations: vec![
304 "Consider using gradient accumulation instead".to_string(),
305 "Review memory management strategy".to_string(),
306 ],
307 });
308 }
309
310 if code.contains("grad_norm") && code.contains("clip") {
312 patterns.push(DetectedPattern {
313 pattern_type: PatternType::GoodPattern,
314 name: "Gradient Clipping".to_string(),
315 description: "Proper gradient clipping implementation detected".to_string(),
316 severity: Severity::Info,
317 confidence: 0.9,
318 recommendations: vec!["Consider adaptive gradient clipping".to_string()],
319 });
320 }
321
322 if code.contains("detach()") && code.contains("requires_grad") {
324 patterns.push(DetectedPattern {
325 pattern_type: PatternType::OptimizationOpportunity,
326 name: "Gradient Computation Inefficiency".to_string(),
327 description: "Potential inefficient gradient computation detected".to_string(),
328 severity: Severity::Medium,
329 confidence: 0.75,
330 recommendations: vec![
331 "Consider using torch.no_grad() context".to_string(),
332 "Review gradient requirements".to_string(),
333 ],
334 });
335 }
336
337 Ok(patterns)
338 }
339
340 async fn perform_deep_analysis(
344 &self,
345 code: &str,
346 _context: &ModelContext,
347 ) -> Result<Vec<IdentifiedIssue>> {
348 debug!("Performing rule-based deep code analysis");
349
350 let mut issues = Vec::new();
351
352 if code.contains("log") && !code.contains("log1p") && code.contains("softmax") {
354 issues.push(IdentifiedIssue {
355 issue_type: IssueType::NumericalStability,
356 title: "Potential Numerical Instability in Log-Softmax".to_string(),
357 description: "Using log(softmax(x)) can cause numerical instability. Consider using log_softmax directly.".to_string(),
358 severity: Severity::High,
359 confidence: 0.88,
360 suggested_fix: "Replace log(softmax(x)) with log_softmax(x)".to_string(),
361 code_location: locate_in_code(code, "softmax"),
362 });
363 }
364
365 if code.contains("attention") && code.contains("matmul") && !code.contains("flash") {
367 issues.push(IdentifiedIssue {
368 issue_type: IssueType::Performance,
369 title: "Inefficient Attention Implementation".to_string(),
370 description:
371 "Standard attention implementation may be inefficient for large sequences."
372 .to_string(),
373 severity: Severity::Medium,
374 confidence: 0.75,
375 suggested_fix:
376 "Consider using Flash Attention or other optimized attention mechanisms"
377 .to_string(),
378 code_location: locate_in_code(code, "attention"),
379 });
380 }
381
382 if code.contains("accumulate") && !code.contains("zero_grad") {
384 issues.push(IdentifiedIssue {
385 issue_type: IssueType::MemoryLeak,
386 title: "Potential Gradient Accumulation Memory Leak".to_string(),
387 description: "Gradient accumulation without zero_grad() can cause memory leaks."
388 .to_string(),
389 severity: Severity::High,
390 confidence: 0.82,
391 suggested_fix: "Ensure optimizer.zero_grad() is called appropriately".to_string(),
392 code_location: locate_in_code(code, "accumulate"),
393 });
394 }
395
396 Ok(issues)
397 }
398
399 async fn generate_optimization_suggestions(
400 &self,
401 code: &str,
402 context: &ModelContext,
403 ) -> Result<Vec<OptimizationSuggestion>> {
404 debug!("Generating optimization suggestions");
405
406 let mut suggestions = Vec::new();
407
408 if context.model_type == ModelType::Training && !code.contains("autocast") {
410 suggestions.push(OptimizationSuggestion {
411 optimization_type: OptimizationType::MixedPrecision,
412 title: "Enable Mixed Precision Training".to_string(),
413 description: "Mixed precision training can significantly speed up training and reduce memory usage.".to_string(),
414 potential_speedup: 1.5,
415 memory_savings: 0.4,
416 implementation_effort: ImplementationEffort::Low,
417 confidence: 0.9,
418 code_example: Some("with torch.autocast(device_type='cuda', dtype=torch.float16):".to_string()),
419 });
420 }
421
422 if context.model_type == ModelType::Production && !code.contains("compile") {
424 suggestions.push(OptimizationSuggestion {
425 optimization_type: OptimizationType::ModelCompilation,
426 title: "Enable Model Compilation".to_string(),
427 description: "Model compilation can provide significant inference speedups."
428 .to_string(),
429 potential_speedup: 2.0,
430 memory_savings: 0.0,
431 implementation_effort: ImplementationEffort::Low,
432 confidence: 0.85,
433 code_example: Some("model = torch.compile(model)".to_string()),
434 });
435 }
436
437 if context.model_size > 1_000_000_000 && !code.contains("checkpoint") {
439 suggestions.push(OptimizationSuggestion {
440 optimization_type: OptimizationType::MemoryOptimization,
441 title: "Enable Gradient Checkpointing".to_string(),
442 description:
443 "Gradient checkpointing can significantly reduce memory usage for large models."
444 .to_string(),
445 potential_speedup: 0.8, memory_savings: 0.6,
447 implementation_effort: ImplementationEffort::Medium,
448 confidence: 0.88,
449 code_example: Some("torch.utils.checkpoint.checkpoint(layer, x)".to_string()),
450 });
451 }
452
453 Ok(suggestions)
454 }
455
456 async fn detect_vulnerabilities(
457 &self,
458 code: &str,
459 context: &ModelContext,
460 ) -> Result<Vec<SecurityIssue>> {
461 debug!("Detecting security vulnerabilities");
462
463 let mut vulnerabilities = Vec::new();
464
465 if code.contains("pickle.load") && !code.contains("safe_load") {
467 vulnerabilities.push(SecurityIssue {
468 vulnerability_type: VulnerabilityType::CodeExecution,
469 title: "Unsafe Pickle Loading".to_string(),
470 description:
471 "Loading pickle files can execute arbitrary code. Use safe alternatives."
472 .to_string(),
473 severity: Severity::Critical,
474 confidence: 0.95,
475 mitigation: "Use torch.load with weights_only=True or safetensors".to_string(),
476 cve_references: vec!["CWE-502".to_string()],
477 });
478 }
479
480 if code.contains("state_dict")
482 && code.contains("save")
483 && context.model_type == ModelType::Production
484 {
485 vulnerabilities.push(SecurityIssue {
486 vulnerability_type: VulnerabilityType::DataExposure,
487 title: "Potential Model Parameter Exposure".to_string(),
488 description: "Saving full model state may expose sensitive parameters.".to_string(),
489 severity: Severity::Medium,
490 confidence: 0.7,
491 mitigation: "Consider differential privacy or parameter encryption".to_string(),
492 cve_references: vec![],
493 });
494 }
495
496 if code.contains("input") && !code.contains("validate") && !code.contains("sanitize") {
498 vulnerabilities.push(SecurityIssue {
499 vulnerability_type: VulnerabilityType::InputValidation,
500 title: "Missing Input Validation".to_string(),
501 description: "Input validation is important for preventing adversarial attacks."
502 .to_string(),
503 severity: Severity::Medium,
504 confidence: 0.65,
505 mitigation: "Implement input validation and sanitization".to_string(),
506 cve_references: vec![],
507 });
508 }
509
510 Ok(vulnerabilities)
511 }
512
513 async fn predict_performance_characteristics(
517 &self,
518 code: &str,
519 context: &ModelContext,
520 ) -> Result<PerformancePredictions> {
521 debug!("Predicting performance characteristics");
522
523 let mut predictions = PerformancePredictions::new();
524
525 predictions.estimated_memory_usage = self.estimate_memory_usage(code, context);
527 predictions.estimated_training_time = self.estimate_training_time(code, context);
528 predictions.estimated_inference_latency = self.estimate_inference_latency(code, context);
529 predictions.scaling_characteristics = self.predict_scaling_behavior(code, context);
530
531 let mut signal_count: u32 = 0;
536 if code.contains("attention") && !code.contains("flash") {
537 predictions.predicted_bottlenecks.push(
538 "Attention computation may become a bottleneck for long sequences (no Flash \
539 Attention detected in this code)"
540 .to_string(),
541 );
542 signal_count += 1;
543 }
544 if context.model_size > 1_000_000_000 {
545 predictions.predicted_bottlenecks.push(
546 "Memory bandwidth may limit performance for large batch sizes (model exceeds \
547 1B parameters)"
548 .to_string(),
549 );
550 signal_count += 1;
551 }
552
553 predictions.confidence_score = match signal_count {
558 0 => 0.0,
559 1 => 0.6,
560 _ => 0.75,
561 };
562
563 Ok(predictions)
564 }
565
566 async fn detect_fusion_opportunities(
567 &self,
568 operations: &[TensorOperation],
569 ) -> Result<Vec<FusionOpportunity>> {
570 let mut opportunities = Vec::new();
571
572 for window in operations.windows(2) {
574 if let [op1, op2] = window {
575 if matches!(op1.op_type, OperationType::MatMul)
576 && matches!(op2.op_type, OperationType::Add)
577 {
578 opportunities.push(FusionOpportunity {
579 operations: vec![op1.clone(), op2.clone()],
580 fusion_type: FusionType::GEMM,
581 estimated_speedup: 1.3,
582 description: "MatMul + Add can be fused into GEMM operation".to_string(),
583 });
584 }
585 }
586 }
587
588 for window in operations.windows(2) {
590 if let [op1, op2] = window {
591 if matches!(op1.op_type, OperationType::Linear)
592 && matches!(op2.op_type, OperationType::Activation)
593 {
594 opportunities.push(FusionOpportunity {
595 operations: vec![op1.clone(), op2.clone()],
596 fusion_type: FusionType::LinearActivation,
597 estimated_speedup: 1.2,
598 description: "Linear + Activation can be fused".to_string(),
599 });
600 }
601 }
602 }
603
604 Ok(opportunities)
605 }
606
607 async fn detect_memory_optimizations(
608 &self,
609 operations: &[TensorOperation],
610 ) -> Result<Vec<MemoryOptimization>> {
611 let mut optimizations = Vec::new();
612
613 for op in operations {
615 if op.can_be_inplace() && !op.is_inplace {
616 optimizations.push(MemoryOptimization {
617 operation: op.clone(),
618 optimization_type: MemoryOptimizationType::InPlace,
619 memory_savings: op.output_size_bytes,
620 description: format!("Operation {} can be performed in-place", op.name),
621 });
622 }
623 }
624
625 let mut tensor_usage = HashMap::new();
627 for op in operations {
628 for input in &op.inputs {
629 *tensor_usage.entry(input.clone()).or_insert(0) += 1;
630 }
631 }
632
633 for (tensor, usage_count) in tensor_usage {
634 if usage_count == 1 {
635 optimizations.push(MemoryOptimization {
636 operation: TensorOperation::default(),
637 optimization_type: MemoryOptimizationType::TensorReuse,
638 memory_savings: 0, description: format!("Tensor {} can be reused", tensor),
640 });
641 }
642 }
643
644 Ok(optimizations)
645 }
646
647 async fn detect_parallelization_opportunities(
648 &self,
649 operations: &[TensorOperation],
650 ) -> Result<Vec<ParallelizationOpportunity>> {
651 let mut opportunities = Vec::new();
652
653 for (i, op1) in operations.iter().enumerate() {
655 for op2 in operations.iter().skip(i + 1) {
656 if self.operations_are_independent(op1, op2) {
657 opportunities.push(ParallelizationOpportunity {
658 operations: vec![op1.clone(), op2.clone()],
659 parallelization_type: ParallelizationType::DataParallel,
660 estimated_speedup: 1.8,
661 description: "Operations can run in parallel".to_string(),
662 });
663 }
664 }
665 }
666
667 Ok(opportunities)
668 }
669
670 async fn detect_redundant_operations(
671 &self,
672 operations: &[TensorOperation],
673 ) -> Result<Vec<RedundantOperation>> {
674 let mut redundant = Vec::new();
675
676 for (i, op1) in operations.iter().enumerate() {
678 for (_j, op2) in operations.iter().enumerate().skip(i + 1) {
679 if self.operations_are_equivalent(op1, op2) {
680 redundant.push(RedundantOperation {
681 original_operation: op1.clone(),
682 redundant_operation: op2.clone(),
683 redundancy_type: RedundancyType::Duplicate,
684 description: "Operations produce identical results".to_string(),
685 });
686 }
687 }
688 }
689
690 Ok(redundant)
691 }
692
693 fn operations_are_independent(&self, op1: &TensorOperation, op2: &TensorOperation) -> bool {
696 for input1 in &op1.inputs {
698 for output2 in &op2.outputs {
699 if input1 == output2 {
700 return false;
701 }
702 }
703 }
704 for input2 in &op2.inputs {
705 for output1 in &op1.outputs {
706 if input2 == output1 {
707 return false;
708 }
709 }
710 }
711 true
712 }
713
714 fn operations_are_equivalent(&self, op1: &TensorOperation, op2: &TensorOperation) -> bool {
715 op1.op_type == op2.op_type && op1.inputs == op2.inputs && op1.parameters == op2.parameters
716 }
717
718 fn compute_code_hash(&self, code: &str) -> String {
719 use std::collections::hash_map::DefaultHasher;
720 use std::hash::{Hash, Hasher};
721
722 let mut hasher = DefaultHasher::new();
723 code.hash(&mut hasher);
724 format!("{:x}", hasher.finish())
725 }
726
727 fn get_cached_analysis(&self, code_hash: &str) -> Option<&CachedAnalysis> {
728 self.analysis_cache.get(code_hash).and_then(|cached| {
729 let age = std::time::SystemTime::now()
730 .duration_since(cached.timestamp)
731 .unwrap_or_default();
732
733 if age.as_secs() < self.config.cache_expiration_hours * 3600 {
734 Some(cached)
735 } else {
736 None
737 }
738 })
739 }
740
741 fn cache_analysis(&mut self, code_hash: String, result: &CodeAnalysisResult) {
742 self.analysis_cache.insert(
743 code_hash.clone(),
744 CachedAnalysis {
745 result: result.clone(),
746 timestamp: std::time::SystemTime::now(),
747 code_hash,
748 },
749 );
750 }
751
752 fn calculate_quality_score(&self, result: &CodeAnalysisResult) -> f64 {
753 let mut score: f64 = 100.0;
754
755 for issue in &result.identified_issues {
757 match issue.severity {
758 Severity::Critical => score -= 20.0,
759 Severity::High => score -= 10.0,
760 Severity::Medium => score -= 5.0,
761 Severity::Low => score -= 2.0,
762 Severity::Info => score -= 0.0,
763 }
764 }
765
766 for vulnerability in &result.security_issues {
768 match vulnerability.severity {
769 Severity::Critical => score -= 25.0,
770 Severity::High => score -= 15.0,
771 Severity::Medium => score -= 8.0,
772 Severity::Low => score -= 3.0,
773 Severity::Info => score -= 0.0,
774 }
775 }
776
777 for pattern in &result.detected_patterns {
779 if pattern.pattern_type == PatternType::GoodPattern {
780 score += 2.0;
781 }
782 }
783
784 score.max(0.0).min(100.0)
785 }
786
787 fn calculate_confidence_score(&self, result: &CodeAnalysisResult) -> f64 {
788 let mut total_confidence = 0.0;
789 let mut count = 0;
790
791 for issue in &result.identified_issues {
792 total_confidence += issue.confidence;
793 count += 1;
794 }
795
796 for pattern in &result.detected_patterns {
797 total_confidence += pattern.confidence;
798 count += 1;
799 }
800
801 if count > 0 {
802 total_confidence / count as f64
803 } else {
804 1.0
805 }
806 }
807
808 fn estimate_memory_usage(&self, code: &str, context: &ModelContext) -> f64 {
809 let base_memory = context.model_size as f64 * 4.0; let mut multiplier = 1.0;
813 if code.contains("gradient_accumulation") {
814 multiplier += 0.5;
815 }
816 if code.contains("mixed_precision") {
817 multiplier *= 0.7;
818 }
819
820 base_memory * multiplier / 1_000_000.0 }
822
823 fn estimate_training_time(&self, code: &str, context: &ModelContext) -> f64 {
824 let base_time = (context.model_size as f64).log10() * 10.0;
826
827 let mut multiplier = 1.0;
828 if code.contains("mixed_precision") {
829 multiplier *= 0.6;
830 }
831 if code.contains("gradient_checkpointing") {
832 multiplier *= 1.3;
833 }
834
835 base_time * multiplier
836 }
837
838 fn estimate_inference_latency(&self, code: &str, context: &ModelContext) -> f64 {
839 let base_latency = (context.model_size as f64).log10() * 5.0;
841
842 let mut multiplier = 1.0;
843 if code.contains("compile") {
844 multiplier *= 0.5;
845 }
846 if code.contains("quantization") {
847 multiplier *= 0.7;
848 }
849
850 base_latency * multiplier
851 }
852
853 fn predict_scaling_behavior(
854 &self,
855 _code: &str,
856 context: &ModelContext,
857 ) -> ScalingCharacteristics {
858 ScalingCharacteristics {
859 batch_size_scaling: if context.model_size > 1_000_000_000 {
860 ScalingBehavior::Sublinear
861 } else {
862 ScalingBehavior::Linear
863 },
864 sequence_length_scaling: ScalingBehavior::Quadratic, model_size_scaling: ScalingBehavior::Linear,
866 memory_scaling: ScalingBehavior::Linear,
867 }
868 }
869
870 fn estimate_optimization_speedup(&self, report: &TensorOptimizationReport) -> f64 {
871 let mut speedup = 1.0;
872
873 for fusion in &report.fusion_opportunities {
874 speedup *= fusion.estimated_speedup;
875 }
876
877 for parallel in &report.parallelization_opportunities {
878 speedup *= parallel.estimated_speedup;
879 }
880
881 speedup.min(10.0) }
883
884 fn estimate_memory_savings(&self, report: &TensorOptimizationReport) -> f64 {
885 let total_savings: u64 =
886 report.memory_optimizations.iter().map(|opt| opt.memory_savings).sum();
887
888 total_savings as f64 / 1_000_000.0 }
890
891 async fn analyze_error_patterns(
892 &self,
893 error_context: &ErrorContext,
894 ) -> Result<Vec<ProbableCause>> {
895 let mut causes = Vec::new();
896
897 match error_context.error_type.as_str() {
898 "OutOfMemoryError" => {
899 causes.push(ProbableCause {
900 cause: "Batch size too large".to_string(),
901 probability: 0.8,
902 evidence: vec!["GPU memory limit exceeded".to_string()],
903 });
904 causes.push(ProbableCause {
905 cause: "Model too large for available memory".to_string(),
906 probability: 0.6,
907 evidence: vec!["Model parameter count".to_string()],
908 });
909 },
910 "GradientExplosion" => {
911 causes.push(ProbableCause {
912 cause: "Learning rate too high".to_string(),
913 probability: 0.7,
914 evidence: vec!["Gradient norm increasing rapidly".to_string()],
915 });
916 },
917 _ => {
918 causes.push(ProbableCause {
919 cause: "Unknown error pattern".to_string(),
920 probability: 0.3,
921 evidence: vec![],
922 });
923 },
924 }
925
926 Ok(causes)
927 }
928
929 async fn generate_suggested_fixes(
930 &self,
931 error_context: &ErrorContext,
932 ) -> Result<Vec<SuggestedFix>> {
933 let mut fixes = Vec::new();
934
935 match error_context.error_type.as_str() {
936 "OutOfMemoryError" => {
937 fixes.push(SuggestedFix {
938 description: "Reduce batch size".to_string(),
939 implementation: "batch_size = batch_size // 2".to_string(),
940 confidence: 0.9,
941 estimated_impact: "Should free ~50% of memory".to_string(),
942 });
943 fixes.push(SuggestedFix {
944 description: "Enable gradient checkpointing".to_string(),
945 implementation: "model.gradient_checkpointing_enable()".to_string(),
946 confidence: 0.8,
947 estimated_impact: "Reduces memory by ~40% with 10-20% speed penalty"
948 .to_string(),
949 });
950 },
951 "GradientExplosion" => {
952 fixes.push(SuggestedFix {
953 description: "Add gradient clipping".to_string(),
954 implementation:
955 "torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)"
956 .to_string(),
957 confidence: 0.95,
958 estimated_impact: "Prevents gradient explosion".to_string(),
959 });
960 },
961 _ => {},
962 }
963
964 Ok(fixes)
965 }
966
967 async fn generate_debugging_steps(
968 &self,
969 error_context: &ErrorContext,
970 ) -> Result<Vec<DebuggingStep>> {
971 let mut steps = Vec::new();
972
973 steps.push(DebuggingStep {
974 step_number: 1,
975 description: "Check system resources".to_string(),
976 command: Some("nvidia-smi".to_string()),
977 expected_output: "GPU memory usage and availability".to_string(),
978 });
979
980 steps.push(DebuggingStep {
981 step_number: 2,
982 description: "Verify model configuration".to_string(),
983 command: Some("print(model)".to_string()),
984 expected_output: "Model architecture and parameter count".to_string(),
985 });
986
987 if error_context.error_type.as_str() == "OutOfMemoryError" {
988 steps.push(DebuggingStep {
989 step_number: 3,
990 description: "Check tensor shapes and batch size".to_string(),
991 command: Some(
992 "print(f'Batch size: {batch_size}, Input shape: {input.shape}')".to_string(),
993 ),
994 expected_output: "Current batch size and input dimensions".to_string(),
995 });
996 }
997
998 Ok(steps)
999 }
1000
1001 async fn find_related_documentation(
1002 &self,
1003 error_context: &ErrorContext,
1004 ) -> Result<Vec<DocumentationReference>> {
1005 let mut references = Vec::new();
1006
1007 match error_context.error_type.as_str() {
1008 "OutOfMemoryError" => {
1009 references.push(DocumentationReference {
1010 title: "Memory Management Best Practices".to_string(),
1011 url: "https://docs.trustformers.ai/memory-management".to_string(),
1012 relevance_score: 0.95,
1013 });
1014 references.push(DocumentationReference {
1015 title: "Gradient Checkpointing Guide".to_string(),
1016 url: "https://docs.trustformers.ai/gradient-checkpointing".to_string(),
1017 relevance_score: 0.8,
1018 });
1019 },
1020 "GradientExplosion" => {
1021 references.push(DocumentationReference {
1022 title: "Training Stability Guide".to_string(),
1023 url: "https://docs.trustformers.ai/training-stability".to_string(),
1024 relevance_score: 0.9,
1025 });
1026 },
1027 _ => {},
1028 }
1029
1030 Ok(references)
1031 }
1032
1033 fn calculate_debugging_confidence(&self, assistance: &DebuggingAssistance) -> f64 {
1034 let avg_cause_probability =
1035 assistance.probable_causes.iter().map(|cause| cause.probability).sum::<f64>()
1036 / assistance.probable_causes.len().max(1) as f64;
1037
1038 let avg_fix_confidence =
1039 assistance.suggested_fixes.iter().map(|fix| fix.confidence).sum::<f64>()
1040 / assistance.suggested_fixes.len().max(1) as f64;
1041
1042 (avg_cause_probability + avg_fix_confidence) / 2.0
1043 }
1044}
1045
1046#[derive(Debug)]
1050struct ModelPatternDatabase {
1051 patterns: HashMap<String, PatternDefinition>,
1052}
1053
1054impl ModelPatternDatabase {
1055 fn new() -> Self {
1056 let mut patterns = HashMap::new();
1057
1058 patterns.insert(
1060 "gradient_clipping".to_string(),
1061 PatternDefinition {
1062 name: "Gradient Clipping".to_string(),
1063 pattern_type: PatternType::GoodPattern,
1064 keywords: vec![
1065 "clip_grad_norm".to_string(),
1066 "gradient".to_string(),
1067 "clip".to_string(),
1068 ],
1069 severity: Severity::Info,
1070 description: "Proper gradient clipping prevents gradient explosion".to_string(),
1071 },
1072 );
1073
1074 Self { patterns }
1075 }
1076}
1077
1078#[derive(Debug, Clone)]
1079struct PatternDefinition {
1080 name: String,
1081 pattern_type: PatternType,
1082 keywords: Vec<String>,
1083 severity: Severity,
1084 description: String,
1085}
1086
1087#[derive(Debug, Clone)]
1089pub struct ModelContext {
1090 pub model_type: ModelType,
1091 pub model_size: u64, pub framework: String,
1093 pub target_hardware: String,
1094 pub training_stage: TrainingStage,
1095}
1096
1097#[derive(Debug, Clone, PartialEq)]
1098pub enum ModelType {
1099 Training,
1100 Inference,
1101 Production,
1102 Development,
1103}
1104
1105#[derive(Debug, Clone)]
1106pub enum TrainingStage {
1107 Training,
1108 Development,
1109 Pretraining,
1110 Finetuning,
1111 Evaluation,
1112 Inference,
1113}
1114
1115#[derive(Debug, Clone, Serialize, Deserialize)]
1117pub struct CodeAnalysisResult {
1118 pub quality_score: f64,
1119 pub detected_patterns: Vec<DetectedPattern>,
1120 pub identified_issues: Vec<IdentifiedIssue>,
1121 pub optimization_suggestions: Vec<OptimizationSuggestion>,
1122 pub security_issues: Vec<SecurityIssue>,
1123 pub performance_predictions: PerformancePredictions,
1124 pub analysis_metadata: AnalysisMetadata,
1125}
1126
1127impl CodeAnalysisResult {
1128 fn new() -> Self {
1129 Self {
1130 quality_score: 0.0,
1131 detected_patterns: Vec::new(),
1132 identified_issues: Vec::new(),
1133 optimization_suggestions: Vec::new(),
1134 security_issues: Vec::new(),
1135 performance_predictions: PerformancePredictions::new(),
1136 analysis_metadata: AnalysisMetadata::default(),
1137 }
1138 }
1139}
1140
1141#[derive(Debug, Clone, Serialize, Deserialize)]
1142pub struct DetectedPattern {
1143 pub pattern_type: PatternType,
1144 pub name: String,
1145 pub description: String,
1146 pub severity: Severity,
1147 pub confidence: f64,
1148 pub recommendations: Vec<String>,
1149}
1150
1151#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1152pub enum PatternType {
1153 GoodPattern,
1154 AntiPattern,
1155 OptimizationOpportunity,
1156 SecurityConcern,
1157}
1158
1159#[derive(Debug, Clone, Serialize, Deserialize)]
1160pub struct IdentifiedIssue {
1161 pub issue_type: IssueType,
1162 pub title: String,
1163 pub description: String,
1164 pub severity: Severity,
1165 pub confidence: f64,
1166 pub suggested_fix: String,
1167 pub code_location: Option<CodeLocation>,
1168}
1169
1170#[derive(Debug, Clone, Serialize, Deserialize)]
1171pub enum IssueType {
1172 NumericalStability,
1173 Performance,
1174 MemoryLeak,
1175 LogicError,
1176 TypeMismatch,
1177 ResourceLeak,
1178}
1179
1180#[derive(Debug, Clone, Serialize, Deserialize)]
1181pub struct CodeLocation {
1182 pub file: String,
1183 pub line: u32,
1184 pub column: u32,
1185}
1186
1187fn locate_in_code(code: &str, needle: &str) -> Option<CodeLocation> {
1196 let byte_offset = code.find(needle)?;
1197 let mut line: u32 = 1;
1198 let mut last_newline_offset: Option<usize> = None;
1199 for (idx, ch) in code[..byte_offset].char_indices() {
1200 if ch == '\n' {
1201 line += 1;
1202 last_newline_offset = Some(idx);
1203 }
1204 }
1205 let column = match last_newline_offset {
1206 Some(nl_idx) => (byte_offset - nl_idx) as u32,
1207 None => byte_offset as u32 + 1,
1208 };
1209
1210 Some(CodeLocation {
1211 file: "<analyzed source>".to_string(),
1212 line,
1213 column,
1214 })
1215}
1216
1217#[derive(Debug, Clone, Serialize, Deserialize)]
1218pub struct OptimizationSuggestion {
1219 pub optimization_type: OptimizationType,
1220 pub title: String,
1221 pub description: String,
1222 pub potential_speedup: f64,
1223 pub memory_savings: f64,
1224 pub implementation_effort: ImplementationEffort,
1225 pub confidence: f64,
1226 pub code_example: Option<String>,
1227}
1228
1229#[derive(Debug, Clone, Serialize, Deserialize)]
1230pub enum OptimizationType {
1231 MixedPrecision,
1232 ModelCompilation,
1233 MemoryOptimization,
1234 ComputationOptimization,
1235 IOOptimization,
1236 ParallelizationOptimization,
1237}
1238
1239#[derive(Debug, Clone, Serialize, Deserialize)]
1240pub enum ImplementationEffort {
1241 Low,
1242 Medium,
1243 High,
1244}
1245
1246#[derive(Debug, Clone, Serialize, Deserialize)]
1247pub struct SecurityIssue {
1248 pub vulnerability_type: VulnerabilityType,
1249 pub title: String,
1250 pub description: String,
1251 pub severity: Severity,
1252 pub confidence: f64,
1253 pub mitigation: String,
1254 pub cve_references: Vec<String>,
1255}
1256
1257#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1258pub enum VulnerabilityType {
1259 CodeExecution,
1260 DataExposure,
1261 InputValidation,
1262 AuthenticationBypass,
1263 PrivilegeEscalation,
1264}
1265
1266#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1267pub enum Severity {
1268 Critical,
1269 High,
1270 Medium,
1271 Low,
1272 Info,
1273}
1274
1275#[derive(Debug, Clone, Serialize, Deserialize)]
1276pub struct PerformancePredictions {
1277 pub estimated_memory_usage: f64, pub estimated_training_time: f64, pub estimated_inference_latency: f64, pub scaling_characteristics: ScalingCharacteristics,
1281 pub predicted_bottlenecks: Vec<String>,
1282 pub confidence_score: f64,
1283}
1284
1285impl PerformancePredictions {
1286 fn new() -> Self {
1287 Self {
1288 estimated_memory_usage: 0.0,
1289 estimated_training_time: 0.0,
1290 estimated_inference_latency: 0.0,
1291 scaling_characteristics: ScalingCharacteristics::default(),
1292 predicted_bottlenecks: Vec::new(),
1293 confidence_score: 0.0,
1294 }
1295 }
1296}
1297
1298#[derive(Debug, Clone, Serialize, Deserialize)]
1299pub struct ScalingCharacteristics {
1300 pub batch_size_scaling: ScalingBehavior,
1301 pub sequence_length_scaling: ScalingBehavior,
1302 pub model_size_scaling: ScalingBehavior,
1303 pub memory_scaling: ScalingBehavior,
1304}
1305
1306impl Default for ScalingCharacteristics {
1307 fn default() -> Self {
1308 Self {
1309 batch_size_scaling: ScalingBehavior::Linear,
1310 sequence_length_scaling: ScalingBehavior::Linear,
1311 model_size_scaling: ScalingBehavior::Linear,
1312 memory_scaling: ScalingBehavior::Linear,
1313 }
1314 }
1315}
1316
1317#[derive(Debug, Clone, Serialize, Deserialize)]
1318pub enum ScalingBehavior {
1319 Constant,
1320 Linear,
1321 Quadratic,
1322 Exponential,
1323 Sublinear,
1324}
1325
1326#[derive(Debug, Clone, Serialize, Deserialize)]
1327pub struct AnalysisMetadata {
1328 pub analysis_duration: Duration,
1329 pub confidence_score: f64,
1330 pub analyzer_version: String,
1331 pub timestamp: std::time::SystemTime,
1332}
1333
1334impl Default for AnalysisMetadata {
1335 fn default() -> Self {
1336 Self {
1337 analysis_duration: Duration::from_secs(0),
1338 confidence_score: 0.0,
1339 analyzer_version: "1.0.0".to_string(),
1340 timestamp: std::time::SystemTime::now(),
1341 }
1342 }
1343}
1344
1345#[derive(Debug, Clone)]
1348pub struct TensorOperation {
1349 pub name: String,
1350 pub op_type: OperationType,
1351 pub inputs: Vec<String>,
1352 pub outputs: Vec<String>,
1353 pub parameters: HashMap<String, String>,
1354 pub output_size_bytes: u64,
1355 pub is_inplace: bool,
1356}
1357
1358impl Default for TensorOperation {
1359 fn default() -> Self {
1360 Self {
1361 name: String::new(),
1362 op_type: OperationType::Unknown,
1363 inputs: Vec::new(),
1364 outputs: Vec::new(),
1365 parameters: HashMap::new(),
1366 output_size_bytes: 0,
1367 is_inplace: false,
1368 }
1369 }
1370}
1371
1372impl TensorOperation {
1373 fn can_be_inplace(&self) -> bool {
1374 matches!(
1375 self.op_type,
1376 OperationType::Add | OperationType::Mul | OperationType::Activation
1377 )
1378 }
1379}
1380
1381#[derive(Debug, Clone, PartialEq)]
1382pub enum OperationType {
1383 MatMul,
1384 Add,
1385 Mul,
1386 Conv2D,
1387 Linear,
1388 Activation,
1389 Pooling,
1390 BatchNorm,
1391 LayerNorm,
1392 Attention,
1393 Unknown,
1394}
1395
1396#[derive(Debug, Clone)]
1397pub struct TensorOptimizationReport {
1398 pub fusion_opportunities: Vec<FusionOpportunity>,
1399 pub memory_optimizations: Vec<MemoryOptimization>,
1400 pub parallelization_opportunities: Vec<ParallelizationOpportunity>,
1401 pub redundant_operations: Vec<RedundantOperation>,
1402 pub estimated_speedup: f64,
1403 pub estimated_memory_savings: f64,
1404}
1405
1406impl TensorOptimizationReport {
1407 fn new() -> Self {
1408 Self {
1409 fusion_opportunities: Vec::new(),
1410 memory_optimizations: Vec::new(),
1411 parallelization_opportunities: Vec::new(),
1412 redundant_operations: Vec::new(),
1413 estimated_speedup: 1.0,
1414 estimated_memory_savings: 0.0,
1415 }
1416 }
1417}
1418
1419#[derive(Debug, Clone)]
1420pub struct FusionOpportunity {
1421 pub operations: Vec<TensorOperation>,
1422 pub fusion_type: FusionType,
1423 pub estimated_speedup: f64,
1424 pub description: String,
1425}
1426
1427#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1428pub enum FusionType {
1429 GEMM,
1430 LinearActivation,
1431 ConvBatchNorm,
1432 AttentionQKV,
1433}
1434
1435#[derive(Debug, Clone)]
1436pub struct MemoryOptimization {
1437 pub operation: TensorOperation,
1438 pub optimization_type: MemoryOptimizationType,
1439 pub memory_savings: u64,
1440 pub description: String,
1441}
1442
1443#[derive(Debug, Clone)]
1444pub enum MemoryOptimizationType {
1445 InPlace,
1446 TensorReuse,
1447 MemoryPool,
1448 GradientCheckpointing,
1449}
1450
1451#[derive(Debug, Clone)]
1452pub struct ParallelizationOpportunity {
1453 pub operations: Vec<TensorOperation>,
1454 pub parallelization_type: ParallelizationType,
1455 pub estimated_speedup: f64,
1456 pub description: String,
1457}
1458
1459#[derive(Debug, Clone)]
1460pub enum ParallelizationType {
1461 DataParallel,
1462 ModelParallel,
1463 PipelineParallel,
1464 TensorParallel,
1465}
1466
1467#[derive(Debug, Clone)]
1468pub struct RedundantOperation {
1469 pub original_operation: TensorOperation,
1470 pub redundant_operation: TensorOperation,
1471 pub redundancy_type: RedundancyType,
1472 pub description: String,
1473}
1474
1475#[derive(Debug, Clone)]
1476pub enum RedundancyType {
1477 Duplicate,
1478 Subsumed,
1479 Unnecessary,
1480}
1481
1482#[derive(Debug, Clone)]
1485pub struct ErrorContext {
1486 pub error_type: String,
1487 pub error_message: String,
1488 pub stack_trace: Option<String>,
1489 pub system_info: SystemInfo,
1490 pub model_info: Option<ModelContext>,
1491}
1492
1493#[derive(Debug, Clone)]
1494pub struct SystemInfo {
1495 pub gpu_memory_total: u64,
1496 pub gpu_memory_used: u64,
1497 pub cpu_count: u32,
1498 pub ram_total: u64,
1499 pub ram_used: u64,
1500}
1501
1502#[derive(Debug, Clone)]
1503pub struct DebuggingAssistance {
1504 pub probable_causes: Vec<ProbableCause>,
1505 pub suggested_fixes: Vec<SuggestedFix>,
1506 pub debugging_steps: Vec<DebuggingStep>,
1507 pub related_documentation: Vec<DocumentationReference>,
1508 pub confidence_score: f64,
1509}
1510
1511impl DebuggingAssistance {
1512 fn new() -> Self {
1513 Self {
1514 probable_causes: Vec::new(),
1515 suggested_fixes: Vec::new(),
1516 debugging_steps: Vec::new(),
1517 related_documentation: Vec::new(),
1518 confidence_score: 0.0,
1519 }
1520 }
1521}
1522
1523#[derive(Debug, Clone)]
1524pub struct ProbableCause {
1525 pub cause: String,
1526 pub probability: f64,
1527 pub evidence: Vec<String>,
1528}
1529
1530#[derive(Debug, Clone)]
1531pub struct SuggestedFix {
1532 pub description: String,
1533 pub implementation: String,
1534 pub confidence: f64,
1535 pub estimated_impact: String,
1536}
1537
1538#[derive(Debug, Clone)]
1539pub struct DebuggingStep {
1540 pub step_number: u32,
1541 pub description: String,
1542 pub command: Option<String>,
1543 pub expected_output: String,
1544}
1545
1546#[derive(Debug, Clone)]
1547pub struct DocumentationReference {
1548 pub title: String,
1549 pub url: String,
1550 pub relevance_score: f64,
1551}
1552
1553#[derive(Debug, Serialize, Deserialize)]
1556pub struct AnalysisPerformanceMetrics {
1557 pub total_analyses: u64,
1558 pub average_analysis_time: Duration,
1559 pub cache_hit_rate: f64,
1560 pub cached_results: usize,
1561}
1562
1563#[macro_export]
1565macro_rules! ai_analyze {
1566 ($code:expr, $context:expr) => {{
1567 let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1568 analyzer.analyze_model_code($code, $context).await
1569 }};
1570}
1571
1572#[path = "ai_code_analyzer_tests.rs"]
1573mod ai_code_analyzer_tests;
1574
1575#[cfg(test)]
1576mod tests {
1577 use super::*;
1578
1579 #[tokio::test]
1580 async fn test_ai_code_analyzer_creation() {
1581 let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1582 assert!(analyzer.config.enable_deep_analysis);
1583 }
1584
1585 #[tokio::test]
1586 async fn test_pattern_detection() {
1587 let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1588
1589 let code = r#"
1590 import torch
1591
1592 def train_step(model, data):
1593 torch.cuda.empty_cache() # Should trigger anti-pattern
1594 grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # Good pattern
1595 return grad_norm
1596 "#;
1597
1598 let context = ModelContext {
1599 model_type: ModelType::Production,
1600 model_size: 1_000_000,
1601 framework: "PyTorch".to_string(),
1602 target_hardware: "CUDA".to_string(),
1603 training_stage: TrainingStage::Training,
1604 };
1605
1606 let result = analyzer
1607 .analyze_model_code(code, context)
1608 .await
1609 .expect("async operation failed");
1610 assert!(!result.detected_patterns.is_empty());
1611 }
1612
1613 #[tokio::test]
1614 async fn test_security_vulnerability_detection() {
1615 let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1616
1617 let code = r#"
1618 import pickle
1619
1620 def load_model(path):
1621 with open(path, 'rb') as f:
1622 model = pickle.load(f) # Should trigger security warning
1623 return model
1624 "#;
1625
1626 let context = ModelContext {
1627 model_type: ModelType::Production,
1628 model_size: 1_000_000,
1629 framework: "PyTorch".to_string(),
1630 target_hardware: "CUDA".to_string(),
1631 training_stage: TrainingStage::Inference,
1632 };
1633
1634 let result = analyzer
1635 .analyze_model_code(code, context)
1636 .await
1637 .expect("async operation failed");
1638 assert!(!result.security_issues.is_empty());
1639 assert_eq!(
1640 result.security_issues[0].vulnerability_type,
1641 VulnerabilityType::CodeExecution
1642 );
1643 }
1644
1645 #[tokio::test]
1646 async fn test_tensor_operation_analysis() {
1647 let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1648
1649 let operations = vec![
1650 TensorOperation {
1651 name: "matmul1".to_string(),
1652 op_type: OperationType::MatMul,
1653 inputs: vec!["a".to_string(), "b".to_string()],
1654 outputs: vec!["c".to_string()],
1655 parameters: HashMap::new(),
1656 output_size_bytes: 1024,
1657 is_inplace: false,
1658 },
1659 TensorOperation {
1660 name: "add1".to_string(),
1661 op_type: OperationType::Add,
1662 inputs: vec!["c".to_string(), "bias".to_string()],
1663 outputs: vec!["d".to_string()],
1664 parameters: HashMap::new(),
1665 output_size_bytes: 1024,
1666 is_inplace: false,
1667 },
1668 ];
1669
1670 let report = analyzer
1671 .analyze_tensor_operations(&operations)
1672 .await
1673 .expect("tensor operation failed");
1674 assert!(!report.fusion_opportunities.is_empty());
1675 assert_eq!(report.fusion_opportunities[0].fusion_type, FusionType::GEMM);
1676 }
1677
1678 #[tokio::test]
1679 async fn test_performance_metrics() {
1680 let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1681
1682 let code = "print('hello')";
1684 let context = ModelContext {
1685 model_type: ModelType::Development,
1686 model_size: 1000,
1687 framework: "PyTorch".to_string(),
1688 target_hardware: "CPU".to_string(),
1689 training_stage: TrainingStage::Development,
1690 };
1691
1692 analyzer
1693 .analyze_model_code(code, context.clone())
1694 .await
1695 .expect("async operation failed");
1696 analyzer
1697 .analyze_model_code(code, context)
1698 .await
1699 .expect("async operation failed"); let metrics = analyzer.get_performance_metrics();
1702 assert_eq!(metrics.total_analyses, 2);
1703 assert!(metrics.cache_hit_rate > 0.0);
1704 }
1705
1706 #[tokio::test]
1707 async fn test_debugging_assistance() {
1708 let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1709
1710 let error_context = ErrorContext {
1711 error_type: "OutOfMemoryError".to_string(),
1712 error_message: "CUDA out of memory".to_string(),
1713 stack_trace: None,
1714 system_info: SystemInfo {
1715 gpu_memory_total: 8_000_000_000,
1716 gpu_memory_used: 7_500_000_000,
1717 cpu_count: 8,
1718 ram_total: 32_000_000_000,
1719 ram_used: 16_000_000_000,
1720 },
1721 model_info: None,
1722 };
1723
1724 let assistance = analyzer
1725 .automated_debugging_assistance(&error_context)
1726 .await
1727 .expect("async operation failed");
1728 assert!(!assistance.probable_causes.is_empty());
1729 assert!(!assistance.suggested_fixes.is_empty());
1730 assert!(assistance.confidence_score > 0.0);
1731 }
1732
1733 fn make_context(model_size: u64) -> ModelContext {
1734 ModelContext {
1735 model_type: ModelType::Training,
1736 model_size,
1737 framework: "PyTorch".to_string(),
1738 target_hardware: "CUDA".to_string(),
1739 training_stage: TrainingStage::Training,
1740 }
1741 }
1742
1743 #[tokio::test]
1748 async fn test_identified_issue_code_location_is_real_not_none() {
1749 let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1750 let code = "def f(x):\n return log(softmax(x))\n";
1751 let context = make_context(1_000);
1752
1753 let issues = analyzer
1754 .perform_deep_analysis(code, &context)
1755 .await
1756 .expect("deep analysis should succeed");
1757 let issue = issues
1758 .iter()
1759 .find(|i| matches!(i.issue_type, IssueType::NumericalStability))
1760 .expect("the log-softmax rule should have fired");
1761
1762 let location = issue
1763 .code_location
1764 .as_ref()
1765 .expect("code_location must be Some (a real position), not the old hardcoded None");
1766 assert_eq!(location.line, 2);
1769 assert_eq!(location.column, " return log(".len() as u32 + 1);
1770 }
1771
1772 #[tokio::test]
1778 async fn test_performance_predictions_are_honest_with_no_signals() {
1779 let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1780 let code = "def f(x):\n return x + 1\n";
1781 let context = make_context(1_000); let predictions = analyzer
1784 .predict_performance_characteristics(code, &context)
1785 .await
1786 .expect("prediction should succeed");
1787
1788 assert!(
1789 predictions.predicted_bottlenecks.is_empty(),
1790 "must not report bottlenecks when no real signal was detected: {:?}",
1791 predictions.predicted_bottlenecks
1792 );
1793 assert_eq!(
1794 predictions.confidence_score, 0.0,
1795 "must not be the old fixed 0.75 when nothing was actually detected"
1796 );
1797 }
1798
1799 #[tokio::test]
1804 async fn test_performance_predictions_reflect_real_code_signals() {
1805 let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1806 let attention_code = "out = attention(q, k, v, matmul_impl=True)";
1807 let plain_context = make_context(1_000);
1808
1809 let with_attention = analyzer
1810 .predict_performance_characteristics(attention_code, &plain_context)
1811 .await
1812 .expect("prediction should succeed");
1813 assert_eq!(with_attention.predicted_bottlenecks.len(), 1);
1814 assert!(with_attention.predicted_bottlenecks[0].contains("Attention"));
1815 assert_eq!(with_attention.confidence_score, 0.6);
1816
1817 let large_model_context = make_context(2_000_000_000);
1818 let with_both = analyzer
1819 .predict_performance_characteristics(attention_code, &large_model_context)
1820 .await
1821 .expect("prediction should succeed");
1822 assert_eq!(
1823 with_both.predicted_bottlenecks.len(),
1824 2,
1825 "a large model plus an attention signal must report both real bottlenecks"
1826 );
1827 assert_eq!(with_both.confidence_score, 0.75);
1828 }
1829
1830 #[tokio::test]
1836 async fn test_analyze_model_code_has_no_artificial_latency() {
1837 let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1838 let code = "out = attention(q, k, v, matmul_impl=True)\nlog(softmax(x))";
1839 let context = make_context(2_000_000_000);
1840
1841 let start = std::time::Instant::now();
1842 analyzer
1843 .analyze_model_code(code, context)
1844 .await
1845 .expect("analysis should succeed");
1846 let elapsed = start.elapsed();
1847
1848 assert!(
1853 elapsed < std::time::Duration::from_millis(50),
1854 "analysis took {elapsed:?}; the old implementation's artificial sleeps must be gone"
1855 );
1856 }
1857}