kirin-toolchain 0.13.0__py3-none-any.whl

kirin/interp/impl.py

@@ -0,0 +1,197 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Type,
+    Union,
+    Generic,
+    TypeVar,
+    Callable,
+    TypeAlias,
+    overload,
+)
+from dataclasses import dataclass
+
+from kirin import types
+from kirin.ir import Attribute, Statement
+from kirin.interp.value import StatementResult
+
+if TYPE_CHECKING:
+    from kirin.interp.base import FrameABC, BaseInterpreter
+    from kirin.interp.table import MethodTable
+
+MethodTableSelf = TypeVar("MethodTableSelf", bound="MethodTable")
+InterpreterType = TypeVar("InterpreterType", bound="BaseInterpreter")
+FrameType = TypeVar("FrameType", bound="FrameABC")
+StatementType = TypeVar("StatementType", bound=Statement)
+AttributeType = TypeVar("AttributeType", bound=Attribute)
+MethodFunction: TypeAlias = Callable[
+    [MethodTableSelf, InterpreterType, FrameType, StatementType], StatementResult
+]
+AttributeFunction: TypeAlias = Callable[
+    [MethodTableSelf, InterpreterType, AttributeType], Any
+]
+
+
+@dataclass(frozen=True)
+class Signature:
+    """Signature of a statement."""
+
+    stmt: Type[Statement]
+    args: tuple[types.TypeAttribute, ...] | None = None
+
+    def __repr__(self):
+        if self.args:
+            return f"{self.stmt.__name__}[{', '.join(map(repr, self.args))}]"
+        else:
+            return f"{self.stmt.__name__}[...]"
+
+
+SigType = TypeVar("SigType")
+ImplType = TypeVar("ImplType")
+
+
+@dataclass
+class Def(Generic[SigType, ImplType]):
+    """Base class of an interpreter implementation definition."""
+
+    signature: SigType
+    impl: ImplType
+
+
+@dataclass
+class ImplDef(Def[tuple[Signature, ...], "MethodFunction"]):
+    """Definition of an interpreter implementation for a statement."""
+
+    parent: Type[Statement]
+
+    def __repr__(self):
+        if self.parent.dialect:
+            return f"interp {self.parent.dialect.name}.{self.parent.name}"
+        else:
+            return f"interp {self.parent.name}"
+
+    def __call__(self, *args, **kwargs):
+        return self.impl(*args, **kwargs)
+
+
+@dataclass
+class AttributeImplDef(Def[type[Attribute], "AttributeFunction"]):
+    """Definition of an interpreter implementation for an attribute."""
+
+    def __repr__(self):
+        if self.signature.dialect:
+            return f"attribute impl {self.signature.dialect.name}.{self.signature.name}"
+        else:
+            return f"attribute impl {self.signature.name}"
+
+
+StatementType = TypeVar("StatementType", bound=Statement)
+HeadType = TypeVar("HeadType")
+
+
+class impl(Generic[HeadType]):
+    """Decorator to define an interpreter implementation for a statement or attribute.
+
+    !!! note
+        While the `impl` decorator accepts both statements and attributes, and optionally
+        statements with its type signature, unlike a programming language, the actual
+        dispatch behavior given an instance of a statement or attribute is defined by the
+        implementation of the interpreter
+        (via [`lookup_registry`][kirin.interp.base.BaseInterpreter.lookup_registry]).
+
+    # Example
+
+    ```python
+    @dialect.register
+    class MyMethods(interp.MethodTable):
+        @impl(Add)
+        def interp_add(
+            self,
+            interp: Interpreter,
+            frame: Frame,
+            stmt: Add,
+        ) -> StatementResult:
+            ...
+    ```
+    """
+
+    # TODO: validate only concrete types are allowed here
+
+    def __init__(
+        self, stmt_or_attribute: Type[HeadType], *args: types.TypeAttribute
+    ) -> None:
+        if args and issubclass(stmt_or_attribute, Attribute):
+            raise ValueError("Attributes do not take arguments")
+        self.stmt_or_attribute: type[HeadType] = stmt_or_attribute
+        self.args = args
+
+    @overload
+    def __call__(
+        self,
+        func: Union[
+            Callable[
+                [MethodTableSelf, InterpreterType, FrameType, StatementType],
+                StatementResult,
+            ],
+            ImplDef,
+        ],
+    ) -> ImplDef: ...
+
+    @overload
+    def __call__(
+        self,
+        func: Union[
+            Callable[
+                [MethodTableSelf, InterpreterType, AttributeType],
+                Any,
+            ],
+            AttributeImplDef,
+        ],
+    ) -> AttributeImplDef: ...
+
+    def __call__(
+        self,
+        func: Union[
+            Callable[
+                [MethodTableSelf, InterpreterType, FrameType, StatementType],
+                StatementResult,
+            ],
+            Callable[
+                [MethodTableSelf, InterpreterType, AttributeType],
+                Any,
+            ],
+            ImplDef,
+            AttributeImplDef,
+        ],
+    ) -> Def:
+        if issubclass(self.stmt_or_attribute, Statement):
+            return self._impl_statement(self.stmt_or_attribute, func)
+        elif issubclass(self.stmt_or_attribute, Attribute):
+            return self._impl_attribute(self.stmt_or_attribute, func)
+        else:
+            raise ValueError(f"Invalid statement type {self.stmt_or_attribute}")
+
+    def _impl_attribute(
+        self,
+        attr: Type[Attribute],
+        func: Union[Callable, Def],
+    ) -> AttributeImplDef:
+        if isinstance(func, Def):
+            return AttributeImplDef(attr, func.impl)
+        else:
+            return AttributeImplDef(attr, func)
+
+    def _impl_statement(
+        self,
+        stmt: Type[Statement],
+        func: Union[Callable, Def],
+    ) -> ImplDef:
+        if self.args:
+            sig = Signature(stmt, self.args)
+        else:
+            sig = Signature(stmt)
+
+        if isinstance(func, Def):
+            return ImplDef(func.signature + (sig,), func.impl, stmt)
+        else:
+            return ImplDef((sig,), func, stmt)

