[mlir] Support dimension permutations in ProgressiveVectorToSCF

[springerm]

Summary

This commit adds support for dimension permutations in permutation maps of vector transfer ops.

Depends On D100946

[springerm]

(no change)

[springerm]

(no change)

[springerm]

(no change)

[springerm]

rebase

[springerm]

(no change)

[springerm]

(no change)

[nicolasvasilache]

Looks good, note that there is also a pass that rewrites the transfer into memref.transpose and identity permutation map.

Would be good to start thinking (in the background) about a mechanism where we can more easily expose the decision of either:

1. what you do here.
2. what I mention above.
3. lowering to vector.transpose

We will be interested in the interaction of the above with vector.shape_cast and vector.reshape as well as masking.

[springerm]

(rebase)

[springerm]

(new export)

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

static Value setAllocAtFunctionEntry(MemRefType type, Operation *op) {
   b.setInsertionPointToStart(&scope->getRegion(0).front());
   Value res = memref_alloca(type);
   return res;
 }
 
 /// Given a vector transfer op, calculate which dimension of the `source`
 /// memref should be unpacked in the next application of TransferOpConversion.
 template <typename OpTy>
-static int64_t unpackedDim(OpTy xferOp) {
-  return xferOp.getShapedType().getRank() - xferOp.getVectorType().getRank();
+static unsigned unpackedDim(OpTy xferOp) {
+  auto map = xferOp.permutation_map();
+  // TODO: Handle broadcast
+  auto expr = map.getResult(0).template dyn_cast<AffineDimExpr>();
+  assert(expr && "Expected AffineDimExpr in permutation map result");
+  return expr.getPosition();
+}
+
+/// Compute the permutation map for the new (N-1)-D vector transfer op. This
+/// map is identical to the current permutation map, but the first result is
+/// omitted.
+template <typename OpTy>
+static AffineMap unpackedPermutationMap(OpTy xferOp, OpBuilder &builder) {
+  auto map = xferOp.permutation_map();
+  return AffineMap::get(map.getNumDims(), 0, map.getResults().drop_front(),
+                        builder.getContext());
 }
 
 /// Calculate the indices for the new vector transfer op.
 ///
 /// E.g.: transfer_read %A[%a, %b, %c, %d] ... : vector<5x4x3xf32> ...
 ///       --> transfer_read %A[%a, %b + iv, %c, %d] ... vector<4x3f32>
 ///                                 ^^^^^^
 ///              `iv` is the iteration variable of the (new) surrounding loop.
 template <typename OpTy>
 static void getXferIndices(OpTy xferOp, Value iv,
                            SmallVector<Value, 8> &indices) {
   typename OpTy::Adaptor adaptor(xferOp);
   // Corresponding memref dim of the vector dim that is unpacked.
   auto dim = unpackedDim(xferOp);
   auto prevIndices = adaptor.indices();
   indices.append(prevIndices.begin(), prevIndices.end());
   using edsc::op::operator+;
   indices[dim] = adaptor.indices()[dim] + iv;
 }
 
-static void maybeYieldValue(bool hasRetVal, OpBuilder builder, Location loc,
-                            Value value) {
+static void maybeYieldValue(
+    bool hasRetVal, OpBuilder builder, Location loc, Value value) {
   if (hasRetVal) {
     builder.create<scf::YieldOp>(loc, value);
   } else {
     builder.create<scf::YieldOp>(loc);
   }
 }
 
 /// Helper function TransferOpConversion and Strided1dTransferOpConversion.

 ///   (out-of-bounds case)
 /// }
 /// ```
 /// This function variant returns the value returned by `inBoundsCase` or
 /// `outOfBoundsCase`. The MLIR type of the return value must be specified in
 /// `resultTypes`.
 template <typename OpTy>
 static Value generateInBoundsCheck(
-    OpTy xferOp, Value iv, OpBuilder &builder, int64_t dim,
+    OpTy xferOp, Value iv, OpBuilder &builder, unsigned dim,
     TypeRange resultTypes,
     function_ref<Value(OpBuilder &, Location)> inBoundsCase,
     function_ref<Value(OpBuilder &, Location)> outOfBoundsCase = nullptr) {
   bool hasRetVal = !resultTypes.empty();
   if (!xferOp.isDimInBounds(0)) {
     auto memrefDim = memref_dim(xferOp.source(), std_constant_index(dim));
     using edsc::op::operator+;
     auto memrefIdx = xferOp.indices()[dim] + iv;
     auto cond = std_cmpi_sgt(memrefDim.value, memrefIdx);
     auto check = builder.create<scf::IfOp>(
         xferOp.getLoc(), resultTypes, cond,
-        /*thenBuilder=*/
-        [&](OpBuilder &builder, Location loc) {
-          maybeYieldValue(hasRetVal, builder, loc, inBoundsCase(builder, loc));
-        },
-        /*elseBuilder=*/
-        [&](OpBuilder &builder, Location loc) {
-          if (outOfBoundsCase) {
-            maybeYieldValue(hasRetVal, builder, loc,
-                            outOfBoundsCase(builder, loc));
-          } else {
-            builder.create<scf::YieldOp>(loc);
-          }
-        });
+        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
+    }, /*elseBuilder=*/[&](OpBuilder &builder, Location loc) {
+        maybeYieldValue(hasRetVal, builder, loc, outOfBoundsCase(builder, loc));
 
     return hasRetVal ? check.getResult(0) : Value();
   }
 
   // No runtime check needed if dim is guaranteed to be in-bounds.
   return inBoundsCase(builder, xferOp.getLoc());
 }
 
 /// In this function variant, `inBoundsCase` and `outOfBoundsCase` do not have
 /// a return value. Consequently, this function does not have a return value.
 template <typename OpTy>
 static void generateInBoundsCheck(
     OpTy xferOp, Value iv, OpBuilder &builder, int64_t dim,
     function_ref<void(OpBuilder &, Location)> inBoundsCase,
     function_ref<void(OpBuilder &, Location)> outOfBoundsCase = nullptr) {
   generateInBoundsCheck(
       xferOp, iv, builder, dim, /*resultTypes=*/TypeRange(),
-      /*inBoundsCase=*/
-      [&](OpBuilder &builder, Location loc) {
+      /*inBoundsCase=*/[&](OpBuilder &builder, Location loc) {
         inBoundsCase(builder, loc);
         return Value();
       },
-      /*outOfBoundsCase=*/
-      [&](OpBuilder &builder, Location loc) {
+      /*outOfBoundsCase=*/[&](OpBuilder &builder, Location loc) {
         if (outOfBoundsCase)
             outOfBoundsCase(builder, loc);
         return Value();
       });
 }
 
 /// Given an ArrayAttr, return a copy where the first element is dropped.
 static ArrayAttr dropFirstElem(PatternRewriter &rewriter, ArrayAttr attr) {
   if (!attr)
       return attr;
   return ArrayAttr::get(rewriter.getContext(), attr.getValue().drop_front());
 }
 
 /// Codegen strategy, depending on the operation.
 template <typename OpTy>
 struct Strategy;
 
 /// Code strategy for vector TransferReadOp.
-template <>
+template<>
 struct Strategy<TransferReadOp> {
   /// Find the StoreOp that is used for writing the current TransferReadOp's
   /// result to the temporary buffer allocation.
   static memref::StoreOp getStoreOp(TransferReadOp xferOp) {
     assert(xferOp->hasOneUse() && "Expected exactly one use of TransferReadOp");
-    auto storeOp = dyn_cast<memref::StoreOp>((*xferOp->use_begin()).getOwner());
+    auto storeOp = dyn_cast<memref::StoreOp>(
+        (*xferOp->use_begin()).getOwner());
     assert(storeOp && "Expected TransferReadOp result used by StoreOp");
     return storeOp;
   }
 
   /// Find the temporary buffer allocation. All labeled TransferReadOps are
   /// used like this, where %buf is either the buffer allocation or a type cast
   /// of the buffer allocation:
   /// ```
   /// %vec = vector.transfer_read ... { __vector_to_scf_lowering__ } ...
   /// memref.store %vec, %buf[...] ...
   /// ```
   static Value getBuffer(TransferReadOp xferOp) {
     return getStoreOp(xferOp).getMemRef();
   }
 
   /// Retrieve the indices of the current StoreOp.
   static void getStoreIndices(TransferReadOp xferOp,
-                              SmallVector<Value, 8> &indices) {
     auto storeOp = getStoreOp(xferOp);
     auto prevIndices = memref::StoreOpAdaptor(storeOp).indices();
     indices.append(prevIndices.begin(), prevIndices.end());
   }
 
   /// Rewrite the TransferReadOp, assuming that there are no out-of-bounds
   /// accesses on the to-be-unpacked dimension.
   ///
static void rewriteOp(PatternRewriter &rewriter, TransferReadOp xferOp,
     getStoreIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
     auto vecType = bufferType.getElementType().dyn_cast<VectorType>();
-    auto map = getTransferMinorIdentityMap(xferOp.getShapedType(), vecType);
     auto inBoundsAttr = dropFirstElem(rewriter, xferOp.in_boundsAttr());
-    auto newXfer = vector_transfer_read(vecType, xferOp.source(), xferIndices,
-                                        AffineMapAttr::get(map),
-                                        xferOp.padding(), Value(), inBoundsAttr)
-                       .value;
+    auto newXfer =
+        vector_transfer_read(
+            vecType, xferOp.source(), xferIndices,
+            AffineMapAttr::get(unpackedPermutationMap(xferOp, rewriter)),
 
     if (vecType.getRank() > kTargetRank)
         newXfer.getDefiningOp()->setAttr(kPassLabel, rewriter.getUnitAttr());
 
     memref_store(newXfer, buffer, storeIndices);
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension: Write
   /// padding value to the temporary buffer.
-  static void handleOutOfBoundsDim(PatternRewriter &rewriter,
-                                   TransferReadOp xferOp, Value buffer,
-                                   Value iv) {
+  static void handleOutOfBoundsDim(
+      PatternRewriter &rewriter, TransferReadOp xferOp, Value buffer,
     SmallVector<Value, 8> storeIndices;
     getStoreIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
     auto vecType = bufferType.getElementType().dyn_cast<VectorType>();
     auto vec = std_splat(vecType, xferOp.padding());
     memref_store(vec, buffer, storeIndices);
   }
 
   /// Cleanup after rewriting the op.
   static void cleanup(PatternRewriter &rewriter, TransferReadOp xferOp) {
     rewriter.eraseOp(getStoreOp(xferOp));
     rewriter.eraseOp(xferOp);
   }
 };
 
 /// Codegen strategy for vector TransferWriteOp.
-template <>
+template<>
 struct Strategy<TransferWriteOp> {
   /// Find the temporary buffer allocation. All labeled TransferWriteOps are
   /// used like this, where %buf is either the buffer allocation or a type cast
   /// of the buffer allocation:
   /// ```
   /// %vec = memref.load %buf[...] ...
   /// vector.transfer_write %vec ... { __vector_to_scf_lowering__ } ...
   /// ```
static void rewriteOp(PatternRewriter &rewriter, TransferWriteOp xferOp,
     getLoadIndices(xferOp, loadIndices);
     loadIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto vec = memref_load(buffer, loadIndices);
     auto vecType = vec.value.getType().dyn_cast<VectorType>();
-    auto map = getTransferMinorIdentityMap(xferOp.getShapedType(), vecType);
     auto inBoundsAttr = dropFirstElem(rewriter, xferOp.in_boundsAttr());
-    auto newXfer =
-        vector_transfer_write(Type(), vec, xferOp.source(), xferIndices,
-                              AffineMapAttr::get(map), Value(), inBoundsAttr);
+    auto newXfer = vector_transfer_write(
+        Type(), vec, xferOp.source(), xferIndices,
+        AffineMapAttr::get(unpackedPermutationMap(xferOp, rewriter)), Value(),
+        inBoundsAttr);
 
     if (vecType.getRank() > kTargetRank)
         newXfer.op->setAttr(kPassLabel, rewriter.getUnitAttr());
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension.
-  static void handleOutOfBoundsDim(PatternRewriter &rewriter,
-                                   TransferWriteOp xferOp, Value buffer,
-                                   Value iv) {}
+  static void handleOutOfBoundsDim(
+      PatternRewriter &rewriter, TransferWriteOp xferOp, Value buffer,
 
   /// Cleanup after rewriting the op.
   static void cleanup(PatternRewriter &rewriter, TransferWriteOp xferOp) {
     rewriter.eraseOp(xferOp);
   }
 };
 
 template <typename OpTy>
 LogicalResult checkPrepareXferOp(OpTy xferOp) {
   if (xferOp->hasAttr(kPassLabel))
       return failure();
   if (xferOp.getVectorType().getRank() <= kTargetRank)
       return failure();
   if (xferOp.mask())
       return failure();
-  if (!xferOp.permutation_map().isMinorIdentity())
-    return failure();
   return success();
 }
 
 /// Prepare a TransferReadOp for progressive lowering.
 ///
 /// 1. Allocate a temporary buffer.
 /// 2. Label the TransferReadOp, marking it eligible for progressive lowering.
 /// 3. Store the result of the TransferReadOp into the temporary buffer.
 /// 4. Load the result from the temporary buffer and replace all uses of the
 ///    original TransferReadOp with this load.
 ///
 /// E.g.:
 /// ```
 /// %vec = vector.transfer_read %A[%a, %b, %c], %cst
 ///     : vector<5x4xf32>, memref<?x?x?xf32>
 /// ```
 /// is rewritten to:
 /// ```
 /// %0 = memref.alloca() : memref<vector<5x4xf32>>
 /// %1 = vector.transfer_read %A[%a, %b, %c], %cst
 ///     { __vector_to_scf_lowering__ } : vector<5x4xf32>, memref<?x?x?xf32>
 /// memref.store %1, %0[] : memref<vector<5x4xf32>>
 /// %vec = memref.load %0[] : memref<vector<5x4xf32>>
 /// ```
-struct PrepareTransferReadConversion : public OpRewritePattern<TransferReadOp> {
+struct PrepareTransferReadConversion
+    : public OpRewritePattern<TransferReadOp> {
   using OpRewritePattern<TransferReadOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TransferReadOp xferOp,
                                 PatternRewriter &rewriter) const override {
     if (checkPrepareXferOp(xferOp).failed())
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());

 /// 3. Clean up according to the corresponding Strategy<OpTy>.
 template <typename OpTy>
 struct TransferOpConversion : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy xferOp,
                                 PatternRewriter &rewriter) const override {
     if (!xferOp->hasAttr(kPassLabel))
         return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     // How the buffer can be found depends on OpTy.
     auto buffer = Strategy<OpTy>::getBuffer(xferOp);
     auto bufferType = buffer.getType().template dyn_cast<MemRefType>();
     auto castedType = unpackOneDim(bufferType);
     auto casted = vector_type_cast(castedType, buffer);
 
     auto lb = std_constant_index(0).value;
-    auto ub =
-        std_constant_index(castedType.getDimSize(castedType.getRank() - 1))
-            .value;
+    auto ub = std_constant_index(
+        castedType.getDimSize(castedType.getRank() - 1)).value;
     affineLoopBuilder(lb, ub, 1, [&](Value iv) {
       generateInBoundsCheck(
           xferOp, iv, rewriter, unpackedDim(xferOp),
-          /*inBoundsCase=*/
-          [&](OpBuilder & /*b*/, Location loc) {
-            Strategy<OpTy>::rewriteOp(rewriter, xferOp, casted, iv);
-          },
-          /*outOfBoundsCase=*/
-          [&](OpBuilder & /*b*/, Location loc) {
-            Strategy<OpTy>::handleOutOfBoundsDim(rewriter, xferOp, casted, iv);
-          });
+          /*inBoundsCase=*/[&](OpBuilder& /*b*/, Location loc) {
+      }, /*outOfBoundsCase=*/[&](OpBuilder& /*b*/, Location loc) {
     });
 
     Strategy<OpTy>::cleanup(rewriter, xferOp);
     return success();
   }
 };
 
 /// Compute the indices into the memref for the LoadOp/StoreOp generated as

 /// Codegen strategy for Strided1dTransferOpConversion, depending on the
 /// operation.
 template <typename OpTy>
 struct Strategy1d;
 
 /// Codegen strategy for TransferReadOp.
 template <>
 struct Strategy1d<TransferReadOp> {
-  static void generateForLoopBody(OpBuilder &builder, Location loc,
-                                  TransferReadOp xferOp, Value iv,
-                                  ValueRange loopState) {
+  static void generateForLoopBody(
+      OpBuilder &builder, Location loc, TransferReadOp xferOp, Value iv,
     SmallVector<Value, 8> indices;
     auto dim = get1dMemrefIndices(xferOp, iv, indices);
-    auto ivI32 = std_index_cast(IntegerType::get(builder.getContext(), 32), iv);
+    auto ivI32 = std_index_cast(
+        IntegerType::get(builder.getContext(), 32), iv);
     auto vec = loopState[0];
 
     // In case of out-of-bounds access, leave `vec` as is (was initialized with
     // padding value).
     auto nextVec = generateInBoundsCheck(
         xferOp, iv, builder, dim, TypeRange(xferOp.getVectorType()),
-        /*inBoundsCase=*/
-        [&](OpBuilder & /*b*/, Location loc) {
-          auto val = memref_load(xferOp.source(), indices);
-          return vector_insert_element(val, vec, ivI32.value).value;
-        },
-        /*outOfBoundsCase=*/
-        [&](OpBuilder & /*b*/, Location loc) { return vec; });
+        /*inBoundsCase=*/[&](OpBuilder& /*b*/, Location loc) {
+    }, /*outOfBoundsCase=*/[&](OpBuilder& /*b*/, Location loc) {
+      return vec;
+    });
     builder.create<scf::YieldOp>(loc, nextVec);
   }
 
   static Value initialLoopState(TransferReadOp xferOp) {
     // Inititalize vector with padding value.
     return std_splat(xferOp.getVectorType(), xferOp.padding()).value;
   }
 };
 
 /// Codegen strategy for TransferWriteOp.
 template <>
 struct Strategy1d<TransferWriteOp> {
-  static void generateForLoopBody(OpBuilder &builder, Location loc,
-                                  TransferWriteOp xferOp, Value iv,
-                                  ValueRange /*loopState*/) {
+  static void generateForLoopBody(
+      OpBuilder &builder, Location loc, TransferWriteOp xferOp, Value iv,
     SmallVector<Value, 8> indices;
     auto dim = get1dMemrefIndices(xferOp, iv, indices);
-    auto ivI32 = std_index_cast(IntegerType::get(builder.getContext(), 32), iv);
+    auto ivI32 = std_index_cast(
+        IntegerType::get(builder.getContext(), 32), iv);
 
     // Nothing to do in case of out-of-bounds access.
     generateInBoundsCheck(
         xferOp, iv, builder, dim,
-        /*inBoundsCase=*/[&](OpBuilder & /*b*/, Location loc) {
-          auto val = vector_extract_element(xferOp.vector(), ivI32.value);
-          memref_store(val, xferOp.source(), indices);
-        });
+        /*inBoundsCase=*/[&](OpBuilder& /*b*/, Location loc) {
     builder.create<scf::YieldOp>(loc);
   }
 
-  static Value initialLoopState(TransferWriteOp xferOp) { return Value(); }
+  static Value initialLoopState(TransferWriteOp xferOp) {
+    return Value();
+  }
 };
 
 /// Lower a 1D vector transfer op that operates on a dimension different from
 /// the last one. Instead of accessing contiguous chunks (vectors) of memory,
 /// such ops access memory in a strided fashion.
 ///
 /// 1. Generate a for loop iterating over each vector element.
 /// 2. Inside the loop, generate a InsertElementOp or ExtractElementOp,
struct Strided1dTransferOpConversion : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy xferOp,
                                 PatternRewriter &rewriter) const override {
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto map = xferOp.permutation_map();
 
     if (xferOp.getVectorType().getRank() != 1)
         return failure();
-    if (map.isMinorIdentity()) // Handled by ConvertVectorToLLVM
+    if (map.isMinorIdentity())  // Handled by ConvertVectorToLLVM
         return failure();
     if (xferOp.mask())
         return failure();
 
     // Loop bounds, step, state...
     auto vecType = xferOp.getVectorType();
     auto lb = std_constant_index(0);
     auto ub = std_constant_index(vecType.getDimSize(0));
     auto step = std_constant_index(1);
     auto loopState = Strategy1d<OpTy>::initialLoopState(xferOp);
 
     // Generate for loop.
     rewriter.replaceOpWithNewOp<scf::ForOp>(
         xferOp, lb, ub, step, loopState ? ValueRange(loopState) : ValueRange(),
         [&](OpBuilder &builder, Location loc, Value iv, ValueRange loopState) {
-          ScopedContext nestedScope(builder, loc);
-          Strategy1d<OpTy>::generateForLoopBody(builder, loc, xferOp, iv,
-                                                loopState);
-        });
+      Strategy1d<OpTy>::generateForLoopBody(
+          builder, loc, xferOp, iv, loopState);
 
     return success();
   }
 };
 
-} // namespace
+}  // namespace
 
 namespace mlir {
 
 void populateProgressiveVectorToSCFConversionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<PrepareTransferReadConversion, PrepareTransferWriteConversion,
                TransferOpConversion<TransferReadOp>,
-               TransferOpConversion<TransferWriteOp>,
-               Strided1dTransferOpConversion<TransferReadOp>,
+               TransferOpConversion<TransferWriteOp>>(patterns.getContext());
+
+  if (kTargetRank == 1) {
+    patterns.add<Strided1dTransferOpConversion<TransferReadOp>,
                  Strided1dTransferOpConversion<TransferWriteOp>>(
         patterns.getContext());
   }
+}
 
 struct ConvertProgressiveVectorToSCFPass
     : public ConvertVectorToSCFBase<ConvertProgressiveVectorToSCFPass> {
   void runOnFunction() override {
     RewritePatternSet patterns(getFunction().getContext());
     populateProgressiveVectorToSCFConversionPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
-} // namespace mlir
+}  // namespace mlir
 
-std::unique_ptr<Pass> mlir::createProgressiveConvertVectorToSCFPass() {
+std::unique_ptr<Pass>
+mlir::createProgressiveConvertVectorToSCFPass() {
   return std::make_unique<ConvertProgressiveVectorToSCFPass>();
 }

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-2d.mlir

func @transfer_read_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     memref<?x?xf32>, vector<4x9xf32>
   vector.print %f: vector<4x9xf32>
   return
 }
 
+func @transfer_read_2d_transposed(
+    %A : memref<?x?xf32>, %base1: index, %base2: index) {
+  %fm42 = constant -42.0: f32
+  %f = vector.transfer_read %A[%base1, %base2], %fm42
+      {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} :
+    memref<?x?xf32>, vector<4x9xf32>
+  vector.print %f: vector<4x9xf32>
+  return
+}
+
 func @transfer_write_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fn1 = constant -1.0 : f32
   %vf0 = splat %fn1 : vector<1x4xf32>
   vector.transfer_write %vf0, %A[%base1, %base2]
     {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     vector<1x4xf32>, memref<?x?xf32>
   return
 }
scf.for %j = %c0 to %second step %c1 {
         %fj = sitofp %j32 : i32 to f32
         %fres = addf %fi10, %fj : f32
         memref.store %fres, %A[%i, %j] : memref<?x?xf32>
     }
   }
   // On input, memory contains [[ 0, 1, 2, ...], [10, 11, 12, ...], ...]
   // Read shifted by 2 and pad with -42:
   call @transfer_read_2d(%A, %c1, %c2) : (memref<?x?xf32>, index, index) -> ()
+  // Same as above, but transposed
+  call @transfer_read_2d_transposed(%A, %c1, %c2)
+      : (memref<?x?xf32>, index, index) -> ()
   // Write into memory shifted by 3
   call @transfer_write_2d(%A, %c3, %c1) : (memref<?x?xf32>, index, index) -> ()
   // Read shifted by 0 and pad with -42:
   call @transfer_read_2d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()
+  // Same as above, but transposed
+  call @transfer_read_2d_transposed(%A, %c0, %c0)
+      : (memref<?x?xf32>, index, index) -> ()
   return
 }
 
 // CHECK: ( ( 12, 13, -42, -42, -42, -42, -42, -42, -42 ), ( 22, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+// CHECK: ( ( 12, 22, -42, -42, -42, -42, -42, -42, -42 ), ( 13, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
 // CHECK: ( ( 0, 1, 2, 3, -42, -42, -42, -42, -42 ), ( 10, 11, 12, 13, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+// CHECK: ( ( 0, 10, 20, -42, -42, -42, -42, -42, -42 ), ( 1, 11, 21, -42, -42, -42, -42, -42, -42 ), ( 2, 12, 22, -42, -42, -42, -42, -42, -42 ), ( 3, 13, 23, -42, -42, -42, -42, -42, -42 ) )

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-3d.mlir

func @transfer_read_3d(%A : memref<?x?x?x?xf32>,
                        %o: index, %a: index, %b: index, %c: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
       : memref<?x?x?x?xf32>, vector<2x5x3xf32>
   vector.print %f: vector<2x5x3xf32>
   return
 }
 
+func @transfer_read_3d_transposed(%A : memref<?x?x?x?xf32>,
+                                  %o: index, %a: index, %b: index, %c: index) {
+  %fm42 = constant -42.0: f32
+  %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
+      {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1)>}
+      : memref<?x?x?x?xf32>, vector<3x5x3xf32>
+  vector.print %f: vector<3x5x3xf32>
+  return
+}
+
 func @transfer_write_3d(%A : memref<?x?x?x?xf32>,
                         %o: index, %a: index, %b: index, %c: index) {
   %fn1 = constant -1.0 : f32
   %vf0 = splat %fn1 : vector<2x9x3xf32>
   vector.transfer_write %vf0, %A[%o, %a, %b, %c]
       : vector<2x9x3xf32>, memref<?x?x?x?xf32>
   return
 }
func @entry() {
   }
 
   call @transfer_read_3d(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
   call @transfer_write_3d(%A, %c0, %c0, %c1, %c1)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
   call @transfer_read_3d(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  call @transfer_read_3d_transposed(%A, %c0, %c0, %c0, %c0)
+      : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
   return
 }
 
 // CHECK: ( ( ( 0, 0, -42 ), ( 2, 3, -42 ), ( 4, 6, -42 ), ( 6, 9, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, 33, -42 ), ( 24, 36, -42 ), ( 26, 39, -42 ), ( -42, -42, -42 ) ) )
 // CHECK: ( ( ( 0, 0, -42 ), ( 2, -1, -42 ), ( 4, -1, -42 ), ( 6, -1, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, -1, -42 ), ( 24, -1, -42 ), ( 26, -1, -42 ), ( -42, -42, -42 ) ) )
+// CHECK: ( ( ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ) ), ( ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ) ), ( ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ) ) )

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-to-loops.mlir

 // RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext,%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
+// RUN: mlir-opt %s -test-progressive-convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void  \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext,%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
 #map0 = affine_map<(d0, d1) -> (d1, d0)>
 #map1 = affine_map<(d0, d1) -> (d1)>
 
 func private @print_memref_f32(memref<*xf32>)
 
 func @alloc_2d_filled_f32(%arg0: index, %arg1: index) -> memref<?x?xf32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index