tinygrad 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

tinygrad/dtype.py

@@ -0,0 +1,113 @@
+from typing import Final, Optional, ClassVar, Set, Tuple, Dict, Union
+from dataclasses import dataclass
+import numpy as np  # TODO: remove numpy
+import functools
+from tinygrad.helpers import getenv
+
+ConstType = Union[float, int, bool]
+
+@dataclass(frozen=True, order=True)
+class DType:
+  priority: int  # this determines when things get upcasted
+  itemsize: int
+  name: str
+  fmt: Optional[str]
+  count: int
+  def __repr__(self): return f"dtypes.{'_'*(c:=self.count!=1)}{INVERSE_DTYPES_DICT[self.name if not c else self.scalar().name]}{str(self.count)*c}"
+  def vec(self, sz:int):
+    assert sz > 1 and self.count == 1, f"can't vectorize {self} with size {sz}"
+    return DType(self.priority, self.itemsize*sz, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz)
+  def scalar(self): return DTYPES_DICT[self.name[:-len(str(self.count))]] if self.count > 1 else self
+  # TODO: someday this will be removed with the "remove numpy" project
+  @property
+  def np(self) -> Optional[type]: return np.dtype(self.fmt).type if self.fmt is not None else None
+
+# dependent typing?
+@dataclass(frozen=True, repr=False)
+class ImageDType(DType):
+  shape: Tuple[int, ...]   # arbitrary arg for the dtype, used in image for the shape
+  base: DType
+  def scalar(self): return self.base
+  def vec(self, sz:int): return self.base.vec(sz)
+  def __repr__(self): return f"dtypes.{self.name}({self.shape})"
+
+# @dataclass(frozen=True, init=False, repr=False, eq=False)
+class PtrDType(DType):
+  def __init__(self, dt:DType): super().__init__(dt.priority, dt.itemsize, dt.name, dt.fmt, dt.count)
+  def __repr__(self): return f"ptr.{super().__repr__()}"
+  def __hash__(self): return super().__hash__()
+  def __eq__(self, dt): return self.priority==dt.priority and self.itemsize==dt.itemsize and self.name==dt.name and self.count==dt.count
+  def __ne__(self, dt): return not (self == dt)
+
+class dtypes:
+  @staticmethod
+  def is_float(x: DType) -> bool: return x.scalar() in (dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64)
+  @staticmethod # static methds on top, or bool in the type info will refer to dtypes.bool
+  def is_int(x: DType) -> bool: return x.scalar() in (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64) or dtypes.is_unsigned(x)
+  @staticmethod
+  def is_unsigned(x: DType) -> bool: return x.scalar() in (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+  @staticmethod
+  def from_np(x: type) -> DType: return DTYPES_DICT[np.dtype(x).name]
+  @staticmethod  # NOTE: isinstance(True, int) is True in python
+  def from_py(x) -> DType: return dtypes.default_float if isinstance(x, float) else dtypes.bool if isinstance(x, bool) else dtypes.default_int
+  @staticmethod
+  def as_const(val: ConstType, dtype:DType): return int(val) if dtypes.is_int(dtype) else float(val) if dtypes.is_float(dtype) else bool(val)
+  @staticmethod
+  def fields() -> Dict[str, DType]: return DTYPES_DICT
+  bool: Final[DType] = DType(0, 1, "bool", '?', 1)
+  int8: Final[DType] = DType(1, 1, "char", 'b', 1)
+  uint8: Final[DType] = DType(2, 1, "unsigned char", 'B', 1)
+  int16: Final[DType] = DType(3, 2, "short", 'h', 1)
+  uint16: Final[DType] = DType(4, 2, "unsigned short", 'H', 1)
+  int32: Final[DType] = DType(5, 4, "int", 'i', 1)
+  uint32: Final[DType] = DType(6, 4, "unsigned int", 'I', 1)
+  int64: Final[DType] = DType(7, 8, "long", 'l', 1)
+  uint64: Final[DType] = DType(8, 8, "unsigned long", 'L', 1)
+  float16: Final[DType] = DType(9, 2, "half", 'e', 1)
+  # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
+  bfloat16: Final[DType] = DType(10, 2, "__bf16", None, 1)
+  float32: Final[DType] = DType(11, 4, "float", 'f', 1)
+  float64: Final[DType] = DType(12, 8, "double", 'd', 1)
+
+  # dtype aliases
+  half = float16; float = float32; double = float64 # noqa: E702
+  uchar = uint8; ushort = uint16; uint = uint32; ulong = uint64 # noqa: E702
+  char = int8; short = int16; int = int32; long = int64 # noqa: E702
+
+  # NOTE: these are image dtypes
+  @staticmethod
+  def imageh(shp): return ImageDType(100, 2, "imageh", 'e', 1, shape=shp, base=dtypes.float32)
+  @staticmethod
+  def imagef(shp): return ImageDType(100, 4, "imagef", 'f', 1, shape=shp, base=dtypes.float32)
+
+  default_float: ClassVar[DType] = float32
+  default_int: ClassVar[DType] = int32
+
+if (env_default_float := getenv("DEFAULT_FLOAT", "")):
+  dtypes.default_float = getattr(dtypes, env_default_float.lower())
+  assert dtypes.is_float(dtypes.default_float), f"{env_default_float} is not a float dtype"
+
+# https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
+# we don't support weak type and complex type
+promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
+  dtypes.int64: [dtypes.float16, dtypes.bfloat16], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
+  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.float16, dtypes.bfloat16],
+  dtypes.float16: [dtypes.float32], dtypes.bfloat16: [dtypes.float32], dtypes.float32: [dtypes.float64], }
+
+@functools.lru_cache(None)
+def _get_recursive_parents(dtype:DType) -> Set[DType]:
+  return set.union(*[_get_recursive_parents(d) for d in promo_lattice[dtype]], {dtype}) if dtype != dtypes.float64 else {dtypes.float64}
+@functools.lru_cache(None)
+def least_upper_dtype(*ds:DType) -> DType:
+  return min(set.intersection(*[_get_recursive_parents(d) for d in ds])) if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0]
+def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.float32)
+
+# HACK: staticmethods are not callable in 3.8 so we have to compare the class
+DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if not (k.startswith(('__', 'default')) or v.__class__ is staticmethod)}
+INVERSE_DTYPES_DICT = {v.name:k for k,v in DTYPES_DICT.items()}
+
+def sum_acc_dtype(dt:DType):
+  # default acc dtype for sum
+  if dtypes.is_unsigned(dt): return least_upper_dtype(dt, dtypes.uint)
+  if dtypes.is_int(dt) or dt == dtypes.bool: return least_upper_dtype(dt, dtypes.int)
+  return least_upper_dtype(dt, dtypes.float)

