kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/py/cmp/__init__.py

@@ -0,0 +1,16 @@
+"""The cmp dialect for Python.
+
+This module contains the dialect for comparison semantics in Python, including:
+
+- The `Eq`, `NotEq`, `Lt`, `LtE`, `Gt`, `GtE`, `Is`, and `IsNot` statement classes.
+- The lowering pass for comparison operations.
+- The concrete implementation of comparison operations.
+- The Julia emitter for comparison operations.
+
+This dialect maps `ast.Compare` nodes to the `Eq`, `NotEq`, `Lt`, `LtE`,
+`Gt`, `GtE`, `Is`, and `IsNot` statements.
+"""
+
+from . import julia as julia, interp as interp, lowering as lowering
+from .stmts import *  # noqa: F403
+from ._dialect import dialect as dialect

kirin/dialects/py/cmp/_dialect.py

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

kirin/dialects/py/cmp/interp.py

@@ -0,0 +1,48 @@
+from kirin import interp
+
+from . import stmts as cmp
+from ._dialect import dialect
+
+
+@dialect.register
+class CmpMethod(interp.MethodTable):
+
+    @interp.impl(cmp.Eq)
+    def eq(self, interp, frame: interp.Frame, stmt: cmp.Eq):
+        return (frame.get(stmt.lhs) == frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.NotEq)
+    def not_eq(self, interp, frame: interp.Frame, stmt: cmp.NotEq):
+        return (frame.get(stmt.lhs) != frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.Lt)
+    def lt(self, interp, frame: interp.Frame, stmt: cmp.Lt):
+        return (frame.get(stmt.lhs) < frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.LtE)
+    def lt_eq(self, interp, frame: interp.Frame, stmt: cmp.LtE):
+        return (frame.get(stmt.lhs) <= frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.Gt)
+    def gt(self, interp, frame: interp.Frame, stmt: cmp.Gt):
+        return (frame.get(stmt.lhs) > frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.GtE)
+    def gt_eq(self, interp, frame: interp.Frame, stmt: cmp.GtE):
+        return (frame.get(stmt.lhs) >= frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.In)
+    def in_(self, interp, frame: interp.Frame, stmt: cmp.In):
+        return (frame.get(stmt.lhs) in frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.NotIn)
+    def not_in(self, interp, frame: interp.Frame, stmt: cmp.NotIn):
+        return (frame.get(stmt.lhs) not in frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.Is)
+    def is_(self, interp, frame: interp.Frame, stmt: cmp.Is):
+        return (frame.get(stmt.lhs) is frame.get(stmt.rhs),)
+
+    @interp.impl(cmp.IsNot)
+    def is_not(self, interp, frame: interp.Frame, stmt: cmp.IsNot):
+        return (frame.get(stmt.lhs) is not frame.get(stmt.rhs),)

kirin/dialects/py/cmp/julia.py

@@ -0,0 +1,33 @@
+from kirin import interp
+from kirin.emit.julia import EmitJulia, EmitStrFrame
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="emit.julia")
+class JuliaTable(interp.MethodTable):
+
+    @interp.impl(stmts.Eq)
+    def emit_Eq(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Eq):
+        return emit.emit_binaryop(frame, "==", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.GtE)
+    def emit_GtE(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.GtE):
+        return emit.emit_binaryop(frame, ">=", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.LtE)
+    def emit_LtE(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.LtE):
+        return emit.emit_binaryop(frame, "<=", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.NotEq)
+    def emit_NotEq(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.NotEq):
+        return emit.emit_binaryop(frame, "!=", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Gt)
+    def emit_Gt(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Gt):
+        return emit.emit_binaryop(frame, ">", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Lt)
+    def emit_Lt(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Lt):
+        return emit.emit_binaryop(frame, "<", stmt.lhs, stmt.rhs, stmt.result)

kirin/dialects/py/cmp/lowering.py

