[mlir][vector] Expose the VectorUnroll transform as a function and a Transform IR op.

[christopherbate]

Summary

This change is NFC with regards to the VectorUnroll transform patterns.
Previously, the vector unrolling transform could only be applied by running a
set of rewrite patterns, and parameters such as the target shape need to be
provided via a callback. This can make it awkward to write transformations that
could compose with unrolling directly.

This change exposes the unrolling transform as a function, vector::unroll,
that accepts any VectorUnrollOpInterface op and returns (1) a set of operations
corresponding to the unrolled target and (2) the replacement value. The various
different patterns in the VectorUnroll transform are replaced by a single
rewrite pattern operating on the VectorUnrollOpInterface.

A transform operation transform.vector.unroll is added.

Depends on D150000.

[nicolasvasilache]

Nice improvements, thank you @christopherbate !

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

31

Can we return OpFoldResults here and let the called materialize with getValueOrCreateConstantIndexOp ?
The need for Values is because the vector transfer operations do not take the mixed form.

35

while we are at it, can we reorder parameters to be builder, loc, rest ?
this will improve consistency to match standalone transform APIs.

513

TypeSwitch should be nicer here.

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

: Op<Transform_Dialect, opname,
                       TransformWithPatternsOpTrait,
                       MemoryEffectsOpInterface,
                       FunctionalStyleTransformOpTrait], traits)> {
   let extraClassDeclaration = [{
     void populatePatterns(RewritePatternSet &patterns);
   }];
 }
 
