add region op decorator

Author: makslevental

mlir_utils/dialects/.gitignore

@@ -1,22 +1,22 @@
-arith.py
-async_dialect.py
-bufferization.py
-builtin.py
-cf.py
-complex.py
-func.py
-gpu.py
-linalg.py
-math.py
-memref.py
-ml_program.py
-pdl.py
-quant.py
-scf.py
-shape.py
-sparse_tensor.py
-tensor.py
-torch.py
-tosa.py
-transform.py
-vector.py
+/arith.py
+/async_dialect.py
+/bufferization.py
+/builtin.py
+/cf.py
+/complex.py
+/func.py
+/gpu.py
+/linalg.py
+/math.py
+/memref.py
+/ml_program.py
+/pdl.py
+/quant.py
+/scf.py
+/shape.py
+/sparse_tensor.py
+/tensor.py
+/torch.py
+/tosa.py
+/transform.py
+/vector.py

mlir_utils/dialects/ext/__init__.py

@@ -0,0 +1,57 @@
+import inspect
+from functools import wraps
+
+from mlir.dialects.func import FuncOp, ReturnOp
+from mlir.ir import InsertionPoint, FunctionType, StringAttr, TypeAttr
+
+from mlir_utils.dialects.util import (
+    get_result_or_results,
+    make_maybe_no_args_decorator,
+)
+
+
+@make_maybe_no_args_decorator
+def func(sym_visibility=None, arg_attrs=None, res_attrs=None, loc=None, ip=None):
+    ip = ip or InsertionPoint.current
+
+    def builder_wrapper(body_builder):
+        @wraps(body_builder)
+        def wrapper(*args):
+            sig = inspect.signature(body_builder)
+            implicit_return = sig.return_annotation is inspect._empty
+            function_type = TypeAttr.get(
+                FunctionType.get(
+                    inputs=[a.type for a in args],
+                    results=[] if implicit_return else sig.return_annotation,
+                )
+            )
+            # FuncOp is extended but we do really want the base
+            op = FuncOp.__base__(
+                body_builder.__name__,
+                function_type,
+                sym_visibility=StringAttr.get(str(sym_visibility))
+                if sym_visibility is not None
+                else None,
+                arg_attrs=arg_attrs,
+                res_attrs=res_attrs,
+                loc=loc,
+                ip=ip,
+            )
+            op.regions[0].blocks.append(*[a.type for a in args])
+            with InsertionPoint(op.regions[0].blocks[0]):
+                r = get_result_or_results(
+                    body_builder(*op.regions[0].blocks[0].arguments)
+                )
+                if r is not None:
+                    if isinstance(r, (tuple, list)):
+                        ReturnOp(list(r))
+                    else:
+                        ReturnOp([r])
+                else:
+                    ReturnOp([])
+                return r
+
+        # wrapper.op = op
+        return wrapper
+
+    return builder_wrapper

mlir_utils/dialects/ext/func.py

@@ -21,7 +21,9 @@ def ast_call(name, args=None, keywords=None):
     )
 
 
-def make_fun(op_name, op_class):
+# TODO(max): ops that have symboltables need to be classes but that requires some upstream support for statically
+# identifying such ops
+def generate_free_fun(op_class):
     _mod = ast.parse(dedent(inspect.getsource(op_class.__init__)))
     init_fn = next(n for n in _mod.body if isinstance(n, ast.FunctionDef))
     args = init_fn.args
@@ -39,17 +41,33 @@ def make_fun(op_name, op_class):
 
     for a in args.args + args.kwonlyargs:
         a.annotation = None
-    ret = ast.parse(
-        f"return get_result_or_results({ast.unparse(ast_call(op_name, args.args, keywords))})"
-    ).body[0]
     fun_name = op_class.OPERATION_NAME.split(".")[-1]
     if keyword.iskeyword(fun_name):
         fun_name = fun_name + "_"
