never mind, i don't like this

* write scan in uops

* ops range

* no need for variable

* meh, later

* shorter

.github/workflows/benchmark.yml

@@ -643,14 +643,14 @@ jobs:
       run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: benchmark MobileNetV2 on DSP
-      run: |
-        # generate quantized weights
-        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-        ln -s /data/home/tiny/tinygrad/testsig-*.so .
-        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
-        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+    # - name: benchmark MobileNetV2 on DSP
+    #   run: |
+    #     # generate quantized weights
+    #     ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+    #     ln -s /data/home/tiny/tinygrad/testsig-*.so .
+    #     PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+    #     # benchmark on DSP with NOOPT=1, the devectorizer has issues
+    #     PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

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

@@ -0,0 +1,543 @@
+import pickle
+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 = {
+  # 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",
+
+  # ------------------------------------------------------------------------
+  # 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
+
+  # ------------------------------------------------------------------------
+  # 0x10–0x19: timestamps, layout headers, events, perf
+  # ------------------------------------------------------------------------
+  0x10: "PSEUDO_NEED_MORE_BITS",    # not a real packet; decoder refill hint
+
+  0x11: "TS_WAVE_STATE_SAMPLE",     # wave stall/termination sample (byte at +10)
+  0x13: "EVT_SMALL_GENERIC",        # same structural family as 0x08/0x12/0x19
+
+  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
+# rocprof_trace_decoder_parse_data-0x11c6a0
+# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
+
+# ---------- 1. local_138: 256-byte state->token table ----------
+
+STATE_TO_TOKEN: bytes = bytes([
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+])
+
+
+# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
+
+NIBBLE_BUDGET = [
+  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40,
+  0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
+  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40,
+  0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+]
+assert len(NIBBLE_BUDGET) == 32
+
+
+# ---------- 3. delta_map from your hash nodes ----------
+
+# opcode -> (shift, width)
+DELTA_MAP_DEFAULT = {
+  0x01: (3,  3),   # shift=3,  end=6
+  0x02: (4,  2),   # shift=4,  end=6
+  0x03: (4,  2),   # shift=4,  end=6
+  0x04: (4,  3),   # shift=4,  end=7
+  0x05: (5,  3),   # shift=5,  end=8
+  0x06: (5,  3),   # shift=5,  end=8
+  0x07: (8,  3),   # shift=8,  end=11
+  0x08: (5,  3),   # shift=5,  end=8
+  0x09: (5,  2),   # shift=5,  end=7
+  0x0A: (5,  2),   # shift=5,  end=7
+  0x0B: (5,  3),   # shift=5,  end=8
+  0x0C: (5,  3),   # shift=5,  end=8
+  0x0D: (5,  3),   # shift=5,  end=8
+  0x0E: (7,  2),   # shift=7,  end=9
+  0x0F: (4,  4),   # shift=4,  end=8
+  0x10: (0,  0),   # shift=0,  end=0  (no delta)
+  0x11: (7,  9),   # shift=7,  end=16
+  0x12: (8,  3),   # shift=8,  end=11
+  0x13: (8,  3),   # shift=8,  end=11
+  0x14: (4,  3),   # shift=4,  end=7
+  0x15: (7,  3),   # shift=7,  end=10
+  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
+  0x17: (0,  0),   # shift=0,  end=0  (no delta)
+  0x18: (4,  3),   # shift=4,  end=7
+  0x19: (7,  2),   # shift=7,  end=9
+}
+
+# ---------- 4. One-line-per-packet parser ----------
+
+def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
+  """
+  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 real packet bits ---------------------------------
+  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
+  if nb_bits <= 0 or nb_bits >= 64:
+    pkt = reg & ((1 << 64) - 1)
+  else:
+    pkt = reg & ((1 << nb_bits) - 1)
+
+  fields: list[str] = []
+
+  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
+
+  # =====================================================================
+  # 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:  # 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"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:  # 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:
+      mode = "delta"
+    elif not bit8:
+      mode = "marker"
+    else:
+      mode = "other"
+    val36 = (pkt >> 12) & ((1 << 36) - 1)
+    fields.append(f"mode={mode}")
+    if mode != "delta":
+      fields.append(f"val36=0x{val36:x}")
+    return ", ".join(fields)
+
+  # 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)
+
+  # =====================================================================
+  # 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 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, 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)
+
+  # =====================================================================
+  # 4. Opcode 0x17: layout / mode header
+  # =====================================================================
+
+  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
+    #   sel_a  = (w >> 0x1c) & 0xf
+    #   sel_b  = (w >> 0x21) & 7
+    #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
+    layout = (pkt >> 7)  & 0x3F
+    mode   = (pkt >> 13) & 0x3
+    group  = (pkt >> 15) & 0x7
+    sel_a  = (pkt >> 0x1C) & 0xF
+    sel_b  = (pkt >> 0x21) & 0x7
+    flag4  = (pkt >> 0x3B) & 0x1
+
+    fields.append(f"layout={layout}")
+    fields.append(f"group={group}")
+    fields.append(f"mode={mode}")
+    fields.append(f"sel_a={sel_a}")
+    fields.append(f"sel_b={sel_b}")
+    if layout == 4:
+      fields.append(f"layout4_flag={flag4}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 5. Opcode 0x09: state / route config record
+  # =====================================================================
+
+  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}")
+    fields.append(f"slot4=0x{slot4:x}")
+    fields.append(f"idx_lo5=0x{idx_lo:x}")
+    fields.append(f"idx_hi5=0x{idx_hi:x}")
+    fields.append(f"id7=0x{id7:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 6. Opcode 0x18: perf/event selector (FUN_0010aba0)
+  # =====================================================================
+
+  if opcode == 0x18:  # PERF_EVENT_SELECT
+    # From case 0x18:
+    #   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
+    flag_b6  = (pkt >> 6) & 0x1
+    flag_b7  = (pkt >> 7) & 0x1
+    idx5_a   = (pkt >> 7) & 0x1F
+    idx5_b   = (pkt >> 8) & 0x1F
+    hi8      = (pkt >> 12) & 0xFF
+    hi7      = (pkt >> 13) & 0x7F
+
+    fields.append(f"low3=0x{low3:x}")
+    fields.append(f"grp3_a=0x{grp3_a:x}")
+    fields.append(f"grp3_b=0x{grp3_b:x}")
+    fields.append(f"flag_b6={flag_b6}")
+    fields.append(f"flag_b7={flag_b7}")
+    fields.append(f"idx5_a=0x{idx5_a:x}")
+    fields.append(f"idx5_b=0x{idx5_b:x}")
+    fields.append(f"hi8=0x{hi8:02x}")
+    fields.append(f"hi7=0x{hi7:02x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 7. Opcode 0x15: perfcounter snapshot
+  # =====================================================================
+
+  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}")
+    if payload:
+      fields.append(f"payload=0x{payload:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 9. Pseudo opcode 0x10: never a "real" packet
+  # =====================================================================
+
+  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).
+
+  Now prints only the actual nibbles that belong to each packet, instead of
+  the full 64-bit shift register.
+  """
+  n = len(data)
+  time = 0
+  reg = 0          # shift register
+  offset = 0       # bit offset, in steps of 4 (one nibble)
+  nib_budget = 0x40
+  flags = 0
+  token_index = 0
+
+  while (offset >> 3) < n and token_index < max_tokens:
+    # Remember where we started refilling for this step (bit offset),
+    # but the *logical* start of the current packet is last_real_offset.
+    refill_start = offset
+
+    # 1) Fill register with nibbles according to nib_budget
+    if nib_budget != 0:
+      target = refill_start + 4 + ((nib_budget - 1) & ~3)
+      cur = refill_start
+      while cur != target and (cur >> 3) < n:
+        byte_index = cur >> 3
+        byte = data[byte_index]
+        shift = 4 if (cur & 4) else 0  # low then high nibble
+        nib = (byte >> shift) & 0xF
+        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
+        cur += 4
+      offset = cur
+
+    # 2) Decode token from low 8 bits
+    state = reg & 0xFF
+    opcode = STATE_TO_TOKEN[state]
+
+    # 3) Handle pseudo-token 0x10: need more bits, don't print. Looks like a NOP.
+    if opcode == 0x10:
+      # "need more bits" pseudo-token: adjust nibble budget and continue
+      nib_budget = 4
+      if (offset >> 3) >= n:
+        break
+      # Do NOT count this as a real packet; do not update last_real_offset.
+      continue
+
+    # 4) Set next nibble budget
+    nb_index = opcode & 0x1F
+    nib_budget = NIBBLE_BUDGET[nb_index]
+    time_before = time
+    note = ""
+    # 5) Special opcode 0x16 (timestamp / marker)
+    if opcode == 0x16:
+      two_bits = (reg >> 8) & 0x3
+      if two_bits == 1:
+        flags |= 0x01
+
+      # Common 36-bit field at bits [12..47]
+
+      if (reg & 0x200) == 0:
+        # delta mode: add 36-bit delta to time
+        delta = (reg >> 12) & ((1 << 36) - 1)
+        time += delta
+      else:
+        # marker / other modes: no time advance
+        if (reg & 0x100) == 0:
+          # real marker: bit9=1, bit8=0, non-zero payload
+          # "other" 0x16 variants, ignored for timing
+          delta = 0
+    else:
+      # 6) Generic opcode (including 0x0F)
+      shift, width = DELTA_MAP_DEFAULT[opcode]
+      mask = (1 << width) - 1
+      delta = (reg >> shift) & mask
+
+      # TODO: add more opcode parsers here that add notes to other opcodes
+      if opcode == 0x0F:
+        delta_with_fix = delta + 4
+        time += delta_with_fix
+        delta = delta_with_fix
+      else:
+        time += delta
+
+    # Append extra decoded fields into the note string
+    note = decode_packet_fields(opcode, reg, delta)
+
+    if filter is None or opcode not in filter:
+      my_reg = reg
+      my_reg &= (1 << nib_budget) - 1
+      print(
+        f"{token_index:4d}  "
+        f"off={offset//4:5d}  "
+        f"op=0x{opcode:02x} "
+        f"{OPCODE_NAMES[opcode]:24s} "
+        f" time={time_before:8d}+{delta:8d}  "
+        f"{my_reg:16X} "
+        f"{note}"
+      )
+
+    token_index += 1
+
+  # Optional summary at the end
+  print(f"# done: tokens={token_index}, final_time={time}, flags=0x{flags:02x}")
+
+def parse(fn:str):
+  dat = pickle.load(open(fn, "rb"))
+  ctx = decode(dat)
+  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
+  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
+  return dat_sqtt
+
+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")
+  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]
@@ -63,10 +64,10 @@ class _ROCParseCtx:
     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):
