Respond to review comments

Author: ekalda

python/tvm/relay/backend/contrib/ethosu/legalize.py

@@ -16,7 +16,7 @@
 # under the License.
 # pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter
 """A set of passes to legalize some of operations for the NPU"""
-from typing import List, Type
+from typing import List, Type, Callable
 import math
 
 import numpy as np  # type: ignore
@@ -125,7 +125,9 @@ def __call__(self, *args, **kwargs):
         pass
 
 
-def get_lut_from_func(ifm_scale, ifm_zp, ofm_scale, ofm_zp, func):
+def get_lut_from_func(
+    ifm_scale: float, ifm_zp: int, ofm_scale: float, ofm_zp: int, func: Callable[[float], float]
+) -> List[int]:
     """Method to calculate the values of the lookup table based on the calculation function"""
     lut_values = list()
     # Only int8 is currently supported
@@ -144,13 +146,15 @@ def get_lut_from_func(ifm_scale, ifm_zp, ofm_scale, ofm_zp, func):
 class LutActivationRewriter(DFPatternCallback):
     """A class to create an identity operator with the LUT"""
 
-    def __init__(self, params_class, activation_type, calc_func):
+    def __init__(
+        self, params_class: Type, activation_type: str, calc_func: Callable[[float], float]
+    ):
         super().__init__(require_type=True, rewrite_once=True)
         self.pattern = (wildcard().has_attr({"Composite": params_class.composite_name}))(wildcard())
         self.activation_type = activation_type
         self.calc_func = calc_func
 
-    def callback(self, pre, post, node_map):
+    def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map):
         id_input = post.args[0]
 
         quantize_args = post.op.body.args
@@ -205,7 +209,7 @@ def __call__(self, *args, **kwargs):
         pass
 
 
-def sigmoid_calc_func(x):
+def sigmoid_calc_func(x: float) -> float:
     """Function to calculate the values for sigmoid"""
     # Thse limits are inherited from TFLite
     upper_limit = 8.0

python/tvm/relay/op/contrib/ethosu.py

@@ -918,27 +918,30 @@ def abs_pattern() -> tvm.relay.dataflow_pattern.DFPattern:
     return pattern
 
 
-class TanhParams:
+class LutActivationParams:
     """
-    This class will parse a call to a ethos-u.tanh composite function
-    and extract the parameter information.
+    A parent class for LUT based activation functions that extract the input and
+    output tensors and check whether they are valid.
     """
 
-    composite_name = "ethos-u.tanh"
-
     def __init__(self, func_body: Call):
         self.ofm = TensorParams(func_body)
         self.ifm = TensorParams(func_body.args[0].args[0].args[0])
 
     def is_valid(self):
         """
-        This function checks whether reshape has compatible attributes with the NPU
+        This function checks whether activation has compatible attributes with the NPU
         """
         if not check_valid_dtypes([self.ifm, self.ofm], supported_dtypes=[np.int8]):
             return False
         return True
 
 
+class TanhParams(LutActivationParams):
+
+    composite_name = "ethos-u.tanh"
+
+
 def tanh_pattern():
     """Create pattern for tanh"""
     dequant = is_op("qnn.dequantize")(wildcard(), is_constant(), is_constant())
@@ -947,6 +950,23 @@ def tanh_pattern():
     return quant
 
 
+class SigmoidParams(LutActivationParams):
+    """
+    This class will parse a call to a ethos-u.sigmoid composite function
+    and extract the parameter information.
+    """
+
+    composite_name = "ethos-u.sigmoid"
+
+
+def sigmoid_pattern():
+    """Create pattern for sigmoid"""
+    dequant = is_op("qnn.dequantize")(wildcard(), is_constant(), is_constant())
+    sigmoid = is_op("sigmoid")(dequant)
+    quant = is_op("qnn.quantize")(sigmoid, is_constant(), is_constant())
+    return quant
+
+
 class MeanParams:
     """
     This class will parse a call to ethosu.mean composite function
@@ -1087,35 +1107,6 @@ def concat_pattern():
     return concat
 
 
-class SigmoidParams:
-    """
-    This class will parse a call to a ethos-u.sigmoid composite function
-    and extract the parameter information.
-    """
-
-    composite_name = "ethos-u.sigmoid"
-
-    def __init__(self, func_body: Call):
-        self.ofm = TensorParams(func_body)
-        self.ifm = TensorParams(func_body.args[0].args[0].args[0])
-
-    def is_valid(self):
-        """
-        This function checks whether sigmoid has compatible attributes with the NPU
-        """
-        if not check_valid_dtypes([self.ifm, self.ofm], supported_dtypes=[np.int8]):
-            return False
-        return True
-
-
-def sigmoid_pattern():
-    """Create pattern for sigmoid"""
-    dequant = is_op("qnn.dequantize")(wildcard(), is_constant(), is_constant())
-    sigmoid = is_op("sigmoid")(dequant)
-    quant = is_op("qnn.quantize")(sigmoid, is_constant(), is_constant())
-    return quant
-
-
 @register_pattern_table("ethos-u")
 def pattern_table() -> List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Callable]]:
     return [