Migrate pt2e qualcomm

Differential Revision: D75166700

Pull Request resolved: pytorch#11049

Author: metascroy

.lintrunner.toml

@@ -388,15 +388,12 @@ exclude_patterns = [
     # backends
     "backends/vulkan/quantizer/**",
     "backends/vulkan/test/**",
-    "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",
 ]
 

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -135,8 +135,8 @@ def get_to_edge_transform_passes(
         from executorch.backends.qualcomm.builders import node_visitor
         from executorch.exir.dialects._ops import ops as exir_ops
 
-        node_visitor.q_ops.add(exir_ops.edge.pt2e_quant.quantize_affine.default)
-        node_visitor.dq_ops.add(exir_ops.edge.pt2e_quant.dequantize_affine.default)
+        node_visitor.q_ops.add(exir_ops.edge.torchao.quantize_affine.default)
+        node_visitor.dq_ops.add(exir_ops.edge.torchao.dequantize_affine.default)
 
         passes_job = (
             passes_job if passes_job is not None else get_capture_program_passes()

backends/qualcomm/builders/node_visitor.py

@@ -265,8 +265,8 @@ def get_quant_encoding_conf(
         )
         # TODO: refactor this when target could be correctly detected
         per_block_encoding = {
-            exir_ops.edge.pt2e_quant.quantize_affine.default,
-            exir_ops.edge.pt2e_quant.dequantize_affine.default,
+            exir_ops.edge.torchao.quantize_affine.default,
+            exir_ops.edge.torchao.dequantize_affine.default,
         }
         if quant_attrs[QCOM_ENCODING] in per_block_encoding:
             return self.make_qnn_per_block_config(node, quant_attrs)

backends/qualcomm/partition/utils.py

@@ -57,7 +57,7 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]:
         torch.ops.aten.upsample_bicubic2d.vec,
         # This request is ignored because it is in a blocklist. Refer to exir/program/_program.py
         torch.ops.aten.unbind.int,
-        torch.ops.pt2e_quant.quantize_affine.default,
-        torch.ops.pt2e_quant.dequantize_affine.default,
+        torch.ops.torchao.quantize_affine.default,
+        torch.ops.torchao.dequantize_affine.default,
     ]
     return do_not_decompose

backends/qualcomm/quantizer/annotators.py

@@ -12,20 +12,17 @@
 from torch._ops import OpOverload
 
 from torch._subclasses import FakeTensor
-from torch.ao.quantization.fake_quantize import FixedQParamsFakeQuantize
+from torch.fx import Node
 
-from torch.ao.quantization.observer import FixedQParamsObserver
-from torch.ao.quantization.quantizer import (
+from torchao.quantization.pt2e import FixedQParamsFakeQuantize, FixedQParamsObserver
+from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec,
     DerivedQuantizationSpec,
     QuantizationAnnotation,
     QuantizationSpec,
     SharedQuantizationSpec,
 )
-from torch.ao.quantization.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
-)
-from torch.fx import Node
 
 from .qconfig import (
     get_16a16w_qnn_ptq_config,
@@ -643,19 +640,19 @@ def annotate_rms_norm(node: Node, quantization_config: QuantizationConfig) -> No
         return
 
     # TODO current only support 16a16w
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         quantization_config.input_activation,
     )
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quantization_config.input_activation,
     )
     nodes_to_mark_annotated = [node]
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 
@@ -844,25 +841,25 @@ def annotate_group_norm(node: Node, quantization_config: QuantizationConfig) ->
     if _is_annotated([node]):
         return
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         quantization_config.input_activation,
     )
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quantization_config.weight,
     )
     nodes_to_mark_annotated = [node, weight_node]
     if bias_node:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             bias_node,
             quantization_config.bias,
         )
         nodes_to_mark_annotated.append(bias_node)
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 
@@ -1027,12 +1024,12 @@ def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None
     if _is_annotated([node]):
         return
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         quantization_config.input_activation,
     )
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quantization_config.weight,
@@ -1043,9 +1040,9 @@ def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None
             bias_config = quantization_config.bias(node)
         else:
             bias_config = quantization_config.bias
-        _annotate_input_qspec_map(node, bias_node, bias_config)
+        annotate_input_qspec_map(node, bias_node, bias_config)
         nodes_to_mark_annotated.append(bias_node)
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
     # We use get_source_partition in pass, but it is the same source for MultiheadAttention, so we need to change its source_fn_stack.
