tico 0.1.0__py3-none-any.whl

tico/utils/graph.py

@@ -0,0 +1,282 @@
+# Portions of this file are adapted from code originally authored by
+# Meta Platforms, Inc. and affiliates, licensed under the BSD-style
+# license found in the LICENSE file in the root directory of their source tree.
+
+# 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, Tuple, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.export import ExportedProgram
+from torch.export.exported_program import InputKind, InputSpec, TensorArgument
+
+from tico.utils.utils import get_fake_mode, set_new_meta_val
+
+
+def is_torch_param(node: torch.fx.Node, ep: ExportedProgram):
+    assert node.op == "placeholder"
+
+    return node.name in ep.graph_signature.inputs_to_parameters
+
+
+def is_torch_buffer(node: torch.fx.Node, ep: ExportedProgram):
+    assert node.op == "placeholder"
+
+    return node.name in ep.graph_signature.inputs_to_buffers
+
+
+def get_torch_param_value(node: torch.fx.Node, ep: ExportedProgram):
+    assert isinstance(node, torch.fx.Node)
+    assert node.op == "placeholder"
+    assert (
+        node.name in ep.graph_signature.inputs_to_parameters
+    ), "Node {node.name} is not in the parameters"  # FIX CALLER UNLESS
+
+    param_name = ep.graph_signature.inputs_to_parameters[node.name]
+    named_params = dict(ep.named_parameters())
+    assert param_name in named_params
+
+    return named_params[param_name].data
+
+
+def get_torch_buffer_value(node: torch.fx.Node, ep: ExportedProgram):
+    assert isinstance(node, torch.fx.Node)
+    assert node.op == "placeholder"
+    assert (
+        node.name in ep.graph_signature.inputs_to_buffers
+    ), "Node {node.name} is not in the buffers"  # FIX CALLER UNLESS
+
+    buf_name = ep.graph_signature.inputs_to_buffers[node.name]
+    named_buf = dict(ep.named_buffers())
+    assert buf_name in named_buf
+
+    return named_buf[buf_name]
+
+
+def get_first_user_input(exported_program: ExportedProgram) -> Optional[torch.fx.Node]:
+    """Returns the first user input node in the graph."""
+    first_user_input: Optional[torch.fx.Node] = None
+    graph_module = exported_program.graph_module
+    graph: torch.fx.Graph = graph_module.graph
+    for node in graph.nodes:
+        if (
+            node.op == "placeholder"
+            and node.name in exported_program.graph_signature.user_inputs
+        ):
+            first_user_input = node
+            break
+
+    return first_user_input
+
+
+def generate_fqn(prefix: str, exported_program: ExportedProgram):
+    """
+    Generate fully-qualized name for constants.
+
+    This function prevents `exported_program.constants` from having duplicate keys.
+    """
+    cnt = len(exported_program.constants)
+    while True:
+        if f"{prefix}{cnt}" in exported_program.constants:
+            cnt += 1
+            continue
+        break
+    return f"{prefix}{cnt}"
+
+
+def create_input_spec(node, input_kind: InputKind):
+    """
+    @ref https://pytorch.org/docs/stable/export.ir_spec.html#placeholder
+    """
+    if input_kind == InputKind.CONSTANT_TENSOR:
+        return InputSpec(
+            kind=InputKind.CONSTANT_TENSOR,
+            arg=TensorArgument(name=node.name),
+            target=node.target,  # type: ignore[arg-type]
+            persistent=True,
+        )
+    else:
+        raise NotImplementedError("NYI")
+
+
+def validate_input_specs(exported_program):
+    name_to_spec_dict = {
+        s.arg.name: s for s in exported_program.graph_signature.input_specs
+    }
+
+    for node in exported_program.graph.nodes:
+        if node.op != "placeholder":
+            continue
+
+        if node.name not in name_to_spec_dict:
+            raise RuntimeError(
+                "Placeholder node {node.name} does not have corresponding input spec!"
+            )
+
+
+def add_placeholder(
+    exported_program: ExportedProgram,
+    tensor: torch.Tensor,
+    prefix: str,
+) -> torch.fx.Node:
+    """
+    Add a placeholder to the graph and update the exported program.
+    """
+    fqn_name = generate_fqn(prefix, exported_program)
+
+    # Get fake mode before adding placeholder
+    fake_mode = get_fake_mode(exported_program)
+
+    first_user_input = get_first_user_input(exported_program)
+    if not first_user_input:
+        # Placeholder nodes must be the first N nodes in the nodes list of a graph.
+        # Therefore, insert the newly created placeholders at the start of the node list.
+        assert exported_program.graph.nodes
+        first_node = list(exported_program.graph.nodes)[0]
+        first_user_input = first_node
+
+    # Add a placeholder to the graph.
+    with exported_program.graph.inserting_before(first_user_input):
+        const_node = exported_program.graph.placeholder(fqn_name)
+
+    const_node.meta["val"] = fake_mode.from_tensor(tensor, static_shapes=True)
+    const_node.meta["val"].constant = tensor
+
+    # Add a new constant to the exported program.
+    exported_program.constants[const_node.name] = tensor
+
+    # Use update (instead of append) if this assert is violated
+    assert const_node.name not in [
+        s.arg.name for s in exported_program.graph_signature.input_specs
+    ]
+
+    # Append the new input spec.
+    exported_program.graph_signature.input_specs.append(
+        create_input_spec(const_node, InputKind.CONSTANT_TENSOR)
+    )
+
+    # Get old input specs
+    name_to_spec_dict = {
+        s.arg.name: s for s in exported_program.graph_signature.input_specs
+    }
+
+    # Add the new constants to input specs dict.
+    name_to_spec_dict.update(
+        {const_node.name: create_input_spec(const_node, InputKind.CONSTANT_TENSOR)}
+    )
+
+    # Generate new input spec *in the same order of nodes*
+    # IMPORTANT Input specs and their placeholder nodes must have the same order.
+    new_input_specs = []
+    for node in exported_program.graph.nodes:
+        if node.op != "placeholder":
+            continue
+        new_input_specs.append(name_to_spec_dict[node.name])
+    exported_program.graph_signature.input_specs = new_input_specs
+
+    return const_node
+
+
+def is_single_value_tensor(t: torch.Tensor):
+    if len(t.size()) == 0:
+        return True
+    if len(t.size()) == 1 and t.size()[0] == 1:
+        return True
+
+    return False
+
+
+def get_module_name_chain(node: Optional[torch.fx.Node]) -> str:
+    """
+    Returns a slash-separated string of module names representing the
+    hierarchical path of the FX node within the original model.
+
+    If the node has no `nn_module_stack` metadata, "unknown" is returned.
+
+    Example:
+        "encoder/layer1/linear"
+
+    Parameters
+    ----------
+    node: torch.fx.Node
+        A node from an ExportedProgram graph.
+
+    Returns
+    -------
+    str
+        A human-readable string that describes the full module path.
+    """
+    if node is None:
+        return "unknown"
+    # Let's prefix "tico" for graph inputs
+    if node.op == "placeholder" and "nn_module_stack" not in node.meta:
+        return "tico"
+
+    assert isinstance(node, torch.fx.Node)
+    stack = node.meta.get("nn_module_stack")
+    if stack:
+        assert isinstance(stack, dict)
+        # Retrieving the last element is enough.
+        return next(reversed(stack.values()))[1]
+    else:
+        return "unknown"
+
+
+def create_node(
+    graph: torch.fx.Graph,
+    target: torch._ops.OpOverload,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    *,
+    origin: Optional[torch.fx.Node] = None,
+) -> torch.fx.Node:
+    """
+    Insert a new node into graph and propagate metadata from *origin*.
+
+    Parameters
+    ----------
+    graph : torch.fx.Graph
+        The graph that will own the newly-created node.
+
+    target : torch._ops.OpOverload
+        The op to call (e.g. `torch.add` or "call_function" target).
+
+    args : Tuple[Any, ...], optional
+        Positional arguments for the new node.
+
+    kwargs : Dict[str, Any], optional
+        Keyword arguments for the new node.
+
+    origin : torch.fx.Node, optional
+        If given, every key in `origin.meta` **except** "val" is copied
+        onto the new node.  "val" is recomputed from *args* /*kwargs* using
+        the internal meta-inference helper.
+
+    Returns
+    -------
+    torch.fx.Node
+        The freshly inserted node with fully-populated `.meta`.
+    """
+    new_node = graph.call_function(target, args=args, kwargs=kwargs)
+    if origin:
+        assert isinstance(origin, torch.fx.Node), type(origin)
+        # Propagate "nn_module_stack" to retain the originating module context
+        #  for meaningful node names.
+        if "nn_module_stack" in origin.meta:
+            new_node.meta["nn_module_stack"] = origin.meta["nn_module_stack"]
+
+    return new_node

