Move callbacks into _compare_outputs

This is to make run_method_and_compare_outputs less complex since lintrunner was complaining.
Additionally moves out previous info dumps in compare_output into a new callback function to handle all error handling in the same way.

Change-Id: If7c7cfd515d4870c018a34785dd80be15d4fbcef

Author: AdrianLundell

backends/arm/test/tester/analyze_output_utils.py

@@ -4,8 +4,16 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+import tempfile
 
 import torch
+from executorch.backends.arm.test.runner_utils import (
+    _get_input_quantization_params,
+    _get_output_node,
+    _get_output_quantization_params,
+)
+
+from executorch.backends.xnnpack.test.tester.tester import Export, Quantize
 
 logger = logging.getLogger(__name__)
 
@@ -71,6 +79,7 @@ def _print_elements(result, reference, C, H, W, rtol, atol):
 
 
 def print_error_diffs(
+    tester,
     result: torch.Tensor | tuple,
     reference: torch.Tensor | tuple,
     quantization_scale=None,
@@ -193,7 +202,50 @@ def print_error_diffs(
     end_separator = "#" * separator_length + "\n"
     output_str += end_separator
 
-    logger.info(output_str)
+    logger.error(output_str)
+
+
+def dump_error_output(
+    tester,
+    reference_output,
+    stage_output,
+    quantization_scale=None,
+    atol=1e-03,
+    rtol=1e-03,
+    qtol=0,
+):
+    """
+    Prints Quantization info and error tolerances, and saves the differing tensors to disc.
+    """
+    # Capture assertion error and print more info
+    banner = "=" * 40 + "TOSA debug info" + "=" * 40
+    logger.error(banner)
+    path_to_tosa_files = tester.runner_util.intermediate_path
+
+    if path_to_tosa_files is None:
+        path_to_tosa_files = tempfile.mkdtemp(prefix="executorch_result_dump_")
+
+    export_stage = tester.stages.get(tester.stage_name(Export), None)
+    quantize_stage = tester.stages.get(tester.stage_name(Quantize), None)
+    if export_stage is not None and quantize_stage is not None:
+        output_node = _get_output_node(export_stage.artifact)
+        qp_input = _get_input_quantization_params(export_stage.artifact)
+        qp_output = _get_output_quantization_params(export_stage.artifact, output_node)
+        logger.error(f"Input QuantArgs: {qp_input}")
+        logger.error(f"Output QuantArgs: {qp_output}")
+
+    logger.error(f"{path_to_tosa_files=}")
+    import os
+
+    torch.save(
+        stage_output,
+        os.path.join(path_to_tosa_files, "torch_tosa_output.pt"),
+    )
+    torch.save(
+        reference_output,
+        os.path.join(path_to_tosa_files, "torch_ref_output.pt"),
+    )
+    logger.error(f"{atol=}, {rtol=}, {qtol=}")
 
 
 if __name__ == "__main__":

backends/arm/test/tester/arm_tester.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
-import tempfile
 
 from collections import Counter
 from pprint import pformat
@@ -25,14 +24,11 @@
 )
 from executorch.backends.arm.test.common import get_target_board
 
-from executorch.backends.arm.test.runner_utils import (
-    _get_input_quantization_params,
-    _get_output_node,
-    _get_output_quantization_params,
-    dbg_tosa_fb_to_json,
-    RunnerUtil,
+from executorch.backends.arm.test.runner_utils import dbg_tosa_fb_to_json, RunnerUtil
+from executorch.backends.arm.test.tester.analyze_output_utils import (
+    dump_error_output,
+    print_error_diffs,
 )
-from executorch.backends.arm.test.tester.analyze_output_utils import print_error_diffs
 from executorch.backends.arm.tosa_mapping import extract_tensor_meta
 
 from executorch.backends.xnnpack.test.tester import Tester
@@ -279,7 +275,7 @@ def run_method_and_compare_outputs(
         atol=1e-03,
         rtol=1e-03,
         qtol=0,
-        callback=print_error_diffs,
+        error_callbacks=None,
     ):
         """
         Compares the run_artifact output of 'stage' with the output of a reference stage.
@@ -367,20 +363,15 @@ def run_method_and_compare_outputs(
             ):
                 test_output = self.transpose_data_format(test_output, "NCHW")
 
-            try:
-                self._compare_outputs(
-                    reference_output, test_output, quantization_scale, atol, rtol, qtol
-                )
-            except AssertionError as e:
-                callback(
-                    reference_output,
-                    test_output,
-                    quantization_scale,
-                    atol,
-                    rtol,
-                    qtol,
-                )
-                raise e
+            self._compare_outputs(
+                reference_output,
+                test_output,
+                quantization_scale,
+                atol,
+                rtol,
+                qtol,
+                error_callbacks,
+            )
 
         return self
 
@@ -528,42 +519,25 @@ def _compare_outputs(
         atol=1e-03,
         rtol=1e-03,
         qtol=0,
+        error_callbacks=None,
     ):
         try:
             super()._compare_outputs(
                 reference_output, stage_output, quantization_scale, atol, rtol, qtol
             )
         except AssertionError as e:
-            # Capture assertion error and print more info
-            banner = "=" * 40 + "TOSA debug info" + "=" * 40
-            logger.error(banner)
-            path_to_tosa_files = self.runner_util.intermediate_path
-            if path_to_tosa_files is None:
-                path_to_tosa_files = tempfile.mkdtemp(prefix="executorch_result_dump_")
-
-            export_stage = self.stages.get(self.stage_name(tester.Export), None)
-            quantize_stage = self.stages.get(self.stage_name(tester.Quantize), None)
-            if export_stage is not None and quantize_stage is not None:
-                output_node = _get_output_node(export_stage.artifact)
-                qp_input = _get_input_quantization_params(export_stage.artifact)
-                qp_output = _get_output_quantization_params(
-                    export_stage.artifact, output_node
+            if error_callbacks is None:
+                error_callbacks = [print_error_diffs, dump_error_output]
+            for callback in error_callbacks:
+                callback(
+                    self,
+                    reference_output,
+                    stage_output,
+                    quantization_scale=None,
+                    atol=1e-03,
+                    rtol=1e-03,
+                    qtol=0,
                 )
-                logger.error(f"Input QuantArgs: {qp_input}")
-                logger.error(f"Output QuantArgs: {qp_output}")
-
-            logger.error(f"{path_to_tosa_files=}")
-            import os
-
-            torch.save(
-                stage_output,
-                os.path.join(path_to_tosa_files, "torch_tosa_output.pt"),
-            )
-            torch.save(
-                reference_output,
-                os.path.join(path_to_tosa_files, "torch_ref_output.pt"),
-            )
-            logger.error(f"{atol=}, {rtol=}, {qtol=}")
             raise e