kirin/interp/result.py

@@ -0,0 +1,93 @@
+from typing import Generic, TypeVar
+from dataclasses import dataclass
+
+from kirin.ir.traits import SymbolOpInterface, CallableStmtInterface
+
+from .frame import FrameABC
+from .exceptions import WrapException
+
+ValueType = TypeVar("ValueType")
+FrameType = TypeVar("FrameType", bound=FrameABC)
+
+
+@dataclass
+class Result(Generic[ValueType]):
+    """Result type for the interpreter.
+
+    This is a generic result type that represents the result of interpretation.
+    The result can be either an `Ok` value or an `Err` value. The `Ok` value
+    represents a successful interpretation result, while the `Err` value
+    represents an error during interpretation with a stack trace. One can use
+    the `expect` method to extract the value from the result, which will raise
+    an exception and print the stack trace if the result is an `Err`.
+    """
+
+    def expect(self) -> ValueType:
+        raise NotImplementedError("unreachable")
+
+
+@dataclass
+class Err(Result[ValueType], Generic[FrameType, ValueType]):
+    exception: Exception
+    frames: list[FrameType]
+
+    def __init__(self, exception: Exception, frames: list[FrameType]):
+        super().__init__()
+        self.exception = exception
+        self.frames = frames
+
+    def __repr__(self) -> str:
+        return f"Err({self.exception.__class__.__name__}: {self.exception})"
+
+    def print_stack(self):
+        """Print the stack trace of the error."""
+        top_method_code = self.frames[0].code
+        if (call_trait := top_method_code.get_trait(CallableStmtInterface)) is None:
+            raise ValueError(f"Method code {top_method_code} is not callable")
+
+        region = call_trait.get_callable_region(top_method_code)
+        name = (
+            top_method_code.get_trait(SymbolOpInterface)
+            .get_sym_name(top_method_code)  # type: ignore
+            .data
+        )
+        args = ",".join(
+            [
+                (
+                    f"{arg.name}"
+                    if arg.type is arg.type.top()
+                    else f"{arg.name}:{arg.type}"
+                )
+                for arg in region.blocks[0].args[1:]
+            ]
+        )
+        print("Traceback (most recent call last):")
+        print(f"  {name}({args})")
+        for frame in reversed(self.frames):
+            if frame.code:
+                frame.code.print()
+        print(f"{self.exception.__class__.__name__}: {self.exception}")
+        print(
+            "================================ Python Stacktrace ================================"
+        )
+
+    def expect(self) -> ValueType:
+        self.print_stack()
+        if isinstance(self.exception, WrapException):
+            raise self.exception.exception from self.exception
+        else:
+            raise self.exception from None
+
+
+@dataclass
+class Ok(Result[ValueType]):
+    value: ValueType
+
+    def __len__(self) -> int:
+        return 1
+
+    def __repr__(self) -> str:
+        return f"Ok({self.value})"
+
+    def expect(self) -> ValueType:
+        return self.value

