pickle

* add optional compiler in Renderer [pr]

* fix

* late init

* remove precompiled

* cleanup

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
 

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 list of UOps.
+Transform the optimized ast into a linearized and rendered program.
 
-::: tinygrad.codegen.full_rewrite
+::: tinygrad.codegen.get_program
     options:
         members: false
         show_labels: false

eslint.config.mjs

@@ -1,9 +0,0 @@
-import globals from "globals";
-import pluginJs from "@eslint/js";
-import pluginHtml from "eslint-plugin-html";
-
-export default [
-  {files: ["**/*.html"], plugins: {html: pluginHtml}, rules:{"max-len": ["error", {"code": 150}]}},
-  {languageOptions: {globals: globals.browser}},
-  pluginJs.configs.recommended,
-];

examples/mlperf/model_train.py

@@ -1438,28 +1438,34 @@ def train_llama3():
   iter = get_train_iter()
   i, sequences_seen = resume_ckpt, 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
-    t = time.perf_counter()
     GlobalCounters.reset()
-    loss, lr = train_step(model, tokens)
-    loss = loss.float().item()
+    if getenv("TRAIN", 1):
+      t = time.perf_counter()
+      loss, lr = train_step(model, tokens)
+      loss = loss.float().item()
+      lr = lr.item()
 
-    i += 1
-    sequences_seen += tokens.shape[0]
+      i += 1
+      sequences_seen += tokens.shape[0]
 
-    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
-    if (fname:=getenv("LOSS_FILE", "")):
-      with open(fname, "a") as f:
-        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
+      sec = time.perf_counter()-t
+      tqdm.write(
+        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, "
+        f"{GlobalCounters.global_ops * 1e-9 / sec:9.2f} GFLOPS")
 
-    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
-      tqdm.write("saving checkpoint")
-      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/llama3_{i}.safe"
-      safe_save(get_state_dict(model), fn)
+      if (fname:=getenv("LOSS_FILE", "")):
+        with open(fname, "a") as f:
+          f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
 
-      tqdm.write("saving optim checkpoint")
-      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-      safe_save(get_state_dict(scheduler), fn)
+      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
+        tqdm.write("saving checkpoint")
+        if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+        fn = f"{ckpt_dir}/llama3_{i}.safe"
+        safe_save(get_state_dict(model), fn)
+
+        tqdm.write("saving optim checkpoint")
+        fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
+        safe_save(get_state_dict(scheduler), fn)
 
     if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
       tqdm.write(f"evaluating after {sequences_seen} sequences")

extra/amdpci/am_smi.py

@@ -184,7 +184,7 @@ class SMICtx:
         if compact: return {k: temps[k] for k in ("Hotspot", "HBM") if temps.get(k, 0) != 0}
         return {k: v for k, v in temps.items() if v != 0}
       case _:
-        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_TEMP_e__enumvalues.items()
+        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.TEMP_e.items()
                       if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
         if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
         return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
@@ -193,7 +193,7 @@ class SMICtx:
     match dev.ip_ver[am.MP1_HWIP]:
       case (13,0,6): return {}
       case _:
-        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
+        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
                         if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
         return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
 

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/sqtt/roc.py

@@ -5,29 +5,7 @@ from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEven
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
 from tinygrad.runtime.autogen import llvm, rocprof
 from tinygrad.runtime.support.elf import elf_loader
-
-# to pass NULL to callbacks
-llvm.LLVMCreateDisasmCPUFeatures.argtypes = tuple(llvm.LLVMCreateDisasmCPUFeatures.argtypes[:5]) + (ctypes.c_void_p, ctypes.c_void_p)
-def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
-  llvm.LLVMInitializeAMDGPUTargetInfo()
-  llvm.LLVMInitializeAMDGPUTargetMC()
-  llvm.LLVMInitializeAMDGPUAsmParser()
-  llvm.LLVMInitializeAMDGPUDisassembler()
-  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, None, None)
-
-  image, sections, relocs = elf_loader(lib)
-  text = next((sh.header for sh in sections if sh.name == ".text"), None)
-  off, sz = unwrap(text).sh_addr, unwrap(text).sh_size
-
-  addr_table:dict[int, tuple[str, int]] = {}
-  out = ctypes.create_string_buffer(128)
-  cur_off = off
-  while cur_off < sz + off:
-    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
-    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
-    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
-    cur_off += instr_sz
-  return addr_table
+from tinygrad.viz.serve import llvm_disasm
 
 @dataclasses.dataclass(frozen=True)
 class InstExec:

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

opencode.json

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

test/external/external_benchmark_op_conv.py

@@ -1,10 +1,9 @@
 # ruff: noqa: E501 E712 F401
+from dataclasses import replace
 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.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
 from tinygrad.dtype import ImageDType, Invalid
@@ -86,16 +85,11 @@ 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
 
-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)
+ps = get_program(ast, renderer)
+cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
 
 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_hip_compiler_bug.py

@@ -1,225 +0,0 @@
-# [<buf device:HIP size:1605632 dtype:dtypes.float>, <buf device:HIP size:301506 dtype:dtypes.float>, <buf device:HIP size:9408 dtype:dtypes.float>]
-from tinygrad import Device, dtypes
-from tinygrad.device import Buffer, CompiledRunner
-
-import ctypes
-import gpuctypes.hip as hip
-from tinygrad.helpers import to_char_p_p, init_c_var
-def get_bytes(arg, get_sz, get_str, check) -> bytes: return (sz := init_c_var(ctypes.c_size_t(), lambda x: check(get_sz(arg, ctypes.byref(x)))), ctypes.string_at(init_c_var(ctypes.create_string_buffer(sz.value), lambda x: check(get_str(arg, x))), size=sz.value))[1]  # noqa: E501
-def check(status):
-  if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
-def compile_hip(prg:str, arch="gfx1100") -> bytes:
-  check(hip.hiprtcCreateProgram(ctypes.byref(prog := hip.hiprtcProgram()), prg.encode(), "<null>".encode(), 0, None, None))
-  compile_options = [f'--offload-arch={arch}', '-I/opt/rocm/include']
-  status = hip.hiprtcCompileProgram(prog, len(compile_options), to_char_p_p([o.encode() for o in compile_options]))
-  if status != 0: raise RuntimeError(f"compile failed: {get_bytes(prog, hip.hiprtcGetProgramLogSize, hip.hiprtcGetProgramLog, check).decode()}")
-  return get_bytes(prog, hip.hiprtcGetCodeSize, hip.hiprtcGetCode, check)
-
-prefix = """
-typedef long unsigned int size_t;
-extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_id(unsigned int);
-extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_group_id(unsigned int);
-extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_size(unsigned int);
-typedef float float2 __attribute__((ext_vector_type(2)));
-static inline __attribute__((device)) float2 make_float2(float x, float y) { return {x, y}; }
-"""
-
-code = """
-extern "C" __attribute__((global))void r_2_8_7_7_4_8_3_7_7_4_4_2_2(float* data0, const float* data1, const float* data2) {
-  int gidx0 = __ockl_get_group_id(2); /* 2 */
-  int gidx1 = __ockl_get_group_id(1); /* 8 */
-  int gidx2 = __ockl_get_group_id(0); /* 49 */
-  int lidx4 = __ockl_get_local_id(1); /* 4 */
-  int lidx5 = __ockl_get_local_id(0); /* 8 */
-  float2 acc0 = make_float2(0.0f,0.0f);
-  float2 acc1 = make_float2(0.0f,0.0f);
-  float2 acc2 = make_float2(0.0f,0.0f);
-  float2 acc3 = make_float2(0.0f,0.0f);
-  float2 acc4 = make_float2(0.0f,0.0f);
-  float2 acc5 = make_float2(0.0f,0.0f);
-  float2 acc6 = make_float2(0.0f,0.0f);
-  float2 acc7 = make_float2(0.0f,0.0f);
-  float2 acc8 = make_float2(0.0f,0.0f);
-  float2 acc9 = make_float2(0.0f,0.0f);
-  float2 acc10 = make_float2(0.0f,0.0f);
-  float2 acc11 = make_float2(0.0f,0.0f);
-  float2 acc12 = make_float2(0.0f,0.0f);
-  float2 acc13 = make_float2(0.0f,0.0f);
-  float2 acc14 = make_float2(0.0f,0.0f);
-  float2 acc15 = make_float2(0.0f,0.0f);
-  float2 acc16 = make_float2(0.0f,0.0f);
-  float2 acc17 = make_float2(0.0f,0.0f);
-  float2 acc18 = make_float2(0.0f,0.0f);
-  float2 acc19 = make_float2(0.0f,0.0f);
-  float2 acc20 = make_float2(0.0f,0.0f);
-  float2 acc21 = make_float2(0.0f,0.0f);
-  float2 acc22 = make_float2(0.0f,0.0f);
-  float2 acc23 = make_float2(0.0f,0.0f);
-  float2 acc24 = make_float2(0.0f,0.0f);
-  float2 acc25 = make_float2(0.0f,0.0f);
-  float2 acc26 = make_float2(0.0f,0.0f);
-  float2 acc27 = make_float2(0.0f,0.0f);
-  float2 acc28 = make_float2(0.0f,0.0f);
-  float2 acc29 = make_float2(0.0f,0.0f);
-  float2 acc30 = make_float2(0.0f,0.0f);
-  float2 acc31 = make_float2(0.0f,0.0f);
-  int alu0 = (gidx2/7);
-  int alu1 = (gidx2%7);
-  int alu2 = (alu1*32);
-  int alu3 = (lidx5*4);
-  int alu4 = ((gidx0*802816)+(gidx1*100352)+(alu0*1792)+(alu1*16)+(lidx4*448)+(lidx5*2));
-  for (int ridx0 = 0; ridx0 < 3; ridx0++) {
-    for (int ridx1 = 0; ridx1 < 7; ridx1++) {
-      int alu5 = ((alu0*(-32))+(lidx4*(-8))+(ridx1*(-1)));
-      bool alu6 = (alu5<(-2));
-      bool alu7 = (alu5<0);
-      bool alu8 = (((alu0*32)+(lidx4*8)+ridx1)<221);
-      for (int ridx2 = 0; ridx2 < 7; ridx2++) {
-        int alu9 = ((gidx0*150528)+(ridx0*50176)+(alu0*7168)+(lidx4*1792)+(ridx1*224)+alu2+alu3+ridx2);
-        int alu10 = ((alu1*(-32))+(lidx5*(-4))+(ridx2*(-1)));
-        bool alu11 = (alu10<(-2));
-        float val0 = 0.0f;
-        if ((alu6*alu11)) { val0 = data1[alu9+(-675)]; }
-        float val1 = 0.0f;
-        if ((alu7*alu11)) { val1 = data1[alu9+(-227)]; }
-        float val2 = 0.0f;
-        if (alu11) { val2 = data1[alu9+221]; }
-        float val3 = 0.0f;
-        if ((alu8*alu11)) { val3 = data1[alu9+669]; }
-        bool alu12 = (alu10<0);
-        bool alu13 = ((alu2+alu3+ridx2)<225);
-        float val4 = 0.0f;
-        if ((alu6*alu12*alu13)) { val4 = data1[alu9+(-673)]; }
-        float val5 = 0.0f;
-        if ((alu7*alu12*alu13)) { val5 = data1[alu9+(-225)]; }
-        float val6 = 0.0f;
-        if ((alu12*alu13)) { val6 = data1[alu9+223]; }
-        float val7 = 0.0f;
-        if ((alu8*alu12*alu13)) { val7 = data1[alu9+671]; }
-        int alu14 = ((gidx1*1176)+(ridx0*49)+(ridx1*7)+ridx2);
-        float val8 = data2[alu14];
-        float val9 = data2[alu14+147];
-        float val10 = data2[alu14+294];
-        float val11 = data2[alu14+441];
-        float val12 = data2[alu14+588];
-        float val13 = data2[alu14+735];
-        float val14 = data2[alu14+882];
-        float val15 = data2[alu14+1029];
-        (acc0).x = ((val0*val8)+(acc0).x);
-        (acc1).x = ((val0*val9)+(acc1).x);
-        (acc2).x = ((val0*val10)+(acc2).x);
-        (acc3).x = ((val0*val11)+(acc3).x);
-        (acc4).x = ((val1*val8)+(acc4).x);
-        (acc5).x = ((val1*val9)+(acc5).x);
-        (acc6).x = ((val1*val10)+(acc6).x);
-        (acc7).x = ((val1*val11)+(acc7).x);
-        (acc8).x = ((val2*val8)+(acc8).x);
-        (acc9).x = ((val2*val9)+(acc9).x);
-        (acc10).x = ((val2*val10)+(acc10).x);
-        (acc11).x = ((val2*val11)+(acc11).x);
-        (acc12).x = ((val3*val8)+(acc12).x);
-        (acc13).x = ((val3*val9)+(acc13).x);
-        (acc14).x = ((val3*val10)+(acc14).x);
-        (acc15).x = ((val3*val11)+(acc15).x);
-        (acc16).x = ((val0*val12)+(acc16).x);
-        (acc17).x = ((val0*val13)+(acc17).x);
-        (acc18).x = ((val0*val14)+(acc18).x);
-        (acc19).x = ((val0*val15)+(acc19).x);
-        (acc20).x = ((val1*val12)+(acc20).x);
-        (acc21).x = ((val1*val13)+(acc21).x);
-        (acc22).x = ((val1*val14)+(acc22).x);
-        (acc23).x = ((val1*val15)+(acc23).x);
-        (acc24).x = ((val2*val12)+(acc24).x);
-        (acc25).x = ((val2*val13)+(acc25).x);
-        (acc26).x = ((val2*val14)+(acc26).x);
-        (acc27).x = ((val2*val15)+(acc27).x);
-        (acc28).x = ((val3*val12)+(acc28).x);
-        (acc29).x = ((val3*val13)+(acc29).x);
-        (acc30).x = ((val3*val14)+(acc30).x);
-        (acc31).x = ((val3*val15)+(acc31).x);
-        (acc0).y = ((val4*val8)+(acc0).y);
-        (acc1).y = ((val4*val9)+(acc1).y);
-        (acc2).y = ((val4*val10)+(acc2).y);
-        (acc3).y = ((val4*val11)+(acc3).y);
-        (acc4).y = ((val5*val8)+(acc4).y);
-        (acc5).y = ((val5*val9)+(acc5).y);
-        (acc6).y = ((val5*val10)+(acc6).y);
-        (acc7).y = ((val5*val11)+(acc7).y);
-        (acc8).y = ((val6*val8)+(acc8).y);
-        (acc9).y = ((val6*val9)+(acc9).y);
-        (acc10).y = ((val6*val10)+(acc10).y);
-        (acc11).y = ((val6*val11)+(acc11).y);
-        (acc12).y = ((val7*val8)+(acc12).y);
-        (acc13).y = ((val7*val9)+(acc13).y);
-        (acc14).y = ((val7*val10)+(acc14).y);
-        (acc15).y = ((val7*val11)+(acc15).y);
-        (acc16).y = ((val4*val12)+(acc16).y);
-        (acc17).y = ((val4*val13)+(acc17).y);
-        (acc18).y = ((val4*val14)+(acc18).y);
-        (acc19).y = ((val4*val15)+(acc19).y);
-        (acc20).y = ((val5*val12)+(acc20).y);
-        (acc21).y = ((val5*val13)+(acc21).y);
-        (acc22).y = ((val5*val14)+(acc22).y);
-        (acc23).y = ((val5*val15)+(acc23).y);
-        (acc24).y = ((val6*val12)+(acc24).y);
-        (acc25).y = ((val6*val13)+(acc25).y);
-        (acc26).y = ((val6*val14)+(acc26).y);
-        (acc27).y = ((val6*val15)+(acc27).y);
-        (acc28).y = ((val7*val12)+(acc28).y);
-        (acc29).y = ((val7*val13)+(acc29).y);
-        (acc30).y = ((val7*val14)+(acc30).y);
-        (acc31).y = ((val7*val15)+(acc31).y);
-      }
-    }
-  }
-  *((float2*)(data0+alu4)) = acc0;
-  *((float2*)(data0+alu4+12544)) = acc1;
-  *((float2*)(data0+alu4+25088)) = acc2;
-  *((float2*)(data0+alu4+37632)) = acc3;
-  *((float2*)(data0+alu4+112)) = acc4;
-  *((float2*)(data0+alu4+12656)) = acc5;
-  *((float2*)(data0+alu4+25200)) = acc6;
-  *((float2*)(data0+alu4+37744)) = acc7;
-  *((float2*)(data0+alu4+224)) = acc8;
-  *((float2*)(data0+alu4+12768)) = acc9;
-  *((float2*)(data0+alu4+25312)) = acc10;
-  *((float2*)(data0+alu4+37856)) = acc11;
-  *((float2*)(data0+alu4+336)) = acc12;
-  *((float2*)(data0+alu4+12880)) = acc13;
-  *((float2*)(data0+alu4+25424)) = acc14;
-  *((float2*)(data0+alu4+37968)) = acc15;
-  *((float2*)(data0+alu4+50176)) = acc16;
-  *((float2*)(data0+alu4+62720)) = acc17;
-  *((float2*)(data0+alu4+75264)) = acc18;
-  *((float2*)(data0+alu4+87808)) = acc19;
-  *((float2*)(data0+alu4+50288)) = acc20;
-  *((float2*)(data0+alu4+62832)) = acc21;
-  *((float2*)(data0+alu4+75376)) = acc22;
-  *((float2*)(data0+alu4+87920)) = acc23;
-  *((float2*)(data0+alu4+50400)) = acc24;
-  *((float2*)(data0+alu4+62944)) = acc25;
-  *((float2*)(data0+alu4+75488)) = acc26;
-  *((float2*)(data0+alu4+88032)) = acc27;
-  *((float2*)(data0+alu4+50512)) = acc28;
-  *((float2*)(data0+alu4+63056)) = acc29;
-  *((float2*)(data0+alu4+75600)) = acc30;
-  *((float2*)(data0+alu4+88144)) = acc31;
-}
-"""
-
-dev = "HIP"
-lib = Device[dev].compiler.compile(prefix+code)
-#lib = compile_hip(code)
-b0 = Buffer(dev, 1605632, dtypes.float)
-b1 = Buffer(dev, 301506, dtypes.float)
-b2 = Buffer(dev, 9408, dtypes.float)
-print(hex(b0._buf.value), hex(b0._buf.value+1605632*4))
-print(hex(b1._buf.value))
-print(hex(b2._buf.value))
-#prg = CompiledRunner("r_2_8_7_7_4_8_3_7_7_4_4_2_2", "", dev, [7, 1, 1], [8, 4, 1], precompiled=lib)
-prg = CompiledRunner("r_2_8_7_7_4_8_3_7_7_4_4_2_2", "", dev, [49, 8, 2], [8, 4, 1], precompiled=lib)
-print("compiled")
-prg([b0, b1, b2], {})
-print("ran")
-Device[dev].synchronize()
-print("sync")

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
+from tinygrad.runtime.support.memory import PageTableTraverseContext, AddrSpace
 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)])
