kirin-toolchain 0.13.0__py3-none-any.whl

kirin/dialects/py/assign.py

@@ -0,0 +1,103 @@
+"""Assignment dialect for Python.
+
+This module contains the dialect for the Python assignment statement, including:
+
+- Statements: `Alias`, `SetItem`.
+- The lowering pass for the assignments.
+- The concrete implementation of the assignment statements.
+
+This dialects maps Python assignment syntax.
+"""
+
+import ast
+
+from kirin import ir, types, interp, lowering, exceptions
+from kirin.decl import info, statement
+from kirin.print import Printer
+
+dialect = ir.Dialect("py.assign")
+
+T = types.TypeVar("T")
+
+
+@statement(dialect=dialect)
+class Alias(ir.Statement):
+    name = "alias"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    value: ir.SSAValue = info.argument(T)
+    target: ir.PyAttr[str] = info.attribute()
+    result: ir.ResultValue = info.result(T)
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        with printer.rich(style="symbol"):
+            printer.plain_print(self.target.data)
+
+        with printer.rich(style="keyword"):
+            printer.plain_print(" = ")
+
+        printer.print(self.value)
+
+
+@statement(dialect=dialect)
+class SetItem(ir.Statement):
+    name = "setitem"
+    traits = frozenset({ir.FromPythonCall()})
+    obj: ir.SSAValue = info.argument(print=False)
+    value: ir.SSAValue = info.argument(print=False)
+    index: ir.SSAValue = info.argument(print=False)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Alias)
+    def alias(self, interp, frame: interp.Frame, stmt: Alias):
+        return (frame.get(stmt.value),)
+
+    @interp.impl(SetItem)
+    def setindex(self, interp, frame: interp.Frame, stmt: SetItem):
+        frame.get(stmt.obj)[frame.get(stmt.index)] = frame.get(stmt.value)
+        return (None,)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Assign(
+        self, state: lowering.LoweringState, node: ast.Assign
+    ) -> lowering.Result:
+        results: lowering.Result = state.visit(node.value)
+        assert len(node.targets) == len(
+            results
+        ), "number of targets and results do not match"
+
+        current_frame = state.current_frame
+        match node:
+            case ast.Assign(
+                targets=[ast.Name(lhs_name, ast.Store())], value=ast.Name(_, ast.Load())
+            ):
+                stmt = Alias(
+                    value=results[0], target=ir.PyAttr(lhs_name)
+                )  # NOTE: this is guaranteed to be one result
+                stmt.result.name = lhs_name
+                current_frame.defs[lhs_name] = state.append_stmt(stmt).result
+            case _:
+                for target, value in zip(node.targets, results.values):
+                    match target:
+                        # NOTE: if the name exists new ssa value will be
+                        # used in the future to shadow the old one
+                        case ast.Name(name, ast.Store()):
+                            value.name = name
+                            current_frame.defs[name] = value
+                        case ast.Subscript(obj, slice):
+                            obj = state.visit(obj).expect_one()
+                            slice = state.visit(slice).expect_one()
+                            stmt = SetItem(obj=obj, index=slice, value=value)
+                            state.append_stmt(stmt)
+                        case _:
+                            raise exceptions.DialectLoweringError(
+                                f"unsupported target {target}"
+                            )
+        return lowering.Result()  # python assign does not have value

kirin/dialects/py/attr.py

@@ -0,0 +1,59 @@
+"""Attribute access dialect for Python.
+
+This module contains the dialect for the Python attribute access statement, including:
+
+- The `GetAttr` statement class.
+- The lowering pass for the attribute access statement.
+- The concrete implementation of the attribute access statement.
+
+This dialect maps `ast.Attribute` nodes to the `GetAttr` statement.
+"""
+
+import ast
+
+from kirin import ir, interp, lowering, exceptions
+from kirin.decl import info, statement
+
+dialect = ir.Dialect("py.attr")
+
+
+@statement(dialect=dialect)
+class GetAttr(ir.Statement):
+    name = "getattr"
+    traits = frozenset({ir.FromPythonCall()})
+    obj: ir.SSAValue = info.argument(print=False)
+    attrname: str = info.attribute()
+    result: ir.ResultValue = info.result()
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(GetAttr)
+    def getattr(self, interp: interp.Interpreter, frame: interp.Frame, stmt: GetAttr):
+        return (getattr(frame.get(stmt.obj), stmt.attrname),)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Attribute(
+        self, state: lowering.LoweringState, node: ast.Attribute
+    ) -> lowering.Result:
+        from kirin.dialects.py import Constant
+
+        if not isinstance(node.ctx, ast.Load):
+            raise exceptions.DialectLoweringError(
+                f"unsupported attribute context {node.ctx}"
+            )
+
+        # NOTE: eagerly load global variables
+        value = state.get_global_nothrow(node)
+        if value is not None:
+            stmt = state.append_stmt(Constant(value.unwrap()))
+            return lowering.Result(stmt)
+
+        value = state.visit(node.value).expect_one()
+        stmt = GetAttr(obj=value, attrname=node.attr)
+        state.append_stmt(stmt)
+        return lowering.Result(stmt)

