enhanced_fusion.py

#!/usr/bin/env python3
"""
Enhanced Fusion Pass for ITC Framework
Supports Conv2D+ReLU, Conv2D+BatchNorm+ReLU, and other fused operations
"""

import torch
import torch.nn as nn
from typing import Dict, Any, List, Tuple, Optional

class EnhancedFusionPass:
    """Enhanced fusion pass supporting multiple fusion patterns"""
    
    @staticmethod
    def find_conv_relu_patterns(nodes: List[Dict[str, Any]]) -> List[Tuple[int, int]]:
        """Find Conv2D + ReLU fusion patterns"""
        patterns = []
        
        for i in range(len(nodes) - 1):
            conv_node = nodes[i]
            relu_node = nodes[i + 1]
            
            if (conv_node['op_type'] == 'conv2d' and 
                relu_node['op_type'] == 'relu'):
                
                # Check if they're connected
                conv_outputs = conv_node.get('outputs', [])
                relu_inputs = relu_node.get('inputs', [])
                
                if len(set(conv_outputs) & set(relu_inputs)) > 0:
                    patterns.append((i, i + 1))
        
        return patterns
    
    @staticmethod
    def find_conv_bn_relu_patterns(nodes: List[Dict[str, Any]]) -> List[Tuple[int, int, int]]:
        """Find Conv2D + BatchNorm + ReLU fusion patterns"""
        patterns = []
        
        for i in range(len(nodes) - 2):
            conv_node = nodes[i]
            bn_node = nodes[i + 1]
            relu_node = nodes[i + 2]
            
            if (conv_node['op_type'] == 'conv2d' and 
                bn_node['op_type'] == 'batchnorm2d' and
                relu_node['op_type'] == 'relu'):
                
                # Check connections
                conv_outputs = conv_node.get('outputs', [])
                bn_inputs = bn_node.get('inputs', [])
                bn_outputs = bn_node.get('outputs', [])
                relu_inputs = relu_node.get('inputs', [])
                
                conv_to_bn = len(set(conv_outputs) & set(bn_inputs)) > 0
                bn_to_relu = len(set(bn_outputs) & set(relu_inputs)) > 0
                
                if conv_to_bn and bn_to_relu:
                    patterns.append((i, i + 1, i + 2))
        
        return patterns
    
    @staticmethod
    def find_linear_relu_patterns(nodes: List[Dict[str, Any]]) -> List[Tuple[int, int]]:
        """Find Linear + ReLU fusion patterns (existing functionality)"""
        patterns = []
        
        for i in range(len(nodes) - 1):
            linear_node = nodes[i]
            relu_node = nodes[i + 1]
            
            if (linear_node['op_type'] == 'linear' and 
                relu_node['op_type'] == 'relu'):
                
                # Check if they're connected
                linear_outputs = linear_node.get('outputs', [])
                relu_inputs = relu_node.get('inputs', [])
                
                if len(set(linear_outputs) & set(relu_inputs)) > 0:
                    patterns.append((i, i + 1))
        
        return patterns
    
    @staticmethod
    def apply_conv_relu_fusion(nodes: List[Dict[str, Any]], conv_idx: int, relu_idx: int) -> bool:
        """Apply Conv2D + ReLU fusion"""
        try:
            conv_node = nodes[conv_idx]
            relu_node = nodes[relu_idx]
            
            # Create fused node
            fused_node = conv_node.copy()
            fused_node['op_type'] = 'conv2d_relu_fused'
            fused_node['id'] = f"{conv_node['id']}_fused_relu"
            fused_node['outputs'] = relu_node.get('outputs', [])
            
            # Preserve all weight information
            if 'weights' in conv_node:
                fused_node['weights'] = conv_node['weights'].copy()
            
            # Preserve quantization information
            if 'per_channel_quantization' in conv_node:
                fused_node['per_channel_quantization'] = conv_node['per_channel_quantization']
            if 'activation_quantization' in conv_node:
                fused_node['activation_quantization'] = conv_node['activation_quantization']
            
            # Add fusion metadata
            fused_node['fusion_info'] = {
                'type': 'conv2d_relu',
                'original_nodes': [conv_node['id'], relu_node['id']],
                'activation': 'relu'
            }
            
            # Replace conv node with fused node
            nodes[conv_idx] = fused_node
            
            return True
            
        except Exception as e:
            print(f"      ❌ Conv2D+ReLU fusion failed: {e}")
            return False
    
    @staticmethod
    def apply_conv_bn_relu_fusion(nodes: List[Dict[str, Any]], conv_idx: int, bn_idx: int, relu_idx: int) -> bool:
        """Apply Conv2D + BatchNorm + ReLU fusion"""
        try:
            conv_node = nodes[conv_idx]
            bn_node = nodes[bn_idx]
            relu_node = nodes[relu_idx]
            
            # Create fused node
            fused_node = conv_node.copy()
            fused_node['op_type'] = 'conv2d_bn_relu_fused'
            fused_node['id'] = f"{conv_node['id']}_fused_bn_relu"
            fused_node['outputs'] = relu_node.get('outputs', [])
            
            # Preserve BatchNorm parameters for folding
            if 'weights' not in fused_node:
                fused_node['weights'] = {}
            
            # Copy BatchNorm parameters
            if 'weights' in bn_node:
                fused_node['weights'].update({
                    f"bn_{k}": v for k, v in bn_node['weights'].items()
                })
            
            # Add fusion metadata
            fused_node['fusion_info'] = {
                'type': 'conv2d_bn_relu',
                'original_nodes': [conv_node['id'], bn_node['id'], relu_node['id']],
                'batch_norm': True,
                'activation': 'relu'
            }
            
            # Replace conv node with fused node
            nodes[conv_idx] = fused_node
            
            return True
            
        except Exception as e:
            print(f"      ❌ Conv2D+BN+ReLU fusion failed: {e}")
            return False
    
    @staticmethod
    def apply_linear_relu_fusion(nodes: List[Dict[str, Any]], linear_idx: int, relu_idx: int) -> bool:
        """Apply Linear + ReLU fusion (existing functionality)"""
        try:
            linear_node = nodes[linear_idx]
            relu_node = nodes[relu_idx]
            
            # Create fused node
            fused_node = linear_node.copy()
            fused_node['op_type'] = 'linear_relu_fused'
            fused_node['id'] = f"{linear_node['id']}_fused_relu"
            fused_node['outputs'] = relu_node.get('outputs', [])
            
            # Preserve all weight information
            if 'weights' in linear_node:
                fused_node['weights'] = linear_node['weights'].copy()
            
            # Preserve quantization information
            if 'per_channel_quantization' in linear_node:
                fused_node['per_channel_quantization'] = linear_node['per_channel_quantization']
            if 'activation_quantization' in linear_node:
                fused_node['activation_quantization'] = linear_node['activation_quantization']
            
            # Add fusion metadata
            fused_node['fusion_info'] = {
                'type': 'linear_relu',
                'original_nodes': [linear_node['id'], relu_node['id']],
                'activation': 'relu'
            }
            
            # Replace linear node with fused node
            nodes[linear_idx] = fused_node
            
            return True
            
        except Exception as e:
            print(f"      ❌ Linear+ReLU fusion failed: {e}")
            return False

