from typing import Any, Optional import networkx as nx from torch.fx import Graph, Node class GraphInfoProvider: """ This class provides information about the graph, such as the nodes, edges, and their runtime and memory requirements. It also provides methods to create graphs from the information provided. """ __RECOMPUTABLE_NODE_ONLY_GRAPH = "recomputable_node_only_graph" __RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT = ( "recomputable_node_only_graph_with_larger_graph_context" ) __FULL_NX_JOINT_GRAPH = "full_nx_joint_graph" __SIMPLIFIED_FX_JOINT_GRAPH = "fx_joint_graph" def __init__( self, graph_nodes_in_order: list[str], graph_edges: list[tuple[str, str]], all_recomputable_banned_nodes: list[str], all_node_runtimes: Optional[dict[str, float]] = None, all_node_memories: Optional[dict[str, float]] = None, recorded_knapsack_input_memories: Optional[list[float]] = None, recorded_knapsack_input_runtimes: Optional[list[float]] = None, joint_graph: Optional[Graph] = None, ): self.graph_nodes_in_order = graph_nodes_in_order self.graph_edges = graph_edges self.all_node_runtimes: dict[str, float] = dict() if all_node_runtimes is None: if recorded_knapsack_input_runtimes is None: raise ValueError( "Either all_node_runtimes or recorded_knapsack_input_runtimes must be provided." ) self.all_node_runtimes = { node: recorded_knapsack_input_runtimes[i] for i, node in enumerate(all_recomputable_banned_nodes) } else: self.all_node_runtimes.update(all_node_runtimes) self.all_node_memories: dict[str, float] = dict() if all_node_memories is None: if recorded_knapsack_input_memories is None: raise ValueError( "Either all_node_memories or recorded_knapsack_input_memories must be provided." ) self.all_node_memories = { node: recorded_knapsack_input_memories[i] for i, node in enumerate(all_recomputable_banned_nodes) } else: self.all_node_memories.update(all_node_memories) self.all_recomputable_banned_nodes = all_recomputable_banned_nodes self.all_recomputable_banned_nodes_set = set(all_recomputable_banned_nodes) self.recorded_knapsack_input_memories = recorded_knapsack_input_memories self.recorded_knapsack_input_runtimes = recorded_knapsack_input_runtimes self._lazily_initialized_graphs: dict[str, Any] = { self.__RECOMPUTABLE_NODE_ONLY_GRAPH: None, self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT: None, self.__FULL_NX_JOINT_GRAPH: None, self.__SIMPLIFIED_FX_JOINT_GRAPH: None, } @classmethod def inialize_from_graph( cls, joint_graph: Graph, all_recomputable_banned_nodes: list[Node], recorded_knapsack_input_memories: list[float], recorded_knapsack_input_runtimes: list[float], ) -> "GraphInfoProvider": """ Enables initialization from a joint graph. """ graph_nodes_in_order = [node.name for node in joint_graph.nodes] graph_edges = [ (node.name, user.name) for node in joint_graph.nodes for user in node.users ] all_recomputable_banned_node_names = [ node.name for node in all_recomputable_banned_nodes ] return cls( graph_nodes_in_order=graph_nodes_in_order, graph_edges=graph_edges, all_recomputable_banned_nodes=all_recomputable_banned_node_names, recorded_knapsack_input_memories=recorded_knapsack_input_memories, recorded_knapsack_input_runtimes=recorded_knapsack_input_runtimes, joint_graph=joint_graph, ) @property def recomputable_node_only_graph(self) -> nx.DiGraph: if self._lazily_initialized_graphs[self.__RECOMPUTABLE_NODE_ONLY_GRAPH] is None: self._lazily_initialized_graphs[ self.__RECOMPUTABLE_NODE_ONLY_GRAPH ] = self._create_recomputable_node_only_graph() return self._lazily_initialized_graphs[self.__RECOMPUTABLE_NODE_ONLY_GRAPH] @property def recomputable_node_only_graph_with_larger_graph_context(self) -> nx.DiGraph: if ( self._lazily_initialized_graphs[ self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT ] is None ): self._lazily_initialized_graphs[ self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT ] = self._create_recomputable_node_only_graph_with_larger_graph_context() return self._lazily_initialized_graphs[ self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT ] @property def full_joint_nx_graph(self) -> nx.DiGraph: if self._lazily_initialized_graphs[self.__FULL_NX_JOINT_GRAPH] is None: self._lazily_initialized_graphs[ self.__FULL_NX_JOINT_GRAPH ] = self._create_full_joint_graph() return self._lazily_initialized_graphs[self.__FULL_NX_JOINT_GRAPH] @property def simplified_fx_joint_graph(self) -> Graph: if self._lazily_initialized_graphs[self.__SIMPLIFIED_FX_JOINT_GRAPH] is None: self._lazily_initialized_graphs[ self.__SIMPLIFIED_FX_JOINT_GRAPH ] = self._recreate_psuedo_joint_graph() return self._lazily_initialized_graphs[self.__SIMPLIFIED_FX_JOINT_GRAPH] def get_non_ac_peak_memory(self) -> float: return sum( self.all_node_memories[node_name] for node_name in self.all_recomputable_banned_nodes_set ) def get_theoretical_max_runtime(self) -> float: return sum( self.all_node_runtimes[node_name] for node_name in self.all_recomputable_banned_nodes_set ) def get_knapsack_memory_input(self) -> list[float]: return ( self.recorded_knapsack_input_memories if self.recorded_knapsack_input_memories else [ self.all_node_memories[node_name] for node_name in self.all_recomputable_banned_nodes ] ) def get_knapsack_runtime_input(self) -> list[float]: return ( self.recorded_knapsack_input_runtimes if self.recorded_knapsack_input_runtimes else [ self.all_node_runtimes[node_name] for node_name in self.all_recomputable_banned_nodes ] ) def _create_recomputable_node_only_graph(self) -> nx.DiGraph: graph = nx.DiGraph() for recomputable_node in self.all_recomputable_banned_nodes: graph.add_node(recomputable_node) for a, b in self.graph_edges: if ( a in self.all_recomputable_banned_nodes_set and b in self.all_recomputable_banned_nodes_set ): graph.add_edge(a, b) return graph def _create_recomputable_node_only_graph_with_larger_graph_context( self, ) -> nx.DiGraph: # Create a dictionary to store the reachable nodes for each node all_recomputable_banned_nodes_set = set(self.all_recomputable_banned_nodes) reachable_nodes = {} for node in all_recomputable_banned_nodes_set: # Use BFS to find all reachable nodes predecessors = dict(nx.bfs_predecessors(self.full_joint_nx_graph, node)) reachable_recomputable_nodes = set(predecessors.keys()).intersection( all_recomputable_banned_nodes_set ) reachable_nodes[node] = reachable_recomputable_nodes # Create the candidate graph candidate_graph = nx.DiGraph() candidate_graph.add_nodes_from(all_recomputable_banned_nodes_set) for node1 in all_recomputable_banned_nodes_set: for node2 in reachable_nodes[node1]: # Check if there is an overlapping path overlapping_path = False for intermediate_node in reachable_nodes[node1]: if ( intermediate_node != node2 and node2 in reachable_nodes[intermediate_node] ): overlapping_path = True break if not overlapping_path: candidate_graph.add_edge(node1, node2) return candidate_graph def _create_full_joint_graph(self) -> nx.DiGraph: graph = nx.DiGraph() for node in self.graph_nodes_in_order: if node == "output": continue graph.add_node(node) for a, b in self.graph_edges: if a == "output" or b == "output": continue graph.add_edge(a, b) return graph def _recreate_psuedo_joint_graph(self) -> Graph: # Create a dictionary to store the dependencies of each node node_dependencies: dict[str, list[str]] = { node: [] for node in self.graph_nodes_in_order } for a, b in self.graph_edges: if a not in node_dependencies or b not in node_dependencies: raise ValueError(f"Edge ({a}, {b}) references a non-existent node.") node_dependencies[b].append(a) joint_graph = Graph() # Create nodes in the graph nodes: dict[str, Node] = {} for node_name in self.graph_nodes_in_order: input_nodes = [nodes[dep] for dep in node_dependencies[node_name]] if input_nodes: node = joint_graph.call_function(lambda *x: x, tuple(input_nodes)) node.name = node_name else: node = joint_graph.placeholder(node_name) nodes[node_name] = node return joint_graph def _visualize_recomputable_candidate_graph_with_larger_context( self, layout_k: float = 0.5, layout_iterations: int = 30, ) -> None: """ Visualize the recomputable candidate graph with larger context. """ from matplotlib import cm, colors as mcolors, pyplot as plt pos = nx.spring_layout( self.recomputable_node_only_graph_with_larger_graph_context, k=layout_k, iterations=layout_iterations, ) # pos = nx.spectral_layout(graph_with_indirect_edges) plt.figure(figsize=(20, 15)) # Create a dictionary for node labels using the index labels = { node: self.recomputable_node_only_graph_with_larger_graph_context.nodes[ node ].get("index", node) for node in self.recomputable_node_only_graph_with_larger_graph_context.nodes } # Extract memory values and normalize them norm = mcolors.Normalize( vmin=min(self.get_knapsack_memory_input()), vmax=max(self.get_knapsack_memory_input()), ) cmap = cm.viridis # type: ignore[attr-defined] # Assign colors based on memory node_colors = [ cmap( norm( float( self.recomputable_node_only_graph_with_larger_graph_context.nodes[ node ][ "memory" ] ) ) ) for node in self.recomputable_node_only_graph_with_larger_graph_context.nodes ] # Draw the graph with parsed nodes only nx.draw_networkx_nodes( self.recomputable_node_only_graph_with_larger_graph_context, pos, node_color=node_colors, node_size=300, label="Parsed Nodes", ) nx.draw_networkx_edges( self.recomputable_node_only_graph_with_larger_graph_context, pos, arrows=True, arrowsize=10, ) nx.draw_networkx_labels( self.recomputable_node_only_graph_with_larger_graph_context, pos, labels=labels, font_size=8, font_weight="bold", ) plt.title("Memory Colour Coded Dependency Graph for Recomputable Nodes") plt.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), label="Memory") plt.show()