[PIR] add slice_array_dense api (#60433)

zhangbo9674 · web-flow · commit df4ca858521e · 2024-01-02T15:25:45.000+08:00
* fix

* fix
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.cc b/paddle/fluid/pir/dialect/operator/ir/manual_api.cc
@@ -165,5 +165,12 @@ std::tuple<pir::OpResult, pir::OpResult> array_to_tensor(pir::Value x,
   return std::make_tuple(array_to_tensor.result(0), array_to_tensor.result(1));
 }
 
+pir::OpResult slice_array_dense(pir::Value input, pir::Value starts) {
+  auto op = ApiBuilder::Instance()
+                .GetBuilder()
+                ->Build<paddle::dialect::SliceArrayDenseOp>(input, starts);
+  return op.result(0);
+}
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.h b/paddle/fluid/pir/dialect/operator/ir/manual_api.h
@@ -72,5 +72,7 @@ std::tuple<pir::OpResult, pir::OpResult> array_to_tensor(pir::Value x,
                                                          int axis,
                                                          bool use_stack);
 
+pir::OpResult slice_array_dense(pir::Value input, pir::Value starts);
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -1807,7 +1807,7 @@ OpInfoTuple SliceArrayDenseOp::GetOpInfo() {
       paddle::dialect::OpOutputInfo(
           "out", "paddle::dialect::DenseTensorType", false, false)};
   paddle::dialect::OpRunTimeInfo run_time_info =
-      paddle::dialect::OpRunTimeInfo("SliceArrayInferMeta",
+      paddle::dialect::OpRunTimeInfo("SliceArrayDenseInferMeta",
                                      {"input", "starts"},
                                      "slice_array_dense",
                                      {"input", "starts"},
@@ -1855,6 +1855,71 @@ void SliceArrayDenseOp::VerifySig() {
   VLOG(4) << "End Verifying for: SliceArrayOp.";
 }
 
+void SliceArrayDenseOp::Build(pir::Builder &builder,             // NOLINT
+                              pir::OperationArgument &argument,  // NOLINT
+                              pir::Value input,
+                              pir::Value starts) {
+  VLOG(4) << "Start build SliceArrayDenseOp";
+  VLOG(4) << "Builder construction inputs";
+  argument.AddInputs({input, starts});
+  VLOG(4) << "Builder construction attributes";
+  VLOG(4) << "Builder construction outputs";
+  paddle::dialect::DenseTensorArrayType input_type =
+      input.type().dyn_cast<paddle::dialect::DenseTensorArrayType>();
+  paddle::dialect::IrTensor dense_input(
+      paddle::dialect::TransToPhiDataType(input_type.dtype()),
+      {},
+      input_type.data_layout(),
+      {});
+  paddle::dialect::IrMetaTensor meta_input(&dense_input);
+
+  phi::IntArray starts_list;
+  if (starts.dyn_cast<pir::OpResult>()
+          .owner()
+          ->isa<paddle::dialect::FullIntArrayOp>()) {
+    starts_list = std::move(phi::IntArray(paddle::dialect::GetInt64Vector(
+        starts.dyn_cast<pir::OpResult>()
+            .owner()
+            ->dyn_cast<paddle::dialect::FullIntArrayOp>()
+            .attribute("value"))));
+  } else if (starts.type().isa<pir::VectorType>()) {
+    size_t starts_size = starts.type().dyn_cast<pir::VectorType>().size();
+    starts_list =
+        std::move(phi::IntArray(std::vector<int64_t>(starts_size, -1)));
+    starts_list.SetFromTensor(true);
+  } else if (starts.type().isa<paddle::dialect::DenseTensorType>()) {
+    common::DDim starts_dim =
+        starts.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+    size_t starts_size = common::product(starts_dim);
+    if (common::contain_unknown_dim(starts_dim)) {
+      starts_size = 1;
+    }
+    starts_list =
+        std::move(phi::IntArray(std::vector<int64_t>(starts_size, -1)));
+    starts_list.SetFromTensor(true);
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "Only support VectorType or DenseTensorType"));
+  }
+
+  paddle::dialect::IrTensor dense_out;
+  paddle::dialect::IrMetaTensor meta_out(&dense_out);
+
+  phi::SliceArrayDenseInferMeta(
+      meta_input, starts_list, &meta_out, 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);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+}
+
 void SliceArrayDenseOp::InferMeta(phi::InferMetaContext *infer_meta) {
   auto fn = PD_INFER_META(phi::SliceArrayDenseInferMeta);
   fn(infer_meta);
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_op.h
@@ -334,6 +334,11 @@ class SliceArrayDenseOp
   static OpInfoTuple GetOpInfo();
   void VerifySig();
 
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value input,
+                    pir::Value starts);
+
   static phi::DataType GetKernelTypeForVar(
       const std::string &var_name,
       const phi::DataType &tensor_dtype,
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -274,6 +274,43 @@ static PyObject *static_api_array_to_tensor(PyObject *self,
   }
 }
 
+static PyObject *static_api_slice_array_dense(PyObject *self,
+                                              PyObject *args,
+                                              PyObject *kwargs) {
+  try {
+    VLOG(6) << "Add slice_array_dense 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, "slice_array_dense", 0);
+
+    PyObject *starts_obj = PyTuple_GET_ITEM(args, 1);
+    pir::Value starts;
+    if (PyObject_CheckIROpResult(starts_obj)) {
+      starts = CastPyArg2Value(starts_obj, "slice_array_dense", 1);
+    } else if (PyObject_CheckIRVectorOfOpResult(starts_obj)) {
+      std::vector<pir::Value> starts_tmp =
+          CastPyArg2VectorOfValue(starts_obj, "slice_array_dense", 1);
+      starts = paddle::dialect::stack(starts_tmp, /*axis*/ 0);
+
+    } else {
+      std::vector<int64_t> starts_tmp =
+          CastPyArg2Longs(starts_obj, "slice_array_dense", 1);
+      starts = paddle::dialect::full_int_array(
+          starts_tmp, phi::DataType::INT64, phi::CPUPlace());
+    }
+
+    // Call ir static api
+    auto static_api_out = paddle::dialect::slice_array_dense(input, starts);
+
+    return ToPyObject(static_api_out);
+  } catch (...) {
+    ThrowExceptionToPython(std::current_exception());
+    return nullptr;
+  }
+}
+
 static PyMethodDef ManualOpsAPI[] = {
     {"set_parameter",
      (PyCFunction)(void (*)(void))static_api_set_parameter,
@@ -303,6 +340,10 @@ static PyMethodDef ManualOpsAPI[] = {
      (PyCFunction)(void (*)(void))static_api_array_to_tensor,
      METH_VARARGS | METH_KEYWORDS,
      "C++ interface function for array_to_tensor."},
+    {"slice_array_dense",
+     (PyCFunction)(void (*)(void))static_api_slice_array_dense,
+     METH_VARARGS | METH_KEYWORDS,
+     "C++ interface function for slice_array_dense."},
     {nullptr, nullptr, 0, nullptr}};
 
 }  // namespace pybind
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3687,7 +3687,8 @@ void SliceArrayDenseInferMeta(const MetaTensor& input,
   if (config.is_runtime) {
     return;
   }
-  out->set_dims(input.dims());
+  // out->set_dims(input.dims());
+  out->set_dtype(input.dtype());
 }
 
 void SliceRawInferMeta(const MetaTensor& input,

Original file line number	Diff line number	Diff line change
`@@ -3687,7 +3687,8 @@ void SliceArrayDenseInferMeta(const MetaTensor& input,`
`3687`	`3687`	`if (config.is_runtime) {`
`3688`	`3688`	`return;`
`3689`	`3689`	`}`
`3690`		`- out->set_dims(input.dims());`
	`3690`	`+ // out->set_dims(input.dims());`
	`3691`	`+ out->set_dtype(input.dtype());`
`3691`	`3692`	`}`
`3692`	`3693`
`3693`	`3694`	`void SliceRawInferMeta(const MetaTensor& input,`