Improve ArmTester logging (#5629)

Summary:
See commits for description

Pull Request resolved: #5629

Reviewed By: digantdesai

Differential Revision: D63637272

Pulled By: mergennachin

fbshipit-source-id: c0025dcbe41b0e8c7c20252af4f31531d230712e

Author: Erik-Lundell

backends/arm/test/common.py

@@ -4,9 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 import os
 import shutil
 import subprocess
+import sys
 import tempfile
 
 import pytest
@@ -37,6 +39,7 @@ def pytest_configure(config):
                 "Tests are run with --arm_run_corstone300 but corstone300 FVP is not installed."
             )
         _enabled_options.append("corstone300")
+    logging.basicConfig(level=logging.INFO, stream=sys.stdout)
 
 
 def pytest_collection_modifyitems(config, items):

backends/arm/test/misc/test_debug_feats.py

@@ -127,29 +127,52 @@ def test_numerical_diff_prints(self):
             self.fail()
 
 
-class TestDumpOperatorsAndDtypes(unittest.TestCase):
-    def test_dump_ops_and_dtypes(self):
-        model = Linear(20, 30)
-        (
-            ArmTester(
-                model,
-                example_inputs=model.get_inputs(),
-                compile_spec=common.get_tosa_compile_spec(),
-            )
-            .quantize()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
-            .export()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
-            .to_edge()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
-            .partition()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
+def test_dump_ops_and_dtypes():
+    model = Linear(20, 30)
+    (
+        ArmTester(
+            model,
+            example_inputs=model.get_inputs(),
+            compile_spec=common.get_tosa_compile_spec(),
+        )
+        .quantize()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+        .export()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+        .to_edge()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+        .partition()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+    )
+    # Just test that there are no execptions.
+
+
+def test_dump_ops_and_dtypes_parseable():
+    model = Linear(20, 30)
+    (
+        ArmTester(
+            model,
+            example_inputs=model.get_inputs(),
+            compile_spec=common.get_tosa_compile_spec(),
         )
-        # Just test that there are no execeptions.
+        .quantize()
+        .dump_dtype_distribution(print_table=False)
+        .dump_operator_distribution(print_table=False)
+        .export()
+        .dump_dtype_distribution(print_table=False)
+        .dump_operator_distribution(print_table=False)
+        .to_edge()
+        .dump_dtype_distribution(print_table=False)
+        .dump_operator_distribution(print_table=False)
+        .partition()
+        .dump_dtype_distribution(print_table=False)
+        .dump_operator_distribution(print_table=False)
+    )
+    # Just test that there are no execptions.
 
 
 class TestCollateTosaTests(unittest.TestCase):
@@ -186,3 +209,41 @@ def test_collate_tosa_BI_tests(self):
 
         os.environ.pop("TOSA_TESTCASES_BASE_PATH")
         shutil.rmtree("test_collate_tosa_tests", ignore_errors=True)
+
+
+def test_dump_tosa_ops(caplog):
+    caplog.set_level(logging.INFO)
+    model = Linear(20, 30)
+    (
+        ArmTester(
+            model,
+            example_inputs=model.get_inputs(),
+            compile_spec=common.get_tosa_compile_spec(),
+        )
+        .quantize()
+        .export()
+        .to_edge()
+        .partition()
+        .dump_operator_distribution()
+    )
+    assert "TOSA operators:" in caplog.text
+
+
+def test_fail_dump_tosa_ops(caplog):
+    caplog.set_level(logging.INFO)
+
+    class Add(torch.nn.Module):
+        def forward(self, x):
+            return x + x
+
+    model = Add()
+    compile_spec = common.get_u55_compile_spec()
+    (
+        ArmTester(model, example_inputs=(torch.ones(5),), compile_spec=compile_spec)
+        .quantize()
+        .export()
+        .to_edge()
+        .partition()
+        .dump_operator_distribution()
+    )
+    assert "Can not get operator distribution for Vela command stream." in caplog.text