[PIR] Add array_to_tensor op (PaddlePaddle#59270)

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

paddle/fluid/pir/dialect/op_generator/api_gen.py

@@ -172,6 +172,7 @@ def _parse_yaml(self, op_yaml_files, op_compat_yaml_file):
         # replace old ir ops with pir ops
         if need_update_ops:
             update_ops(op_yaml_items, update_yaml_file)
+
         op_info_items = []
         for op in op_yaml_items:
             op_compat_item = op_compat_parser.get_compat(op['name'])

paddle/fluid/pir/dialect/operator/ir/manual_api.cc

@@ -148,5 +148,15 @@ pir::OpResult array_write_(pir::Value array, pir::Value x, pir::Value i) {
   return array_write_op.out();
 }
 
+std::tuple<pir::OpResult, pir::OpResult> array_to_tensor(pir::Value x,
+                                                         int axis,
+                                                         bool use_stack) {
+  auto array_to_tensor =
+      ApiBuilder::Instance()
+          .GetBuilder()
+          ->Build<paddle::dialect::ArrayToTensorOp>(x, axis, use_stack);
+  return std::make_tuple(array_to_tensor.result(0), array_to_tensor.result(1));
+}
+
 }  // namespace dialect
 }  // namespace paddle

paddle/fluid/pir/dialect/operator/ir/manual_api.h

@@ -68,5 +68,9 @@ pir::OpResult array_read(pir::Value array, pir::Value i);
 
 pir::OpResult array_write_(pir::Value array, pir::Value x, pir::Value i);
 
+std::tuple<pir::OpResult, pir::OpResult> array_to_tensor(pir::Value x,
+                                                         int axis,
+                                                         bool use_stack);
+
 }  // namespace dialect
 }  // namespace paddle

paddle/fluid/pir/dialect/operator/ir/manual_op.cc

@@ -1533,6 +1533,153 @@ void ArrayWrite_Op::InferMeta(phi::InferMetaContext *infer_meta) {
   fn(infer_meta);
 }
 
