typed

docs/developer/developer.md

@@ -13,7 +13,7 @@ There's also a [doc describing speed](../developer/speed.md)
 
 Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).
 
-The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view. Inputs to a base can be either base or view, inputs to a view can only be a single base.
+The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
 
 ## Scheduling
 

extra/models/llama.py

@@ -245,11 +245,6 @@ def convert_from_huggingface(weights:dict[str, Tensor], n_layers: int, n_heads:
       continue
     sd[keymap[k]] = v
   for k,v in experts.items(): sd[k] = Tensor.stack(*[v[i] for i in range(len(v))])
-
-  # Handle tied embeddings (e.g., Llama 3.2 1B Instruct where lm_head shares weights with embed_tokens)
-  if "output.weight" not in sd and "tok_embeddings.weight" in sd:
-    sd["output.weight"] = sd["tok_embeddings.weight"]
-
   return sd
 
 def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):

extra/optimization/extract_policynet.py

@@ -0,0 +1,114 @@
+import os, sys, sqlite3, pickle, random
+from tqdm import tqdm, trange
+from copy import deepcopy
+from tinygrad.nn import Linear
+from tinygrad.tensor import Tensor
+from tinygrad.nn.optim import Adam
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+from tinygrad.codegen.opt.search import actions
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.helpers import getenv
+
+# stuff needed to unpack a kernel
+from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.uop.ops import Variable
+inf, nan = float('inf'), float('nan')
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+
+INNER = 256
+class PolicyNet:
+  def __init__(self):
+    self.l1 = Linear(1021,INNER)
+    self.l2 = Linear(INNER,INNER)
+    self.l3 = Linear(INNER,1+len(actions))
+  def __call__(self, x):
+    x = self.l1(x).relu()
+    x = self.l2(x).relu().dropout(0.9)
+    return self.l3(x).log_softmax()
+
+def dataset_from_cache(fn):
+  conn = sqlite3.connect(fn)
+  cur = conn.cursor()
+  cur.execute("SELECT * FROM beam_search")
+  X,A = [], []
+  for f in tqdm(cur.fetchall()):
+    Xs,As = [], []
+    try:
+      lin = Kernel(eval(f[0]))
+      opts = pickle.loads(f[-1])
+      for o in opts:
+        Xs.append(lin_to_feats(lin, use_sts=True))
+        As.append(actions.index(o))
+        lin.apply_opt(o)
+      Xs.append(lin_to_feats(lin, use_sts=True))
+      As.append(0)
+    except Exception:
+      pass
+    X += Xs
+    A += As
+  return X,A
+
+if __name__ == "__main__":
+  if getenv("REGEN"):
+    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
+    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset_policy")
+  else:
+    ld = safe_load("/tmp/dataset_policy")
+    X,V = ld['X'].numpy(), ld['V'].numpy()
+
+  print(X.shape, V.shape)
+  order = list(range(X.shape[0]))
+  random.shuffle(order)
+  X, V = X[order], V[order]
+
+  ratio = -256
+  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
+  X,V = X[:ratio], V[:ratio]
+  print(X.shape, V.shape)
+
+  net = PolicyNet()
+  #if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
+  optim = Adam(get_parameters(net))
+
+  def get_minibatch(X,Y,bs):
+    xs, ys = [], []
+    for _ in range(bs):
+      sel = random.randint(0, len(X)-1)
+      xs.append(X[sel])
+      ys.append(Y[sel])
+    return Tensor(xs), Tensor(ys)
+
+  Tensor.training = True
+  losses = []
+  test_losses = []
+  test_accuracy = 0
+  test_loss = float('inf')
+  for i in (t:=trange(500)):
+    x,y = get_minibatch(X,V,bs=256)
+    out = net(x)
+    loss = out.sparse_categorical_crossentropy(y)
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+    cat = out.argmax(axis=-1)
+    accuracy = (cat == y).mean()
+    t.set_description(f"loss {loss.numpy():7.2f} accuracy {accuracy.numpy()*100:7.2f}%, test loss {test_loss:7.2f} test accuracy {test_accuracy*100:7.2f}%")
+
+    losses.append(loss.numpy().item())
+    test_losses.append(test_loss)
+    if i % 10:
+      out = net(X_test)
+      test_loss = out.sparse_categorical_crossentropy(V_test).square().mean().numpy().item()
+      cat = out.argmax(axis=-1)
+      test_accuracy = (cat == y).mean().numpy()
+
+  safe_save(get_state_dict(net), "/tmp/policynet.safetensors")
+
+  import matplotlib.pyplot as plt
+  plt.plot(losses[10:])
+  plt.plot(test_losses[10:])
+  plt.show()

extra/optimization/extract_sa_pairs.py

