kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/math/stmts.py

@@ -0,0 +1,369 @@
+# This file is generated by gen.py
+from kirin import ir, types
+from kirin.decl import info, statement
+from kirin.dialects.math.dialect import dialect
+
+
+@statement(dialect=dialect)
+class acos(ir.Statement):
+    """acos statement, wrapping the math.acos function"""
+
+    name = "acos"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class asin(ir.Statement):
+    """asin statement, wrapping the math.asin function"""
+
+    name = "asin"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class asinh(ir.Statement):
+    """asinh statement, wrapping the math.asinh function"""
+
+    name = "asinh"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class atan(ir.Statement):
+    """atan statement, wrapping the math.atan function"""
+
+    name = "atan"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class atan2(ir.Statement):
+    """atan2 statement, wrapping the math.atan2 function"""
+
+    name = "atan2"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    y: ir.SSAValue = info.argument(types.Float)
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class atanh(ir.Statement):
+    """atanh statement, wrapping the math.atanh function"""
+
+    name = "atanh"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class ceil(ir.Statement):
+    """ceil statement, wrapping the math.ceil function"""
+
+    name = "ceil"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class copysign(ir.Statement):
+    """copysign statement, wrapping the math.copysign function"""
+
+    name = "copysign"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    y: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class cos(ir.Statement):
+    """cos statement, wrapping the math.cos function"""
+
+    name = "cos"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class cosh(ir.Statement):
+    """cosh statement, wrapping the math.cosh function"""
+
+    name = "cosh"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class degrees(ir.Statement):
+    """degrees statement, wrapping the math.degrees function"""
+
+    name = "degrees"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class erf(ir.Statement):
+    """erf statement, wrapping the math.erf function"""
+
+    name = "erf"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class erfc(ir.Statement):
+    """erfc statement, wrapping the math.erfc function"""
+
+    name = "erfc"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class exp(ir.Statement):
+    """exp statement, wrapping the math.exp function"""
+
+    name = "exp"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class expm1(ir.Statement):
+    """expm1 statement, wrapping the math.expm1 function"""
+
+    name = "expm1"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class fabs(ir.Statement):
+    """fabs statement, wrapping the math.fabs function"""
+
+    name = "fabs"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class floor(ir.Statement):
+    """floor statement, wrapping the math.floor function"""
+
+    name = "floor"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class fmod(ir.Statement):
+    """fmod statement, wrapping the math.fmod function"""
+
+    name = "fmod"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    y: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class gamma(ir.Statement):
+    """gamma statement, wrapping the math.gamma function"""
+
+    name = "gamma"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class isfinite(ir.Statement):
+    """isfinite statement, wrapping the math.isfinite function"""
+
+    name = "isfinite"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class isinf(ir.Statement):
+    """isinf statement, wrapping the math.isinf function"""
+
+    name = "isinf"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class isnan(ir.Statement):
+    """isnan statement, wrapping the math.isnan function"""
+
+    name = "isnan"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class lgamma(ir.Statement):
+    """lgamma statement, wrapping the math.lgamma function"""
+
+    name = "lgamma"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class log10(ir.Statement):
+    """log10 statement, wrapping the math.log10 function"""
+
+    name = "log10"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class log1p(ir.Statement):
+    """log1p statement, wrapping the math.log1p function"""
+
+    name = "log1p"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class log2(ir.Statement):
+    """log2 statement, wrapping the math.log2 function"""
+
+    name = "log2"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class pow(ir.Statement):
+    """pow statement, wrapping the math.pow function"""
+
+    name = "pow"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    y: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class radians(ir.Statement):
+    """radians statement, wrapping the math.radians function"""
+
+    name = "radians"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class remainder(ir.Statement):
+    """remainder statement, wrapping the math.remainder function"""
+
+    name = "remainder"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    y: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class sin(ir.Statement):
+    """sin statement, wrapping the math.sin function"""
+
+    name = "sin"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class sinh(ir.Statement):
+    """sinh statement, wrapping the math.sinh function"""
+
+    name = "sinh"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class sqrt(ir.Statement):
+    """sqrt statement, wrapping the math.sqrt function"""
+
+    name = "sqrt"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class tan(ir.Statement):
+    """tan statement, wrapping the math.tan function"""
+
+    name = "tan"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class tanh(ir.Statement):
+    """tanh statement, wrapping the math.tanh function"""
+
+    name = "tanh"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class trunc(ir.Statement):
+    """trunc statement, wrapping the math.trunc function"""
+
+    name = "trunc"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)
+
+
+@statement(dialect=dialect)
+class ulp(ir.Statement):
+    """ulp statement, wrapping the math.ulp function"""
+
+    name = "ulp"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    x: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(types.Float)

