tico 0.1.0.dev250411__py3-none-any.whl

tico/experimental/quantization/evaluation/utils.py

@@ -0,0 +1,185 @@
+# 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, List, Tuple
+
+import numpy as np
+import torch
+
+from circle_schema import circle
+from numpy.typing import DTypeLike
+
+from tico.utils import logging
+from tico.utils.model import CircleModel
+
+
+def quantize(
+    data: np.ndarray, scale: float, zero_point: int, dtype: DTypeLike
+) -> np.ndarray:
+    """
+    Quantize the given data using the specified scale, zero point, and data type.
+    This function takes input data and applies quantization using the formula:
+        round(data / scale) + zero_point
+    The result is clamped to the range of the specified data type.
+    """
+    logger = logging.getLogger(__name__)
+    dtype = np.dtype(dtype)
+    assert dtype == np.uint8 or dtype == np.int16, f"Invalid dtype: {dtype}"
+    if dtype == np.int16:
+        assert zero_point == 0
+
+    # Convert input to Numpy array if necessary
+    if not isinstance(data, np.ndarray):
+        data = np.array(data)
+    # Perfrom quantization
+    if not scale:
+        logger.warn("WARNING: scale value is 0. 1e-7 will be used instead.")
+        scale = 1e-7
+    rescaled = np.round(data / scale) + zero_point
+    # Clamp the values
+    clipped = np.clip(rescaled, np.iinfo(dtype).min, np.iinfo(dtype).max)
+    # Convert to the specified dtype
+    return clipped.astype(dtype)
+
+
+def dequantize(
+    data: np.ndarray, scale: float, zero_point: int, dtype: DTypeLike
+) -> np.ndarray:
+    """
+    Dequantize the given quantized data using the specified scale and zero point.
+    This function reverses the quantization process by applying the formula:
+        (quantized_value - zero_point) * scale
+    """
+    dtype = np.dtype(dtype)
+    assert dtype == np.uint8 or dtype == np.int16, f"Invalid dtype: {dtype}"
+    if dtype == np.int16:
+        assert zero_point == 0
+
+    # Convert input to Numpy array if necessary
+    if not isinstance(data, np.ndarray):
+        data = np.array(data)
+    # Perform dequantization
+    ret = (data.astype(np.float32) - zero_point) * scale
+    # np.float32 * np.int64 = np.float64
+    return ret.astype(np.float32)
+
+
+def get_graph_input_output(
+    circle_model: CircleModel,
+) -> Tuple[List[circle.Tensor.Tensor], List[circle.Tensor.Tensor]]:
+    """
+    Retrieve the inputs and the outputs from the circle model, and return them
+     as two lists.
+    """
+    circle_buf: bytes = circle_model.circle_binary
+    circle_fb: circle.Model.Model = circle.Model.Model.GetRootAs(circle_buf, 0)
+    assert circle_fb.SubgraphsLength() == 1, "Only support single graph."
+    circle_graph = circle_fb.Subgraphs(0)
+    circle_inputs: List[circle.Tensor.Tensor] = [
+        circle_graph.Tensors(circle_graph.Inputs(i))
+        for i in range(circle_graph.InputsLength())
+    ]
+    circle_outputs: List[circle.Tensor.Tensor] = [
+        circle_graph.Tensors(circle_graph.Outputs(o))
+        for o in range(circle_graph.OutputsLength())
+    ]
+
+    return circle_inputs, circle_outputs
+
+
+def find_invalid_types(
+    input: List[torch.Tensor] | List[np.ndarray], allowed_types: List
+) -> List:
+    """
+    Indentifies the types of items in a list that are not allowed and removes duplicates.
+
+    Parameters
+    -----------
+        input
+            List of itmes to check.
+        allowed_types
+            List of allowed types (e.g. [int, str])
+    Returns
+    --------
+        A list of unique types that are not allowed in the input list.
+    """
+    # Use set comprehension for uniqueness
+    invalid_types = {
+        type(item) for item in input if not isinstance(item, tuple(allowed_types))
+    }
+    return list(invalid_types)
+
+
+def plot_two_outputs(x_values: torch.Tensor, y_values: torch.Tensor):
+    """
+    Plot two values on a 2D graph using plotext.
+
+    Returns
+    --------
+        A figure built from plotext.
+
+    Example
+    --------
+        >>> x_values = torch.tensor([1, 2, 3, 4, 5])
+        >>> y_values = torch.tensor([10, 20, 30, 40, 50])
+        >>> fig = plot_two_outputs(x_values, y_values)
+        >>> print(fig)
+    """
+    x_np = x_values.numpy().reshape(-1)
+    y_np = y_values.numpy().reshape(-1)
+    min_value = min([x_np.min(), y_np.min()])
+    max_value = max([x_np.max(), y_np.max()])
+
+    interval = max_value - min_value
+    interval = 1.0 if interval == 0.0 else interval  # Avoid zero interval
+
+    # Enlarge axis
+    axis_min = min_value - interval * 0.05
+    axis_max = max_value + interval * 0.05
+
+    import plotext as plt
+
+    plt.clear_data()
+    plt.xlim(axis_min, axis_max)
+    plt.ylim(axis_min, axis_max)
+    plt.plotsize(width=50, height=25)
+    plt.scatter(x_np, y_np, marker="dot")
+    plt.theme("clear")
+
+    return plt.build()
+
+
+def ensure_list(inputs: Any | Tuple[Any] | List[Any]) -> List[Any]:
+    """
+    Ensures that the given inputs is converted into a list.
+
+    - If the input is a single element, it wraps it into a list.
+    - If the input is a tuple, it converts the tuple to a list.
+    - If the input is already a list, it returns the input unchanged.
+
+    Example
+    --------
+        >>> ensure_list(42)
+        >>> [42]
+        >>> ensure_list((1, 2, 3))
+        >>> [1, 2, 3]
+        >>> ensure_list([4, 5, 6])
+        >>> [4, 5, 6]
+    """
+    if isinstance(inputs, list):
+        return inputs
+    elif isinstance(inputs, tuple):
+        return list(inputs)
+    else:
+        return [inputs]

tico/experimental/quantization/passes/__init__.py

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

tico/experimental/quantization/passes/fold_quant_ops.py

@@ -0,0 +1,97 @@
+# 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 torch
+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 (
+    DequantizePerTensorArgs,
+    QuantizePerTensorArgs,
+)
+
+
+@trace_graph_diff_on_pass
+class FoldQuantOps(PassBase):
+    """
+    This pass folds (Q - DQ) pattern to previous op. After quantization from torch, activation ops
+     have (op - Q - DQ) pattern.
+
+    To export quantized circle, this pass removes (Q - DQ) nodes and saves those quantization info
+     to previous op's metadata.
+
+    [BEFORE]
+      op (float) - Quantize - Dequantize - (float)
+
+    [AFTER]
+      op (float with meta[QPARAM_KEY])
+    """
+
+    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 dq.op != "call_function":
+                continue
+            if (
+                dq.target
+                != torch.ops.quantized_decomposed.dequantize_per_tensor.default
+            ):
+                continue
+            dq_args = DequantizePerTensorArgs(*dq.args, **dq.kwargs)
+
+            q = dq_args.input
+            if q.target != torch.ops.quantized_decomposed.quantize_per_tensor.default:
+                continue
+            q_args = QuantizePerTensorArgs(*q.args, **q.kwargs)  # type: ignore[arg-type]
+            op = q_args.tensor
+
+            # Check if Q and DQ have same quant param
+            if q_args.scale != dq_args.scale:
+                continue
+            if q_args.zero_p != dq_args.zero_point:
+                continue
+            if q_args.dtype != dq_args.dtype:
+                continue
+
+            if QPARAM_KEY not in op.meta:
+                qparam = QuantParam()
+                qparam.scale = [q_args.scale]
+                qparam.zero_point = [q_args.zero_p]
+                assert "val" in q.meta and hasattr(q.meta["val"], "dtype")
+                qparam.dtype = to_qparam_dtype(q.meta["val"].dtype)
+                op.meta[QPARAM_KEY] = qparam
+
+            dq.replace_all_uses_with(op, propagate_meta=False)
+
+            logger.debug(f"{q.name} and {dq.name} are folded to {op.name}.")
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py