-def ApplyRankReducingSubviewPatternsOp : 
+def ApplyRankReducingSubviewPatternsOp :
   TransformWithPatternsOp<"vector.apply_rank_reducing_subview_patterns"> {
   let description = [{
     Apply opt-in vector transfer permutation patterns that include:
       - TransferReadDropUnitDimsPattern
       - TransferWriteDropUnitDimsPattern
-    
-    These patterns have the effect of rewriting a vector.transfer with unit 
+    These patterns have the effect of rewriting a vector.transfer with unit
     dimensions into a rank-reduced version thanks to subview operations.
     This is complemented by shape_cast folding patterns.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$results);
 
   let assemblyFormat = [{
     $target
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
-def ApplyTransferPermutationPatternsOp : 
+def ApplyTransferPermutationPatternsOp :
   TransformWithPatternsOp<"vector.apply_transfer_permutation_patterns"> {
   let description = [{
     Apply opt-in vector transfer permutation patterns that include:
       - TransferReadPermutationLowering
       - TransferWritePermutationLowering
       - TransferOpReduceRank
       - TransferWriteNonPermutationLowering
-    
-    These patterns have the effect of rewriting a vector.transfer with an 
+    These patterns have the effect of rewriting a vector.transfer with an
     arbitrary permutation_map to a vector.transfer with a permutation_map that is
     a minor identity followed by a vector.transpose.
 
     In other words, this makes the vector.transfer contiguous on the most minor
     dimensions and materializes the permutation_map as a vector.transpose.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
let assemblyFormat = [{
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 // TODO: evolve lowering_strategy to proper enums.
 def LowerContractionOp : TransformWithPatternsOp<"vector.lower_contraction"> {
   let description = [{
-    Indicates that the vector contraction-like operations nested under the 
+    Indicates that the vector contraction-like operations nested under the
     isolated from above op `target` should be lowered to finer-grained vector
     primitives.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
let assemblyFormat = [{
     $target
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 def MaterializeMasksOp : TransformWithPatternsOp<"vector.materialize_masks"> {
   let description = [{
-    Indicates that mask operations nested under the isolated from above op 
+    Indicates that mask operations nested under the isolated from above op
     `target` should be lowered to fine-grained arithemtic operations.
 
     This is usually the last step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$results);
 
   let assemblyFormat = [{
     $target
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 // TODO: evolve lowering_strategy to proper enums.
 def LowerMultiReductionOp
     : TransformWithPatternsOp<"vector.lower_multi_reduction"> {
   let description = [{
-    Indicates that the vector multi_reduction-like operations nested under the 
+    Indicates that the vector multi_reduction-like operations nested under the
     isolated from above op `target` should be lowered to finer-grained vector
     primitives.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
let assemblyFormat = [{
     $target
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 def LowerShapeCastOp : TransformWithPatternsOp<"vector.lower_shape_cast"> {
   let description = [{
-    Indicates that the vector shape_cast operations nested under the 
+    Indicates that the vector shape_cast operations nested under the
     isolated from above op `target` should be lowered to finer-grained vector
     primitives.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$results);
 
   let assemblyFormat = [{
     $target
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 def LowerTransferOp : TransformWithPatternsOp<"vector.lower_transfer"> {
   let description = [{
-    Indicates that the vector transfer operations nested under the 
+    Indicates that the vector transfer operations nested under the
     isolated from above op `target` should be lowered to finer-grained vector
     primitives.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
      DefaultValuedAttr<I64Attr, "1">:$max_transfer_rank
   );
   let results = (outs TransformHandleTypeInterface:$results);
 
   let assemblyFormat = [{
     $target
     (`max_transfer_rank` `=` $max_transfer_rank^)?
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 // TODO: evolve lowering_strategy to proper enums.
 def LowerTransposeOp : TransformWithPatternsOp<"vector.lower_transpose"> {
   let description = [{
-    Indicates that the vector transpose-like operations nested under the 
+    Indicates that the vector transpose-like operations nested under the
     isolated from above op `target` should be lowered to finer-grained vector
     primitives.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
let assemblyFormat = [{
     `:` functional-type($target, results)
   }];
 }
 
 // TODO: evolve split_transfer_strategy to proper enums.
 def SplitTransferFullPartialOp
     : TransformWithPatternsOp<"vector.split_transfer_full_partial"> {
   let description = [{
-    Indicates that the vector transfer operations nested under the 
+    Indicates that the vector transfer operations nested under the
     isolated from above op `target` should be split to full and partial parts.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
      DefaultValuedAttr<VectorTransferSplitAttr,
        "vector::VectorTransferSplit::LinalgCopy">:$split_transfer_strategy
   );
   let results = (outs TransformHandleTypeInterface:$results);
 
   let assemblyFormat = [{
     $target
     (`split_transfer_strategy` `=` $split_transfer_strategy^)?
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
 def TransferToScfOp : TransformWithPatternsOp<"vector.transfer_to_scf"> {
   let description = [{
-    Indicates that the vector transfer operations nested under the 
+    Indicates that the vector transfer operations nested under the
     isolated from above op `target` should be rewritten with scf.for loops over
     finer-grained vector primitives.
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
      DefaultValuedAttr<I64Attr, "1">:$max_transfer_rank,
      DefaultValuedAttr<BoolAttr, "false">:$full_unroll
   );
   let results = (outs TransformHandleTypeInterface:$results);
 
   let assemblyFormat = [{
     $target
     oilist (
         `max_transfer_rank` `=` $max_transfer_rank
       | `full_unroll` `=` $full_unroll
     )
     attr-dict
     `:` functional-type($target, results)
   }];
 }
 
+def UnrollOp : Op<Transform_Dialect, "vector.unroll",
+         [TransformOpInterface,
+          TransformEachOpTrait,
+          MemoryEffectsOpInterface,
+          FunctionalStyleTransformOpTrait]> {
+  let description = [{
+    Unrolls the `target` using the provided `unroll_shape` and `unroll_order`.
+    The `unroll_shape` size must be less-than-or-equal-to the size of the
+    target's canonical shape provided by `target.getUnrollShape()`. The
+    `unroll_order` controls the order in which the unrolling occurs by
+    specifying a permutation on the dimensions (in the order of slowest to
+    fastest varying).
+  }];
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       DenseI64ArrayAttr:$unroll_shape,
+                       OptionalAttr<DenseI64ArrayAttr>:$unroll_order);
+  let results = (outs );
+  let assemblyFormat = "$target attr-dict `:` type($target)";
+  let extraClassDeclaration = [{
+    DiagnosedSilenceableFailure applyToOne(
+      VectorUnrollOpInterface target,
+      transform::ApplyToEachResultList &transformResults,
+      transform::TransformState &state
+    );
+  }];
+}
+
 #endif // VECTOR_TRANSFORM_OPS

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

 //===- VectorTransforms.h - Vector transformations as patterns --*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORTRANSFORMS_H
 #define MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORTRANSFORMS_H
 
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Interfaces/VectorInterfaces.h"
 
 namespace mlir {
 class MLIRContext;
 class VectorTransferOpInterface;
 class RewritePatternSet;
 class RewriterBase;
 
 namespace scf {
LogicalResult splitFullAndPartialTransfer(
     RewriterBase &b, VectorTransferOpInterface xferOp,
     VectorTransformsOptions options = VectorTransformsOptions(),
     scf::IfOp *ifOp = nullptr);
 
 /// Implements transfer op write to read forwarding and dead transfer write
 /// optimizations.
 void transferOpflowOpt(RewriterBase &rewriter, Operation *rootOp);
 
+/// Encapsulates the result of an unrolling transformation. The `unrolledOps`
+/// are of the unrolled ops of the same operation type as the target of the
+/// unrolling function. The `replacements` is what should be used to replace the
+/// result of the unrolled operation.a
+struct UnrollResult {
+  SmallVector<Operation *> unrolledOps;
+  Value replacement;
+  UnrollResult(ArrayRef<Operation *> unrolledOps, Value replacement)
+      : unrolledOps(unrolledOps), replacement(replacement) {}
+};
+
+/// Unrolls the `target` using the provided target shape and unroll order. The
+/// `targetShape` size must be less-than-or-equal-to the size of the target's
+/// canonical shape provided by `target.getUnrollShape()`. The `unrollOrder`
+/// controls the order in which the unrolling occurs by specifying a permutation
+/// on the dimensions (in the order of slowest to fastest varying).
+FailureOr<UnrollResult> unroll(RewriterBase &rewriter,
+                               VectorUnrollOpInterface target,
+                               ArrayRef<int64_t> targetShape,
+                               ArrayRef<int64_t> unrollOrder);
+
 } // namespace vector
 } // namespace mlir
 
 #endif // MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORTRANSFORMS_H

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Dialect/X86Vector/Transforms.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/VectorInterfaces.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::vector;
 using namespace mlir::transform;
 
 //===----------------------------------------------------------------------===//
 // ApplyRankReducingSubviewPatternsOp
void transform::LowerMaskedTransfersOp::populatePatterns(
     RewritePatternSet &patterns) {
   populateVectorMaskLoweringPatternsForSideEffectingOps(patterns);
 }
 
 //===----------------------------------------------------------------------===//
 // MaterializeMasksOp
 //===----------------------------------------------------------------------===//
 
-void transform::MaterializeMasksOp::populatePatterns(RewritePatternSet &patterns) {
+void transform::MaterializeMasksOp::populatePatterns(
+    RewritePatternSet &patterns) {
   populateVectorMaskMaterializationPatterns(patterns,
                                             /*force32BitVectorIndices=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // LowerMultiReductionOp
 //===----------------------------------------------------------------------===//
 
void transform::TransferToScfOp::populatePatterns(RewritePatternSet &patterns) {
   VectorTransferToSCFOptions vectorTransferToSCFOptions =
       VectorTransferToSCFOptions()
           .enableFullUnroll(getFullUnroll())
           .setTargetRank(getMaxTransferRank());
   populateVectorToSCFConversionPatterns(patterns, vectorTransferToSCFOptions);
 }
 
 //===----------------------------------------------------------------------===//
+// UnrollOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::UnrollOp::applyToOne(
+    VectorUnrollOpInterface target,
+    transform::ApplyToEachResultList &transformResults,
+    transform::TransformState &state) {
+  IRRewriter rewriter(target->getContext());
+  rewriter.setInsertionPoint(target);
+  SmallVector<int64_t> unrollOrder =
+      getUnrollOrder()
+          ? llvm::to_vector(*getUnrollOrder())
+          : llvm::to_vector(llvm::seq<int64_t>(0, getUnrollShape().size()));
+  FailureOr<UnrollResult> result =
+      vector::unroll(rewriter, target, getUnrollShape(), unrollOrder);
+  if (failed(result)) {
+    target->emitOpError("failed to apply unrolling transformation");
+    return emitDefaultDefiniteFailure(target);
+  }
+  if (target->getNumResults() == 0)
+    rewriter.eraseOp(target);
+  else
+    rewriter.replaceOp(target, result->replacement);
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the additional
 /// ops are using PDL types for operands and results.
 class VectorTransformDialectExtension
     : public transform::TransformDialectExtension<

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

 //===- VectorUnrollDistribute.cpp - patterns to do vector unrolling -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements patterns to do vector unrolling and vector distribution.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/Utils/Utils.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/ADT/STLExtras.h"
 #include <numeric>
 #include <optional>
 
 #define DEBUG_TYPE "vector-unrolling"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 /// Compute the indices of the slice `index` for a tranfer op.
 static SmallVector<Value> sliceTransferIndices(ArrayRef<int64_t> elementOffsets,
                                                ArrayRef<Value> indices,
                                                AffineMap permutationMap,
                                                Location loc,
-                                               OpBuilder &builder) {
+                                               RewriterBase &builder) {
   MLIRContext *ctx = builder.getContext();
   auto isBroadcast = [](AffineExpr expr) {
     if (auto constExpr = expr.dyn_cast<AffineConstantExpr>())
       return constExpr.getValue() == 0;
     return false;
   };
   // Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'.
-  SmallVector<Value> slicedIndices(indices.begin(), indices.end());
+  SmallVector<OpFoldResult> 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);
     auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
-    slicedIndices[pos] =
-        builder.create<affine::AffineApplyOp>(loc, map, indices[pos]);
+    slicedIndices[pos] = affine::makeComposedFoldedAffineApply(
+        builder, loc, map, OpFoldResult(indices[pos]));
   }
-  return slicedIndices;
+  return getValueOrCreateConstantIndexOp(builder, loc, slicedIndices);
 }
 
 // Clones `op` into a new operations that takes `operands` and returns
 // `resultTypes`.
-static Operation *cloneOpWithOperandsAndTypes(OpBuilder &builder, Location loc,
-                                              Operation *op,
+static Operation *cloneOpWithOperandsAndTypes(RewriterBase &builder,
+                                              Location loc, Operation *op,
                                               ArrayRef<Value> operands,
                                               ArrayRef<Type> resultTypes) {
   return builder.create(loc, op->getName().getIdentifier(), operands,
                         resultTypes, op->getAttrs());
 }
 
+/// Returns true if `op` implements VectorUnrollOpInterface and  `targetShape`
+/// is valid for unrolling `op`. If `trivialUnrollValid` is true, then an unroll
+/// that will result in just one operation (for example, if the shapes are
+/// equal) is considered valid, otherwise it is invalid.
+static bool isValidTargetUnrollShape(ArrayRef<int64_t> targetShape,
+                                     Operation *op,
+                                     bool trivialUnrollValid = true) {
+  auto unrollableVectorOp = dyn_cast<VectorUnrollOpInterface>(op);
+  if (!unrollableVectorOp)
+    return false;
+  auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
+  if (!maybeUnrollShape)
+    return false;
+  auto maybeShapeRatio = computeShapeRatio(*maybeUnrollShape, targetShape);
+  if (!maybeShapeRatio)
+    return false;
+  // Reject the trivial case if required. This is required in pattern rewrites
+  // where we want to abort when the unrolling is identity transform.
+  if (!trivialUnrollValid &&
+      llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
+    return false;
+  return true;
+}
+
 /// Return the target shape for unrolling for the given `op`. Return
 /// std::nullopt if the op shouldn't be or cannot be unrolled.
 static std::optional<SmallVector<int64_t>>
-getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
+getTargetShapeFromOptions(const vector::UnrollVectorOptions &options,
+                          Operation *op) {
   if (options.filterConstraint && failed(options.filterConstraint(op)))
     return std::nullopt;
   assert(options.nativeShape &&
          "vector unrolling expects the native shape or native"
          "shape call back function to be set");
-  auto unrollableVectorOp = dyn_cast<VectorUnrollOpInterface>(op);
-  if (!unrollableVectorOp)
-    return std::nullopt;
-  auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
-  if (!maybeUnrollShape)
-    return std::nullopt;
   std::optional<SmallVector<int64_t>> targetShape = options.nativeShape(op);
-  if (!targetShape)
-    return std::nullopt;
-  auto maybeShapeRatio = computeShapeRatio(*maybeUnrollShape, *targetShape);
-  if (!maybeShapeRatio ||
-      llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
+  if (!targetShape ||
+      !isValidTargetUnrollShape(*targetShape, op, /*equalShapeValid=*/false))
     return std::nullopt;
   return targetShape;
 }
 
+/// Return the unroll order that should be used for `op`, or if no unroll order
+/// is specified via the callback, use the identity order.
 static SmallVector<int64_t>
-getUnrollOrder(unsigned numLoops, Operation *op,
+getUnrollOrderFromOptions(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) {
     std::optional<SmallVector<int64_t>> order =
         options.traversalOrderCallback(op);
-    if (order) {
-      loopOrder = std::move(*order);
-    }
+    if (order)
+      return *order;
   }
-  return loopOrder;
-}
-
-namespace {
-
-struct UnrollTransferReadPattern
-    : public OpRewritePattern<vector::TransferReadOp> {
-  UnrollTransferReadPattern(MLIRContext *context,
-                            const vector::UnrollVectorOptions &options,
-                            PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::TransferReadOp>(context, benefit),
-        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> strides(targetShape->size(), 1);
-    Location loc = readOp.getLoc();
-    ArrayRef<int64_t> originalSize = readOp.getVectorType().getShape();
+  auto unrollShape = dyn_cast<VectorUnrollOpInterface>(op).getShapeForUnroll();
+  assert(unrollShape && "expected valid unroll shape");
+  return llvm::to_vector(
+      llvm::seq<int64_t>(0, static_cast<int64_t>(unrollShape->size())));
+}
+
+/// Unroll a `vector.transfer_read`.
+static FailureOr<UnrollResult> unrollTransferRead(RewriterBase &rewriter,
+                                                  TransferReadOp readOp,
+                                                  ArrayRef<int64_t> targetShape,
+                                                  ArrayRef<int64_t> loopOrder) {
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(readOp);
+  if (!isValidTargetUnrollShape(targetShape, readOp))
+    return failure();
+  if (readOp.getTransferRank() == 0 && targetShape.empty())
+    return UnrollResult({readOp}, Value(readOp));
+  if (readOp.getTransferRank() == 0 || readOp.getMask())
+  SmallVector<int64_t> strides(targetShape.size(), 1);
 
-    // Prepare the result vector;
-    Value result = rewriter.create<arith::ConstantOp>(
-        loc, sourceVectorType, rewriter.getZeroAttr(sourceVectorType));
-    auto targetType =
-        VectorType::get(*targetShape, sourceVectorType.getElementType());
-    SmallVector<Value> originalIndices(readOp.getIndices().begin(),
-                                       readOp.getIndices().end());
-    SmallVector<int64_t> loopOrder =
-        getUnrollOrder(originalSize.size(), readOp, options);
-    for (SmallVector<int64_t> elementOffsets :
-         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
-      SmallVector<Value> indices =
-          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());
-
-      result = rewriter.create<vector::InsertStridedSliceOp>(
-          loc, slicedRead, result, elementOffsets, strides);
-    }
-    rewriter.replaceOp(readOp, result);
-    return success();
-  }
-
-private:
-  vector::UnrollVectorOptions options;
-};
-
-struct UnrollTransferWritePattern
-    : public OpRewritePattern<vector::TransferWriteOp> {
-  UnrollTransferWritePattern(MLIRContext *context,
-                             const vector::UnrollVectorOptions &options,
-                             PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::TransferWriteOp>(context, benefit),
-        options(options) {}
-
-  LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
-                                PatternRewriter &rewriter) const override {
-    // TODO: support 0-d corner case.
-    if (writeOp.getTransferRank() == 0)
-      return failure();
-
-    if (writeOp.getMask())
-      return failure();
-    auto targetShape = getTargetShape(options, writeOp);
-    if (!targetShape)
-      return failure();
-    auto sourceVectorType = writeOp.getVectorType();
-    SmallVector<int64_t> strides(targetShape->size(), 1);
-    Location loc = writeOp.getLoc();
-    ArrayRef<int64_t> originalSize = sourceVectorType.getShape();
-    SmallVector<Value> originalIndices(writeOp.getIndices().begin(),
-                                       writeOp.getIndices().end());
-    SmallVector<int64_t> loopOrder =
-        getUnrollOrder(originalSize.size(), writeOp, options);
-    Value resultTensor;
-    for (SmallVector<int64_t> elementOffsets :
-         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
-      Value slicedVector = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, writeOp.getVector(), elementOffsets, *targetShape, strides);
-      SmallVector<Value> indices =
-          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);
-    }
-    if (resultTensor)
-      rewriter.replaceOp(writeOp, resultTensor);
-    else
-      rewriter.eraseOp(writeOp);
-    return success();
-  }
+      VectorType::get(targetShape, sourceVectorType.getElementType());
+  SmallVector<Operation *> transferReadOps;
+       StaticTileOffsetRange(originalSize, targetShape, loopOrder)) {
+    result = rewriter.create<InsertStridedSliceOp>(loc, slicedRead, result,
+                                                   elementOffsets, strides);
+    transferReadOps.push_back(slicedRead);
+  return UnrollResult(transferReadOps, result);
+}
+
+/// Unroll a `vector.transfer_write` op.
+static FailureOr<UnrollResult>
+unrollTransferWrite(RewriterBase &rewriter, TransferWriteOp writeOp,
+                    ArrayRef<int64_t> targetShape,
+                    ArrayRef<int64_t> loopOrder) {
+  if (!isValidTargetUnrollShape(targetShape, writeOp))
+    return failure();
+  /// We are already unrolled.
+  if (writeOp.getTransferRank() == 0 && targetShape.empty())
+    return UnrollResult({writeOp}, Value(writeOp->getResult(0)));
+  SmallVector<int64_t> strides(targetShape.size(), 1);
+  SmallVector<Operation *> writeOps;
+       StaticTileOffsetRange(originalSize, targetShape, loopOrder)) {
+    Value slicedVector = rewriter.create<ExtractStridedSliceOp>(
+        loc, writeOp.getVector(), elementOffsets, targetShape, strides);
+    auto slicedWrite = rewriter.create<TransferWriteOp>(
+    writeOps.push_back(slicedWrite);
+  return UnrollResult(writeOps, resultTensor);
 
-private:
-  vector::UnrollVectorOptions options;
-};
-
+namespace {
 struct OffsetMapInfo {
   static SmallVector<int64_t> getEmptyKey() { return {int64_t(-1)}; }
 
   static SmallVector<int64_t> getTombstoneKey() { return {int64_t(-2)}; }
 
   static unsigned getHashValue(const SmallVector<int64_t> &v) {
     return static_cast<unsigned>(llvm::hash_combine_range(v.begin(), v.end()));
   }
 
   static bool isEqual(const SmallVector<int64_t> &lhs,
                       const SmallVector<int64_t> &rhs) {
     return lhs == rhs;
   }
 };
+} // namespace
 
-struct UnrollContractionPattern
-    : public OpRewritePattern<vector::ContractionOp> {
-  UnrollContractionPattern(MLIRContext *context,
-                           const vector::UnrollVectorOptions &options,
-                           PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::ContractionOp>(context, benefit),
-        options(options) {}
-
-  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> originalSize = *contractOp.getShapeForUnroll();
+/// Unroll a `vector.contract` operation.
+static FailureOr<UnrollResult> unrollContraction(RewriterBase &rewriter,
+                                                 ContractionOp contractOp,
+                                                 ArrayRef<int64_t> targetShape,
+                                                 ArrayRef<int64_t> loopOrder) {
 
-    Location loc = contractOp.getLoc();
-    unsigned accIndex = vector::ContractionOp::getAccOperandIndex();
-    AffineMap dstAffineMap = contractOp.getIndexingMapsArray()[accIndex];
-    llvm::MapVector<
-        SmallVector<int64_t>, Value,
+  llvm::MapVector<SmallVector<int64_t>, Value,
                   llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
-        accCache;
-
-    SmallVector<int64_t> loopOrder = getUnrollOrder(
-        contractOp.getIteratorTypes().size(), contractOp, options);
-
-    for (SmallVector<int64_t> offsets :
-         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
-      SmallVector<Value> slicesOperands(contractOp.getNumOperands());
+  SmallVector<Operation *> newContractOps;
+       StaticTileOffsetRange(originalSize, targetShape, loopOrder)) {
 
-      // Helper to compute 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));
-        SmallVector<int64_t> operandStrides(operandOffets.size(), 1);
-        slicesOperands[index] = rewriter.create<vector::ExtractStridedSliceOp>(
-            loc, operand, operandOffets, operandShape, operandStrides);
-      };
+      SmallVector<int64_t> operandShape =
+          applyPermutationMap(permutationMap, ArrayRef<int64_t>(targetShape));
 
-      // Extract the new lhs operand.
-      AffineMap lhsPermutationMap = contractOp.getIndexingMapsArray()[0];
-      SmallVector<int64_t> lhsOffets =
-          applyPermutationMap(lhsPermutationMap, ArrayRef<int64_t>(offsets));
-      extractOperand(0, contractOp.getLhs(), lhsPermutationMap, lhsOffets);
 
-      // Extract the new rhs operand.
-      AffineMap rhsPermutationMap = contractOp.getIndexingMapsArray()[1];
-      SmallVector<int64_t> rhsOffets =
-          applyPermutationMap(rhsPermutationMap, ArrayRef<int64_t>(offsets));
-      extractOperand(1, contractOp.getRhs(), rhsPermutationMap, rhsOffets);
 
-      AffineMap accPermutationMap = contractOp.getIndexingMapsArray()[2];
-      SmallVector<int64_t> accOffets =
-          applyPermutationMap(accPermutationMap, ArrayRef<int64_t>(offsets));
-      // If a version of the accumulator has already been computed, use it
-      // otherwise extract the first version from the original operand.
-      auto accIt = accCache.find(accOffets);
-      if (accIt != accCache.end())
-        slicesOperands[2] = accIt->second;
-      else
-        extractOperand(2, contractOp.getAcc(), accPermutationMap, accOffets);
 
-      SmallVector<int64_t> dstShape =
-          applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(*targetShape));
-      auto targetType = VectorType::get(dstShape, dstVecType.getElementType());
-      Operation *newOp = cloneOpWithOperandsAndTypes(
-          rewriter, loc, contractOp, slicesOperands, targetType);
+        applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(targetShape));
+    Operation *newOp = cloneOpWithOperandsAndTypes(rewriter, loc, contractOp,
+                                                   slicesOperands, targetType);
+    newContractOps.push_back(newOp);
 
-      SmallVector<int64_t> dstOffets =
-          applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(offsets));
-      // Save the accumulated value untill all the loops are unrolled since
-      // reduction loop keep updating the accumulator.
-      accCache[dstOffets] = newOp->getResult(0);
-    }
-    // Assemble back the accumulator into a single vector.
-    Value result = rewriter.create<arith::ConstantOp>(
-        loc, dstVecType, rewriter.getZeroAttr(dstVecType));
-    for (const auto &it : accCache) {
-      SmallVector<int64_t> dstStrides(it.first.size(), 1);
-      result = rewriter.create<vector::InsertStridedSliceOp>(
-          loc, it.second, result, it.first, dstStrides);
-    }
-    rewriter.replaceOp(contractOp, result);
-    return success();
-  }
+  return UnrollResult(newContractOps, result);
 
-private:
-  vector::UnrollVectorOptions options;
-};
-
-struct UnrollMultiReductionPattern
-    : public OpRewritePattern<vector::MultiDimReductionOp> {
-  UnrollMultiReductionPattern(MLIRContext *context,
-                              const vector::UnrollVectorOptions &options,
-                              PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::MultiDimReductionOp>(context, benefit),
-        options(options) {}
-
-  LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
-                                PatternRewriter &rewriter) const override {
-    std::optional<SmallVector<int64_t>> targetShape =
-        getTargetShape(options, reductionOp);
-    if (!targetShape)
-      return failure();
-    SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
-    llvm::MapVector<
-        SmallVector<int64_t>, Value,
+/// Unroll a `vector.multi_dim_reduction`.
+static FailureOr<UnrollResult>
+unrollMultiDimReduce(RewriterBase &rewriter, MultiDimReductionOp reductionOp,
+                     ArrayRef<int64_t> targetShape,
+                     ArrayRef<int64_t> unrollOrder) {
+  if (!isValidTargetUnrollShape(targetShape, reductionOp))
+  llvm::MapVector<SmallVector<int64_t>, Value,
                   llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
-        accCache;
-    Location loc = reductionOp.getLoc();
 
-    // Stride of the ratios, this gives us the offsets of sliceCount in a basis
-    // of multiples of the targetShape.
-    for (SmallVector<int64_t> offsets :
-         StaticTileOffsetRange(originalSize, *targetShape)) {
-      SmallVector<Value> operands;
-      SmallVector<int64_t> operandStrides(offsets.size(), 1);
-      Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, reductionOp.getSource(), offsets, *targetShape, operandStrides);
-      operands.push_back(slicedOperand);
-      SmallVector<int64_t> dstShape;
-      SmallVector<int64_t> destOffset;
-      for (size_t i : llvm::seq(size_t(0), targetShape->size())) {
-        if (!reductionOp.isReducedDim(i)) {
-          destOffset.push_back(offsets[i]);
-          dstShape.push_back((*targetShape)[i]);
-        }
-      }
-      Value acc;
-      SmallVector<int64_t> accStrides(destOffset.size(), 1);
-      // If a version of the accumulator has already been computed, use it
-      // otherwise extract the first version from the original operand.
-      auto accIt = accCache.find(destOffset);
-      if (accIt != accCache.end())
-        acc = accIt->second;
-      else
-        acc = rewriter.create<vector::ExtractStridedSliceOp>(
-            loc, reductionOp.getAcc(), destOffset, dstShape, accStrides);
-      operands.push_back(acc);
-      auto targetType = VectorType::get(
-          dstShape, reductionOp.getSourceVectorType().getElementType());
-      Operation *newOp = cloneOpWithOperandsAndTypes(rewriter, loc, reductionOp,
-                                                     operands, targetType);
-      Value result = newOp->getResult(0);
-      accCache[destOffset] = result;
-    }
-    // Assemble back the accumulator into a single vector.
-    Value result = rewriter.create<arith::ConstantOp>(
-        loc, reductionOp.getDestType(),
-        rewriter.getZeroAttr(reductionOp.getDestType()));
-    for (const auto &it : accCache) {
-      SmallVector<int64_t> dstStrides(it.first.size(), 1);
-      result = rewriter.create<vector::InsertStridedSliceOp>(
-          loc, it.second, result, it.first, dstStrides);
-    }
-    rewriter.replaceOp(reductionOp, result);
-    return success();
-  }
+  SmallVector<Operation *> newReduceOps;
+       StaticTileOffsetRange(originalSize, targetShape)) {
+        loc, reductionOp.getSource(), offsets, targetShape, operandStrides);
+    for (size_t i : llvm::seq(size_t(0), targetShape.size())) {
+        dstShape.push_back(targetShape[i]);
+    newReduceOps.push_back(newOp);
+  return UnrollResult(newReduceOps, result);
 
-private:
-  vector::UnrollVectorOptions options;
-};
-
-struct UnrollElementwisePattern : public RewritePattern {
-  UnrollElementwisePattern(MLIRContext *context,
-                           const vector::UnrollVectorOptions &options,
-                           PatternBenefit benefit = 1)
-      : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
-        options(options) {}
-
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
-    if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
-      return failure();
-    auto targetShape = getTargetShape(options, op);
-    if (!targetShape)
-      return failure();
-    auto dstVecType = op->getResult(0).getType().cast<VectorType>();
-    SmallVector<int64_t> originalSize =
-        *cast<VectorUnrollOpInterface>(op).getShapeForUnroll();
-    Location loc = op->getLoc();
-    // Prepare the result vector.
-    Value result = rewriter.create<arith::ConstantOp>(
-        loc, dstVecType, rewriter.getZeroAttr(dstVecType));
-    SmallVector<int64_t> strides(targetShape->size(), 1);
-    VectorType newVecType =
-        VectorType::get(*targetShape, dstVecType.getElementType());
+/// Unroll any op with the ElementwiseMappable trait that is performing some
+/// computation on vector types.
+static FailureOr<UnrollResult>
+unrollElementwise(RewriterBase &rewriter, Operation *op,
+                  ArrayRef<int64_t> targetShape,
+                  ArrayRef<int64_t> unrollOrder) {
+  SmallVector<int64_t> strides(targetShape.size(), 1);
+      VectorType::get(targetShape, dstVecType.getElementType());
 
-    // Create the unrolled computation.
-    for (SmallVector<int64_t> offsets :
-         StaticTileOffsetRange(originalSize, *targetShape)) {
-      SmallVector<Value> extractOperands;
-      for (OpOperand &operand : op->getOpOperands()) {
-        auto vecType = operand.get().getType().template dyn_cast<VectorType>();
-        if (!vecType) {
-          extractOperands.push_back(operand.get());
-          continue;
-        }
-        extractOperands.push_back(
-            rewriter.create<vector::ExtractStridedSliceOp>(
-                loc, operand.get(), offsets, *targetShape, strides));
-      }
-      Operation *newOp = cloneOpWithOperandsAndTypes(
-          rewriter, loc, op, extractOperands, newVecType);
-      result = rewriter.create<vector::InsertStridedSliceOp>(
-          loc, newOp->getResult(0), result, offsets, strides);
-    }
-    rewriter.replaceOp(op, result);
-    return success();
-  }
+  SmallVector<Operation *> newOps;
+       StaticTileOffsetRange(originalSize, targetShape)) {
+      extractOperands.push_back(rewriter.create<vector::ExtractStridedSliceOp>(
+          loc, operand.get(), offsets, targetShape, strides));
+    Operation *newOp = cloneOpWithOperandsAndTypes(rewriter, loc, op,
+                                                   extractOperands, newVecType);
+    newOps.push_back(newOp);
+  return UnrollResult(newOps, result);
 
-private:
-  vector::UnrollVectorOptions options;
-};
-
-struct UnrollReductionPattern : public OpRewritePattern<vector::ReductionOp> {
-  UnrollReductionPattern(MLIRContext *context,
-                         const vector::UnrollVectorOptions &options,
-                         PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::ReductionOp>(context, benefit),
-        options(options) {}
-
-  LogicalResult matchAndRewrite(vector::ReductionOp reductionOp,
-                                PatternRewriter &rewriter) const override {
-    std::optional<SmallVector<int64_t>> targetShape =
-        getTargetShape(options, reductionOp);
-    if (!targetShape)
-      return failure();
-    SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
-
-    // Create unrolled vector reduction.
-    Location loc = reductionOp.getLoc();
-    Value accumulator = nullptr;
-    for (SmallVector<int64_t> offsets :
-         StaticTileOffsetRange(originalSize, *targetShape)) {
-      SmallVector<int64_t> strides(offsets.size(), 1);
-      Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, reductionOp.getVector(), offsets, *targetShape, strides);
-      Operation *newOp = cloneOpWithOperandsAndTypes(
-          rewriter, loc, reductionOp, slicedOperand, reductionOp.getType());
-      Value result = newOp->getResult(0);
+/// Unroll a `vector.reduce` operation.
+static FailureOr<UnrollResult> unrollReduction(RewriterBase &rewriter,
+                                               ReductionOp reductionOp,
+                                               ArrayRef<int64_t> targetShape,
+                                               ArrayRef<int64_t> unrollOrder) {
+  SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
+       StaticTileOffsetRange(originalSize, targetShape)) {
+        loc, reductionOp.getVector(), offsets, targetShape, strides);
 
-      if (!accumulator) {
-        // This is the first reduction.
-        accumulator = result;
-      } else {
-        // On subsequent reduction, combine with the accumulator.
-        accumulator = makeArithReduction(rewriter, loc, reductionOp.getKind(),
-                                         accumulator, result);
-      }
-    }
-
-    rewriter.replaceOp(reductionOp, accumulator);
-    return success();
-  }
+  return UnrollResult({}, accumulator);
 
-private:
-  const vector::UnrollVectorOptions options;
-};
-
-struct UnrollTransposePattern : public OpRewritePattern<vector::TransposeOp> {
-  UnrollTransposePattern(MLIRContext *context,
-                         const vector::UnrollVectorOptions &options,
-                         PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::TransposeOp>(context, benefit),
-        options(options) {}
-
-  LogicalResult matchAndRewrite(vector::TransposeOp transposeOp,
-                                PatternRewriter &rewriter) const override {
-    if (transposeOp.getResultVectorType().getRank() == 0)
-      return failure();
-    auto targetShape = getTargetShape(options, transposeOp);
-    if (!targetShape)
-      return failure();
-    auto originalVectorType = transposeOp.getResultVectorType();
-    SmallVector<int64_t> strides(targetShape->size(), 1);
-    Location loc = transposeOp.getLoc();
-    ArrayRef<int64_t> originalSize = originalVectorType.getShape();
+static FailureOr<UnrollResult> unrollTranspose(RewriterBase &rewriter,
+                                               TransposeOp transposeOp,
+                                               ArrayRef<int64_t> targetShape,
+                                               ArrayRef<int64_t> loopOrder) {
+  SmallVector<int64_t> strides(targetShape.size(), 1);
 
-    // Prepare the result vector;
-    Value result = rewriter.create<arith::ConstantOp>(
-        loc, originalVectorType, rewriter.getZeroAttr(originalVectorType));
-    SmallVector<int64_t> permutation;
-    transposeOp.getTransp(permutation);
 
-    // Unroll the computation.
-    for (SmallVector<int64_t> elementOffsets :
-         StaticTileOffsetRange(originalSize, *targetShape)) {
-      SmallVector<int64_t> permutedOffsets(elementOffsets.size());
-      SmallVector<int64_t> permutedShape(elementOffsets.size());
-      // Compute the source offsets and shape.
-      for (auto indices : llvm::enumerate(permutation)) {
-        permutedOffsets[indices.value()] = elementOffsets[indices.index()];
-        permutedShape[indices.value()] = (*targetShape)[indices.index()];
-      }
-      Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, transposeOp.getVector(), permutedOffsets, permutedShape,
-          strides);
-      Value transposedSlice =
-          rewriter.create<vector::TransposeOp>(loc, slicedOperand, permutation);
-      result = rewriter.create<vector::InsertStridedSliceOp>(
-          loc, transposedSlice, result, elementOffsets, strides);
-    }
-    rewriter.replaceOp(transposeOp, result);
-    return success();
-  }
+  SmallVector<Operation *> newTransposeOps;
+       StaticTileOffsetRange(originalSize, targetShape)) {
+      permutedShape[indices.value()] = targetShape[indices.index()];
+        loc, transposeOp.getVector(), permutedOffsets, permutedShape, strides);
+    auto transposedSlice =
+    newTransposeOps.push_back(transposedSlice);
+  return UnrollResult(newTransposeOps, result);
 
-private:
-  vector::UnrollVectorOptions options;
-};
-
-struct UnrollGatherPattern : public OpRewritePattern<vector::GatherOp> {
-  UnrollGatherPattern(MLIRContext *context,
-                      const vector::UnrollVectorOptions &options,
-                      PatternBenefit benefit = 1)
-      : OpRewritePattern<vector::GatherOp>(context, benefit), options(options) {
-  }
-
-  LogicalResult matchAndRewrite(vector::GatherOp gatherOp,
-                                PatternRewriter &rewriter) const override {
-    VectorType sourceVectorType = gatherOp.getVectorType();
-    if (sourceVectorType.getRank() == 0)
-      return failure();
-    auto targetShape = getTargetShape(options, gatherOp);
-    if (!targetShape)
-      return failure();
-    SmallVector<int64_t> strides(targetShape->size(), 1);
-    Location loc = gatherOp.getLoc();
-    ArrayRef<int64_t> originalSize = gatherOp.getVectorType().getShape();
+/// Unroll a `vector.gather` operation.
+static FailureOr<UnrollResult> unrollGather(RewriterBase &rewriter,
+                                            GatherOp gatherOp,
+                                            ArrayRef<int64_t> targetShape,
+                                            ArrayRef<int64_t> loopOrder) {
+  SmallVector<int64_t> strides(targetShape.size(), 1);
 
-    // Prepare the result vector;
-    Value result = rewriter.create<arith::ConstantOp>(
-        loc, sourceVectorType, rewriter.getZeroAttr(sourceVectorType));
-    auto targetType =
-        VectorType::get(*targetShape, sourceVectorType.getElementType());
-
-    SmallVector<int64_t> loopOrder =
-        getUnrollOrder(originalSize.size(), gatherOp, options);
-    for (SmallVector<int64_t> elementOffsets :
-         StaticTileOffsetRange(originalSize, *targetShape, loopOrder)) {
-      // To get the unrolled gather, extract the same slice based on the
-      // decomposed shape from each of the index, mask, and pass-through
-      // vectors.
-      Value indexSubVec = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, gatherOp.getIndexVec(), elementOffsets, *targetShape, strides);
-      Value maskSubVec = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, gatherOp.getMask(), elementOffsets, *targetShape, strides);
-      Value passThruSubVec = rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, gatherOp.getPassThru(), elementOffsets, *targetShape, strides);
-      auto slicedGather = rewriter.create<vector::GatherOp>(
-          loc, targetType, gatherOp.getBase(), gatherOp.getIndices(),
-          indexSubVec, maskSubVec, passThruSubVec);
+      VectorType::get(targetShape, sourceVectorType.getElementType());
+  SmallVector<Operation *> newGatherOps;
+       StaticTileOffsetRange(originalSize, targetShape, loopOrder)) {
+        loc, gatherOp.getIndexVec(), elementOffsets, targetShape, strides);
+        loc, gatherOp.getMask(), elementOffsets, targetShape, strides);
+        loc, gatherOp.getPassThru(), elementOffsets, targetShape, strides);
+        loc, targetType, gatherOp.getBase(), gatherOp.getIndices(), indexSubVec,
+        maskSubVec, passThruSubVec);
+    newGatherOps.push_back(slicedGather);
 
-      result = rewriter.create<vector::InsertStridedSliceOp>(
-          loc, slicedGather, result, elementOffsets, strides);
-    }
-    rewriter.replaceOp(gatherOp, result);
+  return UnrollResult(newGatherOps, result);
+}
+
+FailureOr<UnrollResult> vector::unroll(RewriterBase &rewriter,
+                                       VectorUnrollOpInterface target,
+                                       ArrayRef<int64_t> targetShape,
+                                       ArrayRef<int64_t> unrollOrder) {
+  FailureOr<UnrollResult> result = failure();
+  if (auto readOp = dyn_cast<vector::TransferReadOp>(*target))
+    result = unrollTransferRead(rewriter, readOp, targetShape, unrollOrder);
+  if (auto writeOp = dyn_cast<vector::TransferWriteOp>(*target))
+    result = unrollTransferWrite(rewriter, writeOp, targetShape, unrollOrder);
+  if (auto contractOp = dyn_cast<vector::ContractionOp>(*target))
+    result = unrollContraction(rewriter, contractOp, targetShape, unrollOrder);
+  if (auto redOp = dyn_cast<vector::MultiDimReductionOp>(*target))
+    result = unrollMultiDimReduce(rewriter, redOp, targetShape, unrollOrder);
+  if (auto reductionOp = dyn_cast<vector::ReductionOp>(*target))
+    result = unrollReduction(rewriter, reductionOp, targetShape, unrollOrder);
+  if (auto transposeOp = dyn_cast<vector::TransposeOp>(*target))
+    result = unrollTranspose(rewriter, transposeOp, targetShape, unrollOrder);
+  if (auto gatherOp = dyn_cast<vector::GatherOp>(*target))
+    return unrollGather(rewriter, gatherOp, targetShape, unrollOrder);
+  if (OpTrait::hasElementwiseMappableTraits(target))
+    result = unrollElementwise(rewriter, target, targetShape, unrollOrder);
+  return result;
+}
+
+namespace {
+struct UnrollVectorRewrite
+    : public OpInterfaceRewritePattern<VectorUnrollOpInterface> {
+  UnrollVectorRewrite(MLIRContext *context,
+                      const vector::UnrollVectorOptions &options,
+                      PatternBenefit benefit = 1)
+      : OpInterfaceRewritePattern<VectorUnrollOpInterface>(context, benefit),
+        options(options) {}
+  LogicalResult matchAndRewrite(VectorUnrollOpInterface op,
+                                PatternRewriter &rewriter) const override {
+    std::optional<SmallVector<int64_t>> targetShape =
+        getTargetShapeFromOptions(options, op);
+    if (!targetShape)
+      return rewriter.notifyMatchFailure(
+          op, "failed to get a valid unroll target shape");
+    SmallVector<int64_t> unrollOrder = getUnrollOrderFromOptions(op, options);
+    FailureOr<UnrollResult> result =
+        unroll(rewriter, op, *targetShape, unrollOrder);
+    if (failed(result))
+      return rewriter.notifyMatchFailure(
+          op, "failed to apply vector unrolling transformation");
+
+    if (op->getNumResults() == 0) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+    assert(op->getNumResults() == 1 && result->replacement &&
+           "expected single-result op and a valid replacement");
+    rewriter.replaceOp(op, result->replacement);
     return success();
   }
 
 private:
   vector::UnrollVectorOptions options;
 };
-
 } // namespace
 
 void mlir::vector::populateVectorUnrollPatterns(
     RewritePatternSet &patterns, const UnrollVectorOptions &options,
     PatternBenefit benefit) {
-  patterns.add<UnrollTransferReadPattern, UnrollTransferWritePattern,
-               UnrollContractionPattern, UnrollElementwisePattern,
-               UnrollReductionPattern, UnrollMultiReductionPattern,
-               UnrollTransposePattern, UnrollGatherPattern>(
-      patterns.getContext(), options, benefit);
+  patterns.add<UnrollVectorRewrite>(patterns.getContext(), options, benefit);
 }

mlir/test/Dialect/Vector/unroll-transform-op.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter -canonicalize -cse --split-input-file | FileCheck %s
+
+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>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["vector.transfer_read"]} in %arg0 : (!pdl.operation) -> !pdl.operation
+  transform.vector.unroll %0 {
+    unroll_shape = array<i64: 2, 2>
+  } : !pdl.operation
+}
+
+// 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>
+// -----
+
+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>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["vector.transfer_read"]} in %arg0 : (!pdl.operation) -> !pdl.operation
+  transform.vector.unroll %0 {
+    unroll_shape = array<i64: 2, 2>,
+    unroll_order = array<i64: 1, 0>
+  } : !pdl.operation
+}
+
+// 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 {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  CHECK-NEXT:   %[[VEC1:.*]] = vector.insert_strided_slice %[[VTR1]], %[[VEC0]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+//  CHECK-NEXT:   %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  CHECK-NEXT:   %[[VEC2:.*]] = vector.insert_strided_slice %[[VTR2]], %[[VEC1]] {offsets = [0, 2], 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>
+
+// -----
+
+
+func.func @transfer_readwrite_unroll(%arg0 : memref<4x4xf32>) {
+  %c0 = arith.constant 0 : index
+  %cf0 = arith.constant 0.0 : f32
+  %0 = vector.transfer_read %arg0[%c0, %c0], %cf0 : memref<4x4xf32>, vector<4x4xf32>
+  vector.transfer_write %0, %arg0[%c0, %c0] : vector<4x4xf32>, memref<4x4xf32>
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["vector.transfer_read"]} in %arg0 : (!pdl.operation) -> !pdl.operation
+  transform.vector.unroll %0 {
+    unroll_shape = array<i64: 2, 2>
+  } : !pdl.operation
+  %1 = transform.structured.match ops{["vector.transfer_write"]} in %arg0 : (!pdl.operation) -> !pdl.operation
+  transform.vector.unroll %1 {
+    unroll_shape = array<i64: 2, 2>
+  } : !pdl.operation
+}
+
+// CHECK-LABEL: func @transfer_readwrite_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:   %[[VTR1:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  CHECK-NEXT:   %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  CHECK-NEXT:   %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
+//  CHECK-NEXT:   vector.transfer_write %[[VTR0]], {{.*}}[%[[C0]], %[[C0]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  CHECK-NEXT:   vector.transfer_write %[[VTR1]], {{.*}}[%[[C0]], %[[C2]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  CHECK-NEXT:   vector.transfer_write %[[VTR2]], {{.*}}[%[[C2]], %[[C0]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  CHECK-NEXT:   vector.transfer_write %[[VTR3]], {{.*}}[%[[C2]], %[[C2]]] {{.*}} : vector<2x2xf32>, memref<4x4xf32>
+//  CHECK-NEXT:   return
+
+// -----
+
+func.func @vector_tranpose(%v : vector<2x4x3x8xf32>) -> vector<2x3x8x4xf32> {
+  %t = vector.transpose %v, [0, 2, 3, 1] : vector<2x4x3x8xf32> to vector<2x3x8x4xf32>
+  return %t : vector<2x3x8x4xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["vector.transpose"]} in %arg0 : (!pdl.operation) -> !pdl.operation
+  transform.vector.unroll %0 {
+    unroll_shape = array<i64: 1, 3, 4, 2>
+  } : !pdl.operation
+}
+
+// CHECK-LABEL: func @vector_tranpose
+//       CHECK:   %[[VI:.*]] = arith.constant dense<0.000000e+00> : vector<2x3x8x4xf32>
+//       CHECK:   %[[E0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 0, 0], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
+//       CHECK:   %[[T0:.*]] = vector.transpose %[[E0]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
+//       CHECK:   %[[V0:.*]] = vector.insert_strided_slice %[[T0]], %[[VI]] {offsets = [0, 0, 0, 0], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
+//       CHECK:   %[[E1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2, 0, 0], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
+//       CHECK:   %[[T1:.*]] = vector.transpose %[[E1]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
+//       CHECK:   %[[V1:.*]] = vector.insert_strided_slice %[[T1]], %[[V0]] {offsets = [0, 0, 0, 2], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
+//       CHECK:   %[[E2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 0, 4], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
+//       CHECK:   %[[T2:.*]] = vector.transpose %[[E2]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
+//       CHECK:   %[[V2:.*]] = vector.insert_strided_slice %[[T2]], %[[V1]] {offsets = [0, 0, 4, 0], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
+//       CHECK:   %[[E3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2, 0, 4], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
+//       CHECK:   %[[T3:.*]] = vector.transpose %[[E3]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
+//       CHECK:   %[[V3:.*]] = vector.insert_strided_slice %[[T3]], %[[V2]] {offsets = [0, 0, 4, 2], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
+//       CHECK:   %[[E4:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 0, 0], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
+//       CHECK:   %[[T4:.*]] = vector.transpose %[[E4]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
+//       CHECK:   %[[V4:.*]] = vector.insert_strided_slice %[[T4]], %[[V3]] {offsets = [1, 0, 0, 0], strides = [1, 1, 1, 1]} : vector<1x3x4x2xf32> into vector<2x3x8x4xf32>
+//       CHECK:   %[[E5:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 2, 0, 0], sizes = [1, 2, 3, 4], strides = [1, 1, 1, 1]} : vector<2x4x3x8xf32> to vector<1x2x3x4xf32>
+//       CHECK:   %[[T5:.*]] = vector.transpose %[[E5]], [0, 2, 3, 1] : vector<1x2x3x4xf32> to vector<1x3x4x2xf32>
+//       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>