tico 0.1.0.dev250722__py3-none-any.whl → 0.1.0.dev250724__py3-none-any.whl

tico/serialize/circle_serializer.py

@@ -18,12 +18,7 @@ from typing import Dict
 import flatbuffers
 import torch
 from circle_schema import circle
-from torch.export.exported_program import (
-    ConstantArgument,
-    ExportedProgram,
-    InputKind,
-    TensorArgument,
-)
+from torch.export.exported_program import ConstantArgument, ExportedProgram, InputKind
 
 from tico.serialize.circle_mapping import to_circle_dtype
 from tico.serialize.operators import *
@@ -39,147 +34,38 @@ multiple_output_ops = [
     torch.ops.aten.max.dim,
 ]
 
-# Build circle model from ExportedProgram
-# Return raw bytes of circle model
-def build_circle(edge_program: ExportedProgram) -> bytes:
-    logger = logging.getLogger(__name__)
 
-    builder = flatbuffers.Builder()
+def _initialize_model() -> tuple[CircleModel, CircleSubgraph]:
+    """Initialize a new Circle model and subgraph.
 
-    # Init Model
+    Returns:
+        Tuple containing the model and subgraph
+    """
     model = CircleModel()
-
-    # Add empty buffer at the front (convention)
-    model.add_buffer(circle.Buffer.BufferT())
-
-    # Create an empty subgraph (assume a single subgraph)
+    model.add_buffer(circle.Buffer.BufferT())  # Add empty buffer at the front
     graph = CircleSubgraph(model)
+    return model, graph
 
-    # Export tensors
-    logger.debug("---------------Export tensors--------------")
-    buf_name_to_data = {name: buf for name, buf in edge_program.named_buffers()}
-    for node in edge_program.graph.nodes:
-        if node.op == "call_function":
-            if node.target in multiple_output_ops:
-                continue
-            node_val = node.meta["val"]
-            if node_val.layout != torch.strided:
-                raise RuntimeError(
-                    f"Only support dense tensors (node layout: {node_val.layout})"
-                )
-            graph.add_tensor_from_node(node)
-            logger.debug(f"call_function: {node.name} tensor exported.")
-
-        # placeholder: function input (including parameters, buffers, constant tensors)
-        elif node.op == "placeholder":
-            # placeholder invariants
-            assert node.args is None or len(node.args) == 0  # Not support default param
-
-            # parameters
-            if node.name in edge_program.graph_signature.inputs_to_parameters:
-                param_name = edge_program.graph_signature.inputs_to_parameters[
-                    node.name
-                ]
-                param_data = edge_program.state_dict[param_name]
-
-                assert isinstance(
-                    param_data, torch.Tensor
-                ), "Expect parameters to be a tensor"
-                param_value = param_data.cpu().detach().numpy()
-
-                graph.add_tensor_from_node(node, param_value)
-                logger.debug(f"placeholder(param): {node.name} tensor exported.")
-            elif node.name in edge_program.graph_signature.inputs_to_buffers:
-                buffer_name = edge_program.graph_signature.inputs_to_buffers[node.name]
-                assert buffer_name in buf_name_to_data
-                buffer_data = buf_name_to_data[buffer_name]
-                assert isinstance(
-                    buffer_data, torch.Tensor
-                ), "Expect buffers to be a tensor"
-                buffer_value = buffer_data.cpu().detach().numpy()
-
-                graph.add_tensor_from_node(node, buffer_value)
-                logger.debug(f"placeholder(buffer): {node.name} tensor exported.")
-            elif (
-                node.name
-                in edge_program.graph_signature.inputs_to_lifted_tensor_constants
-            ):
-                ctensor_name = (
-                    edge_program.graph_signature.inputs_to_lifted_tensor_constants[
-                        node.name
-                    ]
-                )
-                ctensor_data = edge_program.constants[ctensor_name]
-
-                assert isinstance(
-                    ctensor_data, torch.Tensor
-                ), "Expect constant tensor to be a tensor"
-                ctensor_value = ctensor_data.cpu().detach().numpy()
-
-                graph.add_tensor_from_node(node, ctensor_value)
-                logger.debug(
-                    f"placeholder(constant tensor): {node.name} tensor exported."
-                )
-            else:
-                user_inputs = [
-                    specs
-                    for specs in edge_program.graph_signature.input_specs
-                    if specs.kind == InputKind.USER_INPUT
-                ]
-                constant_inputs = [
-                    specs
-                    for specs in user_inputs
-                    if isinstance(specs.arg, ConstantArgument)
-                ]
-                name_to_value = {
-                    specs.arg.name: specs.arg.value for specs in constant_inputs
-                }
-                # NoneType ConstantArgument is ignored.
-                if node.name in name_to_value and name_to_value[node.name] == None:
-                    continue
-                graph.add_tensor_from_node(node)
-                logger.debug(f"placeholder: {node.name} tensor exported.")
-
-        # get_attr: retrieve parameter
-        elif node.op == "get_attr":
-            # node.name: Place where fetched attribute is saved
-            # node.target: Attribute in the module
-            attr_tensor = getattr(node.graph.owning_module, node.target)
-            assert isinstance(attr_tensor, torch.Tensor)
 
