IQ.Pilot Prebuilt Release @ 658635c

tinygrad_repo/extra/sqtt/README.md

@@ -0,0 +1,14 @@
+# SQTT Profiling
+
+## Getting SQ Thread Trace
+
+`VIZ=2` to enable SQTT profiling.
+
+`SQTT_ITRACE_SE_MASK=X` to select shader engines for instruction tracing, -1 = all, 0 = disabled, >0 = SE bitmask, default 0b11.
+
+`SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
+
+## Viewing the traces
+
+- Web UI: `tinygrad/viz/serve.py`
+- Command line: `python -m tinygrad.renderer.amd.sqtt`

tinygrad_repo/extra/sqtt/examples/generate_examples.py

@@ -0,0 +1,27 @@
+import os, subprocess, sys, shlex
+from pathlib import Path
+from tinygrad.helpers import temp, getenv
+
+EXAMPLES_DIR = Path(__file__).parent
+PROFILE_PATH = Path(temp("profile.pkl", append_user=True))
+
+EXAMPLES = {
+  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
+  "plus":"test/test_tiny.py TestTiny.test_plus",
+  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
+  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_lds_sync",
+  "handwritten":"test/amd/test_custom_kernel.py TestCustomKernel.test_handwritten",
+}
+
+if __name__ == "__main__":
+  arch = subprocess.check_output(["python", "-c", "from tinygrad import Device; print(Device['AMD'].arch)"], text=True,
+                                 env={**os.environ, "DEBUG":"0"}).rstrip()
+  (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
+  for name,test in EXAMPLES.items():
+    if getenv("NAME", name) != name: continue
+    for i in range(2):
+      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
+      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
+                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
+      print(f"saved SQTT trace to {dest}")

tinygrad_repo/extra/sqtt/install_rocprof_decoder.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+import os, platform, shutil, subprocess
+from pathlib import Path
+from tinygrad.helpers import fetch, OSX
+
+VERSION = "0.1.6"
+DEST = Path("/usr/local/lib")
+DEST.mkdir(exist_ok=True)
+
+if __name__ == "__main__":
+  if OSX:
+    arch = "arm64" if platform.machine() == "arm64" else "x86_64"
+    dmg = fetch(f"https://github.com/ROCm/rocprof-trace-decoder/releases/download/{VERSION}/rocprof-trace-decoder-macos-{arch}-{VERSION}-Darwin.dmg")
+    mnt = Path(subprocess.check_output(["hdiutil", "attach", "-nobrowse", "-readonly", "-mountrandom", "/tmp", str(dmg)],
+                                       text=True).split("\t")[-1].strip())
+    try: shutil.copy2(next(mnt.rglob("librocprof-trace-decoder.dylib")), DEST)
+    finally: subprocess.run(["hdiutil", "detach", str(mnt)], check=True)
+    lib = DEST/"librocprof-trace-decoder.dylib"
+  else:
+    lib = DEST/"librocprof-trace-decoder.so"
+    os.system(f"sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/{VERSION}/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o {lib}")
+    os.system("sudo ldconfig")
+  print(f"Installed {lib.name} ({VERSION}) to", DEST)

tinygrad_repo/extra/sqtt/rgptool.py

@@ -0,0 +1,347 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+import argparse, ctypes, struct, hashlib, pickle, code, typing, functools
+import tinygrad.runtime.autogen.sqtt as sqtt
+from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+from tinygrad.helpers import round_up, flatten, all_same, temp
+from dataclasses import dataclass
+
+CHUNK_CLASSES = {
+  sqtt.SQTT_FILE_CHUNK_TYPE_ASIC_INFO: sqtt.struct_sqtt_file_chunk_asic_info,
+  sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DESC: sqtt.struct_sqtt_file_chunk_sqtt_desc,
+  sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DATA: sqtt.struct_sqtt_file_chunk_sqtt_data,
+  sqtt.SQTT_FILE_CHUNK_TYPE_API_INFO: sqtt.struct_sqtt_file_chunk_api_info,
+  sqtt.SQTT_FILE_CHUNK_TYPE_QUEUE_EVENT_TIMINGS: sqtt.struct_sqtt_file_chunk_queue_event_timings,
+  sqtt.SQTT_FILE_CHUNK_TYPE_CLOCK_CALIBRATION: sqtt.struct_sqtt_file_chunk_clock_calibration,
+  sqtt.SQTT_FILE_CHUNK_TYPE_CPU_INFO: sqtt.struct_sqtt_file_chunk_cpu_info,
+  sqtt.SQTT_FILE_CHUNK_TYPE_SPM_DB: sqtt.struct_sqtt_file_chunk_spm_db,
+  sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_DATABASE: sqtt.struct_sqtt_file_chunk_code_object_database,
+  sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_LOADER_EVENTS: sqtt.struct_sqtt_file_chunk_code_object_loader_events,
+  sqtt.SQTT_FILE_CHUNK_TYPE_PSO_CORRELATION: sqtt.struct_sqtt_file_chunk_pso_correlation,
+}
+
+def pretty(val, pad=0) -> str:
+  if isinstance(val, (ctypes.Structure, ctypes.Union)):
+    nl = '\n' # old python versions don't support \ in f-strings
+    return f"{val.__class__.__name__}({nl}{' '*(pad+2)}{(f', {nl}'+' '*(pad+2)).join([f'{field[0]}={pretty(getattr(val, field[0]), pad=pad+2)}' for field in val._fields_])}{nl}{' '*pad})"
+  if isinstance(val, ctypes.Array):
+    return f"[{', '.join(map(pretty, val))}]"
+  if isinstance(val, int) and val >= 1024: return hex(val)
+  return repr(val)
+
+@dataclass(frozen=True)
+class RGPChunk:
+  header: sqtt.Structure
+  data: list[typing.Any]|list[tuple[typing.Any, bytes]]|bytes|None = None
+  def print(self):
+    print(pretty(self.header))
+    # if isinstance(self.data, bytes): print(repr(self.data))
+    if isinstance(self.data, list):
+      for dchunk in self.data:
+        if isinstance(dchunk, tuple):
+          print(pretty(dchunk[0]))
+          # print(repr(dchunk[1]))
+        else:
+          print(pretty(dchunk))
+  # TODO: `def fixup` and true immutability
+  def to_bytes(self, offset:int) -> bytes:
+    cid = self.header.header.chunk_id.type
+    match cid:
+      case _ if cid in {sqtt.SQTT_FILE_CHUNK_TYPE_ASIC_INFO, sqtt.SQTT_FILE_CHUNK_TYPE_CPU_INFO, sqtt.SQTT_FILE_CHUNK_TYPE_API_INFO, sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DESC}:
+        self.header.header.size_in_bytes = ctypes.sizeof(self.header)
+        return bytes(self.header)
+      case sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DATA:
+        assert isinstance(self.data, bytes)
+        self.header.header.size_in_bytes = ctypes.sizeof(self.header) + len(self.data)
+        self.header.offset = offset+ctypes.sizeof(self.header)
+        self.header.size = len(self.data)
+        return bytes(self.header) + self.data
+      case sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_DATABASE:
+        assert isinstance(self.data, list)
+        data_codb = typing.cast(list[tuple[sqtt.struct_sqtt_code_object_database_record, bytes]], self.data)
+        ret = bytearray()
+        sz = ctypes.sizeof(self.header)+sum([ctypes.sizeof(record_hdr)+round_up(len(record_blob), 4) for record_hdr,record_blob in data_codb])
+        self.header.header.size_in_bytes = sz
+        self.header.offset = offset
+        self.header.record_count = len(data_codb)
+        self.header.size = sz
+        ret += self.header
+        for record_hdr,record_blob in data_codb:
+          record_hdr.size = round_up(len(record_blob), 4)
+          ret += record_hdr
+          ret += record_blob.ljust(4, b'\x00')
+        return ret
+      case sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_LOADER_EVENTS:
+        assert isinstance(self.data, list)
+        data_lev = typing.cast(list[tuple[sqtt.struct_sqtt_code_object_loader_events_record]], self.data)
+        self.header.header.size_in_bytes = ctypes.sizeof(self.header)+ctypes.sizeof(sqtt.struct_sqtt_code_object_loader_events_record)*len(data_lev)
+        self.header.offset = offset
+        self.header.record_size = ctypes.sizeof(sqtt.struct_sqtt_code_object_loader_events_record)
+        self.header.record_count = len(data_lev)
+        return bytes(self.header) + b''.join(map(bytes, data_lev))
+      case sqtt.SQTT_FILE_CHUNK_TYPE_PSO_CORRELATION:
+        assert isinstance(self.data, list)
+        data_pso = typing.cast(list[tuple[sqtt.struct_sqtt_pso_correlation_record]], self.data)
+        self.header.header.size_in_bytes = ctypes.sizeof(self.header)+ctypes.sizeof(sqtt.struct_sqtt_pso_correlation_record)*len(data_pso)
+        self.header.offset = offset
+        self.header.record_size = ctypes.sizeof(sqtt.struct_sqtt_pso_correlation_record)
+        self.header.record_count = len(data_pso)
+        return bytes(self.header) + b''.join(map(bytes, data_pso))
+      case _: raise NotImplementedError(pretty(self.header))
+
+@dataclass(frozen=True)
+class RGP:
+  header: sqtt.struct_sqtt_file_header
+  chunks: list[RGPChunk]
+  @staticmethod
+  def from_bytes(blob: bytes) -> RGP:
+    file_header = sqtt.struct_sqtt_file_header.from_buffer_copy(blob)
+    assert file_header.magic_number == sqtt.SQTT_FILE_MAGIC_NUMBER and file_header.version_major == sqtt.SQTT_FILE_VERSION_MAJOR
+    i = file_header.chunk_offset
+    chunks = []
+    while i < len(blob):
+      assert i%4==0, hex(i)
+      hdr = sqtt.struct_sqtt_file_chunk_header.from_buffer_copy(blob, i)
+      cid = hdr.chunk_id.type
+      header: ctypes.Structure
+      match cid:
+        case _ if cid in {sqtt.SQTT_FILE_CHUNK_TYPE_RESERVED, sqtt.SQTT_FILE_CHUNK_TYPE_QUEUE_EVENT_TIMINGS, sqtt.SQTT_FILE_CHUNK_TYPE_CLOCK_CALIBRATION, sqtt.SQTT_FILE_CHUNK_TYPE_SPM_DB}:
+          chunk = None
+        case sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_DATABASE:
+          header = sqtt.struct_sqtt_file_chunk_code_object_database.from_buffer_copy(blob, i)
+          j = header.offset + ctypes.sizeof(header)
+          data: list = []
+          while j < header.offset + header.size:
+            rec_hdr: ctypes.Structure = sqtt.struct_sqtt_code_object_database_record.from_buffer_copy(blob, j)
+            data.append((rec_hdr, elf:=blob[j+ctypes.sizeof(rec_hdr):j+ctypes.sizeof(rec_hdr)+rec_hdr.size]))
+            assert elf[:4] == b'\x7fELF', repr(elf[:16])
+            j += ctypes.sizeof(rec_hdr)+rec_hdr.size
+          assert len(data) == header.record_count
+          chunk = RGPChunk(header, data)
+        case sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_LOADER_EVENTS:
+          header = sqtt.struct_sqtt_file_chunk_code_object_loader_events.from_buffer_copy(blob, i)
+          data = [sqtt.struct_sqtt_code_object_loader_events_record.from_buffer_copy(blob, header.offset+ctypes.sizeof(header)+j*header.record_size)
+                                                                                                                for j in range(header.record_count)]
+          chunk = RGPChunk(header, data)
+        case sqtt.SQTT_FILE_CHUNK_TYPE_PSO_CORRELATION:
+          header = sqtt.struct_sqtt_file_chunk_pso_correlation.from_buffer_copy(blob, i)
+          data = [sqtt.struct_sqtt_pso_correlation_record.from_buffer_copy(blob, header.offset+ctypes.sizeof(header)+j*header.record_size)
+                                                                                                                for j in range(header.record_count)]
+          chunk = RGPChunk(header, data)
+        case sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DATA:
+          header = sqtt.struct_sqtt_file_chunk_sqtt_data.from_buffer_copy(blob, i)
+          chunk = RGPChunk(header, blob[header.offset:header.offset+header.size])
+        case _ if cid in {sqtt.SQTT_FILE_CHUNK_TYPE_ASIC_INFO, sqtt.SQTT_FILE_CHUNK_TYPE_CPU_INFO, sqtt.SQTT_FILE_CHUNK_TYPE_API_INFO,
+                          sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DESC}:
+          chunk = RGPChunk(CHUNK_CLASSES[cid].from_buffer_copy(blob, i))
+        case _:
+          chunk = None
+          print(f"unknown chunk id {cid}")
+      if chunk is not None: chunks.append(chunk)
+      i += hdr.size_in_bytes
+    assert i == len(blob), f'{i} != {len(blob)}'
+    return RGP(file_header, chunks)
+  @staticmethod
+  def from_profile(profile_pickled, device:str|None=None):
+    profile: list[ProfileEvent] = pickle.loads(profile_pickled)
+    def _is_base_dev(d): return all(p.isdigit() for p in d.split(":")[1:])
+    device_events = {x.device:x for x in profile if isinstance(x, ProfileDeviceEvent) and x.device.startswith('AMD') and _is_base_dev(x.device)}
+    if device is None:
+      if len(device_events) == 0: raise RuntimeError('No supported devices found in profile')
+      if len(device_events) > 1: raise RuntimeError(f"More than one supported device found, select which one to export: {', '.join(device_events.keys())}")
+      _, device_event = device_events.popitem()
+    else:
+      if device not in device_events: raise RuntimeError(f"Device {device} not found in profile, devices in profile: {', '.join(device_events.keys())} ")
+      device_event = device_events[device]
+    sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
+    device_props = device_event.props
+    # merge events per SE
+    merged_sqtt_events:dict[int, ProfileSQTTEvent] = {}
+    for ev in sqtt_events:
+      if ev.se not in merged_sqtt_events: merged_sqtt_events[ev.se] = ev
+      else:
+        merged_sqtt_events[ev.se] = ProfileSQTTEvent(
+          device=ev.device,
+          kern=ev.kern,
+          se=ev.se,
+          itrace=merged_sqtt_events[ev.se].itrace or ev.itrace,
+          blob=merged_sqtt_events[ev.se].blob + ev.blob,
+          exec_tag=0,
+        )
+    sqtt_events = list(merged_sqtt_events.values())
+
+    if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
+    gfx_ver = device_props['gfx_target_version'] // 10000
+    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
+                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
+    sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
+    sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
+    sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
+    load_events = [x for x in profile if isinstance(x, ProfileProgramEvent) and x.device == device_event.device]
+    loads = [(event.base, struct.unpack('<Q', hashlib.md5(event.lib).digest()[:8])*2) for event in load_events if event.base is not None and event.lib is not None]
+    code_objects = list(dict.fromkeys([x.lib for x in load_events if x.lib is not None]).keys())
+    if len(loads) == 0: raise RuntimeError('No load events in profile')
+    # TODO: tons of stuff hardcoded for 7900xtx
+    file_header = sqtt.struct_sqtt_file_header(
+      magic_number=sqtt.SQTT_FILE_MAGIC_NUMBER,
+      version_major=sqtt.SQTT_FILE_VERSION_MAJOR,
+      version_minor=sqtt.SQTT_FILE_VERSION_MINOR,
+      flags=sqtt.struct_sqtt_file_header_flags(value=1,),
+      chunk_offset=ctypes.sizeof(sqtt.struct_sqtt_file_header),
+    )
+    chunks = [
+      RGPChunk(sqtt.struct_sqtt_file_chunk_cpu_info(
+        header=sqtt.struct_sqtt_file_chunk_header(
+          chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_CPU_INFO),
+          major_version=0, minor_version=0,
+        ),
+        cpu_timestamp_freq=1000000000,
+        clock_speed=2994, # in mhz???
+        num_logical_cores=64,
+        num_physical_cores=32,
+        system_ram_size=256*1024, # in mb???
+      )),
+      RGPChunk(sqtt.struct_sqtt_file_chunk_asic_info(
+        header=sqtt.struct_sqtt_file_chunk_header(
+          chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_ASIC_INFO),
+          major_version=0, minor_version=5,
+        ),
+        flags=0,
+        trace_shader_core_clock=0x93f05080,
+        trace_memory_clock=0x4a723a40,
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
+        device_revision_id=0xc8,
+        vgprs_per_simd=1536,
+        sgprs_per_simd=128*16,
+        shader_engines=device_props['array_count'] // device_props['simd_arrays_per_engine'],
+        compute_unit_per_shader_engine=device_props['simd_count'] // device_props['simd_per_cu'] // (device_props['array_count'] // device_props['simd_arrays_per_engine']),
+        simd_per_compute_unit=device_props['simd_per_cu'],
+        wavefronts_per_simd=device_props['max_waves_per_simd'],
+        minimum_vgpr_alloc=4,
+        vgpr_alloc_granularity=8,
+        minimum_sgpr_alloc=128,
+        sgpr_alloc_granularity=128,
+        hardware_contexts=8,
+        gpu_type=sqtt.SQTT_GPU_TYPE_DISCRETE,
+        gfxip_level=gfx_iplvl,
+        gpu_index=0,
+        gds_size=0,
+        gds_per_shader_engine=0,
+        ce_ram_size=0,
+        ce_ram_size_graphics=0,
+        ce_ram_size_compute=0,
+        max_number_of_dedicated_cus=0,
+        vram_size=24 * 1024 * 1024 * 1024,  # 24 GB
+        vram_bus_width=384, # 384-bit
+        l2_cache_size=6 * 1024 * 1024, # 6 MB
+        l1_cache_size=32 * 1024, # 32 KB per SIMD (?)
+        lds_size=device_props['lds_size_in_kb'] * 1024,
+        gpu_name=b'NAVI31',
+        alu_per_clock=0,
+        texture_per_clock=0,
+        prims_per_clock=6,
+        pixels_per_clock=0,
+        gpu_timestamp_frequency=100000000, # 100 MHz
+        max_shader_core_clock=2500000000, # 2.5 GHz (boost clock)
+        max_memory_clock=1250000000,  # 1.25 GHz
+        memory_ops_per_clock=16,
+        memory_chip_type=sqtt.SQTT_MEMORY_TYPE_GDDR6,
+        lds_granularity=512,
+        cu_mask=((255, 255),)*6 + ((0,0),)*(32-6),
+        gl1_cache_size=256 * 1024, # 256 KB
+        instruction_cache_size=32 * 1024, # 32 KB
+        scalar_cache_size=16 * 1024, # 16 KB
+        mall_cache_size=96 * 1024 * 1024, # 96 MB
+      )),
+      RGPChunk(sqtt.struct_sqtt_file_chunk_api_info(
+        header=sqtt.struct_sqtt_file_chunk_header(
+          chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_API_INFO),
+          major_version=0,
+          minor_version=2,
+        ),
+        api_type=5, # HIP, not in enum
+        major_version=12, minor_version=0,
+        profiling_mode=sqtt.SQTT_PROFILING_MODE_PRESENT,
+        instruction_trace_mode=sqtt.SQTT_INSTRUCTION_TRACE_FULL_FRAME if sqtt_itrace_enabled else sqtt.SQTT_INSTRUCTION_TRACE_DISABLED,
+        instruction_trace_data=sqtt.union_sqtt_instruction_trace_data(
+          shader_engine_filter=sqtt.union_sqtt_instruction_trace_data_shader_engine_filter(mask=sqtt_itrace_se_mask),
+        ),
+      )),
+      *flatten([(
+        RGPChunk(sqtt.struct_sqtt_file_chunk_sqtt_desc(
+          header=sqtt.struct_sqtt_file_chunk_header(
+            chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DESC, index=sqtt_event.se),
+            major_version=0, minor_version=2,
+          ),
+          shader_engine_index=sqtt_event.se,
+          sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
+          v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_v1(
+            instrumentation_spec_version=1,
+            instrumentation_api_version=0,
+            compute_unit_index=0,
+          )
+        )),
+        RGPChunk(sqtt.struct_sqtt_file_chunk_sqtt_data(
+          header=sqtt.struct_sqtt_file_chunk_header(
+            chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_SQTT_DATA, index=sqtt_event.se),
+            major_version=0, minor_version=0,
+          ),
+        ), sqtt_event.blob),
+      ) for sqtt_event in sqtt_events]),
+      RGPChunk(sqtt.struct_sqtt_file_chunk_code_object_database(
+        header=sqtt.struct_sqtt_file_chunk_header(
+          chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_DATABASE),
+          major_version=0, minor_version=0,
+        ),
+      ), [(sqtt.struct_sqtt_code_object_database_record(), lib) for lib in code_objects]),
+      RGPChunk(sqtt.struct_sqtt_file_chunk_code_object_loader_events(
+        header=sqtt.struct_sqtt_file_chunk_header(
+          chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_CODE_OBJECT_LOADER_EVENTS),
+          major_version=1, minor_version=0,
+        ),
+      ), [sqtt.struct_sqtt_code_object_loader_events_record(base_address=base, code_object_hash=hash) for base,hash in loads]),
+      RGPChunk(sqtt.struct_sqtt_file_chunk_pso_correlation(
+        header=sqtt.struct_sqtt_file_chunk_header(
+          chunk_id=sqtt.struct_sqtt_file_chunk_id(type=sqtt.SQTT_FILE_CHUNK_TYPE_PSO_CORRELATION),
+          major_version=0, minor_version=0,
+        ),
+      ), [sqtt.struct_sqtt_pso_correlation_record(api_pso_hash=hash[0], pipeline_hash=hash) for _,hash in loads])
+    ]
+    return RGP(file_header, chunks)
+  def to_bytes(self) -> bytes:
+    ret = bytearray()
+    ret += self.header
+    for chunk in self.chunks:
+      ret += chunk.to_bytes(len(ret))
+    return bytes(ret)
+  def print(self):
+    print(pretty(self.header))
+    for chunk in self.chunks: chunk.print()
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser(prog='rgptool', description='A tool to create (from pickled tinygrad profile), inspect and modify Radeon GPU Profiler files')
+  parser.add_argument('command')
+  parser.add_argument('input', nargs='?', default=temp("profile.pkl", append_user=True))
+  parser.add_argument('-d', '--device')
+  parser.add_argument('-o', '--output')
+  args = parser.parse_args()
+
+  with open(args.input, 'rb') as fd: input_bytes = fd.read()
+
+  match args.command:
+    case 'print':
+      rgp = RGP.from_bytes(input_bytes)
+      rgp.print()
+    case 'create':
+      rgp = RGP.from_profile(input_bytes, device=args.device)
+      # rgp.to_bytes() # fixup
+      # rgp.print()
+    case 'repl':
+      rgp = RGP.from_bytes(input_bytes)
+      code.interact(local=locals())
+    case _: raise RuntimeError(args.command)
+
+  if args.output is not None:
+    with open(args.output, 'wb+') as fd: fd.write(rgp.to_bytes())
+    print(f"Saved to {args.output}")

tinygrad_repo/extra/sqtt/roc.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+import ctypes, pathlib, argparse, pickle, dataclasses, threading, itertools
+from decimal import Decimal
+from typing import Generator
+from tinygrad.helpers import temp, unwrap, DEBUG
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+from tinygrad.runtime.autogen import rocprof
+from tinygrad.renderer.amd.dsl import Inst
+from tinygrad.helpers import ProfileEvent, ProfileRangeEvent, ProfilePointEvent
+from tinygrad.device import ProfileProgramEvent
+from test.amd.disasm import disasm
+
+@dataclasses.dataclass(frozen=True)
+class InstExec:
+  typ:str
+  pc:int
+  stall:int
+  dur:int
+  time:int
+
+@dataclasses.dataclass(frozen=True)
+class WaveSlot:
+  wave_id:int
+  cu:int
+  simd:int
+  se:int
+  @property
+  def cu_loc(self) -> str: return f"SE:{self.se} CU:{self.cu}"
+  @property
+  def wave_loc(self) -> str: return f"{self.cu_loc} SIMD:{self.simd} W:{self.wave_id}"
+
+@dataclasses.dataclass(frozen=True)
+class WaveExec(WaveSlot):
+  begin_time:int
+  end_time:int
+  insts:bytearray
+  def unpack_insts(self) -> Generator[InstExec, None, None]:
+    sz = ctypes.sizeof(struct:=rocprof.rocprofiler_thread_trace_decoder_inst_t)
+    insts_array = (struct*(len(self.insts)//sz)).from_buffer(self.insts)
+    for inst in insts_array:
+      inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst.category)
+      yield InstExec(inst_typ, inst.pc.address, inst.stall, inst.duration, inst.time)
+
+@dataclasses.dataclass(frozen=True)
+class OccEvent(WaveSlot):
+  time:int
+  start:int
+
+RunKey = tuple[str, int]
+
+class _ROCParseCtx:
+  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]]):
+    self.sqtt_evs, self.disasms = iter(sqtt_evs), {k:{k2:(disasm(v2), v2.size()) for k2,v2 in v.items()} for k,v in disasms.items()}
+    self.inst_execs:dict[RunKey, list[WaveExec]] = {}
+    self.occ_events:dict[RunKey, list[OccEvent]] = {}
+
+  def next_sqtt(self):
+    x = next(self.sqtt_evs, None)
+    self.active_run = (x.kern, x.exec_tag) if x is not None else None
+    self.active_se = x.se if x is not None else None
+    self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
+    return self.active_blob
+
+  def on_occupancy_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_occupancy_t):
+    if DEBUG >= 5: print(f"OCC {ev.time=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.wave_id=} {ev.start=}")
+    self.occ_events.setdefault(unwrap(self.active_run), []).append(OccEvent(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.time, ev.start))
+
+  def on_wave_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_wave_t):
+    if DEBUG >= 5: print(f"WAVE {ev.wave_id=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.contexts=} {ev.begin_time=} {ev.end_time=}")
+    # Skip wave events without instruction timings, occupancy events give the start and duration.
+    if ev.instructions_size == 0: return
+
+    insts_blob = bytearray(sz:=ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t))
+    ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
+
+    self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
+                                                                             ev.end_time, insts_blob))
+
+def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]]) -> _ROCParseCtx:
+  ROCParseCtx = _ROCParseCtx(sqtt_evs, disasms)
+
+  @rocprof.rocprof_trace_decoder_se_data_callback_t
+  def copy_cb(buf, buf_size, _):
+    if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
+    buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
+    buf_size[0] = len(prof_info)
+    return len(prof_info)
+
+  @rocprof.rocprof_trace_decoder_trace_callback_t
+  def trace_cb(record_type, events_ptr, n, _):
+    match record_type:
+      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
+        for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
+      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
+        for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
+      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME:
+        if DEBUG >= 5:
+          pairs = [(ev.shader_clock, ev.realtime_clock) for ev in (rocprof.rocprofiler_thread_trace_decoder_realtime_t * n).from_address(events_ptr)]
+          print(f"REALTIME {pairs}")
+      case _:
+        if DEBUG >= 5: print(rocprof.enum_rocprofiler_thread_trace_decoder_record_type_t.get(record_type), events_ptr, n)
+    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
+
+  @rocprof.rocprof_trace_decoder_isa_callback_t
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run)[0]][pc.address]
+
+    # this is the number of bytes to next instruction, set to 0 for end_pgm
+    if instr == "s_endpgm": mem_size_ptr[0] = 0
+    if (max_sz:=size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
+
+    # truncate the instr if it doesn't fit
+    if (str_sz:=len(instr_bytes:=instr.encode()))+1 > max_sz: str_sz = max_sz
+    ctypes.memmove(instr_ptr, instr_bytes, str_sz)
+    size_ptr[0] = str_sz
+
+    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
+
+  exc:Exception|None = None
+  def worker():
+    nonlocal exc
+    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+    except AttributeError as e:
+      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_rocprof_decoder.py to install")
+      exc.__cause__ = e
+  (t:=threading.Thread(target=worker, daemon=True)).start()
+  t.join()
+  if exc is not None:
+    raise exc
+  return ROCParseCtx
+
+def unpack_occ(viz_data, i:int, j:int, key:tuple[str, int], data:list, p:ProfileProgramEvent, target:str) -> dict:
+  from tinygrad.viz.serve import amd_decode, create_step, row_tuple
+  steps = viz_data.ctxs[i]["steps"]
+  if len(steps[j+1:]) > 0: return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
+  base = unwrap(p.base)
+  disasm:dict[int, Inst] = {addr+base:inst for addr,inst in amd_decode(unwrap(p.lib), target).items()}
+  rctx = decode(data, {p.tag:disasm})
+  cu_events:dict[str, list[ProfileEvent]] = {}
+  # ** inst traces
+  wave_insts:dict[str, dict[str, dict]] = {}
+  inst_units:dict[str, itertools.count] = {}
+  for w in rctx.inst_execs.get(key, []):
+    if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
+    n = next(inst_units[u])
+    if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
+    events.append(ProfileRangeEvent(f"SIMD:{w.simd}", loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
+    wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "prg":p, "run_number":n, "loc":loc}
+  # ** occ traces (only WAVESTART/WAVEEND)
+  units:dict[str, itertools.count] = {}
+  wave_start:dict[str, int] = {}
+  for occ in rctx.occ_events.get(key, []):
+    if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
+    if u in inst_units: continue
+    if occ.start: wave_start[u] = occ.time
+    else:
+      if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
+      events.append(ProfileRangeEvent(f"SIMD:{occ.simd}", f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)),Decimal(occ.time)))
+  # ** split graph by CU
+  for cu in sorted(cu_events, key=row_tuple):
+    steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/cu-sqtt", i, len(steps)), depth=1,
+                             data=[ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]))
+    for k in sorted(wave_insts.get(cu, []), key=row_tuple):
+      wd = wave_insts[cu][k]
+      steps.append(create_step(k.replace(cu, ""), ("/amd-sqtt-insts", i, len(steps)), loc=wd["loc"], depth=2,
+                               data={"fxn":unpack_insts, "args":(wd,)}))
+  return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
+
+def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
+  columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
+  inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]
+  # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
+  #           The idle time can be caused by:
+  #             * Arbiter loss
+  #             * Source or destination register dependency
+  #             * Instruction cache miss
+  # Stall:    The total number of cycles the hardware pipe couldn't issue an instruction.
+  # Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
+  prev_instr = (w:=data["wave"]).begin_time
+  pc_to_inst = data["disasm"]
+  start_pc = None
+  rows:dict[int, dict] = {}
+  for pc, inst in pc_to_inst.items():
+    if start_pc is None: start_pc = pc
+    rows[pc] = {"pc":pc-start_pc, "inst":str(inst), "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
+  for e in w.unpack_insts():
+    if not (inst:=rows[e.pc]).get("type"): inst["type"] = str(e.typ).split("_")[-1]
+    inst["hit_count"] += 1
+    inst["dur"] += e.dur
+    inst["stall"] += e.stall
+    inst["hits"]["rows"].append((inst["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
+    prev_instr = max(prev_instr, e.time + e.dur)
+  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
+             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
+
+def print_data(data:dict) -> None:
+  from tabulate import tabulate
+  # plaintext
+  if "src" in data: print(data["src"])
+  # table format
+  elif "cols" in data:
+    print(tabulate([r[:len(data["cols"])] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
+
+def main() -> None:
+  import tinygrad.viz.serve as viz
+  from tinygrad.uop.ops import RewriteTrace
+  data = viz.VizData()
+
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
+                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Kernel to focus on (optional name, default: all kernels)')
+  parser.add_argument('-n', type=int, default=3, metavar="NUM", help='Max traces to print (optional number, default: 3 traces)')
+  args = parser.parse_args()
+
+  with args.profile.open("rb") as f: profile = pickle.load(f)
+
+  viz.get_profile(profile, data=data)
+
+  # List all kernels
+  if args.kernel is None:
+    for c in data.ctxs:
+      print(c["name"])
+      for s in c["steps"]: print("  "+s["name"])
+    return None
+
+  # Find kernel trace
+  trace = next((c for c in data.ctxs if c["name"] == f"SQTT {args.kernel}"), None)
+  if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
+  n = 0
+  for s in trace["steps"]:
+    if "PKTS" in s["name"]: continue
+    print(s["name"])
+    ret = viz.get_render(data, s["query"])
+    print_data(ret)
+    n += 1
+    if n > args.n: break
+
+if __name__ == "__main__":
+  main()