github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-25 23:08:06 -05:00
Code Issues Packages Projects Releases Wiki Activity

hotfix: use captured uops in viz render [pr] (#9593)

Browse Source 
* hotfix: use captured uops in viz render [pr]

* better error
This commit is contained in:
qazal
2025-03-27 11:52:12 +08:00
committed by GitHub
parent e5ff7b23d7
commit 0877497bad
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tinygrad/viz/serve.py
View File

@@ -21,7 +21,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
# NOTE: if any extra rendering in VIZ fails, we don't crash
def pcall(fxn:Callable[..., str], *args, **kwargs) -> str:
try: return fxn(*args, **kwargs)
except Exception as e: return f"ERROR: {e}"
except Exception as e: return f"ERROR in {fxn.__name__}: {e}"
# ** Metadata for a track_rewrites scope
@@ -33,10 +33,10 @@ class GraphRewriteMetadata(TypedDict):
name: str|None # optional name of the rewrite
@functools.lru_cache(None)
def _prg(k:Kernel): return k.to_program().src
def render_program(k:Kernel): return k.opts.render(k.uops)
def to_metadata(k:Any, v:TrackedGraphRewrite) -> GraphRewriteMetadata:
return {"loc":v.loc, "match_count":len(v.matches), "code_line":lines(v.loc[0])[v.loc[1]-1].strip(),
"kernel_code":pcall(_prg, k) if isinstance(k, Kernel) else None, "name":v.name}
"kernel_code":pcall(render_program, k) if isinstance(k, Kernel) else None, "name":v.name}
def get_metadata(keys:list[Any], contexts:list[list[TrackedGraphRewrite]]) -> list[tuple[str, list[GraphRewriteMetadata]]]:
return [(k.name if isinstance(k, Kernel) else str(k), [to_metadata(k, v) for v in vals]) for k,vals in zip(keys, contexts)]