machine-dialect 0.1.0a1__py3-none-any.whl

machine_dialect/mir/register_allocation.py

@@ -0,0 +1,302 @@
+"""Virtual register allocation for MIR.
+
+This module implements register allocation using linear scan algorithm.
+"""
+
+from dataclasses import dataclass
+
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.mir_instructions import MIRInstruction
+from machine_dialect.mir.mir_values import MIRValue, Temp, Variable
+
+
+@dataclass
+class LiveInterval:
+    """Represents the live interval of a value.
+
+    Attributes:
+        value: The MIR value.
+        start: Start position of the interval.
+        end: End position of the interval.
+        register: Allocated register/slot number.
+    """
+
+    value: MIRValue
+    start: int
+    end: int
+    register: int | None = None
+
+
+@dataclass
+class RegisterAllocation:
+    """Result of register allocation.
+
+    Attributes:
+        allocations: Mapping from MIR values to register numbers.
+        spilled_values: Set of values that need to be spilled to memory.
+        max_registers: Maximum number of registers used.
+    """
+
+    allocations: dict[MIRValue, int]
+    spilled_values: set[MIRValue]
+    max_registers: int
+
+
+class RegisterAllocator:
+    """Allocates virtual registers for MIR values using linear scan."""
+
+    def __init__(self, function: MIRFunction, max_registers: int = 256) -> None:
+        """Initialize the register allocator.
+
+        Args:
+            function: The MIR function to allocate registers for.
+            max_registers: Maximum number of available registers.
+        """
+        self.function = function
+        self.max_registers = max_registers
+        self.live_intervals: list[LiveInterval] = []
+        self.active_intervals: list[LiveInterval] = []
+        self.free_registers: list[int] = list(range(max_registers))
+        self.instruction_positions: dict[MIRInstruction, int] = {}
+        self.spilled_values: set[MIRValue] = set()
+        self.next_spill_slot = 0  # Track spill slot allocation
+
+    def allocate(self) -> RegisterAllocation:
+        """Perform register allocation.
+
+        Returns:
+            The register allocation result.
+        """
+        # Build instruction positions
+        self._build_instruction_positions()
+
+        # Compute live intervals
+        self._compute_live_intervals()
+
+        # Sort intervals by start position
+        self.live_intervals.sort(key=lambda x: x.start)
+
+        # Perform linear scan allocation
+        allocations = self._linear_scan()
+
+        # Calculate the actual number of registers used
+        max_reg_used = 0
+        for reg in allocations.values():
+            if reg >= 0:  # Only count actual registers, not spill slots
+                max_reg_used = max(max_reg_used, reg + 1)
+
+        return RegisterAllocation(
+            allocations=allocations, spilled_values=self.spilled_values, max_registers=max_reg_used
+        )
+
+    def _build_instruction_positions(self) -> None:
+        """Build a mapping from instructions to positions."""
+        position = 0
+        for block in self.function.cfg.blocks.values():
+            for inst in block.instructions:
+                self.instruction_positions[inst] = position
+                position += 1
+
+    def _compute_live_intervals(self) -> None:
+        """Compute live intervals for all values."""
+        # Track first definition and last use for each value
+        first_def: dict[MIRValue, int] = {}
+        last_use: dict[MIRValue, int] = {}
+
+        for block in self.function.cfg.blocks.values():
+            for inst in block.instructions:
+                position = self.instruction_positions[inst]
+
+                # Process definitions
+                for def_val in inst.get_defs():
+                    if self._should_allocate(def_val):
+                        if def_val not in first_def:
+                            first_def[def_val] = position
+                        last_use[def_val] = position  # Def is also a use
+
+                # Process uses
+                for use_val in inst.get_uses():
+                    if self._should_allocate(use_val):
+                        last_use[use_val] = position
+                        if use_val not in first_def:
+                            # Value used before defined (parameter or external)
+                            first_def[use_val] = 0
+
+        # Create intervals
+        for value in first_def:
+            interval = LiveInterval(value=value, start=first_def[value], end=last_use.get(value, first_def[value]))
+            self.live_intervals.append(interval)
+
+    def _should_allocate(self, value: MIRValue) -> bool:
+        """Check if a value needs register allocation.
+
+        Args:
+            value: The value to check.
+
+        Returns:
+            True if the value needs a register.
+        """
+        # Allocate registers for temps and variables
+        return isinstance(value, Temp | Variable)
+
+    def _linear_scan(self) -> dict[MIRValue, int]:
+        """Perform linear scan register allocation.
+
+        Returns:
+            Mapping from values to register numbers.
+        """
+        allocations: dict[MIRValue, int] = {}
+
+        for interval in self.live_intervals:
+            # Expire old intervals
+            self._expire_old_intervals(interval.start)
+
+            # Try to allocate a register
+            if self.free_registers:
+                # Allocate from free registers
+                register = self.free_registers.pop(0)
+                interval.register = register
+                allocations[interval.value] = register
+                self.active_intervals.append(interval)
+                # Sort active intervals by end position
+                self.active_intervals.sort(key=lambda x: x.end)
+            else:
+                # Need to spill - all registers are in use
+                self._spill_at_interval(interval)
+                if interval.register is not None:
+                    # Got a register through spilling
+                    allocations[interval.value] = interval.register
+                    self.active_intervals.append(interval)
+                    self.active_intervals.sort(key=lambda x: x.end)
+                else:
+                    # This interval was spilled to memory
+                    self.spilled_values.add(interval.value)
+                    # Assign a spill slot (using negative numbers for spill slots)
+                    self.next_spill_slot += 1
+                    allocations[interval.value] = -(self.max_registers + self.next_spill_slot)
+
+        return allocations
+
+    def _expire_old_intervals(self, current_position: int) -> None:
+        """Expire intervals that are no longer live.
+
+        Args:
+            current_position: The current position in the program.
+        """
+        expired = []
+        for interval in self.active_intervals:
+            if interval.end >= current_position:
+                break  # Sorted by end, so we can stop
+            expired.append(interval)
+
+        for interval in expired:
+            self.active_intervals.remove(interval)
+            if interval.register is not None and interval.register >= 0:
+                self.free_registers.append(interval.register)
+                self.free_registers.sort()
+
+    def _spill_at_interval(self, interval: LiveInterval) -> None:
+        """Spill a value when no registers are available.
+
+        Args:
+            interval: The interval that needs a register.
+        """
+        if not self.active_intervals:
+            # No active intervals, must spill current
+            self.spilled_values.add(interval.value)
+            interval.register = None
+            return
+
+        # Find the interval with the furthest end point
+        # (this is the last one since active_intervals is sorted by end)
+        spill_candidate = self.active_intervals[-1]
+
+        if spill_candidate.end > interval.end:
+            # Spill the furthest interval and give its register to current
+            self.active_intervals.remove(spill_candidate)
+            interval.register = spill_candidate.register
+            self.spilled_values.add(spill_candidate.value)
+            spill_candidate.register = None
+        else:
+            # Current interval ends later, spill it instead
+            self.spilled_values.add(interval.value)
+            interval.register = None
+
+
+class LifetimeAnalyzer:
+    """Analyzes the lifetime of temporaries for optimization."""
+
+    def __init__(self, function: MIRFunction) -> None:
+        """Initialize the lifetime analyzer.
+
+        Args:
+            function: The function to analyze.
+        """
+        self.function = function
+        self.lifetimes: dict[MIRValue, tuple[int, int]] = {}
+
+    def analyze(self) -> dict[MIRValue, tuple[int, int]]:
+        """Analyze lifetimes of all values.
+
+        Returns:
+            Mapping from values to (first_use, last_use) positions.
+        """
+        position = 0
+
+        for block in self.function.cfg.blocks.values():
+            for inst in block.instructions:
+                # Track definitions
+                for def_val in inst.get_defs():
+                    if isinstance(def_val, Temp | Variable):
+                        if def_val not in self.lifetimes:
+                            self.lifetimes[def_val] = (position, position)
+                        else:
+                            start, _ = self.lifetimes[def_val]
+                            self.lifetimes[def_val] = (start, position)
+
+                # Track uses
+                for use_val in inst.get_uses():
+                    if isinstance(use_val, Temp | Variable):
+                        if use_val not in self.lifetimes:
+                            self.lifetimes[use_val] = (position, position)
+                        else:
+                            start, _ = self.lifetimes[use_val]
+                            self.lifetimes[use_val] = (start, position)
+
+                position += 1
+
+        return self.lifetimes
+
+    def find_reusable_slots(self) -> list[set[MIRValue]]:
+        """Find sets of values that can share the same stack slot.
+
+        Returns:
+            List of sets where each set contains values that can share a slot.
+        """
+        reusable_groups: list[set[MIRValue]] = []
+
+        # Sort values by start of lifetime
+        sorted_values = sorted(self.lifetimes.items(), key=lambda x: x[1][0])
+
+        for value, (start, end) in sorted_values:
+            # Find a group where this value doesn't overlap with any member
+            placed = False
+            for group in reusable_groups:
+                can_share = True
+                for other in group:
+                    other_start, other_end = self.lifetimes[other]
+                    # Check for overlap
+                    if not (end < other_start or start > other_end):
+                        can_share = False
+                        break
+
+                if can_share:
+                    group.add(value)
+                    placed = True
+                    break
+
+            if not placed:
+                # Create a new group
+                reusable_groups.append({value})
+
+        return reusable_groups