@@ -0,0 +1,45 @@
+import ast
+
+from kirin.ir import SSAValue
+from kirin.lowering import Result, FromPythonAST, LoweringState
+from kirin.exceptions import DialectLoweringError
+from kirin.dialects.py import boolop
+
+from . import stmts as cmp
+from ._dialect import dialect
+
+
+@dialect.register
+class PythonLowering(FromPythonAST):
+
+    def lower_Compare(self, state: LoweringState, node: ast.Compare) -> Result:
+        # NOTE: a key difference here is we need to lower
+        # the multi-argument comparison operators into binary operators
+        # since low-level comparision operators are binary + we need a static
+        # number of arguments in each instruction
+        lhs = state.visit(node.left).expect_one()
+
+        comparators = [
+            state.visit(comparator).expect_one() for comparator in node.comparators
+        ]
+
+        cmp_results: list[SSAValue] = []
+        for op, rhs in zip(node.ops, comparators):
+            if op := getattr(cmp, op.__class__.__name__, None):
+                stmt = op(lhs=lhs, rhs=rhs)
+            else:
+                raise DialectLoweringError(f"unsupported compare operator {op}")
+            state.append_stmt(stmt)
+            cmp_results.append(Result(stmt).expect_one())
+            lhs = rhs
+
+        if len(cmp_results) == 1:
+            return Result(cmp_results)
+
+        lhs = cmp_results[0]
+        for op in cmp_results[1:]:
+            stmt = boolop.And(lhs=lhs, rhs=op)
+            state.append_stmt(stmt)
+            lhs = stmt.result
+
+        return Result(lhs)

kirin/dialects/py/cmp/stmts.py

@@ -0,0 +1,62 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from ._dialect import dialect
+
+
+@statement
+class Cmp(ir.Statement):
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument()
+    rhs: ir.SSAValue = info.argument()
+    result: ir.ResultValue = info.result(types.Bool)
+
+
+@statement(dialect=dialect)
+class Eq(Cmp):
+    name = "eq"
+
+
+@statement(dialect=dialect)
+class NotEq(Cmp):
+    name = "ne"
+
+
+@statement(dialect=dialect)
+class Lt(Cmp):
+    name = "lt"
+
+
+@statement(dialect=dialect)
+class Gt(Cmp):
+    name = "gt"
+
+
+@statement(dialect=dialect)
+class LtE(Cmp):
+    name = "lte"
+
+
+@statement(dialect=dialect)
+class GtE(Cmp):
+    name = "gte"
+
+
+@statement(dialect=dialect)
+class Is(Cmp):
+    name = "is"
+
+
+@statement(dialect=dialect)
+class IsNot(Cmp):
+    name = "is_not"
+
+
+@statement(dialect=dialect)
+class In(Cmp):
+    name = "in"
+
+
+@statement(dialect=dialect)
+class NotIn(Cmp):
+    name = "not_in"

kirin/dialects/py/constant.py

@@ -0,0 +1,79 @@
+"""Constant statement for Python dialect.
+
+This module contains the dialect for the Python `constant` statement, including:
+
+- The `Constant` statement class.
+- The lowering pass for the `constant` statement.
+- The concrete implementation of the `constant` statement.
+- The Julia emitter for the `constant` statement.
+
+This dialect maps `ast.Constant` nodes to the `Constant` statement.
+"""
+
+import ast
+from typing import Generic, TypeVar
+
+from kirin import ir, types, interp, lowering, exceptions
+from kirin.decl import info, statement
+from kirin.print import Printer
+from kirin.emit.julia import EmitJulia, EmitStrFrame
+
+dialect = ir.Dialect("py.constant")
+
+T = TypeVar("T", covariant=True)
+
+
+@statement(dialect=dialect)
+class Constant(ir.Statement, Generic[T]):
+    name = "constant"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), ir.FromPythonCall()})
+    value: T = info.attribute()
+    result: ir.ResultValue = info.result()
+
+    # NOTE: we allow py.Constant take data.PyAttr too
+    def __init__(self, value: T | ir.PyAttr[T]) -> None:
+        if not isinstance(value, ir.PyAttr):
+            value = ir.PyAttr(value)
+        super().__init__(
+            attributes={"value": value},
+            result_types=(value.type,),
+        )
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+    def typecheck(self) -> None:
+        if not isinstance(self.result.type, types.TypeAttribute):
+            raise exceptions.VerificationError(
+                self, f"Expected result type to be PyType, got {self.result.type}"
+            )
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Constant(
+        self, state: lowering.LoweringState, node: ast.Constant
+    ) -> lowering.Result:
+        return lowering.Result(state.append_stmt(Constant(node.value)))
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Constant)
+    def constant(self, interp, frame: interp.Frame, stmt: Constant):
+        return (stmt.value,)
+
+
+@dialect.register(key="emit.julia")
+class JuliaTable(interp.MethodTable):
+
+    @interp.impl(Constant)
+    def emit_Constant(self, emit: EmitJulia, frame: EmitStrFrame, stmt: Constant):
+        return (emit.emit_attribute(ir.PyAttr(stmt.value)),)

