Spaces:
Running
Running
| import torch | |
| from torch.fx.node import Node | |
| from torch.fx._symbolic_trace import symbolic_trace | |
| from torch.fx.passes.tools_common import legalize_graph | |
| import itertools | |
| import operator | |
| from typing import Dict, List, Tuple | |
| def split_result_tensors( | |
| result: torch.Tensor, inputs: List[torch.Tensor] | |
| ) -> Tuple[torch.Tensor, ...]: | |
| """ | |
| A free function for use in the merge_matmul graph transformation below that | |
| splits the output from a merged matmul into the individual results for each | |
| input tensor. | |
| Arguments: | |
| result: The merged matmul result tensor. | |
| inputs: The list of inputs that were merged into one for the matmul. | |
| Returns: | |
| List of matmul results for each input tensor. | |
| """ | |
| # When fx tracer is running, x.shape[0] will be torch.fx.Attribute but we | |
| # need an int even when tracing | |
| if isinstance(result, torch.fx.Proxy): | |
| splits = [0] * len(inputs) | |
| else: | |
| splits = [x.shape[0] for x in inputs] | |
| return torch.split(result, splits) | |
| def may_depend_on(a: Node, b: Node, search_depth: int = 6): | |
| """ | |
| Determine if one node depends on another in a torch.fx.Graph. | |
| Arguments: | |
| a: The node that may have a dependency on b. | |
| b: The node that a may have a dependency on. | |
| search_depth: In the case of an indirect dependency, this function | |
| searches upto this many nodes away in search of a | |
| data dependency. If none is found, the function | |
| makes the conservative assumption that there is a | |
| dependency. | |
| Returns: | |
| True if a may depend on b, False if it definitely does not. | |
| """ | |
| # Equivalence is defined as dependence. | |
| if a == b: | |
| return True | |
| # If a has no inputs, it cannot depend on b. | |
| if len(a.all_input_nodes) == 0: | |
| return False | |
| # If the search depth has been exhausted and no conclusion has been | |
| # reached, assume that there is a data dependency. | |
| if search_depth == 0: | |
| return True | |
| # Recursively check all inputs of a. | |
| for inp in a.all_input_nodes: | |
| if may_depend_on(inp, b, search_depth - 1): | |
| return True | |
| return False | |
| def are_nodes_independent(nodes: List[Node]): | |
| """ | |
| Check if all of the given nodes are pairwise-data independent. | |
| Arguments: | |
| nodes: The nodes to check for data dependencies. | |
| Returns: | |
| True if any pair in nodes has a data dependency. | |
| """ | |
| # For each pair in nodes: | |
| for i, j in itertools.combinations(nodes, 2): | |
| if may_depend_on(i, j) or may_depend_on(j, i): | |
| return False | |
| return True | |
| def merge_matmul(in_mod: torch.nn.Module): | |
| """ | |
| A graph transformation that merges matrix multiplication operations that share the same right-hand | |
| side operand into one large matrix multiplication. | |
| ____ _________ _________ | |
| ---- | | | | M| A * C | | |
| M| A | T| B | * K| C | = |---------| | |
| ---- , | | | | T| B * C | | |
| K ---- --------- --------- | |
| K R R | |
| """ | |
| gm = symbolic_trace(in_mod) | |
| rhs_users: Dict[Node, List[Node]] = {} | |
| lhs_users: Dict[Node, List[Node]] = {} | |
| # Populate rhs_users and lhs_users - maps from LHS/RHS matrix multiply operands to | |
| # the matmul of which they are the LHS/RHS. | |
| for node in gm.graph.nodes: | |
| if node.op != "call_function" or node.target is not torch.matmul: | |
| continue | |
| lhs, rhs = node.args | |
| # TODO: Properly handle aliasing caused by get_attr. For now, | |
| # use the attribute name as the operand if the node is a | |
| # get_attr. | |
| lhs = lhs.target if lhs.op == "get_attr" else lhs | |
| rhs = rhs.target if rhs.op == "get_attr" else rhs | |
| lhs_users.setdefault(lhs, []).append(node) | |
| rhs_users.setdefault(rhs, []).append(node) | |
| for rhs, mms in rhs_users.items(): | |
| # There must be at least matmuls for a merge to make sense. | |
| if len(mms) < 2: | |
| continue | |
| # All matmuls must not depend on each other directly or indirectly | |
| # in order for the merge to be possible. | |
| if not are_nodes_independent(mms): | |
| continue | |
| lhs_vals = [mm.args[0] for mm in mms] | |
| # Merge the matmul. | |
| # Collect a list of LHS operands and the single RHS operand. | |
| lhs = [gm.graph.get_attr(l) if isinstance(l, str) else l for l in lhs_vals] | |
| rhs = gm.graph.get_attr(rhs) if isinstance(rhs, str) else rhs | |
| # Concatenate all the LHS operands. | |
| merge_mm_cat = gm.graph.call_function(torch.cat, (lhs,), {}) | |
| # Multiply the concatenated LHS operands with the one RHS. This will produce | |
| # the same results as all the individual matmuls involving rhs in the original graph, | |
| # but they will all be concatenated together. | |
| merge_mm = gm.graph.call_function(torch.matmul, (merge_mm_cat, rhs,), {}) | |
| # Split the result of the merged matmul using the shapes of the LHS operands | |
| # to ascertain how large each chunk should be. | |
| merge_mm_split = gm.graph.call_function( | |
| split_result_tensors, (merge_mm, lhs), {} | |
| ) | |
| merge_mm_res = [ | |
| gm.graph.call_function(operator.getitem, (merge_mm_split, out), {}) | |
| for out in range(len(lhs)) | |
| ] | |
| # Replace all uses of the original, unmerged matmuls with the equivalent split chunk from the merged matmul. | |
| for old, new in zip(mms, merge_mm_res): | |
| old.replace_all_uses_with(new) | |
| gm.graph.erase_node(old) | |
| # All of the new nodes created above were inserted at the end, so we need to sort | |
| # the nodes topologically to make sure all definitions precede uses. | |
| legalize_graph(gm) | |
| gm.recompile() | |
| gm.graph.lint() | |
| return gm | |