[hop] move insert_deferred_runtime_asserts under subtracer (#161416)

ydwu4 · pytorchmergebot · commit 16ce6a4aad75 · 2025-08-27T17:43:02.000Z
Pull Request resolved: #161416 Approved by: https://github.com/pianpwk ghstack dependencies: #160548
diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py
@@ -114,9 +114,7 @@ def _maybe_compile_and_run_fn(fn, *args):
 
 
 def reenter_make_fx(fn):
-    from torch._guards import detect_fake_mode
     from torch.fx.experimental.proxy_tensor import _CURRENT_MAKE_FX_TRACER
-    from torch.fx.passes.runtime_assert import insert_deferred_runtime_asserts
 
     @functools.wraps(fn)
     def wrapped(*args):
@@ -126,9 +124,6 @@ def wrapped(*args):
         gm = _CURRENT_MAKE_FX_TRACER.trace_subgraph(
             _maybe_run_with_interpreter(fn), *args
         )
-        if (fake_mode := detect_fake_mode()) and fake_mode.shape_env is not None:
-            insert_deferred_runtime_asserts(gm, fake_mode.shape_env, "reenter_make_fx")
-            gm.recompile()
         return gm
 
     return wrapped
diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py
@@ -2033,6 +2033,7 @@ def __init__(
         _allow_fake_constant: bool,
         _error_on_data_dependent_ops: bool,
         record_stack_traces: bool = False,
+        parent_tracer: Optional[_MakefxTracer] = None,
     ) -> None:
         # Configurations that are used to initialize the context managers and their states.
         # Should not modify them during tracing.
@@ -2064,6 +2065,7 @@ def __init__(
             nullcontext()
         )
         self.record_stack_traces = record_stack_traces
+        self.parent_tracer: Optional[_MakefxTracer] = parent_tracer
 
     def _checkpoint_modes(self) -> list[Any]:
         return [
@@ -2312,6 +2314,16 @@ def _wrap_func(f: Callable[_P, R], phs: Sequence[PHBase]) -> Callable[_P, R]:
                 )
                 raise
 
+        if (
+            self.is_hop_subgraph_tracer()
+            and (fake_mode := torch._guards.detect_fake_mode(args))
+            and fake_mode.shape_env is not None
+        ):
+            from torch.fx.passes.runtime_assert import insert_deferred_runtime_asserts
+
+            insert_deferred_runtime_asserts(t, fake_mode.shape_env, "reenter_make_fx")
+            t.recompile()
+
         # TODO: kind of a bad way to do it, should maybe figure out a better way
         if self.tracing_mode == "symbolic":
             assert self.fake_tensor_mode is not None
@@ -2322,6 +2334,9 @@ def trace(self, f: Callable, *args: object) -> fx.GraphModule:
         with self._init_modes_from_inputs(f, args):
             return self._trace_inner(f, *args)
 
+    def is_hop_subgraph_tracer(self) -> bool:
+        return self.parent_tracer is not None
+
     def trace_subgraph(self, f: Callable, *args: object) -> GraphModule:
         # Create a new tracer based on parent's config
         sub_tracer = _MakefxTracer(
@@ -2332,6 +2347,7 @@ def trace_subgraph(self, f: Callable, *args: object) -> GraphModule:
             self.record_module_stack,
             self._allow_fake_constant,
             self._error_on_data_dependent_ops,
+            parent_tracer=self,
         )
         with sub_tracer._init_modes_from_parent(self):
             return sub_tracer._trace_inner(f, *args)
