kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/cf/constprop.py

@@ -0,0 +1,51 @@
+from kirin.interp import Successor, MethodTable, impl
+from kirin.analysis import const
+from kirin.dialects.cf.stmts import Branch, ConditionalBranch
+from kirin.dialects.cf.dialect import dialect
+
+
+@dialect.register(key="constprop")
+class DialectConstProp(MethodTable):
+
+    @impl(Branch)
+    def branch(self, interp: const.Propagate, frame: const.Frame, stmt: Branch):
+        interp.state.current_frame().worklist.append(
+            Successor(stmt.successor, *frame.get_values(stmt.arguments))
+        )
+        return ()
+
+    @impl(ConditionalBranch)
+    def conditional_branch(
+        self,
+        interp: const.Propagate,
+        frame: const.Frame,
+        stmt: ConditionalBranch,
+    ):
+        frame = interp.state.current_frame()
+        cond = frame.get(stmt.cond)
+        if isinstance(cond, const.Value):
+            else_successor = Successor(
+                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+            )
+            then_successor = Successor(
+                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+            )
+            if cond.data:
+                frame.worklist.append(then_successor)
+            else:
+                frame.worklist.append(else_successor)
+        else:
+            frame.entries[stmt.cond] = const.Value(True)
+            then_successor = Successor(
+                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+            )
+            frame.worklist.append(then_successor)
+
+            frame.entries[stmt.cond] = const.Value(False)
+            else_successor = Successor(
+                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+            )
+            frame.worklist.append(else_successor)
+
+            frame.entries[stmt.cond] = cond
+        return ()

kirin/dialects/cf/dialect.py

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

kirin/dialects/cf/emit.py

@@ -0,0 +1,58 @@
+from typing import IO, TypeVar
+
+from kirin import emit
+from kirin.interp import Successor, MethodTable, impl
+from kirin.emit.julia import EmitJulia
+
+from .stmts import Branch, ConditionalBranch
+from .dialect import dialect
+
+IO_t = TypeVar("IO_t", bound=IO)
+
+
+@dialect.register(key="emit.julia")
+class JuliaMethodTable(MethodTable):
+
+    @impl(Branch)
+    def emit_branch(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Branch
+    ):
+        interp.writeln(frame, f"@goto {interp.block_id[stmt.successor]};")
+        frame.worklist.append(
+            Successor(stmt.successor, frame.get_values(stmt.arguments))
+        )
+        return ()
+
+    @impl(ConditionalBranch)
+    def emit_cbr(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: ConditionalBranch
+    ):
+        cond = frame.get(stmt.cond)
+        interp.writeln(frame, f"if {cond}")
+        frame.indent += 1
+        values = frame.get_values(stmt.then_arguments)
+        block_values = tuple(interp.ssa_id[x] for x in stmt.then_successor.args)
+        frame.set_values(stmt.then_successor.args, block_values)
+        for x, y in zip(block_values, values):
+            interp.writeln(frame, f"{x} = {y};")
+        interp.writeln(frame, f"@goto {interp.block_id[stmt.then_successor]};")
+        frame.indent -= 1
+        interp.writeln(frame, "else")
+        frame.indent += 1
+
+        values = frame.get_values(stmt.else_arguments)
+        block_values = tuple(interp.ssa_id[x] for x in stmt.else_successor.args)
+        frame.set_values(stmt.else_successor.args, block_values)
+        for x, y in zip(block_values, values):
+            interp.writeln(frame, f"{x} = {y};")
+        interp.writeln(frame, f"@goto {interp.block_id[stmt.else_successor]};")
+        frame.indent -= 1
+        interp.writeln(frame, "end")
+
+        frame.worklist.append(
+            Successor(stmt.then_successor, frame.get_values(stmt.then_arguments))
+        )
+        frame.worklist.append(
+            Successor(stmt.else_successor, frame.get_values(stmt.else_arguments))
+        )
+        return ()

kirin/dialects/cf/interp.py

