IQ.Pilot Prebuilt Release @ 7e87bc7
This commit is contained in:
IQ.Lvbs CI [bot]
296
artifacts/package_runtime/tinygrad/engine/jit.py
Normal file
296
artifacts/package_runtime/tinygrad/engine/jit.py
Normal file
@@ -0,0 +1,296 @@
|
||||
from typing import TypeVar, Generic, Callable, Any, overload
|
||||
import functools
|
||||
from tinygrad.tensor import Tensor, all_tensors
|
||||
from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc
|
||||
from tinygrad.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker
|
||||
from tinygrad.dtype import DType
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite
|
||||
from tinygrad.renderer import Estimates
|
||||
from tinygrad.engine.realize import capturing, compile_linear, link_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.uop.movement 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, src=(UOp(Ops.LINEAR, src=tuple(batch)),), arg="graph")
|
||||
return cf.call(*input_list)
|
||||
|
||||
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.max_numel(), u.dtype), u),)))
|
||||
return new
|
||||
|
||||
@rewrite_group(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 = memory_plan_rewrite(linear, held_bufs)
|
||||
linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
|
||||
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
|
||||
self.deps = DepsTracker()
|
||||
|
||||
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):
|
||||
return self.deps.access_resources(bufs, write, new_dependency)
|
||||
|
||||
@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
|
||||
|
||||
@functools.cached_property
|
||||
def linear(self) -> UOp: return link_linear(self._linear)
|
||||
|
||||
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[0]] if t.uop.op is Ops.UNSHARD else [t.uop] for t in tensors])
|
||||
if any(u.is_virtual 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, u.base.dtype)}, 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
|
||||
|
||||
@disable_gc()
|
||||
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)
|
||||
|
||||
# hold all buffers reachable from live Tensors (e.g. lazy .grad created during capture), the memory planner can't suballocate those
|
||||
held_bufs = set(buffers) | {u for tref in list(all_tensors) if (t:=tref()) is not None for u in t.uop.toposort() if u.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
|
||||
|
||||
# overload signatures support both @TinyJit and @TinyJit(prune=True) syntax
|
||||
@overload
|
||||
def TinyJit(fxn:Callable[..., ReturnType], *, prune:bool=False) -> _TinyJit[ReturnType]: ...
|
||||
@overload
|
||||
def TinyJit(fxn:None=None, *, prune:bool=False) -> Callable[[Callable[..., ReturnType]], _TinyJit[ReturnType]]: ...
|
||||
def TinyJit(fxn=None, **kwargs): return (lambda f: _TinyJit(f, **kwargs)) if fxn is None else _TinyJit(fxn, **kwargs)
|
||||
306
artifacts/package_runtime/tinygrad/engine/realize.py
Normal file
306
artifacts/package_runtime/tinygrad/engine/realize.py
Normal file
@@ -0,0 +1,306 @@
|
||||
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, prod, flatten, Context, getenv, to_tuple
|
||||
from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
|
||||
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
|
||||
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:Sequence[Buffer|UOp], 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 {}))
|
||||
def _dev_str(buf:Buffer|UOp) -> str: return ', '.join(d[:7] for d in to_tuple(buf.device))
|
||||
|
||||
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].val * 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}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):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")
|
||||
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return cast(str, call.arg.name)
|
||||
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
|
||||
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.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), {"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})"))
|
||||
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 = to_tuple(prg.device)[0]
|
||||
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:
|
||||
MAX_WORKGROUP = 1024
|
||||
# divisors only: a non-divisor candidate yields a fractional global size, a malformed
|
||||
# dispatch some GPUs (adreno a630) execute anyway and stall on
|
||||
local_dims = [[x for x in set([sz, 1, 2, 4, 8, 16, 32, 64, 128, 256, MAX_WORKGROUP]) if x<=sz and sz%x==0] for sz in prg.arg.global_size]
|
||||
if getenv("NO_LOCAL_SEARCH"):
|
||||
# deterministic bake without timing dispatches on the GPU (one >timeout kernel
|
||||
# permanently poisons the device error_state); largest divisor within limits
|
||||
rt = get_runtime(device, prg, cache=False)
|
||||
max_local = min(MAX_WORKGROUP, getattr(rt, 'max_threads', MAX_WORKGROUP))
|
||||
if getattr(rt, 'pvtmem', 0) > 0: max_local = 1
|
||||
best, budget = [], max_local
|
||||
for cands in local_dims:
|
||||
pick = max(x for x in cands if x <= budget)
|
||||
best.append(pick); budget //= pick
|
||||
local_size = local_size_cache[prg.key] = tuple(best)
|
||||
else:
|
||||
# reuse one loaded runtime across candidates, only launch dims vary
|
||||
bufs, runtime = [b.allocate() for b in bufs_from_ast(prg.src[0], device)], get_runtime(device, prg, cache=False)
|
||||
def try_exec(local_size):
|
||||
try:
|
||||
new_gs = tuple(g//l for g,l in zip(prg.arg.global_size, local_size))
|
||||
return runtime(*[bufs[i].get_buf(device) for i in prg.arg.globals], global_size=new_gs, local_size=(*local_size,),
|
||||
vals=prg.arg.vals({}), wait=True)
|
||||
except Exception: return float('inf')
|
||||
|
||||
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):
|
||||
if (runtime:=runtime_cache.get(key:=(ast.key, device))) is None:
|
||||
runtime = Device[device].runtime(ast.to_elf())
|
||||
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:]))
|
||||
if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
|
||||
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:
|
||||
# the DEVICE axis is bound per device at launch: it's a RANGE in the AST and the _device_num variable after codegen
|
||||
has_dnum = any((x.op is Ops.RANGE and x.arg[-1] is AxisType.DEVICE) or (x.op is Ops.PARAM and x.arg.name == '_device_num')
|
||||
for x in call.src[0].toposort())
|
||||
for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_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].max_numel(), ast.dtype, ast.src[1].val*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)
|
||||
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 hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
|
||||
else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
|
||||
return None
|
||||
|
||||
def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
|
||||
et = None
|
||||
for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), 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, 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.val 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)
|
||||
return t[0]
|
||||
|
||||
def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
|
||||
if (inputs:=call.arg.aux.inputs) is not None:
|
||||
bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
|
||||
table = call.src[1+inputs].buffer
|
||||
for j,dev in enumerate(call.arg.aux.device):
|
||||
addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
|
||||
mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
|
||||
wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
|
||||
mv[:len(addrs)] = addrs
|
||||
|
||||
exec_kernel(replace(ctx, update_stats=False), call, ast)
|
||||
|
||||
st = time.perf_counter()
|
||||
for d in call.arg.aux.device:
|
||||
with track_stats(ctx, call, d, [], ctx.var_vals):
|
||||
if ctx.wait: cast(Any, Device[d]).synchronize(timeout=ctx.timeout)
|
||||
return time.perf_counter() - st
|
||||
|
||||
# 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) 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, 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="hcq", src=(UPat(Ops.PROGRAM, name="ast"),)),), name="call", allow_any_len=True), exec_hcq),
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
|
||||
])
|
||||
|
||||
if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above
|
||||
|
||||
def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None) -> 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"): linear = hcq_compile(linear, input_uops)
|
||||
return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
|
||||
|
||||
def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
|
||||
|
||||
def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
|
||||
inputs = list(input_uops)
|
||||
if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs))
|
||||
ctx = ExecContext(var_vals or {}, tuple(inputs), 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[1].device], '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)
|
||||
ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
|
||||
linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache)
|
||||
return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)
|
||||
Reference in New Issue
Block a user