kirin/dialects/module.py

@@ -0,0 +1,139 @@
+"""Module dialect provides a simple module
+that is roughly a list of function statements.
+
+This dialect provides the dialect necessary for compiling a function into
+lower-level IR with all its callee functions.
+"""
+
+from kirin import ir, types, interp
+from kirin.decl import info, statement
+from kirin.print import Printer
+from kirin.analysis import TypeInference
+from kirin.exceptions import VerificationError
+
+from ._pprint_helper import pprint_calllike
+
+dialect = ir.Dialect("module")
+
+
+@statement(dialect=dialect)
+class Module(ir.Statement):
+    traits = frozenset(
+        {ir.IsolatedFromAbove(), ir.SymbolTable(), ir.SymbolOpInterface()}
+    )
+    sym_name: str = info.attribute()
+    entry: str = info.attribute()
+    body: ir.Region = info.region(multi=False)
+
+
+@statement(dialect=dialect)
+class Invoke(ir.Statement):
+    """A special statement that represents
+    a function calling functions by symbol name.
+
+    Note:
+        This statement is here for completeness, for interpretation,
+        it is recommended to rewrite this statement into a `func.Invoke`
+        after looking up the symbol table.
+    """
+
+    callee: str = info.attribute()
+    inputs: tuple[ir.SSAValue, ...] = info.argument()
+    kwargs: tuple[str, ...] = info.attribute()
+    result: ir.ResultValue = info.result()
+
+    def print_impl(self, printer: Printer) -> None:
+        pprint_calllike(self, self.callee, printer)
+
+    def verify(self) -> None:
+        if self.kwargs:
+            for name in self.kwargs:
+                if name not in self.callee:
+                    raise VerificationError(
+                        self,
+                        f"method {self.callee} does not have argument {name}",
+                    )
+        elif len(self.callee) - 1 != len(self.args):
+            raise VerificationError(
+                self,
+                f"expected {len(self.callee)} arguments, got {len(self.args)}",
+            )
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Module)
+    def interp_Module(
+        self, interp: interp.Interpreter, frame: interp.Frame, stmt: Module
+    ):
+        for stmt_ in stmt.body.blocks[0].stmts:
+            if (trait := stmt.get_trait(ir.SymbolOpInterface)) is not None:
+                interp.symbol_table[trait.get_sym_name(stmt_).data] = stmt_
+        return ()
+
+    @interp.impl(Invoke)
+    def interp_Invoke(
+        self, interpreter: interp.Interpreter, frame: interp.Frame, stmt: Invoke
+    ):
+        callee = interpreter.symbol_table.get(stmt.callee)
+        if callee is None:
+            raise interp.InterpreterError(f"symbol {stmt.callee} not found")
+
+        trait = callee.get_trait(ir.CallableStmtInterface)
+        if trait is None:
+            raise interp.InterpreterError(
+                f"{stmt.callee} is not callable, got {callee.__class__.__name__}"
+            )
+
+        body = trait.get_callable_region(callee)
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=stmt.callee,
+            arg_names=[
+                arg.name or str(idx) for idx, arg in enumerate(body.blocks[0].args)
+            ],
+            dialects=interpreter.dialects,
+            code=stmt,
+        )
+        return interpreter.run_method(mt, frame.get_values(stmt.inputs))
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Module)
+    def typeinfer_Module(
+        self, interp: TypeInference, frame: interp.Frame, stmt: Module
+    ):
+        for stmt_ in stmt.body.blocks[0].stmts:
+            if (trait := stmt.get_trait(ir.SymbolOpInterface)) is not None:
+                interp.symbol_table[trait.get_sym_name(stmt_).data] = stmt_
+        return ()
+
+    @interp.impl(Invoke)
+    def typeinfer_Invoke(
+        self, interp: TypeInference, frame: interp.Frame, stmt: Invoke
+    ):
+        callee = interp.symbol_table.get(stmt.callee)
+        if callee is None:
+            return (types.Bottom,)
+
+        trait = callee.get_trait(ir.CallableStmtInterface)
+        if trait is None:
+            return (types.Bottom,)
+
+        body = trait.get_callable_region(callee)
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=stmt.callee,
+            arg_names=[
+                arg.name or str(idx) for idx, arg in enumerate(body.blocks[0].args)
+            ],
+            dialects=interp.dialects,
+            code=stmt,
+        )
+        interp.run_method(mt, mt.arg_types)
+        return tuple(result.type for result in callee.results)