+  def on_occupancy_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_occupancy_t):
     if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
 
-  def on_wave_ev(self, ev):
+  def on_wave_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_wave_t):
     if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
 
     inst_execs:list[InstExec] = []
@@ -75,9 +76,11 @@ class _ROCParseCtx:
       inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst_ev.category)
       inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
       inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
+      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] = {}

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
@@ -39,10 +37,10 @@ def save_sqtt():
   sqtt:dict[str, list[WaveExec]] = {}
   yield sqtt
   # decode sqtt
-  if os.environ["DEV"] == "AMD":
-    rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
-    assert len(rctx.inst_execs) > 0, "empty sqtt output"
-    sqtt.update(rctx.inst_execs)
+  if os.environ["DEV"] != "AMD": return
+  rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
+  assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  sqtt.update(rctx.inst_execs)
 
 class TestTiming(unittest.TestCase):
   def test_v_add(self):
@@ -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"):
@@ -84,13 +81,17 @@ class TestTiming(unittest.TestCase):
 
   def test_wmma(self):
     with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_wmma_f32_16x16x16_f16 v[16:23], v[0:7], v[8:15], v[16:23]",
-        "v_add_f32_e32 v0 v16 v0",
-      ], l=32*4).realize()
-    assert len(sqtt) == 2, f"expected two waves, got {len(sqtt)} {list(sqtt.keys())}"
-    wmma = list(sqtt.values())[0][0]
-    self.assertGreater(wmma.dur, 1) # rgp says 32 clocks
+      for tc in dev.renderer.get_tensor_cores(dev.arch):
+        M, K, N = tc.dims
+        s = 32
+        a = Tensor.empty(M*s, K*s, dtype=tc.dtype_in)@Tensor.empty(K*s, N*s, dtype=tc.dtype_in)
+        a.realize()
+        print(a)
+    for p,waves in sqtt.items():
+      for e in waves[0].insts:
+        if (e.inst.startswith("v_wmma")):
+          instruction = e.inst.split(" ")[0]
+          print(f"{instruction:<29} : {e.dur} cycles")
 
   def test_sleep(self):
     n = 1
