tico 0.1.0.dev250615__py3-none-any.whl → 0.1.0.dev250617__py3-none-any.whl

tico/passes/decompose_slice_scatter.py

@@ -23,6 +23,7 @@ from torch.export import ExportedProgram
 from tico.serialize.circle_mapping import extract_shape
 
 from tico.utils import logging
+from tico.utils.graph import create_node
 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, is_target_node
@@ -130,16 +131,19 @@ class DecomposeSliceScatter(PassBase):
                 slices = []
 
                 if 0 < start:
-                    slice_0 = graph.call_function(
+                    slice_0 = create_node(
+                        graph,
                         torch.ops.aten.slice_copy.Tensor,
                         args=(input, dim, 0, start, 1),
+                        origin=node,
                     )
                     slices.append(slice_0)
 
                 slices.append(src)
 
                 if start + end < extract_shape(input)[dim]:
-                    slice_1 = graph.call_function(
+                    slice_1 = create_node(
+                        graph,
                         torch.ops.aten.slice_copy.Tensor,
                         args=(
                             input,
@@ -148,13 +152,13 @@ class DecomposeSliceScatter(PassBase):
                             extract_shape(input)[dim],
                             1,
                         ),
+                        origin=node,
                     )
                     slices.append(slice_1)
 
-                concat = graph.call_function(
-                    torch.ops.aten.cat.default, args=(slices, dim)
+                concat = create_node(
+                    graph, torch.ops.aten.cat.default, args=(slices, dim)
                 )
-                # Not set meta for propagating replacing node's meta.
                 node.replace_all_uses_with(concat, propagate_meta=True)
 
             modified = True

tico/passes/fuse_leading_unsqueeze_reshape.py

@@ -20,6 +20,7 @@ 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.graph import create_node
 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
@@ -84,15 +85,19 @@ class FuseLeadingUnsqueezeReshape(PassBase):
             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(
+                r_new = create_node(
+                    graph,
                     torch.ops.aten.reshape.default,
                     args=(reshape_front_input, new_shape),
+                    origin=reshape_back,
                 )
                 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)
+                p_new = create_node(
+                    graph,
+                    torch.ops.aten.permute.default,
+                    args=(r_new, new_p_dims),
                 )
 
             reshape_back.replace_all_uses_with(p_new, propagate_meta=True)

tico/passes/legalize_predefined_layout_operators.py

@@ -23,6 +23,7 @@ from torch.export import ExportedProgram
 from tico.serialize.circle_graph import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
+from tico.utils.graph import create_node
 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
@@ -124,9 +125,11 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # TODO Introduce a method that inserts permute op.
         # input permute
         with graph.inserting_after(input):
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input, NCHW_to_NHWC),
+                origin=input,
             )
             node.update_arg(node.args.index(input), input_permute)
 
@@ -142,9 +145,11 @@ class LegalizePreDefinedLayoutOperators(PassBase):
             else:
                 assert groups == 1 or groups == input_shape[1]  # Cannot reach here
 
-            weight_permute = graph_module.graph.call_function(
+            weight_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(weight, perm),
+                origin=weight,
             )
             if args.weight.target in [
                 torch.ops.quantized_decomposed.dequantize_per_channel.default,
@@ -171,18 +176,16 @@ class LegalizePreDefinedLayoutOperators(PassBase):
             else:
                 assert groups == 1 or groups == input_shape[1]  # Cannot reach here
 
-            circle_op = graph_module.graph.call_function(
-                legalized_op,
-                args=node.args,
-                kwargs=node.kwargs,
+            circle_op = create_node(
+                graph, legalized_op, args=node.args, kwargs=node.kwargs, origin=node
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            conv_out_permute = graph_module.graph.call_function(
+            conv_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_op, NHWC_to_NCHW),
             )
-            # Not set meta for propagating replacing node's meta.
         node.replace_all_uses_with(conv_out_permute, propagate_meta=True)
 
         logger.debug(f"{node.name} is replaced with {circle_op.name}")
@@ -224,25 +227,29 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         with graph.inserting_after(input):
             # input permute
             NCHW_to_NHWC = [0, 2, 3, 1]
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input, NCHW_to_NHWC),
+                origin=input,
             )
             node.update_arg(node.args.index(input), input_permute)
         with graph.inserting_before(node):
             # circle instnorm
-            circle_instnorm = graph_module.graph.call_function(
+            circle_instnorm = create_node(
+                graph,
                 torch.ops.circle_custom.instance_norm,
                 args=node.args,
                 kwargs=node.kwargs,
+                origin=node,
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            instnorm_out_permute = graph_module.graph.call_function(
+            instnorm_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_instnorm, NHWC_to_NCHW),
             )
-            # Not set meta for propagating replacing node's meta.
         node.replace_all_uses_with(instnorm_out_permute, propagate_meta=True)
 
         logger.debug(f"{node.name} is replaced with {circle_instnorm.name}")
@@ -275,25 +282,25 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # TODO Introduce a method that inserts permute op.
         # input permute
         with graph.inserting_after(input_):
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input_, NCHW_to_NHWC),
+                origin=input_,
             )
             node.update_arg(node.args.index(input_), input_permute)
         with graph.inserting_before(node):
             legalized_op = torch.ops.circle_custom.maxpool2d
