tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev251102__py3-none-any.whl

tico/quantization/wrapq/wrappers/quant_elementwise.py

@@ -0,0 +1,111 @@
+# 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
+
+import torch
+import torch.nn as nn
+
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import register
+
+
+class QuantElementwise(QuantModuleBase):
+    """
+    Generic wrapper for any 1-to-1 element-wise op `y = f(x)`.
+
+    Sub-classes only need to implement:
+        • `FUNC`: a Callable that maps tensor→tensor
+    """
+
+    # subclass must set this
+    FUNC: Callable[[torch.Tensor], torch.Tensor] | None = None
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if cls is QuantElementwise:
+            return
+        if cls.FUNC is None:
+            raise NotImplementedError(
+                f"{cls.__name__} must define a staticmethod `FUNC(tensor) -> tensor`"
+            )
+
+    def __init__(
+        self,
+        fp_module: nn.Module,
+        *,
+        qcfg: Optional[PTQConfig] = None,
+        fp_name: Optional[str] = None,
+    ):
+        super().__init__(qcfg, fp_name=fp_name)
+        self.module = fp_module
+        self.act_in_obs = self._make_obs("act_in")
+        self.act_out_obs = self._make_obs("act_out")
+
+    # ------------------------------------------------------------
+    def forward(self, x):
+        x_q = self._fq(x, self.act_in_obs)
+        assert self.FUNC is not None
+        y = self.FUNC(x_q)  # element-wise op
+        y_q = self._fq(y, self.act_out_obs)
+        return y_q
+
+    # ------------------------------------------------------------
+    def _all_observers(self):
+        return (self.act_in_obs, self.act_out_obs)
+
+
+"""
+Why `FUNC` is a `staticmethod`
+
+- Prevents automatic binding: calling `self.FUNC(x)` will not inject `self`,
+  so the callable keeps the expected signature `Tensor -> Tensor`
+  (e.g., `torch.sigmoid(x)`), avoiding TypeErrors.
+
+- Expresses purity and statelessness: `FUNC` is a pure, element-wise transform
+  that must not read or mutate module state (params, buffers, config).
+
+- Tracing/export friendly (FX / TorchScript): the call is captured as
+  `call_function(torch.*)` instead of a bound `call_method`, which makes graph
+  rewrites/pattern-matching and backends' substitutions more reliable.
+
+- Avoids submodule pollution: we keep a functional op (`torch.relu`) rather
+  than an `nn.Module` instance that would appear in the module tree.
+
+- Small perf/alloc win: no bound-method objects are created on each call.
+"""
+
+# Sigmoid
+@register(nn.Sigmoid)
+class QuantSigmoid(QuantElementwise):
+    FUNC = staticmethod(torch.sigmoid)
+
+
+# Tanh
+@register(nn.Tanh)
+class QuantTanh(QuantElementwise):
+    FUNC = staticmethod(torch.tanh)
+
+
+# ReLU
+@register(nn.ReLU)
+class QuantReLU(QuantElementwise):
+    FUNC = staticmethod(torch.relu)
+
+
+# GELU (approximate)
+@register(nn.GELU)
+class QuantGELU(QuantElementwise):
+    FUNC = staticmethod(torch.nn.functional.gelu)

