PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+178-3Lines changed: 178 additions & 3 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+178-3Lines changed: 178 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+34Lines changed: 34 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+34Lines changed: 34 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/manual_static_op_function.h
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions b/‎paddle/fluid/pybind/manual_static_op_function.h
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions
diff --git a/‎paddle/phi/core/tensor_array.cc
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎paddle/phi/core/tensor_array.cc
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
diff --git a/‎paddle/phi/core/tensor_array.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/core/tensor_array.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+16-1Lines changed: 16 additions & 1 deletion b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+16-1Lines changed: 16 additions & 1 deletion
@@ -225,5 +225,17 @@ pir::Value assign(const pir::Value& x) {
   }
 }
 
+pir::Value array_pop(pir::Value input, int index) {
+  if (input.type().isa<paddle::dialect::DenseTensorArrayType>()) {
+    paddle::dialect::ArrayPopOp array_pop_op =
+        ApiBuilder::Instance().GetBuilder()->Build<paddle::dialect::ArrayPopOp>(
+            input, index);
+    return array_pop_op.result(1);
+  } else {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "pop only supports DenseTensorArrayType."));
+  }
+}
+
 }  // namespace dialect
 }  // namespace paddle
@@ -87,5 +87,7 @@ pir::Value slice_array_dense(pir::Value input, pir::Value starts);
 
 pir::Value assign(const pir::Value& x);
 
+pir::Value array_pop(pir::Value input, int index);
+
 }  // namespace dialect
 }  // namespace paddle
@@ -24,7 +24,8 @@ paddle::dialect::AddNOp, paddle::dialect::AddN_Op,
     paddle::dialect::AssignArrayOp, paddle::dialect::AssignArray_Op,
     paddle::dialect::ArrayToTensorOp, paddle::dialect::TensorToArrayOp,
     paddle::dialect::IncrementOp, paddle::dialect::Increment_Op,
-    paddle::dialect::ShapeBroadcastOp, paddle::dialect::MemcpyD2hMultiIoOp
+    paddle::dialect::ShapeBroadcastOp, paddle::dialect::MemcpyD2hMultiIoOp,
+    paddle::dialect::ArrayPopOp
 #else
 
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
@@ -2052,7 +2053,7 @@ void ArrayReadOp::VerifySig() {
         phi::errors::PreconditionNotMet(
             "Type validation failed for the 0th output."));
   }
