[Mlir-commits] [mlir] [MLIR][XeGPU] Move arch-specific topology under uArch check (PR #207449)

[Artem Kroviakov]

https://github.com/akroviakov created https://github.com/llvm/llvm-project/pull/207449

This PR preserves `order` in dpas layout propagation for inst_data and moves the uarch-specific assumptions under the uarch check.

>From 3606148988591ffbe501e41a524c7372710205f8 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Fri, 3 Jul 2026 19:02:09 +0000
Subject: [PATCH] [MLIR][XeGPU] Move arch-specific topology under uArch check

---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 43 +++++++++++--------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 678f585bdcca5..6ea55f4c78dd6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2003,31 +2003,40 @@ xegpu::completeDpasLaneLayoutFromInstData(xegpu::DistributeLayoutAttr aLayout,
   if (!uArchInstruction)
     return std::nullopt;
   auto subgroupSize = uArch->getSubgroupSize();
-
-  auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
-      aTy.getShape(), subgroupSize,
-      aTy.getElementType().getIntOrFloatBitWidth(),
-      uArchInstruction->getPackedFormatBitSizeA());
-  auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
-      bTy.getShape(), subgroupSize,
-      bTy.getElementType().getIntOrFloatBitWidth(),
-      uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
-  auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
-      cdTy.getShape(), subgroupSize,
-      cdTy.getElementType().getIntOrFloatBitWidth(),
-      cdTy.getElementType().getIntOrFloatBitWidth());
+  llvm::SmallVector<int64_t> laneLayoutA, laneDataA, laneLayoutB, laneDataB,
+      laneLayoutCD, laneDataCD;
   SmallVector<int64_t> instDataA = aLayout.getEffectiveInstDataAsInt();
   SmallVector<int64_t> instDataB = bLayout.getEffectiveInstDataAsInt();
   SmallVector<int64_t> instDataCD = cdLayout.getEffectiveInstDataAsInt();
+
+  if (isa<xegpu::uArch::Xe2, xegpu::uArch::Xe3>(uArch)) {
+    std::tie(laneLayoutA, laneDataA) = compute2DBlockIOLaneLayoutAndData(
+        aTy.getShape(), subgroupSize,
+        aTy.getElementType().getIntOrFloatBitWidth(),
+        uArchInstruction->getPackedFormatBitSizeA());
+    std::tie(laneLayoutB, laneDataB) = compute2DBlockIOLaneLayoutAndData(
+        bTy.getShape(), subgroupSize,
+        bTy.getElementType().getIntOrFloatBitWidth(),
+        uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+    std::tie(laneLayoutCD, laneDataCD) = compute2DBlockIOLaneLayoutAndData(
+        cdTy.getShape(), subgroupSize,
+        cdTy.getElementType().getIntOrFloatBitWidth(),
+        cdTy.getElementType().getIntOrFloatBitWidth());
+  } else {
+    assert(false && "Unsupported uArch for DPAS lane layout completion");
+  }
+
   if (!isValidLaneLayout(instDataA, laneLayoutA, laneDataA) ||
       !isValidLaneLayout(instDataB, laneLayoutB, laneDataB) ||
       !isValidLaneLayout(instDataCD, laneLayoutCD, laneDataCD))
     return std::nullopt;
   return std::make_tuple(
-      buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
-      buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
-      buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
-                                  laneDataCD));
+      buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA,
+                                  aLayout.getOrder()),
+      buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB,
+                                  bLayout.getOrder()),
+      buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD, laneDataCD,
+                                  cdLayout.getOrder()));
 }
 
 /// Like completeDpasLaneLayoutFromInstData, but for dpas_mx: also re-derives