intel/llvm
1
0
Fork 0
mirror of https://github.com/intel/llvm.git synced 2026-01-14 03:50:17 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Rolf Morel
2025-11-07 14:04:53 +00:00
committed by GitHub
parent 3ee2f07e17
commit d78e0ded52
2 changed files with 302 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
mlir/python/mlir/dialects/transform/__init__.py
View File

@@ -39,16 +39,32 @@ class CastOp(CastOp):
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
@@ -56,6 +72,25 @@ class ApplyPatternsOp(ApplyPatternsOp):
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 @@ class GetParentOp(GetParentOp):
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 @@ class GetParentOp(GetParentOp):
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,6 +119,31 @@ class GetParentOp(GetParentOp):
)
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__(
@@ -89,17 +151,32 @@ class MergeHandlesOp(MergeHandlesOp):
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 @@ class ReplicateOp(ReplicateOp):
)
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 @@ class SequenceOp(SequenceOp):
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 @@ class SequenceOp(SequenceOp):
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 @@ class NamedSequenceOp(NamedSequenceOp):
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 @@ class YieldOp(YieldOp):
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 @@ class ApplyRegisteredPassOp(ApplyRegisteredPassOp):
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):