[mlir][Linalg] Add pad vectorization patterns into LinalgStrategyVectorize passes.

[mravishankar]

Summary

Add an option to control whether these patterns are added to the
pattern list or not.

[mravishankar]

Fix build errors.

[mravishankar]

Rebase.

mlir/include/mlir/Dialect/Linalg/Passes.h

createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange = {},
                                         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgStrategyVectorizePass(StringRef opName = "",
                                   linalg::LinalgVectorizationOptions opt =
                                       linalg::LinalgVectorizationOptions(),
                                   linalg::LinalgTransformationFilter filter =
-                                      linalg::LinalgTransformationFilter());
+                                      linalg::LinalgTransformationFilter(),
+                                  bool padVectorize = false);
 
 /// Create a LinalgStrategyEnablePass.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyEnablePass(
     linalg::LinalgEnablingOptions opt = linalg::LinalgEnablingOptions(),
     linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyLowerVectorsPass.

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgStrategyVectorizePass
   let summary = "Configurable pass to apply pattern-based linalg vectorization.";
   let constructor = "mlir::createLinalgStrategyVectorizePass()";
   let dependentDialects = ["linalg::LinalgDialect"];
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,
     Option<"anchorOpName", "anchor-op", "std::string", /*default=*/"",
       "Which linalg op within the func is the anchor to latch on.">,
+    Option<"vectorizePadding", "vectorize-padding", "bool", "false",
+      "Enable vectorization of padding ops.">,
   ];
 }
 
 def LinalgStrategyEnablePass
     : FunctionPass<"linalg-strategy-enable-pass"> {
   let summary = "Configurable pass to enable the application of other "
     "pattern-based linalg passes.";
   let constructor = "mlir::createLinalgStrategyEnablePass()";

mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h

 
 private:
   SmallVector<int64_t> iteratorInterchange;
 };
 
 /// Represent one application of createLinalgStrategyVectorizePass.
 struct Vectorize : public Transformation {
   explicit Vectorize(linalg::LinalgVectorizationOptions options,
-                     LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(f), opName(), options(options) {}
+                     LinalgTransformationFilter::FilterFunction f = nullptr,
+                     bool padVectorize = false)
+      : Transformation(f), opName(), options(options),
+        vectorizePadding(padVectorize) {}
 
   Vectorize(StringRef name, linalg::LinalgVectorizationOptions options,
-            LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(f), opName(name), options(options) {}
+            LinalgTransformationFilter::FilterFunction f = nullptr,
+            bool padVectorize = false)
+      : Transformation(f), opName(name), options(options),
+        vectorizePadding(padVectorize) {}
 
   void addToPassPipeline(OpPassManager &pm,
                          LinalgTransformationFilter m) const override {
-    pm.addPass(createLinalgStrategyVectorizePass(opName, options, m));
+    pm.addPass(createLinalgStrategyVectorizePass(opName, options, m,
+                                                 vectorizePadding));
   }
 
 private:
   std::string opName;
   linalg::LinalgVectorizationOptions options;
+  bool vectorizePadding;
 };
 
 /// Represent one application of createLinalgStrategyLowerVectorsPass.
 struct VectorLowering : public Transformation {
   explicit VectorLowering(
       linalg::LinalgVectorLoweringOptions options,
       LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), options(options) {}
struct CodegenStrategy {
   interchangeIf(bool b, ArrayRef<int64_t> iteratorInterchange,
                 LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? interchange(iteratorInterchange, f) : *this;
     return *this;
   }
   /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
   CodegenStrategy &
   vectorize(StringRef opName,
-            LinalgTransformationFilter::FilterFunction f = nullptr) {
+            LinalgTransformationFilter::FilterFunction f = nullptr,
+            bool vectorizePadding = false) {
     assert(!opName.empty() && "expected an op name");
     transformationSequence.emplace_back(std::make_unique<Vectorize>(
-        opName, linalg::LinalgVectorizationOptions(), f));
+        opName, linalg::LinalgVectorizationOptions(), f, vectorizePadding));
     return *this;
   }
   /// Conditionally append a pattern to rewrite `LinalgOpType` as a vector
   /// operation.
   CodegenStrategy &
   vectorizeIf(bool b, StringRef opName,
-              LinalgTransformationFilter::FilterFunction f = nullptr) {
-    return b ? vectorize(opName, f) : *this;
+              LinalgTransformationFilter::FilterFunction f = nullptr,
+              bool vectorizePadding = false) {
+    return b ? vectorize(opName, f, vectorizePadding) : *this;
     return *this;
   }
   /// Append a pattern to lower all vector operations.
   CodegenStrategy &vectorLowering(LinalgVectorLoweringOptions options) {
     transformationSequence.emplace_back(
         std::make_unique<VectorLowering>(options));
     return *this;
   }

mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp

 
 /// Configurable pass to apply pattern-based linalg vectorization.
 struct LinalgStrategyVectorizePass
     : public LinalgStrategyVectorizePassBase<LinalgStrategyVectorizePass> {
 
   LinalgStrategyVectorizePass() = default;
 
   LinalgStrategyVectorizePass(StringRef opName, LinalgVectorizationOptions opt,
-                              LinalgTransformationFilter filt)
+                              LinalgTransformationFilter filt,
+                              bool padVectorize = false)
       : options(opt), filter(filt) {
     this->anchorOpName.setValue(opName.str());
+    this->vectorizePadding.setValue(padVectorize);
   }
 
   void runOnFunction() override {
     auto funcOp = getFunction();
     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
       return;
 
     RewritePatternSet vectorizationPatterns(funcOp.getContext());
     if (!anchorOpName.empty()) {
       vectorizationPatterns.add<LinalgVectorizationPattern>(
           anchorOpName, funcOp.getContext(), options, filter);
     } else {
       vectorizationPatterns.add<LinalgVectorizationPattern>(funcOp.getContext(),
                                                             filter, options);
     }
     vector::populateVectorTransferPermutationMapLoweringPatterns(
         vectorizationPatterns);
     vector::populateVectorReductionToContractPatterns(vectorizationPatterns);
     vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
                               linalg::LinalgCopyVTWForwardingPattern>(
         funcOp.getContext(), /*benefit=*/2);
+    if (vectorizePadding) {
+      linalg::populatePadTensorOpVectorizationPatterns(vectorizationPatterns);
+    }
     (void)applyPatternsAndFoldGreedily(funcOp,
                                        std::move(vectorizationPatterns));
   }
 
   LinalgVectorizationOptions options;
   LinalgTransformationFilter filter;
 };
 

 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
                                           LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
                                                          filter);
 }
 
 /// Create a LinalgStrategyVectorizePass.
-std::unique_ptr<OperationPass<FuncOp>>
-mlir::createLinalgStrategyVectorizePass(StringRef opName,
-                                        LinalgVectorizationOptions opt,
-                                        LinalgTransformationFilter filter) {
-  return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
+std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgStrategyVectorizePass(
+    StringRef opName, LinalgVectorizationOptions opt,
+    LinalgTransformationFilter filter, bool padVectorize) {
+  return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter,
+                                                       padVectorize);
 }
 
 /// Create a LinalgStrategyEnablePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyEnablePass(LinalgEnablingOptions opt,
                                      LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyEnablePass>(opt, filter);
 }