【BUAA】【Infer Symbolic Shape No.39】Add FakeChannelWiseQuantizeDequantizeAbsMax Operator for CINN compiler (#66619)

Jeff114514 · web-flow · commit 147572938e3f · 2024-08-14T14:23:39.000+08:00
* summon

* add test

* addTest

* open

* Update test_fake_quantize_op.py
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
@@ -720,6 +720,42 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
   return true;
 }
 
+bool FakeChannelWiseQuantizeDequantizeAbsMaxOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+
+  int bit_length = op->attribute<pir::Int32Attribute>("bit_length").data();
+  int quant_axis = op->attribute<pir::Int32Attribute>("quant_axis").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));
+  PADDLE_ENFORCE_EQ(
+      quant_axis == 0 || quant_axis == 1,
+      true,
+      common::errors::InvalidArgument("'quant_axis' should be 0 or 1, but "
+                                      "the received is %d",
+                                      quant_axis));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(x_shape_or_data.shape())});
+
+  std::vector<symbol::DimExpr> out_scale_shape = {
+      x_shape_or_data.shape()[quant_axis]};
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_scale_shape)});
+
+  return true;
+}
+
 bool EigOpInferSymbolicShape(pir::Operation *op,
                              pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
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
@@ -60,6 +60,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigvals)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool3D)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeChannelWiseQuantizeAbsMax)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeChannelWiseQuantizeDequantizeAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2c)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2r)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftR2c)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1695,6 +1695,7 @@
   kernel :
     func : fake_channel_wise_quantize_dequantize_abs_max
     data_type : x
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
   backward : fake_channel_wise_quantize_dequantize_abs_max_grad
 
 - op : fake_dequantize_max_abs
diff --git a/test/legacy_test/test_fake_quantize_op.py b/test/legacy_test/test_fake_quantize_op.py
@@ -188,6 +188,79 @@ def test_fake_channel_wise_quantize_abs_max(self):
                 )
 
 
+class TestFakeChannelWiseQuantizeDequantizeAbsMaxOp(OpTest):
+    def setUp(self):
+        self.op_type = 'fake_channel_wise_quantize_dequantize_abs_max'
+        self.attrs = {'bit_length': 8}
+
+    def _fake_channel_wise_quantize_dequantize_abs_max(
+        self,
+        dtype,
+        input_shape,
+        quant_axis,
+        distribution,
+        round_type='TiesToEven',
+    ):
+        assert quant_axis in [0, 1], 'quant_axis should be 0 or 1.'
+        input_data = distribution(input_shape).astype(dtype)
+        compute_type = get_compute_type(dtype)
+        bnt = (1 << (self.attrs['bit_length'] - 1)) - 1
+        output_data = input_data.copy().astype(compute_type)
+        compute_axis = tuple(
+            i for i in range(len(input_shape)) if i != quant_axis
+        )
+        scale_broadcast = np.amax(input_data, axis=compute_axis, keepdims=True)
+        if round_type == 'TiesToEven':
+            round_out = np.round(bnt * output_data / scale_broadcast)
+            output_data = (
+                np.clip(round_out, -bnt - 1, bnt) * scale_broadcast / bnt
+            )
+            self.attrs['round_type'] = 0
+        else:
+            output_data = (
+                round_c(bnt * output_data / scale_broadcast)
+                * scale_broadcast
+                / bnt
+            )
+            self.attrs['round_type'] = 1
+        if quant_axis == 1:
+            scale_broadcast = np.transpose(scale_broadcast, (1,) + compute_axis)
+        scale = scale_broadcast.reshape(input_shape[quant_axis], -1)[:, 0]
+        self.python_api = fake_channel_wise_quantize_dequantize_abs_max_wrapper
+        self.inputs = {'X': input_data}
+        self.outputs = {'Out': output_data, 'OutScale': scale}
+        self.dtype = dtype
+        self.attrs['quant_axis'] = quant_axis
+        self.check_output(check_dygraph=False, check_pir=True)
+        gradient = [np.ones(input_data.shape) / np.prod(input_data.shape)]
+        self.check_grad(['X'], 'Out', user_defined_grads=gradient)
+
+    def test_channel_wise_fake_quant_dequant_abs_max(self):
+        input_shape_quant_axis_options = [
+            [(3, 4, 64, 64), 0],
+            [(15, 20, 5, 5), 1],
+            [(30, 15), 0],
+            [(30, 15), 1],
+        ]
+        round_type_options = ['TiesToEven', 'TiesAwayFromZero']
+        for input_shape_quant_axis, round_type in itertools.product(
+            input_shape_quant_axis_options, round_type_options
+        ):
+            input_shape, quant_axis = input_shape_quant_axis
+            with self.subTest(
+                input_shape=input_shape,
+                quant_axis=quant_axis,
+                round_type=round_type,
+            ):
+                self._fake_channel_wise_quantize_dequantize_abs_max(
+                    np.float32,
+                    input_shape,
+                    quant_axis,
+                    np.random.random,
+                    round_type=round_type,
+                )
+
+
 class TestFakeQuantizeRangeAbsMaxOp(OpTest):
     def setUp(self):
         self.op_type = 'fake_quantize_range_abs_max'
