diff --git a/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv b/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv index bc32e7ae18f..748d6aecda2 100644 --- a/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv +++ b/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 diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py index c780905e9ea..b50e95a4357 100644 --- a/torch/_dynamo/symbolic_convert.py +++ b/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) diff --git a/torch/_guards.py b/torch/_guards.py index 9e724808172..8a763ff68e3 100644 --- a/torch/_guards.py +++ b/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