【BUAA】【Infer Symbolic Shape No.40、41、45】Add three operator for CINN compiler (#66827)

Jeff114514 · web-flow · commit 55c547e02331 · 2024-08-13T20:40:53.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
@@ -809,6 +809,36 @@ bool ConcatOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool FakeQuantizeRangeAbsMaxOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+
+  int window_size = op->attribute<pir::Int32Attribute>("window_size").data();
+  int bit_length = op->attribute<pir::Int32Attribute>("bit_length").data();
+
+  PADDLE_ENFORCE_EQ(bit_length >= 1 && bit_length <= 16,
+                    true,
+                    common::errors::InvalidArgument(
+                        "'bit_length' should be between 1 and 16, but "
+                        "the received is %d",
+                        bit_length));
+  const auto &out_scale = symbol::DimExpr(1);
+  const auto &out_scales = symbol::DimExpr(window_size);
+  infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data);
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({out_scale})});
+  if (op->result(2)) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs({out_scales})});
+  }
+  return true;
+}
+
 bool EditDistanceOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &hyps_shape_or_data =
@@ -899,6 +929,11 @@ bool FakeQuantizeMovingAverageAbsMaxOpInferSymbolicShape(
   return true;
 }
 
+bool FakeQuantizeRangeAbsMax_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return FakeQuantizeRangeAbsMaxOpInferSymbolicShape(op, infer_context);
+}
+
 bool FakeQuantizeMovingAverageAbsMax_OpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   return FakeQuantizeMovingAverageAbsMaxOpInferSymbolicShape(op, infer_context);
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
@@ -44,6 +44,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeMovingAverageAbsMax_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CoalesceTensor)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CoalesceTensor_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(EditDistance)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeRangeAbsMax)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeRangeAbsMax_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullWithTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttn)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnUnpadded)
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
@@ -78,6 +78,7 @@ 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(FakeDequantizeMaxAbs)
 OP_SAME_OPERANDS_AND_RESULT(Fill)
 OP_SAME_OPERANDS_AND_RESULT(Fill_)
 OP_SAME_OPERANDS_AND_RESULT(Fetch)
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
@@ -70,6 +70,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exp_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exponential_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeDequantizeMaxAbs)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fetch)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fill)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fill_)
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
@@ -807,6 +807,32 @@ bool FractionalMaxPool2dOpInferSymbolicShape(
   return FractionalMaxPoolOpInferSymbolicShape(op, infer_context);
 }
 
+bool FakeQuantizeAbsMaxOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &x_dims = x_shape_or_data.shape();
+
+  int bit_length = op->attribute<pir::Int32Attribute>("bit_length").data();
+
+  PADDLE_ENFORCE_EQ(bit_length >= 1 && bit_length <= 16,
+                    true,
+                    common::errors::InvalidArgument(
+                        "'bit_length' should be between 1 and 16, but "
+                        "the received is %d",
+                        bit_length));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_dims)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)})});
+
+  return true;
+}
+
 bool FakeChannelWiseQuantizeAbsMaxOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
@@ -821,6 +847,7 @@ bool FakeChannelWiseQuantizeAbsMaxOpInferSymbolicShape(
                         "'bit_length' should be between 1 and 16, but "
                         "the received is %d",
                         bit_length));
+
   PADDLE_ENFORCE_EQ(
       quant_axis == 0 || quant_axis == 1,
       true,
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
@@ -58,6 +58,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigvals)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool2D)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool3D)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeChannelWiseQuantizeAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2c)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2r)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1702,6 +1702,7 @@
   kernel :
     func : fake_dequantize_max_abs
     data_type : x
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fake_quantize_abs_max
   args : (Tensor x, int bit_length = 8, int round_type = 1)
@@ -1711,6 +1712,7 @@
   kernel :
     func : fake_quantize_abs_max
     data_type : x
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fake_quantize_dequantize_abs_max
   args : (Tensor x, int bit_length = 8, int round_type = 1)
@@ -1758,6 +1760,7 @@
     data_type : x
   optional : iter, out_scales
   inplace: (in_scale -> out_scale)
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fft_c2c
   args : (Tensor x, int64_t[] axes, str normalization, bool forward)
