[FX] Fix retracing wrapped functions (pytorch#58061)

Summary: Pull Request resolved: pytorch#58061

Test Plan: Imported from OSS

Reviewed By: yuhc

Differential Revision: D28358801

Pulled By: jamesr66a

fbshipit-source-id: c7c9a8a80e5bfe1eb1f6d2cf858ac7e57153a860

Author: James Reed

test/test_fx.py

@@ -327,6 +327,18 @@ def forward(self, x: torch.Tensor):
         ref_batchnorm1d = torch.nn.BatchNorm1d(2, affine=False)
         self.assertEqual(ref_batchnorm1d(input), m(input))
 
+    def test_wrapped_retrace(self):
+        def to_trace(y):
+            return wrapped_via_decorator(y)
+
+        m = symbolic_trace(to_trace)
+        self.assertIn('wrapped_via_decorator', m.code)
+        self.assertEqual(m(0), 1)
+
+        retraced = symbolic_trace(m)
+        self.assertIn('wrapped_via_decorator', retraced.code)
+        self.assertEqual(retraced(0), 1)
+
     def test_graph_edit_with_proxy(self):
         class M(torch.nn.Module):
             def forward(self, a, b):
@@ -2547,7 +2559,7 @@ def f(b, a):
 
 
 def run_getitem_target():
-    from torch.fx.symbolic_trace import _wrapped_methods_to_patch
+    from torch.fx._symbolic_trace import _wrapped_methods_to_patch
     _wrapped_methods_to_patch.append((torch.Tensor, "__getitem__"))
     try:
         TestFX().getitem_inner()

test/test_fx_experimental.py

@@ -5,7 +5,7 @@
 import math
 import numbers
 from typing import Callable, Dict, Union, List, Optional
-from torch.fx.symbolic_trace import symbolic_trace
+from torch.fx._symbolic_trace import symbolic_trace
 from torch.fx.graph_module import GraphModule
 from torch.fx.node import Node
 from torch.fx.experimental import graph_manipulation

torch/fx/__init__.py

@@ -82,7 +82,7 @@ def forward(self, x):
 '''
 
 from .graph_module import GraphModule
-from .symbolic_trace import symbolic_trace, Tracer, wrap, PH, ProxyableClassMeta
+from ._symbolic_trace import symbolic_trace, Tracer, wrap, PH, ProxyableClassMeta
 from .graph import Graph
 from .node import Node, map_arg
 from .proxy import Proxy

torch/fx/__init__.pyi

@@ -2,6 +2,6 @@ from .graph import Graph as Graph
 from .graph_module import GraphModule as GraphModule
 from .node import Node as Node, map_arg as map_arg
 from .proxy import Proxy as Proxy
-from .symbolic_trace import Tracer as Tracer, symbolic_trace as symbolic_trace, wrap as wrap
+from ._symbolic_trace import Tracer as Tracer, symbolic_trace as symbolic_trace, wrap as wrap
 from .interpreter import Interpreter as Interpreter, Transformer as Transformer
 from .subgraph_rewriter import replace_pattern as replace_pattern

torch/fx/symbolic_trace.py renamed to torch/fx/_symbolic_trace.py

@@ -620,7 +620,9 @@ def wrapped(*args, **kwargs):
         """
         proxy = _find_proxy(args, kwargs)
         if proxy is not None:
-            return proxy.tracer.create_proxy('call_function', orig_fn, args, kwargs)
+            return_proxy = proxy.tracer.create_proxy('call_function', orig_fn, args, kwargs)
+            return_proxy.node.meta['is_wrapped'] = True
+            return return_proxy
         return orig_fn(*args, **kwargs)
 
     return wrapped

torch/fx/experimental/merge_matmul.py

@@ -3,7 +3,7 @@
 from torch.fx.graph import Graph
 from torch.fx.graph_module import GraphModule
 from torch.fx.node import Node
-from torch.fx.symbolic_trace import symbolic_trace
+from torch.fx._symbolic_trace import symbolic_trace
 
 import itertools
 import operator

torch/fx/experimental/rewriter.py

@@ -4,7 +4,7 @@
 import copy
 from types import FunctionType
 from typing import cast, Union, Callable, Dict, Optional, Any
-from torch.fx.symbolic_trace import Tracer
+from torch.fx._symbolic_trace import Tracer
 from torch.fx.graph import Graph
 from torch.jit.frontend import normalize_source_lines
 import torch

torch/fx/graph.py

@@ -15,7 +15,7 @@
 
 
 if TYPE_CHECKING:
-    from .graph_module import GraphModule
+    from .graph_module import GraphModule  # noqa: F401
 
 
 # Mapping of builtins to their `typing` equivalent.
@@ -805,6 +805,7 @@ def _python_code(self, root_module: str, namespace: _Namespace) -> PythonCode:
         free_vars: List[str] = []
         body: List[str] = []
         globals_: Dict[str, Any] = {}
+        wrapped_fns: Dict[str, None] = {}
 
         # Wrap string in list to pass by reference
         maybe_return_annotation : List[str] = ['']
@@ -919,6 +920,8 @@ def emit_node(node : Node):
                     body.append(f'{repr(node)}{maybe_type_annotation} = {_format_target(repr(node.args[0]), node.args[1])}')
                     return
                 body.append(f'{repr(node)}{maybe_type_annotation} = {global_name}({_format_args(node.args, node.kwargs)})')
+                if node.meta.get('is_wrapped', False):
+                    wrapped_fns.setdefault(global_name)
                 return
             elif node.op == 'call_module':
                 assert isinstance(node.target, str)
@@ -960,18 +963,24 @@ def emit_node(node : Node):
         else:
             orig_args = free_vars
 
+        if len(wrapped_fns) > 0:
+            wrap_name = add_global('wrap', torch.fx.wrap)
+            wrap_stmts = '\n'.join([f'{wrap_name}("{name}")' for name in wrapped_fns])
+        else:
+            wrap_stmts = ''
+
         # If the original function didn't have self as its first argument, we
         # would have added it.
         if len(orig_args) == 0 or orig_args[0] != 'self':
             orig_args.insert(0, 'self')
         code = ''.join(body)
         code = '\n'.join('    ' + line for line in code.split('\n'))
         fn_code = f"""
+{wrap_stmts}
+
 def forward({', '.join(orig_args)}){maybe_return_annotation[0]}:
 {code}"""
-
-        return PythonCode(fn_code,
-                          globals_)
+        return PythonCode(fn_code, globals_)
 
     def __str__(self) -> str:
         """

torch/fx/graph_module.py

@@ -96,7 +96,7 @@ def __init__(self, body):
     # Try to retrieve the forward source in a backward-compatible way
     CodeOnlyModule.forward = forward
 
-    from .symbolic_trace import Tracer
+    from ._symbolic_trace import Tracer
 
     # we shouldn't trace into any of the submodules, they were not
     # because they were not traced in the original GraphModule

torch/fx/interpreter.py

@@ -2,7 +2,7 @@
 from .graph import Graph
 from .node import Argument, Node, Target, map_arg, map_aggregate
 from .proxy import Proxy
-from .symbolic_trace import Tracer
+from ._symbolic_trace import Tracer
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 class Interpreter:

torch/fx/subgraph_rewriter.py

@@ -1,7 +1,7 @@
 from .graph_module import GraphModule
 from .graph import Graph
 from .node import Node
-from .symbolic_trace import symbolic_trace
+from ._symbolic_trace import symbolic_trace
 
 import copy
 from typing import Callable, Dict, List, NamedTuple, Optional, Set

torch/quantization/quantize_fx.py

@@ -1,6 +1,6 @@
 import torch
 from torch.fx import GraphModule
-from torch.fx.symbolic_trace import Tracer
+from torch.fx._symbolic_trace import Tracer
 from torch.fx.node import Target, Node, Argument
 from .fx import Fuser  # noqa: F401
 from .fx import Quantizer  # noqa: F401