kirin-toolchain 0.13.0__py3-none-any.whl

kirin/ir/group.py

@@ -0,0 +1,249 @@
+import inspect
+from types import ModuleType
+from typing import (
+    TYPE_CHECKING,
+    Union,
+    Generic,
+    TypeVar,
+    Callable,
+    ParamSpec,
+    Concatenate,
+    overload,
+)
+from functools import update_wrapper
+from dataclasses import dataclass
+from collections.abc import Iterable
+
+from kirin.ir.method import Method
+from kirin.exceptions import CompilerError
+
+if TYPE_CHECKING:
+    from kirin.registry import Registry
+    from kirin.ir.dialect import Dialect
+
+PassParams = ParamSpec("PassParams")
+RunPass = Callable[Concatenate[Method, PassParams], None]
+RunPassGen = Callable[["DialectGroup"], RunPass[PassParams]]
+
+
+@dataclass(init=False)
+class DialectGroup(Generic[PassParams]):
+    # method wrapper params
+    Param = ParamSpec("Param")
+    RetType = TypeVar("RetType")
+    MethodTransform = Callable[[Callable[Param, RetType]], Method[Param, RetType]]
+
+    data: frozenset["Dialect"]
+    """The set of dialects in the group."""
+    # NOTE: this is used to create new dialect groups from existing one
+    run_pass_gen: RunPassGen[PassParams] | None = None
+    """the function that generates the `run_pass` function.
+
+    This is used to create new dialect groups from existing ones, while
+    keeping the same `run_pass` function.
+    """
+    run_pass: RunPass[PassParams] | None = None
+    """the function that runs the passes on the method."""
+
+    def __init__(
+        self,
+        dialects: Iterable[Union["Dialect", ModuleType]],
+        run_pass: RunPassGen[PassParams] | None = None,
+    ):
+        def identity(code: Method):
+            pass
+
+        self.data = frozenset(self.map_module(dialect) for dialect in dialects)
+        if run_pass is None:
+            self.run_pass_gen = None
+            self.run_pass = None
+        else:
+            self.run_pass_gen = run_pass
+            self.run_pass = run_pass(self)
+
+    def __iter__(self):
+        return iter(self.data)
+
+    def __repr__(self) -> str:
+        names = ", ".join(each.name for each in self.data)
+        return f"DialectGroup([{names}])"
+
+    @staticmethod
+    def map_module(dialect: Union["Dialect", ModuleType]) -> "Dialect":
+        """map the module to the dialect if it is a module.
+        It assumes that the module has a `dialect` attribute
+        that is an instance of [`Dialect`][kirin.ir.Dialect].
+        """
+        if isinstance(dialect, ModuleType):
+            return getattr(dialect, "dialect")
+        return dialect
+
+    def add(self, dialect: Union["Dialect", ModuleType]) -> "DialectGroup":
+        """add a dialect to the group.
+
+        Args:
+            dialect (Union[Dialect, ModuleType]): the dialect to add
+
+        Returns:
+            DialectGroup: the new dialect group with the added
+        """
+        return self.union([dialect])
+
+    def union(self, dialect: Iterable[Union["Dialect", ModuleType]]) -> "DialectGroup":
+        """union a set of dialects to the group.
+
+        Args:
+            dialect (Iterable[Union[Dialect, ModuleType]]): the dialects to union
+
+        Returns:
+            DialectGroup: the new dialect group with the union.
+        """
+        return DialectGroup(
+            dialects=self.data.union(frozenset(self.map_module(d) for d in dialect)),
+            run_pass=self.run_pass_gen,  # pass the run_pass_gen function
+        )
+
+    def discard(self, dialect: Union["Dialect", ModuleType]) -> "DialectGroup":
+        """discard a dialect from the group.
+
+        !!! note
+            This does not raise an error if the dialect is not in the group.
+
+        Args:
+            dialect (Union[Dialect, ModuleType]): the dialect to discard
+
+        Returns:
+            DialectGroup: the new dialect group with the discarded dialect.
+        """
+        dialect_ = self.map_module(dialect)
+        return DialectGroup(
+            dialects=frozenset(
+                each for each in self.data if each.name != dialect_.name
+            ),
+            run_pass=self.run_pass_gen,  # pass the run_pass_gen function
+        )
+
+    @property
+    def registry(self) -> "Registry":
+        """return the registry for the dialect group. This
+        returns a proxy object that can be used to select
+        the lowering interpreters, interpreters, and codegen
+        for the dialects in the group.
+
+        Returns:
+            Registry: the registry object.
+        """
+        from kirin.registry import Registry
+
+        return Registry(self)
+
+    @overload
+    def __call__(
+        self,
+        py_func: Callable[Param, RetType],
+        *args: PassParams.args,
+        **options: PassParams.kwargs,
+    ) -> Method[Param, RetType]: ...
+
+    @overload
+    def __call__(
+        self,
+        py_func: None = None,
+        *args: PassParams.args,
+        **options: PassParams.kwargs,
+    ) -> MethodTransform[Param, RetType]: ...
+
+    def __call__(
+        self,
+        py_func: Callable[Param, RetType] | None = None,
+        *args: PassParams.args,
+        **options: PassParams.kwargs,
+    ) -> Method[Param, RetType] | MethodTransform[Param, RetType]:
+        """create a method from the python function.
+
+        Args:
+            py_func (Callable): the python function to create the method from.
+            args (PassParams.args): the arguments to pass to the run_pass function.
+            options (PassParams.kwargs): the keyword arguments to pass to the run_pass function.
+
+        Returns:
+            Method: the method created from the python function.
+        """
+        from kirin.lowering import Lowering
+
+        emit_ir = Lowering(self)
+
+        def wrapper(py_func: Callable) -> Method:
+            if py_func.__name__ == "<lambda>":
+                raise ValueError("Cannot compile lambda functions")
+
+            lineno_offset, file = 0, ""
+            frame = inspect.currentframe()
+            if frame and frame.f_back is not None and frame.f_back.f_back is not None:
+                call_site_frame = frame.f_back.f_back
+                if py_func.__name__ in call_site_frame.f_locals:
+                    raise CompilerError(
+                        f"overwriting function definition of `{py_func.__name__}`"
+                    )
+
+                lineno_offset = call_site_frame.f_lineno - 1
+                file = call_site_frame.f_code.co_filename
+
+            code = emit_ir.run(py_func, lineno_offset=lineno_offset)
+            mt = Method(
+                mod=inspect.getmodule(py_func),
+                py_func=py_func,
+                sym_name=py_func.__name__,
+                arg_names=["#self#"] + inspect.getfullargspec(py_func).args,
+                dialects=self,
+                code=code,
+                file=file,
+            )
+            if doc := inspect.getdoc(py_func):
+                mt.__doc__ = doc
+
+            if self.run_pass is not None:
+                self.run_pass(mt, *args, **options)
+            return mt
+
+        if py_func is not None:
+            return wrapper(py_func)
+        return wrapper
+
+
+def dialect_group(
+    dialects: Iterable[Union["Dialect", ModuleType]]
+) -> Callable[[RunPassGen[PassParams]], DialectGroup[PassParams]]:
+    """Create a dialect group from the given dialects based on the
+    definition of `run_pass` function.
+
+    Args:
+        dialects (Iterable[Union[Dialect, ModuleType]]): the dialects to include in the group.
+
+    Returns:
+        Callable[[RunPassGen[PassParams]], DialectGroup[PassParams]]: the dialect group.
+
+    Example:
+        ```python
+        from kirin.dialects import cf, fcf, func, math
+
+        @dialect_group([cf, fcf, func, math])
+        def basic_no_opt(self):
+            # initializations
+            def run_pass(mt: Method) -> None:
+                # how passes are applied to the method
+                pass
+
+            return run_pass
+        ```
+    """
+
+    # NOTE: do not alias the annotation below
+    def wrapper(
+        transform: RunPassGen[PassParams],
+    ) -> DialectGroup[PassParams]:
+        ret = DialectGroup(dialects, run_pass=transform)
+        update_wrapper(ret, transform)
+        return ret
+
+    return wrapper

