tico 0.1.0__py3-none-any.whl

tico/utils/mx/formats.py

@@ -0,0 +1,125 @@
+"""
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+"""
+
+from enum import Enum, IntEnum
+
+FP32_EXPONENT_BIAS = 127
+FP32_MIN_NORMAL = 2 ** (-FP32_EXPONENT_BIAS + 1)
+
+# Enum for rounding modes
+class RoundingMode(IntEnum):
+    nearest = 0
+    floor = 1
+    even = 2
+
+    @staticmethod
+    def string_enums():
+        return [s.name for s in list(RoundingMode)]
+
+# Enum for scalar data formats
+class ElemFormat(Enum):
+    int8 = 1
+    int4 = 2
+    int2 = 3
+    fp8_e5m2 = 4
+    fp8_e4m3 = 5
+    fp6_e3m2 = 6
+    fp6_e2m3 = 7
+    fp4 = 8
+    fp4_e2m1 = 8
+    float16 = 9
+    fp16 = 9
+    bfloat16 = 10
+    bf16 = 10
+
+    @staticmethod
+    def from_str(s):
+        assert(s != None), "String elem_format == None"
+        s = s.lower()
+        if hasattr(ElemFormat, s):
+            return getattr(ElemFormat, s)
+        else:
+            raise Exception("Undefined elem format", s)
+
+
+def _get_min_norm(ebits):
+    """ Valid for all float formats """
+    emin = 2 - (2 ** (ebits - 1))
+    return 0 if ebits == 0 else 2 ** emin
+
+
+def _get_max_norm(ebits, mbits):
+    """ Valid only for floats that define NaN """
+    assert(ebits >= 5), "invalid for floats that don't define NaN"
+    emax = 0 if ebits==0 else 2**(ebits - 1) - 1
+    return 2**emax * float(2**(mbits-1) - 1) / 2**(mbits-2)
+
+
+_FORMAT_CACHE = {}
+def _get_format_params(fmt):
+    """ Allowed formats:
+        - intX:         2 <= X <= 32, assume sign-magnitude, 1.xxx representation
+        - floatX/fpX:   16 <= X <= 28, assume top exp is used for NaN/Inf
+        - bfloatX/bfX:  9 <= X <= 32
+        - fp4,                  no NaN/Inf
+        - fp6_e3m2/e2m3,        no NaN/Inf 
+        - fp8_e4m3/e5m2,        e5m2 normal NaN/Inf, e4m3 special behavior
+
+        Returns:
+          ebits: exponent bits
+          mbits: mantissa bits: includes sign and implicit bits
+          emax: max normal exponent
+          max_norm: max normal number
+          min_norm: min normal number
+    """
+    if type(fmt) is str:
+        fmt = ElemFormat.from_str(fmt)
+
+    if fmt in _FORMAT_CACHE:
+        return _FORMAT_CACHE[fmt]
+
+    if fmt == ElemFormat.int8:
+        ebits, mbits = 0, 8
+        emax = 0
+    elif fmt == ElemFormat.int4:
+        ebits, mbits = 0, 4
+        emax = 0
+    elif fmt == ElemFormat.int2:
+        ebits, mbits = 0, 2
+        emax = 0
+    elif fmt == ElemFormat.fp8_e5m2:
+        ebits, mbits = 5, 4
+        emax = 2**(ebits - 1) - 1
+    elif fmt == ElemFormat.fp8_e4m3:
+        ebits, mbits = 4, 5
+        emax = 2**(ebits - 1)
+    elif fmt == ElemFormat.fp6_e3m2:
+        ebits, mbits = 3, 4
+        emax = 2**(ebits - 1)
+    elif fmt == ElemFormat.fp6_e2m3:
+        ebits, mbits = 2, 5
+        emax = 2**(ebits - 1)
+    elif fmt == ElemFormat.fp4:
+        ebits, mbits = 2, 3
+        emax = 2**(ebits - 1)
+    elif fmt == ElemFormat.float16:
+        ebits, mbits = 5, 12
+        emax = 2**(ebits - 1) - 1
+    elif fmt == ElemFormat.bfloat16:
+        ebits, mbits = 8, 9
+        emax = 2**(ebits - 1) - 1
+    else:
+        raise Exception("Unknown element format %s" % fmt)
+    
+    if fmt != ElemFormat.fp8_e4m3:
+        max_norm = 2**emax * float(2**(mbits-1) - 1) / 2**(mbits-2)
+    else:
+        max_norm = 2**emax * 1.75  # FP8 has custom max_norm
+
+    min_norm = _get_min_norm(ebits)
+    
+    _FORMAT_CACHE[fmt] = (ebits, mbits, emax, max_norm, min_norm)
+
+    return ebits, mbits, emax, max_norm, min_norm