kirin/dialects/py/base.py

@@ -0,0 +1,34 @@
+"""Base dialect for Python.
+
+This dialect does not contain statements. It only contains
+lowering rules for `ast.Name` and `ast.Expr`.
+"""
+
+import ast
+
+from kirin import ir, lowering, exceptions
+
+dialect = ir.Dialect("py.base")
+
+
+@dialect.register
+class PythonLowering(lowering.FromPythonAST):
+
+    def lower_Name(
+        self, state: lowering.LoweringState, node: ast.Name
+    ) -> lowering.Result:
+        name = node.id
+        if isinstance(node.ctx, ast.Load):
+            value = state.current_frame.get(name)
+            if value is None:
+                raise exceptions.DialectLoweringError(f"{name} is not defined")
+            return lowering.Result(value)
+        elif isinstance(node.ctx, ast.Store):
+            raise exceptions.DialectLoweringError("unhandled store operation")
+        else:  # Del
+            raise exceptions.DialectLoweringError("unhandled del operation")
+
+    def lower_Expr(
+        self, state: lowering.LoweringState, node: ast.Expr
+    ) -> lowering.Result:
+        return state.visit(node.value)

kirin/dialects/py/binop/__init__.py

@@ -0,0 +1,23 @@
+"""The binop dialect for Python.
+
+This module contains the dialect for binary operation semantics in Python, including:
+
+- The `Add`, `Sub`, `Mult`, `Div`, `FloorDiv`, `Mod`, `Pow`,
+    `LShift`, `RShift`, `BitOr`, `BitXor`, and `BitAnd` statement classes.
+- The lowering pass for binary operations.
+- The concrete implementation of binary operations.
+- The type inference implementation of binary operations.
+- The Julia emitter for binary operations.
+
+This dialect maps `ast.BinOp` nodes to the `Add`, `Sub`, `Mult`, `Div`, `FloorDiv`,
+`Mod`, `Pow`, `LShift`, `RShift`, `BitOr`, `BitXor`, and `BitAnd` statements.
+"""
+
+from . import (
+    julia as julia,
+    interp as interp,
+    lowering as lowering,
+    typeinfer as typeinfer,
+)
+from .stmts import *  # noqa: F403
+from ._dialect import dialect as dialect

kirin/dialects/py/binop/_dialect.py

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

kirin/dialects/py/binop/interp.py

