[Mlir-commits] [mlir] [MLIR][XeGPUToXeVM] Add vector-shuffle for interleaved row-values (PR #233509)

[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 covers the transposed-block-load case with `32x8xi32, inst_data=[32, 8]` shape ([uArch link](https://github.com/llvm/llvm-project/blob/980be5863a8419ae9aa146229e804b99e4e4f79c/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h#L351)). 

With this instruction each lane gets 2 rows of data (`lane_i` and `lane_i+15` (8 + 8 = 16 elements)). The result is a 1D vector of 16 elements:
```mlir
%res = xevm.blockload2d ... <{
    cache_control = #xevm.load_cache_control<Use_Default>,
    elem_size_in_bits = 32 : i32,
    pack_register = false,
    tile_height = 32 : i32, tile_width = 8 : i32,
    transpose = true, v_blocks = 1 : i32
}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<16xi32>
```

Before this PR `xegpu-to-xevm` suggested that the hardware places rows data contiguous in the result vector, e.g. first all the elements of row lane_i and then all elements of lane_i+15:
```
[r0_c0, r0_c1, r0_c2, ... r0_c7, r15_c0, r15_c1, ... r15_c7]

// r5_c3 means: (row-5, column-3) of the loaded tile
```
however in reality row-elements are interleaved:
```
[r0_c0, r15_c0, r0_c1, r15_c1, ... r0_c7, r15_c7]
```

The PR fixes this by adding a shuffle to reorder rows-data to rely contiguous in the vector.

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


2 Files Affected:

- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+13) 
- (modified) mlir/test/Conversion/XeGPUToXeVM/loadstore_nd_transpose.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index c0b3adf8402e03..4a4cd45b37141f 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -733,6 +733,19 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
               baseShapeH, basePitchBytes, offsetW, offsetH, elemBitSize, tileW,
               tileH, vblocks, transpose, vnni,
               xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
+          // A transposed block taller than the subgroup gives each lane
+          // several source rows, which the HW interleaves column by column.
+          // Reorder them row by row, as the lane-level value expects.
+          int64_t numLoaded = loadedTy.getNumElements();
+          if (transpose && numLoaded > tileW && numLoaded % tileW == 0) {
+            int64_t rowsPerLane = numLoaded / tileW;
+            SmallVector<int64_t> rowMajor;
+            for (int64_t r = 0; r < rowsPerLane; ++r)
+              for (int64_t c = 0; c < tileW; ++c)
+                rowMajor.push_back(c * rowsPerLane + r);
+            resultFlatVec = vector::ShuffleOp::create(
+                rewriter, loc, resultFlatVec, resultFlatVec, rowMajor);
+          }
           resultFlatVec = vector::BitCastOp::create(
               rewriter, loc,
               encodeVectorTypeTo(loadedTy, dstVecTy.getElementType()),
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd_transpose.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd_transpose.mlir
index d77a72119613e9..761c4736449f67 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd_transpose.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd_transpose.mlir
@@ -43,4 +43,21 @@ gpu.module @load_check {
         vector.store %loaded, %dstte[%c0, %c0] : memref<16x16xf16>, vector<16xf16>
         gpu.return
     }
+
+    // Each lane gets two source rows interleaved by the HW; they are reordered
+    // row by row.
+    // CHECK-LABEL: gpu.func @load_nd_transpose_32_rows(
+    gpu.func @load_nd_transpose_32_rows(%src: memref<32x8xi32, 1>, %dst: memref<16xi32, 1>) kernel {
+        %srcce = memref.memory_space_cast %src : memref<32x8xi32, 1> to memref<32x8xi32>
+        %src_tdesc = xegpu.create_nd_tdesc %srcce : memref<32x8xi32> -> !xegpu.tensor_desc<32x8xi32>
+        // CHECK: %[[LD:.*]] = xevm.blockload2d
+        // CHECK-SAME: tile_height = 32 : i32, tile_width = 8 : i32, transpose = true
+        // CHECK-SAME: -> vector<16xi32>
+        // CHECK: vector.shuffle %[[LD]], %[[LD]] [0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15] : vector<16xi32>, vector<16xi32>
+        %loaded = xegpu.load_nd %src_tdesc[0, 0] <{transpose = array<i64: 1, 0>}> : !xegpu.tensor_desc<32x8xi32> -> vector<16xi32>
+        %dstte = memref.memory_space_cast %dst : memref<16xi32, 1> to memref<16xi32>
+        %c0 = arith.constant 0 : index
+        vector.store %loaded, %dstte[%c0] : memref<16xi32>, vector<16xi32>
+        gpu.return
+    }
 }

``````````

</details>


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