[F2QI] Fix a rounding error on AvgPool when input and output affine scales differ (#12577)

cc @sfvaroglu @AndrewZhaoLuo

Author: Matthew Brookhart

python/tvm/relay/transform/fake_quantization_to_integer.py

@@ -114,33 +114,79 @@ def adaptive_avgpool1d(expr, type_map):
     """Rewrite an adaptive avgpool op"""
     arg = expr.args[0]
     t = type_map[arg]
-    arg = relay.op.cast(arg, "int32")
+    out_t = type_map[expr]
+    if not (
+        approx_equal(t.scale, out_t.scale)
+        and approx_equal(t.zero_point, out_t.zero_point)
+        and tvm.ir.structural_equal(t.dtype, out_t.dtype)
+    ):
+        arg = relay.qnn.op.requantize(
+            arg,
+            t.scale,
+            t.zero_point,
+            out_t.scale,
+            out_t.zero_point,
+            out_dtype="int32",
+            axis=t.axis,
+        )
+    else:
+        arg = relay.op.cast(arg, "int32")
     output_size = expr.attrs.output_size
     out = relay.op.nn.adaptive_avg_pool1d(arg, output_size)
-    out = relay.op.cast(out, t.dtype)
-    return [out, t]
+    return [out, TensorAffineType(out_t.scale, out_t.zero_point, "int32", out_t.axis)]
 
 
 @register_fake_quantization_to_integer("nn.avg_pool2d")
 def avgpool2d(expr, type_map):
     """Rewrite a avgpool op"""
     arg = expr.args[0]
     t = type_map[arg]
-    arg = relay.op.cast(arg, "int32")
+    out_t = type_map[expr]
+    if not (
+        approx_equal(t.scale, out_t.scale)
+        and approx_equal(t.zero_point, out_t.zero_point)
+        and tvm.ir.structural_equal(t.dtype, out_t.dtype)
+    ):
+        arg = relay.qnn.op.requantize(
+            arg,
+            t.scale,
+            t.zero_point,
+            out_t.scale,
+            out_t.zero_point,
+            out_dtype="int32",
+            axis=t.axis,
+        )
+    else:
+        arg = relay.op.cast(arg, "int32")
     out = relay.op.nn.avg_pool2d(arg, **expr.attrs)
-    out = relay.op.cast(out, t.dtype)
-    return [out, t]
+    return [out, TensorAffineType(out_t.scale, out_t.zero_point, "int32", out_t.axis)]
 
 
 @register_fake_quantization_to_integer("nn.global_avg_pool2d")
 def global_avgpool2d(expr, type_map):
     """Rewrite a global_avgpool op"""
     arg = expr.args[0]
     t = type_map[arg]
-    arg = relay.op.cast(arg, "int32")
+    out_t = type_map[expr]
+    out_t = type_map[expr]
+    if not (
+        approx_equal(t.scale, out_t.scale)
+        and approx_equal(t.zero_point, out_t.zero_point)
+        and tvm.ir.structural_equal(t.dtype, out_t.dtype)
+    ):
+        arg = relay.qnn.op.requantize(
+            arg,
+            t.scale,
+            t.zero_point,
+            out_t.scale,
+            out_t.zero_point,
+            out_dtype="int32",
+            axis=t.axis,
+        )
+    else:
+        arg = relay.op.cast(arg, "int32")
     out = relay.op.nn.global_avg_pool2d(arg)
-    out = relay.op.cast(out, t.dtype)
-    return [out, t]
+    return [out, TensorAffineType(out_t.scale, out_t.zero_point, "int32", out_t.axis)]
 
 
 @register_fake_quantization_to_integer("broadcast_to")

tests/python/relay/test_pass_fake_quantization_to_integer.py

@@ -281,10 +281,9 @@ def test_fake_quantize_maxpool():
 def test_fake_quantize_adaptive_avgpool1d(output_size):
     x = relay.var("x", shape=[1, 128, 768], dtype="int8")
 
-    zero = relay.const(0)
-    x = relay.qnn.op.dequantize(x, relay.const(2.0), zero)
+    x = relay.qnn.op.dequantize(x, relay.const(2.0), relay.const(-12))
     op = relay.op.nn.adaptive_avg_pool1d(x, output_size)
-    op = relay.qnn.op.quantize(op, relay.const(2.0), zero)
+    op = relay.qnn.op.quantize(op, relay.const(0.5), relay.const(10))
 
     x_np = np.random.randint(-128, 127, size=[1, 128, 768], dtype="int8")
 
@@ -294,10 +293,9 @@ def test_fake_quantize_adaptive_avgpool1d(output_size):
 def test_fake_quantize_avgpool():
     x = relay.var("x", shape=[1, 3, 224, 224], dtype="int8")
 
-    zero = relay.const(0)
-    x = relay.qnn.op.dequantize(x, relay.const(2.0), zero)
+    x = relay.qnn.op.dequantize(x, relay.const(2.0), relay.const(-12))
     op = relay.op.nn.avg_pool2d(x, [3, 3])
-    op = relay.qnn.op.quantize(op, relay.const(2.0), zero)
+    op = relay.qnn.op.quantize(op, relay.const(0.5), relay.const(10))
 
     x_np = np.random.randint(-128, 127, size=[1, 3, 224, 224], dtype="int8")
 
@@ -307,10 +305,9 @@ def test_fake_quantize_avgpool():
 def test_fake_quantize_global_avg_pool():
     x = relay.var("x", shape=[1, 3, 224, 224], dtype="int8")
 
-    zero = relay.const(0)
-    x = relay.qnn.op.dequantize(x, relay.const(2.0), zero)
+    x = relay.qnn.op.dequantize(x, relay.const(2.0), relay.const(-12))
     op = relay.op.nn.global_avg_pool2d(x)
-    op = relay.qnn.op.quantize(op, relay.const(2.0), zero)
+    op = relay.qnn.op.quantize(op, relay.const(0.5), relay.const(10))
 
     x_np = np.random.randint(-128, 127, size=[1, 3, 224, 224], dtype="int8")