@@ -0,0 +1,24 @@
+from kirin.interp import Frame, Successor, Interpreter, MethodTable, impl
+from kirin.dialects.cf.stmts import Branch, ConditionalBranch
+from kirin.dialects.cf.dialect import dialect
+
+
+@dialect.register
+class CfInterpreter(MethodTable):
+
+    @impl(Branch)
+    def branch(self, interp: Interpreter, frame: Frame, stmt: Branch):
+        return Successor(stmt.successor, *frame.get_values(stmt.arguments))
+
+    @impl(ConditionalBranch)
+    def conditional_branch(
+        self, interp: Interpreter, frame: Frame, stmt: ConditionalBranch
+    ):
+        if frame.get(stmt.cond):
+            return Successor(
+                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+            )
+        else:
+            return Successor(
+                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+            )

kirin/dialects/cf/stmts.py

@@ -0,0 +1,68 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+from kirin.print.printer import Printer
+from kirin.dialects.cf.dialect import dialect
+
+
+@statement(dialect=dialect)
+class Branch(ir.Statement):
+    name = "br"
+    traits = frozenset({ir.IsTerminator()})
+
+    arguments: tuple[ir.SSAValue, ...]
+    successor: ir.Block = info.block()
+
+    def verify(self) -> None:
+        return
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="keyword"):
+            printer.print_name(self)
+
+        printer.plain_print(" ")
+        printer.plain_print(printer.state.block_id[self.successor])
+        printer.print_seq(
+            self.arguments,
+            delim=", ",
+            prefix="(",
+            suffix=")",
+        )
+
+
+@statement(dialect=dialect)
+class ConditionalBranch(ir.Statement):
+    name = "cond_br"
+    traits = frozenset({ir.IsTerminator()})
+
+    cond: ir.SSAValue = info.argument(types.Bool)
+    then_arguments: tuple[ir.SSAValue, ...]
+    else_arguments: tuple[ir.SSAValue, ...]
+
+    then_successor: ir.Block = info.block()
+    else_successor: ir.Block = info.block()
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="keyword"):
+            printer.print_name(self)
+
+        printer.plain_print(" ")
+        printer.print(self.cond)
+
+        with printer.rich(style="keyword"):
+            printer.plain_print(" goto ")
+
+        printer.plain_print(printer.state.block_id[self.then_successor])
+        printer.plain_print("(")
+        printer.print_seq(self.then_arguments, delim=", ")
+        printer.plain_print(")")
+
+        with printer.rich(style="keyword"):
+            printer.plain_print(" else ")
+
+        printer.plain_print(printer.state.block_id[self.else_successor])
+        printer.plain_print("(")
+        printer.print_seq(self.else_arguments, delim=", ")
+        printer.plain_print(")")
+
+    def verify(self) -> None:
+        return

kirin/dialects/cf/typeinfer.py

@@ -0,0 +1,27 @@
+from kirin.interp import Successor, MethodTable, AbstractFrame, impl
+from kirin.dialects.cf.stmts import Branch, ConditionalBranch
+from kirin.analysis.typeinfer import TypeInference
+from kirin.dialects.cf.dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(MethodTable):
+
+    @impl(Branch)
+    def branch(self, interp: TypeInference, frame: AbstractFrame, stmt: Branch):
+        frame.worklist.append(
+            Successor(stmt.successor, *frame.get_values(stmt.arguments))
+        )
+        return ()
+
+    @impl(ConditionalBranch)
+    def conditional_branch(
+        self, interp: TypeInference, frame: AbstractFrame, stmt: ConditionalBranch
+    ):
+        frame.worklist.append(
+            Successor(stmt.else_successor, *frame.get_values(stmt.else_arguments))
+        )
+        frame.worklist.append(
+            Successor(stmt.then_successor, *frame.get_values(stmt.then_arguments))
+        )
+        return ()

kirin/dialects/eltype.py

