machine-dialect 0.1.0a1__py3-none-any.whl

machine_dialect/mir/analyses/alias_analysis.py

@@ -0,0 +1,315 @@
+"""Alias analysis for MIR.
+
+This module provides alias analysis to determine when different variables
+may refer to the same memory location, supporting optimizations like
+escape analysis and redundant load elimination.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from machine_dialect.mir.basic_block import BasicBlock
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.mir_instructions import (
+    Copy,
+    GetAttr,
+    LoadVar,
+    MIRInstruction,
+    Phi,
+    SetAttr,
+    StoreVar,
+)
+from machine_dialect.mir.mir_types import MIRType
+from machine_dialect.mir.mir_values import Variable
+from machine_dialect.mir.optimization_pass import (
+    FunctionAnalysisPass,
+    PassInfo,
+    PassType,
+    PreservationLevel,
+)
+
+
+class AliasType(Enum):
+    """Type of aliasing relationship."""
+
+    NO_ALIAS = "no_alias"  # Definitely different locations
+    MAY_ALIAS = "may_alias"  # Might be the same location
+    MUST_ALIAS = "must_alias"  # Definitely the same location
+    PARTIAL_ALIAS = "partial_alias"  # Overlapping but not identical
+
+
+@dataclass
+class AliasSet:
+    """Set of potentially aliasing values.
+
+    Attributes:
+        members: Values that may alias each other.
+        alias_type: Type of aliasing relationship.
+        attributes: Attributes accessed on these values.
+    """
+
+    members: set[Variable] = field(default_factory=set)
+    alias_type: AliasType = AliasType.MAY_ALIAS
+    attributes: set[str] = field(default_factory=set)
+
+    def add_member(self, var: Variable) -> None:
+        """Add a variable to the alias set.
+
+        Args:
+            var: Variable to add.
+        """
+        self.members.add(var)
+
+    def merge(self, other: "AliasSet") -> None:
+        """Merge another alias set into this one.
+
+        Args:
+            other: Alias set to merge.
+        """
+        self.members.update(other.members)
+        self.attributes.update(other.attributes)
+        # Downgrade alias type to be conservative
+        if other.alias_type == AliasType.MAY_ALIAS or self.alias_type == AliasType.MAY_ALIAS:
+            self.alias_type = AliasType.MAY_ALIAS
+
+
+class AliasInfo:
+    """Container for alias analysis results."""
+
+    def __init__(self) -> None:
+        """Initialize alias information."""
+        self.alias_sets: list[AliasSet] = []
+        self.var_to_set: dict[str, AliasSet] = {}
+
+    def get_aliases(self, var: Variable) -> set[Variable]:
+        """Get all variables that may alias with the given variable.
+
+        Args:
+            var: Variable to query.
+
+        Returns:
+            Set of potentially aliasing variables.
+        """
+        alias_set = self.var_to_set.get(var.name)
+        if alias_set:
+            return alias_set.members - {var}
+        return set()
+
+    def may_alias(self, var1: Variable, var2: Variable) -> bool:
+        """Check if two variables may alias.
+
+        Args:
+            var1: First variable.
+            var2: Second variable.
+
+        Returns:
+            True if variables may alias.
+        """
+        set1 = self.var_to_set.get(var1.name)
+        set2 = self.var_to_set.get(var2.name)
+        return set1 is set2 if (set1 and set2) else False
+
+    def must_alias(self, var1: Variable, var2: Variable) -> bool:
+        """Check if two variables must alias.
+
+        Args:
+            var1: First variable.
+            var2: Second variable.
+
+        Returns:
+            True if variables must alias.
+        """
+        set1 = self.var_to_set.get(var1.name)
+        set2 = self.var_to_set.get(var2.name)
+        if set1 is set2 and set1:
+            return set1.alias_type == AliasType.MUST_ALIAS
+        return False
+
+
+class AliasAnalysis(FunctionAnalysisPass):
+    """Analyze pointer aliasing relationships."""
+
+    def get_info(self) -> PassInfo:
+        """Get pass information.
+
+        Returns:
+            Pass information.
+        """
+        return PassInfo(
+            name="alias-analysis",
+            description="Analyze variable aliasing relationships",
+            pass_type=PassType.ANALYSIS,
+            requires=[],
+            preserves=PreservationLevel.ALL,
+        )
+
+    def run_on_function(self, function: MIRFunction) -> AliasInfo:
+        """Analyze aliasing in a function.
+
+        Args:
+            function: The function to analyze.
+
+        Returns:
+            Alias analysis results.
+        """
+        alias_info = AliasInfo()
+
+        # Initialize each variable in its own set
+        self._initialize_variables(function, alias_info)
+
+        # Analyze instructions to find aliasing
+        for block in function.cfg.blocks.values():
+            self._analyze_block(block, alias_info)
+
+        # Handle phi nodes specially
+        self._analyze_phi_nodes(function, alias_info)
+
+        return alias_info
+
+    def _initialize_variables(self, function: MIRFunction, alias_info: AliasInfo) -> None:
+        """Initialize alias sets for all variables.
+
+        Args:
+            function: The function being analyzed.
+            alias_info: Alias information to populate.
+        """
+        # Create initial alias sets for locals
+        for var_name in function.locals:
+            var = Variable(var_name, MIRType.INT)
+            alias_set = AliasSet()
+            alias_set.add_member(var)
+            alias_set.alias_type = AliasType.NO_ALIAS  # Initially no aliases
+            alias_info.alias_sets.append(alias_set)
+            alias_info.var_to_set[var_name] = alias_set
+
+        # Create initial sets for parameters
+        for param in function.params:
+            alias_set = AliasSet()
+            alias_set.add_member(param)
+            # Parameters may alias with external data
+            alias_set.alias_type = AliasType.MAY_ALIAS
+            alias_info.alias_sets.append(alias_set)
+            alias_info.var_to_set[param.name] = alias_set
+
+    def _analyze_block(self, block: BasicBlock, alias_info: AliasInfo) -> None:
+        """Analyze aliasing in a block.
+
+        Args:
+            block: The block to analyze.
+            alias_info: Alias information to update.
+        """
+        for inst in block.instructions:
+            self._analyze_instruction(inst, alias_info)
+
+    def _analyze_instruction(self, inst: MIRInstruction, alias_info: AliasInfo) -> None:
+        """Analyze aliasing effects of an instruction.
+
+        Args:
+            inst: The instruction to analyze.
+            alias_info: Alias information to update.
+        """
+        # Copy creates must-alias relationship
+        if isinstance(inst, Copy):
+            if isinstance(inst.source, Variable) and isinstance(inst.dest, Variable):
+                self._merge_alias_sets(inst.source, inst.dest, AliasType.MUST_ALIAS, alias_info)
+
+        # Load creates may-alias relationship
+        elif isinstance(inst, LoadVar):
+            if isinstance(inst.dest, Variable) and isinstance(inst.var, Variable):
+                self._merge_alias_sets(inst.var, inst.dest, AliasType.MAY_ALIAS, alias_info)
+
+        # Store can create aliasing through memory
+        elif isinstance(inst, StoreVar):
+            if isinstance(inst.source, Variable) and isinstance(inst.var, Variable):
+                self._merge_alias_sets(inst.var, inst.source, AliasType.MAY_ALIAS, alias_info)
+
+        # GetAttr tracks attribute access
+        elif isinstance(inst, GetAttr):
+            if isinstance(inst.obj, Variable) and isinstance(inst.dest, Variable):
+                obj_set = alias_info.var_to_set.get(inst.obj.name)
+                if obj_set:
+                    obj_set.attributes.add(inst.attr)
+                # Result may alias with the field
+                self._merge_alias_sets(inst.obj, inst.dest, AliasType.PARTIAL_ALIAS, alias_info)
+
+        # SetAttr tracks attribute modification
+        elif isinstance(inst, SetAttr):
+            if isinstance(inst.obj, Variable):
+                obj_set = alias_info.var_to_set.get(inst.obj.name)
+                if obj_set:
+                    obj_set.attributes.add(inst.attr)
+
+    def _merge_alias_sets(
+        self,
+        var1: Variable,
+        var2: Variable,
+        alias_type: AliasType,
+        alias_info: AliasInfo,
+    ) -> None:
+        """Merge alias sets for two variables.
+
+        Args:
+            var1: First variable.
+            var2: Second variable.
+            alias_type: Type of aliasing.
+            alias_info: Alias information to update.
+        """
+        set1 = alias_info.var_to_set.get(var1.name)
+        set2 = alias_info.var_to_set.get(var2.name)
+
+        if not set1 and not set2:
+            # Create new set for both
+            new_set = AliasSet()
+            new_set.add_member(var1)
+            new_set.add_member(var2)
+            new_set.alias_type = alias_type
+            alias_info.alias_sets.append(new_set)
+            alias_info.var_to_set[var1.name] = new_set
+            alias_info.var_to_set[var2.name] = new_set
+        elif set1 and not set2:
+            # Add var2 to set1
+            set1.add_member(var2)
+            if alias_type == AliasType.MAY_ALIAS or set1.alias_type != alias_type:
+                set1.alias_type = AliasType.MAY_ALIAS
+            alias_info.var_to_set[var2.name] = set1
+        elif not set1 and set2:
+            # Add var1 to set2
+            set2.add_member(var1)
+            if alias_type == AliasType.MAY_ALIAS or set2.alias_type != alias_type:
+                set2.alias_type = AliasType.MAY_ALIAS
+            alias_info.var_to_set[var1.name] = set2
+        elif set1 and set2 and set1 is not set2:
+            # Merge sets
+            set1.merge(set2)
+            if alias_type == AliasType.MAY_ALIAS:
+                set1.alias_type = AliasType.MAY_ALIAS
+            # Update all members of set2 to point to set1
+            for member in set2.members:
+                alias_info.var_to_set[member.name] = set1
+            # Remove set2 from list
+            alias_info.alias_sets.remove(set2)
+
+    def _analyze_phi_nodes(self, function: MIRFunction, alias_info: AliasInfo) -> None:
+        """Analyze phi nodes for aliasing.
+
+        Args:
+            function: The function being analyzed.
+            alias_info: Alias information to update.
+        """
+        for block in function.cfg.blocks.values():
+            for inst in block.instructions:
+                if isinstance(inst, Phi):
+                    # All phi inputs may alias with the output
+                    if isinstance(inst.dest, Variable):
+                        for value, _ in inst.incoming:
+                            if isinstance(value, Variable):
+                                self._merge_alias_sets(
+                                    inst.dest,
+                                    value,
+                                    AliasType.MAY_ALIAS,
+                                    alias_info,
+                                )
+
+    def finalize(self) -> None:
+        """Finalize the pass."""
+        pass

machine_dialect/mir/analyses/dominance_analysis.py

@@ -0,0 +1,49 @@
+"""Dominance analysis for MIR.
+
+This module provides a pass wrapper for dominance analysis.
+"""
+
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.optimization_pass import (
+    FunctionAnalysisPass,
+    PassInfo,
+    PassType,
+    PreservationLevel,
+)
+from machine_dialect.mir.ssa_construction import DominanceInfo
+
+
+class DominanceAnalysis(FunctionAnalysisPass):
+    """Analysis pass that computes dominance information."""
+
+    def get_info(self) -> PassInfo:
+        """Get pass information.
+
+        Returns:
+            Pass information.
+        """
+        return PassInfo(
+            name="dominance",
+            description="Compute dominance information",
+            pass_type=PassType.ANALYSIS,
+            requires=[],
+            preserves=PreservationLevel.ALL,
+        )
+
+    def run_on_function(self, function: MIRFunction) -> DominanceInfo:
+        """Compute dominance for a function.
+
+        Args:
+            function: The function to analyze.
+
+        Returns:
+            Dominance information.
+        """
+        return DominanceInfo(function.cfg)
+
+    def finalize(self) -> None:
+        """Finalize the analysis pass.
+
+        Nothing to clean up for dominance analysis.
+        """
+        pass

machine_dialect/mir/analyses/escape_analysis.py

@@ -0,0 +1,286 @@
+"""Escape analysis for MIR variables.
+
+This module analyzes variable escape behavior to determine which variables
+can be safely allocated on the stack instead of the heap, improving memory
+efficiency and cache locality.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from machine_dialect.mir.basic_block import BasicBlock
+from machine_dialect.mir.mir_function import MIRFunction
+from machine_dialect.mir.mir_instructions import (
+    Call,
+    Copy,
+    MIRInstruction,
+    Return,
+    SetAttr,
+    StoreVar,
+)
+from machine_dialect.mir.mir_types import MIRType
+from machine_dialect.mir.mir_values import Variable
+from machine_dialect.mir.optimization_pass import (
+    FunctionAnalysisPass,
+    PassInfo,
+    PassType,
+    PreservationLevel,
+)
+
+
+class EscapeState(Enum):
+    """Escape state of a variable."""
+
+    NO_ESCAPE = "no_escape"  # Can be stack allocated
+    ARG_ESCAPE = "arg_escape"  # Escapes as function argument
+    RETURN_ESCAPE = "return_escape"  # Escapes via return
+    HEAP_ESCAPE = "heap_escape"  # Stored in heap object
+    GLOBAL_ESCAPE = "global_escape"  # Escapes globally
+
+
+@dataclass
+class VariableEscapeInfo:
+    """Information about a variable's escape behavior.
+
+    Attributes:
+        variable: The variable being analyzed.
+        state: Current escape state.
+        escape_sites: Instructions where escape occurs.
+        aliases: Other variables that alias this one.
+        stack_eligible: Whether eligible for stack allocation.
+    """
+
+    variable: Variable
+    state: EscapeState = EscapeState.NO_ESCAPE
+    escape_sites: list[MIRInstruction] = field(default_factory=list)
+    aliases: set[Variable] = field(default_factory=set)
+    stack_eligible: bool = True
+
+    def mark_escape(self, state: EscapeState, site: MIRInstruction) -> None:
+        """Mark variable as escaping.
+
+        Args:
+            state: New escape state.
+            site: Instruction causing escape.
+        """
+        # Upgrade escape state (more restrictive wins)
+        if state == EscapeState.GLOBAL_ESCAPE or self.state == EscapeState.GLOBAL_ESCAPE:
+            self.state = EscapeState.GLOBAL_ESCAPE
+            self.stack_eligible = False
+        elif state == EscapeState.HEAP_ESCAPE or self.state == EscapeState.HEAP_ESCAPE:
+            self.state = EscapeState.HEAP_ESCAPE
+            self.stack_eligible = False
+        elif state == EscapeState.RETURN_ESCAPE or self.state == EscapeState.RETURN_ESCAPE:
+            self.state = EscapeState.RETURN_ESCAPE
+            self.stack_eligible = False
+        elif state == EscapeState.ARG_ESCAPE or self.state == EscapeState.ARG_ESCAPE:
+            self.state = EscapeState.ARG_ESCAPE
+            self.stack_eligible = False
+
+        self.escape_sites.append(site)
+
+
+class EscapeInfo:
+    """Container for escape analysis results."""
+
+    def __init__(self) -> None:
+        """Initialize escape information."""
+        self.variable_info: dict[str, VariableEscapeInfo] = {}
+        self.stack_eligible: set[Variable] = set()
+        self.escaping_vars: set[Variable] = set()
+
+    def get_info(self, var: Variable) -> VariableEscapeInfo | None:
+        """Get escape info for a variable.
+
+        Args:
+            var: The variable to query.
+
+        Returns:
+            Escape information or None.
+        """
+        return self.variable_info.get(var.name)
+
+    def is_stack_eligible(self, var: Variable) -> bool:
+        """Check if variable can be stack allocated.
+
+        Args:
+            var: The variable to check.
+
+        Returns:
+            True if stack eligible.
+        """
+        info = self.get_info(var)
+        return info.stack_eligible if info else False
+
+    def does_escape(self, var: Variable) -> bool:
+        """Check if variable escapes.
+
+        Args:
+            var: The variable to check.
+
+        Returns:
+            True if variable escapes.
+        """
+        info = self.get_info(var)
+        return info.state != EscapeState.NO_ESCAPE if info else False
+
+
+class EscapeAnalysis(FunctionAnalysisPass):
+    """Analyze variable escape behavior for stack allocation optimization."""
+
+    def get_info(self) -> PassInfo:
+        """Get pass information.
+
+        Returns:
+            Pass information.
+        """
+        return PassInfo(
+            name="escape-analysis",
+            description="Analyze variable escape behavior",
+            pass_type=PassType.ANALYSIS,
+            requires=[],
+            preserves=PreservationLevel.ALL,
+        )
+
+    def run_on_function(self, function: MIRFunction) -> EscapeInfo:
+        """Analyze escape behavior in a function.
+
+        Args:
+            function: The function to analyze.
+
+        Returns:
+            Escape analysis results.
+        """
+        escape_info = EscapeInfo()
+
+        # Initialize info for all variables
+        self._initialize_variables(function, escape_info)
+
+        # Analyze each block
+        for block in function.cfg.blocks.values():
+            self._analyze_block(block, escape_info)
+
+        # Propagate escape through aliases
+        self._propagate_escapes(escape_info)
+
+        # Collect final results
+        self._finalize_results(escape_info)
+
+        return escape_info
+
+    def _initialize_variables(self, function: MIRFunction, escape_info: EscapeInfo) -> None:
+        """Initialize escape info for all variables.
+
+        Args:
+            function: The function being analyzed.
+            escape_info: Escape information to populate.
+        """
+        # Add local variables
+        for var_name in function.locals:
+            var = Variable(var_name, MIRType.INT)  # Default to INT for simplicity
+            escape_info.variable_info[var_name] = VariableEscapeInfo(variable=var)
+
+        # Parameters are more complex - they're already "from outside"
+        # but we track if they escape further
+        for param in function.params:
+            info = VariableEscapeInfo(variable=param)
+            # Parameters themselves don't escape just by existing
+            escape_info.variable_info[param.name] = info
+
+    def _analyze_block(self, block: BasicBlock, escape_info: EscapeInfo) -> None:
+        """Analyze escape behavior in a block.
+
+        Args:
+            block: The block to analyze.
+            escape_info: Escape information to update.
+        """
+        for inst in block.instructions:
+            self._analyze_instruction(inst, escape_info)
+
+    def _analyze_instruction(self, inst: MIRInstruction, escape_info: EscapeInfo) -> None:
+        """Analyze escape behavior of an instruction.
+
+        Args:
+            inst: The instruction to analyze.
+            escape_info: Escape information to update.
+        """
+        # Check for returns
+        if isinstance(inst, Return):
+            if inst.value and isinstance(inst.value, Variable):
+                info = escape_info.variable_info.get(inst.value.name)
+                if info:
+                    info.mark_escape(EscapeState.RETURN_ESCAPE, inst)
+
+        # Check for function calls
+        elif isinstance(inst, Call):
+            # Arguments to calls may escape
+            for arg in inst.args:
+                if isinstance(arg, Variable):
+                    info = escape_info.variable_info.get(arg.name)
+                    if info:
+                        # Conservative: assume args escape
+                        # Could be refined with interprocedural analysis
+                        info.mark_escape(EscapeState.ARG_ESCAPE, inst)
+
+        # Check for stores to heap objects
+        elif isinstance(inst, SetAttr):
+            if isinstance(inst.value, Variable):
+                info = escape_info.variable_info.get(inst.value.name)
+                if info:
+                    info.mark_escape(EscapeState.HEAP_ESCAPE, inst)
+
+        # Check for copies (creates aliases)
+        elif isinstance(inst, Copy):
+            if isinstance(inst.source, Variable) and isinstance(inst.dest, Variable):
+                src_info = escape_info.variable_info.get(inst.source.name)
+                dest_info = escape_info.variable_info.get(inst.dest.name)
+                if src_info and dest_info:
+                    # Track aliasing
+                    src_info.aliases.add(inst.dest)
+                    dest_info.aliases.add(inst.source)
+
+        # Check for stores to global/external locations
+        elif isinstance(inst, StoreVar):
+            if isinstance(inst.source, Variable):
+                info = escape_info.variable_info.get(inst.source.name)
+                if info:
+                    # Storing to another variable might be escape
+                    # Conservative for now
+                    if inst.var.name not in escape_info.variable_info:
+                        # Storing to unknown location
+                        info.mark_escape(EscapeState.GLOBAL_ESCAPE, inst)
+
+    def _propagate_escapes(self, escape_info: EscapeInfo) -> None:
+        """Propagate escape states through aliases.
+
+        Args:
+            escape_info: Escape information to update.
+        """
+        # If a variable escapes, all its aliases escape
+        changed = True
+        while changed:
+            changed = False
+            for var_info in escape_info.variable_info.values():
+                if var_info.state != EscapeState.NO_ESCAPE:
+                    for alias in var_info.aliases:
+                        alias_info = escape_info.variable_info.get(alias.name)
+                        if alias_info and alias_info.state == EscapeState.NO_ESCAPE:
+                            alias_info.state = var_info.state
+                            alias_info.stack_eligible = False
+                            changed = True
+
+    def _finalize_results(self, escape_info: EscapeInfo) -> None:
+        """Finalize escape analysis results.
+
+        Args:
+            escape_info: Escape information to finalize.
+        """
+        for var_info in escape_info.variable_info.values():
+            if var_info.stack_eligible:
+                escape_info.stack_eligible.add(var_info.variable)
+            else:
+                escape_info.escaping_vars.add(var_info.variable)
+
+    def finalize(self) -> None:
+        """Finalize the pass."""
+        pass