PyPI - kirin-toolchain - Versions diffs - 0.13.0__py3-none-any.whl - Mend

kirin-toolchain 0.13.0__py3-none-any.whl

Files changed (225) hide show

kirin/__init__.py +7 -0
kirin/analysis/__init__.py +24 -0
kirin/analysis/callgraph.py +61 -0
kirin/analysis/cfg.py +112 -0
kirin/analysis/const/__init__.py +20 -0
kirin/analysis/const/_visitor.py +2 -0
kirin/analysis/const/_visitor.pyi +8 -0
kirin/analysis/const/lattice.py +219 -0
kirin/analysis/const/prop.py +116 -0
kirin/analysis/forward.py +100 -0
kirin/analysis/typeinfer/__init__.py +5 -0
kirin/analysis/typeinfer/analysis.py +90 -0
kirin/analysis/typeinfer/solve.py +141 -0
kirin/decl/__init__.py +108 -0
kirin/decl/base.py +65 -0
kirin/decl/camel2snake.py +2 -0
kirin/decl/emit/__init__.py +0 -0
kirin/decl/emit/_create_fn.py +29 -0
kirin/decl/emit/_set_new_attribute.py +22 -0
kirin/decl/emit/dialect.py +8 -0
kirin/decl/emit/init.py +277 -0
kirin/decl/emit/name.py +10 -0
kirin/decl/emit/property.py +182 -0
kirin/decl/emit/repr.py +31 -0
kirin/decl/emit/traits.py +13 -0
kirin/decl/emit/typecheck.py +77 -0
kirin/decl/emit/verify.py +51 -0
kirin/decl/info.py +346 -0
kirin/decl/scan_fields.py +157 -0
kirin/decl/verify.py +69 -0
kirin/dialects/__init__.py +14 -0
kirin/dialects/_pprint_helper.py +53 -0
kirin/dialects/cf/__init__.py +20 -0
kirin/dialects/cf/constprop.py +51 -0
kirin/dialects/cf/dialect.py +3 -0
kirin/dialects/cf/emit.py +58 -0
kirin/dialects/cf/interp.py +24 -0
kirin/dialects/cf/stmts.py +68 -0
kirin/dialects/cf/typeinfer.py +27 -0
kirin/dialects/eltype.py +23 -0
kirin/dialects/func/__init__.py +20 -0
kirin/dialects/func/attrs.py +39 -0
kirin/dialects/func/constprop.py +138 -0
kirin/dialects/func/dialect.py +3 -0
kirin/dialects/func/emit.py +80 -0
kirin/dialects/func/interp.py +68 -0
kirin/dialects/func/stmts.py +233 -0
kirin/dialects/func/typeinfer.py +124 -0
kirin/dialects/ilist/__init__.py +33 -0
kirin/dialects/ilist/_dialect.py +3 -0
kirin/dialects/ilist/_wrapper.py +51 -0
kirin/dialects/ilist/interp.py +85 -0
kirin/dialects/ilist/lowering.py +25 -0
kirin/dialects/ilist/passes.py +32 -0
kirin/dialects/ilist/rewrite/__init__.py +3 -0
kirin/dialects/ilist/rewrite/const.py +45 -0
kirin/dialects/ilist/rewrite/list.py +38 -0
kirin/dialects/ilist/rewrite/unroll.py +131 -0
kirin/dialects/ilist/runtime.py +63 -0
kirin/dialects/ilist/stmts.py +102 -0
kirin/dialects/ilist/typeinfer.py +120 -0
kirin/dialects/lowering/__init__.py +7 -0
kirin/dialects/lowering/call.py +48 -0
kirin/dialects/lowering/cf.py +206 -0
kirin/dialects/lowering/func.py +134 -0
kirin/dialects/math/__init__.py +41 -0
kirin/dialects/math/_gen.py +176 -0
kirin/dialects/math/dialect.py +3 -0
kirin/dialects/math/interp.py +190 -0
kirin/dialects/math/stmts.py +369 -0
kirin/dialects/module.py +139 -0
kirin/dialects/py/__init__.py +40 -0
kirin/dialects/py/assertion.py +91 -0
kirin/dialects/py/assign.py +103 -0
kirin/dialects/py/attr.py +59 -0
kirin/dialects/py/base.py +34 -0
kirin/dialects/py/binop/__init__.py +23 -0
kirin/dialects/py/binop/_dialect.py +3 -0
kirin/dialects/py/binop/interp.py +60 -0
kirin/dialects/py/binop/julia.py +33 -0
kirin/dialects/py/binop/lowering.py +22 -0
kirin/dialects/py/binop/stmts.py +79 -0
kirin/dialects/py/binop/typeinfer.py +108 -0
kirin/dialects/py/boolop.py +84 -0
kirin/dialects/py/builtin.py +78 -0
kirin/dialects/py/cmp/__init__.py +16 -0
kirin/dialects/py/cmp/_dialect.py +3 -0
kirin/dialects/py/cmp/interp.py +48 -0
kirin/dialects/py/cmp/julia.py +33 -0
kirin/dialects/py/cmp/lowering.py +45 -0
kirin/dialects/py/cmp/stmts.py +62 -0
kirin/dialects/py/constant.py +79 -0
kirin/dialects/py/indexing.py +251 -0
kirin/dialects/py/iterable.py +90 -0
kirin/dialects/py/len.py +57 -0
kirin/dialects/py/list/__init__.py +15 -0
kirin/dialects/py/list/_dialect.py +3 -0
kirin/dialects/py/list/interp.py +21 -0
kirin/dialects/py/list/lowering.py +25 -0
kirin/dialects/py/list/stmts.py +22 -0
kirin/dialects/py/list/typeinfer.py +54 -0
kirin/dialects/py/range.py +76 -0
kirin/dialects/py/slice.py +120 -0
kirin/dialects/py/tuple.py +109 -0
kirin/dialects/py/unary/__init__.py +24 -0
kirin/dialects/py/unary/_dialect.py +3 -0
kirin/dialects/py/unary/constprop.py +20 -0
kirin/dialects/py/unary/interp.py +24 -0
kirin/dialects/py/unary/julia.py +21 -0
kirin/dialects/py/unary/lowering.py +22 -0
kirin/dialects/py/unary/stmts.py +33 -0
kirin/dialects/py/unary/typeinfer.py +23 -0
kirin/dialects/py/unpack.py +90 -0
kirin/dialects/scf/__init__.py +23 -0
kirin/dialects/scf/_dialect.py +3 -0
kirin/dialects/scf/absint.py +64 -0
kirin/dialects/scf/constprop.py +140 -0
kirin/dialects/scf/interp.py +35 -0
kirin/dialects/scf/lowering.py +123 -0
kirin/dialects/scf/stmts.py +250 -0
kirin/dialects/scf/trim.py +36 -0
kirin/dialects/scf/typeinfer.py +58 -0
kirin/dialects/scf/unroll.py +92 -0
kirin/emit/__init__.py +3 -0
kirin/emit/abc.py +89 -0
kirin/emit/abc.pyi +38 -0
kirin/emit/exceptions.py +5 -0
kirin/emit/julia.py +63 -0
kirin/emit/str.py +51 -0
kirin/exceptions.py +59 -0
kirin/graph.py +34 -0
kirin/idtable.py +57 -0
kirin/interp/__init__.py +39 -0
kirin/interp/abstract.py +253 -0
kirin/interp/base.py +438 -0
kirin/interp/concrete.py +62 -0
kirin/interp/exceptions.py +26 -0
kirin/interp/frame.py +151 -0
kirin/interp/impl.py +197 -0
kirin/interp/result.py +93 -0
kirin/interp/state.py +71 -0
kirin/interp/table.py +40 -0
kirin/interp/value.py +73 -0
kirin/ir/__init__.py +46 -0
kirin/ir/attrs/__init__.py +20 -0
kirin/ir/attrs/_types.py +8 -0
kirin/ir/attrs/_types.pyi +13 -0
kirin/ir/attrs/abc.py +46 -0
kirin/ir/attrs/py.py +45 -0
kirin/ir/attrs/types.py +522 -0
kirin/ir/dialect.py +125 -0
kirin/ir/group.py +249 -0
kirin/ir/method.py +118 -0
kirin/ir/nodes/__init__.py +7 -0
kirin/ir/nodes/base.py +149 -0
kirin/ir/nodes/block.py +458 -0
kirin/ir/nodes/region.py +337 -0
kirin/ir/nodes/stmt.py +713 -0
kirin/ir/nodes/view.py +142 -0
kirin/ir/ssa.py +204 -0
kirin/ir/traits/__init__.py +36 -0
kirin/ir/traits/abc.py +42 -0
kirin/ir/traits/basic.py +78 -0
kirin/ir/traits/callable.py +51 -0
kirin/ir/traits/lowering/__init__.py +2 -0
kirin/ir/traits/lowering/call.py +37 -0
kirin/ir/traits/lowering/context.py +120 -0
kirin/ir/traits/region/__init__.py +2 -0
kirin/ir/traits/region/ssacfg.py +22 -0
kirin/ir/traits/symbol.py +57 -0
kirin/ir/use.py +17 -0
kirin/lattice/__init__.py +13 -0
kirin/lattice/abc.py +128 -0
kirin/lattice/empty.py +25 -0
kirin/lattice/mixin.py +51 -0
kirin/lowering/__init__.py +7 -0
kirin/lowering/binding.py +65 -0
kirin/lowering/core.py +72 -0
kirin/lowering/dialect.py +35 -0
kirin/lowering/dialect.pyi +183 -0
kirin/lowering/frame.py +171 -0
kirin/lowering/result.py +68 -0
kirin/lowering/state.py +441 -0
kirin/lowering/stream.py +53 -0
kirin/passes/__init__.py +3 -0
kirin/passes/abc.py +44 -0
kirin/passes/aggressive/__init__.py +1 -0
kirin/passes/aggressive/fold.py +43 -0
kirin/passes/fold.py +45 -0
kirin/passes/inline.py +25 -0
kirin/passes/typeinfer.py +25 -0
kirin/prelude.py +197 -0
kirin/print/__init__.py +15 -0
kirin/print/printable.py +141 -0
kirin/print/printer.py +415 -0
kirin/py.typed +0 -0
kirin/registry.py +105 -0
kirin/registry.pyi +52 -0
kirin/rewrite/__init__.py +14 -0
kirin/rewrite/abc.py +43 -0
kirin/rewrite/aggressive/__init__.py +1 -0
kirin/rewrite/aggressive/fold.py +43 -0
kirin/rewrite/alias.py +16 -0
kirin/rewrite/apply_type.py +47 -0
kirin/rewrite/call2invoke.py +34 -0
kirin/rewrite/chain.py +39 -0
kirin/rewrite/compactify.py +288 -0
kirin/rewrite/cse.py +48 -0
kirin/rewrite/dce.py +19 -0
kirin/rewrite/fixpoint.py +34 -0
kirin/rewrite/fold.py +57 -0
kirin/rewrite/getfield.py +21 -0
kirin/rewrite/getitem.py +37 -0
kirin/rewrite/inline.py +143 -0
kirin/rewrite/result.py +15 -0
kirin/rewrite/walk.py +83 -0
kirin/rewrite/wrap_const.py +55 -0
kirin/source.py +21 -0
kirin/symbol_table.py +27 -0
kirin/types.py +34 -0
kirin/worklist.py +30 -0
kirin_toolchain-0.13.0.dist-info/METADATA +42 -0
kirin_toolchain-0.13.0.dist-info/RECORD +225 -0
kirin_toolchain-0.13.0.dist-info/WHEEL +4 -0
kirin_toolchain-0.13.0.dist-info/licenses/LICENSE +234 -0