tico/quantization/wrapq/wrappers/quant_module_base.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 abc import ABC, abstractmethod
+from typing import Iterable, Optional, Tuple
+
+import torch.nn as nn
+
+from tico.quantization.config.ptq import PTQConfig
+
+from tico.quantization.wrapq.mode import Mode
+from tico.quantization.wrapq.observers.base import ObserverBase
+
+
+class QuantModuleBase(nn.Module, ABC):
+    """
+    Abstract parent for EVERY wrapper.
+
+    Responsibilities
+    ----------------
+    • Own *one* Mode enum (`NO_QUANT / CALIB / QUANT`)
+    • Own a PTQConfig describing default / per-observer dtypes
+    • Expose a canonical lifecycle:
+          enable_calibration()
+          freeze_qparams()
+    • Provide helper `_fq(x, observer)` (“fake-quant or collect”) so
+      subclasses write arithmetic code without boilerplate.
+    """
+
+    def __init__(
+        self, qcfg: Optional[PTQConfig] = None, *, fp_name: Optional[str] = None
+    ) -> None:
+        super().__init__()
+        self.qcfg = qcfg or PTQConfig()
+        self._mode: Mode = Mode.NO_QUANT  # default state
+        self.fp_name = fp_name
+
+    def _child_quant_modules(self):
+        """
+        Yield immediate QuantModuleBase *descendants*, skipping over pure containers
+        (e.g., ModuleList/Sequential/ModuleDict). Once a QuantModuleBase is found,
+        do NOT descend into it here—let recursion happen level by level.
+        """
+        seen = set()
+        stack = list(self.children())  # start from direct children
+
+        while stack:
+            m = stack.pop()
+            if isinstance(m, QuantModuleBase):
+                if id(m) not in seen:
+                    seen.add(id(m))
+                    yield m
+                # IMPORTANT: do not recurse into `m` here; its own call will handle its subtree
+            elif isinstance(m, (nn.ModuleList, nn.ModuleDict, nn.Sequential)):
+                # `m` is a container or a non-quant leaf: keep descending until we hit quant modules
+                stack.extend(list(m.children()))
+
+    def enable_calibration(self) -> None:
+        self._mode = Mode.CALIB
+        for obs in self._all_observers():
+            obs.enabled = True
+            obs.reset()
+
+        # propagate to children
+        for child in self._child_quant_modules():
+            child.enable_calibration()
+
+    def freeze_qparams(self) -> None:
+        self._mode = Mode.QUANT
+        for obs in self._all_observers():
+            obs.enabled = False
+            obs.compute_qparams()
+
+        # propagate to children
+        for child in self._child_quant_modules():
+            child.freeze_qparams()
+
+    def _fq(self, x, obs: ObserverBase):
+        """Fake-quant or collect."""
+        if self._mode is Mode.CALIB:
+            obs.collect(x.detach())
+            return x
+        if self._mode is Mode.QUANT:
+            return obs.fake_quant(x)
+        return x  # NO_QUANT
+
+    @abstractmethod
+    def _all_observers(self) -> Iterable[ObserverBase]:
+        """Return every observer owned by this module."""
+        ...
+
+    def named_observers(self) -> Iterable[Tuple[str, ObserverBase]]:
+        for obs in self._all_observers():
+            yield obs.name, obs
+
+    def get_observer(self, name: str) -> Optional[ObserverBase]:
+        for obs in self._all_observers():
+            if obs.name == name:
+                return obs
+        return None
+
+    def _make_obs(
+        self,
+        name: str,
+        **default_kwargs,
+    ) -> ObserverBase:
+        """
+        Instantiate an observer named *name*.
+
+        Precedence (3-tier) for keys:
+           • observer:  user > wrapper-default > PTQConfig.default_observer
+           • dtype:     user > wrapper-default > PTQConfig.default_dtype
+           • qscheme:   user > wrapper-default > PTQConfig.default_qscheme
+
+        Other kwargs (e.g., qscheme, channel_axis, etc.) remain:
+           user override > wrapper-default
+        """
+        _UNSPEC = object()
+
+        wrapper_defaults = default_kwargs.copy()
+        user_cfg = self.qcfg.get_kwargs(name).copy()
+
+        def pick3(user_val, wrap_val, global_val):
+            return (
+                user_val
+                if user_val is not _UNSPEC
+                else wrap_val
+                if wrap_val is not _UNSPEC
+                else global_val
+            )
+
+        # 1) resolve observer class
+        user_observer = user_cfg.pop("observer", _UNSPEC)
+        wrapper_observer = wrapper_defaults.pop("observer", _UNSPEC)
+        obs_cls = pick3(user_observer, wrapper_observer, self.qcfg.default_observer)
+
+        # 2) resolve dtype
+        user_dtype = user_cfg.pop("dtype", _UNSPEC)
+        wrapper_dtype = wrapper_defaults.pop("dtype", _UNSPEC)
+        final_dtype = pick3(user_dtype, wrapper_dtype, self.qcfg.default_dtype)
+
+        # 3) resolve qscheme
+        user_qscheme = user_cfg.pop("qscheme", _UNSPEC)
+        wrapper_qscheme = wrapper_defaults.pop("qscheme", _UNSPEC)
+        final_qscheme = pick3(user_qscheme, wrapper_qscheme, self.qcfg.default_qscheme)
+
+        # 4) merge remaining kwargs: user_cfg wins
+        final_kw = wrapper_defaults
+        final_kw.update(user_cfg)
+        final_kw["dtype"] = final_dtype
+        final_kw["qscheme"] = final_qscheme
+
+        return obs_cls(**final_kw, name=name)
+
+    # nice repr
+    def extra_repr(self) -> str:
+        return f"mode={self._mode.name.lower()}"

