[Dy2St][PIR] Handle OutletType in backward inputs (#62256)

SigureMo · web-flow · commit 368c04bc01d8 · 2024-03-05T00:02:15.000+08:00
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -85,14 +85,72 @@ static std::vector<std::string> GetTensorsName(
   return in_names;
 }
 
+static bool IsVariableRefArray(const Tensor &tensor) {
+  return paddle::framework::VariableRefArray::classof(tensor.impl().get());
+}
+
+static auto GetNameFromValue(const ::pir::Block *block,
+                             const std::vector<::pir::Value> &values,
+                             bool is_input) {
+  // we use name here, later value is used directly.
+  std::unordered_map<::pir::Value, std::string> value2name;
+  if (is_input) {
+    for (auto &kwarg : block->kwargs()) {
+      value2name[kwarg.second] = kwarg.first;
+    }
+  }
+  for (auto &op : *block) {
+    std::string name;
+    if (is_input && op.name() == "pd_op.data") {
+      name =
+          op.attributes().at("name").dyn_cast<pir::StrAttribute>().AsString();
+      value2name[op.results()[0].Value::impl()] = name;
+    } else if (!is_input && op.name() == "builtin.set_parameter") {
+      name = op.attributes()
+                 .at("parameter_name")
+                 .dyn_cast<pir::StrAttribute>()
+                 .AsString();
+      value2name[op.operand(0).source()] = name;
+    } else if (!is_input && op.name() == "builtin.shadow_output") {
+      name = op.attributes()
+                 .at("output_name")
+                 .dyn_cast<pir::StrAttribute>()
+                 .AsString();
+      value2name[op.operand(0).source()] = name;
+    } else if (is_input && op.name() == "builtin.parameter") {
+      name = op.attributes()
+                 .at("parameter_name")
+                 .dyn_cast<pir::StrAttribute>()
+                 .AsString();
+      value2name[op.result(0).Value::impl()] = name;
+    } else if (is_input && op.name() == "builtin.constant") {
+      if (op.isa<pir::ConstantTensorOp>()) {
+        name = op.dyn_cast<pir::ConstantTensorOp>().tensor_name();
+        value2name[op.result(0).Value::impl()] = name;
+      }
+    }
+  }
+  std::vector<std::string> names;
+  std::transform(values.begin(),
+                 values.end(),
+                 std::back_inserter(names),
+                 [&value2name](const ::pir::Value &v) {
+                   if (!value2name.count(v))
+                     return std::string(paddle::framework::kFakeVarName);
+                   return value2name.at(v);
+                 });
+  return names;
+}
+
 static void CheckInputVarStatus(const Tensor &tensor) {
-  PADDLE_ENFORCE_EQ(tensor.defined() && tensor.is_dense_tensor(),
-                    true,
-                    paddle::platform::errors::InvalidArgument(
-                        "The input tensor %s of "
-                        "RunProgram(Grad)Op holds "
-                        "wrong type. Expect type is DenseTensor.",
-                        tensor.name()));
+  PADDLE_ENFORCE_EQ(
+      tensor.defined() &&
+          (tensor.is_dense_tensor() || IsVariableRefArray(tensor)),
+      true,
+      paddle::platform::errors::InvalidArgument(
+          "The input tensor %s of RunProgram(Grad)Op holds "
+          "wrong type. Expect type is DenseTensor or VariableRefArray.",
+          tensor.name()));
 }
 
 static void CheckOutputVarStatus(const paddle::framework::Variable &src_var,
@@ -121,8 +179,7 @@ static void CheckOutputVarStatus(const paddle::framework::Variable &src_var,
                           "RunProgram(Grad)Op's internal scope holds "
                           "wrong type. Expect type is SelectedRows",
                           name));
-  } else if (paddle::framework::VariableRefArray::classof(
-                 dst_tensor.impl().get())) {
+  } else if (IsVariableRefArray(dst_tensor)) {
     auto &src_tensor = src_var.Get<paddle::framework::VariableRefArray>();
     PADDLE_ENFORCE_EQ(paddle::framework::VariableRefArray::classof(&src_tensor),
                       true,
@@ -139,38 +196,15 @@ static void CheckOutputVarStatus(const paddle::framework::Variable &src_var,
   }
 }
 
-static void ShareTensorsIntoScope(const std::vector<Tensor> &tensors,
-                                  paddle::framework::Scope *scope) {
-  for (size_t i = 0; i < tensors.size(); ++i) {
-    VLOG(4) << "Share Tensor Into Scope: " << i;
-    auto name = tensors[i].name();
-    if (name == paddle::framework::kFakeVarName ||
-        name == paddle::framework::kEmptyVarName) {
-      continue;
-    }
-    auto *var = scope->Var(name);
-    CheckInputVarStatus(tensors[i]);
-    // share tensor
-    auto tensor_base = tensors[i].impl();
-    if (phi::DenseTensor::classof(tensor_base.get())) {
-      auto *dst_tensor = var->GetMutable<phi::DenseTensor>();
-      auto t = std::dynamic_pointer_cast<phi::DenseTensor>(tensor_base);
-      *dst_tensor = *t;
-    } else if (phi::SelectedRows::classof(tensor_base.get())) {
-      auto *dst_tensor = var->GetMutable<phi::SelectedRows>();
-      auto t = std::dynamic_pointer_cast<phi::SelectedRows>(tensor_base);
-      *dst_tensor = *t;
-    }
-  }
-}
-
 static void ShareTensorsIntoScopeWithName(
     const std::vector<Tensor> &tensors,
     const std::vector<std::string> &tensor_names,
     paddle::framework::Scope *scope) {
   for (size_t i = 0; i < tensors.size(); ++i) {
+    VLOG(4) << "Share Tensor Into Scope: " << i;
     auto name = tensor_names[i];
-    if (name == paddle::framework::kFakeVarName) {
+    if (name == paddle::framework::kFakeVarName ||
+        name == paddle::framework::kEmptyVarName) {
       continue;
     }
     auto *var = scope->Var(name);
@@ -185,102 +219,28 @@ static void ShareTensorsIntoScopeWithName(
       auto *dst_tensor = var->GetMutable<phi::SelectedRows>();
       auto t = std::dynamic_pointer_cast<phi::SelectedRows>(tensor_base);
       *dst_tensor = *t;
+    } else if (paddle::framework::VariableRefArray::classof(
+                   tensor_base.get())) {
+      auto *dst_tensor = var->GetMutable<paddle::framework::VariableRefArray>();
+      auto t = std::dynamic_pointer_cast<paddle::framework::VariableRefArray>(
+          tensor_base);
+      *dst_tensor = *t;
     }
   }
 }
 
-static auto GetNameFromValue(const ::pir::Block *block,
-                             const std::vector<::pir::Value> &values,
-                             bool is_input) {
-  // we use name here, later value is used directly.
-  std::unordered_map<::pir::Value, std::string> value2name;
-  if (is_input) {
-    for (auto &kwarg : block->kwargs()) {
-      value2name[kwarg.second] = kwarg.first;
-    }
-  }
-  for (auto &op : *block) {
-    std::string name;
-    if (is_input && op.name() == "pd_op.data") {
-      name =
-          op.attributes().at("name").dyn_cast<pir::StrAttribute>().AsString();
-      value2name[op.results()[0].Value::impl()] = name;
-    } else if (!is_input && op.name() == "builtin.set_parameter") {
-      name = op.attributes()
-                 .at("parameter_name")
-                 .dyn_cast<pir::StrAttribute>()
-                 .AsString();
-      value2name[op.operand(0).source()] = name;
-    } else if (!is_input && op.name() == "builtin.shadow_output") {
-      name = op.attributes()
-                 .at("output_name")
-                 .dyn_cast<pir::StrAttribute>()
-                 .AsString();
-      value2name[op.operand(0).source()] = name;
-    } else if (is_input && op.name() == "builtin.parameter") {
-      name = op.attributes()
-                 .at("parameter_name")
-                 .dyn_cast<pir::StrAttribute>()
-                 .AsString();
-      value2name[op.result(0).Value::impl()] = name;
-    } else if (is_input && op.name() == "builtin.constant") {
-      if (op.isa<pir::ConstantTensorOp>()) {
-        name = op.dyn_cast<pir::ConstantTensorOp>().tensor_name();
-        value2name[op.result(0).Value::impl()] = name;
-      }
-    }
-  }
-  std::vector<std::string> names;
-  std::transform(values.begin(),
-                 values.end(),
-                 std::back_inserter(names),
-                 [&value2name](const ::pir::Value &v) {
-                   if (!value2name.count(v))
-                     return std::string(paddle::framework::kFakeVarName);
-                   return value2name.at(v);
-                 });
-  return names;
-}
+static void ShareTensorsIntoScope(const std::vector<Tensor> &tensors,
+                                  paddle::framework::Scope *scope) {
+  const std::vector<std::string> names =
+      [&](const std::vector<Tensor> &tensors) {
+        std::vector<std::string> names;
+        for (auto &t : tensors) {
+          names.push_back(t.name());
+        }
+        return names;
+      }(tensors);
 
-static void ShareTensorsFromScope(
-    const std::vector<Tensor *> &tensors,
-    const paddle::framework::BlockDesc &global_block,
-    paddle::framework::Scope *scope) {
-  for (size_t i = 0; i < tensors.size(); ++i) {
-    // NOTE: In case of setting out_tmp.stop_gradient = True in model code, all
-    // parameters before generating out_tmp have no @GRAD, it will raise error
-    // because we can't find them in scope. So we skip sharing these vars or
-    // var@GRAD if they don't appear in global block.
-    auto &name = tensors[i]->name();
-    if (name == paddle::framework::kEmptyVarName ||
-        name == paddle::framework::kFakeVarName || !global_block.HasVar(name)) {
-      VLOG(2) << "find tensor name is " << name << ", skip it!";
-      continue;
-    }
-    // NOTE: Here skip not found var is dangerous, if a bug is caused here,
-    // the result is grad calculation error, which will be very hidden!
-    auto *var = scope->FindVar(name);
-    PADDLE_ENFORCE_NOT_NULL(
-        var,
-        paddle::platform::errors::NotFound("The output tensor %s is not in "
-                                           "RunProgram(Grad)Op'"
-                                           "s internal scope.",
-                                           name));
-    CheckOutputVarStatus(*var, *tensors[i]);
-    // share tensor
-    if (var->IsType<phi::DenseTensor>()) {
-      auto &src_tensor = var->Get<phi::DenseTensor>();
-      auto *dst_tensor = const_cast<phi::DenseTensor *>(
-          dynamic_cast<const phi::DenseTensor *>(tensors[i]->impl().get()));
-      VLOG(4) << "share " << name << " from scope";
-      *dst_tensor = src_tensor;
-    } else if (var->IsType<phi::SelectedRows>()) {
-      auto &src_tensor = var->Get<phi::SelectedRows>();
-      auto *dst_tensor = const_cast<phi::SelectedRows *>(
-          dynamic_cast<const phi::SelectedRows *>(tensors[i]->impl().get()));
-      *dst_tensor = src_tensor;
-    }
-  }
+  ShareTensorsIntoScopeWithName(tensors, names, scope);
 }
 
 static void ShareTensorsIntoScopeByValue(
@@ -372,6 +332,17 @@ static void ShareTensorsFromScopeWithPartialBlock(
       auto *dst_tensor = const_cast<phi::SelectedRows *>(
           dynamic_cast<const phi::SelectedRows *>(tensors[i]->impl().get()));
       *dst_tensor = src_tensor;
+    } else if (var->IsType<paddle::framework::VariableRefArray>()) {
+      auto &src_tensor = var->Get<paddle::framework::VariableRefArray>();
+      auto *dst_tensor = const_cast<paddle::framework::VariableRefArray *>(
+          dynamic_cast<const paddle::framework::VariableRefArray *>(
+              tensors[i]->impl().get()));
+      *dst_tensor = src_tensor;
+    } else {
+      PADDLE_THROW(paddle::platform::errors::InvalidArgument(
+          "The RunProgram(Grad)Op only support output "
+          "variable of type DenseTensor, SelectedRows or VariableRefArray",
+          name));
     }
   }
 }
@@ -1541,12 +1512,19 @@ class PirGradNodeRunProgram : public egr::GradNodeBase {
             x_grad_values.size()));
 
     // TODO(dev): Need an elegant way to determine information of grad_tensor,
-    // such as: name, tensor type(DenseTensor or SelectedRows).
+    // such as: name, tensor type (DenseTensor, SelectedRows or
+    // VariableRefArray).
     for (size_t i = 0; i < x.size(); i++) {
       if (x[i].is_dense_tensor()) {
         x_grad->emplace_back(std::make_shared<phi::DenseTensor>());
       } else if (x[i].is_selected_rows()) {
         x_grad->emplace_back(std::make_shared<phi::SelectedRows>());
+      } else if (details::IsVariableRefArray(x[i])) {
+        x_grad->emplace_back(
+            std::make_shared<paddle::framework::VariableRefArray>());
+      } else {
+        PADDLE_THROW(paddle::platform::errors::InvalidArgument(
+            "The grad tensor type is not supported."));
       }
     }
   }
diff --git a/test/dygraph_to_static/test_ifelse.py b/test/dygraph_to_static/test_ifelse.py
@@ -554,6 +554,7 @@ def forward(self, a, b, c):
         a = paddle.matmul(a, self.param)
         a = paddle.reshape(a, (2, 4))
         cond = paddle.to_tensor([10])
+        b = b.broadcast_to(self.param.shape)
         if paddle.equal(cond, 10):
             a_argmax = a.argmax(axis=-1)
             b = b + self.param