-            circle_maxpool2d = graph_module.graph.call_function(
-                legalized_op,
-                args=node.args,
-                kwargs=node.kwargs,
+            circle_maxpool2d = create_node(
+                graph, legalized_op, args=node.args, kwargs=node.kwargs, origin=node
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            maxpool_out_permute = graph_module.graph.call_function(
+            maxpool_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_maxpool2d, NHWC_to_NCHW),
             )
-            # Not set meta for propagating replacing get_item's meta.
         get_item, *_ = node.users.keys()
         get_item.replace_all_uses_with(maxpool_out_permute, propagate_meta=True)
 
@@ -327,21 +334,22 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # TODO Introduce a method that inserts permute op.
         # input permute
         with graph.inserting_after(input_):
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input_, NCHW_to_NHWC),
+                origin=input_,
             )
             node.update_arg(node.args.index(input_), input_permute)
         with graph.inserting_before(node):
             legalized_op = torch.ops.circle_custom.avgpool2d
-            circle_avgpool2d = graph_module.graph.call_function(
-                legalized_op,
-                args=node.args,
-                kwargs=node.kwargs,
+            circle_avgpool2d = create_node(
+                graph, legalized_op, args=node.args, kwargs=node.kwargs, origin=node
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            avgpool_out_permute = graph_module.graph.call_function(
+            avgpool_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_avgpool2d, NHWC_to_NCHW),
             )

tico/passes/lower_pow2_to_mul.py

@@ -20,6 +20,7 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.utils import logging
+from tico.utils.graph import create_node
 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
@@ -55,7 +56,8 @@ class LowerPow2ToMul(PassBase):
 
             lhs = rhs = in_
             with graph.inserting_after(node):
-                new_mul = graph.call_function(
+                new_mul = create_node(
+                    graph,
                     torch.ops.aten.mul.Tensor,
                     args=(lhs, rhs),
                     kwargs={},

tico/passes/lower_to_resize_nearest_neighbor.py

@@ -22,6 +22,7 @@ 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 create_node
 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
@@ -108,18 +109,23 @@ class LowerToResizeNearestNeighbor(PassBase):
         assert expected_shape == list(extract_shape(node))
 
         with graph.inserting_before(node):
-            nchw_to_nhwc = graph.call_function(
-                torch.ops.aten.permute.default, args=(input_tensor, [0, 2, 3, 1])
+            nchw_to_nhwc = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(input_tensor, [0, 2, 3, 1]),
+                origin=input_tensor,
             )
-            resize_nearest_neighbor = graph.call_function(
+            resize_nearest_neighbor = create_node(
+                graph,
                 torch.ops.circle_custom.resize_nearest_neighbor,
                 args=(nchw_to_nhwc, [len(expected_H_index), len(expected_W_index)]),
+                origin=node,
             )
-            nhwc_to_nchw = graph.call_function(
+            nhwc_to_nchw = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(resize_nearest_neighbor, [0, 3, 1, 2]),
             )
-            # Not set meta for propagating replacing node's meta.
             node.replace_all_uses_with(nhwc_to_nchw, propagate_meta=True)
 
         return resize_nearest_neighbor
@@ -171,18 +177,23 @@ class LowerToResizeNearestNeighbor(PassBase):
             )
 
         with graph.inserting_before(node):
-            nchw_to_nhwc = graph.call_function(
-                torch.ops.aten.permute.default, args=(input_tensor, [0, 2, 3, 1])
+            nchw_to_nhwc = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(input_tensor, [0, 2, 3, 1]),
+                origin=input_tensor,
             )
-            resize_nearest_neighbor = graph.call_function(
+            resize_nearest_neighbor = create_node(
+                graph,
                 torch.ops.circle_custom.resize_nearest_neighbor,
                 args=(nchw_to_nhwc, [expected_H, expected_W]),
+                origin=node,
             )
-            nhwc_to_nchw = graph.call_function(
+            nhwc_to_nchw = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(resize_nearest_neighbor, [0, 3, 1, 2]),
             )
-            # Not set meta for propagating replacing node's meta.
             node.replace_all_uses_with(nhwc_to_nchw, propagate_meta=True)
             return resize_nearest_neighbor
 

tico/passes/lower_to_slice.py

@@ -30,7 +30,7 @@ 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.graph import create_node, 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
@@ -103,17 +103,22 @@ class LowerSelectCopyToSlice(PassBase):
 
             with graph.inserting_after(node):
                 # slice
-                slice_node = graph.call_function(
-                    torch.ops.aten.slice.Tensor, args=slice_copy_args
+                slice_node = create_node(
+                    graph,
+                    torch.ops.aten.slice.Tensor,
+                    args=slice_copy_args,
+                    origin=node,
                 )
                 node_shape = extract_shape(node)
             with graph.inserting_after(slice_node):
                 # reshape
                 reshape_args = (slice_node, list(node_shape))
-                reshape_node = graph.call_function(
-                    torch.ops.aten.reshape.default, args=reshape_args
+                reshape_node = create_node(
+                    graph,
+                    torch.ops.aten.reshape.default,
+                    args=reshape_args,
                 )
-                node.replace_all_uses_with(reshape_node, propagate_meta=False)
+                node.replace_all_uses_with(reshape_node, propagate_meta=True)
 
             modified = True
             logger.debug(
@@ -196,17 +201,22 @@ class LowerIndexSelectToSlice(PassBase):
 
             with graph.inserting_after(node):
                 # slice
-                slice_node = graph.call_function(
-                    torch.ops.aten.slice.Tensor, args=slice_copy_args
+                slice_node = create_node(
+                    graph,
+                    torch.ops.aten.slice.Tensor,
+                    args=slice_copy_args,
+                    origin=node,
                 )
                 node_shape = extract_shape(node)
             with graph.inserting_after(slice_node):
                 # reshape
                 reshape_args = (slice_node, list(node_shape))
-                reshape_node = graph.call_function(
-                    torch.ops.aten.reshape.default, args=reshape_args
+                reshape_node = create_node(
+                    graph,
+                    torch.ops.aten.reshape.default,
+                    args=reshape_args,
                 )
-                node.replace_all_uses_with(reshape_node, propagate_meta=False)
+                node.replace_all_uses_with(reshape_node, propagate_meta=True)
 
             modified = True
             logger.debug(

tico/passes/remove_redundant_permute.py

@@ -14,7 +14,6 @@
 
 from typing import TYPE_CHECKING
 
-
 if TYPE_CHECKING:
     import torch.fx
 import torch
@@ -23,6 +22,7 @@ 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.graph import create_node
 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
@@ -106,8 +106,10 @@ class RemoveRedundantPermutePattern1(PassBase):
             else:
                 with graph.inserting_after(permute2):
                     new_args = (permute1_input, fused_dims)
-                    fused_permute = graph.call_function(
-                        torch.ops.aten.permute.default, args=new_args
+                    fused_permute = create_node(
+                        graph,
+                        torch.ops.aten.permute.default,
+                        args=new_args,
                     )
                     permute2.replace_all_uses_with(fused_permute, propagate_meta=True)
                     logger.debug(f"{permute1.name} and {permute2.name} are fused.")

tico/passes/remove_redundant_reshape.py

@@ -22,6 +22,7 @@ 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.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import broadcastable, is_target_node, set_new_meta_val
@@ -369,8 +370,10 @@ class RemoveRedundantReshapePattern4(PassBase):
                 assert isinstance(s, int), type(s)
 
             with graph.inserting_before(reshape1):
-                fused_reshape = graph.call_function(
-                    reshape1.target, (reshape1_input, reshape2_size)
+                fused_reshape = create_node(
+                    graph,
+                    reshape1.target,
+                    (reshape1_input, reshape2_size),
                 )
 
             reshape2.replace_all_uses_with(fused_reshape, propagate_meta=True)

tico/passes/remove_redundant_to_copy.py

@@ -70,6 +70,10 @@ class RemoveRedundantToCopy(PassBase):
                 if input_dtype != target_dtype:
                     continue
 
+            if hasattr(args, "memory_format") and args.memory_format is not None:
+                if args.memory_format != torch.contiguous_format:
+                    continue
+
             node.replace_all_uses_with(input_, propagate_meta=False)
 
             modified = True

tico/passes/restore_linear.py

@@ -16,6 +16,7 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 
@@ -74,11 +75,12 @@ class RestoreLinear(PassBase):
 
                 addmm_args = (input, weight, bias)
                 with graph.inserting_after(node):
-                    linear_node = graph.call_function(
-                        torch.ops.aten.linear.default, args=addmm_args
+                    linear_node = create_node(
+                        graph,
+                        torch.ops.aten.linear.default,
+                        args=addmm_args,
                     )
                     node.replace_all_uses_with(linear_node, propagate_meta=True)
-                    graph.erase_node(node)
 
             elif node.target == torch.ops.aten.mm.default:
                 assert len(node.args) == 2
@@ -97,8 +99,8 @@ class RestoreLinear(PassBase):
 
                 mm_args = (input, weight)
                 with graph.inserting_after(node):
-                    linear_node = graph.call_function(
-                        torch.ops.aten.linear.default, args=mm_args
+                    linear_node = create_node(
+                        graph, torch.ops.aten.linear.default, args=mm_args
                     )
                     node.replace_all_uses_with(linear_node, propagate_meta=True)
 

tico/passes/segment_index_select.py

@@ -31,7 +31,7 @@ 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 add_placeholder, is_single_value_tensor
+from tico.utils.graph import add_placeholder, create_node, 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
@@ -116,18 +116,22 @@ class SegmentIndexSelectConst(PassBase):
                     exported_program, torch.tensor([i]), prefix="segm_index"
                 )
                 with graph.inserting_before(node):
-                    index_select_node = graph.call_function(
+                    index_select_node = create_node(
+                        graph,
                         torch.ops.aten.index_select.default,
                         args=(input, dim, index_node),
+                        origin=node,
                     )
                     index_select_node_list.append(index_select_node)
 
             with graph.inserting_before(node):
-                concat_node = graph.call_function(
-                    torch.ops.aten.cat.default, args=(index_select_node_list, dim)
+                concat_node = create_node(
+                    graph,
+                    torch.ops.aten.cat.default,
+                    args=(index_select_node_list, dim),
                 )
 
-                node.replace_all_uses_with(concat_node, propagate_meta=False)
+                node.replace_all_uses_with(concat_node, propagate_meta=True)
 
             modified = True
             logger.debug(

tico/utils/graph.py

@@ -16,7 +16,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, TYPE_CHECKING
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx
@@ -24,7 +24,7 @@ import torch
 from torch.export import ExportedProgram
 from torch.export.exported_program import InputKind, InputSpec, TensorArgument
 
-from tico.utils.utils import get_fake_mode
+from tico.utils.utils import get_fake_mode, set_new_meta_val
 
 
 def is_torch_param(node: torch.fx.Node, ep: ExportedProgram):
@@ -234,3 +234,49 @@ def get_module_name_chain(node: Optional[torch.fx.Node]) -> str:
         return next(reversed(stack.values()))[1]
     else:
         return "unknown"
+
+
+def create_node(
+    graph: torch.fx.Graph,
+    target: torch._ops.OpOverload,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    *,
+    origin: Optional[torch.fx.Node] = None,
+) -> torch.fx.Node:
+    """
+    Insert a new node into graph and propagate metadata from *origin*.
+
+    Parameters
+    ----------
+    graph : torch.fx.Graph
+        The graph that will own the newly-created node.
+
+    target : torch._ops.OpOverload
+        The op to call (e.g. `torch.add` or "call_function" target).
+
+    args : Tuple[Any, ...], optional
+        Positional arguments for the new node.
+
+    kwargs : Dict[str, Any], optional
+        Keyword arguments for the new node.
+
+    origin : torch.fx.Node, optional
+        If given, every key in `origin.meta` **except** "val" is copied
+        onto the new node.  "val" is recomputed from *args* /*kwargs* using
+        the internal meta-inference helper.
+
+    Returns
+    -------
+    torch.fx.Node
+        The freshly inserted node with fully-populated `.meta`.
+    """
+    new_node = graph.call_function(target, args=args, kwargs=kwargs)
+    if origin:
+        assert isinstance(origin, torch.fx.Node), type(origin)
+        # Propagate "nn_module_stack" to retain the originating module context
+        #  for meaningful node names.
+        if "nn_module_stack" in origin.meta:
+            new_node.meta["nn_module_stack"] = origin.meta["nn_module_stack"]
+
+    return new_node

{tico-0.1.0.dev250615.dist-info → tico-0.1.0.dev250617.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250615
+Version: 0.1.0.dev250617
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN
@@ -105,10 +105,15 @@ You can convert a torch module to a circle model with these steps.
 torch_module = AddModule()
 example_inputs = (torch.ones(4), torch.ones(4))
 
-circle_model = tico.convert(torch_module, example_inputs)
+circle_model = tico.convert(torch_module.eval(), example_inputs)
 circle_model.save('add.circle')
 ```
 
+**NOTE**
+Please make sure to call `eval()` on the PyTorch module before passing it to our API.
+This ensures the model runs in inference mode, disabling layers like dropout and
+batch normalization updates.
+
 **Compile with configuration**
 
 ```python