sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

viz: addrspace little colored box (#16427)

Browse Source 
* return addrspace

* layout

* render

* addrspace encodes color

* update colors

* in input_ast all are params are green

* update stroke
This commit is contained in:
qazal
2026-05-29 17:25:07 +09:00
committed by GitHub
parent 814d414f41
commit 54cfb794b8
4 changed files with 20 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_viz.py
+5 -2
View File
@@ -5,13 +5,13 @@ from typing import Generator
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
from tinygrad.uop.symbolic import sym
from tinygrad.dtype import dtypes
from tinygrad.dtype import dtypes, AddrSpace
from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap
from tinygrad.device import Buffer
from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render
from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors
from tinygrad.codegen import to_program_cache
from tinygrad.codegen import to_program
@@ -348,6 +348,9 @@ class TestVizIntegration(unittest.TestCase):
self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1")
self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
self.assertEqual(lst[2]["name"], prg.arg.name)
input_ast = next(viz.get_details(2, 0))["graph"].values()
for u in input_ast:
if u["label"].startswith("PARAM\n"): self.assertEqual(u["addrspace"], addrspace_colors[AddrSpace.GLOBAL])
# schedule graph CALL nodes have a link to jump to codegen
def test_link_sched_codegen(self):
tinygrad/viz/js/index.js
+4 -2
View File
@@ -57,9 +57,9 @@ function intersectRect(r1, r2) {
}
function addTags(root, path) {
root.selectAll("circle").data(d => d.rect ? [] : [d]).join("circle").attr("r", 5).style("fill", d => d.fill ?? null);
root.selectAll("circle").data(d => d.rect ? [] : [d]).join("circle").attr("r", 5).style("fill", d => d.fill ?? null).style("stroke", d => d.stroke ?? null);
root.selectAll("rect").data(d => d.rect ? [d] : []).join("rect").attr("x", d => -d.width/2).attr("y", d => -d.height/2)
.attr("width", d => d.width).attr("height", d => d.height).style("fill", d => d.fill ?? null);
.attr("width", d => d.width).attr("height", d => d.height).style("fill", d => d.fill ?? null).style("stroke", d => d.stroke ?? null);
if (path != null) root.selectAll("path").data(d => [d]).join("path").attr("d", path);
else root.selectAll("text").data(d => [d]).join("text").text(d => d.text).attr("dy", "0.35em");
}
@@ -115,6 +115,8 @@ const drawGraph = (data) => {
});
addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
.attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => ({ text:e.tag })));
addTags(nodes.selectAll("g.addrspace").data(d => d.addrspace != null ? [d] : []).join("g").attr("class", "tag addrspace")
.attr("transform", d => `translate(${d.width/2-8}, ${-d.height/2+8})`).datum(e => ({ rect:true, width:10, height:10, fill:e.addrspace, stroke:"none" })));
const CALL_TAG_WIDTH = 14;
addTags(nodes.selectAll("g.type").data(d => d.collapsible ? [d] : []).join("g").attr("class", d => `tag clickable ${d.collapsed ? 'collapsed' : 'expanded'}`)
.attr("transform", d => d.callNode ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
tinygrad/viz/js/worker.js
+4 -4
View File
@@ -31,7 +31,7 @@ const layoutCfg = (g, { blocks, paths, pc_tokens }) => {
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" });
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"};
@@ -45,9 +45,9 @@ 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, className:"overlay"});
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 }] of Object.entries(graph)) {
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")) {
@@ -56,7 +56,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
}
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});
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;
tinygrad/viz/serve.py
+7 -2
View File
@@ -43,7 +43,7 @@ from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupO
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
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",
@@ -56,6 +56,8 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
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
@@ -160,8 +162,11 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
# limit SOURCE labels line count
if u.op is Ops.SOURCE and len(lines:=label.split("\n")) > 40:
label = "\n".join(lines[:30]) + "\n..."
try: addrspace = u.addrspace
except Exception: addrspace = None
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}
"ref":ref, "tag":repr(u.tag) if u.tag is not None else None,
"addrspace":addrspace_colors.get(addrspace, None) if addrspace is not None else None}
return graph
def _reconstruct(data:VizData, a:int, depth:int|None=None):