callbacks working

Author: makslevental

mlir_utils/dialects/ext/func.py

@@ -1,4 +1,5 @@
 import inspect
+from typing import Union, Optional
 
 from mlir.dialects.func import FuncOp, ReturnOp, CallOp
 from mlir.ir import (
@@ -8,7 +9,7 @@
     TypeAttr,
     FlatSymbolRefAttr,
     Type,
-    Location,
+    Value,
 )
 
 from mlir_utils.util import (
@@ -19,16 +20,90 @@
 )
 
 
+def call(
+    callee_or_results: Union[FuncOp, list[Type]],
+    arguments_or_callee: Union[list[Value], FlatSymbolRefAttr, str],
+    arguments: Optional[list] = None,
+    *,
+    call_op_ctor=CallOp.__base__,
+    loc=None,
+    ip=None,
+):
+    """Creates an call operation.
+
+    The constructor accepts three different forms:
+
+      1. A function op to be called followed by a list of arguments.
+      2. A list of result types, followed by the name of the function to be
+         called as string, following by a list of arguments.
+      3. A list of result types, followed by the name of the function to be
+         called as symbol reference attribute, followed by a list of arguments.
+
+    For example
+
+        f = func.FuncOp("foo", ...)
+        func.CallOp(f, [args])
+        func.CallOp([result_types], "foo", [args])
+
+    In all cases, the location and insertion point may be specified as keyword
+    arguments if not provided by the surrounding context managers.
+    """
+    if loc is None:
+        loc = get_user_code_loc()
+    if isinstance(callee_or_results, FuncOp.__base__):
+        if not isinstance(arguments_or_callee, (list, tuple)):
+            raise ValueError(
+                "when constructing a call to a function, expected "
+                + "the second argument to be a list of call arguments, "
+                + f"got {type(arguments_or_callee)}"
+            )
+        if arguments is not None:
+            raise ValueError(
+                "unexpected third argument when constructing a call" + "to a function"
+            )
+        return call_op_ctor(
+            callee_or_results.function_type.value.results,
+            FlatSymbolRefAttr.get(callee_or_results.sym_name.value),
+            arguments_or_callee,
+            loc=loc,
+            ip=ip,
+        )
+
+    if isinstance(arguments_or_callee, list):
+        raise ValueError(
+            "when constructing a call to a function by name, "
+            + "expected the second argument to be a string or a "
+            + f"FlatSymbolRefAttr, got {type(arguments_or_callee)}"
+        )
+
+    if isinstance(arguments_or_callee, FlatSymbolRefAttr):
+        return call_op_ctor(
+            callee_or_results, arguments_or_callee, arguments, loc=loc, ip=ip
+        )
+    elif isinstance(arguments_or_callee, str):
+        return call_op_ctor(
+            callee_or_results,
+            FlatSymbolRefAttr.get(arguments_or_callee),
+            arguments,
+            loc=loc,
+            ip=ip,
+        )
+    else:
+        raise ValueError(f"unexpected type {callee_or_results=}")
+
+
 class FuncBase:
     def __init__(
         self,
         body_builder,
         func_op_ctor,
         return_op_ctor,
         call_op_ctor,
+        return_types=None,
         sym_visibility=None,
         arg_attrs=None,
         res_attrs=None,
+        func_attrs=None,
         loc=None,
         ip=None,
     ):
