From 6080ee728fd3c13fac6888438a41c1c1d41859b2 Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Thu, 12 Sep 2024 14:44:05 -0700
Subject: [PATCH] making event name related with cuda graph more informational
 by marking its graph id

PiperOrigin-RevId: 674017391
---
 xla/backends/profiler/gpu/cupti_collector.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/xla/backends/profiler/gpu/cupti_collector.cc b/xla/backends/profiler/gpu/cupti_collector.cc
index 6191849b0d0944..937b85f9a6c8fa 100644
--- a/xla/backends/profiler/gpu/cupti_collector.cc
+++ b/xla/backends/profiler/gpu/cupti_collector.cc
@@ -157,6 +157,11 @@ class PerDeviceCollector {
     if (kernel_name.empty()) {
       kernel_name = GetTraceEventTypeName(event.type);
     }
+    // For CPU events like cuGraphLaunch(), add the graph id to the name.
+    if (event.graph_id != 0 && event.type == CuptiTracerEventType::CudaGraph &&
+        event.source == CuptiTracerEventSource::DriverCallback) {
+      absl::StrAppend(&kernel_name, " (CudaGraph:", event.graph_id, ")");
+    }
     XEventMetadata* event_metadata =
         plane->GetOrCreateEventMetadata(std::move(kernel_name));
     XEventBuilder xevent = line->AddEvent(*event_metadata);