kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/lowering/func.py

@@ -0,0 +1,134 @@
+import ast
+
+from kirin import ir, types, lowering
+from kirin.dialects import cf, func
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("lowering.func")
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Return(
+        self, state: lowering.LoweringState, node: ast.Return
+    ) -> lowering.Result:
+        if node.value is None:
+            stmt = func.Return(state.append_stmt(func.ConstantNone()).result)
+            state.append_stmt(stmt)
+        else:
+            result = state.visit(node.value).expect_one()
+            stmt = func.Return(result)
+            state.append_stmt(stmt)
+        return lowering.Result()
+
+    def lower_FunctionDef(
+        self, state: lowering.LoweringState, node: ast.FunctionDef
+    ) -> lowering.Result:
+        self.assert_simple_arguments(node.args)
+        signature = func.Signature(
+            inputs=tuple(
+                self.get_hint(state, arg.annotation) for arg in node.args.args
+            ),
+            output=self.get_hint(state, node.returns),
+        )
+        frame = state.current_frame
+
+        entries: dict[str, ir.SSAValue] = {}
+        entr_block = ir.Block()
+        fn_self = entr_block.args.append_from(
+            types.Generic(
+                ir.Method, types.Tuple.where(signature.inputs), signature.output
+            ),
+            node.name + "_self",
+        )
+        entries[node.name] = fn_self
+        for arg, type in zip(node.args.args, signature.inputs):
+            entries[arg.arg] = entr_block.args.append_from(type, arg.arg)
+
+        def callback(frame: lowering.Frame, value: ir.SSAValue):
+            first_stmt = entr_block.first_stmt
+            stmt = func.GetField(obj=fn_self, field=len(frame.captures) - 1)
+            if value.name:
+                stmt.result.name = value.name
+            stmt.result.type = value.type
+            stmt.source = state.source
+            if first_stmt:
+                stmt.insert_before(first_stmt)
+            else:
+                entr_block.stmts.append(stmt)
+            return stmt.result
+
+        func_frame = state.push_frame(
+            lowering.Frame.from_stmts(
+                node.body,
+                state,
+                entr_block=entr_block,
+                globals=frame.globals,
+                capture_callback=callback,
+            )
+        )
+        func_frame.defs.update(entries)
+        state.exhaust()
+
+        for block in func_frame.curr_region.blocks:
+            if not block.last_stmt or not block.last_stmt.has_trait(ir.IsTerminator):
+                block.stmts.append(
+                    cf.Branch(arguments=(), successor=func_frame.next_block)
+                )
+
+        none_stmt = func.ConstantNone()
+        rtrn_stmt = func.Return(none_stmt.result)
+        func_frame.next_block.stmts.append(none_stmt)
+        func_frame.next_block.stmts.append(rtrn_stmt)
+        state.pop_frame()
+
+        if state.current_frame.parent is None:  # toplevel function
+            stmt = frame.append_stmt(
+                func.Function(
+                    sym_name=node.name,
+                    signature=signature,
+                    body=func_frame.curr_region,
+                )
+            )
+            return lowering.Result(stmt)
+
+        if node.decorator_list:
+            raise DialectLoweringError(
+                "decorators are not supported on nested functions"
+            )
+
+        # nested function, lookup unknown variables
+        first_stmt = func_frame.curr_region.blocks[0].first_stmt
+        if first_stmt is None:
+            raise DialectLoweringError("empty function body")
+
+        captured = [value for value in func_frame.captures.values()]
+        lambda_stmt = func.Lambda(
+            tuple(captured),
+            sym_name=node.name,
+            signature=signature,
+            body=func_frame.curr_region,
+        )
+        lambda_stmt.result.name = node.name
+        # NOTE: Python automatically assigns the lambda to the name
+        frame.defs[node.name] = frame.append_stmt(lambda_stmt).result
+        return lowering.Result(lambda_stmt)
+
+    def assert_simple_arguments(self, node: ast.arguments) -> None:
+        if node.kwonlyargs:
+            raise DialectLoweringError("keyword-only arguments are not supported")
+
+        if node.posonlyargs:
+            raise DialectLoweringError("positional-only arguments are not supported")
+
+    @staticmethod
+    def get_hint(state: lowering.LoweringState, node: ast.expr | None):
+        if node is None:
+            return types.Any
+
+        try:
+            t = state.get_global(node).unwrap()
+            return types.hint2type(t)
+        except:  # noqa: E722
+            raise DialectLoweringError(f"expect a type hint, got {ast.unparse(node)}")

