[CINN] Add symbol_infer_interface for create_array_like && has_elements (#69055)

ooooo-create · web-flow · commit 3dccf91eb555 · 2024-11-19T17:41:05.000+08:00
* Add symbol_infer_interface for create_array_like &amp;&amp; has_elements

* fix codestyle

* refine code
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -950,6 +950,14 @@ void HasElementsOp::VerifySig() {
                         "The type of cf.has_elements' output is not correct."));
 }
 
+bool HasElementsOp::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  infer_context->SetShapeOrDataForValue(
+      out(),
+      symbol::ShapeOrDataDimExprs(
+          symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)})));
+  return true;
+}
 const char *AssertOp::attributes_name[1] = {"summarize"};    // NOLINT
 const char AssertOp::ERROR_INFO_ATTR_NAME[] = "error_info";  // NOLINT
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -120,7 +120,8 @@ struct TuplePushOpVjpInterfaceModel : public VjpInterface::Concept {
 ///   ...
 /// }
 ///
-class HasElementsOp : public pir::Op<HasElementsOp> {
+class HasElementsOp
+    : public pir::Op<HasElementsOp, InferSymbolicShapeInterface> {
  public:
   using Op::Op;
   static const char *name() { return "cf.has_elements"; }
@@ -133,6 +134,7 @@ class HasElementsOp : public pir::Op<HasElementsOp> {
   void VerifySig();
   pir::Value input() { return operand_source(0); }
   pir::Value out() { return result(0); }
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 };
 
 ///
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -1638,6 +1638,15 @@ std::vector<pir::Type> CreateArrayLikeOp::InferMeta(
   return argument_outputs;
 }
 
+bool CreateArrayLikeOp::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape_or_data =
+      infer_context->GetShapeOrDataForValue(input())
+          .dyn_cast<symbol::RankedTensorArrayShapeOrDataDimExprs>();
+  infer_context->SetShapeOrDataForValue(
+      out(), symbol::ShapeOrDataDimExprs{input_shape_or_data});
+  return true;
+}
 OpInfoTuple ArrayLengthOp::GetOpInfo() {
   std::vector<paddle::dialect::OpInputInfo> inputs = {
       OpInputInfo("x",
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_op.h
@@ -202,6 +202,7 @@ class TEST_API CreateArrayLikeOp
     : public pir::Op<CreateArrayLikeOp,
                      OpYamlInfoInterface,
                      InferMetaInterface,
+                     InferSymbolicShapeInterface,
                      paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
@@ -220,6 +221,7 @@ class TEST_API CreateArrayLikeOp
   static std::vector<pir::Type> InferMeta(
       const std::vector<pir::Value> &input_values,
       pir::AttributeMap *p_attributes);
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 };
 
 class TEST_API ArrayLengthOp