@@ -98,15 +99,35 @@ 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)
     diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
     with save_sqtt() as sqtt:
       diff_hw_reg.realize()
-    diff_sqtt = list(sqtt.values())[0][2]
-    self.assertEqual(diff_sqtt.dur, diff_hw_reg.item()-1) # 1 cycle for reading the counter register
+    sleep = next((e for e in sqtt[f"sleep_{n}"][0].insts if e.inst.startswith("s_sleep")))
+    # 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_arange.py

@@ -5,6 +5,8 @@ from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.uop.ops import Ops
+from tinygrad.renderer import Estimates
+from tinygrad.renderer.ptx import PTXRenderer
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -29,6 +31,14 @@ class TestArange(unittest.TestCase):
       # NOTE: not every backend supports CMPEQ
       self.assertLessEqual(self._get_flops(Tensor.eye(2560).contiguous(), np.eye(2560)), 2*2560*2560)
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
+  def test_tri_complexity(self):
+    with Context(NOOPT=1):
+      t = Tensor.ones(256, 256).contiguous().realize()
+      sched = t.triu().schedule()
+      p = get_program(sched[-1].ast)
+      self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)
+
 DSET, DDIM = 2048, 32
 
 class TestIndexing(unittest.TestCase):

test/test_custom_kernel.py