+      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       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)])
-      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
+      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.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)])
+      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       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)])
+      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)], aspace=AddrSpace.PHYS)
 
       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)])
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
     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)])
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
     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))
+      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
       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))
+      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
   def test_frag_size(self):

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

@@ -2,18 +2,27 @@ 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
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.engine.realize import Runner
+from tinygrad.engine.realize import Runner, get_program
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, CI
-from tinygrad.codegen import full_rewrite
+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
 
 # 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))
@@ -51,13 +60,12 @@ 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, *data)))
+    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
-  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)))
+  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
-  return out_buf.cast(uop.dtype.fmt).tolist()[0]
+  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 
 def not_support_multi_device():
   # CL and CUDA don't support multi device if in CI

test/test_renderer_failures.py

@@ -1,30 +1,26 @@
 import unittest
 import numpy as np
+from dataclasses import replace
 from tinygrad.device import Buffer, Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup, flatten, prod
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.helpers import 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], stores:list[UOp], local_size=None):
+def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
   for x in inputs: x.realize()
-  # 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))
+  uops = prg.uops
   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]
-  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))
+  prg = replace(prg, device=Device.DEFAULT)
+  if local_size is not None: prg = replace(prg, local_size=local_size)
+  ei = CompiledRunner(prg)
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
@@ -37,8 +33,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,))
-  uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-  return _test_uop_result([Tensor([input_val])], uops)[0]
+  prg = get_program(sink, Device[Device.DEFAULT].renderer)
+  return _test_uop_result([Tensor([input_val])], prg)[0]
 
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -47,8 +43,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,))
-    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
+    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], prg, 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")
@@ -58,8 +54,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,))
-    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], uops, local_size=[4, 2, 1])[0]
+    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], prg, 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")
@@ -104,8 +100,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,))
-    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
+    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], prg, 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.codegen import full_rewrite
+from tinygrad.engine.realize import get_program
 from tinygrad.dtype import DType
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -869,9 +869,8 @@ class TestIdxUpcast(unittest.TestCase):
     for s in schedule:
       if s.ast.op is Ops.SINK:
         renderer = Device[s.bufs[0].device].renderer
-        uops = full_rewrite(s.ast, renderer)
-        renderer.render(uops)
-        return uops
+        prg = get_program(s.ast, renderer)
+        return prg.uops
 
   def _assert(self, dtype: DType, a: Tensor):
     uops = self._schedule_render(a)

test/test_uops.py

@@ -7,30 +7,27 @@ 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
+from dataclasses import replace
 
 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 = full_rewrite(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  ret = get_uops(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
   assert ret[-1].op is Ops.SINK
   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))
+  prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
+  return CompiledRunner(replace(prg, device=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))

test/test_uops_stats.py

@@ -4,7 +4,6 @@ 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
@@ -146,7 +145,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 = full_rewrite(u5.sink())
+    uops = list(u5.toposort())
 
     globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
     o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
@@ -155,7 +154,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 = full_rewrite(u4.sink())
+    uops_fma = list(u4.toposort())
 
     self.assertEqual(flops_mem(uops), flops_mem(uops_fma))
 

test/testextra/test_cfg_viz.py

@@ -0,0 +1,172 @@
+import unittest
+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
+
+# TODO: use the RDNA3 renderer when it's in master
+template = """.text
+.globl fn_name
+.p2align 8
+.type fn_name,@function
+fn_name:
+  INSTRUCTION
+
+.rodata
+.p2align 6
+.amdhsa_kernel fn_name
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
+  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
+  .amdhsa_wavefront_size32 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: fn_name
+    .symbol: fn_name.kd
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 8
+    .max_flat_workgroup_size: 1024
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+...
+.end_amdgpu_metadata
+"""
+
+@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
+
+@unittest.skipUnless(Device.DEFAULT == "AMD", "only on AMD")
+class TestCfg(unittest.TestCase):
+  def setUp(self):
+    arch = Device["AMD"].arch
+    if not any(arch.startswith(a) for a in {"gfx11", "gfx12"}):
+      self.skipTest(f"tests written for RDNA, got arch {arch}")
+
+  def test_simple(self):
+    run_asm("simple", """
+      entry:
+        s_branch bb1
+      bb1:
+        s_endpgm
+    """)
+
+  def test_diamond(self):
+    run_asm("diamond", """
+      entry:
+        s_cmp_eq_i32 s0, 0
+        s_cbranch_scc1 if
+        s_branch else
+      if:
+        s_nop 1
+        s_branch end
+      else:
+        s_nop 0
+      end:
+        s_endpgm
+    """)
+
+  def test_loop(self):
+    run_asm("simple_loop", """
+      entry:
+        s_mov_b32 s1, 4
+      loop:
+        s_add_u32 s1, s1, -1
+        s_cmp_eq_i32 s1, 0
+        s_cbranch_scc0 loop
+        s_endpgm
+    """)
+
+  def test_loop_branch(self):
+    run_asm("loop_if", """
+      entry:
+        s_mov_b32 s1, 4
+      loop:
+        s_add_u32 s1, s1, -1
+        s_cmp_eq_i32 s1, 2
+        s_cbranch_scc1 cond
+        s_branch cont
+      cond:
+        s_add_u32 s1, s1, -2
+      cont:
+        s_cmp_eq_i32 s1, 0
+        s_cbranch_scc0 loop
+        s_endpgm
+    """)
+
+  def test_loop_break(self):
+    run_asm("loop_break", """
+      entry:
+        s_mov_b32 s1, 8
+      loop:
+        s_add_u32 s1, s1, -1
+        s_cmp_eq_i32 s1, 5
+        s_cbranch_scc1 break
+        s_cmp_eq_i32 s1, 0
+        s_cbranch_scc0 loop
+      break:
+        s_endpgm
+    """)
+
+  def test_switch(self):
+    run_asm("switch_case", """
+      entry:
+        s_cmp_eq_i32 s0, 0
+        s_cbranch_scc1 case0
+        s_cmp_eq_i32 s0, 1
+        s_cbranch_scc1 case1
+        s_branch case2
+      case0:
+        s_nop 0
+        s_branch join
+      case1:
+        s_nop 1
+        s_branch join
+      case2:
+        s_nop 2
+        s_branch join
+      join:
+        s_endpgm
+    """)
+
+  def test_ping_pong(self):
+    run_asm("ping_pong", """
+      entry:
+        s_cmp_eq_i32 s0, 0
+        s_cbranch_scc1 ping
+        s_branch pong
+      ping:
+        s_cmp_eq_i32 s1, 0
+        s_cbranch_scc1 pong
+        s_branch end
+      pong:
+        s_cmp_eq_i32 s2, 0
+        s_cbranch_scc1 ping
+      end:
+        s_endpgm
+    """)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_dtype.py