kirin/ir/traits/lowering/context.py ADDED Viewed

@@ -0,0 +1,120 @@
+"""Traits for customizing lowering of Python `with` syntax to a statement.
+"""
+import ast
+from typing import TYPE_CHECKING, TypeVar
+from dataclasses import dataclass
+from kirin.exceptions import DialectLoweringError
+from ..abc import PythonLoweringTrait
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+    from kirin.lowering import Result, LoweringState
+StatementType = TypeVar("StatementType", bound="Statement")
+@dataclass(frozen=True)
+class FromPythonWith(PythonLoweringTrait[StatementType, ast.With]):
+    """Trait for customizing lowering of Python with statements to a statement.
+    Subclassing this trait allows for customizing the lowering of Python with
+    statements to the statement. The `lower` method should be implemented to parse
+    the arguments from the Python with statement and construct the statement instance.
+    """
+    pass
+@dataclass(frozen=True)
+class FromPythonWithSingleItem(FromPythonWith[StatementType]):
+    """Trait for customizing lowering of the following Python with syntax to a statement:
+    ```python
+    with <stmt>[ as <name>]:
+        <body>
+    ```
+    where `<stmt>` is the statement being lowered, `<name>` is an optional name for the result
+    of the statement, and `<body>` is the body of the with statement. The optional `as <name>`
+    is not valid when the statement has no results.
+    This syntax is slightly different from the standard Python `with` statement in that
+    `<name>` refers to the result of the statement, not the context manager. Thus typically
+    one sould access `<name>` in `<body>` to use the result of the statement.
+    In some cases, however, `<name>` may be used as a reference of a special value `self` that
+    is passed to the `<body>` of the statement. This is useful for statements that have a similar
+    behavior to a closure.
+    """
+    def lower(
+        self, stmt: type[StatementType], state: "LoweringState", node: ast.With
+    ) -> "Result":
+        from kirin import ir, lowering
+        from kirin.decl import fields
+        from kirin.dialects import cf
+        fs = fields(stmt)
+        if len(fs.regions) != 1:
+            raise DialectLoweringError(
+                "Expected exactly one region in statement declaration"
+            )
+        if len(node.items) != 1:
+            raise DialectLoweringError("Expected exactly one item in statement")
+        item, body = node.items[0], node.body
+        if not isinstance(item.context_expr, ast.Call):
+            raise DialectLoweringError(
+                f"Expected context expression to be a call for with {stmt.name}"
+            )
+        body_frame = lowering.Frame.from_stmts(body, state, parent=state.current_frame)
+        state.push_frame(body_frame)
+        state.exhaust()
+        region_name, region_info = next(iter(fs.regions.items()))
+        if region_info.multi:  # branch to exit block if not terminated
+            for block in body_frame.curr_region.blocks:
+                if block.last_stmt is None or not block.last_stmt.has_trait(
+                    ir.IsTerminator
+                ):
+                    block.stmts.append(
+                        cf.Branch(arguments=(), successor=body_frame.next_block)
+                    )
+            state.pop_frame()
+        else:
+            if len(body_frame.curr_region.blocks) != 1:
+                raise DialectLoweringError(
+                    f"Expected exactly one block in region {region_name}"
+                )
+            state.pop_frame(finalize_next=False)
+        args, kwargs = state.default_Call_inputs(stmt, item.context_expr)
+        kwargs[region_name] = body_frame.curr_region
+        results = state.append_stmt(stmt(*args.values(), **kwargs)).results
+        if len(results) == 0:
+            return lowering.Result()
+        elif len(results) > 1:
+            raise DialectLoweringError(
+                f"Expected exactly one result or no result from statement {stmt.name}"
+            )
+        result = results[0]
+        if item.optional_vars is not None and isinstance(item.optional_vars, ast.Name):
+            result.name = item.optional_vars.id
+            state.current_frame.defs[result.name] = result
+        return lowering.Result(result)
+    def verify(self, stmt: "Statement"):
+        assert (
+            len(stmt.regions) == 1
+        ), "FromPythonWithSingleItem statements must have one region"
+        assert (
+            len(stmt.successors) == 0
+        ), "FromPythonWithSingleItem statements cannot have successors"
+        assert (
+            len(stmt.results) <= 1
+        ), "FromPythonWithSingleItem statements can have at most one result"

