Merge branch 'main' into hexagon_qhl

python/tvm/tir/schedule/_type_checker.py

@@ -17,7 +17,7 @@
 """Type checking functionality"""
 import functools
 import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
 import typing
 
 
@@ -216,7 +216,10 @@ def _type_check(v: Any, name: str, type_: Any) -> Optional[str]:
     return _TYPE_CHECK[key](v, name, *subtypes)
 
 
-def type_checked(func: Callable) -> Callable:
+FType = TypeVar("FType", bound=Callable[..., Any])
+
+
+def type_checked(func: FType) -> FType:
     """Type check the input arguments of a function."""
     sig = inspect.signature(func)
 
@@ -236,4 +239,4 @@ def wrap(*args, **kwargs):
                     raise TypeError(error_msg)
         return func(*args, **kwargs)
 
-    return wrap
+    return wrap  # type: ignore

python/tvm/tir/schedule/schedule.py

@@ -2319,10 +2319,10 @@ def _normalize_buffer_arg(
         self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
     ) -> Tuple[str, int, Buffer]:
 
-        block_name = self.get(block).name_hint
+        block_obj: Block = self.get(block)
+        block_name = block_obj.name_hint
 
         def iter_buffers():
-            block_obj = self.get(block)
             for i, read in enumerate(block_obj.reads):
                 yield "read", i, read.buffer
             for i, write in enumerate(block_obj.writes):
@@ -2358,9 +2358,7 @@ def iter_buffers():
                 f"Expected 'read' or 'write', "
                 f"but received {buffer_index_type}"
             )
-            buffer_list = (
-                self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
-            )
+            buffer_list = block_obj.reads if buffer_index_type == "read" else block_obj.writes
             assert 0 <= buffer_index < len(buffer_list), (
                 f"Invalid buffer_index {buffer_index}.  "
                 f"Block {block_name} has only "

src/relay/backend/aot_executor_codegen.cc

@@ -494,7 +494,6 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     }
 
     tir::Stmt body = tir::SeqStmt({func_call});
-    LOG(INFO) << "CreateFuncCall: " << call_lowered_props.lowered_func->name_hint << " -> " << body;
     stmts_.push_back(body);
   }
 

src/relay/backend/contrib/cmsisnn/compiler_attrs.cc

@@ -29,7 +29,8 @@ namespace contrib {
 namespace cmsisnn {
 
 static const char* mveCPUs[] = {"cortex-m55"};
-static const char* dspCPUs[] = {"cortex-m4", "cortex-m7", "cortex-m33", "cortex-m35p"};
+static const char* dspCPUs[] = {"cortex-m55", "cortex-m4", "cortex-m7", "cortex-m33",
+                                "cortex-m35p"};
 
 TVM_REGISTER_NODE_TYPE(CMSISNNCompilerConfigNode);
 TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.cmsisnn.options", CMSISNNCompilerConfig);

tests/cpp/relay/backend/contrib/cmsisnn/compiler_attrs_test.cc

@@ -69,7 +69,7 @@ TEST_P(CMSISNNFlagsMVECPUs, CheckMVESet) {
 TEST_P(CMSISNNFlagsMVECPUs, CheckMVEOverrideCPU) {
   std::string mcpu = GetParam();
   CMSISNNFlags flags = GetFlagsWithCompilerAttrs(mcpu + "+nomve", "");
-  ASSERT_EQ(flags.dsp, false);
+  ASSERT_EQ(flags.dsp, true);
   ASSERT_EQ(flags.mve, false);
 }
 
@@ -92,7 +92,7 @@ TEST_P(CMSISNNFlagsMVECPUs, CheckCombinedOverrideCPU) {
 
 TEST_P(CMSISNNFlagsMVECPUs, CheckMVEOverrideMAttr) {
   CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), "+nomve");
-  ASSERT_EQ(flags.dsp, false);
+  ASSERT_EQ(flags.dsp, true);
   ASSERT_EQ(flags.mve, false);
 }
 

tests/python/contrib/test_cmsisnn/test_binary_ops.py

@@ -37,6 +37,7 @@
     get_range_for_dtype_str,
     assert_partitioned_function,
     assert_no_external_function,
+    create_test_runner,
 )
 
 
@@ -98,13 +99,22 @@ def make_model(
     ],
     [[0.256, 33, 0.256, 33], [0.0128, -64, 0.0128, -64], [0.0128, -64, 0.256, 33]],
 )