tinygrad/engine/graph.py

@@ -0,0 +1,100 @@
+import os, atexit, functools
+from collections import defaultdict
+from typing import List, Any, DefaultDict
+from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, LoadOps, BufferOps, TernaryOps, LazyOp
+from tinygrad.device import Device
+from tinygrad.helpers import GRAPHPATH, DEBUG, GlobalCounters, getenv
+from tinygrad.codegen.linearizer import UOps, UOp
+from tinygrad.shape.symbolic import NumNode
+from tinygrad.lazy import LazyBuffer
+
+try: import networkx as nx
+except ImportError: pass
+
+# **** debugging and graphing ****
+
+if DEBUG >= 2:
+  def print_globalcounters():
+    if GlobalCounters.time_sum_s == 0: return
+    print(f"avg: {GlobalCounters.global_ops*1e-9/GlobalCounters.time_sum_s:8.2f} GFLOPS {GlobalCounters.global_mem*1e-9/GlobalCounters.time_sum_s:8.2f} GB/s",  # noqa: E501
+          f"{' '*10}total: {GlobalCounters.kernel_count:5d} kernels {GlobalCounters.global_ops*1e-9:8.2f} GOPS {GlobalCounters.global_mem*1e-9:8.2f} GB {GlobalCounters.time_sum_s*1e3:8.2f} ms")  # noqa: E501
+  atexit.register(print_globalcounters)
+
+def save_graph(G, fn, opt=""):
+  print("saving", G, f"to {fn}.svg")
+  nx.drawing.nx_pydot.write_dot(G, f'{fn}.dot')
+  os.system(f'dot {opt} -Tsvg {fn}.dot -o {fn}.svg')
+
+G:Any = None
+def init_graph():
+  global G
+  if G is not None: return
+  G = nx.DiGraph()
+  atexit.register(functools.partial(save_graph, G, GRAPHPATH)) # -Gnslimit=100 can make it finish, but you won't like results
+
+counts: DefaultDict[type, int] = defaultdict(int)
+def nm(x):
+  if not hasattr(x, 'node_id'):
+    setattr(x, 'node_id', counts[type(x)])
+    counts[type(x)] += 1
+  return x.node_id
+
+def realized_lazybuffer(lb:'LazyBuffer', num):
+  init_graph()
+  G.nodes[nm(lb)]['style'] = '"filled,bold"'
+  G.nodes[nm(lb)]['fillcolor'] = G.nodes[nm(lb)]['fillcolor'][:-2]
+  G.nodes[nm(lb)]['label'] = '"' + G.nodes[nm(lb)]["label"].replace('"', '') + f'\nK:{num}"'
+
+top_colors = {LoadOps: '#FFFFa0', UnaryOps: "#c0c0c0", ReduceOps: "#FFA0A0", BinaryOps: "#c0c0c0",
+              TernaryOps: "#c0c0c0", BufferOps: '#a0a0ff'}
+def log_lazybuffer(lb:'LazyBuffer', scheduled=False):
+  init_graph()
+  if lb.base.realized is None and lb.base.op is LoadOps.CONST: return
+  if lb.base != lb:
+    offset = lb.st.expr_idxs([NumNode(0)] * len(lb.st.shape))[0]
+    label = f"{lb.st.shape}\n{lb.st.real_strides()}" + (f"\n{offset}" if offset != 0 else "")
+    G.add_node(nm(lb), style='"filled,dashed"', fillcolor="#80ff8080", color="black", label=label)
+    G.add_edge(nm(lb.base), nm(lb), color='#00000060')
+    lb = lb.base
+  if lb.realized is None:
+    label_append = []
+    for idx,x in enumerate(lb.srcs):
+      if nm(x) not in G.nodes: log_lazybuffer(x)
+      if x.base.realized is None and x.base.op is LoadOps.CONST:
+        label_append.append(f"\nCONST{idx} {x.base.arg}")
+      else:
+        G.add_edge(nm(x), nm(lb), color='#a0a0a0')
+    label = '"' + \
+      (str(set(x.shape for x in lb.srcs))+"\n"+str(lb.shape) if lb.op in ReduceOps else str(lb.shape)) + \
+      (f"\n{lb.dtype.name}" if lb.dtype.name != "float" else "")+f"\n{lb.op}"+(f"\n{lb.arg}" if lb.op in {LoadOps.CONST, UnaryOps.CAST} else "") + \
+      (f"\n{lb.device}" if lb.device != Device.DEFAULT else "") + ''.join(label_append) + '"'
+    G.add_node(nm(lb), style='"filled,dashed"', fillcolor=[v for k,v in top_colors.items() if lb.op in k][0] + "80", color="black", label=label)
+    if scheduled: G.nodes[nm(lb)]['shape'] = 'box'
+  else:
+    if nm(lb) not in G.nodes:
+      # realized but unseen?
+      G.add_node(nm(lb), label=f'"{str(lb.base.realized)[5:-1].replace(" ", chr(10))}\nb:{nm(lb.realized)}"', style='filled', fillcolor="#f0c08080")
+
+def _tree(lazyop:LazyOp, cycles, cnt, prefix=""):
+  cnt[0] += 1
+  if len(lazyop.src) == 0: return [f"━━ {prefix}{lazyop.op.name} {lazyop.arg if lazyop.arg else ''}"]
+  if (lid := id(lazyop)) in cycles and cycles[lid][1] > (tcnt := getenv("TREE_CYCLE_CNT", 5)) and tcnt >= 0:
+    return [f"━⬆︎ goto {cycles[id(lazyop)][0]}: {lazyop.op.name}"]
+  cycles[lid] = (cnt[0], 1 if lid not in cycles else cycles[lid][1]+1)
+  lines = [f"━┳ {prefix}{lazyop.op.name} {lazyop.arg if lazyop.arg else ''}"]
+  childs = [_tree(c, cycles, cnt) for c in lazyop.src[:]]
+  for c in childs[:-1]: lines += [f" ┣{c[0]}"] + [f" ┃{l}" for l in c[1:]]
+  return lines + [" ┗"+childs[-1][0]] + ["  "+l for l in childs[-1][1:]]
+
+def print_tree(lazyop:LazyOp): print("\n".join([f"{str(i).rjust(3)} {s}" for i,s in enumerate(_tree(lazyop, {}, [-1]))]))
+
+def graph_uops(uops:List[UOp]):
+  colors = {UOps.ALU: "#ffffc0", UOps.LOAD: "#ffc0c0", UOps.STORE: "#c0ffc0", UOps.SPECIAL: "#c0c0ff", UOps.CONST: "#e0e0e0",
+            UOps.DEFINE_GLOBAL: "#ffe0b0", UOps.DEFINE_LOCAL: "#ffe0d0", UOps.DEFINE_ACC: "#f0ffe0",
+            UOps.RANGE: "#c8a0e0", UOps.PHI: "#e0ffc0", UOps.BARRIER: "#ff8080", UOps.IF: "#c8b0c0"}
+  G = nx.DiGraph()
+  for u in uops:
+    if u.uop in {UOps.ENDRANGE, UOps.ENDIF}: continue
+    G.add_node(uops.index(u), label=f"{str(u.uop)[5:]}{(' '+str(u.arg)) if u.arg is not None else ''}\n{str(u.dtype)}", style="filled", fillcolor=colors.get(u.uop, "#ffffff"))  # noqa: E501
+    for v in u.vin: G.add_edge(uops.index(v), uops.index(u))
+  save_graph(G, f'{GRAPHPATH}.uops', '-Grankdir=LR')

