experiments with multi being range

test/test_rangeify.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, nn
-from tinygrad.helpers import Context, GlobalCounters, CI, CPU_LVP, getenv
+from tinygrad.helpers import Context, GlobalCounters, CI
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
 
 class TestRangeifyAssign(unittest.TestCase):
@@ -28,35 +28,6 @@ class TestRangeifyEdgeCase(unittest.TestCase):
     res = Tensor.cat(a, c, dim=0)
     self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
 
-@unittest.skipIf(CPU_LVP, "broken in LVP")
-class TestPcontig(unittest.TestCase):
-  def test_flash_attention(self):
-    if getenv("BIG") > 1:
-      # llama 8B
-      BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-    elif getenv("BIG") > 0:
-      # bigger
-      BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-    else:
-      BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(PCONTIG=2, DEBUG=2):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(DEBUG=2):
-      GlobalCounters.reset()
-      cmp = fa()
-    with Context(DEBUG=0):
-      mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-
 # *** non CI rangeify tests below this line ***
 
 N = 256
@@ -244,6 +215,33 @@ class TestRangeify(unittest.TestCase):
     out = blk._feed_forward(x)
     out.realize()
 
+  @unittest.skip("RANGEIFY=0 does nothing")
+  def test_flash_attention(self):
+    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+
+    # bigger
+    #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
+
+    # llama 8B
+    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
+
+    def fa():
+      Tensor.manual_seed(1337)
+      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+      return q.scaled_dot_product_attention(k, v).realize()
+
+    with Context(DEBUG=4):
+      GlobalCounters.reset()
+      ret = fa()
+    with Context(RANGEIFY=0):
+      with Context(DEBUG=2):
+        GlobalCounters.reset()
+        cmp = fa()
+      with Context(DEBUG=0):
+        mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
 # contiguous + reduce can support ranges?
 
 @unittest.skip("pm_rangeify no longer exists. test this in a different way")

test/unit/test_viz.py

@@ -333,7 +333,7 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
       for _ in range(event_count):
         alloc, ts, key = u("<BII")
         if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IBB") for _ in range(u("<I")[0])]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg":{"users":[u("<I")[0] for _ in range(u("<I")[0])]}})
   return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 class TestVizProfiler(unittest.TestCase):
@@ -498,19 +498,6 @@ class TestVizMemoryLayout(BaseTestViz):
     buffers = profile["layout"]["NULL Memory"]["events"]
     user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
     self.assertEqual(max(user_cnt), n)
-    input_buf = buffers.pop()
-    assert all(u[2] == 0 for u in input_buf["arg"]["users"])
-
-  def test_annotate_read_write(self):
-    a = Tensor.ones(4, device="NULL").contiguous().realize()
-    b = a.assign(a+2)
-    c = a+1
-    Tensor.realize(b, c)
-    buf_events = load_profile(cpu_events+Buffer.profile_events)["layout"]["NULL Memory"]["events"]
-    users = next((b["arg"]["users"] for b in buf_events if len(b["arg"].get("users",[])) == 3))
-    self.assertEqual(users[0][2], 1) # write Tensor.ones
-    self.assertEqual(users[1][2], 2) # read+write Tensor.assign
-    self.assertEqual(users[2][2], 0) # readonly
 
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/opt/__init__.py

@@ -17,9 +17,9 @@ class Opt:
   def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
 
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.MULTI: "m"}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta", AxisType.MULTI: "GREEN"}
 
 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):

tinygrad/codegen/opt/postrange.py