+@pytest.mark.parametrize(
+    "compiler_cpu, cpu_flags", [("cortex-m55", "+nomve"), ("cortex-m55", ""), ("cortex-m7", "")]
+)
 def test_op_int8(
-    op, relu_type, input_0_scale, input_0_zero_point, input_1_scale, input_1_zero_point
+    op,
+    relu_type,
+    input_0_scale,
+    input_0_zero_point,
+    input_1_scale,
+    input_1_zero_point,
+    compiler_cpu,
+    cpu_flags,
 ):
     """Tests QNN binary operator for CMSIS-NN"""
     interface_api = "c"
     use_unpacked_api = True
-    test_runner = AOT_USMP_CORSTONE300_RUNNER
 
     dtype = "int8"
     shape = [1, 16, 16, 3]
@@ -139,7 +149,7 @@ def test_op_int8(
             outputs=output_list,
             output_tolerance=1,
         ),
-        test_runner,
+        create_test_runner(compiler_cpu, cpu_flags),
         interface_api,
         use_unpacked_api,
     )

tests/python/contrib/test_cmsisnn/test_conv2d.py

@@ -39,6 +39,7 @@
     make_qnn_relu,
     assert_partitioned_function,
     assert_no_external_function,
+    create_test_runner,
 )
 
 
@@ -227,6 +228,9 @@ def test_conv2d_number_primfunc_args(
     "input_zero_point, input_scale, kernel_scale, out_channels",
     [(10, 0.0128, [0.11, 0.22], 2), (-64, 1, [1, 0.0256, 1.37], 3)],
 )
+@pytest.mark.parametrize(
+    "compiler_cpu, cpu_flags", [("cortex-m55", "+nomve"), ("cortex-m55", ""), ("cortex-m7", "")]
+)
 def test_conv2d_symmetric_padding_int8(
     padding,
     enable_bias,
@@ -235,11 +239,12 @@ def test_conv2d_symmetric_padding_int8(
     input_scale,
     kernel_scale,
     out_channels,
+    compiler_cpu,
+    cpu_flags,
 ):
     """Tests QNN Conv2D where the padding is symmetric on both sides of input"""
     interface_api = "c"
     use_unpacked_api = True
-    test_runner = AOT_USMP_CORSTONE300_RUNNER
 
     ifm_shape = (1, 64, 100, 4)
     kernel_size = (3, 3)
@@ -303,7 +308,7 @@ def test_conv2d_symmetric_padding_int8(
             params=params,
             output_tolerance=1,
         ),
-        test_runner,
+        create_test_runner(compiler_cpu, cpu_flags),
         interface_api,
         use_unpacked_api,
     )
@@ -456,6 +461,9 @@ def test_conv2d_int8_tflite(ifm_shape, kernel_shape, strides, dilation, padding,
     "input_zero_point, input_scale, kernel_scale, out_channels",
     [(10, 0.0128, [0.11, 0.22], 2), (-64, 1, [1, 0.0256, 1.37], 3)],
 )
+@pytest.mark.parametrize(
+    "compiler_cpu, cpu_flags", [("cortex-m55", "+nomve"), ("cortex-m55", ""), ("cortex-m7", "")]
+)
 def test_depthwise_int8(
     ifm_shape,
     kernel_size,
@@ -469,11 +477,12 @@ def test_depthwise_int8(
     kernel_scale,
     out_channels,
     depth_multiplier,
+    compiler_cpu,
+    cpu_flags,
 ):
     """Tests QNN Depthwise int8 op via CMSIS-NN"""
     interface_api = "c"
     use_unpacked_api = True
-    test_runner = AOT_USMP_CORSTONE300_RUNNER
 
     dtype = "int8"
     groups = 1
@@ -541,7 +550,7 @@ def test_depthwise_int8(
             params=params,
             output_tolerance=1,
         ),
-        test_runner,
+        create_test_runner(compiler_cpu, cpu_flags),
         interface_api,
         use_unpacked_api,
     )