@@ -0,0 +1,129 @@
+import sys, sqlite3, pickle, math
+from collections import defaultdict
+from tqdm import tqdm, trange
+import numpy as np
+
+# stuff needed to unpack a kernel
+from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.uop.ops import Variable
+inf, nan = float('inf'), float('nan')
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+
+# more stuff
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import actions
+from extra.optimization.helpers import lin_to_feats
+from extra.optimization.pretrain_valuenet import ValueNet
+from tinygrad.nn.optim import Adam
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+import random
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv
+
+def dataset_from_cache(fn):
+  conn = sqlite3.connect(fn)
+  cur = conn.cursor()
+  cur.execute("SELECT * FROM time_linearizer")
+  grouped = defaultdict(dict)
+  for f in tqdm(cur.fetchall()): grouped[f[0]][f[1:-1]] = pickle.loads(f[-1])
+
+  opts_to_outcome = {}
+
+  for ast,sk in grouped.items():
+    cnts = defaultdict(int)
+    for sks,tm in sk.items():
+      if sks[1] != 1: continue
+      opts = eval(sks[0])
+      cnts[(len(opts), sks[1])] += 1
+      opts_to_outcome[(ast, tuple(opts))] = tm
+    #print(cnts)
+
+  S,A,V = [], [], []
+  for ast,k in tqdm(opts_to_outcome):
+    if len(k) == 0: continue
+    old_tm = min(opts_to_outcome[(ast,k[:-1])])
+    new_tm = min(opts_to_outcome[(ast,k)])
+    if math.isinf(old_tm) or math.isinf(new_tm) or old_tm < 1e-9 or new_tm < 1e-9: continue
+    try:
+      lin = Kernel(eval(ast))
+    except Exception:
+      continue
+    lin.apply_opts(k[:-1])
+    act = k[-1]
+    log_ratio = math.log(old_tm/new_tm)
+    #print(f"ratio: {old_tm/new_tm:6.2f}x (log {log_ratio:5.2f}) from {str(act):50s} on {lin.colored_shape()}")
+    S.append(lin_to_feats(lin, use_sts=True))
+    A.append(actions.index(act))
+    V.append([log_ratio])  # NOTE: i have written the bug many times with this having the wrong dim
+
+  S, A, V = np.array(S), np.array(A), np.array(V, dtype=np.float32)
+  X = np.zeros((S.shape[0], S.shape[1]+len(actions)), dtype=np.float32)
+  X[:, :S.shape[1]] = S
+  X[range(S.shape[0]), S.shape[1]+A] = 1.0
+  return X, V
+
+def log_likelihood(x:Tensor, mu:Tensor, log_sigma:Tensor):
+  #print(x.shape, mu.shape, log_sigma.shape)
+  #return (x-mu).abs() * (-log_sigma).exp() + log_sigma
+  return (x-mu).square() * (-2*log_sigma).exp() / 2 + log_sigma
+
+if __name__ == "__main__":
+  if getenv("REGEN"):
+    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
+    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset")
+  else:
+    ld = safe_load("/tmp/dataset")
+    X,V = ld['X'].numpy(), ld['V'].numpy()
+
+  print(X.shape, V.shape)
+  order = list(range(X.shape[0]))
+  random.shuffle(order)
+  X, V = X[order], V[order]
+
+  ratio = -512
+  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
+  X,V = X[:ratio], V[:ratio]
+  print(X.shape, V.shape)
+
+  #print(X[0], V[0])
+  #print(X[-1], V[-1])
+  print(X.shape)
+
+  net = ValueNet(X.shape[1], 2)
+  optim = Adam(get_parameters(net))
+
+  def get_minibatch(X,Y,bs):
+    xs, ys = [], []
+    #random.seed(1337)
+    for _ in range(bs):
+      sel = random.randint(0, len(X)-1)
+      xs.append(X[sel])
+      ys.append(Y[sel])
+    return Tensor(xs), Tensor(ys)
+
+  Tensor.training = True
+  losses = []
+  test_losses = []
+  test_loss = float('inf')
+  for i in (t:=trange(2000)):
+    x,y = get_minibatch(X,V,bs=256)
+    out = net(x)
+    #loss = (out-y).square().mean()
+    loss = log_likelihood(y, out[:, 0:1], out[:, 1:2]).mean()
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
+    losses.append(loss.numpy().item())
+    test_losses.append(test_loss)
+    if i % 10: test_loss = (net(X_test)[:, 0:1]-V_test).square().mean().numpy().item()
+
+  safe_save(get_state_dict(net), "/tmp/qnet.safetensors")
+
+  import matplotlib.pyplot as plt
+  plt.plot(losses[20:])
+  plt.plot(test_losses[20:])
+  plt.show()

extra/optimization/helpers.py

@@ -0,0 +1,124 @@
+# stuff needed to unpack a kernel
+from tinygrad import Variable
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.dtype import dtypes, PtrDType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.helpers import getenv
+from tinygrad.engine.realize import get_program
+inf, nan = float('inf'), float('nan')
+UOps = Ops
+
+# kernel unpacker
+from tinygrad.codegen.opt.kernel import Kernel
+def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
+def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
+def kern_str_to_lin(kern_str:str, opts=None):
+  (ast, applied_opts,) = eval(kern_str)
+  k = Kernel(ast, opts=opts)
+  k.apply_opts(applied_opts)
+  return k
+
+# load worlds, a dataset of about 12k kernels
+import gzip
+from pathlib import Path
+import random
+from tinygrad.helpers import dedup, DEBUG
+def load_worlds(filter_reduce=True, filter_noimage=True, filter_novariable=True):
+  fn = Path(__file__).parent.parent / "datasets/sops.gz"
+  ast_strs = dedup(gzip.open(fn).read().decode('utf-8').strip().split("\n"))
+  assert len(ast_strs) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
+  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} before filters")
+  if filter_reduce: ast_strs = [x for x in ast_strs if "REDUCE_AXIS" in x]
+  if filter_noimage: ast_strs = [x for x in ast_strs if "dtypes.image" not in x]
+  if filter_novariable: ast_strs = [x for x in ast_strs if "DEFINE_VAR" not in x]
+  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} after filters {filter_reduce=}, {filter_noimage=}, {filter_novariable=}")
+  random.seed(1337)
+  random.shuffle(ast_strs)
+  return ast_strs
+
+def assert_same_lin(l1, l2):
+  assert l1.colored_shape() == l2.colored_shape()
+  assert all(x==y for x,y in zip(l1.sts, l2.sts))
+
+# get features
+import math
+
+MAX_DIMS = 16
+MAX_BUFS = 9
+def lin_to_feats(lin:Kernel, use_sts=True):
+  assert lin.shape_len < MAX_DIMS, "too many dims"
+
+  all_colors = ["blue", "cyan", "white", "green", "red", "magenta", "yellow"]
+  lc = [all_colors.index(x) for x in lin.colors()]
+
+  ret = []
+  # before, some generic linearizer stuff
+  ret.append(lin.upcasted)
+  ret.append(lin.local_dims)
+
+  # first, the full shape, including the colors
+  for s,os,c in zip(lin.full_shape,lin.output_shape,lc):
+    if isinstance(s, UOp):
+      ret.append(False)
+      ret += [0]*9
+    else:
+      ret.append(True)
+      ret.append(math.log2(s))
+      ret.append(min(33, s))
+      ret.append(math.log2(os))
+      ret.append(min(33, os))
+      ret.append(s%2 == 0)
+      ret.append(s%3 == 0)
+      ret.append(s%4 == 0)
+      ret.append(s%8 == 0)
+      ret.append(s%16 == 0)
+    cc = [0]*7
+    cc[c] = 1
+    ret += cc
+  ret += [0] * (17*(MAX_DIMS-len(lin.full_shape)))
+  ret = [float(x) for x in ret]
+
+  if use_sts:
+    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
+    assert len(my_sts) < MAX_BUFS
+    sts_len = 3 + 5*MAX_DIMS
+    for s in my_sts:
+      ret.append(s[0])  # reduce
+      ret.append(s[2])  # has mask
+      ret.append(s[3])  # len views
+      for d in s[1]:
+        ret.append(d is None)
+        ret.append(d == 0)
+        ret.append(d == 1)
+        ret.append(min(33, d) if d is not None else -1)
+        if d is not None and d >= 1: ret.append(math.log2(d))
+        else: ret.append(-1)
+      ret += [0] * (5*(MAX_DIMS - len(s[1])))
+    ret += [0] * (sts_len*(MAX_BUFS - len(my_sts)))
+    assert len(ret) == 1021, f"wrong len {len(ret)}"
+  else:
+    assert len(ret) == 274, f"wrong len {len(ret)}"
+  return ret
+
+from tinygrad.device import Device, Buffer
+from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
+from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
+
+def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501
+  key = {"ast": lin.ast.key, "opts": str(lin.applied_opts), "allow_test_size": allow_test_size,
+         "max_global_size": max_global_size, "clear_l2": clear_l2, "device": lin.opts.device, "suffix": lin.opts.suffix}
+  if not disable_cache and CACHELEVEL >= 2 and (val:=diskcache_get("time_linearizer", key)) is not None: return min(val)
+
+  dev = Device[lin.opts.device]
+  assert dev.compiler is not None
+
+  rawbufs = _ensure_buffer_alloc(rawbufs)
+  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  p = get_program(lin.get_optimized_ast(), lin.opts)
+  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
+                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
+
+  if CACHELEVEL >= 2: diskcache_put("time_linearizer", key, tms)
+  return min(tms)

