[DRR][Inference] Fix a drr rewrite bug, Adjust the order of basic pass required by framework in pir inference (#63062)

yuanlehome · web-flow · commit 75a3f48db8b5 · 2024-03-28T20:29:34.000+08:00
* refine some code

* update

* update

* fix drr rewrite
diff --git a/paddle/fluid/inference/api/analysis_config.cc b/paddle/fluid/inference/api/analysis_config.cc
@@ -593,7 +593,9 @@ AnalysisConfig::AnalysisConfig(const AnalysisConfig &other) {
   CP_MEMBER(use_new_executor_);
   CP_MEMBER(use_pir_);
   CP_MEMBER(custom_passes_);
+  CP_MEMBER(custom_pass_only_);
   CP_MEMBER(pm_opt_level_);
+  CP_MEMBER(ir_debug_passes_);
 
   if (use_gpu_) {
     PADDLE_ENFORCE_EQ(use_xpu_,
@@ -1326,8 +1328,10 @@ NativeConfig AnalysisConfig::ToNativeConfig() const {
   return config;
 }
 
-void AnalysisConfig::SwitchIrDebug(int x) {
+void AnalysisConfig::SwitchIrDebug(int x,
+                                   const std::vector<std::string> &passes) {
   ir_debug_ = x;
+  ir_debug_passes_ = passes;
   Update();
 }
 
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -115,7 +115,6 @@
 #include "paddle/fluid/pir/transforms/general/constant_folding_pass.h"
 #include "paddle/fluid/pir/transforms/general/dead_code_elimination_pass.h"
 #include "paddle/fluid/pir/transforms/general/inplace_pass.h"
-#include "paddle/fluid/pir/transforms/general/matmul_transpose_fuse_pass.h"
 #include "paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.h"
 #include "paddle/fluid/pir/transforms/general/replace_fetch_with_shadow_output_pass.h"
 #include "paddle/fluid/pir/transforms/passes.h"
@@ -886,6 +885,16 @@ bool AnalysisPredictor::PrepareExecutor() {
       pir_program_ =
           paddle::TranslateLegacyProgramToProgram(*inference_program_);
 
+      auto ir_printing_conditions = [this](::pir::Pass *pass,
+                                           ::pir::Operation *op) {
+        if (this->config_.ir_debug_passes_.empty()) {
+          return true;
+        }
+        return std::find(this->config_.ir_debug_passes_.begin(),
+                         this->config_.ir_debug_passes_.end(),
+                         pass->name()) != this->config_.ir_debug_passes_.end();
+      };
+
 #ifdef PADDLE_WITH_CINN
       if (paddle::prim::PrimCommonUtils::IsFwdPrimEnabled()) {
         VLOG(4) << "[Prim] Decomp program in predictor begin.";
@@ -911,13 +920,16 @@ bool AnalysisPredictor::PrepareExecutor() {
             pass_manager->EnablePrintStatistics();
           }
           if (config_.ir_debug_) {
-            pass_manager->EnableIRPrinting();
+            pass_manager->EnableIRPrinting(
+                std::make_unique<pir::PassManager::IRPrinterOption>(
+                    ir_printing_conditions, ir_printing_conditions));
           }
           return pass_manager;
         });
       }
 #endif
 
+      // Apply some optimization passes required by the inference
       ::pir::PassManager pass_pm(::pir::IrContext::Instance(),
                                  config_.pm_opt_level_);
       if (!config_.custom_passes_.empty()) {
@@ -934,14 +946,6 @@ bool AnalysisPredictor::PrepareExecutor() {
                 std::move(pir::PassRegistry::Instance().Get(gpu_pass)));
           }
         }
-        // Basic pass required by the framework
-        auto params_sync_among_devices_pass =
-            ::pir::CreateParamsSyncAmongDevicesPass();
-        params_sync_among_devices_pass->SetNotOwned(pir::kPlaceAttr, &place_);
-        params_sync_among_devices_pass->SetNotOwned(pir::kParamScopeAttr,
-                                                    sub_scope_);
-        pass_pm.AddPass(std::move(params_sync_among_devices_pass));
-
 #ifdef PADDLE_WITH_DNNL
       } else if (config_.mkldnn_enabled()) {
         // mkldnn
@@ -961,21 +965,46 @@ bool AnalysisPredictor::PrepareExecutor() {
           }
         }
       }
-      auto constant_folding_pass = ::pir::CreateConstantFoldingPass();
-      constant_folding_pass->SetNotOwned(pir::kPlaceAttr, &place_);
-      constant_folding_pass->SetNotOwned(pir::kParamScopeAttr, sub_scope_);
-      pass_pm.AddPass(std::move(constant_folding_pass));
-      pass_pm.AddPass(::pir::CreateDeadCodeEliminationPass());
-      pass_pm.AddPass(::pir::CreateReplaceFetchWithShadowOutputPass());
-      //----------------------------------------------------------------------------------------------//
+
       if (!config_.glog_info_disabled()) {
         pass_pm.EnablePrintStatistics();
       }
       if (config_.ir_debug_) {
-        pass_pm.EnableIRPrinting();
+        pass_pm.EnableIRPrinting(
+            std::make_unique<pir::PassManager::IRPrinterOption>(
+                ir_printing_conditions, ir_printing_conditions));
       }
       pass_pm.Run(pir_program_.get());
 
+      // Apply some basic passes required by the framework
+      ::pir::PassManager basic_pass_pm(::pir::IrContext::Instance(),
+                                       config_.pm_opt_level_);
+
+      auto params_sync_among_devices_pass =
+          ::pir::CreateParamsSyncAmongDevicesPass();
+      params_sync_among_devices_pass->SetNotOwned(pir::Pass::kPlaceAttr,
+                                                  &place_);
+      params_sync_among_devices_pass->SetNotOwned(pir::Pass::kParamScopeAttr,
+                                                  sub_scope_);
+      basic_pass_pm.AddPass(std::move(params_sync_among_devices_pass));
+      auto constant_folding_pass = ::pir::CreateConstantFoldingPass();
+      constant_folding_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place_);
+      constant_folding_pass->SetNotOwned(pir::Pass::kParamScopeAttr,
+                                         sub_scope_);
+      basic_pass_pm.AddPass(std::move(constant_folding_pass));
+      basic_pass_pm.AddPass(::pir::CreateDeadCodeEliminationPass());
+      basic_pass_pm.AddPass(::pir::CreateReplaceFetchWithShadowOutputPass());
+      if (!config_.glog_info_disabled()) {
+        basic_pass_pm.EnablePrintStatistics();
+      }
+      if (config_.ir_debug_) {
+        basic_pass_pm.EnableIRPrinting(
+            std::make_unique<pir::PassManager::IRPrinterOption>(
+                ir_printing_conditions, ir_printing_conditions));
+      }
+      basic_pass_pm.Run(pir_program_.get());
+      //----------------------------------------------------------------------------------------------//
+
       pir_program_ =
           paddle::dialect::PdOpLowerToKernelPass(pir_program_.get(), place_);
 
