machine-dialect 0.1.0a1__py3-none-any.whl

machine_dialect/mir/optimizations/cse.py

@@ -0,0 +1,398 @@
+"""Common Subexpression Elimination (CSE) optimization pass.
+
+This module implements CSE at the MIR level, eliminating redundant
+computations by reusing previously computed values.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+from machine_dialect.mir.basic_block import BasicBlock
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.mir_instructions import (
+    BinaryOp,
+    Call,
+    Copy,
+    LoadConst,
+    MIRInstruction,
+    StoreVar,
+    UnaryOp,
+)
+from machine_dialect.mir.mir_transformer import MIRTransformer
+from machine_dialect.mir.mir_values import Constant, MIRValue, Temp, Variable
+from machine_dialect.mir.optimization_pass import (
+    OptimizationPass,
+    PassInfo,
+    PassType,
+    PreservationLevel,
+)
+
+
+@dataclass(frozen=True)
+class Expression:
+    """Represents an expression for CSE.
+
+    Attributes:
+        op: Operation type.
+        operands: Tuple of operands.
+    """
+
+    op: str
+    operands: tuple[Any, ...]
+
+    def __hash__(self) -> int:
+        """Hash the expression."""
+        return hash((self.op, self.operands))
+
+
+class AvailableExpressions:
+    """Tracks available expressions in a block."""
+
+    def __init__(self) -> None:
+        """Initialize available expressions."""
+        # Map from expression to the value containing it
+        self.expressions: dict[Expression, MIRValue] = {}
+        # Map from value to expressions it defines
+        self.definitions: dict[MIRValue, set[Expression]] = {}
+
+    def add(self, expr: Expression, value: MIRValue) -> None:
+        """Add an available expression.
+
+        Args:
+            expr: The expression.
+            value: The value containing the expression.
+        """
+        self.expressions[expr] = value
+        if value not in self.definitions:
+            self.definitions[value] = set()
+        self.definitions[value].add(expr)
+
+    def get(self, expr: Expression) -> MIRValue | None:
+        """Get the value for an expression.
+
+        Args:
+            expr: The expression to look up.
+
+        Returns:
+            The value containing the expression or None.
+        """
+        return self.expressions.get(expr)
+
+    def invalidate(self, value: MIRValue) -> None:
+        """Invalidate expressions involving a value.
+
+        Args:
+            value: The value that changed.
+        """
+        # Remove expressions that use this value
+        to_remove = []
+        for expr in self.expressions:
+            if value in expr.operands:
+                to_remove.append(expr)
+
+        for expr in to_remove:
+            del self.expressions[expr]
+
+        # Remove expressions defined by this value
+        if value in self.definitions:
+            for expr in self.definitions[value]:
+                if expr in self.expressions:
+                    del self.expressions[expr]
+            del self.definitions[value]
+
+    def copy(self) -> "AvailableExpressions":
+        """Create a copy of available expressions.
+
+        Returns:
+            A copy of this available expressions set.
+        """
+        new = AvailableExpressions()
+        new.expressions = self.expressions.copy()
+        new.definitions = {k: v.copy() for k, v in self.definitions.items()}
+        return new
+
+
+class CommonSubexpressionElimination(OptimizationPass):
+    """Common subexpression elimination optimization pass."""
+
+    def get_info(self) -> PassInfo:
+        """Get pass information.
+
+        Returns:
+            Pass information.
+        """
+        return PassInfo(
+            name="cse",
+            description="Eliminate common subexpressions",
+            pass_type=PassType.OPTIMIZATION,
+            requires=[],
+            preserves=PreservationLevel.CFG,
+        )
+
+    def run_on_function(self, function: MIRFunction) -> bool:
+        """Run CSE on a function.
+
+        Args:
+            function: The function to optimize.
+
+        Returns:
+            True if the function was modified.
+        """
+        transformer = MIRTransformer(function)
+
+        # Perform local CSE within each block
+        for block in function.cfg.blocks.values():
+            self._local_cse(block, transformer)
+
+        # Perform global CSE across blocks
+        self._global_cse(function, transformer)
+
+        return transformer.modified
+
+    def _local_cse(self, block: BasicBlock, transformer: MIRTransformer) -> None:
+        """Perform local CSE within a block.
+
+        Args:
+            block: The block to optimize.
+            transformer: MIR transformer.
+        """
+        available = AvailableExpressions()
+
+        for inst in list(block.instructions):
+            # Check if this instruction computes an expression
+            expr = self._get_expression(inst)
+
+            if expr:
+                # Check if expression is already available
+                existing = available.get(expr)
+                if existing and existing != self._get_result(inst):
+                    # Replace with copy of existing value
+                    result = self._get_result(inst)
+                    if result:
+                        new_inst = Copy(result, existing, inst.source_location)
+                        transformer.replace_instruction(block, inst, new_inst)
+                        self.stats["local_cse_eliminated"] = self.stats.get("local_cse_eliminated", 0) + 1
+                else:
+                    # Add expression to available set
+                    result = self._get_result(inst)
+                    if result:
+                        available.add(expr, result)
+
+            # Update available expressions based on side effects
+            self._update_available(inst, available)
+
+    def _global_cse(self, function: MIRFunction, transformer: MIRTransformer) -> None:
+        """Perform global CSE across blocks.
+
+        Args:
+            function: The function to optimize.
+            transformer: MIR transformer.
+        """
+        # Compute available expressions at entry of each block
+        block_available: dict[BasicBlock, AvailableExpressions] = {}
+
+        # Initialize with empty sets
+        for block in function.cfg.blocks.values():
+            block_available[block] = AvailableExpressions()
+
+        # Iterate until fixed point
+        changed = True
+        while changed:
+            changed = False
+
+            for block in function.cfg.blocks.values():
+                # Compute available at entry as intersection of predecessors
+                if block.predecessors:
+                    # Start with copy of first predecessor
+                    if block.predecessors[0] in block_available:
+                        new_available = self._intersect_available(
+                            [block_available.get(p, AvailableExpressions()) for p in block.predecessors]
+                        )
+                    else:
+                        new_available = AvailableExpressions()
+                else:
+                    new_available = AvailableExpressions()
+
+                # Check if changed
+                if self._available_changed(block_available[block], new_available):
+                    block_available[block] = new_available
+                    changed = True
+
+                # Compute available at exit
+                available = new_available.copy()
+                for inst in block.instructions:
+                    expr = self._get_expression(inst)
+                    if expr:
+                        result = self._get_result(inst)
+                        if result:
+                            available.add(expr, result)
+                    self._update_available(inst, available)
+
+        # Apply CSE based on available expressions
+        for block in function.cfg.blocks.values():
+            available = block_available[block].copy()
+
+            for inst in list(block.instructions):
+                expr = self._get_expression(inst)
+
+                if expr:
+                    existing = available.get(expr)
+                    if existing and existing != self._get_result(inst):
+                        # Replace with copy
+                        result = self._get_result(inst)
+                        if result:
+                            source_loc = inst.source_location if hasattr(inst, "source_location") else (0, 0)
+                            new_inst = Copy(result, existing, source_loc)
+                            transformer.replace_instruction(block, inst, new_inst)
+                            self.stats["global_cse_eliminated"] = self.stats.get("global_cse_eliminated", 0) + 1
+                    else:
+                        result = self._get_result(inst)
+                        if result:
+                            available.add(expr, result)
+
+                self._update_available(inst, available)
+
+    def _get_expression(self, inst: MIRInstruction) -> Expression | None:
+        """Extract expression from an instruction.
+
+        Args:
+            inst: The instruction.
+
+        Returns:
+            The expression or None.
+        """
+        if isinstance(inst, BinaryOp):
+            # Normalize commutative operations
+            if inst.op in ["+", "*", "==", "!=", "and", "or"]:
+                # Sort operands for commutative ops
+                operands = tuple(sorted([self._normalize_value(inst.left), self._normalize_value(inst.right)], key=str))
+            else:
+                operands = (self._normalize_value(inst.left), self._normalize_value(inst.right))
+            return Expression(f"binary_{inst.op}", operands)
+
+        elif isinstance(inst, UnaryOp):
+            return Expression(f"unary_{inst.op}", (self._normalize_value(inst.operand),))
+
+        elif isinstance(inst, LoadConst):
+            # Constants are their own expressions
+            return Expression("const", (inst.constant.value, inst.constant.type))
+
+        return None
+
+    def _normalize_value(self, value: MIRValue) -> Any:
+        """Normalize a value for expression comparison.
+
+        Args:
+            value: The value to normalize.
+
+        Returns:
+            Normalized representation.
+        """
+        if isinstance(value, Constant):
+            return ("const", value.value, value.type)
+        elif isinstance(value, Variable):
+            return ("var", value.name)
+        elif isinstance(value, Temp):
+            return ("temp", value.id)
+        else:
+            return str(value)
+
+    def _get_result(self, inst: MIRInstruction) -> MIRValue | None:
+        """Get the result value of an instruction.
+
+        Args:
+            inst: The instruction.
+
+        Returns:
+            The result value or None.
+        """
+        defs = inst.get_defs()
+        if defs and len(defs) == 1:
+            return defs[0]
+        return None
+
+    def _update_available(
+        self,
+        inst: MIRInstruction,
+        available: AvailableExpressions,
+    ) -> None:
+        """Update available expressions after an instruction.
+
+        Args:
+            inst: The instruction.
+            available: Available expressions to update.
+        """
+        # Invalidate expressions if instruction has side effects
+        if isinstance(inst, StoreVar):
+            # Invalidate expressions using this variable
+            available.invalidate(inst.var)
+        elif isinstance(inst, Call):
+            # Conservative: invalidate all expressions with variables
+            # (calls might modify globals or have other side effects)
+            for value in list(available.definitions.keys()):
+                if isinstance(value, Variable):
+                    available.invalidate(value)
+
+    def _intersect_available(
+        self,
+        sets: list[AvailableExpressions],
+    ) -> AvailableExpressions:
+        """Compute intersection of available expression sets.
+
+        Args:
+            sets: List of available expression sets.
+
+        Returns:
+            The intersection.
+        """
+        if not sets:
+            return AvailableExpressions()
+
+        # Start with first set
+        result = AvailableExpressions()
+        if not sets[0].expressions:
+            return result
+
+        # Find expressions available in all sets
+        for expr, value in sets[0].expressions.items():
+            available_in_all = True
+            for s in sets[1:]:
+                if expr not in s.expressions:
+                    available_in_all = False
+                    break
+                # Check if same value
+                if s.expressions[expr] != value:
+                    available_in_all = False
+                    break
+
+            if available_in_all:
+                result.add(expr, value)
+
+        return result
+
+    def _available_changed(
+        self,
+        old: AvailableExpressions,
+        new: AvailableExpressions,
+    ) -> bool:
+        """Check if available expressions changed.
+
+        Args:
+            old: Old available expressions.
+            new: New available expressions.
+
+        Returns:
+            True if changed.
+        """
+        if len(old.expressions) != len(new.expressions):
+            return True
+
+        for expr, value in old.expressions.items():
+            if expr not in new.expressions or new.expressions[expr] != value:
+                return True
+
+        return False
+
+    def finalize(self) -> None:
+        """Finalize the pass."""
+        pass

machine_dialect/mir/optimizations/dce.py

@@ -0,0 +1,288 @@
+"""Dead code elimination optimization pass.
+
+This module implements dead code elimination (DCE) at the MIR level,
+removing instructions and blocks that have no effect on program output.
+"""
+
+from machine_dialect.mir.analyses.use_def_chains import UseDefChains
+from machine_dialect.mir.basic_block import BasicBlock
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.mir_instructions import (
+    Assert,
+    Call,
+    ConditionalJump,
+    Jump,
+    MIRInstruction,
+    Print,
+    Return,
+    Scope,
+    StoreVar,
+)
+from machine_dialect.mir.mir_transformer import MIRTransformer
+from machine_dialect.mir.mir_values import Temp, Variable
+from machine_dialect.mir.optimization_pass import (
+    OptimizationPass,
+    PassInfo,
+    PassType,
+    PreservationLevel,
+)
+
+
+class DeadCodeElimination(OptimizationPass):
+    """Dead code elimination optimization pass."""
+
+    def get_info(self) -> PassInfo:
+        """Get pass information.
+
+        Returns:
+            Pass information.
+        """
+        return PassInfo(
+            name="dce",
+            description="Eliminate dead code and unreachable blocks",
+            pass_type=[PassType.OPTIMIZATION, PassType.CLEANUP],
+            requires=["use-def-chains"],
+            preserves=PreservationLevel.CFG,
+        )
+
+    def run_on_function(self, function: MIRFunction) -> bool:
+        """Run dead code elimination on a function.
+
+        Args:
+            function: The function to optimize.
+
+        Returns:
+            True if the function was modified.
+        """
+        transformer = MIRTransformer(function)
+
+        # Get use-def chains
+        use_def_chains: UseDefChains = self.get_analysis("use-def-chains", function)
+
+        # Phase 1: Remove dead instructions
+        dead_instructions = self._find_dead_instructions(function, use_def_chains)
+        for block, inst in dead_instructions:
+            transformer.remove_instruction(block, inst)
+            self.stats["dead_instructions_removed"] = self.stats.get("dead_instructions_removed", 0) + 1
+
+        # Phase 2: Remove dead stores
+        dead_stores = self._find_dead_stores(function, use_def_chains)
+        for block, inst in dead_stores:
+            # When removing a StoreVar, replace uses of the variable with the source value
+            if isinstance(inst, StoreVar):
+                # Replace all uses of the variable with the source value
+                transformer.replace_uses(inst.var, inst.source)
+            transformer.remove_instruction(block, inst)
+            self.stats["dead_stores_removed"] = self.stats.get("dead_stores_removed", 0) + 1
+
+        # Phase 3: Remove unreachable blocks
+        num_unreachable = transformer.eliminate_unreachable_blocks()
+        self.stats["unreachable_blocks_removed"] = num_unreachable
+
+        # Phase 4: Simplify control flow
+        transformer.simplify_cfg()
+
+        return transformer.modified
+
+    def _find_dead_instructions(
+        self,
+        function: MIRFunction,
+        use_def_chains: UseDefChains,
+    ) -> list[tuple[BasicBlock, MIRInstruction]]:
+        """Find dead instructions that can be removed.
+
+        Args:
+            function: The function to analyze.
+            use_def_chains: Use-def chain information.
+
+        Returns:
+            List of (block, instruction) pairs to remove.
+        """
+        dead = []
+
+        for block in function.cfg.blocks.values():
+            for inst in block.instructions:
+                # Skip instructions with side effects
+                if self._has_side_effects(inst):
+                    continue
+
+                # Check if all values defined by this instruction are dead
+                defs = inst.get_defs()
+                if not defs:
+                    continue
+
+                all_dead = True
+                for value in defs:
+                    if isinstance(value, Temp | Variable):
+                        if not use_def_chains.is_dead(value):
+                            all_dead = False
+                            break
+
+                if all_dead:
+                    dead.append((block, inst))
+
+        return dead
+
+    def _find_dead_stores(
+        self,
+        function: MIRFunction,
+        use_def_chains: UseDefChains,
+    ) -> list[tuple[BasicBlock, MIRInstruction]]:
+        """Find dead store instructions.
+
+        Args:
+            function: The function to analyze.
+            use_def_chains: Use-def chain information.
+
+        Returns:
+            List of (block, instruction) pairs to remove.
+        """
+        dead_stores = []
+
+        for block in function.cfg.blocks.values():
+            # Track last store to each variable in the block
+            last_stores: dict[Variable, MIRInstruction] = {}
+
+            for inst in block.instructions:
+                if isinstance(inst, StoreVar):
+                    # Check if there's a previous store to the same variable
+                    if inst.var in last_stores:
+                        # Previous store might be dead if not used between stores
+                        prev_store = last_stores[inst.var]
+                        if self._is_dead_store(
+                            prev_store,
+                            inst,
+                            block,
+                            use_def_chains,
+                        ):
+                            dead_stores.append((block, prev_store))
+
+                    last_stores[inst.var] = inst
+
+            # Check if final stores are dead (no uses after block)
+            for _var, store in last_stores.items():
+                if self._is_store_dead_at_end(store, block, use_def_chains):
+                    dead_stores.append((block, store))
+
+        return dead_stores
+
+    def _is_dead_store(
+        self,
+        store1: MIRInstruction,
+        store2: MIRInstruction,
+        block: BasicBlock,
+        use_def_chains: UseDefChains,
+    ) -> bool:
+        """Check if store1 is dead because of store2.
+
+        Args:
+            store1: First store instruction.
+            store2: Second store instruction.
+            block: Containing block.
+            use_def_chains: Use-def chain information.
+
+        Returns:
+            True if store1 is dead.
+        """
+        # Find instructions between the two stores
+        idx1 = block.instructions.index(store1)
+        idx2 = block.instructions.index(store2)
+
+        if idx1 >= idx2:
+            return False
+
+        # Check if variable is used between the stores
+        if isinstance(store1, StoreVar):
+            var = store1.var
+            for i in range(idx1 + 1, idx2):
+                inst = block.instructions[i]
+                if var in inst.get_uses():
+                    return False
+
+        return True
+
+    def _is_store_dead_at_end(
+        self,
+        store: MIRInstruction,
+        block: BasicBlock,
+        use_def_chains: UseDefChains,
+    ) -> bool:
+        """Check if a store at the end of a block is dead.
+
+        Args:
+            store: Store instruction.
+            block: Containing block.
+            use_def_chains: Use-def chain information.
+
+        Returns:
+            True if the store is dead.
+        """
+        if not isinstance(store, StoreVar):
+            return False
+
+        # Check if variable is used after this block
+        var = store.var
+
+        # Check uses in successor blocks
+        visited = set()
+        worklist = list(block.successors)
+
+        while worklist:
+            succ = worklist.pop()
+            if succ in visited:
+                continue
+            visited.add(succ)
+
+            # Check phi nodes
+            for phi in succ.phi_nodes:
+                for val, pred_label in phi.incoming:
+                    if val == var and pred_label == block.label:
+                        return False
+
+            # Check instructions
+            for inst in succ.instructions:
+                if var in inst.get_uses():
+                    return False
+                # If we see another store to the same variable, we're done
+                if isinstance(inst, StoreVar) and inst.var == var:
+                    break
+            else:
+                # No store found, check successors
+                worklist.extend(succ.successors)
+
+        return True
+
+    def _has_side_effects(self, inst: MIRInstruction) -> bool:
+        """Check if an instruction has side effects.
+
+        Args:
+            inst: Instruction to check.
+
+        Returns:
+            True if the instruction has side effects.
+        """
+        # Control flow instructions
+        if isinstance(inst, Jump | ConditionalJump | Return):
+            return True
+
+        # I/O operations
+        if isinstance(inst, Print):
+            return True
+
+        # Function calls (conservative - assume all calls have side effects)
+        if isinstance(inst, Call):
+            return True
+
+        # Memory operations
+        if isinstance(inst, StoreVar):
+            return True
+
+        # Assertions and scopes
+        if isinstance(inst, Assert | Scope):
+            return True
+
+        return False
+
+    def finalize(self) -> None:
+        """Finalize the pass."""
+        pass