extra/optimization/pretrain_valuenet.py

@@ -0,0 +1,88 @@
+from tinygrad.codegen.opt.kernel import Kernel
+from tqdm import tqdm, trange
+import math
+import random
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Linear
+from tinygrad.nn.optim import Adam
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+
+# stuff needed to unpack a kernel
+from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.uop.ops import Variable
+inf, nan = float('inf'), float('nan')
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+
+from extra.optimization.helpers import lin_to_feats, MAX_DIMS
+
+# NOTE: this is not real value of the state, it's just a prediction of the runtime
+INNER = 512
+class ValueNet:
+  def __init__(self, feats=240, out=1):
+    self.l1 = Linear(feats,INNER)
+    self.l2 = Linear(INNER,INNER)
+    self.l3 = Linear(INNER,INNER)
+    self.l4 = Linear(INNER,out)
+  def __call__(self, x):
+    x = self.l1(x).relu()
+    x = self.l2(x).relu()
+    x = self.l3(x).relu().dropout(0.8)
+    return self.l4(x)
+
+if __name__ == "__main__":
+  net = ValueNet()
+  optim = Adam(get_parameters(net))
+
+  TEST_SIZE = 256
+
+  dset = open("/tmp/logtm").read().strip().split("\n")
+  random.seed(1337)
+  random.shuffle(dset)
+
+  X,Y = [], []
+  for i,x in enumerate(tqdm(dset)):
+    ast, opts, tms = eval(x)
+    lin = Kernel(ast)
+    for o in opts: lin.apply_opt(o)
+    if lin.shape_len >= MAX_DIMS: continue
+    if min(tms) == float('inf'): continue
+    X.append(lin_to_feats(lin))
+    Y.append([math.log(min(tms))])
+  print(f"got {len(X)} samples")
+
+  X_test,Y_test = Tensor(X[-TEST_SIZE:]), Tensor(Y[-TEST_SIZE:])
+  X,Y = X[:-TEST_SIZE], Y[:-TEST_SIZE]
+
+  def get_minibatch(X,Y,bs):
+    xs, ys = [], []
+    for _ in range(bs):
+      sel = random.randint(0, len(X)-1)
+      xs.append(X[sel])
+      ys.append(Y[sel])
+    return Tensor(xs), Tensor(ys)
+
+  Tensor.training = True
+  losses = []
+  test_losses = []
+  test_loss = float('inf')
+  for i in (t:=trange(2000)):
+    x,y = get_minibatch(X,Y,bs=256)
+    out = net(x)
+    loss = (out-y).square().mean()
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
+    losses.append(loss.numpy().item())
+    test_losses.append(test_loss)
+    if i % 10: test_loss = (net(X_test)-Y_test).square().mean().numpy().item()
+
+  safe_save(get_state_dict(net), "/tmp/valuenet.safetensors")
+
+  import matplotlib.pyplot as plt
+  plt.plot(losses[200:])
+  plt.plot(test_losses[200:])
+  plt.show()

extra/test_pyrender.py

@@ -0,0 +1,20 @@
+from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from tinygrad.helpers import tqdm
+from tinygrad.uop.ops import pyrender, UOp, Ops
+from tinygrad import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker, View
+inf, nan = float('inf'), float('nan')
+
+if __name__ == "__main__":
+  ast_strs = load_worlds()
+  for i, ast_str in enumerate(tqdm(ast_strs)):
+    good_ast = ast_str_to_ast(ast_str)
+    code = '\n'.join(pyrender(good_ast))
+    print("\n***************\n\n"+code)
+    exec(code)
+    if str(good_ast) != str(ast):
+      print(code)
+      print("MISMATCH")
+      print(good_ast)
+      print(ast)
+      break

test/external/external_test_nv.py

@@ -0,0 +1,60 @@
+import unittest, struct, array, ctypes
+from tinygrad import Device, dtypes, Tensor
+from tinygrad.helpers import to_mv
+from tinygrad.runtime.ops_nv import NVDevice, HWQueue
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
+from test.external.fuzz_linearizer import get_fuzz_rawbufs
+
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+
+@unittest.skipUnless(Device.DEFAULT == "NV", "NV specific tests/fixes")
+class TestNV(unittest.TestCase):
+  @classmethod
+  def setUpClass(self):
+    TestNV.d0: NVDevice = Device["NV"]
+    TestNV.a = Tensor([0.,1.], device="NV").realize()
+    TestNV.b = self.a + 1
+    si = self.b.schedule()[-1]
+    TestNV.d0_runner = get_runner(TestNV.d0.device, si.ast)
+    TestNV.b.uop.buffer.allocate()
+    TestNV.addr = struct.pack("QQ", TestNV.b.uop.buffer._buf.va_addr, TestNV.a.uop.buffer._buf.va_addr)
+
+  def test_error_on_huge_dims(self):
+    ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=Ops.MUL, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 683), strides=(0, 0, 0, 1), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 683), strides=(0, 0, 683, 1), offset=0, mask=None, contiguous=True),))))), arg=None),), arg=dtypes.float),), arg=(3,)),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 1), strides=(0, 0, 1, 0), offset=0, mask=None, contiguous=True),)))) # noqa: E501
+    opts = [Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.PADTO, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=2), Opt(op=OptOps.LOCAL, axis=0, arg=2)] # noqa: E501
+    with self.assertRaises(RuntimeError) as cm:
+      lin = Kernel(ast)
+      lin.apply_opts(opts)
+      rawbufs = get_fuzz_rawbufs(lin)
+      prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+      prg(rawbufs, {}, wait=True)
+    self.assertEqual(str(cm.exception), "This is a runtime error message")
+
+  def test_buf4_usage(self):
+    TestNV.along = Tensor([105615], device="NV").realize()
+    ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=Ops.SIN, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.ulong, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), arg=dtypes.float),), arg=None),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))) # noqa: E501
+    temp_runner = get_runner(TestNV.d0.device, (ast,))
+    temp_runner([TestNV.b.uop.buffer, TestNV.along.uop.buffer], var_vals={})
+    val = TestNV.b.uop.buffer.as_buffer().cast("f")[0]
+    assert abs(val - 0.80647) < 0.001, f"got val {val}"
+
+  def test_kernargs_no_oob_access(self):
+    kernargs_start = TestNV.d0._gpu_alloc((2 << 20), map_to_cpu=True).va_addr
+    kernargs = kernargs_start + ((2 << 20) - TestNV.d0_runner._prg.kernargs_alloc_size)
+    to_mv(kernargs, 0x160).cast('I')[:] = array.array('I', TestNV.d0_runner._prg.constbuffer_0)
+    ctypes.memmove(kernargs + TestNV.d0_runner._prg.kernargs_offset, TestNV.addr, len(TestNV.addr))
+
+    q = HWQueue()
+    q.exec(TestNV.d0_runner._prg, kernargs, TestNV.d0_runner.global_size, TestNV.d0_runner.local_size)
+    q.signal(TestNV.d0.timeline_signal, TestNV.d0.timeline_value).submit(TestNV.d0)
+    TestNV.d0._wait_signal(TestNV.d0.timeline_signal, TestNV.d0.timeline_value)
+    TestNV.d0.timeline_value += 1
+    val = TestNV.b.uop.buffer.as_buffer().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+
+if __name__ == "__main__":
+  unittest.main()

