Skip to content

Commit

Permalink
Include the graph ID in the host CUDA graph launch event, add cuda gr…
Browse files Browse the repository at this point in the history
…aph id stat in the derived Hlo Op event

PiperOrigin-RevId: 674017391
  • Loading branch information
Google-ML-Automation committed Sep 18, 2024
1 parent ac14cd9 commit 43edfa8
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions xla/backends/profiler/gpu/cupti_collector.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 43edfa8

Please sign in to comment.