kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/ilist/rewrite/unroll.py

@@ -0,0 +1,131 @@
+from kirin import ir, types
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+from kirin.dialects.py.tuple import New as TupleNew
+from kirin.dialects.func.stmts import Call
+from kirin.dialects.ilist.stmts import Map, New, Scan, Foldl, Foldr, ForEach, IListType
+from kirin.dialects.py.constant import Constant
+from kirin.dialects.py.indexing import GetItem
+
+
+class Unroll(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        return getattr(
+            self, f"rewrite_{node.__class__.__name__}", self.rewrite_fallback
+        )(node)
+
+    def rewrite_fallback(self, node: ir.Statement) -> RewriteResult:
+        return RewriteResult()
+
+    def _get_collection_len(self, collection: ir.SSAValue):
+        coll_type = collection.type
+        if not isinstance(coll_type, types.Generic):
+            return None
+
+        if not coll_type.is_subseteq(IListType):
+            return None
+
+        if not (
+            isinstance(coll_type.vars[1], types.Literal)
+            and isinstance(coll_type.vars[1].data, int)
+        ):
+            return None
+
+        return coll_type.vars[1].data
+
+    def rewrite_Map(self, node: Map) -> RewriteResult:
+        # NOTE: if node.collection is a constant, we can
+        # just leave it because Map is pure, and this will
+        # be folded.
+        if (coll_len := self._get_collection_len(node.collection)) is None:
+            return RewriteResult()
+
+        new_elems: list[ir.SSAValue] = []
+        for elt_idx in range(coll_len):
+            index = Constant(elt_idx)
+            index.insert_before(node)
+            elt = GetItem(node.collection, index.result)
+            elt.insert_before(node)
+            fn_call = Call(node.fn, (elt.result,))
+            fn_call.insert_before(node)
+            new_elems.append(fn_call.result)
+
+        node.replace_by(New(values=tuple(new_elems)))
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_Scan(self, node: Scan) -> RewriteResult:
+        if (coll_len := self._get_collection_len(node.collection)) is None:
+            return RewriteResult()
+
+        index_0 = Constant(0)
+        index_1 = Constant(1)
+        # index_0.result.name = "idx0"
+        # index_1.result.name = "idx1"
+        index_0.insert_before(node)
+        index_1.insert_before(node)
+        carry = node.init
+        ys: list[ir.SSAValue] = []
+        for elem_idx in range(coll_len):
+            index = Constant(elem_idx)
+            # index.result.name = f"idx_{elem_idx}"
+            elt = GetItem(node.collection, index.result)
+            fn_call = Call(node.fn, (carry, elt.result))
+            carry_stmt = GetItem(fn_call.result, index_0.result)
+            y_stmt = GetItem(fn_call.result, index_1.result)
+            carry = carry_stmt.result
+            ys.append(y_stmt.result)
+
+            index.insert_before(node)
+            elt.insert_before(node)
+            fn_call.insert_before(node)
+            carry_stmt.insert_before(node)
+            y_stmt.insert_before(node)
+
+        ys_stmt = New(values=tuple(ys))
+        ys_stmt.insert_before(node)
+        ret = TupleNew(values=(carry, ys_stmt.result))
+        node.replace_by(ret)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_Foldr(self, node: Foldr) -> RewriteResult:
+        return self._rewrite_fold(node, True)
+
+    def rewrite_Foldl(self, node: Foldl) -> RewriteResult:
+        return self._rewrite_fold(node, False)
+
+    def _rewrite_fold(self, node: Foldr | Foldl, reversed: bool) -> RewriteResult:
+        if (coll_len := self._get_collection_len(node.collection)) is None:
+            return RewriteResult()
+
+        acc = node.init
+        for elem_idx in range(coll_len):
+            if reversed:
+                elem_idx = coll_len - elem_idx - 1
+            index = Constant(elem_idx)
+            index.insert_before(node)
+            elt = GetItem(node.collection, index.result)
+            elt.insert_before(node)
+
+            acc_stmt = Call(node.fn, (acc, elt.result))
+            acc_stmt.insert_before(node)
+            acc = acc_stmt.result
+
+        node.result.replace_by(acc)
+        node.delete()
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_ForEach(self, node: ForEach) -> RewriteResult:
+        if (coll_len := self._get_collection_len(node.collection)) is None:
+            return RewriteResult()
+
+        for elem_idx in range(coll_len):
+            index = Constant(elem_idx)
+            index.insert_before(node)
+            elt = GetItem(node.collection, index.result)
+            elt.insert_before(node)
+            fn_call = Call(node.fn, (elt.result,))
+            fn_call.insert_before(node)
+
+        node.delete()
+        return RewriteResult(has_done_something=True)

kirin/dialects/ilist/runtime.py

@@ -0,0 +1,63 @@
+# TODO: replace with something faster
+from typing import Any, Generic, TypeVar, overload
+from dataclasses import dataclass
+from collections.abc import Sequence
+
+T = TypeVar("T")
+L = TypeVar("L")
+
+
+@dataclass
+class IList(Generic[T, L]):
+    """A simple immutable list."""
+
+    data: Sequence[T]
+
+    def __hash__(self) -> int:
+        return id(self)  # do not hash the data
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    @overload
+    def __add__(self, other: "IList[T, Any]") -> "IList[T, Any]": ...
+
+    @overload
+    def __add__(self, other: list[T]) -> "IList[T, Any]": ...
+
+    def __add__(self, other):
+        return IList(self.data + other)
+
+    @overload
+    def __radd__(self, other: "IList[T, Any]") -> "IList[T, Any]": ...
+
+    @overload
+    def __radd__(self, other: list[T]) -> "IList[T, Any]": ...
+
+    def __radd__(self, other):
+        return IList(other + self.data)
+
+    def __repr__(self) -> str:
+        return f"IList({self.data})"
+
+    def __str__(self) -> str:
+        return f"IList({self.data})"
+
+    def __iter__(self):
+        return iter(self.data)
+
+    @overload
+    def __getitem__(self, index: slice) -> "IList[T, Any]": ...
+
+    @overload
+    def __getitem__(self, index: int) -> T: ...
+
+    def __getitem__(self, index: int | slice) -> T | "IList[T, Any]":
+        if isinstance(index, slice):
+            return IList(self.data[index])
+        return self.data[index]
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, IList):
+            return False
+        return self.data == value.data

