[mlir][vector][NFC] Change UnrollVectorPattern to not be statically dependent on an op type

[ThomasRaoux]

Summary

Make UnrollVectorPattern inherit from RewritePattern instead of OpRewritePattern so that we don't need to create many patterns when applying to many different type of ops. Since we may want to apply the pattern to all arithmetic op, it is more convenient to filter dynamically.

[bondhugula]

mlir/test/lib/Transforms/TestVectorTransforms.cpp

38

Doc comment please.

154

This method is wrongly named: Contraint -> Constraint. (I understand this is due to a mix up with its French spelling contrainte. :-))

155

Nit side comment: Not sure why the setFilterConstraint argument function just doesn't return a bool instead of LogicalResult.

[andydavis1]

Thanks!

[ThomasRaoux]

Address review comments.

[ThomasRaoux]

mlir/test/lib/Transforms/TestVectorTransforms.cpp

154

Fixed it :)

155

Not sure either but I'm leaving it for now as I don't know the reason behind it. Should be easy to do a clean up anytime if we want to.

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

LogicalResult unrollTransferWriteOp(OpBuilder &builder, Operation *op,
                                     ArrayRef<int64_t> targetShape);
 
 /// Options that control the vector unrolling.
 struct UnrollVectorOptions {
   using FilterConstraintFnType = std::function<LogicalResult(Operation *op)>;
   /// Callback function that indicates whether vector unrolling should be
   /// attempted on the operation.
   FilterConstraintFnType filterConstraint = nullptr;
-  UnrollVectorOptions &setFilterContraint(FilterConstraintFnType constraint) {
+  UnrollVectorOptions &setFilterConstraint(FilterConstraintFnType constraint) {
     filterConstraint = constraint;
     return *this;
   }
 
   using NativeShapeFnType =
       std::function<Optional<SmallVector<int64_t, 4>>(Operation *op)>;
   /// Function that returns the shape of the vector to unroll to for a given
   /// operation. The unrolling is aborted if the function returns `llvm::None`.
UnrollVectorOptions &setNativeShape(ArrayRef<int64_t> shape) {
     nativeShape = [=](Operation *) -> Optional<SmallVector<int64_t, 4>> {
       return tsShape;
     };
     return *this;
   }
 };
 /// 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());
+    auto unrollableVectorOp = dyn_cast<VectorUnrollOpInterface>(op);
     if (!unrollableVectorOp)
       return failure();
     auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
     if (!maybeUnrollShape)
       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);
       return success();
     }
     if (op->getNumResults() != 1)
       return failure();
     auto resultVector = unrollSingleResultVectorOp(rewriter, op, *targetShape);

mlir/test/lib/Transforms/TestVectorTransforms.cpp

 using namespace mlir::vector;
 namespace {
 
 struct TestVectorToVectorConversion
     : public PassWrapper<TestVectorToVectorConversion, FunctionPass> {
   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:
+  // Return the target shape based on op type.
+  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
     : public PassWrapper<TestVectorSlicesConversion, FunctionPass> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     populateVectorSlicesLoweringPatterns(patterns, &getContext());
     applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));

 
 struct TestVectorUnrollingPatterns
     : public PassWrapper<TestVectorUnrollingPatterns, FunctionPass> {
   TestVectorUnrollingPatterns() = default;
   TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass) {}
   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})
+                 .setFilterConstraint(
+                     [](Operation *op) { return success(isa<AddFOp>(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;
           }
         }
         return nativeShape;
       };
-      patterns.insert<UnrollVectorPattern<vector::ContractionOp>>(
-          ctx, UnrollVectorOptions().setNativeShapeFn(nativeShapeFn));
+      patterns.insert<UnrollVectorPattern>(
+          ctx, UnrollVectorOptions()
+                   .setNativeShapeFn(nativeShapeFn)
+                   .setFilterConstraint([](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})
+                   .setFilterConstraint([](Operation *op) {
+                     return success(isa<ContractionOp>(op));
+                   }));
     }
     populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 
   Option<bool> unrollBasedOnType{
       *this, "unroll-based-on-type",

 struct TestVectorTransferUnrollingPatterns
     : public PassWrapper<TestVectorTransferUnrollingPatterns, FunctionPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<AffineDialect>();
   }
   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})
+            .setFilterConstraint([](Operation *op) {
+              return success(
+                  isa<vector::TransferReadOp, vector::TransferWriteOp>(op));
+            }));
     populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
 struct TestVectorTransferFullPartialSplitPatterns
     : public PassWrapper<TestVectorTransferFullPartialSplitPatterns,