[CINN]Fix bug int the InferSymbolicShape interface of stack op (#65951)

gongshaotian · web-flow · commit 953aeea73c19 · 2024-07-12T14:36:23.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -722,40 +722,61 @@ bool StackOpInferSymbolicShape(pir::Operation *op,
 
   const auto &attributes = op->attributes();
   int axis = attributes.at("axis").dyn_cast<pir::Int32Attribute>().data();
-
   const symbol::TensorListShapeOrDataDimExprs &shape_data_list =
       infer_context->GetShapeOrDataForValue(operand_source)
           .dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
 
-  int rank = shape_data_list[0].shape().size();
+  size_t rank = shape_data_list.at(0).shape().size();
   if (axis < 0) axis += rank + 1;
-
   const symbol::ShapeOrDataDimExprs shape_data = [&] {
-    std::vector<symbol::DimExpr> shape_dim_exprs;
-    std::vector<symbol::DimExpr> data_dim_exprs;
-    for (const auto &shape_data : shape_data_list) {
-      if (shape_data.data().has_value() && axis == 0) {
-        data_dim_exprs.emplace_back(shape_data.data().value()[0]);
+    std::vector<symbol::DimExpr> result_shape = {};
+    std::vector<symbol::DimExpr> result_data = {};
+    const symbol::TensorShapeOrDataDimExprs &x_shape_data =
+        shape_data_list.at(0);
+
+    const bool data_flag = [&] {
+      for (const auto &shape_data : shape_data_list) {
+        if (!shape_data.data().has_value()) {
+          return false;
+        }
       }
-    }
+      return true;
+    }();
 
-    if (!data_dim_exprs.empty()) {
-      shape_dim_exprs.emplace_back(
+    if (data_flag) {
+      // case 1: data is not empty, eg: shape_data_list =
+      // [[shape:{3},data:{S0,6,7}],...]
+      if (axis == 0 && x_shape_data.data().value().size() <= 1) {
+        for (const auto &shape_data : shape_data_list) {
+          result_data.emplace_back(shape_data.data().value().at(0));
+        }
+      } else {
+        PADDLE_THROW(phi::errors::Unimplemented(
+            op->name() +
+            " 's InferSymbolicShape can NOT deal with data size > 1 now."));
+      }
+      result_shape.emplace_back(
           static_cast<std::int64_t>(shape_data_list.size()));
     } else {
-      for (int i = 0; i < rank; ++i) {
+      // case 2: data is empty, eg: shape_data_list =
+      // [[shape:{5,6,7},data:{}],...]
+      for (size_t i = 0; i < rank; ++i) {
         details::BuildCstrEqForTensorListAlongAxis(
             infer_context, shape_data_list, i);
       }
-      shape_dim_exprs.insert(shape_dim_exprs.begin() + axis,
-                             static_cast<std::int64_t>(shape_data_list.size()));
+      for (const symbol::DimExpr &dim : x_shape_data.shape()) {
+        result_shape.emplace_back(dim);
+      }
+      result_shape.insert(result_shape.begin() + axis,
+                          static_cast<std::int64_t>(shape_data_list.size()));
     }
-    if (data_dim_exprs.empty()) {
+
+    if (result_data.empty()) {
       return symbol::ShapeOrDataDimExprs(
-          symbol::TensorShapeOrDataDimExprs(shape_dim_exprs));
+          symbol::TensorShapeOrDataDimExprs(result_shape));
     }
     return symbol::ShapeOrDataDimExprs(
-        symbol::TensorShapeOrDataDimExprs(shape_dim_exprs, data_dim_exprs));
+        symbol::TensorShapeOrDataDimExprs(result_shape, result_data));
   }();
 
   pir::Value res = op->result(0);
