diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index 1d8189203d..c253d2674d 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -250,7 +250,7 @@ class TestVizProfiler(unittest.TestCase): j = json.loads(get_profile(prof)) - dev_events = j['layout']['NV']['timeline']['shapes'] + dev_events = j['layout']['NV']['shapes'] self.assertEqual(len(dev_events), 1) event = dev_events[0] self.assertEqual(event['name'], 'E_2') @@ -263,7 +263,7 @@ class TestVizProfiler(unittest.TestCase): j = json.loads(get_profile(prof)) - event = j['layout']['NV']['timeline']['shapes'][0] + event = j['layout']['NV']['shapes'][0] self.assertEqual(event['name'], 'COPYxx') self.assertEqual(event['st'], 900) # diff clock self.assertEqual(event['dur'], 10) @@ -278,23 +278,23 @@ class TestVizProfiler(unittest.TestCase): j = json.loads(get_profile(prof)) - devices = list(j['layout']) - self.assertEqual(devices[0], 'NV Graph') - self.assertEqual(devices[1], 'NV') - self.assertEqual(devices[2], 'NV:1') + tracks = list(j['layout']) + self.assertEqual(tracks[0], 'NV Graph') + self.assertEqual(tracks[2], 'NV') + self.assertEqual(tracks[4], 'NV:1') - nv_events = j['layout']['NV']['timeline']['shapes'] + nv_events = j['layout']['NV']['shapes'] self.assertEqual(nv_events[0]['name'], 'E_25_4n2') self.assertEqual(nv_events[0]['st'], 0) self.assertEqual(nv_events[0]['dur'], 2) #self.assertEqual(j['devEvents'][6]['pid'], j['devEvents'][0]['pid']) - nv1_events = j['layout']['NV:1']['timeline']['shapes'] + nv1_events = j['layout']['NV:1']['shapes'] self.assertEqual(nv1_events[0]['name'], 'NV -> NV:1') self.assertEqual(nv1_events[0]['st'], 954) #self.assertEqual(j['devEvents'][7]['pid'], j['devEvents'][3]['pid']) - graph_events = j['layout']['NV Graph']['timeline']['shapes'] + graph_events = j['layout']['NV Graph']['shapes'] self.assertEqual(graph_events[0]['st'], nv_events[0]['st']) self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur']) @@ -308,7 +308,7 @@ class TestVizMemoryLayout(BaseTestViz): a = _alloc(1) _b = _alloc(1) profile_ret = json.loads(get_profile(Buffer.profile_events)) - ret = profile_ret["layout"][a.device]["mem"] + ret = profile_ret["layout"][f"{a.device} Memory"] self.assertEqual(ret["peak"], 2) self.assertEqual(ret["shapes"][0]["x"], [0, 2]) self.assertEqual(ret["shapes"][1]["x"], [1, 2]) @@ -318,7 +318,7 @@ class TestVizMemoryLayout(BaseTestViz): del a b = _alloc(1) profile_ret = json.loads(get_profile(Buffer.profile_events)) - ret = profile_ret["layout"][b.device]["mem"] + ret = profile_ret["layout"][f"{b.device} Memory"] self.assertEqual(ret["peak"], 1) self.assertEqual(ret["shapes"][0]["x"], [0, 2]) self.assertEqual(ret["shapes"][1]["x"], [2, 3]) @@ -331,7 +331,7 @@ class TestVizMemoryLayout(BaseTestViz): del a c = _alloc(1) profile_ret = json.loads(get_profile(Buffer.profile_events)) - ret = profile_ret["layout"][c.device]["mem"] + ret = profile_ret["layout"][f"{c.device} Memory"] self.assertEqual(ret["peak"], 2) self.assertEqual(ret["shapes"][0]["x"], [0, 3]) self.assertEqual(ret["shapes"][1]["x"], [1, 3, 3, 4]) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index 50d952ef0e..dd1368d9ed 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -228,7 +228,7 @@ } #device-list > div { min-height: 32px; - max-width: 100px; + max-width: 132px; overflow-x: auto; overflow-y: hidden; white-space: nowrap; diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 96c09f61a3..8772568ef5 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -135,6 +135,20 @@ const createPolygons = (source, area) => { return shapes; } +const rescaleTrack = (source, tid, k) => { + for (const e of source.shapes) { + for (let i=0; i { ctx.beginPath(); ctx.moveTo(x[0], y[0]); @@ -162,61 +176,49 @@ async function renderProfiler() { // color by key (name/category/device) const colorMap = new Map(); data = {tracks:new Map(), axes:{}, st, et}; - const areaScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.mem.peak), 0)]).range([4,maxArea=100]); - for (const [k, { timeline, mem }] of Object.entries(layout)) { - if (timeline.shapes.length === 0 && mem.shapes.length == 0) continue; - const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px").node(); - div.onclick = () => { // TODO: make this feature more visible - const prevScroll = profiler.node().scrollTop; - let newOffset = null; - for (const [track, v] of data.tracks) { - if (track === `${k} memory`) { - // expand the y axis or reset to default size - const pick = [areaScale(mem.peak), maxArea*4]; - const expand = k !== focusedDevice; - const [newArea, prevArea] = expand ? pick.reverse() : pick; - focusedDevice = expand ? k : null; - data.axes.y = expand ? { domain:[0, mem.peak], range:[v.offsetY+newArea, v.offsetY], fmt:"B" } : null; - // either way update all offsets - v.shapes = createPolygons(mem, newArea); - newOffset = newArea-prevArea; - v.div.style.height = rect(v.div).height+newOffset+"px"; - } else if (newOffset != null) v.offsetY += newOffset; - } - d3.select(canvas).call(canvasZoom.transform, zoomLevel); - if (prevScroll) profiler.node().scrollTop = prevScroll; - } - const { y:baseY, height:baseHeight } = rect(div); - const levelHeight = baseHeight-padding; + const areaScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.peak||0), 0)]).range([4,maxArea=100]); + for (const [k, v] of Object.entries(layout)) { + if (v.shapes.length === 0) continue; + const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px"); + const { y:baseY, height:baseHeight } = rect(div.node()); const offsetY = baseY-canvasTop+padding/2; - const shapes = []; - data.tracks.set(k, { shapes, offsetY }); - let colorKey, ref; - for (const e of timeline.shapes) { - if (e.depth === 0) colorKey = e.cat ?? e.name; - if (!colorMap.has(colorKey)) colorMap.set(colorKey, cycleColors(colorScheme[k] ?? colorScheme.DEFAULT, colorMap.size)); - const fillColor = d3.color(colorMap.get(colorKey)).brighter(e.depth).toString(); - const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width })); - if (e.ref != null) ref = {ctx:e.ref, step:0}; - else if (ref != null) { - const start = ref.step>0 ? ref.step+1 : 0; - const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name); - ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx}; + if (v.shapes[0].dur != null) { + const levelHeight = baseHeight-padding; + const shapes = []; + data.tracks.set(k, { shapes, offsetY }); + let colorKey, ref; + for (const e of v.shapes) { + if (e.depth === 0) colorKey = e.cat ?? e.name; + if (!colorMap.has(colorKey)) colorMap.set(colorKey, cycleColors(colorScheme[k] ?? colorScheme.DEFAULT, colorMap.size)); + const fillColor = d3.color(colorMap.get(colorKey)).brighter(e.depth).toString(); + const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width })); + if (e.ref != null) ref = {ctx:e.ref, step:0}; + else if (ref != null) { + const start = ref.step>0 ? ref.step+1 : 0; + const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name); + ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx}; + } + const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref }; + // offset y by depth + shapes.push({x:e.st-st, y:levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor }); } - const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref }; - // offset y by depth - shapes.push({x:e.st-st, y:levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor }); + div.style("height", levelHeight*v.maxDepth+padding+"px").style("pointerEvents", "none"); + } else { + const area = areaScale(v.peak); + data.tracks.set(k, { shapes:createPolygons(v, area), offsetY, area, peak:v.peak, scaleFactor:maxArea*4/area }); + div.style("height", area+padding+"px").style("cursor", "pointer").on("click", (e) => { + const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id; + let offset = 0; + for (const [tid, track] of data.tracks) { + track.offsetY += offset; + if (tid === newFocus) offset += rescaleTrack(track, tid, track.scaleFactor); + else if (tid === focusedDevice) offset += rescaleTrack(track, tid, 1/track.scaleFactor); + } + data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.area, t.offsetY], fmt:"B" } : null; + focusedDevice = newFocus; + return resize(); + }); } - // position shapes on the canvas and scale to fit fixed area - let area = mem.shapes.length === 0 ? 0 : areaScale(mem.peak); - if (area === 0) div.style.pointerEvents = "none"; - else { - const startY = offsetY+(levelHeight*timeline.maxDepth)+padding/2; - data.tracks.set(`${k} memory`, { shapes:createPolygons(mem, area), offsetY:startY, div }); - div.style.cursor = "pointer"; - } - // lastly, adjust device rect by number of levels - div.style.height = `${Math.max(levelHeight*timeline.maxDepth, baseHeight)+area+padding}px`; } updateProgress({ "show":false }); // draw events on a timeline diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 01cb8025ac..e9ad038c9e 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -185,9 +185,12 @@ def get_profile(profile:list[ProfileEvent]): if min_ts is None or st < min_ts: min_ts = st if max_ts is None or et > max_ts: max_ts = et # return layout of per device events - for events in dev_events.values(): events.sort(key=lambda v:v[0]) - 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") + layout:dict[str, dict] = {} + for k,v in dev_events.items(): + v.sort(key=lambda e:e[0]) + layout[k] = timeline_layout(v) + layout[f"{k} Memory"] = mem_layout(v) + return json.dumps({"layout":layout, "st":min_ts, "et":max_ts}).encode("utf-8") def get_runtime_stats(key) -> list[dict]: ret:list[dict] = []