Add transmission/delay budgets to XProf MXLA trace/graph viewer.

PiperOrigin-RevId: 675340215
openxla · Sep 16, 2024 · 90300c7 · 90300c7
1 parent 0cc9959
commit 90300c7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/third_party/tsl/tsl/profiler/utils/xplane_schema.cc b/third_party/tsl/tsl/profiler/utils/xplane_schema.cc
@@ -365,6 +365,8 @@ const MegaScaleStatTypeMap& GetMegaScaleStatTypeMap() {
       {"chunk", kMegaScaleChunk},
       {"launch_id", kMegaScaleLaunchId},
       {"loop_iteration", kMegaScaleLoopIteration},
+      {"transmission_budget_us", kMegaScaleTransmissionBudgetUs},
+      {"delay_budget_us", kMegaScaleDelayBudgetUs},
       {"graph_protos", kMegaScaleGraphProtos},
       {"network_transport_latency_us", kMegaScaleNetworkTransportLatency},
   });

diff --git a/third_party/tsl/tsl/profiler/utils/xplane_schema.h b/third_party/tsl/tsl/profiler/utils/xplane_schema.h
@@ -353,7 +353,9 @@ enum MegaScaleStatType : uint8_t {
   kMegaScaleLoopIteration,
   kMegaScaleGraphProtos,
   kMegaScaleNetworkTransportLatency,
-  kLastMegaScaleStatType = kMegaScaleNetworkTransportLatency,
+  kMegaScaleTransmissionBudgetUs,
+  kMegaScaleDelayBudgetUs,
+  kLastMegaScaleStatType = kMegaScaleDelayBudgetUs,
 };
 
 enum TaskEnvStatType {