@@ -13,8 +13,8 @@ from tinygrad.renderer import Renderer
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+axis_to_pos = {AxisType.MULTI: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
+               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 class Scheduler:
   def __init__(self, ast:UOp, opts:Renderer):

tinygrad/codegen/simplify.py

@@ -69,24 +69,18 @@ pm_split_ranges = PatternMatcher([
 
 def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)
 
-def reduce_unparented(red:UOp):
-  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
-  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
-  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
-  if len(reduce_unparented) == 0: return None
-  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
+def reduce_rangeless(red:UOp):
+  # TODO: share code with reduce_unparented
+  if red.arg not in {Ops.ADD, Ops.MAX}: return None
+  if red.src[0].dtype != red.dtype: return None
+  if not no_range(red.src[0]): return None
+  ret = red.src[0]
   if red.arg is Ops.ADD:
-    for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
-  if red.arg is Ops.MUL:
-    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+    for r in red.src[1:]:
+      ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
   return ret
 
-pm_reduce_unparented = PatternMatcher([
-  # remove any ranges from a REDUCE that aren't referenced in the reduce source
-  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
-])
-
-pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
+pm_reduce_collapse = PatternMatcher([
   # lift x+y out of reduce on lt
   ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
   # lift x*y out of reduce
@@ -112,6 +106,8 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
   # AND on WHERE
   ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
     lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
+  # remove REDUCEs that no longer have a RANGE in the src
+  (UPat(Ops.REDUCE, name="red"), reduce_rangeless),
 ])+sym
 
 def reduce_collapse(red:UOp):
@@ -126,6 +122,23 @@ def reduce_collapse(red:UOp):
   sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse")
   return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None
 
+def reduce_unparented(red:UOp):
+  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
+  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
+  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
+  if len(reduce_unparented) == 0: return None
+  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
+  if red.arg is Ops.ADD:
+    for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  if red.arg is Ops.MUL:
+    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  return ret
+
+pm_reduce_unparented = PatternMatcher([
+  # remove any ranges from a REDUCE that aren't referenced in the reduce source
+  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
+])
+
 pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
   # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
   (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),

tinygrad/engine/realize.py

@@ -167,7 +167,6 @@ class ExecItem:
     bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
     if PROFILE:
       payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs]}
-      payload["outputs"], payload["inputs"] = (self.prg.p.outs, self.prg.p.ins) if isinstance(self.prg, CompiledRunner) else ([0], [1])
       cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, payload))
     et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
     if do_update_stats:
@@ -181,11 +180,10 @@ class ExecItem:
         header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
         ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
         flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
-        flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
-        mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
-          colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
+        flops_str = f"{flops*1e-9:9.2f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:9.2f} TFLOPS", 'green')
+        mem_str = f"{membw*1e-9:6.1f}|{ldsbw*1e-9:<7.1f} GB/s" if membw < 1e13 else colored(f"{membw*1e-12:6.1f}|{ldsbw*1e-12:<7.1f} TB/s", 'green')
         print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
+          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB"+
           ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
           f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
       self.prg.first_run = False

tinygrad/helpers.py

@@ -169,7 +169,6 @@ VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
-PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/runtime/ops_null.py

@@ -17,7 +17,7 @@ class NullRenderer(CStyleLanguage):
 class NullProgram:
   def __init__(self, device:str, name:str, lib:bytes): self.device, self.name = device, name
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
-    with cpu_profile(self.name, self.device): return 1e-3
+    with cpu_profile(self.name, self.device): return 1e-4
 
 class NullAllocator(Allocator['NullDevice']):
   def _alloc(self, size, options): pass
@@ -28,7 +28,7 @@ class NullAllocator(Allocator['NullDevice']):
   def _offset(self, buf, offset:int, size:int): pass
 
 class NullGraph(MultiGraphRunner):
-  def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-1
+  def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-3
 
 class NullDevice(Compiled):
   def __init__(self, device:str):

tinygrad/schedule/indexing.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey, PCONTIG, colored
+from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
@@ -40,13 +40,12 @@ class BufferizeOpts:
 
 @dataclass
 class IndexingContext:
-  realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
+  realize_map: dict[UOp, None] = field(default_factory=dict)
   range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
 
   # create ranges
   range_idx: Iterator[int] = field(default_factory=itertools.count)
   def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP):
-    # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
@@ -58,13 +57,8 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
     if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.ASSIGN and s.src[1].op is Ops.KERNEL):
       if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
     elif s in ctx.realize_map:
-      realized_ranges = ctx.realize_map[s]
-      assert isinstance(realized_ranges, list), "realize map must contain range list"
-      closed_ranges = tuple([r for i,r in enumerate(ctx.range_map[s][1]) if i in realized_ranges])
-      # None in the device assigns it a number later
-      opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
-      new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
-      if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
+      new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+tuple(ctx.range_map[s][1]), arg=BufferizeOpts(device=s.device), tag=s.tag)
+      if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
     new_srcs.append(new_src)
   # NOTE: do we need this?
   return x.replace(src=tns) if x.src != (tns:=tuple(new_srcs)) else None
@@ -157,8 +151,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     ending_ranges[x] = any(ending_ranges[u] for u in consumer_map[x])
 
     # if this element has weight and it's ending a range, we (force) realize it