@@ -0,0 +1,289 @@
+# 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, QuantParam
+from tico.utils import logging
+from tico.utils.errors import NotYetSupportedError
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import quant_min_max, set_new_meta_val
+from tico.utils.validate_args_kwargs import (
+    BmmArgs,
+    LinearArgs,
+    MulTensorArgs,
+    PermuteArgs,
+)
+
+
+def qparam_dtype(node: torch.fx.Node) -> str:
+    assert QPARAM_KEY in node.meta
+    return node.meta[QPARAM_KEY].dtype
+
+
+# Convert i16 qparam to u8 qparam
+# scale and zero_point are inferred from i16 qparam
+def _i16_to_u8(qparam: QuantParam) -> QuantParam:
+    # Assume per-tensor quantization
+    assert qparam.scale is not None and len(qparam.scale) == 1
+    assert qparam.dtype == "int16"
+
+    s16_scale = qparam.scale[0]
+    max_ = s16_scale * 32767  # numeric_limits<int16>
+    min_ = -max_
+
+    u8_scale = (max_ - min_) / 255
+    u8_zerop = round(-min_ / u8_scale)
+
+    new_qparam = QuantParam()
+    new_qparam.scale = [u8_scale]
+    new_qparam.zero_point = [u8_zerop]
+    new_qparam.dtype = "uint8"
+
+    return new_qparam
+
+
+# Convert u8 qparam to i16 qparam
+# scale is inferred from u8 qparam
+def _u8_to_i16(qparam: QuantParam) -> QuantParam:
+    # Assume per-tensor quantization
+    assert qparam.scale is not None and len(qparam.scale) == 1
+    assert qparam.zero_point is not None and len(qparam.zero_point) == 1
+    assert qparam.dtype == "uint8"
+
+    u8_scale = qparam.scale[0]
+    u8_zerop = qparam.zero_point[0]
+    max_ = u8_scale * (255 - u8_zerop)
+    min_ = u8_scale * (-u8_zerop)
+
+    abs_max = max([max_, min_], key=abs)
+    s16_scale = abs_max / 32767
+    s16_zerop = 0
+
+    new_qparam = QuantParam()
+    new_qparam.scale = [s16_scale]
+    new_qparam.zero_point = [s16_zerop]
+    new_qparam.dtype = "int16"
+
+    return new_qparam
+
+
+@trace_graph_diff_on_pass
+class InsertQuantizeOnDtypeMismatch(PassBase):
+    """
+    Insert quantize Op in the operators where circle's type inference is violated.
+    Example. FullyConnected
+    [BEFORE]
+      Op (uint8) -  aten.linear.default (int16)
+    [AFTER]
+      Op (uint8) -  aten.linear.default (uint8) - quantized_decomposed.quantize_per_tensor.default (int16)
+    Why is this pass necessary?
+    - For some operators, circle's type inference pass overwrites the input's dtype to
+    the output's dtype. For the above example, fully-connected layer (aten.linear.default)'s
+    output dtype (int16) is updated to the input dtype (uint8), which breaks the semantics.
+    This problem can occur in the tools (ex: circle2circle) that automatically apply type inference.
+    - To resolve the issue, we insert quantize operators not to violate circle's type inference logic.
+    """
+
+    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
+
+        def _insert_quantize_op_before(node, inp):
+            qparam: QuantParam = node.meta[QPARAM_KEY]
+            assert qparam.scale is not None
+            assert qparam.zero_point is not None
+            scale = qparam.scale[0]
+            zerop = qparam.zero_point[0]
+            min_, max_ = quant_min_max(qparam.dtype)
+            dtype = getattr(torch, qparam.dtype)
+
+            with graph.inserting_before(node):
+                q_args = (inp, scale, zerop, min_, max_, dtype)
+                quantize = graph.call_function(
+                    torch.ops.quantized_decomposed.quantize_per_tensor.default,
+                    args=q_args,
+                )
+                quantize.meta[QPARAM_KEY] = copy.deepcopy(qparam)
+                set_new_meta_val(quantize)
+
+            node.replace_input_with(inp, quantize)
+
+            return quantize
+
+        def _insert_quantize_op_after(node):
+            qparam: QuantParam = node.meta[QPARAM_KEY]
+            assert qparam.scale is not None
+            assert qparam.zero_point is not None
+            scale = qparam.scale[0]
+            zerop = qparam.zero_point[0]
+            min_, max_ = quant_min_max(qparam.dtype)
+            dtype = getattr(torch, qparam.dtype)
+            with graph.inserting_after(node):
+                q_args = (node, scale, zerop, min_, max_, dtype)
+                quantize = graph.call_function(
+                    torch.ops.quantized_decomposed.quantize_per_tensor.default,
+                    args=q_args,
+                )
+
+            node.replace_all_uses_with(quantize, propagate_meta=True)
+            quantize.replace_input_with(quantize, node)
+
+            quantize.meta[QPARAM_KEY] = copy.deepcopy(qparam)
+
+            return quantize
+
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.linear.default:
+                lin_args = LinearArgs(*node.args, **node.kwargs)
+                inp = lin_args.input
+
+                if QPARAM_KEY not in inp.meta:
+                    continue
+
+                if QPARAM_KEY not in node.meta:
+                    continue
+
+                if qparam_dtype(inp) == qparam_dtype(node):
+                    continue
+
+                if qparam_dtype(inp) == "uint8" and qparam_dtype(node) == "int16":
+                    quantize = _insert_quantize_op_after(node)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+
+                    # Update node's qparam from i16 to u8
+                    # NOTE This would severely degrade accuracy. It is
+                    # important to mitigate this accuracy drop in backend.
+                    node.meta[QPARAM_KEY] = _i16_to_u8(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted after {node.name}."
+                    )
+                else:
+                    raise NotYetSupportedError("Unsupported dtype")
+
+            elif node.target == torch.ops.aten.mul.Tensor:
+                mul_args = MulTensorArgs(*node.args, **node.kwargs)
+                x = mul_args.input
+                y = mul_args.other
+
+                if not isinstance(x, torch.fx.Node):
+                    continue
+                if not isinstance(y, torch.fx.Node):
+                    continue
+
+                if QPARAM_KEY not in x.meta:
+                    continue
+                if QPARAM_KEY not in y.meta:
+                    continue
+                if QPARAM_KEY not in node.meta:
+                    continue
+
+                if qparam_dtype(x) == qparam_dtype(node):
+                    continue
+
+                if qparam_dtype(x) == "int16" and qparam_dtype(node) == "uint8":
+                    quantize = _insert_quantize_op_after(node)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+                    node.meta[QPARAM_KEY] = _u8_to_i16(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted after {node.name}."
+                    )
+                else:
+                    raise NotYetSupportedError("Unsupported dtype")
+
+            elif node.target == torch.ops.aten.bmm.default:
+                bmm_args = BmmArgs(*node.args, **node.kwargs)
+                x = bmm_args.input
+                y = bmm_args.mat2
+
+                if QPARAM_KEY not in x.meta:
+                    continue
+                if QPARAM_KEY not in y.meta:
+                    continue
+                if QPARAM_KEY not in node.meta:
+                    continue
+
+                if qparam_dtype(x) == qparam_dtype(node):
+                    continue
+
+                if qparam_dtype(x) == "int16" and qparam_dtype(node) == "uint8":
+                    quantize = _insert_quantize_op_after(node)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+                    node.meta[QPARAM_KEY] = _u8_to_i16(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted after {node.name}."
+                    )
+                else:
+                    raise NotYetSupportedError("Unsupported dtype")
+
+            elif node.target == torch.ops.aten.permute.default:
+                per_args = PermuteArgs(*node.args, **node.kwargs)
+                inp = per_args.input
+
+                if QPARAM_KEY not in inp.meta:
+                    continue
+
+                if QPARAM_KEY not in node.meta:
+                    continue
+
+                if qparam_dtype(inp) == qparam_dtype(node):
+                    continue
+
+                if qparam_dtype(inp) == "int16" and qparam_dtype(node) == "uint8":
+                    # A new Quantize Op (s16 to u8) is inserted before (not after)
+                    # permute Op to reduce tensor size ealier
+                    quantize = _insert_quantize_op_before(node, inp)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+                    node.meta[QPARAM_KEY] = _u8_to_i16(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted before {node.name}."
+                    )
+                elif qparam_dtype(inp) == "uint8" and qparam_dtype(node) == "int16":
+                    quantize = _insert_quantize_op_after(node)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+                    node.meta[QPARAM_KEY] = _i16_to_u8(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted after {node.name}."
+                    )
+                else:
+                    raise NotYetSupportedError("Unsupported dtype")
+
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/propagate_qparam_backward.py

@@ -0,0 +1,91 @@
+# 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, PermuteArgs, ReshapeArgs
+
+
+@trace_graph_diff_on_pass
+class PropagateQParamBackward(PassBase):
+    """
+    This pass propagates quantization parameters backward.
+
+    BEFORE)
+
+    node -> reshape (with meta[QPARAM_KEY])
+
+    AFTER)
+
+    node (with meta[QPARAM_KEY]) -> reshape (with meta[QPARAM_KEY])
+    """
+
+    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
+
+        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}.")
+
+        # Do reverse-order traversal for backward propagation
+        for node in reversed(graph.nodes):
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.cat.default:
+                concat_args = CatArgs(*node.args, **node.kwargs)
+                concat_inputs = concat_args.tensors
+
+                for concat_input in concat_inputs:
+                    _propagate_qparam_if_possible(node, concat_input)
+            elif node.target == torch.ops.aten.reshape.default:
+                args = ReshapeArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(node, args.input)
+            elif node.target == torch.ops.aten.permute.default:
+                permute_args = PermuteArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(node, permute_args.input)
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)