kirin/ir/method.py

@@ -0,0 +1,118 @@
+import typing
+from types import ModuleType
+
+# from typing import TYPE_CHECKING, Generic, TypeVar, Callable, ParamSpec
+from dataclasses import field, dataclass
+
+from kirin.ir.traits import HasSignature, CallableStmtInterface
+from kirin.exceptions import VerificationError
+from kirin.ir.nodes.stmt import Statement
+from kirin.print.printer import Printer
+from kirin.ir.attrs.types import Generic
+from kirin.print.printable import Printable
+
+if typing.TYPE_CHECKING:
+    from kirin.ir.group import DialectGroup
+
+Param = typing.ParamSpec("Param")
+RetType = typing.TypeVar("RetType")
+
+
+@dataclass
+class Method(Printable, typing.Generic[Param, RetType]):
+    mod: ModuleType | None  # ref
+    py_func: typing.Callable[Param, RetType] | None  # ref
+    sym_name: str
+    arg_names: list[str]
+    dialects: "DialectGroup"  # own
+    code: Statement  # own, the corresponding IR, a func.func usually
+    # values contained if closure
+    fields: tuple = field(default_factory=tuple)  # own
+    file: str = ""
+    inferred: bool = False
+    """if typeinfer has been run on this method
+    """
+    verified: bool = False
+    """if `code.verify` has been run on this method
+    """
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def __call__(self, *args: Param.args, **kwargs: Param.kwargs) -> RetType:
+        from kirin.interp.concrete import Interpreter
+
+        if len(args) + len(kwargs) != len(self.arg_names) - 1:
+            raise ValueError("Incorrect number of arguments")
+        # NOTE: multi-return values will be wrapped in a tuple for Python
+        interp = Interpreter(self.dialects)
+        return interp.run(self, args=args, kwargs=kwargs).expect()
+
+    @property
+    def args(self):
+        """Return the arguments of the method. (excluding self)"""
+        return tuple(arg for arg in self.callable_region.blocks[0].args[1:])
+
+    @property
+    def arg_types(self):
+        """Return the types of the arguments of the method. (excluding self)"""
+        return tuple(arg.type for arg in self.args)
+
+    @property
+    def self_type(self):
+        """Return the type of the self argument of the method."""
+        trait = self.code.get_trait(HasSignature)
+        if trait is None:
+            raise ValueError("Method body must implement HasSignature")
+        signature = trait.get_signature(self.code)
+        return Generic(Method, Generic(tuple, *signature.inputs), signature.output)
+
+    @property
+    def callable_region(self):
+        trait = self.code.get_trait(CallableStmtInterface)
+        if trait is None:
+            raise ValueError("Method body must implement CallableStmtInterface")
+        return trait.get_callable_region(self.code)
+
+    @property
+    def return_type(self):
+        trait = self.code.get_trait(HasSignature)
+        if trait is None:
+            raise ValueError("Method body must implement HasSignature")
+        return trait.get_signature(self.code).output
+
+    def __repr__(self) -> str:
+        return f'Method("{self.sym_name}")'
+
+    def print_impl(self, printer: Printer) -> None:
+        return printer.print(self.code)
+
+    def verify(self) -> None:
+        """verify the method body."""
+        try:
+            self.code.verify()
+        except VerificationError as e:
+            msg = f'File "{self.file}"'
+            if isinstance(e.node, Statement):
+                if e.node.source:
+                    msg += f", line {e.node.source.lineno}"
+                msg += f", in {e.node.name}"
+
+            msg += f":\n    Verification failed for {self.sym_name}: {e.args[0]}"
+            raise Exception(msg) from e
+        self.verified = True
+        return
+
+    def similar(self, dialects: typing.Optional["DialectGroup"] = None):
+        return Method(
+            self.mod,
+            self.py_func,
+            self.sym_name,
+            self.arg_names,
+            dialects or self.dialects,
+            self.code.from_stmt(self.code, regions=[self.callable_region.clone()]),
+            self.fields,
+            self.file,
+            self.inferred,
+            self.verified,
+        )

