[fx] Add Graph option for replace_pattern

dulinriley · facebook-github-bot · commit 43149f37133b · 2023-10-30T09:59:42.000-07:00
Summary:
Allow doing pattern replacement with just an fx.Graph instead of a fx.GraphModule,
which can let callers avoid paying the cost of `recompile()` for a small graph if they
don't need the module.

This is a significant speedup if you use hundreds of small patterns for replacement.

Test Plan: Tested in a diff stacked on top of this: {D50756722}

Reviewed By: SherlockNoMad, angelayi

Differential Revision: D50756723
diff --git a/torch/fx/subgraph_rewriter.py b/torch/fx/subgraph_rewriter.py
@@ -200,8 +200,8 @@ def forward(self, x, w1, w2):
 @compatibility(is_backward_compatible=False)
 def replace_pattern_with_filters(
     gm: GraphModule,
-    pattern: Union[Callable, GraphModule],
-    replacement: Union[Callable, GraphModule],
+    pattern: Union[Callable, Graph, GraphModule],
+    replacement: Union[Callable, Graph, GraphModule],
     match_filters: Optional[List[Callable[["InternalMatch", Graph, Graph], bool]]] = None,  # type: ignore[name-defined]
     ignore_literals: bool = False,
 ) -> List[ReplacedPatterns]:
@@ -220,8 +220,8 @@ def replace_pattern_with_filters(
 
 def _replace_pattern(
     gm: GraphModule,
-    pattern: Union[Callable, GraphModule],
-    replacement: Union[Callable, GraphModule],
+    pattern: Union[Callable, Graph, GraphModule],
+    replacement: Union[Callable, Graph, GraphModule],
     match_filters: Optional[List[Callable[["InternalMatch", Graph, Graph], bool]]] = None,  # type: ignore[name-defined]
     ignore_literals: bool = False,
 ) -> List[ReplacedPatterns]:
@@ -236,11 +236,15 @@ def _replace_pattern(
 
     if isinstance(pattern, GraphModule):
         pattern_graph = pattern.graph
+    elif isinstance(pattern, Graph):
+        pattern_graph = 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
 
