From 8085bd57ece792dd04c09527fd6568e9ee338411 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 8 Jul 2026 17:11:15 -0400 Subject: [PATCH] remove old stale metadata codes [pr] (#16936) --- test/backend/test_graph.py | 54 +++++++++++++++++------------------ tinygrad/engine/jit.py | 4 +-- tinygrad/engine/realize.py | 7 ++--- tinygrad/schedule/__init__.py | 2 +- tinygrad/uop/ops.py | 13 ++++----- tinygrad/viz/serve.py | 1 - 6 files changed, 39 insertions(+), 42 deletions(-) diff --git a/test/backend/test_graph.py b/test/backend/test_graph.py index 554c3db646..e9c5b915e5 100644 --- a/test/backend/test_graph.py +++ b/test/backend/test_graph.py @@ -73,8 +73,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()), + get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)), + get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c)), ] zero_bufs([b[0]]) @@ -92,8 +92,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()), + get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)), + get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c)), ] zero_bufs([b[0], b[1]]) @@ -111,8 +111,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c), metadata=()), + get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)), + get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c)), ] zero_bufs([b[0], b[1]]) @@ -131,8 +131,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()), - UOp(Ops.COPY).call(get_buf_uop(b[3],c), get_buf_uop(b[0],c), metadata=()), + get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)), + UOp(Ops.COPY).call(get_buf_uop(b[3],c), get_buf_uop(b[0],c)), ] zero_bufs([b[0], b[3]]) @@ -151,8 +151,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - UOp(Ops.COPY).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()), + UOp(Ops.COPY).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c)), + get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)), ] zero_bufs([b[1], b[3]]) @@ -169,9 +169,9 @@ class TestGraph(unittest.TestCase): b = [make_buffer(d0, fill=True) for _ in range(8)] c: dict[Buffer,UOp] = {} - calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=())] - calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c), metadata=())] - calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c), metadata=())] + calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c))] + calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c))] + calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c))] out = [b[3], b[4], b[5]] zero_bufs(out) @@ -194,8 +194,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c), metadata=()), + UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c)), + get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c)), ] out = [b1[0], b0[2]] @@ -219,8 +219,8 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - UOp(Ops.COPY).call(get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()), + UOp(Ops.COPY).call(get_buf_uop(b0,c), get_buf_uop(b2,c)), + get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c)), ] zero_bufs([b0]) @@ -245,9 +245,9 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()), - UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()), + UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c)), + UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c)), + get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c)), ] zero_bufs([base, out]) @@ -272,9 +272,9 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - UOp(Ops.COPY).call(get_buf_uop(copy_dst,c), get_buf_uop(base,c), metadata=()), - UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c), metadata=()), + UOp(Ops.COPY).call(get_buf_uop(copy_dst,c), get_buf_uop(base,c)), + UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c)), + get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c)), ] zero_bufs([copy_dst, base]) @@ -299,10 +299,10 @@ class TestGraph(unittest.TestCase): c: dict[Buffer,UOp] = {} calls = [ - UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()), - UOp(Ops.COPY).call(get_buf_uop(v_mid,c), get_buf_uop(copy_src_mid,c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c), metadata=()), - get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()), + UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c)), + UOp(Ops.COPY).call(get_buf_uop(v_mid,c), get_buf_uop(copy_src_mid,c)), + get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c)), + get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c)), ] outs = [base, out1, out2] diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index fca73d24dc..45c6f3ca92 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -27,7 +27,7 @@ def create_graph_call(batch:list[UOp]) -> UOp: # all external inputs are PARAMs input_list = dedup(u for si in batch for b in si.src[1:] for u in b.toposort() if u.op is Ops.PARAM) cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(UOp(Ops.LINEAR, src=tuple(batch)),), arg="graph") - return cf.call(*input_list, metadata=tuple(m for si in batch for m in si.arg.metadata)) + return cf.call(*input_list) def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp: new_src: list[UOp] = [] @@ -61,7 +61,7 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp: return linear.replace(src=tuple(new_src)) def _copy_input(u:UOp) -> UOp: - run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u, metadata=()),))) + run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u),))) return new @track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}") diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 6346774dda..5420333086 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import cast, Iterator, Any, Sequence import time, random, itertools, math, contextlib, weakref, array from dataclasses import dataclass, replace, field -from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, TRACEMETA, prod, flatten, Context, getenv, to_tuple +from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events from tinygrad.dtype import dtypes from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo @@ -54,7 +54,7 @@ first_run_cache:set[bytes] = set() def track_stats(ctx:ExecContext, call:UOp, device:str, bufs:list[Buffer], var_vals:dict[str, int]): if PROFILE: outputs, inputs = get_call_outs_ins(call) - cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"metadata": call.arg.metadata, "var_vals": var_vals, + cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": var_vals, "bufs": [b.trace_num for b in bufs], "name": get_call_name(call, bufs, var_vals), "outputs": outputs, "inputs": inputs})) et: list[float|None] = [None] if DEBUG >= 2: st = time.perf_counter() @@ -81,8 +81,7 @@ def track_stats(ctx:ExecContext, call:UOp, device:str, bufs:list[Buffer], var_va colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green') print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+ f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+ - ("" if et[0] 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 call.arg.metadata] if call.arg.metadata else ''}") + ("" if et[0] is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")) first_run_cache.add(call.src[0].key) local_size_cache: dict[bytes, tuple[int, ...]] = {} diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 47b10c3aee..bf657a1148 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -72,7 +72,7 @@ def create_schedule(sched_sink:UOp) -> UOp: k = rk.src[0] if rk.op is Ops.END else rk assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}" buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND) - linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata)) + linearized.append(k.src[0].call(*buf_uops)) for x in children.get(rk, []): in_degree[x] -= 1 if in_degree[x] == 0: queue.append(x) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index f0b254b370..93bf4d2c02 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1036,14 +1036,14 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE) _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION} - def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), + def call(self, *srcs:UOp, grad_fxn:Callable|None=None, name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp: assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}" if self.op in UOp._OPAQUE_CALL_BODIES: - return UOp(Ops.CALL, dtypes.void, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward, aux)) + return UOp(Ops.CALL, dtypes.void, (self,)+srcs, CallInfo(grad_fxn, name, precompile, precompile_backward, aux)) # value-producing bodies are always wrapped in TUPLE so FUNCTION dtype is always void body = self if self.op is Ops.TUPLE else UOp.maketuple(self) - return UOp(Ops.FUNCTION, dtypes.void, (body,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward, aux)) + return UOp(Ops.FUNCTION, dtypes.void, (body,)+srcs, CallInfo(grad_fxn, name, precompile, precompile_backward, aux)) def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]: contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs) placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)] @@ -1112,16 +1112,15 @@ class ProgramInfo: @dataclass(frozen=True) class CallInfo: grad_fxn: Callable|None = None - metadata: tuple[Metadata, ...] = () name: str|None = None precompile: bool = False precompile_backward: bool = False aux: Any = None - # grad_fxn can't be pickled, but metadata can - def __reduce__(self): return (CallInfo, (None, self.metadata, self.name, self.precompile, self.precompile_backward, self.aux)) + # grad_fxn can't be pickled + def __reduce__(self): return (CallInfo, (None, self.name, self.precompile, self.precompile_backward, self.aux)) def __repr__(self): gf = id(self.grad_fxn) if self.grad_fxn else None - return f"CallInfo({gf}, {self.metadata}, {repr(self.name)}, {self.precompile}, {self.precompile_backward})" + return f"CallInfo({gf}, {repr(self.name)}, {self.precompile}, {self.precompile_backward})" # ******** ops in python ******** diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 18ccddfaa1..5a8536b7bb 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -239,7 +239,6 @@ def timeline_layout(data:VizData, dev_events:list[tuple[int, int, float, DevEven if (ki:=data.ctxs[ref].get("ki")) is not None and ki.estimates is not None and ei is not None: fmt["FLOPS"] = int(sym_infer(ki.estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6)) fmt["B/s mem"], fmt["B/s lds"] = int(sym_infer(ki.estimates.mem, var_vals)/t), int(sym_infer(ki.estimates.lds, var_vals)/t) - if ei.arg["metadata"]: fmt["metadata"] = ",".join([str(m) for m in ei.arg['metadata']+["batched" if isinstance(e,ProfileGraphEntry) else ""]]) key = ei.key elif isinstance(e.name, TracingKey): name = e.name.display_name