kirin/dialects/py/__init__.py

@@ -0,0 +1,40 @@
+"""Python dialects module.
+
+This module contains a set of dialects that represent
+different fractions of the Python language. The dialects
+are designed to be used in a union to represent the
+entire Python language.
+"""
+
+from . import (
+    cmp as cmp,
+    len as len,
+    attr as attr,
+    base as base,
+    list as list,
+    binop as binop,
+    range as range,
+    slice as slice,
+    tuple as tuple,
+    unary as unary,
+    assign as assign,
+    boolop as boolop,
+    unpack as unpack,
+    builtin as builtin,
+    constant as constant,
+    indexing as indexing,
+    iterable as iterable,
+)
+from .len import Len as Len
+from .attr import GetAttr as GetAttr
+from .range import Range as Range
+from .slice import Slice as Slice
+from .assign import Alias as Alias, SetItem as SetItem
+from .boolop import Or as Or, And as And
+from .builtin import Abs as Abs, Sum as Sum
+from .constant import Constant as Constant
+from .indexing import GetItem as GetItem, PyGetItemLike as PyGetItemLike
+from .cmp.stmts import *  # noqa: F403
+from .list.stmts import Append as Append
+from .binop.stmts import *  # noqa: F403
+from .unary.stmts import *  # noqa: F403

kirin/dialects/py/assertion.py

@@ -0,0 +1,91 @@
+"""Assertion dialect for Python.
+
+This module contains the dialect for the Python `assert` statement, including:
+
+- The `Assert` statement class.
+- The lowering pass for the `assert` statement.
+- The concrete implementation of the `assert` statement.
+- The type inference implementation of the `assert` statement.
+- The Julia emitter for the `assert` statement.
+
+This dialect maps `ast.Assert` nodes to the `Assert` statement.
+"""
+
+import ast
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.emit import EmitStrFrame, julia
+from kirin.print import Printer
+
+dialect = ir.Dialect("py.assert")
+
+
+@statement(dialect=dialect)
+class Assert(ir.Statement):
+    condition: ir.SSAValue
+    message: ir.SSAValue = info.argument(types.String)
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="keyword"):
+            printer.print_name(self)
+
+        printer.plain_print(" ")
+        printer.print(self.condition)
+
+        if self.message:
+            printer.plain_print(", ")
+            printer.print(self.message)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Assert(
+        self, state: lowering.LoweringState, node: ast.Assert
+    ) -> lowering.Result:
+        from kirin.dialects.py.constant import Constant
+
+        cond = state.visit(node.test).expect_one()
+        if node.msg:
+            message = state.visit(node.msg).expect_one()
+            state.append_stmt(Assert(condition=cond, message=message))
+        else:
+            message_stmt = state.append_stmt(Constant(""))
+            state.append_stmt(Assert(condition=cond, message=message_stmt.result))
+        return lowering.Result()
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Assert)
+    def assert_stmt(
+        self, interp_: interp.Interpreter, frame: interp.Frame, stmt: Assert
+    ):
+        if frame.get(stmt.condition) is True:
+            return ()
+
+        if stmt.message:
+            raise interp.WrapException(AssertionError(frame.get(stmt.message)))
+        else:
+            raise interp.WrapException(AssertionError("Assertion failed"))
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Assert)
+    def assert_stmt(self, interp, frame, stmt: Assert):
+        return (types.Bottom,)
+
+
+@dialect.register(key="emit.julia")
+class EmitJulia(interp.MethodTable):
+
+    @interp.impl(Assert)
+    def emit_assert(self, interp: julia.EmitJulia, frame: EmitStrFrame, stmt: Assert):
+        interp.writeln(
+            frame, f"@assert {frame.get(stmt.condition)} {frame.get(stmt.message)}"
+        )
+        return ()