Add multi-GPU track support

Add per-GPU grouping in the UI when a trace contains multiple GPUs.
Previously, all GPU tracks (render stages, counters, frequency, etc.)
were merged into a single flat "GPU" group regardless of how many
physical GPUs existed.

When multiple GPUs exist, separate top-level groups "GPU (1)",
"GPU (2)", etc. are created. GPUs are enumerated sequentially rather
than using the raw gpu_id from the trace (which may be a hash).
Single-GPU traces look identical to before.

Changes:
- Add gpu dimension to render stage track blueprint so gpu_id is
  stored as a track dimension (matching GPU counters).
- Add getOrCreateGpuGroup() to StandardGroupsPlugin to create
  top-level per-GPU groups.
- Extract gpu_id from dimension_arg_set_id in TraceProcessorTrackPlugin
  queries and route GPU tracks to per-GPU groups.
- Update GpuFreq plugin to use per-GPU grouping.

Bug: #5097

Author: dreveman

src/trace_processor/importers/proto/gpu_event_parser.cc

@@ -115,6 +115,7 @@ using protos::pbzero::VulkanMemoryEvent;
 constexpr auto kRenderStageBlueprint = TrackCompressor::SliceBlueprint(
     "gpu_render_stage",
     tracks::DimensionBlueprints(
+        tracks::kGpuDimensionBlueprint,
         tracks::StringDimensionBlueprint("render_stage_source"),
         tracks::UintDimensionBlueprint("hwqueue_id"),
         tracks::StringIdDimensionBlueprint("hwqueue_name")),
@@ -283,6 +284,7 @@ StringId GpuEventParser::GetFullStageName(
 }
 
 void GpuEventParser::InsertTrackForUninternedRenderStage(
+    uint32_t gpu_id,
     uint32_t hw_queue_id,
     const GpuRenderStageEvent::Specifications::Description::Decoder& hw_queue) {
   if (!hw_queue.has_name()) {
@@ -314,7 +316,7 @@ void GpuEventParser::InsertTrackForUninternedRenderStage(
 
   auto factory = context_->track_compressor->CreateTrackFactory(
       kRenderStageBlueprint,
-      tracks::Dimensions("id", hw_queue_id, kNullStringId),
+      tracks::Dimensions(gpu_id, "id", hw_queue_id, kNullStringId),
       tracks::DynamicName(name),
       [&, this](ArgsTracker::BoundInserter& inserter) {
         inserter.AddArg(description_id_, Variadic::String(description));
@@ -379,12 +381,14 @@ void GpuEventParser::ParseGpuRenderStageEvent(
   GpuRenderStageEvent::Decoder event(blob);
 
   int32_t pid = 0;
+  uint32_t gpu_id =
+      event.has_gpu_id() ? static_cast<uint32_t>(event.gpu_id()) : 0;
   if (event.has_specifications()) {
     GpuRenderStageEvent::Specifications::Decoder spec(event.specifications());
     uint32_t hw_queue_id = 0;
     for (auto it = spec.hw_queue(); it; ++it) {
       GpuRenderStageEvent::Specifications::Description::Decoder hw_queue(*it);
-      InsertTrackForUninternedRenderStage(hw_queue_id++, hw_queue);
+      InsertTrackForUninternedRenderStage(gpu_id, hw_queue_id++, hw_queue);
     }
     for (auto it = spec.stage(); it; ++it) {
       GpuRenderStageEvent::Specifications::Description::Decoder stage(*it);
@@ -477,7 +481,7 @@ void GpuEventParser::ParseGpuRenderStageEvent(
                                       : kNullStringId;
     TrackId track_id = context_->track_compressor->InternScoped(
         kRenderStageBlueprint,
-        tracks::Dimensions(base::StringView(source),
+        tracks::Dimensions(gpu_id, base::StringView(source),
                            static_cast<uint32_t>(hw_queue_id), dimension_name),
         ts, static_cast<int64_t>(event.duration()),
         tracks::DynamicName(track_name),

src/trace_processor/importers/proto/gpu_event_parser.h

@@ -84,6 +84,7 @@ class GpuEventParser {
       PacketSequenceStateGeneration* sequence_state,
       const protos::pbzero::GpuRenderStageEvent_Decoder& event) const;
   void InsertTrackForUninternedRenderStage(
+      uint32_t gpu_id,
       uint32_t id,
       const protos::pbzero::GpuRenderStageEvent::Specifications::Description::
           Decoder&);

ui/src/plugins/dev.perfetto.GpuFreq/index.ts

@@ -44,25 +44,29 @@ export default class implements PerfettoPlugin {
 
     if (tracks.length === 0) return;
 
-    const gpuGroup = ctx.plugins
-      .getPlugin(StandardGroupsPlugin)
-      .getOrCreateStandardGroup(ctx.defaultWorkspace, 'GPU');
+    const standardGroupsPlugin = ctx.plugins.getPlugin(StandardGroupsPlugin);
 
-    // Only create a sub-group if there's more than one track.
-    let parent: TrackNode;
-    if (tracks.length > 1) {
-      parent = new TrackNode({
-        name: 'GPU Frequency',
-        isSummary: true,
-      });
-      gpuGroup.addChildInOrder(parent);
-    } else {
-      parent = gpuGroup;
+    // Query distinct GPU IDs across all GPU tracks and build a mapping from
+    // raw gpu_id to a 1-based enumerated index.
+    const gpuIdsResult = await ctx.engine.query(`
+      select distinct extract_arg(dimension_arg_set_id, 'gpu') as gpu_id
+      from track
+      where extract_arg(dimension_arg_set_id, 'gpu') is not null
+      order by gpu_id
+    `);
+    const gpuIdToIndex = new Map<number, number>();
+    const gpuIt = gpuIdsResult.iter({gpu_id: NUM});
+    let idx = 1;
+    for (; gpuIt.valid(); gpuIt.next()) {
+      gpuIdToIndex.set(gpuIt.gpu_id, idx++);
     }
+    const gpuCount = gpuIdToIndex.size;
 
     for (const {id, gpuId, unit} of tracks) {
+      const gpuIndex = gpuIdToIndex.get(gpuId) ?? 1;
       const uri = `/gpu_frequency_${gpuId}`;
-      const name = `Gpu ${gpuId} Frequency`;
+      const name =
+        gpuCount > 1 ? `GPU (${gpuIndex}) Frequency` : 'GPU Frequency';
       ctx.tracks.registerTrack({
         uri,
         tags: {
@@ -77,8 +81,16 @@ export default class implements PerfettoPlugin {
           name,
         ),
       });
+
+      // Determine the parent group for this frequency track.
+      const gpuGroup = standardGroupsPlugin.getOrCreateGpuGroup(
+        ctx.defaultWorkspace,
+        gpuIndex,
+        gpuCount,
+      );
+
       const track = new TrackNode({uri, name, sortOrder: -20});
-      parent.addChildInOrder(track);
+      gpuGroup.addChildInOrder(track);
     }
   }
 }

ui/src/plugins/dev.perfetto.StandardGroups/index.ts

@@ -50,6 +50,8 @@ export default class implements PerfettoPlugin {
     HYPERVISOR: makeGroupNode('Hypervisor'),
   };
 
+  private readonly gpuSubGroups = new Map<number, TrackNode>();
+
   async onTraceLoad() {}
 
   /**
@@ -70,6 +72,37 @@ export default class implements PerfettoPlugin {
 
     return node;
   }
+
+  /**
+   * Gets or creates a per-GPU group for multi-GPU traces.
+   *
+   * When gpuCount <= 1, returns the standard "GPU" group (no change).
+   * When gpuCount > 1, creates/returns a top-level "GPU (N)" group where N
+   * is a 1-based enumerated index (not the raw gpu_id from the trace).
+   *
+   * @param workspace - The workspace on which to create the group.
+   * @param gpuIndex - 1-based enumerated GPU index.
+   * @param gpuCount - Total number of distinct GPUs in the trace.
+   */
+  getOrCreateGpuGroup(
+    workspace: Workspace,
+    gpuIndex: number,
+    gpuCount: number,
+  ): TrackNode {
+    if (gpuCount <= 1) {
+      return this.getOrCreateStandardGroup(workspace, 'GPU');
+    }
+
+    const existing = this.gpuSubGroups.get(gpuIndex);
+    if (existing) {
+      return existing;
+    }
+
+    const group = makeGroupNode(`GPU (${gpuIndex})`);
+    workspace.addChildInOrder(group);
+    this.gpuSubGroups.set(gpuIndex, group);
+    return group;
+  }
 }
 
 function makeGroupNode(name: string, collapsed = true) {

ui/src/plugins/dev.perfetto.TraceProcessorTrack/index.ts

@@ -92,6 +92,8 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
 
   private groups = new Map<string, TrackNode>();
   private store?: Store<TraceProcessorTrackPluginState>;
+  private gpuCount = 0;
+  private gpuIdToIndex = new Map<number, number>();
 
   private migrateTraceProcessorTrackPluginState(
     init: unknown,
@@ -104,6 +106,23 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
     this.store = ctx.mountStore(TraceProcessorTrackPlugin.id, (init) =>
       this.migrateTraceProcessorTrackPluginState(init),
     );
+
+    // Query distinct GPU IDs in the trace and build a mapping from raw gpu_id
+    // to a 1-based enumerated index. The raw gpu_id may be a hash or arbitrary
+    // value, so we enumerate sequentially for display purposes.
+    const gpuIdsResult = await ctx.engine.query(`
+      select distinct extract_arg(dimension_arg_set_id, 'gpu') as gpu_id
+      from track
+      where extract_arg(dimension_arg_set_id, 'gpu') is not null
+      order by gpu_id
+    `);
+    const gpuIt = gpuIdsResult.iter({gpu_id: NUM});
+    let gpuIndex = 1;
+    for (; gpuIt.valid(); gpuIt.next()) {
+      this.gpuIdToIndex.set(gpuIt.gpu_id, gpuIndex++);
+    }
+    this.gpuCount = this.gpuIdToIndex.size;
+
     await this.addCounters(ctx);
     await this.addSlices(ctx);
     this.addAggregations(ctx);
@@ -124,6 +143,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
           ct.machine_id as machine,
           extract_arg(ct.dimension_arg_set_id, 'utid') as utid,
           extract_arg(ct.dimension_arg_set_id, 'upid') as upid,
+          extract_arg(ct.dimension_arg_set_id, 'gpu') as gpu_id,
           extract_arg(ct.source_arg_set_id, 'description') as description
         from counter_track ct
         join _counter_track_summary using (id)
@@ -153,6 +173,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
       unit: STR_NULL,
       utid: NUM_NULL,
       upid: NUM_NULL,
+      gpu_id: NUM_NULL,
       threadName: STR_NULL,
       processName: STR_NULL,
       tid: LONG_NULL,
@@ -170,6 +191,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
         unit,
         utid,
         upid,
+        gpu_id: gpuId,
         threadName,
         processName,
         tid,
@@ -232,6 +254,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
         group,
         upid,
         utid,
+        gpuId,
         new TrackNode({
           uri,
           name: trackName,
@@ -264,6 +287,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
             lower(min(t.name)) as lower_name,
             extract_arg(t.dimension_arg_set_id, 'utid') as utid,
             extract_arg(t.dimension_arg_set_id, 'upid') as upid,
+            extract_arg(t.dimension_arg_set_id, 'gpu') as gpu_id,
             extract_arg(t.source_arg_set_id, 'description') as description,
             min(t.id) minTrackId,
             group_concat(t.id) as trackIds,
@@ -277,13 +301,14 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
           from _slice_track_summary s
           join track t using (id)
           left join _track_event_tracks_with_callstacks cs on cs.track_id = t.id
-          group by type, upid, utid, t.track_group_id, ifnull(t.track_group_id, t.id)
+          group by type, upid, utid, gpu_id, t.track_group_id, ifnull(t.track_group_id, t.id)
         )
         select
           s.type,
           s.name,
           s.utid,
           ifnull(s.upid, tp.upid) as upid,
+          s.gpu_id,
           s.minTrackId as minTrackId,
           s.trackIds as trackIds,
           s.trackCount,
@@ -330,6 +355,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
       name: STR_NULL,
       utid: NUM_NULL,
       upid: NUM_NULL,
+      gpu_id: NUM_NULL,
       trackIds: STR,
       maxDepth: NUM,
       tid: LONG_NULL,
@@ -351,6 +377,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
         maxDepth,
         utid,
         upid,
+        gpu_id: gpuId,
         threadName,
         processName,
         tid,
@@ -417,6 +444,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
         group,
         upid,
         utid,
+        gpuId,
 
         new TrackNode({
           uri,
@@ -436,6 +464,7 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
     group: string | TrackGroupSchema | undefined,
     upid: number | null,
     utid: number | null,
+    gpuId: number | null,
     track: TrackNode,
   ) {
     switch (topLevelGroup) {
@@ -466,10 +495,28 @@ export default class TraceProcessorTrackPlugin implements PerfettoPlugin {
         break;
       }
       default: {
-        const standardGroup = ctx.plugins
-          .getPlugin(StandardGroupsPlugin)
-          .getOrCreateStandardGroup(ctx.defaultWorkspace, topLevelGroup);
-        this.getGroupByName(standardGroup, group, null).addChildInOrder(track);
+        const standardGroupsPlugin =
+          ctx.plugins.getPlugin(StandardGroupsPlugin);
+
+        // For GPU tracks with a gpu_id, use per-GPU groups when multiple
+        // GPUs exist in the trace.
+        if (topLevelGroup === 'GPU' && gpuId !== null) {
+          const gpuIndex = this.gpuIdToIndex.get(gpuId) ?? 1;
+          const gpuGroup = standardGroupsPlugin.getOrCreateGpuGroup(
+            ctx.defaultWorkspace,
+            gpuIndex,
+            this.gpuCount,
+          );
+          this.getGroupByName(gpuGroup, group, gpuId).addChildInOrder(track);
+        } else {
+          const standardGroup = standardGroupsPlugin.getOrCreateStandardGroup(
+            ctx.defaultWorkspace,
+            topLevelGroup,
+          );
+          this.getGroupByName(standardGroup, group, null).addChildInOrder(
+            track,
+          );
+        }
         break;
       }
     }