tico 0.1.0.dev250609__py3-none-any.whl → 0.1.0.dev250611__py3-none-any.whl

tico/passes/decompose_slice_scatter.py

@@ -25,7 +25,7 @@ from tico.serialize.circle_mapping import extract_shape
 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.utils import enforce_type
+from tico.utils.utils import enforce_type, is_target_node
 
 
 @trace_graph_diff_on_pass
@@ -82,9 +82,7 @@ class DecomposeSliceScatter(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op != "call_function":
-                continue
-            if node.target != torch.ops.aten.slice_scatter.default:
+            if not is_target_node(node, torch.ops.aten.slice_scatter.default):
                 continue
 
             @enforce_type

tico/passes/extract_dtype_kwargs.py

@@ -23,6 +23,7 @@ from torch.utils import _pytree as pytree
 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.utils import is_target_node
 
 
 def _extract_to_output(node: torch.fx.Node, graph: torch.fx.Graph) -> bool:
@@ -107,7 +108,7 @@ class ExtractDtypeKwargsPass(PassBase):
         graph: torch.fx.Graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function" or node.target not in self.target_ops:
+            if not is_target_node(node, list(self.target_ops.keys())):
                 continue
             if "dtype" not in node.kwargs:
                 continue

tico/passes/fuse_leading_unsqueeze_reshape.py

@@ -0,0 +1,107 @@
+# 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 Sequence
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.passes import ops
+from tico.serialize.circle_mapping import extract_shape
+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.utils import is_target_node
+from tico.utils.validate_args_kwargs import PermuteArgs, ReshapeArgs
+
+
+def _is_leading_unsqueeze(target: Sequence[int], permuted: Sequence[int]) -> bool:
+    """
+    True if `target` == [1]*k + permuted,  k>=1.
+    """
+    k = len(target) - len(permuted)
+    return (
+        k > 0 and all(d == 1 for d in target[:k]) and list(target[k:]) == list(permuted)
+    )
+
+
+@trace_graph_diff_on_pass
+class FuseLeadingUnsqueezeReshape(PassBase):
+    """
+    Fuse reshape → permute → reshape where the second reshape only
+    prepends one-sized dims (unsqueeze) to the permuted tensor.
+
+    [BEFORE]
+        x - aten.reshape(s1) - aten.permute(p) - aten.reshape([1]*k + p(s1))
+    [AFTER]
+        x - aten.reshape([1]*k + s1) - aten.permute(list(range(k)) + [d+k for d in p])
+    """
+
+    def call(self, ep: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        gm = ep.graph_module
+        graph = gm.graph
+        modified = False
+        for reshape_back in graph.nodes:
+            if not is_target_node(reshape_back, ops.aten.reshape):
+                continue
+            reshape_back_args = ReshapeArgs(*reshape_back.args, **reshape_back.kwargs)  # type: ignore[arg-type]
+            permute, reshape_back_size = reshape_back_args.input, reshape_back_args.size
+
+            if not is_target_node(permute, ops.aten.permute):
+                continue
+            permute_args = PermuteArgs(*permute.args, **permute.kwargs)  # type: ignore[arg-type]
+            reshape_front, permute_dims = permute_args.input, permute_args.dims
+
+            if not is_target_node(reshape_front, ops.aten.reshape):
+                continue
+            reshape_front_args = ReshapeArgs(*reshape_front.args, **reshape_front.kwargs)  # type: ignore[arg-type]
+            reshape_front_input, reshape_front_size = (
+                reshape_front_args.input,
+                reshape_front_args.size,
+            )
+
+            # ---- condition: only leading unsqueeze ------------------------
+            back_shape = extract_shape(reshape_back)
+            permute_shape = extract_shape(permute)
+
+            if not _is_leading_unsqueeze(back_shape, permute_shape):
+                continue
+
+            # ---- create new reshape & new permute -------------------------
+            k = len(back_shape) - len(permute_shape)
+            with graph.inserting_before(permute):
+                new_shape = [1] * k + list(reshape_front_size)
+                r_new = graph.call_function(
+                    torch.ops.aten.reshape.default,
+                    args=(reshape_front_input, new_shape),
+                )
+                new_p_dims = list(range(k)) + [
+                    d + k for d in permute_dims
+                ]  # shift by k
+                p_new = graph.call_function(
+                    torch.ops.aten.permute.default, args=(r_new, new_p_dims)
+                )
+
+            reshape_back.replace_all_uses_with(p_new, propagate_meta=True)
+            modified = True
+            logger.debug(f"{reshape_back.name} is fused to {r_new.name}")
+
+        if modified:
+            graph.eliminate_dead_code()
+            graph.lint()
+            gm.recompile()
+
+        return PassResult(modified)

tico/passes/fuse_redundant_reshape_to_mean.py

@@ -25,6 +25,7 @@ from tico.serialize.circle_mapping import extract_shape
 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.utils import is_target_node
 
 
 @trace_graph_diff_on_pass
@@ -45,10 +46,7 @@ class FuseRedundantReshapeToMean(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if node.target != torch.ops.aten.mean.dim:
+            if not is_target_node(node, torch.ops.aten.mean.dim):
                 continue
 
             # If mean is being used in other nodes, do not fuse it.
@@ -56,7 +54,7 @@ class FuseRedundantReshapeToMean(PassBase):
                 continue
 
             user_node = next(iter(node.users))
-            if user_node.target not in ops.aten.reshape:
+            if not is_target_node(user_node, ops.aten.reshape):
                 continue
 
             mean_args, mean_kwargs = pytree.tree_map_only(

tico/passes/legalize_causal_mask_value.py

@@ -20,10 +20,10 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.passes import ops
-
 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.utils import is_target_node
 from tico.utils.validate_args_kwargs import AddTensorArgs
 
 
@@ -53,14 +53,9 @@ class LegalizeCausalMaskValue(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.add:
+            if not is_target_node(node, ops.aten.add):
                 continue
 
-            assert len(node.args) == 2
-
             args = AddTensorArgs(*node.args, **node.kwargs)
             input = args.input
             other = args.other

tico/passes/legalize_predefined_layout_operators.py

@@ -25,6 +25,7 @@ from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import (
     AvgPool2dArgs,
     Conv2DArgs,
@@ -363,11 +364,9 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, list(target_to_legalize_func.keys())):
                 continue
 
-            if node.target not in target_to_legalize_func:
-                continue
             modified |= target_to_legalize_func[node.target](exported_program, node)
 
         graph.eliminate_dead_code()

tico/passes/lower_pow2_to_mul.py

@@ -22,6 +22,8 @@ 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.utils import is_target_node
+from tico.utils.validate_args_kwargs import PowTensorScalarArgs
 
 
 @trace_graph_diff_on_pass
@@ -42,15 +44,11 @@ class LowerPow2ToMul(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, torch.ops.aten.pow.Tensor_Scalar):
                 continue
 
-            if node.target != torch.ops.aten.pow.Tensor_Scalar:
-                continue
-
-            assert len(node.args) == 2, len(node.args)
-            in_, exp = node.args
-            assert isinstance(in_, torch.fx.Node), type(in_)
+            args = PowTensorScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            in_, exp = args.input, args.exponent
 
             if exp != 2:
                 continue

tico/passes/lower_to_resize_nearest_neighbor.py

@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx
-from typing import Optional
-
 import torch
 from torch.export import ExportedProgram
 
@@ -26,6 +24,7 @@ from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import IndexArgs, UpsampleNearest2DVecArgs
 
 
@@ -194,13 +193,10 @@ class LowerToResizeNearestNeighbor(PassBase):
         graph_module = exported_program.graph_module
         graph = graph_module.graph
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if node.target not in [
-                torch.ops.aten.index.Tensor,
-                torch.ops.aten.upsample_nearest2d.vec,
-            ]:
+            if not is_target_node(
+                node,
+                [torch.ops.aten.index.Tensor, torch.ops.aten.upsample_nearest2d.vec],
+            ):
                 continue
 
             resize_nearest_neighbor = None

tico/passes/lower_to_slice.py

@@ -28,13 +28,12 @@ from torch._export.utils import (
 from torch.export import ExportedProgram
 
 from tico.passes import ops
-
 from tico.serialize.circle_graph import extract_shape
 from tico.utils import logging
-
 from tico.utils.graph import is_single_value_tensor
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_const_diff_on_pass
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import IndexSelectArgs, SelectCopyIntArgs
 
 
@@ -85,10 +84,7 @@ class LowerSelectCopyToSlice(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.select:
+            if not is_target_node(node, ops.aten.select):
                 continue
 
             args = SelectCopyIntArgs(*node.args, **node.kwargs)
@@ -163,11 +159,9 @@ class LowerIndexSelectToSlice(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, ops.aten.index_select):
                 continue
 
-            if not node.target in ops.aten.index_select:
-                continue
             args = IndexSelectArgs(*node.args, **node.kwargs)
             input = args.input
             dim = args.dim

tico/passes/merge_consecutive_cat.py

@@ -18,6 +18,7 @@ from tico.passes import ops
 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.utils import is_target_node
 from tico.utils.validate_args_kwargs import CatArgs
 
 
@@ -37,10 +38,7 @@ class MergeConsecutiveCat(PassBase):
         graph = graph_module.graph
         modified = False
         for cat in graph.nodes:
-            if not cat.op == "call_function":
-                continue
-
-            if not cat.target in ops.aten.cat:
+            if not is_target_node(cat, ops.aten.cat):
                 continue
 
             args = CatArgs(*cat.args, **cat.kwargs)  # type: ignore[arg-type]

tico/passes/remove_nop.py

@@ -23,6 +23,7 @@ from tico.passes import ops
 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.utils import is_target_node
 
 
 @trace_graph_diff_on_pass
@@ -51,11 +52,9 @@ class RemoveNop(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, RemoveNop.target_ops):
                 continue
 
-            if not node.target in RemoveNop.target_ops:
-                continue
             # TODO Consider memory format
             if node.target in ops.aten.clone and "memory_format" in node.kwargs:
                 if node.kwargs["memory_format"] not in [

tico/passes/remove_redundant_assert_nodes.py

@@ -17,6 +17,7 @@ from torch.export import ExportedProgram
 
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import is_target_node
 
 
 assert_node_targets = [
@@ -39,7 +40,7 @@ class RemoveRedundantAssertionNodes(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if node.op == "call_function" and node.target in assert_node_targets:
+            if is_target_node(node, assert_node_targets):
                 graph.erase_node(node)
                 modified = True
 

tico/passes/remove_redundant_expand.py

@@ -24,6 +24,8 @@ from tico.serialize.circle_mapping import extract_shape
 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.utils import is_target_node
+from tico.utils.validate_args_kwargs import ExpandArgs
 
 
 @trace_graph_diff_on_pass
@@ -42,17 +44,11 @@ class RemoveRedundantExpand(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, ops.aten.expand):
                 continue
 
-            if not node.target in ops.aten.expand:
-                continue
-
-            assert len(node.args) == 2
-
-            input, size = list(node.args)
-            assert isinstance(input, torch.fx.Node), type(input)
-            assert isinstance(size, list), type(size)
+            args = ExpandArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            input, size = args.input, args.size
 
             input_shape = extract_shape(input)
             if list(input_shape) != size:

tico/passes/remove_redundant_permute.py

@@ -14,16 +14,18 @@
 
 from typing import TYPE_CHECKING
 
+
 if TYPE_CHECKING:
     import torch.fx
 import torch
 from torch.export import ExportedProgram
 
 from tico.passes import ops
-from tico.serialize.circle_mapping import extract_shape, extract_stride
+from tico.serialize.circle_mapping import extract_shape
 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.utils import is_target_node
 from tico.utils.validate_args_kwargs import PermuteArgs
 
 
@@ -75,16 +77,15 @@ class RemoveRedundantPermutePattern1(PassBase):
         graph = graph_module.graph
         modified = False
         for permute2 in graph.nodes:
-            if not permute2.op == "call_function":
-                continue
-            if not permute2.target in ops.aten.permute:
+            if not is_target_node(permute2, ops.aten.permute):
                 continue
+
             if len(permute2.users) != 1:
                 continue
             permute2_args = PermuteArgs(*permute2.args, **permute2.kwargs)  # type: ignore[arg-type]
             permute1, permute2_dims = permute2_args.input, permute2_args.dims
 
-            if not permute1.target in ops.aten.permute:
+            if not is_target_node(permute1, ops.aten.permute):
                 continue
             if len(permute1.users) != 1:
                 continue

tico/passes/remove_redundant_reshape.py

@@ -24,7 +24,7 @@ from tico.serialize.circle_mapping import extract_shape
 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.utils import broadcastable, set_new_meta_val
+from tico.utils.utils import broadcastable, is_target_node, set_new_meta_val
 from tico.utils.validate_args_kwargs import (
     AddTensorArgs,
     PermuteArgs,
@@ -67,12 +67,10 @@ class RemoveRedundantReshapePattern1(PassBase):
         graph = graph_module.graph
         modified = False
         for reshape1 in graph.nodes:
-            if not reshape1.op == "call_function":
-                continue
-
             ### first reshape
-            if not reshape1.target in ops.aten.reshape:
+            if not is_target_node(reshape1, ops.aten.reshape):
                 continue
+
             # Assumes that other node do not use ops in the pattern for simplisity.
             if len(reshape1.users) != 1:
                 continue
@@ -86,7 +84,7 @@ class RemoveRedundantReshapePattern1(PassBase):
 
             ### permute
             permute = next(iter(reshape1.users))
-            if not permute.target in ops.aten.permute:
+            if not is_target_node(permute, ops.aten.permute):
                 continue
             if len(permute.users) != 1:
                 continue
@@ -98,14 +96,14 @@ class RemoveRedundantReshapePattern1(PassBase):
 
             ### mul
             mul = next(iter(permute.users))
-            if not mul.target in RemoveRedundantReshapePattern1.mul_ops:
+            if not is_target_node(mul, RemoveRedundantReshapePattern1.mul_ops):
                 continue
             if len(mul.users) != 1:
                 continue
 
             ### second reshape
             reshape2 = next(iter(mul.users))
-            if not reshape2.target in ops.aten.reshape:
+            if not is_target_node(reshape2, ops.aten.reshape):
                 continue
             if len(reshape2.users) != 1:
                 continue
@@ -153,11 +151,8 @@ class RemoveRedundantReshapePattern2(PassBase):
         graph = graph_module.graph
         modified = False
         for reshape1 in graph.nodes:
-            if not reshape1.op == "call_function":
-                continue
-
             ### first reshape
-            if not reshape1.target in ops.aten.reshape:
+            if not is_target_node(reshape1, ops.aten.reshape):
                 continue
             if len(reshape1.users) != 1:
                 continue
@@ -171,7 +166,7 @@ class RemoveRedundantReshapePattern2(PassBase):
 
             ### permute
             permute = next(iter(reshape1.users))
-            if not permute.target in ops.aten.permute:
+            if not is_target_node(permute, ops.aten.permute):
                 continue
             if len(permute.users) != 1:
                 continue
@@ -183,7 +178,7 @@ class RemoveRedundantReshapePattern2(PassBase):
 
             ### second reshape
             reshape2 = next(iter(permute.users))
-            if not reshape2.target in ops.aten.reshape:
+            if not is_target_node(reshape2, ops.aten.reshape):
                 continue
             if len(reshape2.users) != 1:
                 continue
@@ -239,20 +234,14 @@ class RemoveRedundantReshapePattern3(PassBase):
         graph = graph_module.graph
         modified = False
         for reshape_1 in graph.nodes:
-            assert isinstance(reshape_1, torch.fx.Node), type(reshape_1)
             # reshape_1
-            if not reshape_1.op == "call_function":
-                continue
-            if not reshape_1.target in ops.aten.reshape:
+            if not is_target_node(reshape_1, ops.aten.reshape):
                 continue
             reshape_1_args = ReshapeArgs(*reshape_1.args, **reshape_1.kwargs)  # type: ignore[arg-type]
-            softmax, reshape_1_size = reshape_1_args.input, reshape_1_args.size
+            softmax = reshape_1_args.input
 
             # softmax
-            assert isinstance(softmax, torch.fx.Node), type(softmax)
-            if not softmax.op == "call_function":
-                continue
-            if not softmax.target in ops.aten.softmax:
+            if not is_target_node(softmax, ops.aten.softmax):
                 continue
             if softmax.target == torch.ops.aten._softmax.default:
                 softmax_args = SoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type, assignment]
@@ -354,10 +343,9 @@ class RemoveRedundantReshapePattern4(PassBase):
         modified = False
         for reshape1 in graph.nodes:
             # reshape_1
-            if not reshape1.op == "call_function":
-                continue
-            if not reshape1.target in ops.aten.reshape:
+            if not is_target_node(reshape1, ops.aten.reshape):
                 continue
+
             reshape1_args = ReshapeArgs(*reshape1.args, **reshape1.kwargs)  # type: ignore[arg-type]
             reshape1_input, size = reshape1_args.input, reshape1_args.size
             assert isinstance(reshape1_input, torch.fx.Node), type(reshape1_input)
@@ -370,12 +358,10 @@ class RemoveRedundantReshapePattern4(PassBase):
 
             # reshape_2
             reshape2 = next(iter(reshape1.users))
-            if not reshape2.op == "call_function":
+            if not is_target_node(reshape2, ops.aten.reshape):
                 continue
-            if not reshape2.target in ops.aten.reshape:
-                continue
-            reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
 
+            reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
             reshape2_input, reshape2_size = reshape2_args.input, reshape2_args.size
             assert isinstance(reshape2_input, torch.fx.Node), type(reshape2_input)
             assert isinstance(reshape2_size, list), type(reshape2_size)
@@ -420,10 +406,7 @@ class RemoveRedundantReshapePattern5(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.reshape:
+            if not is_target_node(node, ops.aten.reshape):
                 continue
 
             args = ReshapeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]

tico/passes/remove_redundant_slice.py

@@ -19,6 +19,7 @@ from tico.serialize.circle_mapping import extract_shape
 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.utils import is_target_node
 from tico.utils.validate_args_kwargs import SliceArgs
 
 
@@ -38,10 +39,7 @@ class RemoveRedundantSlice(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.slice:
+            if not is_target_node(node, ops.aten.slice):
                 continue
 
             args = SliceArgs(*node.args, **node.kwargs)

tico/passes/remove_redundant_to_copy.py

@@ -22,6 +22,7 @@ from tico.serialize.circle_mapping import extract_torch_dtype
 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.utils import is_target_node
 from tico.utils.validate_args_kwargs import ToCopyArgs, ToDtypeArgs, ToDtypeLayoutArgs
 
 
@@ -41,10 +42,7 @@ class RemoveRedundantToCopy(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.to_copy:
+            if not is_target_node(node, ops.aten.to_copy):
                 continue
 
             args: Union[ToCopyArgs, ToDtypeArgs, ToDtypeLayoutArgs]

tico/passes/segment_index_select.py

@@ -34,6 +34,7 @@ from tico.utils import logging
 from tico.utils.graph import add_placeholder, is_single_value_tensor
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_const_diff_on_pass
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import IndexSelectArgs
 
 
@@ -78,10 +79,7 @@ class SegmentIndexSelectConst(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.index_select:
+            if not is_target_node(node, ops.aten.index_select):
                 continue
 
             args = IndexSelectArgs(*node.args, **node.kwargs)