Fix Perfetto trace overlap handling

Author: drisspg

test/test_perfetto.py

@@ -0,0 +1,152 @@
+import gzip
+import json
+
+from transformer_nuggets.utils.perfetto import (
+    default_trace_path,
+    read_trace,
+    split_overlapping_slices,
+    write_trace,
+)
+
+
+def _duration_events(trace):
+    return [event for event in trace["traceEvents"] if event.get("ph") == "X"]
+
+
+def test_split_overlapping_slices_creates_adjacent_lanes():
+    trace = {
+        "traceEvents": [
+            {
+                "ph": "M",
+                "name": "thread_name",
+                "pid": 0,
+                "tid": 7,
+                "args": {"name": "stream 7"},
+            },
+            {"ph": "X", "name": "a", "pid": 0, "tid": 7, "ts": 0, "dur": 10},
+            {"ph": "X", "name": "b", "pid": 0, "tid": 7, "ts": 5, "dur": 10},
+            {"ph": "X", "name": "c", "pid": 0, "tid": 7, "ts": 10, "dur": 1},
+        ]
+    }
+
+    fixed = split_overlapping_slices(trace, track_pattern="stream.*")
+    durations = _duration_events(fixed)
+
+    assert [event["tid"] for event in durations] == [700, 701, 700]
+    thread_names = {
+        event["tid"]: event["args"]["name"]
+        for event in fixed["traceEvents"]
+        if event.get("ph") == "M" and event.get("name") == "thread_name"
+    }
+    assert thread_names[700] == "stream 7 #0"
+    assert thread_names[701] == "stream 7 #1"
+
+
+def test_split_overlapping_slices_leaves_non_overlapping_tracks_unchanged():
+    trace = {
+        "traceEvents": [
+            {
+                "ph": "M",
+                "name": "thread_name",
+                "pid": 0,
+                "tid": 3,
+                "args": {"name": "stream 3"},
+            },
+            {"ph": "X", "name": "a", "pid": 0, "tid": 3, "ts": 0, "dur": 10},
+            {"ph": "X", "name": "b", "pid": 0, "tid": 3, "ts": 10, "dur": 10},
+        ]
+    }
+
+    fixed = split_overlapping_slices(trace, track_pattern="stream.*")
+
+    assert fixed == trace
+
+
+def test_split_overlapping_slices_remaps_flow_by_correlation_and_timestamp():
+    trace = {
+        "traceEvents": [
+            {
+                "ph": "M",
+                "name": "thread_name",
+                "pid": 0,
+                "tid": 7,
+                "args": {"name": "stream 7"},
+            },
+            {
+                "ph": "X",
+                "name": "a",
+                "pid": 0,
+                "tid": 7,
+                "ts": 0,
+                "dur": 10,
+                "args": {"correlation": 42},
+            },
+            {
+                "ph": "X",
+                "name": "b",
+                "pid": 0,
+                "tid": 7,
+                "ts": 5,
+                "dur": 10,
+                "args": {"correlation": 42},
+            },
+            {"ph": "f", "pid": 0, "tid": 7, "ts": 0, "id": 42},
+            {"ph": "f", "pid": 0, "tid": 7, "ts": 5, "id": 42},
+        ]
+    }
+
+    fixed = split_overlapping_slices(trace, track_pattern="stream.*")
+    flow_tids = [
+        event["tid"] for event in fixed["traceEvents"] if event.get("ph") == "f"
+    ]
+
+    assert flow_tids == [700, 701]
+
+
+def test_split_overlapping_slices_keeps_same_tid_in_different_pids_separate():
+    trace = {
+        "traceEvents": [
+            {
+                "ph": "M",
+                "name": "thread_name",
+                "pid": 0,
+                "tid": 7,
+                "args": {"name": "stream 7"},
+            },
+            {
+                "ph": "M",
+                "name": "thread_name",
+                "pid": 1,
+                "tid": 7,
+                "args": {"name": "stream 7"},
+            },
+            {"ph": "X", "name": "a", "pid": 0, "tid": 7, "ts": 0, "dur": 10},
+            {"ph": "X", "name": "b", "pid": 0, "tid": 7, "ts": 5, "dur": 10},
+            {"ph": "X", "name": "c", "pid": 1, "tid": 7, "ts": 0, "dur": 10},
+            {"ph": "X", "name": "d", "pid": 1, "tid": 7, "ts": 5, "dur": 10},
+        ]
+    }
+
+    fixed = split_overlapping_slices(trace, track_pattern="stream.*")
+    by_pid = {}
+    for event in _duration_events(fixed):
+        by_pid.setdefault(event["pid"], []).append(event["tid"])
+
+    assert by_pid == {0: [700, 701], 1: [700, 701]}
+
+
+def test_gzip_trace_roundtrip(tmp_path):
+    path = tmp_path / "trace.json.gz"
+    trace = {"traceEvents": [{"ph": "X", "name": "a", "pid": 0, "tid": 0}]}
+
+    write_trace(path, trace)
+
+    with gzip.open(path, "rt", encoding="utf-8") as f:
+        assert json.load(f) == trace
+    assert read_trace(path) == trace
+
+
+def test_default_trace_path_prefers_gzip_for_stems_and_respects_explicit_gzip():
+    assert default_trace_path("foo").as_posix() == "foo.json.gz"
+    assert default_trace_path("foo.json.gz").as_posix() == "foo.json.gz"
+    assert default_trace_path("foo", gzip_by_default=False).as_posix() == "foo.json"

transformer_nuggets/cute/profiler/README.md

@@ -20,7 +20,7 @@ with profile_session(
     max_events_per_unit=64,
     num_units=(num_blocks, "Block"),
     tag_names=["compute"],
-    trace_path="trace.json",
+    trace_path="trace.json.gz",
 ) as (prof, _):
     my_kernel(output, prof.tensor, prof.max_events_per_unit)
 ```
@@ -69,7 +69,7 @@ You set `max_events_per_unit` to something larger than you need; the decoder sca
 | `profile_session(...)` | Context manager: allocate, yield, decode, write trace |
 | `allocate_profile_buffer(max_events_per_unit, num_units, device)` | Allocate buffer |
 | `decode_events(buf, tag_table)` | Decode to `Event` list |
-| `events_to_perfetto(events, path)` | Write Chrome trace JSON |
+| `events_to_perfetto(events, path)` | Write Chrome trace JSON/JSON.GZ and split overlapping slices into adjacent lanes by default |
 | `TagTable(names)` | Map tag names ↔ integer IDs |
 | `PostProcessContext` | Context passed to post-processing callbacks |
 
@@ -83,6 +83,8 @@ You set `max_events_per_unit` to something larger than you need; the decoder sca
 
 ## Post-Processing
 
+Trace paths ending in `.gz` are written compressed. Overlapping duration slices on the same Perfetto track are split into adjacent `#0`, `#1`, ... lanes by default; pass `split_overlaps=False` to `profile_session` or `events_to_perfetto` to keep raw tracks.
+
 You can pass callbacks to `profile_session` to mutate events or the Perfetto trace before writing:
 
 - `post_process_events(events, ctx) -> events`: Rename, filter, or regroup events.

transformer_nuggets/cute/profiler/__init__.py

@@ -21,7 +21,7 @@ def my_kernel(output, prof_buf, max_events):
         max_events_per_unit=64,
         num_units=(num_blocks, "Block"),
         tag_names=["compute"],
-        trace_path="trace.json",
+        trace_path="trace.json.gz",
     ) as (prof, tag_table):
         my_kernel(output, prof.tensor, prof.max_events_per_unit)
 

transformer_nuggets/cute/profiler/host.py

@@ -22,23 +22,24 @@
     self.kernel(x_c, prof.to_cute(), prof.max_events_per_unit, produce_tag, consume_tag).launch(...)
 
     events = decode_events(prof, tag_table)
-    events_to_perfetto(events, "trace.json", pid=0)
+    events_to_perfetto(events, "trace.json.gz", pid=0)
 
 Or use the context manager:
-    with profile_session(64, num_units=4, tag_names=["produce", "consume"], trace_path="trace.json") as (prof, tags):
+    with profile_session(64, num_units=4, tag_names=["produce", "consume"], trace_path="trace.json.gz") as (prof, tags):
         produce_tag = tags.id("produce")
         self.kernel(x_c, prof.to_cute(), prof.max_events_per_unit, produce_tag).launch(...)
 """
 
 from __future__ import annotations
 
-import json
 from contextlib import contextmanager
 from dataclasses import dataclass
 from collections.abc import Callable, Iterator
 
 import torch
 
+from transformer_nuggets.utils.perfetto import split_overlapping_slices, write_trace
+
 try:
     import cutlass.cute as cute
     from cutlass.cute.runtime import from_dlpack
@@ -305,6 +306,7 @@ def events_to_perfetto(
     pid: int = 0,
     tid_prefix: int | None = None,
     unit_name: str = "Unit",
+    split_overlaps: bool = True,
 ) -> dict:
     """Export events to Perfetto-compatible Chrome trace JSON format.
 
@@ -323,15 +325,16 @@ def events_to_perfetto(
         tid_prefix: Optional prefix for thread IDs. If provided, tid = tid_prefix + event.tid.
                     This can satisfy Perfetto's pid/tid quirk where tid should start with pid.
         unit_name: Name for units in trace (e.g., "Block", "Warp"). Defaults to "Unit".
+        split_overlaps: If true, split overlapping duration slices on the same
+            track into adjacent lanes before returning/writing the trace.
 
     Returns:
         The trace dict (can be further modified before writing).
     """
     if not events:
         trace = {"traceEvents": []}
         if trace_path is not None:
-            with open(trace_path, "w") as f:
-                json.dump(trace, f)
+            write_trace(trace_path, trace)
         return trace
 
     trace_events = []
@@ -412,9 +415,11 @@ def events_to_perfetto(
         "traceEvents": trace_events,
     }
 
+    if split_overlaps:
+        trace = split_overlapping_slices(trace)
+
     if trace_path is not None:
-        with open(trace_path, "w") as f:
-            json.dump(trace, f, indent=2)
+        write_trace(trace_path, trace)
 
     return trace
 
@@ -430,6 +435,7 @@ def profile_session(
     pid: int = 0,
     post_process_events: Callable[[list[Event], PostProcessContext], list[Event]] | None = None,
     post_process_trace: Callable[[dict, PostProcessContext], dict] | None = None,
+    split_overlaps: bool = True,
 ) -> Iterator[tuple[ProfileBuf, TagTable]]:
     """Context manager for profiling a kernel session.
 
@@ -441,7 +447,7 @@ def profile_session(
             max_events_per_unit=64,
             num_units=(num_blocks, "Block"),  # Named units for nicer traces
             tag_names=["produce", "consume"],
-            trace_path="trace.json"
+            trace_path="trace.json.gz"
         ) as (prof, tags):
             TAG_PRODUCE = tags.id("produce")
             TAG_CONSUME = tags.id("consume")
@@ -475,6 +481,8 @@ def group_by_block_warp(events, ctx):
             Signature: (events, context) -> events. Can rename, filter, or regroup events.
         post_process_trace: Optional callback to mutate the Perfetto trace dict before writing.
             Signature: (trace_dict, context) -> trace_dict. Can add flow events, counters, etc.
+        split_overlaps: If true, split overlapping duration slices on the same
+            Perfetto track into adjacent lanes before writing the trace.
 
     Yields:
         Tuple of (ProfileBuf, TagTable).
@@ -506,10 +514,18 @@ def group_by_block_warp(events, ctx):
         events = post_process_events(events, ctx)
 
     if trace_path is not None:
-        trace = events_to_perfetto(events, trace_path=None, pid=pid, unit_name=prof.unit_name)
+        trace = events_to_perfetto(
+            events,
+            trace_path=None,
+            pid=pid,
+            unit_name=prof.unit_name,
+            split_overlaps=False,
+        )
 
         if post_process_trace is not None:
             trace = post_process_trace(trace, ctx)
 
-        with open(trace_path, "w") as f:
-            json.dump(trace, f, indent=2)
+        if split_overlaps:
+            trace = split_overlapping_slices(trace)
+
+        write_trace(trace_path, trace)

transformer_nuggets/utils/__init__.py

@@ -17,4 +17,10 @@
 from transformer_nuggets.utils.triton import print_sass
 from transformer_nuggets.utils.merge_traces import merge_traces
 from transformer_nuggets.utils.memory_viz import generate_memory_comparison_html
+from transformer_nuggets.utils.perfetto import (
+    default_trace_path,
+    read_trace,
+    split_overlapping_slices,
+    write_trace,
+)
 # from transformer_nuggets.utils.model_extraction import extract_attention_data

transformer_nuggets/utils/benchmark.py

@@ -19,6 +19,13 @@
 from torch.cuda._memory_viz import profile_plot  # type: ignore
 from torch.profiler import profile, ProfilerActivity, record_function, schedule
 
+from transformer_nuggets.utils.perfetto import (
+    default_trace_path,
+    read_trace,
+    split_overlapping_slices,
+    write_trace,
+)
+
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
@@ -161,6 +168,9 @@ class ProfileConfig:
     extra_kwargs: dict = field(default_factory=dict)
     memory_profile_path: str | None = None
     row_limit: int = 10
+    gzip_trace: bool = True
+    fix_overlapping_events: bool = True
+    overlap_track_pattern: str | None = "stream.*"
 
 
 @dataclass(frozen=True)
@@ -521,8 +531,24 @@ def profile_function(
                     torch.cuda.synchronize()
 
     if config.file_path is not None:
-        trace_path = Path(config.file_path).with_suffix(".json")
-        prof.export_chrome_trace(str(trace_path))
+        trace_path = default_trace_path(config.file_path, gzip_by_default=config.gzip_trace)
+        export_path = trace_path
+        if trace_path.suffix == ".gz":
+            export_path = trace_path.with_name(f"{trace_path.name}.tmp.json")
+
+        prof.export_chrome_trace(str(export_path))
+
+        if config.fix_overlapping_events or export_path != trace_path:
+            trace = read_trace(export_path)
+            if config.fix_overlapping_events:
+                trace = split_overlapping_slices(
+                    trace,
+                    track_pattern=config.overlap_track_pattern,
+                )
+            write_trace(trace_path, trace)
+            if export_path != trace_path:
+                export_path.unlink()
+
         logger.info(f"💾 Trace file 📄 saved to: {bcolors.OKGREEN}{trace_path}{bcolors.ENDC}")
 
     if profile_memory and config.memory_profile_path is not None: