pytorch · jbschlosser · May 15, 2025 · May 15, 2025 · May 19, 2025
diff --git a/test/dynamo/test_utils.py b/test/dynamo/test_utils.py
@@ -1,5 +1,6 @@
 # Owner(s): ["module: dynamo"]
 import dataclasses
+import os
 import pprint
 import sys
 from unittest import mock
@@ -141,6 +142,63 @@ def break_it2(x):
             compilation_events = [arg[0][0] for arg in log_event.call_args_list]
             self.assertEqual(compilation_events[-1].num_graph_breaks, 2)
 
+    def test_frame_traced_hook(self):
+        from utils import add, break_it
+
+        traced_file_sets = []
+
+        def get_traced_files(s):
+            nonlocal traced_file_sets
+            traced_file_sets.append(s)
+
+        utils_path = os.path.join(os.path.dirname(__file__), "utils.py")
+
+        # === no inlining ===
+        @torch.compile(options={"frame_traced_fn": get_traced_files})
+        def fn(x):
+            return x * 2
+
+        x = torch.randn(3)
+        traced_file_sets = []
+        fn(x)
+        # expect hook to be called once with this file
+        self.assertEqual(traced_file_sets, [{__file__}])
+
+        # === successful inlining ===
+        @torch.compile(options={"frame_traced_fn": get_traced_files})
+        def fn(x):
+            return add(x) * 2
+
+        x = torch.randn(3)
+        traced_file_sets = []
+        fn(x)
+        utils_path = os.path.join(os.path.dirname(__file__), "utils.py")
+        # expect hook to be called once with both this file and file of inlined func
+        self.assertEqual(traced_file_sets, [{__file__, utils_path}])
+
+        # === graph break occurs during inlining ===
+        @torch.compile(options={"frame_traced_fn": get_traced_files})
+        def fn(x):
+            y = break_it(x)
+            return y * 2
+
+        x = torch.randn(3)
+        traced_file_sets = []
+        fn(x)
+        # expect hook to be called twice; once for this file one for file of inlined func
+        self.assertEqual(traced_file_sets, [{__file__}, {utils_path}])
+
+        # === empty graph ===
+        @torch.compile(options={"frame_traced_fn": get_traced_files})
+        def fn(x):
+            return x
+
+        x = torch.randn(3)
+        traced_file_sets = []
+        fn(x)
+        # hook is not expected to be called at all for an empty graph
+        self.assertEqual(traced_file_sets, [])
+
 
 class TestModel(torch.nn.Module):
     def __init__(self):

diff --git a/test/dynamo/utils.py b/test/dynamo/utils.py
@@ -39,6 +39,10 @@ def add(x):
     return x + 1
 
 
+def break_it(x):
+    return x.sum().item()
+
+
 def create_dummy_module_and_function():
     module = types.ModuleType("dummy_module")
     module.__spec__ = importlib.machinery.ModuleSpec(

diff --git a/torch/__init__.py b/torch/__init__.py
@@ -2597,6 +2597,10 @@ def fn(model: _Callable[_InputT, _RetT]) -> _Callable[_InputT, _RetT]:
     if options and isinstance(options, dict):
         guard_filter_fn = options.pop("guard_filter_fn", None)
 
+    frame_traced_fn = None
+    if options and isinstance(options, dict):
+        frame_traced_fn = options.pop("frame_traced_fn", None)
+
     if backend == "inductor":
         backend = _TorchCompileInductorWrapper(mode, options, dynamic)
     else:
@@ -2608,6 +2612,7 @@ def fn(model: _Callable[_InputT, _RetT]) -> _Callable[_InputT, _RetT]:
         dynamic=dynamic,
         disable=disable,
         guard_filter_fn=guard_filter_fn,
+        frame_traced_fn=frame_traced_fn,
     )(model)  # type: ignore[return-value]
 
 

