fix pir vector grad (#68811)

wanghuancoder · web-flow · commit 35dcc8285f62 · 2024-10-21T14:52:22.000+08:00
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -1330,7 +1330,8 @@ void PirInterpreter::CalculateLastLiveOps() {
         // skip no_need_buffer input vars
         if ((ins.count(item.first) &&
              instr->NoNeedBuffer().count(item.first)) ||
-            instr->Name() == "builtin_combine_instruction") {
+            instr->Name() == "builtin_combine_instruction" ||
+            instr->Name() == "pd_op.shadow_feed_tensors") {
           continue;
         }
         gc_check_vars.insert(var_id);
@@ -1343,8 +1344,10 @@ void PirInterpreter::CalculateLastLiveOps() {
           value_exe_info_->GetNameById(static_cast<int>(var_id)));
       PADDLE_ENFORCE_NOT_NULL(
           var,
-          common::errors::NotFound("Var(id=%d) should not be nullptr.",
-                                   static_cast<int>(var_id)));
+          common::errors::NotFound(
+              "Var(id=%d,%s) should not be nullptr.",
+              static_cast<int>(var_id),
+              value_exe_info_->GetNameById(static_cast<int>(var_id))));
       if (var->IsType<phi::DenseTensor>() || var->IsType<phi::SelectedRows>() ||
           var->IsType<phi::TensorArray>() ||
           var->IsType<phi::SparseCooTensor>() ||
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.cc b/paddle/fluid/pir/dialect/operator/ir/manual_api.cc
@@ -37,6 +37,11 @@ pir::Value builtin_combine(const std::vector<pir::Value>& x) {
   return combine_op.out();
 }
 
+std::vector<pir::Value> builtin_split(const pir::Value& x) {
+  auto split_op = ApiBuilder::Instance().GetBuilder()->Build<pir::SplitOp>(x);
+  return split_op.outputs();
+}
+
 std::vector<pir::Value> add_n_grad(const std::vector<pir::Value>& inputs,
                                    const pir::Value& out_grad) {
   std::vector<pir::Value> inputs_grad;
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.h b/paddle/fluid/pir/dialect/operator/ir/manual_api.h
@@ -26,6 +26,7 @@ namespace paddle {
 namespace dialect {
 
 pir::Value builtin_combine(const std::vector<pir::Value>& x);
+std::vector<pir::Value> builtin_split(const pir::Value& x);
 
 std::vector<pir::Value> add_n_grad(const std::vector<pir::Value>& inputs,
                                    const pir::Value& out_grad);
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -940,12 +940,10 @@ static phi::Backend GetKernelBackendByYaml(
 
   for (const auto& slot_name : backend_info) {
     auto& input_map = op_info_parser->InputName2Id();
-
     if (input_map.count(slot_name)) {
       // parse from input
       int in_index = static_cast<int>(input_map.at(slot_name));
       auto type = map_value_pair.at(op->operand_source(in_index)).type();
-
       if (type.isa<AllocatedDenseTensorType>()) {
         kernel_backend = paddle::experimental::ParseBackend(
             type.dyn_cast<AllocatedDenseTensorType>().place());
@@ -2183,6 +2181,8 @@ void HandleForSpecialOp(
   // only deal with single output
   if (op_item->num_results() > 0) {
     for (size_t i = 0; i < op_item->num_results(); ++i) {
+      VLOG(6) << "2816:" << op_item->result(i).type();
+      VLOG(6) << "2817:" << op->result(i).type();
       (*map_value_pair)[op_item->result(i)] = op->result(i);
     }
   }
@@ -3027,9 +3027,7 @@ void RemoveRedundantMemcpyAfterShadowFeed(pir::Block* block,
                                           pir::IrContext* ctx) {
   for (auto it = block->begin(); it != block->end(); ++it) {
     if (it->isa<PhiKernelOp>() &&
-        (it->dyn_cast<PhiKernelOp>().op_name() == "pd_op.shadow_feed" ||
-         it->dyn_cast<PhiKernelOp>().op_name() ==
-             "pd_op.shadow_feed_tensors")) {
+        (it->dyn_cast<PhiKernelOp>().op_name() == "pd_op.shadow_feed")) {
       pir::Value shadow_value = it->result(0);
       if (shadow_value.use_count() == 1) {
         pir::Operation* next_op = shadow_value.first_use().owner();
@@ -3320,11 +3318,11 @@ void ProcessBlock(
     for (auto& [name, attr] : arg.dyn_cast<pir::BlockArgument>().attributes()) {
       new_arg.set_attribute(name, attr);
     }
-    (*map_value_pair)[arg] = new_arg;
     if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
       new_arg.set_type(
           AllocatedDenseTensorType::get(ctx, phi::Place(), dense_tensor_type));
     }
+    (*map_value_pair)[arg] = new_arg;
   }
   if (phi::is_accelerat_place(place)) {
     for (auto& [keyword, arg] : block->kwargs()) {
@@ -3351,6 +3349,46 @@ void ProcessBlock(
 
         new_block->push_back(shadow_op);
         (*map_value_pair)[arg] = shadow_op->result(0);
+      } else if (auto vec_type = arg.type().dyn_cast<pir::VectorType>()) {
+        for (size_t i = 0; i < vec_type.size(); ++i) {
+          PADDLE_ENFORCE_EQ(
+              vec_type[i].isa<DenseTensorType>(),
+              true,
+              common::errors::PreconditionNotMet(
+                  "AddShadowFeedTensors only support DenseTensorType Now"));
+        }
+        // Add ShadowFeedTensors Op
+        phi::KernelKey shadow_key{
+            paddle::experimental::get_accelerat_backend(),
+            phi::DataLayout::ANY,
+            TransToPhiDataType(
+                vec_type[0].dyn_cast<DenseTensorType>().dtype())};
+
+        std::unordered_map<std::string, pir::Attribute> attr_map{
+            {"op_name",
+             pir::StrAttribute::get(ctx, "pd_op.shadow_feed_tensors")},
+            {"kernel_name", pir::StrAttribute::get(ctx, "shadow_feed_tensors")},
+            {"kernel_key", KernelAttribute::get(ctx, shadow_key)},
+            {"dst_place_type", pir::Int32Attribute::get(ctx, 1)}};
+
+        pir::OpInfo phi_kernel_op_info =
+            ctx->GetRegisteredOpInfo(PhiKernelOp::name());
+
+        std::vector<pir::Type> vec_out_types;
+        for (size_t i = 0; i < vec_type.size(); ++i) {
+          vec_out_types.push_back(AllocatedDenseTensorType::get(
+              ctx,
+              phi::TransToPhiPlace(shadow_key.backend()),
+              vec_type[i].dyn_cast<DenseTensorType>()));
+        }
+        auto out_type = pir::VectorType::get(ctx, vec_out_types);
+        pir::Operation* shadow_tensors_op = pir::Operation::Create(
+            {(*map_value_pair)[arg]}, attr_map, {out_type}, phi_kernel_op_info);
+        shadow_tensors_op->set_attribute(
+            "origin_id",
+            pir::Int64Attribute::get(ctx, shadow_tensors_op->id()));
+        new_block->push_back(shadow_tensors_op);
+        (*map_value_pair)[arg] = shadow_tensors_op->result(0);
       }
     }
   }
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -982,6 +982,26 @@ static PyObject *builtin_combine_op(PyObject *self,
   }
 }
 
+static PyObject *builtin_split_op(PyObject *self,
+                                  PyObject *args,
+                                  PyObject *kwargs) {
+  try {
+    VLOG(6) << "Add builtin_split op into program";
+    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
+    // Get Value from args
+    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
+    auto x = CastPyArg2Value(x_obj, "builtin_split", 0, false);
+    CallStackRecorder callstack_recoder("builtin_builtin_split");
+    callstack_recoder.Record();
+    auto static_api_out = paddle::dialect::builtin_split(x);
+    callstack_recoder.AttachToOps();
+    return ToPyObject(static_api_out);
+  } catch (...) {
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
 static PyObject *static_api_fused_gemm_epilogue(PyObject *self,
                                                 PyObject *args,
                                                 PyObject *kwargs) {
@@ -1216,6 +1236,10 @@ static PyMethodDef ManualOpsAPI[] = {
      (PyCFunction)(void (*)(void))builtin_combine_op,
      METH_VARARGS | METH_KEYWORDS,
      "C++ interface function for builtin_combine_op."},
+    {"builtin_split",
+     (PyCFunction)(void (*)(void))builtin_split_op,
+     METH_VARARGS | METH_KEYWORDS,
+     "C++ interface function for builtin_split_op."},
     {"tensorrt_engine",
      (PyCFunction)(void (*)(void))static_api_tensorrt_engine,
      METH_VARARGS | METH_KEYWORDS,
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -79,6 +79,21 @@ def append_full_like(float_value, copy_value, value, state, backward_ops):
             )
             full_like_op = value_grad.get_defining_op()
             backward_ops_ = [full_like_op]
+        elif copy_value.is_combine():
+            values = paddle._C_ops.builtin_split(copy_value)
+            value_grad = []
+            backward_ops_ = []
+            for v in values:
+                grad = paddle.full_like(
+                    v,
+                    float_value,
+                    dtype=v.dtype,
+                )
+                value_grad.append(grad)
+                full_like_op = grad.get_defining_op()
+                full_op = full_like_op.operand_source(1).get_defining_op()
+                backward_ops_.append(full_like_op)
+                backward_ops_.append(full_op)
         else:
             value_grad = paddle.full_like(
                 copy_value,
@@ -93,7 +108,10 @@ def append_full_like(float_value, copy_value, value, state, backward_ops):
             state.op_to_opgrad[value.get_defining_op()],
             backward_ops_,
         )
-        state.value_to_valuegrad[value] = [[value_grad]]
+        if copy_value.is_combine():
+            state.value_to_valuegrad[value] = [value_grad]
+        else:
+            state.value_to_valuegrad[value] = [[value_grad]]
         return value_grad
 
 
