kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/scf/stmts.py

@@ -0,0 +1,250 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+from kirin.exceptions import VerificationError, DialectLoweringError
+from kirin.print.printer import Printer
+
+from ._dialect import dialect
+
+
+@statement(dialect=dialect, init=False)
+class IfElse(ir.Statement):
+    """Python-like if-else statement.
+
+    This statement has a condition, then body, and else body.
+
+    Then body either terminates with a yield statement or `scf.return`.
+    """
+
+    name = "if"
+    traits = frozenset({ir.MaybePure()})
+    purity: bool = info.attribute(default=False)
+    cond: ir.SSAValue = info.argument(types.Any)
+    # NOTE: we don't enforce the type here
+    # because anything implements __bool__ in Python
+    # can be used as a condition
+    then_body: ir.Region = info.region(multi=False)
+    else_body: ir.Region = info.region(multi=False, default_factory=ir.Region)
+
+    def __init__(
+        self,
+        cond: ir.SSAValue,
+        then_body: ir.Region | ir.Block,
+        else_body: ir.Region | ir.Block | None = None,
+    ):
+        if isinstance(then_body, ir.Region):
+            if len(then_body.blocks) != 1:
+                raise DialectLoweringError(
+                    "if-else statement must have a single block in the then region"
+                )
+            then_body_region = then_body
+            then_body = then_body_region.blocks[0]
+        elif isinstance(then_body, ir.Block):
+            then_body_region = ir.Region(then_body)
+
+        if isinstance(else_body, ir.Region):
+            if not else_body.blocks:  # empty region
+                else_body_region = else_body
+                else_body = None
+            elif len(else_body.blocks) != 1:
+                raise DialectLoweringError(
+                    "if-else statement must have a single block in the else region"
+                )
+            else:
+                else_body_region = else_body
+                else_body = else_body_region.blocks[0]
+        elif isinstance(else_body, ir.Block):
+            else_body_region = ir.Region(else_body)
+        else:
+            else_body_region = ir.Region()
+
+        # if either then or else body has yield, we generate results
+        # we assume if both have yields, they have the same number of results
+        then_yield = then_body.last_stmt
+        else_yield = else_body.last_stmt if else_body is not None else None
+        if then_yield is not None and isinstance(then_yield, Yield):
+            results = then_yield.values
+        elif else_yield is not None and isinstance(else_yield, Yield):
+            results = else_yield.values
+        else:
+            results = ()
+
+        result_types = tuple(value.type for value in results)
+        super().__init__(
+            args=(cond,),
+            regions=(then_body_region, else_body_region),
+            result_types=result_types,
+            args_slice={"cond": 0},
+            attributes={"purity": ir.PyAttr(False)},
+        )
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.print(self.cond)
+        printer.plain_print(" ")
+        printer.print(self.then_body)
+        if self.else_body.blocks and not (
+            len(self.else_body.blocks[0].stmts) == 1
+            and isinstance(else_term := self.else_body.blocks[0].last_stmt, Yield)
+            and not else_term.values  # empty yield
+        ):
+            printer.plain_print(" else ", style="keyword")
+            printer.print(self.else_body)
+
+        with printer.rich(style="comment"):
+            printer.plain_print(f" -> purity={self.purity}")
+
+    def verify(self) -> None:
+        from kirin.dialects.func import Return
+
+        if len(self.then_body.blocks) != 1:
+            raise VerificationError(self, "then region must have a single block")
+
+        if len(self.else_body.blocks) != 1:
+            raise VerificationError(self, "else region must have a single block")
+
+        then_block = self.then_body.blocks[0]
+        else_block = self.else_body.blocks[0]
+        if len(then_block.args) != 1:
+            raise VerificationError(
+                self, "then block must have a single argument for condition"
+            )
+
+        if len(else_block.args) != 1:
+            raise VerificationError(
+                self, "else block must have a single argument for condition"
+            )
+
+        then_stmt = then_block.last_stmt
+        else_stmt = else_block.last_stmt
+        if then_stmt is None or not isinstance(then_stmt, (Yield, Return)):
+            raise VerificationError(
+                self, "then block must terminate with a yield or return"
+            )
+
+        if else_stmt is None or not isinstance(else_stmt, (Yield, Return)):
+            raise VerificationError(
+                self, "else block must terminate with a yield or return"
+            )
+
+
+@statement(dialect=dialect, init=False)
+class For(ir.Statement):
+    name = "for"
+    traits = frozenset({ir.MaybePure()})
+    purity: bool = info.attribute(default=False)
+    iterable: ir.SSAValue = info.argument(types.Any)
+    body: ir.Region = info.region(multi=False)
+    initializers: tuple[ir.SSAValue, ...] = info.argument(types.Any)
+
+    def __init__(
+        self,
+        iterable: ir.SSAValue,
+        body: ir.Region,
+        *initializers: ir.SSAValue,
+    ):
+        stmt = body.blocks[0].last_stmt
+        if isinstance(stmt, Yield):
+            result_types = tuple(value.type for value in stmt.values)
+        else:
+            result_types = ()
+        super().__init__(
+            args=(iterable, *initializers),
+            regions=(body,),
+            result_types=result_types,
+            args_slice={"iterable": 0, "initializers": slice(1, None)},
+            attributes={"purity": ir.PyAttr(False)},
+        )
+
+    def verify(self) -> None:
+        from kirin.dialects.func import Return
+
+        if len(self.body.blocks) != 1:
+            raise VerificationError(self, "for loop body must have a single block")
+
+        if len(self.body.blocks[0].args) != len(self.initializers) + 1:
+            raise VerificationError(
+                self,
+                "for loop body must have arguments for all initializers and the loop variable",
+            )
+
+        stmt = self.body.blocks[0].last_stmt
+        if stmt is None or not isinstance(stmt, (Yield, Return)):
+            raise VerificationError(
+                self, "for loop body must terminate with a yield or return"
+            )
+
+        if isinstance(stmt, Return):
+            return
+
+        if len(stmt.values) != len(self.initializers):
+            raise VerificationError(
+                self,
+                "for loop body must have the same number of results as initializers",
+            )
+        if len(self.results) != len(stmt.values):
+            raise VerificationError(
+                self,
+                "for loop must have the same number of results as the yield in the body",
+            )
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        block = self.body.blocks[0]
+        printer.print(block.args[0])
+        printer.plain_print(" in ", style="keyword")
+        printer.print(self.iterable)
+        if self.results:
+            with printer.rich(style="comment"):
+                printer.plain_print(" -> ")
+                printer.print_seq(
+                    tuple(result.type for result in self.results),
+                    delim=", ",
+                    style="comment",
+                )
+
+        with printer.indent():
+            if self.initializers:
+                printer.print_newline()
+                printer.plain_print("iter_args(")
+                for idx, (arg, val) in enumerate(
+                    zip(block.args[1:], self.initializers)
+                ):
+                    printer.print(arg)
+                    printer.plain_print(" = ")
+                    printer.print(val)
+                    if idx < len(self.initializers) - 1:
+                        printer.plain_print(", ")
+                printer.plain_print(")")
+
+            printer.plain_print(" {")
+            if printer.analysis is not None:
+                with printer.rich(style="warning"):
+                    for arg in block.args:
+                        printer.print_newline()
+                        printer.print_analysis(
+                            arg, prefix=f"{printer.state.ssa_id[arg]} --> "
+                        )
+            with printer.align(printer.result_width(block.stmts)):
+                for stmt in block.stmts:
+                    printer.print_newline()
+                    printer.print_stmt(stmt)
+        printer.print_newline()
+        printer.plain_print("}")
+        with printer.rich(style="comment"):
+            printer.plain_print(f" -> purity={self.purity}")
+
+
+@statement(dialect=dialect)
+class Yield(ir.Statement):
+    name = "yield"
+    traits = frozenset({ir.IsTerminator()})
+    values: tuple[ir.SSAValue, ...] = info.argument(types.Any)
+
+    def __init__(self, *values: ir.SSAValue):
+        super().__init__(args=values, args_slice={"values": slice(None)})
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.print_seq(self.values, prefix=" ", delim=", ")