-  VLOG(4) << "End Verifying for: ArrayWrite_Op.";
+  VLOG(4) << "End Verifying for: ArrayReadOp.";
 }
 
 void ArrayReadOp::InferMeta(phi::InferMetaContext *infer_meta) {
@@ -2860,7 +2861,7 @@ void SliceArrayOp::Build(pir::Builder &builder,             // NOLINT
                          pir::Value input,
                          pir::Value starts,
                          pir::Value ends) {
-  VLOG(4) << "Start build SliceArrayDenseOp";
+  VLOG(4) << "Start build SliceArrayOp";
   VLOG(4) << "Builder construction inputs";
   std::vector<pir::Value> argument_inputs = {input, starts, ends};
   argument.AddInputs(argument_inputs);
@@ -4435,6 +4436,179 @@ phi::DataType MemcpyD2hMultiIoOp::GetKernelTypeForVar(
   return expected_kernel_dtype;
 }
 
+const char *ArrayPopOp::attributes_name[1] = {"index"};
+
+OpInfoTuple ArrayPopOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo("array",
+                                   "paddle::dialect::DenseTensorArrayType",
+                                   false,
+                                   false,
+                                   false,
+                                   false)};
+
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo("index", "pir::Int32Attribute", ""),
+  };
+
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      paddle::dialect::OpOutputInfo(
+          "array_out", "paddle::dialect::DenseTensorArrayType", false, false),
+      paddle::dialect::OpOutputInfo(
+          "out", "paddle::dialect::DenseTensorType", false, false)};
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      paddle::dialect::OpRunTimeInfo("ArrayPopInferMeta",
+                                     {"array", "index"},
+                                     "array_pop",
+                                     {"array", "index"},
+                                     {"array"},
+                                     {"array"},
+                                     {{"array_out", "array"}},
+                                     {});
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "array_pop");
+}
+
+void ArrayPopOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: ArrayPopOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    IR_ENFORCE(input_size == 1u,
+               "The size %d of inputs must be equal to 1.",
+               input_size);
+    IR_ENFORCE((*this)
+                   ->operand_source(0)
+                   .type()
+                   .isa<paddle::dialect::DenseTensorArrayType>(),
+               "Type validation failed for the 0th input, got %s.",
+               (*this)->operand_source(0).type());
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    IR_ENFORCE(attributes.count("index") > 0, "index does not exist.");
+    IR_ENFORCE(attributes.at("index").isa<pir::Int32Attribute>(),
+               "Type of attribute: index is not pir::Int32Attribute.");
+  }
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    IR_ENFORCE(output_size == 2u,
+               "The size %d of outputs must be equal to 2.",
+               output_size);
+    IR_ENFORCE(
+        (*this)->result(0).type().isa<paddle::dialect::DenseTensorArrayType>(),
+        "Type validation failed for the 0th output.");
+    IR_ENFORCE(
+        (*this)->result(1).type().isa<paddle::dialect::DenseTensorType>(),
+        "Type validation failed for the 1st output.");
+  }
+  VLOG(4) << "End Verifying for: ArrayPopOp.";
+}
+
+void ArrayPopOp::Build(pir::Builder &builder,             // NOLINT
+                       pir::OperationArgument &argument,  // NOLINT
+                       pir::Value array,
+                       int index) {
+  VLOG(4) << "Start build ArrayPopOp";
+  VLOG(4) << "Builder construction inputs";
+  std::vector<pir::Value> argument_inputs = {array};
+  argument.AddInputs({array});
+  VLOG(4) << "Builder construction attributes";
+  pir::AttributeMap argument_attributes = {};
+  pir::Attribute attr_index =
+      pir::Int32Attribute::get(pir::IrContext::Instance(), index);
+  argument.AddAttribute("index", attr_index);
+  argument_attributes.insert({"index", attr_index});
+  std::vector<pir::Type> argument_outputs =
+      ArrayPopOp::InferMeta(argument_inputs, argument_attributes);
+
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
+}
+
+void ArrayPopOp::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::ArrayPopInferMeta);
+  fn(infer_meta);
+}
+
+std::vector<pir::Type> ArrayPopOp::InferMeta(
+    const std::vector<pir::Value> &input_values,
+    const pir::AttributeMap &attributes) {
+  VLOG(4) << "Start infermeta ArrayPopOp";
+  IR_ENFORCE(input_values.size() == 1,
+             "Num of inputs is expected to be 1 but got %d.",
+             input_values.size());
+  pir::Value input = input_values[0];
+
+  VLOG(4) << "Builder construction outputs";
+  paddle::dialect::DenseTensorArrayType input_type;
+  if (input.type().isa<paddle::dialect::DenseTensorArrayType>()) {
+    input_type = input.type().dyn_cast<paddle::dialect::DenseTensorArrayType>();
+    (void)input_type;
+  } else if (input.type()
+                 .isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
+    paddle::dialect::AllocatedDenseTensorArrayType allocated_input =
+        input.type().dyn_cast<paddle::dialect::AllocatedDenseTensorArrayType>();
+    input_type = paddle::dialect::DenseTensorArrayType::get(
+        pir::IrContext::Instance(),
+        allocated_input.dtype(),
+        allocated_input.data_layout());
+    (void)input_type;
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "Only support paddle::dialect::DenseTensorArrayType or "
+        "paddle::dialect::AllocatedDenseTensorArrayType"));
+  }
+
+  IR_ENFORCE(attributes.find("index") != attributes.end(),
+             "'index' Attribute is expected for ArrayPopOp. ");
+  int index = attributes.at("index").dyn_cast<pir::Int32Attribute>().data();
+
+  paddle::dialect::IrTensor dense_input(
+      paddle::dialect::TransToPhiDataType(input_type.dtype()),
+      {},
+      input_type.data_layout(),
+      {});
+  paddle::dialect::IrMetaTensor meta_input(&dense_input);
+  paddle::dialect::IrTensor array_out;
+  paddle::dialect::IrMetaTensor meta_array_out(&array_out);
+  paddle::dialect::IrTensor dense_out;
+  paddle::dialect::IrMetaTensor meta_out(&dense_out);
+
+  phi::ArrayPopInferMeta(meta_input,
+                         index,
+                         &meta_array_out,
+                         &meta_out,
+                         phi::MetaConfig(false, false));
+
+  std::vector<pir::Type> argument_outputs;
+  pir::Type out_array_type = paddle::dialect::DenseTensorArrayType::get(
+      pir::IrContext::Instance(),
+      paddle::dialect::TransToIrDataType(array_out.dtype()),
+      array_out.layout());
+  pir::Type out_dense_tensor_type = paddle::dialect::DenseTensorType::get(
+      pir::IrContext::Instance(),
+      paddle::dialect::TransToIrDataType(dense_out.dtype()),
+      dense_out.dims(),
+      dense_out.layout(),
+      dense_out.lod(),
+      dense_out.offset());
+  argument_outputs.push_back(out_array_type);
+  argument_outputs.push_back(out_dense_tensor_type);
+  return argument_outputs;
+}
+
+phi::DataType ArrayPopOp::GetKernelTypeForVar(
+    const std::string &var_name,
+    const phi::DataType &tensor_dtype,
+    const phi::DataType &expected_kernel_dtype) {
+  VLOG(4) << "Get KernelType for Var of op: ArrayPopOp";
+
+  return expected_kernel_dtype;
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
@@ -4462,4 +4636,5 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::Increment_Op)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::MemcpyD2hMultiIoOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ShapeBroadcastOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayPopOp)
 #endif