test/external/fuzz_shapetracker.py

@@ -0,0 +1,61 @@
+import random
+from tinygrad.helpers import DEBUG, getenv
+from test.unit.test_shapetracker import CheckingShapeTracker
+
+def do_permute(st):
+  perm = list(range(0, len(st.shape)))
+  random.shuffle(perm)
+  perm = tuple(perm)
+  if DEBUG >= 1: print("st.permute(", perm, ")")
+  st.permute(perm)
+
+def do_pad(st):
+  c = random.randint(0, len(st.shape)-1)
+  pad = tuple((random.randint(0,2), random.randint(0,2)) if i==c else (0,0) for i in range(len(st.shape)))
+  if DEBUG >= 1: print("st.pad(", pad, ")")
+  st.pad(pad)
+
+def do_reshape_split_one(st):
+  c = random.randint(0, len(st.shape)-1)
+  poss = [n for n in [1,2,3,4,5] if st.shape[c]%n == 0]
+  spl = random.choice(poss)
+  shp = st.shape[0:c] + (st.shape[c]//spl, spl) + st.shape[c+1:]
+  if DEBUG >= 1: print("st.reshape(", shp, ")")
+  st.reshape(shp)
+
+def do_reshape_combine_two(st):
+  if len(st.shape) < 2: return
+  c = random.randint(0, len(st.shape)-2)
+  shp = st.shape[:c] + (st.shape[c] * st.shape[c+1], ) + st.shape[c+2:]
+  if DEBUG >= 1: print("st.reshape(", shp, ")")
+  st.reshape(shp)
+
+def do_shrink(st):
+  c = random.randint(0, len(st.shape)-1)
+  while 1:
+    shrink = tuple((random.randint(0,s), random.randint(0,s)) if i == c else (0,s) for i,s in enumerate(st.shape))
+    if all(x<y for (x,y) in shrink): break
+  if DEBUG >= 1: print("st.shrink(", shrink, ")")
+  st.shrink(shrink)
+
+def do_flip(st):
+  flip = tuple(random.random() < 0.5 for _ in st.shape)
+  if DEBUG >= 1: print("st.flip(", flip, ")")
+  st.flip(flip)
+
+def do_expand(st):
+  c = [i for i,s in enumerate(st.shape) if s==1]
+  if len(c) == 0: return
+  c = random.choice(c)
+  expand = tuple(random.choice([2,3,4]) if i==c else s for i,s in enumerate(st.shape))
+  if DEBUG >= 1: print("st.expand(", expand, ")")
+  st.expand(expand)
+
+shapetracker_ops = [do_permute, do_pad, do_shrink, do_reshape_split_one, do_reshape_combine_two, do_flip, do_expand]
+
+if __name__ == "__main__":
+  random.seed(42)
+  for _ in range(getenv("CNT", 200)):
+    st = CheckingShapeTracker((random.randint(2, 10), random.randint(2, 10), random.randint(2, 10)))
+    for i in range(8): random.choice(shapetracker_ops)(st)
+    st.assert_same()

test/external/fuzz_shapetracker_math.py

@@ -0,0 +1,34 @@
+import random
+from tinygrad.helpers import getenv, DEBUG, colored, trange
+from tinygrad.shape.shapetracker import ShapeTracker
+from test.external.fuzz_shapetracker import shapetracker_ops
+from test.unit.test_shapetracker_math import st_equal, MultiShapeTracker
+
+def fuzz_plus() -> tuple[ShapeTracker, ShapeTracker]:
+  m = MultiShapeTracker([ShapeTracker.from_shape((random.randint(1, 10), random.randint(1, 10), random.randint(1, 10)))])
+  for _ in range(4): random.choice(shapetracker_ops)(m)
+  backup = m.sts[0]
+  m.sts.append(ShapeTracker.from_shape(m.sts[0].shape))
+  for _ in range(4): random.choice(shapetracker_ops)(m)
+  st_sum = backup + m.sts[1]
+  return m.sts[0], st_sum
+
+if __name__ == "__main__":
+  if seed:=getenv("SEED"): random.seed(seed)
+  total = getenv("CNT", 1000)
+  for fuzz in [globals()[f'fuzz_{x}'] for x in getenv("FUZZ", "plus").split(",")]:
+    same_but_neq = 0
+    for _ in trange(total, desc=f"{fuzz}"):
+      st1, st2 = fuzz()
+      eq = st_equal(st1, st2)
+      if getenv("CHECK_NEQ") and eq and st1.simplify() != st2.simplify():
+        print(colored("same but unequal", "yellow"))
+        print(st1.simplify())
+        print(st2.simplify())
+        same_but_neq += 1
+      if DEBUG >= 1:
+        print(f"EXP: {st1}")
+        print(f"GOT: {st2}")
+        print(colored("****", "green" if eq else "red"))
+      if not eq: exit(0)
+    if getenv("CHECK_NEQ"): print(f"same but unequal {same_but_neq}/{total} = {(same_but_neq/total)*100:.2f}%")

test/opt/test_tensor_cores.py

@@ -26,14 +26,14 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
   opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
 
   if ensure_triggered:
-    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, Device.DEFAULT, opts=opts_to_apply)
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
     wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
     tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
     assert wmmas > 0, "tensor core not triggered"
     assert tcs == 1, "tensor core opt not included"
   else:
     try:
-      program = get_program(realized_ast, Device[Device.DEFAULT].renderer, Device.DEFAULT, opts=opts_to_apply)
+      program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
       assert False, "OptOps.TC triggered, expected KernelOptError"
     except KernelOptError: pass
 
@@ -44,7 +44,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   if dtype_in == dtypes.bfloat16: r = r.float()
   realized_ast, bufs = helper_realized_ast(r)
   opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, Device.DEFAULT, opts=opts), device=Device.DEFAULT))
+  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts), device=Device.DEFAULT))
   if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
   assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
   prg.exec(bufs)

test/test_ops.py

@@ -267,7 +267,7 @@ class TestOps(unittest.TestCase):
       for tor_i, ten_i in zip(tor, ten):
         helper_test_op([], lambda: tor_i, lambda: ten_i)
 
-    self.helper_test_exception([], lambda: torch.meshgrid(x, indexing="bad"), lambda: xt.meshgrid(indexing="bad"), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.meshgrid(x, indexing="bad"), lambda: xt.meshgrid(indexing="bad"), expected=Exception)
 
   def test_arange(self):
     helper_test_op([], lambda: torch.arange(10, dtype=torch.int32), lambda: Tensor.arange(10), forward_only=True)
@@ -587,7 +587,7 @@ class TestOps(unittest.TestCase):
         helper_test_op(None, lambda x,y: x.div(y, rounding_mode="trunc"), forward_only=True, vals=[[numerator], [denominator]])
         helper_test_op(None, lambda x,y: x.div(y, rounding_mode="floor"), forward_only=True, vals=[[numerator], [denominator]])
 
-    self.helper_test_exception(None, lambda x,y: x.div(y, rounding_mode="typo"), forward_only=True, vals=[[5], [0]], expected=RuntimeError)
+    self.helper_test_exception(None, lambda x,y: x.div(y, rounding_mode="typo"), forward_only=True, vals=[[5], [0]], expected=Exception)
 
   def test_div_int(self):
     helper_test_op(None, lambda x,y: x/y, Tensor.div, forward_only=True, vals=[[5, 6, 7],[1, 2, 3]])
@@ -2989,7 +2989,7 @@ class TestOps(unittest.TestCase):
     self.helper_test_exception([(4,5,6), (4,5,6)],
       lambda x,src: x.scatter_reduce(dim=0, index=b, src=src, reduce="INVALID"),
       lambda x,src: x.scatter_reduce(dim=0, index=a, src=src, reduce="INVALID"),
-      RuntimeError)
+      Exception)
     # dtype mismatch
     self.helper_test_exception([(4,5,6), (4,5,6)],
       lambda x,src: x.half().scatter_reduce(dim=0, index=b, src=src, reduce="sum"),
@@ -3068,7 +3068,7 @@ class TestOps(unittest.TestCase):
       helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y, reduction=r),
                                          lambda x,y: x.cross_entropy(y, reduction=r))
     self.helper_test_exception([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y, reduction="typo"),
-                                                   lambda x,y: x.cross_entropy(y, reduction="typo"), expected=ValueError)
+                                                   lambda x,y: x.cross_entropy(y, reduction="typo"), expected=Exception)
 
   def test_cross_entropy_smoothing(self):
     for ls in (0., 0.3, 0.7, 1.):
@@ -3131,7 +3131,7 @@ class TestOps(unittest.TestCase):
         lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), reduction=r))
     self.helper_test_exception([(32,10)],
       lambda x: torch.nn.functional.nll_loss(x, torch.tensor(target), reduction="typo"),
-      lambda x: x.nll_loss(Tensor(target), reduction="typo"), expected=ValueError)
+      lambda x: x.nll_loss(Tensor(target), reduction="typo"), expected=Exception)
 
   def test_nll_loss_weight(self):
     target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()

test/test_renderer_failures.py

@@ -23,9 +23,8 @@ def _test_uop_result(inputs:list[Tensor], stores:list[UOp], local_size=None):
       initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
   inbufs = [x.uop.base.buffer for x in inputs]
   src = Device[Device.DEFAULT].renderer.render(uops)
-  lib = Device[Device.DEFAULT].compiler.compile_cached(src)
   ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
-                                  src, Device.DEFAULT, uops[-1], uops=uops, lib=lib, local_size=local_size))
+                                  src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 

test/test_tensor.py

@@ -456,7 +456,7 @@ class TestTinygrad(unittest.TestCase):
 
   def test_tensor_dtype_errors(self):
     with self.assertRaises(AttributeError): Tensor([3], dtype="typo")
-    with self.assertRaises(AttributeError): Tensor([3], dtype=(dtypes.int,))
+    with self.assertRaises(Exception): Tensor([3], dtype=(dtypes.int,))  # AttributeError or TypeCheckError with TYPED=1
 
   def test_tensor_bytes(self):
     data = b"abc123"

test/test_uops.py

@@ -7,6 +7,7 @@ from tinygrad.dtype import dtypes, DType, AddrSpace
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
 from tinygrad.uop.spec import shared_spec
+from tinygrad.renderer import ProgramSpec
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
 from tinygrad.engine.schedule import ExecItem
@@ -25,7 +26,11 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   return ret[:-1]
 
 def _uops_to_prg(uops_list):
-  return CompiledRunner(get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer, Device.DEFAULT))
+  uops = full_rewrite(ast:=UOp.sink(*uops_list), ren=Device[Device.DEFAULT].renderer)
+  src = Device[Device.DEFAULT].renderer.render(uops)
+  has_local = Device[Device.DEFAULT].renderer.has_local
+  return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
+                                global_size=[1,1,1] if has_local else None, local_size=[1,1,1] if has_local else None))
 
 def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   uops.append(UOp(uop, dtype, tuple(src), arg))

tinygrad/codegen/__init__.py