-    if ending_ranges[x] and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}) and not (PCONTIG>1):
-      rctx.realize_map[x] = None
+    if ending_ranges[x] and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}): rctx.realize_map[x] = None
 
     # *** the ranges on the output are
     #  1. new if this op is realized
@@ -171,9 +164,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       out_rngs = tuple(rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape)
       # all ranges are ended now
       ending_ranges[x] = False
-      # mark all ranges as ended
-      assert rctx.realize_map[x] is None
-      rctx.realize_map[x] = list(range(len(out_rngs)))
     elif x.op in {Ops.MSTACK, Ops.MSELECT}:
       # treat MSTACK/MSELECT like SINK
       continue
@@ -185,29 +175,29 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       out_rngs = consumer_rngs[0]
     elif len(consumer_rngs) > 1:
       # if this has two consumers, we have to merge the ranges and might create new ones
-      all_rngs: list[tuple[UOp, ...]] = list(zip(*consumer_rngs))
+      all_rngs = list(zip(*consumer_rngs))
       rngs_valids = []
       for valid_rngs in all_rngs:
         local_rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs])
-        rngs_valids.append((local_rngs, valids))
+        # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
+        same_rngs = [x if x.op is not Ops.RANGE or resolve(x.src[0] != 1) else UOp.const(dtypes.index, 0) for x in local_rngs]
+        rngs_valids.append((local_rngs, valids, all_same(same_rngs)))
 
       # TODO: in RANGEIFY > 1 all_all_same isn't required
-      all_all_same = all(all_same(local_rngs) for local_rngs,_ in rngs_valids)
+      all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids)
       _out_rngs = []
-      _new_rngs = []
-      for i,(local_rngs,valids) in enumerate(rngs_valids):
+      for i,(local_rngs,valids,same_rngs) in enumerate(rngs_valids):
         # we compare the ranges without their valids
-        if all_all_same or (PCONTIG and all_same(local_rngs)):
+        if all_all_same:
           # the new valid is the OR of all the children valids
           minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
           _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
         else:
           _out_rngs.append(rctx.new_range(x.shape[i]))
-          _new_rngs.append(i)
       out_rngs = tuple(_out_rngs)
 
-      # we have to (partially) realize here if there's new ranges
-      if len(_new_rngs): rctx.realize_map[x] = _new_rngs
+      # we have to realize here if there's new ranges
+      if not all_all_same: rctx.realize_map[x] = None
 
     # TODO: some ops don't have shape, enable this after the `.st` property is removed
     #assert len(out_rngs) == len(x.shape), \
@@ -223,7 +213,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # apply movement ops
     if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.marg, rngs)
     # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
-    # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
     if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] = True
 
     # REDUCE_AXIS creates ranges for the axes it is reducing
@@ -231,13 +220,8 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
 
     if debug:
-      realized_ranges = rctx.realize_map.get(x, None)
-      disp = []
-      for i, (ri, ro) in enumerate(zip([r.render() for r in rngs], [r.render() for r in out_rngs])):
-        rng = f"{ri}" if ri == ro else f"{ri} -> {ro}"
-        if realized_ranges is not None and i in realized_ranges: rng = colored(rng, "yellow")
-        disp.append("["+rng+"]")
-      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}", ''.join(disp))
+      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}",
+            UOp.sink().index(*rngs).render(), " -> ", UOp.sink().index(*out_rngs).render())
 
     # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
     rctx.range_map[x] = (rngs, out_rngs)

tinygrad/schedule/multi.py

@@ -1,7 +1,7 @@
 from typing import cast
 import functools, itertools, operator
 from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv
-from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, track_rewrites, graph_rewrite_map, graph_rewrite
+from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, track_rewrites, graph_rewrite_map
 from tinygrad.device import Device
 
 # *** allreduce implementation ***
@@ -219,8 +219,4 @@ multi_pm = PatternMatcher([
 ])+replace_allreduce
 
 @track_rewrites()
-def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]:
-  if getenv("VIZ"): graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
-  ret = graph_rewrite_map(big_sink, multi_pm, name="multi_pm")
-  if getenv("VIZ"): graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
-  return ret
+def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]: return graph_rewrite_map(big_sink, multi_pm, name="multi_pm")

tinygrad/schedule/rangeify.py

@@ -108,6 +108,18 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
    (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes),
 ])
 