kirin/dialects/math/__init__.py

@@ -0,0 +1,41 @@
+"math dialect, modeling functions in python's `math` stdlib"  # This file is generated by gen.py
+from kirin.dialects.math.stmts import (
+    cos as cos,
+    erf as erf,
+    exp as exp,
+    pow as pow,
+    sin as sin,
+    tan as tan,
+    ulp as ulp,
+    acos as acos,
+    asin as asin,
+    atan as atan,
+    ceil as ceil,
+    cosh as cosh,
+    erfc as erfc,
+    fabs as fabs,
+    fmod as fmod,
+    log2 as log2,
+    sinh as sinh,
+    sqrt as sqrt,
+    tanh as tanh,
+    asinh as asinh,
+    atan2 as atan2,
+    atanh as atanh,
+    expm1 as expm1,
+    floor as floor,
+    gamma as gamma,
+    isinf as isinf,
+    isnan as isnan,
+    log1p as log1p,
+    log10 as log10,
+    trunc as trunc,
+    lgamma as lgamma,
+    degrees as degrees,
+    radians as radians,
+    copysign as copysign,
+    isfinite as isfinite,
+    remainder as remainder,
+)
+from kirin.dialects.math.interp import MathMethodTable as MathMethodTable
+from kirin.dialects.math.dialect import dialect as dialect

kirin/dialects/math/_gen.py

