[logging] Add dynamo_compile fields for pre-dispatch/joint/post-dispatch times (#140306)

Tested internally: P1679622670

Differential Revision: [D65986059](https://our.internmc.facebook.com/intern/diff/D65986059)
Pull Request resolved: pytorch/pytorch#140306
Approved by: https://github.com/ezyang

Author: masnesral

test/dynamo/test_utils.py

@@ -156,6 +156,7 @@ def test_dynamo_timed(self, mock_time, mock_time_ns):
  'Scheduler.__init__': [0.0, 0.0],
  'Scheduler.codegen': [0.0, 0.0],
  '_compile.compile_inner': [0.0],
+ '_recursive_joint_graph_passes': [0.0],
  '_recursive_post_grad_passes': [0.0, 0.0],
  '_recursive_pre_grad_passes': [0.0],
  'async_compile.wait': [0.0, 0.0],
@@ -173,7 +174,10 @@ def test_dynamo_timed(self, mock_time, mock_time_ns):
         self.assertExpectedInline(
             pprint.pformat(time_spent),
             """\
-{'backend_compile': 0.0,
+{'_recursive_joint_graph_passes': 0.0,
+ '_recursive_post_grad_passes': 0.0,
+ '_recursive_pre_grad_passes': 0.0,
+ 'backend_compile': 0.0,
  'code_gen': 0.0,
  'entire_backward_compile': 0.0,
  'entire_frame_compile': 0.0,
@@ -235,9 +239,12 @@ def test_dynamo_timed(self, mock_time, mock_time_ns):
  'inductor_compile_time_s': 0.0,
  'inductor_cumulative_compile_time_us': 0,
  'is_forward': True,
+ 'joint_graph_pass_time_us': 0,
  'log_format_version': 3,
  'non_compliant_ops': set(),
  'num_triton_bundles': None,
+ 'post_grad_pass_time_us': 0,
+ 'pre_grad_pass_time_us': 0,
  'remote_cache_time_saved_s': None,
  'remote_fx_graph_cache_get_time_ms': None,
  'remote_fx_graph_cache_get_time_us': None,
@@ -295,9 +302,12 @@ def test_dynamo_timed(self, mock_time, mock_time_ns):
  'inductor_compile_time_s': 0.0,
  'inductor_cumulative_compile_time_us': 0,
  'is_forward': False,
+ 'joint_graph_pass_time_us': None,
  'log_format_version': 3,
  'non_compliant_ops': None,
  'num_triton_bundles': None,
+ 'post_grad_pass_time_us': 0,
+ 'pre_grad_pass_time_us': None,
  'remote_cache_time_saved_s': None,
  'remote_fx_graph_cache_get_time_ms': None,
  'remote_fx_graph_cache_get_time_us': None,

torch/_dynamo/utils.py

@@ -859,6 +859,9 @@ class CompilationMetrics:
     remote_fx_graph_cache_put_time_us: Optional[int] = None
     backward_cumulative_compile_time_us: Optional[int] = None
     end_time_us: Optional[int] = None
+    pre_grad_pass_time_us: Optional[int] = None
+    post_grad_pass_time_us: Optional[int] = None
+    joint_graph_pass_time_us: Optional[int] = None
     log_format_version: int = LOG_FORMAT_VERSION
 
 

torch/_inductor/compile_fx.py

@@ -325,7 +325,11 @@ def _get_subgraph_names(gm: GraphModule) -> Generator[str, None, None]:
 def _recursive_pre_grad_passes(
     gm: GraphModule, example_inputs: Sequence[InputType]
 ) -> GraphModule:
-    with dynamo_timed("_recursive_pre_grad_passes", log_pt2_compile_event=True):
+    with dynamo_timed(
+        "_recursive_pre_grad_passes",
+        log_pt2_compile_event=True,
+        dynamo_compile_column_us="pre_grad_pass_time_us",
+    ):
         for subgraph_name in _get_subgraph_names(gm):
             subgraph = getattr(gm, subgraph_name)
             # as we don't have recursive example inputs, passing empty set here
@@ -335,14 +339,23 @@ def _recursive_pre_grad_passes(
 
 
 def _recursive_joint_graph_passes(gm: GraphModule) -> None:
-    for subgraph_name in _get_subgraph_names(gm):
-        subgraph = getattr(gm, subgraph_name)
-        _recursive_joint_graph_passes(subgraph)
-    joint_graph_passes(gm)
+    with dynamo_timed(
+        "_recursive_joint_graph_passes",
+        log_pt2_compile_event=True,
+        dynamo_compile_column_us="joint_graph_pass_time_us",
+    ):
+        for subgraph_name in _get_subgraph_names(gm):
+            subgraph = getattr(gm, subgraph_name)
+            _recursive_joint_graph_passes(subgraph)
+        joint_graph_passes(gm)
 
 
 def _recursive_post_grad_passes(gm: GraphModule, is_inference: bool = False) -> None:
-    with dynamo_timed("_recursive_post_grad_passes", log_pt2_compile_event=True):
+    with dynamo_timed(
+        "_recursive_post_grad_passes",
+        log_pt2_compile_event=True,
+        dynamo_compile_column_us="post_grad_pass_time_us",
+    ):
         for subgraph_name in _get_subgraph_names(gm):
             subgraph = getattr(gm, subgraph_name)
             _recursive_post_grad_passes(subgraph, is_inference)