[PIR] refactor pir GetValueName - part 1 (#66363)

gouzil · web-flow · commit 8036ef3ac2ad · 2024-07-25T15:03:11.000+08:00
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -1978,8 +1978,8 @@ paddle::Tensor CreateTensorFromValue(const pir::Value& value) {
 
   auto dims = phi::vectorize(GetValueDims(value));
   auto ddims = phi::make_ddim(dims);
-  if (HasValueName(value)) {
-    tensor.set_name(GetValueName(value));
+  if (auto name = name_analysis::TryGetValueFirstName(value)) {
+    tensor.set_name(name.value());
   }
   auto autograd_meta = egr::EagerUtils::autograd_meta(&tensor);
   autograd_meta->SetPersistable(false);
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -23,7 +23,9 @@
 #include <unordered_map>
 #include <unordered_set>
 #include <utility>
+#include <vector>
 
+#include "paddle/common/enforce.h"
 #include "paddle/common/flags.h"
 #include "paddle/fluid/framework/executor.h"
 #include "paddle/fluid/framework/ir/pass.h"
@@ -198,6 +200,7 @@ std::string GetValueInfo(Value v) {
   return ss.str();
 }
 
+namespace name_analysis {
 Value GetOutputValueByName(const Program &program, const std::string &name) {
   auto &block = *program.block();
   StrAttribute name_attr = StrAttribute::get(IrContext::Instance(), name);
@@ -234,7 +237,7 @@ Value GetParameterValueByName(const Program &program, const std::string &name) {
   return value;
 }
 
-void SetValueName(Value value, const std::string name) {
+void SetValueAllNamesWith(Value value, const std::string name) {
   pir::Operation *define_op = value.defining_op();
   if (define_op->isa<pir::ParameterOp>()) {
     define_op->set_attribute(
@@ -262,71 +265,77 @@ void SetValueName(Value value, const std::string name) {
   }
 }
 
-bool IsUsedByShadowOutput(const pir::Value &value) {
-  for (auto iter = value.use_begin(); iter != value.use_end(); ++iter) {
-    if (iter->owner()->isa<::pir::ShadowOutputOp>()) {
-      return true;
-    }
-  }
-  return false;
-}
-
-bool HasValueName(const Value &value) {
-  if (IsFakeValue(value)) {
-    return false;
-  }
-  if (value.defining_op()->isa<::pir::ParameterOp>() ||
-      value.defining_op()->isa<paddle::dialect::DataOp>() ||
-      value.isa<BlockArgument>() || IsUsedByShadowOutput(value)) {
-    return true;
-  } else {
-    return false;
-  }
-}
-
 std::optional<std::string> GetValueInputName(Value value) {
   std::optional<std::string> name;
-  if (auto param_op = value.defining_op<::pir::ParameterOp>()) {
-    name = param_op.param_name();
-  } else if (auto data_op = value.defining_op<paddle::dialect::DataOp>()) {
-    name = data_op.attribute<StrAttribute>("name").AsString();
-  } else if (auto block_arg = value.dyn_cast<BlockArgument>()) {
+  if (auto block_arg = value.dyn_cast<BlockArgument>()) {
     if (block_arg.is_kwarg()) {
       name = block_arg.keyword();
     } else {
       name = "arg_" + std::to_string(block_arg.index());
     }
+  } else if (auto param_op = value.defining_op<::pir::ParameterOp>()) {
+    name = param_op.param_name();
+  } else if (auto data_op = value.defining_op<paddle::dialect::DataOp>()) {
+    name = data_op.attribute<StrAttribute>("name").AsString();
+  } else if (auto constant_op = value.defining_op<::pir::ConstantTensorOp>()) {
+    name = constant_op.tensor_name();
   }
   return name;
 }
 
-std::optional<std::string> GetValueOutputName(Value value) {
-  std::optional<std::string> name;
-  if (IsUsedByShadowOutput(value)) {
-    for (auto iter = value.use_begin(); iter != value.use_end(); ++iter) {
-      if (iter->owner()->isa<::pir::ShadowOutputOp>()) {
-        name = iter->owner()->attribute<StrAttribute>("output_name").AsString();
-        break;
-      }
+std::vector<std::string> GetValueOutputNames(Value value) {
+  std::vector<std::string> names;
+  for (auto iter = value.use_begin(); iter != value.use_end(); ++iter) {
+    if (iter->owner()->isa<::pir::ShadowOutputOp>()) {
+      names.push_back(
+          iter->owner()->attribute<StrAttribute>("output_name").AsString());
+    } else if (iter->owner()->isa<::pir::SetParameterOp>()) {
+      names.push_back(
+          iter->owner()->attribute<StrAttribute>("parameter_name").AsString());
     }
   }
-  return name;
+  return names;
 }
 
-std::string GetValueName(Value value) {
+std::vector<std::string> GetValueAllNames(Value value) {
+  std::vector<std::string> names;
   std::optional<std::string> input_name = GetValueInputName(value);
-  if (input_name.has_value()) return input_name.value();
+  if (input_name.has_value()) {
+    names.push_back(input_name.value());
+  }
+
+  std::vector<std::string> output_name = GetValueOutputNames(value);
+  for (auto &name : output_name) {
+    names.push_back(name);
+  }
+
+  return names;
+}
+
+std::string GetValueFirstName(Value value) {
+  auto name = TryGetValueFirstName(value);
 
-  VLOG(5) << "GetValueInputName(value) return empty, try to "
-             "GetValueOutputName(value).";
-  std::optional<std::string> output_name = GetValueOutputName(value);
-  if (output_name.has_value()) return output_name.value();
+  PADDLE_ENFORCE(name.has_value(),
+                 phi::errors::InvalidArgument(
+                     "Currently, we can only get name of Value from "
+                     "DataOp/ParameterOp/BlockArgument/ConstantTensorOp/"
+                     "SetParameterOp and ShadowOutputOp."));
 
-  PADDLE_THROW(phi::errors::InvalidArgument(
-      "Currently, we can only get name of Value from "
-      "DataOp/ParameterOp/BlockArgument and ShadowOutputOp."));
+  return name.value();
 }
 
+std::optional<std::string> TryGetValueFirstName(Value value) {
+  std::optional<std::string> name;
+
+  auto names = GetValueAllNames(value);
+  if (!names.empty()) {
+    return names[0];
+  }
+
+  return name;
+}
+}  // namespace name_analysis
+
 phi::DataType GetTensorDtype(Type type) {
   if (!type) {
     PADDLE_THROW(phi::errors::InvalidArgument("The type of value is nullptr."));
@@ -434,12 +443,12 @@ void PruneWithInput(const std::vector<pir::Value> &input_vars,
     std::vector<pir::Value> new_input_vars;
     for (uint64_t idx = 0; idx < input_vars.size(); idx++) {
       auto input = input_vars[idx];
-      auto orgin_op = input.defining_op();
-      std::string name = "input_" + idx;
-      if (HasValueName(input)) {
-        name = GetValueName(input);
+      auto origin_op = input.defining_op();
+      std::string name = "input_" + std::to_string(idx);
+      if (auto names = name_analysis::TryGetValueFirstName(input)) {
+        name = names.value();
       }
-      auto new_input = AppendDataOp(global_block, input, name, *orgin_op);
+      auto new_input = AppendDataOp(global_block, input, name, *origin_op);
       input.ReplaceAllUsesWith(new_input);
       new_input_vars.push_back(new_input);
     }
@@ -649,11 +658,11 @@ void BindProgram(py::module *m) {
           return_value_policy::reference)
       .def("get_output_value_by_name",
            [](Program &self, const std::string &name) {
-             return GetOutputValueByName(self, name);
+             return name_analysis::GetOutputValueByName(self, name);
            })
       .def("get_parameter_value_by_name",
            [](Program &self, const std::string &name) {
-             return GetParameterValueByName(self, name);
+             return name_analysis::GetParameterValueByName(self, name);
            })
       .def("num_ops", [](Program &self) { return self.num_ops(); })
       .def(
@@ -670,14 +679,16 @@ void BindProgram(py::module *m) {
                     var.attribute<BoolAttribute>(kAttrIsPersistable);
                 if (is_persistable && is_persistable.data()) {
                   if (var.defining_op()->isa<::pir::ParameterOp>()) {
-                    std::string var_name = GetValueName(var);
+                    std::string var_name =
+                        name_analysis::GetValueAllNames(var)[0];
                     auto tensor =
                         scope.FindVar(var_name)->GetMutable<phi::DenseTensor>();
                     state_dict_param[var_name] = *tensor;
                     state_dict_all[var_name] = *tensor;
                   } else if (var.defining_op()
                                  ->isa<paddle::dialect::DataOp>()) {
-                    std::string var_name = GetValueName(var);
+                    std::string var_name =
+                        name_analysis::GetValueAllNames(var)[0];
                     auto tensor =
                         scope.FindVar(var_name)->GetMutable<phi::DenseTensor>();
                     state_dict_opt[var_name] = *tensor;
@@ -1336,24 +1347,26 @@ void BindValue(py::module *m) {
           })
       .def_property(
           "name",
-          [](Value self) { return GetValueName(self); },
-          [](Value self, const std::string &name) { SetValueName(self, name); })
-      .def_property_readonly("has_name",
-                             [](Value self) { return HasValueName(self); })
+          [](Value self) -> std::string {
+            return name_analysis::GetValueFirstName(self);
+          },
+          [](Value self, const std::string &name) {
+            name_analysis::SetValueAllNamesWith(self, name);
+          })
+      .def_property_readonly(
+          "has_name",
+          [](Value self) {
+            return name_analysis::TryGetValueFirstName(self).has_value();
+          })
       // Return all Maybe names of given Value, for example:
       // DataOp("var_1") -> %0 -> shadow_output("output_2")
       // Return ["var_1", "output_2"]
-      .def_property_readonly(
-          "_names",
-          [](Value self) -> py::list {
-            py::list all_names;
-            std::optional<std::string> input_name = GetValueInputName(self);
-            if (input_name.has_value()) all_names.append(input_name.value());
-
-            std::optional<std::string> output_name = GetValueOutputName(self);
-            if (output_name.has_value()) all_names.append(output_name.value());
-            return all_names;
-          })
+      .def_property_readonly("_names",
+                             [](Value self) -> py::list {
+                               std::vector<std::string> names =
+                                   name_analysis::GetValueAllNames(self);
+                               return py::cast(names);
+                             })
       .def_property(
           "shape",
           [](Value self) { return phi::vectorize(GetValueDims(self)); },
@@ -1889,8 +1902,9 @@ int AppendShadowOutputs(Program *program,
   for (const auto &value : outputs) {
     if (!added_value.count(value) || IsFakeValue(value)) {
       std::string shadow_output_name = name_prefix + std::to_string(counter);
-      if (HasValueName(value)) {
-        shadow_output_name = GetValueName(value);
+      if (auto names = name_analysis::GetValueOutputNames(value);
+          !names.empty()) {
+        shadow_output_name = names[0];
       }
       AppendShadowOutput(
           program, value, shadow_output_name, start_point + counter);
@@ -2003,8 +2017,8 @@ SplitedResult SplitForwardBackward(
         }
         std::string shadow_output_name =
             std::string("output_") + std::to_string(counter);
-        if (HasValueName(v)) {
-          shadow_output_name = GetValueName(v);
+        if (auto names = name_analysis::TryGetValueFirstName(v)) {
+          shadow_output_name = names.value();
         }
         auto op_info = ctx->GetRegisteredOpInfo(pir::ShadowOutputOp::name());
         pir::AttributeMap attribute_map = {
@@ -2204,7 +2218,7 @@ static void inline CreateVariableIfNotExist(
 
   for (size_t i = 0; i < len; ++i) {
     pir::Value value = var_list[i];
-    std::string para_name = GetValueName(value);
+    std::string para_name = name_analysis::GetValueFirstName(value);
     auto var = scope->FindVar(para_name);
     if (var == nullptr) {
       PADDLE_ENFORCE_NOT_NULL(exe,
diff --git a/paddle/fluid/pybind/pir.h b/paddle/fluid/pybind/pir.h
@@ -26,8 +26,14 @@ using pir::Value;
 void BindPir(pybind11::module *m);
 const phi::DDim &GetValueDims(Value value);
 bool GetValueBoolAttr(Value value, const std::string &attr_name);
-std::string GetValueName(Value value);
-bool HasValueName(const Value &value);
+
+namespace name_analysis {
+
+std::vector<std::string> GetValueAllNames(Value value);
+std::string GetValueFirstName(Value value);
+std::optional<std::string> TryGetValueFirstName(Value value);
+
+}  // namespace name_analysis
 
 }  // namespace pybind
 }  // namespace paddle
