kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/py/list/__init__.py

@@ -0,0 +1,15 @@
+"""The list dialect for Python.
+
+This module contains the dialect for list semantics in Python, including:
+
+- The `New` and `Append` statement classes.
+- The lowering pass for list operations.
+- The concrete implementation of list operations.
+- The type inference implementation of list operations.
+
+This dialect maps `list()`, `ast.List` and `append()` calls to the `New` and `Append` statements.
+"""
+
+from . import interp as interp, lowering as lowering, typeinfer as typeinfer
+from .stmts import New as New, Append as Append
+from ._dialect import dialect as dialect

kirin/dialects/py/list/_dialect.py

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

kirin/dialects/py/list/interp.py

@@ -0,0 +1,21 @@
+from kirin import types, interp
+from kirin.dialects.py.binop import Add
+
+from .stmts import New, Append
+from ._dialect import dialect
+
+
+@dialect.register
+class ListMethods(interp.MethodTable):
+
+    @interp.impl(New)
+    def new(self, interp, frame: interp.Frame, stmt: New):
+        return (list(frame.get_values(stmt.values)),)
+
+    @interp.impl(Add, types.PyClass(list), types.PyClass(list))
+    def add(self, interp, frame: interp.Frame, stmt: Add):
+        return (frame.get(stmt.lhs) + frame.get(stmt.rhs),)
+
+    @interp.impl(Append)
+    def append(self, interp, frame: interp.Frame, stmt: Append):
+        return (frame.get(stmt.list_).append(frame.get(stmt.value)),)

kirin/dialects/py/list/lowering.py

@@ -0,0 +1,25 @@
+import ast
+
+from kirin import types
+from kirin.lowering import Result, FromPythonAST, LoweringState
+
+from .stmts import New
+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 = New(values=tuple(elts))
+        state.append_stmt(stmt)
+        return Result(stmt)

kirin/dialects/py/list/stmts.py

@@ -0,0 +1,22 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from ._dialect import dialect
+
+T = types.TypeVar("T")
+
+
+@statement(dialect=dialect)
+class New(ir.Statement):
+    name = "list"
+    traits = frozenset({ir.FromPythonCall()})
+    values: tuple[ir.SSAValue, ...] = info.argument(T)
+    result: ir.ResultValue = info.result(types.List[T])
+
+
+@statement(dialect=dialect)
+class Append(ir.Statement):
+    name = "append"
+    traits = frozenset({ir.FromPythonCall()})
+    list_: ir.SSAValue = info.argument(types.List[T])
+    value: ir.SSAValue = info.argument(T)

kirin/dialects/py/list/typeinfer.py

@@ -0,0 +1,54 @@
+from kirin import types, interp
+from kirin.dialects.eltype import ElType
+from kirin.dialects.py.binop import Add
+from kirin.dialects.py.indexing import GetItem
+
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(ElType, types.PyClass(list))
+    def eltype_list(self, interp, frame: interp.Frame, stmt: ElType):
+        list_type = frame.get(stmt.container)
+        if isinstance(list_type, types.Generic):
+            return (list_type.vars[0],)
+        else:
+            return (types.Any,)
+
+    @interp.impl(Add, types.PyClass(list), types.PyClass(list))
+    def add(self, interp, frame: interp.Frame, stmt: Add):
+        lhs_type = frame.get(stmt.lhs)
+        rhs_type = frame.get(stmt.rhs)
+        if isinstance(lhs_type, types.Generic):
+            lhs_elem_type = lhs_type.vars[0]
+        else:
+            lhs_elem_type = types.Any
+
+        if isinstance(rhs_type, types.Generic):
+            rhs_elem_type = rhs_type.vars[0]
+        else:
+            rhs_elem_type = types.Any
+
+        return (types.List[lhs_elem_type.join(rhs_elem_type)],)
+
+    @interp.impl(GetItem, types.PyClass(list), types.Int)
+    def getitem_list_int(
+        self, interp, frame: interp.Frame[types.TypeAttribute], stmt: GetItem
+    ):
+        obj_type = frame.get(stmt.obj)
+        if isinstance(obj_type, types.Generic):
+            return (obj_type.vars[0],)
+        else:
+            return (types.Any,)
+
+    @interp.impl(GetItem, types.PyClass(list), types.PyClass(slice))
+    def getitem_list_slice(
+        self, interp, frame: interp.Frame[types.TypeAttribute], stmt: GetItem
+    ):
+        obj_type = frame.get(stmt.obj)
+        if isinstance(obj_type, types.Generic):
+            return (types.List[obj_type.vars[0]],)
+        else:
+            return (types.Any,)