@@ -0,0 +1,23 @@
+"""This dialect offers a statement `eltype` for other dialects'
+type inference to query/implement the element type of a value.
+For example, the `ilist` dialect implements the `eltype` statement
+on the `ilist.IList` type to return the element type.
+"""
+
+from kirin import ir, types
+from kirin.decl import info, statement
+
+dialect = ir.Dialect("eltype")
+
+
+@statement(dialect=dialect)
+class ElType(ir.Statement):
+    """Returns the element type of a value.
+
+    This statement is used by other dialects to query the element type of a value.
+    """
+
+    container: ir.SSAValue = info.argument(types.Any)
+    """The value to query the element type of."""
+    elem: ir.ResultValue = info.result(types.PyClass(types.TypeAttribute))
+    """The element type of the value."""

kirin/dialects/func/__init__.py

@@ -0,0 +1,20 @@
+"""A function dialect that is compatible with python semantics.
+"""
+
+from kirin.dialects.func import (
+    emit as emit,
+    interp as interp,
+    constprop as constprop,
+    typeinfer as typeinfer,
+)
+from kirin.dialects.func.attrs import Signature as Signature, MethodType as MethodType
+from kirin.dialects.func.stmts import (
+    Call as Call,
+    Invoke as Invoke,
+    Lambda as Lambda,
+    Return as Return,
+    Function as Function,
+    GetField as GetField,
+    ConstantNone as ConstantNone,
+)
+from kirin.dialects.func.dialect import dialect as dialect

kirin/dialects/func/attrs.py

@@ -0,0 +1,39 @@
+from typing import Generic, TypeVar
+from dataclasses import dataclass
+
+from kirin import types
+from kirin.ir import Method, Attribute
+from kirin.print.printer import Printer
+from kirin.dialects.func.dialect import dialect
+
+TypeofMethodType = types.PyClass[Method]
+MethodType = types.Generic(
+    Method, types.TypeVar("Params", types.Tuple), types.TypeVar("Ret")
+)
+TypeLatticeElem = TypeVar("TypeLatticeElem", bound="types.TypeAttribute")
+
+
+@dialect.register
+@dataclass
+class Signature(Generic[TypeLatticeElem], Attribute):
+    """function body signature.
+
+    This is not a type attribute because it just stores
+    the signature of a function at its definition site.
+    We don't perform type inference on this directly.
+
+    The type of a function is the type of `inputs[0]`, which
+    typically is a `MethodType`.
+    """
+
+    name = "Signature"
+    inputs: tuple[TypeLatticeElem, ...]
+    output: TypeLatticeElem  # multi-output must be tuple
+
+    def __hash__(self) -> int:
+        return hash((self.inputs, self.output))
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_seq(self.inputs, delim=", ", prefix="(", suffix=")")
+        printer.plain_print(" -> ")
+        printer.print(self.output)

kirin/dialects/func/constprop.py

