tico 0.1.0.dev250411__py3-none-any.whl

tico/passes/convert_to_relu6.py

@@ -0,0 +1,180 @@
+# 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 List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.export import ExportedProgram
+
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import ClampArgs, HardTanhArgs
+
+
+class Converter:  # type: ignore[empty-body]
+    def __init__(self):
+        super().__init__()
+
+    def match(self, node) -> bool:  # type: ignore[empty-body]
+        return False
+
+    def convert(self, exported_program, node) -> None:  # type: ignore[empty-body]
+        pass
+
+
+class ConvertHardTanhToReLU6(Converter):
+    def __init__(self):
+        super().__init__()
+
+    def match(self, node) -> bool:
+        if node.target == torch.ops.aten.hardtanh.default:
+            args = HardTanhArgs(*node.args, **node.kwargs)
+            min_val = args.min_val
+            max_val = args.max_val
+
+            # NOTE: int and float are both covered by pytorch implicit type conversion
+            return min_val == 0.0 and max_val == 6.0
+
+        return False
+
+    def convert(self, exported_program, node):
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        args = HardTanhArgs(*node.args, **node.kwargs)
+        input = args.input
+
+        with graph.inserting_after(node):
+            relu_node = graph.call_function(torch.ops.aten.relu6.default, args=(input,))
+            node.replace_all_uses_with(relu_node, propagate_meta=True)
+
+
+class ConvertClampToReLU6(Converter):
+    def __init__(self):
+        super().__init__()
+
+    def match(self, node) -> bool:
+        if node.target == torch.ops.aten.clamp.default:
+            args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            min_val = args.min
+            max_val = args.max
+
+            # NOTE: int and float are both covered by pytorch implicit type conversion
+            return min_val == 0 and max_val == 6
+
+        return False
+
+    def convert(self, exported_program, node):
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        with graph.inserting_after(node):
+            relu_node = graph.call_function(torch.ops.aten.relu6.default, args=(input,))
+            node.replace_all_uses_with(relu_node, propagate_meta=True)
+
+
+class ConvertDoubleClampsToReLU6(Converter):
+    def __init__(self):
+        super().__init__()
+
+    def match(self, node) -> bool:
+        """
+        This pass matches the pattern of two clamps where it equals to clamp which has a min value of 0 and a max value of 6.
+
+                    (equivalent)
+        input                        input
+        |                            |
+        node_prev (min, max)         node   (0, 6)
+        |                            |
+        node      (min', max')       |
+        |                            |
+        output                       output
+
+        *where max(min, min') == 0 and min(max, max') == 6 so that it equivalents to clamp(input, 0, 6)
+
+        TODO Make this step more generic. For now we only support the case above.
+        """
+        if not node.target == torch.ops.aten.clamp.default:
+            return False
+
+        args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        node_prev = args.input
+        min_val = args.min if args.min is not None else float("-inf")
+        max_val = args.max if args.max is not None else float("inf")
+
+        if not node_prev.target == torch.ops.aten.clamp.default:
+            return False
+
+        prev_args = ClampArgs(*node_prev.args, **node_prev.kwargs)  # type: ignore[arg-type]
+        min_val_prev = prev_args.min if prev_args.min is not None else float("-inf")
+        max_val_prev = prev_args.max if prev_args.max is not None else float("inf")
+
+        # NOTE: int and float are both covered by pytorch implicit type conversion
+        if max(min_val, min_val_prev) == 0 and min(max_val, max_val_prev) == 6:
+            return True
+
+        return False
+
+    def convert(self, exported_program, node):
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        prev_node = args.input
+        prev_args = ClampArgs(*prev_node.args, **prev_node.kwargs)  # type: ignore[arg-type]
+        input = prev_args.input
+
+        with graph.inserting_after(node):
+            relu_node = graph.call_function(torch.ops.aten.relu6.default, args=(input,))
+            node.replace_all_uses_with(relu_node, propagate_meta=True)
+
+
+@trace_graph_diff_on_pass
+class ConvertToReLU6(PassBase):
+    def __init__(self):
+        super().__init__()
+        self.converters: List[Converter] = [
+            ConvertHardTanhToReLU6(),
+            ConvertClampToReLU6(),
+            ConvertDoubleClampsToReLU6(),
+        ]
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+        for node in graph.nodes:
+            if not node.op == "call_function":
+                continue
+
+            for converter in self.converters:
+                if not converter.match(node):
+                    continue
+
+                converter.convert(exported_program, node)
+                modified = True
+                logger.debug(f"{node.name} is replaced with ReLU6 operator")
+                break
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/passes/decompose_addmm.py

