quantized_kernels.py

"""
Fast quantized execution using PyTorch's optimized int8 kernels
Falls back to simulation when PyTorch kernels are unavailable
"""
import torch
import torch.nn.functional as F
from typing import Dict, Any

def quantize_tensor_pytorch(tensor: torch.Tensor, scale: float, zero_point: int = 0, 
                           dtype: torch.dtype = torch.qint8) -> torch.Tensor:
    """Quantize tensor using PyTorch's built-in quantization"""
    try:
        return torch.quantize_per_tensor(tensor, scale, zero_point, dtype)
    except Exception:
        # Fallback to manual quantization
        qmin = -128 if dtype == torch.qint8 else 0
        qmax = 127 if dtype == torch.qint8 else 255
        quantized = ((tensor / scale) + zero_point).round().clamp(qmin, qmax)
        return quantized.to(torch.int8 if dtype == torch.qint8 else torch.uint8)

def dequantize_tensor_pytorch(qtensor: torch.Tensor, scale: float, 
                             zero_point: int = 0) -> torch.Tensor:
    """Dequantize tensor using PyTorch's built-in dequantization"""
    try:
        if hasattr(qtensor, 'dequantize'):
            return qtensor.dequantize()
        else:
            return (qtensor.float() - zero_point) * scale
    except Exception:
        return (qtensor.float() - zero_point) * scale

def create_quantized_linear(weight: torch.Tensor, bias: torch.Tensor, 
                           weight_scale: float, bias_scale: float) -> torch.nn.Module:
    """Create a quantized linear layer using PyTorch's quantized modules"""
    try:
        # Try to create PyTorch quantized linear
        import torch.nn.quantized as nnq
        
        # Quantize weights and bias
        qweight = quantize_tensor_pytorch(weight, weight_scale, 0, torch.qint8)
        qbias = quantize_tensor_pytorch(bias, bias_scale, 0, torch.qint32) if bias is not None else None
        
        # Create quantized linear layer
        qlinear = nnq.Linear(weight.shape[1], weight.shape[0])
        qlinear.set_weight_bias(qweight, qbias)
        qlinear.scale = weight_scale
        qlinear.zero_point = 0
        
        return qlinear
    except Exception as e:
        print(f"PyTorch quantized linear creation failed: {e}")
        return None

def run_mir_quantized_fast(mir: Dict[str, Any], model: torch.nn.Module, 
                          input_tensor: torch.Tensor) -> torch.Tensor:
    """
    Fast quantized execution using PyTorch's optimized kernels when available.
    Falls back to simulation when kernels are not supported.
    """
    print("🚀 FAST QUANTIZED EXECUTION (PyTorch Kernels)")
    print("=" * 50)
    
    try:
        # First try to check if PyTorch quantization is working
        test_tensor = torch.randn(2, 2)
        test_scale = 0.1
        test_quantized = quantize_tensor_pytorch(test_tensor, test_scale)
        test_dequantized = dequantize_tensor_pytorch(test_quantized, test_scale)
        
        # If we get here, PyTorch quantization is working
        return _run_with_pytorch_quantized(mir, model, input_tensor)
    except Exception as e:
        print(f"⚠️  PyTorch quantized kernels not available: {type(e).__name__}")
        print("📋 Falling back to validated simulation...")
        # Import the working simulation function
        from executor.run import run_mir_quantized
        return run_mir_quantized(mir, model, input_tensor)