tico/quantization/wrapq/wrappers/registry.py

@@ -0,0 +1,125 @@
+# 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 importlib
+from typing import Callable, Dict, Type
+
+import torch.nn as nn
+
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+
+_WRAPPERS: Dict[Type[nn.Module], Type[QuantModuleBase]] = {}
+_IMPORT_ONCE = False
+_CORE_MODULES = (
+    "tico.quantization.wrapq.wrappers.quant_elementwise",
+    "tico.quantization.wrapq.wrappers.nn.quant_layernorm",
+    "tico.quantization.wrapq.wrappers.nn.quant_linear",
+    "tico.quantization.wrapq.wrappers.nn.quant_silu",
+    # llama
+    "tico.quantization.wrapq.wrappers.llama.quant_attn",
+    "tico.quantization.wrapq.wrappers.llama.quant_decoder_layer",
+    "tico.quantization.wrapq.wrappers.llama.quant_mlp",
+    # fairseq
+    "tico.quantization.wrapq.wrappers.fairseq.quant_decoder_layer",
+    "tico.quantization.wrapq.wrappers.fairseq.quant_encoder",
+    "tico.quantization.wrapq.wrappers.fairseq.quant_encoder_layer",
+    "tico.quantization.wrapq.wrappers.fairseq.quant_mha",
+    # add future core wrappers here
+)
+
+
+def _lazy_init():
+    """
+    Deferred one-shot import of "core wrapper modules".
+
+    Why not import everything when the program first starts?
+    --------------------------------------------------
+    * **Avoid circular-import hell**
+      Core wrappers often import `PTQWrapper`, which in turn calls
+      `registry.lookup()`.  Importing those files eagerly here would create a
+      cycle (`registry → wrapper → registry`).  Delaying the import until the
+      *first* `lookup()` call lets Python finish constructing the registry
+      module before any wrapper files are touched.
+
+    * **Cold-start speed**
+      Most user code never wraps layers explicitly; they only hit
+      `PTQWrapper` if they are doing quantization.  Deferring half-a-dozen
+      heavyweight `import torch …` files until they are really needed
+      reduces library start-up latency in the common path.
+
+    * **Optional dependencies**
+      Core wrappers listed in `_CORE_MODULES` are chosen to be dependency-free
+      (pure PyTorch).  Anything that needs `transformers`, `torchvision`,
+      etc. uses the `@try_register()` decorator inside its own module.  Those
+      optional modules are *not* imported here, so users without the extra
+      packages still get a clean import.
+
+    Implementation notes
+    --------------------
+    * `_IMPORT_ONCE` guard ensures we execute the import loop only once,
+      even if `lookup()` is called from multiple threads.
+    * Each path in `_CORE_MODULES` is a "fully-qualified module string"
+      (e.g. "ptq.wrappers.linear_quant").  Importing the module runs all
+      its `@register(nn.Layer)` decorators, populating `_WRAPPERS`.
+    * After the first call the function becomes a cheap constant-time no-op.
+    """
+    global _IMPORT_ONCE
+    if _IMPORT_ONCE:
+        return
+    for mod in _CORE_MODULES:
+        __import__(mod)  # triggers decorators
+    _IMPORT_ONCE = True
+
+
+# ───────────────────────────── decorator for always-present classes
+def register(
+    fp_cls: Type[nn.Module],
+) -> Callable[[Type[QuantModuleBase]], Type[QuantModuleBase]]:
+    def _decorator(quant_cls: Type[QuantModuleBase]):
+        _WRAPPERS[fp_cls] = quant_cls
+        return quant_cls
+
+    return _decorator
+
+
+# ───────────────────────────── conditional decorator
+def try_register(
+    *paths: str,
+) -> Callable[[Type[QuantModuleBase]], Type[QuantModuleBase]]:
+    """
+    @try_register("transformers.models.llama.modeling_llama.LlamaMLP")
+
+    • If import succeeds → behave like `@register`
+    • If module/class not found → become a NO-OP
+    """
+
+    def _decorator(quant_cls: Type[QuantModuleBase]):
+        for path in paths:
+            module_name, _, cls_name = path.rpartition(".")
+            try:
+                mod = importlib.import_module(module_name)
+                fp_cls = getattr(mod, cls_name)
+                _WRAPPERS[fp_cls] = quant_cls
+            except (ModuleNotFoundError, AttributeError):
+                # optional dep missing or class renamed – skip silently
+                pass
+        return quant_cls
+
+    return _decorator
+
+
+# ───────────────────────────── lookup
+def lookup(fp_cls: Type[nn.Module]) -> Type[QuantModuleBase] | None:
+    _lazy_init()
+    return _WRAPPERS.get(fp_cls)

