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
 

docs/developer/layout.md

@@ -26,9 +26,9 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 ## tinygrad/codegen
 
-Transform the optimized ast into a linearized and rendered program.
+Transform the optimized ast into a linearized list of UOps.
 
-::: tinygrad.codegen.full_rewrite_to_program
+::: tinygrad.codegen.full_rewrite
     options:
         members: false
         show_labels: false

extra/export_model.py

@@ -1,11 +1,12 @@
 from typing import Tuple, Dict, List, Optional
 from tinygrad.dtype import DType
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv, to_function_name
+from tinygrad.helpers import Context, to_mv
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict
 
@@ -14,13 +15,8 @@ EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
   functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
   for ji in run.jit_cache:
-    prg: UOp = ji.prg.p
-    name = prg.src[0].arg.name
-    function_name = to_function_name(name)
-    src = prg.src[3].arg
-    global_size, local_size = prg.sizes
-    prg_vars = prg.variables()
-    functions[function_name] = src   # NOTE: this assumes all with the same name are the same
+    fxn: ProgramSpec = ji.prg.p
+    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
     cargs = []
     for i,arg in enumerate(ji.bufs):
       key = id(arg)
@@ -32,8 +28,8 @@ def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str]
           bufnum += 1
           if i > 0: bufs_to_save[bufs[key][0]] = arg   # if first usage of a buffer is not an output, and it's not a special name
       cargs.append(bufs[key][0])
-    cargs += [var for var in prg_vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
-    statements.append((function_name, cargs, list(global_size) if global_size else None, list(local_size) if local_size else None))
+    cargs += [var for var in fxn.vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
+    statements.append((fxn.function_name, cargs, fxn.global_size, fxn.local_size))
 
   return functions, statements, {name:(size, dtype, key) for (name,size,dtype,key) in bufs.values()}, bufs_to_save
 

extra/gemm/triton_nv_matmul.py

@@ -4,8 +4,7 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)
 
@@ -86,12 +85,9 @@ if __name__ == "__main__":
   # remove debug sections
   src = src.split("\t.file")[0]
   assert '.extern .shared' not in src
-  # Create linearized uops with SPECIAL for global/local sizes
-  global_size = [M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1]
-  local_size = [32*compiled.metadata.num_warps, 1, 1]
-  uops = [UOp(Ops.SPECIAL, arg=('g', i), src=(UOp.const(dtypes.int, global_size[i]),)) for i in range(3)]
-  uops += [UOp(Ops.SPECIAL, arg=('l', i), src=(UOp.const(dtypes.int, local_size[i]),)) for i in range(3)]
-  prg = UOp.new_program("matmul_kernel", src, Device.DEFAULT, si.ast, uops)
+  prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
+                global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
+                mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
   ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
   tflops = []
   for i in range(5):

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/remu/test/hwtest.py

@@ -1,10 +1,10 @@
 import numpy as np
 import unittest
 import subprocess, struct, math
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad import Tensor, dtypes, Device, UOp
 from tinygrad.helpers import getenv
 from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
+from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner
 
 def get_output(asm:str, n_threads:int=1):
@@ -22,9 +22,7 @@ def get_output(asm:str, n_threads:int=1):
     *(data0_1+l) = res;
   }}"""
   t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
-  # Create linearized uops with SPECIAL for local size
-  uops = [UOp(Ops.SPECIAL, arg=('l', 0), src=(UOp.const(dtypes.int, n_threads),))]
-  prg = UOp.new_program("test", src, Device.DEFAULT, UOp.sink(t.uop), uops)
+  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
   car = CompiledRunner(prg)
   if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
   car([t.uop.buffer], {}, wait=True)

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

opencode.json

@@ -1 +1 @@
-{"$schema": "https://opencode.ai/config.json", "formatter": false, "lsp": false}
+{"$schema": "https://opencode.ai/config.json", "formatter": false}

test/external/external_benchmark_op_conv.py

@@ -1,8 +1,10 @@
 # ruff: noqa: E501 E712 F401
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
+from tinygrad.codegen import full_rewrite
 from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
 from tinygrad.dtype import ImageDType, Invalid
@@ -84,13 +86,16 @@ def dm_conv_172():
 
 ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
 
+compiler = Device.default.compiler
 renderer = Device.default.renderer
 allocator = Device.default.allocator
 
-ps = get_program(ast, renderer)
-# update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-ps = ps.replace(src=(ps.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *ps.src[2:]))
-cr = CompiledRunner(ps)
+uops = full_rewrite(ast, renderer)
+src = renderer.render(uops)
+
+lib = compiler.compile(src)
+ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
+cr = CompiledRunner(ps, precompiled=lib)
 
 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
 # print(len(gs))

test/external/external_test_am.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_GMC
 from tinygrad.runtime.support.hcq import MMIOInterface
-from tinygrad.runtime.support.memory import PageTableTraverseContext, AddrSpace
+from tinygrad.runtime.support.memory import PageTableTraverseContext
 from tinygrad.runtime.autogen.am import am
 from tinygrad.helpers import mv_address
 
@@ -70,7 +70,7 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
                   (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
-      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)])
 
       ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
       results = list(ctx.next(sz))
@@ -102,8 +102,8 @@ class TestAMPageTable(unittest.TestCase):
     mm0 = self.d[0].mm
 
     for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)], aspace=AddrSpace.PHYS)
-      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
       mm0.unmap_range(helper_va(va2), sz2)
       mm0.unmap_range(helper_va(va1), sz1)
 
@@ -112,24 +112,24 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
       exteranl_va = helper_va(va)
-      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)])
 
       mm0.unmap_range(vaddr=exteranl_va, size=sz)
 
       # Finally can map and check paddrs
-      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)])
 
       ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
       for tup in ctx.next(0x100000):
@@ -147,13 +147,13 @@ class TestAMPageTable(unittest.TestCase):
     with self.assertRaises(AssertionError):
       mm0.unmap_range(helper_va(0x10000), 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
     mm0.unmap_range(helper_va(0x10000), 0x3000)
 
     with self.assertRaises(AssertionError):
       mm0.unmap_range(helper_va(0x10000), 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
     mm0.unmap_range(helper_va(0x10000), 0x3000)
 
     with self.assertRaises(AssertionError):
@@ -164,16 +164,16 @@ class TestAMPageTable(unittest.TestCase):
 
     # offset from start
     for off in [0, 0x3000, 0x10000]:
-      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
       mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
     # offset from end
     for off in [0x1000, 0x20000]:
-      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
       mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
   def test_frag_size(self):

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/external/process_replay/process_replay.py

@@ -9,7 +9,7 @@ if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
 
 try:
   from tinygrad.schedule.rangeify import get_rangeify_map
-  from tinygrad.renderer import Renderer
+  from tinygrad.renderer import Renderer, ProgramSpec
   from tinygrad.engine.realize import get_program
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen.opt import Opt
@@ -51,20 +51,15 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
 
-def replay_get_program(p:UOp, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
-  # p is a PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  p_name = p.src[0].arg.name
-  p_applied_opts = p.src[0].arg.applied_opts
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p_applied_opts if BEAM>=1 else None)
-  input_ast = ast.replace(arg=replace(sink_arg, name=p_name))
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
+  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
   p2 = get_program(input_ast, renderer=renderer)
-  def to_str(ret:UOp) -> str:
+  def to_str(ret:ProgramSpec) -> str:
     # PYTHON renderer pickles UOps, first unpickle and decode here
-    ret_src = ret.src[3].arg
-    ret_device = ret.device
-    if ret_device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret_src))])
-    return ret_src
+    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
+    return ret.src
   # properly color the name arg
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
   return to_str(p2), to_str(p), (ast_repr, renderer)

test/helpers.py

@@ -2,27 +2,18 @@ import os, time, struct, functools, unittest
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.engine.realize import Runner, get_program
+from tinygrad.engine.realize import Runner
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, CI
-from tinygrad.renderer import Renderer
-from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups
-from tinygrad.codegen.late.linearizer import linearize
+from tinygrad.codegen import full_rewrite
 
 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
-def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
-  """Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
-  if ren is None: ren = Renderer()
-  if sink.arg is None: sink = sink.replace(arg=KernelInfo())
-  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
-
 def derandomize_model(model):
   for p in get_parameters(model):
     p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
