kirin-toolchain 0.13.0__py3-none-any.whl

kirin/rewrite/apply_type.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+
+from kirin import ir, types
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.func.attrs import Signature
+
+
+@dataclass
+class ApplyType(RewriteRule):
+    results: dict[ir.SSAValue, types.TypeAttribute]
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        has_done_something = False
+        for arg in node.args:
+            if arg in self.results:
+                arg.type = self.results[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 result in self.results:
+                result.type = self.results[result]
+                has_done_something = True
+
+        if (trait := node.get_trait(ir.HasSignature)) is not None and (
+            callable_trait := node.get_trait(ir.CallableStmtInterface)
+        ) is not None:
+            callable_region = callable_trait.get_callable_region(node)
+            inputs = tuple(
+                self.results.get(arg, arg.type)
+                for arg in callable_region.blocks[0].args
+            )
+
+            if (
+                len(node._results) == 1
+                and isinstance(
+                    output_ := self.results.get(node._results[0]), types.Generic
+                )
+                and output_.is_subseteq(types.MethodType)
+            ):
+                output_ = output_.vars[1]
+                trait.set_signature(node, Signature(inputs, output_))
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)

kirin/rewrite/call2invoke.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.func import Call, Invoke
+
+
+@dataclass
+class Call2Invoke(RewriteRule):
+    """Rewrite a `Call` statement to an `Invoke` statement."""
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, Call):
+            return RewriteResult()
+
+        if (mt := node.callee.hints.get("const")) is None:
+            return RewriteResult()
+
+        if not isinstance(mt, const.Value):
+            return RewriteResult()
+
+        if not isinstance(mt.data, ir.Method):
+            return RewriteResult()
+
+        stmt = Invoke(inputs=node.inputs, callee=mt.data, kwargs=node.kwargs)
+        for result, new_result in zip(node.results, stmt.results):
+            new_result.name = result.name
+            new_result.type = result.type
+            if result_hint := result.hints.get("const"):
+                new_result.hints["const"] = result_hint
+
+        node.replace_by(stmt)
+        return RewriteResult(has_done_something=True)

kirin/rewrite/chain.py

@@ -0,0 +1,39 @@
+from typing import Iterable
+from dataclasses import dataclass
+
+from kirin.ir import IRNode
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+
+
+@dataclass
+class Chain(RewriteRule):
+    """Chain multiple rewrites together.
+
+    The chain will apply each rewrite in order until one of the rewrites terminates.
+    """
+
+    rules: list[RewriteRule]
+
+    def __init__(self, rule: RewriteRule | Iterable[RewriteRule], *others: RewriteRule):
+        if isinstance(rule, RewriteRule):
+            self.rules = [rule, *others]
+        else:
+            assert (
+                others == ()
+            ), "Cannot pass multiple positional arguments if the first argument is an iterable"
+            self.rules = list(rule)
+
+    def rewrite(self, node: IRNode) -> RewriteResult:
+        has_done_something = False
+        for rule in self.rules:
+            result = rule.rewrite(node)
+            if result.terminated:
+                return result
+
+            if result.has_done_something:
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)
+
+    def __repr__(self):
+        return " -> ".join(map(str, self.rules))

kirin/rewrite/compactify.py