@@ -66,5 +66,21 @@ class TestDtypeTolist(unittest.TestCase):
     # 57344
     self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
 
+class TestCanLosslessCast(unittest.TestCase):
+  def test_can_lossless_cast(self):
+    from tinygrad.dtype import can_lossless_cast
+    # signed -> unsigned is NOT lossless (negative values wrap)
+    self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.uint64))
+    self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.uint32))
+    # unsigned -> larger signed is lossless
+    self.assertTrue(can_lossless_cast(dtypes.uint8, dtypes.int16))
+    self.assertTrue(can_lossless_cast(dtypes.uint32, dtypes.int64))
+    # large ints don't fit in floats
+    self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.float))
+    self.assertFalse(can_lossless_cast(dtypes.int64, dtypes.double))
+    # half has more mantissa bits
+    self.assertTrue(can_lossless_cast(dtypes.int8, dtypes.half))
+    self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.bfloat16))
+
 if __name__ == "__main__":
   unittest.main()

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=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
 
 core_dtypes = list(DTYPES_DICT.values())

test/unit/test_elf.py

@@ -1,5 +1,5 @@
 import unittest, subprocess, platform
-from tinygrad.runtime.ops_cpu import ClangJITCompiler
+from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
 from tinygrad.runtime.support.elf import elf_loader
 
 class TestElfLoader(unittest.TestCase):

test/unit/test_llm_moe.py

@@ -0,0 +1,53 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor
+
+class TestMoEFeedForward(unittest.TestCase):
+  def test_moe_feed_forward(self):
+    from tinygrad.apps.llm import TransformerBlock
+    dim, hidden, n_heads = 8, 16, 2
+    num_experts, k = 4, 2
+
+    block = TransformerBlock(dim, hidden, n_heads, n_heads, norm_eps=1e-5, head_dim=dim//n_heads,
+                             rope_theta=10000, max_context=16, num_experts=num_experts, num_experts_per_tok=k)
+
+    # set up weights: gate scales by (expert_id+1), up/down are identity-ish, router picks experts 0,2
+    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
+    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
+    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
+    block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T  # router strongly prefers experts 0 and 2
+    block.ffn_norm.weight = Tensor.ones(dim)  # identity norm
+
+    # input of ones -> after norm still ~ones -> experts 0,2 selected -> weighted sum of silu outputs
+    h = Tensor.ones(1, 1, dim)
+    out = block._feed_forward(h)
+
+    # expected: residual + moe_output ≈ 1 + avg(silu(1), silu(3))
+    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
+    np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
+
+  def test_moe_feed_forward_batched(self):
+    from tinygrad.apps.llm import TransformerBlock
+    dim, hidden, n_heads = 8, 16, 2
+    num_experts, k = 4, 2
+
+    block = TransformerBlock(dim, hidden, n_heads, n_heads, norm_eps=1e-5, head_dim=dim//n_heads,
+                             rope_theta=10000, max_context=16, num_experts=num_experts, num_experts_per_tok=k)
+
+    # same setup as BS=1 test
+    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
+    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
+    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
+    block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T
+    block.ffn_norm.weight = Tensor.ones(dim)
+
+    # test with BS=2, T=3
+    h = Tensor.ones(2, 3, dim)
+    out = block._feed_forward(h)
+
+    # all outputs should match the BS=1 expected value
+    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
+    np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
+
+if __name__ == '__main__':
+  unittest.main()

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.port = 11435
-    cls.server = TCPServerWithReuse(('127.0.0.1', cls.port), Handler)
+    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
+    cls.port = cls.server.server_address[1]
     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 = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
+    uops = get_uops(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)))

test/web/test_viz.js

@@ -1,35 +0,0 @@
-const { spawn } = require("child_process");
-const puppeteer = require("puppeteer");
-
-async function main() {
-  // ** start viz server
-  const proc = spawn("python", ["-u", "-c", "from tinygrad import Tensor; Tensor.arange(4).realize()"], { env: { ...process.env, VIZ:"1" },
-                      stdio: ["inherit", "pipe", "inherit"]});
-  await new Promise(resolve => proc.stdout.on("data", r => {
-    if (r.includes("ready")) resolve();
-  }));
-
-  // ** run browser tests
-  let browser, page;
-  try {
-    browser = await puppeteer.launch({ headless: true });
-    page = await browser.newPage();
-    const res = await page.goto("http://localhost:8000", { waitUntil:"domcontentloaded" });
-    if (res.status() !== 200) throw new Error("Failed to load page");
-    const scheduleSelector = await page.waitForSelector("ul:nth-of-type(2)");
-    scheduleSelector.click();
-    await page.waitForSelector("rect");
-    await page.waitForFunction(() => {
-      const nodes = document.querySelectorAll("#nodes > g").length;
-      const edges = document.querySelectorAll("#edges > path").length;
-      return nodes > 0 && edges > 0;
-    });
-  } finally {
-    // ** cleanups
-    if (page != null) await page.close();
-    if (browser != null) await browser.close();
-    proc.kill();
-  }
-}
-
-main();

tinygrad/apps/llm.py

@@ -5,14 +5,14 @@ from tinygrad.helpers import partition, TCPServerWithReuse, HTTPRequestHandler,
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3"):
-    if preset not in ("llama3","llama-v3","llama-bpe","qwen2"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
     # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
     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
-    # 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))
+    # 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))
     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}]+")
@@ -52,9 +52,13 @@ class SimpleTokenizer:
 
   def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode(errors='replace')
   def role(self, role:str):
+    if self.preset == 'olmo': return self.encode("<|" + role + "|>\n")  # OLMoE Instruct format
     if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
     return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
-  def end_turn(self, eos_id:int): return [eos_id] + self.encode("\n") if self.preset == 'qwen2' else [eos_id]
+  def end_turn(self, eos_id:int):
+    if self.preset == 'olmo': return self.encode("\n")
+    if self.preset == 'qwen2': return [eos_id] + self.encode("\n")
+    return [eos_id]
 
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
@@ -62,6 +66,14 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
   freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
   return freqs.cos().cat(freqs.sin(), dim=-1).contiguous()
 
+class ExpertWeights:
+  """Like nn.Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features)."""
+  def __init__(self, num_experts:int, in_features:int, out_features:int):
+    self.weight = Tensor.zeros(num_experts, out_features, in_features)
+  def __call__(self, sel:Tensor, x:Tensor) -> Tensor:
+    # sel: (B, T, k), x: (B, T, 1, in) or (B, T, k, in) -> output: (B, T, k, out)
+    return (x.unsqueeze(-2) @ self.weight[sel].transpose(-1, -2)).squeeze(-2)
+
 def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
   assert x.shape[-1] % 2 == 0
   cos, sin = freqs_cis.reshape(1, 1, x.shape[2], -1).chunk(2, dim=-1)
@@ -70,12 +82,13 @@ def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:bool=False):
+               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
     self.n_heads      = n_heads
     self.n_kv_heads   = n_kv_heads
     self.head_dim     = head_dim
-    self.max_context  = max_context
     self.rope_theta   = rope_theta
+    self.max_context  = max_context
+    self.qk_norm      = qk_norm
 
     # --- attention projections (all linear, bias-free) ------------------
     q_proj_out       = self.head_dim * n_heads
@@ -88,23 +101,30 @@ class TransformerBlock:
     # --- RMSNorms --------------------------------------------------------
     self.attn_norm   = nn.RMSNorm(dim, norm_eps)
     self.ffn_norm    = nn.RMSNorm(dim, norm_eps)
-    if qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(self.head_dim, norm_eps), nn.RMSNorm(self.head_dim, norm_eps)
+    if qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(qk_norm, norm_eps), nn.RMSNorm(qk_norm, norm_eps)
 
-    # --- feed-forward ----------------------------------------------------
-    self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
-    self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
-    self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
+    # --- feed-forward (MoE or dense) -------------------------------------
+    if num_experts > 0:
+      self.num_experts_per_tok = num_experts_per_tok
+      self.ffn_gate_inp = nn.Linear(dim, num_experts, bias=False)  # router
+      self.ffn_gate_exps = ExpertWeights(num_experts, dim, hidden_dim)
+      self.ffn_up_exps = ExpertWeights(num_experts, dim, hidden_dim)
+      self.ffn_down_exps = ExpertWeights(num_experts, hidden_dim, dim)
+    else:
+      self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
+      self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
+      self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     x_norm = self.attn_norm(x)                       # (B,T,D)
     q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
+    if self.qk_norm and self.qk_norm != self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
 
     B, T, _ = x.shape
     q = q.reshape(B, T, self.n_heads,    self.head_dim).transpose(1, 2)  # (B,H,T,Hd)
     k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
     v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
-
-    if hasattr(self, 'attn_q_norm'): q, k = self.attn_q_norm(q), self.attn_k_norm(k)
+    if self.qk_norm == self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
 
     # TODO: make UOp have SupportsIndex
     freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)[start_pos:start_pos+T]  # type: ignore
@@ -127,6 +147,11 @@ class TransformerBlock:
 
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
+    if hasattr(self, 'ffn_gate_exps'):
+      x = h_norm.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
+      probs, sel = self.ffn_gate_inp(h_norm).softmax(-1).topk(self.num_experts_per_tok)  # (B, T, k) each
+      x_down = self.ffn_down_exps(sel, self.ffn_gate_exps(sel, x).silu() * self.ffn_up_exps(sel, x))  # (B, T, k, D)
+      return h + (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
     # TODO: remove the need for this contiguous
     gated  = self.ffn_gate(h_norm).silu().contiguous() * self.ffn_up(h_norm)
     return h + self.ffn_down(gated)
@@ -136,9 +161,9 @@ class TransformerBlock:
 
 class Transformer:
   def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:bool=False):
-    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, qk_norm)
-                for _ in range(num_blocks)]
+               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
+    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, qk_norm,
+                                 num_experts, num_experts_per_tok) for _ in range(num_blocks)]
     self.token_embd  = nn.Embedding(vocab_size, dim)
     self.output_norm = nn.RMSNorm(dim, norm_eps)
     self.output = nn.Linear(dim, vocab_size, bias=False)
@@ -170,16 +195,20 @@ class Transformer:
     max_context = min(max_context, kv[f'{arch}.context_length']) if max_context is not None else kv[f'{arch}.context_length']
     n_heads, n_kv_heads = kv[f'{arch}.attention.head_count'], kv[f'{arch}.attention.head_count_kv']
 
-    # permute Q/K weights from interleaved to half-split RoPE layout: [0,1,2,3,4,5...] -> [0,2,4,...,1,3,5,...]
-    if arch != 'qwen3':
+    # Permute Q/K weights from interleaved to half-split RoPE layout (llama-style models only)
+    if arch == 'llama':
       for name in state_dict:
         if 'attn_q.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_heads, two=2)
         if 'attn_k.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_kv_heads, two=2)
 
-    model = Transformer(num_blocks=kv[f'{arch}.block_count'], dim=kv[f'{arch}.embedding_length'], hidden_dim=kv[f'{arch}.feed_forward_length'],
+    model = Transformer(num_blocks=kv[f'{arch}.block_count'], dim=kv[f'{arch}.embedding_length'],
+                        hidden_dim=kv.get(f'{arch}.expert_feed_forward_length', kv[f'{arch}.feed_forward_length']),
                         n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'],
-                        vocab_size=len(kv['tokenizer.ggml.tokens']), head_dim=kv[f'{arch}.attention.key_length'],
-                        rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context, qk_norm='blk.0.attn_q_norm.weight' in state_dict)
+                        vocab_size=len(kv['tokenizer.ggml.tokens']),
+                        head_dim=kv.get(f'{arch}.attention.key_length', kv[f'{arch}.embedding_length'] // n_heads),
+                        rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context,
+                        qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
+                        num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0))
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
     for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
