From f90f34c27038bffb37b626a51bb68a40069c6796 Mon Sep 17 00:00:00 2001
From: Aleksei-grovety <113356454+Aleksei-grovety@users.noreply.github.com>
Date: Fri, 5 May 2023 16:50:54 +0400
Subject: [PATCH] [microNPU][ETHOSU] Fix Softmax quantization parameters

Fix zero point and scale values for operations according to the values in Vela, the test is updated to check case with different input and output zero point.
---
 .../contrib/ethosu/softmax_rewriter.py        | 28 +++++++++----------
 .../contrib/test_ethosu/test_codegen.py       |  2 +-
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/python/tvm/relay/backend/contrib/ethosu/softmax_rewriter.py b/python/tvm/relay/backend/contrib/ethosu/softmax_rewriter.py
index 16067fed951d..6c0a1dffc30c 100644
--- a/python/tvm/relay/backend/contrib/ethosu/softmax_rewriter.py
+++ b/python/tvm/relay/backend/contrib/ethosu/softmax_rewriter.py
@@ -101,7 +101,7 @@ def callback(
             ifm_scale=float(params.ifm.q_params.scale_f32),
             ifm_zero_point=int(params.ifm.q_params.zero_point),
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             pool_shape=(1, depth),
             ofm_channels=1,
             ofm_dtype=ifm_dtype,
@@ -141,7 +141,7 @@ def callback(
             ifm2_scale=0.0,
             ifm2_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ifm_channels=params.ifm.shape[-1],
             ifm2_channels=1,
             reversed_operands=False,
@@ -160,7 +160,7 @@ def callback(
             ifm_scale=0.0,
             ifm_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             pool_shape=(1, 1),
             ofm_channels=1,
             upscale="NONE",
@@ -175,7 +175,7 @@ def callback(
             ifm_scale=0.0,
             ifm_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ofm_channels=1,
         )
 
@@ -186,12 +186,12 @@ def callback(
             ifm2=headroom_plus_one,
             lut=lut,
             operator_type="SUB",
-            ifm_scale=1.0,
+            ifm_scale=0.0,
             ifm_zero_point=0,
             ifm2_scale=0.0,
             ifm2_zero_point=0,
-            ofm_scale=1.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_scale=0.0,
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ifm_channels=1,
             ifm2_channels=1,
             reversed_operands=False,
@@ -210,7 +210,7 @@ def callback(
             ifm2_scale=0.0,
             ifm2_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ifm_channels=1,
             ifm2_channels=1,
             reversed_operands=False,
@@ -228,7 +228,7 @@ def callback(
             ifm2_scale=0.0,
             ifm2_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ifm_channels=depth,
             ifm2_channels=1,
             reversed_operands=False,
@@ -250,7 +250,7 @@ def callback(
             ifm2_scale=0.0,
             ifm2_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ifm_channels=1,
             ifm2_channels=1,
             reversed_operands=False,
@@ -268,7 +268,7 @@ def callback(
             ifm2_scale=0.0,
             ifm2_zero_point=0,
             ofm_scale=0.0,
-            ofm_zero_point=int(params.ofm.q_params.zero_point),
+            ofm_zero_point=int(params.ifm.q_params.zero_point),
             ifm_channels=1,
             ifm2_channels=1,
             reversed_operands=False,
@@ -378,9 +378,9 @@ def callback(
                 ifm2=half_denominator_times_x,
                 lut=lut,
                 operator_type="SUB",
-                ifm_scale=0.0,
+                ifm_scale=2.0,
                 ifm_zero_point=0,
-                ifm2_scale=2.0,
+                ifm2_scale=0.0,
                 ifm2_zero_point=0,
                 ofm_scale=1.0,
                 ofm_zero_point=0,
@@ -422,7 +422,7 @@ def callback(
                 ifm2_scale=0.0,
                 ifm2_zero_point=0,
                 ofm_scale=0.0,
-                ofm_zero_point=int(params.ofm.q_params.zero_point),
+                ofm_zero_point=int(params.ifm.q_params.zero_point),
                 ifm_channels=1,
                 ifm2_channels=1,
                 reversed_operands=False,
diff --git a/tests/python/contrib/test_ethosu/test_codegen.py b/tests/python/contrib/test_ethosu/test_codegen.py
index 86f64d648320..e9a3e82a281c 100644
--- a/tests/python/contrib/test_ethosu/test_codegen.py
+++ b/tests/python/contrib/test_ethosu/test_codegen.py
@@ -330,7 +330,7 @@ def test_ethosu_softmax(
     def softmax(x):
         return tf.nn.softmax(x)
 
-    infra.compare_tvm_with_tflite(softmax, [ifm_shape], accel_type)
+    infra.compare_tvm_with_tflite(softmax, [ifm_shape], accel_type, ranges=[(-1, 1)])
 
 
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)