tico/utils/logging.py

@@ -0,0 +1,45 @@
+# 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 logging
+import os
+
+
+def _loggerLevel():
+    TICO_LOG = os.environ.get("TICO_LOG")
+    if TICO_LOG == "1":
+        log_level = logging.FATAL
+    elif TICO_LOG == "2":
+        log_level = logging.WARNING
+    elif TICO_LOG == "3":
+        log_level = logging.INFO
+    elif TICO_LOG == "4":
+        log_level = logging.DEBUG
+    else:
+        log_level = logging.WARNING
+    return log_level
+
+
+LOG_LEVEL = _loggerLevel()
+
+
+def getLogger(name: str):
+    """
+    Get logger with setting log level according to the `TICO_LOG` environment variable.
+    """
+    logging.basicConfig()
+    logger = logging.getLogger(name)
+    logger.setLevel(LOG_LEVEL)
+
+    return logger

tico/utils/model.py

@@ -0,0 +1,37 @@
+# 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 __future__ import annotations
+
+from typing import Any
+
+from tico.interpreter import infer
+
+
+class CircleModel:
+    def __init__(self, circle_binary: bytes):
+        self.circle_binary = circle_binary
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return infer.infer(self.circle_binary, *args, **kwargs)
+
+    @staticmethod
+    def load(circle_path: str) -> CircleModel:
+        with open(circle_path, "rb") as f:
+            buf = bytes(f.read())
+        return CircleModel(buf)
+
+    def save(self, circle_path: str) -> None:
+        with open(circle_path, "wb") as f:
+            f.write(self.circle_binary)

