[CINN] add some op's infer_symbol_shape (#66400)

zyfncg · web-flow · commit 2354c7c998a0 · 2024-08-02T10:17:08.000+08:00
* add ome op's infer_symbol_shape

* fix bug

* fix bug of cross
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
@@ -238,6 +238,42 @@ bool Conv3dOpInferSymbolicShape(pir::Operation *op,
   return Conv2dOpInferSymbolicShape(op, infer_context);
 }
 
+bool CrossOpInferSymbolicShape(pir::Operation *op,
+                               pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &y_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  size_t x_dim = x_shape.shape().size();
+  size_t y_dim = y_shape.shape().size();
+
+  PADDLE_ENFORCE_EQ(x_dim,
+                    y_dim,
+                    common::errors::InvalidArgument(
+                        "The 'shape' of Input(X) should be equal to "
+                        "the 'shape' of Input(Y). But received "
+                        "Input(X).dimensions = [%d], "
+                        "Input(Y).dimensions = [%d]",
+                        x_dim,
+                        y_dim));
+
+  for (size_t i = 0; i < x_dim; i++) {
+    infer_context->AddEqualCstr(x_shape.shape()[i], y_shape.shape()[i]);
+  }
+
+  const int axis = op->attribute<pir::Int32Attribute>("axis").data();
+  if (axis != common::DDim::kMaxRank) {
+    const int dim = axis < 0 ? axis + x_dim : axis;
+    infer_context->AddEqualCstr(x_shape.shape()[dim], symbol::DimExpr{3});
+    infer_context->AddEqualCstr(y_shape.shape()[dim], symbol::DimExpr{3});
+  }
+
+  infer_context->SetShapeOrDataForValue(op->result(0), x_shape);
+
+  return true;
+}
+
 bool EmbeddingOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const std::vector<symbol::DimExpr> &x_dims =
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
@@ -24,6 +24,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cross)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Embedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseWeightEmbedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ExpandAs)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/element_wise_binary.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/element_wise_binary.cc
@@ -140,6 +140,8 @@ OP_ELEMENT_WISE_BINARY(BitwiseAnd_)
 OP_ELEMENT_WISE_BINARY(BitwiseXor)
 OP_ELEMENT_WISE_BINARY(BitwiseXor_)
 OP_ELEMENT_WISE_BINARY(Complex)
+OP_ELEMENT_WISE_BINARY(Copysign)
+OP_ELEMENT_WISE_BINARY(Copysign_)
 OP_ELEMENT_WISE_BINARY(Divide_)
 OP_ELEMENT_WISE_BINARY(ElementwisePow)
 OP_ELEMENT_WISE_BINARY(Equal)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/element_wise_binary.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/element_wise_binary.h
@@ -24,6 +24,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(BitwiseAnd_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BitwiseXor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BitwiseXor_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Complex)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Copysign)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Copysign_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Divide)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Divide_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ElementwisePow)
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
@@ -50,9 +50,11 @@ OP_SAME_OPERANDS_AND_RESULT(Cast)
 OP_SAME_OPERANDS_AND_RESULT(Cast_)
 OP_SAME_OPERANDS_AND_RESULT(Ceil)
 OP_SAME_OPERANDS_AND_RESULT(Ceil_)
+OP_SAME_OPERANDS_AND_RESULT(Celu)
 OP_SAME_OPERANDS_AND_RESULT(Clip)
 OP_SAME_OPERANDS_AND_RESULT(Clip_)
 OP_SAME_OPERANDS_AND_RESULT(Conj)
+OP_SAME_OPERANDS_AND_RESULT(CopyTo)
 OP_SAME_OPERANDS_AND_RESULT(Cos)
 OP_SAME_OPERANDS_AND_RESULT(Cos_)
 OP_SAME_OPERANDS_AND_RESULT(Cosh)
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
@@ -42,9 +42,11 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Ceil)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Ceil_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Celu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Clip)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Clip_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conj)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(CopyTo)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cos)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cos_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cosh)
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
@@ -366,6 +366,27 @@ bool BipartiteMatchOpInferSymbolicShape(
   return true;
 }
 
+bool CholeskyOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+
+  auto rank = x_shape.shape().size();
+  PADDLE_ENFORCE_GE(rank,
+                    2,
+                    common::errors::InvalidArgument(
+                        "The Input(X) should have at least 2 dimensions. But "
+                        "received a %d dimension tensor.",
+                        rank));
+
+  infer_context->AddEqualCstr(x_shape.shape()[rank - 2],
+                              x_shape.shape()[rank - 1]);
+
+  infer_context->SetShapeOrDataForValue(op->result(0), x_shape);
+
+  return true;
+}
+
 bool CummaxOpInferSymbolicShape(pir::Operation *op,
                                 pir::InferSymbolicShapeContext *infer_context) {
   pir::Value operand_source = op->operand_source(0);
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
@@ -26,6 +26,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Argmin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AsComplex)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AsReal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BipartiteMatch)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cholesky)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cummax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cummin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumprod)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -822,6 +822,7 @@
   kernel :
     func : celu
   backward : celu_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : channel_shuffle
   args : (Tensor x, int groups, str data_format="NCHW")
@@ -861,6 +862,7 @@
   kernel :
     func : cholesky
   backward : cholesky_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : cholesky_solve
   args : (Tensor x, Tensor y, bool upper=false)
@@ -1013,6 +1015,7 @@
   args : (Tensor x, Place place, bool blocking)
   output : Tensor(out)
   invoke : copy_to_impl(x, place, blocking)
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : copysign
   args : (Tensor x, Tensor y)
@@ -1023,6 +1026,7 @@
     func : copysign
   inplace: (x -> out)
   backward : copysign_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : correlation
   args : (Tensor input1, Tensor input2, int pad_size, int kernel_size, int max_displacement, int stride1, int stride2, int corr_type_multiply=1)
@@ -1086,6 +1090,7 @@
     func : cross
     data_type : x
   backward : cross_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 # Part of python API paddle.nn.functional.cross_entropy
 - op : cross_entropy_with_softmax
