Revert "[Dynamo] Trace enter/exit of TorchFunctionModes (#135422) (#137114)"

pytorchmergebot · pytorchmergebot · commit d34b617bb98c · 2024-10-08T20:33:17.000Z
This reverts commit 51bc839. Reverted #137114 on behalf of https://github.com/huydhn due to The top of the stack has been reverted but it leaves trunk in a broken state, so I try to revert the rest of the stack ([comment](#137114 (comment)))
diff --git a/test/dynamo/test_modes.py b/test/dynamo/test_modes.py
@@ -461,94 +461,6 @@ def fn(x):
 
         self.assertEqual(expected, actual)
 
-    def test_torch_function_mode_enter_exit(self):
-        def fn(x, y):
-            with TestMode():
-                o = torch.add(x, 3)
-
-            return torch.add(o, y)
-
-        inp = (torch.ones(2, 2) + 1, torch.ones(2, 2) + 2)
-        fn_opt = torch.compile(fn, fullgraph=True)
-
-        expected = fn(*inp)
-        actual = fn_opt(*inp)
-
-        self.assertEqual(expected, actual)
-
-    def test_torch_function_mode_graph_break(self):
-        def fn(x, y):
-            with TestMode():
-                torch._dynamo.graph_break()
-                o = torch.add(x, 3)
-
-            return torch.add(o, y)
-
-        inp = (torch.ones(2, 2) + 1, torch.ones(2, 2) + 2)
-        fn_opt = torch.compile(fn)
-
-        expected = fn(*inp)
-        actual = fn_opt(*inp)
-
-        self.assertEqual(expected, actual)
-
-    def test_torch_function_mode_and_pop_graph_break(self):
-        def fn(x, y):
-            with TestMode():
-                z = _pop_torch_function_stack()
-                torch._dynamo.graph_break()
-                _push_on_torch_function_stack(z)
-                o = torch.add(x, 3)
-
-            return torch.add(o, y)
-
-        inp = (torch.ones(2, 2) + 1, torch.ones(2, 2) + 2)
-        fn_opt = torch.compile(fn)
-
-        expected = fn(*inp)
-        actual = fn_opt(*inp)
-
-        self.assertEqual(expected, actual)
-
-    def test_torch_function_mode_restore_on_exc(self):
-        @torch._dynamo.disable()
-        def err():
-            raise RuntimeError("test")
-
-        @torch.compile()
-        def fn(x):
-            with TestMode():
-                x += 1
-                err()
-                x += 2
-                return x
-
-        try:
-            fn(torch.ones(2, 2))
-        except RuntimeError:
-            pass
-        self.assertEqual(_len_torch_function_stack(), 0)
-
-    def test_torch_function_mode_and_pop_graph_break_mutation(self):
-        def fn(x, y):
-            with TestMode():
-                z = _pop_torch_function_stack()
-                z.y = 5
-                torch._dynamo.graph_break()
-                _push_on_torch_function_stack(z)
-                o = torch.add(x, 3)
-                o = torch.mul(o, z.y)
-
-            return torch.add(o, y)
-
-        inp = (torch.ones(2, 2) + 1, torch.ones(2, 2) + 2)
-        fn_opt = torch.compile(fn)
-
-        expected = fn(*inp)
-        actual = fn_opt(*inp)
-
-        self.assertEqual(expected, actual)
-
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/torch/_dynamo/convert_frame.py b/torch/_dynamo/convert_frame.py
@@ -120,7 +120,6 @@
     troubleshooting_url,
     write_record_to_file,
 )
-from .variables.torch_function import torch_function_mode_stack_state_mgr
 
 
 np: Optional[ModuleType]
@@ -219,18 +218,15 @@ def _fn(*args: _P.args, **kwargs: _P.kwargs) -> _T:
             prior_fwd_from_src = torch.fx.graph_module._forward_from_src
             torch.fx.graph_module._forward_from_src = fx_forward_from_src_skip_result
             cleanup = setup_compile_debug()
+
             exit_stack = contextlib.ExitStack()
             exit_stack.enter_context(
                 torch.fx._symbolic_trace._maybe_revert_all_patches()
             )