kirin/dialects/py/indexing.py

@@ -0,0 +1,251 @@
+"""The indexing dialect for Python.
+
+This module contains the dialect for the Python indexing syntax, including:
+
+- The `GetItem` statement class.
+- A base class `Subscript` for indexing statements.
+- A trait `GetItemLike` for indexing statements.
+- The lowering pass for the indexing statement.
+- The concrete implementation of the indexing statement.
+- The constant propagation implementation (special case) of the indexing statement.
+- The type inference implementation of the indexing statement.
+- A canonical rewrite rule for the rewriting of a given getitem-like
+    statement to another getitem-like statement.
+"""
+
+import ast
+from abc import abstractmethod
+from typing import Generic, TypeVar
+from dataclasses import dataclass
+
+from kirin import ir, types, interp, lowering, exceptions
+from kirin.decl import info, statement
+from kirin.analysis import const
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.analysis.typeinfer import TypeInference
+
+dialect = ir.Dialect("py.indexing")
+
+GetItemLikeStmt = TypeVar("GetItemLikeStmt", bound=ir.Statement)
+
+
+@dataclass(frozen=True, eq=False)
+class GetItemLike(ir.StmtTrait, Generic[GetItemLikeStmt]):
+
+    @abstractmethod
+    def get_object(self, stmt: GetItemLikeStmt) -> ir.SSAValue: ...
+
+    @abstractmethod
+    def get_index(self, stmt: GetItemLikeStmt) -> ir.SSAValue: ...
+
+    @abstractmethod
+    def new(
+        self, stmt_type: type[GetItemLikeStmt], obj: ir.SSAValue, index: ir.SSAValue
+    ) -> GetItemLikeStmt: ...
+
+
+PyGetItemLikeStmt = TypeVar("PyGetItemLikeStmt", bound="GetItem")
+
+
+@dataclass(frozen=True, eq=False)
+class PyGetItemLike(GetItemLike[PyGetItemLikeStmt]):
+
+    def get_object(self, stmt: PyGetItemLikeStmt) -> ir.SSAValue:
+        return stmt.obj
+
+    def get_index(self, stmt: PyGetItemLikeStmt) -> ir.SSAValue:
+        return stmt.index
+
+    def new(
+        self, stmt_type: type[PyGetItemLikeStmt], obj: ir.SSAValue, index: ir.SSAValue
+    ) -> PyGetItemLikeStmt:
+        return stmt_type(obj=obj, index=index)
+
+
+# NOTE: in IR setindex is very different from getindex
+# taking Julia's semantics as reference here
+@statement
+class Subscript(ir.Statement):
+    pass
+
+
+@statement(dialect=dialect)
+class GetItem(Subscript):
+    name = "getitem"
+    traits = frozenset({ir.Pure(), PyGetItemLike(), ir.FromPythonCall()})
+    obj: ir.SSAValue = info.argument(print=False)
+    index: ir.SSAValue = info.argument(print=False)
+    result: ir.ResultValue = info.result(types.Any)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Subscript(
+        self, state: lowering.LoweringState, node: ast.Subscript
+    ) -> lowering.Result:
+        value = state.visit(node.value).expect_one()
+        slice = state.visit(node.slice).expect_one()
+        if isinstance(node.ctx, ast.Load):
+            stmt = GetItem(obj=value, index=slice)
+        else:
+            raise exceptions.DialectLoweringError(
+                f"unsupported subscript context {node.ctx}"
+            )
+        state.append_stmt(stmt)
+        return lowering.Result(stmt)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(GetItem)
+    def getindex(self, interp, frame: interp.Frame, stmt: GetItem):
+        return (frame.get(stmt.obj)[frame.get(stmt.index)],)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(GetItem)
+    def getitem(
+        self,
+        interp: TypeInference,
+        frame: interp.Frame[types.TypeAttribute],
+        stmt: GetItem,
+    ):
+        obj = frame.get(stmt.obj)
+        index: types.TypeAttribute = frame.get(stmt.index)
+        # TODO: replace this when we can multiple dispatch
+        if obj.is_subseteq(types.Tuple):
+            return self.getitem_tuple(interp, stmt, obj, index)
+        elif obj.is_subseteq(types.String):
+            return (types.String,)
+        else:
+            return (types.Any,)
+
+    def getitem_tuple(
+        self,
+        interp,
+        stmt: GetItem,
+        obj: types.TypeAttribute,
+        index: types.TypeAttribute,
+    ):
+        if isinstance(obj, types.Generic):
+            if index.is_subseteq(types.Int):
+                return self.getitem_tuple_index(interp, stmt, obj, index)
+            elif index.is_subseteq(types.Slice):
+                return self.getitem_tuple_slice(interp, stmt, obj, index)
+            else:
+                return (types.Bottom,)
+        elif isinstance(obj, types.PyClass):
+            return (types.Any,)
+        else:
+            return (types.Bottom,)
+
+    def getitem_tuple_index(
+        self,
+        interp: TypeInference,
+        stmt: GetItem,
+        obj: types.Generic,
+        index: types.TypeAttribute,
+    ):
+        if index_ := interp.maybe_const(stmt.index, int):
+            if obj.vararg and (index_ >= len(obj.vars) or -len(obj.vars) <= index_ < 0):
+                return (obj.vararg.typ,)
+            elif obj.vars and (
+                0 <= index_ < len(obj.vars) or -len(obj.vars) <= index_ < 0
+            ):
+                return (obj.vars[index_],)
+            else:
+                return (types.Bottom,)
+        else:
+            return (self.getitem_tuple_union(obj),)
+
+    def getitem_tuple_slice(
+        self,
+        interp: TypeInference,
+        stmt: GetItem,
+        obj: types.Generic,
+        index: types.TypeAttribute,
+    ):
+        if index_ := interp.maybe_const(stmt.index, slice):
+            if obj.vararg and index_.stop >= len(obj.vars):
+                return (
+                    types.Union(
+                        *obj.vars[slice(index_.start, len(obj.vars), index_.step)],
+                        obj.vararg.typ,
+                    ),
+                )
+            elif index_.stop is None or index_.stop < len(obj.vars):
+                return (
+                    types.Tuple.where(
+                        obj.vars[slice(index_.start, index_.stop, index_.step)]
+                    ),
+                )
+            else:  # out of bounds
+                return (types.Bottom,)
+        else:
+            return (types.Tuple[types.Vararg(self.getitem_tuple_union(obj))],)
+
+    def getitem_tuple_union(self, obj: types.Generic):
+        if obj.vararg:
+            return types.Union(*obj.vars, obj.vararg.typ)
+        else:
+            return types.Union(*obj.vars)
+
+
+@dialect.register(key="constprop")
+class ConstProp(interp.MethodTable):
+
+    @interp.impl(GetItem)
+    def getitem(
+        self,
+        _: const.Propagate,
+        frame: const.Frame,
+        stmt: GetItem,
+    ) -> interp.StatementResult[const.Result]:
+        obj = frame.get(stmt.obj)
+        index = frame.get(stmt.index)
+        if not isinstance(index, const.Value):
+            return (const.Unknown(),)
+
+        if isinstance(obj, const.PartialTuple):
+            obj = obj.data
+            if isinstance(index.data, int) and 0 <= index.data < len(obj):
+                return (obj[index.data],)
+            elif isinstance(index.data, slice):
+                start, stop, step = index.data.indices(len(obj))
+                return (const.PartialTuple(obj[start:stop:step]),)
+        return (const.Unknown(),)
+
+
+GetItemLikeStmt = TypeVar("GetItemLikeStmt", bound=ir.Statement)
+
+
+@dataclass(init=False)
+class RewriteGetItem(RewriteRule, Generic[GetItemLikeStmt]):
+    target_stmt_type: type[GetItemLikeStmt]
+    obj_type: types.TypeAttribute
+    getitem_like: GetItemLike[GetItemLikeStmt]
+
+    def __init__(self, stmt_type: type[GetItemLikeStmt], obj_type: types.TypeAttribute):
+        trait = stmt_type.get_trait(GetItemLike)
+        if trait is None:
+            raise ValueError(f"{stmt_type} does not have GetItemLike trait")
+
+        self.obj_type = obj_type
+        self.target_stmt_type = stmt_type
+        self.getitem_like = trait
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, GetItem):
+            return RewriteResult()
+
+        if not node.obj.type.is_subseteq(self.obj_type):
+            return RewriteResult()
+
+        node.replace_by(
+            self.getitem_like.new(self.target_stmt_type, node.obj, node.index)
+        )
+        return RewriteResult(has_done_something=True)

