[CI] Apply linting rules to AOT tests

tests/lint/pylint.sh

@@ -21,4 +21,5 @@ python3 -m pylint python/tvm --rcfile="$(dirname "$0")"/pylintrc
 python3 -m pylint vta/python/vta --rcfile="$(dirname "$0")"/pylintrc
 python3 -m pylint tests/python/unittest/test_tvmscript_type.py --rcfile="$(dirname "$0")"/pylintrc
 python3 -m pylint tests/python/contrib/test_cmsisnn --rcfile="$(dirname "$0")"/pylintrc
+python3 -m pylint tests/python/relay/aot/*.py --rcfile="$(dirname "$0")"/pylintrc
 

tests/python/relay/aot/test_c_device_api.py

@@ -14,32 +14,38 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+"""AOT with C Device API Tests"""
 
-import sys
+import re
 from collections import OrderedDict
 
 import numpy as np
 import pytest
-import re
-import tvm.testing
 
+import tvm.testing
 from tvm import relay
 from tvm.ir.module import IRModule
 from tvm.testing.aot import AOTTestModel, generate_ref_data, compile_models
 from tvm.micro.testing.aot_test_utils import AOT_DEFAULT_RUNNER
 
 
-@pytest.fixture
-def device_api_main_func():
+@pytest.fixture(name="device_api_main_func")
+def fixture_device_api_main_func():
+    """Test function generator which generates C Device API calls"""
+
     # Ideally we should have a sample Target registered here
     # but we're going to re-use this for now
     pytest.importorskip("ethosu.vela")
+
+    # pylint: disable=import-outside-toplevel
     import tensorflow as tf
     import tflite.Model
 
     from tests.python.contrib.test_ethosu.infra import create_test_runner, generate_ref_data_tflite
     from tvm.relay.op.contrib.ethosu import partition_for_ethosu
 
+    # pylint: enable=import-outside-toplevel
+
     tf.config.run_functions_eagerly(True)
 
     class Model(tf.Module):
@@ -97,8 +103,9 @@ def compile_to_main_func(interface_api="c", use_unpacked_api=True):
     return compile_to_main_func
 
 
-@pytest.fixture
-def non_device_api_main_func():
+@pytest.fixture(name="non_device_api_main_func")
+def fixture_non_device_api_main_func():
+    """Test function generator which does not generate C Device API calls"""
     x = relay.var("x", shape=(10, 10))
     y = relay.var("y", shape=(1, 10))
     func = relay.Function([x, y], relay.multiply(x, y))
@@ -151,7 +158,10 @@ def test_device_api_hooks_unpacked_api(device_api_main_func):
     # We dont need to check exact input and output var names in this test.
     # Hence, using a regex to cover any legal I/O name.
     regex = re.compile(
-        'tir\.tvm_check_return\(0, -1, tir\.call_extern\("tvmgen_default_ethos_u_main_0", \w+, \w+, device_context_ethos_u\)\)'
+        r"tir\.tvm_check_return\("
+        r"0, -1, "
+        r'tir\.call_extern\("tvmgen_default_ethos_u_main_0", '
+        r"\w+, \w+, device_context_ethos_u\)\)"
     )
     assert regex.match(str(main_func.body[1][0][0][1]))
     # Close Device
@@ -171,7 +181,9 @@ def test_device_api_hooks_unpacked_api(device_api_main_func):
 
 
 @pytest.mark.skip(
-    "Skipping this test as this is incorrectly using Arm(R) Ethos(TM)-U NPU with packed calling convention which is not supported by the NPU codegen's TIR to Runtime Hook. We need to use a different target to test this feature"
+    "Skipping this test as this is incorrectly using Arm(R) Ethos(TM)-U NPU "
+    "with packed calling convention which is not supported by the NPU codegen's "
+    "TIR to Runtime Hook. We need to use a different target to test this feature"
 )
 def test_device_api_hooks_packed_api(device_api_main_func):
     """Check for Device API hooks with packed internal calls"""
@@ -236,11 +248,12 @@ def test_without_device_api_packed_api(non_device_api_main_func):
     """Test a graph without the Device API with the packed internal calls"""
 
     main_func = non_device_api_main_func(interface_api="packed", use_unpacked_api=False)
+
     assert str(main_func.body) == (
         'tir.tvm_call_cpacked("tvmgen_default_fused_multiply", '
-        "tir.tvm_stack_make_array(x_buffer_var, tir.tvm_stack_make_shape(10, 10), tir.reinterpret((uint64)0), (uint32)2, float32(0), 0), "
-        "tir.tvm_stack_make_array(y_buffer_var, tir.tvm_stack_make_shape(1, 10), tir.reinterpret((uint64)0), (uint32)2, float32(0), 0), "
-        "tir.tvm_stack_make_array(output_buffer_var, tir.tvm_stack_make_shape(10, 10), tir.reinterpret((uint64)0), (uint32)2, float32(0), 0), "
+        "tir.tvm_stack_make_array(x_buffer_var, tir.tvm_stack_make_shape(10, 10), tir.reinterpret((uint64)0), (uint32)2, float32(0), 0), "  # pylint: disable=line-too-long
+        "tir.tvm_stack_make_array(y_buffer_var, tir.tvm_stack_make_shape(1, 10), tir.reinterpret((uint64)0), (uint32)2, float32(0), 0), "  # pylint: disable=line-too-long
+        "tir.tvm_stack_make_array(output_buffer_var, tir.tvm_stack_make_shape(10, 10), tir.reinterpret((uint64)0), (uint32)2, float32(0), 0), "  # pylint: disable=line-too-long
         "tir.reinterpret((uint64)0))\n"
     )
 

tests/python/relay/aot/test_cpp_aot.py

@@ -14,10 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
+"""AOT with C++ Runtime Tests"""
 
 import re
-import sys
 import textwrap
 
 import numpy as np
@@ -28,13 +27,10 @@
 from tvm import relay
 from tvm.relay import backend, testing
 from tvm.testing.aot import generate_ref_data
-from tvm.micro.testing.aot_test_utils import AOT_DEFAULT_RUNNER
 
 
 def test_error_c_interface():
-    interface_api = "c"
-    use_unpacked_api = False
-    test_runner = AOT_DEFAULT_RUNNER
+    """Checks that an error occurs when using the packed API in combination with C interface"""
 
     two = relay.add(relay.const(1), relay.const(1))
     func = relay.Function([], two)
@@ -53,12 +49,11 @@ def test_error_c_interface():
         )
 
 
-enable_usmp = tvm.testing.parameter(True, False)
-target_kind = tvm.testing.parameter("c", "llvm")
-
-
+@pytest.mark.parametrize("enable_usmp", [True, False])
+@pytest.mark.parametrize("target_kind", ["c", "llvm"])
 def test_conv2d(enable_usmp, target_kind):
-    RELAY_MODEL = textwrap.dedent(
+    """Tests compilation of convolutions"""
+    relay_model = textwrap.dedent(
         """\
         #[version = "0.0.5"]
         def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(3, 3, 5, 5), int8]) {
@@ -86,7 +81,7 @@ def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(3, 3, 5, 5),
         }
     """
     )
-    ir_mod = tvm.parser.fromtext(RELAY_MODEL)
+    ir_mod = tvm.parser.fromtext(relay_model)
 
     main_func = ir_mod["main"]
     shape_dict = {p.name_hint: p.checked_type.concrete_shape for p in main_func.params}
@@ -119,7 +114,10 @@ def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(3, 3, 5, 5),
     assert (runner.get_output(0).asnumpy() == list(ref_outputs.values())[0]).all()
 
 
+@pytest.mark.parametrize("enable_usmp", [True, False])
+@pytest.mark.parametrize("target_kind", ["c", "llvm"])
 def test_mobilenet(enable_usmp, target_kind):
+    """Full network test with Mobilenet"""
     ir_mod, params = testing.mobilenet.get_workload(batch_size=1)
     data_shape = [int(x) for x in ir_mod["main"].checked_type.arg_types[0].shape]
     data = np.random.uniform(size=data_shape).astype("float32")
@@ -147,10 +145,11 @@ def test_mobilenet(enable_usmp, target_kind):
 
 
 def test_module_list():
-    x = tvm.relay.var("x", tvm.relay.TensorType([1], dtype="float32"))
-    expr = tvm.relay.add(x, tvm.relay.Constant(tvm.nd.array(np.array([1], dtype="float32"))))
+    """Checks the correct list of module names is generated"""
+    input_x = tvm.relay.var("x", tvm.relay.TensorType([1], dtype="float32"))
+    expr = tvm.relay.add(input_x, tvm.relay.Constant(tvm.nd.array(np.array([1], dtype="float32"))))
     mod = tvm.relay.build(
-        tvm.IRModule.from_expr(tvm.relay.Function([x], expr)),
+        tvm.IRModule.from_expr(tvm.relay.Function([input_x], expr)),
         target="c",
         executor=tvm.relay.backend.Executor("aot", {"interface-api": "packed"}),
         mod_name="unusual_module_name_fred",
@@ -177,6 +176,7 @@ def test_create_executor():
 
 
 def test_pass_wrong_device_arg():
+    """Ensure an error is generated if the incorrect number of devices are passed"""
     x = tvm.relay.var("x", tvm.relay.TensorType([1], dtype="float32"))
     expr = tvm.relay.add(x, tvm.relay.Constant(tvm.nd.array(np.array([1], dtype="float32"))))
     with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
@@ -191,12 +191,12 @@ def test_pass_wrong_device_arg():
     mod.export_library(test_so_path, cc="gcc", options=["-std=c11"])
     loaded_mod = tvm.runtime.load_module(test_so_path)
 
-    with pytest.raises(tvm.TVMError) as cm:
+    with pytest.raises(tvm.TVMError) as error:
         tvm.runtime.executor.AotModule(loaded_mod["default"](tvm.cpu(0), tvm.cpu(0)))
 
         assert (
             "Check failed: devices_.size() == 1 (2 vs. 1) : Expect exactly 1 device passed."
-            in str(cm.exception)
+            in str(error.exception)
         )
     # TODO write asserts for # and type of device.