Arm backend: Dump partition report to file if intermediate dir is given (#19545)

Dump to <intermediate_dir>/partition_report.txt
Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/arm/ethosu/partitioner.py

@@ -34,3 +34,4 @@ def __init__(
         self.additional_checks = additional_checks
         self.tosa_spec = compile_spec.tosa_spec
         self._custom_partition_ops: set[torch._ops.OpOverload] = set()
+        self.intermediate_path = compile_spec._get_intermediate_path()

backends/arm/test/misc/test_debug_feats.py

@@ -219,9 +219,13 @@ def test_collate_tests_tosa_INT(test_data: input_t1):
     )
     # test that the output directory is created and contains the expected files
     assert os.path.exists(test_collate_dir)
+    found_file = False
     for file in os.listdir(test_collate_dir):
         file_name_prefix = "TOSA-1.0+INT"
-        assert file.endswith((f"{file_name_prefix}.json", f"{file_name_prefix}.tosa"))
+        found_file |= file.endswith(
+            (f"{file_name_prefix}.json", f"{file_name_prefix}.tosa")
+        )
+    assert found_file
 
     os.environ.pop("TOSA_TESTCASES_BASE_PATH")
     shutil.rmtree("test_collate_tosa_tests", ignore_errors=True)

backends/arm/tosa/partitioner.py

@@ -16,6 +16,7 @@
 import logging
 import operator
 from itertools import count
+from pathlib import Path
 from typing import Callable, cast, List, Optional, Sequence, Tuple
 
 import torch
@@ -177,6 +178,7 @@ def __init__(
         self.tosa_spec = compile_spec.tosa_spec
         self.additional_checks = additional_checks
         self._custom_partition_ops: set[torch._ops.OpOverload] = set()
+        self.intermediate_path = compile_spec._get_intermediate_path()
 
     def register_custom_partition_op(self, op: torch._ops.OpOverload) -> None:
         """Register a custom op to be considered supported by this
@@ -430,6 +432,13 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         partition_tags = {tag: self.delegation_spec for tag in tags}
 
         tag_constant_data(exported_program)
+        if self.intermediate_path is not None and logger.level <= logging.INFO:
+            intermediate_path = Path(self.intermediate_path)
+            intermediate_path.mkdir(parents=True, exist_ok=True)
+            file_handler = logging.FileHandler(
+                intermediate_path / "partition_report.txt"
+            )
+            logger.addHandler(file_handler)
         logger.info(f"The following nodes were rejected for {self.tosa_spec}:")
         logger.info("\n" + reporter.get_table_report())
         logger.info("(Placeholders and outputs are not included in this list)")

backends/arm/vgf/partitioner.py

@@ -34,3 +34,4 @@ def __init__(
         self.additional_checks = additional_checks
         self.tosa_spec = compile_spec.tosa_spec
         self._custom_partition_ops: set[torch._ops.OpOverload] = set()
+        self.intermediate_path = compile_spec._get_intermediate_path()