@@ -207,6 +236,8 @@ models = {
   "qwen3:0.6b": "https://huggingface.co/Qwen/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q8_0.gguf",
   "qwen3:1.7b": "https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/main/Qwen3-1.7B-Q4_K_M.gguf",
   "qwen3:8b": "https://huggingface.co/Qwen/Qwen3-8B-GGUF/resolve/main/Qwen3-8B-Q4_K_M.gguf",
+  "qwen3:30b-a3b": "https://huggingface.co/Qwen/Qwen3-30B-A3B-GGUF/resolve/main/Qwen3-30B-A3B-Q4_K_M.gguf",
+  "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
 }
 
 # *** simple OpenAI compatible server on 11434 to match ollama ***

tinygrad/codegen/__init__.py

@@ -1,11 +1,12 @@
 from typing import cast
 import itertools
-from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
+from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
-from tinygrad.renderer import Renderer
+from tinygrad.renderer import Renderer, ProgramSpec
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.helpers import panic
+from tinygrad.codegen.opt import Opt
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -28,6 +29,8 @@ pm_syntactic_sugar = PatternMatcher([
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
+  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
+  if DEBUG >= 5: print(pyrender(sink))
   if SPEC: type_verify(sink, kernel_spec)
 
   # preprocess
@@ -123,20 +126,49 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher) -> list[UOp]:
     newlst.extend(ret[1])
   return newlst
 
-def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
+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:Renderer, prg:UOp, lin:UOp) -> UOp:
+  src = ctx.render(list(lin.src))
+  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),))
+
+def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
+  if ctx.compiler is None: return None
+  lib = ctx.compiler.compile_cached(source.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"), UPat(Ops.DEVICE)), name="prg"), do_linearize),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
+])
+
+@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:
   """
-  Function to transform the Kernel UOp graph into a linearized program.
+  Transform an AST into a ProgramSpec. May trigger BEAM search.
 
   Args:
-    sink: The Ops.SINK rooting the Kernel graph.
-    ren: The Renderer (can change how things are processed, fix this).
+    ast: The Ops.SINK rooted AST
+    renderer: The renderer used to generate the code
 
   Returns:
-    Linear program in UOps.
+    The ProgramSpec of the program.
   """
 
-  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  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
+  # fix up KernelInfo
+  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())
+
+  # rewrite to prg
+  full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
+  prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
+  prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
+
+  # create the ProgramSpec
+  return ProgramSpec.from_uop(prg)

tinygrad/codegen/opt/search.py

@@ -6,7 +6,8 @@ from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, di
 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
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.codegen import get_program
 from tinygrad.renderer import ProgramSpec
 from tinygrad.codegen.opt.postrange import Scheduler
 
@@ -37,10 +38,10 @@ def get_test_global_size(global_size, max_global_size, var_vals):
 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
-  if allow_test_size and p.global_size is not None and max_global_size is not None:
+  if allow_test_size 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]
@@ -71,7 +72,7 @@ def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[
       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)
+    prog = p.lib if p.lib is not None else compiler.compile(p.src)
     et = time.perf_counter() - st
     ret = (p, prog, et)
   except RuntimeError:

tinygrad/device.py

@@ -278,7 +278,7 @@ class Compiler:
   def disassemble(self, lib:bytes): pass
 
 @dataclass(frozen=True)
-class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial; ctrl_var:ContextVar|None = None # noqa: E702
+class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial|None; ctrl_var:ContextVar|None = None # noqa: E702
 
 @dataclass(frozen=True)
 class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
@@ -290,21 +290,25 @@ class Compiled:
     self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
 
     self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
-    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial]]] = {}
-    self.cached_pair:dict[Any, tuple[Renderer, Compiler]] = {}
+    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial|None]]] = {}
+    self.cached_pair:dict[Any, tuple[Renderer, Compiler|None]] = {}
     for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
-      self.comp_sets[self._compiler_name(cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
+      self.comp_sets[self._compiler_name(cpair.renderer, cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
 
   @property
   def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
 
   @property
-  def compiler(self) -> Compiler: return self._select_compiler_pair()[1]
+  def compiler(self) -> Compiler:
+    if (ret:=self.renderer.compiler or self._select_compiler_pair()[1]) is None: raise RuntimeError(f"no compiler for {self.device}")
+    return ret
 
-  def _compiler_name(self, c:type[Compiler]|functools.partial) -> str:
-    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
+  def _compiler_name(self, r:type[Renderer]|functools.partial, c:type[Compiler]|functools.partial|None) -> str:
+    devname = self.device.split(':')[0].upper()
+    if c is None: return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname) or devname
+    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname) or devname
 
-  def _select_compiler_pair(self) -> tuple[Renderer, Compiler]:
+  def _select_compiler_pair(self) -> tuple[Renderer, Compiler|None]:
     # select forced compiler from global env var.
     forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])
 
@@ -398,7 +402,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
             # d.renderer, d.compiler = r(), c()
             with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
             if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-            set_text = f'({cc_ctrl_var.key}={d._compiler_name(c)} to make default)' if cc_ctrl_var is not None else ''
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r, c)} to make default)' if cc_ctrl_var is not None else ''
             default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
             compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
             any_works = True

tinygrad/dtype.py

@@ -218,17 +218,19 @@ DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType)
 INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "index":"index"}
 
 @functools.cache
-def can_safe_cast(dt0:DType, dt1:DType) -> bool:
+def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
   # return if dt1 preserves value of dt0
-  # https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
+  # similar to https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
   if dt0 == dt1 or dt0 == dtypes.bool: return True
   match dt1:
     case dtypes.index: return dt0 in dtypes.ints
     case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16, *dtypes.fp8s,
       dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
     case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
+    case dtypes.half: return dt0 in (*dtypes.fp8s, dtypes.uint8, dtypes.int8)
     case dtypes.uint64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8)
     case dtypes.uint32: return dt0 in (dtypes.uint16, dtypes.uint8)
+    case dtypes.uint16: return dt0 in (dtypes.uint8,)
     case dtypes.int64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
     case dtypes.int32: return dt0 in (dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
     case dtypes.int16: return dt0 in (dtypes.uint8, dtypes.int8)
@@ -319,11 +321,8 @@ def fp8_to_float(x: int, dtype: DType) -> float:
 truncate: dict[DType, Callable] = {dtypes.bool: bool,
   dtypes.float16: float_to_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)),
   **{fp8: (lambda x, dtype=fp8: fp8_to_float(float_to_fp8(x, dtype), dtype)) for fp8 in dtypes.fp8s},
-  dtypes.float32: lambda x: ctypes.c_float(x).value, dtypes.float64: lambda x: ctypes.c_double(x).value,
-  dtypes.uint8: lambda x: ctypes.c_uint8(x).value, dtypes.uint16: lambda x: ctypes.c_uint16(x).value,
-  dtypes.uint32: lambda x: ctypes.c_uint32(x).value, dtypes.uint64: lambda x: ctypes.c_uint64(x).value,
-  dtypes.int8: lambda x: ctypes.c_int8(x).value, dtypes.int16: lambda x: ctypes.c_int16(x).value, dtypes.int32: lambda x: ctypes.c_int32(x).value,
-  dtypes.int64: lambda x: ctypes.c_int64(x).value}
+  **{getattr(dtypes, n): (lambda x, c=getattr(ctypes, f'c_{n}'): c(x).value)
+     for n in ('float', 'double', 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64')}}
 
 # numpy and torch dtype interop
 

tinygrad/engine/jit.py

@@ -97,7 +97,7 @@ class GraphRunner(Runner):
         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 ji.prg.p.global_size is not None and ji.prg.p.local_size is not None
+          assert 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

tinygrad/engine/realize.py

@@ -2,52 +2,12 @@ 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, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
 from tinygrad.helpers import unwrap
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer
 from tinygrad.device import Device, Buffer
-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.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer, 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
-
-  Returns:
-    The ProgramSpec of the program.
-  """
-
-  if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
-  if DEBUG >= 5: print(pyrender(ast))
-
-  # linearize
-  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"
-
-  # print and render
-  if DEBUG >= 6: print_uops(uops)
-  src = renderer.render(uops)
-
-  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)
+from tinygrad.renderer import ProgramSpec, Estimates
+from tinygrad.codegen import get_program
 
 # **************** Runners ****************
 
@@ -76,36 +36,29 @@ 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)
-    self.p:ProgramSpec = p
-    if precompiled is not None: self.lib = precompiled
-    else:
+    if p.lib is None:
       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
+        p = replace(p, lib=Device[p.device].compiler.compile_cached(p.src))
+    self.p:ProgramSpec = p
+    assert self.p.lib is not None
+    if DEBUG >= 7: Device[p.device].compiler.disassemble(self.p.lib)
+    self._prg = Device[p.device].runtime(p.function_name, self.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 = {}
-    has_local = Device[self.p.device].renderer.has_local
     global_size, local_size = self.p.launch_dims(var_vals)
-    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]
+    if Device[self.p.device].renderer.has_local 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 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)
-      assert len(global_size) == 3, "global size must have len 3"
-    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.vars), wait=wait)
+    return self._prg(*[x._buf for x in rawbufs], global_size=tuple(global_size), local_size=tuple(local_size) if local_size else None,
+                     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)
@@ -162,7 +115,7 @@ 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)
     method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))

tinygrad/engine/schedule.py

@@ -90,23 +90,29 @@ 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:dict[UOp, UOp], b:UOp):
-  if (ret:=ctx.get(b, None)) is None:
+def replace_input_buffer(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
+  if (ret:=ctx[0].get(b, None)) is None:
     if b.op is Ops.BUFFER:
-      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
+      ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx[0])), b.src[1]))
     else:
       # TODO: flip args in CONST
       assert b.op is Ops.CONST
-      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
+      ctx[0][b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx[0]))))
   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"), lambda ctx,b: ctx.setdefault(b, b.replace(src=(b.src[0],)))),
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), strip_bind),
 ])
 
 def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
@@ -129,9 +135,10 @@ 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
+  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key, extract var_vals
   input_buffers: dict[UOp, UOp] = {}
-  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
+  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")
   sched_cache_key = big_sink_cache.key
 
   if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
@@ -139,7 +146,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] = {}
@@ -191,17 +198,8 @@ 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
+  return tensor_map, schedule, var_vals if schedule else {}

tinygrad/helpers.py

@@ -115,12 +115,12 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
-def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg:str, cache:dict|None=None) -> tuple[T,...]|T:
+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]|None]], err_msg:str, cache:dict|None=None):
   excs = []
   for typ in candidates:
     if cache is not None and typ in cache: return cache[typ]
     try:
-      x = tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+      x = tuple([cast(Callable, t)() if t is not None else None for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
       if cache is not None: cache[typ] = x
       return x
     except Exception as e: excs.append(e)

tinygrad/renderer/__init__.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
-from typing import Callable, cast
+from typing import Callable, cast, TYPE_CHECKING
 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.helpers import to_function_name, dedup, prod, DEBUG
+from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher, print_uops
 from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.codegen.opt import Opt
+if TYPE_CHECKING: from tinygrad.device import Compiler
 
 @dataclass(frozen=True)
 class Estimates:
@@ -64,36 +65,15 @@ class ProgramSpec:
   device:str
   ast:UOp  # save the base ast (this is method cache key)
   uops:list[UOp]|None=None
+  lib:bytes|None=None
 
-  # filled in from uops (if we have uops)
-  global_size:list[int]|None=None
+  # filled in from uops (via from_uop)
+  global_size:list[int]=field(default_factory=lambda: [1,1,1])
   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:
@@ -109,10 +89,43 @@ class ProgramSpec:
     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
+    global_size = [sym_infer(sz, var_vals) for sz in self.global_size]
     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
 
+  @staticmethod
+  def from_uop(prg:UOp) -> ProgramSpec:
+    """Construct ProgramSpec from a PROGRAM UOp."""
+    assert prg.op is Ops.PROGRAM, f"expected PROGRAM, got {prg.op}"
+    # SINK/DEVICE/LINEAR/SOURCE/BINARY?
+    sink, device, linear, source = prg.src[:4]
+    lib = prg.src[4].arg if len(prg.src) > 4 else None
+    uops = list(linear.src)
+    if DEBUG >= 6: print_uops(uops)  # LINEAR is src[2]
+
+    # single pass through the uops to extract metadata
+    _vars: list[Variable] = []
+    _globals: list[int] = []
+    outs: list[int] = []
+    ins: list[int] = []
+    global_size: list[int] = [1, 1, 1]
+    local_size: list[int]|None = [1, 1, 1]
+    for u in uops:
+      if u.op is Ops.DEFINE_VAR: _vars.append(u)
+      if u.op is Ops.DEFINE_GLOBAL: _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: (outs if u.op is Ops.STORE else ins).append(buf.arg)
+        # TODO: can else happen?
+      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
+        # 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())
+
+    return ProgramSpec(sink.arg.name, source.arg, device.arg, sink, uops, lib, global_size, local_size,
+                       sorted(_vars, key=lambda v: v.arg), sorted(dedup(_globals)), sorted(dedup(outs)), sorted(dedup(ins)))
+
 class Renderer:
   device: str = ""
   suffix: str = ""
@@ -129,6 +142,7 @@ class Renderer:
   pre_matcher: PatternMatcher|None = None
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
+  compiler: Compiler|None = None
 
   def __reduce__(self): return self.__class__, ()
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")

tinygrad/renderer/cstyle.py

@@ -8,6 +8,7 @@ from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, trunc
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.late.devectorizer import no_vectorized_alu
 
+
 base_rewrite = PatternMatcher([
   (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
@@ -278,6 +279,11 @@ class ClangRenderer(CStyleLanguage):
     defines = '\n'.join(self._render_defines(uops))
     return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs)
 
+class ClangJITRenderer(ClangRenderer):
+  def __init__(self):
+    from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
+    self.compiler = ClangJITCompiler()
+
 class OpenCLRenderer(CStyleLanguage):
   device = "CL"
 
@@ -328,7 +334,10 @@ class IntelRenderer(OpenCLRenderer):
 class MetalRenderer(CStyleLanguage):
   device = "METAL"
   shared_max = 32768
-  def __init__(self): self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
+  def __init__(self):
+    self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
+    from tinygrad.runtime.ops_metal import MetalCompiler
+    self.compiler = MetalCompiler()
 
   # language options
   kernel_typedef = "kernel void"
@@ -440,6 +449,18 @@ class CUDARenderer(CStyleLanguage):
 
     return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix)
 
+class CUDACUDARenderer(CUDARenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch)
+    from tinygrad.runtime.support.compiler_cuda import CUDACompiler
+    self.compiler = CUDACompiler(arch)
+
+class CUDANVCCRenderer(CUDARenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch)
+    from tinygrad.runtime.support.compiler_cuda import NVCCCompiler
+    self.compiler = NVCCCompiler(arch)
+
 class AMDRenderer(CStyleLanguage):
   device = "AMD"
   shared_max = 65536
@@ -532,6 +553,32 @@ class AMDRenderer(CStyleLanguage):
   for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""")
     return super().render_kernel(function_name, kernel, bufs, uops, prefix)
 
+class AMDHIPRenderer(AMDRenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch)
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    self.compiler = HIPCompiler(arch)
+
+class AMDHIPCCRenderer(AMDRenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch)
+    from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+    self.compiler = HIPCCCompiler(arch)
+
 class NVRenderer(CUDARenderer): device = "NV"
+
+class NVNVRenderer(NVRenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch)
+    from tinygrad.runtime.support.compiler_cuda import NVCompiler
+    self.compiler = NVCompiler(arch)
+
 class HIPRenderer(AMDRenderer): device = "HIP"
+
+class HIPHIPRenderer(HIPRenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch)
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    self.compiler = HIPCompiler(arch)
+
 class QCOMRenderer(OpenCLRenderer): device = "QCOM"

tinygrad/renderer/llvmir.py

@@ -143,6 +143,9 @@ class LLVMRenderer(Renderer):
   if AMX: tensor_cores = tc.amx
 
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
+  def __init__(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+    self.compiler = CPULLVMCompiler()
   def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
   def _render_footer(self, uops: list[UOp]) -> str: return 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }'
   def _render_fn(self, name:str, args:list[tuple[str,DType]], kernel:list[str], prefix:list[str]|None=None) -> str:
@@ -254,7 +257,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
                   f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
     return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
   def __init__(self, arch:str):
+    from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
     self.arch = arch
+    self.compiler = AMDLLVMCompiler(arch)
     self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
     self.is_cdna = AMDRenderer.is_cdna(arch)
     self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])