diff --git a/torch/_dynamo/convert_frame.py b/torch/_dynamo/convert_frame.py
@@ -936,13 +936,17 @@ def count_args(code: CodeType) -> int:
             annotation_str,
         )
 
-        if not output.is_empty_graph() and hooks.guard_export_fn is not None:
-            # We should not run the guard_export_fn when Dynamo does not
-            # generate any graph. This can happen in export when TorchDynamo
-            # generated bytecode has some reconstruction logic for mutated
-            # variables which can trigger TorchDynamo on the children frames but
-            # they are benign and do not generate any new graphs.
-            hooks.guard_export_fn(output.guards)
+        if not output.is_empty_graph():
+            if hooks.guard_export_fn is not None:
+                # We should not run the guard_export_fn when Dynamo does not
+                # generate any graph. This can happen in export when TorchDynamo
+                # generated bytecode has some reconstruction logic for mutated
+                # variables which can trigger TorchDynamo on the children frames but
+                # they are benign and do not generate any new graphs.
+                hooks.guard_export_fn(output.guards)
+            if hooks.frame_traced_fn is not None:
+                output.tracing_context.traced_files.add(output.co_fields["co_filename"])
+                hooks.frame_traced_fn(output.tracing_context.traced_files)
 
         return wrap_guarded_code(guarded_code)
 

diff --git a/torch/_dynamo/eval_frame.py b/torch/_dynamo/eval_frame.py
@@ -994,6 +994,7 @@ def _optimize(
     guard_export_fn=None,
     guard_fail_fn=None,
     guard_filter_fn=None,
+    frame_traced_fn=None,
     disable=False,
     dynamic=None,
 ) -> Union[OptimizeContext, _NullDecorator]:
@@ -1033,6 +1034,7 @@ def toy_example(a, b): ...
         guard_export_fn=guard_export_fn,
         guard_fail_fn=guard_fail_fn,
         guard_filter_fn=guard_filter_fn,
+        frame_traced_fn=frame_traced_fn,
     )
     torch._C._log_api_usage_once("torch._dynamo.optimize")
     if (

diff --git a/torch/_dynamo/hooks.py b/torch/_dynamo/hooks.py
@@ -6,7 +6,8 @@
 The Hooks class manages two types of hook functions:
 - guard_export_fn: Called when guards need to be exported, taking a GuardsSet as input
 - guard_fail_fn: Called when a guard check fails, taking a GuardFail object as input
-
+- frame_traced_fn: Called when a frame has finished tracing, resulting in a non-empty graph.
+                   This hook will be passed the set of filenames containing traced code
 These hooks enable customization of guard export and failure handling behaviors.
 """
 
@@ -23,3 +24,4 @@ class Hooks:
     guard_export_fn: Optional[Callable[[GuardsSet], None]] = None
     guard_fail_fn: Optional[Callable[[GuardFail], None]] = None
     guard_filter_fn: Optional[Callable[[list[GuardFilterEntry]], list[bool]]] = None
+    frame_traced_fn: Optional[Callable[[set[str]], None]] = None
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -3921,6 +3921,7 @@ def inline_call_(self):
         parent.inconsistent_side_effects |= self.inconsistent_side_effects
 
         log.debug("DONE INLINING %s", code)
+        self.output.tracing_context.traced_files.add(code.co_filename)
 
         if config.enable_faithful_generator_behavior or (
             isinstance(self, InliningGeneratorInstructionTranslator)

diff --git a/torch/_guards.py b/torch/_guards.py
@@ -865,6 +865,8 @@ def __init__(self, fake_mode):
         # see note: [Returning Fake Tensors on First AOT Autograd Call]
         self.fakify_first_call = False
         self.hop_dispatch_set_cache = HopDispatchSetCache()
+        # set of filenames for inlined functions
+        self.traced_files: set[str] = set()
 
     def clear(self):
         # Look at the note in output_graph.py in function `save_global_state`