def enhanced_fusion_pass(mir: Dict[str, Any]) -> Dict[str, Any]:
    """
    Apply enhanced fusion optimization pass
    
    Supports:
    - Conv2D + ReLU
    - Conv2D + BatchNorm + ReLU  
    - Linear + ReLU
    
    Args:
        mir: MIR graph dictionary
        
    Returns:
        Optimized MIR graph with fused operations
    """
    
    print("🔧 Applying enhanced fusion optimizations...")
    
    mir_copy = mir.copy()
    nodes = mir_copy['graph']['nodes']
    original_count = len(nodes)
    
    total_fusions = 0
    indices_to_remove = []
    
    # 1. Find and apply Conv2D + BatchNorm + ReLU fusions (3-way)
    conv_bn_relu_patterns = EnhancedFusionPass.find_conv_bn_relu_patterns(nodes)
    
    for conv_idx, bn_idx, relu_idx in conv_bn_relu_patterns:
        if EnhancedFusionPass.apply_conv_bn_relu_fusion(nodes, conv_idx, bn_idx, relu_idx):
            print(f"  ✅ Fused conv2d + batchnorm + relu → {nodes[conv_idx]['id']}")
            indices_to_remove.extend([bn_idx, relu_idx])
            total_fusions += 1
    
    # 2. Find and apply Conv2D + ReLU fusions (2-way) 
    conv_relu_patterns = EnhancedFusionPass.find_conv_relu_patterns(nodes)
    
    for conv_idx, relu_idx in conv_relu_patterns:
        # Skip if already part of a 3-way fusion
        if relu_idx not in indices_to_remove:
            if EnhancedFusionPass.apply_conv_relu_fusion(nodes, conv_idx, relu_idx):
                print(f"  ✅ Fused conv2d + relu → {nodes[conv_idx]['id']}")
                indices_to_remove.append(relu_idx)
                total_fusions += 1
    
    # 3. Find and apply Linear + ReLU fusions (existing)
    linear_relu_patterns = EnhancedFusionPass.find_linear_relu_patterns(nodes)
    
    for linear_idx, relu_idx in linear_relu_patterns:
        # Skip if already part of another fusion
        if relu_idx not in indices_to_remove:
            if EnhancedFusionPass.apply_linear_relu_fusion(nodes, linear_idx, relu_idx):
                print(f"  ✅ Fused linear + relu → {nodes[linear_idx]['id']}")
                indices_to_remove.append(relu_idx)
                total_fusions += 1
    
    # Remove fused nodes (in reverse order to maintain indices)
    for idx in sorted(set(indices_to_remove), reverse=True):
        removed_node = nodes.pop(idx)
        print(f"   Removed fused node: {removed_node['id']}")
    
    final_count = len(nodes)
    
    if total_fusions > 0:
        print(f"✅ Enhanced fusion complete: {total_fusions} fusions applied")
        print(f"   Nodes reduced: {original_count} → {final_count} ({original_count - final_count} fewer)")
    else:
        print("ℹ️  No fusion opportunities found")
    
    return mir_copy

