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

[llvmlistbot at llvm.org]

Author: Andrey Pavlenko
Date: 2026-07-13T15:38:51+02:00
New Revision: f88150e1d5fc8591b0d71619d02e5b553161c1f8

URL: https://github.com/llvm/llvm-project/commit/f88150e1d5fc8591b0d71619d02e5b553161c1f8
DIFF: https://github.com/llvm/llvm-project/commit/f88150e1d5fc8591b0d71619d02e5b553161c1f8.diff

LOG: [MLIR][XeGPU][VectorToXeGPU] Fixed transposed transfer_read/write (#197281)

The problem was that adjustStridesForPermutation() permutes strides, but
they are required to compute base offset from transfer read indices.
Added optional argument to computeMemrefMeta() to preserve the original
strides.

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
    mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 75ffbe58edd6c..9a994e87697f6 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -256,13 +256,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};
 }
 
@@ -312,13 +307,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 db7928724f082..3354868cd6cb9 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -703,3 +703,48 @@ gpu.func @load_0D_vector_unsupported(%source: memref<3xf32>,
 // LOAD-GATHER: 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 {{.*}}
+
+}