[CINN]Add more infer shape funcs (#65889)

Hongqing-work · web-flow · commit 913b257c55c9 · 2024-07-25T19:45:48.000+08:00
* [CINN]Add more infer_symbol_shape

* add interface in yaml

* fix
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -264,6 +264,24 @@ bool ExpandAsOpInferSymbolicShape(
   return true;
 }
 
+bool FillDiagonalTensorOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_dims = x_shape_or_data.shape();
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_dims)});
+
+  return true;
+}
+
+bool FillDiagonalTensor_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return FillDiagonalTensorOpInferSymbolicShape(op, infer_context);
+}
+
 bool GatherOpInferSymbolicShape(pir::Operation *op,
                                 pir::InferSymbolicShapeContext *infer_context) {
   const auto &input_shape_or_data =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -26,6 +26,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Embedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseWeightEmbedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ExpandAs)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonalTensor)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonalTensor_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gather)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherNd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -71,6 +71,8 @@ OP_SAME_OPERANDS_AND_RESULT(Exp_)
 OP_SAME_OPERANDS_AND_RESULT(Expm1)
 OP_SAME_OPERANDS_AND_RESULT(Expm1_)
 OP_SAME_OPERANDS_AND_RESULT(Exponential_)
+OP_SAME_OPERANDS_AND_RESULT(Fill)
+OP_SAME_OPERANDS_AND_RESULT(Fill_)
 OP_SAME_OPERANDS_AND_RESULT(Fetch)
 OP_SAME_OPERANDS_AND_RESULT(Flip)
 OP_SAME_OPERANDS_AND_RESULT(Floor)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -64,6 +64,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exponential_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fetch)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fill)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fill_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flip)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Floor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Floor_)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -339,6 +339,91 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
   return true;
 }
 
+bool FftC2cOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_dims = x_shape_or_data.shape();
+
+  // Set the output shape to be the same as the input shape
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_dims)});
+
+  return true;
+}
+
+bool FftC2rOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_dims = x_shape_or_data.shape();
+
+  auto axes = paddle::dialect::details::GetVectorAttr<int64_t>(op, "axes");
+  int64_t last_dim_size =
+      op->attribute<pir::Int64Attribute>("last_dim_size").data();
+  int last_fft_axis = static_cast<int>(axes.back());
+
+  std::vector<symbol::DimExpr> out_dims = x_dims;
+
+  if (last_dim_size > 0) {
+    out_dims[last_fft_axis] = symbol::DimExpr(last_dim_size);
+  } else {
+    symbol::DimExprBuilder builder;
+    out_dims[last_fft_axis] =
+        builder.Mul(x_dims[last_fft_axis], 2) - symbol::DimExpr{1};
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+
+  return true;
+}
+
+bool FftR2cOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_dims = x_shape_or_data.shape();
+
+  auto axes = paddle::dialect::details::GetVectorAttr<int64_t>(op, "axes");
+  bool onesided = op->attribute<pir::BoolAttribute>("onesided").data();
+
+  std::vector<symbol::DimExpr> out_dims = x_dims;
+
+  if (onesided) {
+    int last_fft_axis = static_cast<int>(axes.back());
+    symbol::DimExprBuilder builder;
+    out_dims[last_fft_axis] =
+        builder.Add(builder.Div(x_dims[last_fft_axis], 2), 1);
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+
+  return true;
+}
+
+bool FillDiagonalOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_dims = x_shape_or_data.shape();
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_dims)});
+
+  return true;
+}
+
+bool FillDiagonal_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return FillDiagonalOpInferSymbolicShape(op, infer_context);
+}
+
 bool FlattenOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &attributes = op->attributes();
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -34,6 +34,11 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DistributeFpnProposals)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2c)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2r)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftR2c)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonal)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonal_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kthvalue)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1726,6 +1726,7 @@
   kernel :
     func : fft_c2c
   backward : fft_c2c_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fft_c2r
   args : (Tensor x, int64_t[] axes, str normalization, bool forward, int64_t last_dim_size=0L)
@@ -1735,6 +1736,7 @@
   kernel :
     func : fft_c2r
   backward : fft_c2r_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fft_r2c
   args : (Tensor x, int64_t[] axes, str normalization, bool forward, bool onesided)
@@ -1744,6 +1746,7 @@
   kernel :
     func : fft_r2c
   backward : fft_r2c_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fill
   args : (Tensor x, Scalar(double) value=0)
@@ -1755,6 +1758,7 @@
     func : fill
   inplace : (x -> out)
   backward: fill_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fill_diagonal
   args : (Tensor x, float value=0, int offset=0, bool wrap=false)
@@ -1766,6 +1770,7 @@
     data_type : x
   inplace : (x -> out)
   backward : fill_diagonal_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fill_diagonal_tensor
   args : (Tensor x, Tensor y, int64_t offset = 0, int dim1 = 0, int dim2 = 1)
@@ -1776,6 +1781,7 @@
     func : fill_diagonal_tensor
   inplace : (x -> out)
   backward : fill_diagonal_tensor_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : flash_attn
   args : (Tensor q, Tensor k, Tensor v, Tensor fixed_seed_offset, Tensor attn_mask, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