tinygrad/renderer/nir.py

@@ -245,6 +245,11 @@ class LVPRenderer(NIRRenderer):
     srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
     setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
 
+  def __init__(self):
+    from tinygrad.runtime.support.compiler_mesa import LVPCompiler
+    super().__init__()
+    self.compiler = LVPCompiler()
+
   def prerender(self, uops:list[UOp]):
     super().prerender(uops)
     self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])

tinygrad/renderer/ptx.py

@@ -240,3 +240,17 @@ class PTXRenderer(Renderer):
 
       if u.op is Ops.SPECIAL: kernel = [f".reg .u32 %{u.arg};"] + kernel
     return self.render_kernel(kernel, name, bufs, c.items(), uops)
+
+class CUDAPTXRenderer(PTXRenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch, "CUDA")
+    from tinygrad.runtime.support.compiler_cuda import PTXCompiler
+    self.compiler = PTXCompiler(arch)
+  def __reduce__(self): return self.__class__, (self.arch,)
+
+class NVPTXRenderer(PTXRenderer):
+  def __init__(self, arch:str):
+    super().__init__(arch, "NV")
+    from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler
+    self.compiler = NVPTXCompiler(arch)
+  def __reduce__(self): return self.__class__, (self.arch,)

tinygrad/renderer/wgsl.py

@@ -1,6 +1,7 @@
 from tinygrad.dtype import DType, PtrDType, dtypes, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite, extra_pm
+from tinygrad.device import Compiler
 from tinygrad.helpers import strip_parens
 
 def sign_extend(val:UOp, sext_am:int):
@@ -46,6 +47,7 @@ class WGSLRenderer(CStyleLanguage):
   global_max = (65535, 65535, 65535)
   local_max = (256, 256, 64)
   code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[int(x)]})", "l": lambda x: f"i32(lindex.{'xyz'[int(x)]})"}
+  def __init__(self): self.compiler = Compiler()
   extra_matcher = wgsl_matcher
   supports_float4 = False
   barrier = "workgroupBarrier();"

tinygrad/runtime/ops_amd.py

@@ -9,15 +9,15 @@ from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar
 from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv
-from tinygrad.renderer.cstyle import AMDRenderer
+from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
 from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler
 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)
@@ -830,8 +830,8 @@ class PCIIface(PCIIfaceBase):
                                     doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
     else:
       pv = self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
-        eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0, queue=0,
-        aql=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL))
+        eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0,
+        queue=int(is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL)), aql=is_aql)
 
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
       read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')], put_value=pv)
@@ -859,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)], system=True, uncached=True)
+    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)
     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:
@@ -930,9 +930,9 @@ class AMDDevice(HCQCompiled):
     max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
 
-    compilers = CompilerSet([CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
-                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch), AMD_LLVM),
-                             CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))], ctrl_var=AMD_CC)
+    compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
+                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
+                             CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC)
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),

tinygrad/runtime/ops_cpu.py

@@ -5,10 +5,9 @@ from tinygrad.helpers import CPU_CC, CPU_LVP, CPU_LLVM
 from tinygrad.device import BufferSpec, DMACPURef, CompilerSet, CompilerPair
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
-from tinygrad.renderer.cstyle import ClangRenderer
+from tinygrad.renderer.cstyle import ClangJITRenderer
 from tinygrad.renderer.llvmir import LLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
-from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
 from tinygrad.runtime.support.compiler_mesa import LVPCompiler
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
@@ -136,6 +135,6 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = CompilerSet([CompilerPair(ClangRenderer, ClangJITCompiler), CompilerPair(LLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
-                             CompilerPair(LVPRenderer, LVPCompiler, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
+    compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(LLVMRenderer, None, ctrl_var=CPU_LLVM),
+                             CompilerPair(LVPRenderer, None, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
     super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)

tinygrad/runtime/ops_cuda.py

@@ -2,10 +2,10 @@ from __future__ import annotations
 import ctypes, functools
 from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing, CUDA_CC, CUDA_PTX
 from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPair, CompilerSet
-from tinygrad.renderer.cstyle import CUDARenderer
-from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.cstyle import CUDACUDARenderer, CUDANVCCRenderer
+from tinygrad.renderer.ptx import CUDAPTXRenderer
 from tinygrad.runtime.autogen import cuda
-from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint: disable=unused-import
 if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported
 
@@ -117,9 +117,9 @@ class CUDADevice(Compiled):
     CUDADevice.devices.append(self)
 
     from tinygrad.runtime.graph.cuda import CUDAGraph
-    compilers = CompilerSet([CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)),
-                             CompilerPair(functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch), CUDA_PTX),
-                             CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))], ctrl_var=CUDA_CC)
+    compilers = CompilerSet([CompilerPair(functools.partial(CUDACUDARenderer, self.arch), None),
+                             CompilerPair(functools.partial(CUDAPTXRenderer, self.arch), None, CUDA_PTX),
+                             CompilerPair(functools.partial(CUDANVCCRenderer, self.arch), None)], ctrl_var=CUDA_CC)
     super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)
 
   def synchronize(self):

tinygrad/runtime/ops_dsp.py

@@ -83,7 +83,7 @@ class DSPProgram:
   def __init__(self, dev:DSPDevice, name:str, lib:bytes):
     self.dev, self.lib = dev, lib
 
-  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
 
     pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
@@ -289,7 +289,7 @@ class MockDSPRenderer(DSPRenderer):
 
 class MockDSPProgram:
   def __init__(self, name:str, lib:bytes): self.lib = lib
-  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
     with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
       dsp_lib.write(self.lib)
       dsp_lib.flush()

tinygrad/runtime/ops_hip.py

@@ -2,8 +2,7 @@ import ctypes, functools
 from tinygrad.helpers import init_c_var, mv_address, init_c_struct_t, getenv
 from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.runtime.autogen import hip
-from tinygrad.runtime.support.compiler_amd import HIPCompiler
-from tinygrad.renderer.cstyle import HIPRenderer
+from tinygrad.renderer.cstyle import HIPHIPRenderer
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
 def check(status):
@@ -15,7 +14,7 @@ class HIPDevice(Compiled):
     self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
     self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
 
-    compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), functools.partial(HIPCompiler, self.arch))])
+    compilers = CompilerSet([CompilerPair(functools.partial(HIPHIPRenderer, self.arch), None)])
     super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self))
   def synchronize(self):
     check(hip.hipSetDevice(self.device_id))

tinygrad/runtime/ops_metal.py

@@ -44,7 +44,7 @@ class MetalDevice(Compiled):
     from tinygrad.runtime.graph.metal import MetalGraph
     # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph.
     # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
-    super().__init__(device, MetalAllocator(self), CompilerSet([CompilerPair(MetalRenderer, MetalCompiler), CompilerPair(MetalRenderer, Compiler)]),
+    super().__init__(device, MetalAllocator(self), CompilerSet([CompilerPair(MetalRenderer, None)]),
       functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)
 
   def synchronize(self):

tinygrad/runtime/ops_nv.py

@@ -8,9 +8,8 @@ from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU
 from tinygrad.uop.ops import sint
 from tinygrad.device import BufferSpec, CompilerPair, CompilerSet
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, NV_CC, NV_PTX, NV_NAK
-from tinygrad.renderer.ptx import PTXRenderer
-from tinygrad.renderer.cstyle import NVRenderer
-from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
+from tinygrad.renderer.ptx import CUDAPTXRenderer, NVPTXRenderer
+from tinygrad.renderer.cstyle import NVNVRenderer, CUDACUDARenderer
 from tinygrad.runtime.support.compiler_mesa import NAKCompiler
 from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa
 from tinygrad.runtime.support.elf import elf_loader
@@ -583,9 +582,9 @@ class NVDevice(HCQCompiled[HCQSignal]):
     self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
-    cucc, ptxcc = (CUDACompiler, PTXCompiler) if MOCKGPU else (NVCompiler, NVPTXCompiler)
-    compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(NVRenderer, self.arch),functools.partial(cucc, self.arch)),
-       CompilerPair(functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(ptxcc, self.arch), NV_PTX),
+    nvr, ptxr = (CUDACUDARenderer, CUDAPTXRenderer) if MOCKGPU else (NVNVRenderer, NVPTXRenderer)
+    compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(nvr, self.arch), None),
+       CompilerPair(functools.partial(ptxr, self.arch), None, NV_PTX),
        CompilerPair(functools.partial(NAKRenderer, dev=self), functools.partial(NAKCompiler, self.arch, self.max_warps_per_sm), NV_NAK)])
     super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue)
 

tinygrad/runtime/ops_python.py

@@ -213,10 +213,14 @@ class PythonProgram:
         i += 1
     return time.perf_counter() - st
 
+class PythonCompiler(Compiler):
+  def compile(self, src:str) -> bytes: return base64.b64decode(src)
+
 class PythonRenderer(Renderer):
   device = "PYTHON"
   code_for_op = python_alu
   def __init__(self):