@@ -0,0 +1,60 @@
+from kirin import interp
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class PyMethodTable(interp.MethodTable):
+
+    @interp.impl(stmts.Add)
+    def add(self, interp, frame: interp.Frame, stmt: stmts.Add):
+        return (frame.get(stmt.lhs) + frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Sub)
+    def sub(self, interp, frame: interp.Frame, stmt: stmts.Sub):
+        return (frame.get(stmt.lhs) - frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Mult)
+    def mult(self, interp, frame: interp.Frame, stmt: stmts.Mult):
+        return (frame.get(stmt.lhs) * frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Div)
+    def div(self, interp, frame: interp.Frame, stmt: stmts.Div):
+        return (frame.get(stmt.lhs) / frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Mod)
+    def mod(self, interp, frame: interp.Frame, stmt: stmts.Mod):
+        return (frame.get(stmt.lhs) % frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.BitAnd)
+    def bit_and(self, interp, frame: interp.Frame, stmt: stmts.BitAnd):
+        return (frame.get(stmt.lhs) & frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.BitOr)
+    def bit_or(self, interp, frame: interp.Frame, stmt: stmts.BitOr):
+        return (frame.get(stmt.lhs) | frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.BitXor)
+    def bit_xor(self, interp, frame: interp.Frame, stmt: stmts.BitXor):
+        return (frame.get(stmt.lhs) ^ frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.LShift)
+    def lshift(self, interp, frame: interp.Frame, stmt: stmts.LShift):
+        return (frame.get(stmt.lhs) << frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.RShift)
+    def rshift(self, interp, frame: interp.Frame, stmt: stmts.RShift):
+        return (frame.get(stmt.lhs) >> frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.FloorDiv)
+    def floor_div(self, interp, frame: interp.Frame, stmt: stmts.FloorDiv):
+        return (frame.get(stmt.lhs) // frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.Pow)
+    def pow(self, interp, frame: interp.Frame, stmt: stmts.Pow):
+        return (frame.get(stmt.lhs) ** frame.get(stmt.rhs),)
+
+    @interp.impl(stmts.MatMult)
+    def mat_mult(self, interp, frame: interp.Frame, stmt: stmts.MatMult):
+        return (frame.get(stmt.lhs) @ frame.get(stmt.rhs),)

kirin/dialects/py/binop/julia.py

@@ -0,0 +1,33 @@
+from kirin import interp
+from kirin.emit.julia import EmitJulia, EmitStrFrame
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="emit.julia")
+class JuliaTable(interp.MethodTable):
+
+    @interp.impl(stmts.Add)
+    def emit_Add(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Add):
+        return emit.emit_binaryop(frame, "+", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Sub)
+    def emit_Sub(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Sub):
+        return emit.emit_binaryop(frame, "-", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Mult)
+    def emit_Mult(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Mult):
+        return emit.emit_binaryop(frame, "*", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Div)
+    def emit_Div(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Div):
+        return emit.emit_binaryop(frame, "/", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Mod)
+    def emit_Mod(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Mod):
+        return emit.emit_binaryop(frame, "%", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(stmts.Pow)
+    def emit_Pow(self, emit: EmitJulia, frame: EmitStrFrame, stmt: stmts.Pow):
+        return emit.emit_binaryop(frame, "^", stmt.lhs, stmt.rhs, stmt.result)

kirin/dialects/py/binop/lowering.py

@@ -0,0 +1,22 @@
+import ast
+
+from kirin import lowering, exceptions
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_BinOp(
+        self, state: lowering.LoweringState, node: ast.BinOp
+    ) -> lowering.Result:
+        lhs = state.visit(node.left).expect_one()
+        rhs = state.visit(node.right).expect_one()
+
+        if op := getattr(stmts, node.op.__class__.__name__, None):
+            stmt = op(lhs=lhs, rhs=rhs)
+        else:
+            raise exceptions.DialectLoweringError(f"unsupported binop {node.op}")
+        return lowering.Result(state.append_stmt(stmt))

kirin/dialects/py/binop/stmts.py

@@ -0,0 +1,79 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from ._dialect import dialect
+
+T = types.TypeVar("T")
+
+
+@statement
+class BinOp(ir.Statement):
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(T, print=False)
+    rhs: ir.SSAValue = info.argument(T, print=False)
+    result: ir.ResultValue = info.result(T)
+
+
+@statement(dialect=dialect)
+class Add(BinOp):
+    name = "add"
+
+
+@statement(dialect=dialect)
+class Sub(BinOp):
+    name = "sub"
+
+
+@statement(dialect=dialect)
+class Mult(BinOp):
+    name = "mult"
+
+
+@statement(dialect=dialect)
+class Div(BinOp):
+    name = "div"
+
+
+@statement(dialect=dialect)
+class Mod(BinOp):
+    name = "mod"
+
+
+@statement(dialect=dialect)
+class Pow(BinOp):
+    name = "pow"
+
+
+@statement(dialect=dialect)
+class LShift(BinOp):
+    name = "lshift"
+
+
+@statement(dialect=dialect)
+class RShift(BinOp):
+    name = "rshift"
+
+
+@statement(dialect=dialect)
+class BitAnd(BinOp):
+    name = "bitand"
+
+
+@statement(dialect=dialect)
+class BitOr(BinOp):
+    name = "bitor"
+
+
+@statement(dialect=dialect)
+class BitXor(BinOp):
+    name = "bitxor"
+
+
+@statement(dialect=dialect)
+class FloorDiv(BinOp):
+    name = "floordiv"
+
+
+@statement(dialect=dialect)
+class MatMult(BinOp):
+    name = "matmult"

kirin/dialects/py/binop/typeinfer.py

@@ -0,0 +1,108 @@
+from kirin import types, interp
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(stmts.Add, types.Float, types.Float)
+    @interp.impl(stmts.Add, types.Float, types.Int)
+    @interp.impl(stmts.Add, types.Int, types.Float)
+    def addf(self, interp, frame, stmt):
+        return (types.Float,)
+
+    @interp.impl(stmts.Add, types.Int, types.Int)
+    def addi(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.Sub, types.Float, types.Float)
+    @interp.impl(stmts.Sub, types.Float, types.Int)
+    @interp.impl(stmts.Sub, types.Int, types.Float)
+    def subf(self, *_):
+        return (types.Float,)
+
+    @interp.impl(stmts.Sub, types.Int, types.Int)
+    def subi(self, *_):
+        return (types.Int,)
+
+    @interp.impl(stmts.Mult, types.Float, types.Float)
+    @interp.impl(stmts.Mult, types.Float, types.Int)
+    @interp.impl(stmts.Mult, types.Int, types.Float)
+    def multf(self, *_):
+        return (types.Float,)
+
+    @interp.impl(stmts.Mult, types.Int, types.Int)
+    def multi(self, *_):
+        return (types.Int,)
+
+    @interp.impl(stmts.Div)
+    def divf(self, typeinfer_, frame, stmt):
+        return (types.Float,)
+
+    @interp.impl(stmts.Mod, types.Float, types.Float)
+    @interp.impl(stmts.Mod, types.Float, types.Int)
+    @interp.impl(stmts.Mod, types.Int, types.Float)
+    def modf(self, *_):
+        return (types.Float,)
+
+    @interp.impl(stmts.Mod, types.Int, types.Int)
+    def modi(self, *_):
+        return (types.Int,)
+
+    @interp.impl(stmts.BitAnd, types.Int, types.Int)
+    def bit_andi(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.BitAnd, types.Bool, types.Bool)
+    def bit_andb(self, interp, frame, stmt):
+        return (types.Bool,)
+
+    @interp.impl(stmts.BitOr, types.Int, types.Int)
+    def bit_ori(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.BitOr, types.Bool, types.Bool)
+    def bit_orb(self, interp, frame, stmt):
+        return (types.Bool,)
+
+    @interp.impl(stmts.BitXor, types.Int, types.Int)
+    def bit_xori(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.BitXor, types.Bool, types.Bool)
+    def bit_xorb(self, interp, frame, stmt):
+        return (types.Bool,)
+
+    @interp.impl(stmts.LShift, types.Int)
+    def lshift(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.RShift, types.Int)
+    def rshift(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.FloorDiv, types.Float, types.Float)
+    @interp.impl(stmts.FloorDiv, types.Int, types.Float)
+    @interp.impl(stmts.FloorDiv, types.Float, types.Int)
+    def floor_divf(self, interp, frame, stmt):
+        return (types.Float,)
+
+    @interp.impl(stmts.FloorDiv, types.Int, types.Int)
+    def floor_divi(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.Pow, types.Float, types.Float)
+    @interp.impl(stmts.Pow, types.Float, types.Int)
+    @interp.impl(stmts.Pow, types.Int, types.Float)
+    def powf(self, interp, frame, stmt):
+        return (types.Float,)
+
+    @interp.impl(stmts.Pow, types.Int, types.Int)
+    def powi(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(stmts.MatMult)
+    def mat_mult(self, interp, frame, stmt):
+        raise NotImplementedError("np.array @ np.array not implemented")

kirin/dialects/py/boolop.py

@@ -0,0 +1,84 @@
+"""Boolean operators for Python dialect.
+
+This module contains the dialect for the Python boolean operators, including:
+
+- The `And` and `Or` statement classes.
+- The lowering pass for the boolean operators.
+- The concrete implementation of the boolean operators.
+- The Julia emitter for the boolean operators.
+
+This dialect maps `ast.BoolOp` nodes to the `And` and `Or` statements.
+"""
+
+import ast
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.emit.julia import EmitJulia, EmitStrFrame
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("py.boolop")
+
+
+@statement
+class BoolOp(ir.Statement):
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(print=False)
+    rhs: ir.SSAValue = info.argument(print=False)
+    result: ir.ResultValue = info.result(types.Bool)
+
+
+@statement(dialect=dialect)
+class And(BoolOp):
+    name = "and"
+
+
+@statement(dialect=dialect)
+class Or(BoolOp):
+    name = "or"
+
+
+@dialect.register
+class PythonLowering(lowering.FromPythonAST):
+
+    def lower_BoolOp(
+        self, state: lowering.LoweringState, node: ast.BoolOp
+    ) -> lowering.Result:
+        lhs = state.visit(node.values[0]).expect_one()
+        match node.op:
+            case ast.And():
+                boolop = And
+            case ast.Or():
+                boolop = Or
+            case _:
+                raise DialectLoweringError(f"unsupported boolop {node.op}")
+
+        for value in node.values[1:]:
+            lhs = state.append_stmt(
+                boolop(lhs=lhs, rhs=state.visit(value).expect_one())
+            ).result
+        return lowering.Result(lhs)
+
+
+@dialect.register
+class BoolOpMethod(interp.MethodTable):
+
+    @interp.impl(And)
+    def and_(self, interp, frame: interp.Frame, stmt: And):
+        return (frame.get(stmt.lhs) and frame.get(stmt.rhs),)
+
+    @interp.impl(Or)
+    def or_(self, interp, frame: interp.Frame, stmt: Or):
+        return (frame.get(stmt.lhs) or frame.get(stmt.rhs),)
+
+
+@dialect.register(key="emit.julia")
+class JuliaTable(interp.MethodTable):
+
+    @interp.impl(And)
+    def emit_And(self, emit: EmitJulia, frame: EmitStrFrame, stmt: And):
+        return emit.emit_binaryop(frame, "&&", stmt.lhs, stmt.rhs, stmt.result)
+
+    @interp.impl(Or)
+    def emit_Or(self, emit: EmitJulia, frame: EmitStrFrame, stmt: Or):
+        return emit.emit_binaryop(frame, "||", stmt.lhs, stmt.rhs, stmt.result)

kirin/dialects/py/builtin.py

@@ -0,0 +1,78 @@
+"""builtin dialect for python builtins
+
+This dialect provides implementations for builtin functions like abs and sum.
+
+- Statements: `Abs`, `Sum`.
+- The lowering pass for the builtin functions.
+- The concrete implementation of the builtin functions.
+- The type inference implementation of the builtin functions.
+
+This dialect maps `ast.Call` nodes of builtin functions to the `Abs` and `Sum` statements.
+"""
+
+from ast import Call
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+
+dialect = ir.Dialect("py.builtin")
+
+T = types.TypeVar("T", bound=types.Int | types.Float)
+
+
+@statement(dialect=dialect)
+class Abs(ir.Statement):
+    name = "abs"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    value: ir.SSAValue = info.argument(T, print=False)
+    result: ir.ResultValue = info.result(T)
+
+
+@statement(dialect=dialect)
+class Sum(ir.Statement):
+    name = "sum"
+    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    value: ir.SSAValue = info.argument(types.Any, print=False)
+    result: ir.ResultValue = info.result(types.Any)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Call_abs(
+        self, state: lowering.LoweringState, node: Call
+    ) -> lowering.Result:
+        return lowering.Result(
+            state.append_stmt(Abs(state.visit(node.args[0]).expect_one()))
+        )
+
+    def lower_Call_sum(
+        self, state: lowering.LoweringState, node: Call
+    ) -> lowering.Result:
+        return lowering.Result(
+            state.append_stmt(Sum(state.visit(node.args[0]).expect_one()))
+        )
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Abs)
+    def abs(self, interp, frame: interp.Frame, stmt: Abs):
+        return (abs(frame.get(stmt.value)),)
+
+    @interp.impl(Sum)
+    def _sum(self, interp, frame: interp.Frame, stmt: Sum):
+        return (sum(frame.get(stmt.value)),)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Abs, types.Int)
+    def absi(self, interp, frame, stmt):
+        return (types.Int,)
+
+    @interp.impl(Abs, types.Float)
+    def absf(self, interp, frame, stmt):
+        return (types.Float,)