machine_dialect/mir/reporting/__init__.py

@@ -0,0 +1,17 @@
+"""MIR optimization reporting infrastructure."""
+
+from machine_dialect.mir.reporting.optimization_reporter import OptimizationReporter
+from machine_dialect.mir.reporting.report_formatter import (
+    HTMLReportFormatter,
+    JSONReportFormatter,
+    ReportFormatter,
+    TextReportFormatter,
+)
+
+__all__ = [
+    "HTMLReportFormatter",
+    "JSONReportFormatter",
+    "OptimizationReporter",
+    "ReportFormatter",
+    "TextReportFormatter",
+]

machine_dialect/mir/reporting/optimization_reporter.py

@@ -0,0 +1,314 @@
+"""Optimization reporter for collecting and aggregating pass statistics.
+
+This module provides infrastructure for collecting optimization statistics
+from various passes and generating comprehensive reports.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class MetricType(Enum):
+    """Types of metrics collected."""
+
+    COUNT = "count"  # Simple count (e.g., instructions removed)
+    PERCENTAGE = "percentage"  # Percentage value
+    SIZE = "size"  # Size in bytes
+    TIME = "time"  # Time in milliseconds
+    RATIO = "ratio"  # Ratio between two values
+
+
+@dataclass
+class PassMetrics:
+    """Metrics collected from a single optimization pass.
+
+    Attributes:
+        pass_name: Name of the optimization pass.
+        phase: Optimization phase (early, middle, late).
+        metrics: Dictionary of metric name to value.
+        before_stats: Statistics before the pass.
+        after_stats: Statistics after the pass.
+        time_ms: Time taken to run the pass in milliseconds.
+    """
+
+    pass_name: str
+    phase: str = "main"
+    metrics: dict[str, int] = field(default_factory=dict)
+    before_stats: dict[str, int] = field(default_factory=dict)
+    after_stats: dict[str, int] = field(default_factory=dict)
+    time_ms: float = 0.0
+
+    def get_improvement(self, metric: str) -> float:
+        """Calculate improvement percentage for a metric.
+
+        Args:
+            metric: Metric name.
+
+        Returns:
+            Improvement percentage (positive means reduction).
+        """
+        before = self.before_stats.get(metric, 0)
+        after = self.after_stats.get(metric, 0)
+        if before == 0:
+            return 0.0
+        return ((before - after) / before) * 100
+
+
+@dataclass
+class ModuleMetrics:
+    """Metrics for an entire module.
+
+    Attributes:
+        module_name: Name of the module.
+        function_metrics: Metrics for each function.
+        pass_metrics: Metrics from each pass.
+        total_time_ms: Total optimization time.
+        optimization_level: Optimization level used.
+    """
+
+    module_name: str
+    function_metrics: dict[str, dict[str, Any]] = field(default_factory=dict)
+    pass_metrics: list[PassMetrics] = field(default_factory=list)
+    total_time_ms: float = 0.0
+    optimization_level: int = 0
+
+    def add_pass_metrics(self, metrics: PassMetrics) -> None:
+        """Add metrics from a pass.
+
+        Args:
+            metrics: Pass metrics to add.
+        """
+        self.pass_metrics.append(metrics)
+        self.total_time_ms += metrics.time_ms
+
+    def get_summary(self) -> dict[str, Any]:
+        """Get summary statistics.
+
+        Returns:
+            Dictionary of summary statistics.
+        """
+        summary = {
+            "module_name": self.module_name,
+            "optimization_level": self.optimization_level,
+            "total_passes": len(self.pass_metrics),
+            "total_time_ms": self.total_time_ms,
+            "passes_applied": [m.pass_name for m in self.pass_metrics],
+        }
+
+        # Aggregate improvements
+        total_improvements = {}
+        for metrics in self.pass_metrics:
+            for key in metrics.before_stats:
+                if key in metrics.after_stats:
+                    improvement = metrics.get_improvement(key)
+                    if key not in total_improvements:
+                        total_improvements[key] = 0.0
+                    total_improvements[key] += improvement
+
+        summary["improvements"] = total_improvements
+
+        # Calculate total metrics
+        total_metrics = {}
+        for metrics in self.pass_metrics:
+            for key, value in metrics.metrics.items():
+                if key not in total_metrics:
+                    total_metrics[key] = 0
+                total_metrics[key] += value
+
+        summary["total_metrics"] = total_metrics
+
+        return summary
+
+
+class OptimizationReporter:
+    """Collects and reports optimization statistics.
+
+    This class aggregates statistics from multiple optimization passes
+    and generates comprehensive reports about the optimization process.
+    """
+
+    def __init__(self, module_name: str = "unknown") -> None:
+        """Initialize the reporter.
+
+        Args:
+            module_name: Name of the module being optimized.
+        """
+        self.module_metrics = ModuleMetrics(module_name=module_name)
+        self.current_pass: PassMetrics | None = None
+
+    def start_pass(
+        self,
+        pass_name: str,
+        phase: str = "main",
+        before_stats: dict[str, int] | None = None,
+    ) -> None:
+        """Start tracking a new pass.
+
+        Args:
+            pass_name: Name of the pass.
+            phase: Optimization phase.
+            before_stats: Statistics before the pass.
+        """
+        self.current_pass = PassMetrics(
+            pass_name=pass_name,
+            phase=phase,
+            before_stats=before_stats or {},
+        )
+
+    def end_pass(
+        self,
+        metrics: dict[str, int] | None = None,
+        after_stats: dict[str, int] | None = None,
+        time_ms: float = 0.0,
+    ) -> None:
+        """End tracking the current pass.
+
+        Args:
+            metrics: Pass-specific metrics.
+            after_stats: Statistics after the pass.
+            time_ms: Time taken by the pass.
+        """
+        if self.current_pass:
+            self.current_pass.metrics = metrics or {}
+            self.current_pass.after_stats = after_stats or {}
+            self.current_pass.time_ms = time_ms
+            self.module_metrics.add_pass_metrics(self.current_pass)
+            self.current_pass = None
+
+    def add_function_metrics(self, func_name: str, metrics: dict[str, Any]) -> None:
+        """Add metrics for a specific function.
+
+        Args:
+            func_name: Function name.
+            metrics: Function metrics.
+        """
+        self.module_metrics.function_metrics[func_name] = metrics
+
+    def add_custom_stats(self, pass_name: str, stats: dict[str, int]) -> None:
+        """Add custom statistics for a pass.
+
+        Args:
+            pass_name: Name of the pass.
+            stats: Statistics to add.
+        """
+        # Create a pass metrics entry for custom stats
+        metrics = PassMetrics(pass_name=pass_name, phase="bytecode", metrics=stats)
+        self.module_metrics.add_pass_metrics(metrics)
+
+    def set_optimization_level(self, level: int) -> None:
+        """Set the optimization level.
+
+        Args:
+            level: Optimization level (0-3).
+        """
+        self.module_metrics.optimization_level = level
+
+    def get_report_data(self) -> ModuleMetrics:
+        """Get the collected metrics.
+
+        Returns:
+            Module metrics.
+        """
+        return self.module_metrics
+
+    def generate_summary(self) -> str:
+        """Generate a text summary of optimizations.
+
+        Returns:
+            Text summary.
+        """
+        summary = self.module_metrics.get_summary()
+        lines = []
+
+        lines.append(f"Module: {summary['module_name']}")
+        lines.append(f"Optimization Level: {summary['optimization_level']}")
+        lines.append(f"Total Passes: {summary['total_passes']}")
+        lines.append(f"Total Time: {summary['total_time_ms']:.2f}ms")
+        lines.append("")
+
+        if summary["passes_applied"]:
+            lines.append("Passes Applied:")
+            for pass_name in summary["passes_applied"]:
+                lines.append(f"  - {pass_name}")
+            lines.append("")
+
+        if summary["improvements"]:
+            lines.append("Improvements:")
+            for metric, improvement in summary["improvements"].items():
+                if improvement > 0:
+                    lines.append(f"  {metric}: {improvement:.1f}% reduction")
+            lines.append("")
+
+        if summary["total_metrics"]:
+            lines.append("Total Changes:")
+            for metric, value in summary["total_metrics"].items():
+                if value > 0:
+                    lines.append(f"  {metric}: {value}")
+
+        return "\n".join(lines)
+
+    def generate_detailed_report(self) -> str:
+        """Generate a detailed report with per-pass statistics.
+
+        Returns:
+            Detailed text report.
+        """
+        lines = []
+        lines.append("=" * 60)
+        lines.append("OPTIMIZATION REPORT")
+        lines.append("=" * 60)
+        lines.append("")
+
+        # Summary
+        lines.append(self.generate_summary())
+        lines.append("")
+        lines.append("=" * 60)
+        lines.append("DETAILED PASS STATISTICS")
+        lines.append("=" * 60)
+
+        # Per-pass details
+        for metrics in self.module_metrics.pass_metrics:
+            lines.append("")
+            lines.append(f"Pass: {metrics.pass_name}")
+            lines.append(f"Phase: {metrics.phase}")
+            lines.append(f"Time: {metrics.time_ms:.2f}ms")
+
+            if metrics.metrics:
+                lines.append("Metrics:")
+                for key, value in metrics.metrics.items():
+                    if value > 0:
+                        lines.append(f"  {key}: {value}")
+
+            # Show improvements
+            improvements = []
+            for key in metrics.before_stats:
+                if key in metrics.after_stats:
+                    improvement = metrics.get_improvement(key)
+                    if improvement > 0:
+                        improvements.append(
+                            f"  {key}: {metrics.before_stats[key]} → "
+                            f"{metrics.after_stats[key]} "
+                            f"({improvement:.1f}% reduction)"
+                        )
+
+            if improvements:
+                lines.append("Improvements:")
+                lines.extend(improvements)
+
+            lines.append("-" * 40)
+
+        # Function-specific metrics if available
+        if self.module_metrics.function_metrics:
+            lines.append("")
+            lines.append("=" * 60)
+            lines.append("FUNCTION METRICS")
+            lines.append("=" * 60)
+
+            for func_name, func_metrics in self.module_metrics.function_metrics.items():
+                lines.append("")
+                lines.append(f"Function: {func_name}")
+                for key, value in func_metrics.items():
+                    lines.append(f"  {key}: {value}")
+
+        return "\n".join(lines)