[Mlir-commits] [mlir] [MLIR][XeGPU] Fix order remapping in layout transpose (PR #205212)

[Jianhui Li]

https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/205212

>From 46ea8667e7835dd9eca378a9e3a036d2dc06967d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 22 Jun 2026 22:58:04 +0000
Subject: [PATCH] [MLIR][XeGPU] Fix order remapping in layout transpose

LayoutAttr::transposeDims and LayoutAttr::isTransposeOf treated the
`order` field like the size-valued fields (sg_layout, sg_data, ...),
reindexing it by position (new[i] = orig[perm[i]]). But `order`'s
*values* are dimension indices, so a transpose relabels those values
through the inverse permutation while preserving the sequence:
newOrder[i] = inversePerm[origOrder[i]].

The position-reindex bug was self-consistent across both functions, so
it only surfaced for non-trivial (e.g. 3D) permutations: a broadcast +
transpose chain recovered the default order instead of the correct one,
which then made convert_layout fall onto the SLM redistribution path
unnecessarily.

Fix both functions to remap `order` values through the inverse
permutation, and add a 3D transpose propagation test that guards the
recovered source order.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
 mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp    | 36 ++++++++++++++++---
 .../XeGPU/propagate-layout-subgroup.mlir      | 19 ++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 4aa1c0d666a94..7e2117b5fc7a3 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -933,8 +933,17 @@ DistributeLayoutAttr LayoutAttr::transposeDims(ArrayRef<int64_t> permutation) {
       sgLayout.push_back(static_cast<int32_t>(origSgLayout[idx]));
       sgData.push_back(static_cast<int32_t>(origSgData[idx]));
     }
-    order.push_back(static_cast<int32_t>(origOrder[idx]));
   }
+
+  // `order` is distinct from the size-valued fields above: its *values* are
+  // dimension indices (order[0] is the fastest-varying dim), not per-position
+  // sizes. A transpose relabels dimensions (source dim d becomes result dim
+  // inversePerm[d]) so the dimension values are remapped through the inverse
+  // permutation: newOrder[i] = inversePerm[origOrder[i]].
+  SmallVector<int64_t> inversePermutation =
+      invertPermutationVector(permutation);
+  for (int64_t dim : origOrder)
+    order.push_back(static_cast<int32_t>(inversePermutation[dim]));
   if (origLaneLayout.empty() && origSgLayout.empty())
     order.clear();
 
@@ -968,13 +977,30 @@ bool LayoutAttr::isTransposeOf(const xegpu::DistributeLayoutAttr &other,
     }
     return true;
   };
+  // `order` is different: its *values* are dimension indices, so a transpose
+  // relabels them through the inverse permutation rather than reindexing by
+  // position. `this` (= dst) is a transpose of `other` (= src) iff
+  // dst.order[i] == inversePerm[src.order[i]] for all i. This matches the
+  // convention produced by `transposeDims`.
+  auto checkOrderTranspose = [](ArrayRef<int64_t> dstOrder,
+                                ArrayRef<int64_t> srcOrder,
+                                ArrayRef<int64_t> perm) {
+    if (dstOrder.size() != srcOrder.size())
+      return false;
+    SmallVector<int64_t> inversePerm = invertPermutationVector(perm);
+    for (auto [d, s] : llvm::zip_equal(dstOrder, srcOrder)) {
+      if (d != inversePerm[s])
+        return false;
+    }
+    return true;
+  };
   if (kind == xegpu::LayoutKind::Subgroup)
     return checkTranspose(getEffectiveSgLayoutAsInt(),
                           other.getEffectiveSgLayoutAsInt(), perm) &&
            checkTranspose(getEffectiveSgDataAsInt(),
                           other.getEffectiveSgDataAsInt(), perm) &&
-           checkTranspose(getEffectiveOrderAsInt(),
-                          other.getEffectiveOrderAsInt(), perm);
+           checkOrderTranspose(getEffectiveOrderAsInt(),
+                               other.getEffectiveOrderAsInt(), perm);
   if (kind == xegpu::LayoutKind::InstData)
     return checkTranspose(getEffectiveInstDataAsInt(),
                           other.getEffectiveInstDataAsInt(), perm);
@@ -983,8 +1009,8 @@ bool LayoutAttr::isTransposeOf(const xegpu::DistributeLayoutAttr &other,
                           other.getEffectiveLaneLayoutAsInt(), perm) &&
            checkTranspose(getEffectiveLaneDataAsInt(),
                           other.getEffectiveLaneDataAsInt(), perm) &&
-           checkTranspose(getEffectiveOrderAsInt(),
-                          other.getEffectiveOrderAsInt(), perm);
+           checkOrderTranspose(getEffectiveOrderAsInt(),
+                               other.getEffectiveOrderAsInt(), perm);
 
   return false;
 }
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index d44497d0bba34..4c6353e45cfbe 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -589,3 +589,22 @@ gpu.module @test {
     gpu.return
   }
 }
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: gpu.func @transpose_3d_order_remap(
+// CHECK: %[[TD_LD:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<32x2x32xf16> ->
+// CHECK-SAME: !xegpu.tensor_desc<32x2x32xf16, #xegpu.layout<sg_layout = [1, 2, 2], sg_data = [32, 1, 16], order = [0, 2, 1]>>
+// CHECK: %[[LD:.*]] = xegpu.load_nd %[[TD_LD]][%{{.*}}] <{layout = #xegpu.layout<sg_layout = [1, 2, 2], sg_data = [32, 1, 16], order = [0, 2, 1]>}>
+// CHECK: %[[TR:.*]] = vector.transpose %[[LD]], [1, 0, 2] {layout_result_0 = #xegpu.layout<sg_layout = [2, 1, 2], sg_data = [1, 32, 16], order = [1, 2, 0]>}
+// CHECK-SAME: : vector<32x2x32xf16> to vector<2x32x32xf16>
+  gpu.func @transpose_3d_order_remap(%src: memref<32x2x32xf16>, %dst: memref<2x32x32xf16>) kernel {
+    %c0 = arith.constant 0 : index
+    %td_in = xegpu.create_nd_tdesc %src : memref<32x2x32xf16> -> !xegpu.tensor_desc<32x2x32xf16>
+    %ld = xegpu.load_nd %td_in[%c0, %c0, %c0] : !xegpu.tensor_desc<32x2x32xf16> -> vector<32x2x32xf16>
+    %tr = vector.transpose %ld, [1, 0, 2] : vector<32x2x32xf16> to vector<2x32x32xf16>
+    %td_out = xegpu.create_nd_tdesc %dst : memref<2x32x32xf16> -> !xegpu.tensor_desc<2x32x32xf16>
+    xegpu.store_nd %tr, %td_out[%c0, %c0, %c0] <{layout = #xegpu.layout<sg_layout = [2, 1, 2], sg_data = [1, 32, 16], order = [1, 2, 0]>}> : vector<2x32x32xf16>, !xegpu.tensor_desc<2x32x32xf16>
+    gpu.return
+  }
+}