@@ -60,12 +51,13 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
   bufs = []
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
-    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
+    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + buf_dt.fmt, *data)))
   g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
-  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))  # source code is in src[3].arg
+  opts = PythonRenderer()
+  lst = full_rewrite(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), opts)
+  prog = PythonProgram("run", PythonCompiler().compile(opts.render(lst)))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
-  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
+  return out_buf.cast(uop.dtype.fmt).tolist()[0]
 
 def not_support_multi_device():
   # CL and CUDA don't support multi device if in CI

test/opt/test_tensor_cores.py

@@ -1,9 +1,10 @@
 import numpy as np
 import unittest
+from dataclasses import replace
 
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops
 from tinygrad.dtype import DType
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import AMX, AMD_LLVM, CPU_LLVM, Context
@@ -43,10 +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))]
-  p = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts)
-  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  p = p.replace(src=(p.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *p.src[2:]))
-  prg = CompiledRunner(p)
+  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_fusion_op.py

@@ -28,7 +28,7 @@ class TestFusionOp(unittest.TestCase):
     sched = a.schedule()
     sched[-1].lower()
     self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(sched[-1].prg.p.src[3].arg.splitlines()) < 250  # source code is in src[3].arg
+    assert len(sched[-1].prg.p.src.splitlines()) < 250
 
   def test_recursive_add_cmp(self):
     st = time.perf_counter()

test/test_linearizer.py

@@ -1,5 +1,6 @@
 import numpy as np
 import unittest
+from dataclasses import replace
 
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
@@ -167,8 +168,8 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)])
-    self.assertEqual(len(prg.src[3].arg.split("for")), 5)  # source code is in src[3].arg
+    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
+    self.assertEqual(len(prg.src.split("for")), 5)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -516,11 +517,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 
-  def get_prg(opts):
-    prg = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
-    # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-    prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=device), *prg.src[2:]))
-    return CompiledRunner(prg)
+  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

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_opts.py

@@ -14,7 +14,7 @@ class TestOpts(unittest.TestCase):
     self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
     if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
       prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
-      self.assertIn('float4', prg.src[3].arg)  # source code is in src[3].arg
+      self.assertIn('float4', prg.src)
 
 if __name__ == '__main__':
   unittest.main()

test/test_renderer_failures.py

@@ -2,24 +2,29 @@ import unittest
 import numpy as np
 from tinygrad.device import Buffer, Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.helpers import prod
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import dedup, flatten, prod
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, python_alu
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.codegen import full_rewrite
 
-def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
+def _test_uop_result(inputs:list[Tensor], stores:list[UOp], local_size=None):
   for x in inputs: x.realize()
-  uops = prg.uops
+  # NOTE: we only toposort the stores
+  uops: list[UOp] = []
+  def _recursive_add(uop:UOp) -> list[UOp]: return flatten([_recursive_add(x) for x in uop.src])+[uop]
+  uops = dedup(flatten(_recursive_add(st) for st in stores))
   outbufs = [Buffer(Device.DEFAULT, sz:=(1 if local_size is None else prod(local_size)), (dtype:=u.src[1].dtype), \
       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]
-  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *prg.src[2:]))
-  ei = CompiledRunner(prg)
+  src = Device[Device.DEFAULT].renderer.render(uops)
+  ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
+                                  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]
 
@@ -32,8 +37,8 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   alu = ld.alu(alu_op, *alu_src_uops)
   store = UOp.store(a.index(idx), alu)
   sink = UOp(Ops.SINK, dtypes.void, (store,))
-  prg = get_program(sink, Device[Device.DEFAULT].renderer)
-  return _test_uop_result([Tensor([input_val])], prg)[0]
+  uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+  return _test_uop_result([Tensor([input_val])], uops)[0]
 
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -42,8 +47,8 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
+    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -53,8 +58,8 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0]
+    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], uops, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
 
 @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
@@ -71,7 +76,7 @@ class TestCStyleFailures(unittest.TestCase):
     schedule = ret.schedule()
     assert len(schedule) == 1
     schedule[0].lower()
-    src = schedule[0].prg.p.src[3].arg  # source code is in src[3].arg
+    src = schedule[0].prg.p.src
     self.assertEqual("("*5 not in src, should_strip_paren)
 
   def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
@@ -99,8 +104,8 @@ class TestPTXFailures(unittest.TestCase):
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
+    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")

test/test_tensor.py

@@ -10,7 +10,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.engine.realize import get_program
+from tinygrad.codegen import full_rewrite
 from tinygrad.dtype import DType
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -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"
@@ -869,8 +869,9 @@ class TestIdxUpcast(unittest.TestCase):
     for s in schedule:
       if s.ast.op is Ops.SINK:
         renderer = Device[s.bufs[0].device].renderer
-        prg = get_program(s.ast, renderer)
-        return prg.uops
+        uops = full_rewrite(s.ast, renderer)
+        renderer.render(uops)
+        return uops
 
   def _assert(self, dtype: DType, a: Tensor):
     uops = self._schedule_render(a)

test/test_uops.py

@@ -7,28 +7,30 @@ 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
+from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import get_uops
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   sink = UOp.group(*u)
   for r in sink.ranges: sink = sink.end(r)
   # we strip the SINK here for legacy reasons
-  ret = get_uops(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  ret = full_rewrite(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
   assert ret[-1].op is Ops.SINK
   return ret[:-1]
 
 def _uops_to_prg(uops_list):
-  prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
-  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *prg.src[2:]))
-  return CompiledRunner(prg)
+  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))

test/test_uops_stats.py

@@ -2,7 +2,9 @@ import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import get_program
+from tinygrad.renderer import ProgramSpec
 from tinygrad.renderer import Estimates
+from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
@@ -144,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):
     u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
     u4 = UOp(Ops.MUL, dtypes.int, (u1,u2))
     u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
-    uops = list(u5.toposort())
+    uops = full_rewrite(u5.sink())
 
     globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
     o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
@@ -153,7 +155,7 @@ class TestUOpsStats(unittest.TestCase):
     u2 = globl.index(o2)
     u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
     u4 = UOp(Ops.MULACC, dtypes.int, (u1,u2,u3))
-    uops_fma = list(u4.toposort())
+    uops_fma = full_rewrite(u4.sink())
 
     self.assertEqual(flops_mem(uops), flops_mem(uops_fma))
 
@@ -165,26 +167,24 @@ class TestStatsOptimized(unittest.TestCase):
     cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast
     cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule()[-1].ast
 
