Call RewriteONNXToONNX from the driver

Signed-off-by: Tung D. Le <[email protected]>
onnx · Sep 11, 2023 · e55713b · e55713b
1 parent d432ee5
commit e55713b
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/src/Compiler/CompilerPasses.cpp b/src/Compiler/CompilerPasses.cpp
@@ -76,8 +76,6 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU) {
     pm.addPass(mlir::createCanonicalizerPass());
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
   }
-  // Rewrite ONNX operators.
-  pm.addPass(onnx_mlir::createRewriteONNXToONNXPass());
   // Convolution Optimization for CPU: enable when there are no accelerators.
   if (targetCPU && enableConvOptPass) {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createConvOptONNXToONNXPass(
@@ -97,16 +95,14 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU) {
     // Statically add extra passes
     for (int i = 0; i < repeatOnnxTransform; i++) {
       pm.addPass(mlir::createCanonicalizerPass());
+      pm.addPass(onnx_mlir::createSimplifyShapeRelatedOpsPass());
       pm.addPass(onnx_mlir::createRewriteONNXToONNXPass());
       pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
       pm.addNestedPass<func::FuncOp>(
           onnx_mlir::createConstPropONNXToONNXPass());
     }
   }
 
-  // Simplify shape-related ops.
-  pm.addPass(onnx_mlir::createSimplifyShapeRelatedOpsPass());
-
   // One more call to ONNX shape inference/canonicalization/... to update shape
   // if possible.
   if (enableONNXHybridPass) {
@@ -116,7 +112,6 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU) {
   } else {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
     pm.addPass(mlir::createCanonicalizerPass());
-    pm.addPass(onnx_mlir::createRewriteONNXToONNXPass());
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
   }
 

diff --git a/src/Transform/ONNX/CMakeLists.txt b/src/Transform/ONNX/CMakeLists.txt
@@ -62,6 +62,7 @@ add_onnx_mlir_library(OMOpTransform
   OMONNXOps
   MLIRPass
   OMONNXRewrite
+  OMONNXSimplifyShapeRelatedOps
   OMShapeInferencePass
   MLIRTransforms
   )

diff --git a/src/Transform/ONNX/ONNXOpTransformPass.cpp b/src/Transform/ONNX/ONNXOpTransformPass.cpp
@@ -88,6 +88,9 @@ void ONNXOpTransformPass::runOnOperation() {
     }
     dynamicPM.addNestedPass<func::FuncOp>(
         onnx_mlir::createConstPropONNXToONNXPass());
+    // Simplify shape-related ops.
+    dynamicPM.addPass(onnx_mlir::createSimplifyShapeRelatedOpsPass());
+    // Rewrite ONNX operators.
     dynamicPM.addPass(onnx_mlir::createRewriteONNXToONNXPass());
     if (failed(runPipeline(dynamicPM, module)))
       return signalPassFailure();