Open source core_id_to_details map for adding host names in pod viewer.
PiperOrigin-RevId: 336933370
Change-Id: I3f3dfa55359eace4590107696bc2cfde47c1bdcc
diff --git a/tensorflow/core/profiler/convert/op_stats_combiner.cc b/tensorflow/core/profiler/convert/op_stats_combiner.cc
index d4ce7ec..5c492a6 100644
--- a/tensorflow/core/profiler/convert/op_stats_combiner.cc
+++ b/tensorflow/core/profiler/convert/op_stats_combiner.cc
@@ -157,6 +157,10 @@
// Combine tf-function stats.
CombineTfFunctionDb(src.tf_function_db(), dst->mutable_tf_function_db());
+
+ // Combine the mapping from core ID to details.
+ CombineCoreIdMap(src_host_id, src.core_id_to_details(),
+ dst->mutable_core_id_to_details());
}
} // namespace
diff --git a/tensorflow/core/profiler/protobuf/op_stats.proto b/tensorflow/core/profiler/protobuf/op_stats.proto
index 500de69..53d6499 100644
--- a/tensorflow/core/profiler/protobuf/op_stats.proto
+++ b/tensorflow/core/profiler/protobuf/op_stats.proto
@@ -90,6 +90,17 @@
uint32 host_trace_level = 12;
}
+// Next ID: 7
+message CoreDetails {
+ string hostname = 1;
+ uint32 device_ordinal = 2; // unique within host, TPU core only
+ uint32 core_num = 3; // unique within chip per core type
+ uint32 local_chip_id = 4; // unique within host
+ uint32 global_chip_id = 5; // unique within mesh
+ uint32 global_core_id = 6; // unique within mesh, TPU core only
+}
+
+// Next ID: 12
// Operator Statistics.
message OpStats {
// The database for the op metrics collected from the host over the entire
@@ -110,6 +121,8 @@
KernelStatsDb kernel_stats_db = 6;
// Statistics for all tf-functions.
TfFunctionDb tf_function_db = 8;
+ // A map from core ID to details.
+ map<uint32, CoreDetails> core_id_to_details = 11;
// Error and warning messages for diagnosing profiling issues.
Diagnostics diagnostics = 9;
reserved 7;
