empty

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '11'
+  DOWNLOAD_CACHE_VERSION: '10'
   PYTHON_CACHE_VERSION: '2'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -46,11 +46,6 @@ jobs:
       with:
         deps: docs
         pydeps: "capstone"
-    - name: Build wheel and show size
-      run: |
-        pip install build
-        python -m build --wheel --outdir dist
-        ls -lh dist/*.whl
     - name: Use as an external package
       run: |
         mkdir $HOME/test_external_dir

test/test_softmax_fusion.py

@@ -117,14 +117,6 @@ class TestFuse(unittest.TestCase):
     c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
     self.assertListEqual(c.tolist(), [30]*16)
 
-  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
-  def test_fuse_and_tc_opt(self):
-    A = Tensor.randn(8, 8).realize()
-    B = Tensor.randn(8, 8).realize()
-    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
-    out = (C + (A @ B)).fuse()
-    out.realize()
-
 class TestSoftmaxFusion(unittest.TestCase):
   @classmethod
   def setUpClass(cls):

test/unit/test_dtype.py

@@ -53,13 +53,5 @@ class TestCastConvenienceMethod(unittest.TestCase):
       self.assertEqual(t.float().dtype, dtypes.float)
       self.assertEqual(t.double().dtype, dtypes.double)
 
-class TestDtypeTolist(unittest.TestCase):
-  def test_bfloat16(self):
-    self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
-    # 448
-    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
-    # 57344
-    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
-
 if __name__ == "__main__":
   unittest.main()

test/unit/test_viz.py

@@ -250,7 +250,7 @@ class TestVizProfiler(unittest.TestCase):
 
     j = json.loads(get_profile(prof))
 
-    dev_events = j['layout']['NV']['shapes']
+    dev_events = j['layout']['NV']['timeline']['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']['shapes'][0]
+    event = j['layout']['NV']['timeline']['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))
 
-    tracks = list(j['layout'])
-    self.assertEqual(tracks[0], 'NV Graph')
-    self.assertEqual(tracks[2], 'NV')
-    self.assertEqual(tracks[4], 'NV:1')
+    devices = list(j['layout'])
+    self.assertEqual(devices[0], 'NV Graph')
+    self.assertEqual(devices[1], 'NV')
+    self.assertEqual(devices[2], 'NV:1')
 
-    nv_events = j['layout']['NV']['shapes']
+    nv_events = j['layout']['NV']['timeline']['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']['shapes']
+    nv1_events = j['layout']['NV:1']['timeline']['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']['shapes']
+    graph_events = j['layout']['NV Graph']['timeline']['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"][f"{a.device} Memory"]
+    ret = profile_ret["layout"][a.device]["mem"]
     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"][f"{b.device} Memory"]
+    ret = profile_ret["layout"][b.device]["mem"]
     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"][f"{c.device} Memory"]
+    ret = profile_ret["layout"][c.device]["mem"]
     self.assertEqual(ret["peak"], 2)
     self.assertEqual(ret["shapes"][0]["x"], [0, 3])
     self.assertEqual(ret["shapes"][1]["x"], [1, 3, 3, 4])

tinygrad/codegen/__init__.py

@@ -87,7 +87,7 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
 
   # decompositions
   pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)
-  ret.append(RewriteStep(pm_decomp, lambda _: opts.device, name="decompositions"))
+  ret.append(RewriteStep(pm_decomp, name="decompositions"))
 
   # final rules for the renderer (without sym)
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher

tinygrad/codegen/opt/kernel.py

@@ -378,9 +378,9 @@ class Kernel:
       tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
       for tc in tensor_cores:
         tensor_core_opts = [self._create_tc_opts(reduceop, tc, axis, opt_level) for reduceop in self.reduceops]
-        if tensor_core_opts[0] is None: continue
         # can only fuse reduces with the same tc options
         assert all_same(tensor_core_opts)
+        if tensor_core_opts[0] is None: continue
         self.tensor_core_opts = tc_opts = tensor_core_opts[0]
 
         # attempt to pad the tensor axes that require it

tinygrad/runtime/ops_null.py

@@ -1,11 +1,12 @@
-from tinygrad.device import Compiled, Compiler, Allocator
+from tinygrad.device import Compiled, Compiler, Renderer, Allocator
+from tinygrad.uop.ops import Ops
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.renderer.cstyle import CStyleLanguage
 
-class NullRenderer(CStyleLanguage):
+class NullRenderer(Renderer):
   device = "NULL"
+  code_for_op = {k:lambda:None for k in [Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT]}
   has_local = False
-  float4 = "float4"
+  def render(self, uops:list) -> str: return ""
 
 class NullProgram:
   def __init__(self, name:str, lib:bytes): pass

tinygrad/runtime/support/am/amdev.py

@@ -29,7 +29,7 @@ class AMFirmware:
     # Load SOS firmware
     self.sos_fw = {}
 
-    blob, sos_hdr = self.load_fw(f"psp_{fmt_ver(am.MP0_HWIP)}_sos.bin", versioned_header='struct_psp_firmware_header')
+    blob, sos_hdr = self.load_fw(f"psp_{fmt_ver(am.MP0_HWIP)}_sos.bin", am.struct_psp_firmware_header_v2_0)
     fw_bin = sos_hdr.psp_fw_bin
 
     for fw_i in range(sos_hdr.psp_fw_bin_count):
@@ -45,11 +45,11 @@ class AMFirmware:
     self.smu_psp_desc = self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.header.ucode_size_bytes, am.GFX_FW_TYPE_SMU)
 
     # SDMA firmware
-    blob, hdr = self.load_fw(f"sdma_{fmt_ver(am.SDMA0_HWIP)}.bin", versioned_header='struct_sdma_firmware_header')
+    blob, hdr, hdr_v3 = self.load_fw(f"sdma_{fmt_ver(am.SDMA0_HWIP)}.bin", am.struct_sdma_firmware_header_v2_0, am.struct_sdma_firmware_header_v3_0)
     if hdr.header.header_version_major < 3:
       self.descs += [self.desc(blob, hdr.ctl_ucode_offset, hdr.ctl_ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH1)]
       self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ctx_ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
-    else: self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
+    else: self.descs += [self.desc(blob, hdr_v3.header.ucode_array_offset_bytes, hdr_v3.ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
 
     # PFP, ME, MEC firmware
     for (fw_name, fw_cnt) in ([('PFP', 1), ('ME', 1)] if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else []) + [('MEC', 1)]:
@@ -83,13 +83,10 @@ class AMFirmware:
 
     self.descs += [self.desc(blob, hdr0.header.ucode_array_offset_bytes, hdr0.header.ucode_size_bytes, am.GFX_FW_TYPE_RLC_G)]
 
-  def load_fw(self, fname:str, *headers, versioned_header:str|None=None):
+  def load_fw(self, fname:str, *headers):
     fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/45f59212aebd226c7630aff4b58598967c0c8c91/amdgpu/{fname}", subdir="fw")
     blob = memoryview(bytearray(fpath.read_bytes()))
     if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}")
-    if versioned_header:
-      chdr = am.struct_common_firmware_header.from_address(mv_address(blob))
-      headers += (getattr(am, versioned_header + f"_v{chdr.header_version_major}_{chdr.header_version_minor}"),)
     return tuple([blob] + [hdr.from_address(mv_address(blob)) for hdr in headers])
 
   def desc(self, blob:memoryview, offset:int, size:int, *types:int) -> tuple[list[int], memoryview]: return (list(types), blob[offset:offset+size])
@@ -226,7 +223,7 @@ class AMDev(PCIDevImplBase):
 
     self.bhdr = am.struct_binary_header.from_buffer(bytearray(self.vram.view(tmr_offset, tmr_size)[:]))
     ihdr = am.struct_ip_discovery_header.from_address(ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.IP_DISCOVERY].offset)
-    assert self.bhdr.binary_signature == am.BINARY_SIGNATURE and ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE, "discovery signatures mismatch"
+    assert ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE and not ihdr.base_addr_64_bit, f"0x{ihdr.signature:X} != 0x{am.DISCOVERY_TABLE_SIGNATURE:X}"
 
     # Mapping of HW IP to Discovery HW IP
     hw_id_map = {am.__dict__[x]: int(y) for x,y in am.hw_id_map}
@@ -259,3 +256,4 @@ class AMDev(PCIDevImplBase):
     for prefix, hwip in mods:
       self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip])))
     self.__dict__.update(import_asic_regs('mp', (11, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))
+

tinygrad/schedule/multi.py

@@ -13,10 +13,6 @@ def handle_allreduce_multirank(buf:UOp, red:UOp) -> UOp|None:
   for i,dev in enumerate(buf.device):
     groups.setdefault(Device[dev].group_id, []).append(buf.mselect(i))
 
-  # Put reduce leader of each group first
-  reduce_leaders = set(getenv("REDUCE_LEADERS", "").split(","))
-  groups = {gid: sorted(bufs, key=lambda x: (x.device not in reduce_leaders, x.device)) for gid,bufs in groups.items()}
-
   # Skip if only one group or if every group has only one buffer
   if len(groups) <= 1 or not any(len(g) > 1 for g in groups.values()): return None
 

tinygrad/tensor.py

@@ -345,7 +345,6 @@ class Tensor(MathTrait):
     print(t.tolist())
     ```
     """
-    if self.dtype in (dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
     return self.data().tolist()
 
   def numpy(self) -> 'np.ndarray':  # type: ignore [name-defined] # noqa: F821

tinygrad/uop/__init__.py

@@ -12,14 +12,11 @@ class Ops(FastEnum):
   NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
 
   # track children
-  CHILD = auto(); CHILDREN = auto() # noqa: E702
+  CHILD = auto()
 
   # buffer ops
   COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
 
-  # create buffer
-  BUFFERIZE = auto()
-
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702
 

tinygrad/uop/ops.py

@@ -658,9 +658,6 @@ class UPat(MathTrait):
   @staticmethod
   def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b)
 
-  # lil helper
-  def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs)
-
   # copied from UOp
   def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def index(self, idx:UPat, valid:UPat|None=None): return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx))
@@ -674,7 +671,6 @@ class UPat(MathTrait):
   def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.dtype, src=(self,)+src, **kwargs)
   def fuse(self): return self.alu(Ops.FUSE)
   def or_broadcasted(self, **kwargs): return UPat.any(self, UPat(Ops.VECTORIZE, self.dtype, src=self, **kwargs))
-  def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
 
   def const_like(self, b:ConstLike): return UPat.const(self.dtype, cast(ConstType, b))
   def alu(self, op:Ops, *src:UPat):

tinygrad/uop/symbolic.py

@@ -242,18 +242,18 @@ def gep_through_wmma(gep:UOp, wmma:UOp):
 
 gep_pushing = PatternMatcher([
   # GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST
-  (UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'),
+  (UPat(Ops.GEP, src=(UPat(Ops.GEP, name='g2'),), name='g1'),
    lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))),
-  (UPat(Ops.VECTORIZE, name='vec').f(Ops.GEP, name='gep'),
+  (UPat(Ops.GEP, src=(UPat(Ops.VECTORIZE, name="vec"),), name="gep"),
    lambda gep, vec: UOp(Ops.VECTORIZE, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
-  (UPat.cvar("c", vec=False).f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
-  (UPat(Ops.VCONST, name="c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
+  (UPat(Ops.GEP, src=(UPat.cvar("c", vec=False),), name="gep"), lambda gep, c: gep.const_like(c.arg)),
+  (UPat(Ops.GEP, src=(UPat(Ops.VCONST, name="c"),), name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
   # GEP on void is skipped
   (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
   # GEP in order is removed
   (UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
   # push all GEPs through ALUs (fix arange stuff)
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, name='gep'),
+  (UPat(Ops.GEP, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu'),), name='gep'),
    lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \
      if not isinstance(gep.dtype, PtrDType) else None),
   # CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later)
@@ -262,7 +262,7 @@ gep_pushing = PatternMatcher([
   # VECTORIZE on same GEP
   (UPat(Ops.VECTORIZE, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))),
   # push some GEPs through WMMAs
-  (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
+  (UPat(Ops.GEP, src=(UPat(Ops.WMMA, name="wmma"),), name="gep"), gep_through_wmma),
 ])
 
 commutative = PatternMatcher([

tinygrad/viz/index.html

@@ -228,7 +228,7 @@
   }
   #device-list > div {
     min-height: 32px;
-    max-width: 132px;
+    max-width: 100px;
     overflow-x: auto;
     overflow-y: hidden;
     white-space: nowrap;

tinygrad/viz/js/index.js

@@ -122,11 +122,11 @@ const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#1d2e62", "#63b0cd"
   CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
 const cycleColors = (lst, i) => lst[i%lst.length];
 
-const createPolygons = (source, height) => {
+const createPolygons = (source, area) => {
   const shapes = [];
-  const yscale = d3.scaleLinear().domain([0, source.peak]).range([height, 0]);
+  const yscale = d3.scaleLinear().domain([0, source.peak]).range([area, 0]);
   for (const [i,e] of source.shapes.entries()) {
-    const x = e.x.map((i,_) => source.timestamps[i]-data.st);
+    const x = e.x.map((i,_) => (source.timestamps[i] ?? data.et)-data.st);
     const y0 = e.y.map(yscale);
     const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
     const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
@@ -135,20 +135,6 @@ const createPolygons = (source, height) => {
   return shapes;
 }
 
-const rescaleTrack = (source, tid, k) => {
-  for (const e of source.shapes) {
-    for (let i=0; i<e.y0.length; i++) {
-      e.y0[i] = e.y0[i]*k;
-      e.y1[i] = e.y1[i]*k;
-    }
-  }
-  const change = (source.height*k)-source.height;
-  const div = document.getElementById(tid);
-  div.style.height = rect(div).height+change+"px";
-  source.height = source.height*k;
-  return change;
-}
-
 const drawLine = (ctx, x, y) => {
   ctx.beginPath();
   ctx.moveTo(x[0], y[0]);
@@ -164,61 +150,77 @@ async function renderProfiler() {
   // layout once!
   if (data != null) return;
   const profiler = d3.select(".profiler").html("");
-  const { layout, st, et } = await (await fetch("/get_profile")).json();
-  // place devices on the y axis and set vertical positions
-  const [tickSize, padding] = [10, 8];
-  const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+"px");
+  const deviceList = profiler.append("div").attr("id", "device-list").node();
   const canvas = profiler.append("canvas").attr("id", "timeline").node();
   // NOTE: scrolling via mouse can only zoom the graph
   canvas.addEventListener("wheel", e => (e.stopPropagation(), e.preventDefault()), { passive:false });
+  const profileRet = await (await fetch("/get_profile")).json()
+  const { layout, st, et } = profileRet;
+  // place devices on the y axis and set vertical positions
+  const [tickSize, padding] = [10, 8];
+  deviceList.style.paddingTop = `${tickSize+padding}px`;
   const ctx = canvas.getContext("2d");
   const canvasTop = rect(canvas).top;
   // color by key (name/category/device)
   const colorMap = new Map();
   data = {tracks:new Map(), axes:{}, st, et};
-  const heightScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.peak||0), 0)]).range([4,maxheight=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;
-    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 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.appendChild(document.createElement("div"));
+    div.innerText = k;
+    div.style.padding = `${padding}px`;
+    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;
       }
-      div.style("height", levelHeight*v.maxDepth+padding+"px").style("pointerEvents", "none");
-    } else {
-      const height = heightScale(v.peak);
-      data.tracks.set(k, { shapes:createPolygons(v, height), offsetY, height, peak:v.peak, scaleFactor:maxheight*4/height });
-      div.style("height", height+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.height, t.offsetY], fmt:"B" } : null;
-        focusedDevice = newFocus;
-        return resize();
-      });
+      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 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};
+      }
+      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 });
+    }
+    // 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

tinygrad/viz/serve.py

@@ -19,7 +19,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
                Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
-               Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e"}
+               Ops.CHILD: "#80fff0", Ops.REWRITE_ERROR: "#ff2e2e"}
 
 # VIZ API
 
@@ -73,8 +73,8 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
     if u.dtype != dtypes.void: label += f"\n{u.dtype}"
     for idx,x in enumerate(u.src):
       if x in excluded:
-        arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(u.dtype) else f"{x.arg}"
-        label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
+        if x.op is Ops.CONST and dtypes.is_float(u.dtype): label += f"\nCONST{idx} {x.arg:g}"
+        else: label += f"\n{x.op.name}{idx} {x.arg}"
     try:
       if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None:
         label += f"\n{shape_to_str(u.shape)}"
@@ -144,7 +144,7 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
     shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
   return {"shapes":shapes, "maxDepth":len(levels)}
 
-def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dict:
+def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
   step, peak, mem = 0, 0, 0
   shps:dict[int, dict] = {}
   temp:dict[int, dict] = {}
@@ -170,10 +170,9 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dic
   for v in temp.values():
     v["x"].append(step)
     v["y"].append(v["y"][-1])
-  timestamps.append(max_ts)
   return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
 
-def get_profile(profile:list[ProfileEvent]) -> bytes|None:
+def get_profile(profile:list[ProfileEvent]):
   # start by getting the time diffs
   for ev in profile:
     if isinstance(ev,ProfileDeviceEvent): device_ts_diffs[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
@@ -185,14 +184,10 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
     dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e))
     if min_ts is None or st < min_ts: min_ts = st
     if max_ts is None or et > max_ts: max_ts = et
-  if min_ts is None: return None
   # return layout of per device events
-  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, unwrap(max_ts))
-  return json.dumps({"layout":layout, "st":min_ts, "et":max_ts}).encode("utf-8")
+  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")
 
 def get_runtime_stats(key) -> list[dict]:
   ret:list[dict] = []