kirin/dialects/scf/trim.py

@@ -0,0 +1,36 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+
+from .stmts import For, Yield, IfElse
+
+
+class UnusedYield(RewriteRule):
+    """Trim unused results from `For` and `IfElse` statements."""
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, (For, IfElse)):
+            return RewriteResult()
+
+        any_unused = False
+        uses: list[int] = []
+        results: list[ir.ResultValue] = []
+        for idx, result in enumerate(node.results):
+            if result.uses:
+                uses.append(idx)
+                results.append(result)
+            else:
+                any_unused = True
+
+        if not any_unused:
+            return RewriteResult()
+
+        node._results = results
+        for region in node.regions:
+            for block in region.blocks:
+                if not isinstance(block.last_stmt, Yield):
+                    continue
+
+                block.last_stmt.args = [block.last_stmt.args[idx] for idx in uses]
+
+        return RewriteResult(has_done_something=True)

kirin/dialects/scf/typeinfer.py

@@ -0,0 +1,58 @@
+from kirin import ir, types, interp
+from kirin.analysis import ForwardFrame, TypeInference
+from kirin.dialects import func
+from kirin.dialects.eltype import ElType
+
+from . import absint
+from .stmts import For, IfElse
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(absint.Methods):
+
+    @interp.impl(IfElse)
+    def if_else_(
+        self,
+        interp_: TypeInference,
+        frame: ForwardFrame[types.TypeAttribute],
+        stmt: IfElse,
+    ):
+        frame.set(
+            stmt.cond, frame.get(stmt.cond).meet(types.Bool)
+        )  # set cond backwards
+        return super().if_else(self, interp_, frame, stmt)
+
+    @interp.impl(For)
+    def for_loop(
+        self,
+        interp_: TypeInference,
+        frame: ForwardFrame[types.TypeAttribute],
+        stmt: For,
+    ):
+        iterable = frame.get(stmt.iterable)
+        loop_vars = frame.get_values(stmt.initializers)
+        body_block = stmt.body.blocks[0]
+        block_args = body_block.args
+
+        eltype = interp_.run_stmt(ElType(ir.TestValue()), (iterable,))
+        if not isinstance(eltype, tuple):  # error
+            return
+        item = eltype[0]
+        frame.set_values(block_args, (item,) + loop_vars)
+
+        if isinstance(body_block.last_stmt, func.Return):
+            frame.worklist.append(interp.Successor(body_block, item, *loop_vars))
+            return  # if terminate is Return, there is no result
+
+        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+            body_frame.entries.update(frame.entries)
+            loop_vars_ = interp_.run_ssacfg_region(body_frame, stmt.body)
+
+        frame.entries.update(body_frame.entries)
+        if isinstance(loop_vars_, interp.ReturnValue):
+            return loop_vars_
+        elif isinstance(loop_vars_, tuple):
+            return interp_.join_results(loop_vars, loop_vars_)
+        else:  # None, loop has no result
+            return

