warpscope 0.1.0__tar.gz

warpscope-0.1.0/LICENSE

@@ -0,0 +1,37 @@
+MIT License
+
+Copyright (c) 2026 warpscope contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+Credits / Attribution
+
+This project's in-kernel CUDA profiling design and wire format originate from:
+
+  * 侯博涵 (Hou Bohan) — original idea and write-up:
+    https://zhuanlan.zhihu.com/p/2054305616391304228
+
+  * Apache TVM, TIRx `CudaProfiler` (Apache License 2.0) — wire format and the
+    host-side decode / Perfetto export logic are adapted from:
+    https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py
+    https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html
+
+The CUDA implementation in this repository was written by Claude Opus (Anthropic).

warpscope-0.1.0/PKG-INFO

@@ -0,0 +1,156 @@
+Metadata-Version: 2.4
+Name: warpscope
+Version: 0.1.0
+Summary: In-kernel %globaltimer profiler for warp-specialized CUDA kernels (Perfetto/Chrome timelines).
+Author: warpscope contributors
+License: MIT
+Project-URL: Homepage, https://github.com/YangWang92/warpscope
+Keywords: cuda,profiler,globaltimer,perfetto,gpu,warp-specialized,tracing
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.21
+Provides-Extra: torch
+Requires-Dist: torch; extra == "torch"
+Provides-Extra: perfetto
+Requires-Dist: tg4perfetto; extra == "perfetto"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+# warpscope
+
+English | [中文](README.zh.md)
+
+> Credits: the idea and design all come from **侯博涵 (Hou Bohan)**'s write-up
+> ([zhihu](https://zhuanlan.zhihu.com/p/2054305616391304228)); the wire format and the
+> host-side decode / Perfetto export are adapted from **Apache TVM TIRx `CudaProfiler`**
+> ([bench.py](https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py),
+> [docs](https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html)).
+> The implementation in this repo was **written by Claude Opus (Anthropic)**.
+
+In-kernel `%globaltimer` profiler for **warp-specialized CUDA kernels**. Bracket the
+logical stages inside a kernel (TMA load, MMA, softmax, epilogue, ...) with
+`start`/`end` markers; one leader thread per logical group stamps the GPU global timer
+into a buffer you pass as an ordinary kernel argument. Decode it on the host into
+per-`(block, group)` durations or a **Perfetto / Chrome trace** to see how the
+producer and consumer warp-groups actually overlap — something total launch time and
+SM-level counters can't show.
+
+It is **not** zero cost (a timer read + a global store + a block fence per event), so
+it is a debugging/analysis tool. Build with the profiler disabled for production.
+
+## Layout
+
+```
+warpscope/
+  include/
+    warpscope.cuh         # device header (header-only, NVRTC-safe)
+    warpscope_host.hpp    # host decoder + Chrome-trace writer (header-only, pure C++)
+  *.py                    # Python: Profiler buffer mgmt, decode, trace export
+examples/                 # toy CUDA program (pure C++ path) + python driver
+tests/                    # wire-format + decode tests
+```
+
+## Install
+
+```bash
+pip install -e .            # core (numpy only)
+pip install -e ".[torch]"   # + GPU buffer allocation
+pip install -e ".[dev]"     # + pytest
+```
+
+## Device side (CUDA C++)
+
+```cpp
+#include <warpscope.cuh>
+enum : uint32_t { EvWait = 0, EvWork = 1 };
+
+__global__ void k(..., uint64_t* prof, uint32_t stride,
+                  uint32_t num_groups, uint32_t num_blocks, uint32_t max_rec) {
+    ws::Profiler<true> p;     // <false> compiles to a no-op
+    const uint32_t warp = threadIdx.x / 32, lane = threadIdx.x % 32;
+
+    if (warp == 0) {                                  // e.g. TMA producer = group 0
+        p.init(prof, stride, /*group=*/0, num_groups, num_blocks,
+               /*leader=*/lane == 0, max_rec);
+        { WS_REGION(p, EvWait); /* barrier wait */ }  // RAII start/end
+        { WS_REGION(p, EvWork); /* issue work  */ }
+        p.finalize();
+    }
+    // ... other warp-groups: init with their own group id + one leader each ...
+}
+```
+
+Build: `nvcc -I"$(warpscope --include)" -arch=sm_100a my.cu`
+
+## Host side — pick one
+
+**Pure C++ (header-only, no Python):**
+
+```cpp
+#include <warpscope_host.hpp>
+std::vector<uint64_t> h(slots);                 // cudaMemcpy buffer back into h
+ws::write_chrome_trace(h.data(), h.size(),
+    /*events*/ {"wait", "work"},
+    /*groups*/ {"tma", "umma", "utccp", "epilogue"},
+    "trace.json");                              // open in chrome://tracing or perfetto
+```
+
+**Python:**
+
+```python
+import warpscope as ws
+prof = ws.Profiler(num_blocks=num_sms, num_groups=4, max_records_per_lane=64)
+launch(..., prof.ptr)        # pass the device pointer
+torch.cuda.synchronize()
+res = prof.decode(event_names={0: "wait", 1: "work"},
+                  group_names={0: "tma", 1: "umma", 2: "utccp", 3: "epilogue"})
+res.print_durations()
+res.to_perfetto("trace.json")   # Chrome JSON; opens in ui.perfetto.dev too
+```
+
+## Output
+
+The raw output is a `uint64` buffer. Both host paths turn it into a **Chrome Trace
+Event JSON** file (`pid = block`, `tid = group`, `ts/dur` in microseconds) that opens
+directly in `chrome://tracing` and <https://ui.perfetto.dev>. A native
+`.perfetto-trace` writer is available via the optional `tg4perfetto` dependency.
+
+## Wire format (v1, shared ABI)
+
+```
+record = (globaltimer_lo32 << 32) | tag32
+tag32  = (block_group << 12) | (event_id << 2) | event_type
+block_group = block_idx * num_groups + group_id
+event_type : 0=begin 1=end 2=instant 3=finalize
+buf[0] header = (num_groups << 32) | num_blocks
+```
+
+Identical to the format used by TIRx/flashinfer, so traces are cross-tool compatible.
+
+## Caveats
+
+- Zero the buffer before launch (the decoder treats 0 as empty).
+- Exactly one leader thread per `(block, group)` lane (two writers clobber the cursor).
+- `%globaltimer_lo` is 32-bit ns: ~tens-of-ns resolution and a ~4.29 s wrap.
+- Persistent grids stream records — cap with `max_records_per_lane` (host) which is
+  also enforced device-side via `init(..., max_records_per_lane=...)`.
+- The fence + store perturb tight pipelines; keep events coarse and compare against an
+  unprofiled (`ws::Profiler<false>`) build.
+
+## Credits & License
+
+Licensed under the **MIT License** (see [LICENSE](LICENSE)).
+
+- **侯博涵 (Hou Bohan)** — original idea and write-up:
+  <https://zhuanlan.zhihu.com/p/2054305616391304228>
+- **Apache TVM TIRx `CudaProfiler`** (Apache-2.0) — wire format + host decode/Perfetto
+  export are adapted from it:
+  <https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py> ·
+  <https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html>
+- The CUDA/Python implementation in this repository was **written by Claude Opus (Anthropic)**.

warpscope-0.1.0/README.md

@@ -0,0 +1,132 @@
+# warpscope
+
+English | [中文](README.zh.md)
+
+> Credits: the idea and design all come from **侯博涵 (Hou Bohan)**'s write-up
+> ([zhihu](https://zhuanlan.zhihu.com/p/2054305616391304228)); the wire format and the
+> host-side decode / Perfetto export are adapted from **Apache TVM TIRx `CudaProfiler`**
+> ([bench.py](https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py),
+> [docs](https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html)).
+> The implementation in this repo was **written by Claude Opus (Anthropic)**.
+
+In-kernel `%globaltimer` profiler for **warp-specialized CUDA kernels**. Bracket the
+logical stages inside a kernel (TMA load, MMA, softmax, epilogue, ...) with
+`start`/`end` markers; one leader thread per logical group stamps the GPU global timer
+into a buffer you pass as an ordinary kernel argument. Decode it on the host into
+per-`(block, group)` durations or a **Perfetto / Chrome trace** to see how the
+producer and consumer warp-groups actually overlap — something total launch time and
+SM-level counters can't show.
+
+It is **not** zero cost (a timer read + a global store + a block fence per event), so
+it is a debugging/analysis tool. Build with the profiler disabled for production.
+
+## Layout
+
+```
+warpscope/
+  include/
+    warpscope.cuh         # device header (header-only, NVRTC-safe)
+    warpscope_host.hpp    # host decoder + Chrome-trace writer (header-only, pure C++)
+  *.py                    # Python: Profiler buffer mgmt, decode, trace export
+examples/                 # toy CUDA program (pure C++ path) + python driver
+tests/                    # wire-format + decode tests
+```
+
+## Install
+
+```bash
+pip install -e .            # core (numpy only)
+pip install -e ".[torch]"   # + GPU buffer allocation
+pip install -e ".[dev]"     # + pytest
+```
+
+## Device side (CUDA C++)
+
+```cpp
+#include <warpscope.cuh>
+enum : uint32_t { EvWait = 0, EvWork = 1 };
+
+__global__ void k(..., uint64_t* prof, uint32_t stride,
+                  uint32_t num_groups, uint32_t num_blocks, uint32_t max_rec) {
+    ws::Profiler<true> p;     // <false> compiles to a no-op
+    const uint32_t warp = threadIdx.x / 32, lane = threadIdx.x % 32;
+
+    if (warp == 0) {                                  // e.g. TMA producer = group 0
+        p.init(prof, stride, /*group=*/0, num_groups, num_blocks,
+               /*leader=*/lane == 0, max_rec);
+        { WS_REGION(p, EvWait); /* barrier wait */ }  // RAII start/end
+        { WS_REGION(p, EvWork); /* issue work  */ }
+        p.finalize();
+    }
+    // ... other warp-groups: init with their own group id + one leader each ...
+}
+```
+
+Build: `nvcc -I"$(warpscope --include)" -arch=sm_100a my.cu`
+
+## Host side — pick one
+
+**Pure C++ (header-only, no Python):**
+
+```cpp
+#include <warpscope_host.hpp>
+std::vector<uint64_t> h(slots);                 // cudaMemcpy buffer back into h
+ws::write_chrome_trace(h.data(), h.size(),
+    /*events*/ {"wait", "work"},
+    /*groups*/ {"tma", "umma", "utccp", "epilogue"},
+    "trace.json");                              // open in chrome://tracing or perfetto
+```
+
+**Python:**
+
+```python
+import warpscope as ws
+prof = ws.Profiler(num_blocks=num_sms, num_groups=4, max_records_per_lane=64)
+launch(..., prof.ptr)        # pass the device pointer
+torch.cuda.synchronize()
+res = prof.decode(event_names={0: "wait", 1: "work"},
+                  group_names={0: "tma", 1: "umma", 2: "utccp", 3: "epilogue"})
+res.print_durations()
+res.to_perfetto("trace.json")   # Chrome JSON; opens in ui.perfetto.dev too
+```
+
+## Output
+
+The raw output is a `uint64` buffer. Both host paths turn it into a **Chrome Trace
+Event JSON** file (`pid = block`, `tid = group`, `ts/dur` in microseconds) that opens
+directly in `chrome://tracing` and <https://ui.perfetto.dev>. A native
+`.perfetto-trace` writer is available via the optional `tg4perfetto` dependency.
+
+## Wire format (v1, shared ABI)
+
+```
+record = (globaltimer_lo32 << 32) | tag32
+tag32  = (block_group << 12) | (event_id << 2) | event_type
+block_group = block_idx * num_groups + group_id
+event_type : 0=begin 1=end 2=instant 3=finalize
+buf[0] header = (num_groups << 32) | num_blocks
+```
+
+Identical to the format used by TIRx/flashinfer, so traces are cross-tool compatible.
+
+## Caveats
+
+- Zero the buffer before launch (the decoder treats 0 as empty).
+- Exactly one leader thread per `(block, group)` lane (two writers clobber the cursor).
+- `%globaltimer_lo` is 32-bit ns: ~tens-of-ns resolution and a ~4.29 s wrap.
+- Persistent grids stream records — cap with `max_records_per_lane` (host) which is
+  also enforced device-side via `init(..., max_records_per_lane=...)`.
+- The fence + store perturb tight pipelines; keep events coarse and compare against an
+  unprofiled (`ws::Profiler<false>`) build.
+
+## Credits & License
+
+Licensed under the **MIT License** (see [LICENSE](LICENSE)).
+
+- **侯博涵 (Hou Bohan)** — original idea and write-up:
+  <https://zhuanlan.zhihu.com/p/2054305616391304228>
+- **Apache TVM TIRx `CudaProfiler`** (Apache-2.0) — wire format + host decode/Perfetto
+  export are adapted from it:
+  <https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py> ·
+  <https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html>
+- The CUDA/Python implementation in this repository was **written by Claude Opus (Anthropic)**.

warpscope-0.1.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["setuptools>=64", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "warpscope"
+version = "0.1.0"
+description = "In-kernel %globaltimer profiler for warp-specialized CUDA kernels (Perfetto/Chrome timelines)."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [{ name = "warpscope contributors" }]
+keywords = ["cuda", "profiler", "globaltimer", "perfetto", "gpu", "warp-specialized", "tracing"]
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Software Development :: Libraries",
+    "Environment :: GPU :: NVIDIA CUDA",
+]
+dependencies = ["numpy>=1.21"]
+
+[project.optional-dependencies]
+torch = ["torch"]
+perfetto = ["tg4perfetto"]
+dev = ["pytest"]
+
+[project.scripts]
+warpscope = "warpscope.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/YangWang92/warpscope"
+
+[tool.setuptools]
+packages = ["warpscope"]
+
+[tool.setuptools.package-data]
+warpscope = ["include/*.cuh", "include/*.hpp"]

warpscope-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

warpscope-0.1.0/tests/test_decode_golden.py

@@ -0,0 +1,84 @@
+"""Build a synthetic buffer the way the device would, then check decode + export."""
+
+import json
+import os
+
+import numpy as np
+
+from warpscope import decode
+from warpscope._wire import (
+    EVENT_TYPE_BEGIN,
+    EVENT_TYPE_END,
+    EVENT_TYPE_FINALIZE,
+    pack_header,
+    pack_record,
+    pack_tag,
+)
+
+
+def _build_buffer(num_blocks, num_groups, lane_records):
+    """lane_records: {(block, group): [(event_id, type, ts), ...]} in cursor order."""
+    write_stride = num_blocks * num_groups
+    max_rec = max((len(v) for v in lane_records.values()), default=0) + 1
+    buf = np.zeros(1 + write_stride * max_rec, dtype=np.uint64)
+    buf[0] = pack_header(num_groups, num_blocks)
+    for (block, group), recs in lane_records.items():
+        bg = block * num_groups + group
+        cursor = 1 + bg
+        for (ev, typ, ts) in recs:
+            tag = pack_tag(block, group, num_groups, ev, typ)
+            buf[cursor] = pack_record(ts, tag)
+            cursor += write_stride
+    return buf
+
+
+def test_decode_two_groups():
+    lanes = {
+        (0, 0): [(0, EVENT_TYPE_BEGIN, 100), (0, EVENT_TYPE_END, 196),
+                 (1, EVENT_TYPE_BEGIN, 200), (1, EVENT_TYPE_END, 3240),
+                 (0, EVENT_TYPE_FINALIZE, 3300)],
+        (0, 1): [(1, EVENT_TYPE_BEGIN, 200), (1, EVENT_TYPE_END, 11016),
+                 (0, EVENT_TYPE_FINALIZE, 11100)],
+    }
+    buf = _build_buffer(num_blocks=1, num_groups=2, lane_records=lanes)
+    res = decode(buf, event_names={0: "load", 1: "compute"},
+                 group_names={0: "light", 1: "heavy"})
+
+    assert res.num_blocks == 1 and res.num_groups == 2
+    durs = {(s.group, s.event_id): s.dur_ns for s in res.spans}
+    assert durs[(0, 0)] == 96      # group 0 load
+    assert durs[(0, 1)] == 3040    # group 0 compute
+    assert durs[(1, 1)] == 10816   # group 1 compute (heavier)
+
+
+def test_finalize_drops_later_records():
+    lanes = {
+        (0, 0): [(0, EVENT_TYPE_BEGIN, 10), (0, EVENT_TYPE_END, 20),
+                 (0, EVENT_TYPE_FINALIZE, 30),
+                 (0, EVENT_TYPE_BEGIN, 40), (0, EVENT_TYPE_END, 99)],
+    }
+    buf = _build_buffer(1, 1, lanes)
+    res = decode(buf)
+    assert len(res.spans) == 1
+    assert res.spans[0].dur_ns == 10
+
+
+def test_timestamp_wrap():
+    lanes = {(0, 0): [(0, EVENT_TYPE_BEGIN, 0xFFFFFFF0), (0, EVENT_TYPE_END, 0x10)]}
+    buf = _build_buffer(1, 1, lanes)
+    res = decode(buf)
+    assert res.spans[0].dur_ns == 0x20  # (0x10 - 0xFFFFFFF0) & 0xFFFFFFFF
+
+
+def test_chrome_trace_export(tmp_path):
+    lanes = {(0, 0): [(0, EVENT_TYPE_BEGIN, 1000), (0, EVENT_TYPE_END, 2000)]}
+    buf = _build_buffer(1, 1, lanes)
+    res = decode(buf, event_names={0: "work"}, group_names={0: "main"})
+    out = os.path.join(tmp_path, "t.json")
+    res.to_perfetto(out)
+    with open(out) as f:
+        data = json.load(f)
+    xs = [e for e in data["traceEvents"] if e["ph"] == "X"]
+    assert len(xs) == 1
+    assert xs[0]["name"] == "work"
+    assert xs[0]["ts"] == 1.0 and xs[0]["dur"] == 1.0  # microseconds

warpscope-0.1.0/tests/test_wire_roundtrip.py

@@ -0,0 +1,47 @@
+import numpy as np
+import pytest
+
+from warpscope._wire import (
+    EVENT_TYPE_BEGIN,
+    EVENT_TYPE_END,
+    decode_tag,
+    pack_header,
+    pack_record,
+    pack_tag,
+)
+
+
+def test_tag_roundtrip():
+    for block in (0, 1, 7, 147):
+        for group in range(4):
+            for ev in (0, 1, 9, 1023):
+                for typ in range(4):
+                    tag = pack_tag(block, group, 4, ev, typ)
+                    b, g, e, t = decode_tag(tag, 4)
+                    assert (b, g, e, t) == (block, group, ev, typ)
+
+
+def test_record_layout():
+    tag = pack_tag(2, 1, 4, 6, EVENT_TYPE_BEGIN)
+    rec = pack_record(0x1234ABCD, tag)
+    assert rec >> 32 == 0x1234ABCD
+    assert rec & 0xFFFFFFFF == tag
+
+
+def test_header():
+    h = pack_header(num_groups=4, num_blocks=148)
+    assert h >> 32 == 4
+    assert h & 0xFFFFFFFF == 148
+
+
+def test_event_id_overflow():
+    with pytest.raises(ValueError):
+        pack_tag(0, 0, 1, 1024, EVENT_TYPE_END)
+
+
+def test_view_as_uint64():
+    # records must survive an int64<->uint64 view (torch stores as int64)
+    tag = pack_tag(1, 0, 2, 3, EVENT_TYPE_END)
+    rec = pack_record(0xFFFFFFFF, tag)
+    arr = np.array([rec], dtype=np.uint64)
+    assert int(arr.view(np.int64).view(np.uint64)[0]) == rec

warpscope-0.1.0/warpscope/__init__.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 warpscope contributors
+#
+# Credits:
+#   - Original idea & design: 侯博涵 (Hou Bohan)
+#       https://zhuanlan.zhihu.com/p/2054305616391304228
+#   - Wire format / decode / Perfetto export adapted from Apache TVM TIRx CudaProfiler:
+#       https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py
+#       https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html
+#   - Implementation written by Claude Opus (Anthropic).
+"""warpscope - in-kernel %globaltimer profiler for warp-specialized CUDA kernels.
+
+Bracket logical stages inside a CUDA kernel with the device header
+``warpscope.cuh``, pass a zeroed uint64 buffer, then decode it here into per-
+(block, group) durations or a Perfetto/Chrome trace file.
+"""
+
+from ._decode import Instant, Span, TraceResult, decode
+from ._wire import (
+    EVENT_TYPE_BEGIN,
+    EVENT_TYPE_END,
+    EVENT_TYPE_FINALIZE,
+    EVENT_TYPE_INSTANT,
+    WIRE_VERSION,
+    decode_tag,
+    pack_tag,
+)
+from .buffer import Profiler
+from .cli import include_dir
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Profiler",
+    "decode",
+    "TraceResult",
+    "Span",
+    "Instant",
+    "include_dir",
+    "decode_tag",
+    "pack_tag",
+    "WIRE_VERSION",
+    "EVENT_TYPE_BEGIN",
+    "EVENT_TYPE_END",
+    "EVENT_TYPE_INSTANT",
+    "EVENT_TYPE_FINALIZE",
+    "__version__",
+]

warpscope-0.1.0/warpscope/_decode.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 warpscope contributors
+#
+# Credits:
+#   - Original idea & design: 侯博涵 (Hou Bohan)
+#       https://zhuanlan.zhihu.com/p/2054305616391304228
+#   - Decode logic adapted from Apache TVM TIRx CudaProfiler (Apache-2.0):
+#       https://github.com/apache/tvm/blob/main/python/tvm/tirx/bench.py
+#       https://tvm.apache.org/docs/tirx/native_basics/cuda/profiling.html
+#   - Implementation written by Claude Opus (Anthropic).
+"""Decode a warpscope profiler buffer into spans/instants."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import numpy as np
+
+from ._wire import (
+    EVENT_TYPE_BEGIN,
+    EVENT_TYPE_END,
+    EVENT_TYPE_FINALIZE,
+    EVENT_TYPE_INSTANT,
+    decode_tag,
+)
+
+
+@dataclass
+class Span:
+    block: int
+    group: int
+    event_id: int
+    begin_ns: int
+    dur_ns: int
+
+
+@dataclass
+class Instant:
+    block: int
+    group: int
+    event_id: int
+    ts_ns: int
+
+
+@dataclass
+class TraceResult:
+    num_blocks: int
+    num_groups: int
+    spans: List[Span] = field(default_factory=list)
+    instants: List[Instant] = field(default_factory=list)
+    event_names: Optional[Dict[int, str]] = None
+    group_names: Optional[Dict[int, str]] = None
+
+    def event_name(self, event_id: int) -> str:
+        if self.event_names and event_id in self.event_names:
+            return self.event_names[event_id]
+        return f"event_{event_id}"
+
+    def group_name(self, group: int) -> str:
+        if self.group_names and group in self.group_names:
+            return self.group_names[group]
+        return f"group_{group}"
+
+    def print_durations(self) -> None:
+        lanes: Dict[tuple, List[Span]] = {}
+        for s in self.spans:
+            lanes.setdefault((s.block, s.group), []).append(s)
+        for (block, group) in sorted(lanes):
+            parts = ", ".join(f"{self.event_name(s.event_id)}={s.dur_ns}ns" for s in lanes[(block, group)])
+            print(f"block {block} {self.group_name(group)}: {parts}")
+
+    def summary(self) -> Dict[tuple, Dict[str, float]]:
+        """Mean/max duration per (group, event) across all blocks."""
+        agg: Dict[tuple, List[int]] = {}
+        for s in self.spans:
+            agg.setdefault((s.group, s.event_id), []).append(s.dur_ns)
+        out = {}
+        for (group, ev), vals in agg.items():
+            arr = np.asarray(vals, dtype=np.float64)
+            out[(self.group_name(group), self.event_name(ev))] = {
+                "count": int(arr.size),
+                "mean_ns": float(arr.mean()),
+                "max_ns": float(arr.max()),
+            }
+        return out
+
+    def to_perfetto(self, path: str) -> str:
+        """Alias for to_chrome_trace (Perfetto UI also reads Chrome JSON)."""
+        return self.to_chrome_trace(path)
+
+    def to_chrome_trace(self, path: str) -> str:
+        from ._perfetto import write_chrome_trace
+
+        write_chrome_trace(self, path)
+        return path
+
+
+def decode(
+    buf,
+    event_names: Optional[Dict[int, str]] = None,
+    group_names: Optional[Dict[int, str]] = None,
+) -> TraceResult:
+    """Decode a buffer (numpy array / torch tensor / sequence of uint64)."""
+    if hasattr(buf, "detach"):  # torch tensor
+        buf = buf.detach().cpu().numpy()
+    arr = np.ascontiguousarray(buf).view(np.uint64).ravel()
+    if arr.size == 0:
+        return TraceResult(0, 0, event_names=event_names, group_names=group_names)
+
+    header = int(arr[0])
+    num_groups = header >> 32
+    num_blocks = header & 0xFFFFFFFF
+    result = TraceResult(int(num_blocks), int(num_groups),
+                         event_names=event_names, group_names=group_names)
+    if num_groups == 0:
+        return result
+
+    opens: Dict[tuple, int] = {}
+    finished = set()
+    for i in range(1, arr.size):
+        w = int(arr[i])
+        if w == 0:
+            continue
+        ts = w >> 32
+        tag = w & 0xFFFFFFFF
+        block, group, ev, typ = decode_tag(tag, num_groups)
+        if typ == EVENT_TYPE_FINALIZE:
+            finished.add((block, group))
+            continue
+        if (block, group) in finished:
+            continue
+        if typ == EVENT_TYPE_BEGIN:
+            opens[(block, group, ev)] = ts
+        elif typ == EVENT_TYPE_END:
+            t0 = opens.pop((block, group, ev), None)
+            if t0 is not None:
+                result.spans.append(Span(block, group, ev, t0, (ts - t0) & 0xFFFFFFFF))
+        elif typ == EVENT_TYPE_INSTANT:
+            result.instants.append(Instant(block, group, ev, ts))
+    return result