kirin/ir/traits/region/__init__.py ADDED Viewed

	@@ -0,0 +1,2 @@
1	+ """Builtin region traits.
2	+ """

kirin/ir/traits/region/ssacfg.py ADDED Viewed

@@ -0,0 +1,22 @@
+"""SSACFG region trait.
+This module defines the SSACFGRegion trait, which is used to indicate that a
+region has an SSACFG graph.
+"""
+from typing import TYPE_CHECKING
+from dataclasses import dataclass
+from kirin.ir.traits.abc import RegionTrait
+if TYPE_CHECKING:
+    from kirin.ir import Region
+@dataclass(frozen=True)
+class SSACFGRegion(RegionTrait):
+    def get_graph(self, region: "Region"):
+        from kirin.analysis.cfg import CFG
+        return CFG(region)

kirin/ir/traits/symbol.py ADDED Viewed

@@ -0,0 +1,57 @@
+from typing import TYPE_CHECKING
+from dataclasses import dataclass
+from kirin.exceptions import VerificationError
+from kirin.ir.attrs.py import PyAttr
+from kirin.ir.traits.abc import StmtTrait
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+@dataclass(frozen=True)
+class SymbolOpInterface(StmtTrait):
+    """A trait that indicates that a statement is a symbol operation.
+    A symbol operation is a statement that has a symbol name attribute.
+    """
+    def get_sym_name(self, stmt: "Statement") -> "PyAttr[str]":
+        sym_name: PyAttr[str] | None = stmt.get_attr_or_prop("sym_name")  # type: ignore
+        # NOTE: unlike MLIR or xDSL we do not allow empty symbol names
+        if sym_name is None:
+            raise ValueError(f"Statement {stmt.name} does not have a symbol name")
+        return sym_name
+    def verify(self, stmt: "Statement"):
+        from kirin.types import String
+        sym_name = self.get_sym_name(stmt)
+        if not (isinstance(sym_name, PyAttr) and sym_name.type.is_subseteq(String)):
+            raise ValueError(f"Symbol name {sym_name} is not a string attribute")
+@dataclass(frozen=True)
+class SymbolTable(StmtTrait):
+    """
+    Statement with SymbolTable trait can only have one region with one block.
+    """
+    @staticmethod
+    def walk(stmt: "Statement"):
+        return stmt.regions[0].blocks[0].stmts
+    def verify(self, stmt: "Statement"):
+        if len(stmt.regions) != 1:
+            raise VerificationError(
+                stmt,
+                f"Statement {stmt.name} with SymbolTable trait must have exactly one region",
+            )
+        if len(stmt.regions[0].blocks) != 1:
+            raise VerificationError(
+                stmt,
+                f"Statement {stmt.name} with SymbolTable trait must have exactly one block",
+            )
+        # TODO: check uniqueness of symbol names

