[Mlir-commits] [mlir] [MLIR][XeGPU] Refactoring Transpose OP Layout Propagation (PR #184702)

[Jianhui Li]

================
@@ -636,6 +636,96 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
   return collapsedLayout;
 }
 
+// Derive a new layout by transpose the layout using `permutation`.
+DistributeLayoutAttr LayoutAttr::transposeDims(ArrayRef<int64_t> permutation) {
+
+  SmallVector<int64_t> origSgLayout = getEffectiveSgLayoutAsInt();
+  SmallVector<int64_t> origSgData = getEffectiveSgDataAsInt();
+  SmallVector<int64_t> origInstData = getEffectiveInstDataAsInt();
+  SmallVector<int64_t> origLaneLayout = getEffectiveLaneLayoutAsInt();
+  SmallVector<int64_t> origLaneData = getEffectiveLaneDataAsInt();
+  SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
+
+  SmallVector<int32_t> sgLayout;
+  SmallVector<int32_t> sgData;
+  SmallVector<int32_t> instData;
+  SmallVector<int32_t> laneLayout;
+  SmallVector<int32_t> laneData;
+  SmallVector<int32_t> order;
+  xegpu::LayoutAttr layoutAttr;
+
+  for (int64_t idx : permutation) {
+    if (!origLaneLayout.empty()) {
+      laneLayout.push_back(static_cast<int32_t>(origLaneLayout[idx]));
+      laneData.push_back(static_cast<int32_t>(origLaneData[idx]));
+    }
+    if (!origInstData.empty())
+      instData.push_back(static_cast<int32_t>(origInstData[idx]));
+    if (!origSgData.empty()) {
+      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]));
+  }
+
+  if (!origLaneLayout.empty()) {
+    layoutAttr = xegpu::LayoutAttr::get(
+        getContext(),
+        /*sg_lane =*/nullptr, /*sg_data =*/nullptr, /*inst_data =*/nullptr,
+        DenseI32ArrayAttr::get(getContext(), laneLayout),
+        DenseI32ArrayAttr::get(getContext(), laneData),
+        DenseI32ArrayAttr::get(getContext(), order));
+  }
+  if (!origInstData.empty()) {
+    layoutAttr = xegpu::LayoutAttr::get(getContext(), instData);
+  }
+  if (!origSgData.empty()) {
+    layoutAttr = xegpu::LayoutAttr::get(
----------------
Jianhui-Li wrote:

good catch! I am fixing it. 

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