kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/py/unpack.py

@@ -0,0 +1,90 @@
+"""The unpack dialect for Python.
+
+This module contains the dialect for the Python unpack semantics, including:
+
+- The `Unpack` statement class.
+- The lowering pass for the unpack statement.
+- The concrete implementation of the unpack statement.
+- The type inference implementation of the unpack statement.
+- A helper function `unpacking` for unpacking Python AST nodes during lowering.
+"""
+
+import ast
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.print import Printer
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("py.unpack")
+
+
+@statement(dialect=dialect, init=False)
+class Unpack(ir.Statement):
+    value: ir.SSAValue = info.argument(types.Any)
+    names: tuple[str | None, ...] = info.attribute()
+
+    def __init__(self, value: ir.SSAValue, names: tuple[str | None, ...]):
+        result_types = [types.Any] * len(names)
+        super().__init__(
+            args=(value,),
+            result_types=result_types,
+            args_slice={"value": 0},
+            attributes={"names": ir.PyAttr(names)},
+        )
+        for result, name in zip(self.results, names):
+            result.name = name
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.print(self.value)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Unpack)
+    def unpack(self, interp: interp.Interpreter, frame: interp.Frame, stmt: Unpack):
+        return tuple(frame.get(stmt.value))
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Unpack)
+    def unpack(self, interp, frame: interp.Frame[types.TypeAttribute], stmt: Unpack):
+        value = frame.get(stmt.value)
+        if isinstance(value, types.Generic) and value.is_subseteq(types.Tuple):
+            if value.vararg:
+                rest = tuple(value.vararg.typ for _ in stmt.names[len(value.vars) :])
+                return tuple(value.vars) + rest
+            else:
+                return value.vars
+        # TODO: support unpacking other types
+        return tuple(types.Any for _ in stmt.names)
+
+
+def unpacking(state: lowering.LoweringState, node: ast.expr, value: ir.SSAValue):
+    if isinstance(node, ast.Name):
+        state.current_frame.defs[node.id] = value
+        value.name = node.id
+        return
+    elif not isinstance(node, ast.Tuple):
+        raise DialectLoweringError(f"unsupported unpack node {node}")
+
+    names: list[str | None] = []
+    continue_unpack: list[int] = []
+    for idx, item in enumerate(node.elts):
+        if isinstance(item, ast.Name):
+            names.append(item.id)
+        else:
+            names.append(None)
+            continue_unpack.append(idx)
+    stmt = state.append_stmt(Unpack(value, tuple(names)))
+    for name, result in zip(names, stmt.results):
+        if name is not None:
+            state.current_frame.defs[name] = result
+
+    for idx in continue_unpack:
+        unpacking(state, node.elts[idx], stmt.results[idx])

kirin/dialects/scf/__init__.py

@@ -0,0 +1,23 @@
+"""A Python-like structural Control Flow dialect.
+
+This dialect provides constructs for expressing control flow in a structured
+manner. The dialect provides constructs for expressing loops and conditionals.
+Unlike MLIR SCF dialect, this dialect does not restrict the control flow to
+statically analyzable forms. This dialect is designed to be compatible with
+Python native control flow constructs.
+
+This dialect depends on the following dialects:
+- `eltype`: for obtaining the element type of a value.
+"""
+
+from . import (
+    trim as trim,
+    absint as absint,
+    interp as interp,
+    unroll as unroll,
+    lowering as lowering,
+    constprop as constprop,
+    typeinfer as typeinfer,
+)
+from .stmts import For as For, Yield as Yield, IfElse as IfElse
+from ._dialect import dialect as dialect

kirin/dialects/scf/_dialect.py

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

kirin/dialects/scf/absint.py

