ghostflow_nn/

inference.rs

//! Inference optimization utilities
//!
//! This module provides optimizations for model inference including:
//! - Operator fusion
//! - Constant folding
//! - Memory optimization
//! - Batch inference

use ghostflow_core::{Result, Tensor, GhostError};
use std::collections::HashMap;

/// Inference mode configuration
#[derive(Debug, Clone)]
pub struct InferenceConfig {
    /// Enable operator fusion
    pub enable_fusion: bool,
    /// Enable constant folding
    pub enable_constant_folding: bool,
    /// Batch size for inference
    pub batch_size: usize,
    /// Use mixed precision (FP16)
    pub use_mixed_precision: bool,
    /// Number of threads for CPU inference
    pub num_threads: usize,
}

impl Default for InferenceConfig {
    fn default() -> Self {
        Self {
            enable_fusion: true,
            enable_constant_folding: true,
            batch_size: 1,
            use_mixed_precision: false,
            num_threads: num_cpus::get(),
        }
    }
}

/// Inference optimizer
pub struct InferenceOptimizer {
    config: InferenceConfig,
    fused_ops: Vec<FusedOp>,
}

/// Fused operation
#[derive(Debug, Clone)]
pub struct FusedOp {
    pub name: String,
    pub ops: Vec<String>,
}

impl InferenceOptimizer {
    /// Create a new inference optimizer
    pub fn new(config: InferenceConfig) -> Self {
        Self {
            config,
            fused_ops: Vec::new(),
        }
    }

    /// Optimize a model for inference
    pub fn optimize(&mut self) -> Result<()> {
        if self.config.enable_fusion {
            self.fuse_operators()?;
        }
        
        if self.config.enable_constant_folding {
            self.fold_constants()?;
        }
        
        Ok(())
    }

    /// Fuse common operator patterns
    fn fuse_operators(&mut self) -> Result<()> {
        // Conv + BatchNorm + ReLU fusion
        self.fused_ops.push(FusedOp {
            name: "ConvBNReLU".to_string(),
            ops: vec!["Conv2d".to_string(), "BatchNorm".to_string(), "ReLU".to_string()],
        });
        
        // Linear + ReLU fusion
        self.fused_ops.push(FusedOp {
            name: "LinearReLU".to_string(),
            ops: vec!["Linear".to_string(), "ReLU".to_string()],
        });
        
        // MatMul + Add fusion (GEMM)
        self.fused_ops.push(FusedOp {
            name: "GEMM".to_string(),
            ops: vec!["MatMul".to_string(), "Add".to_string()],
        });
        
        Ok(())
    }

    /// Fold constant operations
    fn fold_constants(&mut self) -> Result<()> {
        // Constant folding would pre-compute operations on constant tensors
        // This is a placeholder for the actual implementation
        Ok(())
    }

    /// Get fused operations
    pub fn get_fused_ops(&self) -> &[FusedOp] {
        &self.fused_ops
    }
}

/// Batch inference helper
pub struct BatchInference {
    batch_size: usize,
    buffer: Vec<Tensor>,
}

impl BatchInference {
    /// Create a new batch inference helper
    pub fn new(batch_size: usize) -> Self {
        Self {
            batch_size,
            buffer: Vec::new(),
        }
    }

    /// Add a sample to the batch
    pub fn add(&mut self, sample: Tensor) {
        self.buffer.push(sample);
    }

    /// Check if batch is ready
    pub fn is_ready(&self) -> bool {
        self.buffer.len() >= self.batch_size
    }

    /// Get the current batch and clear buffer
    pub fn get_batch(&mut self) -> Result<Option<Tensor>> {
        if !self.is_ready() {
            return Ok(None);
        }
        
        // Stack tensors into a batch
        let batch = self.stack_tensors()?;
        self.buffer.clear();
        Ok(Some(batch))
    }

    /// Flush remaining samples (even if batch is not full)
    pub fn flush(&mut self) -> Result<Option<Tensor>> {
        if self.buffer.is_empty() {
            return Ok(None);
        }
        
        let batch = self.stack_tensors()?;
        self.buffer.clear();
        Ok(Some(batch))
    }

    fn stack_tensors(&self) -> Result<Tensor> {
        if self.buffer.is_empty() {
            return Err(GhostError::InvalidShape("Empty buffer".to_string()));
        }
        
        let first_shape = self.buffer[0].dims();
        let batch_size = self.buffer.len();
        
        // Create new shape with batch dimension
        let mut new_shape = vec![batch_size];
        new_shape.extend_from_slice(first_shape);
        
        // Collect all data
        let mut all_data = Vec::new();
        for tensor in &self.buffer {
            all_data.extend(tensor.data_f32());
        }
        
        Tensor::from_slice(&all_data, &new_shape)
    }
}