+# *****************
+
+pm_where_is_multi = PatternMatcher([
+  # move *0 through where
+  (UPat.var("gate").where(UPat.var("a"), 0) * UPat.var("b"), lambda gate,a,b: gate.where(a*b, 0)),
+  # move *0 through unary op
+  (UPat(Ops.CONTIGUOUS, src=(UPat.var("gate").where(UPat.var("a"), 0),), name="u"), lambda gate,a,u: gate.where(u.replace(src=(a,)), 0)),
+  # move where 0 through reduce if the reduce ranges are not in the gate
+  (UPat(Ops.REDUCE, src=(UPat.var("gate").where(UPat.var("a"), 0),), name="red", allow_any_len=True),
+   lambda gate,a,red: gate.where(red.replace(src=(a,)+red.src[1:]), 0) if all(r not in gate.ranges for r in red.src[1:]) else None),
+])
+
 # *****************
 # 3.5 cleanups
 
@@ -178,11 +190,8 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced
   # NOTE: if buf src is a const, we don't replace it
-  if getenv("REAL_SUBSTITUTE"):
-    return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST})
-  else:
-    replaces = flatten([(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST])
-    return UOp(Ops.SUBSTITUTE, dtype=src.dtype, src=(src, UOp(Ops.NOOP, src=tuple(replaces[0::2])), UOp(Ops.NOOP, src=tuple(replaces[1::2]))))
+  replaces = flatten([(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST])
+  return UOp(Ops.SUBSTITUTE, dtype=src.dtype, src=(src, UOp(Ops.NOOP, src=tuple(replaces[0::2])), UOp(Ops.NOOP, src=tuple(replaces[1::2]))))
 
 def pre_bufferize(b:UOp, x:UOp, copy:UOp):
   nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:])
@@ -343,6 +352,7 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
   if r.tag is not None: return None
+  if r.arg[-1] is AxisType.MULTI: return None
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
   ctx.range += 1
   return ret
@@ -417,7 +427,7 @@ class Kernel:
     return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"
 
 def split_store(ctx:list[UOp], x:UOp):
-  if len(x.ranges): return None
+  if len([r for r in x.ranges if r.arg[-1] != AxisType.MULTI]): return None
   if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None
 
   # local kernel rewrite
@@ -499,6 +509,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
   tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding
+
+  tsink = graph_rewrite(tsink, pm_where_is_multi, name="where_is_multi")
+
   tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
   # TODO: can you substitute and remove costly buffers at the same time?
   tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")

tinygrad/tensor.py

@@ -231,9 +231,9 @@ class Tensor(MathTrait):
     # verify Tensors match the spec
     if __debug__: type_verify(list(big_sink.toposort()), tensor_uop_spec)
 
-    if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
-      _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
-      big_sink = UOp.sink(*flatten([x.uop.src if x.uop.op is Ops.MULTI else [x.uop] for x in (self,)+lst]))
+    #if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
+    #  _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
+    #  big_sink = UOp.sink(*flatten([x.uop.src if x.uop.op is Ops.MULTI else [x.uop] for x in (self,)+lst]))
 
     becomes_map = get_rangeify_map(big_sink)
     _apply_map_to_tensors(becomes_map, name="Apply Kernelize Map")

tinygrad/uop/ops.py

@@ -15,7 +15,7 @@ if TYPE_CHECKING:
 class AxisType(Enum):
   def __repr__(self): return str(self)
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto()
+  THREAD = auto(); MULTI = auto() # noqa: E702
 
 range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
 
@@ -114,7 +114,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   @functools.cached_property
   def key(self) -> bytes:
     return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest()
-  def __repr__(self): return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=(%s))")
+  def __repr__(self):
+    if self.dtype == dtypes.index: return srender(self)  # makes shapes print nicely
+    return pretty_print(self, lambda x: f"{type(self).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=(%s))")
   def argstr(self): return f'({", ".join(map(str, self.arg))})' if self.op is Ops.REDUCE_AXIS else repr(self.arg)
   def tagstr(self): return f", tag={self.tag}" if self.tag is not None else ""
 
@@ -220,7 +222,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       match self.op:
         case Ops.RESHAPE:
           if not all(x >= 0 for x in self.marg): raise ValueError(f"shape can't contain negative numbers {self.marg}")
-          if prod(ps) != prod(self.marg): raise ValueError(f"bad reshape: {ps} -> {self.marg}")
+          #if prod(ps) != prod(self.marg): raise ValueError(f"bad reshape: {ps} -> {self.marg}")
           return self.marg
         case Ops.EXPAND:
           if len(ps) != len(self.marg) or not all(s==ns or (s==1 and ns>=0) for s,ns in zip(ps, self.marg)):
