[Mlir-commits] [mlir] [MLIR][XeGPU] Decompose unsupported 'vector.transfer_read'-transpose-permutations (PR #182875)

[Adam Siemieniuk]

================
@@ -841,6 +862,54 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
   }
 };
 
+// Splits 'vector.transfer_read' with unsupported transpose-permutations
+// into 'transfer_read() + transpose()'.
+struct TransferReadDecomposeUnsupportedTranspose
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
+                                PatternRewriter &rewriter) const override {
+    Location loc = readOp.getLoc();
+    AffineMap readMap = readOp.getPermutationMap();
+
+    bool isTransposeLoad = !readMap.isMinorIdentity();
+    if (!isTransposeLoad)
+      return failure();
+
+    bool isTransposeSupported = isTransferReadTransposeSupported(readOp);
+    if (isTransposeSupported)
+      return failure();
+
+    auto resultType = cast<VectorType>(readOp.getResult().getType());
+    if (!resultType)
+      return failure();
+
+    // 'Revert' permutation for the transfer_read result shape to make it
+    // 'untransposed'
+    auto newShape = applyPermutationMap(readMap, resultType.getShape());
+    auto newTransferReadRes =
+        VectorType::get(newShape, resultType.getElementType());
+
+    // Step 1. Create 'plain' transfer_read without transpose
+    auto newReadOp = vector::TransferReadOp::create(
+        rewriter, loc, newTransferReadRes, readOp.getBase(),
+        readOp.getIndices(), AffineMap::get(loc.getContext()),
+        readOp.getPadding(), readOp.getMask(), readOp.getInBoundsAttr());
+
+    // Step 2. Transpose the result of the 'plain' transfer_read
+    auto range = llvm::seq<int64_t>(0, readMap.getResults().size());
----------------
adam-smnk wrote:

+1 for merging this logic with transfer op lowering.

Looking from a perspective of a larger pipeline, you probably want to greedily fold all transposes before going to XeGPU to expose opportunities for hw-supported fused transposes.
In this case, a separate pattern brings extra burden to maintain its precondition whether given transfer should be modified or not due to possibly random pattern ordering when fused transpose is already supported.

You could also consider if existing `vector-to-xegpu` patterns could be combined with `vector::populateVectorTransferPermutationMapLoweringPatterns` that also can decompose transposes.

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