@@ -0,0 +1,64 @@
+from kirin import ir, interp
+from kirin.analysis import const
+from kirin.dialects import func
+
+from .stmts import Yield, IfElse
+from ._dialect import dialect
+
+
+@dialect.register(key="absint")
+class Methods(interp.MethodTable):
+
+    @interp.impl(Yield)
+    def yield_stmt(
+        self,
+        interp_: interp.AbstractInterpreter,
+        frame: interp.AbstractFrame,
+        stmt: Yield,
+    ):
+        return interp.YieldValue(frame.get_values(stmt.values))
+
+    @interp.impl(IfElse)
+    def if_else(
+        self,
+        interp_: interp.AbstractInterpreter,
+        frame: interp.AbstractFrame,
+        stmt: IfElse,
+    ):
+        if isinstance(hint := stmt.cond.hints.get("const"), const.Value):
+            if hint.data:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+            else:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+        then_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+        else_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+
+        match (then_results, else_results):
+            case (interp.ReturnValue(then_value), interp.ReturnValue(else_value)):
+                return interp.ReturnValue(then_value.join(else_value))
+            case (interp.ReturnValue(then_value), _):
+                return then_results
+            case (_, interp.ReturnValue(else_value)):
+                return else_results
+            case _:
+                return interp_.join_results(then_results, else_results)
+
+    def _infer_if_else_cond(
+        self,
+        interp_: interp.AbstractInterpreter,
+        frame: interp.AbstractFrame,
+        stmt: IfElse,
+        body: ir.Region,
+    ):
+        body_block = body.blocks[0]
+        body_term = body_block.last_stmt
+        if isinstance(body_term, func.Return):
+            frame.worklist.append(interp.Successor(body_block, frame.get(stmt.cond)))
+            return
+
+        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+            body_frame.entries.update(frame.entries)
+            body_frame.set(body_block.args[0], frame.get(stmt.cond))
+            ret = interp_.run_ssacfg_region(body_frame, body)
+            frame.entries.update(body_frame.entries)
+            return ret

kirin/dialects/scf/constprop.py

@@ -0,0 +1,140 @@
+from collections.abc import Iterable
+
+from kirin import ir, interp
+from kirin.analysis import const
+from kirin.dialects import func
+
+from .stmts import For, Yield, IfElse
+from ._dialect import dialect
+
+# NOTE: unlike concrete interpreter, we need to use a new frame
+# for each iteration because otherwise join two constant values
+# will result in bottom (error) element.
+
+
+@dialect.register(key="constprop")
+class DialectConstProp(interp.MethodTable):
+
+    @interp.impl(Yield)
+    def yield_stmt(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: Yield,
+    ):
+        return interp.YieldValue(frame.get_values(stmt.values))
+
+    @interp.impl(IfElse)
+    def if_else(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: IfElse,
+    ):
+        cond = frame.get(stmt.cond)
+        if isinstance(cond, const.Value):
+            if cond.data:
+                body = stmt.then_body
+            else:
+                body = stmt.else_body
+            body_frame, ret = self._prop_const_cond_ifelse(
+                interp_, frame, stmt, cond, body
+            )
+            frame.entries.update(body_frame.entries)
+            if not body_frame.frame_is_not_pure and not isinstance(
+                body.blocks[0].last_stmt, func.Return
+            ):
+                frame.should_be_pure.add(stmt)
+            return ret
+        else:
+            then_frame, then_results = self._prop_const_cond_ifelse(
+                interp_, frame, stmt, const.Value(True), stmt.then_body
+            )
+            else_frame, else_results = self._prop_const_cond_ifelse(
+                interp_, frame, stmt, const.Value(False), stmt.else_body
+            )
+            # NOTE: then_frame and else_frame do not change
+            # parent frame variables value except cond
+            frame.entries.update(then_frame.entries)
+            frame.entries.update(else_frame.entries)
+            # TODO: pick the non-return value
+            if isinstance(then_results, interp.ReturnValue) and isinstance(
+                else_results, interp.ReturnValue
+            ):
+                return interp.ReturnValue(then_results.value.join(else_results.value))
+            elif isinstance(then_results, interp.ReturnValue):
+                ret = else_results
+            elif isinstance(else_results, interp.ReturnValue):
+                ret = then_results
+            else:
+                if not (
+                    then_frame.frame_is_not_pure is True
+                    or else_frame.frame_is_not_pure is True
+                ):
+                    frame.should_be_pure.add(stmt)
+                ret = interp_.join_results(then_results, else_results)
+        return ret
+
+    def _prop_const_cond_ifelse(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: IfElse,
+        cond: const.Value,
+        body: ir.Region,
+    ):
+        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+            body_frame.entries.update(frame.entries)
+            body_frame.set(body.blocks[0].args[0], cond)
+            results = interp_.run_ssacfg_region(body_frame, body)
+        return body_frame, results
+
+    @interp.impl(For)
+    def for_loop(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: For,
+    ):
+        iterable = frame.get(stmt.iterable)
+        if isinstance(iterable, const.Value):
+            return self._prop_const_iterable_forloop(interp_, frame, stmt, iterable)
+        else:  # TODO: support other iteration
+            return tuple(interp_.lattice.top() for _ in stmt.results)
+
+    def _prop_const_iterable_forloop(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: For,
+        iterable: const.Value,
+    ):
+        frame_is_not_pure = False
+        if not isinstance(iterable.data, Iterable):
+            raise interp.InterpreterError(
+                f"Expected iterable, got {type(iterable.data)}"
+            )
+
+        loop_vars = frame.get_values(stmt.initializers)
+        body_block = stmt.body.blocks[0]
+        block_args = body_block.args
+
+        for value in iterable.data:
+            with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+                body_frame.entries.update(frame.entries)
+                body_frame.set_values(
+                    block_args,
+                    (const.Value(value),) + loop_vars,
+                )
+                loop_vars = interp_.run_ssacfg_region(body_frame, stmt.body)
+
+            if body_frame.frame_is_not_pure:
+                frame_is_not_pure = True
+            if loop_vars is None:
+                loop_vars = ()
+            elif isinstance(loop_vars, interp.ReturnValue):
+                return loop_vars
+
+        if not frame_is_not_pure:
+            frame.should_be_pure.add(stmt)
+        return loop_vars

