extract calculate_time_spent from print_time_report (#127362)
Fixes #ISSUE_NUMBER
wrap certain steps in a separate function for easier TTFB instrumentation (fb internal use case)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127362
Approved by: https://github.com/yanboliang, https://github.com/mengluy0125
diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py
index 56340a3..5876895 100644
--- a/torch/_dynamo/utils.py
+++ b/torch/_dynamo/utils.py
@@ -154,12 +154,9 @@
op_count += cnt
-# Print a report of time spent so far
-# Ex:
-# TIMING:
-# entire_frame_compile:8.574629999999999
-# backend_compile:5.26806
-def print_time_report():
+# Calculate total time spent so far for each phase
+# For example, {'entire_frame_compile':8.574629999999999, 'backend_compile':5.26806}
+def calculate_time_spent():
total = 0.0
total_by_key = {}
for timings in frame_phase_timing.values():
@@ -170,6 +167,17 @@
else:
total_by_key[key] += timing
+ return total_by_key
+
+
+# Print a report of time spent so far
+# Ex:
+# TIMING:
+# entire_frame_compile:8.574629999999999
+# backend_compile:5.26806
+def print_time_report():
+ total_by_key = calculate_time_spent()
+
out = "TIMING:"
for key, value in total_by_key.items():
out = f"{out} {key}:{round(value, 5)}"
