tico 0.1.0.dev250803__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_serializer.py

@@ -20,6 +20,7 @@ import torch
 from circle_schema import circle
 from torch.export.exported_program import ConstantArgument, ExportedProgram, InputKind
 
+from tico.config import CompileConfigBase, get_default_config
 from tico.serialize.circle_mapping import to_circle_dtype, to_circle_shape
 from tico.serialize.operators import *
 from tico.serialize.circle_graph import CircleModel, CircleSubgraph
@@ -47,7 +48,9 @@ def _initialize_model() -> tuple[CircleModel, CircleSubgraph]:
     return model, graph
 
 
-def build_circle(ep: ExportedProgram) -> bytes:
+def build_circle(
+    ep: ExportedProgram, config: CompileConfigBase = get_default_config()
+) -> bytes:
     """Convert ExportedProgram to Circle format.
 
     Args:
@@ -68,9 +71,13 @@ def build_circle(ep: ExportedProgram) -> bytes:
     for in_spec in ep.graph_signature.input_specs:
         if in_spec.kind != InputKind.USER_INPUT:
             continue
-        # NoneType ConstantArgument is ignored.
-        if isinstance(in_spec.arg, ConstantArgument) and in_spec.arg.value == None:
-            continue
+        if isinstance(in_spec.arg, ConstantArgument):
+            # ConstantArgument is ignored when option is given
+            if config.get("remove_constant_input"):
+                continue
+            # NoneType ConstantArgument is ignored.
+            if in_spec.arg.value == None:
+                continue
         arg_name = in_spec.arg.name
         graph.add_input(arg_name)
         logger.debug(f"Registered input: {arg_name}")

tico/serialize/operators/adapters/__init__.py

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

tico/serialize/operators/adapters/llama_rmsnorm.py

@@ -0,0 +1,35 @@
+# 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 contextlib import contextmanager
+
+import torch
+
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+
+
+def llama_rmsnorm_forward_adapter(self: LlamaRMSNorm, hidden_states: torch.Tensor):
+    return torch.ops.circle_custom.rms_norm(
+        hidden_states, self.weight, self.variance_epsilon
+    )
+
+
+@contextmanager
+def patched_llama_rmsnorm():
+    orig = LlamaRMSNorm.forward
+    LlamaRMSNorm.forward = llama_rmsnorm_forward_adapter
+    try:
+        yield
+    finally:
+        LlamaRMSNorm.forward = orig

tico/serialize/operators/op_constant_pad_nd.py

@@ -28,6 +28,42 @@ from tico.utils.errors import InvalidArgumentError
 from tico.utils.validate_args_kwargs import ConstantPadNdArgs
 
 
+def convert_to_circle_padding(pad, input_shape_len):
+    MAX_RANK = 4
+
+    if not (1 <= input_shape_len <= MAX_RANK):
+        raise InvalidArgumentError(
+            f"Input rank must be between 1 and {MAX_RANK}, got {input_shape_len}"
+        )
+
+    if len(pad) % 2 != 0 or len(pad) < 2 or len(pad) > 8:
+        raise InvalidArgumentError(
+            f"Pad length must be an even number between 2 and 8, got {len(pad)}"
+        )
+
+    if len(pad) == 2:
+        padding = [[pad[0], pad[1]]]
+    elif len(pad) == 4:
+        padding = [[pad[2], pad[3]], [pad[0], pad[1]]]
+    elif len(pad) == 6:
+        padding = [[pad[4], pad[5]], [pad[2], pad[3]], [pad[0], pad[1]]]
+    elif len(pad) == 8:
+        padding = [
+            [pad[6], pad[7]],
+            [pad[4], pad[5]],
+            [pad[2], pad[3]],
+            [pad[0], pad[1]],
+        ]
+    else:
+        assert False, "Cannot reach here"
+
+    # Fill [0, 0] padding for the rest of dimension
+    while len(padding) < input_shape_len:
+        padding.insert(0, [0, 0])
+
+    return padding
+
+
 @register_node_visitor
 class ConstantPadNdVisitor(NodeVisitor):
     target: List[torch._ops.OpOverload] = [torch.ops.aten.constant_pad_nd.default]
@@ -45,19 +81,13 @@ class ConstantPadNdVisitor(NodeVisitor):
         val = args.value
 
         if val != 0:
-            raise InvalidArgumentError("Only support 0 value padding.")
+            raise InvalidArgumentError(f"Only support 0 value padding. pad:{pad}")
 
         input_shape_len = len(extract_shape(input_))
-        padding_size = [[pad[2], pad[3]], [pad[0], pad[1]]]
-        if input_shape_len == 3:
-            padding_size = [[0, 0]] + padding_size
-        elif input_shape_len == 4:
-            padding_size = [[0, 0], [0, 0]] + padding_size
-        else:
-            raise InvalidArgumentError("Only support 3D/4D inputs.")
-
-        paddings = torch.tensor(padding_size, dtype=torch.int32)
-        inputs = [input_, paddings]
+
+        padding = convert_to_circle_padding(pad, input_shape_len)
+
+        inputs = [input_, torch.tensor(padding, dtype=torch.int32)]
         outputs = [node]
 
         op_index = get_op_index(

tico/serialize/operators/op_le.py

@@ -0,0 +1,54 @@
+# 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, TYPE_CHECKING
+
+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.validate_args_kwargs import LeArgs
+
+
+@register_node_visitor
+class LeVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.le.Scalar,
+        torch.ops.aten.le.Tensor,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(self, node: torch.fx.Node) -> circle.Operator.OperatorT:
+        args = LeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        other = args.other
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.LESS_EQUAL, self._op_codes
+        )
+
+        inputs = [input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        return operator