[mlir][vector] Allow unroll of vector.contract, vector.transfer_read/write in arbitrary order

[christopherbate]

Summary

Adds supprot for vector unroll transformations to unroll in different
orders. For example, the vector.contract can be unrolled into a
smaller set of contractions. There is a choice of how to unroll the
decomposition based on the traversal order of (dim0, dim1, dim2).
The choice of traversal order can now be specified by a callback which
given by the caller of the transform. For now, only the
vector.contract, vector.transfer_read/transfer_write operations provide
support for the callback.

[christopherbate]

Thomas had some comments from offline discussion (paraphrase)

• The functor utility DecomposeShapeIterator is more efficient but we could just use getVectorOffset be pre-permuting the shapes.

Personally I feel that the functor utility is more readable than the getVectorOffset which links into helper functions in a different file. Can we keep this and then slowly remove the getVectorOffset utility?

• Add vector.transfer_read/transfer_write support in this commit as well?

[christopherbate]

Add support for changing unroll order of transfer_read/transfer_write

[christopherbate]

Update commit message.

[ThomasRaoux]

Thanks! Looks good to me, I added few style related comments.

mlir/lib/Dialect/Vector/Transforms/VectorUnrollDistribute.cpp

92

I think having an explicit function name would make it easier to read the code than overloading operator().

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

522

nit: skip braces on single line if:
https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

[christopherbate]

Address reviewer comments and cleanup dangling comments in code.

[christopherbate]

Address comments

[okkwon]

Hi, just for heads-up, this change caused many failures in IREE tests including compilation errors and accuracy errors.

[christopherbate]

In D127004#3564403, @okkwon wrote:

Hi, just for heads-up, this change caused many failures in IREE tests including compilation errors and accuracy errors.

Ok, thanks. Do you mean C++ compilation or mlir compilation errors?

There were build failures on some configs/compilers that I fixed with one-line change a392a39f75af [mlir][vector] fix typo in vector unroll transform.

Do we need t o revert?
Let me know if you can share the compilation/accuracy errors.

[okkwon]

It would be nice if you don't mind reverting.

There are compilation errors in compiling some mlir sources. There is a crash too.

The PR for integration: https://github.com/google/iree/pull/9369
(You can find the failing tests there.)

Crash
https://source.cloud.google.com/results/invocations/727a0eb1-7000-4b8e-ab2f-36012c00f858/targets/iree%2Fgcp_ubuntu%2Fcmake%2Flinux%2Fx86-swiftshader%2Fpresubmit/log

[christopherbate]

In D127004#3564660, @okkwon wrote:

It would be nice if you don't mind reverting.

There are compilation errors in compiling some mlir sources. There is a crash too.

The PR for integration: https://github.com/google/iree/pull/9369
(You can find the failing tests there.)

Crash
https://source.cloud.google.com/results/invocations/727a0eb1-7000-4b8e-ab2f-36012c00f858/targets/iree%2Fgcp_ubuntu%2Fcmake%2Flinux%2Fx86-swiftshader%2Fpresubmit/log

Ok yeah definitely looks like I made a mistake somewhere. I'll revert now, thanks for sharing those tests

[christopherbate]

Fix mistake in default unroll permutation for vector.contract

[ThomasRaoux]

Thanks, LGTM

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