-            graph.add_tensor_from_scratch(
-                prefix=node.name,
-                shape=list(attr_tensor.shape),
-                dtype=to_circle_dtype(attr_tensor.dtype),
-                source_node=node,
-            )
+def build_circle(ep: ExportedProgram) -> bytes:
+    """Convert ExportedProgram to Circle format.
 
-            logger.debug(f"get_attr: {node.name} tensor exported.")
+    Args:
+        ep: The exported PyTorch program to convert
 
-        # output: function output
-        elif node.op == "output":
-            # output node itself does not need a buffer
-            # argument of output node is assumed to be exported beforehand
-            for output in node.args[0]:
-                if isinstance(output, torch.fx.Node):
-                    assert graph.has_tensor(output.name)
-            continue
-
-        # call_method: call method
-        elif node.op == "call_method":
-            raise AssertionError("Not yet implemented")
-
-        # call_module: call 'forward' of module
-        elif node.op == "call_module":
-            raise AssertionError("Not yet implemented")
+    Returns:
+        bytes: Raw bytes of the Circle model
+    """
+    logger = logging.getLogger(__name__)
+    builder = flatbuffers.Builder()
+    model, graph = _initialize_model()
 
-        else:
-            # Add more if fx.Node is extended
-            raise AssertionError(f"Unknown fx.Node op {node.op}")
+    # Export tensors
+    _export_tensors(graph, ep)
 
     # Register inputs
     logger.debug("---------------Register inputs--------------")
-    for in_spec in edge_program.graph_signature.input_specs:
+    for in_spec in ep.graph_signature.input_specs:
         if in_spec.kind != InputKind.USER_INPUT:
             continue
         # NoneType ConstantArgument is ignored.
@@ -191,9 +77,9 @@ def build_circle(edge_program: ExportedProgram) -> bytes:
 
     # Register outputs
     logger.debug("---------------Register outputs--------------")
-    for user_output in edge_program.graph_signature.user_outputs:
+    for user_output in ep.graph_signature.user_outputs:
         if user_output == None:
-            logger.debug(f"Ignore 'None' output")
+            logger.debug("Ignore 'None' output")
             continue
 
         graph.add_output(user_output)
@@ -203,7 +89,7 @@ def build_circle(edge_program: ExportedProgram) -> bytes:
     logger.debug("---------------Export operators--------------")
     op_codes: Dict[OpCode, int] = {}
     visitors = get_node_visitors(op_codes, graph)
-    for node in edge_program.graph.nodes:
+    for node in ep.graph.nodes:
         if node.op != "call_function":
             continue
 
@@ -227,10 +113,8 @@ def build_circle(edge_program: ExportedProgram) -> bytes:
         code for code, _ in sorted(op_codes.items(), key=lambda x: x[1])
     ]
 
-    # Description
+    # Final model settings
     model.description = "circle"
-
-    # Set version
     model.version = 0
 
     # Finish model
@@ -238,3 +122,212 @@ def build_circle(edge_program: ExportedProgram) -> bytes:
     buf = builder.Output()
 
     return bytes(buf)
+
+
+def _export_tensors(graph: CircleSubgraph, ep: ExportedProgram) -> None:
+    """Export all tensors from the exported program to the circle graph.
+
+    Args:
+        graph: The CircleSubgraph to add tensors to
+        ep: The exported PyTorch program
+    """
+    logger = logging.getLogger(__name__)
+    logger.debug("---------------Export tensors--------------")
+    buf_name_to_data = {name: buf for name, buf in ep.named_buffers()}
+
+    for node in ep.graph.nodes:
+        if node.op == "call_function":
+            if node.target in multiple_output_ops:
+                continue
+            node_val = node.meta["val"]
+            if node_val.layout != torch.strided:
+                raise RuntimeError(
+                    f"Only support dense tensors (node layout: {node_val.layout})"
+                )
+            graph.add_tensor_from_node(node)
+            logger.debug(f"call_function: {node.name} tensor exported.")
+
+        elif node.op == "placeholder":
+            _handle_placeholder_node(graph, node, ep, buf_name_to_data)
+
+        elif node.op == "get_attr":
+            _handle_get_attr_node(graph, node)
+
+        elif node.op == "output":
+            for output in node.args[0]:
+                if isinstance(output, torch.fx.Node):
+                    assert graph.has_tensor(output.name)
+            continue
+
+        elif node.op == "call_method":
+            raise AssertionError("Not yet implemented")
+
+        elif node.op == "call_module":
+            raise AssertionError("Not yet implemented")
+
+        else:
+            raise AssertionError(f"Unknown fx.Node op {node.op}")
+
+
+def _handle_placeholder_node(
+    graph: CircleSubgraph,
+    node: torch.fx.Node,
+    ep: ExportedProgram,
+    buf_name_to_data: dict,
+) -> None:
+    """Handle a placeholder node during tensor export."""
+    # placeholder invariants
+    assert node.args is None or len(node.args) == 0  # Not support default param
+
+    if node.name in ep.graph_signature.inputs_to_parameters:
+        _handle_parameter_node(graph, node, ep)
+    elif node.name in ep.graph_signature.inputs_to_buffers:
+        _handle_buffer_node(graph, node, ep, buf_name_to_data)
+    elif node.name in ep.graph_signature.inputs_to_lifted_tensor_constants:
+        _handle_constant_tensor_node(graph, node, ep)
+    else:
+        _handle_user_input_node(graph, node, ep)
+
+
+def _handle_parameter_node(
+    graph: CircleSubgraph,
+    node: torch.fx.Node,
+    ep: ExportedProgram,
+) -> None:
+    """Handle a parameter placeholder node by exporting its tensor data.
+
+    Args:
+        graph: CircleSubgraph to add tensor to
+        node: The parameter node to process
+        ep: ExportedProgram containing parameter data
+    """
+    param_name = ep.graph_signature.inputs_to_parameters[node.name]
+    param_data = ep.state_dict[param_name]
+
+    if not isinstance(param_data, torch.Tensor):
+        raise ValueError(f"Parameter {param_name} is not a tensor")
+
+    tensor_value = param_data.cpu().detach().numpy()
+    graph.add_tensor_from_node(node, tensor_value)
+
+    logger = logging.getLogger(__name__)
+    logger.debug(f"Exported parameter tensor: {node.name}")
+
+
+def _handle_buffer_node(
+    graph: CircleSubgraph,
+    node: torch.fx.Node,
+    ep: ExportedProgram,
+    buf_name_to_data: dict,
+) -> None:
+    """Handle a buffer placeholder node by exporting its tensor data.
+
+    Args:
+        graph: CircleSubgraph to add tensor to
+        node: The buffer node to process
+        ep: ExportedProgram containing buffer info
+        buf_name_to_data: Mapping of buffer names to data
+    """
+    buffer_name = ep.graph_signature.inputs_to_buffers[node.name]
+
+    if buffer_name not in buf_name_to_data:
+        raise ValueError(f"Buffer {buffer_name} not found in buffer data")
+
+    buffer_data = buf_name_to_data[buffer_name]
+
+    if not isinstance(buffer_data, torch.Tensor):
+        raise ValueError(f"Buffer {buffer_name} is not a tensor")
+
+    tensor_value = buffer_data.cpu().detach().numpy()
+    graph.add_tensor_from_node(node, tensor_value)
+
+    logger = logging.getLogger(__name__)
+    logger.debug(f"Exported buffer tensor: {node.name}")
+
+
+def _handle_constant_tensor_node(
+    graph: CircleSubgraph,
+    node: torch.fx.Node,
+    ep: ExportedProgram,
+) -> None:
+    """Handle a constant tensor placeholder node by exporting its tensor data.
+
+    Args:
+        graph: CircleSubgraph to add tensor to
+        node: The constant tensor node to process
+        ep: ExportedProgram containing constant data
+    """
+    ctensor_name = ep.graph_signature.inputs_to_lifted_tensor_constants[node.name]
+
+    if ctensor_name not in ep.constants:
+        raise ValueError(f"Constant tensor {ctensor_name} not found")
+
+    ctensor_data = ep.constants[ctensor_name]
+
+    if not isinstance(ctensor_data, torch.Tensor):
+        raise ValueError(f"Constant tensor {ctensor_name} is not a tensor")
+
+    tensor_value = ctensor_data.cpu().detach().numpy()
+    graph.add_tensor_from_node(node, tensor_value)
+
+    logger = logging.getLogger(__name__)
+    logger.debug(f"Exported constant tensor: {node.name}")
+
+
+def _handle_user_input_node(
+    graph: CircleSubgraph,
+    node: torch.fx.Node,
+    ep: ExportedProgram,
+) -> None:
+    """Handle a user input placeholder node by exporting its tensor data.
+
+    Args:
+        graph: CircleSubgraph to add tensor to
+        node: The user input node to process
+        ep: ExportedProgram containing input specs
+    """
+    user_inputs = [
+        specs
+        for specs in ep.graph_signature.input_specs
+        if specs.kind == InputKind.USER_INPUT
+    ]
+    constant_inputs = [
+        specs for specs in user_inputs if isinstance(specs.arg, ConstantArgument)
+    ]
+    name_to_value = {specs.arg.name: specs.arg.value for specs in constant_inputs}
+
+    # Skip NoneType ConstantArgument
+    if node.name in name_to_value and name_to_value[node.name] is None:
+        return
+
+    graph.add_tensor_from_node(node)
+
+    logger = logging.getLogger(__name__)
+    logger.debug(f"Exported user input tensor: {node.name}")
+
+
+def _handle_get_attr_node(
+    graph: CircleSubgraph,
+    node: torch.fx.Node,
+) -> None:
+    """Handle a get_attr node by exporting its tensor data.
+
+    Args:
+        graph: CircleSubgraph to add tensor to
+        node: The get_attr node to process
+    """
+    assert isinstance(node.target, str)
+    attr_tensor = getattr(node.graph.owning_module, node.target)
+
+    if not isinstance(attr_tensor, torch.Tensor):
+        raise ValueError(f"Attribute {node.target} is not a tensor")
+
+    graph.add_tensor_from_scratch(
+        prefix=node.name,
+        shape=list(attr_tensor.shape),
+        dtype=to_circle_dtype(attr_tensor.dtype),
+        source_node=node,
+    )
+
+    logger = logging.getLogger(__name__)
+    logger.debug(f"Exported attribute tensor: {node.name}")

tico/serialize/operators/op_any.py

@@ -22,7 +22,6 @@ from circle_schema import circle
 from tico.serialize.circle_graph import CircleSubgraph
 from tico.serialize.circle_mapping import (
     circle_legalize_dtype_to,
-    extract_circle_dtype,
     extract_shape,
     extract_torch_dtype,
 )
@@ -100,8 +99,6 @@ class AnyVisitor(NodeVisitor):
         keepdim = args.keepdim
 
         input_shape = list(extract_shape(input))
-        output_shape = list(extract_shape(node))
-
         dim_i32 = None
         if dim is None:
             dims = tuple(i for i in range(0, len(input_shape)))

tico/serialize/operators/op_clamp.py

@@ -21,12 +21,9 @@ import torch
 from circle_schema import circle
 
 from tico.passes import ops
+from tico.serialize.circle_graph import CircleSubgraph
 
-from tico.serialize.circle_graph import (
-    CircleSubgraph,
-    extract_circle_dtype,
-    extract_shape,
-)
+from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index

tico/serialize/operators/op_full_like.py

@@ -21,10 +21,8 @@ import torch
 from circle_schema import circle
 
 from tico.serialize.circle_graph import CircleSubgraph
-from tico.serialize.circle_mapping import to_circle_dtype
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
-from tico.serialize.operators.utils import create_builtin_operator, get_op_index
 from tico.utils.validate_args_kwargs import FullLikeArgs
 
 

