[Inference] execution_config.used_for_inference set to true in while op and refine some code (#60203)

yuanlehome · web-flow · commit c22d840f6b21 · 2023-12-21T21:20:43.000+08:00
* fix some code

* update

* rewrite dce pass

* update

* update
diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
@@ -638,6 +638,9 @@ void BuildOpFuncList(const platform::Place& place,
       if (op->Type() == "while") {
         op->SetInputHooks(input_hookfuncs);
         op->SetOutputHooks(output_hookfuncs);
+        auto runtime_attrs = op->RuntimeAttrs();
+        runtime_attrs.insert(std::make_pair("used_for_inference", true));
+        op->SetRuntimeAttributeMap(runtime_attrs);
       }
     }
 
diff --git a/paddle/fluid/framework/new_executor/program_interpreter.cc b/paddle/fluid/framework/new_executor/program_interpreter.cc
@@ -918,6 +918,9 @@ void ProgramInterpreter::RunOperator(const Instruction& instr_node) {
     if (op->Type() == "while") {
       op->SetInputHooks(input_hookfuncs_);
       op->SetOutputHooks(output_hookfuncs_);
+      auto runtime_attrs = op->RuntimeAttrs();
+      runtime_attrs.insert(std::make_pair("used_for_inference", true));
+      op->SetRuntimeAttributeMap(runtime_attrs);
     }
   }
 
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -1931,6 +1931,9 @@ CreatePaddlePredictor<AnalysisConfig, PaddleEngineKind::kAnalysis>(
         if (std::getenv("FLAGS_initial_cpu_memory_in_mb") == nullptr) {
           SetGflag("initial_cpu_memory_in_mb", "0");
         }
+        if (std::getenv("FLAGS_cache_inference_while_scope") == nullptr) {
+          SetGflag("cache_inference_while_scope", "1");
+        }
       });
 
       if (config.thread_local_stream_enabled() &&
diff --git a/paddle/fluid/operators/controlflow/while_op.cc b/paddle/fluid/operators/controlflow/while_op.cc
@@ -173,6 +173,9 @@ class WhileOp : public framework::OperatorBase {
       framework::Scope placeholder;  // Don't care if it's valid, just for
                                      // initialize InterpreterCore
       framework::interpreter::ExecutionConfig execution_config;
+      if (HasAttr("used_for_inference") && Attr<bool>("used_for_inference")) {
+        execution_config.used_for_inference = true;
+      }
       execution_config.create_local_scope = false;
       execution_config.used_for_control_flow_op = true;
       execution_config.skip_gc_vars =
diff --git a/paddle/fluid/pir/drr/drr_rewrite_pattern.cc b/paddle/fluid/pir/drr/drr_rewrite_pattern.cc
@@ -267,12 +267,18 @@ bool DrrRewritePattern::MatchFromOutputToInput(
       source_pattern_match_ctx->BindIrValue(
           drr_input_tensors[i]->name(),
           std::make_shared<IrValue>(ir_node->operand(i).source()));
+      if (ir_node->operand_source(i).isa<pir::BlockArgument>()) {
+        matched = false;
+        VLOG(8) << drr_node->name()
+                << " Match failed: it's input value is a block argument.";
+        break;
+      }
+
       auto* drr_producer_op = drr_input_tensors[i]->producer();
       if (drr_producer_op == nullptr) {
         continue;
       }
-      auto* ir_producer_op =
-          ir_node->operand(i).source().dyn_cast<pir::OpResult>().owner();
+
       if (drr_input_tensors[i]->consumers().size() !=
           ir_node->operand(i).source().use_count()) {
         matched = false;
@@ -282,11 +288,14 @@ bool DrrRewritePattern::MatchFromOutputToInput(
                 << ir_node->operand(i).source().use_count() << " }.";
         break;
       }
+
+      auto* ir_producer_op = ir_node->operand_source(i).defining_op();
       // bfs producer_op of current_op
       if (drr_visited.count(drr_producer_op) &&
           ir_visited.count(ir_producer_op)) {
         continue;
       }
+
       if (!drr_visited.count(drr_producer_op) &&
           !ir_visited.count(ir_producer_op)) {
         drr_q.push(drr_producer_op);
@@ -396,7 +405,10 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
       auto ir_val = res_match_ctx.GetIrValue(input->name());
       if (ir_val) {
         Operation* ir_input_op = ir_val.dyn_cast<pir::OpResult>().owner();
-        if (max_input_op_index < op_2_temp_program_index.at(ir_input_op)) {
+        if (op_2_temp_program_index.count(ir_input_op) == 0) {
+          max_input_op_index = 0UL;
+        } 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);
           max_index_op = ir_input_op;
         } else if (max_input_op_index ==
diff --git a/paddle/fluid/pir/transforms/constant_folding_pass.cc b/paddle/fluid/pir/transforms/constant_folding_pass.cc
@@ -83,8 +83,9 @@ class ConstantFoldingPattern : public pir::RewritePattern {
         continue;
       }
       // 2. inputs must come from parameter op or constant op
-      if (!(pir::GetDefiningOpForInput(op, i)->isa<pir::ParameterOp>() ||
-            pir::GetDefiningOpForInput(op, i)->isa<pir::ConstantTensorOp>())) {
+      auto* prev_op = pir::GetDefiningOpForInput(op, i);
+      if (!prev_op || !(prev_op->isa<pir::ParameterOp>() ||
+                        prev_op->isa<pir::ConstantTensorOp>())) {
         return false;
       }
       // 3. inputs must be a dense tensor type
@@ -140,7 +141,9 @@ class ConstantFoldingPattern : public pir::RewritePattern {
 
     core.Run({});
 
-    rewriter.SetInsertionPointToStart(rewriter.block());
+    // ParameterOp and ConstantTensorOp should be created in the top-level block
+    rewriter.SetInsertionPointToStart(
+        rewriter.block()->parent_program()->block());
 
     bool use_parameter_op = ReplaceResultByParameterOp(op);
 
@@ -168,14 +171,15 @@ class ConstantFoldingPattern : public pir::RewritePattern {
             paddle::framework::TensorCopySync(
                 temp_tensor, place_, output_tensor);
           }
-          auto parameter_op = rewriter.Build<pir::ParameterOp>(
-              output_var_name, op->result(i).type());
-          parameter_op->set_attribute(
-              kAttrIsPersisable,
-              rewriter.array_attr({rewriter.bool_attr(true)}));
-
-          rewriter.ReplaceAllUsesWith(op->result(i), parameter_op->result(0));
         }
+
+        auto parameter_op = rewriter.Build<pir::ParameterOp>(
+            output_var_name, op->result(i).type());
+        parameter_op->set_attribute(
+            kAttrIsPersisable, rewriter.array_attr({rewriter.bool_attr(true)}));
+
+        rewriter.ReplaceAllUsesWith(op->result(i), parameter_op->result(0));
+
       } else {
         if (output_var->IsType<phi::DenseTensor>()) {
           auto* output_tensor = output_var->GetMutable<phi::DenseTensor>();
@@ -336,11 +340,6 @@ class ConstantFoldingPass : public pir::Pass {
     paddle::memory::Release(phi::CPUPlace{});
   }
 
-  bool CanApplyOn(pir::Operation* op) const override {
-    // TODO(liuyuanle): remove op->isa<::pir::ModuleOp>()
-    return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
-  }
-
  private:
   size_t counter_{0};
   phi::Place place_;
diff --git a/paddle/fluid/pir/transforms/dead_code_elimination_pass.cc b/paddle/fluid/pir/transforms/dead_code_elimination_pass.cc
@@ -15,63 +15,50 @@
 #include "paddle/fluid/pir/transforms/dead_code_elimination_pass.h"
 
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/pir/core/block.h"
 #include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/core/op_trait.h"
-#include "paddle/pir/core/program.h"
-#include "paddle/pir/dialect/control_flow/ir/cf_op.h"
 #include "paddle/pir/pass/pass.h"
 #include "paddle/pir/pass/pass_registry.h"
-#include "paddle/pir/pattern_rewrite/frozen_rewrite_pattern_set.h"
-#include "paddle/pir/pattern_rewrite/pattern_match.h"
-#include "paddle/pir/pattern_rewrite/pattern_rewrite_driver.h"
 
 namespace {
 
-class DeadCodeEliminationPattern : public pir::RewritePattern {
- public:
-  DeadCodeEliminationPattern(
-      pir::IrContext* context,
-      pir::PatternBenefit benefit = 1,
-      const std::vector<std::string>& generated_names = {})
-      : RewritePattern(MatchAnyOpTypeTag(), benefit, context, generated_names) {
-  }
-
-  bool Match(pir::Operation* op) const override {
-    if (op->HasTrait<pir::SideEffectTrait>()) return false;
-
-    if (op->isa<paddle::dialect::DataOp>()) {
-      return false;
-    }
-    return op->use_empty();
-  }
-
-  void Rewrite(pir::Operation* op,
-               pir::PatternRewriter& rewriter) const override {  // NOLINT
-    rewriter.EraseOp(op);
-  }
-};
-
 class DeadCodeEliminationPass : public pir::Pass {
  public:
   DeadCodeEliminationPass() : pir::Pass("dead_code_elimination_pass", 0) {}
 
-  bool Initialize(pir::IrContext* context) override {
-    pir::RewritePatternSet ps(context);
-    ps.Add<DeadCodeEliminationPattern>(context);
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
-  }
-
   void Run(pir::Operation* op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    auto [_, num_rewrites] = pir::ApplyPatternsGreedily(op, patterns_, cfg);
-    PrintStatistics(num_rewrites);
+    VLOG(6) << "apply dead_code_elimination_pass";
+    int64_t num_erasers{0};
+    EraseOp(*op->GetParentProgram()->block(), &num_erasers);
+    PrintStatistics(num_erasers);
   }
 
  private:
-  pir::FrozenRewritePatternSet patterns_;
+  void EraseOp(const pir::Block& block, int64_t* num_erasers) {
+    std::vector<pir::Operation*> deleted_ops;
+    for (auto& op : block) {
+      if (op.HasTrait<pir::SideEffectTrait>() ||
+          op.isa<paddle::dialect::DataOp>()) {
+        continue;
+      }
+      if (op.use_empty()) {
+        deleted_ops.push_back(&op);
+      }
+    }
+    for (auto* op : deleted_ops) {
+      op->Erase();
+      (*num_erasers)++;
+    }
+    for (auto& op : block) {
+      for (size_t i = 0; i < op.num_regions(); ++i) {
+        auto& inner_region = op.region(i);
+        for (auto& inner_block : inner_region) {
+          EraseOp(inner_block, num_erasers);
+        }
+      }
+    }
+  }
 };
 
 }  // namespace
diff --git a/paddle/fluid/pir/transforms/params_sync_among_devices_pass.cc b/paddle/fluid/pir/transforms/params_sync_among_devices_pass.cc
@@ -46,7 +46,7 @@ class ParamsSyncAmongDevicesPass : public pir::Pass {
   }
 
   void Run(pir::Operation* op) override {
-    VLOG(6) << "apply ParamsSyncAmongDevicesPass";
+    VLOG(6) << "apply params_sync_among_devices_pass";
     auto module_op = op->dyn_cast<pir::ModuleOp>();
     PADDLE_ENFORCE_NOT_NULL(
         module_op,
diff --git a/paddle/pir/pass/pass.cc b/paddle/pir/pass/pass.cc
@@ -35,16 +35,24 @@ Pass::~Pass() = default;
 bool Pass::CanApplyOn(Operation* op) const { return op->num_regions() > 0; }
 
 void Pass::PrintStatistics(int64_t match_count) const {
-  LOG(INFO) << "--- detected [" << match_count << "] subgraphs!";
+  if (match_count > 0) {
+    LOG(INFO) << "--- detected [" << match_count << "] subgraphs!";
+  }
 }
 
 void Pass::PrintStatistics(int64_t match_count, int64_t all_count) const {
-  LOG(INFO) << "--- detected [" << match_count << "/" << all_count
-            << "] subgraphs!";
+  IR_ENFORCE(match_count < all_count,
+             "match_count should smaller than all_count");
+  if (match_count > 0) {
+    LOG(INFO) << "--- detected [" << match_count << "/" << all_count
+              << "] subgraphs!";
+  }
 }
 
 void Pass::PrintStatistics(const std::string& custom_log) const {
-  LOG(INFO) << custom_log;
+  if (!custom_log.empty()) {
+    LOG(INFO) << custom_log;
+  }
 }
 
 detail::PassExecutionState& Pass::pass_state() {
diff --git a/test/cpp/pir/pattern_rewrite/drr_fuse_linear_param_grad_add_test.cc b/test/cpp/pir/pattern_rewrite/drr_fuse_linear_param_grad_add_test.cc
@@ -175,7 +175,7 @@ TEST(DrrTest, FusedLinearParamGradAdd0) {
 
   pir::PassManager pm(ctx);
   pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
@@ -194,7 +194,7 @@ TEST(DrrTest, FusedLinearParamGradAdd1) {
 
   pir::PassManager pm(ctx);
   pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
@@ -213,7 +213,7 @@ TEST(DrrTest, FusedLinearParamGradAdd2) {
 
   pir::PassManager pm(ctx);
   pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
@@ -232,7 +232,7 @@ TEST(DrrTest, FusedLinearParamGradAdd3) {
 
   pir::PassManager pm(ctx);
   pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
diff --git a/test/cpp/pir/pattern_rewrite/drr_fuse_linear_test.cc b/test/cpp/pir/pattern_rewrite/drr_fuse_linear_test.cc
@@ -136,7 +136,7 @@ TEST(DrrTest, FusedLinear) {
 
   pir::PassManager pm(ctx);
   pm.AddPass(pir::CreateFusedGemmEpiloguePass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
diff --git a/test/cpp/pir/pattern_rewrite/drr_same_type_binding_test.cc b/test/cpp/pir/pattern_rewrite/drr_same_type_binding_test.cc
@@ -310,7 +310,7 @@ TEST(DrrTest, drr_demo) {
   pir::PassManager pm(ctx);
   pm.AddPass(std::make_unique<DrrPatternRewritePass>());
   pm.AddPass(pir::CreateDeadCodeEliminationPass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
diff --git a/test/cpp/pir/pattern_rewrite/drr_test.cc b/test/cpp/pir/pattern_rewrite/drr_test.cc
@@ -215,7 +215,7 @@ TEST(DrrTest, drr_demo) {
   pir::PassManager pm(ctx);
   pm.AddPass(std::make_unique<DrrPatternRewritePass>());
   pm.AddPass(pir::CreateDeadCodeEliminationPass());
-  // pm.EnablePassTiming();
+  pm.EnablePassTiming();
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);

Original file line number	Diff line number	Diff line change
`@@ -638,6 +638,9 @@ void BuildOpFuncList(const platform::Place& place,`
`638`	`638`	`if (op->Type() == "while") {`
`639`	`639`	`op->SetInputHooks(input_hookfuncs);`
`640`	`640`	`op->SetOutputHooks(output_hookfuncs);`
	`641`	`+ auto runtime_attrs = op->RuntimeAttrs();`
	`642`	`+ runtime_attrs.insert(std::make_pair("used_for_inference", true));`
	`643`	`+ op->SetRuntimeAttributeMap(runtime_attrs);`
`641`	`644`	`}`
`642`	`645`	`}`
`643`	`646`
Original file line number	Diff line number	Diff line change
`@@ -918,6 +918,9 @@ void ProgramInterpreter::RunOperator(const Instruction& instr_node) {`
`918`	`918`	`if (op->Type() == "while") {`
`919`	`919`	`op->SetInputHooks(input_hookfuncs_);`
`920`	`920`	`op->SetOutputHooks(output_hookfuncs_);`
	`921`	`+ auto runtime_attrs = op->RuntimeAttrs();`
	`922`	`+ runtime_attrs.insert(std::make_pair("used_for_inference", true));`
	`923`	`+ op->SetRuntimeAttributeMap(runtime_attrs);`
`921`	`924`	`}`
`922`	`925`	`}`
`923`	`926`
Original file line number	Diff line number	Diff line change
`@@ -1931,6 +1931,9 @@ CreatePaddlePredictor<AnalysisConfig, PaddleEngineKind::kAnalysis>(`
`1931`	`1931`	`if (std::getenv("FLAGS_initial_cpu_memory_in_mb") == nullptr) {`
`1932`	`1932`	`SetGflag("initial_cpu_memory_in_mb", "0");`
`1933`	`1933`	`}`
	`1934`	`+ if (std::getenv("FLAGS_cache_inference_while_scope") == nullptr) {`
	`1935`	`+ SetGflag("cache_inference_while_scope", "1");`
	`1936`	`+ }`
`1934`	`1937`	`});`
`1935`	`1938`
`1936`	`1939`	`if (config.thread_local_stream_enabled() &&`
Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ class ParamsSyncAmongDevicesPass : public pir::Pass {`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`void Run(pir::Operation* op) override {`
`49`		`- VLOG(6) << "apply ParamsSyncAmongDevicesPass";`
	`49`	`+ VLOG(6) << "apply params_sync_among_devices_pass";`
`50`	`50`	`auto module_op = op->dyn_cast<pir::ModuleOp>();`
`51`	`51`	`PADDLE_ENFORCE_NOT_NULL(`
`52`	`52`	`module_op,`