Make v1 quantsim export v2-compatible

TrainingExtensions/torch/src/python/aimet_torch/experimental/v2/quantization/quantization_mixin.py

@@ -0,0 +1,59 @@
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2023-2023, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+# pylint: skip-file
+""" Placeholder for _QuantizationMixin definition, to be deleted/moved/updated """
+
+from abc import ABC, abstractmethod
+
+class _QuantizationMixin(ABC):
+    """ Base class for quantized modules """
+
+    @abstractmethod
+    def export_input_encodings(self):
+        ...
+
+    @abstractmethod
+    def export_output_encodings(self):
+        ...
+
+    @abstractmethod
+    def export_param_encodings(self):
+        ...
+
+    @abstractmethod
+    def get_original_module(self):
+        ...

TrainingExtensions/torch/src/python/aimet_torch/qc_quantize_op.py

@@ -40,7 +40,7 @@
 # pylint: disable=too-many-lines
 import abc
 from enum import Enum
-from typing import Dict, Tuple, Union, List, Callable, Type, Any
+from typing import Dict, Tuple, Union, List, Callable, Type, Any, Optional
 import os
 import torch
 from torch import nn
@@ -52,7 +52,7 @@
 from aimet_torch.custom import custom_tensor_utils
 from aimet_torch import utils
 from aimet_torch.tensor_quantizer import StaticGridPerTensorQuantizer, StaticGridPerChannelQuantizer, TensorQuantizer, \
-    LearnedGridTensorQuantizer, set_encoding_min_max_gating_threshold
+    LearnedGridTensorQuantizer, set_encoding_min_max_gating_threshold, StaticGridTensorQuantizer
 from aimet_torch.torch_quantizer import TorchQuantizer
 import aimet_torch.quantsim_straight_through_grad as ste
 
@@ -490,6 +490,30 @@ def should_perform_quant_dequant(tensor: torch.Tensor, tensor_quantizer: TensorQ
             return False
         return True
 
+    def get_original_module(self) -> torch.nn.Module:
+        """
+        Returns the wrapped torch.nn.Module
+        """
+        return self._module_to_wrap
+
+    def export_param_encodings(self) -> Dict[str, List]:
+        """
+        Returns the layer's parameter encodings in an exportable format
+        """
+        return {name: export_quantizer_encoding(quantizer) for name, quantizer in self.param_quantizers.items()}
+
+    def export_output_encodings(self) -> List[List[Dict]]:
+        """
+        Returns the layer's output encodings in an exportable format
+        """
+        return [export_quantizer_encoding(quantizer) for quantizer in self.output_quantizers]
+
+    def export_input_encodings(self) -> List[List[Dict]]:
+        """
+        Returns the layer's input encodings in an exportable format
+        """
+        return [export_quantizer_encoding(quantizer) for quantizer in self.input_quantizers]
+
 
 class StaticGridQuantWrapper(QcQuantizeWrapper):
     """ A custom PyTorch module that derives from QcQuantizeWrapper and quantizes modules """
@@ -1324,3 +1348,37 @@ def backward(ctx, grad):
             dloss_by_db = grad.sum(dim=0)
 
         return dloss_by_dx, dloss_by_dW, dloss_by_db, dloss_by_dmin, dloss_by_dmax, None
+
+
+def get_encoding_by_quantizer(quantizer: Union[StaticGridTensorQuantizer, LearnedGridTensorQuantizer]) \
+        -> Optional[Union[libpymo.TfEncoding, List[libpymo.TfEncoding]]]:
+    """
+    Retrieve encoding object by quantizer type (StaticGridTensorQuantizer or LearnedGridTensorQuantizer)
+    In particular, LearnedGridTensorQuantizer should use get_effective_encoding to achieve true encoding
+
+    :param quantizer: TensorQuantizer (StaticGridTensorQuantizer or LearnedGridTensorQuantizer)
+    :return: TfEncoding or list of TfEncoding. None if quantizer is not enabled
+    """
+    if isinstance(quantizer, LearnedGridTensorQuantizer):
+        return quantizer.get_effective_encoding()
+
+    return quantizer.encoding
+
+
+def export_quantizer_encoding(quantizer: Union[StaticGridTensorQuantizer, LearnedGridTensorQuantizer]) \
+        -> Optional[List[Dict]]:
+    """
+    Returns the encoding of a quantizer in exportable form.
+
+    :param quantizer: Quantizer from which to export the encoding
+    :return: List of encoding dictionaries for the quantizer
+    """
+    if not quantizer.enabled:
+        return None
+    encoding = get_encoding_by_quantizer(quantizer)
+    if isinstance(encoding, List):
+        encoding = [utils.create_encoding_dict(enc, quantizer, False) for enc in encoding]
+    else:
+        encoding = utils.create_encoding_dict(encoding, quantizer, False)
+        encoding = [encoding] if encoding else None
+    return encoding