tico/serialize/operators/op_instance_norm.py

@@ -73,12 +73,6 @@ class InstanceNormVisitor(NodeVisitor):
         eps = args.eps
 
         # Ignore training-related args
-        running_mean = args.running_mean
-        running_var = args.running_var
-        use_input_stats = args.use_input_stats
-        momentum = args.momentum
-        cudnn_enabled = args.cudnn_enabled
-
         input_shape = list(extract_shape(input))
         assert len(input_shape) == 4, len(input_shape)
 

tico/serialize/operators/op_mul.py

@@ -66,10 +66,7 @@ class MulTensorVisitor(BaseMulVisitor):
         self,
         node: torch.fx.Node,
     ) -> circle.Operator.OperatorT:
-        args = MulTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-        input = args.input
-        other = args.other
-
+        _ = MulTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         operator = super().define_node(
             node,
         )
@@ -88,10 +85,7 @@ class MulScalarVisitor(BaseMulVisitor):
         self,
         node: torch.fx.Node,
     ) -> circle.Operator.OperatorT:
-        args = MulScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-        input = args.input
-        other = args.other
-
+        _ = MulScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         operator = super().define_node(
             node,
         )

tico/serialize/operators/op_transpose_conv.py

@@ -76,9 +76,7 @@ class TransposeConvVisitor(NodeVisitor):
         bias = args.bias
         stride = args.stride
         padding = args.padding
-        output_padding = args.output_padding
         groups = args.groups
-        dilation = args.dilation
 
         assert groups == 1, "Only support group 1"
 

tico/serialize/quant_param.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
 """
 This is a key for torch.fx.Node's meta dict to save QuantParam
 
@@ -19,11 +24,6 @@ QuantParam can be retrieved as node.meta[QPARAM_KEY]
 """
 QPARAM_KEY = "_quantization_parameters_"
 
-from dataclasses import dataclass
-from typing import List, Optional
-
-import torch
-
 
 @dataclass
 class QuantParam:

tico/utils/convert.py

@@ -157,7 +157,7 @@ def check_unsupported_target(exported_program: ExportedProgram):
     for n in exported_program.graph.nodes:
         if n.op != "call_function":
             continue
-        if not n.target in supported_target:
+        if n.target not in supported_target:
             unsupported.append(n)
 
     if unsupported:

tico/utils/graph.py

@@ -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, set_new_meta_val
+from tico.utils.utils import get_fake_mode
 
 
 def is_torch_param(node: torch.fx.Node, ep: ExportedProgram):

tico/utils/padding.py

@@ -15,8 +15,6 @@
 from enum import IntEnum
 from typing import NamedTuple, Optional, Sequence, Tuple, Union
 
-import torch
-
 from tico.utils.errors import InvalidArgumentError
 
 

tico/utils/serialize.py

@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
-
-import torch
 
 from tico.serialize.circle_graph import CircleSubgraph
 from tico.utils.graph import get_module_name_chain

tico/utils/utils.py

@@ -21,7 +21,6 @@ from typing import List
 
 import torch
 from circle_schema import circle
-from packaging.version import Version
 from torch._guards import detect_fake_mode
 from torch.export import ExportedProgram
 from torch.utils import _pytree as pytree
@@ -131,7 +130,7 @@ def enforce_type(callable):
 
                     return True
 
-                if typing.get_origin(type_hint) == dict:
+                if typing.get_origin(type_hint) is dict:
                     if not isinstance(value, typing.get_origin(type_hint)):
                         return False
 

tico/utils/validate_args_kwargs.py

@@ -16,10 +16,8 @@ from dataclasses import dataclass, field
 from typing import List, Optional, TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
-    import torch._ops
     import torch.fx
 import torch
-import torch.fx.node
 
 from tico.utils.utils import enforce_type
 
@@ -137,7 +135,7 @@ class AvgPool2dArgs:
             assert len(self.padding) == 2, len(self.padding)
         if self.divisor_override is not None:
             assert isinstance(self.divisor_override, int), type(self.divisor_override)
-            assert self.divisor_override != 0, f"Divisor must be not zero."
+            assert self.divisor_override != 0, "Divisor must be not zero."
 
 
 @enforce_type

{tico-0.1.0.dev250722.dist-info → tico-0.1.0.dev250724.dist-info}/METADATA

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