From c8faeb1edd8447fb67ed7ef04158a07582aa8771 Mon Sep 17 00:00:00 2001
From: Tobias Gysi <gysit@google.com>
Date: Tue, 12 Oct 2021 14:53:33 +0000
Subject: [PATCH] [mlir][linalg] Add switch to disable/enable vector transfer
 lowering.

Add a switch to code gen strategy to disable/enable the vector transfer lowering and disable it by default.

Differential Revision: https://reviews.llvm.org/D111647
---
 mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h | 4 ++++
 mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h      | 2 ++
 mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp        | 2 ++
 mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp   | 6 ++++--
 4 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h b/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
index 251b850..7454c91 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
@@ -228,6 +228,10 @@ struct CodegenStrategy {
     this->lateCodegenStrategyOptions.maxTransferRank = val;
     return *this;
   }
+  CodegenStrategy &setEnableVectorTransferLowering(bool val) {
+    this->lateCodegenStrategyOptions.enableVectorTransferLowering = val;
+    return *this;
+  }
   CodegenStrategy &setEnableVectorTransferPartialRewrite(bool val) {
     this->lateCodegenStrategyOptions.enableVectorTransferPartialRewrite = val;
     return *this;
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 6d8422f..c29e364 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -836,6 +836,7 @@ struct LateCodegenStrategyOptions {
   /// Vector lowering operations may result in surprising behavior when
   /// composing multiple codegen strategies and must be enabled explicitly.
   int64_t maxTransferRank = 1;
+  bool enableVectorTransferLowering = false;
   bool enableVectorTransferPartialRewrite = false;
   bool enableVectorContractLowering = false;
   bool enableVectorToSCFConversion = false;
@@ -854,6 +855,7 @@ struct LinalgEnablingOptions {
 /// form.
 struct LinalgVectorLoweringOptions {
   int64_t maxTransferRank = 1;
+  bool enableVectorTransferLowering = false;
   bool enableVectorTransferPartialRewrite = false;
   bool enableVectorContractLowering = false;
   bool enableVectorToSCFConversion = false;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp b/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
index 04bdc49..1770cd9f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
@@ -49,6 +49,8 @@ void mlir::linalg::CodegenStrategy::configurePassPipeline(
   LinalgVectorLoweringOptions vectorLoweringOptions;
   vectorLoweringOptions.maxTransferRank =
       lateCodegenStrategyOptions.maxTransferRank;
+  vectorLoweringOptions.enableVectorTransferLowering =
+      lateCodegenStrategyOptions.enableVectorTransferLowering;
   vectorLoweringOptions.enableVectorTransferPartialRewrite =
       lateCodegenStrategyOptions.enableVectorTransferPartialRewrite;
   vectorLoweringOptions.enableVectorContractLowering =
diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp
index bb1793d..617ea0a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp
@@ -260,8 +260,10 @@ struct LinalgStrategyLowerVectorsPass
 
     MLIRContext *context = funcOp.getContext();
     RewritePatternSet patterns(context);
-    vector::populateVectorTransferLoweringPatterns(patterns,
-                                                   options.maxTransferRank);
+    if (options.enableVectorTransferLowering) {
+      vector::populateVectorTransferLoweringPatterns(patterns,
+                                                     options.maxTransferRank);
+    }
     if (options.enableVectorTransferPartialRewrite) {
       patterns.add<vector::VectorTransferFullPartialRewriter>(
           context, options.vectorTransformOptions);
-- 
2.7.4