[Mlir-commits] [mlir] [MLIR][Python] Add structured.fuseop to generator. (PR #120601)

[Hugo Trachino]

https://github.com/nujaa updated https://github.com/llvm/llvm-project/pull/120601

>From dd7406b21c30e98b86c39a07ab4c99cd4b1f7dda Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Thu, 19 Dec 2024 19:30:52 +0800
Subject: [PATCH 1/3] [MLIR][Python] Add structured.fuseop to generator.

---
 .../mlir/dialects/transform/structured.py     | 52 +++++++++++++++++++
 .../dialects/transform_structured_ext.py      | 21 ++++++++
 2 files changed, 73 insertions(+)

diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 9121aa8e40237b..e715872a16ce48 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -139,6 +139,58 @@ def __init__(
             ip=ip,
         )
 
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class FuseOp(FuseOp):
+    """Specialization for FuseOp class."""
+
+    @overload
+    def __init__(
+        self,
+        target: Union[Operation, Value, OpView],
+        *,
+        sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+        interchange: OptionalIntList = None,
+        loc=None,
+        ip=None,
+    ):
+        ...
+
+    def __init__(
+        self,
+        loop_types_or_target: Union[Type, List[Type], Operation, Value],
+        target_or_none: Optional[Union[Operation, Value, OpView]] = None,
+        *,
+        sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+        interchange: OptionalIntList = None,
+        loc=None,
+        ip=None,
+    ):
+        sizes = sizes if sizes else []
+        num_loops = sum(v if v == 0 else 1 for v in sizes)
+
+        if isinstance(loop_types_or_target, (Operation, Value, OpView)):
+            loop_types = [transform.AnyOpType.get()] * num_loops
+            target = loop_types_or_target
+            assert (
+                target_or_none is None
+            ), "Cannot construct FuseOp with two targets."
+        else:
+            loop_types = (
+                ([loop_types_or_target] * num_loops)
+                if isinstance(loop_types_or_target, Type)
+                else loop_types_or_target
+            )
+            target = target_or_none
+        super().__init__(
+            target.type,
+            loop_types,
+            target,
+            tile_sizes=sizes,
+            tile_interchange=interchange,
+            loc=loc,
+            ip=ip,
+        )
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class GeneralizeOp(GeneralizeOp):
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index fb4c75b5337928..69cac29b1fbc91 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -101,6 +101,27 @@ def testFuseIntoContainingOpCompact(target):
     # CHECK-SAME: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
 
 
+ at run
+ at create_sequence
+def testFuseOpCompact(target):
+    structured.FuseOp(target, sizes=[4, 8], interchange=[0, 1])
+    # CHECK-LABEL: TEST: testFuseOpCompact
+    # CHECK: transform.sequence
+    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
+    # CHECK-SAME: interchange [0, 1]
+    # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+
+ at run
+ at create_sequence
+def testFuseOpNoArg(target):
+    structured.FuseOp(target)
+    # CHECK-LABEL: TEST: testFuseOpNoArg
+    # CHECK: transform.sequence
+    # CHECK: %{{.+}} = transform.structured.fuse %{{.*}} :
+    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
+
+
 @run
 @create_sequence
 def testGeneralize(target):

>From a4a76ae4da326299d67997e17e795366b0485b49 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Fri, 20 Dec 2024 00:53:27 +0800
Subject: [PATCH 2/3] fixup! [MLIR][Python] Add structured.fuseop to generator.

---
 mlir/python/mlir/dialects/transform/structured.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index e715872a16ce48..e30500f512cd7b 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -171,9 +171,7 @@ def __init__(
         if isinstance(loop_types_or_target, (Operation, Value, OpView)):
             loop_types = [transform.AnyOpType.get()] * num_loops
             target = loop_types_or_target
-            assert (
-                target_or_none is None
-            ), "Cannot construct FuseOp with two targets."
+            assert target_or_none is None, "Cannot construct FuseOp with two targets."
         else:
             loop_types = (
                 ([loop_types_or_target] * num_loops)

>From b564d573398041f5ea6a81848852605b8617e4d6 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Fri, 20 Dec 2024 01:09:11 +0800
Subject: [PATCH 3/3] fixup! fixup! [MLIR][Python] Add structured.fuseop to
 generator.

---
 mlir/python/mlir/dialects/transform/structured.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index e30500f512cd7b..1d0c71639d07b4 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -139,6 +139,7 @@ def __init__(
             ip=ip,
         )
 
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class FuseOp(FuseOp):
     """Specialization for FuseOp class."""