@@ -0,0 +1,288 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import cf
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.analysis.cfg import CFG
+from kirin.rewrite.walk import Walk
+from kirin.rewrite.chain import Chain
+from kirin.rewrite.fixpoint import Fixpoint
+
+
+@dataclass
+class DeadBlock(RewriteRule):
+    """Compactify the CFG by removing dead blocks."""
+
+    cfg: CFG
+
+    def rewrite_Region(self, node: ir.Region) -> RewriteResult:
+        # remove non-entry blocks that are not reachable from the entry block
+        # TODO: check if this region is using SSACFG convention?
+        has_done_something = False
+        for block in node.blocks[1:]:
+            predecessors = self.cfg.predecessors.get(block)
+            if not predecessors:  # empty predecessors
+                successors = self.cfg.successors.get(block, set())
+                for successor in successors:
+                    self.cfg.predecessors[successor].discard(block)
+                self.cfg.successors.pop(block, None)
+                self.cfg.predecessors.pop(block, None)
+                block.delete()
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)
+
+
+@dataclass
+class CFGEdge(RewriteRule):
+    """Merge non-branching blocks on the edge of the CFG.
+
+    Example:
+
+        /---> [B] --> [D] --> [E]
+    [A]-----> [C] -------------^
+
+    [B] and [D] are non-branching blocks on the same edge. They can be merged into one block.
+
+        /---> [B,D] --> [E]
+    [A]-----> [C] -------^
+    """
+
+    cfg: CFG
+
+    def rewrite_Region(self, node: ir.Region) -> RewriteResult:
+        result = RewriteResult()
+        for block in node.blocks:
+            result = self.rewrite_Block(block).join(result)
+        return result
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        successors = self.cfg.successors.get(node, None)
+        if (
+            successors is None or len(successors) > 1 or len(successors) == 0
+        ):  # multiple outgoing edges
+            return RewriteResult()
+
+        successor = next(iter(successors))
+        if len(self.cfg.predecessors[successor]) > 1:  # multiple incoming edges
+            return RewriteResult()
+
+        if not ((last_stmt := node.last_stmt) and isinstance(last_stmt, cf.Branch)):
+            return RewriteResult()
+
+        # merge the two blocks
+        for arg, input in zip(successor.args, last_stmt.arguments):
+            arg.replace_by(input)
+        last_stmt.delete()
+        for stmt in successor.stmts:
+            stmt.detach()
+            node.stmts.append(stmt)
+        successor.delete()
+
+        # update the CFG
+        new_successors = self.cfg.successors[successor]
+        self.cfg.successors[node] = new_successors
+        for new_successor in new_successors:
+            self.cfg.predecessors[new_successor].discard(successor)
+            self.cfg.predecessors[new_successor].add(node)
+        del self.cfg.successors[successor]
+        del self.cfg.predecessors[successor]  # this is just [node]
+        return RewriteResult(has_done_something=True)
+
+
+class DuplicatedBranch(RewriteRule):
+    """Merge duplicated branches into a single branch.
+
+    Example:
+
+        [A]-->[B]
+          -----^
+
+    Merge the two branches into one without changing the CFG:
+
+        [A]-->[B]
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if (
+            not isinstance(node, cf.ConditionalBranch)
+            or node.then_successor is not node.else_successor
+        ):
+            return RewriteResult()
+
+        for then_x, else_x in zip(node.then_arguments, node.else_arguments):
+            if then_x is not else_x:
+                return RewriteResult()
+
+        node.replace_by(
+            cf.Branch(arguments=node.then_arguments, successor=node.then_successor)
+        )
+        return RewriteResult(has_done_something=True)
+
+
+@dataclass
+class SkipBlock(RewriteRule):
+    """Simplify a block that only contains a branch statement."""
+
+    cfg: CFG
+
+    def rewrite_Region(self, node: ir.Region) -> RewriteResult:
+        result = RewriteResult()
+        for block in node.blocks:
+            result = self.rewrite_Block(block).join(result)
+        return result
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        if len(node.stmts) != 1:
+            return RewriteResult()
+
+        stmt = node.last_stmt
+        if not isinstance(stmt, cf.Branch):
+            return RewriteResult()
+
+        has_done_something = False
+        predecessors = self.cfg.predecessors.get(node, set())
+        # only if there is one predecessor and no uses of the arguments
+        if len(predecessors) == 1 and all(
+            self.can_skip(stmt, each) for each in node.args
+        ):
+            has_done_something = self.rewrite_pred(node, stmt, next(iter(predecessors)))
+        return RewriteResult(has_done_something=has_done_something)
+
+    def can_skip(self, terminator: cf.Branch, value: ir.SSAValue) -> bool:
+        for use in value.uses:
+            if use.stmt is terminator:
+                continue
+            return False
+        return True
+
+    def rewrite_pred(
+        self, node: ir.Block, node_terminator: cf.Branch, predecessor: ir.Block
+    ) -> bool:
+        terminator = predecessor.last_stmt
+        if isinstance(terminator, cf.Branch):
+            return self.rewrite_pred_Branch(
+                node, node_terminator, predecessor, terminator
+            )
+        elif isinstance(terminator, cf.ConditionalBranch):
+            return self.rewrite_pred_ConditionalBranch(
+                node, node_terminator, predecessor, terminator
+            )
+        return False
+
+    def rewrite_pred_Branch(
+        self,
+        node: ir.Block,
+        node_terminator: cf.Branch,
+        predecessor: ir.Block,
+        pred_terminator: cf.Branch,
+    ) -> bool:
+        ssamap = self._block_inputs(node, pred_terminator.arguments)
+        pred_terminator.replace_by(
+            cf.Branch(
+                # NOTE: the argument can also be SSAs from previous blocks (non-phi)
+                arguments=tuple(
+                    ssamap.get(arg, arg) for arg in node_terminator.arguments
+                ),
+                successor=node_terminator.successor,
+            )
+        )
+
+        self.fix_cfg(predecessor, node, node_terminator.successor)
+        return True
+
+    def rewrite_pred_ConditionalBranch(
+        self,
+        node: ir.Block,
+        node_terminator: cf.Branch,
+        predecessor: ir.Block,
+        pred_terminator: cf.ConditionalBranch,
+    ) -> bool:
+        then_arguments = pred_terminator.then_arguments
+        else_arguments = pred_terminator.else_arguments
+        then_successor = pred_terminator.then_successor
+        else_successor = pred_terminator.else_successor
+
+        has_done_something = False
+        if pred_terminator.then_successor is node:
+            ssamap = self._block_inputs(node, pred_terminator.then_arguments)
+            then_arguments = tuple(
+                ssamap.get(arg, arg) for arg in node_terminator.arguments
+            )
+            then_successor = node_terminator.successor
+            has_done_something = True
+            self.fix_cfg(predecessor, node, then_successor)
+
+        if pred_terminator.else_successor is node:
+            ssamap = self._block_inputs(node, pred_terminator.else_arguments)
+            else_arguments = tuple(
+                ssamap.get(arg, arg) for arg in node_terminator.arguments
+            )
+            else_successor = node_terminator.successor
+            has_done_something = True
+            self.fix_cfg(predecessor, node, else_successor)
+
+        pred_terminator.replace_by(
+            cf.ConditionalBranch(
+                cond=pred_terminator.cond,
+                then_arguments=then_arguments,
+                then_successor=then_successor,
+                else_arguments=else_arguments,
+                else_successor=else_successor,
+            )
+        )
+        return has_done_something
+
+    def fix_cfg(self, predecessor: ir.Block, node: ir.Block, successor: ir.Block):
+        node_pred_succ = self.cfg.successors.setdefault(predecessor, set())
+        node_pred_succ.discard(node)
+        node_pred_succ.add(successor)
+
+        node_succ_pred = self.cfg.predecessors.setdefault(successor, set())
+        node_succ_pred.add(predecessor)
+
+        node_pred = self.cfg.predecessors.setdefault(node, set())
+        node_pred.discard(predecessor)
+
+    def _block_inputs(
+        self, block: ir.Block, arguments: tuple[ir.SSAValue, ...]
+    ) -> dict[ir.SSAValue, ir.SSAValue]:
+        return dict(zip(block.args, arguments))
+
+
+@dataclass
+class CompactifyRegion(RewriteRule):
+    """Wrapper to share the CFG object with same CFG region."""
+
+    cfg: CFG
+
+    def __init__(self, cfg: CFG):
+        self.cfg = cfg
+        self.rule = Fixpoint(
+            Chain(
+                DeadBlock(cfg), Walk(DuplicatedBranch()), SkipBlock(cfg), CFGEdge(cfg)
+            )
+        )
+
+    def rewrite(self, node: ir.IRNode) -> RewriteResult:
+        return self.rule.rewrite(node)
+
+
+@dataclass
+class CFGCompactify(RewriteRule):
+    """Compactify the CFG by removing dead blocks and merging blocks
+    if the statement uses the SSACFG convention. Do nothing if given
+    `ir.Region` or `ir.Block` due to no context of the region.
+
+    To compactify hierarchical CFG, combine this rule with `kirin.rewrite.Walk`
+    to recursively apply this rule to all statements.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        result = RewriteResult()
+        if not (trait := node.get_trait(ir.SSACFGRegion)):
+            return result
+
+        for region in node.regions:
+            cfg = trait.get_graph(region)
+            result = CompactifyRegion(cfg).rewrite(region).join(result)
+        return result

