From aa3360298f990b462b71995ca20799626d7f9970 Mon Sep 17 00:00:00 2001
From: Hanhan Wang <hanchung@google.com>
Date: Thu, 24 Aug 2023 14:33:47 -0700
Subject: [PATCH] Move more passes to GlobalOptimization

---
 .../Dialect/Flow/Transforms/Passes.cpp        | 13 -------------
 .../compiler/GlobalOptimization/BUILD.bazel   |  1 +
 .../GlobalOptimization/CMakeLists.txt         |  1 +
 .../compiler/GlobalOptimization/Passes.cpp    | 19 +++++++++++++++++--
 4 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index eb019e6f51667..d60bf446e2e29 100644
--- 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
index ffa3218cb7230..0c4af96873f9b 100644
--- 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
index 99bfbb82711ff..f1846f82c5aac 100644
--- 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
index ff00f4c7592c8..cf28b348058ec 100644
--- 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());