Migrate pt2e arm (#11053)

Migrate arm backend to use pt2e from torchao

Author: metascroy

.lintrunner.toml

@@ -386,8 +386,6 @@ exclude_patterns = [
     "third-party/**",
     # TODO: remove exceptions as we migrate
     # backends
-    "backends/arm/quantizer/**",
-    "backends/arm/test/ops/**",
     "backends/vulkan/quantizer/**",
     "backends/vulkan/test/**",
     "backends/qualcomm/quantizer/**",

backends/arm/quantizer/TARGETS

@@ -6,6 +6,7 @@ python_library(
     srcs = ["quantization_config.py"],
     deps = [
         "//caffe2:torch",
+        "//pytorch/ao:torchao",
     ],
 )
 
@@ -18,6 +19,7 @@ python_library(
         ":quantization_annotator",
         "//caffe2:torch",
         "//executorch/exir:lib",
+        "//pytorch/ao:torchao",
     ],
 )
 
@@ -28,6 +30,7 @@ python_library(
         ":arm_quantizer_utils",
         ":quantization_config",
         "//caffe2:torch",
+        "//pytorch/ao:torchao",
     ],
 )
 
@@ -36,6 +39,7 @@ python_library(
     srcs = ["arm_quantizer_utils.py"],
     deps = [
         ":quantization_config",
+        "//pytorch/ao:torchao",
     ],
 )
 

backends/arm/quantizer/arm_quantizer.py

@@ -30,25 +30,26 @@
     is_vgf,
 )  # usort: skip
 from executorch.exir.backend.compile_spec_schema import CompileSpec
-from torch.ao.quantization.fake_quantize import (
+
+from torch.fx import GraphModule, Node
+from torchao.quantization.pt2e import (
     FakeQuantize,
     FusedMovingAvgObsFakeQuantize,
-)
-from torch.ao.quantization.observer import (
     HistogramObserver,
     MinMaxObserver,
     MovingAverageMinMaxObserver,
     MovingAveragePerChannelMinMaxObserver,
+    ObserverOrFakeQuantizeConstructor,
     PerChannelMinMaxObserver,
     PlaceholderObserver,
 )
-from torch.ao.quantization.qconfig import _ObserverOrFakeQuantizeConstructor
-from torch.ao.quantization.quantizer import QuantizationSpec, Quantizer
-from torch.ao.quantization.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
+
+from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec,
+    QuantizationSpec,
+    Quantizer,
 )
-from torch.fx import GraphModule, Node
 
 __all__ = [
     "TOSAQuantizer",
@@ -97,7 +98,7 @@ def get_symmetric_quantization_config(
     weight_qscheme = (
         torch.per_channel_symmetric if is_per_channel else torch.per_tensor_symmetric
     )
-    weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = (
+    weight_observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor = (
         MinMaxObserver
     )
     if is_qat:
@@ -337,14 +338,14 @@ def _annotate_io(
             if is_annotated(node):
                 continue
             if node.op == "placeholder" and len(node.users) > 0:
-                _annotate_output_qspec(
+                annotate_output_qspec(
                     node,
                     quantization_config.get_output_act_qspec(),
                 )
                 mark_node_as_annotated(node)
             if node.op == "output":
                 parent = node.all_input_nodes[0]
-                _annotate_input_qspec_map(
+                annotate_input_qspec_map(
                     node, parent, quantization_config.get_input_act_qspec()
                 )
                 mark_node_as_annotated(node)

backends/arm/quantizer/arm_quantizer_utils.py

@@ -15,10 +15,10 @@
 
 import torch
 from torch._subclasses import FakeTensor
-
-from torch.ao.quantization.quantizer import QuantizationAnnotation
 from torch.fx import GraphModule, Node
 
+from torchao.quantization.pt2e.quantizer import QuantizationAnnotation
+
 
 def is_annotated(node: Node) -> bool:
     """Given a node return whether the node is annotated."""

backends/arm/quantizer/quantization_annotator.py

@@ -13,12 +13,14 @@
 import torch.nn.functional as F
 from executorch.backends.arm.quantizer import 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 (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
-)
+
 from torch.fx import Node
+from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec,
+    QuantizationSpecBase,
+    SharedQuantizationSpec,
+)
 
 from .arm_quantizer_utils import (
     is_annotated,
@@ -119,7 +121,7 @@ def _annotate_input(node: Node, quant_property: _QuantProperty):
         strict=True,
     ):
         assert isinstance(n_arg, Node)
-        _annotate_input_qspec_map(node, n_arg, qspec)
+        annotate_input_qspec_map(node, n_arg, qspec)
         if quant_property.mark_annotated:
             mark_node_as_annotated(n_arg)  # type: ignore[attr-defined]
 
@@ -130,7 +132,7 @@ def _annotate_output(node: Node, quant_property: _QuantProperty):
     assert not quant_property.optional
     assert quant_property.index == 0, "Only one output annotation supported currently"
 
-    _annotate_output_qspec(node, quant_property.qspec)
+    annotate_output_qspec(node, quant_property.qspec)
 
 
 def _match_pattern(

backends/arm/quantizer/quantization_config.py

@@ -9,9 +9,9 @@
 from dataclasses import dataclass
 
 import torch
-from torch.ao.quantization import ObserverOrFakeQuantize
+from torchao.quantization.pt2e import ObserverOrFakeQuantize
 
-from torch.ao.quantization.quantizer import (
+from torchao.quantization.pt2e.quantizer import (
     DerivedQuantizationSpec,
     FixedQParamsQuantizationSpec,
     QuantizationSpec,

backends/arm/test/ops/test_add.py

@@ -19,8 +19,8 @@
 )
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.xnnpack.test.tester import Quantize
-from torch.ao.quantization.observer import HistogramObserver
-from torch.ao.quantization.quantizer import QuantizationSpec
+from torchao.quantization.pt2e import HistogramObserver
+from torchao.quantization.pt2e.quantizer import QuantizationSpec
 
 aten_op = "torch.ops.aten.add.Tensor"
 exir_op = "executorch_exir_dialects_edge__ops_aten_add_Tensor"

backends/arm/test/ops/test_sigmoid_16bit.py

@@ -18,8 +18,8 @@
 )
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.xnnpack.test.tester import Quantize
-from torch.ao.quantization.observer import HistogramObserver
-from torch.ao.quantization.quantizer import QuantizationSpec
+from torchao.quantization.pt2e import HistogramObserver
+from torchao.quantization.pt2e.quantizer import QuantizationSpec
 
 
 def _get_16_bit_quant_config():

backends/arm/test/ops/test_sigmoid_32bit.py

@@ -14,8 +14,8 @@
 )
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.xnnpack.test.tester import Quantize
-from torch.ao.quantization.observer import HistogramObserver
-from torch.ao.quantization.quantizer import QuantizationSpec
+from torchao.quantization.pt2e import HistogramObserver
+from torchao.quantization.pt2e.quantizer import QuantizationSpec
 
 
 def _get_16_bit_quant_config():