kirin/dialects/py/range.py

@@ -0,0 +1,76 @@
+"""The range dialect for Python.
+
+This dialect models the builtin `range()` function in Python.
+
+The dialect includes:
+- The `Range` statement class.
+- The lowering pass for the `range()` function.
+
+This dialect does not include a concrete implementation or type inference
+for the `range()` function. One needs to use other dialect for the concrete
+implementation and type inference, e.g., `ilist` dialect.
+"""
+
+import ast
+from dataclasses import dataclass
+
+from kirin import ir, types, interp, lowering, exceptions
+from kirin.decl import info, statement
+from kirin.dialects import eltype
+
+dialect = ir.Dialect("py.range")
+
+
+@dataclass(frozen=True)
+class RangeLowering(ir.FromPythonCall["Range"]):
+
+    def lower(
+        self, stmt: type["Range"], state: lowering.LoweringState, node: ast.Call
+    ) -> lowering.Result:
+        return _lower_range(state, node)
+
+
+@statement(dialect=dialect)
+class Range(ir.Statement):
+    name = "range"
+    traits = frozenset({ir.Pure(), RangeLowering()})
+    start: ir.SSAValue = info.argument(types.Int)
+    stop: ir.SSAValue = info.argument(types.Int)
+    step: ir.SSAValue = info.argument(types.Int)
+    result: ir.ResultValue = info.result(types.PyClass(range))
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Call_range(
+        self, state: lowering.LoweringState, node: ast.Call
+    ) -> lowering.Result:
+        return _lower_range(state, node)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(eltype.ElType, types.PyClass(range))
+    def eltype_range(self, interp_, frame: interp.Frame, stmt: eltype.ElType):
+        return (types.Int,)
+
+
+def _lower_range(state: lowering.LoweringState, node: ast.Call) -> lowering.Result:
+    if len(node.args) == 1:
+        start = state.visit(ast.Constant(0)).expect_one()
+        stop = state.visit(node.args[0]).expect_one()
+        step = state.visit(ast.Constant(1)).expect_one()
+    elif len(node.args) == 2:
+        start = state.visit(node.args[0]).expect_one()
+        stop = state.visit(node.args[1]).expect_one()
+        step = state.visit(ast.Constant(1)).expect_one()
+    elif len(node.args) == 3:
+        start = state.visit(node.args[0]).expect_one()
+        stop = state.visit(node.args[1]).expect_one()
+        step = state.visit(node.args[2]).expect_one()
+    else:
+        raise exceptions.DialectLoweringError("range() takes 1-3 arguments")
+
+    return lowering.Result(state.append_stmt(Range(start, stop, step)))

kirin/dialects/py/slice.py

