From ca9045f7f05ef75ae77858ceb75f56cd01da4de7 Mon Sep 17 00:00:00 2001
From: kangguangli <kangguangli@hotmail.com>
Date: Thu, 7 Sep 2023 15:04:07 +0800
Subject: [PATCH] [NewIR] optimize GetNameById in check_gc (#56964)

* optimize GetNameById in check_gc

* remove debug code

* fix

* fix
---
 .../new_executor/interpreter/interpreter_util.cc     | 10 ++++++++++
 .../new_executor/interpreter/interpreter_util.h      |  2 ++
 .../framework/new_executor/new_ir_interpreter.cc     | 12 +++++++-----
 .../framework/new_executor/new_ir_interpreter.h      |  1 +
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
index be87549ecbf..9d114ed0c19 100644
--- a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
+++ b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
@@ -1218,6 +1218,16 @@ std::unordered_set<std::string> GetSpecialOpNames() {
   };
 }
 
+void BuildId2VarName(const std::map<std::string, int>& var_name_2_id,
+                     std::unordered_map<int, std::string>* id_2_var_name) {
+  PADDLE_ENFORCE_NOT_NULL(
+      id_2_var_name,
+      phi::errors::InvalidArgument("id2_var_name can not be null"));
+  for (auto [var_name, id] : var_name_2_id) {
+    id_2_var_name->insert({id, var_name});
+  }
+}
+
 }  // namespace interpreter
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.h b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.h
index 186f9459fba..33b89cac542 100644
--- a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.h
+++ b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.h
@@ -116,6 +116,8 @@ void BuildOpFuncList(
 void BuildVariableScope(const framework::BlockDesc& block,
                         const ExecutionConfig& execution_config,
                         VariableScope* var_scope);
+void BuildId2VarName(const std::map<std::string, int>& var_name_2_id,
+                     std::unordered_map<int, std::string>* id_2_var_name);
 
 void LogDeviceMemoryStats(const platform::Place& place);
 
diff --git a/paddle/fluid/framework/new_executor/new_ir_interpreter.cc b/paddle/fluid/framework/new_executor/new_ir_interpreter.cc
index e026e914adb..37cecdad130 100644
--- a/paddle/fluid/framework/new_executor/new_ir_interpreter.cc
+++ b/paddle/fluid/framework/new_executor/new_ir_interpreter.cc
@@ -194,11 +194,10 @@ std::string NewIRInterpreter::GetNameById(int id) const {
   // typically when the target variable is not existed in the original program
   // desc, but created by interpretercore.
   // For example, created and used by d2h_copy or h2d_copy operator.
-  auto it = std::find_if(var_name_2_id_.begin(),
-                         var_name_2_id_.end(),
-                         [id](const auto& pair) { return pair.second == id; });
-  if (it != var_name_2_id_.end()) {
-    return it->first;
+
+  auto it = id_2_var_name_.find(id);
+  if (it != id_2_var_name_.end()) {
+    return it->second;
   }
   return "";
 }
@@ -925,6 +924,9 @@ FetchList NewIRInterpreter::Run(const std::vector<std::string>& feed_names,
                      &variable_2_var_name_,
                      &var_name_2_id_,
                      &variable_list_);
+
+    interpreter::BuildId2VarName(var_name_2_id_, &id_2_var_name_);
+
     VLOG(4) << "Done BuildScope";
     VLOG(4) << DebugValueInfo();
 
diff --git a/paddle/fluid/framework/new_executor/new_ir_interpreter.h b/paddle/fluid/framework/new_executor/new_ir_interpreter.h
index 841e9136a2e..b37b26d1075 100644
--- a/paddle/fluid/framework/new_executor/new_ir_interpreter.h
+++ b/paddle/fluid/framework/new_executor/new_ir_interpreter.h
@@ -208,6 +208,7 @@ class NewIRInterpreter : public InterpreterBaseImpl {
       variable_2_var_name_;
 
   std::map<std::string, int> var_name_2_id_;
+  std::unordered_map<int, std::string> id_2_var_name_;
 
   std::vector<Variable*> variable_list_;
 
-- 
GitLab