-  def check_gemm(self, p:UOp, extra_flops=0):
+  def check_gemm(self, p:ProgramSpec, extra_flops=0):
     #p.uops.print()
-    #print(p.src[3].arg)
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    print(p.src[0].arg.name, estimates.ops, estimates.mem, estimates.lds)
-    self.assertEqual(estimates.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
-    self.assertEqual(estimates.mem, 3*N*N*4) # 3 NxN mats with floats
+    #print(p.src)
+    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
+    self.assertEqual(p.estimates.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
+    self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
     p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[])
     self.check_gemm(p)
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    self.assertEqual(estimates.lds, 2*N*N*N*4 + 4*N*N)
+    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
   def test_gemm_tc_unroll(self):
     try:
       p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
     except KernelOptError:
       raise unittest.SkipTest("no tensor cores")
-    print(p.src[3].arg)
+    print(p.src)
     self.check_gemm(p)
 
   # this is a good lesson about why UPCASTing is a good idea
@@ -192,15 +192,13 @@ class TestStatsOptimized(unittest.TestCase):
   def test_gemm_one_upcasted(self):
     p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4)])
     self.check_gemm(p)
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    self.assertEqual(estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
+    self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
     p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer,
                     opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
     self.check_gemm(p)
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    self.assertEqual(estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
     try:
@@ -209,8 +207,7 @@ class TestStatsOptimized(unittest.TestCase):
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     self.check_gemm(p)
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    self.assertEqual(estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_group(self):
     try:
@@ -220,15 +217,13 @@ class TestStatsOptimized(unittest.TestCase):
     SZ = N*N*4
     # NOTE: these are sort of wrong. they aren't honoring the IF statement
     self.check_gemm(p, extra_flops=SZ*4)
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    self.assertEqual(estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
+    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
 
   def test_reduce(self):
     p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[])
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    print(p.src[0].arg.name, estimates.ops, estimates.mem, estimates.lds)
-    self.assertEqual(estimates.ops, N*N)
-    self.assertEqual(estimates.mem, N*N*4 + 4)
+    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
+    self.assertEqual(p.estimates.ops, N*N)
+    self.assertEqual(p.estimates.mem, N*N*4 + 4)
 
   def test_reduce_group(self):
     try:
@@ -236,8 +231,7 @@ class TestStatsOptimized(unittest.TestCase):
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     # NOTE: these are wrong, they don't respect the if statement
-    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-    print(p.src[0].arg.name, estimates.ops, estimates.mem, estimates.lds)
+    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
 
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/testextra/test_cfg_viz.py

@@ -3,6 +3,7 @@ import textwrap
 
 from tinygrad import Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, track_rewrites
+from tinygrad.renderer import ProgramSpec
 from tinygrad.helpers import TracingKey
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 
@@ -50,9 +51,9 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
 
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.src[0].arg.name, ret=ret))
-def run_asm(name:str, src:str) -> UOp:
-  prg = UOp.new_program(name, template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK), [])
+@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, ret=ret))
+def run_asm(name:str, src:str) -> ProgramSpec:
+  prg = ProgramSpec(name, template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK))
   ei = ExecItem(UOp(Ops.SINK), [Tensor.empty(1).uop.buffer.ensure_allocated()], prg=CompiledRunner(prg))
   ei.run()
   return prg

test/unit/test_dtype_spec.py

@@ -8,7 +8,7 @@ from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
 
-settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
 
 core_dtypes = list(DTYPES_DICT.values())

test/unit/test_llm_server.py

@@ -27,8 +27,8 @@ class TestLLMServer(unittest.TestCase):
     from tinygrad.apps.llm import Handler
     from tinygrad.helpers import TCPServerWithReuse
 
-    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
-    cls.port = cls.server.server_address[1]
+    cls.port = 11435
+    cls.server = TCPServerWithReuse(('127.0.0.1', cls.port), Handler)
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
     time.sleep(0.1)

test/unit/test_uop_symbolic.py

@@ -3,8 +3,8 @@ import unittest, pickle, functools, math
 import z3
 
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
+from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
-from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
 from tinygrad.uop.validate import uops_to_z3
@@ -747,7 +747,7 @@ class TestSymbolic(unittest.TestCase):
 
     # TODO: copied from render, render does not support cast
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
-    uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
+    uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
     rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
 
     self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))

tinygrad/apps/llm.py

@@ -11,8 +11,8 @@ class SimpleTokenizer:
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    # TODO: ucat_range is slow
+    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
     self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
       f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")

tinygrad/codegen/__init__.py

@@ -123,21 +123,20 @@ 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 full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
+  """
+  Function to transform the Kernel UOp graph into a linearized program.
 
-def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
-  src = ctx.render(list(lin.src))
-  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),))
+  Args:
+    sink: The Ops.SINK rooting the Kernel graph.
+    ren: The Renderer (can change how things are processed, fix this).
 
-pm_to_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE)), name="prg"), do_linearize),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
-])
+  Returns:
+    Linear program in UOps.
+  """
 
-def full_rewrite_to_program(sink:UOp, ren:Renderer) -> UOp:
   full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  sink = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=ren.device)))
-  return graph_rewrite(sink, pm_to_program, ctx=ren, name="linearize/render")
+  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
+  lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
+  if SPEC: type_verify(lst, program_spec)
+  return lst

tinygrad/codegen/opt/search.py

@@ -1,11 +1,13 @@
 import functools, math, time, multiprocessing, traceback, signal, atexit
-from tinygrad.uop.ops import sym_infer, AxisType, pyrender, UOp
+from dataclasses import replace
+from tinygrad.uop.ops import sym_infer, AxisType, pyrender
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap
 from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
-from tinygrad.engine.realize import CompiledRunner, get_program, Estimates
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.renderer import ProgramSpec
 from tinygrad.codegen.opt.postrange import Scheduler
 
 actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)]
@@ -32,22 +34,19 @@ def get_test_global_size(global_size, max_global_size, var_vals):
         break
   return test_global_size, input_size / prod(test_global_size)
 
-def _time_program(p:UOp, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
+def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
                   allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
   factor = 1
-  global_size, local_size = p.sizes
-  if allow_test_size and global_size is not None and max_global_size is not None:
-    test_global_size, factor = get_test_global_size(global_size, max_global_size, var_vals)
-    # NOTE: we can't modify p.sizes directly, but optimize_local_size doesn't run in BEAM so this is ok
+  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(p, precompiled=lib)
   except AssertionError: return [math.inf] * cnt
   tms = []
   input_bufs = [rawbufs[i] for i in car.p.globals]
   for _ in range(cnt):
     if clear_l2:
-      p_device = p.device
-      assert isinstance(p_device, str), f"PROGRAM device must be a string, not {type(p_device)}"
-      if hasattr(dev:=Device[p_device], 'invalidate_caches'): dev.invalidate_caches()
+      if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
       else:
         with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
     tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
@@ -59,7 +58,7 @@ def timeout_handler(signum, frame):
   if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT")
   raise TimeoutException()
 
-def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[UOp, bytes, float]|None]:
+def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
   if hasattr(signal, "alarm"):
     signal.signal(getattr(signal, 'SIGALRM'), timeout_handler)
     # set timeout
@@ -67,12 +66,12 @@ def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[
   ret = None
   try:
     p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren)
