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

[llvmlistbot at llvm.org]

Author: Jianhui Li
Date: 2026-06-24T17:51:35-07:00
New Revision: 1fadc09c1985a6f3bb0ff993fc876384797116a3

URL: https://github.com/llvm/llvm-project/commit/1fadc09c1985a6f3bb0ff993fc876384797116a3
DIFF: https://github.com/llvm/llvm-project/commit/1fadc09c1985a6f3bb0ff993fc876384797116a3.diff

LOG: [MLIR][XeGPU] Fix order remapping in layout transpose (#205212)

LayoutAttr::transposeDims and LayoutAttr::isTransposeOf mishandled the
`order` field when transposing a layout. The `order` field is
fundamentally different from the size-valued fields (sg_layout, sg_data,
inst_data, lane_layout, lane_data): its values are dimension indices
(order[0] is the fastest-varying dim), not per-position sizes. The two
require different transpose rules:
    - Size fields — reindex by position: new[i] = orig[perm[i]]
- order — relabel values through the inverse permutation: newOrder[i] =
inversePerm[origOrder[i]]

Both functions incorrectly applied the size-field rule to `order`.
Because the bug was applied consistently in both places, it stayed
hidden for trivial/symmetric (e.g. 2D [1,0]) permutations, where the two
rules happen to coincide. It only surfaces for non-trivial permutations
such as the 3D [1,0,2] produced by a broadcast→transpose chain.
   
   Assist-by-Claude

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
    mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 311cf9a64c0c4..5d55342afea15 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -933,8 +933,22 @@ 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]].
+  //
+  // The linearization order this describes is invariant under transpose: a
+  // transpose only renames dimensions, so the subgroup ID assigned to a given
+  // block of data must stay the same. Remapping the values through the inverse
+  // permutation is exactly what preserves that order.
+  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 +982,30 @@ bool LayoutAttr::isTransposeOf(const xegpu::DistributeLayoutAttr &other,
     }
     return true;
   };
+  // `order` is 
diff erent: 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 +1014,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
+  }
+}