power metric specific proto

PiperOrigin-RevId: 561091977
tensorflow · Aug 29, 2023 · 9258563 · 9258563
1 parent 7b68b94
commit 9258563
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/plugin/tensorboard_plugin_profile/protobuf/BUILD b/plugin/tensorboard_plugin_profile/protobuf/BUILD
@@ -15,6 +15,7 @@ proto_library(
         "input_pipeline.proto",
         "kernel_stats.proto",
         "overview_page.proto",
+        "power_metrics.proto",
         "tf_data_stats.proto",
         "tf_stats.proto",
         "trace_events.proto",
@@ -30,6 +31,7 @@ py_proto_library(
         "input_pipeline.proto",
         "kernel_stats.proto",
         "overview_page.proto",
+        "power_metrics.proto",
         "tf_data_stats.proto",
         "tf_stats.proto",
         "trace_events.proto",

diff --git a/plugin/tensorboard_plugin_profile/protobuf/overview_page.proto b/plugin/tensorboard_plugin_profile/protobuf/overview_page.proto
@@ -5,6 +5,7 @@ package tensorflow.profiler;
 import "google/protobuf/any.proto";
 import "plugin/tensorboard_plugin_profile/protobuf/diagnostics.proto";
 import "plugin/tensorboard_plugin_profile/protobuf/input_pipeline.proto";
+import "plugin/tensorboard_plugin_profile/protobuf/power_metrics.proto";
 
 // Overview result for a TensorFlow Op.
 message OverviewTfOp {
@@ -187,6 +188,8 @@ message OverviewPageRunEnvironment {
   int32 num_cores_per_replica = 10;
   // Whether it is a training analysis or inference analysis.
   bool is_training = 11;
+  // Power Metrics for TPU.
+  PowerMetrics power_metrics = 12;
   reserved 6;
 }
 

diff --git a/plugin/tensorboard_plugin_profile/protobuf/power_metrics.proto b/plugin/tensorboard_plugin_profile/protobuf/power_metrics.proto
@@ -0,0 +1,16 @@
+syntax = "proto3";
+
+package tensorflow.profiler;
+
+message PowerComponentMetrics {
+  // power rail or component name, e.g. HBM, Core.
+  string component_name = 1;
+  // maximum watts monitored.
+  double max_power = 2;
+  // average watts monitored.
+  double avg_power = 3;
+}
+
+message PowerMetrics {
+  repeated PowerComponentMetrics power_component_metrics = 1;
+}