【pir】modify test_while_op.py (#60753)

xiaoguoguo626807 · zhangbo9674 · web-flow · commit 446f7a72fe52 · 2024-01-16T19:56:53.000+08:00
* test

* fix

* fix

* fix

* modify add_n block_arg

* modify increment return value

* merge

* modfiy whiel_op.py

* clear pop push instruction

* optimizer code

* modify test_while_op

* modidy * param

* comment

* delete std::cout

* modify

---------

Co-authored-by: zhangbo9674 &lt;zhangbo54@baidu.com&gt;
diff --git a/paddle/fluid/framework/new_executor/instruction/control_flow/tuple_pop_instruction.cc b/paddle/fluid/framework/new_executor/instruction/control_flow/tuple_pop_instruction.cc
@@ -106,20 +106,20 @@ void ShareVarData(const Variable* src_var, Variable* dst_var) {
       }
       dst_tensor_array->push_back(*tmp_dst_tensor);
     }
+  } else if (src_var->IsType<VariableRefArray>()) {
+    auto src_var_array = src_var->Get<VariableRefArray>();
+    auto* dst_var_array = dst_var->GetMutable<VariableRefArray>();
+    for (size_t i = 0; i < src_var_array.size(); ++i) {
+      Variable* copy_var = const_cast<Variable*>(dst_var_array->at(i));
+      ShareVarData(src_var_array.at(i), copy_var);
+    }
   } else {
     PADDLE_THROW(phi::errors::PreconditionNotMet(
         "Output only support DenseTensorType "
         "or SelectedRowsType or TensorArrayType"));
   }
 }
 
