[Mlir-commits] [mlir] [mlir][python] Fix mixed tile interchange dispatch (PR #215996)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Shenghang Tsai (jackalcooper)

<details>
<summary>Changes</summary>

Handle absent, mixed, and packed interchange values separately.

Test: `transform_structured_ext.py`


---
Full diff: https://github.com/llvm/llvm-project/pull/215996.diff


2 Files Affected:

- (modified) mlir/python/mlir/dialects/transform/structured.py (+17-5) 
- (modified) mlir/test/python/dialects/transform_structured_ext.py (+11) 


``````````diff
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 67eca6886c5dc..43ebdef8384e4 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -610,11 +610,23 @@ def __init__(
             if packed_tile_sizes is not None
             else sum(v if v == 0 else 1 for v in static_sizes)
         )
-        (
-            dynamic_interchange,
-            packed_interchange,
-            static_interchange,
-        ) = _dispatch_mixed_values(interchange)
+        if interchange is None:
+            dynamic_interchange = []
+            packed_interchange = None
+            static_interchange = []
+        elif isinstance(interchange, (Operation, Value, OpView)):
+            (
+                dynamic_interchange,
+                packed_interchange,
+                static_interchange,
+            ) = _dispatch_mixed_values(interchange)
+        else:
+            (
+                dynamic_interchange,
+                static_interchange,
+                _,
+            ) = _dispatch_dynamic_index_list(interchange)
+            packed_interchange = None
 
         if isinstance(loop_types_or_target, (Operation, Value, OpView)):
             loop_types = [transform.AnyOpType.get()] * num_loops
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index dfe798932a238..678d39d78b512 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -523,6 +523,17 @@ def testTileInterchangeMixed(target):
     # CHECK-SAME: [4, 8] interchange = [%[[I0]], 1]
 
 
+ at run
+ at create_sequence
+def testTileInterchangePacked(target):
+    interchange = structured.MatchOp.match_op_names(target, ["arith.constant"])
+    structured.TileUsingForOp(target, sizes=[4, 8], interchange=interchange)
+    # CHECK-LABEL: TEST: testTileInterchangePacked
+    # CHECK: %[[I:.*]] = transform.structured.match
+    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for
+    # CHECK-SAME: [4, 8] interchange = *(%[[I]])
+
+
 @run
 @create_sequence
 def testTileZero(target):

``````````

</details>


https://github.com/llvm/llvm-project/pull/215996