Fix quantization for input to reference model

Add the zerpoint instead of subtracting. This worked since the tests so
far used the ones as inputs which quantize to a zp of -128 which gives
the same np.int8 result in both cases since the int8 wraps.
Also needs to round and clip the scaled values to the int8 range.

Signed-off-by: Per Åstrand <[email protected]>

Change-Id: Ideaed6d072a4065573b38fb7476c7dbe8ba814fd

Author: per

backends/arm/test/arm_tosa_reference.py

@@ -139,9 +139,17 @@ def tosa_ref_dump_inputs(
         # Torch is doing Input[FP32]->Q[INT8]->DQ[FP32]->Operator[FP32]->Q[INT]->DQ[FP32]->[Output]FP32
         # Need to quantize the input to INT8 for TOSA comsumption
         if profile is TosaProfile.BI:
+            int8_max = np.iinfo(np.int8).max
+            int8_min = np.iinfo(np.int8).min
             data = (
-                (data / input_quantization_scales[name]) - input_quantization_zps[name]
-            ).astype(np.int8)
+                (
+                    (data / np.float32(input_quantization_scales[name]))
+                    + input_quantization_zps[name]
+                )
+                .round()
+                .clip(int8_min, int8_max)
+                .astype(np.int8)
+            )
 
         if save_on_disk:
             file_path = os.path.join(path, name + ".npy")

backends/arm/test/ops/test_add.py

@@ -9,7 +9,7 @@
 import shutil
 import unittest
 
-from typing import Optional, Tuple
+from typing import Tuple
 
 import torch
 from executorch.backends.arm.test.test_models import TosaProfile
@@ -30,6 +30,12 @@
 
 class TestSimpleAdd(unittest.TestCase):
     class Add(torch.nn.Module):
+        test_parameters = [
+            (torch.ones(5),),
+            (3 * torch.ones(8),),
+            (10 * torch.randn(8),),
+        ]
+
         def __init__(self):
             super().__init__()
             self.permute_memory_to_nhwc = False
@@ -38,6 +44,12 @@ def forward(self, x):
             return x + x
 
     class Add2(torch.nn.Module):
+        test_parameters = [
+            (torch.ones(1, 1, 4, 4), torch.ones(1, 1, 4, 4)),
+            (torch.randn(1, 1, 4, 4), torch.ones(1, 1, 4, 1)),
+            (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1)),
+        ]
+
         def __init__(self):
             super().__init__()
             self.permute_memory_to_nhwc = False
@@ -118,40 +130,40 @@ def _test_add_u55_BI_pipeline(
             .to_executorch()
         )
 
-    def test_add_tosa_MI(self):
-        test_data = (torch.randn(4, 4, 4),)
+    @parameterized.expand(Add.test_parameters)
+    def test_add_tosa_MI(self, test_data: torch.Tensor):
+        test_data = (test_data,)
         self._test_add_tosa_MI_pipeline(self.Add(), test_data)
 
-    @parameterized.expand(
-        [
-            (torch.ones(5),),  # test_data
-            (3 * torch.ones(8),),
-        ]
-    )
-    def test_add_tosa_BI(self, test_data: Optional[Tuple[torch.Tensor]]):
+    @parameterized.expand(Add.test_parameters)
+    def test_add_tosa_BI(self, test_data: torch.Tensor):
         test_data = (test_data,)
         self._test_add_tosa_BI_pipeline(self.Add(), test_data)
 
     @unittest.skipIf(
         not VELA_INSTALLED,
         "There is no point in running U55 tests if the Vela tool is not installed",
     )
-    def test_add_u55_BI(self):
-        test_data = (3 * torch.ones(5),)
+    @parameterized.expand(Add.test_parameters)
+    def test_add_u55_BI(self, test_data: torch.Tensor):
+        test_data = (test_data,)
         self._test_add_u55_BI_pipeline(self.Add(), test_data)
 
-    def test_add2_tosa_MI(self):
-        test_data = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1))
+    @parameterized.expand(Add2.test_parameters)
+    def test_add2_tosa_MI(self, operand1: torch.Tensor, operand2: torch.Tensor):
+        test_data = (operand1, operand2)
         self._test_add_tosa_MI_pipeline(self.Add2(), test_data)
 
-    def test_add2_tosa_BI(self):
-        test_data = (torch.ones(1, 1, 4, 4), torch.ones(1, 1, 4, 1))
+    @parameterized.expand(Add2.test_parameters)
+    def test_add2_tosa_BI(self, operand1: torch.Tensor, operand2: torch.Tensor):
+        test_data = (operand1, operand2)
         self._test_add_tosa_BI_pipeline(self.Add2(), test_data)
 
     @unittest.skipIf(
         not VELA_INSTALLED,
         "There is no point in running U55 tests if the Vela tool is not installed",
     )
-    def test_add2_u55_BI(self):
-        test_data = (torch.ones(1, 1, 4, 4), torch.ones(1, 1, 4, 1))
+    @parameterized.expand(Add2.test_parameters)
+    def test_add2_u55_BI(self, operand1: torch.Tensor, operand2: torch.Tensor):
+        test_data = (operand1, operand2)
         self._test_add_u55_BI_pipeline(self.Add2(), test_data)