kirin/interp/state.py

@@ -0,0 +1,71 @@
+from typing import Generic, TypeVar
+from contextlib import contextmanager
+from dataclasses import field, dataclass
+
+from kirin.interp.frame import FrameABC
+
+FrameType = TypeVar("FrameType", bound=FrameABC)
+
+
+@dataclass
+class InterpreterState(Generic[FrameType]):
+    """Interpreter state.
+
+    This class represents the state of the interpreter. It contains the
+    stack of frames for the interpreter. The stack of frames is used to
+    store the current state of the interpreter during interpretation.
+    """
+
+    frames: list[FrameType] = field(default_factory=list)
+
+    def push_frame(self, frame: FrameType):
+        """Push a frame onto the stack.
+
+        Args:
+            frame(FrameType): The frame to push onto the stack.
+
+        Returns:
+            FrameType: The frame that was pushed.
+        """
+        self.frames.append(frame)
+        return frame
+
+    def pop_frame(self) -> FrameType:
+        """Pop a frame from the stack.
+
+        Returns:
+            FrameType: The frame that was popped.
+        """
+        return self.frames.pop()
+
+    @contextmanager
+    def new_frame(self, frame: FrameType):
+        """Context manager to push and pop a frame.
+
+        This context manager is used to push a frame onto the stack before
+        executing a block of code and pop the frame from the stack after
+        executing the block of code.
+
+        Args:
+            frame(FrameType): The frame to push onto the stack.
+
+        Yields:
+            FrameType: The frame that was pushed.
+        """
+        self.push_frame(frame)
+        try:
+            yield frame
+        finally:
+            self.pop_frame()
+
+    def current_frame(self) -> FrameType:
+        """Get the current frame.
+
+        Returns:
+            FrameType: The current frame.
+        """
+        if not self.frames:
+            raise ValueError(
+                "unable to retrieve the current frame because no frames were pushed"
+            )
+        return self.frames[-1]

kirin/interp/table.py

@@ -0,0 +1,40 @@
+import inspect
+from abc import ABC
+from typing import TYPE_CHECKING, TypeVar, ClassVar
+from dataclasses import dataclass
+
+from kirin.interp.base import BaseInterpreter
+from kirin.interp.impl import ImplDef, AttributeImplDef
+
+if TYPE_CHECKING:
+    from kirin.ir import Attribute
+    from kirin.interp.base import BaseInterpreter
+    from kirin.interp.impl import Signature, MethodFunction, AttributeFunction
+
+
+InterpreterType = TypeVar("InterpreterType", bound="BaseInterpreter")
+ValueType = TypeVar("ValueType")
+
+
+@dataclass
+class MethodTable(ABC):
+    """Base class to define lookup tables for interpreting code for IR nodes in a dialect."""
+
+    table: ClassVar[dict["Signature", "MethodFunction"]]
+    """Lookup table for interpreting code for IR nodes in a dialect."""
+    attribute: ClassVar[dict[type["Attribute"], "AttributeFunction"]]
+    """Lookup table for interpreting code for IR attributes in a dialect."""
+
+    def __init_subclass__(cls) -> None:
+        # init the subclass first
+        super().__init_subclass__()
+        cls.table = {}
+        for _, value in inspect.getmembers(cls):
+            if isinstance(value, ImplDef):
+                for sig in value.signature:
+                    cls.table[sig] = value.impl
+
+        cls.attribute = {}
+        for _, value in inspect.getmembers(cls):
+            if isinstance(value, AttributeImplDef):
+                cls.attribute[value.signature] = value.impl