kirin/ir/use.py ADDED Viewed

@@ -0,0 +1,17 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from dataclasses import dataclass
+if TYPE_CHECKING:
+    from kirin.ir.nodes.stmt import Statement
+@dataclass(frozen=True)
+class Use:
+    """A use of an SSA value in a statement."""
+    stmt: Statement
+    """The statement that uses the SSA value."""
+    index: int
+    """The index of the use in the statement."""

kirin/lattice/__init__.py ADDED Viewed

@@ -0,0 +1,13 @@
+from kirin.lattice.abc import (
+    Lattice as Lattice,
+    UnionMeta as UnionMeta,
+    LatticeMeta as LatticeMeta,
+    SingletonMeta as SingletonMeta,
+    BoundedLattice as BoundedLattice,
+)
+from kirin.lattice.empty import EmptyLattice as EmptyLattice
+from kirin.lattice.mixin import (
+    IsSubsetEqMixin as IsSubsetEqMixin,
+    SimpleJoinMixin as SimpleJoinMixin,
+    SimpleMeetMixin as SimpleMeetMixin,
+)

kirin/lattice/abc.py ADDED Viewed

@@ -0,0 +1,128 @@
+from abc import ABC, ABCMeta, abstractmethod
+from typing import Generic, TypeVar, Iterable
+class LatticeMeta(ABCMeta):
+    pass
+class SingletonMeta(LatticeMeta):
+    """
+    Singleton metaclass for lattices. It ensures that only one instance of a lattice is created.
+    See https://stackoverflow.com/questions/674304/why-is-init-always-called-after-new/8665179#8665179
+    """
+    def __init__(cls, name, bases, attrs):
+        super().__init__(name, bases, attrs)
+        cls._instance = None
+    def __call__(cls):
+        if cls._instance is None:
+            cls._instance = super().__call__()
+        return cls._instance
+LatticeType = TypeVar("LatticeType", bound="Lattice")
+class Lattice(ABC, Generic[LatticeType], metaclass=LatticeMeta):
+    """ABC for lattices as Python class.
+    While `Lattice` is only an interface, `LatticeABC` is an abstract
+    class that can be inherited from. This provides a few default
+    implementations for the lattice operations.
+    """
+    @abstractmethod
+    def join(self, other: LatticeType) -> LatticeType:
+        """Join operation."""
+        ...
+    @abstractmethod
+    def meet(self, other: LatticeType) -> LatticeType:
+        """Meet operation."""
+        ...
+    @abstractmethod
+    def is_subseteq(self, other: LatticeType) -> bool:
+        """Subseteq operation."""
+        ...
+    def is_equal(self, other: LatticeType) -> bool:
+        """Check if two lattices are equal."""
+        if self is other:
+            return True
+        else:
+            return self.is_subseteq(other) and other.is_subseteq(self)
+    def is_subset(self, other: LatticeType) -> bool:
+        return self.is_subseteq(other) and not other.is_subseteq(self)
+    def __eq__(self, value: object) -> bool:
+        raise NotImplementedError(
+            "Equality is not implemented for lattices, use is_equal instead"
+        )
+    def __hash__(self) -> int:
+        raise NotImplementedError("Hash is not implemented for lattices")
+BoundedLatticeType = TypeVar("BoundedLatticeType", bound="BoundedLattice")
+class BoundedLattice(Lattice[BoundedLatticeType]):
+    """ABC for bounded lattices as Python class.
+    `BoundedLattice` is an abstract class that can be inherited from.
+    It requires the implementation of the `bottom` and `top` methods.
+    """
+    @classmethod
+    @abstractmethod
+    def bottom(cls) -> BoundedLatticeType: ...
+    @classmethod
+    @abstractmethod
+    def top(cls) -> BoundedLatticeType: ...
+class UnionMeta(LatticeMeta):
+    """Meta class for union types. It simplifies the union if possible."""
+    def __call__(
+        self,
+        typ: Iterable[LatticeType] | LatticeType,
+        *others: LatticeType,
+    ):
+        from kirin.lattice.abc import Lattice
+        if isinstance(typ, Lattice):
+            typs: Iterable[LatticeType] = (typ, *others)
+        elif not others:
+            typs = typ
+        else:
+            raise ValueError(
+                "Expected an iterable of types or variadic arguments of types"
+            )
+        # try if the union can be simplified
+        params: list[LatticeType] = []
+        for typ in typs:
+            contains = False
+            for idx, other in enumerate(params):
+                if typ.is_subseteq(other):
+                    contains = True
+                    break
+                elif other.is_subseteq(typ):
+                    params[idx] = typ
+                    contains = True
+                    break
+            if not contains:
+                params.append(typ)
+        if len(params) == 1:
+            return params[0]
+        return super(UnionMeta, self).__call__(*params)