-    uops = list(p.src[2].src)  # LINEAR is src[2]
-    if len(uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
-      if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(uops)=}, {uops_max=}")
+    assert p.uops is not None, "uop list wasn't generated?"
+    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
+      if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}")
       raise RuntimeError("too many uops")
     st = time.perf_counter()
-    prog = compiler.compile(p.src[3].arg)  # SOURCE is src[3]
+    prog = compiler.compile(p.src)
     et = time.perf_counter() - st
     ret = (p, prog, et)
   except RuntimeError:
@@ -155,8 +154,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
         p, lib, compile_et = proc
         if lib in seen_libs: continue
         # filter out kernels that use 1000x more compute than the smallest
-        estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-        least_compute_ops = min(this_compute_ops:=sym_infer(estimates.ops, var_vals), least_compute_ops)
+        least_compute_ops = min(this_compute_ops:=sym_infer(p.estimates.ops, var_vals), least_compute_ops)
         if least_compute_ops*1000 < this_compute_ops:
           if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too much compute. {this_compute_ops} when least is {least_compute_ops}")
           continue
@@ -169,7 +167,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
           raise
         timed.append((candidates[i], min(tms)))
         if BEAM_DEBUG > 1:
-          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(p.src[2].src):5d} uops",
+          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(unwrap(p.uops)):5d} uops",
                 f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed[-1][1], w=12)} run",
                 f"      {len(timed):4d}/{len(candidates):4d}         {timed[-1][0].colored_shape()}")
         elif DEBUG >= 2:

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/jit.py

@@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, JIT_BATCH_SIZE, dedup, partition, unwrap
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType
-from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops, sint
+from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
 from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
 from tinygrad.engine.memory import _internal_memory_planner
 from tinygrad.nn.state import get_parameters
@@ -78,18 +78,13 @@ class GraphRunner(Runner):
     self.input_replace:dict[tuple[int, int], int] = get_input_replace(jit_cache, input_rawbuffers)
     self.var_vals_replace:dict[int, list[tuple[int, int]]] = {}
     self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {}
-    self.launch_dims_base:dict[int, tuple[tuple[sint, ...], tuple[sint, ...]]] = {}
+    self.launch_dims_base:dict[int, tuple[tuple[int, ...], tuple[int, ...]]] = {}
 
-    def is_sym_dim(dim) -> bool: return dim is not None and not all(isinstance(d, (int, float)) for d in dim)
+    def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim)
 
     self.vars = sorted(var_vals.keys())
-    sym_dims: list[tuple[sint, ...]] = []
-    for ji in jit_cache:
-      if isinstance(ji.prg, CompiledRunner):
-        global_size, local_size = ji.prg.p.sizes
-        if local_size is not None and is_sym_dim(local_size): sym_dims.append(tuple(local_size))
-        if global_size is not None and is_sym_dim(global_size): sym_dims.append(tuple(global_size))
-    self.symbolic_dims = dedup(sym_dims)
+    self.symbolic_dims = dedup([tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.local_size) and is_sym_dim(d)] +
+                               [tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.global_size) and is_sym_dim(d)])
     def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None
 
     estimates = Estimates()
@@ -97,15 +92,13 @@ class GraphRunner(Runner):
       assert ji.prg is not None
       estimates += ji.prg.estimates
       if isinstance(ji.prg, CompiledRunner):
-        prg_vars = ji.prg.p.variables()
-        if prg_vars: self.var_vals_replace[j] = [(i, self.vars.index(v.expr)) for i, v in enumerate(prg_vars) if v.expr not in ji.fixedvars]
+        if ji.prg.p.vars: self.var_vals_replace[j] = [(i, self.vars.index(v.expr)) for i, v in enumerate(ji.prg.p.vars) if v.expr not in ji.fixedvars]
 
-        global_size, local_size = ji.prg.p.sizes
-        global_dim_idx, local_dim_idx = find_symbolic_dim(global_size), find_symbolic_dim(local_size)
+        global_dim_idx, local_dim_idx = find_symbolic_dim(ji.prg.p.global_size), find_symbolic_dim(ji.prg.p.local_size)
         if global_dim_idx is not None or local_dim_idx is not None:
           self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx)
-          assert global_size is not None and local_size is not None
-          self.launch_dims_base[j] = (tuple(global_size), tuple(local_size))
+          assert ji.prg.p.global_size is not None and ji.prg.p.local_size is not None
+          self.launch_dims_base[j] = (tuple(ji.prg.p.global_size), tuple(ji.prg.p.local_size))
 
     # used in MultiGraphRunner. the ints are id() of _bufs
     self.w_dependency_map: dict[int, Any] = {}

tinygrad/engine/realize.py

@@ -1,29 +1,28 @@
 from typing import cast, Callable
 import time, pprint, random, itertools, math
-from dataclasses import dataclass, field
+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, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
-from tinygrad.helpers import unwrap, to_function_name
+from tinygrad.helpers import unwrap
 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, Estimates
-from tinygrad.codegen import full_rewrite_to_program
+from tinygrad.renderer import Renderer, ProgramSpec, Estimates
+from tinygrad.codegen import full_rewrite
 from tinygrad.codegen.opt import Opt
 
 # **************** Program Creation ****************
 
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.src[0].arg.name, (to_function_name(ret.src[0].arg.name), ret.src[0]), ret=ret),
-                replay=True)
-def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> UOp:
+@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, opts:list[Opt]|tuple[Opt, ...]|None=None) -> ProgramSpec:
   """
-  Transform an AST into a PROGRAM UOp. May trigger BEAM search.
+  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
 
   Returns:
-    The PROGRAM UOp with structure (SINK, DEVICE, LINEAR, SOURCE).
+    The ProgramSpec of the program.
   """
 
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
@@ -33,14 +32,22 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> UOp:
   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)
+  # print and render
+  if DEBUG >= 6: print_uops(uops)
+  src = renderer.render(uops)
 
-  # print
-  if DEBUG >= 6: print_uops(list(prg.src[2].src))  # LINEAR is src[2]
-
-  return prg
+  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)
 
 # **************** Runners ****************
 
@@ -69,36 +76,28 @@ def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffe
   return ret[1]
 
 class CompiledRunner(Runner):
-  def __init__(self, p:UOp, precompiled:bytes|None=None, prg=None):
-    assert p.op is Ops.PROGRAM, f"CompiledRunner requires PROGRAM UOp, not {p.op}"
-    self.p:UOp = p
-    dev = p.device
-    assert isinstance(dev, str), f"PROGRAM device must be a string, not {type(dev)}"
-    name = p.src[0].arg.name
-    src = p.src[3].arg
-    function_name = to_function_name(name)
-    if DEBUG >= 3: print(p.src[0].arg.applied_opts)
-    if DEBUG >= 4: print(src)
+  def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None):
+    if DEBUG >= 3: print(p.applied_opts)
+    if DEBUG >= 4: print(p.src)
+    self.p:ProgramSpec = p
     if precompiled is not None: self.lib = precompiled
     else:
-      with cpu_profile(TracingKey(f"compile {name}", (function_name,)), "TINY"):
-        self.lib = Device[dev].compiler.compile_cached(src)
-    if DEBUG >= 7: Device[dev].compiler.disassemble(self.lib)
-    self._prg = Device[dev].runtime(function_name, self.lib) if prg is None else prg
-    super().__init__(name, dev, Estimates.from_uops(list(p.src[2].src), ignore_indexing=True))
+      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, 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 = {}
-    dev = self.p.device
-    assert isinstance(dev, str), f"PROGRAM device must be a string, not {type(dev)}"
-    has_local = Device[dev].renderer.has_local
+    has_local = Device[self.p.device].renderer.has_local
     global_size, local_size = self.p.launch_dims(var_vals)
