Migrate pt2e backend batch1

Differential Revision: D75119434

Pull Request resolved: #11033

Author: metascroy

.lintrunner.toml

@@ -384,6 +384,23 @@ code = "TORCH_AO_IMPORT"
 include_patterns = ["**/*.py"]
 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/cadence/aot/quantizer/**",
+    "backends/qualcomm/quantizer/**",
+    "examples/qualcomm/**",
+    "backends/xnnpack/quantizer/**",
+    "backends/xnnpack/test/**",
+    "exir/tests/test_passes.py",
+    "extension/llm/export/builder.py",
+    "extension/llm/export/quantizer_lib.py",
+    "exir/tests/test_memory_planning.py",
+    "backends/transforms/duplicate_dynamic_quant_chain.py",
+    "exir/backend/test/demos/test_xnnpack_qnnpack.py",
 ]
 
 command = [
@@ -392,7 +409,7 @@ command = [
   "lintrunner_adapters",
   "run",
   "grep_linter",
-  "--pattern=\\bfrom torch\\.ao\\.quantization\\.(?:quantize_pt2e)(?:\\.[A-Za-z0-9_]+)*\\b",
+  "--pattern=\\bfrom torch\\.ao\\.quantization\\.(?:quantizer|observer|quantize_pt2e|pt2e)(?:\\.[A-Za-z0-9_]+)*\\b",
   "--linter-name=TorchAOImport",
   "--error-name=Prohibited torch.ao.quantization import",
   """--error-description=\

backends/cortex_m/test/test_replace_quant_nodes.py

@@ -16,15 +16,15 @@
     ReplaceQuantNodesPass,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
-from torch.ao.quantization.observer import HistogramObserver
-from torch.ao.quantization.quantizer.quantizer import (
+from torch.export import export, export_for_training
+from torch.fx import GraphModule
+from torchao.quantization.pt2e.observer import HistogramObserver
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+from torchao.quantization.pt2e.quantizer import (
     QuantizationAnnotation,
     QuantizationSpec,
     Quantizer,
 )
-from torch.export import export, export_for_training
-from torch.fx import GraphModule
-from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 
 @dataclass(eq=True, frozen=True)

backends/example/example_backend_delegate_passes/permute_memory_formats_pass.py

@@ -11,7 +11,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dim_order_utils import get_dim_order
 from executorch.exir.pass_base import ExportPass, PassResult
-from torch.ao.quantization.pt2e.graph_utils import find_sequential_partitions
+from torchao.quantization.pt2e import find_sequential_partitions
 
 
 class PermuteMemoryFormatsPass(ExportPass):

backends/example/example_operators/utils.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torch.ao.quantization.quantizer.quantizer import QuantizationAnnotation
+from torchao.quantization.pt2e.quantizer import QuantizationAnnotation
 
 
 def _nodes_are_annotated(node_list):

backends/example/example_partitioner.py

@@ -21,7 +21,7 @@
 from executorch.exir.graph_module import get_control_flow_submodules
 from torch.export import ExportedProgram
 from torch.fx.passes.operator_support import OperatorSupportBase
-from torchao.quantization.pt2e.graph_utils import find_sequential_partitions
+from torchao.quantization.pt2e import find_sequential_partitions
 
 
 @final

backends/mediatek/quantizer/annotator.py

@@ -10,18 +10,18 @@
 from torch._ops import OpOverload
 from torch._subclasses import FakeTensor
 
-from torch.ao.quantization.quantizer import QuantizationAnnotation
-from torch.ao.quantization.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
-)
-
 from torch.export import export_for_training
 from torch.fx import Graph, Node
 from torch.fx.passes.utils.matcher_with_name_node_map_utils import (
     SubgraphMatcherWithNameNodeMap,
 )
 
+from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec as _annotate_output_qspec,
+    QuantizationAnnotation,
+)
+
 from .qconfig import QuantizationConfig
 
 
@@ -108,7 +108,7 @@ def _annotate_fused_activation_pattern(
             torch.ops.aten.linear.default,
         ]:
             weight_node = producer_node.args[1]
-            _annotate_input_qspec_map(
+            annotate_input_qspec_map(
                 producer_node,
                 weight_node,
                 quant_config.weight,
@@ -201,7 +201,7 @@ def annotate_affine_ops(node: Node, quant_config: QuantizationConfig) -> None:
         return
 
     weight_node = node.args[1]
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quant_config.weight,
@@ -260,5 +260,5 @@ def annotate_embedding_op(node: Node, quant_config: QuantizationConfig) -> None:
         return
 
     wgt_node = node.args[0]
-    _annotate_input_qspec_map(node, wgt_node, quant_config.activation)
+    annotate_input_qspec_map(node, wgt_node, quant_config.activation)
     _mark_as_annotated([node])

backends/mediatek/quantizer/qconfig.py

@@ -10,9 +10,9 @@
 
 import torch
 
-from torch.ao.quantization.fake_quantize import FakeQuantize
-from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver
-from torch.ao.quantization.quantizer import QuantizationSpec
+from torchao.quantization.pt2e.fake_quantize import FakeQuantize
+from torchao.quantization.pt2e.observer import MinMaxObserver, PerChannelMinMaxObserver
+from torchao.quantization.pt2e.quantizer import QuantizationSpec
 
 
 @unique

backends/mediatek/quantizer/quantizer.py

@@ -4,8 +4,8 @@
 # except in compliance with the License. See the license file in the root
 # directory of this source tree for more details.
 
-from torch.ao.quantization.quantizer import Quantizer
 from torch.fx import GraphModule
+from torchao.quantization.pt2e.quantizer import Quantizer
 
 from .._passes.decompose_scaled_dot_product_attention import (
     DecomposeScaledDotProductAttention,

backends/nxp/quantizer/neutron_quantizer.py

@@ -35,9 +35,9 @@
     QuantizationSpec,
 )
 from torch import fx
-from torch.ao.quantization.observer import HistogramObserver, MinMaxObserver
-from torch.ao.quantization.quantizer import DerivedQuantizationSpec, Quantizer
-from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
+from torchao.quantization.pt2e.observer import HistogramObserver, MinMaxObserver
+from torchao.quantization.pt2e.quantizer import DerivedQuantizationSpec, Quantizer
+from torchao.quantization.pt2e.quantizer.composable_quantizer import ComposableQuantizer
 
 
 class NeutronAtenQuantizer(Quantizer):

backends/nxp/quantizer/patterns.py

@@ -14,7 +14,7 @@
 from executorch.backends.nxp.quantizer.utils import get_bias_qparams
 from torch import fx
 from torch._ops import OpOverload
-from torch.ao.quantization.quantizer import (
+from torchao.quantization.pt2e.quantizer import (
     DerivedQuantizationSpec,
     FixedQParamsQuantizationSpec,
     SharedQuantizationSpec,