kirin-toolchain 0.13.0__py3-none-any.whl

kirin/print/printer.py

@@ -0,0 +1,415 @@
+import io
+from typing import (
+    IO,
+    TYPE_CHECKING,
+    Any,
+    Union,
+    Literal,
+    TypeVar,
+    Callable,
+    Iterable,
+    Generator,
+)
+from contextlib import contextmanager
+from dataclasses import field, dataclass
+
+from rich.theme import Theme
+from rich.console import Console
+
+from kirin.idtable import IdTable
+from kirin.print.printable import Printable
+
+if TYPE_CHECKING:
+    from kirin import ir
+
+
+DEFAULT_THEME = {
+    "dark": Theme(
+        {
+            "dialect": "dark_blue",
+            "type": "dark_blue",
+            "comment": "bright_black",
+            "keyword": "red",
+            "symbol": "cyan",
+            "warning": "yellow",
+            "string": "green",
+            "irrational": "default",
+            "number": "default",
+        }
+    ),
+    "light": Theme(
+        {
+            "dialect": "blue",
+            "type": "blue",
+            "comment": "bright_black",
+            "keyword": "red",
+            "symbol": "cyan",
+            "warning": "yellow",
+            "string": "green",
+            "irrational": "magenta",
+        }
+    ),
+}
+
+
+@dataclass
+class PrintState:
+    ssa_id: IdTable["ir.SSAValue"] = field(default_factory=IdTable["ir.SSAValue"])
+    block_id: IdTable["ir.Block"] = field(
+        default_factory=lambda: IdTable["ir.Block"](prefix="^")
+    )
+    indent: int = 0
+    result_width: int = 0
+    indent_marks: list[int] = field(default_factory=list)
+    result_width: int = 0
+    "SSA-value column width in printing"
+    rich_style: str | None = None
+    rich_highlight: bool | None = False
+    messages: list[str] = field(default_factory=list)
+
+
+IOType = TypeVar("IOType", bound=IO)
+
+
+def _default_console():
+    return Console(force_jupyter=False)
+
+
+@dataclass
+class Printer:
+    """A IR pretty printer build on top of Rich console."""
+
+    console: Console = field(default_factory=_default_console)
+    """Rich console"""
+    analysis: dict["ir.SSAValue", Any] | None = None
+    """Analysis results"""
+    hint: str | None = None
+    """key of the SSAValue hint to print"""
+    state: PrintState = field(default_factory=PrintState)
+    """Printing state"""
+    show_indent_mark: bool = field(default=True, kw_only=True)
+    "Whether to show indent marks, e.g │"
+    theme: Theme | dict | Literal["dark", "light"] = field(default="dark", kw_only=True)
+    "Theme to use for printing"
+
+    def __post_init__(self):
+        if isinstance(self.theme, dict):
+            self.theme = Theme(self.theme)
+        elif isinstance(self.theme, str):
+            self.theme = DEFAULT_THEME[self.theme]
+        self.console.push_theme(self.theme)
+
+    def print(self, object):
+        """entry point for printing an object
+
+        Args:
+            object: object to print.
+        """
+        if isinstance(object, Printable):
+            object.print_impl(self)
+        else:
+            fn = getattr(self, f"print_{object.__class__.__name__}", None)
+            if fn is None:
+                raise NotImplementedError(
+                    f"Printer for {object.__class__.__name__} not found"
+                )
+            fn(object)
+
+    def print_name(
+        self, node: Union["ir.Attribute", "ir.Statement"], prefix: str = ""
+    ) -> None:
+        """print the name of a node
+
+        Args:
+            node(ir.Attribute | ir.Statement): node to print
+            prefix(str): prefix to print before the name, default to ""
+        """
+        self.print_dialect_path(node, prefix=prefix)
+        if node.dialect:
+            self.plain_print(".")
+        self.plain_print(node.name)
+
+    def print_stmt(self, node: "ir.Statement"):
+        if node._results:
+            result_str = self.result_str(node._results)
+            self.plain_print(result_str.rjust(self.state.result_width), " = ")
+        elif self.state.result_width:
+            self.plain_print(" " * self.state.result_width, "   ")
+        with self.indent(self.state.result_width + 3, mark=True):
+            self.print(node)
+            with self.rich(style="warning"):
+                self.print_analysis(*node._results, prefix=" # ---> ")
+            with self.rich(style="comment"):
+                self.print_hint(*node._results)
+
+    def print_hint(
+        self,
+        *values: "ir.SSAValue",
+        prefix: str = " // hint<",
+        suffix: str = ">",
+    ):
+        if not self.hint or not values:
+            return
+
+        self.plain_print(prefix)
+        self.plain_print(self.hint)
+        for idx, item in enumerate(values):
+            if idx > 0:
+                self.plain_print(", ")
+
+            self.plain_print("=")
+            if item.hints.get(self.hint) is not None:
+                self.plain_print(repr(item.hints.get(self.hint)))
+            else:
+                self.plain_print("missing")
+        self.plain_print(suffix)
+
+    def print_analysis(
+        self,
+        *values: "ir.SSAValue",
+        prefix: str = "",
+    ):
+        if self.analysis is None or not values:
+            return
+
+        self.plain_print(prefix)
+        for idx, value in enumerate(values):
+            if idx > 0:
+                self.plain_print(", ")
+
+            if result := self.analysis.get(value):
+                self.plain_print(repr(result))
+            else:
+                self.plain_print("missing")
+
+    def print_dialect_path(
+        self, node: Union["ir.Attribute", "ir.Statement"], prefix: str = ""
+    ) -> None:
+        """print the dialect path of a node.
+
+        Args:
+            node(ir.Attribute | ir.Statement): node to print
+            prefix(str): prefix to print before the dialect path, default to ""
+        """
+        if node.dialect:  # not None
+            self.plain_print(prefix)
+            self.plain_print(node.dialect.name, style="dialect")
+        else:
+            self.plain_print(prefix)
+
+    def print_newline(self):
+        """print a newline character.
+
+        This method also prints any messages in the state for debugging.
+        """
+        self.plain_print("\n")
+
+        if self.state.messages:
+            for message in self.state.messages:
+                self.plain_print(message)
+                self.plain_print("\n")
+            self.state.messages.clear()
+        self.print_indent()
+
+    def print_indent(self):
+        """print the current indentation level optionally with indent marks."""
+        indent_str = ""
+        if self.show_indent_mark and self.state.indent_marks:
+            indent_str = "".join(
+                "│" if i in self.state.indent_marks else " "
+                for i in range(self.state.indent)
+            )
+            with self.rich(style="comment"):
+                self.plain_print(indent_str)
+        else:
+            indent_str = " " * self.state.indent
+            self.plain_print(indent_str)
+
+    def plain_print(self, *objects, sep="", end="", style=None, highlight=None):
+        """print objects without any formatting.
+
+        Args:
+            *objects: objects to print
+
+        Keyword Args:
+            sep(str): separator between objects, default to ""
+            end(str): end character, default to ""
+            style(str): style to use, default to None
+            highlight(bool): whether to highlight the text, default to None
+        """
+        self.console.out(
+            *objects,
+            sep=sep,
+            end=end,
+            style=style or self.state.rich_style,
+            highlight=highlight or self.state.rich_highlight,
+        )
+
+    ElemType = TypeVar("ElemType")
+
+    def print_seq(
+        self,
+        seq: Iterable[ElemType],
+        *,
+        emit: Callable[[ElemType], None] | None = None,
+        delim: str = ", ",
+        prefix: str = "",
+        suffix: str = "",
+        style=None,
+        highlight=None,
+    ) -> None:
+        """print a sequence of objects.
+
+        Args:
+            seq(Iterable[ElemType]): sequence of objects to print
+
+        Keyword Args:
+            emit(Callable[[ElemType], None]): function to print each element, default to None
+            delim(str): delimiter between elements, default to ", "
+            prefix(str): prefix to print before the sequence, default to ""
+            suffix(str): suffix to print after the sequence, default to ""
+            style(str): style to use, default to None
+            highlight(bool): whether to highlight the text, default to None
+        """
+        emit = emit or self.print
+        self.plain_print(prefix, style=style, highlight=highlight)
+        for idx, item in enumerate(seq):
+            if idx > 0:
+                self.plain_print(delim, style=style)
+            emit(item)
+        self.plain_print(suffix, style=style, highlight=highlight)
+
+    KeyType = TypeVar("KeyType")
+    ValueType = TypeVar("ValueType", bound=Printable)
+
+    def print_mapping(
+        self,
+        elems: dict[KeyType, ValueType],
+        *,
+        emit: Callable[[ValueType], None] | None = None,
+        delim: str = ", ",
+    ) -> None:
+        """print a mapping of key-value pairs.
+
+        Args:
+            elems(dict[KeyType, ValueType]): mapping to print
+
+        Keyword Args:
+            emit(Callable[[ValueType], None]): function to print each value, default to None
+            delim(str): delimiter between key-value pairs, default to ", "
+        """
+        emit = emit or self.print
+        for i, (key, value) in enumerate(elems.items()):
+            if i > 0:
+                self.plain_print(delim)
+            self.plain_print(f"{key}=")
+            emit(value)
+
+    def result_width(self, stmts: Iterable["ir.Statement"]) -> int:
+        """return the maximum width of the result column for a sequence of statements.
+
+        Args:
+            stmts(Iterable[ir.Statement]): sequence of statements
+
+        Returns:
+            int: maximum width of the result column
+        """
+        result_width = 0
+        for stmt in stmts:
+            result_width = max(result_width, len(self.result_str(stmt._results)))
+        return result_width
+
+    @contextmanager
+    def align(self, width: int) -> Generator[PrintState, Any, None]:
+        """align the result column width, and restore it after the context.
+
+        Args:
+            width(int): width of the column
+
+        Yields:
+            PrintState: the state with the new column width
+        """
+        old_width = self.state.result_width
+        self.state.result_width = width
+        try:
+            yield self.state
+        finally:
+            self.state.result_width = old_width
+
+    @contextmanager
+    def indent(
+        self, increase: int = 2, mark: bool | None = None
+    ) -> Generator[PrintState, Any, None]:
+        """increase the indentation level, and restore it after the context.
+
+        Args:
+            increase(int): amount to increase the indentation level, default to 2
+            mark(bool): whether to mark the indentation level, default to None
+
+        Yields:
+            PrintState: the state with the new indentation level.
+        """
+        mark = mark if mark is not None else self.show_indent_mark
+        self.state.indent += increase
+        if mark:
+            self.state.indent_marks.append(self.state.indent)
+        try:
+            yield self.state
+        finally:
+            self.state.indent -= increase
+            if mark:
+                self.state.indent_marks.pop()
+
+    @contextmanager
+    def rich(
+        self, style: str | None = None, highlight: bool = False
+    ) -> Generator[PrintState, Any, None]:
+        """set the rich style and highlight, and restore them after the context.
+
+        Args:
+            style(str | None): style to use, default to None
+            highlight(bool): whether to highlight the text, default to False
+
+        Yields:
+            PrintState: the state with the new style and highlight.
+        """
+        old_style = self.state.rich_style
+        old_highlight = self.state.rich_highlight
+        self.state.rich_style = style
+        self.state.rich_highlight = highlight
+        try:
+            yield self.state
+        finally:
+            self.state.rich_style = old_style
+            self.state.rich_highlight = old_highlight
+
+    @contextmanager
+    def string_io(self) -> Generator[io.StringIO, Any, None]:
+        """Temporary string IO for capturing output.
+
+        Yields:
+            io.StringIO: the string IO object.
+        """
+        stream = io.StringIO()
+        old_file = self.console.file
+        self.console.file = stream
+        try:
+            yield stream
+        finally:
+            self.console.file = old_file
+            stream.close()
+
+    def result_str(self, results: list["ir.ResultValue"]) -> str:
+        """return the string representation of a list of result values.
+
+        Args:
+            results(list[ir.ResultValue]): list of result values to print
+        """
+        with self.string_io() as stream:
+            self.print_seq(results, delim=", ")
+            result_str = stream.getvalue()
+        return result_str
+
+    def debug(self, message: str):
+        """Print a debug message."""
+        self.state.messages.append(f"DEBUG: {message}")