@@ -0,0 +1,138 @@
+from kirin import ir, types
+from kirin.interp import MethodTable, ReturnValue, StatementResult, impl
+from kirin.analysis import const
+from kirin.dialects.func.stmts import Call, Invoke, Lambda, Return, GetField
+from kirin.dialects.func.dialect import dialect
+
+
+@dialect.register(key="constprop")
+class DialectConstProp(MethodTable):
+
+    @impl(Return)
+    def return_(
+        self, interp: const.Propagate, frame: const.Frame, stmt: Return
+    ) -> StatementResult[const.Result]:
+        return ReturnValue(frame.get(stmt.value))
+
+    @impl(Call)
+    def call(
+        self, interp: const.Propagate, frame: const.Frame, stmt: Call
+    ) -> StatementResult[const.Result]:
+        # give up on dynamic method calls
+        callee = frame.get(stmt.callee)
+        if isinstance(callee, const.PartialLambda):
+            call_frame, ret = self._call_lambda(
+                interp,
+                callee,
+                interp.permute_values(
+                    callee.argnames, frame.get_values(stmt.inputs), stmt.kwargs
+                ),
+            )
+            if not call_frame.frame_is_not_pure:
+                frame.should_be_pure.add(stmt)
+            return (ret,)
+
+        if not (isinstance(callee, const.Value) and isinstance(callee.data, ir.Method)):
+            return (const.Result.bottom(),)
+
+        mt: ir.Method = callee.data
+        call_frame, ret = interp.run_method(
+            mt,
+            interp.permute_values(
+                mt.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+            ),
+        )
+        if not call_frame.frame_is_not_pure:
+            frame.should_be_pure.add(stmt)
+        return (ret,)
+
+    def _call_lambda(
+        self,
+        interp: const.Propagate,
+        callee: const.PartialLambda,
+        args: tuple[const.Result, ...],
+    ):
+        # NOTE: we still use PartialLambda because
+        # we want to gurantee what we receive here in captured
+        # values are all lattice elements and not just obtain via
+        # Const(Method(...)) which is Any.
+        if (trait := callee.code.get_trait(ir.SymbolOpInterface)) is not None:
+            name = trait.get_sym_name(callee.code).data
+        else:
+            name = "lambda"
+
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=name,
+            arg_names=callee.argnames,
+            dialects=interp.dialects,
+            code=callee.code,
+            fields=callee.captured,
+        )
+        return interp.run_method(mt, args)
+
+    @impl(Invoke)
+    def invoke(
+        self,
+        interp: const.Propagate,
+        frame: const.Frame,
+        stmt: Invoke,
+    ) -> StatementResult[const.Result]:
+        invoke_frame, ret = interp.run_method(
+            stmt.callee,
+            interp.permute_values(
+                stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+            ),
+        )
+        if not invoke_frame.frame_is_not_pure:
+            frame.should_be_pure.add(stmt)
+        return (ret,)
+
+    @impl(Lambda)
+    def lambda_(
+        self, interp: const.Propagate, frame: const.Frame, stmt: Lambda
+    ) -> StatementResult[const.Result]:
+        captured = frame.get_values(stmt.captured)
+        arg_names = [
+            arg.name or str(idx) for idx, arg in enumerate(stmt.body.blocks[0].args)
+        ]
+        if stmt.body.blocks and types.is_tuple_of(captured, const.Value):
+            return (
+                const.Value(
+                    ir.Method(
+                        mod=None,
+                        py_func=None,
+                        sym_name=stmt.sym_name,
+                        arg_names=arg_names,
+                        dialects=interp.dialects,
+                        code=stmt,
+                        fields=tuple(each.data for each in captured),
+                    )
+                ),
+            )
+
+        return (
+            const.PartialLambda(
+                arg_names,
+                stmt,
+                tuple(each for each in captured),
+            ),
+        )
+
+    @impl(GetField)
+    def getfield(
+        self,
+        interp: const.Propagate,
+        frame: const.Frame,
+        stmt: GetField,
+    ) -> StatementResult[const.Result]:
+        callee_self = frame.get(stmt.obj)
+        if isinstance(callee_self, const.Value) and isinstance(
+            callee_self.data, ir.Method
+        ):
+            mt: ir.Method = callee_self.data
+            return (const.Value(mt.fields[stmt.field]),)
+        elif isinstance(callee_self, const.PartialLambda):
+            return (callee_self.captured[stmt.field],)
+        return (const.Unknown(),)

kirin/dialects/func/dialect.py

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

kirin/dialects/func/emit.py

