IQ.Pilot Prebuilt Release @ ab07000

tinygrad_repo/tinygrad/engine/jit.py

@@ -0,0 +1,301 @@
+from typing import TypeVar, Generic, Callable, Any
+import functools, collections
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ
+from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
+from tinygrad.dtype import DType, dtypes
+from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
+from tinygrad.engine.realize import capturing, Estimates, compile_linear, run_linear, graph_cache, estimate_uop, get_runtime
+from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
+from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
+from tinygrad.nn.state import get_parameters
+from tinygrad.schedule.rangeify import mop_cleanup
+from dataclasses import dataclass
+
+def prune_linear(linear:UOp, needed:set[UOp]) -> tuple[UOp, UOp]:
+  kept, onetime = [], []
+  for si in linear.src:
+    si_bufs = {b for src in si.src[1:] for b in _collect_bufs(src)}
+    if not si_bufs.isdisjoint(needed):
+      kept.append(si)
+      needed |= si_bufs
+    else: onetime.append(si)
+  return linear.replace(src=tuple(kept)), linear.replace(src=tuple(onetime))
+
+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))
+
+def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
+  new_src: list[UOp] = []
+  current_batch: list[UOp] = []
+  current_batch_devs: list[Compiled] = []
+
+  def flush_batch():
+    nonlocal current_batch, current_batch_devs, max_batch_size, new_src
+    if len(current_batch) <= 1 and not getenv("GRAPH_ONE_KERNEL"): new_src.extend(current_batch)
+    else:
+      new_src.append(create_graph_call(current_batch))
+      max_batch_size *= 2
+      if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels")
+    current_batch, current_batch_devs = [], []
+
+  for si in linear.src:
+    if si.src[0].op is Ops.SLICE: continue
+
+    devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
+    graph_t = graph_class(devs[0]) if devs[0].graph is not None else None
+
+    can_graph = graph_t is not None and graph_t.supports_uop(devs, si)
+    can_extend = can_graph and graph_t is not None and (not current_batch_devs or graph_t.supports_uop(current_batch_devs, si)) \
+      and (max_batch_size == 0 or len(current_batch) < max_batch_size)
+    if not can_extend and current_batch: flush_batch()
+
+    # append this si and update devs
+    (current_batch if can_graph else new_src).append(si)
+    current_batch_devs = dedup(current_batch_devs + devs) if can_graph else []
+  if current_batch: flush_batch()
+  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.arg, u.dtype), u, metadata=()),)))
+  return new
+
+@track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}")
+def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
+  if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")
+
+  # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
+  linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
+  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
+  linear = memory_plan_rewrite(linear, held_bufs)
+  if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
+  if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
+  return linear
+
+class GraphException(Exception): pass
+class JitError(Exception): pass
+
+def _check_no_non_tensor_return(ret):
+  if ret is None or isinstance(ret, Tensor): return
+  if isinstance(ret, (tuple, list, dict)):
+    for item in (ret.values() if isinstance(ret, dict) else ret): _check_no_non_tensor_return(item)
+    return
+  raise JitError(f"JIT return contains non-Tensor value of type {type(ret).__name__}")
+
+def graph_class(dev): return dev.graph.func if isinstance(dev.graph, functools.partial) else dev.graph
+
+class GraphRunner:
+  def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
+    self.linear = linear.src[0]
+    self.calls: list[tuple[int, UOp, list[Buffer], dict[str, int]]] = []
+    self.runtimes: list[Any|None] = []
+    self.uop_replace: list[list[tuple[int, int]]] = []
+    for call in self.linear.src:
+      replace = [(p, b.arg.slot) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM]
+      for dev_idx, (bufs, device_vars) in enumerate(unwrap_multi(call, resolve_params(call, input_uops))):
+        self.calls.append((dev_idx, call.src[0], [b.ensure_allocated() for b in bufs], device_vars))
+        self.runtimes.append(get_runtime(bufs[0].device, call.src[0]) if call.src[0].op is Ops.PROGRAM else None)
+        self.uop_replace.append(replace)
+
+    self.var_vals_replace:dict[int, list[tuple[int, int]]] = {}
+    self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {}
+    self.launch_dims_base:dict[int, tuple[tuple[int|float, ...], tuple[int, ...]]] = {}
+
+    def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim)
+
+    crs = [(j, self.calls[j][1].arg, self.calls[j][3]) for j in range(len(self.calls)) if self.calls[j][1].op is Ops.PROGRAM]
+    self.vars = sorted({v.expr for _,p,dv in crs for v in p.vars if v.expr not in dv | p.runtimevars})
+    self.symbolic_dims = dedup(tuple(d) for _,p,_ in crs for d in (p.local_size, p.global_size) if d and is_sym_dim(d))
+
+    def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None
+
+    for j,p,dv in crs:
+      if (replace:=[(i, self.vars.index(v.expr)) for i, v in enumerate(p.vars) if v.expr not in dv | p.runtimevars]):
+        self.var_vals_replace[j] = replace
+      global_dim_idx, local_dim_idx = find_symbolic_dim(p.global_size), find_symbolic_dim(p.local_size)
+      if global_dim_idx is not None or local_dim_idx is not None:
+        self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx)
+        assert p.local_size is not None
+        self.launch_dims_base[j] = (tuple(p.global_size), tuple(p.local_size))
+
+    estimates = sum((estimate_uop(call) for call in self.linear.src), Estimates())
+
+    # used in MultiGraphRunner. tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
+    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+
+    self.device, self.estimates = self.calls[0][2][0].device.split(":")[0], estimates.simplify()
+
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: raise NotImplementedError("override this")
+
+  def updated_vars(self, var_vals: dict[str, int]):
+    vals = [var_vals[v] for v in self.vars]
+    for j, vidxs in self.var_vals_replace.items():
+      for i, v in vidxs: yield j, i, vals[v]
+
+  def updated_launch_dims(self, var_vals: dict[str, int]):
+    dims = [tuple(sym_infer(s, var_vals) for s in dim) for dim in self.symbolic_dims]
+    for j, (gl, lc) in self.launch_dims_replace.items():
+      yield j, (dims[gl] if gl is not None else self.launch_dims_base[j][0]), (dims[lc] if lc is not None else self.launch_dims_base[j][1])
+
+  def _access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
+    wait_nodes = []
+    for i,buf in enumerate(bufs):
+      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
+      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
+      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
+    for i,buf in enumerate(bufs):
+      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
+      if i in write:
+        for dmap in [self.w_dependency_map, self.r_dependency_map]:
+          kept = []
+          for st,en,dep in dmap[key]:
+            if st < min(s, en): kept.append((st, min(s, en), dep))
+            if max(e, st) < en: kept.append((max(e, st), en, dep))
+          dmap[key] = kept
+        self.w_dependency_map[key].append((s, e, new_dependency))
+      else: self.r_dependency_map[key].append((s, e, new_dependency))
+    return list({id(x):x for x in wait_nodes}.values())
+
+  @staticmethod
+  def _all_devs(batch_devs:list[Compiled], new_call:UOp) -> list[Compiled]:
+    return dedup(batch_devs + [Device[x] for b in get_call_arg_uops(new_call)
+                 for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
+
+  @staticmethod
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
+    return new_call.src[0].op is Ops.PROGRAM and len(GraphRunner._all_devs(batch_devs, new_call)) == 1
+
+# a marker for your graph supporting multiple devices of the same type
+class MultiGraphRunner(GraphRunner):
+  @staticmethod
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
+    # Devices must be the same type
+    return new_call.src[0].op in (Ops.PROGRAM, Ops.COPY) and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1
+
+ReturnType = TypeVar('ReturnType')
+@dataclass
+class CapturedJit(Generic[ReturnType]):
+  ret: Any  # includes the Tensors or any other returned object
+  linear: UOp
+  expected_names: list[int|str]
+  expected_input_info: list[tuple[UOp, tuple[Variable, ...], DType, str]]  # (view, variables, dtype, device) per input
+
+  def __reduce__(self): return self.__class__, (self.ret, self.linear, self.expected_names, self.expected_input_info)
+
+  @functools.cached_property
+  def _written_uops(self) -> set[UOp]:
+    out: set[UOp] = set()
+    for call in self.linear.toposort():
+      if call.op is not Ops.CALL: continue
+      arg_uops = get_call_arg_uops(call)
+      outs, ins = get_call_outs_ins(call)
+      out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)}
+    return out
+
+  def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
+    concrete = tuple(_copy_input(u) if u in self._written_uops else u for u in input_uops)
+    if DEBUG >= 1 and len(self.linear.src) >= 10: print(f"jit execs {len(self.linear.src)} calls")
+    run_linear(self.linear, var_vals, input_uops=concrete, jit=True)
+    return self.ret
+
+  def free_intermediates(self):
+    # drop graph runners
+    for call in self.linear.src:
+      if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "graph": graph_cache.pop(call.src[0], None)
+    for u in self._written_uops:
+      if (buf:=buffers.get(u)) is None: continue
+      for b in (buf.bufs if isinstance(buf, MultiBuffer) else (buf,)):
+        if b.is_initialized(): b.deallocate()
+        if (base:=b._base) is not None and base.allocated_views == 0 and base.is_allocated(): base.deallocate()
+
+def _prepare_jit_inputs(args, kwargs):
+  input_tensors: list[tuple[int|str, Tensor]] = [(name,t) for name,t in list(enumerate(args))+sorted(kwargs.items()) if t.__class__ is Tensor]
+  names, tensors = [name for name,_ in input_tensors], [t for _,t in input_tensors]
+  # extract tensors from containers (shallow, not recursive to avoid grabbing model weights)
+  for x in args + tuple(kwargs.values()):
+    it = x if isinstance(x, (tuple,list)) else x.values() if isinstance(x, dict) else []
+    tensors += [t for t in it if t.__class__ is Tensor and not any(t is y for y in tensors)]
+  def get_input_uops() -> list[UOp]: return flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors])
+  # TODO: drop the CONST branch once all CONST are deviceless
+  if any(u.device is None or u.base.op is Ops.CONST for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()")
+  if len(unrealized_tensors := [x for x in tensors if not x.uop.is_realized]): Tensor.realize(*unrealized_tensors)
+  input_uops = get_input_uops()
+  # collect buffer UOps (including MultiBuffer)
+  input_buf_uops: list[UOp] = [u.base for u in input_uops if u.base.realized is not None]
+  if len(set(input_buf_uops)) != len(input_buf_uops): raise JitError("duplicate inputs to JIT")
+  inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
+  _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
+  var_vals = {k.expr:v for k,v in _var_vals.items()}
+  expected_input_info = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in inputs]
+  return input_buf_uops, var_vals, names, expected_input_info
+
+class TinyJit(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType]|None, captured:CapturedJit|None=None, prune=False):
+    assert fxn or captured, "need either a function or a CapturedJit"
+    self.fxn = fxn
+    self.captured: CapturedJit|None = captured
+    self.cnt: int = 2 if self.fxn is None else 0
+    self.prune = prune
+
+  def add_linear(self, linear:UOp, var_vals:dict[str, int]): self._linears.append(linear)
+
+  def reset(self):
+    assert self.fxn is not None, "can't reset without function"
+    self.cnt = 0
+    self.captured = None
+
+  def __reduce__(self):
+    assert self.captured is not None, "can't pickle an uncaptured JIT"
+    return self.__class__, (None, self.captured)
+
+  def __get__(self, obj, objtype): return functools.partial(self.__call__, obj) # add support for instance methods
+
+  def __call__(self, *args, **kwargs) -> ReturnType:
+    input_buf_uops, var_vals, names, expected_input_info = _prepare_jit_inputs(args, kwargs)
+    if not JIT or self.cnt == 0:
+      # jit ignore
+      assert self.fxn is not None
+      with Context(BEAM=0 if getenv("IGNORE_JIT_FIRST_BEAM") else BEAM.value):
+        ret = self.fxn(*args, **kwargs)
+        if len(params:=get_parameters(ret)): Tensor.realize(*params)
+    elif self.cnt == 1:
+      # jit capture
+      assert self.fxn is not None
+      if capturing: raise RuntimeError(f"having TinyJit inside another TinyJit is not supported {len(capturing)=} {capturing=}")
+      self._linears: list[UOp] = []
+      capturing.append(self)
+      try:
+        ret = self.fxn(*args, **kwargs)
+        if len(params:=get_parameters(ret)): Tensor.realize(*params)
+      finally: capturing.clear()
+      if not len(self._linears): raise JitError("didn't JIT anything!")
+      _check_no_non_tensor_return(ret)
+      if DEBUG >= 1: print(f"JIT captured {len(self._linears)} linears with {len(input_buf_uops)} inputs")
+
+      # combine all captured linears into one, memory plan, and graph split
+      big_linear = UOp(Ops.LINEAR, src=tuple(flatten([l.src for l in self._linears])))
+      del self._linears
+
+      if self.prune:
+        big_linear, onetime_linear = prune_linear(big_linear, set(input_buf_uops))
+        if DEBUG >= 1: print(f"pruned from {len(big_linear.src) + len(onetime_linear.src)} -> {len(big_linear.src)} kernels")
+        run_linear(onetime_linear, var_vals)
+
+      held_bufs = set(buffers) | {t.uop.buf_uop for t in get_parameters(ret) if t.uop.buf_uop.op is Ops.BUFFER}
+      linear = jit_lower(big_linear, held_bufs, input_buf_uops)
+      self.captured = CapturedJit(ret, linear, names, expected_input_info)
+      ret = self.captured(input_buf_uops, var_vals)
+    elif self.cnt >= 2:
+      # jit exec
+      assert self.captured is not None
+      if self.captured.expected_names != names: raise JitError(f"args mismatch in JIT: {self.captured.expected_names=} != {names}")
+      if self.captured.expected_input_info != expected_input_info:
+        raise JitError(f"args mismatch in JIT: {self.captured.expected_input_info=} != {expected_input_info=}")
+      ret = self.captured(input_buf_uops, var_vals)
+
+    self.cnt += 1
+    return ret