kirin/lattice/empty.py ADDED Viewed

@@ -0,0 +1,25 @@
+from kirin.lattice.abc import SingletonMeta, BoundedLattice
+class EmptyLattice(BoundedLattice["EmptyLattice"], metaclass=SingletonMeta):
+    """Empty lattice."""
+    def join(self, other: "EmptyLattice") -> "EmptyLattice":
+        return self
+    def meet(self, other: "EmptyLattice") -> "EmptyLattice":
+        return self
+    @classmethod
+    def bottom(cls):
+        return cls()
+    @classmethod
+    def top(cls):
+        return cls()
+    def __hash__(self) -> int:
+        return id(self)
+    def is_subseteq(self, other: "EmptyLattice") -> bool:
+        return True

kirin/lattice/mixin.py ADDED Viewed

@@ -0,0 +1,51 @@
+from typing import TypeVar
+from .abc import BoundedLattice
+BoundedLatticeType = TypeVar("BoundedLatticeType", bound="BoundedLattice")
+class IsSubsetEqMixin(BoundedLattice[BoundedLatticeType]):
+    """A special mixin for lattices that provides a default implementation for `is_subseteq`
+    by using the visitor pattern. This is useful if the lattice has a lot of different
+    subclasses that need to be compared.
+    Must be used before `BoundedLattice` in the inheritance chain.
+    """
+    def is_subseteq(self, other: BoundedLatticeType) -> bool:
+        if other is self.top():
+            return True
+        elif other is self.bottom():
+            return False
+        method = getattr(
+            self,
+            "is_subseteq_" + other.__class__.__name__,
+            getattr(self, "is_subseteq_fallback", None),
+        )
+        if method is not None:
+            return method(other)
+        return False
+class SimpleJoinMixin(BoundedLattice[BoundedLatticeType]):
+    """A mixin that provides a simple implementation for the join operation."""
+    def join(self, other: BoundedLatticeType) -> BoundedLatticeType:
+        if self.is_subseteq(other):
+            return other
+        elif other.is_subseteq(self):
+            return self  # type: ignore
+        return self.top()
+class SimpleMeetMixin(BoundedLattice[BoundedLatticeType]):
+    """A mixin that provides a simple implementation for the meet operation."""
+    def meet(self, other: BoundedLatticeType) -> BoundedLatticeType:
+        if self.is_subseteq(other):
+            return self  # type: ignore
+        elif other.is_subseteq(self):
+            return other
+        return self.bottom()

