Arm backend: Populate __init__.py for quantizer and Arm root (#10351)

Populate backends/arm/quantizer/__init__.py and backends/arm/__init__.py

Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/__init__.py

@@ -0,0 +1,10 @@
+# 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 .arm_backend import ArmCompileSpecBuilder  # noqa  # usort: skip
+from .tosa_backend import TOSABackend  # noqa  # usort: skip
+from .tosa_partitioner import TOSAPartitioner  # noqa  # usort: skip
+from .ethosu_backend import EthosUBackend  # noqa  # usort: skip
+from .ethosu_partitioner import EthosUPartitioner  # noqa  # usort: skip

backends/arm/ethosu_backend.py

@@ -14,9 +14,9 @@
 import logging
 from typing import final, List
 
-from executorch.backends.arm.arm_vela import vela_compile
+from executorch.backends.arm import TOSABackend
 
-from executorch.backends.arm.tosa_backend import TOSABackend
+from executorch.backends.arm.arm_vela import vela_compile
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.export.exported_program import ExportedProgram

backends/arm/ethosu_partitioner.py

@@ -10,8 +10,7 @@
 from executorch.backends.arm.arm_backend import (
     is_ethosu,
 )  # usort: skip
-from executorch.backends.arm.ethosu_backend import EthosUBackend
-from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
+from executorch.backends.arm import EthosUBackend, TOSAPartitioner
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import DelegationSpec
 from torch.fx.passes.operator_support import OperatorSupportBase

backends/arm/quantizer/__init__.py

@@ -1,4 +1,12 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-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 .quantization_config import QuantizationConfig  # noqa  # usort: skip
+from .arm_quantizer import (  # noqa
+    EthosUQuantizer,
+    get_symmetric_quantization_config,
+    TOSAQuantizer,
+)

backends/arm/quantizer/arm_quantizer.py

@@ -19,15 +19,14 @@
 import torch
 from executorch.backends.arm._passes import ArmPassManager
 
-from executorch.backends.arm.quantizer import arm_quantizer_utils
+from executorch.backends.arm.quantizer import arm_quantizer_utils, QuantizationConfig
 from executorch.backends.arm.quantizer.arm_quantizer_utils import (  # type: ignore[attr-defined]
     mark_node_as_annotated,
 )
 from executorch.backends.arm.quantizer.quantization_annotator import (  # type: ignore[import-not-found]
     annotate_graph,
 )
 
-from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.arm.arm_backend import (
     get_tosa_spec,

backends/arm/quantizer/quantization_annotator.py

@@ -10,8 +10,7 @@
 
 import torch
 import torch.fx
-from executorch.backends.arm.quantizer import arm_quantizer_utils
-from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
+from executorch.backends.arm.quantizer import arm_quantizer_utils, QuantizationConfig
 from executorch.backends.arm.tosa_utils import get_node_debug_info
 from torch.ao.quantization.quantizer import QuantizationSpecBase, SharedQuantizationSpec
 from torch.ao.quantization.quantizer.utils import (

backends/arm/test/ops/test_expand.py

@@ -15,7 +15,7 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_hardtanh.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_max_pool.py

@@ -12,7 +12,7 @@
 import pytest
 
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_permute.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_relu.py

@@ -10,7 +10,7 @@
 from typing import Tuple
 
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_repeat.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_sigmoid_16bit.py

@@ -6,7 +6,7 @@
 import pytest
 
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     get_symmetric_quantization_config,
     TOSAQuantizer,
 )

backends/arm/test/ops/test_sigmoid_32bit.py

@@ -5,7 +5,7 @@
 
 import pytest
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import TOSAQuantizer
+from executorch.backends.arm.quantizer import TOSAQuantizer
 from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (

backends/arm/test/ops/test_var.py

@@ -11,7 +11,7 @@
 import unittest
 
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/ops/test_where.py

@@ -9,7 +9,7 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/test/tester/arm_tester.py

@@ -18,6 +18,7 @@
 import torch.utils._pytree as pytree
 
 import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore[import-untyped]
+from executorch.backends.arm import EthosUPartitioner, TOSAPartitioner
 from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
 
 from executorch.backends.arm.arm_backend import (
@@ -26,8 +27,7 @@
     is_ethosu,
     is_tosa,
 )
-from executorch.backends.arm.ethosu_partitioner import EthosUPartitioner
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,
@@ -47,7 +47,6 @@
     print_error_diffs,
 )
 from executorch.backends.arm.tosa_mapping import extract_tensor_meta
-from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
 
 from executorch.backends.xnnpack.test.tester import Tester
 from executorch.devtools.backend_debug import get_delegation_info

backends/arm/test/tester/test_pipeline.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,

backends/arm/tosa_partitioner.py

@@ -13,11 +13,11 @@
     get_tosa_spec,
     is_tosa,
 )  # usort: skip
+from executorch.backends.arm import TOSABackend
 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.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import (
     DelegationSpec,

backends/arm/tosa_quant_utils.py

@@ -10,8 +10,6 @@
 import math
 from typing import cast, List, NamedTuple, Tuple
 
-import executorch.backends.arm.tosa_mapping
-
 import torch.fx
 import torch.fx.node
 
@@ -234,7 +232,7 @@ def build_rescale(
 
 def build_rescale_to_int32(
     tosa_fb: ts.TosaSerializer,
-    input_arg: executorch.backends.arm.tosa_mapping.TosaArg,
+    input_arg: TosaArg,
     input_zp: int,
     rescale_scale: list[float],
     is_scale32: bool = True,

examples/arm/aot_arm_compiler.py

@@ -17,19 +17,17 @@
 
 import torch
 from examples.devtools.scripts.export_bundled_program import save_bundled_program
-from executorch.backends.arm.arm_backend import (
+from executorch.backends.arm import (
     ArmCompileSpecBuilder,
-    get_tosa_spec,
-    is_ethosu,
-    is_tosa,
+    EthosUPartitioner,
+    TOSAPartitioner,
 )
-from executorch.backends.arm.ethosu_partitioner import EthosUPartitioner
-from executorch.backends.arm.quantizer.arm_quantizer import (
+from executorch.backends.arm.arm_backend import get_tosa_spec, is_ethosu, is_tosa
+from executorch.backends.arm.quantizer import (
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,
 )
-from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
 from executorch.backends.arm.util.arm_model_evaluator import (

examples/arm/ethos_u_minimal_example.ipynb

@@ -79,8 +79,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder\n",
-    "from executorch.backends.arm.quantizer.arm_quantizer import (\n",
+    "from executorch.backends.arm import ArmCompileSpecBuilder\n",
+    "from executorch.backends.arm.quantizer import (\n",
     "    EthosUQuantizer,\n",
     "    get_symmetric_quantization_config,\n",
     ")\n",
@@ -89,7 +89,7 @@
     "target = \"ethos-u55-128\"\n",
     "\n",
     "# Create a compilation spec describing the target for configuring the quantizer\n",
-    "# Some args are used by the Arm Vela graph compiler later in the example. Refer to Arm Vela documentation for an \n",
+    "# Some args are used by the Arm Vela graph compiler later in the example. Refer to Arm Vela documentation for an\n",
     "# explanation of its flags: https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela/-/blob/main/OPTIONS.md\n",
     "spec_builder = ArmCompileSpecBuilder().ethosu_compile_spec(\n",
     "            target,\n",
@@ -100,12 +100,12 @@
     "compile_spec = spec_builder.build()\n",
     "\n",
     "# Create and configure quantizer to use a symmetric quantization config globally on all nodes\n",
-    "quantizer = EthosUQuantizer(compile_spec) \n",
+    "quantizer = EthosUQuantizer(compile_spec)\n",
     "operator_config = get_symmetric_quantization_config(is_per_channel=False)\n",
     "quantizer.set_global(operator_config)\n",
     "\n",
     "# Post training quantization\n",
-    "quantized_graph_module = prepare_pt2e(graph_module, quantizer) \n",
+    "quantized_graph_module = prepare_pt2e(graph_module, quantizer)\n",
     "quantized_graph_module(*example_inputs) # Calibrate the graph module with the example input\n",
     "quantized_graph_module = convert_pt2e(quantized_graph_module)\n",
     "\n",
@@ -128,8 +128,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import subprocess \n",
-    "import os \n",
+    "import subprocess\n",
+    "import os\n",
     "\n",
     "# Setup paths\n",
     "cwd_dir = os.getcwd()\n",
@@ -163,16 +163,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from executorch.backends.arm.ethosu_partitioner import EthosUPartitioner\n",
+    "from executorch.backends.arm import EthosUPartitioner\n",
     "from executorch.exir import (\n",
     "    EdgeCompileConfig,\n",
     "    ExecutorchBackendConfig,\n",
     "    to_edge_transform_and_lower,\n",
     ")\n",
     "from executorch.extension.export_util.utils import save_pte_program\n",
-    "import platform \n",
+    "import platform\n",
     "\n",
-    "# Create partitioner from compile spec \n",
+    "# Create partitioner from compile spec\n",
     "partitioner = EthosUPartitioner(compile_spec)\n",
     "\n",
     "# Lower the exported program to the Ethos-U backend\n",
@@ -185,8 +185,8 @@
     "        )\n",
     "\n",
     "# Load quantization ops library\n",
-    "os_aot_lib_names = {\"Darwin\" : \"libquantized_ops_aot_lib.dylib\", \n",
-    "                \"Linux\"  : \"libquantized_ops_aot_lib.so\", \n",
+    "os_aot_lib_names = {\"Darwin\" : \"libquantized_ops_aot_lib.dylib\",\n",
+    "                \"Linux\"  : \"libquantized_ops_aot_lib.so\",\n",
     "                \"Windows\": \"libquantized_ops_aot_lib.dll\"}\n",
     "aot_lib_name = os_aot_lib_names[platform.system()]\n",
     "\n",
@@ -226,7 +226,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Build executorch \n",
+    "# Build executorch\n",
     "subprocess.run(os.path.join(script_dir, \"build_executorch.sh\"), shell=True, cwd=et_dir)\n",
     "\n",
     "# Build portable kernels\n",