kirin-toolchain 0.13.0__py3-none-any.whl

kirin/prelude.py

@@ -0,0 +1,197 @@
+"""This module contains some common eDSLs that can be used to build
+more complex eDSLs. The eDSLs in this module are built on top of the
+basic dialects provided by the `kirin.dialects` module.
+"""
+
+from typing_extensions import Doc, Annotated
+
+from kirin.ir import Method, dialect_group
+from kirin.passes import aggressive
+from kirin.dialects import cf, scf, func, math, ilist, lowering
+from kirin.dialects.py import (
+    cmp,
+    len,
+    attr,
+    base,
+    list,
+    binop,
+    range,
+    slice,
+    tuple,
+    unary,
+    assign,
+    boolop,
+    builtin,
+    constant,
+    indexing,
+    iterable,
+    assertion,
+)
+from kirin.passes.fold import Fold
+from kirin.passes.typeinfer import TypeInfer
+
+
+@dialect_group(
+    [
+        base,
+        binop,
+        cmp,
+        unary,
+        assign,
+        attr,
+        boolop,
+        builtin,
+        constant,
+        indexing,
+        len,
+        tuple,
+        assertion,
+        iterable,
+    ]
+)
+def python_basic(self):
+    """The basic Python dialect without list, range, and slice."""
+
+    def run_pass(mt: Method) -> None:
+        pass
+
+    return run_pass
+
+
+@dialect_group(
+    python_basic.union(
+        [list, range, slice, cf, func, lowering.cf, lowering.func, lowering.call, math]
+    )
+)
+def python_no_opt(self):
+    """The Python dialect without optimization passes."""
+
+    def run_pass(mt: Method) -> None:
+        pass
+
+    return run_pass
+
+
+@dialect_group(
+    python_basic.union(
+        [ilist, range, slice, cf, func, math, lowering.cf, lowering.func, lowering.call]
+    )
+)
+def basic_no_opt(self):
+    """The basic kernel without optimization passes. This is a builtin
+    eDSL that includes the basic dialects that are commonly used in
+    Python-like eDSLs.
+
+    This eDSL includes the basic dialects without any optimization passes.
+    Other eDSL can usually be built on top of this eDSL by utilizing the
+    `basic_no_opt.add` method to add more dialects and optimization passes.
+
+    Note that unlike Python, list in this eDSL is immutable, and the
+    `append` method is not available. Use `+` operator to concatenate lists
+    instead. Immutable list is easier to optimize and reason about.
+
+    See also [`basic`][kirin.prelude.basic] for the basic kernel with optimization passes.
+    See also [`ilist`][kirin.dialects.ilist] for the immutable list dialect.
+    """
+    ilist_desugar = ilist.IListDesugar(self)
+
+    def run_pass(mt: Method) -> None:
+        ilist_desugar.fixpoint(mt)
+
+    return run_pass
+
+
+@dialect_group(basic_no_opt)
+def basic(self):
+    """The basic kernel.
+
+    This eDSL includes the basic dialects and the basic optimization passes.
+    Other eDSL can usually be built on top of this eDSL by utilizing the
+    `basic.add` method to add more dialects and optimization passes.
+
+    See also [`basic_no_opt`][kirin.prelude.basic_no_opt] for the basic kernel without optimization passes.
+
+    ## Example
+
+    ```python
+    from kirin.prelude import basic
+
+    @basic(typeinfer=True)
+    def main(x: int) -> int:
+        return x + 1 + 1
+
+    main.print() # main is a Method!
+    ```
+    """
+    fold_pass = Fold(self)
+    ilist_desugar = ilist.IListDesugar(self)
+    aggressive_fold_pass = aggressive.Fold(self)
+    typeinfer_pass = TypeInfer(self)
+
+    def run_pass(
+        mt: Annotated[Method, Doc("The method to run pass on.")],
+        *,
+        verify: Annotated[
+            bool, Doc("run `verify` before running passes, default is `True`")
+        ] = True,
+        typeinfer: Annotated[
+            bool,
+            Doc(
+                "run type inference and apply the inferred type to IR, default `False`"
+            ),
+        ] = False,
+        fold: Annotated[bool, Doc("run folding passes")] = True,
+        aggressive: Annotated[
+            bool, Doc("run aggressive folding passes if `fold=True`")
+        ] = False,
+    ) -> None:
+        if verify:
+            mt.verify()
+
+        ilist_desugar.fixpoint(mt)
+
+        if fold:
+            if aggressive:
+                aggressive_fold_pass.fixpoint(mt)
+            else:
+                fold_pass(mt)
+
+        if typeinfer:
+            typeinfer_pass(mt)
+
+    return run_pass
+
+
+@dialect_group(
+    python_basic.union(
+        [ilist, range, slice, scf, cf, func, math, lowering.func, lowering.call]
+    )
+)
+def structural_no_opt(self):
+    """Structural kernel without optimization passes."""
+
+    def run_pass(method: Method) -> None:
+        pass
+
+    return run_pass
+
+
+@dialect_group(
+    python_basic.union(
+        [ilist, range, slice, scf, cf, func, math, lowering.func, lowering.call]
+    )
+)
+def structural(self):
+    """Structural kernel without optimization passes."""
+    typeinfer_pass = TypeInfer(self)
+
+    def run_pass(
+        method: Method, *, verify: bool = True, typeinfer: bool = True
+    ) -> None:
+        if verify:
+            method.verify()
+
+        if typeinfer:
+            typeinfer_pass(method)
+
+    return run_pass