+const char *ArrayToTensorOp::attributes_name[2] = {"axis", "use_stack"};
+
+OpInfoTuple ArrayToTensorOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo("x",
+                                   "paddle::dialect::DenseTensorArrayType",
+                                   false,
+                                   false,
+                                   false,
+                                   true)};
+
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo("axis", "pir::Int32Attribute", ""),
+      paddle::dialect::OpAttributeInfo("use_stack", "pir::BoolAttribute", "")};
+
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      paddle::dialect::OpOutputInfo(
+          "out", "paddle::dialect::DenseTensorType", false, false),
+      paddle::dialect::OpOutputInfo(
+          "out_index", "paddle::dialect::DenseTensorType", false, false)};
+
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      paddle::dialect::OpRunTimeInfo("ArrayToTensorInferMeta",
+                                     {"x", "axis", "use_stack"},
+                                     "array_to_tensor",
+                                     {"x", "axis", "use_stack"},
+                                     {"x"},
+                                     {},
+                                     {},
+                                     {});
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "array_to_tensor");
+}
+
+void ArrayToTensorOp::Build(pir::Builder &builder,             // NOLINT
+                            pir::OperationArgument &argument,  // NOLINT
+                            pir::Value x,
+                            int axis,
+                            bool use_stack) {
+  VLOG(4) << "Start build ArrayToTensorOp";
+  VLOG(4) << "Builder construction inputs";
+  argument.AddInputs({x});
+
+  VLOG(4) << "Builder construction attributes";
+  pir::Attribute attr_axis =
+      pir::Int32Attribute::get(pir::IrContext::Instance(), axis);
+  argument.AddAttribute("axis", attr_axis);
+  pir::Attribute attr_use_stack =
+      pir::BoolAttribute::get(pir::IrContext::Instance(), use_stack);
+  argument.AddAttribute("use_stack", attr_use_stack);
+
+  VLOG(4) << "Builder construction outputs";
+  paddle::dialect::DenseTensorArrayType x_type =
+      x.type().dyn_cast<paddle::dialect::DenseTensorArrayType>();
+  paddle::dialect::IrTensor dense_x(
+      paddle::dialect::TransToPhiDataType(x_type.dtype()),
+      {},
+      x_type.data_layout(),
+      {});
+  paddle::dialect::IrMetaTensor meta_x(&dense_x);
+
+  paddle::dialect::IrTensor dense_out;
+  paddle::dialect::IrMetaTensor meta_out(&dense_out);
+
+  paddle::dialect::IrTensor dense_out_index;
+  paddle::dialect::IrMetaTensor meta_out_index(&dense_out_index);
+
+  phi::ArrayToTensorInferMeta(meta_x,
+                              axis,
+                              use_stack,
+                              &meta_out,
+                              &meta_out_index,
+                              phi::MetaConfig(false, false));
+
+  std::vector<pir::Type> argument_outputs;
+  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_dense_tensor_type);
+  pir::Type out_index_dense_tensor_type = paddle::dialect::DenseTensorType::get(
+      pir::IrContext::Instance(),
+      paddle::dialect::TransToIrDataType(dense_out_index.dtype()),
+      dense_out_index.dims(),
+      dense_out_index.layout(),
+      dense_out_index.lod(),
+      dense_out_index.offset());
+  argument_outputs.push_back(out_index_dense_tensor_type);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+}
+
+void ArrayToTensorOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: "
+             "ArrayToTensorOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    PADDLE_ENFORCE_EQ(
+        input_size,
+        1u,
+        phi::errors::PreconditionNotMet(
+            "The size %d of inputs must be equal to 1.", input_size));
+
+    PADDLE_ENFORCE((*this)
+                       ->operand_source(0)
+                       .type()
+                       .isa<paddle::dialect::DenseTensorArrayType>(),
+                   phi::errors::PreconditionNotMet(
+                       "Type validation failed for the 0th input."));
+  }
+
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    PADDLE_ENFORCE(attributes.count("axis") > 0, "axis does not exist.");
+    PADDLE_ENFORCE(attributes.count("use_stack") > 0,
+                   "use_stack does not exist.");
+  }
+
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    PADDLE_ENFORCE_EQ(
+        output_size,
+        2u,
+        phi::errors::PreconditionNotMet(
+            "The size %d of outputs must be equal to 1.", output_size));
+    PADDLE_ENFORCE(
+        (*this)->result(0).type().isa<paddle::dialect::DenseTensorType>(),
+        phi::errors::PreconditionNotMet(
+            "Type validation failed for the 0th output."));
+    PADDLE_ENFORCE(
+        (*this)->result(1).type().isa<paddle::dialect::DenseTensorType>(),
+        phi::errors::PreconditionNotMet(
+            "Type validation failed for the 0th output."));
+  }
+  VLOG(4) << "End Verifying for: ArrayToTensorOp.";
+}
+
+void ArrayToTensorOp::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::ArrayToTensorInferMeta);
+  fn(infer_meta);
+}
+
 OpInfoTuple ExpandOp::GetOpInfo() {
   std::vector<paddle::dialect::OpInputInfo> inputs = {
       paddle::dialect::OpInputInfo(
@@ -1807,4 +1954,5 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::CreateArrayOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayLengthOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayReadOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayWrite_Op)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayToTensorOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)

paddle/fluid/pir/dialect/operator/ir/manual_op.h

@@ -261,6 +261,26 @@ class ArrayWrite_Op : public pir::Op<ArrayWrite_Op,
   static void InferMeta(phi::InferMetaContext *infer_meta);
 };
 
+class ArrayToTensorOp
+    : public pir::Op<ArrayToTensorOp, OpYamlInfoInterface, InferMetaInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.array_to_tensor"; }
+  static constexpr uint32_t attributes_num = 2;
+  static const char *attributes_name[attributes_num];
+  static OpInfoTuple GetOpInfo();
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x,
+                    int axis,
+                    bool use_stack);
+  void VerifySig();
+  pir::Value x() { return operand_source(0); }
+  pir::OpResult out() { return result(0); }
+  pir::OpResult out_index() { return result(2); }
+  static void InferMeta(phi::InferMetaContext *infer_meta);
+};
+
 class ExpandOp : public pir::Op<ExpandOp,
                                 paddle::dialect::OpYamlInfoInterface,
                                 paddle::dialect::InferMetaInterface,
@@ -320,4 +340,5 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::CreateArrayOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayLengthOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayReadOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayWrite_Op)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayToTensorOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)

paddle/fluid/pir/dialect/operator/ir/op_dialect.cc

@@ -58,10 +58,12 @@ void OperatorDialect::initialize() {
 #define GET_OP_LIST
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.cc"  // NOLINT
       >();
+
   RegisterOps<
 #define GET_OP_LIST
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc"  // NOLINT
       >();
+
   RegisterOps<paddle::dialect::AddNOp,
               paddle::dialect::AddN_Op,
               paddle::dialect::AddNWithKernelOp,
@@ -72,7 +74,8 @@ void OperatorDialect::initialize() {
               paddle::dialect::CreateArrayOp,
               paddle::dialect::ArrayLengthOp,
               paddle::dialect::ArrayReadOp,
-              paddle::dialect::ArrayWrite_Op>();
+              paddle::dialect::ArrayWrite_Op,
+              paddle::dialect::ArrayToTensorOp>();
 
   RegisterInterfaces<ParameterConvertInterface>();
 }

paddle/fluid/pybind/manual_static_op_function.h

@@ -235,6 +235,45 @@ static PyObject *static_api_array_write_(PyObject *self,
   }
 }
 