kirin/interp/value.py

@@ -0,0 +1,73 @@
+from typing import Tuple, Generic, TypeVar, TypeAlias, final
+from dataclasses import dataclass
+
+from kirin.ir import Block
+
+ValueType = TypeVar("ValueType")
+
+
+@final
+@dataclass
+class ReturnValue(Generic[ValueType]):
+    """Return value from a statement evaluation.
+
+    This class represents a return value from a statement evaluation. It is used
+    to indicate that the statement evaluation should later pop the frame and
+    return the value. Kirin does not allow multiple return values to follow Python
+    semantics. If you want to return multiple values, you should return a tuple.
+    """
+
+    value: ValueType
+
+    def __len__(self) -> int:
+        return 0
+
+
+@final
+@dataclass
+class YieldValue(Generic[ValueType]):
+    """Yield value from a statement evaluation.
+
+    This class represents values returned from a statement that terminates current
+    region execution and returns the values to the caller. Unlike `ReturnValue`, this
+    class won't pop the frame and return the value to the caller.
+    """
+
+    values: tuple[ValueType, ...]
+
+    def __len__(self) -> int:
+        return len(self.values)
+
+    def __iter__(self):
+        return iter(self.values)
+
+    def __getitem__(self, index: int) -> ValueType:
+        return self.values[index]
+
+
+@final
+@dataclass(init=False)
+class Successor(Generic[ValueType]):
+    """Successor block from a statement evaluation."""
+
+    block: Block
+    block_args: Tuple[ValueType, ...]
+
+    def __init__(self, block: Block, *block_args: ValueType):
+        super().__init__()
+        self.block = block
+        self.block_args = block_args
+
+    def __hash__(self) -> int:
+        return hash(self.block)
+
+    def __len__(self) -> int:
+        return 0
+
+
+SpecialValue: TypeAlias = (
+    None | ReturnValue[ValueType] | YieldValue[ValueType] | Successor[ValueType]
+)
+"""Special value for statement evaluation."""
+StatementResult: TypeAlias = tuple[ValueType, ...] | SpecialValue[ValueType]
+"""Type alias for the result of a statement evaluation."""

kirin/ir/__init__.py

@@ -0,0 +1,46 @@
+"""IR module for kirin.
+
+This module contains the data structure definition
+of the intermediate representation (IR) for kirin.
+"""
+
+from kirin.ir import attrs as attrs, traits as traits
+from kirin.ir.ssa import (
+    SSAValue as SSAValue,
+    TestValue as TestValue,
+    ResultValue as ResultValue,
+    BlockArgument as BlockArgument,
+    DeletedSSAValue as DeletedSSAValue,
+)
+from kirin.ir.use import Use as Use
+from kirin.ir.group import DialectGroup as DialectGroup, dialect_group as dialect_group
+from kirin.ir.nodes import (
+    Block as Block,
+    IRNode as IRNode,
+    Region as Region,
+    Statement as Statement,
+)
+from kirin.ir.method import Method as Method
+from kirin.ir.traits import (
+    Pure as Pure,
+    HasParent as HasParent,
+    MaybePure as MaybePure,
+    StmtTrait as StmtTrait,
+    RegionTrait as RegionTrait,
+    SymbolTable as SymbolTable,
+    ConstantLike as ConstantLike,
+    HasSignature as HasSignature,
+    IsTerminator as IsTerminator,
+    NoTerminator as NoTerminator,
+    SSACFGRegion as SSACFGRegion,
+    FromPythonCall as FromPythonCall,
+    FromPythonWith as FromPythonWith,
+    IsolatedFromAbove as IsolatedFromAbove,
+    SymbolOpInterface as SymbolOpInterface,
+    PythonLoweringTrait as PythonLoweringTrait,
+    CallableStmtInterface as CallableStmtInterface,
+    FromPythonWithSingleItem as FromPythonWithSingleItem,
+)
+from kirin.ir.dialect import Dialect as Dialect
+from kirin.ir.attrs.py import PyAttr as PyAttr
+from kirin.ir.attrs.abc import Attribute as Attribute, AttributeMeta as AttributeMeta

