From 90300c74ecc2009fec5fd176e9097be722affa4e Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Mon, 16 Sep 2024 16:52:15 -0700
Subject: [PATCH] Add transmission/delay budgets to XProf MXLA trace/graph
 viewer.

PiperOrigin-RevId: 675340215
---
 third_party/tsl/tsl/profiler/utils/xplane_schema.cc | 2 ++
 third_party/tsl/tsl/profiler/utils/xplane_schema.h  | 4 +++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/third_party/tsl/tsl/profiler/utils/xplane_schema.cc b/third_party/tsl/tsl/profiler/utils/xplane_schema.cc
index f5e8be8b6e5c2..da16a8704187b 100644
--- 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
index bdfc9256a5c36..96ebf29d4feff 100644
--- 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 {