[Inference] New executor support input hook and fix shape file collection in trt (#59466)

yuanlehome · web-flow · commit b9d91e7ef850 · 2023-11-30T09:42:19.000+08:00
* [Inference] new executor support input hook

* update

* update
diff --git a/paddle/fluid/framework/naive_executor.cc b/paddle/fluid/framework/naive_executor.cc
@@ -226,6 +226,9 @@ void NaiveExecutor::RegisterOutputHook(const HookFunc &hookfunc) {
 
 void NaiveExecutor::RegisterInputHook(const HookFunc &hookfunc) {
   input_hookfuncs_.push_back(hookfunc);
+  if (interpreter_core_) {
+    interpreter_core_->SetInputHooks(input_hookfuncs_);
+  }
 }
 
 void NaiveExecutor::MakeReusePlan(
diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
@@ -558,6 +558,8 @@ void BuildOpFuncList(const platform::Place& place,
                      std::vector<OpFuncNode>* vec_func_list,
                      VariableScope* var_scope,
                      const ExecutionConfig& execution_config,
+                     const std::vector<HookFunc>& input_hookfuncs,
+                     const std::vector<HookFunc>& output_hookfuncs,
                      bool use_local_scope,
                      bool static_build) {
   Scope* local_scope = use_local_scope ? var_scope->GetMutableLocalScope()
@@ -602,6 +604,14 @@ void BuildOpFuncList(const platform::Place& place,
       if (op_type == "feed" || op_type == "fetch") {
         continue;
       }
+      for (auto& hook : input_hookfuncs) {
+        hook(op, local_scope);
+      }
+
+      if (op->Type() == "while") {
+        op->SetInputHooks(input_hookfuncs);
+        op->SetOutputHooks(output_hookfuncs);
+      }
     }
 
     VLOG(6) << "Build OpFuncNode from : " << op_type;
@@ -980,6 +990,12 @@ void BuildOpFuncList(const platform::Place& place,
       }
     }
 
+    if (execution_config.used_for_inference) {
+      for (auto& hook : output_hookfuncs) {
+        hook(op, local_scope);
+      }
+    }
+
     VLOG(4) << "End run " << place << " "
             << op_func_node.operator_base_->DebugStringEx(local_scope);
 
diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.h b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.h
@@ -101,6 +101,8 @@ void BuildOpFuncList(const platform::Place& place,
                      std::vector<OpFuncNode>* vec_func_list,
                      VariableScope* scope,
                      const ExecutionConfig& execution_config,
+                     const std::vector<HookFunc>& input_hookfuncs,
+                     const std::vector<HookFunc>& output_hookfuncs,
                      bool use_local_scope = true,
                      bool static_build = false);
 
diff --git a/paddle/fluid/framework/new_executor/interpreter_base_impl.h b/paddle/fluid/framework/new_executor/interpreter_base_impl.h
@@ -57,7 +57,6 @@ constexpr const char* kTaskCompletion = "TaskCompletion";
 
 namespace paddle {
 namespace framework {
-using HookFunc = std::function<void(OperatorBase*, Scope*)>;
 
 /// @brief InterpreterBaseImpl is a abstract Base Class and define necessary
 /// interface with virtual keywords for Derived class.
@@ -97,6 +96,8 @@ class InterpreterBaseImpl {
 
   virtual void SetOutputHooks(const std::vector<HookFunc>& hookfuncs) = 0;
 
+  virtual void SetInputHooks(const std::vector<HookFunc>& hookfuncs) = 0;
+
   virtual std::shared_ptr<std::vector<size_t>> GetDependencyCount() const = 0;
 
   virtual bool IsSharedResultsBuild() const = 0;
diff --git a/paddle/fluid/framework/new_executor/interpretercore.cc b/paddle/fluid/framework/new_executor/interpretercore.cc
@@ -119,6 +119,10 @@ const platform::Place& InterpreterCore::GetPlace() const {
   return impl_->GetPlace();
 }
 
+void InterpreterCore::SetInputHooks(const std::vector<HookFunc>& hookfuncs) {
+  impl_->SetInputHooks(hookfuncs);
+}
+
 void InterpreterCore::SetOutputHooks(const std::vector<HookFunc>& hookfuncs) {
   impl_->SetOutputHooks(hookfuncs);
 }
diff --git a/paddle/fluid/framework/new_executor/interpretercore.h b/paddle/fluid/framework/new_executor/interpretercore.h
@@ -76,6 +76,8 @@ class InterpreterCore {
 
   void SetOutputHooks(const std::vector<HookFunc>& hookfuncs);
 
+  void SetInputHooks(const std::vector<HookFunc>& hookfuncs);
+
   void Build(const std::vector<std::string>& feed_names,
              std::vector<paddle::framework::OpFuncNode>* op_func_nodes);
 
diff --git a/paddle/fluid/framework/new_executor/new_executor_defs.h b/paddle/fluid/framework/new_executor/new_executor_defs.h
@@ -42,6 +42,8 @@ namespace framework {
 
 using OpKernelComputeFunc = std::function<void(const ExecutionContext&)>;
 
+using HookFunc = std::function<void(OperatorBase*, Scope*)>;
+
 using SchedulingPriority = int64_t;
 
 constexpr const char* kCoalesceTensor = "coalesce_tensor";
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.h b/paddle/fluid/framework/new_executor/pir_interpreter.h
@@ -88,7 +88,11 @@ class PirInterpreter : public InterpreterBaseImpl {
   const platform::Place& GetPlace() const override { return place_; }
 
   void SetOutputHooks(const std::vector<HookFunc>& hookfuncs) override {
-    hookfuncs_ = hookfuncs;
+    output_hookfuncs_ = hookfuncs;
+  }
+
+  void SetInputHooks(const std::vector<HookFunc>& hookfuncs) override {
+    input_hookfuncs_ = hookfuncs;
   }
 
   std::string GetNameByValue(::pir::Value value) const;
@@ -168,7 +172,8 @@ class PirInterpreter : public InterpreterBaseImpl {
   int64_t nccl_op_num_{-1};
   std::vector<size_t> trace_execute_order_;
 
-  std::vector<HookFunc> hookfuncs_;
+  std::vector<HookFunc> output_hookfuncs_;
+  std::vector<HookFunc> input_hookfuncs_;
 
   /// ======================== ///
   ///        For new ir        ///
diff --git a/paddle/fluid/framework/new_executor/program_interpreter.cc b/paddle/fluid/framework/new_executor/program_interpreter.cc
@@ -210,6 +210,8 @@ void ProgramInterpreter::Build(
         op_func_nodes,
         &var_scope_,
         execution_config_,
+        input_hookfuncs_,
+        output_hookfuncs_,
         HasLocalScope(),
         static_build_);
   }
@@ -893,8 +895,15 @@ void ProgramInterpreter::RunOperator(const Instruction& instr_node) {
                                        : var_scope_.GetMutableScope();
   VLOG(4) << "Start run " << place << " " << op->DebugStringEx(local_scope);
 
-  if (op->Type() == "while") {
-    op->SetOutputHooks(hookfuncs_);
+  if (execution_config_.used_for_inference) {
+    for (auto& hook : input_hookfuncs_) {
+      hook(op, local_scope);
+    }
+
+    if (op->Type() == "while") {
+      op->SetInputHooks(input_hookfuncs_);
+      op->SetOutputHooks(output_hookfuncs_);
+    }
   }
 
   auto op_with_kernel = dynamic_cast<const framework::OperatorWithKernel*>(op);
@@ -1002,8 +1011,10 @@ void ProgramInterpreter::RunOperator(const Instruction& instr_node) {
 #endif
   }
 
-  for (auto& hook : hookfuncs_) {
-    hook(op, local_scope);
+  if (execution_config_.used_for_inference) {
+    for (auto& hook : output_hookfuncs_) {
+      hook(op, local_scope);
+    }
   }
 
   // for debug nan/inf
@@ -1440,6 +1451,8 @@ void ProgramInterpreter::Prepare(
         &op_func_nodes,
         &var_scope_,
         execution_config_,
+        input_hookfuncs_,
+        output_hookfuncs_,
         HasLocalScope(),
         static_build_);
     SetFeedVarsInplaceSkip(feed_names);
diff --git a/paddle/fluid/framework/new_executor/program_interpreter.h b/paddle/fluid/framework/new_executor/program_interpreter.h
@@ -89,7 +89,11 @@ class ProgramInterpreter : public InterpreterBaseImpl {
   const platform::Place& GetPlace() const override { return place_; }
 
   void SetOutputHooks(const std::vector<HookFunc>& hookfuncs) override {
-    hookfuncs_ = hookfuncs;
+    output_hookfuncs_ = hookfuncs;
+  }
+
+  void SetInputHooks(const std::vector<HookFunc>& hookfuncs) override {
+    input_hookfuncs_ = hookfuncs;
   }
 
   std::unordered_map<std::string, std::shared_ptr<EventInter>>*
@@ -218,7 +222,8 @@ class ProgramInterpreter : public InterpreterBaseImpl {
 
   InstructionSchedulingPriorityLess instruction_scheduling_priority_less;
 
-  std::vector<HookFunc> hookfuncs_;
+  std::vector<HookFunc> output_hookfuncs_;
+  std::vector<HookFunc> input_hookfuncs_;
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   std::unique_ptr<phi::CalculateStreamTimer> calculate_stream_timer_;
diff --git a/paddle/fluid/framework/operator.h b/paddle/fluid/framework/operator.h
@@ -378,7 +378,10 @@ class TEST_API OperatorBase {
 
   using HookFunc = std::function<void(OperatorBase*, Scope*)>;
   void SetOutputHooks(const std::vector<HookFunc>& hookfuncs) {
-    hookfuncs_ = hookfuncs;
+    output_hookfuncs_ = hookfuncs;
+  }
+  void SetInputHooks(const std::vector<HookFunc>& hookfuncs) {
+    input_hookfuncs_ = hookfuncs;
   }
 
  protected:
@@ -409,7 +412,8 @@ class TEST_API OperatorBase {
   // Whether this operator executes in an Executor.
   bool run_by_executor_{true};
 
-  std::vector<HookFunc> hookfuncs_;
+  std::vector<HookFunc> output_hookfuncs_;
+  std::vector<HookFunc> input_hookfuncs_;
 
  private:
   void GenerateTemporaryNames();
diff --git a/paddle/fluid/operators/controlflow/while_op.cc b/paddle/fluid/operators/controlflow/while_op.cc
@@ -182,7 +182,8 @@ class WhileOp : public framework::OperatorBase {
           dev_place, *block, &placeholder, execution_config));
     }
 
-    core_->SetOutputHooks(hookfuncs_);
+    core_->SetOutputHooks(output_hookfuncs_);
+    core_->SetInputHooks(input_hookfuncs_);
 
     if (!is_test) {
       while (cond_data) {

Original file line number	Diff line number	Diff line change
`@@ -226,6 +226,9 @@ void NaiveExecutor::RegisterOutputHook(const HookFunc &hookfunc) {`
`226`	`226`
`227`	`227`	`void NaiveExecutor::RegisterInputHook(const HookFunc &hookfunc) {`
`228`	`228`	`input_hookfuncs_.push_back(hookfunc);`
	`229`	`+ if (interpreter_core_) {`
	`230`	`+ interpreter_core_->SetInputHooks(input_hookfuncs_);`
	`231`	`+ }`
`229`	`232`	`}`
`230`	`233`
`231`	`234`	`void NaiveExecutor::MakeReusePlan(`
Original file line number	Diff line number	Diff line change
`@@ -119,6 +119,10 @@ const platform::Place& InterpreterCore::GetPlace() const {`
`119`	`119`	`return impl_->GetPlace();`
`120`	`120`	`}`
`121`	`121`
	`122`	`+void InterpreterCore::SetInputHooks(const std::vector<HookFunc>& hookfuncs) {`
	`123`	`+ impl_->SetInputHooks(hookfuncs);`
	`124`	`+}`
	`125`	`+`
`122`	`126`	`void InterpreterCore::SetOutputHooks(const std::vector<HookFunc>& hookfuncs) {`
`123`	`127`	`impl_->SetOutputHooks(hookfuncs);`
`124`	`128`	`}`