@@ -0,0 +1,127 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.circle_mapping import extract_shape
+from tico.utils import logging
+from tico.utils.graph import add_placeholder
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import set_new_meta_val
+from tico.utils.validate_args_kwargs import AddmmArgs
+
+
+@trace_graph_diff_on_pass
+class DecomposeAddmm(PassBase):
+    """
+    Let's decompose addmm to add + mul + matmul.
+
+    [BEFORE]
+
+    input   mat1    mat2    beta    alpha
+    |       |       |       |       |
+    --------------addmm--------------
+                    |
+                    out
+
+    [AFTER]
+
+    input   beta    mat1    mat2    alpha
+    |       |       |       |        |
+    ---mul---       ---mm----        |
+        |               |            |
+        |               -----mul-----
+        |                     |
+        ---------add----------
+                  |
+                  out
+
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        gm = exported_program.graph_module
+        graph: torch.fx.Graph = gm.graph
+        modified = False
+
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+
+            if node.target in [
+                torch.ops.aten.addmm.default,
+            ]:
+                args = AddmmArgs(*node.args, **node.kwargs)
+                input = args.input
+                mat1 = args.mat1
+                mat2 = args.mat2
+                beta = args.beta
+                alpha = args.alpha
+
+                with graph.inserting_before(node):
+                    # out = beta * input + alpha * (mat1 @ mat2)
+                    matmul = graph.call_function(
+                        torch.ops.aten.mm.default, (mat1, mat2)
+                    )
+                    set_new_meta_val(matmul)
+
+                    if beta == 1:
+                        bias: torch.fx.Node | torch.Tensor = input
+                    elif beta == 0:
+                        bias = add_placeholder(
+                            exported_program,
+                            torch.zeros(extract_shape(input)),
+                            f"{node.name}_beta_zeros",
+                        )
+                    else:
+                        bias = graph.call_function(
+                            torch.ops.aten.mul.Tensor, (input, beta)
+                        )
+
+                    if alpha == 1:
+                        scaled_matmul: torch.fx.Node | torch.Tensor = matmul
+                    elif alpha == 0:
+                        scaled_matmul = add_placeholder(
+                            exported_program,
+                            torch.zeros(extract_shape(matmul)),
+                            f"{node.name}_alpha_zeros",
+                        )
+                    else:
+                        scaled_matmul = graph.call_function(
+                            torch.ops.aten.mul.Tensor, (matmul, alpha)
+                        )
+
+                    result = graph.call_function(
+                        torch.ops.aten.add.Tensor, (bias, scaled_matmul)
+                    )
+
+                node.replace_all_uses_with(result, propagate_meta=True)
+
+                modified = True
+
+        gm.graph.eliminate_dead_code()
+        gm.graph.lint()
+        gm.recompile()
+
+        return PassResult(modified)

tico/passes/decompose_batch_norm.py

@@ -0,0 +1,198 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.circle_mapping import extract_shape
+from tico.utils import logging
+from tico.utils.errors import NotYetSupportedError
+from tico.utils.graph import (
+    add_placeholder,
+    get_first_user_input,
+    get_torch_buffer_value,
+    get_torch_param_value,
+    is_torch_buffer,
+    is_torch_param,
+)
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import fill_meta_val
+from tico.utils.validate_args_kwargs import NativeBatchNormLegitNoTrainingArgs
+
+
+def insert_node(graph: torch.fx.Graph, operation, args):
+    new_node = graph.call_function(operation, args)
+
+    return new_node
+
+
+@trace_graph_diff_on_pass
+class DecomposeBatchNorm(PassBase):
+    """
+    [BatchNorm]
+
+    The op can be decomposed to a single aten.mul and a single aten.add because mean and
+      var are fixed during evaluation.
+
+      W = (weight / sqrt(var + eps))
+      B = bias - (mean * weight) / sqrt(var + eps)
+      Y = X * W + B
+
+    [before]
+
+        input (tensor, weight, bias, running_mean, running_var, momentum, eps)
+          |
+      BatchNorm
+          |
+        output
+
+    [after]
+
+        input
+       (tensor)
+          |    W
+          |   /
+         mul
+          |   B
+          |  /
+         add
+          |
+        output
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        gm = exported_program.graph_module
+        graph: torch.fx.Graph = gm.graph
+        modified = False
+
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+
+            if node.target in [
+                torch.ops.aten._native_batch_norm_legit_no_training.default,
+            ]:
+                args = NativeBatchNormLegitNoTrainingArgs(*node.args)
+                input_ = args.input
+                weight = args.weight
+                bias = args.bias
+                running_mean = args.running_mean
+                running_var = args.running_var
+                eps = args.eps
+
+                if not running_mean:
+                    raise NotYetSupportedError(
+                        f"running_mean=None is not supported yet"
+                    )
+                if not running_var:
+                    raise NotYetSupportedError(f"running_var=None is not supported yet")
+
+                """
+                Only support the cases generated from torch.nn.BatchNorm2d module,
+                 for which, let's checks if weight and bias are parameters and 
+                 running_mean and running_var are buffers.
+                """
+                if weight and not is_torch_param(weight, exported_program):
+                    continue
+                if bias and not is_torch_param(bias, exported_program):
+                    continue
+                if not is_torch_buffer(running_mean, exported_program):
+                    continue
+                if not is_torch_buffer(running_var, exported_program):
+                    continue
+
+                input_shape = extract_shape(input_)
+                assert len(input_shape) == 4
+                C = input_shape[1]
+
+                weight_value = (
+                    get_torch_param_value(weight, exported_program)
+                    if weight
+                    else torch.tensor([1] * C)
+                )
+                bias_value = (
+                    get_torch_param_value(bias, exported_program)
+                    if bias
+                    else torch.tensor([0] * C)
+                )
+                mean_value = get_torch_buffer_value(running_mean, exported_program)
+                var_value = get_torch_buffer_value(running_var, exported_program)
+
+                assert isinstance(weight_value, torch.Tensor)
+                assert isinstance(bias_value, torch.Tensor)
+                assert isinstance(mean_value, torch.Tensor)
+                assert isinstance(var_value, torch.Tensor)
+
+                assert (
+                    weight_value.shape
+                    == bias_value.shape
+                    == mean_value.shape
+                    == var_value.shape
+                )
+                # Calculate constants for mul and add
+                mul_const = weight_value / torch.sqrt(var_value + eps)
+                add_const = bias_value - (mul_const * mean_value)
+                # N, C, H, W
+                assert len(mul_const) == len(add_const) == C
+                # reshape along with channel dimension
+                mul_const = mul_const.view(1, mul_const.shape[0], 1, 1)
+                add_const = add_const.view(1, add_const.shape[0], 1, 1)
+
+                # 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.
+                with exported_program.graph.inserting_before(
+                    get_first_user_input(exported_program)
+                ):
+                    mul_const_node = add_placeholder(
+                        exported_program,
+                        mul_const,
+                        prefix=f"{node.name}_mul_const",
+                    )
+                    add_const_node = add_placeholder(
+                        exported_program,
+                        add_const,
+                        prefix=f"{node.name}_add_const",
+                    )
+
+                with gm.graph.inserting_before(node):
+                    mul = graph.call_function(
+                        torch.ops.aten.mul.Tensor,
+                        args=(input_, mul_const_node),
+                    )
+                    add = graph.call_function(
+                        torch.ops.aten.add.Tensor,
+                        args=(mul, add_const_node),
+                    )
+                    # Not set meta for propagating replacing get_item's meta.
+                get_item, *_ = node.users.keys()
+                get_item.replace_all_uses_with(add, propagate_meta=True)
+
+                fill_meta_val(exported_program)
+                modified = True
+
+        gm.graph.eliminate_dead_code()
+        gm.graph.lint()
+        gm.recompile()
+
+        return PassResult(modified)

