From 6c79cfabb4bb469bfc34daa460e55f8a772442d6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 3 Jul 2024 11:57:10 -0700
Subject: [PATCH] Support Cuda Graph in XProf

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

diff --git a/tsl/profiler/utils/xplane_schema.cc b/tsl/profiler/utils/xplane_schema.cc
index 15d77c874..33de2b0f6 100644
--- a/tsl/profiler/utils/xplane_schema.cc
+++ b/tsl/profiler/utils/xplane_schema.cc
@@ -335,6 +335,8 @@ const StatTypeMap& GetStatTypeMap() {
       {"dcn_chunk", kDcnChunk},
       {"dcn_loop_index", kDcnLoopIndex},
       {"dropped_traces", kDroppedTraces},
+      {"cuda_graph_id", kCudaGraphId},
+      {"cuda_graph_details", kCudaGraphDetails},
   });
   DCHECK_EQ(stat_type_map->size(), kNumStatTypes);
   return *stat_type_map;
diff --git a/tsl/profiler/utils/xplane_schema.h b/tsl/profiler/utils/xplane_schema.h
index 02ea6de1a..2e693b447 100644
--- a/tsl/profiler/utils/xplane_schema.h
+++ b/tsl/profiler/utils/xplane_schema.h
@@ -320,7 +320,9 @@ enum StatType {
   kEdgeTpuModelProfileInfo,
   kEdgeTpuMlir,
   kDroppedTraces,
-  kLastStatType = kDroppedTraces,
+  kCudaGraphId,
+  kCudaGraphDetails,
+  kLastStatType = kCudaGraphDetails,
 };
 
 enum MegaScaleStatType : uint8_t {