Arm backend: Add complie spec factories (pytorch#14376)

Signed-off-by: Erik Lundell <[email protected]>
Co-authored-by: Digant Desai <[email protected]>

Author: Erik-Lundell

backends/arm/TARGETS

@@ -106,3 +106,17 @@ runtime.python_library(
         "//caffe2:torch",
     ]
 )
+runtime.python_library(
+    name = "_factory",
+    srcs = [
+        "util/_factory.py"
+    ],
+    deps = [
+        ":ethosu",
+        ":vgf",
+        ":arm_compile_spec",
+        "//executorch/backends/arm/quantizer:lib",
+        "//executorch/exir/backend:operator_support",
+        "//executorch/exir/backend:compile_spec_schema",
+    ]
+)

backends/arm/test/TARGETS

@@ -1,3 +1,8 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load(":targets.bzl", "define_arm_tests")
 
@@ -58,6 +63,7 @@ runtime.python_library(
         "//executorch/backends/arm/quantizer:lib",
         "//executorch/backends/arm/tosa:mapping",
         "//executorch/backends/arm:vgf",
+        "//executorch/backends/arm:_factory",
         "//executorch/devtools/backend_debug:delegation_info",
         "//executorch/exir/backend:operator_support",
         "fbsource//third-party/pypi/tabulate:tabulate",

backends/arm/test/common.py

@@ -14,6 +14,7 @@
 
 import pytest
 from executorch.backends.arm.ethosu import EthosUCompileSpec
+
 from executorch.backends.arm.test.runner_utils import (
     arm_executor_runner_exists,
     corstone300_installed,

backends/arm/test/tester/arm_tester.py

@@ -28,17 +28,11 @@
 
 import torch.fx
 import torch.utils._pytree as pytree
-
 from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
 
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
-from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner
-from executorch.backends.arm.quantizer import (
-    EthosUQuantizer,
-    get_symmetric_quantization_config,
-    TOSAQuantizer,
-    VgfQuantizer,
-)
+from executorch.backends.arm.ethosu import EthosUCompileSpec
+from executorch.backends.arm.quantizer import get_symmetric_quantization_config
 from executorch.backends.arm.test.runner_utils import (
     dbg_tosa_fb_to_json,
     get_output_quantization_params,
@@ -53,9 +47,13 @@
 from executorch.backends.arm.tosa import TosaSpecification
 from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec
 from executorch.backends.arm.tosa.mapping import extract_tensor_meta
-from executorch.backends.arm.tosa.partitioner import TOSAPartitioner
 
-from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner
+from executorch.backends.arm.util._factory import (
+    create_partitioner,
+    create_quantizer,
+    parse_compile_spec,
+)
+from executorch.backends.arm.vgf import VgfCompileSpec
 
 from executorch.backends.test.harness.error_statistics import ErrorStatistics
 from executorch.backends.test.harness.stages import Stage, StageType
@@ -83,7 +81,6 @@
     _copy_module,
     _update_exported_program_graph_module,
 )
-
 from tabulate import tabulate
 
 from torch.export.graph_signature import ExportGraphSignature, InputSpec, OutputSpec
@@ -103,26 +100,20 @@ def _dump_lowered_modules_artifact(
         artifact.exported_program().graph_signature
     )
 
-    def get_output_format(lowered_module) -> str | None:
-        for spec in lowered_module.compile_specs:
-            if spec.key == "output_format":
-                return spec.value.decode()
-        return None
-
     for node in graph_module.graph.nodes:
         if node.op == "get_attr" and node.name.startswith("lowered_module_"):
             lowered_module = getattr(graph_module, node.name)
             assert isinstance(
                 lowered_module, LoweredBackendModule
             ), f"Attribute {node.name} must be of type LoweredBackendModule."
 
-            output_format = get_output_format(lowered_module)
-            if output_format == "tosa":
+            compile_spec = parse_compile_spec(lowered_module.compile_specs)
+            if isinstance(compile_spec, TosaCompileSpec):
                 tosa_fb = lowered_module.processed_bytes
                 to_print = dbg_tosa_fb_to_json(tosa_fb)
                 to_print = pformat(to_print, compact=True, indent=1)
                 output += f"\nTOSA deserialized {node.name}: \n{to_print}\n"
-            elif output_format == EthosUCompileSpec.get_output_format():
+            elif isinstance(compile_spec, EthosUCompileSpec):
                 vela_cmd_stream = lowered_module.processed_bytes
                 output += f"\nVela command stream {node.name}: \n{vela_cmd_stream}\n"
             else:
@@ -284,13 +275,7 @@ def quantize(
         quantize_stage: Optional[tester.Quantize] = None,
     ):
         if quantize_stage is None:
-            quantizer = None
-            if isinstance(self.compile_spec, TosaCompileSpec):
-                quantizer = TOSAQuantizer(self.compile_spec)
-            elif isinstance(self.compile_spec, EthosUCompileSpec):
-                quantizer = EthosUQuantizer(self.compile_spec)
-            elif isinstance(self.compile_spec, VgfCompileSpec):
-                quantizer = VgfQuantizer(self.compile_spec)
+            quantizer = create_quantizer(self.compile_spec)
             quantize_stage = tester.Quantize(
                 quantizer,
                 get_symmetric_quantization_config(),
@@ -312,14 +297,7 @@ def to_edge(
 
     def partition(self, partition_stage: Optional[Partition] = None):
         if partition_stage is None:
-            if isinstance(self.compile_spec, TosaCompileSpec):
-                arm_partitioner = TOSAPartitioner(self.compile_spec)
-            elif isinstance(self.compile_spec, EthosUCompileSpec):
-                arm_partitioner = EthosUPartitioner(self.compile_spec)
-            elif isinstance(self.compile_spec, VgfCompileSpec):
-                arm_partitioner = VgfPartitioner(self.compile_spec)
-            else:
-                raise ValueError("compile spec doesn't target any Arm Partitioner")
+            arm_partitioner = create_partitioner(self.compile_spec)
             partition_stage = Partition(arm_partitioner)
         return super().partition(partition_stage)
 
@@ -329,7 +307,7 @@ def to_edge_transform_and_lower(
         partitioners: Optional[List[Partitioner]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
         additional_checks: Optional[
-            List[Union[DontPartition | DontPartitionModule | DontPartitionName]]
+            List[DontPartition | DontPartitionModule | DontPartitionName]
         ] = None,
         transform_passes: Optional[
             Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
@@ -343,20 +321,9 @@ def to_edge_transform_and_lower(
 
         if to_edge_and_lower_stage is None:
             if partitioners is None:
-                if isinstance(self.compile_spec, TosaCompileSpec):
-                    arm_partitioner = TOSAPartitioner(
-                        self.compile_spec, additional_checks
-                    )
-                elif isinstance(self.compile_spec, EthosUCompileSpec):
-                    arm_partitioner = EthosUPartitioner(
-                        self.compile_spec, additional_checks
-                    )
-                elif isinstance(self.compile_spec, VgfCompileSpec):
-                    arm_partitioner = VgfPartitioner(
-                        self.compile_spec, additional_checks
-                    )
-                else:
-                    raise ValueError("compile spec doesn't target any Arm Partitioner")
+                arm_partitioner = create_partitioner(
+                    self.compile_spec, additional_checks
+                )
                 partitioners = [arm_partitioner]
             to_edge_and_lower_stage = ToEdgeTransformAndLower(
                 partitioners,
@@ -743,22 +710,19 @@ def _get_tosa_operator_distribution(
     op_list = []
     id = 0
     while lowered_module := getattr(graph_module, f"lowered_module_{id}", None):
-        for spec in lowered_module.compile_specs:
-            if spec.key != "output_format":
-                continue
-            if spec.value == b"tosa":
-                tosa_fb = lowered_module.processed_bytes
-                tosa_json = dbg_tosa_fb_to_json(tosa_fb)
-                for region in tosa_json["regions"]:
-                    for block in region["blocks"]:
-                        op_list.extend(
-                            [operator["op"] for operator in block["operators"]]
-                        )
-                break
-            elif spec.value == EthosUCompileSpec.get_output_format().encode():
-                return "Can not get operator distribution for Vela command stream."
-            else:
-                return f"Unknown output format '{spec.value}'."
+        compile_spec = parse_compile_spec(lowered_module.compile_specs)
+        if isinstance(compile_spec, TosaCompileSpec):
+            tosa_fb = lowered_module.processed_bytes
+            tosa_json = dbg_tosa_fb_to_json(tosa_fb)
+            for region in tosa_json["regions"]:
+                for block in region["blocks"]:
+                    op_list.extend([operator["op"] for operator in block["operators"]])
+        elif isinstance(compile_spec, EthosUCompileSpec):
+            return "Can not get operator distribution for Vela command stream."
+        elif isinstance(compile_spec, VgfCompileSpec):
+            return "Can not get operator distribution for VGF."
+        else:
+            return f"Unknown output format '{compile_spec.get_output_format()}'."
         id += 1
     if id == 0:
         return "No delegate with name 'lowered_module_0 found in graph module."

backends/arm/tosa/backend.py

@@ -206,8 +206,8 @@ def filter_tosa_compile_specs(
         hardware.
         """
 
-        new_compile_spec = TosaCompileSpec.__new__(TosaCompileSpec)
-        new_compile_spec._set_compile_specs(
-            compile_spec.tosa_spec, [], compile_spec.get_intermediate_path()
+        return (
+            TosaCompileSpec(compile_spec.tosa_spec)
+            .dump_intermediate_artifacts_to(compile_spec.get_intermediate_path())
+            .dump_debug_info(compile_spec.tosa_debug_mode)
         )
-        return new_compile_spec

backends/arm/util/_factory.py

@@ -0,0 +1,59 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
+from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner
+from executorch.backends.arm.quantizer import (
+    EthosUQuantizer,
+    TOSAQuantizer,
+    VgfQuantizer,
+)
+from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec
+from executorch.backends.arm.tosa.partitioner import TOSAPartitioner
+from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from torch.fx.passes.operator_support import OperatorSupportBase
+
+
+def parse_compile_spec(compile_specs: list[CompileSpec]) -> ArmCompileSpec:
+    output_format = None
+    for spec in compile_specs:
+        if spec.key == "output_format":
+            output_format = spec.value.decode()
+            break
+    else:
+        raise ValueError("Compile spec without output format.")
+    if output_format == TosaCompileSpec.get_output_format():
+        return TosaCompileSpec.from_list(compile_specs)
+    if output_format == EthosUCompileSpec.get_output_format():
+        return EthosUCompileSpec.from_list(compile_specs)
+    if output_format == VgfCompileSpec.get_output_format():
+        return VgfCompileSpec.from_list(compile_specs)
+    raise ValueError(f"Unknown output format {output_format}")
+
+
+def create_partitioner(
+    compile_spec: ArmCompileSpec,
+    additional_checks: list[OperatorSupportBase] | None = None,
+):
+    if isinstance(compile_spec, TosaCompileSpec):
+        return TOSAPartitioner(compile_spec, additional_checks)
+    elif isinstance(compile_spec, EthosUCompileSpec):
+        return EthosUPartitioner(compile_spec, additional_checks)
+    elif isinstance(compile_spec, VgfCompileSpec):
+        return VgfPartitioner(compile_spec, additional_checks)
+    else:
+        raise ValueError("compile spec doesn't target any Arm Partitioner")
+
+
+def create_quantizer(compile_spec: ArmCompileSpec):
+    if isinstance(compile_spec, TosaCompileSpec):
+        return TOSAQuantizer(compile_spec)
+    elif isinstance(compile_spec, EthosUCompileSpec):
+        return EthosUQuantizer(compile_spec)
+    elif isinstance(compile_spec, VgfCompileSpec):
+        return VgfQuantizer(compile_spec)
+    else:
+        raise ValueError("compile spec doesn't target any Arm Quantizer")

examples/arm/aot_arm_compiler.py

@@ -18,23 +18,18 @@
 import torch
 from examples.devtools.scripts.export_bundled_program import save_bundled_program
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
-from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner
-from executorch.backends.arm.quantizer import (
-    EthosUQuantizer,
-    get_symmetric_quantization_config,
-    TOSAQuantizer,
-    VgfQuantizer,
-)
+from executorch.backends.arm.ethosu import EthosUCompileSpec
+from executorch.backends.arm.quantizer import get_symmetric_quantization_config
 from executorch.backends.arm.tosa import TosaSpecification
 from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec
-from executorch.backends.arm.tosa.partitioner import TOSAPartitioner
+from executorch.backends.arm.util._factory import create_partitioner, create_quantizer
 
 from executorch.backends.arm.util.arm_model_evaluator import (
     evaluate_model,
     evaluator_calibration_data,
 )
 
-from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner
+from executorch.backends.arm.vgf import VgfCompileSpec
 
 # To use Cortex-M backend
 from executorch.backends.cortex_m.passes.quantized_linear_fusion_pass import (
@@ -158,15 +153,8 @@ def quantize(
     export"""
     logging.info("Quantizing Model...")
     logging.debug(f"Original model: {model}")
-    quantizer = None
-    if isinstance(compile_specs, EthosUCompileSpec):
-        quantizer = EthosUQuantizer(compile_specs)
-    elif isinstance(compile_specs, TosaCompileSpec):
-        quantizer = TOSAQuantizer(compile_specs)
-    elif isinstance(compile_specs, VgfCompileSpec):
-        quantizer = VgfQuantizer(compile_specs)
-    else:
-        raise RuntimeError("Unsupported compilespecs for quantization!")
+
+    quantizer = create_quantizer(compile_specs)
 
     operator_config = get_symmetric_quantization_config()
     quantizer.set_global(operator_config)
@@ -649,14 +637,7 @@ def to_edge_TOSA_delegate(
             args, model, example_inputs, compile_spec
         )
 
-    if isinstance(compile_spec, EthosUCompileSpec):
-        partitioner = EthosUPartitioner(compile_spec)
-    elif isinstance(compile_spec, TosaCompileSpec):
-        partitioner = TOSAPartitioner(compile_spec)
-    elif isinstance(compile_spec, VgfCompileSpec):
-        partitioner = VgfPartitioner(compile_spec)
-    else:
-        raise RuntimeError(f"Unhandled compile spec: {compile_spec}")
+    partitioner = create_partitioner(compile_spec)
 
     edge = to_edge_transform_and_lower(
         exported_program,