never mind, i don't like this

* write scan in uops

* ops range

* no need for variable

* meh, later

* shorter

examples/llama3.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
-import tiktoken
-from tiktoken.load import load_tiktoken_bpe
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
@@ -12,6 +10,8 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 class Tokenizer:
   pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
   def __init__(self, model_path: str):
+    import tiktoken
+    from tiktoken.load import load_tiktoken_bpe
     mergeable_ranks = load_tiktoken_bpe(model_path)
     self.num_base_tokens = len(mergeable_ranks)
     special_tokens = [

extra/gemm/amd_uop_matmul.py

@@ -4,9 +4,9 @@ from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
 
-N = 4096
+N = getenv("N", 4096)
 M = K = N
-run_count = 5
+run_count = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
@@ -155,14 +155,15 @@ def test_matmul(sink:UOp, N=N):
       ets.append(ei.run(wait=True))
   print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
 
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    tc = (a @ b).realize()
-  with Context(DEBUG=0):
-    err = (hc - tc).square().mean().item()
-  print(f"mean squared error {err}")
-  if err > 1e-06:
-    raise RuntimeError("matmul is wrong!")
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2):
+      tc = (a @ b).realize()
+    with Context(DEBUG=0):
+      err = (hc - tc).square().mean().item()
+    print(f"mean squared error {err}")
+    if err > 1e-06:
+      raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
   test_matmul(hand_spec_kernel3(), N=N)

extra/mmapeak/mmapeak.py

@@ -1,8 +1,10 @@
-import pathlib
+import os, pathlib
+
+# TODO: there is a timing bug without this
+os.environ["AMD_AQL"] = "1"
+
 from tinygrad.device import Device
 from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
-import time
-import os
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
@@ -44,9 +46,9 @@ if __name__=="__main__":
     raise RuntimeError("Error while initiating AMD device")
 
   COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103'}:
-    if DEV.arch == 'gfx1103':
-      NUM_WORKGROUPS = 8
+  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
+    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
+    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 40
     launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
     launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
     launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))

extra/sqtt/active_sqtt_parse.py

@@ -0,0 +1,131 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+os.environ["PROFILE"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+from dataclasses import replace
+import atexit, contextlib
+from tinygrad import Tensor
+from tinygrad.helpers import system, OSX
+from tinygrad.runtime.ops_amd import AMDProgram
+from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
+from tinygrad.device import Device, ProfileDeviceEvent
+
+from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+
+# TODO: should really check for AM driver / USB
+if not OSX:
+  def set_power(x): system(f"sudo /opt/rocm/bin/amd-smi set -l {x}")
+  @atexit.register
+  def reset_power(): set_power("auto")
+  set_power("stable_std")
+
+dev = Device["AMD"]
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[str, list[WaveExec]] = {}
+  yield sqtt
+  events = dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())]
+
+  rctx = decode(events)
+  assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  sqtt.update(rctx.inst_execs)
+
+  for e in events:
+    if isinstance(e, ProfileSQTTEvent):
+      print(replace(e, blob=b''))
+      if e.se == 0:
+        parse_sqtt_print_packets(e.blob)
+
+template = """.text
+.globl matmul
+.p2align 8
+.type matmul,@function
+matmul:
+  INSTRUCTION
+  s_endpgm
+
+.rodata
+.p2align 6
+.amdhsa_kernel matmul
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
+  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
+  .amdhsa_wavefront_size32 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: matmul
+    .symbol: matmul.kd
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 32
+    .max_flat_workgroup_size: 1024
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+...
+.end_amdgpu_metadata
+"""
+
+def run_asm(src):
+  NUM_WORKGROUPS = 1
+  WAVE_SIZE = 32
+  NUM_WAVES = 1
+  t = Tensor.empty(0x1000).realize()
+  buf = t.uop.buffer.ensure_allocated()
+  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
+  dev.compiler.disassemble(lib)
+  fxn = AMDProgram(dev, "matmul", lib)
+  fxn(buf._buf, global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
+
+if __name__ == "__main__":
+  with save_sqtt() as sqtt:
+    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
+    Tensor.empty(1).elu().realize()
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    # what's in v0?
+    run_asm([
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v1, 0",
+      "s_clause 0x1",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+    ]+[
+      "global_load_b32 v1, v0, s[0:1]",
+    ]*10+[
+      "global_load_b32 v10, v1, s[0:1]",
+      "s_waitcnt vmcnt(0)",
+
+      #"v_rcp_f32 v1, v0"
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v5 v4 v4",
+      #"v_add_f32_e32 v7 v6 v6",
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v2 v1 v1",
+      #"s_nop 1"
+    ]*5+[
+      "v_add_f32_e32 v3 v2 v2",
+    ]*5+[
+      "v_mul_f32_e32 v3 v2 v2",
+    ]*7)

extra/sqtt/attempt_sqtt_parse.py

@@ -1,44 +1,57 @@
 import pickle
-from hexdump import hexdump
-from extra.sqtt.roc import decode, ProfileSQTTEvent
 from tinygrad.helpers import getenv
+from extra.sqtt.roc import decode, ProfileSQTTEvent
 
 # Instruction packets (one per ISA op)
 # NOTE: these are bad guesses and may be wrong! feel free to update if you know better
+# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
 
 OPCODE_NAMES = {
-  # Small metadata / structural packets (NOT ISA op kinds)
-  0x01: "META_SMALL_ID",          # 12-bit identifier / slot tag
-  0x02: "META_FLAG",              # 1-byte flag/mode (CF/AF/8F/DF...)
-  0x03: "META_SUBEVENT_CODE",     # 1-byte sub-event/classification code
-  0x04: "META_BASE_INDEX_TAG",    # 12-bit base index/tag (..D, 9D, 10D, 58D...)
+  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
+  0x02: "VMEMEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
+  0x03: "ALUEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
+  0x01: "VALUINST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
+  0x06: "WAVERDY",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
+  0x08: "WAVEEND",
+  0x09: "WAVESTART",
+  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
+  0x04: "IMMEDIATE_4",
+  0x05: "IMMEDIATE_5",
+  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there
+  0x14: "REG",
+  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
+  0x12: "EVENT",
+  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
+  0x18: "INST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
+  0x19: "UTILCTR",
 
-  # Instruction / timing / timestamp packets
-  0x0F: "TIME_SHORT_DELTA_PLUS4", # short ts, raw_delta+4
-  0x11: "TIME_WAVE_STATE",        # compact wave timing/stall state record
-  0x14: "INST_EXEC_RECORD",       # per-instruction execution record
-  0x16: "TIME_LONG_OR_MARKER",    # long delta / marker with 6-byte payload
+  # ------------------------------------------------------------------------
+  # 0x07–0x0F: pure timestamp-ish deltas
+  # ------------------------------------------------------------------------
+  0x07: "TS_DELTA_S8_W3",           # shift=8,  width=3  (small delta)
+  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
+  0x0B: "TS_DELTA_S5_W3_A",         # shift=5,  width=3
+  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
+  0x0D: "TS_DELTA_S5_W3_C",         # shift=5,  width=3
+  0x0E: "TS_DELTA_S7_W2",           # shift=7,  width=2
+  0x0F: "TS_DELTA_SHORT_PLUS4",     # short delta; ROCm adds +4 before accumulate
 
-  # State / control / perf snapshots
-  0x09: "CONTROL_CONFIG_32B",     # 32-bit control/config word (bursts of FE88..., C488...)
-  0x15: "PERFCOUNTER_SNAPSHOT",   # perf / TT configuration snapshot (8-byte)
+  # ------------------------------------------------------------------------
+  # 0x10–0x19: timestamps, layout headers, events, perf
+  # ------------------------------------------------------------------------
+  0x10: "PSEUDO_NEED_MORE_BITS",    # not a real packet; decoder refill hint
 
-  # Extra descriptors / events / metrics
-  0x06: "META_DESCRIPTOR_24B",    # 24-bit descriptor (seen in complex kernels like GEMM)
-  0x08: "EVENT_SMALL",            # small in-stream event (5-nibble payload)
-  0x12: "TIME_SECONDARY_METRIC",  # 3-byte secondary timing/latency/perf metric
-  0x18: "EVENT_SMALL_PAYLOAD",    # generic small side-band payload (5 nibbles)
-  0x19: "EVENT_SUMMARY_48B",      # rare 6-byte summary/aggregate metric
+  0x11: "TS_WAVE_STATE_SAMPLE",     # wave stall/termination sample (byte at +10)
+  0x13: "EVT_SMALL_GENERIC",        # same structural family as 0x08/0x12/0x19
 
-  # Pseudo / unknown / not yet observed
-  0x07: "UNK_DELTA",              # unknown
-  0x0A: "UNK_DELTA2",             # unknown
-  0x0B: "UNK_DELTA3",             # unknown
-  0x0C: "UNK_DELTA4",             # unknown
-  0x0D: "UNK_DELTA5",             # unknown
-  0x0E: "UNK_DELTA6",             # unknown
-  0x10: "UNK_PSEUDO",             # not seen; pseudo/placeholder
-  0x17: "UNK_NO_DELTA",           # unknown, likely non-timing event
+  0x15: "PERFCOUNTER_SNAPSHOT",     # small delta + 50-ish bits of snapshot
+  0x16: "TS_DELTA36_OR_MARK",       # 36-bit long delta or 36-bit marker
+  0x17: "LAYOUT_MODE_HEADER",       # layout/mode/group + selectors A/B
 }
 
 # these tables are from rocprof trace decoder
@@ -113,17 +126,13 @@ DELTA_MAP_DEFAULT = {
 
 def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
   """
-  Conservative decoding of a few packet types.
-
-  Rules:
-    - We first mask the 64-bit shift register down to the actual packet
-      width using NIBBLE_BUDGET[opcode & 0x1F], so we never read bits
-      that aren't really part of the packet.
-    - Only layouts that are clearly visible from the decompiled C are
-      decoded, and names are kept generic (cfg_*, idx_*, id_*, etc).
+  Decode packet payloads conservatively, using:
+    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
+    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
+    - Per-opcode layouts derived from rocprof's decompiled consumers.
   """
-  # --- 0. Restrict to the real packet bits for this opcode -------------
-  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]  # this table is in bits
+  # --- 0. Restrict to real packet bits ---------------------------------
+  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
   if nb_bits <= 0 or nb_bits >= 64:
     pkt = reg & ((1 << 64) - 1)
   else:
@@ -131,23 +140,46 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
 
   fields: list[str] = []
 
-  # --- 1. Timestamp-ish opcodes ----------------------------------------
+  shift, width = DELTA_MAP_DEFAULT.get(opcode, (0, 0))
+  if width:
+    field_mask = (1 << width) - 1
+    shaped_field = (pkt >> shift) & field_mask
+  else:
+    field_mask = 0
+    shaped_field = 0
 
-  if opcode == 0x0F:  # TIME_SHORT_DELTA_PLUS4
-    # By the time we get here, `delta` is already raw_delta+4.
+  # =====================================================================
+  # 1. Timestamp-centric opcodes (actually drive 'time')
+  # =====================================================================
+
+  if opcode == 0x0F:  # TS_DELTA_SHORT_PLUS4
+    # In the caller, delta already has +4 applied.
+    raw_delta = shaped_field
+    fields.append(f"raw_delta={raw_delta}")
     fields.append(f"ts_short_plus4={delta}")
     return ", ".join(fields)
 
-  if opcode == 0x11:  # TIME_WAVE_STATE (medium/large delta)
-    shift, width = DELTA_MAP_DEFAULT[opcode]
-    raw_delta = (pkt >> shift) & ((1 << width) - 1)
-    coarse    = (pkt >> (shift + width)) & 0xFF  # next byte above delta
+  if opcode == 0x11:  # TS_WAVE_STATE_SAMPLE
+    # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
+    raw_delta = shaped_field
+    coarse    = (pkt >> (shift + width)) & 0xFF  # matches byte at +10 in C
     fields.append(f"raw_delta={raw_delta}")
     if coarse:
-      fields.append(f"raw_coarse=0x{coarse:02x}")
+      fields.append(f"coarse_state=0x{coarse:02x}")
+    # From decomp:
+    #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
+    #  - when coarse&8, it marks all live waves as "terminated"
+    if coarse & 0x01:
+      fields.append("flag_wave_interest=1")
+    if coarse & 0x08:
+      fields.append("flag_terminate_all=1")
     return ", ".join(fields)
 
-  if opcode == 0x16:  # TIME_LONG_OR_MARKER
+  if opcode == 0x16:  # TS_DELTA36_OR_MARK
+    # Bits:
+    #   bit8  -> 0x100
+    #   bit9  -> 0x200
+    #   bits 12..47 -> 36-bit field used as delta or marker
     bit8 = bool(pkt & 0x100)
     bit9 = bool(pkt & 0x200)
     if not bit9:
@@ -158,41 +190,95 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
       mode = "other"
     val36 = (pkt >> 12) & ((1 << 36) - 1)
     fields.append(f"mode={mode}")
-    fields.append(f"val36=0x{val36:x}")
+    if mode != "delta":
+      fields.append(f"val36=0x{val36:x}")
     return ", ".join(fields)
 
-  # --- 2. Opcode 0x14: exec/config record ------------------------------
+  # For 0x07, 0x0A–0x0E, we know they drive time (via DELTA_MAP_DEFAULT),
+  # but we don't see any other fields used in the decomp.
+  if opcode in (0x07, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E):
+    if width:
+      raw_delta = shaped_field
+      leftover  = pkt & ~(field_mask << shift)
+      fields.append(f"raw_delta={raw_delta}")
+      if leftover:
+        fields.append(f"payload=0x{leftover:x}")
+    return ", ".join(fields)
 
-  if opcode == 0x14:
-    subop   = (pkt >> 16) & 0xFFFF       # matches (short)(w >> 0x10)
-    val32   = (pkt >> 32) & 0xFFFFFFFF   # matches (uint)(w >> 0x20)
-    slot    = (pkt >> 7) & 0x7           # used as (idx & 4) + (idx & 3)
-    hi_byte = (pkt >> 8) & 0xFF
+  # =====================================================================
+  # 2. Small "meta + tiny delta" packets (0x01–0x06)
+  # =====================================================================
+
+  if opcode == 0x01:  # META_ID12_TS_SMALL
+    id12 = pkt & 0xFFF
+    fields.append(f"id12=0x{id12:03x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode == 0x02:  # META_FLAG8_TS_SMALL
+    flag8 = pkt & 0xFF
+    fields.append(f"flag8=0x{flag8:02x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode == 0x03:  # META_SUBEVENT8_TS_SMALL
+    sub8 = pkt & 0xFF
+    fields.append(f"subevent8=0x{sub8:02x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode == 0x04:  # META_BASE_INDEX12_TS
+    idx12 = pkt & 0xFFF
+    fields.append(f"base_index12=0x{idx12:03x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode in (0x05, 0x06):  # META_DESC24_TS_A/B
+    desc24 = pkt & 0xFFFFFF
+    fields.append(f"desc24=0x{desc24:06x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 3. Opcode 0x14: exec/config record (+ COR marker)
+  # =====================================================================
+
+  if opcode == 0x14:  # INST_EXEC_OR_CFG
+    subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
+    val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
+    slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
+    hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
 
     fields.append(f"subop=0x{subop:04x}")
     fields.append(f"slot={slot}")
     fields.append(f"val32=0x{val32:08x}")
 
     if hi_byte & 0x80:
-      # "config" flavour, writes into local_168[...] etc.
+      # Config flavour: writes config words into per-slot state arrays.
       fields.append("kind=config")
       if subop == 0x000C:
         fields.append("cfg_target=local_168[slot].lo")
       elif subop == 0x000D:
         fields.append("cfg_target=local_168[slot].hi")
     else:
-      # COR marker: subop 0xC342, val32==0x434F5200 ("COR\0")
+      # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
       if subop == 0xC342:
         fields.append("kind=cor_stream")
         if val32 == 0x434F5200:
           fields.append("cor_magic='COR\\0'")
-
     return ", ".join(fields)
 
-  # --- 3. Opcode 0x17: mode/layout header ------------------------------
+  # =====================================================================
+  # 4. Opcode 0x17: layout / mode header
+  # =====================================================================
 
-  if opcode == 0x17:
-    # From case 0x17:
+  if opcode == 0x17:  # LAYOUT_MODE_HEADER
+    # From decomp (two sites with identical logic):
     #   layout = (w >> 7) & 0x3f
     #   mode   = (w >> 0xd) & 3
     #   group  = (w >> 0xf) & 7
@@ -213,27 +299,26 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
     fields.append(f"sel_b={sel_b}")
     if layout == 4:
       fields.append(f"layout4_flag={flag4}")
-
     return ", ".join(fields)
 
-  # --- 4. Opcode 0x09: state-ish / indirection record ------------------
+  # =====================================================================
+  # 5. Opcode 0x09: state / route config record
+  # =====================================================================
 
-  if opcode == 0x09:
-    # From case 9 on puVar58[1] (here pkt):
-    #
-    #   uVar41 = (w & 0xffffffff) >> 7;  local_520 = uVar41 & 1
-    #   local_4a0 = (w >> 8) & 3;
-    #   local_4a8 = (w >> 10) & (7 or 0xf)   (depends on local_494)
-    #   uVar69    = (w >> 0xd) or (w >> 0xf) (depends on local_494)
-    #   local_518 = (w >> 0x19) & 0x7f;
-    #
-    # We *don’t* know local_494 here, so we just expose the raw slices.
-    flag7   = (pkt >> 7) & 0x1          # low bit of uVar41
-    cls2    = (pkt >> 8) & 0x3          # local_4a0
-    slot4   = (pkt >> 10) & 0xF         # superset of 3-bit local_4a8
-    idx_lo  = (pkt >> 13) & 0x1F        # matches uVar69&0x1F when layout<4
-    idx_hi  = (pkt >> 15) & 0x1F        # matches uVar69&0x1F when layout>=4
-    id7     = (pkt >> 0x19) & 0x7F      # local_518
+  if opcode == 0x09:  # PERF_ROUTE_CONFIG
+    # From case 9 in multiple consumers:
+    #   flag7  = (w >> 7) & 1        (low bit of uVar41)
+    #   cls2   = (w >> 8) & 3        (class / group)
+    #   slot4  = (w >> 10) & 0xf     (slot / group index)
+    #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
+    #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
+    #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
+    flag7   = (pkt >> 7) & 0x1
+    cls2    = (pkt >> 8) & 0x3
+    slot4   = (pkt >> 10) & 0xF
+    idx_lo  = (pkt >> 13) & 0x1F
+    idx_hi  = (pkt >> 15) & 0x1F
+    id7     = (pkt >> 0x19) & 0x7F
 
     fields.append(f"flag7={flag7}")
     fields.append(f"cls2={cls2}")
@@ -243,18 +328,18 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
     fields.append(f"id7=0x{id7:x}")
     return ", ".join(fields)
 
-  # --- 5. Opcode 0x18: perf/event trigger ------------------------------
+  # =====================================================================
+  # 6. Opcode 0x18: perf/event selector (FUN_0010aba0)
+  # =====================================================================
 
-  if opcode == 0x18:
+  if opcode == 0x18:  # PERF_EVENT_SELECT
     # From case 0x18:
-    #  - low 3 bits:    (w & 7)
-    #  - mid 3 bits:    (w >> 3) & 7   or (w >> 4) & 7  (layout–dependent)
-    #  - hi id:         (w >> 0xc) & 0xff   OR (w >> 0xd) & 0x7f
-    #  - flag bits at 6 / 7
-    #
-    # The *real* semantics depend on global local_494 and accumulated
-    # local_500, so we keep this as a raw view that’s still useful for
-    # debugging, but not layout-dependent.
+    #   low3   = w & 7
+    #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
+    #   flags  = bits 6 (B6) and 7 (B7)
+    #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
+    #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
+    #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
     low3     = pkt & 0x7
     grp3_a   = (pkt >> 3) & 0x7
     grp3_b   = (pkt >> 4) & 0x7
@@ -276,13 +361,34 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
     fields.append(f"hi7=0x{hi7:02x}")
     return ", ".join(fields)
 
-  # --- 6. Generic tiny event-ish packets -------------------------------
+  # =====================================================================
+  # 7. Opcode 0x15: perfcounter snapshot
+  # =====================================================================
 
-  if opcode in (0x08, 0x12, 0x19):
-    # These are all "small event" style tokens. The exact layout depends
-    # on global state (local_494 etc), so we just show:
-    #   - low 8 bits as a kind/flag byte
-    #   - the rest as an opaque payload.
+  if opcode == 0x15:  # PERFCOUNTER_SNAPSHOT
+    # NIBBLE_BUDGET gives full 64 bits here.
+    # DELTA_MAP_DEFAULT: shift=7, width=3 → tiny delta field.
+    raw_delta = shaped_field if width else 0
+    # low bits below the delta field
+    snap_low  = pkt & ((1 << shift) - 1) if shift else 0
+    # everything above delta field
+    snap_hi   = pkt >> (shift + width) if width else (pkt >> shift)
+
+    fields.append(f"raw_delta={raw_delta}")
+    fields.append(f"snap_low_s{shift}=0x{snap_low:x}")
+    fields.append(f"snap_hi=0x{snap_hi:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 8. Small event-ish packets (0x08 / 0x12 / 0x13 / 0x19)
+  # =====================================================================
+
+  if opcode in (0x08, 0x12, 0x13, 0x19):
+    # These are all "small event / metric" style tokens. The exact semantics
+    # depend on layout (0x17) and accumulated state (local_500 etc), so we
+    # expose:
+    #   - low 8 bits as kind byte
+    #   - rest as opaque payload.
     kind    = pkt & 0xFF
     payload = pkt >> 8
     fields.append(f"kind_byte=0x{kind:02x}")
@@ -290,10 +396,35 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
       fields.append(f"payload=0x{payload:x}")
     return ", ".join(fields)
 
-  # --- 7. Everything else: no extra decode -----------------------------
-  return ""
+  # =====================================================================
+  # 9. Pseudo opcode 0x10: never a "real" packet
+  # =====================================================================
 
-def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000) -> None:
+  if opcode == 0x10:  # PSEUDO_NEED_MORE_BITS
+    # The main loop never prints these; they're just a control token.
+    return ""
+
+  # =====================================================================
+  # 10. Generic fallback: expose the DELTA_MAP_DEFAULT field + leftover
+  # =====================================================================
+
+  if width:
+    fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    leftover = pkt & ~(field_mask << shift)
+    if leftover:
+      fields.append(f"payload=0x{leftover:x}")
+
+  return ", ".join(fields)
+
+# 0xb is time something
+# 0xd is time something
+# 0xf is small time advance
+# 0x11 is time advance
+# 0x16 is big time advance + markers
+# 0x14 is REG
+DEFAULT_FILTER = (0xb, 0xd, 0xf, 0x11, 0x16, 0x14) if getenv("FILTER", 1) else None
+
+def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000, filter=DEFAULT_FILTER) -> None:
   """
   Minimal debug: print ONE LINE per decoded token (packet).
 
@@ -350,20 +481,18 @@ def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000) -> None:
       if two_bits == 1:
         flags |= 0x01
 
+      # Common 36-bit field at bits [12..47]
+
       if (reg & 0x200) == 0:
-        # delta mode: 36-bit delta at bits [12..47]
+        # delta mode: add 36-bit delta to time
         delta = (reg >> 12) & ((1 << 36) - 1)
         time += delta
-        note = "0x16-delta"
       else:
-        # marker mode if bit9==1 and bit8==0
+        # marker / other modes: no time advance
         if (reg & 0x100) == 0:
-          val = (reg >> 12) & ((1 << 36) - 1)
+          # real marker: bit9=1, bit8=0, non-zero payload
+          # "other" 0x16 variants, ignored for timing
           delta = 0
-          note = f"0x16-marker val=0x{val:x}"
-        else:
-          delta = 0
-          note = "0x16-other"
     else:
       # 6) Generic opcode (including 0x0F)
       shift, width = DELTA_MAP_DEFAULT[opcode]
@@ -373,22 +502,15 @@ def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000) -> None:
       # TODO: add more opcode parsers here that add notes to other opcodes
       if opcode == 0x0F:
         delta_with_fix = delta + 4
-        note = f"0x0f (+4) raw_delta={delta}"
         time += delta_with_fix
         delta = delta_with_fix
       else:
         time += delta
 
-    # ONE-LINE PRINT PER PACKET
-    #assert last_real_offset%8 == 0
-    #assert (offset)%8 == 0, f"misalign offset {offset}"
-
     # Append extra decoded fields into the note string
-    extra = decode_packet_fields(opcode, reg, delta)
-    if extra: note = (note + " ; " + extra) if note else extra
+    note = decode_packet_fields(opcode, reg, delta)
 
-    BORING_OPCODES = {0x11, 0x14}
-    if opcode not in BORING_OPCODES or getenv("BORING", 1):
+    if filter is None or opcode not in filter:
       my_reg = reg
       my_reg &= (1 << nib_budget) - 1
       print(
@@ -415,7 +537,7 @@ def parse(fn:str):
 
 if __name__ == "__main__":
   #dat_sqtt = parse("extra/sqtt/examples/profile_empty_run_0.pkl")
-  dat_sqtt = parse("extra/sqtt/examples/profile_plus_run_0.pkl")
-  #dat_sqtt = parse("extra/sqtt/examples/profile_gemm_run_0.pkl")
+  #dat_sqtt = parse("extra/sqtt/examples/profile_plus_run_0.pkl")
+  dat_sqtt = parse("extra/sqtt/examples/profile_gemm_run_0.pkl")
   blob_0 = dat_sqtt[0].blob
   parse_sqtt_print_packets(blob_0[8:])

extra/sqtt/roc.py

@@ -41,6 +41,7 @@ class WaveExec:
   wave_id:int
   cu:int
   simd:int
+  se:int
   begin_time:int
   end_time:int
   insts:list[InstExec]
@@ -78,7 +79,8 @@ class _ROCParseCtx:
       if DEBUG >= 8: print(inst_execs[-1])
 
     if ev.instructions_size > 0:
-      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, ev.begin_time, ev.end_time, inst_execs))
+      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
+                                                                               ev.end_time, inst_execs))
 
 def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   dev_events:dict[str, ProfileDeviceEvent] = {}
@@ -111,7 +113,6 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
   def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    if DEBUG >= 8: print(f"isa_cb {pc.address=} {pc.code_object_id=}")
     instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm

extra/sqtt/test_timing.py

@@ -7,11 +7,9 @@ os.environ["AMD_LLVM"] = "0"
 
 import unittest
 import sys, contextlib
-from tinygrad import Tensor
-from tinygrad.dtype import dtypes
-from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.device import Device, ProfileDeviceEvent
 
 from extra.sqtt.roc import decode, WaveExec
@@ -75,7 +73,6 @@ class TestTiming(unittest.TestCase):
     inp = Tensor([-2.0]).realize()
     with save_sqtt() as sqtt:
       Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
-
     wave = list(sqtt.values())[0][0]
     for i in range(len(wave.insts)):
       if wave.insts[i].inst.startswith("global_store"):
@@ -102,7 +99,7 @@ class TestTiming(unittest.TestCase):
       assert data0.dtype.base == dtypes.ulong
       op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
       op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
-      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
+      op = custom("unsigned long long t1 = __builtin_readcyclecounter();", op)
       op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
       return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
     diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
@@ -113,5 +110,24 @@ class TestTiming(unittest.TestCase):
     # cycles = sleep dur + overhead of storing hi/lo REG_SHADER_CYCLES
     self.assertGreaterEqual(diff_hw_reg.item(), sleep.dur)
 
+  def test_nop(self):
+    with save_sqtt() as sqtt:
+      asm_kernel(["s_nop 1"]*10).realize()
+    wave = list(sqtt.values())[0][0]
+    for e in wave.insts:
+      print(f"{e.inst} {e.dur=} {e.stall=}")
+
+  def test_wave_sched(self):
+    num_waves = getenv("NUM_WAVES", 16)
+    num_wgps = getenv("NUM_WGPS", 2)
+    num_vgpr = getenv("NUM_VGPR", 256)
+    with save_sqtt() as sqtt:
+      # 1 cycle decode, no stall
+      asm_kernel([f"v_mov_b32_e32 v{i} {i}" for i in range(num_vgpr)], l=32*num_waves, g=num_wgps).realize()
+    waves = list(sqtt.values())[0]
+    print(len(waves), "waves decoded")
+    for w in waves:
+      print(f"{w.wave_id:<2} {w.simd=} {w.cu=} {w.se=} @ clk {w.begin_time}")
+
 if __name__ == "__main__":
   unittest.main()

extra/thunder/tiny/tk/__init__.py

@@ -1 +1,6 @@
-WARP_THREADS = 32
+from tinygrad.device import Device
+
+if Device.DEFAULT == "AMD":
+  WARP_THREADS = 64
+else:
+  WARP_THREADS = 32

extra/thunder/tiny/tk/group.py

@@ -162,7 +162,7 @@ class Group:
 
   # ops that can work across multiple warps
 
-  LOAD_INNER = 8
+  LOAD_INNER = 4
   def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
     dst, src = cast(UOp, dst), cast(UOp, src)
     assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
@@ -225,7 +225,7 @@ class Group:
 
     return dst.after(dst_store.barrier()).reshape(dst.shape)
 
-  STORE_INNER = 8
+  STORE_INNER = 4
   def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
     dst, src = cast(UOp, dst), cast(UOp, src)
     assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)

test/external/external_test_example.py

@@ -1,13 +1,14 @@
 import unittest
 from tinygrad import Device
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, CI
+from tinygrad.helpers import getenv, CI, OSX
 
 def multidevice_test(fxn):
   exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
   def ret(self):
     for device in Device._devices:
-      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"]: continue
+      # broken on OSX USB AMD, why?
+      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
       if not CI: print(device)
       if device in exclude_devices:
         if not CI: print(f"WARNING: {device} test is excluded")

test/mockgpu/amd/amdgpu.py

@@ -124,6 +124,7 @@ class PM4Executor(AMDQueue):
     elif mem_data_sel == 3:
       if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
       else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")
+    elif mem_data_sel == 0: pass # no write
     else: raise RuntimeError(f"Unknown {mem_data_sel=}")
 
   def _exec_copy_data(self, n):

test/test_outerworld.py

@@ -50,8 +50,52 @@ class TestOuterRange(unittest.TestCase):
     # 3 matmuls with outer world range
     i = UOp.range(3, -100, AxisType.OUTER)
     vec_i = Tensor(vec.uop.after(i))
-    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
-    out = Tensor(vec.uop.after(vec_i.uop.store((vec_i.contiguous() @ mats[vi]).uop).end(i)))
+    comp = vec_i.contiguous() @ mats[i]
+    store = vec_i.uop.store(comp.uop).end(i)
+    out = Tensor(vec.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+class TestOuterScan(unittest.TestCase):
+  def _test_scan(self):
+    vec = Tensor.randn(1, 10).realize()
+    mats = Tensor.randn(3, 10, 10).realize()
+
+    # 3 matmuls in "scan"
+    vec1 = vec @ mats[0]
+    vec2 = vec1 @ mats[1]
+    vec3 = vec2 @ mats[2]
+    ref = Tensor.stack(vec1, vec2, vec3)
+    ref.realize()
+    return vec, mats, ref
+
+  def test_uop_fold_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with FOLD
+    i = UOp.range(3, -100, AxisType.OUTER)
+    out = Tensor.empty(1, 10)
+    phi = Tensor(i.eq(0).where(vec.uop, out.uop))
+    comp = phi @ mats[i]
+    store = out.uop.store(comp.uop).end(i)
+    out = Tensor(out.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref[2], out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+  def test_uop_scan_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with SCAN
+    i = UOp.range(3, -100, AxisType.OUTER)
+    out = Tensor.empty(3, 1, 10)
+    phi = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop))
+    comp = phi @ mats[i]
+    store = out[i].uop.store(comp.uop).end(i)
+    out = Tensor(out.uop.after(store))
     out.realize()
 
     # TODO: testing allclose
@@ -116,5 +160,84 @@ class TestOuterworld(unittest.TestCase):
     out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
     self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
 
+class TestVmap(unittest.TestCase):
+  def test_vmap_inner(self):
+    x = Tensor.ones(3, 2).contiguous()
+
+    # vmap across axis 0
+    a = UOp.range(3, -1)
+    out = x[a]*2
+    out = out.end(a)
+    out.realize()
+
+    self.assertTrue((out==2).all().item())
+
+  def test_vmap_outer(self):
+    x = Tensor.ones(3, 2).contiguous()
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, AxisType.OUTER)
+    out = x[a]*2
+    out = out.end(a)
+    out.realize()
+
+    self.assertTrue((out==2).all().item())
+
+  def test_fancy_vmap(self):
+    def f(x,y): return x+y
+
+    x = Tensor.arange(9).reshape(3,3).contiguous()
+    y = Tensor.arange(9).reshape(3,3).contiguous()
+
+    a = UOp.range(3, -1)
+    out = f(x[:,a], y[a,:])
+    out = out.end(a).realize()
+    self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
+
+  def test_vmap_inner_fusion(self):
+    x = Tensor.ones(3, 10, 2).contiguous()
+
+    # vmap across axis 0
+    a = UOp.range(3, -1)
+    out = x[a].sum(axis=0)*2
+    out = out.end(a)*4
+    out.realize()
+
+    self.assertTrue((out==10*2*4).all().item())
+
+  def test_vmap_outer_fusion(self):
+    x = Tensor.ones(3, 10, 2).contiguous()
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, AxisType.OUTER)
+    out = x[a].sum(axis=0)*2
+    out = out.end(a)*4
+    out.realize()
+
+    self.assertTrue((out==10*2*4).all().item())
+
+  def test_vmap_outer_matmul(self):
+    x = Tensor.ones(1, 10).contiguous().requires_grad_()
+    mats = Tensor.ones(3, 10, 10).contiguous()
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, AxisType.OUTER)
+    out = x @ mats[a]
+    out = out.end(a)
+
+    out.realize()
+
+  def test_vmap_outer_matmul_grad(self):
+    x = Tensor.ones(1, 10).contiguous().requires_grad_()
+    mats = Tensor.ones(3, 10, 10).contiguous().requires_grad_()
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, AxisType.OUTER)
+    out = x @ mats[a]
+    out = out.end(a)
+    out.mean().backward()
+
+    mats.grad.realize()
+
 if __name__ == '__main__':
   unittest.main()

test/test_profiler.py

@@ -17,7 +17,7 @@ def helper_collect_profile(*devs):
   cpu_events.clear()
 
   profile_list = []
-  with Context(VIZ=1):
+  with Context(VIZ=1, PROFILE=1):
     yield profile_list
     for dev in devs: dev.synchronize()
     for dev in devs: dev._at_profile_finalize()

test/testextra/test_tk.py

@@ -9,7 +9,7 @@ import numpy as np
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
 
-@unittest.skipUnless(Device.DEFAULT in ["CUDA", "NV"], "only cuda")
+@unittest.skipIf(CI and Device.DEFAULT not in ["CUDA", "NV"], "only cuda")
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "no ptx")
 class TestTK(unittest.TestCase):
   @unittest.skipIf(CI, "no wmma in ci")

test/unit/test_tqdm.py

@@ -66,6 +66,7 @@ class TestProgressBar(unittest.TestCase):
       tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
       self._compare_bars(tinytqdm_output, tqdm_output)
 
+  @unittest.skip("this is flaky")
   @patch('sys.stderr', new_callable=StringIO)
   @patch('shutil.get_terminal_size')
   def test_unit_scale(self, mock_terminal_size, mock_stderr):

test/unit/test_viz.py

@@ -6,6 +6,7 @@ from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatch
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
+from tinygrad.helpers import VIZ
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -33,11 +34,14 @@ class BaseTestViz(unittest.TestCase):
     cpu_events.clear()
     self.tms = TRACK_MATCH_STATS.value
     self.profile = PROFILE.value
+    self.viz = VIZ.value
     TRACK_MATCH_STATS.value = 2
     PROFILE.value = 1
+    VIZ.value = 1
   def tearDown(self):
     TRACK_MATCH_STATS.value = self.tms
     PROFILE.value = self.profile
+    VIZ.value = self.viz
 
 class TestViz(BaseTestViz):
   def test_simple(self):

tinygrad/codegen/opt/postrange.py

@@ -52,8 +52,10 @@ class Scheduler:
   def get_optimized_ast(self, name_override:str|None=None):
     if name_override is not None: name = name_override
     else:
-      kernel_type = "r" if self.reduceop is not None else "E"
-      name = kernel_type + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
+      k_type = "r" if self.reduceop is not None else "E"
+      special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg)
+      special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops]
+      name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
       Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
       num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
       name += colored(num, 'BLACK')

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited, VIZ
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -355,8 +355,9 @@ if PROFILE:
 
     with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)
 
-    from tinygrad.uop.ops import launch_viz
-    launch_viz("PROFILE", fn)
+    if VIZ:
+      from tinygrad.uop.ops import launch_viz
+      launch_viz("PROFILE", fn)
 
 def enumerate_devices_str() -> Generator[str, None, None]:
   from tinygrad import Tensor, Device

tinygrad/engine/schedule.py

@@ -23,10 +23,12 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
   in_degree: dict[UOp, int] = {}
   var_vals: dict[str, int] = {}
   for u in sched_sink.toposort():
-    if u.op is not Ops.AFTER: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
+    if u.op is Ops.RANGE:
+      in_degree.setdefault(u, 0)
+      continue
+    if u.op is not Ops.AFTER or u.src[1].op is Ops.RANGE: continue
     k = u.src[1]
     in_degree.setdefault(k, 0)
-    if k.op is Ops.RANGE: continue
     for s in k.src[0].src if k.op is Ops.END else k.src:
       if s.op is Ops.AFTER:
         children[s.src[1]].append(k)
@@ -88,7 +90,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
       if rk.op is Ops.END: schedule.append(rk)
     else:
       raise RuntimeError(f"can't schedule {k.op}")
-    for x in children[k]:
+    for x in children[rk]:
       in_degree[x] -= 1
       if in_degree[x] == 0: queues[_heuristic(x)].append(x)
 

tinygrad/gradient.py

@@ -43,6 +43,8 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
+  # this only works on single ends of outer ranges
+  (UPat(Ops.END, name="e"), lambda ctx, e: (ctx.shrink(((e.src[1],e.src[1]+1),)+(None,)*(len(ctx.shape)-1)).reshape(ctx.shape[1:]), None)),
 ])
 
 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:

tinygrad/helpers.py

@@ -179,7 +179,9 @@ ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), Conte
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
-VIZ = PROFILE = ContextVar("VIZ", 0)
+# VIZ implies PROFILE, but you can run PROFILE without VIZ
+VIZ = ContextVar("VIZ", 0)
+PROFILE = ContextVar("PROFILE", VIZ.value)
 SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)

tinygrad/runtime/ops_amd.py

@@ -12,7 +12,7 @@ from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
 from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
+from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
@@ -357,6 +357,7 @@ class AMDComputeQueue(HWQueue):
 
   def timestamp(self, signal:AMDSignal):
     with self.pred_exec(xcc_mask=0b1):
+      self.release_mem(cache_flush=False) # ensure all prior writes are done
       self.release_mem(signal.timestamp_addr, 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter, self.pm4.int_sel__mec_release_mem__none)
       self.acquire_mem() # ensure timestamp is written
     return self
@@ -908,7 +909,8 @@ class AMDDevice(HCQCompiled):
     self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
 
     compilers:list[CompilerPairT] = [(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
-                                     (functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch))]
+                                     (functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch)),
+                                     (functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))]
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),

tinygrad/runtime/support/compiler_amd.py

@@ -1,4 +1,4 @@
-import ctypes
+import ctypes, hashlib, tempfile, subprocess, pathlib
 from tinygrad.helpers import system
 from tinygrad.runtime.autogen import comgr
 try:
@@ -13,7 +13,7 @@ from tinygrad.runtime.support.compiler_cpu import LLVMCompiler
 from tinygrad.helpers import OSX, to_char_p_p
 
 def amdgpu_disassemble(lib:bytes):
-  asm = system(f"{'llvm-objdump' if OSX else '/opt/rocm/llvm/bin/llvm-objdump'} -d -", input=lib).splitlines()
+  asm = system(f"{'/opt/homebrew/opt/llvm/bin/llvm-objdump' if OSX else '/opt/rocm/llvm/bin/llvm-objdump'} -d -", input=lib).splitlines()
   while asm and ("s_nop 0" in asm[-1] or "s_code_end" in asm[-1]): asm.pop()
   print("\n".join(asm))
 
@@ -90,6 +90,24 @@ class HIPCompiler(Compiler):
     except RuntimeError as e: raise CompileError(e) from e
   def disassemble(self, lib:bytes): amdgpu_disassemble(lib)
 
+class HIPCCCompiler(Compiler):
+  def __init__(self, arch:str, extra_options:list[str]=[]):
+    self.arch, self.extra_options = arch, extra_options
+    super().__init__(f"compile_hipcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+  def compile(self, src:str) -> bytes:
+    with tempfile.NamedTemporaryFile(suffix=".cpp") as srcf, tempfile.NamedTemporaryFile(suffix=".bc") as bcf:
+      with tempfile.NamedTemporaryFile(suffix=".hsaco") as libf:
+        srcf.write(src.encode())
+        srcf.flush()
+
+        subprocess.run(["hipcc", "-c", "-emit-llvm", "--cuda-device-only", "-O3", "-mcumode",
+                        f"--offload-arch={self.arch}", "-I/opt/rocm/include/hip", "-o", bcf.name, srcf.name] + self.extra_options, check=True)
+        subprocess.run(["hipcc", "-target", "amdgcn-amd-amdhsa", f"-mcpu={self.arch}",
+                        "-O3", "-mllvm", "-amdgpu-internalize-symbols", "-c", "-o", libf.name, bcf.name], check=True)
+
+        return pathlib.Path(libf.name).read_bytes()
+  def disassemble(self, lib:bytes): amdgpu_disassemble(lib)
+
 class AMDLLVMCompiler(LLVMCompiler):
   jit = False
   target_arch = "AMDGPU"

tinygrad/schedule/indexing.py

@@ -24,8 +24,8 @@ def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
 pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize COPY/BUFFER_VIEW/CONTIGUOUS/STORE
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
+  # always realize COPY/BUFFER_VIEW/CONTIGUOUS/STORE/END
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.END}, name="tr"), realize),
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
@@ -66,6 +66,9 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
         new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
         del ctx.realize_map[s]
       else:
+        if new_src.op is Ops.END:
+          # skip END
+          new_src = new_src.src[0]
         # None in the device assigns it a number later
         opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
         new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
@@ -173,8 +176,12 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
     consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
     if x in rctx.realize_map:
-      # if this is in the realize_map, we create new ranges (at the output)
-      out_rngs = tuple(rctx.new_range(s) for s in x.shape)
+      if x.op is Ops.END:
+        # for END, we use the ranges in the src as the early ones
+        out_rngs = x.src[1:]+tuple(rctx.new_range(s) for s in x.src[0].shape)
+      else:
+        # if this is in the realize_map, we create new ranges (at the output)
+        out_rngs = tuple(rctx.new_range(s) for s in x.shape)
       # all ranges are ended now
       ending_ranges[x] = []
       # mark all ranges as ended
@@ -249,6 +256,10 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     if x.op is Ops.REDUCE_AXIS:
       rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
 
+    # END ends ranges
+    if x.op is Ops.END:
+      rngs = rngs[len(x.src)-1:]
+
     if debug:
       realized_ranges = rctx.realize_map.get(x, None)
       if x.op is Ops.RESHAPE or len(rngs) != len(out_rngs):

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
-from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags
+from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags, range_str
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
 from tinygrad.helpers import PCONTIG, partition, get_single_element, unwrap, disable_gc
@@ -397,6 +397,9 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
   if r.tag != (): return None
+  if r.arg[-1] == AxisType.OUTER:
+    # for outer range, we replace with a bound variable
+    return UOp.variable("range_"+range_str(r), r.vmin, r.vmax).bind(r.replace(tag=None))
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
   ctx.range += 1
   return ret
@@ -469,6 +472,7 @@ pm_add_range_tags = PatternMatcher([
 ])
 
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
+  # if we have any outer ranges open here, we don't split
   if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
 
   # ends of outer range don't go in kernels
@@ -496,7 +500,12 @@ def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
     raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop.buffer for b in kernel.src)}")
   return kernel
 
+def split_inner_and_outer_end(x: UOp):
+  outer_ranges, inner_ranges = partition(x.src[1:], lambda r: r.arg[-1] == AxisType.OUTER)
+  if len(outer_ranges) and len(inner_ranges): return x.src[0].end(*inner_ranges).end(*outer_ranges)
+
 split_kernels = PatternMatcher([
+  (UPat(Ops.END, name="x"), split_inner_and_outer_end),
   (UPat((Ops.STORE, Ops.END), name="x"), split_store),
 ])
 
@@ -507,7 +516,7 @@ def tag_uop(ctx:list[UOp], x:UOp):
   return x.replace(tag=(len(ctx)-1,))
 add_tags = PatternMatcher([
   # don't tag BUFFERs, they are global
-  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.KERNEL, Ops.END,
+  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.KERNEL,
                      Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
   (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])
@@ -571,12 +580,12 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
       assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
   if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
 
-  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-
   # TODO: we can probably get this earlier
   sink_tags = [s.tag for s in tsink.src]
   tsink = graph_rewrite(tsink, _remove_all_tags, name="remove all tags")
 
+  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+
   becomes_map: dict[UOp, UOp] = {}
   for tag, s in zip(sink_tags, tsink.src):
     assert tag is not None

tinygrad/tensor.py

@@ -481,6 +481,9 @@ class Tensor(OpMixin):
     if y.op is Ops.ADD: return Tensor.from_uop(y.src[0]) + Tensor.from_uop(y.src[1])
     raise RuntimeError(f"unhandled UOp {y}")
 
+  def end(self, *rngs:UOp):
+    return self._apply_uop(UOp.end, extra_args=rngs, dtype=self.dtype)
+
   # ***** creation entrypoint *****
 
   @staticmethod

tinygrad/uop/ops.py

@@ -219,9 +219,13 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
 
       # passthrough ops
-      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
+      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER:
         return self.src[0]._shape
 
+      # end adds dims to the front
+      case Ops.END:
+        return None if self.src[0]._shape is None else (tuple(x.vmax+1 for x in self.src[1:]) + self.src[0]._shape)
+
       # ops with custom handling
       case Ops.KERNEL: return self.arg.ast._shape
 
@@ -398,9 +402,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
   def store(self, src:UOp|ConstType, **kwargs):
     return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self, UOp.const(self.dtype, src) if not isinstance(src, UOp) else src), **kwargs)
-  def end(self, *src:UOp):
+  def end(self, *src:UOp, **kwargs):
     if len(src) == 0: return self
-    return UOp(Ops.END, src=(self,)+src)
+    return UOp(Ops.END, src=(self,)+src, **kwargs)
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
@@ -1265,15 +1269,12 @@ pm_lower_index_dtype = PatternMatcher([
   (UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.index),), name="u"), lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.index)),
   (UPat(Ops.DEFINE_VAR, dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.index)),
   (UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.index), UPat.cvar("val").cast(dtypes.index))), lambda var,val: var.bind(val).cast(dtypes.index)),
-  (UPat(Ops.CAST, src=(UPat(name="x").cast(dtypes.index),), name="c"), lambda x,c: x.cast(c.dtype)),
   # lower Invalid
   (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond, ptr=True)),
   # remove hanging casts
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx, ptr=True)),
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))),
    lambda buf,idx,valid: buf.index(idx, valid, ptr=True)),
-  (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
-    lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
   (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
    lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.index else s for s in n.src))),
 ])

tinygrad/uop/spec.py

@@ -42,7 +42,7 @@ shared_spec = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
 
   # RANGE/SPECIAL define loops, END closes them
-  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE))), lambda: True),
 ])
 
 # ***** UOp spec in the Tensor graph *****
@@ -171,7 +171,7 @@ kernel_spec = PatternMatcher([
   (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
 
   # END can end multiple axes here
-  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
+  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
 
   # bufferize can be on anything
   (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True),

tinygrad/viz/js/index.js

@@ -149,7 +149,7 @@ function renderDag(graph, additions, recenter, layoutOpts) {
 
 // ** profiler graph
 
-function formatTime(ts, dur=ts) {
+function formatMicroseconds(ts, dur=ts) {
   if (dur<=1e3) return `${ts.toFixed(2)}us`;
   if (dur<=1e6) return `${(ts*1e-3).toFixed(2)}ms`;
   return `${(ts*1e-6).toFixed(2)}s`;
@@ -158,7 +158,7 @@ const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
 
 const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#1d2e62", "#63b0cd"],
   DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
-  BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"],
+  BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"], SIMD:["#3600f0"],
   CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
 const cycleColors = (lst, i) => lst[i%lst.length];
 
@@ -198,13 +198,15 @@ function focusShape(shape) {
   return metadata.replaceChildren(shapeMetadata.get(focusedShape) ?? "");
 }
 
-async function renderProfiler() {
+async function renderProfiler(path, unit) {
   displaySelection("#profiler");
   metadata.replaceChildren(shapeMetadata.get(focusedShape) ?? "");
   // layout once!
-  if (data != null) return updateProgress({ start:false });
+  if (data != null && data.path === path) return updateProgress({ start:false });
+  // support non realtime x axis units
+  const formatTime = unit === "realtime" ? formatMicroseconds : (s) => `${s} ${unit}`;
   const profiler = d3.select("#profiler").html("");
-  const buf = await (await fetch("/get_profile")).arrayBuffer();
+  const buf = await (await fetch(path)).arrayBuffer();
   const view = new DataView(buf);
   let offset = 0;
   const u8 = () => { const ret = view.getUint8(offset); offset += 1; return ret; }
@@ -227,7 +229,7 @@ async function renderProfiler() {
   const colorMap = new Map();
   // map shapes by event key
   const shapeMap = new Map();
-  data = {tracks:new Map(), axes:{}};
+  data = {tracks:new Map(), axes:{}, path};
   const heightScale = d3.scaleLinear().domain([0, tracePeak]).range([4,maxheight=100]);
   for (let i=0; i<layoutsLen; i++) {
     const nameLen = view.getUint8(offset, true); offset += 1;
@@ -479,6 +481,7 @@ async function renderProfiler() {
     d3.select(canvas).call(canvasZoom.transform, zoomLevel);
   }
 
+  zoomLevel = d3.zoomIdentity;
   canvasZoom = d3.zoom().filter(vizZoomFilter).scaleExtent([1, Infinity]).translateExtent([[0,0], [Infinity,0]]).on("zoom", e => render(e.transform));
   d3.select(canvas).call(canvasZoom);
   document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
@@ -691,13 +694,14 @@ async function main() {
     if (url.pathname+url.search !== ckey) e.close();
     else if (e.readyState === EventSource.OPEN) activeSrc = e;
   }
-  if (ctx.name === "Profiler") return renderProfiler();
+  if (ctx.name === "Profiler") return renderProfiler("/get_profile", "realtime");
   if (workerUrl == null) await initWorker();
   if (ckey in cache) {
     ret = cache[ckey];
   }
   // ** Disassembly view
   if (ckey.startsWith("/render")) {
+    if (step.fmt === "timeline") return renderProfiler(ckey, "clk"); // cycles on the x axis
     if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
     displaySelection("#custom");
     metadata.innerHTML = "";

tinygrad/viz/serve.py

@@ -215,10 +215,12 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
   except Exception: return err("DECODER ERROR")
   if not rctx.inst_execs: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
   steps:list[dict] = []
+  units:set[str] = set()
   for name,waves in rctx.inst_execs.items():
+    events:list[ProfileEvent] = []
     prg = trace.keys[r].ret if (r:=ref_map.get(name)) else None
-    steps.append({"name":prg.name if prg is not None else name, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
-                  "depth":0, "data":{"src":prg.src if prg is not None else name, "lang":"cpp"}})
+    steps.append(first:={"name":prg.name if prg is not None else name, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                         "depth":0, "fmt":"timeline"})
 
     # 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:
@@ -228,14 +230,18 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
     # 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+.
     for w in waves:
+      units.add(row:=f"SIMD:{w.simd} CU:{w.cu} SE:{w.se}")
+      events.append(ProfileRangeEvent(row, wave_name:=f"wave {w.wave_id}", Decimal(w.begin_time), Decimal(w.end_time)))
       rows, prev_instr = [], w.begin_time
       for i,e in enumerate(w.insts):
         rows.append((e.inst, e.time, max(0, e.time-prev_instr), e.dur, e.stall, str(e.typ).split("_")[-1]))
         prev_instr = max(prev_instr, e.time + e.dur)
-      summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"CU", "value":w.cu},
-                 {"label":"SIMD", "value":w.simd}]
-      steps.append({"name":f"Wave {w.wave_id}", "depth":1, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+      summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SIMD", "value":w.simd}, {"label":"CU", "value":w.cu},
+                 {"label":"SE", "value":w.se}]
+      steps.append({"name":wave_name, "depth":1, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
                     "data":{"rows":rows, "cols":["Instruction", "Clk", "Idle", "Duration", "Stall", "Type"], "summary":summary}})
+    events = [ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+events
+    first["data"] = {"value":get_profile(events), "content_type":"application/octet-stream"}
   ctxs.append({"name":"Counters", "steps":steps})
 
 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
@@ -302,9 +308,9 @@ def get_stdout(f: Callable) -> str:
   except Exception: traceback.print_exc(file=buf)
   return buf.getvalue()
 
-def get_render(i:int, j:int, fmt:str) -> dict|None:
+def get_render(i:int, j:int, fmt:str) -> dict:
   if fmt == "counters": return ctxs[i]["steps"][j]["data"]
-  if not isinstance(prg:=trace.keys[i].ret, ProgramSpec): return None
+  if not isinstance(prg:=trace.keys[i].ret, ProgramSpec): return {}
   if fmt == "uops": return {"src":get_stdout(lambda: print_uops(prg.uops or [])), "lang":"txt"}
   if fmt == "src": return {"src":prg.src, "lang":"cpp"}
   compiler = Device[prg.device].compiler
@@ -336,11 +342,14 @@ class Handler(BaseHTTPRequestHandler):
     elif (query:=parse_qs(url.query)):
       if url.path == "/render":
         render_src = get_render(get_int(query, "ctx"), get_int(query, "step"), query["fmt"][0])
-        ret, content_type = json.dumps(render_src).encode(), "application/json"
+        if "content_type" in render_src: ret, content_type = render_src["value"], render_src["content_type"]
+        else: ret, content_type = json.dumps(render_src).encode(), "application/json"
       else:
         try: return self.stream_json(get_full_rewrite(trace.rewrites[i:=get_int(query, "ctx")][get_int(query, "idx")], i))
         except (KeyError, IndexError): status_code = 404
-    elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
+    elif url.path == "/ctxs":
+      lst = [{**c, "steps":[{k:v for k, v in s.items() if k != "data"} for s in c["steps"]]} for c in ctxs]
+      ret, content_type = json.dumps(lst).encode(), "application/json"
     elif url.path == "/get_profile" and profile_ret: ret, content_type = profile_ret, "application/octet-stream"
     else: status_code = 404