Files
rustytorch/crates/training/rtx-evolution/src/knowledge.rs
T
2026-03-04 00:08:42 +00:00

537 lines
18 KiB
Rust

//! Knowledge Graph
//!
//! Meta-learning storage, pattern relationships, and success prediction
use crate::{ExecutionResult, ProposalSpec, Result};
use petgraph::Graph;
use petgraph::graph::{NodeIndex, UnGraph};
use std::collections::HashMap;
use uuid::Uuid;
/// Knowledge graph for meta-learning
pub struct KnowledgeGraph {
graph: UnGraph<Pattern, Relationship>,
pattern_index: HashMap<String, NodeIndex>,
success_history: Vec<SuccessRecord>,
ready: bool,
}
/// Pattern stored in knowledge graph
#[derive(Debug, Clone)]
pub struct Pattern {
pub id: String,
pub name: String,
pub pattern_type: PatternType,
pub success_rate: f64,
pub average_improvement: f64,
pub usage_count: u64,
pub confidence: f64,
}
/// Type of pattern in the knowledge graph
#[derive(Debug, Clone, PartialEq)]
pub enum PatternType {
OptimizationTechnique, // General optimization approach
ParameterSetting, // Specific parameter configurations
AlgorithmChoice, // Algorithm selection patterns
ResourceUtilization, // Resource usage patterns
PerformanceCorrelation, // Performance correlation patterns
}
/// Relationship between patterns
#[derive(Debug, Clone)]
pub struct Relationship {
pub kind: RelationshipKind,
pub strength: f64,
pub evidence_count: u64,
}
/// Types of relationships between patterns
#[derive(Debug, Clone, PartialEq)]
pub enum RelationshipKind {
Synergy, // Patterns work well together
Conflict, // Patterns interfere with each other
Prerequisite, // One pattern requires another
Alternative, // Patterns are alternatives to each other
Correlation, // Patterns often co-occur
}
/// Record of successful optimizations
#[derive(Debug, Clone)]
struct SuccessRecord {
proposal_id: Uuid,
patterns_used: Vec<String>,
performance_improvement: f64,
timestamp: std::time::SystemTime,
}
impl KnowledgeGraph {
/// Create new knowledge graph
pub fn new() -> Self {
let mut kg = Self {
graph: Graph::new_undirected(),
pattern_index: HashMap::new(),
success_history: Vec::new(),
ready: true,
};
// Initialize with some basic patterns
kg.initialize_base_patterns();
kg
}
/// Check if knowledge graph is ready
pub fn is_ready(&self) -> bool {
self.ready
}
/// Update knowledge graph from execution results
pub async fn update_from_execution(
&mut self,
proposal: &ProposalSpec,
result: &ExecutionResult,
) -> Result<()> {
if result.safety_check_passed && result.performance_delta > 0.0 {
// Extract patterns from successful proposal
let patterns = self.extract_patterns_from_proposal(proposal).await?;
// Update pattern success rates
for pattern_name in &patterns {
self.update_pattern_success(pattern_name, result.performance_delta)
.await?;
}
// Learn relationships between patterns
self.learn_pattern_relationships(&patterns, result.performance_delta)
.await?;
// Record successful execution
self.record_success(proposal, result, patterns).await?;
} else {
// Learn from failures too
let patterns = self.extract_patterns_from_proposal(proposal).await?;
for pattern_name in &patterns {
self.update_pattern_failure(pattern_name).await?;
}
}
Ok(())
}
/// Initialize knowledge graph with base patterns
fn initialize_base_patterns(&mut self) {
let base_patterns = vec![
Pattern {
id: "kernel_tile_optimization".to_string(),
name: "Kernel Tile Size Optimization".to_string(),
pattern_type: PatternType::ParameterSetting,
success_rate: 0.7,
average_improvement: 0.12,
usage_count: 0,
confidence: 0.8,
},
Pattern {
id: "memory_layout_optimization".to_string(),
name: "Memory Layout Optimization".to_string(),
pattern_type: PatternType::OptimizationTechnique,
success_rate: 0.6,
average_improvement: 0.08,
usage_count: 0,
confidence: 0.7,
},
Pattern {
id: "compiler_aggressive_opts".to_string(),
name: "Aggressive Compiler Optimizations".to_string(),
pattern_type: PatternType::AlgorithmChoice,
success_rate: 0.5,
average_improvement: 0.05,
usage_count: 0,
confidence: 0.6,
},
Pattern {
id: "algorithm_switching".to_string(),
name: "Algorithm Switching".to_string(),
pattern_type: PatternType::AlgorithmChoice,
success_rate: 0.8,
average_improvement: 0.15,
usage_count: 0,
confidence: 0.85,
},
// AI-powered optimization patterns
Pattern {
id: "ai_powered_optimization".to_string(),
name: "AI-Powered Code Optimization".to_string(),
pattern_type: PatternType::OptimizationTechnique,
success_rate: 0.85,
average_improvement: 0.22,
usage_count: 0,
confidence: 0.90,
},
// RTX 5090 hardware-specific patterns
Pattern {
id: "hardware_specific_optimization".to_string(),
name: "Hardware-Specific RTX 5090 Optimization".to_string(),
pattern_type: PatternType::OptimizationTechnique,
success_rate: 0.90,
average_improvement: 0.28,
usage_count: 0,
confidence: 0.95,
},
// CUDA 13.0 feature enablement patterns
Pattern {
id: "cuda_feature_enablement".to_string(),
name: "CUDA 13.0 Feature Enablement".to_string(),
pattern_type: PatternType::OptimizationTechnique,
success_rate: 0.82,
average_improvement: 0.20,
usage_count: 0,
confidence: 0.88,
},
];
for pattern in base_patterns {
let node_index = self.graph.add_node(pattern.clone());
self.pattern_index.insert(pattern.id, node_index);
}
// Add some initial relationships
self.add_initial_relationships();
}
/// Add initial known relationships between patterns
fn add_initial_relationships(&mut self) {
// Kernel tile optimization and memory layout often work well together (synergy)
if let (Some(&tile_idx), Some(&memory_idx)) = (
self.pattern_index.get("kernel_tile_optimization"),
self.pattern_index.get("memory_layout_optimization"),
) {
self.graph.add_edge(
tile_idx,
memory_idx,
Relationship {
kind: RelationshipKind::Synergy,
strength: 0.7,
evidence_count: 5,
},
);
}
// Aggressive compiler opts and algorithm switching can conflict
if let (Some(&compiler_idx), Some(&algorithm_idx)) = (
self.pattern_index.get("compiler_aggressive_opts"),
self.pattern_index.get("algorithm_switching"),
) {
self.graph.add_edge(
compiler_idx,
algorithm_idx,
Relationship {
kind: RelationshipKind::Conflict,
strength: 0.4,
evidence_count: 3,
},
);
}
}
/// Extract patterns from a proposal
async fn extract_patterns_from_proposal(&self, proposal: &ProposalSpec) -> Result<Vec<String>> {
let mut patterns = Vec::new();
for change in &proposal.changes {
match change {
crate::Change::KernelParameter { .. } => {
patterns.push("kernel_tile_optimization".to_string());
}
crate::Change::MemoryLayout { .. } => {
patterns.push("memory_layout_optimization".to_string());
}
crate::Change::CompilerFlag { .. } => {
patterns.push("compiler_aggressive_opts".to_string());
}
crate::Change::AlgorithmSwitch { .. } => {
patterns.push("algorithm_switching".to_string());
}
crate::Change::CodeOptimization {
optimization_type, ..
} => {
patterns.push(format!("ai_code_optimization_{}", optimization_type));
patterns.push("ai_powered_optimization".to_string());
}
crate::Change::Rtx5090Optimization {
optimization_type, ..
} => {
patterns.push(format!("rtx5090_optimization_{}", optimization_type));
patterns.push("hardware_specific_optimization".to_string());
}
crate::Change::Cuda13Feature { feature_name, .. } => {
patterns.push(format!("cuda13_feature_{}", feature_name));
patterns.push("cuda_feature_enablement".to_string());
}
}
}
Ok(patterns)
}
/// Update pattern success statistics
async fn update_pattern_success(&mut self, pattern_name: &str, improvement: f64) -> Result<()> {
if let Some(&node_idx) = self.pattern_index.get(pattern_name) {
if let Some(pattern) = self.graph.node_weight_mut(node_idx) {
// Update success rate using exponential moving average
let alpha = 0.1; // Learning rate
pattern.success_rate = (1.0 - alpha) * pattern.success_rate + alpha * 1.0;
// Update average improvement
pattern.average_improvement =
(1.0 - alpha) * pattern.average_improvement + alpha * improvement;
// Increment usage count
pattern.usage_count += 1;
// Update confidence based on usage count
pattern.confidence =
(pattern.usage_count as f64 / (pattern.usage_count as f64 + 10.0)).min(0.95);
}
}
Ok(())
}
/// Update pattern failure statistics
async fn update_pattern_failure(&mut self, pattern_name: &str) -> Result<()> {
if let Some(&node_idx) = self.pattern_index.get(pattern_name) {
if let Some(pattern) = self.graph.node_weight_mut(node_idx) {
// Update success rate with failure
let alpha = 0.1;
pattern.success_rate = (1.0 - alpha) * pattern.success_rate + alpha * 0.0;
// Increment usage count
pattern.usage_count += 1;
// Adjust confidence down slightly for failures
pattern.confidence *= 0.98;
}
}
Ok(())
}
/// Learn relationships between patterns based on co-occurrence
async fn learn_pattern_relationships(
&mut self,
patterns: &[String],
improvement: f64,
) -> Result<()> {
// Learn relationships between all pairs of patterns used together
for i in 0..patterns.len() {
for j in (i + 1)..patterns.len() {
let pattern_a = &patterns[i];
let pattern_b = &patterns[j];
if let (Some(&idx_a), Some(&idx_b)) = (
self.pattern_index.get(pattern_a),
self.pattern_index.get(pattern_b),
) {
self.update_relationship(idx_a, idx_b, improvement).await?;
}
}
}
Ok(())
}
/// Update relationship strength between two patterns
async fn update_relationship(
&mut self,
idx_a: NodeIndex,
idx_b: NodeIndex,
improvement: f64,
) -> Result<()> {
// Find existing edge or create new one
let edge_idx = self.graph.find_edge(idx_a, idx_b);
match edge_idx {
Some(edge_idx) => {
// Update existing relationship
if let Some(relationship) = self.graph.edge_weight_mut(edge_idx) {
relationship.evidence_count += 1;
// Determine relationship kind based on improvement
let relationship_kind = if improvement > 0.1 {
RelationshipKind::Synergy
} else if improvement < 0.0 {
RelationshipKind::Conflict
} else {
RelationshipKind::Correlation
};
// Update strength using exponential moving average
let alpha = 0.2;
let new_strength = improvement.abs();
relationship.strength =
(1.0 - alpha) * relationship.strength + alpha * new_strength;
// Update relationship kind if evidence is strong
if relationship.evidence_count >= 3 {
relationship.kind = relationship_kind;
}
}
}
None => {
// Create new relationship
let relationship_kind = if improvement > 0.1 {
RelationshipKind::Synergy
} else if improvement < 0.0 {
RelationshipKind::Conflict
} else {
RelationshipKind::Correlation
};
let new_relationship = Relationship {
kind: relationship_kind,
strength: improvement.abs(),
evidence_count: 1,
};
self.graph.add_edge(idx_a, idx_b, new_relationship);
}
}
Ok(())
}
/// Record successful execution in history
async fn record_success(
&mut self,
proposal: &ProposalSpec,
result: &ExecutionResult,
patterns: Vec<String>,
) -> Result<()> {
let success_record = SuccessRecord {
proposal_id: proposal.id,
patterns_used: patterns,
performance_improvement: result.performance_delta,
timestamp: std::time::SystemTime::now(),
};
// In a real implementation, this would be mutable
// For now, we simulate recording the success
tracing::info!(
"Recording success: proposal {} with {:.2}% improvement using {} patterns",
proposal.id,
result.performance_delta * 100.0,
success_record.patterns_used.len()
);
Ok(())
}
/// Get recommendations for optimization patterns
pub async fn get_recommendations(
&self,
context: &OptimizationContext,
) -> Result<Vec<PatternRecommendation>> {
let mut recommendations = Vec::new();
// Analyze all patterns and rank by potential success
for (pattern_id, &node_idx) in &self.pattern_index {
if let Some(pattern) = self.graph.node_weight(node_idx) {
let score = self
.calculate_recommendation_score(pattern, context)
.await?;
if score > 0.3 {
// Threshold for recommendations
recommendations.push(PatternRecommendation {
pattern_id: pattern_id.clone(),
pattern_name: pattern.name.clone(),
confidence: pattern.confidence,
expected_improvement: pattern.average_improvement,
score,
rationale: self.generate_rationale(pattern, context).await?,
});
}
}
}
// Sort by score (highest first)
recommendations.sort_by(|a, b| b.score.total_cmp(&a.score));
// Limit to top 5 recommendations
recommendations.truncate(5);
Ok(recommendations)
}
/// Calculate recommendation score for a pattern given context
async fn calculate_recommendation_score(
&self,
pattern: &Pattern,
context: &OptimizationContext,
) -> Result<f64> {
let base_score = pattern.success_rate * pattern.confidence;
// Adjust based on context
let context_multiplier = match (&pattern.pattern_type, &context.optimization_target) {
(PatternType::ParameterSetting, OptimizationTarget::Performance) => 1.2,
(PatternType::AlgorithmChoice, OptimizationTarget::Performance) => 1.5,
(PatternType::OptimizationTechnique, OptimizationTarget::Memory) => 1.3,
(PatternType::ResourceUtilization, OptimizationTarget::Power) => 1.4,
_ => 1.0,
};
let final_score = base_score * context_multiplier;
Ok(final_score.max(0.0).min(1.0))
}
/// Generate rationale for why a pattern is recommended
async fn generate_rationale(
&self,
pattern: &Pattern,
context: &OptimizationContext,
) -> Result<String> {
let success_percentage = (pattern.success_rate * 100.0) as i32;
let avg_improvement_percentage = (pattern.average_improvement * 100.0) as i32;
let rationale = format!(
"Pattern '{}' has {:.0}% success rate with average {:.0}% improvement. Used successfully {} times. Confidence: {:.0}%",
pattern.name,
success_percentage,
avg_improvement_percentage,
pattern.usage_count,
pattern.confidence * 100.0
);
Ok(rationale)
}
}
/// Context for optimization recommendations
#[derive(Debug, Clone)]
pub struct OptimizationContext {
pub optimization_target: OptimizationTarget,
pub current_metrics: HashMap<String, f64>,
pub constraints: Vec<String>,
}
/// Target for optimization
#[derive(Debug, Clone, PartialEq)]
pub enum OptimizationTarget {
Performance,
Memory,
Power,
Latency,
}
/// Pattern recommendation from knowledge graph
#[derive(Debug, Clone)]
pub struct PatternRecommendation {
pub pattern_id: String,
pub pattern_name: String,
pub confidence: f64,
pub expected_improvement: f64,
pub score: f64,
pub rationale: String,
}