kirin/dialects/py/iterable.py

@@ -0,0 +1,90 @@
+"""This module provides access to Python iterables.
+
+This is used to lower Python loops into `cf` dialect.
+
+This module contains the common methods for the Python iterable:
+
+- The `Iter` statement class.
+- The `Next` statement class.
+- The lowering pass for the iterable.
+- The concrete implementation of the iterable.
+
+This dialect maps `iter()` and `next()` calls to the `Iter` and `Next` statements.
+"""
+
+from ast import Call
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("py.iterable")
+
+PyRangeIterType = types.PyClass(type(iter(range(0))))
+
+
+@statement(dialect=dialect)
+class Iter(ir.Statement):
+    """This is equivalent to `iter(value)` in Python."""
+
+    traits = frozenset({ir.Pure()})
+    value: ir.SSAValue = info.argument(types.Any)
+    iter: ir.ResultValue = info.result(types.Any)
+
+
+@statement(dialect=dialect)
+class Next(ir.Statement):
+    """This is equivalent to `next(iterable, None)` in Python."""
+
+    iter: ir.SSAValue = info.argument(types.Any)
+    value: ir.ResultValue = info.result(types.Any)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Iter)
+    def iter_(self, interp, frame: interp.Frame, stmt: Iter):
+        return (iter(frame.get(stmt.value)),)
+
+    @interp.impl(Next)
+    def next_(self, interp, frame: interp.Frame, stmt: Next):
+        return (next(frame.get(stmt.iter), None),)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Iter, types.PyClass(range))
+    def iter_(self, interp, frame: interp.Frame, stmt: Iter):
+        return (PyRangeIterType,)
+
+    @interp.impl(Next, PyRangeIterType)
+    def next_(self, interp, frame: interp.Frame, stmt: Next):
+        return (types.Int,)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Call_iter(
+        self, state: lowering.LoweringState, node: Call
+    ) -> lowering.Result:
+        if len(node.args) != 1:
+            raise DialectLoweringError("iter() takes exactly 1 argument")
+        return lowering.Result(
+            state.append_stmt(Iter(state.visit(node.args[0]).expect_one()))
+        )
+
+    def lower_Call_next(
+        self, state: lowering.LoweringState, node: Call
+    ) -> lowering.Result:
+        if len(node.args) == 2:
+            raise DialectLoweringError(
+                "next() does not throw StopIteration inside kernel"
+            )
+        if len(node.args) != 1:
+            raise DialectLoweringError("next() takes exactly 1 argument")
+        return lowering.Result(
+            state.append_stmt(Next(state.visit(node.args[0]).expect_one()))
+        )

