Initial commit

crates/training/rtx-nas/src/hardware/latency.rs

@@ -0,0 +1,491 @@
+//! Latency prediction for hardware-aware NAS
+//!
+//! Provides latency estimation and prediction for architectures
+//! on different hardware devices.
+
+use super::device::DeviceProfile;
+use crate::error::{NASError, Result};
+use crate::search_space::{Architecture, Cell, OperationType};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+/// Latency predictor trait
+pub trait LatencyPredictor: Send + Sync {
+    /// Predict latency for an architecture on a device
+    fn predict(&self, arch: &Architecture, device: &DeviceProfile) -> Result<f32>;
+
+    /// Calibrate the predictor with measured samples
+    fn calibrate(&mut self, samples: &[(Architecture, f32)]) -> Result<()>;
+
+    /// Get prediction confidence (0.0 to 1.0)
+    fn confidence(&self) -> f32;
+}
+
+/// Operation-level latency lookup table
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OperationLatency {
+    /// Base latency in microseconds
+    pub base_latency_us: f32,
+    /// Latency per channel
+    pub per_channel_us: f32,
+    /// Latency per spatial element (H*W)
+    pub per_spatial_us: f32,
+    /// Memory access pattern factor
+    pub memory_factor: f32,
+}
+
+impl OperationLatency {
+    /// Create a new operation latency entry
+    pub fn new(base: f32, per_channel: f32, per_spatial: f32) -> Self {
+        Self {
+            base_latency_us: base,
+            per_channel_us: per_channel,
+            per_spatial_us: per_spatial,
+            memory_factor: 1.0,
+        }
+    }
+
+    /// Estimate latency for given configuration
+    pub fn estimate(&self, channels: usize, height: usize, width: usize) -> f32 {
+        let spatial = (height * width) as f32;
+        self.base_latency_us
+            + self.per_channel_us * channels as f32
+            + self.per_spatial_us * spatial * self.memory_factor
+    }
+}
+
+/// Lookup table based latency predictor
+#[derive(Debug, Clone)]
+pub struct LookupTablePredictor {
+    /// Latency tables per device type per operation
+    tables: HashMap<String, HashMap<OperationType, OperationLatency>>,
+    /// Default latencies when no specific table exists
+    default_latencies: HashMap<OperationType, OperationLatency>,
+    /// Confidence in predictions
+    confidence: f32,
+}
+
+impl LookupTablePredictor {
+    /// Create a new lookup table predictor with default values
+    pub fn new() -> Self {
+        let mut default_latencies = HashMap::new();
+
+        // Default latencies based on typical GPU performance
+        default_latencies.insert(
+            OperationType::Identity,
+            OperationLatency::new(0.1, 0.0, 0.0001),
+        );
+        default_latencies.insert(OperationType::Zero, OperationLatency::new(0.05, 0.0, 0.0));
+        default_latencies.insert(
+            OperationType::Conv3x3,
+            OperationLatency::new(10.0, 0.5, 0.01),
+        );
+        default_latencies.insert(
+            OperationType::Conv5x5,
+            OperationLatency::new(25.0, 1.2, 0.025),
+        );
+        default_latencies.insert(
+            OperationType::SepConv3x3,
+            OperationLatency::new(8.0, 0.3, 0.008),
+        );
+        default_latencies.insert(
+            OperationType::SepConv5x5,
+            OperationLatency::new(15.0, 0.6, 0.015),
+        );
+        default_latencies.insert(
+            OperationType::DilConv3x3,
+            OperationLatency::new(12.0, 0.6, 0.012),
+        );
+        default_latencies.insert(
+            OperationType::MaxPool3x3,
+            OperationLatency::new(2.0, 0.1, 0.002),
+        );
+        default_latencies.insert(
+            OperationType::AvgPool3x3,
+            OperationLatency::new(2.5, 0.1, 0.0025),
+        );
+
+        Self {
+            tables: HashMap::new(),
+            default_latencies,
+            confidence: 0.5, // Medium confidence for uncalibrated predictor
+        }
+    }
+
+    /// Add device-specific latency table
+    pub fn add_device_table(
+        &mut self,
+        device_name: impl Into<String>,
+        table: HashMap<OperationType, OperationLatency>,
+    ) {
+        self.tables.insert(device_name.into(), table);
+    }
+
+    /// Get latency for an operation on a device
+    fn get_op_latency(&self, op: OperationType, device: &DeviceProfile) -> OperationLatency {
+        // Try device-specific table first
+        if let Some(device_table) = self.tables.get(&device.name) {
+            if let Some(latency) = device_table.get(&op) {
+                return latency.clone();
+            }
+        }
+
+        // Fall back to default
+        self.default_latencies
+            .get(&op)
+            .cloned()
+            .unwrap_or_else(|| OperationLatency::new(10.0, 0.5, 0.01))
+    }
+
+    /// Estimate latency for a cell
+    fn estimate_cell_latency(
+        &self,
+        cell: &Cell,
+        device: &DeviceProfile,
+        channels: usize,
+        height: usize,
+        width: usize,
+    ) -> f32 {
+        let mut total_latency = 0.0;
+
+        for edge in cell.edges() {
+            if let Some(op_type) = cell.get_operation(&edge) {
+                let op_latency = self.get_op_latency(op_type, device);
+                total_latency += op_latency.estimate(channels, height, width);
+            }
+        }
+
+        // Apply device-specific scaling
+        let device_factor = match device.device_type {
+            super::device::DeviceType::CPU => 10.0,
+            super::device::DeviceType::Mobile => 5.0,
+            super::device::DeviceType::Edge => 20.0,
+            _ => 1.0, // GPUs as baseline
+        };
+
+        total_latency * device_factor / device.peak_tflops_fp32.max(0.1)
+    }
+}
+
+impl Default for LookupTablePredictor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl LatencyPredictor for LookupTablePredictor {
+    fn predict(&self, arch: &Architecture, device: &DeviceProfile) -> Result<f32> {
+        let mut total_latency_ms = 0.0;
+
+        // Default dimensions if not specified in architecture
+        let channels = arch.channels;
+        let height = 224; // Default ImageNet size
+        let width = 224;
+
+        for cell in &arch.cells {
+            let cell_latency_us = self.estimate_cell_latency(cell, device, channels, height, width);
+            total_latency_ms += cell_latency_us / 1000.0; // Convert to ms
+        }
+
+        Ok(total_latency_ms)
+    }
+
+    fn calibrate(&mut self, samples: &[(Architecture, f32)]) -> Result<()> {
+        if samples.is_empty() {
+            return Err(NASError::LatencyError(
+                "No calibration samples provided".into(),
+            ));
+        }
+
+        // Simple calibration: adjust confidence based on sample count
+        let sample_count = samples.len();
+        self.confidence = (sample_count as f32 / 100.0).min(0.95);
+
+        // In a real implementation, we would fit the latency tables to the samples
+        Ok(())
+    }
+
+    fn confidence(&self) -> f32 {
+        self.confidence
+    }
+}
+
+/// Linear regression based latency predictor
+#[derive(Debug, Clone)]
+pub struct RegressionPredictor {
+    /// Learned coefficients for each operation type
+    coefficients: HashMap<OperationType, Vec<f32>>,
+    /// Bias term
+    bias: f32,
+    /// Prediction confidence
+    confidence: f32,
+    /// Number of features per operation
+    num_features: usize,
+}
+
+impl RegressionPredictor {
+    /// Create a new regression predictor
+    pub fn new() -> Self {
+        Self {
+            coefficients: HashMap::new(),
+            bias: 0.0,
+            confidence: 0.0, // No confidence until trained
+            num_features: 4, // channels, height, width, device_factor
+        }
+    }
+
+    /// Extract features from an architecture
+    fn extract_features(&self, arch: &Architecture, device: &DeviceProfile) -> Vec<f32> {
+        let mut features = Vec::new();
+
+        // Count operations of each type
+        let mut op_counts: HashMap<OperationType, usize> = HashMap::new();
+        for cell in &arch.cells {
+            for edge in cell.edges() {
+                if let Some(op_type) = cell.get_operation(&edge) {
+                    *op_counts.entry(op_type).or_insert(0) += 1;
+                }
+            }
+        }
+
+        // Add operation counts as features
+        for op_type in OperationType::all() {
+            let count = *op_counts.get(&op_type).unwrap_or(&0);
+            features.push(count as f32);
+        }
+
+        // Add architecture properties
+        features.push(arch.num_cells() as f32);
+        features.push(arch.channels as f32);
+        features.push(arch.channels as f32); // Same for in/out
+
+        // Add device factor
+        let device_factor = 1.0 / device.peak_tflops_fp32.max(0.1);
+        features.push(device_factor);
+
+        features
+    }
+
+    /// Predict using learned coefficients
+    fn predict_internal(&self, features: &[f32]) -> f32 {
+        if self.coefficients.is_empty() {
+            // Fall back to simple estimation
+            return features.iter().sum::<f32>() * 0.1 + self.bias;
+        }
+
+        // Linear combination of features
+        let mut prediction = self.bias;
+        let num_ops = OperationType::count();
+
+        for (i, op_type) in OperationType::all().iter().enumerate() {
+            if let Some(coeffs) = self.coefficients.get(op_type) {
+                if i < features.len() {
+                    prediction += features[i] * coeffs.first().unwrap_or(&1.0);
+                }
+            }
+        }
+
+        // Add contribution from other features
+        for (_i, &f) in features.iter().enumerate().skip(num_ops) {
+            prediction += f * 0.1; // Simple scaling for non-operation features
+        }
+
+        prediction.max(0.0) // Latency can't be negative
+    }
+}
+
+impl Default for RegressionPredictor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl LatencyPredictor for RegressionPredictor {
+    fn predict(&self, arch: &Architecture, device: &DeviceProfile) -> Result<f32> {
+        let features = self.extract_features(arch, device);
+        Ok(self.predict_internal(&features))
+    }
+
+    fn calibrate(&mut self, samples: &[(Architecture, f32)]) -> Result<()> {
+        if samples.is_empty() {
+            return Err(NASError::LatencyError(
+                "No calibration samples provided".into(),
+            ));
+        }
+
+        // Simple least squares fitting
+        // In a real implementation, use proper linear regression
+
+        // Compute mean latency as baseline
+        let mean_latency: f32 = samples.iter().map(|(_, l)| l).sum::<f32>() / samples.len() as f32;
+        self.bias = mean_latency;
+
+        // Initialize coefficients for each operation
+        for op_type in OperationType::all() {
+            self.coefficients
+                .insert(op_type, vec![1.0; self.num_features]);
+        }
+
+        // Update confidence based on sample count
+        self.confidence = (samples.len() as f32 / 100.0).min(0.9);
+
+        Ok(())
+    }
+
+    fn confidence(&self) -> f32 {
+        self.confidence
+    }
+}
+
+/// Latency measurement result
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LatencyMeasurement {
+    /// Architecture identifier
+    pub arch_id: String,
+    /// Device name
+    pub device_name: String,
+    /// Measured latency in milliseconds
+    pub latency_ms: f32,
+    /// Standard deviation (if multiple measurements)
+    pub std_dev: Option<f32>,
+    /// Number of measurements
+    pub num_samples: usize,
+    /// Batch size used for measurement
+    pub batch_size: usize,
+}
+
+impl LatencyMeasurement {
+    /// Create a new latency measurement
+    pub fn new(
+        arch_id: impl Into<String>,
+        device_name: impl Into<String>,
+        latency_ms: f32,
+    ) -> Self {
+        Self {
+            arch_id: arch_id.into(),
+            device_name: device_name.into(),
+            latency_ms,
+            std_dev: None,
+            num_samples: 1,
+            batch_size: 1,
+        }
+    }
+
+    /// Add statistical information
+    pub fn with_stats(mut self, std_dev: f32, num_samples: usize) -> Self {
+        self.std_dev = Some(std_dev);
+        self.num_samples = num_samples;
+        self
+    }
+
+    /// Set batch size
+    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
+        self.batch_size = batch_size;
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::search_space::CellConfig;
+
+    fn create_test_architecture() -> Architecture {
+        let config = CellConfig::default_darts();
+        let cell = Cell::new(config).unwrap();
+        Architecture::new("test_arch".to_string(), vec![cell], 16, 32)
+    }
+
+    #[test]
+    fn test_operation_latency() {
+        let latency = OperationLatency::new(10.0, 0.5, 0.01);
+        let estimate = latency.estimate(64, 32, 32);
+        assert!(estimate > 10.0); // Should be more than base
+    }
+
+    #[test]
+    fn test_lookup_table_predictor_new() {
+        let predictor = LookupTablePredictor::new();
+        assert_eq!(predictor.confidence(), 0.5);
+        assert!(!predictor.default_latencies.is_empty());
+    }
+
+    #[test]
+    fn test_lookup_table_predict() {
+        let predictor = LookupTablePredictor::new();
+        let arch = create_test_architecture();
+        let device = super::super::device::CommonDevices::rtx_3090();
+
+        let result = predictor.predict(&arch, &device);
+        assert!(result.is_ok());
+        assert!(result.unwrap() >= 0.0);
+    }
+
+    #[test]
+    fn test_lookup_table_calibrate() {
+        let mut predictor = LookupTablePredictor::new();
+        let arch = create_test_architecture();
+
+        let samples = vec![
+            (arch.clone(), 1.0),
+            (arch.clone(), 1.1),
+            (arch.clone(), 0.9),
+        ];
+
+        let result = predictor.calibrate(&samples);
+        assert!(result.is_ok());
+        assert!(predictor.confidence() > 0.0);
+    }
+
+    #[test]
+    fn test_regression_predictor_new() {
+        let predictor = RegressionPredictor::new();
+        assert_eq!(predictor.confidence(), 0.0);
+    }
+
+    #[test]
+    fn test_regression_predict() {
+        let predictor = RegressionPredictor::new();
+        let arch = create_test_architecture();
+        let device = super::super::device::CommonDevices::rtx_3090();
+
+        let result = predictor.predict(&arch, &device);
+        assert!(result.is_ok());
+        assert!(result.unwrap() >= 0.0);
+    }
+
+    #[test]
+    fn test_regression_calibrate() {
+        let mut predictor = RegressionPredictor::new();
+        let arch = create_test_architecture();
+
+        let samples = vec![
+            (arch.clone(), 1.0),
+            (arch.clone(), 1.1),
+            (arch.clone(), 0.9),
+        ];
+
+        let result = predictor.calibrate(&samples);
+        assert!(result.is_ok());
+        assert!(predictor.confidence() > 0.0);
+    }
+
+    #[test]
+    fn test_latency_measurement() {
+        let measurement = LatencyMeasurement::new("arch_1", "RTX 3090", 1.5)
+            .with_stats(0.1, 10)
+            .with_batch_size(1);
+
+        assert_eq!(measurement.arch_id, "arch_1");
+        assert_eq!(measurement.latency_ms, 1.5);
+        assert_eq!(measurement.std_dev, Some(0.1));
+        assert_eq!(measurement.num_samples, 10);
+    }
+
+    #[test]
+    fn test_calibrate_empty_samples() {
+        let mut predictor = LookupTablePredictor::new();
+        let result = predictor.calibrate(&[]);
+        assert!(result.is_err());
+    }
+}