def _run_with_pytorch_quantized(mir: Dict[str, Any], model: torch.nn.Module, 
                               input_tensor: torch.Tensor) -> torch.Tensor:
    """Internal function using PyTorch quantized operations"""
    model.eval()
    named_modules = dict(model.named_modules())
    nodes = mir["graph"]["nodes"]
    
    # Get quantization metadata
    weight_quant = mir.get("metadata", {}).get("quantization", {})
    act_quant = mir.get("metadata", {}).get("activation_quantization", {})
    
    # Quantize input
    input_scale = 0.01  # Dynamic scale for input
    input_range = input_tensor.max() - input_tensor.min()
    if input_range > 0:
        input_scale = input_range.item() / 255.0
    
    print(f"Input: torch.float32, range=[{input_tensor.min():.4f}, {input_tensor.max():.4f}]")
    print(f"Input scale: {input_scale:.6f}")
    
    # Try to quantize input with PyTorch
    try:
        x = quantize_tensor_pytorch(input_tensor, input_scale, 0, torch.qint8)
        using_pytorch_quant = True
        print("✅ Using PyTorch quantized tensors")
    except Exception:
        # Fall back to manual quantization
        x = input_tensor
        using_pytorch_quant = False
        print("📋 Using manual quantization simulation")
    
    # Execute nodes
    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 fast node {i}: {node_id} ({op_type})")
        
        if op_type == "linear" 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
            
            # Get quantization scales
            weight_key = f"{origin_name}.weight"
            bias_key = f"{origin_name}.bias"
            weight_scale = weight_quant.get(weight_key, {}).get("scale", 0.01)
            bias_scale = weight_quant.get(bias_key, {}).get("scale", 0.01)
            
            if using_pytorch_quant:
                try:
                    # Try to use PyTorch quantized linear
                    qlinear = create_quantized_linear(weight, bias, weight_scale, bias_scale)
                    if qlinear is not None:
                        # Dequantize input for PyTorch quantized layer
                        x_float = dequantize_tensor_pytorch(x, input_scale)
                        x_float = qlinear(x_float)
                        # Requantize output
                        output_scale = act_quant.get(node_id, input_scale)
                        x = quantize_tensor_pytorch(x_float, output_scale, 0, torch.qint8)
                        print(f"  ✅ PyTorch quantized linear, output scale: {output_scale:.6f}")
                        continue
                except Exception as e:
                    print(f"  ⚠️  PyTorch quantized linear failed: {e}")
            
            # Fall back to manual quantized operation
            if using_pytorch_quant:
                x = dequantize_tensor_pytorch(x, input_scale)
            
            # Quantize weights manually
            weight_q = torch.round(weight / weight_scale).clamp(-127, 127)
            weight_dq = weight_q * weight_scale
            
            if bias is not None:
                bias_q = torch.round(bias / bias_scale).clamp(-127, 127)
                bias_dq = bias_q * bias_scale
            else:
                bias_dq = None
            
            # Perform linear operation
            x = F.linear(x, weight_dq, bias_dq)
            
            # Quantize output for next layer
            if using_pytorch_quant:
                output_scale = act_quant.get(node_id, input_scale)
                x = quantize_tensor_pytorch(x, output_scale, 0, torch.qint8)
                input_scale = output_scale
            
            print(f"  Weight quantized with scale={weight_scale:.6f}")
            if bias is not None:
                print(f"  Bias quantized with scale={bias_scale:.6f}")
            
        elif op_type == "relu":
            if using_pytorch_quant:
                x_float = dequantize_tensor_pytorch(x, input_scale)
                x_float = F.relu(x_float)
                output_scale = act_quant.get(node_id, input_scale)
                x = quantize_tensor_pytorch(x_float, output_scale, 0, torch.qint8)
                input_scale = output_scale
            else:
                x = F.relu(x)
            print(f"  ReLU applied")
        
        # Print output info
        if using_pytorch_quant:
            x_float = dequantize_tensor_pytorch(x, input_scale)
            print(f"  Output: torch.qint8, scale={input_scale:.6f}, range=[{x_float.min():.4f}, {x_float.max():.4f}]")
        else:
            print(f"  Output: torch.float32, range=[{x.min():.4f}, {x.max():.4f}]")
    
    # Return final result as float
    if using_pytorch_quant:
        result = dequantize_tensor_pytorch(x, input_scale)
    else:
        result = x
    
    print(f"\n✅ Fast quantized result: {result}")
    return result