parameters: (in_mod: torch.nn.Module)
return_statement: return gm
docstring: 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
docstring_summary: 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
docstring_tokens: 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
function: 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_sizes = [
              gm.graph.call_function(get_first_dim, (l,), {}) for l in lhs
          ]
          merge_mm_split = gm.graph.call_function(
              torch.split, (merge_mm, merge_mm_sizes), {}
          )
          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
function_tokens: def merge_matmul ( in_mod : torch . nn . Module ) : 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_sizes = [ gm . graph . call_function ( get_first_dim , ( l , ) , { } ) for l in lhs ] merge_mm_split = gm . graph . call_function ( torch . split , ( merge_mm , merge_mm_sizes ) , { } ) 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