kirin/registry.py

@@ -0,0 +1,105 @@
+from typing import TYPE_CHECKING, Callable, Iterable
+from dataclasses import dataclass
+
+if TYPE_CHECKING:
+    from kirin.ir import Attribute
+    from kirin.ir.group import DialectGroup
+    from kirin.ir.nodes import Statement
+    from kirin.lowering import FromPythonAST
+    from kirin.interp.impl import Signature
+    from kirin.interp.table import MethodTable
+
+
+@dataclass
+class StatementImpl:
+    parent: "MethodTable"
+    impl: Callable
+
+    def __call__(self, interp, frame, stmt: "Statement"):
+        return self.impl(self.parent, interp, frame, stmt)
+
+    def __repr__(self) -> str:
+        return f"method impl `{self.impl.__name__}` in {repr(self.parent.__class__)}"
+
+
+@dataclass
+class AttributeImpl:
+    parent: "MethodTable"
+    impl: Callable
+
+    def __call__(self, interp, attr: "Attribute"):
+        return self.impl(self.parent, interp, attr)
+
+    def __repr__(self) -> str:
+        return f"attribute impl `{self.impl.__name__}` in {repr(self.parent.__class__)}"
+
+
+@dataclass
+class InterpreterRegistry:
+    attributes: dict[type["Attribute"], "AttributeImpl"]
+    statements: dict["Signature", "StatementImpl"]
+
+
+@dataclass
+class Registry:
+    """Proxy class to build different registries from a dialect group."""
+
+    dialects: "DialectGroup"
+    """The dialect group to build the registry from."""
+
+    def ast(self, keys: Iterable[str]) -> dict[str, "FromPythonAST"]:
+        """select the dialect lowering interpreters for the given key.
+
+        Args:
+            keys (Iterable[str]): the keys to search for in the dialects
+
+        Returns:
+            a map of dialects to their lowering interpreters
+        """
+        ret: dict[str, "FromPythonAST"] = {}
+        from_ast = None
+        for dialect in self.dialects.data:
+            for key in keys:
+                if key in dialect.lowering:
+                    from_ast = dialect.lowering[key]
+                    break
+
+            if from_ast is None:
+                msg = ",".join(keys)
+                raise KeyError(f"Lowering not found for {msg}")
+
+            for name in from_ast.names:
+                if name in ret:
+                    raise KeyError(f"Lowering {name} already exists")
+
+                ret[name] = from_ast
+        return ret
+
+    def interpreter(self, keys: Iterable[str]):
+        """select the dialect interpreter for the given key.
+
+        Args:
+            keys (Iterable[str]): the keys to search for in the dialects
+
+        Returns:
+            a map of statement signatures to their interpretation functions,
+            and a map of dialects to their fallback interpreters.
+        """
+        attributes: dict[type["Attribute"], "AttributeImpl"] = {}
+        table: dict["Signature", "StatementImpl"] = {}
+        for dialect in self.dialects.data:
+            dialect_table = None
+            for key in keys:
+                if key not in dialect.interps:
+                    continue
+
+                dialect_table = dialect.interps[key]
+                for sig, func in dialect_table.attribute.items():
+                    if sig not in attributes:
+                        attributes[sig] = AttributeImpl(dialect_table, func)
+
+                for sig, func in dialect_table.table.items():
+                    if sig not in table:
+                        table[sig] = StatementImpl(dialect_table, func)
+
+        return InterpreterRegistry(attributes, table)