@@ -987,7 +1016,9 @@ bool AnalysisPredictor::PrepareExecutor() {
         lowered_pm.EnablePrintStatistics();
       }
       if (config_.ir_debug_) {
-        lowered_pm.EnableIRPrinting();
+        lowered_pm.EnableIRPrinting(
+            std::make_unique<pir::PassManager::IRPrinterOption>(
+                ir_printing_conditions, ir_printing_conditions));
       }
       lowered_pm.Run(pir_program_.get());
 
diff --git a/paddle/fluid/inference/api/paddle_analysis_config.h b/paddle/fluid/inference/api/paddle_analysis_config.h
@@ -967,7 +967,7 @@ struct PD_INFER_DECL AnalysisConfig {
   ///
   /// \param x whether to debug IR graph analysis phase.
   ///
-  void SwitchIrDebug(int x = true);
+  void SwitchIrDebug(int x = true, const std::vector<std::string>& passes = {});
 
   ///
   /// \brief Turn on MKLDNN.
@@ -1485,6 +1485,7 @@ struct PD_INFER_DECL AnalysisConfig {
   std::vector<std::string> custom_passes_;
   bool custom_pass_only_{false};
   int pm_opt_level_{2};
+  std::vector<std::string> ir_debug_passes_;
 };
 
 }  // namespace paddle
