[PIR] Refactor pir GetNameMap and assert value unique (#66738)

gouzil · SigureMo · web-flow · commit 3454eccbda1c · 2024-08-01T18:32:00.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -475,7 +475,6 @@ inline void PirRunProgramAPI(
     // no_need_buffers + outputs)
     std::vector<std::string> skip_names;
     // update interpretercore skip_gc_var
-    std::vector<pir::Value> kwargs_values;
     for (auto &kwarg : backward_program->block()->kwargs()) {
       skip_names.push_back(kwarg.first);
     }
diff --git a/paddle/fluid/pir/utils/name_analysis.h b/paddle/fluid/pir/utils/name_analysis.h
@@ -14,6 +14,7 @@
 
 #pragma once
 #include <optional>
+#include <set>
 #include <string>
 #include <vector>
 #include "paddle/pir/include/core/program.h"
@@ -39,6 +40,11 @@ std::vector<std::string> GetValueOutputNames(pir::Value value);
 pir::Value GetOutputValueByName(const pir::Program &program,
                                 const std::string &name);
 
+inline bool HasOnlyOneValueName(pir::Value value) {
+  std::vector<std::string> names = GetValueAllNames(value);
+  return std::set<std::string>(names.begin(), names.end()).size() <= 1;
+}
+
 }  // namespace name_analysis
 }  // namespace utils
 }  // namespace pir
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1350,6 +1350,10 @@ void BindValue(py::module *m) {
       .def("is_same", &Value::operator==)
       .def("hash", [](Value self) { return std::hash<pir::Value>{}(self); })
       .def("_rename", &name_analysis::RenameValue)
+      .def("_has_only_one_name",
+           [](Value self) -> bool {
+             return name_analysis::HasOnlyOneValueName(self);
+           })
       .def("detach",
            [](Value self) {
              auto share_data_op =
@@ -1795,43 +1799,6 @@ int AppendShadowOutputs(Program *program,
   return counter;
 }
 
-std::unordered_map<::pir::Value, std::string> GetNameMap(
-    const ::pir::Block *block) {
-  std::unordered_map<::pir::Value, std::string> value2name;
-  for (auto &kwarg : block->kwargs()) {
-    value2name[kwarg.second] = kwarg.first;
-  }
-  for (auto &op : *block) {
-    std::string name;
-    if (op.name() == "pd_op.data") {
-      name = op.attributes().at("name").dyn_cast<StrAttribute>().AsString();
-      value2name[op.results()[0].Value::impl()] = name;
-    } else if (op.name() == "builtin.set_parameter") {
-      name = op.attributes()
-                 .at("parameter_name")
-                 .dyn_cast<StrAttribute>()
-                 .AsString();
-      value2name[op.operand(0).source()] = name;
-    } else if (op.name() == "builtin.shadow_output") {
-      name =
-          op.attributes().at("output_name").dyn_cast<StrAttribute>().AsString();
-      value2name[op.operand(0).source()] = name;
-    } else if (op.name() == "builtin.parameter") {
-      name = op.attributes()
-                 .at("parameter_name")
-                 .dyn_cast<StrAttribute>()
-                 .AsString();
-      value2name[op.result(0).Value::impl()] = name;
-    } else if (op.name() == "builtin.constant") {
-      if (op.isa<pir::ConstantTensorOp>()) {
-        name = op.dyn_cast<pir::ConstantTensorOp>().tensor_name();
-        value2name[op.result(0).Value::impl()] = name;
-      }
-    }
-  }
-  return value2name;
-}
-
 SplitedResult SplitForwardBackward(
     const Program &program,
     const std::vector<pir::Value> &forward_inputs,
@@ -1919,22 +1886,16 @@ SplitedResult SplitForwardBackward(
   std::for_each(
       forward_outputs.begin(), forward_outputs.end(), create_output_fn_forward);
 
-  pir::Block *forward_block = forward_program->block();
-  const auto &forward_name_map = GetNameMap(forward_block);
   auto create_kwarg_fn = [&backward_block,
                           &backward_inputs,
                           &backward_value_map,
                           &forward_value_map,
-                          &forward_name_map,
                           &counter](const pir::Value &v) {
     if (v && !backward_value_map.count(v) && (backward_inputs.count(v))) {
-      auto forward_value = forward_value_map[v];
-      std::string name = "input_" + std::to_string(counter++);
-      if (forward_name_map.count(forward_value)) {
-        name = forward_name_map.at(forward_value);
-      }
-
-      backward_value_map[v] = backward_block.AddKwarg(name, v.type());
+      backward_value_map[v] = backward_block.AddKwarg(
+          name_analysis::TryGetValueFirstName(forward_value_map[v])
+              .value_or("input_" + std::to_string(counter++)),
+          v.type());
     }
   };
 
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -342,6 +342,12 @@ def unify_value_names(
             rename_mapping.update(
                 value._rename(new_name, program.global_block())
             )
+        # Get all values again because some values has been erased.
+        for value in RunnableProgram._get_program_all_values(program):
+            if value.has_name:
+                assert (
+                    value._has_only_one_name()
+                ), f"Expected all values in Program have only one name, but {value} has multiple names: {value._names}"
         return rename_mapping
 
     @cached_property

Original file line number	Diff line number	Diff line change
`@@ -475,7 +475,6 @@ inline void PirRunProgramAPI(`
`475`	`475`	`// no_need_buffers + outputs)`
`476`	`476`	`std::vector<std::string> skip_names;`
`477`	`477`	`// update interpretercore skip_gc_var`
`478`		`- std::vector<pir::Value> kwargs_values;`
`479`	`478`	`for (auto &kwarg : backward_program->block()->kwargs()) {`
`480`	`479`	`skip_names.push_back(kwarg.first);`
`481`	`480`	`}`