struct UnrollVectorOptions {
   /// Set the native shape to use for unrolling.
   UnrollVectorOptions &setNativeShape(ArrayRef<int64_t> shape) {
     SmallVector<int64_t, 4> tsShape(shape.begin(), shape.end());
     nativeShape = [=](Operation *) -> Optional<SmallVector<int64_t, 4>> {
       return tsShape;
     };
     return *this;
   }
+
+  /// Function that returns the traversal order (in terms of "for loop order",
+  /// i.e. slowest varying dimension to fastest varying dimension) that shoudl
+  /// be used when unrolling the given operation into units of the native vector
+  /// size.
+  using UnrollTraversalOrderFnType =
+      std::function<Optional<SmallVector<int64_t>>(Operation *op)>;
+  UnrollTraversalOrderFnType traversalOrderCallback = nullptr;
+  UnrollVectorOptions &
+  setUnrollTraversalOrderFn(UnrollTraversalOrderFnType traversalOrderFn) {
+    traversalOrderCallback = std::move(traversalOrderFn);
+    return *this;
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // Vector transformation exposed as populate functions over rewrite patterns.
 //===----------------------------------------------------------------------===//
 
 /// Insert TransposeLowering patterns into extraction/insertion.
 void populateVectorTransposeLoweringPatterns(

mlir/lib/Dialect/Vector/Transforms/VectorUnrollDistribute.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
+#include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/MapVector.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/ADT/STLExtras.h"
+#include <numeric>
 
 #define DEBUG_TYPE "vector-unrolling"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 /// During unrolling from `originalShape` to `targetShape` return the offset for
 /// the slice `index`.
 static SmallVector<int64_t, 4> getVectorOffset(ArrayRef<int64_t> originalShape,
                                                ArrayRef<int64_t> targetShape,
                                                int64_t index) {
   SmallVector<int64_t, 4> dstSliceStrides =
       computeStrides(originalShape, targetShape);
   SmallVector<int64_t, 4> vectorOffsets = delinearize(dstSliceStrides, index);
   SmallVector<int64_t, 4> elementOffsets =
       computeElementOffsetsFromVectorSliceOffsets(targetShape, vectorOffsets);
   return elementOffsets;
 }
 
+/// A functor that accomplishes the same thing as `getVectorOffset` but allows
+/// for reordering the traversal of the dimensions. The order of traversal is
+/// given in "for loop order" (outer to inner).
+namespace {
+class DecomposeShapeIterator {
+private:
+  SmallVector<int64_t, 4> vectorShape;
+  SmallVector<int64_t> loopOrder;
+  SmallVector<int64_t> sliceStrides;
+  int64_t maxIndexVal{1};
+
+public:
+  DecomposeShapeIterator(ArrayRef<int64_t> originalShape,
+                         ArrayRef<int64_t> targetShape,
+                         ArrayRef<int64_t> loopOrder)
+      : vectorShape(targetShape.begin(), targetShape.end()),
+        loopOrder(loopOrder.begin(), loopOrder.end()),
+        sliceStrides(originalShape.size()) {
+    assert(originalShape.size() == targetShape.size());
+    assert(loopOrder.size() == targetShape.size());
+
+    // Compute the count for each dimension.
+    SmallVector<int64_t> sliceDimCounts(originalShape.size());
+    for (unsigned r = 0; r < originalShape.size(); ++r) {
+      sliceDimCounts[r] = ceilDiv(originalShape[r], targetShape[r]);
+      maxIndexVal *= sliceDimCounts[r];
+    }
+
+    // Reversing "loop order" gives dimensions from fastest varying to slowest
+    // varying (smallest stride to largest stride).
+    int64_t accum = 1;
+    for (auto idx : llvm::reverse(loopOrder)) {
+      sliceStrides[idx] = accum;
+      accum *= sliceDimCounts[idx];
+    }
+  }
+
+  // Turn the linear index into a d-tuple based on units of vectors of size
+  // `vectorShape`. The linear index is assumed to represent traversal of the
+  // dimensions based on `order`.
+  SmallVector<int64_t> delinearize(int64_t index) const {
+    // Traverse in for loop order (largest stride to smallest stride).
+    SmallVector<int64_t> vectorOffsets(sliceStrides.size());
+    for (auto idx : loopOrder) {
+      vectorOffsets[idx] = index / sliceStrides[idx];
+      index %= sliceStrides[idx];
+    }
+    return vectorOffsets;
+  }
+
+  int64_t maxIndex() const { return maxIndexVal; }
+
+  /// Return the offset within d-tuple based on the ordering given by
+  /// `loopOrder`.
+  SmallVector<int64_t> getVectorOffset(int64_t index) const {
+    SmallVector<int64_t> vectorOffsets = delinearize(index);
+    SmallVector<int64_t> elementOffsets =
+        computeElementOffsetsFromVectorSliceOffsets(vectorShape, vectorOffsets);
+    return elementOffsets;
+  }
+};
+} // namespace
+
 /// Compute the indices of the slice `index` for a tranfer op.
-static SmallVector<Value>
-sliceTransferIndices(int64_t index, ArrayRef<int64_t> originalShape,
-                     ArrayRef<int64_t> targetShape, ArrayRef<Value> indices,
-                     AffineMap permutationMap, Location loc,
+static SmallVector<Value> sliceTransferIndices(ArrayRef<int64_t> elementOffsets,
+                                               ArrayRef<Value> indices,
+                                               AffineMap permutationMap,
+                                               Location loc,
                                                OpBuilder &builder) {
   MLIRContext *ctx = builder.getContext();
   auto isBroadcast = [](AffineExpr expr) {
     if (auto constExpr = expr.dyn_cast<AffineConstantExpr>())
       return constExpr.getValue() == 0;
     return false;
   };
-  SmallVector<int64_t, 4> elementOffsets =
-      getVectorOffset(originalShape, targetShape, index);
   // Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'.
   SmallVector<Value> slicedIndices(indices.begin(), indices.end());
   for (const auto &dim : llvm::enumerate(permutationMap.getResults())) {
     if (isBroadcast(dim.value()))
       continue;
     unsigned pos = dim.value().cast<AffineDimExpr>().getPosition();
     auto expr = getAffineDimExpr(0, builder.getContext()) +
                 getAffineConstantExpr(elementOffsets[dim.index()], ctx);
if (!targetShape)
     return llvm::None;
   auto maybeShapeRatio = shapeRatio(*maybeUnrollShape, *targetShape);
   if (!maybeShapeRatio ||
       llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
     return llvm::None;
   return targetShape;
 }
 
+static SmallVector<int64_t>
+getUnrollOrder(unsigned numLoops, Operation *op,
+               const vector::UnrollVectorOptions &options) {
+  SmallVector<int64_t> loopOrder =
+      llvm::to_vector(llvm::seq<int64_t>(0, static_cast<int64_t>(numLoops)));
+  if (options.traversalOrderCallback != nullptr) {
+    Optional<SmallVector<int64_t>> order = options.traversalOrderCallback(op);
+    if (order.hasValue()) {
+      loopOrder = std::move(*order);
+    }
+  }
+  return loopOrder;
+}
+
 namespace {
 
 struct UnrollTransferReadPattern
     : public OpRewritePattern<vector::TransferReadOp> {
   UnrollTransferReadPattern(MLIRContext *context,
                             const vector::UnrollVectorOptions &options)
       : OpRewritePattern<vector::TransferReadOp>(context, /*benefit=*/1),
         options(options) {}
   LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                 PatternRewriter &rewriter) const override {
     // TODO: support 0-d corner case.
     if (readOp.getTransferRank() == 0)
       return failure();
     if (readOp.getMask())
       return failure();
     auto targetShape = getTargetShape(options, readOp);
     if (!targetShape)
       return failure();
     auto sourceVectorType = readOp.getVectorType();
     SmallVector<int64_t, 4> strides(targetShape->size(), 1);
     Location loc = readOp.getLoc();
     ArrayRef<int64_t> originalSize = readOp.getVectorType().getShape();
-    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
-    // Compute shape ratio of 'shape' and 'sizes'.
-    int64_t sliceCount = computeMaxLinearIndex(ratio);
+
     // Prepare the result vector;
     Value result = rewriter.create<arith::ConstantOp>(
         loc, sourceVectorType, rewriter.getZeroAttr(sourceVectorType));
     auto targetType =
         VectorType::get(*targetShape, sourceVectorType.getElementType());
     SmallVector<Value, 4> originalIndices(readOp.getIndices().begin(),
                                           readOp.getIndices().end());
-    for (int64_t i = 0; i < sliceCount; i++) {
+
+    SmallVector<int64_t> loopOrder =
+        getUnrollOrder(originalSize.size(), readOp, options);
+    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
+                                          loopOrder);
+    for (int64_t i = 0; i < indexToOffsets.maxIndex(); i++) {
+      SmallVector<int64_t, 4> elementOffsets =
+          indexToOffsets.getVectorOffset(i);
       SmallVector<Value, 4> indices =
-          sliceTransferIndices(i, originalSize, *targetShape, originalIndices,
+          sliceTransferIndices(elementOffsets, originalIndices,
                                readOp.getPermutationMap(), loc, rewriter);
       auto slicedRead = rewriter.create<vector::TransferReadOp>(
           loc, targetType, readOp.getSource(), indices,
           readOp.getPermutationMapAttr(), readOp.getPadding(), readOp.getMask(),
           readOp.getInBoundsAttr());
 
-      SmallVector<int64_t, 4> elementOffsets =
-          getVectorOffset(originalSize, *targetShape, i);
       result = rewriter.create<vector::InsertStridedSliceOp>(
           loc, slicedRead, result, elementOffsets, strides);
     }
     rewriter.replaceOp(readOp, result);
     return success();
   }
 
 private:
if (writeOp.getMask())
       return failure();
     auto targetShape = getTargetShape(options, writeOp);
     if (!targetShape)
       return failure();
     auto sourceVectorType = writeOp.getVectorType();
     SmallVector<int64_t, 4> strides(targetShape->size(), 1);
     Location loc = writeOp.getLoc();
     ArrayRef<int64_t> originalSize = sourceVectorType.getShape();
-    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
-    // Compute shape ratio of 'shape' and 'sizes'.
-    int64_t sliceCount = computeMaxLinearIndex(ratio);
     SmallVector<Value, 4> originalIndices(writeOp.getIndices().begin(),
                                           writeOp.getIndices().end());
+
+    SmallVector<int64_t> loopOrder =
+        getUnrollOrder(originalSize.size(), writeOp, options);
+    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
+                                          loopOrder);
     Value resultTensor;
-    for (int64_t i = 0; i < sliceCount; i++) {
+    for (int64_t i = 0; i < indexToOffsets.maxIndex(); i++) {
       SmallVector<int64_t, 4> elementOffsets =
-          getVectorOffset(originalSize, *targetShape, i);
+          indexToOffsets.getVectorOffset(i);
       Value slicedVector = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, writeOp.getVector(), elementOffsets, *targetShape, strides);
-
       SmallVector<Value, 4> indices =
-          sliceTransferIndices(i, originalSize, *targetShape, originalIndices,
+          sliceTransferIndices(elementOffsets, originalIndices,
                                writeOp.getPermutationMap(), loc, rewriter);
       Operation *slicedWrite = rewriter.create<vector::TransferWriteOp>(
           loc, slicedVector, resultTensor ? resultTensor : writeOp.getSource(),
           indices, writeOp.getPermutationMapAttr(), writeOp.getInBoundsAttr());
       // For the tensor case update the destination for the next transfer write.
       if (!slicedWrite->getResults().empty())
         resultTensor = slicedWrite->getResult(0);
     }
struct UnrollContractionPattern
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
     auto targetShape = getTargetShape(options, contractOp);
     if (!targetShape)
       return failure();
     auto dstVecType = contractOp.getResultType().cast<VectorType>();
     SmallVector<int64_t, 4> originalSize = *contractOp.getShapeForUnroll();
-    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
 
-    // Compute shape ratio of 'shape' and 'sizes'.
-    int64_t sliceCount = computeMaxLinearIndex(ratio);
     Location loc = contractOp.getLoc();
     unsigned accIndex = vector::ContractionOp::getAccOperandIndex();
     AffineMap dstAffineMap = contractOp.getIndexingMaps()[accIndex];
     llvm::MapVector<
         SmallVector<int64_t>, Value,
         llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
         accCache;
+
+    SmallVector<int64_t> loopOrder = getUnrollOrder(
+        contractOp.getIteratorTypes().size(), contractOp, options);
+    DecomposeShapeIterator indexToOffsets(originalSize, *targetShape,
+                                          loopOrder);
+    const int64_t sliceCount = indexToOffsets.maxIndex();
     for (int64_t i = 0; i < sliceCount; i++) {
-      SmallVector<int64_t, 4> offsets =
-          getVectorOffset(originalSize, *targetShape, i);
+      SmallVector<int64_t, 4> offsets = indexToOffsets.getVectorOffset(i);
       SmallVector<Value, 4> slicesOperands(contractOp.getNumOperands());
 
       // Helper to coompute the new shape of each operand and extract the slice.
       auto extractOperand = [&](unsigned index, Value operand,
                                 AffineMap permutationMap,
                                 ArrayRef<int64_t> operandOffets) {
         SmallVector<int64_t> operandShape = applyPermutationMap(
             permutationMap, ArrayRef<int64_t>(*targetShape));

mlir/test/Dialect/Vector/vector-transfer-unroll.mlir

 // RUN: mlir-opt %s -test-vector-transfer-unrolling-patterns --split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-vector-transfer-unrolling-patterns=reverse-unroll-order --split-input-file | FileCheck %s --check-prefix=ORDER
 
 // CHECK-LABEL: func @transfer_read_unroll
 //       CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
 //       CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC0:.*]] = vector.insert_strided_slice %[[VTR0]], %{{.*}} {offsets = [0, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
 //  CHECK-NEXT:   %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC1:.*]] = vector.insert_strided_slice %[[VTR1]], %[[VEC0]] {offsets = [0, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
 //  CHECK-NEXT:   %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC2:.*]] = vector.insert_strided_slice %[[VTR2]], %[[VEC1]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
 //  CHECK-NEXT:   %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC3:.*]] = vector.insert_strided_slice %[[VTR3]], %[[VEC2]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
 //  CHECK-NEXT:   return %[[VEC3]] : vector<4x4xf32>
 
+// ORDER-LABEL: func @transfer_read_unroll
+//       ORDER-DAG:   %[[C2:.*]] = arith.constant 2 : index
+//       ORDER-DAG:   %[[C0:.*]] = arith.constant 0 : index
+//       ORDER:   %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC0:.*]] = vector.insert_strided_slice %[[VTR0]], %{{.*}} {offsets = [0, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+//  ORDER-NEXT:   %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC1:.*]] = vector.insert_strided_slice %[[VTR1]], %[[VEC0]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+//  ORDER-NEXT:   %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC2:.*]] = vector.insert_strided_slice %[[VTR2]], %[[VEC1]] {offsets = [0, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+//  ORDER-NEXT:   %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC3:.*]] = vector.insert_strided_slice %[[VTR3]], %[[VEC2]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+//  ORDER-NEXT:   return %[[VEC3]] : vector<4x4xf32>
+
 func.func @transfer_read_unroll(%arg0 : memref<4x4xf32>) -> vector<4x4xf32> {
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
   %0 = vector.transfer_read %arg0[%c0, %c0], %cf0 : memref<4x4xf32>, vector<4x4xf32>
   return %0 : vector<4x4xf32>
 }
 
 // CHECK-LABEL: func @transfer_write_unroll
 //       CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
 //       CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[S0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 //  CHECK-NEXT:   vector.transfer_write %[[S0]], {{.*}}[%[[C0]], %[[C0]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
 //  CHECK-NEXT:   %[[S1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 //  CHECK-NEXT:   vector.transfer_write %[[S1]], {{.*}}[%[[C0]], %[[C2]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
 //  CHECK-NEXT:   %[[S2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 //  CHECK-NEXT:   vector.transfer_write %[[S2]], {{.*}}[%[[C2]], %[[C0]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
 //  CHECK-NEXT:   %[[S3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 //  CHECK-NEXT:   vector.transfer_write %[[S3]], {{.*}}[%[[C2]], %[[C2]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
 //  CHECK-NEXT:   return
 
+// ORDER-LABEL: func @transfer_write_unroll
+//       ORDER-DAG:   %[[C2:.*]] = arith.constant 2 : index
+//       ORDER-DAG:   %[[C0:.*]] = arith.constant 0 : index
+//       ORDER:   %[[S0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
+//  ORDER-NEXT:   vector.transfer_write %[[S0]], {{.*}}[%[[C0]], %[[C0]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  ORDER-NEXT:   %[[S1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
+//  ORDER-NEXT:   vector.transfer_write %[[S1]], {{.*}}[%[[C2]], %[[C0]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  ORDER-NEXT:   %[[S2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
+//  ORDER-NEXT:   vector.transfer_write %[[S2]], {{.*}}[%[[C0]], %[[C2]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  ORDER-NEXT:   %[[S3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
+//  ORDER-NEXT:   vector.transfer_write %[[S3]], {{.*}}[%[[C2]], %[[C2]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  ORDER-NEXT:   return
+
 func.func @transfer_write_unroll(%arg0 : memref<4x4xf32>, %arg1 : vector<4x4xf32>) {
   %c0 = arith.constant 0 : index
   vector.transfer_write %arg1, %arg0[%c0, %c0] : vector<4x4xf32>, memref<4x4xf32>
   return
 }
 
 // CHECK-LABEL: func @transfer_readwrite_unroll
 //       CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index

 //  CHECK-NEXT:   %[[VEC2:.*]] = vector.insert_strided_slice %[[VTR2]], %[[VEC1]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
 //  CHECK-NEXT:   %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C2]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC3:.*]] = vector.insert_strided_slice %[[VTR3]], %[[VEC2]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
 //  CHECK-NEXT:   %[[VTR4:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C4]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC4:.*]] = vector.insert_strided_slice %[[VTR4]], %[[VEC3]] {offsets = [4, 0], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
 //  CHECK-NEXT:   %[[VTR5:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C4]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VEC5:.*]] = vector.insert_strided_slice %[[VTR5]], %[[VEC4]] {offsets = [4, 2], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
 //  CHECK-NEXT:   return %[[VEC5]] : vector<6x4xf32>
+
+// ORDER-LABEL: func @transfer_read_unroll_different_rank
+//       ORDER-DAG:   %[[C4:.*]] = arith.constant 4 : index
+//       ORDER-DAG:   %[[C2:.*]] = arith.constant 2 : index
+//       ORDER-DAG:   %[[C0:.*]] = arith.constant 0 : index
+//       ORDER:   %[[VTR0:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC0:.*]] = vector.insert_strided_slice %[[VTR0]], %{{.*}} {offsets = [0, 0], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
+//  ORDER-NEXT:   %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C2]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC1:.*]] = vector.insert_strided_slice %[[VTR1]], %[[VEC0]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
+//  ORDER-NEXT:   %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]], %[[C4]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC2:.*]] = vector.insert_strided_slice %[[VTR2]], %[[VEC1]] {offsets = [4, 0], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
+//  ORDER-NEXT:   %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C0]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC3:.*]] = vector.insert_strided_slice %[[VTR3]], %[[VEC2]] {offsets = [0, 2], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
+//  ORDER-NEXT:   %[[VTR4:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C2]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC4:.*]] = vector.insert_strided_slice %[[VTR4]], %[[VEC3]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
+//  ORDER-NEXT:   %[[VTR5:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]], %[[C4]]], %{{.*}} : memref<?x?x?xf32>, vector<2x2xf32>
+//  ORDER-NEXT:   %[[VEC5:.*]] = vector.insert_strided_slice %[[VTR5]], %[[VEC4]] {offsets = [4, 2], strides = [1, 1]} : vector<2x2xf32> into vector<6x4xf32>
+//  ORDER-NEXT:   return %[[VEC5]] : vector<6x4xf32>
+
 #map0 = affine_map<(d0, d1, d2) -> (d2, d0)>
 func.func @transfer_read_unroll_different_rank(%arg0 : memref<?x?x?xf32>) -> vector<6x4xf32> {
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
   %0 = vector.transfer_read %arg0[%c0, %c0, %c0], %cf0 {permutation_map = #map0} : memref<?x?x?xf32>, vector<6x4xf32>
   return %0 : vector<6x4xf32>
 }

mlir/test/Dialect/Vector/vector-unroll-options.mlir

 // RUN: mlir-opt %s -test-vector-unrolling-patterns=unroll-based-on-type | FileCheck %s
+// RUN: mlir-opt %s -test-vector-unrolling-patterns="unroll-based-on-type unroll-order=2,0,1"  | FileCheck %s --check-prefix=ORDER
+// RUN: mlir-opt %s -test-vector-unrolling-patterns="unroll-based-on-type unroll-order=0,3,1,2" | FileCheck %s --check-prefix=BATCHED
 
-func.func @vector_contract_f32(%lhs : vector<8x8xf32>, %rhs : vector<8x8xf32>,
+func.func @vector_contract_f32(%lhs : vector<8x4xf32>, %rhs : vector<8x4xf32>,
                           %init : vector<8x8xf32>) -> vector<8x8xf32> {
   %0 = vector.contract
          {indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
                            affine_map<(i, j, k) -> (j, k)>,
                            affine_map<(i, j, k) -> (i, j)>],
           iterator_types = ["parallel", "parallel", "reduction"]}
-       %lhs, %rhs, %init : vector<8x8xf32>, vector<8x8xf32> into vector<8x8xf32>
+       %lhs, %rhs, %init : vector<8x4xf32>, vector<8x4xf32> into vector<8x8xf32>
   return %0 : vector<8x8xf32>
 }
 // CHECK-LABEL: func @vector_contract_f32
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
-//  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+// CHECK-SAME: [[arg0:%.+]]: vector<8x4xf32>, [[arg1:%.+]]: vector<8x4xf32>, [[arg2:%.+]]: vector<8x8xf32>
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [0, 0]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [0, 0]
+//       CHECK:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  CHECK-SAME:   offsets = [0, 0]
+//       CHECK:   [[accum1:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [0, 2]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [0, 2]
+//       CHECK:   [[accum2:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum1]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [0, 0]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [4, 0]
+//       CHECK:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  CHECK-SAME:   offsets = [0, 4]
+//       CHECK:   [[accum3:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [0, 2]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [4, 2]
+//       CHECK:   [[accum4:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum3]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [4, 0]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [0, 0]
+//       CHECK:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  CHECK-SAME:   offsets = [4, 0]
+//       CHECK:   [[accum5:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [4, 2]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [0, 2]
+//       CHECK:   [[accum6:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum5]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [4, 0]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [4, 0]
+//       CHECK:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  CHECK-SAME:   offsets = [4, 4]
+//       CHECK:   [[accum7:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
-//       CHECK:   vector.contract {
+
+//       CHECK:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  CHECK-SAME:   offsets = [4, 2]
+//       CHECK:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  CHECK-SAME:   offsets = [4, 2]
+//       CHECK:   [[accum8:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum7]]
 //  CHECK-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
 //       CHECK:   return
 
+// ORDER-LABEL: func @vector_contract_f32
+// ORDER-SAME: [[arg0:%.+]]: vector<8x4xf32>, [[arg1:%.+]]: vector<8x4xf32>, [[arg2:%.+]]: vector<8x8xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [0, 0]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [0, 0]
+//       ORDER:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  ORDER-SAME:   offsets = [0, 0]
+//       ORDER:   [[accum1:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [0, 0]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [4, 0]
+//       ORDER:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  ORDER-SAME:   offsets = [0, 4]
+//       ORDER:   [[accum2:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [4, 0]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [0, 0]
+//       ORDER:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  ORDER-SAME:   offsets = [4, 0]
+//       ORDER:   [[accum3:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [4, 0]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [4, 0]
+//       ORDER:   [[c:%.+]] = vector.extract_strided_slice [[arg2]] 
+//  ORDER-SAME:   offsets = [4, 4]
+//       ORDER:   [[accum4:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[c]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [0, 2]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [0, 2]
+//       ORDER:   [[accum5:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum1]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [0, 2]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [4, 2]
+//       ORDER:   [[accum6:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum2]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [4, 2]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [0, 2]
+//       ORDER:   [[accum7:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum3]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   [[a:%.+]] = vector.extract_strided_slice [[arg0]] 
+//  ORDER-SAME:   offsets = [4, 2]
+//       ORDER:   [[b:%.+]] = vector.extract_strided_slice [[arg1]] 
+//  ORDER-SAME:   offsets = [4, 2]
+//       ORDER:   [[accum8:%.+]] = vector.contract {{{.*}}} [[a]], [[b]], [[accum4]]
+//  ORDER-SAME:     vector<4x2xf32>, vector<4x2xf32> into vector<4x4xf32>
+
+//       ORDER:   return
+
+
+
 func.func @vector_contract_f16(%lhs : vector<8x8xf16>, %rhs : vector<8x8xf16>,
                           %init : vector<8x8xf16>) -> vector<8x8xf16> {
   %0 = vector.contract
          {indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
                            affine_map<(i, j, k) -> (j, k)>,
                            affine_map<(i, j, k) -> (i, j)>],
           iterator_types = ["parallel", "parallel", "reduction"]}
        %lhs, %rhs, %init : vector<8x8xf16>, vector<8x8xf16> into vector<8x8xf16>

 //       CHECK:   %[[V5:.*]] = vector.insert_strided_slice %[[T5]], %[[V4]] {offsets = [1, 0, 0, 2], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
 //       CHECK:   %[[E6:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 0, 4], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
 //       CHECK:   %[[T6:.*]] = vector.transpose %[[E6]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
 //       CHECK:   %[[V6:.*]] = vector.insert_strided_slice %[[T6]], %[[V5]] {offsets = [1, 0, 4, 0], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
 //       CHECK:   %[[E7:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 2, 0, 4], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
 //       CHECK:   %[[T7:.*]] = vector.transpose %[[E7]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
 //       CHECK:   %[[V7:.*]] = vector.insert_strided_slice %[[T7]], %[[V6]] {offsets = [1, 0, 4, 2], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
 //       CHECK:   return %[[V7]] : vector<2x3x8x4xf32>
+
+// -----
+
+func.func @vector_contract_batched(%lhs: vector<8x8x4xf32>, %rhs: vector<8x8x4xf32>, %init: vector<8x8x8xf32>) -> vector<8x8x8xf32> {    
+  %0 = vector.contract
+         {indexing_maps = [affine_map<(d0,d1,d2,c0) -> (d0,d1,c0)>,
+                           affine_map<(d0,d1,d2,c0) -> (d0,d2,c0)>,
+                           affine_map<(d0,d1,d2,c0) -> (d0,d1,d2)>],
+          iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+       %lhs, %rhs, %init : vector<8x8x4xf32>, vector<8x8x4xf32> into vector<8x8x8xf32>
+  return %0 : vector<8x8x8xf32>
+}
+
+
+//    CHECK-LABEL: vector_contract_batched
+// CHECK-COUNT-16: vector.contract
+//      CHECK-NOT: vector.contract
+//          CHECK: return 
+
+//    UNROLL-LABEL: vector_contract_batched
+//  UNROLL-COUNT-1: vector.contract
+//      UNROLL-NOT: vector.contract
+//          UNROLL: return
+
+
+//    BATCHED-LABEL: vector_contract_batched
+// BATCHED-COUNT-16: vector.contract
+//      BATCHED-NOT: vector.contract
+//          BATCHED: return

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::vector;
 
 namespace {
 
populateVectorUnrollPatterns(
                       .setFilterConstraint([](Operation *op) {
                         return success(isa<vector::TransposeOp>(op));
                       }));
 
     if (unrollBasedOnType) {
       UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
           [](Operation *op) -> Optional<SmallVector<int64_t, 4>> {
         vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
-        SmallVector<int64_t, 4> nativeShape = {4, 4, 2};
-        if (auto floatType = contractOp.getLhsType()
-                                 .getElementType()
-                                 .dyn_cast<FloatType>()) {
-          if (floatType.getWidth() == 16) {
-            nativeShape[2] = 4;
-          }
-        }
+        SmallVector<int64_t, 4> nativeShape(
+            contractOp.getIteratorTypes().size(), 4);
+        Type lhsType = contractOp.getLhsType().getElementType();
+        nativeShape[nativeShape.size() - 1] = lhsType.isF16() ? 4 : 2;
         return nativeShape;
       };
+
+      UnrollVectorOptions opts;
+      opts.setNativeShapeFn(nativeShapeFn)
+          .setFilterConstraint(
+              [](Operation *op) { return success(isa<ContractionOp>(op)); });
+
+      if (!unrollOrder.empty()) {
+        opts.setUnrollTraversalOrderFn([this](Operation *op)
+                                           -> Optional<SmallVector<int64_t>> {
+          vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
+          if (contractOp.getIteratorTypes().size() == unrollOrder.size())
+            return SmallVector<int64_t>(unrollOrder.begin(), unrollOrder.end());
+          return None;
+        });
+      }
+      populateVectorUnrollPatterns(patterns, opts);
+    } else {
+      auto nativeShapeFn =
+          [](Operation *op) -> Optional<SmallVector<int64_t, 4>> {
+        auto contractOp = dyn_cast<ContractionOp>(op);
+        if (!contractOp)
+          return None;
+        return SmallVector<int64_t, 4>(contractOp.getIteratorTypes().size(), 2);
+      };
       populateVectorUnrollPatterns(patterns,
                                    UnrollVectorOptions()
                                        .setNativeShapeFn(nativeShapeFn)
                                        .setFilterConstraint([](Operation *op) {
                                          return success(isa<ContractionOp>(op));
                                        }));
-    } else {
-      populateVectorUnrollPatterns(
-          patterns, UnrollVectorOptions()
-                        .setNativeShape(ArrayRef<int64_t>{2, 2, 2})
-                        .setFilterConstraint([](Operation *op) {
-                          return success(isa<ContractionOp>(op));
-                        }));
     }
     populateVectorToVectorCanonicalizationPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 
+  ListOption<int64_t> unrollOrder{*this, "unroll-order",
+                                  llvm::cl::desc("set the unroll order"),
+                                  llvm::cl::ZeroOrMore};
+
   Option<bool> unrollBasedOnType{
       *this, "unroll-based-on-type",
       llvm::cl::desc("Set the unroll factor based on type of the operation"),
       llvm::cl::init(false)};
 };
 
 struct TestVectorDistributePatterns
     : public PassWrapper<TestVectorDistributePatterns,

 };
 
 struct TestVectorTransferUnrollingPatterns
     : public PassWrapper<TestVectorTransferUnrollingPatterns,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestVectorTransferUnrollingPatterns)
 
+  TestVectorTransferUnrollingPatterns() = default;
+  TestVectorTransferUnrollingPatterns(
+      const TestVectorTransferUnrollingPatterns &pass)
+      : PassWrapper(pass) {}
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<AffineDialect>();
   }
   StringRef getArgument() const final {
     return "test-vector-transfer-unrolling-patterns";
   }
   StringRef getDescription() const final {
     return "Test lowering patterns to unroll transfer ops in the vector "
            "dialect";
   }
   void runOnOperation() override {
     MLIRContext *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    populateVectorUnrollPatterns(
-        patterns,
-        UnrollVectorOptions()
-            .setNativeShape(ArrayRef<int64_t>{2, 2})
-            .setFilterConstraint([](Operation *op) {
-              return success(
-                  isa<vector::TransferReadOp, vector::TransferWriteOp>(op));
-            }));
+    UnrollVectorOptions opts;
+    opts.setNativeShape(ArrayRef<int64_t>{2, 2})
+        });
+    if (reverseUnrollOrder.getValue()) {
+      opts.setUnrollTraversalOrderFn(
+          [](Operation *op) -> Optional<SmallVector<int64_t>> {
+            int64_t numLoops = 0;
+            if (auto readOp = dyn_cast<vector::TransferReadOp>(op))
+              numLoops = readOp.getVectorType().getRank();
+            else if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op))
+              numLoops = writeOp.getVectorType().getRank();
+            else
+              return None;
+            auto order = llvm::reverse(llvm::seq<int64_t>(0, numLoops));
+            return llvm::to_vector(order);
+          });
+    }
+    populateVectorUnrollPatterns(patterns, opts);
     populateVectorToVectorCanonicalizationPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
+
+  Option<bool> reverseUnrollOrder{
+      *this, "reverse-unroll-order",
+      llvm::cl::desc(
+          "reverse the order of unrolling of vector transfer operations"),
+      llvm::cl::init(false)};
 };
 
 struct TestVectorTransferFullPartialSplitPatterns
     : public PassWrapper<TestVectorTransferFullPartialSplitPatterns,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestVectorTransferFullPartialSplitPatterns)