PaddlePaddle
diff --git a/‎paddle/fluid/framework/new_executor/instruction/control_flow/tuple_pop_instruction.cc
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/framework/new_executor/instruction/control_flow/tuple_pop_instruction.cc
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/framework/new_executor/instruction/instruction_util.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/fluid/framework/new_executor/instruction/instruction_util.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/ops.yaml
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/ops.yaml
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/transforms/inplace_pass.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/transforms/inplace_pass.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+167-38Lines changed: 167 additions & 38 deletions b/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+167-38Lines changed: 167 additions & 38 deletions
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/cpu/data_kernel.cc
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions b/‎paddle/phi/kernels/cpu/data_kernel.cc
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions
@@ -28,8 +28,8 @@ TuplePopInstruction::TuplePopInstruction(size_t id,
     : InstructionBase(id, place), op_(op), value_exe_info_(value_exe_info) {
   tuple_pop_op_ = op->dyn_cast<pir::TuplePopOp>();
   VLOG(6) << "construct tuple_pop instruction for: " << tuple_pop_op_->name();
-  auto stack_value = tuple_pop_op_.container();
-  auto var_array = value_exe_info_->GetVarByValue(stack_value);
+  auto outlet_value = tuple_pop_op_.outlet();
+  auto var_array = value_exe_info_->GetVarByValue(outlet_value);
   stack_element_var_array_ = var_array->GetMutable<VariableRefArray>();
 
   std::unordered_map<pir::Value, std::vector<int>> inputs;
 
@@ -281,7 +281,9 @@ std::unordered_set<pir::Value> GetInternalInputs(pir::Block* block) {
     }
     if (op.isa<pir::TuplePopOp>()) {
       auto tuple_pop_op = op.dyn_cast<pir::TuplePopOp>();
-      inner_inputs.insert(tuple_pop_op.container());
+      if (tuple_pop_op.has_container()) {
+        inner_inputs.insert(tuple_pop_op.container());
+      }
     }
     for (size_t i = 0; i < op.num_operands(); ++i) {
       inner_inputs.insert(op.operand_source(i));
 
@@ -1342,6 +1342,7 @@ void PrintValuesAndVariables(
         GetOriginOutputNames(op_name);
 
     // 1. output string
+    VLOG(10) << "Generate output string ...";
     std::string ret_value_str = "Value   : (";
     std::string ret_variable_str = "Variable: (";
     if (!op.results().empty()) {
@@ -1387,10 +1388,12 @@ void PrintValuesAndVariables(
     ret_variable_str += ") = ";
 
     // 2. op name
+    VLOG(10) << "Generate op name ...";
     ret_value_str += op_name;
     ret_variable_str += op_name;
 
     // 3. input string
+    VLOG(10) << "Generate input string ...";
     ret_value_str += "(";
     ret_variable_str += "(";
     if (!op.operands().empty()) {
 
@@ -158,6 +158,7 @@
     'seed',
     'send_v2',
     'shadow_feed',
+    'shadow_feed_tensors',
     'shuffle_batch',
     'sparse_momentum',
     'tdm_sampler',
 
@@ -1415,6 +1415,16 @@
     func: shadow_feed
     param: [x]
 
+- op : shadow_feed_tensors
+  args : (Tensor[] x)
+  output : Tensor[](out){x.size()}
+  infer_meta:
+    func: UnchangedVectorInferMeta
+    param: [x]
+  kernel:
+    func: shadow_feed_tensors
+    param: [x]
+
 - op : share_data
   args : (Tensor x)
   output : Tensor(out)
 
@@ -213,6 +213,7 @@ std::unordered_set<pir::Value> GetSkipDeletionValues(const pir::Block& block) {
       skip_dels.insert(op.result(0));
       continue;
     }
+    // TODO(chenxi67) add logic for shadow_feed_tensors op
     if (upper_op_name == "pd_op.fetch" ||
         upper_op_name == "builtin.shadow_output") {
       skip_dels.insert(op.operand_source(0));
 
@@ -28,6 +28,7 @@
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
 #include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
 #include "paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h"
+#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
@@ -75,6 +76,14 @@ pir::Type ConvertOpTypeToKernelType(pir::IrContext* ctx,
   } else if (op_type.isa<SelectedRowsType>()) {
     return AllocatedSelectedRowsType::get(
         ctx, place, op_type.dyn_cast<SelectedRowsType>());
+  } else if (op_type.isa<pir::VectorType>()) {
+    auto vec_type = op_type.dyn_cast<pir::VectorType>();
+    std::vector<pir::Type> vec_target_type;
+    for (size_t i = 0; i < vec_type.size(); ++i) {
+      vec_target_type.push_back(
+          ConvertOpTypeToKernelType(ctx, vec_type[i], place));
+    }
+    return pir::VectorType::get(ctx, vec_target_type);
   }
   PADDLE_THROW(platform::errors::Unimplemented(
       "Not support op type %s in ConvertOpTypeToKernelType.", op_type));
@@ -1357,6 +1366,119 @@ phi::DataType ParsePhiDType(pir::Type type) {
   }
 }
 
+void AddShadowFeedForValue(
+    size_t index,
+    pir::Operation* op_item,
+    pir::Operation* op_item_with_place,
+    pir::Block* block,
+    pir::IrContext* ctx,
+    std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
+    std::unordered_map<pir::Value, pir::Value>* map_value_pair) {
+  if (op_item->result(index).type().isa<DenseTensorType>()) {
+    phi::KernelKey shadow_key{
+        phi::Backend::GPU,
+        phi::DataLayout::ANY,
+        TransToPhiDataType(
+            op_item->result(index).type().dyn_cast<DenseTensorType>().dtype())};
+    std::unordered_map<std::string, pir::Attribute> attr_map{
+        {"op_name", pir::StrAttribute::get(ctx, "pd_op.shadow_feed")},
+        {"kernel_name", pir::StrAttribute::get(ctx, "shadow_feed")},
+        {"kernel_key", KernelAttribute::get(ctx, shadow_key)}};
+
+    auto out_type = AllocatedDenseTensorType::get(
+        ctx,
+        phi::TransToPhiPlace(shadow_key.backend()),
+        op_item->result(index).type().dyn_cast<DenseTensorType>());
+
+    pir::OpInfo phi_kernel_op_info =
+        ctx->GetRegisteredOpInfo(PhiKernelOp::name());
+    pir::Operation* shadow_op =
+        pir::Operation::Create({op_item_with_place->result(index)},
+                               attr_map,
+                               {out_type},
+                               phi_kernel_op_info);
+    block->push_back(shadow_op);
+    (*map_op_pair)[op_item] = shadow_op;
+    (*map_value_pair)[op_item->result(index)] = shadow_op->result(0);
+  } else if (op_item->result(index).type().isa<pir::VectorType>()) {
+    auto vec_type = op_item->result(index).type().dyn_cast<pir::VectorType>();
+    for (size_t i = 0; i < vec_type.size(); ++i) {
+      PADDLE_ENFORCE_EQ(
+          vec_type[i].isa<DenseTensorType>(),
+          true,
+          phi::errors::PreconditionNotMet(
+              "AddShadowFeedTensors only support DenseTensorType Now"));
+    }
+    // Add ShadowFeedTensors Op
+    phi::KernelKey shadow_key{
+        phi::Backend::GPU,
+        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)}};
+
+    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({op_item_with_place->result(index)},
+                               attr_map,
+                               {out_type},
+                               phi_kernel_op_info);
+    block->push_back(shadow_tensors_op);
+    (*map_op_pair)[op_item] = shadow_tensors_op;
+    (*map_value_pair)[op_item->result(index)] = shadow_tensors_op->result(0);
+  } else {
+    PADDLE_THROW(
+        phi::errors::Unimplemented("AddShadowFeed for value only support "
+                                   "DenseTensorType and VectorType Now"));
+  }
+}
+
+void AddShadowFeedForTuplePopOp(
+    const phi::Place& place,
+    pir::Operation* op_item,
+    pir::Operation* op_item_with_undefined_place,
+    pir::Block* block,
+    pir::IrContext* ctx,
+    std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
+    std::unordered_map<pir::Value, pir::Value>* map_value_pair) {
+  VLOG(4) << "Add AddShadowFeed for op " << op_item->name();
+
+  bool add_shadow_feed = true;
+  if (op_item->attributes().count("place")) {
+    add_shadow_feed = (op_item->attributes()
+                           .at("place")
+                           .dyn_cast<PlaceAttribute>()
+                           .data()
+                           .GetType()) == phi::AllocationType::UNDEFINED;
+  }
+
+  // if value place not gpu, add shadow feed op
+  if (platform::is_gpu_place(place) && add_shadow_feed) {
+    for (size_t i = 0; i < op_item->num_results(); ++i) {
+      AddShadowFeedForValue(i,
+                            op_item,
+                            op_item_with_undefined_place,
+                            block,
+                            ctx,
+                            map_op_pair,
+                            map_value_pair);
+    }
+  }
+}
+
 void HandleForSpecialOp(
     const phi::Place& place,
     pir::Operation* op_item,
@@ -1627,17 +1749,46 @@ void HandleForSpecialOp(
     }
 
     auto pop_back_op = op_item->dyn_cast<::pir::TuplePopOp>();
-    for (size_t i = 0; i < op_item->num_results(); ++i) {
-      auto cur_inlet_element = pop_back_op.inlet_element(i);
-      PADDLE_ENFORCE_EQ(map_value_pair->count(cur_inlet_element),
-                        true,
-                        phi::errors::PreconditionNotMet(
-                            "[%d]'s output of [%s] op MUST be in map pair",
-                            i,
-                            op_item->name()));
-      auto new_inlet_element = map_value_pair->at(cur_inlet_element);
 
-      op_output_types.push_back(new_inlet_element.type());
+    if (pop_back_op.has_container()) {
+      // if TuplePopOp and TuplePushOp are in the same sub_program
+      for (size_t i = 0; i < op_item->num_results(); ++i) {
+        auto cur_inlet_element = pop_back_op.inlet_element(i);
+        PADDLE_ENFORCE_EQ(map_value_pair->count(cur_inlet_element),
+                          true,
+                          phi::errors::PreconditionNotMet(
+                              "[%d]'s output of [%s] op MUST be in map pair",
+                              i,
+                              op_item->name()));
+        auto new_inlet_element = map_value_pair->at(cur_inlet_element);
+
+        op_output_types.push_back(new_inlet_element.type());
+      }
+    } else {
+      VLOG(4) << "TuplePopOp and TuplePushOp are in different sub_program.";
+      for (size_t i = 0; i < op_item->num_results(); ++i) {
+        auto cur_inlet_element = op_item->result(i);
+        auto out_place = phi::TransToPhiPlace(phi::Backend::UNDEFINED);
+        pir::Type new_inlet_element_type =
+            ConvertOpTypeToKernelType(ctx, cur_inlet_element.type(), out_place);
+        op_output_types.push_back(new_inlet_element_type);
+      }
+
+      pir::OpInfo op_info = ctx->GetRegisteredOpInfo(op_item->name());
+      pir::Operation* op = pir::Operation::Create(
+          vec_inputs, op_item->attributes(), op_output_types, op_info);
+
+      block->push_back(op);
+      (*map_op_pair)[op_item] = op;
+      // only deal with single output
+      if (op_item->num_results() > 0) {
+        for (size_t i = 0; i < op_item->num_results(); ++i) {
+          (*map_value_pair)[op_item->result(i)] = op->result(i);
+        }
+      }
+      AddShadowFeedForTuplePopOp(
+          place, op_item, op, block, ctx, map_op_pair, map_value_pair);
+      return;
     }
   }
 
@@ -2311,34 +2462,12 @@ void AddShadowFeedOpForDataOrFeed(
            .GetType() == phi::AllocationType::UNDEFINED);
   bool add_shadow_feed = feed_op_add_shadow_feed || data_op_add_shadow_feed;
   if (add_shadow_feed) {
-    // if shadow data op place not gpu,add shadow feed op
-    phi::KernelKey shadow_key{
-        phi::Backend::GPU,
-        phi::DataLayout::ANY,
-        TransToPhiDataType(
-            op_item->result(0).type().dyn_cast<DenseTensorType>().dtype())};
-    std::unordered_map<std::string, pir::Attribute> attr_map{
-        {"op_name", pir::StrAttribute::get(ctx, "pd_op.shadow_feed")},
-        {"kernel_name", pir::StrAttribute::get(ctx, "shadow_feed")},
-        {"kernel_key", KernelAttribute::get(ctx, shadow_key)}};
-
-    auto out_type = AllocatedDenseTensorType::get(
-        ctx,
-        phi::TransToPhiPlace(shadow_key.backend()),
-        op_item->result(0).type().dyn_cast<DenseTensorType>());
-
-    pir::OpInfo phi_kernel_op_info =
-        ctx->GetRegisteredOpInfo(PhiKernelOp::name());
-    pir::Operation* shadow_op = pir::Operation::Create(
-        {kernel_op->result(0)}, attr_map, {out_type}, phi_kernel_op_info);
-
-    (*map_op_pair)[op_item] = shadow_op;
-    block->push_back(shadow_op);
-    if (op_item->num_results() > 0) {
-      for (size_t i = 0; i < shadow_op->num_results(); ++i) {
-        (*map_value_pair)[op_item->result(i)] = shadow_op->result(i);
-      }
-    }
+    PADDLE_ENFORCE(op_item->num_results() == 1,
+                   phi::errors::PreconditionNotMet(
+                       "op_item should have only one result, but got %d",
+                       op_item->num_results()));
+    AddShadowFeedForValue(
+        0, op_item, kernel_op, block, ctx, map_op_pair, map_value_pair);
   }
 }
 
 
@@ -4885,6 +4885,14 @@ void UnchangedArrayInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_layout(x.layout());
 }
 
+void UnchangedVectorInferMeta(const std::vector<const MetaTensor*>& xs,
+                              std::vector<MetaTensor*> outs) {
+  for (size_t i = 0; i < xs.size(); ++i) {
+    outs[i]->set_dtype(xs[i]->dtype());
+    outs[i]->set_layout(xs[i]->layout());
+  }
+}
+
 // meta x -> out without change, check if axis in range [-Rank(x), Rank(x)-1]
 void UnchangedInferMetaCheckAxis(const MetaTensor& x,
                                  int axis,
 
@@ -755,6 +755,8 @@ void UnchangedExceptLayoutInferMeta(const MetaTensor& x, MetaTensor* out);
 void UnchangedExceptDtypeInferMeta(const MetaTensor& x, MetaTensor* out);
 void UnchangedInferMeta(const MetaTensor& x, MetaTensor* out);
 void UnchangedArrayInferMeta(const MetaTensor& x, MetaTensor* out);
+void UnchangedVectorInferMeta(const std::vector<const MetaTensor*>& xs,
+                              std::vector<MetaTensor*> outs);
 
 // meta x -> out without change, check if axis in range [-Rank(x), Rank(x)-1]
 void UnchangedInferMetaCheckAxis(const MetaTensor& x,
 
@@ -70,6 +70,23 @@ PD_REGISTER_KERNEL(shadow_feed,
                    phi::complex64,
                    phi::complex128) {}
 
+PD_REGISTER_KERNEL(shadow_feed_tensors,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ShadowFeedTensorsKernel,
+                   bool,
+                   uint8_t,
+                   float,
+                   int8_t,
+                   int16_t,
+                   int32_t,
+                   int64_t,
+                   double,
+                   phi::float16,
+                   phi::bfloat16,
+                   phi::complex64,
+                   phi::complex128) {}
+
 PD_REGISTER_KERNEL(print_kernel,
                    CPU,
                    ALL_LAYOUT,
Original file line number	Diff line number	Diff line change
`@@ -281,7 +281,9 @@ std::unordered_set<pir::Value> GetInternalInputs(pir::Block* block) {`
`281`	`281`	`}`
`282`	`282`	`if (op.isa<pir::TuplePopOp>()) {`
`283`	`283`	`auto tuple_pop_op = op.dyn_cast<pir::TuplePopOp>();`
`284`		`- inner_inputs.insert(tuple_pop_op.container());`
	`284`	`+ if (tuple_pop_op.has_container()) {`
	`285`	`+ inner_inputs.insert(tuple_pop_op.container());`
	`286`	`+ }`
`285`	`287`	`}`
`286`	`288`	`for (size_t i = 0; i < op.num_operands(); ++i) {`
`287`	`289`	`inner_inputs.insert(op.operand_source(i));`
Original file line number	Diff line number	Diff line change
`@@ -213,6 +213,7 @@ std::unordered_set<pir::Value> GetSkipDeletionValues(const pir::Block& block) {`
`213`	`213`	`skip_dels.insert(op.result(0));`
`214`	`214`	`continue;`
`215`	`215`	`}`
	`216`	`+ // TODO(chenxi67) add logic for shadow_feed_tensors op`
`216`	`217`	`if (upper_op_name == "pd_op.fetch" \|\|`
`217`	`218`	`upper_op_name == "builtin.shadow_output") {`
`218`	`219`	`skip_dels.insert(op.operand_source(0));`