+static PyObject *static_api_array_to_tensor(PyObject *self,
+                                            PyObject *args,
+                                            PyObject *kwargs) {
+  try {
+    VLOG(6) << "Add array_to_tensor op into program";
+    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
+
+    // Get Value from args
+    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
+    pir::Value x;
+    if (PyObject_CheckIROpResult(x_obj)) {
+      x = CastPyArg2Value(x_obj, "array_to_tensor", 0);
+    } else if (PyObject_CheckIRVectorOfOpResult(x_obj)) {
+      std::vector<pir::Value> x_tmp =
+          CastPyArg2VectorOfValue(x_obj, "array_to_tensor", 0);
+      if (x_tmp.size() != 1) {
+        PADDLE_THROW(platform::errors::InvalidArgument(
+            "Input x expects only one input, but %d are given.",
+            x_tmp.size()));  // NOLINT
+      }
+      x = x_tmp[0];
+    }
+
+    PyObject *axis_obj = PyTuple_GET_ITEM(args, 1);
+    auto axis = CastPyArg2Int(axis_obj, "array_to_tensor", 1);
+
+    PyObject *use_stack_obj = PyTuple_GET_ITEM(args, 2);
+    auto use_stack = CastPyArg2Boolean(use_stack_obj, "array_to_tensor", 2);
+
+    // Call ir static api
+    auto static_api_out = paddle::dialect::array_to_tensor(x, axis, use_stack);
+
+    return ToPyObject(static_api_out);
+  } catch (...) {
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
 static PyMethodDef ManualOpsAPI[] = {
     {"set_parameter",
      (PyCFunction)(void (*)(void))static_api_set_parameter,
@@ -260,6 +299,10 @@ static PyMethodDef ManualOpsAPI[] = {
      (PyCFunction)(void (*)(void))static_api_array_write_,
      METH_VARARGS | METH_KEYWORDS,
      "C++ interface function for array_write_."},
+    {"array_to_tensor",
+     (PyCFunction)(void (*)(void))static_api_array_to_tensor,
+     METH_VARARGS | METH_KEYWORDS,
+     "C++ interface function for array_to_tensor."},
     {nullptr, nullptr, 0, nullptr}};
 
 }  // namespace pybind

paddle/phi/infermeta/unary.cc

@@ -149,6 +149,32 @@ void ArrayLengthInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_dims(make_ddim({1}));
 }
 
+void ArrayToTensorInferMeta(const MetaTensor& x,
+                            int axis,
+                            bool use_stack,
+                            MetaTensor* out,
+                            MetaTensor* out_index,
+                            MetaConfig config) {
+  if (config.is_runtime) return;
+  auto dims = x.dims();
+  // if the shape is empty
+  if (dims == phi::make_ddim({0UL})) return;
+  // otherwise, suppose the shape of array is the shape of tensor in the
+  // array, which is consistent with what tensor_array_read_write dose
+  if (use_stack) {
+    auto dim_vec = phi::vectorize<int>(dims);
+    // use -1 for the stack dim size
+    dim_vec.insert(dim_vec.begin() + axis, -1);
+    dims = phi::make_ddim(dim_vec);
+  } else {
+    // use -1 for the concat dim size
+    dims[axis] = -1;
+  }
+  out->set_dims(dims);
+  out_index->set_dtype(DataType::INT32);
+  out_index->set_dims(phi::make_ddim({-1}));
+}
+
 void ArgMinMaxInferMeta(const MetaTensor& x,
                         const Scalar& axis,
                         bool keepdims,

paddle/phi/infermeta/unary.h

@@ -61,6 +61,13 @@ void ArgsortInferMeta(const MetaTensor& input,
 
 void ArrayLengthInferMeta(const MetaTensor& x, MetaTensor* out);
 
+void ArrayToTensorInferMeta(const MetaTensor& x,
+                            int axis,
+                            bool use_stack,
+                            MetaTensor* out,
+                            MetaTensor* out_index,
+                            MetaConfig config = MetaConfig());
+
 void AsRealInferMeta(const MetaTensor& input, MetaTensor* output);
 
 void AsComplexInferMeta(const MetaTensor& input, MetaTensor* output);