[dynamo][compile-time] Compute logging related flags once (#153426)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/153426
Approved by: https://github.com/jansel

benchmarks/dynamo/pr_time_benchmarks/expected_results.csv

@@ -1,4 +1,4 @@
-add_loop_eager,compile_time_instruction_count,3051000000,0.015
+add_loop_eager,compile_time_instruction_count,2987000000,0.015
 
 
 
@@ -6,7 +6,7 @@ add_loop_eager_dynamic,compile_time_instruction_count,5928000000,0.025
 
 
 
-add_loop_inductor,compile_time_instruction_count,29380000000,0.015
+add_loop_inductor,compile_time_instruction_count,29370000000,0.015
 
 
 
@@ -18,7 +18,7 @@ add_loop_inductor_gpu,compile_time_instruction_count,25900000000,0.015
 
 
 
-basic_modules_ListOfLinears_eager,compile_time_instruction_count,945000000,0.015
+basic_modules_ListOfLinears_eager,compile_time_instruction_count,939900000,0.015
 
 
 
@@ -74,4 +74,4 @@ aotdispatcher_training_nosubclass_cpu,compile_time_instruction_count,3818000000,
 
 
 
-aotdispatcher_training_subclass_cpu,compile_time_instruction_count,10290000000,0.015
+aotdispatcher_training_subclass_cpu,compile_time_instruction_count,10350000000,0.015

torch/_dynamo/symbolic_convert.py

@@ -51,6 +51,7 @@ import torch
 import torch._logging
 from torch._dynamo.exc import TensorifyScalarRestartAnalysis
 from torch._guards import tracing, TracingContext
+from torch._logging.structured import dump_file
 from torch.fx.experimental.symbolic_shapes import guard_bool
 from torch.utils._functools import cache_method
 
@@ -1217,10 +1218,8 @@ class InstructionTranslatorBase(
         TracingContext.set_current_loc(
             self.f_code.co_filename, lineno, self.f_code.co_name
         )
-        from torch._logging.structured import dump_file
 
-        dump_file(self.f_code.co_filename)
-        if trace_source_log.isEnabledFor(logging.DEBUG):
+        if self.is_trace_source_log_enabled:
             trace_source_log.debug("%s", LazyString(self.get_log_starts_line_log_str))
 
     def step(self):
@@ -1243,7 +1242,7 @@ class InstructionTranslatorBase(
             if self.current_speculation.failed:
                 return self.step_graph_break(inst)
 
-        if trace_bytecode_log.isEnabledFor(logging.DEBUG):
+        if self.is_trace_bytecode_log_enabled:
             trace_bytecode_log.debug(
                 "TRACE %s %s %s", inst.opname, inst.argval, self.stack
             )
@@ -1339,6 +1338,7 @@ class InstructionTranslatorBase(
 
     def run(self):
         with self.run_ctx_mgr():
+            dump_file(self.f_code.co_filename)
             try:
                 self.output.push_tx(self)
                 self.start_point = self.instruction_pointer
@@ -3283,6 +3283,13 @@ class InstructionTranslatorBase(
         self._constants_cache: list[Optional[VariableTracker]] = [None] * len(
             f_code.co_consts
         )
+
+        self.is_trace_bytecode_log_enabled: Optional[bool] = (
+            trace_bytecode_log.isEnabledFor(logging.DEBUG)
+        )
+        self.is_trace_source_log_enabled: Optional[bool] = (
+            trace_source_log.isEnabledFor(logging.DEBUG)
+        )
         linecache.lazycache(f_code.co_filename, f_globals)
 
 

torch/_guards.py

@@ -899,9 +899,14 @@ class TracingContext:
             return traceback.StackSummary()
         stack = self.frame_summary_stack
         if self.loc_in_frame is not None:
-            stack = stack + [self.loc_in_frame]
+            stack = stack + [self._populate_loc_in_frame_summary()]
         return traceback.StackSummary.from_list(stack)
 
+    def _populate_loc_in_frame_summary(self):
+        assert self.loc_in_frame is not None
+        filename, lineno, frame_name = self.loc_in_frame
+        return traceback.FrameSummary(filename, lineno, frame_name, lookup_line=False)
+
     # Call this when you want to call into some code that isn't necessarily
     # associated with the current frame state
     @staticmethod
@@ -973,9 +978,9 @@ class TracingContext:
 
     @staticmethod
     def set_current_loc(filename, lineno, frame_name):
-        TracingContext.get().loc_in_frame = traceback.FrameSummary(
-            filename, lineno, frame_name, lookup_line=False
-        )
+        # Save the current location in the frame. Lazily generate the
+        # framesummary.
+        TracingContext.get().loc_in_frame = (filename, lineno, frame_name)
 
 
 @contextmanager