-void ShareVarData(const VariableRefArray* src_var, VariableRefArray* dst_var) {
-  for (size_t i = 0; i < src_var->size(); ++i) {
-    Variable* copy_var = const_cast<Variable*>(dst_var->at(i));
-    ShareVarData(src_var->at(i), copy_var);
-  }
-}
-
 void TuplePopInstruction::Run() {
   VLOG(6) << "run tuple_pop instruction";
   if (tuple_pop_op_.tuple_size() == 0) {
diff --git a/paddle/fluid/framework/new_executor/instruction/control_flow/tuple_push_instruction.cc b/paddle/fluid/framework/new_executor/instruction/control_flow/tuple_push_instruction.cc
@@ -20,6 +20,39 @@
 
 namespace paddle {
 namespace framework {
+bool ParsePlace(const pir::Type& type, OpFuncType* type_) {
+  if (type.isa<paddle::dialect::AllocatedDenseTensorType>()) {
+    auto place =
+        type.dyn_cast<paddle::dialect::AllocatedDenseTensorType>().place();
+    place == phi::CPUPlace();
+    if (place == phi::GPUPlace()) {
+      *type_ = OpFuncType::kGpuAsync;
+      return true;
+    }
+  } else if (type.isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
+    auto place =
+        type.dyn_cast<paddle::dialect::AllocatedDenseTensorArrayType>().place();
+    if (place == phi::GPUPlace()) {
+      *type_ = OpFuncType::kGpuAsync;
+      return true;
+    }
+  } else if (type.isa<pir::VectorType>()) {
+    pir::VectorType inlet_element_type = type.dyn_cast<pir::VectorType>();
+    for (size_t i = 0; i < static_cast<size_t>(inlet_element_type.size());
+         i++) {
+      if (ParsePlace(inlet_element_type[i], type_)) {
+        return true;
+      }
+    }
+  } else {
+    PADDLE_THROW(phi::errors::PreconditionNotMet(
+        "Only support AllocatedDenseTensorType and "
+        "AllocatedDenseTensorArrayType in vectortype now, but get: %s",
+        type));
+  }
+  return false;
+}
+
 TuplePushInstruction::TuplePushInstruction(size_t id,
                                            const platform::Place& place,
                                            ::pir::Operation* op,
@@ -50,63 +83,8 @@ TuplePushInstruction::TuplePushInstruction(size_t id,
   type_ = OpFuncType::kCpuSync;
   for (size_t i = 0; i < tuple_push_op_.tuple_size(); ++i) {
     auto inlet_element_value = tuple_push_op_.inlet_element(i);
-
-    if (inlet_element_value.type()
-            .isa<paddle::dialect::AllocatedDenseTensorType>()) {
-      auto place = inlet_element_value.type()
-                       .dyn_cast<paddle::dialect::AllocatedDenseTensorType>()
-                       .place();
-      if (place == phi::GPUPlace()) {
-        type_ = OpFuncType::kGpuAsync;
-        break;
-      }
-    } else if (inlet_element_value.type()
-                   .isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
-      auto place =
-          inlet_element_value.type()
-              .dyn_cast<paddle::dialect::AllocatedDenseTensorArrayType>()
-              .place();
-      if (place == phi::GPUPlace()) {
-        type_ = OpFuncType::kGpuAsync;
-        break;
-      }
-    } else if (inlet_element_value.type().isa<pir::VectorType>()) {
-      pir::VectorType inlet_element_type =
-          inlet_element_value.type().dyn_cast<pir::VectorType>();
-      for (size_t i = 0; i < static_cast<size_t>(inlet_element_type.size());
-           i++) {
-        if (inlet_element_type[i]
-                .isa<paddle::dialect::AllocatedDenseTensorType>()) {
-          auto place =
-              inlet_element_type[i]
-                  .dyn_cast<paddle::dialect::AllocatedDenseTensorType>()
-                  .place();
-          if (place == phi::GPUPlace()) {
-            type_ = OpFuncType::kGpuAsync;
-            break;
-          }
-        } else if (inlet_element_type[i]
-                       .isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
-          auto place =
-              inlet_element_value.type()
-                  .dyn_cast<paddle::dialect::AllocatedDenseTensorArrayType>()
-                  .place();
-          if (place == phi::GPUPlace()) {
-            type_ = OpFuncType::kGpuAsync;
-            break;
-          }
-        } else {
-          PADDLE_THROW(phi::errors::PreconditionNotMet(
-              "Only support AllocatedDenseTensorType and "
-              "AllocatedDenseTensorArrayType in vectortype now, but get: %s",
-              inlet_element_type[i]));
-        }
-      }
-    } else {
-      PADDLE_THROW(phi::errors::PreconditionNotMet(
-          "Only support AllocatedDenseTensorType and "
-          "AllocatedDenseTensorArrayType in vectortype now, but get: %s",
-          inlet_element_value.type()));
+    if (ParsePlace(inlet_element_value.type(), &type_)) {
+      break;
     }
   }
 }
@@ -119,6 +97,7 @@ void TuplePushInstruction::Run() {
     auto& value_2_var_name = value_exe_info_->GetValue2VarName();
     // TODO(zhangbo): Performance optimization: static acquisition of TuplePush
     // input variables and name.
+    std::map<Variable*, Variable*> src_to_dst;
     for (size_t i = 0; i < tuple_push_op_.tuple_size(); i++) {
       auto inlet_element_value = tuple_push_op_.inlet_element(i);
       Variable* var = value_exe_info_->GetVarByValue(inlet_element_value);
@@ -130,7 +109,8 @@ void TuplePushInstruction::Run() {
       auto* copy_var = value_exe_info_->GetScope()->Var(new_name);
       bool is_optional = (inlet_element_value.impl() == nullptr ||
                           !inlet_element_value.type());
-      DeepCopyVariable(var, copy_var, value_exe_info_, stack_size, is_optional);
+      DeepCopyVariable(
+          var, copy_var, value_exe_info_, stack_size, is_optional, &src_to_dst);
       VLOG(10) << "done DeepCopyVariable " << new_name;
       stack_element_var_array_->emplace_back(copy_var);
       VLOG(6) << "push back var: " << new_name << "[" << copy_var << "]"
diff --git a/paddle/fluid/framework/new_executor/instruction/instruction_base.cc b/paddle/fluid/framework/new_executor/instruction/instruction_base.cc
@@ -239,7 +239,17 @@ const std::vector<size_t>& InstructionBase::GCCheckVars() const {
 }
 
 void InstructionBase::AddEagerGCVar(Variable* var) {
-  eager_gc_vars_.push_back(var);
+  if (var->IsType<VariableRefArray>()) {
+    auto array = var->Get<VariableRefArray>();
+    for (size_t i = 0; i < array.size(); ++i) {
+      AddEagerGCVar(const_cast<Variable*>(array.at(i)));
+    }
+  } else {
+    if (std::find(eager_gc_vars_.begin(), eager_gc_vars_.end(), var) ==
+        eager_gc_vars_.end()) {
+      eager_gc_vars_.push_back(var);
+    }
+  }
 }
 
 const std::vector<Variable*>& InstructionBase::EagerGCVars() const {
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
@@ -289,10 +289,16 @@ void CheckInputVars(pir::Operation* op,
 }
 
 void DeepCopyVariable(const Variable* src_var,
-                      Variable* dst_var,
+                      Variable*& dst_var,  // NOLINT
                       ValueExecutionInfo* value_exe_info,
                       uint32_t stack_size,
-                      bool is_optional) {
+                      bool is_optional,
+                      std::map<Variable*, Variable*>* src_to_dst_map) {
+  if (src_to_dst_map->count(const_cast<Variable*>(src_var))) {
+    dst_var = (*src_to_dst_map)[const_cast<Variable*>(src_var)];
+    return;
+  }
+
   if (src_var->IsType<phi::DenseTensor>()) {
     auto& src_tensor = src_var->Get<phi::DenseTensor>();
     auto* tmp_dst_tensor = dst_var->GetMutable<phi::DenseTensor>();
@@ -366,8 +372,12 @@ void DeepCopyVariable(const Variable* src_var,
       std::string new_name = "copied_" + std::to_string(stack_size) + '_' +
                              value_exe_info->GetVarName(src_ref_var);
       auto tmp_dst_var = value_exe_info->GetScope()->Var(new_name);
-      DeepCopyVariable(
-          src_ref_var, tmp_dst_var, value_exe_info, stack_size, is_optional);
+      DeepCopyVariable(src_ref_var,
+                       tmp_dst_var,
+                       value_exe_info,
+                       stack_size,
+                       is_optional,
+                       src_to_dst_map);
       dst_ref_array->emplace_back(tmp_dst_var);
     }
 
@@ -376,6 +386,7 @@ void DeepCopyVariable(const Variable* src_var,
         "Output only support DenseTensorType "
         "or SelectedRowsType or TensorArrayType or VariableRefArrayType"));
   }
+  (*src_to_dst_map)[const_cast<Variable*>(src_var)] = dst_var;
 }
 
 void BuildValue(pir::Value value,
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h
@@ -134,10 +134,11 @@ void BuildScope(const pir::Block& block,
                 ValueExecutionInfo* value_exe_info = nullptr);
 
 void DeepCopyVariable(const Variable* src_var,
-                      Variable* dst_var,
+                      Variable*& dst_var,  // NOLINT
                       ValueExecutionInfo* value_exe_info,
                       uint32_t stack_size,
-                      bool is_optional);
+                      bool is_optional,
+                      std::map<Variable*, Variable*>* src_to_dst_map);
 
 void BuildRuntimeContext(pir::Operation* op,
                          const ValueExecutionInfo& value_exec_info,
diff --git a/test/legacy_test/test_while_op.py b/test/legacy_test/test_while_op.py
@@ -20,6 +20,7 @@
 import paddle
 from paddle import base
 from paddle.base import core
+from paddle.base.backward import append_backward
 from paddle.base.executor import Executor
 from paddle.base.framework import in_pir_mode
 from paddle.incubate.layers.nn import shuffle_batch
@@ -94,7 +95,7 @@ def test_simple_net(self):
         startup_program = base.Program()
         with base.program_guard(main_program, startup_program):
             loss, sum_result = self.simple_net()
-            # append_backward(loss)
+            append_backward(loss)
 
             cpu = core.CPUPlace()
             exe = Executor(cpu)
