Fixed lint errors and turned on the CMSIS_NN flag for config.cmake

Change-Id: I333eb8be30ad5a74dcb0f4cb732b7fe29a5de74c
apache · Aug 26, 2021 · 6a63601 · 6a63601
1 parent 3a21249
commit 6a63601
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 5 deletions.
diff --git a/python/tvm/relay/backend/contrib/cmsisnn/codegen.py b/python/tvm/relay/backend/contrib/cmsisnn/codegen.py
@@ -46,6 +46,12 @@ def is_quantized_softmax(self, call):
         """
         if not self.call_contains_op(call, "qnn.quantize"):
             return False
+        softmax_call = call.args[0]
+        if not self.call_contains_op(softmax_call, "nn.softmax"):
+            return False
+        dequantize_call = softmax_call.args[0]
+        if not self.call_contains_op(dequantize_call, "qnn.dequantize"):
+            return False
         self.scale = call.args[1].data.numpy().item(0)
         return True
 
@@ -56,8 +62,8 @@ def emit_softmax_tir(self, call):
         ir_builder = tvm.tir.ir_builder.create()
         in_buf = tvm.tir.decl_buffer(shape=shape, dtype=dtype)
         out_buf = tvm.tir.decl_buffer(shape=shape, dtype=dtype)
-        
-        trailing_dim = len(shape)-1
+
+        trailing_dim = len(shape) - 1
         num_rows = 1
         for dim in range(trailing_dim):
             num_rows *= shape[dim]

diff --git a/tests/python/driver/tvmc/conftest.py b/tests/python/driver/tvmc/conftest.py
@@ -190,6 +190,7 @@ def tflite_mobilenet_v1_0_25_128(tmpdir_factory):
 
     return model_file
 
+
 @pytest.fixture(scope="session")
 def tflite_cnn_s_quantized(tmpdir_factory):
     base_url = "https://github.com/ARM-software/ML-zoo/raw/master/models/keyword_spotting/cnn_small/tflite_int8/"

diff --git a/tests/python/driver/tvmc/test_compiler.py b/tests/python/driver/tvmc/test_compiler.py
@@ -316,9 +316,8 @@ def test_compile_tflite_module_with_external_codegen_cmsisnn(
     pytest.importorskip("tflite")
 
     output_dir = tmpdir_factory.mktemp("mlf")
-    output_dir = "./mlf"
     tvmc_model = tvmc.load(tflite_cnn_s_quantized)
-    
+
     output_file_name = f"{output_dir}/file.tar"
 
     tvmc_package = tvmc.compiler.compile_model(

diff --git a/tests/scripts/task_config_build_cpu.sh b/tests/scripts/task_config_build_cpu.sh
@@ -43,7 +43,7 @@ echo set\(USE_TENSORFLOW_PATH \"/tensorflow\"\) >> config.cmake
 echo set\(USE_FLATBUFFERS_PATH \"/flatbuffers\"\) >> config.cmake
 echo set\(USE_ETHOSN /opt/arm/ethosn-driver\) >> config.cmake
 echo set\(USE_ETHOSN_HW OFF\) >> config.cmake
-echo set\(USE_CMSISNN OFF\) >> config.cmake
+echo set\(USE_CMSISNN ON\) >> config.cmake
 echo set\(USE_VITIS_AI ON\) >> config.cmake
 echo set\(USE_VERILATOR ON\) >> config.cmake
 echo set\(USE_LIBBACKTRACE ON\) >> config.cmake