tinygrad_repo/tinygrad/engine/realize.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+from typing import cast, Iterator, Any
+import time, random, itertools, math, contextlib, weakref
+from dataclasses import dataclass, replace, field
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, TRACEMETA, prod, flatten, Context, getenv
+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
+from tinygrad.device import Device, Buffer, MultiBuffer
+from tinygrad.renderer import Estimates
+from tinygrad.codegen import to_program
+from tinygrad.codegen.opt.postrange import bufs_from_ast
+
+# **************** Helpers ****************
+
+def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if s.op is not Ops.BIND)
+
+def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
+  ast = call.src[0]
+  if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
+  if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,)
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
+  return (), ()
+
+def get_call_name(call:UOp, bufs:list[Buffer], var_vals:dict[str, int]|None=None) -> str:
+  def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
+
+  ast, arg_uops = call.src[0], get_call_arg_uops(call)
+  if ast.op is Ops.PROGRAM: return ast.arg.name
+  if ast.op is Ops.SLICE:
+    offset = ast.src[1].arg * arg_uops[1].dtype.itemsize
+    return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
+  if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {bufs[0].device[:7]:>7s} <- {bufs[1].device[:7]:7s}", "yellow")
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
+  raise NotImplementedError("get_call_name is not implemented")
+
+# **************** Stat ****************
+
+def estimate_uop(call:UOp) -> Estimates:
+  ast = call.src[0]
+  if ast.op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
+  if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
+    nbytes = prod(call.src[1].shape) * call.src[1].dtype.itemsize
+    return Estimates(lds=nbytes, mem=nbytes)
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return get_graph_runtime(ast).estimates
+  return Estimates()
+
+first_run_cache:set[bytes] = set()
+@contextlib.contextmanager
+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,
+      "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()
+  yield et
+  if not ctx.update_stats: return
+
+  if DEBUG >= 2 and et[0] is None:
+    Device[device].synchronize()
+    et[0] = time.perf_counter() - st
+
+  estimates = estimate_uop(call)
+  GlobalCounters.kernel_count += 1
+  GlobalCounters.global_ops += (op_est:=sym_infer(estimates.ops, var_vals))
+  GlobalCounters.global_mem += (mem_est:=sym_infer(estimates.mem, var_vals))
+  if et[0] is not None: GlobalCounters.time_sum_s += et[0]
+  if DEBUG >= 2:
+    display_name = get_call_name(call, bufs, var_vals)
+    lds_est = sym_infer(estimates.lds, var_vals)
+    header_color = 'magenta' if ctx.jit else ('green' if call.src[0].key not in first_run_cache else None)
+    ptm = colored(time_to_str(et[0], w=9), "yellow" if et[0] > 0.01 else None) if et[0] is not None else ""
+    flops, membw, ldsbw = op_est/(et[0] or 1e-20), mem_est/(et[0] or 1e-20), lds_est/(et[0] 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')
+    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 ''}")
+    first_run_cache.add(call.src[0].key)
+
+local_size_cache: dict[bytes, tuple[int, ...]] = {}
+def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
+  device = prg.src[1].arg
+  if prg.arg.local_size is not None or not Device[device].renderer.has_local or not all_int(prg.arg.global_size): return None
+
+  if (local_size:=local_size_cache.get(prg.key)) is None:
+    bufs = [UOp.from_buffer(b.allocate()) for b in bufs_from_ast(prg.src[0], device)]
+    def try_exec(local_size):
+      try:
+        new_gs = tuple(g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size))
+        return time_call(prg.replace(arg=replace(prg.arg, global_size=new_gs, local_size=tuple(local_size))).call(*bufs))
+      except Exception: return float('inf')
+
+    MAX_WORKGROUP = 1024
+    local_dims = [[x for x in set([sz, 1, 2, 4, 8, 16, 32, 64, 128, 256, MAX_WORKGROUP]) if x<=sz] for sz in prg.arg.global_size]
+    local_sizes = [list(x) for x in itertools.product(*local_dims) if prod(x) <= MAX_WORKGROUP] * 2  # try each valid size twice
+    best_time, best = min([(try_exec(ls), ls) for ls in random.sample(local_sizes, len(local_sizes))])
+    assert not math.isinf(best_time), "all optimize_local_size exec failed"
+    local_size = local_size_cache[prg.key] = tuple(best)
+
+  new_global = tuple(g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size))
+  return call.replace(src=(prg.replace(arg=replace(prg.arg, global_size=new_global, local_size=local_size)), *call.src[1:]))
+
+# **************** runtime cache ****************
+
+runtime_cache: dict[tuple[bytes, str], Any] = {}
+def get_runtime(device:str, ast:UOp, cache=True):
+  assert ast.op is Ops.PROGRAM and isinstance(ast.arg, ProgramInfo), "get_runtime should only be called with a PROGRAM ast"
+  if (runtime:=runtime_cache.get(key:=(ast.key, device))) is None:
+    runtime = Device[device].runtime(ast.arg.function_name, ast.src[4].arg, *ast.arg.aux, runtimevars=ast.arg.runtimevars, prg=ast)
+    if cache: runtime_cache[key] = runtime
+  return runtime
+
+graph_cache:weakref.WeakKeyDictionary[UOp, Any] = weakref.WeakKeyDictionary()
+def get_graph_runtime(ast:UOp, input_uops:tuple[UOp, ...]|None=None):
+  assert ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph", "get_graph_runtime should only be called with a graph ast"
+  if (runtime:=graph_cache.get(ast)) is None and input_uops is not None:
+    graph_cache[ast] = runtime = Device[ast.device if isinstance(ast.device, str) else ast.device[0]].graph(ast, input_uops=input_uops)
+  return runtime
+
+# **************** run linear ****************
+
+capturing: list = []  # put classes with an add_linear method in here
+
+@dataclass
+class ExecContext:
+  var_vals: dict[str, int] = field(default_factory=dict)
+  input_uops: tuple[UOp, ...] = ()
+  update_stats: bool = True
+  jit: bool = False
+  wait: bool = False
+  timeout: int|None = None
+  cache: bool = True
+
+def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
+  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
+  return inputs[b.arg.slot] if b.op is Ops.PARAM else b
+def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
+
+def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], dict[str, int]]]:
+  bufs = [b.buffer for b in resolved]
+  if not any(isinstance(b, MultiBuffer) for b in bufs): yield cast(list[Buffer], bufs), {}
+  else:
+    dnum = next((x.expr for x in call.src[0].variables() if x.expr == '_device_num'), None)
+    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {dnum: j} if dnum else {}
+
+def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  resolved = resolve_params(call, ctx.input_uops)
+  bufs = [cast(Buffer, b.buffer) for b in resolved]
+  bv = bufs[1].view(resolved[0].arg, ast.dtype, ast.src[1].arg*bufs[1].dtype.itemsize)
+  with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
+  return None
+
+def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
+    dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
+    with track_stats(ctx, call, dest.device, [dest, src], ctx.var_vals):
+      if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
+        dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev) # type:ignore[attr-defined]
+      elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
+           and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
+        dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
+      elif src.device.startswith(("DISK", "TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
+        src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
+      else: dest.copyin(src.as_memoryview(allow_zero_copy=True))
+  return None
+
+def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  et = None
+  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
+    var_vals = {**ctx.var_vals, **device_vars}
+    prg_bufs = [bufs[i].ensure_allocated() for i in ast.arg.globals]
+    rt = get_runtime(device:=bufs[0].device, ast, cache=ctx.cache)
+    global_size, local_size = ast.arg.launch_dims(var_vals)
+    with track_stats(ctx, call, device, prg_bufs, var_vals) as tm:
+      et = tm[0] = rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
+                       wait=ctx.wait, timeout=ctx.timeout)
+  return et
+
+def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  import numpy as np
+  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
+    bufs, dev_bufs = bufs[:len(bufs)//2], bufs[len(bufs)//2:]
+    var_vals = {**ctx.var_vals, **device_vars}
+    cpu_rt = get_runtime("CPU", prg:=to_program(ast.src[0], Device["CPU"].renderer))
+    global_size, local_size = prg.arg.launch_dims(var_vals)
+    cpu_rt(*[bufs[i].ensure_allocated()._buf for i in prg.arg.globals], global_size=global_size, local_size=local_size, vals=prg.arg.vals(var_vals))
+    for i in prg.arg.outs: np.testing.assert_allclose(dev_bufs[i].ensure_allocated().numpy(), bufs[i].numpy(), rtol=1e-3, atol=1e-3)
+  return None
+
+def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)]
+  shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
+  with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
+    bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
+  return None
+
+def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  rt = get_graph_runtime(ast, ctx.input_uops)
+  with track_stats(ctx, call, rt.device, [], ctx.var_vals) as t: t[0] = rt(ctx.input_uops, ctx.var_vals, wait=ctx.wait) # type: ignore[call-arg]
+  return t[0]
+
+# flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
+pm_flatten_linear = PatternMatcher([
+  (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="lin"),
+   lambda lin: lin.replace(src=tuple(flatten(c.src if c.op is Ops.LINEAR else (c,) for c in lin.src)))),
+])
+
+def _validate(call:UOp, sink:UOp) -> UOp:
+  params = get_call_arg_uops(call)
+  shadows = tuple(UOp.new_buffer(("CPU",)*len(p.device) if isinstance(p.device, tuple) else "CPU", prod(p.max_shape), p.dtype.base) for p in params)
+  copies = tuple(p.copy_to_device(s.device).call(s, p) for s, p in zip(shadows, params))
+  return UOp(Ops.LINEAR, src=copies + (call, UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(sink,), arg="validate").call(*shadows, *params)))
+pm_validate = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), name="call", allow_any_len=True), _validate)]) + pm_flatten_linear
+
+# ctx is beam value
+pm_beam = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), name="call", allow_any_len=True),
+   lambda ctx,call,sink: call.replace(src=(sink.replace(arg=replace(sink.arg, beam=ctx)), *call.src[1:])) if sink.arg.beam == 0 else None),
+])
+
+pm_compile = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), lambda call,ast:
+    call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))),
+])
+
+pm_optimize_local_size = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size),
+])
+
+pm_exec = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="graph", name="ast"),), name="call", allow_any_len=True), exec_graph),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
+])
+
+def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp:
+  if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
+  if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
+  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
+  if getenv("HCQ2"):
+    from extra.hcq2.hcq2 import hcq_schedule
+    linear = hcq_schedule(linear)
+  return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
+
+def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), update_stats=True, jit=False, wait=False):
+  if not jit: linear = compile_linear(linear, validate=VALIDATE_WITH_CPU)
+  ctx = ExecContext(var_vals or {}, input_uops, update_stats, jit, wait or DEBUG>=2)
+  for call in linear.src: pm_exec.rewrite(call, ctx)
+
+def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> float:
+  if clear_l2:
+    if hasattr(dev:=Device[call.src[0].src[1].arg], 'invalidate_caches'): dev.invalidate_caches()
+    else:
+      from tinygrad.tensor import Tensor
+      with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
+  call = compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0).src[0]
+  return cast(float, pm_exec.rewrite(call, ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)))