@@ -436,16 +438,30 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   def _unshard(self, axis:int) -> UOp:
     bsz, dcount = self.shape[axis], len(self.device)
-    dnum = UOp.variable("_device_num", 0, dcount-1)
+    #dnum = UOp.variable("_device_num", 0, dcount-1)
+    dnum = UOp.range(dcount, -10, AxisType.MULTI)
     return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))
 
   def _shard(self, axis:int) -> UOp:
     dcount = len(self.device)
-    dnum = UOp.variable("_device_num", 0, dcount-1)
+    dnum = UOp.range(dcount, -10, AxisType.MULTI)
     if self.shape[axis] % dcount != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {dcount=}")
     sz = self.shape[axis] // dcount
-    return self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
-  def shard(self, devices:tuple[str, ...], axis:int) -> UOp: return self.copy_to_device(devices)._shard(axis).multi(axis)
+    #ret = self.reshape(tuple(s if i != axis else dnum*sz for i,s in enumerate(self.shape)))
+    #return ret
+
+    #flatten([[s] if i != axis else [dcount, sz] for i,s in enumerate(self.shape)])))
+
+    #ret = self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
+    #print(ret.shape)
+    #print(dnum)
+    #dnum = UOp.variable("_device_num", 0, dcount-1)
+    # TODO: 0 isn't correct here
+    ret = self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
+    ret = ret.pad(tuple((0,0) if a != axis else (sz*dnum, sz*(dcount-1) - sz*dnum) for a in range(len(self.shape))))
+    return ret
+
+  def shard(self, devices:tuple[str, ...], axis:int) -> UOp: return self.copy_to_device(devices)._shard(axis) #.multi(axis)
 
   # *** from LazyBuffer ***
 