/// Inference session for optimized model execution
pub struct InferenceSession {
    config: InferenceConfig,
    optimizer: InferenceOptimizer,
    cache: HashMap<String, Tensor>,
}

impl InferenceSession {
    /// Create a new inference session
    pub fn new(config: InferenceConfig) -> Self {
        let optimizer = InferenceOptimizer::new(config.clone());
        Self {
            config,
            optimizer,
            cache: HashMap::new(),
        }
    }

    /// Initialize the session
    pub fn initialize(&mut self) -> Result<()> {
        self.optimizer.optimize()?;
        Ok(())
    }

    /// Run inference on a single input
    pub fn run(&mut self, _input: &Tensor) -> Result<Tensor> {
        // Placeholder for actual inference
        // In a real implementation, this would execute the optimized model
        Err(GhostError::NotImplemented("Inference execution not yet implemented".to_string()))
    }

    /// Run batch inference
    pub fn run_batch(&mut self, _inputs: &[Tensor]) -> Result<Vec<Tensor>> {
        // Placeholder for batch inference
        Err(GhostError::NotImplemented("Batch inference not yet implemented".to_string()))
    }

    /// Cache a tensor for reuse
    pub fn cache_tensor(&mut self, name: String, tensor: Tensor) {
        self.cache.insert(name, tensor);
    }

    /// Get a cached tensor
    pub fn get_cached(&self, name: &str) -> Option<&Tensor> {
        self.cache.get(name)
    }

    /// Clear the cache
    pub fn clear_cache(&mut self) {
        self.cache.clear();
    }

    /// Get the configuration
    pub fn config(&self) -> &InferenceConfig {
        &self.config
    }
}

/// Warmup helper for inference
pub fn warmup_model<F>(mut inference_fn: F, input_shape: &[usize], num_iterations: usize) -> Result<f64>
where
    F: FnMut(&Tensor) -> Result<Tensor>,
{
    use std::time::Instant;
    
    // Create dummy input
    let numel: usize = input_shape.iter().product();
    let dummy_data = vec![0.0f32; numel];
    let dummy_input = Tensor::from_slice(&dummy_data, input_shape)?;
    
    // Warmup iterations
    for _ in 0..3 {
        let _ = inference_fn(&dummy_input)?;
    }
    
    // Timed iterations
    let start = Instant::now();
    for _ in 0..num_iterations {
        let _ = inference_fn(&dummy_input)?;
    }
    let elapsed = start.elapsed();
    
    // Return average time in milliseconds
    Ok(elapsed.as_secs_f64() * 1000.0 / num_iterations as f64)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_inference_config() {
        let config = InferenceConfig::default();
        assert!(config.enable_fusion);
        assert!(config.enable_constant_folding);
        assert_eq!(config.batch_size, 1);
    }

    #[test]
    fn test_inference_optimizer() {
        let config = InferenceConfig::default();
        let mut optimizer = InferenceOptimizer::new(config);
        
        optimizer.optimize().unwrap();
        
        let fused_ops = optimizer.get_fused_ops();
        assert!(!fused_ops.is_empty());
    }

    #[test]
    fn test_batch_inference() {
        let mut batch = BatchInference::new(2);
        
        let t1 = Tensor::from_slice(&[1.0f32, 2.0], &[2]).unwrap();
        let t2 = Tensor::from_slice(&[3.0f32, 4.0], &[2]).unwrap();
        
        batch.add(t1);
        assert!(!batch.is_ready());
        
        batch.add(t2);
        assert!(batch.is_ready());
        
        let batched = batch.get_batch().unwrap().unwrap();
        assert_eq!(batched.dims(), &[2, 2]);
    }

    #[test]
    fn test_batch_flush() {
        let mut batch = BatchInference::new(3);
        
        let t1 = Tensor::from_slice(&[1.0f32, 2.0], &[2]).unwrap();
        batch.add(t1);
        
        let flushed = batch.flush().unwrap().unwrap();
        assert_eq!(flushed.dims(), &[1, 2]);
    }

    #[test]
    fn test_inference_session() {
        let config = InferenceConfig::default();
        let mut session = InferenceSession::new(config);
        
        session.initialize().unwrap();
        
        // Test caching
        let tensor = Tensor::from_slice(&[1.0f32, 2.0, 3.0], &[3]).unwrap();
        session.cache_tensor("test".to_string(), tensor);
        
        assert!(session.get_cached("test").is_some());
        
        session.clear_cache();
        assert!(session.get_cached("test").is_none());
    }
}