[Mlir-commits] [mlir] 1dccdf7 - [mlir][linalg][transform][python] Add type arg to MatchOp extension.

[Ingo Müller]

Author: Ingo Müller
Date: 2023-07-19T09:15:41Z
New Revision: 1dccdf7f49a0cdad7121913c900ce9cb8b6e9fdc

URL: https://github.com/llvm/llvm-project/commit/1dccdf7f49a0cdad7121913c900ce9cb8b6e9fdc
DIFF: https://github.com/llvm/llvm-project/commit/1dccdf7f49a0cdad7121913c900ce9cb8b6e9fdc.diff

LOG: [mlir][linalg][transform][python] Add type arg to MatchOp extension.

The extension class to MatchOp has a class method called match_op_names.
The previous version of that function did not allow to specify the
result type. This, however, may be useful/necessary if the op consuming
the resulting handle requires a particular type (such as the
bufferization.EmptyTensorToAllocTensorOp). This patch adds an overload
to match_op_names that allows to specify the result type.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D155567

Added: 
    

Modified: 
    mlir/python/mlir/dialects/_structured_transform_ops_ext.py
    mlir/test/python/dialects/transform_structured_ext.py

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/dialects/_structured_transform_ops_ext.py b/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
index b754034c8a4809..640730997f93e1 100644
--- a/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
+++ b/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
@@ -85,17 +85,52 @@ def __init__(
 class MatchOp:
     """Specialization for MatchOp class."""
 
+    @overload
     @classmethod
     def match_op_names(
-        MatchOp,
+        cls,
         target: Union[Operation, Value],
         names: Sequence[str],
+        *,
         loc=None,
         ip=None,
     ):
-        pdl_operation_type = pdl.OperationType.get()
-        return MatchOp(
-            pdl_operation_type,
+       ...
+
+    @overload
+    @classmethod
+    def match_op_names(
+        cls,
+        result_type: Type,
+        target: Union[Operation, Value],
+        names: Sequence[str],
+        *,
+        loc=None,
+        ip=None,
+    ):
+       ...
+
+    @classmethod
+    def match_op_names(
+        cls,
+        result_type_or_target: Union[Type, Operation, Value],
+        target_or_names: Union[Operation, Value,  Sequence[str]],
+        names_or_none: Optional[Sequence[str]] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        if isinstance(result_type_or_target, Type):
+           result_type = result_type_or_target
+           target = target_or_names
+           names = names_or_none
+        else:
+           result_type = transform.AnyOpType.get()
+           target = result_type_or_target
+           names = target_or_names
+
+        return cls(
+            result_type,
             _get_op_result_or_value(target),
             ops=ArrayAttr.get(list(map(lambda s: StringAttr.get(s), names))),
             loc=loc,

diff  --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index 2dfae47bdfb492..03a47166c824c6 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -57,6 +57,38 @@ def testInterchange():
     # CHECK: iterator_interchange = [1, 0]
 
 
+ at run
+def testMatchOpNames():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.MatchOp.match_op_names(sequence.bodyTarget, ["test.dummy"])
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testMatchOpNames
+    # CHECK: transform.structured.match ops
+    # CHECK-SAME: ["test.dummy"]
+    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
+
+
+ at run
+def testMatchOpNamesTyped():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.MatchOp.match_op_names(
+            transform.OperationType.get("test.dummy"),
+            sequence.bodyTarget,
+            ["test.dummy"],
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testMatchOpNamesTyped
+    # CHECK: transform.structured.match ops
+    # CHECK-SAME: ["test.dummy"]
+    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
+
+
 @run
 def testMultitileSizes():
     sequence = transform.SequenceOp(