[Science] Fix dpa bug

paddle/fluid/framework/new_executor/instruction/control_flow/yield_instruction.cc

@@ -17,7 +17,11 @@
 #include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
 #include "paddle/fluid/framework/new_executor/new_executor_defs.h"
 #include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
+#include "paddle/fluid/framework/op_kernel_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
+#include "paddle/phi/api/include/context_pool.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/pir/include/core/builtin_type.h"
 
 namespace paddle {
 namespace framework {
@@ -26,11 +30,9 @@ YieldInstruction::YieldInstruction(size_t id,
                                    const phi::Place &place,
                                    ::pir::Operation *op,
                                    ValueExecutionInfo *value_exe_info)
-    : InstructionBase(id, place), op_(op) {
+    : InstructionBase(id, place), op_(op), value_exe_info_(value_exe_info) {
   VLOG(6) << "construct yield instruction";
-
   auto parent_op = op->GetParentOp();
-
   std::unordered_map<pir::Value, std::vector<int>> inputs;
   for (size_t i = 0; i < op->num_operands(); ++i) {
     // Skip the first input (cond) when the parent op is a while op.
@@ -39,8 +41,24 @@ YieldInstruction::YieldInstruction(size_t id,
     }
     auto in = op->operand_source(i);
     if (in) {
-      inputs.emplace(in, GetValueIds(in, *value_exe_info));
-      input_vars_.push_back(value_exe_info->GetVarByValue(in));
+      inputs.emplace(in, GetValueIds(in, *value_exe_info_));
+      input_vars_.push_back(value_exe_info_->GetVarByValue(in));
+    } else {
+      // value 为空的时候根据 parent op 输出 value 的 meta 信息填一个全 0
+      // tensor。Build instruction 的时候先创建 var
+      if (parent_op->result(i) && parent_op->result(i).type()) {
+        auto out_type = parent_op->result(i).type();
+        std::string new_name = "_fake_var_op_" + std::to_string(op->id()) +
+                               "_input_" + std::to_string(i) + "_";
+        Variable *fake_var = value_exe_info_->GetScope()->Var(new_name);
+        if (out_type.isa<paddle::dialect::AllocatedDenseTensorType>()) {
+          fake_var->GetMutable<phi::DenseTensor>();
+          input_vars_.push_back(fake_var);
+        } else {
+          PADDLE_THROW(common::errors::Unimplemented(
+              "unsupported type %d", out_type.dyn_cast<pir::Type>().type_id()));
+        }
+      }
     }
   }
   SetInputs(inputs);
@@ -62,11 +80,52 @@ YieldInstruction::YieldInstruction(size_t id,
                                       output_vars_.size()));
 }
 
+template <typename T0, typename T1>
+void FullFakeTensor(const pir::Value &output_value, Variable *output_var) {
+  if (!output_value || !output_var) {
+    output_var = nullptr;
+    return;
+  }
+  auto out_tensor_type = output_value.type().dyn_cast<T0>();
+  auto abs_dims = out_tensor_type.dims();
+  for (int i = 0; i < abs_dims.size(); ++i) {
+    // dynamic shape, set to 1
+    if (abs_dims[i] == -1) {
+      abs_dims[i] = 1;
+    }
+  }
+#ifdef PADDLE_WITH_CUDA
+  phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+  auto *dev_ctx = pool.Get(phi::GPUPlace());
+  phi::DataType dtype =
+      paddle::dialect::TransToPhiDataType(out_tensor_type.dtype());
+  phi::FullKernel<float, phi::GPUContext>(
+      *(static_cast<phi::GPUContext *>(dev_ctx)),
+      phi::IntArray(common::vectorize(abs_dims)),
+      0.0,
+      dtype,
+      output_var->GetMutable<T1>());
+#else
+  VLOG(1) << "unsupported Device for Fake Tensor.";
+  output_var = nullptr;
+  return;
+#endif
+}
 void YieldInstruction::Run() {
   for (size_t i = 0; i < input_vars_.size(); ++i) {
-    if (input_vars_[i]->IsType<phi::DenseTensor>()) {
-      output_vars_[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
-          input_vars_[i]->Get<phi::DenseTensor>());
+    if (input_vars_[i] == nullptr) {
+      output_vars_[i] = nullptr;
+    } else if (input_vars_[i]->IsType<phi::DenseTensor>()) {
+      if (input_vars_[i]->IsInitialized() &&
+          !input_vars_[i]->Get<phi::DenseTensor>().initialized()) {
+        // 对应 input 为 NULL VALUE 的情况，fake tensor
+        FullFakeTensor<paddle::dialect::AllocatedDenseTensorType,
+                       phi::DenseTensor>(
+            value_exe_info_->GetValueByVar(output_vars_[i]), output_vars_[i]);
+      } else {
+        output_vars_[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
+            input_vars_[i]->Get<phi::DenseTensor>());
+      }
     } else if (input_vars_[i]->IsType<phi::TensorArray>()) {
       const auto &inner_array = input_vars_[i]->Get<phi::TensorArray>();
       auto *output_array = output_vars_[i]->GetMutable<phi::TensorArray>();

paddle/fluid/framework/new_executor/instruction/control_flow/yield_instruction.h

@@ -41,6 +41,8 @@ class YieldInstruction : public InstructionBase {
   std::vector<Variable*> input_vars_;
 
   std::vector<Variable*> output_vars_;
+
+  ValueExecutionInfo* value_exe_info_;
 };
 
 }  // namespace framework

paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc

@@ -136,6 +136,16 @@ Variable* ValueExecutionInfo::GetVarByValue(pir::Value value) const {
   return scope_->FindVar(GetVarName(value));
 }
 
+::pir::Value ValueExecutionInfo::GetValueByVar(const Variable* var) const {
+  for (const auto& pair : value_2_var_name_) {
+    if (pair.second == GetVarName(var)) {
+      return pair.first;
+    }
+  }
+  PADDLE_THROW(::common::errors::Unimplemented("Cannot find value by var %s",
+                                               GetVarName(var)));
+}
+
 const std::unordered_map<::pir::Value, std::string>&
 ValueExecutionInfo::GetValue2VarName() const {
   return value_2_var_name_;

paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h

@@ -69,6 +69,8 @@ class ValueExecutionInfo {
 
   Variable* GetVarByValue(pir::Value value) const;
 
+  ::pir::Value GetValueByVar(const Variable* var) const;
+
   const std::unordered_map<::pir::Value, std::string>& GetValue2VarName() const;
 
   void AddValue2VarName(::pir::Value value, const std::string& var_name);