Emit profiler traces as native Perfetto TrackEvents

Author: drisspg

examples/warp_specialized_static.py

@@ -110,7 +110,7 @@ def run():
         device=device,
     )
 
-    trace_path = transformer_nuggets.DATA_DIR / "profiler_warp_static_trace.json"
+    trace_path = transformer_nuggets.DATA_DIR / "profiler_warp_static_trace.pftrace"
 
     process_names = {i: f"CTA {i}" for i in range(NUM_BLOCKS)}
 

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "seaborn",
     "typer",
     "jinja2",
+    "perfetto",
 ]
 
 [project.optional-dependencies]

test/test_perfetto.py

@@ -1,8 +1,12 @@
 import gzip
 import json
 
+import pytest
+
 from transformer_nuggets.utils.perfetto import (
+    chrome_trace_to_track_event_trace,
     default_trace_path,
+    default_track_event_path,
     read_trace,
     split_overlapping_slices,
     write_trace,
@@ -146,5 +150,152 @@ def test_gzip_trace_roundtrip(tmp_path):
 
 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").as_posix() == "foo.json"
     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"
+
+
+def test_default_track_event_path_uses_native_perfetto_suffix():
+    assert default_track_event_path("foo").as_posix() == "foo.pftrace"
+    assert default_track_event_path("foo.json").as_posix() == "foo.pftrace"
+    assert default_track_event_path("foo.json.gz").as_posix() == "foo.pftrace"
+    assert default_track_event_path("foo.pftrace").as_posix() == "foo.pftrace"
+
+
+def test_track_event_conversion_preserves_instants_counters_and_warns_on_unsupported():
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+
+    trace = {
+        "traceEvents": [
+            {"ph": "M", "name": "thread_name", "pid": 0, "tid": 1, "args": {"name": "worker"}},
+            {"ph": "i", "name": "marker", "pid": 0, "tid": 1, "ts": 1},
+            {"ph": "C", "name": "memory", "pid": 0, "tid": 1, "ts": 2, "args": {"bytes": 42}},
+            {"ph": "B", "name": "unsupported", "pid": 0, "tid": 1, "ts": 3},
+        ]
+    }
+
+    with pytest.warns(RuntimeWarning, match="unsupported Chrome trace phases: B"):
+        payload = chrome_trace_to_track_event_trace(trace)
+
+    parsed = Trace()
+    parsed.ParseFromString(payload)
+    event_types = [
+        packet.track_event.type for packet in parsed.packet if packet.HasField("track_event")
+    ]
+    assert TrackEvent.TYPE_INSTANT in event_types
+    assert TrackEvent.TYPE_COUNTER in event_types
+
+
+def test_track_event_conversion_puts_gpu_annotations_on_separate_track():
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+
+    trace = {
+        "traceEvents": [
+            {"ph": "M", "name": "thread_name", "pid": 0, "tid": 7, "args": {"name": "stream 7"}},
+            {"ph": "X", "cat": "kernel", "name": "kernel", "pid": 0, "tid": 7, "ts": 0, "dur": 10},
+            {
+                "ph": "X",
+                "cat": "gpu_user_annotation",
+                "name": "burst_0",
+                "pid": 0,
+                "tid": 7,
+                "ts": 0,
+                "dur": 10,
+            },
+        ]
+    }
+
+    parsed = Trace()
+    parsed.ParseFromString(chrome_trace_to_track_event_trace(trace))
+    names_by_uuid = {
+        packet.track_descriptor.uuid: packet.track_descriptor.name
+        for packet in parsed.packet
+        if packet.HasField("track_descriptor") and packet.track_descriptor.name
+    }
+    event_tracks = {
+        packet.track_event.name: names_by_uuid[packet.track_event.track_uuid]
+        for packet in parsed.packet
+        if packet.HasField("track_event")
+        and packet.track_event.type == TrackEvent.TYPE_SLICE_BEGIN
+    }
+
+    assert event_tracks["kernel"] == "stream 7"
+    assert event_tracks["burst_0"] == "GPU annotations stream 7"
+
+
+def test_track_event_conversion_attaches_paired_flows_to_slices():
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+
+    trace = {
+        "traceEvents": [
+            {"ph": "M", "name": "thread_name", "pid": 0, "tid": 1, "args": {"name": "cpu"}},
+            {"ph": "M", "name": "thread_name", "pid": 0, "tid": 2, "args": {"name": "gpu"}},
+            {"ph": "X", "name": "cudaLaunchKernel", "pid": 0, "tid": 1, "ts": 0, "dur": 10},
+            {"ph": "X", "name": "kernel", "pid": 0, "tid": 2, "ts": 20, "dur": 5},
+            {"ph": "s", "name": "ac2g", "pid": 0, "tid": 1, "ts": 1, "id": 99},
+            {"ph": "f", "name": "ac2g", "pid": 0, "tid": 2, "ts": 20, "id": 99},
+            {"ph": "f", "name": "single-ended-noise", "pid": 0, "tid": 2, "ts": 22, "id": 100},
+        ]
+    }
+
+    parsed = Trace()
+    parsed.ParseFromString(chrome_trace_to_track_event_trace(trace))
+    begins = [
+        packet.track_event
+        for packet in parsed.packet
+        if packet.HasField("track_event")
+        and packet.track_event.type == TrackEvent.TYPE_SLICE_BEGIN
+    ]
+    flow_ids_by_name = {event.name: tuple(event.flow_ids) for event in begins}
+
+    assert flow_ids_by_name["cudaLaunchKernel"] == (99,)
+    assert flow_ids_by_name["kernel"] == (99,)
+    assert all(100 not in flow_ids for flow_ids in flow_ids_by_name.values())
+
+
+def test_track_event_conversion_splits_crossing_slices_and_keeps_nested_slices():
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import (
+        TrackDescriptor,
+        TrackEvent,
+        Trace,
+    )
+
+    trace = {
+        "traceEvents": [
+            {"ph": "M", "name": "thread_name", "pid": 0, "tid": 1, "args": {"name": "stream 1"}},
+            {"ph": "X", "name": "outer", "pid": 0, "tid": 1, "ts": 0, "dur": 10},
+            {"ph": "X", "name": "inner", "pid": 0, "tid": 1, "ts": 2, "dur": 2},
+            {"ph": "X", "name": "crossing", "pid": 0, "tid": 1, "ts": 5, "dur": 10},
+        ]
+    }
+
+    parsed = Trace()
+    parsed.ParseFromString(chrome_trace_to_track_event_trace(trace, track_pattern="stream.*"))
+
+    descriptors_by_uuid = {
+        packet.track_descriptor.uuid: packet.track_descriptor
+        for packet in parsed.packet
+        if packet.HasField("track_descriptor") and packet.track_descriptor.name
+    }
+    names_by_uuid = {uuid: descriptor.name for uuid, descriptor in descriptors_by_uuid.items()}
+    begin_events = [
+        packet.track_event
+        for packet in parsed.packet
+        if packet.HasField("track_event")
+        and packet.track_event.type == TrackEvent.TYPE_SLICE_BEGIN
+    ]
+    event_tracks = {event.name: names_by_uuid[event.track_uuid] for event in begin_events}
+    event_track_uuids = {event.name: event.track_uuid for event in begin_events}
+
+    assert set(names_by_uuid.values()) >= {"stream 1"}
+    assert event_tracks["outer"] == event_tracks["inner"] == event_tracks["crossing"] == "stream 1"
+    assert event_track_uuids["outer"] == event_track_uuids["inner"]
+    assert event_track_uuids["crossing"] != event_track_uuids["outer"]
+
+    outer_desc = descriptors_by_uuid[event_track_uuids["outer"]]
+    crossing_desc = descriptors_by_uuid[event_track_uuids["crossing"]]
+    assert outer_desc.sibling_merge_behavior == (
+        TrackDescriptor.SIBLING_MERGE_BEHAVIOR_BY_SIBLING_MERGE_KEY
+    )
+    assert crossing_desc.sibling_merge_behavior == outer_desc.sibling_merge_behavior
+    assert crossing_desc.sibling_merge_key == outer_desc.sibling_merge_key

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.gz",
+    trace_path="trace.pftrace",
 ) 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/JSON.GZ and split overlapping slices into adjacent lanes by default |