tico/serialize/circle_graph.py

@@ -24,9 +24,10 @@ from torch._subclasses.fake_tensor import FakeTensor
 
 from tico.serialize.circle_mapping import (
     extract_circle_dtype,
-    extract_shape,
+    extract_circle_shape,
     str_to_circle_dtype,
     to_circle_dtype,
+    to_circle_shape,
 )
 from tico.serialize.pack import pack_buffer
 from tico.serialize.quant_param import QPARAM_KEY, QuantParam
@@ -151,7 +152,8 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         self.name_to_node[tensor.name] = node
         assert node.meta.get("val") is not None
         tensor.type = extract_circle_dtype(node)
-        tensor.shape = list(extract_shape(node))
+        tensor.shape, tensor.shapeSignature = extract_circle_shape(node)  # type: ignore[assignment]
+
         if QPARAM_KEY in node.meta:
             tensor.quantization = to_circle_qparam(node.meta[QPARAM_KEY])
             tensor.type = str_to_circle_dtype(node.meta[QPARAM_KEY].dtype)
@@ -185,7 +187,7 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         torch_t = torch.as_tensor(data=data)
         torch_t_shape = list(torch_t.size())
         tensor.type = to_circle_dtype(torch_dtype=torch_t.dtype)
-        tensor.shape = torch_t_shape
+        tensor.shape, tensor.shapeSignature = to_circle_shape(torch_t_shape)
 
         buffer = circle.Buffer.BufferT()
         buffer.data = torch_t.flatten().cpu().numpy().view(np.uint8)  # type: ignore[assignment]
@@ -199,6 +201,7 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         self,
         prefix: str,
         shape: List[int],
+        shape_signature: Optional[List[int]],
         dtype: int,
         qparam: Optional[QuantParam] = None,
         source_node: Optional[torch.fx.Node] = None,
@@ -221,6 +224,8 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
             A name prefix used to generate a unique tensor name.
         shape : List[int]
             The shape of the tensor.
+        shape_signature : Optional[List[int]]
+            The shape signature of the tensor to express Dynamic Shape. Defaults to `None` for Static Shape.
         dtype : int
             The Circle-compatible dtype of the tensor. Use `to_circle_dtype()` to convert.
         qparam : Optional[QuantParam]
@@ -241,6 +246,9 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         if source_node is not None:
             self.name_to_node[tensor.name] = source_node
         tensor.shape = shape