tico/passes/decompose_fake_quantize.py

@@ -0,0 +1,126 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+# To import torch.ops.quantized_decomposed related operator
+from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib
+from torch.export import ExportedProgram
+
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import FakeQuantizePerChannelArgs
+
+
+def get_quant_type(min: int, max: int) -> torch.dtype:
+    if min == 0 and max == 255:
+        return torch.uint8
+    if min == -32768 and max == 32767:
+        return torch.int16
+    if min == -32767 and max == 32767:
+        return torch.int16
+
+    raise RuntimeError("Not supported min/max values")
+
+
+@trace_graph_diff_on_pass
+class DecomposeFakeQuantize(PassBase):
+    """
+    Decompose fake quantize operator to quant/dequant operators.
+    Otherwise, it can't be converted to the edge IR because fake quantize operator is not Aten Canonical.
+
+    [Before]
+    def forward(self, x):
+        fake_quantize_per_tensor_affine = torch.ops.aten.fake_quantize_per_tensor_affine.default(tensor, scale, zero_p, quant_min, quant_max);  x = None
+        return (fake_quantize_per_tensor_affine,)
+
+    [After]
+    def forward(self, x):
+        quantize_per_tensor_default = torch.ops.quantized_decomposed.quantize_per_tensor.default(tensor, scale, zero_p, quant_min, quant_max, dtype = ${torch.dtype});  x = None
+        dequantize_per_tensor_default = torch.ops.quantized_decomposed.dequantize_per_tensor.default(quantize_per_tensor_default, scale, zero_p, quant_min, quant_max, dtype = ${torch.dtype});  quantize_per_tensor_default = None
+        return (dequantize_per_tensor_default,)
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+        modified = False
+
+        gm = exported_program.graph_module
+        qd = torch.ops.quantized_decomposed  # type: ignore[return]
+        for node in gm.graph.nodes:
+            if node.op == "call_function" and node.target in [
+                torch.ops.aten.fake_quantize_per_tensor_affine.default
+            ]:
+                # tensor, scale, zero_p, quant_min, quant_max
+                assert len(node.args) == 5
+                _, _, _, quant_min, quant_max = node.args
+
+                quant_kwargs = {
+                    **node.kwargs,
+                    **{"dtype": get_quant_type(quant_min, quant_max)},
+                }
+                with gm.graph.inserting_before(node):
+                    quant = gm.graph.call_function(
+                        qd.quantize_per_tensor.default,
+                        args=node.args,
+                        kwargs=quant_kwargs,
+                    )
+                    dequnt = gm.graph.call_function(
+                        qd.dequantize_per_tensor.default,
+                        args=(quant, *quant.args[1:]),
+                        kwargs=quant.kwargs,
+                    )
+                    # Not set meta for propagating replacing node's meta.
+                    node.replace_all_uses_with(dequnt, propagate_meta=True)
+                modified = True
+
+            if node.op == "call_function" and node.target in [
+                torch.ops.aten.fake_quantize_per_channel_affine.default
+            ]:
+                fq_args = FakeQuantizePerChannelArgs(*node.args, **node.kwargs)
+                quant_min = fq_args.quant_min
+                quant_max = fq_args.quant_max
+
+                quant_kwargs = {
+                    **node.kwargs,
+                    **{"dtype": get_quant_type(quant_min, quant_max)},
+                }
+                with gm.graph.inserting_before(node):
+                    quant = gm.graph.call_function(
+                        qd.quantize_per_channel.default,
+                        args=node.args,
+                        kwargs=quant_kwargs,
+                    )
+                    dequnt = gm.graph.call_function(
+                        qd.dequantize_per_channel.default,
+                        args=(quant, *quant.args[1:]),
+                        kwargs=quant.kwargs,
+                    )
+                    # Not set meta for propagating replacing node's meta.
+                    node.replace_all_uses_with(dequnt, propagate_meta=True)
+                modified = True
+
+        gm.graph.eliminate_dead_code()
+        gm.graph.lint()
+        gm.recompile()
+
+        return PassResult(modified)