machine-dialect 0.1.0a1__py3-none-any.whl

machine_dialect/mir/optimization_pass.py

@@ -0,0 +1,251 @@
+"""Base classes and interfaces for MIR optimization and analysis passes.
+
+This module provides the foundation for all optimization and analysis passes
+in the MIR optimization framework.
+"""
+
+import abc
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.mir_module import MIRModule
+
+
+class PassType(Enum):
+    """Type of pass."""
+
+    ANALYSIS = "analysis"
+    OPTIMIZATION = "optimization"
+    CLEANUP = "cleanup"  # Cleanup passes that can run after optimizations
+    UTILITY = "utility"
+
+
+class PreservationLevel(Enum):
+    """Level of preservation for analyses."""
+
+    NONE = "none"  # Invalidates everything
+    CFG = "cfg"  # Preserves CFG structure
+    DOMINANCE = "dominance"  # Preserves dominance info
+    ALL = "all"  # Preserves all analyses
+
+
+@dataclass
+class PassInfo:
+    """Information about a pass.
+
+    Attributes:
+        name: Pass name.
+        description: Pass description.
+        pass_type: Type(s) of pass - can be single type or list of types.
+        requires: List of required analysis passes.
+        preserves: What analyses this pass preserves.
+    """
+
+    name: str
+    description: str
+    pass_type: PassType | list[PassType]
+    requires: list[str]
+    preserves: PreservationLevel
+
+
+class Pass(ABC):
+    """Base class for all passes."""
+
+    def __init__(self) -> None:
+        """Initialize the pass."""
+        self.stats: dict[str, int] = {}
+        self.debug_mode = False
+
+    @abstractmethod
+    def get_info(self) -> PassInfo:
+        """Get information about this pass.
+
+        Returns:
+            Pass information.
+        """
+        pass
+
+    def initialize(self) -> None:
+        """Initialize the pass before running.
+
+        Override this for pass-specific initialization.
+        """
+        self.stats.clear()
+
+    @abc.abstractmethod
+    def finalize(self) -> None:
+        """Finalize the pass after running.
+
+        Override this for pass-specific finalization.
+        """
+
+    def get_stats(self) -> dict[str, int]:
+        """Get pass statistics.
+
+        Returns:
+            Dictionary of statistics.
+        """
+        return self.stats.copy()
+
+
+class ModulePass(Pass):
+    """Base class for module-level passes."""
+
+    @abstractmethod
+    def run_on_module(self, module: MIRModule) -> bool:
+        """Run the pass on a module.
+
+        Args:
+            module: The module to process.
+
+        Returns:
+            True if the module was modified.
+        """
+        pass
+
+
+class FunctionPass(Pass):
+    """Base class for function-level passes."""
+
+    @abstractmethod
+    def run_on_function(self, function: MIRFunction) -> bool:
+        """Run the pass on a function.
+
+        Args:
+            function: The function to process.
+
+        Returns:
+            True if the function was modified.
+        """
+        pass
+
+    def run_on_module(self, module: MIRModule) -> bool:
+        """Run the pass on all functions in a module.
+
+        Args:
+            module: The module to process.
+
+        Returns:
+            True if any function was modified.
+        """
+        modified = False
+        for function in module.functions.values():
+            if self.run_on_function(function):
+                modified = True
+        return modified
+
+
+class AnalysisPass(Pass):
+    """Base class for analysis passes."""
+
+    def __init__(self) -> None:
+        """Initialize the analysis pass."""
+        super().__init__()
+        self._cache: dict[str, Any] = {}
+        self._valid = False
+
+    def invalidate(self) -> None:
+        """Invalidate the analysis cache."""
+        self._cache.clear()
+        self._valid = False
+
+    def is_valid(self) -> bool:
+        """Check if the analysis is valid.
+
+        Returns:
+            True if the analysis is valid.
+        """
+        return self._valid
+
+
+class ModuleAnalysisPass(AnalysisPass):
+    """Base class for module-level analysis passes."""
+
+    @abstractmethod
+    def run_on_module(self, module: MIRModule) -> Any:
+        """Run analysis on a module.
+
+        Args:
+            module: The module to analyze.
+
+        Returns:
+            Analysis results.
+        """
+        pass
+
+    def get_analysis(self, module: MIRModule) -> Any:
+        """Get cached analysis results or compute them.
+
+        Args:
+            module: The module to analyze.
+
+        Returns:
+            Analysis results.
+        """
+        if not self._valid:
+            result = self.run_on_module(module)
+            self._cache[module.name] = result
+            self._valid = True
+            return result
+        return self._cache.get(module.name)
+
+
+class FunctionAnalysisPass(AnalysisPass):
+    """Base class for function-level analysis passes."""
+
+    @abstractmethod
+    def run_on_function(self, function: MIRFunction) -> Any:
+        """Run analysis on a function.
+
+        Args:
+            function: The function to analyze.
+
+        Returns:
+            Analysis results.
+        """
+        pass
+
+    def get_analysis(self, function: MIRFunction) -> Any:
+        """Get cached analysis results or compute them.
+
+        Args:
+            function: The function to analyze.
+
+        Returns:
+            Analysis results.
+        """
+        if not self._valid or function.name not in self._cache:
+            result = self.run_on_function(function)
+            self._cache[function.name] = result
+            self._valid = True
+            return result
+        return self._cache[function.name]
+
+
+class OptimizationPass(FunctionPass):
+    """Base class for optimization passes."""
+
+    def __init__(self) -> None:
+        """Initialize the optimization pass."""
+        super().__init__()
+        self.analysis_manager: Any = None  # Set by pass manager
+
+    def get_analysis(self, analysis_name: str, function: MIRFunction) -> Any:
+        """Get analysis results from the analysis manager.
+
+        Args:
+            analysis_name: Name of the analysis.
+            function: The function to analyze.
+
+        Returns:
+            Analysis results.
+
+        Raises:
+            RuntimeError: If analysis manager is not set.
+        """
+        if self.analysis_manager is None:
+            raise RuntimeError("Analysis manager not set")
+        return self.analysis_manager.get_analysis(analysis_name, function)