kirin/registry.pyi

@@ -0,0 +1,52 @@
+from typing import Generic, TypeVar, Callable, Iterable, TypeAlias
+from dataclasses import dataclass
+
+from kirin.ir.group import DialectGroup
+from kirin.ir.nodes import Statement
+from kirin.lowering import FromPythonAST
+from kirin.interp.base import FrameABC, BaseInterpreter
+from kirin.interp.impl import Signature
+from kirin.interp.table import MethodTable
+from kirin.interp.value import StatementResult
+from kirin.ir.attrs.abc import Attribute
+
+MethodTableSelf = TypeVar("MethodTableSelf", bound="MethodTable")
+InterpreterType = TypeVar("InterpreterType", bound="BaseInterpreter")
+FrameType = TypeVar("FrameType", bound="FrameABC")
+StatementType = TypeVar("StatementType", bound="Statement")
+MethodFunction: TypeAlias = Callable[
+    [MethodTableSelf, InterpreterType, FrameType, StatementType], StatementResult
+]
+
+@dataclass
+class StatementImpl(Generic[InterpreterType, FrameType]):
+    parent: "MethodTable"
+    impl: MethodFunction["MethodTable", InterpreterType, FrameType, "Statement"]
+
+    def __call__(
+        self, interp: InterpreterType, frame: FrameType, stmt: "Statement"
+    ) -> StatementResult: ...
+    def __repr__(self) -> str: ...
+
+@dataclass
+class AttributeImpl:
+    parent: "MethodTable"
+    impl: Callable
+
+    def __call__(self, interp, attr: "Attribute"): ...
+    def __repr__(self) -> str: ...
+
+@dataclass
+class InterpreterRegistry:
+    attributes: dict[type["Attribute"], "AttributeImpl"]
+    statements: dict["Signature", "StatementImpl"]
+
+@dataclass
+class Registry:
+    """Proxy class to build different registries from a dialect group."""
+
+    dialects: "DialectGroup"
+    """The dialect group to build the registry from."""
+
+    def ast(self, keys: Iterable[str]) -> dict[str, "FromPythonAST"]: ...
+    def interpreter(self, keys: Iterable[str]) -> InterpreterRegistry: ...

