kirin-toolchain 0.13.0__py3-none-any.whl

kirin/analysis/typeinfer/analysis.py

@@ -0,0 +1,90 @@
+from typing import TypeVar, final
+
+from kirin import ir, types, interp
+from kirin.decl import fields
+from kirin.analysis import const
+from kirin.interp.impl import Signature
+from kirin.analysis.forward import Forward, ForwardFrame
+
+from .solve import TypeResolution
+
+
+@final
+class TypeInference(Forward[types.TypeAttribute]):
+    """Type inference analysis for kirin.
+
+    This analysis uses the forward dataflow analysis framework to infer the types of
+    the IR. The analysis uses the type information within the IR to determine the
+    method dispatch.
+
+    The analysis will fallback to a type resolution algorithm if the type information
+    is not available in the IR but the type information is available in the abstract
+    values.
+    """
+
+    keys = ["typeinfer"]
+    lattice = types.TypeAttribute
+
+    def run_analysis(
+        self, method: ir.Method, args: tuple[types.TypeAttribute, ...] | None = None
+    ) -> tuple[ForwardFrame[types.TypeAttribute], types.TypeAttribute]:
+        if args is None:
+            args = method.arg_types
+        return super().run_analysis(method, args)
+
+    # NOTE: unlike concrete interpreter, instead of using type information
+    # within the IR. Type inference will use the interpreted
+    # value (which is a type) to determine the method dispatch.
+    def build_signature(
+        self, frame: ForwardFrame[types.TypeAttribute], stmt: ir.Statement
+    ) -> Signature:
+        _args = ()
+        for x in frame.get_values(stmt.args):
+            # TODO: remove this after we have multiple dispatch...
+            if isinstance(x, types.Generic):
+                _args += (x.body,)
+            else:
+                _args += (x,)
+        return Signature(stmt.__class__, _args)
+
+    def eval_stmt_fallback(
+        self, frame: ForwardFrame[types.TypeAttribute], stmt: ir.Statement
+    ) -> tuple[types.TypeAttribute, ...] | interp.SpecialValue[types.TypeAttribute]:
+        resolve = TypeResolution()
+        fs = fields(stmt)
+        for f, value in zip(fs.args.values(), frame.get_values(stmt.args)):
+            resolve.solve(f.type, value)
+
+        for arg, f in zip(stmt.args, fs.args.values()):
+            frame.set(arg, frame.get(arg).meet(resolve.substitute(f.type)))
+        return tuple(resolve.substitute(result.type) for result in stmt.results)
+
+    def run_method(
+        self, method: ir.Method, args: tuple[types.TypeAttribute, ...]
+    ) -> tuple[ForwardFrame[types.TypeAttribute], types.TypeAttribute]:
+        return self.run_callable(method.code, (method.self_type,) + args)
+
+    T = TypeVar("T")
+
+    @classmethod
+    def maybe_const(cls, value: ir.SSAValue, type_: type[T]) -> T | None:
+        """Get a constant value of a given type.
+
+        If the value is not a constant or the constant is not of the given type, return
+        `None`.
+        """
+        hint = value.hints.get("const")
+        if isinstance(hint, const.Value) and isinstance(hint.data, type_):
+            return hint.data
+
+    @classmethod
+    def expect_const(cls, value: ir.SSAValue, type_: type[T]):
+        """Expect a constant value of a given type.
+
+        If the value is not a constant or the constant is not of the given type, raise
+        an `InterpreterError`.
+        """
+        hint = cls.maybe_const(value, type_)
+        if hint is None:
+            raise interp.InterpreterError(f"expected {type_}, got {hint}")
+        return hint

kirin/analysis/typeinfer/solve.py