+    self.compiler = PythonCompiler()
     match cast(str, EMULATE.value):
       case "METAL": self.device, self.tensor_cores = "METAL", tc.metal
       case "AMD": self.device, self.tensor_cores = "AMD", tc.amd_rdna3
@@ -235,9 +239,6 @@ class PythonRenderer(Renderer):
     lops = [(u.op, u.dtype, [uops.index(v) for v in u.src if u.op is not Ops.SPECIAL], u.arg) for u in uops]
     return base64.b64encode(pickle.dumps(lops)).decode()
 
-class PythonCompiler(Compiler):
-  def compile(self, src:str) -> bytes: return base64.b64decode(src)
-
 class PythonAllocator(Allocator['PythonDevice']):
   def _alloc(self, size, options): return memoryview(bytearray(size))
   def _copyin(self, dest, src:memoryview): dest[:] = src
@@ -245,4 +246,4 @@ class PythonAllocator(Allocator['PythonDevice']):
 
 class PythonDevice(Compiled):
   def __init__(self, device:str):
-    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, PythonCompiler)]), PythonProgram)
+    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, None)]), PythonProgram)

tinygrad/runtime/ops_webgpu.py

@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Allocator, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.helpers import round_up, suppress_finalizing
 from tinygrad.runtime.autogen import webgpu
@@ -215,7 +215,7 @@ class WebGpuDevice(Compiled):
     device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
       webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
 
-    super().__init__(device, WebGpuAllocator(device_res), CompilerSet([CompilerPair(WGSLRenderer, Compiler)]),
+    super().__init__(device, WebGpuAllocator(device_res), CompilerSet([CompilerPair(WGSLRenderer, None)]),
       functools.partial(WebGPUProgram, (device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))
 
   def synchronize(self):

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
+from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 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,10 +120,11 @@ 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, system=False, snooped=False, frag=0, valid=True):
-    if not system: paddr = self.adev.paddr2xgmi(paddr)
+  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)
     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, system, snooped, valid) | (paddr & 0x0000FFFFFFFFF000)
+    self.entries[entry_id] = self.adev.gmc.get_pte_flags(self.lv, table, frag, uncached, is_sys, 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
@@ -142,7 +143,7 @@ class AMMemoryManager(MemoryManager):
     self.dev.gmc.flush_tlb(ip='MM', vmid=0)
 
 class AMDev(PCIDevImplBase):
-  Version = 0xA0000006
+  Version = 0xA0000007
 
   def __init__(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None, reset_mode=False):
     self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions

tinygrad/runtime/support/am/ip.py

@@ -3,6 +3,7 @@ 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
@@ -212,14 +213,19 @@ class AM_SMU(AM_IP):
     return (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).read() if read_back_arg else None
 
 class AM_GFX(AM_IP):
-  def init_sw(self): self.xccs = len(self.adev.regs_offset[am.GC_HWIP])
+  def init_sw(self):
+    self.xccs = len(self.adev.regs_offset[am.GC_HWIP])
+    self.mqd_paddr = [self.adev.mm.palloc(0x1000 * self.xccs, zero=False, boot=True) for i in range(2)]
+    self.mqd_mc = [self.adev.paddr2mc(mqd_paddr) for mqd_paddr in self.mqd_paddr]
 
   def init_hw(self):
     # Wait for RLC autoload to complete
     while self.adev.regCP_STAT.read() != 0 and self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] != 0: pass
 
-    self._config_gfx_rs64()
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
+    if self.adev.partial_boot: return
+
+    self._config_gfx_rs64()
 
     # NOTE: Golden reg for gfx11. No values for this reg provided. The kernel just ors 0x20000000 to this reg.
     for xcc in range(self.xccs): self.adev.regTCP_CNTL.write(self.adev.regTCP_CNTL.read() | 0x20000000, inst=xcc)
@@ -265,49 +271,61 @@ class AM_GFX(AM_IP):
     if self.xccs > 1 and not self.adev.partial_boot: self.adev.psp._spatial_partition_cmd(1)
 
   def fini_hw(self):
-    for xcc in range(self.xccs):
-      self._grbm_select(me=1, pipe=0, queue=0, inst=xcc)
-      if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1: self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
-      self._grbm_select(inst=xcc)
+    # NOTE: For aqls with xccs (queue=1), will continue from the saved state.
+    for q in range(2 if self.xccs == 1 else 1):
+      for xcc in range(self.xccs):
+        self._grbm_select(me=1, pipe=0, queue=q, inst=xcc)
+        if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1: self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
+        self._grbm_select(inst=xcc)
     for xcc in range(self.xccs): self.adev.regGCVM_CONTEXT0_CNTL.write(0, inst=xcc)
 
   def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, doorbell:int, pipe:int, queue:int,
                  aql:bool) -> int:
-    for xcc in range(self.xccs if aql else 1):
-      mqd = self.adev.mm.valloc(0x1000, uncached=True, contiguous=True)
+    self._grbm_select(me=1, pipe=pipe, queue=queue, inst=0)
+    restore_queue = aql and self.xccs > 1 and self.adev.partial_boot and (self.adev.regCP_HQD_ACTIVE.read(inst=0) & 1)
+    restore_ptr = (self.adev.regCP_HQD_PQ_WPTR_LO.read(inst=0) | (self.adev.regCP_HQD_PQ_WPTR_HI.read(inst=0) << 32)) if restore_queue else 0
+    if DEBUG >= 2 and restore_queue: print(f"am {self.adev.devfmt}: GFX queue already active, continuing from saved state {restore_ptr=:#x}.")
 
+    for xcc in range(self.xccs if aql else 1):
       struct_t = getattr(am, f"struct_v{self.adev.ip_ver[am.GC_HWIP][0]}{'_compute' if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else ''}_mqd")
-      mqd_struct = struct_t(header=0xC0310800, cp_mqd_base_addr_lo=lo32(mqd.va_addr), cp_mqd_base_addr_hi=hi32(mqd.va_addr),
+      mqd_struct = struct_t(header=0xC0310800, cp_mqd_base_addr_lo=lo32(self.mqd_mc[queue] + 0x1000*xcc),
+        cp_mqd_base_addr_hi=hi32(self.mqd_mc[queue] + 0x1000*xcc), cp_hqd_pipe_priority=0x2, cp_hqd_queue_priority=0xf, cp_hqd_quantum=0x111,
         cp_hqd_persistent_state=self.adev.regCP_HQD_PERSISTENT_STATE.encode(preload_size=0x55, preload_req=1),
-        cp_hqd_pipe_priority=0x2, cp_hqd_queue_priority=0xf, cp_hqd_quantum=0x111,
         cp_hqd_pq_base_lo=lo32(ring_addr>>8), cp_hqd_pq_base_hi=hi32(ring_addr>>8),
         cp_hqd_pq_rptr_report_addr_lo=lo32(rptr_addr), cp_hqd_pq_rptr_report_addr_hi=hi32(rptr_addr),
         cp_hqd_pq_wptr_poll_addr_lo=lo32(wptr_addr), cp_hqd_pq_wptr_poll_addr_hi=hi32(wptr_addr),
         cp_hqd_pq_doorbell_control=self.adev.regCP_HQD_PQ_DOORBELL_CONTROL.encode(doorbell_offset=doorbell*2, doorbell_en=1),
         cp_hqd_pq_control=self.adev.regCP_HQD_PQ_CONTROL.encode(rptr_block_size=5, unord_dispatch=0, queue_size=(ring_size//4).bit_length()-2,
-          **({'queue_full_en':1, 'slot_based_wptr':2, 'wpp_clamp_en':1, 'no_update_rptr':xcc!=0 or self.xccs==1} if aql else {})),
+          **({'queue_full_en':1, 'slot_based_wptr':2, 'no_update_rptr':xcc!=0 or self.xccs==1} if aql else {})),
         cp_hqd_ib_control=self.adev.regCP_HQD_IB_CONTROL.encode(min_ib_avail_size=0x3), cp_hqd_hq_status0=0x20004000,
         cp_mqd_control=self.adev.regCP_MQD_CONTROL.encode(priv_state=1), cp_hqd_vmid=0, cp_hqd_aql_control=int(aql),
         cp_hqd_eop_base_addr_lo=lo32(eop_addr>>8), cp_hqd_eop_base_addr_hi=hi32(eop_addr>>8),
         cp_hqd_eop_control=self.adev.regCP_HQD_EOP_CONTROL.encode(eop_size=(eop_size//4).bit_length()-2),
-        **({'compute_tg_chunk_size':1, 'compute_current_logic_xcc_id':xcc} if aql and self.xccs > 1 else {}))
+        **({'compute_tg_chunk_size':1, 'compute_current_logic_xcc_id':xcc, 'cp_mqd_stride_size':0x1000} if aql and self.xccs > 1 else {}))
       for se in range(8 if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else 4): setattr(mqd_struct, f'compute_static_thread_mgmt_se{se}', 0xffffffff)
 
       # Copy mqd into memory
-      self.adev.vram.view(mqd.paddrs[0][0], ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
-      self.adev.gmc.flush_hdp()
-
       self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)
 
-      mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
-      for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[xcc], self.adev.regCP_HQD_PQ_WPTR_HI.addr[xcc] + 1)):
-        self.adev.wreg(reg, mqd_st_mv[0x80 + i])
-      self.adev.regCP_HQD_ACTIVE.write(0x1, inst=xcc)
+      if restore_queue:
+        for r in [self.adev.regCP_HQD_PQ_RPTR_REPORT_ADDR, self.adev.regCP_HQD_EOP_BASE_ADDR, self.adev.regCP_HQD_EOP_BASE_ADDR_HI,
+                  self.adev.regCP_HQD_PQ_RPTR_REPORT_ADDR_HI, self.adev.regCP_HQD_PQ_WPTR_POLL_ADDR, self.adev.regCP_HQD_PQ_WPTR_POLL_ADDR_HI]:
+          val = memoryview(bytes(mqd_struct)).cast('I')[0x80 + (off:=r.addr[xcc] - self.adev.regCP_MQD_BASE_ADDR.addr[xcc])]
+          self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct), fmt='I')[0x80 + off] = val
+          r.write(val, inst=xcc)
+      else:
+        self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
 
+        mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
+        for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[xcc], self.adev.regCP_HQD_PQ_WPTR_HI.addr[xcc] + 1)):
+          self.adev.wreg(reg, mqd_st_mv[0x80 + i])
+        self.adev.regCP_HQD_ACTIVE.write(0x1, inst=xcc)
+
+      self.adev.gmc.flush_hdp()
       self._grbm_select(inst=xcc)
 
       self.adev.reg(f"regCP_ME1_PIPE{pipe}_INT_CNTL").update(time_stamp_int_enable=1, generic0_int_enable=1, inst=xcc)
-    return 0
+    return restore_ptr // 16
 
   def set_clockgating_state(self):
     if hasattr(self.adev, 'regMM_ATC_L2_MISC_CG'): self.adev.regMM_ATC_L2_MISC_CG.write(enable=1, mem_ls_enable=1)
@@ -451,7 +469,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)], system=True, uncached=True, boot=True)
+      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)
     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/memory.py

@@ -1,4 +1,4 @@
-import collections, functools, dataclasses
+import collections, functools, dataclasses, enum
 from typing import Any, ClassVar
 from tinygrad.helpers import round_up, getenv
 
@@ -107,8 +107,10 @@ 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]]; uncached:bool=False; system:bool=False; snooped:bool=False # noqa: E702
+class VirtMapping: va_addr:int; size:int; paddrs:list[tuple[int, int]]; aspace:AddrSpace; uncached:bool=False; snooped:bool=False # noqa: E702
 
 class PageTableTraverseContext:
   def __init__(self, dev, pt, vaddr, create_pts=False, free_pts=False, boot=False):
@@ -190,7 +192,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]], uncached=False, system=False, snooped=False, boot=False) -> VirtMapping:
+  def map_range(self, vaddr:int, size:int, paddrs:list[tuple[int, int]], aspace:AddrSpace, uncached=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)=}"
@@ -200,11 +202,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, system=system, snooped=snooped,
+          pt.set_entry(pte_idx + pte_off, paddr + off + pte_off * pte_covers, uncached=uncached, aspace=aspace, snooped=snooped,
                        frag=self._frag_size(ctx.vaddr+off, pte_cnt * pte_covers), valid=True)
 
     self.on_range_mapped()
-    return VirtMapping(vaddr, size, paddrs, uncached=uncached, system=system, snooped=snooped)
+    return VirtMapping(vaddr, size, paddrs, aspace=aspace, uncached=uncached, 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})")
@@ -243,7 +245,7 @@ class MemoryManager:
           continue
         rem_size -= self.palloc_ranges[nxt_range][0]
 