-            exit_stack.enter_context(torch_function_mode_stack_state_mgr)
             try:
                 return fn(*args, **kwargs)
             finally:
                 cleanup.close()
-                assert (
-                    torch._C._len_torch_function_stack() == 0
-                ), "Torch function mode stack state changed while dynamo tracing, please report a bug"
                 exit_stack.close()
                 torch._C._set_grad_enabled(prior_grad_mode)
                 torch.autograd.grad_mode._enter_inference_mode(prior_inference_mode)
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
@@ -78,6 +78,7 @@
     get_instruction_source_311,
     get_locals_to_steal,
     get_static_address_type,
+    get_torch_function_mode_stack,
     graph_break_reasons,
     increment_op_count,
     lazy_format_graph_code,
@@ -249,7 +250,6 @@ def __init__(
         local_scope: Scope,
         global_scope: Scope,
         f_code,
-        torch_function_mode_stack,
     ):
         super().__init__()
         self.tracers = [SubgraphTracer(self, export_root=export)]
@@ -368,7 +368,7 @@ def __init__(
         # This returns false if TF Overall (both mode and subclass) is disabled OR that TF Mode stack is empty
         self.torch_function_mode_enabled = torch._C._is_torch_function_mode_enabled()
         # This records the initial torch function mode stack for guarding
-        self.torch_function_mode_stack = torch_function_mode_stack
+        self.torch_function_mode_stack = get_torch_function_mode_stack()
 
         # Tracks if the output graph has a user defined allowed function in the
         # graph. This is used later to determine if we should fallback to eager
@@ -1021,7 +1021,7 @@ def append_prefix_insts():
             prefix_insts.clear()
 
         for block in reversed(tx.block_stack):
-            block.exit(tx, is_graph_break=reason.graph_break)
+            block.exit(tx)
 
         self.cleanup_graph()
         tx.prune_dead_locals()
diff --git a/torch/_dynamo/polyfills/__init__.py b/torch/_dynamo/polyfills/__init__.py
@@ -25,26 +25,6 @@
         sys as sys,
     )
 
-from torch.overrides import BaseTorchFunctionMode
-
-
-# These classes handle support for TorchFunctionModes across
-# graph breaks
-# Today the TorchFunctionMode enter (for the classes we support)
-# simply pushes the mode onto the stack. Since after this occurs
-# the stack is mutated, and we replay these mutations, we don't need
-# any cleanup logic to be run once the graph break occurs, we simply replay
-# these mutations to ensure at the graph break the torch function mode stack is correct
-#  and reconstruct the torch function mode stack normally
-# when we compile the resume function on the other side of the break.
-# However, to ensure we exit properly
-# in the resume function, we need to re-enter the contexts as we do other contexts.
-# These contexts do nothing on enter, but provide the correct exit logic to ensure
-# the stack state is correct.
-class NoEnterTorchFunctionMode(BaseTorchFunctionMode):
-    def __enter__(self):
-        pass
-
 
 def index(iterator, item, start=0, end=None):
     from itertools import islice
diff --git a/torch/_dynamo/resume_execution.py b/torch/_dynamo/resume_execution.py
@@ -90,25 +90,27 @@ class ReenterWith:
     stack_index: int
     target_values: Optional[Tuple[Any, ...]] = None
 
-    def try_except_torch_function_mode(self, code_options, cleanup: List[Instruction]):
-        """
-        Codegen based off of:
-        try:
-            (rest)
-        except:
-            (restore previous tf mode stack)
-            raise
-        """
-        from .variables.torch_function import get_prev_stack_var_name
-
-        setup_try_except, epilogue = _bytecode_from_template_with_split(
-            _try_except_tf_mode_template,
-            self.stack_index,
-            varname_map={"stack_var_name": get_prev_stack_var_name()},
-        )
-        cleanup[:] = epilogue + cleanup
-
-        return setup_try_except
+    # TODO(mlazos) - Uncomment with the reland of torch function mode support
+    # def try_except_torch_function_mode(self, code_options, cleanup: List[Instruction]):
+    #     """
+    #     Codegen based off of:
+    #     try:
+    #         (rest)
+    #     except:
+    #         (restore previous tf mode stack)
+    #         raise
+
+    #     """
+    #     from .variables.torch_function import get_prev_stack_var_name
+
+    #     setup_try_except, epilogue = _bytecode_from_template_with_split(
+    #         _try_except_tf_mode_template,
+    #         self.stack_index,
+    #         varname_map={"stack_var_name": get_prev_stack_var_name()},
+    #     )
+    #     cleanup[:] = epilogue + cleanup
+
+    #     return setup_try_except
 
     # If we do not want to destroy the stack, we can do the same thing as a
     # `SETUP_WITH` block, only that we store the context manager in a local_symbol