@@ -102,6 +102,11 @@ def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
 # **** tests ****
 
 class TestCustomKernel(unittest.TestCase):
+  def test_empty(self):
+    a = Tensor.empty(1)
+    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink())[0]
+    a.realize()
+
   def test_simple(self):
     a = Tensor.ones(16, 16).contiguous()
     b = Tensor.ones(16, 16).contiguous()

test/test_outerworld.py

@@ -50,47 +50,56 @@ 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()=}"
 
-  def test_range_grad(self):
-    def range_matmul(vec, mats):
-      # vec: (1, 10), mats: (3, 10, 10)
-      # assume vec, mats already have requires_grad set however you like
+class TestOuterScan(unittest.TestCase):
+  def _test_scan(self):
+    vec = Tensor.randn(1, 10).realize()
+    mats = Tensor.randn(3, 10, 10).realize()
 
-      i = UOp.range(3, -100, AxisType.OUTER)      # loop axis
-      vec_i = Tensor(vec.uop.after(i))            # "loop-carried" vector
-      vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
+    # 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
 
-      body = (vec_i.contiguous() @ mats[vi])      # matmul using loop index
-      out = Tensor(vec.uop.after(vec_i.uop.store(body.uop).end(i)))
-      return out
+  def test_uop_fold_matmul(self):
+    vec, mats, ref = self._test_scan()
 
-    vec = Tensor.randn(1, 3, requires_grad=True)
-    mats = Tensor.randn(3, 3, 3, requires_grad=True)
-    Tensor.realize(vec, mats)
+    # 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()
 
-    ref = ((vec @ mats[0]) @ mats[1]) @ mats[2]
-    loss = (1.0 - ref).square().mean()
-    loss.backward()
-    Tensor.realize(vec.grad, mats.grad)
-    print(vec.grad.numpy())
-    print(mats.grad.numpy())
-    vec.grad = None
-    mats.grad = None
+    # TODO: testing allclose
+    assert Tensor.allclose(ref[2], out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
 
-    out = range_matmul(vec, mats)
-    loss = (1.0 - out).square().mean()
-    loss.backward()
-    Tensor.realize(vec.grad, mats.grad)
+  def test_uop_scan_matmul(self):
+    vec, mats, ref = self._test_scan()
 
-    print(vec.grad, mats.grad)   # should be non-None and finite
-    print(vec.grad.numpy())
-    print(mats.grad.numpy())
+    # 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
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
 
 class TestOuterworld(unittest.TestCase):
   def test_range_plus_1(self):
@@ -151,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,17 +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,)),
-
-  # RANGE: loop index / axis, not a differentiable quantity
-  (UPat(Ops.RANGE), lambda: (None,)),
-
-  # STORE: buffer write. Gradient flows only into the value being stored.
-  # src layout is roughly (buffer, value, *axes_or_indices)
-  (UPat(Ops.STORE), lambda ctx: (None, ctx)),
-
-  # END: loop terminator / "end of range" node.
-  # Just pass the gradient into the body (first src), ignore the ranges.
-  (UPat(Ops.END, name="ret"), lambda ctx, ret: (ctx, *[None]*(len(ret.src) - 1))),
+  # 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/renderer/nir.py

