[Mlir-commits] [mlir] [MLIR][XeGPU][VectorToXeGPU] Fixed transposed transfer_read/write (PR #197281)

[Andrey Pavlenko]

https://github.com/AndreyPavlenko updated https://github.com/llvm/llvm-project/pull/197281

>From 63501e958730dc81c5d28c5404838954500a8eb9 Mon Sep 17 00:00:00 2001
From: Andrey Pavlenko <andrey.a.pavlenko at gmail.com>
Date: Tue, 12 May 2026 19:35:41 +0000
Subject: [PATCH] Fixed transposed transfer_read/write

The problem was that adjustStridesForPermutation() permutes strides,
but they are required to compute base offset from transfer read indices.
---
 .../VectorToXeGPU/VectorToXeGPU.cpp           | 18 ++++----
 .../VectorToXeGPU/transfer-read-to-xegpu.mlir | 45 +++++++++++++++++++
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 8dcc9771d8db4..ba5f010cdd862 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -254,13 +254,8 @@ computeMemrefMeta(OpType xferOp, PatternRewriter &rewriter) {
       offsetVal = meta.getOffset();
   }
 
-  if constexpr (llvm::is_one_of<std::decay_t<OpType>, vector::TransferReadOp,
-                                vector::TransferWriteOp>::value) {
-    AffineMap permMap = xferOp.getPermutationMap();
-    // Adjust strides according to the permutation map (e.g., for transpose)
-    adjustStridesForPermutation(permMap, strides);
-  }
-
+  // Strides are returned in original memref order; permutation is applied in
+  // computeOffsets only where offsets are indexed in vector order.
   return {strides, offsetVal};
 }
 
@@ -310,13 +305,18 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
     return stepOp;
   });
 
+  // Local offsets are indexed in vector order, so permute strides; the base
+  // offset below uses the original memref-order strides.
+  SmallVector<Value> permutedStrides(strides.begin(), strides.end());
+  adjustStridesForPermutation(xferOp.getPermutationMap(), permutedStrides);
+
   // Multiply step vectors by corresponding strides
-  size_t memrefRank = strides.size();
+  size_t memrefRank = permutedStrides.size();
   size_t vectorRank = vectorShape.size();
   SmallVector<Value> strideMultiplied;
   for (size_t i = 0; i < vectorRank; ++i) {
     size_t memrefDim = memrefRank - vectorRank + i;
-    Value strideValue = strides[memrefDim];
+    Value strideValue = permutedStrides[memrefDim];
     auto mulType = dyn_cast<VectorType>(stepVectors[i].getType());
     auto bcastOp =
         vector::BroadcastOp::create(rewriter, loc, mulType, strideValue);
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index f64bcc27940ea..b319efa97554f 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -679,4 +679,49 @@ gpu.func @load_0D_memref_unsupported(%source: memref<f16>) -> vector<f16> {
 // CHECK-LABEL: @load_0D_memref_unsupported
 // CHECK: vector.transfer_read
 
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @transpose_1x1024x24x64(
+    %arg0: memref<1x1024x24x64xf16>,
+    %arg1: memref<1x24x1024x64xf16>) kernel
+    attributes {known_block_size = array<i32: 256, 1, 1>} {
+  %pad = ub.poison : f16
+  %c0 = arith.constant 0 : index
+  %block_id_x = gpu.block_id x
+  %block_id_y = gpu.block_id y
+  %block_id_z = gpu.block_id z
+  %seq_off = affine.apply affine_map<()[s0] -> (s0 * 16)>()[%block_id_y]
+  %hid_off = affine.apply affine_map<()[s0] -> (s0 * 8)>()[%block_id_z]
+  %vec = vector.transfer_read %arg0[%c0, %seq_off, %block_id_x, %hid_off], %pad
+    {in_bounds = [true, true, true, true]}
+    : memref<1x1024x24x64xf16>, vector<1x16x1x8xf16>
+  %transposed = vector.transpose %vec, [0, 2, 1, 3]
+    : vector<1x16x1x8xf16> to vector<1x1x16x8xf16>
+  vector.transfer_write %transposed, %arg1[%c0, %block_id_x, %seq_off, %hid_off]
+    {in_bounds = [true, true, true, true]}
+    : vector<1x1x16x8xf16>, memref<1x24x1024x64xf16>
+  gpu.return
+}
+
+// CHECK-LABEL: @transpose_1x1024x24x64
+// CHECK-DAG: %[[C1536:.+]] = arith.constant 1536 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C65536:.+]] = arith.constant 65536 : index
+
+// Read from memref<1x1024x24x64xf16>, strides [1572864, 1536, 64, 1].
+// Scalar base offset: seq_off * 1536 (original dim1 stride),
+//                    block_id_x * 64  (original dim2 stride).
+// CHECK:     arith.muli %{{.+}}, %[[C1536]] : index
+// CHECK:     arith.muli %block_id_x, %[[C64]] : index
+// CHECK:     xegpu.load {{.*}} -> vector<1x1x16x8xf16>
+
+// Write to memref<1x24x1024x64xf16>, strides [1572864, 65536, 64, 1].
+// Scalar base offset: block_id_x * 65536 (original dim1 stride),
+//                    seq_off * 64        (original dim2 stride).
+// CHECK:     arith.muli %block_id_x, %[[C65536]] : index
+// CHECK:     arith.muli %{{.+}}, %[[C64]] : index
+// CHECK:     xegpu.store {{.*}}
+
 }
\ No newline at end of file