[CINN] Add InferSymbolicShapeInterface for (eye) op (#68548)

ooooo-create · web-flow · commit 3db5a57665ad · 2024-10-28T17:37:48.000+08:00
* add infer symbolic shape for op (eye)

* update

* update

* update

* update

* refine logic

* apply review
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
@@ -246,6 +246,52 @@ bool EmptyOpInferSymbolicShape(pir::Operation *op,
   }
 }
 
+bool EyeOpInferSymbolicShape(pir::Operation *op,
+                             pir::InferSymbolicShapeContext *infer_context) {
+  std::vector<symbol::DimExpr> out_shape;
+  symbol::DimExpr num_columns_dim;
+  symbol::DimExpr num_rows_dim;
+  if (op->HasAttribute("num_rows")) {
+    int num_rows_int = op->attribute<pir::Int64Attribute>("num_rows").data();
+    num_rows_dim = symbol::DimExpr(num_rows_int);
+  } else if (op->operand_source(0)) {
+    const auto &num_rows_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(0));
+    if (num_rows_shape_or_data.data().has_value()) {
+      num_rows_dim =
+          symbol::DimExpr(num_rows_shape_or_data.data().value().at(0));
+    } else {
+      num_rows_dim = symbol::DimExpr(infer_context->GetNextSymName());
+    }
+  }
+
+  if (op->HasAttribute("num_columns")) {
+    int num_columns_int =
+        op->attribute<pir::Int64Attribute>("num_columns").data();
+    if (num_columns_int == -1) {
+      num_columns_dim = num_rows_dim;
+    } else {
+      num_columns_dim = symbol::DimExpr(num_columns_int);
+    }
+  } else if (op->operand_source(1)) {
+    const auto &num_columns_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(1));
+    if (num_columns_shape_or_data.data().has_value()) {
+      num_columns_dim =
+          symbol::DimExpr(num_columns_shape_or_data.data().value().at(0));
+    } else {
+      num_columns_dim = symbol::DimExpr(infer_context->GetNextSymName());
+    }
+  }
+  out_shape.emplace_back(num_rows_dim);
+  out_shape.emplace_back(num_columns_dim);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  return true;
+}
 bool FeedOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
   std::string name =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h
@@ -23,6 +23,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(AssignValue_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CudnnLstm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Data)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Empty)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eye)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Feed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Full)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Full_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1771,6 +1771,7 @@
     param : [num_rows, num_columns, dtype]
     data_type : dtype
     backend : place
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
   traits : paddle::dialect::ForwardOnlyTrait
 
 - op : fake_channel_wise_dequantize_max_abs
