diff --git a/tinygrad/device.py b/tinygrad/device.py index 3693c1973b..0c6b5e6b1f 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -336,7 +336,7 @@ if PROFILE: if not getenv("SQTT", 0): from tinygrad.uop.ops import launch_viz - launch_viz("PROFILE", fn) + launch_viz(PROFILE, fn) if __name__ == "__main__": for device in ALL_DEVICES: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index c39c992182..36f4dbe569 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -859,7 +859,7 @@ if TRACK_MATCH_STATS or PROFILE: with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f: print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}") pickle.dump((tracked_keys, tracked_ctxs, uop_fields), f) - if VIZ: launch_viz("VIZ", temp("rewrites.pkl", append_user=True)) + if VIZ: launch_viz(VIZ, temp("rewrites.pkl", append_user=True)) if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value): ret = [0,0,0.0,0.0] for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]+x[1][3]): @@ -869,9 +869,10 @@ if TRACK_MATCH_STATS or PROFILE: print(f"{ret[0]:6d} / {ret[1]:7d} -- {ret[3]*1000.:9.2f} / {(ret[2]+ret[3])*1000.:9.2f} ms -- TOTAL") print(f"{len(match_stats)} rules, {sum(v[0] > 0 for v in match_stats.values())} matched once") - def launch_viz(env_str:str, data:str): - os.environ[env_str] = "0" + def launch_viz(var:ContextVar, data:str): + os.environ[(env_str:=var.key)] = "0" os.environ[f"{env_str}_DATA"] = data + os.environ[f"{env_str}_VALUE"] = str(var.value) if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")): args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else [] args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else [] diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 45279f754c..e4a8ec3ce6 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -362,7 +362,7 @@ document.getElementById("zoom-to-fit-btn").addEventListener("click", () => { // **** main VIZ interfacae -function codeBlock(st, language, { loc, wrap }) { +function codeBlock(st, language, { loc, wrap }={}) { const code = document.createElement("code"); code.innerHTML = hljs.highlight(st, { language }).value; code.className = "hljs"; @@ -505,6 +505,15 @@ async function main() { const metadata = document.querySelector(".metadata"); const [code, lang] = ctx.fmt != null ? [ctx.fmt, "cpp"] : [ret[currentRewrite].uop, "python"]; metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeBlock(code, lang, { wrap:false })); + if (ctx.runtime_stats != null) { + const div = metadata.appendChild(document.createElement("div")); + div.style.maxHeight = "200px"; + div.style.overflow = "auto"; + for (const [i, s] of ctx.runtime_stats.entries()) { + const p = div.appendChild(document.createElement("p")); + p.innerText = `Run ${i+1} ${formatTime(s.duration)}`; + } + } // ** rewrite steps if (step.match_count >= 1) { const rewriteList = metadata.appendChild(document.createElement("div")); diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 51a542a292..a301d706c9 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -26,7 +26,8 @@ def get_metadata(keys:list[TracingKey], contexts:list[list[TrackedGraphRewrite]] ret = [] 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)} for s in v] - ret.append({"name":k.display_name, "fmt":k.fmt, "steps":steps}) + ret.append(r:={"name":k.display_name, "fmt":k.fmt, "steps":steps}) + if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0]) for key in k.keys: ref_map[key] = i return ret @@ -172,6 +173,13 @@ def get_profile(profile:list[ProfileEvent]): dev_layout = {k:{"timeline":timeline_layout(v), "mem":mem_layout(v)} for k,v in dev_events.items()} return json.dumps({"layout":dev_layout, "st":min_ts, "et":max_ts}).encode("utf-8") +def get_runtime_stats(key) -> list[dict]: + ret:list[dict] = [] + for e in profile: + if isinstance(e, ProfileRangeEvent) and e.en is not None and e.name == key: + ret.append({"device":e.device, "duration":float(e.en-e.st)}) + return ret + # ** HTTP server class Handler(BaseHTTPRequestHandler):