[quant][pt2e] Rename _pt2e to pt2e (pytorch#104668)

jerryzh168 · facebook-github-bot · commit 33cb60d909af · 2023-07-10T18:11:27.000-07:00
Summary: Pull Request resolved: pytorch#104668 X-link: pytorch/executorch#3 att Test Plan: Imported from OSS Reviewed By: andrewor14 Differential Revision: D47202807 fbshipit-source-id: d6cdda700e7fecbdf2cc8b7940ef9245bd0da891
diff --git a/docs/source/quantization-support.rst b/docs/source/quantization-support.rst
@@ -120,6 +120,13 @@ This module contains a few CustomConfig classes that's used in both eager mode a
     ConvertCustomConfig
     StandaloneModuleConfigEntry
 
+torch.ao.quantization.pt2e (quantization in pytorch 2.0 export)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: torch.ao.quantization.pt2e
+.. automodule:: torch.ao.quantization.pt2e.quantizer
+.. automodule:: torch.ao.quantization.pt2e.representation
+
 torch (quantization related functions)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/test/inductor/test_inductor_freezing.py b/test/inductor/test_inductor_freezing.py
@@ -12,13 +12,13 @@
 import torch
 
 import torch._dynamo as torchdynamo
-import torch.ao.quantization._pt2e.quantizer.x86_inductor_quantizer as xiq
+import torch.ao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq
 from torch import nn
 from torch._inductor import config
 from torch._inductor.compile_fx import compile_fx
 from torch._inductor.utils import override_lowering, run_and_get_code
-from torch.ao.quantization._pt2e.quantizer import X86InductorQuantizer
 from torch.ao.quantization._quantize_pt2e import convert_pt2e, prepare_pt2e_quantizer
+from torch.ao.quantization.pt2e.quantizer import X86InductorQuantizer
 from torch.testing import FileCheck
 from torch.testing._internal.common_quantization import (
     skipIfNoDynamoSupport,
diff --git a/test/quantization/pt2e/test_graph_utils.py b/test/quantization/pt2e/test_graph_utils.py
@@ -5,7 +5,7 @@
 import torch
 import torch._dynamo as torchdynamo
 
-from torch.ao.quantization._pt2e.graph_utils import (
+from torch.ao.quantization.pt2e.graph_utils import (
     find_sequential_partitions,
     get_equivalent_types,
     update_equivalent_types_dict,
diff --git a/test/quantization/pt2e/test_quantize_pt2e.py b/test/quantization/pt2e/test_quantize_pt2e.py
@@ -15,7 +15,7 @@
     ObserverOrFakeQuantize,
     QConfigMapping,
 )
-from torch.ao.quantization._pt2e.quantizer import (
+from torch.ao.quantization.pt2e.quantizer import (
     ComposableQuantizer,
     DerivedQuantizationSpec,
     EmbeddingQuantizer,
@@ -27,10 +27,10 @@
     Quantizer,
     SharedQuantizationSpec,
 )
-from torch.ao.quantization._pt2e.quantizer.composable_quantizer import (  # noqa: F811
+from torch.ao.quantization.pt2e.quantizer.composable_quantizer import (  # noqa: F811
     ComposableQuantizer,
 )
-from torch.ao.quantization._pt2e.quantizer.qnnpack_quantizer import (
+from torch.ao.quantization.pt2e.quantizer.qnnpack_quantizer import (
     get_symmetric_quantization_config,
 )
 from torch.ao.quantization._quantize_pt2e import (
@@ -1774,7 +1774,7 @@ def __init__(self):
             def forward(self, x, y):
                 return x + y
 
-        import torch.ao.quantization._pt2e.quantizer.qnnpack_quantizer as qq
+        import torch.ao.quantization.pt2e.quantizer.qnnpack_quantizer as qq
 
         quantizer = QNNPackQuantizer()
         operator_config = qq.get_symmetric_quantization_config(is_per_channel=True)
@@ -1799,7 +1799,7 @@ def __init__(self):
             def forward(self, x, y):
                 return x + y
 
-        import torch.ao.quantization._pt2e.quantizer.qnnpack_quantizer as qq
+        import torch.ao.quantization.pt2e.quantizer.qnnpack_quantizer as qq
 
         quantizer = QNNPackQuantizer()
         operator_config = qq.get_symmetric_quantization_config(is_per_channel=True)
diff --git a/test/quantization/pt2e/test_x86inductor_quantizer.py b/test/quantization/pt2e/test_x86inductor_quantizer.py
@@ -3,7 +3,7 @@
 import torch
 import torch._dynamo as torchdynamo
 import torch.nn as nn
-from torch.ao.quantization._pt2e.quantizer import (
+from torch.ao.quantization.pt2e.quantizer import (
     X86InductorQuantizer,
 )
 from torch.ao.quantization._quantize_pt2e import (
@@ -19,7 +19,7 @@
 from torch.testing._internal.common_quantized import override_quantized_engine
 from enum import Enum
 import itertools
-import torch.ao.quantization._pt2e.quantizer.x86_inductor_quantizer as xiq
+import torch.ao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq
 from torch.testing._internal.common_utils import skip_but_pass_in_sandcastle
 
 
diff --git a/torch/_dynamo/skipfiles.py b/torch/_dynamo/skipfiles.py
@@ -139,9 +139,9 @@ def _module_dir(m: types.ModuleType):
 # TODO: find a better way to express this path without having to import
 # `torch.ao.quantization._pt2e`, which interferes with memory profiling
 FILENAME_ALLOWLIST |= {
-    _module_dir(torch) + "ao/quantization/_pt2e/qat_utils.py",
-    _module_dir(torch) + "ao/quantization/_pt2e/quantizer/qnnpack_quantizer.py",
-    _module_dir(torch) + "ao/quantization/_pt2e/representation/rewrite.py",
+    _module_dir(torch) + "ao/quantization/pt2e/qat_utils.py",
+    _module_dir(torch) + "ao/quantization/pt2e/quantizer/qnnpack_quantizer.py",
+    _module_dir(torch) + "ao/quantization/pt2e/representation/rewrite.py",
 }
 
 # TODO (zhxchen17) Make exportdb importable here.
diff --git a/torch/_inductor/freezing.py b/torch/_inductor/freezing.py
@@ -13,7 +13,7 @@
 from torch._inductor.compile_fx import fake_tensor_prop
 from torch._inductor.fx_passes.freezing_patterns import freezing_passes
 from torch._inductor.fx_passes.post_grad import view_to_reshape
-from torch.ao.quantization._pt2e.utils import _fuse_conv_bn_
+from torch.ao.quantization.pt2e.utils import _fuse_conv_bn_
 from torch.fx.experimental.proxy_tensor import make_fx
 from . import config
 from .decomposition import select_decomp_table
diff --git a/torch/ao/quantization/_quantize_pt2e.py b/torch/ao/quantization/_quantize_pt2e.py
@@ -1,21 +1,21 @@
 from torch.fx import GraphModule
 
-from ._pt2e.prepare import prepare
-from ._pt2e._propagate_annotation import propagate_annotation
-from ._pt2e.qat_utils import (
+from .pt2e.prepare import prepare
+from .pt2e._propagate_annotation import propagate_annotation
+from .pt2e.qat_utils import (
     _fuse_conv_bn_qat,
     _fold_conv_bn_qat,
 )
-from ._pt2e.utils import (
+from .pt2e.utils import (
     _get_node_name_to_scope,
     _fuse_conv_bn_,
     _rearrange_weight_observer_for_decomposed_linear,
 )
-from ._pt2e.representation import reference_representation_rewrite
+from .pt2e.representation import reference_representation_rewrite
 from .fx.prepare import prepare as fx_prepare
 from .quantize_fx import _convert_to_reference_decomposed_fx
 from torch.ao.quantization import QConfigMapping
-from torch.ao.quantization._pt2e.quantizer import Quantizer
+from torch.ao.quantization.pt2e.quantizer import Quantizer
 from torch.ao.quantization.backend_config import BackendConfig
 
 from typing import Any, Tuple
diff --git a/torch/ao/quantization/fx/prepare.py b/torch/ao/quantization/fx/prepare.py
@@ -106,7 +106,7 @@
     PrepareCustomConfig,
     StandaloneModuleConfigEntry,
 )
-from torch.ao.quantization._pt2e.quantizer import (
+from torch.ao.quantization.pt2e.quantizer import (
     EdgeOrNode,
     QuantizationSpec,
     FixedQParamsQuantizationSpec,
diff --git a/torch/ao/quantization/pt2e/__init__.py b/torch/ao/quantization/pt2e/__init__.py
diff --git a/torch/ao/quantization/pt2e/_propagate_annotation.py b/torch/ao/quantization/pt2e/_propagate_annotation.py
@@ -1,7 +1,7 @@
 from typing import Callable
 
 import torch
-from torch.ao.quantization._pt2e.quantizer import (
+from torch.ao.quantization.pt2e.quantizer import (
     QuantizationAnnotation,
     SharedQuantizationSpec,
 )
diff --git a/torch/ao/quantization/pt2e/graph_utils.py b/torch/ao/quantization/pt2e/graph_utils.py
@@ -10,6 +10,12 @@
     SourcePartition,
 )
 
+__all__ = [
+    "find_sequential_partitions",
+    "get_equivalent_types",
+    "update_equivalent_types_dict",
+]
+
 _EQUIVALENT_TYPES: List[Set] = [
     {torch.nn.Conv2d, torch.nn.functional.conv2d},
     {torch.nn.AdaptiveAvgPool2d, torch.nn.functional.adaptive_avg_pool2d},
diff --git a/torch/ao/quantization/pt2e/prepare.py b/torch/ao/quantization/pt2e/prepare.py
@@ -19,7 +19,7 @@
 from torch.ao.quantization.qconfig import QConfigAny
 from torch.ao.quantization.fx.custom_config import PrepareCustomConfig
 from typing import Dict, Tuple, Union, Any
-from torch.ao.quantization._pt2e.quantizer import (
+from torch.ao.quantization.pt2e.quantizer import (
     QuantizationAnnotation,
     EdgeOrNode,
 )
diff --git a/torch/ao/quantization/pt2e/qat_utils.py b/torch/ao/quantization/pt2e/qat_utils.py
@@ -14,8 +14,8 @@
     SharedQuantizationSpec,
     QuantizationSpecBase,
 )
-from .utils import _fold_bn_weights_into_conv_node
-from .utils import _get_aten_graph_module
+from .utils import fold_bn_weights_into_conv_node
+from .utils import get_aten_graph_module
 
 # Example inputs for `_conv2d_bn_pattern`, `_qat_conv2d_bn_pattern`, and `_qat_conv2d_bn_pattern_no_bias`
 _conv2d_bn_pattern_example_inputs = (
@@ -494,15 +494,15 @@ def _fuse_conv_bn_qat(m: GraphModule) -> GraphModule:
     m.graph.eliminate_dead_code()
     m.recompile()
     example_inputs = _conv2d_bn_pattern_example_inputs
-    match_pattern = _get_aten_graph_module(_conv2d_bn_pattern, example_inputs)
+    match_pattern = get_aten_graph_module(_conv2d_bn_pattern, example_inputs)
 
     # Step (1): Replace patterns with conv bias
     #
     # Here we do replacement separately for cases with and without conv bias, since
     # the replacement patterns for these two cases are substantially different.
     # TODO: use the public replace_pattern API once it also returns replacement nodes
 
-    replacement_pattern_with_conv_bias = _get_aten_graph_module(
+    replacement_pattern_with_conv_bias = get_aten_graph_module(
         _qat_conv2d_bn_pattern,
         example_inputs,
     )
@@ -517,7 +517,7 @@ def _fuse_conv_bn_qat(m: GraphModule) -> GraphModule:
 
     # Step (2): Replace patterns without conv bias
 
-    replacement_pattern_no_conv_bias = _get_aten_graph_module(
+    replacement_pattern_no_conv_bias = get_aten_graph_module(
         _qat_conv2d_bn_pattern_no_conv_bias,
         example_inputs,
     )
@@ -650,11 +650,11 @@ def _fold_conv_bn_qat(m: GraphModule) -> GraphModule:
         match_pattern = _get_quantized_qat_conv2d_bn_pattern(
             is_per_channel, has_relu, has_bias, relu_is_inplace,
         )
-        match_pattern = _get_aten_graph_module(match_pattern, example_inputs, **kwargs)
+        match_pattern = get_aten_graph_module(match_pattern, example_inputs, **kwargs)
         replacement_pattern = _get_folded_quantized_qat_conv2d_bn_pattern(
             is_per_channel, has_relu, has_bias, relu_is_inplace,
         )
-        replacement_pattern = _get_aten_graph_module(replacement_pattern, example_inputs, **kwargs)
+        replacement_pattern = get_aten_graph_module(replacement_pattern, example_inputs, **kwargs)
         replacements.extend(
             replace_pattern_with_filters(
                 m,
@@ -718,7 +718,7 @@ def _fold_conv_bn_qat(m: GraphModule) -> GraphModule:
                 )
 
         # fold bn weights into conv
-        _fold_bn_weights_into_conv_node(conv_node, conv_weight, conv_bias, bn_node, m)
+        fold_bn_weights_into_conv_node(conv_node, conv_weight, conv_bias, bn_node, m)
 
         # Copy over literal args for conv
         for _, original_node in _filter_nodes_map(r.nodes_map).items():
diff --git a/torch/ao/quantization/pt2e/quantizer/__init__.py b/torch/ao/quantization/pt2e/quantizer/__init__.py
diff --git a/torch/ao/quantization/pt2e/quantizer/composable_quantizer.py b/torch/ao/quantization/pt2e/quantizer/composable_quantizer.py
diff --git a/torch/ao/quantization/pt2e/quantizer/embedding_quantizer.py b/torch/ao/quantization/pt2e/quantizer/embedding_quantizer.py
@@ -5,7 +5,7 @@
 
 import torch
 import torch.nn.functional as F
-from torch.ao.quantization._pt2e.quantizer.quantizer import (
+from torch.ao.quantization.pt2e.quantizer.quantizer import (
     OperatorConfig,
     OperatorPatternType,
     QuantizationAnnotation,
@@ -15,6 +15,10 @@
 )
 from torch.ao.quantization.observer import PerChannelMinMaxObserver
 
+__all__ = [
+    "get_embedding_operators_config",
+    "EmbeddingQuantizer",
+]
 
 def get_embedding_operators_config() -> OperatorConfig:
     weight_quantization_spec = QuantizationSpec(
diff --git a/torch/ao/quantization/pt2e/quantizer/qnnpack_quantizer.py b/torch/ao/quantization/pt2e/quantizer/qnnpack_quantizer.py
@@ -11,9 +11,9 @@
 import torch._dynamo as torchdynamo
 import torch.nn.functional as F
 
-from torch.ao.quantization._pt2e.graph_utils import find_sequential_partitions
+from torch.ao.quantization.pt2e.graph_utils import find_sequential_partitions
 
-from torch.ao.quantization._pt2e.quantizer.utils import (
+from torch.ao.quantization.pt2e.quantizer.utils import (
     _annotate_input_qspec_map,
     _annotate_output_qspec,
     _is_sym_size_node,
@@ -84,7 +84,7 @@ def linear_op(act, weight, bias=None):
     return [pattern_w_bias, pattern_wo_bias]
 
 
-def supported_symmetric_quantized_operators() -> Dict[str, List[OperatorPatternType]]:
+def _supported_symmetric_quantized_operators() -> Dict[str, List[OperatorPatternType]]:
     supported_operators: Dict[str, List[OperatorPatternType]] = {
         # Both conv and linear should be able to handle relu + hardtanh fusion since
         # those are clamp ops
@@ -107,15 +107,15 @@ def supported_symmetric_quantized_operators() -> Dict[str, List[OperatorPatternT
     return copy.deepcopy(supported_operators)
 
 
-def get_supported_symmetric_config_and_operators() -> List[OperatorConfig]:
+def _get_supported_symmetric_config_and_operators() -> List[OperatorConfig]:
     supported_config_and_operators: List[OperatorConfig] = []
     for quantization_config in [
         get_symmetric_quantization_config(),
         get_symmetric_quantization_config(is_qat=True),
         get_symmetric_quantization_config(is_per_channel=True),
         get_symmetric_quantization_config(is_per_channel=True, is_qat=True),
     ]:
-        ops = supported_symmetric_quantized_operators()
+        ops = _supported_symmetric_quantized_operators()
         for op_string, pattern_list in ops.items():
             supported_config_and_operators.append(
                 OperatorConfig(quantization_config, pattern_list)
@@ -205,8 +205,8 @@ def get_symmetric_quantization_config(
     return quantization_config
 
 
-def get_supported_config_and_operators() -> List[OperatorConfig]:
-    return get_supported_symmetric_config_and_operators()
+def _get_supported_config_and_operators() -> List[OperatorConfig]:
+    return _get_supported_symmetric_config_and_operators()
 
 
 def _is_annotated(nodes: List[Node]):
@@ -225,7 +225,7 @@ def _is_annotated(nodes: List[Node]):
 
 
 class QNNPackQuantizer(Quantizer):
-    supported_config_and_operators = get_supported_config_and_operators()
+    supported_config_and_operators = _get_supported_config_and_operators()
 
     def __init__(self):
         super().__init__()
diff --git a/torch/ao/quantization/pt2e/quantizer/quantizer.py b/torch/ao/quantization/pt2e/quantizer/quantizer.py
@@ -13,9 +13,12 @@
     "QuantizationSpecBase",
     "QuantizationSpec",
     "FixedQParamsQuantizationSpec",
+    "EdgeOrNode",
     "SharedQuantizationSpec",
     "DerivedQuantizationSpec",
     "QuantizationAnnotation",
+    "QuantizationConfig",
+    "OperatorConfig",
 ]
 
 # TODO: maybe remove torch.float32
@@ -84,17 +87,19 @@ class FixedQParamsQuantizationSpec(QuantizationSpecBase):
     quant_max: Optional[int] = None
     qscheme: Optional[torch.qscheme] = None
 
+"""
+The way we refer to other points of quantization in the graph will be either
+an input edge or an output value
+input edge is the connection between input node and the node consuming the input, so it's a Tuple[Node, Node]
+output value is an fx Node
+"""
 EdgeOrNode = Union[Tuple[Node, Node], Node]
+EdgeOrNode.__module__ = "torch.ao.quantization.pt2e.quantizer.quantizer"
 
 @dataclass(eq=True, frozen=True)
 class SharedQuantizationSpec(QuantizationSpecBase):
     """
     Quantization spec for the Tensors whose quantization parameters are shared with other Tensors
-
-    The way we refer to other points of quantization in the graph will be either
-    an input edge or an output value
-    input edge is the connection between input node and the node consuming the input, so it's a Tuple[Node, Node]
-    output value is an fx Node
     """
     edge_or_node: EdgeOrNode
 
diff --git a/torch/ao/quantization/pt2e/quantizer/utils.py b/torch/ao/quantization/pt2e/quantizer/utils.py
@@ -1,13 +1,19 @@
 from typing import List
 
 import torch
-from torch.ao.quantization._pt2e.quantizer.quantizer import (
+from torch.ao.quantization.pt2e.quantizer.quantizer import (
     QuantizationAnnotation,
     QuantizationConfig,
     QuantizationSpec,
 )
 from torch.fx import Node
 
+__all__ = [
+    "get_input_act_qspec",
+    "get_output_act_qspec",
+    "get_weight_qspec",
+    "get_bias_qspec",
+]
 
 def get_input_act_qspec(quantization_config: QuantizationConfig):
     if quantization_config is None:
diff --git a/torch/ao/quantization/pt2e/quantizer/x86_inductor_quantizer.py b/torch/ao/quantization/pt2e/quantizer/x86_inductor_quantizer.py
diff --git a/torch/ao/quantization/pt2e/representation/__init__.py b/torch/ao/quantization/pt2e/representation/__init__.py
diff --git a/torch/ao/quantization/pt2e/representation/rewrite.py b/torch/ao/quantization/pt2e/representation/rewrite.py
diff --git a/torch/ao/quantization/pt2e/utils.py b/torch/ao/quantization/pt2e/utils.py

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@`
`106`	`106`	`PrepareCustomConfig,`
`107`	`107`	`StandaloneModuleConfigEntry,`
`108`	`108`	`)`
`109`		`-from torch.ao.quantization._pt2e.quantizer import (`
	`109`	`+from torch.ao.quantization.pt2e.quantizer import (`
`110`	`110`	`EdgeOrNode,`
`111`	`111`	`QuantizationSpec,`
`112`	`112`	`FixedQParamsQuantizationSpec,`