@@ -0,0 +1,120 @@
+"""The slice dialect for Python.
+
+This dialect provides a `Slice` statement that represents a slice object in Python:
+
+- The `Slice` statement class.
+- The lowering pass for the `slice` call.
+- The concrete implementation of the `slice` call.
+- The type inference implementation of the `slice` call.
+"""
+
+import ast
+from dataclasses import dataclass
+
+from kirin import ir, types, interp, lowering, exceptions
+from kirin.decl import info, statement
+from kirin.dialects.py.constant import Constant
+
+dialect = ir.Dialect("py.slice")
+
+
+@dataclass(frozen=True)
+class SliceLowering(ir.FromPythonCall["Slice"]):
+
+    def lower(
+        self, stmt: type["Slice"], state: lowering.LoweringState, node: ast.Call
+    ) -> lowering.Result:
+        return _lower_slice(state, node)
+
+
+T = types.TypeVar("T")
+
+
+@statement(dialect=dialect, init=False)
+class Slice(ir.Statement):
+    name = "slice"
+    traits = frozenset({ir.Pure(), SliceLowering()})
+    start: ir.SSAValue = info.argument(T | types.NoneType)
+    stop: ir.SSAValue = info.argument(T | types.NoneType)
+    step: ir.SSAValue = info.argument(T | types.NoneType)
+    result: ir.ResultValue = info.result(types.Slice[T])
+
+    def __init__(
+        self, start: ir.SSAValue, stop: ir.SSAValue, step: ir.SSAValue
+    ) -> None:
+        if not (
+            isinstance(stop.type, types.TypeAttribute)
+            and isinstance(start.type, types.TypeAttribute)
+        ):
+            result_type = types.Bottom
+        elif start.type.is_subseteq(types.NoneType):
+            if stop.type.is_subseteq(types.NoneType):
+                result_type = types.Bottom
+            else:
+                result_type = types.Slice[stop.type]
+        else:
+            result_type = types.Slice[start.type]
+
+        super().__init__(
+            args=(start, stop, step),
+            result_types=[result_type],
+            args_slice={"start": 0, "stop": 1, "step": 2},
+        )
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Slice)
+    def _slice(self, interp, frame: interp.Frame, stmt: Slice):
+        start, stop, step = frame.get_values(stmt.args)
+        if start is None and step is None:
+            return (slice(stop),)
+        elif step is None:
+            return (slice(start, stop),)
+        else:
+            return (slice(start, stop, step),)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Slice(
+        self, state: lowering.LoweringState, node: ast.Slice
+    ) -> lowering.Result:
+        def value_or_none(expr: ast.expr | None) -> ir.SSAValue:
+            if expr is not None:
+                return state.visit(expr).expect_one()
+            else:
+                return state.append_stmt(Constant(None)).result
+
+        lower = value_or_none(node.lower)
+        upper = value_or_none(node.upper)
+        step = value_or_none(node.step)
+        return lowering.Result(
+            state.append_stmt(Slice(start=lower, stop=upper, step=step))
+        )
+
+    def lower_Call_slice(
+        self, state: lowering.LoweringState, node: ast.Call
+    ) -> lowering.Result:
+        return _lower_slice(state, node)
+
+
+def _lower_slice(state: lowering.LoweringState, node: ast.Call) -> lowering.Result:
+    if len(node.args) == 1:
+        start = state.visit(ast.Constant(None)).expect_one()
+        stop = state.visit(node.args[0]).expect_one()
+        step = state.visit(ast.Constant(None)).expect_one()
+    elif len(node.args) == 2:
+        start = state.visit(node.args[0]).expect_one()
+        stop = state.visit(node.args[1]).expect_one()
+        step = state.visit(ast.Constant(None)).expect_one()
+    elif len(node.args) == 3:
+        start = state.visit(node.args[0]).expect_one()
+        stop = state.visit(node.args[1]).expect_one()
+        step = state.visit(node.args[2]).expect_one()
+    else:
+        raise exceptions.DialectLoweringError("slice() takes 1-3 arguments")
+
+    return lowering.Result(state.append_stmt(Slice(start, stop, step)))

kirin/dialects/py/tuple.py

