diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index 4ce5911ec1f..cd7e8926292 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/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() diff --git a/backends/arm/test/misc/test_debug_feats.py b/backends/arm/test/misc/test_debug_feats.py index b032a5f9d63..2becccf4911 100644 --- a/backends/arm/test/misc/test_debug_feats.py +++ b/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) diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 0f4a3c6e550..97ce7f75248 100644 --- a/backends/arm/tosa/partitioner.py +++ b/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)") diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index eea49e8bc03..3810ba750ef 100644 --- a/backends/arm/vgf/partitioner.py +++ b/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()