kirin/dialects/scf/interp.py

@@ -0,0 +1,35 @@
+from kirin import interp
+
+from .stmts import For, Yield, IfElse
+from ._dialect import dialect
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Yield)
+    def yield_stmt(self, interp_: interp.Interpreter, frame: interp.Frame, stmt: Yield):
+        return interp.YieldValue(frame.get_values(stmt.values))
+
+    @interp.impl(IfElse)
+    def if_else(self, interp_: interp.Interpreter, frame: interp.Frame, stmt: IfElse):
+        cond = frame.get(stmt.cond)
+        if cond:
+            body = stmt.then_body
+        else:
+            body = stmt.else_body
+        return interp_.run_ssacfg_region(frame, body)
+
+    @interp.impl(For)
+    def for_loop(self, interpreter: interp.Interpreter, frame: interp.Frame, stmt: For):
+        iterable = frame.get(stmt.iterable)
+        loop_vars = frame.get_values(stmt.initializers)
+        block_args = stmt.body.blocks[0].args
+        for value in iterable:
+            frame.set_values(block_args, (value,) + loop_vars)
+            loop_vars = interpreter.run_ssacfg_region(frame, stmt.body)
+            if isinstance(loop_vars, interp.ReturnValue):
+                return loop_vars
+            elif loop_vars is None:
+                loop_vars = ()
+        return loop_vars

kirin/dialects/scf/lowering.py