@@ -18,7 +18,6 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
-from tinygrad.device import Compiler
 
 pm_syntactic_sugar = PatternMatcher([
   # INDEX on ptr INDEX concats them
@@ -124,30 +123,6 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher) -> list[UOp]:
     newlst.extend(ret[1])
   return newlst
 
-def do_linearize(prg:UOp, sink:UOp) -> UOp:
-  lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
-  if SPEC: type_verify(lst, program_spec)
-  return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
-
-def do_render(ctx:tuple[Renderer, Compiler], prg:UOp, lin:UOp) -> UOp:
-  src = ctx[0].render(list(lin.src))
-  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),))
-
-def do_compile(ctx:tuple[Renderer, Compiler], prg:UOp, src:UOp) -> UOp:
-  lib = ctx[1].compile_cached(src.arg)
-  return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
-
-pm_to_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"),), name="prg"), do_linearize),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(Ops.SOURCE, name="src")), name="prg"), do_compile),
-])
-
-def full_rewrite_to_program(sink:UOp, ren:Renderer, compiler:Compiler) -> UOp:
-  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  sink = UOp(Ops.PROGRAM, src=(full_sink,))
-  return graph_rewrite(sink, pm_to_program, ctx=(ren, compiler), name="linearize/render/compile")
-
 def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   """
   Function to transform the Kernel UOp graph into a linearized program.

tinygrad/codegen/opt/search.py

@@ -40,7 +40,7 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:lis
   if allow_test_size and p.global_size is not None and max_global_size is not None:
     global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals)
     p = replace(p, global_size=global_size)
-  try: car = CompiledRunner(replace(p, lib=lib))
+  try: car = CompiledRunner(p, precompiled=lib)
   except AssertionError: return [math.inf] * cnt
   tms = []
   input_bufs = [rawbufs[i] for i in car.p.globals]

tinygrad/device.py

@@ -340,7 +340,8 @@ class Compiled:
     # override this in your device implementation
 
 # TODO: move this to each Device
-def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
+def is_dtype_supported(dtype:DType|None, device:str|None=None) -> bool:
+  if dtype is None: return True
   if dtype == dtypes.index: return False
   if device is None: device = Device.DEFAULT
   if dtype == dtypes.bfloat16:

tinygrad/dtype.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
-from typing import Final, ClassVar, Callable, Literal
+from typing import Final, ClassVar, Callable, Literal, TYPE_CHECKING, Any
+if TYPE_CHECKING: import numpy as np
 import math, struct, ctypes, functools
 from dataclasses import dataclass, fields
 from tinygrad.helpers import getenv, prod
@@ -123,7 +124,7 @@ class dtypes:
     if x.__class__ is list or x.__class__ is tuple: return max(dtypes.from_py(xi) for xi in x) if x else dtypes.default_float
     raise RuntimeError(f"Could not infer dtype of {x} with type {type(x)}")
   @staticmethod
-  def as_const(val: tuple[ConstType|InvalidType, ...]|ConstType|InvalidType, dtype:DType):
+  def as_const(val:Any, dtype:DType):
     if isinstance(val, tuple):
       assert len(val) == dtype.count, f"mismatch {val} {dtype}"
       return tuple(dtypes.as_const(x, dtype) for x in val)

tinygrad/engine/realize.py

@@ -2,30 +2,28 @@ from typing import cast, Callable
 import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, PROFILE, ProfilePointEvent, cpu_events, prod, Context
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
 from tinygrad.helpers import unwrap
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, track_rewrites, KernelInfo, pyrender
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
-from tinygrad.codegen import full_rewrite_to_program
+from tinygrad.codegen import full_rewrite
 from tinygrad.codegen.opt import Opt
 
 # **************** Program Creation ****************
 
 @track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer, device:str|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|tuple[Opt, ...]|None=None) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
 
   Args:
     ast: The Ops.SINK rooted AST
     renderer: The renderer used to generate the code
-    device: The device to compile for (defaults to renderer.device)
 
   Returns:
     The ProgramSpec of the program.
   """
-  device = device or renderer.device
 
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -34,14 +32,20 @@ def get_program(ast:UOp, renderer:Renderer, device:str|None=None, opts:list[Opt]
   if opts is not None:
     assert ast.arg is None, "can't apply opts if sink has an arg"
     ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
-  if ast.arg is None: ast = ast.replace(arg=KernelInfo())
+  try:
+    uops = full_rewrite(ast, renderer)
+  except RuntimeError as e:
+    print("***** LINEARIZE FAILURE *****")
+    print(e)
+    print(pyrender(ast))
+    raise
+  assert uops[-1].op is Ops.SINK, "last uop must be sink"
 
-  prg = full_rewrite_to_program(ast, renderer, Device[device].compiler)
-  # SINK/LINEAR/SOURCE/BINARY
-  sink, linear, source, binary = prg.src
+  # print and render
+  if DEBUG >= 6: print_uops(uops)
+  src = renderer.render(uops)
 
-  # legacy
-  return ProgramSpec(sink.arg.name, source.arg, device, sink, list(linear.src), binary.arg,
+  return ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, renderer.device, ast, uops,
                      global_size=[1,1,1] if renderer.has_local or renderer.has_threads else None,
                      local_size=[1,1,1] if renderer.has_local else None)
 
@@ -72,16 +76,19 @@ def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffe
   return ret[1]
 
 class CompiledRunner(Runner):
-  def __init__(self, p:ProgramSpec, prg=None):
+  def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None):
     if DEBUG >= 3: print(p.applied_opts)
     if DEBUG >= 4: print(p.src)
-    assert p.lib is not None, "lib must be provided"
     self.p:ProgramSpec = p
-    if DEBUG >= 7: Device[p.device].compiler.disassemble(p.lib)
-    self._prg = Device[p.device].runtime(p.function_name, p.lib) if prg is None else prg
+    if precompiled is not None: self.lib = precompiled
+    else:
+      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"):
+        self.lib = Device[p.device].compiler.compile_cached(p.src)
+    if DEBUG >= 7: Device[p.device].compiler.disassemble(self.lib)
+    self._prg = Device[p.device].runtime(p.function_name, self.lib) if prg is None else prg
     super().__init__(p.name, p.device, p.estimates)
 
-  def __reduce__(self): return self.__class__, (self.p,)
+  def __reduce__(self): return self.__class__, (self.p, self.lib)
 
   def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
     if var_vals is None: var_vals = {}
@@ -155,9 +162,9 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner:
   if cret:=method_cache.get(ckey): return cret
   bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True)
   if bret:=method_cache.get(bkey):
-    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device))
+    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib)
   else:
-    prg: ProgramSpec = get_program(ast, Device[device].renderer, device)
+    prg: ProgramSpec = get_program(ast, Device[device].renderer)
     method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))
   return ret
 
@@ -207,11 +214,11 @@ class ExecItem:
     et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
     if do_update_stats:
       GlobalCounters.kernel_count += 1
