sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove old stale metadata codes [pr] (#16936)

Browse Source
This commit is contained in:
chenyu
2026-07-08 17:11:15 -04:00
committed by GitHub
parent 52b84adc2a
commit 8085bd57ec
6 changed files with 39 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_graph.py
+27 -27
View File
@@ -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]
tinygrad/engine/jit.py
+2 -2
View File
@@ -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))}")
tinygrad/engine/realize.py
+3 -4
View File
@@ -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, ...]] = {}
tinygrad/schedule/__init__.py
+1 -1
View File
@@ -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)
tinygrad/uop/ops.py
+6 -7
View File
@@ -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 ********
tinygrad/viz/serve.py
-1
View File
@@ -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