split py dialect (#138)

this PR splits the python dialect into many smaller dialects allow
better reusibility in other DSLs. Most importantly we split an immutable
subset of list semantics from Python which is mostly used in various
dialects in QuEra.

Author: Roger-luo

src/kirin/decl/emit/from_python_call.py

@@ -28,11 +28,11 @@ class EmitFromPythonCall(BaseModifier):
     _AST_CONSTANT = "_kirin_ast_Constant"
 
     def __init__(self, cls: type, **kwargs: Unpack[StatementOptions]) -> None:
-        from kirin.dialects.py import data
+        from kirin.dialects.py.data import PyAttr
 
         super().__init__(cls, **kwargs)
         self.globals[self._KIRIN_RESULT] = Result
-        self.globals[self._KIRIN_PYATTR] = data.PyAttr
+        self.globals[self._KIRIN_PYATTR] = PyAttr
         self.globals[self._KIRIN_ERROR] = DialectLoweringError
         self.globals[self._AST_TUPLE] = ast.Tuple
         self.globals[self._AST_CONSTANT] = ast.Constant

src/kirin/decl/emit/property.py

@@ -111,7 +111,7 @@ def _emit_result_property(self, index: int, f: info.ResultField):
         return getter, setter
 
     def _emit_attribute_property(self, f: info.AttributeField):
-        from kirin.dialects.py import data
+        from kirin.dialects.py.data import PyAttr
 
         storage = "properties" if f.property else "attributes"
         attr = f"{self._self_name}.{storage}['{f.name}']"
@@ -131,7 +131,7 @@ def _emit_attribute_property(self, f: info.AttributeField):
                     f"raise AttributeError('attribute property {f.name} is read-only')"
                 ],
                 globals=self.globals,
-                locals={"_value_hint": data.PyAttr if f.pytype else f.annotation},
+                locals={"_value_hint": PyAttr if f.pytype else f.annotation},
                 return_type=None,
             )
         else:

src/kirin/dialects/cf/lower.py

@@ -4,7 +4,6 @@
 from kirin.lowering import Frame, Result, FromPythonAST, LoweringState
 from kirin.exceptions import DialectLoweringError
 from kirin.dialects.cf import stmts as cf
-from kirin.dialects.py import stmts
 from kirin.dialects.cf.dialect import dialect
 
 
@@ -19,12 +18,14 @@ def lower_Pass(self, state: LoweringState, node: ast.Pass) -> Result:
         return Result()
 
     def lower_Assert(self, state: LoweringState, node: ast.Assert) -> Result:
+        from kirin.dialects.py.constant import Constant
+
         cond = state.visit(node.test).expect_one()
         if node.msg:
             message = state.visit(node.msg).expect_one()
             state.append_stmt(cf.Assert(condition=cond, message=message))
         else:
-            message_stmt = state.append_stmt(stmts.Constant(""))
+            message_stmt = state.append_stmt(Constant(""))
             state.append_stmt(cf.Assert(condition=cond, message=message_stmt.result))
         return Result()
 

src/kirin/dialects/py/__init__.py

@@ -1,24 +1,39 @@
-from . import data as data, stmts as stmts
-
-# from kirin.dialects.pytypes import types as types
-# from kirin.dialects.pytypes.data import PyAttr as PyAttr
-# from kirin.dialects.pytypes.types.base import (
-#     PyBottomType as PyBottomType,
-#     PyClass as PyClass,
-#     PyGeneric as PyGeneric,
-#     PyTypeVar as PyTypeVar,
-#     PyUnion as PyUnion,
-#     PyVararg as PyVararg,
-# )
-# from kirin.dialects.pytypes.types.builtin import (
-#     Bool as Bool,
-#     Dict as Dict,
-#     Float as Float,
-#     FunctionType as FunctionType,
-#     Int as Int,
-#     List as List,
-#     NoneType as NoneType,
-#     String as String,
-#     Tuple as Tuple,
-#     TypeofFunctionType as TypeofFunctionType,
-# )
+from . import (
+    cmp as cmp,
+    len as len,
+    attr as attr,
+    data as data,
+    binop as binop,
+    ilist as ilist,
+    range as range,
+    slice as slice,
+    tuple as tuple,
+    unary as unary,
+    append as append,
+    assign as assign,
+    boolop as boolop,
+    builtin as builtin,
+    constant as constant,
+    indexing as indexing,
+)
+from .len import Len as Len
+from .attr import GetAttr as GetAttr
+from .data import PyAttr as PyAttr
+from .range import Range as Range
+from .slice import Slice as Slice
+from .append import Append as Append
+from .assign import Alias as Alias, SetItem as SetItem
+from .boolop import Or as Or, And as And
+from .builtin import Abs as Abs, Sum as Sum
+from .constant import Constant as Constant
+from .indexing import GetItem as GetItem, PyGetItemLike as PyGetItemLike
+from .cmp.stmts import *  # noqa: F403
+from .binop.stmts import *  # noqa: F403
+from .ilist.stmts import (
+    Map as Map,
+    Scan as Scan,
+    FoldL as FoldL,
+    FoldR as FoldR,
+    ForEach as ForEach,
+)
+from .unary.stmts import *  # noqa: F403