@@ -0,0 +1,109 @@
+"""The tuple dialect for Python.
+
+This dialect provides a way to work with Python tuples in the IR, including:
+
+- The `New` statement class.
+- The lowering pass for the tuple statement.
+- The concrete implementation of the tuple statement.
+- The type inference implementation of the tuple addition with `py.binop.Add`.
+- The constant propagation implementation of the tuple statement.
+- The Julia emitter for the tuple statement.
+
+This dialect maps `ast.Tuple` nodes to the `New` statement.
+"""
+
+import ast
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.analysis import const
+from kirin.emit.julia import EmitJulia, EmitStrFrame
+from kirin.dialects.eltype import ElType
+from kirin.dialects.py.binop import Add
+
+dialect = ir.Dialect("py.tuple")
+
+
+@statement(dialect=dialect)
+class New(ir.Statement):
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    result: ir.ResultValue = info.result()
+
+    def __init__(self, values: tuple[ir.SSAValue, ...]) -> None:
+        result_type = types.Generic(tuple, *tuple(value.type for value in values))
+        super().__init__(
+            args=values,
+            result_types=[
+                result_type,
+            ],
+        )
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(New)
+    def new(self, interp: interp.Interpreter, frame: interp.Frame, stmt: New):
+        return (frame.get_values(stmt.args),)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(ElType, types.PyClass(tuple))
+    def eltype_tuple(self, interp, frame: interp.Frame, stmt: ElType):
+        tuple_type = frame.get(stmt.container)
+        if isinstance(tuple_type, types.Generic):
+            ret = tuple_type.vars[0]
+            for var in tuple_type.vars[1:]:
+                ret = ret.join(var)
+            return (ret,)
+        else:
+            return (types.Any,)
+
+    @interp.impl(Add, types.PyClass(tuple), types.PyClass(tuple))
+    def add(self, interp, frame: interp.Frame[types.TypeAttribute], stmt):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        if isinstance(lhs, types.Generic) and isinstance(rhs, types.Generic):
+            return (types.Generic(tuple, *(lhs.vars + rhs.vars)),)
+        else:
+            return (types.PyClass(tuple),)  # no type param, so unknown
+
+
+@dialect.register(key="constprop")
+class ConstPropTable(interp.MethodTable):
+
+    @interp.impl(New)
+    def new_tuple(
+        self,
+        _: const.Propagate,
+        frame: const.Frame,
+        stmt: New,
+    ) -> interp.StatementResult[const.Result]:
+        return (const.PartialTuple(tuple(x for x in frame.get_values(stmt.args))),)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Tuple(
+        self, state: lowering.LoweringState, node: ast.Tuple
+    ) -> lowering.Result:
+        return lowering.Result(
+            state.append_stmt(
+                stmt=New(tuple(state.visit(elem).expect_one() for elem in node.elts))
+            )
+        )
+
+
+@dialect.register(key="emit.julia")
+class JuliaTable(interp.MethodTable):
+
+    @interp.impl(New)
+    def emit_NewTuple(self, emit: EmitJulia, frame: EmitStrFrame, stmt: New):
+        return (
+            emit.write_assign(
+                frame, stmt.result, "(" + ", ".join(frame.get_values(stmt.args)) + ")"
+            ),
+        )

kirin/dialects/py/unary/__init__.py

@@ -0,0 +1,24 @@
+"""The unary dialect for Python.
+
+This module contains the dialect for unary semantics in Python, including:
+
+- The `UnaryOp` base class for unary operations.
+- The `UAdd`, `USub`, `Not`, and `Invert` statement classes.
+- The lowering pass for unary operations.
+- The concrete implementation of unary operations.
+- The type inference implementation of unary operations.
+- The constant propagation implementation of unary operations.
+- The Julia emitter for unary operations.
+
+This dialect maps `ast.UnaryOp` nodes to the `UAdd`, `USub`, `Not`, and `Invert` statements.
+"""
+
+from . import (
+    julia as julia,
+    interp as interp,
+    lowering as lowering,
+    constprop as constprop,
+    typeinfer as typeinfer,
+)
+from .stmts import *  # noqa: F403
+from ._dialect import dialect as dialect

kirin/dialects/py/unary/_dialect.py

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

kirin/dialects/py/unary/constprop.py

