【Infer Symbolic Shape BUAA No13、66、91、122】Add 4 ops (#66877)

Guanhuachen2003 · web-flow · commit 7fb77f9e8bcd · 2024-08-08T18:13:38.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -634,6 +634,36 @@ bool BilinearInterpOpInferSymbolicShape(
   return BicubicInterpOpInferSymbolicShape(op, infer_context);
 }
 
+bool CheckFiniteAndUnscaleOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  // Retrieve the shape information of the input list
+  pir::Value operand_source = op->operand_source(0);
+  const symbol::TensorListShapeOrDataDimExprs &xs_shapes =
+      infer_context->GetShapeOrDataForValue(operand_source)
+          .dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
+
+  // Set the shapes for the output tensor list
+  symbol::TensorListShapeOrDataDimExprs outs_shapes;
+  outs_shapes.reserve(xs_shapes.size());
+  for (const auto &input_shape : xs_shapes) {
+    outs_shapes.emplace_back(input_shape.shape());
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::ShapeOrDataDimExprs{outs_shapes});
+
+  // Set the shape for the found_infinite output
+  symbol::TensorShapeOrDataDimExprs found_infinite_shape({symbol::DimExpr(1)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1), symbol::ShapeOrDataDimExprs{found_infinite_shape});
+
+  return true;
+}
+
+bool CheckFiniteAndUnscale_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return CheckFiniteAndUnscaleOpInferSymbolicShape(op, infer_context);
+}
+
 // bool CrfDecodingOpInferSymbolicShape(pir::Operation *op,
 //                                      pir::InferSymbolicShapeContext
 //                                      *infer_context) {
@@ -1289,6 +1319,31 @@ bool MeshgridOpInferSymbolicShape(
   return true;
 }
 
+bool MovingAverageAbsMaxScaleOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  // Get shapes of input tensors
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &in_state_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+  const auto &in_accum_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  // Set shapes for output tensors
+  infer_context->SetShapeOrDataForValue(op->result(0), x_shape);
+  infer_context->SetShapeOrDataForValue(
+      op->result(1), symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)}));
+  infer_context->SetShapeOrDataForValue(op->result(2), in_state_shape);
+  infer_context->SetShapeOrDataForValue(op->result(3), in_accum_shape);
+
+  return true;
+}
+
+bool MovingAverageAbsMaxScale_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return MovingAverageAbsMaxScaleOpInferSymbolicShape(op, infer_context);
+}
+
 // bool NceOpInferSymbolicShape(pir::Operation *op,
 //                              pir::InferSymbolicShapeContext *infer_context) {
 //   // pass
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -31,6 +31,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchNorm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BicubicInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bilinear)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BilinearInterp)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(CheckFiniteAndUnscale)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(CheckFiniteAndUnscale_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CrfDecoding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Concat)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CrossEntropyWithSoftmax)
@@ -66,6 +68,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNMS3)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MemoryEfficientAttention)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MovingAverageAbsMaxScale)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MovingAverageAbsMaxScale_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Nce)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PsroiPool)
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
@@ -1031,6 +1031,22 @@ bool KthvalueOpInferSymbolicShape(
   return true;
 }
 
+bool L1NormOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  // The output is a scalar, set the output shape accordingly
+  std::vector<symbol::DimExpr> output_shape;
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+  return true;
+}
+
+bool L1Norm_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return L1NormOpInferSymbolicShape(op, infer_context);
+}
+
 bool InverseOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &input_shape =
@@ -1526,6 +1542,52 @@ bool ShapeOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool RreluOpInferSymbolicShape(pir::Operation *op,
+                               pir::InferSymbolicShapeContext *infer_context) {
+  float lower = op->attribute<pir::FloatAttribute>("lower").data();
+  float upper = op->attribute<pir::FloatAttribute>("upper").data();
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &x_shape = x_shape_or_data.shape();
+
+  // Check constraints for the attributes lower and upper
+  PADDLE_ENFORCE_GE(lower,
+                    0,
+                    phi::errors::InvalidArgument(
+                        "The lower value should be greater than or equal to 0. "
+                        "But received lower value = %f.",
+                        lower));
+  PADDLE_ENFORCE_LE(upper,
+                    1,
+                    phi::errors::InvalidArgument(
+                        "The upper value should be less than or equal to 1. "
+                        "But received upper value = %f.",
+                        upper));
+  PADDLE_ENFORCE_GE(
+      upper,
+      lower,
+      phi::errors::InvalidArgument(
+          "The upper value should be greater than or equal to lower value. "
+          "But received upper value = %f, lower value = %f.",
+          upper,
+          lower));
+
+  // Set the shape for the output tensor out
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+
+  // Set the shape for the output tensor noise if it exists
+  if (op->num_results() > 1 && op->result(1) != nullptr) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(1),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_shape)});
+  }
+
+  return true;
+}
+
 bool ShapeSrOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   return ShapeOpInferSymbolicShape(op, infer_context);
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
@@ -68,6 +68,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Inverse)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IdentityLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IdentityLoss_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kthvalue)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(L1Norm)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(L1Norm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LpPool2d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)
@@ -93,6 +95,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Prod)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RepeatInterleave)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reshape)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reshape_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Rrelu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Shape)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShapeSr)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Slice)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -940,6 +940,7 @@
     param: [x, in_accum, in_state, moving_rate, is_test]
   optional : in_accum, in_state, out, out_state, out_accum
   inplace : (in_accum -> out_accum), (in_state -> out_state)
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op: nce
   args: (Tensor input, Tensor label, Tensor weight, Tensor bias, Tensor sample_weight, Tensor custom_dist_probs, Tensor custom_dist_alias, Tensor custom_dist_alias_probs, int num_total_classes, int[] custom_neg_classes={}, int num_neg_samples=10, int sampler=0, int seed=0, bool is_sparse=false, bool remote_prefetch=false, bool is_test=false)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -855,6 +855,7 @@
     param : [x, scale]
     data_type : x
   inplace : (x -> out)
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : check_numerics
   args : (Tensor tensor, str op_type = "", str var_name = "", int check_nan_inf_level = 0, int stack_height_limit = -1, str output_dir = "")
@@ -2645,6 +2646,7 @@
     data_type : x
   inplace: (x -> out)
   backward : l1_norm_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : label_smooth
   args : (Tensor label, Tensor prior_dist, float epsilon = 0.0f)
@@ -3963,6 +3965,7 @@
     data_type : x
   intermediate : noise
   backward : rrelu_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : rsqrt
   args : (Tensor x)
