Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1 @@
|
||||
from .common import compare_graphs, HolderModule, lift_subgraph_as_module
|
||||
@@ -0,0 +1,95 @@
|
||||
# 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
|
||||
@@ -0,0 +1,303 @@
|
||||
import copy
|
||||
import heapq
|
||||
|
||||
import torch.fx
|
||||
from torch.fx._compatibility import compatibility
|
||||
from torch.fx.graph import Graph
|
||||
from torch.fx.graph_module import GraphModule
|
||||
from torch.fx.node import Node
|
||||
from torch.fx.passes.tools_common import legalize_graph, NodeList, NodeSet # noqa: F401
|
||||
from torch.fx.passes.utils import lift_subgraph_as_module # type: ignore[attr-defined]
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def topo_sort(nodes: NodeList) -> NodeList:
|
||||
# Stable topological sort: among nodes with no dependency between them,
|
||||
# preserve their relative order in the input list. This uses a min-heap
|
||||
# keyed by original position instead of a FIFO queue.
|
||||
indegree_map = dict.fromkeys(nodes, 0)
|
||||
position = {node: i for i, node in enumerate(nodes)}
|
||||
candidates: list[tuple[int, Node]] = []
|
||||
|
||||
for node in nodes:
|
||||
for n in node.all_input_nodes:
|
||||
if n in indegree_map:
|
||||
indegree_map[node] += 1
|
||||
if indegree_map[node] == 0:
|
||||
heapq.heappush(candidates, (position[node], node))
|
||||
|
||||
sorted_nodes: NodeList = []
|
||||
while candidates:
|
||||
_, node = heapq.heappop(candidates)
|
||||
sorted_nodes.append(node)
|
||||
|
||||
for n in node.users:
|
||||
if n in indegree_map:
|
||||
indegree_map[n] -= 1
|
||||
if indegree_map[n] == 0:
|
||||
heapq.heappush(candidates, (position[n], n))
|
||||
|
||||
if len(nodes) != len(sorted_nodes):
|
||||
raise AssertionError(
|
||||
"topological sorted nodes doesn't have same length as input nodes"
|
||||
)
|
||||
|
||||
return sorted_nodes
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def validate_partition(partition: NodeList) -> bool:
|
||||
# verify the partition doesn't form a dependency cycle in the original graph
|
||||
# returns True for valid partition, False for invalid
|
||||
|
||||
partition_set = set(partition)
|
||||
|
||||
outputs: NodeList = []
|
||||
for node in partition_set:
|
||||
for user_node in node.users:
|
||||
if user_node not in partition_set:
|
||||
# external user node, need to expose as an output
|
||||
outputs.append(user_node)
|
||||
|
||||
# Perform BFS on the partition outputs.
|
||||
# If it reaches a node within the partition, then it found a cycle.
|
||||
# This function takes the ownership of `root_nodes` and may modify it.
|
||||
def bfs_find_cycle(root_nodes: NodeList) -> bool:
|
||||
# Set used to exclude nodes that have already been visited.
|
||||
# If a node has been visited, that node and all its children have
|
||||
# been checked for cycles.
|
||||
visited: NodeSet = set()
|
||||
|
||||
# Start with `root_nodes` and traverse through (toward child nodes)
|
||||
# their connected sub-graph. Nodes in `visited` won't be added
|
||||
# to `queue` again.
|
||||
queue: NodeList = root_nodes
|
||||
while queue:
|
||||
current = queue.pop()
|
||||
visited.add(current)
|
||||
if current in partition_set:
|
||||
# Started from partition's `output` nodes, and reached
|
||||
# another node in partition. Cycle!
|
||||
return True
|
||||
for user_node in current.users:
|
||||
if user_node in visited:
|
||||
continue
|
||||
queue.append(user_node)
|
||||
# `root_nodes` don't cause cycle.
|
||||
return False
|
||||
|
||||
# Use all output nodes as roots to traverse
|
||||
# the graph to check cycles.
|
||||
if bfs_find_cycle(outputs):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def fuse_as_graphmodule(
|
||||
gm: GraphModule,
|
||||
nodes: NodeList,
|
||||
module_name: str,
|
||||
partition_lookup_table: dict[Node, int | None] | None = None,
|
||||
*,
|
||||
always_return_tuple: bool = False,
|
||||
) -> tuple[GraphModule, tuple[Node, ...], tuple[Node, ...]]:
|
||||
"""
|
||||
Fuse nodes in graph_module into a GraphModule.
|
||||
|
||||
Args:
|
||||
gm (GraphModule): target graph_module
|
||||
|
||||
nodes (List[Node]): list of nodes in `gm` to fuse, where the node must be topologically sorted
|
||||
|
||||
module_name: class name for the fused GraphModule
|
||||
|
||||
partition_lookup_table (Optional[Dict[Node, None]]): optional dict of nodes to speed up lookup
|
||||
|
||||
always_return_tuple (bool): whether to always return a tuple, even if there is only one output
|
||||
|
||||
Returns:
|
||||
fused_gm (GraphModule): fused graph module, where its node is a copy of `nodes` in `gm`
|
||||
|
||||
original_inputs (Tuple[Node, ...]): input nodes to `nodes` in original `gm`
|
||||
|
||||
original_outputs (Tuple[Node, ...]): consumer nodes of `nodes` in original `gm`
|
||||
|
||||
"""
|
||||
|
||||
# assumption: nodes are already sorted in topo order
|
||||
|
||||
for node in nodes:
|
||||
if node.graph.owning_module is not gm:
|
||||
raise AssertionError(
|
||||
f"{node} doesn't belong to passed in graph module {gm._get_name()}"
|
||||
)
|
||||
if node._erased:
|
||||
raise AssertionError(f"{node} has been removed from owning graph")
|
||||
if node not in gm.graph._find_nodes_lookup_table:
|
||||
raise AssertionError(
|
||||
f"{node} is not found in graph module {gm._get_name()}"
|
||||
)
|
||||
|
||||
# validates partition doesn't introduce dependency circles in the graph
|
||||
if not validate_partition(nodes):
|
||||
raise AssertionError("Invalid partition, found dependency cycles")
|
||||
|
||||
# if no dict of partition nodes is provided, reconstruct it by nodes list to reduce lookup time
|
||||
if partition_lookup_table is None:
|
||||
partition_lookup_table = dict.fromkeys(nodes)
|
||||
|
||||
subgraph = Graph()
|
||||
|
||||
node_to_placeholder: dict[
|
||||
Node, Node
|
||||
] = {} # mapping of nodes from old graph to placeholder in new graph
|
||||
node_map: dict[Node, Node] = {} # mapping of nodes from old graph to new graph
|
||||
|
||||
# handles inputs through graph.node_copy's arg_transform functions
|
||||
def remap_inputs(x: Node) -> Node:
|
||||
if x.op == "get_attr":
|
||||
# TODO: do we really need copy the get_attr node into the graph?
|
||||
# do something here
|
||||
pass
|
||||
|
||||
if x in partition_lookup_table:
|
||||
# x is inside subgraph, return the copied node
|
||||
# the node should have been copied already, as we are copying graph in the topological order
|
||||
return node_map[x]
|
||||
|
||||
if x not in node_to_placeholder:
|
||||
# x is not in subgraph, create a new placeholder for subgraph
|
||||
placeholder_node = subgraph.placeholder(x.name, type_expr=x.type)
|
||||
# copy all meta fields, even if some fields might be irrelevant for the placeholder node
|
||||
placeholder_node.meta = copy.copy(x.meta)
|
||||
node_to_placeholder[x] = placeholder_node
|
||||
|
||||
return node_to_placeholder[x]
|
||||
|
||||
# copy nodes in topological order
|
||||
for node in nodes:
|
||||
new_node = subgraph.node_copy(node, remap_inputs)
|
||||
node_map[node] = new_node
|
||||
|
||||
# handles outputs
|
||||
output_mapping: dict[Node, Node] = {} # mapping from old output to new outputs
|
||||
|
||||
for node in nodes:
|
||||
for user_node in node.users:
|
||||
if user_node not in partition_lookup_table:
|
||||
# external user node, need to expose as an output
|
||||
output_mapping[node] = node_map[node]
|
||||
|
||||
# outs contain nodes in the new subgraph
|
||||
outs = tuple(output_mapping.values())
|
||||
|
||||
if always_return_tuple:
|
||||
# always return a tuple, even if there is only one output
|
||||
subgraph.output(outs)
|
||||
else:
|
||||
# If there's a single output then return it directly, otherwise return a tuple.
|
||||
subgraph.output(outs[0] if len(outs) == 1 else outs)
|
||||
|
||||
# lint to ensure correctness
|
||||
subgraph.lint() # type: ignore[no-untyped-call]
|
||||
fused_gm: GraphModule
|
||||
fused_gm, _ = lift_subgraph_as_module(
|
||||
gm, subgraph, comp_name="", class_name=module_name
|
||||
)
|
||||
|
||||
# sub_gm's input nodes in the original module
|
||||
original_inputs: tuple[Node, ...] = tuple(node_to_placeholder.keys())
|
||||
|
||||
# sub_gm's outputs node in the original module
|
||||
original_outputs: tuple[Node, ...] = tuple(output_mapping.keys())
|
||||
|
||||
return fused_gm, original_inputs, original_outputs
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def insert_subgm(
|
||||
gm: GraphModule,
|
||||
sub_gm: GraphModule,
|
||||
orig_inputs: tuple[Node, ...],
|
||||
orig_outputs: tuple[Node, ...],
|
||||
insertion_point: Node | None = None,
|
||||
) -> GraphModule:
|
||||
# add sub_gm into gm
|
||||
submodule_name = sub_gm.__class__.__name__
|
||||
gm.add_submodule(submodule_name, sub_gm)
|
||||
|
||||
# Use provided insertion point, or fall back to last output node for backwards compat
|
||||
if insertion_point is None:
|
||||
for node in reversed(gm.graph.nodes):
|
||||
if node in orig_outputs:
|
||||
insertion_point = node
|
||||
break
|
||||
if insertion_point is None:
|
||||
raise AssertionError(
|
||||
"Cannot determine insertion point: no insertion_point provided and "
|
||||
"orig_outputs is empty. Pass the last partition node as insertion_point."
|
||||
)
|
||||
|
||||
# Create a call_module node in main graph.
|
||||
with gm.graph.inserting_after(insertion_point):
|
||||
module_node = gm.graph.call_module(
|
||||
submodule_name, args=orig_inputs, kwargs=None
|
||||
)
|
||||
output_node = sub_gm.graph.output_node()
|
||||
|
||||
# Replace uses of original outputs with the fused module outputs.
|
||||
# If there are no external outputs, skip replacement (nothing to replace).
|
||||
if orig_outputs:
|
||||
next_node = module_node.next
|
||||
with gm.graph.inserting_before(next_node):
|
||||
if len(orig_outputs) == 1 and not isinstance(output_node.args[0], tuple):
|
||||
# main_remapping[comp.orig_outputs[0]] = module_node
|
||||
orig_outputs[0].replace_all_uses_with(module_node, propagate_meta=True)
|
||||
else:
|
||||
for i, orig_output in enumerate(orig_outputs):
|
||||
# Use Proxy to record getitem access.
|
||||
proxy_out = torch.fx.Proxy(module_node)[i].node # type: ignore[index]
|
||||
orig_output.replace_all_uses_with(proxy_out, propagate_meta=True)
|
||||
|
||||
module_node.meta["val"] = tuple(
|
||||
orig_output.meta.get("val", None) for orig_output in orig_outputs
|
||||
)
|
||||
return gm
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def erase_nodes(gm: GraphModule, nodes: NodeList) -> None:
|
||||
# erase original nodes in inversed topological order
|
||||
for node in reversed(nodes):
|
||||
gm.graph.erase_node(node)
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def fuse_by_partitions(
|
||||
gm: GraphModule,
|
||||
partitions: list[dict[Node, int | None]],
|
||||
prefix: str = "fused_",
|
||||
always_return_tuple: bool = False,
|
||||
) -> GraphModule:
|
||||
for partition_id, partition in enumerate(partitions):
|
||||
sorted_nodes = topo_sort(list(partition))
|
||||
|
||||
submodule_name = prefix + str(partition_id)
|
||||
sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(
|
||||
gm,
|
||||
sorted_nodes,
|
||||
submodule_name,
|
||||
partition,
|
||||
always_return_tuple=always_return_tuple,
|
||||
)
|
||||
|
||||
insert_subgm(gm, sub_gm, orig_inputs, orig_outputs, sorted_nodes[-1])
|
||||
|
||||
erase_nodes(gm, sorted_nodes)
|
||||
|
||||
torch.fx.passes.tools_common.stable_topological_sort(gm)
|
||||
gm.graph.lint()
|
||||
|
||||
return gm
|
||||
@@ -0,0 +1,449 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.fx import Graph, Node
|
||||
from torch.fx._compatibility import compatibility
|
||||
|
||||
|
||||
__all__ = ["SubgraphMatcher", "InternalMatch"]
|
||||
|
||||
|
||||
# Set`PYTORCH_MATCHER_LOGLEVEL=INFO` to see debug logs
|
||||
def _init_logger():
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
level = os.environ.get("PYTORCH_MATCHER_LOGLEVEL", "WARNING").upper()
|
||||
logger.setLevel(level)
|
||||
console = logging.StreamHandler()
|
||||
formatter = logging.Formatter("%(filename)s > %(message)s")
|
||||
console.setFormatter(formatter)
|
||||
console.setLevel(level)
|
||||
# add the handlers to the logger
|
||||
logger.addHandler(console)
|
||||
logger.propagate = False
|
||||
return logger
|
||||
|
||||
|
||||
logger = _init_logger()
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
@dataclass
|
||||
class InternalMatch:
|
||||
# Nodes from which the match was found
|
||||
anchors: list[Node]
|
||||
# Maps nodes in the pattern subgraph to nodes in the larger graph
|
||||
nodes_map: dict[Node, Node] = field(default_factory=dict)
|
||||
|
||||
# nodes in target graph that are matched placeholder in pattern
|
||||
placeholder_nodes: list[Node] = field(default_factory=list)
|
||||
|
||||
# nodes in matched subgraph returned by output
|
||||
returning_nodes: list[Node] = field(default_factory=list)
|
||||
|
||||
# map from a string name to a node in the target graph
|
||||
# only available if the matcher is `SubgraphMatcherWithNameNodesMap`
|
||||
name_node_map: dict[str, Node] = field(default_factory=dict)
|
||||
|
||||
def __copy__(self):
|
||||
return InternalMatch(
|
||||
anchors=self.anchors,
|
||||
nodes_map=self.nodes_map.copy(),
|
||||
placeholder_nodes=self.placeholder_nodes.copy(),
|
||||
returning_nodes=self.returning_nodes.copy(),
|
||||
)
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
class SubgraphMatcher:
|
||||
def __init__(
|
||||
self,
|
||||
pattern: Graph,
|
||||
match_output: bool = False,
|
||||
match_placeholder: bool = False,
|
||||
remove_overlapping_matches: bool = True,
|
||||
ignore_literals: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
pattern: the targeted matching pattern, represented in fx.Graph.
|
||||
match_output: If True, output node in the pattern graph will be treated as a part of the targeted pattern.
|
||||
If False, output node is ignored during match.
|
||||
match_placeholder: If True, placeholder node in the pattern graph will be treated as a part of
|
||||
the targeted pattern. If False, placeholder nodes will be used a wildcard.
|
||||
remove_overlapping_matches: If True, in the case of overlapping matches, only the first match
|
||||
will be returned.
|
||||
ignore_literals: If True, will not check if literals are equal and
|
||||
will instead treat them as wildcards.
|
||||
"""
|
||||
|
||||
self.pattern = pattern
|
||||
self.match_output = match_output
|
||||
self.match_placeholder = match_placeholder
|
||||
self.remove_overlapping_matches = remove_overlapping_matches
|
||||
self.ignore_literals = ignore_literals
|
||||
|
||||
if len(pattern.nodes) == 0:
|
||||
raise ValueError(
|
||||
"SubgraphMatcher cannot be initialized with an empty pattern"
|
||||
)
|
||||
|
||||
for node in pattern.nodes:
|
||||
if node.op != "output" and not node.is_impure():
|
||||
if len(node.users) == 0:
|
||||
raise AssertionError(
|
||||
"SubgraphMatcher cannot be initialized with an pattern with dead code"
|
||||
)
|
||||
|
||||
# TODO: assert pattern is a connected graph
|
||||
|
||||
self.pattern_placeholder_nodes = [
|
||||
n for n in pattern.nodes if n.op == "placeholder"
|
||||
]
|
||||
output_node = next(iter(reversed(pattern.nodes)))
|
||||
# nodes returned by outputs
|
||||
self.pattern_returning_nodes: list[Node] = output_node.all_input_nodes
|
||||
|
||||
self.pattern_anchors: list[Node] = []
|
||||
if match_output:
|
||||
self.pattern_anchors = [output_node]
|
||||
else:
|
||||
# If a node has output_node as the ONLY user, then this node is a graph sink,
|
||||
# and should be matched against as an anchor
|
||||
self.pattern_anchors = [
|
||||
n for n in output_node.all_input_nodes if len(n.users) == 1
|
||||
]
|
||||
|
||||
def _match_attributes(self, pn: Node, gn: Node) -> bool:
|
||||
# Attributes matching is complicated. Right now we only support matching constant tensor
|
||||
if not isinstance(pn.target, str):
|
||||
raise AssertionError(f"pn.target {pn.target} must be a string.")
|
||||
if not isinstance(gn.target, str):
|
||||
raise AssertionError(f"gn.target {gn.target} must be a string.")
|
||||
|
||||
pn_value = torch.fx.graph_module._get_attr(pn.graph.owning_module, pn.target)
|
||||
gn_value = torch.fx.graph_module._get_attr(gn.graph.owning_module, gn.target)
|
||||
|
||||
if type(pn_value) is not type(gn_value):
|
||||
return False
|
||||
|
||||
# Don't require exact match on tensor values.
|
||||
if isinstance(pn_value, torch.Tensor):
|
||||
return isinstance(gn_value, torch.Tensor)
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported type {pn_value} when matching attributes")
|
||||
# pyrefly: ignore [unreachable]
|
||||
return False
|
||||
|
||||
def _nodes_are_equal(self, pn: Node, gn: Node, node_name_match: str = "") -> bool:
|
||||
# if exact match for placeholder is not required, then use placeholder as a wildcard
|
||||
if not self.match_placeholder and pn.op == "placeholder":
|
||||
return True
|
||||
|
||||
if node_name_match and node_name_match in gn.name:
|
||||
return True
|
||||
|
||||
if pn.op == gn.op:
|
||||
if pn.op == "placeholder" or pn.op == "output":
|
||||
return True
|
||||
elif pn.op == "get_attr":
|
||||
return self._match_attributes(pn, gn)
|
||||
return pn.target == gn.target
|
||||
return False
|
||||
|
||||
def _is_contained(self, nodes_map: dict[Node, Node]) -> bool:
|
||||
# `lookup` represents all the nodes in `original_graph`
|
||||
# that are part of `pattern`
|
||||
|
||||
# Placeholders can be used by other nodes in the graphs
|
||||
lookup: dict[Node, Node] = {
|
||||
gn: pn for pn, gn in nodes_map.items() if pn.op != "placeholder"
|
||||
}
|
||||
|
||||
for gn, pn in lookup.items():
|
||||
# nodes returned by output are allowed to be used in other areas of the graph
|
||||
if pn in self.pattern_returning_nodes:
|
||||
continue
|
||||
|
||||
for user in gn.users:
|
||||
# If this node has users that were not in `lookup`, then it must leak out of the
|
||||
# pattern subgraph
|
||||
if user not in lookup:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _remove_overlapping_matches(
|
||||
self, matches: list[InternalMatch]
|
||||
) -> list[InternalMatch]:
|
||||
non_overlapping_matches: list[InternalMatch] = []
|
||||
nodes_matched: set[Node] = set()
|
||||
|
||||
for match in matches:
|
||||
found_overlap = False
|
||||
for pn, gn in match.nodes_map.items():
|
||||
if pn.op not in {"placeholder", "output"} and gn in nodes_matched:
|
||||
found_overlap = True
|
||||
break
|
||||
|
||||
if not found_overlap:
|
||||
non_overlapping_matches.append(match)
|
||||
for pn, gn in match.nodes_map.items():
|
||||
if pn.op not in {"placeholder", "output"}:
|
||||
nodes_matched.add(gn)
|
||||
return non_overlapping_matches
|
||||
|
||||
def _match_literals(self, pn: Any, gn: Any, match: InternalMatch) -> bool:
|
||||
if isinstance(pn, Node) and isinstance(gn, Node):
|
||||
raise AssertionError("pn and gn cannot both be Node")
|
||||
|
||||
if isinstance(pn, Node) and not isinstance(gn, Node):
|
||||
if pn.op == "placeholder":
|
||||
# Check if we've already matched these nodes in the current
|
||||
# traversal
|
||||
if pn in match.nodes_map:
|
||||
return match.nodes_map[pn] == gn
|
||||
|
||||
match.nodes_map[pn] = gn
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
elif not isinstance(pn, Node) and isinstance(gn, Node):
|
||||
return False
|
||||
else:
|
||||
return type(gn) is type(pn) and gn == pn
|
||||
|
||||
def _match_nodes(
|
||||
self, pn: Node, gn: Node, match: InternalMatch, node_name_match: str = ""
|
||||
) -> bool:
|
||||
logger.info(" matching %s to %s", pn, gn)
|
||||
|
||||
if not (isinstance(pn, Node) and isinstance(gn, Node)):
|
||||
raise AssertionError(f"pn and gn must be Node, pn: {pn}, gn: {gn}")
|
||||
|
||||
# Check if we've already matched these nodes in the current
|
||||
# traversal
|
||||
if pn in match.nodes_map:
|
||||
return match.nodes_map[pn] == gn
|
||||
|
||||
# TODO: use a more efficient way to check if gn is matched before: two-way dict
|
||||
if gn in match.nodes_map.values():
|
||||
return False
|
||||
|
||||
if not self._nodes_are_equal(pn, gn, node_name_match):
|
||||
return False
|
||||
|
||||
# Optimistically mark `pn` as a match for `gn`, and save a local copy of match
|
||||
saved_match = copy.copy(match)
|
||||
match.nodes_map[pn] = gn
|
||||
|
||||
# Placeholder is a wildcard and can be matched with any python object
|
||||
# (including list/tuple)
|
||||
if pn.op == "placeholder":
|
||||
return True
|
||||
|
||||
# Recursively traverse upwards to check if `pn` is a true
|
||||
# match for `gn`
|
||||
match_found = True
|
||||
|
||||
def _match_args(args1: list | tuple, args2: list | tuple) -> bool:
|
||||
if len(args1) != len(args2):
|
||||
return False
|
||||
|
||||
for a1, a2 in zip(args1, args2):
|
||||
if isinstance(a1, Node) and isinstance(a2, Node):
|
||||
matched = self._match_nodes(a1, a2, match)
|
||||
elif isinstance(a1, (list, tuple)) and isinstance(a2, (list, tuple)):
|
||||
matched = _match_args(a1, a2)
|
||||
else:
|
||||
matched = (
|
||||
self._match_literals(a1, a2, match) or self.ignore_literals
|
||||
)
|
||||
|
||||
if not matched:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
# Flatten all args/kwargs into 1 list of args
|
||||
pn_args, gn_args = None, None
|
||||
if (
|
||||
(
|
||||
len(pn.args) != len(gn.args)
|
||||
or list(pn.kwargs.keys()) != list(gn.kwargs.keys())
|
||||
)
|
||||
and pn.op == "call_function"
|
||||
and isinstance(pn.target, torch._ops.OpOverload)
|
||||
):
|
||||
args_schema = pn.target._schema.arguments
|
||||
|
||||
def get_all_arguments(orig_args, orig_kwargs):
|
||||
all_args = []
|
||||
for i, schema in enumerate(args_schema):
|
||||
if schema.name in orig_kwargs:
|
||||
all_args.append(orig_kwargs[schema.name])
|
||||
elif not schema.kwarg_only and i < len(orig_args):
|
||||
all_args.append(orig_args[i])
|
||||
else:
|
||||
all_args.append(schema.default_value)
|
||||
return all_args
|
||||
|
||||
pn_args = get_all_arguments(pn.args, pn.kwargs)
|
||||
gn_args = get_all_arguments(gn.args, gn.kwargs)
|
||||
|
||||
elif len(pn.args) == len(gn.args) and list(pn.kwargs.keys()) == list(
|
||||
gn.kwargs.keys()
|
||||
):
|
||||
pn_args = list(pn.args)
|
||||
gn_args = list(gn.args)
|
||||
pn_args.extend(list(pn.kwargs.values()))
|
||||
gn_args.extend(list(gn.kwargs.values()))
|
||||
else:
|
||||
match_found = False
|
||||
|
||||
match_found = (
|
||||
match_found
|
||||
and pn_args is not None
|
||||
and gn_args is not None
|
||||
and _match_args(pn_args, gn_args)
|
||||
)
|
||||
|
||||
if not match_found:
|
||||
# revert to saved_match before matching with current node
|
||||
match = copy.copy(saved_match)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def match(self, graph: Graph, node_name_match: str = "") -> list[InternalMatch]:
|
||||
"""
|
||||
Returns:
|
||||
The matched subgraphs.
|
||||
The returned subgraph would be fully self-contained, meaning the nodes (except placeholder
|
||||
and nodes returned by output) can only be consumed by nodes within the matched subgraph.
|
||||
|
||||
Subgraph pattern matcher is implemented with the backtracking style in the following steps:
|
||||
|
||||
1. We first identify all the anchor nodes in the pattern graph. The anchor nodes
|
||||
are the "sinks" (nodes with no user other than the output node) of the pattern graph.
|
||||
One pattern graph could have multiple anchors if it has multiple return values.
|
||||
|
||||
2. In the target graph, we identify the potential candidate nodes that can be matched
|
||||
with each anchor. These anchor-candidate pairs are the starting points for
|
||||
pairwise per-node matching.
|
||||
|
||||
3. For each anchor-candidate pair, we simultaneously traverse backwards (DFS) in both
|
||||
pattern and target graphs. For every pattern nodes along traversal path, we compare it
|
||||
against the target nodes. In case any comparison failed, the match for this anchor-candidate
|
||||
pair fails. A match is found when DFS completes traversing the graph. See `self._match_nodes`
|
||||
for more details.
|
||||
|
||||
4. In the case of multiple anchors, every anchor will need to find a match using step 3.
|
||||
In addition, the matches found between anchors need to have a common intersection node
|
||||
in order for the match to be valid. This is implemented with backtracking. See `backtracking`
|
||||
for more details.
|
||||
|
||||
Notice: graph traversal must be done in the reverser order because a tensor can have multiple
|
||||
consumers, but can only have a single producer. Only with reverser order, we can we jointly
|
||||
traverse the pattern and target graph in a deterministic path.
|
||||
|
||||
Warning: In theory, this backtracking algorithm have an **exponential** time complexity. However,
|
||||
in practice, it's unlikely to blow up.
|
||||
|
||||
"""
|
||||
from torch.fx.passes.utils.fuser_utils import validate_partition
|
||||
|
||||
# find candidate nodes to match with pattern anchors
|
||||
match_candidates: dict[Node, list[Node]] = defaultdict(list)
|
||||
for pattern_anchor in self.pattern_anchors:
|
||||
for node in graph.nodes:
|
||||
if self._nodes_are_equal(pattern_anchor, node, node_name_match):
|
||||
match_candidates[pattern_anchor].append(node)
|
||||
match_candidates_list = list(match_candidates.items())
|
||||
|
||||
logger.info("Initial match_candidates_list: %s\n", match_candidates_list)
|
||||
|
||||
matches: list[InternalMatch] = []
|
||||
|
||||
def backtracking(anchor_index, match):
|
||||
if anchor_index == len(match_candidates_list):
|
||||
match.placeholder_nodes = [
|
||||
match.nodes_map[pn] for pn in self.pattern_placeholder_nodes
|
||||
]
|
||||
match.returning_nodes = [
|
||||
match.nodes_map[pn] for pn in self.pattern_returning_nodes
|
||||
]
|
||||
matches.append(match)
|
||||
|
||||
logger.info("Found a match: %s\n", match)
|
||||
return
|
||||
|
||||
pattern_anchor, candidate_nodes = match_candidates_list[anchor_index]
|
||||
saved_match = copy.copy(match)
|
||||
|
||||
for node in candidate_nodes:
|
||||
logger.info("Trying to match anchor %s to %s", pattern_anchor, node)
|
||||
|
||||
match_found = self._match_nodes(
|
||||
pattern_anchor, node, match, node_name_match
|
||||
)
|
||||
if match_found:
|
||||
# match next anchor
|
||||
backtracking(anchor_index + 1, match)
|
||||
else:
|
||||
logger.info(
|
||||
"Failed to match anchor %s to %s\n", pattern_anchor, node
|
||||
)
|
||||
|
||||
# revert to saved_match before matching with current anchor
|
||||
match = copy.copy(saved_match)
|
||||
|
||||
match = InternalMatch(anchors=self.pattern_anchors)
|
||||
if match_candidates_list:
|
||||
backtracking(0, match)
|
||||
|
||||
# filter out the matches where the subgraph is not fully_contained
|
||||
before = len(matches)
|
||||
matches = [match for match in matches if self._is_contained(match.nodes_map)]
|
||||
after = len(matches)
|
||||
if before != after:
|
||||
logger.info(
|
||||
"Filtered out %s matches because they are not fully contained",
|
||||
before - after,
|
||||
)
|
||||
|
||||
# filter out the matches that form a cycle if the subgraph is fused
|
||||
valid_matches = []
|
||||
for match in matches:
|
||||
matched_compute_nodes = [
|
||||
gn
|
||||
for pn, gn in match.nodes_map.items()
|
||||
if pn.op not in {"placeholder", "output"}
|
||||
]
|
||||
if validate_partition(matched_compute_nodes):
|
||||
valid_matches.append(match)
|
||||
if len(valid_matches) != len(matches):
|
||||
logger.info(
|
||||
"Filtered out %s matches because \
|
||||
matched subgraph would form a cycle if fused",
|
||||
len(matches) - len(valid_matches),
|
||||
)
|
||||
|
||||
if self.remove_overlapping_matches:
|
||||
before = len(valid_matches)
|
||||
matches = self._remove_overlapping_matches(valid_matches)
|
||||
after = len(matches)
|
||||
if before != after:
|
||||
logger.info(
|
||||
"Filtered out %s matches because matched subgraphs are overlapping",
|
||||
before - after,
|
||||
)
|
||||
|
||||
logger.info("Matches returned: %s", matches)
|
||||
|
||||
return matches
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
from torch.fx import Graph, GraphModule, Node
|
||||
from torch.fx._compatibility import compatibility
|
||||
|
||||
from .matcher_utils import InternalMatch, SubgraphMatcher
|
||||
|
||||
|
||||
__all__ = ["SubgraphMatcherWithNameNodeMap"]
|
||||
|
||||
|
||||
def _split_to_graph_and_name_node_map(
|
||||
gm: GraphModule,
|
||||
) -> tuple[GraphModule, dict[str, Node]]:
|
||||
from torch.fx.graph import _PyTreeInfo
|
||||
from torch.utils._pytree import tree_flatten, tree_unflatten
|
||||
|
||||
name_node_map = {}
|
||||
for n in gm.graph.nodes:
|
||||
if n.op == "output":
|
||||
if gm._out_spec is None:
|
||||
raise AssertionError("gm._out_spec is None")
|
||||
output = tree_unflatten(n.args[0], gm._out_spec)
|
||||
if not isinstance(output, tuple):
|
||||
raise AssertionError("Expecting the pattern graph to return a tuple")
|
||||
if len(output) < 2:
|
||||
raise AssertionError(
|
||||
"Expecting the pattern graph to have at least two outputs"
|
||||
)
|
||||
*out, name_node_map = output
|
||||
flattened, out_spec = tree_flatten(out)
|
||||
if not isinstance(name_node_map, dict):
|
||||
raise AssertionError(
|
||||
"Expecting the input graph to have a dict output as the last element"
|
||||
)
|
||||
n.args = (flattened,)
|
||||
orig_pytree_info = gm._graph._codegen.pytree_info # type: ignore[attr-defined]
|
||||
gm._graph._codegen.pytree_info = _PyTreeInfo( # type: ignore[attr-defined]
|
||||
orig_pytree_info.orig_args, orig_pytree_info.in_spec, out_spec
|
||||
)
|
||||
gm.recompile()
|
||||
return gm, name_node_map
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
class SubgraphMatcherWithNameNodeMap(SubgraphMatcher):
|
||||
"""Extends SubgraphMatcher to support querying the matched subgraph nodes through node name,
|
||||
this requires pattern to have specific format (returning and additional dictionary at the output,
|
||||
that has node name as key, and the node in the pattern graph as value, see Example for more details)
|
||||
|
||||
Difference with SubgraphMatcher is that it takes a `pattern_gm` GraphModule as input during
|
||||
initialization since we need to modify the graph (which requires `recompile` the GraphModule)
|
||||
|
||||
Example::
|
||||
def pattern(x, weight):
|
||||
conv = F.conv2d(x, weight)
|
||||
relu = F.relu(conv)
|
||||
return relu, {"conv": conv, "relu": relu}
|
||||
|
||||
|
||||
def target_graph(x, weight):
|
||||
conv = F.conv2d(x, weight)
|
||||
relu = F.relu(conv)
|
||||
relu *= 2
|
||||
return relu
|
||||
|
||||
|
||||
pattern_gm = export(pattern, example_inputs).module()
|
||||
target_gm = export(target_graph, example_inputs).module()
|
||||
matcher = SubgraphMatcherWithNameNodeMap(pattern_gm)
|
||||
matches = matcher.match(target_gm)
|
||||
for match in matches:
|
||||
match.name_node_map["conv"].meta["annotation"] = ...
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pattern_gm: GraphModule,
|
||||
match_output: bool = False,
|
||||
match_placeholder: bool = False,
|
||||
remove_overlapping_matches: bool = True,
|
||||
ignore_literals: bool = False,
|
||||
) -> None:
|
||||
pattern_gm, name_node_map = _split_to_graph_and_name_node_map(pattern_gm)
|
||||
self.name_node_map = name_node_map
|
||||
super().__init__(
|
||||
pattern_gm.graph,
|
||||
match_output,
|
||||
match_placeholder,
|
||||
remove_overlapping_matches,
|
||||
ignore_literals,
|
||||
)
|
||||
|
||||
def match(self, graph: Graph, node_name_match: str = "") -> list[InternalMatch]:
|
||||
"""The returned InternalMatch will have name_node_map populated with a map
|
||||
from node name (str) to the target node, e.g.
|
||||
{"conv": target_conv_ndoe, "relu": target_relu_node}
|
||||
|
||||
this requires the pattern graph returns an additional
|
||||
output of node name to node, e.g. instead of:
|
||||
```
|
||||
def pattern(...):
|
||||
...
|
||||
return relu
|
||||
```
|
||||
we should do:
|
||||
```
|
||||
def pattern(...):
|
||||
...
|
||||
return relu, {"conv": conv, "relu": relu}
|
||||
``` instead
|
||||
"""
|
||||
internal_matches = super().match(graph, node_name_match)
|
||||
for internal_match in internal_matches:
|
||||
for k, n in self.name_node_map.items():
|
||||
internal_match.name_node_map[k] = internal_match.nodes_map[n]
|
||||
return internal_matches
|
||||
+191
@@ -0,0 +1,191 @@
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from torch.fx._compatibility import compatibility
|
||||
from torch.fx.graph import Graph
|
||||
from torch.fx.node import Node
|
||||
|
||||
|
||||
__all__ = ["get_source_partitions", "check_subgraphs_connected", "SourcePartition"]
|
||||
|
||||
|
||||
# Set`PYTORCH_MATCHER_LOGLEVEL=INFO` to see debug logs
|
||||
def _init_logger() -> logging.Logger:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
level = os.environ.get("PYTORCH_MATCHER_LOGLEVEL", "WARNING").upper()
|
||||
logger.setLevel(level)
|
||||
console = logging.StreamHandler()
|
||||
formatter = logging.Formatter("%(filename)s > %(message)s")
|
||||
console.setFormatter(formatter)
|
||||
console.setLevel(level)
|
||||
# add the handlers to the logger
|
||||
logger.addHandler(console)
|
||||
logger.propagate = False
|
||||
return logger
|
||||
|
||||
|
||||
logger = _init_logger()
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
@dataclass
|
||||
class SourcePartition:
|
||||
# Nodes in a particular partition
|
||||
nodes: list[Node]
|
||||
|
||||
# The source these nodes decomposed from
|
||||
source: Any
|
||||
|
||||
# Nodes in the graph that are needed as inputs to the partition
|
||||
# These do not include the params of the partition
|
||||
input_nodes: list[Node] = field(default_factory=list)
|
||||
|
||||
# Nodes in the partition that are being used by nodes outside of the
|
||||
# partition
|
||||
output_nodes: list[Node] = field(default_factory=list)
|
||||
|
||||
# Parameters that are being used
|
||||
params: list[Node] = field(default_factory=list)
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False) # type: ignore[misc]
|
||||
def get_source_partitions(
|
||||
graph: Graph,
|
||||
wanted_sources: list[Any],
|
||||
filter_fn: Callable[[Node], bool] | None = None,
|
||||
) -> dict[Any, list[SourcePartition]]:
|
||||
"""
|
||||
Args:
|
||||
graph: The graph we want to partition
|
||||
wanted_sources: List of sources of nodes that were decomposed from this
|
||||
source. This can be a function (ex. torch.nn.functional.linear) or a
|
||||
leaf module type (ex. torch.nn.Linear).
|
||||
|
||||
Returns:
|
||||
Dictionary mapping sources that were given to a list of SourcePartitions
|
||||
that correspond to the list of nodes that were decomposed from the given
|
||||
source.
|
||||
"""
|
||||
modules: dict[type, dict[str, list[Node]]] = {}
|
||||
|
||||
def add_to_partition(src: Any, fqn: str, node: Node) -> None:
|
||||
diff_modules = modules.setdefault(src, {})
|
||||
partition = diff_modules.setdefault(fqn, [])
|
||||
partition.append(node)
|
||||
|
||||
for node in graph.nodes:
|
||||
# The metadata source_fn should contain a tuple of a unique name for the
|
||||
# source, and the source function if the node is decomposed from a
|
||||
# function, or the type of module if the node is decomposed from a leaf
|
||||
# module
|
||||
|
||||
# TODO: Bypass "torch_fn" when "source_fn_stack" because now "torch_fn" can
|
||||
# be different from "source_fn_stack", for example for the add_ node
|
||||
# decomposed from batch norm. We should remove the check on "source_fn_stack"
|
||||
# after we fix "torch_fn". T199561090
|
||||
source_fn_st = node.meta.get("source_fn_stack", None)
|
||||
if source_fn_st is None:
|
||||
matched = False
|
||||
torch_fn = node.meta.get("torch_fn", None)
|
||||
if torch_fn is not None:
|
||||
node_fqn, source_fn = torch_fn
|
||||
source_fn_name = source_fn.split(".")[1]
|
||||
if source_fn_name in wanted_sources:
|
||||
add_to_partition(source_fn_name, node_fqn, node)
|
||||
matched = True
|
||||
# Fallback: when source_fn_stack is not populated (e.g. strict=False export),
|
||||
# use nn_module_stack to resolve the originating module type.
|
||||
# Only apply to call_function nodes to avoid incorrectly including
|
||||
# placeholder, get_attr, or output nodes in partitions.
|
||||
if not matched and node.op == "call_function":
|
||||
nn_module_stack = node.meta.get("nn_module_stack", None)
|
||||
if nn_module_stack:
|
||||
# Get the innermost module (last entry in the ordered dict)
|
||||
innermost_fqn, innermost_cls = list(nn_module_stack.values())[-1]
|
||||
for src in wanted_sources:
|
||||
if isinstance(src, type):
|
||||
if isinstance(innermost_cls, type) and issubclass(
|
||||
innermost_cls, src
|
||||
):
|
||||
add_to_partition(src, innermost_fqn, node)
|
||||
break
|
||||
elif isinstance(innermost_cls, str):
|
||||
src_str = src.__module__ + "." + src.__qualname__
|
||||
if innermost_cls == src_str:
|
||||
add_to_partition(src, innermost_fqn, node)
|
||||
break
|
||||
elif innermost_cls == src:
|
||||
add_to_partition(src, innermost_fqn, node)
|
||||
break
|
||||
|
||||
if source_fn_st is not None:
|
||||
source_fn = source_fn_st[-1]
|
||||
if source_fn[1] in wanted_sources:
|
||||
add_to_partition(source_fn[1], source_fn[0], node)
|
||||
|
||||
def make_partition(nodes: list[Node], module_type: type) -> SourcePartition:
|
||||
input_nodes = set()
|
||||
output_nodes = set()
|
||||
params = set()
|
||||
for node in nodes:
|
||||
for arg in node.args:
|
||||
if isinstance(arg, Node) and arg not in nodes and arg.op != "get_attr":
|
||||
input_nodes.add(arg)
|
||||
|
||||
if node.op == "get_attr":
|
||||
params.add(node)
|
||||
# get_attr nodes won't be output nodes
|
||||
continue
|
||||
|
||||
for user in node.users:
|
||||
if user not in nodes:
|
||||
output_nodes.add(node)
|
||||
|
||||
return SourcePartition(
|
||||
nodes,
|
||||
module_type,
|
||||
list(input_nodes),
|
||||
list(output_nodes),
|
||||
list(params), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
ret: dict[type[Any], list[SourcePartition]] = {}
|
||||
|
||||
if filter_fn:
|
||||
# for each partition, we apply filter_fn to filter out all partitions that doesn't satisfy the
|
||||
# filter condition
|
||||
filtered_modules = {}
|
||||
for tp, name_to_partition in modules.items():
|
||||
filtered_name_to_partition = {
|
||||
name: partition
|
||||
for name, partition in name_to_partition.items()
|
||||
if all(map(filter_fn, partition))
|
||||
}
|
||||
filtered_modules[tp] = filtered_name_to_partition
|
||||
modules = filtered_modules
|
||||
|
||||
for k, v in modules.items():
|
||||
ret[k] = [make_partition(partition, k) for partition in v.values()]
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False) # type: ignore[misc]
|
||||
def check_subgraphs_connected(
|
||||
subgraph1: SourcePartition, subgraph2: SourcePartition
|
||||
) -> bool:
|
||||
"""
|
||||
Given two subgraphs A and B (in the form of a list of nodes), checks if
|
||||
A has nodes connecting to at least one node in B -- aka there exists a node
|
||||
in B that uses a node in A (not the other way around).
|
||||
"""
|
||||
|
||||
for node in reversed(subgraph1.nodes):
|
||||
for user in node.users:
|
||||
if user in subgraph2.nodes:
|
||||
return True
|
||||
return False
|
||||
Reference in New Issue
Block a user