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

[Shenghang Tsai]

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

>From ba60e9f324fad026b13fc233268e84739ac54455 Mon Sep 17 00:00:00 2001
From: Shenghang Tsai <jackalcooper at gmail.com>
Date: Thu, 13 Aug 2026 17:48:05 +0800
Subject: [PATCH] [mlir][python] Fix tile interchange ArrayAttr dispatch

---
 .../mlir/dialects/transform/structured.py     | 27 +++++++++++++++----
 .../dialects/transform_structured_ext.py      | 21 +++++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 67eca6886c5dc..4c7d76e70c5d9 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -610,11 +610,28 @@ 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)
+        dynamic_interchange = []
+        packed_interchange = None
+        static_interchange = []
+        if isinstance(interchange, (Operation, Value, OpView)):
+            (
+                dynamic_interchange,
+                packed_interchange,
+                static_interchange,
+            ) = _dispatch_mixed_values(interchange)
+        elif (
+            isinstance(interchange, Sequence)
+            and not isinstance(interchange, (str, bytes))
+        ) or isinstance(interchange, (ArrayAttr, DenseI64ArrayAttr)):
+            (
+                dynamic_interchange,
+                static_interchange,
+                _,
+            ) = _dispatch_dynamic_index_list(interchange)
+        elif interchange is not None:
+            raise ValueError(
+                f"expected {interchange=} to be an MLIR object or sequence of mixed values"
+            )
 
         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..584ec79179cd7 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -523,6 +523,27 @@ def testTileInterchangeMixed(target):
     # CHECK-SAME: [4, 8] interchange = [%[[I0]], 1]
 
 
+ at run
+ at create_sequence
+def testTileInterchangeArrayAttr(target):
+    interchange = ArrayAttr.get(
+        [IntegerAttr.get(IndexType.get(), 0), IntegerAttr.get(IndexType.get(), 1)]
+    )
+    structured.TileUsingForOp(target, sizes=[4, 8], interchange=interchange)
+    # CHECK-LABEL: TEST: testTileInterchangeArrayAttr
+    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for
+    # CHECK-SAME: [4, 8] interchange = [0, 1]
+
+
+ at run
+ at create_sequence
+def testTileInterchangeTuple(target):
+    structured.TileUsingForOp(target, sizes=[4, 8], interchange=(0, 1))
+    # CHECK-LABEL: TEST: testTileInterchangeTuple
+    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for
+    # CHECK-SAME: [4, 8] interchange = [0, 1]
+
+
 @run
 @create_sequence
 def testTileZero(target):