onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/graph_builder.py

@@ -0,0 +1,634 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections import deque
+from typing import Any, Callable, Dict, List, Optional, Tuple, Sequence, Union
+
+import numpy as np
+
+import torch
+import torch.fx
+import onnx
+from onnx import numpy_helper
+
+from .exceptions import UnsupportedDTypeError, UnsupportedOpError, ValueNotFoundError
+from .op_registry import get_handler
+from .utils.dtype import DTYPE_MAP
+from .utils.external_data import resolve_external_data
+from .utils.names import sanitize_name
+
+# Import ops module to register all operators
+from . import ops  # noqa: F401
+
+
+def _collect_all_inputs(node: onnx.NodeProto) -> set:
+    """Recursively collect all inputs from a node including subgraph inputs.
+
+    For control flow nodes like If and Loop, this also collects inputs that
+    are referenced by the subgraphs (then_branch, else_branch, body).
+
+    Parameters
+    ----------
+    node : onnx.NodeProto
+        The ONNX node to collect inputs from.
+
+    Returns
+    -------
+    set
+        Set of all input names referenced by this node and its subgraphs.
+    """
+    inputs = set(node.input)
+
+    # Collect inputs from subgraphs (for If, Loop, etc.)
+    for attr in node.attribute:
+        if attr.type == onnx.AttributeProto.GRAPH:
+            subgraph = attr.g
+            # Collect subgraph's own initializers and inputs as local values
+            local_values = set()
+            for init in subgraph.initializer:
+                local_values.add(init.name)
+            for inp in subgraph.input:
+                local_values.add(inp.name)
+
+            # Recursively collect inputs from subgraph nodes
+            for sub_node in subgraph.node:
+                sub_inputs = _collect_all_inputs(sub_node)
+                # Add outputs of this subgraph node to local values
+                for out in sub_node.output:
+                    if out:
+                        local_values.add(out)
+                # Inputs not satisfied locally are outer references
+                for sub_inp in sub_inputs:
+                    if sub_inp and sub_inp not in local_values:
+                        inputs.add(sub_inp)
+
+    return inputs
+
+
+def _topological_sort(
+    nodes: List[onnx.NodeProto],
+    graph_inputs: set,
+    initializers: set,
+) -> List[onnx.NodeProto]:
+    """Topologically sort ONNX graph nodes using Kahn's algorithm.
+
+    Some ONNX models have nodes in non-topological order (e.g., Cast nodes
+    at the end of the graph but their outputs used earlier). This function
+    reorders nodes so dependencies are processed before their consumers.
+
+    This function also considers inputs referenced by subgraphs (for nodes
+    like If and Loop) to ensure proper ordering.
+
+    Parameters
+    ----------
+    nodes : List[onnx.NodeProto]
+        The list of ONNX nodes to sort.
+    graph_inputs : set
+        Set of graph input names.
+    initializers : set
+        Set of initializer names.
+
+    Returns
+    -------
+    List[onnx.NodeProto]
+        Topologically sorted list of nodes.
+    """
+    if not nodes:
+        return []
+
+    # Pre-compute all inputs for each node (including subgraph inputs)
+    node_all_inputs: Dict[int, set] = {}
+    for node in nodes:
+        node_all_inputs[id(node)] = _collect_all_inputs(node)
+
+    # Build output->node mapping (which node produces each output)
+    output_to_node: Dict[str, onnx.NodeProto] = {}
+    for node in nodes:
+        for output in node.output:
+            if output:  # Skip empty outputs
+                output_to_node[output] = node
+
+    # Available values: graph inputs + initializers
+    available = graph_inputs | initializers
+
+    # Compute in-degree for each node (number of unsatisfied dependencies)
+    in_degree: Dict[int, int] = {}
+    node_id: Dict[int, onnx.NodeProto] = {}
+    for i, node in enumerate(nodes):
+        node_id[id(node)] = node
+        # Count inputs that are neither available nor empty
+        deps = 0
+        all_inputs = node_all_inputs[id(node)]
+        for inp in all_inputs:
+            if inp and inp not in available:
+                deps += 1
+        in_degree[id(node)] = deps
+
+    # Initialize queue with nodes that have no dependencies
+    queue = deque()
+    for node in nodes:
+        if in_degree[id(node)] == 0:
+            queue.append(node)
+
+    sorted_nodes: List[onnx.NodeProto] = []
+    while queue:
+        node = queue.popleft()
+        sorted_nodes.append(node)
+
+        # Mark this node's outputs as available
+        for output in node.output:
+            if output:
+                available.add(output)
+
+        # Reduce in-degree for nodes that depend on this node's outputs
+        for candidate in nodes:
+            if in_degree[id(candidate)] > 0:
+                # Check if any of candidate's inputs are now satisfied
+                all_inputs = node_all_inputs[id(candidate)]
+                for inp in all_inputs:
+                    if inp in node.output and inp:
+                        in_degree[id(candidate)] -= 1
+                        if in_degree[id(candidate)] == 0:
+                            queue.append(candidate)
+
+    # If we couldn't sort all nodes, there's a cycle or missing dependency
+    # Fall back to original order
+    if len(sorted_nodes) != len(nodes):
+        return list(nodes)
+
+    return sorted_nodes
+
+
+class GraphBuilder:
+    """Builds a PyTorch FX GraphModule from an ONNX model.
+
+    This class handles the conversion of ONNX graph structure to PyTorch FX,
+    including initializer loading, placeholder creation, node conversion,
+    and output creation.
+
+    Parameters
+    ----------
+    model : onnx.ModelProto
+        The ONNX model to convert.
+
+    Attributes
+    ----------
+    model : onnx.ModelProto
+        The input ONNX model (with shape inference if successful).
+    graph : torch.fx.Graph
+        The FX graph being constructed.
+    env : Dict[str, torch.fx.Node]
+        Mapping from ONNX tensor names to FX nodes.
+    opset_version : int
+        The opset version for the default ONNX domain.
+    """
+
+    def __init__(
+        self,
+        model: onnx.ModelProto,
+        *,
+        base_dir: Optional[str] = None,
+        memmap_external_data: bool = False,
+    ) -> None:
+        # Try shape inference but preserve original model if it fails
+        # (shape_inference may drop graph contents for large models with external data)
+        try:
+            inferred_model = onnx.shape_inference.infer_shapes(model)
+            # Check if shape inference preserved the model structure
+            if len(inferred_model.graph.node) > 0:
+                model = inferred_model
+            # If nodes were lost, keep original model
+        except Exception:
+            pass
+        self.model: onnx.ModelProto = model
+        self.graph: torch.fx.Graph = torch.fx.Graph()
+        self._base_dir = base_dir
+        self._memmap_external_data = memmap_external_data
+        self.value_info_map = self._create_value_info_map()
+        self.initializer_map = self._create_initializer_map()
+        self.input_names: List[str] = []
+        self.env: Dict[str, torch.fx.Node] = {}
+        self._constants: Dict[str, torch.Tensor] = {}
+        self._submodules: Dict[str, torch.nn.Module] = {}
+        self._opset_versions: Dict[str, int] = self._extract_opset_versions()
+
+    def _extract_opset_versions(self) -> Dict[str, int]:
+        """Extract opset versions for all domains from the model.
+
+        Returns
+        -------
+        Dict[str, int]
+            Dictionary mapping domain names to their opset versions.
+            Empty string "" represents the default ONNX domain.
+        """
+        versions: Dict[str, int] = {}
+        for opset in self.model.opset_import:
+            domain = opset.domain if opset.domain else ""
+            versions[domain] = opset.version
+        return versions
+
+    def _resolve_handler(
+        self, node: onnx.NodeProto
+    ) -> tuple[Callable[["GraphBuilder", onnx.NodeProto], Any], str, int]:
+        """Resolve the handler for an ONNX node and return handler, domain, opset."""
+        domain = node.domain if node.domain else ""
+        opset = self.get_opset_version(domain)
+        handler = get_handler(node.op_type, domain, opset)
+        if handler is None:
+            raise UnsupportedOpError(node.op_type, domain=domain, opset_version=opset)
+        return handler, domain, opset
+
+    def _tag_operator_node(
+        self, node: onnx.NodeProto, fx_node: torch.fx.Node, domain: str
+    ) -> None:
+        """Attach ONNX metadata to an operator node."""
+        if fx_node is not None and hasattr(fx_node, "meta"):
+            self._set_onnx_metadata(
+                fx_node,
+                op_type=node.op_type,
+                name=node.name,
+                domain=domain,
+            )
+
+    def _register_outputs(
+        self, node: onnx.NodeProto, fx_node: torch.fx.Node, domain: str
+    ) -> None:
+        """Register node outputs in the environment."""
+        if len(node.output) == 1:
+            self.env[node.output[0]] = fx_node
+            return
+
+        for i, output_name in enumerate(node.output):
+            if output_name:  # Skip empty output names
+                getitem_node = self.graph.call_function(
+                    lambda x, idx=i: x[idx] if isinstance(x, (tuple, list)) else x,
+                    args=(fx_node, i),
+                )
+                self._set_onnx_metadata(
+                    getitem_node,
+                    op_type=node.op_type,
+                    name=node.name,
+                    domain=domain,
+                    output_index=i,
+                )
+                self.env[output_name] = getitem_node
+
+    @property
+    def opset_version(self) -> int:
+        """Get the opset version for the default ONNX domain.
+
+        Returns
+        -------
+        int
+            The opset version number. Defaults to 1 if not specified.
+        """
+        return self._opset_versions.get("", 1)
+
+    def get_opset_version(self, domain: str = "") -> int:
+        """Get the opset version for a specific domain.
+
+        Parameters
+        ----------
+        domain : str, optional
+            The ONNX domain. Default is "" (standard ONNX domain).
+
+        Returns
+        -------
+        int
+            The opset version number. Defaults to 1 if not specified.
+        """
+        return self._opset_versions.get(domain, 1)
+
+    def build(self) -> torch.fx.GraphModule:
+        self._load_initializers()
+        self._create_placeholders()
+        self._convert_nodes()
+        self._create_outputs()
+        root_module = torch.nn.Module()
+        # Register constants as buffers
+        for name, tensor in self._constants.items():
+            root_module.register_buffer(sanitize_name(name), tensor)
+        # Register submodules
+        for name, submod in self._submodules.items():
+            root_module.add_module(name, submod)
+        module = torch.fx.GraphModule(root_module, self.graph)
+        if self._memmap_external_data:
+            module._onnx2fx_inference_only = True
+        module.graph.lint()
+        return module
+
+    @staticmethod
+    def _set_onnx_metadata(
+        node: torch.fx.Node,
+        *,
+        op_type: str,
+        name: Optional[str] = None,
+        domain: Optional[str] = None,
+        shape: Optional[List[Optional[int]]] = None,
+        dtype: Optional[torch.dtype] = None,
+        output_index: Optional[int] = None,
+    ) -> None:
+        """Populate standard ONNX metadata on an FX node."""
+        node.meta["onnx_op_type"] = op_type
+        if name is not None:
+            node.meta["onnx_name"] = name
+        if domain is not None:
+            node.meta["onnx_domain"] = domain
+        if shape is not None:
+            node.meta["onnx_shape"] = shape
+        if dtype is not None:
+            node.meta["onnx_dtype"] = dtype
+        if output_index is not None:
+            node.meta["onnx_output_index"] = output_index
+
+    def get_value(self, name: str) -> torch.fx.Node:
+        """Get a value (node) by name from the environment.
+
+        Parameters
+        ----------
+        name : str
+            The name of the value.
+
+        Returns
+        -------
+        torch.fx.Node
+            The corresponding FX node.
+
+        Raises
+        ------
+        KeyError
+            If the name is not found in the environment.
+        """
+        if name not in self.env:
+            raise ValueNotFoundError(name, available=list(self.env.keys()))
+        return self.env[name]
+
+    def has_value(self, name: str) -> bool:
+        """Check if a value exists in the environment."""
+        return name in self.env
+
+    def call_function(
+        self,
+        func: Callable[..., Any],
+        args: Sequence[Union[torch.fx.Node, Any]] = (),
+        kwargs: Optional[Dict[str, Any]] = None,
+    ) -> torch.fx.Node:
+        """Create a function call node in the FX graph.
+
+        Parameters
+        ----------
+        func : Callable[..., Any]
+            The function to call. Can be a PyTorch function, lambda, or any callable.
+        args : Sequence[Union[torch.fx.Node, Any]], optional
+            Positional arguments to the function. Can include FX nodes or constants.
+        kwargs : Optional[Dict[str, Any]], optional
+            Keyword arguments to the function.
+
+        Returns
+        -------
+        torch.fx.Node
+            The FX node representing this function call.
+        """
+        fx_node = self.graph.call_function(func, args=tuple(args), kwargs=kwargs or {})
+        return fx_node
+
+    def register_submodule(self, name: str, module: torch.nn.Module) -> str:
+        """Register a submodule for use in the graph.
+
+        Parameters
+        ----------
+        name : str
+            Base name for the submodule.
+        module : torch.nn.Module
+            The submodule to register.
+
+        Returns
+        -------
+        str
+            The actual name used (may be modified to avoid conflicts).
+        """
+        # Sanitize name
+        safe_name = sanitize_name(name)
+        # Ensure unique name
+        if safe_name in self._submodules:
+            counter = 0
+            while f"{safe_name}_{counter}" in self._submodules:
+                counter += 1
+            safe_name = f"{safe_name}_{counter}"
+        self._submodules[safe_name] = module
+        return safe_name
+
+    def call_module(
+        self,
+        module_name: str,
+        args: Sequence[Union[torch.fx.Node, Any]] = (),
+        kwargs: Optional[Dict[str, Any]] = None,
+    ) -> torch.fx.Node:
+        """Create a module call node in the FX graph.
+
+        Parameters
+        ----------
+        module_name : str
+            The name of a registered submodule.
+        args : Sequence[Union[torch.fx.Node, Any]], optional
+            Positional arguments to the module.
+        kwargs : Optional[Dict[str, Any]], optional
+            Keyword arguments to the module.
+
+        Returns
+        -------
+        torch.fx.Node
+            The FX node representing this module call.
+        """
+        fx_node = self.graph.call_module(
+            module_name, args=tuple(args), kwargs=kwargs or {}
+        )
+        return fx_node
+
+    def _create_value_info_map(
+        self,
+    ) -> Dict[str, Tuple[Optional[List[Optional[int]]], Optional[torch.dtype]]]:
+        """Build a mapping from value names to their shape and dtype info."""
+
+        def extract_tensor_shape(
+            value: onnx.ValueInfoProto,
+        ) -> Optional[List[Optional[int]]]:
+            """Extract a list-based representation of a tensor shape from a value info."""
+
+            tensor_type = value.type.tensor_type
+            if not tensor_type.HasField("shape"):
+                return None
+            dims: List[Optional[int]] = []
+            for dim in tensor_type.shape.dim:
+                if dim.HasField("dim_value"):
+                    dims.append(int(dim.dim_value))
+                elif dim.HasField("dim_param"):
+                    dims.append(None)
+                else:
+                    dims.append(None)
+            return dims
+
+        def extract_tensor_dtype(value: onnx.ValueInfoProto) -> Optional[torch.dtype]:
+            """Extract the Torch dtype that corresponds to a value info."""
+
+            onnx_dtype = value.type.tensor_type.elem_type
+            if onnx_dtype == 0:
+                return None
+            torch_dtype = DTYPE_MAP.get(onnx_dtype)
+            if torch_dtype is None:
+                if onnx_dtype == onnx.TensorProto.STRING:
+                    return None
+                raise UnsupportedDTypeError(
+                    onnx_dtype=onnx_dtype,
+                    tensor_name=value.name,
+                    details="value_info dtype not supported",
+                )
+            return torch_dtype
+
+        info_map = {}
+        for value_info in (
+            list(self.model.graph.input)
+            + list(self.model.graph.value_info)
+            + list(self.model.graph.output)
+        ):
+            info_map[value_info.name] = (
+                extract_tensor_shape(value_info),
+                extract_tensor_dtype(value_info),
+            )
+        return info_map
+
+    def is_optional_type(self, name: str) -> bool:
+        """Check if a value has optional type in the ONNX model.
+
+        Parameters
+        ----------
+        name : str
+            The name of the value to check.
+
+        Returns
+        -------
+        bool
+            True if the value has optional type, False otherwise.
+        """
+        # Search in graph inputs
+        for value_info in self.model.graph.input:
+            if value_info.name == name:
+                return value_info.type.HasField("optional_type")
+        # Search in value_info
+        for value_info in self.model.graph.value_info:
+            if value_info.name == name:
+                return value_info.type.HasField("optional_type")
+        # Search in outputs
+        for value_info in self.model.graph.output:
+            if value_info.name == name:
+                return value_info.type.HasField("optional_type")
+        return False
+
+    def _create_initializer_map(self) -> Dict[str, torch.Tensor]:
+        """Build a mapping from initializer names to PyTorch tensors."""
+        init_map = {}
+        for initializer in self.model.graph.initializer:
+            init_map[initializer.name] = self.load_tensor(initializer)
+        return init_map
+
+    def load_tensor(self, tensor: onnx.TensorProto) -> torch.Tensor:
+        """Load an ONNX TensorProto into a Torch tensor."""
+        onnx_dtype = tensor.data_type
+        if DTYPE_MAP.get(onnx_dtype) is None:
+            raise UnsupportedDTypeError(
+                onnx_dtype=onnx_dtype,
+                tensor_name=tensor.name or "<unnamed>",
+                details="initializer dtype not supported",
+            )
+
+        if self._memmap_external_data and (
+            tensor.data_location == onnx.TensorProto.EXTERNAL or tensor.external_data
+        ):
+            info = resolve_external_data(
+                tensor,
+                base_dir=self._base_dir,
+                strict=True,
+            )
+            memmap_array = np.memmap(
+                info.path,
+                dtype=info.numpy_dtype,
+                mode="r",
+                offset=info.offset,
+                shape=info.shape,
+            )
+            return torch.from_numpy(memmap_array)
+
+        np_array = numpy_helper.to_array(tensor)
+        return torch.from_numpy(np_array.copy())
+
+    def _load_initializers(self) -> None:
+        """Load ONNX initializers as constant nodes in the FX graph."""
+        for name, tensor in self.initializer_map.items():
+            # Store in constants dict for later registration as buffers
+            safe_name = sanitize_name(name)
+            self._constants[safe_name] = tensor
+
+            # Create a get_attr node to access the buffer
+            fx_node = self.graph.get_attr(safe_name)
+            self._set_onnx_metadata(
+                fx_node,
+                op_type="Initializer",
+                name=name,
+                shape=list(tensor.shape),
+                dtype=tensor.dtype,
+            )
+            self.env[name] = fx_node
+
+    def _create_placeholders(self) -> None:
+        """Create FX placeholder nodes for graph inputs.
+
+        Note: Inputs that are already loaded as initializers are skipped.
+        """
+        for value in self.model.graph.input:
+            # Skip if already loaded as initializer
+            if value.name in self.env:
+                continue
+
+            # Sanitize name for valid Python identifier
+            safe_name = sanitize_name(value.name)
+            placeholder = self.graph.placeholder(safe_name)
+            info = self.value_info_map.get(value.name)
+            self._set_onnx_metadata(
+                placeholder,
+                op_type="Input",
+                name=value.name,
+                shape=info[0] if info else None,
+                dtype=info[1] if info else None,
+            )
+            self.env[value.name] = placeholder
+            self.input_names.append(value.name)
+
+    def _convert_nodes(self) -> None:
+        # Get graph inputs and initializers for topological sort
+        graph_inputs = {inp.name for inp in self.model.graph.input}
+        initializers = set(self.initializer_map.keys())
+
+        # Topologically sort nodes to handle out-of-order dependencies
+        sorted_nodes = _topological_sort(
+            list(self.model.graph.node),
+            graph_inputs,
+            initializers,
+        )
+
+        for node in sorted_nodes:
+            # Get handler with domain and opset version support
+            handler, domain, _opset = self._resolve_handler(node)
+            fx_node = handler(self, node)
+
+            # Add ONNX metadata to the operator node
+            # Some handlers return a list of nodes (e.g., gradient ops)
+            self._tag_operator_node(node, fx_node, domain)
+            self._register_outputs(node, fx_node, domain)
+
+    def _create_outputs(self) -> None:
+        output_nodes = [self.get_value(value.name) for value in self.model.graph.output]
+        if len(output_nodes) == 1:
+            self.graph.output(output_nodes[0])
+        else:
+            self.graph.output(tuple(output_nodes))
+
+
+__all__ = ["GraphBuilder"]