fusion.py

"""
Operator fusion pass for performance optimization
Combines consecutive operations into fused kernels
"""
import torch
import torch.nn.functional as F
from typing import Dict, Any, List, Tuple

def can_fuse_linear_relu(nodes: List[Dict], start_idx: int) -> bool:
    """Check if Linear + ReLU can be fused at the given position"""
    if start_idx + 1 >= len(nodes):
        return False
    
    linear_node = nodes[start_idx]
    relu_node = nodes[start_idx + 1]
    
    # Check if first node is linear and second is relu
    if linear_node["op_type"] != "linear" or relu_node["op_type"] != "relu":
        return False
    
    # Check if relu input matches linear output
    linear_output = linear_node["outputs"][0] if linear_node["outputs"] else f"t_{start_idx}"
    relu_input = relu_node["inputs"][0] if relu_node["inputs"] else f"t_{start_idx}"
    
    return linear_output == relu_input

def can_fuse_conv_relu(nodes: List[Dict], start_idx: int) -> bool:
    """Check if Conv2D + ReLU can be fused at the given position"""
    if start_idx + 1 >= len(nodes):
        return False
    
    conv_node = nodes[start_idx]
    relu_node = nodes[start_idx + 1]
    
    # Check if first node is conv and second is relu
    if conv_node["op_type"] != "conv2d" or relu_node["op_type"] != "relu":
        return False
    
    # Check if relu input matches conv output
    conv_output = conv_node["outputs"][0] if conv_node["outputs"] else f"t_{start_idx}"
    relu_input = relu_node["inputs"][0] if relu_node["inputs"] else f"t_{start_idx}"
    
    return conv_output == relu_input

def create_fused_linear_relu(linear_node: Dict, relu_node: Dict) -> Dict:
    """Create a fused Linear+ReLU node"""
    fused_node = {
        "id": f"{linear_node['id']}_fused_relu",
        "op_type": "linear_relu_fused",
        "inputs": linear_node["inputs"],
        "outputs": relu_node["outputs"],
        "params": linear_node["params"].copy(),
        "compression": linear_node.get("compression", {}),
        "hints": {
            "fused_operations": ["linear", "relu"],
            "original_nodes": [linear_node["id"], relu_node["id"]]
        },
        "metadata": {
            "origin_name": linear_node.get("metadata", {}).get("origin_name"),
            "fused_from": [linear_node["id"], relu_node["id"]],
            "fusion_type": "linear_relu"
        }
    }
    return fused_node

def create_fused_conv_relu(conv_node: Dict, relu_node: Dict) -> Dict:
    """Create a fused Conv2D+ReLU node"""
    fused_node = {
        "id": f"{conv_node['id']}_fused_relu",
        "op_type": "conv2d_relu_fused",
        "inputs": conv_node["inputs"],
        "outputs": relu_node["outputs"],
        "params": conv_node["params"].copy(),
        "compression": conv_node.get("compression", {}),
        "hints": {
            "fused_operations": ["conv2d", "relu"],
            "original_nodes": [conv_node["id"], relu_node["id"]]
        },
        "metadata": {
            "origin_name": conv_node.get("metadata", {}).get("origin_name"),
            "fused_from": [conv_node["id"], relu_node["id"]],
            "fusion_type": "conv2d_relu"
        }
    }
    return fused_node

