[CINN] add assign_array && assign_array_ InferSymbolicShape (#69331)

ooooo-create · web-flow · commit d31f52c69121 · 2024-11-19T17:41:18.000+08:00
* add assign_array &amp;&amp; assign_array_

* Refine logic

* refine code

* refine code
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -3125,6 +3125,15 @@ OpInfoTuple AssignArray_Op::GetOpInfo() {
   return std::make_tuple(
       inputs, attributes, outputs, run_time_info, "assign_array");
 }
+bool AssignArrayOp::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(x())
+          .dyn_cast<symbol::RankedTensorArrayShapeOrDataDimExprs>();
+  infer_context->SetShapeOrDataForValue(
+      out(), symbol::ShapeOrDataDimExprs{x_shape_or_data});
+  return true;
+}
 
 void AssignArray_Op::VerifySig() {
   VLOG(4)
@@ -3223,6 +3232,16 @@ phi::DataType AssignArray_Op::GetKernelTypeForVar(
   return expected_kernel_dtype;
 }
 
+bool AssignArray_Op::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(x())
+          .dyn_cast<symbol::RankedTensorArrayShapeOrDataDimExprs>();
+  infer_context->SetShapeOrDataForValue(
+      out(), symbol::ShapeOrDataDimExprs{x_shape_or_data});
+  return true;
+}
+
 OpInfoTuple ExpandOp::GetOpInfo() {
   std::vector<paddle::dialect::OpInputInfo> inputs = {
       paddle::dialect::OpInputInfo(
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_op.h
@@ -495,7 +495,8 @@ class TEST_API AssignArrayOp
                      paddle::dialect::OpYamlInfoInterface,
                      paddle::dialect::InferMetaInterface,
                      paddle::dialect::GetKernelTypeForVarInterface,
-                     paddle::dialect::ForwardOnlyTrait> {
+                     paddle::dialect::ForwardOnlyTrait,
+                     paddle::dialect::InferSymbolicShapeInterface> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.assign_array"; }
@@ -520,14 +521,16 @@ class TEST_API AssignArrayOp
   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 AssignArray_Op
     : public pir::Op<AssignArray_Op,
                      paddle::dialect::OpYamlInfoInterface,
                      paddle::dialect::InferMetaInterface,
                      paddle::dialect::GetKernelTypeForVarInterface,
-                     paddle::dialect::ForwardOnlyTrait> {
+                     paddle::dialect::ForwardOnlyTrait,
+                     paddle::dialect::InferSymbolicShapeInterface> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.assign_array_"; }
@@ -549,6 +552,7 @@ class TEST_API AssignArray_Op
   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 ExpandOp
