[quant][refactor] Move some util functions from torch/quantization/fx/utils.py to torch/quantization/utils.py (pytorch#48107)

Summary: Pull Request resolved: pytorch#48107

Test Plan: Imported from OSS

Reviewed By: supriyar

Differential Revision: D25026495

fbshipit-source-id: 3634b6b95a18670232600874b1e593180ea9f44c

Author: jerryzh168

torch/quantization/fx/quantization_patterns.py

@@ -13,20 +13,24 @@
     get_static_quant_module_class,
     get_quantized_operator,
 )
+from ..utils import (
+    get_swapped_custom_module_class,
+    activation_is_statically_quantized,
+    weight_is_statically_quantized,
+    weight_dtype,
+    get_qconfig_dtypes,
+)
+
 from .pattern_utils import (
     register_quant_pattern,
     mark_input_output_not_observed,
 )
+
 from .utils import (
     _parent_name,
     quantize_node,
     get_per_tensor_qparams,
-    get_swapped_custom_module_class,
-    activation_is_statically_quantized,
-    weight_is_quantized,
-    weight_dtype,
     get_linear_prepack_op_for_dtype,
-    get_qconfig_dtypes,
 )
 
 from abc import ABC, abstractmethod
@@ -339,7 +343,7 @@ def convert(self, quantizer, node, load_arg, debug=False, convert_custom_config_
                 quantized_input_idxs = []
                 if activation_statically_quantized:
                     quantized_input_idxs.append(0)
-                if weight_is_quantized(qconfig):
+                if weight_is_statically_quantized(qconfig):
                     quantized_input_idxs.append(1)
                 args = load_arg(quantized=quantized_input_idxs)(self.linear_node.args)
                 args = load_arg(quantized=False)(self.linear_node.args)
@@ -360,7 +364,7 @@ def convert(self, quantizer, node, load_arg, debug=False, convert_custom_config_
             else:  # non-debug option
                 # linear args
                 # (x, weight, bias, ...)
-                weight_quantized = weight_is_quantized(qconfig)
+                weight_quantized = weight_is_statically_quantized(qconfig)
                 linear_weight = load_arg(quantized=weight_quantized)(self.linear_node.args[1])
 
                 # get other arguments

torch/quantization/fx/quantize.py

@@ -25,7 +25,9 @@
 )
 
 from ..utils import (
-    get_combined_dict
+    get_combined_dict,
+    get_swapped_custom_module_class,
+    activation_is_statically_quantized,
 )
 
 from .pattern_utils import (
@@ -48,8 +50,6 @@
     _parent_name,
     quantize_node,
     get_custom_module_class_keys,
-    get_swapped_custom_module_class,
-    activation_is_statically_quantized,
 )
 
 from collections import OrderedDict

torch/quantization/fx/utils.py

@@ -1,6 +1,6 @@
 import re
 import torch
-from ..quant_type import QuantType, quant_type_to_str
+from ..utils import is_per_tensor, is_per_channel
 
 # turn foo.bar -> ['foo', 'bar']
 def _parent_name(target):
@@ -76,15 +76,6 @@ def graph_pretty_str(g, shorten=True) -> str:
         res_str += "*obs_{n} = activation_post_process_{n}\n"
     return res_str
 
-def is_per_tensor(qscheme):
-    return qscheme == torch.per_tensor_affine or \
-        qscheme == torch.per_tensor_symmetric
-
-def is_per_channel(qscheme):
-    return qscheme in [torch.per_channel_affine,
-                       torch.per_channel_affine_float_qparams,
-                       torch.per_channel_symmetric]
-
 def get_per_tensor_qparams(activation_post_process):
     assert is_per_tensor(activation_post_process.qscheme), 'Only per tensor quantization is supported'
     scale, zero_point = activation_post_process.calculate_qparams()
@@ -171,73 +162,6 @@ def get_custom_module_class_keys(custom_config_dict, custom_config_dict_key):
         float_custom_module_classes |= quant_mode_custom_module_classes
     return list(float_custom_module_classes)
 
-def get_swapped_custom_module_class(custom_module, custom_module_class_mapping, qconfig):
-    """ Get the observed/quantized custom module class that we need
-    to swap `custom_module` to
-    Input:
-        custom_module: input, can be an instance of either a float or observed custom module
-        custom_module_class_mapping: the float to observed or observed to quantized custom module class mapping
-        qconfig: qconfig configured for the custom module
-
-    Output:
-        corresponding observed/quantized custom module class for input custom module instance
-    """
-    quant_type = get_quant_type(qconfig)
-    quant_type_str = quant_type_to_str(quant_type)
-    class_mapping = custom_module_class_mapping.get(quant_type_str, {})
-    assert type(custom_module) in class_mapping, "did not found corresponding observed " \
-        "module class for {} in mapping: {}".format(type(custom_module), class_mapping)
-    return class_mapping[type(custom_module)]
-
-def activation_is_statically_quantized(qconfig):
-    """ Given a qconfig, decide if the activation needs to be
-    statically quantized or not
-    """
-    assert qconfig is not None
-    activation = qconfig.activation()
-    return activation.dtype in [torch.quint8, torch.qint8]
-
-def weight_dtype(qconfig):
-    assert qconfig is not None
-    weight = qconfig.weight()
-    return weight.dtype
-
-def weight_is_quantized(qconfig):
-    """ Given a qconfig, decide if the activation needs to be
-    quantized or not
-    """
-    return weight_dtype(qconfig) in [torch.quint8, torch.qint8]
-
-def get_qconfig_dtypes(qconfig):
-    r""" returns the qconfig tuple for qconfig:
-    (activation_dtype, weight_dtype, activation_compute_dtype)
-    """
-    assert qconfig is not None
-    activation = qconfig.activation()
-    weight = qconfig.weight()
-    compute_dtype = activation.compute_dtype if hasattr(activation, 'compute_dtype') else None
-    return (activation.dtype, weight.dtype, compute_dtype)
-
-def get_quant_type(qconfig):
-    assert qconfig is not None
-    activation = qconfig.activation()
-    weight = qconfig.weight()
-    static_dtypes = [torch.quint8, torch.qint8]
-    if weight.dtype in static_dtypes:
-        if activation.dtype in static_dtypes:
-            return QuantType.STATIC
-        elif hasattr(activation, 'compute_dtype') and activation.compute_dtype in static_dtypes:
-            return QuantType.DYNAMIC
-        else:
-            return QuantType.WEIGHT_ONLY
-
-    if weight.dtype == torch.float16:
-        if activation.dtype == torch.float:
-            return QuantType.DYNAMIC
-
-    raise Exception("Unrecognized dtype combination in get_quant_type: activation({}),"
-                    "weight({})".format(activation.dtype, weight.dtype))
-
 def get_linear_prepack_op_for_dtype(dtype):
     if dtype == torch.float16:
         return torch.ops.quantized.linear_prepack_fp16

torch/quantization/utils.py

@@ -1,9 +1,87 @@
 """
 Utils shared by different modes of quantization (eager/graph)
 """
+import torch
+from .quant_type import QuantType, quant_type_to_str
 
 def get_combined_dict(default_dict, additional_dict):
     d = default_dict.copy()
     for k, v in additional_dict.items():
         d[k] = v
     return d
+
+def is_per_tensor(qscheme):
+    return qscheme == torch.per_tensor_affine or \
+        qscheme == torch.per_tensor_symmetric
+
+def is_per_channel(qscheme):
+    return qscheme in [torch.per_channel_affine,
+                       torch.per_channel_affine_float_qparams,
+                       torch.per_channel_symmetric]
+
+def get_swapped_custom_module_class(custom_module, custom_module_class_mapping, qconfig):
+    """ Get the observed/quantized custom module class that we need
+    to swap `custom_module` to
+    Input:
+        custom_module: input, can be an instance of either a float or observed custom module
+        custom_module_class_mapping: the float to observed or observed to quantized custom module class mapping
+        qconfig: qconfig configured for the custom module
+
+    Output:
+        corresponding observed/quantized custom module class for input custom module instance
+    """
+    quant_type = get_quant_type(qconfig)
+    quant_type_str = quant_type_to_str(quant_type)
+    class_mapping = custom_module_class_mapping.get(quant_type_str, {})
+    assert type(custom_module) in class_mapping, "did not found corresponding observed " \
+        "module class for {} in mapping: {}".format(type(custom_module), class_mapping)
+    return class_mapping[type(custom_module)]
+
+def activation_is_statically_quantized(qconfig):
+    """ Given a qconfig, decide if the activation needs to be
+    statically quantized or not
+    """
+    assert qconfig is not None
+    activation = qconfig.activation()
+    return activation.dtype in [torch.quint8, torch.qint8]
+
+def weight_dtype(qconfig):
+    assert qconfig is not None
+    weight = qconfig.weight()
+    return weight.dtype
+
+def weight_is_statically_quantized(qconfig):
+    """ Given a qconfig, decide if the weight needs to be
+    quantized or not
+    """
+    return weight_dtype(qconfig) in [torch.quint8, torch.qint8]
+
+def get_qconfig_dtypes(qconfig):
+    r""" returns the qconfig tuple for qconfig:
+    (activation_dtype, weight_dtype, activation_compute_dtype)
+    """
+    assert qconfig is not None
+    activation = qconfig.activation()
+    weight = qconfig.weight()
+    compute_dtype = activation.compute_dtype if hasattr(activation, 'compute_dtype') else None
+    return (activation.dtype, weight.dtype, compute_dtype)
+
+def get_quant_type(qconfig):
+    assert qconfig is not None
+    activation = qconfig.activation()
+    weight = qconfig.weight()
+    static_dtypes = [torch.quint8, torch.qint8]
+    if weight.dtype in static_dtypes:
+        if activation.dtype in static_dtypes:
+            return QuantType.STATIC
+        elif hasattr(activation, 'compute_dtype') and activation.compute_dtype in static_dtypes:
+            return QuantType.DYNAMIC
+        else:
+            return QuantType.WEIGHT_ONLY
+
+    if weight.dtype == torch.float16:
+        if activation.dtype == torch.float:
+            return QuantType.DYNAMIC
+
+    raise Exception("Unrecognized dtype combination in get_quant_type: activation({}),"
+                    "weight({})".format(activation.dtype, weight.dtype))