-      GlobalCounters.global_ops += (op_est:=sym_infer(self.prg.estimates.ops, var_vals))
-      GlobalCounters.global_mem += (mem_est:=sym_infer(self.prg.estimates.mem, var_vals))
+      GlobalCounters.global_ops += (op_est:=int(sym_infer(self.prg.estimates.ops, var_vals)))
+      GlobalCounters.global_mem += (mem_est:=int(sym_infer(self.prg.estimates.mem, var_vals)))
       if et is not None: GlobalCounters.time_sum_s += et
       if DEBUG >= 2:
-        lds_est = sym_infer(self.prg.estimates.lds, var_vals)
+        lds_est = int(sym_infer(self.prg.estimates.lds, var_vals))
         mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed
         header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
         ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""

tinygrad/nn/__init__.py

@@ -249,7 +249,7 @@ class LayerNorm:
   print(t.mean().item(), t.std().item())
   ```
   """
-  def __init__(self, normalized_shape:int|tuple[int, ...], eps:float=1e-5, elementwise_affine:bool=True):
+  def __init__(self, normalized_shape:int|tuple[int, ...]|list[int], eps:float=1e-5, elementwise_affine:bool=True):
     self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1)
     self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps
     self.weight: Tensor|None = Tensor.ones(*self.normalized_shape) if elementwise_affine else None

tinygrad/nn/state.py

@@ -84,7 +84,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
 
 # state dict
 
-def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:
+def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Any]:
   """
   Returns a `state_dict` of the object, with optional prefix.
 
@@ -203,7 +203,7 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
 
 # TODO: this should use tar_extract and zip_extract
 @accept_filename
-def torch_load(t:Tensor) -> dict[str, Tensor]:
+def torch_load(t:Tensor) -> dict[str, Any]:
   """
   ```python
   torch_load(fn: Tensor | str | Path) -> dict[str, Tensor]

tinygrad/renderer/__init__.py

@@ -64,7 +64,6 @@ class ProgramSpec:
   device:str
   ast:UOp  # save the base ast (this is method cache key)
   uops:list[UOp]|None=None
-  lib:bytes|None=None  # compiled binary
 
   # filled in from uops (if we have uops)
   global_size:list[int]|None=None

tinygrad/runtime/support/hcq.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import cast, Callable, Type, TypeVar, Generic, Any
+from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence
 import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
@@ -279,14 +279,14 @@ def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]
     if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, (queue_type or type(queue)) is dev.hw_copy_queue_t))
 
 class HCQArgsState(Generic[ProgramType]):
-  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint, ...]=()):
+  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:Sequence[HCQBuffer], vals:Sequence[sint]=()):
     self.buf, self.prg, self.bufs, self.vals = buf, prg, bufs, vals
     self.bind_data:list[tuple[tuple[sint, ...], MMIOInterface, str]] = []
 
   def bind_sints_to_buf(self, *vals:sint, buf:HCQBuffer, fmt, offset=0): self.bind_data.append((vals, buf.cpu_view().view(offset=offset), fmt))
 
 class CLikeArgsState(HCQArgsState[ProgramType]):
-  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint, ...]=(), prefix:list[int]|None=None):
+  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:Sequence[HCQBuffer], vals:Sequence[sint]=(), prefix:list[int]|None=None):
     super().__init__(buf, prg, bufs, vals=vals)
 
     if prefix is not None: self.buf.cpu_view().view(size=len(prefix) * 4, fmt='I')[:] = array.array('I', prefix)
@@ -302,7 +302,7 @@ class HCQProgram(Generic[HCQDeviceType]):
   @staticmethod
   def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec)
 