@@ -0,0 +1,176 @@
+import os
+import math
+import inspect
+import textwrap
+from pathlib import Path
+
+import black
+
+# NOTE: typeinfer and lowering should be the default, so we don't generate them.
+
+
+def builtin_math_functions():
+    for name, obj in inspect.getmembers(math):
+        # skip some special cases for now
+        if name in (
+            "prod",
+            "perm",
+            "modf",
+            "ldexp",
+            "lcm",
+            "isqrt",
+            "isclose",
+            "gcd",
+            "fsum",
+            "frexp",
+            "factorial",
+            "acosh",
+            "comb",
+            "dist",
+            "sumprod",
+            "nextafter",
+            # 3.10 compat
+            "cbrt",
+            "exp2",
+        ):
+            continue
+
+        if inspect.isbuiltin(obj):
+            try:
+                sig = inspect.signature(obj)
+                yield name, obj, sig
+            except:  # noqa: E722
+                continue
+
+
+with open(os.path.join(os.path.dirname(__file__), "stmts.py"), "w") as f:
+    f.write("# This file is generated by gen.py\n")
+    f.write("from kirin import ir, types\n")
+    f.write("from kirin.decl import statement, info\n")
+    f.write("from kirin.dialects.math.dialect import dialect\n")
+    f.write("\n")
+    for name, obj, sig in builtin_math_functions():
+        fields = "\n".join(
+            [
+                f"    {arg} : ir.SSAValue = info.argument(types.Float)"
+                for arg in sig.parameters.keys()
+            ]
+        )
+        f.write(
+            textwrap.dedent(
+                f"""
+@statement(dialect=dialect)
+class {name}(ir.Statement):
+    \"\"\"{name} statement, wrapping the math.{name} function
+    \"\"\"
+    name = "{name}"
+    traits = frozenset({{ir.Pure(), ir.FromPythonCall()}})
+{fields}
+    result: ir.ResultValue = info.result(types.Float)
+"""
+            )
+        )
+
+
+with open(os.path.join(os.path.dirname(__file__), "interp.py"), "w") as f:
+    f.write("# This file is generated by gen.py\n")
+    f.write("import math\n")
+    f.write("from kirin.dialects.math.dialect import dialect\n")
+    f.write("from kirin.dialects.math import stmts\n")
+    f.write("from kirin.interp import MethodTable, Frame, impl\n")
+    f.write("\n")
+
+    implements = []
+    for name, obj, sig in builtin_math_functions():
+        fields = ", ".join(
+            [f"values[{idx}]" for idx, _ in enumerate(sig.parameters.keys())]
+        )
+        implements.append(
+            f"""
+    @impl(stmts.{name})
+    def {name}(self, interp, frame: Frame, stmt: stmts.{name}):
+        values = frame.get_values(stmt.args)
+        return (math.{name}({fields}),)"""
+        )
+
+    # Write the interpreter class
+    methods = "\n\n".join(implements)
+    f.write(
+        f"""
+@dialect.register
+class MathMethodTable(MethodTable):
+{methods}
+"""
+    )
+
+# __init__.py
+with open(os.path.join(os.path.dirname(__file__), "__init__.py"), "w") as f:
+    f.write('"math dialect, modeling functions in python\'s `math` stdlib"')
+    f.write("# This file is generated by gen.py\n")
+    f.write("from kirin.dialects.math.dialect import dialect as dialect\n")
+    f.write("from kirin.dialects.math.stmts import (\n")
+    for name, obj, sig in builtin_math_functions():
+        f.write(f"    {name} as {name},\n")
+    f.write(")\n")
+    f.write(
+        "from kirin.dialects.math.interp import MathMethodTable as MathMethodTable\n"
+    )
+    f.write("\n")
+
+for file in ["__init__.py", "interp.py", "stmts.py"]:
+    # format the file in place + using the project config
+    black.format_file_in_place(
+        Path(os.path.join(os.path.dirname(__file__), file)),
+        fast=False,
+        mode=black.FileMode(),
+    )
+
+
+# import math as pymath
+
+# from kirin.compile import compile
+# from kirin.dialects import math
+
+
+# # print(math.sin(x=TestValue()))
+# # print(inspect.getargspec(math.sin.__init__))
+# # print(math.sin.__init__)
+# @basic
+# def complicated_math_expr(x):
+#     return math.sin(math.cos(x) + math.tan(0.5))
+
+
+# def test_math():
+#     complicated_math_expr.code.print()
+#     complicated_math_expr.narrow_types()
+#     truth = pymath.sin(pymath.cos(1) + pymath.tan(0.5))
+#     assert (complicated_math_expr(1) - truth) / truth < 1e-6
+
+# test_basic.py
+project_dir = Path(__file__).parent.parent.parent.parent.parent
+with open(project_dir.joinpath("test", "dialects", "math", "test_basic.py"), "w") as f:
+    f.write("# type: ignore\n")
+    f.write("# This file is generated by gen.py\n")
+    f.write("import math as pymath\n")
+    f.write("from kirin.prelude import basic\n")
+    f.write("from kirin.dialects import math\n")
+    f.write("\n")
+    f.write("\n")
+
+    for name, obj, sig in builtin_math_functions():
+        args = ", ".join(arg for arg in sig.parameters.keys())
+        inputs = ", ".join("0.42" for _ in sig.parameters.keys())
+
+        f.write(
+            textwrap.dedent(
+                f"""
+                @basic
+                def {name}_func({args}):
+                    return math.{name}({args})
+
+                def test_{name}():
+                    truth = pymath.{name}({inputs})
+                    assert ({name}_func({inputs}) - truth) < 1e-6
+                """
+            )
+        )

kirin/dialects/math/dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("math")

kirin/dialects/math/interp.py

