Move more passes to GlobalOptimization

iree-org · Aug 24, 2023 · aa33602 · aa33602
1 parent c477184
commit aa33602
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 15 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -179,19 +179,6 @@ void buildFlowTransformPassPipeline(OpPassManager &passManager,
   }
 
   FunctionLikeNest(passManager)
-      // Preprocess the input to a form more amenable for fusion
-      // - Convert all elementwise ops to Linalg
-      // - Remove unit-extent dimensions.
-      .addPass(mlir::createConvertElementwiseToLinalgPass)
-      .addPass(createGeneralizeLinalgNamedOpsPass)
-      .addPass(createFuseDequantizationMatmulPass)
-      .addPass(createFoldUnitExtentDimsPass)
-      .addPass(createRaiseSpecialOps)
-      .addPass(createInterchangeGenericOpsPass)
-      .addPass(createCollapseDimsPass)
-      .addPass(memref::createResolveShapedTypeResultDimsPass)
-      .addPass(mlir::createCanonicalizerPass)
-      .addPass(mlir::createCSEPass)
       // Elementwise fusion.
       .addPass(
           []() { return createFusionOfTensorOpsPass(clEnableFuseMultiUse); })

diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -28,6 +28,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgTransforms",
+        "@llvm-project//mlir:MemRefTransforms",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Transforms",
     ],

diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -22,6 +22,7 @@ iree_cc_library(
     MLIRFuncDialect
     MLIRIR
     MLIRLinalgTransforms
+    MLIRMemRefTransforms
     MLIRPass
     MLIRTransforms
     iree::compiler::Dialect::Flow::Transforms

diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -10,6 +10,7 @@
 #include "iree/compiler/Utils/PassUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Transforms/Passes.h"
 
 namespace mlir {
@@ -27,10 +28,24 @@ void buildGlobalOptimizationPassPipeline(
       .addPass(mlir::createLinalgNamedOpConversionPass)
       .addPass(IREE::Flow::createConvert1X1FilterConv2DToMatmulPass);
   mainPassManager.addPass(IREE::Flow::createEraseUnusedLinalgOperands());
+  FunctionLikeNest(mainPassManager)
+      // Preprocess the input to a form more amenable for fusion
+      // - Convert all elementwise ops to Linalg
+      // - Remove unit-extent dimensions.
+      .addPass(mlir::createConvertElementwiseToLinalgPass)
+      .addPass(IREE::Flow::createGeneralizeLinalgNamedOpsPass)
+      .addPass(IREE::Flow::createFuseDequantizationMatmulPass)
+      .addPass(IREE::Flow::createFoldUnitExtentDimsPass)
+      .addPass(IREE::Flow::createRaiseSpecialOps)
+      .addPass(IREE::Flow::createInterchangeGenericOpsPass)
+      .addPass(IREE::Flow::createCollapseDimsPass)
+      .addPass(memref::createResolveShapedTypeResultDimsPass)
+      .addPass(mlir::createCanonicalizerPass)
+      .addPass(mlir::createCSEPass);
 
   // Expand tensor shapes into SSA values and optimize the whole program.
-  // The more we are able to equate shape dimensions at this level the better
-  // our fusions will be.
+  // The more we are able to equate shape dimensions at this level the
+  // better our fusions will be.
   FunctionLikeNest(mainPassManager)
       .addPass(IREE::Flow::createTopLevelSCFToCFGPass);
   mainPassManager.addPass(IREE::Flow::createExpandTensorShapesPass());