@@ -1063,29 +1060,29 @@ def annotate_batch_and_instance_norm(
         return
 
     annotated_args = [act]
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act,
         quantization_config.input_activation,
     )
     # QNN requires uint8 instead of int8 in 'weight' config
     if weight is not None:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             weight,
             quantization_config.input_activation,
         )
         annotated_args.append(weight)
 
     if bias is not None:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             bias,
             quantization_config.bias,
         )
         annotated_args.append(bias)
 
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated([node, *annotated_args])
 
 
@@ -1095,7 +1092,7 @@ def annotate_getitem(node: Node, quantization_config: QuantizationConfig) -> Non
         return
 
     if _is_float_tensor(node):
-        _annotate_output_qspec(node, quantization_config.output_activation)
+        annotate_output_qspec(node, quantization_config.output_activation)
         _mark_nodes_as_annotated([node])
 
 
@@ -1111,32 +1108,32 @@ def annotate_layer_norm(node: Node, quantization_config: QuantizationConfig) ->
         return
     input_act_qspec = quantization_config.input_activation
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         input_act_qspec,
     )
     if input_act_qspec.dtype == torch.int32:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             weight_node,
             get_16a16w_qnn_ptq_config().weight,
         )
     else:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             weight_node,
             input_act_qspec,
         )
     nodes_to_mark_annotated = [node, weight_node]
     if bias_node:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             bias_node,
             quantization_config.bias,
         )
         nodes_to_mark_annotated.append(bias_node)
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 

backends/qualcomm/quantizer/custom_annotation.py

@@ -17,13 +17,13 @@
     QuantizationConfig,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
-from torch.ao.quantization.observer import FixedQParamsObserver, MinMaxObserver
-from torch.ao.quantization.quantizer import (
+from torch.fx import Node
+from torchao.quantization.pt2e import FixedQParamsObserver, MinMaxObserver
+from torchao.quantization.pt2e.quantizer import (
     QuantizationAnnotation,
     QuantizationSpec,
     SharedQuantizationSpec,
 )
-from torch.fx import Node
 
 
 def annotate_mimi_decoder(gm: torch.fx.GraphModule):

backends/qualcomm/quantizer/observers/per_block_param_observer.py

@@ -7,8 +7,8 @@
 from typing import Tuple
 
 import torch
-from torch.ao.quantization.observer import MappingType, PerBlock
-from torch.ao.quantization.pt2e._affine_quantization import (
+from torchao.quantization.pt2e import MappingType, PerBlock
+from torchao.quantization.pt2e._affine_quantization import (
     _get_reduction_params,
     AffineQuantizedMinMaxObserver,
     choose_qparams_affine_with_min_max,

backends/qualcomm/quantizer/observers/per_channel_param_observer.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torch.ao.quantization.observer import UniformQuantizationObserverBase
+from torchao.quantization.pt2e import UniformQuantizationObserverBase
 
 
 # TODO move to torch/ao/quantization/observer.py.

backends/qualcomm/quantizer/qconfig.py

@@ -7,18 +7,19 @@
     PerBlockParamObserver,
 )
 from torch import Tensor
-from torch.ao.quantization.fake_quantize import (
+from torch.fx import Node
+from torchao.quantization.pt2e import (
     FakeQuantize,
     FusedMovingAvgObsFakeQuantize,
-)
-from torch.ao.quantization.observer import (
     MinMaxObserver,
     MovingAverageMinMaxObserver,
     MovingAveragePerChannelMinMaxObserver,
     PerChannelMinMaxObserver,
 )
-from torch.ao.quantization.quantizer import DerivedQuantizationSpec, QuantizationSpec
-from torch.fx import Node
+from torchao.quantization.pt2e.quantizer import (
+    DerivedQuantizationSpec,
+    QuantizationSpec,
+)
 
 
 @dataclass(eq=True)

backends/qualcomm/quantizer/quantizer.py

@@ -12,8 +12,9 @@
 from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager
 
 from torch._ops import OpOverload
-from torch.ao.quantization.quantizer import Quantizer
 from torch.fx import GraphModule
+from torchao.quantization.pt2e import UniformQuantizationObserverBase
+from torchao.quantization.pt2e.quantizer import Quantizer
 
 from .annotators import OP_ANNOTATOR
 
@@ -130,9 +131,7 @@ class ModuleQConfig:
     is_qat: bool = False
     is_conv_per_channel: bool = False
     is_linear_per_channel: bool = False
-    act_observer: Optional[
-        torch.ao.quantization.observer.UniformQuantizationObserverBase
-    ] = None
+    act_observer: Optional[UniformQuantizationObserverBase] = None
 
     def __post_init__(self):
         if (self.quant_dtype, self.is_qat) not in QUANT_CONFIG_DICT: