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/sqtt/active_sqtt_parse.py

@@ -7,17 +7,20 @@ os.environ["AMD_LLVM"] = "0"
 
 from dataclasses import replace
 import atexit, contextlib
-from tinygrad.helpers import system, getenv
+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
 
-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")
+# 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"]
 
@@ -37,8 +40,7 @@ def save_sqtt():
     if isinstance(e, ProfileSQTTEvent):
       print(replace(e, blob=b''))
       if e.se == 0:
-        parse_sqtt_print_packets(e.blob, filter=[0xf, 0x11, 0x12, 0x14] if getenv("FILTER", 1) else None)
-
+        parse_sqtt_print_packets(e.blob)
 
 template = """.text
 .globl matmul
@@ -51,6 +53,7 @@ matmul:
 .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
@@ -64,14 +67,21 @@ amdhsa.version:
 amdhsa.kernels:
   - .name: matmul
     .symbol: matmul.kd
-    .kernarg_segment_size: 0
     .group_segment_fixed_size: 0
     .private_segment_fixed_size: 0
-    .kernarg_segment_align: 4
     .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
 """
@@ -80,20 +90,42 @@ 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(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
+  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 v3 v2 v2",
-      "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 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"
-    ]*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,26 +1,39 @@
 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 = {
-  # ------------------------------------------------------------------------
-  # 0x01–0x06: small “meta + maybe tiny delta” packets
-  # ------------------------------------------------------------------------
-  0x01: "META_ID12_TS_SMALL",       # 12-bit ID + 3-bit delta field
-  0x02: "META_FLAG8_TS_SMALL",      # 8-bit flag/mode + small delta
-  0x03: "META_SUBEVENT8_TS_SMALL",  # 8-bit subevent/class + small delta
-  0x04: "META_BASE_INDEX12_TS",     # 12-bit base index + small delta
-  0x05: "META_DESC24_TS_A",         # 24-bit descriptor-ish + delta field
-  0x06: "META_DESC24_TS_B",         # second flavour, 24-bit, delta field
+  # 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)
-  0x08: "EVT_MATCH_SMALL",          # event-ish, see fields below
-  0x09: "PERF_ROUTE_CONFIG",        # routing/indirection config
   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)
@@ -34,15 +47,11 @@ OPCODE_NAMES = {
   0x10: "PSEUDO_NEED_MORE_BITS",    # not a real packet; decoder refill hint
 
   0x11: "TS_WAVE_STATE_SAMPLE",     # wave stall/termination sample (byte at +10)
-  0x12: "EVT_SECONDARY_METRIC24",   # 24-bit secondary timing/perf metric
   0x13: "EVT_SMALL_GENERIC",        # same structural family as 0x08/0x12/0x19
 
-  0x14: "INST_EXEC_OR_CFG",         # instruction exec record / config write / COR marker
   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
-  0x18: "PERF_EVENT_SELECT",        # packed selector → FUN_0010aba0
-  0x19: "EVT_SUMMARY_48B",          # 6-byte summary/aggregate metric
 }
 
 # these tables are from rocprof trace decoder
@@ -181,9 +190,8 @@ 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"delta36={delta}")
+    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),
@@ -408,7 +416,15 @@ def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
 
   return ", ".join(fields)
 
-def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000, filter=None) -> None:
+# 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).
 
@@ -466,23 +482,17 @@ def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000, filter=None)
         flags |= 0x01
 
       # Common 36-bit field at bits [12..47]
-      val36 = (reg >> 12) & ((1 << 36) - 1)
 
       if (reg & 0x200) == 0:
         # delta mode: add 36-bit delta to time
-        delta = val36
+        delta = (reg >> 12) & ((1 << 36) - 1)
         time += delta
-        note = "0x16-delta"
       else:
         # marker / other modes: no time advance
-        if (reg & 0x100) == 0 and val36 != 0:
+        if (reg & 0x100) == 0:
           # real marker: bit9=1, bit8=0, non-zero payload
-          delta = 0
-          note = f"0x16-marker val=0x{val36:x}"
-        else:
           # "other" 0x16 variants, ignored for timing
           delta = 0
-          note = "0x16-other"
     else:
       # 6) Generic opcode (including 0x0F)
       shift, width = DELTA_MAP_DEFAULT[opcode]
@@ -492,19 +502,13 @@ def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000, filter=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)
 
     if filter is None or opcode not in filter:
       my_reg = reg
@@ -533,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] = {}

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/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")

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/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/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