Add max infeed time core name to InputPipelineAnalysis.

Add max infeed table conversion in python.

PiperOrigin-RevId: 716386581

plugin/tensorboard_plugin_profile/convert/BUILD

@@ -69,6 +69,7 @@ py_test(
     srcs_version = "PY2AND3",
     deps = [
         ":input_pipeline_proto_to_gviz",
+        "//google/protobuf:any_py_pb2",
         requirement("absl-py"),  # build_cleaner: keep; go/disable_tf2
         "//third_party/py/google/protobuf:use_fast_cpp_protos",  # Automatically added go/proto_python_upb_flip
         requirement("gviz_api"),

plugin/tensorboard_plugin_profile/convert/input_pipeline_proto_to_gviz.py

@@ -474,6 +474,26 @@ def get_recommendation_table_args(ipa):
   return (table_description, data, None)
 
 
+def generate_max_infeed_core_table(ipa):
+  """Creates a table that reports the name of the TPU core that has the maximum infeed time at each step."""
+  table_description = [
+      ("index", "string", "Index"),
+  ]
+  step_number_row = ["Step Number"]
+  core_name_row = ["Core Name"]
+  for idx, raw_step_details in enumerate(ipa.step_details):
+    table_description.append((str(idx), "string", str(idx)))
+    step_details = tpu_input_pipeline_pb2.PerTpuStepDetails()
+    if not raw_step_details.Unpack(step_details):
+      continue
+    step_number_row.append(step_details.step_number)
+    core_name_row.append(step_details.max_infeed_time_core_name)
+
+  return gviz_api.DataTable(
+      table_description, (step_number_row, core_name_row), None
+  )
+
+
 def generate_step_breakdown_table_for_tpu(ipa):
   table_description, data, custom_properties = (
       get_step_breakdown_table_args_for_tpu(ipa)
@@ -502,14 +522,14 @@ def generate_recommendation_table(ipa):
 def generate_all_chart_tables(ipa):
   """Generates a list of gviz tables from InputPipelineAnalysisResult."""
 
-  step_breakdown_table = (
-      generate_step_breakdown_table(ipa)
-      if ipa.tag
-      else generate_step_breakdown_table_for_tpu(ipa)
-  )
+  tables = []
+  if ipa.tag:
+    tables.append(generate_step_breakdown_table_for_tpu(ipa))
+    tables.append(generate_max_infeed_core_table(ipa))
+  else:
+    tables.append(generate_step_breakdown_table(ipa))
 
-  return [
-      step_breakdown_table,
+  return tables + [
       generate_input_op_table(ipa),
       generate_recommendation_table(ipa),
       diag.generate_diagnostics_table(ipa.diagnostics),

plugin/tensorboard_plugin_profile/convert/input_pipeline_proto_to_gviz_test.py

@@ -455,6 +455,35 @@ def test_input_pipeline_step_breakdown_for_tpu_with_legacy_sc(self):
           else:
             self.assertEqual(str(expected[cc]), cell_str)
 
+  def test_generate_max_infeed_core_table(self):
+    ipa = input_pipeline_pb2.InputPipelineAnalysisResult()
+    ipa.hardware_type = "TPU"
+    ipa.tag = True
+    for step_number in range(0, 3):
+      step_details = tpu_input_pipeline_pb2.PerTpuStepDetails()
+      step_details.step_number = step_number
+      step_details.max_infeed_time_core_name = "core_name"
+      step_details_any = Any()
+      step_details_any.Pack(step_details)
+      ipa.step_details.append(step_details_any)
+    data_table = input_pipeline_proto_to_gviz.generate_max_infeed_core_table(
+        ipa
+    )
+    self.assertEqual(2, data_table.NumberOfRows())
+    for row in csv.reader(io.StringIO(data_table.ToCsv())):
+      self.assertIn(row[0], ["Index", "Step Number", "Core Name"])
+      if row[0] == "Index":
+        # Process column headers.
+        for idx, col in enumerate(row[1:]):
+          self.assertEqual(col, str(idx))
+      elif row[0] == "Step Number":
+        for idx, step_number in enumerate(row[1:]):
+          self.assertEqual(step_number, str(idx))
+      else:
+        self.assertEqual(row[0], "Core Name")
+        for core_name in row[1:]:
+          self.assertEqual(core_name, "core_name")
+
 
 if __name__ == "__main__":
   tf.test.main()

plugin/tensorboard_plugin_profile/protobuf/tpu_input_pipeline.proto

@@ -5,7 +5,7 @@ package tensorflow.profiler;
 import "plugin/tensorboard_plugin_profile/protobuf/input_pipeline.proto";
 
 // Per-step details on TPU.
-// Next ID: 25
+// Next ID: 26
 message PerTpuStepDetails {
   // The step number of a step.
   int32 step_number = 1;
@@ -36,6 +36,9 @@ message PerTpuStepDetails {
   // The core with the maximum infeed time in this step.
   uint32 coreid_max_infeed_time = 7;
 
+  // The name of the core with the maximum infeed time in this step.
+  string max_infeed_time_core_name = 25;
+
   // The part of a step (in ms) that is spent on the all-reduce compute.
   double all_reduce_compute_time_ms = 11;