[Mlir-commits] [mlir] ba2ac9c - [mli][linalg] Add flag to control CodegenStrategy enable pass.

[Tobias Gysi]

Author: Tobias Gysi
Date: 2021-11-10T14:11:40Z
New Revision: ba2ac9c97c0c1f497d5bf7895163819d158219e5

URL: https://github.com/llvm/llvm-project/commit/ba2ac9c97c0c1f497d5bf7895163819d158219e5
DIFF: https://github.com/llvm/llvm-project/commit/ba2ac9c97c0c1f497d5bf7895163819d158219e5.diff

LOG: [mli][linalg] Add flag to control CodegenStrategy enable pass.

Add a flag to control if CodegenStrategy runs the EnablePass between the transformations.

Depends On D113382

Reviewed By: nicolasvasilache

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
    mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
    mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h b/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
index bff0c888eb85e..f4579a6150489 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h
@@ -257,7 +257,8 @@ struct CodegenStrategy {
 
   /// Apply the transformation patterns in sequence with cleanup
   /// transformations interleaved.
-  void configurePassPipeline(OpPassManager &pm, MLIRContext *context) const;
+  void configurePassPipeline(OpPassManager &pm, MLIRContext *context,
+                             bool addEnablePass = true) const;
 
 private:
   LogicalResult postPatternTransforms(Operation *func) const;

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp b/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
index 862eac7fa8e97..d2d1752499adb 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp
@@ -29,7 +29,7 @@ using namespace mlir::linalg;
 #define DEBUG_TYPE "linalg-codegen-strategy"
 
 void mlir::linalg::CodegenStrategy::configurePassPipeline(
-    OpPassManager &pm, MLIRContext *context) const {
+    OpPassManager &pm, MLIRContext *context, bool addEnablePass) const {
   for (unsigned stepCount = 0, e = transformationSequence.size(); stepCount < e;
        ++stepCount) {
     const std::unique_ptr<Transformation> &t =
@@ -44,7 +44,8 @@ void mlir::linalg::CodegenStrategy::configurePassPipeline(
                       : linalg::LinalgTransformationFilter(
                             t->filter, currentState, nextState);
     t->addToPassPipeline(pm, filter);
-    pm.addPass(createLinalgStrategyEnablePass(linalgEnablingOptions));
+    if (addEnablePass)
+      pm.addPass(createLinalgStrategyEnablePass(linalgEnablingOptions));
   }
   pm.addPass(createLinalgStrategyRemoveMarkersPass());
 }

diff  --git a/mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp b/mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp
index 249217fa58c41..49b2bb9b4bcd6 100644
--- a/mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp
+++ b/mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp
@@ -125,6 +125,10 @@ struct TestLinalgCodegenStrategy
       *this, "unroll-vector-transfers",
       llvm::cl::desc("Enable full unrolling of vector.transfer operations"),
       llvm::cl::init(false)};
+  Option<bool> runEnablePass{
+      *this, "run-enable-pass",
+      llvm::cl::desc("Run the enable pass between transformations"),
+      llvm::cl::init(true)};
   Option<std::string> anchorOpName{
       *this, "anchor-op",
       llvm::cl::desc(
@@ -178,11 +182,10 @@ void TestLinalgCodegenStrategy::runStrategy(
               .enableTransferPartialRewrite()
               .enableContractionLowering()
               .enableTransferToSCFConversion());
-
   // Created a nested OpPassManager and run.
   FuncOp funcOp = getFunction();
   OpPassManager dynamicPM("builtin.func");
-  strategy.configurePassPipeline(dynamicPM, funcOp.getContext());
+  strategy.configurePassPipeline(dynamicPM, funcOp.getContext(), runEnablePass);
   if (failed(runPipeline(dynamicPM, funcOp)))
     return signalPassFailure();
 }