kirin/dialects/scf/unroll.py

@@ -0,0 +1,92 @@
+from kirin import ir
+from kirin.analysis import const
+from kirin.dialects import func
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+from kirin.dialects.py.constant import Constant
+
+from .stmts import For, Yield, IfElse
+
+
+class PickIfElse(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, IfElse):
+            return RewriteResult()
+
+        if not isinstance(hint := node.cond.hints.get("const"), const.Value):
+            return RewriteResult()
+
+        if hint.data:
+            return self.insert_body(node, node.then_body)
+        else:
+            return self.insert_body(node, node.else_body)
+
+    def insert_body(self, node: IfElse, body: ir.Region):
+        body_block = body.blocks[0]
+        body_block.args[0].replace_by(node.cond)
+        block_stmt = body_block.first_stmt
+        while block_stmt and not block_stmt.has_trait(ir.IsTerminator):
+            block_stmt.detach()
+            block_stmt.insert_before(node)
+            block_stmt = body_block.first_stmt
+
+        terminator = body_block.last_stmt
+        if isinstance(terminator, Yield):
+            for result, output in zip(node.results, terminator.values):
+                result.replace_by(output)
+            node.delete()
+            return RewriteResult(has_done_something=True)
+        elif isinstance(terminator, func.Return):
+            block = node.parent
+            assert block is not None
+            stmt = block.last_stmt
+            while stmt is not None and stmt is not node:  # remove the rest of the block
+                delete_stmt = stmt
+                stmt = stmt.prev_stmt
+                delete_stmt.delete()
+
+            terminator.detach()
+            node.replace_by(terminator)
+            return RewriteResult(has_done_something=True)
+        else:
+            return RewriteResult()
+
+
+class ForLoop(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, For):
+            return RewriteResult()
+
+        # TODO: support for PartialTuple and IList with known length
+        if not isinstance(hint := node.iterable.hints.get("const"), const.Value):
+            return RewriteResult()
+
+        loop_vars = node.initializers
+        for item in hint.data:
+            body = node.body.clone()
+            block = body.blocks[0]
+            item_stmt = Constant(item)
+            item_stmt.insert_before(node)
+            block.args[0].replace_by(item_stmt.result)
+            for var, input in zip(block.args[1:], loop_vars):
+                var.replace_by(input)
+
+            block_stmt = block.first_stmt
+            while block_stmt and not block_stmt.has_trait(ir.IsTerminator):
+                block_stmt.detach()
+                block_stmt.insert_before(node)
+                block_stmt = block.first_stmt
+
+            terminator = block.last_stmt
+            # we assume Yield has the same # of values as initializers
+            # TODO: check this in validation
+            if isinstance(terminator, Yield):
+                loop_vars = terminator.values
+                terminator.delete()
+
+        for result, output in zip(node.results, loop_vars):
+            result.replace_by(output)
+        node.delete()
+        return RewriteResult(has_done_something=True)