if __name__ == "__main__":
    # Test enhanced fusion
    import sys
    import os
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    
    from parser.pytorch_parser import parse_pytorch_model
    from passes.lowerer import hir_to_mir
    
    # Create test model with various fusion patterns
    class TestFusionCNN(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
            self.bn1 = nn.BatchNorm2d(32)
            self.relu1 = nn.ReLU()
            
            self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
            self.relu2 = nn.ReLU()
            
            self.fc = nn.Linear(64, 10)
            self.relu3 = nn.ReLU()
            
        def forward(self, x):
            # Conv2D + BN + ReLU pattern
            x = self.conv1(x)
            x = self.bn1(x)
            x = self.relu1(x)
            
            # Conv2D + ReLU pattern
            x = self.conv2(x)
            x = self.relu2(x)
            
            # Global average pooling
            x = torch.mean(x, dim=[2, 3])
            
            # Linear + ReLU pattern
            x = self.fc(x)
            x = self.relu3(x)
            
            return x
    
    model = TestFusionCNN()
    input_tensor = torch.randn(1, 3, 32, 32)
    
    # Parse and convert to MIR
    hir = parse_pytorch_model(model, input_tensor)
    mir = hir_to_mir(hir.to_dict())
    
    print(f"Original MIR nodes: {len(mir['graph']['nodes'])}")
    for node in mir['graph']['nodes']:
        print(f"  {node['id']}: {node['op_type']}")
    
    # Apply enhanced fusion
    fused_mir = enhanced_fusion_pass(mir)
    
    print(f"\nFused MIR nodes: {len(fused_mir['graph']['nodes'])}")
    for node in fused_mir['graph']['nodes']:
        print(f"  {node['id']}: {node['op_type']}")