Open source TPU-specific input pipeline analysis.

PiperOrigin-RevId: 722815891
makesoftwaresafe · Feb 3, 2025 · b0667ce · b0667ce
1 parent d195dd3
commit b0667ce
Show file tree

Hide file tree

Showing 5 changed files with 1,075 additions and 21 deletions.
diff --git a/tensorflow/core/profiler/convert/BUILD b/tensorflow/core/profiler/convert/BUILD
@@ -275,7 +275,9 @@ cc_library(
     copts = tf_profiler_copts(),
     deps = [
         ":op_metrics_to_record",
+        ":profile_time_breakdown",
         ":step_events_to_steps_db",
+        ":tpu_input_pipeline_analysis_constants",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core/platform:logging",
@@ -284,20 +286,43 @@ cc_library(
         "//tensorflow/core/profiler/protobuf:op_metrics_proto_cc",
         "//tensorflow/core/profiler/protobuf:op_stats_proto_cc",
         "//tensorflow/core/profiler/protobuf:steps_db_proto_cc",
+        "//tensorflow/core/profiler/protobuf:tpu_input_pipeline_proto_cc",
         "//tensorflow/core/profiler/utils:diagnostics",
         "//tensorflow/core/profiler/utils:event_span",
-        "//tensorflow/core/profiler/utils:hardware_type_utils",
         "//tensorflow/core/profiler/utils:html_utils",
         "//tensorflow/core/profiler/utils:math_utils",
         "//tensorflow/core/profiler/utils:op_metrics_db_utils",
+        "//tensorflow/core/profiler/utils:tpu_step_breakdown_utils",
+        "//tensorflow/core/profiler/utils:tpu_step_details_utils",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@local_xla//xla/hlo/ir:hlo",
+        "@local_xla//xla/tsl/profiler/convert:xla_op_utils",
         "@local_xla//xla/tsl/profiler/utils:format_utils",
+        "@local_xla//xla/tsl/profiler/utils:math_utils",
         "@local_xla//xla/tsl/profiler/utils:tf_op_utils",
         "@local_xla//xla/tsl/util:stats_calculator_portable",
     ],
 )
 
+tf_cc_test(
+    name = "op_stats_to_input_pipeline_analysis_test",
+    srcs = ["op_stats_to_input_pipeline_analysis_test.cc"],
+    deps = [
+        ":op_stats_to_input_pipeline_analysis",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/profiler/protobuf:steps_db_proto_cc",
+        "//tensorflow/core/profiler/utils:event_span",
+        "//tensorflow/core/profiler/utils:op_metrics_db_utils",
+        "@local_xla//xla/hlo/ir:hlo",
+        "@local_xla//xla/tsl/profiler/utils:timespan",
+    ],
+)
+
 cc_library(
     name = "op_stats_to_tf_stats",
     srcs = ["op_stats_to_tf_stats.cc"],