@@ -779,6 +779,39 @@ class IR_API ShapeBroadcastOp
   bool InferSymbolicShape(pir::ShapeConstraintIRAnalysis *shape_analysis);
 };
 
+class ArrayPopOp : public pir::Op<ArrayPopOp,
+                                  paddle::dialect::OpYamlInfoInterface,
+                                  paddle::dialect::InferMetaInterface,
+                                  paddle::dialect::GetKernelTypeForVarInterface,
+                                  InplaceTrait> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.array_pop"; }
+  static const char *attributes_name[1];
+  static constexpr uint32_t attributes_num = 1;
+  static OpInfoTuple GetOpInfo();
+  void VerifySig();
+
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value input,
+                    int index);
+
+  static phi::DataType GetKernelTypeForVar(
+      const std::string &var_name,
+      const phi::DataType &tensor_dtype,
+      const phi::DataType &expected_kernel_dtype);
+
+  pir::Value input() { return operand_source(0); }
+  pir::Value array_out() { return result(0); }
+  pir::Value out() { return result(1); }
+
+  static void InferMeta(phi::InferMetaContext *infer_meta);
+  static std::vector<pir::Type> InferMeta(
+      const std::vector<pir::Value> &input_values,
+      const pir::AttributeMap &attributes);
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
@@ -806,3 +839,4 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::Increment_Op)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ShapeBroadcastOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::MemcpyD2hMultiIoOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayPopOp)
@@ -754,6 +754,30 @@ static PyObject *run_custom_op(PyObject *self,
   }
 }
 
+static PyObject *static_api_array_pop(PyObject *self,
+                                      PyObject *args,
+                                      PyObject *kwargs) {
+  try {
+    VLOG(6) << "Add array_pop op into program";
+    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
+
+    // Get Value from args
+    PyObject *input_obj = PyTuple_GET_ITEM(args, 0);
+    auto input = CastPyArg2Value(input_obj, "array_pop", 0);
+
+    PyObject *index_obj = PyTuple_GET_ITEM(args, 1);
+    auto index = CastPyArg2Int(index_obj, "array_pop", 1);
+
+    // Call ir static api
+    auto static_api_out = paddle::dialect::array_pop(input, index);
+
+    return ToPyObject(static_api_out);
+  } catch (...) {
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
 static PyMethodDef ManualOpsAPI[] = {
     {"set_parameter",
      (PyCFunction)(void (*)(void))static_api_set_parameter,
@@ -803,6 +827,10 @@ static PyMethodDef ManualOpsAPI[] = {
      (PyCFunction)(void (*)(void))run_custom_op,
      METH_VARARGS | METH_KEYWORDS,
      "C++ interface function for run_custom_op."},
+    {"array_pop",
+     (PyCFunction)(void (*)(void))static_api_array_pop,
+     METH_VARARGS | METH_KEYWORDS,
+     "C++ interface function for array_pop."},
     {nullptr, nullptr, 0, nullptr}};
 
 }  // namespace pybind
 
@@ -103,6 +103,16 @@ void TensorArray::push_back(const DenseTensor& tensor) {
   tensors_.push_back(tensor);
 }
 
+void TensorArray::pop(size_t i) {
+  PADDLE_ENFORCE_LT(i,
+                    tensors_.size(),
+                    errors::OutOfRange("The size of TensorArray is %d, "
+                                       "but the received index is %d.",
+                                       tensors_.size(),
+                                       i));
+  tensors_.erase(tensors_.begin() + i);
+}
+
 void TensorArray::emplace_back(const DenseTensor& tensor) {
   tensors_.emplace_back(tensor);
 }
 
@@ -111,6 +111,8 @@ class TensorArray : public TensorBase,
 
   void emplace_back(const DenseTensor& tensor);
 
+  void pop(size_t i);
+
   /// \brief Return the last tensor in TensorArray
   DenseTensor& back() { return tensors_.back(); }
 
 
@@ -3755,7 +3755,7 @@ void NumelInferMeta(const MetaTensor& input, MetaTensor* out) {
 }
 
 // This logic is copied from
-// `paddle/fluid/operators/slice_op.cc/Slcieop::InferShape`
+// `paddle/fluid/operators/slice_op.cc/Sliceop::InferShape`
 void SliceArrayInferMeta(const MetaTensor& input,
                          const IntArray& starts,
                          const IntArray& ends,
@@ -5462,6 +5462,21 @@ void LrnInferMeta(const MetaTensor& x,
   mid_out->share_lod(x);
 }
 
+void ArrayPopInferMeta(const MetaTensor& array,
+                       int index,
+                       MetaTensor* array_out,
+                       MetaTensor* out,
+                       MetaConfig config) {
+  array_out->set_dtype(array.dtype());
+  array_out->set_layout(array.layout());
+  if (config.is_runtime) {
+    return;
+  }
+  // TODO(dev): Support get a tensor dims from TensorArray.
+  // out->set_dims(x.dims());
+  out->set_dtype(array.dtype());
+}
+
 }  // namespace phi
 
 PD_REGISTER_INFER_META_FN(flatten, phi::FlattenInferMeta);