diff --git a/paddle/fluid/pir/drr/src/rewrite_pattern.cc b/paddle/fluid/pir/drr/src/rewrite_pattern.cc
@@ -473,7 +473,7 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
   GraphTopo graph_topo_visit(&result_pattern_graph);
   graph_topo_visit.WalkGraphNodesTopoOrder([&](const OpCall& op_call) {
     // set insert point
-    size_t max_input_op_index = 0;
+    size_t max_input_op_index = 0UL;
     pir::Operation* max_index_op = nullptr;
     for (const Tensor* input : op_call.inputs()) {
       if (input->is_none()) {
@@ -483,7 +483,7 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
       if (ir_val) {
         pir::Operation* ir_input_op = ir_val.defining_op();
         if (op_2_temp_program_index.count(ir_input_op) == 0) {
-          max_input_op_index = 0UL;
+          // do nothing
         } else if (max_input_op_index <
                    op_2_temp_program_index.at(ir_input_op)) {
           max_input_op_index = op_2_temp_program_index.at(ir_input_op);
diff --git a/paddle/fluid/pir/transforms/general/auto_mixed_precision_pass.cc b/paddle/fluid/pir/transforms/general/auto_mixed_precision_pass.cc
@@ -62,7 +62,7 @@ class AutoMixedPrecisionPass : public pir::Pass {
 
   bool Initialize(pir::IrContext* context) override {
     PADDLE_ENFORCE_EQ(
-        Has(pir::kPlaceAttr),
+        Has(pir::Pass::kPlaceAttr),
         true,
         phi::errors::InvalidArgument(
             "Pass initialize failed."
@@ -77,7 +77,7 @@ class AutoMixedPrecisionPass : public pir::Pass {
             "required!"
             "Use Set method to set the scope attribute."));
 
-    place_ = Get<phi::Place>(pir::kPlaceAttr);
+    place_ = Get<phi::Place>(pir::Pass::kPlaceAttr);
     precision_mode_ = Get<phi::DataType>("__mixed_precision_mode__");
     context_ = context;
     enable_low_precision_io_ = false;
diff --git a/paddle/fluid/pir/transforms/general/constant_folding_pass.cc b/paddle/fluid/pir/transforms/general/constant_folding_pass.cc
@@ -461,30 +461,27 @@ class ConstantFoldingPatternForTrain : public ConstantFoldingPattern {
 
 class ConstantFoldingPass : public pir::Pass {
  public:
-  ConstantFoldingPass()
-      : pir::Pass("constant_folding_pass", 1),
-        place_(phi::CPUPlace{}),
-        scope_(nullptr) {}
+  ConstantFoldingPass() : pir::Pass("constant_folding_pass", 1) {}
 
  private:
   bool Initialize(pir::IrContext* context) override {
     PADDLE_ENFORCE_EQ(
-        Has(pir::kPlaceAttr),
+        Has(pir::Pass::kPlaceAttr),
         true,
         phi::errors::InvalidArgument(
             "Pass initialize failed."
             "When using ConstantFoldingPass, place attribute is required!"
             "Use Set method to set the place attribute."));
     PADDLE_ENFORCE_EQ(
-        Has(pir::kParamScopeAttr),
+        Has(pir::Pass::kParamScopeAttr),
         true,
         phi::errors::InvalidArgument(
             "Pass initialize failed."
             "When using ConstantFoldingPass, scope attribute is required!"
             "Use Set method to set the scope attribute."));
 
-    place_ = Get<phi::Place>(pir::kPlaceAttr);
-    scope_ = &Get<paddle::framework::Scope>(pir::kParamScopeAttr);
+    place_ = Get<phi::Place>(pir::Pass::kPlaceAttr);
+    scope_ = &Get<paddle::framework::Scope>(pir::Pass::kParamScopeAttr);
 
     PADDLE_ENFORCE_NOT_NULL(
         scope_, phi::errors::InvalidArgument("scope can not be nullptr"));
@@ -529,7 +526,7 @@ class ConstantFoldingPass : public pir::Pass {
 
  private:
   size_t suffix_{0};
-  phi::Place place_;
+  phi::Place place_{phi::CPUPlace{}};
   paddle::framework::Scope* scope_{nullptr};
   paddle::framework::interpreter::ExecutionConfig exe_config_{};
   std::vector<std::string> deleted_vars_;
diff --git a/paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.cc b/paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.cc
@@ -38,30 +38,22 @@ class ParamsSyncAmongDevicesPass : public pir::Pass {
 
   bool Initialize(pir::IrContext* context) override {
     PADDLE_ENFORCE_EQ(
-        Has(pir::kPlaceAttr),
+        Has(pir::Pass::kPlaceAttr),
         true,
         phi::errors::InvalidArgument(
             "Pass initialize failed."
             "When using ConstantFoldingPass, place attribute is required!"
             "Use Set method to set the place attribute."));
     PADDLE_ENFORCE_EQ(
-        Has(pir::kParamScopeAttr),
+        Has(pir::Pass::kParamScopeAttr),
         true,
         phi::errors::InvalidArgument(
             "Pass initialize failed."
             "When using ConstantFoldingPass, scope attribute is required!"
             "Use Set method to set the scope attribute."));
 
-    place_ = Get<phi::Place>(pir::kPlaceAttr);
-    scope_ = &Get<paddle::framework::Scope>(pir::kParamScopeAttr);
-
-    PADDLE_ENFORCE_NOT_NULL(
-        scope_, phi::errors::InvalidArgument("scope can not be nullptr"));
-    PADDLE_ENFORCE(
-        paddle::platform::is_gpu_place(place_) ||
-            paddle::platform::is_cpu_place(place_),
-        phi::errors::PreconditionNotMet(
-            "params_sync_among_devices_pass should run on cpu or gpu."));
+    place_ = Get<phi::Place>(pir::Pass::kPlaceAttr);
+    scope_ = &Get<paddle::framework::Scope>(pir::Pass::kParamScopeAttr);
     return true;
   }
 
@@ -106,11 +98,21 @@ class ParamsSyncAmongDevicesPass : public pir::Pass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
+    PADDLE_ENFORCE_NOT_NULL(
+        scope_, phi::errors::InvalidArgument("scope can not be nullptr"));
+    PADDLE_ENFORCE(paddle::platform::is_gpu_place(place_) ||
+                       paddle::platform::is_cpu_place(place_),
+                   phi::errors::PreconditionNotMet(
+                       "The Place attr in params_sync_among_devices_pass "
+                       "should be cpu or gpu."));
+    if (paddle::platform::is_cpu_place(place_)) {
+      return false;
+    }
     return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
   }
 
  private:
-  phi::Place place_;
+  phi::Place place_{phi::CPUPlace{}};
   paddle::framework::Scope* scope_{nullptr};
 };
 
diff --git a/paddle/fluid/pybind/inference_api.cc b/paddle/fluid/pybind/inference_api.cc
@@ -981,7 +981,8 @@ void BindAnalysisConfig(py::module *m) {
       .def("lite_engine_enabled", &AnalysisConfig::lite_engine_enabled)
       .def("switch_ir_debug",
            &AnalysisConfig::SwitchIrDebug,
-           py::arg("x") = true)
+           py::arg("x") = true,
+           py::arg("passes") = std::vector<std::string>())
       .def("enable_mkldnn", &AnalysisConfig::EnableMKLDNN)
       .def("disable_mkldnn", &AnalysisConfig::DisableMKLDNN)
       .def("mkldnn_enabled", &AnalysisConfig::mkldnn_enabled)
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1796,7 +1796,9 @@ void BindPassManager(pybind11::module *m) {
       .def("empty", &PassManager::empty)
       .def("clear", &PassManager::clear)
       .def("enable_ir_printing",
-           [](PassManager &self) { self.EnableIRPrinting(); });
+           [](PassManager &self) { self.EnableIRPrinting(); })
+      .def("enable_print_statistics",
+           [](PassManager &self) { self.EnablePrintStatistics(); });
 }
 
 void BindPir(pybind11::module *module) {
diff --git a/paddle/pir/include/pass/pass.h b/paddle/pir/include/pass/pass.h
@@ -71,12 +71,12 @@ struct PassInfo {
 
 }  // namespace detail
 
-static const char kParamScopeAttr[] = "__param_scope__";
-static const char kPlaceAttr[] = "__place__";
-
 /// We can access pass only from PassManager.
 class IR_API Pass {
  public:
+  inline static const char kParamScopeAttr[] = "__param_scope__";
+  inline static const char kPlaceAttr[] = "__place__";
+
   explicit Pass(const std::string& name,
                 uint8_t opt_level,
                 const std::vector<std::string>& dependents = {})
diff --git a/test/cpp/pir/pattern_rewrite/drr_attention_fuse_test.cc b/test/cpp/pir/pattern_rewrite/drr_attention_fuse_test.cc
@@ -153,8 +153,9 @@ TEST(DrrTest, AttentionFuse) {
   pm.AddPass(pir::CreateMultiHeadMatmulFusePass());
   std::unique_ptr<pir::Pass> constant_folding_pass =
       pir::CreateConstantFoldingPass();
-  constant_folding_pass->Set(pir::kPlaceAttr, new phi::Place{phi::GPUPlace{}});
-  constant_folding_pass->Set(pir::kParamScopeAttr,
+  constant_folding_pass->Set(pir::Pass::kPlaceAttr,
+                             new phi::Place{phi::GPUPlace{}});
+  constant_folding_pass->Set(pir::Pass::kParamScopeAttr,
                              new paddle::framework::Scope{});
   pm.AddPass(std::move(constant_folding_pass));
   pm.AddPass(pir::CreateDeadCodeEliminationPass());
diff --git a/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc b/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc
diff --git a/test/ir/pir/fused_pass/onednn/pass_test.py b/test/ir/pir/fused_pass/onednn/pass_test.py
diff --git a/test/ir/pir/fused_pass/pass_test.py b/test/ir/pir/fused_pass/pass_test.py