tico/utils/mx/mx_ops.py

@@ -0,0 +1,270 @@
+# 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.
+#
+# This file was copied from https://github.com/microsoft/microxcaling/tree/v1.1.0
+# and modified for our purpose.
+"""
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+
+Name:    mx_ops.py
+
+Pytorch methods for MX 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_mx_op - quantizes a tensor to MX format.
+
+Internal Methods:
+    _safe_lshift, _safe_rshift - fp16 compatible shifts
+    _shared_exponents - Returns MX shared exponent for the passed tensor
+    _reshape_to_blocks - tiles a tensor by splitting one dim into two
+    _undo_reshape_to_blocks - undos the above reshaping
+    _quantize_mx - quantizes a tensor to MX format
+"""
+
+import torch
+
+from .elemwise_ops import _quantize_elemwise_core
+
+from .formats import (
+    _get_format_params,
+    FP32_EXPONENT_BIAS,
+    FP32_MIN_NORMAL,
+    RoundingMode,
+)
+
+
+# -------------------------------------------------------------------------
+# Helper funcs
+# -------------------------------------------------------------------------
+def _shared_exponents(A, method="max", axes=None, ebits=0):
+    """
+    Get shared exponents for the passed matrix A.
+    Args:
+      A      {PyTorch tensor} -- Input tensor
+      method {str}            -- Exponent selection method.
+                                 "max" uses the max absolute value
+                                 "none" uses an exponent for each value (i.e., no sharing)
+      axes   {list(int)}      -- List of integers which specifies the axes across which
+                                 shared exponents are calculated.
+    Returns:
+      shared_exp {PyTorch tensor} -- Tensor of shared exponents
+    """
+
+    if method == "max":
+        if axes is None:
+            shared_exp = torch.max(torch.abs(A))
+        else:
+            shared_exp = A
+            for axis in axes:
+                shared_exp, _ = torch.max(torch.abs(shared_exp), dim=axis, keepdim=True)
+    elif method == "none":
+        shared_exp = torch.abs(A)
+    else:
+        raise Exception("Unrecognized shared exponent selection method %s" % (method))
+
+    # log2(shared_exp) and truncate to integer
+    shared_exp = torch.floor(
+        torch.log2(
+            shared_exp + FP32_MIN_NORMAL * (shared_exp == 0).type(shared_exp.dtype)
+        )
+    )
+
+    # Restrict to [-emax, emax] range
+    if ebits > 0:
+        emax = 2 ** (ebits - 1) - 1
+        # shared_exp = torch.clamp(shared_exp, -emax, emax)
+        # Overflow to Inf
+        shared_exp[shared_exp > emax] = float("NaN")
+        # Underflows are set to -127 which causes them to be
+        # flushed to 0 later
+        shared_exp[shared_exp < -emax] = -emax
+
+    return shared_exp
+
+
+def _reshape_to_blocks(A, axes, block_size):
+    if axes is None:
+        raise Exception(
+            "axes required in order to determine which "
+            "dimension toapply block size to"
+        )
+    if block_size == 0:
+        raise Exception("block_size == 0 in _reshape_to_blocks")
+
+    # Fix axes to be positive and sort them
+    axes = [(x + len(A.shape) if x < 0 else x) for x in axes]
+    assert all(x >= 0 for x in axes)
+    axes = sorted(axes)
+
+    # Add extra dimension for tiles
+    for i in range(len(axes)):
+        axes[i] += i  # Shift axes due to added dimensions
+        A = torch.unsqueeze(A, dim=axes[i] + 1)
+
+    # Pad to block_size
+    orig_shape = A.size()
+    pad = []
+    for i in range(len(orig_shape)):
+        pad += [0, 0]
+
+    do_padding = False
+    for axis in axes:
+        pre_pad_size = orig_shape[axis]
+        if isinstance(pre_pad_size, torch.Tensor):
+            pre_pad_size = int(pre_pad_size.value)
+        # Don't pad if the axis is short enough to fit inside one tile
+        if pre_pad_size % block_size == 0:
+            pad[2 * axis] = 0
+        else:
+            pad[2 * axis] = block_size - pre_pad_size % block_size
+            do_padding = True
+
+    if do_padding:
+        pad = list(reversed(pad))
+        A = torch.nn.functional.pad(A, pad, mode="constant")
+
+    def _reshape(shape, reshape_block_size):
+        for axis in axes:
+            # Reshape to tiles if axis length > reshape_block_size
+            if shape[axis] >= reshape_block_size:
+                assert shape[axis] % reshape_block_size == 0
+                shape[axis + 1] = reshape_block_size
+                shape[axis] = shape[axis] // reshape_block_size
+            # Otherwise preserve length and insert a 1 into the shape
+            else:
+                shape[axis + 1] = shape[axis]
+                shape[axis] = 1
+        return shape
+
+    # Reshape to tiles
+    padded_shape = A.size()
+    reshape = _reshape(list(padded_shape), block_size)
+
+    A = A.view(reshape)
+    return A, axes, orig_shape, padded_shape
+
+
+def _undo_reshape_to_blocks(A, padded_shape, orig_shape, axes):
+    # Undo tile reshaping
+    A = A.view(padded_shape)
+    # Undo padding
+    if not list(padded_shape) == list(orig_shape):
+        slices = [slice(0, x) for x in orig_shape]
+        A = A[slices]
+    for axis in reversed(axes):
+        # Remove extra dimension
+        A = torch.squeeze(A, dim=axis + 1)
+    return A
+
+
+# -------------------------------------------------------------------------
+# Main funcs
+# -------------------------------------------------------------------------
+def _quantize_mx(
+    A,
+    scale_bits,
+    elem_format,  # can be None for no quantization
+    shared_exp_method="max",
+    axes=None,
+    block_size=0,
+    round="nearest",
+    flush_fp32_subnorms=False,
+    custom_cuda=False,
+):
+    """Function used for MX* quantization"""
+    # Shortcut for no quantization
+    if elem_format == None:
+        return A
+
+    assert scale_bits > 0
+
+    # Make sure axes is a list of non-negative numbers
+    axes = [axes] if type(axes) == int else axes
+    axes = [x + A.ndim if x < 0 else x for x in axes]
+
+    # Custom CUDA only supports limited rounding modes
+    custom_cuda = custom_cuda and round in RoundingMode.string_enums()
+
+    ebits, mbits, emax, max_norm, _ = _get_format_params(elem_format)
+
+    # Perform tiling to the hardware vector size
+    if block_size > 0:
+        A, axes, orig_shape, padded_shape = _reshape_to_blocks(A, axes, block_size)
+
+    ####################
+    # Quantize
+    ####################
+    shared_exp_axes = [x + 1 for x in axes] if block_size > 0 else axes
+
+    # Get shared exponents
+    shared_exp = _shared_exponents(
+        A,
+        method=shared_exp_method,
+        axes=shared_exp_axes,
+        ebits=0,
+    )
+
+    # Flush subnormal FP32 inputs to zero
+    if flush_fp32_subnorms:
+        A = A * (shared_exp > -FP32_EXPONENT_BIAS).type(A.dtype)
+
+    # Offset the max exponent by the largest representable exponent
+    # in the element data format
+    shared_exp = shared_exp - emax
+
+    scale_emax = 2 ** (scale_bits - 1) - 1
+    shared_exp[shared_exp > scale_emax] = float("NaN")
+    shared_exp[shared_exp < -scale_emax] = -scale_emax
+
+    A = A / (2**shared_exp)
+
+    A = _quantize_elemwise_core(
+        A,
+        mbits,
+        ebits,
+        max_norm,
+        round=round,
+        allow_denorm=True,
+        saturate_normals=True,
+        custom_cuda=custom_cuda,
+    )
+
+    A = A * (2**shared_exp)
+
+    # Undo tile reshaping
+    if block_size:
+        A = _undo_reshape_to_blocks(A, padded_shape, orig_shape, axes)
+
+    return A
+
+
+# Wrapper function of circle_custom::quantize_mx
+def quantize_mx(
+    input_: torch.Tensor,
+    elem_format: str,
+    axis: int,
+    shared_exp_method: str = "max",
+    round: str = "nearest",
+) -> torch.Tensor:
+    return torch.ops.circle_custom.quantize_mx(
+        input_, elem_format, axis, shared_exp_method=shared_exp_method, round=round
+    )

