per_channel_quantize.py

"""
Per-channel quantization for improved accuracy
Computes separate scales for each output channel of weights
"""
import torch
from typing import Dict, Any, Tuple

def compute_per_channel_scales(weight: torch.Tensor, num_bits: int = 8) -> torch.Tensor:
    """
    Compute per-channel quantization scales for weight tensor.
    
    Args:
        weight: Weight tensor of shape (out_channels, in_channels, ...)
        num_bits: Number of quantization bits
    
    Returns:
        scales: Tensor of shape (out_channels,) with per-channel scales
    """
    # Reshape weight to (out_channels, -1) for per-channel analysis
    out_channels = weight.shape[0]
    weight_reshaped = weight.view(out_channels, -1)
    
    # Compute max absolute value per output channel
    max_vals = torch.abs(weight_reshaped).max(dim=1)[0]
    
    # Avoid division by zero
    max_vals = torch.clamp(max_vals, min=1e-8)
    
    # Compute scale: max_val / (2^(bits-1) - 1)
    max_int = (2 ** (num_bits - 1)) - 1
    scales = max_vals / max_int
    
    return scales

def quantize_per_channel(weight: torch.Tensor, scales: torch.Tensor, 
                        num_bits: int = 8) -> torch.Tensor:
    """
    Quantize weight tensor using per-channel scales.
    
    Args:
        weight: Weight tensor of shape (out_channels, ...)
        scales: Per-channel scales of shape (out_channels,)
        num_bits: Number of quantization bits
    
    Returns:
        quantized_weight: Quantized weight tensor
    """
    out_channels = weight.shape[0]
    
    # Reshape for broadcasting
    scales_broadcast = scales.view(out_channels, *([1] * (weight.dim() - 1)))
    
    # Quantize: weight / scale, round, clamp
    qmin = -(2 ** (num_bits - 1))
    qmax = (2 ** (num_bits - 1)) - 1
    
    quantized = (weight / scales_broadcast).round().clamp(qmin, qmax)
    
    return quantized

def dequantize_per_channel(quantized_weight: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    """
    Dequantize weight tensor using per-channel scales.
    
    Args:
        quantized_weight: Quantized weight tensor of shape (out_channels, ...)
        scales: Per-channel scales of shape (out_channels,)
    
    Returns:
        dequantized_weight: Dequantized weight tensor
    """
    out_channels = quantized_weight.shape[0]
    
    # Reshape for broadcasting
    scales_broadcast = scales.view(out_channels, *([1] * (quantized_weight.dim() - 1)))
    
    # Dequantize: quantized * scale
    dequantized = quantized_weight * scales_broadcast
    
    return dequantized

def per_channel_quantization_pass(mir: Dict[str, Any], model: torch.nn.Module, 
                                 num_bits: int = 8) -> Dict[str, Any]:
    """
    Enhanced quantization pass using per-channel quantization for weights.
    
    Args:
        mir: MIR dictionary to be modified
        model: Original PyTorch model
        num_bits: Number of quantization bits
    
    Returns:
        Updated MIR with per-channel quantization metadata and embedded weights
    """
    print(f"🔧 Applying per-channel quantization ({num_bits}-bit)...")
    
    # Ensure weights are embedded first
    if not _weights_embedded_in_mir(mir):
        from passes.quantize import _embed_weights_in_mir_nodes
        _embed_weights_in_mir_nodes(mir, model)
    
    state_dict = model.state_dict()
    per_channel_meta = {}
    quantized_layers = 0
    
    for name, tensor in state_dict.items():
        if not isinstance(tensor, torch.Tensor):
            continue
            
        # Apply per-channel quantization to weight tensors
        if 'weight' in name and tensor.dim() >= 2:
            # Compute per-channel scales
            scales = compute_per_channel_scales(tensor, num_bits)
            
            # Store metadata
            per_channel_meta[name] = {
                "quantized": True,
                "bits": num_bits,
                "scales": scales.tolist(),  # Convert to list for JSON serialization
                "dtype": "int8",
                "per_channel": True,
                "channels": tensor.shape[0]
            }
            
            print(f"  {name}: {tensor.shape[0]} channels, scale range [{scales.min():.6f}, {scales.max():.6f}]")
            quantized_layers += 1
            
        # Use per-tensor quantization for biases and other parameters
        elif tensor.numel() > 0:
            max_val = tensor.abs().max().item()
            if max_val == 0:
                scale = 1.0
            else:
                scale = max_val / ((2 ** (num_bits - 1)) - 1)
            
            per_channel_meta[name] = {
                "quantized": True,
                "bits": num_bits,
                "scale": float(scale),
                "dtype": "int8",
                "per_channel": False
            }
            
            print(f"  {name}: per-tensor scale={scale:.6f}")
    
    # Update MIR nodes with per-channel quantization info
    for node in mir['graph']['nodes']:
        node_id = node['id']
        op_type = node['op_type']
        
        if op_type in ['conv2d', 'linear'] and 'weights' in node and 'weight' in node['weights']:
            # Find corresponding weight tensor
            possible_names = [node_id, node_id.replace('_', '.')]
            
            for base_name in possible_names:
                weight_key = f"{base_name}.weight"
                if weight_key in per_channel_meta and per_channel_meta[weight_key].get('per_channel', False):
                    node['per_channel_quantization'] = {
                        'enabled': True,
                        'num_bits': num_bits,
                        'channel_scales': per_channel_meta[weight_key]['scales'],
                        'quantization_axis': 0
                    }
                    break
    
    # Update MIR metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    
    mir["metadata"]["per_channel_quantization"] = per_channel_meta
    mir["metadata"]["per_channel_num_bits"] = num_bits
    
    print(f"✅ Per-channel quantization applied to {quantized_layers} weight tensors")
    
    return mir

def _weights_embedded_in_mir(mir: Dict[str, Any]) -> bool:
    """Check if weights are already embedded in MIR nodes"""
    for node in mir['graph']['nodes']:
        if node['op_type'] in ['conv2d', 'linear'] and 'weights' in node:
            return True
    return False

def linear_per_channel_quantized(input_tensor: torch.Tensor, weight: torch.Tensor, 
                                bias: torch.Tensor, weight_scales: torch.Tensor, 
                                bias_scale: float = None) -> torch.Tensor:
    """
    Perform quantized linear operation using per-channel quantized weights.
    
    Args:
        input_tensor: Input tensor
        weight: Original weight tensor
        bias: Original bias tensor (can be None)
        weight_scales: Per-channel scales for weight
        bias_scale: Scale for bias quantization
    
    Returns:
        Output tensor after quantized linear operation
    """
    # Quantize weights per-channel
    weight_quantized = quantize_per_channel(weight, weight_scales)
    weight_dequantized = dequantize_per_channel(weight_quantized, weight_scales)
    
    # Quantize bias per-tensor if present
    bias_dequantized = None
    if bias is not None and bias_scale is not None:
        bias_quantized = torch.round(bias / bias_scale).clamp(-127, 127)
        bias_dequantized = bias_quantized * bias_scale
    elif bias is not None:
        bias_dequantized = bias
    
    # Perform linear operation with quantized weights
    output = torch.nn.functional.linear(input_tensor, weight_dequantized, bias_dequantized)
    
    return output