kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/func/stmts.py

@@ -0,0 +1,233 @@
+from kirin import types
+from kirin.ir import (
+    Pure,
+    Method,
+    Region,
+    HasParent,
+    MaybePure,
+    Statement,
+    ResultValue,
+    ConstantLike,
+    HasSignature,
+    IsTerminator,
+    SSACFGRegion,
+    IsolatedFromAbove,
+    SymbolOpInterface,
+    CallableStmtInterface,
+)
+from kirin.decl import info, statement
+from kirin.ir.ssa import SSAValue
+from kirin.exceptions import VerificationError
+from kirin.print.printer import Printer
+from kirin.dialects.func.attrs import Signature, MethodType
+from kirin.dialects.func.dialect import dialect
+
+from .._pprint_helper import pprint_calllike
+
+
+class FuncOpCallableInterface(CallableStmtInterface["Function"]):
+
+    @classmethod
+    def get_callable_region(cls, stmt: "Function") -> Region:
+        return stmt.body
+
+
+@statement(dialect=dialect)
+class Function(Statement):
+    name = "func"
+    traits = frozenset(
+        {
+            IsolatedFromAbove(),
+            SymbolOpInterface(),
+            HasSignature(),
+            FuncOpCallableInterface(),
+            SSACFGRegion(),
+        }
+    )
+    sym_name: str = info.attribute()
+    """The symbol name of the function."""
+    signature: Signature = info.attribute()
+    body: Region = info.region(multi=True)
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="keyword"):
+            printer.print_name(self)
+            printer.plain_print(" ")
+
+        with printer.rich(style="symbol"):
+            printer.plain_print(self.sym_name)
+
+        printer.print_seq(self.signature.inputs, prefix="(", suffix=")", delim=", ")
+
+        with printer.rich(style="comment"):
+            printer.plain_print(" -> ")
+            printer.print(self.signature.output)
+            printer.plain_print(" ")
+
+        printer.print(self.body)
+
+        with printer.rich(style="comment"):
+            printer.plain_print(f" // func.func {self.sym_name}")
+
+
+@statement(dialect=dialect)
+class Call(Statement):
+    name = "call"
+    traits = frozenset({MaybePure()})
+    # not a fixed type here so just any
+    callee: SSAValue = info.argument()
+    inputs: tuple[SSAValue, ...] = info.argument()
+    kwargs: tuple[str, ...] = info.attribute(default_factory=lambda: ())
+    result: ResultValue = info.result()
+    purity: bool = info.attribute(default=False)
+
+    def print_impl(self, printer: Printer) -> None:
+        pprint_calllike(self, printer.state.ssa_id[self.callee], printer)
+
+
+@statement(dialect=dialect)
+class ConstantNone(Statement):
+    """A constant None value.
+
+    This is mainly used to represent the None return value of a function
+    to match Python semantics.
+    """
+
+    name = "const.none"
+    traits = frozenset({Pure(), ConstantLike()})
+    result: ResultValue = info.result(types.NoneType)
+
+
+@statement(dialect=dialect, init=False)
+class Return(Statement):
+    name = "return"
+    traits = frozenset({IsTerminator(), HasParent((Function,))})
+    value: SSAValue = info.argument()
+
+    def __init__(self, value_or_stmt: SSAValue | Statement | None = None) -> None:
+        if isinstance(value_or_stmt, SSAValue):
+            args = [value_or_stmt]
+        elif isinstance(value_or_stmt, Statement):
+            if len(value_or_stmt._results) == 1:
+                args = [value_or_stmt._results[0]]
+            else:
+                raise ValueError(
+                    f"expected a single result, got {len(value_or_stmt._results)} results from {value_or_stmt.name}"
+                )
+        elif value_or_stmt is None:
+            args = []
+        else:
+            raise ValueError(f"expected SSAValue or Statement, got {value_or_stmt}")
+
+        super().__init__(args=args, args_slice={"value": 0})
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="keyword"):
+            printer.print_name(self)
+
+        if self.args:
+            printer.plain_print(" ")
+            printer.print_seq(self.args, delim=", ")
+
+    def verify(self) -> None:
+        if not self.args:
+            raise VerificationError(
+                self, "return statement must have at least one value"
+            )
+
+        if len(self.args) > 1:
+            raise VerificationError(
+                self,
+                "return statement must have at most one value"
+                ", wrap multiple values in a tuple",
+            )
+
+
+@statement(dialect=dialect)
+class Lambda(Statement):
+    name = "lambda"
+    traits = frozenset(
+        {
+            Pure(),
+            HasSignature(),
+            SymbolOpInterface(),
+            FuncOpCallableInterface(),
+            SSACFGRegion(),
+        }
+    )
+    sym_name: str = info.attribute()
+    signature: Signature = info.attribute()
+    captured: tuple[SSAValue, ...] = info.argument()
+    body: Region = info.region(multi=True)
+    result: ResultValue = info.result(MethodType)
+
+    def verify(self) -> None:
+        if self.body.blocks.isempty():
+            raise VerificationError(self, "lambda body must not be empty")
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="keyword"):
+            printer.print_name(self)
+        printer.plain_print(" ")
+
+        with printer.rich(style="symbol"):
+            printer.plain_print(self.sym_name)
+
+        printer.print_seq(self.captured, prefix="(", suffix=")", delim=", ")
+
+        with printer.rich(style="bright_black"):
+            printer.plain_print(" -> ")
+            printer.print(self.signature.output)
+
+        printer.plain_print(" ")
+        printer.print(self.body)
+
+        with printer.rich(style="black"):
+            printer.plain_print(f" // func.lambda {self.sym_name}")
+
+
+@statement(dialect=dialect)
+class GetField(Statement):
+    name = "getfield"
+    traits = frozenset({Pure()})
+    obj: SSAValue = info.argument(MethodType)
+    field: int = info.attribute()
+    # NOTE: mypy somehow doesn't understand default init=False
+    result: ResultValue = info.result(init=False)
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(
+            "(", printer.state.ssa_id[self.obj], ", ", str(self.field), ")"
+        )
+        with printer.rich(style="black"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class Invoke(Statement):
+    name = "invoke"
+    traits = frozenset({MaybePure()})
+    callee: Method = info.attribute()
+    inputs: tuple[SSAValue, ...] = info.argument()
+    kwargs: tuple[str, ...] = info.attribute()
+    result: ResultValue = info.result()
+    purity: bool = info.attribute(default=False)
+
+    def print_impl(self, printer: Printer) -> None:
+        pprint_calllike(self, self.callee.sym_name, printer)
+
+    def verify(self) -> None:
+        if self.kwargs:
+            for name in self.kwargs:
+                if name not in self.callee.arg_names:
+                    raise VerificationError(
+                        self,
+                        f"method {self.callee.sym_name} does not have argument {name}",
+                    )
+        elif len(self.callee.arg_names) - 1 != len(self.args):
+            raise VerificationError(
+                self,
+                f"expected {len(self.callee.arg_names)} arguments, got {len(self.args)}",
+            )

kirin/dialects/func/typeinfer.py

@@ -0,0 +1,124 @@
+from typing import Iterable, cast
+
+from kirin import ir, types
+from kirin.interp import Frame, MethodTable, ReturnValue, impl
+from kirin.analysis import const
+from kirin.analysis.typeinfer import TypeInference, TypeResolution
+from kirin.dialects.func.stmts import (
+    Call,
+    Invoke,
+    Lambda,
+    Return,
+    GetField,
+    ConstantNone,
+)
+from kirin.dialects.func.dialect import dialect
+
+
+# NOTE: a lot of the type infer rules are same as the builtin dialect
+@dialect.register(key="typeinfer")
+class TypeInfer(MethodTable):
+
+    @impl(ConstantNone)
+    def const_none(self, interp: TypeInference, frame: Frame, stmt: ConstantNone):
+        return (types.NoneType,)
+
+    @impl(Return)
+    def return_(self, interp: TypeInference, frame: Frame, stmt: Return) -> ReturnValue:
+        if (
+            isinstance(hint := stmt.value.hints.get("const"), const.Value)
+            and hint.data is not None
+        ):
+            return ReturnValue(types.Literal(hint.data))
+        return ReturnValue(frame.get(stmt.value))
+
+    @impl(Call)
+    def call(self, interp: TypeInference, frame: Frame, stmt: Call):
+        # give up on dynamic method calls
+        mt = interp.maybe_const(stmt.callee, ir.Method)
+        if mt is None:
+            return self._solve_method_type(interp, frame, stmt)
+        return self._invoke_method(
+            interp,
+            frame,
+            mt,
+            stmt.args[1:],
+            interp.permute_values(
+                mt.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+            ),
+        )
+
+    def _solve_method_type(self, interp: TypeInference, frame: Frame, stmt: Call):
+        mt_inferred = frame.get(stmt.callee)
+        if not isinstance(mt_inferred, types.Generic):
+            return (types.Bottom,)
+
+        if len(mt_inferred.vars) != 2:
+            return (types.Bottom,)
+
+        args = mt_inferred.vars[0]
+        result = mt_inferred.vars[1]
+        if not args.is_subseteq(types.Tuple):
+            return (types.Bottom,)
+
+        resolve = TypeResolution()
+        args = cast(types.Generic, args)
+        for arg, value in zip(args.vars, frame.get_values(stmt.inputs)):
+            resolve.solve(arg, value)
+        return (resolve.substitute(result),)
+
+    @impl(Invoke)
+    def invoke(self, interp: TypeInference, frame: Frame, stmt: Invoke):
+        return self._invoke_method(
+            interp,
+            frame,
+            stmt.callee,
+            stmt.inputs,
+            interp.permute_values(
+                stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+            ),
+        )
+
+    def _invoke_method(
+        self,
+        interp: TypeInference,
+        frame: Frame,
+        mt: ir.Method,
+        args: Iterable[ir.SSAValue],
+        values: tuple,
+    ):
+        if mt.inferred:  # so we don't end up in infinite loop
+            return (mt.return_type,)
+
+        # NOTE: narrowing the argument type based on method signature
+        inputs = tuple(
+            typ.meet(input_typ) for typ, input_typ in zip(mt.arg_types, values)
+        )
+
+        # NOTE: we use lower bound here because function call contains an
+        # implicit type check at call site. This will be validated either compile time
+        # or runtime.
+        # update the results with the narrowed types
+        frame.set_values(args, inputs)
+        _, ret = interp.run_method(mt, inputs)
+        return (ret,)
+
+    @impl(Lambda)
+    def lambda_(
+        self, interp_: TypeInference, frame: Frame[types.TypeAttribute], stmt: Lambda
+    ):
+        body_frame, ret = interp_.run_callable(
+            stmt,
+            (types.MethodType,)
+            + tuple(arg.type for arg in stmt.body.blocks[0].args[1:]),
+        )
+        argtypes = tuple(arg.type for arg in stmt.body.blocks[0].args[1:])
+        ret = types.MethodType[[*argtypes], ret]
+        frame.entries.update(body_frame.entries)  # pass results back to upper frame
+        self_ = stmt.body.blocks[0].args[0]
+        frame.set(self_, ret)
+        return (ret,)
+
+    @impl(GetField)
+    def getfield(self, interp: TypeInference, frame, stmt: GetField):
+        return (stmt.result.type,)

kirin/dialects/ilist/__init__.py

@@ -0,0 +1,33 @@
+"""
+Immutable list dialect for Python.
+
+This dialect provides a simple, immutable list dialect similar
+to Python's built-in list type.
+"""
+
+from . import (
+    interp as interp,
+    rewrite as rewrite,
+    lowering as lowering,
+    typeinfer as typeinfer,
+)
+from .stmts import (
+    Map as Map,
+    New as New,
+    Push as Push,
+    Scan as Scan,
+    Foldl as Foldl,
+    Foldr as Foldr,
+    ForEach as ForEach,
+    IListType as IListType,
+)
+from .passes import IListDesugar as IListDesugar
+from .runtime import IList as IList
+from ._dialect import dialect as dialect
+from ._wrapper import (
+    map as map,
+    scan as scan,
+    foldl as foldl,
+    foldr as foldr,
+    for_each as for_each,
+)

kirin/dialects/ilist/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("py.ilist")

kirin/dialects/ilist/_wrapper.py

@@ -0,0 +1,51 @@
+import typing
+
+from kirin.lowering import wraps
+
+from . import stmts
+from .runtime import IList
+
+ElemT = typing.TypeVar("ElemT")
+OutElemT = typing.TypeVar("OutElemT")
+LenT = typing.TypeVar("LenT")
+ResultT = typing.TypeVar("ResultT")
+
+# NOTE: we use Callable here to make nested function work.
+
+
+@wraps(stmts.Map)
+def map(
+    fn: typing.Callable[[ElemT], OutElemT],
+    collection: IList[ElemT, LenT] | list[ElemT],
+) -> IList[OutElemT, LenT]: ...
+
+
+@wraps(stmts.Foldr)
+def foldr(
+    fn: typing.Callable[[ElemT, OutElemT], OutElemT],
+    collection: IList[ElemT, LenT] | list[ElemT],
+    init: OutElemT,
+) -> OutElemT: ...
+
+
+@wraps(stmts.Foldl)
+def foldl(
+    fn: typing.Callable[[OutElemT, ElemT], OutElemT],
+    collection: IList[ElemT, LenT] | list[ElemT],
+    init: OutElemT,
+) -> OutElemT: ...
+
+
+@wraps(stmts.Scan)
+def scan(
+    fn: typing.Callable[[OutElemT, ElemT], tuple[OutElemT, ResultT]],
+    collection: IList[ElemT, LenT] | list[ElemT],
+    init: OutElemT,
+) -> tuple[OutElemT, IList[ResultT, LenT]]: ...
+
+
+@wraps(stmts.ForEach)
+def for_each(
+    fn: typing.Callable[[ElemT], typing.Any],
+    collection: IList[ElemT, LenT] | list[ElemT],
+) -> None: ...

kirin/dialects/ilist/interp.py

@@ -0,0 +1,85 @@
+from kirin import ir, types
+from kirin.interp import Frame, Interpreter, MethodTable, impl
+from kirin.dialects.py.len import Len
+from kirin.dialects.py.binop import Add
+from kirin.dialects.py.range import Range
+
+from .stmts import Map, New, Push, Scan, Foldl, Foldr, ForEach
+from .runtime import IList
+from ._dialect import dialect
+
+
+@dialect.register
+class IListInterpreter(MethodTable):
+
+    @impl(Range)
+    def _range(self, interp, frame: Frame, stmt: Range):
+        return (IList(range(*frame.get_values(stmt.args))),)
+
+    @impl(New)
+    def new(self, interp, frame: Frame, stmt: New):
+        return (IList(list(frame.get_values(stmt.values))),)
+
+    @impl(Len, types.PyClass(IList))
+    def len(self, interp, frame: Frame, stmt: Len):
+        return (len(frame.get(stmt.value).data),)
+
+    @impl(Add, types.PyClass(IList), types.PyClass(IList))
+    def add(self, interp, frame: Frame, stmt: Add):
+        return (IList(frame.get(stmt.lhs).data + frame.get(stmt.rhs).data),)
+
+    @impl(Push)
+    def push(self, interp, frame: Frame, stmt: Push):
+        return (IList(frame.get(stmt.lst).data + [frame.get(stmt.value)]),)
+
+    @impl(Map)
+    def map(self, interp: Interpreter, frame: Frame, stmt: Map):
+        fn: ir.Method = frame.get(stmt.fn)
+        coll: IList = frame.get(stmt.collection)
+        ret = []
+        for elem in coll.data:
+            # NOTE: assume fn has been type checked
+            _, item = interp.run_method(fn, (elem,))
+            ret.append(item)
+        return (IList(ret),)
+
+    @impl(Scan)
+    def scan(self, interp: Interpreter, frame: Frame, stmt: Scan):
+        fn: ir.Method = frame.get(stmt.fn)
+        init = frame.get(stmt.init)
+        coll: IList = frame.get(stmt.collection)
+
+        carry = init
+        ys = []
+        for elem in coll.data:
+            # NOTE: assume fn has been type checked
+            _, (carry, y) = interp.run_method(fn, (carry, elem))
+            ys.append(y)
+        return ((carry, IList(ys)),)
+
+    @impl(Foldr)
+    def foldr(self, interp: Interpreter, frame: Frame, stmt: Foldr):
+        return self.fold(interp, frame, stmt, reversed(frame.get(stmt.collection).data))
+
+    @impl(Foldl)
+    def foldl(self, interp: Interpreter, frame: Frame, stmt: Foldl):
+        return self.fold(interp, frame, stmt, frame.get(stmt.collection).data)
+
+    def fold(self, interp: Interpreter, frame: Frame, stmt: Foldr | Foldl, coll):
+        fn: ir.Method = frame.get(stmt.fn)
+        init = frame.get(stmt.init)
+
+        acc = init
+        for elem in coll:
+            # NOTE: assume fn has been type checked
+            _, acc = interp.run_method(fn, (acc, elem))
+        return (acc,)
+
+    @impl(ForEach)
+    def for_each(self, interp: Interpreter, frame: Frame, stmt: ForEach):
+        fn: ir.Method = frame.get(stmt.fn)
+        coll: IList = frame.get(stmt.collection)
+        for elem in coll.data:
+            # NOTE: assume fn has been type checked
+            interp.run_method(fn, (elem,))
+        return (None,)

kirin/dialects/ilist/lowering.py

@@ -0,0 +1,25 @@
+import ast
+
+from kirin import types
+from kirin.lowering import Result, FromPythonAST, LoweringState
+
+from . import stmts as ilist
+from ._dialect import dialect
+
+
+@dialect.register
+class PythonLowering(FromPythonAST):
+
+    def lower_List(self, state: LoweringState, node: ast.List) -> Result:
+        elts = tuple(state.visit(each).expect_one() for each in node.elts)
+
+        if len(elts):
+            typ = elts[0].type
+            for each in elts:
+                typ = typ.join(each.type)
+        else:
+            typ = types.Any
+
+        stmt = ilist.New(values=tuple(elts))
+        state.append_stmt(stmt)
+        return Result(stmt)

kirin/dialects/ilist/passes.py

@@ -0,0 +1,32 @@
+from kirin import ir, types
+from kirin.rewrite import Walk, Chain, Fixpoint
+from kirin.passes.abc import Pass
+from kirin.rewrite.result import RewriteResult
+from kirin.dialects.ilist.rewrite import List2IList, ConstList2IList
+
+
+class IListDesugar(Pass):
+    """This pass desugars the Python list dialect
+    to the immutable list dialect by rewriting all
+    constant `list` type into `IList` type.
+    """
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        for arg in mt.args:
+            _check_list(arg.type, arg.type)
+        return Fixpoint(Walk(Chain(ConstList2IList(), List2IList()))).rewrite(mt.code)
+
+
+def _check_list(total: types.TypeAttribute, type_: types.TypeAttribute):
+    if isinstance(type_, types.Generic):
+        _check_list(total, type_.body)
+        for var in type_.vars:
+            _check_list(total, var)
+        if type_.vararg:
+            _check_list(total, type_.vararg.typ)
+    elif isinstance(type_, types.PyClass):
+        if issubclass(type_.typ, list):
+            raise TypeError(
+                f"Invalid type {total} for this kernel, use IList instead of {type_}."
+            )
+    return

kirin/dialects/ilist/rewrite/__init__.py

@@ -0,0 +1,3 @@
+from .list import List2IList as List2IList
+from .const import ConstList2IList as ConstList2IList
+from .unroll import Unroll as Unroll

kirin/dialects/ilist/rewrite/const.py

@@ -0,0 +1,45 @@
+from kirin import ir, types
+from kirin.analysis import const
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+
+from ..stmts import IListType
+from ..runtime import IList
+
+
+class ConstList2IList(RewriteRule):
+    """Rewrite type annotation for SSAValue with constant `IList`
+    in `Hinted` type. This should be run after constant folding and
+    `WrapConst` rule.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        has_done_something = False
+        for result in node.results:
+            if not isinstance(hint := result.hints.get("const"), const.Value):
+                continue
+
+            typ = result.type
+            data = hint.data
+            if isinstance(typ, types.PyClass) and typ.is_subseteq(types.PyClass(IList)):
+                has_done_something = self._rewrite_IList_type(result, data)
+            elif isinstance(typ, types.Generic) and typ.body.is_subseteq(
+                types.PyClass(IList)
+            ):
+                has_done_something = self._rewrite_IList_type(result, data)
+        return RewriteResult(has_done_something=has_done_something)
+
+    def _rewrite_IList_type(self, result: ir.SSAValue, data):
+        if not isinstance(data, IList):
+            return False
+
+        if not data.data:
+            return False
+
+        elem_type = types.PyClass(type(data[0]))
+        for elem in data.data[1:]:
+            elem_type = elem_type.join(types.PyClass(type(elem)))
+
+        result.type = IListType[elem_type, types.Literal(len(data.data))]
+        result.hints["const"] = const.Value(data)
+        return True

kirin/dialects/ilist/rewrite/list.py

@@ -0,0 +1,38 @@
+from kirin import ir, types
+from kirin.dialects.py import constant
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+from kirin.dialects.ilist.stmts import IListType
+from kirin.dialects.ilist.runtime import IList
+
+
+class List2IList(RewriteRule):
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        has_done_something = False
+        for arg in node.args:
+            has_done_something = self._rewrite_SSAValue_type(arg)
+        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:
+            has_done_something = self._rewrite_SSAValue_type(result)
+
+        if has_done_something and isinstance(node, constant.Constant):
+            node.replace_by(constant.Constant(value=IList(data=node.value)))
+
+        return RewriteResult(has_done_something=has_done_something)
+
+    def _rewrite_SSAValue_type(self, value: ir.SSAValue):
+        # NOTE: cannot use issubseteq here because type can be Bottom
+        if isinstance(value.type, types.Generic) and issubclass(
+            value.type.body.typ, list
+        ):
+            value.type = IListType[value.type.vars[0], types.Any]
+            return True
+
+        elif isinstance(value.type, types.PyClass) and issubclass(value.type.typ, list):
+            value.type = IListType[types.Any, types.Any]
+            return True
+        return False