src/kirin/dialects/py/append.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from kirin import ir
+from kirin.decl import info, statement
+from kirin.interp import Frame, Interpreter, MethodTable, impl
+
+dialect = ir.Dialect("py.list")
+
+ElemT = ir.types.TypeVar("ElemT")
+
+
+@statement(dialect=dialect)
+class Append(ir.Statement):
+    lst: ir.SSAValue = info.argument(ir.types.List[ElemT])
+    value: ir.SSAValue = info.argument(ir.types.Any)
+
+
+@dialect.register
+class MutableListMethod(MethodTable):
+
+    @impl(Append)
+    def append(self, interp: Interpreter, frame: Frame[Any], stmt: Append):
+        lst: list = frame.get(stmt.lst)
+        value = frame.get(stmt.value)
+        lst.append(value)
+        return (lst,)

src/kirin/dialects/py/assign.py

@@ -0,0 +1,92 @@
+import ast
+
+from kirin import ir, interp, lowering, exceptions
+from kirin.decl import info, statement
+from kirin.print import Printer
+from kirin.dialects.py import data
+
+dialect = ir.Dialect("py.assign")
+
+T = ir.types.TypeVar("T")
+
+
+@statement(dialect=dialect)
+class Alias(ir.Statement):
+    name = "alias"
+    traits = frozenset({ir.Pure()})
+    value: ir.SSAValue = info.argument(T)
+    target: data.PyAttr[str] = info.attribute(property=True)
+    result: ir.ResultValue = info.result(T)
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        with printer.rich(style=printer.color.symbol):
+            printer.plain_print(self.target.data)
+
+        with printer.rich(style=printer.color.keyword):
+            printer.plain_print(" = ")
+
+        printer.print(self.value)
+
+
+@statement(dialect=dialect)
+class SetItem(ir.Statement):
+    name = "setitem"
+    obj: ir.SSAValue = info.argument(print=False)
+    value: ir.SSAValue = info.argument(print=False)
+    index: ir.SSAValue = info.argument(print=False)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Alias)
+    def alias(self, interp, frame: interp.Frame, stmt: Alias):
+        return (frame.get(stmt.value),)
+
+    @interp.impl(SetItem)
+    def setindex(self, interp, frame: interp.Frame, stmt: SetItem):
+        frame.get(stmt.obj)[frame.get(stmt.index)] = frame.get(stmt.value)
+        return (None,)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Assign(
+        self, state: lowering.LoweringState, node: ast.Assign
+    ) -> lowering.Result:
+        results: lowering.Result = state.visit(node.value)
+        assert len(node.targets) == len(
+            results
+        ), "number of targets and results do not match"
+
+        current_frame = state.current_frame
+        match node:
+            case ast.Assign(
+                targets=[ast.Name(lhs_name, ast.Store())], value=ast.Name(_, ast.Load())
+            ):
+                stmt = Alias(
+                    value=results[0], target=data.PyAttr(lhs_name)
+                )  # NOTE: this is guaranteed to be one result
+                stmt.result.name = lhs_name
+                current_frame.defs[lhs_name] = state.append_stmt(stmt).result
+            case _:
+                for target, value in zip(node.targets, results.values):
+                    match target:
+                        # NOTE: if the name exists new ssa value will be
+                        # used in the future to shadow the old one
+                        case ast.Name(name, ast.Store()):
+                            value.name = name
+                            current_frame.defs[name] = value
+                        case ast.Subscript(obj, slice):
+                            obj = state.visit(obj).expect_one()
+                            slice = state.visit(slice).expect_one()
+                            stmt = SetItem(obj=obj, index=slice, value=value)
+                            state.append_stmt(stmt)
+                        case _:
+                            raise exceptions.DialectLoweringError(
+                                f"unsupported target {target}"
+                            )
+        return lowering.Result()  # python assign does not have value