@@ -0,0 +1,190 @@
+# This file is generated by gen.py
+import math
+
+from kirin.interp import Frame, MethodTable, impl
+from kirin.dialects.math import stmts
+from kirin.dialects.math.dialect import dialect
+
+
+@dialect.register
+class MathMethodTable(MethodTable):
+
+    @impl(stmts.acos)
+    def acos(self, interp, frame: Frame, stmt: stmts.acos):
+        values = frame.get_values(stmt.args)
+        return (math.acos(values[0]),)
+
+    @impl(stmts.asin)
+    def asin(self, interp, frame: Frame, stmt: stmts.asin):
+        values = frame.get_values(stmt.args)
+        return (math.asin(values[0]),)
+
+    @impl(stmts.asinh)
+    def asinh(self, interp, frame: Frame, stmt: stmts.asinh):
+        values = frame.get_values(stmt.args)
+        return (math.asinh(values[0]),)
+
+    @impl(stmts.atan)
+    def atan(self, interp, frame: Frame, stmt: stmts.atan):
+        values = frame.get_values(stmt.args)
+        return (math.atan(values[0]),)
+
+    @impl(stmts.atan2)
+    def atan2(self, interp, frame: Frame, stmt: stmts.atan2):
+        values = frame.get_values(stmt.args)
+        return (math.atan2(values[0], values[1]),)
+
+    @impl(stmts.atanh)
+    def atanh(self, interp, frame: Frame, stmt: stmts.atanh):
+        values = frame.get_values(stmt.args)
+        return (math.atanh(values[0]),)
+
+    @impl(stmts.ceil)
+    def ceil(self, interp, frame: Frame, stmt: stmts.ceil):
+        values = frame.get_values(stmt.args)
+        return (math.ceil(values[0]),)
+
+    @impl(stmts.copysign)
+    def copysign(self, interp, frame: Frame, stmt: stmts.copysign):
+        values = frame.get_values(stmt.args)
+        return (math.copysign(values[0], values[1]),)
+
+    @impl(stmts.cos)
+    def cos(self, interp, frame: Frame, stmt: stmts.cos):
+        values = frame.get_values(stmt.args)
+        return (math.cos(values[0]),)
+
+    @impl(stmts.cosh)
+    def cosh(self, interp, frame: Frame, stmt: stmts.cosh):
+        values = frame.get_values(stmt.args)
+        return (math.cosh(values[0]),)
+
+    @impl(stmts.degrees)
+    def degrees(self, interp, frame: Frame, stmt: stmts.degrees):
+        values = frame.get_values(stmt.args)
+        return (math.degrees(values[0]),)
+
+    @impl(stmts.erf)
+    def erf(self, interp, frame: Frame, stmt: stmts.erf):
+        values = frame.get_values(stmt.args)
+        return (math.erf(values[0]),)
+
+    @impl(stmts.erfc)
+    def erfc(self, interp, frame: Frame, stmt: stmts.erfc):
+        values = frame.get_values(stmt.args)
+        return (math.erfc(values[0]),)
+
+    @impl(stmts.exp)
+    def exp(self, interp, frame: Frame, stmt: stmts.exp):
+        values = frame.get_values(stmt.args)
+        return (math.exp(values[0]),)
+
+    @impl(stmts.expm1)
+    def expm1(self, interp, frame: Frame, stmt: stmts.expm1):
+        values = frame.get_values(stmt.args)
+        return (math.expm1(values[0]),)
+
+    @impl(stmts.fabs)
+    def fabs(self, interp, frame: Frame, stmt: stmts.fabs):
+        values = frame.get_values(stmt.args)
+        return (math.fabs(values[0]),)
+
+    @impl(stmts.floor)
+    def floor(self, interp, frame: Frame, stmt: stmts.floor):
+        values = frame.get_values(stmt.args)
+        return (math.floor(values[0]),)
+
+    @impl(stmts.fmod)
+    def fmod(self, interp, frame: Frame, stmt: stmts.fmod):
+        values = frame.get_values(stmt.args)
+        return (math.fmod(values[0], values[1]),)
+
+    @impl(stmts.gamma)
+    def gamma(self, interp, frame: Frame, stmt: stmts.gamma):
+        values = frame.get_values(stmt.args)
+        return (math.gamma(values[0]),)
+
+    @impl(stmts.isfinite)
+    def isfinite(self, interp, frame: Frame, stmt: stmts.isfinite):
+        values = frame.get_values(stmt.args)
+        return (math.isfinite(values[0]),)
+
+    @impl(stmts.isinf)
+    def isinf(self, interp, frame: Frame, stmt: stmts.isinf):
+        values = frame.get_values(stmt.args)
+        return (math.isinf(values[0]),)
+
+    @impl(stmts.isnan)
+    def isnan(self, interp, frame: Frame, stmt: stmts.isnan):
+        values = frame.get_values(stmt.args)
+        return (math.isnan(values[0]),)
+
+    @impl(stmts.lgamma)
+    def lgamma(self, interp, frame: Frame, stmt: stmts.lgamma):
+        values = frame.get_values(stmt.args)
+        return (math.lgamma(values[0]),)
+
+    @impl(stmts.log10)
+    def log10(self, interp, frame: Frame, stmt: stmts.log10):
+        values = frame.get_values(stmt.args)
+        return (math.log10(values[0]),)
+
+    @impl(stmts.log1p)
+    def log1p(self, interp, frame: Frame, stmt: stmts.log1p):
+        values = frame.get_values(stmt.args)
+        return (math.log1p(values[0]),)
+
+    @impl(stmts.log2)
+    def log2(self, interp, frame: Frame, stmt: stmts.log2):
+        values = frame.get_values(stmt.args)
+        return (math.log2(values[0]),)
+
+    @impl(stmts.pow)
+    def pow(self, interp, frame: Frame, stmt: stmts.pow):
+        values = frame.get_values(stmt.args)
+        return (math.pow(values[0], values[1]),)
+
+    @impl(stmts.radians)
+    def radians(self, interp, frame: Frame, stmt: stmts.radians):
+        values = frame.get_values(stmt.args)
+        return (math.radians(values[0]),)
+
+    @impl(stmts.remainder)
+    def remainder(self, interp, frame: Frame, stmt: stmts.remainder):
+        values = frame.get_values(stmt.args)
+        return (math.remainder(values[0], values[1]),)
+
+    @impl(stmts.sin)
+    def sin(self, interp, frame: Frame, stmt: stmts.sin):
+        values = frame.get_values(stmt.args)
+        return (math.sin(values[0]),)
+
+    @impl(stmts.sinh)
+    def sinh(self, interp, frame: Frame, stmt: stmts.sinh):
+        values = frame.get_values(stmt.args)
+        return (math.sinh(values[0]),)
+
+    @impl(stmts.sqrt)
+    def sqrt(self, interp, frame: Frame, stmt: stmts.sqrt):
+        values = frame.get_values(stmt.args)
+        return (math.sqrt(values[0]),)
+
+    @impl(stmts.tan)
+    def tan(self, interp, frame: Frame, stmt: stmts.tan):
+        values = frame.get_values(stmt.args)
+        return (math.tan(values[0]),)
+
+    @impl(stmts.tanh)
+    def tanh(self, interp, frame: Frame, stmt: stmts.tanh):
+        values = frame.get_values(stmt.args)
+        return (math.tanh(values[0]),)
+
+    @impl(stmts.trunc)
+    def trunc(self, interp, frame: Frame, stmt: stmts.trunc):
+        values = frame.get_values(stmt.args)
+        return (math.trunc(values[0]),)
+
+    @impl(stmts.ulp)
+    def ulp(self, interp, frame: Frame, stmt: stmts.ulp):
+        values = frame.get_values(stmt.args)
+        return (math.ulp(values[0]),)