-  def fill_kernargs(self, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
+  def fill_kernargs(self, bufs:Sequence[HCQBuffer], vals:Sequence[sint]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
     """
     Fills arguments for the kernel, optionally allocating space from the device if `kernargs_ptr` is not provided.
     Args:

tinygrad/tensor.py

@@ -326,9 +326,7 @@ class Tensor(OpMixin):
     assert self.numel() == 1, "must have one element for item"
     return self.data()[(0,) * len(self.shape)]
 
-  # TODO: should be Tensor.tolist() -> Union[list[ConstType], ConstType]. The list is Sequence because mypy expects memoryview.tolist() -> list[int]
-  # src: https://github.com/python/mypy/blob/release-1.6/mypy/typeshed/stdlib/builtins.pyi#L803
-  def tolist(self) -> Sequence[ConstType]|ConstType:
+  def tolist(self) -> list|ConstType:
     """
     Returns the value of this tensor as a nested list.
     Returns single value for const tensor.
@@ -612,7 +610,7 @@ class Tensor(OpMixin):
   # ***** creation helper functions *****
 
   @staticmethod
-  def full(shape:tuple[sint, ...], fill_value:ConstType, **kwargs) -> Tensor:
+  def full(shape:tuple[sint, ...]|int, fill_value:ConstType, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with the given value.
 
@@ -1249,7 +1247,7 @@ class Tensor(OpMixin):
     """
     return self._getitem(indices)
 
-  def __setitem__(self, indices, v:Tensor|ConstType) -> None:
+  def __setitem__(self, indices, v:Tensor|ConstType|list) -> None:
     if isinstance(self.device, str) and self.device.startswith("DISK"):
       self.realize()._getitem(indices).assign(v)
       return
@@ -1289,7 +1287,7 @@ class Tensor(OpMixin):
     x = self.shrink(tuple((0, i) if d != dim else None for d,i in enumerate(index.shape))).unsqueeze(-1).transpose(-1, dim)
     return (index.unsqueeze(-1)._one_hot_along_dim(self.shape[dim]).where(x, 0)).sum(-1, dtype=self.dtype)
 
-  def cat(self:Tensor, *args:Tensor, dim:int=0) -> Tensor:
+  def cat(self:Tensor|tuple[Tensor, ...]|list[Tensor], *args:Tensor, dim:int=0) -> Tensor:
     """
     Concatenates self with other `Tensor` in `args` along an axis specified by `dim`.
     All tensors must have the same shape except in the concatenating dimension.
@@ -1302,14 +1300,15 @@ class Tensor(OpMixin):
     print(t0.cat(t1, t2, dim=1).numpy())
     ```
     """
+    if isinstance(self, (tuple, list)): self, args = self[0], tuple(self[1:]) + args  # type: ignore[arg-type]
     dim = self._resolve_dim(dim)
     for arg in args: assert arg.ndim==self.ndim and all(ti==ai for i,(ti,ai) in enumerate(zip(self.shape, arg.shape)) if i!=dim)
-    tensors = [self, *args]
+    tensors:list[Tensor] = [self, *args]
     dim_cumsum = list(itertools.accumulate([t.shape[dim] for t in tensors], initial=0))
     for i,t in enumerate(tensors): tensors[i] = t.pad([(dim_cumsum[i], dim_cumsum[-1]-dim_cumsum[i+1]) if j==dim else None for j in range(t.ndim)])
     return functools.reduce(Tensor.add, tensors)
 
-  def stack(self:Tensor, *args:Tensor, dim:int=0) -> Tensor:
+  def stack(self:Tensor|tuple[Tensor, ...]|list[Tensor], *args:Tensor, dim:int=0) -> Tensor:
     """
     Concatenates self with other `Tensor` in `args` along a new dimension specified by `dim`.
 
@@ -2114,7 +2113,7 @@ class Tensor(OpMixin):
     return pads
 
   # NOTE: these work for more than 2D
-  def avg_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
+  def avg_pool2d(self, kernel_size:int|tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]|list[int]=0,
                  ceil_mode=False, count_include_pad=True) -> Tensor:
     """
     Applies average pooling over a tensor.
@@ -2160,7 +2159,7 @@ class Tensor(OpMixin):
     if not ceil_mode: return pool(self, reg_pads).mean(axis)
     return pool(self, ceil_pads).sum(axis) / pool(self.pad(reg_pads).ones_like(), tuple(cp-rp for cp,rp in zip(ceil_pads, reg_pads))).sum(axis)
 
-  def max_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
+  def max_pool2d(self, kernel_size:int|tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]|list[int]=0,
                  ceil_mode=False, return_indices=False) -> Tensor | tuple[Tensor, Tensor]:
     """
     Applies max pooling over a tensor.
@@ -2204,7 +2203,8 @@ class Tensor(OpMixin):
     idx = m * idx.pad(pads, value=dtypes.min(idx.dtype))._pool(k_, stride if stride is not None else k_, dilation)
     return pooled.max(axis), spatial_sz - idx.max(axis)
 
-  def max_unpool2d(self, indices:Tensor, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, output_size=None):
+  def max_unpool2d(self, indices:Tensor, kernel_size:int|tuple[int, ...]=(2,2), stride=None, dilation=1,
+                   padding:int|tuple[int, ...]|list[int]=0, output_size=None):
     """
     Performs a partial inverse of `max_pool2d` using the indices from the argmax.
 
@@ -2235,7 +2235,7 @@ class Tensor(OpMixin):
     ret = (indices.reshape(bs,c,1,-1)._one_hot_along_dim(prod(output_size), 2).where(self.reshape(bs,c,1,-1), 0)).sum(3)
     return ret.reshape(bs,c,*output_size)
 
-  def conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding:int|tuple[int, ...]=0,
+  def conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding:int|tuple[int, ...]|list[int]=0,
              dtype:DTypeLike|None=None) -> Tensor:
     """
     Applies a convolution over a tensor with a given `weight` and optional `bias`.
@@ -3614,7 +3614,7 @@ class Tensor(OpMixin):
     nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight
     return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction)
 
-  def newton_schulz(self, steps:int, params:tuple[int, ...], eps:float=1.0e-7) -> Tensor:
+  def newton_schulz(self, steps:int, params:tuple[float|int, ...], eps:float=1.0e-7) -> Tensor:
     """
     Performs the newton-schulz algorithm for odd polynomials. The degree of the odd polynomial depends on the number of params.
 

tinygrad/uop/__init__.py

@@ -27,10 +27,6 @@ class Ops(FastEnum):
   # uops that aren't rendered
   NOOP = auto(); REWRITE_ERROR = auto()
 
-  # renderer/compiler
-  # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has a bytes arg that's not
-  PROGRAM = auto(); LINEAR = auto(); SOURCE = auto(); BINARY = auto()
-
   # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
   # GROUP is a NOOP that just merges things together
   SINK = auto(); AFTER = auto(); GROUP = auto()

tinygrad/uop/ops.py

@@ -46,7 +46,7 @@ def smin(*lst) -> sint: return _suop(argfix(*lst), UOp.minimum, min)
 def srender(x:sint) -> str: return x.render() if isinstance(x, UOp) else str(x)
 
 def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
-def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop
+def sym_infer(uop: UOp|int|float, var_vals: dict[str, int]) -> int|float: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop
 
 def range_str(u:UOp, color=False) -> str:
   ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
@@ -218,8 +218,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | \
-           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL:
         return None
 
       case Ops.INDEX:
@@ -700,11 +699,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     uval = self.const_like(val) if isinstance(val, int) else val
     assert self.arg[1] <= uval.vmin and uval.vmax <= self.arg[2], f"bind {val} not in range [{self.arg[1]}, {self.arg[2]}]"
     return UOp(Ops.BIND, self.dtype, (self, uval))
-  def unbind(self) -> tuple[Variable, int]:
+  def unbind(self) -> tuple[UOp, int]:
     assert self.op is Ops.BIND and self.src[0].op is Ops.DEFINE_VAR and self.src[1].op is Ops.CONST, f"can't unbind {self}"
     return self.src[0], self.src[1].arg
-  def unbind_all(self) -> tuple[UOp, dict[Variable, int]]:
-    ret:dict[Variable, int] = {}
+  def unbind_all(self) -> tuple[UOp, dict[UOp, int]]:
+    ret:dict[UOp, int] = {}
     return graph_rewrite(self, pm_unbind, ctx=ret), ret
   @property
   def val(self) -> int: return self.unbind()[1]
@@ -713,7 +712,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     bound_var_base = set(x.src[0] for x in bound_vars)
     all_vars = set([x for x in self.toposort() if x.op is Ops.DEFINE_VAR])
     return bound_vars.union(set([x for x in all_vars if x not in bound_var_base]))
-  def variables(self) -> list[Variable]:
+  def variables(self) -> list[UOp]:
     return sorted(set([x.unbind()[0] if x.op is not Ops.DEFINE_VAR else x for x in self.vars()]), key=lambda v: v.arg)
 
   # *** uop symbolic stuff ***
@@ -1321,7 +1320,7 @@ def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
-def do_unbind(ctx:dict[Variable, int], x:UOp):
+def do_unbind(ctx:dict[UOp, int], x:UOp):
   v,i = x.unbind()
   ctx[v] = i
   return v

tinygrad/uop/spec.py

@@ -249,16 +249,6 @@ full_spec = PatternMatcher([
   # in progress MSTACK may lose device
   (UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True),
 
-  # codegen: PROGRAM with progressive sources through the pipeline
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK),)), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
-  # codegen: standalone LINEAR/SOURCE/BINARY
-  (UPat(Ops.LINEAR, dtypes.void), lambda: True),
-  (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
-  (UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
-
   # temp VECTORIZE/INDEX during rewrite have the wrong dtype
   (UPat(Ops.VECTORIZE), lambda: True),
   (UPat(Ops.INDEX), lambda: True),