kirin-toolchain 0.13.0__py3-none-any.whl

kirin/rewrite/inline.py

@@ -0,0 +1,143 @@
+from typing import Callable
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import cf, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+# TODO: use func.Constant instead of kirin.dialects.py.stmts.Constant
+from kirin.dialects.py.constant import Constant
+
+# NOTE: this only inlines func dialect
+
+
+@dataclass
+class Inline(RewriteRule):
+    heuristic: Callable[[ir.IRNode], bool]
+    """inline heuristic that determines whether a function should be inlined
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if isinstance(node, func.Invoke):
+            return self.rewrite_func_Invoke(node)
+        elif isinstance(node, func.Call):
+            return self.rewrite_func_Call(node)
+        else:
+            return RewriteResult()
+
+    def rewrite_func_Call(self, node: func.Call) -> RewriteResult:
+        if not isinstance(lambda_stmt := node.callee.owner, func.Lambda):
+            return RewriteResult()
+
+        # NOTE: a lambda statement is defined and used in the same scope
+        self.inline_call_like(node, tuple(node.args), lambda_stmt.body)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_func_Invoke(self, node: func.Invoke) -> RewriteResult:
+        has_done_something = False
+        callee = node.callee
+
+        if (
+            isinstance(callee, ir.Method)
+            and self.heuristic(callee.code)
+            and (call_trait := callee.code.get_trait(ir.CallableStmtInterface))
+            is not None
+        ):
+            region = call_trait.get_callable_region(callee.code)
+            func_self = Constant(node.callee)
+            func_self.result.name = node.callee.sym_name
+            func_self.insert_before(node)
+            self.inline_call_like(
+                node, (func_self.result,) + tuple(arg for arg in node.args), region
+            )
+            has_done_something = True
+
+        return RewriteResult(has_done_something=has_done_something)
+
+    def inline_call_like(
+        self,
+        call_like: ir.Statement,
+        args: tuple[ir.SSAValue, ...],
+        region: ir.Region,
+    ):
+        """
+        Inline a function call-like statement
+
+        Args:
+            call_like (ir.Statement): the call-like statement
+            args (tuple[ir.SSAValue, ...]): the arguments of the call (first one is the callee)
+            region (ir.Region): the region of the callee
+        """
+        # <stmt>
+        # <stmt>
+        # <br (a, b, c)>
+
+        # <block (a, b,c)>:
+        # <block>:
+        # <block>:
+        # <br>
+
+        # ^<block>:
+        # <stmt>
+        # <stmt>
+
+        # 1. we insert the entry block of the callee function
+        # 2. we insert the rest of the blocks into the parent region
+        # 3.1 if the return is in the entry block, means no control flow,
+        #     replace the call results with the return values
+        # 3.2 if the return is some of the blocks, means control flow,
+        #     split the current block into two, and replace the return with
+        #     the branch instruction
+        # 4. remove the call
+        if not call_like.parent_block:
+            return
+
+        if not call_like.parent_region:
+            return
+
+        # NOTE: we cannot change region because it may be used elsewhere
+        inline_region: ir.Region = region.clone()
+        parent_block: ir.Block = call_like.parent_block
+        parent_region: ir.Region = call_like.parent_region
+
+        # wrap what's after invoke into a block
+        after_block = ir.Block()
+        stmt = call_like.next_stmt
+        while stmt is not None:
+            stmt.detach()
+            after_block.stmts.append(stmt)
+            stmt = call_like.next_stmt
+
+        for result in call_like.results:
+            block_arg = after_block.args.append_from(result.type, result.name)
+            result.replace_by(block_arg)
+
+        parent_block_idx = parent_region._block_idx[parent_block]
+        entry_block = inline_region.blocks.popfirst()
+        idx, block = 0, entry_block
+        while block is not None:
+            idx += 1
+
+            if block.last_stmt and isinstance(block.last_stmt, func.Return):
+                block.last_stmt.replace_by(
+                    cf.Branch(
+                        arguments=(block.last_stmt.value,),
+                        successor=after_block,
+                    )
+                )
+
+            parent_region.blocks.insert(parent_block_idx + idx, block)
+            block = inline_region.blocks.popfirst()
+
+        parent_region.blocks.append(after_block)
+
+        # NOTE: we expect always to have an entry block
+        # but we still check for it cuz we are not gonna
+        # error for empty regions here.
+        if entry_block:
+            cf.Branch(
+                arguments=args,
+                successor=entry_block,
+            ).insert_before(call_like)
+        call_like.delete()
+        return

kirin/rewrite/result.py

@@ -0,0 +1,15 @@
+from dataclasses import field, dataclass
+
+
+@dataclass
+class RewriteResult:
+    terminated: bool = field(default=False, kw_only=True)
+    has_done_something: bool = field(default=False, kw_only=True)
+    exceeded_max_iter: bool = field(default=False, kw_only=True)
+
+    def join(self, other: "RewriteResult") -> "RewriteResult":
+        return RewriteResult(
+            terminated=self.terminated or other.terminated,
+            has_done_something=self.has_done_something or other.has_done_something,
+            exceeded_max_iter=self.exceeded_max_iter or other.exceeded_max_iter,
+        )

kirin/rewrite/walk.py

@@ -0,0 +1,83 @@
+from typing import Callable
+from dataclasses import field, dataclass
+
+from kirin.ir import Block, Region, Statement
+from kirin.worklist import WorkList
+from kirin.rewrite.abc import RewriteRule
+from kirin.ir.nodes.base import IRNode
+from kirin.rewrite.result import RewriteResult
+
+
+@dataclass
+class Walk(RewriteRule):
+    """Walk through the IR nodes and apply a rewrite rule.
+
+    The walk will apply the rewrite rule to each node in the IR tree in a depth-first order.
+
+    ### Parameters
+    - `map`: The rewrite rule to apply.
+    - `reverse`: Whether to traverse the IR tree in reverse order. Default is `False`.
+    - `region_first`: Whether to traverse the regions before the blocks. Default is `False`.
+    """
+
+    rule: RewriteRule
+    worklist: WorkList[IRNode] = field(default_factory=WorkList)
+    skip: Callable[[IRNode], bool] = field(default=lambda _: False)
+
+    # options
+    reverse: bool = field(default=False)
+    region_first: bool = field(default=False)
+
+    def rewrite(self, node: IRNode) -> RewriteResult:
+        # NOTE: because the rewrite pass may mutate the node
+        # thus we need to save the list of nodes to be processed
+        # first before we start processing them
+        self.populate_worklist(node)
+        has_done_something = False
+        subnode = self.worklist.pop()
+        while subnode is not None:
+            result = self.rule.rewrite(subnode)
+            if result.terminated:
+                return result
+
+            if result.has_done_something:
+                has_done_something = True
+            subnode = self.worklist.pop()
+        return RewriteResult(has_done_something=has_done_something)
+
+    def populate_worklist(self, node: IRNode) -> None:
+        if self.skip(node):
+            return
+
+        if isinstance(node, Statement):
+            self.populate_worklist_Statement(node)
+        elif isinstance(node, Region):
+            self.populate_worklist_Region(node)
+        elif isinstance(node, Block):
+            self.populate_worklist_Block(node)
+        else:
+            raise NotImplementedError(f"populate_worklist_{node.__class__.__name__}")
+
+    def populate_worklist_Statement(self, node: Statement) -> None:
+        if self.region_first:
+            self.worklist.append(node)
+
+        if node.regions:
+            for region in reversed(node.regions) if not self.reverse else node.regions:
+                self.populate_worklist(region)
+
+        if not self.region_first:
+            self.worklist.append(node)
+
+    def populate_worklist_Region(self, node: Region) -> None:
+        self.worklist.append(node)
+        if node.blocks:
+            for block in reversed(node.blocks) if not self.reverse else node.blocks:
+                self.populate_worklist(block)
+
+    def populate_worklist_Block(self, node: Block) -> None:
+        self.worklist.append(node)
+        stmt = node.last_stmt
+        while stmt is not None:
+            self.populate_worklist(stmt)
+            stmt = stmt.prev_stmt

kirin/rewrite/wrap_const.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class WrapConst(RewriteRule):
+    """Insert constant hints into the SSA values.
+
+    !!! note
+        This pass is not exactly the same as ConstantFold. ConstantFold
+        only rewrites the SSAValue with `const.Value` into a constant
+        statement. This rule, however, inserts the entire constant lattice
+        into the SSAValue's hints. Thus enabling other rules/analysis to
+        utilize the constant information beyond just `const.Value`.
+    """
+
+    frame: const.Frame
+
+    def wrap(self, value: ir.SSAValue) -> bool:
+        result = self.frame.entries.get(value)
+        if not result:
+            return False
+
+        const_hint = value.hints.get("const")
+        if const_hint and isinstance(const_hint, const.Result):
+            const_result = result.join(const_hint)
+            if const_result.is_equal(const_hint):
+                return False
+        else:
+            const_result = result
+        value.hints["const"] = const_result
+        return True
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        has_done_something = False
+        for arg in node.args:
+            if self.wrap(arg):
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        has_done_something = False
+        for result in node.results:
+            if self.wrap(result):
+                has_done_something = True
+
+        if (
+            trait := node.get_trait(ir.MaybePure)
+        ) and node in self.frame.should_be_pure:
+            trait.set_pure(node)
+            has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)

kirin/source.py

@@ -0,0 +1,21 @@
+import ast
+from dataclasses import dataclass
+
+
+@dataclass
+class SourceInfo:
+    lineno: int
+    col_offset: int
+    end_lineno: int | None
+    end_col_offset: int | None
+
+    @classmethod
+    def from_ast(cls, node: ast.AST, lineno_offset: int = 0, col_offset: int = 0):
+        end_lineno = getattr(node, "end_lineno", None)
+        end_col_offset = getattr(node, "end_col_offset", None)
+        return cls(
+            getattr(node, "lineno", 0) + lineno_offset,
+            getattr(node, "col_offset", 0) + col_offset,
+            end_lineno + lineno_offset if end_lineno is not None else None,
+            end_col_offset + col_offset if end_col_offset is not None else None,
+        )

kirin/symbol_table.py

@@ -0,0 +1,27 @@
+from typing import Generic, TypeVar
+from dataclasses import field, dataclass
+
+T = TypeVar("T")
+
+
+@dataclass
+class SymbolTable(Generic[T]):
+    names: dict[str, T] = field(default_factory=dict)
+    """The table that maps names to values."""
+    prefix: str = field(default="", kw_only=True)
+    name_count: dict[str, int] = field(default_factory=dict, kw_only=True)
+    """The count of names that have been requested."""
+
+    def __getitem__(self, name: str) -> T:
+        return self.names[name]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.names
+
+    def __setitem__(self, name: str, value: T) -> None:
+        count = self.name_count.setdefault(name, 0)
+        self.name_count[name] = count + 1
+        self.names[f"{self.prefix}_{name}_{count}"] = value
+
+    def __delitem__(self, name: str) -> None:
+        del self.names[name]

kirin/types.py

@@ -0,0 +1,34 @@
+"""Bindings for built-in types.
+"""
+
+from kirin.ir.method import Method
+from kirin.ir.attrs.types import (
+    Union as Union,
+    Vararg as Vararg,
+    AnyType as AnyType,
+    Generic as Generic,
+    Literal as Literal,
+    PyClass as PyClass,
+    TypeVar as TypeVar,
+    BottomType as BottomType,
+    TypeAttribute as TypeAttribute,
+    hint2type as hint2type,
+    is_tuple_of as is_tuple_of,
+)
+
+Any = AnyType()
+Bottom = BottomType()
+Int = PyClass(int)
+Float = PyClass(float)
+String = PyClass(str)
+Bool = PyClass(bool)
+NoneType = PyClass(type(None))
+List = Generic(list, TypeVar("T"))
+Slice = Generic(slice, TypeVar("T"))
+Tuple = Generic(tuple, Vararg(TypeVar("T")))
+Dict = Generic(dict, TypeVar("K"), TypeVar("V"))
+Set = Generic(set, TypeVar("T"))
+FrozenSet = Generic(frozenset, TypeVar("T"))
+TypeofFunctionType = Generic[type(lambda: None)]
+FunctionType = Generic(type(lambda: None), Tuple, Vararg(Any))
+MethodType = Generic(Method, TypeVar("Params", Tuple), TypeVar("Ret"))

kirin/worklist.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import Generic, TypeVar, Iterable
+from dataclasses import field, dataclass
+
+ElemType = TypeVar("ElemType")
+
+
+@dataclass
+class WorkList(Generic[ElemType]):
+    """The worklist data structure.
+
+    The worklist is a stack that allows for O(1) removal of elements from the stack.
+    """
+
+    _stack: list[ElemType] = field(default_factory=list, init=False)
+
+    def is_empty(self) -> bool:
+        return len(self._stack) == 0
+
+    def append(self, item: ElemType) -> None:
+        self._stack.append(item)
+
+    def extend(self, items: Iterable[ElemType]) -> None:
+        self._stack.extend(items)
+
+    def pop(self) -> ElemType | None:
+        if self._stack:
+            return self._stack.pop()
+        return None

kirin_toolchain-0.13.0.dist-info/METADATA

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: kirin-toolchain
+Version: 0.13.0
+Summary: The Kirin Toolchain for building compilers and interpreters.
+Author-email: Roger-luo <rluo@quera.com>
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: beartype>=0.17.2
+Requires-Dist: rich>=13.7.1
+Requires-Dist: typing-extensions>=4.11.0
+Description-Content-Type: text/markdown
+
+# KIRIN
+
+[![CI](https://github.com/QuEraComputing/kirin/actions/workflows/ci.yml/badge.svg)](https://github.com/QuEraComputing/kirin/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/QuEraComputing/kirin/graph/badge.svg?token=lkUZ9DTqy4)](https://codecov.io/gh/QuEraComputing/kirin)
+[![Supported Python versions](https://img.shields.io/pypi/pyversions/kirin-toolchain.svg?color=%2334D058)](https://pypi.org/project/fastapi)
+
+*K*ernel *I*ntermediate *R*epresentation *IN*frastructure
+
+> [!IMPORTANT]
+>
+> This project is in the early stage of development. API and features are subject to change.
+> If you are concerned about the stability of the APIs, consider pin the version of Kirin in your project.
+
+## Installation
+
+### Install via `uv` (Recommended)
+
+```py
+uv add kirin-toolchain
+```
+
+### Install via pip
+
+```bash
+pip install kirin-toolchain
+```
+
+## License
+
+Apache License 2.0 with LLVM Exceptions