kirin-toolchain 0.13.0__py3-none-any.whl

kirin/__init__.py

@@ -0,0 +1,7 @@
+# re-exports the public API of the kirin package
+from kirin import ir
+from kirin.decl import info, statement
+
+from . import types as types
+
+__all__ = ["ir", "types", "statement", "info"]

kirin/analysis/__init__.py

@@ -0,0 +1,24 @@
+"""Analysis module for kirin.
+
+This module contains the analysis framework for kirin. The analysis framework is
+built on top of the interpreter framework. This module provides a set of base classes
+and frameworks for implementing compiler analysis passes on the IR.
+
+The analysis framework contains the following modules:
+
+- [`cfg`][kirin.analysis.cfg]: Control flow graph for a given IR.
+- [`forward`][kirin.analysis.forward]: Forward dataflow analysis.
+- [`callgraph`][kirin.analysis.callgraph]: Call graph for a given IR.
+- [`typeinfer`][kirin.analysis.typeinfer]: Type inference analysis.
+- [`const`][kirin.analysis.const]: Constants used in the analysis framework.
+"""
+
+from kirin.analysis import const as const
+from kirin.analysis.cfg import CFG as CFG
+from kirin.analysis.forward import (
+    Forward as Forward,
+    ForwardExtra as ForwardExtra,
+    ForwardFrame as ForwardFrame,
+)
+from kirin.analysis.callgraph import CallGraph as CallGraph
+from kirin.analysis.typeinfer import TypeInference as TypeInference

kirin/analysis/callgraph.py

@@ -0,0 +1,61 @@
+from typing import Iterable
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.print import Printable
+from kirin.dialects import func
+from kirin.print.printer import Printer
+
+
+@dataclass
+class CallGraph(Printable):
+    """Call graph for a given [`ir.Method`][kirin.ir.Method].
+
+    This class implements the [`kirin.graph.Graph`][kirin.graph.Graph] protocol.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
+
+    defs: dict[str, ir.Method] = field(default_factory=dict)
+    """Mapping from symbol names to methods."""
+    backedges: dict[str, set[str]] = field(default_factory=dict)
+    """Mapping from symbol names to backedges."""
+
+    def __init__(self, mt: ir.Method):
+        self.defs = {}
+        self.backedges = {}
+        self.__build(mt)
+
+    def __build(self, mt: ir.Method):
+        self.defs[mt.sym_name] = mt
+        for stmt in mt.callable_region.walk():
+            if isinstance(stmt, func.Invoke):
+                backedges = self.backedges.setdefault(stmt.callee.sym_name, set())
+                backedges.add(mt.sym_name)
+                self.__build(stmt.callee)
+
+    def get_neighbors(self, node: str) -> Iterable[str]:
+        """Get the neighbors of a node in the call graph."""
+        return self.backedges.get(node, ())
+
+    def get_edges(self) -> Iterable[tuple[str, str]]:
+        """Get the edges of the call graph."""
+        for node, neighbors in self.backedges.items():
+            for neighbor in neighbors:
+                yield node, neighbor
+
+    def get_nodes(self) -> Iterable[str]:
+        """Get the nodes of the call graph."""
+        return self.defs.keys()
+
+    def print_impl(self, printer: Printer) -> None:
+        for idx, (caller, callee) in enumerate(self.backedges.items()):
+            printer.plain_print(caller)
+            printer.plain_print(" -> ")
+            printer.print_seq(
+                callee, delim=", ", prefix="[", suffix="]", emit=printer.plain_print
+            )
+            if idx < len(self.backedges) - 1:
+                printer.print_newline()

kirin/analysis/cfg.py

@@ -0,0 +1,112 @@
+from typing import Iterable
+from functools import cached_property
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.print import Printer, Printable
+from kirin.worklist import WorkList
+
+
+@dataclass
+class CFG(Printable):
+    """Control Flow Graph of a given IR statement.
+
+    This class implements the [`kirin.graph.Graph`][kirin.graph.Graph] protocol.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
+
+    parent: ir.Region
+    """Parent IR statement.
+    """
+    entry: ir.Block | None = None
+    """Entry block of the CFG.
+    """
+
+    def __post_init__(self):
+        if self.parent.blocks.isempty():
+            self.entry = None
+        else:
+            self.entry = self.parent.blocks[0]
+
+    @cached_property
+    def predecessors(self):
+        """CFG data, mapping a block to its predecessors."""
+        graph: dict[ir.Block, set[ir.Block]] = {}
+        for block, neighbors in self.successors.items():
+            for neighbor in neighbors:
+                graph.setdefault(neighbor, set()).add(block)
+        return graph
+
+    @cached_property
+    def successors(self):
+        """CFG data, mapping a block to its neighbors."""
+        graph: dict[ir.Block, set[ir.Block]] = {}
+        visited: set[ir.Block] = set()
+        worklist: WorkList[ir.Block] = WorkList()
+        if self.parent.blocks.isempty():
+            return graph
+
+        block = self.entry
+        while block is not None:
+            neighbors = graph.setdefault(block, set())
+            if block.last_stmt is not None:
+                neighbors.update(block.last_stmt.successors)
+                worklist.extend(block.last_stmt.successors)
+            visited.add(block)
+
+            block = worklist.pop()
+            while block is not None and block in visited:
+                block = worklist.pop()
+        return graph
+
+    # graph interface
+    def get_neighbors(self, node: ir.Block) -> Iterable[ir.Block]:
+        return self.successors[node]
+
+    def get_edges(self) -> Iterable[tuple[ir.Block, ir.Block]]:
+        for block, neighbors in self.successors.items():
+            for neighbor in neighbors:
+                yield block, neighbor
+
+    def get_nodes(self) -> Iterable[ir.Block]:
+        return self.successors.keys()
+
+    # printable interface
+    def print_impl(self, printer: Printer) -> None:
+        # NOTE: this make sure we use the same name
+        # as the printing of CFG parent.
+        with printer.string_io():
+            self.parent.print(printer)
+
+        printer.plain_print("Successors:")
+        printer.print_newline()
+        for block, neighbors in self.successors.items():
+            printer.plain_print(f"{printer.state.block_id[block]} -> ", end="")
+            printer.print_seq(
+                neighbors,
+                delim=", ",
+                prefix="[",
+                suffix="]",
+                emit=lambda block: printer.plain_print(printer.state.block_id[block]),
+            )
+            printer.print_newline()
+
+        if self.predecessors:
+            printer.print_newline()
+            printer.plain_print("Predecessors:")
+            printer.print_newline()
+            for block, neighbors in self.predecessors.items():
+                printer.plain_print(f"{printer.state.block_id[block]} <- ", end="")
+                printer.print_seq(
+                    neighbors,
+                    delim=", ",
+                    prefix="[",
+                    suffix="]",
+                    emit=lambda block: printer.plain_print(
+                        printer.state.block_id[block]
+                    ),
+                )
+                printer.print_newline()