kirin/rewrite/__init__.py

@@ -0,0 +1,14 @@
+from .cse import CommonSubexpressionElimination as CommonSubexpressionElimination
+from .dce import DeadCodeElimination as DeadCodeElimination
+from .fold import ConstantFold as ConstantFold
+from .walk import Walk as Walk
+from .alias import InlineAlias as InlineAlias
+from .chain import Chain as Chain
+from .inline import Inline as Inline
+from .getitem import InlineGetItem as InlineGetItem
+from .fixpoint import Fixpoint as Fixpoint
+from .getfield import InlineGetField as InlineGetField
+from .apply_type import ApplyType as ApplyType
+from .compactify import CFGCompactify as CFGCompactify
+from .wrap_const import WrapConst as WrapConst
+from .call2invoke import Call2Invoke as Call2Invoke

kirin/rewrite/abc.py

@@ -0,0 +1,43 @@
+from abc import ABC
+from dataclasses import dataclass
+
+from kirin.ir import Pure, Block, IRNode, Region, MaybePure, Statement
+from kirin.rewrite.result import RewriteResult
+
+
+@dataclass(repr=False)
+class RewriteRule(ABC):
+    """A rewrite rule that matches and rewrites IR nodes.
+
+    The rewrite rule is applied to an IR node by calling the instance with the node as an argument.
+    The rewrite rule should mutate the node instead of returning a new node. A `RewriteResult` should
+    be returned to indicate whether the rewrite rule has done something, whether the rewrite rule
+    should terminate, and whether the rewrite rule has exceeded the maximum number of iterations.
+    """
+
+    def rewrite(self, node: IRNode) -> RewriteResult:
+        if isinstance(node, Region):
+            return self.rewrite_Region(node)
+        elif isinstance(node, Block):
+            return self.rewrite_Block(node)
+        elif isinstance(node, Statement):
+            return self.rewrite_Statement(node)
+        else:
+            return RewriteResult()
+
+    def rewrite_Region(self, node: Region) -> RewriteResult:
+        return RewriteResult()
+
+    def rewrite_Block(self, node: Block) -> RewriteResult:
+        return RewriteResult()
+
+    def rewrite_Statement(self, node: Statement) -> RewriteResult:
+        return RewriteResult()
+
+    def is_pure(self, node: Statement):
+        if node.has_trait(Pure):
+            return True
+
+        if (trait := node.get_trait(MaybePure)) and trait.is_pure(node):
+            return True
+        return False

