[CINN]Add default new data method for infer symbolic (#66849)

Hongqing-work · web-flow · commit 9eaf052e49ea · 2024-08-01T12:53:09.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc
@@ -111,6 +111,71 @@ symbol::ShapeOrDataDimExprs CreateShapeOrDataForXShape(
   return symbol::TensorShapeOrDataDimExprs(InsertZeros(x_dims));
 }
 
+ExprVec GetOrCreateExprVecFromData(
+    const ShapeOrData &shapeordata,
+    pir::InferSymbolicShapeContext *infer_context) {
+  if (!HasCompleteData(shapeordata)) {
+    LOG(WARNING) << "ShapeOrDataDimExprs with incomplete data info. Need to "
+                    "create new DimExpr symbol.";
+  }
+
+  const auto &GetTensorItemNumFromShape =
+      [](const std::vector<symbol::DimExpr> &shape) -> int64_t {
+    const auto &optional_int64_shape = VecExpr2Int64(shape);
+    PADDLE_ENFORCE_EQ(
+        optional_int64_shape.has_value(),
+        true,
+        phi::errors::InvalidArgument(
+            "The shape of tensor should be known when GetExprVecFromData."));
+    return std::accumulate(optional_int64_shape->begin(),
+                           optional_int64_shape->end(),
+                           1,
+                           std::multiplies<int64_t>());
+  };
+
+  const auto &GetNewDataDimExpr = [&]() -> symbol::DimExpr {
+    return symbol::DimExpr{infer_context->GetNextSymName()};
+  };
+
+  ExprVec result;
+  const auto &GetTensorData =
+      [&](const symbol::TensorShapeOrDataDimExprs &tensor_shape_or_data) {
+        if (tensor_shape_or_data.data().has_value()) {
+          for (const auto &expr : tensor_shape_or_data.data().value()) {
+            result.emplace_back(expr);
+          }
+        } else {
+          for (int i = 0;
+               i < GetTensorItemNumFromShape(tensor_shape_or_data.shape());
+               ++i) {
+            result.emplace_back(GetNewDataDimExpr());
+          }
+        }
+      };
+
+  shapeordata.Match(
+      [&](const symbol::TensorShapeOrDataDimExprs &impl) {
+        GetTensorData(impl);
+      },
+      [&](const symbol::TensorListShapeOrDataDimExprs &impl) {
+        for (const auto &tensor_shape_or_data : impl) {
+          GetTensorData(tensor_shape_or_data);
+        }
+      },
+      [&](const symbol::RankedTensorArrayShapeOrDataDimExprs &impl) {
+        PADDLE_THROW(phi::errors::Fatal(
+            "Dead code, RankedTensorArrayShapeOrDataDimExprs can not get "
+            "data"));
+        return;
+      },
+      [&](const symbol::NullShapeOrDataDimExpr &impl) {
+        PADDLE_THROW(phi::errors::Fatal(
+            "Dead code, NullShapeOrDataDimExpr can not get data"));
+        return;
+      });
+  return result;
+}
+
 void BuildCstrEqForTensorListAlongAxis(
     pir::InferSymbolicShapeContext *infer_context,
     const symbol::TensorListShapeOrDataDimExprs &shape_data_list,
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h
@@ -125,6 +125,10 @@ inline ExprVec GetExprVecFromData(const ShapeOrData &shapeordata) {
   return result;
 }
 
+ExprVec GetOrCreateExprVecFromData(
+    const ShapeOrData &shapeordata,
+    pir::InferSymbolicShapeContext *infer_context);
+
 inline ExprVec GetExprVecFromShape(const ShapeOrData &shapeordata) {
   const auto GetShapeExprsFromList = [&]() {
     ExprVec result;
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
@@ -364,25 +364,15 @@ bool BicubicInterpOpInferSymbolicShape(
         common::errors::InvalidArgument(
             "The size_tensor of Interpolation should be type of "
             "TensorListShapeOrDataDimExprs"));
-    const auto &size_tensor_list_shape =
-        size_tensor_shape.dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
-    std::vector<symbol::DimExpr> result;
-    result.reserve(size_tensor_list_shape.size());
-    for (const auto &tensor_shape : size_tensor_list_shape) {
-      result.push_back(tensor_shape.data()->front());
-    }
-    return result;
+    return details::GetOrCreateExprVecFromData(size_tensor_shape,
+                                               infer_context);
   };
   auto GetOutSizeDataExpr =
       [&](pir::Value value) -> std::vector<symbol::DimExpr> {
     const symbol::ShapeOrDataDimExprs &out_size_tensor_shape =
         infer_context->GetShapeOrDataForValue(value);
-    PADDLE_ENFORCE_EQ(
-        out_size_tensor_shape.data().has_value(),
-        true,
-        common::errors::InvalidArgument(
-            "The output size of Interpolation should have data value."));
-    return out_size_tensor_shape.data().value();
+    return details::GetOrCreateExprVecFromData(out_size_tensor_shape,
+                                               infer_context);
   };
   auto GetOutDimByScale = [&](const symbol::DimExpr &in_dim,
                               float scale) -> symbol::DimExpr {
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -28,6 +28,7 @@ paddle::dialect::AddN_Op, paddle::dialect::AddNArrayOp,
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/ir/ir_meta_tensor.h"
 #include "paddle/fluid/pir/dialect/operator/ir/ir_selected_rows.h"
@@ -3317,9 +3318,13 @@ bool ExpandOp::InferSymbolicShape(
       infer_context->GetShapeOrDataForValue(shape());
 
   const std::vector<symbol::DimExpr> &x_dims = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> expand_shape =
+      details::GetOrCreateExprVecFromData(expand_shape_shape_or_data,
+                                          infer_context);
 
   const auto &DealWithMinusOneAndSetOutput =
-      [&](std::vector<symbol::DimExpr> &expand_shape) {
+      [&](const std::vector<symbol::DimExpr> &expand_shape) {
+        std::vector<symbol::DimExpr> result_shape(expand_shape);
         for (size_t i = 0; i < expand_shape.size(); i++) {
           if (expand_shape[i] == symbol::DimExpr{-1}) {  // copy the dim from x
             // the shape is right aligned
@@ -3332,62 +3337,17 @@ bool ExpandOp::InferSymbolicShape(
                                              "but got %d",
                                              index));
 
-            expand_shape[i] = x_dims[index];
+            result_shape[i] = x_dims[index];
           }
         }
 
         infer_context->SetShapeOrDataForValue(
             out(),
             symbol::ShapeOrDataDimExprs{
-                symbol::TensorShapeOrDataDimExprs(expand_shape)});
+                symbol::TensorShapeOrDataDimExprs(result_shape)});
       };
 
-  const auto &InferWithTensorShapeOrDataDimExprs =
-      [&](const symbol::TensorShapeOrDataDimExprs &shape_or_data) {
-        if (shape_or_data.data()) {
-          std::vector<symbol::DimExpr> expand_shape =
-              shape_or_data.data().value();
-          DealWithMinusOneAndSetOutput(expand_shape);
-        } else {
-          infer_context->SetSymbolForValueByStaticShape(out());
-        }
-      };
-
-  const auto &InferWithTensorListShapeOrDataDimExprs =
-      [&](const symbol::TensorListShapeOrDataDimExprs &shape_or_data_list) {
-        if (shape_or_data_list.size() == 1) {
-          InferWithTensorShapeOrDataDimExprs(shape_or_data_list.at(0));
-        } else {
-          std::vector<symbol::DimExpr> expand_shape;
-          for (const auto &shape_data : shape_or_data_list) {
-            if (shape_data.data()) {
-              expand_shape.emplace_back(shape_data.data()->at(0));
-            } else {
-              expand_shape.emplace_back(
-                  symbol::DimExpr{infer_context->GetNextSymName()});
-            }
-          }
-          DealWithMinusOneAndSetOutput(expand_shape);
-        }
-      };
-
-  expand_shape_shape_or_data.Match(
-      [&](const symbol::TensorShapeOrDataDimExprs &impl) {
-        InferWithTensorShapeOrDataDimExprs(impl);
-      },
-      [&](const symbol::TensorListShapeOrDataDimExprs &impl) {
-        InferWithTensorListShapeOrDataDimExprs(impl);
-      },
-      [&](const symbol::RankedTensorArrayShapeOrDataDimExprs &impl) {
-        PADDLE_THROW(
-            phi::errors::Fatal("Dead code, TensorArray should not be "
-                               "shape value for expand."));
-      },
-      [&](const symbol::NullShapeOrDataDimExpr &impl) {
-        PADDLE_THROW(
-            phi::errors::Fatal("Dead code, null value should not be "
-                               "shape value for expand."));
-      });
+  DealWithMinusOneAndSetOutput(expand_shape);
   return true;
 }
 