kirin/rewrite/cse.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+
+from kirin.ir import Pure, Block, Statement
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class CommonSubexpressionElimination(RewriteRule):
+
+    def rewrite_Block(self, node: Block) -> RewriteResult:
+        seen: dict[int, Statement] = {}
+
+        for stmt in node.stmts:
+            if not stmt.has_trait(Pure):
+                continue
+
+            if stmt.regions:
+                continue
+
+            hash_value = hash(
+                (type(stmt),)
+                + tuple(stmt.args)
+                + tuple(stmt.attributes.values())
+                + tuple(stmt.successors)
+                + tuple(stmt.regions)
+            )
+            if hash_value in seen:
+                old_stmt = seen[hash_value]
+                for result in stmt._results:
+                    result.replace_by(old_stmt._results[0])
+                stmt.delete()
+                return RewriteResult(has_done_something=True)
+            else:
+                seen[hash_value] = stmt
+        return RewriteResult()
+
+    def rewrite_Statement(self, node: Statement) -> RewriteResult:
+        if not node.regions:
+            return RewriteResult()
+
+        has_done_something = False
+        for region in node.regions:
+            for block in region.blocks:
+                result = self.rewrite_Block(block)
+                if result.has_done_something:
+                    has_done_something = True
+
+        return RewriteResult(has_done_something=has_done_something)