@@ -0,0 +1,141 @@
+"""Type resolution for type inference.
+
+This module contains the type resolution algorithm for type inference.
+A simple algorithm is used to resolve the types of the IR by comparing
+the input types with the output types.
+"""
+
+from dataclasses import field, dataclass
+
+from kirin import types
+
+
+@dataclass
+class TypeResolutionResult:
+    """Base class for type resolution results."""
+
+    pass
+
+
+@dataclass
+class ResolutionOk(TypeResolutionResult):
+    """Type resolution result for successful resolution."""
+
+    def __bool__(self):
+        return True
+
+
+Ok = ResolutionOk()
+
+
+@dataclass
+class ResolutionError(TypeResolutionResult):
+    """Type resolution result for failed resolution."""
+
+    expr: types.TypeAttribute
+    value: types.TypeAttribute
+
+    def __bool__(self):
+        return False
+
+    def __str__(self):
+        return f"expected {self.expr}, got {self.value}"
+
+
+@dataclass
+class TypeResolution:
+    """Type resolution algorithm for type inference."""
+
+    vars: dict[types.TypeVar, types.TypeAttribute] = field(default_factory=dict)
+
+    def substitute(self, typ: types.TypeAttribute) -> types.TypeAttribute:
+        """Substitute type variables in the type with their values.
+
+        This method substitutes type variables in the given type with their
+        values. If the type is a generic type, the method recursively
+        substitutes the type variables in the type arguments.
+
+        Args:
+            typ: The type to substitute.
+
+        Returns:
+            The type with the type variables substituted.
+        """
+        if isinstance(typ, types.TypeVar):
+            return self.vars.get(typ, typ)
+        elif isinstance(typ, types.Generic):
+            return types.Generic(
+                typ.body, *tuple(self.substitute(var) for var in typ.vars)
+            )
+        elif isinstance(typ, types.Union):
+            return types.Union(self.substitute(t) for t in typ.types)
+        return typ
+
+    def solve(
+        self, annot: types.TypeAttribute, value: types.TypeAttribute
+    ) -> TypeResolutionResult:
+        """Solve the type resolution problem.
+
+        This method compares the expected type `annot` with the actual
+        type `value` and returns a result indicating whether the types
+        match or not.
+
+        Args:
+            annot: The expected type.
+            value: The actual type.
+
+        Returns:
+            A `TypeResolutionResult` object indicating the result of the
+            resolution.
+        """
+        if isinstance(annot, types.TypeVar):
+            return self.solve_TypeVar(annot, value)
+        elif isinstance(annot, types.Generic):
+            return self.solve_Generic(annot, value)
+        elif isinstance(annot, types.Union):
+            return self.solve_Union(annot, value)
+
+        if annot.is_subseteq(value):
+            return Ok
+        else:
+            return ResolutionError(annot, value)
+
+    def solve_TypeVar(self, annot: types.TypeVar, value: types.TypeAttribute):
+        if annot in self.vars:
+            if value.is_subseteq(self.vars[annot]):
+                self.vars[annot] = value
+            elif self.vars[annot].is_subseteq(value):
+                pass
+            else:
+                return ResolutionError(annot, value)
+        else:
+            self.vars[annot] = value
+        return Ok
+
+    def solve_Generic(self, annot: types.Generic, value: types.TypeAttribute):
+        if not isinstance(value, types.Generic):
+            return ResolutionError(annot, value)
+
+        if not value.body.is_subseteq(annot.body):
+            return ResolutionError(annot.body, value.body)
+
+        for var, val in zip(annot.vars, value.vars):
+            result = self.solve(var, val)
+            if not result:
+                return result
+
+        if not annot.vararg:
+            return Ok
+
+        for val in value.vars[len(annot.vars) :]:
+            result = self.solve(annot.vararg.typ, val)
+            if not result:
+                return result
+        return Ok
+
+    def solve_Union(self, annot: types.Union, value: types.TypeAttribute):
+        for typ in annot.types:
+            result = self.solve(typ, value)
+            if result:
+                return Ok
+        return ResolutionError(annot, value)

kirin/decl/__init__.py

