tico 0.1.0.dev250411__py3-none-any.whl

tico/utils/register_custom_op.py

@@ -0,0 +1,562 @@
+# 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, Optional
+
+import torch
+from torch._subclasses.fake_tensor import FakeTensor
+from torch.library import custom_op, register_fake
+
+
+# Note that an operator assumes input tensor has NHWC format.
+def CircleResizeNearestNeighbor():
+    @custom_op("circle_custom::resize_nearest_neighbor", mutates_args=())
+    def resize_nearest_neighbor(input_: torch.Tensor, size: List[int]) -> torch.Tensor:
+        input_size = input_.size()
+        H = input_size[1]
+        W = input_size[2]
+        H_scale_factor = size[1] / H
+        W_scale_factor = size[2] / W
+        if H_scale_factor != W_scale_factor:
+            raise RuntimeError("Scale factor of H and W should be same.")
+        return torch.nn.functional.interpolate(
+            input_, scale_factor=H_scale_factor, mode="nearest"
+        )
+
+    @register_fake("circle_custom::resize_nearest_neighbor")
+    def _(input_: torch.Tensor, size: List[int]):
+        shape = list(input_.size())
+        new_shape = [shape[0]] + list(size) + [shape[3]]
+        result = torch.empty(new_shape, dtype=input_.dtype)
+        return result
+
+
+def CircleConv2d():
+    """
+    Note that this op follows the input spec of `aten.conv2d.default` whose number
+     of arguments meets (2 <= node.args <= 7) condition.
+
+    [RESTRICTION]
+      Therefore, I tried to define a spec of conv2d as conv2d(input, weight, *args).
+      But, custom operators in torch do not support positional-only args. So, I set it
+       them as None by default.
+    """
+
+    @custom_op("circle_custom::conv2d", mutates_args=())
+    def conv2d(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = [0, 0] if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+        groups = 1 if groups is None else groups
+
+        if groups != 1:
+            raise RuntimeError(
+                f"CircleConv2d only supports 1 'groups'. the node's groups: {groups}"
+            )
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHWI_to_OIHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        OIHW_weight = torch.ops.aten.permute.default(weight, OHWI_to_OIHW)
+
+        args = [NCHW_input, OIHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::conv2d")
+    def _(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ):
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = [0, 0] if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+        groups = 1 if groups is None else groups
+        if groups != 1:
+            raise RuntimeError(
+                f"CircleConv2d only supports 1 'groups'. the node's groups: {groups}"
+            )
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHWI_to_OIHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        OIHW_weight = torch.ops.aten.permute.default(weight, OHWI_to_OIHW)
+
+        args = [NCHW_input, OIHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+
+def CircleConv2dPadding():
+    """
+    Almost same with `CircleConv2d` except padding arugment is a string type.
+
+    Q) Why create another custom op rather than make `CircleConv2d` cover multiple padding type?
+    A) `padding` with Optional[Union[List[int], str]] type is not allowed in torch.
+    """
+
+    @custom_op("circle_custom::conv2d.padding", mutates_args=())
+    def conv2d_padding(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[str] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = "valid" if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+        groups = 1 if groups is None else groups
+        if groups != 1:
+            raise RuntimeError(
+                f"CircleConv2d only supports 1 'groups'. the node's groups: {groups}"
+            )
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHWI_to_OIHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        OIHW_weight = torch.ops.aten.permute.default(weight, OHWI_to_OIHW)
+
+        args = [NCHW_input, OIHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.padding(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::conv2d.padding")
+    def _(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[str] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ):
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = "valid" if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+        groups = 1 if groups is None else groups
+        if groups != 1:
+            raise RuntimeError(
+                f"CircleConv2d only supports 1 'groups'. the node's groups: {groups}"
+            )
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHWI_to_OIHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        OIHW_weight = torch.ops.aten.permute.default(weight, OHWI_to_OIHW)
+
+        args = [NCHW_input, OIHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.padding(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+
+def CircleDepthwiseConv2d():
+    """
+    Note that this op follows the input spec of `aten.conv2d.default` whose number
+     of arguments meets (2 <= node.args <= 7) condition.
+
+    [RESTRICTION]
+      Therefore, I tried to define a spec of conv2d as conv2d(input, weight, *args).
+      But, custom operators in torch do not support positional-only args. So, I set it
+       them as None by default.
+    """
+
+    @custom_op("circle_custom::depthwise_conv2d", mutates_args=())
+    def depthwise_conv2d(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = [0, 0] if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+
+        assert groups and groups > 1
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHW1_to_1OHW = [3, 0, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        _1OHW_weight = torch.ops.aten.permute.default(weight, OHW1_to_1OHW)
+
+        args = [NCHW_input, _1OHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::depthwise_conv2d")
+    def _(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ):
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = [0, 0] if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+
+        assert groups and groups > 1
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHW1_to_1OHW = [3, 0, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        _1OHW_weight = torch.ops.aten.permute.default(weight, OHW1_to_1OHW)
+
+        args = [NCHW_input, _1OHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+
+def CircleDepthwiseConv2dPadding():
+    @custom_op("circle_custom::depthwise_conv2d.padding", mutates_args=())
+    def depthwise_conv2d_padding(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[str] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = "valid" if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+
+        assert groups and groups > 1
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHW1_to_1OHW = [3, 0, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        _1OHW_weight = torch.ops.aten.permute.default(weight, OHW1_to_1OHW)
+
+        args = [NCHW_input, _1OHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.padding(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::depthwise_conv2d.padding")
+    def _(
+        input_: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Optional[List[int]] = None,
+        padding: Optional[str] = None,
+        dilation: Optional[List[int]] = None,
+        groups: Optional[int] = None,
+    ):
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        stride = [1, 1] if stride is None else stride
+        padding = "valid" if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+
+        assert groups and groups > 1
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        OHW1_to_1OHW = [3, 0, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+        _1OHW_weight = torch.ops.aten.permute.default(weight, OHW1_to_1OHW)
+
+        args = [NCHW_input, _1OHW_weight, bias, stride, padding, dilation, groups]
+        NCHW_output = torch.ops.aten.conv2d.padding(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+
+def CircleMaxPool2D():
+    """
+    Note that this op follows the input spec of `aten.max_pool2d_with_indices.default` whose number
+     of arguments meets (3 <= node.args <= 6) condition.
+
+    [RESTRICTION]
+      Custom operators in torch do not support positional-only args. So, I set it
+       them as None by default.
+    """
+
+    @custom_op("circle_custom::maxpool2d", mutates_args=())
+    def maxpool2d(
+        input_: torch.Tensor,
+        kernel_size: List[int],
+        stride: List[int],
+        padding: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        ceil_mode: Optional[bool] = None,
+    ) -> torch.Tensor:
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        padding = [0, 0] if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+        ceil_mode = False if ceil_mode is None else ceil_mode
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+
+        args = [NCHW_input, kernel_size, stride, padding, dilation, ceil_mode]
+        NCHW_output = torch.ops.aten.max_pool2d_with_indices.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        # use first output
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output[0], NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::maxpool2d")
+    def _(
+        input_: torch.Tensor,
+        kernel_size: List[int],
+        stride: List[int],
+        padding: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        ceil_mode: Optional[bool] = None,
+    ):
+        """
+        Set default values.
+
+        Custom operators have limited types when it comes to default values.
+        So, let's set them by None in input specs, and then, set it by default values.
+        https://github.com/pytorch/pytorch/blob/6b05aafc/torch/_library/infer_schema.py#L131-L144
+        """
+        padding = [0, 0] if padding is None else padding
+        dilation = [1, 1] if dilation is None else dilation
+        ceil_mode = False if ceil_mode is None else ceil_mode
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+
+        args = [NCHW_input, kernel_size, stride, padding, dilation, ceil_mode]
+        NCHW_output = torch.ops.aten.max_pool2d_with_indices.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        # use first output
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output[0], NCHW_to_NHWC)
+
+        return NHWC_output
+
+
+def CircleAvgPool2D():
+    @custom_op("circle_custom::avgpool2d", mutates_args=())
+    def avgpool2d(
+        input_: torch.Tensor,
+        kernel_size: List[int],
+        stride: List[int],
+        padding: Optional[List[int]] = None,
+        ceil_mode: Optional[bool] = None,
+        count_include_pad: Optional[bool] = None,
+        divisor_override: Optional[int] = None,
+    ) -> torch.Tensor:
+        padding = [0, 0] if padding is None else padding
+        ceil_mode = False if ceil_mode is None else ceil_mode
+        count_include_pad = True if count_include_pad is None else count_include_pad
+        divisor_override = None if divisor_override is None else divisor_override
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+
+        args = [
+            NCHW_input,
+            kernel_size,
+            stride,
+            padding,
+            ceil_mode,
+            count_include_pad,
+            divisor_override,
+        ]
+        NCHW_output = torch.ops.aten.avg_pool2d.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::avgpool2d")
+    def _(
+        input_: torch.Tensor,
+        kernel_size: List[int],
+        stride: List[int],
+        padding: Optional[List[int]] = None,
+        ceil_mode: Optional[bool] = None,
+        count_include_pad: Optional[bool] = None,
+        divisor_override: Optional[int] = None,
+    ):
+        padding = [0, 0] if padding is None else padding
+        ceil_mode = False if ceil_mode is None else ceil_mode
+        count_include_pad = True if count_include_pad is None else count_include_pad
+        divisor_override = None if divisor_override is None else divisor_override
+
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+
+        args = [
+            NCHW_input,
+            kernel_size,
+            stride,
+            padding,
+            ceil_mode,
+            count_include_pad,
+            divisor_override,
+        ]
+        NCHW_output = torch.ops.aten.avg_pool2d.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+
+def CircleInstanceNorm():
+    @custom_op("circle_custom::instance_norm", mutates_args=())
+    def instance_norm(
+        input_: torch.Tensor,
+        weight: Optional[torch.Tensor] = None,
+        bias: Optional[torch.Tensor] = None,
+        running_mean: Optional[torch.Tensor] = None,
+        running_var: Optional[torch.Tensor] = None,
+        use_input_stats: bool = False,
+        momentum: float = 0.1,
+        eps: float = 1e-05,
+        cudnn_enabled: bool = False,
+    ) -> torch.Tensor:
+        NHWC_to_NCHW = [0, 3, 1, 2]
+        NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
+
+        args = [NCHW_input, weight, bias, None, None, False, momentum, eps, False]
+        NCHW_output = torch.ops.aten.instance_norm.default(*args)
+        NCHW_to_NHWC = [0, 2, 3, 1]
+        NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
+
+        return NHWC_output
+
+    @register_fake("circle_custom::instance_norm")
+    def _(
+        input: FakeTensor,
+        weight: Optional[FakeTensor] = None,
+        bias: Optional[FakeTensor] = None,
+        running_mean: Optional[FakeTensor] = None,
+        running_var: Optional[FakeTensor] = None,
+        use_input_stats: bool = False,
+        momentum: float = 0.1,
+        eps: float = 1e-05,
+        cudnn_enabled: bool = False,
+    ):
+        # shape is preserved
+        return input.new_empty(input.size())
+
+
+# Add custom ops to the torch namespace
+def RegisterOps():
+    CircleResizeNearestNeighbor()
+    CircleDepthwiseConv2d()
+    CircleDepthwiseConv2dPadding()
+    CircleConv2d()
+    CircleConv2dPadding()
+    CircleMaxPool2D()
+    CircleAvgPool2D()
+    CircleInstanceNorm()

tico/utils/trace_decorators.py

@@ -0,0 +1,101 @@
+# 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 functools import wraps
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.utils.diff_graph import capture, capture_const, log, log_const
+from tico.utils.passes import PassBase
+
+
+def trace_const_diff_on_pass(cls):
+    """Decorator for PassBase to trace const diff"""
+
+    assert issubclass(cls, PassBase), type(cls)
+
+    def _call_traced(fn):
+        @wraps(fn)
+        def wrapped(*args):
+            _, exported_program = args
+            assert isinstance(exported_program, ExportedProgram)
+            graph_module = exported_program.graph_module
+            assert isinstance(graph_module, torch.fx.GraphModule), type(graph_module)
+            capture_const(exported_program)
+            ret = fn(*args)
+            log_const(exported_program, title=str(cls.__name__), recapture=False)
+            return ret
+
+        return wrapped
+
+    # replace call function it with traced version
+    for key, val in vars(cls).items():
+        if key == "call":
+            setattr(cls, key, _call_traced(val))
+    return cls
+
+
+def trace_graph_diff_on_pass(cls):
+    """Decorator for PassBase to trace graph diff"""
+
+    assert issubclass(cls, PassBase), type(cls)
+
+    def _call_traced(fn):
+        @wraps(fn)
+        def wrapped(*args):
+            _, exported_program = args
+            assert isinstance(exported_program, ExportedProgram)
+            graph_module = exported_program.graph_module
+            assert isinstance(graph_module, torch.fx.GraphModule), type(graph_module)
+            capture(graph_module.graph)
+            ret = fn(*args)
+            log(graph_module.graph, title=str(cls.__name__), recapture=False)
+            return ret
+
+        return wrapped
+
+    # replace call function it with traced version
+    for key, val in vars(cls).items():
+        if key == "call":
+            setattr(cls, key, _call_traced(val))
+    return cls
+
+
+def trace_const_diff_on_func(fn):
+    """Decorator for function to trace const diff"""
+
+    @wraps(fn)
+    def wrapped(ep: torch.export.ExportedProgram):
+        assert isinstance(ep, torch.export.ExportedProgram)
+        capture_const(ep)
+        ret = fn(ep)
+        log_const(ret, title=str(fn.__name__), recapture=False)
+        return ret
+
+    return wrapped
+
+
+def trace_graph_diff_on_func(fn):
+    """Decorator for function to trace graph diff"""
+
+    @wraps(fn)
+    def wrapped(ep: torch.export.ExportedProgram):
+        assert isinstance(ep, torch.export.ExportedProgram)
+        capture(ep.graph)
+        ret = fn(ep)
+        log(ret.graph, title=str(fn.__name__), recapture=False)
+        return ret
+
+    return wrapped