Index: mlir/include/mlir/Dialect/Vector/VectorTransforms.h
===================================================================
--- mlir/include/mlir/Dialect/Vector/VectorTransforms.h
+++ mlir/include/mlir/Dialect/Vector/VectorTransforms.h
@@ -117,21 +117,20 @@
 };
 /// Pattern to apply `unrollSingleResultVectorOp` to a `targetShape`
 /// declaratively.
-template <typename OpTy>
-struct UnrollVectorPattern : public OpRewritePattern<OpTy> {
-  using FilterConstraintType = std::function<LogicalResult(OpTy op)>;
+struct UnrollVectorPattern : public RewritePattern {
+  using FilterConstraintType = std::function<LogicalResult(Operation *op)>;
   UnrollVectorPattern(MLIRContext *context, UnrollVectorOptions options)
-      : OpRewritePattern<OpTy>(context), options(options) {}
-  LogicalResult matchAndRewrite(OpTy op,
+      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()), options(options) {}
+  LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     if (options.filterConstraint && failed(options.filterConstraint(op)))
       return failure();
     if (!options.nativeShape) {
-      return op.emitError("vector unrolling expects the native shape or native"
+      return op->emitError("vector unrolling expects the native shape or native"
                           "shape call back function to be set");
     }
     auto unrollableVectorOp =
-        dyn_cast<VectorUnrollOpInterface>(op.getOperation());
+        dyn_cast<VectorUnrollOpInterface>(op);
     if (!unrollableVectorOp)
       return failure();
     auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
@@ -139,12 +138,12 @@
       return failure();
     Optional<SmallVector<int64_t, 4>> targetShape = options.nativeShape(op);
     if (!targetShape)
-      return op.emitError("failed to get target shape for vector unroll");
+      return op->emitError("failed to get target shape for vector unroll");
     auto maybeShapeRatio = shapeRatio(*maybeUnrollShape, *targetShape);
     if (!maybeShapeRatio ||
         llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
       return failure();
-    if (std::is_same<OpTy, TransferWriteOp>::value) {
+    if (isa<TransferWriteOp>(op)) {
       if (failed(unrollTransferWriteOp(rewriter, op, *targetShape)))
         return failure();
       rewriter.eraseOp(op);
Index: mlir/test/lib/Transforms/TestVectorTransforms.cpp
===================================================================
--- mlir/test/lib/Transforms/TestVectorTransforms.cpp
+++ mlir/test/lib/Transforms/TestVectorTransforms.cpp
@@ -27,14 +27,21 @@
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     auto *ctx = &getContext();
-    patterns.insert<UnrollVectorPattern<AddFOp>>(
-        ctx, UnrollVectorOptions().setNativeShape(ArrayRef<int64_t>{2, 2}));
-    patterns.insert<UnrollVectorPattern<vector::ContractionOp>>(
-        ctx, UnrollVectorOptions().setNativeShape(ArrayRef<int64_t>{2, 2, 2}));
+    patterns.insert<UnrollVectorPattern>(
+        ctx, UnrollVectorOptions().setNativeShapeFn(getShape));
     populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
+
+private:
+  static Optional<SmallVector<int64_t, 4>> getShape(Operation *op) {
+    if (isa<AddFOp>(op))
+      return SmallVector<int64_t, 4>(2, 2);
+    if (isa<vector::ContractionOp>(op))
+      return SmallVector<int64_t, 4>(3, 2);
+    return llvm::None;
+  }
 };
 
 struct TestVectorSlicesConversion
@@ -120,8 +127,11 @@
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
-    patterns.insert<UnrollVectorPattern<AddFOp>>(
-        ctx, UnrollVectorOptions().setNativeShape(ArrayRef<int64_t>{2, 2}));
+    patterns.insert<UnrollVectorPattern>(
+        ctx, UnrollVectorOptions()
+                 .setNativeShape(ArrayRef<int64_t>{2, 2})
+                 .setFilterContraint(
+                     [](Operation *op) { return success(isa<AddFOp>(op)); }));
 
     if (unrollBasedOnType) {
       UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
@@ -137,12 +147,19 @@
         }
         return nativeShape;
       };
-      patterns.insert<UnrollVectorPattern<vector::ContractionOp>>(
-          ctx, UnrollVectorOptions().setNativeShapeFn(nativeShapeFn));
+      patterns.insert<UnrollVectorPattern>(
+          ctx, UnrollVectorOptions()
+                   .setNativeShapeFn(nativeShapeFn)
+                   .setFilterContraint([](Operation *op) {
+                     return success(isa<ContractionOp>(op));
+                   }));
     } else {
-      patterns.insert<UnrollVectorPattern<vector::ContractionOp>>(
-          ctx,
-          UnrollVectorOptions().setNativeShape(ArrayRef<int64_t>{2, 2, 2}));
+      patterns.insert<UnrollVectorPattern>(
+          ctx, UnrollVectorOptions()
+                   .setNativeShape(ArrayRef<int64_t>{2, 2, 2})
+                   .setFilterContraint([](Operation *op) {
+                     return success(isa<ContractionOp>(op));
+                   }));
     }
     populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
@@ -273,10 +290,14 @@
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
-    patterns.insert<UnrollVectorPattern<vector::TransferReadOp>>(
-        ctx, UnrollVectorOptions().setNativeShape(ArrayRef<int64_t>{2, 2}));
-    patterns.insert<UnrollVectorPattern<vector::TransferWriteOp>>(
-        ctx, UnrollVectorOptions().setNativeShape(ArrayRef<int64_t>{2, 2}));
+    patterns.insert<UnrollVectorPattern>(
+        ctx,
+        UnrollVectorOptions()
+            .setNativeShape(ArrayRef<int64_t>{2, 2})
+            .setFilterContraint([](Operation *op) {
+              return success(
+                  isa<vector::TransferReadOp, vector::TransferWriteOp>(op));
+            }));
     populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));