Ethos-u: update quantizer/__init__

Summary:
Expose Quantizer and quant config through __init__.py

Relanding Oscar's #10351 (but in parts)

Differential Revision: D73530601

Author: digantdesai

backends/arm/quantizer/TARGETS

@@ -1,5 +1,15 @@
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 
+# Exposed through __init__.py
+python_library(
+    name = "quantization_config",
+    srcs = ["quantization_config.py"],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
+# Exposed through __init__.py
 python_library(
     name = "arm_quantizer",
     srcs = ["arm_quantizer.py"],
@@ -22,17 +32,18 @@ python_library(
 )
 
 python_library(
-    name = "quantization_config",
-    srcs = ["quantization_config.py"],
+    name = "arm_quantizer_utils",
+    srcs = ["arm_quantizer_utils.py"],
     deps = [
-        "//caffe2:torch",
+        ":quantization_config",
     ],
 )
 
 python_library(
-    name = "arm_quantizer_utils",
-    srcs = ["arm_quantizer_utils.py"],
+    name = "lib",
+    srcs = ["__init__.py"],
     deps = [
+        ":arm_quantizer",
         ":quantization_config",
-    ],
+    ]
 )

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 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/TARGETS

@@ -42,7 +42,7 @@ python_library(
         ":common",
         "//executorch/backends/xnnpack/test/tester:tester",
         "//executorch/backends/arm:arm_partitioner",
-        "//executorch/backends/arm/quantizer:arm_quantizer",
+        "//executorch/backends/arm/quantizer:lib",
         "//executorch/backends/arm:tosa_mapping",
         "//executorch/devtools/backend_debug:delegation_info",
         "fbsource//third-party/pypi/tabulate:tabulate",

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

@@ -27,7 +27,7 @@
     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,

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_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

@@ -24,7 +24,7 @@
     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,

examples/arm/ethos_u_minimal_example.ipynb

@@ -80,7 +80,7 @@
    "outputs": [],
    "source": [
     "from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder\n",
-    "from executorch.backends.arm.quantizer.arm_quantizer import (\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",
@@ -170,9 +170,9 @@
     "    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",