@@ -0,0 +1,108 @@
+from typing import TypeVar, Callable
+
+from typing_extensions import Unpack, dataclass_transform
+
+from kirin.ir import Statement
+from kirin.decl import info
+from kirin.decl.base import StatementOptions
+from kirin.decl.verify import Verify
+from kirin.decl.emit.init import EmitInit
+from kirin.decl.emit.name import EmitName
+from kirin.decl.emit.repr import EmitRepr
+from kirin.decl.emit.traits import EmitTraits
+from kirin.decl.emit.verify import EmitVerify
+from kirin.decl.scan_fields import ScanFields
+from kirin.decl.emit.dialect import EmitDialect
+from kirin.decl.emit.property import EmitProperty
+from kirin.decl.emit.typecheck import EmitTypeCheck
+
+
+class StatementDecl(
+    ScanFields,
+    Verify,
+    EmitInit,
+    EmitProperty,
+    EmitDialect,
+    EmitName,
+    EmitRepr,
+    EmitTraits,
+    EmitVerify,
+    EmitTypeCheck,
+):
+    pass
+
+
+StmtType = TypeVar("StmtType", bound=Statement)
+
+
+@dataclass_transform(
+    field_specifiers=(
+        info.attribute,
+        info.argument,
+        info.region,
+        info.result,
+        info.block,
+    )
+)
+def statement(
+    cls=None,
+    **kwargs: Unpack[StatementOptions],
+) -> Callable[[type[StmtType]], type[StmtType]]:
+    """Declare a new statement class.
+
+    This decorator is used to declare a new statement class. It is used to
+    generate the necessary boilerplate code for the class. The class should
+    inherit from `kirin.ir.Statement`.
+
+    Args:
+        init(bool): Whether to generate an `__init__` method.
+        repr(bool): Whether to generate a `__repr__` method.
+        kw_only(bool): Whether to use keyword-only arguments in the `__init__`
+            method.
+        dialect(Optional[Dialect]): The dialect of the statement.
+        property(bool): Whether to generate property methods for attributes.
+
+    Example:
+        The following is an example of how to use the `statement` decorator.
+
+        ```python
+        @statement
+        class MyStatement(ir.Statement):
+            name = "some_name"
+            traits = frozenset({TraitA(), TraitB()})
+            some_input: ir.SSAValue = info.argument()
+            some_output: ir.ResultValue = info.result()
+            body: ir.Region = info.region()
+            successor: ir.Block = info.block()
+        ```
+
+        If the `name` field is not specified, a lowercase name field will be auto generated.
+
+        In addition, one can optionally register the statement to a dialect
+        by providing the `dialect` argument to the decorator.
+
+        The following example register the statement to a dialect `my_dialect_object`, and `name = "myfoo"` field is autogenerated
+
+        ```python
+        @statement(dialect=my_dialect_object)
+        class MyFoo(ir.Statement):
+            traits = frozenset({ir.FromPythonCall()})
+            value: str = info.attribute()
+        ```
+    """
+
+    def wrap(cls):
+        decl = StatementDecl(cls, **kwargs)
+        decl.scan_fields()
+        decl.verify()
+        decl.emit()
+        decl.register()
+        return cls
+
+    if cls is None:
+        return wrap
+    return wrap(cls)
+
+
+def fields(cls: type[Statement] | Statement) -> info.StatementFields:
+    return getattr(cls, ScanFields._FIELDS)

kirin/decl/base.py