kirin/dialects/ilist/stmts.py

@@ -0,0 +1,102 @@
+from typing import Sequence
+
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from .runtime import IList
+from ._dialect import dialect
+
+ElemT = types.TypeVar("ElemT")
+ListLen = types.TypeVar("ListLen")
+IListType = types.Generic(IList, ElemT, ListLen)
+
+
+@statement(dialect=dialect, init=False)
+class New(ir.Statement):
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    values: tuple[ir.SSAValue, ...] = info.argument(ElemT)
+    result: ir.ResultValue = info.result(IListType[ElemT])
+
+    def __init__(
+        self,
+        values: Sequence[ir.SSAValue],
+    ) -> None:
+        # get elem type
+        if not values:
+            elem_type = types.Any
+        else:
+            elem_type = values[0].type
+            for v in values:
+                elem_type = elem_type.join(v.type)
+
+        result_type = IListType[elem_type, types.Literal(len(values))]
+        super().__init__(
+            args=values,
+            result_types=(result_type,),
+            args_slice={"values": slice(0, len(values))},
+        )
+
+
+@statement(dialect=dialect)
+class Push(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
+    lst: ir.SSAValue = info.argument(IListType[ElemT])
+    value: ir.SSAValue = info.argument(IListType[ElemT])
+    result: ir.ResultValue = info.result(IListType[ElemT])
+
+
+OutElemT = types.TypeVar("OutElemT")
+
+
+@statement(dialect=dialect)
+class Map(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
+    fn: ir.SSAValue = info.argument(types.MethodType[[ElemT], OutElemT])
+    collection: ir.SSAValue = info.argument(IListType[ElemT, ListLen])
+    result: ir.ResultValue = info.result(IListType[OutElemT, ListLen])
+
+
+@statement(dialect=dialect)
+class Foldr(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
+    fn: ir.SSAValue = info.argument(
+        types.Generic(ir.Method, [ElemT, OutElemT], OutElemT)
+    )
+    collection: ir.SSAValue = info.argument(IListType[ElemT])
+    init: ir.SSAValue = info.argument(OutElemT)
+    result: ir.ResultValue = info.result(OutElemT)
+
+
+@statement(dialect=dialect)
+class Foldl(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
+    fn: ir.SSAValue = info.argument(
+        types.Generic(ir.Method, [OutElemT, ElemT], OutElemT)
+    )
+    collection: ir.SSAValue = info.argument(IListType[ElemT])
+    init: ir.SSAValue = info.argument(OutElemT)
+    result: ir.ResultValue = info.result(OutElemT)
+
+
+CarryT = types.TypeVar("CarryT")
+ResultT = types.TypeVar("ResultT")
+
+
+@statement(dialect=dialect)
+class Scan(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
+    fn: ir.SSAValue = info.argument(
+        types.Generic(ir.Method, [OutElemT, ElemT], types.Tuple[OutElemT, ResultT])
+    )
+    collection: ir.SSAValue = info.argument(IListType[ElemT, ListLen])
+    init: ir.SSAValue = info.argument(OutElemT)
+    result: ir.ResultValue = info.result(
+        types.Tuple[OutElemT, IListType[ResultT, ListLen]]
+    )
+
+
+@statement(dialect=dialect)
+class ForEach(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
+    fn: ir.SSAValue = info.argument(types.Generic(ir.Method, [ElemT], types.NoneType))
+    collection: ir.SSAValue = info.argument(IListType[ElemT])

kirin/dialects/ilist/typeinfer.py

@@ -0,0 +1,120 @@
+from kirin import types
+from kirin.interp import Frame, MethodTable, impl
+from kirin.dialects.eltype import ElType
+from kirin.dialects.py.binop import Add
+from kirin.analysis.typeinfer import TypeInference
+from kirin.dialects.py.indexing import GetItem
+
+from .stmts import New, Push, IListType
+from .runtime import IList
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(MethodTable):
+
+    @staticmethod
+    def _get_list_len(typ: types.Generic):
+        if isinstance(typ.vars[1], types.Literal) and isinstance(typ.vars[1].data, int):
+            return typ.vars[1].data
+        else:
+            return types.Any
+
+    @impl(ElType, types.PyClass(IList))
+    def eltype_list(
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: ElType
+    ):
+        list_type = frame.get(stmt.container)
+        if isinstance(list_type, types.Generic):
+            return (list_type.vars[0],)
+        else:
+            return (types.Any,)
+
+    @impl(New)
+    def new(self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: New):
+        values = frame.get_values(stmt.values)
+        if not values:
+            return (IListType[types.Any, types.Literal(0)],)
+
+        elem_type = values[0]
+        for v in values:
+            elem_type = elem_type.join(v)
+
+        return (IListType[elem_type, types.Literal(len(values))],)
+
+    @impl(Push)
+    def push(
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: Push
+    ):
+        lst_type: types.Generic = frame.get(stmt.lst)  # type: ignore
+        value_type = frame.get(stmt.value)
+        if not lst_type.is_subseteq(IListType):
+            return (types.Bottom,)
+
+        if not lst_type.vars[0].is_subseteq(value_type):
+            return (types.Bottom,)
+
+        lst_len = self._get_list_len(lst_type)
+        if not isinstance(lst_len, int):
+            return (IListType[lst_type.vars[0], types.Any],)
+
+        return (IListType[lst_type.vars[0], types.Literal(lst_len + 1)],)
+
+    @impl(Add, types.PyClass(IList), types.PyClass(IList))
+    def add(self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: Add):
+        lhs_type = frame.get(stmt.lhs)
+        rhs_type = frame.get(stmt.rhs)
+        if not lhs_type.is_subseteq(IListType) or not rhs_type.is_subseteq(IListType):
+            return (types.Bottom,)
+
+        if not isinstance(lhs_type, types.Generic):  # just annotated with list
+            lhs_type = IListType[types.Any, types.Any]
+
+        if not isinstance(rhs_type, types.Generic):
+            rhs_type = IListType[types.Any, types.Any]
+
+        if len(lhs_type.vars) != 2 or len(rhs_type.vars) != 2:
+            raise TypeError("missing type argument for list")
+
+        elem_type = lhs_type.vars[0].join(rhs_type.vars[0])
+
+        lhs_len = self._get_list_len(lhs_type)
+        rhs_len = self._get_list_len(rhs_type)
+        if isinstance(lhs_len, int) and isinstance(rhs_len, int):
+            return (IListType[elem_type, types.Literal(lhs_len + rhs_len)],)
+        return (IListType[elem_type, types.Any],)
+
+    @impl(GetItem, types.PyClass(IList), types.PyClass(int))
+    def getitem(
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: GetItem
+    ):
+        obj_type = frame.get(stmt.obj)
+        if not obj_type.is_subseteq(IListType):
+            raise TypeError(f"Expected list, got {obj_type}")
+
+        # just list type
+        if not isinstance(obj_type, types.Generic):
+            return (types.Any,)
+        else:
+            return (obj_type.vars[0],)
+
+    @impl(GetItem, types.PyClass(IList), types.PyClass(slice))
+    def getitem_slice(
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: GetItem
+    ):
+        obj_type = frame.get(stmt.obj)
+        if not obj_type.is_subseteq(IListType):
+            raise TypeError(f"Expected list, got {obj_type}")
+
+        # just list type
+        if not isinstance(obj_type, types.Generic):
+            return (IListType[types.Any, types.Any],)
+        elif index_ := interp.maybe_const(stmt.index, slice):
+            # TODO: actually calculate the size
+            obj_len = obj_type.vars[1]
+            if not isinstance(obj_len, types.Literal):
+                return (IListType[obj_type.vars[0], types.Any],)
+            LenT = types.Literal(len(range(obj_len.data)[index_]))
+            return (IListType[obj_type.vars[0], LenT],)
+        else:
+            return (IListType[obj_type.vars[0], types.Any],)

kirin/dialects/lowering/__init__.py

@@ -0,0 +1,7 @@
+"""This module contains the dialects for choosing different lowering strategies.
+
+The dialects defined inside this module do not provide any new statements, it only
+provide different lowering strategies for existing statements.
+"""
+
+from . import cf as cf, call as call, func as func

kirin/dialects/lowering/call.py

@@ -0,0 +1,48 @@
+import ast
+
+from kirin import ir, types, lowering
+from kirin.dialects import func
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("lowering.call")
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Call_local(
+        self, state: lowering.LoweringState, callee: ir.SSAValue, node: ast.Call
+    ) -> lowering.Result:
+        args, keywords = self.__lower_Call_args_kwargs(state, node)
+        stmt = func.Call(callee, args, kwargs=keywords)
+        return lowering.Result(state.append_stmt(stmt))
+
+    def lower_Call_global_method(
+        self,
+        state: lowering.LoweringState,
+        method: ir.Method,
+        node: ast.Call,
+    ) -> lowering.Result:
+        args, keywords = self.__lower_Call_args_kwargs(state, node)
+        stmt = func.Invoke(args, callee=method, kwargs=keywords)
+        stmt.result.type = method.return_type or types.Any
+        return lowering.Result(state.append_stmt(stmt))
+
+    def __lower_Call_args_kwargs(
+        self,
+        state: lowering.LoweringState,
+        node: ast.Call,
+    ):
+        args: list[ir.SSAValue] = []
+        for arg in node.args:
+            if isinstance(arg, ast.Starred):  # TODO: support *args
+                raise DialectLoweringError("starred arguments are not supported")
+            else:
+                args.append(state.visit(arg).expect_one())
+
+        keywords = []
+        for kw in node.keywords:
+            keywords.append(kw.arg)
+            args.append(state.visit(kw.value).expect_one())
+
+        return tuple(args), tuple(keywords)

kirin/dialects/lowering/cf.py

@@ -0,0 +1,206 @@
+"""Lowering Python AST to cf dialect.
+"""
+
+import ast
+
+from kirin import ir, types
+from kirin.dialects import cf, py
+from kirin.lowering import Frame, Result, FromPythonAST, LoweringState
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("lowering.cf")
+
+
+@dialect.register
+class CfLowering(FromPythonAST):
+
+    def lower_Pass(self, state: LoweringState, node: ast.Pass) -> Result:
+        state.append_stmt(
+            cf.Branch(arguments=(), successor=state.current_frame.next_block)
+        )
+        return Result()
+
+    def lower_For(self, state: LoweringState, node: ast.For) -> Result:
+        yields: list[str] = []
+
+        def new_block_arg_if_inside_loop(frame: Frame, capture: ir.SSAValue):
+            if not capture.name:
+                raise DialectLoweringError("unexpected loop variable captured")
+            yields.append(capture.name)
+            return frame.entr_block.args.append_from(capture.type, capture.name)
+
+        frame = state.current_frame
+        iterable = state.visit(node.iter).expect_one()
+        iter_stmt = frame.append_stmt(py.iterable.Iter(iterable))
+        none_stmt = frame.append_stmt(py.Constant(None))
+
+        body_frame = state.push_frame(
+            Frame.from_stmts(
+                node.body,
+                state,
+                region=state.current_frame.curr_region,
+                globals=state.current_frame.globals,
+                capture_callback=new_block_arg_if_inside_loop,
+            )
+        )
+        next_value = body_frame.entr_block.args.append_from(types.Any, "next_value")
+        py.unpack.unpacking(state, node.target, next_value)
+        state.exhaust(body_frame)
+        self.branch_next_if_not_terminated(body_frame)
+        yield_args = tuple(body_frame.get_scope(name) for name in yields)
+        next_stmt = py.iterable.Next(iter_stmt.iter)
+        cond_stmt = py.cmp.Is(next_stmt.value, none_stmt.result)
+        body_frame.next_block.stmts.append(next_stmt)
+        body_frame.next_block.stmts.append(cond_stmt)
+        body_frame.next_block.stmts.append(
+            cf.ConditionalBranch(
+                cond_stmt.result,
+                yield_args,
+                (next_stmt.value,) + yield_args,
+                then_successor=frame.next_block,
+                else_successor=body_frame.entr_block,
+            )
+        )
+        state.pop_frame()
+
+        # insert the branch to the entrance of the loop (the code block before loop)
+        next_stmt = frame.append_stmt(py.iterable.Next(iter_stmt.iter))
+        cond_stmt = frame.append_stmt(py.cmp.Is(next_stmt.value, none_stmt.result))
+        yield_args = tuple(frame.get_scope(name) for name in yields)
+        frame.append_stmt(
+            cf.ConditionalBranch(
+                cond_stmt.result,
+                yield_args,
+                (next_stmt.value,) + yield_args,
+                then_successor=frame.next_block,  # empty iterator
+                else_successor=body_frame.entr_block,
+            )
+        )
+
+        frame.jump_next()
+        for name, arg in zip(yields, yield_args):
+            input = frame.curr_block.args.append_from(arg.type, name)
+            frame.defs[name] = input
+        return Result()
+
+    def lower_If(self, state: LoweringState, node: ast.If) -> Result:
+        cond = state.visit(node.test).expect_one()
+        frame = state.current_frame
+        before_block = frame.curr_block
+        if_frame = state.push_frame(
+            Frame.from_stmts(
+                node.body,
+                state,
+                region=frame.curr_region,
+                globals=frame.globals,
+            )
+        )
+        true_cond = if_frame.entr_block.args.append_from(types.Bool, cond.name)
+        if cond.name:
+            if_frame.defs[cond.name] = true_cond
+        state.exhaust()
+        self.branch_next_if_not_terminated(if_frame)
+        state.pop_frame()
+
+        else_frame = state.push_frame(
+            Frame.from_stmts(
+                node.orelse,
+                state,
+                region=frame.curr_region,
+                globals=frame.globals,
+            )
+        )
+        true_cond = else_frame.entr_block.args.append_from(types.Bool, cond.name)
+        if cond.name:
+            else_frame.defs[cond.name] = true_cond
+        state.exhaust()
+        self.branch_next_if_not_terminated(else_frame)
+        state.pop_frame()
+
+        after_frame = state.push_frame(
+            Frame.from_stmts(
+                frame.stream.split(),
+                state,
+                region=frame.curr_region,
+                globals=frame.globals,
+            )
+        )
+
+        after_frame.defs.update(frame.defs)
+        phi: set[str] = set()
+        for name in if_frame.defs.keys():
+            if frame.get(name):
+                phi.add(name)
+            elif name in else_frame.defs:
+                phi.add(name)
+
+        for name in else_frame.defs.keys():
+            if frame.get(name):  # not defined in if_frame
+                phi.add(name)
+
+        for name in phi:
+            after_frame.defs[name] = after_frame.entr_block.args.append_from(
+                types.Any, name
+            )
+
+        state.exhaust()
+        self.branch_next_if_not_terminated(after_frame)
+        after_frame.next_block.stmts.append(
+            cf.Branch(arguments=(), successor=frame.next_block)
+        )
+        state.pop_frame()
+
+        if_args = []
+        for name in phi:
+            if value := if_frame.get(name):
+                if_args.append(value)
+            else:
+                raise DialectLoweringError(f"undefined variable {name} in if branch")
+
+        else_args = []
+        for name in phi:
+            if value := else_frame.get(name):
+                else_args.append(value)
+            else:
+                raise DialectLoweringError(f"undefined variable {name} in else branch")
+
+        if_frame.next_block.stmts.append(
+            cf.Branch(
+                arguments=tuple(if_args),
+                successor=after_frame.entr_block,
+            )
+        )
+        else_frame.next_block.stmts.append(
+            cf.Branch(
+                arguments=tuple(else_args),
+                successor=after_frame.entr_block,
+            )
+        )
+        before_block.stmts.append(
+            cf.ConditionalBranch(
+                cond=cond,
+                then_arguments=(cond,),
+                then_successor=if_frame.entr_block,
+                else_arguments=(cond,),
+                else_successor=else_frame.entr_block,
+            )
+        )
+        frame.jump_next()
+        return Result()
+
+    def branch_next_if_not_terminated(self, frame: Frame):
+        """Branch to the next block if the current block is not terminated.
+
+        This must be used after exhausting the current frame and before popping the frame.
+        """
+        if not frame.curr_block.last_stmt or not frame.curr_block.last_stmt.has_trait(
+            ir.IsTerminator
+        ):
+            frame.curr_block.stmts.append(
+                cf.Branch(arguments=(), successor=frame.next_block)
+            )
+
+    def current_block_terminated(self, frame: Frame):
+        return frame.curr_block.last_stmt and frame.curr_block.last_stmt.has_trait(
+            ir.IsTerminator
+        )