def fusion_pass(mir: Dict[str, Any]) -> Dict[str, Any]:
    """
    Apply operator fusion optimizations to the MIR.
    
    Currently supports:
    - Linear + ReLU → FusedLinearReLU
    - Conv2D + ReLU → FusedConv2DReLU
    
    Args:
        mir: MIR dictionary to be optimized
    
    Returns:
        Updated MIR with fused operations
    """
    print("🔧 Applying operator fusion optimizations...")
    
    nodes = mir["graph"]["nodes"]
    fused_nodes = []
    fusions_applied = 0
    
    i = 0
    while i < len(nodes):
        # Try Linear + ReLU fusion
        if can_fuse_linear_relu(nodes, i):
            fused_node = create_fused_linear_relu(nodes[i], nodes[i + 1])
            fused_nodes.append(fused_node)
            print(f"  ✅ Fused {nodes[i]['id']} + {nodes[i + 1]['id']} → {fused_node['id']}")
            fusions_applied += 1
            i += 2  # Skip both nodes
        
        # Try Conv2D + ReLU fusion
        elif can_fuse_conv_relu(nodes, i):
            fused_node = create_fused_conv_relu(nodes[i], nodes[i + 1])
            fused_nodes.append(fused_node)
            print(f"  ✅ Fused {nodes[i]['id']} + {nodes[i + 1]['id']} → {fused_node['id']}")
            fusions_applied += 1
            i += 2  # Skip both nodes
        
        else:
            # No fusion possible, keep original node
            fused_nodes.append(nodes[i])
            i += 1
    
    # Update MIR with fused nodes
    mir["graph"]["nodes"] = fused_nodes
    
    # Add fusion metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    
    mir["metadata"]["fusion_applied"] = True
    mir["metadata"]["fusions_count"] = fusions_applied
    mir["metadata"]["original_nodes_count"] = len(nodes)
    mir["metadata"]["fused_nodes_count"] = len(fused_nodes)
    
    print(f"✅ Fusion complete: {fusions_applied} fusions applied")
    print(f"   Nodes reduced: {len(nodes)} → {len(fused_nodes)} ({len(nodes) - len(fused_nodes)} fewer)")
    
    return mir

def execute_fused_linear_relu(input_tensor: torch.Tensor, weight: torch.Tensor, 
                             bias: torch.Tensor, weight_scales: torch.Tensor = None, 
                             bias_scale: float = None, per_channel: bool = False) -> torch.Tensor:
    """
    Execute fused Linear + ReLU operation with quantization.
    
    Args:
        input_tensor: Input tensor
        weight: Weight tensor
        bias: Bias tensor (can be None)
        weight_scales: Quantization scales (per-channel or per-tensor)
        bias_scale: Bias quantization scale
        per_channel: Whether to use per-channel quantization
    
    Returns:
        Output tensor after fused linear+relu operation
    """
    if weight_scales is not None and per_channel:
        # Use per-channel quantized linear
        from passes.per_channel_quantize import linear_per_channel_quantized
        output = linear_per_channel_quantized(input_tensor, weight, bias, weight_scales, bias_scale)
    elif weight_scales is not None:
        # Use per-tensor quantized linear
        weight_scale = weight_scales if isinstance(weight_scales, float) else weight_scales[0]
        weight_quantized = torch.round(weight / weight_scale).clamp(-127, 127)
        weight_dequantized = weight_quantized * weight_scale
        
        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
        
        output = F.linear(input_tensor, weight_dequantized, bias_dequantized)
    else:
        # Standard floating-point linear
        output = F.linear(input_tensor, weight, bias)
    
    # Apply ReLU activation
    output = F.relu(output)
    
    return output

def execute_fused_conv_relu(input_tensor: torch.Tensor, weight: torch.Tensor, 
                           bias: torch.Tensor, stride: int = 1, padding: int = 0,
                           weight_scales: torch.Tensor = None, bias_scale: float = None, 
                           per_channel: bool = False) -> torch.Tensor:
    """
    Execute fused Conv2D + ReLU operation with quantization.
    
    Args:
        input_tensor: Input tensor
        weight: Conv2D weight tensor
        bias: Conv2D bias tensor (can be None)
        stride: Convolution stride
        padding: Convolution padding
        weight_scales: Quantization scales (per-channel or per-tensor)
        bias_scale: Bias quantization scale
        per_channel: Whether to use per-channel quantization
    
    Returns:
        Output tensor after fused conv2d+relu operation
    """
    if weight_scales is not None and per_channel:
        # Use per-channel quantized conv2d
        from passes.per_channel_quantize import quantize_per_channel, dequantize_per_channel
        
        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 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
        
        output = F.conv2d(input_tensor, weight_dequantized, bias_dequantized, stride, padding)
    
    elif weight_scales is not None:
        # Use per-tensor quantized conv2d
        weight_scale = weight_scales if isinstance(weight_scales, float) else weight_scales[0]
        weight_quantized = torch.round(weight / weight_scale).clamp(-127, 127)
        weight_dequantized = weight_quantized * weight_scale
        
        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
        
        output = F.conv2d(input_tensor, weight_dequantized, bias_dequantized, stride, padding)
    else:
        # Standard floating-point conv2d
        output = F.conv2d(input_tensor, weight, bias, stride, padding)
    
    # Apply ReLU activation
    output = F.relu(output)
    
    return output