[Mlir-commits] [mlir] [mlir][python] remove mixins (PR #68853)

[Maksim Levental]

================
@@ -2,4 +2,118 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+from typing import Union
+
 from ._ml_program_ops_gen import *
+from ._ml_program_ops_gen import _Dialect
+
+try:
+    from ..ir import *
+    from ._ods_common import (
+        get_default_loc_context as _get_default_loc_context,
+        _cext as _ods_cext,
+    )
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+
+ARGUMENT_ATTRIBUTE_NAME = "arg_attrs"
+RESULT_ATTRIBUTE_NAME = "res_attrs"
+
+
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class FuncOp(FuncOp):
+    """Specialization for the func op class."""
+
+    def __init__(
+        self, name, type, *, visibility=None, body_builder=None, loc=None, ip=None
+    ):
+        """
+        Create a FuncOp with the provided `name`, `type`, and `visibility`.
+        - `name` is a string representing the function name.
+        - `type` is either a FunctionType or a pair of list describing inputs and
+          results.
+        - `visibility` is a string matching `public`, `private`, or `nested`. None
+          implies private visibility.
+        - `body_builder` is an optional callback, when provided a new entry block
+          is created and the callback is invoked with the new op as argument within
+          an InsertionPoint context already set for the block. The callback is
+          expected to insert a terminator in the block.
+        """
+        sym_name = StringAttr.get(str(name))
+
+        # If the type is passed as a tuple, build a FunctionType on the fly.
+        if isinstance(type, tuple):
+            type = FunctionType.get(inputs=type[0], results=type[1])
+
+        type = TypeAttr.get(type)
+        sym_visibility = (
+            StringAttr.get(str(visibility)) if visibility is not None else None
+        )
+        super().__init__(sym_name, type, sym_visibility=sym_visibility, loc=loc, ip=ip)
+        if body_builder:
+            entry_block = self.add_entry_block()
+            with InsertionPoint(entry_block):
+                body_builder(self)
+
+    @property
+    def is_external(self):
+        return len(self.regions[0].blocks) == 0
+
+    @property
+    def body(self):
+        return self.regions[0]
+
+    @property
+    def type(self):
+        return FunctionType(TypeAttr(self.attributes["function_type"]).value)
+
+    @property
+    def visibility(self):
+        return self.attributes["sym_visibility"]
+
+    @property
+    def name(self) -> StringAttr:
----------------
makslevental wrote:

These are indeed fragile I agree but they're also opt-in (to some extent): you can always import the `<DIALECT>_ops_gen.py` file directly. 

Additionally we could defend against drift by making sure each such property is exercised in a test (I'm hoping most already are...).

But adding some of these to the generated builders should be straightforward - I'll take a shot at it soon.

https://github.com/llvm/llvm-project/pull/68853