diff --git a/torch/_dynamo/side_effects.py b/torch/_dynamo/side_effects.py
@@ -629,22 +629,11 @@ def codegen_update_mutated(self, cg: PyCodegen):
             elif isinstance(
                 var, variables.torch_function.TorchFunctionModeStackVariable
             ):
-                # Needed in the finally block for stack restoration
-                cg.add_push_null(
-                    lambda: cg.load_import_from(
-                        utils.__name__, "get_torch_function_mode_stack"
-                    )
-                )
-                cg.call_function(0, False)
-                name = variables.torch_function.get_prev_stack_var_name()
-                cg.code_options["co_varnames"] += (name,)
-                cg.append_output(create_instruction("STORE_FAST", argval=name))
                 cg.add_push_null(
                     lambda: cg.load_import_from(
                         utils.__name__, "set_torch_function_mode_stack"
                     )
                 )
-
                 cg.foreach(var.symbolic_stack)
                 cg.append_output(
                     create_instruction("BUILD_LIST", arg=len(var.symbolic_stack))
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -267,12 +267,13 @@ def resume_fn(self):
         else:
             return ReenterWith(self.stack_index)
 
-    def exit(self, tx, is_graph_break):
+    def exit(self, tx):
+        if hasattr(self, "graph_break") and isinstance(
+            self.with_context, TorchFunctionModeVariable
+        ):
+            return
         assert self.with_context is not None
-        if (
-            is_graph_break and self.with_context.exit_on_graph_break()
-        ) or not is_graph_break:
-            return self.with_context.exit(tx)
+        return self.with_context.exit(tx)
 
 
 class ReturnValueOp(Exception):
@@ -656,17 +657,10 @@ def handle_graph_break(
             cleanup: List[Instruction] = []
             # Reconstruct the context variable CLASS in the block stack
             for b in self.block_stack:
-                # Don't exit any modes we have entered,
-                # output bytecode will mutate the tf mode stack accordingly
-                if isinstance(b.with_context, TorchFunctionModeVariable):
-                    cg.extend_output(
-                        b.resume_fn().try_except_torch_function_mode(
-                            cg.code_options, cleanup
-                        )
-                    )
-                    continue
                 assert b.with_context is not None
-                assert isinstance(b.with_context, (ContextWrappingVariable))
+                assert isinstance(
+                    b.with_context, (ContextWrappingVariable, TorchFunctionModeVariable)
+                )
                 b.with_context.reconstruct_type(cg)
                 cg.extend_output(b.resume_fn().try_finally(cg.code_options, cleanup))
             self.output.add_output_instructions(cg.get_instructions())
@@ -2320,10 +2314,7 @@ def setup_or_before_with(self, inst):
         ):
             unimplemented(f"{inst.opname} {ctx}")
 
-        if (
-            isinstance(ctx, GenericContextWrappingVariable)
-            and not ctx.supports_graph_breaks()
-        ):
+        if isinstance(ctx, GenericContextWrappingVariable):
             self.generic_context_manager_depth += 1
 
         # Need this redundant check for mypy
@@ -2696,7 +2687,6 @@ def __init__(
                 local_scope=f_locals,
                 global_scope=f_globals,
                 f_code=f_code,
-                torch_function_mode_stack=torch_function_mode_stack,
             ),
             instructions=instructions,
             f_locals=f_locals,
diff --git a/torch/_dynamo/testing.py b/torch/_dynamo/testing.py
@@ -187,7 +187,6 @@ def insert_nops(instructions: List[Any], code_options: Any) -> None:
         local_scope=locals(),
         global_scope=globals(),
         f_code=frame.f_code,
-        torch_function_mode_stack=[],
     )
 
     return GuardedCode(code, CheckFunctionManager(graph).check_fn, CompileId(0, 0))
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
@@ -304,7 +304,6 @@
     "torch.fx.experimental.symbolic_shapes.guard_size_oblivious": TorchInGraphFunctionVariable,
     "torch.cuda._get_device_properties": TorchInGraphFunctionVariable,
     "torch.utils.hooks.BackwardHook": TorchInGraphFunctionVariable,