tico/utils/mx/__init__.py

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

tico/utils/mx/elemwise_ops.py

@@ -0,0 +1,267 @@
+"""
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+
+Name:    elemwise_ops.py
+
+Pytorch functions for elementwise (i.e. bfloat) quantization.
+
+Usage Notes:
+ - Use the "Exposed Methods" below to implement autograd functions
+ - Use autograd functions to then implement torch.nn.Module(s)
+ - Do *not* use methods in this file in Modules, they have no defined
+   backwards pass and will block gradient computation.
+ - Avoid importing internal function if at all possible.
+
+Exposed Methods:
+    quantize_elemwise_op - quantizes a tensor to bfloat or other
+                           custom float format
+"""
+import torch
+
+from .formats import RoundingMode, _get_format_params
+from .formats import _get_min_norm, _get_max_norm
+
+
+# -------------------------------------------------------------------------
+# Helper funcs
+# -------------------------------------------------------------------------
+# Never explicitly compute 2**(-exp) since subnorm numbers have
+# exponents smaller than -126
+def _safe_lshift(x, bits, exp):
+    if exp is None:
+        return x * (2**bits)
+    else:
+        return x / (2 ** exp) * (2**bits)
+
+
+def _safe_rshift(x, bits, exp):
+    if exp is None:
+        return x / (2**bits)
+    else:
+        return x / (2**bits) * (2 ** exp)
+
+
+def _round_mantissa(A, bits, round, clamp=False):
+    """
+    Rounds mantissa to nearest bits depending on the rounding method 'round'
+    Args:
+      A     {PyTorch tensor} -- Input tensor
+      round {str}            --  Rounding method
+                                 "floor" rounds to the floor
+                                 "nearest" rounds to ceil or floor, whichever is nearest
+    Returns:
+      A {PyTorch tensor} -- Tensor with mantissas rounded
+    """
+
+    if round == "dither":
+        rand_A = torch.rand_like(A, requires_grad=False)
+        A = torch.sign(A) * torch.floor(torch.abs(A) + rand_A)
+    elif round == "floor":
+        A = torch.sign(A) * torch.floor(torch.abs(A))
+    elif round == "nearest":
+        A = torch.sign(A) * torch.floor(torch.abs(A) + 0.5)
+    elif round == "even":
+        absA = torch.abs(A)
+        # find 0.5, 2.5, 4.5 ...
+        maskA = ((absA - 0.5) % 2 == torch.zeros_like(A)).type(A.dtype)
+        A = torch.sign(A) * (torch.floor(absA + 0.5) - maskA)
+    else:
+        raise Exception("Unrecognized round method %s" % (round))
+
+    # Clip values that cannot be expressed by the specified number of bits
+    if clamp:
+        max_mantissa = 2 ** (bits - 1) - 1
+        A = torch.clamp(A, -max_mantissa, max_mantissa)
+    return A
+
+
+# -------------------------------------------------------------------------
+# Main funcs
+# -------------------------------------------------------------------------
+def _quantize_elemwise_core(A, bits, exp_bits, max_norm, round='nearest',
+                            saturate_normals=False, allow_denorm=True,
+                            custom_cuda=False):
+    """ Core function used for element-wise quantization
+    Arguments:
+      A         {PyTorch tensor} -- A tensor to be quantized
+      bits      {int}            -- Number of mantissa bits. Includes
+                                    sign bit and implicit one for floats
+      exp_bits  {int}            -- Number of exponent bits, 0 for ints
+      max_norm  {float}          -- Largest representable normal number
+      round     {str}            -- Rounding mode: (floor, nearest, even)
+      saturate_normals {bool}    -- If True, normal numbers (i.e., not NaN/Inf)
+                                    that exceed max norm are clamped.
+                                    Must be True for correct MX conversion.
+      allow_denorm     {bool}    -- If False, flush denorm numbers in the
+                                    elem_format to zero.
+      custom_cuda      {str}     -- If True, use custom CUDA kernels
+    Returns:
+      quantized tensor {PyTorch tensor} -- A tensor that has been quantized
+    """
+    A_is_sparse = A.is_sparse
+    if A_is_sparse:
+        if A.layout != torch.sparse_coo:
+            raise NotImplementedError("Only COO layout sparse tensors are currently supported.")
+
+        sparse_A = A.coalesce()
+        A = sparse_A.values().clone()
+
+    # custom cuda only support floor and nearest rounding modes
+    custom_cuda = custom_cuda and round in RoundingMode.string_enums()
+
+    if custom_cuda:
+        A = A.contiguous()
+
+        from . import custom_extensions
+        if A.device.type == "cuda":
+            A = custom_extensions.funcs.quantize_elemwise_func_cuda(
+                A, bits, exp_bits, max_norm, RoundingMode[round],
+                saturate_normals, allow_denorm)
+        elif A.device.type == "cpu":
+            A = custom_extensions.funcs.quantize_elemwise_func_cpp(
+                A, bits, exp_bits, max_norm, RoundingMode[round],
+                saturate_normals, allow_denorm)
+        return A
+
+    # Flush values < min_norm to zero if denorms are not allowed
+    if not allow_denorm and exp_bits > 0:
+        min_norm = _get_min_norm(exp_bits)
+        out = (torch.abs(A) >= min_norm).type(A.dtype) * A
+    else:
+        out = A
+
+    if exp_bits != 0:
+        private_exp = torch.floor(torch.log2(
+            torch.abs(A) + (A == 0).type(A.dtype)))
+
+        # The minimum representable exponent for 8 exp bits is -126
+        min_exp = -(2**(exp_bits-1)) + 2
+        private_exp = private_exp.clip(min=min_exp)
+    else:
+        private_exp = None
+
+    # Scale up so appropriate number of bits are in the integer portion of the number
+    out = _safe_lshift(out, bits - 2, private_exp)
+
+    out = _round_mantissa(out, bits, round, clamp=False)
+
+    # Undo scaling
+    out = _safe_rshift(out, bits - 2, private_exp)
+
+    # Set values > max_norm to Inf if desired, else clamp them
+    if saturate_normals or exp_bits == 0:
+        out = torch.clamp(out, min=-max_norm, max=max_norm)
+    else:
+        out = torch.where((torch.abs(out) > max_norm),
+                           torch.sign(out) * float("Inf"), out)
+
+    # handle Inf/NaN
+    if not custom_cuda:
+        out[A == float("Inf")] = float("Inf")
+        out[A == -float("Inf")] = -float("Inf")
+        out[A == float("NaN")] = float("NaN")
+
+    if A_is_sparse:
+        output = torch.sparse_coo_tensor(sparse_A.indices(), output,
+                sparse_A.size(), dtype=sparse_A.dtype, device=sparse_A.device,
+                requires_grad=sparse_A.requires_grad)
+
+    return out
+
+
+def _quantize_elemwise(A, elem_format, round='nearest', custom_cuda=False,
+                       saturate_normals=False, allow_denorm=True):
+    """ Quantize values to a defined format. See _quantize_elemwise_core()
+    """
+    if elem_format == None:
+        return A
+
+    ebits, mbits, _, max_norm, _ = _get_format_params(elem_format)
+
+    output = _quantize_elemwise_core(
+            A, mbits, ebits, max_norm,
+            round=round, allow_denorm=allow_denorm,
+            saturate_normals=saturate_normals,
+            custom_cuda=custom_cuda)
+
+    return output
+
+
+def _quantize_bfloat(A, bfloat, round='nearest', custom_cuda=False, allow_denorm=True):
+    """ Quantize values to bfloatX format
+    Arguments:
+      bfloat      {int}       -- Total number of bits for bfloatX format,
+                                 Includes 1 sign, 8 exp bits, and variable
+                                 mantissa bits. Must be >= 9.
+    """
+    # Shortcut for no quantization
+    if bfloat == 0 or bfloat == 32:
+        return A
+
+    max_norm = _get_max_norm(8, bfloat-7)
+
+    return _quantize_elemwise_core(
+            A, bits=bfloat-7, exp_bits=8, max_norm=max_norm, round=round,
+            allow_denorm=allow_denorm, custom_cuda=custom_cuda)
+
+
+def _quantize_fp(A, exp_bits=None, mantissa_bits=None,
+                 round='nearest', custom_cuda=False, allow_denorm=True):
+    """ Quantize values to IEEE fpX format. The format defines NaN/Inf
+        and subnorm numbers in the same way as FP32 and FP16.
+    Arguments:
+        exp_bits        {int} -- number of bits used to store exponent
+        mantissa_bits   {int} -- number of bits used to store mantissa, not
+                                 including sign or implicit 1
+        round           {str} -- Rounding mode, (floor, nearest, even)
+    """
+    # Shortcut for no quantization
+    if exp_bits is None or mantissa_bits is None:
+        return A
+
+    max_norm = _get_max_norm(exp_bits, mantissa_bits+2)
+
+    output = _quantize_elemwise_core(
+            A, bits=mantissa_bits + 2, exp_bits=exp_bits,
+            max_norm=max_norm, round=round, allow_denorm=allow_denorm,
+            custom_cuda=custom_cuda)
+
+    return output
+
+
+def quantize_elemwise_op(A, mx_specs, round=None):
+    """A function used for element-wise quantization with mx_specs
+    Arguments:
+      A          {PyTorch tensor} -- a tensor that needs to be quantized
+      mx_specs {dictionary}     -- dictionary to specify mx_specs
+      round      {str}            -- Rounding mode, choose from (floor, nearest, even)
+                                     (default: "nearest")
+    Returns:
+      quantized value {PyTorch tensor} -- a tensor that has been quantized
+    """
+    if mx_specs is None:
+        return A
+    elif round is None:
+        round = mx_specs['round']
+
+    if mx_specs['bfloat'] == 16 and round == 'even'\
+        and torch.cuda.is_bf16_supported() \
+        and mx_specs['bfloat_subnorms'] == True:
+        return A.to(torch.bfloat16)
+
+    if mx_specs['bfloat'] > 0 and mx_specs['fp'] > 0:
+        raise ValueError("Cannot set both [bfloat] and [fp] in mx_specs.")
+    elif mx_specs['bfloat'] > 9:
+        A = _quantize_bfloat(A, bfloat=mx_specs['bfloat'], round=round,
+                             custom_cuda=mx_specs['custom_cuda'],
+                             allow_denorm=mx_specs['bfloat_subnorms'])
+    elif mx_specs['bfloat'] > 0 and mx_specs['bfloat'] <= 9:
+        raise ValueError("Cannot set [bfloat] <= 9 in mx_specs.")
+    elif mx_specs['fp'] > 6:
+        A = _quantize_fp(A, exp_bits=5, mantissa_bits=mx_specs['fp'] - 6,
+                         round=round, custom_cuda=mx_specs['custom_cuda'],
+                         allow_denorm=mx_specs['bfloat_subnorms'])
+    elif mx_specs['fp'] > 0 and mx_specs['fp'] <= 6:
+        raise ValueError("Cannot set [fp] <= 6 in mx_specs.")
+    return A