src/core/cpu_optimizer.py

"""
Advanced CPU Optimizer for training on CPU-only systems
Optimized for maximum performance on limited hardware
"""

import os
import logging
import threading
from typing import Dict, Any, Optional, List
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
from .memory_manager import AdvancedMemoryManager

logger = logging.getLogger(__name__)

class CPUOptimizer:
    """
    Advanced CPU optimization for training and inference
    """
    
    def __init__(self, memory_manager: AdvancedMemoryManager):
        """
        Initialize CPU optimizer
        
        Args:
            memory_manager: Memory manager instance
        """
        self.memory_manager = memory_manager
        self.cpu_count = os.cpu_count()
        self.optimizations_applied = []
        
        # Apply initial optimizations
        self._apply_global_optimizations()
        
        logger.info(f"CPU Optimizer initialized for {self.cpu_count} cores")
    
    def _apply_global_optimizations(self):
        """Apply global CPU optimizations"""
        
        # Set optimal thread count for PyTorch
        optimal_threads = min(self.cpu_count, 8)  # Cap at 8 for stability
        torch.set_num_threads(optimal_threads)
        self.optimizations_applied.append(f"PyTorch threads: {optimal_threads}")
        
        # Set thread count for inter-op parallelism
        torch.set_num_interop_threads(min(self.cpu_count // 2, 4))
        self.optimizations_applied.append("Inter-op parallelism configured")
        
        # Enable Intel MKL optimizations if available
        try:
            import intel_extension_for_pytorch as ipex
            self.optimizations_applied.append("Intel Extension for PyTorch enabled")
        except ImportError:
            logger.warning("Intel Extension for PyTorch not available")
        
        # Set environment variables for CPU optimization
        os.environ['OMP_NUM_THREADS'] = str(optimal_threads)
        os.environ['MKL_NUM_THREADS'] = str(optimal_threads)
        os.environ['NUMEXPR_NUM_THREADS'] = str(optimal_threads)
        os.environ['OPENBLAS_NUM_THREADS'] = str(optimal_threads)
        self.optimizations_applied.append("Environment variables optimized")
        
        # Enable CPU-specific optimizations
        torch.backends.mkl.enabled = True
        torch.backends.mkldnn.enabled = True
        self.optimizations_applied.append("MKL and MKLDNN enabled")
        
        logger.info(f"Applied optimizations: {', '.join(self.optimizations_applied)}")
    
    def optimize_model(self, model: nn.Module, 
                      use_jit: bool = True,
                      use_channels_last: bool = True) -> nn.Module:
        """
        Optimize model for CPU inference/training
        
        Args:
            model: PyTorch model to optimize
            use_jit: Whether to use TorchScript JIT compilation
            use_channels_last: Whether to use channels-last memory format
            
        Returns:
            Optimized model
        """
        with self.memory_manager.memory_context("optimize_model"):
            logger.info("Optimizing model for CPU")
            
            # Set model to CPU
            model = model.cpu()
            
            # Set to evaluation mode for optimization
            was_training = model.training
            model.eval()
            
            try:
                # Apply Intel Extension optimizations if available
                try:
                    import intel_extension_for_pytorch as ipex
                    model = ipex.optimize(model, dtype=torch.float32)
                    logger.info("Applied Intel Extension optimizations")
                except ImportError:
                    pass
                
                # Apply channels-last memory format for conv models
                if use_channels_last and self._has_conv_layers(model):
                    model = model.to(memory_format=torch.channels_last)
                    logger.info("Applied channels-last memory format")
                
                # Apply TorchScript JIT compilation
                if use_jit:
                    try:
                        # Create dummy input for tracing
                        dummy_input = self._create_dummy_input(model)
                        if dummy_input is not None:
                            model = torch.jit.trace(model, dummy_input)
                            logger.info("Applied TorchScript JIT compilation")
                    except Exception as e:
                        logger.warning(f"JIT compilation failed: {e}")
                
                # Restore training mode if needed
                if was_training:
                    model.train()
                
                return model
                
            except Exception as e:
                logger.error(f"Model optimization failed: {e}")
                return model
    
    def _has_conv_layers(self, model: nn.Module) -> bool:
        """Check if model has convolutional layers"""
        for module in model.modules():
            if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
                return True
        return False
    
    def _create_dummy_input(self, model: nn.Module) -> Optional[torch.Tensor]:
        """Create dummy input for model tracing"""
        try:
            # Try to infer input shape from model
            for name, param in model.named_parameters():
                if 'embedding' in name.lower() and param.dim() == 2:
                    # Text model - create token input
                    vocab_size = param.shape[0]
                    return torch.randint(0, min(vocab_size, 1000), (1, 32))
                elif 'conv' in name.lower() and param.dim() == 4:
                    # Vision model - create image input
                    channels = param.shape[1]
                    return torch.randn(1, channels, 224, 224)
            
            # Default fallback
            return torch.randn(1, 512)
            
        except Exception:
            return None
    
    def optimize_dataloader(self, dataloader: DataLoader) -> DataLoader:
        """
        Optimize DataLoader for CPU training
        
        Args:
            dataloader: Original DataLoader
            
        Returns:
            Optimized DataLoader
        """
        # Calculate optimal number of workers
        optimal_workers = min(self.cpu_count // 2, 4)
        
        # Create new DataLoader with optimized settings
        optimized_loader = DataLoader(
            dataloader.dataset,
            batch_size=dataloader.batch_size,
            shuffle=dataloader.drop_last if hasattr(dataloader, 'drop_last') else False,
            num_workers=optimal_workers,
            pin_memory=False,  # Not needed for CPU
            persistent_workers=True if optimal_workers > 0 else False,
            prefetch_factor=2 if optimal_workers > 0 else 2,
        )
        
        logger.info(f"Optimized DataLoader with {optimal_workers} workers")
        return optimized_loader
    
    def optimize_optimizer(self, optimizer: optim.Optimizer, 
                          model: nn.Module) -> optim.Optimizer:
        """
        Optimize optimizer settings for CPU training
        
        Args:
            optimizer: PyTorch optimizer
            model: Model being optimized
            
        Returns:
            Optimized optimizer
        """
        # Apply gradient clipping
        for param_group in optimizer.param_groups:
            if 'weight_decay' not in param_group:
                param_group['weight_decay'] = 0.01
        
        logger.info("Applied optimizer optimizations")
        return optimizer
    
    def enable_mixed_precision(self) -> bool:
        """
        Enable mixed precision training for CPU (if supported)
        
        Returns:
            Whether mixed precision was enabled
        """
        try:
            # Check if CPU supports mixed precision
            if torch.cpu.amp.autocast is not None:
                logger.info("CPU mixed precision available")
                return True
        except AttributeError:
            pass
        
        logger.warning("CPU mixed precision not available")
        return False
    
    def optimize_batch_size(self, base_batch_size: int, 
                           model_size_mb: float) -> int:
        """
        Calculate optimal batch size based on available memory
        
        Args:
            base_batch_size: Base batch size to start from
            model_size_mb: Model size in MB
            
        Returns:
            Optimized batch size
        """
        memory_info = self.memory_manager.get_memory_info()
        available_memory_mb = memory_info['system_memory_available_gb'] * 1024
        
        # Reserve memory for model and overhead
        usable_memory_mb = available_memory_mb - model_size_mb - 2000  # 2GB overhead
        
        # Estimate memory per sample (rough approximation)
        memory_per_sample_mb = model_size_mb * 0.1  # 10% of model size per sample
        
        if memory_per_sample_mb > 0:
            max_batch_size = int(usable_memory_mb / memory_per_sample_mb)
            optimal_batch_size = min(base_batch_size, max_batch_size, 32)  # Cap at 32
        else:
            optimal_batch_size = min(base_batch_size, 8)  # Conservative fallback
        
        optimal_batch_size = max(1, optimal_batch_size)  # At least 1
        
        logger.info(f"Optimized batch size: {optimal_batch_size} (was {base_batch_size})")
        return optimal_batch_size
    
    def get_performance_recommendations(self, model: nn.Module) -> List[str]:
        """
        Get performance recommendations for the current setup
        
        Args:
            model: Model to analyze
            
        Returns:
            List of recommendations
        """
        recommendations = []
        
        # Check model size
        param_count = sum(p.numel() for p in model.parameters())
        model_size_mb = param_count * 4 / (1024**2)  # Assume float32
        
        if model_size_mb > 2000:  # > 2GB
            recommendations.append("Consider using model sharding for large models")
            recommendations.append("Use gradient checkpointing to reduce memory usage")
        
        # Check CPU utilization
        if self.cpu_count > 8:
            recommendations.append("Consider using distributed training across CPU cores")
        
        # Check memory
        memory_info = self.memory_manager.get_memory_info()
        if memory_info['system_memory_percent'] > 80:
            recommendations.append("Reduce batch size to lower memory usage")
            recommendations.append("Enable gradient accumulation instead of large batches")
        
        # Check for optimization opportunities
        if not any('Intel Extension' in opt for opt in self.optimizations_applied):
            recommendations.append("Install Intel Extension for PyTorch for better CPU performance")
        
        return recommendations
    
    def benchmark_performance(self, model: nn.Module, 
                            input_shape: tuple,
                            num_iterations: int = 100) -> Dict[str, float]:
        """
        Benchmark model performance
        
        Args:
            model: Model to benchmark
            input_shape: Input tensor shape
            num_iterations: Number of iterations to run
            
        Returns:
            Performance metrics
        """
        model.eval()
        dummy_input = torch.randn(*input_shape)
        
        # Warmup
        with torch.no_grad():
            for _ in range(10):
                _ = model(dummy_input)
        
        # Benchmark
        import time
        start_time = time.time()
        
        with torch.no_grad():
            for _ in range(num_iterations):
                _ = model(dummy_input)
        
        end_time = time.time()
        
        total_time = end_time - start_time
        avg_time_per_inference = total_time / num_iterations
        throughput = 1.0 / avg_time_per_inference
        
        return {
            'total_time_seconds': total_time,
            'avg_time_per_inference_ms': avg_time_per_inference * 1000,
            'throughput_inferences_per_second': throughput,
            'iterations': num_iterations
        }