kirin/emit/__init__.py

@@ -0,0 +1,3 @@
+from .abc import EmitABC as EmitABC, EmitFrame as EmitFrame
+from .str import EmitStr as EmitStr, EmitStrFrame as EmitStrFrame
+from .exceptions import EmitError as EmitError

kirin/emit/abc.py

@@ -0,0 +1,89 @@
+from abc import ABC
+from typing import TypeVar
+from dataclasses import field, dataclass
+
+from kirin import ir, interp
+from kirin.worklist import WorkList
+
+ValueType = TypeVar("ValueType")
+
+
+@dataclass
+class EmitFrame(interp.Frame[ValueType]):
+    worklist: WorkList[interp.Successor] = field(default_factory=WorkList)
+    block_ref: dict[ir.Block, ValueType] = field(default_factory=dict)
+
+
+FrameType = TypeVar("FrameType", bound=EmitFrame)
+
+
+@dataclass
+class EmitABC(interp.BaseInterpreter[FrameType, ValueType], ABC):
+
+    def run_callable_region(
+        self, frame: FrameType, code: ir.Statement, region: ir.Region
+    ) -> ValueType:
+        results = self.eval_stmt(frame, code)
+        if isinstance(results, tuple):
+            if len(results) == 0:
+                return self.void
+            elif len(results) == 1:
+                return results[0]
+        raise interp.InterpreterError(f"Unexpected results {results}")
+
+    def run_ssacfg_region(
+        self, frame: FrameType, region: ir.Region
+    ) -> tuple[ValueType, ...]:
+        frame.worklist.append(
+            interp.Successor(region.blocks[0], frame.get_values(region.blocks[0].args))
+        )
+        while (succ := frame.worklist.pop()) is not None:
+            block_header = self.emit_block(frame, succ.block)
+            frame.block_ref[succ.block] = block_header
+        return ()
+
+    def emit_attribute(self, attr: ir.Attribute) -> ValueType:
+        return getattr(
+            self, f"emit_type_{type(attr).__name__}", self.emit_attribute_fallback
+        )(attr)
+
+    def emit_attribute_fallback(self, attr: ir.Attribute) -> ValueType:
+        if (method := self.registry.attributes.get(type(attr))) is not None:
+            return method(self, attr)
+        raise NotImplementedError(f"Attribute {type(attr)} not implemented")
+
+    def emit_stmt_begin(self, frame: FrameType, stmt: ir.Statement) -> None:
+        return
+
+    def emit_stmt_end(self, frame: FrameType, stmt: ir.Statement) -> None:
+        return
+
+    def emit_block_begin(self, frame: FrameType, block: ir.Block) -> None:
+        return
+
+    def emit_block_end(self, frame: FrameType, block: ir.Block) -> None:
+        return
+
+    def emit_block(self, frame: FrameType, block: ir.Block) -> ValueType:
+        self.emit_block_begin(frame, block)
+        stmt = block.first_stmt
+        while stmt is not None:
+            if self.consume_fuel() == self.FuelResult.Stop:
+                raise interp.FuelExhaustedError("fuel exhausted")
+
+            self.emit_stmt_begin(frame, stmt)
+            stmt_results = self.eval_stmt(frame, stmt)
+            self.emit_stmt_end(frame, stmt)
+
+            match stmt_results:
+                case tuple(values):
+                    frame.set_values(stmt._results, values)
+                case interp.ReturnValue(_) | interp.YieldValue(_):
+                    pass
+                case _:
+                    raise ValueError(f"Unexpected result {stmt_results}")
+
+            stmt = stmt.next_stmt
+
+        self.emit_block_end(frame, block)
+        return frame.block_ref[block]

