Log aot and idx waitcounters. (#152444)

jovianjaison · pytorchmergebot · commit ea9ea029595a · 2025-05-05T17:35:29.000Z
Summary: Added for create_aot_dispatcher_function and compile_fx_inner. Note: Log wait counters flag is already set for: 1. async_compile.precompile 2. remote_fx_graph_cache_get 3. remote_fx_graph_cache_put Test Plan: contbuild Differential Revision: D73866124 Pull Request resolved: #152444 Approved by: https://github.com/ppanchalia, https://github.com/masnesral
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
@@ -1547,6 +1547,8 @@ def call_user_compiler(self, gm: fx.GraphModule) -> CompiledFn:
             "OutputGraph.call_user_compiler",
             phase_name="backend_compile",
             log_pt2_compile_event=True,
+            log_waitcounter=True,
+            waitcounter_name_override="compile_aot_autograd",
             dynamo_compile_column_us="aot_autograd_cumulative_compile_time_us",
         ):
             return self._call_user_compiler(gm)
diff --git a/torch/_inductor/async_compile.py b/torch/_inductor/async_compile.py
@@ -387,6 +387,7 @@ def get_result() -> CachingAutotuner:
                 log_pt2_compile_event=True,
                 dynamo_compile_column_us="triton_compile_time_us",
                 log_waitcounter=True,
+                waitcounter_name_override="compile_async_triton",
             ):
                 start_ns = time_ns()
                 _set_triton_ptxas_path()
@@ -475,6 +476,7 @@ def wait(self, scope: dict[str, Any]) -> None:
                 log_pt2_compile_event=True,
                 dynamo_compile_column_us="triton_compile_time_us",
                 log_waitcounter=True,
+                waitcounter_name_override="compile_async_triton",
             ):
                 self._wait_futures(scope)
 
diff --git a/torch/_inductor/compile_fx.py b/torch/_inductor/compile_fx.py
@@ -714,6 +714,8 @@ def compile_fx_inner(
                 "compile_fx_inner",
                 phase_name="inductor_compile",
                 log_pt2_compile_event=True,
+                log_waitcounter=True,
+                waitcounter_name_override="compile_inductor",
                 dynamo_compile_column_us="inductor_cumulative_compile_time_us",
             )
         )

Original file line number	Diff line number	Diff line change
`@@ -714,6 +714,8 @@ def compile_fx_inner(`
`714`	`714`	`"compile_fx_inner",`
`715`	`715`	`phase_name="inductor_compile",`
`716`	`716`	`log_pt2_compile_event=True,`
	`717`	`+ log_waitcounter=True,`
	`718`	`+ waitcounter_name_override="compile_inductor",`
`717`	`719`	`dynamo_compile_column_us="inductor_cumulative_compile_time_us",`
`718`	`720`	`)`
`719`	`721`	`)`