[mlir][linalg] Add vector transfer lowering patters to codegen strategy.

Add the vector transfer patterns and introduce the max transfer rank option on the codegen strategy.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D111635

Author: Tobias Gysi

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

@@ -224,6 +224,10 @@ struct CodegenStrategy {
         .enableHoistRedundantVectorTransfersOnTensor = val;
     return *this;
   }
+  CodegenStrategy &setMaxTransferRank(int64_t val) {
+    this->lateCodegenStrategyOptions.maxTransferRank = val;
+    return *this;
+  }
   CodegenStrategy &setEnableVectorTransferPartialRewrite(bool val) {
     this->lateCodegenStrategyOptions.enableVectorTransferPartialRewrite = val;
     return *this;

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

@@ -835,6 +835,7 @@ struct LateCodegenStrategyOptions {
   bool enableHoistRedundantVectorTransfersOnTensor = true;
   /// Vector lowering operations may result in surprising behavior when
   /// composing multiple codegen strategies and must be enabled explicitly.
+  int64_t maxTransferRank = 1;
   bool enableVectorTransferPartialRewrite = false;
   bool enableVectorContractLowering = false;
   bool enableVectorToSCFConversion = false;
@@ -852,6 +853,7 @@ struct LinalgEnablingOptions {
 /// Vector lowering options control how ops are lowered down to 1-D and scf.for
 /// form.
 struct LinalgVectorLoweringOptions {
+  int64_t maxTransferRank = 1;
   bool enableVectorTransferPartialRewrite = false;
   bool enableVectorContractLowering = false;
   bool enableVectorToSCFConversion = false;

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

@@ -47,6 +47,8 @@ void mlir::linalg::CodegenStrategy::configurePassPipeline(
     pm.addPass(createLinalgStrategyEnablePass());
   }
   LinalgVectorLoweringOptions vectorLoweringOptions;
+  vectorLoweringOptions.maxTransferRank =
+      lateCodegenStrategyOptions.maxTransferRank;
   vectorLoweringOptions.enableVectorTransferPartialRewrite =
       lateCodegenStrategyOptions.enableVectorTransferPartialRewrite;
   vectorLoweringOptions.enableVectorContractLowering =

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

@@ -260,6 +260,8 @@ struct LinalgStrategyLowerVectorsPass
 
     MLIRContext *context = funcOp.getContext();
     RewritePatternSet patterns(context);
+    vector::populateVectorTransferLoweringPatterns(patterns,
+                                                   options.maxTransferRank);
     if (options.enableVectorTransferPartialRewrite) {
       patterns.add<vector::VectorTransferFullPartialRewriter>(
           context, options.vectorTransformOptions);