one compiler path

* rewrites for renderer and compiler

* full_rewrite_to_program

* fix pre-commit

* compiler passed into get_program

* no pkl compiler

* lib on program spec

* fix spec

* fix test

* no device

* compiler_device

* nm

* fix nir

* fix

* simplest

* fix tests

* revert

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 (specified by a ShapeTracker). 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. Inputs to a base can be either base or view, inputs to a view can only be a single base.
 
 ## Scheduling
 

extra/models/llama.py

@@ -245,6 +245,11 @@ 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

@@ -1,114 +0,0 @@
-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

@@ -1,129 +0,0 @@
-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

@@ -1,124 +0,0 @@
-# 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

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

extra/test_pyrender.py

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

test/external/external_test_nv.py

@@ -1,60 +0,0 @@
-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

@@ -1,61 +0,0 @@
-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

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

test/opt/test_tensor_cores.py

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

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

test/test_uops.py

@@ -7,7 +7,6 @@ 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
@@ -26,11 +25,7 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   return ret[:-1]
 
 def _uops_to_prg(uops_list):
-  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))
+  return CompiledRunner(get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer, Device.DEFAULT))
 
 def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   uops.append(UOp(uop, dtype, tuple(src), arg))

tinygrad/codegen/__init__.py

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

tinygrad/codegen/opt/search.py

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

tinygrad/engine/realize.py

@@ -2,28 +2,30 @@ from typing import cast, Callable
 import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, PROFILE, ProfilePointEvent, cpu_events, prod, Context
 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.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
-from tinygrad.codegen import full_rewrite
+from tinygrad.codegen import full_rewrite_to_program
 from tinygrad.codegen.opt import Opt
 
 # **************** Program Creation ****************
 
 @track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer, device:str|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
 
   Args:
     ast: The Ops.SINK rooted AST
     renderer: The renderer used to generate the code
+    device: The device to compile for (defaults to renderer.device)
 
   Returns:
     The ProgramSpec of the program.
   """
+  device = device or renderer.device
 
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -32,20 +34,14 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
   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)))
-  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"
+  if ast.arg is None: ast = ast.replace(arg=KernelInfo())
 
-  # print and render
-  if DEBUG >= 6: print_uops(uops)
-  src = renderer.render(uops)
+  prg = full_rewrite_to_program(ast, renderer, Device[device].compiler)
+  # SINK/LINEAR/SOURCE/BINARY
+  sink, linear, source, binary = prg.src
 
-  return ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, renderer.device, ast, uops,
+  # legacy
+  return ProgramSpec(sink.arg.name, source.arg, device, sink, list(linear.src), binary.arg,
                      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)
 
@@ -76,19 +72,16 @@ def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffe
   return ret[1]
 
 class CompiledRunner(Runner):
-  def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None):
+  def __init__(self, p:ProgramSpec, prg=None):
     if DEBUG >= 3: print(p.applied_opts)
     if DEBUG >= 4: print(p.src)
+    assert p.lib is not None, "lib must be provided"
     self.p:ProgramSpec = p
-    if precompiled is not None: self.lib = precompiled
-    else:
-      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"):
-        self.lib = Device[p.device].compiler.compile_cached(p.src)
-    if DEBUG >= 7: Device[p.device].compiler.disassemble(self.lib)
-    self._prg = Device[p.device].runtime(p.function_name, self.lib) if prg is None else prg
+    if DEBUG >= 7: Device[p.device].compiler.disassemble(p.lib)
+    self._prg = Device[p.device].runtime(p.function_name, p.lib) if prg is None else prg
     super().__init__(p.name, p.device, p.estimates)
 
-  def __reduce__(self): return self.__class__, (self.p, self.lib)
+  def __reduce__(self): return self.__class__, (self.p,)
 
   def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
     if var_vals is None: var_vals = {}
@@ -162,9 +155,9 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner:
   if cret:=method_cache.get(ckey): return cret
   bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True)
   if bret:=method_cache.get(bkey):
-    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib)
+    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device))
   else:
-    prg: ProgramSpec = get_program(ast, Device[device].renderer)
+    prg: ProgramSpec = get_program(ast, Device[device].renderer, device)
     method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))
   return ret
 

tinygrad/renderer/__init__.py

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

tinygrad/uop/__init__.py

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

tinygrad/uop/ops.py

@@ -218,7 +218,8 @@ 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.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | \
+           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
         return None
 
       case Ops.INDEX:

tinygrad/uop/spec.py

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