@@ -0,0 +1,65 @@
+import sys
+import inspect
+from typing import Any, TypedDict
+
+from typing_extensions import Unpack, Optional
+
+from kirin.ir import Dialect
+from kirin.decl.info import StatementFields
+
+
+class StatementOptions(TypedDict, total=False):
+    init: bool
+    repr: bool
+    kw_only: bool
+    dialect: Optional[Dialect]
+    property: bool
+
+
+class BaseModifier:
+    _PARAMS = "__kirin_stmt_params"
+
+    def __init__(self, cls: type, **kwargs: Unpack[StatementOptions]) -> None:
+        self.cls = cls
+        self.cls_module = sys.modules.get(cls.__module__)
+
+        if "dialect" in kwargs:
+            self.dialect = kwargs["dialect"]
+        else:
+            self.dialect = None
+        self.params = kwargs
+        setattr(cls, self._PARAMS, self.params)
+
+        if cls.__module__ in sys.modules:
+            self.globals = sys.modules[cls.__module__].__dict__
+        else:
+            # Theoretically this can happen if someone writes
+            # a custom string to cls.__module__.  In which case
+            # such dataclass won't be fully introspectable
+            # (w.r.t. typing.get_type_hints) but will still function
+            # correctly.
+            self.globals: dict[str, Any] = {}
+
+        # analysis state, used by scan_field, etc.
+        self.fields = StatementFields()
+        self.has_statement_bases = False
+        self.kw_only = self.params.get("kw_only", False)
+        self.KW_ONLY_seen = False
+
+    def register(self) -> None:
+        if self.dialect is None:
+            return
+        self.dialect.register(self.cls)
+
+    def emit(self):
+        self._self_name = "__kirin_stmt_self" if "self" in self.fields else "self"
+        self._class_name = "__kirin_stmt_cls" if "cls" in self.fields else "cls"
+        self._run_passes("emit_")
+
+    def verify(self):
+        self._run_passes("verify_")
+
+    def _run_passes(self, prefix: str):
+        for name, method in inspect.getmembers(self, inspect.ismethod):
+            if name.startswith(prefix):
+                method()

kirin/decl/camel2snake.py

@@ -0,0 +1,2 @@
+def camel2snake(name: str) -> str:
+    return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_")

kirin/decl/emit/_create_fn.py

@@ -0,0 +1,29 @@
+"""This module provides a function to create a function dynamically.
+
+Copied from `dataclasses._create_fn` in Python 3.10.13.
+"""
+
+from dataclasses import MISSING
+
+
+def create_fn(name, args, body, *, globals=None, locals=None, return_type=MISSING):
+    # Note that we may mutate locals. Callers beware!
+    # The only callers are internal to this module, so no
+    # worries about external callers.
+    if locals is None:
+        locals = {}
+    return_annotation = ""
+    if return_type is not MISSING:
+        locals["_return_type"] = return_type
+        return_annotation = "->_return_type"
+    args = ",".join(args)
+    body = "\n".join(f"        {b}" for b in body)
+
+    # Compute the text of the entire function.
+    txt = f"    def {name}({args}){return_annotation}:\n{body}"
+
+    local_vars = ", ".join(locals.keys())
+    txt = f"def __create_fn__({local_vars}):\n{txt}\n    return {name}"
+    ns = {}
+    exec(txt, globals, ns)
+    return ns["__create_fn__"](**locals)

kirin/decl/emit/_set_new_attribute.py

@@ -0,0 +1,22 @@
+"""Copied from dataclasses in Python 3.10.13.
+"""
+
+from types import FunctionType
+
+
+def set_qualname(cls: type, value):
+    # Ensure that the functions returned from _create_fn uses the proper
+    # __qualname__ (the class they belong to).
+    if isinstance(value, FunctionType):
+        value.__qualname__ = f"{cls.__qualname__}.{value.__name__}"
+    return value
+
+
+def set_new_attribute(cls: type, name: str, value):
+    # Never overwrites an existing attribute.  Returns True if the
+    # attribute already exists.
+    if name in cls.__dict__:
+        return True
+    set_qualname(cls, value)
+    setattr(cls, name, value)
+    return False

kirin/decl/emit/dialect.py

@@ -0,0 +1,8 @@
+from kirin.decl.base import BaseModifier
+
+
+class EmitDialect(BaseModifier):
+
+    def emit_dialect(self):
+        setattr(self.cls, "dialect", self.dialect)
+        return