-    return self.map_range(va, size, paddrs, uncached=uncached)
+    return self.map_range(va, size, paddrs, aspace=AddrSpace.PHYS, 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
+from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 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, system=False, snooped=False, frag=0, valid=True):
+  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, aspace=AddrSpace.PHYS, snooped=False, frag=0, valid=True):
     if not table:
-      x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if system else 0, kind=6,
+      x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if aspace is AddrSpace.SYS 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
+from tinygrad.runtime.support.memory import MemoryManager, VirtMapping, AddrSpace
 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], system=True, snooped=True, uncached=True)
+      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, 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,19 +271,23 @@ 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 not b.meta.mapping.system: self.dev_impl.mm.vfree(b.meta.mapping)
+    if b.meta.mapping.aspace is AddrSpace.PHYS: 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, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, True
+      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
     elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, LNXPCIIfaceBase):
-      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
+      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
     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, system=True, snooped=snooped, uncached=uncached)
+    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
 
 class APLPCIIfaceBase(LNXPCIIfaceBase):
   def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):

tinygrad/uop/__init__.py

@@ -27,6 +27,10 @@ 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, BINARY has bytes arg that's compiled
+  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 (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
+  (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),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), 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),

tinygrad/uop/symbolic.py

@@ -2,7 +2,7 @@
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import ConstType, dtypes, PtrDType, can_safe_cast, Invalid
+from tinygrad.dtype import ConstType, dtypes, PtrDType, can_lossless_cast, Invalid
 from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, unwrap, IMAGE, dedup
 from tinygrad.uop.decompositions import xpow
 from tinygrad.uop.divandmod import div_and_mod_symbolic
@@ -68,6 +68,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   # ** zero folding **
   (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
   (UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
+  (UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
   (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)) != UPat.var("x"),
    lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
   # ** constant folding **
@@ -97,7 +98,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
   (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
   # b.cast(a).cast(b) -> b if a preserves all values in b
-  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
+  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None),
   # ** pow **
   (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
   # positive const ** x
@@ -242,7 +243,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat(Ops.RANGE, src=UPat.var("end"), name="r")//UPat.var("end"), lambda r,end: r.const_like(0)),
   # cast/long folding
   # if the intermediate cast doesnt narrow we can do it in one cast
-  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) else None),
+  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None),
   (UPat.var('x', dtypes.ints+(dtypes.index,)).cast(dtypes.ints+(dtypes.index,), name="a").cast(name="b"),
     lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
   # try to do math in int instead of long

tinygrad/viz/index.html

@@ -124,7 +124,6 @@
     fill: rgba(26, 27, 38, 0.5);
   }
   .edgePath {
-    stroke: #4a4b57;
     fill: none;
     stroke-width: 1.4px;
   }
@@ -322,14 +321,6 @@
   td.Instruction {
     font-family: monospace;
   }
-  td.pct-row > div {
-    height: 12px;
-    width: 100%;
-    display: flex;
-  }
-  td.pct-row > div > div {
-    height: 100%;
-  }
   thead {
     position: sticky;
     top: 0;
@@ -350,12 +341,11 @@
   tr.nested-row table tr.main-row:hover {
     background-color: unset;
   }
-  tr.main-row.has-children > td:first-child {
-    white-space: pre;
+  tr.main-row.has-children > td:first-child > p {
+    display: inline-block;
   }
   tr.main-row.has-children > td:first-child::before {
     content: "▸ ";
-    display: inline-block;
     width: 1em;
     margin-left: -0.25em;
   }

tinygrad/viz/js/index.js

@@ -76,23 +76,19 @@ const drawGraph = (data) => {
     .attr("x", d => -d.width/2).attr("y", d => -d.height/2);
   const STROKE_WIDTH = 1.4;
   const labels = nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label");
-  const hasLabelDims = data.nodes[0]?.value.labelWidth != null;
-  if (hasLabelDims) labels.attr("transform", d => `translate(-${d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`);
+  labels.attr("transform", d => `translate(-${d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`);
   labels.selectAll("text").data(d => {
+    if (Array.isArray(d.label)) return [d.label];
     const ret = [[]];
-    for (const { st, color } of parseColors(d.label, defaultColor="initial")) {
-      const lines = st.split("\n");
+    for (const s of parseColors(d.label, defaultColor="initial")) {
+      const color = darkenHex(s.color, 25);
+      const lines = s.st.split("\n");
       ret.at(-1).push({ st:lines[0], color });
       for (let i=1; i<lines.length; i++) ret.push([{ st:lines[i], color }]);
     }
     return [ret];
   }).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
-    .attr("fill", d => darkenHex(d.color, 25)).text(d => d.st).attr("xml:space", "preserve");
-  // recenter after drawing texts if needed
-  if (!hasLabelDims) labels.attr("transform", (_,i,els) => {
-    const b = els[i].getBBox();
-    return `translate(${-b.x-b.width/2}, ${-b.y-b.height/2})`
-  });
+    .attr("fill", d => d.color).text(d => d.st).attr("xml:space", "preserve").style("font-family", g.graph().font);
   addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
     .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
   // draw edges
@@ -103,7 +99,7 @@ const drawGraph = (data) => {
     points.unshift(intersectRect(g.node(e.v), points[0]));
     points.push(intersectRect(g.node(e.w), points[points.length-1]));
     return line(points);
-  }).attr("marker-end", "url(#arrowhead)");
+  }).attr("marker-end", "url(#arrowhead)").attr("stroke", e => g.edge(e).color || "#4a4b57");
 }
 
 // ** UOp graph
@@ -114,12 +110,12 @@ async function initWorker() {
   workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
 }
 
-function renderDag(graph, additions, recenter, layoutOpts) {
+function renderDag(layoutSpec, { recenter }) {
   // start calculating the new layout (non-blocking)
   updateProgress(Status.STARTED, "Rendering new graph...");
   if (worker != null) worker.terminate();
   worker = new Worker(workerUrl);
-  worker.postMessage({graph, additions, opts:layoutOpts });
+  worker.postMessage(layoutSpec);
   worker.onmessage = (e) => {
     displaySelection("#graph");
     updateProgress(Status.COMPLETE);
@@ -158,8 +154,7 @@ const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
 
 const colorScheme = {TINY:new Map([["Schedule","#1b5745"],["get_program","#1d2e62"],["compile","#63b0cd"],["DEFAULT","#354f52"]]),
   DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
-  BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"], SE:new Map([["OCC", "#101725"], ["INST", "#0A2042"]]),
-  CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
+  BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"], SE:new Map([["OCC", "#101725"], ["INST", "#0A2042"]]),}
 const cycleColors = (lst, i) => lst[i%lst.length];
 
 const rescaleTrack = (source, tid, k) => {
@@ -635,9 +630,6 @@ hljs.registerLanguage("cpp", (hljs) => ({
   ...hljs.getLanguage('cpp'),
   contains: [{ begin: '\\b(?:float|half)[0-9]+\\b', className: 'type' }, ...hljs.getLanguage('cpp').contains]
 }));
-hljs.registerLanguage("amdgpu", (hljs) => ({
-  contains: [hljs.COMMENT("//", "$"), { begin:/\b(?:s_|v_|global_|buffer_|scratch_|flat_|ds_)[a-z0-9_]*\b/, className:"code" }]
-}));
 
 async function fetchValue(path) {
   const res = await fetch(path);
@@ -799,23 +791,17 @@ async function main() {
           }
           const td = tr.append("td").classed(ret.cols[i], true);
           // string format scalar values
-          if (!Array.isArray(value)) { td.text(typeof value === "string" ? value : ret.cols[i] === "Duration" ? formatMicroseconds(value) : formatUnit(value)); continue; }
-          // display arrays in a bar graph
-          td.classed("pct-row", true);
-          const bar = td.append("div");
-          value.forEach(([k, v, width]) => bar.append("div").style("width", width+"%").attr("title", `${ret.cols[i].labels[k]} ${v}`)
-            .style("background", cycleColors(colorScheme.CATEGORICAL, parseInt(k))))
+          td.append(() => typeof value === "string" ? colored(value) : d3.create("p").text(ret.cols[i] === "Duration" ? formatMicroseconds(value) : formatUnit(value)).node());
         }
       }
       return table;
     }
-    if (ret.cols != null) {
-      renderTable(root, ret);
-    } else root.append(() => codeBlock(ret.src, ret.lang));
+    if (ret.cols != null) renderTable(root, ret);
+    else if (ret.data != null) renderDag(ret, { recenter:true });
+    else if (ret.src != null) root.append(() => codeBlock(ret.src, ret.lang));
     ret.metadata?.forEach(m => {
-      if (Array.isArray(m)) return metadata.appendChild(tabulate(m.map(({ label, value, idx }) => {
-        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, idx)).style("width", "100%").style("height", "100%");
-        return [label.trim(), div.text(typeof value === "string" ? value : formatUnit(value)).node()];
+      if (Array.isArray(m)) return metadata.appendChild(tabulate(m.map(({ label, value }) => {
+        return [label.trim(), typeof value === "string" ? value : formatUnit(value)];
       })).node());
       metadata.appendChild(codeBlock(m.src)).classList.add("full-height")
     });
@@ -842,7 +828,7 @@ async function main() {
   if (ret.length === 0) return;
   // ** center graph
   const data = ret[currentRewrite];
-  const render = (opts) => renderDag(data.graph, data.changed_nodes ?? [], currentRewrite === 0, opts);
+  const render = (opts) => renderDag({ data, opts }, { recenter:currentRewrite === 0 });
   render({ showIndexing:toggle.checked });
   toggle.onchange = (e) => render({ showIndexing:e.target.checked });
   // ** right sidebar metadata

tinygrad/viz/js/worker.js

@@ -1,27 +1,57 @@
 const NODE_PADDING = 10;
+const rectDims = (lw, lh) => ({ width:lw+NODE_PADDING*2, height:lh+NODE_PADDING*2, labelWidth:lw, labelHeight:lh });
+
 const LINE_HEIGHT = 14;
 const canvas = new OffscreenCanvas(0, 0);
 const ctx = canvas.getContext("2d");
-ctx.font = `350 ${LINE_HEIGHT}px sans-serif`;
 
 onmessage = (e) => {
-  const { graph, additions, opts } = e.data;
-  const g = new dagre.graphlib.Graph({ compound: true });
-  g.setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
-  if (additions.length !== 0) g.setNode("addition", {label:"", labelWidth:0, labelHeight:0, className:"overlay"});
-  for (let [k, {label, src, ref, ...rest }] of Object.entries(graph)) {
+  const { data, opts } = e.data;
+  const g = new dagre.graphlib.Graph({ compound: true }).setDefaultEdgeLabel(function() { return {}; });
+  (data.blocks != null ? layoutCfg : layoutUOp)(g, data, opts);
+  postMessage(dagre.graphlib.json.write(g));
+  self.close();
+}
+
+const layoutCfg = (g, { blocks, paths, pc_table, colors }) => {
+  g.setGraph({ rankdir:"TD", font:"monospace" });
+  ctx.font = `350 ${LINE_HEIGHT}px ${g.graph().font}`;
+  // basic blocks render the assembly in nodes
+  for (const [lead, members] of Object.entries(blocks)) {
+    let [width, height, label] = [0, 0, []];
+    for (const m of members) {
+      const text = pc_table[m][0];
+      width = Math.max(width, ctx.measureText(text).width);
+      height += LINE_HEIGHT;
+      const [inst, ...operands] = text.split(" ");
+      label.push([{st:inst+" ", color:"#7aa2f7"}, {st:operands.join(" "), color:"#9aa5ce"}]);
+    }
+    g.setNode(lead, { ...rectDims(width, height), label, id:lead, color:"#1a1b26" });
+  }
+  // paths become edges between basic blocks
+  for (const [lead, value] of Object.entries(paths)) {
+    for (const [id, color] of Object.entries(value)) g.setEdge(lead, id, {label:{type:"port", text:""}, color:colors[color]});
+  }
+  dagre.layout(g);
+}
+
+const layoutUOp = (g, { graph, change }, opts) => {
+  g.setGraph({ rankdir: "LR", font:"sans-serif" });
+  ctx.font = `350 ${LINE_HEIGHT}px ${g.graph().font}`;
+  if (change?.length) g.setNode("overlay", {label:"", labelWidth:0, labelHeight:0, className:"overlay"});
+  for (const [k, {label, src, ref, color, tag }] of Object.entries(graph)) {
     // adjust node dims by label size (excluding escape codes) + add padding
     let [width, height] = [0, 0];
     for (line of label.replace(/\u001B\[(?:K|.*?m)/g, "").split("\n")) {
       width = Math.max(width, ctx.measureText(line).width);
       height += LINE_HEIGHT;
     }
-    g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, label, labelHeight:height, labelWidth:width, ref, id:k, ...rest});
+    g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag});
     // add edges
-    const edgeCounts = {}
+    const edgeCounts = {};
     for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
     for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port}});
-    if (additions.includes(parseInt(k))) g.setParent(k, "addition");
+    if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
   }
   // optionally hide nodes from the layuot
   if (!opts.showIndexing) {
@@ -31,8 +61,6 @@ onmessage = (e) => {
     }
   }
   dagre.layout(g);
-  // remove additions overlay if it's empty
-  if (!g.node("addition")?.width) g.removeNode("addition");
-  postMessage(dagre.graphlib.json.write(g));
-  self.close();
+  // remove overlay node if it's empty
+  if (!g.node("overlay")?.width) g.removeNode("overlay");
 }

tinygrad/viz/serve.py

@@ -6,7 +6,7 @@ from decimal import Decimal
 from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, TypeVar, Generator, Callable
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
-from tinygrad.helpers import printable, system, TCPServerWithReuse, HTTPRequestHandler
+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
@@ -53,7 +53,7 @@ class GraphRewriteDetails(TypedDict):
   graph: dict                            # JSON serialized UOp for this rewrite step
   uop: str                               # strigified UOp for this rewrite step
   diff: list[str]|None                   # diff of the single UOp that changed
-  changed_nodes: list[int]|None          # the changed UOp id + all its parents ids
+  change: list[int]|None                 # the new UOp id + all its parents ids
   upat: tuple[tuple[str, int], str]|None # [loc, source_code] of the matched UPat
 
 def shape_to_str(s:tuple[sint, ...]): return "(" + ','.join(srender(x) for x in s) + ")"
@@ -115,14 +115,14 @@ def _reconstruct(a:int):
 def get_full_rewrite(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
   next_sink = _reconstruct(ctx.sink)
   # in the schedule graph we don't show indexing ops (unless it's in a kernel AST or rewriting dtypes.index sink)
-  yield {"graph":uop_to_json(next_sink), "uop":pystr(next_sink), "changed_nodes":None, "diff":None, "upat":None}
+  yield {"graph":uop_to_json(next_sink), "uop":pystr(next_sink), "change":None, "diff":None, "upat":None}
   replaces: dict[UOp, UOp] = {}
   for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
     replaces[u0:=_reconstruct(u0_num)] = u1 = _reconstruct(u1_num)
     try: new_sink = next_sink.substitute(replaces)
     except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
     match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
-    yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":pystr(new_sink), "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
+    yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":pystr(new_sink), "change":[id(x) for x in u1.toposort() if id(x) in sink_json],
            "diff":list(difflib.unified_diff(pystr(u0).splitlines(), pystr(u1).splitlines())), "upat":(upat_loc, match_repr)}
     if not ctx.bottom_up: next_sink = new_sink
 
@@ -248,23 +248,25 @@ def load_counters(profile:list[ProfileEvent]) -> None:
       durations.setdefault(str(e.name), []).append(float(e.en-e.st))
     if isinstance(e, ProfileProgramEvent): prg_events[str(e.name)] = e
     if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
-  ctxs.append({"name":"All Counters", "steps":[create_step("PMC", ("/all-pmc", len(ctxs), 0), \
-      (durations, {k:v[ProfilePMCEvent][0] for k,v in counter_events.items()}))]})
+  if len(counter_events) == 0: return None
+  ctxs.append({"name":"All Counters", "steps":[create_step("PMC", ("/all-pmc", len(ctxs), 0), (durations, all_counters:={}))]})
   run_number = {n:0 for n,_ in counter_events}
-  for k,v in counter_events.items():
-    prg = trace.keys[r].ret if (r:=ref_map.get(k[0])) else None
-    name = prg.name if prg is not None else k[0]
-    run_number[k[0]] += 1
+  for (k, tag),v in counter_events.items():
+    # use the colored name if it exists
+    name = trace.keys[r].ret.name if (r:=ref_map.get(k)) is not None else k
+    run_number[k] += 1
     steps:list[dict] = []
-    if (pmc:=v.get(ProfilePMCEvent)): steps.append(create_step("PMC", ("/prg-pmc", len(ctxs), len(steps)), pmc))
+    if (pmc:=v.get(ProfilePMCEvent)):
+      steps.append(create_step("PMC", ("/prg-pmc", len(ctxs), len(steps)), pmc))
+      all_counters[(name, run_number[k], k)] = pmc[0]
     if (sqtt:=v.get(ProfileSQTTEvent)):
       # to decode a SQTT trace, we need the raw stream, program binary and device properties
-      steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), (k, [*sqtt, prg_events[k[0]], dev_events[sqtt[0].device]])))
+      steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), ((k, tag), [*sqtt, prg_events[k], dev_events[sqtt[0].device]])))
       if getenv("SQTT_PARSE"):
         # run our decoder on startup, we don't use this since it only works on gfx11
         from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
         for e in sqtt: parse_sqtt_print_packets(e.blob)