tinygrad/engine/jit.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+from typing import TypeVar, Generic, Callable, List, Tuple, Union, Dict, cast, Optional, Any
+import functools, itertools, collections
+from tinygrad.tensor import Tensor
+from tinygrad.lazy import LazyBuffer
+from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, GRAPH, BEAM, getenv, all_int, GraphException, colored, JIT
+from tinygrad.device import Buffer, Compiled, Device
+from tinygrad.dtype import DType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.symbolic import Variable, sint
+from tinygrad.engine.realize import ExecItem, capturing, EmptyOp, ViewOp, BufferXfer, CompiledRunner, Runner
+from tinygrad.engine.schedule import _internal_memory_planner
+from tinygrad.nn.state import get_parameters
+from weakref import WeakKeyDictionary
+
+def apply_graph_to_jit(jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]) -> List[ExecItem]:
+  # Split JIT cache into batches for faster graph execution.
+  # This allows the accelerator to run some batches while subsequent graphs are still being updated.
+  max_batch_size = getenv("JIT_BATCH_SIZE", 32)
+  graphed_jit_cache: List[ExecItem] = []
+  current_batch: List[ExecItem] = []
+  current_device: Optional[Compiled] = None
+
+  def flush_batch():
+    nonlocal current_batch, current_device, max_batch_size
+    try:
+      if len(current_batch) <= 1 or current_device is None: raise GraphException("only one kernel doesn't graph")
+      graph_runner = current_device.graph(current_batch, input_rawbuffers, var_vals)
+      # clear jit inputs to allow their memory to be freed/reused
+      for (j,i) in graph_runner.input_replace.keys(): graph_runner.jit_cache[j].bufs[i] = None
+      graphed_jit_cache.append(ExecItem(graph_runner, cast(List[Optional[Buffer]], input_rawbuffers)))
+      max_batch_size *= 2
+      if DEBUG >= 2: print(f"\tJIT GRAPHing batch with {len(current_batch)} kernels on device {current_device}")
+    except GraphException as e:
+      graphed_jit_cache.extend(current_batch)
+      if DEBUG >= 2: print(f"\tJIT GRAPHing failed batch with {len(current_batch)} kernels on device {current_device}: {e}")
+    current_batch = []
+    current_device = None
+
+  for ji in jit_cache:
+    if ji.prg.__class__ in {EmptyOp, ViewOp}: continue
+    ji_graph_dev: Optional[Compiled] = None # device on which the ji will be graphed. Not graphed if None.
+    if isinstance(ji.prg, CompiledRunner): ji_graph_dev = ji.prg.device
+    elif isinstance(ji.prg, BufferXfer) and ji.bufs[0] and ji.bufs[0].device.split(":", 1)[0] in {"HSA", "CUDA", "NV", "AMD"}:
+      ji_graph_dev = Device[ji.bufs[0].device]
+
+    graph_class = (ji_graph_dev.graph.func if isinstance(ji_graph_dev.graph, functools.partial) else ji_graph_dev.graph) if ji_graph_dev else None #type: ignore
+    can_be_graphed = ji_graph_dev and ji_graph_dev.graph
+    can_share_graph = (ji_graph_dev == current_device or (isinstance(graph_class, type) and issubclass(graph_class, MultiGraphRunner)) and
+                       type(ji_graph_dev) == type(current_device))
+    can_extend_graph_batch = can_be_graphed and len(current_batch) < max_batch_size and can_share_graph
+    if not can_extend_graph_batch and len(current_batch) > 0: flush_batch()
+
+    if can_be_graphed: current_batch.append(ji)
+    else: graphed_jit_cache.append(ji)
+
+    current_device = ji_graph_dev
+
+  if len(current_batch) > 0: flush_batch()
+  return graphed_jit_cache
+
+def get_input_replace(jit_cache: List[ExecItem], input_rawbuffers:List[Buffer]) -> Dict[Tuple[int, int], int]:
+  input_replace: Dict[Tuple[int, int], int] = {}
+  for j,ji in enumerate(jit_cache):
+    for i,a in enumerate(ji.bufs):
+      if a in input_rawbuffers:
+        input_replace[(j,i)] = input_rawbuffers.index(a)
+  return input_replace
+
+class GraphRunner(Runner):  # pylint: disable=abstract-method
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    self.jit_cache = jit_cache
+    self.input_replace = get_input_replace(jit_cache, input_rawbuffers)
+    self.jc_idx_with_updatable_launch_dims = []
+    self.jc_idx_with_updatable_var_vals = []
+    op_estimate: sint = 0
+    mem_estimate: sint = 0
+    for j,ji in enumerate(jit_cache):
+      op_estimate += ji.prg.op_estimate
+      mem_estimate += ji.prg.mem_estimate
+      if isinstance(ji.prg, CompiledRunner):
+        if ji.prg.p.vars: self.jc_idx_with_updatable_var_vals.append(j)
+        if (ji.prg.p.global_size and not all_int(ji.prg.p.global_size)) or (ji.prg.p.local_size and not all_int(ji.prg.p.local_size)):
+          self.jc_idx_with_updatable_launch_dims.append(j)
+    self.vars = list(var_vals.keys())
+    super().__init__(colored(f"<batched {len(self.jit_cache)}>", "cyan"), jit_cache[0].prg.dname.split(":")[0], op_estimate, mem_estimate)
+
+class MultiGraphRunner(GraphRunner):  # pylint: disable=abstract-method
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    self.w_dependency_map: Dict[Any, Any] = {}
+    self.r_dependency_map: Dict[Any, List[Any]] = collections.defaultdict(list)
+    super().__init__(jit_cache, input_rawbuffers, var_vals)
+
+  def _access_resources(self, read, write, new_dependency:Any):
+    # To synchronize access to resources, we monitor the necessary prerequisites for accessing each resource,
+    # whether for write or read operations. A resource can be accessed by either a single writer or multiple readers.
+    wait_nodes = []
+
+    for rawbuf in read + write:
+      if id(rawbuf._buf) in self.w_dependency_map: wait_nodes.append(self.w_dependency_map[id(rawbuf._buf)])
+    for rawbuf in write:
+      if id(rawbuf._buf) in self.r_dependency_map: wait_nodes.extend(self.r_dependency_map.pop(id(rawbuf._buf)))
+
+    for rawbuf in read: self.r_dependency_map[id(rawbuf._buf)].append(new_dependency)
+    for rawbuf in write: self.w_dependency_map[id(rawbuf._buf)] = new_dependency
+    return list({id(x):x for x in wait_nodes}.values())
+
+ReturnType = TypeVar('ReturnType')
+class TinyJit(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType]):
+    self.fxn = fxn
+    self.reset()
+
+  def add_buffer(self, b:Buffer) -> Buffer:
+    if found:=self.buffer_replace.get(b, None): return found
+    if b.is_allocated() or b.lb_refcount > 0: return b
+    if b._base is not None:
+      self.buffer_replace[b] = ret = Buffer(b.device, b.size, b.dtype, base=self.buffer_replace.get(b._base, b._base), offset=b.offset)
+    else:
+      self.buffer_replace[b] = ret = Buffer(b.device, b.size, b.dtype, options=b.options)
+    return ret
+
+  def add(self, ei:ExecItem):
+    self.jit_cache.append(ExecItem(ei.prg, [self.add_buffer(buf) for buf in ei.bufs if buf is not None]))
+
+  def reset(self):
+    self.jit_cache: List[ExecItem] = []
+    self.input_replace: Dict[Tuple[int, int], int] = {}
+    self.extra_view_inputs: List[Tuple[int, int, str, int, DType]] = []
+    self.buffer_replace: WeakKeyDictionary[Buffer, Buffer] = WeakKeyDictionary()
+    self.cnt: int = 0
+
+  def __get__(self, obj, objtype): return functools.partial(self.__call__, obj) # add support for instance methods
+
+  def __call__(self, *args, **kwargs) -> ReturnType:
+    input_tensors: List[Tuple[Union[int, str], Tensor]] = \
+      [(cast(Union[int, str], k),v) for k,v in itertools.chain(enumerate(args), sorted(kwargs.items())) if v.__class__ is Tensor]
+    if len(input_tensors): Tensor.realize(*[x[1] for x in input_tensors])
+    lbs: List[LazyBuffer] = flatten([v.lazydata.lbs for _,v in input_tensors])
+    expected_sts_var_dtype_device = [(*x.st.unbind(), x.dtype, x.device) for x in lbs]
+    input_rawbuffers: List[Buffer] = [v.base.realized for v in lbs if v.base.realized is not None]
+    assert len(set(input_rawbuffers)) == len(input_rawbuffers), "duplicate inputs to JIT"
+    var_vals: Dict[Variable, int] = merge_dicts([x[1] for x in expected_sts_var_dtype_device] + \
+                                                [dict(x.unbind() for x in itertools.chain(args, kwargs.values()) if isinstance(x, Variable))])
+
+    expected_names, expected_lbs = [x[0] for x in input_tensors], [(x[0], tuple(x[1].keys()), x[2], x[3]) for x in expected_sts_var_dtype_device]
+    if self.cnt == 0:
+      # jit ignore
+      with Context(BEAM=0 if getenv("IGNORE_JIT_FIRST_BEAM") else BEAM.value):
+        self.ret = self.fxn(*args, **kwargs)
+        if len(params:=get_parameters(self.ret)): Tensor.realize(params[0], *params[1:])
+    elif self.cnt == 1:
+      # jit capture
+      self.expected_names: List[Union[int, str]] = expected_names
+      self.expected_lbs: List[Tuple[ShapeTracker, Tuple[Variable, ...], DType, str]] = expected_lbs
+      with Context(GRAPH=getenv("JITGRAPH", GRAPH.value), BEAM=getenv("JITBEAM", BEAM.value)):
+        capturing.append(self)
+        self.ret = self.fxn(*args, **kwargs)
+        if len(params:=get_parameters(self.ret)): Tensor.realize(params[0], *params[1:])
+        capturing.clear()
+      del self.buffer_replace
+      assert len(self.jit_cache), "didn't JIT anything!"
+      if DEBUG >= 1: print(f"JIT captured {len(self.jit_cache)} kernels with {len(input_rawbuffers)} inputs")
+
+      # track inputs that are views of buffers
+      for ji in self.jit_cache:
+        for b in ji.bufs:
+          if b is not None and b._base is not None and b._base in input_rawbuffers:
+            input_rawbuffers.append(b)
+            self.extra_view_inputs.append((input_rawbuffers.index(b.base), b.offset, b.device, b.size, b.dtype))
+
+      # memory planning (optional)
+      assigned = _internal_memory_planner([cast(List[Buffer], x.bufs) for x in self.jit_cache], debug_prefix="JIT ")
+      self.jit_cache = [ExecItem(ei.prg, [assigned.get(x,x).ensure_allocated() for x in ei.bufs if x is not None]) for ei in self.jit_cache]
+
+      # Condense the items into a graph executor.
+      if JIT < 2: self.jit_cache = apply_graph_to_jit(self.jit_cache, input_rawbuffers, var_vals)
+
+      self.input_replace = get_input_replace(self.jit_cache, input_rawbuffers)
+      if DEBUG >= 1 and len(set(self.input_replace.values())) != len(input_rawbuffers): print("WARNING: some input tensors not found")
+    elif self.cnt >= 2:
+      # jit exec
+      assert self.expected_names == expected_names, f"args mismatch in JIT: {self.expected_names=} != {expected_names}"
+      assert self.expected_lbs == expected_lbs, f"args mismatch in JIT: {self.expected_lbs=} != {expected_lbs=}"
+      for idx, offset, device, size, dtype in self.extra_view_inputs:
+        input_rawbuffers.append(Buffer(device, size, dtype, base=input_rawbuffers[idx], offset=offset).ensure_allocated())
+      for (j,i),input_idx in self.input_replace.items(): self.jit_cache[j].bufs[i] = input_rawbuffers[input_idx]
+      if DEBUG >= 1 and len(self.jit_cache) >= 10: print(f"jit execs {len(self.jit_cache)} kernels")
+      for ei in self.jit_cache: ei.run(var_vals, jit=True)
+
+    # clear jit inputs
+    for (j,i) in self.input_replace.keys(): self.jit_cache[j].bufs[i] = None
+
+    self.cnt += 1
+    return self.ret