kirin/lowering/__init__.py ADDED Viewed

@@ -0,0 +1,7 @@
+from kirin.lowering.core import Lowering as Lowering
+from kirin.lowering.frame import Frame as Frame
+from kirin.lowering.state import LoweringState as LoweringState
+from kirin.lowering.result import Result as Result
+from kirin.lowering.stream import StmtStream as StmtStream
+from kirin.lowering.binding import wraps as wraps
+from kirin.lowering.dialect import FromPythonAST as FromPythonAST

kirin/lowering/binding.py ADDED Viewed

@@ -0,0 +1,65 @@
+from typing import TYPE_CHECKING, Generic, TypeVar, Callable, ParamSpec
+from dataclasses import dataclass
+if TYPE_CHECKING:
+    from kirin.ir.nodes.stmt import Statement
+Params = ParamSpec("Params")
+RetType = TypeVar("RetType")
+@dataclass(frozen=True)
+class Binding(Generic[Params, RetType]):
+    parent: type["Statement"]
+    def __call__(self, *args: Params.args, **kwargs: Params.kwargs) -> RetType:
+        raise NotImplementedError(
+            f"Binding of {self.parent.name} can \
+            only be called from a kernel"
+        )
+def wraps(parent: type["Statement"]):
+    """Wraps a [`Statement`][kirin.ir.nodes.stmt.Statement] to a `Binding` object
+    which will be special cased in the lowering process.
+    This is useful for providing type hints by faking the call signature of a
+    [`Statement`][kirin.ir.nodes.stmt.Statement].
+    ## Example
+    Directly writing a function with the statement will let Python linter think
+    you intend to call the constructor of the statement class. However, given the
+    context of a kernel, our intention is to actually "call" the statement, e.g
+    the following will produce type errors with pyright or mypy:
+    ```python
+    from kirin.dialects import math
+    from kirin.prelude import basic_no_opt
+    @basic_no_opt
+    def main(x: float):
+        return math.sin(x) # this is a statement, not a function
+    ```
+    the `@lowering.wraps` decorator allows us to provide a type hint for the
+    statement, e.g:
+    ```python
+    from kirin import lowering
+    @lowering.wraps(math.sin)
+    def sin(value: float) -> float: ...
+    @basic_no_opt
+    def main(x: float):
+        return sin(x) # linter now thinks this is a function
+    sin(1.0) # this will raise a NotImplementedError("Binding of sin can only be called from a kernel")
+    ```
+    """
+    def wrapper(func: Callable[Params, RetType]) -> Binding[Params, RetType]:
+        return Binding(parent)
+    return wrapper

