diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 5d26125992..e04d61c2ac 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -451,9 +451,6 @@ pm_add_range_tags = PatternMatcher([ class Kernel: ast: UOp metadata: tuple[Metadata, ...] = () - def __repr__(self): - ast_rep = f"SINK{tuple(s.op for s in self.ast.src)}" if self.ast.op is Ops.SINK else repr(self.ast.op) - return f"" def split_store(ctx:list[UOp], x:UOp) -> UOp|None: if len(x.ranges): return None diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 0f8e9a928f..6990280b0a 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -69,6 +69,9 @@ def uop_to_json(x:UOp, ignore_indexing=False) -> dict[int, dict]: if u in excluded: continue argst = codecs.decode(str(u.arg), "unicode_escape") if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg) + if u.op is Ops.KERNEL: + ast_str = f"SINK{tuple(s.op for s in u.arg.ast.src)}" if u.arg.ast.op is Ops.SINK else repr(u.arg.ast.op) + argst = f"" label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" for idx,x in enumerate(u.src[:1] if u.op in {Ops.BUFFERIZE, Ops.INDEX} else (u.src if u.op is not Ops.END else [])): @@ -88,7 +91,7 @@ def uop_to_json(x:UOp, ignore_indexing=False) -> dict[int, dict]: label += "\n" if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}" # NOTE: kernel already has metadata in arg - if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+repr(u.metadata) + if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+str(u.metadata) graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src) if x not in excluded], "color":uops_colors.get(u.op, "#ffffff"), "ref":ref, "tag":repr(u.tag) if u.tag is not None else None} return graph @@ -153,7 +156,7 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts: name = ctxs[ref]["name"] if isinstance(p:=trace.keys[ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None: info = f"{sym_infer(p.estimates.ops, ei.arg['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei.arg['var_vals'])/t:4.1f}"+ \ - f"|{sym_infer(p.estimates.lds,ei.arg['var_vals'])/t:.1f} GB/s\n{ei.arg['metadata']}" + f"|{sym_infer(p.estimates.lds,ei.arg['var_vals'])/t:.1f} GB/s\n{[str(m) for m in ei.arg['metadata']]}" key = ei.key elif isinstance(e.name, TracingKey): name = e.name.display_name