kirin/analysis/const/__init__.py

@@ -0,0 +1,20 @@
+"""Const analysis module.
+
+This module contains the constant analysis framework for kirin. The constant
+analysis framework is built on top of the interpreter framework.
+
+This module provides a lattice for constant propagation analysis and a
+propagation algorithm for computing the constant values for each SSA value in
+the IR.
+"""
+
+from .prop import Frame as Frame, Propagate as Propagate
+from .lattice import (
+    Value as Value,
+    Bottom as Bottom,
+    Result as Result,
+    Unknown as Unknown,
+    PartialConst as PartialConst,
+    PartialTuple as PartialTuple,
+    PartialLambda as PartialLambda,
+)

kirin/analysis/const/_visitor.py

@@ -0,0 +1,2 @@
+class _ElemVisitor:
+    pass

kirin/analysis/const/_visitor.pyi

@@ -0,0 +1,8 @@
+from .lattice import Value, Bottom, Unknown, PartialTuple, PartialLambda
+
+class _ElemVisitor:
+    def is_subseteq_Value(self, other: Value) -> bool: ...
+    def is_subseteq_NotConst(self, other: Unknown) -> bool: ...
+    def is_subseteq_Unknown(self, other: Bottom) -> bool: ...
+    def is_subseteq_PartialTuple(self, other: PartialTuple) -> bool: ...
+    def is_subseteq_PartialLambda(self, other: PartialLambda) -> bool: ...

kirin/analysis/const/lattice.py