@@ -0,0 +1,80 @@
+from typing import IO, TypeVar
+
+from kirin import emit
+from kirin.interp import MethodTable, InterpreterError, impl
+from kirin.emit.julia import EmitJulia
+
+from .stmts import Call, Invoke, Lambda, Return, Function, GetField, ConstantNone
+from .dialect import dialect
+
+IO_t = TypeVar("IO_t", bound=IO)
+
+
+@dialect.register(key="emit.julia")
+class JuliaMethodTable(MethodTable):
+
+    @impl(Function)
+    def emit_function(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Function
+    ):
+        fn_args = stmt.body.blocks[0].args[1:]
+        argnames = frame.get_values(fn_args)
+        argtypes = tuple(interp.emit_attribute(x.type) for x in fn_args)
+        args = [f"{name}::{type}" for name, type in zip(argnames, argtypes)]
+        interp.write(f"function {stmt.sym_name}({', '.join(args)})")
+        frame.indent += 1
+        interp.run_ssacfg_region(frame, stmt.body)
+        frame.indent -= 1
+        interp.writeln(frame, "end")
+        return ()
+
+    @impl(Return)
+    def emit_return(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Return
+    ):
+        interp.writeln(frame, f"return {frame.get(stmt.value)}")
+        return ()
+
+    @impl(ConstantNone)
+    def emit_constant_none(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: ConstantNone
+    ):
+        return ("nothing",)
+
+    @impl(Call)
+    def emit_call(self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Call):
+        if stmt.kwargs:
+            raise InterpreterError("cannot emit kwargs for dyanmic calls")
+        return (
+            f"{frame.get(stmt.callee)}({', '.join(frame.get_values(stmt.inputs))})",
+        )
+
+    @impl(Invoke)
+    def emit_invoke(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Invoke
+    ):
+        args = interp.permute_values(
+            stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+        )
+        return (f"{stmt.callee.sym_name}({', '.join(args)})",)
+
+    @impl(Lambda)
+    def emit_lambda(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Lambda
+    ):
+        args = tuple(interp.ssa_id[x] for x in stmt.body.blocks[0].args[1:])
+        frame.set_values(stmt.body.blocks[0].args, args)
+        frame.set_values((stmt.body.blocks[0].args[0],), (stmt.sym_name,))
+        frame.captured[stmt.body.blocks[0].args[0]] = frame.get_values(stmt.captured)
+        interp.writeln(frame, f"function {stmt.sym_name}({', '.join(args[1:])})")
+        frame.indent += 1
+        interp.run_ssacfg_region(frame, stmt.body)
+        frame.indent -= 1
+        interp.writeln(frame, "end")
+        return (stmt.sym_name,)
+
+    @impl(GetField)
+    def emit_getfield(
+        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: GetField
+    ):
+        return (frame.captured[stmt.obj][stmt.field],)

kirin/dialects/func/interp.py

@@ -0,0 +1,68 @@
+from kirin.ir import Method
+from kirin.interp import Frame, MethodTable, ReturnValue, impl, concrete
+from kirin.dialects.func.stmts import (
+    Call,
+    Invoke,
+    Lambda,
+    Return,
+    GetField,
+    ConstantNone,
+)
+from kirin.dialects.func.dialect import dialect
+
+
+@dialect.register
+class Interpreter(MethodTable):
+
+    @impl(Call)
+    def call(self, interp: concrete.Interpreter, frame: Frame, stmt: Call):
+        mt: Method = frame.get(stmt.callee)
+        _, result = interp.run_method(
+            mt,
+            interp.permute_values(
+                mt.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+            ),
+        )
+        return (result,)
+
+    @impl(Invoke)
+    def invoke(self, interp: concrete.Interpreter, frame: Frame, stmt: Invoke):
+        _, result = interp.run_method(
+            stmt.callee,
+            interp.permute_values(
+                stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
+            ),
+        )
+        return (result,)
+
+    @impl(Return)
+    def return_(self, interp: concrete.Interpreter, frame: Frame, stmt: Return):
+        return ReturnValue(frame.get(stmt.value))
+
+    @impl(ConstantNone)
+    def const_none(
+        self, interp: concrete.Interpreter, frame: Frame, stmt: ConstantNone
+    ):
+        return (None,)
+
+    @impl(GetField)
+    def getfield(self, interp: concrete.Interpreter, frame: Frame, stmt: GetField):
+        mt: Method = frame.get(stmt.obj)
+        return (mt.fields[stmt.field],)
+
+    @impl(Lambda)
+    def lambda_(self, interp: concrete.Interpreter, frame: Frame, stmt: Lambda):
+        return (
+            Method(
+                mod=None,
+                py_func=None,
+                sym_name=stmt.name,
+                arg_names=[
+                    arg.name or str(idx)
+                    for idx, arg in enumerate(stmt.body.blocks[0].args)
+                ],
+                dialects=interp.dialects,
+                code=stmt,
+                fields=frame.get_values(stmt.captured),
+            ),
+        )