tico/utils/padding.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.
+
+import torch
+
+from tico.utils.errors import InvalidArgumentError
+
+SAME = 0
+VALID = 1
+
+
+def is_valid_padding(padding: str | list):
+    if isinstance(padding, str):
+        return padding == "valid"
+
+    if isinstance(padding, list):
+        assert len(padding) == 2, "Padding should be a list of length 2."
+        return padding == [0, 0]
+
+    raise InvalidArgumentError("Invalid padding.")
+
+
+def is_same_padding(
+    padding: str | list, input_shape: list | torch.Size, output_shape: list | torch.Size
+):
+    if isinstance(padding, str):
+        return padding == "same"
+
+    if isinstance(padding, list):
+        assert len(padding) == 2, "Padding should be a list of length 2."
+
+        input_HW = input_shape[1:2]  # N H W C
+        output_HW = output_shape[1:2]  # N H W C
+        return input_HW == output_HW
+
+    raise InvalidArgumentError("Invalid padding.")

tico/utils/passes.py

@@ -0,0 +1,76 @@
+# 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 dataclasses import dataclass
+from enum import Enum
+from typing import List
+
+from torch.export import ExportedProgram
+
+
+@dataclass
+class PassResult:
+    modified: bool
+
+
+class PassBase(ABC):
+    """
+    Base interface for passes.
+    """
+
+    @abstractmethod
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        pass
+
+
+class PassStrategy(Enum):
+    # Run passes until there are no changes.
+    UNTIL_NO_CHANGE = (1,)
+    # Same as `UNTIL_NO_CHANGE` but it starts agian from the beginning.
+    RESTART = (2,)
+
+
+class PassManager:
+    def __init__(
+        self,
+        passes: List[PassBase],
+        strategy: PassStrategy = PassStrategy.RESTART,
+    ):
+        self.passes: List[PassBase] = passes
+        self.strategy: PassStrategy = strategy
+
+    def run(self, exported_program: ExportedProgram):
+        MAXIMUM_STEP_COUNT = 1000
+        step = 0
+        while True:
+            modified = False
+            for _pass in self.passes:
+                # Automatically update the signatures of the input and output.
+                # https://github.com/pytorch/executorch/issues/4013#issuecomment-2187161844
+                with exported_program.graph_module._set_replace_hook(
+                    exported_program.graph_signature.get_replace_hook()
+                ):
+                    result = _pass.call(exported_program)
+                modified = modified or result.modified
+                if modified and self.strategy == PassStrategy.RESTART:
+                    break
+
+            if not modified:
+                break
+            step += 1
+
+            assert (
+                step < MAXIMUM_STEP_COUNT
+            ), f"Loop iterated for {MAXIMUM_STEP_COUNT} times. Circular loop is suspected in {self.passes}"