Fix CPU-GPU linking when ac2g start events are missing.

TraceLens/Trace2Tree/trace_to_tree.py

@@ -1194,28 +1194,26 @@ def _find_corresponding_output_event(self, input_event):
         # 1. Get the linking id from the input event
         # 2. Find the corresponding start and end ac2g events for the linking id
         # 3. Find the output event using the pid, tid, and linking id of the end ac2g event
+        # 4. Some runtimes (e.g. hipDrvLaunchKernelEx, hipMemsetAsync) emit only
+        #    the finish half of the ac2g flow, so fall back to an unambiguous
+        #    correlation-id match.
         link_id = input_event.get(TraceEventUtils.TraceKeys.Args, {}).get(
             self.linking_key
         )
         ac2g_start_event = self.ac2g_event_map["start"].get(link_id)
         ac2g_end_event = self.ac2g_event_map["end"].get(link_id)
 
-        if not ac2g_start_event:
-            return None
-
-        if not ac2g_end_event:
-            # print(f"Warning: start ac2g event found for {self.linking_key}={link_id} but no corresponding end ac2g event found.")
-            # print(f"Input event name: {input_event[TraceEventUtils.TraceKeys.Name]}")
-            # print(('-'*64))
+        if ac2g_start_event and ac2g_end_event:
+            pid = ac2g_end_event.get(TraceEventUtils.TraceKeys.PID)
+            tid = ac2g_end_event.get(TraceEventUtils.TraceKeys.TID)
+            end_link_id = ac2g_end_event.get("id")
+            return self.pid_tid_event_map.get((pid, tid, end_link_id))
+        else:
+            gpu_events = self.linking_id_to_gpu_events.get(link_id, [])
+            if len(gpu_events) == 1:
+                return gpu_events[0]
             return None
 
-        pid = ac2g_end_event.get(TraceEventUtils.TraceKeys.PID)
-        tid = ac2g_end_event.get(TraceEventUtils.TraceKeys.TID)
-        link_id = ac2g_end_event.get("id")
-
-        output_event = self.pid_tid_event_map.get((pid, tid, link_id))
-        return output_event
-
     def get_nn_module_children(self, nn_module_event: Dict[str, Any]):
         """
         Get the UIDs of the nn.Module children of the provided nn.Module event.

tests/test_trace2tree.py

@@ -634,6 +634,74 @@ def test_label_non_gpu_paths(self):
         assert "non_gpu_path" not in gpu_op_evt
         assert cpu_only.get("non_gpu_path") is True
 
+    def test_links_kernel_when_ac2g_start_is_missing(self):
+        def _launch_events(corr, gpu_events, launcher="hipDrvLaunchKernelEx"):
+            events = [
+                _mk_event("cpu_op", "aten::mm", ts=0, dur=100, pid=1, tid=1, args={}),
+                _mk_event(
+                    "cuda_runtime",
+                    launcher,
+                    ts=5,
+                    dur=5,
+                    pid=1,
+                    tid=1,
+                    args={"correlation": corr},
+                ),
+            ]
+            for idx, (cat, name) in enumerate(gpu_events):
+                events.append(
+                    _mk_event(
+                        cat,
+                        name,
+                        ts=20 + idx * 20,
+                        dur=10,
+                        pid=0,
+                        tid=7,
+                        args={"correlation": corr, "stream": 7},
+                    )
+                )
+            events.append(_mk_ac2g(corr, pid=0, tid=7, ts=20, phase="f"))
+            return events
+
+        # A single unambiguous kernel is recovered from the correlation id.
+        unique = _build_tree(_launch_events(26391, [("kernel", "Cijk_Alik_Bljk")]))
+        mm = next(e for e in unique.events if e["name"] == "aten::mm")
+        gpu_events = unique.get_gpu_events(mm)
+        assert len(gpu_events) == 1
+        assert gpu_events[0]["name"] == "Cijk_Alik_Bljk"
+
+        # Several kernels share the correlation id, so the match is ambiguous.
+        ambiguous = _build_tree(
+            _launch_events(42, [("kernel", "kernel_a"), ("kernel", "kernel_b")])
+        )
+        mm = next(e for e in ambiguous.events if e["name"] == "aten::mm")
+        assert ambiguous.get_gpu_events(mm) == []
+
+        # Unique memsets with only the ac2g finish event are linked too.
+        # Real traces also contain kernel launches, which is how linking_key
+        # is set to "correlation"; a memset-only trace would fall back to
+        # "External id" and never take this path.
+        memset_events = _launch_events(
+            43, [("gpu_memset", "Memset (Device)")], launcher="hipMemsetAsync"
+        )
+        memset_events.insert(
+            1,
+            _mk_event(
+                "cuda_runtime",
+                "hipLaunchKernel",
+                ts=1,
+                dur=1,
+                pid=1,
+                tid=1,
+                args={"correlation": 1},
+            ),
+        )
+        memset = _build_tree(memset_events)
+        mm = next(e for e in memset.events if e["name"] == "aten::mm")
+        gpu_events = memset.get_gpu_events(mm)
+        assert len(gpu_events) == 1
+        assert gpu_events[0]["name"] == "Memset (Device)"
+
     def test_linking_key_uses_correlation_when_present(self):
         events = [
             _mk_event("cpu_op", "aten::add", ts=0, dur=10, pid=1, tid=1, args={}),

tests/traces/inference/xdit_hunyuanvideo/perf_csvs/gpu_timeline.csv

@@ -1,9 +1,9 @@
 type,time ms,percent
-computation_time,2102.9709057617188,87.4578204938859
+computation_time,2102.980433105469,87.45821671464064
 exposed_comm_time,0.0,0.0
 exposed_memcpy_time,6.673357421875,0.27752989539457523
-busy_time,2109.644263183594,87.73535038928048
-idle_time,294.9101767578125,12.264649610719514
+busy_time,2109.653790527344,87.73574661003522
+idle_time,294.9006494140625,12.264253389964779
 total_time,2404.554439941406,100.0
 total_comm_time,0.0,0.0
 total_memcpy_time,6.673357421875,0.27752989539457523