@@ -0,0 +1,219 @@
+"""Lattice for constant analysis.
+"""
+
+from typing import Any, final
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.lattice import (
+    BoundedLattice,
+    IsSubsetEqMixin,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+from kirin.ir.attrs.abc import LatticeAttributeMeta, SingletonLatticeAttributeMeta
+from kirin.print.printer import Printer
+
+from ._visitor import _ElemVisitor
+
+
+@dataclass
+class Result(
+    ir.Attribute,
+    IsSubsetEqMixin["Result"],
+    SimpleJoinMixin["Result"],
+    SimpleMeetMixin["Result"],
+    BoundedLattice["Result"],
+    _ElemVisitor,
+    metaclass=LatticeAttributeMeta,
+):
+    """Base class for constant analysis results."""
+
+    @classmethod
+    def top(cls) -> "Result":
+        return Unknown()
+
+    @classmethod
+    def bottom(cls) -> "Result":
+        return Bottom()
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.plain_print(repr(self))
+
+
+@final
+@dataclass
+class Unknown(Result, metaclass=SingletonLatticeAttributeMeta):
+    """Unknown constant value. This is the top element of the lattice."""
+
+    def is_subseteq(self, other: Result) -> bool:
+        return isinstance(other, Unknown)
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+@final
+@dataclass
+class Bottom(Result, metaclass=SingletonLatticeAttributeMeta):
+    """Bottom element of the lattice."""
+
+    def is_subseteq(self, other: Result) -> bool:
+        return True
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+@final
+@dataclass
+class Value(Result):
+    """Constant value. Wraps any hashable Python value."""
+
+    data: Any
+
+    def is_subseteq_Value(self, other: "Value") -> bool:
+        return self.data == other.data
+
+    def is_equal(self, other: Result) -> bool:
+        if isinstance(other, Value):
+            return self.data == other.data
+        return False
+
+    def __hash__(self) -> int:
+        # NOTE: we use id here because the data
+        # may not be hashable. This is fine because
+        # the data is guaranteed to be unique.
+        return id(self)
+
+
+@dataclass
+class PartialConst(Result):
+    """Base class for partial constant values."""
+
+    pass
+
+
+@final
+class PartialTupleMeta(LatticeAttributeMeta):
+    """Metaclass for PartialTuple.
+
+    This metaclass canonicalizes PartialTuple instances with all Value elements
+    into a single Value instance.
+    """
+
+    def __call__(cls, data: tuple[Result, ...]):
+        if all(isinstance(x, Value) for x in data):
+            return Value(tuple(x.data for x in data))  # type: ignore
+        return super().__call__(data)
+
+
+@final
+@dataclass
+class PartialTuple(PartialConst, metaclass=PartialTupleMeta):
+    """Partial tuple constant value."""
+
+    data: tuple[Result, ...]
+
+    def join(self, other: Result) -> Result:
+        if other.is_subseteq(self):
+            return self
+        elif self.is_subseteq(other):
+            return other
+        elif isinstance(other, PartialTuple):
+            return PartialTuple(tuple(x.join(y) for x, y in zip(self.data, other.data)))
+        elif isinstance(other, Value) and isinstance(other.data, tuple):
+            return PartialTuple(
+                tuple(x.join(Value(y)) for x, y in zip(self.data, other.data))
+            )
+        return Unknown()
+
+    def meet(self, other: Result) -> Result:
+        if self.is_subseteq(other):
+            return self
+        elif other.is_subseteq(self):
+            return other
+        elif isinstance(other, PartialTuple):
+            return PartialTuple(tuple(x.meet(y) for x, y in zip(self.data, other.data)))
+        elif isinstance(other, Value) and isinstance(other.data, tuple):
+            return PartialTuple(
+                tuple(x.meet(Value(y)) for x, y in zip(self.data, other.data))
+            )
+        return self.bottom()
+
+    def is_equal(self, other: Result) -> bool:
+        if isinstance(other, PartialTuple):
+            return all(x.is_equal(y) for x, y in zip(self.data, other.data))
+        elif isinstance(other, Value) and isinstance(other.data, tuple):
+            return all(x.is_equal(Value(y)) for x, y in zip(self.data, other.data))
+        return False
+
+    def is_subseteq_PartialTuple(self, other: "PartialTuple") -> bool:
+        return all(x.is_subseteq(y) for x, y in zip(self.data, other.data))
+
+    def is_subseteq_Value(self, other: Value) -> bool:
+        if isinstance(other.data, tuple):
+            return all(x.is_subseteq(Value(y)) for x, y in zip(self.data, other.data))
+        return False
+
+    def __hash__(self) -> int:
+        return hash(self.data)
+
+
+@final
+@dataclass
+class PartialLambda(PartialConst):
+    """Partial lambda constant value.
+
+    This represents a closure with captured variables.
+    """
+
+    argnames: list[str]
+    code: ir.Statement
+    captured: tuple[Result, ...]
+
+    def __hash__(self) -> int:
+        return hash((self.argnames, self.code, self.captured))
+
+    def is_subseteq_PartialLambda(self, other: "PartialLambda") -> bool:
+        if self.code is not other.code:
+            return False
+        if len(self.captured) != len(other.captured):
+            return False
+
+        return all(x.is_subseteq(y) for x, y in zip(self.captured, other.captured))
+
+    def join(self, other: Result) -> Result:
+        if other is other.bottom():
+            return self
+
+        if not isinstance(other, PartialLambda):
+            return Unknown().join(other)  # widen self
+
+        if self.code is not other.code:
+            return Unknown()  # lambda stmt is pure
+
+        if len(self.captured) != len(other.captured):
+            return self.bottom()  # err
+
+        return PartialLambda(
+            self.argnames,
+            self.code,
+            tuple(x.join(y) for x, y in zip(self.captured, other.captured)),
+        )
+
+    def meet(self, other: Result) -> Result:
+        if not isinstance(other, PartialLambda):
+            return Unknown().meet(other)
+
+        if self.code is not other.code:
+            return self.bottom()
+
+        if len(self.captured) != len(other.captured):
+            return Unknown()
+
+        return PartialLambda(
+            self.argnames,
+            self.code,
+            tuple(x.meet(y) for x, y in zip(self.captured, other.captured)),
+        )

kirin/analysis/const/prop.py

@@ -0,0 +1,116 @@
+from dataclasses import field, dataclass
+
+from kirin import ir, types, interp
+from kirin.analysis.forward import ForwardExtra, ForwardFrame
+
+from .lattice import Value, Result, Unknown
+
+
+@dataclass
+class Frame(ForwardFrame[Result]):
+    should_be_pure: set[ir.Statement] = field(default_factory=set)
+    """If any ir.MaybePure is actually pure."""
+    frame_is_not_pure: bool = False
+    """If we hit any non-pure statement."""
+
+
+@dataclass
+class Propagate(ForwardExtra[Frame, Result]):
+    """Forward dataflow analysis for constant propagation.
+
+    This analysis is a forward dataflow analysis that propagates constant values
+    through the program. It uses the `Result` lattice to track the constant
+    values and purity of the values.
+
+    The analysis is implemented as a forward dataflow analysis, where the
+    `eval_stmt` method is overridden to handle the different types of statements
+    in the IR. The analysis uses the `interp.Interpreter` to evaluate the
+    statements and propagate the constant values.
+
+    When a statement is registered under the "constprop" key in the method table,
+    the analysis will call the method to evaluate the statement instead of using
+    the interpreter. This allows for custom handling of statements.
+    """
+
+    keys = ["constprop"]
+    lattice = Result
+
+    _interp: interp.Interpreter = field(init=False)
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self._interp = interp.Interpreter(
+            self.dialects,
+            fuel=self.fuel,
+            debug=self.debug,
+            max_depth=self.max_depth,
+            max_python_recursion_depth=self.max_python_recursion_depth,
+        )
+
+    def initialize(self):
+        super().initialize()
+        self._interp.initialize()
+        return self
+
+    def new_frame(self, code: ir.Statement) -> Frame:
+        return Frame.from_func_like(code)
+
+    def _try_eval_const_pure(
+        self,
+        frame: Frame,
+        stmt: ir.Statement,
+        values: tuple[Value, ...],
+    ) -> interp.StatementResult[Result]:
+        _frame = self._interp.new_frame(frame.code)
+        _frame.set_values(stmt.args, tuple(x.data for x in values))
+        method = self._interp.lookup_registry(frame, stmt)
+        if method is not None:
+            value = method(self._interp, _frame, stmt)
+        else:
+            return (Unknown(),)
+        match value:
+            case tuple():
+                return tuple(Value(each) for each in value)
+            case interp.ReturnValue(ret):
+                return interp.ReturnValue(Value(ret))
+            case interp.YieldValue(yields):
+                return interp.YieldValue(tuple(Value(each) for each in yields))
+            case interp.Successor(block, args):
+                return interp.Successor(
+                    block,
+                    *tuple(Value(each) for each in args),
+                )
+
+    def eval_stmt(
+        self, frame: Frame, stmt: ir.Statement
+    ) -> interp.StatementResult[Result]:
+        if stmt.has_trait(ir.ConstantLike):
+            return self._try_eval_const_pure(frame, stmt, ())
+        elif stmt.has_trait(ir.Pure):
+            values = frame.get_values(stmt.args)
+            if types.is_tuple_of(values, Value):
+                return self._try_eval_const_pure(frame, stmt, values)
+
+        method = self.lookup_registry(frame, stmt)
+        if method is None:
+            if stmt.has_trait(ir.Pure):
+                return (Unknown(),)  # no implementation but pure
+            # not pure, and no implementation, let's say it's not pure
+            frame.frame_is_not_pure = True
+            return (Unknown(),)
+
+        ret = method(self, frame, stmt)
+        if stmt.has_trait(ir.IsTerminator):
+            return ret
+        elif not stmt.has_trait(ir.MaybePure):  # cannot be pure at all
+            frame.frame_is_not_pure = True
+        elif (
+            stmt not in frame.should_be_pure
+        ):  # implementation cannot decide if it's pure
+            frame.frame_is_not_pure = True
+        return ret
+
+    def run_method(
+        self, method: ir.Method, args: tuple[Result, ...]
+    ) -> tuple[Frame, Result]:
+        return self.run_callable(method.code, (Value(method),) + args)

kirin/analysis/forward.py

@@ -0,0 +1,100 @@
+import sys
+from abc import ABC
+from typing import TypeVar, Iterable
+from dataclasses import dataclass
+
+from kirin import ir, interp
+from kirin.interp import AbstractFrame, AbstractInterpreter
+from kirin.lattice import BoundedLattice
+
+ExtraType = TypeVar("ExtraType")
+LatticeElemType = TypeVar("LatticeElemType", bound=BoundedLattice)
+
+
+@dataclass
+class ForwardFrame(AbstractFrame[LatticeElemType]):
+    pass
+
+
+ForwardFrameType = TypeVar("ForwardFrameType", bound=ForwardFrame)
+
+
+@dataclass
+class ForwardExtra(
+    AbstractInterpreter[ForwardFrameType, LatticeElemType],
+    ABC,
+):
+    """Abstract interpreter but record results for each SSA value.
+
+    Params:
+        LatticeElemType: The lattice element type.
+        ExtraType: The type of extra information to be stored in the frame.
+    """
+
+    def run_analysis(
+        self,
+        method: ir.Method,
+        args: tuple[LatticeElemType, ...] | None = None,
+    ) -> tuple[ForwardFrameType, LatticeElemType]:
+        """Run the forward dataflow analysis.
+
+        Args:
+            method(ir.Method): The method to analyze.
+            args(tuple[LatticeElemType]): The arguments to the method. Defaults to tuple of top values.
+
+        Returns:
+            ForwardFrameType: The results of the analysis contained in the frame.
+            LatticeElemType: The result of the analysis for the method return value.
+        """
+        args = args or tuple(self.lattice.top() for _ in method.args)
+
+        if self._eval_lock:
+            raise interp.InterpreterError(
+                "recursive eval is not allowed, use run_method instead"
+            )
+
+        self._eval_lock = True
+        self.initialize()
+        current_recursion_limit = sys.getrecursionlimit()
+        sys.setrecursionlimit(self.max_python_recursion_depth)
+        try:
+            frame, ret = self.run_method(method, args)
+        except interp.InterpreterError:
+            # NOTE: initialize will create new State
+            # so we don't need to copy the frames.
+            return self.new_frame(method.code), self.lattice.bottom()
+        finally:
+            self._eval_lock = False
+            sys.setrecursionlimit(current_recursion_limit)
+        return frame, ret
+
+    def set_values(
+        self,
+        frame: AbstractFrame[LatticeElemType],
+        ssa: Iterable[ir.SSAValue],
+        results: Iterable[LatticeElemType],
+    ):
+        """Set the abstract values for the given SSA values in the frame.
+
+        This method is used to customize how the abstract values are set in
+        the frame. By default, the abstract values are set directly in the
+        frame. This method is overridden to join the results if the SSA value
+        already exists in the frame.
+        """
+        for ssa_value, result in zip(ssa, results):
+            if ssa_value in frame.entries:
+                frame.entries[ssa_value] = frame.entries[ssa_value].join(result)
+            else:
+                frame.entries[ssa_value] = result
+
+
+class Forward(ForwardExtra[ForwardFrame[LatticeElemType], LatticeElemType], ABC):
+    """Forward dataflow analysis.
+
+    This is the base class for forward dataflow analysis. If your analysis
+    requires extra information per frame, you should subclass
+    [`ForwardExtra`][kirin.analysis.forward.ForwardExtra] instead.
+    """
+
+    def new_frame(self, code: ir.Statement) -> ForwardFrame[LatticeElemType]:
+        return ForwardFrame.from_func_like(code)

kirin/analysis/typeinfer/__init__.py

@@ -0,0 +1,5 @@
+"""Type inference analysis for kirin.
+"""
+
+from .solve import TypeResolution as TypeResolution
+from .analysis import TypeInference as TypeInference