| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495 |
- # mypy: allow-untyped-defs
- from torch.fx._compatibility import compatibility
- from torch.fx.graph import Graph
- from torch.fx.graph_module import GraphModule
- from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
- from torch.nn import Module
- __all__ = ["HolderModule", "lift_subgraph_as_module", "compare_graphs"]
- @compatibility(is_backward_compatible=False)
- class HolderModule(Module):
- """
- HolderModule is used to copy all the attributes from original module to submodules
- that uses the attributes
- """
- def __init__(self, d):
- super().__init__()
- for k, v in d.items():
- self.add_module(k, v)
- @compatibility(is_backward_compatible=False)
- def lift_subgraph_as_module(
- gm: GraphModule,
- subgraph: Graph,
- comp_name: str = "",
- class_name: str = "GraphModule",
- ) -> tuple[GraphModule, dict[str, str]]:
- """
- Create a GraphModule for subgraph, which copies the necessary attributes from the original parent graph_module.
- Args:
- gm (GraphModule): parent graph module
- subgraph (Graph): a valid subgraph that contains copied nodes from the parent graph
- comp_name (str): name for the new component
- class_name (str): name for the submodule
- """
- # Loop through all module calls (call_module) and param fetches (get_attr)
- # in this component, creating HolderModules as necessary to match the path.
- # e.g. if in the original module there's a get_attr node fetches "conv.weight".
- # We create a HolderModule as root -> add a HolderModule named "conv" ->
- # make "weight" a attribute of "conv" HolderModule and point to conv.weight in
- # the original module.
- submodule = HolderModule({})
- orig_to_split_fqn_mapping: dict[str, str] = {}
- for n in subgraph.nodes:
- if n.op not in ("call_module", "get_attr"):
- continue
- target = n.target
- if not isinstance(target, str):
- raise AssertionError(f"Expected str target, got {type(target)}")
- target_name_parts = target.split(".")
- curr = submodule
- orig_gm = gm
- for name in target_name_parts[:-1]:
- if not hasattr(curr, name):
- curr.add_module(name, HolderModule({}))
- curr = getattr(curr, name)
- orig_gm = getattr(orig_gm, name)
- leaf_node_name = target_name_parts[-1]
- leaf_node = getattr(orig_gm, leaf_node_name)
- orig_to_split_fqn_mapping[target] = f"{comp_name}.{target}"
- # Relies on custom __setattr__ magic.
- setattr(curr, leaf_node_name, leaf_node)
- return GraphModule(submodule, subgraph, class_name), orig_to_split_fqn_mapping
- @compatibility(is_backward_compatible=False)
- def compare_graphs(left: Graph, right: Graph) -> bool:
- """
- Return True if two graphs are identical, i.e they
- - have the same number of outputs in the same order
- - have the same number of inputs in the same order
- - have the same set of nodes, and identical connectivity
- """
- matcher = SubgraphMatcher(left, match_output=True, match_placeholder=True)
- matches = matcher.match(right)
- return len(matches) > 0
|