kirin/ir/nodes/__init__.py

@@ -0,0 +1,7 @@
+"""Definition of Kirin's Intermediate Representation (IR) nodes.
+"""
+
+from kirin.ir.nodes.base import IRNode as IRNode
+from kirin.ir.nodes.stmt import Statement as Statement
+from kirin.ir.nodes.block import Block as Block
+from kirin.ir.nodes.region import Region as Region

kirin/ir/nodes/base.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, TypeVar, Iterator
+from dataclasses import dataclass
+
+from typing_extensions import Self
+
+from kirin.print import Printer, Printable
+from kirin.ir.ssa import SSAValue
+
+if TYPE_CHECKING:
+    from kirin.ir.nodes.stmt import Statement
+
+
+ParentType = TypeVar("ParentType", bound="IRNode")
+
+
+@dataclass
+class IRNode(Generic[ParentType], ABC, Printable):
+    """Base class for all IR nodes. All IR nodes are hashable and can be compared
+    for equality. The hash of an IR node is the same as the id of the object.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
+
+    def assert_parent(self, type_: type[IRNode], parent) -> None:
+        assert (
+            isinstance(parent, type_) or parent is None
+        ), f"Invalid parent, expect {type_} or None, got {type(parent)}"
+
+    @property
+    @abstractmethod
+    def parent_node(self) -> ParentType | None:
+        """Parent node of the current node."""
+        ...
+
+    @parent_node.setter
+    @abstractmethod
+    def parent_node(self, parent: ParentType | None) -> None: ...
+
+    def is_ancestor(self, op: IRNode) -> bool:
+        """Check if the given node is an ancestor of the current node."""
+        if op is self:
+            return True
+        if (parent := op.parent_node) is None:
+            return False
+        return self.is_ancestor(parent)
+
+    def get_root(self) -> IRNode:
+        """Get the root node of the current node."""
+        if (parent := self.parent_node) is None:
+            return self
+        return parent.get_root()
+
+    def is_equal(self, other: IRNode, context: dict = {}) -> bool:
+        """Check if the current node is equal to the other node.
+
+        Args:
+            other: The other node to compare.
+            context: The context to store the visited nodes. Defaults to {}.
+
+        Returns:
+            True if the nodes are equal, False otherwise.
+
+        !!! note
+            This method is not the same as the `==` operator. It checks for
+            structural equality rather than identity. To change the behavior
+            of structural equality, override the `is_structurally_equal` method.
+        """
+        if not isinstance(other, type(self)):
+            return False
+        return self.is_structurally_equal(other, context)
+
+    def attach(self, parent: ParentType) -> None:
+        """Attach the current node to the parent node."""
+        assert isinstance(parent, IRNode), f"Expected IRNode, got {type(parent)}"
+
+        if self.parent_node:
+            raise ValueError("Node already has a parent")
+        if self.is_ancestor(parent):
+            raise ValueError("Node is an ancestor of the parent")
+        self.parent_node = parent
+
+    @abstractmethod
+    def detach(self) -> None:
+        """Detach the current node from the parent node."""
+        ...
+
+    @abstractmethod
+    def drop_all_references(self) -> None:
+        """Drop all references to other nodes."""
+        ...
+
+    @abstractmethod
+    def delete(self, safe: bool = True) -> None:
+        """Delete the current node.
+
+        Args:
+            safe: If True, check if the node has any references before deleting.
+        """
+        ...
+
+    @abstractmethod
+    def is_structurally_equal(
+        self,
+        other: Self,
+        context: dict[IRNode | SSAValue, IRNode | SSAValue] | None = None,
+    ) -> bool:
+        """Check if the current node is structurally equal to the other node.
+
+        !!! note
+            This method is for tweaking the behavior of structural equality.
+            To check if two nodes are structurally equal, use the `is_equal` method.
+
+        Args:
+            other: The other node to compare.
+            context: The context to store the visited nodes.
+
+        Returns:
+            True if the nodes are structurally equal, False otherwise.
+        """
+        ...
+
+    def __eq__(self, other) -> bool:
+        return self is other
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    @abstractmethod
+    def walk(
+        self, *, reverse: bool = False, region_first: bool = False
+    ) -> Iterator[Statement]: ...
+
+    @abstractmethod
+    def print_impl(self, printer: Printer) -> None: ...
+
+    @abstractmethod
+    def typecheck(self) -> None:
+        """check if types are correct."""
+        ...
+
+    @abstractmethod
+    def verify(self) -> None:
+        """run mandatory validation checks. This is not same as typecheck, which may be optional."""
+        ...