[fx] Have replace_pattern return replaced nodes (pytorch#90244)

Summary: Modified replace_pattern in the subgraph rewriter to return a list of pairs of matches along with their corresponding replacement nodes in the modified graph (`List[Tuple[Match, List[Node]]]`). This allows us to easily modify the replaced nodes, including setting the metadata.

Test Plan: CI

Differential Revision: D41737056

Pull Request resolved: pytorch#90244
Approved by: https://github.com/SherlockNoMad

Author: angelayi

test/fx/test_subgraph_rewriter.py

@@ -819,17 +819,22 @@ def second_input_is_scalar(match, original_graph, pattern_graph):
                 return False
             return True
 
-        def num_repalcement_node_found(traced):
-            return sum(1 for node in traced.graph.nodes if node.target == torch.mul)
+        def check_replacement_nodes(self, traced, matches):
+            replacement_nodes_in_graph = [node for node in traced.graph.nodes if node.target == torch.mul]
+            replacement_nodes_in_res = [r for m in matches for r in m.replacements]
+            self.assertEqual(len(replacement_nodes_in_graph), len(replacement_nodes_in_res))
+            self.assertEqual(replacement_nodes_in_graph, replacement_nodes_in_res)
+            return len(replacement_nodes_in_graph)
 
         # match without filter, should find 2 match
         traced = symbolic_trace(M())
-        matches = subgraph_rewriter.replace_pattern(
+        matches = subgraph_rewriter.replace_pattern_with_filters(
             traced,
             BinaryOpScalarReLUPattern,
-            BinaryOpScalarReLUReplacement)
+            BinaryOpScalarReLUReplacement,
+            None)
         self.assertEqual(len(matches), 2)
-        self.assertEqual(num_repalcement_node_found(traced), 2)
+        self.assertEqual(check_replacement_nodes(self, traced, matches), 2)
 
         # match with filter, should find 1 match
         traced = symbolic_trace(M())
@@ -839,7 +844,7 @@ def num_repalcement_node_found(traced):
             BinaryOpScalarReLUReplacement,
             [second_input_is_scalar])
         self.assertEqual(len(matches), 1)
-        self.assertEqual(num_repalcement_node_found(traced), 1)
+        self.assertEqual(check_replacement_nodes(self, traced, matches), 1)
 
     def test_matching_pattern_with_list_type_arg(self):
         class M(torch.nn.Module):

torch/fx/subgraph_rewriter.py

@@ -5,10 +5,11 @@
 from ._compatibility import compatibility
 
 import copy
+from dataclasses import dataclass
 from typing import Callable, Dict, List, NamedTuple, Optional, Set, Union
 import torch
 
-__all__ = ['Match', 'replace_pattern', 'replace_pattern_with_filters']
+__all__ = ['Match', 'replace_pattern', 'replace_pattern_with_filters', "ReplacedPatterns"]
 
 @compatibility(is_backward_compatible=True)
 class Match(NamedTuple):
@@ -17,6 +18,15 @@ class Match(NamedTuple):
     # Maps nodes in the pattern subgraph to nodes in the larger graph
     nodes_map: Dict[Node, Node]
 
+@compatibility(is_backward_compatible=False)
+@dataclass
+class ReplacedPatterns:
+    # Node from which the match was found
+    anchor: Node
+    # Maps nodes in the pattern subgraph to nodes in the larger graph
+    nodes_map: Dict[Node, Node]
+    # List of nodes that were added into the graph
+    replacements: List[Node]
 
 def _replace_submodules(gm: GraphModule, replacement: torch.nn.Module) -> None:
     gm.delete_all_unused_submodules()
@@ -183,7 +193,8 @@ def forward(self, x, w1, w2):
             add_2 = add_1 + max_2
             return add_2
     """
-    return _replace_pattern(gm, pattern, replacement)
+    match_and_replacements = _replace_pattern(gm, pattern, replacement)
+    return [Match(anchor=m.anchor, nodes_map=m.nodes_map) for m in match_and_replacements]
 
 
 # Experimental API, not backward compatible
@@ -193,7 +204,7 @@ def replace_pattern_with_filters(
     pattern: Union[Callable, GraphModule],
     replacement: Union[Callable, GraphModule],
     match_filters: List[Callable[["InternalMatch", Graph, Graph], bool]],  # type: ignore[name-defined]
-) -> List[Match]:
+) -> List[ReplacedPatterns]:
     """
     See replace_pattern for documentation. This function is an overload with an additional match_filter argument.
 
@@ -211,8 +222,8 @@ def _replace_pattern(
     gm: GraphModule,
     pattern: Union[Callable, GraphModule],
     replacement: Union[Callable, GraphModule],
-    match_filters: List[Callable[["InternalMatch", Graph, Graph], bool]] = None  # type: ignore[name-defined]
-) -> List[Match]:
+    match_filters: List[Callable[["InternalMatch", Graph, Graph], bool]] = None,  # type: ignore[name-defined]
+) -> List[ReplacedPatterns]:
 
     from torch.fx.passes.utils.matcher_utils import SubgraphMatcher, InternalMatch
 
@@ -248,6 +259,7 @@ def _replace_pattern(
     # As we progressively replace nodes, we'll need to keep track of how the match results should change
     match_changed_node: Dict[Node, Node] = {}
 
+    match_and_replacements = []
     for match in _matches:
 
         # Build connecting between replacement graph's input and original graph input producer node
@@ -285,6 +297,20 @@ def _replace_pattern(
         if isinstance(copied_returning_nodes, Node):
             copied_returning_nodes = (copied_returning_nodes, )
 
+        # Get a list of nodes that have been replaced into the graph
+        replacement_nodes = []
+
+        def get_replacement_nodes(curr_node: Node):
+            nonlocal replacement_nodes
+            for arg in curr_node.args:
+                if isinstance(arg, Node):
+                    if arg not in val_map.values():
+                        get_replacement_nodes(arg)
+            replacement_nodes.append(curr_node)
+
+        for ret_node in copied_returning_nodes:
+            get_replacement_nodes(ret_node)
+
         # Hook the output Node of the replacement subgraph in to the
         # original Graph at the correct location
         assert len(match.returning_nodes) == len(copied_returning_nodes)
@@ -297,6 +323,14 @@ def _replace_pattern(
                 gn = match.nodes_map[node]
                 gm.graph.erase_node(gn)
 
+        match_and_replacements.append(
+            ReplacedPatterns(
+                anchor=match.anchors[0],
+                nodes_map=match.nodes_map,
+                replacements=replacement_nodes
+            )
+        )
+
     # Update the passed-in GraphModule to reflect the new state of
     # `original_graph`
     gm.recompile()
@@ -306,6 +340,4 @@ def _replace_pattern(
     if isinstance(replacement, torch.nn.Module):
         _replace_submodules(gm, replacement)
 
-    # Convert _matches: InternalMatch to Match to comply with backward compatibility of this function
-    matches: List[Match] = [Match(anchor=match.anchors[0], nodes_map=match.nodes_map) for match in _matches]
-    return matches
+    return match_and_replacements