IQ.Pilot Prebuilt Release @ ab07000

2026-08-22 23:42:42 -05:00
commit 9f9c9a70cc
3729 changed files with 778697 additions and 0 deletions
--- a/tinygrad_repo/tinygrad/viz/README
+++ b/tinygrad_repo/tinygrad/viz/README
@@ -0,0 +1,61 @@
+VIZ is a tool for inspecting tinygrad's compilation process and performance profiling.
+
+to use:
+1. Run tinygrad with VIZ=1 (this saves the pkls and launches the server in interactive shells)
+2. That's it!
+
+This can:
+1. See all schedules
+2. See all graphs and how they were rewritten
+3. See generated code
+4. See profile
+5. (AMD only) see instruction level SQTT profile
+
+VIZ pkls can be viewed in two ways:
+
+1. Web browser: python -m tinygrad.viz.serve
+2. Command line: python -m tinygrad.viz.cli
+
+By default, VIZ UIs automatically load the latest files.
+
+user story: viewing profiling data
+* tinygrad ran 32 LLM decode steps:
+web: click "profiler", view the timeline of all python codegen and GPU kernels.
+cli: Run `DEBUG=3 python -m tinygrad.viz.cli --json` to extract kernel timing info and ASTs in JSON format.
+  - note: Make sure to add NO_COLOR=1 to disable colored output.
+
+user story: viewing code
+* tinygrad ran 3 schedules: init the model + first train step, train step, test step
+  * schedule 1 (123) = main.py:97
+  * schedule 2 (97) = main.py:97
+  * schedule 3 (10) = main.py:145
+* web: click "schedule 1", get list of kernels (like DEBUG=2)
+* cli: `python -m tinygrad.viz.cli -s TINY "Schedule 3 Kernels n1"`
+  * kernel 1 "E_34_34" -- 'sin'
+  * kernel 2 "R_4545"
+* web: click "E_34_34"
+* cli: `python -m tinygrad.viz.cli -s TINY "do_to_program for E_34_34" "initial symbolic"`
+  * pre-rewritten UOp graph (step through rewrite here)
+  * post-rewritten UOp graph
+  * UOp list
+  * generated code
+
+user story: debugging scheduler
+* tinygrad ran 3 schedules: init the model + first train step, train step, test step
+  * ...
+* click "schedule 1 graph", get a graph of the schedule in UOps
+  * step through rewrite rules
+  * see how things are broken into kernels
+  * see why two kernels didn't fuse
+
+user story: SQTT / PMC profiling
+
+note: SQTT has additional overhead, to enable it, set VIZ=2.
+
+* tinygrad ran custom assembly GEMM kernel.
+* web: click "SQTT gemm SE:1 PKTS", see wave instruction scheduling and CU execution unit occupancy at every clock cycle.
+* cli: python -m tinygrad.viz.cli -s "kernel SQTT SE:0 PKTS"
+
+* get bank conflicts:
+* web: click "gemm PMC"
+* cli: python -m tinygrad.viz.cli -s "gemm PMC" | rg -A 16 SQC_LDS_BANK_CONFLICT
--- a/tinygrad_repo/tinygrad/viz/init.py
+++ b/tinygrad_repo/tinygrad/viz/init.py
--- a/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js
+++ b/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js
--- a/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js
+++ b/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js
@@ -0,0 +1,47 @@
+/*! `cpp` grammar compiled for Highlight.js 11.10.0 */
+(()=>{var e=(()=>{"use strict";return e=>{const t=e.regex,a=e.COMMENT("//","$",{
+contains:[{begin:/\\\n/}]
+}),n="decltype\\(auto\\)",r="[a-zA-Z_]\\w*::",i="(?!struct)("+n+"|"+t.optional(r)+"[a-zA-Z_]\\w*"+t.optional("<[^<>]+>")+")",s={
+className:"type",begin:"\\b[a-z\\d_]*_t\\b"},c={className:"string",variants:[{
+begin:'(u8?|U|L)?"',end:'"',illegal:"\\n",contains:[e.BACKSLASH_ESCAPE]},{
+begin:"(u8?|U|L)?'(\\\\(x[0-9A-Fa-f]{2}|u[0-9A-Fa-f]{4,8}|[0-7]{3}|\\S)|.)",
+end:"'",illegal:"."},e.END_SAME_AS_BEGIN({
+begin:/(?:u8?|U|L)?R"([^()\\ ]{0,16})\(/,end:/\)([^()\\ ]{0,16})"/})]},o={
+className:"number",variants:[{
+begin:"[+-]?(?:(?:[0-9](?:'?[0-9])*\\.(?:[0-9](?:'?[0-9])*)?|\\.[0-9](?:'?[0-9])*)(?:[Ee][+-]?[0-9](?:'?[0-9])*)?|[0-9](?:'?[0-9])*[Ee][+-]?[0-9](?:'?[0-9])*|0[Xx](?:[0-9A-Fa-f](?:'?[0-9A-Fa-f])*(?:\\.(?:[0-9A-Fa-f](?:'?[0-9A-Fa-f])*)?)?|\\.[0-9A-Fa-f](?:'?[0-9A-Fa-f])*)[Pp][+-]?[0-9](?:'?[0-9])*)(?:[Ff](?:16|32|64|128)?|(BF|bf)16|[Ll]|)"
+},{
+begin:"[+-]?\\b(?:0[Bb][01](?:'?[01])*|0[Xx][0-9A-Fa-f](?:'?[0-9A-Fa-f])*|0(?:'?[0-7])*|[1-9](?:'?[0-9])*)(?:[Uu](?:LL?|ll?)|[Uu][Zz]?|(?:LL?|ll?)[Uu]?|[Zz][Uu]|)"
+}],relevance:0},l={className:"meta",begin:/#\s*[a-z]+\b/,end:/$/,keywords:{
+keyword:"if else elif endif define undef warning error line pragma _Pragma ifdef ifndef include"
+},contains:[{begin:/\\\n/,relevance:0},e.inherit(c,{className:"string"}),{
+className:"string",begin:/<.*?>/},a,e.C_BLOCK_COMMENT_MODE]},u={
+className:"title",begin:t.optional(r)+e.IDENT_RE,relevance:0
+},d=t.optional(r)+e.IDENT_RE+"\\s*\\(",p={
+type:["bool","char","char16_t","char32_t","char8_t","double","float","int","long","short","void","wchar_t","unsigned","signed","const","static"],
+keyword:["alignas","alignof","and","and_eq","asm","atomic_cancel","atomic_commit","atomic_noexcept","auto","bitand","bitor","break","case","catch","class","co_await","co_return","co_yield","compl","concept","const_cast|10","consteval","constexpr","constinit","continue","decltype","default","delete","do","dynamic_cast|10","else","enum","explicit","export","extern","false","final","for","friend","goto","if","import","inline","module","mutable","namespace","new","noexcept","not","not_eq","nullptr","operator","or","or_eq","override","private","protected","public","reflexpr","register","reinterpret_cast|10","requires","return","sizeof","static_assert","static_cast|10","struct","switch","synchronized","template","this","thread_local","throw","transaction_safe","transaction_safe_dynamic","true","try","typedef","typeid","typename","union","using","virtual","volatile","while","xor","xor_eq"],
+literal:["NULL","false","nullopt","nullptr","true"],built_in:["_Pragma"],
+_type_hints:["any","auto_ptr","barrier","binary_semaphore","bitset","complex","condition_variable","condition_variable_any","counting_semaphore","deque","false_type","future","imaginary","initializer_list","istringstream","jthread","latch","lock_guard","multimap","multiset","mutex","optional","ostringstream","packaged_task","pair","promise","priority_queue","queue","recursive_mutex","recursive_timed_mutex","scoped_lock","set","shared_future","shared_lock","shared_mutex","shared_timed_mutex","shared_ptr","stack","string_view","stringstream","timed_mutex","thread","true_type","tuple","unique_lock","unique_ptr","unordered_map","unordered_multimap","unordered_multiset","unordered_set","variant","vector","weak_ptr","wstring","wstring_view"]
+},_={className:"function.dispatch",relevance:0,keywords:{
+_hint:["abort","abs","acos","apply","as_const","asin","atan","atan2","calloc","ceil","cerr","cin","clog","cos","cosh","cout","declval","endl","exchange","exit","exp","fabs","floor","fmod","forward","fprintf","fputs","free","frexp","fscanf","future","invoke","isalnum","isalpha","iscntrl","isdigit","isgraph","islower","isprint","ispunct","isspace","isupper","isxdigit","labs","launder","ldexp","log","log10","make_pair","make_shared","make_shared_for_overwrite","make_tuple","make_unique","malloc","memchr","memcmp","memcpy","memset","modf","move","pow","printf","putchar","puts","realloc","scanf","sin","sinh","snprintf","sprintf","sqrt","sscanf","std","stderr","stdin","stdout","strcat","strchr","strcmp","strcpy","strcspn","strlen","strncat","strncmp","strncpy","strpbrk","strrchr","strspn","strstr","swap","tan","tanh","terminate","to_underlying","tolower","toupper","vfprintf","visit","vprintf","vsprintf"]
+},
+begin:t.concat(/\b/,/(?!decltype)/,/(?!if)/,/(?!for)/,/(?!switch)/,/(?!while)/,e.IDENT_RE,t.lookahead(/(<[^<>]+>|)\s*\(/))
+},m=[_,l,s,a,e.C_BLOCK_COMMENT_MODE,o,c],f={variants:[{begin:/=/,end:/;/},{
+begin:/\(/,end:/\)/},{beginKeywords:"new throw return else",end:/;/}],
+keywords:p,contains:m.concat([{begin:/\(/,end:/\)/,keywords:p,
+contains:m.concat(["self"]),relevance:0}]),relevance:0},g={className:"function",
+begin:"("+i+"[\\*&\\s]+)+"+d,returnBegin:!0,end:/[{;=]/,excludeEnd:!0,
+keywords:p,illegal:/[^\w\s\*&:<>.]/,contains:[{begin:n,keywords:p,relevance:0},{
+begin:d,returnBegin:!0,contains:[u],relevance:0},{begin:/::/,relevance:0},{
+begin:/:/,endsWithParent:!0,contains:[c,o]},{relevance:0,match:/,/},{
+className:"params",begin:/\(/,end:/\)/,keywords:p,relevance:0,
+contains:[a,e.C_BLOCK_COMMENT_MODE,c,o,s,{begin:/\(/,end:/\)/,keywords:p,
+relevance:0,contains:["self",a,e.C_BLOCK_COMMENT_MODE,c,o,s]}]
+},s,a,e.C_BLOCK_COMMENT_MODE,l]};return{name:"C++",
+aliases:["cc","c++","h++","hpp","hh","hxx","cxx"],keywords:p,illegal:"</",
+classNameAliases:{"function.dispatch":"built_in"},
+contains:[].concat(f,g,_,m,[l,{
+begin:"\\b(deque|list|queue|priority_queue|pair|stack|vector|map|set|bitset|multiset|multimap|unordered_map|unordered_set|unordered_multiset|unordered_multimap|array|tuple|optional|variant|function)\\s*<(?!<)",
+end:">",keywords:p,contains:["self",s]},{begin:e.IDENT_RE+"::",keywords:p},{
+match:[/\b(?:enum(?:\s+(?:class|struct))?|class|struct|union)/,/\s+/,/\w+/],
+className:{1:"keyword",3:"title.class"}}])}}})();hljs.registerLanguage("cpp",e)
+})();
--- a/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js
+++ b/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js
@@ -0,0 +1,42 @@
+/*! `python` grammar compiled for Highlight.js 11.10.0 */
+(()=>{var e=(()=>{"use strict";return e=>{
+const n=e.regex,a=/[\p{XID_Start}_]\p{XID_Continue}*/u,s=["and","as","assert","async","await","break","case","class","continue","def","del","elif","else","except","finally","for","from","global","if","import","in","is","lambda","match","nonlocal|10","not","or","pass","raise","return","try","while","with","yield"],t={
+$pattern:/[A-Za-z]\w+|__\w+__/,keyword:s,
+built_in:["__import__","abs","all","any","ascii","bin","bool","breakpoint","bytearray","bytes","callable","chr","classmethod","compile","complex","delattr","dict","dir","divmod","enumerate","eval","exec","filter","float","format","frozenset","getattr","globals","hasattr","hash","help","hex","id","input","int","isinstance","issubclass","iter","len","list","locals","map","max","memoryview","min","next","object","oct","open","ord","pow","print","property","range","repr","reversed","round","set","setattr","slice","sorted","staticmethod","str","sum","super","tuple","type","vars","zip"],
+literal:["__debug__","Ellipsis","False","None","NotImplemented","True"],
+type:["Any","Callable","Coroutine","Dict","List","Literal","Generic","Optional","Sequence","Set","Tuple","Type","Union"]
+},i={className:"meta",begin:/^(>>>|\.\.\.) /},r={className:"subst",begin:/\{/,
+end:/\}/,keywords:t,illegal:/#/},l={begin:/\{\{/,relevance:0},o={
+className:"string",contains:[e.BACKSLASH_ESCAPE],variants:[{
+begin:/([uU]|[bB]|[rR]|[bB][rR]|[rR][bB])?'''/,end:/'''/,
+contains:[e.BACKSLASH_ESCAPE,i],relevance:10},{
+begin:/([uU]|[bB]|[rR]|[bB][rR]|[rR][bB])?"""/,end:/"""/,
+contains:[e.BACKSLASH_ESCAPE,i],relevance:10},{
+begin:/([fF][rR]|[rR][fF]|[fF])'''/,end:/'''/,
+contains:[e.BACKSLASH_ESCAPE,i,l,r]},{begin:/([fF][rR]|[rR][fF]|[fF])"""/,
+end:/"""/,contains:[e.BACKSLASH_ESCAPE,i,l,r]},{begin:/([uU]|[rR])'/,end:/'/,
+relevance:10},{begin:/([uU]|[rR])"/,end:/"/,relevance:10},{
+begin:/([bB]|[bB][rR]|[rR][bB])'/,end:/'/},{begin:/([bB]|[bB][rR]|[rR][bB])"/,
+end:/"/},{begin:/([fF][rR]|[rR][fF]|[fF])'/,end:/'/,
+contains:[e.BACKSLASH_ESCAPE,l,r]},{begin:/([fF][rR]|[rR][fF]|[fF])"/,end:/"/,
+contains:[e.BACKSLASH_ESCAPE,l,r]},e.APOS_STRING_MODE,e.QUOTE_STRING_MODE]
+},b="[0-9](_?[0-9])*",c=`(\\b(${b}))?\\.(${b})|\\b(${b})\\.`,d="\\b|"+s.join("|"),g={
+className:"number",relevance:0,variants:[{
+begin:`(\\b(${b})|(${c}))[eE][+-]?(${b})[jJ]?(?=${d})`},{begin:`(${c})[jJ]?`},{
+begin:`\\b([1-9](_?[0-9])*|0+(_?0)*)[lLjJ]?(?=${d})`},{
+begin:`\\b0[bB](_?[01])+[lL]?(?=${d})`},{begin:`\\b0[oO](_?[0-7])+[lL]?(?=${d})`
+},{begin:`\\b0[xX](_?[0-9a-fA-F])+[lL]?(?=${d})`},{begin:`\\b(${b})[jJ](?=${d})`
+}]},p={className:"comment",begin:n.lookahead(/# type:/),end:/$/,keywords:t,
+contains:[{begin:/# type:/},{begin:/#/,end:/\b\B/,endsWithParent:!0}]},m={
+className:"params",variants:[{className:"",begin:/\(\s*\)/,skip:!0},{begin:/\(/,
+end:/\)/,excludeBegin:!0,excludeEnd:!0,keywords:t,
+contains:["self",i,g,o,e.HASH_COMMENT_MODE]}]};return r.contains=[o,g,i],{
+name:"Python",aliases:["py","gyp","ipython"],unicodeRegex:!0,keywords:t,
+illegal:/(<\/|\?)|=>/,contains:[i,g,{scope:"variable.language",match:/\bself\b/
+},{beginKeywords:"if",relevance:0},{match:/\bor\b/,scope:"keyword"
+},o,p,e.HASH_COMMENT_MODE,{match:[/\bdef/,/\s+/,a],scope:{1:"keyword",
+3:"title.function"},contains:[m]},{variants:[{
+match:[/\bclass/,/\s+/,a,/\s*/,/\(\s*/,a,/\s*\)/]},{match:[/\bclass/,/\s+/,a]}],
+scope:{1:"keyword",3:"title.class",6:"title.class.inherited"}},{
+className:"meta",begin:/^[\t ]*@/,end:/(?=#)|$/,contains:[g,m,o]}]}}})()
+;hljs.registerLanguage("python",e)})();
--- a/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/x86asm.min.js
+++ b/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/x86asm.min.js
--- a/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/styles/default.min.css
+++ b/tinygrad_repo/tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/styles/default.min.css
@@ -0,0 +1,9 @@
+/*!
+  Theme: Default
+  Description: Original highlight.js style
+  Author: (c) Ivan Sagalaev <maniac@softwaremaniacs.org>
+  Maintainer: @highlightjs/core-team
+  Website: https://highlightjs.org/
+  License: see project LICENSE
+  Touched: 2021
+*/pre code.hljs{display:block;overflow-x:auto;padding:1em}code.hljs{padding:3px 5px}.hljs{background:#f3f3f3;color:#444}.hljs-comment{color:#697070}.hljs-punctuation,.hljs-tag{color:#444a}.hljs-tag .hljs-attr,.hljs-tag .hljs-name{color:#444}.hljs-attribute,.hljs-doctag,.hljs-keyword,.hljs-meta .hljs-keyword,.hljs-name,.hljs-selector-tag{font-weight:700}.hljs-deletion,.hljs-number,.hljs-quote,.hljs-selector-class,.hljs-selector-id,.hljs-string,.hljs-template-tag,.hljs-type{color:#800}.hljs-section,.hljs-title{color:#800;font-weight:700}.hljs-link,.hljs-operator,.hljs-regexp,.hljs-selector-attr,.hljs-selector-pseudo,.hljs-symbol,.hljs-template-variable,.hljs-variable{color:#ab5656}.hljs-literal{color:#695}.hljs-addition,.hljs-built_in,.hljs-bullet,.hljs-code{color:#397300}.hljs-meta{color:#1f7199}.hljs-meta .hljs-string{color:#38a}.hljs-emphasis{font-style:italic}.hljs-strong{font-weight:700}
--- a/tinygrad_repo/tinygrad/viz/assets/d3js.org/d3.v7.min.js
+++ b/tinygrad_repo/tinygrad/viz/assets/d3js.org/d3.v7.min.js
--- a/tinygrad_repo/tinygrad/viz/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js
+++ b/tinygrad_repo/tinygrad/viz/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js
--- a/tinygrad_repo/tinygrad/viz/assets/unpkg.com/@highlightjs/cdn-assets@11.10.0/styles/tokyo-night-dark.min.css
+++ b/tinygrad_repo/tinygrad/viz/assets/unpkg.com/@highlightjs/cdn-assets@11.10.0/styles/tokyo-night-dark.min.css
@@ -0,0 +1,8 @@
+pre code.hljs{display:block;overflow-x:auto;padding:1em}code.hljs{padding:3px 5px}/*!
+  Theme: Tokyo-night-Dark
+  origin: https://github.com/enkia/tokyo-night-vscode-theme
+  Description: Original highlight.js style
+  Author: (c) Henri Vandersleyen <hvandersleyen@gmail.com>
+  License: see project LICENSE
+  Touched: 2022
+*/.hljs-comment,.hljs-meta{color:#565f89}.hljs-deletion,.hljs-doctag,.hljs-regexp,.hljs-selector-attr,.hljs-selector-class,.hljs-selector-id,.hljs-selector-pseudo,.hljs-tag,.hljs-template-tag,.hljs-variable.language_{color:#f7768e}.hljs-link,.hljs-literal,.hljs-number,.hljs-params,.hljs-template-variable,.hljs-type,.hljs-variable{color:#ff9e64}.hljs-attribute,.hljs-built_in{color:#e0af68}.hljs-keyword,.hljs-property,.hljs-subst,.hljs-title,.hljs-title.class_,.hljs-title.class_.inherited__,.hljs-title.function_{color:#7dcfff}.hljs-selector-tag{color:#73daca}.hljs-addition,.hljs-bullet,.hljs-quote,.hljs-string,.hljs-symbol{color:#9ece6a}.hljs-code,.hljs-formula,.hljs-section{color:#7aa2f7}.hljs-attr,.hljs-char.escape_,.hljs-keyword,.hljs-name,.hljs-operator{color:#bb9af7}.hljs-punctuation{color:#c0caf5}.hljs{background:#1a1b26;color:#9aa5ce}.hljs-emphasis{font-style:italic}.hljs-strong{font-weight:700}
--- a/tinygrad_repo/tinygrad/viz/cli.py
+++ b/tinygrad_repo/tinygrad/viz/cli.py
@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+import argparse, pathlib, signal, struct, json, os, itertools, heapq
+os.environ["VIZ"] = "0"
+if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+from typing import Iterator
+from tinygrad.viz import serve as viz
+from tinygrad.viz.serve import fmt_colored
+from tinygrad.uop.ops import RewriteTrace
+from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen, ProfilePointEvent, ProfileRangeEvent, TracingKey, unwrap, NO_COLOR, DEBUG
+
+# profile decoder used in CLI and tests
+def decode_profile(data:bytes) -> dict:
+  ret, off = data, 0
+  def u(fmt:str) -> tuple:
+    nonlocal off
+    vals = struct.unpack_from(fmt, ret, off)
+    off += struct.calcsize(fmt)
+    return vals
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
+  off += index_len
+  layout:dict[str, dict] = {}
+  # 0 means None, otherwise it's an enum value
+  def option(i:int) -> int|None: return None if i == 0 else i-1
+  for _ in range(layout_len):
+    klen = u("<B")[0]
+    k = ret[off:off+klen].decode()
+    off += klen
+    event_type, event_count = u("<BI")
+    layout[k] = v = {"event_type":event_type, "events":[]}
+    if event_type == 0:
+      for _ in range(event_count):
+        name, ref, key, st, dur, fmt = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":json.loads(strings[fmt])})
+    else:
+      v["linear"] = u("<B")[0]
+      v["peak"] = u("<Q")[0]
+      for _ in range(event_count):
+        if v["linear"]:
+          ts, value = u("<IQ")
+          v["events"].append({"event":"freq", "ts":ts, "value":value})
+        else:
+          alloc, ts, key = u("<BII")
+          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[(k, strings[rep], num, mode) \
+              for k,rep,num,mode in [u("<IIIB") for _ in range(u("<I")[0])]]}})
+  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
+
+def to_str(k:str, v) -> str:
+  if k == "FLOPS" or k.startswith("B/s"): return f"{v*1e-9:.0f} G{k}" if v < 1e13 else f"{v*1e-12:.0f} T{k}"
+  if k == "B": return next((f"{v/s:.0f} {u}" for s,u in ((1e9,"GB"),(1e6,"MB"),(1e3,"KB")) if v>=s), f"{v:.0f} B")
+  return f"{k}={v}"
+def fmt_data(data:dict) -> str: return "  ".join((p:=to_str(k, v))+" "*max(0, 14-ansilen(p)) for k,v in data.items())
+
+def marker_st(markers:list[dict], name:str) -> int:
+  try: return next(e["ts"] for e in markers if e["name"] == name)
+  except StopIteration: raise RuntimeError(f"marker not found: {name}") from None
+
+def get(data:dict, key:str):
+  for k,v in data.items():
+    if ansistrip(k) == key: return v
+  import difflib
+  match = difflib.get_close_matches(key, [ansistrip(k) for k in data], n=1, cutoff=0.6)
+  raise RuntimeError(f'item "{key}" not found in list'+(f", did you mean {match[0]!r}?" if match else ''))
+
+def main(args) -> None:
+  viz.load_rewrites(viz_data:=viz.VizData(viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))))
+
+  def emit(val, to_str=str) -> str: return json.dumps(val if isinstance(val, dict) else {"value":val}) if args.json else to_str(val)
+
+  def print_step(step:dict, print_graph=False, reconstruct_matches=False) -> None:
+    data = viz.get_render(viz_data, step["query"])
+    if isinstance(data.get("value"), Iterator):
+      for m in data["value"]:
+        if print_graph and "graph" in m and not args.json:
+          for k,v in m["graph"].items():
+            print(f"[{k}] {' '.join((lines:=v['label'].splitlines())[:5])}{'...' if len(lines) > 5 else ''}"+(f" tag={v['tag']}" if v['tag'] else ''))
+            if v["src"]:
+              print("  src: "+", ".join([f"{i}->[{x}]" for i,x in v["src"][:5]])+(f", ... and {len(v['src'])-5} more" if len(v["src"]) > 5 else ""))
+        elif "uop" in m: print(emit(m["graph"] if print_graph else m["uop"]))
+        if not reconstruct_matches: return None
+        if m.get("diff"):
+          loc = pathlib.Path(m["upat"][0][0])
+          print(emit(f"{loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}"))
+          for line in m["diff"]: print(emit(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None)))
+    if data.get("src") is not None: print(emit(data["src"]))
+
+  profile_bytes = viz.get_profile(viz_data, viz.load_pickle(args.profile_path, default=[]))
+  if profile_bytes is None: raise RuntimeError(f"empty profile in {args.profile_path}")
+  profile = decode_profile(profile_bytes)
+  profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in viz_data.ctxs
+                            if c["name"].startswith("SQTT") for s in c["steps"] if s["name"].endswith(("PMC", "PKTS"))])
+  if args.list and not args.src: return print("\n".join(emit(fmt_colored(k)) for k in ["ALL"]+list(profile["layout"])))
+
+  # ** SQTT printer
+  data = None if not args.src else get(profile["layout"], args.src[0])
+  if args.src and "SQTT" in args.src[0]:
+    # modern terminals support 24-bit color
+    def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16)};{int(color[3:5],16)};{int(color[5:7],16)}m{st}\x1b[0m"
+    print(emit(f"{'Clk':<12} {'Unit':<20} {'Op':<22} {'Dur':<4} {'Delay':<4} {'Info'}"))
+    print(emit("-" * 100))
+    pc_map:dict[int, str] = {}
+    pkt_idxs:dict[str, itertools.count] = {}
+    dispatch_to_inst:dict[str, tuple[str, int]] = {}
+    inst_st:int|None = None
+    for e in viz.sqtt_timeline(*unwrap(data)):
+      if isinstance(e, ProfilePointEvent) and e.key == 'pcMap': pc_map = e.arg
+      if not isinstance(e, ProfileRangeEvent): continue
+      if inst_st is None: inst_st = int(e.st)
+      assert isinstance(e.name, TracingKey)
+      op_name, ret, info = e.name.display_name, json.loads(e.name.ret[4:]) if e.name.ret else {}, ""
+      color = next((v for k,v in viz.wave_colors.items() if k in op_name), None)
+      op_str = hex_colored(op_name, color) if color and not NO_COLOR else op_name
+      inst, phase, delay = None, None, 0
+      idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
+      if e.device.startswith("WAVE"):
+        inst = f"0x{pc:05x} {pc_map[pc]}" if (pc:=ret.get("pc")) is not None else f"{'':7} {op_name}"
+        dispatch_to_inst[f"{e.device}-{idx}"] = (inst, int(e.st))
+        phase = "DISPATCH"
+      if (link:=ret.get("link")) is not None:
+        inst, dispatch_st = dispatch_to_inst[link]
+        phase, delay = "EXEC", int(e.st) - dispatch_st
+      if inst and phase: info = f"{phase:<8} {inst}"
+      unit = e.device.replace(" ", "-")
+      row = {"clk":int(e.st)-inst_st, "cycle":int(e.st), "unit":unit, "op":op_name, "dur":int(unwrap(e.en)-e.st), "delay":delay or "", "info":info}
+      print(emit(row, lambda _: f"{row['clk']:<12} {unit:<20} {op_str}{' '*(22-ansilen(op_str))} {row['dur']:<4} {str(row['delay']):<4} {info}"))
+
+  # ** PMC printer
+  elif args.src and "PMC" in args.src[0]:
+    pmc = viz.unpack_pmc(unwrap(data))
+    pmc_fmt:list[str] = []
+    for name,val,*detail in pmc["rows"]:
+      pmc_fmt += [f"{name} {val}"]+([" ".join(f"{k}={v}" for k,v in zip(detail[0]["cols"], r)) for r in detail[0]["rows"]] if detail else [])
+    print(emit(pmc, lambda _: "\n".join(pmc_fmt)))
+
+  # ** Memory printer
+  elif data is not None and data["event_type"] == 1:
+    print(emit({"peak":data["peak"]}, lambda _: f"Peak: {data['peak']}"+"\n"+f"{'TS':<10}  {'Event':<6}  {'Key':>8}  Info"))
+    for e in data["events"]:
+      info = str(arg:=e.pop("arg", {}))
+      if e["event"] == "free": info = ', '.join([f"{fmt_colored(k)} {['read','write','write+read'][m]}@data{n}" for _,k,n,m in arg["users"]])
+      print(emit({**e, "info":info}, lambda _: f"{e['ts']:<10}  {e['event']:<6}  {e.get('key', ''):>8}  {info}"))
+
+  # ** Profiler printer
+  else:
+    timelines = [(n,l) for n,l in profile["layout"].items() if isinstance(l, dict) and l.get("event_type") == 0]
+    markers = profile.get("markers", [])
+    interval:tuple[int, int]|None = None if not args.interval else (marker_st(markers, args.interval[0]), marker_st(markers, args.interval[1]))
+    def produce_top_kernels() -> Iterator[dict]:
+      tagged = ((n,e) for n,l in timelines for e in l["events"]) if not args.src else ((args.src[0],e) for e in unwrap(data)["events"])
+      agg:dict[tuple[str,str], tuple[float, int, int|None, dict[str, float]]] = {} # map (device, kernel name) to (total time, count, ref, est)
+      est_keys = ("FLOPS", "B/s mem", "B/s lds")
+      total = 0
+      for dev,e in tagged:
+        if interval and not interval[0] <= e["st"] <= interval[1]: continue
+        et = e["dur"] * 1e-3
+        t, c, ref, est = agg.get((dev,e["name"]), (0.0, 0, None, {}))
+        est.update({k:est.get(k, 0.0)+e["fmt"][k]*e["dur"]*1e-6 for k in est_keys if k in e["fmt"]})
+        agg[(dev,e["name"])] = (t+et, c+1, e["ref"], est)
+        total += et
+      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
+      num_rows = len(items) if args.t < 0 else args.t
+      for (dev,name),(t,c,ref,est) in items[:num_rows]:
+        display = f"{dev[:7]:7s} {fmt_colored(name)}" if not args.src else fmt_colored(name)
+        yield {"name":display, "dur_ms":t, "count":c, "pct":t/total*100.0, "ref":ref, "fmt":{k:int(est[k]/(t*1e-3)) for k in est_keys if k in est}}
+      if num_rows > 0 and items[num_rows:]:
+        other_t = sum(t for _,(t,_,_,_) in items[num_rows:])
+        other_c = sum(c for _,(_,c,_,_) in items[num_rows:])
+        yield {"name":"Other", "dur_ms":other_t, "count":other_c, "pct":other_t/total*100.0, "ref":None, "fmt":None}
+    def produce_all_kernels() -> Iterator[dict]:
+      event_streams = [[(e["st"], n, e) for e in l["events"]] for n,l in timelines] if not args.src \
+                      else [[(e["st"], args.src[0], e) for e in unwrap(data)["events"]]]
+      if not args.src:
+        for n,l in profile["layout"].items():
+          if not isinstance(l, dict) or l.get("event_type") != 0: yield {"device":"SOURCE", "name":n, "st_ms":0, "ref":None, "ext":None}
+      marker_stream = sorted([(m["ts"], "MARKER", m) for m in markers], key=lambda t:t[0])
+      for ts,dev,e in heapq.merge(*event_streams, marker_stream, key=lambda t:t[0]):
+        if interval is not None and not interval[0] <= ts <= interval[1]: continue
+        if dev == "MARKER":
+          yield {"device":dev, "name":fmt_colored(e["name"]), "st_ms":ts*1e-3, "ref":None, "ext":None}
+          continue
+        ext, fmt = [], e["fmt"]
+        if (tb:=fmt.pop("tb", [])):
+          while tb:
+            file, lineno, fxn, code = tb.pop()
+            line = f"{file.split('/')[-1]}:{lineno} {fxn}"
+            if fmt: ext.append(f"{line} {code}")
+            elif not file.startswith("<") and not fxn.startswith("<"): fmt["loc"] = line
+        yield {"device":dev, "name":fmt_colored(e["name"]), "dur_ms":e["dur"]*1e-3, "st_ms":e["st"]*1e-3, "fmt":fmt, "ref":e["ref"],
+               "ext":"\n".join(ext)}
+    def fmt_top(k:dict) -> str:
+      return f"{fmt_colored(k['name'])}{' ' * max(0, 38-ansilen(k['name']))} {time_to_str(k['dur_ms']*1e-3, w=9)} {k['count']:7d} {k['pct']:6.2f}%"+\
+          (" "*4+fmt_data(k['fmt']) if k['fmt'] else "")
+    def fmt_all(k:dict) -> str:
+      if k["device"] in {"MARKER", "SOURCE"}: return f"--- {k['device']} {k['name']}"+(f"/{k['st_ms']:9.2f}ms" if k['st_ms'] else "")
+      ptm = colored(time_to_str(k["dur_ms"]*1e-3, w=9), "yellow" if k["dur_ms"] > 10 else None)
+      name = f"*** {k['device'][:7]:7s} "+k["name"]+" "*(46-ansilen(k["name"]))
+      return f"{name} tm {ptm}/{k['st_ms']:9.2f}ms"+(f" ({fmt_data(k['fmt'])})" if k["fmt"] else "")
+    fmt_row = fmt_top if args.t else fmt_all
+    seen_refs:set[int] = set()
+    def render_event(k:dict, ls=args.list) -> None:
+      if len(args.src) > 1 and ansistrip(k["name"]) not in args.src: return None
+      print(emit(k, to_str=fmt_row))
+      if k["ref"] is not None and k["ref"] not in seen_refs:
+        seen_refs.add(k["ref"])
+        for i,s in enumerate(viz_data.ctxs[k["ref"]]["steps"]):
+          if DEBUG >= 3 and s["name"] == "View Base AST": print_step(s)
+          if DEBUG >= 4 and s["name"] == "View Source": print_step(s)
+          if DEBUG >= 5 or ls: print(emit(" "*s["depth"]+s["name"]+(f" - {s['match_count']}" if s.get('match_count', 0) else '')))
+          if DEBUG >= 6 or (DEBUG >= 5 and s["name"] == "View Kernel Graph") or (s["name"] in args.src):
+            print_step(s, print_graph=True, reconstruct_matches=s["name"] in args.src)
+          if DEBUG >= 7: print_step(s, reconstruct_matches=True)
+      elif DEBUG >= 3 and k.get("ext"): print(emit(k["ext"]))
+    for k in (produce_top_kernels if args.t else produce_all_kernels)(): render_event(k)
+
+def get_arg_parser() -> argparse.ArgumentParser:
+  parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli")
+  parser.add_argument("-s", "--src", nargs="+", default=[], metavar="NAME", help="Select a data source (default: all)")
+  parser.add_argument("--list", "--ls", dest="list", action="store_true", help="List sources")
+  parser.add_argument("--interval", nargs=2, metavar=("START", "END"), help="Optional start and end marker")
+  parser.add_argument("-t", nargs="?", type=int, const=20, metavar="COUNT", help="Aggregate top kernels (optional count, default 20)")
+  parser.add_argument("--profile-path", type=str, metavar="PATH", help="Optional path to profile.pkl (default: latest profile)",
+                      default=temp("profile.pkl", append_user=True))
+  parser.add_argument("--rewrites-path", type=str, metavar="PATH", help="Optional path to rewrites.pkl (default: latest rewrites)",
+                      default=temp("rewrites.pkl", append_user=True))
+  parser.add_argument("--json", action="store_true", help="Emit profiler output as JSON")
+  return parser
+
+if __name__ == "__main__":
+  try: main(get_arg_parser().parse_args())
+  except KeyboardInterrupt: pass
--- a/tinygrad_repo/tinygrad/viz/fetch_assets.sh
+++ b/tinygrad_repo/tinygrad/viz/fetch_assets.sh
@@ -0,0 +1,14 @@
+#!/bin/bash
+fetch() {
+  echo "fetch $1"
+  mkdir -p assets/$1
+  rmdir assets/$1
+  curl -o assets/$1 https://$1
+}
+fetch "d3js.org/d3.v7.min.js"
+fetch "dagrejs.github.io/project/dagre/latest/dagre.min.js"
+fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/styles/default.min.css"
+fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js"
+fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js"
+fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js"
+fetch "unpkg.com/@highlightjs/cdn-assets@11.10.0/styles/tokyo-night-dark.min.css"
--- a/tinygrad_repo/tinygrad/viz/index.html
+++ b/tinygrad_repo/tinygrad/viz/index.html
@@ -0,0 +1,437 @@
+<!DOCTYPE html>
+<html>
+<head>
+  <title>tinygrad viz</title>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <link rel="icon" href="data:;base64,iVBORw0KGgo=">
+  <script src="assets/d3js.org/d3.v7.min.js" charset="utf-8"></script>
+  <script src="assets/dagrejs.github.io/project/dagre/latest/dagre.min.js"></script>
+  <link rel="stylesheet" href="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/styles/default.min.css">
+  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js"></script>
+  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js"></script>
+  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js"></script>
+  <link rel="stylesheet" href="assets/unpkg.com/@highlightjs/cdn-assets@11.10.0/styles/tokyo-night-dark.min.css" />
+  <style>
+  * {
+    box-sizing: border-box;
+    margin: 0;
+    padding: 0;
+  }
+  html, body {
+    color: #f0f0f5;
+    margin: 0;
+    padding: 0;
+    width: 100%;
+    height: 100%;
+    font-family: sans-serif;
+    font-optical-sizing: auto;
+    font-weight: 400;
+    font-style: normal;
+    font-variation-settings: "wdth" 100;
+    font-size: 14px;
+    overflow: hidden;
+    background-color: #08090e;
+    scrollbar-color: #686977 #151720;
+  }
+  ::-webkit-scrollbar-track { background: #151720; }
+  ::-webkit-scrollbar-thumb { background: #686977; }
+  a {
+    color: #4a90e2;
+    text-decoration: underline;
+    cursor: pointer;
+    display: block;
+  }
+  ul {
+    padding: 0;
+    white-space: nowrap;
+    cursor: pointer;
+  }
+  ul > p {
+    opacity: 0.6;
+  }
+  ul.active > p {
+    opacity: 1;
+  }
+  ul > ul {
+    display: none;
+    margin-left: 6px;
+  }
+  ul.has-children ul {
+    margin-left: calc(6px + 1ch);
+  }
+  ul.has-children > p::before {
+    content:"▸ ";
+  }
+  ul.has-children.expanded > p::before {
+    content:"▾ ";
+  }
+  ul.expanded > ul {
+    display: block;
+  }
+  ul.disabled > p {
+    opacity: 0.4;
+  }
+  ul.disabled {
+    pointer-events: none;
+  }
+  label {
+    display: flex;
+    align-items: center;
+    gap: 4px;
+    line-height: 1;
+    user-select: none;
+    cursor: pointer;
+  }
+  input {
+    outline: none;
+  }
+  #graph svg {
+    width: 100%;
+    height: 100%;
+  }
+  svg * {
+    cursor: default;
+    user-select: none;
+  }
+  g.clickable * {
+    cursor: pointer;
+    user-select: auto;
+  }
+  g.tag circle {
+    fill: #FFD700;
+    stroke: #B8860B;
+  }
+  g.tag.collapsed circle, g.tag.collapsed rect {
+    fill: #5CD68D;
+    stroke: #4a4b57;
+  }
+  g.tag.expanded circle, g.tag.expanded rect {
+    fill: #9FDDE6;
+    stroke: #4a4b57;
+  }
+  g.port circle {
+    fill: #b3dcc2;
+  }
+  g.tag circle, g.tag rect, #edge-labels circle {
+    stroke-width: 0.8;
+  }
+  g.tag text, #edge-labels text {
+    font-family: monospace;
+    text-anchor: middle;
+    font-size: 6px;
+    fill: #08090e;
+  }
+  g.tag.ref * {
+    cursor: pointer;
+  }
+  g.tag.ref circle {
+    fill: #9FDDE6;
+    stroke: #4a4b57;
+  }
+  g.tag.ref path {
+    fill: none;
+    stroke: #08090e;
+    stroke-width: 0.7;
+    stroke-linejoin: miter;
+  }
+  .label :is(text, p) {
+    font-weight: 350;
+  }
+  g.node rect {
+    stroke-width: 1.4;
+    stroke: #4a4b57;
+  }
+  g.overlay rect {
+    fill: rgba(26, 27, 38, 0.5);
+  }
+  .edgePath {
+    fill: none;
+    stroke-width: 1.4px;
+  }
+  g.label rect.bg {
+    fill: transparent;
+    stroke: none;
+  }
+  g.label rect.bg.highlight {
+    fill: #5f0059;
+  }
+  #insts .line {
+    display: flex;
+    flex-direction: column;
+    margin-bottom: 8px;
+  }
+  #insts .left {
+    width: fit-content;
+    display: flex;
+    gap: 4px;
+  }
+  #insts .left.highlight {
+    background-color: rgba(0, 199, 47, 0.2);
+  }
+  #insts .pc {
+    color: #73daca;
+  }
+  g.node.highlight rect.node, .edgePath.highlight, g.port circle {
+    stroke: #89C9A2;
+  }
+  g.highlight.child rect.node, .edgePath.highlight.child {
+    stroke: #C888B0;
+  }
+  #edge-labels g.port.highlight {
+    display: block
+  }
+  #edge-labels g.port {
+    display: none
+  }
+  #arrowhead {
+    fill: #4a4b57;
+  }
+  .main-container {
+    display: flex;
+    width: 100%;
+    height: 100%;
+    position: relative;
+  }
+  .container {
+    flex: 0 0 auto;
+    background-color: #0f1018;
+    padding: 20px;
+    z-index: 2;
+    position: relative;
+    height: 100%;
+  }
+  .metadata > * + *, .info > * + *, .rewrite-container > * + *, .ctx-list > * + * {
+    margin-top: 12px;
+  }
+  ul > * + *, .args > * + * {
+    margin-top: 4px;
+  }
+  #graph {
+    position: absolute;
+    inset: 0;
+    z-index: 1;
+  }
+  #profiler, #custom {
+    flex: 1 1 auto;
+    min-width: 0;
+    width: 100%;
+    height: calc(100% - 50px);
+    margin-top: 50px;
+    overflow-y: auto;
+    overflow-x: hidden;
+    scrollbar-gutter: stable;
+  }
+  .ctx-list-parent {
+    width: 15%;
+    padding-top: 50px;
+    border-right: 1px solid #4a4b56;
+  }
+  .ctx-list, .metadata {
+    width: 100%;
+    height: 100%;
+    overflow-y: auto;
+    scrollbar-gutter: stable;
+  }
+  .metadata-parent {
+    width: 20%;
+    border-left: 1px solid #4a4b56;
+    margin-left: auto;
+  }
+  .resize-handle {
+    position: absolute;
+    top: 0;
+    bottom: 0;
+    width: 20px;
+    height: 100%;
+    cursor: col-resize;
+    z-index: 3;
+    background-color: transparent;
+  }
+  .floating-container {
+    position: fixed;
+    top: 10px;
+    left: 20px;
+    z-index: 4;
+    display: flex;
+    flex-direction: row;
+    gap: 8px;
+  }
+  .btn {
+    outline: none;
+    background-color: #1a1b26;
+    border: 1px solid #4a4b56;
+    color: #f0f0f5;
+    border-radius: 4px;
+    padding: 6px;
+    cursor: pointer;
+    height: 32px;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    text-decoration: none;
+  }
+  .btn:hover {
+    background-color: #2a2b36;
+  }
+  .collapsed .container {
+    display: none;
+  }
+  .rewrite-list {
+    display: flex;
+    flex-wrap: wrap;
+  }
+  .rewrite-list > ul {
+    padding: 2px;
+  }
+  .wrap {
+    word-wrap: break-word;
+    white-space: pre-wrap;
+  }
+  pre code.hljs {
+    overflow-y: auto;
+    max-height: 30vh;
+    padding: 8px;
+  }
+  pre.full-height code.hljs {
+    max-height: none;
+  }
+  #progress-message {
+    position: absolute;
+    z-index: 2;
+    left: 50%;
+    top: 2%;
+    color: #ffd230;
+    display: none;
+  }
+  #tooltip {
+    position: absolute;
+    z-index: 4;
+    background-color: #1e2029;
+    padding: 4px 8px;
+    max-width: 164px;
+    border-radius: 4px;
+    pointer-events: none;
+    display: none;
+    font-size: 10px;
+  }
+  #device-list > div {
+    overflow-x: auto;
+    overflow-y: hidden;
+    white-space: nowrap;
+    display: flex;
+    min-height: 32px;
+  }
+  #device-list > div:hover {
+    background-color: rgba(20, 23, 35, 0.3);
+  }
+  #device-list {
+    height: fit-content;
+  }
+  .raw-text {
+    padding-left: 15px;
+    width: 100%;
+    height: 100%;
+    max-height: 100vh;
+    overflow-x: auto;
+  }
+  .raw-text > pre {
+    display: inline-block;
+    min-width: 100%;
+  }
+  .raw-text code {
+    max-height: none !important;
+  }
+  table {
+    width: 100%;
+    border-collapse: separate;
+    border-spacing: 0;
+    background-color: #1a1b26;
+    color: #f0f0f5;
+    font-size: 0.95em;
+  }
+  table td {
+    border-bottom: 1px solid #4a4b56;
+    vertical-align: top;
+  }
+  table tr:last-child > td {
+     border-bottom: none;
+  }
+  tr.main-row:hover, tr.main-row.expanded, tr.nested-row > td > table, tr.nested-row thead {
+    background-color: #2a2d3a;
+  }
+  tr.main-row > td {
+    padding: 8px 12px;
+  }
+  td.Instruction {
+    font-family: monospace;
+  }
+  thead {
+    position: sticky;
+    top: 0;
+    z-index: 10;
+    background-color: #20222e;
+  }
+  thead th {
+    text-align: left;
+    padding: 10px 12px;
+    font-weight: 600;
+    border-bottom: 1px solid #4a4b56;
+    font-size: 0.95em;
+    letter-spacing: 0.03em;
+  }
+  tr.nested-row > td {
+    border-bottom: none;
+  }
+  tr.nested-row table tr.main-row:hover {
+    background-color: unset;
+  }
+  tr.main-row.has-children > td:first-child > p {
+    display: inline-block;
+  }
+  tr.main-row.has-children > td:first-child::before {
+    content: "▸ ";
+    width: 1em;
+    margin-left: -0.25em;
+  }
+  tr.main-row.has-children.expanded > td:first-child::before {
+    content: "▾ ";
+  }
+  </style>
+</head>
+<body>
+  <div class="main-container">
+    <div class="floating-container">
+      <button class="btn collapse-btn">
+        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" width="20"><path d="M15 19l-7-7 7-7"/></svg>
+      </button>
+      <button class="btn" id="zoom-to-fit-btn" aria-label="Fit graph">
+        <svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" stroke-width="1.5" stroke="currentColor" width="20">
+          <path stroke-linecap="round" stroke-linejoin="round" d="M7.5 3.75H6A2.25 2.25 0 0 0 3.75 6v1.5M16.5 3.75H18A2.25 2.25 0 0 1 20.25 6v1.5m0 9V18A2.25 2.25 0 0 1 18 20.25h-1.5m-9 0H6A2.25 2.25 0 0 1 3.75 18v-1.5M15 12a3 3 0 1 1-6 0 3 3 0 0 1 6 0Z" />
+        </svg>
+      </button>
+    </div>
+    <div id="progress-message"></div>
+    <div class="container ctx-list-parent"><div class="ctx-list"></div></div>
+    <div class="view" id="profiler"></div>
+    <div class="view" id="graph">
+      <svg id="graph-svg" preserveAspectRatio="xMidYMid meet">
+        <g id="render">
+          <g id="edges"></g>
+          <g id="nodes"></g>
+          <g id="edge-labels"></g> <!-- NOTE: this ensures edge labels are always on top -->
+        </g>
+        <defs>
+          <marker id="arrowhead" viewBox="0 -5 10 10" refX="10" refY="0" markerWidth="6" markerHeight="6" orient="auto">
+            <path d="M0,-5L10,0L0,5" fill="context-stroke"></path>
+          </marker>
+        </defs>
+      </svg>
+    </div>
+    <div class="view" id="custom"></div>
+    <div class="container metadata-parent"><div class="metadata"></div></div>
+  </div>
+  <div id="tooltip" class="wrap"></div>
+  <script src="/js/index.js"></script>
+</body>
+</html>
--- a/tinygrad_repo/tinygrad/viz/js/index.js
+++ b/tinygrad_repo/tinygrad/viz/js/index.js
--- a/tinygrad_repo/tinygrad/viz/js/worker.js
+++ b/tinygrad_repo/tinygrad/viz/js/worker.js
@@ -0,0 +1,119 @@
+const NODE_PADDING = 10;
+const rectDims = (lw, lh) => ({ width:lw+NODE_PADDING*2, height:lh+NODE_PADDING*2, labelWidth:lw, labelHeight:lh });
+
+const canvas = new OffscreenCanvas(0, 0);
+const ctx = canvas.getContext("2d");
+
+onmessage = (e) => {
+  try {
+    const { data, opts } = e.data;
+    const g = new dagre.graphlib.Graph({ compound: true }).setDefaultEdgeLabel(function() { return {}; });
+    (data.blocks != null ? layoutCfg : layoutUOp)(g, data, opts);
+    postMessage({result: dagre.graphlib.json.write(g)});
+    self.close();
+  } catch (err) {
+    postMessage({error: err.stack || err.message || String(err)});
+    self.close();
+  }
+}
+
+const layoutCfg = (g, { blocks, paths, pc_tokens }) => {
+  const lineHeight = 18;
+  g.setGraph({ rankdir:"TD", font:"monospace", lh:lineHeight, textSpace:"1ch" });
+  ctx.font = `350 ${lineHeight}px ${g.graph().font}`;
+  // basic blocks render the assembly in nodes
+  const tokenColors = {0:"#7aa2f7", 1:"#9aa5ce"};
+  for (const [lead, members] of Object.entries(blocks)) {
+    let [width, height, label] = [0, 0, []];
+    for (const m of members) {
+      const tokens = pc_tokens[m];
+      label.push(tokens.map((t, i) => ({st:t.st, keys:t.keys, color:tokenColors[t.kind]})));
+      width = Math.max(width, ctx.measureText(tokens.map((t) => t.st).join("")).width);
+      height += lineHeight;
+    }
+    g.setNode(lead, { ...rectDims(width, height), label, labelX:0, id:lead, color:"#1a1b26", addrspace:null });
+  }
+  // paths become edges between basic blocks
+  const pathColors = {0:"#3f7564", 1:"#7a4540", 2:"#3b5f7e"};
+  for (const [lead, value] of Object.entries(paths)) {
+    for (const [id, color] of Object.entries(value)) g.setEdge(lead, id, {label:{type:"port", text:""}, color:pathColors[color]});
+  }
+  dagre.layout(g);
+}
+
+const layoutUOp = (g, { graph, change }, opts) => {
+  const lineHeight = 14;
+  g.setGraph({ rankdir: "LR", font:"sans-serif", lh:lineHeight });
+  ctx.font = `350 ${lineHeight}px ${g.graph().font}`;
+  if (change?.length) g.setNode("overlay", {label:"", labelWidth:0, labelHeight:0, labelX:0, className:"overlay"});
+  let callCount = 0;
+  for (const [k, {label, src, ref, color, tag, exclude, addrspace}] of Object.entries(graph)) {
+    // adjust node dims by label size (excluding escape codes) + add padding
+    let [width, height] = [0, 0];
+    for (line of label.replace(/\u001B\[(?:K|.*?m)/g, "").split("\n")) {
+      width = Math.max(width, ctx.measureText(line).width);
+      height += lineHeight;
+    }
+    const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n");
+    if (callNode) callCount++;
+    g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude, addrspace});
+    // add edges
+    const edgeCounts = {};
+    for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
+    for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
+      ...(callNode && port === 0 && {color:"#a0a1b8"})});
+    if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
+  }
+  // optionally hide nodes from the layout
+  if (!opts.showSink) {
+    for (const n of g.nodes()) {
+      const node = g.node(n);
+      if ((node.label === "SINK" || node.label.startsWith("SINK\n")) && (g.successors(n) || []).length === 0) g.removeNode(n);
+    }
+  }
+  if (!opts.showIndexing) {
+    for (const n of g.nodes()) {
+      const node = g.node(n);
+      if (node.label.includes("dtypes.weakint")) g.removeNode(n);
+    }
+  }
+  // optionally remove node srcs, track affected nodes
+  const disconnected = new Set();
+  const CALL_TAG_WIDTH = 14;
+  for (const n of g.nodes()) {
+    const node = g.node(n);
+    for (const consumerId of (g.successors(n) || [])) {
+      const consumer = g.node(consumerId);
+      // add +- toggle if this consumer has collapsible sources
+      const edge = g.edge(n, consumerId);
+      const collapsible = consumer.callNode ? edge?.label?.text === 0 : node.exclude;
+      if (!collapsible) continue;
+      consumer.collapsible = true;
+      // increase width of call/function nodes to make space for a toggle
+      if (consumer.callNode) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
+      // make sources invisible if UI has toggled it off
+      const collapsed = consumer.callNode ? opts.showCallSrc === opts.callSrcMask.has(consumerId) : !opts.expandedNodes.has(consumerId);
+      if (!collapsed) continue;
+      consumer.collapsed = true;
+      g.removeEdge(n, consumerId);
+      disconnected.add(n);
+    }
+  }
+  // remove nodes that are now disconnected (no successors), only from affected subtree
+  let changed = true;
+  while (changed) {
+    changed = false;
+    for (const n of disconnected) {
+      if (!g.hasNode(n)) continue;
+      if ((g.successors(n) || []).length === 0) {
+        for (const pred of (g.predecessors(n) || [])) disconnected.add(pred);
+        g.removeNode(n);
+        changed = true;
+      }
+    }
+  }
+  g.graph().callCount = callCount;
+  dagre.layout(g);
+  // remove overlay node if it's empty
+  if (!g.node("overlay")?.width) g.removeNode("overlay");
+}
--- a/tinygrad_repo/tinygrad/viz/serve.py
+++ b/tinygrad_repo/tinygrad/viz/serve.py
@@ -0,0 +1,723 @@
+#!/usr/bin/env python3
+import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, codecs, io, struct, re, traceback, itertools
+import socketserver
+from contextlib import redirect_stdout, redirect_stderr, contextmanager
+from decimal import Decimal
+from dataclasses import dataclass, field
+from urllib.parse import parse_qs, urlparse
+from http.server import BaseHTTPRequestHandler
+from typing import Any, TypedDict, TypeVar, Generator, Callable
+from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
+from tinygrad.helpers import printable, Context, START_TIME, NO_COLOR, ansistrip
+from tinygrad.renderer.amd.dsl import Inst
+from tinygrad.renderer.amd import detect_format
+
+# NOTE: using HTTPServer forces a potentially slow socket.getfqdn
+class TCPServerWithReuse(socketserver.TCPServer):
+  allow_reuse_address = True
+  def __init__(self, server_address, RequestHandlerClass):
+    print(f"*** started server on http://127.0.0.1:{server_address[1]} at {time.perf_counter()-START_TIME:.2f} s")
+    super().__init__(server_address, RequestHandlerClass)
+
+class HTTPRequestHandler(BaseHTTPRequestHandler):
+  def send_data(self, data:bytes, content_type:str="application/json", status_code:int=200):
+    self.send_response(status_code)
+    self.send_header("Content-Type", content_type)
+    self.send_header("Content-Length", str(len(data)))
+    self.end_headers()
+    return self.wfile.write(data)
+  def stream_json(self, source:Generator):
+    try:
+      self.send_response(200)
+      self.send_header("Content-Type", "text/event-stream")
+      self.send_header("Cache-Control", "no-cache")
+      self.end_headers()
+      for r in source:
+        self.wfile.write(f"data: {json.dumps(r)}\n\n".encode("utf-8"))
+        self.wfile.flush()
+      self.wfile.write("data: [DONE]\n\n".encode("utf-8"))
+    # pass if client closed connection
+    except (BrokenPipeError, ConnectionResetError): return
+
+from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, range_start, multirange_str
+from tinygrad.uop.ops import KernelInfo
+from tinygrad.uop.render import print_uops, pyrender
+from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfileProgramEvent
+from tinygrad.dtype import dtypes, AddrSpace
+
+uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
+               **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.SHAPED_WMMA: "#FF5B5B",
+               Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
+               Ops.INDEX: "#D8F9E4", Ops.STACK: "#D8F9E4",
+               Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
+               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
+               Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
+               Ops.LINEAR: "#7DF4FF",
+               Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
+               Ops.STAGE: "#AC640D", Ops.REWRITE_ERROR: "#ff2e2e", Ops.AFTER: "#8A7866", Ops.END: "#524C46"}
+
+addrspace_colors = {AddrSpace.REG:"#e68181", AddrSpace.LOCAL:"#e7c86a", AddrSpace.GLOBAL:"#75bd7b"}
+
+# VIZ API
+
+# A step is a lightweight descriptor for a trace entry
+# Includes a name, metadata and a URL path for fetching the full data
+
+def create_step(name:str, query:tuple[str, int, int], data=None, depth:int=0, **kwargs) -> dict:
+  return {"name":name, "query":f"{query[0]}?ctx={query[1]}&step={query[2]}", "data":data, "depth":depth, **kwargs}
+
+@dataclass(frozen=True)
+class VizData:
+  trace:RewriteTrace = field(default_factory=lambda: RewriteTrace([], [], {}))
+  ctxs:list[dict] = field(default_factory=list)
+  ref_map:dict[Any, int] = field(default_factory=dict)
+  all_uops:dict[int, UOp] = field(default_factory=dict)
+
+# ** load all saved rewrites
+
+def load_rewrites(data:VizData) -> None:
+  assert not data.ctxs and not data.ref_map, "load_rewrites called multiple times"
+  for i,k in enumerate(data.trace.keys):
+    steps:list[dict] = []
+    p:UOp|None = None
+    for j,s in enumerate(data.trace.rewrites[i]):
+      steps.append(create_step(s.name, ("/graph-rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc),
+                               trace=k.tb if j==0 else None, depth=s.depth))
+      # get source and binary from Ops.PROGRAM
+      if s.name == "View Program":
+        p = _reconstruct(data, s.sink, depth=1)
+        steps.append(create_step("View UOp List", ("/uops", i, len(steps))))
+        steps.append(create_step("View Source", ("/code", i, len(steps)), p.src[3].arg))
+        steps.append(create_step("View Disassembly", ("/asm", i, len(steps)), (k.ret, p.src[4].arg)))
+    for key in k.keys: data.ref_map[canonicalize_ast(key) if isinstance(key, UOp) else key] = i
+    data.ctxs.append({"name":k.display_name, "steps":steps, "prg":p})
+
+# ** get the complete UOp graphs for one rewrite
+
+class GraphRewriteDetails(TypedDict):
+  graph: dict                            # JSON serialized UOp for this rewrite step
+  uop: str                               # strigified UOp for this rewrite step
+  diff: list[str]|None                   # diff of the single UOp that changed
+  change: list[int]|None                 # the new UOp id + all its parents ids
+  upat: tuple[tuple[str, int], str]|None # [loc, source_code] of the matched UPat
+
+def shape_to_str(s:tuple[sint, ...]): return "(" + ','.join(srender(x) for x in s) + ")"
+def mask_to_str(s:tuple[tuple[sint, sint], ...]): return "(" + ','.join(shape_to_str(x) for x in s) + ")"
+def pystr(u:UOp) -> str:
+   # pyrender may check for shape mismatch
+  try: return pyrender(u)
+  except Exception: return str(u)
+
+def fmt_colored(s:str) -> str: return ansistrip(s) if NO_COLOR else s
+
+def canonicalize_ast(u:UOp) -> UOp: return u.replace(arg=KernelInfo()) if u.op is Ops.SINK and isinstance(u.arg, KernelInfo) else u
+
+def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
+  assert isinstance(x, UOp)
+  graph: dict[int, dict] = {}
+  excluded: set[UOp] = set()
+  for u in (toposort:=x.toposort()):
+    # always exclude DEVICE/CONST/UNIQUE
+    if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE, Ops.LUNIQUE} and u is not x: excluded.add(u)
+    if u.op is Ops.CONST and len(u.src) and u.src[0].op in {Ops.UNIQUE, Ops.LUNIQUE}: excluded.remove(u)
+    if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u)
+    # exclude RESHAPE/EXPAND that only serve to broadcast a CONST
+    if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
+    if u.op in GroupOp.Movement: excluded.update(s for s in u.src if s.op is Ops.STACK and all(x.op is Ops.CONST for x in s.src))
+  for u in toposort:
+    argst = codecs.decode(str(u.arg), "unicode_escape")
+    if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg)
+    if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>"
+    if u.op is Ops.CONST and dtypes.is_float(u.dtype): argst = f"{u.arg:g}"
+    wrap_len = 200 if u.op is Ops.SOURCE else 80
+    label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''), wrap=wrap_len)) if u.arg is not None else ''}"
+    if u.dtype != dtypes.void: label += f"\n{u.dtype}"
+    for idx,x in enumerate(u.src[:1] if u.op in {Ops.STAGE, Ops.INDEX} else (u.src if u.op is not Ops.END else [])):
+      if x in excluded:
+        # walk through excluded movement ops to find the underlying CONST
+        cx = x
+        while cx.op in GroupOp.Movement and len(cx.src) >= 1 and cx.src[0] in excluded: cx = cx.src[0]
+        arg = f"{cx.arg:g}" if cx.op is Ops.CONST and dtypes.is_float(cx.dtype) else f"{cx.arg}"
+        label += f"\n{cx.op.name}{idx} {arg}" + (f" {cx.src[0].op}" if len(cx.src) else "")
+    try:
+      if len(rngs:=u.ranges):
+        label += f"\n({multirange_str(rngs, color=True)})"
+      if u._shape is not None:
+        label += f"\n{shape_to_str(u.shape)}"
+      if u.op in {Ops.CALL, Ops.FUNCTION}:
+        label += f"\n{u.src[0].key.hex()[:8]}"
+      if u.op in {Ops.INDEX, Ops.STAGE}:
+        if len(u.toposort()) < 30: label += f"\n{u.render()}"
+        ranges: list[UOp] = []
+        for us in u.src[1:]: ranges += [s for s in us.toposort() if s.op in {Ops.RANGE, Ops.SPECIAL}]
+        if ranges: label += "\n"+' '.join([f"{s.render()}={s.vmax+1}" for s in ranges])
+      if u.op in {Ops.END, Ops.REDUCE} and len(trngs:=list(UOp.sink(*u.src[range_start[u.op]:]).ranges)):
+        label += "\n"+' '.join([f"{range_str(s, color=True)}({s.vmax+1})" for s in trngs])
+    except Exception:
+      label += "\n<ISSUE GETTING LABEL>"
+    ref = data.ref_map.get(canonicalize_ast(u.src[0])) if u.op in {Ops.CALL, Ops.FUNCTION} else None
+    if ref is not None: label += f"\ncodegen@{fmt_colored(data.ctxs[ref]['name'])}"
+    # NOTE: kernel already has metadata in arg
+    if TRACEMETA >= 2 and u.metadata is not None and u.op not in {Ops.CALL, Ops.FUNCTION}: label += "\n"+str(u.metadata)
+    # limit SOURCE labels line count
+    if u.op is Ops.SOURCE and len(lines:=label.split("\n")) > 40:
+      label = "\n".join(lines[:30]) + "\n..."
+    graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src)], "exclude":u in excluded, "color":uops_colors.get(u.op, "#ffffff"),
+                    "ref":ref, "tag":repr(u.tag) if u.tag is not None else None,
+                    "addrspace":addrspace_colors.get(u.addrspace, None) if u.addrspace is not None else None}
+  return graph
+
+def _reconstruct(data:VizData, a:int, depth:int|None=None):
+  if depth is None and a in data.all_uops: return data.all_uops[a]
+  op, dtype, src, arg, *rest = data.trace.uop_fields[a]
+  if depth is not None and depth <= 0: return UOp(op, dtype, (), arg, *rest)
+  ret = UOp(op, dtype, tuple(_reconstruct(data, s, None if depth is None else depth-1) for s in src), arg, *rest)
+  if depth is None: data.all_uops[a] = ret
+  return ret
+
+def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
+  next_sink = _reconstruct(data, ctx.sink)
+  yield {"graph":uop_to_json(data, next_sink), "uop":pystr(next_sink), "change":None, "diff":None, "upat":None}
+  replaces: dict[UOp, UOp] = {}
+  for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches, disable=not ctx.matches):
+    replaces[u0:=_reconstruct(data, u0_num)] = u1 = _reconstruct(data, u1_num)
+    try: new_sink = next_sink.substitute(replaces)
+    except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
+    match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
+    yield {"graph":(sink_json:=uop_to_json(data, new_sink)), "uop":pystr(new_sink), "change":[id(x) for x in u1.toposort() if id(x) in sink_json],
+           "diff":list(difflib.unified_diff(pystr(u0).splitlines(), pystr(u1).splitlines())), "upat":(upat_loc, match_repr)}
+    if not ctx.bottom_up: next_sink = new_sink
+
+# encoder helpers
+
+def enum_str(s, cache:dict[str, int]) -> int:
+  if (cret:=cache.get(s)) is not None: return cret
+  cache[s] = ret = len(cache)
+  return ret
+
+def option(s:int|None) -> int: return 0 if s is None else s+1
+
+def rel_ts(ts:int|Decimal, start_ts:int, ctx:str="") -> int:
+  val = int(ts) - start_ts
+  if val < 0 or val > 0xFFFFFFFF: raise ValueError(f"timestamp out of range: {ctx} diff={val} (ts={ts} start={start_ts})")
+  return val
+
+# Profiler API
+
+def cpu_ts_diff(device_ts_diffs:dict[str, Decimal], device:str) -> Decimal: return device_ts_diffs.get(device, Decimal(0))
+
+DevEvent = ProfileRangeEvent|ProfileGraphEntry|ProfilePointEvent
+def flatten_events(profile:list[ProfileEvent], device_ts_diffs:dict[str, Decimal]) -> Generator[tuple[Decimal, Decimal, DevEvent], None, None]:
+  for e in profile:
+    if isinstance(e, ProfileRangeEvent): yield (e.st+(diff:=cpu_ts_diff(device_ts_diffs, e.device)), (e.en if e.en is not None else e.st)+diff, e)
+    elif isinstance(e, ProfilePointEvent): yield (e.ts, e.ts, e)
+    elif isinstance(e, ProfileGraphEvent):
+      cpu_ts = []
+      for ent in e.ents: cpu_ts += [e.sigs[ent.st_id]+(diff:=cpu_ts_diff(device_ts_diffs, ent.device)), e.sigs[ent.en_id]+diff]
+      yield (st:=min(cpu_ts)), (et:=max(cpu_ts)), ProfileRangeEvent(f"{e.ents[0].device.split(':')[0]} Graph", f"batched {len(e.ents)}", st, et)
+      for i,ent in enumerate(e.ents): yield (cpu_ts[i*2], cpu_ts[i*2+1], ent)
+
+# normalize event timestamps and attach kernel metadata
+def timeline_layout(data:VizData, dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, scache:dict[str, int]) -> bytes|None:
+  events:list[bytes] = []
+  ei:ProfilePointEvent|None = None
+  for st,et,dur,e in dev_events:
+    if isinstance(e, ProfilePointEvent) and e.name == "exec": ei = e
+    if dur == 0: continue
+    name, key = e.name, None
+    fmt:dict = {}
+    if (ref:=data.ref_map.get(name)) is not None and ref < len(data.ctxs):
+      name = data.ctxs[ref]["name"]
+      if (p:=data.ctxs[ref].get("prg")) is not None and (ki:=p.src[0].arg).estimates is not None and ei is not None:
+        fmt["FLOPS"] = int(sym_infer(ki.estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6))
+        fmt["B/s mem"], fmt["B/s lds"] = int(sym_infer(ki.estimates.mem, var_vals)/t), int(sym_infer(ki.estimates.lds, var_vals)/t)
+        if ei.arg["metadata"]: fmt["metadata"] = ",".join([str(m) for m in ei.arg['metadata']+["batched" if isinstance(e,ProfileGraphEntry) else ""]])
+        key = ei.key
+    elif isinstance(e.name, TracingKey):
+      name = e.name.display_name
+      ref = next((v for k in e.name.keys if (v:=data.ref_map.get(k)) is not None), None)
+      if isinstance(e.name.ret, str): fmt.update(json.loads(e.name.ret[4:]) if e.name.ret.startswith("JSON") else {"metadata":e.name.ret})
+      elif isinstance(e.name.ret, int): fmt["B/s"], fmt["B"] = int(e.name.ret/(dur*1e-6)), e.name.ret
+      elif e.name.tb: fmt["tb"] = e.name.tb
+    events.append(struct.pack("<IIIIfI", enum_str(name, scache), option(ref), option(key), rel_ts(st,start_ts, f"'{name}' on {e.device}"),
+                              dur, enum_str(json.dumps(fmt),scache)))
+  return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None
+
+def encode_mem_free(key:int, ts:int, execs:list[ProfilePointEvent], scache:dict) -> bytes:
+  ei_encoding:list[tuple[int, int, int, int]] = [] # <[u32, u32, u32, u8] [run id, display name, buffer number and mode (2 = r/w, 1 = w, 0 = r)]
+  for e in execs:
+    num = next(i for i,k in enumerate(e.arg["bufs"]) if k == key)
+    mode = 2 if (num in e.arg["inputs"] and num in e.arg["outputs"]) else 1 if (num in e.arg["outputs"]) else 0
+    ei_encoding.append((e.key, enum_str(e.arg["name"], scache), num, mode))
+  return struct.pack("<BIII", 0, ts, key, len(ei_encoding))+b"".join(struct.pack("<IIIB", *t) for t in ei_encoding)
+
+def graph_layout(k:str, dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
+                 scache:dict[str, int]) -> tuple[str, bytes|None]:
+  if k.startswith("LINE:"):
+    xy = [(rel_ts(e.ts, start_ts, f"line '{k}' on {e.device}"), e.key) for st,_,_,e in dev_events if isinstance(e, ProfilePointEvent)]
+    peaks.append(peak:=max([y for _,y in xy]))
+    return k.replace("LINE:", ""), struct.pack("<BIBQ", 1, len(xy), 1, peak)+b"".join(struct.pack("<IQ", x, y) for x,y in xy)
+  peak, mem = 0, 0
+  temp:dict[int, int] = {}
+  events:list[bytes] = []
+  buf_ei:dict[int, list[ProfilePointEvent]] = {}
+  for st,_,_,e in dev_events:
+    if not isinstance(e, ProfilePointEvent): continue
+    if e.name == "alloc":
+      safe_sz = min(1_000_000_000_000, e.arg["sz"])
+      events.append(struct.pack("<BIIIQ", 1, rel_ts(e.ts, start_ts, f"alloc on {e.device}"), e.key, enum_str(e.arg["dtype"].name, scache), safe_sz))
+      dtype_size.setdefault(e.arg["dtype"].name, e.arg["dtype"].itemsize)
+      temp[e.key] = nbytes = safe_sz*e.arg["dtype"].itemsize
+      mem += nbytes
+      if mem > peak: peak = mem
+    if e.name == "exec" and e.arg["bufs"]:
+      for b in e.arg["bufs"]: buf_ei.setdefault(b, []).append(e)
+    if e.name == "free":
+      events.append(encode_mem_free(e.key, rel_ts(e.ts, start_ts, f"free on {e.device}"), buf_ei.pop(e.key, []), scache))
+      mem -= temp.pop(e.key)
+  for t in temp: events.append(encode_mem_free(t, rel_ts(end_ts, start_ts, f"end_ts for {k}"), buf_ei.pop(t, []), scache))
+  peaks.append(peak)
+  return f"{k} Memory", struct.pack("<BIBQ", 1, len(events), 0, peak)+b"".join(events) if events else None
+
+# by default, VIZ does not start when there is an error
+# use this to instead display the traceback to the user
+@contextmanager
+def soft_err(fn:Callable):
+  try: yield
+  except Exception: fn({"src":traceback.format_exc()})
+
+def row_tuple(row:str) -> tuple[tuple[int, int], ...]:
+  return ((0, 0),) if "Clock" in row else tuple((ord(ss[0][0]), int(ss[1])) if len(ss:=x.split(":"))>1 else (999,999) for x in row.split())
+
+# *** Performance counters
+
+metrics:dict[str, Callable[[dict[str, tuple[int, int, int]]], str]] = {
+  "VALU utilization": lambda s: f"{100 * (s['SQ_INSTS_VALU'][0] / s['SQ_INSTS_VALU'][2]) / (s['GRBM_GUI_ACTIVE'][1] * 4):.1f}%",
+  "SALU utilization": lambda s: f"{100 * (s['SQ_INSTS_SALU'][0] / s['SQ_INSTS_SALU'][2]) / (s['GRBM_GUI_ACTIVE'][1] * 4):.1f}%",
+}
+
+def unpack_pmc(e) -> dict:
+  agg_cols = ["Name", "Sum"]
+  rows:list[list] = []
+  stats:dict[str, tuple[int, int, int]] = {}  # name -> (sum, max, count)
+  view, ptr = memoryview(e.blob).cast('Q'), 0
+  for s in e.sched:
+    sample_cols = ["XCC", "INST", "SE", "SA"] + [f"WGP:{i}" for i in range(s.wgp)]
+    row:list = [s.name, 0, {"cols":sample_cols, "rows":[]}]
+    max_val, cnt = 0, 0
+    for sample in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa)):
+      vals:list[int] = []
+      # pack work group processors on the same se
+      for _ in range(s.wgp):
+        row[1] += (val:=int(view[ptr]))
+        max_val, cnt = max(max_val, val), cnt + 1
+        vals.append(val)
+        ptr += 1
+      row[2]["rows"].append(sample+tuple(vals))
+    stats[s.name] = (row[1], max_val, cnt)
+    rows.append(row)
+  for name, fn in metrics.items():
+    try: rows.append([name, fn(stats)])
+    except KeyError: pass
+  return {"rows":rows, "cols":agg_cols}
+
+# ** on startup, list all the performance counter traces
+
+def load_amd_counters(data:VizData, profile:list) -> None:
+  counter_events:dict[tuple[int, int], dict] = {}
+  durations:dict[str, list[float]] = {}
+  prg_events:dict[int, ProfileProgramEvent] = {}
+  arch = ""
+  for e in profile:
+    if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent"}:
+      counter_events.setdefault((e.kern, e.exec_tag), {}).setdefault(type(e).__name__, []).append(e)
+    if isinstance(e, ProfileRangeEvent) and e.device.startswith("AMD") and e.en is not None:
+      durations.setdefault(str(e.name), []).append(float(e.en-e.st))
+    if isinstance(e, ProfileProgramEvent) and e.tag is not None: prg_events[e.tag] = e
+    if isinstance(e, ProfileDeviceEvent) and e.device.startswith("AMD"): arch = f"gfx{unwrap(e.props)['gfx_target_version']//1000}"
+  if len(counter_events) == 0: return None
+  data.ctxs.append({"name":"All Counters", "steps":[create_step("PMC", ("/all-pmc", len(data.ctxs), 0), (durations, all_counters:={}))]})
+  run_number = {n:0 for n,_ in counter_events}
+  for (k, tag),v in counter_events.items():
+    # use the colored name if it exists
+    name = data.ctxs[r]["prg"].src[0].arg.name if (r:=data.ref_map.get(pname:=prg_events[k].name)) is not None else pname
+    run_number[k] += 1
+    steps:list[dict] = []
+    if (pmc:=v.get("ProfilePMCEvent")):
+      steps.append(create_step("PMC", ("/prg-pmc", len(data.ctxs), len(steps)), pmc[0]))
+      all_counters[(name, run_number[k], pname)] = pmc[0]
+    # to decode a SQTT trace, we need the raw stream, program binary and device properties
+    if (sqtt:=v.get("ProfileSQTTEvent")):
+      for e in sqtt:
+        if e.itrace: steps.append(create_step(f"SE:{e.se} PKTS", (f"/sqtt-{e.se}",len(data.ctxs),len(steps)), data=(e.blob,prg_events[k].lib,arch)))
+      try:
+        with Context(DEBUG=0): from extra.sqtt.roc import unpack_occ
+        steps.append(create_step("OCC", ("/amd-sqtt-occ", len(data.ctxs), len(steps)),
+                                 data={"fxn":unpack_occ, "args":((k, tag), sqtt, prg_events[k], arch)}))
+      except Exception: pass
+    data.ctxs.append({"name":f"SQTT {name}"+(f" n{run_number[k]}" if run_number[k] > 1 else ""), "steps":steps})
+
+wave_colors = {"WMMA": "#1F7857", **{x:"#ffffc0" for x in ["VALU", "VINTERP"]}, "SALU": "#cef263", "SMEM": "#ffc0c0", "STORE": "#4fa3cc",
+               **{x:"#b2b7c9" for x in ["VMEM", "SGMEM"]}, "LDS": "#9fb4a6", "IMMEDIATE": "#f3b44a", "BARRIER": "#d00000",
+               "JUMP_NO": "#fb8500", "JUMP": "#ffb703", "WAVERDY": "#1a2a2a"}
+
+def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent, None, None]:
+  from tinygrad.renderer.amd.sqtt import (map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC,
+                                          ALUEXEC, INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4, CDNA_INST, InstOpCDNA,
+                                          WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND, WAVERDY)
+  pc_map = {addr:str(inst) for addr,inst in amd_decode(lib, target).items()}
+  row_ends:dict[str, Decimal] = {}
+  row_counts:dict[str, itertools.count] = {}
+  curr_barrier:dict[int, ProfileRangeEvent] = {}
+  exec_pending:dict[str, list[tuple[str, str]]] = {}
+  dispatch_to_exec = {"WMMA":"VALU", "VALU":"VALU", "VALU1":"VALU", "VALUT":"VALU", "VALUB":"VALU", "VALUINST":"VALU", "VINTERP":"VALU",
+                      "SGMEM":"VMEM", "FLAT":"VMEM", "LDS":"LDS", "SALU":"SALU", "SMEM":"SALU", "VMEM":"VMEM"}
+  def add(name:str, p:PacketType, wave:int|None=None, info:InstructionInfo|None=None) -> Generator[ProfileEvent, None, None]:
+    row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name.replace('_ALT', '')}"
+    # by default we extend the packet to one cycle after timestamp
+    start_time, end_time = p._time, p._time+1
+    # exec links to dispatch, dispatch links to PC
+    link:dict|None = {"pc":info.pc} if info else None
+    if isinstance(p, (ALUEXEC, VMEMEXEC)):
+      dispatch_id, op_type = exec_pending[name].pop(0)
+      # wmma exec gets its own color and its own row on rdna4
+      if op_type.startswith("WMMA"):
+        name = name+"_WMMA"
+        if not op_type.startswith("WMMA_VALU"): row = "ALUEXEC:0 WMMA"
+      # transcendental valu gets its own row
+      if op_type.startswith("VALUT"): row = "ALUEXEC:0 TFU"
+      # extend execs by the op type's known duration, p._time marks the first or last cycle based on the op type
+      duration = int(dur_match.group(1)) if (dur_match:=re.match(r".*_(\d+)$", op_type)) else 1
+      if any(ss in row for ss in ("SALU", "TFU", "VMEM", "LDS")): start_time, end_time = p._time, p._time+duration
+      else: start_time, end_time = p._time-duration, p._time
+      link = {"link":dispatch_id}
+    # queue inst dispatches
+    idx = next(row_counts.setdefault(row, itertools.count(0)))
+    if isinstance(p, (VALUINST, INST, INST_RDNA4)) and (exec_type:=dispatch_to_exec.get(name.replace("OTHER_", "").split("_")[0])) is not None:
+      if name.startswith("OTHER_"): exec_type = f"{exec_type}_ALT"
+      # detect rdna3 wmma from the asm, only rdna4 has an op type for it
+      if isinstance(p, VALUINST) and (asm:=getattr(unwrap(info).inst, "op_name", "")).startswith("V_WMMA"):
+        name = f"WMMA_VALU_{16 if 'IU4' in asm else 32}"
+      exec_pending.setdefault(exec_type, []).append((f"{row}-{idx}", name))
+    # construct and yield the event for this packet
+    if row not in row_ends: yield ProfilePointEvent(row, "JSON", "pcMap", pc_map, ts=Decimal(0))
+    yield (e:=ProfileRangeEvent(row, TracingKey(name, ret="JSON"+json.dumps(link) if link else None), Decimal(start_time), Decimal(end_time)))
+    row_ends[row] = unwrap(e.en)
+    # barrier on this wave extends to fill the time it was waiting
+    if wave is not None:
+      if (barrier:=curr_barrier.pop(wave, None)) is not None: barrier.en = Decimal(p._time)
+      if name in {"BARRIER", "BARRIER_SIGNAL"}: curr_barrier[wave] = e
+  NS_PER_TICK = 10  # 100MHz
+  prev_pair:tuple[int, int]|None = None # (shader, realtime)
+  yield ProfilePointEvent("", "JSON", "waveColors", list(wave_colors.items()), ts=Decimal(0))
+  for p, info in map_insts(data, lib, target):
+    if isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) and p.is_marker:
+      pair = (p._time, p.delta)
+      if prev_pair is None: prev_pair = pair
+      else:
+        (s0, r0), (s1, r1) = prev_pair, pair
+        freq_hz = (s1 - s0) * 1_000_000_000 // ((r1 - r0) * NS_PER_TICK)
+        yield ProfilePointEvent("LINE:Shader Clock", "freq_hz", freq_hz, ts=Decimal(p._time))
+        prev_pair = pair
+    if isinstance(p, (INST, INST_RDNA4, CDNA_INST)):
+      name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4, InstOpCDNA)) else f"0x{p.op:02x}"
+      yield from add(name, p, info=info)
+    if isinstance(p, (VALUINST, IMMEDIATE, WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)): yield from add(p.__class__.__name__, p, info=info)
+    if isinstance(p, IMMEDIATE_MASK): yield from add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
+    if isinstance(p, WAVERDY):
+      for wave in range(16):
+        if p.mask & (1 << wave):
+          if wave in curr_barrier: yield from add("WAVERDY", p, wave=wave)
+    if isinstance(p, (VMEMEXEC, ALUEXEC)):
+      name = str(p.src).split('.')[1]
+      if name == "VALU_SALU":
+        yield from add("VALU", p)
+        yield from add("SALU", p)
+      else:
+        yield from add(name, p)
+
+def device_sort_fn(k:str) -> tuple:
+  special = {"GC": 0, "USER": 1, "TINY": 2, "ALLDEVS":100, "DISK": 999}
+  is_memory = k.endswith(" Memory")
+  p = k.split(" ")[0].split(":")
+  dev_base = p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"
+  return (is_memory, special.get(p[0], special['ALLDEVS']), dev_base, k)
+
+def get_profile(data:VizData, profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_sort_fn) -> bytes|None:
+  # start by getting the time diffs
+  device_ts_diffs:dict[str, Decimal] = {}
+  device_decoders:dict[str, Callable[[VizData, list[ProfileEvent]], None]] = {}
+  for ev in profile:
+    if isinstance(ev, ProfileDeviceEvent):
+      device_ts_diffs[ev.device] = ev.tdiff
+      if (d:=ev.device.split(":")[0]) == "AMD": device_decoders[d] = load_amd_counters
+      if d == "NV": device_decoders[d] = load_nv_counters
+  # load device specific counters
+  for fxn in device_decoders.values(): fxn(data, profile)
+  # map events per device
+  dev_events:dict[str, list[tuple[int, int, float, DevEvent]]] = {}
+  markers:list[ProfilePointEvent] = []
+  ext_data:dict[str, Any] = {}
+  start_ts:int|None = None
+  end_ts:int|None = None
+  for ts,en,e in flatten_events(profile, device_ts_diffs):
+    dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e))
+    if start_ts is None or st < start_ts: start_ts = st
+    if end_ts is None or et > end_ts: end_ts = et
+    if isinstance(e, ProfilePointEvent) and e.name == "marker": markers.append(e)
+    if isinstance(e, ProfilePointEvent) and e.name == "JSON": ext_data[e.key] = e.arg
+  if start_ts is None: return None
+  # return layout of per device events
+  layout:dict[str, bytes|None] = {}
+  scache:dict[str, int] = {}
+  peaks:list[int] = []
+  dtype_size:dict[str, int] = {}
+  for k,v in dev_events.items():
+    v.sort(key=lambda e:e[0])
+    layout[k] = timeline_layout(data, v, start_ts, scache)
+    layout.update([graph_layout(k, v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)])
+  sorted_layout = sorted([k for k,v in layout.items() if v is not None], key=sort_fn)
+  ret = [b"".join([struct.pack("<B", len(k)), k.encode(), unwrap(layout[k])]) for k in sorted_layout]
+  index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size,
+                      "markers":[{"ts":rel_ts(e.ts, start_ts, f"marker '{e.arg.get('name','?')}'"), **e.arg} for e in markers],
+                      **ext_data}).encode()
+  return struct.pack("<IQII", rel_ts(unwrap(end_ts), start_ts, "end_ts"), max(peaks,default=0), len(index), len(ret))+index+b"".join(ret)
+
+# ** PMA counters
+
+def load_nv_counters(data:VizData, profile:list) -> None:
+  steps:list[dict] = []
+  sm_version = {e.device:e.props.get("sm_version", 0x800) for e in profile if isinstance(e, ProfileDeviceEvent) and e.props is not None}
+  run_number:dict[str, int] = {}
+  for e in profile:
+    if type(e).__name__ == "ProfilePMAEvent":
+      run_number[e.kern] = run_num = run_number.get(e.kern, 0)+1
+      steps.append(create_step(f"PMA {e.kern}"+(f"n{run_num}" if run_num>1 else ""), ("/prg-pma-pkts", len(data.ctxs), len(steps)),
+                               data=(e.blob, sm_version[e.device])))
+  if steps: data.ctxs.append({"name":"All Counters", "steps":steps})
+
+def pma_timeline(blob:bytes, sm_version:int) -> list[ProfileEvent]:
+  from extra.nv_pma.decode import decode, decode_tpc_id
+  ret:list[ProfileEvent] = []
+  rows:dict[str, None] = {}
+  tpc_count:dict[int, int] = {}
+  # assume every sample is 32 cycles
+  cycles_per_sample = 32
+  for s, tpc_id in decode(blob, sm_version):
+    if len(ret) > getenv("MAX_SQTT_PKTS", 50_000): break
+    gpc, tpc, sm = decode_tpc_id(tpc_id)
+    tpc_count[tpc_id] = (n:=tpc_count.get(tpc_id,0)) + 1
+    rows.setdefault(row:=f"GPC:{gpc} TPC:{tpc} SM:{sm} WAVE:{s.wave_id}")
+    ret.append(ProfileRangeEvent(row, TracingKey(s.stall_reason.name, ret=f"pc=0x{s.pc_offset:06x} active={s.active}"),
+                                 Decimal(n*cycles_per_sample), Decimal((n+1)*cycles_per_sample)))
+  return [ProfilePointEvent(r, "start", r, ts=Decimal(0)) for r in rows]+ret
+
+# ** Assembly static analyzers
+
+def get_stdout(f: Callable) -> str:
+  buf = io.StringIO()
+  try:
+    with redirect_stdout(buf), redirect_stderr(buf): f()
+  except Exception: traceback.print_exc(file=buf)
+  return buf.getvalue()
+
+def get_elf_section(lib:bytes, name:str):
+  from tinygrad.runtime.support.elf import elf_loader
+  return next((sh for sh in elf_loader(lib)[1] if sh.name == name))
+
+def amd_decode(lib:bytes, target:str) -> dict[int, Inst]:
+  text = get_elf_section(lib, ".text")
+  off, buf = text.header.sh_addr, text.content
+  arch = "rdna3" if target.startswith("gfx11") else "rdna4" if target.startswith("gfx12") else "cdna"
+  addr_table:dict[int, Inst] = {}
+  offset = 0
+  while offset < len(buf):
+    remaining = buf[offset:]
+    fmt = detect_format(remaining, arch)
+    decoded = fmt.from_bytes(remaining)
+    addr_table[off+offset] = decoded
+    offset += decoded.size()
+  return addr_table
+
+def parse_branch(inst) -> int|None:
+  if "branch" in getattr(inst, "op_name", "").lower():
+    x = inst.simm16 & 0xffff
+    return (x - 0x10000 if x & 0x8000 else x)*4
+  return None
+
+COND_TAKEN, COND_NOT_TAKEN, UNCOND = range(3)
+def amdgpu_cfg(lib:bytes, target:str) -> dict:
+  # decode
+  pc_table = amd_decode(lib, target)
+  # get leaders
+  leaders:set[int] = {next(iter(pc_table))}
+  for pc, inst in pc_table.items():
+    if (offset:=parse_branch(inst)) is not None: leaders.update((pc+inst.size()+offset, pc+inst.size()))
+  # build the cfg
+  curr:int|None = None
+  blocks:dict[int, list[int]] = {}
+  paths:dict[int, dict[int, int]] = {}
+  for pc, inst in pc_table.items():
+    if pc in leaders:
+      paths[curr:=pc] = {}
+      blocks[pc] = []
+    else: assert curr is not None, f"no basic block found for {pc}"
+    blocks[curr].append(pc)
+    # otherwise a basic block can have exactly one or two paths
+    nx = pc+inst.size()
+    if (offset:=parse_branch(inst)) is not None:
+      if inst.op_name == "S_BRANCH": paths[curr][nx+offset] = UNCOND
+      else: paths[curr].update([(nx+offset, COND_TAKEN), (nx, COND_NOT_TAKEN)])
+    elif nx in leaders: paths[curr][nx] = UNCOND
+  pc_tokens:dict[int, list[dict]] = {}
+  from tinygrad.renderer.amd.dsl import Reg
+  for pc, inst in pc_table.items():
+    pc_tokens[pc] = tokens = []
+    for name, f in inst._fields:
+      if isinstance(val:=getattr(inst, name), Reg): tokens.append({"st":val.fmt(), "keys":[f"r{val.offset+i}" for i in range(val.sz)], "kind":1})
+      elif name in {"op","opx","opy"}: tokens.append({"st":(op_name:=val.name.lower()), "keys":[op_name], "kind":0})
+      elif name != "encoding" and val != f.default: tokens.append({"st":(s:=repr(val)), "keys":[s], "kind":1})
+  # show a smaller view for repeated instructions in the graph
+  lines:list[str] = []
+  disasm = {pc:str(inst) for pc,inst in pc_table.items()}
+  asm_width = max(len(asm) for asm in disasm.values())
+  for pcs in blocks.values():
+    new_pcs:list[int] = []
+    i, n = 0, len(pcs)
+    while i < n:
+      j = i+1
+      while j<n and pc_table[pcs[j]] == pc_table[pcs[i]]: j += 1
+      new_pcs.append(pcs[i])
+      if j-i>1:
+        pc_tokens[pcs[i]].append({"st":f"({j-i}x)", "keys":[], "kind":0})
+        for k in range(i+1, j): del pc_tokens[pcs[k]]
+      lines.append(f"{disasm[pcs[i]]:<{asm_width}}  # {pcs[i]:012X}"+(f"...{pcs[j-1]:012X} ({j-i}x)" if j-i>1 else ""))
+      i = j
+    pcs[:] = new_pcs
+  from tinygrad.runtime.autogen import amdgpu_kd
+  kd = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytearray(get_elf_section(lib, ".rodata").content))
+  vgpr_gran = kd.compute_pgm_rsrc1 & amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT
+  return {"data":{"blocks":blocks, "paths":paths, "pc_tokens":pc_tokens}, "src":"\n".join(lines), "lang":"python",
+          "metadata":[[{"label":f"{r} Alloc", "value":v} for r,v in [("VGPR", (vgpr_gran+1)*8-7), ("LDS", kd.group_segment_fixed_size),
+                                                                     ("Scratch", kd.private_segment_fixed_size)] if v>0]]}
+
+# ** Main render function to get the complete details about a trace event
+
+def get_render(viz_data:VizData, query:str) -> dict:
+  url = urlparse(query)
+  i, j, fmt = get_int(qs:=parse_qs(url.query), "ctx"), get_int(qs, "step"), url.path.lstrip("/")
+  data = viz_data.ctxs[i]["steps"][j]["data"]
+  if fmt == "graph-rewrites": return {"value":get_full_rewrite(viz_data, viz_data.trace.rewrites[i][j]), "content_type":"text/event-stream"}
+  if fmt == "uops": return {"src":get_stdout(lambda: print_uops(_reconstruct(viz_data, viz_data.trace.rewrites[i][j-1].sink).src[2].src))}
+  if fmt == "code": return {"src":data, "lang":"cpp"}
+  if fmt == "asm":
+    ret:dict = {}
+    renderer, lib = data
+    if renderer.target.arch.startswith("gfx"):
+      with soft_err(lambda err: ret.update(err)): ret.update(amdgpu_cfg(lib, renderer.target.arch))
+    else: ret["src"] = get_stdout(lambda: renderer.compiler.disassemble(lib))
+    return ret
+  if fmt == "all-pmc":
+    durations, pmc = data
+    ret = {"cols":{}, "rows":[]}
+    for (name, n, k),events in pmc.items():
+      pmc_table = unpack_pmc(events)
+      ret["cols"].update([(r[0], None) for r in pmc_table["rows"]])
+      ret["rows"].append((name, durations[k][n-1], *[r[1] for r in pmc_table["rows"]]))
+    ret["cols"] = ["Kernel", "Duration", *ret["cols"]]
+    return ret
+  if fmt == "prg-pmc": return unpack_pmc(data)
+  if fmt.startswith("sqtt"):
+    ret = {}
+    with soft_err(lambda err:ret.update(err)):
+      if (events:=get_profile(viz_data, list(itertools.islice(sqtt_timeline(*data), getenv("MAX_SQTT_PKTS", 50_000))), sort_fn=row_tuple)):
+        ret = {"value":events, "content_type":"application/octet-stream"}
+      else: ret = {"src":"No SQTT trace on this SE."}
+    return ret
+  # viewers for the amd decoder in extra
+  if fmt.startswith("amd-sqtt"): return data["fxn"](viz_data, i, j, *data["args"])
+  if fmt == "cu-sqtt": return {"value":get_profile(viz_data, data, sort_fn=row_tuple), "content_type":"application/octet-stream"}
+  if fmt == "prg-pma-pkts":
+    ret = {}
+    with soft_err(lambda err:ret.update(err)):
+      if (events:=get_profile(viz_data, pma_timeline(*data), sort_fn=row_tuple)): ret = {"value":events, "content_type":"application/octet-stream"}
+      else: ret = {"src":"No PMA samples found."}
+    return ret
+  return data
+
+# ** HTTP server
+
+def get_int(query:dict[str, list[str]], k:str) -> int: return int(query.get(k,["0"])[0])
+
+class Handler(HTTPRequestHandler):
+  def do_GET(self):
+    ret, status_code, content_type = b"", 200, "text/html"
+
+    if (url:=urlparse(self.path)).path == "/":
+      with open(os.path.join(os.path.dirname(__file__), "index.html"), "rb") as f: ret = f.read()
+    elif self.path.startswith(("/assets/", "/js/")) and '/..' not in self.path:
+      try:
+        with open(os.path.join(os.path.dirname(__file__), self.path.strip('/')), "rb") as f: ret = f.read()
+        if url.path.endswith(".js"): content_type = "application/javascript"
+        if url.path.endswith(".css"): content_type = "text/css"
+      except FileNotFoundError: status_code = 404
+
+    elif url.path == "/ctxs":
+      lst = [{"name":c["name"], "steps":[{k:v for k, v in s.items() if k != "data"} for s in c["steps"]]} for c in data.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:
+      if not (render_src:=get_render(data, self.path)): status_code = 404
+      else:
+        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"
+        if content_type == "text/event-stream": return self.stream_json(render_src["value"])
+
+    return self.send_data(ret, content_type, status_code)
+
+# ** main loop
+
+def reloader():
+  mtime = os.stat(__file__).st_mtime
+  while not stop_reloader.is_set():
+    if mtime != os.stat(__file__).st_mtime:
+      print("reloading server...")
+      os.execv(sys.executable, [sys.executable] + sys.argv)
+    time.sleep(0.1)
+
+T = TypeVar("T")
+# unpickling may load libraries, turn off DEBUG=3 output
+@Context(DEBUG=0)
+def load_pickle(path:str, default:T) -> T:
+  if not os.path.exists(path): return default
+  with open(path, "rb") as f: return pickle.load(f)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--rewrites-path', type=str, help='Path to rewrites', default=temp("rewrites.pkl", append_user=True))
+  parser.add_argument('--profile-path', type=str, help='Path to profile', default=temp("profile.pkl", append_user=True))
+  args = parser.parse_args()
+
+  with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+    if s.connect_ex(((HOST:="http://127.0.0.1").replace("http://", ""), PORT:=getenv("PORT", 8000))) == 0:
+      raise RuntimeError(f"{HOST}:{PORT} is occupied! use PORT= to change.")
+  stop_reloader = threading.Event()
+  multiprocessing.current_process().name = "VizProcess"
+  Context(ALLOW_DEVICE_USAGE=0).__enter__()                # disallow opening of devices
+  st = time.perf_counter()
+  print("*** viz is starting")
+
+  data = VizData(load_pickle(args.rewrites_path, default=RewriteTrace([], [], {})))
+  load_rewrites(data)
+  profile_ret = get_profile(data, load_pickle(args.profile_path, default=[]))
+
+  server = TCPServerWithReuse(('', PORT), Handler)
+  reloader_thread = threading.Thread(target=reloader)
+  reloader_thread.start()
+  print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"), flush=True)
+  if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}")
+  try: server.serve_forever()
+  except KeyboardInterrupt:
+    print("*** viz is shutting down...")
+    stop_reloader.set()