[CINN] Change add_cinn_pass to apply_cinn_pass (#62071)

zyfncg · web-flow · commit ab485e4e1744 · 2024-02-28T10:19:24.000+08:00
* refector lower cinn pass

* fix bug

* fix complie bug

* fix complie bug

* fix bug

* refine

* fix bug
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
@@ -51,8 +51,8 @@ COMMON_DECLARE_bool(check_infer_symbolic);
 namespace cinn::dialect::ir {
 
 namespace {
-bool HasDynamicShape(const pir::Program &program) {
-  for (const auto &op : *program.block()) {
+bool HasDynamicShape(const pir::Program& program) {
+  for (const auto& op : *program.block()) {
     if (op.isa<pir::CombineOp>()) {
       continue;
     }
@@ -70,18 +70,12 @@ bool HasDynamicShape(const pir::Program &program) {
 }
 }  // namespace
 
-void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
-                 pir::Program &program) {                          // NOLINT
-  pir::IrContext *ctx = pir::IrContext::Instance();
-  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
-  ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
-  ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
-
-  bool has_dynamic_shape = HasDynamicShape(program);
-
-  if (FLAGS_print_ir) {
-    pass_manager->EnableIRPrinting();
-  }
+void ApplyCinnPreprocessPass(
+    ::pir::Program* program,
+    const std::function<std::shared_ptr<::pir::PassManager>()>&
+        CreatePassManager) {
+  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
+  bool has_dynamic_shape = HasDynamicShape(*program);
 
   pass_manager->AddPass(cinn::dialect::ir::CreateConvert0DTo1DPass());
   if (!has_dynamic_shape && FLAGS_check_infer_symbolic) {
@@ -116,6 +110,17 @@ void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
   pass_manager->AddPass(cinn::dialect::ir::CreateReplaceDynamicExpandOpPass());
   pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
 
+  pass_manager->Run(program);
+}
+
+void ApplyCinnLowerPass(
+    ::pir::Program* program,
+    const std::function<std::shared_ptr<pir::PassManager>()>&
+        CreatePassManager) {
+  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
+
+  bool has_dynamic_shape = HasDynamicShape(*program);
+
   bool force_static_shape = false;
   if (auto pass = cinn::dialect::ir::CreateConvertDynamicToStaticDimPass()) {
     pass_manager->AddPass(std::move(pass.value()));
@@ -129,9 +134,19 @@ void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
     pass_manager->AddPass(
         cinn::dialect::ir::CreateLowerCinnDyShapeFusionOpPass());
   }
+
   pass_manager->AddPass(cinn::dialect::ir::CreateLowerCinnFusionOpPass());
   pass_manager->AddPass(
       cinn::dialect::ir::CreateSplitGenerateShapeIntoShapeOpsPass());
+
+  pass_manager->Run(program);
+}
+
+void ApplyCinnPass(::pir::Program* program,
+                   const std::function<std::shared_ptr<pir::PassManager>()>&
+                       CreatePassManager) {
+  ApplyCinnPreprocessPass(program, CreatePassManager);
+  ApplyCinnLowerPass(program, CreatePassManager);
 }
 
 }  // namespace cinn::dialect::ir
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <functional>
 #include <memory>
 
 namespace pir {
@@ -25,7 +26,8 @@ class Program;
 
 namespace cinn::dialect::ir {
 
-void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
-                 pir::Program &program);                           // NOLINT
+void ApplyCinnPass(::pir::Program* program,
+                   const std::function<std::shared_ptr<pir::PassManager>()>&
+                       CreatePassManager);
 
 }  // namespace cinn::dialect::ir
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/divide_group_op_to_fusion_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/divide_group_op_to_fusion_op_pass.cc
@@ -199,7 +199,7 @@ class DivideGroupOpToFusionOpPass : public pir::PatternRewritePass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+    return op->num_regions() > 0;
   }
 };
 
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc
@@ -751,7 +751,7 @@ class LowerCinnFusionOpPass : public pir::PatternRewritePass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+    return op->num_regions() > 0;
   }
 };
 
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -106,7 +106,9 @@
 #endif
 
 #ifdef PADDLE_WITH_CINN
+#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h"
+#include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
 #endif
 
 #include "paddle/common/flags.h"
@@ -880,17 +882,21 @@ bool AnalysisPredictor::PrepareExecutor() {
       }
 #ifdef PADDLE_WITH_CINN
       if (config_.cinn_enabled()) {
-        VLOG(4) << "[CINN] Begin AddCinnPass";
-        auto cinn_pm = std::make_shared<::pir::PassManager>(
-            ::pir::IrContext::Instance(), 2);
-        cinn::dialect::ir::AddCinnPass(cinn_pm, *pir_program_.get());
-        if (!config_.glog_info_disabled()) {
-          cinn_pm->EnablePrintStatistics();
-        }
-        if (config_.ir_debug_) {
-          cinn_pm->EnableIRPrinting();
-        }
-        cinn_pm->Run(pir_program_.get());
+        VLOG(4) << "[CINN] Begin ApplyCinnPass";
+        cinn::dialect::ir::ApplyCinnPass(pir_program_.get(), [&] {
+          pir::IrContext *ctx = pir::IrContext::Instance();
+          ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
+          ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
+          auto pass_manager = std::make_shared<::pir::PassManager>(
+              ::pir::IrContext::Instance(), 2);
+          if (!config_.glog_info_disabled()) {
+            pass_manager->EnablePrintStatistics();
+          }
+          if (config_.ir_debug_) {
+            pass_manager->EnableIRPrinting();
+          }
+          return pass_manager;
+        });
       }
 #endif
 
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -83,6 +83,7 @@
 #include "pybind11/stl.h"
 
 #ifdef PADDLE_WITH_CINN
+#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h"
 #include "paddle/cinn/hlir/framework/pir_compiler.h"
 #endif
@@ -1551,10 +1552,19 @@ bool HasDynamicShape(const pir::Program &program) {
   return false;
 }
 
-void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
-                 pir::Program &program) {                          // NOLINT
+void ApplyCinnPass(Program &program) {  // NOLINT
 #ifdef PADDLE_WITH_CINN
-  cinn::dialect::ir::AddCinnPass(pass_manager, program);
+  cinn::dialect::ir::ApplyCinnPass(&program, [] {
+    pir::IrContext *ctx = pir::IrContext::Instance();
+    ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+    ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
+    ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
+    auto pass_manager = std::make_shared<pir::PassManager>(ctx);
+    if (FLAGS_print_ir) {
+      pass_manager->EnableIRPrinting();
+    }
+    return pass_manager;
+  });
 #else
   PADDLE_THROW(common::errors::Unimplemented(
       "Currently we only support CINN Pass for Pir under @to_static, please "
@@ -1575,7 +1585,7 @@ void InferSymbolicShapePass(
 }
 
 void BindIrPass(pybind11::module *m) {
-  m->def("add_cinn_pass", AddCinnPass);
+  m->def("apply_cinn_pass", ApplyCinnPass);
   m->def("infer_symbolic_shape_pass", InferSymbolicShapePass);
 
   py::class_<Pass, std::shared_ptr<Pass>> pass(*m,
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -552,13 +552,17 @@ def _create_program(self, is_infer_mode=False):
         if is_infer_mode:
 
             def pass_fn(forward_program, backward_program):
+                # common pass
                 pm = paddle.base.libpaddle.pir.PassManager()
                 paddle.base.libpaddle.pir.infer_symbolic_shape_pass(
                     pm, forward_program
                 )
-                if self._build_strategy.build_cinn_pass:
-                    paddle.base.libpaddle.pir.add_cinn_pass(pm, forward_program)
                 pm.run(forward_program)
+
+                # if-else pass
+                if self._build_strategy.build_cinn_pass:
+                    paddle.base.libpaddle.pir.apply_cinn_pass(forward_program)
+
                 return forward_program, backward_program
 
             # TODO(xiongkun) who to transfer the pruning program?
@@ -574,18 +578,9 @@ def pass_fn(forward_program, backward_program):
             self._set_grad_type(self._params, train_program)
 
             def pass_fn(forward_program, backward_program):
-                fwd_pm = paddle.base.libpaddle.pir.PassManager()
-                bwd_pm = paddle.base.libpaddle.pir.PassManager()
-
                 if self._build_strategy.build_cinn_pass:
-                    paddle.base.libpaddle.pir.add_cinn_pass(
-                        fwd_pm, forward_program
-                    )
-                    paddle.base.libpaddle.pir.add_cinn_pass(
-                        bwd_pm, backward_program
-                    )
-                    fwd_pm.run(forward_program)
-                    bwd_pm.run(backward_program)
+                    paddle.base.libpaddle.pir.apply_cinn_pass(forward_program)
+                    paddle.base.libpaddle.pir.apply_cinn_pass(backward_program)
                 return forward_program, backward_program
 
             train_program.apply_pir_program_pass(pass_fn)
diff --git a/test/ir/pir/my_task.py b/test/ir/pir/my_task.py
@@ -183,9 +183,7 @@ def check_infer(self, enable_cinn):
     def run_program(self, program, feed, fetch_list, enable_cinn):
         if enable_cinn:
             paddle.decomposition.decomp.decompose(program, [])
-            fwd_pm = paddle.base.libpaddle.pir.PassManager()
-            paddle.base.libpaddle.pir.add_cinn_pass(fwd_pm, program)
-            fwd_pm.run(program)
+            paddle.base.libpaddle.pir.apply_cinn_pass(program)
 
         exe = paddle.static.Executor(paddle.CUDAPlace(0))
         outs = exe._run_pir_impl(

Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,7 @@ class DivideGroupOpToFusionOpPass : public pir::PatternRewritePass {`
`199`	`199`	`}`
`200`	`200`
`201`	`201`	`bool CanApplyOn(pir::Operation* op) const override {`
`202`		`- return op->isa<pir::ModuleOp>() && op->num_regions() > 0;`
	`202`	`+ return op->num_regions() > 0;`
`203`	`203`	`}`
`204`	`204`	`};`
`205`	`205`
Original file line number	Diff line number	Diff line change
`@@ -751,7 +751,7 @@ class LowerCinnFusionOpPass : public pir::PatternRewritePass {`
`751`	`751`	`}`
`752`	`752`
`753`	`753`	`bool CanApplyOn(pir::Operation* op) const override {`
`754`		`- return op->isa<pir::ModuleOp>() && op->num_regions() > 0;`
	`754`	`+ return op->num_regions() > 0;`
`755`	`755`	`}`
`756`	`756`	`};`
`757`	`757`