+| `events_to_perfetto(events, path)` | Write native Perfetto TrackEvent `.pftrace` by default, or Chrome JSON/JSON.GZ with `trace_format="chrome_json"` |
 | `TagTable(names)` | Map tag names ↔ integer IDs |
 | `PostProcessContext` | Context passed to post-processing callbacks |
 
@@ -81,9 +81,17 @@ You set `max_events_per_unit` to something larger than you need; the decoder sca
 | `warp_start/warp_stop(...)` | Low-level start/stop (lane 0 of target_warp) |
 | `warp_atomic_alloc(...)` | Allocate event index atomically |
 
+## Trace Formats
+
+Two output formats are supported:
+
+- `track_event` (default): native Perfetto protobuf (`.pftrace` / `.perfetto-trace`). This is the preferred format for programmatically generated traces. Crossing overlaps on one logical track are encoded as multiple backing TrackEvent tracks with the same merge key, so Perfetto can display them as one logical row.
+- `chrome_json`: legacy Chrome JSON/JSON.GZ. Use this only for compatibility with tools that require Chrome JSON. Perfetto handles this format on a best-effort basis and requires duration events on a track to nest cleanly.
+
+Pass `trace_format="chrome_json"` to `profile_session` or `events_to_perfetto` to opt into legacy JSON output. Pass `split_overlaps=False` to keep raw tracks.
+
 ## 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:
 

transformer_nuggets/cute/profiler/__init__.py

@@ -1,7 +1,7 @@
 """NVIDIA Intra-Kernel Profiling for CUTE DSL Kernels.
 
 This package provides utilities for profiling code regions *inside* GPU kernels,
-generating Chrome trace format files viewable in Perfetto (https://ui.perfetto.dev/).
+generating native Perfetto TrackEvent traces viewable in Perfetto (https://ui.perfetto.dev/).
 
 Two modes are supported:
 - Atomic mode: No event_idx needed, indices allocated via atomics (simple)
@@ -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.gz",
+        trace_path="trace.pftrace",
     ) as (prof, tag_table):
         my_kernel(output, prof.tensor, prof.max_events_per_unit)
 

transformer_nuggets/cute/profiler/example.py

@@ -7,7 +7,7 @@
 Run with:
     python -m transformer_nuggets.cute.profiler.example
 
-View the generated trace.json at https://ui.perfetto.dev/
+View the generated .pftrace files at https://ui.perfetto.dev/
 """
 
 import torch
@@ -158,7 +158,7 @@ def run_atomic_mode():
 
     import transformer_nuggets
 
-    trace_path = transformer_nuggets.DATA_DIR / "profiler_atomic_trace.json"
+    trace_path = transformer_nuggets.DATA_DIR / "profiler_atomic_trace.pftrace"
 
     with profile_session(
         max_events_per_unit=3 * NUM_ITERATIONS + 2,
@@ -191,7 +191,7 @@ def run_static_mode():
 
     import transformer_nuggets
 
-    trace_path = transformer_nuggets.DATA_DIR / "profiler_static_trace.json"
+    trace_path = transformer_nuggets.DATA_DIR / "profiler_static_trace.pftrace"
 
     with profile_session(
         max_events_per_unit=3 * NUM_ITERATIONS + 2,