Rebase and fix the review commits

Author: ekalda

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

@@ -141,14 +141,14 @@ def get_lut_from_func(ifm_scale, ifm_zp, ofm_scale, ofm_zp, func):
     return lut_values
 
 
-class TanhRewriter(DFPatternCallback):
-    """This pass adds tanh as a LUT to the identity operator"""
+class LutActivationRewriter(DFPatternCallback):
+    """A class to create an identity operator with the LUT"""
 
-    def __init__(self):
+    def __init__(self, params_class, activation_type, calc_func):
         super().__init__(require_type=True, rewrite_once=True)
-        self.pattern = (
-            wildcard().has_attr({"Composite": ethosu_patterns.TanhParams.composite_name})
-        )(wildcard())
+        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):
         id_input = post.args[0]
@@ -161,7 +161,9 @@ def callback(self, pre, post, node_map):
         input_scale = float(dequantize_args[1].data.asnumpy())
         input_zp = int(dequantize_args[2].data.asnumpy())
 
-        lut_values = get_lut_from_func(input_scale, input_zp, output_scale, output_zp, math.tanh)
+        lut_values = get_lut_from_func(
+            input_scale, input_zp, output_scale, output_zp, self.calc_func
+        )
         lut = relay.const(lut_values, dtype="uint8")
 
         # We baked the requantization into the LUT, so we don't requantize the identity operator
@@ -172,12 +174,21 @@ def callback(self, pre, post, node_map):
             ifm_zero_point=input_zp,
             ofm_scale=input_scale,
             ofm_zero_point=input_zp,
-            activation="TANH",
+            activation=self.activation_type,
         )
 
         return identity
 
 
+class TanhRewriter(LutActivationRewriter):
+    """This pass adds tanh as a LUT to the identity operator"""
+
+    def __init__(self):
+        super().__init__(
+            params_class=ethosu_patterns.TanhParams, activation_type="TANH", calc_func=math.tanh
+        )
+
+
 @ir.transform.module_pass(opt_level=1)
 class LegalizeTanh:
     """This is the pass that wraps TanhRewriter"""
@@ -209,44 +220,16 @@ def sigmoid_calc_func(x):
     return y
 
 
-class SigmoidRewriter(DFPatternCallback):
+class SigmoidRewriter(LutActivationRewriter):
     """This pass adds sigmoid as a LUT for identity op"""
 
     def __init__(self):
-        super().__init__(require_type=True, rewrite_once=True)
-        self.pattern = (
-            wildcard().has_attr({"Composite": ethosu_patterns.SigmoidParams.composite_name})
-        )(wildcard())
-
-    def callback(self, pre, post, node_map):
-        inp = post.args[0]
-
-        quantize_args = post.op.body.args
-        output_scale = float(quantize_args[1].data.asnumpy())
-        output_zp = int(quantize_args[2].data.asnumpy())
-
-        dequantize_args = quantize_args[0].args[0].args
-        input_scale = float(dequantize_args[1].data.asnumpy())
-        input_zp = int(dequantize_args[2].data.asnumpy())
-
-        lut_values = get_lut_from_func(
-            input_scale, input_zp, output_scale, output_zp, sigmoid_calc_func
-        )
-        lut = relay.const(lut_values, dtype="uint8")
-
-        # We baked the requantization into the LUT, so we don't requantize the identity operator
-        identity = ethosu_ops.ethosu_identity(
-            ifm=inp,
-            lut=lut,
-            ifm_scale=input_scale,
-            ifm_zero_point=input_zp,
-            ofm_scale=input_scale,
-            ofm_zero_point=input_zp,
-            activation="SIGMOID",
+        super().__init__(
+            params_class=ethosu_patterns.SigmoidParams,
+            activation_type="SIGMOID",
+            calc_func=sigmoid_calc_func,
         )
 
-        return identity
-
 
 @ir.transform.module_pass(opt_level=1)
 class LegalizeSigmoid:

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

@@ -1049,7 +1049,7 @@ def __init__(self, func_body: Call):
 
     def is_valid(self):
         """
-        This function checks whether reshape has compatible attributes with the NPU
+        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

tests/python/contrib/test_ethosu/test_codegen.py

@@ -826,131 +826,26 @@ def clz_comp(n):
 
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 def test_tflite_tanh(accel_type):
-    dtype = "int8"
     ifm_shape = [1, 115, 32, 7]
 
-    def create_tflite_graph():
-        class Model(tf.Module):
-            @tf.function
-            def tanh_function(self, x):
-                op = tf.nn.tanh(x)
-                return op
-
-        model = Model()
-        concrete_func = model.tanh_function.get_concrete_function(
-            tf.TensorSpec(ifm_shape, dtype=tf.float32)
-        )
-
-        # Convert the model
-        def representative_dataset():
-            for _ in range(100):
-                data = np.random.rand(*tuple(ifm_shape))
-                yield [data.astype(np.float32)]
-
-        converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
-        converter.optimizations = [tf.lite.Optimize.DEFAULT]
-        converter.representative_dataset = representative_dataset
-        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
-        converter.inference_input_type = tf.int8
-        converter.inference_output_type = tf.int8
-        tflite_model = converter.convert()
-        return tflite_model
-
-    tflite_graph = create_tflite_graph()
-
-    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
-
-    relay_module, params = relay.frontend.from_tflite(
-        tflite_model,
-        shape_dict={"input": ifm_shape},
-        dtype_dict={"input": dtype},
-    )
-    mod = partition_for_ethosu(relay_module, params)
-
-    # Generate reference data
-    input_data, output_data = infra.generate_ref_data_tflite(tflite_graph)
-
-    compiled_models = infra.build_source(
-        mod,
-        input_data,
-        output_data,
-        accel_type,
-    )
-
-    # Assumes only two runtime.Modules are created -- i.e. single offload module
-    ethosu_module = compiled_models[0].executor_factory.lib.imported_modules[0].imported_modules[0]
+    @tf.function
+    def tanh_func(x):
+        op = tf.nn.tanh(x)
+        return op
 
-    # Verify generated C source
-    get_artifacts = tvm._ffi.get_global_func("runtime.module.ethos-u.get_artifacts")
-    compilation_artifacts = get_artifacts(ethosu_module)
-    cmms = bytes.fromhex(compilation_artifacts[0].command_stream)
-    infra.print_payload(cmms)
-    infra.verify_source(compiled_models, accel_type)
+    _compare_tvm_with_tflite(tanh_func, [ifm_shape], accel_type)
 
 
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
-@pytest.mark.parametrize("ifm_shape", [[1, 115, 32, 7], [1, 4, 5, 2]])
-def test_tflite_sigmoid(accel_type, ifm_shape):
-    dtype = "int8"
-
-    def create_tflite_graph():
-        tf.config.run_functions_eagerly(True)
-
-        class Model(tf.Module):
-            @tf.function
-            def tanh_function(self, x):
-                op = tf.nn.sigmoid(x)
-                return op
-
-        model = Model()
-        concrete_func = model.tanh_function.get_concrete_function(
-            tf.TensorSpec(ifm_shape, dtype=tf.float32)
-        )
-
-        # Convert the model
-        def representative_dataset():
-            for _ in range(100):
-                data = np.random.rand(*tuple(ifm_shape))
-                yield [data.astype(np.float32)]
-
-        converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
-        converter.optimizations = [tf.lite.Optimize.DEFAULT]
-        converter.representative_dataset = representative_dataset
-        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
-        converter.inference_input_type = tf.int8
-        converter.inference_output_type = tf.int8
-        tflite_model = converter.convert()
-        return tflite_model
-
-    tflite_graph = create_tflite_graph()
-    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
-
-    relay_module, params = relay.frontend.from_tflite(
-        tflite_model,
-        shape_dict={"input": ifm_shape},
-        dtype_dict={"input": dtype},
-    )
-    mod = partition_for_ethosu(relay_module, params)
-
-    # Generate reference data
-    input_data, output_data = infra.generate_ref_data_tflite(tflite_graph)
-
-    compiled_models = infra.build_source(
-        mod,
-        input_data,
-        output_data,
-        accel_type,
-    )
+def test_tflite_sigmoid(accel_type):
+    ifm_shape = [1, 135, 41, 6]
 
-    # Assumes only two runtime.Modules are created -- i.e. single offload module
-    ethosu_module = compiled_models[0].executor_factory.lib.imported_modules[0].imported_modules[0]
+    @tf.function
+    def sigmoid_function(x):
+        op = tf.nn.sigmoid(x)
+        return op
 
-    # Verify generated C source
-    get_artifacts = tvm._ffi.get_global_func("runtime.module.ethos-u.get_artifacts")
-    compilation_artifacts = get_artifacts(ethosu_module)
-    cmms = bytes.fromhex(compilation_artifacts[0].command_stream)
-    infra.print_payload(cmms)
-    infra.verify_source(compiled_models, accel_type)
+    _compare_tvm_with_tflite(sigmoid_function, [ifm_shape], accel_type)
 
 
 if __name__ == "__main__":