[MLIR][Transform][Python] Sync derived classes and their wrappers (#166871)

Updates the derived Op-classes for the main transform ops to have all
the arguments, etc, from the auto-generated classes. Additionally
updates and adds missing snake_case wrappers for the derived classes
which shadow the snake_case wrappers of the auto-generated classes,
which were hitherto exposed alongside the derived classes.

Author: rolfmorel

mlir/python/mlir/dialects/transform/__init__.py

@@ -39,23 +39,58 @@ def __init__(
         super().__init__(result_type, _get_op_result_or_value(target), loc=loc, ip=ip)
 
 
+def cast(
+    result_type: Type, target: Union[Operation, Value], *, loc=None, ip=None
+) -> OpResult:
+    return CastOp(result_type=result_type, target=target, loc=loc, ip=ip).result
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class ApplyPatternsOp(ApplyPatternsOp):
     def __init__(
         self,
         target: Union[Operation, Value, OpView],
+        apply_cse: bool = False,
+        max_iterations: Optional[Union[IntegerAttr, int]] = None,
+        max_num_rewrites: Optional[Union[IntegerAttr, int]] = None,
         *,
         loc=None,
         ip=None,
     ):
-        super().__init__(target, loc=loc, ip=ip)
+        super().__init__(
+            target,
+            apply_cse=apply_cse,
+            max_iterations=max_iterations,
+            max_num_rewrites=max_num_rewrites,
+            loc=loc,
+            ip=ip,
+        )
         self.regions[0].blocks.append()
 
     @property
     def patterns(self) -> Block:
         return self.regions[0].blocks[0]
 
 
+def apply_patterns(
+    target: Union[Operation, Value, OpView],
+    apply_cse: bool = False,
+    max_iterations: Optional[Union[IntegerAttr, int]] = None,
+    max_num_rewrites: Optional[Union[IntegerAttr, int]] = None,
+    *,
+    loc=None,
+    ip=None,
+) -> ApplyPatternsOp:
+    return ApplyPatternsOp(
+        target=target,
+        apply_cse=apply_cse,
+        max_iterations=max_iterations,
+        max_num_rewrites=max_num_rewrites,
+        loc=loc,
+        ip=ip,
+    )
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class GetParentOp(GetParentOp):
     def __init__(
@@ -64,6 +99,7 @@ def __init__(
         target: Union[Operation, Value],
         *,
         isolated_from_above: bool = False,
+        allow_empty_results: bool = False,
         op_name: Optional[str] = None,
         deduplicate: bool = False,
         nth_parent: int = 1,
@@ -74,6 +110,7 @@ def __init__(
             result_type,
             _get_op_result_or_value(target),
             isolated_from_above=isolated_from_above,
+            allow_empty_results=allow_empty_results,
             op_name=op_name,
             deduplicate=deduplicate,
             nth_parent=nth_parent,
@@ -82,24 +119,64 @@ def __init__(
         )
 
 
+def get_parent_op(
+    result_type: Type,
+    target: Union[Operation, Value],
+    *,
+    isolated_from_above: bool = False,
+    allow_empty_results: bool = False,
+    op_name: Optional[str] = None,
+    deduplicate: bool = False,
+    nth_parent: int = 1,
+    loc=None,
+    ip=None,
+) -> OpResult:
+    return GetParentOp(
+        result_type=result_type,
+        target=target,
+        isolated_from_above=isolated_from_above,
+        allow_empty_results=allow_empty_results,
+        op_name=op_name,
+        deduplicate=deduplicate,
+        nth_parent=nth_parent,
+        loc=loc,
+        ip=ip,
+    ).result
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MergeHandlesOp(MergeHandlesOp):
     def __init__(
         self,
         handles: Sequence[Union[Operation, Value]],
         *,
         deduplicate: bool = False,
+        results: Optional[Sequence[Type]] = None,
         loc=None,
         ip=None,
     ):
         super().__init__(
             [_get_op_result_or_value(h) for h in handles],
             deduplicate=deduplicate,
+            results=results,
             loc=loc,
             ip=ip,
         )
 
 
+def merge_handles(
+    handles: Sequence[Union[Operation, Value]],
+    *,
+    deduplicate: bool = False,
+    results: Optional[Sequence[Type]] = None,
+    loc=None,
+    ip=None,
+) -> OpResult:
+    return MergeHandlesOp(
+        handles=handles, deduplicate=deduplicate, results=results, loc=loc, ip=ip
+    ).result
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class ReplicateOp(ReplicateOp):
     def __init__(
@@ -119,16 +196,31 @@ def __init__(
         )
 
 
+def replicate(
+    pattern: Union[Operation, Value],
+    handles: Sequence[Union[Operation, Value]],
+    *,
+    loc=None,
+    ip=None,
+) -> Union[OpResult, OpResultList, ReplicateOp]:
+    op = ReplicateOp(pattern=pattern, handles=handles, loc=loc, ip=ip)
+    results = op.results
+    return results if len(results) > 1 else (results[0] if len(results) == 1 else op)
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class SequenceOp(SequenceOp):
     def __init__(
         self,
-        failure_propagation_mode,
+        failure_propagation_mode: FailurePropagationMode,
         results: Sequence[Type],
         target: Union[Operation, Value, Type],
         extra_bindings: Optional[
             Union[Sequence[Value], Sequence[Type], Operation, OpView]
         ] = None,
+        *,
+        loc=None,
+        ip=None,
     ):
         root = (
             _get_op_result_or_value(target)
@@ -155,6 +247,8 @@ def __init__(
             failure_propagation_mode=failure_propagation_mode,
             root=root,
             extra_bindings=extra_bindings,
+            loc=loc,
+            ip=ip,
         )
         self.regions[0].blocks.append(*tuple([root_type] + extra_binding_types))
 
@@ -171,16 +265,42 @@ def bodyExtraArgs(self) -> BlockArgumentList:
         return self.body.arguments[1:]
 
 
+def sequence(
+    failure_propagation_mode: FailurePropagationMode,
+    results: Sequence[Type],
+    target: Union[Operation, Value, Type],
+    extra_bindings: Optional[
+        Union[Sequence[Value], Sequence[Type], Operation, OpView]
+    ] = None,
+    *,
+    loc=None,
+    ip=None,
+) -> Union[OpResult, OpResultList, SequenceOp]:
+    op = SequenceOp(
+        results=results,
+        failure_propagation_mode=failure_propagation_mode,
+        extra_bindings=extra_bindings,
+        target=target,
+        loc=loc,
+        ip=ip,
+    )
+    results = op.results
+    return results if len(results) > 1 else (results[0] if len(results) == 1 else op)
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class NamedSequenceOp(NamedSequenceOp):
     def __init__(
         self,
-        sym_name,
+        sym_name: Union[str, SymbolRefAttr],
         input_types: Sequence[Type],
         result_types: Sequence[Type],
-        sym_visibility=None,
-        arg_attrs=None,
-        res_attrs=None,
+        *,
+        sym_visibility: Optional[Union[str, StringAttr]] = None,
+        arg_attrs: Optional[Union[Sequence[dict], "DictArrayAttr"]] = None,
+        res_attrs: Optional[Union[Sequence[dict], "DictArrayAttr"]] = None,
+        loc=None,
+        ip=None,
     ):
         function_type = FunctionType.get(input_types, result_types)
         super().__init__(
@@ -205,6 +325,29 @@ def bodyExtraArgs(self) -> BlockArgumentList:
         return self.body.arguments[1:]
 
 
+def named_sequence(
+    sym_name: Union[str, SymbolRefAttr],
+    input_types: Sequence[Type],
+    result_types: Sequence[Type],
+    *,
+    sym_visibility: Optional[Union[str, StringAttr]] = None,
+    arg_attrs: Optional[Union[Sequence[dict], "DictArrayAttr"]] = None,
+    res_attrs: Optional[Union[Sequence[dict], "DictArrayAttr"]] = None,
+    loc=None,
+    ip=None,
+) -> NamedSequenceOp:
+    return NamedSequenceOp(
+        sym_name=sym_name,
+        input_types=input_types,
+        result_types=result_types,
+        sym_visibility=sym_visibility,
+        arg_attrs=arg_attrs,
+        res_attrs=res_attrs,
+        loc=loc,
+        ip=ip,
+    )
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class YieldOp(YieldOp):
     def __init__(
@@ -219,6 +362,12 @@ def __init__(
         super().__init__(_get_op_results_or_values(operands), loc=loc, ip=ip)
 
 
+def yield_(
+    operands: Optional[Union[Operation, Sequence[Value]]] = None, *, loc=None, ip=None
+) -> YieldOp:
+    return YieldOp(operands=operands, loc=loc, ip=ip)
+
+
 OptionValueTypes = Union[
     Sequence["OptionValueTypes"], Attribute, Value, Operation, OpView, str, int, bool
 ]
@@ -247,7 +396,7 @@ def __init__(
         def option_value_to_attr(value):
             nonlocal cur_param_operand_idx
             if isinstance(value, (Value, Operation, OpView)):
-                dynamic_options.append(_get_op_result_or_value(value))
+                dynamic_options.append(value)
                 cur_param_operand_idx += 1
                 return ParamOperandAttr(cur_param_operand_idx - 1, context)
             elif isinstance(value, Attribute):