[Mlir-commits] [mlir] [MLIR][XeGPU][LayoutPropagation] Only use 'transformed' load when uArch allows it (PR #233485)

[Dmitry Chigarev]

https://github.com/dchigarev updated https://github.com/llvm/llvm-project/pull/233485

>From 7e16bac0144e0ecb57516ae89654e2681f89fbf4 Mon Sep 17 00:00:00 2001
From: Dmitry Chigarev <dmitry.chigarev at intel.com>
Date: Mon, 19 Oct 2026 16:29:33 +0000
Subject: [PATCH] [MLIR][XeGPU][LayoutPropagation] Only use 'transformed' load
 when uArch allows it

Signed-off-by: Dmitry Chigarev <dmitry.chigarev at intel.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      |  8 +++++++-
 .../XeGPU/propagate-layout-inst-data.mlir     | 20 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 649231342e51bd..fc919c49ca820f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1930,10 +1930,16 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
     // spread packed lanes over (width 16 for an 8-bit scale needing 16 lanes of
     // 2 elements) cannot be loaded regularly; transformed, it packs down the
     // column instead. A vertical lane_layout is already transposed.
+    // Only switch if a transformed load is possible (checking via a call
+    // to 'get2DBlockIOInstDataLayout') otherwise keep the regularone and leave
+    // the conflict to a convert_layout.
     if (!hasTranspose && !hasTransform &&
         consumerInstData[rank - 1] %
                 (laneLayout[rank - 1] * laneData[rank - 1]) !=
-            0) {
+            0 &&
+        get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction,
+                                   /*transform=*/true,
+                                   /*transpose=*/false)) {
       hasTransform = true;
       std::tie(laneLayout, laneData, order) = compute2DBlockIOLaneLayout(
           dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(),
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index a0b261f12402d6..3b8fb379e7160e 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -787,3 +787,23 @@ func.func @shape_cast_insert_unit_dim_replicated_lanes() {
   return
 }
 }
+
+// -----
+// A [1, 64] i4 row is too short for a transformed block (64 rows for i4), so
+// the load stays a regular one with its own lane layout ([1, 16] / [1, 4])
+// even though the consumer asks for [1, 16] / [1, 1]; the conflict is left to a
+// convert_layout.
+gpu.module @test {
+  // CHECK-LABEL: load_nd_i4_row_stays_regular
+  // CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [1, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
+  // CHECK-SAME: -> vector<1x64xi4>
+  // CHECK: arith.extui {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+  gpu.func @load_nd_i4_row_stays_regular(%src: memref<32x512xi4>, %dst: memref<32x512xi32>) kernel {
+    %t = xegpu.create_nd_tdesc %src : memref<32x512xi4> -> !xegpu.tensor_desc<1x64xi4>
+    %v = xegpu.load_nd %t[0, 0] : !xegpu.tensor_desc<1x64xi4> -> vector<1x64xi4>
+    %e = arith.extui %v : vector<1x64xi4> to vector<1x64xi32>
+    %o = xegpu.create_nd_tdesc %dst : memref<32x512xi32> -> !xegpu.tensor_desc<1x64xi32>
+    xegpu.store_nd %e, %o[0, 0] <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<1x64xi32>, !xegpu.tensor_desc<1x64xi32>
+    gpu.return
+  }
+}