[torch][fx] Add new replacement_callback to materialize a replacement just in time (pytorch#135553)

Summary:
Sometimes we only want to generate a replacement for a matched pattern
once we know some information about the nodes in the pattern.

So far, we have found this the most useful to do matches based on specific
shapes of tensors flowing into functions.
Use a callback function similar to `match_filters`. By default this isn't used.

Had to make `replacement` a None-able parameter because Callable was
already used to detect a case where a graph needed to be traced.

Differential Revision: D62412628

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

Author: dulinriley

test/fx/test_subgraph_rewriter.py

@@ -980,3 +980,38 @@ def check_replacement_nodes(self, traced, matches):
             return len(replacement_nodes_in_graph)
 
         self.assertEqual(check_replacement_nodes(self, traced, matches), 2)
+
+    def test_replace_pattern_with_callback(self) -> None:
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.add(x, y)
+
+        def pattern(x, y):
+            return torch.add(x, y)
+
+        def replacement(x, y):
+            return torch.sub(torch.mul(x, y), y)
+
+        traced = symbolic_trace(M())
+        # Return the same replacement graph for all matches, but have it be a unique
+        # object each time.
+        matches = subgraph_rewriter.replace_pattern_with_filters(
+            traced,
+            pattern,
+            replacement_callback=lambda *args: symbolic_trace(replacement).graph,
+        )
+
+        def check_replacement_nodes(self, traced, matches):
+            replacement_nodes_in_graph = [
+                node
+                for node in traced.graph.nodes
+                if node.target in {torch.sub, 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)
+
+        self.assertEqual(check_replacement_nodes(self, traced, matches), 2)

torch/fx/subgraph_rewriter.py

@@ -207,9 +207,11 @@ def forward(self, x, w1, w2):
 def replace_pattern_with_filters(
     gm: GraphModule,
     pattern: Union[Callable, Graph, GraphModule],
-    replacement: Union[Callable, Graph, GraphModule],
+    replacement: Union[Callable, Graph, GraphModule, None] = None,
     match_filters: Optional[List[Callable[["InternalMatch", Graph, Graph], bool]]] = None,
     ignore_literals: bool = False,
+    # Placed at the end to avoid breaking backward compatibility
+    replacement_callback: Optional[Callable[["InternalMatch", Graph, Graph], Graph]] = None,
 ) -> List[ReplacedPatterns]:
     """
     See replace_pattern for documentation. This function is an overload with an additional match_filter argument.
@@ -219,17 +221,22 @@ def replace_pattern_with_filters(
             (match: InternalMatch, original_graph: Graph, pattern_graph: Graph) and return a boolean indicating
             whether the match satisfies the condition.
             See matcher_utils.py for definition of InternalMatch.
+        ``replacement_callback``: A function that takes in a match and returns a
+            Graph to be used as the replacement. This allows you to construct a
+            replacement graph based on the match.
     """
 
-    return _replace_pattern(gm, pattern, replacement, match_filters, ignore_literals)
+    return _replace_pattern(gm, pattern, replacement, match_filters, ignore_literals, replacement_callback)
 
 
 def _replace_pattern(
     gm: GraphModule,
     pattern: Union[Callable, Graph, GraphModule],
-    replacement: Union[Callable, Graph, GraphModule],
+    replacement: Union[Callable, Graph, GraphModule, None] = None,
     match_filters: Optional[List[Callable[["InternalMatch", Graph, Graph], bool]]] = None,
     ignore_literals: bool = False,
+    # Placed at the end to avoid breaking backward compatibility
+    replacement_callback: Optional[Callable[["InternalMatch", Graph, Graph], Graph]] = None,
 ) -> List[ReplacedPatterns]:
 
     from torch.fx.passes.utils.matcher_utils import SubgraphMatcher, InternalMatch
@@ -247,13 +254,6 @@ def _replace_pattern(
     else:
         pattern_graph = symbolic_trace(pattern).graph
 
-    if isinstance(replacement, GraphModule):
-        replacement_graph = replacement.graph
-    elif isinstance(replacement, Graph):
-        replacement_graph = replacement
-    else:
-        replacement_graph = symbolic_trace(replacement).graph
-
     matcher = SubgraphMatcher(pattern_graph, match_output=False, match_placeholder=False,
                               remove_overlapping_matches=True, ignore_literals=ignore_literals)
     _matches: List[InternalMatch] = matcher.match(original_graph)
@@ -265,13 +265,27 @@ def _replace_pattern(
                for match_filter in match_filters)
     ]
 
-    replacement_placeholders = [n for n in replacement_graph.nodes if n.op == "placeholder"]
+    if isinstance(replacement, GraphModule):
+        common_replacement_graph = replacement.graph
+    elif isinstance(replacement, Graph):
+        common_replacement_graph = replacement
+    elif callable(replacement):
+        common_replacement_graph = symbolic_trace(replacement).graph
+    else:
+        assert replacement_callback is not None, "Must provide either a replacement GraphModule or a replacement callback"
+        common_replacement_graph = None
 
     # 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:
+    for i, match in enumerate(_matches):
+        if replacement_callback is not None:
+            replacement_graph = replacement_callback(match, original_graph, pattern_graph)
+        else:
+            assert common_replacement_graph is not None, "Must provide either a replacement GraphModule or a replacement callback"
+            replacement_graph = common_replacement_graph
+        replacement_placeholders = [n for n in replacement_graph.nodes if n.op == "placeholder"]
 
         # Build connecting between replacement graph's input and original graph input producer node