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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Dmitry Chigarev (dchigarev)

<details>
<summary>Changes</summary>

The PR fixes behavior where 'transformed' load could have been picked for `dataShape` that doesn't support it.

Before performing 'mxfp-fallback' to the 'transform-load' we now query `uArch` for [supported `dataShape`s](https://github.com/llvm/llvm-project/blob/cbbb5acd6db3c9ccf3bae60b7ffbe56ab213f110/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h#L337-L356) via `get2DBlockIOInstDataLayout`. If transform-load is not possible we proceed with the regular load and let the `xegpu.convert_layout` handle the conflict.

---
Full diff: https://github.com/llvm/llvm-project/pull/233485.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+7-1) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+20) 


``````````diff
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
+  }
+}

``````````

</details>


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