kirin/emit/abc.pyi

@@ -0,0 +1,38 @@
+from typing import TypeVar
+from dataclasses import field, dataclass
+
+from kirin import ir, types, interp
+from kirin.worklist import WorkList
+
+ValueType = TypeVar("ValueType")
+
+@dataclass
+class EmitFrame(interp.Frame[ValueType]):
+    worklist: WorkList[interp.Successor] = field(default_factory=WorkList)
+    block_ref: dict[ir.Block, ValueType] = field(default_factory=dict)
+
+FrameType = TypeVar("FrameType", bound=EmitFrame)
+
+class EmitABC(interp.BaseInterpreter[FrameType, ValueType]):
+    def run_callable_region(
+        self, frame: FrameType, code: ir.Statement, region: ir.Region
+    ) -> ValueType: ...
+    def run_ssacfg_region(
+        self, frame: FrameType, region: ir.Region
+    ) -> tuple[ValueType, ...]: ...
+    def emit_attribute(self, attr: ir.Attribute) -> ValueType: ...
+    def emit_type_Any(self, attr: types.AnyType) -> ValueType: ...
+    def emit_type_Bottom(self, attr: types.BottomType) -> ValueType: ...
+    def emit_type_Literal(self, attr: types.Literal) -> ValueType: ...
+    def emit_type_Union(self, attr: types.Union) -> ValueType: ...
+    def emit_type_TypeVar(self, attr: types.TypeVar) -> ValueType: ...
+    def emit_type_Vararg(self, attr: types.Vararg) -> ValueType: ...
+    def emit_type_Generic(self, attr: types.Generic) -> ValueType: ...
+    def emit_type_PyClass(self, attr: types.PyClass) -> ValueType: ...
+    def emit_type_PyAttr(self, attr: ir.PyAttr) -> ValueType: ...
+    def emit_attribute_fallback(self, attr: ir.Attribute) -> ValueType: ...
+    def emit_stmt_begin(self, frame: FrameType, stmt: ir.Statement) -> None: ...
+    def emit_stmt_end(self, frame: FrameType, stmt: ir.Statement) -> None: ...
+    def emit_block_begin(self, frame: FrameType, block: ir.Block) -> None: ...
+    def emit_block_end(self, frame: FrameType, block: ir.Block) -> None: ...
+    def emit_block(self, frame: FrameType, block: ir.Block) -> ValueType: ...

kirin/emit/exceptions.py

@@ -0,0 +1,5 @@
+from kirin.interp import InterpreterError
+
+
+class EmitError(InterpreterError):
+    pass