-    ctxs.append({"name":f"Exec {name} n{run_number[k[0]]}", "steps":steps})
+    ctxs.append({"name":f"Exec {name} n{run_number[k]}", "steps":steps})
 
 # ** SQTT OCC only unpacks wave start, end time and SIMD location
 
@@ -337,26 +339,6 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
 
 # ** Assembly static analyzers
 
-def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
-  target_args = f"-mtriple={mtriple} -mcpu={mcpu}"
-  # disassembly output can include headers / metadata, skip if llvm-mca can't parse those lines
-  data = json.loads(system("llvm-mca -skip-unsupported-instructions=parse-failure --json -"+target_args, input=asm.encode()))
-  cr = data["CodeRegions"][0]
-  resource_labels = [repr(x)[1:-1] for x in data["TargetInfo"]["Resources"]]
-  rows:list = [[instr] for instr in cr["Instructions"]]
-  # add scheduler estimates
-  for info in cr["InstructionInfoView"]["InstructionList"]: rows[info["Instruction"]].append(info["Latency"])
-  # map per instruction resource usage
-  instr_usage:dict[int, dict[int, int]] = {}
-  for d in cr["ResourcePressureView"]["ResourcePressureInfo"]:
-    instr_usage.setdefault(i:=d["InstructionIndex"], {}).setdefault(r:=d["ResourceIndex"], 0)
-    instr_usage[i][r] += d["ResourceUsage"]
-  # last row is the usage summary
-  summary = [{"idx":k, "label":resource_labels[k], "value":v} for k,v in instr_usage.pop(len(rows), {}).items()]
-  max_usage = max([sum(v.values()) for i,v in instr_usage.items() if i<len(rows)], default=0)
-  for i,usage in instr_usage.items(): rows[i].append([[k, v, (v/max_usage)*100] for k,v in usage.items()])
-  return {"rows":rows, "cols":["Instruction", "Latency", {"title":"HW Resources", "labels":resource_labels}], "metadata":[summary]}
-
 def get_stdout(f: Callable) -> str:
   buf = io.StringIO()
   try:
@@ -376,6 +358,67 @@ def amd_readelf(lib:bytes) -> list[dict]:
           ".group_segment_fixed_size":"LDS size", ".private_segment_fixed_size":"Scratch size"}
   return [{"label":label, "value":v} for k,label in keys.items() if (v:=notes["amdhsa.kernels"][0][k]) > 0]
 
+def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
+  from tinygrad.runtime.autogen import llvm
+  from tinygrad.runtime.support.elf import elf_loader
+  llvm.LLVMInitializeAMDGPUTargetInfo()
+  llvm.LLVMInitializeAMDGPUTargetMC()
+  llvm.LLVMInitializeAMDGPUAsmParser()
+  llvm.LLVMInitializeAMDGPUDisassembler()
+  # pass NULL to callbacks
+  cbs = [ctypes.cast(0, llvm.LLVMCreateDisasmCPUFeatures.argtypes[i]) for i in {5,6}]
+  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, *cbs)
+  image, sections, _ = elf_loader(lib)
+  text = next((sh.header for sh in sections if sh.name == ".text"), None)
+  assert text is not None, "no .text section found in ELF"
+  off, sz = text.sh_addr, text.sh_size
+  addr_table:dict[int, tuple[str, int]] = {}
+  out = ctypes.create_string_buffer(128)
+  cur_off = off
+  while cur_off < sz + off:
+    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
+    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
+    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
+    cur_off += instr_sz
+  return addr_table
+
+SOPP_INSTS = {"s_branch", "s_cbranch_scc0", "s_cbranch_scc1", "s_cbranch_vccz", "s_cbranch_vccnz", "s_cbranch_execz", "s_cbranch_execnz"}
+def parse_branch(asm:str) -> int|None:
+  inst, *operands = asm.split(" ")
+  if inst in SOPP_INSTS:
+    x = int(operands[0]) & 0xffff
+    return (x - 0x10000 if x & 0x8000 else x)*4
+  return None
+
+COND_TAKEN, COND_NOT_TAKEN, UNCOND = range(3)
+cfg_colors = {COND_TAKEN: "#3f7564", COND_NOT_TAKEN: "#7a4540", UNCOND: "#3b5f7e"}
+def amdgpu_cfg(lib:bytes, arch:str) -> dict:
+  # disassemble
+  pc_table = llvm_disasm(arch, lib)
+  # get leaders
+  leaders:set[int] = {next(iter(pc_table))}
+  for pc, (asm, sz) in pc_table.items():
+    if (offset:=parse_branch(asm)) is not None: leaders.update((pc+sz+offset, pc+sz))
+  # build the cfg
+  curr:int|None = None
+  blocks:dict[int, list[int]] = {}
+  paths:dict[int, dict[int, int]] = {}
+  for pc, (asm, sz) in pc_table.items():
+    if pc in leaders:
+      paths[curr:=pc] = {}
+      blocks[pc] = []
+    else: assert curr is not None, f"no basic block found for {pc}"
+    blocks[curr].append(pc)
+    # control flow ends in endpgm
+    if asm == "s_endpgm": break
+    # otherwise a basic block can have exactly one or two paths
+    nx = pc+sz
+    if (offset:=parse_branch(asm)) is not None:
+      if asm.startswith("s_branch"): paths[curr][nx+offset] = UNCOND
+      else: paths[curr].update([(nx+offset, COND_TAKEN), (nx, COND_NOT_TAKEN)])
+    elif nx in leaders: paths[curr][nx] = UNCOND
+  return {"blocks":blocks, "paths":paths, "pc_table":pc_table, "colors":cfg_colors}
+
 # ** Main render function to get the complete details about a trace event
 
 def get_render(i:int, j:int, fmt:str) -> dict:
@@ -386,22 +429,19 @@ def get_render(i:int, j:int, fmt:str) -> dict:
   if fmt == "asm":
     compiler = Device[data.device].compiler
     disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(data.src)))
-    from tinygrad.runtime.support.compiler_cpu import llvm, LLVMCompiler
-    if isinstance(compiler, LLVMCompiler):
-      return get_llvm_mca(disasm_str, ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode(),
-                          ctypes.string_at(llvm.LLVMGetTargetMachineCPU(tm)).decode())
-    metadata:list = []
+    ret:dict = {"src":disasm_str}
     if data.device.startswith("AMD"):
-      with soft_err(lambda err: metadata.append(err)):
-        metadata.append(amd_readelf(compiler.compile(data.src)))
-    return {"src":disasm_str, "lang":"amdgpu" if data.device.startswith("AMD") else None, "metadata":metadata}
+      with soft_err(lambda err: ret.update(err)):
+        metadata = amd_readelf(lib:=compiler.compile(data.src))
+        ret = {"data":amdgpu_cfg(lib, getattr(compiler, "arch")), "metadata":[metadata]}
+    return ret
   if fmt == "all-pmc":
     durations, pmc = data
-    ret:dict = {"cols":{}, "rows":[]}
-    for (prg,_),events in pmc.items():
+    ret = {"cols":{}, "rows":[]}
+    for (name, n, k),events in data[1].items():
       pmc_table = unpack_pmc(events)
       ret["cols"].update([(r[0], None) for r in pmc_table["rows"]])
-      ret["rows"].append((prg, durations[prg].pop(0), *[r[1] for r in pmc_table["rows"]]))
+      ret["rows"].append((name, durations[k][n-1], *[r[1] for r in pmc_table["rows"]]))
     ret["cols"] = ["Kernel", "Duration", *ret["cols"]]
     return ret
   if fmt == "prg-pmc": return unpack_pmc(data[0])