machine_dialect/mir/optimization_pipeline.py

@@ -0,0 +1,355 @@
+"""Optimization pipeline configuration for MIR.
+
+This module defines optimization levels and pass pipelines for different
+compilation scenarios.
+"""
+
+from enum import Enum
+from typing import Any
+
+from machine_dialect.mir.analyses.dominance_analysis import DominanceAnalysis
+from machine_dialect.mir.analyses.loop_analysis import LoopAnalysis
+from machine_dialect.mir.analyses.use_def_chains import UseDefChainsAnalysis
+from machine_dialect.mir.mir_module import MIRModule
+from machine_dialect.mir.optimization_pass import ModulePass, Pass
+from machine_dialect.mir.optimizations.constant_propagation import ConstantPropagation
+from machine_dialect.mir.optimizations.cse import CommonSubexpressionElimination
+from machine_dialect.mir.optimizations.dce import DeadCodeElimination
+from machine_dialect.mir.optimizations.inlining import FunctionInlining
+from machine_dialect.mir.optimizations.jump_threading import JumpThreadingPass
+from machine_dialect.mir.optimizations.licm import LoopInvariantCodeMotion
+
+# from machine_dialect.mir.optimizations.peephole_optimizer import PeepholePass  # Disabled
+from machine_dialect.mir.optimizations.strength_reduction import StrengthReduction
+from machine_dialect.mir.pass_manager import PassManager
+
+
+class OptimizationLevel(Enum):
+    """Optimization levels for compilation."""
+
+    O0 = "O0"  # No optimization
+    O1 = "O1"  # Basic optimization
+    O2 = "O2"  # Standard optimization
+    O3 = "O3"  # Aggressive optimization
+    Os = "Os"  # Optimize for size
+
+
+class OptimizationPipeline:
+    """Manages optimization pipelines for different optimization levels."""
+
+    def __init__(self) -> None:
+        """Initialize the optimization pipeline.
+
+        Creates a pass manager and registers all available optimization
+        and analysis passes.
+        """
+        self.pass_manager = PassManager()
+        self._register_all_passes()
+        self.stats: dict[str, Any] = {}
+
+    def _register_all_passes(self) -> None:
+        """Register all available passes with the pass manager.
+
+        Registers both analysis passes (dominance, loop, use-def chains)
+        and optimization passes (constant propagation, DCE, CSE, etc.).
+        """
+        # Register analysis passes
+        self.pass_manager.register_pass(DominanceAnalysis)
+        self.pass_manager.register_pass(LoopAnalysis)
+        self.pass_manager.register_pass(UseDefChainsAnalysis)
+
+        # Register optimization passes
+        self.pass_manager.register_pass(ConstantPropagation)
+        self.pass_manager.register_pass(DeadCodeElimination)
+        self.pass_manager.register_pass(CommonSubexpressionElimination)
+        self.pass_manager.register_pass(StrengthReduction)
+        self.pass_manager.register_pass(JumpThreadingPass)
+        # self.pass_manager.register_pass(PeepholePass)  # Disabled
+        self.pass_manager.register_pass(LoopInvariantCodeMotion)
+        self.pass_manager.register_pass(FunctionInlining)
+
+    def get_passes_for_level(self, level: OptimizationLevel) -> list[Pass]:
+        """Get the list of passes for an optimization level.
+
+        Args:
+            level: The optimization level.
+
+        Returns:
+            List of pass instances to run.
+        """
+        passes: list[Pass] = []
+
+        if level == OptimizationLevel.O0:
+            # No optimization
+            return passes
+
+        elif level == OptimizationLevel.O1:
+            # Basic optimization
+            # Focus on simple, fast optimizations
+            o1_passes = [
+                self.pass_manager.registry.get_pass("constant-propagation"),
+                self.pass_manager.registry.get_pass("strength-reduction"),
+                self.pass_manager.registry.get_pass("dce"),
+                self.pass_manager.registry.get_pass("peephole"),
+            ]
+            passes.extend([p for p in o1_passes if p is not None])
+
+        elif level == OptimizationLevel.O2:
+            # Standard optimization
+            # Add more expensive optimizations
+            o2_passes: list[Pass | None] = [
+                # First pass: basic cleanup
+                self.pass_manager.registry.get_pass("constant-propagation"),
+                self.pass_manager.registry.get_pass("strength-reduction"),
+                self.pass_manager.registry.get_pass("dce"),
+                # Inlining (small functions only)
+                FunctionInlining(size_threshold=30),
+                # After inlining, more opportunities
+                self.pass_manager.registry.get_pass("constant-propagation"),
+                self.pass_manager.registry.get_pass("cse"),
+                self.pass_manager.registry.get_pass("licm"),
+                self.pass_manager.registry.get_pass("dce"),
+                # Final cleanup
+                self.pass_manager.registry.get_pass("jump-threading"),
+                self.pass_manager.registry.get_pass("peephole"),
+            ]
+            passes.extend([p for p in o2_passes if p is not None])
+
+        elif level == OptimizationLevel.O3:
+            # Aggressive optimization
+            # More aggressive thresholds and multiple iterations
+            o3_passes: list[Pass | None] = [
+                # First pass: aggressive inlining
+                FunctionInlining(size_threshold=100),
+                # Full optimization suite
+                self.pass_manager.registry.get_pass("constant-propagation"),
+                self.pass_manager.registry.get_pass("strength-reduction"),
+                self.pass_manager.registry.get_pass("cse"),
+                self.pass_manager.registry.get_pass("licm"),
+                self.pass_manager.registry.get_pass("dce"),
+                # Second iteration after first round
+                self.pass_manager.registry.get_pass("constant-propagation"),
+                self.pass_manager.registry.get_pass("cse"),
+                self.pass_manager.registry.get_pass("dce"),
+                # Control flow optimization
+                self.pass_manager.registry.get_pass("jump-threading"),
+                self.pass_manager.registry.get_pass("peephole"),
+                # Final DCE to clean up
+                self.pass_manager.registry.get_pass("dce"),
+            ]
+            passes.extend([p for p in o3_passes if p is not None])
+
+        elif level == OptimizationLevel.Os:
+            # Optimize for size
+            # Focus on reducing code size
+            os_passes = [
+                # No inlining (increases size)
+                self.pass_manager.registry.get_pass("constant-propagation"),
+                self.pass_manager.registry.get_pass("cse"),  # Reduces duplicate code
+                self.pass_manager.registry.get_pass("dce"),  # Removes dead code
+                self.pass_manager.registry.get_pass("jump-threading"),  # Simplifies control flow
+                self.pass_manager.registry.get_pass("peephole"),
+            ]
+            passes.extend([p for p in os_passes if p is not None])
+
+        # Already filtered out None passes
+        return passes
+
+    def optimize(self, module: MIRModule, level: OptimizationLevel = OptimizationLevel.O2) -> bool:
+        """Run optimization pipeline on a module.
+
+        Args:
+            module: The module to optimize.
+            level: The optimization level.
+
+        Returns:
+            True if the module was modified.
+        """
+        passes = self.get_passes_for_level(level)
+        modified = False
+
+        # Reset statistics
+        self.stats = {
+            "level": level.value,
+            "passes_run": [],
+            "total_modifications": 0,
+            "pass_stats": {},
+        }
+
+        for pass_instance in passes:
+            pass_info = pass_instance.get_info()
+            pass_name = pass_info.name
+
+            # Run the pass
+            pass_modified = False
+            if isinstance(pass_instance, ModulePass):
+                pass_modified = pass_instance.run_on_module(module)
+
+            # Track statistics
+            self.stats["passes_run"].append(pass_name)
+            if pass_modified:
+                modified = True
+                self.stats["total_modifications"] += 1
+
+            # Get pass-specific statistics if available
+            if hasattr(pass_instance, "get_statistics"):
+                self.stats["pass_stats"][pass_name] = pass_instance.get_statistics()
+
+        return modified
+
+    def optimize_with_custom_pipeline(self, module: MIRModule, pass_names: list[str]) -> bool:
+        """Run a custom optimization pipeline.
+
+        Args:
+            module: The module to optimize.
+            pass_names: List of pass names to run in order.
+
+        Returns:
+            True if the module was modified.
+        """
+        modified = False
+
+        for pass_name in pass_names:
+            pass_instance = self.pass_manager.registry.get_pass(pass_name)
+            if pass_instance:
+                if isinstance(pass_instance, ModulePass):
+                    if pass_instance.run_on_module(module):
+                        modified = True
+
+        return modified
+
+    def get_statistics(self) -> dict[str, Any]:
+        """Get statistics from the last optimization run.
+
+        Returns:
+            Dictionary containing optimization level, passes run,
+            total modifications made, and per-pass statistics.
+        """
+        return self.stats
+
+
+class PipelineBuilder:
+    """Builder for creating custom optimization pipelines."""
+
+    def __init__(self) -> None:
+        """Initialize the pipeline builder.
+
+        Creates an empty pipeline ready for pass configuration.
+        """
+        self.passes: list[str] = []
+        self.pass_configs: dict[str, dict[str, Any]] = {}
+
+    def add_pass(self, pass_name: str, **config: Any) -> "PipelineBuilder":
+        """Add a pass to the pipeline.
+
+        Args:
+            pass_name: Name of the pass.
+            **config: Configuration for the pass.
+
+        Returns:
+            Self for chaining.
+        """
+        self.passes.append(pass_name)
+        if config:
+            self.pass_configs[pass_name] = config
+        return self
+
+    def add_cleanup_passes(self) -> "PipelineBuilder":
+        """Add standard cleanup passes.
+
+        Returns:
+            Self for chaining.
+        """
+        self.passes.extend(["dce", "jump-threading", "peephole"])
+        return self
+
+    def add_algebraic_passes(self) -> "PipelineBuilder":
+        """Add algebraic optimization passes.
+
+        Returns:
+            Self for chaining.
+        """
+        self.passes.extend(["constant-propagation", "strength-reduction", "cse"])
+        return self
+
+    def add_loop_passes(self) -> "PipelineBuilder":
+        """Add loop optimization passes.
+
+        Returns:
+            Self for chaining.
+        """
+        self.passes.append("licm")
+        return self
+
+    def repeat(self, times: int = 2) -> "PipelineBuilder":
+        """Repeat the current pipeline multiple times.
+
+        Args:
+            times: Number of times to repeat.
+
+        Returns:
+            Self for chaining.
+        """
+        current_passes = self.passes.copy()
+        for _ in range(times - 1):
+            self.passes.extend(current_passes)
+        return self
+
+    def build(self) -> list[str]:
+        """Build the pipeline.
+
+        Returns:
+            List of pass names.
+        """
+        return self.passes.copy()
+
+
+# Convenience functions
+def create_o0_pipeline() -> OptimizationPipeline:
+    """Create a pipeline with no optimization.
+
+    Returns:
+        Pipeline configured for O0 (no optimization passes).
+    """
+    pipeline = OptimizationPipeline()
+    return pipeline
+
+
+def create_o1_pipeline() -> OptimizationPipeline:
+    """Create a pipeline with basic optimization.
+
+    Returns:
+        Pipeline configured for O1 (fast, simple optimizations).
+    """
+    pipeline = OptimizationPipeline()
+    return pipeline
+
+
+def create_o2_pipeline() -> OptimizationPipeline:
+    """Create a pipeline with standard optimization.
+
+    Returns:
+        Pipeline configured for O2 (balanced performance/compile time).
+    """
+    pipeline = OptimizationPipeline()
+    return pipeline
+
+
+def create_o3_pipeline() -> OptimizationPipeline:
+    """Create a pipeline with aggressive optimization.
+
+    Returns:
+        Pipeline configured for O3 (maximum performance optimization).
+    """
+    pipeline = OptimizationPipeline()
+    return pipeline
+
+
+def create_size_pipeline() -> OptimizationPipeline:
+    """Create a pipeline optimized for code size.
+
+    Returns:
+        Pipeline configured for Os (minimize code size).
+    """
+    pipeline = OptimizationPipeline()
+    return pipeline