@@ -652,8 +668,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       new_count.subtract(div_fac.split_uop(Ops.MUL))
       if const%div_const==0 and all(v>=0 for v in new_count.values()): return math.prod([*new_count.elements(), self.const_like(const//div_const)])
     return None # generic None if we aren't sure
-  def sum(self:UOp, *uops:UOp) -> UOp: return functools.reduce(operator.or_ if self.dtype is dtypes.bool else operator.add, uops, self)
-  def prod(self:UOp, *uops:UOp) -> UOp: return functools.reduce(operator.and_ if self.dtype is dtypes.bool else operator.mul, uops, self)
   @property
   def vmin(self) -> ConstType: return self._min_max[0]
   @property

tinygrad/uop/symbolic.py

@@ -130,7 +130,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
 def lt_folding(x:UOp, c:int) -> UOp|None:
   p, np = partition(x.split_uop(Ops.ADD), lambda u: u.const_factor() == 1)
   if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d:
-    return cast(UOp, UOp.sum(*np).divides(d))<(c//d)
+    return cast(UOp, functools.reduce(operator.add, np).divides(d))<(c//d)
   return None
 
 def canonicalize_simplex(X:UOp) -> UOp|None:
@@ -144,7 +144,7 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
       u = u.src[0]
     if not (u.op in GroupOp.Irreducible and u.vmin >= 0): return None
     ret.append(u)
-  return UOp.sum(*ret) if changed else None
+  return functools.reduce(operator.add, ret) if changed else None
 
 def cancel_divmod(d: UOp, x: UOp, y: UOp) -> UOp|None:
   # simple cancel div/mod case when the range of the numerator lies within a single denominator interval
@@ -167,7 +167,7 @@ def remove_nested_mod(m: UOp, x: UOp, y: UOp) -> UOp|None:
         something_changed = True
         u = u.src[0]
     new_xs.append(u)
-  new_x: UOp = UOp.sum(*new_xs)
+  new_x: UOp = functools.reduce(operator.add, new_xs)
   if something_changed and new_x.vmin>=0: return new_x % y
   return None
 
@@ -300,7 +300,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
   ((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
   (-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)),  # -(x+c) -> -x + -c
-  (UPat.cvar("y") * (UPat.var("x", dtype=dtypes.index) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)),  # -(x+c) -> -x + -c
   # ** where folding **
   (UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), lambda cond, t, f: cond.where(f,t)
     if f.arg is not Invalid else None),
@@ -453,9 +452,9 @@ def simplify_valid(valid:UOp) -> UOp|None:
   something_changed = False
   valids = list(valid.split_uop(Ops.AND))
   for stmt in sorted(valids, key=lambda v: _valid_priority(v, valids)):
-    ret.append(uop_given_valid(UOp.prod(*ret), stmt) if ret else stmt)
+    ret.append(uop_given_valid(functools.reduce(operator.and_, ret), stmt) if ret else stmt)
     if ret[-1] is not stmt: something_changed = True
-  return UOp.prod(*ret) if something_changed else None
+  return functools.reduce(operator.and_, ret) if something_changed else None
 
 # ******** phase 3 is the complete symbolic, and deals with very complex things like loop rewriting and threefry transform ********
 
@@ -472,7 +471,7 @@ def reduce_mul_chain(r:UOp):
 
 def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None:
   if not (dropped_clauses:=[c for c in cond.split_uop(Ops.AND) if not any(r in x.ranges for r in c.ranges)]): return None
-  return UOp.const(dtypes.bool, True).prod(*[c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses]).where(x, i)
+  return functools.reduce(operator.and_, [c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses], UOp.const(dtypes.bool, True)).where(x, i)
 pm_drop_and_clauses = PatternMatcher([(UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), drop_and_clauses)])
 
 def where_on_load(l, c1, buf, x):
@@ -484,7 +483,7 @@ def where_on_load(l, c1, buf, x):
     and not c.op_in_backward_slice_with_self(Ops.LOAD)]
   if not (removed:=moved_clauses+duplicate_clauses): return None
   # aditionally we can drop the clause on the where if it already exists in the load
-  remaining_clause = UOp.const(dtypes.bool, True).prod(*[c for c in c1.split_uop(Ops.AND) if c not in removed])
+  remaining_clause = functools.reduce(operator.and_, [c for c in c1.split_uop(Ops.AND) if c not in removed], UOp.const(dtypes.bool, True))
   return remaining_clause.where(UOp.load(buf.index(x.get_idx().valid(functools.reduce(operator.and_, moved_clauses, c2)), *l.src[1:])), 0)
 pm_move_where_on_load = PatternMatcher([
   (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l"), 0), where_on_load),

tinygrad/viz/js/index.js

@@ -51,7 +51,7 @@ function addTags(root) {
   root.selectAll("text").data(d => [d]).join("text").text(d => d).attr("dy", "0.35em");
 }
 
-let workerUrl = null, worker = null;
+let [workerUrl, worker] = [null, null];
 async function initWorker() {
   const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/js/worker.js"].map(u => fetch(u)));
   workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
@@ -259,7 +259,7 @@ async function renderProfiler() {
           x += 1; y += nbytes; valueMap.set(ts, y);
         } else {
           const free = buf_shapes.get(key);
-          free.users = Array.from({ length: u32() }, () => ({name:strings[u32()], num:u8(), mode:u8()}));
+          free.users = Array.from({ length: u32() }, () => strings[u32()]);
           timestamps.push(ts); valueMap.set(ts, y);
           x += 1; y -= free.nbytes;
           free.x.push(x);
@@ -284,7 +284,7 @@ async function renderProfiler() {
         const info = html.appendChild(tabulate(rows).node());
         for (let u=0; u<users?.length; u++) {
           const p = html.appendChild(document.createElement("p")); p.style.marginTop = "4px"; p.style.cursor = "pointer";
-          const { name, num, mode } = users[u]; p.appendChild(colored(`[${u}] ${name} ${mode == 2 ? 'read+write' : mode == 1 ? 'write' : 'read'}@data${num}`));
+          const name = users[u]; p.appendChild(colored(`[${u}] ${name}`));
           p.onclick = () => {
             const cid = ctxs.findIndex(c => c.name === name);
             if (cid != null) setCtxWithHistory(cid-1);
@@ -345,10 +345,11 @@ async function renderProfiler() {
     for (const [_, { offsetY, shapes, visible, valueMap }] of data.tracks) {
       visible.length = 0;
       for (const e of shapes) {
-        const p = new Path2D();
-        if (e.width == null) { // generic polygon
+        // generic polygon
+        if (e.width == null) {
           if (e.x[0]>et || e.x.at(-1)<st) continue;
           const x = e.x.map(xscale);
+          const p = new Path2D();
           p.moveTo(x[0], offsetY+e.y0[0]);
           for (let i=1; i<x.length; i++) {
             p.lineTo(x[i], offsetY+e.y0[i]);
@@ -359,29 +360,32 @@ async function renderProfiler() {
           for (let i=x.length-1; i>=0; i--) p.lineTo(x[i], offsetY+e.y1[i]);
           p.closePath();
           ctx.fillStyle = e.fillColor; ctx.fill(p);
-        } else { // contiguous rect
-          if (e.x>et || e.x+e.width<st) continue;
-          const x = xscale(e.x);
-          const y = offsetY+e.y;
-          const width = xscale(e.x+e.width)-x;
-          p.rect(x, y, width, e.height);
-          visible.push({ y0:y, y1:y+e.height, x0:x, x1:x+width, arg:e.arg });
-          ctx.fillStyle = e.fillColor; ctx.fill(p);
-          // add label
-          let lw = 0;
-          const lx = x+2, ly = y+e.height/2;
-          for (let li=0; li<e.label?.length; li++) {
-            if (lw+e.label[li].width+(li===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
-              if (lw>0) ctx.fillText("...", lx+lw, ly);
-              break;
-            }
-            ctx.textBaseline = "middle";
-            ctx.fillStyle = e.label[li].color;
-            ctx.fillText(e.label[li].st, lx+lw, ly);
-            lw += e.label[li].width;
+          if (focusedShape?.key && e.arg?.key === focusedShape.key) { paths.push(p); }
+          continue;
+        }
+        // contiguous rect
+        if (e.x>et || e.x+e.width<st) continue;
+        const x = xscale(e.x);
+        const y = offsetY+e.y;
+        const width = xscale(e.x+e.width)-x;
+        ctx.fillStyle = e.fillColor; ctx.fillRect(x, y, width, e.height);
+        visible.push({ y0:y, y1:y+e.height, x0:x, x1:x+width, arg:e.arg });
+        // add label
+        if (e.label == null) continue;
+        ctx.textAlign = "left";
+        ctx.textBaseline = "middle";
+        let labelX = x+2, labelWidth = 0;
+        const labelY = y+e.height/2;
+        for (const [i,l] of e.label.entries()) {
+          if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
+            if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
+            break;
           }
+          ctx.fillStyle = l.color;
+          ctx.fillText(l.st, labelX, labelY);
+          labelWidth += l.width;
+          labelX += l.width;
         }
-        if (focusedShape?.key && e.arg?.key === focusedShape.key) { paths.push(p); }
       }
     }
     // draw axes
@@ -512,6 +516,11 @@ function codeBlock(st, language, { loc, wrap }={}) {
   return ret;
 }
 
+function appendTd(tr, value, unit=null) {
+  const fmt = (typeof value === "number" && !Number.isInteger(value)) ? value.toFixed(2) : value;
+  tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? "");
+}
+
 function setActive(e) {
   if (e == null) return;
   e.classList.add("active");
@@ -636,7 +645,7 @@ async function main() {
         tr.className = "main-row code-row";
         for (const [i,value] of r.entries()) {
           // string format scalar values
-          if (!Array.isArray(value)) tr.appendChild(document.createElement("td")).innerText = value;
+          if (!Array.isArray(value)) appendTd(tr, value);
           // display arrays in a bar graph
           else {
             const segmentsTd = tr.appendChild(document.createElement("td"));

tinygrad/viz/serve.py

@@ -153,12 +153,8 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
   return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None
 
 def encode_mem_free(key:int, ts:int, execs:list[ProfilePointEvent], scache:dict) -> bytes:
-  ei_encoding:list[tuple[int, int, int]] = [] # <[u32, u8, u8] [function name, buffer number and mode (2 = r/w, 1 = w, 0 = r)]
-  for e in execs:
-    num = next(i for i,k in enumerate(e.arg["bufs"]) if k == key)
-    mode = 2 if (num in e.arg["inputs"] and num in e.arg["outputs"]) else 1 if (num in e.arg["outputs"]) else 0
-    ei_encoding.append((enum_str(e.key, scache), num, mode))
-  return struct.pack("<BIII", 0, ts, key, len(ei_encoding))+b"".join(struct.pack("<IBB", *t) for t in ei_encoding)
+  kernel_names = [enum_str(ei.key, scache) for ei in execs]
+  return struct.pack(f"<BIII{len(kernel_names)}I", 0, ts, key, len(kernel_names), *kernel_names)
 
 def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
                scache:dict[str, int]) -> bytes|None: