tico 0.1.0__py3-none-any.whl

tico/experimental/quantization/algorithm/pt2e/annotation/op/conv2d.py

@@ -0,0 +1,92 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.ao.quantization.quantizer import DerivedQuantizationSpec
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import Conv2DArgs
+
+
+@annot_spec.register_annotator(
+    [torch.ops.aten.conv2d.default, torch.ops.aten.conv2d.padding]
+)
+def _annotate_conv2d(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    for node in gm.graph.nodes:
+        if node.op != "call_function" or node.target not in [
+            torch.ops.aten.conv2d.default,
+            torch.ops.aten.conv2d.padding,
+        ]:
+            continue
+        if filter_fn and not filter_fn(node):
+            continue
+        if quant_utils.is_annotated(node):
+            continue
+
+        args = Conv2DArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input_ = args.input
+        weight = args.weight
+        bias = args.bias
+
+        input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+        weight_qspec = quant_utils.get_weight_qspec(quantization_config)
+        annot_utils.annotate_input_qspec_map(node, input_, input_act_qspec)
+        annot_utils.annotate_input_qspec_map(node, weight, weight_qspec)
+        nodes_to_mark_annotated = [input_, weight, node]
+        if bias:
+
+            def _derive_bias_qparams_from_act_and_weight_qparams(obs_or_fqs):
+                act_scale, _ = obs_or_fqs[0].calculate_qparams()
+                weight_scale, _ = obs_or_fqs[1].calculate_qparams()
+                bias_scale = act_scale * weight_scale
+                bias_zero_point = torch.zeros_like(bias_scale, dtype=torch.int32)
+                return bias_scale, bias_zero_point
+
+            bias_qspec = DerivedQuantizationSpec(
+                derived_from=[
+                    (input_, node),
+                    (weight, node),
+                ],
+                derive_qparams_fn=_derive_bias_qparams_from_act_and_weight_qparams,
+                dtype=torch.int32,
+                quant_min=-(2**31),
+                quant_max=2**31 - 1,
+                qscheme=weight_qspec.qscheme,
+                ch_axis=0 if weight_qspec.qscheme == torch.per_channel_affine else None,
+            )
+            annot_utils.annotate_input_qspec_map(
+                node,
+                bias,
+                bias_qspec,
+            )
+            nodes_to_mark_annotated.append(bias)
+
+        output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+        annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+        annot_utils.mark_nodes_as_annotated(nodes_to_mark_annotated)

tico/experimental/quantization/algorithm/pt2e/annotation/op/div.py

@@ -0,0 +1,57 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import DivTensorArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.div.Tensor])
+def _annotate_div(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.div.Tensor:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = DivTensorArgs(*node.args)  # type: ignore[arg-type]
+    input = args.input
+    other = args.other
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    if isinstance(input, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+    if isinstance(other, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/op/linear.py

@@ -0,0 +1,94 @@
+# 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 DerivedQuantizationSpec
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import LinearArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.linear.default])
+def _annotate_linear(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.linear.default:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = LinearArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+    input_ = args.input
+    weight = args.weight
+    bias = args.bias
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    weight_qspec = quant_utils.get_weight_qspec(quantization_config)
+    bias_qspec = quant_utils.get_bias_qspec(quantization_config)
+
+    annot_utils.annotate_input_qspec_map(
+        node,
+        input_,
+        input_act_qspec,
+    )
+    annot_utils.annotate_input_qspec_map(
+        node,
+        weight,
+        weight_qspec,
+    )
+    nodes_to_mark_annotated = [node, weight]
+    if bias:
+
+        def _derive_bias_qparams_from_act_and_weight_qparams(obs_or_fqs):
+            act_scale, _ = obs_or_fqs[0].calculate_qparams()
+            weight_scale, _ = obs_or_fqs[1].calculate_qparams()
+            bias_scale = act_scale * weight_scale
+            bias_zero_point = torch.zeros_like(bias_scale, dtype=torch.int32)
+            return bias_scale, bias_zero_point
+
+        bias_qspec = DerivedQuantizationSpec(
+            derived_from=[
+                (input_, node),
+                (weight, node),
+            ],
+            derive_qparams_fn=_derive_bias_qparams_from_act_and_weight_qparams,
+            dtype=torch.int32,
+            quant_min=-(2**31),
+            quant_max=2**31 - 1,
+            qscheme=weight_qspec.qscheme,
+            ch_axis=0 if weight_qspec.qscheme == torch.per_channel_affine else None,
+        )
+        annot_utils.annotate_input_qspec_map(
+            node,
+            bias,
+            bias_qspec,
+        )
+        nodes_to_mark_annotated.append(bias)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+    annot_utils.mark_nodes_as_annotated(nodes_to_mark_annotated)

tico/experimental/quantization/algorithm/pt2e/annotation/op/mean.py

@@ -0,0 +1,53 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import MeanDimArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.mean.dim])
+def _annotate_mean(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.mean.dim:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = MeanDimArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+    input = args.input
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/op/mul.py

@@ -0,0 +1,57 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import MulTensorArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.mul.Tensor])
+def _annotate_mul(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.mul.Tensor:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = MulTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+    input = args.input
+    other = args.other
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    if isinstance(input, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+    if isinstance(other, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/op/relu6.py

@@ -0,0 +1,53 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import Relu6Args
+
+
+@annot_spec.register_annotator([torch.ops.aten.relu6.default])
+def _annotate_relu6(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.relu6.default:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = Relu6Args(*node.args, **node.kwargs)  # type: ignore
+    input = args.input
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/op/rsqrt.py

@@ -0,0 +1,53 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import RsqrtArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.rsqrt.default])
+def _annotate_rsqrt(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.rsqrt.default:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = RsqrtArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+    input = args.input
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/op/sub.py

@@ -0,0 +1,57 @@
+# 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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import SubTensorArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.sub.Tensor])
+def _annotate_sub(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.sub.Tensor:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = SubTensorArgs(*node.args)  # type: ignore[arg-type]
+    input = args.input
+    other = args.other
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    if isinstance(input, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+    if isinstance(other, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/spec.py

@@ -0,0 +1,47 @@
+# 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, Dict, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+
+AnnotatorType = Callable[
+    [
+        torch.fx.GraphModule,
+        torch.fx.Node,
+        Optional[QuantizationConfig],
+        Optional[Callable[[torch.fx.Node], bool]],
+    ],
+    None,
+]
+OP_TO_ANNOTATOR: Dict[torch._ops.OpOverload, AnnotatorType] = {}
+OP_TO_SHARE_QUANT_SPEC: List[Callable] = [
+    torch.ops.aten.view_copy.default,
+    torch.ops.aten.view.default,
+]
+
+
+def register_annotator(target: List[torch._ops.OpOverload]):
+    def decorator(annotator: AnnotatorType):
+        for t in target:
+            OP_TO_ANNOTATOR[t] = annotator
+        return annotator
+
+    return decorator

tico/experimental/quantization/algorithm/pt2e/annotation/utils.py

@@ -0,0 +1,88 @@
+# 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
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.ao.quantization.quantizer import (
+    QuantizationAnnotation,
+    SharedQuantizationSpec,
+)
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+
+
+def annotate_input_qspec_map(node: torch.fx.Node, input_node: torch.fx.Node, qspec):
+    quantization_annotation: QuantizationAnnotation = node.meta.get(
+        "quantization_annotation", QuantizationAnnotation()
+    )
+    quantization_annotation.input_qspec_map[input_node] = qspec
+    node.meta["quantization_annotation"] = quantization_annotation
+
+
+def annotate_output_qspec(node: torch.fx.Node, qspec):
+    quantization_annotation: QuantizationAnnotation = node.meta.get(
+        "quantization_annotation", QuantizationAnnotation()
+    )
+    quantization_annotation.output_qspec = qspec
+    node.meta["quantization_annotation"] = quantization_annotation
+
+
+def mark_nodes_as_annotated(nodes: List[torch.fx.Node] | torch.fx.Node):
+    if isinstance(nodes, torch.fx.Node):
+        nodes = [nodes]
+    for node in nodes:
+        if node is not None:
+            if "quantization_annotation" not in node.meta:
+                node.meta["quantization_annotation"] = QuantizationAnnotation()
+            node.meta["quantization_annotation"]._annotated = True
+
+
+def propagate_annotation_forward(model: torch.fx.GraphModule) -> None:
+    for n in model.graph.nodes:
+        if n.op != "call_function" or n.target not in annot_spec.OP_TO_SHARE_QUANT_SPEC:
+            continue
+
+        prev_node = n.args[0]
+        if not isinstance(prev_node, torch.fx.Node):
+            continue
+
+        quantization_annotation: Optional[QuantizationAnnotation] = prev_node.meta.get(
+            "quantization_annotation", None
+        )
+        if not quantization_annotation:
+            continue
+
+        output_qspec = quantization_annotation.output_qspec
+        if not output_qspec:
+            continue
+
+        # Make sure current node is not annotated
+        if (
+            "quantization_annotation" in n.meta
+            and n.meta["quantization_annotation"]._annotated
+        ):
+            continue
+
+        shared_qspec = SharedQuantizationSpec(prev_node)
+        # Propagate the previous output_qspec to the current node
+        n.meta["quantization_annotation"] = QuantizationAnnotation(
+            input_qspec_map={
+                prev_node: shared_qspec,
+            },
+            output_qspec=shared_qspec,
+            _annotated=True,
+        )