machine_dialect/mir/optimizations/__init__.py

@@ -0,0 +1,84 @@
+"""MIR optimization passes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from machine_dialect.mir.optimizations.algebraic_simplification import AlgebraicSimplification
+from machine_dialect.mir.optimizations.branch_prediction import BranchPredictionOptimization
+from machine_dialect.mir.optimizations.constant_propagation import ConstantPropagation
+from machine_dialect.mir.optimizations.cse import CommonSubexpressionElimination
+from machine_dialect.mir.optimizations.dce import DeadCodeElimination
+from machine_dialect.mir.optimizations.inlining import FunctionInlining
+from machine_dialect.mir.optimizations.jump_threading import JumpThreadingOptimizer, JumpThreadingPass
+from machine_dialect.mir.optimizations.licm import LoopInvariantCodeMotion
+from machine_dialect.mir.optimizations.loop_unrolling import LoopUnrolling
+
+# from machine_dialect.mir.optimizations.peephole_optimizer import PeepholeOptimizer  # Disabled - needs update
+from machine_dialect.mir.optimizations.strength_reduction import StrengthReduction
+from machine_dialect.mir.optimizations.tail_call import TailCallOptimization
+from machine_dialect.mir.optimizations.type_narrowing import TypeNarrowing
+from machine_dialect.mir.optimizations.type_specialization import TypeSpecialization
+from machine_dialect.mir.optimizations.type_specific import TypeSpecificOptimization
+
+if TYPE_CHECKING:
+    from machine_dialect.mir.pass_manager import PassManager
+
+__all__ = [
+    "AlgebraicSimplification",
+    "BranchPredictionOptimization",
+    "CommonSubexpressionElimination",
+    "ConstantPropagation",
+    "DeadCodeElimination",
+    "FunctionInlining",
+    "JumpThreadingOptimizer",
+    "JumpThreadingPass",
+    "LoopInvariantCodeMotion",
+    "LoopUnrolling",
+    # "PeepholeOptimizer",  # Disabled
+    # "PeepholePass",  # Disabled
+    "StrengthReduction",
+    "TailCallOptimization",
+    "TypeNarrowing",
+    "TypeSpecialization",
+    "TypeSpecificOptimization",
+]
+
+
+def register_all_passes(pass_manager: PassManager) -> None:
+    """Register all optimization passes with the pass manager.
+
+    Args:
+        pass_manager: Pass manager to register with.
+    """
+    from machine_dialect.mir.analyses.alias_analysis import AliasAnalysis
+    from machine_dialect.mir.analyses.dominance_analysis import DominanceAnalysis
+    from machine_dialect.mir.analyses.escape_analysis import EscapeAnalysis
+    from machine_dialect.mir.analyses.loop_analysis import LoopAnalysis
+    from machine_dialect.mir.analyses.type_analysis import TypeAnalysis
+    from machine_dialect.mir.analyses.use_def_chains import UseDefChainsAnalysis
+
+    # Register analysis passes
+    pass_manager.register_pass(DominanceAnalysis)
+    pass_manager.register_pass(UseDefChainsAnalysis)
+    pass_manager.register_pass(LoopAnalysis)
+    pass_manager.register_pass(AliasAnalysis)
+    pass_manager.register_pass(EscapeAnalysis)
+    pass_manager.register_pass(TypeAnalysis)
+
+    # Register optimization passes
+    pass_manager.register_pass(TypeSpecificOptimization)  # Run early to benefit other passes
+    pass_manager.register_pass(TypeNarrowing)  # Run after type-specific to narrow union types
+    pass_manager.register_pass(ConstantPropagation)
+    pass_manager.register_pass(CommonSubexpressionElimination)
+    pass_manager.register_pass(DeadCodeElimination)
+    pass_manager.register_pass(StrengthReduction)
+    pass_manager.register_pass(AlgebraicSimplification)
+    pass_manager.register_pass(TailCallOptimization)
+    pass_manager.register_pass(TypeSpecialization)
+    pass_manager.register_pass(FunctionInlining)
+    pass_manager.register_pass(LoopInvariantCodeMotion)
+    pass_manager.register_pass(LoopUnrolling)
+    pass_manager.register_pass(BranchPredictionOptimization)
+    pass_manager.register_pass(JumpThreadingPass)
+    # pass_manager.register_pass(PeepholePass)  # Disabled - needs update for new opcodes