kirin/lowering/core.py ADDED Viewed

@@ -0,0 +1,72 @@
+import ast
+import inspect
+import textwrap
+from types import ModuleType
+from typing import Any, Callable, Iterable
+from dataclasses import dataclass
+from kirin.ir import Dialect, DialectGroup
+from kirin.exceptions import DialectLoweringError
+from kirin.lowering.state import LoweringState
+from kirin.lowering.dialect import FromPythonAST
+@dataclass(init=False)
+class Lowering(ast.NodeVisitor):
+    dialects: DialectGroup
+    registry: dict[str, FromPythonAST]
+    state: LoweringState | None = None
+    # max lines to show in error hint
+    max_lines: int = 3
+    def __init__(
+        self,
+        dialects: DialectGroup | Iterable[Dialect | ModuleType],
+        keys: list[str] | None = None,
+        max_lines: int = 3,
+    ):
+        if isinstance(dialects, DialectGroup):
+            self.dialects = dialects
+        else:
+            self.dialects = DialectGroup(dialects)
+        self.max_lines = max_lines
+        self.registry: dict[str, FromPythonAST] = self.dialects.registry.ast(
+            keys=keys or ["main", "default"]
+        )
+        self.state = None
+    def run(
+        self,
+        stmt: ast.stmt | Callable,
+        source: str | None = None,
+        globals: dict[str, Any] | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+    ):
+        if isinstance(stmt, Callable):
+            source = source or textwrap.dedent(inspect.getsource(stmt))
+            globals = globals or stmt.__globals__
+            try:
+                nonlocals = inspect.getclosurevars(stmt).nonlocals
+            except Exception:
+                nonlocals = {}
+            globals.update(nonlocals)
+            stmt = ast.parse(source).body[0]
+        state = LoweringState.from_stmt(
+            self, stmt, source, globals, self.max_lines, lineno_offset, col_offset
+        )
+        try:
+            state.visit(stmt)
+        except DialectLoweringError as e:
+            e.args = (f"{e.args[0]}\n\n{state.error_hint()}",) + e.args[1:]
+            raise e
+        if compactify:
+            from kirin.rewrite import Walk, CFGCompactify
+            Walk(CFGCompactify()).rewrite(state.code)
+        return state.code

kirin/lowering/dialect.py ADDED Viewed

@@ -0,0 +1,35 @@
+# NOTE: this module is only interface, will be used inside
+# the `ir` module try to minimize the dependencies as much
+# as possible
+from __future__ import annotations
+import ast
+from abc import ABC
+from typing import TYPE_CHECKING
+from kirin.exceptions import DialectLoweringError
+from kirin.lowering.result import Result
+if TYPE_CHECKING:
+    from kirin.lowering.state import LoweringState
+class FromPythonAST(ABC):
+    @property
+    def names(self) -> list[str]:  # show the name without lower_
+        return [name[6:] for name in dir(self) if name.startswith("lower_")]
+    def lower(self, state: LoweringState, node: ast.AST) -> Result:
+        """Entry point of dialect specific lowering."""
+        return getattr(self, f"lower_{node.__class__.__name__}", self.unreachable)(
+            state, node
+        )
+    def unreachable(self, state: LoweringState, node: ast.AST) -> Result:
+        raise DialectLoweringError(f"unreachable reached for {node.__class__.__name__}")
+class NoSpecialLowering(FromPythonAST):
+    pass