tests/python/contrib/test_cmsisnn/test_fully_connected.py

@@ -24,16 +24,14 @@
 from tvm.relay.op.contrib import cmsisnn
 
 from tvm.testing.aot import generate_ref_data, AOTTestModel, compile_and_run
-from tvm.micro.testing.aot_test_utils import (
-    AOT_USMP_CORSTONE300_RUNNER,
-)
 from .utils import (
     make_module,
     get_range_for_dtype_str,
     get_conv2d_qnn_params,
     make_qnn_relu,
     assert_partitioned_function,
     assert_no_external_function,
+    create_test_runner,
 )
 
 
@@ -100,18 +98,22 @@ def make_model(
     "input_zero_point, input_scale, kernel_scale",
     [(10, 0.0128, 0.11), (-64, 0.0256, 1.37)],
 )
+@pytest.mark.parametrize(
+    "compiler_cpu, cpu_flags", [("cortex-m55", "+nomve"), ("cortex-m55", ""), ("cortex-m7", "")]
+)
 def test_op_int8(
     in_shape,
     enable_bias,
     input_zero_point,
     input_scale,
     kernel_scale,
     out_channels,
+    compiler_cpu,
+    cpu_flags,
 ):
     """Test QNN fully connected layer"""
     interface_api = "c"
     use_unpacked_api = True
-    test_runner = AOT_USMP_CORSTONE300_RUNNER
 
     dtype = "int8"
     kernel_zero_point = 0
@@ -160,7 +162,7 @@ def test_op_int8(
             params=params,
             output_tolerance=1,
         ),
-        test_runner,
+        create_test_runner(compiler_cpu, cpu_flags),
         interface_api,
         use_unpacked_api,
     )

tests/python/contrib/test_cmsisnn/test_pooling.py

@@ -31,6 +31,7 @@
     make_qnn_relu,
     assert_partitioned_function,
     assert_no_external_function,
+    create_test_runner,
 )
 
 
