optimized_executor.py

"""
Enhanced quantized executor with PyTorch kernels, per-channel quantization, and fusion
"""
import torch
import torch.nn.functional as F
from typing import Dict, Any
from executor.quantized_kernels import run_mir_quantized_fast
from passes.per_channel_quantize import linear_per_channel_quantized
from passes.fusion import execute_fused_linear_relu, execute_fused_conv_relu

def run_mir_optimized(mir: Dict[str, Any], model: torch.nn.Module, 
                     input_tensor: torch.Tensor) -> torch.Tensor:
    """
    Run MIR with all optimizations: PyTorch kernels, per-channel quantization, and fusion.
    
    Args:
        mir: MIR dictionary with optimizations applied
        model: Original PyTorch model
        input_tensor: Input tensor
    
    Returns:
        Output tensor after optimized execution
    """
    print("🚀 OPTIMIZED QUANTIZED EXECUTION")
    print("=" * 50)
    print("Features: PyTorch Kernels + Per-Channel Quantization + Operator Fusion")
    
    model.eval()
    named_modules = dict(model.named_modules())
    nodes = mir["graph"]["nodes"]
    
    # Get quantization metadata
    per_channel_quant = mir.get("metadata", {}).get("per_channel_quantization", {})
    weight_quant = mir.get("metadata", {}).get("quantization", {})
    act_quant = mir.get("metadata", {}).get("activation_quantization", {})
    
    # Check if fusion was applied
    fusion_applied = mir.get("metadata", {}).get("fusion_applied", False)
    if fusion_applied:
        fusions_count = mir.get("metadata", {}).get("fusions_count", 0)
        print(f"🔧 Fusion optimizations: {fusions_count} fusions applied")
    
    # Check if per-channel quantization is available
    per_channel_available = len(per_channel_quant) > 0
    if per_channel_available:
        per_channel_count = len([k for k, v in per_channel_quant.items() if v.get('per_channel', False)])
        print(f"📊 Per-channel quantization: {per_channel_count} layers optimized")
    
    # Quantize input
    input_range = input_tensor.max() - input_tensor.min()
    input_scale = input_range.item() / 255.0 if input_range > 0 else 0.01
    
    print(f"Input: {input_tensor.shape}, range=[{input_tensor.min():.4f}, {input_tensor.max():.4f}]")
    print(f"Input scale: {input_scale:.6f}")
    
    x = input_tensor.clone()
    
    # Execute nodes with optimizations
    for i, node in enumerate(nodes):
        node_id = node["id"]
        op_type = node["op_type"]
        origin_name = node.get("metadata", {}).get("origin_name")
        
        print(f"Executing optimized node {i}: {node_id} ({op_type})")
        
        # Handle fused operations
        if op_type == "linear_relu_fused":
            x = _execute_fused_linear_relu_optimized(x, node, named_modules, 
                                                   per_channel_quant, weight_quant)
        
        elif op_type == "conv2d_relu_fused":
            x = _execute_fused_conv_relu_optimized(x, node, named_modules, 
                                                 per_channel_quant, weight_quant)
        
        # Handle individual operations
        elif op_type == "linear":
            x = _execute_linear_optimized(x, node, named_modules, 
                                        per_channel_quant, weight_quant)
        
        elif op_type == "conv2d":
            x = _execute_conv2d_optimized(x, node, named_modules, 
                                        per_channel_quant, weight_quant)
        
        elif op_type == "relu":
            x = F.relu(x)
            print(f"  ✅ ReLU applied")
        
        elif op_type == "adaptive_avg_pool2d":
            # Handle AdaptiveAvgPool2d operations
            params = node.get("params", {})
            output_size = params.get("output_size", (1, 1))
            x = F.adaptive_avg_pool2d(x, output_size)
            print(f"  ✅ AdaptiveAvgPool2d applied, output_size={output_size}")
            
        elif op_type == "max_pool2d":
            # Handle MaxPool2d operations
            params = node.get("params", {})
            kernel_size = params.get("kernel_size", 2)
            stride = params.get("stride", None)
            padding = params.get("padding", 0)
            dilation = params.get("dilation", 1)
            x = F.max_pool2d(x, kernel_size, stride, padding, dilation)
            print(f"  ✅ MaxPool2d applied, kernel_size={kernel_size}")
            
        elif op_type == "avg_pool2d":
            # Handle AvgPool2d operations  
            params = node.get("params", {})
            kernel_size = params.get("kernel_size", 2)
            stride = params.get("stride", None)
            padding = params.get("padding", 0)
            x = F.avg_pool2d(x, kernel_size, stride, padding)
            print(f"  ✅ AvgPool2d applied, kernel_size={kernel_size}")
            
        elif op_type == "flatten":
            # Handle Flatten operations
            params = node.get("params", {})
            start_dim = params.get("start_dim", 1)
            end_dim = params.get("end_dim", -1)
            x = torch.flatten(x, start_dim, end_dim)
            print(f"  ✅ Flatten applied, start_dim={start_dim}, end_dim={end_dim}")
        
        else:
            print(f"  ⚠️  Unsupported operation: {op_type}")
        
        # Apply activation quantization if available
        if node_id in act_quant:
            act_scale = act_quant[node_id]
            x_quantized = torch.round(x / act_scale).clamp(-127, 127)
            x = x_quantized * act_scale
            print(f"  📊 Activation quantized with scale={act_scale:.6f}")
        
        print(f"  Output: {x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
    
    print(f"\n✅ Optimized execution complete: {x}")
    return x

def _execute_fused_linear_relu_optimized(x: torch.Tensor, node: Dict, named_modules: Dict, 
                                       per_channel_quant: Dict, weight_quant: Dict) -> torch.Tensor:
    """Execute fused Linear+ReLU with optimizations"""
    origin_name = node.get("metadata", {}).get("origin_name")
    
    if origin_name and origin_name in named_modules:
        module = named_modules[origin_name]
        weight = module.weight.data
        bias = module.bias.data if module.bias is not None else None
        
        # Check for per-channel quantization
        weight_key = f"{origin_name}.weight"
        bias_key = f"{origin_name}.bias"
        
        if weight_key in per_channel_quant and per_channel_quant[weight_key].get("per_channel", False):
            # Use per-channel quantization
            weight_scales = torch.tensor(per_channel_quant[weight_key]["scales"])
            bias_scale = per_channel_quant.get(bias_key, {}).get("scale", 0.01)
            
            x = execute_fused_linear_relu(x, weight, bias, weight_scales, bias_scale, per_channel=True)
            print(f"  ✅ Fused Linear+ReLU (per-channel), channels={len(weight_scales)}")
            
        elif weight_key in weight_quant:
            # Use per-tensor quantization
            weight_scale = weight_quant[weight_key]["scale"]
            bias_scale = weight_quant.get(bias_key, {}).get("scale", 0.01)
            
            x = execute_fused_linear_relu(x, weight, bias, weight_scale, bias_scale, per_channel=False)
            print(f"  ✅ Fused Linear+ReLU (per-tensor), scale={weight_scale:.6f}")
            
        else:
            # No quantization metadata, use float
            x = execute_fused_linear_relu(x, weight, bias)
            print(f"  ✅ Fused Linear+ReLU (float)")
    
    return x

def _execute_fused_conv_relu_optimized(x: torch.Tensor, node: Dict, named_modules: Dict, 
                                     per_channel_quant: Dict, weight_quant: Dict) -> torch.Tensor:
    """Execute fused Conv2D+ReLU with optimizations"""
    origin_name = node.get("metadata", {}).get("origin_name")
    
    if origin_name and origin_name in named_modules:
        module = named_modules[origin_name]
        weight = module.weight.data
        bias = module.bias.data if module.bias is not None else None
        stride = module.stride[0] if hasattr(module, 'stride') else 1
        padding = module.padding[0] if hasattr(module, 'padding') else 0
        
        # Check for per-channel quantization
        weight_key = f"{origin_name}.weight"
        bias_key = f"{origin_name}.bias"
        
        if weight_key in per_channel_quant and per_channel_quant[weight_key].get("per_channel", False):
            # Use per-channel quantization
            weight_scales = torch.tensor(per_channel_quant[weight_key]["scales"])
            bias_scale = per_channel_quant.get(bias_key, {}).get("scale", 0.01)
            
            x = execute_fused_conv_relu(x, weight, bias, stride, padding, 
                                      weight_scales, bias_scale, per_channel=True)
            print(f"  ✅ Fused Conv2D+ReLU (per-channel), channels={len(weight_scales)}")
            
        elif weight_key in weight_quant:
            # Use per-tensor quantization
            weight_scale = weight_quant[weight_key]["scale"]
            bias_scale = weight_quant.get(bias_key, {}).get("scale", 0.01)
            
            x = execute_fused_conv_relu(x, weight, bias, stride, padding, 
                                      weight_scale, bias_scale, per_channel=False)
            print(f"  ✅ Fused Conv2D+ReLU (per-tensor), scale={weight_scale:.6f}")
            
        else:
            # No quantization metadata, use float
            x = execute_fused_conv_relu(x, weight, bias, stride, padding)
            print(f"  ✅ Fused Conv2D+ReLU (float)")
    
    return x

def _execute_linear_optimized(x: torch.Tensor, node: Dict, named_modules: Dict, 
                            per_channel_quant: Dict, weight_quant: Dict) -> torch.Tensor:
    """Execute Linear operation with optimizations"""
    origin_name = node.get("metadata", {}).get("origin_name")
    
    if origin_name and origin_name in named_modules:
        module = named_modules[origin_name]
        weight = module.weight.data
        bias = module.bias.data if module.bias is not None else None
        
        # Check for per-channel quantization
        weight_key = f"{origin_name}.weight"
        bias_key = f"{origin_name}.bias"
        
        if weight_key in per_channel_quant and per_channel_quant[weight_key].get("per_channel", False):
            # Use per-channel quantization
            weight_scales = torch.tensor(per_channel_quant[weight_key]["scales"])
            bias_scale = per_channel_quant.get(bias_key, {}).get("scale", 0.01)
            
            x = linear_per_channel_quantized(x, weight, bias, weight_scales, bias_scale)
            print(f"  ✅ Linear (per-channel), channels={len(weight_scales)}")
            
        elif weight_key in weight_quant:
            # Use per-tensor quantization
            weight_scale = weight_quant[weight_key]["scale"]
            bias_scale = weight_quant.get(bias_key, {}).get("scale", 0.01)
            
            # Quantize weights
            weight_quantized = torch.round(weight / weight_scale).clamp(-127, 127)
            weight_dequantized = weight_quantized * weight_scale
            
            bias_dequantized = None
            if bias is not None:
                bias_quantized = torch.round(bias / bias_scale).clamp(-127, 127)
                bias_dequantized = bias_quantized * bias_scale
            
            x = F.linear(x, weight_dequantized, bias_dequantized)
            print(f"  ✅ Linear (per-tensor), scale={weight_scale:.6f}")
            
        else:
            # No quantization metadata, use float
            x = F.linear(x, weight, bias)
            print(f"  ✅ Linear (float)")
    
    return x

def _execute_conv2d_optimized(x: torch.Tensor, node: Dict, named_modules: Dict, 
                            per_channel_quant: Dict, weight_quant: Dict) -> torch.Tensor:
    """Execute Conv2D operation with optimizations"""
    origin_name = node.get("metadata", {}).get("origin_name")
    
    if origin_name and origin_name in named_modules:
        module = named_modules[origin_name]
        weight = module.weight.data
        bias = module.bias.data if module.bias is not None else None
        stride = module.stride[0] if hasattr(module, 'stride') else 1
        padding = module.padding[0] if hasattr(module, 'padding') else 0
        
        # Check for per-channel quantization
        weight_key = f"{origin_name}.weight"
        bias_key = f"{origin_name}.bias"
        
        if weight_key in per_channel_quant and per_channel_quant[weight_key].get("per_channel", False):
            # Use per-channel quantization
            from passes.per_channel_quantize import quantize_per_channel, dequantize_per_channel
            
            weight_scales = torch.tensor(per_channel_quant[weight_key]["scales"])
            bias_scale = per_channel_quant.get(bias_key, {}).get("scale", 0.01)
            
            weight_quantized = quantize_per_channel(weight, weight_scales)
            weight_dequantized = dequantize_per_channel(weight_quantized, weight_scales)
            
            bias_dequantized = None
            if bias is not None:
                bias_quantized = torch.round(bias / bias_scale).clamp(-127, 127)
                bias_dequantized = bias_quantized * bias_scale
            
            x = F.conv2d(x, weight_dequantized, bias_dequantized, stride, padding)
            print(f"  ✅ Conv2D (per-channel), channels={len(weight_scales)}")
            
        elif weight_key in weight_quant:
            # Use per-tensor quantization
            weight_scale = weight_quant[weight_key]["scale"]
            bias_scale = weight_quant.get(bias_key, {}).get("scale", 0.01)
            
            # Quantize weights
            weight_quantized = torch.round(weight / weight_scale).clamp(-127, 127)
            weight_dequantized = weight_quantized * weight_scale
            
            bias_dequantized = None
            if bias is not None:
                bias_quantized = torch.round(bias / bias_scale).clamp(-127, 127)
                bias_dequantized = bias_quantized * bias_scale
            
            x = F.conv2d(x, weight_dequantized, bias_dequantized, stride, padding)
            print(f"  ✅ Conv2D (per-tensor), scale={weight_scale:.6f}")
            
        else:
            # No quantization metadata, use float
            x = F.conv2d(x, weight, bias, stride, padding)
            print(f"  ✅ Conv2D (float)")
    
    return x