add stack if (#61416)

chen2016013 · web-flow · commit 9aa32d2f9cd6 · 2024-02-01T19:38:33.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/op_infermeta_gen.py b/paddle/fluid/pir/dialect/op_generator/op_infermeta_gen.py
@@ -98,6 +98,9 @@ def get_infermeta_inputs_str(
     infermeta_inputs_str += "\n"
 
     infermeta_inputs_str += '  VLOG(4) << "Builder construction outputs";\n'
+    infermeta_inputs_str += (
+        '  bool is_from_tensor = false; (void) is_from_tensor;\n'
+    )
     # Prepare input type
     for idx in range(len(op_input_name_list)):
         if op_input_name_list[idx] not in inuse_infer_meta_args:
@@ -250,35 +253,9 @@ def GenBuildOutputsPart2(
 
 """
 
-    CREATE_INTARRAY_MUTABLE_ATTRIBUE_WITH_UNKONW_DATA_TEMPLATE = """  phi::IntArray {name};
-  if ({name}_.isa<pir::OpResult>() && {name}_.defining_op()->isa<paddle::dialect::FullIntArrayOp>()) {{
-    {name} = std::move(phi::IntArray(paddle::dialect::GetInt64Vector(
-                          {name}_.defining_op()
-                          ->dyn_cast<paddle::dialect::FullIntArrayOp>()
-                          .attribute("value"))));
-  }} else if ({name}_.type().isa<pir::VectorType>()) {{
-    size_t {name}_size = {name}_.type().dyn_cast<pir::VectorType>().size();
-    {name} = std::move(phi::IntArray(std::vector<int64_t>({name}_size, -1)));
-    {name}.SetFromTensor(true);
-  }} else if ({name}_.type().isa<paddle::dialect::DenseTensorType>()) {{
-    common::DDim {name}_dim = {name}_.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
-    size_t {name}_size = common::product({name}_dim);
-    if (common::contain_unknown_dim({name}_dim)) {{
-      {name}_size = 1;
-    }}
-    {name} = std::move(phi::IntArray(std::vector<int64_t>({name}_size, -1)));
-    {name}.SetFromTensor(true);
-  }} else if ({name}_.type().isa<paddle::dialect::AllocatedDenseTensorType>()) {{
-    common::DDim {name}_dim = {name}_.type().dyn_cast<paddle::dialect::AllocatedDenseTensorType>().dims();
-    size_t {name}_size = common::product({name}_dim);
-    if (common::contain_unknown_dim({name}_dim)) {{
-      {name}_size = 1;
-    }}
-    {name} = std::move(phi::IntArray(std::vector<int64_t>({name}_size, -1)));
-    {name}.SetFromTensor(true);
-  }} else {{
-    PADDLE_THROW(phi::errors::Unimplemented("Only support VectorType or DenseTensorType or AllocatedDenseTensorType"));
-  }}\n"""
+    CREATE_INTARRAY_MUTABLE_ATTRIBUE_WITH_UNKONW_DATA_TEMPLATE = """  is_from_tensor = false;
+  phi::IntArray {name} = std::move(phi::IntArray(paddle::dialect::ParseValueShape({name}_, &is_from_tensor)));
+  if (is_from_tensor) {name}.SetFromTensor(true);\n"""
 
     CREATE_VECTOR_INT_MUTABLE_ATTRIBUE_WITH_UNKONW_DATA_TEMPLATE = """  std::vector<int64_t> {name};
   if ({name}_.isa<pir::OpResult>() && {name}_.defining_op()->isa<paddle::dialect::FullIntArrayOp>()) {{
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -439,5 +439,56 @@ void CheckDataTypeOrValue(const phi::DataType& dtype,
   }
 }
 
+std::vector<int64_t> ParseValueShape(const pir::Value& shape_,
+                                     bool* is_from_tensor) {
+  std::vector<int64_t> vec_shape;
+  if (shape_.isa<pir::OpResult>() &&
+      shape_.defining_op()->isa<paddle::dialect::FullIntArrayOp>()) {
+    vec_shape = paddle::dialect::GetInt64Vector(
+        shape_.defining_op()
+            ->dyn_cast<paddle::dialect::FullIntArrayOp>()
+            .attribute("value"));
+  } else if (shape_.isa<pir::OpResult>() &&
+             shape_.defining_op()->isa<paddle::dialect::StackOp>()) {
+    std::vector<pir::Value> inputs = shape_.defining_op()
+                                         ->operand_source(0)
+                                         .defining_op()
+                                         ->operands_source();
+    for (auto item : inputs) {
+      auto tmp = ParseValueShape(item, is_from_tensor);
+      vec_shape.insert(vec_shape.end(), tmp.begin(), tmp.end());
+    }
+  } else if (shape_.type().isa<pir::VectorType>()) {
+    size_t shape_size = shape_.type().dyn_cast<pir::VectorType>().size();
+    vec_shape = std::vector<int64_t>(shape_size, -1);
+    *is_from_tensor = true;
+  } else if (shape_.type().isa<paddle::dialect::DenseTensorType>()) {
+    common::DDim shape_dim =
+        shape_.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+    size_t shape_size = common::product(shape_dim);
+    if (common::contain_unknown_dim(shape_dim)) {
+      shape_size = 1;
+    }
+    vec_shape = std::vector<int64_t>(shape_size, -1);
+    *is_from_tensor = true;
+  } else if (shape_.type().isa<paddle::dialect::AllocatedDenseTensorType>()) {
+    common::DDim shape_dim =
+        shape_.type()
+            .dyn_cast<paddle::dialect::AllocatedDenseTensorType>()
+            .dims();
+    size_t shape_size = common::product(shape_dim);
+    if (common::contain_unknown_dim(shape_dim)) {
+      shape_size = 1;
+    }
+    vec_shape = std::vector<int64_t>(shape_size, -1);
+    *is_from_tensor = true;
+  } else {
+    PADDLE_THROW(
+        phi::errors::Unimplemented("Only support VectorType or DenseTensorType "
+                                   "or AllocatedDenseTensorType"));
+  }
+  return vec_shape;
+}
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.h b/paddle/fluid/pir/dialect/operator/utils/utils.h
@@ -162,5 +162,8 @@ void CheckDataTypeOrValue(const phi::DataType& dtype,
 
 phi::DataType GetValueDataType(const pir::Value& value);
 
+std::vector<int64_t> ParseValueShape(const pir::Value& shape_,
+                                     bool* is_from_tensor);
+
 }  // namespace dialect
 }  // namespace paddle
