tico 0.1.0.dev250411__py3-none-any.whl

tico/experimental/quantization/passes/propagate_qparam_forward.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import copy
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import (
+    CatArgs,
+    NegArgs,
+    PermuteArgs,
+    ReshapeArgs,
+    SliceArgs,
+)
+
+
+@trace_graph_diff_on_pass
+class PropagateQParamForward(PassBase):
+    """
+    A pass propagates quantization parameters through operations that do not alter them.
+
+    This pass identifies and propagates quantization parameters through operations that
+     do not change their values, such as `permute`, `reshape`, `transpose`, `view` and
+    similar tensor transformations.
+
+    By ensuring that quantization parameters remain consistent across such operations,
+    this pass helps maintain correctness in quantization-aware representations.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        def _propagate_qparam_if_possible(src: torch.fx.Node, dst: torch.fx.Node):
+            if QPARAM_KEY not in src.meta:
+                return
+
+            if (
+                QPARAM_KEY in dst.meta
+                and src.meta[QPARAM_KEY].dtype != dst.meta[QPARAM_KEY].dtype
+            ):
+                return
+
+            dst.meta[QPARAM_KEY] = copy.deepcopy(src.meta[QPARAM_KEY])
+
+            logger.debug(f"{src.name}'s quantparam is propagated to {dst.name}.")
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.permute.default:
+                permute_args = PermuteArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(permute_args.input, node)
+            elif node.target == torch.ops.aten.reshape.default:
+                reshape_args = ReshapeArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(reshape_args.input, node)
+            elif node.target == torch.ops.aten.slice.Tensor:
+                slice_args = SliceArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(slice_args.input, node)
+            elif node.target == torch.ops.aten.neg.default:
+                neg_args = NegArgs(*node.args, **node.kwargs)
+
+                if QPARAM_KEY not in neg_args.input.meta:
+                    continue
+                # Only support int16 for now
+                if neg_args.input.meta[QPARAM_KEY].dtype != "int16":
+                    continue
+
+                _propagate_qparam_if_possible(neg_args.input, node)
+
+            elif node.target == torch.ops.aten.cat.default:
+                concat_args = CatArgs(*node.args, **node.kwargs)
+                concat_inputs = concat_args.tensors
+
+                cond = True
+                for concat_input in concat_inputs:
+                    # Check all inputs have qparam
+                    if QPARAM_KEY not in concat_input.meta:
+                        cond = False
+                        break
+
+                    # Only support int16 for now
+                    if concat_input.meta[QPARAM_KEY].dtype != "int16":
+                        cond = False
+                        break
+
+                    if concat_input.meta[QPARAM_KEY].scale is None:
+                        cond = False
+                        break
+
+                    if len(concat_input.meta[QPARAM_KEY].scale) != 1:
+                        cond = False
+                        break
+
+                if not cond:
+                    continue
+
+                # Find max scale node
+                max_scale = 0.0
+                max_scale_node = None
+                for concat_input in concat_inputs:
+                    scale = concat_input.meta[QPARAM_KEY].scale[0]
+                    if max_scale < scale:
+                        max_scale = scale
+                        max_scale_node = concat_input
+
+                assert max_scale_node is not None
+                _propagate_qparam_if_possible(max_scale_node, node)
+
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/remove_weight_dequant_op.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch._subclasses.fake_tensor import FakeTensor
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY, QuantParam, to_qparam_dtype
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import (
+    DequantizePerChannelArgs,
+    DequantizePerTensorArgs,
+)
+
+
+def get_constant(exported_program: ExportedProgram, node: torch.fx.Node):
+    assert isinstance(node, torch.fx.Node)
+    if node.name in exported_program.constants:
+        return exported_program.constants[node.name]
+    elif node.name in exported_program.graph_signature.inputs_to_buffers:
+        buffer_name = exported_program.graph_signature.inputs_to_buffers[node.name]
+        named_buffer = dict(exported_program.named_buffers())
+        return named_buffer[buffer_name]
+    else:
+        raise RuntimeError("NYI constant")
+
+
+class ValRange:
+    def __init__(self, val: Union[torch.Tensor, List[int]]):
+        if isinstance(val, torch.Tensor):
+            self.max = torch.max(val).item()
+            self.min = torch.min(val).item()
+        elif type(val) == list:
+            self.max = max(val)
+            self.min = min(val)
+        else:
+            raise RuntimeError("Wrong dtype (val)")
+
+    def within(self, min_val, max_val):
+        return self.min >= min_val and self.max <= max_val
+
+
+# Infer dtype using weight, zero point, and dtype
+def infer_dtype(weight: torch.Tensor, zerop: List[int], dtype: torch.dtype) -> str:
+    weight_val = ValRange(weight)
+    zp_val = ValRange(zerop)
+
+    if weight_val.within(0, 15) and zp_val.within(0, 15) and dtype == torch.uint8:
+        return "uint4"
+    else:
+        return to_qparam_dtype(dtype)
+
+
+@trace_graph_diff_on_pass
+class RemoveWeightDequantOp(PassBase):
+    """
+    This pass identifies and removes any remaining Dequantize ops associated with
+     quantized weights.
+
+    Since weights already quantized earlier (and possibly kept in float by
+     attaching a DQ), the final stage of the quantization pipeline typically
+    does not require those DQ ops anymore.
+
+    NOTE Removing 'DQ' causes a sementic change: f32 -> quantized
+
+    [BEFORE]
+      W (quantized) - Dequantize (float)
+
+    [AFTER]
+      W (quantized)
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+        for dq in graph.nodes:
+            if not dq.op == "call_function":
+                continue
+
+            if dq.target not in [
+                torch.ops.quantized_decomposed.dequantize_per_channel.default,
+                torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            ]:
+                continue
+            dq_args: Optional[DequantizePerChannelArgs | DequantizePerTensorArgs] = None
+
+            if (
+                dq.target
+                == torch.ops.quantized_decomposed.dequantize_per_channel.default
+            ):
+                dq_args = DequantizePerChannelArgs(*dq.args, *dq.kwargs)
+            elif (
+                dq.target
+                == torch.ops.quantized_decomposed.dequantize_per_tensor.default
+            ):
+                dq_args = DequantizePerTensorArgs(*dq.args, *dq.kwargs)
+            else:
+                raise RuntimeError(f"Invalid DQ target: {dq.target}")
+
+            q_weight = dq_args.input
+            # All weights are placehoders.
+            if q_weight.op != "placeholder":
+                continue
+            # Check if DQ already has quant param because DQ can be shared.
+            if QPARAM_KEY in q_weight.meta:
+                continue
+
+            q_weight_meta = q_weight.meta["val"]
+            assert isinstance(q_weight_meta, FakeTensor)
+            # Weight should have quantized values.
+            assert q_weight_meta.dtype != torch.float
+
+            q_weight_val = get_constant(exported_program, q_weight)
+            assert isinstance(q_weight_val, torch.Tensor)
+
+            quant_param = QuantParam()
+            if isinstance(dq_args, DequantizePerChannelArgs):
+                scales = get_constant(exported_program, dq_args.scales)
+                zero_ps = get_constant(exported_program, dq_args.zero_points)
+                quant_param.scale = scales.tolist()
+                quant_param.zero_point = zero_ps.tolist()
+                assert quant_param.zero_point is not None  # To avoid mypy error
+                quant_param.quantized_dimension = dq_args.axis
+                quant_param.dtype = infer_dtype(
+                    q_weight_val, quant_param.zero_point, q_weight_meta.dtype
+                )
+            elif isinstance(dq_args, DequantizePerTensorArgs):
+                quant_param.scale = [dq_args.scale]
+                quant_param.zero_point = [dq_args.zero_point]
+                assert quant_param.zero_point is not None  # To avoid mypy error
+                quant_param.dtype = infer_dtype(
+                    q_weight_val, quant_param.zero_point, q_weight_meta.dtype
+                )
+            else:
+                raise RuntimeError(f"Invalid DQ target: {dq.target}")
+
+            q_weight.meta[QPARAM_KEY] = quant_param
+            dq.replace_all_uses_with(q_weight, propagate_meta=False)
+            logger.debug(f"{dq.name} is removed.")
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/public_interface.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from tico.experimental.quantization.algorithm.gptq.quantizer import GPTQQuantizer
+from tico.experimental.quantization.algorithm.pt2e.quantizer import PT2EQuantizer
+from tico.experimental.quantization.algorithm.smoothquant.quantizer import (
+    SmoothQuantQuantizer,
+)
+from tico.experimental.quantization.config import BaseConfig
+from tico.experimental.quantization.quantizer import BaseQuantizer
+
+
+config_to_quantizer: Dict[str, Type[BaseQuantizer]] = {
+    "pt2e": PT2EQuantizer,
+    "gptq": GPTQQuantizer,
+    "smooth_quant": SmoothQuantQuantizer,
+}
+
+QUANTIZER_ATTRIBUTE_NAME = "tico_quantizer"
+
+
+def prepare(
+    model: torch.nn.Module,
+    quant_config: BaseConfig,
+    args: Optional[Any] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    inplace: Optional[bool] = False,
+):
+    """
+    Prepare the model for quantization using the provided configuration.
+
+    Determines the appropriate quantizer based on the type of `quant_config` and
+    prepares the model accordingly.
+
+    Parameters:
+        model: The PyTorch model to be quantized.
+        quant_config (BaseConfig): The quantization configuration.
+        args (Any, optional): Positional example inputs required for activation quantization.
+        kwargs (Dict[str, Any], optional): Keyword example inputs required for activation quantization.
+        inplace (bool, optional): If true, the model will be modified in place;
+                                   otherwise, a new prepared model is returned.
+
+    Returns:
+        The model prepared for quantization.
+    """
+    if quant_config.name == "pt2e" and inplace:
+        raise RuntimeError(
+            "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
+        )
+
+    model = model if inplace else copy.deepcopy(model)
+
+    quantizer = config_to_quantizer[quant_config.name](quant_config)
+    model = quantizer.prepare(model, args, kwargs)
+    setattr(model, QUANTIZER_ATTRIBUTE_NAME, quantizer)
+
+    return model
+
+
+def convert(model, inplace: Optional[bool] = False):
+    """
+    Convert the prepared model to a quantized model using the provided configuration.
+
+    Determines the appropriate quantizer based on the type of quant_config and
+    converts the model accordingly.
+
+    Parameters:
+        model: The prepared PyTorch model.
+        inplace (bool, optional): If true, the model will be modified in place;
+                                   otherwise, a new prepared model is returned.
+
+    Returns:
+        The quantized model.
+    """
+    # Get quantizer first before calling deepcopy that does not copy attributes properly.
+    if hasattr(model, QUANTIZER_ATTRIBUTE_NAME):
+        quantizer = getattr(model, QUANTIZER_ATTRIBUTE_NAME)
+        delattr(model, QUANTIZER_ATTRIBUTE_NAME)
+    else:
+        raise RuntimeError("Call prepare() function first.")
+
+    if isinstance(quantizer, PT2EQuantizer) and inplace:
+        raise RuntimeError(
+            "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
+        )
+
+    model = model if inplace else copy.deepcopy(model)
+
+    assert isinstance(quantizer, BaseQuantizer)
+    model = quantizer.convert(model)
+
+    return model