-    sym_global_size, sym_local_size = self.p.sizes
-    if has_local and global_size is not None and local_size is None and sym_global_size is not None and all_int(sym_global_size):
+    if has_local and global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
       local_size = optimize_local_size(self._prg, global_size, rawbufs)
-      global_size = [g//l if g%l == 0 else int(g/l) for g,l in zip(global_size, local_size)]
+      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'] = tuple(global_size)
@@ -106,7 +105,7 @@ class CompiledRunner(Runner):
     if local_size:
       lra['local_size'] = tuple(local_size)
       assert len(local_size) == 3, "local size must have len 3"
-    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k.expr] for k in self.p.variables()), wait=wait)
+    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k.expr] for k in self.p.vars), wait=wait)
 
 class ViewOp(Runner):
   def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
@@ -163,14 +162,10 @@ 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):
-    # update device in PROGRAM UOp
-    new_p = bret.p.replace(src=(bret.p.src[0], UOp(Ops.DEVICE, arg=device), *bret.p.src[2:]))
-    method_cache[ckey] = ret = CompiledRunner(new_p, bret.lib)
+    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib)
   else:
-    prg = get_program(ast, Device[device].renderer)
-    # update device in PROGRAM UOp to match the actual device
-    prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=device), *prg.src[2:]))
-    method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(prg)
+    prg: ProgramSpec = get_program(ast, Device[device].renderer)
+    method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))
   return ret
 
 # **************** lowering functions ****************
@@ -219,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/engine/schedule.py

@@ -90,29 +90,23 @@ from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map
 
-def replace_input_buffer(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
-  if (ret:=ctx[0].get(b, None)) is None:
+def replace_input_buffer(ctx:dict[UOp, UOp], b:UOp):
+  if (ret:=ctx.get(b, None)) is None:
     if b.op is Ops.BUFFER:
-      ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx[0])), b.src[1]))
+      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
     else:
       # TODO: flip args in CONST
       assert b.op is Ops.CONST
-      ctx[0][b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx[0]))))
+      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
   return ret
 
-def strip_bind(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
-  var, val = b.src[0], b.src[1].arg
-  assert var.expr not in ctx[1] or ctx[1][var.expr] == val, f"bind mismatch on {var}, {ctx[1][var.expr]} != {val}"
-  ctx[1][var.expr] = val
-  return ctx[0].setdefault(b, b.replace(src=(b.src[0],)))
-
 pm_pre_sched_cache = PatternMatcher([
   # replace input buffers
   (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
   # remove unique consts
   (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
   # strip value from BIND for cache key normalization, so different values hit same cache
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), strip_bind),
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), lambda ctx,b: ctx.setdefault(b, b.replace(src=(b.src[0],)))),
 ])
 
 def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
@@ -135,10 +129,9 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
   # big_sink srcs are all the Tensors
   st = time.perf_counter()
 
-  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key, extract var_vals
+  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key
   input_buffers: dict[UOp, UOp] = {}
-  var_vals: dict[str, int] = {}
-  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=(input_buffers, var_vals), name="rewrite for sched cache")
+  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
   sched_cache_key = big_sink_cache.key
 
   if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
@@ -146,7 +139,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
     if SPEC: type_verify(big_sink, tensor_spec)
 
     # hack to preserve metadata
-    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx=({}, {}), name="preserve metadata")
+    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata")
 
     # tensor map is what we return
     tensor_map: dict[UOp, UOp] = {}
@@ -198,8 +191,17 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
       schedule.append(ExecItem(si.ast, list(ubufs), si.metadata, si.fixedvars))
   with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
 
+  # extract var_vals from BINDs that were stripped (only if there are kernels)
+  var_vals: dict[str, int] = {}
+  if schedule:
+    for u in input_buffers:
+      if u.op is Ops.BIND:
+        var, val = u.unbind()
+        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+        var_vals[var.expr] = val
+
   if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
     print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
           f" | {' cache hit' if sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
           f" | {len(UOpMetaClass.ucache)} uops in cache")
-  return tensor_map, schedule, var_vals if schedule else {}
+  return tensor_map, schedule, var_vals

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

@@ -1,10 +1,12 @@
 from __future__ import annotations
 from typing import Callable, cast
-from dataclasses import dataclass
-from tinygrad.helpers import prod
-from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, GroupOp, PatternMatcher
+import functools
+from dataclasses import dataclass, field
+from tinygrad.helpers import to_function_name, dedup, prod
+from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.codegen.opt.tc import TensorCore
+from tinygrad.codegen.opt import Opt
 
 @dataclass(frozen=True)
 class Estimates:
@@ -55,6 +57,62 @@ class Estimates:
       elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
     return Estimates(flops, lds, sum(mem.values()))
 
+@dataclass
+class ProgramSpec:
+  name:str
+  src:str
+  device:str
+  ast:UOp  # save the base ast (this is method cache key)
+  uops:list[UOp]|None=None
+
+  # filled in from uops (if we have uops)
+  global_size:list[int]|None=None
+  local_size:list[int]|None=None
+  vars:list[Variable]=field(default_factory=list)
+  globals:list[int]=field(default_factory=list)
+  outs:list[int]=field(default_factory=list)
+  ins:list[int]=field(default_factory=list)
+  _ran_post_init:bool=False  # NOTE: this is needed if you call replace on the Program
+
+  def __post_init__(self):
+    if not self._ran_post_init and self.uops is not None:
+      # single pass through the uops
+      for u in self.uops:
+        if u.op is Ops.DEFINE_VAR: self.vars.append(u)
+        if u.op is Ops.DEFINE_GLOBAL: self.globals.append(u.arg)
+        if u.op in (Ops.STORE, Ops.LOAD):
+          if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
+            if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: (self.outs if u.op is Ops.STORE else self.ins).append(buf.arg)
+          # TODO: can else happen?
+        if u.op is Ops.SPECIAL:
+          # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
+          if u.arg[0] == 'i': self.local_size = None
+          special_size = self.local_size if u.arg[0] == 'l' else self.global_size
+          # TODO: this cast is wrong, u.src[0].ssimplify() can be sint
+          if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify())
+      self.vars = sorted(self.vars, key=lambda v: v.arg)
+      self.outs = sorted(dedup(self.outs))
+      self.ins = sorted(dedup(self.ins))
+      self._ran_post_init = True
+
+  @functools.cached_property
+  def estimates(self) -> Estimates:
+    return Estimates() if self.uops is None else Estimates.from_uops(self.uops, ignore_indexing=True)
+
+  @functools.cached_property
+  def function_name(self) -> str: return to_function_name(self.name)
+
+  @property
+  def applied_opts(self) -> tuple[Opt, ...]|None:
+    if self.uops is None: return None
+    assert self.uops[-1].op is Ops.SINK, self.uops[-1].op
+    return self.uops[-1].arg.applied_opts
+
+  def launch_dims(self, var_vals:dict[str, int]):
+    global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else None
+    local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else None
+    return global_size, local_size
+
 class Renderer:
   device: str = ""
   suffix: str = ""

tinygrad/runtime/graph/cuda.py