kirin/ir/attrs/__init__.py

@@ -0,0 +1,20 @@
+"""Compile-time values in Kirin IR.
+
+This module contains the following:
+
+## `abc` module
+`abc.AttributeMeta`: The metaclass for all attributes.
+`abc.Attribute`: The base class for all attributes.
+
+## `types` module
+`types.TypeAttribute`: The base class for all type attributes.
+`types.PyClass`: A type attribute representing a Python class.
+`types.TypeVar`: A type attribute representing a type variable.
+`types.Literal`: A type attribute representing a literal type.
+`types.Generic`: A type attribute representing a generic type.
+`types.Union`: A type attribute representing a union type.
+
+## `py` module
+
+`py.PyAttr`: An attribute representing a Python value.
+"""

kirin/ir/attrs/_types.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+
+from .abc import Attribute
+
+
+@dataclass
+class _TypeAttribute(Attribute):
+    pass

kirin/ir/attrs/_types.pyi

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+from .abc import Attribute
+from .types import Union, Generic, Literal, PyClass, TypeVar, TypeAttribute
+
+@dataclass
+class _TypeAttribute(Attribute):
+    def is_subseteq_Union(self, other: Union) -> bool: ...
+    def is_subseteq_Literal(self, other: Literal) -> bool: ...
+    def is_subseteq_TypeVar(self, other: TypeVar) -> bool: ...
+    def is_subseteq_PyClass(self, other: PyClass) -> bool: ...
+    def is_subseteq_Generic(self, other: Generic) -> bool: ...
+    def is_subseteq_fallback(self, other: TypeAttribute) -> bool: ...

kirin/ir/attrs/abc.py

@@ -0,0 +1,46 @@
+from abc import ABC, ABCMeta, abstractmethod
+from typing import TYPE_CHECKING, ClassVar, Optional
+from dataclasses import field, dataclass
+
+from kirin.print import Printable
+from kirin.lattice.abc import LatticeMeta, SingletonMeta
+
+if TYPE_CHECKING:
+    from kirin.ir.dialect import Dialect
+
+
+class AttributeMeta(ABCMeta):
+    """Metaclass for attributes."""
+
+    pass
+
+
+class LatticeAttributeMeta(LatticeMeta, AttributeMeta):
+    """Metaclass for lattice attributes."""
+
+    pass
+
+
+class SingletonLatticeAttributeMeta(LatticeAttributeMeta, SingletonMeta):
+    """Metaclass for singleton lattice attributes."""
+
+    pass
+
+
+@dataclass(eq=False)
+class Attribute(ABC, Printable, metaclass=AttributeMeta):
+    """ABC for compile-time values. All attributes are hashable
+    and thus need to implement the `__hash__` method.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
+
+    dialect: ClassVar[Optional["Dialect"]] = field(default=None, init=False, repr=False)
+    """Dialect of the attribute. (default: None)"""
+    name: ClassVar[str] = field(init=False, repr=False)
+    """Name of the attribute in printing and other text format."""
+
+    @abstractmethod
+    def __hash__(self) -> int: ...

kirin/ir/attrs/py.py

@@ -0,0 +1,45 @@
+from typing import Generic, TypeVar
+from dataclasses import dataclass
+
+from kirin.print import Printer
+
+from .abc import Attribute
+from .types import PyClass, TypeAttribute
+
+T = TypeVar("T")
+
+
+@dataclass
+class PyAttr(Generic[T], Attribute):
+    """Python attribute for compile-time values.
+    This is a generic attribute that holds a Python value.
+
+    The constructor takes a Python value and an optional type attribute.
+    If the type attribute is not provided, the type of the value is inferred
+    as `PyClass(type(value))`.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
+
+    name = "PyAttr"
+    data: T
+    type: TypeAttribute
+
+    def __init__(self, data: T, pytype: TypeAttribute | None = None):
+        self.data = data
+
+        if pytype is None:
+            self.type = PyClass(type(data))
+        else:
+            self.type = pytype
+
+    def __hash__(self):
+        return hash(self.data) + hash(self.type)
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.plain_print(repr(self.data))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.type)