tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 18:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

viz: disassembly viewer (#11393)

Browse Source 
* test

* CPU=1 disasm works

* METAL=1 disasm works

* fix that

* work

* can unwrap

* work p2

* don't crash
This commit is contained in:
qazal
2025-07-27 18:44:28 +03:00
committed by GitHub
parent 3bb232eb29
commit 3466a220de
8 changed files with 62 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/realize.py
+1 -1
View File
@@ -13,7 +13,7 @@ from tinygrad.uop.spec import type_verify
# **************** Program Creation ****************
@track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret.src))
@track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret.src, ret=ret))
def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
"""
Transform an AST into a ProgramSpec. May trigger BEAM search.
tinygrad/helpers.py
+1
View File
@@ -191,6 +191,7 @@ class TracingKey:
keys:tuple[str, ...]=() # optional keys to search for related traces
fmt:str|None=None # optional detailed formatting
cat:str|None=None # optional category to color this by
ret:Any=None
class ProfileEvent: pass
tinygrad/runtime/ops_metal.py
+6 -3
View File
@@ -1,6 +1,6 @@
import os, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
from typing import Any, cast
from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile
from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent
from tinygrad.renderer.cstyle import MetalRenderer
@@ -144,7 +144,10 @@ class MetalCompiler(Compiler):
with tempfile.NamedTemporaryFile(delete=True) as shader:
shader.write(lib)
shader.flush()
ret = os.system(f"cd {pathlib.Path(__file__).parents[2]}/extra/disassemblers/applegpu && python3 compiler_explorer.py {shader.name}")
proc = subprocess.Popen(f"cd {pathlib.Path(__file__).parents[2]}/extra/disassemblers/applegpu && python3 compiler_explorer.py {shader.name}",
stdout=subprocess.PIPE, shell=True, text=True, bufsize=1)
for line in unwrap(proc.stdout): print(line, end="")
ret = proc.wait()
if ret: print("Disassembler Error: Make sure you have https://github.com/dougallj/applegpu cloned to tinygrad/extra/disassemblers/applegpu")
class MetalProgram:
tinygrad/viz/assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/x86asm.min.js
Vendored
+19
View File
File diff suppressed because one or more lines are too long
tinygrad/viz/fetch_assets.sh
+1
View File
@@ -10,5 +10,6 @@ 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/x86asm.min.js"
fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js"
fetch "unpkg.com/@highlightjs/[email protected]/styles/tokyo-night-dark.min.css"
tinygrad/viz/index.html
+9
View File
@@ -10,6 +10,7 @@
<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>
<script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/x86asm.min.js"></script>
<link rel="stylesheet" href="assets/unpkg.com/@highlightjs/[email protected]/styles/tokyo-night-dark.min.css" />
<style>
* {
@@ -223,6 +224,14 @@
#device-list > div:hover {
background-color: rgba(20, 23, 35, 0.3);
}
.raw-text {
padding: 0 8px;
width: 100%;
height: 100%;
}
.raw-text code {
max-height: none !important;
}
</style>
</head>
<body>
tinygrad/viz/js/index.js
+12 -1
View File
@@ -456,7 +456,7 @@ async function main() {
for (const [j,u] of steps.entries()) {
const inner = ul.appendChild(document.createElement("ul"));
inner.id = `step-${i}-${j}`;
inner.innerText = `${u.name ?? u.loc[0].replaceAll("\\", "/").split("/").pop()+':'+u.loc[1]} - ${u.match_count}`;
inner.innerText = `${u.name ?? u.loc[0].replaceAll("\\", "/").split("/").pop()+':'+u.loc[1]}`+(u.match_count != null ? `- ${u.match_count}` : '');
inner.style.marginLeft = `${8*u.depth}px`;
inner.onclick = (e) => {
e.stopPropagation();
@@ -483,6 +483,17 @@ async function main() {
if (ckey in cache) {
ret = cache[ckey];
}
// ** Disassembly view
if (ckey.startsWith("/disasm")) {
if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
displayGraph("profiler");
document.querySelector(".metadata").innerHTML = "";
const root = document.createElement("div");
root.className = "raw-text";
root.appendChild(codeBlock(ret.src, "x86asm"));
return document.querySelector(".profiler").replaceChildren(root);
}
// ** UOp view (default)
// if we don't have a complete cache yet we start streaming rewrites in this step
if (!(ckey in cache) || (cache[ckey].length !== step.match_count+1 && activeSrc == null)) {
ret = [];
tinygrad/viz/serve.py
+13 -3
View File
@@ -1,12 +1,14 @@
#!/usr/bin/env python3
import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs
import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io
from contextlib import redirect_stdout
from decimal import Decimal
from http.server import BaseHTTPRequestHandler
from urllib.parse import parse_qs, urlparse
from typing import Any, TypedDict, Generator
from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey
from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint
from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfilePointEvent
from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfilePointEvent, Device
from tinygrad.renderer import ProgramSpec
from tinygrad.dtype import dtypes
uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
@@ -27,6 +29,7 @@ def get_metadata(keys:list[TracingKey], contexts:list[list[TrackedGraphRewrite]]
for i,(k,v) in enumerate(zip(keys, contexts)):
steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc),
"query":f"/ctxs?ctx={i}&idx={j}"} for j,s in enumerate(v)]
if isinstance(k.ret, ProgramSpec): steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
ret.append(r:={"name":k.display_name, "fmt":k.fmt, "steps":steps})
# use the first key to get runtime profiling data about this context
if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0])
@@ -182,6 +185,12 @@ def get_runtime_stats(key) -> list[dict]:
ret.append({"device":e.device, "duration":float(e.en-e.st)})
return ret
def get_disassembly(ctx:list[str]):
if not isinstance(prg:=contexts[0][int(ctx[0])].ret, ProgramSpec): return
lib = Device[prg.device].compiler.compile(prg.src)
with redirect_stdout(buf:=io.StringIO()): Device[prg.device].compiler.disassemble(lib)
return json.dumps({"src":buf.getvalue()}).encode()
# ** HTTP server
class Handler(BaseHTTPRequestHandler):
@@ -197,7 +206,8 @@ class Handler(BaseHTTPRequestHandler):
if url.path.endswith(".css"): content_type = "text/css"
except FileNotFoundError: status_code = 404
elif (query:=parse_qs(url.query)):
return self.stream_json(get_details(contexts[1][int(query["ctx"][0])][int(query["idx"][0])]))
if url.path == "/disasm": ret, content_type = get_disassembly(**query), "application/json"
else: return self.stream_json(get_details(contexts[1][int(query["ctx"][0])][int(query["idx"][0])]))
elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
elif url.path == "/get_profile" and profile_ret is not None: ret, content_type = profile_ret, "application/json"
else: status_code = 404