from collections import defaultdict from typing import Optional import torch import torch.fx as fx from torch.utils._ordered_set import OrderedSet class AugmentedGraphHelper: """ Graph helper that augments the original graph with additional dependencies and uses, plus tracks node equivalences for coalescing. TODO: if this becomes too large of compile time, consider binding graphcycles.cc """ def __init__( self, graph: fx.Graph, node_ancestors: Optional[dict[fx.Node, OrderedSet[fx.Node]]] = None, ): # Each node starts in its own singleton set self.graph = graph self.merge_sets = {node: OrderedSet([node]) for node in graph.nodes} # Extra dependencies: node depends on dep (dep must come before node) self.extra_deps: dict[fx.Node, OrderedSet[fx.Node]] = defaultdict(OrderedSet) # Extra uses: reverse of extra_deps (node is used by user) self.extra_uses: dict[fx.Node, OrderedSet[fx.Node]] = defaultdict(OrderedSet) # Note: only reflect original ancestors, not maintained through additional deps # or merge sets self.node_ancestors = node_ancestors def add_extra_dep(self, *, n: fx.Node, dep: fx.Node) -> None: """Add extra dependency: node depends on dep.""" self.extra_deps[n].add(dep) self.extra_uses[dep].add(n) def remove_extra_dep(self, *, n: fx.Node, dep: fx.Node) -> None: if dep in self.extra_deps[n]: self.extra_deps[n].discard(dep) self.extra_uses[dep].discard(n) def merge_to_set(self, existing_node: fx.Node, new_node: fx.Node) -> None: """ Merge new_node into existing_node's set. The new node must be a singleton set. """ existing_set = self.merge_sets[existing_node] new_set = self.merge_sets[new_node] assert len(new_set) == 1 # Add all nodes from new_set to existing_set existing_set.update(new_set) # Update all nodes from new_set to point to existing_set for node in new_set: self.merge_sets[node] = existing_set def unmerge_node(self, node: fx.Node) -> None: """Remove a node from its merge set, making it singleton.""" old_set = self.merge_sets[node] # If already singleton, nothing to do if len(old_set) == 1: return # Remove from old set old_set.remove(node) # Make node singleton self.merge_sets[node] = OrderedSet([node]) def get_merged_deps(self, node: fx.Node) -> OrderedSet[fx.Node]: """ Get all dependencies of a node considering merges and extra deps. Combines: 1. Direct deps (all_input_nodes) of node and its merge equivalents 2. Extra deps of node and its merge equivalents """ deps: OrderedSet[fx.Node] = OrderedSet() # For each node in the merge set for merged_node in self.merge_sets[node]: # Add direct dependencies from all_input_nodes deps.update(merged_node.all_input_nodes) # Add extra dependencies deps.update(self.extra_deps[merged_node]) return deps def has_cycle(self) -> bool: merged_deps = {n: self.get_merged_deps(n) for n in self.graph.nodes} return torch._dynamo.graph_deduplication._has_cycle(self.graph, merged_deps) def has_path(self, source: fx.Node, target: fx.Node) -> bool: """Check if there's a path from source to target.""" # we should not be checking path from node to itself assert self.merge_sets[source] is not self.merge_sets[target] # search backwards from target to source visited: OrderedSet[fx.Node] = OrderedSet() queue = [target] visited.add(target) while queue: current = queue.pop() for dep in self.get_merged_deps(current): # Check if we reached source or its equivalent if dep in self.merge_sets[source]: return True if dep in visited: continue # We are searching from target, so this node is necessarily an ancestor # of target. # If dep is an ancestor of source, any path through dep to source would imply a cycle if self.node_ancestors: source_set = self.merge_sets[source] is_ancestor_of_source = any( dep in self.node_ancestors[s] for s in source_set ) # Add to visited to avoid recomputing this check if we see dep again if is_ancestor_of_source: visited.add(dep) continue visited.add(dep) queue.append(dep) return False def transfer_erased_node_deps( self, erased_to_new: dict[fx.Node, fx.Node | None] ) -> None: """ Transfer all extra dependencies from erased nodes to their replacements, handling cross-dependencies between erased nodes correctly. """ erased_merge_sets: dict[fx.Node, fx.Node | None] = {} for replaced, new in erased_to_new.items(): for equiv in self.merge_sets[replaced]: erased_merge_sets[equiv] = new # Transfer dependencies for old_node, new_node in erased_merge_sets.items(): if new_node is None: # Clean up references to removed node for extra_use in list(self.extra_uses[old_node]): updated_use = erased_merge_sets.get(extra_use, extra_use) if updated_use is not None: self.extra_deps[updated_use].discard(old_node) for extra_dep in list(self.extra_deps[old_node]): updated_dep = erased_merge_sets.get(extra_dep, extra_dep) if updated_dep is not None: self.extra_uses[updated_dep].discard(old_node) else: # Transfer dependencies FROM old_node (what old_node depended on) for extra_dep in self.extra_deps[old_node]: updated_dep = erased_merge_sets.get(extra_dep, extra_dep) if updated_dep is not None and updated_dep != new_node: self.extra_deps[new_node].add(updated_dep) self.extra_uses[updated_dep].discard(old_node) self.extra_uses[updated_dep].add(new_node) # Transfer dependencies TO old_node (what depended on old_node) for extra_use in self.extra_uses[old_node]: updated_use = erased_merge_sets.get(extra_use, extra_use) if updated_use is not None and updated_use != new_node: self.extra_deps[updated_use].discard(old_node) self.extra_deps[updated_use].add(new_node) self.extra_uses[new_node].add(updated_use) # Clean up erased nodes for old_node in erased_merge_sets: self.extra_deps[old_node].clear() self.extra_uses[old_node].clear() del self.merge_sets[old_node] def get_all_extra_deps(self) -> dict[fx.Node, OrderedSet[fx.Node]]: """ Get all extra dependencies in a format suitable for topological sort. Returns a copy to avoid external modifications. """ return { node: OrderedSet(deps) for node, deps in self.extra_deps.items() if deps # Only include nodes with non-empty deps }