@@ -5,7 +5,7 @@ from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa
-import base64, ctypes, ctypes.util, struct, functools, inspect
+import base64, contextlib, ctypes, ctypes.util, struct, functools, inspect
 
 def g(s:str): return getattr(mesa, s)
 def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))
@@ -157,8 +157,7 @@ class NIRRenderer(Renderer):
   def __init__(self): mesa.glsl_type_singleton_init_or_ref()
 
   def __del__(self):
-    try: mesa.glsl_type_singleton_decref()
-    except FileNotFoundError: pass
+    with contextlib.suppress(AttributeError):mesa.glsl_type_singleton_decref()
 
   @property
   def nir_options(self): raise NotImplementedError("needs nir_options")

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/ops_qcom.py

@@ -9,7 +9,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, adreno
 from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
-from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE
+from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE, suppress_finalizing
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
 
@@ -317,6 +317,7 @@ class QCOMAllocator(HCQAllocatorBase):
     self.dev.synchronize()
     return to_mv(cast(int, src.va_addr), src.size)
 
+  @suppress_finalizing
   def _free(self, opaque, options:BufferSpec):
     self.dev.synchronize()
     self.dev._gpu_free(opaque)

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/runtime/support/objc.py

@@ -4,12 +4,14 @@ from typing import TYPE_CHECKING, Any
 if TYPE_CHECKING: id_ = ctypes.c_void_p
 else:
   class id_(ctypes.c_void_p):
+    _is_finalizing = sys.is_finalizing # FIXME: why is this needed
+
     retain: bool = False
     # This prevents ctypes from converting response to plain int, and dict.fromkeys() can use it to dedup
     def __hash__(self): return hash(self.value)
     def __eq__(self, other): return self.value == other.value
     def __del__(self):
-      if self.retain and not sys.is_finalizing(): self.release()
+      if self.retain and not self._is_finalizing(): self.release()
     def release(self): msg("release")(self)
     def retained(self):
       setattr(self, 'retain', True)

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
@@ -2456,14 +2459,9 @@ class Tensor(OpMixin):
     return self._split_cumalu(axis, Ops.MAX)
 
   @staticmethod
-  def _tri(r:sint, c:sint, diagonal:int=0, **kwargs) -> Tensor:
+  def _tri(r:sint, c:sint, diagonal:int=0, device=None, requires_grad:bool|None=None) -> Tensor:
     assert isinstance(r, int) and isinstance(c, int), f"does not support symbolic, getting {r=}, {c=}"
-    if r == 0 or c == 0 or diagonal >= c: return Tensor.zeros(r,c,**kwargs)
-    if r+diagonal <= 0: return Tensor.ones(r,c,**kwargs)
-    s = r+c-1
-    # build a (s, s) upper triangle
-    t = Tensor.ones(s,s,**kwargs).pad((None,(0,s))).flatten().shrink(((0,s*(2*s-1)),)).reshape(s,-1).shrink((None,(0,s)))
-    return t[:r,-diagonal:c-diagonal] if diagonal <= 0 else t[diagonal:r+diagonal,:c]
+    return (Tensor.arange(r, device=device).unsqueeze(-1) + diagonal <= Tensor.arange(c, device=device)).requires_grad_(requires_grad)
 
   def triu(self, diagonal:int=0) -> Tensor:
     """
@@ -2486,7 +2484,7 @@ class Tensor(OpMixin):
     print(t.triu(diagonal=-1).numpy())
     ```
     """
-    return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal, device=self.device, dtype=dtypes.bool).where(self, self.zeros_like())
+    return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal, device=self.device).where(self, self.zeros_like())
 
   def tril(self, diagonal:int=0) -> Tensor:
     """
@@ -2509,7 +2507,7 @@ class Tensor(OpMixin):
     print(t.tril(diagonal=-1).numpy())
     ```
     """
-    return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal+1, device=self.device, dtype=dtypes.bool).where(self.zeros_like(), self)
+    return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal+1, device=self.device).where(self.zeros_like(), self)
 
   def interpolate(self, size:tuple[int, ...], mode:str="linear", align_corners:bool=False) -> Tensor:
     """

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