[CINN]Fix DataOp and ConcatOp infer symbolic shape problem (#65231)

Aurelius84 · web-flow · commit 20179a0d6829 · 2024-06-18T09:42:00.000+08:00
* [CINN]Fix DataOp and ConcatOp infer symbolic shape problem

* fix comment
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc
@@ -43,9 +43,16 @@ bool ConcatOpInferSymbolicShape(pir::Operation *op,
   const auto input_values = op->operands_source();
   const auto input_size = input_values.size();
 
-  if (infer_context->GetShapeOrDataForValue(input_values[0])
-          .data()
-          .has_value()) {
+  const auto IsAllDataValue = [&]() -> bool {
+    for (const auto &value : input_values) {
+      if (!infer_context->GetShapeOrDataForValue(value).data().has_value()) {
+        return false;
+      }
+    }
+    return true;
+  };
+
+  if (IsAllDataValue()) {
     std::vector<symbol::DimExpr> out_data;
     for (const auto &value : input_values) {
       const auto &shape_or_data = infer_context->GetShapeOrDataForValue(value);
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc
@@ -117,9 +117,11 @@ bool DataOpInferSymbolicShape(pir::Operation *op,
     return sym_dims;
   }();
 
-  auto IsOneNumel = [&](pir::Value value) {
+  auto IsNumelLEKMaxRank = [](pir::Value value) {
     const auto &dims = value.type().dyn_cast<pir::DenseTensorType>().dims();
-    if (dims.size() == 1 && dims[0] == 1) {
+    if (dims.size() == 0) return true;
+    if (dims.size() != 1) return false;
+    if (dims[0] >= 1 && dims[0] <= ::common::DDim::kMaxRank) {
       return true;
     }
     return false;
@@ -130,12 +132,20 @@ bool DataOpInferSymbolicShape(pir::Operation *op,
     return dtype.isa<pir::Int32Type>() || dtype.isa<pir::Int64Type>();
   };
 
+  auto CreateSymForEachNumel = [&](pir::Value value) -> decltype(auto) {
+    const auto &dims = value.type().dyn_cast<pir::DenseTensorType>().dims();
+    const int64_t numel = ::common::product(dims);
+    std::vector<symbol::DimExpr> data;
+    for (int64_t i = 0; i < numel; ++i) {
+      data.push_back(symbol::DimExpr(infer_context->GetNextSymName()));
+    }
+    return data;
+  };
+
   const auto &shape_or_data = [&]() {
-    if (IsOneNumel(op->result(0)) && IsIntType(op->result(0))) {
-      std::vector<symbol::DimExpr> data{
-          symbol::DimExpr(infer_context->GetNextSymName())};
-      return symbol::ShapeOrDataDimExprs{
-          symbol::TensorShapeOrDataDimExprs(sym_dims, data)};
+    if (IsNumelLEKMaxRank(op->result(0)) && IsIntType(op->result(0))) {
+      return symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          sym_dims, CreateSymForEachNumel(op->result(0)))};
     } else {
       return symbol::ShapeOrDataDimExprs{
           symbol::TensorShapeOrDataDimExprs(sym_dims)};
