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

327 lines
10 KiB
Rust

//! Proposal generation and validation for autonomous optimization.
use crate::error::AutoResult;
use rtx_runtime::Runtime;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use uuid::Uuid;
/// Types of optimization proposals that can be generated.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ProposalType {
/// Data layout optimization proposals
DataLayout,
/// Memory access coalescing proposals
MemoryCoalescing,
/// Cache optimization proposals
CacheOptimization,
/// Data parallelization proposals
DataParallel,
/// Pipeline parallelization proposals
PipelineParallel,
/// Tensor parallelization proposals
TensorParallel,
/// Communication optimization proposals
CommunicationOptimization,
/// Quantization proposals
Quantization,
/// Mixed precision proposals
MixedPrecision,
/// Calibration proposals
Calibration,
/// Adaptive quantization proposals
AdaptiveQuantization,
/// Kernel fusion proposals
KernelFusion,
/// Kernel optimization proposals
KernelOptimization,
/// Memory access pattern optimization
MemoryAccess,
/// GPU occupancy optimization
Occupancy,
/// Code generation optimization
CodeGeneration,
}
/// Status of an optimization proposal.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProposalStatus {
/// Proposal has been created but not yet evaluated
Pending,
/// Proposal is currently being evaluated
Evaluating,
/// Proposal has been evaluated and approved for application
Approved,
/// Proposal has been rejected
Rejected,
/// Proposal has been applied successfully
Applied,
/// Proposal application failed
Failed,
/// Proposal has been rolled back
RolledBack,
}
/// An optimization proposal generated by autonomous agents.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Proposal {
id: String,
proposal_type: ProposalType,
description: String,
expected_performance_gain: f32,
status: ProposalStatus,
created_at: u64,
metadata: HashMap<String, String>,
}
impl Proposal {
/// Create a new optimization proposal.
pub fn new(
proposal_type: ProposalType,
description: String,
expected_performance_gain: f32,
) -> Self {
Self {
id: Uuid::new_v4().to_string(),
proposal_type,
description,
expected_performance_gain,
status: ProposalStatus::Pending,
created_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs(),
metadata: HashMap::new(),
}
}
/// Get the proposal ID.
pub fn id(&self) -> &str {
&self.id
}
/// Get the proposal type.
pub fn proposal_type(&self) -> ProposalType {
self.proposal_type
}
/// Get the proposal description.
pub fn description(&self) -> &str {
&self.description
}
/// Get the expected performance gain.
pub fn expected_performance_gain(&self) -> f32 {
self.expected_performance_gain
}
/// Get the current status.
pub fn status(&self) -> ProposalStatus {
self.status
}
/// Get the creation timestamp.
pub fn created_at(&self) -> u64 {
self.created_at
}
/// Set the proposal status.
pub fn set_status(&mut self, status: ProposalStatus) {
self.status = status;
}
/// Add metadata to the proposal.
pub fn add_metadata<K: Into<String>, V: Into<String>>(&mut self, key: K, value: V) {
self.metadata.insert(key.into(), value.into());
}
/// Get metadata value by key.
pub fn get_metadata(&self, key: &str) -> Option<&String> {
self.metadata.get(key)
}
}
/// Validates and scores optimization proposals.
pub struct ProposalValidator {
runtime: Arc<Runtime>,
}
impl ProposalValidator {
/// Create a new proposal validator.
pub fn new(runtime: Arc<Runtime>) -> AutoResult<Self> {
Ok(Self { runtime })
}
/// Validate that the expected performance gain is reasonable.
pub async fn validate_performance_gain(&self, proposal: &Proposal) -> AutoResult<bool> {
// Realistic performance gains depend on proposal type
let max_reasonable_gain = match proposal.proposal_type {
ProposalType::DataLayout => 5.0,
ProposalType::MemoryCoalescing => 3.0,
ProposalType::CacheOptimization => 2.5,
ProposalType::DataParallel => 8.0,
ProposalType::PipelineParallel => 6.0,
ProposalType::TensorParallel => 4.0,
ProposalType::CommunicationOptimization => 3.0,
ProposalType::Quantization => 5.0,
ProposalType::MixedPrecision => 3.5,
ProposalType::Calibration => 2.0,
ProposalType::AdaptiveQuantization => 4.0,
ProposalType::KernelFusion => 8.0,
ProposalType::KernelOptimization => 10.0,
ProposalType::MemoryAccess => 4.0,
ProposalType::Occupancy => 3.0,
ProposalType::CodeGeneration => 5.0,
};
Ok(proposal.expected_performance_gain > 0.0
&& proposal.expected_performance_gain <= max_reasonable_gain)
}
/// Validate the feasibility of applying the proposal.
pub async fn validate_feasibility(&self, proposal: &Proposal) -> AutoResult<bool> {
// Basic feasibility checks
if proposal.description.is_empty() {
return Ok(false);
}
if proposal.expected_performance_gain <= 0.0 {
return Ok(false);
}
// More sophisticated feasibility analysis would go here
// For now, we consider all proposals with reasonable descriptions feasible
Ok(true)
}
/// Score a proposal based on multiple criteria.
pub async fn score_proposal(&self, proposal: &Proposal) -> AutoResult<f32> {
let mut score = 0.0;
// Performance gain contribution (0.0 - 0.4)
let normalized_gain = (proposal.expected_performance_gain.min(10.0)) / 10.0;
score += normalized_gain * 0.4;
// Proposal type priority (0.0 - 0.3)
let type_priority = match proposal.proposal_type {
ProposalType::KernelFusion | ProposalType::KernelOptimization => 0.3,
ProposalType::DataParallel | ProposalType::PipelineParallel => 0.25,
ProposalType::Quantization | ProposalType::MixedPrecision => 0.2,
ProposalType::MemoryCoalescing | ProposalType::CacheOptimization => 0.15,
_ => 0.1,
};
score += type_priority;
// Description quality (0.0 - 0.2)
let description_score = if proposal.description.len() > 50 {
0.2
} else {
0.1
};
score += description_score;
// Metadata completeness (0.0 - 0.1)
let metadata_score = if proposal.metadata.is_empty() {
0.0
} else {
0.1
};
score += metadata_score;
Ok(score.min(1.0))
}
/// Rank a list of proposals by their scores.
pub async fn rank_proposals(&self, proposals: &[Proposal]) -> AutoResult<Vec<(Proposal, f32)>> {
let mut scored_proposals = Vec::new();
for proposal in proposals {
let score = self.score_proposal(proposal).await?;
scored_proposals.push((proposal.clone(), score));
}
// Sort by score in descending order
scored_proposals.sort_by(|a, b| b.1.total_cmp(&a.1));
Ok(scored_proposals)
}
/// Detect conflicts between proposals.
pub async fn detect_conflicts(
&self,
proposals: &[Proposal],
) -> AutoResult<Vec<(usize, usize)>> {
let mut conflicts = Vec::new();
for (i, proposal1) in proposals.iter().enumerate() {
for (j, proposal2) in proposals.iter().enumerate().skip(i + 1) {
if self.proposals_conflict(proposal1, proposal2) {
conflicts.push((i, j));
}
}
}
Ok(conflicts)
}
/// Check if two proposals conflict with each other.
fn proposals_conflict(&self, proposal1: &Proposal, proposal2: &Proposal) -> bool {
match (proposal1.proposal_type, proposal2.proposal_type) {
// Data layout optimizations may conflict
(ProposalType::DataLayout, ProposalType::DataLayout) => {
proposal1.description != proposal2.description
}
// Memory access patterns may conflict
(ProposalType::MemoryAccess, ProposalType::MemoryAccess) => true,
(ProposalType::MemoryCoalescing, ProposalType::MemoryAccess) => true,
(ProposalType::MemoryAccess, ProposalType::MemoryCoalescing) => true,
// Parallelization strategies may conflict
(ProposalType::DataParallel, ProposalType::PipelineParallel) => true,
(ProposalType::PipelineParallel, ProposalType::DataParallel) => true,
// Quantization strategies may conflict
(ProposalType::Quantization, ProposalType::MixedPrecision) => true,
(ProposalType::MixedPrecision, ProposalType::Quantization) => true,
_ => false,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_proposal_creation() {
let proposal = Proposal::new(ProposalType::DataLayout, "Test proposal".to_string(), 2.0);
assert_eq!(proposal.proposal_type, ProposalType::DataLayout);
assert_eq!(proposal.expected_performance_gain, 2.0);
assert_eq!(proposal.status, ProposalStatus::Pending);
}
#[test]
fn test_proposal_status_update() {
let mut proposal =
Proposal::new(ProposalType::KernelFusion, "Fuse kernels".to_string(), 3.0);
proposal.set_status(ProposalStatus::Approved);
assert_eq!(proposal.status, ProposalStatus::Approved);
}
#[test]
fn test_proposal_metadata() {
let mut proposal = Proposal::new(
ProposalType::Quantization,
"Apply quantization".to_string(),
2.5,
);
proposal.add_metadata("precision", "int8");
assert_eq!(
proposal.get_metadata("precision"),
Some(&"int8".to_string())
);
}
}