src/kirin/dialects/py/attr.py

@@ -0,0 +1,38 @@
+import ast
+
+from kirin import ir, interp, lowering, exceptions
+from kirin.decl import info, statement
+
+dialect = ir.Dialect("py.attr")
+
+
+@statement(dialect=dialect)
+class GetAttr(ir.Statement):
+    name = "getattr"
+    obj: ir.SSAValue = info.argument(print=False)
+    attrname: str = info.attribute(property=True)
+    result: ir.ResultValue = info.result()
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(GetAttr)
+    def getattr(self, interp: interp.Interpreter, frame: interp.Frame, stmt: GetAttr):
+        return getattr(frame.get(stmt.obj), stmt.attrname)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Attribute(
+        self, state: lowering.LoweringState, node: ast.Attribute
+    ) -> lowering.Result:
+        if not isinstance(node.ctx, ast.Load):
+            raise exceptions.DialectLoweringError(
+                f"unsupported attribute context {node.ctx}"
+            )
+        value = state.visit(node.value).expect_one()
+        stmt = GetAttr(obj=value, attrname=node.attr)
+        state.append_stmt(stmt)
+        return lowering.Result(stmt)

src/kirin/dialects/py/binop/__init__.py

@@ -0,0 +1,8 @@
+from . import (
+    julia as julia,
+    interp as interp,
+    lowering as lowering,
+    typeinfer as typeinfer,
+)
+from .stmts import *  # noqa: F403
+from ._dialect import dialect as dialect

src/kirin/dialects/py/binop/_dialect.py

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

src/kirin/dialects/py/binop/interp.py

@@ -0,0 +1,60 @@
+from kirin import interp
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class PyMethodTable(interp.MethodTable):
+
+    @interp.impl(stmts.Add)
+    def add(self, interp, frame: interp.Frame, stmt: stmts.Add):
+        return (frame.get(stmt.lhs) + frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Sub)
+    def sub(self, interp, frame: interp.Frame, stmt: stmts.Sub):
+        return (frame.get(stmt.lhs) - frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Mult)
+    def mult(self, interp, frame: interp.Frame, stmt: stmts.Mult):
+        return (frame.get(stmt.lhs) * frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Div)
+    def div(self, interp, frame: interp.Frame, stmt: stmts.Div):
+        return (frame.get(stmt.lhs) / frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Mod)
+    def mod(self, interp, frame: interp.Frame, stmt: stmts.Mod):
+        return (frame.get(stmt.lhs) % frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.BitAnd)
+    def bit_and(self, interp, frame: interp.Frame, stmt: stmts.BitAnd):
+        return (frame.get(stmt.lhs) & frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.BitOr)
+    def bit_or(self, interp, frame: interp.Frame, stmt: stmts.BitOr):
+        return (frame.get(stmt.lhs) | frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.BitXor)
+    def bit_xor(self, interp, frame: interp.Frame, stmt: stmts.BitXor):
+        return (frame.get(stmt.lhs) ^ frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.LShift)
+    def lshift(self, interp, frame: interp.Frame, stmt: stmts.LShift):
+        return (frame.get(stmt.lhs) << frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.RShift)
+    def rshift(self, interp, frame: interp.Frame, stmt: stmts.RShift):
+        return (frame.get(stmt.lhs) >> frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.FloorDiv)
+    def floor_div(self, interp, frame: interp.Frame, stmt: stmts.FloorDiv):
+        return (frame.get(stmt.lhs) // frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Pow)
+    def pow(self, interp, frame: interp.Frame, stmt: stmts.Pow):
+        return (frame.get(stmt.lhs) ** frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.MatMult)
+    def mat_mult(self, interp, frame: interp.Frame, stmt: stmts.MatMult):
+        return (frame.get(stmt.lhs) @ frame.get(stmt.rhs),)