tico/experimental/quantization/quantizer.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+import torch
+
+from tico.experimental.quantization.config import BaseConfig
+
+
+class BaseQuantizer(ABC):
+    """
+    Abstract base class for quantizers that apply a quantization algorithm to a target model.
+    """
+
+    def __init__(self, config: BaseConfig):
+        """
+        Initialize the quantizer with the given configuration.
+
+        Parameters:
+            config (BaseConfig): Quantization configuration parameters.
+        """
+        self.config = config
+
+    @abstractmethod
+    def prepare(
+        self,
+        model: torch.nn.Module,
+        args: Optional[Any] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Prepare the given model for quantization based on the provided algorithm-specific
+         configuration. This involves setting up necessary observers or hooks, and may
+        optionally use example inputs—particularly useful for activation quantization.
+
+        Parameters:
+            model: The target PyTorch model.
+            args (Any, optional): Positional example inputs required for activation quantization.
+            kwargs (Dict[str, Any], optional): Keyword example inputs required for activation quantization.
+
+        Returns:
+            The prepared model.
+        """
+        pass
+
+    @abstractmethod
+    def convert(self, model):
+        """
+        Convert the prepared (or calibrated) model into its quantized form. This function leverages
+         the statistics collected during calibration to perform the quantization transformation.
+
+        Parameters:
+            model: The prepared PyTorch model.
+
+        Returns:
+            The quantized model.
+        """
+        pass

tico/interpreter/__init__.py

@@ -0,0 +1 @@
+# DO NOT REMOVE THIS FILE

tico/interpreter/infer.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import numpy as np
+import torch
+from circle_schema import circle
+
+from tico.interpreter.interpreter import Interpreter
+from tico.serialize.circle_mapping import np_dtype_from_circle_dtype, to_circle_dtype
+
+
+def preprocess_inputs(inputs: Any):
+    """
+    Preprocess user inputs for circle inference.
+
+    1. None inputs are ignored.
+    2. A list/tuple input is flatten when a torch module is exported.
+      e.g. inputs = (torch.Tensor, [2,3,4]) -> inputs = (torch.Tensor, 2, 3, 4)
+    """
+    l = []
+    for value in inputs:
+        if value == None:
+            continue
+        if isinstance(value, (tuple, list)):
+            for val in value:
+                l.append(val)
+        else:
+            l.append(value)
+    # Check if it is a list of a list.
+    if any(isinstance(item, (tuple, list)) for item in l):
+        l = preprocess_inputs(l)
+    return tuple(l)
+
+
+def infer(circle_binary: bytes, *args: Any, **kwargs: Any) -> Any:
+    # When converting a model, it is assumed that the order of keyword arguments is maintained.
+    user_inputs = args + tuple(kwargs.values())
+    user_inputs = preprocess_inputs(user_inputs)
+    # Cast them to torch.Tensor to make it simple.
+    user_inputs = tuple(
+        torch.tensor(user_input) if type(user_input) != torch.Tensor else user_input
+        for user_input in user_inputs
+    )
+
+    # Get input spec from circle binary.
+    model = circle.Model.Model.GetRootAsModel(circle_binary, 0)
+    assert model.SubgraphsLength() == 1
+    graph = model.Subgraphs(0)
+    model_input_tensors = [
+        graph.Tensors(graph.Inputs(o)) for o in range(graph.InputsLength())
+    ]
+    model_input_shapes_np = [t.ShapeAsNumpy() for t in model_input_tensors]
+    model_input_types_cm = [t.Type() for t in model_input_tensors]
+
+    # Check if given inputs' dtype and shape from users match the inputs' from model binary.
+    if len(model_input_shapes_np) != len(user_inputs):
+        raise RuntimeError(
+            f"Mismatch input length: input({len(user_inputs)}) != circle model({len(model_input_shapes_np)})"
+        )
+    for input_idx, user_input in enumerate(user_inputs):
+        # Shape
+        if list(user_input.shape) != list(model_input_shapes_np[input_idx]):
+            raise RuntimeError(
+                f"Mismatch input {input_idx} shape : input({user_input.shape}) != circle model({model_input_shapes_np[input_idx]})"
+            )
+        # Data type
+        user_input_type_cm = to_circle_dtype(user_input.dtype)
+        if user_input_type_cm != model_input_types_cm[input_idx]:
+            raise RuntimeError(
+                f"Mismatch input {input_idx} data type : input({user_input_type_cm}) != circle model({model_input_types_cm[input_idx]})"
+            )
+
+    # Initialize interpreter
+    intp = Interpreter(circle_binary)
+
+    # Set input
+    for input_idx, user_input in enumerate(user_inputs):
+        intp.writeInputTensor(input_idx, user_input)
+
+    # Interpret
+    intp.interpret()
+
+    # Retrieve outputs' dtype and shape from circle model
+    model_output_tensors = [
+        graph.Tensors(graph.Outputs(o)) for o in range(graph.OutputsLength())
+    ]
+    model_output_shapes_np = [t.ShapeAsNumpy() for t in model_output_tensors]
+    model_output_types_cm = [t.Type() for t in model_output_tensors]
+
+    output = []
+    # Get output
+    for output_idx in range(len(model_output_tensors)):
+        result: np.ndarray = np.empty(
+            model_output_shapes_np[output_idx],
+            dtype=np_dtype_from_circle_dtype(model_output_types_cm[output_idx]),
+        )
+        intp.readOutputTensor(output_idx, result)
+        output.append(result)
+
+    if len(output) == 1:
+        return output[0]
+    else:
+        return output