tico 0.1.0.dev250427__py3-none-any.whl → 0.1.0.dev250429__py3-none-any.whl

tico/__init__.py

@@ -22,7 +22,7 @@ from tico.config import CompileConfigV1, get_default_config
 from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = '0.1.0.dev250427'
+__version__ = '0.1.0.dev250429'
 
 
 if Version(torch.__version__) < Version("2.5.0"):

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

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/register_custom_op.py

@@ -18,6 +18,7 @@ import torch
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.library import custom_op, register_fake
 
+from tico.utils.mx.mx_ops import _quantize_mx
 
 # Note that an operator assumes input tensor has NHWC format.
 def CircleResizeNearestNeighbor():
@@ -550,6 +551,47 @@ def CircleInstanceNorm():
         return input.new_empty(input.size())
 
 
+def CircleQuantizeMX():
+    # This operator conducts fake-quantization of microscaling
+    # NOTE Why using "quantize"_mx not "fake_quantize"_mx?
+    # To align with function name of microxcaling repo.
+    # https://github.com/microsoft/microxcaling/blob/v1.1.0/mx/mx_ops.py#L173
+    @custom_op("circle_custom::quantize_mx", mutates_args=())
+    def quantize_mx(
+        input_: torch.Tensor,
+        elem_format: str,
+        axis: int,
+        shared_exp_method: str = "max",
+        round: str = "nearest",
+    ) -> torch.Tensor:
+        if elem_format == "int8":
+            scale_bits = 8
+            block_size = 32
+        else:
+            raise RuntimeError(f"Unsupported elem_format in quantize_mx: {elem_format}")
+
+        result = _quantize_mx(
+            input_,
+            scale_bits=scale_bits,
+            elem_format=elem_format,
+            axes=[axis],
+            block_size=block_size,
+            shared_exp_method=shared_exp_method,
+            round=round,
+        )
+        return result
+
+    @register_fake("circle_custom::quantize_mx")
+    def _(
+        input_: torch.Tensor,
+        elem_format: str,
+        axis: int,
+        shared_exp_method: str = "max",  # Fixed
+        round: str = "nearest",  # Fixed
+    ) -> torch.Tensor:
+        return input_
+
+
 # Add custom ops to the torch namespace
 def RegisterOps():
     CircleResizeNearestNeighbor()
@@ -560,3 +602,4 @@ def RegisterOps():
     CircleMaxPool2D()
     CircleAvgPool2D()
     CircleInstanceNorm()
+    CircleQuantizeMX()

{tico-0.1.0.dev250427.dist-info → tico-0.1.0.dev250429.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250427
+Version: 0.1.0.dev250429
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250427.dist-info → tico-0.1.0.dev250429.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=_2AKXhU80F-fxtXp0EYbYtJt0epyucAnROQ6MZvqQz0,1181
+tico/__init__.py,sha256=VqDTIcftfqxHmepF7PI4HaijESf5C9i2qtOFcdJ9DU8,1181
 tico/pt2_to_circle.py,sha256=PPmFNw20jw2Z2VyM3ln9pX__jTzBOAZiv0gT5a-p-Y8,2666
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=anwOiJFkUxUi7Cef573JgQcjk6S-FSi6O_TLjYASW-g,1244
@@ -184,13 +184,17 @@ tico/utils/logging.py,sha256=IlbBWscsaHidI0dNqro1HEXAbIcbkR3BD5ukLy2m95k,1286
 tico/utils/model.py,sha256=Uqc92AnJXQ2pbvctS2z2F3Ku3yNrwXZ9O33hZVis7is,1250
 tico/utils/padding.py,sha256=GGO27VbaOvtaMYLDrSaKv7uxjeet566aMJD0PyYeMvQ,1484
 tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
-tico/utils/register_custom_op.py,sha256=FbMcrg8o5vKWC_aoVxL2GrIcR14KFi1yKG0mFGqXkPY,21595
+tico/utils/register_custom_op.py,sha256=iRQvdqlBqrJxq_pNkvJyDIJD_SYtCUl88wwbbuvSwlk,22952
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
 tico/utils/utils.py,sha256=pybDU1LoNhjEplANig11lboX9yzYRkvFCSmyYth_2Do,10359
 tico/utils/validate_args_kwargs.py,sha256=krT68b5CfBI9rxBIOsgYSy0LfEJqLfKfRikkp8ep9oQ,24726
-tico-0.1.0.dev250427.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250427.dist-info/METADATA,sha256=zhzQp-7OlEy-RuLApjkGW4eZmf0rEnr_rZwUAcFXNGo,7353
-tico-0.1.0.dev250427.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250427.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250427.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250427.dist-info/RECORD,,
+tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
+tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
+tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
+tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
+tico-0.1.0.dev250429.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250429.dist-info/METADATA,sha256=9CsjedscOI7c1_UCpwOmht40v6n6h-qTsGowGQ5mawo,7353
+tico-0.1.0.dev250429.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250429.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250429.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250429.dist-info/RECORD,,