@@ -0,0 +1,20 @@
+from kirin import interp
+from kirin.analysis import const
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="constprop")
+class ConstProp(interp.MethodTable):
+
+    @interp.impl(stmts.Not)
+    def not_(
+        self, _: const.Propagate, frame: const.Frame, stmt: stmts.Not
+    ) -> interp.StatementResult[const.Result]:
+        hint = frame.get(stmt.value)
+        if isinstance(hint, (const.PartialTuple, const.Value)):
+            ret = const.Value(not hint.data)
+        else:
+            ret = const.Unknown()
+        return (ret,)

kirin/dialects/py/unary/interp.py

@@ -0,0 +1,24 @@
+from kirin import interp
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(stmts.UAdd)
+    def uadd(self, interp, frame: interp.Frame, stmt: stmts.UAdd):
+        return (+frame.get(stmt.value),)
+
+    @interp.impl(stmts.USub)
+    def usub(self, interp, frame: interp.Frame, stmt: stmts.USub):
+        return (-frame.get(stmt.value),)
+
+    @interp.impl(stmts.Not)
+    def not_(self, interp, frame: interp.Frame, stmt: stmts.Not):
+        return (not frame.get(stmt.value),)
+
+    @interp.impl(stmts.Invert)
+    def invert(self, interp, frame: interp.Frame, stmt: stmts.Invert):
+        return (~frame.get(stmt.value),)

kirin/dialects/py/unary/julia.py

@@ -0,0 +1,21 @@
+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.Not)
+    def emit_Not(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Not):
+        return (emit.write_assign(frame, stmt.result, f"!{frame.get(stmt.value)}"),)
+
+    @interp.impl(stmts.USub)
+    def emit_USub(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.USub):
+        return (emit.write_assign(frame, stmt.result, f"-{frame.get(stmt.value)}"),)
+
+    @interp.impl(stmts.UAdd)
+    def emit_UAdd(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.UAdd):
+        return (emit.write_assign(frame, stmt.result, f"+{frame.get(stmt.value)}"),)

kirin/dialects/py/unary/lowering.py

@@ -0,0 +1,22 @@
+import ast
+
+from kirin import lowering, exceptions
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_UnaryOp(
+        self, state: lowering.LoweringState, node: ast.UnaryOp
+    ) -> lowering.Result:
+        if op := getattr(stmts, node.op.__class__.__name__, None):
+            return lowering.Result(
+                state.append_stmt(op(state.visit(node.operand).expect_one()))
+            )
+        else:
+            raise exceptions.DialectLoweringError(
+                f"unsupported unary operator {node.op}"
+            )

kirin/dialects/py/unary/stmts.py

@@ -0,0 +1,33 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from ._dialect import dialect
+
+T = types.TypeVar("T")
+
+
+@statement
+class UnaryOp(ir.Statement):
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    value: ir.SSAValue = info.argument(T, print=False)
+    result: ir.ResultValue = info.result(T)
+
+
+@statement(dialect=dialect)
+class UAdd(UnaryOp):
+    name = "uadd"
+
+
+@statement(dialect=dialect)
+class USub(UnaryOp):
+    name = "usub"
+
+
+@statement(dialect=dialect)
+class Not(UnaryOp):
+    name = "not"
+
+
+@statement(dialect=dialect)
+class Invert(UnaryOp):
+    name = "invert"

kirin/dialects/py/unary/typeinfer.py

@@ -0,0 +1,23 @@
+from kirin import types, interp
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(stmts.UAdd)
+    @interp.impl(stmts.USub)
+    def uadd(
+        self, interp, frame: interp.Frame[types.TypeAttribute], stmt: stmts.UnaryOp
+    ):
+        return (frame.get(stmt.value),)
+
+    @interp.impl(stmts.Not)
+    def not_(self, interp, frame, stmt: stmts.Not):
+        return (types.Bool,)
+
+    @interp.impl(stmts.Invert, types.Int)
+    def invert(self, interp, frame, stmt):
+        return (types.Int,)