+    op_class_name = op_class.__name__
+    body = []
+    if len(args.args) == 1 and args.args[0].arg == "results_":
+        args.defaults.append(ast.Constant(None))
+        body += [ast.parse("results_ = results_ or []").body[0]]
+    if (
+        hasattr(op_class, "_ODS_REGIONS")
+        and op_class._ODS_REGIONS[0] == 1
+        and not op_class.OPERATION_NAME.startswith("linalg")
+    ):
+        decorator_list = [ast.Name(id="region_op", ctx=ast.Load())]
+        body += [ast.Return([ast_call(op_class_name, args.args, keywords)])]
+    else:
+        decorator_list = []
+        body += [
+            ast.parse(
+                f"return get_result_or_results({ast.unparse(ast_call(op_class_name, args.args, keywords))})"
+            ).body[0]
+        ]
     n = ast.FunctionDef(
         name=fun_name,
         args=copy.deepcopy(args),
-        body=[ret],
-        decorator_list=[],
+        body=body,
+        decorator_list=decorator_list,
     )
     ast.fix_missing_locations(n)
     return n
@@ -58,6 +76,7 @@ def make_fun(op_name, op_class):
 def generate_trampoline(input_module, output_file_path, skips=None):
     import mlir_utils
     from mlir_utils.dialects.util import get_result_or_results
+    import mlir.dialects._ods_common
 
     if skips is None:
         skips = set()
@@ -69,6 +88,9 @@ def generate_trampoline(input_module, output_file_path, skips=None):
             and hasattr(obj, "OPERATION_NAME")
             and obj.__name__ not in skips
         ):
+            if obj.__module__ == mlir.dialects._ods_common.__name__:
+                # these are extension classes and we should wrap the generated class instead
+                obj = obj.__base__
             if not inspect.isfunction(obj.__init__):
                 # some builders don't have any __init__ but inherit from opview
                 continue
@@ -77,11 +99,14 @@ def generate_trampoline(input_module, output_file_path, skips=None):
     if not len(init_funs):
         return
 