@@ -0,0 +1,123 @@
+import ast
+
+from kirin import ir, types, lowering
+from kirin.exceptions import DialectLoweringError
+from kirin.dialects.py.unpack import unpacking
+
+from .stmts import For, Yield, IfElse
+from ._dialect import dialect
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_If(self, state: lowering.LoweringState, node: ast.If) -> lowering.Result:
+        cond = state.visit(node.test).expect_one()
+        frame = state.current_frame
+        body_frame = lowering.Frame.from_stmts(node.body, state, globals=frame.globals)
+        then_cond = body_frame.curr_block.args.append_from(types.Bool, cond.name)
+        if cond.name:
+            body_frame.defs[cond.name] = then_cond
+        state.push_frame(body_frame)
+        state.exhaust(body_frame)
+        state.pop_frame(finalize_next=False)  # NOTE: scf does not have multiple blocks
+
+        else_frame = lowering.Frame.from_stmts(
+            node.orelse, state, globals=frame.globals
+        )
+        else_cond = else_frame.curr_block.args.append_from(types.Bool, cond.name)
+        if cond.name:
+            else_frame.defs[cond.name] = else_cond
+        state.push_frame(else_frame)
+        state.exhaust(else_frame)
+        state.pop_frame(finalize_next=False)  # NOTE: scf does not have multiple blocks
+
+        yield_names: list[str] = []
+        body_yields: list[ir.SSAValue] = []
+        else_yields: list[ir.SSAValue] = []
+        if node.orelse:
+            for name in body_frame.defs.keys():
+                if name in else_frame.defs:
+                    yield_names.append(name)
+                    body_yields.append(body_frame.get_scope(name))
+                    else_yields.append(else_frame.get_scope(name))
+        else:
+            for name in body_frame.defs.keys():
+                if name in frame.defs:
+                    yield_names.append(name)
+                    body_yields.append(body_frame.get_scope(name))
+                    value = frame.get(name)
+                    if value is None:
+                        raise DialectLoweringError(f"expected value for {name}")
+                    else_yields.append(value)
+
+        if not (
+            body_frame.curr_block.last_stmt
+            and body_frame.curr_block.last_stmt.has_trait(ir.IsTerminator)
+        ):
+            body_frame.append_stmt(Yield(*body_yields))
+
+        if not (
+            else_frame.curr_block.last_stmt
+            and else_frame.curr_block.last_stmt.has_trait(ir.IsTerminator)
+        ):
+            else_frame.append_stmt(Yield(*else_yields))
+
+        stmt = IfElse(
+            cond,
+            then_body=body_frame.curr_region,
+            else_body=else_frame.curr_region,
+        )
+        for result, name, body, else_ in zip(
+            stmt.results, yield_names, body_yields, else_yields
+        ):
+            result.name = name
+            result.type = body.type.join(else_.type)
+            frame.defs[name] = result
+        state.append_stmt(stmt)
+        return lowering.Result()
+
+    def lower_For(
+        self, state: lowering.LoweringState, node: ast.For
+    ) -> lowering.Result:
+        iter_ = state.visit(node.iter).expect_one()
+
+        yields: list[str] = []
+
+        def new_block_arg_if_inside_loop(frame: lowering.Frame, capture: ir.SSAValue):
+            if not capture.name:
+                raise DialectLoweringError("unexpected loop variable captured")
+            yields.append(capture.name)
+            return frame.curr_block.args.append_from(capture.type, capture.name)
+
+        body_frame = state.push_frame(
+            lowering.Frame.from_stmts(
+                node.body,
+                state,
+                globals=state.current_frame.globals,
+                capture_callback=new_block_arg_if_inside_loop,
+            )
+        )
+        loop_var = body_frame.curr_block.args.append_from(types.Any)
+        unpacking(state, node.target, loop_var)
+        state.exhaust(body_frame)
+        # NOTE: this frame won't have phi nodes
+        if yields and (
+            body_frame.curr_block.last_stmt is None
+            or not body_frame.curr_block.last_stmt.has_trait(ir.IsTerminator)
+        ):
+            body_frame.append_stmt(Yield(*[body_frame.defs[name] for name in yields]))  # type: ignore
+        state.pop_frame(finalize_next=False)  # NOTE: scf does not have multiple blocks
+
+        initializers: list[ir.SSAValue] = []
+        for name in yields:
+            value = state.current_frame.get(name)
+            if value is None:
+                raise DialectLoweringError(f"expected value for {name}")
+            initializers.append(value)
+        stmt = For(iter_, body_frame.curr_region, *initializers)
+        for name, result in zip(yields, stmt.results):
+            state.current_frame.defs[name] = result
+            result.name = name
+        state.append_stmt(stmt)
+        return lowering.Result()