kirin/print/__init__.py

@@ -0,0 +1,15 @@
+"""Pretty printing utilities.
+
+This module provides a pretty printing utility for the IR nodes and other
+objects in the compiler.
+
+The pretty printing utility is implemented using the visitor pattern. The
+[`Printable`][kirin.print.Printable] class is the base class for all objects that can be pretty printed.
+
+The [`Printer`][kirin.print.Printer] class is the visitor that traverses the object and prints the
+object to a string. The [`Printer`][kirin.print.Printer] class provides methods for printing different
+types of objects.
+"""
+
+from kirin.print.printer import Printer as Printer
+from kirin.print.printable import Printable as Printable

kirin/print/printable.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+# NOTE: we don't want to actually load rich here
+if TYPE_CHECKING:
+    from typing import Any, Literal, TypedDict
+
+    from rich.theme import Theme
+    from rich.console import Console
+    from typing_extensions import Unpack
+
+    from kirin import ir
+    from kirin.print import Printer
+
+    class _PrintOptions(TypedDict, total=False):
+        console: Console
+        hint: str
+        analysis: dict["ir.SSAValue", Any]
+        show_indent_mark: bool
+        theme: Theme | dict | Literal["dark", "light"]
+
+
+class Printable:
+    """Base class for all objects that can be pretty printed.
+
+    This class provides an interface for pretty printing objects. The
+    `print` method is the entry point for the printing process. The
+    `print_impl` method is the implementation of the printing process
+    and should be implemented by the derived classes.
+    """
+
+    @staticmethod
+    def __get_printer(
+        printer: Printer | None = None,
+        **options: Unpack["_PrintOptions"],
+    ) -> Printer:
+        if printer is None:
+            from kirin.print import Printer
+
+            return Printer(**options)
+        return printer
+
+    def pager(
+        self,
+        printer: Printer | None = None,
+        **options: Unpack["_PrintOptions"],
+    ) -> None:
+        """Pretty print the object with a pager.
+
+        Args:
+            printer (Printer):
+                `Printer` object to use for printing.
+                If None, a new `Printer` object will be created.
+
+        Keyword Args:
+            console (rich.Console):
+                The console to use for printing. If not provided, a new console
+                will be created.
+            analysis (dict[ir.SSAValue, Printable]):
+                Analysis results to use for printing. If `None`, no analysis results
+            show_indent_mark (bool):
+                Whether to show the indentation mark.
+            theme (Theme | dict | str):
+                The theme to use for printing, defaults to "dark".
+
+        !!! note
+            This function also accepts all other `rich.console.Console` options.
+        """
+        printer = self.__get_printer(printer, **options)
+        with printer.console.pager(styles=True, links=True):
+            self.print(printer)
+
+    def print(
+        self,
+        printer: Printer | None = None,
+        end: str = "\n",
+        **options: Unpack["_PrintOptions"],
+    ) -> None:
+        """
+        Entry point of the printing process.
+
+        Args:
+            printer (Printer):
+                `Printer` object to use for printing.
+                If None, a new `Printer` object will be created.
+
+        Keyword Args:
+            console (rich.Console):
+                The console to use for printing. If not provided, a new console
+                will be created.
+            analysis (dict[ir.SSAValue, Printable]):
+                Analysis results to use for printing. If `None`, no analysis results
+            show_indent_mark (bool):
+                Whether to show the indentation mark.
+            theme (Theme | dict | str):
+                The theme to use for printing, defaults to "dark".
+
+        !!! note
+            This function also accepts all other `rich.console.Console` options.
+        """
+        printer = self.__get_printer(printer, **options)
+        self.print_impl(printer)
+        printer.plain_print(end)
+
+    def print_str(
+        self,
+        printer: Printer | None = None,
+        end: str = "\n",
+        **options: Unpack["_PrintOptions"],
+    ) -> str:
+        """Print the object to a string.
+
+        Args:
+            printer (Printer):
+                `Printer` object to use for printing.
+                If None, a new `Printer` object will be created.
+
+        Keyword Args:
+            console (rich.Console):
+                The console to use for printing. If not provided, a new console
+                will be created.
+            analysis (dict[ir.SSAValue, Printable]):
+                Analysis results to use for printing. If `None`, no analysis results
+            show_indent_mark (bool):
+                Whether to show the indentation mark.
+            theme (Theme | dict | str):
+                The theme to use for printing, defaults to "dark".
+
+        !!! note
+            This function also accepts all other `rich.console.Console` options.
+        """
+        printer = self.__get_printer(printer, **options)
+        with printer.string_io() as stream:
+            self.print(printer, end=end, **options)
+            return stream.getvalue()
+
+    @abstractmethod
+    def print_impl(self, printer: Printer) -> None:
+        raise NotImplementedError(f"print is not implemented for {type(self)}")