@@ -75,7 +76,6 @@ def make_model(
     return op
 
 
-@tvm.testing.requires_corstone300
 @tvm.testing.requires_cmsisnn
 @pytest.mark.parametrize("in_shape", [(1, 28, 28, 12), (1, 64, 100, 4)])
 @pytest.mark.parametrize(
@@ -84,6 +84,9 @@ def make_model(
 @pytest.mark.parametrize("relu_type", ["NONE", "RELU"])
 @pytest.mark.parametrize("pool_type", [relay.nn.max_pool2d, relay.nn.avg_pool2d])
 @pytest.mark.parametrize("zero_point, scale", [(-34, 0.0256)])
+@pytest.mark.parametrize(
+    "compiler_cpu, cpu_flags", [("cortex-m55", "+nomve"), ("cortex-m55", ""), ("cortex-m7", "")]
+)
 def test_op_int8(
     in_shape,
     pool_size,
@@ -93,11 +96,12 @@ def test_op_int8(
     pool_type,
     zero_point,
     scale,
+    compiler_cpu,
+    cpu_flags,
 ):
     """Tests QNN pooling op for int8 inputs"""
     interface_api = "c"
     use_unpacked_api = True
-    test_runner = AOT_USMP_CORSTONE300_RUNNER
 
     dtype = "int8"
 
@@ -133,7 +137,7 @@ def test_op_int8(
             params=None,
             output_tolerance=1,
         ),
-        test_runner,
+        create_test_runner(compiler_cpu, cpu_flags),
         interface_api,
         use_unpacked_api,
     )

tests/python/contrib/test_cmsisnn/test_softmax.py

@@ -25,14 +25,14 @@
 from tvm import relay
 from tvm.relay.op.contrib import cmsisnn
 from tvm.testing.aot import AOTTestModel, compile_and_run, generate_ref_data
-from tvm.micro.testing.aot_test_utils import AOT_USMP_CORSTONE300_RUNNER
 
 from .utils import (
     skip_if_no_reference_system,
     make_module,
     get_range_for_dtype_str,
     assert_partitioned_function,
     assert_no_external_function,
+    create_test_runner,
 )
 
 
@@ -57,13 +57,15 @@ def make_model(
 
 
 @skip_if_no_reference_system
-@pytest.mark.parametrize(["zero_point", "scale"], [[33, 0.256], [-64, 0.0128]])
 @tvm.testing.requires_cmsisnn
-def test_op_int8(zero_point, scale):
+@pytest.mark.parametrize(["zero_point", "scale"], [[33, 0.256], [-64, 0.0128]])
+@pytest.mark.parametrize(
+    "compiler_cpu, cpu_flags", [("cortex-m55", "+nomve"), ("cortex-m55", ""), ("cortex-m7", "")]
+)
+def test_op_int8(zero_point, scale, compiler_cpu, cpu_flags):
     """Tests int8 QNN Softmax for CMSIS-NN"""
     interface_api = "c"
     use_unpacked_api = True
-    test_runner = AOT_USMP_CORSTONE300_RUNNER
 
     dtype = "int8"
     shape = [1, 16, 16, 3]
@@ -84,7 +86,7 @@ def test_op_int8(zero_point, scale):
     output_list = generate_ref_data(orig_mod["main"], inputs, params)
     compile_and_run(
         AOTTestModel(module=cmsisnn_mod, inputs=inputs, outputs=output_list, params=params),
-        test_runner,
+        create_test_runner(compiler_cpu, cpu_flags),
         interface_api,
         use_unpacked_api,
     )

tests/python/contrib/test_cmsisnn/utils.py

@@ -23,6 +23,7 @@
 
 import tvm
 from tvm import relay
+from tvm.testing.aot import AOTTestRunner
 
 
 def skip_if_no_reference_system(func):
@@ -225,3 +226,44 @@ def make_qnn_relu(expr, fused_activation_fn, scale, zero_point, dtype):
     if fused_activation_fn == "RELU":
         return tvm.relay.op.clip(expr, a_min=max(qmin, quantize(0.0)), a_max=qmax)
     raise ValueError("Invalid argument provided with fused_activation_fn")
+
+
+def create_test_runner(compiler_cpu="cortex-m55", cpu_flags=""):
+    """
+    Creates AOT test runner for CMSIS-NN tests.
+
+    Parameters
+    ----------
+    compiler_cpu : str
+       Equivalent of gcc option mcpu
+       Options:  cortex-m55, cortex-m7
+    cpu_flags: str
+        Disable Arm(R) Cortex(R)-M profile vector extension (mve)
+        Options:
+        Arm(R) Cortex(R)-M55: when null +mve is set by default.
+            +nomve disables vector extensions.
+        Arm(R) Cortex(R)-M7 does not support mve.
+    """
+    # cmsis_cpu is used to find out start up code inside CMSIS package
+    cmsis_cpu = "ARMCM7" if compiler_cpu == "cortex-m7" else "ARMCM55"
+    mfloat_abi = "soft" if compiler_cpu == "cortex-m7" else "hard"
+    return AOTTestRunner(
+        makefile="corstone300",
+        prologue="""
+        uart_init();
+        """,
+        includes=["uart.h"],
+        pass_config={
+            "relay.ext.cmsisnn.options": {
+                "mcpu": compiler_cpu + cpu_flags,
+            },
+            "tir.usmp.enable": True,
+            "tir.disable_storage_rewrite": True,
+        },
+        parameters={
+            "ARM_CPU": cmsis_cpu,
+            "MCPU": compiler_cpu,
+            "MCPU_FLAGS": cpu_flags,
+            "MFLOAT_ABI": mfloat_abi,
+        },
+    )