-    functions = [make_fun(n, s) for n, s in init_funs.items()]
+    functions = [
+        generate_free_fun(op_class)
+        for op_class in sorted(init_funs.values(), key=lambda o: o.__name__)
+    ]
 
     ods_imports = ast.ImportFrom(
         module=mlir_utils.dialects.util.__name__,
-        names=[ast.alias(get_result_or_results.__name__)],
+        names=[ast.alias(get_result_or_results.__name__), ast.alias("region_op")],
         level=0,
     )
     op_imports = ast.ImportFrom(

mlir_utils/dialects/generate_trampolines.py

@@ -1,11 +1,93 @@
+from functools import wraps
+
+import numpy as np
+from mlir.dialects import arith
 from mlir.dialects._ods_common import get_op_result_or_value, get_op_results_or_values
+from mlir.ir import (
+    InsertionPoint,
+    IntegerType,
+    F64Type,
+    RankedTensorType,
+    IndexType,
+    F16Type,
+    F32Type,
+)
 
 
 def get_result_or_results(op):
+    if op is None:
+        return
     return (
         get_op_results_or_values(op)
         if len(op.operation.results) > 1
         else get_op_result_or_value(op)
         if len(op.operation.results) > 0
         else None
     )
+
+
+def make_maybe_no_args_decorator(decorator):
+    @wraps(decorator)
+    def maybe_no_args(*args, **kwargs):
+        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+            return decorator()(args[0])
+        else:
+            return decorator(*args, **kwargs)
+
+    return maybe_no_args
+
+
+# builds the decorator
+def region_op(op_constructor):
+    # the decorator itself
+    def op_decorator(*args, **kwargs):
+        op = op_constructor(*args, **kwargs)
+
+        def builder_wrapper(body_builder):
+            @wraps(body_builder)
+            def wrapper(*args):
+                # add a block with block args having types ...
+                op.regions[0].blocks.append(*[a.type for a in args])
+                with InsertionPoint(op.regions[0].blocks[0]):
+                    return get_result_or_results(
+                        body_builder(*op.regions[0].blocks[0].arguments)
+                    )
+
+            wrapper.op = op
+            return wrapper
+
+        return builder_wrapper
+
+    return make_maybe_no_args_decorator(op_decorator)
+
+
+def infer_mlir_type(
+    py_val: int | float | bool | np.ndarray,
+) -> IntegerType | F64Type | RankedTensorType:
+    if isinstance(py_val, bool):
+        return IntegerType.get_signless(1)
+    elif isinstance(py_val, int):
+        return IntegerType.get_signless(64)
+    elif isinstance(py_val, float):
+        return F64Type.get()
+    elif isinstance(py_val, np.ndarray):
+        dtype_ = {
+            np.int8: IntegerType.get_signless(8),
+            np.int16: IntegerType.get_signless(16),
+            np.int32: IntegerType.get_signless(32),
+            np.int64: IntegerType.get_signless(64),
+            np.uintp: IndexType.get(),
+            np.longlong: IndexType.get(),
+            np.float16: F16Type.get(),
+            np.float32: F32Type.get(),
+            np.float64: F64Type.get(),
+        }[py_val.dtype.type]
+        return RankedTensorType.get(py_val.shape, dtype_)
+    else:
+        raise NotImplementedError(
+            f"Unsupported Python value {py_val=} with type {type(py_val)}"
+        )
+
+
+def constant(py_val: int | float | bool | np.ndarray):
+    return arith.ConstantOp(infer_mlir_type(py_val), py_val)

mlir_utils/dialects/util.py

@@ -25,14 +25,14 @@ def filecheck(correct: str, module):
     assert Path(filecheck_path).exists() is not None, "couldn't find FileCheck"
 
     correct = dedent(correct)
+    correct_with_checks = main(correct).replace("CHECK:", "CHECK-NEXT:")
     op = dedent(str(module).strip())
     with tempfile.NamedTemporaryFile() as tmp:
-        correct_with_checks = main(correct)
         tmp.write(correct_with_checks.encode())
         tmp.flush()
         p = Popen([filecheck_path, tmp.name], stdout=PIPE, stdin=PIPE, stderr=PIPE)
         out, err = map(lambda o: o.decode(), p.communicate(input=op.encode()))
-        if len(err):
+        if p.returncode:
             raise ValueError(err)
 
 

mlir_utils/testing/testing.py

@@ -0,0 +1,93 @@
+from textwrap import dedent
+
+import pytest
+
+from mlir_utils.dialects.ext.func import func
+from mlir_utils.dialects.memref import alloca_scope, return_
+from mlir_utils.dialects.scf import execute_region, yield_
+from mlir_utils.dialects.util import constant
+
+# noinspection PyUnresolvedReferences
+from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
+
+# needed since the fix isn't defined here nor conftest.py
+pytest.mark.usefixtures("ctx")
+
+
+def test_simple_region_op(ctx: MLIRContext):
+    @execute_region([])
+    def demo_region():
+        one = constant(1)
+        yield_()
+
+    demo_region()
+
+    ctx.module.operation.verify()
+    filecheck(
+        dedent(
+            """\
+    module {
+      scf.execute_region {
+        %c1_i64 = arith.constant 1 : i64
+        scf.yield
+      }
+    }
+    """
+        ),
+        ctx.module,
+    )
+
+
+def test_no_args_decorator(ctx: MLIRContext):
+    @alloca_scope([])
+    def demo_scope1():
+        one = constant(1)
+        return_()
+
+    @alloca_scope
+    def demo_scope2():
+        one = constant(2)
+        return_()
+
+    demo_scope1()
+    demo_scope2()
+
+    ctx.module.operation.verify()
+    filecheck(
+        dedent(
+            """\
+    module {
+      memref.alloca_scope  {
+        %c1_i64 = arith.constant 1 : i64
+      }
+      memref.alloca_scope  {
+        %c2_i64 = arith.constant 2 : i64
+      }
+    }
+    """
+        ),
+        ctx.module,
+    )
+
+
+def test_func(ctx: MLIRContext):
+    @func
+    def demo_fun1():
+        one = constant(1)
+        return
+
+    demo_fun1()
+    ctx.module.operation.verify()
+    filecheck(
+        dedent(
+            """\
+    module {
+      func.func @demo_fun1() {
+        %c1_i64 = arith.constant 1 : i64
+        return
+      }
+    }
+    """
+        ),
+        ctx.module,
+    )