fixup! fixup! fixup! [MLIR][Python] Add structured.fuseop to generator.

Author: nujaa

mlir/python/mlir/dialects/transform/structured.py

@@ -149,25 +149,28 @@ def __init__(
         self,
         target: Union[Operation, Value, OpView],
         *,
-        sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
-        interchange: OptionalIntList = None,
+        tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+        tile_interchange: OptionalIntList = None,
         loc=None,
         ip=None,
     ):
         ...
 
     def __init__(
         self,
-        loop_types_or_target: Union[Type, List[Type], Operation, Value],
+        loop_types_or_target: Union[Type, Sequence[Type], Operation, OpView, Value],
         target_or_none: Optional[Union[Operation, Value, OpView]] = None,
         *,
-        sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
-        interchange: OptionalIntList = None,
+        tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+        tile_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)
+        tile_sizes = tile_sizes if tile_sizes else []
+        tile_interchange = tile_interchange if tile_interchange else []
+        _, tile_sizes, _ = _dispatch_dynamic_index_list(tile_sizes)
+        _, tile_interchange, _ = _dispatch_dynamic_index_list(tile_interchange)
+        num_loops = sum(0 if v == 0 else 1 for v in tile_sizes)
 
         if isinstance(loop_types_or_target, (Operation, Value, OpView)):
             loop_types = [transform.AnyOpType.get()] * num_loops
@@ -184,8 +187,8 @@ def __init__(
             target.type,
             loop_types,
             target,
-            tile_sizes=sizes,
-            tile_interchange=interchange,
+            tile_sizes=tile_sizes,
+            tile_interchange=tile_interchange,
             loc=loc,
             ip=ip,
         )

mlir/test/python/dialects/transform_structured_ext.py

@@ -104,7 +104,7 @@ def testFuseIntoContainingOpCompact(target):
 @run
 @create_sequence
 def testFuseOpCompact(target):
-    structured.FuseOp(target, sizes=[4, 8], interchange=[0, 1])
+    structured.FuseOp(target, tile_sizes=[4, 8], tile_interchange=[0, 1])
     # CHECK-LABEL: TEST: testFuseOpCompact
     # CHECK: transform.sequence
     # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
@@ -122,6 +122,19 @@ def testFuseOpNoArg(target):
     # CHECK-SAME: (!transform.any_op) -> !transform.any_op
 
 
+@run
+@create_sequence
+def testFuseOpAttributes(target):
+    attr = DenseI64ArrayAttr.get([4, 8])
+    ichange = DenseI64ArrayAttr.get([0, 1])
+    structured.FuseOp(target, tile_sizes=attr, tile_interchange=ichange)
+    # CHECK-LABEL: TEST: testFuseOpAttributes
+    # 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)
+
+
 @run
 @create_sequence
 def testGeneralize(target):