kirin/dialects/py/len.py

@@ -0,0 +1,57 @@
+"""The `Len` dialect.
+
+This dialect maps the `len()` call to the `Len` statement:
+
+- The `Len` statement class.
+- The lowering pass for the `len()` call.
+- The concrete implementation of the `len()` call.
+"""
+
+import ast
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.analysis import const
+
+dialect = ir.Dialect("py.len")
+
+
+@statement(dialect=dialect)
+class Len(ir.Statement):
+    name = "len"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    value: ir.SSAValue = info.argument(types.Any)
+    result: ir.ResultValue = info.result(types.Int)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Len)
+    def len(self, interp, frame: interp.Frame, stmt: Len):
+        return (len(frame.get(stmt.value)),)
+
+
+@dialect.register(key="constprop")
+class ConstProp(interp.MethodTable):
+
+    @interp.impl(Len)
+    def len(self, interp, frame: interp.Frame, stmt: Len):
+        value = frame.get(stmt.value)
+        if isinstance(value, const.Value):
+            return (const.Value(len(value.data)),)
+        elif isinstance(value, const.PartialTuple):
+            return (const.Value(len(value.data)),)
+        else:
+            return (const.Result.top(),)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Call_len(
+        self, state: lowering.LoweringState, node: ast.Call
+    ) -> lowering.Result:
+        return lowering.Result(
+            state.append_stmt(Len(value=state.visit(node.args[0]).expect_one()))
+        )