Arm backend: Move get_tosa_spec to tosa_specification.py (pytorch#13281)

Signed-off-by: Sebastian Larsson <[email protected]>

Author: Sebastian-Larsson

backends/arm/arm_backend.py

@@ -217,13 +217,6 @@ def is_vgf(compile_spec: List[CompileSpec]) -> bool:
     return False
 
 
-def get_tosa_spec(compile_spec: List[CompileSpec]) -> TosaSpecification:
-    for spec in compile_spec:
-        if spec.key == "tosa_spec":
-            return TosaSpecification.create_from_string(spec.value.decode())
-    raise ValueError("Could not find TOSA version in CompileSpec")
-
-
 def get_intermediate_path(compile_spec: List[CompileSpec]) -> Optional[str]:
     for spec in compile_spec:
         if spec.key == "debug_artifact_path":

backends/arm/quantizer/arm_quantizer.py

@@ -20,12 +20,11 @@
 from executorch.backends.arm._passes import ArmPassManager
 
 from executorch.backends.arm.quantizer import QuantizationConfig
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import get_tosa_spec, TosaSpecification
 
 from .arm_quantizer_utils import is_annotated, mark_node_as_annotated
 from .quantization_annotator import annotate_graph
 from executorch.backends.arm.arm_backend import (
-    get_tosa_spec,
     is_ethosu,
     is_vgf,
 )  # usort: skip

backends/arm/test/misc/test_tosa_spec.py

@@ -5,9 +5,11 @@
 
 import unittest
 
-from executorch.backends.arm.arm_backend import get_tosa_spec
-
-from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification
+from executorch.backends.arm.tosa_specification import (
+    get_tosa_spec,
+    Tosa_1_00,
+    TosaSpecification,
+)
 
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from parameterized import parameterized  # type: ignore[import-untyped]

backends/arm/test/ops/test_add.py

@@ -8,7 +8,6 @@
 from typing import Tuple
 
 import torch
-from executorch.backends.arm.arm_backend import get_tosa_spec
 from executorch.backends.arm.quantizer import arm_quantizer
 from executorch.backends.arm.test import common, conftest
 from executorch.backends.arm.test.tester.test_pipeline import (
@@ -18,7 +17,7 @@
     TosaPipelineINT,
     VgfPipeline,
 )
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import get_tosa_spec, TosaSpecification
 from executorch.backends.xnnpack.test.tester import Quantize
 from torchao.quantization.pt2e import HistogramObserver
 from torchao.quantization.pt2e.quantizer import QuantizationSpec

backends/arm/test/runner_utils.py

@@ -18,10 +18,13 @@
 import numpy as np
 import torch
 
-from executorch.backends.arm.arm_backend import get_tosa_spec, is_tosa
+from executorch.backends.arm.arm_backend import is_tosa
 from executorch.backends.arm.test.conftest import is_option_enabled
-from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification
-
+from executorch.backends.arm.tosa_specification import (
+    get_tosa_spec,
+    Tosa_1_00,
+    TosaSpecification,
+)
 from executorch.exir import ExecutorchProgramManager, ExportedProgram
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.lowered_backend_module import LoweredBackendModule

backends/arm/test/tester/arm_tester.py

@@ -34,7 +34,6 @@
 
 from executorch.backends.arm.arm_backend import (
     get_intermediate_path,
-    get_tosa_spec,
     is_ethosu,
     is_tosa,
     is_vgf,
@@ -62,7 +61,7 @@
 )
 from executorch.backends.arm.tosa_mapping import extract_tensor_meta
 from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import get_tosa_spec, TosaSpecification
 
 from executorch.backends.arm.vgf_partitioner import VgfPartitioner
 

backends/arm/tosa_backend.py

@@ -14,8 +14,8 @@
 from typing import cast, final, List
 
 import serializer.tosa_serializer as ts  # type: ignore
-from executorch.backends.arm.arm_backend import get_tosa_spec
 from executorch.backends.arm.operators.node_visitor import get_node_visitors
+from executorch.backends.arm.tosa_specification import get_tosa_spec
 from executorch.backends.arm._passes import (
     ArmPassManager,
 )  # usort: skip

backends/arm/tosa_partitioner.py

@@ -11,14 +11,14 @@
 import torch
 from executorch.backends.arm.constants import DQ_OPS, Q_OPS
 from executorch.backends.arm.arm_backend import (
-    get_tosa_spec,
     is_tosa,
 )  # usort: skip
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     tosa_support_factory,
 )
 from executorch.backends.arm.tosa_backend import TOSABackend
+from executorch.backends.arm.tosa_specification import get_tosa_spec
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import (
     DelegationSpec,

backends/arm/tosa_specification.py

@@ -15,6 +15,10 @@
 import re
 from typing import List
 
+from executorch.exir.backend.compile_spec_schema import (  # type: ignore[import-not-found]
+    CompileSpec,
+)
+
 from packaging.version import Version
 
 
@@ -188,3 +192,10 @@ def get_context_spec() -> TosaSpecification:
         return TosaLoweringContext.tosa_spec_var.get()
     except LookupError:
         raise RuntimeError("Function must be executed within a TosaLoweringContext")
+
+
+def get_tosa_spec(compile_spec: List[CompileSpec]) -> TosaSpecification:
+    for spec in compile_spec:
+        if spec.key == "tosa_spec":
+            return TosaSpecification.create_from_string(spec.value.decode())
+    raise ValueError("Could not find TOSA version in CompileSpec")

examples/arm/aot_arm_compiler.py

@@ -19,7 +19,6 @@
 from examples.devtools.scripts.export_bundled_program import save_bundled_program
 from executorch.backends.arm.arm_backend import (
     ArmCompileSpecBuilder,
-    get_tosa_spec,
     is_ethosu,
     is_tosa,
     is_vgf,
@@ -32,7 +31,7 @@
     VgfQuantizer,
 )
 from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import get_tosa_spec, TosaSpecification
 
 from executorch.backends.arm.util.arm_model_evaluator import (
     GenericModelEvaluator,