tinygrad/engine/realize.py

@@ -0,0 +1,191 @@
+from typing import List, Dict, Optional, cast, Generator, Tuple
+import time
+from dataclasses import dataclass, replace
+from tinygrad.helpers import colored, getenv, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int
+from tinygrad.ops import BufferOps, LoadOps, LazyOp
+from tinygrad.device import Device, Buffer
+from tinygrad.shape.symbolic import Variable, sym_infer, sint
+from tinygrad.renderer import Renderer, Program
+from tinygrad.codegen.linearizer import Linearizer
+from tinygrad.engine.schedule import ScheduleItem
+
+# **************** Program Creation ****************
+
+logkerns, logkerns_level = open(getenv("LOGKERNS", ""), "a") if getenv("LOGKERNS", "") else None, getenv("LOGKERNS_LEVEL", 1)
+def get_linearizer(renderer:Renderer, ast:Tuple[LazyOp, ...]) -> Linearizer:
+  if DEBUG >= 3:
+    from tinygrad.engine.graph import print_tree
+    for op in ast: print_tree(op)
+  k = Linearizer(*ast, opts=renderer)
+  k.required_optimizations()
+  if not NOOPT:
+    if not (used_tensor_cores:=k.apply_tensor_cores(getenv("TC", 1))): k.hand_coded_optimizations()
+    if BEAM >= 1:
+      from tinygrad.engine.search import beam_search, time_linearizer, bufs_from_lin
+      kb, k_opt = Linearizer(*ast, opts=renderer), k
+      kb.required_optimizations()
+      rawbufs = bufs_from_lin(kb, allocate=False)
+      k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+      if getenv("BEAM_COMPARE", 1):
+        # TODO: move the HC/TC/BEAM compare to beam_search so it can be optionally cached which choice is better
+        lins: List[Tuple[str, Linearizer]] = [(f"beam{BEAM.value}", k), (("tc" if used_tensor_cores else "hc"), k_opt)]
+        if used_tensor_cores:
+          lins.append(("hc", Linearizer(*ast, opts=renderer)))
+          lins[-1][1].hand_coded_optimizations()
+        timed = sorted([(nm, tk, time_linearizer(tk, rawbufs, allow_test_size=False, clear_l2=True)) for nm, tk in lins], key=lambda x: x[2])
+        if DEBUG >= 1: print("  <  ".join(f"{nm:6s} : {lin.colored_shape(30, dense=True)} : {tm*1e6:8.2f} us" for nm, lin, tm in timed))
+        k = timed[0][1]
+        if logkerns is not None and logkerns_level > 1: logkerns.writelines([f"{(lin.ast, lin.applied_opts)}\n" for (_,lin,_) in timed[1:]])
+  # TODO: check the correctness inline once compare_linearizer is in core
+  if logkerns is not None: logkerns.writelines([f"{(k.ast, k.applied_opts)}\n"])
+  if DEBUG >= 4: print((k.ast, k.applied_opts)) # print here to show final applied_opts for all kernels instead of just in beam_search
+  return k
+
+# **************** Runners ****************
+
+class Runner:
+  def __init__(self, display_name:str, dname:str, op_estimate:sint=0, mem_estimate:sint=0):
+    self.first_run, self.display_name, self.dname, self.op_estimate, self.mem_estimate = True, display_name, dname, op_estimate, mem_estimate
+  @property
+  def device(self): return Device[self.dname]
+  def exec(self, rawbufs:List[Buffer], var_vals:Optional[Dict[Variable, int]]=None) -> Optional[float]:
+    return self(rawbufs, {} if var_vals is None else var_vals)
+  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False) -> Optional[float]:
+    raise NotImplementedError("override this")
+
+class CompiledRunner(Runner):
+  def __init__(self, p:Program, precompiled:Optional[bytes]=None):
+    if DEBUG >= 4: print(p.src)
+    self.p:Program = p
+    self.lib:bytes = precompiled if precompiled is not None else Device[p.dname].compiler.compile_cached(p.src)
+    self.clprg = Device[p.dname].runtime(p.function_name, self.lib)
+    super().__init__(p.name, p.dname, p.op_estimate, p.mem_estimate)
+
+  def __reduce__(self): return self.__class__, (self.p, self.lib)
+
+  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False) -> Optional[float]:
+    global_size, local_size = self.p.launch_dims(var_vals)
+    if global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
+      # TODO: this is copied from get_program
+      from tinygrad.engine.search import optimize_local_size
+      local_size = optimize_local_size(self.clprg, global_size, rawbufs)
+      global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
+      self.p = replace(self.p, global_size=global_size, local_size=local_size)
+    lra = {}
+    if global_size:
+      lra['global_size'] = global_size
+      assert len(global_size) == 3, "global size must have len 3"
+    if local_size:
+      lra['local_size'] = local_size
+      assert len(local_size) == 3, "local size must have len 3"
+    return self.clprg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k] for k in self.p.vars), wait=wait)
+
+class CustomOp(Runner):
+  def __init__(self, fxn):
+    self.fxn = fxn
+    super().__init__(self.fxn.__name__, "CUSTOM", 0, 0)
+  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False): self.fxn(*rawbufs)
+
+class EmptyOp(Runner):
+  def __init__(self, buf:Buffer): super().__init__(colored(f"empty {buf.size:10d} {buf.dtype}", "yellow"), buf.device)
+  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False): pass
+
+class ViewOp(Runner):
+  def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
+  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False):
+    assert rawbufs[0]._base is not None and rawbufs[0]._base == rawbufs[1].base, f"must be base {rawbufs}"
+
+class BufferCopy(Runner):
+  def __init__(self, total_sz, dest_device, src_device):
+    if total_sz >= 1e6: name = f"{type(self).__name__[6:].lower()} {total_sz/1e6:7.2f}M, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
+    else: name = f"{type(self).__name__[6:].lower()} {total_sz:8d}, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
+    super().__init__(colored(name, "yellow"), dest_device, 0, total_sz)
+  def copy(self, dest, src):
+    if src.device.startswith("DISK") and hasattr(dest.allocator, 'copy_from_fd') and src.nbytes >= 4096 and hasattr(src.allocator.device, 'fd'):
+      dest.allocator.copy_from_fd(dest._buf, src.allocator.device.fd, src._buf.offset, src.nbytes)
+    elif src.device.startswith("DISK") and hasattr(dest.allocator, 'as_buffer'):
+      # fast(ish) path, uses readinto in diskbuffers
+      src.allocator.copyout(dest.allocator.as_buffer(dest._buf), src._buf)
+    else:
+      dest.copyin(src.as_buffer(allow_zero_copy=True))  # may allocate a CPU buffer depending on allow_zero_copy
+  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False):
+    dest, src = rawbufs[0:2]
+    assert dest.size == src.size and dest.dtype == src.dtype, f"buffer copy mismatch, {dest.size} != {src.size}, {dest.dtype} != {src.dtype}"
+    st = time.perf_counter()
+    self.copy(dest, src)
+    if wait:
+      Device[dest.device].synchronize()
+      return time.perf_counter() - st
+
+class BufferXfer(BufferCopy):
+  def copy(self, dest, src):
+    if hasattr(dest.allocator.device, "track_cross_buffer") and hasattr(src.allocator, "track_cross_device"):
+      dest.allocator.device.track_cross_buffer.append(src)
+      src.allocator.track_cross_device.add(dest.allocator.device)
+    dest.allocator.transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.device, dest_dev=dest.allocator.device)
+
+# **************** method cache ****************
+
+method_cache: Dict[Tuple[str, Tuple[LazyOp, ...], int, bool], CompiledRunner] = {}
+def get_runner(dname:str, ast:Tuple[LazyOp, ...]) -> CompiledRunner:
+  ckey = (dname, ast, BEAM.value, False)
+  if cret:=method_cache.get(ckey): return cret
+  bkey = (dname.split(":")[0], ast, BEAM.value, True)
+  if bret:=method_cache.get(bkey):
+    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, dname=dname), bret.lib)
+  else:
+    prg: Program = get_linearizer(Device[dname].renderer, ast).to_program()
+    if hasattr(prg.uops, "fuzz_paths"):
+      from test.external.fuzz_uops import UOpsFuzzerRunner
+      return UOpsFuzzerRunner(replace(prg, dname=dname))
+    method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, dname=dname))
+  return ret
+
+# **************** lowering functions ****************
+
+@dataclass(frozen=True)
+class ExecItem:
+  prg: Runner
+  bufs: List[Optional[Buffer]]
+  def run(self, var_vals:Optional[Dict[Variable, int]]=None, wait=False, jit=False, do_update_stats=True) -> Optional[float]:
+    bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
+    et = self.prg(bufs, var_vals if var_vals is not None else {}, wait=wait or DEBUG >= 2)
+    if do_update_stats:
+      GlobalCounters.kernel_count += 1
+      GlobalCounters.global_ops += (op_estimate:=sym_infer(self.prg.op_estimate, var_vals))
+      GlobalCounters.global_mem += (mem_estimate:=sym_infer(self.prg.mem_estimate, var_vals))
+      if et is not None: GlobalCounters.time_sum_s += et
+      if DEBUG >= 2:
+        ptm = (colored(f"{et*1e3:9.2f}ms", "yellow") if et > 0.01 else f"{et*1e6:9.2f}us") if et is not None else ""
+        print(f"{colored(f'*** {self.prg.dname[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(38-ansilen(self.prg.display_name))} arg {len(self.bufs):3d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
+              (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_estimate/((et or 1e-20)*1e9):8.2f} GFLOPS, {mem_estimate/((et or 1e-20)*1e9):7.2f} GB/s)"))  # noqa: E501
+      self.prg.first_run = False
+    return et
+
+def lower_schedule_item(si:ScheduleItem) -> ExecItem:
+  assert len(set(x.device for x in si.bufs)) == 1 or si.ast[0].op is LoadOps.COPY or getenv("USE_COPY_KERNEL")
+  if si.ast[0].op is BufferOps.STORE:
+    runner = get_runner(si.outputs[0].device, si.ast)
+    return ExecItem(runner, [si.bufs[x[0]] for x in runner.p.globals])
+  out, ast = si.outputs[0], si.ast[0]
+  if ast.op is LoadOps.COPY:
+    kernel_type = BufferCopy
+    if hasattr(Device[out.device].allocator, 'transfer') and out.device.split(":")[0] == si.inputs[0].device.split(":")[0]:
+      kernel_type = BufferXfer
+    return ExecItem(kernel_type(ast.arg, out.device, si.inputs[0].device), list(si.bufs))
+  if ast.op is LoadOps.CUSTOM: return ExecItem(CustomOp(ast.arg), list(si.bufs))
+  if ast.op is LoadOps.EMPTY: return ExecItem(EmptyOp(out), list(si.bufs))
+  if ast.op is LoadOps.VIEW: return ExecItem(ViewOp(out), list(si.bufs))
+  raise RuntimeError(f"don't know how to lower {ast}")
+
+def lower_schedule(schedule:List[ScheduleItem]) -> Generator[ExecItem, None, None]:
+  while len(schedule): yield lower_schedule_item(schedule.pop(0))
+
+# **************** main run function ****************
+
+capturing: List = []  # put classes with an add method in here
+
+def run_schedule(schedule:List[ScheduleItem], var_vals:Optional[Dict[Variable, int]]=None, do_update_stats=True):
+  for ei in lower_schedule(schedule):
+    if len(capturing): capturing[0].add(ei)
+    ei.run(var_vals, do_update_stats=do_update_stats)