-    "torch.set_default_device": UserFunctionVariable,
     "torch.sparse_bsc_tensor": SkipFunctionVariable,
     "torch.sparse_bsr_tensor": SkipFunctionVariable,
     "torch.sparse_csc_tensor": SkipFunctionVariable,
@@ -2803,6 +2802,7 @@
         "torch.random.initial_seed",
         "torch.random.seed",
         "torch.return_types.pytree_register_structseq",
+        "torch.set_default_device",
         "torch.set_default_dtype",
         "torch.set_default_tensor_type",
         "torch.set_deterministic_debug_mode",
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -204,7 +204,6 @@
 from .torch_function import (
     build_torch_function_fn,
     TensorWithTFOverrideVariable,
-    torch_function_mode_stack_state_mgr,
     TorchFunctionModeVariable,
 )
 from .user_defined import (
@@ -1670,16 +1669,15 @@ def wrap_numpy_ndarray(self, value):
                 # but warning is not the end of the world
                 assert isinstance(value.base, np.nditer)
 
-        with torch_function_mode_stack_state_mgr.temp_restore_stack():
-            try:
-                tensor_value = _util._try_convert_to_tensor(value)
-                if readonly:
-                    from torch._prims_common import clone_preserve_strides
-
-                    tensor_value = clone_preserve_strides(tensor_value)
-            except NotImplementedError as e:
-                # failed to convert to tensor, graph break
-                unimplemented(str(e))
+        try:
+            tensor_value = _util._try_convert_to_tensor(value)
+            if readonly:
+                from torch._prims_common import clone_preserve_strides
+
+                tensor_value = clone_preserve_strides(tensor_value)
+        except NotImplementedError as e:
+            # failed to convert to tensor, graph break
+            unimplemented(str(e))
 
         # We do this because we want the full behavior of guarding the numpy ndarray as if it were
         # a tensor. It's a little annoying to make a VT to throw out, but there's so many side effects here
diff --git a/torch/_dynamo/variables/ctx_manager.py b/torch/_dynamo/variables/ctx_manager.py
@@ -125,12 +125,6 @@ def call_function(
         if isinstance(args[0], UserFunctionVariable):
             return WrappedUserFunctionVariable(args[0], self)
 
-    def supports_graph_breaks(self):
-        return True
-
-    def exit_on_graph_break(self):
-        return True
-
 
 class GenericContextWrappingVariable(UserDefinedObjectVariable):
     # Some methods in ContextWrappingVariable assumes the arguments are
@@ -189,12 +183,6 @@ def exit(self, tx: "InstructionTranslator", *args):
         tx.generic_context_manager_depth -= 1
         return x
 
-    def supports_graph_breaks(self):
-        return False
-
-    def exit_on_graph_break(self):
-        return True
-
 
 class GradInplaceRequiresGradCtxManagerVariable(ContextWrappingVariable):
     """represents torch grad requries grad"""
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
diff --git a/torch/_dynamo/variables/torch_function.py b/torch/_dynamo/variables/torch_function.py
diff --git a/torch/_dynamo/variables/user_defined.py b/torch/_dynamo/variables/user_defined.py

Original file line number	Diff line number	Diff line change
`@@ -187,7 +187,6 @@ def insert_nops(instructions: List[Any], code_options: Any) -> None:`
`187`	`187`	`local_scope=locals(),`
`188`	`188`	`global_scope=globals(),`
`189`	`189`	`f_code=frame.f_code,`
`190`		`- torch_function_mode_stack=[],`
`191`	`190`	`)`
`192`	`191`
`193`	`192`	`return GuardedCode(code, CheckFunctionManager(graph).check_fn, CompileId(0, 0))`