@@ -40,6 +115,13 @@ def __init__(
         self.body_builder = body_builder
         self.func_name = self.body_builder.__name__
 
+        if return_types is None:
+            return_types = []
+        sig = inspect.signature(self.body_builder)
+        self.input_types, self.return_types, self.arg_locs = self.prep_func_types(
+            sig, return_types
+        )
+
         self.func_op_ctor = func_op_ctor
         self.return_op_ctor = return_op_ctor
         self.call_op_ctor = call_op_ctor
@@ -48,26 +130,63 @@ def __init__(
         )
         self.arg_attrs = arg_attrs
         self.res_attrs = res_attrs
+        if func_attrs is None:
+            func_attrs = {}
+        self.func_attrs = func_attrs
         self.loc = loc
         self.ip = ip or InsertionPoint.current
-        self.emitted = False
+        self._func_op = None
+
+        if self._is_decl():
+            assert len(self.input_types) == len(
+                sig.parameters
+            ), f"func decl needs all input types annotated"
+            self.sym_visibility = StringAttr.get("private")
+            self.emit()
+
+    def _is_decl(self):
+        # magic constant found from looking at the code for an empty fn
+        return self.body_builder.__code__.co_code == b"\x97\x00d\x00S\x00"
 
     def __str__(self):
         return str(f"{self.__class__} {self.__dict__}")
 
+    def prep_func_types(self, sig, return_types):
+        assert not (
+            not sig.return_annotation is inspect.Signature.empty
+            and len(return_types) > 0
+        ), f"func can use return annotation or explicit return_types but not both"
+        return_types = (
+            sig.return_annotation
+            if not sig.return_annotation is inspect.Signature.empty
+            else return_types
+        )
+        if not isinstance(return_types, (tuple, list)):
+            return_types = [return_types]
+        return_types = list(return_types)
+        assert all(
+            isinstance(r, Type) for r in return_types
+        ), f"all return types must be mlir types {return_types=}"
+
+        input_types = [
+            p.annotation
+            for p in sig.parameters.values()
+            if not p.annotation is inspect.Signature.empty
+        ]
+        assert all(
+            isinstance(r, Type) for r in input_types
+        ), f"all input types must be mlir types {input_types=}"
+        return input_types, return_types, [get_user_code_loc()] * len(sig.parameters)
+
     def body_builder_wrapper(self, *call_args):
-        sig = inspect.signature(self.body_builder)
-        implicit_return = sig.return_annotation is inspect._empty
-        input_types = [p.annotation for p in sig.parameters.values()]
-        if not (
-            len(input_types) == len(sig.parameters)
-            and all(isinstance(t, Type) for t in input_types)
-        ):
+        if len(call_args) == 0:
+            input_types = self.input_types
+        else:
             input_types = [a.type for a in call_args]
         function_type = TypeAttr.get(
             FunctionType.get(
                 inputs=input_types,
-                results=[] if implicit_return else sig.return_annotation,
+                results=self.return_types,
             )
         )
         func_op = self.func_op_ctor(
@@ -79,8 +198,10 @@ def body_builder_wrapper(self, *call_args):
             loc=self.loc,
             ip=self.ip,
         )
-        arg_locs = [get_user_code_loc()] * len(sig.parameters)
-        func_op.regions[0].blocks.append(*input_types, arg_locs=arg_locs)
+        if self._is_decl():
+            return self.return_types, input_types, func_op
+
+        func_op.regions[0].blocks.append(*input_types, arg_locs=self.arg_locs)
         with InsertionPoint(func_op.regions[0].blocks[0]):
             results = get_result_or_results(
                 self.body_builder(
@@ -94,31 +215,23 @@ def body_builder_wrapper(self, *call_args):
                     results = [results]
             else:
                 results = []
+
             self.return_op_ctor(results)
+        return_types = [r.type for r in results]
+        return return_types, input_types, func_op
 
-        return results, input_types, func_op
-
-    def emit(self):
-        self.results, input_types, func_op = self.body_builder_wrapper()
-        return_types = [v.type for v in self.results]
-        function_type = FunctionType.get(inputs=input_types, results=return_types)
-        func_op.attributes["function_type"] = TypeAttr.get(function_type)
-        self.emitted = True
-        # this is the func op itself (funcs never have a resulting ssa value)
-        return maybe_cast(get_result_or_results(func_op))
-
-    def __call__(self, *call_args, loc: Location = None):
-        if loc is None:
-            loc = get_user_code_loc()
-        if not self.emitted:
-            self.emit()
-        call_op = self.call_op_ctor(
-            [r.type for r in self.results],
-            FlatSymbolRefAttr.get(self.func_name),
-            call_args,
-            loc=loc,
-        )
-        return maybe_cast(get_result_or_results(call_op))
+    def emit(self) -> FuncOp:
+        if self._func_op is None:
+            return_types, input_types, func_op = self.body_builder_wrapper()
+            function_type = FunctionType.get(inputs=input_types, results=return_types)
+            func_op.attributes["function_type"] = TypeAttr.get(function_type)
+            for k, v in self.func_attrs.items():
+                func_op.attributes[k] = v
+            self._func_op = func_op
+        return self._func_op
+
+    def __call__(self, *call_args):
+        return call(self.emit(), call_args)
 
 
 @make_maybe_no_args_decorator
@@ -128,9 +241,10 @@ def func(
     sym_visibility=None,
     arg_attrs=None,
     res_attrs=None,
+    func_attrs=None,
     loc=None,
     ip=None,
-):
+) -> FuncBase:
     if loc is None:
         loc = get_user_code_loc()
     return FuncBase(
@@ -141,48 +255,7 @@ def func(
         sym_visibility=sym_visibility,
         arg_attrs=arg_attrs,
         res_attrs=res_attrs,
+        func_attrs=func_attrs,
         loc=loc,
         ip=ip,
     )
-
-
-def call(symbol_name, call_args, return_types, *, loc=None, ip=None):
-    if loc is None:
-        loc = get_user_code_loc()
-    return maybe_cast(
-        get_result_or_results(
-            CallOp.__base__(
-                return_types,
-                FlatSymbolRefAttr.get(symbol_name),
-                call_args,
-                loc=loc,
-                ip=ip,
-            )
-        )
-    )
-
-
-def declare(
-    symbol_name,
-    input_types: list,
-    result_types=None,
-    func_op_ctor=FuncOp,
-):
-    if result_types is None:
-        result_types = []
-    assert all(
-        isinstance(a, Type) for a in input_types
-    ), f"wrong func args {input_types}"
-    assert all(
-        isinstance(a, Type) for a in result_types
-    ), f"wrong func results {result_types}"
-
-    function_type = FunctionType.get(inputs=input_types, results=result_types)
-    sym_name = func_op_ctor(
-        name=symbol_name, type=function_type, visibility="private"
-    ).sym_name
-
-    def callable(*call_args):
-        return call(sym_name.value, call_args, result_types)
-
-    return callable