[CINN] Add array_pop InferSymbolicShape (#69329)

ooooo-create · web-flow · commit 8cf5dd7b08f0 · 2024-11-19T17:30:35.000+08:00
* Add array_pos

* Fix typos
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -4699,6 +4699,34 @@ phi::DataType ArrayPopOp::GetKernelTypeForVar(
   return expected_kernel_dtype;
 }
 
+bool ArrayPopOp::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape_or_data =
+      infer_context->GetShapeOrDataForValue(input())
+          .dyn_cast<symbol::RankedTensorArrayShapeOrDataDimExprs>();
+  std::vector<symbol::DimExpr> input_shape = input_shape_or_data.GetShapeHint();
+  size_t input_rank = input_shape.size();
+  auto gen_shape_with_size = [&](size_t size) {
+    std::vector<symbol::DimExpr> shape;
+    for (size_t i = 0; i < size; ++i) {
+      shape.push_back(infer_context->GetNextSymName());
+    }
+    return shape;
+  };
+  std::vector<symbol::DimExpr> out_shape = gen_shape_with_size(input_rank);
+  std::vector<symbol::DimExpr> array_out_shape =
+      gen_shape_with_size(input_rank);
+  infer_context->SetShapeOrDataForValue(
+      array_out(),
+      symbol::ShapeOrDataDimExprs{
+          symbol::RankedTensorArrayShapeOrDataDimExprs(array_out_shape)});
+  infer_context->SetShapeOrDataForValue(
+      out(),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  return true;
+}
+
 }  // namespace paddle::dialect
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SplitGradOp)
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_op.h
@@ -813,7 +813,8 @@ class TEST_API ArrayPopOp
                      paddle::dialect::InferMetaInterface,
                      paddle::dialect::GetKernelTypeForVarInterface,
                      InplaceTrait,
-                     paddle::dialect::ForwardOnlyTrait> {
+                     paddle::dialect::ForwardOnlyTrait,
+                     paddle::dialect::InferSymbolicShapeInterface> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_pop"; }
@@ -840,6 +841,7 @@ class TEST_API ArrayPopOp
   static std::vector<pir::Type> InferMeta(
       const std::vector<pir::Value> &input_values,
       pir::AttributeMap *p_attributes);
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 };
 
 class ShareVarOp : public pir::Op<ShareVarOp, pir::SideEffectTrait> {
