tico 0.1.0.dev251106__py3-none-any.whl → 0.2.0.dev260122__py3-none-any.whl

tico/quantization/algorithm/pt2e/quantizer.py

@@ -1,81 +0,0 @@
-# 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, Dict, Optional
-
-import torch
-
-from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
-
-from tico.quantization.algorithm.pt2e.annotation.annotator import (
-    get_asymmetric_quantization_config,
-    PT2EAnnotator,
-)
-from tico.quantization.config.pt2e import PT2EConfig
-from tico.quantization.quantizer import BaseQuantizer
-from tico.quantization.quantizer_registry import register_quantizer
-
-
-@register_quantizer(PT2EConfig)
-class PT2EQuantizer(BaseQuantizer):
-    """
-    Quantizer for applying pytorch 2.0 export quantization (typically for activation quantization).
-    """
-
-    def prepare(
-        self,
-        model: torch.nn.Module,
-        args: Optional[Any] = None,
-        kwargs: Optional[Dict[str, Any]] = None,
-    ):
-        """
-        Prepare the model for pt2e quantization.
-
-        Registers activation observers using the provided example inputs.
-
-        Parameters:
-            model: The target PyTorch model.
-            args: Positional example inputs required for capturing graph.
-            kwargs: Keyword example inputs required for capturing graph.
-
-        Returns:
-            The model prepared for pt2e quantization.
-        """
-        # Program capture
-        assert isinstance(args, tuple)
-        model = torch.export.export_for_training(
-            model, args=args, kwargs=kwargs
-        ).module()
-        quantizer = PT2EAnnotator()
-        quantizer = quantizer.set_global(get_asymmetric_quantization_config())
-
-        # Register observers in each nodes
-        assert isinstance(model, torch.fx.GraphModule)
-        model = prepare_pt2e(model, quantizer)
-
-        return model
-
-    def convert(self, model: torch.fx.GraphModule):
-        """
-        Convert the prepared model to its pt2e quantized version.
-
-        Applies the pt2e quantization on activations based on the collected statistics.
-
-        Parameters:
-            model: The prepared PyTorch model.
-
-        Returns:
-            The quantized model.
-        """
-        return convert_pt2e(model)

tico/quantization/algorithm/pt2e/transformation/__init__.py

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

tico/quantization/algorithm/pt2e/transformation/convert_scalars_to_attrs.py

@@ -1,58 +0,0 @@
-# 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 torch
-from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix
-
-
-def convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
-    """
-    Convert scalar values in the graph to `get_attr` nodes.
-
-    This function identifies scalar constants in the graph and transforms them
-     into `get_attr` nodes to ensure compatibility with quantization workflows.
-    """
-    for n in model.graph.nodes:
-        if n.op != "call_function" or n.target not in [
-            # The operators that have scalar parameters.
-            torch.ops.aten.add.Tensor,
-        ]:
-            continue
-        args = list(n.args)
-        new_args = []
-        for arg in args:
-            if isinstance(arg, torch.fx.Node):
-                new_args.append(arg)
-                continue
-
-            assert isinstance(arg, float)
-            prefix = "_tensor_constant_"
-            get_new_attr_name = get_new_attr_name_with_prefix(prefix)
-            tensor_constant_name = get_new_attr_name(model)
-            float_tensor = torch.tensor(float(arg))
-            model.register_buffer(tensor_constant_name, float_tensor)
-
-            fake_mode = n.meta["val"].fake_mode
-            with model.graph.inserting_before(n):
-                get_attr_node = model.graph.create_node(
-                    "get_attr", tensor_constant_name, (), {}
-                )
-                get_attr_node.meta["val"] = fake_mode.from_tensor(
-                    float_tensor, static_shapes=True
-                )
-                new_args.append(get_attr_node)
-        n.args = tuple(new_args)
-    model.recompile()
-
-    return model

tico/quantization/algorithm/pt2e/utils.py

@@ -1,135 +0,0 @@
-# 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 Callable, List, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-from torch.ao.quantization.quantizer import QuantizationSpec
-from torch.ao.quantization.quantizer.utils import _get_module_name_filter
-
-from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
-
-
-def get_module_type_filter(tp: Callable):
-    """
-    Get the module_type_filter function for a given module type.
-
-    The filter accepts a node and checks if the node comes from a module
-     that has certain module type.
-
-    For example:
-        node: linear_op = call_function[...](...)  # comes from a module with type Block -> Sub -> Linear
-
-
-    >> module_type_filter = get_module_type_filter(Sub)  # submodule with type `Sub`, under the `Block` submodule
-    >> print(module_type_filter(node))
-    True  # the node is from the submodule `Sub`
-    """
-
-    tp_str = tp.__module__ + "." + tp.__qualname__
-
-    def module_type_filter(n: torch.fx.Node) -> bool:
-        # example: {
-        #     'L__self___sub': ("L['self'].sub", <class '....Sub'>),
-        #     'L__self___sub_linear': ("L['self'].sub.linear", <class 'torch.nn.modules.linear.Linear'>)
-        # }
-        nn_module_stack = n.meta.get("nn_module_stack", {})
-        types = []
-        for _, t in nn_module_stack.values():
-            # export() returns str, but older APIs (e.g. capture_pre_autograd_graph)
-            # return type. Handle both cases.
-            if isinstance(t, type):
-                t = t.__module__ + "." + t.__qualname__
-            types.append(t)
-        return tp_str in types
-
-    return module_type_filter
-
-
-def get_not_module_type_or_name_filter(
-    tp_list: List[Callable], module_name_list: List[str]
-) -> Callable[[torch.fx.Node], bool]:
-    module_type_filters = [get_module_type_filter(tp) for tp in tp_list]
-    module_name_list_filters = [_get_module_name_filter(m) for m in module_name_list]
-
-    def not_module_type_or_name_filter(n: torch.fx.Node) -> bool:
-        return not any(f(n) for f in module_type_filters + module_name_list_filters)
-
-    return not_module_type_or_name_filter
-
-
-def get_input_act_qspec(quantization_config: Optional[QuantizationConfig]):
-    if quantization_config is None:
-        return None
-    if quantization_config.input_activation is None:
-        return None
-    quantization_spec: QuantizationSpec = quantization_config.input_activation
-    assert quantization_spec.qscheme in [
-        torch.per_tensor_affine,
-    ]
-    return quantization_spec
-
-
-def get_output_act_qspec(quantization_config: Optional[QuantizationConfig]):
-    if quantization_config is None:
-        return None
-    if quantization_config.output_activation is None:
-        return None
-    quantization_spec: QuantizationSpec = quantization_config.output_activation
-    assert quantization_spec.qscheme in [
-        torch.per_tensor_affine,
-    ]
-    return quantization_spec
-
-
-def get_weight_qspec(quantization_config: Optional[QuantizationConfig]):
-    if quantization_config is None:
-        return None
-    if quantization_config.weight is None:
-        return None
-    quantization_spec: QuantizationSpec = quantization_config.weight
-    if quantization_spec.qscheme not in [
-        torch.per_tensor_affine,
-        torch.per_channel_affine,
-    ]:
-        raise ValueError(
-            f"Unsupported quantization_spec {quantization_spec} for weight"
-        )
-    return quantization_spec
-
-
-def get_bias_qspec(quantization_config: Optional[QuantizationConfig]):
-    if quantization_config is None:
-        return None
-    if quantization_config.bias is None:
-        return None
-    quantization_spec: QuantizationSpec = quantization_config.bias
-    return quantization_spec
-
-
-def is_annotated(nodes: List[torch.fx.Node] | torch.fx.Node):
-    """
-    Check if any of the node in the given list is annotated.
-    """
-    annotated = False
-    if isinstance(nodes, torch.fx.Node):
-        nodes = [nodes]
-    for node in nodes:
-        annotated = annotated or (
-            "quantization_annotation" in node.meta
-            and node.meta["quantization_annotation"]._annotated
-        )
-    return annotated

tico/serialize/operators/op_copy.py

@@ -1,187 +0,0 @@
-# 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 Dict, List, Optional, TYPE_CHECKING, Union
-
-if TYPE_CHECKING:
-    import torch._ops
-    import torch.fx
-import torch
-from circle_schema import circle
-
-from tico.serialize.circle_graph import CircleSubgraph
-from tico.serialize.operators.hashable_opcode import OpCode
-from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
-from tico.serialize.operators.utils import create_builtin_operator, get_op_index
-from tico.utils.errors import NotYetSupportedError
-from tico.utils.validate_args_kwargs import CopyArgs
-
-
-@register_node_visitor
-class CopyVisitor(NodeVisitor):
-    """
-    NOTE `torch.Tensor.copy_`'s behavior matches with `Reshape` of CIRCLE.
-    - because `torch.Tensor.copy_` is a in-place operator, so `dst` is converted to `Shape` of CIRCLE.
-    - after that, `dst` converted to `Shape` is connected to shape of `Reshape`.
-    - `src` is connected to tensor of `Reshape`.
-    - if `dst` is not converted to `Shape`.
-      [dst]      [src]
-                   |
-                [Reshape]
-    - if `dst` is converted to `Shape`.
-      [dst]      [src]
-        |          |
-      [Shape]      |
-        \         /
-         [Reshape]
-    """
-
-    target: List[torch._ops.OpOverload] = [torch.ops.aten.copy.default]
-
-    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
-        super().__init__(op_codes, graph)
-
-    def check_to_do_broadcast(
-        self,
-        dst: List[int],
-        dst_sig: Optional[List[int]],
-        src: List[int],
-        src_sig: Optional[List[int]],
-    ) -> bool:
-        assert dst_sig is None
-        assert src_sig is None
-        return dst != src
-
-    def define_broadcast_to_node(
-        self,
-        inputs: List[Union[circle.Tensor.TensorT, torch.Tensor]],
-        outputs: List[circle.Tensor.TensorT],
-    ) -> circle.Operator.OperatorT:
-        op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.BROADCAST_TO, self._op_codes
-        )
-        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
-        operator.builtinOptionsType = (
-            circle.BuiltinOptions.BuiltinOptions.BroadcastToOptions
-        )
-
-        option = circle.BroadcastToOptions.BroadcastToOptionsT()
-        operator.builtinOptions = option
-        return operator
-
-    def define_shape_node(
-        self, inputs: List[torch.fx.Node], outputs: List[circle.Tensor.TensorT]
-    ) -> circle.Operator.OperatorT:
-        op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.SHAPE, self._op_codes
-        )
-        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
-        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.ShapeOptions
-
-        option = circle.ShapeOptions.ShapeOptionsT()
-        option.outType = circle.TensorType.TensorType.INT32
-        operator.builtinOptions = option
-        return operator
-
-    def define_node(
-        self,
-        node: torch.fx.Node,
-    ) -> circle.Operator.OperatorT:
-        if len(node.args) == 3:
-            raise NotYetSupportedError("'non_blocking' is not supported yet.")
-
-        assert len(node.args) == 2, len(node.args)
-
-        args = CopyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-        dst = args.dst
-        src = args.src
-
-        # To connect 'dst' to Reshape node in the graph, 'dst' must be converted to Shape op.
-        dst_tensor: circle.Tensor.TensorT = self.graph.get_tensor(dst)
-        dst_shape: List[int] = dst_tensor.shape
-        dst_shape_signature: Optional[List[int]] = dst_tensor.shapeSignature
-
-        if dst_shape_signature is not None:
-            # TODO: support dynamic shape
-            raise NotYetSupportedError("Dynamic shape is not supported yet.")
-
-        dst_shape_tensor = torch.as_tensor(dst_shape, dtype=torch.int32)
-
-        dst_shape_shape = [len(dst_shape)]
-        dst_name: str = dst.name
-
-        shape_output = self.graph.add_tensor_from_scratch(
-            prefix=f"{dst_name}_shape_output",
-            shape=dst_shape_shape,
-            shape_signature=None,
-            dtype=circle.TensorType.TensorType.INT32,
-            source_node=node,
-        )
-
-        shape_operator = self.define_shape_node([dst], [shape_output])
-        self.graph.add_operator(shape_operator)
-
-        src_tensor: circle.Tensor.TensorT = self.graph.get_tensor(src)
-        src_shape: List[int] = src_tensor.shape
-        src_shape_signature: Optional[List[int]] = src_tensor.shapeSignature
-
-        if src_shape_signature is not None:
-            # TODO: support dynamic shape
-            raise NotYetSupportedError("Dynamic shape is not supported yet.")
-
-        # The src tensor must be broadcastable with the dst tensor.
-        do_broadcast = self.check_to_do_broadcast(
-            dst_shape, dst_shape_signature, src_shape, src_shape_signature
-        )
-        if do_broadcast:
-            # create braodcastTo output tensor
-            src_name: str = src.name
-            src_type: int = src_tensor.type
-
-            broadcast_to_output: circle.Tensor.TensorT = (
-                self.graph.add_tensor_from_scratch(
-                    prefix=f"{src_name}_broadcast_to_output",
-                    shape=dst_shape,
-                    shape_signature=dst_shape_signature,
-                    dtype=src_type,
-                    source_node=node,
-                )
-            )
-
-            broadcast_to_operator: circle.Operator.OperatorT = (
-                self.define_broadcast_to_node(
-                    [src_tensor, dst_shape_tensor], [broadcast_to_output]
-                )
-            )
-            self.graph.add_operator(broadcast_to_operator)
-            inputs: List = [broadcast_to_output, shape_output]
-        else:
-            inputs = [src, shape_output]
-
-        outputs = [node]
-        op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.RESHAPE, self._op_codes
-        )
-
-        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
-
-        # Op-specific option
-        operator.builtinOptionsType = (
-            circle.BuiltinOptions.BuiltinOptions.ReshapeOptions
-        )
-        option = circle.ReshapeOptions.ReshapeOptionsT()
-        option.newShape = dst_shape
-
-        operator.builtinOptions = option
-        return operator