kirin/rewrite/aggressive/__init__.py

@@ -0,0 +1 @@
+from .fold import Fold as Fold

kirin/rewrite/aggressive/fold.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+from kirin.rewrite import Walk, Chain, Fixpoint
+from kirin.analysis import const
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.rewrite.dce import DeadCodeElimination
+from kirin.rewrite.fold import ConstantFold
+from kirin.ir.nodes.base import IRNode
+from kirin.rewrite.inline import Inline
+from kirin.rewrite.getitem import InlineGetItem
+from kirin.rewrite.getfield import InlineGetField
+from kirin.rewrite.compactify import CFGCompactify
+from kirin.rewrite.wrap_const import WrapConst
+from kirin.rewrite.call2invoke import Call2Invoke
+
+
+@dataclass
+class Fold(RewriteRule):
+    rule: RewriteRule
+
+    def __init__(self, frame: const.Frame):
+        rule = Fixpoint(
+            Chain(
+                Walk(WrapConst(frame)),
+                Walk(Inline(lambda _: True)),
+                Walk(ConstantFold()),
+                Walk(Call2Invoke()),
+                Fixpoint(
+                    Walk(
+                        Chain(
+                            InlineGetItem(),
+                            InlineGetField(),
+                            DeadCodeElimination(),
+                        )
+                    )
+                ),
+                Walk(CFGCompactify()),
+            )
+        )
+        self.rule = rule
+
+    def rewrite(self, node: IRNode) -> RewriteResult:
+        return self.rule.rewrite(node)

kirin/rewrite/alias.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.py.assign import Alias
+
+
+@dataclass
+class InlineAlias(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, Alias):
+            return RewriteResult()
+
+        node.result.replace_by(node.value)
+        return RewriteResult(has_done_something=True)