+        if shape_signature is not None:
+            tensor.shapeSignature = shape_signature
+
         if qparam is not None:
             tensor.quantization = to_circle_qparam(qparam)
             tensor.type = str_to_circle_dtype(qparam.dtype)
@@ -305,7 +313,7 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         self, node: Union[torch.fx.Node, circle.Tensor.TensorT, ConstData]
     ) -> int:
         # return -1 if node is None. This is for generating CircleOutputExclude
-        if node == None:
+        if node is None:
             return -1
 
         if hasattr(node, "name") and node.name in self.name_to_tid:

tico/serialize/circle_mapping.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple, TYPE_CHECKING, Union
+from typing import List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
     import torch.fx
@@ -128,6 +128,79 @@ def extract_shape(node: torch.fx.Node) -> torch.Size:
     return val_shape
 
 
+def extract_circle_shape(node: torch.fx.Node) -> Tuple[List[int], Optional[List[int]]]:
+    return to_circle_shape(extract_shape(node))
+
+
+def to_circle_shape(
+    torch_shape: Union[
+        torch.Size, Sequence[int | torch.SymInt]
+    ],  # Sequence[int | torch.SymInt] is added for type covariance
+) -> Tuple[List[int], Optional[List[int]]]:
+
+    if any(isinstance(s, torch.SymInt) for s in torch_shape):
+        # Follow dynamic shape spec
+        shape = []
+        shape_signature = []
+        for s in torch_shape:
+            if isinstance(s, torch.SymInt):
+                shape.append(1)
+                shape_signature.append(-1)
+            elif isinstance(s, int):
+                shape.append(s)
+                shape_signature.append(s)
+            else:
+                raise RuntimeError(f"Unsupported shape {torch_shape}")
+        return shape, shape_signature
+    else:
+        # Follow static shape spec
+        shape = []
+        shape_signature = None
+        for s in torch_shape:
+            if isinstance(s, int):
+                shape.append(s)
+            else:
+                assert False, "Cannot reach here"
+        return shape, shape_signature
+
+
+def validate_circle_shape(shape: List[int], shape_signature: Optional[List[int]]):
+    """
+    Validate circle tensor shape and shape_signature.
+    @ref https://github.com/Samsung/TICO/issues/244
+    """
+    if shape_signature is not None:
+        if len(shape_signature) == 0:
+            raise ValueError(
+                "Invalid circle shape: shape_signature must not be an empty list. "
+                "For static shapes, use None instead of []."
+            )
+        if len(shape) != len(shape_signature):
+            raise ValueError(
+                f"Invalid circle shape: shape and shape_signature must have same length: {shape} {shape_signature}"
+            )
+        if not all(isinstance(s, int) for s in shape_signature):
+            raise ValueError(
+                f"circle tensor shape_signature must be all integer values. {shape_signature}"
+            )
+        for s, ss in zip(shape, shape_signature):
+            if ss == -1:
+                # dynamic shape dimension
+                if s != 1:
+                    raise ValueError(
+                        f"Invalid circle shape: {s} {ss} {shape} {shape_signature}"
+                    )
+            else:
+                # static shape dimension
+                if s != ss:
+                    raise ValueError(
+                        f"Invalid circle shape: {s} {ss} {shape} {shape_signature}"
+                    )
+
+    if not all(isinstance(s, int) for s in shape):
+        raise ValueError(f"circle tensor shape must be all integer values. {shape}")
+
+
 # Return stride of node
 def extract_stride(node: torch.fx.Node) -> Tuple[int, ...]:
     assert node.meta is not None
@@ -157,7 +230,8 @@ def check_if_i32_range(axis: Union[list, int]):
     return all(INT32_MIN <= val <= INT32_MAX for val in values)
 
 
-def circle_legalize_dtype_to(values, *, dtype: torch.dtype):
+# TODO: Revisit this dtype legalization function as it breaks SRP
+def circle_legalize_dtype_to(values, *, dtype: torch.dtype) -> torch.Tensor:
     """
     Legalize data types from `torch.int64` to `torch.int32`.