kirin/rewrite/dce.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class DeadCodeElimination(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if self.is_pure(node):
+            for result in node._results:
+                if result.uses:
+                    return RewriteResult()
+
+            node.delete()
+            return RewriteResult(has_done_something=True)
+
+        return RewriteResult()

kirin/rewrite/fixpoint.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+
+from kirin.ir import IRNode
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+
+
+@dataclass
+class Fixpoint(RewriteRule):
+    """Apply a rewrite rule until a fixpoint is reached.
+
+    The rewrite rule is applied to the node until the rewrite rule does not do anything.
+
+    ### Parameters
+    - `map`: The rewrite rule to apply.
+    - `max_iter`: The maximum number of iterations to apply the rewrite rule. Default is 32.
+    """
+
+    rule: RewriteRule
+    max_iter: int = 32
+
+    def rewrite(self, node: IRNode) -> RewriteResult:
+        has_done_something = False
+        for _ in range(self.max_iter):
+            result = self.rule.rewrite(node)
+            if result.terminated:
+                return result
+
+            if result.has_done_something:
+                has_done_something = True
+            else:
+                return RewriteResult(has_done_something=has_done_something)
+
+        return RewriteResult(exceeded_max_iter=True)

kirin/rewrite/fold.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.dialects import cf
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.py.constant import Constant
+
+
+@dataclass
+class ConstantFold(RewriteRule):
+
+    def get_const(self, value: ir.SSAValue):
+        ret = value.hints.get("const")
+
+        if ret is not None and isinstance(ret, const.Value):
+            return ret
+        return None
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if node.has_trait(ir.ConstantLike):
+            return RewriteResult()
+        elif isinstance(node, cf.ConditionalBranch):
+            return self.rewrite_cf_ConditionalBranch(node)
+
+        if not self.is_pure(node):
+            return RewriteResult()
+
+        has_done_something = False
+        for old_result in node.results:
+            if (value := self.get_const(old_result)) is not None:
+                stmt = Constant(value.data)
+                stmt.insert_before(node)
+                old_result.replace_by(stmt.result)
+                stmt.result.hints["const"] = value
+                if old_result.name:
+                    stmt.result.name = old_result.name
+                has_done_something = True
+        return RewriteResult(has_done_something=has_done_something)
+
+    def rewrite_cf_ConditionalBranch(self, node: cf.ConditionalBranch):
+        if (value := self.get_const(node.cond)) is not None:
+            if value.data is True:
+                cf.Branch(
+                    arguments=node.then_arguments,
+                    successor=node.then_successor,
+                ).insert_before(node)
+            elif value.data is False:
+                cf.Branch(
+                    arguments=node.else_arguments,
+                    successor=node.else_successor,
+                ).insert_before(node)
+            else:
+                raise ValueError(f"Invalid constant value for branch: {value.data}")
+            node.delete()
+            return RewriteResult(has_done_something=True)
+        return RewriteResult()

kirin/rewrite/getfield.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class InlineGetField(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, func.GetField):
+            return RewriteResult()
+
+        if not isinstance(node.obj.owner, func.Lambda):
+            return RewriteResult()
+
+        original = node.obj.owner.captured[node.field]
+        node.result.replace_by(original)
+        node.delete()
+        return RewriteResult(has_done_something=True)

kirin/rewrite/getitem.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class InlineGetItem(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.indexing.GetItem):
+            return RewriteResult()
+
+        if not isinstance(node.obj.owner, py.tuple.New):
+            return RewriteResult()
+
+        if not isinstance(index_value := node.index.hints.get("const"), const.Value):
+            return RewriteResult()
+
+        stmt = node.obj.owner
+        index = index_value.data
+        if isinstance(index, int) and (
+            0 <= index < len(stmt.args) or -len(stmt.args) <= index < 0
+        ):
+            node.result.replace_by(stmt.args[index])
+            return RewriteResult(has_done_something=True)
+        elif isinstance(index, slice):
+            start, stop, step = index.indices(len(stmt.args))
+            new_tuple = py.tuple.New(
+                tuple(stmt.args[start:stop:step]),
+            )
+            node.replace_by(new_tuple)
+            return RewriteResult(has_done_something=True)
+        else:
+            return RewriteResult()