kirin/emit/julia.py

@@ -0,0 +1,63 @@
+from typing import IO, TypeVar
+
+from kirin import ir
+from kirin.ir.attrs.types import PyClass
+from kirin.ir.nodes.block import Block
+
+from .str import EmitStr, EmitStrFrame
+
+IO_t = TypeVar("IO_t", bound=IO)
+
+
+class EmitJulia(EmitStr[IO_t]):
+    keys = ["emit.julia"]
+
+    PYTYPE_MAP = {
+        int: "Int",
+        float: "Real",
+        str: "String",
+        bool: "Bool",
+        type(None): "Nothing",
+        dict: "Dict",
+        list: "Vector",
+        tuple: "Tuple",
+    }
+
+    def emit_block_begin(self, frame: EmitStrFrame, block: Block) -> None:
+        block_id = self.block_id[block]
+        frame.block_ref[block] = block_id
+        self.newline(frame)
+        self.write(f"@label {block_id};")
+
+    def emit_type_PyClass(self, attr: PyClass) -> str:
+        return self.PYTYPE_MAP.get(attr.typ, "Any")
+
+    def write_assign(self, frame: EmitStrFrame, result: ir.SSAValue, *args):
+        result_sym = self.ssa_id[result]
+        frame.set(result, result_sym)
+        self.writeln(frame, result_sym, " = ", *args)
+        return result_sym
+
+    def emit_binaryop(
+        self,
+        frame: EmitStrFrame,
+        sym: str,
+        lhs: ir.SSAValue,
+        rhs: ir.SSAValue,
+        result: ir.ResultValue,
+    ):
+        return (
+            self.write_assign(
+                frame,
+                result,
+                f"{frame.get(lhs)} {sym} {frame.get(rhs)}",
+            ),
+        )
+
+    def emit_type_PyAttr(self, attr: ir.PyAttr) -> str:
+        if isinstance(attr.data, (int, float)):
+            return repr(attr.data)
+        elif isinstance(attr.data, str):
+            return f'"{attr.data}"'
+        else:
+            raise ValueError(f"unsupported type {type(attr.data)}")

kirin/emit/str.py

@@ -0,0 +1,51 @@
+from abc import ABC
+from typing import IO, Generic, TypeVar
+from dataclasses import field, dataclass
+
+from kirin import ir, interp, idtable
+from kirin.emit.abc import EmitABC, EmitFrame
+
+IO_t = TypeVar("IO_t", bound=IO)
+
+
+@dataclass
+class EmitStrFrame(EmitFrame[str]):
+    indent: int = 0
+    captured: dict[ir.SSAValue, tuple[str, ...]] = field(default_factory=dict)
+
+
+@dataclass
+class EmitStr(EmitABC[EmitStrFrame, str], ABC, Generic[IO_t]):
+    void = ""
+    file: IO_t
+    prefix: str = field(default="", kw_only=True)
+    prefix_if_none: str = field(default="var_", kw_only=True)
+
+    def initialize(self):
+        super().initialize()
+        self.ssa_id = idtable.IdTable[ir.SSAValue](
+            prefix=self.prefix, prefix_if_none=self.prefix_if_none
+        )
+        self.block_id = idtable.IdTable[ir.Block](prefix=self.prefix + "block_")
+        return self
+
+    def new_frame(self, code: ir.Statement) -> EmitStrFrame:
+        return EmitStrFrame.from_func_like(code)
+
+    def run_method(
+        self, method: ir.Method, args: tuple[str, ...]
+    ) -> tuple[EmitStrFrame, str]:
+        if len(self.state.frames) >= self.max_depth:
+            raise interp.InterpreterError("maximum recursion depth exceeded")
+        return self.run_callable(method.code, (method.sym_name,) + args)
+
+    def write(self, *args):
+        for arg in args:
+            self.file.write(arg)
+
+    def newline(self, frame: EmitStrFrame):
+        self.file.write("\n" + "  " * frame.indent)
+
+    def writeln(self, frame: EmitStrFrame, *args):
+        self.newline(frame)
+        self.write(*args)