@@ -22,14 +22,12 @@ class CUDAGraph(MultiGraphRunner):
     for j,ji in enumerate(jit_cache):
       if isinstance(ji.prg, CompiledRunner):
         global_size, local_size = ji.prg.p.launch_dims(var_vals)
-        assert global_size is not None and local_size is not None
 
         new_node = cuda.CUgraphNode()
         deps = self._access_resources([x.base for x in ji.bufs if x is not None], ji.prg.p.outs, new_dependency=new_node)
         c_deps = (cuda.CUgraphNode*len(deps))(*deps) if deps else None
 
-        prg_vars = ji.prg.p.variables()
-        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals.get(x.expr, ji.fixedvars.get(x.expr)) for x in prg_vars])
+        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals.get(x.expr, ji.fixedvars.get(x.expr)) for x in ji.prg.p.vars])
         kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(ji.prg._prg.prg, *global_size, *local_size, 0, None, vargs)
         check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(deps), ctypes.byref(kern_params)))
 

tinygrad/runtime/graph/hcq.py

@@ -39,7 +39,7 @@ class HCQGraph(MultiGraphRunner):
       if not isinstance(ji.prg, CompiledRunner): continue
 
       argsbuf = self.kernargs_bufs[ji.prg.dev].offset(kargs_alloc[ji.prg.dev].alloc(ji.prg._prg.kernargs_alloc_size, 16))
-      self.ji_args[j] = ji.prg._prg.fill_kernargs(self.hcq_bufs[j], ji.prg.p.prog_vars(), argsbuf)
+      self.ji_args[j] = ji.prg._prg.fill_kernargs(self.hcq_bufs[j], ji.prg.p.vars, argsbuf)
 
     # Schedule Dependencies.
     # There are two types of queues on each device: copy and compute. Both must synchronize with all external operations before launching any
@@ -159,8 +159,7 @@ class HCQGraph(MultiGraphRunner):
 
       # Encode main commands based on ji type.
       if isinstance(ji.prg, CompiledRunner):
-        global_size, local_size = ji.prg.p.sizes
-        enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(global_size or (1,1,1)), tuple(local_size or (1,1,1)))
+        enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(ji.prg.p.global_size or (1,1,1)), tuple(ji.prg.p.local_size or (1,1,1)))
       elif isinstance(ji.prg, (BufferXfer, BufferCopy)):
         dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
         for bufid, src in enumerate(cast(list[Buffer], ji.bufs)):

tinygrad/runtime/graph/metal.py

@@ -42,11 +42,9 @@ class MetalGraph(GraphRunner):
         if b is not None and b not in input_rawbuffers:
           icb_command.setKernelBuffer_offset_atIndex(b._buf.buf, b._buf.offset, i)
           all_resources.append(b._buf.buf)
-      for i,v in enumerate(prg.p.variables()):
-        icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
+      for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
 
       global_size, local_size = prg.p.launch_dims(var_vals)
-      assert global_size is not None and local_size is not None
       icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
       icb_command.setBarrier()
 

tinygrad/runtime/ops_amd.py

@@ -18,7 +18,6 @@ from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.memory import AddrSpace
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
 SQTT = ContextVar("SQTT", abs(VIZ.value)>=2)
@@ -860,7 +859,7 @@ class USBIface(PCIIface):
     self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x800
 
   def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
+    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], system=True, uncached=True)
     return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:

tinygrad/runtime/support/am/amdev.py

@@ -4,7 +4,7 @@ from tinygrad.helpers import mv_address, getenv, DEBUG, fetch
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
-from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
+from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
 from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
 
@@ -120,11 +120,10 @@ class AMFirmware:
 class AMPageTableEntry:
   def __init__(self, adev, paddr, lv): self.adev, self.paddr, self.lv, self.entries = adev, paddr, lv, adev.vram.view(paddr, 0x1000, fmt='Q')
 
-  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, aspace=AddrSpace.PHYS, snooped=False, frag=0, valid=True):
-    is_sys = aspace is AddrSpace.SYS
-    if aspace is AddrSpace.PHYS: paddr = self.adev.paddr2xgmi(paddr)
+  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, system=False, snooped=False, frag=0, valid=True):
+    if not system: paddr = self.adev.paddr2xgmi(paddr)
     assert paddr & self.adev.gmc.address_space_mask == paddr, f"Invalid physical address {paddr:#x}"
-    self.entries[entry_id] = self.adev.gmc.get_pte_flags(self.lv, table, frag, uncached, is_sys, snooped, valid) | (paddr & 0x0000FFFFFFFFF000)
+    self.entries[entry_id] = self.adev.gmc.get_pte_flags(self.lv, table, frag, uncached, system, snooped, valid) | (paddr & 0x0000FFFFFFFFF000)
 
   def entry(self, entry_id:int) -> int: return self.entries[entry_id]
   def valid(self, entry_id:int) -> bool: return (self.entries[entry_id] & am.AMDGPU_PTE_VALID) != 0

tinygrad/runtime/support/am/ip.py

@@ -3,7 +3,6 @@ from typing import Literal
 from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG, wait_cond, pad_bytes
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.amd import import_soc
-from tinygrad.runtime.support.memory import AddrSpace
 
 class AM_IP:
   def __init__(self, adev): self.adev = adev
@@ -469,7 +468,7 @@ class AM_PSP(AM_IP):
     msg1_region = next((reg for reg in self.adev.dma_regions or [] if reg[1].nbytes >= (512 << 10)), None)
     if msg1_region is not None:
       self.msg1_addr, self.msg1_view = self.adev.mm.alloc_vaddr(size=msg1_region[1].nbytes, align=am.PSP_1_MEG), msg1_region[1]
-      self.adev.mm.map_range(self.msg1_addr, msg1_region[1].nbytes, [(msg1_region[0],msg1_region[1].nbytes)], AddrSpace.SYS, uncached=True, boot=True)
+      self.adev.mm.map_range(self.msg1_addr, msg1_region[1].nbytes, [(msg1_region[0], msg1_region[1].nbytes)], system=True, uncached=True, boot=True)
     else:
       self.msg1_paddr = self.adev.mm.palloc(am.PSP_1_MEG, align=am.PSP_1_MEG, zero=False, boot=True)
       self.msg1_addr, self.msg1_view = self.adev.paddr2mc(self.msg1_paddr), self.adev.vram.view(self.msg1_paddr, am.PSP_1_MEG, 'B')

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/runtime/support/memory.py

@@ -1,4 +1,4 @@
-import collections, functools, dataclasses, enum
+import collections, functools, dataclasses
 from typing import Any, ClassVar
 from tinygrad.helpers import round_up, getenv
 
@@ -107,10 +107,8 @@ class TLSFAllocator:
 
 # Memory Managment
 
-class AddrSpace(enum.Enum): PHYS = enum.auto(); SYS = enum.auto(); PEER = enum.auto() # noqa: E702
-
 @dataclasses.dataclass(frozen=True)
-class VirtMapping: va_addr:int; size:int; paddrs:list[tuple[int, int]]; aspace:AddrSpace; uncached:bool=False; snooped:bool=False # noqa: E702
+class VirtMapping: va_addr:int; size:int; paddrs:list[tuple[int, int]]; uncached:bool=False; system:bool=False; snooped:bool=False # noqa: E702
 
 class PageTableTraverseContext:
   def __init__(self, dev, pt, vaddr, create_pts=False, free_pts=False, boot=False):
@@ -192,7 +190,7 @@ class MemoryManager:
     ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, create_pts=True)
     for _ in ctx.next(size, paddr=0): return [pt for pt, _, _ in ctx.pt_stack]
 
-  def map_range(self, vaddr:int, size:int, paddrs:list[tuple[int, int]], aspace:AddrSpace, uncached=False, snooped=False, boot=False) -> VirtMapping:
+  def map_range(self, vaddr:int, size:int, paddrs:list[tuple[int, int]], uncached=False, system=False, snooped=False, boot=False) -> VirtMapping:
     if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: mapping {vaddr=:#x} ({size=:#x})")
 
     assert size == sum(p[1] for p in paddrs), f"Size mismatch {size=} {sum(p[1] for p in paddrs)=}"
@@ -202,11 +200,11 @@ class MemoryManager:
       for off, pt, pte_idx, pte_cnt, pte_covers in ctx.next(psize, paddr=paddr):
         for pte_off in range(pte_cnt):
           assert not pt.valid(pte_idx + pte_off), f"PTE already mapped: {pt.entry(pte_idx + pte_off):#x}"
-          pt.set_entry(pte_idx + pte_off, paddr + off + pte_off * pte_covers, uncached=uncached, aspace=aspace, snooped=snooped,
+          pt.set_entry(pte_idx + pte_off, paddr + off + pte_off * pte_covers, uncached=uncached, system=system, snooped=snooped,
                        frag=self._frag_size(ctx.vaddr+off, pte_cnt * pte_covers), valid=True)
 
     self.on_range_mapped()
-    return VirtMapping(vaddr, size, paddrs, aspace=aspace, uncached=uncached, snooped=snooped)
+    return VirtMapping(vaddr, size, paddrs, uncached=uncached, system=system, snooped=snooped)
 
   def unmap_range(self, vaddr:int, size:int):
     if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: unmapping {vaddr=:#x} ({size=:#x})")
@@ -245,7 +243,7 @@ class MemoryManager:
           continue
         rem_size -= self.palloc_ranges[nxt_range][0]
 
-    return self.map_range(va, size, paddrs, aspace=AddrSpace.PHYS, uncached=uncached)
+    return self.map_range(va, size, paddrs, uncached=uncached)
 
   def vfree(self, vm:VirtMapping):
     assert self.va_allocator is not None, "must be set it"

tinygrad/runtime/support/nv/nvdev.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes, time, functools, re, gzip, struct
 from tinygrad.helpers import getenv, DEBUG, fetch, getbits
-from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
+from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
 from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
 from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
 
@@ -33,9 +33,9 @@ class NVPageTableEntry:
 
   def _is_dual_pde(self) -> bool: return self.lv == self.nvdev.mm.level_cnt - 2
 
-  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, aspace=AddrSpace.PHYS, snooped=False, frag=0, valid=True):
+  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, system=False, snooped=False, frag=0, valid=True):
     if not table:
-      x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if aspace is AddrSpace.SYS else 0, kind=6,
+      x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if system else 0, kind=6,
         **({'pcf': int(uncached)} if self.nvdev.mmu_ver == 3 else {'vol': uncached}))
     else:
       pde = self.nvdev.dual_pde_t if self._is_dual_pde() else self.nvdev.pde_t

tinygrad/runtime/support/system.py

@@ -3,7 +3,7 @@ from typing import cast, ClassVar
 from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv
 from tinygrad.runtime.autogen import libc, vfio, pci
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
-from tinygrad.runtime.support.memory import MemoryManager, VirtMapping, AddrSpace
+from tinygrad.runtime.support.memory import MemoryManager, VirtMapping
 from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface
 
 MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400
@@ -262,7 +262,7 @@ class LNXPCIIfaceBase:
     if should_use_sysmem:
       vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
       memview, paddrs = System.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
-      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, snooped=True, uncached=True)
+      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], system=True, snooped=True, uncached=True)
       return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
 
     mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access)
@@ -271,23 +271,19 @@ class LNXPCIIfaceBase:
 
   def free(self, b:HCQBuffer):
     for dev in b.mapped_devs[1:]: dev.iface.dev_impl.mm.unmap_range(b.va_addr, b.size)
-    if b.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(b.meta.mapping)
+    if not b.meta.mapping.system: self.dev_impl.mm.vfree(b.meta.mapping)
     if b.owner == self.dev and b.meta.has_cpu_mapping and not OSX: FileIOInterface.munmap(b.va_addr, b.size)
 
   def map(self, b:HCQBuffer):
     if b.owner is not None and b.owner._is_cpu():
       System.lock_memory(cast(int, b.va_addr), b.size)
-      paddrs, aspace = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], AddrSpace.SYS
-      snooped, uncached = True, True
+      paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, True
     elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, LNXPCIIfaceBase):
-      snooped, uncached = True, b.meta.mapping.uncached
-      if b.meta.mapping.aspace is AddrSpace.SYS: paddrs, aspace = b.meta.mapping.paddrs, AddrSpace.SYS
-      elif hasattr(ifa.dev_impl, 'paddr2xgmi') and ifa.dev_impl.gmc.xgmi_seg_sz > 0:
-        paddrs, aspace = [(ifa.dev_impl.paddr2xgmi(p), sz) for p, sz in b.meta.mapping.paddrs], AddrSpace.PEER
-      else: paddrs, aspace = [(p + ifa.p2p_base_addr, sz) for p, sz in b.meta.mapping.paddrs], AddrSpace.SYS
+      paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
+      snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached
     else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
-    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
+    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, system=True, snooped=snooped, uncached=uncached)
 
 class APLPCIIfaceBase(LNXPCIIfaceBase):
   def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):

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
-  # LINEAR is a list of UOps, SOURCE has a str arg that's human readable
-  PROGRAM = auto(); LINEAR = auto(); SOURCE = 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.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL:
         return None
 
       case Ops.INDEX:
@@ -609,17 +608,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
-  @staticmethod
-  def new_program(name:str, src:str, device:str, ast:UOp, uops:list[UOp]):
-    """Create a PROGRAM UOp from raw components."""
-    sink = ast.replace(arg=KernelInfo(name=name)) if ast.arg is None else ast.replace(arg=ast.arg.replace(name=name))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=tuple(uops)), UOp(Ops.SOURCE, arg=src)))
   @property
   def device(self) -> str|tuple[str, ...]: return unwrap(self._device)
   @recursive_property
   def _device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.DEVICE: return self.arg
-    if self.op is Ops.PROGRAM: return self.src[1].arg  # PROGRAM src[1] is DEVICE
     if self.op is Ops.BUFFERIZE: return self.arg.device
     if self.op is Ops.AFTER: return self.src[0]._device
     if self.op is Ops.MSELECT:
@@ -630,104 +623,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     for x in self.src:
       if x._device is not None: return x._device
     return None
-
-  # *** PROGRAM UOp properties ***
-
-  @property
-  def uops(self) -> list[UOp]:
-    """Linearized uops list. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"uops only valid for PROGRAM, not {self.op}"
-    return list(self.src[2].src)
-
-  @property
-  def name(self) -> str:
-    """Kernel name. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"name only valid for PROGRAM, not {self.op}"
-    return self.src[0].arg.name
-
-  @property
-  def applied_opts(self):
-    """Applied optimizations. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"applied_opts only valid for PROGRAM, not {self.op}"
-    return self.src[0].arg.applied_opts
-
-  @functools.cached_property
-  def estimates(self):
-    """Estimates for this program. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"estimates only valid for PROGRAM, not {self.op}"
-    from tinygrad.renderer import Estimates
-    return Estimates.from_uops(self.uops, ignore_indexing=True)
-
-  @property
-  def globals(self) -> list[int]:
-    """DEFINE_GLOBAL arg indices from linearized uops. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"globals only valid for PROGRAM, not {self.op}"
-    return [u.arg for u in self.src[2].src if u.op is Ops.DEFINE_GLOBAL]
-
-  @property
-  def outs(self) -> list[int]:
-    """Buffer indices written to (STORE). Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"outs only valid for PROGRAM, not {self.op}"
-    ret = []
-    for u in self.src[2].src:
-      if u.op is Ops.STORE:
-        idx = u.src[0]
-        if idx.op is Ops.CAST: idx = idx.src[0]
-        if idx.op is Ops.INDEX and idx.src[0].op is Ops.DEFINE_GLOBAL: ret.append(idx.src[0].arg)
-    return sorted(set(ret))
-
-  @property
-  def ins(self) -> list[int]:
-    """Buffer indices read from (LOAD). Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"ins only valid for PROGRAM, not {self.op}"
-    ret = []
-    for u in self.src[2].src:
-      if u.op is Ops.LOAD:
-        idx = u.src[0]
-        if idx.op is Ops.CAST: idx = idx.src[0]
-        if idx.op is Ops.INDEX and idx.src[0].op is Ops.DEFINE_GLOBAL: ret.append(idx.src[0].arg)
-    return sorted(set(ret))
-
-  @functools.cached_property
-  def sizes(self) -> tuple[list[sint]|None, list[sint]|None]:
-    """Get (global_size, local_size) which may contain symbolic values. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"sizes only valid for PROGRAM, not {self.op}"
-    from tinygrad.device import Device
-    dev = self.device
-    assert isinstance(dev, str), f"PROGRAM device must be a string, not {type(dev)}"
-    ren = Device[dev].renderer
-    global_size:list[sint]|None = [1,1,1] if ren.has_local or ren.has_threads else None
-    local_size:list[sint]|None = [1,1,1] if ren.has_local else None
-    for u in self.src[2].src:
-      if u.op is Ops.SPECIAL:
-        if u.arg[0] == 'i': local_size = None
-        special_size = local_size if u.arg[0] == 'l' else global_size
-        if special_size is not None: special_size[int(u.arg[-1])] = cast(sint, u.src[0].ssimplify())
-    return global_size, local_size
-
-  def launch_dims(self, var_vals:dict[str, int]) -> tuple[list[int]|None, list[int]|None]:
-    """Resolve global/local sizes to concrete ints. Only valid for PROGRAM."""
-    global_size, local_size = self.sizes
-    global_ret = [sym_infer(sz, var_vals) for sz in global_size] if global_size is not None else None
-    local_ret = [sym_infer(sz, var_vals) for sz in local_size] if local_size is not None else None
-    return global_ret, local_ret
-
-  @property
-  def global_size(self) -> list[sint]|None:
-    """Global size (may be symbolic). Only valid for PROGRAM."""
-    return self.sizes[0]
-
-  @property
-  def local_size(self) -> list[sint]|None:
-    """Local size (may be symbolic). Only valid for PROGRAM."""
-    return self.sizes[1]
-
-  def prog_vars(self) -> list:
-    """Variables list for this program. Only valid for PROGRAM."""
-    assert self.op is Ops.PROGRAM, f"prog_vars only valid for PROGRAM, not {self.op}"
-    # Get variables from the linearized uops
-    linear_uops = self.src[2]
-    return linear_uops.variables()
   @property
   def buf_uop(self) -> UOp:
     if self.op is Ops.BUFFER: return self
@@ -804,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]
@@ -817,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 ***
@@ -1425,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,14 +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 (SINK, DEVICE, LINEAR, SOURCE)
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
-  # codegen: standalone LINEAR/SOURCE
-  (UPat(Ops.LINEAR, dtypes.void), lambda: True),
-  (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
-
   # temp VECTORIZE/INDEX during rewrite have the wrong dtype
   (UPat(Ops.VECTORIZE), lambda: True),
   (UPat(Ops.INDEX), lambda: True),

tinygrad/viz/serve.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import printable, TCPServerWithReuse, HTTPRequestHandler
 from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, pyrender
 from tinygrad.uop.ops import print_uops, range_start, multirange_str
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device, ProfileProgramEvent
-from tinygrad.renderer import Estimates
+from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
@@ -39,7 +39,7 @@ def get_rewrites(t:RewriteTrace) -> list[dict]:
   for i,(k,v) in enumerate(zip(t.keys, t.rewrites)):
     steps = [create_step(s.name, ("/graph-rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc),
                          trace=k.tb if j==0 else None, depth=s.depth) for j,s in enumerate(v)]
-    if isinstance(k.ret, UOp) and k.ret.op is Ops.PROGRAM:
+    if isinstance(k.ret, ProgramSpec):
       steps.append(create_step("View UOp List", ("/uops", i, len(steps)), k.ret))
       steps.append(create_step("View Program", ("/code", i, len(steps)), k.ret))
       steps.append(create_step("View Disassembly", ("/asm", i, len(steps)), k.ret))
@@ -161,10 +161,9 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
     name, fmt, key = e.name, [], None
     if (ref:=ref_map.get(name)) is not None:
       name = ctxs[ref]["name"]
-      if isinstance(p:=trace.keys[ref].ret, UOp) and p.op is Ops.PROGRAM and (ei:=exec_points.get(p.src[0].arg.name)) is not None:
-        estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
-        flops = sym_infer(estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6)
-        membw, ldsbw = sym_infer(estimates.mem, var_vals)/t, sym_infer(estimates.lds, var_vals)/t
+      if isinstance(p:=trace.keys[ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
+        flops = sym_infer(p.estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6)
+        membw, ldsbw = sym_infer(p.estimates.mem, var_vals)/t, sym_infer(p.estimates.lds, var_vals)/t
         fmt = [f"{flops*1e-9:.0f} GFLOPS" if flops < 1e14 else f"{flops*1e-12:.0f} TFLOPS",
               (f"{membw*1e-9:.0f} GB/s" if membw < 1e13 else f"{membw*1e-12:.0f} TB/s")+" mem",
               (f"{ldsbw*1e-9:.0f} GB/s" if ldsbw < 1e15 else f"{ldsbw*1e-12:.0f} TB/s")+" lds"]
@@ -426,12 +425,10 @@ def get_render(i:int, j:int, fmt:str) -> dict:
   data = ctxs[i]["steps"][j]["data"]
   if fmt == "graph-rewrites": return {"value":get_full_rewrite(trace.rewrites[i][j]), "content_type":"text/event-stream"}
   if fmt == "uops": return {"src":get_stdout(lambda: print_uops(data.uops or [])), "lang":"txt"}
-  # PROGRAM UOp: src[3].arg is source code
-  source_code = data.src[3].arg
-  if fmt == "code": return {"src":source_code, "lang":"cpp"}
+  if fmt == "code": return {"src":data.src, "lang":"cpp"}
   if fmt == "asm":
     compiler = Device[data.device].compiler
-    disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(source_code)))
+    disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(data.src)))
     ret:dict = {"src":disasm_str}
     if data.device.startswith("AMD"):
       with soft_err(lambda err: ret.update(err)):